flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/beta/__init__.py

@@ -0,0 +1,43 @@
+"""Beta: the join between two bound structures, declared rather than implied.
+
+Every structure in the catalog declares its own boundary. What none of
+them declares is the *join* — what has to agree between a structure and
+the one it feeds. Today those agreements exist only inside each impl's
+code, so they are never checked, never negotiated, and never optimised
+across. Every join therefore defaults to "materialise a fresh tensor in
+the host's convention", which is what a hand-written runtime never does:
+there the author holds the whole dataflow and picks, per join, the same
+buffer, the same layout, no re-quantisation.
+
+The vocabulary here is derived, not designed. Each attribute exists
+because one specific join cost real time or gave a wrong verdict, and
+each carries that incident in its docstring. The evidence for the whole
+idea is one controlled comparison: ``dtype`` is the only join attribute
+that was ever declared and negotiated, and it is the only one that
+stopped causing failures. The other five were never declared and every
+one of them bit.
+
+**This is beta.** Three rules keep it honest, and they double as the
+conditions for deleting it:
+
+1. *Descriptive before prescriptive.* It must correctly describe the
+   joins that already work before it is allowed to change any behaviour.
+   A vocabulary that cannot express what the stack already does is the
+   wrong vocabulary.
+2. *No attribute without a negotiator.* An attribute that no join
+   negotiates and no measurement depends on gets deleted, not
+   documented.
+3. *Off by default.* Nothing in the main path consumes this unless asked
+   to, and an attribute a port does not declare is simply not negotiated
+   — so coverage can only grow.
+
+The deprecation signal is stated up front: if each new structure needs a
+new attribute, this is not a vocabulary, it is a junk drawer. The
+attribute set is supposed to converge as structures are added. If it
+does not, delete this package.
+"""
+
+from .negotiate import JoinRefused, negotiate
+from .ports import ATTRIBUTES, Join, Port
+
+__all__ = ["ATTRIBUTES", "Join", "JoinRefused", "Port", "negotiate"]

flashrt_structures/beta/conform.py

