feat: add DFlash-style backbone with target KV injection, 80 tests passing

.pytest_cache/.gitignore

@@ -0,0 +1,2 @@
+# Created by pytest automatically.
+*

.pytest_cache/CACHEDIR.TAG

@@ -0,0 +1,4 @@
+Signature: 8a477f597d28d172789f06886806bc55
+# This file is a cache directory tag created by pytest.
+# For information about cache directory tags, see:
+#	https://bford.info/cachedir/spec.html

.pytest_cache/README.md

@@ -0,0 +1,8 @@
+# pytest cache directory #
+
+This directory contains data from the pytest's cache plugin,
+which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
+
+**Do not** commit this to version control.
+
+See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.

.pytest_cache/v/cache/lastfailed

@@ -0,0 +1 @@
+{}

.pytest_cache/v/cache/nodeids

@@ -0,0 +1,82 @@
+[
+  "tests/test_acceptance.py::TestAcceptanceRule::test_all_accepted",
+  "tests/test_acceptance.py::TestAcceptanceRule::test_batch_independence",
+  "tests/test_acceptance.py::TestAcceptanceRule::test_bonus_token_exists",
+  "tests/test_acceptance.py::TestAcceptanceRule::test_empty_block",
+  "tests/test_acceptance.py::TestAcceptanceRule::test_first_rejected",
+  "tests/test_acceptance.py::TestAcceptanceRule::test_partial_acceptance",
+  "tests/test_acceptance.py::TestExpectedAcceptLength::test_completely_wrong",
+  "tests/test_acceptance.py::TestExpectedAcceptLength::test_perfect_match",
+  "tests/test_backbone.py::TestDFlashAttention::test_forward_shape",
+  "tests/test_backbone.py::TestDFlashAttention::test_gqa_forward",
+  "tests/test_backbone.py::TestDFlashAttention::test_gradient_flow",
+  "tests/test_backbone.py::TestDFlashAttention::test_with_mask",
+  "tests/test_backbone.py::TestDFlashBackbone::test_empty_draft",
+  "tests/test_backbone.py::TestDFlashBackbone::test_forward_shape",
+  "tests/test_backbone.py::TestDFlashBackbone::test_gradient_flow",
+  "tests/test_backbone.py::TestDFlashBackbone::test_many_layers",
+  "tests/test_backbone.py::TestDFlashBackbone::test_output_hidden_states",
+  "tests/test_backbone.py::TestDFlashBackbone::test_single_draft_token",
+  "tests/test_backbone.py::TestDFlashDecoderLayer::test_all_activations",
+  "tests/test_backbone.py::TestDFlashDecoderLayer::test_forward_shape",
+  "tests/test_backbone.py::TestDFlashDecoderLayer::test_residual_connection",
+  "tests/test_backbone.py::TestDSparkAttentionMask::test_context_attention",
+  "tests/test_backbone.py::TestDSparkAttentionMask::test_cross_block_no_attention",
+  "tests/test_backbone.py::TestDSparkAttentionMask::test_intra_block_attention",
+  "tests/test_backbone.py::TestDSparkAttentionMask::test_mask_shape",
+  "tests/test_markov_head.py::TestGatedMarkovHead::test_forward_shape",
+  "tests/test_markov_head.py::TestGatedMarkovHead::test_init",
+  "tests/test_markov_head.py::TestRNNHead::test_apply_block_logits",
+  "tests/test_markov_head.py::TestRNNHead::test_empty_block",
+  "tests/test_markov_head.py::TestRNNHead::test_init",
+  "tests/test_markov_head.py::TestRNNHead::test_step",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_apply_block_logits",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_apply_step_logits",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_build_gated",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_build_vanilla",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_compute_step_bias",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_get_prev_embeddings",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_init",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_sample_block_tokens_empty",
+  "tests/test_markov_head.py::TestVanillaMarkov::test_sample_block_tokens_greedy",
+  "tests/test_sampling.py::TestSampling::test_gather_token_probs",
+  "tests/test_sampling.py::TestSampling::test_logits_to_probs_greedy",
+  "tests/test_sampling.py::TestSampling::test_logits_to_probs_temperature",
+  "tests/test_sampling.py::TestSampling::test_sample_residual",
+  "tests/test_sampling.py::TestSampling::test_sample_residual_identical",
+  "tests/test_sampling.py::TestSampling::test_sample_tokens_2d",
+  "tests/test_sampling.py::TestSampling::test_sample_tokens_greedy",
+  "tests/test_sampling.py::TestSampling::test_sample_tokens_temperature",
+  "tests/test_scheduler.py::TestHardwareAwareScheduler::test_all_zero_confidence",
+  "tests/test_scheduler.py::TestHardwareAwareScheduler::test_monotonic_selection",
+  "tests/test_scheduler.py::TestHardwareAwareScheduler::test_multiple_requests",
+  "tests/test_scheduler.py::TestHardwareAwareScheduler::test_no_future_leakage",
+  "tests/test_scheduler.py::TestHardwareAwareScheduler::test_per_request_validity",
+  "tests/test_scheduler.py::TestHardwareAwareScheduler::test_single_request",
+  "tests/test_scheduler.py::TestStaticScheduler::test_fixed_length",
+  "tests/test_scheduler.py::TestStaticScheduler::test_fixed_length_zero",
+  "tests/test_scheduler.py::TestStaticScheduler::test_static_threshold",
+  "tests/test_scheduler.py::TestStaticScheduler::test_static_threshold_all_below",
+  "tests/test_scheduler.py::TestThroughputProfile::test_clamp_high",
+  "tests/test_scheduler.py::TestThroughputProfile::test_clamp_low",
+  "tests/test_scheduler.py::TestThroughputProfile::test_exact_lookup",
+  "tests/test_scheduler.py::TestThroughputProfile::test_interpolation",
+  "tests/test_shapes.py::TestConfidenceHeadShapes::test_accept_rate_shape",
+  "tests/test_shapes.py::TestConfidenceHeadShapes::test_confidence_head",
+  "tests/test_shapes.py::TestIntegration::test_draft_verify_cycle",
+  "tests/test_shapes.py::TestIntegration::test_train_step_shape",
+  "tests/test_shapes.py::TestLossShapes::test_loss_decay_weights",
+  "tests/test_shapes.py::TestLossShapes::test_loss_shapes",
+  "tests/test_shapes.py::TestLossShapes::test_loss_with_all_terms",
+  "tests/test_shapes.py::TestModelShapes::test_forward_no_confidence",
+  "tests/test_shapes.py::TestModelShapes::test_forward_shapes",
+  "tests/test_shapes.py::TestModelShapes::test_sample_block_shapes",
+  "tests/test_sts.py::TestExpectedCalibrationError::test_miscalibrated",
+  "tests/test_sts.py::TestExpectedCalibrationError::test_perfectly_calibrated",
+  "tests/test_sts.py::TestExpectedCalibrationError::test_shape",
+  "tests/test_sts.py::TestSTSCalibrator::test_fit_transform",
+  "tests/test_sts.py::TestSTSCalibrator::test_shape_preserved",
+  "tests/test_sts.py::TestSTSCalibrator::test_transform_unfitted",
+  "tests/test_sts.py::TestSequentialTemperatureScaling::test_sts_monotonic",
+  "tests/test_sts.py::TestSequentialTemperatureScaling::test_sts_returns_temperatures"
+]

README.md

@@ -58,7 +58,7 @@ sdk_version: "1.0"
   <img src="https://img.shields.io/badge/python-3.10+-blue.svg" alt="Python"/>
   <img src="https://img.shields.io/badge/pytorch-2.1+-orange.svg" alt="PyTorch"/>
   <img src="https://img.shields.io/badge/build-passing-brightgreen.svg" alt="Build"/>
-  <img src="https://img.shields.io/badge/tests-55%20passing-brightgreen.svg" alt="Tests"/>
+  <img src="https://img.shields.io/badge/tests-80%20passing-brightgreen.svg" alt="Tests"/>
 </p>
 
 ---
@@ -85,6 +85,7 @@ UraionSpec/
 │   │   ├── markov_head.py       # Low-rank transition bias (r=256)
 │   │   ├── rnn_head.py          # GRU-like recurrent sequential head
 │   │   ├── confidence_head.py   # Per-position acceptance predictor
+│   │   ├── dflash_backbone.py   # DFlash-style backbone with KV injection ⭐
 │   │   └── draft_model.py       # Combined parallel backbone + heads
 │   ├── decoding/            # Speculative decoding core
 │   │   ├── acceptance.py        # Lossless rejection sampling (min ratio)
@@ -208,6 +209,8 @@ All position-weighted by `w_k = exp(-(k-1)/γ)` emphasizing earlier positions.
 | Component | Status |
 |---|---|
 | 55 unit & integration tests | ✅ All passing |
+| DFlash backbone with KV injection | ✅ 16 tests, all shapes & gradients verified |
+| Sampling utilities (residual, GQA) | ✅ 8 tests |
 | Package import | ✅ Clean |
 | Linting (ruff) | ✅ All checks passed |
 | Smoke training (CPU) | ✅ 3 steps, all losses decreasing |

src/uraionspec.egg-info/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.4
+Name: uraionspec
+Version: 0.1.0
+Summary: Faithful DSpark-style speculative decoding implementation by Uraion Labs
+Author-email: Uraion Labs <uraionlabs@gmail.com>
+License: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.1.0
+Requires-Dist: transformers>=4.38.0
+Requires-Dist: datasets>=2.14.0
+Requires-Dist: accelerate>=0.25.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: tqdm>=4.64.0
+Requires-Dist: sentencepiece>=0.1.99
+Requires-Dist: protobuf>=3.20
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Provides-Extra: eval
+Requires-Dist: vllm>=0.4.0; extra == "eval"

src/uraionspec.egg-info/SOURCES.txt

@@ -0,0 +1,35 @@
+pyproject.toml
+src/uraionspec/__init__.py
+src/uraionspec.egg-info/PKG-INFO
+src/uraionspec.egg-info/SOURCES.txt
+src/uraionspec.egg-info/dependency_links.txt
+src/uraionspec.egg-info/requires.txt
+src/uraionspec.egg-info/top_level.txt
+src/uraionspec/calibration/__init__.py
+src/uraionspec/calibration/sts.py
+src/uraionspec/decoding/__init__.py
+src/uraionspec/decoding/acceptance.py
+src/uraionspec/decoding/scheduler.py
+src/uraionspec/decoding/speculative.py
+src/uraionspec/evaluation/__init__.py
+src/uraionspec/evaluation/benchmark_latency.py
+src/uraionspec/evaluation/eval_acceptance.py
+src/uraionspec/models/__init__.py
+src/uraionspec/models/confidence_head.py
+src/uraionspec/models/draft_model.py
+src/uraionspec/models/markov_head.py
+src/uraionspec/models/rnn_head.py
+src/uraionspec/training/__init__.py
+src/uraionspec/training/cache_targets.py
+src/uraionspec/training/dataset.py
+src/uraionspec/training/losses.py
+src/uraionspec/training/train_drafter.py
+src/uraionspec/utils/__init__.py
+src/uraionspec/utils/hf.py
+src/uraionspec/utils/logging.py
+src/uraionspec/utils/seed.py
+tests/test_acceptance.py
+tests/test_markov_head.py
+tests/test_scheduler.py
+tests/test_shapes.py
+tests/test_sts.py

src/uraionspec.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

src/uraionspec.egg-info/requires.txt

@@ -0,0 +1,16 @@
+torch>=2.1.0
+transformers>=4.38.0
+datasets>=2.14.0
+accelerate>=0.25.0
+numpy>=1.24.0
+tqdm>=4.64.0
+sentencepiece>=0.1.99
+protobuf>=3.20
+
+[dev]
+pytest>=7.0
+ruff>=0.1.0
+pytest-cov>=4.1.0
+
+[eval]
+vllm>=0.4.0

src/uraionspec.egg-info/top_level.txt

@@ -0,0 +1 @@
+uraionspec

src/uraionspec/models/__init__.py

@@ -2,6 +2,7 @@ from .markov_head import VanillaMarkov, GatedMarkovHead, build_markov_head
 from .rnn_head import RNNHead
 from .confidence_head import ConfidenceHead, compute_accept_rate
 from .draft_model import DSparkDraftModel
+from .dflash_backbone import DFlashBackbone, DFlashDecoderLayer, DFlashAttention, DSparkAttentionMask
 
 __all__ = [
     "VanillaMarkov",
@@ -11,4 +12,8 @@ __all__ = [
     "ConfidenceHead",
     "compute_accept_rate",
     "DSparkDraftModel",
+    "DFlashBackbone",
+    "DFlashDecoderLayer",
+    "DFlashAttention",
+    "DSparkAttentionMask",
 ]

src/uraionspec/models/dflash_backbone.py

@@ -0,0 +1,307 @@
+"""
+DFlash-style parallel backbone with target model KV injection.
+
+This is the core parallel backbone used in DSpark (and DFlash).
+Key innovation over standard Transformer decoders:
+- Each layer receives target_hidden_states as context
+- Attention concatenates target KVs with draft KVs:
+    K = [W_K @ target_hidden; W_K @ draft_hidden]
+    V = [W_V @ target_hidden; W_V @ draft_hidden]
+- Draft tokens attend bidirectionally to context and intra-block draft tokens
+- This gives the draft model rich contextual information from the target
+
+Reference: DSpark paper Section 3.1, DFlash (Chen et al., 2026)
+"""
+
+from typing import Optional
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class DFlashAttention(nn.Module):
+    """Multi-head attention with target KV injection.
+
+    Concatenates target context KVs with draft KVs so draft tokens
+    can attend to target representations.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: Optional[int] = None,
+        dropout: float = 0.0,
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.head_dim = hidden_size // num_heads
+        self.dropout = dropout
+
+        self.q_proj = nn.Linear(hidden_size, num_heads * self.head_dim, bias=bias)
+        self.k_proj = nn.Linear(hidden_size, self.num_kv_heads * self.head_dim, bias=bias)
+        self.v_proj = nn.Linear(hidden_size, self.num_kv_heads * self.head_dim, bias=bias)
+        self.o_proj = nn.Linear(num_heads * self.head_dim, hidden_size, bias=bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        target_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: [B, L_draft, D] draft token hidden states
+            target_hidden_states: [B, L_ctx, D] target model context features
+            attention_mask: [B, 1, L_draft, L_ctx + L_draft] or None
+
+        Returns:
+            output: [B, L_draft, D] attended hidden states
+        """
+        B, L_draft, _ = hidden_states.shape
+        L_ctx = target_hidden_states.shape[1]
+
+        # Project Q from draft hidden states
+        q = self.q_proj(hidden_states)
+        q = q.view(B, L_draft, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Project K, V from BOTH target context and draft
+        k_ctx = self.k_proj(target_hidden_states)
+        k_draft = self.k_proj(hidden_states)
+        k = torch.cat([k_ctx, k_draft], dim=1)
+        k = k.view(B, L_ctx + L_draft, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        v_ctx = self.v_proj(target_hidden_states)
+        v_draft = self.v_proj(hidden_states)
+        v = torch.cat([v_ctx, v_draft], dim=1)
+        v = v.view(B, L_ctx + L_draft, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        # Repeat KV heads for GQA
+        if self.num_kv_groups > 1:
+            k = k.repeat_interleave(self.num_kv_groups, dim=1)
+            v = v.repeat_interleave(self.num_kv_groups, dim=1)
+
+        # Scaled dot-product attention
+        scale = self.head_dim ** -0.5
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * scale
+
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
+        attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)
+
+        attn_output = torch.matmul(attn_weights, v)
+
+        # Handle empty draft (reshape with 0 elements)
+        if L_draft == 0:
+            return self.o_proj(attn_output.transpose(1, 2).reshape(B, 0, self.num_heads * self.head_dim))
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(B, L_draft, -1)
+
+        return self.o_proj(attn_output)
+
+
+class DFlashDecoderLayer(nn.Module):
+    """Single decoder layer for the DFlash-style backbone.
+
+    Pre-norm architecture: norm → attention → residual → norm → FFN → residual
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        intermediate_size = intermediate_size or hidden_size * 4
+
+        self.input_layernorm = nn.LayerNorm(hidden_size, eps=1e-6)
+        self.self_attn = DFlashAttention(
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            dropout=dropout,
+            bias=bias,
+        )
+        self.post_attention_layernorm = nn.LayerNorm(hidden_size, eps=1e-6)
+
+        # FFN
+        if activation == "gelu":
+            act_fn = nn.GELU(approximate="tanh")
+        elif activation == "relu":
+            act_fn = nn.ReLU()
+        elif activation == "silu":
+            act_fn = nn.SiLU()
+        else:
+            raise ValueError(f"Unsupported activation: {activation}")
+
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, intermediate_size, bias=bias),
+            act_fn,
+            nn.Linear(intermediate_size, hidden_size, bias=bias),
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        target_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # Pre-norm attention with context injection
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states,
+            target_hidden_states,
+            attention_mask,
+        )
+        hidden_states = residual + hidden_states
+
+        # Pre-norm FFN
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class DFlashBackbone(nn.Module):
+    """DFlash-style parallel backbone for DSpark.
+
+    A stack of DFlashDecoderLayers that all receive target_hidden_states
+    as context for KV injection.
+
+    This processes all draft positions in a single forward pass
+    (parallel, not autoregressive), making drafting latency nearly
+    independent of block size γ.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_layers: int,
+        num_attention_heads: int,
+        num_kv_heads: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_kv_heads = num_kv_heads or num_attention_heads
+
+        self.layers = nn.ModuleList([
+            DFlashDecoderLayer(
+                hidden_size=hidden_size,
+                num_heads=num_attention_heads,
+                num_kv_heads=self.num_kv_heads,
+                intermediate_size=intermediate_size,
+                dropout=dropout,
+                activation=activation,
+                bias=bias,
+            )
+            for _ in range(num_layers)
+        ])
+        self.norm = nn.LayerNorm(hidden_size, eps=1e-6)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        target_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: [B, L_draft, D] draft token embeddings
+            target_hidden_states: [B, L_ctx, D] target model context features
+            attention_mask: optional mask for causal/non-causal attention
+            output_hidden_states: if True, return all hidden layer outputs
+
+        Returns:
+            hidden_states: [B, L_draft, D] after all backbone layers
+        """
+        all_hidden = [hidden_states] if output_hidden_states else None
+
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                target_hidden_states,
+                attention_mask,
+            )
+            if output_hidden_states:
+                all_hidden.append(hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            return hidden_states, all_hidden
+        return hidden_states
+
+
+class DSparkAttentionMask:
+    """Build custom attention masks for DSpark training.
+
+    Draft tokens in the same block attend bidirectionally to each other
+    and to all context tokens, but NOT to draft tokens in other blocks.
+    """
+
+    @staticmethod
+    def create_dspark_attention_mask(
+        *,
+        batch_size: int,
+        seq_len: int,
+        num_blocks: int,
+        block_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Create a block-diagonal attention mask for DSpark.
+
+        Each block's draft tokens can attend to:
+        - All context tokens (bidirectional)
+        - All other draft tokens in the same block (bidirectional)
+        - But NOT to draft tokens in other blocks
+
+        Returns:
+            mask: [B, 1, L_draft, L_ctx + L_draft] float mask
+        """
+        L_ctx = seq_len
+        L_draft = num_blocks * block_size
+        L_total = L_ctx + L_draft
+
+        # Start with all-to-all
+        mask = torch.zeros(batch_size, 1, L_draft, L_total, device=device)
+
+        for b in range(batch_size):
+            for block_idx in range(num_blocks):
+                start = block_idx * block_size
+                end = start + block_size
+
+                # Can attend to all context tokens
+                mask[b, :, start:end, :L_ctx] = 0.0
+
+                # Can attend to all draft tokens in the same block
+                mask[b, :, start:end, L_ctx + start:L_ctx + end] = 0.0
+
+                # Cannot attend to draft tokens in other blocks (already -inf)
+                mask[b, :, start:end, L_ctx:L_ctx + start] = float("-inf")
+                mask[b, :, start:end, L_ctx + end:L_ctx + L_draft] = float("-inf")
+
+        return mask

src/uraionspec/utils/__init__.py

@@ -1,6 +1,7 @@
 from .hf import load_model_and_tokenizer, get_hf_token
 from .logging import setup_logger, add_metric
 from .seed import seed_everything
+from .sampling import logits_to_probs, sample_tokens, sample_residual, gather_token_probs
 
 __all__ = [
     "load_model_and_tokenizer",
@@ -8,4 +9,8 @@ __all__ = [
     "setup_logger",
     "add_metric",
     "seed_everything",
+    "logits_to_probs",
+    "sample_tokens",
+    "sample_residual",
+    "gather_token_probs",
 ]

src/uraionspec/utils/sampling.py

@@ -0,0 +1,75 @@
+"""Sampling utilities for speculative decoding."""
+
+import torch
+
+
+def logits_to_probs(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor:
+    """Convert logits to probabilities with temperature."""
+    if temperature < 1e-5:
+        probs = torch.zeros_like(logits, dtype=torch.float32)
+        probs.scatter_(-1, torch.argmax(logits, dim=-1, keepdim=True), 1.0)
+        return probs
+    return torch.softmax(logits.float() / temperature, dim=-1)
+
+
+def sample_tokens(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor:
+    """Sample tokens from logits with temperature.
+
+    Args:
+        logits: [B, L, V] or [B, V]
+        temperature: 0 = greedy, >0 = multinomial
+
+    Returns:
+        token_ids: same shape as logits minus last dim
+    """
+    if temperature < 1e-5:
+        return torch.argmax(logits, dim=-1)
+
+    bsz = logits.shape[:-1]
+    flat_logits = logits.reshape(-1, logits.size(-1)) / temperature
+    probs = torch.softmax(flat_logits, dim=-1)
+    sampled = torch.multinomial(probs, num_samples=1)
+    return sampled.reshape(*bsz)
+
+
+def sample_residual(
+    target_probs: torch.Tensor,
+    draft_probs: torch.Tensor,
+) -> torch.Tensor:
+    """Sample from residual distribution p_target - p_draft (clamped).
+
+    Used for bonus token sampling in speculative decoding.
+
+    Args:
+        target_probs: [B, V] target probabilities
+        draft_probs: [B, V] draft probabilities
+
+    Returns:
+        token_ids: [B] sampled tokens
+    """
+    residual = torch.clamp(target_probs - draft_probs, min=0.0)
+    residual_mass = residual.sum(dim=-1, keepdim=True)
+    # If residual is near-zero, fall back to target distribution
+    if torch.any(residual_mass <= 1e-8):
+        residual = torch.where(residual_mass <= 1e-8, target_probs, residual)
+        residual_mass = residual.sum(dim=-1, keepdim=True)
+    residual = residual / residual_mass.clamp_min(1e-8)
+    flat = residual.reshape(-1, residual.size(-1))
+    sampled = torch.multinomial(flat, num_samples=1)
+    return sampled.reshape(residual.shape[:-1])
+
+
+def gather_token_probs(
+    probs: torch.Tensor,
+    token_ids: torch.Tensor,
+) -> torch.Tensor:
+    """Gather probabilities of specific tokens from a probability tensor.
+
+    Args:
+        probs: [*, V] probability tensor
+        token_ids: [*] token indices
+
+    Returns:
+        token_probs: [*] probability of each token
+    """
+    return probs.gather(dim=-1, index=token_ids.unsqueeze(-1)).squeeze(-1)

tests/test_backbone.py

@@ -0,0 +1,255 @@
+"""Tests for the DFlash-style parallel backbone with KV injection."""
+
+import torch
+import pytest
+
+from uraionspec.models.dflash_backbone import (
+    DFlashBackbone,
+    DFlashDecoderLayer,
+    DFlashAttention,
+    DSparkAttentionMask,
+)
+
+
+class TestDFlashAttention:
+    """Test the target KV injection attention."""
+
+    @pytest.fixture
+    def attn(self):
+        return DFlashAttention(hidden_size=64, num_heads=4, dropout=0.0)
+
+    def test_forward_shape(self, attn):
+        B, L_draft, L_ctx, D = 2, 4, 8, 64
+        hidden = torch.randn(B, L_draft, D)
+        target = torch.randn(B, L_ctx, D)
+
+        out = attn(hidden, target)
+        assert out.shape == (B, L_draft, D)
+
+    def test_gqa_forward(self):
+        """Test grouped query attention."""
+        attn = DFlashAttention(hidden_size=64, num_heads=4, num_kv_heads=2)
+        B, L_draft, L_ctx, D = 2, 4, 8, 64
+        hidden = torch.randn(B, L_draft, D)
+        target = torch.randn(B, L_ctx, D)
+
+        out = attn(hidden, target)
+        assert out.shape == (B, L_draft, D)
+
+    def test_with_mask(self, attn):
+        B, L_draft, L_ctx, D = 2, 4, 8, 64
+        hidden = torch.randn(B, L_draft, D)
+        target = torch.randn(B, L_ctx, D)
+
+        # Causal mask
+        mask = torch.zeros(B, 1, L_draft, L_ctx + L_draft)
+        mask[:, :, :, L_ctx:] = torch.triu(
+            torch.full((L_draft, L_draft), float("-inf")), diagonal=1
+        ).unsqueeze(0).unsqueeze(0)
+
+        out = attn(hidden, target, attention_mask=mask)
+        assert out.shape == (B, L_draft, D)
+
+    def test_gradient_flow(self, attn):
+        B, L_draft, L_ctx, D = 1, 2, 4, 64
+        hidden = torch.randn(B, L_draft, D, requires_grad=True)
+        target = torch.randn(B, L_ctx, D)
+
+        out = attn(hidden, target)
+        loss = out.sum()
+        loss.backward()
+        assert hidden.grad is not None
+        assert hidden.grad.shape == (B, L_draft, D)
+
+
+class TestDFlashDecoderLayer:
+    """Test a single DFlash decoder layer."""
+
+    @pytest.fixture
+    def layer(self):
+        return DFlashDecoderLayer(
+            hidden_size=64,
+            num_heads=4,
+            intermediate_size=128,
+            dropout=0.0,
+        )
+
+    def test_forward_shape(self, layer):
+        B, L_draft, L_ctx, D = 2, 4, 8, 64
+        hidden = torch.randn(B, L_draft, D)
+        target = torch.randn(B, L_ctx, D)
+
+        out = layer(hidden, target)
+        assert out.shape == (B, L_draft, D)
+
+    def test_residual_connection(self, layer):
+        """Output should differ from input (non-identity transformation)."""
+        B, L_draft, L_ctx, D = 1, 2, 4, 64
+        hidden = torch.randn(B, L_draft, D)
+        target = torch.randn(B, L_ctx, D)
+
+        with torch.no_grad():
+            out = layer(hidden, target)
+        assert not torch.allclose(out, hidden, atol=1e-4)
+
+    def test_all_activations(self):
+        for act in ["gelu", "relu", "silu"]:
+            layer = DFlashDecoderLayer(
+                hidden_size=32, num_heads=2, intermediate_size=64,
+                dropout=0.0, activation=act,
+            )
+            B, L_draft, L_ctx = 1, 2, 4
+            hidden = torch.randn(B, L_draft, 32)
+            target = torch.randn(B, L_ctx, 32)
+            out = layer(hidden, target)
+            assert out.shape == (B, L_draft, 32)
+
+
+class TestDFlashBackbone:
+    """Test the full DFlash backbone stack."""
+
+    @pytest.fixture
+    def backbone(self):
+        return DFlashBackbone(
+            hidden_size=64,
+            num_layers=2,
+            num_attention_heads=4,
+            intermediate_size=128,
+            dropout=0.0,
+        )
+
+    def test_forward_shape(self, backbone):
+        B, L_draft, L_ctx, D = 2, 4, 8, 64
+        hidden = torch.randn(B, L_draft, D)
+        target = torch.randn(B, L_ctx, D)
+
+        out = backbone(hidden, target)
+        assert out.shape == (B, L_draft, D)
+
+    def test_output_hidden_states(self, backbone):
+        B, L_draft, L_ctx, D = 2, 4, 8, 64
+        hidden = torch.randn(B, L_draft, D)
+        target = torch.randn(B, L_ctx, D)
+
+        out, all_hidden = backbone(hidden, target, output_hidden_states=True)
+        assert len(all_hidden) == 3  # input + 2 layers
+        for h in all_hidden:
+            assert h.shape == (B, L_draft, D)
+
+    def test_gradient_flow(self, backbone):
+        B, L_draft, L_ctx, D = 1, 3, 6, 64
+        hidden = torch.randn(B, L_draft, D, requires_grad=True)
+        target = torch.randn(B, L_ctx, D)
+
+        out = backbone(hidden, target)
+        loss = out.sum()
+        loss.backward()
+        assert hidden.grad is not None
+
+    def test_empty_draft(self, backbone):
+        """Edge case: no draft tokens."""
+        B, L_ctx, D = 2, 8, 64
+        hidden = torch.randn(B, 0, D)
+        target = torch.randn(B, L_ctx, D)
+
+        out = backbone(hidden, target)
+        assert out.shape == (B, 0, D)
+
+    def test_single_draft_token(self, backbone):
+        """Edge case: single draft token."""
+        B, L_ctx, D = 2, 8, 64
+        hidden = torch.randn(B, 1, D)
+        target = torch.randn(B, L_ctx, D)
+
+        out = backbone(hidden, target)
+        assert out.shape == (B, 1, D)
+
+    def test_many_layers(self):
+        """Test with more layers."""
+        backbone = DFlashBackbone(
+            hidden_size=32, num_layers=6, num_attention_heads=4,
+        )
+        B, L_draft, L_ctx = 2, 4, 8
+        hidden = torch.randn(B, L_draft, 32)
+        target = torch.randn(B, L_ctx, 32)
+
+        out = backbone(hidden, target)
+        assert out.shape == (B, L_draft, 32)
+
+
+class TestDSparkAttentionMask:
+    """Test the custom DSpark attention mask builder."""
+
+    def test_mask_shape(self):
+        B, seq_len = 2, 10
+        num_blocks, block_size = 3, 4
+        device = "cpu"
+
+        mask = DSparkAttentionMask.create_dspark_attention_mask(
+            batch_size=B,
+            seq_len=seq_len,
+            num_blocks=num_blocks,
+            block_size=block_size,
+            device=torch.device(device),
+        )
+        L_draft = num_blocks * block_size
+        assert mask.shape == (B, 1, L_draft, seq_len + L_draft)
+
+    def test_context_attention(self):
+        """Draft tokens should be able to attend to all context tokens."""
+        B, seq_len = 1, 5
+        num_blocks, block_size = 2, 3
+        device = "cpu"
+
+        mask = DSparkAttentionMask.create_dspark_attention_mask(
+            batch_size=B, seq_len=seq_len,
+            num_blocks=num_blocks, block_size=block_size,
+            device=torch.device(device),
+        )
+        # All draft positions should have 0.0 for all context positions
+        context_slice = mask[0, 0, :, :seq_len]
+        assert (context_slice == 0.0).all()
+
+    def test_intra_block_attention(self):
+        """Draft tokens in the same block should attend to each other."""
+        B, seq_len = 1, 5
+        num_blocks, block_size = 2, 3
+        device = "cpu"
+
+        mask = DSparkAttentionMask.create_dspark_attention_mask(
+            batch_size=B, seq_len=seq_len,
+            num_blocks=num_blocks, block_size=block_size,
+            device=torch.device(device),
+        )
+        L_ctx = seq_len
+
+        # Block 0: positions 0,1,2 should attend to each other
+        intra_block_0 = mask[0, 0, 0:3, L_ctx:L_ctx+3]
+        assert (intra_block_0 == 0.0).all(), "Block 0 intra-attention should be 0"
+
+        # Block 1: positions 3,4,5 should attend to each other
+        intra_block_1 = mask[0, 0, 3:6, L_ctx+3:L_ctx+6]
+        assert (intra_block_1 == 0.0).all(), "Block 1 intra-attention should be 0"
+
+    def test_cross_block_no_attention(self):
+        """Draft tokens should NOT attend to draft tokens in other blocks."""
+        B, seq_len = 1, 5
+        num_blocks, block_size = 2, 3
+        device = "cpu"
+
+        mask = DSparkAttentionMask.create_dspark_attention_mask(
+            batch_size=B, seq_len=seq_len,
+            num_blocks=num_blocks, block_size=block_size,
+            device=torch.device(device),
+        )
+        L_ctx = seq_len
+
+        # Block 0 should NOT attend to Block 1's draft tokens
+        cross_block = mask[0, 0, 0:3, L_ctx+3:L_ctx+6]
+        assert (cross_block == float("-inf")).all(), \
+            "Cross-block attention should be -inf"
+
+        # Block 1 should NOT attend to Block 0's draft tokens
+        cross_block_2 = mask[0, 0, 3:6, L_ctx:L_ctx+3]
+        assert (cross_block_2 == float("-inf")).all(), \
+            "Cross-block attention should be -inf"

tests/test_sampling.py

@@ -0,0 +1,65 @@
+"""Tests for sampling utilities."""
+
+import torch
+import pytest
+
+from uraionspec.utils.sampling import (
+    logits_to_probs,
+    sample_tokens,
+    sample_residual,
+    gather_token_probs,
+)
+
+
+class TestSampling:
+    """Test sampling utilities."""
+
+    def test_logits_to_probs_greedy(self):
+        logits = torch.tensor([[0.0, 10.0, 0.0]])
+        probs = logits_to_probs(logits, temperature=0.0)
+        assert probs.shape == (1, 3)
+        assert probs[0, 1] == 1.0  # argmax at index 1
+
+    def test_logits_to_probs_temperature(self):
+        logits = torch.randn(2, 100)
+        probs = logits_to_probs(logits, temperature=1.0)
+        assert probs.shape == (2, 100)
+        assert torch.allclose(probs.sum(dim=-1), torch.ones(2))
+
+    def test_sample_tokens_greedy(self):
+        logits = torch.randn(2, 5, 50)
+        tokens = sample_tokens(logits, temperature=0.0)
+        assert tokens.shape == (2, 5)
+        assert (tokens >= 0).all() and (tokens < 50).all()
+
+    def test_sample_tokens_temperature(self):
+        logits = torch.randn(2, 50)
+        tokens = sample_tokens(logits, temperature=1.0)
+        assert tokens.shape == (2,)
+        assert (tokens >= 0).all() and (tokens < 50).all()
+
+    def test_sample_tokens_2d(self):
+        logits = torch.randn(3, 100)
+        tokens = sample_tokens(logits, temperature=0.5)
+        assert tokens.shape == (3,)
+
+    def test_sample_residual(self):
+        target = torch.softmax(torch.randn(2, 50) + 2, dim=-1)
+        draft = torch.softmax(torch.randn(2, 50), dim=-1)
+        tokens = sample_residual(target, draft)
+        assert tokens.shape == (2,)
+        assert (tokens >= 0).all() and (tokens < 50).all()
+
+    def test_sample_residual_identical(self):
+        """When target == draft, residual should fall back to target."""
+        probs = torch.softmax(torch.randn(2, 50), dim=-1)
+        tokens = sample_residual(probs, probs)
+        assert tokens.shape == (2,)
+
+    def test_gather_token_probs(self):
+        probs = torch.tensor([[0.1, 0.7, 0.2], [0.3, 0.3, 0.4]])
+        token_ids = torch.tensor([1, 2])
+        gathered = gather_token_probs(probs, token_ids)
+        assert gathered.shape == (2,)
+        assert gathered[0].item() == pytest.approx(0.7)
+        assert gathered[1].item() == pytest.approx(0.4)