diff --git a/README.md b/README.md index 57d7fc2..49340e3 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,18 @@ Once installed, FlashKDA is auto-dispatched from `flash-linear-attention`'s `chu See [BENCHMARK_H20.md](BENCHMARK_H20.md). +To isolate the allocation and latency effects of caller-owned workspace reuse, +run the workspace benchmark with a preallocated output in both modes: + +```bash +python benchmarks/bench_workspace.py --json-out workspace_benchmark.json +``` + +The report includes host enqueue and batched end-to-end p50/p95 latency, CUDA +stream elapsed time, profiler-visible empty operations, incremental peak memory, +output equality, and CUDA Graph replay correctness across fixed, batched, and +variable-length inputs. + ## Tests ```bash @@ -81,10 +93,31 @@ bash tests/test.sh ### `flash_kda.fwd` ```python -flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, - initial_state=None, final_state=None, cu_seqlens=None) +out = flash_kda.fwd( + q, k, v, g, beta, scale, + A_log=A_log, dt_bias=dt_bias, lower_bound=lower_bound, + initial_state=None, final_state=None, cu_seqlens=None, +) ``` +``fwd`` allocates ``out`` and its temporary workspace when they are omitted. +Latency-sensitive callers can allocate the workspace once and reuse it for +sequential calls: + +```python +workspace = flash_kda.allocate_workspace(q, cu_seqlens) +out = torch.empty_like(q) +out = flash_kda.fwd( + q, k, v, g, beta, scale, + A_log=A_log, dt_bias=dt_bias, lower_bound=lower_bound, + out=out, cu_seqlens=cu_seqlens, workspace=workspace, +) +``` + +A workspace sized for a larger input can serve a smaller input on the same +device. Do not share one workspace between overlapping calls on different CUDA +streams; allocate one workspace per concurrent call instead. + **Parameters:** | Parameter | Dtype | Shape | Description | @@ -95,13 +128,14 @@ flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, | `g` | bf16 | `[B, T, H, K]` | Gate before activation | | `beta` | bf16 | `[B, T, H]` | Beta logits (pre-activation; sigmoid applied internally) | | `scale` | float | scalar | scaling factor | -| `out` | bf16 | `[B, T, H, V]` | Output tensor | +| `out` | bf16/None | `[B, T, H, V]` | Optional output tensor; allocated like `q` when omitted | | `A_log` | fp32 | `[H]` | Log-gate parameter | | `dt_bias` | fp32 | `[H, K]` | Gate bias | | `lower_bound` | float | scalar | Gate lower bound (range from -5.0 to 0) | | `initial_state` | bf16/fp32/None | `[B, H, V, K]` or `[N, H, V, K]` | (optional) Initial recurrent state | | `final_state` | bf16/fp32/None | `[B, H, V, K]` or `[N, H, V, K]` | (optional, output) Final recurrent state | | `cu_seqlens` | int64 | `[N+1]` | (optional) Cumulative sequence lengths for variable-length batching | +| `workspace` | uint8/None | `[bytes]` | (optional) Reusable temporary storage from `allocate_workspace` | - Currently requires `K = V = 128`. - `initial_state` / `final_state` accept `None` (stateless), bf16, or fp32 tensors. When both are provided, their dtypes must match. diff --git a/benchmarks/bench_workspace.py b/benchmarks/bench_workspace.py new file mode 100644 index 0000000..8bd2168 --- /dev/null +++ b/benchmarks/bench_workspace.py @@ -0,0 +1,340 @@ +"""Benchmark automatic allocation against caller-owned workspace reuse. + +The output buffer is preallocated in both modes so that the only Python API +difference under test is whether ``flash_kda.fwd`` allocates its workspace. +""" + +import argparse +import gc +import json +import math +import statistics +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch.profiler import ProfilerActivity, profile + +import flash_kda + + +@dataclass(frozen=True) +class Case: + name: str + batch: int + tokens: int + heads: int + seq_lens: Optional[Tuple[int, ...]] = None + + +CASES = ( + Case("fixed-t256-h1", batch=1, tokens=256, heads=1), + Case("fixed-t2048-h8", batch=1, tokens=2048, heads=8), + Case("batched-b4-t512-h8", batch=4, tokens=512, heads=8), + Case( + "varlen-17-33-65-257-h4", + batch=1, + tokens=372, + heads=4, + seq_lens=(17, 33, 65, 257), + ), + Case("fixed-t8192-h32", batch=1, tokens=8192, heads=32), +) + +D = 128 +LOWER_BOUND = -5.0 + + +def percentile(values, percent): + values = sorted(values) + rank = math.ceil(percent / 100 * len(values)) - 1 + return values[max(0, rank)] + + +def make_inputs(case): + torch.manual_seed(123) + shape = (case.batch, case.tokens, case.heads, D) + q = F.normalize(torch.randn(shape, device="cuda"), p=2, dim=-1).bfloat16() + k = F.normalize(torch.randn(shape, device="cuda"), p=2, dim=-1).bfloat16() + v = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + g = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + beta = torch.randn( + (case.batch, case.tokens, case.heads), + device="cuda", + dtype=torch.bfloat16, + ) + A_log = torch.rand(case.heads, device="cuda", dtype=torch.float32) + dt_bias = torch.rand(case.heads, D, device="cuda", dtype=torch.float32) + + cu_seqlens = None + if case.seq_lens is not None: + if case.batch != 1 or sum(case.seq_lens) != case.tokens: + raise ValueError(f"invalid varlen case: {case}") + offsets = [0] + for length in case.seq_lens: + offsets.append(offsets[-1] + length) + cu_seqlens = torch.tensor(offsets, device="cuda", dtype=torch.long) + + kwargs = { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "scale": 1.0 / math.sqrt(D), + "out": torch.empty_like(q), + "A_log": A_log, + "dt_bias": dt_bias, + "lower_bound": LOWER_BOUND, + } + if cu_seqlens is not None: + kwargs["cu_seqlens"] = cu_seqlens + return kwargs, cu_seqlens + + +def warmup(fn, calls): + for _ in range(calls): + fn() + torch.cuda.synchronize() + + +def measure_host_enqueue(fn, iterations): + samples = [] + for _ in range(iterations): + start = time.perf_counter_ns() + fn() + samples.append((time.perf_counter_ns() - start) / 1_000) + torch.cuda.synchronize() + return { + "p50_us": statistics.median(samples), + "p95_us": percentile(samples, 95), + } + + +def measure_batched_e2e(fn, repeats, calls_per_repeat): + samples = [] + for _ in range(repeats): + torch.cuda.synchronize() + start = time.perf_counter_ns() + for _ in range(calls_per_repeat): + fn() + torch.cuda.synchronize() + samples.append( + (time.perf_counter_ns() - start) / 1_000 / calls_per_repeat + ) + return { + "p50_us_per_call": statistics.median(samples), + "p95_us_per_call": percentile(samples, 95), + } + + +def measure_stream_elapsed(fn, iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start.record() + for _ in range(iterations): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1_000 / iterations + + +def count_empty_ops(fn, calls): + with profile(activities=[ProfilerActivity.CPU]) as prof: + for _ in range(calls): + fn() + torch.cuda.synchronize() + return sum( + event.count + for event in prof.key_averages() + if event.key.startswith("aten::empty") + ) + + +def measure_incremental_peak(fn, calls): + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + allocated_before = torch.cuda.memory_allocated() + reserved_before = torch.cuda.memory_reserved() + for _ in range(calls): + fn() + torch.cuda.synchronize() + return { + "allocated_bytes": torch.cuda.max_memory_allocated() - allocated_before, + "reserved_bytes": torch.cuda.max_memory_reserved() - reserved_before, + } + + +def check_cuda_graph(fn): + expected = fn().clone() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = fn() + graph.replay() + torch.cuda.synchronize() + exact = torch.equal(expected, captured) + del graph, captured, expected + return exact + + +def measure_mode(fn, args): + warmup(fn, args.warmup) + return { + "host_enqueue": measure_host_enqueue(fn, args.host_iterations), + "batched_e2e": measure_batched_e2e( + fn, args.e2e_repeats, args.calls_per_repeat + ), + "stream_elapsed_mean_us": measure_stream_elapsed( + fn, args.cuda_iterations + ), + "aten_empty_calls": count_empty_ops(fn, args.profile_calls), + "profiled_calls": args.profile_calls, + "incremental_peak": measure_incremental_peak(fn, args.memory_calls), + } + + +def run_case(case, args): + kwargs, cu_seqlens = make_inputs(case) + workspace = flash_kda.allocate_workspace(kwargs["q"], cu_seqlens) + + def automatic(): + return flash_kda.fwd(**kwargs) + + def reused(): + return flash_kda.fwd(**kwargs, workspace=workspace) + + automatic_out = automatic().clone() + reused_out = reused().clone() + torch.cuda.synchronize() + exact = torch.equal(automatic_out, reused_out) + del automatic_out, reused_out + + modes = { + "automatic": measure_mode(automatic, args), + "reused": measure_mode(reused, args), + } + graph_exact = check_cuda_graph(reused) + + auto = modes["automatic"] + reuse = modes["reused"] + ratios = { + "host_p50_auto_over_reuse": ( + auto["host_enqueue"]["p50_us"] + / reuse["host_enqueue"]["p50_us"] + ), + "e2e_p50_auto_over_reuse": ( + auto["batched_e2e"]["p50_us_per_call"] + / reuse["batched_e2e"]["p50_us_per_call"] + ), + "stream_auto_over_reuse": ( + auto["stream_elapsed_mean_us"] + / reuse["stream_elapsed_mean_us"] + ), + } + return { + "case": asdict(case), + "workspace_bytes": workspace.numel(), + "automatic_equals_reused": exact, + "reused_cuda_graph_exact": graph_exact, + "modes": modes, + "ratios": ratios, + } + + +def format_bytes(value): + return f"{value / (1024 * 1024):.2f} MiB" + + +def print_markdown(results): + print() + print( + "| Case | Mode | Workspace | Host p50/p95 (us) | " + "E2E p50/p95 (us) | Stream mean (us) | empty/call | Peak alloc |" + ) + print("| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |") + for result in results: + for mode_name, mode in result["modes"].items(): + host = mode["host_enqueue"] + e2e = mode["batched_e2e"] + empty_per_call = mode["aten_empty_calls"] / mode["profiled_calls"] + print( + f"| {result['case']['name']} | {mode_name} | " + f"{format_bytes(result['workspace_bytes'])} | " + f"{host['p50_us']:.2f}/{host['p95_us']:.2f} | " + f"{e2e['p50_us_per_call']:.2f}/{e2e['p95_us_per_call']:.2f} | " + f"{mode['stream_elapsed_mean_us']:.2f} | " + f"{empty_per_call:.2f} | " + f"{format_bytes(mode['incremental_peak']['allocated_bytes'])} |" + ) + print() + for result in results: + ratios = result["ratios"] + print( + f"{result['case']['name']}: exact=" + f"{result['automatic_equals_reused']}, graph_exact=" + f"{result['reused_cuda_graph_exact']}, auto/reuse=" + f"host {ratios['host_p50_auto_over_reuse']:.3f}x, " + f"e2e {ratios['e2e_p50_auto_over_reuse']:.3f}x, " + f"stream {ratios['stream_auto_over_reuse']:.3f}x" + ) + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--cases", + nargs="*", + choices=[case.name for case in CASES], + help="case names to run (default: all)", + ) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--host-iterations", type=int, default=200) + parser.add_argument("--e2e-repeats", type=int, default=30) + parser.add_argument("--calls-per-repeat", type=int, default=10) + parser.add_argument("--cuda-iterations", type=int, default=200) + parser.add_argument("--profile-calls", type=int, default=25) + parser.add_argument("--memory-calls", type=int, default=5) + parser.add_argument("--json-out", type=Path) + return parser.parse_args() + + +def main(): + args = parse_args() + selected = [case for case in CASES if not args.cases or case.name in args.cases] + metadata = { + "gpu": torch.cuda.get_device_name(), + "compute_capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "settings": { + key: value + for key, value in vars(args).items() + if key not in {"cases", "json_out"} + }, + } + print(json.dumps(metadata, indent=2)) + + results = [] + with torch.inference_mode(): + for case in selected: + print(f"benchmarking {case.name}...", flush=True) + results.append(run_case(case, args)) + gc.collect() + torch.cuda.empty_cache() + + print_markdown(results) + payload = {"metadata": metadata, "results": results} + if args.json_out is not None: + args.json_out.write_text(json.dumps(payload, indent=2) + "\n") + print(f"wrote {args.json_out}") + + +if __name__ == "__main__": + main() diff --git a/csrc/flash_kda.cpp b/csrc/flash_kda.cpp index 81f5483..85e4cbb 100644 --- a/csrc/flash_kda.cpp +++ b/csrc/flash_kda.cpp @@ -41,10 +41,14 @@ void fwd( std::optional final_state = std::nullopt, std::optional cu_seqlens = std::nullopt ) { - TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && g.is_cuda() && beta.is_cuda() && out.is_cuda() && workspace.is_cuda(), - "all tensors must be on CUDA"); - TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous() && g.is_contiguous() && beta.is_contiguous() && out.is_contiguous() && workspace.is_contiguous(), - "all tensors must be contiguous"); + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && g.is_cuda() && beta.is_cuda() && out.is_cuda(), + "all input and output tensors must be on CUDA"); + TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous() && g.is_contiguous() && beta.is_contiguous() && out.is_contiguous(), + "all input and output tensors must be contiguous"); + TORCH_CHECK(workspace.is_cuda(), "workspace must be a CUDA tensor"); + TORCH_CHECK(workspace.is_contiguous(), "workspace must be contiguous"); + TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must have dtype uint8"); + TORCH_CHECK(workspace.device() == q.device(), "workspace must be on the same device as q"); TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bfloat16"); TORCH_CHECK(k.dtype() == torch::kBFloat16, "k must be bfloat16"); @@ -159,6 +163,12 @@ void fwd( N_val = B; } + int64_t required_workspace_bytes = get_workspace_size(T_total, H, N_val); + TORCH_CHECK( + workspace.numel() >= required_workspace_bytes, + "workspace is too small: expected at least ", required_workspace_bytes, + " bytes, but got ", workspace.numel()); + // Validate state shapes: always [N, H, D, D] if (has_state_in) { auto& is = initial_state.value(); diff --git a/flash_kda/__init__.py b/flash_kda/__init__.py index cc03493..1186d79 100644 --- a/flash_kda/__init__.py +++ b/flash_kda/__init__.py @@ -1,8 +1,42 @@ import torch from flash_kda_C import fwd as _fwd_raw, get_workspace_size +__all__ = ["allocate_workspace", "fwd", "get_workspace_size"] -def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None): + +def _workspace_size_from_inputs(q, cu_seqlens=None): + """Return the workspace size required by a particular forward call.""" + B, T_seq, H = q.shape[:3] + N = cu_seqlens.numel() - 1 if cu_seqlens is not None else B + return get_workspace_size(B * T_seq, H, N) + + +def allocate_workspace(q, cu_seqlens=None): + """Allocate a reusable workspace for :func:`fwd`. + + The returned byte tensor is large enough for ``q`` and ``cu_seqlens``. It + can be passed to multiple sequential calls, including calls with smaller + shapes. A workspace must not be used by overlapping calls on different + CUDA streams because the kernels write to it. + + Args: + q (torch.Tensor): Query tensor whose shape and device determine the + workspace capacity and placement. + cu_seqlens (torch.Tensor, optional): Cumulative sequence lengths for a + variable-length call. + """ + if q.device.type != "cuda": + raise ValueError("q must be on CUDA when allocating a workspace") + return torch.empty( + _workspace_size_from_inputs(q, cu_seqlens), + dtype=torch.uint8, + device=q.device, + ) + + +def fwd(q, k, v, g, beta, scale, out=None, A_log=None, dt_bias=None, + lower_bound=None, initial_state=None, final_state=None, + cu_seqlens=None, workspace=None): """FlashKDA forward (Flash Kimi Delta Attention). Args: @@ -13,8 +47,9 @@ def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state beta (torch.Tensor): Beta logits (pre-activation; sigmoid is applied internally), bf16, shape ``[B, T, H]``. scale (float): Scaling factor. - out (torch.Tensor): Output buffer, bf16, shape ``[B, T, H, V]``. Written - in place. + out (torch.Tensor, optional): Output buffer, bf16, shape + ``[B, T, H, V]``. When omitted, a tensor is allocated with the same + shape, dtype, and device as ``q``. A_log (torch.Tensor): Log-gate parameter, fp32, shape ``[H]``. dt_bias (torch.Tensor): Gate bias, fp32, shape ``[H, K]``. lower_bound (float): Gate lower bound, expected in ``[-5.0, 0]``. @@ -25,17 +60,36 @@ def fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state recurrent state. Same dtype/shape rules as ``initial_state``. cu_seqlens (torch.Tensor, optional): Cumulative sequence lengths, int64, shape ``[N+1]``. When provided, ``B`` must be 1. + workspace (torch.Tensor, optional): Contiguous CUDA byte tensor used as + temporary storage. When omitted, an exactly-sized tensor is + allocated for this call. Use :func:`allocate_workspace` and pass + the same tensor to sequential calls to avoid repeated allocation. + + Returns: + torch.Tensor: ``out`` after it has been written by the kernel. Notes: * Currently requires ``K = V = 128``. * All input tensors must be CUDA, contiguous, and have the dtypes listed above. """ - B, T_seq, H = q.shape[0], q.shape[1], q.shape[2] - T_total = B * T_seq - N = cu_seqlens.numel() - 1 if cu_seqlens is not None else B + missing = [ + name for name, value in ( + ("A_log", A_log), + ("dt_bias", dt_bias), + ("lower_bound", lower_bound), + ) + if value is None + ] + if missing: + names = ", ".join(repr(name) for name in missing) + raise TypeError(f"fwd() missing required argument(s): {names}") - workspace = torch.empty(get_workspace_size(T_total, H, N), dtype=torch.uint8, device=q.device) + if out is None: + out = torch.empty_like(q) + if workspace is None: + workspace = allocate_workspace(q, cu_seqlens) _fwd_raw(q, k, v, g, beta, float(scale), out, workspace, A_log, dt_bias, lower_bound, initial_state=initial_state, final_state=final_state, cu_seqlens=cu_seqlens) + return out diff --git a/tests/test.sh b/tests/test.sh index 68ecba4..a621a07 100644 --- a/tests/test.sh +++ b/tests/test.sh @@ -1,4 +1,5 @@ set -e pip install -e . -pip install "flash-linear-attention>=0.5.0" matplotlib +pip install "flash-linear-attention>=0.5.0" matplotlib pytest +pytest -q tests/test_python_api.py python tests/test_fwd.py diff --git a/tests/test_python_api.py b/tests/test_python_api.py new file mode 100644 index 0000000..18f695d --- /dev/null +++ b/tests/test_python_api.py @@ -0,0 +1,136 @@ +import math + +import pytest +import torch +import torch.nn.functional as F + +import flash_kda + + +def make_inputs(T=33, H=1): + B, D = 1, 128 + torch.manual_seed(0) + q = F.normalize( + torch.randn((B, T, H, D), dtype=torch.float32, device="cuda"), + p=2, + dim=-1, + ).to(torch.bfloat16) + return { + "q": q, + "k": q.clone(), + "v": torch.randn_like(q), + "g": torch.randn_like(q), + "beta": torch.randn((B, T, H), dtype=torch.bfloat16, device="cuda"), + "scale": 1.0 / math.sqrt(D), + "A_log": torch.rand(H, dtype=torch.float32, device="cuda"), + "dt_bias": torch.rand(H, D, dtype=torch.float32, device="cuda"), + "lower_bound": -5.0, + } + + +def test_pythonic_api_allocates_and_returns_out(): + args = make_inputs() + + out = flash_kda.fwd(**args) + expected = torch.empty_like(args["q"]) + flash_kda.fwd(**args, out=expected) + + assert out.shape == args["q"].shape + assert out.dtype == args["q"].dtype + assert out.device == args["q"].device + assert torch.equal(out, expected) + + +def test_provided_out_is_returned(): + args = make_inputs() + provided_out = torch.empty_like(args["q"]) + + returned_out = flash_kda.fwd(**args, out=provided_out) + + assert returned_out is provided_out + + +def test_workspace_can_be_reused_for_sequential_calls(): + args = make_inputs(T=33) + workspace = flash_kda.allocate_workspace(args["q"]) + data_ptr = workspace.data_ptr() + + first = flash_kda.fwd(**args, workspace=workspace).clone() + second = flash_kda.fwd(**args, workspace=workspace).clone() + smaller_args = make_inputs(T=17) + smaller_with_reuse = flash_kda.fwd(**smaller_args, workspace=workspace).clone() + smaller_automatic = flash_kda.fwd(**smaller_args).clone() + + assert workspace.data_ptr() == data_ptr + assert torch.equal(first, second) + assert torch.equal(smaller_with_reuse, smaller_automatic) + + +def test_varlen_workspace_can_be_reused(): + args = make_inputs(T=50) + cu_seqlens = torch.tensor([0, 17, 50], dtype=torch.long, device="cuda") + workspace = flash_kda.allocate_workspace(args["q"], cu_seqlens) + + with_reuse = flash_kda.fwd( + **args, cu_seqlens=cu_seqlens, workspace=workspace + ).clone() + automatic = flash_kda.fwd(**args, cu_seqlens=cu_seqlens).clone() + + assert torch.equal(with_reuse, automatic) + + +def test_reusable_buffers_support_cuda_graph_replay(): + args = make_inputs(T=33) + out = torch.empty_like(args["q"]) + workspace = flash_kda.allocate_workspace(args["q"]) + + expected = flash_kda.fwd(**args, out=out, workspace=workspace).clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = flash_kda.fwd(**args, out=out, workspace=workspace) + + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, expected) + + +@pytest.mark.parametrize( + "workspace_factory, error", + [ + ( + lambda required: torch.empty( + required, dtype=torch.float32, device="cuda" + ), + "dtype uint8", + ), + ( + lambda required: torch.empty( + required - 1, dtype=torch.uint8, device="cuda" + ), + "too small", + ), + ( + lambda required: torch.empty( + required, dtype=torch.uint8, device="cpu" + ), + "CUDA tensor", + ), + ( + lambda required: torch.empty( + required * 2, dtype=torch.uint8, device="cuda" + )[::2], + "contiguous", + ), + ], +) +def test_invalid_workspace_is_rejected(workspace_factory, error): + args = make_inputs(T=17) + required = flash_kda.allocate_workspace(args["q"]).numel() + workspace = workspace_factory(required) + + with pytest.raises(RuntimeError, match=error): + flash_kda.fwd(**args, workspace=workspace) + + +def test_public_exports_are_explicit(): + assert set(flash_kda.__all__) == {"allocate_workspace", "fwd", "get_workspace_size"}