Pass formulae and API¶

This page renders the source documentation for every public canonical pass. The equations are part of the class docstrings, so changing an implemented derivation and changing its code reference happen in the same review surface.

Coverage and provenance¶

Pass	Boundary	Formula provenance	Commit evidence
`DistributeTransformerTrainingPass`	Model → distributed task	Megatron-LM; selective recomputation	executable lineage rules + deterministic replay
`DistributeTransformerInferencePass`	Model → distributed task	Transformer; Megatron-LM; FlashAttention boundary	executable lineage rules + exact work tests
`PlanTransformerTrainingPass`	Distributed task → portable plan	conservation of the derived Transformer work vector	executable work-equality rules
`PlanTransformerInferencePass`	Distributed task → portable plan	work conservation; conservative unfused workspace	executable work-equality rules
`BindReferenceQueueTargetPass`	Portable → concrete	internal deterministic contract; no paper claim	structural verifier + cross-boundary predicates
`BindReferenceSlotTargetPass`	Portable → concrete	internal deterministic contract; no paper claim	extension verifier + dependency-order proof

Formula versus performance evidence

FLOPs, bytes, payloads, shapes, and capacities below are canonical workload facts. Latency, efficiency, overlap, and uncertainty belong to evidence/cost views and are deliberately absent from these pass equations.

Training distribution¶

Bases: DerivationPass[ModelIR, DistributedTaskIR]

Expand one training block into logical TP-local and collective tasks.

The pass interprets a typed TP × PP × DP strategy, but this snapshot materializes one local tensor-parallel block only. Let B be the microbatch, S the sequence length, H the hidden width, F the feed-forward width, t the TP degree, and e bytes per element. A matrix product with shapes [m,n] × [n,k] contributes

\[ W_{\mathrm{gemm}} = 2mnk. \]

Therefore the local attention projections contribute

\[ W_{\mathrm{QKV}}=\frac{6BSH^2}{t},\qquad W_{\mathrm{attn\,matmul}}=\frac{4BS^2H}{t}, \]

and the two MLP projections contribute

\[ W_{\mathrm{MLP}}=\frac{4BSHF}{t}. \]

At a TP semantic boundary, the logical payload and local reduction work are

\[ M_{\mathrm{TP}}=BSH\,e,\qquad W_{\mathrm{reduce}}=BSH\frac{t-1}{t}. \]

The derivation emits explicit forward, recompute, activation-gradient, weight-gradient, optimizer, and collective invocations. It does not attach latency or choose a physical collective algorithm.

References

Shoeybi et al., Megatron-LM.
Narayanan et al., Efficient Large-Scale Language Model Training.
Korthikanti et al., Reducing Activation Recomputation.

run ¶

run(ir: ModelIR, context: PassContext) -> DistributedTaskIR

Source code in src/blueprinting/synthesizer/stages/distributed/passes.py

def run(self, ir: ModelIR, context: PassContext) -> DistributedTaskIR:
    return normalize_training_distribution(ir, context.session)

Inference distribution¶

Bases: DerivationPass[ModelIR, DistributedTaskIR]

Expand one inference phase into logical TP-local and collective tasks.

Let b be batch size, q the query-token count, c the visible KV context, H hidden width, F feed-forward width, h attention-head count, t TP degree, and e bytes per element. Prefill uses q=c; decode uses q=1. The exact local attention-core work is

\[ W_{\mathrm{attention}} = \frac{4bqcH}{t} + 5b\frac{h}{t}qc, \]

where the first term is QKᵀ plus PV and the second is the explicit softmax model. Projection and MLP work are

\[ W_{\mathrm{QKV}}=\frac{6bqH^2}{t},\quad W_{\mathrm{out}}=\frac{2bqH^2}{t},\quad W_{\mathrm{MLP}}=\frac{4bqHF}{t}. \]

The per-rank KV state retained by the derived analysis is

\[ C_{\mathrm{KV}}=2bc\frac{H}{t}e. \]

The pass marks KV reads/writes as effects and leaves fused/paged attention as target-neutral implementation alternatives rather than assuming them.

References

Vaswani et al., Attention Is All You Need.
Shoeybi et al., Megatron-LM.
Dao et al., FlashAttention.

run ¶

run(ir: ModelIR, context: PassContext) -> DistributedTaskIR

Source code in src/blueprinting/synthesizer/stages/distributed/passes.py

def run(self, ir: ModelIR, context: PassContext) -> DistributedTaskIR:
    return normalize_inference_distribution(ir, context.session)

Training portable planning¶

Bases: DerivationPass[DistributedTaskIR, PortablePlanIR]

Materialize exact training work without choosing a hardware target.

This pass is a semantics-preserving reification, not a second estimator. For every distributed invocation i it copies the exact work vector

\[ \mathbf{w}_i=(F_i,R_i,W_i,M_i) \]

into WorkloadFacts(operations, read_bytes, write_bytes, message_bytes). Resource requirements are non-lossy projections of that vector:

\[ Q_{\mathrm{compute}}=F_i,\qquad Q_{\mathrm{memory}}=R_i+W_i,\qquad Q_{\mathrm{network}}=M_i. \]

