colab_bundle/modules/dflash_engine.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DFlash Engine: Block Diffusion for Flash Speculative Decoding
=============================================================
Paper 47: "DFlash: Block Diffusion for Flash Speculative Decoding"
Authors: Jian Chen, Yesheng Liang, Zhijian Liu (UC San Diego)
arXiv: 2602.06036 (Feb 2026, 13 pages, ICML 2026 / PMLR 306)

Key Architectural Innovations:
1. Parallel Block Diffusion Drafting (Section 3.2 & 4.1):
   - Rather than sequential autoregressive drafting incurring T_draft = gamma * t_step,
     DFlash predicts all gamma - 1 draft tokens in parallel in a single forward pass
     (T_draft = t_parallel << gamma * t_step).
   - Allows deeper, more expressive draft models (e.g. 3-5 layers) without latency penalty.

2. Context Feature Extraction & Cross-Layer Fusion (Section 4.1 & Appendix A.3):
   - Concatenates hidden states from selected target layers [H^(l_1); ...; H^(l_5)]
     and projects them into draft hidden dimension:
     H_t = RMSNorm(W_c [H^(l_1); ...; H^(l_5)])
   - Memory overhead is negligible (W_c in R^{5D_t x D_d}, ~42MB for 35B model).

3. KV Injection Mechanism (Section 4.1 & Appendix A.3):
   - At each draft layer i:
     Q_i = W_i^Q H_d
     K_i = [W_i^K H_t; W_i^K H_d]_seq
     V_i = [W_i^V H_t; W_i^V H_d]_seq
   - Target features serve as persistent KV context without passing through Q-projection,
     FFN, or output projection, ensuring draft acceptance length scales with depth.

4. Asymmetric Position-Dependent Loss Weighting (Section 4.2 & Eq. 4):
   - Errors at early positions invalidate all subsequent speculative tokens.
   - Cross-entropy loss is weighted by exponential decay:
     w_k = exp(-(k - 1) / gamma_decay) for k in {1, ..., gamma - 1}
   - Prioritizes early positions, yielding faster convergence and higher acceptance length tau.

5. Parallel Speculative Verification (Section 3.1 & Algorithm):
   - Target model verifies proposed draft block in a single parallel forward pass.
   - Prefix acceptance logic accepts k tokens matching target argmax plus 1 bonus token.
   - Expected acceptance length tau in [1, gamma], achieving up to 6x lossless speedup.

