#!/usr/bin/env python
"""Issue 5 round two: isolate the compressed-tensors W8A8-int8-dynamic path.

Mirrors marlin_isolate.py but for the int-quantized W8A8 dynamic scheme
(int8 weights, static per-output-channel scales, dynamic symmetric per-token
int8 activations) whose kernel on sm_86 is the cutlass int8 scaled-MM —
no Marlin in the path. Loads synthetic int8 tensors through the production
weight-loader machinery and compares layer(x) against a float32 reference
computed from the same dequantized weights with per-token dynamic int8
activation quantization applied to the reference input.

Expected rel-L2 when the chain is correct: a few percent (int8 weight error
~1% plus int8 per-token activation rounding).

Usage: w8a8_isolate.py [--block N]
"""

from __future__ import annotations

import argparse
import json

import torch


def synth(out: int, inp: int, ones: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
    """Per-channel symmetric int8 weight + scale, mirroring the artifact builder.

    ones=True emits an analytically known weight (q=127, scale=1/127 -> W=1):
    with an all-ones input row the layer output must be exactly in_features.
    """
    if ones:
        q = torch.full((out, inp), 127, dtype=torch.int8)
        return q, torch.full((out, 1), 1.0 / 127.0, dtype=torch.float32)
    w = torch.randn(out, inp, dtype=torch.float32)
    scale = w.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
    q = torch.round(w / scale).clamp(-127, 127).to(torch.int8)
    return q, scale.to(torch.float32)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--block", type=int, default=0)
    parser.add_argument("--ones", action="store_true",
                        help="analytic weights (W=1) and ones input; output must equal in_features")
    parser.add_argument("--act-dtype", choices=("bfloat16", "float16"), default="bfloat16")
    args = parser.parse_args()

    import vllm.distributed
    from vllm.config import VllmConfig, set_current_vllm_config
    from vllm.distributed import (
        destroy_distributed_environment,
        destroy_model_parallel,
        initialize_model_parallel,
    )
    from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
    from compressed_tensors.quantization import (
        QuantizationArgs,
        QuantizationConfig,
        QuantizationScheme,
        QuantizationStatus,
        QuantizationStrategy,
        QuantizationType,
    )
    from compressed_tensors.config import CompressionFormat
    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config

    weights = QuantizationArgs(num_bits=8, type=QuantizationType.INT, symmetric=True,
                               strategy=QuantizationStrategy.CHANNEL, dynamic=False, observer="minmax")
    acts = QuantizationArgs(num_bits=8, type=QuantizationType.INT, symmetric=True,
                            strategy=QuantizationStrategy.TOKEN, dynamic=True)
    scheme = QuantizationScheme(targets=["Linear"], weights=weights, input_activations=acts)
    disk = QuantizationConfig(
        config_groups={"group_0": scheme}, format=CompressionFormat.int_quantized,
        ignore=["img_in", "txt_in", "time_text_embed", "modulation", "norm_out", "proj_out"],
        quant_method="compressed-tensors", quantization_status=QuantizationStatus.FROZEN,
    ).to_dict()
    quant_config = resolve_component_quant_config(
        build_quant_config({"transformer": disk}), "transformer"
    )

    vllm.distributed.init_distributed_environment(
        backend="gloo", world_size=1, rank=0, local_rank=0,
        distributed_init_method="tcp://127.0.0.1:29602",
    )
    with set_current_vllm_config(VllmConfig()):
        initialize_model_parallel(tensor_model_parallel_size=1)

    device = torch.device("cuda")
    torch.manual_seed(2301)
    tensors = {
        "q": synth(4096, 4096, args.ones),
        "k": synth(4096, 4096, args.ones),
        "v": synth(4096, 4096, args.ones),
        "to_out": synth(4096, 4096, args.ones),
    }

    b = args.block
    p = f"transformer.transformer_blocks.{b}"
    results = []
    with torch.device(device.type):
        to_qkv = QKVParallelLinear(
            hidden_size=4096, head_size=128, total_num_heads=32, bias=False,
            quant_config=quant_config, prefix=f"{p}.attn.to_qkv",
        )
        for part, shard in (("q", "q"), ("k", "k"), ("v", "v")):
            q, s = tensors[part]
            to_qkv.weight.weight_loader(to_qkv.weight, q.to(device), shard)
            to_qkv.weight_scale.weight_loader(to_qkv.weight_scale, s.to(device), shard)
        to_qkv.quant_method.process_weights_after_loading(to_qkv)

        to_out = RowParallelLinear(
            4096, 4096, bias=False, input_is_parallel=True, return_bias=False,
            quant_config=quant_config, prefix=f"{p}.attn.to_out.0",
        )
        q, s = tensors["to_out"]
        to_out.weight.weight_loader(to_out.weight, q.to(device))
        to_out.weight_scale.weight_loader(to_out.weight_scale, s.to(device))
        loaded_bitwise_ok = bool(torch.equal(to_out.weight.detach().cpu(), q))
        loaded_scale_ok = bool(torch.allclose(to_out.weight_scale.detach().float().cpu(), s))
        to_out.quant_method.process_weights_after_loading(to_out)

    for name, layer, parts, in_features in (
        ("attn.to_qkv (fused q/k/v)", to_qkv, ("q", "k", "v"), 4096),
        ("attn.to_out.0 (row)", to_out, ("to_out",), 4096),
    ):
        if args.ones:
            x = torch.ones(4, in_features, dtype=getattr(torch, args.act_dtype), device=device)
        else:
            x = torch.randn(512, in_features, dtype=getattr(torch, args.act_dtype), device=device)
        result = layer(x)
        out = (result[0] if isinstance(result, tuple) else result).float()

        if args.ones:
            expected = torch.full_like(out, float(in_features))
            rel = float((out - expected).norm() / expected.norm())
            finite = bool(torch.isfinite(out).all())
            results.append({
                "layer": name,
                "rel_l2": round(rel, 6) if finite else "NaN",
                "out_finite": finite,
                "out_row0_first8": [round(v, 2) for v in out[0, :8].tolist()],
            })
            continue

        ref_w = torch.cat([tensors[k][0].float() * tensors[k][1] for k in parts], dim=0)
        # Reference with the same dynamic per-token symmetric int8 activation
        # quantization the kernel applies (dequantized acts, dequantized weights).
        xs = x.float()
        act_scale = xs.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
        x_dq = torch.round(xs / act_scale).clamp(-127, 127) * act_scale
        ref = x_dq @ ref_w.T.to(device)
        finite = bool(torch.isfinite(out).all())
        rel = float((out - ref).norm() / ref.norm())
        row = {"layer": name, "rel_l2": round(rel, 5) if finite else "NaN", "out_finite": finite}
        if name.startswith("attn.to_out"):
            row["loaded_bitwise_ok"] = loaded_bitwise_ok
            row["loaded_scale_ok"] = loaded_scale_ok
            # Reference computed from the layer's own post-repack weight
            # parameter ((K, N) after the transpose) and scales.
            w_layer = to_out.weight.detach().float()  # (K, N)
            s_layer = to_out.weight_scale.detach().float().squeeze(1)  # (N,)
            ref_layer = x_dq @ (w_layer * s_layer.unsqueeze(0)).to(device)
            row["rel_l2_vs_layer_weight"] = round(float((out - ref_layer).norm() / ref_layer.norm()), 5)
        results.append(row)

    destroy_model_parallel()
    destroy_distributed_environment()
    print(json.dumps(results, indent=1))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())