feat(distributed): add HOST-orchestrated builtin.tensor.all_to_all_v

[georgebisbas]

Summary

Closes the last gap in host-orchestrator support for variable-size all-to-all. pld.tensor.all_to_all_v (variable-size all-to-all, MPI_Alltoallv pattern) merged as #2112, but only as an InCore composite — calling it from a host_orch function hit an explicit "not supported in a HOST orchestration function" rejection. This adds the missing HOST dispatch path, cloning the proven builtin.tensor.all_to_all pattern (#1997) across every layer:

• Op registration (collective.cpp): builtin.tensor.all_to_all_v, narrowing input and send_counts to strict window-bound DistributedTensor — forced by EmitBuiltinWindowCollectiveDispatch, which has no dispatch path for a plain Tensor arg at this layer (same narrowing the existing builtin.tensor.all_to_all already applies to its own input).
• Lowering rule (lower_host_tensor_collectives_pass.cpp): MakeBuiltinAllToAllV, deriving MAX_RECV = target.shape[0] / signal.shape[0] and forwarding it as a kwarg/attr alongside dtype.
• Codegen (distributed_ops_codegen.cpp): the variant string mangles both max_recv and dtype (builtin.tensor.all_to_all_v__maxrecv<N>__fp32) — MAX_RECV is baked into the kernel as a compile-time constexpr, and codegen's MarkBuiltinEmitted/RecordBuiltinNextLevel state is program-global, so two call sites with different MAX_RECV would otherwise silently mis-share one kernel instantiation.
• Runtime templates: new all_to_all_v/ builtin package. The kernel always transfers the full MAX_RECV-row block per destination — matching LowerTensorAllToAllVRule's compile-time transfer_shape exactly, for bit-for-bit InCore/HOST parity on the wire — and publishes the runtime-clamped count into peer recv_counts via TNOTIFY inline with the push.
• Comm-domain analysis (materialize_comm_domain_scopes_pass.cpp): device-coverage inheritance for both signal and recv_counts (two CollectiveConsumer entries sharing the same data alloc), plus a HOST-side loop-use guard (repeating_scope_depth_, tracked for both ForStmt and WhileStmt) mirroring the InCore path's CheckAllReduceLoopUse — the HOST path had no equivalent protection before this change, so a caller could previously reuse a single-use signal across loop iterations undetected.
• Removed the explicit IsInCoreOnlyCollective rejection in LowerCompositeOps now that a HOST rule exists for this op.
• Docs (EN+ZH): documents the new HOST path and corrects a pre-existing inaccuracy — the InCore lowering always transfers the full MAX_RECV block; only the published count was ever runtime-gated, not the transfer itself.

Design notes for reviewers