This module is part of the Quillan-Ronin Sovereign AI Architecture.
"""

import math
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, NamedTuple

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class DFlashConfig:
    """Configuration for DFlash Block Diffusion Speculative Decoding."""
    target_hidden_dim: int = 1024
    draft_hidden_dim: int = 512
    vocab_size: int = 32000
    block_size: int = 16              # gamma: total block size (1 anchor + 15 mask tokens)
    num_target_sampled_layers: int = 5  # Uniformly sampled layers from target model
    num_draft_layers: int = 3         # Deeper draft model enabled by parallel generation
    num_heads: int = 8
    loss_decay_gamma: float = 7.0     # gamma_decay in Eq. 4 (7 for block 16, 5 for block 10)
    mask_token_id: int = 0
    temperature: float = 0.0          # 0.0 for greedy argmax verification, >0 for sampling
    dropout: float = 0.0

    @property
    def num_speculative_tokens(self) -> int:
        """Number of draft tokens generated in parallel (gamma - 1)."""
        return max(1, self.block_size - 1)


class DFlashOutput(NamedTuple):
    """Output tuple for DFlash draft block generation."""
    logits: torch.Tensor             # [batch_size, num_speculative_tokens, vocab_size]
    draft_tokens: torch.Tensor       # [batch_size, num_speculative_tokens]
    target_context_norm: torch.Tensor
    attention_entropy: float


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""

    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.pow(2).mean(-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight


class TargetContextFeatureFuser(nn.Module):
    """
    Extracts and fuses hidden representations from selected target model layers.
    Appendix A.3:
    H_t = RMSNorm(W_c [H^(l_1); ...; H^(l_5)])
    """

    def __init__(self, config: DFlashConfig):
        super().__init__()
        self.config = config
        fused_in_dim = config.target_hidden_dim * config.num_target_sampled_layers
        self.W_c = nn.Linear(fused_in_dim, config.draft_hidden_dim, bias=False)
        self.norm = RMSNorm(config.draft_hidden_dim)

    def forward(self, target_layer_hiddens: List[torch.Tensor]) -> torch.Tensor:
        """
        Args:
            target_layer_hiddens: List of tensors, each [B, T, target_hidden_dim]
        Returns:
            H_t: [B, T, draft_hidden_dim] persistent target context features
        """
        assert len(target_layer_hiddens) == self.config.num_target_sampled_layers, (
            f"Expected {self.config.num_target_sampled_layers} target layers, got {len(target_layer_hiddens)}"
        )
        # Concatenate across channel dimension
        cat_hidden = torch.cat(target_layer_hiddens, dim=-1)  # [B, T, num_layers * target_hidden_dim]
        proj = self.W_c(cat_hidden)                           # [B, T, draft_hidden_dim]
        return self.norm(proj)


class DFlashKVInjectionAttention(nn.Module):
    """
    Draft attention with persistent KV injection from target context features.
    Appendix A.3:
    Q_i = W_i^Q H_d
    K_i = [W_i^K H_t; W_i^K H_d]_seq
    V_i = [W_i^V H_t; W_i^V H_d]_seq
    """

    def __init__(self, config: DFlashConfig):
        super().__init__()
        self.config = config
        self.hidden_dim = config.draft_hidden_dim
        self.num_heads = config.num_heads
        self.head_dim = self.hidden_dim // self.num_heads
        assert self.hidden_dim % self.num_heads == 0

        self.W_q = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
        self.W_k = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
        self.W_v = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
        self.W_o = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
        self.dropout = nn.Dropout(config.dropout) if config.dropout > 0.0 else nn.Identity()

    def forward(
        self,
        H_d: torch.Tensor,
        H_t: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, float]:
        """
        Args:
            H_d: [B, S_draft, draft_hidden_dim] draft tokens (anchor + masks)
            H_t: [B, S_target, draft_hidden_dim] persistent target context features
            attn_mask: optional attention mask
        Returns:
            out: [B, S_draft, draft_hidden_dim]
            attn_entropy: scalar float
        """
        B, S_d, _ = H_d.shape
        _, S_t, _ = H_t.shape

        # 1. Queries from draft tokens only: Q_i = W_i^Q H_d
        q = self.W_q(H_d).view(B, S_d, self.num_heads, self.head_dim).transpose(1, 2)

        # 2. Keys and Values from concatenated [H_t; H_d] along sequence dimension
        cat_kv = torch.cat([H_t, H_d], dim=1)  # [B, S_t + S_d, draft_hidden_dim]
        k = self.W_k(cat_kv).view(B, S_t + S_d, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.W_v(cat_kv).view(B, S_t + S_d, self.num_heads, self.head_dim).transpose(1, 2)

        # 3. Scaled dot-product attention
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)  # [B, H, S_d, S_t + S_d]

        if attn_mask is not None:
            scores = scores + attn_mask

        probs = F.softmax(scores, dim=-1)
        probs = self.dropout(probs)

        out = torch.matmul(probs, v)  # [B, H, S_d, head_dim]
        out = out.transpose(1, 2).contiguous().view(B, S_d, self.hidden_dim)
        out = self.W_o(out)

        # Entropy of attention distribution
        with torch.no_grad():
            p_safe = probs.clamp(min=1e-12)
            attn_entropy = float(-(p_safe * p_safe.log()).sum(dim=-1).mean().item())

        return out, attn_entropy


class DFlashDraftBlock(nn.Module):
    """Single transformer layer in DFlash draft model with KV injection and MLP."""

    def __init__(self, config: DFlashConfig):
        super().__init__()
        self.norm1 = RMSNorm(config.draft_hidden_dim)
        self.attn = DFlashKVInjectionAttention(config)
        self.norm2 = RMSNorm(config.draft_hidden_dim)

        ffn_dim = config.draft_hidden_dim * 4
        # SwiGLU MLP
        self.w1 = nn.Linear(config.draft_hidden_dim, ffn_dim, bias=False)
        self.w3 = nn.Linear(config.draft_hidden_dim, ffn_dim, bias=False)
        self.w2 = nn.Linear(ffn_dim, config.draft_hidden_dim, bias=False)

    def forward(
        self,
        H_d: torch.Tensor,
        H_t: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, float]:
        normed_d = self.norm1(H_d)
        attn_out, entropy = self.attn(normed_d, H_t, attn_mask=attn_mask)
        H_d = H_d + attn_out

        normed_ffn = self.norm2(H_d)
        ffn_out = self.w2(F.silu(self.w1(normed_ffn)) * self.w3(normed_ffn))
        H_d = H_d + ffn_out

        return H_d, entropy


class DFlashDraftModel(nn.Module):
    """
    Lightweight Block Diffusion Draft Model for DFlash Speculative Decoding.
    - Fuses cross-layer target context features H_t.
    - Decodes all gamma - 1 speculative tokens in parallel in a single forward pass.
    - Employs deep draft stack (e.g. 3-5 layers) conditioned at every layer via KV injection.
    """

    def __init__(self, config: DFlashConfig):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.draft_hidden_dim)
        self.fuser = TargetContextFeatureFuser(config)

        self.layers = nn.ModuleList([
            DFlashDraftBlock(config) for _ in range(config.num_draft_layers)
        ])
        self.final_norm = RMSNorm(config.draft_hidden_dim)
        self.lm_head = nn.Linear(config.draft_hidden_dim, config.vocab_size, bias=False)

        # Tie weights if dimensions match or keep separate
        self.mask_embedding = nn.Parameter(torch.randn(config.draft_hidden_dim) * 0.02)

    def forward(
        self,
        anchor_tokens: torch.Tensor,
        target_layer_hiddens: List[torch.Tensor],
        mask_token_ids: Optional[torch.Tensor] = None
    ) -> DFlashOutput:
        """
        Parallel Block Forward Pass:
        Args:
            anchor_tokens: [B, 1] clean anchor tokens from target model
            target_layer_hiddens: List of [B, T, target_hidden_dim] extracted from target layers
            mask_token_ids: optional explicit mask ids; if None, uses learned mask embedding
        Returns:
            DFlashOutput with logits [B, gamma-1, vocab_size] and draft token predictions
        """
        B = anchor_tokens.shape[0]
        num_draft = self.config.num_speculative_tokens

        # 1. Fuse target context features H_t once (persistent across draft layers)
        H_t = self.fuser(target_layer_hiddens)  # [B, T, draft_hidden_dim]

        # 2. Construct draft input: [anchor_token, mask_1, ..., mask_{gamma-1}]
        anchor_embed = self.embed_tokens(anchor_tokens)  # [B, 1, draft_hidden_dim]

        if mask_token_ids is not None:
            mask_embed = self.embed_tokens(mask_token_ids)
        else:
            mask_embed = self.mask_embedding.view(1, 1, -1).expand(B, num_draft, -1)

        H_d = torch.cat([anchor_embed, mask_embed], dim=1)  # [B, 1 + num_draft, draft_hidden_dim]

        # 3. Pass through draft layers with persistent KV injection
        entropies = []
        for layer in self.layers:
            H_d, ent = layer(H_d, H_t)
            entropies.append(ent)

        H_d = self.final_norm(H_d)

        # 4. Project speculative positions to vocab logits (slice out anchor token position)
        speculative_hiddens = H_d[:, 1:, :]  # [B, num_draft, draft_hidden_dim]
        logits = self.lm_head(speculative_hiddens)  # [B, num_draft, vocab_size]

        if self.config.temperature == 0.0:
            draft_tokens = torch.argmax(logits, dim=-1)
        else:
            probs = F.softmax(logits / self.config.temperature, dim=-1)
            draft_tokens = torch.multinomial(probs.view(-1, self.config.vocab_size), num_samples=1).view(B, num_draft)

        return DFlashOutput(
            logits=logits,
            draft_tokens=draft_tokens,
            target_context_norm=H_t.norm(dim=-1).mean(),
            attention_entropy=sum(entropies) / max(1, len(entropies))
        )


class DFlashLoss(nn.Module):
    """
    Asymmetric Position-Dependent Loss Weighting (Section 4.2 & Eq. 4):
    w_k = exp(-(k - 1) / gamma_decay) for k in {1, ..., gamma - 1}.
    Emphasizes early positions in the speculative draft block.
    """

    def __init__(self, config: DFlashConfig):
        super().__init__()
        self.config = config
        num_draft = config.num_speculative_tokens
        gamma = config.loss_decay_gamma

        # Precompute weights: k in [1, num_draft]
        k_indices = torch.arange(num_draft, dtype=torch.float32)  # [0, 1, ..., num_draft-1]
        weights = torch.exp(-k_indices / gamma)                   # [num_draft]
        self.register_buffer("weights", weights)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            logits: [B, num_draft, vocab_size]
            targets: [B, num_draft] ground truth future tokens
        Returns:
            weighted_loss: scalar tensor
            unweighted_loss: scalar tensor
        """
        B, num_draft, V = logits.shape
        flat_logits = logits.view(-1, V)
        flat_targets = targets.view(-1)

        token_losses = F.cross_entropy(flat_logits, flat_targets, reduction="none").view(B, num_draft)

        # Apply position-dependent exponential decay weights
        w = self.weights.to(logits.device)
        weighted_loss = (token_losses * w.unsqueeze(0)).sum() / (w.sum() * B)
        unweighted_loss = token_losses.mean()

        return weighted_loss, unweighted_loss


