Repository navigation
[BugFix] Use masked warpsync in in-warp allreduce - #2865
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
📝 WalkthroughWalkthroughThe change updates in-warp ChangesWarp allreduce synchronization
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant TileLangKernel
participant lower_thread_allreduce
participant SharedMemory
participant CUDAWarpSync
TileLangKernel->>lower_thread_allreduce: Lower tvm_thread_allreduce
lower_thread_allreduce->>lower_thread_allreduce: Compute activemask and ballot_sync masks
lower_thread_allreduce->>SharedMemory: Store reduction values
lower_thread_allreduce->>CUDAWarpSync: Emit tl::sync_warp(participant_mask)
CUDAWarpSync-->>SharedMemory: Synchronize participating lanes
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@testing/python/issue/test_tilelang_issue_2628.py`:
- Around line 34-50: Expand test_thread_allreduce_syncwarp into a parameterized
regression matrix covering power-of-two single- and multi-warp thread counts,
every supported reducer combiner, and supported float8 dtypes. Reuse the
existing make_kernel execution and reference validation for each case, and add
explicit rejection coverage for float8 variants when the target does not support
them.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 1ab10aca-30b7-4c14-a397-ea50d1241351
📒 Files selected for processing (2)
src/transform/lower_thread_allreduce.cctesting/python/issue/test_tilelang_issue_2628.py
| def test_thread_allreduce_syncwarp(): | ||
| # Fixed exact-integer inputs so the reported values reproduce byte-for-byte. | ||
| INT_VALS = [10, 25, 40, 30, 15, 42, 20, 20] # sum = 202 (-54 for int8) | ||
| BF16_VALS = [3.0, 8.0, 5.0, 2.0, 7.0, 6.0, 4.0, 6.0] # sum = 41 (exact in bf16) | ||
| for dt, vals in [ | ||
| ("float32", INT_VALS), | ||
| ("int32", INT_VALS), | ||
| ("float16", INT_VALS), | ||
| ("bfloat16", BF16_VALS), | ||
| ("int16", INT_VALS), | ||
| ("int8", INT_VALS), | ||
| ]: | ||
| dtype = getattr(torch, dt) | ||
| A = torch.tensor(vals, device="cuda").to(dtype) | ||
| got = make_kernel(8, dt)(A).to(dtype)[0].item() | ||
| ref = A.sum().to(dtype).item() | ||
| torch.testing.assert_close(got, ref, rtol=1e-3, atol=1e-3) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Expand the regression matrix before merge.
This test executes only an eight-thread + reduction. It does not execute multi-warp reductions, non-additive supported combiners, or float8_* paths. A regression in those required paths can still merge.
Add parameterized power-of-two single- and multi-warp cases for each supported combiner. Test supported float8 variants, or add an explicit rejection test where the target does not support them.
Based on PR objectives, coverage is required for supported reducer operations and single- and multi-warp power-of-two thread counts.
🤖 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 `@testing/python/issue/test_tilelang_issue_2628.py` around lines 34 - 50,
Expand test_thread_allreduce_syncwarp into a parameterized regression matrix
covering power-of-two single- and multi-warp thread counts, every supported
reducer combiner, and supported float8 dtypes. Reuse the existing make_kernel
execution and reference validation for each case, and add explicit rejection
coverage for float8 variants when the target does not support them.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@testing/python/issue/test_tilelang_issue_2628.py`:
- Around line 39-55: Expand test_thread_allreduce_syncwarp_sum to cover 8-, 32-,
and 64-thread power-of-two reductions, creating input values sized for each
thread count and preserving the existing dtype assertions. Extend the dtype
matrix with every supported float8 variant; for unsupported float8 types, add an
explicit test that verifies the documented rejection behavior.
- Around line 35-36: Replace the INT_VALS and BF16_VALS fixtures with values
whose product remains finite and representable for the correctness checks,
avoiding overflow or wrapping in int8, float16, and int16 references. Keep any
intentional overflow coverage in separate dedicated fixtures or cases.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: fa8d4686-cf30-4217-9712-057fdda43e9b
📒 Files selected for processing (1)
testing/python/issue/test_tilelang_issue_2628.py
| INT_VALS = [10, 25, 40, 30, 15, 42, 20, 20] # sum = 202 (-54 for int8) | ||
| BF16_VALS = [3.0, 8.0, 5.0, 2.0, 7.0, 6.0, 4.0, 6.0] # sum = 41 (exact in bf16) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate and inspect the referenced test file and relevant sections.
if [ -f testing/python/issue/test_tilelang_issue_2628.py ]; then
echo "== file exists =="
wc -l testing/python/issue/test_tilelang_issue_2628.py
echo "== lines 1-120 =="
cat -n testing/python/issue/test_tilelang_issue_2628.py | sed -n '1,120p'
else
echo "== missing target file =="
fd -a test_tilelang_issue_2628.py . || true
fi
echo "== related references =="
rg -n "INT_VALS|BF16_VALS|product|float16|int8|reduce|TileLang|tilelang|issue_2628|2628" testing/python/issue -S || trueRepository: tile-ai/tilelang
Length of output: 50372
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
vals = [10, 25, 40, 30, 15, 42, 20, 20]
product = 1
for v in vals:
product *= v
import math
print("product (Python int) =", product)
import struct
def f16_to_float(h):
return math.copysign(2 ** (struct.unpack(">B", struct.pack(">H", h & 0x7fff) >> 10 & 0x1f)[0] - 15), struct.pack(">H", h & 0x8000))
bf16_vals_float = []
for v in vals:
as32 = struct.unpack(">I", struct.pack(">f", float(v)))[0]
bf16 = (as32 >> 16) & 0xffff
bf16_vals_float.append(f16_to_float(bf16))
prod_f16 = 1.0
for v in vals:
prod_f16 *= v
print("bf16 values approximated via Python:", bf16_vals_float)
print("product of bf16 approximations:", prod_f16)
print("is product of bf16 approximations finite:", math.isfinite(prod_f16))
print("INT_VALS sum:", sum(vals))
PYRepository: tile-ai/tilelang
Length of output: 371
🌐 Web query:
bfloat16 float8 format Python struct f16 conversion
💡 Result:
Bfloat16 and float8 are specialized floating-point formats used primarily in machine learning to optimize memory and computational efficiency [1]. As of August 2026, Python's built-in struct module does not provide native support for either format [2][3]. ### Format Overview - Bfloat16 (Brain Floating Point 16): A 16-bit format with 1 sign bit, 8 exponent bits, and 7 mantissa bits [4][1]. Its exponent range matches IEEE 754 single precision (FP32), making it stable for deep learning, though it has lower precision [4][5]. - Float8: A family of 8-bit formats, most commonly E4M3 (4 exponent, 3 mantissa) and E5M2 (5 exponent, 2 mantissa) [6][1]. These provide significant memory bandwidth savings but require careful scaling strategies to manage their narrower dynamic range [1][7]. ### Python struct Limitations The struct module supports standard IEEE 754 formats via specific conversion codes: - e: IEEE 754 binary16 (half-precision) [2][8]. - f: IEEE 754 binary32 (single-precision) [2][8]. - d: IEEE 754 binary64 (double-precision) [2][8]. There is currently no corresponding format character for bfloat16 or float8 in the struct module [2][3]. Conversion between these formats and standard Python types typically requires specialized libraries rather than the standard library's struct module [9][10]. ### How to Handle Conversions Because struct does not support these formats, you must use dedicated tools designed for numerical computing: - For PyTorch users: High-level utilities such as torch.Tensor.to(torch.bfloat16) or specialized modules like torchao.float8 are recommended for model training and inference [11][12][10]. - For numerical conversion: Libraries like ml_dtypes extend NumPy to support bfloat16, allowing you to interface between PyTorch and NumPy more easily [9]. - For low-level binary data: If you need to manually pack or unpack binary data that is not supported by struct, you may need to use bitwise manipulation or third-party packages designed for exotic floating-point representations [13]. Top results: [1][2][3]
Citations:
- 1: https://www.emergentmind.com/topics/fp8-and-bf16-representations
- 2: https://docs.python.org/3/library/struct.html
- 3: https://discuss.python.org/t/add-bfloat16-support-to-struct/106723
- 4: https://acecloud.ai/blog/fp8-vs-bf16-mixed-precision-tensor-cores/
- 5: https://clay-atlas.com/us/blog/2024/09/25/en-deep-learning-float16-bfloat16-fp8/
- 6: https://www.opencompute.org/documents/ocp-8-bit-floating-point-specification-ofp8-revision-1-0-2023-12-01-pdf-1
- 7: https://github.com/eth-easl/torchtitan-mixtera/blob/main/docs/float8.md
- 8: https://github.com/python/cpython/blob/main/Doc/library/struct.rst
- 9: https://stackoverflow.com/questions/78128662/converting-pytorch-bfloat16-tensors-to-numpy-throws-typeerror
- 10: https://runebook.dev/en/docs/pytorch/storage/torch.UntypedStorage.float8_e4m3fnuz
- 11: https://docs.pytorch.org/ao/main/api%5Freference/generated/torchao.float8.convert_to_float8_training.html
- 12: https://docs.pytorch.org/ao/stable/api%5Freference/generated/torchao.float8.convert%5Fto%5Ffloat8%5Ftraining.html
- 13: https://bitstring.readthedocs.io/en/stable/exotic%5Ffloats.html
Use a non-overflowing product fixture.
INT_VALS multiplies to 75_600_000_000. The int8 reference wraps to 0, and the float16/int16 product reference overflows. A correct product reduction can already match these references, so it can pass without catching the synchronization issue. Use a finite, non-overflow product fixture for correctness cases, and keep overflow cases separate.
Applies to lines 35-36 also.
🤖 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 `@testing/python/issue/test_tilelang_issue_2628.py` around lines 35 - 36,
Replace the INT_VALS and BF16_VALS fixtures with values whose product remains
finite and representable for the correctness checks, avoiding overflow or
wrapping in int8, float16, and int16 references. Keep any intentional overflow
coverage in separate dedicated fixtures or cases.
| @pytest.mark.parametrize( | ||
| "dt,vals", | ||
| [ | ||
| ("float32", INT_VALS), | ||
| ("int32", INT_VALS), | ||
| ("float16", INT_VALS), | ||
| ("bfloat16", BF16_VALS), | ||
| ("int16", INT_VALS), | ||
| ("int8", INT_VALS), | ||
| ], | ||
| ) | ||
| def test_thread_allreduce_syncwarp_sum(dt, vals): | ||
| dtype = getattr(torch, dt) | ||
| A = torch.tensor(vals, device="cuda").to(dtype) | ||
| got = make_kernel(8, dt, lambda x, y: x + y, 0)(A).to(dtype)[0].item() | ||
| ref = A.sum().to(dtype).item() | ||
| torch.testing.assert_close(got, ref, rtol=1e-3, atol=1e-3) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Expand the regression matrix before merge.
The previous coverage gap remains partly unresolved. Lines 53 and 72 test only an eight-thread reduction. Add full-warp and multi-warp power-of-two cases, such as 32 and 64 threads, with inputs sized for each reduction.
The matrix also has no float8_* case or explicit unsupported-type rejection. Add coverage for each supported variant, or assert the documented rejection.
Also applies to: 58-74
🤖 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 `@testing/python/issue/test_tilelang_issue_2628.py` around lines 39 - 55,
Expand test_thread_allreduce_syncwarp_sum to cover 8-, 32-, and 64-thread
power-of-two reductions, creating input values sized for each thread count and
preserving the existing dtype assertions. Extend the dtype matrix with every
supported float8 variant; for unsupported float8 types, add an explicit test
that verifies the documented rejection behavior.
#2865 fixed the warp-synchronous allreduce corruption from #2628 by emitting `__syncwarp(mask)` around the in-warp reduction steps. On HIP that call lowers to nothing, so the fix never took effect on ROCm and `test_tilelang_issue_2628.py` fails on gfx942: test_thread_allreduce_syncwarp_sum[float32-vals0] E AssertionError: Scalars are not close! E Expected 202.0 but got 90.0. The existing comment argues that AMD wavefronts execute in lockstep, so no barrier is needed. That is true of *runtime* convergence but not of compile-time scheduling: with nothing marking the synchronization point, the backend is free to reorder or coalesce the LDS accesses that the warp-synchronous reduction depends on, which is exactly the corruption #2628 describes. Emit `__builtin_amdgcn_wave_barrier()` instead. It generates no instructions, so the lockstep reasoning still holds and there is no runtime cost; it only constrains the scheduler. The mask argument remains ignored, since a wavefront is always fully convergent. Verified on 8x MI300X (gfx942), ROCm 7.2, torch 2.14.0.dev+rocm7.2: * before: `test_tilelang_issue_2628.py` -> 3 failed (run aborted by --maxfail=3) * after: `test_tilelang_issue_2628.py` -> 12 passed (sum and mul across float32/int32/float16/bfloat16/int16/int8) * full ROCm suite: 1809 passed, 1227 skipped, 0 failed
The in-warp reduction loop separated each load/store pair with
SyncThread("warp")->tvm_storage_sync("warp"), which is a no-op in CUDA.This leads to datarace and silently produces incorrect results.
Note that simply emitting
__syncthreads()does not help as there might be a surrounding divergentifcondition guarding the sync.To fix this, we compute the participant mask via
__ballot_sync()before entering the branch, and emit__syncwarp(participant_mask)(tl::sync_warp) between stages.Fixes #2628
Summary
tvm_thread_allreduceresults caused by unsynchronized shared-memory stages.tl::sync_warpcalls.__activemask()and__ballot_sync().bfloat16reductions.C++ style / lint notes
docs/developer_guide/cpp_style.mdor introduce a new public API.TLCPP003orTLCPP004findings should not block this change without a clear API, FFI, or maintainability risk.