"""Seeded triton-vs-torch parity sweep for the DSpark kernels.

Guards against toolchain drift (triton/torch upgrades) silently diverging
the triton implementations from their torch references. Every method calls
the production kernel pair directly via `Cls.torch(...)` / `Cls.triton(...)`
(no env-var dispatch) on a small set of adversarial inputs and compares
exactly, or with the tolerance the kernel is specified to meet.
"""

import types
import unittest

import torch

from sglang.kernels.ops.attention import (
    dsv4_attn_metadata_kernels as attn_metadata_kernels,
)
from sglang.kernels.ops.speculative import ragged_verify_kernels
from sglang.kernels.ops.speculative.dspark import (
    dspark_accept,
    dspark_attn_metadata,
    dspark_draft_model,
    dspark_schedule,
    dspark_verify_window,
)
from sglang.srt.speculative.dspark_components.dspark_planner import (
    DSparkScheduleConfig,
)
from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase

register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")

DEVICE = torch.device("cuda")
VOCAB = 129280


def _ri(lo, hi, shape, dtype=torch.int64, g=None):
    return torch.randint(lo, hi, shape, device=DEVICE, dtype=dtype, generator=g)


def _layout(verify_lens, graph_num_tokens):
    return RaggedVerifyLayout.from_verify_lens_device(
        verify_lens=verify_lens, graph_num_tokens=graph_num_tokens
    )


class _Bf16Linear(torch.nn.Module):
    quant_method = None

    def __init__(self, weight):
        super().__init__()
        self.weight = weight

    def forward(self, x):
        return torch.nn.functional.linear(x, self.weight), None


def _case_accept_greedy(tc):
    torch.manual_seed(0)
    bs, t = 8, 6
    candidates = _ri(0, 200, (bs, t))
    target_logits = torch.randn(bs * t, 200, device=DEVICE)
    for cutoff in (None, _ri(1, t + 1, (bs,), torch.int32)):
        tc._parity(
            dspark_accept.AcceptGreedy,
            candidates=candidates,
            target_logits=target_logits,
            verify_num_draft_tokens=t,
            cutoff_verify_lens=cutoff,
        )
    # gather_row_bonus: bonus token at a per-row column index.
    table, idx = _ri(0, VOCAB, (64, t)), _ri(0, t, (64,), torch.int32)
    ref = table[torch.arange(64, device=DEVICE), idx.long()]
    tc._eq(dspark_accept.gather_row_bonus_triton(table=table, idx=idx), ref)


def _case_accept_sampling(tc):
    torch.manual_seed(1)
    bs, t = 64, 6
    accept_index = _ri(0, bs * t, (bs, t))
    predicts = _ri(0, VOCAB, (bs * t,))
    correct_len = _ri(0, t, (bs,), torch.int32)
    rows = torch.arange(bs, device=DEVICE)
    ref = predicts[accept_index[rows, correct_len.long()].long()]
    got = dspark_accept.gather_two_level_bonus_triton(
        accept_index=accept_index, predicts=predicts, correct_len=correct_len
    )
    tc._eq(got, ref)


def _case_build_block_seq_lens_causal(tc):
    torch.manual_seed(2)
    seq_lens = _ri(1, 100000, (128,))
    for block_size in (1, 5, 7):
        tc._parity(
            dspark_attn_metadata.BuildBlockSeqLensCausal,
            seq_lens=seq_lens,
            block_size=block_size,
            device=DEVICE,
        )


def _case_build_out_tokens(tc):
    torch.manual_seed(3)
    bs, gamma = 64, 5
    for cl_dtype in (torch.int32, torch.int64):
        # Bonus insertion swept through every position 0..gamma.
        cl = (torch.arange(bs, device=DEVICE) % (gamma + 1)).to(cl_dtype)
        tc._parity(
            dspark_verify_window.BuildOutTokens,
            draft_tokens=_ri(0, VOCAB, (bs, gamma)),
            correct_len=cl,
            bonus=_ri(0, VOCAB, (bs,)),
            verify_num_draft_tokens=gamma + 1,
            gamma=gamma,
        )


