feat(comm): add selected-record PCIe exchange#64
Open
FujitsuPolycom wants to merge 3 commits into
Open
Conversation
Add a direct CUDA-IPC exchange for destination-selected fixed-width records, expose it through the Sparkinfer PCIe facade, and preserve runtime-JIT CUDA sources in built wheels. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds a CUDA-backed PCIe selected-record exchange with distributed configuration validation, CUDA IPC shared buffers, peer scatter/barrier kernels, public API exports, packaging updates, and CPU/GPU test coverage. ChangesSelected-record exchange
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant PCIeSelectedRecordExchange
participant CUDAExtension
participant PeerSlabs
Caller->>PCIeSelectedRecordExchange: exchange(records, indices, out)
PCIeSelectedRecordExchange->>PCIeSelectedRecordExchange: validate inputs and stream
PCIeSelectedRecordExchange->>CUDAExtension: dispatch exchange
CUDAExtension->>PeerSlabs: scatter selected records
CUDAExtension->>PeerSlabs: synchronize peer barriers
PeerSlabs->>Caller: fill output tensor
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ 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 |
This was referenced Jul 21, 2026
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
voipmonitor
marked this pull request as ready for review
July 22, 2026 09:04
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/comm/test_pcie_selected_records.py (1)
398-432: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test that
close()is idempotent.This layer's description claims idempotent close behavior was added, but no test here calls
close()twice to confirm double-close/double-free safety (e.g., thatcudaIpcCloseMemHandle/cudaFreearen't invoked again, or_owned_bufferis cleared after first close).✅ Suggested addition
def test_close_is_idempotent(monkeypatch): ... runtime.close() operations.clear() runtime.close() # should be a no-op assert operations == []🤖 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 `@tests/comm/test_pcie_selected_records.py` around lines 398 - 432, Add a test alongside test_cleanup_rebinds_the_configured_cuda_device that invokes runtime.close() twice, clearing the recorded operations after the first call, and asserts the second call performs no CUDA cleanup operations. Reuse the existing runtime, _CleanupIPC, and _OwnedSharedBuffer setup to verify close() clears or otherwise invalidates cleanup state and prevents repeated cudaIpcCloseMemHandle or cudaFree calls.
🤖 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 `@sparkinfer/comm/pcie/pcie_selected_records.cu`:
- Around line 76-121: Add a system-scope memory fence at the end of
scatter_records_kernel, after all payload writes complete, so remote peer-slab
stores become visible before the subsequent flag exchange. Ensure the fence is
reached by every participating thread, including threads with no record work.
---
Nitpick comments:
In `@tests/comm/test_pcie_selected_records.py`:
- Around line 398-432: Add a test alongside
test_cleanup_rebinds_the_configured_cuda_device that invokes runtime.close()
twice, clearing the recorded operations after the first call, and asserts the
second call performs no CUDA cleanup operations. Reuse the existing runtime,
_CleanupIPC, and _OwnedSharedBuffer setup to verify close() clears or otherwise
invalidates cleanup state and prevents repeated cudaIpcCloseMemHandle or
cudaFree calls.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 83f65be6-ab65-41f2-ae96-a2498c6f9935
📒 Files selected for processing (7)
pyproject.tomlsparkinfer/comm/pcie/__init__.pysparkinfer/comm/pcie/api.pysparkinfer/comm/pcie/pcie_selected_records.cusparkinfer/comm/pcie/pcie_selected_records.pytests/comm/test_pcie_selected_records.pytests/comm/test_pcie_selected_records_gpu.py
Signed-off-by: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add a record-format-agnostic selected-record PCIe exchange for CUDA-IPC peers.
The caller supplies fixed-width CUDA byte records and a finalized per-destination owner map. An SM kernel writes each owned record into the destination IPC slab at its final selected position, with bounded device barriers for visibility and safe slab reuse.
This transport does not compute top-K, ownership, unions, or remaps.
Contract
uint8records and outputint32orint64destination mapThe process group is used for setup/handshakes; payload data uses CUDA IPC peer stores, not NCCL.
Validation
Focused tests cover:
The paired public integration reported 57 Sparkinfer tests passed and 4 GPU-only tests skipped in CI. The native 368-byte direct transport also passed the four-GPU eager/graph byte-equivalence gate and the FP8-RoPE model integration.
Design
GLM-5.2 Sparse CKV under Decode Context Parallelism
The transport is generic. The paired vLLM integration now enables native 368-byte NVFP4 + FP8-RoPE and 432-byte NVFP4 + BF16-RoPE records.
Credits
The selected-record design builds on Koush's sparse-CKV work and Luke Alonso's B12X/Sparkinfer infrastructure.
Implementation, tests, and documentation were developed with OpenAI Codex assistance and manually reviewed and GPU-validated.
Summary by CodeRabbit
SelectedRecordExchangeandSelectedRecordExchangeInitializationError, and expanded supported dtypes to includeuint8.