code/kernels/fast.py

"""GPU fast path for laya: TileLang fused kernels + bf16 resident weights + CUDA graphs.

    agent = laya.load("convaiinnovations/laya", fast=True)      # or agent.accelerate()

Requires CUDA and `pip install laya[fast]` (tilelang).  Falls back to the stock forward otherwise.
"""
import sys
import torch
import tl_kernels as K

BF = torch.bfloat16


def _bucket_n(n):
    return 1 << max(0, (n - 1).bit_length())


class FastLaya:
    def __init__(self, model, max_len=1024, use_graphs=True, verbose=False):
        self.m = model
        enc = model.encoder
        cfg = enc.config
        dev = next(model.parameters()).device
        self.dev = dev
        self.use_graphs = use_graphs
        self.verbose = verbose
        self.H, self.Dh, self.D = cfg.num_attention_heads, cfg.hidden_size // cfg.num_attention_heads, cfg.hidden_size
        self.F = cfg.intermediate_size
        self.eps = cfg.norm_eps
        self.max_len = max_len
        f32 = lambda t: t.detach().float().contiguous()
        b16 = lambda t: t.detach().to(BF).contiguous()
        zeros = torch.zeros(self.D, device=dev)
        self.zeros = {self.D: zeros, 3 * self.D: torch.zeros(3 * self.D, device=dev), 4 * self.D: torch.zeros(4 * self.D, device=dev),
                      2 * self.F: torch.zeros(2 * self.F, device=dev)}
        # --- encoder weights
        # embeddings are gathered from an exact fp16 copy of the checkpoint values and upcast to fp32, like the stock path
        self.emb_w = enc.embeddings.tok_embeddings.weight.detach().to(torch.float16).contiguous()
        self.emb_ln = f32(enc.embeddings.norm.weight)
        # HF's bidirectional sliding mask keeps keys with |i - j| <= config.sliding_window (= local_attention // 2);
        # ModernBertAttention.sliding_window is that value + 1 (flash-attn's inclusive convention) and must NOT be used here.
        win = getattr(cfg, "sliding_window", None) or cfg.local_attention // 2
        self.layers = []
        for i, lyr in enumerate(enc.layers):
            self.layers.append(dict(
                attn_ln=None if i == 0 else f32(lyr.attn_norm.weight),
                wqkv=b16(lyr.attn.Wqkv.weight), wo=b16(lyr.attn.Wo.weight),
                mlp_ln=f32(lyr.mlp_norm.weight), wi=b16(lyr.mlp.Wi.weight), wo2=b16(lyr.mlp.Wo.weight),
                window=(win if lyr.attention_type == "sliding_attention" else 0), ltype=lyr.attention_type))
        self.final_ln = f32(enc.final_norm.weight)
        # --- rotary tables (rounded through bf16 exactly like HF does before applying)
        rot = enc.rotary_emb
        pos = torch.arange(max_len, device=dev).float()
        self.rope = {}
        for lt in set(cfg.layer_types):
            inv = getattr(rot, f"{lt}_inv_freq").float()
            scl = getattr(rot, f"{lt}_attention_scaling")
            fr = torch.outer(pos, inv)
            self.rope[lt] = ((fr.cos() * scl).to(BF).float().contiguous(), (fr.sin() * scl).to(BF).float().contiguous())
        # --- decision head (nn.TransformerEncoderLayer, norm_first, relu)
        self.type_emb = b16(model.type_emb.weight)
        self.head = []
        for lyr in model.head.layers:
            sa = lyr.self_attn
            self.head.append(dict(
                n1w=f32(lyr.norm1.weight), n1b=f32(lyr.norm1.bias), n2w=f32(lyr.norm2.weight), n2b=f32(lyr.norm2.bias),
                in_w=b16(sa.in_proj_weight), in_b=f32(sa.in_proj_bias), out_w=b16(sa.out_proj.weight), out_b=f32(sa.out_proj.bias),
                l1w=b16(lyr.linear1.weight), l1b=f32(lyr.linear1.bias), l2w=b16(lyr.linear2.weight), l2b=f32(lyr.linear2.bias)))
        # --- kernels (M is dynamic, so these compile once)
        D, F = self.D, self.F
        self.k_qkv = K.gemm_kernel(3 * D, D)
        self.k_o = K.gemm_kernel(D, D)
        self.k_geglu = K.gemm_geglu_kernel(F, D)
        self.k_o2 = K.gemm_kernel(D, F)
        self.k_addln = K.add_ln_kernel(D, residual=True, bias=False, eps=self.eps)
        self.k_addln_b = K.add_ln_kernel(D, residual=True, bias=True, eps=1e-5)
        self.k_ln_b = K.add_ln_kernel(D, residual=False, bias=True, eps=1e-5)
        self.k_in = K.gemm_kernel(3 * D, D, bias=True)
        self.k_out = K.gemm_kernel(D, D, bias=True)
        self.k_ffn1 = K.gemm_kernel(4 * D, D, bias=True, act="relu")
        self.k_ffn2 = K.gemm_kernel(D, 4 * D, bias=True)
        self._rope_k, self._rope_tab, self._attn_k = None, {}, {}
        self.graphs = {}

    # ------------------------------------------------------------------ kernels per shape
    DYNAMIC_MAX_L = 256   # up to here one dynamic-shape attention kernel is as fast as a static one
    LONG_BUCKET = 64      # beyond it, static kernels per (B, L) bucket of this size

    def rope_k(self):
        if self._rope_k is None:
            self._rope_k = K.rope_kernel(self.H, self.Dh)
        return self._rope_k

    def rope_tab(self, ltype, L):
        key = (ltype, L)
        if key not in self._rope_tab:
            cos, sin = self.rope[ltype]
            self._rope_tab[key] = (cos[:L].contiguous(), sin[:L].contiguous())
        return self._rope_tab[key]

    def attn_k(self, B, L, window):
        key = (None, None, window) if L <= self.DYNAMIC_MAX_L else (B, L, window)
        if key not in self._attn_k:
            self._attn_k[key] = K.attn_kernel(key[0], key[1], self.H, self.Dh, window=window)
        return self._attn_k[key]

    # ------------------------------------------------------------------ encoder + head on padded [B, L]
    def _encode(self, ids, lens, qtype):
        """ids [B,L] long (padded), lens [B] int32, qtype [B] long -> hidden [B, L, D] bf16"""
        B, L = ids.shape
        M, D = B * L, self.D
        dev = self.dev
        emb = torch.nn.functional.embedding(ids, self.emb_w).view(M, D).float()
        # residual stream = embeddings.norm(emb), kept in fp32 exactly like the stock autocast path
        X = torch.nn.functional.layer_norm(emb, (D,), self.emb_ln, None, self.eps)
        Y = X.to(BF)                                                            # layer 0 attends to it directly (attn_norm = Identity)
        qkv = torch.empty(M, 3 * D, device=dev, dtype=BF)
        O = torch.empty(M, D, device=dev, dtype=BF)
        G = torch.empty(M, self.F, device=dev, dtype=BF)
        z = self.zeros
        nl = len(self.layers)
        for i, ly in enumerate(self.layers):
            self.k_qkv(Y, ly["wqkv"], z[3 * D], qkv)                            # Y = attn_norm(X) (layer 0: X itself)
            cos, sin = self.rope_tab(ly["ltype"], L)
            self.rope_k()(qkv, cos, sin)
            self.attn_k(B, L, ly["window"])(qkv.view(B, L, 3, self.H, self.Dh), lens, O.view(B, L, D))
            self.k_o(O, ly["wo"], z[D], Y)                                      # Y = attn out
            self.k_addln(X, Y, ly["mlp_ln"], z[D], Y)                           # X += Y ; Y = mlp_norm(X)
            self.k_geglu(Y, ly["wi"], G)
            self.k_o2(G, ly["wo2"], z[D], Y)                                    # Y = mlp out
            nxt = self.layers[i + 1]["attn_ln"] if i + 1 < nl else self.final_ln
            self.k_addln(X, Y, nxt, z[D], Y)                                    # X += Y ; Y = next norm(X)
        # decision head: h = final_norm(x) + type_emb ; 2 x pre-norm transformer layers (relu ffn)
        X = (Y.view(B, L, D).float() + self.type_emb[qtype].float()[:, None, :]).view(M, D).contiguous()   # fp32 stream for the head
        F1 = torch.empty(M, 4 * D, device=dev, dtype=BF)
        for j, h in enumerate(self.head):
            self.k_ln_b(X, Y, h["n1w"], h["n1b"], Y)
            self.k_in(Y, h["in_w"], h["in_b"], qkv)
            self.attn_k(B, L, 0)(qkv.view(B, L, 3, self.H, self.Dh), lens, O.view(B, L, D))
            self.k_out(O, h["out_w"], h["out_b"], Y)
            self.k_addln_b(X, Y, h["n2w"], h["n2b"], Y)                         # X += attn ; Y = norm2(X)
            self.k_ffn1(Y, h["l1w"], h["l1b"], F1)
            self.k_ffn2(F1, h["l2w"], h["l2b"], Y)
            X = X + Y.float()                                                   # residual (fp32, torch, last op)
        return X.view(B, L, D)

    def _encode_graphed(self, ids, lens, qtype):
        key = tuple(ids.shape)
        g = self.graphs.get(key)
        if g is None:
            s_ids, s_lens, s_q = ids.clone(), lens.clone(), qtype.clone()
            st = torch.cuda.Stream()
            st.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(st):
                for _ in range(2):
                    self._encode(s_ids, s_lens, s_q)       # warm-up (compiles kernels, allocs)
            torch.cuda.current_stream().wait_stream(st)
            graph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(graph):
                s_out = self._encode(s_ids, s_lens, s_q)
            g = self.graphs[key] = (graph, s_ids, s_lens, s_q, s_out)
            if self.verbose:
                print(f"[fast_laya] captured CUDA graph for shape {key}", file=sys.stderr)
        graph, s_ids, s_lens, s_q, s_out = g
        s_ids.copy_(ids); s_lens.copy_(lens); s_q.copy_(qtype)
        graph.replay()
        return s_out

    # ------------------------------------------------------------------ DecisionModel.forward replacement
    @torch.no_grad()
    def forward(self, input_ids, attention_mask, marker_pos, marker_mask, qtype, detach_encoder=False):
        m = self.m
        N, L0 = input_ids.shape
        g = 16 if L0 <= self.DYNAMIC_MAX_L else self.LONG_BUCKET
        L = min(self.max_len, ((L0 + g - 1) // g) * g)
        B = _bucket_n(N)
        ids = torch.zeros(B, L, dtype=torch.long, device=self.dev)
        ids[:N, :L0] = input_ids
        lens = torch.zeros(B, dtype=torch.int32, device=self.dev)
        lens[:N] = attention_mask.sum(1).to(torch.int32)
        qt = torch.zeros(B, dtype=torch.long, device=self.dev)
        qt[:N] = qtype
        h = (self._encode_graphed if self.use_graphs else self._encode)(ids, lens, qt)
        h = h[:N, :L0].float()
        # ---- scorer / act head (tiny; identical to laya.common.DecisionModel.forward)
        idx = marker_pos.clamp(min=0)[:, :, None].expand(-1, -1, h.size(-1))
        mk = torch.gather(h, 1, idx)
        logits = m.scorer(mk).squeeze(-1).float()
        logits = logits.masked_fill(~marker_mask, -1e4)
        p = torch.softmax(logits, -1)
        k = marker_mask.sum(-1).clamp(min=2).float()
        ent = -(p * torch.log(p.clamp_min(1e-9))).sum(-1) / torch.log(k)
        # a choice with a single option (e.g. only WAIT left after exclusions): pad the second slot, as upstream does
        top2 = p.topk(2, -1).values if p.shape[-1] >= 2 else torch.stack([p[:, 0], torch.zeros_like(p[:, 0])], dim=-1)
        feats = torch.stack([top2[:, 0], top2[:, 0] - top2[:, 1], ent, k / 255.0], -1)
        pooled = h[:, 0].float()
        act_logits = m.act_head(torch.cat([pooled, feats], -1))
        return logits, act_logits


def accelerate(agent, use_graphs=True, verbose=False):
    """Patch a laya Agent in place so agent.predict() uses the TileLang fast path.  Returns the FastLaya."""
    fast = FastLaya(agent.model, max_len=agent.cfg.get("max_len", 1024), use_graphs=use_graphs, verbose=verbose)
    if not hasattr(agent, "_orig_forward"):
        agent._orig_forward = agent.model.forward
    agent.model.forward = fast.forward
    agent.fast = fast
    return fast


def restore(agent):
    if hasattr(agent, "_orig_forward"):
        agent.model.forward = agent._orig_forward
        del agent._orig_forward