Skip to content

Add CPU nbit-forward tests for empty/pruned-bag zero-fill (#6064) - #6064

Open
knoebelja wants to merge 1 commit into
pytorch:mainfrom
knoebelja:export-D113243672
Open

Add CPU nbit-forward tests for empty/pruned-bag zero-fill (#6064)#6064
knoebelja wants to merge 1 commit into
pytorch:mainfrom
knoebelja:export-D113243672

Conversation

@knoebelja

@knoebelja knoebelja commented Jul 24, 2026

Copy link
Copy Markdown

Summary:

X-link: https://github.com/facebookresearch/FBGEMM/pull/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.

Differential Revision: D113243672

@meta-cla meta-cla Bot added the cla signed label Jul 24, 2026
@meta-codesync

meta-codesync Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

@knoebelja has exported this pull request. If you are a Meta employee, you can view the originating Diff in D113243672.

@meta-codesync meta-codesync Bot changed the title Add CPU nbit-forward tests for empty/pruned-bag zero-fill Add CPU nbit-forward tests for empty/pruned-bag zero-fill (#6064) Jul 24, 2026
@knoebelja
knoebelja force-pushed the export-D113243672 branch 2 times, most recently from 57e0a4c to d62f3e0 Compare July 27, 2026 13:03
@meta-codesync meta-codesync Bot changed the title Add CPU nbit-forward tests for empty/pruned-bag zero-fill (#6064) Add CPU nbit-forward tests for empty/pruned-bag zero-fill Jul 27, 2026
@knoebelja
knoebelja force-pushed the export-D113243672 branch from d62f3e0 to 77e011e Compare July 27, 2026 17:58
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.

Differential Revision: D113243672
@meta-codesync meta-codesync Bot changed the title Add CPU nbit-forward tests for empty/pruned-bag zero-fill Add CPU nbit-forward tests for empty/pruned-bag zero-fill (#6064) Jul 28, 2026
@knoebelja
knoebelja force-pushed the export-D113243672 branch from 77e011e to 8de761f Compare July 28, 2026 18:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant