import torch

import bitsandbytes as bnb
from bitsandbytes import functional as F

k = 20

torch.set_printoptions(precision=5, sci_mode=False, linewidth=120, edgeitems=20, threshold=10000)

@pytest.mark.parametrize(
    ("batch", "seq", "model", "hidden"),
    [
        # pytest.param(1, 128, 6656, 4 * 6656, id="batch=1, seq=128, model=6656, hidden=26k"),
        pytest.param(1, 1, 3584, 512, id="batch=1, seq=128, model=3584, hidden=19k"),
        # pytest.param(4, 128, 6656, 4 * 6656, id="batch=4, seq=128, model=6656, hidden=26k"),
        # pytest.param(16, 256, 6656, 4 * 6656, id="batch=16, seq=256, model=6656, hidden=26k")
    ],
)
@pytest.mark.benchmark
def test_bench_matmul(batch, seq, model, hidden):
    iters = 1000
    formatB = F.get_special_format_str()

    A = torch.randn(batch, seq, model, device="cuda").half()
    B = torch.empty(hidden, model, dtype=torch.float16, device="cuda")
    torch.nn.init.xavier_uniform_(B)

    B_fp4, state = F.quantize_fp4(B)
    B_fp4_c, state_c = F.quantize_fp4(B, compress_statistics=True)

    B_nf4, state_nf4 = F.quantize_nf4(B)
    B_nf4_c, state_nf4_c = F.quantize_nf4(B, compress_statistics=True)

    linear8bit = bnb.nn.Linear8bitLt(model, hidden, False, False).cuda().half()
    linear8bit.eval()

    outliers = torch.randint(0, model, size=(5,)).cuda()
    A[:, :, outliers] = 8.0

    linearMixedBit = bnb.nn.Linear8bitLt(model, hidden, False, False, threshold=6.0).cuda().half()
    # linearMixedBit.eval()

    linear8bit_train = bnb.nn.Linear8bitLt(model, hidden, False).cuda().half()
    linear8bit_train_thresh = bnb.nn.Linear8bitLt(model, hidden, False, threshold=6.0).cuda().half()
    bnb.matmul_4bit(A, B_nf4.t(), quant_state=state_nf4)

    # warmup
    for i in range(iters):
        torch.matmul(A, B.t())
    torch.cuda.synchronize()
    print("")

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(iters):
        torch.matmul(A, B.t())
    torch.cuda.synchronize()
    print(
        f"pytorch fp16: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s",
    )

    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(iters):
    #    bnb.matmul_4bit(A, B_fp4.t(), quant_state=state)
    # torch.cuda.synchronize()
    # print( f"bnb fp4: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s" )

    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(iters):
    #    bnb.matmul_4bit(A, B_fp4.t(), quant_state=state_c)
    # torch.cuda.synchronize()
    # print( f"bnb fp4 + compressed stats: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s" )

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(iters):
        bnb.matmul_4bit(A, B_nf4.t(), quant_state=state_nf4)
    torch.cuda.synchronize()
    print(f"bnb nf4: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s")

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(iters):
        bnb.matmul_4bit(A, B_nf4_c.t(), quant_state=state_nf4_c)
    torch.cuda.synchronize()
    print(
        f"bnb nf4+DQ: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s"
    )

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(iters):
        bnb.matmul(A, B)
    torch.cuda.synchronize()
    print(
        f"B -> CB (each iteration): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s"
    )

    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(iters):
        bnb.matmul(A, B, threshold=6.0)
    torch.cuda.synchronize()
    print(
        f"B -> CB + threshold: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s"
    )

    CA, SCA, _ = F.int8_vectorwise_quant(A, threshold=0.0)
    CB, SCB, _ = F.int8_vectorwise_quant(B)
    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(iters):
        # CA, CAt, SCA, SCAt, coo_tensorA = F.double_quant(A, threshold=0.0)
        out32 = F.int8_linear_matmul(CA, CB)
    torch.cuda.synchronize()
    print(
        f"no overhead int8 [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s"
    )

    # C32A, SA = F.transform(CA, "col32")

    # CxB, SB = F.transform(CB, to_order=formatB)
    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(iters):
    #    out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB)
    # torch.cuda.synchronize()
    # print(f"no overhead matmul-lt: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s")

    # CA, CAt, SCA, SCAt, coo_tensorA = F.double_quant(A, threshold=0.0)
    # C32A, SA = F.transform(CA, "col32")
    # CB, CBt, SCB, SCBt, coo_tensorB = F.double_quant(B)
    # CxB, SB = F.transform(CB, to_order=formatB)
    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(iters):
    #    out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB)
    # torch.cuda.synchronize()
    # print(f"no overhead matmul-lt: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s")

    # BA, statsB = F.vectorwise_quant(B, dim=1)
    # CxB, SB = F.nvidia_transform(CB, to_order=formatB)
    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(iters):
    #    A2 = A.view(-1, A.shape[-1]).contiguous()
    #    CA, statsA = F.vectorwise_quant(A2, dim=1)
    #    C32A, SA = F.nvidia_transform(CA, "col32")
    #    out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB)
    #    Cout, Sout = F.nvidia_transform(out32, "row", state=Sout32)
    #    F.vectorwise_mm_dequant(Cout, statsA, statsB.t())
    # torch.cuda.synchronize()
    # print(f"vector pytorch + nvidia: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s")

    # BA, statsB = F.vectorwise_quant(B, dim=1, quant_type="linear")
    # CxB, SB = F.nvidia_transform(CB, to_order=formatB)
    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(iters):
    #    A2 = A.view(-1, A.shape[-1]).contiguous()
    #    CA, statsA = F.vectorwise_quant(A2, dim=1, quant_type="linear")
    #    C32A, SA = F.nvidia_transform(CA, "col32")
    #    out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB)
    #    Cout, Sout = F.nvidia_transform(out32, "row", state=Sout32)
    #    out = Cout * statsB * statsA * (1.0 / (127 * 127))
    # torch.cuda.synchronize()
    # print(f"linear pytorch + nvidia: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s")

    linear8bit(A)
    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(iters):
        linear8bit(A)
    torch.cuda.synchronize()
    print(
        f"bnb linear8bitlt (eval): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s"
    )

    linearMixedBit(A)
    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(iters):
        linearMixedBit(A)
    torch.cuda.synchronize()
    print(
        f"bnb linear8bitlt with threshold (eval): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s"
    )

    # linear8bit_train(A)
    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(iters):
    #    linear8bit_train(A)
    # torch.cuda.synchronize()
    # print( f"bnb linear8bitlt (training): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s")

    # linear8bit_train_thresh(A)
    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(iters):
    #    linear8bit_train(A)
    # torch.cuda.synchronize()
    # print( f"bnb linear8bitlt with threshold (training): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s")


"""
Inference benchmarking tool.

Requirements:
    transformers
    accelerate
    bitsandbytes
    optimum-benchmark

Usage: python inference_benchmark.py model_id

options:
    -h, --help            show this help message and exit
    --configs {bf16,fp16,nf4,nf4-dq,int8,int8-decomp} [{bf16,fp16,nf4,nf4-dq,int8,int8-decomp} ...]
    --bf16
    --fp16
    --nf4
    --nf4-dq
    --int8
    --int8-decomp
    --batches BATCHES [BATCHES ...]
    --input-length INPUT_LENGTH
    --out-dir OUT_DIR
    --iterations ITERATIONS
    --warmup-runs WARMUP_RUNS
    --output-length OUTPUT_LENGTH
"""

import argparse
from pathlib import Path

from optimum_benchmark import Benchmark, BenchmarkConfig, InferenceConfig, ProcessConfig, PyTorchConfig
from optimum_benchmark.logging_utils import setup_logging
import torch

torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

BFLOAT16_SUPPORT = torch.cuda.get_device_capability()[0] >= 8

WEIGHTS_CONFIGS = {
    "fp16": {"torch_dtype": "float16", "quantization_scheme": None, "quantization_config": {}},
    "bf16": {"torch_dtype": "bfloat16", "quantization_scheme": None, "quantization_config": {}},
    "nf4": {
        "torch_dtype": "bfloat16" if BFLOAT16_SUPPORT else "float16",
        "quantization_scheme": "bnb",
        "quantization_config": {
            "load_in_4bit": True,
            "bnb_4bit_quant_type": "nf4",
            "bnb_4bit_use_double_quant": False,
            "bnb_4bit_compute_dtype": torch.bfloat16 if BFLOAT16_SUPPORT else "float16",
        },
    },
    "nf4-dq": {
        "torch_dtype": "bfloat16" if BFLOAT16_SUPPORT else "float16",
        "quantization_scheme": "bnb",
        "quantization_config": {
            "load_in_4bit": True,
            "bnb_4bit_quant_type": "nf4",
            "bnb_4bit_use_double_quant": True,
            "bnb_4bit_compute_dtype": torch.bfloat16 if BFLOAT16_SUPPORT else "float16",
        },
    },
    "int8-decomp": {
        "torch_dtype": "float16",
        "quantization_scheme": "bnb",
        "quantization_config": {
            "load_in_8bit": True,
            "llm_int8_threshold": 6.0,
        },
    },
    "int8": {
        "torch_dtype": "float16",
        "quantization_scheme": "bnb",
        "quantization_config": {
            "load_in_8bit": True,
            "llm_int8_threshold": 0.0,
        },
    },
}

def parse_args():
    parser = argparse.ArgumentParser(description="bitsandbytes inference benchmark tool")

    parser.add_argument("model_id", type=str, help="The model checkpoint to use.")

    parser.add_argument(
        "--configs",
        nargs="+",
        choices=["bf16", "fp16", "nf4", "nf4-dq", "int8", "int8-decomp"],
        default=["nf4", "int8", "int8-decomp"],
    )
    parser.add_argument("--bf16", dest="configs", action="append_const", const="bf16")
    parser.add_argument("--fp16", dest="configs", action="append_const", const="fp16")
    parser.add_argument("--nf4", dest="configs", action="append_const", const="nf4")
    parser.add_argument("--nf4-dq", dest="configs", action="append_const", const="nf4-dq")
    parser.add_argument("--int8", dest="configs", action="append_const", const="int8")
    parser.add_argument("--int8-decomp", dest="configs", action="append_const", const="int8-decomp")

    parser.add_argument("--batches", nargs="+", type=int, default=[1, 8, 16, 32])
    parser.add_argument("--input-length", type=int, default=64)

    parser.add_argument("--out-dir", type=str, default="reports")

    parser.add_argument("--iterations", type=int, default=10, help="Number of iterations for each benchmark run")
    parser.add_argument(
        "--warmup-runs", type=int, default=10, help="Number of warmup runs to discard before measurement"
    )
    parser.add_argument(
        "--output-length",
        type=int,
        default=64,
        help="If set, `max_new_tokens` and `min_new_tokens` will be set to this value.",
    )

    return parser.parse_args()

