Skip to content

Fuse Triton TBE forward bounds checking (#4485) - #4485

Closed
stashuk-olek wants to merge 2 commits into
meta-pytorch:mainfrom
stashuk-olek:export-D114284802
Closed

Fuse Triton TBE forward bounds checking (#4485)#4485
stashuk-olek wants to merge 2 commits into
meta-pytorch:mainfrom
stashuk-olek:export-D114284802

Conversation

@stashuk-olek

@stashuk-olek stashuk-olek commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 5, 2026
@meta-codesync

meta-codesync Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

@stashuk-olek has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114284802.

stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
@meta-codesync meta-codesync Bot changed the title Fuse Triton TBE forward bounds checking Fuse Triton TBE forward bounds checking (#4485) Aug 5, 2026
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
@stashuk-olek
stashuk-olek force-pushed the export-D114284802 branch 2 times, most recently from 6ac5661 to 6840aee Compare August 6, 2026 23:03
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
Summary:

Now that both of cuda/triton does it bounds check, we fuse that to avoid traffic. Adding as an option

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
@meta-codesync meta-codesync Bot changed the title Fuse Triton TBE forward bounds checking (#4485) Fuse Triton TBE forward bounds checking Aug 7, 2026
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
Summary:

Use BoundsCheckMode.V2 for Triton TBE standalone index validation. This preserves the existing warning and offset-repair behavior while reducing validation overhead before Triton forward.

Reviewed By: axeisghost, TroyGarden

Differential Revision: D114279989
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
@meta-codesync meta-codesync Bot changed the title Fuse Triton TBE forward bounds checking Fuse Triton TBE forward bounds checking (#4485) Aug 8, 2026
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
Summary:

Validate and repair indices inside supported unweighted, non-VBE Triton forward kernels. Unsupported configurations retain the CUDA v2 fallback, while fused checked loads preserve warning accumulation without a separate full-index pass.

Reviewed By: TroyGarden, axeisghost

Differential Revision: D114284802
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 Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant