[yuweih205]

Routed BF16 forward paths currently materialize a.index_select(0, rows) before grouped GEMM. This adds m_grouped_bf16_gemm_nt_contiguous_gathered(a, b, d, grouped_layout, a_row_indices) on SM90, loading A through row indices directly into the existing WGMMA pipeline.

The first version supports K-major BF16, K divisible by 64, contiguous weights, 128-row expert alignment and an int64 CUDA row map. A may have a padded row stride. The A loader uses 16-byte cp.async copies with the existing shared-memory swizzle; B/output retain TMA. It does not add a fused SwiGLU epilogue or change default dispatch.

Current revision

Commit ca1291f7cd2dbe483af7d89f033b2a5a1e96ca1f, based on main 1fc9c66975ee322d628b3288534ca0e9ce12a64d, keeps the change in six existing files:

• Public API and registration in csrc/apis/gemm.hpp; configuration, TMA descriptors and launch in the existing SM90 implementation header.
• Integrated coverage in tests/test_bf16.py, rather than a separate gathered test file. It covers three generated configurations, six boundary cases and three rejected invalid inputs.
• The CUDA kernel's executable tokens are unchanged from the previous published revision. Host checks are reorganized while retaining the dimensionality, dtype, device, row-stride, alignment and empty-input constraints.

The native extension was forcibly rebuilt from this exact source and the focused gathered test passed on the original H200/SM90 stack: PyTorch 2.13.0+cu132, CUDA runtime 13.2, seed 42, strict deterministic algorithms and uninitialized-memory filling disabled. All generated results are checked bitwise against the materialized path. Both paths use CUDA events in the integrated test, but those printed timings are not a new controlled performance result.

This is build/numerical/boundary validation, not a rerun of the historical full integration, sanitizer or performance studies below.

Earlier 2.8.1 measurements

The following results were collected against the 2.8.1 release, #462 (057ca5964aae0879ff2e0eb71ee05a3cb0ba3df7), at PR commit 0ff99780929827f24f19d5e67a10c390c5f77592.

On eight H200s, against activation gather + the existing GEMM with prepacked gate/up weights + the same fused output-compaction/SwiGLU kernel:

D=3072, F=1280, 16 experts	Eager reduction, run 1 / repeat	CUDA Graph reduction
R=8192, M=16384, random routes	6.42% / 6.31%	—
R=108031, M=170003, random routes	6.07% / 5.83%	—
Source-ordered routes, nominal shape	6.18%	5.78%
Source-ordered routes, packet stride D+64	6.39%	6.16%

Eager numbers are median per-GPU ABBA/BAAB reductions; Graph is one GPU. The general eager sweep uses random routes with replacement, while Graph uses source-ordered unique routes; only the explicitly source-ordered rows compare matching routes across modes. Workers and process repeats reuse seed 42. An additional same-GPU, matching-row-map test at R=8192/M=16384 measures 6.25% eager / 6.01% Graph. The small R=1024, M=512, D=F=256 case measures 26.07% eager but regresses 2.31% under Graph, on matching routes. This is a fixed-route forward benchmark, not a full-model speedup.

All tested h/gate/up outputs match both existing-library baselines bitwise. API tests cover partial/empty M, empty experts, source/output stride, aligned offsets and misaligned pointer rejection. An additional 79-case regression covers existing BF16 GEMM/einsum layouts, output dtypes and accumulation modes using exact references and bytewise repeats. Independent FP32 reference checks pass; compute-sanitizer memcheck reports 0 errors. The full 13-phase 2.8.1 run passed, including live EP8 integration and receive-view lifetime checks; combined forward reductions are 16.99%–17.59% across two runs when both this change and DeepEP #769 are enabled.

2.8.1 reproducer, raw cycles, source hashes and limitations. These measurements retain their original test revision and predate the current reorganization; this update makes no additional performance claim. The earlier 2.8.0 measurements remain separately preserved in the same Gist.

Before this reorganization, the PR test scripts also passed a separate eight-H200 correctness rerun against the frozen production-equivalent runtimes. That DeepEP test revision checked an independent input/route oracle, raw deterministic order and two delayed consumer streams under both record-stream policies. Exact test revisions and retained test-development failures.