Conversation
…estore (…" This reverts commit 68ebdbd.
📝 WalkthroughWalkthroughThe changes remove TRT-LLM checkpoint lifecycle APIs and related validation, refactor ChangesSymmetric-memory workspace changes
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
![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 refactors the communication backend by removing the checkpoint/restore lifecycle methods from the TRT-LLM and MNNVL all-reduce workspaces, transitioning instead to torch symmetric memory allocation. Feedback on these changes highlights several critical issues: first, removing torch.cuda.synchronize() before the barrier in the TRT-LLM workspace creation could lead to race conditions during asynchronous initialization; second, performing multicast checks and queries unconditionally in mnnvl.py will cause runtime failures on devices without multicast support when multicast is disabled; and finally, the type annotations for _symm_workspace_refs and symm_refs in trtllm_ar.py are incorrect as they store tuples of tensors and handles rather than just tensors.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if use_symm_dev_mem: | ||
| torch.cuda.synchronize() | ||
| comm_backend.barrier() # must sync after create_workspace |
There was a problem hiding this comment.
The asynchronous CUDA kernel trtllm_lamport_initialize is launched to initialize the lamport buffer on the GPU. Without calling torch.cuda.synchronize() before the host-side barrier, other ranks might proceed and attempt to read from this rank's lamport buffer before the initialization kernel has finished executing on the GPU, leading to race conditions or hangs. We should synchronize the CUDA stream before calling the barrier.
| if use_symm_dev_mem: | |
| torch.cuda.synchronize() | |
| comm_backend.barrier() # must sync after create_workspace | |
| if use_symm_dev_mem: | |
| torch.cuda.synchronize() | |
| comm_backend.barrier() # must sync after create_workspace |
| multicast_supported = checkCudaErrors( | ||
| cuda.cuDeviceGetAttribute( | ||
| cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, | ||
| device_idx, | ||
| ) | ||
| ) | ||
| if multicast_supported == 0: | ||
| raise RuntimeError( | ||
| "[SymmDeviceMemory] Device does not support multicasting." | ||
| ) |
There was a problem hiding this comment.
The multicast support check is now performed unconditionally. If enable_multicast is passed as False (for example, on devices that do not support multicasting), this will raise a RuntimeError and prevent SymmDeviceMemory from being initialized at all. We should wrap this check in if enable_multicast: and also store self._enable_multicast = enable_multicast so that helper methods can check it.
| multicast_supported = checkCudaErrors( | |
| cuda.cuDeviceGetAttribute( | |
| cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, | |
| device_idx, | |
| ) | |
| ) | |
| if multicast_supported == 0: | |
| raise RuntimeError( | |
| "[SymmDeviceMemory] Device does not support multicasting." | |
| ) | |
| self._enable_multicast = enable_multicast | |
| # Check if device supports multicasting | |
| if enable_multicast: | |
| multicast_supported = checkCudaErrors( | |
| cuda.cuDeviceGetAttribute( | |
| cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, | |
| device_idx, | |
| ) | |
| ) | |
| if multicast_supported == 0: | |
| raise RuntimeError( | |
| "[SymmDeviceMemory] Device does not support multicasting." | |
| ) |
| # Set up multicast properties | ||
| mc_prop = cuda.CUmulticastObjectProp() | ||
| mc_prop.numDevices = self.group_size | ||
| mc_prop.size = self.allocation_size | ||
| mc_prop.handleTypes = self._exchanger.handle_type | ||
|
|
||
| # Get multicast granularity and adjust allocation size | ||
| self._mc_granularity = checkCudaErrors( | ||
| cuda.cuMulticastGetGranularity( | ||
| mc_prop, | ||
| cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, | ||
| ) | ||
| self.allocation_size = round_up(self.allocation_size, self._mc_granularity) | ||
| ) | ||
| self.allocation_size = round_up(self.allocation_size, self._mc_granularity) |
There was a problem hiding this comment.
In _get_allocation_prop, the multicast properties are set up and cuMulticastGetGranularity is called unconditionally. If enable_multicast is False (and multicast is not supported by the device), these calls will fail. We should wrap the multicast setup and granularity query in if self._enable_multicast: to avoid executing them when multicast is disabled.
| # Set up multicast properties | |
| mc_prop = cuda.CUmulticastObjectProp() | |
| mc_prop.numDevices = self.group_size | |
| mc_prop.size = self.allocation_size | |
| mc_prop.handleTypes = self._exchanger.handle_type | |
| # Get multicast granularity and adjust allocation size | |
| self._mc_granularity = checkCudaErrors( | |
| cuda.cuMulticastGetGranularity( | |
| mc_prop, | |
| cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, | |
| ) | |
| self.allocation_size = round_up(self.allocation_size, self._mc_granularity) | |
| ) | |
| self.allocation_size = round_up(self.allocation_size, self._mc_granularity) | |
| self._mc_granularity = alloc_granularity | |
| mc_prop = None | |
| if self._enable_multicast: | |
| # Set up multicast properties | |
| mc_prop = cuda.CUmulticastObjectProp() | |
| mc_prop.numDevices = self.group_size | |
| mc_prop.size = self.allocation_size | |
| mc_prop.handleTypes = self._exchanger.handle_type | |
| # Get multicast granularity and adjust allocation size | |
| self._mc_granularity = checkCudaErrors( | |
| cuda.cuMulticastGetGranularity( | |
| mc_prop, | |
| cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, | |
| ) | |
| ) | |
| self.allocation_size = round_up(self.allocation_size, self._mc_granularity) |
| LamportTokenNumThreshold = 16 | ||
|
|
||
| _symm_workspace_refs: dict[int, list[object]] = {} | ||
| _symm_workspace_refs: dict[int, list[torch.Tensor]] = {} |
There was a problem hiding this comment.
The type annotation for _symm_workspace_refs is updated to dict[int, list[torch.Tensor]]. However, the list actually stores tuples of (tensor, handle) where handle is the symmetric memory handle. This mismatch will cause static type checking failures. It should be annotated as dict[int, list[tuple[torch.Tensor, object]]].
| _symm_workspace_refs: dict[int, list[torch.Tensor]] = {} | |
| _symm_workspace_refs: dict[int, list[tuple[torch.Tensor, object]]] = {} |
| if group is not None | ||
| else torch.distributed.group.WORLD.group_name | ||
| ) | ||
| symm_refs: list[torch.Tensor] = [] |
There was a problem hiding this comment.
The type annotation for symm_refs is updated to list[torch.Tensor]. However, the list actually stores tuples of (tensor, handle) returned by _alloc_symm_buffer_bytes. This mismatch will cause static type checking failures. It should be annotated as list[tuple[torch.Tensor, object]].
| symm_refs: list[torch.Tensor] = [] | |
| symm_refs: list[tuple[torch.Tensor, object]] = [] |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/comm/trtllm_ar.py (1)
617-733: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winInitialize the symmetric flag buffer before use
flag_sizeis still allocated withsymm_mem.empty(...)and never cleared here, so the Lamport/barrier flags can start with garbage. Zero that buffer before publishing the workspace; stale flag state can deadlock peers.🤖 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 `@flashinfer/comm/trtllm_ar.py` around lines 617 - 733, In the workspace initialization flow, explicitly zero the symmetric flag buffer allocated for flag_size before publishing workspace_tensor or synchronizing peers. Use the flag buffer tensor returned by _alloc_symm_buffer_bytes for the flag allocation, preserving the existing allocation order and barrier behavior.
🧹 Nitpick comments (2)
flashinfer/comm/trtllm_ar.py (2)
19-19: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType hints for
symm_refs/_symm_workspace_refsdon't match stored tuples.
_symm_workspace_refs: dict[int, list[torch.Tensor]](line 434) andsymm_refs: list[torch.Tensor](line 647) are declared to holdtorch.Tensor, butsymm_refs.append((tensor, handle))(line 669) stores 2-tuples. The line-range details note this narrowed a previously more permissivelist[object]type. Worth correcting for accuracy/static-typing tooling.♻️ Proposed fix
-_symm_workspace_refs: dict[int, list[torch.Tensor]] = {} +_symm_workspace_refs: dict[int, list[tuple[torch.Tensor, Any]]] = {} ... - symm_refs: list[torch.Tensor] = [] + symm_refs: list[tuple[torch.Tensor, Any]] = []Also applies to: 434-434, 642-647
🤖 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 `@flashinfer/comm/trtllm_ar.py` at line 19, Update the type annotations for symm_refs and _symm_workspace_refs to represent their actual contents: two-element tuples containing the tensor and handle values appended in the symmetric workspace flow. Ensure the local declaration and persistent dictionary annotation are consistent with the tuple shape used by symm_refs.append in the surrounding method.
571-596: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueAlign
mem_handleswith the actual handle type
_alloc_symm_buffer_bytesreturns the symmetric-memory rendezvous handle, so the return annotation and docstring should describe that type instead ofList[SymmDeviceMemory].🤖 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 `@flashinfer/comm/trtllm_ar.py` around lines 571 - 596, Update the return annotation and docstring for the surrounding allocation function to describe mem_handles as the symmetric-memory rendezvous handle type returned by _alloc_symm_buffer_bytes, replacing List[SymmDeviceMemory]. Keep the existing tuple variants and other return descriptions 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.
Inline comments:
In `@flashinfer/comm/mnnvl.py`:
- Around line 1023-1033: Gate the multicast support check and all
multicast-specific allocation setup in the constructor around enable_multicast,
including cuDeviceGetAttribute and cuMulticastGetGranularity calls. When
enable_multicast is false, skip these operations entirely and preserve a
unicast-only initialization path that does not require multicast-capable
devices.
In `@flashinfer/comm/trtllm_mnnvl_ar.py`:
- Around line 181-183: Update the Lamport initialization in the relevant
workspace setup method to use the reduction dtype when writing the sentinel,
matching the existing cuMemsetD16/cuMemsetD32 behavior in mnnvl.py. Ensure fp16
and bf16 workspaces receive a uniform 0x8000 pattern while fp32 retains the
negative-zero initialization, and preserve the existing synchronization step.
---
Outside diff comments:
In `@flashinfer/comm/trtllm_ar.py`:
- Around line 617-733: In the workspace initialization flow, explicitly zero the
symmetric flag buffer allocated for flag_size before publishing workspace_tensor
or synchronizing peers. Use the flag buffer tensor returned by
_alloc_symm_buffer_bytes for the flag allocation, preserving the existing
allocation order and barrier behavior.
---
Nitpick comments:
In `@flashinfer/comm/trtllm_ar.py`:
- Line 19: Update the type annotations for symm_refs and _symm_workspace_refs to
represent their actual contents: two-element tuples containing the tensor and
handle values appended in the symmetric workspace flow. Ensure the local
declaration and persistent dictionary annotation are consistent with the tuple
shape used by symm_refs.append in the surrounding method.
- Around line 571-596: Update the return annotation and docstring for the
surrounding allocation function to describe mem_handles as the symmetric-memory
rendezvous handle type returned by _alloc_symm_buffer_bytes, replacing
List[SymmDeviceMemory]. Keep the existing tuple variants and other return
descriptions unchanged.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 864a4680-a31a-410a-a752-caf00814b99f
📒 Files selected for processing (7)
docs/api/comm.rstflashinfer/comm/allreduce.pyflashinfer/comm/mnnvl.pyflashinfer/comm/trtllm_ar.pyflashinfer/comm/trtllm_mnnvl_ar.pytests/comm/test_trtllm_allreduce_checkpoint.pytests/comm/test_trtllm_mnnvl_allreduce.py
💤 Files with no reviewable changes (3)
- docs/api/comm.rst
- tests/comm/test_trtllm_allreduce_checkpoint.py
- flashinfer/comm/allreduce.py
| # Check if device supports multicasting | ||
| if self._enable_multicast: | ||
| multicast_supported = checkCudaErrors( | ||
| cuda.cuDeviceGetAttribute( | ||
| cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, | ||
| device_idx, | ||
| ) | ||
| multicast_supported = checkCudaErrors( | ||
| cuda.cuDeviceGetAttribute( | ||
| cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, | ||
| device_idx, | ||
| ) | ||
| ) | ||
| if multicast_supported == 0: | ||
| raise RuntimeError( | ||
| "[SymmDeviceMemory] Device does not support multicasting." | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the file structure first, then inspect the relevant ranges.
ast-grep outline flashinfer/comm/mnnvl.py --view expanded || true
echo
echo '--- lines 980-1085 ---'
sed -n '980,1085p' flashinfer/comm/mnnvl.py | cat -n
echo
echo '--- lines 1210-1265 ---'
sed -n '1210,1265p' flashinfer/comm/mnnvl.py | cat -n
echo
echo '--- search for enable_multicast and multicast-specific APIs ---'
rg -n "enable_multicast|cuMulticast|multicast_supported|_mc_granularity|_allocate_unicast_buffers|_get_allocation_prop" flashinfer/comm/mnnvl.pyRepository: flashinfer-ai/flashinfer
Length of output: 14662
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline flashinfer/comm/mnnvl.py --view expanded || true
echo
echo '--- lines 980-1085 ---'
sed -n '980,1085p' flashinfer/comm/mnnvl.py | cat -n
echo
echo '--- lines 1210-1265 ---'
sed -n '1210,1265p' flashinfer/comm/mnnvl.py | cat -n
echo
echo '--- search for enable_multicast and multicast-specific APIs ---'
rg -n "enable_multicast|cuMulticast|multicast_supported|_mc_granularity|_allocate_unicast_buffers|_get_allocation_prop" flashinfer/comm/mnnvl.pyRepository: flashinfer-ai/flashinfer
Length of output: 14662
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1000,1055p' flashinfer/comm/mnnvl.py | cat -n
echo '---'
sed -n '1218,1258p' flashinfer/comm/mnnvl.py | cat -nRepository: flashinfer-ai/flashinfer
Length of output: 4608
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- lines 1178-1368 ---'
sed -n '1178,1368p' flashinfer/comm/mnnvl.py | cat -n
echo
echo '--- lines 1368-1392 ---'
sed -n '1368,1392p' flashinfer/comm/mnnvl.py | cat -nRepository: flashinfer-ai/flashinfer
Length of output: 10067
Gate multicast-only setup on enable_multicast
enable_multicast=False still hits cuDeviceGetAttribute(CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED) and cuMulticastGetGranularity(), so construction fails on devices without multicast support instead of falling back to unicast-only mode. Move the support check and multicast allocation setup behind the flag, or split the unicast path entirely.
🤖 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 `@flashinfer/comm/mnnvl.py` around lines 1023 - 1033, Gate the multicast
support check and all multicast-specific allocation setup in the constructor
around enable_multicast, including cuDeviceGetAttribute and
cuMulticastGetGranularity calls. When enable_multicast is false, skip these
operations entirely and preserve a unicast-only initialization path that does
not require multicast-capable devices.
| # lamport initialize tensor to negative zero. | ||
| self.tensor.fill_(-0.0) | ||
| # Wait until the initialization is done |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== outline trtllm_mnnvl_ar.py ==\n'
ast-grep outline flashinfer/comm/trtllm_mnnvl_ar.py --view expanded || true
printf '\n== outline mnnvl.py ==\n'
ast-grep outline flashinfer/comm/mnnvl.py --view expanded || true
printf '\n== relevant slices trtllm_mnnvl_ar.py ==\n'
sed -n '130,230p' flashinfer/comm/trtllm_mnnvl_ar.py | cat -n
printf '\n== relevant slices mnnvl.py ==\n'
sed -n '1,240p' flashinfer/comm/mnnvl.py | cat -nRepository: flashinfer-ai/flashinfer
Length of output: 20722
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import struct
import numpy as np
# float32 -0.0 bytes
f32 = np.array([-0.0], dtype=np.float32).tobytes()
print("float32 -0.0 bytes:", f32.hex())
# reinterpret as two fp16 halfwords
h0 = np.frombuffer(f32, dtype=np.float16)
print("reinterpret as fp16:", [hex(x.view(np.uint16)) for x in h0], h0.tolist())
# bfloat16 isn't native in numpy; emulate by taking top 16 bits of float32 words
u32 = int.from_bytes(f32, "little")
print("float32 word:", hex(u32))
print("low16/high16:", hex(u32 & 0xffff), hex((u32 >> 16) & 0xffff))
# If a 4-byte -0.0 is stored and read as bf16 elements from each 16-bit lane,
# lane patterns are the two halves of the 32-bit word.
PYRepository: flashinfer-ai/flashinfer
Length of output: 292
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '150,210p' flashinfer/comm/trtllm_mnnvl_ar.py | cat -n
printf '\n---\n'
sed -n '1,220p' flashinfer/comm/mnnvl.py | cat -n | sed -n '1,220p'Repository: flashinfer-ai/flashinfer
Length of output: 11118
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== lamport_initialize in mnnvl.py ==\n'
sed -n '1360,1455p' flashinfer/comm/mnnvl.py | cat -n
printf '\n== search for tensor/buffer dtype usage in trtllm_mnnvl_ar.py ==\n'
rg -n "self\.tensor|dtype|buffer_size_bytes|lamport" flashinfer/comm/trtllm_mnnvl_ar.py
printf '\n== search for MNNVLAllReduceFusionWorkspace instantiation ==\n'
rg -n "MNNVLAllReduceFusionWorkspace\\(" -S .Repository: flashinfer-ai/flashinfer
Length of output: 11351
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/mnnvl.py")
lines = p.read_text().splitlines()
for i in range(1360, 1456):
if i <= len(lines):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: flashinfer-ai/flashinfer
Length of output: 4379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "self\.tensor|lamport_initialize|buffer_flags|MNNVLAllReduceFusionWorkspace" flashinfer/comm/trtllm_mnnvl_ar.pyRepository: flashinfer-ai/flashinfer
Length of output: 1915
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '62,130p' flashinfer/comm/trtllm_mnnvl_ar.py | cat -nRepository: flashinfer-ai/flashinfer
Length of output: 4497
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "_alloc_symm_buffer_bytes|create_tensor_from_cuda_memory|view\\(|fill_\\(" flashinfer/comm -g '*.py'
printf '\n== _alloc_symm_buffer_bytes definition ==\n'
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/trtllm_mnnvl_ar.py")
lines = p.read_text().splitlines()
for i, line in enumerate(lines, 1):
if "_alloc_symm_buffer_bytes" in line:
start = max(1, i-25)
end = min(len(lines), i+80)
for j in range(start, end+1):
print(f"{j:4d}: {lines[j-1]}")
break
PYRepository: flashinfer-ai/flashinfer
Length of output: 7061
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/trtllm_mnnvl_ar.py")
lines = p.read_text().splitlines()
for needle in ["def _alloc_symm_buffer_bytes", "def get_allreduce_mnnvl_workspace", "def trtllm_mnnvl_all_reduce", "def trtllm_mnnvl_fused_allreduce_rmsnorm"]:
for i, line in enumerate(lines, 1):
if needle in line:
print(f"\n== {needle} at line {i} ==")
start = max(1, i-20)
end = min(len(lines), i+60)
for j in range(start, end+1):
print(f"{j:4d}: {lines[j-1]}")
break
PYRepository: flashinfer-ai/flashinfer
Length of output: 11578
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/trtllm_mnnvl_ar.py")
lines = p.read_text().splitlines()
for i, line in enumerate(lines, 1):
if line.startswith("def _alloc_symm_buffer_bytes"):
start = max(1, i-10)
end = min(len(lines), i+120)
for j in range(start, end+1):
print(f"{j:4d}: {lines[j-1]}")
break
PYRepository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/trtllm_mnnvl_ar.py")
lines = p.read_text().splitlines()
for i, line in enumerate(lines, 1):
if "def _alloc_symm_buffer_bytes" in line:
for j in range(i, min(len(lines), i+160)+1):
print(f"{j:4d}: {lines[j-1]}")
break
PYRepository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/trtllm_mnnvl_ar.py")
lines = p.read_text().splitlines()
for i, line in enumerate(lines, 1):
if "def _alloc_symm_buffer_bytes" in line:
print(f"== _alloc_symm_buffer_bytes at line {i} ==")
for j in range(i, min(len(lines), i+140)+1):
print(f"{j:4d}: {lines[j-1]}")
break
print("\n== all uses of `tensor` field in MNNVLAllReduceFusionWorkspace ==")
for i, line in enumerate(lines, 1):
if "self.tensor" in line:
print(f"{i:4d}: {line}")
PYRepository: flashinfer-ai/flashinfer
Length of output: 375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "negative zero|0x8000|0x80000000|buffer_flags|poll|wait until|dirty idx|cur idx|lamport" flashinfer/comm/trtllm_mnnvl_ar.py flashinfer/comm/mnnvl.pyRepository: flashinfer-ai/flashinfer
Length of output: 2610
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/trtllm_mnnvl_ar.py")
lines = p.read_text().splitlines()
# print the main protocol methods around the reduction path
for needle in ["def trtllm_mnnvl_all_reduce", "def trtllm_mnnvl_fused_allreduce_add_rmsnorm", "def trtllm_mnnvl_fused_allreduce_rmsnorm"]:
for i, line in enumerate(lines, 1):
if needle in line:
print(f"\n== {needle} at line {i} ==")
for j in range(i, min(len(lines), i+120)+1):
print(f"{j:4d}: {lines[j-1]}")
break
PYRepository: flashinfer-ai/flashinfer
Length of output: 16636
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/torch_symmetric_memory.py")
lines = p.read_text().splitlines()
for i, line in enumerate(lines, 1):
if "def _alloc_symm_buffer_bytes" in line:
print(f"== _alloc_symm_buffer_bytes at line {i} ==")
for j in range(i, min(len(lines), i+140)+1):
print(f"{j:4d}: {lines[j-1]}")
break
PYRepository: flashinfer-ai/flashinfer
Length of output: 1599
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("flashinfer/comm/torch_symmetric_memory.py")
lines = p.read_text().splitlines()
for i, line in enumerate(lines, 1):
if "def _alloc_symm_buffer_bytes" in line:
print(f"== _alloc_symm_buffer_bytes at line {i} ==")
for j in range(i, min(len(lines), i+160)+1):
print(f"{j:4d}: {lines[j-1]}")
break
PYRepository: flashinfer-ai/flashinfer
Length of output: 1599
Use dtype-aware Lamport initialization here
flashinfer/comm/trtllm_mnnvl_ar.py:181-183 always writes float32 negative zero, so fp16/bf16 workspaces get alternating +0.0/-0.0 halfwords instead of a uniform 0x8000 sentinel. Keep the reduction dtype on the workspace and match the cuMemsetD16/cuMemsetD32 logic already used in flashinfer/comm/mnnvl.py.
🤖 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 `@flashinfer/comm/trtllm_mnnvl_ar.py` around lines 181 - 183, Update the
Lamport initialization in the relevant workspace setup method to use the
reduction dtype when writing the sentinel, matching the existing
cuMemsetD16/cuMemsetD32 behavior in mnnvl.py. Ensure fp16 and bf16 workspaces
receive a uniform 0x8000 pattern while fp32 retains the negative-zero
initialization, and preserve the existing synchronization step.
|
close because not a real issue |
2 participants
Reverts #3745
There are some hangs, need further investigation
Summary by CodeRabbit
Breaking Changes
Improvements