slidesparse-icml33765-repro/reproduction-bundle:v1/runner.py

#!/usr/bin/env python3
"""Independent claim checks and scaled CUDA proxies for SlideSparse."""

from __future__ import annotations

import argparse
import csv
import itertools
import json
import math
import random
import re
import statistics
import time
from pathlib import Path


HERE = Path(__file__).resolve().parent
SOURCE_DATA = HERE / "source_data"


def emit(record: dict) -> None:
    print("RESULT_JSON " + json.dumps(record, sort_keys=True), flush=True)


def windows(n: int) -> list[tuple[int, ...]]:
    return [tuple(range(2 * j, 2 * j + 4)) for j in range(n - 1)]


def greedy_assign(mask: tuple[int, ...], n: int) -> list[list[int]]:
    """Mirror the released offline packer's left-to-right allocation."""
    assigned: set[int] = set()
    result: list[list[int]] = []
    for window in windows(n):
        chosen = [i for i in window if mask[i] and i not in assigned][:2]
        result.append(chosen)
        assigned.update(chosen)
    return result


def exact_assignment_exists(mask: tuple[int, ...], n: int) -> bool:
    nonzeros = [i for i, value in enumerate(mask) if value]
    candidates = [[j for j, window in enumerate(windows(n)) if i in window] for i in nonzeros]
    order = sorted(range(len(nonzeros)), key=lambda idx: len(candidates[idx]))
    capacity = [2] * (n - 1)

    def search(depth: int) -> bool:
        if depth == len(order):
            return True
        idx = order[depth]
        for window_id in candidates[idx]:
            if capacity[window_id]:
                capacity[window_id] -= 1
                if search(depth + 1):
                    return True
                capacity[window_id] += 1
        return False

    return search(0)


def coverage_audit(max_n: int, seed: int) -> dict:
    rng = random.Random(seed)
    rows = []
    inner_product_trials = 0
    max_error = 0.0
    for n in range(2, max_n + 1):
        length = 2 * n
        nonzero_count = 2 * n - 2
        patterns = 0
        exact_failures = 0
        greedy_failures = 0
        for nz_positions in itertools.combinations(range(length), nonzero_count):
            mask_list = [0] * length
            for index in nz_positions:
                mask_list[index] = 1
            mask = tuple(mask_list)
            patterns += 1
            if not exact_assignment_exists(mask, n):
                exact_failures += 1
            allocation = greedy_assign(mask, n)
            assigned = {index for group in allocation for index in group}
            if len(assigned) != nonzero_count:
                greedy_failures += 1

            weights = [rng.uniform(-2, 2) if mask[i] else 0.0 for i in range(length)]
            activations = [rng.uniform(-2, 2) for _ in range(length)]
            direct = sum(w * x for w, x in zip(weights, activations))
            lifted = sum(weights[i] * activations[i] for group in allocation for i in group)
            max_error = max(max_error, abs(direct - lifted))
            inner_product_trials += 1

        rows.append(
            {
                "N": n,
                "block": f"{2*n-2}:{2*n}",
                "windows": n - 1,
                "patterns": patterns,
                "exact_failures": exact_failures,
                "greedy_failures": greedy_failures,
                "fewer_windows_capacity": 2 * max(0, n - 2),
                "required_nonzeros": nonzero_count,
            }
        )
    return {
        "claim": 1,
        "status": "verified" if all(r["exact_failures"] == r["greedy_failures"] == 0 for r in rows) else "failed",
        "rows": rows,
        "inner_product_trials": inner_product_trials,
        "max_abs_inner_product_error": max_error,
        "necessity_reason": "K=N-2 has capacity 2N-4 < 2N-2",
    }


def theory_audit() -> dict:
    rows = []
    expected = {3: 1.5, 4: 4 / 3, 5: 1.25}
    for n, target in expected.items():
        gamma = 2 * (n - 1) / n
        speedup = 2 / gamma
        rows.append(
            {
                "N": n,
                "pattern": f"{2*n-2}:{2*n}",
                "expansion_gamma": gamma,
                "derived_speedup": speedup,
                "reported_speedup": target,
                "abs_error": abs(speedup - target),
            }
        )
    return {"claim": 2, "status": "verified", "rows": rows, "limit_as_N_to_infinity": 1.0}


