| |
| """``TtDiffusionPlanner``: the whole plan of Diffusion Planner v5.0 in one trace on a Blackhole p150. |
| |
| Variant ``plan`` (``ttaw.trace.TraceRunner``; inputs :data:`tt.inputs.INPUT_SPECS`, written every plan): |
| encoder + fusion (:mod:`.encoder`) -> cross K / V of the 3 DiT blocks hoisted once -> 11 x (DiT evaluation with |
| the per-step folded adaLN rows + fp32 DPM-Solver++(2M) update + prefix constraint) (:mod:`.decoder`) -> turn head -> |
| one packed readback (``final_x0`` ``[352, 324]`` fp32, ``logit`` ``[5]``, ``ego_steps`` ``[11, 324]``). No host |
| fallback inside the plan, so there are no trace segments. |
| |
| Debug variants (``debug=True``; tests only, also replayed from traces): |
| |
| - ``encoder_taps``: the encoder of ``plan`` returning every encoder tap of ``reference.model.TAP_NAMES``; |
| - ``decode_once``: one decoder evaluation on teacher-forced inputs (``dbg_x`` ``[352, 324]``, ``dbg_enc`` |
| ``[576, 256]``: the reference encoding + 12 zero pad-token rows) with the per-step rows as inputs |
| (``dbg.<i>.<key>``, ``dbg.final.g`` / ``.b``), so one trace serves all 11 evaluation times. |
| |
| Everything here runs under the model lock of ``api.DiffusionPlanner`` (one chip, batch 1). |
| """ |
| from __future__ import annotations |
|
|
| import time |
| from typing import Any, Dict, List, Optional, Sequence |
|
|
| import numpy as np |
|
|
| from ..reference import config as C |
| from ..ttaw.ops import attention as A |
| from ..ttaw.trace import TraceRunner, pack_outputs |
| from . import config as T |
| from . import inputs as I |
| from . import params as P |
| from .decoder import STEP_KEYS, TtDecoder, TtTurnHead |
| from .encoder import TtEncoder |
| from .layers import Build, policy |
|
|
| __all__ = ["TtDiffusionPlanner", "agent_buckets", "needed_rows", "bucket_rows"] |
|
|
|
|
| def agent_buckets() -> tuple: |
| """``COMPACT``: the agent buckets of ``AGENT_BUCKETS`` (multiples of 32 below 352, sorted); () when off.""" |
| k = T.KNOBS.read() |
| if not k.COMPACT: |
| return () |
| out = sorted({int(v) for v in str(k.AGENT_BUCKETS).replace(" ", "").split(",") if v}) |
| bad = [v for v in out if v % T.TILE or not 0 < v < T.AGENTS] |
| if bad: |
| raise ValueError(f"DIFFUSION_PLANNER_AGENT_BUCKETS: {bad} are not multiples of {T.TILE} below {T.AGENTS}") |
| return tuple(out) |
|
|
|
|
| def needed_rows(prepared: Any) -> int: |
| """1 + the last decoder row a plan reads or attends to: the ego (row 0), the valid self-attention keys |
| (``agent_valid``) and the emitted neighbours (``neighbor_rows`` + 1). The rows past it are masked keys whose |
| outputs are never read, so dropping them leaves the needed rows' values unchanged.""" |
| valid = np.flatnonzero(np.asarray(prepared.decoder.agent_valid, bool)) |
| last = int(valid.max()) if valid.size else 0 |
| tok = np.flatnonzero(np.asarray(prepared.features.valid["neighbor"], bool)) |
| if tok.size: |
| last = max(last, int(tok.max()) + 1) |
| emitted = np.asarray(prepared.neighbor_rows) |
| if emitted.size: |
| last = max(last, int(emitted.max()) + 1) |
| return last + 1 |
|
|
|
|
| def bucket_rows(prepared: Any, buckets: Sequence[int]) -> int: |
| n = needed_rows(prepared) |
| return next((int(b) for b in sorted(buckets) if b >= n), T.AGENTS) |
|
|
|
|
| class TtDiffusionPlanner: |
| """Weights on the device, the ``TraceRunner`` with its persistent inputs and variants, and the host glue. |
| |
| ``weights``: ``reference.weights.PlannerWeights``; ``precision``: extra policy rules (``"dec.*=HiFi2+fp32"``) |
| on top of ``DIFFUSION_PLANNER_PRECISION``; ``ln_fp32`` / ``hidden_fp32`` / ``split`` / ``attn_fp32_acc`` / |
| ``attn_matmul``: module globs overriding the knobs of the same names (``tt/config.py``); ``debug``: also register |
| ``encoder_taps`` / ``decode_once``.""" |
|
|
| def __init__(self, device: Any, weights: Any, *, debug: bool = False, precision: Optional[str] = None, |
| ln_fp32: Optional[Sequence[str]] = None, hidden_fp32: Optional[Sequence[str]] = None, |
| split: Optional[Sequence[str]] = None, attn_fp32_acc: Optional[Sequence[str]] = None, |
| attn_matmul: Optional[Sequence[str]] = None, steps: int = C.DPM_SOLVER_STEPS, |
| num_command_queues: Optional[int] = None): |
| import ttnn |
|
|
| t0 = time.perf_counter() |
| self.device = device |
| self.build = Build(device, policy(spec=precision), ln_fp32=ln_fp32, hidden_fp32=hidden_fp32, split=split, |
| attn_fp32_acc=attn_fp32_acc, attn_matmul=attn_matmul) |
| p = weights.params |
| self.tables = P.step_tables(p, steps) |
| self.buckets = agent_buckets() if hasattr(device, "compute_with_storage_grid_size") else () |
| self.encoder = TtEncoder(self.build, p, nb_rows=self.buckets if self.build.compact_enc else ()) |
| self.decoder = TtDecoder(self.build, p, self.tables, rows=(T.AGENTS,) + self.buckets) |
| self.turn = TtTurnHead(self.build, p) |
| self.runner = TraceRunner(device, num_command_queues=num_command_queues, name="diffusion-planner") |
| warm = I.warmup_inputs() |
| for name, shape in I.INPUT_SPECS.items(): |
| dtype = "bfloat16" if name in I.BF16_INPUTS else "float32" |
| self.runner.add_input(name, init=warm[name], dtype=dtype, layout=ttnn.TILE_LAYOUT) |
| self.cs_pad = {} |
| if I.CS_COLS < T.STATE_COLS: |
| for r in (T.AGENTS,) + self.buckets: |
| self.cs_pad[r] = self.build.upload(np.zeros((r, T.STATE_COLS - I.CS_COLS), np.float32), "float32") |
| self.trim = I.CS_COLS < T.STATE_COLS |
| self._zero_inputs: set = set() |
| self.runner.add_variant("plan", self._plan) |
| for r in self.buckets: |
| self.runner.add_variant(f"plan_r{r}", lambda ctx, r=r: self._plan(ctx, r)) |
| self.debug = bool(debug) |
| if self.debug: |
| self.runner.add_input("dbg_x", init=np.zeros((1, 1, T.AGENTS, T.STATE_COLS), np.float32), |
| dtype="float32", layout=ttnn.TILE_LAYOUT) |
| self.runner.add_input("dbg_enc", init=np.zeros((1, 1, T.TOKENS, C.HIDDEN_DIM), np.float32), |
| dtype=self.encoder.fstream, layout=ttnn.TILE_LAYOUT) |
| for name in self._dbg_row_names(): |
| self.runner.add_input(name, init=np.zeros((1, 1, 1, C.HIDDEN_DIM), np.float32), dtype="float32", |
| layout=ttnn.TILE_LAYOUT) |
| self.runner.add_variant("encoder_taps", self._encoder_taps) |
| self.runner.add_variant("decode_once", self._decode_once) |
| self.build_ms = (time.perf_counter() - t0) * 1e3 |
|
|
| |
| def _plan(self, ctx, rows: int = T.AGENTS): |
| """The whole plan; ``rows`` < 352 (``COMPACT``): the decoder on the first ``rows`` agents (the inputs' |
| leading rows, sliced in the trace).""" |
| import ttnn |
|
|
| enc = self.encoder.forward(ctx, nb=rows if rows != T.AGENTS else None) |
| kv = self.decoder.cross_kv(enc) |
| y0, cs, key = ctx["y0"], ctx["cs"], ctx["agent_key_row"] |
| if rows != T.AGENTS: |
| y0 = ttnn.slice(y0, [0, 0, 0, 0], [1, 1, rows, T.STATE_COLS]) |
| cs = ttnn.slice(cs, [0, 0, 0, 0], [1, 1, rows, int(cs.shape[-1])]) |
| key = ttnn.slice(key, [0, 0, 0, 0], [1, 1, 1, rows]) |
| if int(cs.shape[-1]) < T.STATE_COLS: |
| cs = ttnn.concat([cs, self.cs_pad[rows]], dim=-1) |
| self_mask = A.expand_key_bias(key, rows, memory_config=self.build.attn_mem()) |
| final, ego = self.decoder.solve(y0, cs, kv, self_mask) |
| logit = self.turn(final, enc) |
| return pack_outputs({"final_x0": final, "logit": logit, "ego_steps": ego}) |
|
|
| def _encoder_taps(self, ctx): |
| taps: Dict[str, Any] = {} |
| self.encoder.forward(ctx, taps) |
| return taps |
|
|
| @staticmethod |
| def _dbg_row_names() -> List[str]: |
| return [f"dbg.{i}.{k}" for i in range(C.DIT_DEPTH) for k in STEP_KEYS] + ["dbg.final.g", "dbg.final.b"] |
|
|
| def _decode_once(self, ctx): |
| rows = {"blocks": [{k: ctx[f"dbg.{i}.{k}"] for k in STEP_KEYS} for i in range(C.DIT_DEPTH)], |
| "g": ctx["dbg.final.g"], "b": ctx["dbg.final.b"]} |
| kv = self.decoder.cross_kv(ctx["dbg_enc"]) |
| self_mask = A.expand_key_bias(ctx["agent_key_row"], T.AGENTS) |
| return self.decoder.evaluate(ctx["dbg_x"], rows, kv, self_mask) |
|
|
| |
| def capture(self) -> None: |
| self.runner.capture() |
|
|
| def filter_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]: |
| """``INPUT_TRIM``: drop the inputs that are all +0.0 while their device buffer already holds the zeros of this |
| model's previous upload (a value test, not a cache of the previous request: any non-zero array is always |
| uploaded). Every input passed on is recorded as uploaded.""" |
| if not self.trim: |
| return inputs |
| out = {} |
| for name, a in inputs.items(): |
| arr = a if isinstance(a, np.ndarray) else None |
| zero = arr is not None and arr.dtype == np.float32 and not arr.view(np.uint32).any() |
| if zero and name in self._zero_inputs: |
| continue |
| out[name] = a |
| if zero: |
| self._zero_inputs.add(name) |
| else: |
| self._zero_inputs.discard(name) |
| return out |
|
|
| def _run(self, variant: str, inputs: Dict[str, Any], eager: bool): |
| inputs = self.filter_inputs(inputs) |
| return self.runner.run_eager(variant, inputs=inputs) if eager else self.runner(variant, inputs=inputs) |
|
|
| def rows_for(self, prepared: Any) -> int: |
| """The decoder rows a plan needs (``COMPACT``): the smallest agent bucket holding the ego, every valid |
| self-attention key and every emitted neighbour row; 352 without buckets or when none is large enough.""" |
| return bucket_rows(prepared, self.buckets) |
|
|
| def variant_for(self, prepared: Any) -> str: |
| r = self.rows_for(prepared) |
| return "plan" if r == T.AGENTS else f"plan_r{r}" |
|
|
| def unpack(self, out: Dict[str, Any]) -> np.ndarray: |
| """The readback's ``final_x0`` (``[R, 324]``) -> ``[321, 324]``; the rows past R (not computed: masked keys, |
| never emitted) are zero.""" |
| f = np.asarray(out["final_x0"], np.float32).reshape(-1, T.STATE_COLS) |
| if f.shape[0] < T.AGENTS: |
| f = np.concatenate([f, np.zeros((T.AGENTS - f.shape[0], T.STATE_COLS), np.float32)], 0) |
| return f[:C.MAX_NUM_AGENTS] |
|
|
| def forward(self, prepared: Any, *, ego_steps: bool = True, eager: bool = False, |
| variant: Optional[str] = None) -> Dict[str, Any]: |
| """One plan: upload, replay, read (``eager=True``: the same graph without the trace, for bring-up and |
| replay-vs-eager checks). ``final_x0`` ``[321, 81, 4]`` (normalised, prefix-constrained), ``logit`` ``[5]``, |
| ``denoising_steps``: the 11 iterates' ego rows as ``[1, 81, 4]`` arrays. ``variant``: force a plan variant |
| (default :meth:`variant_for`).""" |
| out = self._run(variant or self.variant_for(prepared), I.plan_inputs(prepared), eager) |
| final = self.unpack(out) |
| res = {"final_x0": final.reshape(C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM).astype(np.float32), |
| "logit": out["logit"].reshape(-1)[:C.TURN_INDICATOR_OUTPUT_DIM].astype(np.float32)} |
| if ego_steps: |
| steps = out["ego_steps"].reshape(-1, T.STATE_COLS) |
| res["denoising_steps"] = [s.reshape(1, C.OUTPUT_T + 1, C.POSE_DIM).astype(np.float32) for s in steps] |
| return res |
|
|
| def encoder_taps(self, prepared: Any, *, eager: bool = False) -> Dict[str, np.ndarray]: |
| """Replay of ``encoder_taps``: ``{tap: array}`` with the reference's shapes (``[E, 256]`` category rows, |
| ``[E, 64, 128]`` mixer taps, ``[564, 256]`` token / fusion / encoding rows).""" |
| if not self.debug: |
| raise RuntimeError("encoder_taps needs TtDiffusionPlanner(debug=True)") |
| raw = self._run("encoder_taps", I.plan_inputs(prepared), eager) |
| out = {} |
| for name, a in raw.items(): |
| a = np.asarray(a, np.float32) |
| if name.endswith((".pre", ".mixer")): |
| out[name] = a.reshape(-1, C.MIXER_TOKENS, C.MIXER_CHANNELS) |
| elif name in ("enc.tokens", "enc.encoding") or name.startswith("enc.fusion."): |
| out[name] = a.reshape(-1, C.HIDDEN_DIM)[:T.TOKENS_REAL] |
| else: |
| out[name] = a.reshape(-1, C.HIDDEN_DIM) |
| return out |
|
|
| def decode_once(self, prepared: Any, x: np.ndarray, t: float, *, encoding: np.ndarray, |
| eager: bool = False) -> np.ndarray: |
| """Replay of ``decode_once`` at evaluation time ``t`` (one of ``tables.eval_times``) on the teacher-forced |
| ``x`` ``[321, 81, 4]`` (prefix-constrained) and ``encoding`` ``[564, 256]`` -> ``[321, 81, 4]`` (the t = 0 |
| slot is 0: the masked projection).""" |
| if not self.debug: |
| raise RuntimeError("decode_once needs TtDiffusionPlanner(debug=True)") |
| k = int(np.argmin([abs(float(t) - e) for e in self.tables.eval_times])) |
| if abs(float(t) - self.tables.eval_times[k]) > 1e-6: |
| raise ValueError(f"t={t} is not an evaluation time of the solver plan {self.tables.eval_times}") |
| base = I.plan_inputs(prepared) |
| enc = P.pad_rows(np.asarray(encoding, np.float32).reshape(T.TOKENS_REAL, C.HIDDEN_DIM), T.TOKENS) |
| inputs = {"agent_key_row": base["agent_key_row"], "dbg_x": I.decoder_state(x, x[:, 0]), |
| "dbg_enc": enc.reshape(1, 1, T.TOKENS, C.HIDDEN_DIM)} |
| for i, blk in enumerate(self.tables.blocks[k]): |
| for key in STEP_KEYS: |
| inputs[f"dbg.{i}.{key}"] = np.asarray(blk[key], np.float32).reshape(1, 1, 1, -1) |
| inputs["dbg.final.g"] = np.asarray(self.tables.final[k]["g"], np.float32).reshape(1, 1, 1, -1) |
| inputs["dbg.final.b"] = np.asarray(self.tables.final[k]["b"], np.float32).reshape(1, 1, 1, -1) |
| out = self._run("decode_once", inputs, eager) |
| flat = np.asarray(out, np.float32).reshape(T.AGENTS, T.STATE_COLS)[:C.MAX_NUM_AGENTS] |
| return flat.reshape(C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM) |
|
|
| def trace_buffers_mb(self) -> Optional[float]: |
| """DRAM held by the captured traces' command buffers: the allocated bytes of the device's TRACE region |
| (``ttnn.get_memory_view``; all banks), None when the view is unavailable.""" |
| import ttnn |
|
|
| try: |
| view = ttnn.get_memory_view(self.device, ttnn.BufferType.TRACE) |
| return round(view.num_banks * view.total_bytes_allocated_per_bank / 2 ** 20, 2) |
| except (AttributeError, RuntimeError, TypeError): |
| return None |
|
|
| def describe(self) -> Dict[str, Any]: |
| return {"trace": self.runner.describe(), "trace_buffers_mb": self.trace_buffers_mb(), |
| "precision": self.build.policy.describe(), |
| "options": self.build.options(), |
| "uploaded_mb": round(self.build.uploaded_bytes / 2 ** 20, 2), "build_ms": round(self.build_ms, 1), |
| "tokens": T.TOKENS, "agents": T.AGENTS, "agent_buckets": list(self.buckets), "nfe": self.tables.nfe} |
|
|
| def release(self) -> None: |
| self.runner.release() |
|
|
