Add the KV-canary install API and forward-path wiring (#26809)

python/sglang/srt/kv_canary/api.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import torch
+
+from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
+from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
+from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
+from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
+from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+def install_canary(
+    *,
+    server_args: "ServerArgs",
+    model_runner: "ModelRunner",
+) -> Optional[CanaryManager]:
+    config = CanaryConfig.from_env(server_args)
+    if config.mode is CanaryMode.NONE:
+        return None
+
+    assert server_args.disable_piecewise_cuda_graph, (
+        "kv-canary: piecewise cuda graph is not supported by the current "
+        "SingleForwardManager design; pass --disable-piecewise-cuda-graph "
+        "when canary is enabled"
+    )
+
+    device = torch.device(model_runner.device)
+    buffer_groups = attach_canary_buffers(
+        pool=model_runner.token_to_kv_pool,
+        config=config,
+        device=device,
+        kv_token_id_vs_position_offset=0,
+    )
+    launch_capacities = CanaryLaunchCapacities.from_args(
+        server_args=model_runner.server_args,
+        req_to_token_pool_size=model_runner.req_to_token_pool.size,
+        max_seq_len_per_req=model_runner.req_to_token_pool.req_to_token.shape[1],
+        pool_slot_count=model_runner.max_total_num_tokens,
+    )
+    swa_window_size = model_runner.sliding_window_size or 0
+    manager = CanaryManager(
+        config=config,
+        buffer_groups=buffer_groups,
+        device=device,
+        req_to_token_pool=model_runner.req_to_token_pool,
+        launch_capacities=launch_capacities,
+        swa_window_size=swa_window_size,
+    )
+
+    _patch_model_forward(model_runner=model_runner, manager=manager)
+
+    # Single-line summary of every knob that controls canary behavior at boot time.
+    # Disaggregation mode is included so PD logs are unambiguous about which side this is.
+    logger.info(
+        "install_canary: disaggregation_mode=%s config=%s "
+        "launch_capacities=%s n_buffer_groups=%d buffer_group_kinds=%s "
+        "swa_window_size=%d",
+        server_args.disaggregation_mode,
+        config,
+        launch_capacities,
+        len(buffer_groups),
+        [g.kind.name for g in buffer_groups],
+        swa_window_size,
+    )
+    return manager
+
+
+def _patch_model_forward(
+    *, model_runner: "ModelRunner", manager: CanaryManager
+) -> None:
+    def _with_canary_bracketing(original: Callable, *args: Any, **kwargs: Any) -> Any:
+        forward_batch = _extract_forward_batch(args, kwargs)
+        assert (
+            forward_batch is not None
+        ), "kv-canary: patched model.forward called without a ForwardBatch"
+
+        canary_pre_ops_output = manager.pre_ops_maybe_inside_graph(forward_batch)
+        output = original(*args, **kwargs)
+        manager.post_ops_maybe_inside_graph(forward_batch, canary_pre_ops_output)
+        return output
+
+    wrap_method(model_runner.model, "forward", wrapper=_with_canary_bracketing)
+
+
+def _extract_forward_batch(args, kwargs) -> Optional[ForwardBatch]:
+    if "forward_batch" in kwargs and isinstance(kwargs["forward_batch"], ForwardBatch):
+        return kwargs["forward_batch"]
+    for arg in args:
+        if isinstance(arg, ForwardBatch):
+            return arg
+    return None

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -1122,7 +1122,11 @@ class CudaGraphRunner:
 
             self.deepep_adapter.capture(is_extend_in_batch=False)
 
-            canary_ctx = contextlib.nullcontext()
+            canary_ctx = (
+                c.with_active_single_forward_manager(0)
+                if (c := self.model_runner.canary_manager) is not None
+                else contextlib.nullcontext()
+            )
             with canary_ctx:
                 for _ in range(2):
                     self.device_module.synchronize()

python/sglang/srt/model_executor/model_runner.py

@@ -105,6 +105,8 @@ from sglang.srt.eplb.expert_location import (
 )
 from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
 from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
+from sglang.srt.kv_canary.api import install_canary
+from sglang.srt.kv_canary.runner.canary_manager import context_tuple
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.attention_registry import (
     ATTENTION_BACKENDS,
@@ -751,6 +753,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         # Init memory pool and attention backends
         self.init_memory_pool(pre_model_load_memory)
 
+        # Must be called AFTER init_memory_pool so the pool object exists for
+        # canary to monkey-patch, and BEFORE init_device_graphs so warmup
+        # forwards captured into the graph see the patched pool methods.
+        self.canary_manager = install_canary(
+            server_args=server_args,
+            model_runner=self,
+        )
+
         # Init ngram embedding token table
         self.maybe_init_ngram_embedding()
 
@@ -828,6 +838,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
 
         self.prealloc_symmetric_memory_pool()
 
+        if self.canary_manager is not None and not self.is_draft_worker:
+            self.canary_manager.mark_init_finished()
+
     def adjust_hybrid_swa_layers_for_pp(self):
         if not self.is_hybrid_swa:
             return
@@ -3204,7 +3217,17 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             else contextlib.nullcontext()
         )
 
-        canary_ctx = contextlib.nullcontext()
+        canary_ctx = (
+            context_tuple(
+                c.with_ops_outside_graph(
+                    single_forward_indices=[0],
+                    maybe_inaccurate_forward_batch=forward_batch,
+                ),
+                c.with_active_single_forward_manager(0),
+            )
+            if not self.is_draft_worker and ((c := self.canary_manager) is not None)
+            else contextlib.nullcontext()
+        )
 
         with (
             canary_ctx,

python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py

@@ -310,6 +310,7 @@ class MockModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py

@@ -296,6 +296,7 @@ class DSAMockModelRunner(ModelRunner):
         else:
             spec_num_draft_tokens = 0
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py

@@ -320,6 +320,7 @@ class MockDSV4ModelRunner:
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py

@@ -302,6 +302,7 @@ class DualChunkMockModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py

@@ -205,6 +205,7 @@ class MockGDNModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         speculative_num_draft_tokens = (

python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py

@@ -211,6 +211,7 @@ class MockKDAModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         speculative_num_draft_tokens = (

python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py

@@ -220,6 +220,7 @@ class MockLightningModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         speculative_num_draft_tokens = (

python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py

@@ -315,6 +315,7 @@ class MockMamba2ModelRunner(ModelRunner):
         self.dtype = dtype
         self.kv_cache_dtype = dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         # MambaMixer2 asserts the layer_cache is a `SpeculativeState`

python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py

@@ -222,6 +222,7 @@ class MockMLAModelRunner(ModelRunner):
         # does the BF16->FP8 cast on the way in.
         self.kv_cache_dtype = torch.float8_e4m3fn if fp8_kv_cache else dtype
         self.gpu_id = 0
+        self.canary_manager = None
         self.page_size = case.page_size
         self.model_config = model_config
         self.tp_size = 1

test/registered/kv_canary/test_self_e2e_baseline.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import unittest
+
+from sglang.srt.kv_canary.config import CanaryMode
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kv_canary.e2e_base import CanaryE2EBase
+
+register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
+
+
+class _BaselineBase(CanaryE2EBase):
+    """No perturb, kv-canary=log. Server should run clean with no canary
+    violations and every request must come back 200."""
+
+    kv_canary_mode = CanaryMode.LOG
+    extra_env = {}
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        if cls is _BaselineBase:
+            raise unittest.SkipTest("abstract base; concrete subclasses set model_mode")
+        super().setUpClass()
+
+    def test_no_violation(self) -> None:
+        """Verify the baseline canary run completes without violations."""
+        for _ in range(self.workload_n_batches):
+            self.send_parallel_requests()
+        self.assert_no_violation(wait_seconds=2.0)
+
+
+class TestBaselineMha(_BaselineBase):
+    model_mode = "mha"
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_e2e_bench_speed.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import os
+import unittest
+from pathlib import Path
+from typing import ClassVar, Optional
+
+from sglang.srt.entrypoints.http_server import launch_server
+from sglang.srt.server_args import ServerArgs
+from sglang.test.bench_one_batch_server_internal import (
+    BenchArgs,
+    BenchOneCaseResult,
+    run_benchmark_internal,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER
+
+register_cuda_ci(est_time=600, stage="extra-a", runner_config="1-gpu-large")
+
+
+_QWEN3_MODEL = "Qwen/Qwen3-30B-A3B"
+_QWEN3_SCENARIO_MODEL = "qwen3-30b-a3b"
+
+_PROFILE_DIR_ENV = "SGLANG_KV_CANARY_PROFILE_DIR"
+_PROFILE_STEPS = 30
+_PROFILE_NO_GRAPH_OUTPUT_LEN = 3
+# start_profile blocks until num_steps server steps complete, so it must be <= actual decode steps.
+_PROFILE_NO_GRAPH_STEPS = 3
+
+
+def _make_server_args(
+    *, canary_on: bool, disable_cuda_graph: bool = False
+) -> ServerArgs:
+    # install_canary asserts --disable-piecewise-cuda-graph; pass on both sides for apples-to-apples.
+    extra = [
+        "--model-path",
+        _QWEN3_MODEL,
+        "--disable-piecewise-cuda-graph",
+    ]
+    if disable_cuda_graph:
+        extra.append("--disable-cuda-graph")
+    if canary_on:
+        extra += ["--kv-canary", "raise"]
+    extra += ["--port", str(DEFAULT_PORT_FOR_SRT_TEST_RUNNER)]
+
+    parser = argparse.ArgumentParser()
+    ServerArgs.add_cli_args(parser)
+    namespace = parser.parse_args(extra)
+    return ServerArgs.from_cli_args(namespace)
+
+
+def _make_bench_args(*, batch_size: int, input_len: int, output_len: int) -> BenchArgs:
+    return BenchArgs(
+        run_name=f"canary_self_bench_bs{batch_size}_isl{input_len}_osl{output_len}",
+        batch_size=(batch_size,),
+        input_len=(input_len,),
+        output_len=(output_len,),
+        temperature=0.0,
+        skip_warmup=False,
+        show_report=True,
+        dataset_name="random",
+        seed=42,
+    )
+
+
+def _run_one_canary_setting(
+    *,
+    canary_on: bool,
+    batch_size: int,
+    input_len: int,
+    output_len: int,
+    disable_cuda_graph: bool = False,
+    profile_output_dir: Optional[Path] = None,
+    profile_steps: int = _PROFILE_STEPS,
+) -> BenchOneCaseResult:
+    server_args = _make_server_args(
+        canary_on=canary_on, disable_cuda_graph=disable_cuda_graph
+    )
+    bench_args = _make_bench_args(
+        batch_size=batch_size, input_len=input_len, output_len=output_len
+    )
+    if profile_output_dir is not None:
+        profile_output_dir.mkdir(parents=True, exist_ok=True)
+        bench_args = dataclasses.replace(
+            bench_args,
+            profile=True,
+            profile_steps=profile_steps,
+            profile_output_dir=str(profile_output_dir),
+        )
+
+    results, _server_info = run_benchmark_internal(
+        server_args=server_args,
+        bench_args=bench_args,
+        launch_server_func=launch_server,
+    )
+    if not results:
+        # run_benchmark_internal returns no rows when the bench was skipped
+        # at the token-capacity guard inside it (the Qwen3-30B-A3B model
+        # leaves only ~12GB for KV cache on an H100; this test's bs128 +
+        # 1024 osl needs more than that). Treat that as a hardware-level
+        # skip rather than a test failure: the canary overhead claim is
+        # still meaningful when the runner has enough memory.
+        raise unittest.SkipTest(
+            f"run_benchmark_internal returned no rows (canary_on={canary_on}, "
+            f"bs={batch_size}, isl={input_len}, osl={output_len}); the runner's "
+            f"KV cache is too small to fit this configuration -- nothing to measure."
+        )
+    return results[0]
+
+
+def _make_scenario_key(*, batch_size: int, input_len: int, output_len: int) -> str:
+    workload = "prefill" if output_len == 1 else "decode"
+    return (
+        f"{_QWEN3_SCENARIO_MODEL}/{workload}_bs{batch_size}"
+        f"_isl{input_len}_osl{output_len}"
+    )
+
+
+def _resolve_profile_root() -> Optional[Path]:
+    raw = os.getenv(_PROFILE_DIR_ENV)
+    return Path(raw).expanduser().resolve() if raw else None
+
+
+class TestCanarySelfBenchSpeed(unittest.TestCase):
+    bench_timeout: ClassVar[float] = 1800.0
+
+    def _capture_profiles(
+        self,
+        *,
+        scenario_key: str,
+        profile_root: Path,
+        batch_size: int,
+        input_len: int,
+        output_len: int,
+    ) -> None:
+        scenario_slug = scenario_key.replace("/", "_")
+        scenario_root = profile_root / f"{scenario_slug}_on"
+
+        graph_dir = scenario_root / "cuda_graph"
+        # +3 to cover prefill chunks + tail; capped so long decode runs still stop after 30 steps.
+        graph_profile_steps = min(_PROFILE_STEPS, output_len + 3)
+        graph_run = _run_one_canary_setting(
+            canary_on=True,
+            batch_size=batch_size,
+            input_len=input_len,
+            output_len=output_len,
+            profile_output_dir=graph_dir,
+            profile_steps=graph_profile_steps,
+        )
+        print(
+            f"[canary self-bench] {scenario_key} profile cuda_graph: "
+            f"on={graph_run.latency:.4f}s (trace under {graph_dir})",
+            flush=True,
+        )
+
+        no_graph_dir = scenario_root / "no_cuda_graph_osl3"
+        no_graph_run = _run_one_canary_setting(
+            canary_on=True,
+            batch_size=batch_size,
+            input_len=input_len,
+            output_len=_PROFILE_NO_GRAPH_OUTPUT_LEN,
+            disable_cuda_graph=True,
+            profile_output_dir=no_graph_dir,
+            profile_steps=_PROFILE_NO_GRAPH_STEPS,
+        )
+        print(
+            f"[canary self-bench] {scenario_key} profile no_cuda_graph_osl3: "
+            f"on={no_graph_run.latency:.4f}s (trace under {no_graph_dir}); "
+            f"off baseline + overhead assertion skipped.",
+            flush=True,
+        )
+
+    def _measure_overhead(
+        self,
+        *,
+        batch_size: int,
+        input_len: int,
+        output_len: int,
+        max_overhead_pct: float,
+    ) -> None:
+        scenario_key = _make_scenario_key(
+            batch_size=batch_size, input_len=input_len, output_len=output_len
+        )
+        profile_root = _resolve_profile_root()
+
+        if profile_root is not None:
+            self._capture_profiles(
+                scenario_key=scenario_key,
+                profile_root=profile_root,
+                batch_size=batch_size,
+                input_len=input_len,
+                output_len=output_len,
+            )
+            return
+
+        off = _run_one_canary_setting(
+            canary_on=False,
+            batch_size=batch_size,
+            input_len=input_len,
+            output_len=output_len,
+        )
+        on = _run_one_canary_setting(
+            canary_on=True,
+            batch_size=batch_size,
+            input_len=input_len,
+            output_len=output_len,
+        )
+        overhead_pct = ((on.latency - off.latency) / off.latency) * 100.0
+        summary = (
+            f"[canary self-bench] {scenario_key}: "
+            f"off={off.latency:.4f}s on={on.latency:.4f}s overhead={overhead_pct:.2f}%"
+        )
+        print(summary, flush=True)
+        self.assertLess(
+            overhead_pct,
+            max_overhead_pct,
+            msg=(f"{summary} — exceeds {max_overhead_pct:.1f}% budget"),
+        )
+
+    def test_qwen3_prefill_overhead_bs32_isl16384_osl1(self) -> None:
+        # TODO: tighten further once the per-forward elementwise glue + plan_offsets
+        # single-program kernel are optimized (observed ~2.17% on Qwen3-30B-A3B, H200).
+        self._measure_overhead(
+            batch_size=32,
+            input_len=16384,
+            output_len=1,
+            max_overhead_pct=3.0,
+        )
+
+    def test_qwen3_decode_overhead_bs64_isl256_osl512(self) -> None:
+        # TODO: tighten further once per-forward canary glue is reduced (observed ~0.52% on
+        # Qwen3-30B-A3B, H200 — already amortizes well at large bs). The smaller
+        # 64 * (256+512) = 49K-token budget fits the ~94K KV-cache slice that
+        # extra-a-test-1-gpu-large (H100) leaves after loading the 30B MoE.
+        self._measure_overhead(
+            batch_size=64,
+            input_len=256,
+            output_len=512,
+            max_overhead_pct=1.0,
+        )
+
+    def test_qwen3_decode_overhead_bs1_isl512_osl1024(self) -> None:
+        # TODO: tighten further once the per-forward elementwise glue + plan_offsets
+        # single-program kernel are optimized (observed ~2.10% on Qwen3-30B-A3B, H200).
+        self._measure_overhead(
+            batch_size=1,
+            input_len=512,
+            output_len=1024,
+            max_overhead_pct=3.0,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_runner_per_forward.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag, VerifyPlan
+from sglang.jit_kernel.kv_canary.write import WritePlan
+from sglang.srt.kv_canary import endpoint as endpoint_module
+from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
+from sglang.srt.kv_canary.runner import kernel_launcher as kernel_launcher_module
+from sglang.srt.kv_canary.state import ViolationLog
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kv_canary.fixtures import make_buffer_group, make_forward_batch
+from sglang.test.kv_canary.runner_test_base import (
+    CanaryManagerTestCase,
+    RecordingEndpoint,
+    make_manager,
+)
+
+register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small")
+
+
+class TestManagerPerForward(CanaryManagerTestCase):
+    def test_per_forward_orchestrates_plan_head_tail(self) -> None:
+        """Verify per-forward execution launches plan, head/tail verify kernels, and write kernels in order."""
+        calls: list[object] = []
+        with patch.object(
+            kernel_launcher_module,
+            "launch_canary_plan_kernels",
+            lambda **kwargs: calls.append("plan"),
+        ), patch.object(
+            endpoint_module,
+            "launch_canary_verify_kernel",
+            lambda **kwargs: calls.append(
+                ("verify", kwargs["context"].kernel_kind.name)
+            ),
+        ), patch.object(
+            endpoint_module,
+            "launch_canary_write_kernel",
+            lambda **kwargs: calls.append(
+                ("write", kwargs["context"].kernel_kind.name)
+            ),
+        ):
+            manager = make_manager(device=self.device)
+            forward_batch = make_forward_batch(self.device)
+            with manager.with_ops_outside_graph(
+                single_forward_indices=[0],
+                maybe_inaccurate_forward_batch=forward_batch,
+            ):
+                with manager.with_active_single_forward_manager(0):
+                    pre_ops_output = manager.pre_ops_maybe_inside_graph(forward_batch)
+                    manager.post_ops_maybe_inside_graph(forward_batch, pre_ops_output)
+
+        self.assertEqual(calls[0], "plan")
+        self.assertTrue(
+            any(
+                call[0] == "verify" and "HEAD" in call[1]
+                for call in calls[1:]
+                if isinstance(call, tuple)
+            )
+        )
+        self.assertTrue(
+            any(
+                call[0] == "verify" and "TAIL" in call[1]
+                for call in calls[1:]
+                if isinstance(call, tuple)
+            )
+        )
+
+
+class TestLaunchEndpointsPerForward(CanaryManagerTestCase):
+    def test_launch_endpoints_per_forward_keeps_padded_token_tensors(self) -> None:
+        """Verify endpoint launch preserves contiguous int64 tensor shapes/values through the canonicalizer."""
+        group = make_buffer_group(device=self.device)
+        endpoint = RecordingEndpoint(kernel_kind=CanaryLaunchTag.HEAD_K_FULL)
+        forward_batch = make_forward_batch(self.device, bs=1, seq_lens_list=(1,))
+        forward_batch.input_ids = torch.tensor(
+            [101, 0, 0], dtype=torch.int64, device=self.device
+        )
+        forward_batch.positions = torch.tensor(
+            [10, 0, 0], dtype=torch.int64, device=self.device
+        )
+        forward_batch.out_cache_loc = torch.tensor(
+            [7, 0, 0], dtype=torch.int64, device=self.device
+        )
+        forward_batch.num_token_non_padded_cpu = 1
+
+        kernel_launcher_module.launch_endpoints_per_forward(
+            endpoints=(endpoint,),
+            group=group,
+            tag_filter=lambda tag: True,
+            verify_plan=VerifyPlan.allocate(verify_capacity=1, device=self.device),
+            write_plan=WritePlan.allocate(write_req_capacity=1, device=self.device),
+            forward_batch=forward_batch,
+            expected_inputs=ExpectedInputs.allocate(capacity=3, device=self.device),
+            violation_log=ViolationLog.allocate(ring_capacity=2, device=self.device),
+            enable_write_input_assert=False,
+            enable_verify_token_assert=False,
+        )
+
+        self.assertEqual(len(endpoint.calls), 1)
+        call = endpoint.calls[0]
+        self.assertTrue(
+            torch.equal(
+                call["input_ids"],
+                torch.tensor([101, 0, 0], dtype=torch.int64, device=self.device),
+            )
+        )
+        self.assertTrue(
+            torch.equal(
+                call["positions"],
+                torch.tensor([10, 0, 0], dtype=torch.int64, device=self.device),
+            )
+        )
+        self.assertTrue(
+            torch.equal(
+                call["out_cache_loc"],
+                torch.tensor([7, 0, 0], dtype=torch.int64, device=self.device),
+            )
+        )
+
+    def test_launch_endpoints_per_forward_promotes_int32_boundary_tensors_to_int64(
+        self,
+    ) -> None:
+        """Verify int32 boundary tensors are promoted to int64 at the launch boundary."""
+        group = make_buffer_group(device=self.device)
+        endpoint = RecordingEndpoint(kernel_kind=CanaryLaunchTag.HEAD_K_FULL)
+        forward_batch = make_forward_batch(self.device, bs=1, seq_lens_list=(1,))
+        forward_batch.input_ids = torch.tensor(
+            [101], dtype=torch.int32, device=self.device
+        )
+        forward_batch.positions = torch.tensor(
+            [10], dtype=torch.int32, device=self.device
+        )
+        forward_batch.out_cache_loc = torch.tensor(
+            [7], dtype=torch.int32, device=self.device
+        )
+        forward_batch.num_token_non_padded_cpu = 1
+
+        kernel_launcher_module.launch_endpoints_per_forward(
+            endpoints=(endpoint,),
+            group=group,
+            tag_filter=lambda tag: True,
+            verify_plan=VerifyPlan.allocate(verify_capacity=1, device=self.device),
+            write_plan=WritePlan.allocate(write_req_capacity=1, device=self.device),
+            forward_batch=forward_batch,
+            expected_inputs=ExpectedInputs.allocate(capacity=1, device=self.device),
+            violation_log=ViolationLog.allocate(ring_capacity=2, device=self.device),
+            enable_write_input_assert=False,
+            enable_verify_token_assert=False,
+        )
+
+        self.assertEqual(len(endpoint.calls), 1)
+        call = endpoint.calls[0]
+        self.assertEqual(call["input_ids"].dtype, torch.int64)
+        self.assertEqual(call["positions"].dtype, torch.int64)
+        self.assertEqual(call["out_cache_loc"].dtype, torch.int64)
+
+    def test_launch_endpoints_per_forward_propagates_enable_verify_token_assert_true(
+        self,
+    ) -> None:
+        """Verify enable_verify_token_assert=True is plumbed through to the endpoint kwargs."""
+        group = make_buffer_group(device=self.device)
+        endpoint = RecordingEndpoint(kernel_kind=CanaryLaunchTag.HEAD_K_FULL)
+        forward_batch = make_forward_batch(self.device, bs=1, seq_lens_list=(1,))
+        forward_batch.input_ids = torch.tensor(
+            [101, 0, 0], dtype=torch.int64, device=self.device
+        )
+        forward_batch.positions = torch.tensor(
+            [10, 0, 0], dtype=torch.int64, device=self.device
+        )
+        forward_batch.out_cache_loc = torch.tensor(
+            [7, 0, 0], dtype=torch.int64, device=self.device
+        )
+        forward_batch.num_token_non_padded_cpu = 1
+
+        kernel_launcher_module.launch_endpoints_per_forward(
+            endpoints=(endpoint,),
+            group=group,
+            tag_filter=lambda tag: True,
+            verify_plan=VerifyPlan.allocate(verify_capacity=1, device=self.device),
+            write_plan=WritePlan.allocate(write_req_capacity=1, device=self.device),
+            forward_batch=forward_batch,
+            expected_inputs=ExpectedInputs.allocate(capacity=3, device=self.device),
+            violation_log=ViolationLog.allocate(ring_capacity=2, device=self.device),
+            enable_write_input_assert=False,
+            enable_verify_token_assert=True,
+        )
+
+        self.assertEqual(len(endpoint.calls), 1)
+        call = endpoint.calls[0]
+        self.assertEqual(call["enable_verify_token_assert"], True)
+
+    def test_launch_endpoints_per_forward_materializes_strided_boundary_tensors(
+        self,
+    ) -> None:
+        """Verify non-contiguous boundary views are materialized contiguous at launch."""
+        group = make_buffer_group(device=self.device)
+        endpoint = RecordingEndpoint(kernel_kind=CanaryLaunchTag.HEAD_K_FULL)
+        forward_batch = make_forward_batch(self.device, bs=1, seq_lens_list=(1,))
+        forward_batch.input_ids = torch.tensor(
+            [[101, 102]], dtype=torch.int64, device=self.device
+        )[:, 0]
+        forward_batch.positions = torch.tensor(
+            [[10, 11]], dtype=torch.int64, device=self.device
+        )[:, 0]
+        forward_batch.out_cache_loc = torch.tensor(
+            [[7, 8]], dtype=torch.int64, device=self.device
+        )[:, 0]
+        forward_batch.num_token_non_padded_cpu = 1
+
+        kernel_launcher_module.launch_endpoints_per_forward(
+            endpoints=(endpoint,),
+            group=group,
+            tag_filter=lambda tag: True,
+            verify_plan=VerifyPlan.allocate(verify_capacity=1, device=self.device),
+            write_plan=WritePlan.allocate(write_req_capacity=1, device=self.device),
+            forward_batch=forward_batch,
+            expected_inputs=ExpectedInputs.allocate(capacity=1, device=self.device),
+            violation_log=ViolationLog.allocate(ring_capacity=2, device=self.device),
+            enable_write_input_assert=False,
+            enable_verify_token_assert=False,
+        )
+
+        self.assertEqual(len(endpoint.calls), 1)
+        call = endpoint.calls[0]
+        self.assertTrue(call["input_ids"].is_contiguous())
+        self.assertTrue(call["positions"].is_contiguous())
+        self.assertTrue(call["out_cache_loc"].is_contiguous())
+
+
+class TestManagerBeforeForward(CanaryManagerTestCase):
+    def test_before_forward_does_not_throw_on_oversized_prefix_sum(self) -> None:
+        """Verify oversized prefix sums are handled without host-side errors."""
+        manager = make_manager(device=self.device, per_forward_verify_capacity=4)
+        forward_batch = make_forward_batch(self.device, bs=2, seq_lens_list=(5, 5))
+        _drive_one_cycle(manager, forward_batch)
+
+    def test_before_forward_passes_when_sum_prefix_lens_fits(self) -> None:
+        """Verify prefix sums within capacity pass before-forward handling."""
+        manager = make_manager(device=self.device, per_forward_verify_capacity=16)
+        forward_batch = make_forward_batch(self.device, bs=2, seq_lens_list=(5, 5))
+        _drive_one_cycle(manager, forward_batch)
+
+
+def _drive_one_cycle(manager, forward_batch) -> None:
+    with manager.with_ops_outside_graph(
+        single_forward_indices=[0],
+        maybe_inaccurate_forward_batch=forward_batch,
+    ):
+        with manager.with_active_single_forward_manager(0):
+            pre_ops_output = manager.pre_ops_maybe_inside_graph(forward_batch)
+            manager.post_ops_maybe_inside_graph(forward_batch, pre_ops_output)
+
+
+if __name__ == "__main__":
+    unittest.main()