Module 11. Speculative Decoding and Quantization

import importlib.util
import subprocess
import sys

if importlib.util.find_spec("numpy") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "numpy>=1.24"])

import hashlib
import json
import math
from dataclasses import dataclass
from typing import Sequence

import numpy as np

np.set_printoptions(precision=4, suppress=True)

REFERENCE_JSON = '{"edge_cases":{"identical_distributions":{"acceptance_by_token":[1.0,1.0],"alpha":1.0,"balance":[0.55,0.45],"p":[0.55,0.45],"q":[0.55,0.45],"residual":[0.0,0.0],"residual_mass":0.0,"total_variation":0.0},"q_zero_support_gap":{"acceptance_by_token":[0.5,0.0],"alpha":0.5,"balance":[0.5,0.5],"p":[0.5,0.5],"q":[1.0,0.0],"residual":[0.0,1.0],"residual_mass":0.5,"total_variation":0.5}},"expected_tokens":{"alpha_one_gamma_4":5.0,"alpha_zero_gamma_4":1.0,"calibration_alpha_0_7_gamma_5":2.9411699999999996,"grid":{"0.6":{"1":1.6,"3":2.176,"5":2.3833599999999997,"7":2.4580096},"0.7":{"1":1.7,"3":2.533,"5":2.9411699999999996,"7":3.1411732999999997},"0.8":{"1":1.8,"3":2.9520000000000004,"5":3.6892800000000006,"7":4.161139200000001},"0.9":{"1":1.9,"3":3.439,"5":4.68559,"7":5.6953279000000006}},"trace_alpha_0_85_gamma_4":3.70863125},"monte_carlo":{"cycles":20000,"distribution_atol":0.01,"gamma":4,"mean_length_atol":0.04,"seed":7},"quantization":{"e2m1_dequantized":[0.0,0.5937278925048495,0.0,-7.124734710058194,-0.5937278925048495,-1.187455785009699,0.0,1.187455785009699,-0.5937278925048495,-0.5937278925048495,0.5937278925048495,0.5937278925048495,0.0,-1.187455785009699,0.0,0.5937278925048495],"e2m1_scale":[1.187455785009699],"group_size":16,"int4_dequantized":[0.0,0.0,-0.0,-7.124734710058194,-0.0,-1.0178192442940277,0.0,1.0178192442940277,-0.0,-1.0178192442940277,0.0,0.0,0.0,-1.0178192442940277,-0.0,1.0178192442940277],"int4_over_e2m1_mse_ratio":3.3179967976624045,"int4_scale":[1.0178192442940277],"mse_e2m1":0.02581465276871833,"mse_int4":0.08565293521937434,"seed":7,"vector":[0.0012301533574825742,0.2987455375084699,-0.2741378553622176,-7.124734710058194,-0.45467078517172255,-0.9916465549964624,0.060143602597438485,1.3402152455545335,-0.49220651855132963,-0.6204748998199404,0.4898420501851982,0.35688700816006075,0.10541424899789856,-0.9304680447082047,-0.02925182246327349,0.6953031944582878]},"roofline":{"bandwidth_TB_per_second_decimal":3.35,"bandwidth_bytes_per_second":3350000000000.0,"combined_nvfp4_mtp_scenario_tokens_per_second":1329.3650793650793,"combined_scenario_vs_bf16":6.349206349206349,"formats":{"FP8":{"bits_per_weight":8.0,"payload_GB":8.0,"payload_GiB":7.450580596923828,"payload_bytes":8000000000.0,"weight_stream_ceiling_tokens_per_second":418.75},"MXFP4 (32 + E8M0)":{"bits_per_weight":4.25,"payload_GB":4.25,"payload_GiB":3.9581209421157837,"payload_bytes":4250000000.0,"weight_stream_ceiling_tokens_per_second":788.2352941176471},"NVFP4 (16 + E4M3)":{"bits_per_weight":4.5,"payload_GB":4.5,"payload_GiB":4.190951585769653,"payload_bytes":4500000000.0,"weight_stream_ceiling_tokens_per_second":744.4444444444445},"bf16":{"bits_per_weight":16.0,"payload_GB":16.0,"payload_GiB":14.901161193847656,"payload_bytes":16000000000.0,"weight_stream_ceiling_tokens_per_second":209.375},"ideal W4":{"bits_per_weight":4.0,"payload_GB":4.0,"payload_GiB":3.725290298461914,"payload_bytes":4000000000.0,"weight_stream_ceiling_tokens_per_second":837.5}},"parameter_count":8000000000,"speculative_bf16":{"alpha":0.8,"ceiling_tokens_per_second":594.1869230769231,"draft_cost":0.06,"gamma":5}},"schema_version":1,"speedup":{"external":{"alpha":0.8,"draft_cost":0.1,"gamma_star":6,"speedup_gamma_5":2.4595200000000004,"speedup_gamma_9":2.3490153094736845,"speedup_star":2.4696400000000005},"mtp_scenario":{"alpha":0.875,"draft_cost":0.05,"expected_tokens":1.875,"gamma":1,"speedup":1.7857142857142856}},"trace":{"acceptance_by_token":[1.0,0.8571428571428572,0.7499999999999999,1.0,0.5],"alpha":0.8500000000000001,"balance":[0.4,0.3,0.15,0.1,0.05],"p":[0.4,0.3,0.15,0.1,0.05],"q":[0.25,0.35,0.2,0.1,0.1],"residual":[1.0,0.0,0.0,0.0,0.0],"residual_mass":0.15000000000000002,"total_variation":0.15}}'
EXPECTED_SHA256 = '852f760a5ed5212da775f31a696f262eb9bb641a8ea495360b0330cd77ba400c'
actual_sha256 = hashlib.sha256(REFERENCE_JSON.encode("utf-8")).hexdigest()
assert actual_sha256 == EXPECTED_SHA256
REF = json.loads(REFERENCE_JSON)
MSG = {'reference': 'Embedded M11 contract verified', 'contracts': 'Probability contracts defined', 'trace': 'Rejection balance and edge cases verified', 'expectation': 'E, speedup, and optimal draft length verified', 'sampling': 'Monte Carlo agrees with the target distribution and E', 'quant': 'Grouped quantizers and the zero group verified', 'roofline': 'Weight-stream ceiling and units reproduced', 'negative': 'Negative tests passed'}
assert REF["schema_version"] == 1
print(MSG["reference"], actual_sha256[:12])

