Repository navigation
fix: inference norm missing threadsPerGroup==32 dispatch + ROCm 7.2.2 bf16 ABI - #8906
umechand-amd wants to merge 2 commits into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c83fe5684a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| LAUNCH_FUSED_LN(1, 8, maxThreads); | ||
| } else if (threadsPerGroup == 16) { | ||
| LAUNCH_FUSED_LN(1, 16, maxThreads); | ||
| } else if (threadsPerGroup == 32) { |
There was a problem hiding this comment.
For supported float32 widths such as 80 or 128 through inference v2, InferenceCoreBuilder compiles the separate deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cu and cuda_rms_norm/rms_norm_cuda.cu copies, whose corresponding four subblock dispatches still stop at 16. Those APIs therefore still take the 32-thread case without launching a kernel and return untouched output; mirror this branch into the v2 copies rather than fixing only InferenceBuilder.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
| #include <stdint.h> | ||
|
|
||
| #ifdef BF16_AVAILABLE | ||
| #if defined(BF16_AVAILABLE) && !defined(__HIP_PLATFORM_AMD__) |
There was a problem hiding this comment.
On a ROCm build of an op that leaves enable_bf16 enabled, such as QuantizerBuilder or FPQuantizerBuilder, hipcc itself defines __HIP_PLATFORM_AMD__, so this condition removes the bfloat header from device translation units too. conversion_utils.h nevertheless compiles its BF16_AVAILABLE specializations using __nv_bfloat16, __nv_bfloat162, and their intrinsics; ds_kernel_utils.h only forward-declares the scalar type, and these builders include conversion/reduction before any bfloat definition, so they now fail to compile. Distinguish host-only g++ from hipcc or provide the definition before this header.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
| #include <stdint.h> | ||
|
|
||
| #ifdef BF16_AVAILABLE | ||
| #if defined(BF16_AVAILABLE) && !defined(__HIP_PLATFORM_AMD__) |
There was a problem hiding this comment.
On the stated ROCm 7.2.2 InferenceCoreBuilder build, this edit is not on the include path: op_builder/inference_core_ops.py selects deepspeed/inference/v2/kernels/includes, where both ds_kernel_utils.h and conversion_utils.h still include cuda_bf16.h whenever BF16_AVAILABLE is set. The host .cpp files therefore still expose g++ to amd_hip_bf16.h and the extension continues to fail; apply the guard and forward declaration to the v2 headers actually used by this builder.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
… bf16 ABI
Three dispatch tables in layer_norm.cu and one in rms_norm.cu handle
threadsPerGroup in {1,2,4,8,16} only. For float32 with elems_per_row
in [65,128], T_per_load=4 gives threadsPerGroup=32, which falls through
every branch -- no kernel launches and output holds uninitialized memory.
Also guards cuda_bf16.h in conversion_utils.h with !defined(__HIP_PLATFORM_AMD__)
to match the existing pattern in ds_kernel_utils.h. Without the guard,
plain c++ compiled with -DBF16_AVAILABLE -D__HIP_PLATFORM_AMD__ pulls in
amd_hip_bf16.h which requires Clang-only builtins (__builtin_elementwise_fma
etc.) unavailable under g++, breaking inference_core_ops on ROCm 7.2.2.
Signed-off-by: Umesh Chand <Umesh.Chand@amd.com>
…6 guard, add fp32 small-hidden tests Signed-off-by: Umesh Chand <Umesh.Chand@amd.com>
f27105e to
ac23d4c
Compare
What
Two independent fixes to the v1 inference norm kernels.
1. Missing
threadsPerGroup==32dispatch (layer_norm.cu, rms_norm.cu)launch_fused_ln,launch_fused_residual_ln,launch_fused_residual_ln_store_pre_ln_res, andlaunch_rms_normall dispatch onthreadsPerGroupwith branches for {1, 2, 4, 8, 16} only. Forfloat32inputs withelems_per_rowin [65, 128],T_per_load=4givesthreadsPerGroup=32, which falls through every branch — no kernel is launched and the output tensor holds uninitialized memory (silent numerical corruption).Reproduced on MI350X (gfx950, ROCm 7.2.2): errors of 3–8 units for H=80–128 with
dtype=float32.Fix: add
else if (threadsPerGroup == 32)to all four dispatch tables.Real-world trigger: per-head Q/K RMSNorm in EXAONE 4/4.5 with head_size 64 or 128 (related to #8886).
2. ROCm 7.2.2 bf16 ABI build failure (conversion_utils.h)
conversion_utils.hincludescuda_bf16.hguarded only byBF16_AVAILABLE. Hipify translates this tohip_bf16.hinconversion_utils_hip.h. Whenop_buildercompiles.cppfiles with plaing++and-DBF16_AVAILABLE -D__HIP_PLATFORM_AMD__, it pulls inamd_hip_bf16.h, which on ROCm 7.2.2 requires Clang-only builtins (__builtin_elementwise_fma,__builtin_elementwise_ceil, etc.) and a__bf16vector ABI not available underg++, causing the entireinference_core_opsextension to fail to build.Fix: add
&& !defined(__HIP_PLATFORM_AMD__)to match the identical pattern already inds_kernel_utils.h(which has a comment explaining exactly this problem). Hipify will propagate the guard correctly to the generated_hip.hfile.Test
Repro script for fix 1: verified H=80–128 (float32) produce zero error after the fix on gfx950.
🤖 Generated with Claude Code