Conversation
Host-side TBE launch configurations computed block/grid dims from the compile-time kWarpSize (or a hardcoded 64 for BT_block_size). In a ROCm wheel that serves both wave64 (CDNA) and wave32 (RDNA) archs, the host code is compiled once, so a compile-time warp size is wrong for whichever arch it does not match: on gfx1100 the launches came up with 64-wide x-dims for 32-wide warps. This replaces those host uses with the runtime kWarpSizeHost(), which reports the active device's warp size. On CUDA and on wave64-only ROCm builds kWarpSizeHost() folds to the same values as before, so this is behavior-preserving there. Split out of pytorch#6123 at reviewer request to keep each PR small and independently revertible. The codegen/template reparameterisation that makes the kernel symbols themselves wave-size-agnostic stays in pytorch#6123, which is rebased on top of this. Test Plan: Rendered the TBE codegen for CUDA, ROCm wave64, ROCm wave32, and ROCm wave32+wave64 from this commit and from origin/main, and diffed: the only changes are the intended kWarpSize -> kWarpSizeHost() substitutions in host launch code. Also compiled fbgemm_gpu for gfx90a with the follow-up commit applied (see pytorch#6123). cd fbgemm_gpu/codegen/genscript for flags in "" "--is_rocm --has_wave64" "--is_rocm --has_wave32" \ "--is_rocm --has_wave32 --has_wave64"; do python generate_backward_split.py --opensource $flags python generate_forward_split.py --opensource $flags python generate_forward_quantized.py --opensource $flags python generate_embedding_optimizer.py --opensource $flags python generate_index_select.py --opensource $flags done Authored with assistance from Claude (Anthropic).
|
@q10 has imported this pull request. If you are a Meta employee, you can view this in D119363782. |
The previous commit added #include "fbgemm_gpu/utils/cuda_prelude.cuh" to embedding_backward_split_host_template.cpp so that kWarpSizeHost() would be in scope. That broke every CUDA build: the generated gen_embedding_backward_split_*.cpp files are host sources compiled by the plain host compiler, which has no declaration for the device intrinsics that cuda_prelude.cuh uses, so it fails with "use of undeclared identifier '__ballot_sync'" and "use of undeclared identifier '__syncwarp'". ROCm builds did not see this because their host pass is hipcc, which does declare them. The CPU-only build variant would have failed as well, since cuda_prelude.cuh pulls in <cuda.h>. kWarpSize and kWarpSizeHost move verbatim into a new utils/warp_size.h, which cuda_prelude.cuh includes, so every existing device-side user keeps seeing them unchanged. The host template includes the new header instead, and the kWarpSizeHost() call site stays exactly as it was. 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 and is safe to leave out of the hipify set (.h files are not hipified; running hipify over this block produces no substitutions). The alternative was to keep the include and guard cuda_prelude.cuh's device helpers on __CUDACC__/__HIPCC__. That edits a header included by a large amount of device code and by Meta-internal builds, to fix one host call site, so splitting out the two definitions that host code legitimately needs is the narrower change. Re-implementing the ROCm branch of kWarpSizeHost() at the call site was also considered and rejected: it duplicates the definition and reintroduces a hard-coded host-side warp size, which is what this stack exists to remove. Test Plan: Compiled the new header standalone with the plain host compiler, with and without USE_ROCM, against installed torch headers only (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, so it does not need to be in the hipified source set: ``` 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 no generated host source includes cuda_prelude.cuh, while the kWarpSizeHost() call sites are unchanged: ``` 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 'kWarpSizeHost' *.cpp # 14 matches ``` Full CUDA and ROCm builds are left to CI. Authored with assistance from Claude (Anthropic).
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
|
@q10 heads up: the CUDA builds on this PR went red shortly after the import (D119363782), so the imported diff has the breakage in it. Root cause: this PR added ROCm builds did not catch it because their host pass is hipcc, which does declare them. The CPU-only build variant would have failed as well, since Fix pushed as 60f547b: One thing worth a check on your side: this adds a new file under Could you re-review and re-import when you get a chance? The GitHub workflows are currently sitting in Separately, #6123 (the codegen half of the split) is now rebased onto current main and stacked on this one, and carries the same fix, so it will need a re-import too. Authored with assistance from Claude (Anthropic). |
|
Hi @jeffdaily I think we missed the last commit 60f547b you pushed to the PR when we landed the dif (we had addressed the build issues internally, but they aren't the same code changes as in 60f547b. Could you open 60f547b as a new PR on top of latest main, if the changes are still applicable? |
|
Mostly applicable, with one correction worth flagging. The header split on its own would have been a no-op on current main: your version removed the What is still needed is that the landed diff deleted the whole The other four files landed as a behavior-equivalent refactor, so #6323 does not touch them. Its workflows are sitting in Authored with assistance from Claude (Anthropic). |
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
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
2 participants
Host-side TBE launch configurations computed block/grid dims from the compile-time
kWarpSize(or a hardcoded 64 forBT_block_size). In a ROCm wheel that serves both wave64 (CDNA) and wave32 (RDNA) archs the host code is compiled once, so a compile-time warp size is wrong for whichever arch it does not match: on gfx1100 these launches came up with 64-wide x-dims for 32-wide warps. This replaces those host uses with the runtimekWarpSizeHost(), which reports the active device's warp size. On CUDA and on wave64-only ROCm buildskWarpSizeHost()folds to the same values as before, so this is behavior-preserving there.Split out of #6123 at reviewer request to keep each PR small and independently revertible. The codegen/template reparameterisation that makes the kernel symbols wave-size-agnostic stays in #6123, which is now rebased on top of this PR.
Authored with assistance from Claude (Anthropic).