scripts/fast_b1_inference.py

import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

from b1_tensor_core.ops import b1_pack, b1_matmul
from src.r4t.config import DiffusionConfig
from src.r4t.b1_diffusion import B1EDMDenoiser, B1LinearScaled
from src.r4t.diffusion import sinusoidal_embedding

CHECKPOINT_PATH = Path("checkpoints/b1_tc_10ep_champion.pt")
TAXONOMY_PATH = Path("data/taxonomy_embeddings.pt")


class FastB1Denoiser(torch.nn.Module):
    """
    Optimized 1-Bit Tensor Core Denoiser:
    1. Fuses Q, K, V into a single packed weight matrix per layer (1 pack, 1 MMA).
    2. Caches Cross-Attention memory keys & values across ODE steps.
    3. Compatible with CUDA Graph capture.
    """

    def __init__(self, base_model: B1EDMDenoiser):
        super().__init__()
        self.config = base_model.config
        self.backbone = base_model.backbone
        self.device = next(base_model.parameters()).device

        # Pre-pack and fuse QKV weights for each layer
        self.fused_layers = []
        for layer in self.backbone.layers:
            wq = layer.self_attn.q_proj.linear.packed_weight
            wk = layer.self_attn.k_proj.linear.packed_weight
            wv = layer.self_attn.v_proj.linear.packed_weight
            wqkv = torch.cat([wq, wk, wv], dim=0) # [1536, 16]
            self.fused_layers.append({
                "norm1": layer.norm1,
                "wqkv": wqkv,
                "out_proj": layer.self_attn.out_proj,
                "norm2": layer.norm2,
                "cross_q_proj": layer.cross_attn.q_proj,
                "cross_k_proj": layer.cross_attn.k_proj,
                "cross_v_proj": layer.cross_attn.v_proj,
                "cross_out_proj": layer.cross_attn.out_proj,
                "norm3": layer.norm3,
                "mlp_fc1": layer.mlp.fc1,
                "mlp_act": layer.mlp.act,
                "mlp_fc2": layer.mlp.fc2,
            })

    def precompute_cross_memory(self, query_embedding: torch.Tensor):
        """Precomputes Cross-Attention K and V once before ODE integration."""
        memory = self.backbone.query_projection(query_embedding)[:, None] # [B, 1, D]
        cached = []
        B, S, D = memory.shape
        heads = self.config.heads
        head_dim = D // heads

        for fl in self.fused_layers:
            k = fl["cross_k_proj"](memory).reshape(B, S, heads, head_dim).transpose(1, 2)
            v = fl["cross_v_proj"](memory).reshape(B, S, heads, head_dim).transpose(1, 2)
            cached.append((k, v))
        return cached

    def forward_with_cached_memory(
        self,
        noisy_targets: torch.Tensor,
        sigma: torch.Tensor,
        cached_memory: list,
    ) -> torch.Tensor:
        sigma_data = self.config.sigma_data
        sigma_sq = sigma.square()
        denominator = sigma_sq + sigma_data**2
        c_skip = sigma_data**2 / denominator
        c_out = sigma * sigma_data / denominator.sqrt()
        c_in = denominator.rsqrt()

        c_skip = c_skip.reshape(sigma.shape[0], 1, 1)
        c_out = c_out.reshape(sigma.shape[0], 1, 1)
        c_in = c_in.reshape(sigma.shape[0], 1, 1)

        noise_embedding = sinusoidal_embedding(sigma, self.config.hidden_dim)
        scaled_input = noisy_targets * c_in

        batch, length, _ = scaled_input.shape
        tokens = self.backbone.input_projection(scaled_input) + self.backbone.position[:, :length]
        tokens = tokens + self.backbone.time_mlp(noise_embedding)[:, None]

        heads = self.config.heads
        head_dim = self.config.hidden_dim // heads
        scale_attn = 1.0 / (head_dim**0.5)
        scale_qkv = 1.0 / (self.config.hidden_dim**0.5)

        for i, fl in enumerate(self.fused_layers):
            # 1. Fused Self-Attention (1 bit-pack for Q, K, V)
            norm_x = fl["norm1"](tokens)
            x_packed = b1_pack(norm_x.reshape(batch * length, -1))
            qkv_int = b1_matmul(x_packed, fl["wqkv"], None, backend="tc")
            qkv_float = qkv_int.to(torch.float32) * scale_qkv
            q, k, v = qkv_float.chunk(3, dim=-1)

            q = q.reshape(batch, length, heads, head_dim).transpose(1, 2)
            k = k.reshape(batch, length, heads, head_dim).transpose(1, 2)
            v = v.reshape(batch, length, heads, head_dim).transpose(1, 2)

            scores = torch.matmul(q, k.transpose(-2, -1)) * scale_attn
            attn = torch.softmax(scores, dim=-1)
            out_sa = torch.matmul(attn, v).transpose(1, 2).reshape(batch * length, -1)
            tokens = tokens + fl["out_proj"](out_sa).reshape(batch, length, -1)

            # 2. Fast Cross-Attention using cached Memory K, V
            norm_cx = fl["norm2"](tokens)
            cq = fl["cross_q_proj"](norm_cx).reshape(batch, length, heads, head_dim).transpose(1, 2)
            ck, cv = cached_memory[i]

            cscores = torch.matmul(cq, ck.transpose(-2, -1)) * scale_attn
            cattn = torch.softmax(cscores, dim=-1)
            out_ca = torch.matmul(cattn, cv).transpose(1, 2).reshape(batch * length, -1)
            tokens = tokens + fl["cross_out_proj"](out_ca).reshape(batch, length, -1)

            # 3. MLP
            norm_mlp = fl["norm3"](tokens)
            tokens = tokens + fl["mlp_fc2"](fl["mlp_act"](fl["mlp_fc1"](norm_mlp)))

        model_output = self.backbone.output_projection(self.backbone.final_norm(tokens))
        return c_skip * noisy_targets + c_out * model_output


