Skip to content

[inductor] Proven int32 size arguments for the dynamic-concat kernel in #189940 (opt-in, for review) - #198733

Closed
scasella wants to merge 3 commits into
pytorch:mainfrom
scasella:inductor-proven-int32-size-args
Closed

[inductor] Proven int32 size arguments for the dynamic-concat kernel in #189940 (opt-in, for review)#198733
scasella wants to merge 3 commits into
pytorch:mainfrom
scasella:inductor-proven-int32-size-args

Conversation

@scasella

@scasella scasella commented Sep 26, 2026 •
edited by pytorch-bot Bot
Loading

Copy link
Copy Markdown

What

When config.triton.narrow_proven_size_args is on (default off), a Triton pointwise kernel's ks* size
arguments are declared i32 instead of i64, but only if the kernel is in a proven class and Inductor's own
symbolic state establishes the premises. Every other kernel is unchanged.

Supported class today: exactly the kernel Inductor emits for the nested 6-segment concatenation along a dynamic last
dimension in #189940.

The flag is a staging mechanism for review, not the proposed end state. Following the discussion in #190034 (fix the
root cause; avoid a dedicated knob), the question we want to put to reviewers is whether this proven eligibility
condition can become an automatic compiler decision at the right internal representation.

Why

signature_to_meta widens every ks* to int64, because a product or sum of symbols can overflow even when each
symbol fits. For this kernel that makes every index operation 64-bit, including xindex % ks0 and xindex // ks0,
and roughly doubles the per-element SASS.

Related work

Evidence (two stacks; please read them separately)

  • Performance, measured on torch 2.14.0 / Triton 3.8.0, one L4. A local prototype of this rule produced the
    ks*: i32 kernel. In a paired, pre-registered confirmation over a 14-shape dynamic sequence, it reduced warm
    complete-sequence time by 26% against the emitted kernel (median 601 → 444 ms; paired ratio 1.35×), with
    non-regression shown at all 10 small shapes tested. Output was bitwise identical to eager at every shape.
  • This patch, tested on nightly 2.15.0.dev20260926 (cu130; git 6aa9e2f, the nightly build of main ef166fb),
    one L4.     It applies; the four tests below pass; the 14-shape sequence is bitwise equal to eager with all seven ks
    as i32. That build emits the same kernel source byte for byte. This branch is based on ef166fb, where the two
    modified files are byte-identical to the tested build.
  • This patch, timed flag off versus on, on the same nightly and one L4. It was pre-registered and run once.
    OFF/ON on the 14-shape sequence is 1.388× (95% CI 1.378–1.408); median 592 → 427 ms. Non-regression was shown
    at all 10 small shapes, with every CI lower bound ≥ 1.29. Output was bitwise equal to eager, and there were no
    recompiles during measurement. The executed ON kernel was verified to be the i32 build.
    • Compile cost: the eligibility check takes 14.7 ms per compile of the target kernel and about 95 µs per rejected
      unrelated kernel. It runs once per compile, never per call.
    • Scope: one L4, one program. H100 is untested.

Full report, including negative results and the selection history: https://github.com/scasella/veritile-narrow-cat/blob/stage15-patch-timing/bench/optimizations/inductor_narrow/REPORT.md

Premises

Implemented in torch/_inductor/codegen/triton_size_arg_narrowing.py.

  • C1 (structure): a pointwise kernel on a 1-D grid with no fixed config, reduction or cooperative mode; the
    proved argument list and signature; and a body whose canonical IR (indexing, masks, payload dataflow, dead
    statements) has the proved SHA-256. The recognizer rejects any construct it does not model.
  • C2 (symbolic facts, never size hints):
    • ks1..ks6 are size symbols with a var_to_range lower bound ≥ 1;
    • ks0's precomputed definition equals ks1 + … + ks6;
    • numel == s * ks0 for a size symbol s with lower bound ≥ 1;
    • the index dtype is int32, and the guard numel <= 2**31 - 1 is installed (or statically known).
  • When a premise fails later: if the guard fails, Dynamo recompiles, Inductor chooses int64 indexing, and that
    kernel is not in the class. Exercised at numel = 2**31: bitwise identical to eager.

Why this is sound