def _case_build_ragged_verify_window(tc):
    torch.manual_seed(4)
    gamma, t, bs = 5, 6, 8
    verify_lens = _ri(1, t + 1, (bs,), torch.int32)
    batch = types.SimpleNamespace(
        seq_lens=_ri(1, 20, (bs,)),
        req_pool_indices=torch.randperm(bs + 3, device=DEVICE)[:bs],
    )
    model_runner = types.SimpleNamespace(
        req_to_token_pool=types.SimpleNamespace(
            req_to_token=_ri(0, 1_000_000, (bs + 3, 64), torch.int32)
        )
    )
    for graph_num_tokens in (bs * t, (bs + 3) * t):  # tight and bucket padding
        tc._parity(
            dspark_verify_window.BuildRaggedVerifyWindow,
            batch=batch,
            layout=_layout(verify_lens, graph_num_tokens),
            draft_block_ids=_ri(0, VOCAB, (bs, gamma)),
            draft_tokens=_ri(0, VOCAB, (bs, gamma)),
            bs=bs,
            device=DEVICE,
            verify_num_draft_tokens=t,
            model_runner=model_runner,
        )


def _case_build_step_local(tc):
    torch.manual_seed(5)
    for org_width, per_partition, bias_dtype in (
        (32320, 32384, torch.bfloat16),
        (5000, 8192, torch.float32),
    ):
        bias = (torch.randn(3, org_width, device=DEVICE) * 3.0).to(bias_dtype)
        base = torch.randn(3, per_partition, device=DEVICE)
        got, _ = tc._parity(
            dspark_draft_model.BuildStepLocal, bias=bias, base_local=base
        )
        # Padding columns beyond org_width must stay pure base.
        tc.assertTrue(torch.equal(got[:, org_width:], base[:, org_width:]))


def _case_cap_correct_len(tc):
    torch.manual_seed(6)
    bs, nd = 64, 6
    verify_lens = _ri(1, nd + 1, (bs,), torch.int32)
    for cl_dtype in (torch.int32, torch.int64):
        cl = (torch.arange(bs, device=DEVICE) % (nd + 1)).to(cl_dtype)
        tc._parity(dspark_accept.CapCorrectLen, correct_len=cl, verify_lens=verify_lens)


def _case_causal_swa_page_indices(tc):
    swa, num_pool, pool_len, num_q = 128, 64, 600, 40
    g = torch.Generator(device=DEVICE).manual_seed(7)
    kw = dict(
        req_to_token=_ri(0, 40000, (num_pool, pool_len), torch.int32, g),
        full_to_swa_mapping=_ri(0, 1 << 20, (40000,), torch.int64, g),
        req_pool_indices_repeated=_ri(0, num_pool, (num_q,), torch.int32, g),
        swa_window=swa,
        page_index_aligned_size=96,
    )
    # Lens short of / straddling / beyond the SWA window boundary.
    for lo, hi in ((1, swa), (swa - 4, swa + 4), (swa + 1, pool_len)):
        lens = _ri(lo, hi, (num_q,), torch.int32, g)
        cls = attn_metadata_kernels.BuildCausalSwaPageIndices
        ref = cls.torch(seq_lens_casual=lens, **kw)
        got = cls.triton(seq_lens_casual=lens, **kw)
        tc.assertEqual(got.shape, ref.shape)
        tc.assertEqual(got.dtype, ref.dtype)
        # Parity holds on the attended region; padding slots must be -1.
        col = torch.arange(ref.shape[1], device=DEVICE).view(1, -1)
        attended = col < torch.clamp(lens, max=swa).view(-1, 1)
        tc.assertTrue(torch.equal(got[attended], ref[attended]))
        tc.assertTrue(bool((got[~attended] == -1).all()))


