feat(pcie_dma): add block-INT8 wire modes for AG, ring, and A2A

[yatesdr]

feat(pcie_dma): add block-INT8 wire modes for AG, ring, and A2A

Summary

Adds an opt-in symmetric block-INT8 codec to the existing PCIe-DMA compressed
all-reduce paths:

• i8: BF16 reduce-scatter + INT8 all-gather
• i8_ring: INT8 reduce-scatter ring + INT8 all-gather
• i8_a2a: quantize-once INT8 all-to-all + INT8 broadcast

Each 128-value block uses 128 signed payload bytes and one FP32 scale
(scale = amax / 127), so the wire and staging footprint is exactly the same
as the existing E4M3 codec: 132 bytes per block. The change is opt-in;
BF16 and all existing E4M3 modes retain their current routing and kernels.

This builds on the rank-consistency correction merged in #44. Owners
materialize their shard from the same final wire payload forwarded to peers in
all three INT8 topologies.

Why

On a GLM-5.2 TP4/DCP4 deployment, the rank-consistent E4M3 AG path recovered
needle retrieval through 200k but continued to fail reproducibly at 300k and
350k. Replacing only the one-byte E4M3 payload with symmetric block INT8, while
keeping the block size, FP32 scale, topology, staging allocation, and owner
materialization unchanged, restored deep retrieval:

Wire mode	300k	350k
E4M3 AG	0/3	0/3
block-INT8 AG	3/3	3/3

All six INT8 passes returned the expected needle (738216) with
finish_reason=stop. RestartCount remained zero throughout those retrieval
runs. This isolates the residual quality loss to E4M3 representation rather
than transport topology, wire size, or rank divergence.

Implementation

• Adds CUDA kernels/bindings for:
  • BF16 -> signed INT8 quantization with one scale per 128 values
  • signed INT8 -> BF16 materialization
  • multi-source INT8 dequantization with FP32 accumulation for A2A
  • fused INT8 dequantize-add-requantize for each ring reduce-scatter hop
• Extends the existing mode normalization and routing without changing the
  BF16/E4M3 behavior.
• Keeps the collective route rank-invariant: it depends only on configured
  mode, dtype, and aligned tensor shape.
• Reuses the existing compressed scratch/staging layout; there is no codec
  workspace increase relative to E4M3.
• Extends eager, CUDA-graph, and cross-rank bit-identity gates to all three
  INT8 modes.
• Adds dependency-free codec proofs covering the 132-byte layout, zero blocks,
  the half-quantization-step error bound, and canonical-payload forwarding for
  AG, ring, and A2A.

Accepted aliases include int8, int8-ag, int8-ring, and int8-a2a; the
canonical mode names are i8, i8_ring, and i8_a2a.

Evidence and validation status

Completed:

• All three modes passed four-rank bit-identity checks in eager and CUDA-graph
  execution at rows 512 and 3072. Maximum absolute error versus BF16 was
  0.01758 (i8), 0.01953 (i8_ring), and 0.01563 (i8_a2a).

• Full-depth needle retrieval passed at 50k, 200k, 300k, 350k, and 475k for
  every topology. i8 passed 300k and 350k 3/3. One non-monotonic
  i8_a2a miss at 350k followed a 475k pass and cleared a 3/3 recheck.

• Cold prefill retained the E4M3 performance class:

  Context	i8	i8_ring	i8_a2a
  8k	1,487	1,607	1,463
  50k	1,519	1,641	1,432

• Decode throughput was topology-invariant, as expected because decode
  messages remain below the 6.29 MiB DMA threshold:

  Context 0	C1	C2	C4	C8	C16
  i8	61.9	74.8	96.7	127.9	164.2
  i8_ring	63.2	72.7	95.7	127.2	165.1
  i8_a2a	62.1	75.3	104.7	122.6	162.7

• Clean verified boots reported a 644,864-token KV pool for all three modes,
  with no mode-correlated capacity change.

• Dependency-free codec/topology proofs: 4/4 pass.

• Python compile and CUDA binding/call-site structural checks pass.

• git diff --check passes.

i8_ring is the recommended PCIe mode on the tested no-NVLink topology: its
1,641 tok/s at 50k matches E4M3 ring (1,639) while preserving deep retrieval.

One i8_a2a performance run encountered a downstream
CUBLAS_STATUS_INTERNAL_ERROR in a BF16 MLA projection during decode, where
the PCIe-DMA path is inactive. A clean rerun completed. This is tracked as a
separate v19 concurrency issue and is not attributed to the INT8 codec or A2A
topology.

Files changed

• b12x/distributed/pcie_dma.cu
• b12x/distributed/pcie_dma.py
• tests/distributed/test_pcie_dma_gpu.py
• tests/distributed/test_pcie_dma_rank_consistency_gpu.py
• tests/distributed/test_pcie_dma_int8_codec.py

Summary by CodeRabbit

• New Features

  • Added optional INT8 compressed PCIe DMA wire modes for all-reduce operations.
  • Supports i8, i8_ring, and i8_a2a transport modes alongside existing modes.
  • Added per-block scaling for compressed BF16 transfers, with support for reduce-scatter, ring, and all-to-all workflows.

• Documentation

  • Added configuration guidance, supported mode details, storage overhead, constraints, and example environment settings.

• Tests

  • Expanded GPU coverage and added validation for INT8 accuracy, layout, zero blocks, and collective consistency.

[coderabbitai]

Warning

Review limit reached

@yatesdr, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 35 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 08d4c1ce-9f0c-405c-9dea-d22a02ecda56

📥 Commits

Reviewing files that changed from the base of the PR and between 1a88b38 and 1bc4f1b.

📒 Files selected for processing (6)

• README.md
• b12x/distributed/pcie_dma.cu
• b12x/distributed/pcie_dma.py
• tests/distributed/test_pcie_dma_gpu.py
• tests/distributed/test_pcie_dma_int8_codec.py
• tests/distributed/test_pcie_dma_rank_consistency_gpu.py

📝 Walkthrough

Walkthrough

Adds signed-INT8 PCIe DMA wire codecs, integrates i8, i8_ring, and i8_a2a modes into compressed all-reduce paths, documents configuration, and expands codec, GPU, and rank-consistency coverage.

Changes

INT8 PCIe DMA compression

Layer / File(s)	Summary
INT8 CUDA codec and exports b12x/distributed/pcie_dma.cu	Adds per-128-element INT8 quantization, dequantization, accumulation, requantization, host launchers, and PyBind exports.
Compressed-mode all-reduce integration b12x/distributed/pcie_dma.py	Normalizes INT8 mode aliases, maps modes to INT8 wire codecs, and routes compressed all-reduce stages through the appropriate kernels.
Mode coverage and codec validation README.md, tests/distributed/test_pcie_dma_int8_codec.py, tests/distributed/test_pcie_dma_gpu.py, tests/distributed/test_pcie_dma_rank_consistency_gpu.py	Documents mode configuration and validates INT8 layout, reconstruction, collective patterns, GPU modes, and rank-consistency options.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PCIeDmaAllReduce
  participant PythonExtension
  participant CUDACodec
  PCIeDmaAllReduce->>PythonExtension: select INT8 DMA entry points
  PythonExtension->>CUDACodec: quantize BF16 payload
  CUDACodec->>PythonExtension: return INT8 payload and scales
  PythonExtension->>CUDACodec: dequantize, accumulate, or requantize
  CUDACodec->>PCIeDmaAllReduce: materialize BF16 output

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.63% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: adding block-INT8 wire modes for AG, ring, and A2A.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}