A Lean 4 proof (standard axioms only), NarrowCat.narrow_cat_equiv and companions: https://github.com/scasella/veritile-narrow-cat/blob/stage15-review/bench/optimizations/inductor_narrow/NarrowCat.lean. It models the
kernel's integer semantics (int32/int64 promotion, two's-complement wrap at every operation, truncating % and //,
i32 argument packing). Under the premises, for every lane of the grid, the i64 and i32 kernels agree on the
store mask, store address, stored value and every memory read. Companion lemmas: equal positive divisors (no new
undefined behaviour); every ks_i in [0, 2**31) given positive widths and rows (exact i32 packing); no
xindex wraparound on the grid. The payload is symbolic, so no floating-point claim is made.

Why not assume_32bit_indexing. That flag is a user's promise that every symbol fits. Our boundary tests show
lanes where every ks_i < 2**31 yet the two kernels disagree, because ks1 + ks2 + ks3 overflows. This rule relies
on relations Inductor already knows instead.

Trusted, not proved: SymPy and the sizevars / shape_env facts; the recognizer's parse of the generated source
(binding-tested against the IR the proof is stated over); Triton's lowering; masked-off lanes never dereferencing
their (possibly wrapped) addresses, as the i64 kernel already assumes for past-the-end addresses.

Tests

test/inductor/test_triton_size_arg_narrowing.py, four GPU tests; all passed on one L4 at the pinned revision
(correctness only). They are targeted tests, not a substitute for the usual CI.

  • test_narrows_proven_cat_kernel: flag on → ks* are i32, output bitwise equal to eager at two shapes.
  • test_default_is_unchanged: flag off → ks* stay i64.
  • test_other_kernel_is_unchanged: a 3-segment cat keeps i64.
  • test_recognizer_rejects_mutations: an extra size use, a changed index and a changed cast all fail recognition.

The GPU tests and the timing above ran on the first commit, 406a839. The third commit, 293565c, limits the two tests
that expect recognition to CUDA; see Limitations. The second commit, ad17abc, only applies lint fixes: formatting, OrderedSet in place of set, and type annotations/cast for pyrefly. The test file's AST is
unchanged. We re-checked ad17abc on the pinned nightly's CPU codegen path:

  • the target kernel is accepted and narrowed to i32;
  • four unrelated kernels are rejected and keep i64;
  • the canonical digest still matches the proved kernel;
  • all three recognizer mutations are still rejected.

Limitations

  • One kernel class, recognised by digest. The digest is an identity check, not an explanation of why a class of
    programs is eligible; each new class would currently need its own extraction and proof. The upstream form may
    need to be a structural or range-based rule rather than a list of generated kernels. We would rather get
    direction on that before building it.
  • Safe failure direction. Recognition keys on generated-source structure, including argument naming and
    temporary numbering. A benign codegen change makes the rule stop firing (the kernel keeps i64); it does not
    misfire, provided the recognizer is correct.
  • One device. Measured on L4 only; the issue's H100 is untested.
  • CUDA only. The proven class is pinned to the source Inductor emits on NVIDIA CUDA. On other backends (ROCm,
    XPU) the emitted code may differ, and then the rule does not fire and the kernel keeps i64. The two tests that
    expect recognition therefore run on CUDA only.

Review questions

  1. Is generated-source recognition acceptable as an initial implementation, or should the rule operate earlier, on
    typed indexing expressions?
  2. Does the symbolic-fact extraction establish exactly the theorem's premises, including their maintenance across
    later invocations?
  3. What evidence would justify enabling this automatically for eligible kernels instead of behind a flag?

🤖 Generated with Claude Code

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo

@scasella @claude
[inductor] Opt-in proven int32 size arguments for the dynamic-concat …
406a839 
…kernel (pytorch#189940)

Adds config.triton.narrow_proven_size_args (default off). When on, a Triton
pointwise kernel's ks* size arguments are declared i32 instead of i64 only if
the kernel is in a proven class (currently the 6-segment dynamic concat from
pytorch#189940, recognised by a canonical-IR digest) and Inductor's symbolic state
establishes the premises (ks0 = sum of widths, numel = rows * ks0, lower
bounds >= 1, the int32 numel guard installed). Every other kernel is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@pytorch-bot

pytorch-bot Bot commented Sep 26, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/198733

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 9 Unrelated Failures

As of commit 293565c with merge base ef166fb (image):

NEW FAILURE - The following job has failed:

NEW FAILURES SUPPRESSED BY AI (non-blocking) - The following jobs failed, but the AI CI Advisor judged the failures not to be evidence against this change:

UNSTABLE - The following jobs are marked as unstable, possibly due to flakiness on trunk:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@pytorch-bot

pytorch-bot Bot commented Sep 26, 2026 •
edited
Loading

Copy link
Copy Markdown

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

@pytorch-bot

pytorch-bot Bot commented Sep 26, 2026

Copy link
Copy Markdown

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

  • ciflow/inductor
  • ciflow/torchtitan

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

@linux-foundation-easycla

linux-foundation-easycla Bot commented Sep 26, 2026 •
edited
Loading

Copy link
Copy Markdown

CLA Signed
The committers listed above are authorized under a signed CLA.

  • ✅ Unknown (406a839)
  • ✅ login: scasella / name: Stephen Casella (406a839)

One or more co-authors of this pull request were not found. You must specify co-authors in commit message trailer via:

Co-authored-by: name <email>

Supported Co-authored-by: formats include:

  1. Anything <id+login@users.noreply.github.com> - it will locate your GitHub user by id part.
  2. Anything <login@users.noreply.github.com> - it will locate your GitHub user by login part.
  3. Anything <public-email> - it will locate your GitHub user by public-email part. Note that this email must be made public on Github.
  4. Anything <other-email> - it will locate your GitHub user by other-email part but only if that email was used before for any other CLA as a main commit author.
  5. login <any-valid-email> - it will locate your GitHub user by login part, note that login part must be at least 3 characters long.

Alternatively, if the co-author should not be included, remove the Co-authored-by: line from the commit message.

Please update your commit message(s) by doing git commit --amend and then git push [--force] and then request re-running CLA check via commenting on this pull request:

/easycla

@scasella scasella mentioned this pull request Sep 26, 2026
Open
scasella and others added 2 commits September 26, 2026 18:11
@scasella @claude
Lint fixes: formatting, OrderedSet, type annotations (no behaviour ch…
ad17abc 
…ange)

Applies lintrunner fixes (PYFMT, SET_LINTER) and resolves pyrefly errors in
triton_size_arg_narrowing.py with annotations and typing.cast only. The test
file changes are formatting only (identical AST).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@scasella @claude
Run the recognition tests on CUDA only
293565c 
The proven class is pinned to the source Inductor emits on NVIDIA CUDA; other
backends may emit different code, where the rule simply does not fire. The two
tests that expect recognition now use requires_cuda_and_triton and skipIfRocm;
the two backend-agnostic tests keep requires_gpu.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@scasella

scasella commented Sep 26, 2026

Copy link
Copy Markdown
Author

@pytorchbot label "release notes: inductor"

@scasella
scasella marked this pull request as ready for review September 26, 2026 22:15
@jansel

jansel commented Sep 29, 2026

Copy link
Copy Markdown
Contributor

@claude review these changes

@claude

claude Bot commented Sep 29, 2026 •
edited
Loading

Copy link
Copy Markdown

Claude finished @jansel's task in 2m 3s —— View job

Needs Discussion: the speedup looks real, but matching a SHA-256 of one generated kernel's source isn't something we can maintain in-tree. The same premises (C2) should instead be checked generically at the point where ks* dtypes are decided.

Full review

PR Review: #198733

Summary

The PR adds an opt-in flag, triton.narrow_proven_size_args. When it's on, a kernel whose canonicalized source hashes to one known SHA-256 (the 6-segment dynamic cat from #189940) gets its ks* size args declared i32 instead of i64, provided some symbolic facts hold. The measured win is plausible. The mechanism is the blocker: it's an identity check against one kernel, not a compiler rule. This answers review question 1: the rule should run on the sympy index/mask expressions, not on generated source.

Code Quality

  • Hashing one generated kernel won't hold up in-tree (triton_size_arg_narrowing.py:37, _PROVED_CAT6_SHA256; _canonical). The digest covers temporary numbering, argument order, the eviction_policy kwarg set, the exact tl.full/tl.broadcast_to spellings, and dead statements.
    • Any unrelated codegen change silently turns the optimization off. Nothing in CI would notice, except the recognition test, and then the fix is "recompute the hash".
    • Each new kernel class would need another ~300-line recognizer plus a hash. The PR's own Limitations section says as much.
    • We shouldn't carry a Python re-parser of Triton source (ast.parse of self.body) inside codegen for one pattern.
  • Swallowing every exception hides bugs (triton_size_arg_narrowing.py:309, the bare except Exception: return []; also the except (_Unsupported, AttributeError, IndexError) at line 222). A real bug in _symbolic_facts, such as a KeyError from name_of[k] or a sympy failure, looks exactly like "not eligible". Catch _Unsupported only.
  • H4 compares guards by string (triton_size_arg_narrowing.py:271-274). str(guard) in {str(g.expr) for g in se.guards} depends on exactly how SIMDScheduling.can_use_32bit_indexing (simd.py:4236) phrased and simplified the guard.
    • The string check is also redundant: kernel.index_dtype == "tl.int32" can only happen after select_index_dtype (simd_kernel_features.py:213) installed or statically proved those bounds.
    • Relying on that invariant is better than re-deriving it with fragile string matching.
  • Run-once cost. sympy.simplify on every candidate kernel is where the 14.7 ms per compile comes from. A generic rule based on bounds and value ranges (below) would use statically_known_leq / bound_sympy instead.

Infrastructure

  • The rule runs at the wrong layer. ks* dtypes are decided in signature_to_meta._decide_tl_dtype (triton_utils.py:259). This PR mutates triton_meta["signature"] afterwards, from codegen_kernel (triton.py:7917). That's a second, hidden place where signature dtypes are chosen.
    • The information needed for a generic rule is already there. SIMDKernelFeatures.select_index_dtype / any_index_expr_overflows_int32 (simd_kernel_features.py:213-295) already bound every MemoryDep index expression against int32, under guards Dynamo re-checks on every call.
    • Suggested direction: a ks whose defining expression (via inv_precomputed_replacements) is bounded by a guarded quantity can be i32. The same must hold for every sub-expression it appears in within the emitted indexing and masks. Here that includes the cat's partial sums (ks1+ks2+ks3, ...), which come from the ModularIndexing/where bounds. Check this with value ranges plus the installed int32 guards, and decide it inside _decide_tl_dtype or at the point where kernel size args are created.
    • That covers this kernel. It also covers every cat or slice of this shape, with no digest and no flag, which also resolves the "fix the root cause, avoid a knob" point from inductor: Fall back to ATen for torch.cat with symbolic non-leading dims #190034.
    • It keeps the trusted base inside Inductor's existing sympy reasoning, instead of a hand-written source parser plus an external Lean proof whose premises we'd have to keep matching the recognizer by hand.
  • New config knob (config.py:2259). This is a private on/off flag for one kernel. It shouldn't land as-is; if a generic rule lands, gate it behind the existing index-dtype machinery instead.

Testing

  • The second-shape check doesn't verify its own claim (test_narrows_proven_cat_kernel). The comment says "reuses the same kernel (no recompile)", but nothing asserts no recompile (e.g. a CompileCounter or torch._dynamo.utils.counters), and nothing asserts the i32 kernel is the one that ran.
  • No boundary test. No test covers the soundness-critical case where Dynamo recompiles to int64 indexing once numel > INT_MAX, the case the PR description mentions exercising at numel = 2**31. At minimum, a test on symbolic facts alone (without allocating 4 GB) that H3/H4 reject is needed, plus a unit test of _symbolic_facts with ks0 != sum(ks1..ks6).
  • The tests confirm the brittleness. Two of the four tests are CUDA-only because the hash differs on other backends, which is the core problem above.

Recommendation

Needs Discussion

The measured L4 speedup (1.39×) is worth chasing, but not as a per-kernel source digest behind a private flag. Please rework it as a generic, bounds-based ks* dtype decision in _decide_tl_dtype / SIMDKernelFeatures that reuses the int32 guards select_index_dtype already installs. Then check how it combines with the alignment work in #193614/#193964 on H100.

@albanD

albanD commented Sep 30, 2026

Copy link
Copy Markdown
Collaborator

Thank you very much for sending a PR.
To avoid you spending time preparing a PR for a feature we're not sure we want to add, our usual process is to discuss on the issue before doing the PR.
In particular, make sure that the issue is marked "actionable" before sending a PR!
See https://github.com/pytorch/pytorch/wiki/The-Ultimate-Guide-to-PyTorch-Contributions for details!

@albanD albanD closed this Sep 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

ciflow/inductor ciflow/torchtitan Run TorchTitan integration tests module: inductor open source release notes: inductor

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[inductor] torch.compile of torch.cat with a symbolic (dynamic) concat dim is 2–2.5× slower than eager (symbolic int64 divmod in generated indexing)

4 participants

@scasella @jansel @albanD @pytorchbot