def _case_commit_inject_layout(tc):
    stride, num_pool, pool_len, n_full, bs = 7, 300, 400, 50000, 64
    g = torch.Generator(device=DEVICE).manual_seed(8)
    pool_perm = torch.randperm(num_pool, device=DEVICE, generator=g)
    kw = dict(
        req_pool_indices=pool_perm[:bs],
        req_to_token=_ri(0, n_full, (num_pool, pool_len), torch.int64, g),
        prefix_lens=_ri(1, pool_len - stride, (bs,), torch.int64, g),
        block_pos_offsets=torch.arange(stride, device=DEVICE),
        full_to_swa_mapping=_ri(0, 1 << 20, (n_full,), torch.int64, g),
        commit_lens=_ri(0, stride + 1, (bs,), torch.int32, g),
        stride=stride,
    )
    tc._parity(dspark_verify_window.BuildCommitInjectLayout, **kw)
    # commit_len edges: 0 masks the whole row to -1, stride keeps it all.
    kw.update(
        req_pool_indices=kw["req_pool_indices"][:2],
        prefix_lens=kw["prefix_lens"][:2],
        commit_lens=torch.tensor([0, stride], device=DEVICE, dtype=torch.int32),
    )
    edge = dspark_verify_window.BuildCommitInjectLayout.triton(**kw)
    swa_2d = edge.swa_loc.view(2, stride)
    tc.assertTrue(bool((swa_2d[0] == -1).all()))
    tc.assertTrue(bool((swa_2d[1] >= 0).all()))


def _case_commit_kv_proj(tc):
    hidden, head_dim, num_stages = 1024, 576, 3
    g = torch.Generator(device=DEVICE).manual_seed(9)
    linears = [
        _Bf16Linear(
            (torch.randn(head_dim, hidden, device=DEVICE, generator=g) * 0.02).to(
                torch.bfloat16
            )
        )
        for _ in range(num_stages)
    ]
    main_x = (torch.randn(56, hidden, device=DEVICE, generator=g) * 0.5).to(
        torch.bfloat16
    )
    cls = dspark_draft_model.CommitKvProj
    ref = cls.torch(main_x=main_x, wkv_linears=linears)
    got = cls.triton(main_x=main_x, wkv_linears=linears)
    tc.assertEqual(len(got), num_stages)
    for kv_got, kv_ref in zip(got, ref):
        tc.assertEqual(kv_got.shape, kv_ref.shape)
        tc.assertTrue(kv_got.is_contiguous())
        torch.testing.assert_close(kv_got.float(), kv_ref.float(), rtol=2e-2, atol=2e-3)
    # fp8 blockwise weight dequant path (2x3 grid of 128x128 blocks).
    out_dim, in_dim, block = 192, 384, 128
    w8 = torch.randn(out_dim, in_dim, device=DEVICE, generator=g).to(
        torch.float8_e4m3fn
    )
    scale = torch.rand(2, 3, device=DEVICE, generator=g) + 0.5
    sf = scale.repeat_interleave(block, 0)[:out_dim]
    sf = sf.repeat_interleave(block, 1)[:, :in_dim]
    expected = (w8.to(torch.float32) * sf).to(torch.bfloat16)
    stub = types.SimpleNamespace(weight=w8, weight_scale_inv=scale)
    tc._eq(dspark_draft_model._dequant_linear_weight(stub), expected)


def _case_compact_layout(tc):
    torch.manual_seed(10)
    gamma, t, bs = 5, 6, 64
    verify_lens = _ri(1, t + 1, (bs,), torch.int32)
    total = int(verify_lens.sum().item())
    for padded_total in (total, bs * t):  # exact and bucket padding
        tc._parity(
            dspark_verify_window.CompactRowIndex,
            verify_lens=verify_lens,
            padded_total=padded_total,
            device=DEVICE,
        )
        tc._parity(
            dspark_verify_window.CompactVerifyIds,
            draft_block_ids=_ri(0, VOCAB, (bs, gamma)),
            draft_tokens=_ri(0, VOCAB, (bs, gamma)),
            layout=_layout(verify_lens, padded_total),
            device=DEVICE,
        )


def _case_swa_page_indices(tc):
    torch.manual_seed(11)
    block_size, num_q, max_reqs, n_full = 5, 320, 300, 50000
    _, gather = tc._parity(
        dspark_attn_metadata.ComputeDsparkWindowGather,
        seq_lens_casual=_ri(1, 300, (num_q,), torch.int32),
        req_pool_indices_repeated=_ri(0, max_reqs, (num_q,)),
        block_size=block_size,
        swa_window=128,
    )
    tc._parity(
        dspark_attn_metadata.BuildDsparkSwaPageIndices,
        req_to_token=_ri(0, n_full, (max_reqs, 400), torch.int32),
        full_to_swa_mapping=_ri(0, 20000, (n_full,), torch.int32),
        req_pool_indices_per_request=gather.req_pool_indices_per_request,
        offsets=gather.offsets,
        invalid=gather.invalid,
        out_loc=_ri(0, n_full, (num_q,)),
        context_lens=gather.context_lens,
        block_size=block_size,
        swa_window=128,
        page_index_aligned_size=64,
    )


