[KMCompiler]Fused moe: add thead (PPU-ZW810E) backend

[LittleShun1214]

PR Category

[ Operator | Benchmark ]

Type of Change

[ New Feature | Performance Optimization ]

Description

Add a new _thead backend implementing a pure-Triton fused MoE (fused_experts_impl) specialized for the T-Head PPU (ZW810E, compute_89 / AIU MMA instructions), modeled on the existing _metax/_mthreads vendor backends.

Key points:

• Vendor-specific kernel orchestration: thead variant of moe_align_block_size, transpose-cached weights (reusing the common permute_copy operator synchronized from FlagGems for the transpose), per-tile moe_sum reduction, no deep_gemm dependency in the production path.
• GEMM-style kernel for large M with config tiering by token count (MOE_GEMM_TUNING_MIN_TOKENS, split gemm1/gemm2 stages) and a GEMV-style kernel for small M (fused_moe_gemv_kernel, N >= 1024, row-count gated) that mimics DeepGemm's Gemvt on PPU.
• On-device tuning: stage split, config tiering and BLOCK_SIZE sweep for the PPU; GEMV config BLOCK_N=128, BLOCK_K=32, num_warps=1, num_stages=2.
• Adaptive memory: weight transpose cache with weakref-based lifetime synchronization and config-dependent capacity degradation; the FlagGems-synchronized permute_copy operator is updated with int64 indexing to support tensors with >2^31 elements (DeepSeek-like shapes), which is required by the weight transpose path.

The backend is selected via torch_device_fn (vendor probe); no change to other backends or the default dispatch.

Issue

• Associated with Feature: PPU (T-Head) backend support for fused MoE.

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

Benchmarked against the vLLM production DeepGemm path, cudagraph mode, 24 real shapes (Qwen3-6B-like: M in 1..16384, 2I in 512/2048, E=256, topk=8), 10 successful runs, speedup = mean over runs:

Shape	FlagGems (ms)	DeepGEMM (ms)	Speedup
M=1, 2I=512	0.025	0.033	1.3599
M=1, 2I=2048	0.067	0.080	1.1880
M=8, 2I=512	0.122	0.121	0.9927
M=8, 2I=2048	0.363	0.355	0.9784
M=32, 2I=512	0.309	0.327	1.0596
M=32, 2I=2048	1.01	1.19	1.1824
M=128, 2I=512	0.440	0.496	1.1276
M=128, 2I=2048	1.53	1.82	1.1889
M=512, 2I=512	0.522	0.577	1.1049
M=512, 2I=2048	1.76	2.09	1.1887
M=1036, 2I=512	0.619	0.701	1.1327
M=1036, 2I=2048	1.92	2.23	1.1624
M=2048, 2I=512	0.876	0.931	1.0623
M=2048, 2I=2048	2.80	2.92	1.0432
M=4108, 2I=512	1.44	1.55	1.0705
M=4108, 2I=2048	4.85	5.20	1.0733
M=4435, 2I=512	1.62	1.78	1.1023
M=4435, 2I=2048	5.41	6.20	1.1459
M=13421, 2I=512	4.10	4.16	1.0144
M=13421, 2I=2048	13.43	13.77	1.0253
M=13422, 2I=512	4.11	4.16	1.0131
M=13422, 2I=2048	13.44	13.76	1.0240
M=16384, 2I=512	4.92	4.90	0.9947
M=16384, 2I=2048	16.19	16.03	0.9900

Summary (mean over 10 runs):

Metric	Result
Token-weighted arithmetic-mean speedup	1.0287
Max speedup	1.3599 (M=1, 2I=512)
Min speedup	0.9791 (M=8, 2I=2048)
Shapes with speedup >= 1.0	20/24