@@ -0,0 +1,94 @@
+"""Check the declarations against the impls they claim to describe.
+
+A spec file that nobody checks drifts away from the code and becomes
+worse than no spec, because it is then confidently wrong. Writing these
+ports down already corrected one of my own beliefs about what
+``attention_core`` can emit; that only stays true if something keeps
+checking.
+
+This is deliberately shallow. It asks the impls what they support and
+compares against what the ports claim — it does not run a model. A
+declaration that cannot be checked this way is a declaration that should
+not be made.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+from .joins import DECLARED
+from .ports import Port
+
+
+def _offers(port: Port, attr: str) -> tuple[str, ...]:
+    return tuple(port.offers.get(attr, ()))
+
+
+def check() -> list[str]:
+    """Return one line per disagreement; empty means the ports hold."""
+    problems: list[str] = []
+
+    from ..impls.attention_core import fa2_seqused
+    from ..impls.linear_proj import fp8_static as proj
+    from ..impls.qkv_pack import PackedLinear
+
+    core_cls = fa2_seqused.PackedKVAttention
+
+    # buffer=alias is only real if both sides expose the entries that
+    # make it real, and if the consumer can be told to stop copying
+    for name, (out_port, in_port) in DECLARED.items():
+        if "alias" not in _offers(out_port, "buffer"):
+            continue
+        if out_port.structure == "qkv_pack" and not hasattr(
+                PackedLinear, "alias_stash"):
+            problems.append(
+                f"{name}: qkv_pack declares buffer=alias but has no "
+                "alias_stash entry")
+        if in_port.structure == "attention_core":
+            for entry in ("alias_suffix", "forward_suffix"):
+                if not hasattr(core_cls, entry):
+                    problems.append(
+                        f"{name}: attention_core declares buffer=alias "
+                        f"but has no {entry} entry")
+
+    # layout=bshd on the core's kv port is what forward_suffix takes; the
+    # host-layout entry is a different port and must not be confused
+    kv_in = DECLARED["qkv_pack->attention_core"][1]
+    if _offers(kv_in, "layout") != ("bshd",):
+        problems.append(
+            "qkv_pack->attention_core: the core's kv port takes the "
+            "kernel layout only; forward() is the host-layout entry and "
+            "belongs to a different port")
+
+    # dtype=fp8_static has to be an entry the consumer really has
+    for name, (_, in_port) in DECLARED.items():
+        if "fp8_static" not in _offers(in_port, "dtype"):
+            continue
+        if in_port.structure == "linear_proj":
+            if "fp8_in" not in proj._BAND:
+                problems.append(
+                    f"{name}: linear_proj declares dtype=fp8_static but "
+                    "has no fp8_in form")
+
+    # opacity=must_persist is a claim about a compiler contract, so it
+    # has to point at something the compiler cannot see through
+    style_out = DECLARED["style_broker->producer"][0]
+    if "must_persist" in _offers(style_out, "opacity"):
+        import torch
+        if not hasattr(torch.ops, "flash_rt_structures") or not hasattr(
+                torch.ops.flash_rt_structures, "style_broadcast"):
+            problems.append(
+                "style_broker->producer: opacity=must_persist is claimed "
+                "but the fill is not behind an opaque op, so the "
+                "compiler may inline it back into each reader")
+
+    return problems
+
+
+def report(problems: Sequence[str] | None = None) -> str:
+    problems = check() if problems is None else problems
+    if not problems:
+        return (f"beta: {len(DECLARED)} declared join(s), all consistent "
+                "with the impls")
+    return "beta: declarations disagree with the impls\n  " + "\n  ".join(
+        problems)

flashrt_structures/beta/joins.py

@@ -0,0 +1,113 @@
+"""The joins this stack already has, declared.
+
+Descriptive before prescriptive: these ports describe what the bound
+structures do *today*, including the joins that already work. If the
+vocabulary cannot express the stack as it stands, the vocabulary is
+wrong and no behaviour should be changed on top of it.
+
+Each port carries the shape of the tensor it is about in its note, so a
+reader composing structures by hand — the explicit door — can see the
+agreement without reading the impl.
+"""
+
+from __future__ import annotations
+
+from .ports import Port
+
+# --- adaptive norm producing the packed projections' input ------------
+# Works today. The producer emits fp8 with a shared act scale and the
+# pack skips its own input quantization; the two are bound atomically.
+# Declared here to prove the vocabulary can express a join that already
+# holds, not to change it.
+ADALN_OUT = Port(
+    structure="adaln_producer", name="y", direction="out",
+    offers={"dtype": ("fp8_static", "bf16"),
+            "layout": ("row_major",),
+            "carry": ("gated_residual", "none"),
+            "cadence": ("per_call",)},
+    note="(rows, D) normed activation; carry=gated_residual means the "
+         "producer absorbed the pending residual rather than the host "
+         "closing it with an elementwise add")
+
+QKV_PACK_IN = Port(
+    structure="qkv_pack", name="x", direction="in",
+    offers={"dtype": ("fp8_static", "bf16"),
+            "layout": ("row_major",)},
+    note="(rows, K) shared input of the sibling projections")
+
+DECODER_FFN_IN = Port(
+    structure="decoder_ffn", name="x", direction="in",
+    offers={"dtype": ("fp8_static", "bf16"),
+            "layout": ("row_major",)},
+    note="(rows, D) normed input; the bf16 entry fuses its own quantize, "
+         "which is why this join only pays when the producer is already "
+         "there for another reason")
+
+# --- packed projections feeding the attention core --------------------
+# The join that still does cancelling work. The pack writes k/v into its
+# own stash buffers and the core copies them into its packed KV region;
+# both sides can express the alias, nothing negotiates it yet.
+QKV_PACK_OUT = Port(
+    structure="qkv_pack", name="kv", direction="out",
+    offers={"dtype": ("bf16",),
+            "layout": ("bshd", "row_major"),
+            "buffer": ("alias", "fresh")},
+    note="(rows, N_k) / (rows, N_v) sibling outputs; viewed as "
+         "(B, S, H, D) they are already the kernel's layout, and the "
+         "stash buffer they land in could be the consumer's region")
+
+ATTENTION_CORE_KV_IN = Port(
+    structure="attention_core", name="kv", direction="in",
+    offers={"dtype": ("bf16",),
+            "layout": ("bshd",),
+            "buffer": ("alias", "fresh")},
+    note="(B, S_suffix, H_kv, D) suffix keys/values; the packed region's "
+         "suffix rows are contiguous, so a producer can write into them "
+         "directly")
+
+ATTENTION_CORE_OUT = Port(
+    structure="attention_core", name="out", direction="out",
+    offers={"dtype": ("bf16",), "layout": ("bshd", "row_major")},
+    note="(B, S, H, D). row_major is free *because* the layout is bshd: "
+         "its last two axes are contiguous, so the (.., H*D) reshape the "
+         "output projection wants is a view. From bhsd the same reshape "
+         "would cost a transpose and a copy — which is the whole reason "
+         "layout belongs on the join rather than inside either side")
+
+LINEAR_PROJ_IN = Port(
+    structure="linear_proj", name="x", direction="in",
+    offers={"dtype": ("fp8_static", "bf16"), "layout": ("row_major",)},
+    note="(rows, K); the fp8 form has no quantize to amortise, which is "
+         "why its work band starts an order of magnitude lower")
+
+# --- the stream-scoped style materialisation --------------------------
+# The join that the compiler was allowed to undo until it was declared
+# opaque. Declared here because opacity is otherwise invisible: nothing
+# about the tensors says the arrangement must survive.
+STYLE_BROKER_OUT = Port(
+    structure="adaln_producer", name="style", direction="out",
+    offers={"layout": ("row_major",),
+            "cadence": ("per_step", "per_call"),
+            "opacity": ("must_persist", "fusible")},
+    note="(rows, 3D) style rows for one step; per_step means one fill "
+         "serves every producer on the stream, which only survives when "
+         "the fill is opaque to the compiler")
+
+STYLE_CONSUMER_IN = Port(
+    structure="adaln_producer", name="style", direction="in",
+    offers={"layout": ("row_major",),
+            "cadence": ("per_step", "per_call"),
+            "opacity": ("must_persist", "fusible")},
+    note="(rows, 3D) contiguous, as the kernel requires")
+
+
+#: Every join the stack has today, producer first. The pairs that
+#: already hold are here to be checked against reality; the pairs that
+#: do not are here to be measured.
+DECLARED = {
+    "adaln->qkv_pack": (ADALN_OUT, QKV_PACK_IN),
+    "adaln->decoder_ffn": (ADALN_OUT, DECODER_FFN_IN),
+    "qkv_pack->attention_core": (QKV_PACK_OUT, ATTENTION_CORE_KV_IN),
+    "attention_core->linear_proj": (ATTENTION_CORE_OUT, LINEAR_PROJ_IN),
+    "style_broker->producer": (STYLE_BROKER_OUT, STYLE_CONSUMER_IN),
+}

flashrt_structures/beta/negotiate.py

@@ -0,0 +1,84 @@
+"""Negotiate one join: intersect what both sides offer, name what fails.
+
+This generalises what the fp8 producer chain already does by hand — two
+components agreeing on one attribute and binding atomically — to the
+whole vocabulary. The parts that carried their weight there are kept:
+
+- the intersection is computed before anything binds, so a join that
+  cannot be agreed changes nothing;
+- a refusal names the attribute that failed, not the structure, because
+  "refused" must never read as "this structure cannot be joined";
+- an attribute only one side declares stays unnegotiated and is reported
+  as such, so a partial declaration is a partial gain rather than a
+  silent override.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Mapping, Sequence
+
+from .ports import Join, Port
+
+
+class JoinRefused(ValueError):
+    """The two ports have no common value for some attribute."""
+
+
+def negotiate(
+    producer: Port,
+    consumer: Port,
+    *,
+    prefer: Mapping[str, Sequence[str]] | None = None,
+    cost: Callable[[str, str], float] | None = None,
+) -> Join:
+    """Agree the attributes both ports declare.
+
+    ``prefer`` overrides the producer's own preference order for an
+    attribute — this is where a measured band belongs, so that a choice
+    follows the shapes at hand rather than a fixed opinion baked into
+    the port. ``cost`` picks by number instead when one is available;
+    the lowest wins.
+    """
+    if producer.direction != "out" or consumer.direction != "in":
+        raise JoinRefused(
+            f"{producer.structure}.{producer.name} -> "
+            f"{consumer.structure}.{consumer.name}: a join runs from an "
+            "out port to an in port")
+
+    prefer = dict(prefer or {})
+    chosen: dict[str, str] = {}
+    unconstrained: list[str] = []
+
+    for attr in sorted(set(producer.offers) | set(consumer.offers)):
+        out_side = producer.offers.get(attr)
+        in_side = consumer.offers.get(attr)
+        if out_side is None or in_side is None:
+            # only one side has an opinion: leave the join as the host
+            # already has it, and say so
+            unconstrained.append(attr)
+            continue
+        common = [v for v in out_side if v in in_side]
+        if not common:
+            raise JoinRefused(
+                f"{producer.structure}.{producer.name} -> "
+                f"{consumer.structure}.{consumer.name}: no common {attr} "
+                f"({list(out_side)} vs {list(in_side)}) — the join stays "
+                "as the host has it; the structures themselves are fine")
+        order = [v for v in prefer.get(attr, ()) if v in common] or common
+        if cost is not None:
+            chosen[attr] = min(order, key=lambda value: cost(attr, value))
+        else:
+            chosen[attr] = order[0]
+
+    return Join(producer=producer, consumer=consumer, chosen=chosen,
+                unconstrained=tuple(unconstrained))
+
+
+def describe(joins: Sequence[Join]) -> str:
+    """A readable account of what a set of joins agreed on."""
+    lines = []
+    for join in joins:
+        lines.append(str(join))
+        if join.unconstrained:
+            lines.append(f"    unconstrained: {', '.join(join.unconstrained)}")
+    return "\n".join(lines)