class DFlashSpeculativeVerifier:
    """
    Parallel Speculative Decoding Verifier.
    Performs prefix verification and calculates accepted tokens and acceleration metrics.
    """

    def __init__(self, config: DFlashConfig):
        self.config = config

    def verify_block(
        self,
        draft_tokens: torch.Tensor,
        target_logits: torch.Tensor
    ) -> Tuple[torch.Tensor, int, float]:
        """
        Verifies draft block against parallel target model outputs.
        
        Args:
            draft_tokens: [num_draft] proposed draft tokens
            target_logits: [num_draft + 1, vocab_size] parallel target model logits
                           for the draft positions plus bonus position.
                           
        Returns:
            accepted_tokens: 1D tensor of accepted tokens including 1 bonus token
            num_accepted: int tau in [1, gamma]
            acceptance_rate: float in [0.0, 1.0]
        """
        num_draft = draft_tokens.size(0)
        target_predictions = torch.argmax(target_logits[:num_draft], dim=-1)  # [num_draft]

        # Find prefix agreement length
        k_match = 0
        for i in range(num_draft):
            if draft_tokens[i].item() == target_predictions[i].item():
                k_match += 1
            else:
                break

        # Bonus token from target model at first mismatch position (or end of block)
        bonus_token = torch.argmax(target_logits[k_match], dim=-1).unsqueeze(0)

        if k_match > 0:
            accepted_prefix = draft_tokens[:k_match]
            accepted = torch.cat([accepted_prefix, bonus_token], dim=0)
        else:
            accepted = bonus_token

        num_accepted = k_match + 1  # tau: accepted draft + 1 bonus token
        acceptance_rate = k_match / max(1, num_draft)

        return accepted, num_accepted, acceptance_rate


