sgl-project--sglang/python/sglang/jit_kernel/tests/kv_canary/_differential.py

from __future__ import annotations

from dataclasses import dataclass, replace
from typing import Any, Callable, Iterator, Optional

import torch

from sglang.jit_kernel.kv_canary import consts
from sglang.jit_kernel.kv_canary.plan import launch_canary_plan_kernels
from sglang.jit_kernel.kv_canary.plan_ref import (
    launch_canary_plan_kernels_torch_reference,
)
from sglang.jit_kernel.kv_canary.verify import (
    CanaryLaunchTag,
    RealKvSource,
    VerifyOrWriteContext,
    VerifyPlan,
    launch_canary_verify_kernel,
)
from sglang.jit_kernel.kv_canary.verify_ref import (
    launch_canary_verify_kernel_torch_reference,
)
from sglang.jit_kernel.kv_canary.write import WritePlan, launch_canary_write_kernel
from sglang.jit_kernel.kv_canary.write_ref import (
    launch_canary_write_kernel_torch_reference,
)
from sglang.jit_kernel.tests.kv_canary._canary_helpers import (
    FakeViolationLog,
    assert_canary_buf_equal,
    assert_canary_state_equal,
    make_log_pair,
)

_DEVICE = torch.device("cuda")


def _run_both_plan(
    *,
    triton_verify: VerifyPlan,
    triton_write: WritePlan,
    ref_verify: VerifyPlan,
    ref_write: WritePlan,
    req_pool_indices: torch.Tensor,
    prefix_lens: torch.Tensor,
    extend_seq_lens: torch.Tensor,
    req_to_token: torch.Tensor,
    extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
    swa_window_size: int,
    full_to_swa_index_mapping: Optional[torch.Tensor],
    assert_equal: bool = True,
    active_verify_entries: Optional[int] = None,
    active_write_reqs: Optional[int] = None,
    req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None,
    kv_token_id_vs_position_offset: int = 0,
) -> None:
    _ = extras
    verify_capacity = int(triton_verify.verify_slot_indices.shape[0])
    # Default lens to "no tighter bound than pool width" so existing kernel tests that
    # only care about gather wiring keep their old semantics without each call site
    # explicitly building a per-req lens tensor.
    if (
        req_to_verify_expected_tokens is not None
        and req_to_verify_expected_tokens_valid_lens is None
    ):
        req_to_verify_expected_tokens_valid_lens = torch.full(
            (int(req_pool_indices.shape[0]),),
            int(req_to_verify_expected_tokens.shape[1]),
            dtype=torch.int64,
            device=req_pool_indices.device,
        )
    launch_canary_plan_kernels(
        verify_plan_out=triton_verify,
        write_plan_out=triton_write,
        req_pool_indices=req_pool_indices,
        prefix_lens=prefix_lens,
        extend_seq_lens=extend_seq_lens,
        req_to_token=req_to_token,
        swa_window_size=swa_window_size,
        full_to_swa_index_mapping=full_to_swa_index_mapping,
        verify_capacity=verify_capacity,
        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
        req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
    )
    launch_canary_plan_kernels_torch_reference(
        verify_plan_out=ref_verify,
        write_plan_out=ref_write,
        req_pool_indices=req_pool_indices,
        prefix_lens=prefix_lens,
        extend_seq_lens=extend_seq_lens,
        req_to_token=req_to_token,
        swa_window_size=swa_window_size,
        full_to_swa_index_mapping=full_to_swa_index_mapping,
        verify_capacity=int(ref_verify.verify_slot_indices.shape[0]),
        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
        req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
    )
    torch.cuda.synchronize()

    if assert_equal:
        _assert_plans_byte_equal(
            triton_verify=triton_verify,
            triton_write=triton_write,
            ref_verify=ref_verify,
            ref_write=ref_write,
            active_verify_entries=active_verify_entries,
            active_write_reqs=active_write_reqs,
        )


def _assert_plans_byte_equal(
    *,
    triton_verify: VerifyPlan,
    triton_write: WritePlan,
    ref_verify: VerifyPlan,
    ref_write: WritePlan,
    active_verify_entries: Optional[int] = None,
    active_write_reqs: Optional[int] = None,
) -> None:
    """Byte-equal check on (Triton vs ref) plan outputs.

    Optional ``active_verify_entries`` / ``active_write_reqs`` truncate the comparison to the meaningful
    prefix; tail entries past the active count are kernel-undefined and need not match byte-equal.
    """
    n_verify = (
        active_verify_entries
        if active_verify_entries is not None
        else int(triton_verify.verify_num_valid[0].item())
    )
    n_verify_ref = int(ref_verify.verify_num_valid[0].item())
    assert (
        n_verify == n_verify_ref
    ), f"verify_num_valid diverged: triton={n_verify} ref={n_verify_ref}"
    # When total_verify > VERIFY_CAPACITY the offsets kernel clears verify_enable and
    # plan_entries skips its scatter — leaving verify_slot_indices/positions/prev_slot_indices
    # as whatever the (torch.empty) allocation contained. Skip the byte-equal probe in that
    # case; verify_num_valid being clamped + verify_enable=0 is the contract here.
    triton_enable = int(triton_verify.enable[0].item())
    ref_enable = int(ref_verify.enable[0].item())
    assert (
        triton_enable == ref_enable
    ), f"verify_enable diverged: triton={triton_enable} ref={ref_enable}"
    if n_verify > 0 and triton_enable != 0:
        assert torch.equal(
            triton_verify.verify_slot_indices[:n_verify],
            ref_verify.verify_slot_indices[:n_verify],
        )
        assert torch.equal(
            triton_verify.verify_expected_tokens[:n_verify],
            ref_verify.verify_expected_tokens[:n_verify],
        )
        assert torch.equal(
            triton_verify.verify_expected_positions[:n_verify],
            ref_verify.verify_expected_positions[:n_verify],
        )
        assert torch.equal(
            triton_verify.verify_prev_slot_indices[:n_verify],
            ref_verify.verify_prev_slot_indices[:n_verify],
        )

    n_write = (
        active_write_reqs
        if active_write_reqs is not None
        else int(triton_write.write_num_valid_reqs[0].item())
    )
    n_write_ref = int(ref_write.write_num_valid_reqs[0].item())
    assert (
        n_write == n_write_ref
    ), f"write_num_valid_reqs diverged: triton={n_write} ref={n_write_ref}"
    assert torch.equal(
        triton_write.write_offsets[: n_write + 1],
        ref_write.write_offsets[: n_write + 1],
    )
    if n_write > 0:
        assert torch.equal(
            triton_write.write_seed_slot_indices[:n_write],
            ref_write.write_seed_slot_indices[:n_write],
        )


def _run_both_verify(
    *,
    cuda_canary_buf: torch.Tensor,
    ref_canary_buf: torch.Tensor,
    plan_cuda,
    plan_ref,
    cuda_log: FakeViolationLog,
    ref_log: FakeViolationLog,
    real_kv_sources_cuda: tuple[RealKvSource, ...],
    real_kv_sources_ref: tuple[RealKvSource, ...],
    real_kv_hash_mode: consts.RealKvHashMode,
    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
    assert_equal: bool = True,
    check_verify_expected_token: bool = True,
) -> None:
    launch_canary_verify_kernel(
        context=VerifyOrWriteContext(
            canary_buf=cuda_canary_buf,
            kernel_kind=kernel_kind,
            violation_ring=cuda_log.ring,
            violation_write_index=cuda_log.write_index,
            slot_run_counter=cuda_log.slot_run_counter,
            kernel_run_counter=cuda_log.kernel_run_counter,
            enable_chain_position_assert=cuda_log.enable_chain_position_assert,
            real_kv_sources=real_kv_sources_cuda,
            real_kv_hash_mode=real_kv_hash_mode,
        ),
        plan=plan_cuda,
        check_verify_expected_token=check_verify_expected_token,
    )
    launch_canary_verify_kernel_torch_reference(
        context=VerifyOrWriteContext(
            canary_buf=ref_canary_buf,
            kernel_kind=kernel_kind,
            violation_ring=ref_log.ring,
            violation_write_index=ref_log.write_index,
            slot_run_counter=ref_log.slot_run_counter,
            kernel_run_counter=ref_log.kernel_run_counter,
            enable_chain_position_assert=ref_log.enable_chain_position_assert,
            real_kv_sources=real_kv_sources_ref,
            real_kv_hash_mode=real_kv_hash_mode,
        ),
        plan=plan_ref,
        check_verify_expected_token=check_verify_expected_token,
    )
    torch.cuda.synchronize()

    if assert_equal:
        assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)