def _positive_int(name: str, value: int) -> int:
    """Return a strictly positive, non-boolean integer."""
    if isinstance(value, bool) or not isinstance(value, (int, np.integer)) or int(value) <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return int(value)


def _non_negative_int(name: str, value: int) -> int:
    """Return a non-negative, non-boolean integer."""
    if isinstance(value, bool) or not isinstance(value, (int, np.integer)) or int(value) < 0:
        raise ValueError(f"{name} must be a non-negative integer")
    return int(value)


def _finite_non_negative(name: str, value: float) -> float:
    """Return a finite real value greater than or equal to zero."""
    if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)):
        raise ValueError(f"{name} must be a real number")
    number = float(value)
    if not math.isfinite(number) or number < 0.0:
        raise ValueError(f"{name} must be finite and non-negative")
    return number


def _positive_finite(name: str, value: float) -> float:
    """Return a strictly positive finite real value."""
    number = _finite_non_negative(name, value)
    if number == 0.0:
        raise ValueError(f"{name} must be positive")
    return number


def _probability(name: str, value: float) -> float:
    """Return a finite probability in [0, 1]."""
    number = _finite_non_negative(name, value)
    if number > 1.0:
        raise ValueError(f"{name} must lie in [0, 1]")
    return number


def probability_vector(name: str, values: Sequence[float], *, atol: float = 1e-12) -> np.ndarray:
    """Validate a 1D probability vector without silently renormalizing it."""
    array = np.asarray(values, dtype=np.float64)
    if array.ndim != 1 or array.size < 2:
        raise ValueError(f"{name} must be a 1D vector with at least two entries")
    if not np.isfinite(array).all() or np.any(array < 0.0):
        raise ValueError(f"{name} must contain finite non-negative probabilities")
    total = float(array.sum(dtype=np.float64))
    if not math.isclose(total, 1.0, rel_tol=0.0, abs_tol=atol):
        raise ValueError(f"{name} must sum to 1; got {total}")
    return array.copy()


def validate_probability_pair(p_values: Sequence[float], q_values: Sequence[float]) -> tuple[np.ndarray, np.ndarray]:
    """Validate target and draft vectors and require equal vocabulary shape."""
    p = probability_vector("p", p_values)
    q = probability_vector("q", q_values)
    if p.shape != q.shape:
        raise ValueError(f"p and q must share a shape; got {p.shape} and {q.shape}")
    return p, q


@dataclass(frozen=True)
class RejectionTrace:
    """Exact one-token rejection decomposition for fixed p and q."""
    p: np.ndarray
    q: np.ndarray
    acceptance_by_token: np.ndarray
    alpha: float
    residual: np.ndarray
    residual_mass: float
    total_variation: float
    balance: np.ndarray


def rejection_trace(p_values: Sequence[float], q_values: Sequence[float]) -> RejectionTrace:
    """Build an exact balance, including q[t]=0 and p==q edge cases."""
    p, q = validate_probability_pair(p_values, q_values)
    acceptance = np.zeros_like(p)
    support = q > 0.0
    acceptance[support] = np.minimum(1.0, p[support] / q[support])

    overlap = np.minimum(p, q)
    alpha = float(overlap.sum(dtype=np.float64))
    positive = np.maximum(p - q, 0.0)
    residual_mass = float(positive.sum(dtype=np.float64))
    residual = positive / residual_mass if residual_mass > 0.0 else np.zeros_like(positive)
    balance = overlap + residual_mass * residual
    tv = 0.5 * float(np.abs(p - q).sum(dtype=np.float64))

    assert math.isclose(residual_mass, 1.0 - alpha, abs_tol=1e-12)
    assert np.allclose(balance, p, atol=1e-12, rtol=0.0)
    assert math.isclose(alpha, 1.0 - tv, abs_tol=1e-12)
    return RejectionTrace(p, q, acceptance, alpha, residual, residual_mass, tv, balance)


print(MSG['contracts'])

trace_ref = REF["trace"]
trace = rejection_trace(trace_ref["p"], trace_ref["q"])
assert np.allclose(trace.acceptance_by_token, trace_ref["acceptance_by_token"])
assert math.isclose(trace.alpha, trace_ref["alpha"], abs_tol=1e-12)
assert np.allclose(trace.residual, trace_ref["residual"])
assert np.allclose(trace.balance, trace.p)

support_ref = REF["edge_cases"]["q_zero_support_gap"]
support_gap = rejection_trace(support_ref["p"], support_ref["q"])
assert support_gap.q[1] == 0.0
assert support_gap.acceptance_by_token[1] == 0.0
assert np.allclose(support_gap.residual, [0.0, 1.0])
assert np.allclose(support_gap.balance, support_gap.p)

identical_ref = REF["edge_cases"]["identical_distributions"]
identical = rejection_trace(identical_ref["p"], identical_ref["q"])
assert identical.alpha == 1.0
assert identical.residual_mass == 0.0
assert np.array_equal(identical.residual, np.zeros_like(identical.residual))

print(f"alpha={trace.alpha:.2f} = 1-TV; residual={trace.residual}")
print(MSG["trace"])

def expected_tokens_per_cycle(alpha: float, gamma: int) -> float:
    """Return 1 + alpha + ... + alpha**gamma without an alpha==1 singularity."""
    a = _probability("alpha", alpha)
    g = _non_negative_int("gamma", gamma)
    return math.fsum(a**i for i in range(g + 1))


def speculative_speedup(alpha: float, gamma: int, draft_cost: float) -> float:
    """Return the transparent E/(1+gamma*c) scenario estimate."""
    g = _non_negative_int("gamma", gamma)
    c = _finite_non_negative("draft_cost", draft_cost)
    return expected_tokens_per_cycle(alpha, g) / (1.0 + g * c)


def optimal_gamma(alpha: float, draft_cost: float, gamma_max: int) -> tuple[int, float]:
    """Return the smallest draft length attaining maximal scenario speedup."""
    maximum = _non_negative_int("gamma_max", gamma_max)
    candidates = [(g, speculative_speedup(alpha, g, draft_cost)) for g in range(maximum + 1)]
    best_value = max(value for _, value in candidates)
    best_gamma = min(g for g, value in candidates if math.isclose(value, best_value, rel_tol=0.0, abs_tol=1e-15))
    return best_gamma, best_value


for alpha_text, row in REF["expected_tokens"]["grid"].items():
    for gamma_text, expected in row.items():
        assert math.isclose(
            expected_tokens_per_cycle(float(alpha_text), int(gamma_text)),
            expected,
            abs_tol=1e-12,
        )
assert expected_tokens_per_cycle(1.0, 4) == 5.0
assert expected_tokens_per_cycle(0.0, 4) == 1.0
assert optimal_gamma(0.0, 0.1, 9) == (0, 1.0)
assert optimal_gamma(0.8, 0.1, 0) == (0, 1.0)

external = REF["speedup"]["external"]
gamma_star, speedup_star = optimal_gamma(external["alpha"], external["draft_cost"], 10)
assert gamma_star == external["gamma_star"] == 6
assert math.isclose(speedup_star, external["speedup_star"], abs_tol=1e-12)
mtp = REF["speedup"]["mtp_scenario"]
assert math.isclose(
    speculative_speedup(mtp["alpha"], mtp["gamma"], mtp["draft_cost"]),
    mtp["speedup"],
    abs_tol=1e-12,
)
print(f"gamma*={gamma_star}; scenario={speedup_star:.3f}x; MTP={mtp['speedup']:.3f}x")
print(MSG["expectation"])

@dataclass(frozen=True)
class SpeculativeCycleResult:
    """Tokens and control-flow metadata from one fixed-distribution cycle."""
    tokens: tuple[int, ...]
    accepted_draft_tokens: int
    rejected: bool


def speculative_cycle_constant(
    p_values: Sequence[float],
    q_values: Sequence[float],
    gamma: int,
    rng: np.random.Generator,
) -> SpeculativeCycleResult:
    """Run one exact fixed-distribution speculative sampling cycle."""
    if not isinstance(rng, np.random.Generator):
        raise TypeError("rng must be numpy.random.Generator")
    g = _positive_int("gamma", gamma)
    local = rejection_trace(p_values, q_values)
    output: list[int] = []
    accepted = 0

    for _ in range(g):
        token = int(rng.choice(local.q.size, p=local.q))
        q_token = float(local.q[token])
        if q_token <= 0.0:
            raise AssertionError("sampled a zero-probability draft token")
        acceptance = min(1.0, float(local.p[token]) / q_token)
        if float(rng.random()) < acceptance:
            output.append(token)
            accepted += 1
            continue
        if local.residual_mass <= 0.0:
            raise AssertionError("rejection occurred with zero residual mass")
        output.append(int(rng.choice(local.p.size, p=local.residual)))
        return SpeculativeCycleResult(tuple(output), accepted, True)

    output.append(int(rng.choice(local.p.size, p=local.p)))
    return SpeculativeCycleResult(tuple(output), accepted, False)


mc = REF["monte_carlo"]
rng = np.random.default_rng(mc["seed"])
counts = np.zeros_like(trace.p)
total_tokens = 0
for _ in range(mc["cycles"]):
    result = speculative_cycle_constant(trace.p, trace.q, mc["gamma"], rng)
    for token in result.tokens:
        counts[token] += 1
    total_tokens += len(result.tokens)
