feat(pcie_dma): add INT8 and MXFP8 wire modes#72
Conversation
📝 WalkthroughWalkthroughPCIe DMA compressed wire transport now supports INT8 and MXFP8 codecs across mode parsing, CUDA kernels, all-reduce routing, documentation, and GPU/CPU validation tests. ChangesCompressed PCIe DMA wire codecs
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant PCIeDmaAllReduce
participant PyBindExtension
participant CUDAKernels
participant BF16Output
Caller->>PCIeDmaAllReduce: all_reduce(mode)
PCIeDmaAllReduce->>PyBindExtension: select codec launcher
PyBindExtension->>CUDAKernels: quantize and reduce compressed shards
CUDAKernels-->>PCIeDmaAllReduce: forward compressed payloads
PCIeDmaAllReduce->>PyBindExtension: dequantize final shards
PyBindExtension->>BF16Output: store BF16 results
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 |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![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
🧹 Nitpick comments (1)
tests/comm/test_pcie_dma_int8_codec.py (1)
30-93: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winOptional: add a bit-exact GPU-oracle test for the INT8 kernels. MXFP8 gets a direct
dma_quant_mx/dma_dequant_store_mxvs CPU-reference check (test_pcie_dma_mxfp8_codec_matches_cpu_reference), but the INT8 GPU kernels (dma_quant_i8/dma_dequant_store_i8) are only validated indirectly through the tolerance-based four-GPU gate. A parallel GPU codec test would give the INT8 path the same bit-exactness guarantee.🤖 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_dma_int8_codec.py` around lines 30 - 93, Add a dedicated GPU codec test alongside test_int8_roundtrip_has_half_step_error_bound that runs dma_quant_i8 and dma_dequant_store_i8 on representative INT8 input, then compares payload, scales, and restored values bit-exactly with _quantize and _dequantize. Reuse the existing GPU execution/setup conventions and include edge cases such as zero and out-of-range values.
🤖 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_dma.cu`:
- Around line 173-177: Gate the MXFP8 helper usage in mxfp8_scale_byte and
related E8M0 helpers on CUDA Toolkit 12.6 availability, using the repository’s
existing CUDA-version detection conventions. Ensure older toolkits do not
compile or select the __nv_cvt_float_to_e8m0 path, or explicitly
enforce/document CUDA 12.6 as the minimum if that is the established project
policy.
---
Nitpick comments:
In `@tests/comm/test_pcie_dma_int8_codec.py`:
- Around line 30-93: Add a dedicated GPU codec test alongside
test_int8_roundtrip_has_half_step_error_bound that runs dma_quant_i8 and
dma_dequant_store_i8 on representative INT8 input, then compares payload,
scales, and restored values bit-exactly with _quantize and _dequantize. Reuse
the existing GPU execution/setup conventions and include edge cases such as zero
and out-of-range values.
🪄 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: 314a3ec4-1dc0-44a7-af5e-ba4005ea6b98
📒 Files selected for processing (8)
README.mdsparkinfer/comm/pcie/pcie_dma.cusparkinfer/comm/pcie/pcie_dma.pytests/comm/test_pcie_dma_gpu.pytests/comm/test_pcie_dma_int8_codec.pytests/comm/test_pcie_dma_mode.pytests/comm/test_pcie_dma_mxfp8_codec.pytests/comm/test_pcie_dma_rank_consistency_gpu.py
| __device__ __forceinline__ unsigned char mxfp8_scale_byte(float amax) { | ||
| if (amax == 0.0f) return 127; | ||
| return __nv_cvt_float_to_e8m0( | ||
| amax / kFp8Max, __NV_SATFINITE, cudaRoundPosInf); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
__nv_cvt_float_to_e8m0 CUDA toolkit version introduced cuda_fp8.h
💡 Result:
The CUDA Math API function nv_cvt_float_to_e8m0 and the related E8M0 data type support were introduced in the cuda_fp8.h header in CUDA Toolkit 12.6 [1]. While the cuda_fp8.h header itself has been available in earlier versions of the CUDA Toolkit (e.g., since at least CUDA 11.8 for standard FP8 types like E4M3 and E5M2) [2][3][4], support for the E8M0 format—specifically used for scaling factors—was added in version 12.6 [1]. Third-party projects often implement conditional compilation guards (e.g., checking for CUDA 12.6+) to maintain compatibility with older toolkit versions when using this specific functionality [1].
Citations:
- 1: [Compat] Add CUDA version check for __nv_fp8_e8m0 type tile-ai/tilelang#1537
- 2: https://docs.nvidia.com/cuda/archive/11.8.0/cuda-math-api/group__CUDA__MATH__FP8__MISC.html
- 3: https://docs.nvidia.com/cuda/archive/11.8.0/cuda-math-api/group__CUDA__MATH__INTRINSIC__FP8.html
- 4: https://docs.nvidia.com/cuda/archive/11.8.0/pdf/CUDA_Math_API.pdf
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== repository files =="
git ls-files | rg '(^|/)(cuda\.arch|requirements|pyproject|setup\.py|CMakeLists|cmake|.*cuda.*|.*torch.*ext.*|.*extension.*|README|.*\.cu$)' || true
echo
echo "== relevant references =="
rg -n "__nv_cvt_float_to_e8m0|cudaRoundPosInf|__NV_SATFINITE|cuda_fp8\.h|nv_cvt_float_to_e8m0|load_extension|torch\.utils\.cpp_extension|CMAKE_CUDA|CUDA_ARCH|CUDA_VERSION|cuda" -S --glob '!build/**' --glob '!dist/**' . | head -n 240Repository: local-inference-lab/sparkinfer
Length of output: 28066
Gate E8M0 MXFP8 helpers on CUDA 12.6+ availability. __nv_cvt_float_to_e8m0 is only available starting in CUDA Toolkit 12.6, but this repository does not document or enforce a CUDA >=12.6 minimum. Since pcie_dma.cu includes this directly, builds with older toolkits can fail for the MXFP8 route. Guard these helpers and/or raise the minimum supported toolkit requirement.
🤖 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 `@sparkinfer/comm/pcie/pcie_dma.cu` around lines 173 - 177, Gate the MXFP8
helper usage in mxfp8_scale_byte and related E8M0 helpers on CUDA Toolkit 12.6
availability, using the repository’s existing CUDA-version detection
conventions. Ensure older toolkits do not compile or select the
__nv_cvt_float_to_e8m0 path, or explicitly enforce/document CUDA 12.6 as the
minimum if that is the established project policy.
|
Fresh manual review and validation completed at Validation performed:
The release integration is kept separate from this PR: |
3 participants
Summary
Adds two opt-in compressed wire codecs to
PCIeDmaAllReduce, each available on the existing all-gather, ring, and all-to-all topologies:i8,i8_ring,i8_a2a: signed block INT8 with one FP32 scale per 128 valuesmx,mx_ring,mx_a2a: standard MXFP8 with E4M3 payload values and one E8M0 scale per 32 valuesAll compressed layouts use 132 bytes per 128 values instead of 256 bytes for BF16, reducing wire bytes by 48.4%. Existing BF16 and E4M3 modes are unchanged.
This is an independently maintained combined alternative to #69. It is based directly on current SparkInfer
masterand includes the requestedmx_ringimplementation without modifying the other contributor's PR.Modes
i8i8_ringi8_a2amxmx_ringmx_a2aEvery mode materializes the owner shard from the same canonical payload forwarded to peers. This preserves the all-reduce requirement that every rank receives bit-identical replicated activations.
Unknown non-empty mode strings raise
ValueError; a typo such asmx_rnigcan no longer silently select E4M3ag. Unit tests cover every accepted INT8 and MXFP8 alias, including whitespace and case normalization.Implementation
Validation
Completed on 2026-07-22 against current SparkInfer
master(1a88b389) and the GLM-5.2 v20 stack:git diff --check: passGLM-5.2 TP4/DCP4/MTP0 prefill
agi8_ringmx_ringThe approximately 1% difference between INT8 and MXFP8 is small enough to treat as run/pair variance. Both ring codecs materially outperform the E4M3 AG topology in this prefill profile while using the same wire bytes.
Paired teacher-forced prompt-logprob drift
An 8,191-position comparison against uncompressed BF16 transport on the same model and prompt produced:
i8_ringmx_ringThis is selected-token teacher-forced drift, not full-vocabulary KLD. INT8 was slightly better on this workload; codec selection remains an explicit model-specific choice.
The current long-context needle profile stops being a valid codec discriminator above 200k because its BF16 transport baseline also fails. This PR therefore makes no unsupported long-context retrieval claim.
Configuration
Set
SPARKINFER_PCIE_DMA_FP8or pass the same string as thefp8=constructor argument:Compressed transport requires BF16 input and a per-rank shard divisible by 128 elements. Other shapes retain the BF16 path.
Summary by CodeRabbit
New Features
SPARKINFER_PCIE_DMA_FP8or the constructor’sfp8option.Bug Fixes
Tests