[TRTLLM-15284][feat] add Kimi K3 SiTU MegaMoE support

[longlee0622]

What changed

• Enable the DeepGEMM MegaMoE backend for Kimi K3 routed experts while preserving TRTLLM-Gen as the default backend.
• Add a JIT-specialized SiTU activation path to the DeepGEMM FP8xFP4 MegaMoE kernel and pass Kimi activation parameters through the TRT-LLM integration.
• Keep the existing SwiGLU MegaMoE behavior for models without SiTU configuration.
• Add configuration, wiring, and routed-expert parity coverage.
• Pin the official deepseek-ai/DeepGEMM repository to the upstream SiTU merge commit from feat: add SiTU support to FP8/FP4 Mega MoE deepseek-ai/DeepGEMM#396 and update the attribution metadata.

Upstream DeepGEMM PR: deepseek-ai/DeepGEMM#396

Pinned commit: 8b1392b978f5a03c828dd1711090d7fb50958b8a

Why

Kimi K3 uses SiTU rather than SwiGLU in its routed experts. The existing MegaMoE backend only exposed the SwiGLU path, and the Kimi model forced routed experts back to TRTLLM-Gen even when MEGAMOE_DEEPGEMM was explicitly requested. This change adds the required SiTU specialization and lets an explicit MegaMoE selection reach the fused DeepGEMM dispatch, GEMM, activation, and combine path without changing the historical default.

Impact and constraints

• Applies to the existing SM100-family, BF16 activation, W4A8_MXFP4_MXFP8, expert-parallel MegaMoE path.
• Kimi shared experts remain on the existing replicated shared-MLP path.
• SiTU and SwiGLU are selected at configuration and JIT-specialization time; no activation branch is added to the kernel inner loop.
• DeepGEMM is pinned to an official upstream commit rather than a personal development fork.

Validation

The branch records the following validation of the original Kimi SiTU MegaMoE implementation:

• Targeted configuration and backend-wiring tests.
• B200 single-GPU parity with production Kimi K3 dimensions and real layer-1 expert weights.
• B200 EP8 dispatch, compute, combine smoke and parity coverage.
• B300 production-size EP8 numerical parity with 896 experts, top-16 routing, synthetic and real checkpoint weights, including hot and uniform expert-routing cases.
• Direct MegaMoE versus TRTLLM-Gen real-weight EP8 comparison: cosine similarity 0.99936855, relative L2 3.575343%.

For the official upstream DeepGEMM pin in the latest commit:

• JSON/YAML validation and the relevant TRT-LLM pre-commit checks pass.
• A fresh native build and B200 or B300 runtime validation against the pinned upstream commit are still pending.

[brnguyen2]

The overall structure is good — AUTO resolution keeps the historical default, and the isinstance guard turning create_moe's silent Cutlass fallback into a hard error for explicit MEGAMOE_DEEPGEMM is exactly right. Two correctness gaps and some process items before this merges:

1. The two inline comments on silent SwiGLU selection (explicit CUTLASS passthrough, and the nested text_config miss in _resolve_activation_config) are the important ones — both produce wrong numerics with no error.
2. The description says a fresh native build and B200/B300 runtime validation against the 2.6.1-based pin are still pending. Since the DG kernel call now passes new kwargs (situ_beta/situ_linear_beta) whose acceptance can't be verified from this repo, at least one clean build + parity run against the pinned commit should land in the PR before merge.
3. The dependency now points at a personal GitHub fork (longlee0622/DeepGEMM). The description acknowledges this; please record the concrete plan (upstream PR link or NVIDIA-org fork) so it doesn't linger — a personal fork is a single-owner risk for everyone building this branch.
4. examples/kimi_k3/eval_extra_llm_options.yaml raising moe_config.max_num_tokens 33024→131072 isn't mentioned in the description and also affects the default TRTLLM path (less MoE chunking, higher peak activation memory). Worth a sentence on why.
5. This is a nontrivial feature carrying [None] — it should reference a TRTLLM JIRA ticket per repo convention.

[brnguyen2]

Approving — the comments below are optional touch-ups, not blockers.

All previous review rounds are addressed and verified: backend rejection happens in _routed_moe_model_config before construction, SiTU config is passed explicitly to MegaMoE (with a text_config fallback in the generic resolver), the fused latent-down gate is an explicit isinstance, the NCCL fixture uses a free port and destroys the group, and the example YAML is back to 33024 with the capacity bump scoped to the private routed config.

One residual note on the new parity test (inline): the process-global _MEGA_MOE_SYMM_BUFFER_CACHE is keyed on id(ep_pg) and outlives the per-test process group, so the parametrized runs accumulate GPU-resident SymmBuffers and can theoretically hit a stale entry after address reuse.

Per the description, a fresh native build and B200/B300 runtime validation against the upstream DeepGEMM pin are still pending — fine for feat/kimi_k3, but please close that loop before this branch heads to main.