chore: import upstream snapshot with attribution

python/sglang/jit_kernel/kv_canary/verify_ref.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.kv_canary.consts import splitmix64, splitmix64_mix3
+from sglang.jit_kernel.kv_canary.verify import (
+    RealKvSource,
+    VerifyOrWriteContext,
+    VerifyPlan,
+)
+
+_U64_MASK: int = (1 << 64) - 1
+_I64_SIGN_BIT: int = 1 << 63
+
+
+def launch_canary_verify_kernel_torch_reference(
+    *,
+    context: VerifyOrWriteContext,
+    plan: VerifyPlan,
+    check_verify_expected_token: bool,
+) -> None:
+    canary_buf = context.canary_buf
+    kernel_kind = context.kernel_kind
+    violation_ring = context.violation_ring
+    violation_write_index = context.violation_write_index
+    slot_run_counter = context.slot_run_counter
+    kernel_run_counter = context.kernel_run_counter
+    real_kv_sources = context.real_kv_sources
+    real_kv_hash_mode = context.real_kv_hash_mode
+
+    work_device = torch.device("cpu")
+
+    kernel_run_counter.add_(1)
+
+    enable = int(plan.enable.detach().to("cpu").item())
+    if enable == 0:
+        return
+
+    num_valid = int(
+        plan.verify_slot_indices.new_empty(()).copy_(plan.verify_num_valid[0]).item()
+    )
+    capacity = int(plan.verify_slot_indices.shape[0])
+    active = max(0, min(num_valid, capacity))
+    if active <= 0:
+        return
+
+    slot_indices_host = plan.verify_slot_indices[:active].to(
+        device=work_device, dtype=torch.int64
+    )
+    if check_verify_expected_token:
+        expected_input_ids_host = plan.verify_expected_tokens[:active].to(
+            device=work_device, dtype=torch.int64
+        )
+    else:
+        expected_input_ids_host = torch.full(
+            (active,), -1, dtype=torch.int64, device=work_device
+        )
+    expected_positions_host = plan.verify_expected_positions[:active].to(
+        device=work_device, dtype=torch.int64
+    )
+    prev_slot_indices_host = plan.verify_prev_slot_indices[:active].to(
+        device=work_device, dtype=torch.int64
+    )
+
+    slot_run_counter.add_(active)
+
+    kept_slots: list[int] = []
+    kept_expected_positions: list[int] = []
+    kept_expected_input_ids: list[int] = []
+    kept_prev_slots: list[int] = []
+    for k in range(active):
+        s = int(slot_indices_host[k].item())
+        # Skip SGLang's padded-token dummy KV slot so unfilled req_to_token entries (zero-initialized) do not
+        # produce spurious chain_hash / position violations.
+        if s != consts.TOKEN_TO_KV_SLOT_PADDING:
+            kept_slots.append(s)
+            kept_expected_positions.append(int(expected_positions_host[k].item()))
+            kept_expected_input_ids.append(int(expected_input_ids_host[k].item()))
+            kept_prev_slots.append(int(prev_slot_indices_host[k].item()))
+    active = len(kept_slots)
+    if active <= 0:
+        return
+    slot_indices_list: list[int] = kept_slots
+    expected_positions_list: list[int] = kept_expected_positions
+    expected_input_ids_list: list[int] = kept_expected_input_ids
+    prev_slot_indices_list: list[int] = kept_prev_slots
+
+    buf_i64 = canary_buf.detach().to(device=work_device).contiguous().view(torch.int64)
+    slot_stride_i64 = int(buf_i64.shape[1])
+    if slot_stride_i64 < 4:
+        raise ValueError(
+            f"kv-canary: canary_buf slot stride must hold at least 4 int64 fields, got {slot_stride_i64}"
+        )
+
+    violation_rows: list[list[int]] = []
+
+    for k in range(active):
+        slot_idx = slot_indices_list[k]
+        expected_position = expected_positions_list[k]
+        expected_input_id = expected_input_ids_list[k]
+        prev_slot = prev_slot_indices_list[k]
+
+        stored_token = int(buf_i64[slot_idx, consts.CANARY_FIELD_TOKEN].item())
+        stored_position = int(buf_i64[slot_idx, consts.CANARY_FIELD_POSITION].item())
+        stored_chain_hash = int(buf_i64[slot_idx, consts.CANARY_FIELD_PREV_HASH].item())
+        stored_real_kv_hash = int(
+            buf_i64[slot_idx, consts.CANARY_FIELD_REAL_KV_HASH].item()
+        )
+
+        prev_reachable = prev_slot != consts.TOKEN_TO_KV_SLOT_PADDING
+        if prev_reachable:
+            expected_chain_hash = _to_signed_int64(
+                compute_slot_hash(buf_i64, prev_slot)
+            )
+        else:
+            expected_chain_hash = stored_chain_hash
+
+        expected_real_kv_hash_u64 = _compute_real_kv_hash_scalar(
+            slot_idx=slot_idx,
+            real_kv_sources=real_kv_sources,
+            real_kv_hash_mode=real_kv_hash_mode,
+            work_device=work_device,
+        )
+        expected_real_kv_hash = _to_signed_int64(expected_real_kv_hash_u64)
+
+        fail_reason = consts.FailReason(0)
+        if prev_reachable and stored_chain_hash != expected_chain_hash:
+            fail_reason |= consts.FailReason.VERIFY_CHAIN_HASH_MISMATCH
+        if check_verify_expected_token:
+            if expected_input_id != -1 and stored_token != expected_input_id:
+                fail_reason |= consts.FailReason.VERIFY_TOKEN_MISMATCH
+        if stored_position != expected_position:
+            fail_reason |= consts.FailReason.VERIFY_POSITION_MISMATCH
+        if stored_real_kv_hash != expected_real_kv_hash:
+            fail_reason |= consts.FailReason.VERIFY_REAL_KV_HASH_MISMATCH
+
+        if fail_reason != consts.FailReason(0):
+            row = [0] * consts.VIOLATION_FIELDS
+            row[consts.VIOLATION_FIELD_KERNEL_KIND] = int(kernel_kind)
+            row[consts.VIOLATION_FIELD_SLOT_IDX] = slot_idx
+            row[consts.VIOLATION_FIELD_POSITION] = stored_position
+            row[consts.VIOLATION_FIELD_STORED_TOKEN] = stored_token
+            row[consts.VIOLATION_FIELD_EXPECTED_TOKEN] = expected_input_id
+            row[consts.VIOLATION_FIELD_STORED_CHAIN_HASH] = stored_chain_hash
+            row[consts.VIOLATION_FIELD_EXPECTED_AUX] = expected_chain_hash
+            row[consts.VIOLATION_FIELD_FAIL_REASON_BITS] = int(fail_reason)
+            violation_rows.append(row)
+
+    if len(violation_rows) == 0:
+        return
+
+    num_new_violations = len(violation_rows)
+    base_idx = int(
+        violation_write_index.new_empty(()).copy_(violation_write_index[0]).item()
+    )
+    ring_capacity = int(violation_ring.shape[0])
+
+    new_rows = torch.zeros(
+        (num_new_violations, consts.VIOLATION_FIELDS), dtype=torch.int64
+    )
+    for v, row in enumerate(violation_rows):
+        for f in range(consts.VIOLATION_FIELDS):
+            new_rows[v, f] = row[f]
+
+    write_count_in_ring = max(0, min(num_new_violations, ring_capacity - base_idx))
+    if write_count_in_ring > 0:
+        ring_host = violation_ring.detach().to(device=work_device)
+        ring_host[base_idx : base_idx + write_count_in_ring, :] = new_rows[
+            :write_count_in_ring, :
+        ]
+        violation_ring.copy_(ring_host.to(violation_ring.device))
+
+    violation_write_index[0] = violation_write_index[0] + num_new_violations
+
+
+def _to_signed_int64(value: int) -> int:
+    value &= _U64_MASK
+    if value >= _I64_SIGN_BIT:
+        value -= 1 << 64
+    return value
+
+
+def compute_slot_hash(buf_i64: torch.Tensor, source_slot_idx: int) -> int:
+    if source_slot_idx < 0:
+        return splitmix64(consts.CANARY_CHAIN_ANCHOR)
+    token = int(buf_i64[source_slot_idx, consts.CANARY_FIELD_TOKEN].item())
+    position = int(buf_i64[source_slot_idx, consts.CANARY_FIELD_POSITION].item())
+    prev_hash = int(buf_i64[source_slot_idx, consts.CANARY_FIELD_PREV_HASH].item())
+    return splitmix64_mix3(prev_hash, token, position)
+
+
+def _compute_real_kv_hash_scalar(
+    *,
+    slot_idx: int,
+    real_kv_sources: tuple[RealKvSource, ...],
+    real_kv_hash_mode: consts.RealKvHashMode,
+    work_device: torch.device,
+) -> int:
+    mode = int(real_kv_hash_mode)
+    if mode == int(consts.RealKvHashMode.NONE) or len(real_kv_sources) == 0:
+        return 0
+
+    acc: int = 0
+
+    for source in real_kv_sources:
+        page_size = source.page_size
+        num_bytes_per_token = source.num_bytes_per_token
+        read_bytes = source.read_bytes
+        tensor_u8 = (
+            source.tensor.detach().to(device=work_device).contiguous().view(torch.uint8)
+        )
+
+        row = slot_idx // page_size
+        col_within_page = slot_idx % page_size
+        col_start = col_within_page * num_bytes_per_token
+
+        effective_read_bytes = (
+            16 if mode == int(consts.RealKvHashMode.PARTIAL) else read_bytes
+        )
+        raw_bytes: list[int] = []
+        for b in range(effective_read_bytes):
+            raw_bytes.append(int(tensor_u8[row, col_start + b].item()))
+
+        source_hash = _splitmix64_fold_bytes_scalar(raw_bytes=raw_bytes)
+
+        combined = acc ^ source_hash
+        acc = splitmix64(combined)
+
+    return acc
+
+
+def _splitmix64_fold_bytes_scalar(*, raw_bytes: list[int]) -> int:
+    read_bytes = len(raw_bytes)
+    pad = (8 - read_bytes % 8) % 8
+    padded = raw_bytes + [0] * pad
+    num_words = len(padded) // 8
+
+    acc: int = 0
+    for w in range(num_words):
+        word: int = 0
+        for k in range(8):
+            word |= padded[w * 8 + k] << (8 * k)
+        word &= _U64_MASK
+        acc = splitmix64(acc ^ word)
+
+    return acc