perf(b12x): index exact declared PCIe collective plans

[voipmonitor]

Purpose

Index the exact PCIe collective declarations at preparation time. Eager
prefill calls can then find a prepared plan, or decline the B12X transport,
without scanning every layer's declaration and constructing an exception.
GPU kernels, routing cutoffs and fallback transports are unchanged.

Inspired by fernandaspets/vllm_sm120#2.
Attribution to neolithic5452 and Neo is retained in the commit.
This implementation also preserves the complete existing lookup contract:

• operation, shape, dtype and strides;
• normalization-weight identity and epsilon for fused RMSNorm;
• the first matching declaration when equivalent calls share a key;
• replacement on declaration refresh and removal on communicator close.

The index contains declared keys only. An unsupported shape does not create
a negative-cache entry, so varying request lengths cannot grow a process-wide
cache indefinitely. Direct calls to an undeclared native plan still raise;
normal transport admission still declines them.

No equivalent open LIL or upstream PR was found. The source PR's key omits
normalization identity/epsilon, so it cannot be adopted unchanged.

Validation

Status: qualified for the lookup contract, TP2/TP4 collective replay, and the
bounded serving comparison below. No general model-throughput gain is claimed.

• tests/distributed/test_b12x_pcie_all_reduce.py: 36 passed, one four-GPU
  case skipped on a two-GPU run. Includes real SM120 fused all-reduce replay.

• The composed integration image passes all 45 tests in that file on four
  SM120 GPUs, including TP2 fused and TP4 BF16 two-shot mutated graph replay.
  The composition retains weak collective-owner lifetimes from fix(b12x): release replaced collective descriptor owners #783.

• Tests distinguish equal-valued but different normalization tensors,
  epsilon, dtype, shape and strides; verify first-match precedence, refresh,
  undeclared rejection and bounded index size.

• CPU-only metadata benchmark, same Python/runtime, 5 x 1,000 warmed calls:

  Declared calls	Matching final declaration, scan → index	Undeclared shape, scan → index
  64	6.90 → 0.49 µs	18.94 → 0.68 µs
  512	52.14 → 0.50 µs	139.25 → 0.70 µs
  4096	488.72 → 0.50 µs	1104.88 → 0.69 µs

Reproduce with benchmarks/benchmark_b12x_plan_lookup.py --module-file PATH_TO_PINNED_COMMUNICATOR_SOURCE. These are Python lookup timings, not
GPU kernel or model tok/s gains. The benchmark records source hashes and
all samples. Applicable pre-commit hooks pass.

Model-serving comparison

DeepSeek V4.1, stock RTX PRO 6000 Workstation GPUs 4–7, TP4/DCP1,
adaptive DSpark K7, RAM Engram, 4,096-token budget, temperature 1/top-p .95.
Three warmed 30-second context-zero runs per concurrency; one warmed uncached
32K prefill window. The immutable images differ only in the composed PCIe
communicator source; serving has no source-code mounts.

Measurement	Declaration scan	Declared-plan index	Change
C1 output, tok/s	247.61	254.92	+2.95%
C1 request verifier, steps/s	101.42	99.94	−1.45%
C8 aggregate output, tok/s	799.07	799.57	+0.06%
C8 summed request verifier, steps/s	352.12	350.86	−0.36%
Uncached 32K prefill, tok/s	20,471	20,577	+0.52%

Five API checks and all six decode cells pass in each arm. C1 accepted length
is 2.443 versus 2.551; adaptive verification and stochastic output affect the
rates. C1 output ranges overlap (247.43–266.71 versus 250.35–263.18).
The negative verifier deltas are retained. These samples do not isolate an
end-to-end speedup or prove statistical equivalence; the qualified optimization
is bounded constant-time metadata lookup without changing transport selection.

Local image identities: scan ec81e2005bac38c252329c3c2a6fbcd84188afaef7e54080332cd54359a67ef0,
index 0c4b935622a16713bc56994a409b5cbec84dfd008561f7e2e2f4a4aa5b2ae3ba.

Raw samples, exact runtime arguments, source identities and native checks.

AI-assisted implementation and validation. Human maintainer review is required.
Integration is restricted to beta during serving qualification; no JJ merge is
authorized by this work.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 1be997ca-ee52-434e-a036-65e19bd4671d

📥 Commits

Reviewing files that changed from the base of the PR and between 8e1f1e5 and 5e69387.

📒 Files selected for processing (3)

• benchmarks/benchmark_b12x_plan_lookup.py
• tests/distributed/test_b12x_pcie_all_reduce.py
• vllm/distributed/device_communicators/b12x_pcie_all_reduce.py

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

Changes

B12X plan lookup

Layer / File(s)	Summary
Build and query the declared-plan index vllm/distributed/device_communicators/b12x_pcie_all_reduce.py	B12xPcieAllReduce builds an index for declared plans, uses it for _plan_for and _has_plan_for, and clears it during close().
Validate indexed lookup behavior tests/distributed/test_b12x_pcie_all_reduce.py	Tests cover fused identity, shape/dtype/stride matching, declaration precedence, re-indexing, missing plans, and non-mutating probes.
Measure plan lookup performance benchmarks/benchmark_b12x_plan_lookup.py	A CPU-only benchmark measures warmed hit and miss lookups and emits JSON results with configuration and source hashing.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant B12xPcieAllReduce
  participant PlanIndex
  participant PlanLookupCaller
  B12xPcieAllReduce->>PlanIndex: Index declared plans
  PlanLookupCaller->>B12xPcieAllReduce: Request plan lookup
  B12xPcieAllReduce->>PlanIndex: Query canonical plan key
  PlanIndex-->>B12xPcieAllReduce: Return plan or no match

Loading

Suggested reviewers: lukealonso

Merge Risk: ⚪ Minimal · up to 5e693

No concrete merge-blocking issue is established for the indexed lookup change.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 5.88% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 34 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: indexing exact declared PCIe collective plans for B12X to improve lookup performance.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/pcie-declared-plan-index

---

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.}

[voipmonitor]

@coderabbitai review

[coderabbitai]

⚠️ Action not completed

Review rate limited.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[voipmonitor]

@coderabbitai review

[coderabbitai]

⚠️ Action not completed

Review rate limited.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[voipmonitor]

@coderabbitai full review

[coderabbitai]

✅ Action performed

Full review finished.