feat(pcie_dma): add block-INT8 wire modes to b12x v20#69
Conversation
Signed-off-by: derek <derek.yates@live.com>
📝 WalkthroughWalkthroughChangesPCIe DMA compressed wire support now includes block-scaled signed INT8 quantization across ring, all-gather, and all-to-all modes, with updated CUDA kernels, transport routing, documentation, and tests. PCIe DMA INT8 compressed communication
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant PCIeDmaAllReduce
participant _all_reduce_fp8
participant INT8Kernels
participant BF16Output
Caller->>PCIeDmaAllReduce: select i8, i8_ring, or i8_a2a
PCIeDmaAllReduce->>_all_reduce_fp8: route compressed all-reduce
_all_reduce_fp8->>INT8Kernels: quantize and exchange INT8 payloads
INT8Kernels->>BF16Output: dequantize reduced results
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
sparkinfer/comm/pcie/pcie_dma.py (1)
72-92: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winSilent fallback to
"ag"on unrecognized mode strings.Any misspelled
SPARKINFER_PCIE_DMA_FP8value (e.g.i8_rnig) silently normalizes to"ag"(Line 92) rather than raising, so a typo intending INT8 quantization silently falls back to the FP8 E4M3 codec with only a debug log line as evidence. Now that this function accepts a much larger set of i8/int8 aliases, the odds of an unnoticed typo changing the wire codec in production increase.♻️ Proposed fix: raise on unrecognized non-empty values
- if raw in ("i8_a2a", "i8-a2a", "int8_a2a", "int8-a2a", "a2a_i8"): - return "i8_a2a" - return "ag" + if raw in ("i8_a2a", "i8-a2a", "int8_a2a", "int8-a2a", "a2a_i8"): + return "i8_a2a" + if raw in ("ag", "1"): + return "ag" + raise ValueError(f"unrecognized SPARKINFER_PCIE_DMA_FP8 mode: {value!r}")🤖 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.py` around lines 72 - 92, Update _normalize_fp8_mode to raise an appropriate error for any unrecognized non-empty mode value instead of returning "ag"; preserve the existing normalization for supported aliases and the empty/false values that disable the mode.
🤖 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 `@sparkinfer/comm/pcie/pcie_dma.py`:
- Around line 72-92: Update _normalize_fp8_mode to raise an appropriate error
for any unrecognized non-empty mode value instead of returning "ag"; preserve
the existing normalization for supported aliases and the empty/false values that
disable the mode.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: f09fe0f7-ab7d-4fd8-8944-15e294162013
📒 Files selected for processing (6)
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_rank_consistency_gpu.py
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/comm/test_pcie_dma_mode.py (1)
10-29: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover the complete INT8 alias contract.
The implementation in
sparkinfer/comm/pcie/pcie_dma.py, Lines 72-94, accepts canonicali8,i8_ring, andi8_a2avalues plus additional aliases, but this test does not exercise the canonical INT8 values and covers only a subset of accepted spellings. Parameterize the full alias table, including at least one whitespace/case-normalization case.🤖 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_mode.py` around lines 10 - 29, Expand test_supported_wire_mode_aliases to cover the complete INT8 alias contract from _normalize_fp8_mode, including canonical i8, i8_ring, and i8_a2a values plus every accepted INT8 alias. Add at least one input with surrounding whitespace and mixed case, asserting it normalizes to the expected canonical value while preserving the existing non-INT8 cases.
🤖 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 `@tests/comm/test_pcie_dma_mode.py`:
- Around line 10-29: Expand test_supported_wire_mode_aliases to cover the
complete INT8 alias contract from _normalize_fp8_mode, including canonical i8,
i8_ring, and i8_a2a values plus every accepted INT8 alias. Add at least one
input with surrounding whitespace and mixed case, asserting it normalizes to the
expected canonical value while preserving the existing non-INT8 cases.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 5a702d96-9eb7-472b-81a0-f2e994af5e8a
📒 Files selected for processing (3)
sparkinfer/comm/pcie/pcie_dma.pytests/comm/test_pcie_dma_int8_codec.pytests/comm/test_pcie_dma_mode.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/comm/test_pcie_dma_int8_codec.py
- sparkinfer/comm/pcie/pcie_dma.py
a43ef01 to
d6f0baa
Compare
|
Superseded by #72, which preserves the block-INT8 implementation and adds the validated standard MXFP8 wire modes, including mx_ring. Closing this PR without deleting the contributor branch so its history remains available. |
Summary
Ports the block-INT8 PCIe-DMA wire modes from #46 onto current b12x
masterfor the GLM-5.2 v20 stack. The repository has been renamed back to b12x; its current Python distribution and source namespace remainsparkinfer, so the implementation paths in this PR are undersparkinfer/.i8: BF16 reduce-scatter plus block-INT8 all-gatheri8_ring: block-INT8 reduce-scatter ring plus block-INT8 all-gatheri8_a2a: quantize-once block-INT8 all-to-all plus block-INT8 broadcastEach 128-value block uses 128 signed payload bytes and one FP32 scale, so INT8 uses the same 132-byte wire and staging layout as E4M3. Existing BF16 and E4M3 routes are unchanged.
This is a clean replacement for #46, whose head predates the
b12x/tosparkinfer/source-namespace migration and is now conflicting withmaster.Why
GLM-5.2 long-context validation isolated a residual E4M3 representation loss after the rank-consistency fix in #44:
All INT8 topologies later passed needle retrieval through 475k. On the tested TP4/DCP4 PCIe system,
i8_ringreached 1,641 tok/s cold prefill at 50k, matching E4M3 ring at 1,639 tok/s, while retaining the full 644,864-token KV pool. Full historical evidence remains in #46.V20 integration
The published v20 artifact pins the b12x runtime commit
6a92bcc0f2bf03b13dd03dbc7ce97e26133c580e, which contains E4M3ag,ring, anda2abut not the INT8 codec. The complete patch in this PR passesgit apply --checkagainst that exact release commit without conflicts.The v20 launcher currently validates only
F8_DMA=0|ag|ring; acceptingi8,i8_ring, andi8_a2abelongs in the separateblackwell-llm-dockerintegration repository after this runtime support is available.Implementation
Validation
Completed locally on current
master:git diff --check: pass6a92bccpatch applicability: passPreviously completed on the target four-GPU system for the same codec implementation:
i8,i8_ring, andi8_a2aPending before this draft is marked ready:
wire mode: int8-ring, throughput, KV pool, and zero restartsFiles
sparkinfer/comm/pcie/pcie_dma.cusparkinfer/comm/pcie/pcie_dma.pytests/comm/test_pcie_dma_gpu.pytests/comm/test_pcie_dma_rank_consistency_gpu.pytests/comm/test_pcie_dma_int8_codec.pyREADME.mdSigned-off-by: Derek Yates derek.yates@live.com
Summary by CodeRabbit
New Features
Tests