def run_benchmark(args, config, batch_size):
    launcher_config = ProcessConfig(device_isolation=True, device_isolation_action="warn", start_method="spawn")
    scenario_config = InferenceConfig(
        latency=True,
        memory=True,
        input_shapes={"batch_size": batch_size, "sequence_length": args.input_length},
        iterations=args.iterations,
        warmup_runs=args.warmup_runs,
        # set duration to 0 to disable the duration-based stopping criterion
        # this is IMPORTANT to ensure that all benchmarks run the same number of operations, regardless of hardware speed/bottlenecks
        duration=0,
        # for consistent results, set a fixed min and max for output tokens
        generate_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length},
        forward_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length},
    )

    backend_config = PyTorchConfig(
        device="cuda",
        device_ids="0",
        device_map="auto",
        no_weights=False,
        model=args.model_id,
        **WEIGHTS_CONFIGS[config],
    )

    test_name = (
        f"benchmark-{config}"
        f"-bsz-{batch_size}"
        f"-isz-{args.input_length}"
        f"-osz-{args.output_length}"
        f"-iter-{args.iterations}"
        f"-wrmup-{args.warmup_runs}"
    )
    benchmark_config = BenchmarkConfig(
        name=test_name,
        scenario=scenario_config,
        launcher=launcher_config,
        backend=backend_config,
    )

    out_path = out_dir / (test_name + ".json")
    print(f"[{test_name}] Starting:")
    benchmark_report = Benchmark.launch(benchmark_config)
    benchmark_report.save_json(out_path)

if __name__ == "__main__":
    setup_logging(level="INFO")
    args = parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    for batch_size in args.batches:
        for config in args.configs:
            run_benchmark(args, config, batch_size)


    },
    "int8-decomp": {
        "torch_dtype": "float16",
        "quantization_scheme": "bnb",
        "quantization_config": {
            "load_in_8bit": True,
            "llm_int8_threshold": 6.0,
        },
    },
    "int8": {
        "torch_dtype": "float16",
        "quantization_scheme": "bnb",
        "quantization_config": {
            "load_in_8bit": True,
            "llm_int8_threshold": 0.0,
        },
    },
}

def parse_args():
    parser = argparse.ArgumentParser(description="bitsandbytes inference benchmark tool")

    parser.add_argument("model_id", type=str, help="The model checkpoint to use.")

    parser.add_argument(
        "--configs",
        nargs="+",
        choices=["bf16", "fp16", "nf4", "nf4-dq", "int8", "int8-decomp"],
        default=["nf4", "int8", "int8-decomp"],
    )
    parser.add_argument("--bf16", dest="configs", action="append_const", const="bf16")
    parser.add_argument("--fp16", dest="configs", action="append_const", const="fp16")
    parser.add_argument("--nf4", dest="configs", action="append_const", const="nf4")
    parser.add_argument("--nf4-dq", dest="configs", action="append_const", const="nf4-dq")
    parser.add_argument("--int8", dest="configs", action="append_const", const="int8")
    parser.add_argument("--int8-decomp", dest="configs", action="append_const", const="int8-decomp")

    parser.add_argument("--batches", nargs="+", type=int, default=[1, 8, 16, 32])
    parser.add_argument("--input-length", type=int, default=64)

    parser.add_argument("--out-dir", type=str, default="reports")

    parser.add_argument("--iterations", type=int, default=10, help="Number of iterations for each benchmark run")
    parser.add_argument(
        "--warmup-runs", type=int, default=10, help="Number of warmup runs to discard before measurement"
    )
    parser.add_argument(
        "--output-length",
        type=int,
        default=64,
        help="If set, `max_new_tokens` and `min_new_tokens` will be set to this value.",
    )

    return parser.parse_args()

def run_benchmark(args, config, batch_size):
    launcher_config = ProcessConfig(device_isolation=True, device_isolation_action="warn", start_method="spawn")
    scenario_config = InferenceConfig(
        latency=True,
        memory=True,
        input_shapes={"batch_size": batch_size, "sequence_length": args.input_length},
        iterations=args.iterations,
        warmup_runs=args.warmup_runs,
        # set duration to 0 to disable the duration-based stopping criterion
        # this is IMPORTANT to ensure that all benchmarks run the same number of operations, regardless of hardware speed/bottlenecks
        duration=0,
        # for consistent results, set a fixed min and max for output tokens
        generate_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length},
        forward_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length},
    )

    backend_config = PyTorchConfig(
        device="cuda",


    parser.add_argument("--batches", nargs="+", type=int, default=[1, 8, 16, 32])
    parser.add_argument("--input-length", type=int, default=64)

    parser.add_argument("--out-dir", type=str, default="reports")

    parser.add_argument("--iterations", type=int, default=10, help="Number of iterations for each benchmark run")
    parser.add_argument(
        "--warmup-runs", type=int, default=10, help="Number of warmup runs to discard before measurement"
    )
    parser.add_argument(
        "--output-length",
        type=int,
        default=64,
        help="If set, `max_new_tokens` and `min_new_tokens` will be set to this value.",
    )

    return parser.parse_args()

def run_benchmark(args, config, batch_size):
    launcher_config = ProcessConfig(device_isolation=True, device_isolation_action="warn", start_method="spawn")
    scenario_config = InferenceConfig(
        latency=True,
        memory=True,
        input_shapes={"batch_size": batch_size, "sequence_length": args.input_length},
        iterations=args.iterations,
        warmup_runs=args.warmup_runs,
        # set duration to 0 to disable the duration-based stopping criterion
        # this is IMPORTANT to ensure that all benchmarks run the same number of operations, regardless of hardware speed/bottlenecks
        duration=0,
        # for consistent results, set a fixed min and max for output tokens
        generate_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length},
        forward_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length},
    )

    backend_config = PyTorchConfig(
        device="cuda",
        device_ids="0",
        device_map="auto",
        no_weights=False,
        model=args.model_id,
        **WEIGHTS_CONFIGS[config],
    )

    test_name = (
        f"benchmark-{config}"
        f"-bsz-{batch_size}"
        f"-isz-{args.input_length}"
        f"-osz-{args.output_length}"
        f"-iter-{args.iterations}"
        f"-wrmup-{args.warmup_runs}"
    )
    benchmark_config = BenchmarkConfig(
        name=test_name,
        scenario=scenario_config,
        launcher=launcher_config,
        backend=backend_config,
    )

    out_path = out_dir / (test_name + ".json")
    print(f"[{test_name}] Starting:")
    benchmark_report = Benchmark.launch(benchmark_config)
    benchmark_report.save_json(out_path)

if __name__ == "__main__":
    setup_logging(level="INFO")
    args = parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    for batch_size in args.batches:
        for config in args.configs:
            run_benchmark(args, config, batch_size)


"""
Extracted from tests/test_optim.py

Usage: pytest benchmarking/optimizer_benchmark.py
"""

import time

import pytest
from tests.helpers import describe_dtype, id_formatter
import torch

import bitsandbytes as bnb

str2optimizers = {"paged_adamw": (torch.optim.AdamW, bnb.optim.PagedAdamW)}

@pytest.mark.parametrize("dim1", [2 * 1024], ids=id_formatter("dim1"))
@pytest.mark.parametrize("gtype", [torch.float16], ids=describe_dtype)
@pytest.mark.parametrize("optim_name", ["paged_adamw"], ids=id_formatter("optim_name"))
@pytest.mark.parametrize("mode", ["bnb"], ids=id_formatter("mode"))
@pytest.mark.benchmark
def test_stream_optimizer_bench(dim1, gtype, optim_name, mode):
    layers1 = torch.nn.Sequential(*torch.nn.ModuleList([torch.nn.Linear(dim1, dim1) for i in range(10)]))
    layers1 = layers1.to(gtype)
    layers1 = layers1.cuda()

    large_tensor = None
    if mode == "torch":
        optim = str2optimizers[optim_name][0](layers1.parameters())
    else:
        optim = str2optimizers[optim_name][1](layers1.parameters())
        # 12 GB
        large_tensor = torch.empty((int(4.5e9),), device="cuda")

    torch.cuda.synchronize()
    time.sleep(5)

    num_batches = 5
    batches = torch.randn(num_batches, 128, dim1, device="cuda").to(gtype)
    lbls = torch.randint(0, 10, size=(num_batches, 128)).cuda()

    for i in range(num_batches):
        print(i)
        b = batches[i]
        if i == 2:
            torch.cuda.synchronize()
            t0 = time.time()

        out1 = layers1(b)

        loss1 = torch.nn.functional.cross_entropy(out1, lbls[i]).mean()
        loss1.backward()
        optim.step()
    torch.cuda.synchronize()
    print(mode, time.time() - t0)


Usage: pytest benchmarking/optimizer_benchmark.py
"""

import time

import pytest
from tests.helpers import describe_dtype, id_formatter
import torch

import bitsandbytes as bnb

str2optimizers = {"paged_adamw": (torch.optim.AdamW, bnb.optim.PagedAdamW)}

@pytest.mark.parametrize("dim1", [2 * 1024], ids=id_formatter("dim1"))
@pytest.mark.parametrize("gtype", [torch.float16], ids=describe_dtype)
@pytest.mark.parametrize("optim_name", ["paged_adamw"], ids=id_formatter("optim_name"))
@pytest.mark.parametrize("mode", ["bnb"], ids=id_formatter("mode"))
@pytest.mark.benchmark
def test_stream_optimizer_bench(dim1, gtype, optim_name, mode):
    layers1 = torch.nn.Sequential(*torch.nn.ModuleList([torch.nn.Linear(dim1, dim1) for i in range(10)]))
    layers1 = layers1.to(gtype)
    layers1 = layers1.cuda()

    large_tensor = None
    if mode == "torch":
        optim = str2optimizers[optim_name][0](layers1.parameters())
    else:
        optim = str2optimizers[optim_name][1](layers1.parameters())
        # 12 GB
        large_tensor = torch.empty((int(4.5e9),), device="cuda")

    torch.cuda.synchronize()
    time.sleep(5)

    num_batches = 5
    batches = torch.randn(num_batches, 128, dim1, device="cuda").to(gtype)
    lbls = torch.randint(0, 10, size=(num_batches, 128)).cuda()

    for i in range(num_batches):
        print(i)
        b = batches[i]
        if i == 2:
            torch.cuda.synchronize()
            t0 = time.time()

        out1 = layers1(b)

        loss1 = torch.nn.functional.cross_entropy(out1, lbls[i]).mean()
        loss1.backward()
        optim.step()
    torch.cuda.synchronize()
    print(mode, time.time() - t0)


"""
Extracted from tests/test_functional.py

