| from __future__ import annotations |
|
|
| import math |
| from typing import Optional |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import Tensor, nn |
|
|
| from b1_tensor_core import B1Linear |
| from .config import DiffusionConfig |
| from .diffusion import sinusoidal_embedding |
|
|
|
|
| class B1LinearScaled(nn.Module): |
| """ |
| 1-Bit Tensor Core linear projection normalized by 1/sqrt(in_features) |
| to maintain unit activation variance for continuous diffusion stability. |
| """ |
|
|
| def __init__(self, in_features: int, out_features: int, bias: bool = True, backend: str = "tc"): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| self.backend = backend |
| self.linear = B1Linear(in_features, out_features, bias=bias, backend=backend) |
| self.scale = 1.0 / math.sqrt(in_features) |
|
|
| def freeze_for_inference(self): |
| self.linear.freeze_for_inference() |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| |
| orig_shape = x.shape |
| if x.ndim > 2: |
| x_flat = x.reshape(-1, orig_shape[-1]) |
| out_flat = self.linear(x_flat) |
| out = out_flat.reshape(*orig_shape[:-1], self.out_features) |
| else: |
| out = self.linear(x) |
| return out * self.scale |
|
|
|
|
| class B1SelfAttention(nn.Module): |
| def __init__(self, embed_dim: int, num_heads: int = 8, backend: str = "tc"): |
| super().__init__() |
| self.embed_dim = embed_dim |
| self.num_heads = num_heads |
| self.head_dim = embed_dim // num_heads |
| self.scale = 1.0 / math.sqrt(self.head_dim) |
|
|
| self.q_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend) |
| self.k_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend) |
| self.v_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend) |
| self.out_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend) |
|
|
| def freeze_for_inference(self): |
| self.q_proj.freeze_for_inference() |
| self.k_proj.freeze_for_inference() |
| self.v_proj.freeze_for_inference() |
| self.out_proj.freeze_for_inference() |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| B, L, D = x.shape |
| q = self.q_proj(x).reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(x).reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(x).reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2) |
|
|
| scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale |
| attn = torch.softmax(scores, dim=-1) |
| out = torch.matmul(attn, v) |
| out = out.transpose(1, 2).reshape(B, L, D) |
| return self.out_proj(out) |
|
|
|
|
| class B1CrossAttention(nn.Module): |
| def __init__(self, embed_dim: int, num_heads: int = 8, backend: str = "tc"): |
| super().__init__() |
| self.embed_dim = embed_dim |
| self.num_heads = num_heads |
| self.head_dim = embed_dim // num_heads |
| self.scale = 1.0 / math.sqrt(self.head_dim) |
|
|
| self.q_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend) |
| self.k_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend) |
| self.v_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend) |
| self.out_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend) |
|
|
| def freeze_for_inference(self): |
| self.q_proj.freeze_for_inference() |
| self.k_proj.freeze_for_inference() |
| self.v_proj.freeze_for_inference() |
| self.out_proj.freeze_for_inference() |
|
|
| def forward(self, x: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: |
| B, L, D = x.shape |
| _, S, _ = memory.shape |
|
|
| q = self.q_proj(x).reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(memory).reshape(B, S, self.num_heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(memory).reshape(B, S, self.num_heads, self.head_dim).transpose(1, 2) |
|
|
| scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale |
| attn = torch.softmax(scores, dim=-1) |
| out = torch.matmul(attn, v) |
| out = out.transpose(1, 2).reshape(B, L, D) |
| return self.out_proj(out) |
|
|
|
|
| class B1MLP(nn.Module): |
| def __init__(self, embed_dim: int, mlp_dim: int, backend: str = "tc"): |
| super().__init__() |
| self.fc1 = B1LinearScaled(embed_dim, mlp_dim, backend=backend) |
| self.act = nn.GELU() |
| self.fc2 = B1LinearScaled(mlp_dim, embed_dim, backend=backend) |
|
|
| def freeze_for_inference(self): |
| self.fc1.freeze_for_inference() |
| self.fc2.freeze_for_inference() |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.fc2(self.act(self.fc1(x))) |
|
|
|
|
| class B1TransformerDecoderLayer(nn.Module): |
| def __init__(self, embed_dim: int, num_heads: int, mlp_dim: int, backend: str = "tc"): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(embed_dim) |
| self.self_attn = B1SelfAttention(embed_dim, num_heads, backend=backend) |
| self.norm2 = nn.LayerNorm(embed_dim) |
| self.cross_attn = B1CrossAttention(embed_dim, num_heads, backend=backend) |
| self.norm3 = nn.LayerNorm(embed_dim) |
| self.mlp = B1MLP(embed_dim, mlp_dim, backend=backend) |
|
|
| def freeze_for_inference(self): |
| self.self_attn.freeze_for_inference() |
| self.cross_attn.freeze_for_inference() |
| self.mlp.freeze_for_inference() |
|
|
| def forward(self, x: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: |
| |
| x = x + self.self_attn(self.norm1(x)) |
| x = x + self.cross_attn(self.norm2(x), memory) |
| x = x + self.mlp(self.norm3(x)) |
| return x |
|
|
|
|
| class B1CoherentTransformer(nn.Module): |
| """ |
| Sub-Byte 1-Bit Tensor Core continuous transformer backbone. |
| All internal attention and feedforward projections execute via inline PTX MMA. |
| """ |
|
|
| def __init__(self, config: DiffusionConfig, backend: str = "tc", pure_1bit: bool = False) -> None: |
| super().__init__() |
| self.config = config |
| self.pure_1bit = pure_1bit |
| if pure_1bit: |
| self.input_projection = B1LinearScaled(config.embedding_dim, config.hidden_dim, backend=backend) |
| self.output_projection = B1LinearScaled(config.hidden_dim, config.embedding_dim, backend=backend) |
| self.query_projection = B1LinearScaled(config.embedding_dim, config.hidden_dim, backend=backend) |
| self.time_mlp = nn.Sequential( |
| B1LinearScaled(config.hidden_dim, config.hidden_dim, backend=backend), |
| nn.SiLU(), |
| B1LinearScaled(config.hidden_dim, config.hidden_dim, backend=backend), |
| ) |
| else: |
| self.input_projection = nn.Linear(config.embedding_dim, config.hidden_dim) |
| self.output_projection = nn.Linear(config.hidden_dim, config.embedding_dim) |
| self.query_projection = nn.Linear(config.embedding_dim, config.hidden_dim) |
| self.time_mlp = nn.Sequential( |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| nn.SiLU(), |
| nn.Linear(config.hidden_dim, config.hidden_dim), |
| ) |
| self.position = nn.Parameter(torch.zeros(1, config.sequence_length, config.hidden_dim)) |
| self.layers = nn.ModuleList([ |
| B1TransformerDecoderLayer( |
| embed_dim=config.hidden_dim, |
| num_heads=config.heads, |
| mlp_dim=config.mlp_dim, |
| backend=backend, |
| ) |
| for _ in range(config.layers) |
| ]) |
| self.final_norm = nn.LayerNorm(config.hidden_dim) |
| nn.init.normal_(self.position, std=0.02) |
|
|
| def freeze_for_inference(self): |
| if self.pure_1bit: |
| self.input_projection.freeze_for_inference() |
| self.output_projection.freeze_for_inference() |
| self.query_projection.freeze_for_inference() |
| self.time_mlp[0].freeze_for_inference() |
| self.time_mlp[2].freeze_for_inference() |
| for layer in self.layers: |
| layer.freeze_for_inference() |
|
|
| def forward( |
| self, |
| noisy_targets: Tensor, |
| noise_embedding: Tensor, |
| query_embedding: Tensor, |
| condition_drop_mask: Tensor | None = None, |
| ) -> Tensor: |
| batch, length, _ = noisy_targets.shape |
| if length > self.config.sequence_length: |
| raise ValueError("input sequence exceeds configured sequence_length") |
| if condition_drop_mask is not None: |
| query_embedding = query_embedding.masked_fill(condition_drop_mask[:, None], 0.0) |
| tokens = self.input_projection(noisy_targets) + self.position[:, :length] |
| tokens = tokens + self.time_mlp(noise_embedding)[:, None] |
| memory = self.query_projection(query_embedding)[:, None] |
|
|
| for layer in self.layers: |
| tokens = layer(tokens, memory) |
|
|
| return self.output_projection(self.final_norm(tokens)) |
|
|
|
|
| class B1EDMDenoiser(nn.Module): |
| """1-Bit Hardware Tensor Core EDM-preconditioned denoiser.""" |
|
|
| def __init__(self, config: DiffusionConfig, backend: str = "tc", pure_1bit: bool = False) -> None: |
| super().__init__() |
| self.config = config |
| self.pure_1bit = pure_1bit |
| self.backbone = B1CoherentTransformer(config, backend=backend, pure_1bit=pure_1bit) |
|
|
| def freeze_for_inference(self): |
| self.backbone.freeze_for_inference() |
|
|
| def forward( |
| self, |
| noisy_targets: Tensor, |
| sigma: Tensor, |
| query_embedding: Tensor, |
| condition_drop_mask: Tensor | None = None, |
| ) -> 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) |
| effective_query = ( |
| query_embedding |
| if condition_drop_mask is None |
| else query_embedding.masked_fill(condition_drop_mask[:, None], 0.0) |
| ) |
| scaled_input = noisy_targets * c_in |
| model_output = self.backbone( |
| scaled_input, |
| noise_embedding, |
| effective_query, |
| condition_drop_mask=None, |
| ) |
| return c_skip * noisy_targets + c_out * model_output |
|
|
|
|
