Drop redundant output pre-zero in quantized TBE CPU pooled forward (#6065) - #6065
Open
knoebelja wants to merge 2 commits into
Open
Drop redundant output pre-zero in quantized TBE CPU pooled forward (#6065)#6065knoebelja wants to merge 2 commits into
knoebelja wants to merge 2 commits into
Conversation
Contributor
|
@knoebelja has exported this pull request. If you are a Meta employee, you can view the originating Diff in D113243671. |
knoebelja
added a commit
to knoebelja/FBGEMM
that referenced
this pull request
Jul 27, 2026
…ytorch#6065) Summary: X-link: facebookresearch/FBGEMM#2965 The pooled quantized TBE CPU forward (`int_nbit_split_embedding_codegen_forward_{unweighted,weighted}_cpu`, generated from `embedding_forward_quantized_cpu_template.cpp`) allocates its output with `at::empty(...)` and then unconditionally `output.fill_(0)` for FP32/FP16/BF16. That memset is redundant: the FBGEMM SpMDM kernel writes every output element itself -- each bag's row is emitted via `fbgemm::fill_output`, and empty bags are zeroed via `fbgemm::fillZero` first. For FP32/FP16/BF16, `total_adjusted_D == total_D` and `D_offsets` fully partition `[0, total_D)`, so every column of every row is overwritten. On `ipnext_ads_prod/deployment` this `at::Tensor::fill_` is ~15.83% of the unweighted forward (~\$0.5M/yr; the op also pays the full boxed ATen dispatch around the memset). The win is fleet-wide across CPU-TBE inference services. Safety: an out-of-bounds index is clamped to row 0 by the CPU bounds check before the kernel runs, so the kernel writes every row. If the kernel still returns `!success` (e.g. malformed offsets), the `!success` block now zeroes the output before calling `report_embedding_error` (which TORCH_CHECK-throws), so a partially-written output is never observed. Scope: pooled (non-`nobag`) FP32/FP16/BF16 only. int8/int4 (inline qparams) and the `nobag` path are untouched. The behavior this relies on is pinned by the CPU tests in the parent diff. Differential Revision: D113243671
meta-codesync
Bot
changed the title
knoebelja
force-pushed
the
export-D113243671
branch
from
July 27, 2026 13:03
df06819 to
8d8a77a
Compare
knoebelja
added a commit
to knoebelja/FBGEMM
that referenced
this pull request
Jul 27, 2026
…ytorch#6065) Summary: X-link: facebookresearch/FBGEMM#2965 The pooled quantized TBE CPU forward (`int_nbit_split_embedding_codegen_forward_{unweighted,weighted}_cpu`, generated from `embedding_forward_quantized_cpu_template.cpp`) allocates its output with `at::empty(...)` and then unconditionally `output.fill_(0)` for FP32/FP16/BF16. That memset is redundant: the FBGEMM SpMDM kernel writes every output element itself -- each bag's row is emitted via `fbgemm::fill_output`, and empty bags are zeroed via `fbgemm::fillZero` first. For FP32/FP16/BF16, `total_adjusted_D == total_D` and `D_offsets` fully partition `[0, total_D)`, so every column of every row is overwritten. On `ipnext_ads_prod/deployment` this `at::Tensor::fill_` is ~15.83% of the unweighted forward (~\$0.5M/yr; the op also pays the full boxed ATen dispatch around the memset). The win is fleet-wide across CPU-TBE inference services. Safety: an out-of-bounds index is clamped to row 0 by the CPU bounds check before the kernel runs, so the kernel writes every row. If the kernel still returns `!success` (e.g. malformed offsets), the `!success` block now zeroes the output before calling `report_embedding_error` (which TORCH_CHECK-throws), so a partially-written output is never observed. Scope: pooled (non-`nobag`) FP32/FP16/BF16 only. int8/int4 (inline qparams) and the `nobag` path are untouched. The behavior this relies on is pinned by the CPU tests in the parent diff. Differential Revision: D113243671
knoebelja
force-pushed
the
export-D113243671
branch
from
July 27, 2026 17:58
8d8a77a to
a4b047c
Compare
knoebelja
added a commit
to knoebelja/FBGEMM
that referenced
this pull request
Jul 28, 2026
…ytorch#6065) Summary: X-link: facebookresearch/FBGEMM#2965 The pooled quantized TBE CPU forward (`int_nbit_split_embedding_codegen_forward_{unweighted,weighted}_cpu`, generated from `embedding_forward_quantized_cpu_template.cpp`) allocates its output with `at::empty(...)` and then unconditionally `output.fill_(0)` for FP32/FP16/BF16. That memset is redundant: the FBGEMM SpMDM kernel writes every output element itself -- each bag's row is emitted via `fbgemm::fill_output`, and empty bags are zeroed via `fbgemm::fillZero` first. For FP32/FP16/BF16, `total_adjusted_D == total_D` and `D_offsets` fully partition `[0, total_D)`, so every column of every row is overwritten. On `ipnext_ads_prod/deployment` this `at::Tensor::fill_` is ~15.83% of the unweighted forward (~\$0.5M/yr; the op also pays the full boxed ATen dispatch around the memset). The win is fleet-wide across CPU-TBE inference services. Safety: an out-of-bounds index is clamped to row 0 by the CPU bounds check before the kernel runs, so the kernel writes every row. If the kernel still returns `!success` (e.g. malformed offsets), the `!success` block now zeroes the output before calling `report_embedding_error` (which TORCH_CHECK-throws), so a partially-written output is never observed. Scope: pooled (non-`nobag`) FP32/FP16/BF16 only. int8/int4 (inline qparams) and the `nobag` path are untouched. The behavior this relies on is pinned by the CPU tests in the parent diff. Differential Revision: D113243671
knoebelja
force-pushed
the
export-D113243671
branch
from
July 28, 2026 18:39
a4b047c to
e06220b
Compare
Summary: X-link: facebookresearch/FBGEMM#2964 Add CPU coverage for the pooled quantized TBE forward `int_nbit_split_embedding_codegen_forward_{unweighted,weighted}_cpu`: - `test_nbit_forward_cpu_empty_and_pruned_bags_zero_fill`: single-table SUM pooling with every embedding row set to 1.0 and ragged bag lengths that include empty bags (`L == 0`) and a fully-pruned bag, with ~1/3 of indices pruned to -1. Each output dim must equal the non-pruned count (or weighted sum), and empty/pruned bags must be exactly 0 -- verifying the kernel writes every output row itself. - `test_nbit_forward_cpu_out_of_bounds_index_stays_defined`: an out-of-bounds index is sanitized by the CPU bounds check (clamped to row 0), so the op returns a fully-written, finite output. These pass on current code and pin the behavior that the stacked follow-up relies on when it drops the redundant `output.fill_(0)` pre-zero from the pooled path. Reviewed By: q10 Differential Revision: D113243672
…ytorch#6065) Summary: X-link: facebookresearch/FBGEMM#2965 The pooled quantized TBE CPU forward (`int_nbit_split_embedding_codegen_forward_{unweighted,weighted}_cpu`, generated from `embedding_forward_quantized_cpu_template.cpp`) allocates its output with `at::empty(...)` and then unconditionally `output.fill_(0)` for FP32/FP16/BF16. That memset is redundant: the FBGEMM SpMDM kernel writes every output element itself -- each bag's row is emitted via `fbgemm::fill_output`, and empty bags are zeroed via `fbgemm::fillZero` first. For FP32/FP16/BF16, `total_adjusted_D == total_D` and `D_offsets` fully partition `[0, total_D)`, so every column of every row is overwritten. On `ipnext_ads_prod/deployment` this `at::Tensor::fill_` is ~15.83% of the unweighted forward (~\$0.5M/yr; the op also pays the full boxed ATen dispatch around the memset). The win is fleet-wide across CPU-TBE inference services. Safety: an out-of-bounds index is clamped to row 0 by the CPU bounds check before the kernel runs, so the kernel writes every row. If the kernel still returns `!success` (e.g. malformed offsets), the `!success` block now zeroes the output before calling `report_embedding_error` (which TORCH_CHECK-throws), so a partially-written output is never observed. Scope: pooled (non-`nobag`) FP32/FP16/BF16 only. int8/int4 (inline qparams) and the `nobag` path are untouched. The behavior this relies on is pinned by the CPU tests in the parent diff. Reviewed By: q10, wenxin0319 Differential Revision: D113243671
knoebelja
force-pushed
the
export-D113243671
branch
from
July 31, 2026 15:35
e06220b to
b787c49
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2965
The pooled quantized TBE CPU forward
(
int_nbit_split_embedding_codegen_forward_{unweighted,weighted}_cpu, generatedfrom
embedding_forward_quantized_cpu_template.cpp) allocates its output withat::empty(...)and then unconditionallyoutput.fill_(0)for FP32/FP16/BF16.That memset is redundant: the FBGEMM SpMDM kernel writes every output element
itself -- each bag's row is emitted via
fbgemm::fill_output, and empty bags arezeroed via
fbgemm::fillZerofirst. For FP32/FP16/BF16,total_adjusted_D == total_DandD_offsetsfully partition[0, total_D), so every column of everyrow is overwritten.
On
ipnext_ads_prod/deploymentthisat::Tensor::fill_is15.83% of the$0.5M/yr; the op also pays the full boxed ATen dispatchunweighted forward (
around the memset). The win is fleet-wide across CPU-TBE inference services.
Safety: an out-of-bounds index is clamped to row 0 by the CPU bounds check before
the kernel runs, so the kernel writes every row. If the kernel still returns
!success(e.g. malformed offsets), the!successblock now zeroes the outputbefore calling
report_embedding_error(which TORCH_CHECK-throws), so apartially-written output is never observed.
Scope: pooled (non-
nobag) FP32/FP16/BF16 only. int8/int4 (inline qparams) andthe
nobagpath are untouched. The behavior this relies on is pinned by the CPUtests in the parent diff.
Reviewed By: q10, wenxin0319
Differential Revision: D113243671