Usage: pytest benchmarking/int8/training_benchmark.py
"""

import time

import pytest
import torch

from bitsandbytes import functional as F

k = 20

torch.set_printoptions(precision=5, sci_mode=False, linewidth=120, edgeitems=20, threshold=10000)

@pytest.mark.parametrize(
    ("batch", "seq", "model", "hidden"),
    [
        pytest.param(2, 512, 4 * 1024, 3 * 4 * 1024, id="batch=2, seq=512, model=4k, hidden=12k"),
        pytest.param(2, 512, 5120, 3 * 5120, id="batch=2, seq=512, model=5k, hidden=15k"),
        pytest.param(2, 512, 12 * 1024, 4 * 12 * 1024, id="batch=2, seq=512, model=12k, hidden=48k"),
    ],
)
@pytest.mark.benchmark
def test_bench_8bit_training(batch, seq, model, hidden):
    formatB = F.get_special_format_str()
    A = torch.randn(batch, seq, model, device="cuda").half()
    grad = torch.randn(batch, seq, model, device="cuda").half()
    w1 = torch.randint(-128, 127, size=(hidden, model), device="cuda").half()
    w2 = torch.randint(-128, 127, size=(model, hidden), device="cuda").half()
    print("")

    # torch.cuda.synchronize()
    ## warmup
    # for i in range(100):
    #    torch.matmul(A, w1.t())
    # torch.cuda.synchronize()

    dtype = torch.int8
    A = A.view(-1, A.shape[-1]).contiguous()
    grad = grad.view(-1, grad.shape[-1]).contiguous()
    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(k):
        out1 = torch.matmul(A, w1.t())  # fc1
        # out2 = torch.matmul(out1, w2.t())# fc2

        # d1 = torch.matmul(grad, w2) # delta1
        # d2 = torch.matmul(d1, w1) # delta2

        # grad1 = torch.einsum('bo,bh->oh', out1, grad) # grad w2
        # grad2 = torch.einsum('bh,bo->ho', A, d2) # grad w1

    torch.cuda.synchronize()
    t16 = time.time() - t0
    print(t16)

    # torch.cuda.empty_cache()

    # Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
    # Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2)

    # CTw1, Sw1 = F.transform2(Cw1, formatB)
    # CTw2, Sw2 = F.transform2(Cw2, formatB)
    # CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True)
    # CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True)

    # CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
    # C32A, SA = F.transform2(CA, 'col32')
    ## fc1
    # out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1, dtype=dtype)
    ##out1 = F.mm_dequant(out1_32, Sout1_32, statsAt, statsw1t)

    ## fc2
    # Cout1, Cout1t, statsout1, statsout1t, coo_tensor = F.double_quant(out1)
    # C32out1, Sout1 = F.transform2(Cout1, 'col32')
    # out2_32, Sout2_32 = F.igemmlt(C32out1, CTw2, Sout1, Sw2, dtype=dtype)
    ##out2 = F.mm_dequant(out2_32, Sout2_32, statsout1t, statsw2t)

    ## delta1
    # Cgrad, Cgradt, statsgrad, statsgradt, coo_tensor = F.double_quant(grad)
    # C32grad, Sgrad = F.transform2(Cgrad, 'col32')
    ##d1_32, Sd1_32 = F.igemmlt(C32grad, CTw2t, Sgrad, Sw2t, dtype=dtype)
    ##d1 = F.mm_dequant(d1_32, Sd1_32, statsgradt, statsw2)

    ## delta2
    # Cd1, Cd1t, statsd1, statsd1t, coo_tensor = F.double_quant(d1)
    # C32d1, Sd1 = F.transform2(Cd1, 'col32')
    ##d2_32, Sd2_32 = F.igemmlt(C32d1, CTw1t, Sd1, Sw1t, dtype=dtype)
    ##d2 = F.mm_dequant(d2_32, Sd2_32, statsd1t, statsw1)

    ## grad1
    # C32out1t, Sout1t = F.transform2(Cout1t, 'col32', transpose=True)
    # CTgradt, Sgradt = F.transform2(Cgradt, formatB, transpose=True)
    ##grad1_32, Sgrad1_32 = F.igemmlt(C32out1t, CTgradt, Sout1t, Sgradt, dtype=dtype)
    ##grad1 = F.mm_dequant(grad1_32, Sgrad1_32, statsout1, statsgrad)

    ## grad2
    # C32At, SAt = F.transform2(CAt, 'col32', transpose=True)
    # CTd1t, Sd1t = F.transform2(Cd1t, formatB, transpose=True)
    ##grad2_32, Sgrad2_32 = F.igemmlt(C32At, CTd1t, SAt, Sd1t, dtype=dtype)
    ##grad2 = F.mm_dequant(grad2_32, Sgrad2_32, statsA, statsd1)

    # Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2)

    # Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
    # Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2)

    # CTw1, Sw1 = F.transform2(Cw1, formatB)
    # CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True)
    # CTw2, Sw2 = F.transform2(Cw2, formatB)
    # CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True)
    # torch.cuda.synchronize()
    # t0 = time.time()
    # for i in range(k):
    #    #Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
    #    #CTw1, Sw1 = F.transform2(Cw1, formatB)
    #    #Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
    #    #CTw1, Sw1 = F.transform2(Cw1, formatB)

    #    #CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A, threshold=3.5)
    #    CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
    #    #CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True)
    #    #CTw2, Sw2 = F.transform2(Cw2, formatB)
    #    #CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True)

    #    C32A, SA = F.transform2(CA, 'col32')

    #    # fc1
    #    out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1, dtype=dtype)
    #    #out1dn = F.mm_dequant(out1_32, Sout1_32, statsA, statsw1)

    #    #print(coo_tensor.nnz)
    #    #out1sp = F.spmm_coo(coo_tensor, w1.t())
    #    #print(w1.t().shape)
    #    #out1 = out1dn + out1sp

    #    # fc2
    #    Cout1, Cout1t, statsout1, statsout1t, coo_tensor = F.double_quant(out1)
    #    C32out1, Sout1 = F.transform2(Cout1, 'col32')
    #    out2_32, Sout2_32 = F.igemmlt(C32out1, CTw2, Sout1, Sw2, dtype=dtype)
    #    #out2 = F.mm_dequant(out2_32, Sout2_32, statsout1, statsw2)

    #    # delta1
    #    Cgrad, Cgradt, statsgrad, statsgradt, coo_tensor = F.double_quant(grad)
    #    C32grad, Sgrad = F.transform2(Cgrad, 'col32')
    #    d1_32, Sd1_32 = F.igemmlt(C32grad, CTw2t, Sgrad, Sw2t, dtype=dtype)
    #    #d1 = F.mm_dequant(d1_32, Sd1_32, statsgrad, statsw2t)

    #    # delta2
    #    Cd1, Cd1t, statsd1, statsd1t, coo_tensor = F.double_quant(d1)
    #    C32d1, Sd1 = F.transform2(Cd1, 'col32')
    #    d2_32, Sd2_32 = F.igemmlt(C32d1, CTw1t, Sd1, Sw1t, dtype=dtype)
    #    #d2 = F.mm_dequant(d2_32, Sd2_32, statsd1, statsw1t)

    #    # grad1
    #    #C32out1t, Sout1t = F.transform2(Cout1t, 'col32', transpose=True)
    #    #CTgradt, Sgradt = F.transform2(Cgradt, formatB, transpose=True)
    #    #grad1_32, Sgrad1_32 = F.igemmlt(C32out1t, CTgradt, Sout1t, Sgradt, dtype=dtype)
    #    #grad1 = F.mm_dequant(grad1_32, Sgrad1_32, statsout1t, statsgradt)

    #    ## grad2
    #    #C32At, SAt = F.transform2(CAt, 'col32', transpose=True)
    #    #CTd1t, Sd1t = F.transform2(Cd1t, formatB, transpose=True)
    #    #grad2_32, Sgrad2_32 = F.igemmlt(C32At, CTd1t, SAt, Sd1t, dtype=dtype)
    #    #grad2 = F.mm_dequant(grad2_32, Sgrad2_32, statsAt, statsd1t)

    # torch.cuda.synchronize()
    # t8 = time.time() - t0
    # print(t8)


import pytest
import torch

from bitsandbytes import functional as F

k = 20

torch.set_printoptions(precision=5, sci_mode=False, linewidth=120, edgeitems=20, threshold=10000)

@pytest.mark.parametrize(
    ("batch", "seq", "model", "hidden"),
    [
        pytest.param(2, 512, 4 * 1024, 3 * 4 * 1024, id="batch=2, seq=512, model=4k, hidden=12k"),
        pytest.param(2, 512, 5120, 3 * 5120, id="batch=2, seq=512, model=5k, hidden=15k"),
        pytest.param(2, 512, 12 * 1024, 4 * 12 * 1024, id="batch=2, seq=512, model=12k, hidden=48k"),
    ],
)
@pytest.mark.benchmark
def test_bench_8bit_training(batch, seq, model, hidden):
    formatB = F.get_special_format_str()
    A = torch.randn(batch, seq, model, device="cuda").half()
    grad = torch.randn(batch, seq, model, device="cuda").half()
    w1 = torch.randint(-128, 127, size=(hidden, model), device="cuda").half()
    w2 = torch.randint(-128, 127, size=(model, hidden), device="cuda").half()
    print("")

    # torch.cuda.synchronize()
    ## warmup
    # for i in range(100):
    #    torch.matmul(A, w1.t())
    # torch.cuda.synchronize()

    dtype = torch.int8
    A = A.view(-1, A.shape[-1]).contiguous()
    grad = grad.view(-1, grad.shape[-1]).contiguous()
    torch.cuda.synchronize()
    t0 = time.time()
    for i in range(k):
        out1 = torch.matmul(A, w1.t())  # fc1
        # out2 = torch.matmul(out1, w2.t())# fc2

        # d1 = torch.matmul(grad, w2) # delta1
        # d2 = torch.matmul(d1, w1) # delta2

        # grad1 = torch.einsum('bo,bh->oh', out1, grad) # grad w2
        # grad2 = torch.einsum('bh,bo->ho', A, d2) # grad w1

    torch.cuda.synchronize()
    t16 = time.time() - t0
    print(t16)

    # torch.cuda.empty_cache()

    # Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1)
    # Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2)

    # CTw1, Sw1 = F.transform2(Cw1, formatB)
    # CTw2, Sw2 = F.transform2(Cw2, formatB)
    # CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True)
    # CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True)

    # CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A)
    # C32A, SA = F.transform2(CA, 'col32')
    ## fc1
    # out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1, dtype=dtype)
    ##out1 = F.mm_dequant(out1_32, Sout1_32, statsAt, statsw1t)

    ## fc2
    # Cout1, Cout1t, statsout1, statsout1t, coo_tensor = F.double_quant(out1)
    # C32out1, Sout1 = F.transform2(Cout1, 'col32')


"""
Basic benchmark for text generation.

