flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/adapters/sglang_engine.py

@@ -0,0 +1,207 @@
+"""SGLang engine family: explicit assembly across a process boundary.
+
+SGLang shares vLLM's module lineage — projections return
+``(out, bias)``, seams answer the same structural predicates — but its
+scheduler is a **spawned subprocess**: a patch applied in the launcher
+process never reaches the model. The carrier that does is
+``sitecustomize``: :func:`install` writes a hook module into a
+temporary directory, prepends it to ``PYTHONPATH``, and every
+interpreter the engine spawns runs it at startup. The hook is gated on
+an environment flag and is inert anywhere else.
+
+Two further engine facts ride here:
+
+- **Quantized checkpoints hold FP8 block weights.** A dense binder fed
+  raw FP8 bytes produces garbage that no bind-time smoke can catch
+  (finite, right shape, wrong scale). Dense seating therefore
+  dequantizes through the module's own ``weight_scale_inv`` in row
+  slabs before packing — cross-format regrids carry their scale
+  semantics or they do not run.
+- **The radix prefix cache breaks repeat-determinism on hybrid
+  (linear-attention) models** — measured, host-side, seam exonerated.
+  Serving such a model with seats attached should disable the radix
+  cache until the host fixes the interplay.
+
+Usage, before constructing the engine::
+
+    from flashrt_structures.adapters import sglang_engine
+    sglang_engine.install()
+    llm = sgl.Engine(model_path=...)
+    llm.generate(...)
+
+Scope: dense projection seams (the measured 2x-class win on this
+engine). The fused-MoE and LM-head surfaces differ from vLLM's and are
+refused until profiled, not approximated.
+"""
+
+from __future__ import annotations
+
+import os
+import pathlib
+import sys
+import tempfile
+
+import torch
+
+_ATTACH_FLAG = "FRT_SGLANG_ATTACH"
+_PATH_VAR = "FRT_SGLANG_STRUCTURES_PATH"
+_SEATS_VAR = "FRT_SGLANG_SEATS"
+
+#: dense projection seams by dataflow position; the qwen3_5 family rows
+#: plus the engine-wide decoder conventions. Overridable per install.
+DENSE_SEAT_SUFFIXES = (
+    "linear_attn.out_proj", "linear_attn.in_proj_qkvz",
+    "self_attn.qkv_proj", "self_attn.o_proj",
+    "shared_expert.gate_up_proj", "shared_expert.down_proj",
+    "mlp.gate_up_proj", "mlp.down_proj",
+)
+
+
+def _dense_weight(module):
+    """The module's weight as a dense BF16 matrix.
+
+    FP8 block-quantized modules carry ``weight_scale_inv``; the dequant
+    runs in 4096-row slabs so the transient stays bounded on tight
+    cards. A quantized weight without its scale is refused."""
+    w = module.weight.data
+    fp8 = getattr(torch, "float8_e4m3fn", None)
+    if w.dtype != fp8:
+        return w
+    scale = getattr(module, "weight_scale_inv",
+                    getattr(module, "weight_scale", None))
+    if scale is None:
+        raise ValueError("refused: fp8 weight without a block scale")
+    sd = scale.data.float()
+    n, k = w.shape
+    bn = -(-n // sd.shape[0])
+    bk = -(-k // sd.shape[1])
+    rows = sd.repeat_interleave(bn, 0)[:n]
+    out = torch.empty(n, k, device=w.device, dtype=torch.bfloat16)
+    for i in range(0, n, 4096):
+        j = min(i + 4096, n)
+        out[i:j] = (w[i:j].float()
+                    * rows[i:j].repeat_interleave(bk, 1)[:, :k]
+                    ).to(torch.bfloat16)
+    return out
+
+
+def attach_engine(model, *, seats=DENSE_SEAT_SUFFIXES, use_gemv=None,
+                  release=False, verbose=True):
+    """Seat an SGLang model's dense projections; returns the handle."""
+    from .. import swap as _swap
+    from ..impls.linear_proj import nvfp4_dynamic as _linear
+    from .vllm_engine import _ProjSeat, _is_projection
+
+    if use_gemv is None:
+        use_gemv = torch.cuda.get_device_capability() >= (12, 0)
+    if not use_gemv:
+        orig_init = _linear.LinearProjNvfp4Dynamic.__init__
+
+        def _init(self, *a, **kw):
+            orig_init(self, *a, **kw)
+            self._gemv = None
+        _linear.LinearProjNvfp4Dynamic.__init__ = _init
+
+    refused = []
+    targets = [(n, m) for n, m in model.named_modules()
+               if any(n.endswith(s) for s in seats) and _is_projection(m)]
+    targets.sort(key=lambda t: t[1].weight.numel())
+    model.eval()
+
+    # On a tight card the relief must land while binding continues, not
+    # after it: with release, seats attach in slabs of original bytes
+    # and each slab's originals move to the weight store before the
+    # next slab binds. Without release, one handle carries everything.
+    GROUP = 512 << 20
+    handles, swaps, group_bytes, seated = [], {}, 0, 0
+
+    def flush():
+        nonlocal swaps, group_bytes, seated
+        if not swaps:
+            return
+        handle = _swap.attach(model, swaps)
+        if release:
+            handle.consume()
+        handles.append(handle)
+        seated += len(swaps)
+        swaps, group_bytes = {}, 0
+
+    for name, mod in targets:
+        try:
+            seam, _ = _linear.bind_proj_seam({"w": _dense_weight(mod)})
+        except Exception as e:
+            if not refused:
+                print(f"[structures.sglang] first refusal {name}: "
+                      f"{e!r}"[:180], flush=True)
+            refused.append((name, repr(e)[:120]))
+            continue
+        swaps[name] = _ProjSeat(seam)
+        group_bytes += mod.weight.numel() * mod.weight.element_size()
+        if release and group_bytes >= GROUP:
+            flush()
+    flush()
+    if verbose:
+        print(f"[structures.sglang] {seated} seats "
+              f"({len(handles)} handles), {len(refused)} refused",
+              flush=True)
+    for h in handles:
+        h.notes = {"refused": refused}
+    return handles if len(handles) != 1 else handles[0]
+
+
+def _patch_runner():
+    """Runs inside the spawned scheduler, via the sitecustomize hook."""
+    import sglang.srt.model_executor.model_runner as mr
+
+    orig = mr.ModelRunner.load_model
+
+    def load_model(self, *a, **kw):
+        orig(self, *a, **kw)
+        seats = tuple(s for s in os.environ.get(_SEATS_VAR, "").split(",")
+                      if s) or DENSE_SEAT_SUFFIXES
+        try:
+            attach_engine(self.model, seats=seats,
+                          release=os.environ.get("FRT_SGLANG_RELEASE") == "1")
+        except Exception as e:
+            print(f"[structures.sglang] attach refused: {e!r}", flush=True)
+    mr.ModelRunner.load_model = load_model
+
+
+_HOOK = """\
+import os
+if os.environ.get({flag!r}) == "1":
+    try:
+        import sys
+        sys.path.insert(0, os.environ[{path!r}])
+        from flashrt_structures.adapters import sglang_engine
+        sglang_engine._patch_runner()
+    except Exception as e:
+        print(f"[structures.sglang] hook inert: {{e!r}}", flush=True)
+"""
+
+
+def install(*, seats=None, structures_path=None, release=False):
+    """Arm the spawn hook; call before constructing the engine.
+
+    Writes a ``sitecustomize`` into a temporary directory, prepends it
+    to ``PYTHONPATH`` and flags the attach on — every interpreter the
+    engine spawns picks it up; other processes see a dormant flag."""
+    if structures_path is None:
+        structures_path = str(
+            pathlib.Path(__file__).resolve().parents[3])
+    hook_dir = tempfile.mkdtemp(prefix="frt-sglang-hook-")
+    hook = pathlib.Path(hook_dir) / "sitecustomize.py"
+    hook.write_text(_HOOK.format(flag=_ATTACH_FLAG, path=_PATH_VAR))
+    os.environ[_PATH_VAR] = structures_path
+    os.environ[_ATTACH_FLAG] = "1"
+    if seats:
+        os.environ[_SEATS_VAR] = ",".join(seats)
+    if release:
+        os.environ["FRT_SGLANG_RELEASE"] = "1"
+    prev = os.environ.get("PYTHONPATH", "")
+    os.environ["PYTHONPATH"] = (hook_dir + (":" + prev if prev else ""))
+    # the launcher itself may import sitecustomize-late; patch it too so
+    # single-process embeddings behave the same way
+    if "sglang" in sys.modules:
+        _patch_runner()
+    return hook_dir

flashrt_structures/adapters/transformers_attention_interface.py

@@ -0,0 +1,73 @@
+"""Route transformers-host attention through a hub kernel interface.
+
+Any transformers-generation host resolves its attention through
+``ALL_ATTENTION_FUNCTIONS[config._attn_implementation]`` — the same
+registry the capture lowering's pinned vision attention and the
+per-head rope route already consult. Registering a hub attention
+kernel there and switching the config is therefore a seat like any
+other: capability-qualified (registry present, package resolvable),
+fully revertible (configs restored, registry entry removed), and
+receipted. Which interface actually wins on a box is a measured band
+decision (``decisions.lookup("backbone_attn")``), never a default
+flip: an empty cache keeps the host's own interface.
+"""
+
+from __future__ import annotations
+
+import importlib
+
+import torch
+
+
+class TransformersAttentionInterfaceAdapter:
+    __name__ = "transformers_attention_interface"
+
+    def __call__(self, model, plan=None):
+        from ..decisions import lookup
+        from ..impls import KernelUnavailable, hub_kernel
+
+        if lookup("backbone_attn", default="host") != "fa4":
+            return None
+        from ..impls.graph_lowering.qwen3_vl import _find_qwen3_vl
+        target = _find_qwen3_vl(model)
+        if target is None:
+            return None
+        try:
+            fa4 = hub_kernel("kernels-community/flash-attn4", ">=0")
+        except KernelUnavailable as missing:
+            return {"refused": [("backbone_attn",
+                                 f"fa4 unavailable: {missing}")]}
+
+        def fa4_interface(module, query, key, value,
+                          attention_mask=None, scaling=None,
+                          dropout=0.0, is_causal=False, **kwargs):
+            del module, attention_mask, dropout, kwargs
+            out = fa4.flash_attn_func(
+                query.transpose(1, 2), key.transpose(1, 2),
+                value.transpose(1, 2), softmax_scale=scaling,
+                causal=bool(is_causal))
+            if isinstance(out, tuple):
+                out = out[0]
+            return out, None
+
+        base = target.model
+        modeling = importlib.import_module(type(base.visual).__module__)
+        registry = modeling.ALL_ATTENTION_FUNCTIONS
+        try:
+            registry["flashrt_fa4"] = fa4_interface
+        except TypeError:
+            registry.register("flashrt_fa4", fa4_interface)
+        configs = [base.visual.config, base.language_model.config,
+                   target.config]
+        saved = [(cfg, cfg._attn_implementation) for cfg in configs]
+        for cfg in configs:
+            cfg._attn_implementation = "flashrt_fa4"
+
+        def revert():
+            for cfg, prev in saved:
+                cfg._attn_implementation = prev
+            registry.pop("flashrt_fa4", None) \
+                if hasattr(registry, "pop") else None
+
+        return {"revert": [revert],
+                "notes": {"backbone_attn": "fa4"}}

flashrt_structures/adapters/transformers_gated_delta.py

@@ -0,0 +1,152 @@
+"""Capability adapter for Transformers-style Gated Delta call slots."""
+
+from __future__ import annotations
+
+import torch
+
+from ..impls.gated_delta_core import bind_gated_delta_core
+
+
+def _compatible(module) -> bool:
+    return (
+        callable(getattr(module, "recurrent_gated_delta_rule", None))
+        and callable(getattr(module, "chunk_gated_delta_rule", None))
+        and getattr(module, "num_v_heads", None) in (32, 48)
+        and getattr(module, "head_k_dim", None) == 128
+        and getattr(module, "head_v_dim", None) == 128
+    )
+
+
+class _Recorder:
+    def __init__(self, original, rows, phase):
+        self.original = original
+        self.rows = rows
+        self.phase = phase
+
+    def __call__(self, query, key, value, g, beta, *args, **kwargs):
+        state = kwargs.get("initial_state")
+        row = {
+            "phase": self.phase,
+            "query": query.detach(),
+            "key": key.detach(),
+            "value": value.detach(),
+            "g": g.detach(),
+            "beta": beta.detach(),
+            "state": state.detach() if state is not None else None,
+            "output_final_state": kwargs.get("output_final_state", False),
+            "use_qk_l2norm": kwargs.get(
+                "use_qk_l2norm_in_kernel", False),
+        }
+        if not self.rows:
+            self.rows.append(row)
+        else:
+            expected = tuple(self.rows[0][name].shape
+                             for name in ("query", "key", "value", "g", "beta"))
+            got = tuple(row[name].shape
+                        for name in ("query", "key", "value", "g", "beta"))
+            if got != expected:
+                raise ValueError(
+                    "gated_delta_core: shape moved inside one host call")
+        return self.original(query, key, value, g, beta, *args, **kwargs)
+
+
+class _Route:
+    def __init__(self, core, original):
+        self.core = core
+        self.original = original
+
+    def __call__(self, query, key, value, g, beta, *args, **kwargs):
+        supported_keys = {
+            "initial_state", "output_final_state",
+            "use_qk_l2norm_in_kernel", "chunk_size", "cu_seqlens",
+        }
+        unsupported = set(kwargs).difference(supported_keys)
+        packed = kwargs.get("cu_seqlens") is not None
+        if args or unsupported or packed or query.shape[1] != 1:
+            guard = getattr(self.core, "_frt_guard", None)
+            if guard is not None and not torch.compiler.is_compiling():
+                reason = "unsupported packed or extended GDN call contract"
+                guard.refuse(reason)
+            return self.original(
+                query, key, value, g, beta, *args, **kwargs)
+        initial_state = kwargs.get("initial_state")
+        output_final_state = bool(kwargs.get("output_final_state", False))
+        use_norm = bool(kwargs.get("use_qk_l2norm_in_kernel", False))
+        return self.core(
+            query, key, value, g, beta, initial_state,
+            output_final_state=output_final_state,
+            use_qk_l2norm=use_norm,
+        )
+
+
+class TransformersGatedDeltaAdapter:
+    """Route recurrent/chunk callable slots without matching class names."""
+
+    __name__ = "transformers_gated_delta"
+
+    def __call__(self, model, forward):
+        sites = [(path, module) for path, module in model.named_modules()
+                 if _compatible(module)]
+        if not sites:
+            return None
+        captures = [[] for _ in sites]
+        originals = []
+        for (_, module), rows in zip(sites, captures):
+            recurrent = module.recurrent_gated_delta_rule
+            chunk = module.chunk_gated_delta_rule
+            originals.append((module, recurrent, chunk))
+            module.recurrent_gated_delta_rule = _Recorder(
+                recurrent, rows, "decode_recurrent")
+            module.chunk_gated_delta_rule = _Recorder(
+                chunk, rows, "sequence")
+        try:
+            with torch.no_grad():
+                forward()
+        finally:
+            for module, recurrent, chunk in originals:
+                module.recurrent_gated_delta_rule = recurrent
+                module.chunk_gated_delta_rule = chunk
+
+        if not any(captures):
+            return None
+        routes = []
+        observed = {}
+        for (path, module), rows, (_, recurrent, chunk) in zip(
+                sites, captures, originals):
+            if not rows:
+                continue
+            row = rows[0]
+            if row["phase"] != "decode_recurrent":
+                raise ValueError(
+                    "gated_delta_core: Hub v3 explicit-state executable "
+                    "covers recurrent decode only; a sequence-inout artifact "
+                    "is required for prefill")
+            if not all(row[name].is_contiguous()
+                       for name in ("query", "key", "value")):
+                raise ValueError(
+                    "gated_delta_core: formal Hub v3 artifact requires "
+                    "contiguous Q/K/V, but this host exposes split views; "
+                    "a stride-aware recurrence artifact is required")
+            # Each site owns its output/state scratch. Sharing one core across
+            # equal signatures would make later layers overwrite live buffers
+            # from earlier layers during compiled or captured execution.
+            core = bind_gated_delta_core(row)
+            routes.append((module, recurrent, chunk, core))
+            observed[f"{path}.gated_delta_core"] = core
+
+        def enable():
+            for module, recurrent, chunk, core in routes:
+                module.recurrent_gated_delta_rule = _Route(core, recurrent)
+                module.chunk_gated_delta_rule = _Route(core, chunk)
+
+        def disable():
+            for module, recurrent, chunk, _ in routes:
+                module.recurrent_gated_delta_rule = recurrent
+                module.chunk_gated_delta_rule = chunk
+
+        enable()
+        return {
+            "observed": observed,
+            "revert": [disable],
+            "toggle": (enable, disable),
+        }

flashrt_structures/adapters/transformers_gated_delta_fused.py

@@ -0,0 +1,96 @@
+"""Fused-layer adapter for Transformers-style gated-delta layers.
+
+Where the callable-slot adapter serves the recurrent rule alone, this
+one binds the whole layer's cached-decode step as the fused Hub chain
+(``impls.gated_delta_core.fused_layer``). It is registered ahead of the
+callable-slot adapter and refuses cleanly — letting the ladder fall
+through — when the installed packages predate the chain's entries or a
+layer is outside the fused profile. Recognition is by shape, never by
+class or model names.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ..impls.gated_delta_core import fused_layer
+
+
+def _fusable(module) -> bool:
+    hv = getattr(module, "num_v_heads", None)
+    hk = getattr(module, "num_k_heads", None)
+    return (
+        all(isinstance(getattr(module, name, None), torch.nn.Linear)
+            for name in ("in_proj_qkv", "in_proj_z", "in_proj_b",
+                         "in_proj_a", "out_proj"))
+        and getattr(module, "conv1d", None) is not None
+        and getattr(module, "A_log", None) is not None
+        and getattr(module, "dt_bias", None) is not None
+        and getattr(module, "norm", None) is not None
+        # the chain's own profile envelope: D=128, v-heads a multiple
+        # of k-heads (the 48/16 host keeps its dedicated entries, other
+        # profiles route the head-generic ones; the bind refuses if the
+        # installed build predates them)
+        and isinstance(hv, int) and isinstance(hk, int)
+        and hv > 0 and hk > 0 and hv % hk == 0
+        and getattr(module, "head_k_dim", None) == 128
+        and getattr(module, "head_v_dim", None) == 128
+    )
+
+
+def _layer_index(path: str) -> int | None:
+    parts = path.split(".")
+    for i in range(len(parts) - 1, 0, -1):
+        if parts[i - 1] == "layers" and parts[i].isdigit():
+            return int(parts[i])
+    return None
+
+
+class TransformersGatedDeltaFusedAdapter:
+    """Bind every fusable gated-delta layer module as one fused seam."""
+
+    __name__ = "transformers_gated_delta_fused"
+    scheme_aware = True
+
+    def __call__(self, model, forward, scheme=None):
+        fmt = getattr(scheme, "gdn_projection_format", None)
+        release = bool(getattr(scheme, "gdn_release_host_weights",
+                               False))
+        sites = []
+        for path, module in model.named_modules():
+            for child_name, child in module.named_children():
+                child_path = f"{path}.{child_name}" if path else child_name
+                if not _fusable(child):
+                    continue
+                idx = _layer_index(child_path)
+                if idx is None:
+                    continue
+                sites.append((module, child_name, child, idx))
+        if not sites:
+            return None
+
+        routes = []
+        observed = {}
+        for parent, child_name, child, idx in sites:
+            # a package predating the chain raises here once, and the
+            # whole adapter steps aside for the callable-slot ladder
+            bound = fused_layer.bind_fused_decode_layer(
+                child, idx, projection_format=fmt,
+                release_host_weights=release)
+            routes.append((parent, child_name, child, bound))
+            observed[f"{child_name}@{idx}.gated_delta_fused"] = bound
+
+        def enable():
+            for parent, child_name, _child, bound in routes:
+                setattr(parent, child_name, bound)
+
+        def disable():
+            for parent, child_name, child, _bound in routes:
+                setattr(parent, child_name, child)
+
+        enable()
+        return {
+            "observed": observed,
+            "revert": [disable],
+            "toggle": (enable, disable),
+        }