Chenghao/triton moe dense mlp opt#249
Open
nvchenghaoz wants to merge 25 commits into
Open
Conversation
Add benchmark script and running doc for triton_moe_dense_mlp optimization.
Baseline on H100 (bf16):
- Kernel 1 (activation): 7.5-402 us across shapes
- Kernel 2 (weighted sum): 31-172 us across shapes
- Triton kernels are 5-10% of E2E (BMM dominates)
Shapes: GPT-OSS-120B (E=128) and GPT-OSS-20B (E=32)
with T in {1, 8, 32, 128, 512}.
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…sweep
Sweep 25 combos of num_warps={1,2,4,8,16} x num_stages={1,2,3,4,5}
for both Triton kernels across all 8 benchmark shapes.
Key findings:
- Kernel 2 (weighted_sum): w=16 is universally best, 15-45% faster
- Kernel 1 (activation): w=4 already near-optimal for large T,
w=16 helps at small T (~14%)
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Apply best num_warps/num_stages from sweep to kernel launcher: - Kernel 1: adaptive warps (16 if total_rows<=128, else 4), stages=2 - Kernel 2: warps=16, stages=2 (was 4, 3) Results (Triton kernel total, us): A1: 66.5 (-42.8%), A2: 75.9 (-40.4%), A3: 102.1 (-34.0%) A4: 186.2 (-21.9%), A5: 547.1 (-4.8%) B1: 25.2 (-35.1%), B2: 34.7 (-27.4%), B3: 154.7 (-5.2%) Correctness: 31/31 tests PASS. Also fix bench/sweep scripts with __main__ guard for auto-import safety. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
… kernel 2 Kernel 2 (weighted_sum) now uses a 2D grid (T, cdiv(H,1024)) when T<=128 for better occupancy at low token counts. Falls back to 1D grid for T>128. Results vs baseline (Triton total, us): A1: 60.6 (-47.9%), B1: 14.2 (-63.4%), B2: 23.1 (-51.7%) A5: 547.2 (-4.7%) -- no regression on prefill Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…lesced loads Deinterleave the gate_up tensor before the GLU activation kernel so that gate and up halves are contiguous in memory, enabling coalesced 128-byte loads in the Triton kernel instead of stride-2 scattered loads. Also tighten the K2 BLOCK_H heuristic: use 256 (instead of 1024) for the small-model / low-token regime (E<=32, T<=32) to expose more parallelism. All 31 correctness tests pass. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…K2 adaptive BLOCK_H K1: Rewrite to read contiguous [gate|up] layout (5-7% faster). Launcher deinterleaves gate_up via slicing before kernel launch. K2: Adaptive BLOCK_H (256 for small E+T, 1024 medium, pow2(H) large T). Also tested K1 multi-row (persistent) kernel -- no benefit, dead end. Results vs baseline (Triton total, us): A1: 61.0 (-47.5%), A3: 92.3 (-40.3%), A5: 528.9 (-7.9%) B1: 13.6 (-64.9%), B2: 22.7 (-52.5%), B3: 149.2 (-8.5%) Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Replace _dtype_probe global load with OUTPUT_DTYPE constexpr parameter. Eliminates one wasted global memory access per program invocation. No measurable perf change (probe was already evict_first), but cleaner code. Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…eight rearrange
Iter 6: Replace dtype_probe load with OUTPUT_DTYPE constexpr. Cleanup only.
Iter 7: Tried per-forward weight rearrange -- 4x E2E regression. Reverted.
Weight rearrange must happen at model load time, not per-forward.
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…ights Add zero-weight skip in K2 expert loop. For GPT-OSS top-4/128 routing, this avoids loading 124/128 expert outputs per token. Sparse routing results (top-4 of 128): A1: 59->13 us (-78%), A4: 74->19 us (-74%) B1: 19->8 us (-59%), B2: 20->8 us (-59%) Dense routing (all non-zero): ~10% overhead from branch. Acceptable since real-world GPT-OSS uses sparse top-k routing. Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…e BLOCK_I K1 now uses 2D grid (total_rows, cdiv(I, BLOCK_I)) with BLOCK_I=1024 for small total_rows (<=128), falling back to 1D with BLOCK_I=pow2(I) for large total_rows. Results: ~1% K1 improvement at large T (A5: 383.7->379.6 us). Small T unchanged (already fast from higher warps). Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Add iteration 9 results and comprehensive summary table. Best results vs baseline: - Dense routing: 8-44% Triton kernel speedup across shapes - Sparse top-4 routing (real-world): 50-84% Triton kernel speedup Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Re-sweep revealed num_warps=16 is now optimal for large-T K1 configs (was 4 before coalesced loads). Unify to w=16 for all K1 configs. Results vs iter 9: A3: 31.8->30.5 (-4.1%), A4: 101.4->99.2 (-2.2%) A5: 379.6->375.4 (-1.1%), B3: 101.2->99.2 (-2.0%) Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…se stride-2 kernel Revert to stride-2 kernel reads and remove the launcher deinterleave (slicing+cat+contiguous). The deinterleave was copying ~377MB for A5 which cost more than the ~15% kernel-level speedup from coalesced loads. Net E2E improvement: A4: 3199->2826 (-12%), A5: 7982->6735 (-16%), B3: 2040->1704 (-16%) K1 kernel isolated: ~15% slower (stride-2 vs coalesced) But E2E: ~12-16% faster (no deinterleave copy overhead) Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…iguous() BMM output is already contiguous (verified). Remove the no-op .contiguous() call. Also tested K2 routing preload -- dead end, sparse vs dense E2E is within noise since BMM dominates. Correctness: 31/31 tests PASS (verified in iter 11). Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Remove explicit fp32 upcast in K1 activation kernel. Triton handles dtype promotion internally for sigmoid. Cleaner code, no perf regression. Also tested eviction_policy and BLOCK_I=2048 -- both dead ends. Correctness: 31/31 tests PASS (including fp32 tolerance tests). Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…oc update Bandwidth analysis shows K1 at 74-79% of H100 peak memory BW for large shapes -- approaching hardware limits. K2 with sparse routing is latency-bound (launch overhead), not bandwidth-bound. Updated running doc through iter 14 with comprehensive results. Dead ends: eviction_policy, BLOCK_I=2048, K2 routing preload, K2 num_warps for sparse (w=8 marginal, not worth complexity). Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…parse routing Add _weighted_expert_gather_kernel that uses top-k indices to gather only active experts instead of looping over all E experts. For GPT-OSS top-4/128 sparse routing (measured in isolation): A1: 12.8->5.6 us (-56%), A4: 19.0->7.7 us (-59%) B1: 7.9->5.7 us (-28%) Launcher auto-detects sparse routing and extracts top-k indices. Falls back to loop-based kernel for dense routing. Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
… + BLOCK_I tuning - Replace expensive sparse detection (boolean tensor + reduce) with cheap single-token check: (routing_weights[0] == 0).any() - Add BLOCK_I=2048 tier for medium total_rows (129-1024) for ~4% K1 gain on A2-like shapes Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…ep findings - Add T>=64 threshold for gather kernel: topk extraction overhead outweighs kernel savings for small T. Small T uses zero-skip loop. - K2 gather BLOCK_H/num_stages sweep: BH=512-1024 optimal, stages flat. - K2 gather TOP_K sensitivity: near-free for k<=4, scales linearly after. - K1 stride check: non-contiguous input surprisingly ~6% faster (not exploitable). Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…ge T K1 is bandwidth-saturated (74-79% peak BW). No compute optimization helps (tested: fast_sigmoid, tl.where, paired loads, explicit casts). num_warps=8 gives ~1.5% K1 improvement for large T (>1024 rows). Also tested: torch.compile makes it 3x slower (.item() graph break). Correctness: 31/31 tests PASS. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…+ cleanup Profiled launcher overhead breakdown: BMM1: 60-64%, BMM2: 31-33%, K1: 0.3-6.4%, K2: 2.1-2.8%, overhead: <0.2% Triton kernels are only 3-9% of E2E. BMMs dominate at 94%. Removed redundant .contiguous() before K2 (BMM output is already contiguous). Tested pre-allocation of buffers -- no benefit (3us vs 2-7ms E2E). Dead ends: launcher overhead optimization, buffer pre-allocation. Correctness: verified via prior iterations. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…r 22 Add detailed iteration logs for iter 15-22: - Iter 15: K2 gather kernel (50-59% faster for E=128 sparse) - Iter 16: Sparse detection optimization - Iter 17: Gather threshold T>=64 - Iter 18: K1 BW-saturation experiments (dead ends: fast_sigmoid, etc) - Iter 19-22: Launcher profiling (BMMs=94%, overhead<0.2%) Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Complete optimization loop summary: - K2 dense routing: 108->59 us (-45%) - K2 sparse top-4 routing: 108->5.6 us (-95%) - K1: at 74-79% H100 peak memory BW, 1-4% tuning gains - 13 dead ends tested and documented - Hardware limits analysis confirms K1 is BW-saturated Key wins: K2 num_warps=16, 2D grid, zero-skip, gather kernel Key insight: K1 is memory-BW-bound, no compute optimization helps Future: weight rearrange at load time, fuse BMM+activation Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Add detailed explanation of iter 11 tradeoff: K1 kernel +52 us (stride-2 vs coalesced) saves 1250 us deinterleave copy = 1198 us net E2E gain for A5. Add E2E comparison table showing iter 10 (with deinterleave, worse E2E) vs iter 11+ (without deinterleave, better E2E) to justify the decision. Future fix: rearrange gate_up_w at model load time for both benefits. Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.