def read_csv(name: str) -> list[dict[str, str]]:
    with (SOURCE_DATA / name).open(newline="") as handle:
        return list(csv.DictReader(handle))


def nonempty_float(value: str | None) -> float | None:
    if value is None or value.strip() == "":
        return None
    return float(value)


def source_table_audit() -> list[dict]:
    prefill_abs = read_csv("A100_prefill_absolute_Qwen2.5-7B-INT8.csv")
    prefill_speed = read_csv("A100_prefill_speedup_Qwen2.5-7B-INT8.csv")
    row_abs = next(row for row in prefill_abs if row["M"] == "8192")
    row_speed = next(row for row in prefill_speed if row["M"] == "8192")
    ratio = float(row_abs["cusparse_2_8"]) / float(row_abs["cuBLAS"])
    claim3 = {
        "claim": 3,
        "status": "not_exactly_reproduced",
        "M": 8192,
        "dense_tokens_s": float(row_abs["cuBLAS"]),
        "slidesparse_6_8_tokens_s": float(row_abs["cusparse_2_8"]),
        "recomputed_speedup": ratio,
        "released_rounded_table": float(row_speed["cusparse_2_8"]),
        "paper_value": 1.33,
        "note": "Released source data recomputes to 1.3186x (1.32x at two decimals), not exactly 1.33x.",
    }

    a100 = read_csv("A100_kernel_INT8_total_Qwen2.5-7B.csv")
    rtx = read_csv("RTX4090_kernel_FP8_total_Qwen2.5-7B.csv")
    h100 = read_csv("H100_kernel_FP8_total_Qwen2.5-7B.csv")
    a100_large = [float(r["cuSPARSE_2_8"]) for r in a100 if int(r["M"]) in (8192, 16384)]
    rtx_large = [float(r["cuSPARSE_2_8"]) for r in rtx if int(r["M"]) in (4096, 16384)]
    h100_large = [float(r["cuSPARSE_2_8"]) for r in h100 if int(r["M"]) in (8192, 16384)]
    claim4 = {
        "claim": 4,
        "status": "partly_verified_but_hardware_attribution_refuted",
        "A100_INT8_6_8_M8192_16384": a100_large,
        "RTX4090_FP8_6_8_M4096_16384": rtx_large,
        "H100_FP8_6_8_M8192_16384": h100_large,
        "paper_text": "RTX 4090 FP8 achieves 1.35-1.37x; it does not attribute that range to H100.",
    }

    accuracy_rows = read_csv("reported_Qwen3_reasoning_accuracy.csv")
    recomputed = {}
    for row in accuracy_rows:
        values = [float(row[key]) for key in ("GSM8K", "MATH-500", "HumanEval")]
        recomputed[row["regime"]] = {"mean": statistics.mean(values), "reported_average": float(row["Average"])}
    claim5 = {
        "claim": 5,
        "status": "arithmetic_verified_experiment_not_reproducible_from_release",
        "averages": recomputed,
        "note": "The release lacks the Qwen3 fine-tuned checkpoints and exact training recipe; its accuracy_eval workflow targets different models/tasks.",
    }

    decode_records = []
    for path in sorted(SOURCE_DATA.glob("*_decode_speedup_*.csv")):
        gpu, _, _, model = path.stem.split("_", 3)
        for row in read_csv(path.name):
            value = nonempty_float(row.get("cusparse_2_8"))
            if value is not None:
                decode_records.append({"gpu": gpu, "model": model, "M": int(row["M"]), "speedup": value})
    within_headline = [r for r in decode_records if 1.07 <= r["speedup"] <= 1.21]
    claim6 = {
        "claim": 6,
        "status": "qualified",
        "all_released_decode_6_8_range": [min(r["speedup"] for r in decode_records), max(r["speedup"] for r in decode_records)],
        "records_in_paper_headline_range": len(within_headline),
        "total_records": len(decode_records),
        "note": "The 1.07-1.21x prose summarizes selected configurations; the full released tables include values outside it.",
    }
    return [claim3, claim4, claim5, claim6]