def _case_expand_prefill_causally(tc):
    torch.manual_seed(12)
    # Vectorized branch: ragged extends with padded token count.
    bs = 64
    extend = _ri(1, 8, (bs,))
    num_tokens = int(extend.sum())
    req_pool_indices = torch.randperm(512, device=DEVICE)[:bs]
    seq_lens = _ri(8, 500, (bs,))
    tc._parity(
        attn_metadata_kernels.ExpandPrefillCausally,
        req_pool_indices=req_pool_indices,
        seq_lens=seq_lens,
        extend_seq_lens=extend,
        extend_start_loc=torch.cumsum(extend, dim=0) - extend,
        seq_lens_cpu=None,
        extend_seq_lens_cpu=None,
        num_tokens=num_tokens,
        padded_num_tokens=num_tokens + 5,
    )
    # Loop branch: uniform extend with CPU lens and no padding.
    bs2, block = 8, 6
    tc._parity(
        attn_metadata_kernels.ExpandPrefillCausally,
        req_pool_indices=req_pool_indices[:bs2],
        seq_lens=seq_lens[:bs2],
        extend_seq_lens=torch.full((bs2,), block, device=DEVICE),
        extend_start_loc=None,
        seq_lens_cpu=[int(x) for x in seq_lens[:bs2].tolist()],
        extend_seq_lens_cpu=[block] * bs2,
        num_tokens=bs2 * block,
        padded_num_tokens=None,
    )


def _case_finalize_accept_lens(tc):
    torch.manual_seed(13)
    bs = 64
    for prefix_dtype in (torch.int32, torch.int64):
        tc._parity(
            dspark_accept.FinalizeAcceptLens,
            correct_len=_ri(0, 7, (bs,), torch.int32),
            cap_trim_lens=_ri(0, 4, (bs,)),
            prefix_lens=_ri(1, 4000, (bs,), prefix_dtype),
        )


def _case_mixed_accept_select(tc):
    torch.manual_seed(14)
    bs = 64
    # Mixed dtypes between the greedy and sampling lanes.
    tc._parity(
        dspark_accept.SelectMixedAccept,
        greedy_mask=torch.rand(bs, device=DEVICE) < 0.5,
        greedy_len=_ri(0, 7, (bs,)),
        greedy_bonus=_ri(0, 100000, (bs,)),
        greedy_trim=_ri(0, 4, (bs,)),
        sampling_len=_ri(0, 7, (bs,), torch.int32),
        sampling_bonus=_ri(0, 100000, (bs,)),
        sampling_trim=_ri(0, 4, (bs,), torch.int32),
    )


def _case_padded_to_bucket(tc):
    torch.manual_seed(15)
    for bs, padded_bs, graph_num_tokens in ((3, 6, 16), (2, 8, 16), (8, 128, 768)):
        verify_lens = _ri(1, 7, (bs,), torch.int32)
        if int(verify_lens.sum()) > graph_num_tokens:
            verify_lens = torch.ones(bs, dtype=torch.int32, device=DEVICE)
        got, _ = tc._parity(
            ragged_verify_kernels.PaddedToBucket,
            verify_lens=verify_lens,
            graph_num_tokens=graph_num_tokens,
            bs=bs,
            padded_bs=padded_bs,
        )
        # Padding rows must absorb exactly the leftover budget.
        tc.assertEqual(int(got.to(torch.int64).sum()), graph_num_tokens)
        if padded_bs > bs:
            tc.assertTrue(torch.equal(got[:bs], verify_lens))


