Skip to content

Fix VBE batch sizes not passed to request builder (#5653)#5653

Closed
gregmacnamara wants to merge 1 commit into
pytorch:mainfrom
gregmacnamara:export-D94380814
Closed

Fix VBE batch sizes not passed to request builder (#5653)#5653
gregmacnamara wants to merge 1 commit into
pytorch:mainfrom
gregmacnamara:export-D94380814

Conversation

@gregmacnamara
Copy link
Copy Markdown
Contributor

@gregmacnamara gregmacnamara commented Apr 17, 2026

Summary:

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

When batch_size_per_feature_per_rank is provided externally (e.g. from a
sharding plan), it was not being forwarded to _build_requests_jagged.
Both branches of the VBE/non-VBE conditional were passing the raw
batch_size_per_feature_per_rank parameter (which could be None in the else
branch) instead of the resolved Bs_feature_rank from _generate_batch_sizes.

This fix:

  1. Captures the externally-provided batch sizes into Bs_feature_rank when
    batch_size_per_feature_per_rank is set
  2. Uses _generate_batch_sizes return value otherwise
  3. Removes the duplicate if/else branches that were doing the same thing
  4. Passes the resolved Bs_feature_rank to _build_requests_jagged

Reviewed By: q10

Differential Revision: D94380814

@meta-cla meta-cla Bot added the cla signed label Apr 17, 2026
@meta-codesync
Copy link
Copy Markdown
Contributor

meta-codesync Bot commented Apr 17, 2026

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

@meta-codesync meta-codesync Bot changed the title Fix VBE batch sizes not passed to request builder Fix VBE batch sizes not passed to request builder (#5653) Apr 17, 2026
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
@gregmacnamara gregmacnamara force-pushed the export-D94380814 branch 2 times, most recently from 36ac52f to 0d9ad95 Compare April 17, 2026 17:47
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:
Pull Request resolved: pytorch#5653

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:
Pull Request resolved: pytorch#5653

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
@gregmacnamara gregmacnamara force-pushed the export-D94380814 branch 2 times, most recently from a42d0fc to 9c9b46b Compare April 17, 2026 18:27
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:
Pull Request resolved: pytorch#5653

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 17, 2026
Summary:
Pull Request resolved: pytorch#5653

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Differential Revision: D94380814
@meta-codesync meta-codesync Bot changed the title Fix VBE batch sizes not passed to request builder (#5653) Fix VBE batch sizes not passed to request builder Apr 23, 2026
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 27, 2026
Summary:

X-link: facebookresearch/FBGEMM#2626

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Reviewed By: q10

Differential Revision: D94380814
@meta-codesync meta-codesync Bot changed the title Fix VBE batch sizes not passed to request builder Fix VBE batch sizes not passed to request builder (#5653) Apr 27, 2026
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 27, 2026
Summary:

X-link: facebookresearch/FBGEMM#2626

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Reviewed By: q10

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 27, 2026
Summary:

X-link: facebookresearch/FBGEMM#2626

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Reviewed By: q10

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 27, 2026
Summary:

X-link: facebookresearch/FBGEMM#2626

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Reviewed By: q10

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 27, 2026
Summary:

X-link: facebookresearch/FBGEMM#2626

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Reviewed By: q10

Differential Revision: D94380814
Summary:

X-link: facebookresearch/FBGEMM#2626

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Reviewed By: q10

Differential Revision: D94380814
gregmacnamara added a commit to gregmacnamara/FBGEMM that referenced this pull request Apr 27, 2026
Summary:

X-link: facebookresearch/FBGEMM#2626

When `batch_size_per_feature_per_rank` is provided externally (e.g. from a
sharding plan), it was not being forwarded to `_build_requests_jagged`.
Both branches of the VBE/non-VBE conditional were passing the raw
`batch_size_per_feature_per_rank` parameter (which could be None in the else
branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`.

This fix:
1. Captures the externally-provided batch sizes into `Bs_feature_rank` when
   `batch_size_per_feature_per_rank` is set
2. Uses `_generate_batch_sizes` return value otherwise
3. Removes the duplicate if/else branches that were doing the same thing
4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged`

Reviewed By: q10

Differential Revision: D94380814
@meta-codesync meta-codesync Bot closed this in b8dc9a1 Apr 28, 2026
@meta-codesync
Copy link
Copy Markdown
Contributor

meta-codesync Bot commented Apr 28, 2026

This pull request has been merged in b8dc9a1.

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