Usage: python benchmarking/int8/int8_benchmark.py
"""

import time

import torch
from torch.profiler import ProfilerActivity, profile
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

MAX_NEW_TOKENS = 128
model_name = "meta-llama/Llama-3.1-8B"

text = "Below is a question. I need an answer.\n\nExplain machine learning: "
tokenizer = AutoTokenizer.from_pretrained(model_name)
input_ids = tokenizer([text] * 8, return_tensors="pt").input_ids.to(0)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    quantization_config=BitsAndBytesConfig(
        load_in_8bit=True,
        llm_int8_threshold=6.0,
    ),
    attn_implementation="sdpa",
    torch_dtype=torch.float16,
)

print(model)

# warmup
print("Warmup...")
for i in range(3):
    generated_ids = model.generate(input_ids, max_new_tokens=MAX_NEW_TOKENS)

print("Profiler starting...")
with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    with_modules=True,
    with_stack=True,
) as prof:
    model.generate(input_ids, max_new_tokens=1)

print(
    prof.key_averages().table(
        sort_by="cpu_time_total",
        max_name_column_width=50,
        top_level_events_only=True,
        row_limit=50,
    )
)

torch.cuda.synchronize()

print("Generating...")
num = 0
time_1 = time.time()
for i in range(5):
    generated_ids = model.generate(input_ids, max_new_tokens=MAX_NEW_TOKENS)
    num += len(generated_ids[0])

print("=" * 40)
print(f"Example:\n{tokenizer.decode(generated_ids[0])}")
print("=" * 40)
print(f"Speed: {num / (time.time() - time_1)}token/s")


import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import pandas as pd

cmap = plt.get_cmap("cool")

if __name__ == "__main__":
    fig = plt.figure(tight_layout=True, figsize=(12, 3.5))
    gs = gridspec.GridSpec(1, 2)

    dims_to_consider = [1024, 1280, 1408, 1664, 2048, 4096]
    batch_size_for_plot1 = 32768
    batch_sizes_for_plot2 = [2**14, 2**15, 2**16, 2**17]
    dims_to_xtick = [1024, 2048, 4096]
    logscale_plot1 = True

    ax = fig.add_subplot(gs[0, 0])

    # TODO: change this to what you want.
    rdf = pd.read_json("speed_benchmark/info_a100_py2.jsonl", lines=True)
    df = rdf[rdf.batch_size == batch_size_for_plot1]

    # first plot the time occupied by different operations
    for k, marker, ls, color, name in [
        ("standard_gx+standard_gw+standard_fwd", "s", "-", "C2", "Standard fp16 (sum of parts)"),
        (
            "x_quantize_rowwise+g_quantize_rowwise+w_quantize_global+w_quantize_global_transpose+standard_gw+global_fwd+global_bwd",
            "o",
            "-",
            "C4",
            "SwitchBack int8 (sum of parts)",
        ),
        ("standard_fwd", "^", "--", "C2", "Matmul XW (standard)"),
        ("standard_gw", "^", "-.", "C2", "Matmul GW (standard)"),
        ("standard_gx", "^", ":", "gray", "Matmul GX (both)"),
        ("global_fwd", "^", "--", "C4", "Int8 Matmul XW (switchback)"),
        ("global_bwd", "^", "-.", "C4", "Int8 Matmul GW (switchback)"),
        ("x_quantize_rowwise", "P", "--", "C4", "Quantize rowwise X (switchback)"),
        ("g_quantize_rowwise", "P", "-.", "C4", "Quantize rowwise G (switchback)"),
        ("w_quantize_global", ".", "--", "C4", "Quantize global W (switchback)"),
        ("w_quantize_global_transpose", ".", "-.", "C4", "Quantize global and\ntranspose W (switchback)"),
    ]:
        xs = []
        ys = []
        for embed_dim in dims_to_consider:
            # average over dim -> 4*dim and 4*dim -> dim
            df_ = df[df.dim_in == embed_dim]
            df_ = df_[df_.dim_out == embed_dim * 4]
            xs.append(embed_dim)
            y_ = 0
            for k_ in k.split("+"):
                y_ += df_[k_].values[0]
            df_ = df[df.dim_in == embed_dim * 4]
            df_ = df_[df_.dim_out == embed_dim]
            for k_ in k.split("+"):
                y_ += df_[k_].values[0]
            ys.append(y_ * 0.5)

        ax.plot(
            xs,
            ys,
            color=color,
            label=name,
            marker=marker,
            markersize=5 if marker == "s" else 5,
            linestyle=ls,
            linewidth=2 if "+" in k else 1.0,
        )

    ax.set_xlabel("dim", fontsize=13)
    ax.set_ylabel("time (ms)", fontsize=13)

    ax.grid()

    ax.set_xscale("log")
    if logscale_plot1:
        ax.set_yscale("log")

    ax.tick_params(axis="x", labelsize=11)
    ax.tick_params(axis="y", labelsize=11)

    ax.set_xticks(dims_to_xtick)
    ax.set_xticklabels(dims_to_xtick)
    ax.set_xticks([], minor=True)

    leg = ax.legend(loc="upper center", bbox_to_anchor=(-0.64, 1.0), ncol=1, fontsize=10)
    leg.get_texts()[0].set_fontweight("bold")
    leg.get_texts()[1].set_fontweight("bold")
    plt.subplots_adjust(left=0.1)
    ax.set_title("  Linear layer, batch * sequence length = 32k", fontsize=10, loc="left", y=1.05, pad=-20)

    ax = fig.add_subplot(gs[0, 1])

    # now plot the % speedup for different batch sizes
    for j, batch_size in enumerate(batch_sizes_for_plot2):
        all_xs, all_ys = [], []
        for k, marker, ls, color, name in [
            ("standard_gx+standard_gw+standard_fwd", "s", "-", "C2", "Standard fp16 (total time)"),
            (
                "x_quantize_rowwise+g_quantize_rowwise+w_quantize_global+w_quantize_global_transpose+standard_gw+global_fwd+global_bwd",
                "o",
                "-",
                "C4",
                "SwitchBack int8 (total time)",
            ),
        ]:
            xs, ys = [], []
            df = rdf[rdf.batch_size == batch_size]
            for embed_dim in dims_to_consider:
                df_ = df[df.dim_in == embed_dim]
                df_ = df_[df_.dim_out == embed_dim * 4]
                xs.append(embed_dim)
                y_ = 0
                for k_ in k.split("+"):
                    y_ += df_[k_].values[0]
                df_ = df[df.dim_in == embed_dim * 4]
                df_ = df_[df_.dim_out == embed_dim]
                for k_ in k.split("+"):
                    y_ += df_[k_].values[0]
                ys.append(y_ * 0.5)
            all_xs.append(xs)
            all_ys.append(ys)

        color = cmap(j * 0.25)
        real_ys = [-((all_ys[1][i] - all_ys[0][i]) / all_ys[0][i]) * 100 for i in range(len(all_ys[0]))]
        markers = ["^", "v", "P", "o"]
        ax.plot(
            all_xs[0],
            real_ys,
            color=color,
            label=f"batch * sequence length = {batch_size}",
            marker=markers[j],
            markersize=5 if marker == "s" else 5,
        )

    ax.legend()
    ax.set_xlabel("dim", fontsize=13)
    ax.set_xscale("log")
    ax.grid()
    ax.set_ylabel(r"% speedup", fontsize=13)

    ax.tick_params(axis="x", labelsize=11)
    ax.tick_params(axis="y", labelsize=11)

    ax.set_xticks(dims_to_xtick)
    ax.set_xticklabels(dims_to_xtick)
    ax.set_xticks([], minor=True)

    ax.set_title("  Linear layer summary, varying dimensions", fontsize=10, loc="left", y=1.05, pad=-20)

    plt.savefig("speed_benchmark/plot_with_info.pdf", bbox_inches="tight")


import json
import time

import torch

from bitsandbytes.triton.int8_matmul_mixed_dequantize import (
    int8_matmul_mixed_dequantize,
)
from bitsandbytes.triton.int8_matmul_rowwise_dequantize import (
    int8_matmul_rowwise_dequantize,
)
from bitsandbytes.triton.quantize_columnwise_and_transpose import (
    quantize_columnwise_and_transpose,
)
from bitsandbytes.triton.quantize_global import (
    quantize_global,
    quantize_global_transpose,
)
from bitsandbytes.triton.quantize_rowwise import quantize_rowwise

# KNOW ISSUE: need to optimize "w_quantize_colwise_transpose" when embeddim is too large.

def get_time(k, fn, info_dict):
    for _ in range(repeat // 2):
        fn()

    torch.cuda.synchronize()
    start = time.time()
    for _ in range(repeat):
        fn()

    torch.cuda.synchronize()
    end = time.time()
    ms = (end - start) / repeat * 1000
    print(f"time {k}: {ms:.3f} ms")
    info_dict[k] = ms

if __name__ == "__main__":
    torch.manual_seed(0)
    wm = 4
    for dim in [1024, 1280, 1408, 1664, 2048, 4096]:
        # note "batch_size" is actually "batch_size * embed_dim", which is why it's large
        for batch_size in [256 * 32, 256 * 64, 256 * 128, 256 * 256, 256 * 512]:
            # switch switches dim_in and dim_out
            for switch in [False, True]:
                # hparams
                repeat = 64
                batch_size = batch_size
                dim_out = dim * wm
                dim_in = dim
                if switch:
                    dim_out = dim
                    dim_in = wm * dim

                dim_in = round(dim_in)
                dim_out = round(dim_out)

                # simulate forward pass
                x = torch.randn(batch_size, dim_in, dtype=torch.float16).cuda()
                g = torch.randn(batch_size, dim_out, dtype=torch.float16).cuda()
                w = torch.randn(dim_out, dim_in, dtype=torch.float16).cuda()

                x_int8 = x.clone().to(torch.int8)
                g_int8 = g.clone().to(torch.int8)
                w_int8 = w.clone().to(torch.int8)
                wt_int8 = w.t().contiguous().clone().to(torch.int8)
                state_x_rowwise = x.max(dim=1)[0]
                state_g_rowwise = g.max(dim=1)[0]
                state_w_columnwise = w.max(dim=0)[0]
                state_w_rowwise = w.max(dim=1)[0]
                state_w_global = w.max()

                info = {
                    "repeat": repeat,
                    "batch_size": batch_size,
                    "dim_out": dim_out,
                    "dim_in": dim_in,
                    "wm": wm,
                    "switch": switch,
                }

                get_time("standard_fwd", lambda: x.matmul(w.t()), info)
                get_time("standard_gw", lambda: g.t().matmul(x), info)
                get_time("standard_gx", lambda: g.matmul(w), info)
                get_time(
                    "rowwise_fwd",
                    lambda: int8_matmul_rowwise_dequantize(
                        x_int8,
                        w_int8.t(),
                        state_x_rowwise,
                        state_w_columnwise,
                        None,
                    ),
                    info,
                )
                get_time(
                    "rowwise_bwd",
                    lambda: int8_matmul_rowwise_dequantize(
                        g_int8,
                        wt_int8.t(),
                        state_x_rowwise,
                        state_w_rowwise,
                        None,
                    ),
                    info,
                )
                get_time(
                    "global_fwd",
                    lambda: int8_matmul_mixed_dequantize(x_int8, w_int8.t(), state_x_rowwise, state_w_global, None),
                    info,
                )
                get_time(
                    "global_bwd",
                    lambda: int8_matmul_mixed_dequantize(g_int8, wt_int8.t(), state_x_rowwise, state_w_global, None),
                    info,
                )
                get_time("x_quantize_rowwise", lambda: quantize_rowwise(x), info)
                get_time("g_quantize_rowwise", lambda: quantize_rowwise(g), info)
                get_time("w_quantize_rowwise", lambda: quantize_rowwise(w), info)
                get_time("w_quantize_colwise_transpose", lambda: quantize_columnwise_and_transpose(w), info)
                get_time("w_quantize_global", lambda: quantize_global(w), info)
                get_time("w_quantize_global_transpose", lambda: quantize_global_transpose(w), info)

                time_standard = info["standard_fwd"] + info["standard_gx"] + info["standard_gw"]
                time_rowwise = (
                    info["x_quantize_rowwise"]
                    + info["g_quantize_rowwise"]
                    + info["w_quantize_colwise_transpose"]
                    + info["w_quantize_rowwise"]
                    + info["standard_gw"]
                    + info["rowwise_fwd"]
                    + info["rowwise_bwd"]
                )
                time_global = (
                    info["x_quantize_rowwise"]
                    + info["g_quantize_rowwise"]
                    + info["w_quantize_global"]
                    + info["w_quantize_global_transpose"]
                    + info["standard_gw"]
                    + info["global_fwd"]
                    + info["global_bwd"]
                )

                print("TOTAL STANDARD", time_standard)
                print("TOTAL ROWWISE", time_rowwise)
                print("TOTAL GLOBAL", time_global)

                print("speedup", -100 * (time_global - time_standard) / time_standard)

                info["time_standard"] = time_standard
                info["time_rowwise"] = time_rowwise
                info["time_global"] = time_global

                info_json = json.dumps(info)

                # TODO: change this to what you want.
                with open("speed_benchmark/info.jsonl", "a") as file:
                    file.write(info_json + "\n")


import torch

from bitsandbytes.triton.int8_matmul_mixed_dequantize import (
    int8_matmul_mixed_dequantize,
)
from bitsandbytes.triton.int8_matmul_rowwise_dequantize import (
    int8_matmul_rowwise_dequantize,
)
from bitsandbytes.triton.quantize_columnwise_and_transpose import (
    quantize_columnwise_and_transpose,
)
from bitsandbytes.triton.quantize_global import (
    quantize_global,
    quantize_global_transpose,
)
from bitsandbytes.triton.quantize_rowwise import quantize_rowwise

# KNOW ISSUE: need to optimize "w_quantize_colwise_transpose" when embeddim is too large.

def get_time(k, fn, info_dict):
    for _ in range(repeat // 2):
        fn()

    torch.cuda.synchronize()
    start = time.time()
    for _ in range(repeat):
        fn()

    torch.cuda.synchronize()
    end = time.time()
    ms = (end - start) / repeat * 1000
    print(f"time {k}: {ms:.3f} ms")
    info_dict[k] = ms

if __name__ == "__main__":
    torch.manual_seed(0)
    wm = 4
    for dim in [1024, 1280, 1408, 1664, 2048, 4096]:
        # note "batch_size" is actually "batch_size * embed_dim", which is why it's large
        for batch_size in [256 * 32, 256 * 64, 256 * 128, 256 * 256, 256 * 512]:
            # switch switches dim_in and dim_out
            for switch in [False, True]:
                # hparams
                repeat = 64
                batch_size = batch_size
                dim_out = dim * wm
                dim_in = dim
                if switch:
                    dim_out = dim
                    dim_in = wm * dim

                dim_in = round(dim_in)


import argparse
import time

# import intel_extension_for_pytorch as ipex
import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig

MAX_NEW_TOKENS = 256

get_time = time.time

system_prompt = "You are a helpful assistant"
user_prompt = """Summarize this text please:

