Support multi-dimensional runtime_meta in RES streaming buffers by lazy init (#5643)#5643

Closed

FriedCosey wants to merge 1 commit intopytorch:mainfrom

FriedCosey:export-D100944325

FriedCosey commented Apr 15, 2026 •

edited by meta-codesync Bot

Loading

Summary:

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The res_runtime_meta buffer in SplitTableBatchedEmbeddingBagsCodegen was hardcoded to shape (cache_size, 1). When _hash_zch_runtime_meta has dim > 1 (e.g., feature cache storing 2 cached features via zch_custom_runtime_meta_dim=2), the .copy_() in raw_embedding_stream() crashes with:
RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]

Full output P2274437578

This diff lazy resizes the buffer defaults to (cache_size, 1, torch.long) and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

meta-cla Bot added the cla signed label

Contributor

meta-codesync Bot commented Apr 15, 2026

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

meta-codesync Bot added fb-exported meta-exported labels

meta-codesync Bot changed the title ~~Support multi-dimensional runtime_meta in RES streaming buffers by lazy init~~ Support multi-dimensional runtime_meta in RES streaming buffers by lazy init (#5643)

FriedCosey force-pushed the export-D100944325 branch from 2c863f6 to 3429c2d Compare

April 16, 2026 21:17

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

3429c2d

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

93818b6

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 3429c2d to 93818b6 Compare

April 16, 2026 21:20

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

716155d

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 93818b6 to 716155d Compare

April 16, 2026 21:45

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

d06edd7

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 716155d to d06edd7 Compare

April 16, 2026 21:58

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

74eba68

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from d06edd7 to 74eba68 Compare

April 16, 2026 22:05

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

bbacffe

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 74eba68 to bbacffe Compare

April 16, 2026 22:27

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

3f33c63

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from bbacffe to 3f33c63 Compare

April 16, 2026 22:56

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

b494721

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 3f33c63 to b494721 Compare

April 16, 2026 22:57

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from b494721 to 0530779 Compare

April 16, 2026 22:59

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

30786c1

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 0530779 to 30786c1 Compare

April 16, 2026 23:00

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

705a0e6

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 30786c1 to 705a0e6 Compare

April 16, 2026 23:00

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

9f9e21e

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 705a0e6 to 9f9e21e Compare

April 16, 2026 23:09

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

3ae4bb0

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 9f9e21e to 3ae4bb0 Compare

April 16, 2026 23:17

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

f75438b

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 3ae4bb0 to f75438b Compare

April 16, 2026 23:27

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

a33583a

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from f75438b to a33583a Compare

April 17, 2026 02:08

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

a8a7da6

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from a33583a to a8a7da6 Compare

April 17, 2026 02:10

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

57dc926

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch 2 times, most recently from 57dc926 to 6f26962 Compare

April 17, 2026 02:11

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

6f26962

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

7211bd0

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 6f26962 to 7211bd0 Compare

April 17, 2026 02:12

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

6c69ec7

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 7211bd0 to 6c69ec7 Compare

April 17, 2026 02:16

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

06bdfaa

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch from 6c69ec7 to 06bdfaa Compare

April 17, 2026 02:28

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

49b9861

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

FriedCosey force-pushed the export-D100944325 branch 2 times, most recently from 49b9861 to e091dc0 Compare

April 18, 2026 06:20

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

e091dc0

…zy init (pytorch#5643)

Summary:

X-link: facebookresearch/FBGEMM#2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325


          Support multi-dimensional runtime_meta in RES streaming buffers by la…

8e5c9a5

…zy init (pytorch#5643)

Summary:
Pull Request resolved: pytorch#5643

X-link: https://github.com/facebookresearch/FBGEMM/pull/2591

The `res_runtime_meta` buffer in `SplitTableBatchedEmbeddingBagsCodegen` was hardcoded to shape `(cache_size, 1)`. When `_hash_zch_runtime_meta` has dim > 1 (e.g., feature cache storing 2 cached features via `zch_custom_runtime_meta_dim=2`), the `.copy_()` in `raw_embedding_stream()` crashes with:
`RuntimeError: output with shape [N, 1] doesn't match the broadcast shape [N, 2]`

Full output P2274437578

This diff lazy resizes the buffer defaults to `(cache_size, 1, torch.long)` and auto-corrects on the first iteration when runtime_meta data arrives with a different shape or dtype. This is a one-time operation, after the first resize, dims match and no further reallocation occurs.

Reviewed By: chouxi

Differential Revision: D100944325

Contributor

meta-codesync Bot commented Apr 20, 2026

This pull request has been merged in ffdfb0e.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed fb-exported Merged meta-exported