Upload tests

tests/test_atc_fusion.py

@@ -0,0 +1,187 @@
+import unittest
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from predictor_training.atc_fusion import ATCFusion, MultiInputMLP, StateTransport
+from predictor_training.single_block import SingleBlockPredictor
+from wan.modules.model import rope_params
+
+
+def small_freqs(head_dim: int = 8) -> torch.Tensor:
+    return torch.cat(
+        [
+            rope_params(32, head_dim - 4 * (head_dim // 6)),
+            rope_params(32, 2 * (head_dim // 6)),
+            rope_params(32, 2 * (head_dim // 6)),
+        ],
+        dim=1,
+    )
+
+
+def reference_attention(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+) -> torch.Tensor:
+    scores = torch.einsum("blhd,bshd->bhls", q, k) / q.shape[-1] ** 0.5
+    probability = scores.float().softmax(dim=-1).to(v.dtype)
+    return torch.einsum("bhls,bshd->blhd", probability, v)
+
+
+class DummyBlock(nn.Module):
+    num_heads = 1
+
+    def forward(self, value: torch.Tensor, **_kwargs) -> torch.Tensor:
+        return value
+
+
+class ATCFusionTest(unittest.TestCase):
+    def setUp(self) -> None:
+        torch.manual_seed(0)
+        self.grid = torch.tensor([[2, 2, 2]], dtype=torch.long)
+        self.freqs = small_freqs()
+        self.current = torch.randn(1, 8, 8)
+        self.previous = torch.randn(1, 8, 8)
+
+    def test_multi_input_mlp_matches_explicit_concat(self) -> None:
+        module = MultiInputMLP(3, dim=8, hidden_dim=12)
+        values = [torch.randn(2, 5, 8) for _ in range(3)]
+        combined_weight = torch.cat(
+            [layer.weight for layer in module.inputs], dim=1
+        )
+        combined_bias = sum(layer.bias for layer in module.inputs)
+        reference_hidden = F.linear(
+            torch.cat(values, dim=-1), combined_weight, combined_bias
+        )
+        reference = module.out(module.activation(reference_hidden))
+        torch.testing.assert_close(module(*values), reference)
+
+    @patch(
+        "predictor_training.atc_fusion.attention",
+        side_effect=reference_attention,
+    )
+    def test_last_frame_scope_cannot_read_earlier_frames(self, _mock) -> None:
+        transport = StateTransport(
+            dim=8,
+            num_heads=1,
+            previous_scope="last_frame",
+            gradient_checkpointing=False,
+            diagnostic_queries=4,
+        ).eval()
+        first, _ = transport(
+            self.current, self.previous, self.grid, self.freqs, current_start=8
+        )
+        perturbed_early = self.previous.clone()
+        perturbed_early[:, :4] += torch.randn_like(perturbed_early[:, :4])
+        second, _ = transport(
+            self.current,
+            perturbed_early,
+            self.grid,
+            self.freqs,
+            current_start=8,
+        )
+        torch.testing.assert_close(first, second, rtol=0, atol=0)
+
+        perturbed_last = self.previous.clone()
+        perturbed_last[:, 4:] += torch.randn_like(perturbed_last[:, 4:])
+        third, _ = transport(
+            self.current,
+            perturbed_last,
+            self.grid,
+            self.freqs,
+            current_start=8,
+        )
+        self.assertFalse(torch.equal(first, third))
+
+    @patch(
+        "predictor_training.atc_fusion.attention",
+        side_effect=reference_attention,
+    )
+    def test_zero_initialized_correction_bypass(self, _mock) -> None:
+        fusion = ATCFusion(
+            dim=8,
+            num_heads=1,
+            freq_dim=8,
+            mlp_hidden_dim=16,
+            gate_hidden_dim=4,
+            previous_scope="chunk",
+            gradient_checkpointing=False,
+        ).eval()
+        condition = torch.randn_like(self.current)
+        backbone, gated_correction = fusion(
+            self.current,
+            torch.randn_like(self.current),
+            self.previous,
+            condition,
+            torch.tensor([62.5]),
+            self.grid,
+            self.freqs,
+            current_start=8,
+        )
+        self.assertEqual(backbone.shape, self.current.shape)
+        torch.testing.assert_close(
+            gated_correction,
+            torch.zeros_like(gated_correction),
+            rtol=0,
+            atol=0,
+        )
+        self.assertAlmostEqual(float(fusion.last_gate.mean()), 0.3, delta=0.01)
+
+    def test_atc_token_gate_is_a_stage1_parameter(self) -> None:
+        predictor = SingleBlockPredictor(
+            DummyBlock(),
+            dim=8,
+            gradient_checkpointing=False,
+            input_variant="atc",
+            atc_freq_dim=8,
+            atc_mlp_hidden_dim=16,
+            atc_gate_hidden_dim=4,
+        )
+        token_gate_ids = {
+            id(parameter)
+            for parameter in predictor.fusion.token_gate.parameters()
+        }
+        stage1_ids = {
+            id(parameter) for parameter in predictor.stage1_input_parameters()
+        }
+        self.assertTrue(token_gate_ids)
+        self.assertTrue(token_gate_ids.issubset(stage1_ids))
+        self.assertEqual(predictor.gate_parameters(), [])
+
+    @patch(
+        "predictor_training.atc_fusion.attention",
+        side_effect=reference_attention,
+    )
+    def test_zero_initialized_predictor_returns_anchor(self, _mock) -> None:
+        predictor = SingleBlockPredictor(
+            DummyBlock(),
+            dim=8,
+            gradient_checkpointing=False,
+            input_variant="atc",
+            atc_freq_dim=8,
+            atc_mlp_hidden_dim=16,
+            atc_gate_hidden_dim=4,
+        ).eval()
+        anchor = torch.randn_like(self.current)
+        empty_history = torch.empty(1, 0, 1, 8)
+        prediction = predictor(
+            current_tokens=self.current,
+            anchor_hidden=anchor,
+            previous_hidden=self.previous,
+            timestep_modulation=torch.randn(1, 2, 6, 8),
+            grid_sizes=self.grid,
+            freqs=self.freqs,
+            history_k=empty_history,
+            history_v=empty_history,
+            cross_k=torch.empty(1, 0, 1, 8),
+            cross_v=torch.empty(1, 0, 1, 8),
+            current_start=8,
+            condition_tokens=torch.randn_like(self.current),
+            anchor_distance=torch.tensor([62.5]),
+        )
+        torch.testing.assert_close(prediction, anchor, rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_confidence_lr_schedule.py

@@ -0,0 +1,41 @@
+import unittest
+
+from scripts.train_confidence_token_lazy_ddp import learning_rate_at_step
+
+
+class ConfidenceLearningRateTest(unittest.TestCase):
+    def value(self, step: int) -> float:
+        return learning_rate_at_step(
+            step=step,
+            total_steps=1440,
+            schedule="warmup_cosine",
+            peak_lr=1e-4,
+            warmup_steps=36,
+            warmup_start_lr=1e-5,
+            min_lr=1e-6,
+        )
+
+    def test_warmup_endpoints(self) -> None:
+        self.assertAlmostEqual(self.value(0), 1e-5)
+        self.assertAlmostEqual(self.value(35), 1e-4)
+
+    def test_cosine_is_monotonic_and_reaches_floor(self) -> None:
+        values = [self.value(step) for step in range(35, 1440)]
+        self.assertTrue(all(left >= right for left, right in zip(values, values[1:])))
+        self.assertAlmostEqual(values[-1], 1e-6)
+
+    def test_constant_schedule_is_unchanged(self) -> None:
+        value = learning_rate_at_step(
+            step=503,
+            total_steps=1440,
+            schedule="constant",
+            peak_lr=3e-4,
+            warmup_steps=0,
+            warmup_start_lr=1e-5,
+            min_lr=1e-6,
+        )
+        self.assertEqual(value, 3e-4)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_confidence_token.py

@@ -0,0 +1,60 @@
+import unittest
+
+import torch
+
+from predictor_training.confidence import ConfidenceTokenHead
+
+
+class ConfidenceTokenHeadTest(unittest.TestCase):
+    def setUp(self) -> None:
+        torch.manual_seed(0)
+        self.head = ConfidenceTokenHead(
+            dim=16,
+            token_dim=8,
+            context_dim=4,
+            num_heads=2,
+            ffn_dim=32,
+            dropout=0.0,
+            num_steps=3,
+        )
+
+    def test_output_shape_and_predictor_features_are_detached(self) -> None:
+        transformed = torch.randn(2, 7, 16, requires_grad=True)
+        predicted = torch.randn(2, 7, 16, requires_grad=True)
+        anchor = torch.randn(2, 7, 16, requires_grad=True)
+        output = self.head(
+            transformed_hidden=transformed,
+            pred_hidden=predicted,
+            anchor_hidden=anchor,
+            chunk_position=torch.tensor([0.0, 1.0]),
+            step_id=torch.tensor([1, 3]),
+        )
+        self.assertEqual(output.shape, (2,))
+        output.sum().backward()
+        self.assertIsNone(transformed.grad)
+        self.assertIsNone(predicted.grad)
+        self.assertIsNone(anchor.grad)
+        self.assertTrue(any(parameter.grad is not None for parameter in self.head.parameters()))
+
+    def test_rejects_unsupported_step(self) -> None:
+        value = torch.randn(1, 7, 16)
+        with self.assertRaisesRegex(ValueError, "supports step_id"):
+            self.head(
+                transformed_hidden=value,
+                pred_hidden=value,
+                anchor_hidden=value,
+                chunk_position=torch.tensor([0.0]),
+                step_id=torch.tensor([4]),
+            )
+
+    def test_production_dimensions_and_parameter_count(self) -> None:
+        head = ConfidenceTokenHead()
+        self.assertEqual(head.token_projection[0].in_features, 1536)
+        self.assertEqual(head.token_projection[0].out_features, 512)
+        self.assertEqual(head.output[0].in_features, 1156)
+        count = sum(parameter.numel() for parameter in head.parameters())
+        self.assertEqual(count, 4_872_065)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_disca_input_variant.py

@@ -0,0 +1,60 @@
+import unittest
+
+import torch
+from torch import nn
+
+from predictor_training.single_block import SingleBlockPredictor
+
+
+class EchoBlock(nn.Module):
+    def forward(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        return value
+
+
+def build_model(input_variant: str) -> SingleBlockPredictor:
+    model = SingleBlockPredictor(
+        EchoBlock(),
+        dim=8,
+        gradient_checkpointing=False,
+        input_variant=input_variant,
+    ).eval()
+    with torch.no_grad():
+        model.residual_out.weight.copy_(torch.eye(8))
+    return model
+
+
+def inputs() -> dict[str, object]:
+    return {
+        "current_tokens": torch.randn(2, 3, 8),
+        "anchor_hidden": torch.randn(2, 3, 8),
+        "timestep_modulation": torch.randn(2, 6, 8),
+        "grid_sizes": torch.tensor([[1, 1, 3], [1, 1, 3]]),
+        "freqs": torch.empty(0),
+        "history_k": torch.randn(2, 2, 1, 8),
+        "history_v": torch.randn(2, 2, 1, 8),
+        "cross_k": torch.randn(2, 1, 1, 8),
+        "cross_v": torch.randn(2, 1, 1, 8),
+        "current_start": 2,
+    }
+
+
+class DiscaInputVariantTest(unittest.TestCase):
+    def test_disca_output_ignores_previous_chunk_hidden(self) -> None:
+        torch.manual_seed(0)
+        model = build_model(input_variant="disca")
+        common = inputs()
+        first = model(previous_hidden=torch.randn(2, 3, 8), **common)
+        second = model(previous_hidden=torch.randn(2, 3, 8), **common)
+        self.assertTrue(torch.equal(first, second))
+
+    def test_disca_physically_removes_previous_chunk_channel(self) -> None:
+        baseline = build_model(input_variant="self_forcing")
+        disca = build_model(input_variant="disca")
+        self.assertEqual(baseline.fusion.proj_in.weight.shape, (16, 24))
+        self.assertEqual(disca.fusion.proj_in.weight.shape, (16, 16))
+        self.assertTrue(hasattr(baseline.fusion, "previous_norm"))
+        self.assertFalse(hasattr(disca.fusion, "previous_norm"))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_naive_vbench_policies.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+import unittest
+
+from scripts.naive_vbench_policies import (
+    ALL_STRATEGY_NAMES,
+    EVALUATION_POLICIES,
+    POLICIES_BY_NAME,
+    raw_timestep_schedule,
+)
+
+
+class NaiveVBenchPoliciesTest(unittest.TestCase):
+    def test_timestep_skipping_schedules(self) -> None:
+        self.assertEqual(raw_timestep_schedule(3), (1000, 666, 333))
+        self.assertEqual(raw_timestep_schedule(2), (1000, 500))
+        self.assertEqual(raw_timestep_schedule(1), (1000,))
+        with self.assertRaises(ValueError):
+            raw_timestep_schedule(0)
+
+    def test_strategy_names_are_stable_and_unique(self) -> None:
+        self.assertEqual(
+            ALL_STRATEGY_NAMES,
+            (
+                "ffff",
+                "naive_3step",
+                "naive_2step",
+                "naive_1step",
+                "naive_cache_3step_frff",
+                "naive_cache_2step_frrf",
+                "naive_cache_1step_frrr",
+            ),
+        )
+        self.assertEqual(len(POLICIES_BY_NAME), len(ALL_STRATEGY_NAMES))
+
+    def test_cache_patterns_and_call_counts(self) -> None:
+        expected = {
+            "naive_cache_3step_frff": ("FRFF", 3, 1, 21, 7),
+            "naive_cache_2step_frrf": ("FRRF", 2, 2, 14, 14),
+            "naive_cache_1step_frrr": ("FRRR", 1, 3, 7, 21),
+        }
+        for name, (pattern, full, reuse, full_video, reuse_video) in expected.items():
+            policy = POLICIES_BY_NAME[name]
+            self.assertEqual(policy.cache_pattern, pattern)
+            self.assertEqual(policy.full_calls_per_chunk, full)
+            self.assertEqual(policy.reuse_calls_per_chunk, reuse)
+            self.assertEqual(7 * policy.full_calls_per_chunk, full_video)
+            self.assertEqual(7 * policy.reuse_calls_per_chunk, reuse_video)
+
+    def test_skipping_full_call_counts(self) -> None:
+        expected = {
+            "naive_3step": 21,
+            "naive_2step": 14,
+            "naive_1step": 7,
+        }
+        for policy in EVALUATION_POLICIES[:3]:
+            self.assertEqual(7 * policy.full_calls_per_chunk, expected[policy.name])
+            self.assertEqual(policy.reuse_calls_per_chunk, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_predictor_stage2_dmd.py

@@ -0,0 +1,157 @@
+import json
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+import torch
+from safetensors.torch import save_file
+
+from scripts import train_layer17_predictor_stage2_dmd as stage2
+from scripts.train_confidence_token_lazy_ddp import (
+    read_predictor_config as read_confidence_predictor_config,
+)
+from utils.wan_wrapper import resolve_wan_model_path
+
+
+class PredictorStage2DMDTest(unittest.TestCase):
+    def test_wan_model_path_accepts_names_and_absolute_paths(self) -> None:
+        self.assertEqual(
+            resolve_wan_model_path("Wan2.1-T2V-1.3B"),
+            Path("wan_models/Wan2.1-T2V-1.3B"),
+        )
+        self.assertEqual(
+            resolve_wan_model_path("/data1/chenzhuo/Wan2.1/Wan2.1-T2V-14B"),
+            Path("/data1/chenzhuo/Wan2.1/Wan2.1-T2V-14B"),
+        )
+
+    def test_prompt_pool_uses_seeded_full_file_permutation(self) -> None:
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "prompts.txt"
+            source = [f"prompt {index}" for index in range(20)]
+            path.write_text("\n".join(source) + "\n", encoding="utf-8")
+
+            first = stage2.load_prompt_pool(path, prompt_count=0, prompt_seed=0)
+            second = stage2.load_prompt_pool(path, prompt_count=0, prompt_seed=0)
+
+            self.assertEqual(first, second)
+            self.assertEqual(set(first), set(source))
+            self.assertNotEqual(first, source)
+
+    def test_prompt_draws_do_not_overlap_between_critic_and_student(self) -> None:
+        sample_ids = []
+        for step in range(2):
+            for batch_slot in range(6):
+                for accumulation_index in range(2):
+                    for rank in range(4):
+                        sample_ids.append(
+                            stage2.prompt_sample_id(
+                                student_step=step,
+                                batch_slot=batch_slot,
+                                global_batch_size=8,
+                                batches_per_step=6,
+                                accumulation_index=accumulation_index,
+                                world_size=4,
+                                rank=rank,
+                                prompt_pool_size=248_221,
+                            )
+                        )
+
+        self.assertEqual(sample_ids, list(range(96)))
+
+    def test_reads_atc_metadata_and_validates_variant(self) -> None:
+        config = {
+            "source_layer": 17,
+            "input_variant": "atc",
+            "atc_previous_scope": "last_frame",
+        }
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "predictor.safetensors"
+            save_file(
+                {"weight": torch.zeros(1)},
+                path,
+                metadata={"predictor_config": json.dumps(config)},
+            )
+            self.assertEqual(stage2.read_predictor_config(path, "atc"), config)
+            self.assertEqual(read_confidence_predictor_config(path, "atc"), config)
+            with self.assertRaisesRegex(ValueError, "variant mismatch"):
+                stage2.read_predictor_config(path, "self_forcing")
+
+    def test_stage2_metadata_is_confidence_trainer_compatible(self) -> None:
+        config = {
+            "source_layer": 17,
+            "input_variant": "atc",
+            "atc_previous_scope": "chunk",
+        }
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "predictor_ema.safetensors"
+            save_file(
+                {"weight": torch.zeros(1)},
+                path,
+                metadata=stage2.predictor_metadata(
+                    step=2000, predictor_config=config, ema_decay=0.99
+                ),
+            )
+            self.assertEqual(read_confidence_predictor_config(path, "atc"), config)
+
+    def test_online_step_matches_stage1_anchor_distance_definition(self) -> None:
+        noisy = torch.zeros(1, 3, 1, 1, 1)
+        timestep = torch.full((1, 3), 625.0)
+        anchor_timestep = torch.full((1, 3), 833.0)
+        hidden = torch.zeros(1, 2, 4)
+        history = {
+            "k": torch.zeros(1, 5 * stage2.TOKENS_PER_CHUNK, 1, 1),
+            "v": torch.zeros(1, 5 * stage2.TOKENS_PER_CHUNK, 1, 1),
+            "local_end_index": torch.tensor([5 * stage2.TOKENS_PER_CHUNK]),
+        }
+        cross = {"k": torch.zeros(1), "v": torch.zeros(1)}
+        expected = (hidden, noisy)
+
+        with mock.patch.object(
+            stage2, "forward_predictor", return_value=expected
+        ) as forward:
+            actual = stage2.online_predictor_step(
+                predictor=mock.Mock(),
+                teacher=mock.Mock(),
+                noisy_input=noisy,
+                timestep=timestep,
+                anchor_timestep=anchor_timestep,
+                anchor_hidden=hidden,
+                previous_hidden=hidden,
+                history_cache=history,
+                cross_cache=cross,
+                chunk=4,
+            )
+
+        self.assertIs(actual, expected)
+        batch = forward.call_args.args[1]
+        torch.testing.assert_close(batch["anchor_distance"], torch.tensor([208.0]))
+        self.assertEqual(batch["history_k"].shape[1], 4 * stage2.TOKENS_PER_CHUNK)
+        self.assertEqual(batch["chunk"], 4)
+
+    def test_dmd_loss_excludes_full_only_first_chunk(self) -> None:
+        generated = torch.zeros(
+            1,
+            stage2.NUM_CHUNKS * stage2.FRAMES_PER_CHUNK,
+            1,
+            1,
+            1,
+            dtype=torch.float64,
+            requires_grad=True,
+        )
+        gradient = torch.ones_like(generated)
+        loss = stage2.dmd_loss(generated, gradient)
+        loss.backward()
+
+        self.assertAlmostEqual(float(loss), 0.5)
+        self.assertTrue(
+            torch.equal(
+                generated.grad[:, : stage2.FRAMES_PER_CHUNK],
+                torch.zeros_like(generated.grad[:, : stage2.FRAMES_PER_CHUNK]),
+            )
+        )
+        self.assertTrue(torch.all(generated.grad[:, stage2.FRAMES_PER_CHUNK :] > 0))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_vbench8_protocol.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+import unittest
+
+from scripts.build_vbench8_extended_mapping import build_mapping
+from scripts.summarize_vbench8_generation import (
+    denoise_dit_latency_ms,
+    policy_latency_ms,
+    speedup_percent,
+)
+from scripts.vbench8_protocol import aggregate_selected_score, normalize_scores
+
+
+class VBench8ProtocolTest(unittest.TestCase):
+    def test_mapping_counts_and_extended_prompt_alignment(self) -> None:
+        short = [f"short prompt {index}" for index in range(946)]
+        extended = [f"extended prompt {index}" for index in range(946)]
+        info = [{"prompt_en": prompt, "dimension": ["temporal_style"]} for prompt in short]
+        positions = {
+            "subject_consistency": range(0, 72),
+            "overall_consistency": range(72, 165),
+            "scene": range(165, 251),
+        }
+        for suite, indices in positions.items():
+            for index in indices:
+                info[index] = {
+                    "prompt_en": short[index],
+                    "dimension": [suite],
+                }
+                if suite == "scene":
+                    info[index]["auxiliary_info"] = {
+                        "scene": {"scene": {"scene": f"scene-{index}"}}
+                    }
+        mapping = build_mapping(
+            short_prompts=short,
+            extended_prompts=extended,
+            vbench_info=info,
+        )
+        self.assertEqual(len(mapping), 251)
+        self.assertEqual(
+            {suite: sum(row["prompt_suite"] == suite for row in mapping)
+             for suite in positions},
+            {"subject_consistency": 72, "overall_consistency": 93, "scene": 86},
+        )
+        self.assertEqual(mapping[0]["original_prompt"], "short prompt 0")
+        self.assertEqual(mapping[0]["extended_prompt"], "extended prompt 0")
+        self.assertIn("auxiliary_info", mapping[-1])
+
+    def test_mapping_rejects_order_mismatch(self) -> None:
+        short = [f"short prompt {index}" for index in range(946)]
+        extended = [f"extended prompt {index}" for index in range(946)]
+        info = [{"prompt_en": prompt, "dimension": ["temporal_style"]} for prompt in short]
+        info[10]["prompt_en"] = "wrong order"
+        with self.assertRaises(ValueError):
+            build_mapping(short_prompts=short, extended_prompts=extended, vbench_info=info)
+
+    def test_normalization_motion(self) -> None:
+        raw = {dimension: 0.0 for dimension in (
+            "subject_consistency", "background_consistency", "motion_smoothness",
+            "dynamic_degree", "aesthetic_quality", "imaging_quality", "scene",
+            "overall_consistency",
+        )}
+        raw["motion_smoothness"] = 0.95
+        normalized = normalize_scores(raw)
+        expected = (0.95 - 0.7060) / (0.9975 - 0.7060)
+        self.assertAlmostEqual(normalized["motion_smoothness"], expected)
+
+    def test_aggregate_all_normalized_one(self) -> None:
+        raw = {
+            "subject_consistency": 1.0,
+            "background_consistency": 1.0,
+            "motion_smoothness": 0.9975,
+            "dynamic_degree": 1.0,
+            "aesthetic_quality": 1.0,
+            "imaging_quality": 1.0,
+            "scene": 0.8222,
+            "overall_consistency": 0.3640,
+        }
+        aggregate = aggregate_selected_score(raw)
+        self.assertAlmostEqual(aggregate["quality_score"], 1.0)
+        self.assertAlmostEqual(aggregate["semantic_score"], 1.0)
+        self.assertAlmostEqual(aggregate["selected_vbench_score"], 1.0)
+        self.assertAlmostEqual(aggregate["selected_vbench_percent"], 100.0)
+
+    def test_latency_excludes_context_and_includes_confidence(self) -> None:
+        generation = {
+            "full_dit_time_ms": 100.0,
+            "predictor_time_ms": 20.0,
+            "confidence_head_time_ms": 3.0,
+            "context_dit_time_ms": 500.0,
+        }
+        self.assertEqual(denoise_dit_latency_ms(generation), 120.0)
+        self.assertEqual(policy_latency_ms(generation), 123.0)
+
+    def test_policy_latency_speedup_uses_ratio_of_means(self) -> None:
+        self.assertAlmostEqual(speedup_percent([60.0, 100.0], [100.0, 100.0]), 20.0)
+
+
+if __name__ == "__main__":
+    unittest.main()