Skip to content

[cute] Omit redundant persistent shared loop - #3294

Draft
yushangdi wants to merge 1 commit into
mainfrom
yushangdi/stack/88
Draft

[cute] Omit redundant persistent shared loop#3294
yushangdi wants to merge 1 commit into
mainfrom
yushangdi/stack/88

Conversation

@yushangdi

@yushangdi yushangdi commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Remove the redundant CTA-wide persistent scheduler loop from eligible, fully role-local tcgen05 kernels.

The TMA-load, MMA-execute, and epilogue roles already own independent persistent schedulers. After role partitioning, codegen now inspects the actual residual shared body and omits its loop only when it contains no meaningful work: dependency-only or dead coordinate setup, placeholders, and obsolete CTA barriers.

The decision is fail-closed. A single-root kernel must have all three role-local bodies, and any unrecognized or observable shared statement keeps the loop. Definitions consumed by post-loop cleanup also keep it. Grouped-static kernels retain their specialized replacement-aware classification.

This removes the redundant loop for the pretuned FP8 scale_mm_cute M=64, K=5120, N=5120 configuration while preserving it for the rank-3 RHS group-selection path, whose residual shared body still performs scalar loads and comparisons.

Add focused codegen coverage for both cases, residual-work and post-loop dependency tests, and runtime numerical coverage for the affected persistent role-local kernels.

yushangdi added a commit that referenced this pull request Aug 6, 2026
Remove the redundant CTA-wide persistent shared loop from validated, fully role-local `tcgen05` kernels with `cluster_m=1`.

In these kernels, the TMA-load, MMA-execute, and epilogue roles already own independent persistent scheduler loops. The residual shared loop only repeated tile-coordinate setup, placeholder assignments, scheduler advancement, and CTA barriers; it performed no useful computation or stores.

The existing omission condition now also recognizes `use_validated_cluster_m1_role_local_body`. Omission remains limited to a single-root kernel with all three role-local bodies. Codegen fails closed if the residual shared body contains observable work or defines a value consumed by post-loop cleanup.

stack-info: PR: #3294, branch: yushangdi/stack/88
@yushangdi
yushangdi force-pushed the yushangdi/stack/88 branch from 737301a to f065531 Compare August 6, 2026 01:17
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 6, 2026
@yushangdi

Copy link
Copy Markdown
Contributor Author

Largest affected gains:

  • (64,4096,6144): 6.8%
  • (64,12288,4096): 5.2%
  • (64,5120,5120): 5.0%
  • (64,2048,12288): 5.0%

Before:

GPU: NVIDIA B200
     M       K       N   helion (us)     torch (us)   speedup