def verify_dflash_invariants(config: Optional[DFlashConfig] = None) -> Dict[str, any]:
    """
    Rigorously verifies DFlash mathematical invariants:
    1. Parallel block drafting: generates all gamma - 1 tokens in a single forward pass.
    2. KV injection: target context features are fused once and persisted across all layers.
    3. Asymmetric loss weighting: strictly monotonically decreasing weights w_1 > w_2 > ... > w_{gamma-1}.
    4. Speculative verification: exact prefix agreement + bonus token yielding tau in [1, gamma].
    5. Autograd: backward gradients flow through draft parameters, KV projections, and W_c fuser.
    """
    if config is None:
        config = DFlashConfig(
            target_hidden_dim=128,
            draft_hidden_dim=64,
            vocab_size=256,
            block_size=8,             # gamma = 8 -> 7 speculative tokens
            num_target_sampled_layers=5,
            num_draft_layers=2,
            num_heads=4,
            loss_decay_gamma=4.0
        )

    # Invariant 1: Loss weights monotonicity
    loss_module = DFlashLoss(config)
    weights = loss_module.weights.cpu().tolist()
    assert len(weights) == config.num_speculative_tokens
    for i in range(len(weights) - 1):
        assert weights[i] > weights[i + 1], f"Weights must decay monotonically: {weights[i]} <= {weights[i+1]}"

    # Invariant 2: Target Feature Fusion & Draft Forward Pass
    draft_model = DFlashDraftModel(config)
    B, T = 2, 10
    anchor = torch.randint(0, config.vocab_size, (B, 1))
    target_hiddens = [
        torch.randn(B, T, config.target_hidden_dim, requires_grad=True)
        for _ in range(config.num_target_sampled_layers)
    ]

    out = draft_model(anchor, target_hiddens)
    assert out.logits.shape == (B, config.num_speculative_tokens, config.vocab_size)
    assert out.draft_tokens.shape == (B, config.num_speculative_tokens)

    # Invariant 3: Autograd gradient propagation
    targets = torch.randint(0, config.vocab_size, (B, config.num_speculative_tokens))
    w_loss, u_loss = loss_module(out.logits, targets)
    w_loss.backward()

    # Check gradients
    assert draft_model.fuser.W_c.weight.grad is not None, "W_c must receive gradients"
    for l_idx, layer in enumerate(draft_model.layers):
        assert layer.attn.W_q.weight.grad is not None, f"Layer {l_idx} W_q must receive gradients"
        assert layer.attn.W_k.weight.grad is not None, f"Layer {l_idx} W_k must receive gradients"
        assert layer.attn.W_v.weight.grad is not None, f"Layer {l_idx} W_v must receive gradients"

    for idx, th in enumerate(target_hiddens):
        assert th.grad is not None and not torch.isnan(th.grad).any(), f"Target hidden {idx} must receive gradients"

    # Invariant 4: Parallel Speculative Verification
    verifier = DFlashSpeculativeVerifier(config)
    draft_cand = out.draft_tokens[0]  # [num_draft]
    # Simulate target logits: craft match for first 3 tokens
    mock_target_logits = torch.randn(config.num_speculative_tokens + 1, config.vocab_size)
    mock_target_logits[0, draft_cand[0]] = 100.0
    mock_target_logits[1, draft_cand[1]] = 100.0
    mock_target_logits[2, draft_cand[2]] = 100.0
    # Mismatch at token 3
    mismatch_id = (draft_cand[3] + 1) % config.vocab_size
    mock_target_logits[3, mismatch_id] = 100.0

    accepted, tau, acc_rate = verifier.verify_block(draft_cand, mock_target_logits)
    assert tau == 4, f"Expected tau=4 (3 matches + 1 bonus), got {tau}"
    assert accepted.size(0) == 4
    assert accepted[3].item() == mismatch_id.item()

    return {
        "verified": True,
        "block_size": config.block_size,
        "speculative_tokens": config.num_speculative_tokens,
        "loss_weights": weights,
        "weighted_loss": float(w_loss.item()),
        "unweighted_loss": float(u_loss.item()),
        "tau_tested": tau,
        "attention_entropy": out.attention_entropy
    }


if __name__ == "__main__":
    res = verify_dflash_invariants()
    print("DFlash Invariants Verification Result:")
    for k, v in res.items():
        print(f"  {k}: {v}")