Exact block memory is reduced from the structurally derived layer facts; for example, stored activations are

\[ C_{\mathrm{act}}=\sum_{\ell} \left(A_\ell-O_\ell\,[\neg\mathrm{storeOutput}_\ell] -A_\ell\,[\neg\mathrm{storeActivation}_\ell]\right). \]

The executable pass rules verify work conservation and lineage before the snapshot is committed. The scientific provenance is inherited from the Transformer decomposition rather than introducing a new performance model.

References

Shoeybi et al., Megatron-LM.
Rajbhandari et al., ZeRO.
Korthikanti et al., Selective activation recomputation.

run ¶

run(ir: DistributedTaskIR, context: PassContext) -> PortablePlanIR

Source code in src/blueprinting/synthesizer/stages/portable_plan/passes.py

def run(self, ir: DistributedTaskIR, context: PassContext) -> PortablePlanIR:
    return normalize_training_plan(ir, context.session)

Inference portable planning¶

Bases: DerivationPass[DistributedTaskIR, PortablePlanIR]

Materialize exact inference work without target placement or timing.

Each inference invocation is mapped homomorphically into a portable task:

\[ (F_i,R_i,W_i,M_i)_{\mathrm{distributed}} =(F_i,R_i,W_i,M_i)_{\mathrm{portable}}. \]

The conservative workspace bound intentionally assumes an unfused score materialization until target binding selects an implementation. With boundary D=bqHe and local intermediate element counts 3bqH/t, b(h/t)qc, and bqF/t, it is

\[ C_{\mathrm{workspace}} =D+e\max\left(3bq\frac{H}{t}, b\frac{h}{t}qc, bq\frac{F}{t}\right). \]

This bound is a portable capacity obligation. A target implementation such as tiled exact attention may replace its workspace only during verified target binding; it may not rewrite the canonical operation count.

References

Vaswani et al., Attention Is All You Need.
Dao et al., FlashAttention.

run ¶

run(ir: DistributedTaskIR, context: PassContext) -> PortablePlanIR

Source code in src/blueprinting/synthesizer/stages/portable_plan/passes.py

def run(self, ir: DistributedTaskIR, context: PassContext) -> PortablePlanIR:
    return normalize_inference_plan(ir, context.session)

Reference queue binding¶

Bases: DerivationPass[PortablePlanIR, ConcretePlanIR]

Bind a portable plan to the deterministic queue-reference contract.

This is a contract reference implementation, not a paper-derived scheduler. Buffers are laid out in canonical order. For buffer i with size s_i and required alignment r_i, define a_i=max(16,r_i) and

\[ o_0=0,\qquad o_i=\left\lceil\frac{o_{i-1}+s_{i-1}}{a_i}\right\rceil a_i. \]

Hence o_i mod a_i = 0 and o_i ≥ o_{i-1}+s_{i-1}. Tasks retain topological order and are partitioned into compute, collective, transfer, and host queues; issue order is their stable subsequence in each queue. The commit gate re-evaluates this canonical construction and requires exact equality, so size, alignment, access mode, implementation, task category, dependencies, and target extension are checked as one derivation law.

References

Internal deterministic reference-target contract; no paper-derived performance or scheduling algorithm is claimed.

run ¶

run(ir: PortablePlanIR, context: PassContext) -> ConcretePlanIR

Source code in src/blueprinting/synthesizer/stages/concrete_plan/passes.py

def run(self, ir: PortablePlanIR, context: PassContext) -> ConcretePlanIR:
    return _queue_normal_form(ir, context.session)

Reference slot/dataflow binding¶

Bases: DerivationPass[PortablePlanIR, ConcretePlanIR]

Bind a portable plan to the deterministic slot/dataflow reference contract.

This pass shares the aligned buffer recurrence of BindReferenceQueueTargetPass. Given the portable topological task order T=(t_0,...,t_{n-1}), its reference issue assignment is

\[ \operatorname{issue}(t_i)=(\operatorname{cycle}=i, \operatorname{slot}=0). \]

Therefore every dependency t_j → t_i already proved by PortablePlanIR satisfies j<i and consequently cycle(t_j)<cycle(t_i). Transfers receive explicit source/destination route constraints. This construction exists to exercise typed target extensions and their verifier; it is not a throughput-optimal slot scheduler and has no claimed research-paper result.

References

Internal deterministic slot/dataflow contract; no paper-derived performance or scheduling algorithm is claimed.

run ¶

run(ir: PortablePlanIR, context: PassContext) -> ConcretePlanIR

Source code in src/blueprinting/synthesizer/stages/concrete_plan/passes.py

def run(self, ir: PortablePlanIR, context: PassContext) -> ConcretePlanIR:
    return _slot_normal_form(ir, context.session)

Research sources¶

Shoeybi et al., Megatron-LM.
Narayanan et al., Efficient Large-Scale Language Model Training Using Megatron-LM.
Vaswani et al., Attention Is All You Need.
Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models.
Dao et al., FlashAttention.
Rajbhandari et al., ZeRO.

These papers establish algorithmic provenance; repository verifiers and tests, not citation alone, establish what this implementation actually preserves.