def run_baseline(config: dict) -> None:
    emit({"event": "environment", "mode": "baseline", "python": __import__("sys").version})
    emit(coverage_audit(int(config["coverage_max_n"]), int(config["seed"])))
    emit(theory_audit())
    for record in source_table_audit():
        emit(record)


def cuda_time_ms(fn, warmup: int, repeats: int) -> tuple[float, float]:
    import torch

    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    samples = []
    for _ in range(repeats):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        fn()
        end.record()
        end.synchronize()
        samples.append(float(start.elapsed_time(end)))
    return statistics.median(samples), statistics.mean(samples)


def run_hardware(config: dict) -> None:
    import torch
    from torch.sparse import to_sparse_semi_structured

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for the hardware benchmark")
    h = config["hardware"]
    torch.manual_seed(int(config["seed"]))
    device = torch.device("cuda:0")
    props = torch.cuda.get_device_properties(device)
    requested_dtype = h["dtype"]
    dtype = torch.int8 if requested_dtype == "int8" else torch.float16
    n = int(h["n"])
    k = int(h["k"])
    k_slid = 3 * k // 2
    warmup = int(h["warmup"])
    repeats = int(h["repeats"])

    dense_w = torch.randint(-8, 8, (n, k), device=device, dtype=torch.int8) if dtype == torch.int8 else torch.randn((n, k), device=device, dtype=dtype)
    sparse_w_dense = torch.randint(-8, 8, (n, k_slid), device=device, dtype=torch.int8) if dtype == torch.int8 else torch.randn((n, k_slid), device=device, dtype=dtype)
    sparse_w_dense[:, 2::4] = 0
    sparse_w_dense[:, 3::4] = 0
    sparse_w = to_sparse_semi_structured(sparse_w_dense)

    emit(
        {
            "event": "environment",
            "mode": "hardware",
            "gpu": props.name,
            "compute_capability": f"{props.major}.{props.minor}",
            "memory_gib": props.total_memory / 2**30,
            "torch": torch.__version__,
            "cuda": torch.version.cuda,
            "requested_dtype": requested_dtype,
        }
    )

    for m in h["m_values"]:
        m = int(m)
        if dtype == torch.int8:
            dense_x = torch.randint(-8, 8, (m, k), device=device, dtype=dtype)
            sparse_x = torch.randint(-8, 8, (m, k_slid), device=device, dtype=dtype)
            dense_fn = lambda: torch._int_mm(dense_x, dense_w.t())
        else:
            dense_x = torch.randn((m, k), device=device, dtype=dtype)
            sparse_x = torch.randn((m, k_slid), device=device, dtype=dtype)
            dense_fn = lambda: torch.mm(dense_x, dense_w.t())
        sparse_fn = lambda: torch.mm(sparse_x, sparse_w.t())
        dense_median, dense_mean = cuda_time_ms(dense_fn, warmup, repeats)
        sparse_median, sparse_mean = cuda_time_ms(sparse_fn, warmup, repeats)
        emit(
            {
                "claim": 4 if m >= 512 else 6,
                "status": "scaled_hardware_proxy",
                "M": m,
                "N": n,
                "K_dense": k,
                "K_slid": k_slid,
                "dtype": str(dtype).replace("torch.", ""),
                "dense_median_ms": dense_median,
                "dense_mean_ms": dense_mean,
                "slidesparse_sparse_median_ms": sparse_median,
                "slidesparse_sparse_mean_ms": sparse_mean,
                "speedup": dense_median / sparse_median,
                "theory_6_8": 4 / 3,
                "scope": "PyTorch semi-structured sparse GEMM proxy; excludes fused activation lifting and vLLM.",
            }
        )