def _run_both_write(
    *,
    cuda_canary_buf: torch.Tensor,
    ref_canary_buf: torch.Tensor,
    plan_cuda,
    plan_ref,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    out_cache_loc: torch.Tensor,
    enable_write_verify_inputs: bool,
    expected_input_tokens: torch.Tensor,
    expected_input_positions: torch.Tensor,
    cuda_log: FakeViolationLog,
    ref_log: FakeViolationLog,
    real_kv_sources_cuda: tuple[RealKvSource, ...],
    real_kv_sources_ref: tuple[RealKvSource, ...],
    real_kv_hash_mode: consts.RealKvHashMode,
    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
    assert_equal: bool = True,
) -> None:
    expected_tokens_for_launch = (
        expected_input_tokens if enable_write_verify_inputs else None
    )
    expected_positions_for_launch = (
        expected_input_positions if enable_write_verify_inputs else None
    )
    launch_canary_write_kernel(
        context=VerifyOrWriteContext(
            canary_buf=cuda_canary_buf,
            kernel_kind=kernel_kind,
            violation_ring=cuda_log.ring,
            violation_write_index=cuda_log.write_index,
            slot_run_counter=cuda_log.slot_run_counter,
            kernel_run_counter=cuda_log.kernel_run_counter,
            enable_chain_position_assert=cuda_log.enable_chain_position_assert,
            real_kv_sources=real_kv_sources_cuda,
            real_kv_hash_mode=real_kv_hash_mode,
        ),
        plan=plan_cuda,
        input_ids=input_ids,
        positions=positions,
        out_cache_loc=out_cache_loc,
        enable_write_input_assert=enable_write_verify_inputs,
        expected_input_tokens=expected_tokens_for_launch,
        expected_input_positions=expected_positions_for_launch,
    )
    launch_canary_write_kernel_torch_reference(
        context=VerifyOrWriteContext(
            canary_buf=ref_canary_buf,
            kernel_kind=kernel_kind,
            violation_ring=ref_log.ring,
            violation_write_index=ref_log.write_index,
            slot_run_counter=ref_log.slot_run_counter,
            kernel_run_counter=ref_log.kernel_run_counter,
            enable_chain_position_assert=ref_log.enable_chain_position_assert,
            real_kv_sources=real_kv_sources_ref,
            real_kv_hash_mode=real_kv_hash_mode,
        ),
        plan=plan_ref,
        input_ids=input_ids,
        positions=positions,
        out_cache_loc=out_cache_loc,
        enable_write_input_assert=enable_write_verify_inputs,
        expected_input_tokens=expected_tokens_for_launch,
        expected_input_positions=expected_positions_for_launch,
    )
    torch.cuda.synchronize()

    if assert_equal:
        assert_canary_buf_equal(buf_a=cuda_canary_buf, buf_b=ref_canary_buf)
        assert_canary_state_equal(log_a=cuda_log, log_b=ref_log)


@dataclass(frozen=True, slots=True, kw_only=True)
class ShrinkResult:
    inputs: Any
    mutations_applied: list[str]


def shrink_inputs(
    inputs: Any,
    *,
    check_fn: Callable[[Any], bool],
    max_iterations: int = 50,
) -> ShrinkResult:
    """Greedy 1-step minify for a fuzz inputs dataclass.

    ``check_fn(candidate)`` returns True when ``candidate`` still reproduces the failure. Each round
    yields candidate-simpler-than-current mutations through ``_yield_simpler``; the first accepted
    candidate becomes the new current. Iteration stops when no mutation is accepted or ``max_iterations``
    is reached.
    """
    current = inputs
    applied: list[str] = []
    for _ in range(max_iterations):
        improved = False
        for label, candidate in _yield_simpler(current):
            try:
                still_fails = check_fn(candidate)
            except Exception:
                still_fails = False
            if still_fails:
                current = candidate
                applied.append(label)
                improved = True
                break
        if not improved:
            break
    return ShrinkResult(inputs=current, mutations_applied=applied)


def _yield_simpler(inputs: Any) -> Iterator[tuple[str, Any]]:
    """Yield (label, simpler_candidate) tuples for generic fuzz-input minifiers.

    The candidates touch only well-known field names; an inputs dataclass that lacks a field will simply
    have that mutation skipped. No kernel-specific knowledge is encoded here so the same shrinker drives
    Plan / Verify / Write fuzz failures uniformly.
    """
    fields = {
        f: getattr(inputs, f) for f in inputs.__dataclass_fields__  # type: ignore[attr-defined]
    }

    def emit(label: str, **overrides: Any) -> Iterator[tuple[str, Any]]:
        candidate = replace(inputs, **overrides)
        yield label, candidate

    bs_field = (
        "req_pool_indices"
        if "req_pool_indices" in fields
        else ("input_ids" if "input_ids" in fields else None)
    )
    if bs_field is not None and isinstance(fields[bs_field], torch.Tensor):
        tensor = fields[bs_field]
        if tensor.numel() > 1:
            new_len = tensor.numel() - 1
            related_tensors_overrides: dict[str, Any] = {}
            for name in (
                "req_pool_indices",
                "prefix_lens",
                "extend_seq_lens",
                "input_ids",
                "positions",
                "out_cache_loc",
                "expected_input_tokens",
                "expected_input_positions",
            ):
                t = fields.get(name)
                if (
                    isinstance(t, torch.Tensor)
                    and t.numel() >= new_len
                    and t.dim() == 1
                ):
                    related_tensors_overrides[name] = t[:new_len].contiguous()
            if related_tensors_overrides:
                yield from emit("drop_last_row", **related_tensors_overrides)

    if "swa_window_size" in fields and isinstance(fields["swa_window_size"], int):
        if fields["swa_window_size"] != 0:
            yield from emit(
                "swa_off", swa_window_size=0, full_to_swa_index_mapping=None
            )

    if "extras_count" in fields and isinstance(fields["extras_count"], int):
        if fields["extras_count"] > 0:
            yield from emit("extras_zero", extras_count=0)

    if "real_kv_hash_mode" in fields:
        cur = fields["real_kv_hash_mode"]
        if hasattr(cur, "value"):
            cls = cur.__class__
            if int(cur) == 2:
                yield from emit("hash_mode_bit", real_kv_hash_mode=cls(1))
            elif int(cur) == 1:
                yield from emit("hash_mode_off", real_kv_hash_mode=cls(0))

    if "real_kv_sources" in fields:
        srcs = fields["real_kv_sources"]
        if isinstance(srcs, tuple) and len(srcs) > 1:
            yield from emit("sources_to_one", real_kv_sources=srcs[:1])

    if "enable_write_verify_inputs" in fields:
        cur = fields["enable_write_verify_inputs"]
        if hasattr(cur, "value") and int(cur) != 0:
            cls = cur.__class__
            yield from emit("pseudo_off", enable_write_verify_inputs=cls(0))

    for name in ("verify_capacity", "write_req_capacity"):
        if name in fields and isinstance(fields[name], int):
            current_value = fields[name]
            if current_value > 8:
                yield from emit(f"shrink_{name}", **{name: max(8, current_value // 2)})


def run_verify_diff(
    *,
    buf_pair: tuple[torch.Tensor, torch.Tensor],
    plan_pair: tuple[VerifyPlan, VerifyPlan],
    real_kv_sources_pair: tuple[tuple[RealKvSource, ...], tuple[RealKvSource, ...]] = (
        (),
        (),
    ),
    real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
    device: torch.device = _DEVICE,
    assert_equal: bool = True,
    check_verify_expected_token: bool = True,
) -> tuple[FakeViolationLog, FakeViolationLog]:
    """Thin wrapper around ``_run_both_verify`` that creates a fresh log pair and packs (cuda, ref)
    buf/plan/source arguments into 2-tuples to drop ~8 lines of boilerplate per call site.
    """
    cuda_log, ref_log = make_log_pair(device=device)
    _run_both_verify(
        cuda_canary_buf=buf_pair[0],
        ref_canary_buf=buf_pair[1],
        plan_cuda=plan_pair[0],
        plan_ref=plan_pair[1],
        cuda_log=cuda_log,
        ref_log=ref_log,
        real_kv_sources_cuda=real_kv_sources_pair[0],
        real_kv_sources_ref=real_kv_sources_pair[1],
        real_kv_hash_mode=real_kv_hash_mode,
        kernel_kind=kernel_kind,
        assert_equal=assert_equal,
        check_verify_expected_token=check_verify_expected_token,
    )
    return cuda_log, ref_log


def run_write_diff(
    *,
    buf_pair: tuple[torch.Tensor, torch.Tensor],
    plan_pair: tuple[WritePlan, WritePlan],
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    out_cache_loc: torch.Tensor,
    expected_input_tokens: torch.Tensor,
    expected_input_positions: torch.Tensor,
    enable_write_verify_inputs: bool = False,
    real_kv_sources_pair: tuple[tuple[RealKvSource, ...], tuple[RealKvSource, ...]] = (
        (),
        (),
    ),
    real_kv_hash_mode: consts.RealKvHashMode = consts.RealKvHashMode.NONE,
    kernel_kind: CanaryLaunchTag = CanaryLaunchTag.HEAD_K_FULL,
    device: torch.device = _DEVICE,
    assert_equal: bool = True,
) -> tuple[FakeViolationLog, FakeViolationLog]:
    """Thin wrapper around ``_run_both_write`` that creates a fresh log pair and packs (cuda, ref)
    buf/plan/source arguments into 2-tuples to drop ~10 lines of boilerplate per call site.
    """
    cuda_log, ref_log = make_log_pair(device=device)
    _run_both_write(
        cuda_canary_buf=buf_pair[0],
        ref_canary_buf=buf_pair[1],
        plan_cuda=plan_pair[0],
        plan_ref=plan_pair[1],
        input_ids=input_ids,
        positions=positions,
        out_cache_loc=out_cache_loc,
        enable_write_verify_inputs=enable_write_verify_inputs,
        expected_input_tokens=expected_input_tokens,
        expected_input_positions=expected_input_positions,
        cuda_log=cuda_log,
        ref_log=ref_log,
        real_kv_sources_cuda=real_kv_sources_pair[0],
        real_kv_sources_ref=real_kv_sources_pair[1],
        real_kv_hash_mode=real_kv_hash_mode,
        kernel_kind=kernel_kind,
        assert_equal=assert_equal,
    )
    return cuda_log, ref_log


def run_plan_diff(
    *,
    plan_pair: tuple[tuple[VerifyPlan, WritePlan], tuple[VerifyPlan, WritePlan]],
    req_pool_indices: torch.Tensor,
    prefix_lens: torch.Tensor,
    extend_seq_lens: torch.Tensor,
    req_to_token: torch.Tensor,
    extras: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
    swa_window_size: int = 0,
    full_to_swa_index_mapping: Optional[torch.Tensor] = None,
    assert_equal: bool = True,
    active_verify_entries: Optional[int] = None,
    active_write_reqs: Optional[int] = None,
    req_to_verify_expected_tokens: Optional[torch.Tensor] = None,
    req_to_verify_expected_tokens_valid_lens: Optional[torch.Tensor] = None,
    kv_token_id_vs_position_offset: int = 0,
) -> None:
    """Thin wrapper around ``_run_both_plan`` that unpacks ``((triton_v, triton_w), (ref_v, ref_w))``
    plan pairs to drop the per-call-site ``triton_verify=.../triton_write=.../ref_verify=...`` block.
    """
    (triton_verify, triton_write), (ref_verify, ref_write) = plan_pair
    _run_both_plan(
        triton_verify=triton_verify,
        triton_write=triton_write,
        ref_verify=ref_verify,
        ref_write=ref_write,
        req_pool_indices=req_pool_indices,
        prefix_lens=prefix_lens,
        extend_seq_lens=extend_seq_lens,
        req_to_token=req_to_token,
        extras=extras,
        swa_window_size=swa_window_size,
        full_to_swa_index_mapping=full_to_swa_index_mapping,
        assert_equal=assert_equal,
        active_verify_entries=active_verify_entries,
        active_write_reqs=active_write_reqs,
        req_to_verify_expected_tokens=req_to_verify_expected_tokens,
        req_to_verify_expected_tokens_valid_lens=req_to_verify_expected_tokens_valid_lens,
        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
    )