Skip to content

Optimize Triton TBE forward for small long-bag tables (#4486) - #4486

Closed
stashuk-olek wants to merge 3 commits into
meta-pytorch:mainfrom
stashuk-olek:export-D114273027
Closed

Optimize Triton TBE forward for small long-bag tables (#4486)#4486
stashuk-olek wants to merge 3 commits into
meta-pytorch:mainfrom
stashuk-olek:export-D114273027

Conversation

@stashuk-olek

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

Copy link
Copy Markdown
Contributor

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

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

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

Differential Revision: D114273027
@meta-codesync meta-codesync Bot changed the title Optimize Triton TBE forward for small long-bag tables Optimize Triton TBE forward for small long-bag tables (#4486) Aug 5, 2026
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 5, 2026
)

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

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

Summary:

So there's a ton of gather there that are quite expensive. An interesting trick would be to try to count elements in rows with histogram and then do single TC dot to restore them.

Reviewed By: axeisghost

Differential Revision: D114273027
@meta-codesync meta-codesync Bot changed the title Optimize Triton TBE forward for small long-bag tables (#4486) Optimize Triton TBE forward for small long-bag tables Aug 7, 2026
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
)

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 added a commit to stashuk-olek/torchrec that referenced this pull request Aug 7, 2026
)

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 added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
)

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 added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
)

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 added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
)

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 added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
)

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
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
)

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
@meta-codesync meta-codesync Bot changed the title Optimize Triton TBE forward for small long-bag tables Optimize Triton TBE forward for small long-bag tables (#4486) Aug 8, 2026
stashuk-olek added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
)

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 added a commit to stashuk-olek/torchrec that referenced this pull request Aug 8, 2026
)

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
@meta-codesync meta-codesync Bot closed this in 0c478d9 Aug 8, 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 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