def magnitude_prune_model(model, pattern: str) -> dict:
    import torch

    keep, total = map(int, pattern.split(":"))
    prune = total - keep
    zeroed = 0
    elements = 0
    with torch.no_grad():
        for module in model.modules():
            if not isinstance(module, torch.nn.Linear):
                continue
            weight = module.weight.data
            usable = weight.shape[1] - weight.shape[1] % total
            if usable == 0:
                continue
            view = weight[:, :usable].reshape(weight.shape[0], -1, total)
            indices = view.abs().topk(prune, dim=-1, largest=False).indices
            view.scatter_(-1, indices, 0)
            zeroed += indices.numel()
            elements += view.numel()
    return {"pattern": pattern, "zeroed": zeroed, "eligible_elements": elements, "sparsity": zeroed / elements}


def extract_answer(text: str) -> str | None:
    matches = re.findall(r"-?\d+(?:\.\d+)?", text.replace(",", ""))
    return matches[-1] if matches else None


def run_accuracy_proxy(config: dict) -> None:
    import torch
    from datasets import load_dataset
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for the accuracy proxy")
    a = config["accuracy"]
    model_id = a["model"]
    samples = list(load_dataset("openai/gsm8k", "main", split=f"test[:{int(a['num_samples'])}]"))
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    emit({"event": "environment", "mode": "accuracy_proxy", "gpu": torch.cuda.get_device_name(0), "model": model_id, "samples": len(samples)})

    for pattern in a["patterns"]:
        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="cuda").eval()
        pruning = None if pattern == "dense" else magnitude_prune_model(model, pattern)
        correct = 0
        predictions = []
        answer_nlls = []
        started = time.time()
        for item in samples:
            prompt = "Solve the following problem. Give the final numeric answer clearly.\n\n" + item["question"]
            messages = [{"role": "user", "content": prompt}]
            rendered = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
            inputs = tokenizer(rendered, return_tensors="pt").to("cuda")
            with torch.inference_mode():
                teacher = tokenizer(rendered + "\n" + item["answer"], return_tensors="pt").to("cuda")
                labels = teacher.input_ids.clone()
                labels[:, : inputs.input_ids.shape[1]] = -100
                answer_nll = float(model(**teacher, labels=labels).loss)
                output = model.generate(**inputs, max_new_tokens=int(a["max_new_tokens"]), do_sample=False, pad_token_id=tokenizer.eos_token_id)
            completion = tokenizer.decode(output[0, inputs.input_ids.shape[1] :], skip_special_tokens=True)
            predicted = extract_answer(completion)
            gold = extract_answer(item["answer"].split("####")[-1])
            is_correct = predicted == gold
            correct += int(is_correct)
            answer_nlls.append(answer_nll)
            predictions.append({"predicted": predicted, "gold": gold, "correct": is_correct, "answer_nll": answer_nll})
        mean_answer_nll = statistics.mean(answer_nlls)
        emit(
            {
                "claim": 5,
                "status": "scaled_proxy",
                "model": model_id,
                "pattern": pattern,
                "benchmark": "GSM8K first contiguous test examples, zero-shot",
                "samples": len(samples),
                "correct": correct,
                "accuracy_percent": 100 * correct / len(samples),
                "mean_teacher_forced_answer_nll": mean_answer_nll,
                "answer_perplexity": math.exp(min(mean_answer_nll, 20.0)),
                "pruning": pruning,
                "elapsed_s": time.time() - started,
                "predictions": predictions,
                "scope": "0.6B magnitude-pruned proxy, not the paper's fine-tuned Qwen3 setup or three-task average.",
            }
        )
        del model
        torch.cuda.empty_cache()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, required=True)
    args = parser.parse_args()
    config = json.loads(args.config.read_text())
    mode = config["mode"]
    if mode == "baseline":
        run_baseline(config)
    elif mode == "hardware":
        run_hardware(config)
    elif mode == "accuracy_proxy":
        run_accuracy_proxy(config)
    else:
        raise ValueError(f"Unknown mode: {mode}")


if __name__ == "__main__":
    main()