From 9ecf3149706869a72daf47069c66a5e97eb7539f Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 31 May 2026 09:55:03 +0800 Subject: [PATCH] Add the KV-canary install API and forward-path wiring (#26809) --- python/sglang/srt/kv_canary/api.py | 101 +++++++ .../srt/model_executor/cuda_graph_runner.py | 6 +- .../sglang/srt/model_executor/model_runner.py | 25 +- .../attention_methods/dense_attention.py | 1 + .../attention_methods/dsa_attention.py | 1 + .../attention_methods/dsv4_attention.py | 1 + .../attention_methods/dual_chunk_attention.py | 1 + .../attention_methods/gdn_attention.py | 1 + .../attention_methods/kda_attention.py | 1 + .../attention_methods/lightning_attention.py | 1 + .../attention_methods/mamba2_attention.py | 1 + .../attention_methods/mla_attention.py | 1 + .../kv_canary/test_self_e2e_baseline.py | 37 +++ .../kv_canary/test_self_e2e_bench_speed.py | 256 +++++++++++++++++ .../test_self_unit_runner_per_forward.py | 259 ++++++++++++++++++ 15 files changed, 691 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/kv_canary/api.py create mode 100644 test/registered/kv_canary/test_self_e2e_baseline.py create mode 100644 test/registered/kv_canary/test_self_e2e_bench_speed.py create mode 100644 test/registered/kv_canary/test_self_unit_runner_per_forward.py diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py new file mode 100644 index 000000000..68a739708 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 5f2006edb..4b411c9fa 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/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() diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 11084782b..dc460cdf4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index a9bd11cfa..b26a3f4cd 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index 8cc635ba8..68c7ad9d6 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index e0c54bdd5..75cc359cb 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py index e306524bd..257e14e99 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index 8c9c35356..c404d303b 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/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 = ( diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py index c3f77abdf..5491cf153 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py +++ b/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 = ( diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py index 0b040c0ad..b3cace2b9 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py +++ b/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 = ( diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index ff002fdf4..f86e8f955 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/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` diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index e3b43eccd..908574e28 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/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 diff --git a/test/registered/kv_canary/test_self_e2e_baseline.py b/test/registered/kv_canary/test_self_e2e_baseline.py new file mode 100644 index 000000000..562025302 --- /dev/null +++ b/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() diff --git a/test/registered/kv_canary/test_self_e2e_bench_speed.py b/test/registered/kv_canary/test_self_e2e_bench_speed.py new file mode 100644 index 000000000..9ebfa13f7 --- /dev/null +++ b/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() diff --git a/test/registered/kv_canary/test_self_unit_runner_per_forward.py b/test/registered/kv_canary/test_self_unit_runner_per_forward.py new file mode 100644 index 000000000..f4033d6af --- /dev/null +++ b/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()