empirical = counts / counts.sum()
mean_length = total_tokens / mc["cycles"]
assert np.allclose(empirical, trace.p, atol=mc["distribution_atol"])
assert abs(mean_length - expected_tokens_per_cycle(trace.alpha, mc["gamma"])) < mc["mean_length_atol"]

rng = np.random.default_rng(19)
for _ in range(200):
    result = speculative_cycle_constant(identical.p, identical.q, 4, rng)
    assert not result.rejected and len(result.tokens) == 5

rng = np.random.default_rng(23)
support_counts = np.zeros(2)
for _ in range(10_000):
    result = speculative_cycle_constant(support_gap.p, support_gap.q, 2, rng)
    for token in result.tokens:
        support_counts[token] += 1
assert np.allclose(support_counts / support_counts.sum(), support_gap.p, atol=0.015)

print(f"empirical={empirical}; mean tokens/cycle={mean_length:.3f}")
print(MSG["sampling"])

E2M1_GRID = np.array(
    [-6.0, -4.0, -3.0, -2.0, -1.5, -1.0, -0.5, 0.0,
      0.5,  1.0,  1.5,  2.0,  3.0,  4.0,  6.0],
    dtype=np.float64,
)


def _finite_vector(name: str, values: Sequence[float]) -> np.ndarray:
    """Return a non-empty finite one-dimensional float64 vector."""
    array = np.asarray(values, dtype=np.float64)
    if array.ndim != 1 or array.size == 0:
        raise ValueError(f"{name} must be a non-empty 1D vector")
    if not np.isfinite(array).all():
        raise ValueError(f"{name} must contain finite values")
    return array.copy()


@dataclass(frozen=True)
class GroupedQuantization:
    """Dequantized values and one positive scale per fixed-size group."""
    values: np.ndarray
    scales: np.ndarray
    group_size: int


def quantize_grouped_e2m1(values: Sequence[float], group_size: int = 16) -> GroupedQuantization:
    """Quantize with an educational E2M1 grid and one scale per group."""
    x = _finite_vector("values", values)
    group = _positive_int("group_size", group_size)
    if x.size % group != 0:
        raise ValueError("values length must be divisible by group_size")
    matrix = x.reshape(-1, group)
    max_abs = np.max(np.abs(matrix), axis=1, keepdims=True)
    scales = np.where(max_abs > 0.0, max_abs / 6.0, 1.0)
    normalized = matrix / scales
    indices = np.argmin(np.abs(normalized[..., None] - E2M1_GRID), axis=-1)
    values_q = E2M1_GRID[indices] * scales
    return GroupedQuantization(values_q.reshape(x.shape), scales[:, 0], group)


def quantize_grouped_symmetric_int4(values: Sequence[float], group_size: int = 16) -> GroupedQuantization:
    """Quantize to symmetric signed INT4 levels [-7, 7] per group."""
    x = _finite_vector("values", values)
    group = _positive_int("group_size", group_size)
    if x.size % group != 0:
        raise ValueError("values length must be divisible by group_size")
    matrix = x.reshape(-1, group)
    max_abs = np.max(np.abs(matrix), axis=1, keepdims=True)
    scales = np.where(max_abs > 0.0, max_abs / 7.0, 1.0)
    values_q = np.clip(np.rint(matrix / scales), -7.0, 7.0) * scales
    return GroupedQuantization(values_q.reshape(x.shape), scales[:, 0], group)


def mean_squared_error(reference: Sequence[float], approximation: Sequence[float]) -> float:
    """Return MSE after checking equal finite vector shapes."""
    ref = _finite_vector("reference", reference)
    app = _finite_vector("approximation", approximation)
    if ref.shape != app.shape:
        raise ValueError("reference and approximation must share a shape")
    return float(np.mean(np.square(ref - app), dtype=np.float64))


qref = REF["quantization"]
vector = np.asarray(qref["vector"], dtype=np.float64)
e2m1 = quantize_grouped_e2m1(vector, qref["group_size"])
int4 = quantize_grouped_symmetric_int4(vector, qref["group_size"])
assert np.allclose(e2m1.values, qref["e2m1_dequantized"])
assert np.allclose(int4.values, qref["int4_dequantized"])
assert math.isclose(mean_squared_error(vector, e2m1.values), qref["mse_e2m1"], abs_tol=1e-12)
assert math.isclose(mean_squared_error(vector, int4.values), qref["mse_int4"], abs_tol=1e-12)

zero_e2m1 = quantize_grouped_e2m1(np.zeros(32), 16)
zero_int4 = quantize_grouped_symmetric_int4(np.zeros(32), 16)
assert np.array_equal(zero_e2m1.values, np.zeros(32))
assert np.array_equal(zero_int4.values, np.zeros(32))
assert np.array_equal(zero_e2m1.scales, np.ones(2))
assert np.array_equal(zero_int4.scales, np.ones(2))

print(f"MSE E2M1={qref['mse_e2m1']:.4f}; INT4={qref['mse_int4']:.4f}; ratio={qref['int4_over_e2m1_mse_ratio']:.2f}x")
print(MSG["quant"])

def weight_payload_bytes(parameter_count: int, bits_per_weight: float) -> float:
    """Return the raw weight payload in bytes."""
    parameters = _positive_int("parameter_count", parameter_count)
    bits = _positive_finite("bits_per_weight", bits_per_weight)
    return parameters * bits / 8.0


def memory_bound_ceiling_tokens_per_second(
    bandwidth_bytes_per_second: float,
    parameter_count: int,
    bits_per_weight: float,
) -> float:
    """Return a weight-stream-only arithmetic ceiling in tokens/s."""
    bandwidth = _positive_finite("bandwidth_bytes_per_second", bandwidth_bytes_per_second)
    return bandwidth / weight_payload_bytes(parameter_count, bits_per_weight)


roof = REF["roofline"]
bandwidth = roof["bandwidth_bytes_per_second"]
parameters = roof["parameter_count"]
print(f"bandwidth = {roof['bandwidth_TB_per_second_decimal']:.2f} TB/s (decimal)")
for name, row in roof["formats"].items():
    payload = weight_payload_bytes(parameters, row["bits_per_weight"])
    ceiling = memory_bound_ceiling_tokens_per_second(bandwidth, parameters, row["bits_per_weight"])
    assert math.isclose(payload, row["payload_bytes"], abs_tol=1e-6)
    assert math.isclose(ceiling, row["weight_stream_ceiling_tokens_per_second"], abs_tol=1e-12)
    print(
        f"{name:24s} {row['bits_per_weight']:5.2f} bit/weight | "
        f"{row['payload_GB']:6.3f} GB = {row['payload_GiB']:6.3f} GiB | "
        f"ceiling {ceiling:7.1f} tok/s"
    )

bf16 = roof["formats"]["bf16"]["weight_stream_ceiling_tokens_per_second"]
spec = roof["speculative_bf16"]
actual_spec = speculative_speedup(spec["alpha"], spec["gamma"], spec["draft_cost"]) * bf16
assert math.isclose(actual_spec, spec["ceiling_tokens_per_second"], abs_tol=1e-12)
mtp = REF["speedup"]["mtp_scenario"]
nvfp4 = roof["formats"]["NVFP4 (16 + E4M3)"]["weight_stream_ceiling_tokens_per_second"]
combined = speculative_speedup(mtp["alpha"], mtp["gamma"], mtp["draft_cost"]) * nvfp4
assert math.isclose(combined, roof["combined_nvfp4_mtp_scenario_tokens_per_second"], abs_tol=1e-12)
assert math.isclose(combined / bf16, roof["combined_scenario_vs_bf16"], abs_tol=1e-12)
print(f"scenario: speculative bf16={actual_spec:.1f} tok/s; NVFP4+MTP={combined:.1f} tok/s ({combined/bf16:.2f}x)")
print(MSG["roofline"])

MINIMAL = {}
exec('"""M11 minimal probability and grouped-value teaching kernels."""\nimport numpy as np\n\n\ndef probability_vector(values, *, name="p", atol=1e-12):\n    """Validate a distribution without hidden clipping or renormalization."""\n    x = np.asarray(values, dtype=np.float64)\n    if x.ndim != 1 or x.size < 2 or not np.isfinite(x).all():\n        raise ValueError(f"{name}: expected a finite one-dimensional vector")\n    if np.any(x < 0) or not np.isclose(x.sum(), 1.0, rtol=0.0, atol=atol):\n        raise ValueError(f"{name}: probabilities must be >= 0 and sum to 1")\n    return x\n\n\ndef rejection_kernel(p_values, q_values):\n    """Acceptance probabilities, residual, and exact one-position balance."""\n    p = probability_vector(p_values, name="p")\n    q = probability_vector(q_values, name="q")\n    if p.shape != q.shape:\n        raise ValueError("p and q must use the same vocabulary")\n\n    acceptance = np.zeros_like(p)\n    support = q > 0  # q_t=0 means the proposer cannot emit token t\n    acceptance[support] = np.minimum(1.0, p[support] / q[support])\n\n    overlap = np.minimum(p, q)\n    positive = np.maximum(p - q, 0.0)\n    residual_mass = positive.sum()\n    residual = (\n        positive / residual_mass\n        if residual_mass > 0.0\n        else np.zeros_like(p)\n    )\n    balance = overlap + residual_mass * residual\n    np.testing.assert_allclose(balance, p, rtol=0.0, atol=1e-12)\n    return acceptance, residual, balance\n\n\ndef quantize_grouped_e2m1(values, group_size=16):\n    """Teaching E2M1-like grid with one scale per group."""\n    grid = np.array(\n        [-6, -4, -3, -2, -1.5, -1, -0.5, 0,\n          0.5, 1, 1.5, 2, 3, 4, 6],\n        dtype=np.float64,\n    )\n    x = np.asarray(values, dtype=np.float64)\n    if x.ndim != 1 or x.size == 0 or not np.isfinite(x).all():\n        raise ValueError("expected a non-empty finite one-dimensional vector")\n    if isinstance(group_size, bool) or not isinstance(group_size, (int, np.integer)) or group_size <= 0 or x.size % group_size:\n        raise ValueError("length must be divisible by a positive group_size")\n\n    groups = x.reshape(-1, group_size)\n    max_abs = np.max(np.abs(groups), axis=1, keepdims=True)\n    scale = np.where(max_abs > 0, max_abs / 6.0, 1.0)\n    normalized = groups / scale\n    codes = np.abs(normalized[..., None] - grid).argmin(axis=-1)\n    return (grid[codes] * scale).reshape(x.shape)\n', MINIMAL)
a, r, b = MINIMAL['rejection_kernel']([1-1e-14,1e-14], [1.,0.])
np.testing.assert_array_equal(r, [0.,1.])
for bad in (True, 1.5, 0):
    try:
        MINIMAL['quantize_grouped_e2m1'](np.zeros(16), bad)
    except ValueError:
        pass
    else:
        raise AssertionError(bad)