def _case_page_table_positions(tc):
    num_pool, pool_len = 128, 4096
    g = torch.Generator(device=DEVICE).manual_seed(16)
    req_to_token = _ri(0, 1 << 20, (num_pool, pool_len), torch.int32, g)
    # Large page + non-pool-aligned max_seq_len, then page_size 1.
    for num_q, page_size, max_seq_len in ((300, 64, 4000), (56, 1, 4096)):
        tc._parity(
            attn_metadata_kernels.BuildPageTablePositions,
            req_to_token=req_to_token,
            req_pool_indices_repeated=_ri(0, num_pool, (num_q,), torch.int32, g),
            seq_lens_casual=_ri(1, pool_len, (num_q,), torch.int64, g),
            max_seq_len=max_seq_len,
            page_size=page_size,
            swa_window=128,
        )


def _case_qo_indptr(tc):
    torch.manual_seed(17)
    cls = ragged_verify_kernels.BuildQoIndptr
    for dtype in (torch.int32, torch.int64):
        verify_lens = _ri(1, 8, (129,), dtype)  # straddles the 128 block
        ref = cls.torch(verify_lens=verify_lens)
        got = cls.triton(verify_lens=verify_lens.to(torch.int32))
        tc._eq(got, ref)
    # Aliasing regression: the two outputs must not share storage.
    vl = torch.tensor([3, 1, 5], device=DEVICE, dtype=torch.int32)
    got = cls.triton(verify_lens=vl)
    got.extend_start_loc.fill_(-7)
    tc.assertEqual(got.qo_indptr[:2].tolist(), [0, 3])


def _case_sample_step_tokens(tc):
    torch.manual_seed(18)
    cls = dspark_draft_model.SampleStepTokens
    # Injected noise makes stochastic sampling exactly comparable.
    for vocab, dtype in ((130000, torch.bfloat16), (5003, torch.float32)):
        bs = 3
        tc._parity(
            cls,
            step_logits=(torch.randn(bs, vocab, device=DEVICE) * 4.0).to(dtype),
            temperatures=torch.rand(bs, device=DEVICE) + 0.5,
            greedy_mask=(torch.arange(bs, device=DEVICE) % 2) == 0,
            exp_noise=torch.empty(bs, vocab, device=DEVICE).exponential_(1),
        )
    # Greedy tie straddling a triton block boundary picks the smaller index.
    logits = torch.zeros(1, 2050, device=DEVICE)
    logits[0, 1000] = logits[0, 1100] = 5.0
    tokens = cls.triton(
        step_logits=logits,
        temperatures=torch.tensor([1.0], device=DEVICE),
        greedy_mask=torch.tensor([True], device=DEVICE),
        exp_noise=torch.ones(1, 2050, device=DEVICE),
    )
    tc.assertEqual(tokens.item(), 1000)
    # Non-contiguous strided cropped view must match its contiguous copy.
    view = (torch.randn(2, 129536, device=DEVICE) * 4.0)[:, :VOCAB]
    tc.assertFalse(view.is_contiguous())
    kw = dict(
        temperatures=torch.rand(2, device=DEVICE) + 0.5,
        greedy_mask=torch.tensor([True, False], device=DEVICE),
        exp_noise=torch.empty(2, VOCAB, device=DEVICE).exponential_(1),
    )
    tc._eq(
        cls.triton(step_logits=view, **kw),
        cls.triton(step_logits=view.contiguous(), **kw),
    )


def _case_scatter_compact_to_strided(tc):
    torch.manual_seed(19)
    t, bs, dim = 6, 8, 4096
    verify_lens = _ri(1, t + 1, (bs,), torch.int32)
    total = int(verify_lens.sum().item())
    for graph_num_tokens in (total, bs * t):  # exact and bucket padding
        compact = torch.randn(
            graph_num_tokens, dim, dtype=torch.bfloat16, device=DEVICE
        )
        tc._parity(
            dspark_verify_window.ScatterCompactToStrided,
            compact=compact,
            layout=_layout(verify_lens, graph_num_tokens),
            fill_value=0.0,
            verify_num_draft_tokens=t,
        )


