fix(layout): support strided PSUM scale packing

[morgendave]

Summary

DeepEP expanded FP8 dispatch returns activation scales in TMA-aligned column-major storage. The PSUM scale packer required a contiguous tensor, blocking valid grouped FP8 GEMMs.

This change adds a dedicated stride-aware PSUM pack kernel with two JIT-specialized load modes: four-row vectorized column-major loads for DeepEP and a generic positive-stride fallback. The existing contiguous kernel keeps its original ABI and byte-identical load body, so existing callers do not pay for stride support or compile into a different kernel.

flowchart LR
    D[DeepEP expanded scales] --> V{Input layout}
    V -->|contiguous row-major| R[Original pack kernel]
    V -->|aligned column-major| C[Vectorized strided kernel]
    V -->|other positive strides| G[Generic strided kernel]
    R --> M[PSUM gap mask]
    C --> M
    G --> M
    M --> P[TMA-aligned packed UE8M0]

Loading

Correctness

The B300/SM100 regression matrix covers:

• contiguous row-major, padded TMA column-major, non-vectorizable column-major, and generic positive-stride views;
• alignments 128 and 256;
• empty experts, one/two/many experts, partial BLOCK_MN tails, and scale-K widths 2/3/4/7/48;
• NaN-filled PSUM gaps checked against an independent sanitized PyTorch reference;
• the exact FireTitan PR #42640 contract and a 32-expert GLM production-width scale shape.

Results on current head:

• six-case layout matrix: PASS (1 passed in 5.61s);
• unchanged FireTitan #42640 grouped-GEMM regression: PASS (1 passed in 4.96s);
• independent final branch review: PASS.

Integrated full-MoE simulation

A single-B300 simulated EP8 / 262K-global-token full expert layer also passed on current head. This is not a real NCCL EP8 run; it reproduces the per-rank routed workload and complete expert math:

• global tokens: 262,144; top-k: 8; simulated EP: 8;
• local routed rows: 262,144 across 32 local experts; 263,936 rows after 128-alignment;
• activation scales: (263936, 48), TMA column-major stride (1, 263936);
• full W13: (32, 14336, 6144);
• SwiGLU + masked FP8 requantization;
• full W2: (32, 6144, 7168);
• sampled valid gate/up/output rows finite;
• measured full forward: 24.373 ms; peak allocated memory: 17.48 GiB.

Real-model layer A/B

Measured with real GLM-5.2 checkpoint weights and router from /shared/yuedong/glm-5p2-bf16-8l, layer 3. The benchmark selects 64 local experts from the real 256-expert layer, computes top-8 routes with the real router, quantizes the real W13/W2 weights to the production FP8-block128 format, and runs the integrated expanded W13 → SwiGLU/requantize → W2 path. Old and new DeepGEMM were source-built in the same Torch 2.11 / CUDA 13 image on the same B300; three/five warmups precede 20/50 measured iterations.

Tokens	Routed / padded rows	18190bc + required contiguous copy	PR #11 direct strided scales	Result
8,192	12,257 / 17,280	p50 0.7894 ms, p90 0.7984 ms	p50 0.7990 ms, p90 0.8077 ms	effectively flat; candidate ~1.2% slower
32,768	48,472 / 52,992	p50 1.8767 ms, p90 1.9936 ms	p50 1.7855 ms, p90 1.9829 ms	p50 4.9% lower, p90 0.5% lower

Peak allocated memory was identical within each A/B (2.785 GiB at 8K; 3.748 GiB at 32K). The distributions are mildly bimodal on the shared host, so this supports “no material full-layer regression” rather than a production throughput claim. The primary reason to merge remains correctness: 18190bc rejects the natural expanded DeepEP scale view without a copy.

Performance

Matched A/B methodology: the old 18190bc and this PR were both built from source in the same Torch 2.11 / CUDA 13 trainer image, then measured on the same B300 with CUDA events, 50 warmups, and 500 iterations. The old path must first materialize strided_scales.contiguous(); the candidate consumes the DeepEP view directly.

Case	Shape (M,N,K)	Old copy + GEMM	New direct strided GEMM	Result
FireTitan #42640	(256,128,256)	65.568 us	57.152 us	12.8% lower / 1.147x
GLM representative, 32 experts	(65536,128,6144)	135.296 us	132.288 us	2.2% lower / 1.023x

No extra kernel launch or synchronization is added. The new layout kind becomes part of the generated JIT code, producing at most the two new strided pack variants; runtime stride values do not multiply cache entries.

Testing

• python -m py_compile tests/test_layout.py
• git diff --check origin/main...HEAD
• B300 six-case layout matrix
• B300 FireTitan PR #42640 grouped-GEMM integration test
• matched old/new source-build grouped-GEMM performance A/B
• independent CUDA/JIT branch review

Risk

The new ABI is private to the new strided JIT kernel. DeepEP supplies aligned base storage and positive element strides. Existing contiguous callers retain the original kernel ABI/body and selection path.

---

Note

Medium Risk
Adds a new SM100 CUDA pack kernel and changes grouped-GEMM scale layout handling. Risk is contained by an explicit opt-in flag and an unchanged contiguous kernel path.

Overview
Enables PSUM scale packing from non-contiguous (DeepEP TMA column-major) views so grouped FP8 GEMMs no longer require a contiguous copy of activation scales.

Callers must set allow_strided_psum_scales=True on m_grouped_fp8_fp4_gemm_nt/nn_contiguous and the UE8M0 pack APIs. Without the flag, strided PSUM scales still assert.

A new transpose_and_pack_strided_fp32_into_ue8m0 kernel handles two JIT layouts: vectorized column-major loads and a generic positive-stride fallback. The original contiguous pack kernel keeps its ABI and load path.

Tests cover FireTitan/GLM-style shapes, empty experts, unaligned/generic strides, NaN gap rows, and the opt-in assertion.

^{Reviewed by Cursor Bugbot for commit 1a9338a. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 0f336b1. Configure here.}

[cursor]

Stale comment

Risk: medium. Left a non-blocking comment and assigned reviewers; not approved. The strided PSUM pack kernel and grouped-GEMM API change exceed the low-risk auto-approve bar, and Bugbot skipped with an unresolved finding that the new opt-in GEMM test is never invoked from the script runner.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Risk: medium. Left a non-blocking comment; reviewers were already assigned, so none added. The strided PSUM pack kernel and grouped-GEMM API change still exceed the low-risk auto-approve bar; humans should confirm strided load correctness and the opt-in path.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}