fanout-diffusion / scripts /fast_b1_inference.py

Add training pipelines, consistency distillation scripts, and interactive dashboard server

f08972a verified about 1 month ago

9.91 kB

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()

	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()