```Tell me, O muse, of that ingenious hero who travelled far and wide after he had sacked the famous town of Troy. Many cities did he visit, and many were the nations with whose manners and customs he was acquainted; moreover he suffered much by sea while trying to save his own life and bring his men safely home; but do what he might he could not save his men, for they perished through their own sheer folly in eating the cattle of the Sun-god Hyperion; so the god prevented them from ever reaching home. Tell me, too, about all these things, O daughter of Jove, from whatsoever source you may know them.

So now all who escaped death in battle or by shipwreck had got safely home except Ulysses, and he, though he was longing to return to his wife and country, was detained by the goddess Calypso, who had got him into a large cave and wanted to marry him. But as years went by, there came a time when the gods settled that he should go back to Ithaca; even then, however, when he was among his own people, his troubles were not yet over; nevertheless all the gods had now begun to pity him except Neptune, who still persecuted him without ceasing and would not let him get home.

Now Neptune had gone off to the Ethiopians, who are at the world's end, and lie in two halves, the one looking West and the other East. He had gone there to accept a hecatomb of sheep and oxen, and was enjoying himself at his festival; but the other gods met in the house of Olympian Jove, and the sire of gods and men spoke first. At that moment he was thinking of Aegisthus, who had been killed by Agamemnon's son Orestes; so he said to the other gods:

"See now, how men lay blame upon us gods for what is after all nothing but their own folly. Look at Aegisthus; he must needs make love to Agamemnon's wife unrighteously and then kill Agamemnon, though he knew it would be the death of him; for I sent Mercury to warn him not to do either of these things, inasmuch as Orestes would be sure to take his revenge when he grew up and wanted to return home. Mercury told him this in all good will but he would not listen, and now he has paid for everything in full."

Then Minerva said, "Father, son of Saturn, King of kings, it served Aegisthus right, and so it would any one else who does as he did; but Aegisthus is neither here nor there; it is for Ulysses that my heart bleeds, when I think of his sufferings in that lonely sea-girt island, far away, poor man, from all his friends. It is an island covered with forest, in the very middle of the sea, and a goddess lives there, daughter of the magician Atlas, who looks after the bottom of the ocean, and carries the great columns that keep heaven and earth asunder. This daughter of Atlas has got hold of poor unhappy Ulysses, and keeps trying by every kind of blandishment to make him forget his home, so that he is tired of life, and thinks of nothing but how he may once more see the smoke of his own chimneys. You, sir, take no heed of this, and yet when Ulysses was before Troy did he not propitiate you with many a burnt sacrifice? Why then should you keep on being so angry with him?"

And Jove said, "My child, what are you talking about? How can I forget Ulysses than whom there is no more capable man on earth, nor more liberal in his offerings to the immortal gods that live in heaven? Bear in mind, however, that Neptune is still furious with Ulysses for having blinded an eye of Polyphemus king of the Cyclopes. Polyphemus is son to Neptune by the nymph Thoosa, daughter to the sea-king Phorcys; therefore though he will not kill Ulysses outright, he torments him by preventing him from getting home. Still, let us lay our heads together and see how we can help him to return; Neptune will then be pacified, for if we are all of a mind he can hardly stand out against us."```"""

prompt = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": user_prompt},
]

def get_inputs(tokenizer):
    inputs = tokenizer.apply_chat_template(
        prompt,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    return inputs

def get_streamer(tokenizer):
    streamer = Streamer(tokenizer)
    # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    return streamer

class Streamer:
    def __init__(self, tokenizer, print_median=False):
        self.times = []
        self.print_median = print_median
        self.tokenizer = tokenizer

    def put(self, t):
        self.times.append(get_time())
        if len(self.times) > 1:
            print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms")

        if len(self.times) % 10 == 3 and self.print_median:
            ts = np.array(self.times)
            diff = ts[1:] - ts[:-1]
            # print("Token latency:", 1000 * diff, "ms")
            print("Token latency median:", np.median(1000 * diff), "ms")

    def print_report(self):
        times = np.array(self.times)
        diff = times[1:] - times[:-1]
        print(f"Median latency: {round(np.median(diff) * 1000, 2)}ms")
        percentiles = [10, 25, 50, 75, 90]
        print(
            "Latency percentiles",
            {p: round(1000 * float(np.percentile(diff, p)), 1) for p in percentiles},
        )

    def end(self, *args):
        pass

def parse_arguments():
    parser = argparse.ArgumentParser(description="Run inference benchmark for LLM models")
    parser.add_argument(
        "--device",
        type=str,
        default="xpu",
        help="Device to run inference on (e.g., xpu, cuda, cpu)",
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
        help="Model ID from Hugging Face or local path",
    )
    parser.add_argument(
        "--attn",
        type=str,
        default="eager",
        choices=["eager", "flash_attention", "sdpa"],
        help="Attention implementation to use",
    )
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_arguments()

    device = args.device
    model_id = args.model_id

    print(f"Running inference on {device} with model {model_id}")
    print(f"Using attention implementation: {args.attn}")

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=args.attn)

    inputs = get_inputs(tokenizer)
    streamer = get_streamer(tokenizer)

    inputs = inputs.to(device)
    model = model.to(device)

    generation_config = GenerationConfig(
        use_cache=True,
        forced_eos_token_id=1,
        eos_token_id=1,
        max_new_tokens=MAX_NEW_TOKENS,
        do_sample=False,
    )

    outputs = model.generate(
        **inputs,
        streamer=streamer,
        generation_config=generation_config,
    )

    # Print the final outputs (including the input prompt)
    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    print(r"\Output (including prompt):")
    print("-" * 40)
    print(output_text)
    print("-" * 40)
    print(f"Peak memory usage: {torch.xpu.max_memory_allocated() / 1024**2:.0f}MB")

    streamer.print_report()


user_prompt = """Summarize this text please:

```Tell me, O muse, of that ingenious hero who travelled far and wide after he had sacked the famous town of Troy. Many cities did he visit, and many were the nations with whose manners and customs he was acquainted; moreover he suffered much by sea while trying to save his own life and bring his men safely home; but do what he might he could not save his men, for they perished through their own sheer folly in eating the cattle of the Sun-god Hyperion; so the god prevented them from ever reaching home. Tell me, too, about all these things, O daughter of Jove, from whatsoever source you may know them.

So now all who escaped death in battle or by shipwreck had got safely home except Ulysses, and he, though he was longing to return to his wife and country, was detained by the goddess Calypso, who had got him into a large cave and wanted to marry him. But as years went by, there came a time when the gods settled that he should go back to Ithaca; even then, however, when he was among his own people, his troubles were not yet over; nevertheless all the gods had now begun to pity him except Neptune, who still persecuted him without ceasing and would not let him get home.

Now Neptune had gone off to the Ethiopians, who are at the world's end, and lie in two halves, the one looking West and the other East. He had gone there to accept a hecatomb of sheep and oxen, and was enjoying himself at his festival; but the other gods met in the house of Olympian Jove, and the sire of gods and men spoke first. At that moment he was thinking of Aegisthus, who had been killed by Agamemnon's son Orestes; so he said to the other gods:

"See now, how men lay blame upon us gods for what is after all nothing but their own folly. Look at Aegisthus; he must needs make love to Agamemnon's wife unrighteously and then kill Agamemnon, though he knew it would be the death of him; for I sent Mercury to warn him not to do either of these things, inasmuch as Orestes would be sure to take his revenge when he grew up and wanted to return home. Mercury told him this in all good will but he would not listen, and now he has paid for everything in full."

Then Minerva said, "Father, son of Saturn, King of kings, it served Aegisthus right, and so it would any one else who does as he did; but Aegisthus is neither here nor there; it is for Ulysses that my heart bleeds, when I think of his sufferings in that lonely sea-girt island, far away, poor man, from all his friends. It is an island covered with forest, in the very middle of the sea, and a goddess lives there, daughter of the magician Atlas, who looks after the bottom of the ocean, and carries the great columns that keep heaven and earth asunder. This daughter of Atlas has got hold of poor unhappy Ulysses, and keeps trying by every kind of blandishment to make him forget his home, so that he is tired of life, and thinks of nothing but how he may once more see the smoke of his own chimneys. You, sir, take no heed of this, and yet when Ulysses was before Troy did he not propitiate you with many a burnt sacrifice? Why then should you keep on being so angry with him?"

And Jove said, "My child, what are you talking about? How can I forget Ulysses than whom there is no more capable man on earth, nor more liberal in his offerings to the immortal gods that live in heaven? Bear in mind, however, that Neptune is still furious with Ulysses for having blinded an eye of Polyphemus king of the Cyclopes. Polyphemus is son to Neptune by the nymph Thoosa, daughter to the sea-king Phorcys; therefore though he will not kill Ulysses outright, he torments him by preventing him from getting home. Still, let us lay our heads together and see how we can help him to return; Neptune will then be pacified, for if we are all of a mind he can hardly stand out against us."```"""

prompt = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": user_prompt},
]

def get_inputs(tokenizer):
    inputs = tokenizer.apply_chat_template(
        prompt,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    return inputs

def get_streamer(tokenizer):
    streamer = Streamer(tokenizer)
    # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    return streamer

class Streamer:
    def __init__(self, tokenizer, print_median=False):
        self.times = []
        self.print_median = print_median
        self.tokenizer = tokenizer

    def put(self, t):
        self.times.append(get_time())
        if len(self.times) > 1:
            print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms")

        if len(self.times) % 10 == 3 and self.print_median:


And Jove said, "My child, what are you talking about? How can I forget Ulysses than whom there is no more capable man on earth, nor more liberal in his offerings to the immortal gods that live in heaven? Bear in mind, however, that Neptune is still furious with Ulysses for having blinded an eye of Polyphemus king of the Cyclopes. Polyphemus is son to Neptune by the nymph Thoosa, daughter to the sea-king Phorcys; therefore though he will not kill Ulysses outright, he torments him by preventing him from getting home. Still, let us lay our heads together and see how we can help him to return; Neptune will then be pacified, for if we are all of a mind he can hardly stand out against us."```"""

