[Triton/Gluon] Add MoE weight-gradient kernel

[WuLei-AMD]

Split from #5310 (per reviewer request to separate kernels into individual PRs).

Summary

Adds moe_wgrad: a Triton kernel that computes the weight gradient for a
MoE expert block, using the moe_align_block_size sorting interface
(sorted_token_ids, expert_ids, num_tokens_post_padded).

Independent of the mxfp8/per_token PRs — no shared dependencies,
can be reviewed and merged in any order.

Key design points

• grad must incorporate per-route gating weights for top_k > 1;
  the kernel accumulates as-is without applying gating internally
• Padding slots in sorted_token_ids must be initialized to a sentinel
  >= num_valid_tokens * top_k so the token mask treats them as invalid

Tests

• op_tests/triton_tests/moe/test_moe_wgrad.py: constant-input verification,
  linearity check, zero-padded-token edge case
• All tests pass on MI308X

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5339 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[Copilot]

🟡 Changes recommended

Top-k routes cannot apply distinct gating gradients, and benchmark padding is uninitialized.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds a Triton MoE weight-gradient kernel using sorted routing metadata.

Changes:

• Adds the kernel and public wrapper.
• Exports moe_wgrad.
• Adds correctness tests and benchmarks.

File summaries

File	Description
aiter/ops/triton/_triton_kernels/moe/moe_wgrad.py	Implements atomic weight-gradient accumulation.
aiter/ops/triton/moe/moe_wgrad.py	Adds the public launch wrapper.
aiter/ops/triton/moe/__init__.py	Exports the new operation.
op_tests/triton_tests/moe/test_moe_wgrad.py	Adds correctness tests.
op_tests/op_benchmarks/triton/bench_moe_wgrad.py	Adds representative benchmarks.

Review details

Suppressed comments (1)

aiter/ops/triton/_triton_kernels/moe/moe_wgrad.py:61

• This is a new launchable Triton kernel without the required config-aware repr. The autotune key controls winner caching but does not make trace names identify BLOCK_SIZE_N, BLOCK_SIZE_K, or launch parameters. Because the comment reports an execution issue when combining decorators, resolve that incompatibility through a safe dispatch/workaround rather than omitting the repository-required representation.

# NOTE: repr= is intentionally omitted here. Combining @triton.autotune with
# @triton.jit(repr=...) corrupts the kernel execution on current Triton versions.
# The autotune key ["N", "K"] already encodes the compiled variant identity.
@triton.jit

• Files reviewed: 5/5 changed files
• Comments generated: 4
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.