Pass 公式与 API

本页渲染全部公开 canonical Pass 的源码文档。核心公式直接写在 class docstring 中，因此实现推导发生变化时，代码参考也必须在同一个 review 面中变化。

覆盖与来源

Pass	边界	公式来源	Commit 证据
DistributeTransformerTrainingPass	Model → distributed task	Megatron-LM、selective recomputation	executable lineage rule + deterministic replay
DistributeTransformerInferencePass	Model → distributed task	Transformer、Megatron-LM、FlashAttention implementation boundary	executable lineage rule + exact-work test
PlanTransformerTrainingPass	Distributed task → portable plan	已推导 Transformer work vector 的守恒映射	executable work-equality rule
PlanTransformerInferencePass	Distributed task → portable plan	work conservation、保守 unfused workspace	executable work-equality rule
BindReferenceQueueTargetPass	Portable → concrete	内部 deterministic contract，不声明论文算法	structural verifier + cross-boundary predicate
BindReferenceSlotTargetPass	Portable → concrete	内部 deterministic contract，不声明论文算法	extension verifier + dependency-order proof

公式与性能证据的边界

下列 FLOPs、bytes、payload、shape 和 capacity 是 canonical workload fact。Latency、efficiency、overlap 与 uncertainty 属于 evidence/cost view，不能混入这些 Pass 公式。

训练分布推导

该 Pass 从一个语义 Transformer block 推导 local TP compute、collective、recompute、反向与 optimizer invocation。PP/DP 保留在 typed strategy 中，不伪造尚未物化的跨 stage DAG。

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)

推理分布推导

该 Pass 分别处理 prefill 的 q=c 与 decode 的 q=1，显式推导 attention、projection、MLP、KV-cache 和 collective work。

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)

训练 PortablePlan 推导

该 Pass 不重新估算 work，而是证明 distributed invocation 的 operations/read/write/message 向量被无损写入 WorkloadFacts。

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)

推理 PortablePlan 推导

该 Pass 同时形成 persistent weight、KV state、boundary 和 conservative workspace obligation；target binding 之后才能用已选实现收紧 workspace。

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

这是用于验证 ConcretePlan contract 的确定性构造，不是性能最优 scheduler。公式给出 buffer alignment recurrence 与 stable queue subsequence。

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

这是用于验证 typed target extension 的确定性 cycle/slot 构造，并通过 PortablePlan 的拓扑顺序证明 dependency legality。

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)

研究来源

• 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。

论文用于说明算法来源，不自动证明实现正确；实际 preservation 由 verifier、negative test、deterministic replay 和 baseline gate 证明。