Star258
25
Code
Code
Issues5
Pull Requests20
Actions
Discussion
Wiki
Analysis
Star258
25
Merged
feat(moe): migrate gate operators #9
wangliweiAdd on Sep 17
feat(moe): migrate gate operators #9
Merged
wangliweiAdd on Sep 17
Merged fromwangliwei12/ops-tilelangfeat/moe-gate-operatorsintocann/ops-tilelangmaster
Total: 16 files changed+2129-4
MREADME.md+28-3
@@ -4,7 +4,7 @@
4 4 
5> **硬件支持范围**：本代码仓仅支持 Ascend A5 及以上架构，不保证在更早的 Ascend 架构上可编译或正确运行。5> **硬件支持范围**：本代码仓仅支持 Ascend A5 及以上架构，不保证在更早的 Ascend 架构上可编译或正确运行。
6 6 
7-当前仓库提供基础工程框架：Engram 算子族包含 `engram_fused_weight`、`engram_gate_fwd`、`engram_gate_bwd`、`engram_grad_w_reduce` 和 `engram_sinkhorn_step`，覆盖前向、反向、分片梯度归并及 Sinkhorn 行归一化。所有算子均接入分级参数生成、确定性随机种子、NPU 显存画像、A5 性能回归框架、正则筛选、xdist 多卡绑定以及分阶段测试调度。7+当前仓库提供基础工程框架：Engram 算子族覆盖前向、反向、分片梯度归并及 Sinkhorn 行归一化；MoE 算子族提供 `topk_gate`、`moe_topk_gate` 和 `normalize_weight`。算子测试接入分级参数生成、确定性随机种子、NPU 显存画像、A5 性能回归框架、正则筛选、xdist 多卡绑定以及分阶段测试调度。
8 8 
9## 设计原则9## 设计原则
10 10 
@@ -21,11 +21,13 @@ ops-tilelang/
21├── src/cann_ops_tilelang/21├── src/cann_ops_tilelang/
22│ ├── config.py # Ascend 环境与运行配置22│ ├── config.py # Ascend 环境与运行配置
23│ ├── utils.py # 通用纯 Python 工具23│ ├── utils.py # 通用纯 Python 工具
24-│ └── engram/ # Engram 前向、反向和梯度归并算子24+│ ├── engram/ # Engram 前向、反向和梯度归并算子
25+│ └── moe/ # MoE 路由和权重归一化算子
25├── tests/26├── tests/
26│ ├── testing/ # 参数生成、数值断言和 benchmark 工具27│ ├── testing/ # 参数生成、数值断言和 benchmark 工具
27│ ├── pytest_*_plugin.py # benchmark/显存/seed/xdist/warning 插件28│ ├── pytest_*_plugin.py # benchmark/显存/seed/xdist/warning 插件
28-│ └── engram/ # 算子测试、PyTorch reference、benchmark 与显存画像29+│ ├── engram/ # Engram 测试与 PyTorch reference
30+│ └── moe/ # MoE 测试与 PyTorch reference
29├── docs/ # 架构和新增算子指南31├── docs/ # 架构和新增算子指南
30├── scripts/ # 本地/CI 执行入口32├── scripts/ # 本地/CI 执行入口
31└── pyproject.toml # 构建、依赖和工具配置33└── pyproject.toml # 构建、依赖和工具配置
@@ -115,6 +117,24 @@ row_norm, col_sumsq = engram_sinkhorn_step(matrix, row_scale, col_scale, eps=1e-
115 117 
116三个输入必须位于同一 NPU、连续且为 FP32，并且不能共享存储。`matrix` 的形状为 `(num_rows, head_dim)`，`head_dim` 必须是正的 256 倍数；调用方需保证 `row_scale` 包含非负的 Sinkhorn 缩放系数。已验证的典型维度为 256 和 5120。由于 `SimtVF` 需预留线程上下文，wrapper 会拒绝共享 UB 估算超过 216 KiB 的 shape。`num_rows=0` 时返回空的 `row_norm` 和全零的 `col_sumsq`。118三个输入必须位于同一 NPU、连续且为 FP32，并且不能共享存储。`matrix` 的形状为 `(num_rows, head_dim)`，`head_dim` 必须是正的 256 倍数；调用方需保证 `row_scale` 包含非负的 Sinkhorn 缩放系数。已验证的典型维度为 256 和 5120。由于 `SimtVF` 需预留线程上下文，wrapper 会拒绝共享 UB 估算超过 216 KiB 的 shape。`num_rows=0` 时返回空的 `row_norm` 和全零的 `col_sumsq`。
117 119 
120+MoE 路由算子接收连续的 NPU FP32 分数。`topk_gate` 稳定选择每个 token 的专家索引，分数相同时选择较小索引；`moe_topk_gate` 融合打分、偏置、掩码、固定路由、随机路由和逻辑到物理专家映射；`normalize_weight` 返回每行分母及归一化权重：
ChaoyangJi
ChaoyangJiChaoyangJiSep 20

这里只是举例，不是所有的算子都在这里列一下，前面已经有Engram举例了，所以这里的moe就不用加了。

likedislike
121+ 
122+```python
123+from cann_ops_tilelang import moe_topk_gate, normalize_weight, topk_gate
124+ 
125+topk_idx = topk_gate(scores, num_topk=6)
126+denominator, normalized_weights = normalize_weight(topk_weights)
127+topk_idx, topk_weights = moe_topk_gate(
128+ logits,
129+ num_topk=6,
130+ use_shared_as_routed=False,
131+ num_shared_experts=1,
132+ routed_scaling_factor=1.0,
133+ ep_rank=0,
134+ scoring_func='sigmoid',
135+)
136+```
137+ 
118## 测试138## 测试
119 139 
120```bash140```bash
@@ -135,6 +155,11 @@ pytest tests/engram/test_engram_fused_weight.py -m 'not benchmark' --refresh-mem
135pytest tests/engram -m benchmark --run-benchmark --refresh-baseline155pytest tests/engram -m benchmark --run-benchmark --refresh-baseline
136pytest tests/engram -m 'not benchmark' --refresh-mem156pytest tests/engram -m 'not benchmark' --refresh-mem
137 157 
158+# 运行或刷新 MoE 算子族
ChaoyangJi
ChaoyangJiChaoyangJiSep 20

这里只是举例，不是所有的算子都在这里列一下，前面已经有Engram举例了，所以这里的moe就不用加了。

likedislike
159+pytest tests/moe -m 'npu and not benchmark'
160+pytest tests/moe -m benchmark --run-benchmark --refresh-baseline
161+pytest tests/moe -m 'not benchmark' --refresh-mem
162+ 
138# 标准入口163# 标准入口
139./scripts/run_test.sh164./scripts/run_test.sh
140```165```
Mdocs/architecture.md+1-1
@@ -97,7 +97,7 @@ engram_fused_weight(weight_hidden, weight_embed)
97## 扩展方向97## 扩展方向
98 98 
99- `src/cann_ops_tilelang/quant/`：公共 scaling-factor layout 和 FP8/FP4 类型系统；99- `src/cann_ops_tilelang/quant/`：公共 scaling-factor layout 和 FP8/FP4 类型系统；
100-- `src/cann_ops_tilelang/moe/`：TopK、routing、expert mapping；100+- `src/cann_ops_tilelang/moe/`：在现有 TopK、routing 和 expert mapping 基础上扩展 expert dispatch/combine；
101- `src/cann_ops_tilelang/norm/`：RMSNorm、LayerNorm 和融合残差；101- `src/cann_ops_tilelang/norm/`：RMSNorm、LayerNorm 和融合残差；
102- `src/cann_ops_tilelang/attention/`：RoPE、KV Cache、prefill/decode；102- `src/cann_ops_tilelang/attention/`：RoPE、KV Cache、prefill/decode；
103- `tests/<operator>/benchmark/`：在算子测试目录内按测试文件存储 JSONL 性能基线。103- `tests/<operator>/benchmark/`：在算子测试目录内按测试文件存储 JSONL 性能基线。
Msrc/cann_ops_tilelang/__init__.py+4-0
@@ -6,6 +6,10 @@ __all__ = [
6 'engram_gate_fwd',6 'engram_gate_fwd',
7 'engram_grad_w_reduce',7 'engram_grad_w_reduce',
8 'engram_sinkhorn_step',8 'engram_sinkhorn_step',
9+ 'moe_topk_gate',
10+ 'normalize_weight',
11+ 'topk_gate',
9]12]
10 13 
11from .engram import engram_fused_weight, engram_gate_bwd, engram_gate_fwd, engram_grad_w_reduce, engram_sinkhorn_step14from .engram import engram_fused_weight, engram_gate_bwd, engram_gate_fwd, engram_grad_w_reduce, engram_sinkhorn_step
15+from .moe import moe_topk_gate, normalize_weight, topk_gate
Asrc/cann_ops_tilelang/moe/__init__.py+7-0
@@ -0,0 +1,7 @@
1+"""Mixture-of-Experts operators."""
2+ 
3+__all__ = ['moe_topk_gate', 'normalize_weight', 'topk_gate']
4+ 
5+from .moe_topk_gate import moe_topk_gate
6+from .normalize_weight import normalize_weight
7+from .topk_gate import topk_gate
Asrc/cann_ops_tilelang/moe/moe_topk_gate.py+239-0
@@ -0,0 +1,239 @@
1+"""Public interface for fused Mixture-of-Experts routing."""
2+ 
3+from __future__ import annotations
4+ 
5+import math
6+ 
7+import torch
8+ 
9+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
10+ 
11+from .moe_topk_gate_kernel import get_moe_topk_gate_kernel
12+from .scoring import ScoringFunc
13+ 
14+ 
15+def _check_npu_tensor(
16+ name: str,
17+ tensor: torch.Tensor,
18+ *,
19+ dtype: torch.dtype,
20+ device: torch.device,
21+ ndim: int,
22+ contiguous: bool = True,
23+) -> None:
24+ if not isinstance(tensor, torch.Tensor):
25+ raise TypeError(f'{name} must be a torch.Tensor.')
26+ if tensor.device.type != 'npu':
27+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
28+ if tensor.device != device:
29+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
30+ if tensor.dtype != dtype:
31+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
32+ if tensor.ndim != ndim:
33+ raise ValueError(f'{name} must have {ndim} dimensions, got {tensor.ndim}.')
34+ if contiguous and not tensor.is_contiguous():
35+ raise ValueError(f'{name} must be contiguous.')
36+ 
37+ 
38+def _shares_storage(left: torch.Tensor, right: torch.Tensor) -> bool:
39+ return (
40+ left.numel() > 0
41+ and right.numel() > 0
42+ and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr()
43+ )
44+ 
45+ 
46+def moe_topk_gate(
47+ logits: torch.Tensor,
48+ num_topk: int,
49+ use_shared_as_routed: bool,
50+ num_shared_experts: int,
51+ routed_scaling_factor: float,
52+ ep_rank: int,
53+ scoring_func: str = 'identity',
54+ mask: torch.Tensor | None = None,
55+ bias: torch.Tensor | None = None,
56+ image_bias: torch.Tensor | None = None,
57+ image_token_mask: torch.Tensor | None = None,
58+ fix_routing_mask: torch.Tensor | None = None,
59+ to_physical_map: torch.Tensor | None = None,
60+ logical_count: torch.Tensor | None = None,
61+ unmapped_topk_idx: torch.Tensor | None = None,
62+ force_random: torch.Tensor | None = None,
63+ out: tuple[torch.Tensor, torch.Tensor] | None = None,
64+) -> tuple[torch.Tensor, torch.Tensor]:
65+ """Compute stable top-k expert routing on Ascend."""
66+ 
67+ if not isinstance(logits, torch.Tensor):
68+ raise TypeError('logits must be a torch.Tensor.')
69+ if not isinstance(num_topk, int) or isinstance(num_topk, bool):
70+ raise TypeError(f'num_topk must be an int, got {type(num_topk).__name__}.')
71+ if not isinstance(use_shared_as_routed, bool):
72+ raise TypeError('use_shared_as_routed must be a bool.')
73+ if not isinstance(num_shared_experts, int) or isinstance(num_shared_experts, bool):
74+ raise TypeError(f'num_shared_experts must be an int, got {type(num_shared_experts).__name__}.')
75+ if not isinstance(routed_scaling_factor, (int, float)) or isinstance(routed_scaling_factor, bool):
76+ raise TypeError('routed_scaling_factor must be a real number.')
77+ if not math.isfinite(routed_scaling_factor):
78+ raise ValueError('routed_scaling_factor must be finite.')
79+ if not isinstance(ep_rank, int) or isinstance(ep_rank, bool):
80+ raise TypeError(f'ep_rank must be an int, got {type(ep_rank).__name__}.')
81+ if ep_rank < 0:
82+ raise ValueError(f'ep_rank must be non-negative, got {ep_rank}.')
83+ if not isinstance(scoring_func, str):
84+ raise TypeError('scoring_func must be a string.')
85+ 
86+ require_ascend()
87+ device = logits.device
88+ _check_npu_tensor('logits', logits, dtype=torch.float32, device=device, ndim=2)
89+ num_tokens, num_routed_experts = logits.shape
90+ if num_topk <= 0 or num_topk > min(32, num_routed_experts):
91+ raise ValueError(f'num_topk must be in [1, min(32, {num_routed_experts})], got {num_topk}.')
92+ 
93+ if use_shared_as_routed:
94+ if num_shared_experts not in (1, 2):
95+ raise ValueError(f'num_shared_experts must be 1 or 2, got {num_shared_experts}.')
96+ if num_topk % num_shared_experts != 0:
97+ raise ValueError('num_topk must be divisible by num_shared_experts.')
98+ routed_group_size = num_topk // num_shared_experts
99+ if num_routed_experts % routed_group_size != 0:
100+ raise ValueError('num_routed_experts must be divisible by num_topk // num_shared_experts.')
101+ else:
102+ num_shared_experts = 0
103+ 
104+ has_bias = bias is not None
105+ if bias is not None:
106+ _check_npu_tensor('bias', bias, dtype=torch.float32, device=device, ndim=1)
107+ if bias.shape != (num_routed_experts,):
108+ raise ValueError(f'bias must have shape {(num_routed_experts,)}, got {tuple(bias.shape)}.')
109+ 
110+ has_image_token_mask = image_token_mask is not None
111+ if has_image_token_mask != (image_bias is not None):
112+ raise ValueError('image_bias and image_token_mask must be both provided or both None.')
113+ if image_bias is not None and image_token_mask is not None:
114+ _check_npu_tensor('image_bias', image_bias, dtype=torch.float32, device=device, ndim=1)
115+ _check_npu_tensor('image_token_mask', image_token_mask, dtype=torch.bool, device=device, ndim=1)
116+ if image_bias.shape != (num_routed_experts,):
117+ raise ValueError(f'image_bias must have shape {(num_routed_experts,)}, got {tuple(image_bias.shape)}.')
118+ if image_token_mask.shape != (num_tokens,):
119+ raise ValueError(f'image_token_mask must have shape {(num_tokens,)}, got {tuple(image_token_mask.shape)}.')
120+ 
121+ if mask is not None:
122+ _check_npu_tensor('mask', mask, dtype=torch.bool, device=device, ndim=1)
123+ if mask.shape != (num_tokens,):
124+ raise ValueError(f'mask must have shape {(num_tokens,)}, got {tuple(mask.shape)}.')
125+ if force_random is not None:
126+ _check_npu_tensor('force_random', force_random, dtype=torch.bool, device=device, ndim=1)
127+ if force_random.shape != (num_tokens,):
128+ raise ValueError(f'force_random must have shape {(num_tokens,)}, got {tuple(force_random.shape)}.')
129+ 
130+ if (to_physical_map is None) != (logical_count is None):
131+ raise ValueError('to_physical_map and logical_count must be both provided or both None.')
132+ num_logical_experts = num_routed_experts + num_shared_experts
133+ if to_physical_map is not None and logical_count is not None:
134+ _check_npu_tensor('to_physical_map', to_physical_map, dtype=torch.int32, device=device, ndim=2)
135+ _check_npu_tensor('logical_count', logical_count, dtype=torch.int32, device=device, ndim=1)
136+ if to_physical_map.shape[0] != num_logical_experts or to_physical_map.shape[1] <= 0:
137+ raise ValueError(
138+ 'to_physical_map must have shape '
139+ f'({num_logical_experts}, num_duplicate_experts), got {tuple(to_physical_map.shape)}.'
140+ )
141+ if logical_count.shape != (num_logical_experts,):
142+ raise ValueError(
143+ f'logical_count must have shape {(num_logical_experts,)}, got {tuple(logical_count.shape)}.'
144+ )
145+ 
146+ if unmapped_topk_idx is not None:
147+ _check_npu_tensor(
148+ 'unmapped_topk_idx',
149+ unmapped_topk_idx,
150+ dtype=torch.int64,
151+ device=device,
152+ ndim=2,
153+ contiguous=False,
154+ )
155+ if unmapped_topk_idx.shape != (num_tokens, num_topk):
156+ raise ValueError(
157+ f'unmapped_topk_idx must have shape {(num_tokens, num_topk)}, got {tuple(unmapped_topk_idx.shape)}.'
158+ )
159+ if unmapped_topk_idx.stride(1) != 1:
160+ raise ValueError('unmapped_topk_idx must have a last-dimension stride of 1.')
161+ if fix_routing_mask is not None:
162+ if unmapped_topk_idx is None:
163+ raise ValueError('fix_routing_mask requires unmapped_topk_idx.')
164+ _check_npu_tensor('fix_routing_mask', fix_routing_mask, dtype=torch.bool, device=device, ndim=1)
165+ if fix_routing_mask.shape != (num_tokens,):
166+ raise ValueError(f'fix_routing_mask must have shape {(num_tokens,)}, got {tuple(fix_routing_mask.shape)}.')
167+ 
168+ scoring_type = ScoringFunc.from_str(scoring_func).value
169+ 
170+ output_shape = (num_tokens, num_topk + num_shared_experts)
171+ if out is None:
172+ topk_idx = torch.empty(output_shape, dtype=torch.int64, device=device)
173+ topk_weights = torch.empty(output_shape, dtype=torch.float32, device=device)
174+ else:
175+ if not isinstance(out, tuple) or len(out) != 2:
176+ raise TypeError('out must be a tuple of (topk_idx, topk_weights).')
177+ topk_idx, topk_weights = out
178+ _check_npu_tensor('out[0]', topk_idx, dtype=torch.int64, device=device, ndim=2)
179+ _check_npu_tensor('out[1]', topk_weights, dtype=torch.float32, device=device, ndim=2)
180+ if topk_idx.shape != output_shape:
181+ raise ValueError(f'out[0] must have shape {output_shape}, got {tuple(topk_idx.shape)}.')
182+ if topk_weights.shape != output_shape:
183+ raise ValueError(f'out[1] must have shape {output_shape}, got {tuple(topk_weights.shape)}.')
184+ 
185+ if _shares_storage(topk_idx, topk_weights):
186+ raise ValueError('out[0] and out[1] must not share storage.')
187+ inputs = (
188+ ('logits', logits),
189+ ('mask', mask),
190+ ('bias', bias),
191+ ('image_bias', image_bias),
192+ ('image_token_mask', image_token_mask),
193+ ('fix_routing_mask', fix_routing_mask),
194+ ('to_physical_map', to_physical_map),
195+ ('logical_count', logical_count),
196+ ('unmapped_topk_idx', unmapped_topk_idx),
197+ ('force_random', force_random),
198+ )
199+ for output_name, output in (('out[0]', topk_idx), ('out[1]', topk_weights)):
200+ for input_name, input_tensor in inputs:
201+ if input_tensor is not None and _shares_storage(output, input_tensor):
202+ raise ValueError(f'{output_name} and {input_name} must not share storage.')
203+ 
204+ if num_tokens == 0:
205+ return topk_idx, topk_weights
206+ 
207+ kernel = get_moe_topk_gate_kernel(
208+ scoring_type,
209+ num_topk,
210+ num_routed_experts,
211+ mask is not None,
212+ fix_routing_mask is not None,
213+ unmapped_topk_idx is not None,
214+ to_physical_map is not None,
215+ has_bias,
216+ has_image_token_mask,
217+ force_random is not None,
218+ )
219+ if should_print_kernel_source():
220+ print(kernel.get_kernel_source())
221+ 
222+ kernel(
223+ logits,
224+ bias,
225+ image_bias,
226+ image_token_mask,
227+ mask,
228+ fix_routing_mask,
229+ to_physical_map,
230+ logical_count,
231+ topk_idx,
232+ unmapped_topk_idx,
233+ topk_weights,
234+ force_random,
235+ num_shared_experts,
236+ float(routed_scaling_factor),
237+ ep_rank,
238+ )
239+ return topk_idx, topk_weights
Asrc/cann_ops_tilelang/moe/moe_topk_gate_kernel.py+631-0
@@ -0,0 +1,631 @@
1+import tilelang
2+from tilelang import language as T
3+from tilelang.language import simd as S
4+ 
5+from cann_ops_tilelang.config import get_num_vector_cores
6+from cann_ops_tilelang.utils import align_up, ceil_div
7+ 
8+ 
9+def _balanced_vector_reduce(values, count, mask, op):
10+ """Build a balanced SIMD reduction tree for a static vreg count."""
11+ current = [values[i] for i in range(count)]
12+ while len(current) > 1:
13+ next_level = [op(current[i], current[i + 1], mask) for i in range(0, len(current) - 1, 2)]
14+ if len(current) % 2:
15+ next_level.append(current[-1])
16+ current = next_level
17+ return current[0]
18+ 
19+ 
20+def _balanced_score_index_reduce(scores, indices, count, mask):
21+ """Reduce score/index vreg pairs, preserving the lower index on ties."""
22+ current = [(scores[i], indices[i]) for i in range(count)]
23+ while len(current) > 1:
24+ next_level = []
25+ for i in range(0, len(current) - 1, 2):
26+ left_score, left_index = current[i]
27+ right_score, right_index = current[i + 1]
28+ keep_left = S.vcmp(left_score, right_score, mask, 'ge')
29+ next_level.append((S.vmax(left_score, right_score, mask), S.vsel(left_index, right_index, keep_left)))
30+ if len(current) % 2:
31+ next_level.append(current[-1])
32+ current = next_level
33+ return current[0]
34+ 
35+ 
36+@tilelang.jit
37+def get_moe_topk_gate_kernel(
38+ scoring_type: int,
39+ num_topk: int,
40+ num_routed_experts: int,
41+ mask_exists: bool,
42+ fix_routing_mask_exists: bool,
43+ unmapped_topk_idx_exists: bool,
44+ to_physical_map_exists: bool,
45+ has_bias: bool,
46+ has_image_token_mask: bool,
47+ has_force_random: bool,
48+):
49+ num_cores = get_num_vector_cores()
50+ vector_length = 64
51+ num_stages = 2
52+ # Seven FP32 vectors cover the largest tested expert axis (448) in one
53+ # register-resident TopK scan. Smaller expert counts still allocate only
54+ # ceil_div(num_routed_experts, 64) registers.
55+ max_vregs_per_group = 7
56+ max_physical_topk = num_topk + 2
57+ physical_map_threads = 256
58+ 
59+ # Keep the register footprint bounded, and scan additional expert groups
60+ # instead of silently dropping experts beyond the first six vectors.
61+ num_vregs_per_group = ceil_div(min(num_routed_experts, max_vregs_per_group * vector_length), vector_length)
62+ experts_per_group = num_vregs_per_group * vector_length
63+ num_groups = ceil_div(num_routed_experts, experts_per_group)
64+ num_aligned_experts = num_groups * experts_per_group
65+ num_aligned_topk = align_up(max_physical_topk, vector_length)
66+ 
67+ num_tokens = T.dynamic('num_tokens')
68+ num_physical_topk = T.dynamic('num_physical_topk')
69+ num_logical_experts = T.dynamic('num_logical_experts')
70+ num_duplicate_experts = T.dynamic('num_duplicate_experts')
71+ unmapped_topk_idx_stride = T.dynamic('unmapped_topk_idx_stride')
72+ 
73+ @T.prim_func
74+ def moe_topk_gate_kernel(
75+ logits: T.Tensor[(num_tokens, num_routed_experts), T.float32],
76+ bias: T.Tensor[(num_routed_experts,), T.float32],
77+ image_bias: T.Tensor[(num_routed_experts,), T.float32],
78+ image_token_mask: T.Tensor[(num_tokens,), T.bool],
79+ mask: T.Tensor[(num_tokens,), T.bool],
80+ fix_routing_mask: T.Tensor[(num_tokens,), T.bool],
81+ to_physical_map: T.Tensor[(num_logical_experts, num_duplicate_experts), T.int32],
82+ logical_count: T.Tensor[(num_logical_experts,), T.int32],
83+ topk_idx: T.Tensor[(num_tokens, num_physical_topk), T.int64],
84+ unmapped_topk_idx: T.StridedTensor[(num_tokens, num_topk), (unmapped_topk_idx_stride, 1), T.int64],
85+ topk_weights: T.Tensor[(num_tokens, num_physical_topk), T.float32],
86+ force_random: T.Tensor[(num_tokens,), T.bool],
87+ num_shared_experts: T.int32,
88+ routed_scaling_factor: T.float32,
89+ ep_rank: T.int32,
90+ ):
91+ with T.Kernel(num_cores) as core_id:
92+ logits_ub = T.alloc_shared((num_aligned_experts,), T.float32)
93+ bias_ub = T.alloc_shared((num_aligned_experts,), T.float32)
94+ scores_ub = T.alloc_shared((num_aligned_experts,), T.float32)
95+ ranked_scores_ub = T.alloc_shared((num_aligned_experts,), T.float32)
96+ 
97+ # int64 output is written through an interleaved int32 view. The
98+ # even words are the low halves; clearing the buffer sets all high
99+ # halves to zero for non-negative expert indices.
100+ out_idx_ub = T.alloc_shared((num_aligned_topk * 2,), T.int32)
Asrc/cann_ops_tilelang/moe/normalize_weight.py+43-0
@@ -0,0 +1,43 @@
1+"""Public interface for MoE routing-weight normalization."""
2+ 
3+import torch
4+ 
5+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
6+ 
7+from .normalize_weight_kernel import get_normalize_weight_kernel
8+ 
9+ 
10+def normalize_weight(topk_weights: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
11+ """Normalize each token's FP32 top-k weights on Ascend.
12+ 
13+ ``topk_weights`` must be a contiguous two-dimensional NPU tensor with a
14+ positive top-k dimension. The returned denominator includes the kernel's
15+ ``1e-20`` numerical-stability term.
16+ """
17+ 
18+ if not isinstance(topk_weights, torch.Tensor):
19+ raise TypeError('topk_weights must be a torch.Tensor.')
20+ require_ascend()
21+ if topk_weights.device.type != 'npu':
22+ raise ValueError(f'topk_weights must be an NPU tensor, got {topk_weights.device}.')
23+ if topk_weights.dtype != torch.float32:
24+ raise TypeError(f'topk_weights must use torch.float32, got {topk_weights.dtype}.')
25+ if topk_weights.ndim != 2:
26+ raise ValueError('topk_weights must have shape (num_tokens, num_topk).')
27+ if not topk_weights.is_contiguous():
28+ raise ValueError('topk_weights must be contiguous.')
29+ 
30+ num_tokens, num_topk = topk_weights.shape
31+ if num_topk <= 0:
32+ raise ValueError(f'num_topk must be positive, got {num_topk}.')
33+ 
34+ denominator = torch.empty((num_tokens,), dtype=torch.float32, device=topk_weights.device)
35+ normalized_weights = torch.empty((num_tokens, num_topk), dtype=torch.float32, device=topk_weights.device)
36+ if num_tokens == 0:
37+ return denominator, normalized_weights
38+ 
39+ kernel = get_normalize_weight_kernel(num_topk=num_topk)
40+ if should_print_kernel_source():
41+ print(kernel.get_kernel_source())
42+ kernel(topk_weights, denominator, normalized_weights)
43+ return denominator, normalized_weights
Asrc/cann_ops_tilelang/moe/normalize_weight_kernel.py+53-0
@@ -0,0 +1,53 @@
1+import tilelang
2+from tilelang import language as T
3+ 
4+from cann_ops_tilelang.config import get_num_vector_cores
5+ 
6+ 
7+@tilelang.jit
8+def get_normalize_weight_kernel(num_topk: int):
9+ """Normalize top-k routing weights on Ascend NPU."""
10+ num_cores = get_num_vector_cores()
11+ num_tokens = T.dynamic('num_tokens')
12+ 
13+ tile = 128
14+ num_stages = 4
15+ threads = 128
16+ 
17+ @T.prim_func
18+ def normalize_weight_kernel(
19+ topk_weights: T.Tensor[(num_tokens, num_topk), T.float32],
20+ denominator: T.Tensor[(num_tokens,), T.float32],
21+ normalized_weights: T.Tensor[(num_tokens, num_topk), T.float32],
22+ ):
23+ with T.Kernel(num_cores) as core_id:
24+ for tile_id in T.Persistent(
25+ [T.ceildiv(num_tokens, tile)], num_cores, core_id, group_size=1, num_stages=num_stages
26+ ):
27+ row = tile_id * tile
28+ valid_rows = T.min(tile, num_tokens - row)
29+ weights_ub = T.alloc_shared((tile, num_topk), T.float32)
30+ norm_ub = T.alloc_shared((tile, num_topk), T.float32)
31+ denom_ub = T.alloc_shared((tile,), T.float32)
32+ T.annotate_buffer_versions({weights_ub: num_stages, norm_ub: num_stages, denom_ub: num_stages})
33+ 
34+ T.copy(topk_weights[row : row + valid_rows, :], weights_ub[:valid_rows, :])
35+ 
36+ with T.SimtVF(threads=threads):
37+ weights_frag = T.alloc_fragment((tile, num_topk), T.float32)
38+ denom_frag = T.alloc_fragment((tile,), T.float32)
39+ norm_frag = T.alloc_fragment((tile, num_topk), T.float32)
40+ 
41+ T.copy(weights_ub, weights_frag)
42+ T.reduce_sum(weights_frag, denom_frag, dim=1)
43+ for i in T.Parallel(tile):
44+ denom_frag[i] += 1e-20
45+ for i, j in T.Parallel(tile, num_topk):
46+ norm_frag[i, j] = weights_frag[i, j] / denom_frag[i]
47+ T.copy(denom_frag, denom_ub)
48+ T.copy(norm_frag, norm_ub)
49+ 
50+ T.copy(denom_ub[:valid_rows], denominator[row : row + valid_rows])
51+ T.copy(norm_ub[:valid_rows, :], normalized_weights[row : row + valid_rows, :])
52+ 
53+ return normalize_weight_kernel
Asrc/cann_ops_tilelang/moe/scoring.py+18-0
@@ -0,0 +1,18 @@
1+from enum import IntEnum
2+ 
3+ 
4+class ScoringFunc(IntEnum):
5+ SIGMOID = 0
6+ SQRTSOFTPLUS = 1
7+ SOFTMAX = 2
8+ IDENTITY = 3
9+ 
10+ def __str__(self):
11+ return self.name.lower()
12+ 
13+ @classmethod
14+ def from_str(cls, label: str):
15+ try:
16+ return cls[label.upper()]
17+ except (AttributeError, KeyError):
18+ raise ValueError(f'{label} is not a valid {cls.__name__}') from None
Asrc/cann_ops_tilelang/moe/topk_gate.py+79-0
@@ -0,0 +1,79 @@
1+"""Public interface for stable MoE top-k expert selection."""
2+ 
3+import os
4+ 
5+import torch
6+ 
7+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend, should_print_kernel_source
8+from cann_ops_tilelang.utils import align_up
9+ 
10+from .topk_gate_kernel import get_topk_gate_kernel
11+ 
12+_TOPK_MIN_CORES_ENV = 'OPS_TILELANG_TOPK_MIN_CORES'
13+ 
14+# Cache compiled kernels to bypass JIT decorator overhead on every call.
15+# Key: (num_experts, num_topk, num_cores). Value: compiled kernel callable.
16+_kernel_cache: dict = {}
17+ 
18+ 
19+def _get_topk_min_cores(num_vector_cores: int) -> int:
20+ default = max(1, num_vector_cores // 2)
21+ raw = os.getenv(_TOPK_MIN_CORES_ENV)
22+ if raw is None:
23+ return default
24+ try:
25+ return max(1, min(int(raw), num_vector_cores))
26+ except ValueError:
27+ return default
28+ 
29+ 
30+def topk_gate(scores: torch.Tensor, num_topk: int) -> torch.Tensor:
31+ """Select stable top-k expert indices for each token on Ascend.
32+ 
33+ ``scores`` must be a contiguous FP32 NPU tensor with shape
34+ ``(num_tokens, num_experts)``. Equal scores prefer the lower expert index.
35+ """
36+ 
37+ if not isinstance(scores, torch.Tensor):
38+ raise TypeError('scores must be a torch.Tensor.')
39+ if not isinstance(num_topk, int) or isinstance(num_topk, bool):
40+ raise TypeError(f'num_topk must be an int, got {type(num_topk).__name__}.')
41+ require_ascend()
42+ if scores.device.type != 'npu':
43+ raise ValueError(f'scores must be an NPU tensor, got {scores.device}.')
44+ if scores.dtype != torch.float32:
45+ raise TypeError(f'scores must use torch.float32, got {scores.dtype}.')
46+ if scores.ndim != 2:
47+ raise ValueError('scores must have shape (num_tokens, num_experts).')
48+ if not scores.is_contiguous():
49+ raise ValueError('scores must be contiguous.')
50+ 
51+ num_tokens, num_experts = scores.shape
52+ if num_topk <= 0 or num_topk > num_experts:
53+ raise ValueError(f'num_topk must be in [1, {num_experts}], got {num_topk}.')
54+ 
55+ # The padded backing allocation lets the kernel issue whole-tile stores.
56+ topk_idx = torch.empty(
57+ (align_up(num_tokens, 8), num_topk),
58+ dtype=torch.int64,
59+ device=scores.device,
60+ )[:num_tokens]
61+ if num_tokens == 0:
62+ return topk_idx
63+ 
64+ num_vector_cores = get_num_vector_cores()
65+ num_cores = min(num_vector_cores, num_tokens)
66+ min_cores = _get_topk_min_cores(num_vector_cores)
67+ if 1 < num_cores < min_cores:
68+ num_cores = min_cores
69+ 
70+ cache_key = (num_experts, num_topk, num_cores)
71+ kernel = _kernel_cache.get(cache_key)
72+ if kernel is None:
73+ kernel = get_topk_gate_kernel(num_experts, num_topk, num_cores)
74+ _kernel_cache[cache_key] = kernel
75+ 
76+ if should_print_kernel_source():
77+ print(kernel.get_kernel_source())
78+ kernel(scores, topk_idx)
79+ return topk_idx
Asrc/cann_ops_tilelang/moe/topk_gate_kernel.py+283-0
@@ -0,0 +1,283 @@
1+import tilelang
2+from tilelang import language as T
3+from tilelang.language import simd as S
4+ 
5+from cann_ops_tilelang.config import get_num_vector_cores
6+from cann_ops_tilelang.utils import align_up, ceil_div
7+ 
8+ 
9+@tilelang.jit
10+def _get_topk_gate_kernel_impl(num_experts: int, num_topk: int, num_cores: int):
11+ """Top-k expert selection per token for Ascend NPU.
12+ 
13+ Returns indices of the largest num_topk scores per row (ties broken by
14+ smaller index). Three compile-time specialized select_rows variants are
15+ selected based on num_experts: merged two-row (<=192 experts), single-row
16+ (<=384), or multi-group (>384). Output is int64 indices.
17+ """
18+ # ---------- Compile-time constants and hardware adaptation ----------
19+ num_stages = 6
20+ vector_length = 64
21+ 
22+ # Grouping: split experts into groups that fit integer vector registers.
23+ max_vregs_per_group = 6
24+ num_vregs_per_group = ceil_div(min(num_experts, max_vregs_per_group * vector_length), vector_length)
25+ experts_per_group = num_vregs_per_group * vector_length
26+ num_groups = ceil_div(num_experts, experts_per_group)
27+ if num_groups == 2:
28+ # 2-group layouts break the in-UB masking cycle: every pick after the
29+ # first repeats the initial argmax. Widening to >=3 groups of 3 vregs
30+ # runs the proven multi-group structure (verified at e=769/1025).
31+ num_vregs_per_group = ceil_div(num_vregs_per_group, 2)
32+ experts_per_group = num_vregs_per_group * vector_length
33+ num_groups = ceil_div(num_experts, experts_per_group)
34+ num_aligned_experts = num_groups * experts_per_group
35+ 
36+ # Tiling: single group (<=384 experts) loads 8 rows at once. Multi-group
37+ # tile_rows widens just enough that each row burst is >=32B and 32B-aligned
38+ # (num_topk < 4 rows are 8/16/24B; sub-32B MTE writes must be 32B aligned).
39+ if num_groups == 1:
40+ tile_rows = 8
41+ elif num_topk >= 4:
42+ tile_rows = 1
43+ else:
44+ tile_rows = 4 if num_topk % 2 else 2
45+ # acc_ub stores top‑k indices as int32; output is int64, so each index uses 2 int32 slots.
46+ acc_size = align_up(tile_rows * num_topk * 2, 32)
47+ acc_valid = tile_rows * num_topk * 2
48+ inner_unroll = min(tile_rows, 2)
49+ outer_iters = tile_rows // inner_unroll
50+ 
51+ num_tokens = T.dynamic('num_tokens')
52+ 
53+ # ---------- Utility macros ----------
54+ def _tree_reduce(count, mask, op, value_fn):
55+ """Balanced reduction tree over `count` values built by value_fn.
56+ Called from traced kernel body; no TIR loop constructs inside.
57+ """
58+ current = [value_fn(i) for i in range(count)]
59+ while len(current) > 1:
60+ nxt = []
61+ for j in range(0, len(current) - 1, 2):
62+ nxt.append(op(current[j], current[j + 1], mask))
63+ if len(current) % 2:
64+ nxt.append(current[-1])
65+ current = nxt
66+ return current[0]
67+ 
68+ @T.macro
69+ def _simd_prologue():
70+ """Initialize vector constants: all‑true mask, lane‑1 mask, -inf, int32 max, expert count."""
71+ mask_32b = S.pset(32, 'PAT_ALL')
72+ one = S.pset(32, 'PAT_VL1')
73+ neg_inf = S.vdup(-T.infinity(T.float32), T.float32, mask_32b)
74+ int_max = S.vdup(T.max_value(T.int32), T.int32, mask_32b)
75+ num_exp = S.vdup(T.int32(num_experts), T.int32, mask_32b)
76+ return mask_32b, one, neg_inf, int_max, num_exp
77+ 
78+ @T.macro
79+ def _select_one_k(scores_vec, index_vec, acc_ub, slot, mask_32b, one, neg_inf, int_max):
80+ """Single pick: find max (tie‑break by smallest index), store to acc_ub,
81+ then mask the chosen expert out with -inf.
82+ """
83+ acc_max = _tree_reduce(
84+ num_vregs_per_group, mask_32b, lambda a, b, m: S.vmax(a, b, m), lambda vi: scores_vec[vi]
85+ )
86+ mx = S.vdupv(S.vcmax(acc_max, mask_32b), mask_32b)
87+ acc_min_idx = _tree_reduce(
88+ num_vregs_per_group,
89+ mask_32b,
90+ lambda a, b, m: S.vmin(a, b, m),
91+ lambda vi: S.vsel(index_vec[vi], int_max, S.vcmp(scores_vec[vi], mx, mask_32b, 'eq')),
92+ )
93+ idx = S.vdupv(S.vcmin(acc_min_idx, mask_32b), mask_32b)
94+ S.vsts(acc_ub[slot], idx, one, 'ONEPT_B32')
95+ for i in T.unroll(num_vregs_per_group, explicit=True):
96+ scores_vec[i] = S.vsel(neg_inf, scores_vec[i], S.vcmp(index_vec[i], idx, mask_32b, 'eq'))
97+ 
98+ @T.macro
99+ def _load_group(scores_vec, index_vec, scores_ub, row, group_id, neg_inf, num_exp, mask_32b):
100+ """Load one expert group from UB into registers, masking out invalid padding."""
101+ base = group_id * experts_per_group
102+ for i in T.unroll(num_vregs_per_group, explicit=True):
103+ offset = base + i * vector_length
104+ index_vec[i] = S.vci(T.int32(offset), T.int32)
105+ scores_vec[i] = S.vsel(
106+ S.vld(scores_ub[row, offset]), neg_inf, S.vcmp(index_vec[i], num_exp, mask_32b, 'lt')
107+ )
108+ 
109+ # ---------- Three compile‑time specialized select_rows variants ----------
110+ if num_groups == 1:
111+ if num_vregs_per_group <= 3:
112+ # Variant 1: merged two‑row (≤192 experts). Process two rows per iteration.
113+ @T.macro
114+ def select_rows(base_row, scores_ub, acc_ub):
115+ for tr2 in T.serial(outer_iters):
116+ with T.SimdVF():
117+ mask_32b, one, neg_inf, int_max, num_exp = _simd_prologue()
118+ 
119+ sv0 = S.alloc_local((num_vregs_per_group,), T.float32)
120+ iv0 = S.alloc_local((num_vregs_per_group,), T.int32)
121+ sv1 = S.alloc_local((num_vregs_per_group,), T.float32)
122+ iv1 = S.alloc_local((num_vregs_per_group,), T.int32)
123+ 
124+ # Load two rows and mask invalid lanes.
125+ for i in T.unroll(num_vregs_per_group, explicit=True):
126+ sv0[i] = S.vld(scores_ub[tr2 * 2, i * vector_length])
127+ iv0[i] = S.vci(T.int32(i * vector_length), T.int32)
128+ sv1[i] = S.vld(scores_ub[tr2 * 2 + 1, i * vector_length])
129+ iv1[i] = S.vci(T.int32(i * vector_length), T.int32)
130+ sv0[i] = S.vsel(sv0[i], neg_inf, S.vcmp(iv0[i], num_exp, mask_32b, 'lt'))
131+ sv1[i] = S.vsel(sv1[i], neg_inf, S.vcmp(iv1[i], num_exp, mask_32b, 'lt'))
132+ 
133+ # Interleave top‑k selection for both rows.
134+ for k in T.serial(num_topk):
135+ _select_one_k(
136+ sv0, iv0, acc_ub, (tr2 * 2) * num_topk * 2 + 2 * k, mask_32b, one, neg_inf, int_max
137+ )
138+ _select_one_k(
139+ sv1, iv1, acc_ub, (tr2 * 2 + 1) * num_topk * 2 + 2 * k, mask_32b, one, neg_inf, int_max
140+ )
141+ else:
142+ # Variant 2: single‑row vectorized (192 < experts ≤384). Process one row at a time.
143+ @T.macro
144+ def select_rows(base_row, scores_ub, acc_ub):
145+ for tr in T.serial(tile_rows):
146+ row = base_row + tr
147+ if row < num_tokens:
148+ with T.SimdVF():
149+ mask_32b, one, neg_inf, int_max, num_exp = _simd_prologue()
150+ scores_vec = S.alloc_local((num_vregs_per_group,), T.float32)
151+ index_vec = S.alloc_local((num_vregs_per_group,), T.int32)
152+ 
153+ for i in T.unroll(num_vregs_per_group, explicit=True):
154+ scores_vec[i] = S.vld(scores_ub[tr, i * vector_length])
155+ index_vec[i] = S.vci(T.int32(i * vector_length), T.int32)
156+ 
157+ for i in T.unroll(num_vregs_per_group, explicit=True):
158+ scores_vec[i] = S.vsel(
159+ scores_vec[i], neg_inf, S.vcmp(index_vec[i], num_exp, mask_32b, 'lt')
160+ )
161+ 
162+ for k in T.serial(num_topk):
163+ _select_one_k(
164+ scores_vec,
165+ index_vec,
166+ acc_ub,
167+ tr * num_topk * 2 + 2 * k,
168+ mask_32b,
169+ one,
170+ neg_inf,
171+ int_max,
172+ )
173+ else:
174+ # Variant 3: multi‑group (>384 experts). Experts span multiple groups; process one row at a time.
175+ @T.macro
176+ def select_rows(base_row, scores_ub, acc_ub):
177+ for tr in T.serial(tile_rows):
178+ row = base_row + tr
179+ if row < num_tokens:
180+ with T.SimdVF():
181+ mask_32b, one, neg_inf, int_max, num_exp = _simd_prologue()
182+ scores_vec = S.alloc_local((num_vregs_per_group,), T.float32)
183+ index_vec = S.alloc_local((num_vregs_per_group,), T.int32)
184+ tmp_idx_vec = S.alloc_local((num_vregs_per_group,), T.int32)
185+ acc = S.alloc_var(T.float32)
186+ acc_idx = S.alloc_var(T.int32)
187+ 
188+ for k in T.serial(num_topk):
189+ # ---- Step 1: scan all groups to find global maximum ----
190+ acc = neg_inf
191+ for group_id in T.serial(num_groups):
192+ _load_group(scores_vec, index_vec, scores_ub, tr, group_id, neg_inf, num_exp, mask_32b)
193+ for i in T.unroll(num_vregs_per_group, explicit=True):
194+ acc = S.vmax(acc, scores_vec[i], mask_32b)
195+ mx = S.vdupv(S.vcmax(acc, mask_32b), mask_32b)
196+ 
197+ # ---- Step 2: scan all groups to find smallest index among maxima ----
198+ acc_idx = int_max
199+ for group_id in T.serial(num_groups):
200+ _load_group(scores_vec, index_vec, scores_ub, tr, group_id, neg_inf, num_exp, mask_32b)
201+ for i in T.unroll(num_vregs_per_group, explicit=True):
202+ tmp_idx_vec[i] = S.vsel(
203+ index_vec[i], int_max, S.vcmp(scores_vec[i], mx, mask_32b, 'eq')
204+ )
205+ for i in T.unroll(num_vregs_per_group, explicit=True):
206+ acc_idx = S.vmin(acc_idx, tmp_idx_vec[i], mask_32b)
207+ idx = S.vdupv(S.vcmin(acc_idx, mask_32b), mask_32b)
208+ 
209+ # ---- Store result ----
210+ S.vsts(acc_ub[tr * num_topk * 2 + 2 * k], idx, one, 'ONEPT_B32')
211+ 
212+ # ---- Step 3: mask out the selected expert (write back to scores_ub) ----
213+ for group_id in T.serial(num_groups):
214+ base = group_id * experts_per_group
215+ for i in T.unroll(num_vregs_per_group, explicit=True):
216+ offset = base + i * vector_length
217+ index_vec[i] = S.vci(T.int32(offset), T.int32)
218+ scores_vec[i] = S.vld(scores_ub[tr, offset])
219+ scores_vec[i] = S.vsel(
220+ neg_inf, scores_vec[i], S.vcmp(index_vec[i], idx, mask_32b, 'eq')
221+ )
222+ S.vsts(scores_ub[tr, offset], scores_vec[i], mask_32b)
223+ 
224+ S.mem_bar('VST_VLD')
225+ 
226+ # ---------- Main kernel ----------
227+ @T.prim_func
228+ def topk_gate_kernel(
229+ scores: T.Tensor[(num_tokens, num_experts), T.float32],
230+ topk_idx: T.Tensor[(num_tokens, num_topk), T.int64],
231+ ):
232+ with T.Kernel(num_cores) as core_id:
233+ # Allocate shared memory (UB) buffers
234+ scores_ub = T.alloc_shared((tile_rows, num_aligned_experts), T.float32)
235+ acc_ub = T.alloc_shared((acc_size,), T.int32)
236+ topk_idx_flat = T.view(topk_idx, (num_tokens * num_topk * 2,), T.int32)
237+ T.annotate_buffer_versions({scores_ub: num_stages, acc_ub: num_stages})
238+ 
239+ with T.SimdVF():
240+ T.clear(acc_ub) # zero out all versions; ensures high 32 bits are 0
241+ 
242+ num_outer_iters = T.ceildiv(num_tokens, tile_rows * num_cores)
243+ # Software-pipelined main loop: overlap load, compute, store
244+ for w in T.Pipelined(
245+ num_outer_iters,
246+ num_stages=num_stages,
247+ annotations={'multi_buffer_eligible': [scores_ub, acc_ub]},
248+ ):
249+ base_row = w * tile_rows * num_cores + core_id * tile_rows
250+ 
251+ # ---- Load stage: GM -> UB ----
252+ for tr in T.unroll(tile_rows, explicit=True):
253+ row = base_row + tr
254+ if row < num_tokens:
255+ T.copy(scores[row, :], scores_ub[tr, :num_experts])
256+ 
257+ # ---- Compute stage: top‑k selection within UB ----
258+ select_rows(base_row, scores_ub, acc_ub)
259+ 
260+ # ---- Store stage: UB -> GM ----
261+ # Whole-tile bursts only: sub-32B GM writes (num_topk < 4 row
262+ # tails) hit "MTE non-atomic address not aligned" and fault the
263+ # core. The wrapper pads the output so every started tile is
264+ # whole; pad rows hold garbage and are dropped by the view.
265+ if base_row < num_tokens:
266+ T.copy(
267+ acc_ub[:acc_valid],
268+ topk_idx_flat[base_row * num_topk * 2 : base_row * num_topk * 2 + acc_valid],
269+ )
270+ 
271+ return topk_gate_kernel
272+ 
273+ 
274+def get_topk_gate_kernel(num_experts: int, num_topk: int, num_cores: int = 0):
275+ """Build a kernel with num_cores resolved from the live runtime config.
276+ 
277+ The @tilelang.jit impl is keyed on the concrete num_cores value, so the
278+ resolution must happen here per call: resolving the sentinel inside the
279+ traced body would be frozen into the first compiled specialization.
280+ """
281+ if num_cores <= 0:
282+ num_cores = get_num_vector_cores()
283+ return _get_topk_gate_kernel_impl(num_experts, num_topk, num_cores)
Atests/moe/normalize_weight_ref.py+17-0
@@ -0,0 +1,17 @@
1+import torch
2+ 
3+ 
4+def normalize_weight(topk_weights: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
5+ """Normalize each token's top-k weights so they sum to one.
6+ 
7+ Args:
8+ topk_weights: Float32 tensor of shape ``(num_tokens, num_topk)``.
9+ 
10+ Returns:
11+ Tuple of ``(denominator, normalized_weights)`` where *denominator* has
12+ shape ``(num_tokens,)`` and *normalized_weights* has the same shape as
13+ the input.
14+ """
15+ denominator = torch.sum(topk_weights, dim=1)
16+ normalized_weights = topk_weights / denominator.unsqueeze(1)
17+ return denominator, normalized_weights
Atests/moe/test_moe_topk_gate.py+381-0
@@ -0,0 +1,381 @@
1+import os
2+ 
3+import pytest
4+import torch
5+ 
6+import cann_ops_tilelang
7+from cann_ops_tilelang.config import get_device
8+from cann_ops_tilelang.moe.scoring import ScoringFunc
9+from tests.moe.topk_ref import moe_topk_gate as torch_moe_topk_gate
10+from tests.testing.bench import make_param_id
11+from tests.testing.generator import generate_num_tokens, get_test_level
12+from tests.testing.numeric import assert_equal, count_bytes
13+ 
14+# Disable TileLang prints
15+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
16+pytestmark = pytest.mark.npu
17+ 
18+ 
19+_CORE_CONFIGS = [
20+ (256, 2, 8), # oss-ignore-note: V3 free
21+ (256, 1, 6), # oss-ignore-line: V33 infra V1
22+ (384, 1, 6), # oss-ignore-line: V33 full
23+]
24+ 
25+_FULL_CONFIGS = [
26+ (72, 1, 6), # oss-ignore-note: V3 lite
27+ (32, 2, 6), # oss-ignore-note: V3 MM 0
28+ (64, 2, 6), # oss-ignore-note: V3 MM 1
29+ (96, 2, 6), # oss-ignore-note: V3 MM 2
30+ (16, 2, 6), # oss-ignore-note: V3 MM 3
31+ (36, 2, 6), # oss-ignore-note: V3 MM 4
32+ (108, 2, 6), # oss-ignore-note: V3 MM 5
33+ (128, 2, 6), # oss-ignore-note: V3 MM 6
34+ (144, 2, 6), # oss-ignore-note: V3 MM 7
35+ (256, 2, 8), # oss-ignore-note: V3 free
36+ (8, 1, 2), # oss-ignore-line: V33 lite
37+ (128, 1, 4), # oss-ignore-line: V33 mid
38+ (384, 1, 6), # oss-ignore-line: V33 full
39+ (256, 1, 6), # oss-ignore-line: V33 infra V1
40+]
41+ 
42+ 
43+def generate_test_params(level: int) -> list[dict]:
44+ configs = _CORE_CONFIGS if level == 0 else _FULL_CONFIGS
45+ params = [
46+ {
47+ 'num_tokens': num_tokens,
48+ 'num_padded_tokens': num_padded_tokens,
49+ 'num_topk': num_topk,
50+ 'num_routed_experts': num_routed_experts,
51+ 'num_shared_experts': num_shared_experts,
52+ 'scoring_func': sf,
53+ 'with_bias': with_bias,
54+ 'with_image_bias': with_image_bias,
55+ 'with_force_random': with_force_random,
56+ 'use_shared_as_routed': use_shared_as_routed,
57+ }
58+ for num_tokens in generate_num_tokens(level)
59+ for num_padded_tokens in (0, 10)
60+ for num_routed_experts, num_shared_experts, num_topk in configs
61+ for sf in ScoringFunc
62+ for with_bias in ([True] if level == 0 else [True, False])
63+ for with_image_bias in ([True] if level == 0 else [True, False])
64+ for with_force_random in ([True] if level == 0 else [True, False])
65+ for use_shared_as_routed in (
66+ (False, True)
67+ if num_topk % num_shared_experts == 0 and num_routed_experts % (num_topk // num_shared_experts) == 0
68+ else (False,)
69+ )
70+ ]
71+ return params
72+ 
73+ 
74+# noinspection PyShadowingNames
75+def get_kwargs(params, fix_routing):
76+ num_tokens, num_padded_tokens = params['num_tokens'], params['num_padded_tokens']
77+ num_routed_experts = params['num_routed_experts']
78+ num_shared_experts = params['num_shared_experts']
79+ use_shared_as_routed = params['use_shared_as_routed']
80+ num_topk = params['num_topk']
81+ with_bias, with_image_bias = params['with_bias'], params['with_image_bias']
82+ with_force_random = params['with_force_random']
83+ mask = None
84+ if num_padded_tokens > 0:
85+ mask = torch.ones(num_tokens + num_padded_tokens, dtype=torch.bool, device=get_device())
86+ mask[-num_padded_tokens:] = False
87+ 
88+ to_physical_map = None
89+ logical_count = None
90+ fix_routing_mask = None
91+ force_random = None
92+ unmapped_topk_idx = torch.zeros((num_tokens + num_padded_tokens, num_topk), dtype=torch.int64, device=get_device())
93+ 
94+ bias = torch.randn(num_routed_experts, dtype=torch.float, device=get_device()) if with_bias else None
95+ image_bias = torch.randn(num_routed_experts, dtype=torch.float, device=get_device()) if with_image_bias else None
96+ image_token_mask = (
97+ torch.rand(num_tokens + num_padded_tokens, dtype=torch.float, device=get_device()) < 0.5
98+ if with_image_bias
99+ else None
100+ )
101+ 
102+ if use_shared_as_routed:
103+ to_physical_map = torch.arange(
104+ 0,
105+ num_routed_experts + num_shared_experts,
106+ dtype=torch.int,
107+ device=get_device(),
108+ )
109+ to_physical_map = to_physical_map.view(-1, 1).contiguous()
110+ logical_count = torch.ones((num_routed_experts + num_shared_experts,), dtype=torch.int, device=get_device())
111+ 
112+ if fix_routing:
113+ fix_routing_mask = torch.ones((num_tokens + num_padded_tokens,), dtype=torch.bool, device=get_device())
114+ unmapped_topk_idx = torch.randint(
115+ 0,
116+ num_routed_experts,
117+ (num_tokens + num_padded_tokens, num_topk),
118+ dtype=torch.int64,
119+ device=get_device(),
120+ )
121+ 
122+ if with_force_random:
123+ force_random = torch.rand(num_tokens + num_padded_tokens, device=get_device()) < 0.5
124+ # Force random do not appear with fix_routing, which leads to conflict
125+ if fix_routing:
126+ force_random = torch.logical_and(torch.logical_not(fix_routing_mask), force_random)
127+ 
128+ return dict(
129+ mask=mask,
130+ bias=bias,
131+ image_bias=image_bias,
132+ image_token_mask=image_token_mask,
133+ fix_routing_mask=fix_routing_mask,
134+ to_physical_map=to_physical_map,
135+ logical_count=logical_count,
136+ unmapped_topk_idx=unmapped_topk_idx,
137+ force_random=force_random,
138+ )
139+ 
140+ 
141+def assert_uniform_distribution(values: torch.Tensor, num_bins: int, description: str):
142+ """Assert that values are approximately uniformly distributed across [0, num_bins) using chi-squared test."""
143+ if values.numel() < num_bins * 10:
144+ return
145+ counts = torch.bincount(values.to(torch.int64), minlength=num_bins).float()
146+ expected = values.numel() / num_bins
147+ chi_squared = ((counts - expected) ** 2 / expected).sum().item()
148+ # Use a lenient threshold: chi-squared < 4 * degrees_of_freedom
149+ # (p-value threshold is very loose to avoid flaky tests)
150+ degrees_of_freedom = num_bins - 1
151+ assert chi_squared < 4 * degrees_of_freedom, (
152+ f'Force random expert distribution is not uniform ({description}): '
153+ f'chi_squared={chi_squared:.2f}, threshold={4 * degrees_of_freedom}, '
154+ f'counts={counts.tolist()}, expected={expected:.1f}'
155+ )
156+ 
157+ 
158+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
159+def test_moe_topk_gate(params):
160+ num_tokens, num_padded_tokens = params['num_tokens'], params['num_padded_tokens']
161+ num_routed_experts = params['num_routed_experts']
162+ num_shared_experts = params['num_shared_experts']
163+ num_topk = params['num_topk']
164+ scoring_func = params['scoring_func']
165+ use_shared_as_routed = params['use_shared_as_routed']
166+ 
167+ routed_scaling_factor = 1.5
168+ ep_rank = 0
169+ 
170+ # Correctness test
171+ for fix_routing in (False, True):
172+ logits = torch.randn(
173+ (num_tokens + num_padded_tokens, num_routed_experts), dtype=torch.float, device=get_device()
174+ )
175+ args = (
176+ logits,
177+ num_topk,
178+ use_shared_as_routed,
179+ num_shared_experts,
180+ routed_scaling_factor,
181+ ep_rank,
182+ str(scoring_func),
183+ )
184+ kwargs = get_kwargs(params, fix_routing)
185+ 
186+ def clone(tensor_dict: dict):
187+ return {k: v.clone() if k == 'unmapped_topk_idx' else v for (k, v) in tensor_dict.items()}
188+ 
189+ kwargs_ref = clone(kwargs)
190+ kwargs_out = clone(kwargs)
191+ 
192+ topk_idx, topk_weights = cann_ops_tilelang.moe_topk_gate(*args, **kwargs)
193+ num_physical_topk = num_topk + num_shared_experts if use_shared_as_routed else num_topk
194+ topk_idx_out = torch.empty(
195+ (num_tokens + num_padded_tokens, num_physical_topk), dtype=torch.int64, device=get_device()
196+ )
197+ topk_weights_out = torch.empty(
198+ num_tokens + num_padded_tokens, num_physical_topk, dtype=torch.float32, device=get_device()
199+ )
200+ cann_ops_tilelang.moe_topk_gate(*args, **kwargs_out, out=(topk_idx_out, topk_weights_out))
201+ 
202+ assert_equal(topk_idx, topk_idx_out)
203+ assert_equal(topk_weights, topk_weights_out)
204+ assert_equal(kwargs['unmapped_topk_idx'], kwargs_out['unmapped_topk_idx'])
205+ 
206+ # Compare against torch reference for tokens where force_random is False
207+ kwargs_ref_no_fr = {k: v for k, v in kwargs_ref.items() if k != 'force_random'}
208+ topk_idx_ref, topk_weights_ref = torch_moe_topk_gate(*args, **kwargs_ref_no_fr)
209+ 
210+ # Mask to select tokens that are NOT force_random (i.e. should match ref)
211+ force_random_mask = kwargs['force_random']
212+ mask = kwargs['mask']
213+ logical_count = kwargs['logical_count']
214+ if force_random_mask is not None:
215+ force_random_mask = torch.logical_and(
216+ force_random_mask,
217+ mask
218+ if mask is not None
219+ else torch.ones(num_tokens + num_padded_tokens, dtype=torch.bool, device=get_device()),
220+ )
221+ compare_mask = ~force_random_mask
222+ else:
223+ compare_mask = torch.ones(topk_idx.size(0), dtype=torch.bool, device=get_device())
224+ 
225+ unmapped_topk_idx = kwargs['unmapped_topk_idx']
226+ unmapped_topk_idx_ref = kwargs_ref['unmapped_topk_idx']
227+ 
228+ sorted_topk_idx, _ = topk_idx[compare_mask].sort(dim=1)
229+ sorted_topk_idx_ref, _ = topk_idx_ref[compare_mask].sort(dim=1)
230+ 
231+ sorted_unmapped_topk_idx = unmapped_topk_idx[compare_mask].sort(dim=1)[0]
232+ sorted_unmapped_topk_idx_ref = unmapped_topk_idx_ref[compare_mask].sort(dim=1)[0]
233+ 
234+ sorted_topk_weights, _ = topk_weights[compare_mask].sort(dim=1)
235+ sorted_topk_weights_ref, _ = topk_weights_ref[compare_mask].sort(dim=1)
236+ 
237+ assert_equal(sorted_topk_idx, sorted_topk_idx_ref)
238+ assert_equal(sorted_unmapped_topk_idx, sorted_unmapped_topk_idx_ref)
239+ assert torch.allclose(sorted_topk_weights, sorted_topk_weights_ref), (
240+ f'{sorted_topk_weights=}\n'
241+ f'{sorted_topk_weights_ref=}\n'
242+ f'{scoring_func=}, {num_routed_experts=}, {num_shared_experts=}, {num_topk=}, {fix_routing=}\n'
243+ 'Different topk weights: \n{}'.format(
244+ [
245+ (sorted_topk_weights[i], sorted_topk_weights_ref[i])
246+ for i in range(topk_weights.size(0))
247+ if not torch.equal(sorted_topk_weights[i], sorted_topk_weights_ref[i])
248+ ]
249+ )
250+ )
251+ 
252+ # Check distribution for force random tokens
253+ if force_random_mask is not None and force_random_mask.any():
254+ fr_topk_idx = topk_idx[force_random_mask]
255+ fr_topk_weights = topk_weights[force_random_mask]
256+ num_logical_experts = num_routed_experts + (num_shared_experts if use_shared_as_routed else 0)
257+ num_physical_experts = torch.sum(logical_count).item() if logical_count is not None else num_logical_experts
258+ 
259+ # Weights should be positive
260+ if fr_topk_weights.numel() > 0:
261+ assert (fr_topk_weights > 0).all(), (
262+ 'Force random tokens should have positive weights for valid experts, got: '
263+ f'{fr_topk_weights[fr_topk_weights <= 0]}'
264+ )
265+ 
266+ # Validity: random experts should in range
267+ assert torch.all(torch.logical_and(fr_topk_idx < num_physical_experts, fr_topk_idx >= 0)).item()
268+ 
269+ # Uniformity: random experts should be uniformly distributed
270+ assert_uniform_distribution(fr_topk_idx.view(-1), num_physical_experts, 'rank')
271+ 
272+ # all -1 for unmapped
273+ assert torch.all(unmapped_topk_idx[force_random_mask] == -1).item()
274+ 
275+ # Check sort stability
276+ if not params['with_force_random']:
277+ logits = torch.zeros(
278+ (num_tokens + num_padded_tokens, num_routed_experts), dtype=torch.float, device=get_device()
279+ )
280+ args = (
281+ logits,
282+ num_topk,
283+ use_shared_as_routed,
284+ num_shared_experts,
285+ routed_scaling_factor,
286+ ep_rank,
287+ str(ScoringFunc.SIGMOID),
288+ )
289+ fix_routing = False
290+ kwargs = get_kwargs(params, fix_routing)
291+ kwargs['bias'], kwargs['image_bias'], kwargs['image_token_mask'] = None, None, None
292+ topk_idx, topk_weights = cann_ops_tilelang.moe_topk_gate(*args, **kwargs)
293+ assert torch.all(
294+ topk_idx[:num_tokens, :num_topk] == torch.arange(0, num_topk, dtype=topk_idx.dtype, device=topk_idx.device)
295+ )
296+ 
297+ 
298+def generate_benchmark_func(params):
299+ num_tokens, num_padded_tokens = params['num_tokens'], params['num_padded_tokens']
300+ num_routed_experts = params['num_routed_experts']
301+ routed_scaling_factor = 1.5
302+ ep_rank = 0
303+ fix_routing = False
304+ 
305+ logits = torch.randn((num_tokens + num_padded_tokens, num_routed_experts), dtype=torch.float, device=get_device())
306+ kwargs = get_kwargs(params, fix_routing)
307+ 
308+ args = (
309+ logits,
310+ params['num_topk'],
311+ params['use_shared_as_routed'],
312+ params['num_shared_experts'],
313+ routed_scaling_factor,
314+ ep_rank,
315+ str(params['scoring_func']),
316+ )
317+ 
318+ def run_moe_topk_gate():
319+ return cann_ops_tilelang.moe_topk_gate(*args, **kwargs)
320+ 
321+ num_bytes = count_bytes(logits)
322+ return run_moe_topk_gate, num_bytes
323+ 
324+ 
325+def test_moe_topk_gate_multi_group():
326+ """Experts beyond the first six SIMD vectors must participate in TopK."""
327+ num_routed_experts = 448
328+ logits = torch.zeros((1, num_routed_experts), dtype=torch.float32, device=get_device())
329+ logits[0, -1] = 100.0
330+ 
331+ topk_idx, topk_weights = cann_ops_tilelang.moe_topk_gate(
332+ logits,
333+ num_topk=1,
334+ use_shared_as_routed=False,
335+ num_shared_experts=1,
336+ routed_scaling_factor=1.0,
337+ ep_rank=0,
338+ scoring_func='identity',
339+ )
340+ 
341+ assert_equal(topk_idx, torch.tensor([[447]], dtype=torch.int64, device=get_device()))
342+ assert_equal(topk_weights, torch.tensor([[1.0]], dtype=torch.float32, device=get_device()))
343+ 
344+ 
345+def test_moe_topk_gate_force_random_is_unique():
346+ num_tokens = 64
347+ num_routed_experts = 256
348+ num_topk = 8
349+ logits = torch.zeros((num_tokens, num_routed_experts), dtype=torch.float32, device=get_device())
350+ force_random = torch.ones((num_tokens,), dtype=torch.bool, device=get_device())
351+ 
352+ topk_idx, _ = cann_ops_tilelang.moe_topk_gate(
353+ logits,
354+ num_topk=num_topk,
355+ use_shared_as_routed=False,
356+ num_shared_experts=1,
357+ routed_scaling_factor=1.0,
358+ ep_rank=0,
359+ scoring_func='identity',
360+ force_random=force_random,
361+ )
362+ 
363+ sorted_idx = topk_idx.sort(dim=1).values
364+ assert torch.all(sorted_idx[:, 1:] != sorted_idx[:, :-1])
365+ assert torch.all(torch.logical_and(topk_idx >= 0, topk_idx < num_routed_experts))
366+ 
367+ 
368+@pytest.mark.benchmark
369+@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
370+def test_moe_topk_gate_benchmark(benchmark_timer, benchmark_record, params):
371+ func, num_bytes = generate_benchmark_func(params)
372+ t_us = benchmark_timer(func)
373+ bandwidth_gbs = num_bytes / t_us / 1e3
374+ 
375+ benchmark_record(
376+ kernel='moe_topk_gate',
377+ operation='fwd',
378+ params=params,
379+ time_us=t_us,
380+ bandwidth_gbs=bandwidth_gbs,
381+ )
Atests/moe/test_normalize_weight.py+57-0
@@ -0,0 +1,57 @@
1+import os
2+ 
3+import pytest
4+import torch
5+ 
6+import cann_ops_tilelang
7+from cann_ops_tilelang.config import get_device
8+from tests.moe.normalize_weight_ref import normalize_weight as torch_normalize_weight
9+from tests.testing.bench import make_param_id
10+from tests.testing.generator import generate_moe_params, generate_topk_idx, get_test_level
11+from tests.testing.numeric import count_bytes
12+ 
13+# Disable TileLang prints
14+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
15+pytestmark = pytest.mark.npu
16+ 
17+ 
18+def generate_test_data(params):
19+ num_topk = params['num_topk']
20+ 
21+ topk_idx = generate_topk_idx(params)
22+ num_tokens = topk_idx.shape[0]
23+ topk_weights = torch.rand((num_tokens, num_topk), dtype=torch.float32, device=get_device())
24+ 
25+ return (topk_weights, num_tokens)
26+ 
27+ 
28+@pytest.mark.parametrize('params', list(generate_moe_params(get_test_level())), ids=make_param_id)
29+def test_normalize_weight(params):
30+ topk_weights, _ = generate_test_data(params)
31+ 
32+ denominator, normalized_weights = cann_ops_tilelang.normalize_weight(topk_weights)
33+ 
34+ denom_ref, norm_ref = torch_normalize_weight(topk_weights)
35+ torch.testing.assert_close(denominator, denom_ref, atol=2e-6, rtol=0)
36+ torch.testing.assert_close(normalized_weights, norm_ref, atol=2e-6, rtol=0)
37+ 
38+ 
39+@pytest.mark.benchmark
40+@pytest.mark.parametrize('params', list(generate_moe_params(0)), ids=make_param_id)
41+def test_normalize_weight_benchmark(benchmark_timer, benchmark_record, params):
42+ topk_weights, num_tokens = generate_test_data(params)
43+ 
44+ denominator, normalized_weights = cann_ops_tilelang.normalize_weight(topk_weights)
45+ 
46+ t_us = benchmark_timer(lambda: cann_ops_tilelang.normalize_weight(topk_weights))
47+ num_bytes = count_bytes(topk_weights, denominator, normalized_weights)
48+ bandwidth_gbs = num_bytes / t_us / 1e3
49+ 
50+ params.pop('num_send_tokens')
51+ benchmark_record(
52+ kernel='normalize_weight',
53+ operation='fwd',
54+ params={'num_tokens': num_tokens, **params},
55+ time_us=t_us,
56+ bandwidth_gbs=bandwidth_gbs,
57+ )
Atests/moe/test_topk_gate.py+94-0
@@ -0,0 +1,94 @@
1+import itertools
2+import os
3+ 
4+import pytest
5+import torch
6+ 
7+import cann_ops_tilelang
8+from cann_ops_tilelang.config import get_device
9+from tests.moe.topk_ref import stable_topk as torch_stable_topk
10+from tests.testing.bench import make_param_id
11+from tests.testing.generator import generate_num_tokens, get_test_level
12+from tests.testing.numeric import assert_equal, count_bytes
13+ 
14+# Disable TileLang prints
15+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
16+pytestmark = pytest.mark.npu
17+ 
18+ 
19+_EXPERT_CONFIGS = [
20+ (72, 6), # oss-ignore-note: V3 lite
21+ (32, 6), # oss-ignore-note: V3 MM 0
22+ (64, 6), # oss-ignore-note: V3 MM 1
23+ (96, 6), # oss-ignore-note: V3 MM 2
24+ (16, 6), # oss-ignore-note: V3 MM 3
25+ (36, 6), # oss-ignore-note: V3 MM 4
26+ (108, 6), # oss-ignore-note: V3 MM 5
27+ (128, 6), # oss-ignore-note: V3 MM 6
28+ (144, 6), # oss-ignore-note: V3 MM 7
29+ (256, 8), # oss-ignore-note: V3 full
30+ (8, 2), # oss-ignore-line: V33 lite
31+ (128, 4), # oss-ignore-line: V33 mid
32+ (384, 6), # oss-ignore-line: V33 full
33+ (256, 6), # oss-ignore-line: V33 infra V1
34+] # fmt: off
35+ 
36+ 
37+def generate_test_data(params):
38+ num_tokens = params['num_tokens']
39+ num_experts = params['num_experts']
40+ scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device=get_device())
41+ return scores
42+ 
43+ 
44+def generate_test_params(level: int) -> list[dict]:
45+ return [
46+ {
47+ 'num_tokens': num_tokens,
48+ 'num_experts': num_experts,
49+ 'num_topk': num_topk,
50+ }
51+ for num_tokens in generate_num_tokens(level)
52+ # oss-ignore-note: Researcher may use various combination of num_experts and num_topk, so we try to support a wide range of configurations # noqa: E501
53+ for num_experts, num_topk in itertools.product([1, 35, 1025], [1, 9, 32])
54+ if num_experts >= num_topk
55+ ]
56+ 
57+ 
58+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
59+def test_topk_gate(params):
60+ scores = generate_test_data(params)
61+ num_topk = params['num_topk']
62+ 
63+ topk_idx_ref = torch_stable_topk(scores, num_topk)
64+ topk_idx = cann_ops_tilelang.topk_gate(scores, num_topk)
65+ assert_equal(topk_idx, topk_idx_ref)
66+ 
67+ 
68+@pytest.mark.benchmark
69+@pytest.mark.parametrize(
70+ 'params',
71+ [
72+ {'num_tokens': num_tokens, 'num_experts': num_experts, 'num_topk': num_topk}
73+ for num_tokens in generate_num_tokens(0)
74+ for num_experts, num_topk in _EXPERT_CONFIGS
75+ ],
76+ ids=make_param_id,
77+)
78+def test_topk_gate_benchmark(benchmark_timer, benchmark_record, params):
79+ scores = generate_test_data(params)
80+ num_topk = params['num_topk']
81+ 
82+ topk_idx = cann_ops_tilelang.topk_gate(scores, num_topk)
83+ 
84+ t_us = benchmark_timer(lambda: cann_ops_tilelang.topk_gate(scores, num_topk))
85+ num_bytes = count_bytes(scores, topk_idx)
86+ bandwidth_gbs = num_bytes / t_us / 1e3
87+ 
88+ benchmark_record(
89+ kernel='topk_gate',
90+ operation='fwd',
91+ params=params,
92+ time_us=t_us,
93+ bandwidth_gbs=bandwidth_gbs,
94+ )
Atests/moe/topk_ref.py+194-0
@@ -0,0 +1,194 @@
1+import torch
2+import torch.nn.functional as F
3+ 
4+from cann_ops_tilelang.moe.scoring import ScoringFunc
5+ 
6+ 
7+def stable_topk(scores: torch.Tensor, num_topk: int) -> torch.Tensor:
8+ _, sorted_indices = torch.sort(scores, dim=1, descending=True, stable=True)
9+ topk_idx = sorted_indices[:, :num_topk].contiguous()
10+ return topk_idx.as_strided(topk_idx.shape, (num_topk, 1))
11+ 
12+ 
13+def moe_topk_gate(
14+ logits: torch.Tensor,
15+ num_topk: int,
16+ use_shared_as_routed: bool,
17+ num_shared_experts: int,
18+ routed_scaling_factor: float,
19+ ep_rank: int,
20+ scoring_func: str,
21+ mask: torch.Tensor | None = None,
22+ bias: torch.Tensor | None = None,
23+ image_bias: torch.Tensor | None = None,
24+ image_token_mask: torch.Tensor | None = None,
25+ fix_routing_mask: torch.Tensor | None = None,
26+ to_physical_map: torch.Tensor | None = None,
27+ logical_count: torch.Tensor | None = None,
28+ unmapped_topk_idx: torch.Tensor | None = None,
29+ out: tuple[torch.Tensor, torch.Tensor] | None = None,
30+) -> tuple[torch.Tensor, torch.Tensor]:
31+ """PyTorch reference for top-k expert routing.
32+ 
33+ Args:
34+ logits: Raw token-expert logits, shape ``(num_tokens, num_routed_experts)``,
35+ ``float32``.
36+ bias: Optional per-expert bias added to scores before ranking, shape
37+ ``(num_routed_experts,)``, ``float32``.
38+ num_topk: Number of routed experts to select per token.
39+ use_shared_as_routed: Whether shared experts are appended as extra routed
40+ slots in the output.
41+ num_shared_experts: Number of shared experts.
42+ routed_scaling_factor: Multiplicative scaling applied to normalised weights.
43+ ep_rank: Expert-parallelism rank of this process.
44+ scoring_func: One of ``'sigmoid'``, ``'sqrtsoftplus'``, ``'softmax'``.
45+ mask: Boolean mask, shape ``(num_tokens,)``. ``True`` → route the token,
46+ ``False`` → fill outputs with ``-1`` / ``0``.
47+ fix_routing_mask: Boolean mask, shape ``(num_tokens,)``. When ``True`` for
48+ a token, use the indices already stored in *unmapped_topk_idx* instead of
49+ running the selection algorithm.
50+ to_physical_map: Logical-to-physical expert map, shape
51+ ``(num_logical_experts, num_duplicate_experts + 1)``, ``int32``.
52+ logical_count: Number of active duplicates per logical expert, shape
53+ ``(num_logical_experts,)``, ``int32``.
54+ unmapped_topk_idx: Output tensor (updated in-place) for unmasked expert
55+ indices, shape ``(num_tokens, num_topk)``, ``int64``.
56+ image_bias: Optional expert bias terms for VL of shape
57+ ``(num_routed_experts,)``, ``float32``.
58+ image_token_mask: Optional boolean mask, shape ``(num_tokens,)``.
59+ ``True`` stands for image tokens.
60+ out: Optional tuple of pre-allocated (topk_idx, topk_weights) tensors.
61+ 
62+ Returns:
63+ topk_idx: Selected physical expert indices, shape
64+ ``(num_tokens, num_topk + num_shared_experts)``, ``int64``.
65+ topk_weights: Normalised expert weights, same shape, ``float32``.
66+ """
67+ num_tokens_full, num_routed_experts = logits.shape
68+ scoring = ScoringFunc.from_str(scoring_func)
69+ 
70+ if not use_shared_as_routed:
71+ num_shared_experts = 0
72+ 
73+ num_physical_topk = num_topk + num_shared_experts
74+ num_logical_experts = num_routed_experts + num_shared_experts
75+ device = logits.device
76+ 
77+ if out is None:
78+ topk_idx_out = torch.full((num_tokens_full, num_physical_topk), -1, dtype=torch.int64, device=device)
79+ topk_weights_out = torch.zeros((num_tokens_full, num_physical_topk), dtype=torch.float32, device=device)
80+ else:
81+ topk_idx_out, topk_weights_out = out
82+ topk_idx_out.fill_(-1)
83+ topk_weights_out.zero_()
84+ 
85+ if num_tokens_full == 0:
86+ return topk_idx_out, topk_weights_out
87+ 
88+ active = mask if mask is not None else torch.ones(num_tokens_full, dtype=torch.bool, device=device)
89+ active_indices = active.nonzero(as_tuple=False).squeeze(1)
90+ num_tokens = active_indices.numel()
91+ 
92+ if num_tokens == 0:
93+ if unmapped_topk_idx is not None:
94+ unmapped_topk_idx[~active] = -1
95+ return topk_idx_out, topk_weights_out
96+ 
97+ logits_a = logits[active_indices]
98+ if bias is None:
99+ bias = torch.zeros(num_routed_experts, dtype=torch.float32, device=device)
100+ 
101+ # 1. Apply scoring function
102+ if scoring == ScoringFunc.SIGMOID:
103+ scores_wo_bias = torch.sigmoid(logits_a)
104+ elif scoring == ScoringFunc.SQRTSOFTPLUS:
105+ scores_wo_bias = F.softplus(logits_a).sqrt()
106+ elif scoring == ScoringFunc.SOFTMAX:
107+ scores_wo_bias = torch.softmax(logits_a, dim=-1)
108+ else: # Identity
109+ scores_wo_bias = logits_a
110+ 
111+ # 2. Biased scores for ranking (softmax uses raw logits + bias)
112+ # Use image_bias for image tokens, regular bias otherwise
113+ if image_token_mask is not None:
114+ image_token_mask_a = image_token_mask[active_indices]
115+ bias_for_ranking = torch.where(image_token_mask_a.unsqueeze(-1), image_bias.unsqueeze(0), bias.unsqueeze(0))
116+ else:
117+ bias_for_ranking = bias.unsqueeze(0)
118+ 
119+ scores_biased = (
120+ (logits_a + bias_for_ranking) if scoring == ScoringFunc.SOFTMAX else (scores_wo_bias + bias_for_ranking)
121+ )
122+ 
123+ # 3. Split tokens into normal routing and fix_routing
124+ fix_mask = torch.zeros(num_tokens, dtype=torch.bool, device=device)
125+ if fix_routing_mask is not None and unmapped_topk_idx is not None:
126+ fix_mask = fix_routing_mask[active_indices]
127+ 
128+ topk_idx_local = torch.full((num_tokens, num_topk), -1, dtype=torch.int64, device=device)
129+ topk_score_local = torch.zeros((num_tokens, num_topk), dtype=torch.float32, device=device)
130+ 
131+ # 4. Normal routing: select top-k experts
132+ normal_mask = ~fix_mask
133+ if normal_mask.any():
134+ normal_indices = normal_mask.nonzero(as_tuple=False).squeeze(1)
135+ sb = scores_biased[normal_indices]
136+ 
137+ selected = stable_topk(sb, num_topk)
138+ topk_idx_local[normal_indices] = selected
139+ topk_score_local[normal_indices] = scores_wo_bias[normal_indices].gather(1, selected)
140+ 
141+ # 5. Fix routing: use pre-stored indices
142+ if fix_mask.any() and unmapped_topk_idx is not None:
143+ fix_indices = fix_mask.nonzero(as_tuple=False).squeeze(1)
144+ pre_idx = unmapped_topk_idx[active_indices[fix_indices]]
145+ topk_idx_local[fix_indices] = pre_idx
146+ topk_score_local[fix_indices] = scores_wo_bias[fix_indices].gather(1, pre_idx.clamp(min=0))
147+ 
148+ # 6. Write unmapped_topk_idx
149+ if unmapped_topk_idx is not None:
150+ unmapped_topk_idx[active_indices] = topk_idx_local
151+ if mask is not None:
152+ unmapped_topk_idx[~active] = -1
153+ 
154+ # 7. Normalise weights (top-sum normalisation)
155+ # NOTE: Align with moe_topk_gate kernel implementation, which adds 1e-20 instead of
156+ # clamping. With the identity scoring function the top-k sum can be negative, and
157+ # clamping it to 1e-20 would blow up the weights.
158+ topk_sum = topk_score_local.sum(dim=-1, keepdim=True) + 1e-20
159+ topk_weights_routed = topk_score_local / topk_sum * routed_scaling_factor
160+ 
161+ # 8. Append shared-expert slots
162+ if num_shared_experts > 0:
163+ shared_idx = torch.arange(num_routed_experts, num_logical_experts, dtype=torch.int64, device=device)
164+ topk_idx_all = torch.cat([topk_idx_local, shared_idx.expand(num_tokens, -1)], dim=1)
165+ topk_weights_all = torch.cat(
166+ [
167+ topk_weights_routed,
168+ torch.ones((num_tokens, num_shared_experts), dtype=torch.float32, device=device),
169+ ],
170+ dim=1,
171+ )
172+ else:
173+ topk_idx_all, topk_weights_all = topk_idx_local, topk_weights_routed
174+ 
175+ # 9. Map logical → physical experts
176+ if to_physical_map is not None and logical_count is not None:
177+ physical_map_flat = to_physical_map.reshape(-1)
178+ map_width = to_physical_map.size(1)
179+ global_idx = active_indices.to(torch.int64)
180+ for lane in range(num_physical_topk):
181+ logical = topk_idx_all[:, lane]
182+ valid = logical >= 0
183+ logical_safe = logical.clamp(min=0)
184+ count = torch.index_select(logical_count, 0, logical_safe).to(torch.int64)
185+ dup_idx = (ep_rank + global_idx * 23333) % count
186+ flat_idx = logical_safe * map_width + dup_idx
187+ physical = torch.index_select(physical_map_flat, 0, flat_idx).to(torch.int64)
188+ topk_idx_all[:, lane] = torch.where(valid, physical, logical)
189+ 
190+ # 10. Write outputs
191+ topk_idx_out[active_indices] = topk_idx_all
192+ topk_weights_out[active_indices] = topk_weights_all
193+ 
194+ return topk_idx_out, topk_weights_out