flashrt_structures/beta/ports.py

@@ -0,0 +1,140 @@
+"""The join vocabulary: what two adjacent structures have to agree on.
+
+Every attribute below was extracted from a join that had already gone
+wrong, and the incident is recorded with it. Nothing here is speculative
+— an attribute nobody negotiates is meant to be deleted, not kept as
+documentation.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Mapping, Sequence
+
+# Attribute -> the values a port may declare, most-preferred first where
+# order carries meaning. A port that declares nothing for an attribute
+# leaves it unnegotiated, which is exactly today's behaviour.
+ATTRIBUTES: dict[str, tuple[str, ...]] = {
+    # Who holds the activation in what numeric type, and therefore who
+    # pays for the conversion.
+    #
+    # Incident: an fp8 producer and the pack it feeds, bound
+    # independently, can diverge — a pack expecting fp8 whose producer
+    # did not bind leaves the host to insert a quantize, which the
+    # compiler fuses into the upstream gated residual and turns into fp8
+    # arithmetic the target cannot lower. Binding the pair atomically
+    # fixed it. This is the one join attribute that was ever declared,
+    # and the only one that stopped failing.
+    "dtype": ("fp8_static", "bf16", "int8_static"),
+
+    # How the tensor's axes are arranged when it crosses the join.
+    #
+    # Incident: the host lays q/k/v out as (B, H, S, D) because that is
+    # what eager attention wants; the fused kernel wants (B, S, H, D)
+    # and transposed back. Two transposes that cancel, plus the copies
+    # that make each contiguous, per projection per layer per step —
+    # and neither side could see the other, because each was correct on
+    # its own.
+    "layout": ("row_major", "bshd", "bhsd"),
+
+    # Who allocates the tensor the consumer reads. ``alias`` is the
+    # declarative form of what a hand-written runtime does when it hands
+    # the next stage a pointer: the producer writes straight into the
+    # region the consumer already owns.
+    #
+    # Incident: the packed projection writes its k/v into its own stash
+    # buffers and the attention core then copies them into its packed
+    # KV region — two buffers where one would do.
+    "buffer": ("alias", "caller_provided", "fresh"),
+
+    # Whether a sublayer result is still pending at the join, so the
+    # consumer can absorb it instead of the host closing it first.
+    #
+    # Incident: the adaptive-norm kernel computes ``residual + x * gate``
+    # before it norms. Bound at the norm boundary there is no residual to
+    # hand it, so it was fed zeros and the host kept its own elementwise
+    # add — measured +0.17 ms, refused. The same kernel inside a block
+    # boundary replaces that add instead of adding to it, and turns
+    # positive. Same kernel, opposite verdict, because the join was
+    # declared at the wrong place.
+    "carry": ("gated_residual", "none"),
+
+    # How often the value on this join actually changes, as opposed to
+    # how often it is read.
+    #
+    # Incident: a step's style is fixed for every producer on one
+    # conditioning stream, but was materialised per call — 720 launches
+    # for work that changes 10 times a tick.
+    "cadence": ("per_observation", "per_step", "per_call"),
+
+    # Whether the compiler may rewrite away the arrangement this join
+    # depends on. Not a hardware or dataflow property: a contract with
+    # the compiler underneath.
+    #
+    # Incident: a step-scoped shared buffer, written as plain tensor
+    # work, was legally eliminated — the compiler saw a buffer whose
+    # only consumers were slices of it and inlined the fill back into
+    # each of them. Semantically identical, and the sharing was gone:
+    # -0.164 ms instead of -0.419 ms, with the launch count unchanged.
+    # A structure layer sitting on a compiler has to be able to say
+    # which of its decisions may not be undone.
+    "opacity": ("must_persist", "fusible"),
+}
+
+
+class PortError(ValueError):
+    """A port declared something outside the vocabulary."""
+
+
+@dataclass(frozen=True)
+class Port:
+    """One side of a join: what this structure can accept or emit.
+
+    ``offers`` maps an attribute to the values this port supports, most
+    preferred first. An attribute left out is not negotiated — the join
+    keeps whatever the host already does, which is why declaring a port
+    can only add.
+    """
+
+    structure: str
+    name: str
+    direction: str                       # "in" | "out"
+    offers: Mapping[str, Sequence[str]] = field(default_factory=dict)
+    note: str = ""
+
+    def __post_init__(self) -> None:
+        if self.direction not in ("in", "out"):
+            raise PortError(
+                f"{self.structure}.{self.name}: direction must be in/out")
+        for attr, values in self.offers.items():
+            known = ATTRIBUTES.get(attr)
+            if known is None:
+                raise PortError(
+                    f"{self.structure}.{self.name}: {attr!r} is not a join "
+                    f"attribute. The vocabulary is closed on purpose — "
+                    f"adding one means a new incident to record")
+            unknown = [v for v in values if v not in known]
+            if unknown:
+                raise PortError(
+                    f"{self.structure}.{self.name}: {attr}={unknown} not in "
+                    f"{known}")
+            if not values:
+                raise PortError(
+                    f"{self.structure}.{self.name}: {attr} offers nothing; "
+                    "leave the attribute out to keep it unnegotiated")
+
+
+@dataclass(frozen=True)
+class Join:
+    """A negotiated agreement between two ports."""
+
+    producer: Port
+    consumer: Port
+    chosen: Mapping[str, str]
+    unconstrained: Sequence[str] = ()     # declared by one side only
+    note: str = ""
+
+    def __str__(self) -> str:
+        picks = ", ".join(f"{k}={v}" for k, v in sorted(self.chosen.items()))
+        return (f"{self.producer.structure}.{self.producer.name} -> "
+                f"{self.consumer.structure}.{self.consumer.name} [{picks}]")

flashrt_structures/decisions.py

@@ -0,0 +1,80 @@
+"""Band decisions: measured once per box, cached, consumed by every tier.
+
+The seat-level micro-race was refuted in both directions by
+production-form measurement; what stands is the captured end-to-end
+number. This cache is where those numbers live: a band measurement run
+records the winner per (device, band), and both assembly tiers read it
+— the automatic path to route its formats, the explicit path as its
+default band. An author pin (env) always outranks the cache; an empty
+cache falls to the precision-order default. The cache file is a
+receipt and is transportable: an edge box can ship with its decisions
+measured elsewhere on identical hardware.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import pathlib
+
+import torch
+
+
+def _cache_path() -> pathlib.Path:
+    return pathlib.Path(os.environ.get(
+        "FRT_DECISION_CACHE",
+        os.path.expanduser("~/.cache/flashrt/band_decisions.json")))
+
+
+def _key(band: str) -> str:
+    dev = (torch.cuda.get_device_name(0)
+           if torch.cuda.is_available() else "cpu")
+    return f"{dev}|{band}"
+
+
+def lookup(band: str, default: str | None = None) -> str | None:
+    try:
+        data = json.loads(_cache_path().read_text())
+    except (OSError, ValueError):
+        return default
+    entry = data.get(_key(band))
+    return entry.get("winner", default) if entry else default
+
+
+def record(band: str, winner: str, times_ms: dict) -> pathlib.Path:
+    path = _cache_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        data = json.loads(path.read_text())
+    except (OSError, ValueError):
+        data = {}
+    data[_key(band)] = {"winner": winner, "times_ms": times_ms}
+    path.write_text(json.dumps(data, indent=1))
+    return path
+
+
+def import_decisions(source: str) -> dict:
+    """Merge a decision file measured elsewhere into the local cache.
+
+    The peak-avoidance path for a memory-tight box: adjudicate on a
+    roomy machine of the same device fingerprint, ship the file, and
+    the tight box binds warm from its first run — the race (and its
+    peak) never happens there. Entries whose device fingerprint does
+    not match this box are carried along untouched (the cache is keyed
+    by device, so they are inert here and correct if the file travels
+    on). Local entries win on conflict: a number measured on this box
+    outranks an imported one. Returns a receipt of what was merged.
+    """
+    imported = json.loads(pathlib.Path(source).read_text())
+    path = _cache_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        local = json.loads(path.read_text())
+    except (OSError, ValueError):
+        local = {}
+    added = [k for k in imported if k not in local]
+    merged = dict(imported)
+    merged.update(local)
+    path.write_text(json.dumps(merged, indent=2))
+    return {"imported": len(added), "kept_local": len(local),
+            "entries": sorted(added)}