import torch

from sglang.kernels.jit.benchmark import marker
from sglang.kernels.jit.benchmark.utils import (
    DEFAULT_DEVICE,
    create_empty,
    create_random,
)
from sglang.kernels.ops.kvcache.kvcache import store_cache
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

register_cuda_ci(
    est_time=9, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
)
register_amd_ci(est_time=9, stage="jit-kernel-benchmark", runner_config="amd")


@torch.compile()
def torch_compile_store_cache(
    k: torch.Tensor,
    v: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    indices: torch.Tensor,
) -> None:
    k_cache[indices] = k
    v_cache[indices] = v


alt_stream = torch.cuda.Stream()


def torch_streams_store_cache(
    k: torch.Tensor,
    v: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    indices: torch.Tensor,
) -> None:
    current_stream = torch.cuda.current_stream()
    alt_stream.wait_stream(current_stream)
    k_cache[indices] = k
    with torch.cuda.stream(alt_stream):
        v_cache[indices] = v
    current_stream.wait_stream(alt_stream)


CACHE_SIZE = 2 * 1024 * 1024
FN_MAP = {
    "jit": store_cache,
    "torch_compile": torch_compile_store_cache,
    "torch_streams": torch_streams_store_cache,
}


@marker.parametrize("item_size", [64, 128, 256, 512, 1024], [1024])
@marker.parametrize("batch_size", [2**n for n in range(0, 15)], [16])
@marker.benchmark("impl", ["jit", "torch_compile", "torch_streams"])
def benchmark(batch_size: int, item_size: int, impl: str):
    torch.manual_seed(42)
    k = create_random(batch_size, item_size)
    k_cache = create_empty(CACHE_SIZE, item_size)
    v = create_random(batch_size, item_size)
    v_cache = create_empty(CACHE_SIZE, item_size)
    indices = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size]
    return marker.do_bench(
        FN_MAP[impl],
        input_args=(k, v, k_cache, v_cache, indices),
        graph_clone_args=(0, 1, 4),  # not need to clone cache, which is large
        memory_args=(k, v, indices),  # k_cache / v_cache excluded
        memory_output=(k, v),  # inplace write, size = k + v
    )


# Asymmetric K/V (head_dim != v_head_dim). The item_size sweep above drives both
# rows from one value, so it never reaches the split-prefix-plus-tail path.
# 192/128 and 384/256 both go live in a single MiMoV2 TP=4 deployment, whose
# layers carry either 1 or 2 kv heads per rank; the reversed and wide pairs cover
# a V-side tail and a num_split > 1 shape.
ASYM_ITEM_SIZES = [(192, 128), (384, 256), (128, 192), (1024, 512)]


@marker.parametrize("k_item,v_item", ASYM_ITEM_SIZES, [(192, 128), (1024, 512)])
@marker.parametrize("batch_size", [2**n for n in range(0, 15)], [16])
@marker.benchmark("impl", ["jit", "torch_compile", "torch_streams"])
def benchmark_asymmetric(batch_size: int, k_item: int, v_item: int, impl: str):
    torch.manual_seed(42)
    k = create_random(batch_size, k_item)
    k_cache = create_empty(CACHE_SIZE, k_item)
    v = create_random(batch_size, v_item)
    v_cache = create_empty(CACHE_SIZE, v_item)
    indices = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size]
    return marker.do_bench(
        FN_MAP[impl],
        input_args=(k, v, k_cache, v_cache, indices),
        graph_clone_args=(0, 1, 4),  # not need to clone cache, which is large
        memory_args=(k, v, indices),  # k_cache / v_cache excluded
        memory_output=(k, v),  # inplace write, size = k + v
    )


if __name__ == "__main__":
    benchmark.run()
    benchmark_asymmetric.run()