/home/shangdiy/worktrees/helion-mm/helion/helion/_compiler/compile_environment.py:286: UserWarning: The 'cute' backend is experimental and may have limited functionality.
  warn_once(
     1    4096    4096          7.21          11.91     1.65x  (vs torch)
  4096    4096    4096         62.43          69.90     1.12x  (vs torch)
     1    4096     256          4.92           7.98     1.62x  (vs torch)
   512    2048    4096          9.40           8.49     0.90x  (vs torch)
   512    2048    2048          6.68           7.11     1.06x  (vs torch)
    64    2048    4096          6.73           8.94     1.33x  (vs torch)
    64    2048    2048          6.01           6.45     1.07x  (vs torch)
    64    2048   12288         10.49          11.10     1.06x  (vs torch)
    64    6144    2048          9.02          11.62     1.29x  (vs torch)
    64    4096    6144         10.52          12.65     1.20x  (vs torch)
    64    4096    4096          9.00          12.17     1.35x  (vs torch)
    64    4096   24576         25.95          31.84     1.23x  (vs torch)
    64   12288    4096         16.28          23.84     1.46x  (vs torch)
    64    5120   10240         17.92          17.01     0.95x  (vs torch)
    64    5120    5120         11.35          13.42     1.18x  (vs torch)
    64    5120   51200         57.00          58.81     1.03x  (vs torch)
    64   25600    5120         37.80          49.09     1.30x  (vs torch)
vs torch: wins=15/17 geomean=1.207x best=1.65x

Helion faster on 15/17 shapes vs the best baseline; geomean speedup 1.207x; best speedup 1.65x.

After:

GPU: NVIDIA B200
       M       K       N   helion (us)     torch (us)   speedup
  /home/shangdiy/worktrees/helion-mm/helion/helion/_compiler/compile_environment.py:286: UserWarning: The
  'cute' backend is experimental and may have limited functionality.
    warn_once(
       1    4096    4096          7.25          11.88     1.64x  (vs torch)
    4096    4096    4096         57.89          66.64     1.15x  (vs torch)
       1    4096     256          4.75           7.99     1.68x  (vs torch)
     512    2048    4096          9.13           8.47     0.93x  (vs torch)
     512    2048    2048          6.75           6.95     1.03x  (vs torch)
      64    2048    4096          6.50           9.26     1.42x  (vs torch)
      64    2048    2048          5.95           6.48     1.09x  (vs torch)
      64    2048   12288          9.97          10.57     1.06x  (vs torch)
      64    6144    2048          8.76          11.66     1.33x  (vs torch)
      64    4096    6144          9.80          12.23     1.25x  (vs torch)
      64    4096    4096          8.59          11.98     1.39x  (vs torch)
      64    4096   24576         25.59          31.14     1.22x  (vs torch)
      64   12288    4096         15.44          23.99     1.55x  (vs torch)
      64    5120   10240         17.29          17.31     1.00x  (vs torch)
      64    5120    5120         10.78          13.75     1.28x  (vs torch)
      64    5120   51200         56.19          59.04     1.05x  (vs torch)
      64   25600    5120         36.93          48.71     1.32x  (vs torch)
  vs torch: wins=16/17 geomean=1.240x best=1.68x

yushangdi added a commit that referenced this pull request Aug 6, 2026
Remove the redundant CTA-wide persistent scheduler loop from eligible, fully role-local tcgen05 kernels.

The TMA-load, MMA-execute, and epilogue roles already own independent persistent schedulers. After role partitioning, codegen now inspects the actual residual shared body and omits its loop only when it contains no meaningful work: dependency-only or dead coordinate setup, placeholders, and obsolete CTA barriers.

The decision is fail-closed. A single-root kernel must have all three role-local bodies, and any unrecognized or observable shared statement keeps the loop. Definitions consumed by post-loop cleanup also keep it. Grouped-static kernels retain their specialized replacement-aware classification.

This removes the redundant loop for the pretuned FP8 scale_mm_cute M=64, K=5120, N=5120 configuration while preserving it for the rank-3 RHS group-selection path, whose residual shared body still performs scalar loads and comparisons.

Add focused codegen coverage for both cases, residual-work and post-loop dependency tests, and runtime numerical coverage for the affected persistent role-local kernels.

stack-info: PR: #3294, branch: yushangdi/stack/88
@yushangdi
yushangdi force-pushed the yushangdi/stack/88 branch from f065531 to d94fbae Compare August 6, 2026 02:08
Remove the redundant CTA-wide persistent scheduler loop from eligible, fully role-local tcgen05 kernels.

The TMA-load, MMA-execute, and epilogue roles already own independent persistent schedulers. After role partitioning, codegen now inspects the actual residual shared body and omits its loop only when it contains no meaningful work: dependency-only or dead coordinate setup, placeholders, and obsolete CTA barriers.

The decision is fail-closed. A single-root kernel must have all three role-local bodies, and any unrecognized or observable shared statement keeps the loop. Definitions consumed by post-loop cleanup also keep it. Grouped-static kernels retain their specialized replacement-aware classification.

This removes the redundant loop for the pretuned FP8 scale_mm_cute M=64, K=5120, N=5120 configuration while preserving it for the rank-3 RHS group-selection path, whose residual shared body still performs scalar loads and comparisons.

Add focused codegen coverage for both cases, residual-work and post-loop dependency tests, and runtime numerical coverage for the affected persistent role-local kernels.

stack-info: PR: #3294, branch: yushangdi/stack/88
@yushangdi
yushangdi force-pushed the yushangdi/stack/88 branch from d94fbae to 8097553 Compare August 6, 2026 02:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant