Add the KV-canary verify JIT kernel and reference implementation (#26805)

python/sglang/jit_kernel/tests/kv_canary/_constants.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
+
+# Default fixture sizes — small enough for fast tests, large enough that ring overflow / multi-req cases
+# stay realistic without bloating the assertion surface.
+DEFAULT_RING_CAPACITY: int = 64
+DEFAULT_NUM_SLOTS: int = 32
+DEFAULT_SLOT_STRIDE_BYTES: int = CANARY_SLOT_BYTES
+
+_U64_MASK: int = (1 << 64) - 1
+_I64_SIGN_BIT: int = 1 << 63

python/sglang/jit_kernel/tests/kv_canary/test_const_sync.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
+
+
+_CONSTS_CUH: Path = (
+    Path(__file__).resolve().parents[2] / "csrc" / "kv_canary" / "consts.cuh"
+)
+
+
+def _camel_to_upper_snake(name: str) -> str:
+    return re.sub(r"([A-Z])", r"_\1", name).lstrip("_").upper()
+
+
+def _decode(expr: str) -> int:
+    expr = expr.strip().rstrip("UuLl")
+    if "<<" in expr:
+        return 1 << int(expr.split("<<")[1].strip())
+    return int(expr, 0)
+
+
+def _parse_constexpr_ints(source: str) -> dict[str, int]:
+    pattern = re.compile(r"constexpr\s+(?:[\w:]+)\s+(k[A-Za-z]\w*)\s*=\s*([^;]+);")
+    return {name: _decode(rhs) for name, rhs in pattern.findall(source)}
+
+
+def _parse_enum_class(source: str, enum_name: str) -> dict[str, int]:
+    pattern = re.compile(
+        r"enum\s+class\s+" + re.escape(enum_name) + r"\s*:\s*[^\{]+\{([^}]+)\}"
+    )
+    body = pattern.search(source).group(1)
+    member_re = re.compile(r"(k[A-Za-z]\w*)\s*=\s*([^,]+)")
+    return {name: _decode(rhs) for name, rhs in member_re.findall(body)}
+
+
+def test_int_consts_sync() -> None:
+    cpp = _parse_constexpr_ints(_CONSTS_CUH.read_text(encoding="utf-8"))
+    cpp_normalized = {_camel_to_upper_snake(n[1:]): v for n, v in cpp.items()}
+    py = {
+        n: v
+        for n, v in vars(consts).items()
+        if isinstance(v, int) and not isinstance(v, bool) and not n.startswith("_")
+    }
+    assert cpp_normalized == py
+
+
+def test_enums_sync() -> None:
+    cuh = _CONSTS_CUH.read_text(encoding="utf-8")
+    for enum_name in ("FailReason",):
+        cpp_members = _parse_enum_class(cuh, enum_name)
+        py_enum = getattr(consts, enum_name)
+        cpp_normalized = {
+            _camel_to_upper_snake(n[1:]): v for n, v in cpp_members.items()
+        }
+        py_normalized = {m.name: int(m.value) for m in py_enum}
+        assert cpp_normalized == py_normalized
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main([__file__, "-v"]))

python/sglang/jit_kernel/tests/kv_canary/test_utils.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from sglang.jit_kernel.benchmark.kv_canary.utils import (
+    MAX_EXTEND_TOKENS_PER_FORWARD,
+    build_fast_matrix_cases,
+    build_full_matrix_cases,
+    cases_to_x_vals,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+
+
+def test_fast_matrix_cases_include_e2e_decode_and_chunked_prefill_scenarios() -> None:
+    cases = build_fast_matrix_cases()
+    scenarios = {case.scenario for case in cases}
+
+    assert {
+        "e2e_decode_steady",
+        "e2e_decode_tail",
+        "e2e_prefill_chunk_first",
+        "e2e_prefill_chunk_second",
+        "e2e_prefill_chunk_mid",
+        "e2e_prefill_chunk_last",
+    } <= scenarios
+
+
+def test_extend_cases_are_bounded_to_scheduler_chunk_size() -> None:
+    cases = build_full_matrix_cases()
+    bad_cases = [
+        case
+        for case in cases
+        if case.mode == "extend"
+        and case.bs * case.extend_len > MAX_EXTEND_TOKENS_PER_FORWARD
+    ]
+
+    assert bad_cases == []
+
+
+def test_cases_to_x_vals_includes_scenario_axis() -> None:
+    case = build_fast_matrix_cases()[0]
+
+    x_vals = cases_to_x_vals([case])
+
+    assert x_vals == [
+        (
+            case.scenario,
+            case.bs,
+            case.prefix_len,
+            case.mode,
+            case.extend_len,
+            case.pool_kind,
+        )
+    ]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main([__file__, "-v"]))