prompt = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": user_prompt},
]

def get_inputs(tokenizer):
    inputs = tokenizer.apply_chat_template(
        prompt,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    return inputs

def get_streamer(tokenizer):
    streamer = Streamer(tokenizer)
    # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    return streamer

class Streamer:
    def __init__(self, tokenizer, print_median=False):
        self.times = []
        self.print_median = print_median
        self.tokenizer = tokenizer

    def put(self, t):
        self.times.append(get_time())
        if len(self.times) > 1:
            print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms")

        if len(self.times) % 10 == 3 and self.print_median:
            ts = np.array(self.times)
            diff = ts[1:] - ts[:-1]
            # print("Token latency:", 1000 * diff, "ms")
            print("Token latency median:", np.median(1000 * diff), "ms")

    def print_report(self):


Streamer

]

def get_inputs(tokenizer):
    inputs = tokenizer.apply_chat_template(
        prompt,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    return inputs

def get_streamer(tokenizer):
    streamer = Streamer(tokenizer)
    # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    return streamer

class Streamer:
    def __init__(self, tokenizer, print_median=False):
        self.times = []
        self.print_median = print_median
        self.tokenizer = tokenizer

    def put(self, t):
        self.times.append(get_time())
        if len(self.times) > 1:
            print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms")

        if len(self.times) % 10 == 3 and self.print_median:
            ts = np.array(self.times)
            diff = ts[1:] - ts[:-1]
            # print("Token latency:", 1000 * diff, "ms")
            print("Token latency median:", np.median(1000 * diff), "ms")

    def print_report(self):
        times = np.array(self.times)
        diff = times[1:] - times[:-1]
        print(f"Median latency: {round(np.median(diff) * 1000, 2)}ms")
        percentiles = [10, 25, 50, 75, 90]
        print(
            "Latency percentiles",
            {p: round(1000 * float(np.percentile(diff, p)), 1) for p in percentiles},
        )

    def end(self, *args):
        pass

def parse_arguments():
    parser = argparse.ArgumentParser(description="Run inference benchmark for LLM models")
    parser.add_argument(
        "--device",
        type=str,
        default="xpu",
        help="Device to run inference on (e.g., xpu, cuda, cpu)",
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
        help="Model ID from Hugging Face or local path",
    )
    parser.add_argument(
        "--attn",
        type=str,
        default="eager",


        if len(self.times) % 10 == 3 and self.print_median:
            ts = np.array(self.times)
            diff = ts[1:] - ts[:-1]
            # print("Token latency:", 1000 * diff, "ms")
            print("Token latency median:", np.median(1000 * diff), "ms")

    def print_report(self):
        times = np.array(self.times)
        diff = times[1:] - times[:-1]
        print(f"Median latency: {round(np.median(diff) * 1000, 2)}ms")
        percentiles = [10, 25, 50, 75, 90]
        print(
            "Latency percentiles",
            {p: round(1000 * float(np.percentile(diff, p)), 1) for p in percentiles},
        )

    def end(self, *args):
        pass

def parse_arguments():
    parser = argparse.ArgumentParser(description="Run inference benchmark for LLM models")
    parser.add_argument(
        "--device",
        type=str,
        default="xpu",
        help="Device to run inference on (e.g., xpu, cuda, cpu)",
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
        help="Model ID from Hugging Face or local path",
    )
    parser.add_argument(
        "--attn",
        type=str,
        default="eager",
        choices=["eager", "flash_attention", "sdpa"],
        help="Attention implementation to use",
    )
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_arguments()

    device = args.device
    model_id = args.model_id

    print(f"Running inference on {device} with model {model_id}")
    print(f"Using attention implementation: {args.attn}")

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=args.attn)

    inputs = get_inputs(tokenizer)
    streamer = get_streamer(tokenizer)

    inputs = inputs.to(device)
    model = model.to(device)


def get_inputs(tokenizer):
    inputs = tokenizer.apply_chat_template(
        prompt,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    return inputs

def get_streamer(tokenizer):
    streamer = Streamer(tokenizer)
    # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    return streamer

class Streamer:
    def __init__(self, tokenizer, print_median=False):
        self.times = []
        self.print_median = print_median
        self.tokenizer = tokenizer

    def put(self, t):
        self.times.append(get_time())
        if len(self.times) > 1:
            print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms")

        if len(self.times) % 10 == 3 and self.print_median:
            ts = np.array(self.times)
            diff = ts[1:] - ts[:-1]
            # print("Token latency:", 1000 * diff, "ms")
            print("Token latency median:", np.median(1000 * diff), "ms")

    def print_report(self):
        times = np.array(self.times)
        diff = times[1:] - times[:-1]
        print(f"Median latency: {round(np.median(diff) * 1000, 2)}ms")
        percentiles = [10, 25, 50, 75, 90]
        print(
            "Latency percentiles",
            {p: round(1000 * float(np.percentile(diff, p)), 1) for p in percentiles},


put

        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    return inputs

def get_streamer(tokenizer):
    streamer = Streamer(tokenizer)
    # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    return streamer

class Streamer:
    def __init__(self, tokenizer, print_median=False):
        self.times = []
        self.print_median = print_median
        self.tokenizer = tokenizer

    def put(self, t):
        self.times.append(get_time())
        if len(self.times) > 1:
            print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms")

        if len(self.times) % 10 == 3 and self.print_median:
            ts = np.array(self.times)
            diff = ts[1:] - ts[:-1]
            # print("Token latency:", 1000 * diff, "ms")
            print("Token latency median:", np.median(1000 * diff), "ms")

    def print_report(self):
        times = np.array(self.times)
        diff = times[1:] - times[:-1]
        print(f"Median latency: {round(np.median(diff) * 1000, 2)}ms")
        percentiles = [10, 25, 50, 75, 90]
        print(
            "Latency percentiles",
            {p: round(1000 * float(np.percentile(diff, p)), 1) for p in percentiles},
        )

    def end(self, *args):
        pass

def parse_arguments():
    parser = argparse.ArgumentParser(description="Run inference benchmark for LLM models")
    parser.add_argument(
        "--device",
        type=str,


    return streamer

class Streamer:
    def __init__(self, tokenizer, print_median=False):
        self.times = []
        self.print_median = print_median
        self.tokenizer = tokenizer

    def put(self, t):
        self.times.append(get_time())
        if len(self.times) > 1:
            print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms")

        if len(self.times) % 10 == 3 and self.print_median:
            ts = np.array(self.times)
            diff = ts[1:] - ts[:-1]
            # print("Token latency:", 1000 * diff, "ms")
            print("Token latency median:", np.median(1000 * diff), "ms")

    def print_report(self):
        times = np.array(self.times)
        diff = times[1:] - times[:-1]
        print(f"Median latency: {round(np.median(diff) * 1000, 2)}ms")
        percentiles = [10, 25, 50, 75, 90]
        print(
            "Latency percentiles",
            {p: round(1000 * float(np.percentile(diff, p)), 1) for p in percentiles},
        )

    def end(self, *args):
        pass

def parse_arguments():
    parser = argparse.ArgumentParser(description="Run inference benchmark for LLM models")
    parser.add_argument(
        "--device",
        type=str,
        default="xpu",
        help="Device to run inference on (e.g., xpu, cuda, cpu)",
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
        help="Model ID from Hugging Face or local path",
    )
    parser.add_argument(


end

        self.times.append(get_time())
        if len(self.times) > 1:
            print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms")

        if len(self.times) % 10 == 3 and self.print_median:
            ts = np.array(self.times)
            diff = ts[1:] - ts[:-1]
            # print("Token latency:", 1000 * diff, "ms")
            print("Token latency median:", np.median(1000 * diff), "ms")

    def print_report(self):
        times = np.array(self.times)
        diff = times[1:] - times[:-1]
        print(f"Median latency: {round(np.median(diff) * 1000, 2)}ms")
        percentiles = [10, 25, 50, 75, 90]
        print(
            "Latency percentiles",
            {p: round(1000 * float(np.percentile(diff, p)), 1) for p in percentiles},
        )

    def end(self, *args):
        pass

def parse_arguments():
    parser = argparse.ArgumentParser(description="Run inference benchmark for LLM models")
    parser.add_argument(
        "--device",
        type=str,
        default="xpu",
        help="Device to run inference on (e.g., xpu, cuda, cpu)",
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
        help="Model ID from Hugging Face or local path",
    )
    parser.add_argument(
        "--attn",
        type=str,
        default="eager",


setup.py

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from distutils.errors import DistutilsModuleError
from warnings import warn

from setuptools import find_packages, setup
from setuptools.command.build_py import build_py
from setuptools.dist import Distribution

# Tested with wheel v0.29.0
class BinaryDistribution(Distribution):
    def has_ext_modules(self):
        return True

class ExtBuildPy(build_py):
    def run(self):
        # build_cmake needs to be called prior to build_py, as the latter


import os
import subprocess
import sys
from urllib.request import urlretrieve

cuda_versions = {
    "118": "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run",
    "120": "https://developer.download.nvidia.com/compute/cuda/12.0.1/local_installers/cuda_12.0.1_525.85.12_linux.run",
    "121": "https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run",
    "122": "https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run",
    "123": "https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.run",
    "124": "https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run",
    "125": "https://developer.download.nvidia.com/compute/cuda/12.5.1/local_installers/cuda_12.5.1_555.42.06_linux.run",
    "126": "https://developer.download.nvidia.com/compute/cuda/12.6.2/local_installers/cuda_12.6.2_560.35.03_linux.run",
}

def install_cuda(version, base_path, download_path):
    formatted_version = f"{version[:-1]}.{version[-1]}"
    folder = f"cuda-{formatted_version}"
    install_path = os.path.join(base_path, folder)


import torch

import bitsandbytes as bnb

p = torch.nn.Parameter(torch.rand(10, 10).cuda())
a = torch.rand(10, 10).cuda()

p1 = p.data.sum().item()

adam = bnb.optim.Adam([p])

out = a * p
loss = out.sum()
loss.backward()
adam.step()

p2 = p.data.sum().item()

assert p1 != p2
print("SUCCESS!")
print("Installation was successful!")


tests/helpers.py

import functools
from io import BytesIO
from itertools import product
import os
import random
from typing import Any

import torch

from bitsandbytes.cextension import HIP_ENVIRONMENT

test_dims_rng = random.Random(42)

TRUE_FALSE = (True, False)
BOOLEAN_TRIPLES = list(product(TRUE_FALSE, repeat=3))  # all combinations of (bool, bool, bool)
BOOLEAN_TUPLES = list(product(TRUE_FALSE, repeat=2))  # all combinations of (bool, bool)

@functools.cache
def get_available_devices(no_cpu=False):


tests/conftest.py

import gc
import random

import numpy as np
import pytest
import torch

def _set_seed():
    torch.manual_seed(0)
    torch.cuda.manual_seed_all(0)
    torch.mps.manual_seed(0)
    np.random.seed(0)
    random.seed(0)

def pytest_runtest_call(item):
    try:
        _set_seed()
        item.runtest()
    except AssertionError as ae:


from contextlib import nullcontext
import copy
import os
import pickle
import platform
from tempfile import TemporaryDirectory

import pytest
import torch

import bitsandbytes as bnb
from bitsandbytes.nn.modules import Linear8bitLt
from tests.helpers import (
    TRUE_FALSE,
    get_available_devices,
    id_formatter,
    torch_load_from_buffer,
    torch_save_to_buffer,
)


import pytest
import torch
import torch.nn as nn

from bitsandbytes import functional as F
from bitsandbytes.cextension import HIP_ENVIRONMENT
from bitsandbytes.nn.parametrize import (
    Bnb4bitParametrization,
    replace_parameter_4bit,
    replace_parameter_4bit_prequantized,
)
from tests.helpers import (
    TRUE_FALSE,
    describe_dtype,
    get_available_devices,
    id_formatter,
    is_supported_on_hpu,
)

class ParametrizeTestModule(nn.Module):


import inspect

import pytest
import torch
from torch import nn

import bitsandbytes as bnb
from tests.helpers import get_available_devices, id_formatter, is_supported_on_hpu

class MockArgs:
    def __init__(self, initial_data):
        for key in initial_data:
            setattr(self, key, initial_data[key])

class MLP8bit(torch.nn.Module):
    def __init__(self, dim1, dim2, has_fp16_weights=True, threshold=0.0):
        super().__init__()
        self.fc1 = bnb.nn.Linear8bitLt(
            dim1,


from math import prod

import pytest
import torch

import bitsandbytes
from bitsandbytes.cextension import HIP_ENVIRONMENT
from tests.helpers import TRUE_FALSE, get_available_devices, id_formatter, is_supported_on_hpu

# torch.library.opcheck is only available in torch 2.4 and later.
# When testing with older versions, we will skip it as a no-op.
if torch.__version__ >= (2, 4):
    opcheck = torch.library.opcheck
else:
    opcheck = lambda *args, **kwargs: None

class TestLLMInt8Ops:
    @pytest.mark.parametrize("device", get_available_devices())
    def test_int8_linear_matmul(self, device):
        A = torch.randint(-128, 127, (10, 20), dtype=torch.int8, device=device)


import pytest
import torch

import bitsandbytes as bnb
from tests.helpers import (
    BOOLEAN_TRIPLES,
    TRUE_FALSE,
    describe_dtype,
    get_available_devices,
    id_formatter,
    is_supported_on_hpu,
)

TRANSPOSE_VALS = [(False, True), (False, False)]

@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("dim1", [40], ids=id_formatter("dim1"))
@pytest.mark.parametrize("dim2", [64, 0], ids=id_formatter("dim2"))
@pytest.mark.parametrize("dim3", [32], ids=id_formatter("dim3"))
@pytest.mark.parametrize("dim4", [48], ids=id_formatter("dim4"))


import pytest

from bitsandbytes.cextension import HIP_ENVIRONMENT, get_cuda_bnb_library_path
from bitsandbytes.cuda_specs import CUDASpecs

@pytest.fixture
def cuda120_spec() -> CUDASpecs:
    return CUDASpecs(
        cuda_version_string="120",
        highest_compute_capability=(8, 6),
        cuda_version_tuple=(12, 0),
    )

@pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm")
def test_get_cuda_bnb_library_path(monkeypatch, cuda120_spec):
    monkeypatch.delenv("BNB_CUDA_VERSION", raising=False)
    assert get_cuda_bnb_library_path(cuda120_spec).stem == "libbitsandbytes_cuda120"


import os
from os.path import join
import shutil
import sys
import time
import uuid

from lion_pytorch import Lion
import pytest
import torch

import bitsandbytes as bnb
import bitsandbytes.functional as F
from bitsandbytes.utils import sync_gpu
from tests.helpers import describe_dtype, get_available_devices, id_formatter

# import apex

k = 20


import pytest
import torch

import bitsandbytes as bnb
from bitsandbytes import functional as F
from tests.helpers import BOOLEAN_TRIPLES, describe_dtype, get_test_dims, id_formatter
from tests.test_autograd import TRANSPOSE_VALS

@pytest.mark.deprecated
def test_dynamic_quantization():
    diffs = []
    reldiffs = []
    for i in range(100):
        A1 = torch.randn(1024, 1024, device="cuda")
        C, S = F.quantize(A1)
        A2 = F.dequantize(C, S)
        diff = torch.abs(A1 - A2)
        reldiff = diff / torch.abs(A1 + 1e-8)
        diffs.append(diff.mean().item())
        reldiffs.append(reldiff.mean().item())


from itertools import product
import math

import pytest
import torch

from tests.helpers import TRUE_FALSE, describe_dtype, id_formatter

transformers = pytest.importorskip("transformers")

def get_4bit_config():
    return transformers.BitsAndBytesConfig(
        load_in_4bit=True,
        load_in_8bit=False,
        llm_int8_threshold=6.0,
        llm_int8_has_fp16_weight=False,
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
    )


import pytest
import torch

from bitsandbytes.nn import Linear8bitLt
from bitsandbytes.nn.triton_based_modules import SwitchBackLinear
from bitsandbytes.triton.triton_utils import is_triton_available
from tests.helpers import TRUE_FALSE

@pytest.mark.skipif(
    not is_triton_available() or not torch.cuda.is_available() or not torch.cuda.get_device_capability()[0] >= 8,
    reason="This test requires triton and a GPU with compute capability 8.0 or higher.",
)
@pytest.mark.deprecated
@pytest.mark.parametrize("vector_wise_quantization", TRUE_FALSE)
def test_switchback(vector_wise_quantization):
    for dim in [83]:
        for batch in [13]:
            standard = torch.nn.Linear(dim, 4 * dim).cuda().half()
            switchback = (
                SwitchBackLinear(dim, 4 * dim, vector_wise_quantization=vector_wise_quantization).cuda().half()


import math
import platform
import random
import time

import einops
from packaging import version
import pytest
import torch

import bitsandbytes as bnb
from bitsandbytes import functional as F
from bitsandbytes.cextension import HIP_ENVIRONMENT, ROCM_GPU_ARCH
from tests.helpers import (
    BOOLEAN_TUPLES,
    TRUE_FALSE,
    describe_dtype,
    get_available_devices,
    get_test_dims,
    id_formatter,
    is_supported_on_hpu,


import copy
import os
import pickle
import platform
from tempfile import TemporaryDirectory

import pytest
import torch

import bitsandbytes as bnb
from bitsandbytes.cextension import HIP_ENVIRONMENT
from tests.helpers import (
    TRUE_FALSE,
    describe_dtype,
    get_available_devices,
    id_formatter,
    is_supported_on_hpu,
    torch_load_from_buffer,
    torch_save_to_buffer,
)


import torch
import torch._dynamo
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# torch._dynamo.config.suppress_errors = True

torch.set_float32_matmul_precision("high")

quantization_config = BitsAndBytesConfig(load_in_8bit=True)

# torch._dynamo.config.capture_dynamic_output_shape_ops = True

model_id = "google/gemma-2-2b-it"
# model_id = "Qwen/Qwen2.5-7B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,


import torch
from transformers import LlamaForCausalLM, LlamaTokenizer

MAX_NEW_TOKENS = 128
model_name = "meta-llama/Llama-2-7b-hf"

text = "Hamburg is in which country?\n"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
input_ids = tokenizer(text, return_tensors="pt").input_ids

max_memory = f"{int(torch.cuda.mem_get_info()[0] / 1024**3) - 2}GB"

n_gpus = torch.cuda.device_count()
max_memory = {i: max_memory for i in range(n_gpus)}

model = LlamaForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True, max_memory=max_memory)

generated_ids = model.generate(input_ids, max_length=MAX_NEW_TOKENS)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))


import dataclasses
from functools import lru_cache
import logging
import re
import subprocess
from typing import Optional

import torch

@dataclasses.dataclass(frozen=True)
class CUDASpecs:
    highest_compute_capability: tuple[int, int]
    cuda_version_string: str
    cuda_version_tuple: tuple[int, int]

    @property
    def has_imma(self) -> bool:
        return torch.version.hip or self.highest_compute_capability >= (7, 5)


from collections.abc import Sequence
from math import prod
from typing import Optional

import torch

_IS_TORCH_GTE_24 = False

if hasattr(torch.library, "register_fake"):
    _IS_TORCH_GTE_24 = True
    register_fake = torch.library.register_fake
    register_kernel = torch.library.register_kernel
else:
    # PyTorch <= 2.3
    register_fake = torch.library.impl_abstract
    register_kernel = torch.library.impl

# Int8 mixed precision matmul + dequant + bias
torch.library.define(
    "bitsandbytes::int8_mixed_scaled_mm",
    "(Tensor A, Tensor CA, Tensor CB, Tensor SCA, Tensor SCB, Tensor? outlier_cols=None, Tensor? bias=None) -> (Tensor, Tensor?)",


bitsandbytes/functional.py

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from collections.abc import Iterable
import ctypes as ct
import itertools
from math import prod
from typing import Any, Optional, Union

import numpy as np
import torch
from torch import Tensor
from typing_extensions import deprecated

from bitsandbytes.utils import pack_dict_to_tensor, unpack_tensor_to_dict

from .cextension import HIP_ENVIRONMENT, lib

name2qmap = {}


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import importlib
import sys

import torch

from . import _ops, research, utils
from .autograd._functions import (
    MatmulLtState,
    matmul,
    matmul_4bit,
)
from .backends.cpu import ops as cpu_ops
from .backends.default import ops as default_ops
from .nn import modules
from .optim import adam


if __name__ == "__main__":
    from bitsandbytes.diagnostics.main import main

    main()


bitsandbytes/utils.py

import json
import shlex
import subprocess

import torch

def outlier_hook(module, input):
    assert isinstance(module, torch.nn.Linear)
    tracer = OutlierTracer.get_instance()
    hvalue = tracer.get_hvalue(module.weight)
    if hvalue not in tracer.hvalue2outlier_idx:
        outlier_idx = find_outlier_dims(module.weight)
        tracer.outliers.append(outlier_idx)
        tracer.hvalues.append(hvalue)
        if len(tracer.outliers) > 1:
            # assign the current layer the outlier idx found from the weight
            # of the previous linear layer
            if tracer.outliers[-1].numel() > 0:
                assert tracer.outliers[-1].max() < module.weight.shape[1]
            tracer.hvalue2outlier_idx[hvalue] = tracer.outliers[-1]


bitsandbytes/consts.py

from pathlib import Path
import platform

DYNAMIC_LIBRARY_SUFFIX = {
    "Darwin": ".dylib",
    "Linux": ".so",
    "Windows": ".dll",
}.get(platform.system(), ".so")

PACKAGE_DIR = Path(__file__).parent
PACKAGE_GITHUB_URL = "https://github.com/TimDettmers/bitsandbytes"
NONPYTORCH_DOC_URL = "https://github.com/TimDettmers/bitsandbytes/blob/main/docs/source/nonpytorchcuda.mdx"


bitsandbytes/cextension.py

import ctypes as ct
import functools
import logging
import os
from pathlib import Path
import re
from typing import Optional

import torch

from bitsandbytes.consts import DYNAMIC_LIBRARY_SUFFIX, PACKAGE_DIR
from bitsandbytes.cuda_specs import CUDASpecs, get_cuda_specs, get_cuda_version_tuple, get_rocm_gpu_arch

logger = logging.getLogger(__name__)

def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path:
    """
    Get the disk path to the CUDA BNB native library specified by the
    given CUDA specs, taking into account the `BNB_CUDA_VERSION` override environment variable.


import subprocess

from packaging import version
import torch

try:
    import triton  # noqa: F401
    import triton.language as tl  # noqa: F401

    triton_available = True
except ImportError:
    triton_available = False

_NF4_QUANT_TABLE = torch.tensor(
    [
        -1.0,
        -0.6961928009986877,
        -0.5250730514526367,
        -0.39491748809814453,
        -0.28444138169288635,


from collections.abc import Sequence
from math import prod, sqrt
from typing import Optional

import torch

from ..._ops import register_kernel
from ..utils import CODE

@register_kernel("bitsandbytes::int8_mm_dequant", "default")
def _(
    A: torch.Tensor,
    row_stats: torch.Tensor,
    col_stats: torch.Tensor,
    dtype: Optional[torch.dtype] = None,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    torch._check(A.dtype == torch.int32, lambda: f"A must be int32, got {A.dtype}")
    torch._check(row_stats.dtype == torch.float32, lambda: f"row_stats must be float32, got {row_stats.dtype}")
    torch._check(col_stats.dtype == torch.float32, lambda: f"col_stats must be float32, got {col_stats.dtype}")


from collections.abc import Sequence
import math

import torch

from ..._ops import register_kernel
from ..utils import GAUDI_SW_VER

# convert btw standard 4-bit compression format and ipex compression format
# needed for backward compatibility with older versions of gaudi sw
def _reverse_4bit_compress_format(weight: torch.Tensor):
    out_1 = (weight & 0xF0) >> 4
    out_2 = (weight & 0xF) << 4
    out = out_1 | out_2
    return out

@register_kernel("bitsandbytes::dequantize_4bit", "hpu")
def _(
    A: torch.Tensor,


import torch

import triton
import triton.language as tl

# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dQuantizeFP4
# @triton.autotune(
#     configs=[
#         triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32),
#         triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32),
#         triton.Config({"SPLIT_NUM_BLOCKS": 1}),
#         triton.Config({"SPLIT_NUM_BLOCKS": 2}),
#         triton.Config({"SPLIT_NUM_BLOCKS": 4}),
#         triton.Config({"SPLIT_NUM_BLOCKS": 8}),
#     ],
#     key=["n_elements"],
# )
@triton.jit
def quantize_fp4_blockwise_kernel(
    A_ptr,


from collections.abc import Sequence
from typing import Optional

import torch

from . import kernels_4bit, kernels_8bit_quant, kernels_optim

# currently codes unused, kept for reference
# Should be the same for quant/dequant
# from bitsandbytes.functional import get_4bit_type
# _FP4_QUANT_TABLE = get_4bit_type("fp4", device="xpu")
# _NF4_QUANT_TABLE = get_4bit_type("nf4", device="xpu")
device_type = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda"
torch_accelerator_module = getattr(torch, device_type, torch.cuda)

def quantize_blockwise(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]:
    torch._check_is_size(blocksize)
    # torch._check(A.dtype == torch.float32, lambda: f"A must be float32 on xpu, got {A.dtype}")
    with torch_accelerator_module.device(A.device):
        out, absmax = kernels_8bit_quant.quantize_blockwise_triton(A, code, blocksize)


import math
from typing import Optional

import torch

import triton
import triton.language as tl

# from triton.language.extra import libdevice
from .kernels_8bit_quant import (
    dequant_8bit_blockwise,
    dequant_8bit_blockwise_kernel_util,
    quantize_8bit_blockwise_kernel_util,
    quantize_blockwise_triton,
)

MOMENTUM = 0
RMSPROP = 1
ADAGRAD = 2
ADAM = 3
# LION should be larger than MOMENTUM, RMSPROP, ADAGRAD due to comparison in kernels


import torch

import triton
import triton.language as tl

# @triton.autotune(
#     configs=[
#         # triton.Config({'SPLIT_SIZE': 64}),
#         # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
#         # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
#         # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
#         # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=4, num_warps=32),
#         # triton.Config({'SPLIT_SIZE': 128}),
#         # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
#         # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
#         # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
#         # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=4, num_warps=32),
#         triton.Config({"SPLIT_SIZE": 256}),
#         # triton.Config({'SPLIT_SIZE': 256, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
#         # triton.Config({'SPLIT_SIZE': 256, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),


import ctypes as ct
import logging

import torch

from bitsandbytes.functional import get_ptr

from ..._ops import register_kernel
from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib

logger = logging.getLogger(__name__)

# torch._int_mm for s8@s8->s32 is supported on CPU from torch 2.4+.
# However, we can overflow if we use this without AVX512_VNNI support.
# This is fixed in torch 2.6+, so we set this as the minimum to be safe.
# For more information: https://github.com/pytorch/pytorch/pull/136942
# TODO(matthewdouglas): aarch64?
if torch.__version__ >= (2, 6):

    @register_kernel("bitsandbytes::int8_linear_matmul", "cpu")
    def _(A: torch.Tensor, B: torch.Tensor):


from collections.abc import Sequence
import ctypes as ct
from math import prod
from typing import Optional

import torch

from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, _get_tensor_stream, get_ptr

from ..._ops import register_kernel
from ...cextension import HIP_ENVIRONMENT, lib

@register_kernel("bitsandbytes::int8_linear_matmul", "cuda")
def _(A: torch.Tensor, B: torch.Tensor):
    out = torch.empty((*A.shape[:-1], B.shape[0]), device=A.device, dtype=torch.int32)
    return _int8_linear_matmul_impl(A, B, out)

@register_kernel("bitsandbytes::int8_linear_matmul.out", "cuda")
def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):


from collections.abc import Sequence
import ctypes as ct
import logging

from packaging import version
import torch

from bitsandbytes.functional import _get_tensor_stream, get_ptr

from ..._ops import register_kernel
from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib
from ..utils import triton_available

logger = logging.getLogger(__name__)

# _int_mm is available in torch starting from 2.9 version
if version.parse(torch.__version__).release >= version.parse("2.9").release:

    @register_kernel("bitsandbytes::int8_linear_matmul", "xpu")
    def _(A: torch.Tensor, B: torch.Tensor):
        return torch._int_mm(


import torch

from bitsandbytes.triton.triton_utils import is_triton_available

if not is_triton_available():

    def int8_matmul_rowwise_dequantize(a, b, state_x, state_w, bias):
        return None
else:
    import triton
    import triton.language as tl

    from .matmul_perf_model import early_config_prune, estimate_matmul_time

    # This is a matmul kernel based on triton.ops.matmul
    # It is modified to support rowwise quantized input and columnwise quantized weight
    # It's purpose is fused matmul then dequantize
    # It does support bias.

    def init_to_zero(name):
        return lambda nargs: nargs[name].zero_()


import functools

@functools.lru_cache(None)
def is_triton_available():
    try:
        from torch.utils._triton import has_triton, has_triton_package

        return has_triton_package() and has_triton()
    except Exception:
        return False


import torch

from bitsandbytes.triton.triton_utils import is_triton_available

if not is_triton_available():

    def quantize_global_transpose(input):
        return None

    def quantize_global(x: torch.Tensor):
        return None
else:
    import triton
    import triton.language as tl

    # global quantize
    @triton.autotune(
        configs=[
            triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
            triton.Config({"BLOCK_SIZE": 2048}, num_stages=1),
        ],


import math

import torch

from bitsandbytes.triton.triton_utils import is_triton_available

if not is_triton_available():

    def quantize_columnwise_and_transpose(x: torch.Tensor):
        return None
else:
    import triton
    import triton.language as tl

    # This kernel does fused columnwise quantization and transpose.

    # TODO: autotune this better.
    @triton.autotune(
        configs=[
            triton.Config({}, num_stages=1),
            triton.Config({}, num_stages=2),


import math

import torch

from bitsandbytes.triton.triton_utils import is_triton_available

if not is_triton_available():

    def quantize_rowwise(x: torch.Tensor):
        return None
else:
    import triton
    import triton.language as tl

    # rowwise quantize

    # TODO: autotune this better.
    @triton.autotune(
        configs=[
            triton.Config({}, num_stages=1, num_warps=8),
            triton.Config({}, num_stages=2, num_warps=8),


import math

import torch

from bitsandbytes.triton.triton_utils import is_triton_available

if not is_triton_available():

    def dequantize_rowwise(x: torch.Tensor, state_x: torch.Tensor):
        return None
else:
    import triton
    import triton.language as tl

    # rowwise quantize

    # TODO: autotune this better.
    @triton.autotune(
        configs=[
            triton.Config({}, num_stages=1, num_warps=8),
            triton.Config({}, num_stages=2, num_warps=8),


import torch

from bitsandbytes.triton.triton_utils import is_triton_available

if not is_triton_available():

    def int8_matmul_mixed_dequantize(a, b, state_x, state_w, bias):
        return None
else:
    import triton
    import triton.language as tl

    from .matmul_perf_model import early_config_prune, estimate_matmul_time

    # This is a matmul kernel based on triton.ops.matmul
    # It is modified to support rowwise quantized input and global quantized weight
    # It's purpose is fused matmul then dequantize
    # It does support bias.

    def init_to_zero(name):
        return lambda nargs: nargs[name].zero_()


# Adapted from https://github.com/triton-lang/kernels/blob/eeeebdd8be7d13629de22d600621e6234057eed3/kernels/matmul_perf_model.py
# https://github.com/triton-lang/kernels is licensed under the MIT License.

import functools
import heapq

import torch

from triton import cdiv
from triton.runtime import driver
from triton.testing import (
    get_dram_gbps,
    get_max_simd_tflops,
    get_max_tensorcore_tflops,
    nvsmi,
)

@functools.lru_cache
def get_clock_rate_in_khz():
    try:


from . import nn
from .autograd._functions import (
    matmul_fp8_global,
    matmul_fp8_mixed,
    switchback_bnb,
)


from .modules import LinearFP8Global, LinearFP8Mixed


from typing import TypeVar

import torch
from torch import nn

import bitsandbytes as bnb

T = TypeVar("T", bound="torch.nn.Module")

class LinearFP8Mixed(nn.Linear):
    def __init__(self, input_features, output_features, bias=True):
        super().__init__(input_features, output_features, bias)
        self.bw_code = None
        self.fw_code = None
        array = [4096, 2048, 1024, 512, 256, 128, 64, 0]
        for i, k in enumerate(array):
            if input_features > array[i + 1]:
                self.bsz = k
                break
        for i, k in enumerate(array):


from functools import reduce  # Required in Python 3
import operator
from typing import Optional
import warnings

import torch

from bitsandbytes.autograd._functions import GlobalOutlierPooler, MatmulLtState
import bitsandbytes.functional as F

# math.prod not compatible with python < 3.8
def prod(iterable):
    return reduce(operator.mul, iterable, 1)

class MatMulFP8Mixed(torch.autograd.Function):
    # forward is the same, but we added the fallback for pre-turing GPUs
    # backward is mostly the same, but adds one extra clause (see "elif state.CxB is not None")

    @staticmethod


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from bitsandbytes.optim.optimizer import Optimizer1State

class RMSprop(Optimizer1State):
    def __init__(
        self,
        params,
        lr=1e-2,
        alpha=0.99,
        eps=1e-8,
        weight_decay=0,
        momentum=0,
        centered=False,
        optim_bits=32,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from bitsandbytes.optim.optimizer import Optimizer1State

class SGD(Optimizer1State):
    def __init__(
        self,
        params,
        lr,
        momentum=0,
        dampening=0,
        weight_decay=0,
        nesterov=False,
        optim_bits=32,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        block_wise=True,


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from bitsandbytes.optim.optimizer import Optimizer2State

class AdamW(Optimizer2State):
    def __init__(
        self,
        params,
        lr=1e-3,
        betas=(0.9, 0.999),
        eps=1e-8,
        weight_decay=1e-2,
        amsgrad=False,
        optim_bits=32,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,