Conversation
D119363782 (pytorch#6278) landed with a different resolution of a CUDA host-build break than the one on the PR branch, and that resolution dropped two ROCm launch-config values. The break was that pytorch#6278 added #include "fbgemm_gpu/utils/cuda_prelude.cuh" to embedding_backward_split_host_template.cpp so kWarpSizeHost() would be in scope. The generated gen_embedding_backward_split_*.cpp are host sources compiled by the plain host compiler, which has no declaration for that header's device intrinsics, so every CUDA build failed on __ballot_sync and __syncwarp. The landed version resolved it by deleting the whole `#ifdef USE_ROCM` block instead, which silently moved ROCm from BT_block_size 64 and max_segment_length_per_warp 16384 to the CUDA values of 32 and 32. Both of those predate the warp-size work; max_segment_length_per_warp is the long-run threshold handed to the unique-index counting kernel, so lowering it to 32 reclassifies most segments as long runs and routes them to the cta_per_row kernel on ROCm. This restores the block. max_segment_length_per_warp goes back to its previous 16384. BT_block_size becomes kWarpSizeHost() rather than the previous hardcoded 64, which is what pytorch#6278 set out to do for every other launch site and what already landed in the generated CUDA sources: a ROCm wheel serving both wave64 and wave32 architectures compiles this host code once, so a compile-time warp size is wrong for whichever architecture it does not match. On CUDA and on wave64-only ROCm builds the runtime query folds to the same value as before. Calling kWarpSizeHost() from a host source needs it declared without dragging in device code, so kWarpSize and kWarpSizeHost move verbatim from cuda_prelude.cuh into a new utils/warp_size.h that cuda_prelude.cuh includes. Every existing device-side user keeps seeing them unchanged. Under USE_ROCM, at::cuda::warp_size() is declared by <c10/macros/Macros.h>, so the new header needs no CUDA or HIP header of its own; .h files are not hipified, and running hipify over the moved block produces no substitutions, so it does not need to be in the hipified source set. The alternative was to guard cuda_prelude.cuh's device helpers on __CUDACC__/__HIPCC__ and keep the original include. That edits a header included by a large amount of device code to fix one host call site, so splitting out the two definitions host code legitimately needs is the narrower change. Test Plan: Compiled the new header standalone with the plain host compiler, with and without USE_ROCM, against installed torch headers only, with no CUDA or HIP headers on the include path: ``` g++ -std=c++20 -fsyntax-only -DUSE_ROCM -I fbgemm_gpu/include -I $TORCH_INCLUDE -I $TORCH_INCLUDE/torch/csrc/api/include hdr_probe.cpp g++ -std=c++20 -fsyntax-only -I fbgemm_gpu/include -I $TORCH_INCLUDE -I $TORCH_INCLUDE/torch/csrc/api/include hdr_probe.cpp ``` Confirmed hipify makes no substitutions in the moved block: ``` python3 -c "from torch.utils.hipify.hipify_python import hipify; hipify(project_directory='hipify_probe', output_directory='hipify_probe', includes=['*'], is_pytorch_extension=True)" md5sum hipify_probe/warp_size_probe.cuh fbgemm_gpu/include/fbgemm_gpu/utils/warp_size.h # equal ``` Rendered the TBE backward codegen and confirmed the restored block reaches the generated host sources and that none of them pulls in cuda_prelude.cuh: ``` python3 fbgemm_gpu/codegen/genscript/generate_backward_split.py --opensource --install_dir /tmp/gen cd /tmp/gen grep -l '^#include.*cuda_prelude' *.cpp # no matches grep -l '^#include.*warp_size.h' *.cpp # 21 matches grep -A 8 '^#ifdef USE_ROCM' gen_embedding_backward_split_dense.cpp ``` Full CUDA and ROCm builds are left to CI. Authored with assistance from Claude (Anthropic).
Contributor
|
@q10 has imported this pull request. If you are a Meta employee, you can view this in D120626878. |
q10
added a commit
to q10/FBGEMM
that referenced
this pull request
Sep 18, 2026
Summary: Restores the ROCm V1 TBE backward launch configuration that was dropped while fixing the host-only CUDA build: `max_segment_length_per_warp` returns to 16384, and `BT_block_size` follows the active GPU warp size. Moves `kWarpSize` and `kWarpSizeHost()` from `cuda_prelude.cuh` into the host-safe `utils/warp_size.h`, so generated host sources can use the runtime query without including device intrinsics. Marks the internal-linkage `kWarpSize` declarations `[[maybe_unused]]` because GCC diagnoses them in generated host translation units where only `kWarpSizeHost()` is used. Defers the ROCm runtime query until `dev_weights` is on GPU. CPU dispatches retain the prior fallback of 64 and do not initialize the HIP runtime. This follows up D119363782 and pytorch#6278, corresponding to pytorch#6323. Reviewed By: gchalump Differential Revision: D120626878
meta-codesync Bot
pushed a commit
that referenced
this pull request
Sep 18, 2026
Summary: X-link: https://github.com/facebookresearch/FBGEMM/pull/3199 Pull Request resolved: #6326 Restores the ROCm V1 TBE backward launch configuration that was dropped while fixing the host-only CUDA build: `max_segment_length_per_warp` returns to 16384, and `BT_block_size` follows the active GPU warp size. Moves `kWarpSize` and `kWarpSizeHost()` from `cuda_prelude.cuh` into the host-safe `utils/warp_size.h`, so generated host sources can use the runtime query without including device intrinsics. Marks the internal-linkage `kWarpSize` declarations `[[maybe_unused]]` because GCC diagnoses them in generated host translation units where only `kWarpSizeHost()` is used. Defers the ROCm runtime query until `dev_weights` is on GPU. CPU dispatches retain the prior fallback of 64 and do not initialize the HIP runtime. This follows up D119363782 and #6278, corresponding to #6323. Reviewed By: gchalump Differential Revision: D120626878 fbshipit-source-id: 618c4048c4f5d5e412bbdde010327f480dad7101
Contributor
|
Closed in 659e41f |
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
2 participants
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.
D119363782 (#6278) landed with a different resolution of a CUDA host-build break than the one on the PR branch, and that resolution dropped two ROCm launch-config values.
The break was that #6278 added
#include "fbgemm_gpu/utils/cuda_prelude.cuh"toembedding_backward_split_host_template.cppsokWarpSizeHost()would be in scope. The generatedgen_embedding_backward_split_*.cppare host sources compiled by the plain host compiler, which has no declaration for that header's device intrinsics, so every CUDA build failed on__ballot_syncand__syncwarp. The landed version resolved it by deleting the whole#ifdef USE_ROCMblock instead, which silently moved ROCm fromBT_block_size64 andmax_segment_length_per_warp16384 to the CUDA values of 32 and 32. Both predate the warp-size work;max_segment_length_per_warpis the long-run threshold handed to the unique-index counting kernel, so lowering it to 32 reclassifies most segments as long runs and routes them to thecta_per_rowkernel on ROCm.This restores the block.
max_segment_length_per_warpgoes back to 16384.BT_block_sizebecomeskWarpSizeHost()rather than the previous hardcoded 64, which is what #6278 set out to do for every other launch site and what already landed in the generated CUDA sources: a ROCm wheel serving both wave64 and wave32 architectures compiles this host code once, so a compile-time warp size is wrong for whichever architecture it does not match. On CUDA and on wave64-only ROCm builds the runtime query folds to the same value as before.Calling
kWarpSizeHost()from a host source needs it declared without dragging in device code, sokWarpSizeandkWarpSizeHostmove verbatim fromcuda_prelude.cuhinto a newutils/warp_size.hthatcuda_prelude.cuhincludes. Every existing device-side user keeps seeing them unchanged. UnderUSE_ROCM,at::cuda::warp_size()is declared by<c10/macros/Macros.h>, so the new header needs no CUDA or HIP header of its own;.hfiles are not hipified, and running hipify over the moved block produces no substitutions, so it does not need to be in the hipified source set.The alternative was to guard
cuda_prelude.cuh's device helpers on__CUDACC__/__HIPCC__and keep the original include. That edits a header included by a large amount of device code to fix one host call site, so splitting out the two definitions host code legitimately needs is the narrower change.This is the follow-up @q10 asked for in #6278. The commit is the same change as 60f547b there, now applied on top of main and extended to restore the values the landed diff dropped.
Authored with assistance from Claude (Anthropic).