[KMCompiler] [ttx/npu] add DeepSeek-V4 operator moe_gating_topk_hash

[YangLong114514]

Description

Add a TTX Triton implementation of MojoMoEGatingTopKHash for DeepSeek-V4 Hash-MoE routing on Ascend 910B.

The implementation uses a TLE DSA fast path when triton.experimental.tle is available and falls back to standard Triton otherwise.

Changes

1. Added the MojoMoEGatingTopKHash core operator and Torch reference implementation.
2. Added the Ascend TTX backend operator, kernel exports, and registrations.
3. Added a TLE DSA fast path for sigmoid / sqrt_softplus with out_flag=False.
4. Added a non-TLE fallback that gathers selected logits directly from global memory.
5. Replaced random GM gather with tle.dsa.copy GM-to-UB DMA and UB-side tl.gather.
6. Reduced nonlinear computation from all experts to only the selected Top-K values.
7. Avoided full norm_out allocation and writeback on the optimized path.
8. Used the runtime Vector Core count and a strided row loop for row scheduling.
9. Added accuracy test cases and performance tests for V4-Flash and V4-Pro.

Performance

Test environment: Ascend 910B, TTX Triton backend, bfloat16.

V4-Flash

Configuration: 256 experts, k=6, routed_scaling_factor=1.5, sqrt_softplus.

Tokens	x Shape	AscendC (us)	non-TLE (us)	TLE DSA (us)	Speedup (non-TLE)	Speedup (TLE DSA)
1	[1, 256]	4.78	3.30	1.72	1.45x	2.77x
16	[16, 256]	5.06	3.29	2.85	1.54x	1.77x
64	[64, 256]	8.32	6.31	4.99	1.32x	1.67x
128	[128, 256]	9.46	8.10	6.02	1.17x	1.57x
256	[256, 256]	11.99	10.44	7.55	1.15x	1.59x
512	[512, 256]	15.62	16.63	10.89	0.94x	1.44x
1024	[1024, 256]	22.15	29.51	17.55	0.75x	1.26x
4096	[4096, 256]	59.95	102.66	54.88	0.58x	1.09x
8192	[8192, 256]	113.10	201.60	106.39	0.56x	1.06x

V4-Pro

Configuration: 384 experts, k=6, routed_scaling_factor=2.5, sqrt_softplus.

Tokens	x Shape	AscendC (us)	non-TLE (us)	TLE DSA (us)	Speedup (non-TLE)	Speedup (TLE DSA)
1	[1, 384]	2.71	3.51	1.60	0.77x	1.69x
16	[16, 384]	5.23	3.28	2.85	1.60x	1.84x
64	[64, 384]	8.49	5.52	5.15	1.54x	1.65x
128	[128, 384]	9.49	7.71	6.14	1.23x	1.55x
256	[256, 384]	12.00	11.04	7.60	1.09x	1.58x
512	[512, 384]	15.63	17.10	10.52	0.91x	1.49x
1024	[1024, 384]	22.73	30.19	17.70	0.75x	1.28x
4096	[4096, 384]	62.41	108.17	55.25	0.58x	1.13x
8192	[8192, 384]	117.68	210.56	109.75	0.56x	1.07x

• TLE DSA is faster than AscendC across all measured token counts.
• The largest observed TLE DSA speedup is 2.77x.
• The non-TLE fallback performs well for small token counts but becomes limited by random global-memory gather at larger token counts.
• TLE DSA converts the random GM access into a contiguous GM-to-UB transfer followed by UB-side gather.

Accuracy Test

All 37 accuracy test cases passed with mixed_tol=True.

Coverage includes:

• Dtype: float32, float16, bfloat16
• Norm type: softmax, sigmoid, sqrt_softplus
• Experts: 64, 128, 256, 384, 512
• Top-K: 1, 4, 6, 8, 64
• out_flag: True and False
• Routed scaling factor: 0.5, 1.0, 2.5

[gemini-code-assist]

Code Review

This pull request introduces the MojoMoEGatingTopKHash operator and its NPU-specific Triton backend implementation TTXMoEGatingTopKHash to support hash-based MoE gating for DeepSeek-V4 routing, complete with accuracy tests and performance benchmarks. The review feedback suggests key optimizations and stability enhancements, such as utilizing propagate_nan=tl.PropagateNan.ALL in tl.max to align with NPU hardware behavior, masking logits before exponentiation to avoid overflow, hoisting tle.dsa.alloc out of the loop to eliminate allocation overhead, and adopting F.softplus in the PyTorch reference implementation for better numerical stability.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.