def _case_schedule_verify_lens_topk(tc):
    torch.manual_seed(20)
    gamma, bs = 5, 64
    cfg = DSparkScheduleConfig(gamma=gamma)
    cls = dspark_schedule.ScheduleVerifyLensTopk
    base = torch.rand(bs, gamma, device=DEVICE)
    confidences = (
        torch.full((bs, gamma), 0.5, device=DEVICE),  # all-ties
        (base * 4).floor() / 4,  # coarse quantization
        torch.where(base < 0.3, torch.zeros_like(base), base),  # invalid zeros
    )
    for confidence in confidences:
        for budget in (0, 1, 3, 7, 1000):
            tc._parity(cls, confidence=confidence, budget=budget, cfg=cfg)


def _case_softmax_temp(tc):
    g = torch.Generator(device=DEVICE).manual_seed(21)
    cls = dspark_accept.SoftmaxTemp
    # bf16 logits, non-power-of-two rows_per_request, full vocab.
    logits = (torch.randn(56, VOCAB, device=DEVICE, generator=g) * 8.0).to(
        torch.bfloat16
    )
    temps = (torch.rand(8, device=DEVICE, generator=g) * 1.5 + 0.05).float()
    ref = cls.torch(logits=logits, temperatures=temps, rows_per_request=7)
    got = cls.triton(logits=logits, temperatures=temps, rows_per_request=7)
    tc.assertEqual(got.dtype, torch.float32)
    torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-6)
    torch.testing.assert_close(
        got.sum(dim=-1), torch.ones_like(got.sum(dim=-1)), rtol=1e-5, atol=1e-5
    )
    # Column-shaped (bs, 1) temperatures.
    logits2 = torch.randn(6, 512, device=DEVICE, generator=g).to(torch.bfloat16)
    temps2 = (torch.rand(2, 1, device=DEVICE, generator=g) + 0.3).float()
    ref2 = cls.torch(logits=logits2, temperatures=temps2, rows_per_request=3)
    got2 = cls.triton(logits=logits2, temperatures=temps2, rows_per_request=3)
    torch.testing.assert_close(got2, ref2, rtol=1e-5, atol=1e-7)


_CASES = [
    ("accept_greedy", _case_accept_greedy),
    ("accept_sampling", _case_accept_sampling),
    ("build_block_seq_lens_causal", _case_build_block_seq_lens_causal),
    ("build_out_tokens", _case_build_out_tokens),
    ("build_ragged_verify_window", _case_build_ragged_verify_window),
    ("build_step_local", _case_build_step_local),
    ("cap_correct_len", _case_cap_correct_len),
    ("causal_swa_page_indices", _case_causal_swa_page_indices),
    ("commit_inject_layout", _case_commit_inject_layout),
    ("commit_kv_proj", _case_commit_kv_proj),
    ("compact_layout", _case_compact_layout),
    ("swa_page_indices", _case_swa_page_indices),
    ("expand_prefill_causally", _case_expand_prefill_causally),
    ("finalize_accept_lens", _case_finalize_accept_lens),
    ("mixed_accept_select", _case_mixed_accept_select),
    ("padded_to_bucket", _case_padded_to_bucket),
    ("page_table_positions", _case_page_table_positions),
    ("qo_indptr", _case_qo_indptr),
    ("sample_step_tokens", _case_sample_step_tokens),
    ("scatter_compact_to_strided", _case_scatter_compact_to_strided),
    ("schedule_verify_lens_topk", _case_schedule_verify_lens_topk),
    ("softmax_temp", _case_softmax_temp),
]


class TestDsparkKernelParity(CustomTestCase):
    def _eq(self, got, ref):
        """Exact comparison of tensors, tuples, and msgspec result structs."""
        if isinstance(ref, tuple):
            for g, r in zip(got, ref):
                self._eq(g, r)
        elif hasattr(ref, "__struct_fields__"):
            for name in ref.__struct_fields__:
                self._eq(getattr(got, name), getattr(ref, name))
        elif isinstance(ref, torch.Tensor):
            self.assertEqual(got.dtype, ref.dtype)
            self.assertTrue(torch.equal(got, ref))
        else:
            self.assertEqual(got, ref)

    def _parity(self, cls, **kw):
        got, ref = cls.triton(**kw), cls.torch(**kw)
        self._eq(got, ref)
        return got, ref

    def test_all_kernels_triton_matches_torch(self):
        for name, case in _CASES:
            with self.subTest(kernel=name):
                case(self)


if __name__ == "__main__":
    unittest.main()
