Module tensor_parallel 

════════════════════════════════════════════════════════════════════════════════════════════════ TENSOR PARALLELISM FOR A WHOLE PROGRAM — Megatron over RDUs, as ONE spatial program.

passes::shard splits ONE matmul’s contraction over chips and sums the partials. A model layer is a chain of ops, and Megatron splits the chain, not an op: the projections into a layer are split by COLUMN (each chip owns whole heads, whole FFN columns — its outputs are complete slices), the projections out of it by ROW (each chip holds a PARTIAL of the full output), and the partials are summed by an all-reduce, twice a layer. Between the all-reduces every chip works on its own heads alone; the residuals and norms are the replicated spine.

★ THE DECOMPOSITION IS A PROGRAM, SO IT IS WRITTEN AS ONE. A generator (qwen38::gen_qwen38_tp_decode_*) writes the tensor-parallel layer as ordinary LOGOS: every chip’s slice spelled with its own tensors (Wq3_c1), and every all-reduce as the left fold of plus it means. That source RUNS on the VM, so the decomposition itself is checked against the unsharded layer before the compiler ever sees it. Alongside it comes a TpManifest: the chip of every chip-local tensor, the slice of the full checkpoint tensor each one is (the host’s loading map), and each collective’s partials and fold statements.

★ THE COMPILER THEN LOWERS IT AS IT LOWERS ANYTHING, with two differences. Every section is RELOCATED to the chip of what it computes (relocate); and a collective’s fold statements are not lowered as additions on one chip — ONE allreduce_section replaces them, reading each chip’s partial where it lives and folding them all. The certificate’s shard obligation holds the result to the rule that makes the split mean anything: data derived from a chip’s own tensors reaches the replicated spine only through an all-reduce over EVERY chip. ════════════════════════════════════════════════════════════════════════════════════════════════

Structs

ChipTransfer

A tensor one chip’s statement reads and another chip’s statement writes: it crosses from the writer’s chip to the reader’s, once, before the reader’s first statement that reads it.

SectionTouch

What a section touches, as far as its phase is concerned: the tensors it reads and writes, by name, and its chips.

TensorParallelShim

What a chip allows of a tensor-parallel program: what base allows, less a fused op over more than one chip.

TpCollective

One all-reduce: out = Σ partials, the partials in chip order (partials[c] lives on chip c), spelled in the source as the fold statements folds — the last of which binds out.

TpManifest

How a tensor-parallel program is spread over chips RDUs.

TpSlice

A chip’s tensor as ranges of a full tensor of (of_shape whole) along one axis, concatenated in order: each (start, len) is len rows (or columns) from start. What a host loads onto the chip, and how a test fills the slice with the full tensor’s values.

Enums

RelayStep

One step of the relay a host runs between chips’ programs that halt at their LinkPoints.

SliceAxis

Which axis of a full tensor a chip’s slice cuts.

TpReject

Why a tensor-parallel lowering refused.

Functions

allreduce_section

ONE all-reduce as a section: out = Σ_c partials[c] over a rows × cols tensor, tile by tile — chip c’s partial Loaded on chip c (where it lives) and handed across the chip boundary, all of them folded by an AllReduce on chip 0, the sum stored as the replicated out.

link_points

Where chip chip’s program p halts for transfers: each tensor it sends, after the last section that stores it; each it receives, before the first section that loads it — None for the first section, the host’s to put there before the chip starts. Halts at one boundary are one LinkPoint; they come in program order.

phase_order

Place each section, in source order, in the earliest phase after every section it depends on (a read of what an earlier one writes, a write of what an earlier one reads or writes — the carried state included) where its chips are free, and return the sections’ order, phase by phase, with barrier_after[k] closing the last section of each phase.

phases

★ THE CHIPS RUN AT ONCE. The source writes chip 0’s slice of an op, then chip 1’s; composed in that order with a barrier after each, chip 1 would wait out chip 0’s whole slice. Two sections may share a PHASE — run with no barrier between them — when they are on disjoint chips and neither touches what the other writes. See phase_order.

relay_schedule

The host’s relay over chips’ halts: round by round, each chip that can is run to its next halt — once the tensors the halt it stands at receives have been sent and copied — so chips whose next work needs nothing of each other run at once. Refused, with where each chip stands, when no chip can move before all have finished.

relocate

A scheduled one-chip section of the statement binding binder, moved onto the chip it runs on: every site it names — its stages’, buffers’, placement’s, memory map’s, semaphores’, engine programs’ and the remote semaphore posts in them — and every tensor it touches marked with its chip: the one m gives it, or, for a tensor the section itself stores on behalf of a chip-local statement (a blocked op’s intermediate), that chip; the replicated spine stays unmarked. The section computes exactly what it did; it now does so on another RDU. Its certificate is dropped with the schedule digest it was bound to.

transfers

The tensors that cross between m’s chips when each statement runs on TpManifest::chip_running its binder: for each statement, every operand the latest earlier statement to write it wrote on another chip, in program order, each (tensor, reader's chip) once. An operand no statement writes — a table, a host-written input — crosses nothing: each chip that reads it holds its own.