Skip to content

Expose fused Triton TBE bounds checking controls (#4492) - #4492

Closed
stashuk-olek wants to merge 9 commits into
meta-pytorch:mainfrom
stashuk-olek:export-D114286097
Closed

Expose fused Triton TBE bounds checking controls (#4492)#4492
stashuk-olek wants to merge 9 commits into
meta-pytorch:mainfrom
stashuk-olek:export-D114286097

Conversation

@stashuk-olek

@stashuk-olek stashuk-olek commented Aug 5, 2026 •
edited by meta-codesync Bot
Loading

Copy link
Copy Markdown
Contributor

Summary:

Thread the fused-bounds-check option through TorchRec fused parameters and benchmark configuration. The default remains unchanged, while supported callers can select the fused Triton route without bypassing validation.

Reviewed By: axeisghost

Differential Revision: D114286097

@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 D114286097.

@meta-codesync meta-codesync Bot changed the title Expose fused Triton TBE bounds checking controls Expose fused Triton TBE bounds checking controls (#4492) Aug 5, 2026
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
@stashuk-olek @facebook-github-bot
Expose fused Triton TBE bounds checking controls (meta-pytorch#4492)
d689d09 
Summary:

Just small convenience - expose in the benchmark to test easier.

Reviewed By: axeisghost

Differential Revision: D114286097
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
@stashuk-olek @facebook-github-bot
Expose fused Triton TBE bounds checking controls (meta-pytorch#4492)
82ba611 
Summary:

Just small convenience - expose in the benchmark to test easier.

Reviewed By: axeisghost

Differential Revision: D114286097
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
@stashuk-olek @facebook-github-bot
Expose fused Triton TBE bounds checking controls (meta-pytorch#4492)
82ca9dd 
Summary:

Just small convenience - expose in the benchmark to test easier.

Reviewed By: axeisghost

Differential Revision: D114286097
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
@stashuk-olek @facebook-github-bot
Expose fused Triton TBE bounds checking controls (meta-pytorch#4492)
454d1ff 
Summary:

Just small convenience - expose in the benchmark to test easier.

Reviewed By: axeisghost

Differential Revision: D114286097
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
@stashuk-olek @facebook-github-bot
Expose fused Triton TBE bounds checking controls (meta-pytorch#4492)
044adf1 
Summary:

Just small convenience - expose in the benchmark to test easier.

Reviewed By: axeisghost

Differential Revision: D114286097
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
@stashuk-olek @facebook-github-bot
Expose fused Triton TBE bounds checking controls (meta-pytorch#4492)
03234ec 
Summary:

Just small convenience - expose in the benchmark to test easier.

Reviewed By: axeisghost

Differential Revision: D114286097
@stashuk-olek
stashuk-olek force-pushed the export-D114286097 branch 2 times, most recently from 03234ec to f6ab890 Compare August 6, 2026 23:04
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 6, 2026
@stashuk-olek @facebook-github-bot
Expose fused Triton TBE bounds checking controls (meta-pytorch#4492)
f6ab890 
Summary:

Just small convenience - expose in the benchmark to test easier.

Reviewed By: axeisghost

Differential Revision: D114286097
@meta-codesync meta-codesync Bot changed the title Expose fused Triton TBE bounds checking controls (#4492) Expose fused Triton TBE bounds checking controls Aug 7, 2026
@stashuk-olek @facebook-github-bot
Select CUDA v2 bounds checking for Triton TBE (meta-pytorch#4484)
737ff37 
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
@stashuk-olek @facebook-github-bot
Fuse Triton TBE forward bounds checking (meta-pytorch#4485)
076c527 
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 @facebook-github-bot
Optimize Triton TBE forward for small long-bag tables (meta-pytorch#4486
7b5b197 
)

Summary:

Convert long bags over small tables into per-bag row histograms followed by a tensor-core dot. Generic row, dimension, and bag-size predicates select the route, and checked loads preserve fused bounds behavior.

Reviewed By: axeisghost

Differential Revision: D114273027
@stashuk-olek @facebook-github-bot
Reduce Triton TBE backward reduction pressure (meta-pytorch#4487)
01f1098 
Summary:

Restructure short- and long-run gradient accumulation to reduce serialized reduction pressure. The new reduction schedule exposes more independent work while preserving the existing exact rowwise optimizer update.

Reviewed By: axeisghost

Differential Revision: D114273026
@stashuk-olek @facebook-github-bot
Reduce Triton TBE unweighted backward reduction width (meta-pytorch#4488
d194348 
)

Summary:

Reduce the dout reduction tile from 16 to 8 in the unweighted short-run, separate long-run accumulation, and fused long-run kernels. The smaller tile reduces live vector state and register pressure on B200 while retaining FP32 accumulation and the existing rowwise optimizer update. Weighted kernels remain unchanged.

Reviewed By: axeisghost

Differential Revision: D114403703
@stashuk-olek @facebook-github-bot
Coarsen Triton TBE forward bag scheduling (meta-pytorch#4489)
e979a84 
Summary:

Process multiple bags per Triton program where the workload has enough parallelism. This amortizes program overhead and exposes independent loads without changing the embedding reduction.

Reviewed By: axeisghost

Differential Revision: D114280197
@stashuk-olek @facebook-github-bot
Widen Triton TBE forward gather issue (meta-pytorch#4490)
4bfe857 
Summary:

Issue up to eight independent index gathers per loop iteration on profitable long-bag shapes. The wider loop increases memory-level parallelism while retaining the four-wide path for shapes where register pressure dominates.

Reviewed By: axeisghost

Differential Revision: D114280947
@stashuk-olek @facebook-github-bot
Reuse forward histograms in Triton TBE backward (meta-pytorch#4491)
6539cc4 
Summary:

Reuse exact counts from an eligible leading, unique small-table feature during exact rowwise Adagrad backward. Immutable plans carry actual table geometry, compensated FP16 high/residual matrix products preserve FP32 accuracy, and scalar validity enables safe fallback for long bags, aliases, graph capture, weighted inputs, and VBE.

Reviewed By: axeisghost

Differential Revision: D114284800
@stashuk-olek @facebook-github-bot
Expose fused Triton TBE bounds checking controls (meta-pytorch#4492)
5f11e1e 
Summary:

Thread the fused-bounds-check option through TorchRec fused parameters and benchmark configuration. The default remains unchanged, while supported callers can select the fused Triton route without bypassing validation.

Reviewed By: axeisghost

Differential Revision: D114286097
@meta-codesync meta-codesync Bot changed the title Expose fused Triton TBE bounds checking controls Expose fused Triton TBE bounds checking controls (#4492) Aug 8, 2026
@meta-codesync meta-codesync Bot closed this in 3d0927f Aug 8, 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

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

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@stashuk-olek