• send_counts must be staged through a window buffer at the HOST layer even though it's logically pure-local, per-rank data (never cross-rank-published) — this is a real ergonomic cost of the narrowing above, not a bug. The new HOST system test documents this with an explicit fill_counts_step/fill_counts_orch staging pair.
• Kept the barrier as the existing single-use Set(1)/wait≥1 protocol, matching every other HOST builtin and the InCore all_to_all_v rule this clones — the reusable credit-barrier rework (fix(ir): make composite collective barrier signals reusable #2175) is unmerged and out of scope here.
• Kept FP32-only (CheckSupportedFp32BuiltinVariant), matching every HOST builtin except allreduce. Adding FP16 across all HOST builtins is a separate, broader gap.
• all_to_all/allgather/barrier share the same single-use-signal exposure to host_orch loops and have no equivalent guard today — pre-existing, out of scope for this PR.

Test plan

• Full C++ rebuild in the sim Docker image (pypto3-hw-native-sys:sim) — clean.
• Full unit test suite: 8539 passed, 13 skipped, 0 failed (includes new/updated tests across all 5 touched layers, zero regressions elsewhere).
• pre-commit run --all-files: clean (headers, English-only, EN/ZH docs parity, docs-nav, clang-format, cpplint, markdownlint, ruff check/format, pyright).
• Distributed system tests not run in this environment. tests/st/distributed/test_l3_host_tensor_all_to_all_v.py (new) plus the InCore/HOST all_to_all regression STs all hit ImportError: cannot import name 'RunTiming' from '_task_interface' — a pre-existing local-environment mismatch (the runtime submodule checkout is ahead of what the sim image's simpler runtime was built against), reproducing identically on unmodified, already-shipped sibling tests. Unrelated to this change, but means the actual distributed exchange logic has not been exercised end-to-end here.
• NPU verification pending.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 6cf066a8-2787-47f9-a229-6e69ff21d132

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

pld.tensor.all_to_all_v now supports HOST builtin lowering in addition to InCore lowering. The change adds IR validation, communication-domain materialization, loop checks, generated runtime kernels, compile-time variants, and unit, backend, and end-to-end tests.

Changes

HOST all-to-all-v

Layer / File(s)	Summary
IR contract and host lowering src/ir/op/distributed/collective.cpp, src/ir/transforms/lower_composite_ops_pass.cpp, src/ir/transforms/lower_host_tensor_collectives_pass.cpp, tests/ut/ir/...	Registers and validates builtin.tensor.all_to_all_v. Host lowering emits the builtin with window arguments, directions, dtype, and max_recv.
Communication-domain materialization src/ir/transforms/materialize_comm_domain_scopes_pass.cpp, tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py, docs/en/dev/passes/39-materialize_comm_domain_scopes.md, docs/zh/dev/passes/39-materialize_comm_domain_scopes.md	Resolves all-to-all-v windows, propagates device coverage to signals and receive counts, and rejects calls inside for or while loops.
Runtime templates and variant generation python/pypto/runtime/builtins/collectives/all_to_all_v/..., src/codegen/distributed/distributed_ops_codegen.cpp, tests/ut/codegen/distributed/test_host_orch_distributed.py	Generates variants keyed by max_recv and data type. The kernel transfers fixed-capacity blocks, publishes clamped receive counts, and synchronizes peer signals.
End-to-end coverage and documentation tests/st/distributed/test_l3_host_tensor_all_to_all_v.py, docs/en/dev/distributed_ops.md, docs/zh/dev/distributed_ops.md, docs/en/dev/passes/40-lower_host_tensor_collectives.md, docs/zh/dev/passes/40-lower_host_tensor_collectives.md	Tests HOST execution with two and four ranks. Documentation describes capacity transfers, count handling, window requirements, variants, and loop restrictions.

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

Sequence Diagram(s)

sequenceDiagram
  participant HostOrchestration
  participant AllToAllVKernel
  participant TargetWindow
  participant PeerSignals
  HostOrchestration->>AllToAllVKernel: submit seven kernel arguments
  AllToAllVKernel->>TargetWindow: TPUT fixed MAX_RECV blocks
  AllToAllVKernel->>TargetWindow: publish clamped recv_counts
  AllToAllVKernel->>PeerSignals: notify peer completion
  PeerSignals-->>AllToAllVKernel: GE completion signals

Loading

Possibly related issues

• [Feature] Add orchestration-level collective communication operators to PyPTO compiler #1189 — Covers HOST-level collective support for all_to_all_v, which this PR implements.

Possibly related PRs

• hw-native-sys/pypto#1997 — Adds related HOST collective IR, lowering, codegen, and communication-domain infrastructure.
• hw-native-sys/pypto#2112 — Introduces the existing pld.tensor.all_to_all_v operation extended here with HOST lowering.
• hw-native-sys/pypto#1829 — Extends shared HOST collective infrastructure used by this implementation.

Poem

A rabbit hops through windows wide,
Counts are sent and rows reside.
Full blocks travel, signals ring,
HOST paths make the tensors sing.
MAX_RECV guards the way—
All-to-all-v works today!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 18.87% 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
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.
Title check	✅ Passed	The title clearly summarizes the main change: adding HOST-orchestrated support for builtin.tensor.all_to_all_v.
Description check	✅ Passed	The description directly explains the HOST dispatch implementation, affected layers, tests, and known verification limits.

---

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

[Copilot]

Pull request overview

Adds HOST-orchestrated variable-size all-to-all support across IR lowering, codegen, runtime templates, tests, and documentation.

Changes:

• Registers and lowers builtin.tensor.all_to_all_v.
• Adds MAX_RECV-specific codegen and runtime kernels.
• Extends communication-domain analysis and test coverage.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/ir/op/distributed/collective.cpp	Registers the internal builtin.
src/ir/transforms/lower_host_tensor_collectives_pass.cpp	Adds HOST lowering.
src/ir/transforms/lower_composite_ops_pass.cpp	Defers HOST calls.
src/ir/transforms/materialize_comm_domain_scopes_pass.cpp	Adds domain inheritance and loop checks.
src/codegen/distributed/distributed_ops_codegen.cpp	Emits MAX_RECV-specific variants.
python/pypto/runtime/builtins/collectives/all_to_all_v/__init__.py	Defines the template package.
python/pypto/runtime/builtins/collectives/all_to_all_v/templates/entry.cpp.in	Adds orchestration entrypoint.
python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel.cpp.in	Implements the collective kernel.
python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel_config.py.in	Configures kernel arguments.
tests/ut/ir/test_distributed_ops.py	Tests registration and resources.
tests/ut/ir/transforms/test_lower_composite_ops.py	Tests HOST deferral.
tests/ut/ir/transforms/test_lower_host_tensor_collectives.py	Tests HOST lowering.
tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py	Tests domains and loop rejection.
tests/ut/codegen/distributed/test_host_orch_distributed.py	Tests generated artifacts.
tests/st/distributed/test_l3_host_tensor_all_to_all_v.py	Adds end-to-end system coverage.
docs/en/dev/distributed_ops.md	Documents HOST support.
docs/en/dev/passes/39-materialize_comm_domain_scopes.md	Documents domain analysis.
docs/en/dev/passes/40-lower_host_tensor_collectives.md	Documents HOST lowering.
docs/zh/dev/distributed_ops.md	Synchronizes Chinese operator docs.
docs/zh/dev/passes/39-materialize_comm_domain_scopes.md	Synchronizes Chinese domain docs.
docs/zh/dev/passes/40-lower_host_tensor_collectives.md	Synchronizes Chinese lowering docs.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 72bfa69fb3

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@src/ir/transforms/lower_host_tensor_collectives_pass.cpp`:
- Around line 230-268: Validate the HOST-level input and send_counts types
before HostCollectiveRule::scope_buffers calls GetWindowBuffer, since
MakeBuiltinAllToAllV narrows them too late. Update GetWindowBuffer or add a
user-facing CHECK_SPAN at that boundary to require DistributedTensorType, using
the shared range’s argument name in the diagnostic so plain Tensor inputs
produce a clear error instead of an internal check failure.

🪄 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: c8de8f9a-91aa-4b41-9f1a-0e9e8d4c9214

📥 Commits

Reviewing files that changed from the base of the PR and between 672871b and 72bfa69.

📒 Files selected for processing (21)

• docs/en/dev/distributed_ops.md
• docs/en/dev/passes/39-materialize_comm_domain_scopes.md
• docs/en/dev/passes/40-lower_host_tensor_collectives.md
• docs/zh/dev/distributed_ops.md
• docs/zh/dev/passes/39-materialize_comm_domain_scopes.md
• docs/zh/dev/passes/40-lower_host_tensor_collectives.md
• python/pypto/runtime/builtins/collectives/all_to_all_v/__init__.py
• python/pypto/runtime/builtins/collectives/all_to_all_v/templates/entry.cpp.in
• python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel.cpp.in
• python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel_config.py.in
• src/codegen/distributed/distributed_ops_codegen.cpp
• src/ir/op/distributed/collective.cpp
• src/ir/transforms/lower_composite_ops_pass.cpp
• src/ir/transforms/lower_host_tensor_collectives_pass.cpp
• src/ir/transforms/materialize_comm_domain_scopes_pass.cpp
• tests/st/distributed/test_l3_host_tensor_all_to_all_v.py
• tests/ut/codegen/distributed/test_host_orch_distributed.py
• tests/ut/ir/test_distributed_ops.py
• tests/ut/ir/transforms/test_lower_composite_ops.py
• tests/ut/ir/transforms/test_lower_host_tensor_collectives.py
• tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py

💤 Files with no reviewable changes (1)

• src/ir/transforms/lower_composite_ops_pass.cpp

[Copilot]

Pull request overview

Copilot reviewed 21 out of 21 changed files in this pull request and generated 1 comment.

Suppressed comments (3)

src/ir/transforms/lower_host_tensor_collectives_pass.cpp:213

• This only verifies the IR type, not that the value is actually window-backed. A HOST parameter typed as DistributedTensor has no window_buffer_; when used as input or send_counts, it passes this check and then triggers GetWindowBuffer's internal-invariant failure. Validate window_buffer_ here as well so all documented non-window-bound inputs get a user-facing error.

void CheckHostWindowBoundArg(const ExprPtr& expr, const char* op_name, const char* role) {
  CHECK_SPAN(As<DistributedTensorType>(expr->GetType()) != nullptr, expr->span_)
      << op_name << " " << role
      << " must be a window-bound DistributedTensor when called from a HOST orchestrator "
         "(a plain Tensor is only supported on the InCore composite path)";

src/ir/transforms/lower_host_tensor_collectives_pass.cpp:292

• This range is not guaranteed by the public deducer: a zero-row target yields max_recv_i64 == 0, and a sufficiently large static target can exceed INT32_MAX. Both are user-reachable HOST inputs, so reporting them as an internal compiler invariant is incorrect. Use CHECK_SPAN to produce the intended user-facing validation error.

  INTERNAL_CHECK_SPAN(max_recv_i64 > 0 && max_recv_i64 <= static_cast<int64_t>(INT32_MAX), call->span_)

docs/en/dev/passes/40-lower_host_tensor_collectives.md:51

• The added paragraph now runs directly into the existing lowercase “the pass emits” sentence. Add a paragraph break and capitalize the sentence.

`pld.system.notify`) — all five window args must resolve into the same
`CommDomainScopeStmt`.

[Copilot]

Pull request overview

Copilot reviewed 21 out of 21 changed files in this pull request and generated 1 comment.

Suppressed comments (4)

src/ir/transforms/lower_host_tensor_collectives_pass.cpp:228

• Only control-to-control aliasing is rejected. If input aliases signal or recv_counts, a faster peer's notify can overwrite source rows before this rank TPUTs them; if target aliases any control window, incoming TPUTs race with count reads/notifies or can satisfy/corrupt the barrier. Reject those data/control allocation pairs as well.

void CheckDistinctControlWindows(const CallPtr& call, const char* op_name) {
  auto signal_wb = GetWindowBuffer(call->args_[2], "signal");
  auto counts_wb = GetWindowBuffer(call->args_[3], "send_counts");
  auto recv_wb = GetWindowBuffer(call->args_[4], "recv_counts");
  CHECK_SPAN(signal_wb.get() != recv_wb.get(), call->span_)

src/ir/transforms/lower_host_tensor_collectives_pass.cpp:212

• This only checks the IR type, not whether the value is actually window-bound. A DistributedTensor HOST parameter has no materialized window_buffer_, so it passes this check and then GetWindowBuffer raises an internal-invariant error. Include the back-reference in this user-facing validation.

void CheckHostWindowBoundArg(const ExprPtr& expr, const char* op_name, const char* role) {
  CHECK_SPAN(As<DistributedTensorType>(expr->GetType()) != nullptr, expr->span_)
      << op_name << " " << role
      << " must be a window-bound DistributedTensor when called from a HOST orchestrator "
         "(a plain Tensor is only supported on the InCore composite path)";

tests/st/distributed/test_l3_host_tensor_all_to_all_v.py:227

• These counts never exceed MAX_RECV (nr - d is at most 4), so the new kernel's runtime clamp branch is not exercised end-to-end. Include at least one count above capacity and assert that recv_counts and consumed rows use min(count, MAX_RECV).

        for r in range(nr):
            for d in range(nr):
                n_rows = nr - d  # variable send count, same golden formula as the InCore ST
                send_counts[r, d, 0] = n_rows

docs/en/dev/passes/40-lower_host_tensor_collectives.md:51

• The paragraph now ends at CommDomainScopeStmt, leaving the next line as the lowercase sentence fragment “the pass emits…”. Start a new sentence (or join it to this paragraph) so the pass behavior reads grammatically.

`all_to_all_v` additionally requires `send_counts` (window-bound at this
layer, LOCAL-only) and `recv_counts` (window-bound, published cross-rank via
`pld.system.notify`) — all five window args must resolve into the same
`CommDomainScopeStmt`.

[Copilot]

Pull request overview

Copilot reviewed 21 out of 21 changed files in this pull request and generated no new comments.

Suppressed comments (7)

docs/en/dev/passes/40-lower_host_tensor_collectives.md:51

• Capitalize the sentence after the paragraph break.

the pass emits the corresponding `builtin.tensor.*` dispatch per participating

src/codegen/distributed/distributed_ops_codegen.cpp:347

• The correctness reason for including max_recv in the variant is not covered by the new single-call codegen test: that test would still pass if global deduplication collapsed different capacities. Add one program containing two HOST all_to_all_v calls with different MAX_RECV values and assert that both next-level variants are emitted with their respective kMaxRecv substitutions.

  const std::string variant =
      op->op_->name_ + "__maxrecv" + std::to_string(max_recv) + "__" + Fp32VariantSuffix(dtype);

  if (dist_codegen->MarkBuiltinEmitted(variant)) {
    dist_codegen->RecordBuiltinNextLevel(
        op, variant, {{"max_recv_cpp", std::to_string(max_recv)}, {"dtype_cpp", Fp32TypeCpp(dtype)}});

src/ir/op/distributed/collective.cpp:588

• This description remains internally contradictory: its opening at line 578 says each rank pushes send_counts[dest] rows, while these lines say every rank physically pushes the full MAX_RECV block. Rewrite the opening in terms of logically valid rows versus physical transfer size. The same obsolete contract remains in python/pypto/language/distributed/op/tensor_ops.py:825-853, python/pypto/ir/op/distributed/tensor_ops.py:399-407, and lower_composite_ops_pass.cpp:1787-1825, so those primary API docstrings/comments also need synchronization with the new HOST path.

        "push always transfers the full MAX_RECV-row capacity block per "
        "destination (a compile-time-sized ``pld.tile.put``, independent of "
        "the runtime count) — rows beyond a sender's actual count still cross "
        "the wire, but the receiver skips them using ``recv_counts`` "
        "(MPI_Alltoallv semantics apply to the logical result, not the wire "
        "transfer).  During the same push phase each rank also publishes "

docs/en/dev/distributed_ops.md:312

• The implementation rejects aliasing between every pair of the five operands, but this only documents input/target. State the full user-visible restriction, including the otherwise non-obvious read-only input/send_counts pair.

`builtin.tensor.all_to_all`. `input` and `send_counts` must both be
window-bound `DistributedTensor`s at this layer (narrower than the composite's
`AsTensorTypeLike`, forced by the HOST dispatch codegen, which only supports
window-bound or tile args) — `input` must be distinct from `target`, same
discipline as the symmetric `all_to_all` builtin. `MAX_RECV` is mangled into

docs/en/dev/passes/40-lower_host_tensor_collectives.md:50

• The implementation requires all five all_to_all_v operand windows to use pairwise-distinct allocations, not only distinct input and target windows. Document that restriction here so the pass contract matches CheckAllToAllVDistinctWindows.

`all_to_all_v` additionally requires `send_counts` (window-bound at this
layer, LOCAL-only) and `recv_counts` (window-bound, published cross-rank via
`pld.system.notify`) — all five window args must resolve into the same
`CommDomainScopeStmt`.

docs/zh/dev/distributed_ops.md:274

• 实现会拒绝五个操作数窗口中任意一对共享 allocation，但这里仅记录了 input 与 target 必须不同。请补充完整的两两不同约束，使文档与 CheckAllToAllVDistinctWindows 一致。

在这一层，`input` 与 `send_counts` 都必须是窗口绑定的 `DistributedTensor`（比
composite 的 `AsTensorTypeLike` 更严格，这是 HOST 派发代码生成强制要求的——
它只支持窗口绑定或 tile 参数）——`input` 必须与 `target` 是不同的窗口，与对称
`all_to_all` builtin 的约束相同。`MAX_RECV` 会被混入代码生成的 variant 字符串

docs/zh/dev/passes/40-lower_host_tensor_collectives.md:47

• 实现要求 all_to_all_v 的五个操作数窗口使用两两不同的 allocation，而不仅是 input 与 target 不同。请在此明确该约束，使 pass 文档与 CheckAllToAllVDistinctWindows 一致。

结果窗口。`all_to_all_v` 还额外要求 `send_counts`（在这一层是窗口绑定的，
仅本地使用）和 `recv_counts`（窗口绑定，通过 `pld.system.notify` 跨 rank
发布）——五个窗口参数都必须位于同一个 `CommDomainScopeStmt` 中。

[georgebisbas]

Addressed the suppressed-comment findings from Copilot's latest review round (never posted as inline threads, so nothing to resolve there) in 8b630dbb:

• Fixed the self-contradictory MAX_RECV wording ("only send_counts[dest] rows cross the wire" / "rows beyond count never written") across collective.cpp, lower_composite_ops_pass.cpp (which contradicted its own code below it), and both Python tensor_ops.py docstrings.
• Documented the all-5-pairwise-distinct-window and exact-static-signal-capacity requirements in distributed_ops.md and 40-lower_host_tensor_collectives.md (en + zh), and fixed the capitalization/paragraph-break nit in the latter.
• Added a codegen test asserting two all_to_all_v calls with different MAX_RECV emit two distinct next-level variants rather than collapsing into one mis-instantiated kernel.

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.

Suppressed comments (3)

src/ir/transforms/lower_host_tensor_collectives_pass.cpp:432

• These operands can materialize into different comm-domain scopes when their staging/consume dispatches use different explicit device subsets: input and send_counts retain their own coverage, while signal and recv_counts inherit target's. In that user-reachable case FindScopeForBuffers returns null and LowerCollective reports an INTERNAL_CHECK_SPAN compiler failure. Since the docs make same-domain membership a user requirement, this path should issue an actionable CHECK_SPAN error (or materialization should deliberately unify the five operands' coverage).

            return std::vector<WindowBufferPtr>{
                GetWindowBuffer(call->args_[0], "all_to_all_v input"),
                GetWindowBuffer(call->args_[1], "all_to_all_v target"),
                GetWindowBuffer(call->args_[2], "all_to_all_v signal"),
                GetWindowBuffer(call->args_[3], "all_to_all_v send_counts"),

python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel.cpp.in:124

• Only applying the upper bound lets a negative runtime count be published unchanged into the peer's recv_counts. That value is documented as a valid-row count and is later cast to INDEX/used as a loop bound, so a negative value violates the API contract and may miscompile consumers that rely on INDEX being non-negative. Clamp to [0, kMaxRecv] and make the equivalent change in LowerTensorAllToAllVRule so HOST and InCore remain aligned.

    int32_t raw_count = send_counts_base[dest];
    int64_t rows64 = static_cast<int64_t>(raw_count);
    if (rows64 > kMaxRecv) rows64 = kMaxRecv;
    int32_t rows = static_cast<int32_t>(rows64);

src/ir/op/distributed/collective.cpp:1107

• This argument description still says recv_counts records how many rows were sent, but the new kernel physically sends all MAX_RECV rows. The value is the clamped number of logically valid leading rows; describing it that way avoids contradicting the revised operation description.

    .add_argument("recv_counts",
                  "Window-bound INT32 DistributedTensor [NR, 1] — after the barrier, "
                  "recv_counts[src, 0] holds how many rows src sent to this rank (InOut)")