conditional = np.array([.9,.5])
assert np.isclose(1+np.cumprod(conditional).sum(),2.35)
assert not np.isclose(1+np.mean(conditional)+np.mean(conditional)**2,2.35)

def expect_error(exc_type: type[BaseException], fn, *args, **kwargs) -> str:
    """Require a callable to fail with the expected exception type."""
    try:
        fn(*args, **kwargs)
    except exc_type as exc:
        return str(exc)
    raise AssertionError(f"expected {exc_type.__name__} from {fn.__name__}")


expect_error(ValueError, probability_vector, "bad", [0.7, -0.2, 0.5])
expect_error(ValueError, probability_vector, "bad", [0.6, 0.3])
expect_error(ValueError, probability_vector, "bad", [0.5, np.nan, 0.5])
expect_error(ValueError, validate_probability_pair, [0.5, 0.5], [0.3, 0.3, 0.4])
expect_error(ValueError, expected_tokens_per_cycle, 1.1, 4)
expect_error(ValueError, expected_tokens_per_cycle, 0.8, -1)
expect_error(ValueError, speculative_speedup, 0.8, 4, -0.1)
expect_error(ValueError, quantize_grouped_e2m1, np.ones(15), 16)
expect_error(ValueError, quantize_grouped_symmetric_int4, [0.0, np.inf], 2)
expect_error(ValueError, memory_bound_ceiling_tokens_per_second, "3.35e12", 8_000_000_000, 16)
print(MSG["negative"])

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_rejection_kernel_explicit(
    p: np.ndarray,
    q: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Minimal exact rejection correction, assuming valid probability vectors."""
    acceptance = np.zeros_like(p, dtype=np.float64)
    support = q > 0
    acceptance[support] = np.minimum(1.0, p[support] / q[support])
    overlap = np.minimum(p, q)
    positive = np.maximum(p - q, 0.0)
    residual_mass = positive.sum()
    residual = positive / residual_mass if residual_mass > 0 else np.zeros_like(p)
    reconstructed = overlap + residual_mass * residual
    return acceptance, residual, reconstructed

def torch_rejection_kernel_explicit(
    p: torch.Tensor,
    q: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Minimal tensor counterpart of exact rejection correction."""
    acceptance = torch.zeros_like(p)
    support = q > 0
    acceptance[support] = torch.minimum(torch.ones_like(p[support]), p[support] / q[support])
    overlap = torch.minimum(p, q)
    positive = torch.clamp(p - q, min=0)
    residual_mass = positive.sum()
    residual = positive / residual_mass if bool(residual_mass > 0) else torch.zeros_like(p)
    reconstructed = overlap + residual_mass * residual
    return acceptance, residual, reconstructed

def torch_sample_speculative_token_explicit(
    p: torch.Tensor,
    q: torch.Tensor,
    proposal: int,
    *,
    generator: torch.Generator | None = None,
) -> tuple[int, bool]:
    """Sample one exact speculative step from the explicit rejection kernel."""
    if isinstance(proposal, bool) or not isinstance(proposal, int) or not 0 <= proposal < p.numel() or q[proposal] <= 0:
        raise ValueError("proposal must be an integer in the positive support of q")
    acceptance, residual, _ = torch_rejection_kernel_explicit(p, q)
    draw = torch.rand((), generator=generator, device=p.device, dtype=p.dtype)
    if draw < acceptance[proposal]:
        return int(proposal), True
    return int(torch.multinomial(residual, 1, generator=generator).item()), False

p_np = np.array([0.40,0.30,0.15,0.10,0.05]); q_np = np.array([0.25,0.35,0.20,0.10,0.10])
a_np, r_np, b_np = numpy_rejection_kernel_explicit(p_np, q_np)
a_t, r_t, b_t = torch_rejection_kernel_explicit(torch.tensor(p_np), torch.tensor(q_np))
np.testing.assert_allclose(a_np, a_t.numpy()); np.testing.assert_allclose(r_np, r_t.numpy()); np.testing.assert_allclose(b_np, b_t.numpy()); np.testing.assert_allclose(b_np, p_np)
gen = torch.Generator().manual_seed(18); counts = torch.zeros(5)
q = torch.tensor(q_np); p = torch.tensor(p_np)
for _ in range(12000):
    proposal = int(torch.multinomial(q, 1, generator=gen))
    token, _ = torch_sample_speculative_token_explicit(p, q, proposal, generator=gen)
    counts[token] += 1
assert torch.max(torch.abs(counts / counts.sum() - p)) < 0.02

# A real, tiny residual must remain a probability distribution.
near_q = np.array([1.0, 0.0]); near_p = np.array([1.0-1e-14, 1e-14])
for kernel, args in [
    (numpy_rejection_kernel_explicit, (near_p, near_q)),
    (torch_rejection_kernel_explicit, (torch.tensor(near_p), torch.tensor(near_q))),
]:
    _, residual, balance = kernel(*args)
    np.testing.assert_allclose(np.asarray(residual), [0.0, 1.0], rtol=0, atol=1e-15)
    np.testing.assert_allclose(np.asarray(balance), near_p, rtol=0, atol=1e-15)
# Uniform RNG can return zero: a zero-acceptance proposal must still be rejected.
from unittest.mock import patch
with patch.object(torch, "rand", return_value=torch.tensor(0.0, dtype=torch.float64)):
    token, accepted = torch_sample_speculative_token_explicit(
        torch.tensor([0.,1.], dtype=torch.float64),
        torch.tensor([1.,0.], dtype=torch.float64), 0, generator=gen)
    assert token == 1 and not accepted

print("B18 exact balance, tiny residual, zero draw, Monte Carlo: PASS")

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_grouped_fake_quant_explicit(
    x: np.ndarray,
    *,
    grid: np.ndarray,
    group_size: int,
) -> np.ndarray:
    """Minimal nearest-grid fake quantization along the last dimension."""
    pad = (-x.shape[-1]) % group_size
    padded = np.pad(x, [(0, 0)] * (x.ndim - 1) + [(0, pad)])
    groups = padded.reshape(*x.shape[:-1], -1, group_size)
    max_abs = np.max(np.abs(groups), axis=-1, keepdims=True)
    scales = np.where(max_abs > 0, max_abs / np.max(np.abs(grid)), 1.0)
    normalized = groups / scales
    codes = np.abs(normalized[..., None] - grid).argmin(axis=-1)
    return (grid[codes] * scales).reshape(padded.shape)[..., : x.shape[-1]]

def torch_grouped_fake_quant_explicit(
    x: torch.Tensor,
    *,
    grid: torch.Tensor,
    group_size: int,
) -> torch.Tensor:
    """Minimal tensor counterpart; leading dimensions never share a scale."""
    grid = grid.to(device=x.device, dtype=x.dtype)
    pad = (-x.shape[-1]) % group_size
    padded = torch.nn.functional.pad(x, (0, pad)) if pad else x
    groups = padded.reshape(*x.shape[:-1], -1, group_size)
    max_abs = groups.abs().amax(dim=-1, keepdim=True)
    scales = torch.where(max_abs > 0, max_abs / grid.abs().max(), torch.ones_like(max_abs))
    normalized = groups / scales
    codes = (normalized[..., None] - grid).abs().argmin(dim=-1)
    return (grid[codes] * scales).reshape(padded.shape)[..., : x.shape[-1]]

rng = np.random.default_rng(19)
x_np = rng.standard_normal((3,2,11))
grid_np = np.array([-6,-4,-3,-2,-1.5,-1,-0.5,0,0.5,1,1.5,2,3,4,6.0])
np_q = numpy_grouped_fake_quant_explicit(x_np, grid=grid_np, group_size=4)
t_q = torch_grouped_fake_quant_explicit(torch.tensor(x_np), grid=torch.tensor(grid_np), group_size=4)
np.testing.assert_allclose(np_q, t_q.numpy())
zeros = torch.zeros(2,3,7, dtype=torch.float64)
zero_q = torch_grouped_fake_quant_explicit(zeros, grid=torch.tensor(grid_np), group_size=4)
torch.testing.assert_close(zero_q, zeros)
print("B19 explicit grouped fake quantization: PASS")

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

"""M11 minimal probability and grouped-value teaching kernels."""
import numpy as np


def probability_vector(values, *, name="p", atol=1e-12):
    """Validate a distribution without hidden clipping or renormalization."""
    x = np.asarray(values, dtype=np.float64)
    if x.ndim != 1 or x.size < 2 or not np.isfinite(x).all():
        raise ValueError(f"{name}: expected a finite one-dimensional vector")
    if np.any(x < 0) or not np.isclose(x.sum(), 1.0, rtol=0.0, atol=atol):
        raise ValueError(f"{name}: probabilities must be >= 0 and sum to 1")
    return x


def rejection_kernel(p_values, q_values):
    """Acceptance probabilities, residual, and exact one-position balance."""
    p = probability_vector(p_values, name="p")
    q = probability_vector(q_values, name="q")
    if p.shape != q.shape:
        raise ValueError("p and q must use the same vocabulary")

    acceptance = np.zeros_like(p)
    support = q > 0  # q_t=0 means the proposer cannot emit token t
    acceptance[support] = np.minimum(1.0, p[support] / q[support])

    overlap = np.minimum(p, q)
    positive = np.maximum(p - q, 0.0)
    residual_mass = positive.sum()
    residual = (
        positive / residual_mass
        if residual_mass > 0.0
        else np.zeros_like(p)
    )
    balance = overlap + residual_mass * residual
    np.testing.assert_allclose(balance, p, rtol=0.0, atol=1e-12)
    return acceptance, residual, balance


def quantize_grouped_e2m1(values, group_size=16):
    """Teaching E2M1-like grid with one scale per group."""
    grid = np.array(
        [-6, -4, -3, -2, -1.5, -1, -0.5, 0,
          0.5, 1, 1.5, 2, 3, 4, 6],
        dtype=np.float64,
    )
    x = np.asarray(values, dtype=np.float64)
    if x.ndim != 1 or x.size == 0 or not np.isfinite(x).all():
        raise ValueError("expected a non-empty finite one-dimensional vector")
    if isinstance(group_size, bool) or not isinstance(group_size, (int, np.integer)) or group_size <= 0 or x.size % group_size:
        raise ValueError("length must be divisible by a positive group_size")

    groups = x.reshape(-1, group_size)
    max_abs = np.max(np.abs(groups), axis=1, keepdims=True)
    scale = np.where(max_abs > 0, max_abs / 6.0, 1.0)
    normalized = groups / scale
    codes = np.abs(normalized[..., None] - grid).argmin(axis=-1)
    return (grid[codes] * scale).reshape(x.shape)

# Rows 0..2 depend on the preceding token; row 3 is the initial distribution.
P=np.array([[.7,.2,.1],[.1,.7,.2],[.2,.1,.7],[.4,.35,.25]])
Q=np.array([[.55,.25,.2],[.2,.6,.2],[.3,.15,.55],[.3,.4,.3]])
def row(table, history):
    return table[history[-1] if history else 3]

class Target:
    def __init__(self):
        self.calls=0
    def verify(self, prefix, draft):
        # One stored batch result, including the bonus distribution.
        self.calls+=1
        histories=[prefix+draft[:i] for i in range(len(draft)+1)]
        return np.stack([row(P,h) for h in histories])

def generate(rng, length, gamma=3):
    target=Target(); history=[]; proposed=accepted=0
    while len(history)<length:
        draft=[]; q_rows=[]
        for _ in range(min(gamma,length-len(history))):
            q=row(Q,history+draft);q_rows.append(q.copy())
            draft.append(int(rng.choice(3,p=q)))
        proposed+=len(draft)
        p_rows=target.verify(history,draft)
        emitted=[]
        for i, token in enumerate(draft):
            a,residual,_=rejection_kernel(p_rows[i],q_rows[i])
            if rng.random()<a[token]:
                emitted.append(token);accepted+=1
            else:
                emitted.append(int(rng.choice(3,p=residual)))
                break
        else:
            emitted.append(int(rng.choice(3,p=p_rows[-1])))
        # No EOS in this fixed-length experiment. Remove any bonus beyond the limit.
        history.extend(emitted[:length-len(history)])
    return history,target.calls,proposed,accepted

rng=np.random.default_rng(SEED)
records=[generate(rng,4) for _ in range(12000)]
joint=np.zeros((3,3))
for tokens,_,_,_ in records:
    joint[tokens[0],tokens[1]]+=1
joint/=len(records)
exact=P[3,:,None]*P[:3]
tv=.5*np.abs(joint-exact).sum()
assert tv<.025
calls=sum(r[1] for r in records)
assert len(records)<=calls<=4*len(records)
print("joint-law TV:",tv,"target batch calls:",calls,"AR calls:",4*len(records))
print("emitted tokens / target call:",4*len(records)/calls)
print("This is a call-count experiment, not a GPU speed benchmark.")
# With no proposals, each stored target batch supplies one ordinary token.
baseline = generate(np.random.default_rng(SEED), 4, gamma=0)
assert baseline[1:] == (4, 0, 0)
# Each target row uses the history BEFORE its candidate. A shifted row fails.
wrong=P[3,:,None]*P[:3][[1,2,0]]
assert .5*np.abs(wrong-exact).sum()>.2

rng=np.random.default_rng(SEED+1)
w=rng.normal(size=(16,64));w[:,3]*=6
x=rng.normal(size=(256,64));x[:,0]*=8
grid=np.array([-6,-4,-3,-2,-1.5,-1,-.5,0,.5,1,1.5,2,3,4,6])
def quantize_rows(w, fp4):
    groups=w.reshape(w.shape[0],-1,16)
    peak=np.abs(groups).max(-1,keepdims=True)
    scale=np.where(peak>0,peak/(6 if fp4 else 7),1.)
    normalized=groups/scale
    code=grid[np.abs(normalized[...,None]-grid).argmin(-1)] if fp4 else np.rint(normalized).clip(-7,7)
    return (code*scale).reshape(w.shape),scale.size

y=x@w.T
for fp4 in (False,True):
    restored,n_scales=quantize_rows(w,fp4)
    weight_mse=np.mean((w-restored)**2)
    output_error=np.linalg.norm(x@restored.T-y)/np.linalg.norm(y)
    ideal_bytes=(w.size+1)//2+n_scales*4
    print("E2M1" if fp4 else "INT4","weight MSE:",weight_mse,
          "relative output error:",output_error,"ideal payload + float32 scales B:",ideal_bytes)
    assert np.isfinite(output_error)
assert np.array_equal(quantize_rows(np.zeros_like(w),True)[0],np.zeros_like(w))
# Equal weight MSE can affect different input channels very differently.
e0=np.zeros_like(w);e1=np.zeros_like(w);e0[:,0]=.1;e1[:,1]=.1
assert np.isclose(np.mean(e0**2),np.mean(e1**2))
out0=np.mean((x@e0.T)**2);out1=np.mean((x@e1.T)**2)
assert out0>20*out1
print("equal weight MSE; output-MSE ratio:",out0/out1)
# Decoded arrays are float64; the byte count assumes packing that this code does not implement.