Add the KV-canary install API and forward-path wiring (#26809)
This commit is contained in:
fzyzcjy
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user