Skip to content

Chunk the D2H copy across N copy threads in raw_embedding_streamer - #6068

Open
FriedCosey wants to merge 1 commit into
pytorch:mainfrom
FriedCosey:export-D113594180
Open

Chunk the D2H copy across N copy threads in raw_embedding_streamer#6068
FriedCosey wants to merge 1 commit into
pytorch:mainfrom
FriedCosey:export-D113594180

Conversation

@FriedCosey

Copy link
Copy Markdown

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

Speed up the RES C++ streamer's device->host copy by parallelizing it.

Previously stream() copied all updated rows to CPU on a single thread and enqueued one item. This replaces that serial copy with a chunked copy across up to kNumCopyThreads threads: tensor_copy_chunk() copies a [start, end) row range, and the per-thread tiling is factored into a pure, unit-testable computeChunkRanges().

Differential Revision: D113594180

Summary:
X-link: facebookresearch/FBGEMM#2968

Speed up the RES C++ streamer's device->host copy by parallelizing it.
  
Previously stream() copied all updated rows to CPU on a single thread and enqueued one item. This replaces that serial copy with a chunked copy across up to `kNumCopyThreads` threads: tensor_copy_chunk() copies a [start, end) row range, and the per-thread tiling is factored into a pure, unit-testable computeChunkRanges().

Differential Revision: D113594180
@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

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

FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request Jul 25, 2026
…ytorch#6068)

Summary:

X-link: facebookresearch/FBGEMM#2968

Speed up the RES C++ streamer's device->host copy by parallelizing it.
  
Previously stream() copied all updated rows to CPU on a single thread and enqueued one item. This replaces that serial copy with a chunked copy across up to `kNumCopyThreads` threads: tensor_copy_chunk() copies a [start, end) row range, and the per-thread tiling is factored into a pure, unit-testable computeChunkRanges().

Differential Revision: D113594180
FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request Jul 25, 2026
…ytorch#6068)

Summary:

X-link: facebookresearch/FBGEMM#2968

Speed up the RES C++ streamer's device->host copy by parallelizing it.
  
Previously stream() copied all updated rows to CPU on a single thread and enqueued one item. This replaces that serial copy with a chunked copy across up to `kNumCopyThreads` threads: tensor_copy_chunk() copies a [start, end) row range, and the per-thread tiling is factored into a pure, unit-testable computeChunkRanges().

Differential Revision: D113594180
FriedCosey pushed a commit to FriedCosey/FBGEMM that referenced this pull request Jul 25, 2026
…ytorch#6068)

Summary:

X-link: facebookresearch/FBGEMM#2968

Speed up the RES C++ streamer's device->host copy by parallelizing it.
  
Previously stream() copied all updated rows to CPU on a single thread and enqueued one item. This replaces that serial copy with a chunked copy across up to `kNumCopyThreads` threads: tensor_copy_chunk() copies a [start, end) row range, and the per-thread tiling is factored into a pure, unit-testable computeChunkRanges().

Differential Revision: D113594180
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