fix: MNNVL Allreduce uses bitwise sentinel checking to avoid subnormal value issue (#3053) - #3304
Conversation
📝 WalkthroughWalkthroughThis PR makes negative-zero detection in the MNNVL allreduce kernel bit-exact and adds a parametrized regression test that injects specific sentinel trigger bit-patterns into inputs to exercise and guard the Lamport polling logic. ChangesSentinel Negative-Zero Fix and Regression Test
Sequence DiagramsequenceDiagram
participant Test
participant run_mnnvl_ar_full
participant prepare_test_data
participant _inject_sentinel_trigger_patterns
participant AllreduceKernel
Test->>run_mnnvl_ar_full: call with inject_sentinel_patterns=True
run_mnnvl_ar_full->>prepare_test_data: forward inject_sentinel_patterns=True
prepare_test_data->>_inject_sentinel_trigger_patterns: x_full tensor
_inject_sentinel_trigger_patterns->>prepare_test_data: injected int16 patterns
prepare_test_data->>run_mnnvl_ar_full: initialized inputs (broadcast)
run_mnnvl_ar_full->>AllreduceKernel: execute allreduce (kernel uses isNegZero)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the isNegZero function in the MNNVL allreduce kernel to use bit-exact comparisons for floating-point values, preventing deadlocks caused by Flush-To-Zero (FTZ) behavior on newer hardware where subnormal values could be incorrectly identified as the Lamport sentinel. It also introduces a regression test that injects specific bit patterns to verify the sentinel polling logic. Feedback was provided to use the __float_as_uint intrinsic instead of reinterpret_cast for safer and more idiomatic bit-level conversion in CUDA.
| static inline __device__ bool isNegZero(T val) { | ||
| if constexpr (std::is_same_v<T, float>) { | ||
| return val == 0.F && signbit(val); | ||
| return *reinterpret_cast<uint32_t const*>(&val) == kNEGZERO_FP32; |
There was a problem hiding this comment.
Using reinterpret_cast for type punning can lead to undefined behavior due to strict aliasing rules. In CUDA __device__ code, the idiomatic and safe way to perform bit-level conversion from float to uint32_t is using the __float_as_uint intrinsic. This is more readable and ensures the compiler handles the register move correctly without potential memory round-trips.
return __float_as_uint(val) == kNEGZERO_FP32;
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh (1)
122-123: 💤 Low valueOptional: prefer
__float_as_uintfor the fp32 bit-cast.The fp16/bf16 branch right below uses the
Fp16BitCastunion for type punning, while the new fp32 branch uses a pointerreinterpret_cast. The CUDA intrinsic__float_as_uintis the idiomatic device-side bit-cast forfloat → uint32_t, sidesteps strict-aliasing UB on the pointer cast, and avoids forcingvalto memory just to take its address. Functionally equivalent; purely a cleanup.♻️ Proposed cleanup
if constexpr (std::is_same_v<T, float>) { - return *reinterpret_cast<uint32_t const*>(&val) == kNEGZERO_FP32; + return __float_as_uint(val) == kNEGZERO_FP32; } else if constexpr (std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, __nv_half>) {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh` around lines 122 - 123, Replace the pointer-based bitcast in the float branch with the CUDA intrinsic: in the code path guarded by if constexpr (std::is_same_v<T, float>) (the comparison against kNEGZERO_FP32), use __float_as_uint(val) instead of *reinterpret_cast<uint32_t const*>(&val) to perform the float→uint32_t bit-cast; this matches the fp16/bf16 union approach, avoids strict-aliasing/memory-addressing issues, and keeps the check against kNEGZERO_FP32 unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh`:
- Around line 122-123: Replace the pointer-based bitcast in the float branch
with the CUDA intrinsic: in the code path guarded by if constexpr
(std::is_same_v<T, float>) (the comparison against kNEGZERO_FP32), use
__float_as_uint(val) instead of *reinterpret_cast<uint32_t const*>(&val) to
perform the float→uint32_t bit-cast; this matches the fp16/bf16 union approach,
avoids strict-aliasing/memory-addressing issues, and keeps the check against
kNEGZERO_FP32 unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 16b2bf7b-ecf3-455c-aae8-702e24a32a7e
📒 Files selected for processing (2)
include/flashinfer/comm/trtllm_mnnvl_allreduce.cuhtests/comm/test_trtllm_mnnvl_allreduce.py
|
/bot run |
|
@timlee0212 is not authorized to trigger this CI job. cc: @yzh119, @sricketts, @yongwww |
|
/bot run |
|
Verified to resolve #3053. |
5 participants
📌 Description
This PR fixed an inconsistency when polling communication buffer and checking for sentinel value.
The old code does FP comparison with 0 and checks the sign bit to determine whether the polled value is -0.0; The caveat is that when the valid input contains a negative subnormal value, the FP comparison could cause FTZ behavior and flush the valid input into negative zero, which collides with the sentinel value and makes the kernel stuck at polling.
The solution is simple. This PR uses bitwise comparison to check if the incoming value is negative zero, avoiding such subnormal flushing from happening before polling.
Verified the mentioned hang issues reported in SGLang and VLLM can be solved with this fix.
🔍 Related Issues
#3053
also related to the issue vllm-project/vllm#35772
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Bug Fixes
Tests