[Fused MoE] Extend W13 with peer-indexed EP input

[taoyuanyuan]

Summary

This PR adds intra-NVLink-domain EP communication to the existing HPC-Ops block-FP8 Fused MoE. It extends the W13 input boundary from local contiguous tensors to peer-indexed tensors and returns W2 results by index to the token owner.

Each rank publishes quantized activation, scale, and routing metadata in symmetric GPU memory. During decode, W13 resolves rank-major source-row indices and consumes owner-resident activation tiles directly. Each CTA owns an independent W13 tile and performs both peer load and compute; exposing more resident CTAs lets the GPU warp scheduler hide NVLink load latency across ready work without adding a software wave scheduler or a warp-specialized producer/consumer pipeline. During prefill, routed rows are pulled into the native HPC-Ops layout before the original grouped GEMM. W2 partials are returned by index to token-owner slots and summed locally.

Why W13 owns this capability

The change extends the existing HPC-Ops low-latency execution model. The original local cp.async path turns Top-K routes into row indices and lets W13 consume indexed activation instead of launching a standalone gather. This PR extends that input resolver from one local pointer to a peer pointer table plus rank-major source rows for block-FP8 W13. Decode therefore avoids both a dispatcher-materialized receive tensor and a local gate_up_input staging tensor while preserving the original activation, W2, and reduction paths.

Local BF16 activation + Top-K -> quantize and publish owner rows -> rank-major source indices
Decode  -> peer-indexed W13 directly loads owner-resident tiles -> native activation + W2
Prefill -> pull routed rows into native layout -> original grouped W13 + activation + W2
W2 partials -> indexed return to owner slots -> owner-local Top-K sum

SGLang owns the process-lifetime workspace and selects the direct or pull consumer by phase. The public interface remains --moe-a2a-backend hpc_ops --moe-runner-backend hpc_ops.

Scope

The initial implementation targets Hopper/SM90, BF16 hidden states, 128x128 block-FP8 MoE, DP Attention with EP greater than one, and one GPU load/store-accessible NVLink domain. The first SGLang integration intentionally keeps a single-host guard. A future fabric bootstrap can apply the same W13 input contract to larger load/store-accessible NVLink domains; generic IB/RDMA transport is outside this change.

Performance reference

The serving workload uses one 8x H20-3e NVLink node, GLM-5.2-FP8, DP8 + DP-Attention + EP8, 8K input / 1K forced output, and decode CUDA Graphs. Each point was measured three times after exact-shape warmup and cache flush. The table reports the median of three run-level mean metrics.

Batch	Backend	TTFT	TPOT	E2E	Output tok/s
8	DeepEP + DeepGEMM	6973.52 ms	28.75 ms	35686.14 ms	223.42
8	MegaMoE	6274.91 ms	24.13 ms	30346.29 ms	262.62
8	HPC-Ops peer-indexed	6445.28 ms	24.12 ms	30538.61 ms	261.02
32	DeepEP + DeepGEMM	16816.96 ms	42.37 ms	59169.15 ms	538.41
32	MegaMoE	15066.31 ms	38.19 ms	53216.90 ms	598.31
32	HPC-Ops peer-indexed	15658.29 ms	39.57 ms	55191.77 ms	577.02

HPC-Ops peer-indexed matches MegaMoE at B8 and remains within 4% TPOT/E2E at B32. It also improves over the community-default DeepEP + DeepGEMM reference in this workload. These are end-to-end backend results; the comparison is a performance-safety check rather than a claim that one component universally replaces another.

Follow-up optimization

Future work will explore higher CTA residency across Decode shapes so that more independent W13 tiles remain in flight and peer-load latency can be hidden more effectively. Any production change will remain gated by profiling and a controlled end-to-end A/B.

Validation and reproduction

• Direct and pull consumers match the materialized reference across prepare, epoch, W13, W2, indexed return, and owner sum.
• Tests cover uneven routes, idle owners, repeated generations, CUDA Graph replay, and buffer reuse.
• All measured requests produced exactly 1000 validated output tokens with zero prompt-cache hits.
• Focused HPC-Ops tests, the original Fused MoE regression file, SGLang adapter/multi-GPU tests, formatting, and changed-file static checks passed.
• The serving adapter used for this A/B is available at taoyuanyuan/sglang:feature/hpc-ops-peer-indexed-input-v1, commit bd8a7ca5.
• run_peer_indexed_sglang_server.sh launches the fixed GLM-5.2 DP8 + EP8 configuration for deepep, megamoe, or hpc.
• benchmark_peer_indexed_sglang.sh runs the fixed 8K-input/1K-output B8/B32 workload with exact-shape warmup, cache flush, three measured repetitions, and deterministic output validation.

python3 setup.py build_ext --inplace
export MODEL_PATH=/path/to/GLM-5.2-FP8
export SGLANG_ROOT=/path/to/sglang
MODE=hpc benchmark/fused_moe/run_peer_indexed_sglang_server.sh
# In another shell:
MODE=hpc benchmark/fused_moe/benchmark_peer_indexed_sglang.sh