Skip to content

Cap persistent kernel grids at logical PID count - #3160

Draft
yushangdi wants to merge 1 commit into
mainfrom
fix/persistent-grid-total-pid-cap
Draft

Cap persistent kernel grids at logical PID count#3160
yushangdi wants to merge 1 commit into
mainfrom
fix/persistent-grid-total-pid-cap

Conversation

@yushangdi

Copy link
Copy Markdown
Contributor

Summary

  • cap host-computable persistent launch grids at min(total_pids, num_sms * num_sm_multiplier)
  • preserve uncapped grids for data-dependent bounds and the existing cluster-aware tcgen05 scheduler
  • align the public Config type with the autotuner's supported multipliers through 128
  • add coverage for capped static grids, uncapped dynamic grids, and a 16x persistent multiplier

Motivation

For a 2048-row fused RMS norm on GB300, the autotuner can search num_sm_multiplier=16, but the old launcher emitted _NUM_SM * 16 CTAs. Capping the physical grid emits min(2048, _NUM_SM * 16), making the intended 2048-CTA family available without changing the device-side persistent stride.

With this change, a fresh full LFBO run stayed in the persistent family and selected a 16x configuration. Its canonical benchmark was 15.55 us, compared with 15.68 us for the previous Helion configuration in the same server process.

Testing

  • test/test_persistent_kernels.py: 32 passed, 1 skipped on GB300
  • focused post-rebase run: 6 passed
  • Pyrefly on changed implementation files: 0 errors

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 25, 2026
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