perf: Group source-projection PSUM banks to lift wide-intermediate MFU

[RuifengHua]

Summary

MLP CTE gate/up (source) projection sizes the batch×sequence subtile count as NUM_HW_PSUM_BANKS // src_proj_int_dim_tile_count (mlp_cte_tile_info.py), and the matmul assigns one PSUM bank per (bxs_subtile, int_tile) pair held live across the hidden contraction loop (mlp_cte_projection.py):

psum_bank = bxs_subtile.index * len(int_tiles) + int_tile.index

So subtiles * int_tiles <= 8. Once the intermediate shard reaches int_tiles >= 5 (I_shard > 2048), the subtile count floors to 1 and stays there up to the compile limit (int_tiles = 8, I_shard = 4096).

A single subtile means each weight load feeds only 128 tokens instead of 256+, halving matmul arithmetic intensity and leaving the tensor engine underutilized on wide-intermediate shards — the regime where large dense models (Llama3-70B, Qwen3-32B, Gemma3-27B) run at TP4–TP8.

Fixes #9

Scope

This changes only the standard (non-quantized) source-projection path. The quantized (STATIC/ROW) and MX paths are untouched — they keep the original bank-fitting formula, and group=None reproduces their exact prior behavior. (Extending to quant is in Future work.)

This PR

Processes the intermediate dimension in PSUM-bank-sized groups instead, draining each group to SBUF before reusing its banks for the next. This decouples the subtile count from the total int-tile count while keeping live banks within the 8-bank budget.

• Adds SrcProjPsumGroup + build_src_proj_psum_groups (groups of NUM_HW_PSUM_BANKS // bxs_subtile_count tiles).
• Standard calc_batch_seqlen_dim_tile_size now allows up to NUM_HW_PSUM_BANKS subtiles (bounded by the existing 512 / token-count caps); quantized and MX keep the original bank-fitting formula.
• project_standard_source_tensor_tile takes an optional group; group-scoped weight DMA keeps total weight traffic at one full-width load. group=None reproduces the original flat bank layout exactly.
• Gate/up standard paths run the group loop (_perform_standard_gate_projection / _perform_standard_up_projection); the quantized/MX paths are unchanged.
• Drain sites (apply_source_projection_bias, perform_elementwise_multiply, apply_source_projection_activation) take an optional group with group-local bank indexing; group=None is byte-identical to the original.

Results

Measured on TRN2 (neuronx-cc 2.26.6360.0), before vs after, on the final (this) commit.

• ctrl = already-healthy control (should not regress).
• real_* configs are exact per-rank shapes for dense models at real tensor-parallel degrees (see note below).

config	I_shard	base MFU	opt MFU	base rt	opt rt	speedup
lnc2_H3072_I8192	4096	38.0%	75.6%	3.88ms	1.95ms	1.99×
lnc2_H8192_I4096 (ctrl)	2048	72.9%	84.7%	2.70ms	2.32ms	1.16×
lnc2_H8192_I6144	3072	48.0%	88.2%	6.15ms	3.34ms	1.84×
lnc2_H8192_I8192	4096	39.6%	61.8%	9.93ms	6.36ms	1.56×
real_gemma27b_TP4	3072	39.8%	73.8%	3.39ms	1.83ms	1.85×
real_gemma27b_TP8	3072	42.4%	86.2%	1.59ms	0.78ms	2.04×
real_llama70b_TP4	3584	40.8%	88.5%	5.63ms	2.59ms	2.17×
real_llama70b_TP8 (ctrl)	2048	72.4%	85.4%	1.81ms	1.54ms	1.18×
real_qwen32b_TP4	3584	39.0%	73.1%	3.68ms	1.96ms	1.88×
real_qwen32b_TP8	3584	39.6%	85.4%	1.81ms	0.83ms	2.18×
sc_H4096_I2048 (ctrl)	2048	78.7%	90.5%	1.29ms	1.12ms	1.15×
sc_H4096_I4096	4096	43.9%	81.1%	4.57ms	2.47ms	1.85×
sc_gpt2small_I3072	3072	39.6%	87.1%	0.75ms	0.34ms	2.20×
tok_S3072_H8192_I8192	4096	39.1%	63.1%	20.09ms	12.46ms	1.61×
tok_S512_H8192_I8192	4096	40.5%	62.8%	3.24ms	2.09ms	1.55×

Every config improves with no regressions (1.15–2.2×); the largest gains are on the previously starved wide-intermediate shards. Matmul arithmetic intensity rises from ~125 to ~470 on the wide configs. The three (ctrl) configs (already at 2 subtiles, I_shard 2048) still improve modestly and do not regress.

The real_* configs are the per-rank shapes get_mlp_config() produces in test/integration/nkilib/core/mlp/test_mlp_cte_model_config.py for Llama3-70B, Qwen3-32B, and Gemma3-27B at TP4–TP8 — the degrees where per-rank intermediate falls in the starved regime (higher TP shrinks I_shard below it).

The five starved ones gain 1.85–2.18×; real_llama70b_TP8 (I_shard 2048) is an already-healthy control and correctly shows only 1.18×.

Test plan

Verified with the shipped test_mlp_cte.py::test_mlp_cte_unit suite on TRN2 (real compile + infer):

• 111/111 unit cases pass (standard, gate/up/down bias, skip-gate, RMS/LayerNorm, GELU, ROW and STATIC quant)
• Wide-intermediate standard configs engage 2 subtiles (previously 1)
• Quantized/MX paths unchanged (group=None reproduces original bank layout)
• Before/after profiled on the config sweep above

Future work(Quant Path)

The quantized dual-row path (project_quantized_source_tensor_tile) has the same bxs_subtile * len(int_tiles) + int_tile bank assignment, so the grouping extends to it. I prototyped this and confirmed correctness on TRN2 (wide-I ROW/STATIC cases within tolerance).

The gain is smaller than the standard path: dual-row already feeds ~2 rows per weight load, so even the single-subtile wide case starts at a healthy arithmetic intensity. Grouping lifts the widest shape ~44%→52% MFU, and is flat where it was already fine. In my PoC 2 subtiles beat 3 — pushing to 3 regressed back toward ~44%. I didn't fully root-cause it, but the likely reason is that more subtiles force finer PSUM-bank grouping (group_size = 8 // subtiles, so 2→4 groups here), adding drain/reload traffic that outweighs the higher per-load intensity.

A natural follow-up is to pick the subtile count per config rather than fix it, balanced against the group-count/drain tradeoff. Two ways to get there:

• Simpler: size it from the allocator's live SBUF budget (get_free_space()) minus an estimated reserve for the buffers allocated later in the pipeline. Cheap, but the reserve is an estimate, so it has to stay conservative.

• Correct but more involved: defer the count decision until after those buffers are allocated, so get_free_space() reports the true remaining budget (and probe one subtile's cost instead of modeling it). No estimate, self-correcting across shapes/chips, but it's a structural change to the tiling flow.

I wouldn't expect a large jump from either, but both should recover some perf on the wide-quant tail.

This is about the dual-row (STATIC/ROW) path only. MX is a separate path (project_mx_source_tensor_tile) with its own PSUM bank layout and isn't grouped here; it's also TRN3-only, so I couldn't validate it on the TRN2 setup I profiled on — a natural extension once there's TRN3 access.

Happy to discuss direction if useful.

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.