blackwell_primer.py

#!/usr/bin/env python3
"""
Blackwell Primer — Programmatic Memory and Verification for SM120/SM121

Hardware-first, bottom-up verification and analysis tool for the Blackwell
LLM deployment lab. Covers all GPU architectures (SM120/SM121/SM89) and all
deployment targets (RTX Pro 6000, DGX Spark, RTX 4090).

This script IS the algorithm. It reads hardware, parses source code, checks
configuration, and reports exactly what works, what's blocked, and what to
do next. Every check traces to a file, line number, or hardware query.

Primary function: verify TRT-LLM fork readiness for NVFP4 KV cache on
SM120/SM121. Secondary: model-specific analysis (Qwen3-Next-80B kernel
roadmap, weight inventory, performance ceilings).

State is persisted in data/ directory. Re-run after any code change to
detect regressions. This is a climbing anchor — set it, verify it holds,
build on it.

Commands:
    primer verify <fork_path>       Programmatic code audit of TRT-LLM fork
    primer unknowns                 Show open questions and design decisions
    primer analyze                  Model analysis (architecture + roadmap)
    primer next                     Single instruction: what to do RIGHT NOW
    primer status                   Phase overview with prerequisite checks
    primer record <phase> <key> <value>   Save a measurement
    primer history                  Show all recorded measurements
    primer complete <phase>         Manually mark phase done
    primer note <phase> <text>      Add a note to a phase

Usage:
    python3 blackwell_primer.py verify SM12x.../TensorRT-LLM --arch sm121
    python3 blackwell_primer.py verify SM12x.../TensorRT-LLM --arch sm120
    python3 blackwell_primer.py unknowns
    python3 blackwell_primer.py analyze [config.json] [--gpu sm120]
    python3 blackwell_primer.py status
"""

import json
import sys
import re
import subprocess
import time
import concurrent.futures
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional, Any
from pathlib import Path
from collections import defaultdict


# =============================================================================
# Hardware Specifications (measured, not marketed)
# =============================================================================

@dataclass
class HardwareSpec:
    name: str
    sm_arch: str
    sm_count: int
    memory_bandwidth_gb_s: float       # DRAM bandwidth (GB/s)
    memory_total_gb: float             # Total VRAM (GB)
    registers_per_sm: int              # 32-bit registers
    smem_per_sm_bytes: int             # Shared memory per SM (bytes)
    max_warps_per_sm: int              # Max concurrent warps
    warp_size: int = 32
    kernel_launch_overhead_us: float = 3.5  # Measured CPU→GPU dispatch (μs)
    cuda_graph_launch_overhead_us: float = 0.5  # GPU-side replay (μs)

    @property
    def max_threads_per_sm(self) -> int:
        return self.max_warps_per_sm * self.warp_size

SM120_RTX_PRO_6000 = HardwareSpec(
    name="RTX Pro 6000 (SM120)",
    sm_arch="sm_120a",
    sm_count=188,
    memory_bandwidth_gb_s=1792.0,
    memory_total_gb=96.0,
    registers_per_sm=65536,
    smem_per_sm_bytes=101376,           # 99KB per block (opt-in); 100KB (102400) per SM total
    max_warps_per_sm=48,
    kernel_launch_overhead_us=3.5,
    cuda_graph_launch_overhead_us=0.5,
    # Verified 2026-02-11 via CUDA driver API (cuDeviceGetAttribute).
    # NOT 128KB — LLM research pipelines hallucinate this value.
)

SM121_DGX_SPARK = HardwareSpec(
    name="DGX Spark (SM121)",
    sm_arch="sm_121a",
    sm_count=84,
    memory_bandwidth_gb_s=273.0,       # LPDDR5x
    memory_total_gb=128.0,
    registers_per_sm=65536,
    smem_per_sm_bytes=101376,
    max_warps_per_sm=48,
    kernel_launch_overhead_us=3.5,
    cuda_graph_launch_overhead_us=0.5,
)

GPU_SPECS = {
    "sm120": SM120_RTX_PRO_6000,
    "sm121": SM121_DGX_SPARK,
}


# =============================================================================
# Quantization Format Definitions
# =============================================================================

@dataclass
class QuantFormat:
    name: str
    bits_per_element: float             # Effective bits per weight element
    scale_bits: int                     # Bits per scale factor
    scale_group_size: int               # Elements per scale group
    bytes_per_element: float = 0.0      # Computed

    def __post_init__(self):
        # Total bytes = data bytes + scale bytes
        data_bytes_per_elem = self.bits_per_element / 8.0
        scale_bytes_per_elem = (self.scale_bits / 8.0) / self.scale_group_size
        self.bytes_per_element = data_bytes_per_elem + scale_bytes_per_elem

NVFP4 = QuantFormat(
    name="NVFP4",
    bits_per_element=4,
    scale_bits=8,           # UE4M3
    scale_group_size=16,    # 1 scale per 16 elements
)

FP8_E4M3 = QuantFormat(
    name="FP8",
    bits_per_element=8,
    scale_bits=0,
    scale_group_size=1,
)

BF16 = QuantFormat(
    name="BF16",
    bits_per_element=16,
    scale_bits=0,
    scale_group_size=1,
)

FP16 = QuantFormat(
    name="FP16",
    bits_per_element=16,
    scale_bits=0,
    scale_group_size=1,
)


# =============================================================================
# Weight Matrix Descriptor
# =============================================================================

@dataclass
class WeightMatrix:
    """A single weight matrix in the model."""
    name: str                           # e.g. "layer.0.linear_attn.in_proj_qkvz"
    layer_index: int
    layer_type: str                     # "linear_attention" or "full_attention"
    operation: str                      # "attention_proj", "expert_ffn", "router", etc.
    shape_N: int                        # Output dimension (rows)
    shape_K: int                        # Input dimension (cols)
    quant: QuantFormat                  # NVFP4, BF16, etc.
    count_per_forward: int              # How many times executed per forward pass
    is_routed_expert: bool = False      # True for per-expert weights (grouped)
    expert_index: Optional[int] = None

    @property
    def weight_bytes(self) -> float:
        """Total bytes to read this weight matrix once."""
        return self.shape_N * self.shape_K * self.quant.bytes_per_element

    @property
    def total_bytes_per_forward(self) -> float:
        """Total bytes read from this matrix across the full forward pass."""
        return self.weight_bytes * self.count_per_forward

    @property
    def flops_per_call_m1(self) -> int:
        """FLOPs for M=1 (GEMV): 2*N*K."""
        return 2 * self.shape_N * self.shape_K

    @property
    def arithmetic_intensity_m1(self) -> float:
        """FLOPs per byte at M=1."""
        if self.weight_bytes == 0:
            return 0.0
        return self.flops_per_call_m1 / self.weight_bytes


# =============================================================================
# GEMV Shape Descriptor (aggregated)
# =============================================================================

@dataclass
class GemvShape:
    """Unique (N, K, quant) shape with aggregated count."""
    shape_N: int
    shape_K: int
    quant: QuantFormat
    total_calls_per_forward: int
    operation_names: List[str] = field(default_factory=list)
    is_routed_expert: bool = False

    @property
    def weight_bytes(self) -> float:
        return self.shape_N * self.shape_K * self.quant.bytes_per_element

    @property
    def input_bytes_m1(self) -> float:
        """Input vector x is K elements in FP16."""
        return self.shape_K * 2

    @property
    def total_read_bytes_per_call(self) -> float:
        return self.weight_bytes + self.input_bytes_m1

    @property
    def total_read_bytes_per_forward(self) -> float:
        return self.total_read_bytes_per_call * self.total_calls_per_forward

    @property
    def flops_per_call_m1(self) -> int:
        return 2 * self.shape_N * self.shape_K


# =============================================================================
# Model Config Parser
# =============================================================================

def parse_model_config(config_path: str) -> dict:
    """Load and return the raw HuggingFace config.json."""
    with open(config_path, 'r') as f:
        return json.load(f)


def extract_architecture_parameters(config: dict) -> dict:
    """
    Extract every architecture parameter from config.json into a flat dict.
    Uses exact field names from the config. No guessing.
    """
    params = {}

    # Core dimensions
    params["hidden_size"] = config["hidden_size"]
    params["intermediate_size"] = config.get("intermediate_size", 0)
    params["vocab_size"] = config["vocab_size"]
    params["num_hidden_layers"] = config["num_hidden_layers"]

    # Attention (full)
    params["num_attention_heads"] = config["num_attention_heads"]
    params["num_key_value_heads"] = config["num_key_value_heads"]
    params["head_dim"] = config["head_dim"]

    # Attention (linear / DeltaNet / Mamba)
    params["linear_key_head_dim"] = config.get("linear_key_head_dim", 0)
    params["linear_num_key_heads"] = config.get("linear_num_key_heads", 0)
    params["linear_num_value_heads"] = config.get("linear_num_value_heads", 0)
    params["linear_value_head_dim"] = config.get("linear_value_head_dim", 0)
    params["linear_conv_kernel_dim"] = config.get("linear_conv_kernel_dim", 0)

    # MoE
    params["num_experts"] = config.get("num_experts", 0)
    params["num_experts_per_tok"] = config.get("num_experts_per_tok", 0)
    params["moe_intermediate_size"] = config.get("moe_intermediate_size", 0)
    params["shared_expert_intermediate_size"] = config.get("shared_expert_intermediate_size", 0)
    params["decoder_sparse_step"] = config.get("decoder_sparse_step", 1)

    # Layer types
    params["layer_types"] = config.get("layer_types", [])
    params["full_attention_interval"] = config.get("full_attention_interval", 0)

    # Context
    params["max_position_embeddings"] = config.get("max_position_embeddings", 0)

    # Quantization
    qconfig = config.get("quantization_config", {})
    params["quant_method"] = qconfig.get("quant_method", "none")
    params["quant_algo"] = qconfig.get("quant_algo", "none")
    params["quant_ignore_list"] = qconfig.get("ignore", [])

    # KV cache quantization
    kv_scheme = qconfig.get("kv_cache_scheme", {})
    params["kv_cache_bits"] = kv_scheme.get("num_bits", 16)
    params["kv_cache_type"] = kv_scheme.get("type", "float")

    return params


def classify_layer_types(params: dict) -> Tuple[List[int], List[int]]:
    """
    Return (linear_attention_indices, full_attention_indices) from layer_types.
    """
    linear_indices = []
    full_indices = []
    for i, ltype in enumerate(params["layer_types"]):
        if ltype == "linear_attention":
            linear_indices.append(i)
        elif ltype == "full_attention":
            full_indices.append(i)
        else:
            print(f"  WARNING: Unknown layer type '{ltype}' at index {i}")
    return linear_indices, full_indices


def build_quantization_ignore_set(params: dict) -> set:
    """
    Build a set of layer name patterns that are NOT quantized (stay in BF16).
    The ignore list in config.json contains full layer paths.
    """
    return set(params["quant_ignore_list"])


def determine_weight_quant(layer_idx: int, proj_suffix: str,
                           ignore_set: set, default_quant: QuantFormat) -> QuantFormat:
    """
    Determine if a specific weight matrix is NVFP4 or BF16 by checking
    the quantization ignore list.
    """
    # Build possible ignore patterns
    patterns = [
        f"model.layers.{layer_idx}.{proj_suffix}",
    ]
    for pattern in patterns:
        if pattern in ignore_set:
            return BF16
    return default_quant


# =============================================================================
# Weight Matrix Enumeration — The Complete Forward Pass
# =============================================================================

def enumerate_all_weight_matrices(params: dict) -> List[WeightMatrix]:
    """
    Enumerate EVERY weight matrix read during a single forward pass.
    This is exhaustive. If it reads memory, it's listed here.
    """
    weights = []
    ignore_set = build_quantization_ignore_set(params)
    linear_layers, full_layers = classify_layer_types(params)

    H = params["hidden_size"]                       # 2048
    num_layers = params["num_hidden_layers"]         # 48
    num_experts = params["num_experts"]              # 512
    experts_per_tok = params["num_experts_per_tok"]  # 10
    moe_intermediate = params["moe_intermediate_size"]            # 512
    shared_intermediate = params["shared_expert_intermediate_size"]  # 512
    vocab = params["vocab_size"]                     # 151936

    # Full attention dimensions
    full_q_dim = params["num_attention_heads"] * params["head_dim"]      # 16*256=4096
    full_k_dim = params["num_key_value_heads"] * params["head_dim"]      # 2*256=512
    full_v_dim = params["num_key_value_heads"] * params["head_dim"]      # 2*256=512
    full_o_dim = full_q_dim                                              # 4096

    # Linear attention dimensions
    lin_q_dim = params["linear_num_key_heads"] * params["linear_key_head_dim"]    # 16*128=2048
    lin_k_dim = params["linear_num_key_heads"] * params["linear_key_head_dim"]    # 16*128=2048
    lin_v_dim = params["linear_num_value_heads"] * params["linear_value_head_dim"]  # 32*128=4096
    # Z dimension: typically same as Q for DeltaNet gating
    lin_z_dim = lin_q_dim                                                  # 2048
    lin_qkvz_dim = lin_q_dim + lin_k_dim + lin_v_dim + lin_z_dim          # 10240
    # BA projection (B and A state matrices for DeltaNet)
    # Typically 2 * num_heads * head_dim or similar — check model code
    # For now estimate same as Q+K
    lin_ba_dim = lin_q_dim + lin_k_dim                                     # 4096

    # Linear attention output projection
    lin_o_dim = H                                                          # 2048

    # -------------------------------------------------------------------------
    # Process each layer
    # -------------------------------------------------------------------------
    for layer_idx in range(num_layers):
        layer_type = params["layer_types"][layer_idx]

        # --- Attention Projections ---
        if layer_type == "linear_attention":
            # Fused QKVZ projection: [H] -> [Q+K+V+Z]
            q = determine_weight_quant(layer_idx, "linear_attn.in_proj_qkvz", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.linear_attn.in_proj_qkvz",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation="linear_attention_qkvz_projection",
                shape_N=lin_qkvz_dim,
                shape_K=H,
                quant=q,
                count_per_forward=1,
            ))

            # BA projection: [H] -> [B+A]
            q_ba = determine_weight_quant(layer_idx, "linear_attn.in_proj_ba", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.linear_attn.in_proj_ba",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation="linear_attention_ba_projection",
                shape_N=lin_ba_dim,
                shape_K=H,
                quant=q_ba,
                count_per_forward=1,
            ))

            # Conv1D: small, [kernel_dim * channels]
            # This is a 1D causal conv, not a GEMV. Include for byte accounting.
            conv_kernel = params["linear_conv_kernel_dim"]
            q_conv = determine_weight_quant(layer_idx, "linear_attn.conv1d", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.linear_attn.conv1d",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation="linear_attention_conv1d",
                shape_N=H,
                shape_K=conv_kernel,  # Conv kernel is tiny
                quant=q_conv,
                count_per_forward=1,
            ))

            # Output projection: [V_out] -> [H]
            # This is NOT in the ignore list for linear attention
            q_o = determine_weight_quant(layer_idx, "linear_attn.o_proj", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.linear_attn.o_proj",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation="linear_attention_output_projection",
                shape_N=H,
                shape_K=lin_v_dim,  # Output proj takes V output
                quant=q_o,
                count_per_forward=1,
            ))

        elif layer_type == "full_attention":
            # Q projection: [H] -> [Q_dim]
            q_q = determine_weight_quant(layer_idx, "self_attn.q_proj", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.self_attn.q_proj",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation="full_attention_q_projection",
                shape_N=full_q_dim,
                shape_K=H,
                quant=q_q,
                count_per_forward=1,
            ))

            # K projection: [H] -> [K_dim]
            q_k = determine_weight_quant(layer_idx, "self_attn.k_proj", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.self_attn.k_proj",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation="full_attention_k_projection",
                shape_N=full_k_dim,
                shape_K=H,
                quant=q_k,
                count_per_forward=1,
            ))

            # V projection: [H] -> [V_dim]
            q_v = determine_weight_quant(layer_idx, "self_attn.v_proj", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.self_attn.v_proj",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation="full_attention_v_projection",
                shape_N=full_v_dim,
                shape_K=H,
                quant=q_v,
                count_per_forward=1,
            ))

            # O projection: [Q_dim] -> [H]
            q_o = determine_weight_quant(layer_idx, "self_attn.o_proj", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.self_attn.o_proj",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation="full_attention_output_projection",
                shape_N=H,
                shape_K=full_q_dim,
                quant=q_o,
                count_per_forward=1,
            ))

        # --- MoE: Router ---
        q_gate = determine_weight_quant(layer_idx, "mlp.gate", ignore_set, NVFP4)
        weights.append(WeightMatrix(
            name=f"layer.{layer_idx}.mlp.gate",
            layer_index=layer_idx,
            layer_type=layer_type,
            operation="moe_router",
            shape_N=num_experts,
            shape_K=H,
            quant=q_gate,
            count_per_forward=1,
        ))

        # --- MoE: Routed Experts (only top-K active per token) ---
        for expert_op, (N, K) in [
            ("expert_up_proj",   (moe_intermediate, H)),
            ("expert_gate_proj", (moe_intermediate, H)),
            ("expert_down_proj", (H, moe_intermediate)),
        ]:
            q_exp = determine_weight_quant(
                layer_idx, f"mlp.experts.0.{expert_op.replace('expert_', '')}", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.mlp.experts.*.{expert_op}",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation=f"moe_routed_{expert_op}",
                shape_N=N,
                shape_K=K,
                quant=q_exp,
                count_per_forward=experts_per_tok,  # 10 experts active
                is_routed_expert=True,
            ))

        # --- MoE: Shared Expert ---
        q_shared_gate = determine_weight_quant(
            layer_idx, "mlp.shared_expert_gate", ignore_set, NVFP4)
        # The shared_expert_gate is a scalar gate, not a full FFN
        # But we still account for its weight read
        weights.append(WeightMatrix(
            name=f"layer.{layer_idx}.mlp.shared_expert_gate",
            layer_index=layer_idx,
            layer_type=layer_type,
            operation="moe_shared_expert_gate_scalar",
            shape_N=1,
            shape_K=H,
            quant=q_shared_gate,
            count_per_forward=1,
        ))

        for shared_op, (N, K) in [
            ("shared_expert_up_proj",   (shared_intermediate, H)),
            ("shared_expert_gate_proj", (shared_intermediate, H)),
            ("shared_expert_down_proj", (H, shared_intermediate)),
        ]:
            q_shared = determine_weight_quant(
                layer_idx, f"mlp.shared_expert.{shared_op.replace('shared_expert_', '')}", ignore_set, NVFP4)
            weights.append(WeightMatrix(
                name=f"layer.{layer_idx}.mlp.shared_expert.{shared_op}",
                layer_index=layer_idx,
                layer_type=layer_type,
                operation=f"moe_{shared_op}",
                shape_N=N,
                shape_K=K,
                quant=q_shared,
                count_per_forward=1,
            ))

    # --- LM Head ---
    # lm_head is in the ignore list → BF16
    weights.append(WeightMatrix(
        name="lm_head",
        layer_index=-1,
        layer_type="output",
        operation="lm_head",
        shape_N=vocab,
        shape_K=H,
        quant=BF16,  # Always ignored in quantization
        count_per_forward=1,
    ))

    # --- Embedding ---
    # Embedding table is read once (lookup, not GEMV)
    # Include for total memory accounting but not GEMV time
    weights.append(WeightMatrix(
        name="model.embed_tokens",
        layer_index=-1,
        layer_type="input",
        operation="embedding_lookup",
        shape_N=vocab,
        shape_K=H,
        quant=BF16,
        count_per_forward=1,  # lookup, not matmul
    ))

    return weights


# =============================================================================
# Shape Aggregation
# =============================================================================

def aggregate_gemv_shapes(weights: List[WeightMatrix]) -> List[GemvShape]:
    """
    Group weight matrices by unique (N, K, quant, is_routed) tuples.
    Returns aggregated shapes sorted by total bytes descending.
    """
    shape_map: Dict[Tuple, GemvShape] = {}

    for w in weights:
        # Skip embedding (it's a lookup, not a GEMV)
        if w.operation == "embedding_lookup":
            continue

        key = (w.shape_N, w.shape_K, w.quant.name, w.is_routed_expert)
        if key not in shape_map:
            shape_map[key] = GemvShape(
                shape_N=w.shape_N,
                shape_K=w.shape_K,
                quant=w.quant,
                total_calls_per_forward=0,
                operation_names=[],
                is_routed_expert=w.is_routed_expert,
            )
        shape_map[key].total_calls_per_forward += w.count_per_forward
        if w.operation not in shape_map[key].operation_names:
            shape_map[key].operation_names.append(w.operation)

    shapes = list(shape_map.values())
    shapes.sort(key=lambda s: s.total_read_bytes_per_forward, reverse=True)
    return shapes


# =============================================================================
# Performance Model
# =============================================================================

@dataclass
class ShapePerformanceModel:
    """Performance prediction for a single GEMV shape."""
    shape: GemvShape
    hw: HardwareSpec

    # Computed fields
    blocks_per_call: int = 0            # Threadblocks launched
    sm_utilization: float = 0.0         # blocks / sm_count
    waves: int = 0                      # ceil(blocks / sm_count)
    bandwidth_limited_time_us: float = 0.0  # bytes / bandwidth
    launch_overhead_time_us: float = 0.0    # total launch overhead
    total_time_per_forward_us: float = 0.0  # estimated wall time

    # Classification
    bottleneck: str = ""                # "launch_overhead", "bandwidth", "sm_starvation"

    def compute(self, rows_per_block: int = 8):
        """Compute all performance predictions."""
        s = self.shape
        hw = self.hw

        # Threadblock count: ceil(N / rows_per_block)
        self.blocks_per_call = (s.shape_N + rows_per_block - 1) // rows_per_block
        self.sm_utilization = self.blocks_per_call / hw.sm_count
        self.waves = max(1, (self.blocks_per_call + hw.sm_count - 1) // hw.sm_count)

        # Bandwidth-limited time (pure data movement, no overhead)
        bytes_per_call = s.total_read_bytes_per_call
        self.bandwidth_limited_time_us = bytes_per_call / (hw.memory_bandwidth_gb_s * 1e3)

        # Launch overhead (CPU dispatch per kernel)
        if s.is_routed_expert:
            # Routed experts: if batched into 1 launch, 1 overhead
            # If individual launches, experts_per_tok overheads
            # We model BOTH and show the gap
            self.launch_overhead_time_us = (
                hw.kernel_launch_overhead_us * s.total_calls_per_forward
            )
        else:
            self.launch_overhead_time_us = (
                hw.kernel_launch_overhead_us * s.total_calls_per_forward
            )

        # Total estimate (bandwidth + launch overhead, whichever dominates)
        bw_time_total = self.bandwidth_limited_time_us * s.total_calls_per_forward
        self.total_time_per_forward_us = bw_time_total + self.launch_overhead_time_us

        # Classify bottleneck
        launch_fraction = self.launch_overhead_time_us / max(self.total_time_per_forward_us, 1e-9)
        if launch_fraction > 0.5:
            self.bottleneck = "LAUNCH_OVERHEAD"
        elif self.sm_utilization < 0.5:
            self.bottleneck = "SM_STARVATION"
        else:
            self.bottleneck = "BANDWIDTH"


def compute_full_forward_pass_model(
    shapes: List[GemvShape], hw: HardwareSpec
) -> List[ShapePerformanceModel]:
    """Compute performance model for every shape in the forward pass."""
    models = []
    for s in shapes:
        m = ShapePerformanceModel(shape=s, hw=hw)
        m.compute()
        models.append(m)
    return models


# =============================================================================
# Optimization Phase Generator
# =============================================================================

@dataclass
class OptimizationPhase:
    """A single step in the optimization roadmap."""
    phase_number: int
    name: str
    prerequisite: str
    what_to_build: str
    expected_time_saved_us: float
    expected_tps_after: float
    measurement_command: str
    done_when: str


def generate_optimization_roadmap(
    models: List[ShapePerformanceModel],
    hw: HardwareSpec,
    params: dict,
) -> List[OptimizationPhase]:
    """
    Generate the deterministic optimization roadmap.
    Order is fixed by the performance model: fix the biggest bottleneck first.
    """
    total_time_us = sum(m.total_time_per_forward_us for m in models)
    total_launch_overhead_us = sum(m.launch_overhead_time_us for m in models)
    total_bw_time_us = sum(
        m.bandwidth_limited_time_us * m.shape.total_calls_per_forward
        for m in models
    )
    total_calls = sum(m.shape.total_calls_per_forward for m in models)

    # Count routed expert calls
    routed_calls = sum(
        m.shape.total_calls_per_forward
        for m in models if m.shape.is_routed_expert
    )
    routed_overhead_us = sum(
        m.launch_overhead_time_us
        for m in models if m.shape.is_routed_expert
    )

    # Count non-routed calls
    nonrouted_calls = total_calls - routed_calls
    nonrouted_overhead_us = total_launch_overhead_us - routed_overhead_us

    phases = []
    running_time_us = total_time_us

    # Phase 0: Baseline GEMV kernel (already done)
    phases.append(OptimizationPhase(
        phase_number=0,
        name="BASELINE: Individual GEMV Kernels",
        prerequisite="Working NVFP4 + BF16 GEMV kernels with LUT dequant",
        what_to_build="nvfp4_gemv.cu (NVFP4 path) + bf16_gemv.cu (BF16 path)",
        expected_time_saved_us=0,
        expected_tps_after=1e6 / running_time_us if running_time_us > 0 else 0,
        measurement_command="bench_nvfp4_gemv (all shapes, correctness + timing)",
        done_when="All shapes pass correctness, bandwidth efficiency measured per shape",
    ))

    # Phase 1: Batched GEMV for routed experts
    # Fuse N_experts same-shape kernels into one launch
    if routed_calls > 0:
        # With batching: 1 launch per shape per layer instead of experts_per_tok launches
        batched_routed_launches = routed_calls / params["num_experts_per_tok"]
        saved_launches = routed_calls - batched_routed_launches
        time_saved = saved_launches * hw.kernel_launch_overhead_us
        running_time_us -= time_saved

        phases.append(OptimizationPhase(
            phase_number=1,
            name="BATCHED GEMV: Fuse Routed Expert Launches",
            prerequisite="Phase 0 baseline measured",
            what_to_build=(
                "nvfp4_batched_gemv.cu — single kernel launch processes "
                f"{params['num_experts_per_tok']} expert weight matrices. "
                "Grid: (ceil(N/rows_per_block) * num_active_experts) blocks. "
                "Each block reads expert_id from a dispatch table to select "
                "the correct weight pointer. Same inner loop as Phase 0. "
                "Library alternatives exist (CUTLASS MoEProblemShape + GroupedGemmMoE, "
                "cuBLAS cublasGemmGroupedBatchedEx) but are tile-based (128x128) — "
                f"our custom GEMV handles M=1 with {params['moe_intermediate_size']}-wide N directly."
            ),
            expected_time_saved_us=time_saved,
            expected_tps_after=1e6 / running_time_us if running_time_us > 0 else 0,
            measurement_command=(
                "bench_batched_gemv: compare individual vs batched for "
                f"(N={params['moe_intermediate_size']}, K={params['hidden_size']}) × "
                f"{params['num_experts_per_tok']} experts"
            ),
            done_when=(
                f"Batched launch time < {params['num_experts_per_tok']}× individual launch time. "
                "Correctness matches individual path."
            ),
        ))

    # Phase 2: CUDA Graphs for the full forward pass
    remaining_launches = nonrouted_calls + (routed_calls / params["num_experts_per_tok"] if routed_calls > 0 else 0)
    graph_time_saved = remaining_launches * (hw.kernel_launch_overhead_us - hw.cuda_graph_launch_overhead_us)
    running_time_us -= graph_time_saved

    phases.append(OptimizationPhase(
        phase_number=2,
        name="CUDA GRAPHS: Eliminate CPU Launch Overhead",
        prerequisite="Phase 1 batched GEMV working",
        what_to_build=(
            "forward_pass_graph.cu — Record all GEMV launches "
            f"(~{int(remaining_launches)} kernels after batching) into a single "
            "CUDA graph. Replay is GPU-side, ~0.5μs per launch instead of ~3.5μs. "
            "Requires: all kernel shapes are static (they are at M=1 decode). "
            "The graph captures kernel pointers, grid dims, and shared memory sizes. "
            "Weight pointers are fixed (model weights don't move). "
            "Only the input vector x changes between tokens — update via "
            "cudaGraphExecKernelNodeSetParams or use a persistent x buffer."
        ),
        expected_time_saved_us=graph_time_saved,
        expected_tps_after=1e6 / running_time_us if running_time_us > 0 else 0,
        measurement_command=(
            "bench_cuda_graph: record + replay full forward pass, "
            "compare to unbatched individual launches"
        ),
        done_when=(
            "Graph replay time < sum of individual launches. "
            f"Target: eliminate ~{graph_time_saved:.0f}μs of launch overhead."
        ),
    ))

    # Phase 3: Kernel occupancy tuning
    # Find shapes where SM starvation is the bottleneck
    starved_shapes = [m for m in models if m.bottleneck == "SM_STARVATION"]
    starved_time = sum(m.total_time_per_forward_us for m in starved_shapes)
    # Estimate: with better occupancy, recover ~30% of starvation penalty
    occ_time_saved = starved_time * 0.3
    running_time_us -= occ_time_saved

    phases.append(OptimizationPhase(
        phase_number=3,
        name="OCCUPANCY TUNING: Reduce Register Pressure",
        prerequisite="Phase 2 CUDA Graphs working",
        what_to_build=(
            "Tune launch_bounds and inner loop to reduce registers from 60 to ~48, "
            "enabling 6 blocks/SM (100% occupancy at 256 threads/block). "
            "Currently at 4 blocks/SM (67% occupancy). "
            "Methods: reduce accumulator variables, use __half2 instead of float "
            "for intermediate x values, limit unroll depth. "
            "Profile with: ncu --metrics launch__registers_per_thread"
        ),
        expected_time_saved_us=occ_time_saved,
        expected_tps_after=1e6 / running_time_us if running_time_us > 0 else 0,
        measurement_command=(
            "ncu --metrics sm__warps_active.avg.pct_of_peak_sustained_active "
            "./bench_nvfp4_gemv (before and after)"
        ),
        done_when=(
            "Warp occupancy > 90%. Registers per thread ≤ 48. "
            "No performance regression on any shape."
        ),
    ))

    # Phase 4: LM Head bandwidth optimization
    lm_head_models = [m for m in models if "lm_head" in m.shape.operation_names]
    if lm_head_models:
        lm_time = sum(m.total_time_per_forward_us for m in lm_head_models)
        lm_bw_time = sum(
            m.bandwidth_limited_time_us * m.shape.total_calls_per_forward
            for m in lm_head_models
        )
        # LM head is BF16, biggest single matrix. Target: 70% bandwidth eff
        # Currently ~53% for NVFP4, BF16 LM head will be different
        lm_improvement = lm_time * 0.25  # Estimate 25% improvement
        running_time_us -= lm_improvement

        phases.append(OptimizationPhase(
            phase_number=4,
            name="LM HEAD: BF16 GEMV with Maximum Bandwidth",
            prerequisite="Phase 3 occupancy tuned",
            what_to_build=(
                f"bf16_gemv.cu — Specialized kernel for LM head "
                f"(N={params['vocab_size']}, K={params['hidden_size']}). "
                "BF16 weights (not NVFP4) — no dequant needed. "
                "Use 128-bit vectorized loads (4× BF16 per load). "
                "This is the single largest GEMV in the forward pass "
                f"({params['vocab_size'] * params['hidden_size'] * 2 / 1e6:.0f} MB). "
                "Target: >70% bandwidth efficiency."
            ),
            expected_time_saved_us=lm_improvement,
            expected_tps_after=1e6 / running_time_us if running_time_us > 0 else 0,
            measurement_command=(
                f"bench_bf16_gemv: N={params['vocab_size']}, K={params['hidden_size']}, "
                "measure BW efficiency"
            ),
            done_when=(
                "BF16 GEMV achieves >70% of peak memory bandwidth on LM head shape."
            ),
        ))

    # Phase 5: Non-GEMV operations
    # Research: top-K routing and fused elementwise queries ACCEPTED.
    # RMSNorm and SiLU queries failed (docs insufficient / validation failed).
    phases.append(OptimizationPhase(
        phase_number=5,
        name="NON-GEMV OPS: RMSNorm + SiLU + Routing + Softmax",
        prerequisite="Phase 4 complete",
        what_to_build=(
            "Fused element-wise kernels for the non-GEMV operations: "
            "RMSNorm (per layer), SiLU activation (per expert FFN), "
            "top-K routing with softmax (per MoE layer). "
            "These are tiny in bytes but add launch overhead. "
            "Fuse adjacent operations: e.g., RMSNorm + input broadcast, "
            "expert_up * SiLU(expert_gate) in a single kernel. "
            "Note: top-K routing can reuse CUTLASS MoE epilogue (top_K <= 8 compile-time) "
            f"but our model uses top-{params['num_experts_per_tok']} — needs 2-pass or custom."
        ),
        expected_time_saved_us=0,  # Hard to estimate without profiling
        expected_tps_after=1e6 / running_time_us if running_time_us > 0 else 0,
        measurement_command="Profile full forward pass with nsys to identify remaining gaps",
        done_when="Non-GEMV operations add < 10% overhead to total forward time",
    ))

    # Phase 6: Attention compute (full attention layers only)
    #
    # RESOLVED (Q1, Q2, Q3, Q5 via nvidia-docs-search 2026-02-11):
    #   - mxf8f6f4.block_scale: A=mx_float8_t(FP8) × B=mx_float4_t(FP4), UE8M0 scales, group=32
    #   - mxf4nvf4.block_scale: A=nv_float4_t × B=nv_float4_t, UE4M3 scales, group=16, TN only
    #   - Accumulator: FP32 for all block_scale variants
    #   - SM120 shapes: m16n8k64, m16n8k128
    #
    # THREE VIABLE PATHS (see D1_attention_mma_path in OPEN_QUESTIONS):
    #   Path A: mxf8f6f4 — Q quantized FP16→FP8, K stays FP4. KV cache needs MX format.
    #   Path B: mxf4nvf4 — Q quantized FP16→FP4, K stays NVFP4. Native format. 4x throughput.
    #   Path C: Software dequant — Q stays FP16, K dequanted via LUT. Max precision.
    #
    # UNRESOLVED: Q8 (Q precision loss), Q9 (missing .cuh headers), Q10 (head_dim=256)
    #
    # STRATEGY: Implement Path C first (bandwidth-bound decode, simplest).
    # Then Path B (highest throughput). Then Path A (best Q precision).
    # Benchmark all three. Decode at M=1 is bandwidth-bound, so Path C may win.
    #
    n_full_layers = len([t for t in params["layer_types"] if t == "full_attention"])
    nvfp4_kv_bytes_per_elem = NVFP4.bytes_per_element  # 0.5625 B (4-bit + scale)
    kv_bytes_8k = (n_full_layers * 2 * params['num_key_value_heads']
                   * params['head_dim'] * 8192 * nvfp4_kv_bytes_per_elem)
    phases.append(OptimizationPhase(
        phase_number=6,
        name=f"ATTENTION: NVFP4 KV Cache Decode for {n_full_layers} Full Attention Layers",
        prerequisite="Phase 5 complete",
        what_to_build=(
            f"nvfp4_attention_decode.cu — Only {n_full_layers} of {params['num_hidden_layers']} layers "
            "use full attention with KV cache. "
            f"GQA ratio: {params['num_attention_heads']}:{params['num_key_value_heads']} "
            f"({params['num_attention_heads'] // params['num_key_value_heads']}:1). "
            f"Head dim: {params['head_dim']} (NOTE: existing kernel hardcodes 128 — needs rewrite). "
            "THREE PATHS (benchmark all): "
            "(A) mxf8f6f4.block_scale: Q=FP8×K=MX_FP4, 2x Hopper throughput, "
            "needs KV cache repacking NVFP4→MX (UE4M3→UE8M0, 16→32 group). "
            "(B) mxf4nvf4.block_scale: Q=FP4×K=NVFP4, 4x Hopper throughput (FASTEST), "
            "NVFP4 KV cache used directly, but Q loses precision (FP16→FP4). TN only. "
            "(C) Software LUT dequant: Q stays FP16, K/V dequanted via LUT to float. "
            "No tensor core for attention — pure bandwidth play. May win at M=1. "
            f"At 8K context: NVFP4 KV cache reads = {kv_bytes_8k / 1e6:.1f} MB "
            f"(vs {n_full_layers * 2 * params['num_key_value_heads'] * params['head_dim'] * 8192 / 1e6:.1f} MB for FP8). "
            "BLOCKERS: Must write missing .cuh headers (Q9), support head_dim=256 (Q10), "
            "add GQA config for 16Q/2KV/256dim. Use FlashDecoding split-K across GQA groups."
        ),
        expected_time_saved_us=0,
        expected_tps_after=0,  # Can't predict without attention timing
        measurement_command="bench_attention_decode: measure all 3 paths per-layer at 8K context",
        done_when="Attention adds < 20% overhead to GEMV-only forward time at 8K context",
    ))

    # Phase 7: Linear attention state update
    # CONFIRMED CUSTOM-ONLY: Research pipeline (864 LLM calls, 9-consensus) verified
    # NO cuDNN primitive for Mamba selective scan. 6/9 extractions hallucinated APIs
    # (CUDNN_BACKEND_OPERATION_SEQSCAN_DESCRIPTOR, CUDNN_OP_SCAN — none exist).
    # Linear attention O(1) per-token: all shipped kernels are O(N²). No library path.
    n_linear_layers = len([t for t in params["layer_types"] if t == "linear_attention"])
    phases.append(OptimizationPhase(
        phase_number=7,
        name=f"LINEAR ATTENTION: DeltaNet State Update for {n_linear_layers} Layers",
        prerequisite="Phase 6 complete",
        what_to_build=(
            f"deltanet_state_update.cu — {n_linear_layers} layers use linear attention "
            "(DeltaNet/Mamba hybrid). No KV cache — instead, maintain a recurrent state. "
            "State update is O(1) per token (no sequence length dependency). "
            "This is an element-wise operation on the state matrix, not a GEMV. "
            "CONFIRMED: No cuDNN/cuBLAS primitive exists for selective scan or "
            "linear attention state update. Fully custom CUDA required. "
            "Conv1D causal convolution also needed (kernel_dim="
            f"{params['linear_conv_kernel_dim']}). Conv is short — fuse with state update."
        ),
        expected_time_saved_us=0,
        expected_tps_after=0,
        measurement_command="bench_deltanet_state: measure state update time per layer",
        done_when="State update + conv1d < 5μs per layer",
    ))

    return phases


# =============================================================================
# Report Generation
# =============================================================================

def print_architecture_summary(params: dict):
    """Print a concise architecture summary from parsed config."""
    print("\n" + "=" * 80)
    print("  MODEL ARCHITECTURE SUMMARY")
    print("=" * 80)

    linear_count = sum(1 for t in params["layer_types"] if t == "linear_attention")
    full_count = sum(1 for t in params["layer_types"] if t == "full_attention")

    rows = [
        ("Architecture", "Qwen3NextForCausalLM (MoE + Mamba/DeltaNet hybrid)"),
        ("Total Layers", str(params["num_hidden_layers"])),
        ("  Linear Attention", f"{linear_count} layers (DeltaNet recurrent, no KV cache)"),
        ("  Full Attention", f"{full_count} layers (GQA with KV cache)"),
        ("Layer Pattern", f"[lin, lin, lin, full] × {full_count}"),
        ("Hidden Size", str(params["hidden_size"])),
        ("", ""),
        ("Full Attention", ""),
        ("  Q Heads × Dim", f"{params['num_attention_heads']} × {params['head_dim']} = {params['num_attention_heads'] * params['head_dim']}"),
        ("  KV Heads × Dim", f"{params['num_key_value_heads']} × {params['head_dim']} = {params['num_key_value_heads'] * params['head_dim']}"),
        ("  GQA Ratio", f"{params['num_attention_heads'] // params['num_key_value_heads']}:1"),
        ("", ""),
        ("Linear Attention", ""),
        ("  K Heads × Dim", f"{params['linear_num_key_heads']} × {params['linear_key_head_dim']} = {params['linear_num_key_heads'] * params['linear_key_head_dim']}"),
        ("  V Heads × Dim", f"{params['linear_num_value_heads']} × {params['linear_value_head_dim']} = {params['linear_num_value_heads'] * params['linear_value_head_dim']}"),
        ("  Conv Kernel", str(params["linear_conv_kernel_dim"])),
        ("", ""),
        ("MoE", ""),
        ("  Total Experts", str(params["num_experts"])),
        ("  Active / Token", str(params["num_experts_per_tok"])),
        ("  Expert FFN Dim", str(params["moe_intermediate_size"])),
        ("  Shared Expert Dim", str(params["shared_expert_intermediate_size"])),
        ("", ""),
        ("Quantization", ""),
        ("  Method", f"{params['quant_algo']} via {params['quant_method']}"),
        ("  KV Cache", f"{params['kv_cache_bits']}-bit {params['kv_cache_type']}"),
        ("  Vocab Size", str(params["vocab_size"])),
        ("  Max Context", str(params["max_position_embeddings"])),
    ]

    for label, value in rows:
        if label == "" and value == "":
            print()
        elif value == "":
            print(f"  {label}")
        else:
            print(f"  {label:.<40s} {value}")


def print_weight_matrix_inventory(weights: List[WeightMatrix]):
    """Print complete inventory of weight matrices."""
    print("\n" + "=" * 80)
    print("  WEIGHT MATRIX INVENTORY (per forward pass, M=1)")
    print("=" * 80)

    # Group by operation type
    by_operation: Dict[str, List[WeightMatrix]] = defaultdict(list)
    for w in weights:
        by_operation[w.operation].append(w)

    total_bytes = 0
    total_calls = 0

    for op_name in sorted(by_operation.keys()):
        ws = by_operation[op_name]
        op_bytes = sum(w.total_bytes_per_forward for w in ws)
        op_calls = sum(w.count_per_forward for w in ws)
        n_unique = len(set((w.shape_N, w.shape_K) for w in ws))
        w0 = ws[0]

        print(f"\n  {op_name}")
        print(f"    Shape: ({w0.shape_N}, {w0.shape_K})  Quant: {w0.quant.name}")
        print(f"    Matrices: {len(ws)}  Calls/fwd: {op_calls}  "
              f"Bytes/fwd: {op_bytes / 1e6:.2f} MB")

        total_bytes += op_bytes
        total_calls += op_calls

    print(f"\n  {'─' * 60}")
    print(f"  TOTAL: {total_calls} GEMV calls, {total_bytes / 1e6:.1f} MB read per token")
    print(f"  TOTAL: {total_bytes / 1e9:.3f} GB per token")


def print_gemv_shape_analysis(shapes: List[GemvShape], hw: HardwareSpec):
    """Print aggregated GEMV shapes with performance predictions."""
    print("\n" + "=" * 80)
    print(f"  GEMV SHAPE ANALYSIS — {hw.name}")
    print("=" * 80)

    models = compute_full_forward_pass_model(shapes, hw)

    total_time = sum(m.total_time_per_forward_us for m in models)
    total_bytes = sum(m.shape.total_read_bytes_per_forward for m in models)
    total_calls = sum(m.shape.total_calls_per_forward for m in models)
    total_launch = sum(m.launch_overhead_time_us for m in models)
    total_bw_time = sum(
        m.bandwidth_limited_time_us * m.shape.total_calls_per_forward
        for m in models
    )

    print(f"\n  {'Shape (N×K)':<18s} {'Quant':<6s} {'Calls':>6s} {'Bytes/fwd':>10s} "
          f"{'BW Time':>9s} {'Launch':>9s} {'Total':>9s} {'Bottleneck':<16s}")
    print(f"  {'─' * 18} {'─' * 6} {'─' * 6} {'─' * 10} "
          f"{'─' * 9} {'─' * 9} {'─' * 9} {'─' * 16}")

    for m in models:
        s = m.shape
        bw_time = m.bandwidth_limited_time_us * s.total_calls_per_forward
        print(f"  {s.shape_N:>7d}×{s.shape_K:<7d}  {s.quant.name:<6s} {s.total_calls_per_forward:>6d} "
              f"{s.total_read_bytes_per_forward / 1e6:>9.1f}M "
              f"{bw_time:>8.0f}μs {m.launch_overhead_time_us:>8.0f}μs "
              f"{m.total_time_per_forward_us:>8.0f}μs {m.bottleneck:<16s}")
        for op in s.operation_names:
            print(f"    └─ {op}")

    print(f"\n  {'─' * 80}")
    print(f"  Total GEMV calls per forward pass: {total_calls}")
    print(f"  Total bytes read per forward pass: {total_bytes / 1e6:.1f} MB ({total_bytes / 1e9:.3f} GB)")
    print(f"  Bandwidth-limited time:            {total_bw_time:.0f} μs")
    print(f"  Launch overhead time:              {total_launch:.0f} μs")
    print(f"  Total estimated time:              {total_time:.0f} μs")
    print(f"  Launch overhead fraction:          {total_launch / total_time * 100:.1f}%")
    print(f"  Bandwidth-limited TPS ceiling:     {1e6 / total_bw_time:.0f} TPS")
    print(f"  Current estimated TPS:             {1e6 / total_time:.0f} TPS")


def print_optimization_roadmap(phases: List[OptimizationPhase]):
    """Print the deterministic optimization roadmap."""
    print("\n" + "=" * 80)
    print("  OPTIMIZATION ROADMAP")
    print("=" * 80)
    print("\n  This is not a menu. Execute phases in order.")
    print("  Each phase builds on the previous. Do not skip.\n")

    for p in phases:
        status = "DONE" if p.phase_number == 0 else "TODO"
        print(f"  ┌─ Phase {p.phase_number}: {p.name}")
        print(f"  │  Status: {status}")
        print(f"  │  Prerequisite: {p.prerequisite}")
        print(f"  │")
        print(f"  │  WHAT TO BUILD:")
        # Word-wrap the build description
        desc = p.what_to_build
        for line in _wrap_text(desc, width=68):
            print(f"  │    {line}")
        print(f"  │")
        if p.expected_time_saved_us > 0:
            print(f"  │  Expected time saved: {p.expected_time_saved_us:.0f} μs")
        if p.expected_tps_after > 0:
            print(f"  │  Expected TPS after:  {p.expected_tps_after:.0f}")
        print(f"  │")
        print(f"  │  MEASURE:")
        print(f"  │    {p.measurement_command}")
        print(f"  │")
        print(f"  │  DONE WHEN:")
        print(f"  │    {p.done_when}")
        print(f"  └{'─' * 78}")
        print()


def print_theoretical_ceilings(params: dict, shapes: List[GemvShape], hw: HardwareSpec):
    """Print the hard physics ceilings."""
    print("\n" + "=" * 80)
    print(f"  THEORETICAL CEILINGS — {hw.name}")
    print("=" * 80)

    total_bytes = sum(s.total_read_bytes_per_forward for s in shapes)
    total_calls = sum(s.total_calls_per_forward for s in shapes)

    bw_ceiling_us = total_bytes / (hw.memory_bandwidth_gb_s * 1e3)
    bw_ceiling_tps = 1e6 / bw_ceiling_us

    launch_overhead_individual = total_calls * hw.kernel_launch_overhead_us
    launch_overhead_graph = total_calls * hw.cuda_graph_launch_overhead_us

    print(f"\n  Data movement per token:    {total_bytes / 1e6:.1f} MB")
    print(f"  DRAM bandwidth:             {hw.memory_bandwidth_gb_s:.0f} GB/s")
    print(f"  Bandwidth-limited time:     {bw_ceiling_us:.1f} μs")
    print(f"  Bandwidth-limited TPS:      {bw_ceiling_tps:.0f}")
    print(f"")
    print(f"  Kernel launches per token:  {total_calls}")
    print(f"  Launch overhead (individual): {launch_overhead_individual:.0f} μs")
    print(f"  Launch overhead (CUDA Graph):  {launch_overhead_graph:.0f} μs")
    print(f"")
    print(f"  Floor (BW + graph launch):  {bw_ceiling_us + launch_overhead_graph:.0f} μs = "
          f"{1e6 / (bw_ceiling_us + launch_overhead_graph):.0f} TPS")
    print(f"  Floor (BW only, perfect):   {bw_ceiling_us:.0f} μs = {bw_ceiling_tps:.0f} TPS")

    # KV cache reads at various context lengths
    # Native NVFP4 KV cache: 0.5625 bytes/elem (4-bit E2M1 + UE4M3 scale per 16)
    # FP8 KV cache: 1.0 bytes/elem (shown for comparison)
    n_full = sum(1 for t in params["layer_types"] if t == "full_attention")
    kv_heads = params["num_key_value_heads"]
    head_dim = params["head_dim"]
    nvfp4_kv_bpe = NVFP4.bytes_per_element   # 0.5625
    fp8_kv_bpe = FP8_E4M3.bytes_per_element  # 1.0
    kv_elems_per_token_per_layer = 2 * kv_heads * head_dim  # K + V

    print(f"\n  KV Cache Overhead ({n_full} full attention layers):")
    print(f"    {'Context':>9s}  {'NVFP4 (native)':>14s}  {'FP8 (fallback)':>14s}  {'NVFP4 time':>10s}  {'% of GEMV':>9s}")
    for ctx in [1024, 4096, 8192, 32768, 131072]:
        nvfp4_bytes = n_full * kv_elems_per_token_per_layer * ctx * nvfp4_kv_bpe
        fp8_bytes = n_full * kv_elems_per_token_per_layer * ctx * fp8_kv_bpe
        nvfp4_time = nvfp4_bytes / (hw.memory_bandwidth_gb_s * 1e3)
        pct = nvfp4_time / bw_ceiling_us * 100
        print(f"    {ctx:>7d}:   {nvfp4_bytes / 1e6:>10.1f} MB   {fp8_bytes / 1e6:>10.1f} MB   "
              f"{nvfp4_time:>8.1f} μs   {pct:>6.1f}%")


def _wrap_text(text: str, width: int = 68) -> List[str]:
    """Simple word wrap."""
    words = text.split()
    lines = []
    current = []
    current_len = 0
    for word in words:
        if current_len + len(word) + 1 > width and current:
            lines.append(" ".join(current))
            current = [word]
            current_len = len(word)
        else:
            current.append(word)
            current_len += len(word) + 1
    if current:
        lines.append(" ".join(current))
    return lines


# =============================================================================
# State Persistence — results.json
# =============================================================================

RESULTS_DIR = Path(__file__).parent / "data"
RESULTS_FILE = RESULTS_DIR / "primer_results.json"

DEFAULT_CONFIG = (
    "/home/adeliant/Desktop/archives/huggingface-cache/hub/"
    "models--nvidia--Qwen3-Next-80B-A3B-Thinking-NVFP4/"
    "snapshots/47f93711bad0ce01a49f720ad43c0272687b5d5f/config.json"
)


@dataclass
class PhaseResult:
    """Persistent state for one optimization phase."""
    phase_number: int
    status: str                         # "pending", "in_progress", "done", "blocked"
    measurements: Dict[str, float] = field(default_factory=dict)
    baselines: Dict[str, float] = field(default_factory=dict)  # Locked when done
    started: Optional[str] = None
    completed: Optional[str] = None
    notes: List[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "phase_number": self.phase_number,
            "status": self.status,
            "measurements": self.measurements,
            "baselines": self.baselines,
            "started": self.started,
            "completed": self.completed,
            "notes": self.notes,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "PhaseResult":
        return cls(
            phase_number=d["phase_number"],
            status=d["status"],
            measurements=d.get("measurements", {}),
            baselines=d.get("baselines", {}),
            started=d.get("started"),
            completed=d.get("completed"),
            notes=d.get("notes", []),
        )


def load_results() -> Dict[int, PhaseResult]:
    """Load results from disk. Returns empty dict if no file."""
    if not RESULTS_FILE.exists():
        return {}
    with open(RESULTS_FILE, 'r') as f:
        data = json.load(f)
    results = {}
    for k, v in data.get("phases", {}).items():
        results[int(k)] = PhaseResult.from_dict(v)
    return results


def save_results(results: Dict[int, PhaseResult]):
    """Write results to disk."""
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    data = {
        "phases": {str(k): v.to_dict() for k, v in sorted(results.items())},
        "last_updated": _now(),
    }
    with open(RESULTS_FILE, 'w') as f:
        json.dump(data, f, indent=2)


def _now() -> str:
    from datetime import datetime
    return datetime.now().isoformat(timespec='seconds')


def ensure_phase_exists(results: Dict[int, PhaseResult], phase_num: int) -> PhaseResult:
    """Get or create a phase result entry."""
    if phase_num not in results:
        results[phase_num] = PhaseResult(phase_number=phase_num, status="pending")
    return results[phase_num]


# =============================================================================
# Phase Completion Criteria
# =============================================================================

# Each phase defines required measurements and acceptance thresholds.
# If all required measurements exist and meet thresholds, phase is "done".
# If any measurement regresses a locked baseline, it's flagged.

PHASE_REQUIREMENTS = {
    0: {
        "required_measurements": [
            "nvfp4_lm_head_time_us",
            "nvfp4_lm_head_bw_eff",
            "nvfp4_all_shapes_pass_correctness",
        ],
        "acceptance": {
            "nvfp4_all_shapes_pass_correctness": ("==", 1.0),
        },
        "description": "NVFP4 GEMV baseline — all shapes correct, timing recorded",
    },
    1: {
        "required_measurements": [
            "batched_expert_time_us",
            "individual_expert_time_us",
            "batched_correctness_pass",
        ],
        "acceptance": {
            "batched_correctness_pass": ("==", 1.0),
        },
        "description": "Batched expert GEMV — fused launch faster than individual",
    },
    2: {
        "required_measurements": [
            "graph_replay_total_us",
            "individual_launch_total_us",
        ],
        "acceptance": {},
        "description": "CUDA Graph — replay time < individual launch time",
    },
    3: {
        "required_measurements": [
            "registers_per_thread",
            "warp_occupancy_pct",
        ],
        "acceptance": {
            "registers_per_thread": ("<=", 48.0),
            "warp_occupancy_pct": (">=", 83.0),
        },
        "description": "Occupancy — registers ≤ 48, occupancy ≥ 83%",
    },
    4: {
        "required_measurements": [
            "bf16_lm_head_bw_eff",
        ],
        "acceptance": {
            "bf16_lm_head_bw_eff": (">=", 0.70),
        },
        "description": "BF16 LM head — bandwidth efficiency ≥ 70%",
    },
    5: {
        "required_measurements": [
            "elementwise_overhead_pct",
        ],
        "acceptance": {
            "elementwise_overhead_pct": ("<=", 10.0),
        },
        "description": "Fused elementwise — overhead < 10%",
    },
    6: {
        "required_measurements": [
            "attention_overhead_pct",
        ],
        "acceptance": {
            "attention_overhead_pct": ("<=", 20.0),
        },
        "description": "Attention decode — overhead < 20% at 8K context",
    },
    7: {
        "required_measurements": [
            "deltanet_state_update_per_layer_us",
        ],
        "acceptance": {
            "deltanet_state_update_per_layer_us": ("<=", 5.0),
        },
        "description": "DeltaNet state update — < 5μs/layer",
    },
}


# =============================================================================
# Research Checklist — per-phase queries for the docs vector database
# =============================================================================

PHASE_RESEARCH = {
    0: {
        "description": "Baseline GEMV kernel fundamentals",
        "queries": [
            "NVFP4 E2M1 dequantization lookup table implementation CUDA",
            "SM120 shared memory LUT optimization GEMV kernel",
            "warp shuffle reduction butterfly pattern __shfl_xor_sync",
            "SM12.0 compute capability register file size warps per SM",
        ],
    },
    1: {
        "description": "Batched GEMV for MoE expert dispatch",
        "queries": [
            "CUTLASS SM120 grouped GEMM block-scaled example 82",
            "batched GEMV single kernel launch multiple weight matrices dispatch table",
            "MoE expert parallelism fused kernel single launch CUDA",
            "cuBLAS gemm_grouped_batched_ex API usage",
        ],
    },
    2: {
        "description": "CUDA Graphs for launch overhead elimination",
        "queries": [
            "CUDA Graph stream capture cudaStreamBeginCapture cudaStreamEndCapture",
            "cudaGraphExecKernelNodeSetParams update kernel parameters between replays",
            "CUDA Graph performance SM120 Blackwell launch overhead reduction",
            "CUDA Graph capture GEMM kernel sequence replay static shapes",
        ],
    },
    3: {
        "description": "Occupancy tuning and register pressure on SM120",
        "queries": [
            "SM12.0 occupancy calculator registers per thread warps per SM",
            "__launch_bounds__ maxThreadsPerBlock minBlocksPerMultiprocessor SM120",
            "CUDA register pressure reduction techniques pragma unroll half2 packing",
            "Blackwell tuning guide occupancy optimization SM120",
        ],
    },
    4: {
        "description": "BF16 GEMV for LM head and attention projections",
        "queries": [
            "BF16 vectorized load 128-bit coalesced memory access CUDA kernel",
            "large N GEMV N=151936 bandwidth optimization GPU kernel",
            "SM120 memory bandwidth utilization BF16 matrix vector multiply",
        ],
    },
    5: {
        "description": "Fused elementwise operations (RMSNorm, SiLU, routing)",
        "queries": [
            "RMSNorm CUDA kernel implementation fused with linear projection",
            "SiLU activation gated linear unit fused kernel CUDA",
            "top-K routing softmax MoE expert selection CUDA kernel",
            "fused elementwise operations reducing kernel launch overhead",
        ],
    },
    6: {
        "description": "NVFP4 attention decode — 3 paths: mxf8f6f4(FP8×FP4), mxf4nvf4(FP4×FP4), software dequant",
        "queries": [
            "FlashDecoding GQA grouped query attention decode single token",
            # RESOLVED via nvidia-docs-search 2026-02-11:
            #   mxf8f6f4.block_scale: A=mx_float8_t × B=mx_float4_t (mixed FP8×FP4)
            #   mxf4nvf4.block_scale: nv_float4_t × nv_float4_t (4x Hopper throughput)
            #   Accumulator: FP32. SM120 shapes: m16n8k64, m16n8k128.
            # cuDNN 9.9.x CAN do FP8 KV dequant via Pointwise Multiply graph (confirmed),
            # but that's the fallback, not the primary target.
            "FP8 KV cache dequantization in attention inner loop CUDA",
            "cuDNN attention backend SM120 fused multi-head attention decode",
            "attention decode kernel split-K reduction across heads",
        ],
    },
    7: {
        "description": "DeltaNet linear attention state update",
        "queries": [
            "DeltaNet linear attention recurrent state update CUDA kernel",
            "Mamba state space model selective scan CUDA implementation",
            "linear attention O(1) per token state update kernel",
            "causal convolution Conv1D kernel_size=4 CUDA fused",
        ],
    },
}


# =============================================================================
# Open Questions & Design Decisions
#
# PSEUDOCODE: Track every unknown that blocks implementation.
# Each question has:
#   - status: "resolved" | "unresolved" | "decision_needed"
#   - resolution: what we found (if resolved)
#   - source: where we found it (PTX ISA, docs-db search, etc.)
#   - impact: which phase(s) this blocks
#   - options: for design decisions, the viable paths
#
# The cmd_unknowns() function prints these in a table.
# Agents update status here as they discover answers.
# =============================================================================

OPEN_QUESTIONS = {
    # =========================================================================
    # RESOLVED — answers found via nvidia-docs-search (2026-02-11)
    # =========================================================================
    "Q1_mixed_precision_mma": {
        "question": "Can SM120 block_scale MMA handle mixed FP8×FP4 operands for attention?",
        "status": "resolved",
        "resolution": (
            "YES. mxf8f6f4.block_scale supports A=mx_float8_t × B=mx_float4_t "
            "(and all {mxf4,mxf6,mxf8}×{mxf4,mxf6,mxf8} combos). Uses UE8M0 scale, group=32. "
            "Also mxf4nvf4.block_scale supports nv_float4_t×nv_float4_t (UE4M3 scale, group=16, "
            "TN only, 4x Hopper FP8 throughput). Both confirmed on sm_120a via PTX ISA 8.7+."
        ),
        "source": "CUTLASS blackwell_functionality.html Table 3 rows 1,6; PTX ISA chunks 107,131",
        "impact": ["phase_6"],
        "resolved_date": "2026-02-11",
    },
    "Q2_accumulator_precision": {
        "question": "What accumulator precision does block_scale MMA use?",
        "status": "resolved",
        "resolution": (
            "FP32. All mma.sync.aligned.block_scale variants on SM120 use .f32 accumulators. "
            "Softmax over FP32 logits is numerically safe."
        ),
        "source": "PTX ISA chunk 131: mma.sync.aligned examples show .f32 accumulator",
        "impact": ["phase_6"],
        "resolved_date": "2026-02-11",
    },
    "Q3_scale_factor_formats": {
        "question": "What scale factor formats do the two block_scale kinds use?",
        "status": "resolved",
        "resolution": (
            "mxf8f6f4.block_scale: UE8M0 (power-of-2), 32 elements per block. "
            "mxf4nvf4.block_scale: UE4M3 (137x more range), 16 elements per block. "
            "NVFP4 KV cache uses UE4M3/group=16 → directly compatible with mxf4nvf4, "
            "but needs repacking for mxf8f6f4 (UE4M3→UE8M0, 16→32 group, lossy)."
        ),
        "source": "CUTLASS blackwell_functionality.html 'Block Scaled Narrow Precision Data Types' table",
        "impact": ["phase_6"],
        "resolved_date": "2026-02-11",
    },
    "Q4_deltanet_in_nvidia_docs": {
        "question": "Does NVIDIA documentation cover DeltaNet/linear attention?",
        "status": "resolved",
        "resolution": (
            "NO. Docs-db search returned only LSTM/GRU results (rerank score 0.054 — "
            "effectively no match). Confirms: fully custom CUDA required for all 36 "
            "linear attention layers. No cuDNN primitive, no cuBLAS path."
        ),
        "source": "docs-db search 'DeltaNet linear attention recurrence' → 0.054 rerank",
        "impact": ["phase_7"],
        "resolved_date": "2026-02-11",
    },
    "Q5_mma_shapes_sm120": {
        "question": "What MMA shapes does SM120 block_scale support?",
        "status": "resolved",
        "resolution": (
            "mma.sync.aligned with block_scale on SM120: m16n8k64 and m16n8k128. "
            "mxf4nvf4 and mxf4 kinds: sm_120a and sm_121a only. "
            "mxf8f6f4 kind: sm_120a, sm_120f, sm_121a."
        ),
        "source": "PTX ISA chunk 131",
        "impact": ["phase_6"],
        "resolved_date": "2026-02-11",
    },

    # =========================================================================
    # UNRESOLVED — need external sources (not in NVIDIA docs)
    # =========================================================================
    "Q6_deltanet_recurrence": {
        "question": "What are the DeltaNet recurrence equations for Qwen3-Next-80B?",
        "status": "unresolved",
        "resolution": None,
        "source": "Need: HuggingFace modeling_qwen3_next.py",
        "impact": ["phase_7"],
        "notes": (
            "DeltaNet variants differ (gated, multi-scale, etc.). Must read the actual "
            "model code to get: state transition matrix, forget gate, input gate, "
            "output gate, state size (hidden×head_dim per head? compressed?). "
            "Also need conv1d causal conv details (kernel_dim=4 from config)."
        ),
    },
    "Q7_deltanet_state_shape": {
        "question": "What is the DeltaNet recurrent state tensor shape?",
        "status": "unresolved",
        "resolution": None,
        "source": "Need: HuggingFace modeling_qwen3_next.py",
        "impact": ["phase_7"],
        "notes": (
            "State could be [batch, heads, key_dim, value_dim] or compressed. "
            "For Qwen3-Next: linear_num_key_heads=16, linear_key_head_dim=128, "
            "linear_num_value_heads=32, linear_value_head_dim=128. "
            "State size per layer = unknown until we read the model code."
        ),
    },
    "Q8_q_quantization_precision": {
        "question": "How much precision does Q lose when quantized FP16→FP4 vs FP16→FP8?",
        "status": "unresolved",
        "resolution": None,
        "source": "Need: empirical test or literature survey",
        "impact": ["phase_6"],
        "notes": (
            "For mxf4nvf4: Q must be FP4 (e2m1, range ±6, 3 values per octave). "
            "For mxf8f6f4: Q can be FP8 (e4m3, range ±448, 8 values per octave). "
            "Q vectors are freshly computed per token — full dynamic range of hidden_size=2048. "
            "FP4 quantization of Q may be acceptable if block scales capture the range, "
            "but attention is sensitive to Q precision (softmax amplifies errors). "
            "Key test: measure perplexity impact of FP4 Q vs FP8 Q vs FP16 Q."
        ),
    },
    "Q9_missing_cuh_headers": {
        "question": "The attention kernel references 5 .cuh headers that don't exist on disk.",
        "status": "unresolved",
        "resolution": None,
        "source": "Glob sm12x/**/*.cuh returned empty",
        "impact": ["phase_6"],
        "notes": (
            "Missing: sm12x/attention/gqa_config.cuh, sm12x/attention/nvfp4_kv_cache.cuh, "
            "sm12x/quant/e2m1_types.cuh, sm12x/quant/ue4m3_scale.cuh, sm12x/quant/fp8_convert.cuh. "
            "These were planned but never written. Kernel won't compile without them. "
            "Must write these before any attention kernel work."
        ),
    },
    "Q10_head_dim_256": {
        "question": "Existing attention kernel hardcodes head_dim=128. Model needs 256.",
        "status": "unresolved",
        "resolution": None,
        "source": "nvfp4_attention.cu static_assert(GQA::kHeadDim == 128)",
        "impact": ["phase_6"],
        "notes": (
            "Qwen3-Next-80B full attention: 16 Q heads × 256 dim, 2 KV heads × 256 dim. "
            "Current kernel: warp-parallel flash decoding with 4 warps × 32 threads. "
            "At head_dim=256: each thread must accumulate 256/32 = 8 elements (vs 4 at 128). "
            "Register usage doubles. Need complete redesign of inner loop. "
            "Also need new GQA config: {16Q, 2KV, 256dim} not in current instantiations."
        ),
    },
    "Q11_cuda_graph_variable_kv": {
        "question": "How to handle CUDA Graph capture with variable KV cache length?",
        "status": "unresolved",
        "resolution": None,
        "source": "Phase 2 research accepted but no implementation",
        "impact": ["phase_2"],
        "notes": (
            "KV cache grows by 1 entry per decode step. Graph captured at length L "
            "won't work at length L+1. Options: (a) graph-per-bucket (powers of 2), "
            "(b) cudaGraphExecKernelNodeSetParams to update seq_len, "
            "(c) max-length pre-allocation with masking. Phase 2 research has details."
        ),
    },

    # =========================================================================
    # DESIGN DECISIONS — viable paths identified, need to choose
    # =========================================================================
    "D1_attention_mma_path": {
        "question": "Which block_scale MMA kind to use for attention QK^T and PV?",
        "status": "resolved",
        "resolution": (
            "Path A (mxf8f6f4) ELIMINATED — requires MXFP4 KV cache (UE8M0/group=32), "
            "not NVFP4. After repacking NVFP4→MX, the cache IS MXFP4, breaking the "
            "NVFP4kv project goal. Viable paths: B (mxf4nvf4 tensor cores) and "
            "C (software LUT dequant). Path C already implemented in TRT-LLM fork "
            "(nvfp4KvCacheUtils.h). Path C wins at M=1 decode (bandwidth-bound). "
            "Path B reserved for compute-bound scenarios (large batch/long context)."
        ),
        "resolved_date": "2026-02-12",
        "source": "Resolved Q1 + Q3 + user correction on MXFP4 format",
        "impact": ["phase_6"],
        "options": {
            "A_mxf8f6f4": (
                "ELIMINATED — OUTSIDE PROJECT SCOPE. "
                "mxf8f6f4.block_scale requires MX format KV cache (UE8M0/group=32). "
                "Repacking NVFP4→MX is lossy and changes the cache format to MXFP4. "
                "This violates the NVFP4kv project goal."
            ),
            "B_mxf4nvf4": (
                "mxf4nvf4.block_scale: Q=nv_float4_t × K=nv_float4_t. "
                "Pro: NVFP4 KV cache used directly (UE4M3/group=16), no repacking. "
                "Con: Q quantized FP16→FP4 (precision loss). TN layout only. "
                "Throughput: 4x Hopper FP8 (HIGHEST — 2x faster than option A)."
            ),
            "C_software_dequant": (
                "RECOMMENDED FOR DECODE. Software LUT dequant (TRT-LLM fork). "
                "Pro: Maximum Q precision (stays FP16). Already implemented. "
                "Con: No tensor core acceleration. Bandwidth-bound only. "
                "Throughput: ~70-85% of DRAM bandwidth. Wins at M=1 decode."
            ),
        },
        "recommendation": (
            "Path C (software dequant) is the primary path — already implemented in "
            "TRT-LLM fork, highest precision, and wins at decode (bandwidth-bound). "
            "Path B (mxf4nvf4) is a future optimization for compute-bound scenarios "
            "(large batch, long context prefill). Path A is out of scope."
        ),
    },
    "D2_kv_cache_format": {
        "question": "Should KV cache use NVFP4 (UE4M3/16) or MX (UE8M0/32) format?",
        "status": "decision_needed",
        "resolution": None,
        "source": "Resolved Q3",
        "impact": ["phase_6"],
        "options": {
            "A_nvfp4": (
                "NVFP4: UE4M3 scale, 16 elements/block. "
                "Pro: 137x more scale range, finer granularity. "
                "Compatible with mxf4nvf4 MMA directly. "
                "Con: Incompatible with mxf8f6f4 MMA (needs repacking)."
            ),
            "B_mx": (
                "MX: UE8M0 scale (power-of-2), 32 elements/block. "
                "Pro: Compatible with mxf8f6f4 MMA directly, OCP standard. "
                "Con: Coarser scaling (32 vs 16), power-of-2 only (less range). "
                "Would require changing quantize/dequantize kernels."
            ),
        },
        "recommendation": (
            "Keep NVFP4 (UE4M3/16). Finer scaling preserves more KV cache precision. "
            "If mxf8f6f4 path wins benchmarks (D1 option A), add a lightweight "
            "NVFP4→MX repacking kernel that runs during cache read."
        ),
    },
}


# =============================================================================
# Research Pipeline Configuration
#
# TOKEN BUDGET POLICY: gpt-oss-120b is a reasoning model. Reasoning tokens
# are INVISIBLE but count against max_tokens. A 500-token JSON answer may
# cost 4000+ tokens internally. We set generous limits and detect truncation
# rather than constraining the model. If a response is truncated, RAISE the
# limit — never lower it. The cost of re-running 6 consensus calls because
# one got clipped dwarfs the cost of the extra tokens.
#
# TEMPERATURE POLICY: 0.2 for all structured JSON extraction. Reasoning
# models need determinism for consistent schema compliance, not creativity.
# =============================================================================

DOCS_DB_TOOL = str(Path(__file__).parent / "infrastructure" / "docs-db" / "query_docs.py")
DOCS_DB_SERVER = "http://127.0.0.1:9712"
# Two RTX Pro 6000 GPUs, each running gpt-oss-120b. Round-robin across both.
LLM_ENDPOINTS = [
    "http://localhost:8002/v1/chat/completions",  # GPU 1
    "http://localhost:8003/v1/chat/completions",  # GPU 0
]
LLM_MODEL = "openai/gpt-oss-120b"
RESEARCH_CONSENSUS_N = 9              # 9 parallel LLM calls per stage
RESEARCH_CONFIDENCE_THRESHOLD = 0.667  # 2/3 majority (6/9) to accept
_llm_call_counter = 0  # round-robin counter for load balancing
RESEARCH_RESULTS_FILE = str(Path(__file__).parent / "data" / "primer_research.json")


def _query_docs_db(query: str, n: int = 15) -> List[Dict]:
    """Query the docs vector database via HTTP server API.
    Server must be running: python3 infrastructure/docs-db/query_docs.py serve
    Returns list of {rank, rerank_score, title, path, text}."""
    import urllib.request
    import urllib.parse

    url = f"{DOCS_DB_SERVER}/search?q={urllib.parse.quote(query)}&n={n}"
    try:
        with urllib.request.urlopen(url, timeout=30) as resp:
            data = json.loads(resp.read().decode())
            results = data.get("results", [])
            chunks = []
            for r in results:
                chunks.append({
                    "rank": r.get("rank", 0),
                    "rerank_score": r.get("rerank_score", 0.0),
                    "title": r.get("title", ""),
                    "path": r.get("rel_path", ""),
                    "text": r.get("document", ""),
                })
            return chunks
    except Exception as e:
        return [{"error": str(e)}]


def _find_knee(chunks: List[Dict]) -> List[Dict]:
    """Return chunks above the score knee (biggest drop between consecutive scores)."""
    if len(chunks) <= 2:
        return chunks

    scores = [c.get("rerank_score", 0) for c in chunks]
    max_drop = 0
    knee_idx = len(scores)

    for i in range(1, len(scores)):
        drop = scores[i - 1] - scores[i]
        if drop > max_drop:
            max_drop = drop
            knee_idx = i

    # Always return at least 3 chunks
    knee_idx = max(knee_idx, min(3, len(chunks)))
    return chunks[:knee_idx]


def _call_llm(messages: List[Dict], temperature: float = 0.2,
              json_mode: bool = True, max_tokens: int = 16384,
              **kwargs) -> Dict:
    """Call gpt-oss-120b and return parsed response.

    BURN TOKENS TO SAVE OURS. gpt-oss-120b is a reasoning model — internal
    chain-of-thought tokens are invisible but count against max_tokens. A
    typical extraction response uses ~2200 completion tokens; validation ~400-2500.
    We set max_tokens=16384 (7x headroom) and detect truncation rather than
    constraining the model's reasoning. If finish_reason=="length", the response
    was clipped — raise max_tokens, never lower it.

    Temperature fixed at 0.2 for deterministic structured JSON output.
    json_mode=True forces response_format=json_object.
    Timeout 300s accommodates the reasoning model's think time.
    """
    import urllib.request

    # GENEROUS TOKEN BUDGET — do not reduce. Reasoning models need room to think.
    # The cost of 6 truncated consensus calls >> cost of extra completion tokens.
    payload = {
        "model": LLM_MODEL,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    if json_mode:
        payload["response_format"] = {"type": "json_object"}

    data = json.dumps(payload).encode()
    # Round-robin across both RTX Pro 6000 GPUs for load balancing
    global _llm_call_counter
    endpoint = LLM_ENDPOINTS[_llm_call_counter % len(LLM_ENDPOINTS)]
    _llm_call_counter += 1
    req = urllib.request.Request(
        endpoint,
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST"
    )

    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            result = json.loads(resp.read().decode())
            finish_reason = result["choices"][0].get("finish_reason", "")
            if finish_reason == "length":
                return {"error": f"TRUNCATED at {max_tokens} tokens — raise max_tokens"}
            content = result["choices"][0]["message"]["content"]
            if json_mode:
                return json.loads(content)
            return {"raw": content}
    except json.JSONDecodeError as e:
        return {"error": f"JSON parse failed (likely truncated): {e}"}
    except (urllib.error.URLError, ConnectionRefusedError, OSError) as e:
        # Endpoint down — try the other GPU
        other = [ep for ep in LLM_ENDPOINTS if ep != endpoint]
        if other:
            fallback_req = urllib.request.Request(
                other[0], data=data,
                headers={"Content-Type": "application/json"}, method="POST"
            )
            try:
                with urllib.request.urlopen(fallback_req, timeout=300) as resp:
                    result = json.loads(resp.read().decode())
                    finish_reason = result["choices"][0].get("finish_reason", "")
                    if finish_reason == "length":
                        return {"error": f"TRUNCATED at {max_tokens} tokens — raise max_tokens"}
                    content = result["choices"][0]["message"]["content"]
                    if json_mode:
                        return json.loads(content)
                    return {"raw": content}
            except Exception as e2:
                return {"error": f"Both endpoints failed: {e} / {e2}"}
        return {"error": str(e)}
    except Exception as e:
        return {"error": str(e)}


def _call_llm_batch(messages_list: List[List[Dict]], **kwargs) -> List[Dict]:
    """Call LLM N times concurrently. Returns list of responses."""
    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(messages_list)) as pool:
        futures = [pool.submit(_call_llm, msgs, **kwargs) for msgs in messages_list]
        for f in concurrent.futures.as_completed(futures):
            results.append(f.result())
    return results


def _load_research() -> Dict:
    """Load persisted research results."""
    try:
        with open(RESEARCH_RESULTS_FILE, "r") as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {"phases": {}}


def _save_research(data: Dict):
    """Save research results with merge-on-write for concurrent safety.
    Multiple research processes may run in parallel (one per phase).
    Each process only modifies its own phase key, so we reload the
    current file before writing and merge — no phase gets overwritten
    by a stale snapshot from another process."""
    import fcntl
    Path(RESEARCH_RESULTS_FILE).parent.mkdir(parents=True, exist_ok=True)
    lock_path = RESEARCH_RESULTS_FILE + ".lock"
    with open(lock_path, "w") as lockf:
        fcntl.flock(lockf, fcntl.LOCK_EX)
        try:
            # Reload current state from disk (another process may have updated it)
            try:
                with open(RESEARCH_RESULTS_FILE, "r") as f:
                    current = json.load(f)
            except (FileNotFoundError, json.JSONDecodeError):
                current = {"phases": {}}
            # Merge at QUERY level within each phase — never clobber completed queries
            for phase_key, phase_data in data.get("phases", {}).items():
                existing_phase = current.setdefault("phases", {}).get(phase_key, {})
                if not existing_phase:
                    current["phases"][phase_key] = phase_data
                else:
                    # Merge queries: new results overwrite same query, preserve others
                    existing_queries = existing_phase.setdefault("queries", {})
                    for qk, qv in phase_data.get("queries", {}).items():
                        existing_queries[qk] = qv
                    # Update phase metadata
                    existing_phase["completed"] = phase_data.get("completed", existing_phase.get("completed"))
                    existing_phase["overall_confidence"] = phase_data.get("overall_confidence", 0)
                    current["phases"][phase_key] = existing_phase
            current["last_updated"] = _now()
            with open(RESEARCH_RESULTS_FILE, "w") as f:
                json.dump(current, f, indent=2)
        finally:
            fcntl.flock(lockf, fcntl.LOCK_UN)


def run_research_pipeline(phase_num: int, force: bool = False) -> Dict:
    """
    Run the 3-phase research pipeline for a given phase:
      A) Sufficiency — do the docs answer the question? (6x consensus)
      B) Extraction  — distill into structured JSON (6x)
      C) Validation  — does the extraction faithfully represent the docs? (6x)

    Returns dict with per-query results.
    """
    research_def = PHASE_RESEARCH.get(phase_num)
    if not research_def:
        print(f"  No research queries defined for phase {phase_num}")
        return {}

    existing = _load_research()
    phase_key = str(phase_num)
    if phase_key in existing.get("phases", {}) and not force:
        print(f"  Phase {phase_num} research already completed. Use --force to re-run.")
        return existing["phases"][phase_key]

    print(f"\n  {'═' * 70}")
    print(f"  RESEARCH PIPELINE — Phase {phase_num}: {research_def['description']}")
    print(f"  {'═' * 70}")

    phase_results = {"queries": {}, "completed": _now(), "overall_confidence": 0.0}

    for qi, query in enumerate(research_def["queries"]):
        print(f"\n  ── Query {qi + 1}/{len(research_def['queries'])}: {query}")

        # Step 1: Retrieve docs
        print(f"     Querying docs database...")
        all_chunks = _query_docs_db(query, n=15)
        if not all_chunks or "error" in all_chunks[0]:
            print(f"     ERROR: Docs query failed: {all_chunks}")
            phase_results["queries"][query] = {"status": "docs_failed"}
            continue

        above_knee = _find_knee(all_chunks)
        score_list = [f"{c.get('rerank_score', 0):.4f}" for c in above_knee]
        print(f"     Retrieved {len(all_chunks)} chunks, {len(above_knee)} above knee "
              f"(scores: {score_list})")

        # Combine chunk text for LLM context
        docs_context = ""
        for c in above_knee:
            docs_context += f"\n--- Source: {c.get('path', 'unknown')} (score: {c.get('rerank_score', 0):.4f}) ---\n"
            docs_context += c.get("text", "")[:3000]  # cap per chunk
            docs_context += "\n"

        # Phase A: Sufficiency check (6x concurrent)
        print(f"     Phase A: Sufficiency check ({RESEARCH_CONSENSUS_N}x)...")
        sufficiency_msgs = [[
            {"role": "system", "content": "You are a technical documentation analyst. "
             "Evaluate whether the provided documentation excerpts sufficiently answer the given question. "
             "Return JSON: {\"sufficient\": true/false, \"confidence\": 0.0-1.0, \"reasoning\": \"brief explanation\"}"},
            {"role": "user", "content": f"QUESTION: {query}\n\nDOCUMENTATION:\n{docs_context}\n\n"
             "Does this documentation sufficiently answer the question?"}
        ] for _ in range(RESEARCH_CONSENSUS_N)]

        # All 9 sufficiency calls fire in parallel — no token skimping.
        # Each returns a small JSON (~200 tokens) but the model thinks first.
        suff_results = _call_llm_batch(sufficiency_msgs)
        suff_scores = []
        for r in suff_results:
            if "error" not in r:
                suff_scores.append(r.get("confidence", 0.0))
            else:
                suff_scores.append(0.0)

        avg_suff = sum(suff_scores) / len(suff_scores) if suff_scores else 0.0
        print(f"     Sufficiency scores: {[f'{s:.2f}' for s in suff_scores]} → avg: {avg_suff:.2f}")

        if avg_suff < RESEARCH_CONFIDENCE_THRESHOLD:
            print(f"     INSUFFICIENT (avg {avg_suff:.2f} < {RESEARCH_CONFIDENCE_THRESHOLD}). Flagging for manual review.")
            phase_results["queries"][query] = {
                "status": "insufficient",
                "sufficiency_avg": avg_suff,
                "sufficiency_scores": suff_scores,
                "num_chunks": len(above_knee),
            }
            continue

        # Phase B: Extraction (6x concurrent)
        print(f"     Phase B: Extraction ({RESEARCH_CONSENSUS_N}x)...")
        extract_msgs = [[
            {"role": "system", "content": "You are a CUDA kernel engineering expert. "
             "Extract the key technical facts from the documentation to answer the question. "
             "Return JSON: {\"answer\": \"concise answer\", \"key_constraints\": [\"list of constraints\"], "
             "\"code_pattern\": \"relevant API call or code pattern if any\", "
             "\"gotchas\": [\"potential pitfalls\"], \"source_docs\": [\"doc paths used\"]}"},
            {"role": "user", "content": f"QUESTION: {query}\n\nDOCUMENTATION:\n{docs_context}\n\n"
             "Distill the answer to this question into structured JSON."}
        ] for _ in range(RESEARCH_CONSENSUS_N)]

        # BURN TOKENS HERE. 9 parallel extractions, each ~2200 completion tokens.
        # max_tokens=16384 gives 7x headroom for the reasoning model's CoT.
        # If any return "error" with "TRUNCATED", raise max_tokens in _call_llm.
        extract_results = _call_llm_batch(extract_msgs)
        valid_extractions = [r for r in extract_results if "error" not in r and "answer" in r]
        print(f"     Got {len(valid_extractions)}/{RESEARCH_CONSENSUS_N} valid extractions")

        if len(valid_extractions) < 3:
            print(f"     Too few valid extractions. Flagging for manual review.")
            phase_results["queries"][query] = {
                "status": "extraction_failed",
                "sufficiency_avg": avg_suff,
                "valid_extractions": len(valid_extractions),
            }
            continue

        # Phase C: Validation — feed all extractions back, ask for consensus
        print(f"     Phase C: Validation ({RESEARCH_CONSENSUS_N}x)...")
        extractions_text = json.dumps(valid_extractions, indent=2)
        validate_msgs = [[
            {"role": "system", "content": "You are a technical accuracy validator. "
             "Given multiple extraction attempts from the same documentation, determine if they are "
             "consistent and faithfully represent the source material. "
             "Return JSON: {\"valid\": true/false, \"confidence\": 0.0-1.0, "
             "\"consensus_answer\": \"synthesized answer from all extractions\", "
             "\"consensus_constraints\": [\"agreed-upon constraints\"], "
             "\"consensus_code_pattern\": \"agreed code pattern\", "
             "\"contradictions\": [\"any disagreements found\"]}"},
            {"role": "user", "content": f"QUESTION: {query}\n\nEXTRACTIONS:\n{extractions_text}\n\n"
             "Do these extractions consistently and accurately distill the documentation?"}
        ] for _ in range(RESEARCH_CONSENSUS_N)]

        # 9 parallel validation calls. Input is ALL extractions (~9 JSON blobs),
        # so token usage is higher here (~2500+). Let the model think freely.
        val_results = _call_llm_batch(validate_msgs)
        val_scores = []
        consensus_answers = []
        for r in val_results:
            if "error" not in r:
                val_scores.append(r.get("confidence", 0.0))
                if r.get("consensus_answer"):
                    consensus_answers.append(r)

        avg_val = sum(val_scores) / len(val_scores) if val_scores else 0.0
        print(f"     Validation scores: {[f'{s:.2f}' for s in val_scores]} → avg: {avg_val:.2f}")

        if avg_val >= RESEARCH_CONFIDENCE_THRESHOLD and consensus_answers:
            # Pick the consensus with highest confidence
            best = max(consensus_answers, key=lambda x: x.get("confidence", 0))
            phase_results["queries"][query] = {
                "status": "accepted",
                "sufficiency_avg": avg_suff,
                "validation_avg": avg_val,
                "answer": best.get("consensus_answer", ""),
                "key_constraints": best.get("consensus_constraints", []),
                "code_pattern": best.get("consensus_code_pattern", ""),
                "contradictions": best.get("contradictions", []),
                "num_chunks": len(above_knee),
                "num_extractions": len(valid_extractions),
            }
            print(f"     ACCEPTED: {best.get('consensus_answer', '')[:100]}...")
        else:
            # Log disagreement details so we can diagnose what diverged
            extraction_summaries = [
                e.get("answer", "")[:150] for e in valid_extractions
            ]
            contradiction_details = []
            for r in val_results:
                if "error" not in r and r.get("contradictions"):
                    contradiction_details.extend(r["contradictions"])
            phase_results["queries"][query] = {
                "status": "validation_failed",
                "sufficiency_avg": avg_suff,
                "validation_avg": avg_val,
                "validation_scores": val_scores,
                "num_extractions": len(valid_extractions),
                "extraction_summaries": extraction_summaries,
                "contradictions_found": list(set(contradiction_details))[:10],
            }
            print(f"     VALIDATION FAILED (avg {avg_val:.2f} < {RESEARCH_CONFIDENCE_THRESHOLD}). Manual review needed.")
            if contradiction_details:
                print(f"     Contradictions: {contradiction_details[0][:120]}...")

        # Incremental save after each query — survives crashes, concurrent-safe
        phase_results["overall_confidence"] = (
            sum(1 for q in phase_results["queries"].values() if q.get("status") == "accepted")
            / len(phase_results["queries"])
        )
        _save_research({"phases": {phase_key: phase_results}})

    # Final summary
    query_results = phase_results["queries"]
    accepted = sum(1 for q in query_results.values() if q.get("status") == "accepted")
    total = len(query_results)
    phase_results["overall_confidence"] = accepted / total if total > 0 else 0.0

    print(f"\n  {'─' * 70}")
    print(f"  RESEARCH COMPLETE: {accepted}/{total} queries accepted "
          f"(overall confidence: {phase_results['overall_confidence']:.1%})")

    # Final persist with computed confidence
    _save_research({"phases": {phase_key: phase_results}})

    return phase_results


def check_phase_complete(phase_num: int, result: PhaseResult) -> Tuple[bool, List[str]]:
    """
    Check if a phase meets completion criteria.
    Returns (is_complete, list_of_issues).
    """
    reqs = PHASE_REQUIREMENTS.get(phase_num, {})
    required = reqs.get("required_measurements", [])
    acceptance = reqs.get("acceptance", {})
    issues = []

    # Check required measurements exist
    for key in required:
        if key not in result.measurements:
            issues.append(f"Missing measurement: {key}")

    # Check acceptance thresholds
    for key, (op, threshold) in acceptance.items():
        if key in result.measurements:
            val = result.measurements[key]
            if op == "==" and val != threshold:
                issues.append(f"{key} = {val}, need == {threshold}")
            elif op == "<=" and val > threshold:
                issues.append(f"{key} = {val}, need <= {threshold}")
            elif op == ">=" and val < threshold:
                issues.append(f"{key} = {val}, need >= {threshold}")

    return len(issues) == 0, issues


def check_regression(phase_num: int, key: str, new_value: float,
                     results: Dict[int, PhaseResult]) -> Optional[str]:
    """
    Check if a new measurement regresses a locked baseline.
    Returns warning string or None.
    """
    # Check all completed phases for baselines with this key
    for pnum, result in results.items():
        if result.status == "done" and key in result.baselines:
            baseline = result.baselines[key]
            # Detect regression: time went up or efficiency went down
            if "time" in key or "overhead" in key:
                if new_value > baseline * 1.05:  # 5% tolerance
                    return (f"REGRESSION: {key} = {new_value:.2f} "
                            f"(baseline {baseline:.2f} from phase {pnum}, +{(new_value/baseline - 1)*100:.1f}%)")
            elif "eff" in key or "occupancy" in key:
                if new_value < baseline * 0.95:
                    return (f"REGRESSION: {key} = {new_value:.2f} "
                            f"(baseline {baseline:.2f} from phase {pnum}, {(new_value/baseline - 1)*100:.1f}%)")
    return None


def find_current_phase(results: Dict[int, PhaseResult]) -> int:
    """Find the lowest phase number that is not 'done'."""
    for i in range(8):
        r = results.get(i)
        if r is None or r.status != "done":
            return i
    return 8  # All done


def check_prerequisites_met(phase_num: int, results: Dict[int, PhaseResult]) -> Tuple[bool, str]:
    """Check if the prerequisite phase is complete."""
    if phase_num == 0:
        return True, ""
    prev = results.get(phase_num - 1)
    if prev is None or prev.status != "done":
        return False, f"Phase {phase_num - 1} is not complete (status: {prev.status if prev else 'not started'})"
    return True, ""


# =============================================================================
# Command: next
# =============================================================================

def cmd_next(config_path: str, gpu_key: str):
    """Print the single next action to take."""
    results = load_results()
    current = find_current_phase(results)

    if current >= 8:
        print("\n  ALL PHASES COMPLETE. Forward pass optimization is done.\n")
        return

    # Load model to get phases
    hw = GPU_SPECS[gpu_key]
    config = parse_model_config(config_path)
    params = extract_architecture_parameters(config)
    weights = enumerate_all_weight_matrices(params)
    shapes = aggregate_gemv_shapes(weights)
    models = compute_full_forward_pass_model(shapes, hw)
    phases = generate_optimization_roadmap(models, hw, params)

    phase = phases[current]
    result = ensure_phase_exists(results, current)

    prereq_met, prereq_issue = check_prerequisites_met(current, results)

    print(f"\n  ╔══════════════════════════════════════════════════════════════╗")
    print(f"  ║  NEXT ACTION: Phase {current}                                     ║")
    print(f"  ╚══════════════════════════════════════════════════════════════╝")

    if not prereq_met:
        print(f"\n  BLOCKED: {prereq_issue}")
        print(f"  Complete phase {current - 1} first.\n")
        return

    print(f"\n  Phase: {phase.name}")
    print(f"  Status: {result.status}")

    if result.status == "pending":
        print(f"\n  ACTION: Start building.")
        print(f"\n  WHAT TO BUILD:")
        for line in _wrap_text(phase.what_to_build, width=64):
            print(f"    {line}")

        if phase.expected_time_saved_us > 0:
            print(f"\n  Expected improvement: {phase.expected_time_saved_us:.0f} μs saved")
        if phase.expected_tps_after > 0:
            print(f"  Expected TPS after: {phase.expected_tps_after:.0f}")

        # Auto-transition to in_progress
        result.status = "in_progress"
        result.started = _now()
        save_results(results)
        print(f"\n  [Phase {current} marked as in_progress]")

    elif result.status == "in_progress":
        # Check what measurements are still needed
        reqs = PHASE_REQUIREMENTS.get(current, {})
        required = reqs.get("required_measurements", [])
        missing = [k for k in required if k not in result.measurements]

        if missing:
            print(f"\n  ACTION: Run benchmarks and record measurements.")
            print(f"\n  MEASUREMENT COMMAND:")
            print(f"    {phase.measurement_command}")
            print(f"\n  MISSING MEASUREMENTS ({len(missing)}):")
            for m in missing:
                print(f"    - {m}")
            print(f"\n  Record with:")
            for m in missing:
                print(f"    python3 tools/moe_mamba_blackwell_primer.py record {current} {m} <value>")
        else:
            # All measurements present — check acceptance
            is_complete, issues = check_phase_complete(current, result)
            if is_complete:
                print(f"\n  ACTION: All criteria met. Locking baselines.")
                result.status = "done"
                result.completed = _now()
                result.baselines = dict(result.measurements)  # Lock
                save_results(results)
                print(f"  [Phase {current} COMPLETE]")
                next_phase = current + 1
                if next_phase < len(phases):
                    print(f"\n  Next up: Phase {next_phase}: {phases[next_phase].name}")
            else:
                print(f"\n  ACTION: Fix these issues before phase can complete:")
                for issue in issues:
                    print(f"    - {issue}")

    print()


# =============================================================================
# Command: status
# =============================================================================

def cmd_status(config_path: str, gpu_key: str):
    """Print phase overview."""
    results = load_results()
    hw = GPU_SPECS[gpu_key]
    config = parse_model_config(config_path)
    params = extract_architecture_parameters(config)
    weights = enumerate_all_weight_matrices(params)
    shapes = aggregate_gemv_shapes(weights)
    models = compute_full_forward_pass_model(shapes, hw)
    phases = generate_optimization_roadmap(models, hw, params)

    current = find_current_phase(results)

    print(f"\n  {'═' * 70}")
    print(f"  OPTIMIZATION STATUS — {hw.name}")
    print(f"  {'═' * 70}")
    print(f"\n  {'Phase':<8s} {'Status':<14s} {'Prereq':<10s} {'Measurements':<16s} {'Name'}")
    print(f"  {'─' * 8} {'─' * 14} {'─' * 10} {'─' * 16} {'─' * 30}")

    for phase in phases:
        pn = phase.phase_number
        result = results.get(pn)
        status = result.status if result else "pending"

        prereq_met, _ = check_prerequisites_met(pn, results)
        prereq_str = "OK" if prereq_met else "BLOCKED"

        reqs = PHASE_REQUIREMENTS.get(pn, {})
        required = reqs.get("required_measurements", [])
        if result:
            present = sum(1 for k in required if k in result.measurements)
        else:
            present = 0
        meas_str = f"{present}/{len(required)}"

        marker = ">>>" if pn == current else "   "
        status_display = {
            "pending": "pending",
            "in_progress": "IN PROGRESS",
            "done": "DONE",
            "blocked": "BLOCKED",
        }.get(status, status)

        print(f"  {marker} {pn:<4d} {status_display:<14s} {prereq_str:<10s} {meas_str:<16s} {phase.name}")

    # Summary
    done_count = sum(1 for r in results.values() if r.status == "done")
    total = len(phases)
    print(f"\n  Progress: {done_count}/{total} phases complete")

    # Show current phase detail
    if current < len(phases):
        result = results.get(current)
        if result and result.measurements:
            print(f"\n  Phase {current} measurements:")
            for k, v in sorted(result.measurements.items()):
                print(f"    {k}: {v}")

    print()


# =============================================================================
# Command: record
# =============================================================================

def cmd_record(phase_num: int, key: str, value: float):
    """Record a measurement for a phase."""
    results = load_results()
    result = ensure_phase_exists(results, phase_num)

    # Check for regression
    warning = check_regression(phase_num, key, value, results)
    if warning:
        print(f"\n  WARNING: {warning}")
        print(f"  Recording anyway. Investigate before proceeding.\n")

    old_value = result.measurements.get(key)
    result.measurements[key] = value

    if result.status == "pending":
        result.status = "in_progress"
        result.started = _now()

    save_results(results)

    if old_value is not None:
        print(f"  Phase {phase_num}: {key} = {value} (was {old_value})")
    else:
        print(f"  Phase {phase_num}: {key} = {value} (new)")

    # Check if phase is now complete
    is_complete, issues = check_phase_complete(phase_num, result)
    if is_complete and result.status == "in_progress":
        print(f"\n  All criteria met for phase {phase_num}!")
        print(f"  Run 'primer next' to lock baselines and advance.")
    elif issues:
        reqs = PHASE_REQUIREMENTS.get(phase_num, {})
        required = reqs.get("required_measurements", [])
        missing = [k for k in required if k not in result.measurements]
        if missing:
            print(f"  Still missing: {', '.join(missing)}")


# =============================================================================
# Command: history
# =============================================================================

def cmd_history(phase_filter=None):
    """Show all recorded measurements, optionally filtered to one phase."""
    results = load_results()
    if not results:
        print("\n  No measurements recorded yet.\n")
        return

    print(f"\n  {'═' * 70}")
    print(f"  MEASUREMENT HISTORY")
    print(f"  {'═' * 70}")

    phases = sorted(results.keys())
    if phase_filter is not None:
        phases = [p for p in phases if p == phase_filter]
        if not phases:
            print(f"\n  No measurements for phase {phase_filter}.\n")
            return

    for pn in phases:
        result = results[pn]
        print(f"\n  Phase {pn} [{result.status}]")
        if result.started:
            print(f"    Started:   {result.started}")
        if result.completed:
            print(f"    Completed: {result.completed}")

        if result.measurements:
            print(f"    Measurements:")
            for k, v in sorted(result.measurements.items()):
                baseline_marker = ""
                if k in result.baselines:
                    baseline_marker = " (LOCKED)"
                print(f"      {k}: {v}{baseline_marker}")

        if result.notes:
            print(f"    Notes:")
            for note in result.notes:
                print(f"      - {note}")

    print()


# =============================================================================
# Command: complete (manually mark a phase done)
# =============================================================================

def cmd_complete(phase_num: int):
    """Manually mark a phase as complete, locking baselines."""
    results = load_results()
    result = ensure_phase_exists(results, phase_num)

    if result.status == "done":
        print(f"  Phase {phase_num} is already done.")
        return

    result.status = "done"
    result.completed = _now()
    result.baselines = dict(result.measurements)
    save_results(results)
    print(f"  Phase {phase_num} marked COMPLETE. Baselines locked.")


# =============================================================================
# Command: note
# =============================================================================

def cmd_note(phase_num: int, note_text: str):
    """Add a note to a phase (what was tried, what didn't work, etc.)."""
    results = load_results()
    result = ensure_phase_exists(results, phase_num)
    result.notes.append(f"[{_now()}] {note_text}")
    save_results(results)
    print(f"  Note added to phase {phase_num}.")


# =============================================================================
# Command: analyze (the original full analysis)
# =============================================================================

# =============================================================================
# Command: research
# =============================================================================

def cmd_research(phase_num: int, force: bool = False):
    """Run the research pipeline for a phase."""
    run_research_pipeline(phase_num, force=force)


def cmd_research_status():
    """Show research status for all phases."""
    data = _load_research()
    phases = data.get("phases", {})

    print(f"\n  {'═' * 70}")
    print(f"  RESEARCH STATUS")
    print(f"  {'═' * 70}")

    for pn in range(8):
        research_def = PHASE_RESEARCH.get(pn, {})
        phase_data = phases.get(str(pn))

        desc = research_def.get("description", "—")
        num_queries = len(research_def.get("queries", []))

        if phase_data:
            accepted = sum(1 for q in phase_data.get("queries", {}).values()
                          if q.get("status") == "accepted")
            total = len(phase_data.get("queries", {}))
            conf = phase_data.get("overall_confidence", 0)
            status = f"DONE ({accepted}/{total} accepted, {conf:.0%})"
        else:
            status = "pending"

        print(f"  Phase {pn}: {status:<40s} {desc}")
        if phase_data:
            for query, result in phase_data.get("queries", {}).items():
                s = result.get("status", "?")
                marker = "✓" if s == "accepted" else "✗" if "failed" in s else "?"
                print(f"    {marker} {query[:60]}")

    print()


# =============================================================================
# Command: unknowns
# =============================================================================

def cmd_unknowns():
    """Display all open questions, resolved findings, and design decisions.

    PSEUDOCODE:
    1. Group OPEN_QUESTIONS by status (resolved, unresolved, decision_needed)
    2. For each resolved question: show question, resolution, source
    3. For each unresolved question: show question, notes, where to look
    4. For each design decision: show question, options, recommendation
    5. Print summary counts
    """
    resolved = {k: v for k, v in OPEN_QUESTIONS.items() if v["status"] == "resolved"}
    unresolved = {k: v for k, v in OPEN_QUESTIONS.items() if v["status"] == "unresolved"}
    decisions = {k: v for k, v in OPEN_QUESTIONS.items() if v["status"] == "decision_needed"}

    print(f"\n  {'═' * 74}")
    print(f"  OPEN QUESTIONS & DESIGN DECISIONS")
    print(f"  {'═' * 74}")
    print(f"\n  Resolved: {len(resolved)}  |  Unresolved: {len(unresolved)}  |  Decisions: {len(decisions)}")

    # --- Resolved ---
    if resolved:
        print(f"\n  {'─' * 74}")
        print(f"  RESOLVED (via nvidia-docs-search)")
        print(f"  {'─' * 74}")
        for qid, q in resolved.items():
            print(f"\n  [{qid}] {q['question']}")
            for line in _wrap_text(q['resolution'], width=70):
                print(f"    {line}")
            print(f"    Source: {q['source']}")
            print(f"    Impacts: {', '.join(q['impact'])}")

    # --- Unresolved ---
    if unresolved:
        print(f"\n  {'─' * 74}")
        print(f"  UNRESOLVED (need external sources)")
        print(f"  {'─' * 74}")
        for qid, q in unresolved.items():
            print(f"\n  [{qid}] {q['question']}")
            print(f"    Where to look: {q['source']}")
            if q.get("notes"):
                for line in _wrap_text(q['notes'], width=70):
                    print(f"    {line}")
            print(f"    Impacts: {', '.join(q['impact'])}")

    # --- Design Decisions ---
    if decisions:
        print(f"\n  {'─' * 74}")
        print(f"  DESIGN DECISIONS NEEDED")
        print(f"  {'─' * 74}")
        for qid, q in decisions.items():
            print(f"\n  [{qid}] {q['question']}")
            for opt_name, opt_desc in q.get("options", {}).items():
                print(f"    Option {opt_name}:")
                for line in _wrap_text(opt_desc, width=66):
                    print(f"      {line}")
            if q.get("recommendation"):
                print(f"    RECOMMENDATION:")
                for line in _wrap_text(q['recommendation'], width=66):
                    print(f"      {line}")
            print(f"    Impacts: {', '.join(q['impact'])}")

    print()


# =============================================================================
# Blackwell Primer: Verification Engine
# =============================================================================
#
# DESIGN: This is the programmatic audit layer of the Blackwell primer.
# It is bottom-up verification: hardware → build → kernels → runtime → model.
#
# Purpose:
#   1. ANCHOR — after each code change, re-run to confirm no regressions
#   2. GUIDE — when returning to this project, run verify to see current state
#   3. AUDIT — every check traces to a specific file, line, or hardware query
#
# Each check is deterministic: it reads source files, queries hardware, and
# compares against known-good values. No LLM calls, no heuristics.
#
# Layer hierarchy (bottom-up):
#   Layer 1: HARDWARE    — live GPU query (nvidia-smi, SSH) + static specs
#   Layer 2: BUILD       — CMakeLists.txt: ENABLE_FP4, CUDA version, arch targets
#   Layer 3: KERNEL      — __CUDA_ARCH__ guards: equality, range, exclusion
#   Layer 4: FP4_KV      — MMHA: ENABLE_FP4 ifdefs, nvfp4_kv_cache dispatch, Dh=256
#   Layer 5: RUNTIME     — Python→C++ config bridge: how nvfp4_kv_cache gets set
#   Layer 6: MODEL       — Qwen3-Next-80B: head_dim, Mamba hybrid, MoE routing
#   Layer 7: DOCS        — cross-reference with local documentation library
#
# Reference documentation paths (relative to repo root):
#   docs/reference/cuda/pdf/ptx_isa_9.1.pdf          — PTX ISA for block_scale MMA
#   docs/reference/cuda/pdf/Blackwell_Tuning_Guide.pdf — SM120 optimization
#   docs/reference/gpu/whitepapers/nvidia-rtx-blackwell-gpu-architecture.pdf
#   docs/reference/dgx/dgx-spark/                     — SM121 DGX Spark specs
#   docs/reference/blog/blog/introducing-nvfp4-*/     — NVFP4 format details
#
# Fork source paths (relative to fork root):
#   cpp/kernels/xqa/utils.cuh                          — XQA SMEM tuning constants
#   cpp/kernels/xqa/mha.cu                             — XQA head partition sizes
#   cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/  — MMHA FP4 path
#   cpp/CMakeLists.txt                                 — ENABLE_FP4 compile gate
#   tensorrt_llm/                                      — Python config layer
#
# PSEUDOCODE for future layers (not yet implemented):
#   Layer 8: BUILD_EXECUTION — invoke cmake on target hardware, verify compilation
#     1. SSH to target machine (or local)
#     2. Check CUDA toolkit version: nvcc --version
#     3. Check cmake version
#     4. Run cmake configure with -DCMAKE_CUDA_ARCHITECTURES={arch}a-real
#     5. Build minimum target (MMHA kernels only for fast iteration)
#     6. Report build success/failure with specific error lines
#
#   Layer 9: DEPLOY_EXECUTION — create deployment directory, start inference
#     1. Create SM{tier} deployment directory per /llm-deploy conventions
#     2. Write spec.yaml with kv_cache_config.dtype: "fp4"
#     3. Generate docker-compose.yml or native launch script
#     4. Deploy via ./deploy <dir> or native execution
#     5. Health check: curl /health + /v1/models
#     6. Smoke test: single inference request
#
#   Layer 10: BENCHMARK — measure FP4 KV vs FP8 KV performance
#     1. Single-stream throughput (tok/s) with FP4 KV cache
#     2. Compare against FP8 KV baseline (27.6 TPS on SM121 TRT-LLM)
#     3. Memory usage comparison (model + KV cache)
#     4. Accuracy spot check (known prompt → known output token pattern)
# =============================================================================

@dataclass
class VerifyCheck:
    """
    Result of a single verification check.

    Each check has:
      - name: machine-readable identifier (e.g. "xqa_guard_blocking")
      - layer: which verification layer produced it (e.g. "kernel")
      - status: PASS (confirmed good), FAIL (blocks progress), WARN (review),
                INFO (context only, no judgment)
      - detail: human-readable explanation with specific numbers/paths
      - file_path: source file(s) where evidence was found
      - line_numbers: specific line numbers for direct navigation
      - doc_ref: path to documentation that validates this check
    """
    name: str
    layer: str          # "hardware", "build", "kernel", "fp4_kv", "runtime", "model", "docs"
    status: str         # "PASS", "FAIL", "WARN", "INFO"
    detail: str
    file_path: str = ""
    line_numbers: List[int] = field(default_factory=list)
    doc_ref: str = ""   # path to supporting documentation


def _scan_file_for_pattern(filepath: Path, pattern: str) -> List[Tuple[int, str]]:
    """Return (line_number, line_text) for all lines matching regex pattern."""
    matches = []
    try:
        text = filepath.read_text(errors="replace")
        for i, line in enumerate(text.split("\n"), 1):
            if re.search(pattern, line):
                matches.append((i, line.strip()))
    except FileNotFoundError:
        pass
    return matches


def _scan_dir_for_pattern(dirpath: Path, pattern: str, glob: str = "**/*.cu*") -> List[Tuple[str, int, str]]:
    """Return (relative_path, line_number, line_text) for all matches in directory."""
    results = []
    for f in sorted(dirpath.glob(glob)):
        if f.is_file():
            for lnum, ltxt in _scan_file_for_pattern(f, pattern):
                results.append((str(f.relative_to(dirpath)), lnum, ltxt))
    # Also check .h/.cuh files
    for ext in ("**/*.h", "**/*.cuh"):
        for f in sorted(dirpath.glob(ext)):
            if f.is_file():
                for lnum, ltxt in _scan_file_for_pattern(f, pattern):
                    rel = str(f.relative_to(dirpath))
                    if not any(r[0] == rel and r[1] == lnum for r in results):
                        results.append((rel, lnum, ltxt))
    return results


def _query_live_gpu(target_arch: str) -> Optional[Dict[str, str]]:
    """
    Query live GPU hardware via nvidia-smi or SSH.

    For SM120 (local): nvidia-smi directly
    For SM121 (remote): SSH to edgexpert-user

    Returns dict with keys: gpu_name, driver_version, cuda_version,
    memory_total, memory_used, sm_count, compute_capability.
    Returns None if hardware is unreachable.
    """
    # Build the nvidia-smi command
    smi_cmd = (
        "nvidia-smi --query-gpu=name,driver_version,memory.total,memory.used,"
        "compute_cap --format=csv,noheader,nounits 2>/dev/null"
    )

    if target_arch == "sm121":
        # Remote: SSH to DGX Spark
        cmd = f"ssh -o ConnectTimeout=5 edgexpert-user '{smi_cmd}'"
    else:
        # Local GPU query
        cmd = smi_cmd

    try:
        result = subprocess.run(
            cmd, shell=True, capture_output=True, text=True, timeout=10
        )
        if result.returncode != 0 or not result.stdout.strip():
            return None

        # Parse CSV output (may have multiple GPUs — take first for SM121, filter for SM120)
        lines = [l.strip() for l in result.stdout.strip().split("\n") if l.strip()]
        if not lines:
            return None

        # For SM120: find GPU index 0 or 1 (RTX Pro 6000)
        # For SM121: take first (only GPU)
        parts = [p.strip() for p in lines[0].split(",")]
        if len(parts) >= 5:
            return {
                "gpu_name": parts[0],
                "driver_version": parts[1],
                "memory_total_mb": parts[2],
                "memory_used_mb": parts[3],
                "compute_capability": parts[4],
            }
    except (subprocess.TimeoutExpired, FileNotFoundError):
        pass
    return None


def _query_live_cuda_version(target_arch: str) -> Optional[str]:
    """Query CUDA toolkit version on target hardware."""
    if target_arch == "sm121":
        cmd = "ssh -o ConnectTimeout=5 edgexpert-user 'nvcc --version 2>/dev/null || echo NO_NVCC'"
    else:
        cmd = "nvcc --version 2>/dev/null || echo NO_NVCC"

    try:
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=10)
        if "NO_NVCC" in result.stdout:
            return None
        # Parse: "Cuda compilation tools, release 13.0, V13.0.76"
        m = re.search(r"release\s+([\d.]+)", result.stdout)
        return m.group(1) if m else None
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return None


def verify_trtllm_fork(fork_root: Path, target_arch: str = "sm121") -> List[VerifyCheck]:
    """
    Programmatic audit of TRT-LLM fork for SM120/SM121 NVFP4 KV cache readiness.

    PSEUDOCODE:
      1. Query live hardware (nvidia-smi, SSH) — confirm GPU is reachable
      2. Cross-reference static HardwareSpec with live values
      3. Parse CMakeLists.txt for ENABLE_FP4 gate and CUDA version requirement
      4. Scan all kernel source for __CUDA_ARCH__ guards that block target
      5. Verify MMHA FP4 KV cache code path: ifdefs, dispatch, types, Dh=256
      6. Trace Python config layer for nvfp4_kv_cache parameter propagation
      7. Check model-specific compatibility (head_dim, Mamba, MoE)
      8. Cross-reference with local documentation library

    Each check is deterministic: file reads, regex matches, process output.
    """
    checks = []
    cpp_root = fork_root / "cpp"
    repo_root = fork_root.parent.parent  # SM12x.../TensorRT-LLM -> SM12x... -> repo root
    arch_num = 1200 if target_arch == "sm120" else 1210

    # =========================================================================
    # Layer 1: Hardware (live query + static specs)
    # =========================================================================
    hw = SM121_DGX_SPARK if target_arch == "sm121" else SM120_RTX_PRO_6000

    # 1a. Live GPU query
    live_gpu = _query_live_gpu(target_arch)
    if live_gpu:
        checks.append(VerifyCheck(
            name="hardware_live_gpu",
            layer="hardware",
            status="PASS",
            detail=f"Live GPU: {live_gpu['gpu_name']}, "
                   f"Driver {live_gpu['driver_version']}, "
                   f"CC {live_gpu['compute_capability']}, "
                   f"VRAM {live_gpu['memory_total_mb']} MB "
                   f"(used: {live_gpu['memory_used_mb']} MB)",
        ))
        # Cross-reference compute capability
        expected_cc = "12.0" if target_arch == "sm120" else "12.1"
        cc = live_gpu["compute_capability"]
        checks.append(VerifyCheck(
            name="hardware_compute_capability",
            layer="hardware",
            status="PASS" if cc == expected_cc else "WARN",
            detail=f"Compute capability: {cc} (expected {expected_cc} for {target_arch}). "
                   + ("Match." if cc == expected_cc else f"Mismatch — verify target."),
            doc_ref="docs/reference/cuda/pdf/Blackwell_Compatibility_Guide.pdf",
        ))
    else:
        checks.append(VerifyCheck(
            name="hardware_live_gpu",
            layer="hardware",
            status="WARN",
            detail=f"GPU unreachable for live query "
                   f"({'SSH to edgexpert-user' if target_arch == 'sm121' else 'local nvidia-smi'}). "
                   f"Using static specs only.",
        ))

    # 1b. Live CUDA toolkit version
    live_cuda = _query_live_cuda_version(target_arch)
    if live_cuda:
        checks.append(VerifyCheck(
            name="hardware_cuda_toolkit",
            layer="hardware",
            status="PASS" if float(live_cuda.split(".")[0]) >= 12 else "FAIL",
            detail=f"CUDA toolkit: {live_cuda} "
                   f"(ENABLE_FP4 requires >= 12.8).",
        ))
    else:
        # For DGX Spark, nvcc is only in Docker — note this
        checks.append(VerifyCheck(
            name="hardware_cuda_toolkit",
            layer="hardware",
            status="INFO" if target_arch == "sm121" else "WARN",
            detail="No native CUDA toolkit found. "
                   + ("DGX Spark: nvcc available inside Docker images only (Triton 25.04, TRT-LLM 1.2.0rc6)."
                      if target_arch == "sm121"
                      else "Expected nvcc in PATH for local builds."),
        ))

    # 1c. Static hardware spec
    checks.append(VerifyCheck(
        name="hardware_static_spec",
        layer="hardware",
        status="INFO",
        detail=f"{hw.name}: {hw.sm_count} SMs, {hw.memory_bandwidth_gb_s} GB/s, "
               f"{hw.memory_total_gb} GB, SMEM {hw.smem_per_sm_bytes}B/SM",
        doc_ref=(
            "docs/reference/dgx/dgx-spark/" if target_arch == "sm121"
            else "docs/reference/gpu/datasheets/"
        ),
    ))
    checks.append(VerifyCheck(
        name="hardware_smem_for_xqa",
        layer="hardware",
        status="PASS" if hw.smem_per_sm_bytes >= 99 * 1024 else "FAIL",
        detail=f"XQA needs kMAX_SMEM_SIZE = 99KB (101,376B). "
               f"Hardware has {hw.smem_per_sm_bytes}B ({hw.smem_per_sm_bytes / 1024:.0f}KB). "
               f"Verified via CUDA driver API (cuDeviceGetAttribute), not marketing specs.",
        doc_ref="docs/reference/cuda/pdf/Blackwell_Tuning_Guide.pdf",
    ))

    # =========================================================================
    # Layer 2: Build system
    # =========================================================================
    cmake_file = cpp_root / "CMakeLists.txt"
    if cmake_file.exists():
        cmake_text = cmake_file.read_text()

        # Check ENABLE_FP4 gate
        # This is the critical compile-time flag that enables FP4 KV cache code paths.
        # It must be gated by CUDA toolkit version (>= 12.8), NOT by architecture.
        fp4_match = re.search(
            r'if\s*\(\$\{CUDAToolkit_VERSION\}\s+VERSION_GREATER_EQUAL\s+"(\d+\.\d+)"\)\s*\n\s*add_definitions\("-DENABLE_FP4"\)',
            cmake_text
        )
        if fp4_match:
            cuda_min = fp4_match.group(1)
            checks.append(VerifyCheck(
                name="enable_fp4_gate",
                layer="build",
                status="PASS",
                detail=f"ENABLE_FP4 gated by CUDA >= {cuda_min} (architecture-agnostic). "
                       f"DGX Spark has CUDA 13.0 → will be defined.",
                file_path=str(cmake_file.relative_to(fork_root)),
                doc_ref="docs/reference/cuda/pdf/ptx_isa_9.1.pdf (section: block_scale MMA)",
            ))
        else:
            # Try broader search
            fp4_lines = _scan_file_for_pattern(cmake_file, r"ENABLE_FP4")
            checks.append(VerifyCheck(
                name="enable_fp4_gate",
                layer="build",
                status="WARN" if fp4_lines else "FAIL",
                detail=f"ENABLE_FP4 pattern not matched. Found {len(fp4_lines)} references."
                       + (f" Lines: {[l[0] for l in fp4_lines]}" if fp4_lines else ""),
                file_path=str(cmake_file.relative_to(fork_root)),
            ))

        # Check if sm_121a is in any hardcoded arch list
        arch_refs = _scan_file_for_pattern(cmake_file, r"12[01]a?")
        checks.append(VerifyCheck(
            name="cmake_arch_references",
            layer="build",
            status="INFO",
            detail=f"CMAKE_CUDA_ARCHITECTURES set externally (not hardcoded). "
                   f"Build with -DCMAKE_CUDA_ARCHITECTURES={target_arch.replace('sm', '')}a-real. "
                   f"Found {len(arch_refs)} references to 120/121 in CMakeLists.txt.",
            file_path=str(cmake_file.relative_to(fork_root)),
        ))
    else:
        checks.append(VerifyCheck(
            name="cmake_missing",
            layer="build",
            status="FAIL",
            detail=f"CMakeLists.txt not found at {cmake_file}",
        ))

    # =========================================================================
    # Layer 3: Kernel architecture guards
    # =========================================================================

    # 3a. Hard equality checks: __CUDA_ARCH__ == 1200 (must include target)
    kernels_dir = cpp_root / "kernels"
    tllm_kernels_dir = cpp_root / "tensorrt_llm" / "kernels"

    hard_eq_pattern = r"__CUDA_ARCH__\s*==\s*1200"
    all_hard_eq = []
    for search_dir in [kernels_dir, tllm_kernels_dir]:
        if search_dir.exists():
            all_hard_eq.extend(
                (str(search_dir.relative_to(fork_root)) + "/" + r[0], r[1], r[2])
                for r in _scan_dir_for_pattern(search_dir, hard_eq_pattern)
            )

    # Check if target arch is included alongside 1200
    blocking_guards = []
    patched_guards = []
    for fpath, lnum, ltxt in all_hard_eq:
        if str(arch_num) in ltxt:
            patched_guards.append((fpath, lnum, ltxt))
        else:
            blocking_guards.append((fpath, lnum, ltxt))

    if blocking_guards:
        checks.append(VerifyCheck(
            name="xqa_guard_blocking",
            layer="kernel",
            status="FAIL",
            detail=f"{len(blocking_guards)} guard(s) check == 1200 WITHOUT {arch_num}. "
                   f"SM{target_arch[-3:]} will hit #error or get wrong constants.",
            file_path="; ".join(f"{g[0]}:{g[1]}" for g in blocking_guards),
            line_numbers=[g[1] for g in blocking_guards],
        ))
    else:
        checks.append(VerifyCheck(
            name="xqa_guard_blocking",
            layer="kernel",
            status="PASS",
            detail=f"No guards check == 1200 without also including {arch_num}. "
                   f"{len(patched_guards)} patched guard(s) found.",
            file_path="; ".join(f"{g[0]}:{g[1]}" for g in patched_guards) if patched_guards else "",
        ))

    # 3b. Range checks: __CUDA_ARCH__ >= 900 (SM121 passes)
    range_pattern = r"__CUDA_ARCH__\s*>=\s*(\d+)"
    range_checks = []
    for search_dir in [kernels_dir, tllm_kernels_dir]:
        if search_dir.exists():
            for rel, lnum, ltxt in _scan_dir_for_pattern(search_dir, range_pattern):
                for m in re.finditer(r"__CUDA_ARCH__\s*>=\s*(\d+)", ltxt):
                    threshold = int(m.group(1))
                    range_checks.append((rel, lnum, threshold, ltxt))

    # Group by threshold
    thresholds = {}
    for _, _, t, _ in range_checks:
        thresholds[t] = thresholds.get(t, 0) + 1

    passing = {t: c for t, c in thresholds.items() if arch_num >= t}
    failing = {t: c for t, c in thresholds.items() if arch_num < t}

    checks.append(VerifyCheck(
        name="range_guards_passing",
        layer="kernel",
        status="PASS",
        detail=f"SM{target_arch[-3:]} (arch {arch_num}) passes {sum(passing.values())} range guards. "
               f"Thresholds: {dict(sorted(passing.items()))}",
    ))

    if failing:
        checks.append(VerifyCheck(
            name="range_guards_failing",
            layer="kernel",
            status="WARN",
            detail=f"SM{target_arch[-3:]} fails {sum(failing.values())} range guards "
                   f"(thresholds: {dict(sorted(failing.items()))}). "
                   f"These may be intentional (NVLink, SM100 TMEM features).",
        ))

    # 3c. Exclusion ranges: >= 900 && < 1200 (intentional SM120+ exclusion)
    excl_pattern = r"__CUDA_ARCH__\s*>=\s*900.*<\s*1200|__CUDA_ARCH__\s*<\s*1200.*>=\s*900"
    excl_guards = []
    for search_dir in [kernels_dir, tllm_kernels_dir]:
        if search_dir.exists():
            excl_guards.extend(
                (str(search_dir.relative_to(fork_root)) + "/" + r[0], r[1], r[2])
                for r in _scan_dir_for_pattern(search_dir, excl_pattern)
            )
    checks.append(VerifyCheck(
        name="nvlink_exclusion_guards",
        layer="kernel",
        status="INFO",
        detail=f"{len(excl_guards)} guard(s) use >= 900 && < 1200 pattern "
               f"(NVLink cooperative launch — intentionally excludes SM120+). "
               f"Not a blocker for single-GPU inference.",
    ))

    # =========================================================================
    # Layer 4: FP4 KV cache path
    # =========================================================================
    mmha_dir = tllm_kernels_dir / "decoderMaskedMultiheadAttention"

    if mmha_dir.exists():
        # 4a. Count ENABLE_FP4 ifdefs in MMHA
        fp4_ifdefs = _scan_dir_for_pattern(mmha_dir, r"#ifdef\s+ENABLE_FP4")
        checks.append(VerifyCheck(
            name="mmha_enable_fp4_blocks",
            layer="fp4_kv",
            status="PASS" if len(fp4_ifdefs) >= 5 else "WARN",
            detail=f"{len(fp4_ifdefs)} #ifdef ENABLE_FP4 blocks in MMHA code "
                   f"(template + launch header). Expected >= 8.",
        ))

        # 4b. Runtime dispatch: params.nvfp4_kv_cache
        nvfp4_dispatch = _scan_dir_for_pattern(mmha_dir, r"nvfp4_kv_cache")
        checks.append(VerifyCheck(
            name="mmha_runtime_dispatch",
            layer="fp4_kv",
            status="PASS" if nvfp4_dispatch else "FAIL",
            detail=f"{len(nvfp4_dispatch)} reference(s) to nvfp4_kv_cache in MMHA. "
                   f"Runtime dispatch {'present' if nvfp4_dispatch else 'MISSING'}.",
            file_path="; ".join(f"{d[0]}:{d[1]}" for d in nvfp4_dispatch[:3]),
        ))

        # 4c. No __CUDA_ARCH__ inside FP4 path (verify FP4 is arch-agnostic)
        mmha_arch_guards = _scan_dir_for_pattern(mmha_dir, r"__CUDA_ARCH__")
        # Filter to only guards with hard 1200 equality
        fp4_arch_blocks = [g for g in mmha_arch_guards if "== 1200" in g[2] and str(arch_num) not in g[2]]
        checks.append(VerifyCheck(
            name="mmha_fp4_arch_independent",
            layer="fp4_kv",
            status="PASS" if not fp4_arch_blocks else "FAIL",
            detail=f"MMHA has {len(mmha_arch_guards)} total __CUDA_ARCH__ refs. "
                   f"{'None' if not fp4_arch_blocks else len(fp4_arch_blocks)} block SM{target_arch[-3:]}. "
                   f"FP4 path is {'arch-agnostic' if not fp4_arch_blocks else 'BLOCKED'}.",
        ))

        # 4d. Head dim 256 instantiation
        instantiation_dir = mmha_dir / "instantiation"
        hdim256_files = list(instantiation_dir.glob("*256*")) if instantiation_dir.exists() else []
        checks.append(VerifyCheck(
            name="mmha_hdim256_instantiation",
            layer="fp4_kv",
            status="PASS" if hdim256_files else "FAIL",
            detail=f"Head dim 256 instantiation: {len(hdim256_files)} file(s). "
                   f"{'Present' if hdim256_files else 'MISSING'} — Qwen3-Next-80B needs Dh=256.",
            file_path="; ".join(f.name for f in hdim256_files[:5]),
        ))

        # 4e. FP4 type reference (__nv_fp4_e2m1 or nv_float4_t)
        fp4_type_refs = _scan_dir_for_pattern(mmha_dir, r"__nv_fp4_e2m1|nv_float4_t")
        checks.append(VerifyCheck(
            name="mmha_fp4_type_usage",
            layer="fp4_kv",
            status="PASS" if fp4_type_refs else "WARN",
            detail=f"{len(fp4_type_refs)} reference(s) to NVFP4 data types in MMHA. "
                   f"Types: __nv_fp4_e2m1, nv_float4_t.",
        ))
    else:
        checks.append(VerifyCheck(
            name="mmha_dir_missing",
            layer="fp4_kv",
            status="FAIL",
            detail=f"MMHA directory not found at {mmha_dir}",
        ))

    # =========================================================================
    # Layer 5: Runtime config path
    # =========================================================================

    # Search for the specific config path that enables nvfp4_kv_cache at runtime.
    # We need: (1) a KvCacheConfig or equivalent with dtype/fp4 option,
    #          (2) code that sets params.nvfp4_kv_cache = true based on config.
    python_root = fork_root / "tensorrt_llm"
    if python_root.exists():
        # Targeted search: files that set or reference nvfp4_kv_cache specifically
        kv_config_refs = []
        for pyf in python_root.rglob("*.py"):
            for lnum, ltxt in _scan_file_for_pattern(pyf, r"nvfp4_kv_cache|kv_cache.*dtype.*fp4|fp4.*kv_cache"):
                kv_config_refs.append((str(pyf.relative_to(fork_root)), lnum, ltxt))

        checks.append(VerifyCheck(
            name="python_nvfp4_kv_config_path",
            layer="runtime",
            status="PASS" if kv_config_refs else "WARN",
            detail=f"{len(kv_config_refs)} Python reference(s) to nvfp4_kv_cache parameter."
                   + (f" Files: {', '.join(set(r[0].split('/')[-1] for r in kv_config_refs[:5]))}"
                      if kv_config_refs
                      else " No explicit Python config path found — may need to set via trtllm-serve "
                           "config or manual C++ param injection."),
        ))
    else:
        checks.append(VerifyCheck(
            name="python_root_missing",
            layer="runtime",
            status="WARN",
            detail=f"Python package root not found at {python_root}. Config tracing limited.",
        ))

    # =========================================================================
    # Layer 6: Model compatibility (Qwen3-Next-80B specifics)
    # =========================================================================
    checks.append(VerifyCheck(
        name="model_head_dim",
        layer="model",
        status="INFO",
        detail="Qwen3-Next-80B: head_dim=256, num_kv_heads=2 (GQA 8:1). "
               "MMHA Dh=256 instantiation required (checked in Layer 4).",
    ))
    checks.append(VerifyCheck(
        name="model_mamba_hybrid",
        layer="model",
        status="INFO",
        detail="Qwen3-Next-80B: 36 linear attention (DeltaNet) + 12 full attention layers. "
               "FP4 KV cache applies to 12 full attention layers only. "
               "Mamba layers have no KV cache. enable_block_reuse: false required.",
    ))
    checks.append(VerifyCheck(
        name="model_moe_routing",
        layer="model",
        status="INFO",
        detail="80B total / 3B active (A3B). Top-10 expert routing per token. "
               "MoE GEMM kernels are KV-cache-independent — no FP4 changes needed.",
    ))

    # =========================================================================
    # Layer 7: Documentation (summary check)
    # =========================================================================
    # Key reference docs that the code checks above rely on.
    # Single summary instead of per-file — we know where the docs are.
    key_docs = [
        "docs/reference/cuda/pdf/ptx_isa_9.1.pdf",
        "docs/reference/cuda/pdf/Blackwell_Tuning_Guide.pdf",
        "docs/reference/cuda/pdf/Blackwell_Compatibility_Guide.pdf",
        "docs/reference/gpu/whitepapers/nvidia-rtx-blackwell-gpu-architecture.pdf",
        "docs/reference/blog/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference",
    ]
    if target_arch == "sm121":
        key_docs.append("docs/reference/dgx/dgx-spark")

    found = sum(1 for d in key_docs if (repo_root / d).exists())
    missing = [d.split("/")[-1] for d in key_docs if not (repo_root / d).exists()]
    checks.append(VerifyCheck(
        name="docs_reference_library",
        layer="docs",
        status="PASS" if found == len(key_docs) else "WARN",
        detail=f"{found}/{len(key_docs)} key reference documents found."
               + (f" Missing: {', '.join(missing)}" if missing else ""),
    ))

    return checks


def cmd_verify(fork_path: str, target_arch: str = "sm121"):
    """
    Programmatic audit of TRT-LLM fork for NVFP4 KV cache readiness.

    Scans source code bottom-up: hardware → build → kernels → FP4 path → runtime → model.
    Each check returns PASS/FAIL/WARN/INFO. Re-run after any code change to detect regressions.
    """
    fork_root = Path(fork_path)
    if not fork_root.exists():
        print(f"\n  ERROR: Fork path not found: {fork_root}")
        return

    checks = verify_trtllm_fork(fork_root, target_arch)

    # Print results
    print(f"\n  {'═' * 74}")
    print(f"  TRT-LLM FORK VERIFICATION — SM{target_arch[-3:]} NVFP4 KV CACHE")
    print(f"  {'═' * 74}")
    print(f"  Fork: {fork_root}")
    print(f"  Target: {target_arch}")

    # Group by layer (bottom-up)
    layers = ["hardware", "build", "kernel", "fp4_kv", "runtime", "model", "docs"]
    layer_names = {
        "hardware": "1. HARDWARE (live + static)",
        "build":    "2. BUILD SYSTEM",
        "kernel":   "3. KERNEL GUARDS",
        "fp4_kv":   "4. FP4 KV CACHE PATH",
        "runtime":  "5. RUNTIME CONFIG",
        "model":    "6. MODEL COMPATIBILITY",
        "docs":     "7. DOCUMENTATION",
    }

    status_icons = {"PASS": "✓", "FAIL": "✗", "WARN": "⚠", "INFO": "·"}
    status_counts = {"PASS": 0, "FAIL": 0, "WARN": 0, "INFO": 0}

    for layer in layers:
        layer_checks = [c for c in checks if c.layer == layer]
        if not layer_checks:
            continue
        print(f"\n  {'─' * 74}")
        print(f"  {layer_names.get(layer, layer.upper())}")
        print(f"  {'─' * 74}")
        for c in layer_checks:
            icon = status_icons.get(c.status, "?")
            status_counts[c.status] = status_counts.get(c.status, 0) + 1
            print(f"\n  [{icon} {c.status:4s}] {c.name}")
            for line in _wrap_text(c.detail, width=68):
                print(f"          {line}")
            if c.file_path:
                print(f"          File: {c.file_path}")
            if c.doc_ref:
                print(f"          Ref:  {c.doc_ref}")

    # Summary
    total = sum(status_counts.values())
    print(f"\n  {'═' * 74}")
    print(f"  SUMMARY: {status_counts['PASS']} PASS | {status_counts['FAIL']} FAIL | "
          f"{status_counts['WARN']} WARN | {status_counts['INFO']} INFO  ({total} checks)")

    if status_counts["FAIL"] > 0:
        print(f"  VERDICT: BLOCKED — {status_counts['FAIL']} failing check(s) must be resolved")
    elif status_counts["WARN"] > 0:
        print(f"  VERDICT: READY WITH WARNINGS — review {status_counts['WARN']} warning(s)")
    else:
        print(f"  VERDICT: READY — all checks pass for SM{target_arch[-3:]} NVFP4 KV cache")
    print(f"  {'═' * 74}\n")

    # Save results to JSON
    results_dir = Path(__file__).parent / "data"
    results_dir.mkdir(exist_ok=True)
    results_file = results_dir / f"verify_{target_arch}.json"
    results_data = {
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "fork_path": str(fork_root),
        "target_arch": target_arch,
        "summary": dict(status_counts),
        "checks": [
            {
                "name": c.name,
                "layer": c.layer,
                "status": c.status,
                "detail": c.detail,
                "file_path": c.file_path,
                "line_numbers": c.line_numbers,
                "doc_ref": c.doc_ref,
            }
            for c in checks
        ],
    }
    results_file.write_text(json.dumps(results_data, indent=2))
    print(f"  Results saved: {results_file}")


def cmd_analyze(config_path: str, gpu_key: str):
    """Full analysis: architecture, inventory, shapes, ceilings, roadmap."""
    hw = GPU_SPECS[gpu_key]

    print(f"\n  Config: {config_path}")
    print(f"  GPU:    {hw.name}")

    config = parse_model_config(config_path)
    params = extract_architecture_parameters(config)
    print_architecture_summary(params)

    weights = enumerate_all_weight_matrices(params)
    print_weight_matrix_inventory(weights)

    shapes = aggregate_gemv_shapes(weights)
    print_gemv_shape_analysis(shapes, hw)
    print_theoretical_ceilings(params, shapes, hw)

    models = compute_full_forward_pass_model(shapes, hw)
    phases = generate_optimization_roadmap(models, hw, params)
    print_optimization_roadmap(phases)

    print("\n" + "=" * 80)
    print("  END OF PRIMER")
    print("=" * 80 + "\n")


# =============================================================================
# Dedup — find duplicate codebases/packages across local repo and remote
#
# Detects "intruders": codebases that exist in multiple places, stray repos,
# and diverged copies. Uses .git dirs as primary signal (actual repo roots).
# Automatically SSHs to edgexpert to check remote too.
#
# Output is concise: duplicates grouped by name, with paths and ages.
# No per-file hash dumps. No false positives from common dir names.
# =============================================================================

def _find_git_repos(base_path: str, location: str = "local") -> List[Dict]:
    """Walk base_path, find all .git directories = actual repo roots.
    Returns list of {path, name, mtime, mtime_str, location, rel_path}."""
    import os
    base = os.path.realpath(base_path)
    repos = []
    for dirpath, dirnames, _filenames in os.walk(base, followlinks=False):
        if '.git' in dirnames:
            dirnames.remove('.git')  # don't descend into .git
            git_dir = os.path.join(dirpath, '.git')
            try:
                mtime = os.path.getmtime(git_dir)
            except OSError:
                mtime = 0
            repos.append({
                "path": dirpath,
                "name": os.path.basename(dirpath),
                "mtime": mtime,
                "mtime_str": time.strftime("%Y-%m-%d %H:%M", time.localtime(mtime)),
                "location": location,
                "rel_path": os.path.relpath(dirpath, base),
            })
        # Prune dirs we never want to recurse into
        dirnames[:] = [d for d in dirnames
                       if d not in {'__pycache__', 'node_modules', '.tox',
                                    '.eggs', 'dist', '.mypy_cache', 'build'}]
    return repos

def _find_remote_git_repos(ssh_alias: str, remote_path: str) -> List[Dict]:
    """SSH to remote, find all .git dirs under remote_path.
    Note: remote_path should be a remote-side path. Use '~' for remote home.
    If a local-expanded home path is detected, replace with ~ for remote."""
    # Guard against local shell expanding ~ to /home/localuser
    if remote_path.startswith("/home/") and not remote_path.startswith("/home/luhnatic"):
        remote_path = "~"
    remote_script = (
        f'find {remote_path} -maxdepth 6 -name .git -type d 2>/dev/null | '
        'while read g; do '
        'd=$(dirname "$g"); '
        'mod=$(stat -c "%Y" "$g" 2>/dev/null || echo 0); '
        'echo "$mod $d"; '
        'done'
    )
    try:
        result = subprocess.run(
            ["ssh", "-o", "ConnectTimeout=5", ssh_alias, remote_script],
            capture_output=True, text=True, timeout=30,
        )
        if result.returncode != 0:
            stderr = result.stderr.strip()
            real_errors = [l for l in stderr.split('\n')
                           if l and 'Warning: Permanently added' not in l]
            if real_errors:
                print(f"  [WARN] SSH to {ssh_alias}:{remote_path}: {real_errors[0][:120]}")
    except subprocess.TimeoutExpired:
        print(f"  [WARN] SSH to {ssh_alias} timed out")
        return []

    repos = []
    for line in result.stdout.strip().split("\n"):
        line = line.strip()
        if not line:
            continue
        parts = line.split(" ", 1)
        if len(parts) != 2:
            continue
        try:
            mtime = float(parts[0])
        except ValueError:
            mtime = 0
        dirpath = parts[1]
        repos.append({
            "path": dirpath,
            "name": Path(dirpath).name,
            "mtime": mtime,
            "mtime_str": time.strftime("%Y-%m-%d %H:%M", time.localtime(mtime)),
            "location": f"remote:{ssh_alias}",
            "rel_path": dirpath,
        })
    return repos

def _normalize_repo_name(name: str) -> str:
    """Normalize repo name for duplicate grouping."""
    n = name.lower()
    n = re.sub(r'^\d{8}_', '', n)        # strip date prefixes (20260124_)
    n = re.sub(r'-v[\d.]+\w*$', '', n)   # strip version suffixes (-v4.3.0)
    return n

def _age_str(mtime: float) -> str:
    if mtime <= 0:
        return "unknown"
    days = int((time.time() - mtime) / 86400)
    if days == 0:
        return "today"
    return f"{days}d ago"

def cmd_dedup(ssh_alias: str = "edgexpert-user", remote_path: str = "~"):
    """Find duplicate/stray codebases across local repo and remote machines.

    Uses .git directories as the signal — these are actual repo/codebase roots,
    not every directory with a Makefile. Groups by normalized name, reports
    duplicates with paths and modification dates.
    """
    repo_root = Path(__file__).parent
    print(f"\n{'='*70}")
    print(f"  DEDUP — Duplicate Codebase Detection")
    print(f"{'='*70}")

    # --- Scan local ---
    print(f"  Local: {repo_root} ...", end="", flush=True)
    local_repos = _find_git_repos(str(repo_root), location="local")
    print(f" {len(local_repos)} repos")

    # --- Scan remote ---
    print(f"  Remote: {ssh_alias}:{remote_path} ...", end="", flush=True)
    remote_repos = _find_remote_git_repos(ssh_alias, remote_path)
    print(f" {len(remote_repos)} repos")
    print()

    all_repos = local_repos + remote_repos

    # --- Group by normalized name ---
    groups = defaultdict(list)
    for repo in all_repos:
        key = _normalize_repo_name(repo["name"])
        groups[key].append(repo)

    # --- Separate duplicates from singletons ---
    dupes = {k: v for k, v in groups.items() if len(v) > 1}
    singles = {k: v[0] for k, v in groups.items() if len(v) == 1}

    # --- Report duplicates ---
    if dupes:
        total_dupes = sum(len(v) for v in dupes.values())
        print(f"  DUPLICATES: {len(dupes)} groups ({total_dupes} total copies)")
        print(f"  {'-'*66}")
        for key in sorted(dupes.keys()):
            entries = sorted(dupes[key], key=lambda x: x["mtime"], reverse=True)
            print(f"\n  {key} ({len(entries)} copies):")
            for e in entries:
                loc = e["location"][:16]
                age = _age_str(e["mtime"])
                print(f"    {loc:<16} {e['mtime_str']}  {age:<8}  {e['rel_path']}")
    else:
        print("  No duplicates found.")

    # --- Full inventory ---
    print(f"\n  ALL REPOS ({len(all_repos)} total):")
    print(f"  {'LOC':<10} {'MODIFIED':<18} {'AGE':<8} {'PATH'}")
    print(f"  {'-'*10} {'-'*18} {'-'*8} {'-'*40}")
    for r in sorted(all_repos, key=lambda x: x["name"].lower()):
        loc = r["location"][:10]
        age = _age_str(r["mtime"])
        print(f"  {loc:<10} {r['mtime_str']:<18} {age:<8} {r['rel_path']}")


# =============================================================================
# Main
# =============================================================================

def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="MoE-Mamba-Blackwell Development Primer — Stateful Optimization Tracker"
    )
    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # analyze
    p_analyze = subparsers.add_parser("analyze", help="Full architecture analysis and roadmap")
    p_analyze.add_argument("config", nargs="?", default=DEFAULT_CONFIG)
    p_analyze.add_argument("--gpu", choices=["sm120", "sm121"], default="sm120")

    # next
    p_next = subparsers.add_parser("next", help="What to do right now")
    p_next.add_argument("--config", default=DEFAULT_CONFIG)
    p_next.add_argument("--gpu", choices=["sm120", "sm121"], default="sm120")

    # status
    p_status = subparsers.add_parser("status", help="Phase overview")
    p_status.add_argument("--config", default=DEFAULT_CONFIG)
    p_status.add_argument("--gpu", choices=["sm120", "sm121"], default="sm120")

    # record
    p_record = subparsers.add_parser("record", help="Record a measurement")
    p_record.add_argument("phase", type=int, help="Phase number")
    p_record.add_argument("key", help="Measurement key")
    p_record.add_argument("value", type=float, help="Measurement value")

    # history
    p_history = subparsers.add_parser("history", help="Show all measurements")
    p_history.add_argument("phase", type=int, nargs="?", default=None, help="Optional phase number")

    # complete
    p_complete = subparsers.add_parser("complete", help="Manually mark phase done")
    p_complete.add_argument("phase", type=int)

    # note
    p_note = subparsers.add_parser("note", help="Add a note to a phase")
    p_note.add_argument("phase", type=int)
    p_note.add_argument("text", nargs="+")

    # research (requires external LLM endpoints — commented out for now)
    # p_research = subparsers.add_parser("research", help="Run research pipeline for a phase")
    # p_research.add_argument("phase", type=int, nargs="?", default=None,
    #                         help="Phase number (omit for status)")
    # p_research.add_argument("--force", action="store_true", help="Re-run even if cached")

    # verify
    p_verify = subparsers.add_parser("verify", help="Programmatic code audit of TRT-LLM fork")
    p_verify.add_argument("fork_path", help="Path to TRT-LLM fork root (containing cpp/ and tensorrt_llm/)")
    p_verify.add_argument("--arch", choices=["sm120", "sm121"], default="sm121",
                          help="Target architecture (default: sm121)")

    # unknowns
    subparsers.add_parser("unknowns", help="Show open questions and design decisions")

    # dedup
    p_dedup = subparsers.add_parser("dedup", help="Find duplicate codebases/packages locally and on edgexpert")
    p_dedup.add_argument("--remote", default="edgexpert-user", help="SSH alias for remote machine")
    p_dedup.add_argument("--remote-path", default="~/llm-deploy", help="Remote base path to scan")

    args = parser.parse_args()

    if args.command is None:
        # Default: run verify if a TRT-LLM fork exists nearby, else show status
        default_fork = Path("SM12x.NVFP4w-NVFP4kv.MoE-GQA.TrtLLM-Fork/TensorRT-LLM")
        if default_fork.exists():
            cmd_verify(str(default_fork), "sm121")
        else:
            results = load_results()
            if results:
                cmd_status(DEFAULT_CONFIG, "sm120")
            else:
                print("  No TRT-LLM fork found. Run: primer verify <fork_path>")
    elif args.command == "analyze":
        cmd_analyze(args.config, args.gpu)
    elif args.command == "next":
        cmd_next(args.config, args.gpu)
    elif args.command == "status":
        cmd_status(args.config, args.gpu)
    elif args.command == "record":
        cmd_record(args.phase, args.key, args.value)
    elif args.command == "history":
        cmd_history(args.phase)
    elif args.command == "complete":
        cmd_complete(args.phase)
    elif args.command == "note":
        cmd_note(args.phase, " ".join(args.text))
    # elif args.command == "research":
    #     if args.phase is None:
    #         cmd_research_status()
    #     else:
    #         cmd_research(args.phase, force=args.force)
    elif args.command == "verify":
        cmd_verify(args.fork_path, args.arch)
    elif args.command == "unknowns":
        cmd_unknowns()
    elif args.command == "dedup":
        cmd_dedup(args.remote, args.remote_path)


if __name__ == "__main__":
    main()