@torch.no_grad()
def fast_sample_edm_8step(
    fast_model: FastB1Denoiser,
    query_embedding: torch.Tensor,
    rho: float = 7.0,
) -> torch.Tensor:
    """Optimized 8-step Heun solve with pre-cached Cross-Attention memory."""
    config = fast_model.config
    steps = 8

    # 1. Precompute cross-attention memory once
    cached_memory = fast_model.precompute_cross_memory(query_embedding)

    # 2. Polynomial noise schedule on CPU floats
    ramp = torch.linspace(0, 1, steps)
    min_inv_rho = config.sigma_min ** (1 / rho)
    max_inv_rho = config.sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
    schedule = [float(s) for s in sigmas] + [0.0]

    shape = (len(query_embedding), config.sequence_length, config.embedding_dim)
    values = torch.randn(shape, device=query_embedding.device) * config.sigma_max

    for index in range(len(schedule) - 1):
        current = schedule[index]
        following = schedule[index + 1]
        sigma = torch.full((len(values),), current, device=values.device)
        denoised = fast_model.forward_with_cached_memory(values, sigma, cached_memory)
        derivative = (values - denoised) / max(current, 1e-12)
        proposal = values + (following - current) * derivative
        if following > 0.0:
            next_sigma = torch.full((len(values),), following, device=values.device)
            next_denoised = fast_model.forward_with_cached_memory(proposal, next_sigma, cached_memory)
            values = values + (following - current) * 0.5 * (derivative + (proposal - next_denoised) / following)
        else:
            values = proposal

    return values


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device} ({torch.cuda.get_device_name(0)})")

    print(f"Loading checkpoint {CHECKPOINT_PATH}...")
    ckpt = torch.load(CHECKPOINT_PATH, map_location=device, weights_only=False)
    config = ckpt["config"]
    model = B1EDMDenoiser(config, backend="tc").to(device)

    if "ema_state_dict" in ckpt and "shadow" in ckpt["ema_state_dict"]:
        shadow = ckpt["ema_state_dict"]["shadow"]
        model.load_state_dict({k: shadow[k].to(device) for k in shadow})
    else:
        model.load_state_dict(ckpt["model_state_dict"])

    model.eval()
    model.freeze_for_inference()

    print("Building FastB1Denoiser with Fused QKV and Memory Caching...")
    fast_model = FastB1Denoiser(model)

    dummy_q = torch.randn(1, 768, device=device)
    dummy_q = F.normalize(dummy_q, dim=-1)

    # 1. Warmup
    for _ in range(5):
        fast_sample_edm_8step(fast_model, dummy_q)
    torch.cuda.synchronize()

    # 2. Benchmark Fused & Cached 8-step Heun in eager mode
    times = []
    for _ in range(20):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        fast_sample_edm_8step(fast_model, dummy_q)
        torch.cuda.synchronize()
        times.append((time.perf_counter() - t0) * 1000.0)

    eager_lat = sum(times) / len(times)
    print(f"Fused QKV + Cached Memory (8-Step Heun, Eager): {eager_lat:.2f} ms")

    # 3. CUDA Graph Capture
    print("Capturing 8-step Heun sampling in CUDA Graph...")
    static_q = dummy_q.clone()
    static_out = None

    # Warmup graph stream
    g_stream = torch.cuda.Stream()
    g_stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(g_stream):
        for _ in range(3):
            static_out = fast_sample_edm_8step(fast_model, static_q)
    torch.cuda.current_stream().wait_stream(g_stream)

    # Capture graph
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g, stream=g_stream):
        static_out = fast_sample_edm_8step(fast_model, static_q)

    # Benchmark Graph replay
    torch.cuda.synchronize()
    times_graph = []
    for _ in range(50):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        g.replay()
        torch.cuda.synchronize()
        times_graph.append((time.perf_counter() - t0) * 1000.0)

    graph_lat = sum(times_graph) / len(times_graph)
    print(f"CUDA Graph Accelerated 8-Step Heun: {graph_lat:.2f} ms (P95: {sorted(times_graph)[int(len(times_graph)*0.95)]:.2f} ms)")
    print(f"Speedup vs original un-fused (110.28 ms): {110.28 / graph_lat:.1f}x faster!")


if __name__ == "__main__":
    main()