diff --git a/benchmarks/bench_cake_causal_conv1d_train.py b/benchmarks/bench_cake_causal_conv1d_det_train.py similarity index 95% rename from benchmarks/bench_cake_causal_conv1d_train.py rename to benchmarks/bench_cake_causal_conv1d_det_train.py index 3822cada399..6f5456df730 100644 --- a/benchmarks/bench_cake_causal_conv1d_train.py +++ b/benchmarks/bench_cake_causal_conv1d_det_train.py @@ -53,10 +53,10 @@ Example:: - python benchmarks/bench_cake_causal_conv1d_train.py - python benchmarks/bench_cake_causal_conv1d_train.py --channels 512 --tokens 250340 --json out.json - python benchmarks/bench_cake_causal_conv1d_train.py --weight-dtype bfloat16 --bias bfloat16 - python benchmarks/bench_cake_causal_conv1d_train.py --silu-grad-preactivation fp32 --channels 512 1280 + python benchmarks/bench_cake_causal_conv1d_det_train.py + python benchmarks/bench_cake_causal_conv1d_det_train.py --channels 512 --tokens 250340 --json out.json + python benchmarks/bench_cake_causal_conv1d_det_train.py --weight-dtype bfloat16 --bias bfloat16 + python benchmarks/bench_cake_causal_conv1d_det_train.py --silu-grad-preactivation fp32 --channels 512 1280 """ from __future__ import annotations @@ -69,14 +69,14 @@ import torch -from flashinfer.cake_causal_conv1d_train import ( - cake_causal_conv1d_train_backward, - cake_causal_conv1d_train_backward_workspace, - cake_causal_conv1d_train_forward, - cake_causal_conv1d_train_forward_workspace, - is_cake_causal_conv1d_train_supported, +from flashinfer.cake_causal_conv1d_det_train import ( + cake_causal_conv1d_det_train_backward, + cake_causal_conv1d_det_train_backward_workspace, + cake_causal_conv1d_det_train_forward, + cake_causal_conv1d_det_train_forward_workspace, + is_cake_causal_conv1d_det_train_supported, ) -from flashinfer.cake_causal_conv1d_train_backend import cake_jit as loader +from flashinfer.cake_causal_conv1d_det_train_backend import cake_jit as loader from flashinfer.testing import bench_gpu_time WIDTH = 4 @@ -282,21 +282,21 @@ def make(self, x, w, b, cu, mode: Mode): db = ( None if b is None else torch.empty(channels, dtype=b.dtype, device=x.device) ) - workspace = cake_causal_conv1d_train_backward_workspace( + workspace = cake_causal_conv1d_det_train_backward_workspace( tokens, channels, x.device ) - forward_workspace = cake_causal_conv1d_train_forward_workspace( + forward_workspace = cake_causal_conv1d_det_train_forward_workspace( tokens, channels, x.device ) def fwd(): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, bias=b, out=y, workspace=forward_workspace ) return y def bwd(dy): - return cake_causal_conv1d_train_backward( + return cake_causal_conv1d_det_train_backward( x, w, cu, @@ -560,9 +560,9 @@ def main(argv: Optional[list[str]] = None) -> int: ) device = torch.device("cuda", torch.cuda.current_device()) - if not is_cake_causal_conv1d_train_supported(device): + if not is_cake_causal_conv1d_det_train_supported(device): raise SystemExit( - "cake_causal_conv1d_train has no exported route on this device" + "cake_causal_conv1d_det_train has no exported route on this device" ) arch = loader.arch_for_capability(torch.cuda.get_device_capability(device)) route = mode.backward_route() diff --git a/benchmarks/bench_cake_gdn_det_train_kcp.py b/benchmarks/bench_cake_gdn_det_train_kcp.py index 2795b6b072c..bfa45334345 100644 --- a/benchmarks/bench_cake_gdn_det_train_kcp.py +++ b/benchmarks/bench_cake_gdn_det_train_kcp.py @@ -16,7 +16,7 @@ Benchmark of the context-parallel (KCP) Cake packed GDN training step: one process per rank (``torchrun --nproc_per_node N``), every rank timing the module boundary of its fragments -- the KCP causal convolution -(``flashinfer.cake_causal_conv1d_train.cake_causal_conv1d_kcp``, 10,240 +(``flashinfer.cake_causal_conv1d_det_train.cake_causal_conv1d_kcp``, 10,240 channels) followed by the KCP GDN step (``flashinfer.cake_gdn_det_train_kcp.chunk_gdn_det_kcp``) -- with CUDA events, the slowest rank's time counting for the iteration. @@ -75,8 +75,8 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -from flashinfer.cake_causal_conv1d_train import ( # noqa: E402 - cake_causal_conv1d, +from flashinfer.cake_causal_conv1d_det_train import ( # noqa: E402 + cake_causal_conv1d_det, cake_causal_conv1d_kcp, ) from flashinfer.cake_gdn_det_train_kcp import ( # noqa: E402 @@ -267,7 +267,7 @@ def _local_arm(problem: dict[str, Any], boundary: str) -> Callable[[], None] | N cu = local["cu_seqlens"] def forward(q, k, v, g, beta, x, w): - y = cake_causal_conv1d(x, w, cu) + y = cake_causal_conv1d_det(x, w, cu) out = chunk_gdn(q, k, v, g, beta, cu) return y, out diff --git a/docs/api/mamba.rst b/docs/api/mamba.rst index fb0fe8c3c6b..43279a4795b 100644 --- a/docs/api/mamba.rst +++ b/docs/api/mamba.rst @@ -31,23 +31,39 @@ preactivation (default) or, with ``silu_grad_preactivation="fp32"``, at the unrounded FP32 one for a weight gradient within about 2e-7 rel-L2 of exact, and an autograd wrapper that saves only the forward's inputs. The forward's co-resident CTAs take each channel tile's tokens in chunks from int32 counters -in a forward workspace (``cake_causal_conv1d_train_forward_workspace`` / -``cake_causal_conv1d_train_forward_workspace_bytes``; zero before its first +in a forward workspace (``cake_causal_conv1d_det_train_forward_workspace`` / +``cake_causal_conv1d_det_train_forward_workspace_bytes``; zero before its first launch and left zero by every launch, so one buffer serves every later call of one stream); when the caller passes none the package keeps one per device and stream. -.. currentmodule:: flashinfer.cake_causal_conv1d_train +.. currentmodule:: flashinfer.cake_causal_conv1d_det_train .. autosummary:: :toctree: ../generated - cake_causal_conv1d - cake_causal_conv1d_kcp - cake_causal_conv1d_kcp_multi - KcpRecomputeHandle - cake_causal_conv1d_train_forward - cake_causal_conv1d_train_backward - cake_causal_conv1d_train_backward_workspace_bytes - cake_causal_conv1d_train_backward_workspace - CakeCausalConv1dFunction + cake_causal_conv1d_det + cake_causal_conv1d_det_train_forward + cake_causal_conv1d_det_train_backward + cake_causal_conv1d_det_train_backward_workspace_bytes + cake_causal_conv1d_det_train_backward_workspace + CakeCausalConv1dDetFunction + +Context-parallel training causal conv1d (Cake) +---------------------------------------------- + +The context-parallel (KCP) route of the same convolution: every rank convolves +its fragments of the packed batch, and the first tokens of a fragment that +continues a document read the history rows of the preceding fragment through +the halo collective of the KDA step's context. ``flashinfer.cake_causal_conv1d_det_train`` +keeps the previous names ``cake_causal_conv1d_kcp``, ``cake_causal_conv1d_kcp_multi`` +and ``KcpRecomputeHandle`` as aliases of the same objects. + +.. currentmodule:: flashinfer.cake_causal_conv1d_det_train_kcp + +.. autosummary:: + :toctree: ../generated + + cake_causal_conv1d_det_kcp + cake_causal_conv1d_det_multi_kcp + RecomputeHandleDetKcp diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 93fc24dbb0a..2f6d219c17b 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -171,7 +171,16 @@ ) +# Compatibility module paths of the causal conv1d facade and its backend package; each +# resolves on first attribute access and loads no KCP host module. +_COMPAT_SUBMODULES = frozenset( + {"cake_causal_conv1d_train", "cake_causal_conv1d_train_backend"} +) + + def __getattr__(name: str): + if name in _COMPAT_SUBMODULES: + return importlib.import_module(f"{__name__}.{name}") if name not in _PRIMS_TS_LAZY_EXPORTS: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from . import fused_moe as _fused_moe @@ -334,22 +343,37 @@ def __dir__() -> list[str]: from . import ( cake_rmsnorm_silu_gate_train_backend as cake_rmsnorm_silu_gate_train_backend, ) -from .cake_causal_conv1d_train import ( - CakeCausalConv1dFunction as CakeCausalConv1dFunction, +from .cake_causal_conv1d_det_train import ( + CakeCausalConv1dDetFunction as CakeCausalConv1dDetFunction, +) +from .cake_causal_conv1d_det_train import ( + cake_causal_conv1d_det as cake_causal_conv1d_det, ) -from .cake_causal_conv1d_train import cake_causal_conv1d as cake_causal_conv1d -from .cake_causal_conv1d_train import ( - cake_causal_conv1d_train_backward as cake_causal_conv1d_train_backward, +from .cake_causal_conv1d_det_train import ( + cake_causal_conv1d_det_train_backward as cake_causal_conv1d_det_train_backward, ) -from .cake_causal_conv1d_train import ( - cake_causal_conv1d_train_backward_workspace_bytes as cake_causal_conv1d_train_backward_workspace_bytes, +from .cake_causal_conv1d_det_train import ( + cake_causal_conv1d_det_train_backward_workspace_bytes as cake_causal_conv1d_det_train_backward_workspace_bytes, ) -from .cake_causal_conv1d_train import ( - cake_causal_conv1d_train_forward as cake_causal_conv1d_train_forward, +from .cake_causal_conv1d_det_train import ( + cake_causal_conv1d_det_train_forward as cake_causal_conv1d_det_train_forward, ) -from .cake_causal_conv1d_train import ( - cake_causal_conv1d_train_forward_workspace_bytes as cake_causal_conv1d_train_forward_workspace_bytes, +from .cake_causal_conv1d_det_train import ( + cake_causal_conv1d_det_train_forward_workspace_bytes as cake_causal_conv1d_det_train_forward_workspace_bytes, ) + +# Compatibility names of the causal conv1d entry points, bound to the same objects as the det names. +CakeCausalConv1dFunction = CakeCausalConv1dDetFunction +cake_causal_conv1d = cake_causal_conv1d_det +cake_causal_conv1d_train_backward = cake_causal_conv1d_det_train_backward +cake_causal_conv1d_train_backward_workspace_bytes = ( + cake_causal_conv1d_det_train_backward_workspace_bytes +) +cake_causal_conv1d_train_forward = cake_causal_conv1d_det_train_forward +cake_causal_conv1d_train_forward_workspace_bytes = ( + cake_causal_conv1d_det_train_forward_workspace_bytes +) + from .cake_mhc_det_train import ( CakeMHCDetTrainFinishFunction as CakeMHCDetTrainFinishFunction, ) diff --git a/flashinfer/cake_causal_conv1d_det_train.py b/flashinfer/cake_causal_conv1d_det_train.py new file mode 100644 index 00000000000..3ce6dc399eb --- /dev/null +++ b/flashinfer/cake_causal_conv1d_det_train.py @@ -0,0 +1,382 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Cake packed width-4 depthwise causal convolution + SiLU training forward / +backward (SM100 / SM103 / SM107). + +The input is one packed BF16 ``[S, C]`` (or ``[1, S, C]``) activation whose +documents ``cu_seqlens`` (``int32 [num_docs + 1]`` device prefix sums) +delimit; every document is convolved on its own, with zero history before its +first token. With taps ``w[C, 4]`` (``w[c, 3]`` multiplies the current token) +and an optional per-channel bias ``b[C]``, the forward is, in FP32 with one +BF16 rounding, + + pre[t, c] = sum_k w[c, k] * x[t - 3 + k, c] + b[c] # x of another document is 0 + y[t, c] = bf16(pre * sigmoid(pre)) + +The taps and the bias are FP32 or BF16, in any combination (BF16 is widened to +FP32 on load; the arithmetic is FP32 either way); the bias is added last, after +the four taps, which is flash-linear-attention's order. The backward +recomputes ``pre`` in FP32 from ``x`` and evaluates the SiLU derivative in +FP32, + + dpre[t, c] = dy[t, c] * s * (1 + p * (1 - s)), s = sigmoid(p) + dx[t, c] = bf16(sum_k w[c, k] * dpre[t + 3 - k, c]) # dpre of another document is 0 + dweight[c, k] = sum_t dpre[t, c] * x[t - 3 + k, c] # fixed order + dbias[c] = sum_t dpre[t, c] # with a bias + +at ``p`` chosen by ``silu_grad_preactivation``: ``"bf16"`` (default), ``p = +bf16(pre)``, the rounded value the forward stores and the boundary +flash-linear-attention's backward uses (``dweight`` lands within about 3e-5 +rel-L2 of the exact FP64 gradient); ``"fp32"``, ``p = pre`` unrounded, with +the final ``dweight`` reduction in FP64 (within about 2e-7; a slower program). +``dx`` is within one BF16 ulp of exact under both. The backward is two +ordinary launches on one stream: a persistent main program writes ``dx`` and +one FP32 partial row of ``dweight`` (and ``dbias``) per CTA into a workspace, +and a reduce program launched programmatically dependent on it sums the rows +in a fixed order (no floating-point atomics), so the gradients are bitwise +reproducible for a given ``(S, C, device)`` and route. Neither launch is +cooperative: the backward keeps its speed while kernels of other streams +(NCCL collectives) are resident, which a whole-device cooperative grid had to +wait for. Nothing is saved from the forward but ``x``, ``w``, ``b`` and +``cu_seqlens``. + +The forward is one ordinary launch: the co-resident CTAs of each channel tile +take the tile's tokens in chunks from int32 counters on the device (the CTAs +that finish first take the tail), so ``y`` is the same bits whatever CTA +computes a chunk. The counters live in a forward workspace of +``cake_causal_conv1d_det_train_forward_workspace_bytes(S, C, device)`` bytes that is +zero before its first launch and that every launch leaves zero, so one buffer +serves every later forward call of one stream, CUDA Graph replays included. A +caller may own it (``workspace=``); otherwise the package keeps one zeroed +buffer per (device, current stream), and a call inside a CUDA Graph capture +allocates a zeroed buffer of its own. Two launches that may run at once need +two buffers. + +Inputs and outputs are contiguous (``C`` a multiple of 8); ``dweight`` and +``dbias`` are FP32 or BF16 (the parameter's dtype when not preallocated). +``S``, ``C`` and the document count are runtime scalars: they may change +between calls, the host never reads ``cu_seqlens`` (no synchronization), ``S +== 0`` returns empty outputs without launching, and the calls are CUDA Graph +capturable once the libraries are built. + +A call's arguments (bias, BF16 taps or gradient, the SiLU' policy) select its +route, and within a route the channel count selects one of the architecture's +programs through a host-side rule of the backend's registry (never the data), +so a given ``(S, C, device)`` and route always run the same program. The +workspace size is the maximum over every backward program of every route, so +one buffer of ``cake_causal_conv1d_det_train_backward_workspace_bytes(S, C, +device)`` serves whichever program a call selects. Before capturing a CUDA +Graph, build every library of the device's architecture (the backend's +``cake_jit.specs_for_arch``) so that no JIT build happens inside the capture; +the captured graph holds the program selected for its shape and route. + +Usage:: + + y = cake_causal_conv1d_det_train_forward(x, weight, cu_seqlens) + dx, dweight = cake_causal_conv1d_det_train_backward(x, weight, cu_seqlens, dy) + + # a bias (FP32 or BF16, like the taps) adds a third output + y = cake_causal_conv1d_det_train_forward(x, weight, cu_seqlens, bias=bias) + dx, dweight, dbias = cake_causal_conv1d_det_train_backward(x, weight, cu_seqlens, dy, bias=bias) + + # SiLU' at the unrounded FP32 preactivation for a ~2e-7 rel-L2 weight gradient + dx, dweight = cake_causal_conv1d_det_train_backward( + x, weight, cu_seqlens, dy, silu_grad_preactivation="fp32" + ) + + # autograd (saves x, weight, bias and cu_seqlens only) + y = cake_causal_conv1d_det(x, weight, cu_seqlens, bias) + + # caller-owned backward workspace (per-CTA FP32 partials, no zeroing needed) + ws = cake_causal_conv1d_det_train_backward_workspace(S, C, x.device) + dx, dweight = cake_causal_conv1d_det_train_backward(x, weight, cu_seqlens, dy, workspace=ws) + + # caller-owned forward workspace (zeroed int32 token-chunk counters, left zero + # by every launch: keep one per stream and reuse it for every later call) + fws = cake_causal_conv1d_det_train_forward_workspace(S, C, x.device) + y = cake_causal_conv1d_det_train_forward(x, weight, cu_seqlens, workspace=fws) + + # context-parallel training: this rank's fragments of the packed batch, with the + # history rows of the preceding fragments exchanged through the KCP context + # (FP32 taps, no bias) + y_local = cake_causal_conv1d_kcp(x_local, weight, ctx) + + # a layer's q / k / v convolutions with one halo exchange per direction + # (bitwise three single calls) + q, k, v = cake_causal_conv1d_kcp_multi((xq, xk, xv), (wq, wk, wv), ctx) +""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from . import cake_causal_conv1d_det_train_backend as _backend +from . import cake_causal_conv1d_det_train_kcp as _kcp_facade +from .api_logging import flashinfer_api +from .cake_causal_conv1d_det_train_backend import ( # noqa: F401 + ARCHES, + SUPPORTED_COMPUTE_CAPABILITIES, + CakeCausalConv1dDetFunction, + cake_causal_conv1d_det_train_backward_workspace, + cake_causal_conv1d_det_train_backward_workspace_bytes, + cake_causal_conv1d_det_train_forward_workspace, + cake_causal_conv1d_det_train_forward_workspace_bytes, + is_cake_causal_conv1d_det_train_supported, +) + +# The three stable entry points are real ``def``s here (public-API tooling reads the +# decorated signatures of this module); their bodies live in the self-contained +# backend package next to the generated CUDA sources. The signatures mirror the +# backend's exactly (a test asserts it). + + +@flashinfer_api +def cake_causal_conv1d_det_train_forward( + x: torch.Tensor, + weight: torch.Tensor, + cu_seqlens: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + workspace: Optional[torch.Tensor] = None, +) -> torch.Tensor: + r"""Packed width-4 depthwise causal convolution + SiLU, training forward. + + Parameters + ---------- + x : torch.Tensor + BF16 ``[S, C]`` or ``[1, S, C]`` contiguous packed activations of every + document; ``C`` is a multiple of 8. + weight : torch.Tensor + FP32 or BF16 ``[C, 4]`` contiguous taps; ``weight[c, 3]`` multiplies the + current token, ``weight[c, 0]`` the token three positions back. Any + other width is rejected. + cu_seqlens : torch.Tensor + ``int32 [num_docs + 1]`` contiguous device prefix sums of the document + lengths (``cu_seqlens[0] == 0``, ``cu_seqlens[-1] == S``; a repeated + value is an empty document). The host reads its metadata only: no + element is copied to the host and no stream is synchronized. + bias : torch.Tensor, optional + FP32 or BF16 ``[C]`` contiguous per-channel bias, added after the four + taps; its dtype is independent of ``weight``'s. + out : torch.Tensor, optional + Preallocated BF16 output of ``x``'s shape, contiguous. + workspace : torch.Tensor, optional + Caller-owned ``int32`` buffer from + :func:`cake_causal_conv1d_det_train_forward_workspace` (the token-chunk + counters: zero before its first launch, left zero by every launch, so + one buffer serves every later call of one stream, CUDA Graph replays + included; two launches that may run at once need two buffers). Its + size (:func:`cake_causal_conv1d_det_train_forward_workspace_bytes`) is the + most over every forward route. When omitted the package uses a zeroed + buffer it keeps per (device, current stream), or a zeroed buffer of the + call's own inside a CUDA Graph capture. + + Returns + ------- + torch.Tensor + BF16 ``y`` of ``x``'s shape: ``y[t, c] = bf16(silu(sum_k weight[c, k] * + x[t - 3 + k, c] + bias[c]))`` with zero history before each document's + first token. ``S == 0`` returns the empty output without launching. + """ + + return _backend.cake_causal_conv1d_det_train_forward( + x, weight, cu_seqlens, bias=bias, out=out, workspace=workspace + ) + + +@flashinfer_api +def cake_causal_conv1d_det_train_backward( + x: torch.Tensor, + weight: torch.Tensor, + cu_seqlens: torch.Tensor, + dy: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + silu_grad_preactivation: str = "bf16", + workspace: Optional[torch.Tensor] = None, + dx: Optional[torch.Tensor] = None, + dweight: Optional[torch.Tensor] = None, + dbias: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, ...]: + r"""Packed width-4 depthwise causal convolution + SiLU, training backward. + + Recomputes the FP32 preactivation from ``x``, evaluates the SiLU derivative + at it -- rounded to BF16 (``silu_grad_preactivation="bf16"``, the default) + or unrounded (``"fp32"``) -- and produces ``dx``, the weight gradient and, + with a bias, the bias gradient in two ordinary launches on the current + stream: the main program writes ``dx`` and per-CTA partial rows of the + gradients, and a reduce program launched programmatically dependent on it + sums the rows in a fixed order (no cooperative grid, so kernels resident + on other streams never stall the call). + + The call's arguments select the route (bias, BF16 among ``weight`` / + ``dweight``, the SiLU' policy) and the channel count selects the route's + program through a host-side rule of the backend's registry (never the + data), so a given ``(S, C, device)`` and route always run the same program. + Before capturing a CUDA Graph, build every library of the device's + architecture (the backend's ``cake_jit.specs_for_arch``) so that no JIT + build happens inside the capture; the captured graph holds the program + selected for its shape and route. + + Parameters + ---------- + x, weight, cu_seqlens, bias + The forward's inputs (same contracts as + :func:`cake_causal_conv1d_det_train_forward`); nothing else is saved. + dy : torch.Tensor + BF16 gradient of ``y``, ``x``'s shape, contiguous. + silu_grad_preactivation : {"bf16", "fp32"} + ``"bf16"``: SiLU' at the BF16-rounded recomputed preactivation (the + boundary flash-linear-attention's backward uses; ``dweight`` within + about 3e-5 rel-L2 of the exact FP64 gradient). ``"fp32"``: SiLU' at + the unrounded FP32 preactivation with an FP64 final reduction (about + 2e-7; a slower program). ``dx`` is within one BF16 ulp of exact under + both. + workspace : torch.Tensor, optional + Caller-owned ``uint8`` buffer from + :func:`cake_causal_conv1d_det_train_backward_workspace` (the per-CTA FP32 + ``dweight`` / ``dbias`` partials; fully written by the kernel, so it + needs no zeroing); allocated when omitted. Its size + (:func:`cake_causal_conv1d_det_train_backward_workspace_bytes`) is the + maximum over every backward program of every route, so one buffer + serves whichever program a call selects. + dx, dweight, dbias : torch.Tensor, optional + Preallocated outputs: BF16 of ``x``'s shape, ``[C, 4]`` and ``[C]`` + (each FP32 or BF16; the parameter's dtype when omitted), contiguous. + ``dbias`` needs ``bias``. + + Returns + ------- + ``(dx, dweight)``, with a third element ``dbias`` only when ``bias`` is + given: BF16 ``dx`` of ``x``'s shape, ``dx[t, c] = bf16(sum_k weight[c, k] * + dpre[t + 3 - k, c])`` over the tokens of the same document, the ``[C, 4]`` + ``dweight[c, k] = sum_t dpre[t, c] * x[t - 3 + k, c]`` and the ``[C]`` + ``dbias[c] = sum_t dpre[t, c]``, each reduced in a fixed order, so they are + bitwise reproducible for a given ``(S, C, device)`` and route. ``S == 0`` + returns the empty ``dx`` and zero gradients without launching. + """ + + return _backend.cake_causal_conv1d_det_train_backward( + x, + weight, + cu_seqlens, + dy, + bias=bias, + silu_grad_preactivation=silu_grad_preactivation, + workspace=workspace, + dx=dx, + dweight=dweight, + dbias=dbias, + ) + + +@flashinfer_api +def cake_causal_conv1d_det( + x: torch.Tensor, + weight: torch.Tensor, + cu_seqlens: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + silu_grad_preactivation: str = "bf16", +) -> torch.Tensor: + """Autograd-enabled packed width-4 depthwise causal convolution + SiLU. + + Parameters + ---------- + x : torch.Tensor + BF16 ``[S, C]`` or ``[1, S, C]`` contiguous packed activations on a + supported CUDA device. + weight : torch.Tensor + FP32 or BF16 ``[C, 4]`` contiguous taps. + cu_seqlens : torch.Tensor + ``int32 [num_docs + 1]`` device prefix sums of the document lengths. + bias : torch.Tensor, optional + FP32 or BF16 ``[C]`` contiguous per-channel bias. + silu_grad_preactivation : {"bf16", "fp32"} + The weight gradient's SiLU' policy (see + :func:`cake_causal_conv1d_det_train_backward`). + + Returns + ------- + torch.Tensor + BF16 ``y`` of ``x``'s shape. The backward saves only ``x``, ``weight``, + ``bias`` and ``cu_seqlens`` and launches the fused ``dx`` + + deterministic ``dweight`` (+ ``dbias``) kernel of + :func:`cake_causal_conv1d_det_train_backward`; the parameters' gradients + come in the parameters' dtypes. + """ + + return _backend.cake_causal_conv1d_det( + x, weight, cu_seqlens, bias, silu_grad_preactivation=silu_grad_preactivation + ) + + +# The context-parallel (KCP) entry points live in ``flashinfer.cake_causal_conv1d_det_train_kcp``; their previous +# names stay bound here to the same objects. +cake_causal_conv1d_kcp = _kcp_facade.cake_causal_conv1d_det_kcp +cake_causal_conv1d_kcp_multi = _kcp_facade.cake_causal_conv1d_det_multi_kcp + + +# Compatibility names of the single-rank entry points, each bound to the same object as its det name +# (one ``flashinfer_api`` wrapper per entry point). They are attributes, not ``__all__`` entries. +CakeCausalConv1dFunction = CakeCausalConv1dDetFunction +cake_causal_conv1d = cake_causal_conv1d_det +cake_causal_conv1d_train_backward = cake_causal_conv1d_det_train_backward +cake_causal_conv1d_train_backward_workspace = ( + cake_causal_conv1d_det_train_backward_workspace +) +cake_causal_conv1d_train_backward_workspace_bytes = ( + cake_causal_conv1d_det_train_backward_workspace_bytes +) +cake_causal_conv1d_train_forward = cake_causal_conv1d_det_train_forward +cake_causal_conv1d_train_forward_workspace = ( + cake_causal_conv1d_det_train_forward_workspace +) +cake_causal_conv1d_train_forward_workspace_bytes = ( + cake_causal_conv1d_det_train_forward_workspace_bytes +) +is_cake_causal_conv1d_train_supported = is_cake_causal_conv1d_det_train_supported + + +def __getattr__(name: str) -> Any: + # ``KcpRecomputeHandle`` (the KCP route's checkpoint recompute handle) resolves on first access, so importing the + # facade does not import the KCP host module (torch.distributed, the KCP loader), like the entry point above. + if name == "KcpRecomputeHandle": + from .cake_causal_conv1d_det_train_backend.kcp import RecomputeHandleDetKcp + + return RecomputeHandleDetKcp + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + "ARCHES", + "SUPPORTED_COMPUTE_CAPABILITIES", + "CakeCausalConv1dDetFunction", + "KcpRecomputeHandle", # noqa: F822 - resolved by __getattr__ + "cake_causal_conv1d_det", + "cake_causal_conv1d_det_train_backward", + "cake_causal_conv1d_det_train_backward_workspace", + "cake_causal_conv1d_det_train_backward_workspace_bytes", + "cake_causal_conv1d_det_train_forward", + "cake_causal_conv1d_det_train_forward_workspace", + "cake_causal_conv1d_det_train_forward_workspace_bytes", + "cake_causal_conv1d_kcp", + "cake_causal_conv1d_kcp_multi", + "is_cake_causal_conv1d_det_train_supported", +] diff --git a/flashinfer/cake_causal_conv1d_train_backend/.gitignore b/flashinfer/cake_causal_conv1d_det_train_backend/.gitignore similarity index 100% rename from flashinfer/cake_causal_conv1d_train_backend/.gitignore rename to flashinfer/cake_causal_conv1d_det_train_backend/.gitignore diff --git a/flashinfer/cake_causal_conv1d_det_train_backend/__init__.py b/flashinfer/cake_causal_conv1d_det_train_backend/__init__.py new file mode 100644 index 00000000000..b4918d91ae4 --- /dev/null +++ b/flashinfer/cake_causal_conv1d_det_train_backend/__init__.py @@ -0,0 +1,894 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Backend of :mod:`flashinfer.cake_causal_conv1d_det_train`: the Cake packed width-4 +depthwise causal convolution + SiLU training forward / backward (SM100 / SM103 / +SM107) as one self-contained package (host API bodies, the ``cake_jit`` registry, +the pinned ``cake_build`` loader and the generated CUDA sources under ``csrc/``). +The stable, ``@flashinfer_api``-decorated entry points and the full user-facing +documentation live in :mod:`flashinfer.cake_causal_conv1d_det_train`; import that +module. + +Host design (the knobs are keywords; no environment variable): + +- Taps ``weight[C, 4]`` are FP32 or BF16; an optional per-channel ``bias[C]`` + is FP32 or BF16, independently of the taps. The kernels widen BF16 to FP32 + on load and compute in FP32; the bias is added LAST, after the four taps + (FLA's order). ``dweight`` / ``dbias`` default to the parameter's dtype and + may be preallocated as FP32 or BF16. +- ``silu_grad_preactivation``: ``"bf16"`` (default, the shipped policy: SiLU' + at the BF16-rounded preactivation, the boundary FLA's backward uses) or + ``"fp32"`` (SiLU' at the unrounded FP32 preactivation and an FP64 final + reduction of ``dweight``: a weight gradient within ~2e-7 rel-L2 of exact). +- A call's arguments select its route (``cake_jit.forward_route`` / + ``backward_route``): the default routes run the shipped programs unchanged; + the other routes' programs read the ``flags`` bit mask (``cake_jit.FLAG_*``, + derived here from the actual tensors: taps, the ``dweight`` out tensor, the + bias, the ``dbias`` out tensor) for the BF16 / FP32 load and store forms. + An architecture that lacks a route refuses the call with a ``ValueError`` + naming the route. +- The backward workspace is sized for every backward program of the device's + architecture (the bias programs' partials carry a fifth, dbias, slab), so one + buffer serves every route, token count and CUDA Graph. +- The forward (the dynamic token-chunk programs) takes its token chunks from + int32 counters, two per channel tile, in a forward workspace that is zero + before its first launch and that every launch leaves zero: a caller-owned + buffer (``workspace=``, from :func:`cake_causal_conv1d_det_train_forward_workspace`) + or, when omitted, a zeroed buffer this module keeps per (device, current + stream), grown as needed and never zeroed again; inside a CUDA Graph capture + an omitted workspace is a zeroed buffer of the call's own (its zeroing is + captured with the launch), so a graph never shares counters with eager + calls or with another graph. +""" + +from __future__ import annotations + +import functools +from typing import Optional + +import torch + +from . import cake_jit as _loader +from .cake_jit import ( + ARCH_BY_CAPABILITY, + ARCHES, + CHANNEL_MULTIPLE, + FLAG_B_BF16, + FLAG_DB_BF16, + FLAG_DW_BF16, + FLAG_W_BF16, + INT32_MAX, + REQUIRED_KINDS, + WIDTH, + arch_for_capability, + backward_route, + counter_words, + forward_route, + forward_workspace_words, + launch_grid, + partial_bytes, + partial_rows, + plan, + route_applies, + route_modules, + run_module, + run_modules, + selected_modules, + workspace_layout, + workspace_owners, +) + +SUPPORTED_COMPUTE_CAPABILITIES: tuple[tuple[int, int], ...] = tuple( + sorted(ARCH_BY_CAPABILITY) +) +# The dtypes of the taps, the bias and their gradients (FP32 or BF16, any combination). +PARAMETER_DTYPES = (torch.float32, torch.bfloat16) +# The two SiLU'-evaluation policies of the weight gradient (``silu_grad_preactivation``). +SILU_GRAD_PREACTIVATIONS = ("bf16", "fp32") +# The forward workspaces of the calls that pass none: (device index, CUDA stream +# handle) -> a zeroed int32 buffer, grown as needed. The launches of one stream +# run in order and each leaves its counters zero, so the stream's buffer serves +# every later forward call on it. +_FORWARD_WORKSPACES: dict[tuple[int, int], torch.Tensor] = {} + + +def _device_index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() + + +@functools.cache +def _arch_of_index(device_index: int) -> Optional[str]: + """The exported architecture of one CUDA device index, queried once per process.""" + + return arch_for_capability(tuple(torch.cuda.get_device_capability(device_index))) + + +def _device_arch(device: torch.device) -> str: + if device.type != "cuda": + raise ValueError( + f"Cake causal conv1d training requires CUDA tensors, got {device}" + ) + arch = _arch_of_index(_device_index(device)) + if arch is None: + raise ValueError( + "Cake causal conv1d training supports compute capability " + f"{sorted(ARCH_BY_CAPABILITY)} only, got " + f"{tuple(torch.cuda.get_device_capability(device))}" + ) + return arch + + +def is_cake_causal_conv1d_det_train_supported(device: torch.device) -> bool: + """Whether ``device`` has the exported routes every architecture ships (``fwd`` and ``bwd``). + + The optional routes (bias, BF16 taps / gradients, FP32 SiLU') are checked + per call: a call that needs one the architecture lacks raises a + ``ValueError`` naming it. + """ + + device = torch.device(device) + if device.type != "cuda": + return False + capability = tuple(torch.cuda.get_device_capability(device)) + return all( + route_applies(device_capability=capability, kind=kind) + for kind in REQUIRED_KINDS + ) + + +def _require_tensor(tensor: object, name: str) -> torch.Tensor: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + return tensor + + +def _check_placement( + tensor: torch.Tensor, + name: str, + *, + device: torch.device, + dtype: torch.dtype | tuple[torch.dtype, ...], +) -> None: + """On ``device``, of one of the ``dtype``\\ s, contiguous (``ValueError`` otherwise).""" + + dtypes = (dtype,) if isinstance(dtype, torch.dtype) else tuple(dtype) + if tensor.device != device: + raise ValueError(f"{name} must live on {device}, got {tensor.device}") + if tensor.dtype not in dtypes: + wanted = " or ".join(str(d) for d in dtypes) + raise ValueError(f"{name} must be {wanted}, got {tensor.dtype}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def _tokens( + tensor: object, + name: str, + *, + device: torch.device, + shape: Optional[tuple[int, ...]] = None, +) -> tuple[int, int, torch.Tensor]: + """Validate a BF16 ``[S, C]`` or ``[1, S, C]`` token tensor; return ``(S, C, [S, C] view)``. + + ``shape`` pins the exact shape another token tensor of the call has (``dy`` + against ``x``). Both kernels read the packed rows as 16-byte vectors of eight + channels, so ``C`` is a multiple of eight; ``S`` and ``C`` are 32-bit kernel + scalars (``S x C`` itself may exceed 2^31: the kernels address with 64-bit + offsets). + """ + + t = _require_tensor(tensor, name) + if t.ndim == 3 and t.shape[0] == 1: + S, C = int(t.shape[1]), int(t.shape[2]) + elif t.ndim == 2: + S, C = int(t.shape[0]), int(t.shape[1]) + else: + raise ValueError( + f"{name} must be a [S, C] or [1, S, C] tensor, got shape {tuple(t.shape)}" + ) + if shape is not None and tuple(t.shape) != tuple(shape): + raise ValueError(f"{name} must have shape {tuple(shape)}, got {tuple(t.shape)}") + _check_placement(t, name, device=device, dtype=torch.bfloat16) + if C <= 0 or C % CHANNEL_MULTIPLE: + raise ValueError( + f"{name} must have a channel count C that is a positive multiple of " + f"{CHANNEL_MULTIPLE}, got C={C}" + ) + if S > INT32_MAX or C > INT32_MAX: + raise ValueError( + f"{name}: S={S} and C={C} must each fit a 32-bit kernel scalar" + ) + return S, C, t.view(S, C) + + +def _weight( + weight: object, + C: int, + *, + device: torch.device, + name: str = "weight", + dtypes: tuple[torch.dtype, ...] = (torch.float32,), +) -> torch.Tensor: + """Validate the ``[C, 4]`` taps of one of ``dtypes`` (the training API passes FP32 and BF16; the KCP path + keeps the default, FP32 only).""" + + w = _require_tensor(weight, name) + if w.ndim != 2 or int(w.shape[1]) != WIDTH: + raise ValueError( + f"{name} must be a [C, {WIDTH}] tensor of a width-{WIDTH} convolution " + f"(tap {WIDTH - 1} multiplies the current token), got shape {tuple(w.shape)}" + ) + if int(w.shape[0]) != C: + raise ValueError( + f"{name} must have one row per channel ([{C}, {WIDTH}]), got shape {tuple(w.shape)}" + ) + _check_placement(w, name, device=device, dtype=dtypes) + return w + + +def _bias(bias: object, C: int, *, device: torch.device) -> torch.Tensor: + """Validate the per-channel ``[C]`` bias (FP32 or BF16, independently of the taps).""" + + b = _require_tensor(bias, "bias") + if b.ndim != 1 or int(b.shape[0]) != C: + raise ValueError( + f"bias must be a [C] tensor with one value per channel ([{C}]), got shape {tuple(b.shape)}" + ) + _check_placement(b, "bias", device=device, dtype=PARAMETER_DTYPES) + return b + + +def _cu_seqlens( + cu_seqlens: object, *, device: torch.device +) -> tuple[torch.Tensor, int]: + """Validate the ``int32 [num_docs + 1]`` device prefix sums; return ``(cu_seqlens, num_docs)``. + + The host reads the tensor's metadata only: no element is copied to the host + and no stream is synchronized, so the call is CUDA Graph capturable. + """ + + cu = _require_tensor(cu_seqlens, "cu_seqlens") + if cu.ndim != 1: + raise ValueError( + f"cu_seqlens must be a 1-D [num_docs + 1] tensor, got shape {tuple(cu.shape)}" + ) + _check_placement(cu, "cu_seqlens", device=device, dtype=torch.int32) + num_docs = int(cu.numel()) - 1 + if num_docs < 1: + raise ValueError( + "cu_seqlens must hold at least two prefix sums (one document), got " + f"{int(cu.numel())} elements" + ) + return cu, num_docs + + +def _output( + tensor: Optional[object], + name: str, + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + dtypes: Optional[tuple[torch.dtype, ...]] = None, +) -> torch.Tensor: + """A caller-provided output (exact shape, contiguous, of ``dtype`` or of one of ``dtypes`` when given) or a + fresh ``dtype`` one.""" + + if tensor is None: + return torch.empty(shape, dtype=dtype, device=device) + out = _require_tensor(tensor, name) + if tuple(out.shape) != tuple(shape): + raise ValueError( + f"{name} must have shape {tuple(shape)}, got {tuple(out.shape)}" + ) + _check_placement(out, name, device=device, dtype=dtypes or dtype) + return out + + +def _silu_grad_preactivation(value: object) -> str: + if value not in SILU_GRAD_PREACTIVATIONS: + raise ValueError( + f"silu_grad_preactivation must be one of {SILU_GRAD_PREACTIVATIONS}, got {value!r}" + ) + return str(value) + + +def _scalars(scalars: dict[str, int]) -> dict[str, int]: + """The launch scalars :func:`cake_jit.plan` produced (``flags`` is set beside them by the caller).""" + + return {key: int(scalars[key]) for key in _loader.PLAN_SCALARS} + + +def _flags( + w: torch.Tensor, + dw: Optional[torch.Tensor] = None, + b: Optional[torch.Tensor] = None, + db: Optional[torch.Tensor] = None, +) -> int: + """The ``flags`` bit mask of a ``dtype_flags`` program from the ACTUAL tensors of the call.""" + + flags = 0 + for tensor, bit in ( + (w, FLAG_W_BF16), + (dw, FLAG_DW_BF16), + (b, FLAG_B_BF16), + (db, FLAG_DB_BF16), + ): + if tensor is not None and tensor.dtype == torch.bfloat16: + flags |= bit + return flags + + +def _route(arch: str, kind: str, *, needs: str, C: int, S: int) -> tuple[str, ...]: + """The programs ``kind`` launches on ``arch`` for the call (``ValueError`` naming the route when the + architecture's export lacks it).""" + + if kind not in _loader.ROUTES.get(arch, {}): + raise ValueError( + f"{needs} needs the {kind!r} route, which this package does not export for " + f"{arch}; its routes: {sorted(_loader.ROUTES.get(arch, {}))}" + ) + return selected_modules(arch, kind, C=C, S=S) + + +def _launch( + record: dict, S: int, C: int, num_docs: int, *, device_index: int +) -> tuple[dict[str, int], tuple[int, int, int]]: + """The kernel scalars and the grid of one launch of ``record`` (the host plan of ``cake_jit``). + + CTA ``b`` owns channel tile ``b % n_ch_tiles``, so every channel tile needs + a CTA of the launched program's grid on the device: a ``C`` beyond ``grid + x threads x vec`` channels is refused rather than left partly unwritten. + The caller launches this grid (``run_module`` / ``run_modules`` ``grid``), + so a call plans it once. + """ + + scalars = plan( + S, + C, + num_docs, + geometry=record["grid_rule"]["geometry"], + role=record["role"], + ) + grid = launch_grid(record, device_index=device_index, C=C, S=S) + if scalars["n_ch_tiles"] > grid[0]: + raise ValueError( + f"C={C} spans {scalars['n_ch_tiles']} channel tiles but the launch grid on this " + f"device is {grid[0]} CTAs; every channel tile needs a CTA" + ) + return scalars, grid + + +def cake_causal_conv1d_det_train_backward_workspace_bytes( + S: int, C: int, device: torch.device +) -> int: + """Bytes of the caller-owned backward workspace for ``C`` channels on ``device``. + + The workspace holds the FP32 ``[rows, C, 4]`` ``dweight`` partials (and, for + the bias routes, a ``[rows, C]`` dbias slab after them) of the backward's + main program, one row per ``n_ch_tiles`` CTAs of a program's launch grid + (``rows = ceil(grid / n_ch_tiles)``, ``grid = grid_mult x ctas_per_sm x SM + count``: the bytes grow with the program's over-decomposition factor + ``grid_mult``; a call whose grid drops CTAs to whole channel-tile ranks, + ``cake_jit.main_grid``, writes fewer rows), sized by the largest such + buffer over every backward program of every backward route of the + architecture: the call's arguments select the route and the channel count + the program, every program writes the leading rows of its own grid inside + this buffer, and its reduce + program sums exactly those. So one buffer serves every route (the bias + routes' fifth slab makes it 25 % larger than the dweight partials alone, + for every caller) and, as it depends on ``C`` and the device only, every + token count; the main program fully writes the rows the reduce program + sums, so the buffer needs no zeroing between launches or CUDA Graph + replays. + """ + + device = torch.device(device) + S, C = int(S), int(C) + if S < 0: + raise ValueError("S must be non-negative") + if C <= 0 or C % CHANNEL_MULTIPLE: + raise ValueError( + f"C must be a positive multiple of {CHANNEL_MULTIPLE}, got {C}" + ) + arch = _device_arch(device) + device_index = _device_index(device) + return max( + int(workspace_layout(owner, C, device_index=device_index)["total_bytes"]) + for owner in workspace_owners(arch) + ) + + +def cake_causal_conv1d_det_train_backward_workspace( + S: int, C: int, device: torch.device +) -> torch.Tensor: + """Allocate a backward workspace (see ``..._workspace_bytes``; no zeroing is needed).""" + + nbytes = cake_causal_conv1d_det_train_backward_workspace_bytes(S, C, device) + return torch.empty(nbytes, dtype=torch.uint8, device=torch.device(device)) + + +def cake_causal_conv1d_det_train_forward_workspace_bytes( + S: int, C: int, device: torch.device +) -> int: + """Bytes of the caller-owned forward workspace for ``C`` channels on ``device``. + + The forward's programs take their token chunks from int32 counters, two + per channel tile (the next-chunk counter and the done-CTA counter of the + tile), sized by the most over every forward program of every forward route + of the architecture (``cake_jit.forward_workspace_words``), so one buffer + serves every route; the size depends on ``C`` only (``S`` is checked and + otherwise unused, as for the backward helper). The buffer must be zero + before its first launch; every launch leaves it zero, so one buffer serves + every later forward call of one stream, CUDA Graph replays included, + without zeroing. Two launches that may run at once (on two streams, or + two graphs replayed concurrently) need two buffers. + """ + + device = torch.device(device) + S, C = int(S), int(C) + if S < 0: + raise ValueError("S must be non-negative") + if C <= 0 or C % CHANNEL_MULTIPLE: + raise ValueError( + f"C must be a positive multiple of {CHANNEL_MULTIPLE}, got {C}" + ) + return 4 * forward_workspace_words(_device_arch(device), C) + + +def cake_causal_conv1d_det_train_forward_workspace( + S: int, C: int, device: torch.device +) -> torch.Tensor: + """Allocate a zeroed int32 forward workspace (see ``..._workspace_bytes``); keep it for every later forward + call of one stream.""" + + nbytes = cake_causal_conv1d_det_train_forward_workspace_bytes(S, C, device) + return torch.zeros(nbytes // 4, dtype=torch.int32, device=torch.device(device)) + + +def _forward_work( + workspace: object, words: int, *, device: torch.device, device_index: int +) -> Optional[torch.Tensor]: + """The ``work`` counters of one forward launch needing ``words`` int32 words. + + A caller-owned ``workspace`` (1-D int32, contiguous, on ``device``, at least + ``words`` words, zero before its first launch) is used as is. Without one: + inside a CUDA Graph capture a fresh zeroed buffer (its zeroing is captured + with the launch, so the graph shares no counters with eager calls or other + graphs); else the zeroed buffer this module keeps for the device and the + current stream, replaced by a larger zeroed one when it is too small (the + stream-ordered allocator frees the old one after the launches queued on + the stream). ``None`` when the program takes no counters. + """ + + if workspace is not None: + ws = _require_tensor(workspace, "workspace") + if ws.ndim != 1: + raise ValueError( + f"workspace must be a 1-D int32 CUDA tensor, got shape {tuple(ws.shape)}" + ) + _check_placement(ws, "workspace", device=device, dtype=torch.int32) + if ws.numel() < words: + raise ValueError( + f"workspace holds {ws.numel()} int32 words; this call needs at least " + f"{words}; allocate cake_causal_conv1d_det_train_forward_workspace(S, C, device)" + ) + return ws + if words == 0: + return None + if torch.cuda.is_current_stream_capturing(): + return torch.zeros(words, dtype=torch.int32, device=device) + key = (device_index, int(torch.cuda.current_stream(device_index).cuda_stream)) + held = _FORWARD_WORKSPACES.get(key) + if held is None or held.numel() < words: + held = torch.zeros(words, dtype=torch.int32, device=device) + _FORWARD_WORKSPACES[key] = held + return held + + +def _partial_view( + workspace: object, record: dict, *, C: int, rows: int, device: torch.device +) -> torch.Tensor: + """The FP32 ``[rows, C, taps]`` capacity view of this launch's partial rows inside ``workspace``. + + ``taps`` is the record's (4, or 5 for a bias program whose kernel lays the + ``[rows, C, 4]`` tap slab out first and the ``[rows, C]`` dbias slab after + it); the view only sizes the kernel's buffer. + """ + + ws = _require_tensor(workspace, "workspace") + if ws.ndim != 1: + raise ValueError( + f"workspace must be a 1-D uint8 CUDA tensor, got shape {tuple(ws.shape)}" + ) + _check_placement(ws, "workspace", device=device, dtype=torch.uint8) + if ws.data_ptr() % 16 != 0: + raise ValueError("workspace must be 16-byte aligned") + nbytes = partial_bytes(record, C, rows) + if ws.numel() < nbytes: + raise ValueError( + f"workspace holds {ws.numel()} bytes; this call needs at least " + f"{nbytes}; allocate " + "cake_causal_conv1d_det_train_backward_workspace_bytes(S, C, device) bytes" + ) + taps = int(record["workspace_rule"]["taps"]) + return ws[:nbytes].view(torch.float32).view(rows, C, taps) + + +def cake_causal_conv1d_det_train_forward( + x: torch.Tensor, + weight: torch.Tensor, + cu_seqlens: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + workspace: Optional[torch.Tensor] = None, +) -> torch.Tensor: + r"""Packed width-4 depthwise causal convolution + SiLU, training forward. + + Parameters + ---------- + x : torch.Tensor + BF16 ``[S, C]`` or ``[1, S, C]`` contiguous packed activations of every + document; ``C`` is a multiple of 8. + weight : torch.Tensor + FP32 or BF16 ``[C, 4]`` contiguous taps; ``weight[c, 3]`` multiplies the + current token, ``weight[c, 0]`` the token three positions back. + cu_seqlens : torch.Tensor + ``int32 [num_docs + 1]`` contiguous device prefix sums of the document + lengths (``cu_seqlens[0] == 0``, ``cu_seqlens[-1] == S``; a repeated + value is an empty document). The host reads its metadata only. + bias : torch.Tensor, optional + FP32 or BF16 ``[C]`` contiguous per-channel bias (its dtype is + independent of ``weight``'s), added after the four taps. + out : torch.Tensor, optional + Preallocated BF16 output of ``x``'s shape, contiguous. + workspace : torch.Tensor, optional + Caller-owned int32 buffer from + :func:`cake_causal_conv1d_det_train_forward_workspace` (the token-chunk + counters, zero before its first launch; every launch leaves it zero, + so one buffer serves every later call of one stream). Two launches + that may run at once need two buffers. When omitted the package + uses a zeroed buffer it keeps per (device, current stream), or a + fresh zeroed one inside a CUDA Graph capture. + + Returns + ------- + torch.Tensor + BF16 ``y`` of ``x``'s shape: ``y[t, c] = bf16(silu(sum_k weight[c, k] * + x[t - 3 + k, c] + bias[c]))`` with zero history before each document's + first token. ``S == 0`` returns the empty output without launching. + """ + + x = _require_tensor(x, "x") + device = x.device + arch = _device_arch(device) + S, C, x_flat = _tokens(x, "x", device=device) + w = _weight(weight, C, device=device, dtypes=PARAMETER_DTYPES) + b = None if bias is None else _bias(bias, C, device=device) + cu, num_docs = _cu_seqlens(cu_seqlens, device=device) + y = _output(out, "out", shape=tuple(x.shape), dtype=torch.bfloat16, device=device) + if S == 0: + return y + kind = forward_route(bias=b is not None, bf16_taps=w.dtype == torch.bfloat16) + names = _route(arch, kind, needs="this forward (bias / BF16 taps)", C=C, S=S) + device_index = _device_index(device) + scalars, grid = _launch( + _loader.MODULES[names[0]], S, C, num_docs, device_index=device_index + ) + values: dict[str, object] = { + "x": x_flat, + "w": w, + "cu": cu, + "y": y.view(S, C), + **_scalars(scalars), + "flags": _flags(w, b=b), + } + if b is not None: + values["b"] = b + words = max( + counter_words(_loader.MODULES[name], scalars["n_ch_tiles"]) for name in names + ) + work = _forward_work(workspace, words, device=device, device_index=device_index) + if work is not None: + values["work"] = work + for name in names: + run_module( + name, + values, + device_index=device_index, + grid=grid if name == names[0] else None, + ) + return y + + +def cake_causal_conv1d_det_train_backward( + x: torch.Tensor, + weight: torch.Tensor, + cu_seqlens: torch.Tensor, + dy: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + silu_grad_preactivation: str = "bf16", + workspace: Optional[torch.Tensor] = None, + dx: Optional[torch.Tensor] = None, + dweight: Optional[torch.Tensor] = None, + dbias: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, ...]: + r"""Packed width-4 depthwise causal convolution + SiLU, training backward. + + Parameters + ---------- + x, weight, cu_seqlens, bias + The forward's inputs (same contracts as + :func:`cake_causal_conv1d_det_train_forward`); nothing else is saved. + dy : torch.Tensor + BF16 gradient of ``y``, ``x``'s shape, contiguous. + silu_grad_preactivation : {"bf16", "fp32"} + Where the SiLU derivative of the weight gradient is evaluated: + ``"bf16"`` at the BF16-rounded recomputed preactivation (the shipped + policy and FLA's boundary; ``dweight`` within ~3e-5 rel-L2 of exact), + ``"fp32"`` at the unrounded FP32 preactivation with an FP64 final + reduction (~2e-7). ``dx`` is within one BF16 ulp of exact either way. + workspace : torch.Tensor, optional + Caller-owned ``uint8`` buffer from + :func:`cake_causal_conv1d_det_train_backward_workspace`; allocated when + omitted. + dx, dweight, dbias : torch.Tensor, optional + Preallocated outputs: BF16 of ``x``'s shape, ``[C, 4]`` and ``[C]`` + (FP32 or BF16 each; the parameter's dtype when omitted), contiguous. + ``dbias`` needs ``bias``. + + Returns + ------- + ``(dx, dweight)``, or ``(dx, dweight, dbias)`` when ``bias`` is given: BF16 + ``dx`` of ``x``'s shape and the ``[C, 4]`` weight gradient (and ``[C]`` bias + gradient), each accumulated in a fixed order, so they are bitwise + reproducible for a given ``(S, C, device)`` and route. ``S == 0`` returns + the empty ``dx`` and zero gradients without launching. + + The call's arguments select the route (``cake_jit.backward_route``: + ``bias``, the policy, BF16 among ``weight`` / ``dweight``), and within a + route the channel count selects the program through the registry's + host-side rule (``cake_jit.selected_modules``; the data is never read), so a + given ``(S, C, device)`` and route always run the same program. The + ``workspace`` is sized for every program of every route + (:func:`cake_causal_conv1d_det_train_backward_workspace_bytes`); the selected + program uses the leading rows of its own grid. + + A call is two ordinary launches on the current stream: the selected + program writes ``dx`` and one FP32 partial row of the gradients per CTA, + then the route's reduce program, launched programmatically dependent on + it (the launch latency overlaps the main kernel's tail; the dependency is + a stream-order edge, so the pair captures into a CUDA Graph), sums the + rows per channel in a fixed order into ``dweight`` (and ``dbias``) in the + requested dtypes. Neither launch is cooperative, so the backward keeps + its speed while kernels of other streams (NCCL collectives) are resident: + a cooperative grid could not start until they finished. Build every + library of the device's architecture (``cake_jit.specs_for_arch``) before + capturing a CUDA Graph: the capture holds the programs selected for its + shape and route, and no JIT build may happen inside it. + """ + + policy = _silu_grad_preactivation(silu_grad_preactivation) + x = _require_tensor(x, "x") + device = x.device + arch = _device_arch(device) + S, C, x_flat = _tokens(x, "x", device=device) + w = _weight(weight, C, device=device, dtypes=PARAMETER_DTYPES) + b = None if bias is None else _bias(bias, C, device=device) + cu, num_docs = _cu_seqlens(cu_seqlens, device=device) + _, _, dy_flat = _tokens(dy, "dy", device=device, shape=tuple(x.shape)) + dx_out = _output( + dx, "dx", shape=tuple(x.shape), dtype=torch.bfloat16, device=device + ) + dw_out = _output( + dweight, + "dweight", + shape=(C, WIDTH), + dtype=w.dtype, + device=device, + dtypes=PARAMETER_DTYPES, + ) + if b is None: + if dbias is not None: + raise ValueError( + "dbias is the bias gradient's output; pass bias to produce it" + ) + db_out = None + else: + db_out = _output( + dbias, + "dbias", + shape=(C,), + dtype=b.dtype, + device=device, + dtypes=PARAMETER_DTYPES, + ) + if S == 0: + dw_out.zero_() + if db_out is None: + return dx_out, dw_out + db_out.zero_() + return dx_out, dw_out, db_out + kind = backward_route( + bias=b is not None, + fp32_grad=policy == "fp32", + bf16_taps_or_dweight=torch.bfloat16 in (w.dtype, dw_out.dtype), + ) + # Within the route the program is a pure function of the architecture and C + # (S never enters): its own grid sizes the partial rows this call writes + # inside the workspace. + names = _route( + arch, + kind, + needs=f"this backward (bias / BF16 taps or dweight / silu_grad_preactivation={policy!r})", + C=C, + S=S, + ) + record = _loader.MODULES[names[0]] + if record["workspace_rule"] is None: + raise RuntimeError( + f"the backward program {names[0]} selected for C={C} on {arch} owns no " + "partial workspace" + ) + device_index = _device_index(device) + scalars, grid = _launch(record, S, C, num_docs, device_index=device_index) + grid_x = grid[0] + if workspace is None: + workspace = cake_causal_conv1d_det_train_backward_workspace(S, C, device) + rows = partial_rows(grid_x, scalars["n_ch_tiles"]) + partial = _partial_view(workspace, record, C=C, rows=rows, device=device) + values: dict[str, object] = { + "x": x_flat, + "w": w, + "cu": cu, + "dy": dy_flat, + "dx": dx_out.view(S, C), + "dw": dw_out, + "partial": partial, + **_scalars(scalars), + "flags": _flags(w, dw_out, b, db_out), + "main_ctas": grid_x, + } + if b is not None: + values["b"] = b + values["db"] = db_out + # The main program (an ordinary launch) writes dx and the partial rows; the + # route's reduce program, programmatically dependent on it through its + # binding's launch attribute, sums into dweight (and dbias) on the same + # stream the rows of each channel tile that a grid of main_ctas CTAs wrote. + # Both are issued back to back in one stream context, so the reduce is + # queued while the main program runs. + run_modules( + ((names[0], grid), (record["reduce"], None)), + values, + device_index=device_index, + ) + if db_out is None: + return dx_out, dw_out + return dx_out, dw_out, db_out + + +class CakeCausalConv1dDetFunction(torch.autograd.Function): + """Autograd wrapper saving only ``(x, weight, cu_seqlens, bias)`` for the backward. + + ``forward(x, weight, cu_seqlens, bias=None, silu_grad_preactivation="bf16")`` + returns ``y``; the backward runs the fused ``dx`` + deterministic + ``dweight`` (+ ``dbias``) pair of launches of :func:`cake_causal_conv1d_det_train_backward` + on the incoming gradient (materialized as a contiguous BF16 tensor when + autograd hands over another layout or dtype) and returns the gradients of + ``weight`` and ``bias`` in the parameters' dtypes. + """ + + @staticmethod + def forward(ctx, x, weight, cu_seqlens, bias=None, silu_grad_preactivation="bf16"): + y = cake_causal_conv1d_det_train_forward(x, weight, cu_seqlens, bias=bias) + ctx.save_for_backward(x, weight, cu_seqlens, bias) + ctx.silu_grad_preactivation = _silu_grad_preactivation(silu_grad_preactivation) + return y + + @staticmethod + def backward(ctx, grad_y): + x, weight, cu_seqlens, bias = ctx.saved_tensors + if grad_y is None: + grad_y = torch.zeros_like(x) + dy = grad_y.to(torch.bfloat16).contiguous() + grads = cake_causal_conv1d_det_train_backward( + x, + weight, + cu_seqlens, + dy, + bias=bias, + silu_grad_preactivation=ctx.silu_grad_preactivation, + ) + dbias = grads[2] if bias is not None else None + return grads[0], grads[1], None, dbias, None + + +def cake_causal_conv1d_det( + x: torch.Tensor, + weight: torch.Tensor, + cu_seqlens: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + silu_grad_preactivation: str = "bf16", +) -> torch.Tensor: + """Autograd-enabled packed width-4 depthwise causal convolution + SiLU. + + Parameters + ---------- + x : torch.Tensor + BF16 ``[S, C]`` or ``[1, S, C]`` contiguous packed activations on a + supported CUDA device. + weight : torch.Tensor + FP32 or BF16 ``[C, 4]`` contiguous taps. + cu_seqlens : torch.Tensor + ``int32 [num_docs + 1]`` device prefix sums of the document lengths. + bias : torch.Tensor, optional + FP32 or BF16 ``[C]`` contiguous per-channel bias. + silu_grad_preactivation : {"bf16", "fp32"} + The weight gradient's SiLU' policy (see + :func:`cake_causal_conv1d_det_train_backward`). + + Returns + ------- + torch.Tensor + BF16 ``y`` of ``x``'s shape. The backward saves only ``x``, ``weight``, + ``cu_seqlens`` and ``bias`` and returns BF16 ``dx`` and the ``dweight`` + (and ``dbias``) in the parameters' dtypes. + """ + + return CakeCausalConv1dDetFunction.apply( + x, weight, cu_seqlens, bias, silu_grad_preactivation + ) + + +# Compatibility names of the entry points, bound to the same objects as the det names. They are attributes, not +# ``__all__`` entries. +CakeCausalConv1dFunction = CakeCausalConv1dDetFunction +cake_causal_conv1d = cake_causal_conv1d_det +cake_causal_conv1d_train_backward = cake_causal_conv1d_det_train_backward +cake_causal_conv1d_train_backward_workspace = ( + cake_causal_conv1d_det_train_backward_workspace +) +cake_causal_conv1d_train_backward_workspace_bytes = ( + cake_causal_conv1d_det_train_backward_workspace_bytes +) +cake_causal_conv1d_train_forward = cake_causal_conv1d_det_train_forward +cake_causal_conv1d_train_forward_workspace = ( + cake_causal_conv1d_det_train_forward_workspace +) +cake_causal_conv1d_train_forward_workspace_bytes = ( + cake_causal_conv1d_det_train_forward_workspace_bytes +) +is_cake_causal_conv1d_train_supported = is_cake_causal_conv1d_det_train_supported + + +__all__ = [ + "ARCHES", + "PARAMETER_DTYPES", + "SILU_GRAD_PREACTIVATIONS", + "SUPPORTED_COMPUTE_CAPABILITIES", + "CakeCausalConv1dDetFunction", + "cake_causal_conv1d_det", + "cake_causal_conv1d_det_train_backward", + "cake_causal_conv1d_det_train_backward_workspace", + "cake_causal_conv1d_det_train_backward_workspace_bytes", + "cake_causal_conv1d_det_train_forward", + "cake_causal_conv1d_det_train_forward_workspace", + "cake_causal_conv1d_det_train_forward_workspace_bytes", + "is_cake_causal_conv1d_det_train_supported", +] diff --git a/flashinfer/cake_causal_conv1d_det_train_backend/build.json b/flashinfer/cake_causal_conv1d_det_train_backend/build.json new file mode 100644 index 00000000000..3a391fe62ac --- /dev/null +++ b/flashinfer/cake_causal_conv1d_det_train_backend/build.json @@ -0,0 +1,1745 @@ +{ + "arch_flags": { + "sm_100a": [ + "-gencode=arch=compute_100a,code=sm_100a" + ], + "sm_103a": [ + "-gencode=arch=compute_103a,code=sm_103a" + ], + "sm_107a": [ + "-gencode=arch=compute_107a,code=sm_107a" + ] + }, + "common": { + "include_dirs": [ + "csrc" + ], + "link": { + "cudart": "shared", + "host_linker": "c++", + "tvm_ffi_symbols": "resolved at load; do not pass --no-undefined" + }, + "link_flags": [ + "-shared", + "-lcudart", + "-lcuda" + ], + "nvcc_flags": [ + "-std=c++17", + "-O3", + "-DNDEBUG", + "--compiler-options=-fPIC", + "--expt-relaxed-constexpr", + "-Xfatbin=-compress-all", + "--compress-mode=size" + ], + "nvcc_flags_if": [ + { + "cuda": ">=12.8", + "flags": [ + "-static-global-template-stub=false" + ] + } + ], + "requires": { + "cuda_toolkit": { + "min": "12.8" + }, + "tvm_ffi": { + "include_dirs": [ + "tvm_ffi.libinfo.find_include_path()", + "tvm_ffi.libinfo.find_dlpack_include_path()" + ], + "python": "apache-tvm-ffi>=0.1.11,<0.2" + } + } + }, + "generator": { + "cake_revision": "e312167e6b4f6d8b4193e218da74fa5b17d8ec14", + "cuda": "13.4.92", + "rendered_at": "2026-10-28T23:53:55Z" + }, + "libraries": [ + { + "arch": "sm_100a", + "closure_sha256": "0c024189d7c3bd05da6f1490afc67634ee7e3f58a4e61ff00e6cc506b5af84fb", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu": "9b3a96b5a62d108aeb37775f2c2a655119762d4b8150549dd8dc09171bfd7ddb", + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu": "69bfbc6ca775a99ab51fb672fb184c2fde979ae9a361694397c5f43c4cf6a365" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu", + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "bb51f380752185a4b95d4331fbb29f2ccb7d83a7e9364524e360dbf9e14f09c3", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu": "9b3a96b5a62d108aeb37775f2c2a655119762d4b8150549dd8dc09171bfd7ddb", + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu": "69bfbc6ca775a99ab51fb672fb184c2fde979ae9a361694397c5f43c4cf6a365" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu", + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "7dcdbb7e8d3440254921c72d8583186be8d419cf994a43008119e8d1f18fee19", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu": "9b3a96b5a62d108aeb37775f2c2a655119762d4b8150549dd8dc09171bfd7ddb", + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu": "69bfbc6ca775a99ab51fb672fb184c2fde979ae9a361694397c5f43c4cf6a365" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu", + "csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "9d0241f5bd09d302a6af4cc8ad32c8d4328e48dce1b9b6cce00b914555b5ed2b", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_149ca6f53934ba630033_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_149ca6f53934ba630033_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_binding.cu": "31eec24c0809110fde68c6cda89a6f56566da02df3a5bbea49f6e94a587b13ad", + "csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_kernel.cu": "da476a7b8586b1cfbc8698e2098505a8d62bb83ab12c4d3cf8f9d9d9ee85115d" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_kernel.cu", + "csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "9558a0989345a12d9d8c252d5f7275adb7e68586464c4dfb1ab65729e3324ba9", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_149ca6f53934ba630033_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_149ca6f53934ba630033_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_binding.cu": "31eec24c0809110fde68c6cda89a6f56566da02df3a5bbea49f6e94a587b13ad", + "csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_kernel.cu": "da476a7b8586b1cfbc8698e2098505a8d62bb83ab12c4d3cf8f9d9d9ee85115d" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_kernel.cu", + "csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "1b5c3354f07d7c98d9a46f1fb4f5ff0577465505a18a3482ca2af3182bf7d1f4", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_170d89481a616ae977b3_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_170d89481a616ae977b3_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu": "08e4640bf2f3e3ea36e368e7cec47a4fc11d5cafa44a2112a233abdd275203d2", + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu": "560673dded379b80ca6d345af04ae83b6bd23d136a2e6a95874455b3e70bdbf3" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu", + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "5895d7dc6bd5173b648d04cd11befe94643f6dd2e271c2619aab8cb3ed4ab61d", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_170d89481a616ae977b3_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_170d89481a616ae977b3_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu": "08e4640bf2f3e3ea36e368e7cec47a4fc11d5cafa44a2112a233abdd275203d2", + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu": "560673dded379b80ca6d345af04ae83b6bd23d136a2e6a95874455b3e70bdbf3" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu", + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "7d8f70385e027d68b5d5cc51f2f6b1c3f1e28ea478ae81b945dcbbca4d3be431", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_170d89481a616ae977b3_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_170d89481a616ae977b3_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu": "08e4640bf2f3e3ea36e368e7cec47a4fc11d5cafa44a2112a233abdd275203d2", + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu": "560673dded379b80ca6d345af04ae83b6bd23d136a2e6a95874455b3e70bdbf3" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu", + "csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "0eb214a1983c5cb217d4b3070c9a9b5e19ed0904a8761561ecd2e696a496bbb8", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_binding.cu": "df1bd1bf9f3e5016954fe64b98d11c777c5ae5068a621e2889a1217ea1db271c", + "csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_kernel.cu": "ea4f0cbe026f116a994b655d21566e173d0aacdada6c31aa64356e83aaf1835f" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_kernel.cu", + "csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "26c222026c9702ee06368c4a610764250f909f2c675494e5bce6ec22bb6e6011", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu": "32cc3199c7bb5ecde0352b8226ab8cdb5cb2b5f48af126fb3a7c52787e37ae2c", + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu": "b6473f6b5be6f2a3f817796a6c3bd5bdf7c9cc9c4e4943d985f4531a1dd6efb7" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "da31b5fdd6bc76d6f6fbb8b2653206741311db7819e8865673e7507f16d55250", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu": "32cc3199c7bb5ecde0352b8226ab8cdb5cb2b5f48af126fb3a7c52787e37ae2c", + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu": "b6473f6b5be6f2a3f817796a6c3bd5bdf7c9cc9c4e4943d985f4531a1dd6efb7" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "ab1f258559701e2bf4000feff3df8b082df80c251efcd1dbc9b6d61ae84c53a3", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu": "32cc3199c7bb5ecde0352b8226ab8cdb5cb2b5f48af126fb3a7c52787e37ae2c", + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu": "b6473f6b5be6f2a3f817796a6c3bd5bdf7c9cc9c4e4943d985f4531a1dd6efb7" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "06d38d8dd2cca1f5b61d40f45559251d23b82432b890ae584ddff7697ae86a2d", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_1c21c181af94e1635851_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_1c21c181af94e1635851_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu": "6eb2d702e2cfbb37334048f9a6808f9f24ad78202480f73ef4ccfd6f477b152e", + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu": "5bc13ef14deda70fd44253443d1bb7f55df975bb6cceee7f95eaaaf363face0b" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu", + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "4e20f85d7a0a842fd9c51bee7161257b604a7dd0a71675a6375c20bafc5e91dd", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_1c21c181af94e1635851_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_1c21c181af94e1635851_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu": "6eb2d702e2cfbb37334048f9a6808f9f24ad78202480f73ef4ccfd6f477b152e", + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu": "5bc13ef14deda70fd44253443d1bb7f55df975bb6cceee7f95eaaaf363face0b" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu", + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "2c670612797ea03ea66d51a94c8622ee1db5aec11bf53852fd66a61e3bab0c9c", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_1c21c181af94e1635851_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_1c21c181af94e1635851_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu": "6eb2d702e2cfbb37334048f9a6808f9f24ad78202480f73ef4ccfd6f477b152e", + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu": "5bc13ef14deda70fd44253443d1bb7f55df975bb6cceee7f95eaaaf363face0b" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu", + "csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "0731b332211793ef0d9236ece823140208a06a3ee1ce65cdb357fca2ffc03b37", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_2675cc1cb4efc0937378_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu": "067ab8d02006acb0c128e45b0d3c44047b9fbb6b8d7873b112d04f3e1776ccc8", + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu": "cd193a070b12029c714e974bcaa296c4c0b0488d91834e10963fedc4678ae93f" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu", + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "83bc92a15911e03c30e7c375f2851aba7e1817b069b80f2362d763072f65edbf", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_2675cc1cb4efc0937378_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu": "067ab8d02006acb0c128e45b0d3c44047b9fbb6b8d7873b112d04f3e1776ccc8", + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu": "cd193a070b12029c714e974bcaa296c4c0b0488d91834e10963fedc4678ae93f" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu", + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "b63c8f146586d95c590691580af70caafddce565ac358df71a17012f0f91beb0", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_2675cc1cb4efc0937378_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu": "067ab8d02006acb0c128e45b0d3c44047b9fbb6b8d7873b112d04f3e1776ccc8", + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu": "cd193a070b12029c714e974bcaa296c4c0b0488d91834e10963fedc4678ae93f" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu", + "csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "46f88a0376c3a26d6797f0d0640d8471f78d7f395340f6558353e3fb5a08eef6", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_2d91e5f926238453c78c_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_2d91e5f926238453c78c_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_binding.cu": "d4afcb562dd79e045c2ba752c8639bef70e90fd4e6fb55a74f3b1a911e57f150", + "csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_kernel.cu": "8370b8440c42b4b56442d5b2edda92354738cdf7b34779b3c932561a12d69cf3" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_kernel.cu", + "csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "9657669fa60f2d633c6ae6a7af5e28fc92d90dce8b4910382ce25c042d30c737", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_2d91e5f926238453c78c_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_2d91e5f926238453c78c_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_binding.cu": "d4afcb562dd79e045c2ba752c8639bef70e90fd4e6fb55a74f3b1a911e57f150", + "csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_kernel.cu": "8370b8440c42b4b56442d5b2edda92354738cdf7b34779b3c932561a12d69cf3" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_kernel.cu", + "csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "283aa3ad075169a30ba7ba5e67d3a8c5b49ee1ec7fc080c95d54d00ca13293af", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu": "511ea932e55beb8b2ad8c8de14b0b5147cd2994ab4768110809f005a909f2b25", + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu": "a52e81ee9ca608632bfcf28f4c83bfc38e9bb0c197ae13084c7c519070703862" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu", + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "8e37d99ca44cde06e226e4a990805bd731059a473e8ce1aff02dc7d4b28febb6", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu": "511ea932e55beb8b2ad8c8de14b0b5147cd2994ab4768110809f005a909f2b25", + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu": "a52e81ee9ca608632bfcf28f4c83bfc38e9bb0c197ae13084c7c519070703862" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu", + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "142cca9f31d770bd9204f3409e9c936b768f3d5a2a609699454666c6dbf008e7", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu": "511ea932e55beb8b2ad8c8de14b0b5147cd2994ab4768110809f005a909f2b25", + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu": "a52e81ee9ca608632bfcf28f4c83bfc38e9bb0c197ae13084c7c519070703862" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu", + "csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "d57f439e3d7926029324d6123d6b71385e92d61a700c36051abb41c38c648b20", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu": "d9fee41a409be99804c2da47a6035874656a7ac6e9ef37f0bfec202747815fd3", + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu": "9828599b822db6e3e59eb21b932456816a8f3a74f2bac4c99fc9071dc5a43392" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu", + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "f6808f1d1ccc141086ae21956e659c967487a4cdca6d3cf1624d5dce2488b9e5", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu": "d9fee41a409be99804c2da47a6035874656a7ac6e9ef37f0bfec202747815fd3", + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu": "9828599b822db6e3e59eb21b932456816a8f3a74f2bac4c99fc9071dc5a43392" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu", + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "ac4ac160a4fba6fca97310b8be86cbf0b980b9e5c10047c542517b406b05cb54", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu": "d9fee41a409be99804c2da47a6035874656a7ac6e9ef37f0bfec202747815fd3", + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu": "9828599b822db6e3e59eb21b932456816a8f3a74f2bac4c99fc9071dc5a43392" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu", + "csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "dda4fa5debec7fdae51f473d99a340d5711833cf9c601b3bec010d08b22b2a1f", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_517710d1f0988cf570a9_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_517710d1f0988cf570a9_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu": "260c99fb803c9439b0e7f3ba802732d4c8eeb320c87aa2703262db0df3d44615", + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu": "58419c4ba4b84c2354e9a0e8d3970f71ec709d757d7ba34bca044a6936b6bf2a" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu", + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "1bb898d56a29029f7fff8e4e55900b243fbc94569e3f0c3fdbe6641efdce7d4e", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_517710d1f0988cf570a9_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_517710d1f0988cf570a9_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu": "260c99fb803c9439b0e7f3ba802732d4c8eeb320c87aa2703262db0df3d44615", + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu": "58419c4ba4b84c2354e9a0e8d3970f71ec709d757d7ba34bca044a6936b6bf2a" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu", + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "bced3e69d2c80598ac05345703f49bc988ff98a84a6c774b2203a9c32c835c1b", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_517710d1f0988cf570a9_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_517710d1f0988cf570a9_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu": "260c99fb803c9439b0e7f3ba802732d4c8eeb320c87aa2703262db0df3d44615", + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu": "58419c4ba4b84c2354e9a0e8d3970f71ec709d757d7ba34bca044a6936b6bf2a" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu", + "csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "8ccac0b98cde9e1b521ab5e2235dca0ed915e4145b31282feee3895f15aa8afa", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_binding.cu": "e00d9d0abdb1132c192a64236f4c30c06777c71f0ae40e7a6e1c564388fbd520", + "csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_kernel.cu": "1e20861d7240cc1606a4431b9a749c01c0c16ef648950703213ea30c50ebf304" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "77f52e46f5c1298994c88eaeaa9ba62119f6d1b198954832e9be84bcaef29be1", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_binding.cu": "e00d9d0abdb1132c192a64236f4c30c06777c71f0ae40e7a6e1c564388fbd520", + "csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_kernel.cu": "1e20861d7240cc1606a4431b9a749c01c0c16ef648950703213ea30c50ebf304" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "e0af865de71a2db56b91988633129ed967b5ec72e188bcf656269599b28159a4", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu": "932a2c903f5475e4ce279e7f0b4bd7247c4fbaf47cbd24276ac6c67663c70cf4", + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu": "e6bd58aca087be8358e9daf1073520abcafb0d3dc268781205a7b06da3d82058" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "2181d42e9003674316e9ed8c5a2ecf75c55ecd4e141a05afb8c7000e3c22fd21", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu": "932a2c903f5475e4ce279e7f0b4bd7247c4fbaf47cbd24276ac6c67663c70cf4", + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu": "e6bd58aca087be8358e9daf1073520abcafb0d3dc268781205a7b06da3d82058" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "56d15c3792d1e1b977eaed0eb31cde0850eb8389ae7acfa9c3381d76cde0c30d", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu": "932a2c903f5475e4ce279e7f0b4bd7247c4fbaf47cbd24276ac6c67663c70cf4", + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu": "e6bd58aca087be8358e9daf1073520abcafb0d3dc268781205a7b06da3d82058" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "8e297251f7b1acf865ada0ce212233b166d2e7694b4098245d5d210cef3fcb1e", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_binding.cu": "702586accf2a02a76ff4545f032ff55c52be267cc49e8a5dc8f2d0b59dec9549", + "csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_kernel.cu": "a68547515c30ab67f64534485c7ebec9d914cdcb5f9e4a8d86b53cfe42bc7adb" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "84ab4a570e7c3bf2099145a22882ce28b86da724540078beb94dd4eadffd4d53", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu": "8e915af6d99cfa1d2f84ce59514ed693546e6cae06a8def4bbb1291f1fd42a9d", + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu": "ff9626cd8e86f2ec2c14bfbe6e270400d0b48de31d8ca07cc3c0801d721be7b6" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "9f3b4a003e95870747c6203edd6e2a1fd5dc7059968b10a0dcbe017d3e295a0b", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu": "8e915af6d99cfa1d2f84ce59514ed693546e6cae06a8def4bbb1291f1fd42a9d", + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu": "ff9626cd8e86f2ec2c14bfbe6e270400d0b48de31d8ca07cc3c0801d721be7b6" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "fa8ff764c3c8183dab00424d91c9a9c957cf2e73621f8c13f0c55f2c71936d03", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu": "8e915af6d99cfa1d2f84ce59514ed693546e6cae06a8def4bbb1291f1fd42a9d", + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu": "ff9626cd8e86f2ec2c14bfbe6e270400d0b48de31d8ca07cc3c0801d721be7b6" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu", + "csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "34b3426d38733e0d11e088a9270f964aa5634894eae9a716bffbb81b209f431a", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_843c010db78c79107fd3_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_843c010db78c79107fd3_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_binding.cu": "1ee239829a6dc3fefcdbaededa5993eaa2b3b41256f6a4e9311b01b365a5f819", + "csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_kernel.cu": "dd5b876c00cfab29b5c956028d5f73a41ff5bd1bb67b6e615d4d2b6f22a91374" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_kernel.cu", + "csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "9eeccfc0df8141b2319652d72bda90730a755ae1999f100f63120be6e14442e8", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_binding.cu": "add86dc861723846bef84133814205176a6e7678958b909db1ed535835a21c7c", + "csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_kernel.cu": "ffc0231eca46119f42189e3249005cc71f691bdf8b9c268e062a0efb83a65edb" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_kernel.cu", + "csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "0af47af6d405bd62598ef30be916a50210b24ec02a3168631a0c0934aa76a034", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu": "e355692ffa71f51b206a9aaf1fec11c7a7efa68381713938d06618a65a60d88d", + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu": "6917a5ea38ee17127d562169852373e23fe0080304cf73fadb623e00df2af376" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu", + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "68961d53d8519d89fbbdbd1617173b2972766a97af1bc73a73846b27ba641a32", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu": "e355692ffa71f51b206a9aaf1fec11c7a7efa68381713938d06618a65a60d88d", + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu": "6917a5ea38ee17127d562169852373e23fe0080304cf73fadb623e00df2af376" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu", + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "c4469361fc53fbee928c49b083e41234ee5b389d1492e6035ecf701c42282249", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu": "e355692ffa71f51b206a9aaf1fec11c7a7efa68381713938d06618a65a60d88d", + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu": "6917a5ea38ee17127d562169852373e23fe0080304cf73fadb623e00df2af376" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu", + "csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "f8eae7082863f1b4e377e995231c147e3791be82ee4cfb7eb6a80bed37daf8f1", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_binding.cu": "1fbe56710ed8c7c09e3a4f2305586f577e46bece3c81abdb82a4ef3626e3b8f7", + "csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_kernel.cu": "87a6094b0fc87c2906271f55ee8d8edf4c13e1ce528c0c0b610821f4268c914b" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_kernel.cu", + "csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "a29a8596bf2427e3493d0b8f5173fc591232fdac41143233bf1d3eacce1ac2e7", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_9a07bcaa0223632640ee_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_binding.cu": "a6e0acae8db742cab1f96aa091d24cef12a3c70197e99631538e128b67bb5cf2", + "csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_kernel.cu": "b22e7bd6dd08d167af0d99fd5acba6fcd7dc4a78db1c5519ef9736dcb9433849" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_kernel.cu", + "csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "ad197a06b8cc8c05b73e8740b6de92123c7c58ba9f8f870fb70d96cdb2dd7560", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_binding.cu": "cfff6e32413096e06e0bb4c69e8f841686a3ade9e3060175103dbc80fd383875", + "csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_kernel.cu": "dab9606be3a296333dcc3d040a002de51c87e1b2bbcee8245125d434bf2c34f0" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_kernel.cu", + "csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "e0b4ecee7c50cf1fdd5d67aa10fe4697aaba54b0807d8aca23edadff98d801f3", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_binding.cu": "cfff6e32413096e06e0bb4c69e8f841686a3ade9e3060175103dbc80fd383875", + "csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_kernel.cu": "dab9606be3a296333dcc3d040a002de51c87e1b2bbcee8245125d434bf2c34f0" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_kernel.cu", + "csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "73da8c7ca4b0cb0224d92b93c95914c2b246cf7e74f6e38f184198dd3305a25f", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu": "92de7e88653420496388b3ca085143a6a082f287c6c91606c9be71731d8cd6e8", + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu": "5121b638b113c3b4fabce60d1730aec32ad438ff9843527a6e7e79c05acbed58" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu", + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "3f0d16b5d527ce2b5cca710da1b5a2d981af3d2b9b67ee623ed9d600d063e753", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu": "92de7e88653420496388b3ca085143a6a082f287c6c91606c9be71731d8cd6e8", + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu": "5121b638b113c3b4fabce60d1730aec32ad438ff9843527a6e7e79c05acbed58" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu", + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "7e3984dff60583cef2ef3b0fbd3b31ec8a1cdccc7a4bad8a67025a8e540d613c", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu": "92de7e88653420496388b3ca085143a6a082f287c6c91606c9be71731d8cd6e8", + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu": "5121b638b113c3b4fabce60d1730aec32ad438ff9843527a6e7e79c05acbed58" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu", + "csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "eb140906e4820238476723bd89bf635db61965bc43b7ddaf74e4ee57eb15b281", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_binding.cu": "639cc71fdf7f363fa9817da17dabf28981c587a1c09e5c6350a232db9994bb13", + "csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_kernel.cu": "a89860c6ba8bb9b23133b4f16b3b669d8e82f221520f708365a1598245801681" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_kernel.cu", + "csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "c96c05a1a12b36c2068515c42fde934a756efcd2e43bc6ba401ca8fd25453bde", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_binding.cu": "639cc71fdf7f363fa9817da17dabf28981c587a1c09e5c6350a232db9994bb13", + "csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_kernel.cu": "a89860c6ba8bb9b23133b4f16b3b669d8e82f221520f708365a1598245801681" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_kernel.cu", + "csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "8f7744a95d650be1186d290319653ef0f766d36d9fbbe0a2716ae110b9707d90", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_a681165b3e81204c73ff_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_a681165b3e81204c73ff_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu": "d622dc9614014fe6c31727576eb3ebfb9137ecc7842b69002bf82bba1b1017b6", + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu": "1c423cf1a3287256090b9239e0747cdd98e8bb95ed2167606e2dffc50357106b" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu", + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "b3e253c664d98c069a6aabb158df459c190f00e260d1d9eeb0b244784330cef5", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_a681165b3e81204c73ff_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_a681165b3e81204c73ff_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu": "d622dc9614014fe6c31727576eb3ebfb9137ecc7842b69002bf82bba1b1017b6", + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu": "1c423cf1a3287256090b9239e0747cdd98e8bb95ed2167606e2dffc50357106b" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu", + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "43c4965cf3413c756b1d417b0e99fcc4e4a48a893df5f9c12a14dc05dac86701", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_a681165b3e81204c73ff_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_a681165b3e81204c73ff_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu": "d622dc9614014fe6c31727576eb3ebfb9137ecc7842b69002bf82bba1b1017b6", + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu": "1c423cf1a3287256090b9239e0747cdd98e8bb95ed2167606e2dffc50357106b" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu", + "csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "070c5d08d6462150b3fa923a98076e16519ec5561796b0f7f35e36464990afa5", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_aeff6106e2822f313a0c_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_binding.cu": "dc2d13ae557d050d2f54d1cfc158226d903c4a03c2c4fc0e344a59cecc670fc3", + "csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_kernel.cu": "2c8dade5702621228b0cf1cf506a7db5cbc9eab2381867632261f761924acd10" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_kernel.cu", + "csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "0cb162ea27c314b02fd4fdfdff8e17503ddfce02c043fe1fbc1e7f56fd97723f", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_binding.cu": "eacf0f45dda44978115822db6b2c4a97e78a33b86b61df47b912eb0f7075cdd0", + "csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_kernel.cu": "3fdf90d9f85d7f1f11559569b2f99e7cb27968fd0c17269e73c035cd35136bdb" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_kernel.cu", + "csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "5d4e71d81a1eae9ec5c8c1d7bcd867ca6559da710a3bbc8b7933f7cd7be25aa0", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_binding.cu": "845afa8e438a7fe05c607e0ed405ba36134820d50dade578aafef40d7a5719ad", + "csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_kernel.cu": "c46a4b66cba9dd5a1264fc1af375aeb2b7d149ff1e1ac8fdc99af35c6304975e" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_kernel.cu", + "csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "30f96cd6dd10a788890974f06a85a6ceb270139cbc3d23867cb9e42ea8d48499", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_binding.cu": "845afa8e438a7fe05c607e0ed405ba36134820d50dade578aafef40d7a5719ad", + "csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_kernel.cu": "c46a4b66cba9dd5a1264fc1af375aeb2b7d149ff1e1ac8fdc99af35c6304975e" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_kernel.cu", + "csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "3d273beac5c3a250607c768e1b72133cf9b7e3ee2293cc60100f2ee49fb69d71", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_binding.cu": "ef523f12ee631b1a65091693d2775a36b7be67887a4a321419f8ed57379d496b", + "csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_kernel.cu": "3136b927b7f5eb03c975077847fac73204a98fc2f0230306a29ae58e79f2cebd" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_kernel.cu", + "csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "94edfec807a2624083077221fdb70fa23e8655c45ec008b2a6588171f830241e", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_ce49146b9fef24adb47e_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu": "754035f716a296e48ae411623707f364101dfaf694cbbc22b3517b6777c61ade", + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu": "633e77a585bd2aa52a2280800aeda6b979928f88125aee1b7d3af02b10656398" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu", + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "81431def6d56e1546a29f3a4dbbeb84f74988b0c4c6c4305c8551de2afa92e13", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_ce49146b9fef24adb47e_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu": "754035f716a296e48ae411623707f364101dfaf694cbbc22b3517b6777c61ade", + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu": "633e77a585bd2aa52a2280800aeda6b979928f88125aee1b7d3af02b10656398" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu", + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "4a20879962968f05fe3c9afdc2850308291a589d81ece370c652fc683812a69b", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_ce49146b9fef24adb47e_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu": "754035f716a296e48ae411623707f364101dfaf694cbbc22b3517b6777c61ade", + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu": "633e77a585bd2aa52a2280800aeda6b979928f88125aee1b7d3af02b10656398" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu", + "csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "af5d45fac176cd2297fa5cba56a60c97d8586923e894285beb4dd1683caff056", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_binding.cu": "4a226db4e7b7dc3656e7140403a363d07e940834e4aea9455396f81b9f965da1", + "csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_kernel.cu": "ab0d63386013013d24dc76c210f2ad61eb02d5379891d8b3fb53e08cbc245520" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_kernel.cu", + "csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "22ed821d0cddbeb98baf02f08fd7f52be8290b55ee3cb49a6a204d7cd022f361", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_binding.cu": "4a226db4e7b7dc3656e7140403a363d07e940834e4aea9455396f81b9f965da1", + "csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_kernel.cu": "ab0d63386013013d24dc76c210f2ad61eb02d5379891d8b3fb53e08cbc245520" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_kernel.cu", + "csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "b494eadc9e03346d90a760e1652bdc4561af2153738199c5992c474d2781e6ff", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_d829c95e1d610bc14b23_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_binding.cu": "bc63951666f04dbab400b3807f26a479634e975fe312fcc6ad81028dd2dcfbec", + "csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_kernel.cu": "fe2d7196d071255e4eacbeca508367de048eaf2064afdbc2ba21aa63d5939f4f" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_kernel.cu", + "csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "421981b0c04a3ad005810805a966244a5bc28ad7be100319c4593700d23308d0", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu": "f46a7f40d528dded87c5538df2f91e39b5bc81e8e58958a04798a378671553f0", + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu": "e8c4b8c9f2b529df3461d8b17ef478e60bb007bffbcc7309fbce01a0275ae0b5" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu", + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "dd235c7fb5da9c309fc27ee2202a16759ec18699fd4699dbac50b0b166806d2a", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu": "f46a7f40d528dded87c5538df2f91e39b5bc81e8e58958a04798a378671553f0", + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu": "e8c4b8c9f2b529df3461d8b17ef478e60bb007bffbcc7309fbce01a0275ae0b5" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu", + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "bf80a70806b8a1b93f4fe0353a1bf54ffd46da46233ade88607c5e731baa747a", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu": "f46a7f40d528dded87c5538df2f91e39b5bc81e8e58958a04798a378671553f0", + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu": "e8c4b8c9f2b529df3461d8b17ef478e60bb007bffbcc7309fbce01a0275ae0b5" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu", + "csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "b3170e88e02bc576165abde4679008659a1ee2257da7f9eaef40087601908fae", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_f1545adfd2bb28044419_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_f1545adfd2bb28044419_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_binding.cu": "8563cadb9b9988a145e47bcc369df0257e51c578b41e23c96fb5a0baf066b4c0", + "csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_kernel.cu": "4620f5265725358f226a08d7b7195abab484b08e445ad14dcbe6d2e0ad048cd8" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_kernel.cu", + "csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "dfe135183e395ac6ffe797125408d559e3969a5610aeb3289a024e256e4ddb73", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_causal_conv1d_det_train_f5497330bfd286211643_sm_100a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_f5497330bfd286211643_sm_100a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu": "b1fdaaec4bc4a5abe23e3a0a0e12fc3101e293ea3b08bce6a9c510c7b26ec742", + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu": "e64e76ce6930fda5c1ab92b7c915b0f03198b05398ef2f6ca8962d133588ee0f" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu", + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "8ceea16ddb80df75bcdac950c4a1f14aa08f2d7ba4b353bcab696528c2377231", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_causal_conv1d_det_train_f5497330bfd286211643_sm_103a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_f5497330bfd286211643_sm_103a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu": "b1fdaaec4bc4a5abe23e3a0a0e12fc3101e293ea3b08bce6a9c510c7b26ec742", + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu": "e64e76ce6930fda5c1ab92b7c915b0f03198b05398ef2f6ca8962d133588ee0f" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu", + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "6e27397675c01271529040bd3f207e4bbd65dae21f58d7cd0c1d590573732083", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_causal_conv1d_det_train_f5497330bfd286211643_sm_107a", + "nvcc_flags": [ + "--prec-div=true", + "--prec-sqrt=true", + "--ftz=false" + ], + "nvcc_flags_if": [], + "path": "aot/cake_causal_conv1d_det_train_f5497330bfd286211643_sm_107a.so", + "sha256": { + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu": "b1fdaaec4bc4a5abe23e3a0a0e12fc3101e293ea3b08bce6a9c510c7b26ec742", + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu": "e64e76ce6930fda5c1ab92b7c915b0f03198b05398ef2f6ca8962d133588ee0f" + }, + "sources": [ + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu", + "csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu" + ] + } + ], + "package": "flashinfer.cake_causal_conv1d_det_train_backend", + "schema": "cake-build/1" +} diff --git a/flashinfer/cake_causal_conv1d_det_train_backend/cake_build.py b/flashinfer/cake_causal_conv1d_det_train_backend/cake_build.py new file mode 100644 index 00000000000..00756137eb1 --- /dev/null +++ b/flashinfer/cake_causal_conv1d_det_train_backend/cake_build.py @@ -0,0 +1,324 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Library loader of one Cake-generated FlashInfer package. + +``build.json`` (schema ``cake-build/1``) next to this file is the only record +of what the package compiles: one entry per ``[_]_`` +library with its package-relative CUDA sources, nvcc flags, toolkit-conditional +flags and digests. A library is served from ``aot/.so`` when a builder +(a vendoring consumer's build rules, or FlashInfer's +``scripts/check_cake_build_manifests.py``) has compiled it from the manifest, +and otherwise through FlashInfer's JIT from the same sources and flags. + +The prebuilt path never imports ``flashinfer``, so the package stays usable in +a process without a FlashInfer installation. A builder that fills ``aot/`` +also writes the sidecar ``aot/build_info.json`` (schema ``cake-build-info/1``: +``{"schema", "cuda": "13.0.88", "cudart_version": 13000, "arches": [...]}``, +``arches`` naming every architecture the builder compiled). Every prebuilt +library exports ``build_cuda_version()``, the ``CUDART_VERSION`` it was +compiled against, and :func:`load_library` refuses a library whose value +disagrees with the sidecar. + +Public surface (``__all__``): :func:`manifest`, :func:`libraries_by_name`, +:func:`prebuilt_path`, :func:`library_path`, :func:`load_library`, +:func:`jit_spec`, :func:`jit_specs` and :func:`record_cuda_version`. + +The Cake exporter writes this file byte-identically into every package from +its loader template; do not edit it in place. +""" + +from __future__ import annotations + +import functools +import hashlib +import json +import re +from pathlib import Path +from typing import Any + +__all__ = [ + "jit_spec", + "jit_specs", + "libraries_by_name", + "library_path", + "load_library", + "manifest", + "prebuilt_path", + "record_cuda_version", +] + +_PACKAGE = Path(__file__).parent +_MANIFEST_SCHEMA = "cake-build/1" +_BUILD_INFO_SCHEMA = "cake-build-info/1" +_MANIFEST_PATH = _PACKAGE / "build.json" +_BUILD_INFO_PATH = _PACKAGE / "aot" / "build_info.json" +_CUDA_CONSTRAINT = re.compile(r"(==|>=)([0-9]+)\.([0-9]+)") + + +@functools.cache +def manifest() -> dict[str, Any]: + """The parsed ``build.json`` of this package (read on first use, not at import). + + One cached document shared by every caller: read it, do not modify it. + """ + try: + text = _MANIFEST_PATH.read_text(encoding="utf-8") + except FileNotFoundError as exc: + raise RuntimeError( + f"{_MANIFEST_PATH} is missing: the package was not exported with its {_MANIFEST_SCHEMA} manifest", + ) from exc + document = json.loads(text) + schema = document.get("schema") + if schema != _MANIFEST_SCHEMA: + raise RuntimeError( + f"{_MANIFEST_PATH} has schema {schema!r}; this loader reads {_MANIFEST_SCHEMA!r}", + ) + return document + + +@functools.cache +def libraries_by_name() -> dict[str, dict[str, Any]]: + """``name -> entry`` over the ``libraries`` of ``build.json``.""" + index: dict[str, dict[str, Any]] = {} + for entry in manifest()["libraries"]: + if entry["name"] in index: + raise RuntimeError(f"{_MANIFEST_PATH} lists {entry['name']!r} twice") + index[entry["name"]] = entry + return index + + +def _entry(name: str) -> dict[str, Any]: + try: + return libraries_by_name()[name] + except KeyError: + raise KeyError(f"{name!r} is not a library of {_MANIFEST_PATH}") from None + + +def prebuilt_path(name: str) -> Path: + """Where a builder places the compiled library of ``name`` (``aot/.so``).""" + return _PACKAGE / _entry(name)["path"] + + +def _cuda_constraint_matches(constraint: str, cuda: tuple[int, int]) -> bool: + """Does toolkit ``cuda`` (major, minor) satisfy ``"==X.Y"`` or ``">=X.Y"``?""" + match = _CUDA_CONSTRAINT.fullmatch(constraint) + if match is None: + raise ValueError( + f"unsupported CUDA constraint {constraint!r}; expected '==X.Y' or '>=X.Y'", + ) + operator, major, minor = match.groups() + version = (int(major), int(minor)) + if operator == "==": + return tuple(cuda) == version + return tuple(cuda) >= version + + +def _conditional_flags(items: list[Any], cuda: tuple[int, int]) -> list[str]: + """The flags of every ``nvcc_flags_if`` item whose ``cuda`` constraint matches.""" + flags: list[str] = [] + for item in items: + if _cuda_constraint_matches(item["cuda"], cuda): + flags.extend(item["flags"]) + return flags + + +def _jit_toolchain(name: str) -> tuple[Any, Any]: + """``(gen_jit_spec, get_cuda_version)`` of the local FlashInfer, or the error.""" + try: + from flashinfer.jit.core import gen_jit_spec + from flashinfer.jit.cpp_ext import get_cuda_version + except ImportError as exc: + raise RuntimeError( + f"{name} is not prebuilt ({prebuilt_path(name)}) and flashinfer's JIT is not available", + ) from exc + return gen_jit_spec, get_cuda_version + + +def _local_cuda_version(get_cuda_version: Any) -> tuple[int, int]: + version = get_cuda_version() + return (int(version.major), int(version.minor)) + + +def _version(text: str) -> tuple[int, int]: + """``(major, minor)`` of a dotted version such as ``"13.0"`` or ``"13.0.88"``.""" + parts = str(text).split(".") + if len(parts) < 2 or not all(part.isdigit() for part in parts[:2]): + raise RuntimeError(f"malformed CUDA version {text!r}") + return (int(parts[0]), int(parts[1])) + + +def _toolkit_builds(entry: dict[str, Any], cuda: tuple[int, int]) -> bool: + """Does toolkit ``cuda`` reach the entry's ``min_cuda`` floor?""" + return cuda >= _version(entry["min_cuda"]) + + +def _verified_sources(entry: dict[str, Any]) -> list[Path]: + """The entry's sources as package paths, each checked against its digest.""" + sources: list[Path] = [] + for relative in entry["sources"]: + path = _PACKAGE / relative + digest = hashlib.sha256(path.read_bytes()).hexdigest() + expected = entry["sha256"][relative] + if digest != expected: + raise RuntimeError( + f"{path} does not match its build.json digest for {entry['name']}: {digest} != {expected}", + ) + sources.append(path) + return sources + + +def jit_spec(name: str) -> Any: + """FlashInfer's ``JitSpecNvcc`` of ``name``, built from the manifest entry. + + Sources are digest-verified here and only here, so a modified or partially + delivered source tree fails closed before anything compiles. The spec name + carries the closure digest: a changed closure never reuses a stale JIT cache + entry. Flags are ``arch_flags[arch]`` + the entry's ``nvcc_flags`` + the + ``nvcc_flags_if`` items matching the local toolkit; FlashInfer's JIT adds + its own defaults (``-O3``, ``-fPIC``, ...), so ``common.nvcc_flags`` are not + repeated here. A toolkit below the entry's ``min_cuda`` raises. + """ + entry = _entry(name) + gen_jit_spec, get_cuda_version = _jit_toolchain(name) + document = manifest() + cuda = _local_cuda_version(get_cuda_version) + if not _toolkit_builds(entry, cuda): + raise RuntimeError( + f"{name} needs CUDA {entry['min_cuda']} or newer; the local toolkit is {cuda[0]}.{cuda[1]}", + ) + flags = [ + *document["arch_flags"][entry["arch"]], + *entry["nvcc_flags"], + *_conditional_flags(entry["nvcc_flags_if"], cuda), + ] + include_dirs = [_PACKAGE / d for d in document["common"]["include_dirs"]] + return gen_jit_spec( + f"{name}_{entry['closure_sha256'][:20]}", + _verified_sources(entry), + extra_cuda_cflags=flags, + extra_ldflags=list(document["common"]["link_flags"]), + extra_include_paths=include_dirs, + use_fast_math=False, + ) + + +def _names_of(arch: str) -> list[str]: + index = libraries_by_name() + return [name for name in index if index[name]["arch"] == arch] + + +def jit_specs(arch: str) -> list[Any]: + """JIT specifications of the ``arch`` libraries still to build locally. + + The entries without a prebuilt file that the local toolkit can build: the + list feeds ``flashinfer.jit.build_jit_specs`` warm-ups, so an entry whose + ``min_cuda`` exceeds the local toolkit is left out (the manifest check + script's "skip min_cuda > toolkit" rule), while :func:`jit_spec` on that + entry by name still raises. When every ``arch`` library is prebuilt the + list is empty and ``flashinfer`` is not imported. + """ + pending = [name for name in _names_of(arch) if not prebuilt_path(name).is_file()] + if not pending: + return [] + _, get_cuda_version = _jit_toolchain(pending[0]) + cuda = _local_cuda_version(get_cuda_version) + return [jit_spec(name) for name in pending if _toolkit_builds(_entry(name), cuda)] + + +def library_path(name: str) -> Path: + """The prebuilt library of ``name`` when present, else its JIT cache path.""" + path = prebuilt_path(name) + if path.is_file(): + return path + return jit_spec(name).get_library_path() + + +@functools.cache +def _build_info() -> dict[str, Any]: + """The builder-written ``aot/build_info.json`` sidecar (prebuilt path only).""" + try: + text = _BUILD_INFO_PATH.read_text(encoding="utf-8") + except FileNotFoundError as exc: + raise RuntimeError( + f"{_BUILD_INFO_PATH} is missing although prebuilt libraries exist under {_PACKAGE / 'aot'}: the builder must write the {_BUILD_INFO_SCHEMA} sidecar", + ) from exc + info = json.loads(text) + schema = info.get("schema") + if schema != _BUILD_INFO_SCHEMA: + raise RuntimeError( + f"{_BUILD_INFO_PATH} has schema {schema!r}; this loader reads {_BUILD_INFO_SCHEMA!r}", + ) + return info + + +def _prebuilt_info(path: Path, arch: str) -> dict[str, Any]: + """The sidecar, which must record a build of ``arch`` for the prebuilt ``path``.""" + info = _build_info() + if arch not in info["arches"]: + raise RuntimeError( + f"{path} exists but {_BUILD_INFO_PATH} records no {arch} build (arches {info['arches']})", + ) + return info + + +def load_library(name: str) -> Any: + """The loaded tvm-ffi module of ``name``: the prebuilt library, else a JIT build.""" + entry = _entry(name) + path = _PACKAGE / entry["path"] + if not path.is_file(): + return jit_spec(name).build_and_load() + info = _prebuilt_info(path, entry["arch"]) + import tvm_ffi + + module = tvm_ffi.load_module(str(path)) + built = int(module.build_cuda_version()) + if built != int(info["cudart_version"]): + raise RuntimeError( + f"{path} was compiled against CUDART {built} but {_BUILD_INFO_PATH} records {info['cudart_version']}", + ) + return module + + +def record_cuda_version(arch: str) -> tuple[int, int]: + """``(major, minor)`` of the toolkit that builds the ``arch`` libraries. + + Prebuilt path, taken as soon as any library of the manifest exists under + ``aot/``: the sidecar ``aot/build_info.json`` is mandatory (missing -> + ``RuntimeError``; no local nvcc is consulted) and its ``cuda`` is the + answer whenever its ``arches`` list ``arch`` -- the builder compiled that + architecture, whether or not every one of its libraries is present. An + ``arch`` library under ``aot/`` that the sidecar does not list raises, as + :func:`load_library` does. + + JIT path, when nothing is prebuilt or the builder did not compile + ``arch``: the local toolkit of FlashInfer's JIT (``RuntimeError`` in a + process without a FlashInfer installation). Call it for the architecture + you are about to load. + """ + names = _names_of(arch) + if not names: + raise KeyError(f"{arch!r} has no libraries in {_MANIFEST_PATH}") + index = libraries_by_name() + prebuilt = [name for name in index if (_PACKAGE / index[name]["path"]).is_file()] + if prebuilt: + of_arch = [name for name in prebuilt if index[name]["arch"] == arch] + if of_arch: + return _version(_prebuilt_info(prebuilt_path(of_arch[0]), arch)["cuda"]) + info = _build_info() + if arch in info["arches"]: + return _version(info["cuda"]) + _, get_cuda_version = _jit_toolchain(names[0]) + return _local_cuda_version(get_cuda_version) diff --git a/flashinfer/cake_causal_conv1d_det_train_backend/cake_jit.py b/flashinfer/cake_causal_conv1d_det_train_backend/cake_jit.py new file mode 100644 index 00000000000..34004d19e39 --- /dev/null +++ b/flashinfer/cake_causal_conv1d_det_train_backend/cake_jit.py @@ -0,0 +1,3670 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Host plan and loader glue of the Cake packed width-4 depthwise causal +convolution + SiLU training kernels (SM100 / SM103 / SM107). + +Every module below is one mechanically exported production build: a CUDA +device translation unit plus a tvm-ffi binding compiled together for exactly +one architecture. ``MODULES`` and ``ROUTES`` carry the host plan only: the +kernel's ``library`` stem, the architectures it is delivered for, its launch +record, grid rule, workspace rule and the kernel-parameter -> logical-argument +map. What to compile -- sources, digests, nvcc flags -- is recorded once in +``build.json`` next to this file; :mod:`.cake_build` serves the library +``_`` from a prebuilt ``aot/`` file when a builder compiled it +and otherwise through FlashInfer's JIT from the same sources and flags, +verifying every source digest before the first build. + +A route lists every program one public call may select. The forward routes +produce ``y`` from ``x``, ``weight`` and ``cu_seqlens``: ``fwd`` (FP32 taps), +``fwd_dt`` (BF16 taps) and ``fwd_bias`` (a per-channel bias, FP32 or BF16 taps +and bias). The backward routes produce ``dx`` and the deterministic +``dweight`` (and ``dbias``): ``bwd`` (FP32 taps and +gradient, SiLU' at the BF16-rounded preactivation), ``bwd_dt`` (BF16 taps and +/ or a BF16 ``dweight``), ``bwd_bias``, ``bwd_fp32g`` (SiLU' at the unrounded +FP32 preactivation for the weight gradient) and ``bwd_bias_fp32g``; +:func:`forward_route` / :func:`backward_route` map a call's arguments to its +route and the ``flags`` bit mask (``FLAG_*``) tells a route's programs which +of ``w``, ``dw``, ``b``, ``db`` are BF16. A route, forward or backward, may +carry several programs of its kernel and selects exactly one per call, +chosen by the channel count (and, for one record, the token count): a +record's ``select_rule`` (``channels_at_least``: ``C >= channels[arch]``; +``channels_tokens_at_least``: that and ``S >= tokens[arch]``) names the +calls it owns, of several owners the one of the highest channel threshold +runs, the record without a rule is the route's fallback, and +:func:`selected_modules` applies that rule (a pure function of the +architecture, the route, ``C`` and ``S``; the host never reads the data). +Token and channel counts are runtime scalars, so a changed problem size never +recompiles anything; building every library of the architecture +(:func:`specs_for_arch`) before a CUDA Graph capture keeps the JIT out of the +capture. + +A forward call is one ordinary launch of the program its route selects, a +dynamic token-chunk forward (round 9): the co-resident CTAs of a channel tile +take the tile's token chunks in order from a device counter, so the CTAs that +finish first take the tail instead of an equal share fixed at launch. Each +forward route ships the program of one eight-token stage per sub-tile and +chunk for every channel count on sm_100a and sm_103a and, on sm_107a, below +a per-route channel count (the fallback); from that count (C = 256 on +``fwd``, 512 on ``fwd_dt``, 1024 on ``fwd_bias``) sm_107a runs the program +of two stages per chunk. The counters +are the forward workspace ``work`` (``workspace_rule`` ``tile_counters``: +:func:`counter_words`, two int32 words per channel tile), zero before the first +launch; the last CTA of each tile resets them, so the buffer is zero after +every launch and every CUDA Graph replay and serves every later launch of one +stream. Two launches that may run at once (two streams, two graphs replayed +concurrently) need two buffers. + +A backward call is two ordinary launches on one stream. The selected program +(persistent, ``grid_mult x ctas_per_sm x SM count`` CTAs) writes ``dx`` and +one FP32 partial row of ``dweight`` (and, on the bias routes, of ``dbias``) +per CTA into the workspace and signals its dependents; then the route's reduce +program (``record["reduce"]``), launched programmatically dependent on it +(``launch.use_pdl``: its binding sets the programmatic-stream-serialization +attribute, and the kernel waits for the main grid before its first load), sums +the rows per channel in a fixed order and stores ``dw`` (and ``db``) in the +dtypes ``flags`` names. Until round 6 the backward was one cooperative launch +whose grid-wide barrier preceded the same reduction in-kernel. A cooperative +grid cannot start while any kernel is resident on another stream -- NCCL +collectives are, for most of a training backward (FlashInfer issue #7038) -- +whereas an ordinary launch whose CTAs find no free slot runs them in a later +wave. ``grid_mult`` (``grid_rule["grid_mult"][arch]``) over-decomposes the +token axis so that such a CTA trails by ``1 / grid_mult`` of a stream rather +than a whole one; the partial rows grow with it. The reduction order is +unchanged, so ``dweight`` is deterministic for a given ``(S, C, device)`` and +route and, at ``grid_mult`` 1, bitwise what the cooperative form produced. +Each backward route has its own reduce program: the default one, and builds +of it with ``dtype_flags`` (the ``dw`` store form), ``bias`` (the fifth, dbias, +slab and ``db``) and ``fp64`` (the ``_fp32g`` routes' FP64 sums). + +The host plan is the three functions :func:`plan` (the kernel scalars, with +the channel tile rule :func:`channel_split` of the record's role: the forward +and the backward rules differ where no multiple of 32 divides the channel +vectors), :func:`launch_grid` (the grid: +``co_resident`` records launch ``ctas_per_sm x SM count`` CTAs, ``occupancy`` +records ``grid_mult x ctas_per_sm x SM count``, less the CTAs +:func:`main_grid` drops to whole channel-tile ranks on the rows it names, +``channels_per_cta`` records ``ceil(C / channels_per_cta)``), +:func:`partial_rows` (the backward +workspace rows) and :func:`counter_words` (the forward's ``work`` words); +everything else derives from them. The public workspace is +sized by the backward program of the largest ``grid_mult x ctas_per_sm x +taps`` over every backward route (:func:`workspace_record`, the capacity +owner; the bias programs' partials carry a fifth, dbias, slab); each program +writes the leading rows of its own grid and its reduce program sums exactly +those. The forward workspace is sized by the forward program of the most +counter words over every forward route (:func:`forward_workspace_words`). + +The program name ``cake_causal_conv1d_det_train_<20 hex>`` is the first twenty hex +digits of the SHA-256 of the compact JSON ``{"build": , +"module": }`` of the kernel the program was rendered +from; the default reduce program is built with its kernel file's default +keywords and the forward programs with ``fast=True, sched="dynamic", +chunk_stages=1, pair_math=True`` (the dynamic token-chunk schedule of one +eight-token stage per sub-tile and chunk, the f32x2 pair math, the +branch-free in-document stage; ``y`` is bitwise the default build's) plus +the route's keywords, the sm_107a programs the channel count selects with +``chunk_stages=2`` instead (two eight-token stages per sub-tile and chunk); +every backward program carries ``revision=6`` (the +build revision of the split that retired the cooperative form; the exporter +hashes it so the split programs replaced the cooperative ones under new names); +the default route's fallback is built with ``vec=8, copy="tma2d", +stage_pair=True`` (eight channels, 16-byte vectors, per thread; 2-D TMA row +copies; each eight-token copy stage consumed as two four-token halves, the +dweight fold buffer in the stage ring) on sm_100a and sm_103a and with +``min_blocks=2`` (two CTAs per SM) beside those on sm_107a, where for +256 <= C < 512 from S = 262144 tokens the build of ``vec=8, min_blocks=2`` +(four-token copy stages, no ``stage_pair``) runs instead, and the backwards +the ``channels_at_least`` rules select, builds of ``vec=8, copy="tma2d", +stage_tokens=8, min_blocks=2`` (eight-token stages, two CTAs per SM): as is +on sm_100a from C = 512 and on sm_103a from 768, with ``fold_overlap=True, +fast_roll=True`` (the dweight fold buffer over the stage slots, the stage body +as two rolled passes of four tokens) on sm_103a for 512 <= C < 768 and on +sm_107a for 512 <= C < 1024, as is on sm_107a for 1024 <= C < 2048, and +with ``fast_roll=True, roll_prefetch=True`` (each rolled pass loads the +next pass's rows before it computes) on sm_107a from 2048. The programs of +the six other routes are the same bodies built with ``dtype_flags=True`` (the +``flags`` scalar; untyped tap / bias pointers) plus ``bias=True`` for the +bias routes and ``grad_pre="fp32"`` for the ``_fp32g`` routes, their backward +fallbacks without ``stage_pair`` (``vec=8, copy="tma2d"`` on sm_100a and +sm_103a, ``vec=8, min_blocks=2`` on sm_107a); each record's comment names its +build. +""" + +from __future__ import annotations + +import functools +from typing import Any, Sequence + +import tvm_ffi + +from . import cake_build + +# The derived leaves of every record (``kernel_symbol``, ``launch.block``, +# ``launch.cluster``, ``launch.cooperative``, ``launch.dynamic_smem_bytes``, +# ``launch.use_pdl``, ``arg_plan`` and ``logical``) are rewritten by the +# exporter from the rendered program; every other field is host plan. +MODULES: dict[str, dict[str, Any]] = { + # Route ``fwd`` (round 9): the forward's dynamic token-chunk program with the + # f32x2 pair math and the branch-free in-document stage, built with ``fast=True, + # sched="dynamic", chunk_stages=1, pair_math=True``. The CTAs of a channel tile + # take chunks of ``n_sub x 8`` tokens from a device counter instead of equal + # contiguous shares, so the CTAs of the SMs that draw the most bandwidth take + # the tail; ``y`` is bitwise the static program's. The counters are the + # ``work`` buffer (``workspace_rule``: int32, two words per channel tile: the + # next-chunk counter and the done-CTA counter), zero before the first launch; + # the last CTA of each tile resets them, so ``work`` is zero again after every + # launch and every CUDA Graph replay, and one buffer serves every later launch + # of one stream (two launches that may overlap need two buffers). Where a CTA + # holds several sub-tiles, a chunk's three history rows come from the stage's + # slot (the halo rows its copy brings), not from global memory. It is the + # route's fallback (no ``select_rule``): every channel count on sm_100a and + # sm_103a, and C < 256 on sm_107a. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_f5497330bfd286211643": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "y"], + ["buffer", "work"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # The grid is the co-resident CTA count, ``ctas_per_sm x SM count``: + # CTA ``b`` owns channel tile ``b % n_ch_tiles`` and, with the other + # CTAs of that tile, claims the tile's token chunks in order from + # ``work`` (a CTA's chunks rise; its first chunk past ``S`` ends it). + # Every channel tile needs a CTA (the host refuses a ``C`` of more + # tiles than the grid). ``ctas_per_sm`` 4 is the co-residency of the + # program (128 registers at 128 threads, no stack) on B200 (sm_100a), + # GB300 (sm_103a) and R200 (sm_107a). ``geometry`` is the kernel's + # build geometry the host plan derives its scalars from. + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 4, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_f5497330bfd286211643", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_f5497330bfd286211643", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 46208, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_f5497330bfd286211643", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "search_steps": "search_steps", + "w": "w", + "work": "work", + "x": "x", + "y": "y", + }, + "role": "fwd", + "workspace_rule": { + "kind": "tile_counters", + "words_per_tile": 2, + }, + }, + "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce": { + "arches": [ + "sm_100a", + "sm_103a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # An ordinary launch of ``grid_mult x ctas_per_sm x SM count`` CTAs: + # ``ctas_per_sm`` is the measured co-residency, ``grid_mult`` (per + # architecture) splits the token streams that many times finer so a + # CTA another stream's kernel displaces trails by ``1 / grid_mult`` of + # a stream instead of a whole one. Each CTA writes one partial row; + # ``reduce`` sums them. The grid is a function of the device alone. + # The ``stage_pair`` build: each 8-token copy stage is consumed as two + # 4-token halves, the dweight fold buffer in the stage ring. + # ``ctas_per_sm``: 3 CTAs per SM (168 registers at 128 threads, the + # launch bound; ncu launch__occupancy_limit_registers), the occupancy + # query's answer on B200 and GB300. + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 3, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 66560, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "role": "bwd", + # The FP32 ``partial`` workspace: ``[ceil(grid / n_ch_tiles), C, taps]`` + # rows (row ``b // n_ch_tiles`` belongs to CTA ``b``) the kernel fully + # writes and ``reduce`` sums in a fixed order; nothing is zeroed by the + # host. + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + "cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # An ordinary launch of ``grid_mult x ctas_per_sm x SM count`` CTAs: + # ``ctas_per_sm`` is the measured co-residency, ``grid_mult`` (per + # architecture) splits the token streams that many times finer so a + # CTA another stream's kernel displaces trails by ``1 / grid_mult`` of + # a stream instead of a whole one. Each CTA writes one partial row; + # ``reduce`` sums them. The grid is a function of the device alone. + # The ``copy="tma2d", stage_pair`` build of the sm_100a / sm_103a + # fallback at ``min_blocks`` 2: 2-D TMA stage copies, each 8-token + # stage consumed as two 4-token halves, the dweight fold buffer in the + # stage ring. ``ctas_per_sm``: 2 CTAs per SM, the program's + # ``min_blocks`` 2 launch bound (ncu launch__occupancy_limit_registers), + # the occupancy query's answer on R200; the ``min_blocks`` 3 build of + # the other two architectures measured 10 % slower at C = 128 on its + # 212 SMs. + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 66560, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "role": "bwd", + # The FP32 ``partial`` workspace: ``[ceil(grid / n_ch_tiles), C, taps]`` + # rows (row ``b // n_ch_tiles`` belongs to CTA ``b``) the kernel fully + # writes and ``reduce`` sums in a fixed order; nothing is zeroed by the + # host. + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + "cake_causal_conv1d_det_train_9a07bcaa0223632640ee": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # An ordinary launch of ``grid_mult x ctas_per_sm x SM count`` CTAs: + # ``ctas_per_sm`` is the measured co-residency, ``grid_mult`` (per + # architecture) splits the token streams that many times finer so a + # CTA another stream's kernel displaces trails by ``1 / grid_mult`` of + # a stream instead of a whole one. Each CTA writes one partial row; + # ``reduce`` sums them. The grid is a function of the device alone. + # ``ctas_per_sm``: 2 CTAs per SM, the program's ``min_blocks`` 2 launch + # bound (234 registers at 128 threads; ncu + # launch__occupancy_limit_registers), the occupancy query's answer on + # R200; the ``min_blocks`` 3 build of the other two architectures + # measured 10 % slower at C = 128 on its 212 SMs. + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_9a07bcaa0223632640ee", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 50176, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_9a07bcaa0223632640ee", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + }, + "reduce": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "role": "bwd", + # The four-token build of the sm_107a fallback runs for 256 <= C < 512 + # from S = 262144 tokens; the ``stage_pair`` fallback runs every other + # C < 512 (on those rows its gain was not distinguishable from the + # spread between identical copies of one build). + "select_rule": { + "channels": { + "sm_107a": 256, + }, + "kind": "channels_tokens_at_least", + "tokens": { + "sm_107a": 262144, + }, + }, + # The FP32 ``partial`` workspace: ``[ceil(grid / n_ch_tiles), C, taps]`` + # rows (row ``b // n_ch_tiles`` belongs to CTA ``b``) the kernel fully + # writes and ``reduce`` sums in a fixed order; nothing is zeroed by the + # host. + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Rendered by the exporter from the family adapter's scaffold: the backward + # with stage_tokens 8, copy tma2d, min_blocks 2 (2 CTAs per SM); it runs when + # C >= select_rule.channels[arch] and no routed record of a higher threshold + # admits C; below every threshold the fallback record (no select_rule) runs. + "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 82944, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "role": "bwd", + "select_rule": { + "channels": { + "sm_100a": 512, + "sm_103a": 768, + "sm_107a": 1024, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Rendered by the exporter from the family adapter's scaffold: the backward + # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, fold_overlap, + # fast_roll + # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no + # routed record of a higher threshold admits C; below every threshold the + # fallback record (no select_rule) runs. + "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6": { + "arches": [ + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_103a": 1, + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 66560, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "role": "bwd", + "select_rule": { + "channels": { + "sm_103a": 512, + "sm_107a": 512, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Rendered by the exporter from the family adapter's scaffold: the backward + # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, fast_roll, + # roll_prefetch + # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no + # routed record of a higher threshold admits C; below every threshold the + # fallback record (no select_rule) runs. + "cake_causal_conv1d_det_train_aeff6106e2822f313a0c": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_aeff6106e2822f313a0c", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 82944, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_aeff6106e2822f313a0c", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "role": "bwd", + "select_rule": { + "channels": { + "sm_107a": 2048, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Route ``fwd_dt`` (round 9): the dynamic token-chunk forward with + # ``dtype_flags`` (the ``flags`` scalar selects the tap load form: bit 0 set + # reads BF16 taps), built with ``dtype_flags=True, fast=True, + # sched="dynamic", chunk_stages=1, pair_math=True``; ``work`` as the ``fwd`` + # record's. The route's fallback: every channel count on sm_100a and + # sm_103a, and C < 512 on sm_107a. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "y"], + ["buffer", "work"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 4, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 46208, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "search_steps": "search_steps", + "w": "w", + "work": "work", + "x": "x", + "y": "y", + }, + "role": "fwd", + "workspace_rule": { + "kind": "tile_counters", + "words_per_tile": 2, + }, + }, + # Route ``fwd_bias`` (round 9): the dynamic token-chunk forward with a + # per-channel bias ``b`` added after the four taps (FLA's order) and + # ``dtype_flags`` (bit 0: BF16 taps, bit 2: BF16 bias), built with ``bias=True, + # dtype_flags=True, fast=True, sched="dynamic", chunk_stages=1, + # pair_math=True``; ``work`` as the ``fwd`` record's. The route's fallback: + # every channel count on sm_100a and sm_103a, and C < 1024 on sm_107a. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "y"], + ["buffer", "b"], + ["buffer", "work"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 4, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 46208, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb", + "logical": { + "C": "C", + "S": "S", + "b": "b", + "ct_vec": "ct_vec", + "cu": "cu", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "search_steps": "search_steps", + "w": "w", + "work": "work", + "x": "x", + "y": "y", + }, + "role": "fwd", + "workspace_rule": { + "kind": "tile_counters", + "words_per_tile": 2, + }, + }, + # Route ``fwd`` on sm_107a from C = 256 (round 9): the dynamic token-chunk + # forward of two eight-token stages per sub-tile and chunk, built with + # ``fast=True, sched="dynamic", chunk_stages=2, pair_math=True``. On R200 the + # one-stage program loses to the shipped forward from C = 2048 (one sub-tile + # per CTA over several channel tiles) and wins at C = 128, where two stages + # lose; this program runs from the smallest measured C where it beats the + # one-stage program. ``work`` as the ``fwd`` record's. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_8665d0756b8240ecb4d0": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "y"], + ["buffer", "work"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 4, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_8665d0756b8240ecb4d0", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 33920, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_8665d0756b8240ecb4d0", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "search_steps": "search_steps", + "w": "w", + "work": "work", + "x": "x", + "y": "y", + }, + "role": "fwd", + "select_rule": { + "channels": { + "sm_107a": 256, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "tile_counters", + "words_per_tile": 2, + }, + }, + # Route ``fwd_dt`` on sm_107a from C = 512 (round 9): the two-stage dynamic + # token-chunk forward with ``dtype_flags``, built with ``dtype_flags=True, + # fast=True, sched="dynamic", chunk_stages=2, pair_math=True``; ``work`` as + # the ``fwd`` record's. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "y"], + ["buffer", "work"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 4, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 33920, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "search_steps": "search_steps", + "w": "w", + "work": "work", + "x": "x", + "y": "y", + }, + "role": "fwd", + "select_rule": { + "channels": { + "sm_107a": 512, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "tile_counters", + "words_per_tile": 2, + }, + }, + # Route ``fwd_bias`` on sm_107a from C = 1024 (round 9): the two-stage dynamic + # token-chunk forward with the per-channel bias, built with ``bias=True, + # dtype_flags=True, fast=True, sched="dynamic", chunk_stages=2, + # pair_math=True``; ``work`` as the ``fwd`` record's. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_d829c95e1d610bc14b23": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "y"], + ["buffer", "b"], + ["buffer", "work"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 4, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_d829c95e1d610bc14b23", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 33920, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_d829c95e1d610bc14b23", + "logical": { + "C": "C", + "S": "S", + "b": "b", + "ct_vec": "ct_vec", + "cu": "cu", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "search_steps": "search_steps", + "w": "w", + "work": "work", + "x": "x", + "y": "y", + }, + "role": "fwd", + "select_rule": { + "channels": { + "sm_107a": 1024, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "tile_counters", + "words_per_tile": 2, + }, + }, + # Route ``bwd_dt`` (hand-written fallback, sm_100a / sm_103a): the fallback backward body + # ``vec=8, copy="tma2d"`` with ``dtype_flags=True`` (bit 0: BF16 taps; bit 1, the BF16 + # ``dw`` store, is read by its reduce program); SiLU' at the BF16-rounded preactivation + # as the default route. Like every backward program since the split it is an ordinary + # launch that writes ``dx`` and its partial rows; ``reduce`` names the program that sums + # them (the ``bwd_dt`` reduce: ``dtype_flags`` for the ``dw`` store form). + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa": { + "arches": [ + "sm_100a", + "sm_103a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 3, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 50176, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2", + "role": "bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Route ``bwd_dt`` (hand-written fallback, sm_107a): the sm_107a fallback body ``vec=8, + # min_blocks=2`` with ``dtype_flags=True``. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_8b9b8db0833fd64a578f": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_8b9b8db0833fd64a578f", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 50176, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_8b9b8db0833fd64a578f", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + }, + "reduce": "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2", + "role": "bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Route ``bwd_bias`` (hand-written fallback, sm_100a / sm_103a): the fallback body with + # ``bias=True`` (reads ``b``; the ``partial`` workspace grows a fifth slab, ``taps`` 5, + # of dbias partials that its reduce program sums into ``db``) and ``dtype_flags=True`` + # (input bits 0 and 2; the store bits 1 and 3 are the reduce program's), ``min_blocks=2`` + # (the bias accumulators do not fit the 168-register cap of three CTAs per SM). + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_2d91e5f926238453c78c": { + "arches": [ + "sm_100a", + "sm_103a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["buffer", "b"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_2d91e5f926238453c78c", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 54272, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_2d91e5f926238453c78c", + "logical": { + "C": "C", + "S": "S", + "b": "b", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_1c21c181af94e1635851", + "role": "bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 5, + }, + }, + # Route ``bwd_bias`` (hand-written fallback, sm_107a): the sm_107a fallback body with + # ``bias=True, dtype_flags=True``. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["buffer", "b"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 54272, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4", + "logical": { + "C": "C", + "S": "S", + "b": "b", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + }, + "reduce": "cake_causal_conv1d_det_train_1c21c181af94e1635851", + "role": "bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 5, + }, + }, + # Route ``bwd_fp32g`` (hand-written fallback, sm_100a / sm_103a): the fallback body with + # ``grad_pre="fp32"`` (SiLU' at the unrounded FP32 preactivation; its reduce program + # sums the partial rows in FP64 with one rounding), ``fold_every=4, fold_kahan=True`` + # and ``dtype_flags=True``, built ``min_blocks=2``: two CTAs per SM (the FP32-SiLU' + # body spills at the 168-register cap of three), so ``ctas_per_sm`` 2. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_c8027a64ba3cc38333bf": { + "arches": [ + "sm_100a", + "sm_103a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_c8027a64ba3cc38333bf", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 66560, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_c8027a64ba3cc38333bf", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_ce49146b9fef24adb47e", + "role": "bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Route ``bwd_fp32g`` (hand-written fallback, sm_107a): the sm_107a fallback body with + # ``grad_pre="fp32", dtype_flags=True``. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_f1545adfd2bb28044419": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_f1545adfd2bb28044419", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 66560, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_f1545adfd2bb28044419", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + }, + "reduce": "cake_causal_conv1d_det_train_ce49146b9fef24adb47e", + "role": "bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Route ``bwd_bias_fp32g`` (hand-written fallback, sm_100a / sm_103a): the fallback body + # with ``bias=True, grad_pre="fp32", dtype_flags=True``. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_149ca6f53934ba630033": { + "arches": [ + "sm_100a", + "sm_103a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["buffer", "b"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_149ca6f53934ba630033", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 74752, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_149ca6f53934ba630033", + "logical": { + "C": "C", + "S": "S", + "b": "b", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea", + "role": "bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 5, + }, + }, + # Route ``bwd_bias_fp32g`` (hand-written fallback, sm_107a): the sm_107a fallback body + # with ``bias=True, grad_pre="fp32", dtype_flags=True``. + # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows + # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the + # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. + "cake_causal_conv1d_det_train_843c010db78c79107fd3": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["buffer", "b"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_843c010db78c79107fd3", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 74752, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_843c010db78c79107fd3", + "logical": { + "C": "C", + "S": "S", + "b": "b", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + }, + "reduce": "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea", + "role": "bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 5, + }, + }, + # Rendered by the exporter from the family adapter's scaffold: the backward + # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, dtype_flags, + # fold_layout thread + # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no + # routed record of a higher threshold admits C; below every threshold the + # fallback record (no select_rule) runs. + "cake_causal_conv1d_det_train_2675cc1cb4efc0937378": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_2675cc1cb4efc0937378", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 82944, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_2675cc1cb4efc0937378", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2", + "role": "bwd", + "select_rule": { + "channels": { + "sm_100a": 512, + "sm_103a": 512, + "sm_107a": 1024, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Rendered by the exporter from the family adapter's scaffold: the backward + # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, bias, dtype_flags + # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no + # routed record of a higher threshold admits C; below every threshold the + # fallback record (no select_rule) runs. + "cake_causal_conv1d_det_train_517710d1f0988cf570a9": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["buffer", "b"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_517710d1f0988cf570a9", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 87040, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_517710d1f0988cf570a9", + "logical": { + "C": "C", + "S": "S", + "b": "b", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_1c21c181af94e1635851", + "role": "bwd", + "select_rule": { + "channels": { + "sm_100a": 512, + "sm_103a": 512, + "sm_107a": 1024, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "partial_rows", + "taps": 5, + }, + }, + # Rendered by the exporter from the family adapter's scaffold: the backward + # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, grad_pre fp32, + # fold_every 4, fold_kahan, dtype_flags + # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no + # routed record of a higher threshold admits C; below every threshold the + # fallback record (no select_rule) runs. + "cake_causal_conv1d_det_train_170d89481a616ae977b3": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_170d89481a616ae977b3", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 99328, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_170d89481a616ae977b3", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_ce49146b9fef24adb47e", + "role": "bwd", + "select_rule": { + "channels": { + "sm_100a": 512, + "sm_103a": 512, + "sm_107a": 1024, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + # Rendered by the exporter from the family adapter's scaffold: the backward + # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, bias, grad_pre fp32, + # fold_every 4, fold_kahan, dtype_flags + # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no + # routed record of a higher threshold admits C; below every threshold the + # fallback record (no select_rule) runs. + "cake_causal_conv1d_det_train_a681165b3e81204c73ff": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "partial"], + ["buffer", "b"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["parameter", "flags"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "grid_mult": { + "sm_100a": 1, + "sm_103a": 1, + "sm_107a": 1, + }, + "kind": "occupancy", + }, + "kernel": "cake_causal_conv1d_det_train_a681165b3e81204c73ff", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 107520, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_a681165b3e81204c73ff", + "logical": { + "C": "C", + "S": "S", + "b": "b", + "ct_vec": "ct_vec", + "cu": "cu", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "flags": "flags", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "x_map": "x", + }, + "reduce": "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea", + "role": "bwd", + "select_rule": { + "channels": { + "sm_100a": 512, + "sm_103a": 512, + "sm_107a": 1024, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "partial_rows", + "taps": 5, + }, + }, + # The reduce program every backward route shares: the second launch of a + # backward call. Launched programmatically dependent on the main program + # (``use_pdl``: the binding sets the programmatic-stream-serialization + # launch attribute; the kernel's first statement waits for the main grid, + # so the launch latency overlaps the main kernel's tail and an ordinary + # launch stays correct), it sums the FP32 ``[C, 4]`` partial rows of the + # workspace per channel: channel ``c`` lies in tile ``t = c // (vec x + # ct_vec)``, whose ``ceil((main_ctas - t) / n_ch_tiles)`` rows the main + # grid wrote, in the order the cooperative backward used in-kernel -- one + # warp per channel, lane sums over rows ``lane, lane + 32, ...`` then the + # xor tree 16, 8, 4, 2, 1 -- so ``dweight`` is bitwise what one launch + # produced. Built with its kernel file's default keywords, one program + # for the three architectures. + "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "partial"], + ["buffer", "dw"], + ["parameter", "C"], + ["parameter", "main_ctas"], + ["parameter", "n_ch_tiles"], + ["parameter", "ct_vec"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # One warp per channel, four channels per 128-thread CTA: the grid is + # ``ceil(C / channels_per_cta)`` CTAs, a function of the call's ``C``. + "grid_rule": { + "channels_per_cta": 4, + "cooperative": False, + "kind": "channels_per_cta", + }, + "kernel": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 2048, + "use_pdl": True, + }, + "library": "cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5", + "logical": { + "C": "C", + "ct_vec": "ct_vec", + "dw": "dw", + "main_ctas": "main_ctas", + "n_ch_tiles": "n_ch_tiles", + "partial": "partial", + }, + "role": "reduce", + "workspace_rule": None, + }, + # The reduce program of the ``bwd_dt`` route (hand-written; the render fills the name, + # symbol, argument plan and logical map): the default reduce built with ``dtype_flags=True``; it + # stores ``dw`` as BF16 when ``flags`` bit 1 says so; + # otherwise the default reduce's order and CTA shape (one warp per channel). + "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "partial"], + ["buffer", "dw"], + ["parameter", "C"], + ["parameter", "main_ctas"], + ["parameter", "n_ch_tiles"], + ["parameter", "ct_vec"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "channels_per_cta": 4, + "cooperative": False, + "kind": "channels_per_cta", + }, + "kernel": "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 2048, + "use_pdl": True, + }, + "library": "cake_causal_conv1d_det_train_ea8f91038a9dea452ff2", + "logical": { + "C": "C", + "ct_vec": "ct_vec", + "dw": "dw", + "flags": "flags", + "main_ctas": "main_ctas", + "n_ch_tiles": "n_ch_tiles", + "partial": "partial", + }, + "role": "reduce", + "workspace_rule": None, + }, + # The reduce program of the ``bwd_bias`` route (hand-written; the render fills the name, + # symbol, argument plan and logical map): the default reduce built with ``bias=True, dtype_flags=True``; it + # sums the fifth, dbias, slab of the partial rows into ``db`` as well; + # stores ``dw`` / ``db`` as BF16 when ``flags`` bit 1 / 3 says so; + # otherwise the default reduce's order and CTA shape (one warp per channel). + "cake_causal_conv1d_det_train_1c21c181af94e1635851": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "partial"], + ["buffer", "dw"], + ["buffer", "db"], + ["parameter", "C"], + ["parameter", "main_ctas"], + ["parameter", "n_ch_tiles"], + ["parameter", "ct_vec"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "channels_per_cta": 4, + "cooperative": False, + "kind": "channels_per_cta", + }, + "kernel": "cake_causal_conv1d_det_train_1c21c181af94e1635851", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 2560, + "use_pdl": True, + }, + "library": "cake_causal_conv1d_det_train_1c21c181af94e1635851", + "logical": { + "C": "C", + "ct_vec": "ct_vec", + "db": "db", + "dw": "dw", + "flags": "flags", + "main_ctas": "main_ctas", + "n_ch_tiles": "n_ch_tiles", + "partial": "partial", + }, + "role": "reduce", + "workspace_rule": None, + }, + # The reduce program of the ``bwd_fp32g`` route (hand-written; the render fills the name, + # symbol, argument plan and logical map): the default reduce built with ``fp64=True, dtype_flags=True``; it + # sums in FP64 (lane sums and the xor tree) with one final rounding; + # stores ``dw`` as BF16 when ``flags`` bit 1 says so; + # otherwise the default reduce's order and CTA shape (one warp per channel). + "cake_causal_conv1d_det_train_ce49146b9fef24adb47e": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "partial"], + ["buffer", "dw"], + ["parameter", "C"], + ["parameter", "main_ctas"], + ["parameter", "n_ch_tiles"], + ["parameter", "ct_vec"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "channels_per_cta": 4, + "cooperative": False, + "kind": "channels_per_cta", + }, + "kernel": "cake_causal_conv1d_det_train_ce49146b9fef24adb47e", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 4096, + "use_pdl": True, + }, + "library": "cake_causal_conv1d_det_train_ce49146b9fef24adb47e", + "logical": { + "C": "C", + "ct_vec": "ct_vec", + "dw": "dw", + "flags": "flags", + "main_ctas": "main_ctas", + "n_ch_tiles": "n_ch_tiles", + "partial": "partial", + }, + "role": "reduce", + "workspace_rule": None, + }, + # The reduce program of the ``bwd_bias_fp32g`` route (hand-written; the render fills the name, + # symbol, argument plan and logical map): the default reduce built with ``bias=True, fp64=True, dtype_flags=True``; it + # sums the fifth, dbias, slab of the partial rows into ``db`` as well; + # sums in FP64 (lane sums and the xor tree) with one final rounding; + # stores ``dw`` / ``db`` as BF16 when ``flags`` bit 1 / 3 says so; + # otherwise the default reduce's order and CTA shape (one warp per channel). + "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "partial"], + ["buffer", "dw"], + ["buffer", "db"], + ["parameter", "C"], + ["parameter", "main_ctas"], + ["parameter", "n_ch_tiles"], + ["parameter", "ct_vec"], + ["parameter", "flags"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "grid_rule": { + "channels_per_cta": 4, + "cooperative": False, + "kind": "channels_per_cta", + }, + "kernel": "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 5120, + "use_pdl": True, + }, + "library": "cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea", + "logical": { + "C": "C", + "ct_vec": "ct_vec", + "db": "db", + "dw": "dw", + "flags": "flags", + "main_ctas": "main_ctas", + "n_ch_tiles": "n_ch_tiles", + "partial": "partial", + }, + "role": "reduce", + "workspace_rule": None, + }, +} + +# Architecture -> route -> the records a call of the route selects from (the +# fallback first). The default routes ``fwd`` / ``bwd`` are the shipped +# programs; the six round-6 routes each hold their hand-written fallback +# (above) and, once rendered, the routed programs the exporter scaffolds from +# it; the forward routes' sm_107a two-stage programs are hand-written records. +ROUTES: dict[str, dict[str, list[str]]] = { + "sm_100a": { + "bwd": [ + "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce", + "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b", + ], + "bwd_bias": [ + "cake_causal_conv1d_det_train_2d91e5f926238453c78c", + "cake_causal_conv1d_det_train_517710d1f0988cf570a9", + ], + "bwd_bias_fp32g": [ + "cake_causal_conv1d_det_train_149ca6f53934ba630033", + "cake_causal_conv1d_det_train_a681165b3e81204c73ff", + ], + "bwd_dt": [ + "cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa", + "cake_causal_conv1d_det_train_2675cc1cb4efc0937378", + ], + "bwd_fp32g": [ + "cake_causal_conv1d_det_train_c8027a64ba3cc38333bf", + "cake_causal_conv1d_det_train_170d89481a616ae977b3", + ], + "fwd": [ + "cake_causal_conv1d_det_train_f5497330bfd286211643", + ], + "fwd_bias": [ + "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb", + ], + "fwd_dt": [ + "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e", + ], + }, + "sm_103a": { + "bwd": [ + "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce", + "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b", + "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6", + ], + "bwd_bias": [ + "cake_causal_conv1d_det_train_2d91e5f926238453c78c", + "cake_causal_conv1d_det_train_517710d1f0988cf570a9", + ], + "bwd_bias_fp32g": [ + "cake_causal_conv1d_det_train_149ca6f53934ba630033", + "cake_causal_conv1d_det_train_a681165b3e81204c73ff", + ], + "bwd_dt": [ + "cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa", + "cake_causal_conv1d_det_train_2675cc1cb4efc0937378", + ], + "bwd_fp32g": [ + "cake_causal_conv1d_det_train_c8027a64ba3cc38333bf", + "cake_causal_conv1d_det_train_170d89481a616ae977b3", + ], + "fwd": [ + "cake_causal_conv1d_det_train_f5497330bfd286211643", + ], + "fwd_bias": [ + "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb", + ], + "fwd_dt": [ + "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e", + ], + }, + "sm_107a": { + "bwd": [ + "cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c", + "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b", + "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6", + "cake_causal_conv1d_det_train_aeff6106e2822f313a0c", + "cake_causal_conv1d_det_train_9a07bcaa0223632640ee", + ], + "bwd_bias": [ + "cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4", + "cake_causal_conv1d_det_train_517710d1f0988cf570a9", + ], + "bwd_bias_fp32g": [ + "cake_causal_conv1d_det_train_843c010db78c79107fd3", + "cake_causal_conv1d_det_train_a681165b3e81204c73ff", + ], + "bwd_dt": [ + "cake_causal_conv1d_det_train_8b9b8db0833fd64a578f", + "cake_causal_conv1d_det_train_2675cc1cb4efc0937378", + ], + "bwd_fp32g": [ + "cake_causal_conv1d_det_train_f1545adfd2bb28044419", + "cake_causal_conv1d_det_train_170d89481a616ae977b3", + ], + "fwd": [ + "cake_causal_conv1d_det_train_f5497330bfd286211643", + "cake_causal_conv1d_det_train_8665d0756b8240ecb4d0", + ], + "fwd_bias": [ + "cake_causal_conv1d_det_train_08880ca12e4064bc8cdb", + "cake_causal_conv1d_det_train_d829c95e1d610bc14b23", + ], + "fwd_dt": [ + "cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e", + "cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc", + ], + }, +} + +# One verified build per architecture: B200 (``sm_100a``), B300 (``sm_103a``) +# and R200 (``sm_107a``). Every architecture is compiled with its exact +# ``-gencode`` flag set (``build.json`` ``arch_flags``); there is no family +# fallback. +ARCHES = ("sm_100a", "sm_103a", "sm_107a") +ARCH_BY_CAPABILITY: dict[tuple[int, int], str] = { + (10, 0): "sm_100a", + (10, 3): "sm_103a", + (10, 7): "sm_107a", +} +# The routes (``ROUTES[arch]`` keys) a public call may resolve to. ``fwd`` / +# ``bwd`` are the default routes (FP32 taps, no bias, SiLU' at the BF16-rounded +# preactivation); ``_dt`` routes take BF16 taps and / or store a BF16 gradient +# (the ``flags`` scalar says which), ``_bias`` routes add the per-channel bias +# and its gradient, ``_fp32g`` routes evaluate SiLU' at the unrounded FP32 +# preactivation for the weight gradient. :func:`forward_route` and +# :func:`backward_route` are the host's map from a call's arguments to a route. +KINDS = ( + "fwd", + "fwd_dt", + "fwd_bias", + "bwd", + "bwd_dt", + "bwd_bias", + "bwd_fp32g", + "bwd_bias_fp32g", +) +# The routes every exported architecture ships: the base forward and backward. +# The other routes are optional per architecture; a call that needs one the +# architecture lacks is refused with a ``ValueError`` naming the route. +REQUIRED_KINDS = ("fwd", "bwd") +# The role of a module record: a forward program, a backward program (the +# first launch of a backward call) or the reduce program a backward record +# names in ``record["reduce"]`` (the second launch). +ROLES = ("fwd", "bwd", "reduce") +# The grid-rule vocabulary (``record["grid_rule"]["kind"]``, see +# :func:`launch_grid`): ``co_resident`` launches exactly ``ctas_per_sm x SM +# count`` CTAs, the co-resident count (the forward records and the +# cooperative KCP programs of :mod:`.kcp_jit`, which share this host plan); +# ``occupancy`` launches ``grid_mult[arch] x ctas_per_sm x SM count`` CTAs +# normally (the backward programs since the split; ``grid_mult`` names every +# architecture of its record); ``channels_per_cta`` launches ``ceil(C / +# channels_per_cta)`` CTAs (the reduce programs). +GRID_KINDS = ("co_resident", "occupancy", "channels_per_cta") +# The selection-rule vocabulary of a forward or backward record +# (``record["select_rule"]["kind"]``); of the records whose rule admits a call +# the one of the highest channel threshold runs, and a record without a rule is +# its route's fallback. ``channels_tokens_at_least`` adds a token floor to the +# channel threshold. A reduce record never carries one. +SELECT_KINDS = ("channels_at_least", "channels_tokens_at_least") +# The ``flags`` bit mask of the ``dtype_flags`` programs (every program of the +# six round-6 routes): the host derives it from the tap tensor, the ACTUAL +# ``dweight`` out tensor, the bias tensor and the ACTUAL ``dbias`` out tensor; +# the kernel reads it once in its prologue (load form) and once in its epilogue +# (store form). The default routes' programs take no ``flags``. +FLAG_W_BF16 = 1 # ``w`` holds BF16 taps (else FP32) +FLAG_DW_BF16 = 2 # ``dw`` is stored as BF16 (else FP32) +FLAG_B_BF16 = 4 # ``b`` holds BF16 (else FP32); bias programs +FLAG_DB_BF16 = 8 # ``db`` is stored as BF16 (else FP32); bias programs +# The convolution width: ``weight`` and ``dweight`` are ``[C, WIDTH]`` and tap +# ``WIDTH - 1`` multiplies the current token. +WIDTH = 4 +# A forward thread and a backward thread each own one 16-byte vector of +# eight BF16 channels, so ``C`` is a multiple of eight; the ``vec`` of every +# record's geometry divides this. +CHANNEL_MULTIPLE = 8 +# The kernels take ``S``, ``C`` and the plan scalars as 32-bit integers. +INT32_MAX = 2**31 - 1 +# Logical argument names the public API binds; every module record maps its +# kernel parameter names onto this vocabulary (``record["logical"]``). +LOGICAL_TENSORS = ( + "x", + "w", + "cu", + "y", + "dy", + "dx", + "dw", + "partial", + "b", + "db", + "work", +) +# The workspace-rule vocabulary (``record["workspace_rule"]["kind"]``): a +# backward program's ``partial_rows`` (its FP32 partial rows, :func:`partial_bytes`) +# and a forward program's ``tile_counters`` (its int32 ``work`` counters, +# ``words_per_tile`` per channel tile, :func:`counter_words`); a reduce +# program has none. +WORKSPACE_KINDS = ("partial_rows", "tile_counters") +# The scalars :func:`plan` produces, in kernel order; ``flags`` (the dtype bit +# mask of the ``dtype_flags`` programs) and ``main_ctas`` (the CTA count of the +# selected backward program's launch, ``grid_x``, from which the reduce program +# derives the partial rows of each channel tile, ``ceil((main_ctas - tile) / +# n_ch_tiles)``) are set by the host beside them. +PLAN_SCALARS = ( + "S", + "C", + "num_docs", + "search_steps", + "ct_vec", + "n_sub", + "n_ch_tiles", +) +LOGICAL_SCALARS = (*PLAN_SCALARS, "flags", "main_ctas") + + +def arch_for_capability(device_capability: Sequence[int]) -> str | None: + """Name the exported architecture for one device capability, if any.""" + + values = [int(value) for value in device_capability] + if len(values) != 2: + return None + return ARCH_BY_CAPABILITY.get((values[0], values[1])) + + +def route_applies(*, device_capability: Sequence[int], kind: str) -> bool: + """Whether the verified export owns one training call on a device.""" + + arch = arch_for_capability(device_capability) + if arch is None: + return False + return kind in ROUTES.get(arch, {}) + + +def route_role(kind: str) -> str: + """``"fwd"`` or ``"bwd"``: the kernel a route's programs are built from (``ValueError`` for an unknown route).""" + + if kind not in KINDS: + raise ValueError(f"kind must be one of {KINDS}, got {kind!r}") + return "fwd" if kind.startswith("fwd") else "bwd" + + +def forward_route(*, bias: bool, bf16_taps: bool) -> str: + """The forward route of a call: ``fwd_bias`` with a bias, else ``fwd_dt`` for BF16 taps, else ``fwd``. + + A bias program always has a bias (no runtime switch: adding ``+0.0`` would + flip the sign of ``-0.0`` outputs), so ``bias=None`` never selects one; the + bias programs read FP32 or BF16 taps and bias through ``flags``. + """ + + if bias: + return "fwd_bias" + return "fwd_dt" if bf16_taps else "fwd" + + +def backward_route(*, bias: bool, fp32_grad: bool, bf16_taps_or_dweight: bool) -> str: + """The backward route of a call. + + ``bias`` selects ``bwd_bias`` / ``bwd_bias_fp32g`` (by ``fp32_grad``, the + ``silu_grad_preactivation="fp32"`` policy); without a bias ``fp32_grad`` + selects ``bwd_fp32g``; else BF16 taps or a BF16 ``dweight`` out tensor + select ``bwd_dt``; else the default ``bwd``. Every route but ``bwd`` + reads the tap / gradient dtypes from ``flags``. + """ + + if bias: + return "bwd_bias_fp32g" if fp32_grad else "bwd_bias" + if fp32_grad: + return "bwd_fp32g" + return "bwd_dt" if bf16_taps_or_dweight else "bwd" + + +def route_modules(arch: str, kind: str) -> tuple[str, ...]: + """Ordered module names of one route on ``arch`` (every program the route may launch). + + A call launches the one :func:`selected_modules` picks for its channel + count. + """ + + route_role(kind) + names = ROUTES.get(arch, {}).get(kind) + if not names: + raise ValueError( + f"unsupported Cake causal conv1d training route {kind!r} on {arch!r}; " + f"exported architectures: {ARCHES}" + ) + return tuple(names) + + +def select_rule(record: dict[str, Any]) -> dict[str, Any] | None: + """The selection rule of one module record, ``None`` when the record carries none. + + A record without a rule (no ``select_rule`` key, or ``None``) is its + route's fallback program. ``{"kind": "channels_at_least", "channels": + {arch: C_min, ...}}`` owns the calls whose channel count is at least + ``channels[arch]`` on the architectures it names and never hits on the + others. ``{"kind": "channels_tokens_at_least", "channels": {arch: C_min, + ...}, "tokens": {arch: S_min, ...}}`` owns those calls only when the token + count is also at least ``tokens[arch]``. When several records own a call, + the one of the highest channel threshold runs. + """ + + return record.get("select_rule") + + +def _threshold(rule: dict[str, Any], *, arch: str) -> int | None: + """The channel count from which a selection rule owns the calls on ``arch``, ``None`` where it never does + (``ValueError`` for an unknown kind).""" + + kind = rule["kind"] + if kind in SELECT_KINDS: + channels = rule["channels"] + return int(channels[arch]) if arch in channels else None + raise ValueError( + f"unknown Cake causal conv1d training selection rule: {kind!r} " + f"(known: {SELECT_KINDS})" + ) + + +def _tokens(rule: dict[str, Any], *, arch: str) -> int | None: + """The token count from which a selection rule owns the calls on ``arch`` (its token floor), ``None`` when the + rule has none there (``ValueError`` for an unknown kind).""" + + _threshold(rule, arch=arch) + if rule["kind"] != "channels_tokens_at_least": + return None + tokens = rule["tokens"] + return int(tokens[arch]) if arch in tokens else None + + +def _admits(rule: dict[str, Any], *, arch: str, C: int, S: int | None = None) -> bool: + """Whether a selection rule owns ``C`` channels (and ``S`` tokens) on ``arch``. + + ``ValueError`` for an unknown kind, or when the rule's token floor decides + and ``S`` is ``None``. + """ + + threshold = _threshold(rule, arch=arch) + if threshold is None or int(C) < threshold: + return False + floor = _tokens(rule, arch=arch) + if floor is None: + return True + if S is None: + raise ValueError( + f"the Cake causal conv1d training selection rule {rule!r} on {arch} " + f"for C={int(C)} runs from S {floor}; the selection needs S" + ) + return int(S) >= floor + + +def selected_modules( + arch: str, kind: str, *, C: int, S: int | None = None +) -> tuple[str, ...]: + """Ordered module names one public call launches on ``arch`` for ``C`` channels and ``S`` tokens. + + A route may carry several programs of its kernel and launches exactly one: + of the records whose :func:`select_rule` admits the call on ``arch`` the + one of the highest channel threshold, else the route's single fallback + record (the one without a rule). The choice is a pure function of the + architecture, the route, ``C`` and ``S`` -- ``S`` enters only through a + rule's token floor, and the host never reads the data -- so a given ``(S, + C, device)`` always runs the same program of a route and ``dweight`` + stays bitwise reproducible. Two records of one channel threshold that + both reach ``C``, or zero or several fallbacks, are registry errors + (``RuntimeError``); an unknown rule kind, or ``S`` of ``None`` where a + token floor decides, is a ``ValueError``. + """ + + names = route_modules(arch, kind) + candidates: list[tuple[int, str]] = [] + fallback: list[str] = [] + for name in names: + rule = select_rule(MODULES[name]) + if rule is None: + fallback.append(name) + continue + threshold = _threshold(rule, arch=arch) + if threshold is not None and int(C) >= threshold: + candidates.append((threshold, name)) + if len(fallback) != 1: + raise RuntimeError( + f"the exported route {kind!r} on {arch} declares {len(fallback)} " + "fallback programs (records without a select_rule); expected exactly one" + ) + candidates.sort(reverse=True) + if len(candidates) > 1 and candidates[0][0] == candidates[1][0]: + tied = [name for threshold, name in candidates if threshold == candidates[0][0]] + raise RuntimeError( + f"the exported route {kind!r} on {arch} selects {len(tied)} " + f"programs for C={int(C)} at one threshold {candidates[0][0]} ({tied}); " + "thresholds must be distinct" + ) + for _threshold_c, name in candidates: + if _admits(select_rule(MODULES[name]), arch=arch, C=C, S=S): + return (name,) + return (fallback[0],) + + +def backward_kinds(arch: str) -> tuple[str, ...]: + """The backward routes ``arch`` exports, in :data:`KINDS` order.""" + + return tuple( + kind + for kind in KINDS + if route_role(kind) == "bwd" and ROUTES.get(arch, {}).get(kind) + ) + + +def workspace_owners(arch: str) -> tuple[dict[str, Any], ...]: + """The backward module records of ``arch`` that own a ``partial`` workspace, over every backward route in + :data:`KINDS` order, each route in launch order. + + Every backward program declares a ``workspace_rule``; an architecture + without one is a registry error (``RuntimeError``). + """ + + owners = tuple( + MODULES[name] + for kind in backward_kinds(arch) + for name in route_modules(arch, kind) + if MODULES[name]["workspace_rule"] is not None + ) + if not owners: + raise RuntimeError( + f"the exported backward routes on {arch} declare no workspace owner; " + "expected at least one" + ) + return owners + + +def workspace_record(arch: str) -> dict[str, Any]: + """The capacity owner of the backward ``partial`` workspace on ``arch``. + + The public workspace serves every backward program of every backward route + (:func:`workspace_owners`), each of which writes the leading rows of its + own grid. The owners share one build geometry, so for every ``C`` their + rows are ``ceil(grid_mult x ctas_per_sm x SM count / n_ch_tiles)`` with + one ``n_ch_tiles`` and their bytes ``4 x taps x C x rows``: the record + with the largest :func:`grid_per_sm` ``x taps`` (the bias programs' fifth + slab counts; :data:`KINDS` then route order break ties) bounds them all. + Owners of different geometries are a registry error (``RuntimeError``). + """ + + owners = workspace_owners(arch) + geometries = { + tuple(sorted(owner["grid_rule"]["geometry"].items())) for owner in owners + } + if len(geometries) != 1: + raise RuntimeError( + f"the backward programs on {arch} declare {len(geometries)} build " + "geometries; one workspace serves them only when they share one" + ) + return max( + owners, + key=lambda owner: ( + grid_per_sm(owner, arch) * int(owner["workspace_rule"]["taps"]) + ), + ) + + +def grid_mult(record: dict[str, Any], arch: str) -> int: + """The over-decomposition factor of an ``occupancy`` record on ``arch`` (1 under every other grid rule). + + ``grid_rule["grid_mult"]`` maps each architecture of the record to the + factor its grid multiplies the co-resident CTA count by; the factor per + architecture is the measured trade of tail latency behind a resident + side kernel against partial-row traffic. + """ + + rule = record["grid_rule"] + if rule["kind"] != "occupancy": + return 1 + factor = int(rule["grid_mult"][arch]) + if factor < 1: + raise ValueError( + f"module {record['kernel']}: grid_mult[{arch!r}] must be positive, got {factor}" + ) + return factor + + +def grid_per_sm(record: dict[str, Any], arch: str) -> int: + """CTAs per SM of a ``co_resident`` or ``occupancy`` record's grid on ``arch``: + ``grid_mult x ctas_per_sm`` (the grid is this times the SM count).""" + + return grid_mult(record, arch) * int(record["grid_rule"]["ctas_per_sm"]) + + +def _ceil_div(numerator: int, denominator: int) -> int: + if denominator <= 0: + raise ValueError("denominator must be positive") + return -(-int(numerator) // int(denominator)) + + +# The roles whose programs plan their channel tiles with :func:`channel_split` +# (a reduce program runs on its backward record's plan): each has its own rule. +SPLIT_ROLES = ("fwd", "bwd") + + +def channel_split( + C: int, + *, + geometry: dict[str, int], + role: str | None = None, + legacy: bool = False, +) -> tuple[int, int, int]: + """``(ct_vec, n_sub, n_ch_tiles)`` of ``C`` channels under a record's build geometry and role. + + ``ct_vec`` channel vectors per channel tile (one thread each), ``n_sub = + threads // ct_vec`` token sub-tiles per CTA and ``n_ch_tiles = ceil(C / vec + / ct_vec)`` channel tiles. ``role`` is the record's, ``"fwd"`` or + ``"bwd"`` (a ``ValueError`` without one unless ``legacy``): the two rules + agree except in the fallback below. With ``cv = C / vec`` channel + vectors: ``ct_vec = cv`` when ``cv <= threads``; above that, the multiple + of 32 dividing ``cv`` (at most ``threads``) with the highest lane use + ``ct_vec x (threads // ct_vec)``, the larger ``ct_vec`` on a tie. So C = + 1280 (160 vectors) runs five 32-vector tiles with four token sub-tiles + each, (32, 4, 5), and C = 1536 (64, 2, 3) -- every lane busy -- instead of + two 128-vector tiles at 62 % / 75 % lane use; C <= 1024 and 2048 / 4096 / + 8192 are unchanged. When no multiple of 32 divides ``cv``, the role's + fallback picks the ``ct_vec`` of the highest lane use over the launch, + ``cv / (n_ch_tiles x ct_vec) x (ct_vec x (threads // ct_vec) / threads)`` + rounded to 9 decimals (the larger on a tie), among every multiple of 32 at + most ``threads`` (the last tile partial) and, for the forward, every + divisor of ``cv`` from 40 to ``threads`` (whole tiles). At 128 x 8 the + backward runs C = 1160 as (32, 4, 5), 1344 as (64, 2, 3) and 1600 as (32, + 4, 7), the forward 1160 as (32, 4, 5), 1344 as (42, 3, 4) and 1600 as (40, + 3, 5); both run 1920 as (128, 1, 2). The per element arithmetic does not + depend on the tiling, so ``y`` and ``dx`` are bitwise the legacy plan's; + ``dweight``'s fixed summation order follows the tiling, so its bits differ + from the legacy plan's (both are reproducible). ``legacy=True`` is the + rule the package shipped before round 6, ``ct_vec = min(cv, threads)``, + for either role; the KCP path keeps it (its kernels pick the grouped path + from ``ct_vec`` on the device and are measured under this rule). Each + role's rule is kept in sync with its copies (the kernel generator of + that role and the test plan of that kind). + """ + + return _channel_split( + int(C), int(geometry["vec"]), int(geometry["threads"]), role, bool(legacy) + ) + + +@functools.lru_cache(maxsize=4096) +def _channel_split( + C: int, vec: int, threads: int, role: str | None, legacy: bool +) -> tuple[int, int, int]: + """:func:`channel_split` of plain arguments, memoized: a model's channel + counts are few, so every call after a channel count's first reuses its + plan (a backward call plans its row up to twice).""" + + if (role is None and not legacy) or (role is not None and role not in SPLIT_ROLES): + raise ValueError( + f"role must be one of {SPLIT_ROLES} (or legacy=True), got {role!r}" + ) + if C <= 0 or C % vec: + raise ValueError(f"C={C} must be a positive multiple of {vec}") + channel_vectors = C // vec + if legacy or channel_vectors <= threads: + ct_vec = min(channel_vectors, threads) + else: + candidates = [ + width + for width in range(32, threads + 1, 32) + if channel_vectors % width == 0 + ] + ct_vec = ( + max(candidates, key=lambda width: (width * (threads // width), width)) + if candidates + else _fallback_split(channel_vectors, threads, role) + ) + return ct_vec, threads // ct_vec, _ceil_div(channel_vectors, ct_vec) + + +def _fallback_split(channel_vectors: int, threads: int, role: str) -> int: + """``ct_vec`` of :func:`channel_split` above ``threads`` vectors with no 32-multiple divisor. + + The candidates are every multiple of 32 at most ``threads`` (the last tile + partial: its copies and stores stop at channel ``C``) and, for the + forward, every divisor of ``channel_vectors`` from 40 to ``threads`` (whole + tiles); the one of the highest :func:`_lane_use`, rounded to 9 decimals, + wins, the larger on a tie. No tile is narrower than a warp's 32 vectors + and no whole forward tile narrower than 40. + """ + + candidates = set(range(32, threads + 1, 32)) + if role == "fwd": + candidates |= { + width for width in range(40, threads + 1) if channel_vectors % width == 0 + } + return max( + candidates, + key=lambda width: ( + round(_lane_use(channel_vectors, width, threads), 9), + width, + ), + ) + + +def _lane_use(channel_vectors: int, width: int, threads: int) -> float: + """The busy share of a CTA's lanes over ``ceil(channel_vectors / width)`` tiles of ``width``.""" + + return ( + channel_vectors + / (_ceil_div(channel_vectors, width) * width) + * (width * (threads // width) / threads) + ) + + +def plan( + S: int, + C: int, + num_docs: int, + *, + geometry: dict[str, int], + role: str | None = None, + legacy: bool = False, +) -> dict[str, int]: + """The kernel's launch scalars for one problem under a record's build geometry and role. + + ``geometry`` is ``record["grid_rule"]["geometry"]``: ``threads`` per CTA + and ``vec`` BF16 channels per thread; ``role`` is ``record["role"]`` + (``"fwd"`` or ``"bwd"``). The scalars are ``S``, ``C``, ``num_docs``, + :func:`channel_split`'s ``ct_vec`` / ``n_sub`` / ``n_ch_tiles`` (``role`` + and ``legacy`` are passed on) and ``search_steps``, the number of halvings + of the device-side binary search over ``cu_seqlens`` (``max(1, ceil(log2 + num_docs))``). Every value is a pure function of ``(S, C, num_docs)``, the + geometry and the role: the host never reads ``cu_seqlens``, and the kernel + derives its token ranges from ``S`` and the grid. The ``flags`` scalar of + the ``dtype_flags`` programs is not a plan value: the host derives it from + the tensors of the call. + """ + + S, C, num_docs = int(S), int(C), int(num_docs) + if S < 0 or num_docs < 0: + raise ValueError(f"S and num_docs must be non-negative, got {(S, num_docs)}") + ct_vec, n_sub, n_ch_tiles = channel_split( + C, geometry=geometry, role=role, legacy=legacy + ) + return { + "S": S, + "C": C, + "num_docs": num_docs, + "search_steps": max(1, (num_docs - 1).bit_length()), + "ct_vec": ct_vec, + "n_sub": n_sub, + "n_ch_tiles": n_ch_tiles, + } + + +# Whole-rank trim of the backward's single-wave main grid (:func:`main_grid`). +# A channel tile ``ct`` of a ``grid``-CTA launch gets ``ceil((grid - ct) / +# n_ch_tiles)`` CTAs, so when ``n_ch_tiles`` does not divide the grid two pairs +# of neighbouring tiles stream other token ranges at any instant; a row that is +# not 64-byte aligned shares a cache line (and a 32-byte ``dx`` sector) between +# them, which beyond L2 is then fetched twice. Dropping ``grid % n_ch_tiles`` +# CTAs gives every tile the same CTA count. ``TILE_ALIGNED_GRID[arch]``: the +# co-resident grid the rule was measured at (GB300's 304 CTAs; 0: none), on the +# routes of ``TILE_ALIGNED_ROUTES``, from ``TILE_ALIGNED_MIN_S`` tokens, with a +# last tile of at least ``TILE_ALIGNED_MIN_LAST[ct_vec]`` vectors and at most +# ``TILE_ALIGNED_MAX_TILES[ct_vec]`` tiles (a ``ct_vec`` without an entry: +# never). ``TILE_ALIGNED_CLASSES[arch]`` holds the other devices' measured +# classes, ``(co_resident, ct_vec, n_ch_tiles) -> (min S, min last-tile +# vectors)`` on the routes of ``TILE_ALIGNED_CLASS_ROUTES[arch]`` (B200's 296 +# CTAs; R200 none). Every other launch keeps its grid. ``dx`` does not depend +# on the grid; ``dweight`` / ``dbias`` sum other partial rows, bitwise +# reproducible for the trimmed grid. The rule is kept in sync with its copy in +# the backward kernel generator. +TILE_ALIGNED_GRID: dict[str, int] = {"sm_100a": 0, "sm_103a": 304, "sm_107a": 0} +TILE_ALIGNED_ROUTES: tuple[str, ...] = ("bwd",) +TILE_ALIGNED_MIN_S = 98304 +TILE_ALIGNED_MIN_LAST: dict[int, int] = {32: 15, 64: 49} +TILE_ALIGNED_MAX_TILES: dict[int, int] = {32: 27, 64: 9} +TILE_ALIGNED_CLASSES: dict[str, dict[tuple[int, int, int], tuple[int, int]]] = { + "sm_100a": { + (296, 32, 5): (147456, 27), + (296, 32, 7): (65536, 31), + (296, 64, 3): (147456, 59), + }, + "sm_103a": {}, + "sm_107a": {}, +} +TILE_ALIGNED_CLASS_ROUTES: dict[str, tuple[str, ...]] = { + "sm_100a": ("bwd", "bwd_dt", "bwd_fp32g"), + "sm_103a": (), + "sm_107a": (), +} + + +def main_grid( + co_resident: int, + *, + arch: str, + grid_mult: int = 1, + C: int | None = None, + n_ch_tiles: int = 1, + S: int | None = None, + ct_vec: int | None = None, + route: str | None = None, +) -> int: + """The backward main launch's CTA count: ``grid_mult x co_resident``. + + ``co_resident`` is the program's co-resident CTA count on the device + (``ctas_per_sm x SM count``). A single-wave launch (``grid_mult`` 1) of + a row the caller names (``C``, ``S``, the plan's ``n_ch_tiles`` and + ``ct_vec``, the ``route``) is trimmed to whole channel-tile ranks: + :func:`tile_aligned_grid` on a route of ``TILE_ALIGNED_ROUTES`` when + ``co_resident`` is ``TILE_ALIGNED_GRID[arch]``, else + :func:`class_aligned_grid` on a route of + ``TILE_ALIGNED_CLASS_ROUTES[arch]``. A caller that names no row keeps the + untrimmed grid, the most CTAs any row launches. + """ + + if type(grid_mult) is not int or grid_mult < 1: + raise ValueError(f"grid_mult must be a positive int, got {grid_mult!r}") + grid = int(co_resident) * grid_mult + named = None not in (C, S, ct_vec) + if ( + grid_mult == 1 + and co_resident == TILE_ALIGNED_GRID[arch] + and route in TILE_ALIGNED_ROUTES + and named + ): + grid = tile_aligned_grid(grid, C, n_ch_tiles, S=S, ct_vec=ct_vec) + elif grid_mult == 1 and route in TILE_ALIGNED_CLASS_ROUTES[arch] and named: + grid = class_aligned_grid(grid, C, n_ch_tiles, arch=arch, S=S, ct_vec=ct_vec) + return grid + + +def tile_aligned_grid( + grid: int, C: int, n_ch_tiles: int, *, S: int, ct_vec: int, vec: int = 8 +) -> int: + """``grid`` less ``grid % n_ch_tiles`` CTAs when ``C % 32`` (a row of + ``x`` / ``dy`` / ``dx`` is not 64-byte aligned), ``1 < n_ch_tiles <= + TILE_ALIGNED_MAX_TILES[ct_vec]``, ``grid % n_ch_tiles``, ``grid > + n_ch_tiles``, ``S >= TILE_ALIGNED_MIN_S`` and the last channel tile holds + at least ``TILE_ALIGNED_MIN_LAST[ct_vec]`` vectors (``C / vec - + (n_ch_tiles - 1) x ct_vec``); else ``grid``.""" + + last = C // vec - (n_ch_tiles - 1) * ct_vec + if ( + 1 < n_ch_tiles <= TILE_ALIGNED_MAX_TILES.get(ct_vec, 0) + and C % 32 + and grid % n_ch_tiles + and grid > n_ch_tiles + and S >= TILE_ALIGNED_MIN_S + and last >= TILE_ALIGNED_MIN_LAST.get(ct_vec, ct_vec + 1) + ): + return grid - grid % n_ch_tiles + return grid + + +def class_aligned_grid( + grid: int, + C: int, + n_ch_tiles: int, + *, + arch: str, + S: int, + ct_vec: int, + vec: int = 8, +) -> int: + """``grid`` less ``grid % n_ch_tiles`` CTAs when ``C % 32``, ``grid % + n_ch_tiles`` and ``(grid, ct_vec, n_ch_tiles)`` is a class of + ``TILE_ALIGNED_CLASSES[arch]`` whose bounds the row meets (``S`` and the + last tile's ``C / vec - (n_ch_tiles - 1) x ct_vec`` vectors at least the + class's); else ``grid``.""" + + bound = TILE_ALIGNED_CLASSES[arch].get((grid, ct_vec, n_ch_tiles)) + if bound is None or not C % 32 or not grid % n_ch_tiles or bound[0] > S: + return grid + if C // vec - (n_ch_tiles - 1) * ct_vec < bound[1]: + return grid + return grid - grid % n_ch_tiles + + +def row_may_trim( + record: dict[str, Any], arch: str, co_resident: int, factor: int, C: int, S: int +) -> bool: + """Whether :func:`main_grid` may trim ``record``'s launch for the row. + + Every trim needs a single wave (``factor`` 1), ``C % 32``, the record on a + route of the rule and ``S`` at least the rule's smallest token bound for + ``co_resident``; a call failing any of these launches the untrimmed grid + without planning the row's channel tiles (no host cost beyond the + untrimmed grid's).""" + + if factor != 1 or not C % 32: + return False + library = record["library"] + routes = ROUTES[arch] + if ( + co_resident == TILE_ALIGNED_GRID[arch] + and S >= TILE_ALIGNED_MIN_S + and any(library in routes.get(kind, ()) for kind in TILE_ALIGNED_ROUTES) + ): + return True + return any( + grid == co_resident and bound[0] <= S + for (grid, _, _), bound in TILE_ALIGNED_CLASSES[arch].items() + ) and any( + library in routes.get(kind, ()) for kind in TILE_ALIGNED_CLASS_ROUTES[arch] + ) + + +def record_route(record: dict[str, Any], arch: str) -> str: + """The route of ``ROUTES[arch]`` that lists ``record`` (``ValueError`` if none).""" + + return _library_route(record["library"], arch) + + +@functools.cache +def _library_route(library: str, arch: str) -> str: + for kind, names in ROUTES[arch].items(): + if library in names: + return kind + raise ValueError(f"module {library} is on no route of {arch}") + + +@functools.cache +def sm_count(device_index: int) -> int: + """Streaming multiprocessor count of one CUDA device (immutable per process).""" + + import torch + + return int(torch.cuda.get_device_properties(device_index).multi_processor_count) + + +def launch_grid( + record: dict[str, Any], + *, + device_index: int, + C: int | None = None, + S: int | None = None, +) -> tuple[int, int, int]: + """Launch grid of one module on ``device_index`` under its recorded grid rule. + + ``co_resident`` modules launch exactly ``ctas_per_sm x SM count`` CTAs, + the co-resident count (a cooperative KCP backward never exceeds it). + ``occupancy`` modules launch :func:`main_grid` of that count: + ``grid_mult[arch] x ctas_per_sm x SM count`` CTAs normally + (:func:`grid_per_sm`): ``ctas_per_sm`` co-resident CTAs per SM and + ``grid_mult`` times finer token streams, so a CTA that finds no free slot + beside another stream's kernel trails by a short stream in a later wave. + Given the call's ``C`` and ``S``, a single-wave launch of the rows + :func:`main_grid` names (its route, the record's plan) drops CTAs to whole + channel-tile ranks (the row is planned only where :func:`row_may_trim` + holds). Either grid is a pure function of ``(S, C, device)`` and the + route, so is the deterministic ``dweight`` order, and the untrimmed grid + (``C`` or ``S`` not given) is the most CTAs any call launches, so one + workspace serves every token count. + ``channels_per_cta`` modules (the reduce programs) launch ``ceil(C / + channels_per_cta)`` CTAs for the call's ``C`` channels. + """ + + rule = record["grid_rule"] + kind = rule["kind"] + if kind in ("co_resident", "occupancy"): + arch = device_arch(device_index) + grid_x = grid_per_sm(record, arch) * sm_count(device_index) + if kind == "occupancy" and C is not None and S is not None and C % 32: + co_resident = int(rule["ctas_per_sm"]) * sm_count(device_index) + factor = grid_mult(record, arch) + if row_may_trim(record, arch, co_resident, factor, int(C), int(S)): + ct_vec, _, n_ch_tiles = channel_split( + C, geometry=rule["geometry"], role=record["role"] + ) + grid_x = main_grid( + co_resident, + arch=arch, + grid_mult=factor, + C=int(C), + n_ch_tiles=n_ch_tiles, + S=int(S), + ct_vec=ct_vec, + route=record_route(record, arch), + ) + elif kind == "channels_per_cta": + if C is None: + raise ValueError( + f"the {kind} grid rule of module {record['kernel']} needs the channel count C" + ) + grid_x = _ceil_div(C, int(rule["channels_per_cta"])) + else: + raise ValueError( + f"unknown Cake causal conv1d training grid rule: {kind!r} (known: {GRID_KINDS})" + ) + return (max(1, grid_x), 1, 1) + + +def partial_rows(grid_x: int, n_ch_tiles: int) -> int: + """Rows of the backward ``partial`` workspace for a grid: ``ceil(grid_x / n_ch_tiles)``. + + CTA ``b`` of the grid owns channel tile ``b % n_ch_tiles`` and partial row + ``b // n_ch_tiles``, so tile ``t`` holds ``ceil((grid_x - t) / n_ch_tiles)`` + written rows; the reduce program derives exactly those from ``main_ctas`` + (= ``grid_x``) and ``n_ch_tiles`` and sums nothing else. + """ + + return _ceil_div(grid_x, n_ch_tiles) + + +def partial_bytes(record: dict[str, Any], C: int, rows: int) -> int: + """Bytes of ``rows`` FP32 ``[C, taps]`` partial rows of the module's workspace rule. + + ``taps`` is 4 for the dweight partials; a bias program's rule says 5: its + kernel lays the ``[rows, C, 4]`` tap slab out first and the ``[rows, C]`` + dbias slab after it (the host view is capacity only). + """ + + rule = record["workspace_rule"] + if rule is None or rule["kind"] != "partial_rows": + raise ValueError(f"module {record['kernel_symbol']} owns no backward workspace") + return 4 * int(rule["taps"]) * int(C) * int(rows) + + +def workspace_layout( + record: dict[str, Any], C: int, *, device_index: int +) -> dict[str, int]: + """Byte layout of the caller-owned backward workspace of one module for ``C`` channels. + + The workspace is one CUDA ``uint8`` buffer holding the FP32 ``[rows, C, taps]`` + ``dweight`` (and, ``taps`` 5, dbias) partials, ``rows = partial_rows(grid, + n_ch_tiles)`` for the device's launch grid (``grid_mult x ctas_per_sm x SM + count``: the bytes grow with the record's ``grid_mult``); it depends on + ``C`` and the device only, so one buffer serves every token count. The + kernel fully writes the rows its reduce program sums, so nothing is zeroed + between launches or CUDA Graph replays. + """ + + grid_x = launch_grid(record, device_index=device_index)[0] + _, _, n_ch_tiles = channel_split( + C, geometry=record["grid_rule"]["geometry"], role=record["role"] + ) + rows = partial_rows(grid_x, n_ch_tiles) + return { + "grid": grid_x, + "rows": rows, + "taps": int(record["workspace_rule"]["taps"]), + "total_bytes": partial_bytes(record, C, rows), + } + + +def counter_words(record: dict[str, Any], n_ch_tiles: int) -> int: + """Int32 words of a forward record's ``work`` counters for a launch of ``n_ch_tiles`` channel tiles. + + A ``tile_counters`` rule holds ``words_per_tile`` words per channel tile + (the dynamic forward's next-chunk counter and done-CTA counter); a record + without a workspace rule takes no ``work`` (0 words). + """ + + rule = record["workspace_rule"] + if rule is None: + return 0 + if rule["kind"] != "tile_counters": + raise ValueError( + f"module {record['kernel']}: workspace rule {rule['kind']!r} holds no forward counters" + ) + if n_ch_tiles <= 0: + raise ValueError(f"n_ch_tiles must be positive, got {n_ch_tiles}") + return int(rule["words_per_tile"]) * int(n_ch_tiles) + + +def forward_kinds(arch: str) -> tuple[str, ...]: + """The forward routes ``arch`` exports, in :data:`KINDS` order.""" + + return tuple( + kind + for kind in KINDS + if route_role(kind) == "fwd" and ROUTES.get(arch, {}).get(kind) + ) + + +def forward_workspace_words(arch: str, C: int) -> int: + """Int32 words of the forward workspace for ``C`` channels on ``arch``. + + The largest :func:`counter_words` over every forward program of every + forward route of the architecture (each planned with its own geometry and + role), so one buffer serves every forward route; it depends on ``C`` only + (never on ``S`` or the device's SM count). 0 when no forward program of + the architecture takes ``work``. + """ + + words = 0 + for kind in forward_kinds(arch): + for name in route_modules(arch, kind): + record = MODULES[name] + if record["workspace_rule"] is None: + continue + _, _, n_ch_tiles = channel_split( + C, geometry=record["grid_rule"]["geometry"], role=record["role"] + ) + words = max(words, counter_words(record, n_ch_tiles)) + return words + + +@functools.cache +def device_arch(device_index: int) -> str: + """Exported architecture of one CUDA device (``RuntimeError`` when it is not an exported target).""" + + import torch + + capability = tuple(torch.cuda.get_device_capability(device_index)) + arch = arch_for_capability(capability) + if arch is None: + raise RuntimeError( + f"device {device_index} (capability {capability}) is not an exported Cake causal conv1d training target" + ) + return arch + + +def library_name(name: str, arch: str) -> str: + """``build.json`` library of module record ``name`` built for ``arch``: ``_``.""" + + record = MODULES[name] + if arch not in record["arches"]: + raise RuntimeError( + f"Cake causal conv1d training module {name} is not delivered for {arch}; " + f"its architectures: {record['arches']}" + ) + return f"{record['library']}_{arch}" + + +@functools.cache +def load(name: str, arch: str) -> Any: + """The loaded tvm-ffi module of record ``name`` built for ``arch``. + + The prebuilt ``aot/`` library when a builder compiled it, else FlashInfer's + JIT build of the same sources and flags (``cake_build.load_library``). + """ + + return cake_build.load_library(library_name(name, arch)) + + +def specs_for_arch(arch: str) -> list[Any]: + """JIT specifications of every ``arch`` library without a prebuilt file (for + ahead-of-time warm-up through FlashInfer's ``build_jit_specs``).""" + + return cake_build.jit_specs(arch) + + +_GRID_AXES = {"grid_x": 0, "grid_y": 1, "grid_z": 2} + + +@functools.cache +def _arg_sources(name: str) -> tuple[tuple[bool, Any], ...]: + """``(grid, source)`` per entry of ``MODULES[name]["arg_plan"]``, in + order: a grid slot (``source`` is the axis, 0 / 1 / 2 for ``grid_x`` / + ``grid_y`` / ``grid_z``) or the logical name the kernel parameter binds + (the record's ``logical`` map); the registry is constant, so + :func:`run_modules` reads it once per module.""" + + record = MODULES[name] + logical = record["logical"] + return tuple( + (True, _GRID_AXES[key]) if kind == "grid" else (False, logical[key]) + for kind, key in record["arg_plan"] + ) + + +def _missing_argument(name: str, values: dict[str, Any]) -> ValueError: + """The error of a launch of module ``name`` whose ``values`` lack a + logical name it binds (the first such name, with its kernel parameter).""" + + record = MODULES[name] + logical, key = next( + (record["logical"][key], key) + for kind, key in record["arg_plan"] + if kind != "grid" and record["logical"][key] not in values + ) + return ValueError(f"module {name} needs {logical!r} (kernel parameter {key!r})") + + +@functools.cache +def _ffi_entry(name: str, device_index: int) -> Any: + """The FFI entry of module ``name`` loaded for CUDA device + ``device_index``'s architecture (:func:`load`), looked up once.""" + + module = load(name, device_arch(device_index)) + return getattr(module, MODULES[name]["ffi_entry"]) + + +@functools.cache +def _raw_stream_getter() -> Any: + """``torch._C._cuda_getCurrentRawStream``: the handle of a device's + current torch stream as an integer, without constructing a ``Stream`` + object; or the ``torch.cuda`` form on a torch build without it. Bound on + first use: this module imports torch lazily.""" + + import torch + + getter = getattr(torch._C, "_cuda_getCurrentRawStream", None) + if getter is not None: + return getter + + def fallback(device_index: int) -> int: + return torch.cuda.current_stream(device_index).cuda_stream + + return fallback + + +_FFI_DEVICES: dict[int, Any] = {} + + +def _stream_context(device_index: int) -> Any: + """The tvm-ffi environment-stream context of torch's current stream on + CUDA device ``device_index`` (a CUDA graph's capture stream during + capture): the bindings launch on that stream. The FFI device handle is + kept per index, so entering the context costs a fraction of + ``tvm_ffi.use_torch_stream()``, which builds a ``torch.cuda.Stream`` and + parses a device string on every entry.""" + + device = _FFI_DEVICES.get(device_index) + if device is None: + device = _FFI_DEVICES[device_index] = tvm_ffi.device(f"cuda:{device_index}") + return tvm_ffi.use_raw_stream(device, _raw_stream_getter()(device_index)) + + +def run_modules( + launches: Sequence[tuple[str, tuple[int, int, int] | None]], + values: dict[str, Any], + *, + device_index: int, +) -> None: + """Launch exported modules in order, back to back, with one set of logical + argument ``values``. + + ``launches`` holds ``(name, grid)`` pairs, ``grid`` as in + :func:`run_module`. All of them are issued inside one stream context of + torch's current stream on ``device_index``, each bound (grid and + arguments) right before it is issued, so a programmatically dependent + launch (the backward's reduce program) reaches the stream right behind + the program it depends on, while that one runs, and the first launch + waits for no later one's binding. ``ValueError`` names the first logical + name a module binds that ``values`` lacks. + """ + + with _stream_context(device_index): + for name, grid in launches: + if grid is None: + grid = launch_grid( + MODULES[name], + device_index=device_index, + C=values.get("C"), + S=values.get("S"), + ) + sources = _arg_sources(name) + try: + args = [ + grid[source] if is_grid else values[source] + for is_grid, source in sources + ] + except KeyError: + raise _missing_argument(name, values) from None + _ffi_entry(name, device_index)(*args) + + +def run_module( + name: str, + values: dict[str, Any], + *, + device_index: int, + grid: tuple[int, int, int] | None = None, +) -> None: + """Launch one exported module with logical argument ``values``. + + ``values`` maps the logical names in ``LOGICAL_TENSORS`` / + ``LOGICAL_SCALARS`` to tensors and Python scalars; the module record maps + its kernel parameter names onto that vocabulary. The launch grid is + ``grid`` when the caller planned it (the module's :func:`launch_grid` for + the call, as the public entry points pass it), else the module's + :func:`launch_grid` on ``device_index`` for the call's ``C`` and ``S``, + and the kernel runs on the current torch stream; a + module whose launch record enables programmatic dependent launch + (``use_pdl``) overlaps its launch with the preceding kernel of that stream + and waits for it in-kernel. + """ + + run_modules(((name, grid),), values, device_index=device_index) + + +__all__ = [ + "ARCHES", + "ARCH_BY_CAPABILITY", + "CHANNEL_MULTIPLE", + "FLAG_B_BF16", + "FLAG_DB_BF16", + "FLAG_DW_BF16", + "FLAG_W_BF16", + "GRID_KINDS", + "INT32_MAX", + "KINDS", + "LOGICAL_SCALARS", + "LOGICAL_TENSORS", + "MODULES", + "PLAN_SCALARS", + "REQUIRED_KINDS", + "ROLES", + "ROUTES", + "SELECT_KINDS", + "TILE_ALIGNED_CLASSES", + "TILE_ALIGNED_CLASS_ROUTES", + "TILE_ALIGNED_GRID", + "TILE_ALIGNED_MAX_TILES", + "TILE_ALIGNED_MIN_LAST", + "TILE_ALIGNED_MIN_S", + "TILE_ALIGNED_ROUTES", + "WIDTH", + "WORKSPACE_KINDS", + "arch_for_capability", + "backward_kinds", + "backward_route", + "channel_split", + "class_aligned_grid", + "counter_words", + "device_arch", + "forward_kinds", + "forward_route", + "forward_workspace_words", + "grid_mult", + "grid_per_sm", + "launch_grid", + "library_name", + "load", + "main_grid", + "partial_bytes", + "partial_rows", + "plan", + "record_route", + "route_applies", + "route_modules", + "route_role", + "row_may_trim", + "run_module", + "run_modules", + "select_rule", + "selected_modules", + "sm_count", + "specs_for_arch", + "tile_aligned_grid", + "workspace_layout", + "workspace_owners", + "workspace_record", +] diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/.clang-format b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/.clang-format similarity index 75% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/.clang-format rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/.clang-format index 5b709a08807..fae08c8733d 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/.clang-format +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/.clang-format @@ -1,4 +1,4 @@ -# Receipt-bound generated sources (Cake export of cake_causal_conv1d_train): the source registry attests these bytes; +# Receipt-bound generated sources (Cake export of cake_causal_conv1d_det_train): the source registry attests these bytes; # formatting them would break the receipts that bind them to the exporter. DisableFormat: true SortIncludes: Never diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu similarity index 87% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu index 6cb7025e294..1bfa86720e0 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, void* __restrict__ b, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, void* __restrict__ b, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -110,11 +110,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_y, Tensor config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 46208; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu index e4df5913c3f..83cf7d2e82f 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_08880ca12e4064bc8cdb_kernel.cu @@ -165,7 +165,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, void* __restrict__ b, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) +kernel_cake_causal_conv1d_det_train_08880ca12e4064bc8cdb(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, void* __restrict__ b, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_149ca6f53934ba630033_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_149ca6f53934ba630033_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_binding.cu index 10fbd22e6ed..d983b195409 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_149ca6f53934ba630033_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_149ca6f53934ba630033, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_149ca6f53934ba630033(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_149ca6f53934ba630033 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_149ca6f53934ba630033"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -193,7 +193,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_149ca6f53934ba630033), cudaFuncAttributeMaxDynamicSharedMemorySize, 74752), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033), cudaFuncAttributeMaxDynamicSharedMemorySize, 74752), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -220,11 +220,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 74752; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_149ca6f53934ba630033), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_149ca6f53934ba630033 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_149ca6f53934ba630033::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_149ca6f53934ba630033_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_149ca6f53934ba630033_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_kernel.cu index 0bfcfa93e5c..bf6df9b637e 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_149ca6f53934ba630033_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_149ca6f53934ba630033_kernel.cu @@ -195,7 +195,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_149ca6f53934ba630033(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_149ca6f53934ba630033(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu index 3d18bd6a9b2..1f399680381 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_170d89481a616ae977b3, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_170d89481a616ae977b3(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_170d89481a616ae977b3 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_170d89481a616ae977b3"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -190,7 +190,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_170d89481a616ae977b3), cudaFuncAttributeMaxDynamicSharedMemorySize, 99328), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3), cudaFuncAttributeMaxDynamicSharedMemorySize, 99328), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -216,11 +216,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 99328; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_170d89481a616ae977b3), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_170d89481a616ae977b3 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_170d89481a616ae977b3::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu index 8402e700bce..6c5f9f7c11f 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_170d89481a616ae977b3_kernel.cu @@ -195,7 +195,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_170d89481a616ae977b3(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_170d89481a616ae977b3(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_binding.cu similarity index 87% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_binding.cu index 40e01e8f5f8..efaec38661e 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -106,11 +106,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_y, Tensor config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 33920; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_kernel.cu similarity index 98% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_kernel.cu index 3848a4b0058..0b65b3f3500 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc_kernel.cu @@ -165,7 +165,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) +kernel_cake_causal_conv1d_det_train_176ebb3e14f7ad9c94dc(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu similarity index 85% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu index 111876d2582..557eb0be8f0 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_head, int* __restrict__ head_row_off, int S, int C, int num_docs); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_head, int* __restrict__ head_row_off, int S, int C, int num_docs); -namespace cake_binding_kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -102,11 +102,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 0; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu similarity index 98% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu index a4ef33ca959..d7f0776405d 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp_kernel.cu @@ -76,7 +76,7 @@ __device__ __forceinline__ float mul_f32_rn_noftz(float a, float b) { extern "C" { __global__ __launch_bounds__(128) void -kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_head, int* __restrict__ head_row_off, int S, int C, int num_docs) +kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_head, int* __restrict__ head_row_off, int S, int C, int num_docs) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu index 2acc21d983f..5bebdf40dca 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_binding.cu @@ -14,20 +14,20 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_1c21c181af94e1635851, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851, for the device source that defines it. Do not edit. #include #include "cake_ffi_checks.h" #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_1c21c181af94e1635851(float* __restrict__ partial, void* __restrict__ dw, void* __restrict__ db, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851(float* __restrict__ partial, void* __restrict__ dw, void* __restrict__ db, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_1c21c181af94e1635851 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_1c21c181af94e1635851"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -87,11 +87,11 @@ void Run(TensorView a_partial, TensorView a_dw, TensorView a_db, int64_t a_C, in attrs[0].val.programmaticStreamSerializationAllowed = 1; config.attrs = attrs; config.numAttrs = 1; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_1c21c181af94e1635851), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_1c21c181af94e1635851 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_1c21c181af94e1635851::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu similarity index 97% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu index 11c9d6e6ed9..f0484bb99f2 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_1c21c181af94e1635851_kernel.cu @@ -61,7 +61,7 @@ __device__ __forceinline__ int make_warp_uniform(int x) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_1c21c181af94e1635851(float* __restrict__ partial, void* __restrict__ dw, void* __restrict__ db, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags) +kernel_cake_causal_conv1d_det_train_1c21c181af94e1635851(float* __restrict__ partial, void* __restrict__ dw, void* __restrict__ db, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu index bf764eef665..5635ec1348e 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -190,7 +190,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378), cudaFuncAttributeMaxDynamicSharedMemorySize, 82944), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378), cudaFuncAttributeMaxDynamicSharedMemorySize, 82944), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -216,11 +216,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 82944; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu index 5fedfc7d8e1..64570b74892 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2675cc1cb4efc0937378_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_2675cc1cb4efc0937378(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_binding.cu index 6474d418257..df2eb2712b2 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_2d91e5f926238453c78c, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_2d91e5f926238453c78c(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_2d91e5f926238453c78c { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_2d91e5f926238453c78c"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -193,7 +193,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_2d91e5f926238453c78c), cudaFuncAttributeMaxDynamicSharedMemorySize, 54272), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c), cudaFuncAttributeMaxDynamicSharedMemorySize, 54272), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -220,11 +220,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 54272; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_2d91e5f926238453c78c), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_2d91e5f926238453c78c +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_2d91e5f926238453c78c::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_kernel.cu index 54b91d4fbd6..541b68d19e3 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_2d91e5f926238453c78c_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_2d91e5f926238453c78c(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_2d91e5f926238453c78c(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu index 10f2f50c709..634ff23424a 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -188,7 +188,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b), cudaFuncAttributeMaxDynamicSharedMemorySize, 82944), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b), cudaFuncAttributeMaxDynamicSharedMemorySize, 82944), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -213,11 +213,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 82944; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu index 08ab225c082..65cf1fcb31e 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu index 501eb9a4ecf..275e3cb1550 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_binding.cu @@ -14,20 +14,20 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea, for the device source that defines it. Do not edit. #include #include "cake_ffi_checks.h" #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea(float* __restrict__ partial, void* __restrict__ dw, void* __restrict__ db, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea(float* __restrict__ partial, void* __restrict__ dw, void* __restrict__ db, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -87,11 +87,11 @@ void Run(TensorView a_partial, TensorView a_dw, TensorView a_db, int64_t a_C, in attrs[0].val.programmaticStreamSerializationAllowed = 1; config.attrs = attrs; config.numAttrs = 1; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu similarity index 97% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu index b5476806795..49ed1dd599f 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea_kernel.cu @@ -60,7 +60,7 @@ __device__ __forceinline__ int make_warp_uniform(int x) { extern "C" { __global__ __launch_bounds__(128) void -kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea(float* __restrict__ partial, void* __restrict__ dw, void* __restrict__ db, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags) +kernel_cake_causal_conv1d_det_train_4d291f4b5c2cce6ef1ea(float* __restrict__ partial, void* __restrict__ dw, void* __restrict__ db, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu index 8303b9ddaf4..0a7e690b601 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_517710d1f0988cf570a9, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_517710d1f0988cf570a9(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_517710d1f0988cf570a9 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_517710d1f0988cf570a9"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -193,7 +193,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_517710d1f0988cf570a9), cudaFuncAttributeMaxDynamicSharedMemorySize, 87040), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9), cudaFuncAttributeMaxDynamicSharedMemorySize, 87040), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -220,11 +220,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 87040; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_517710d1f0988cf570a9), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_517710d1f0988cf570a9 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_517710d1f0988cf570a9::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu index 861d29889b3..1a769e2b655 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_517710d1f0988cf570a9_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_517710d1f0988cf570a9(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_517710d1f0988cf570a9(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_binding.cu similarity index 90% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_binding.cu index a88e0e8b19b..b064ffd656b 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_5288d132a5cc927ae729, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_5288d132a5cc927ae729(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_5288d132a5cc927ae729 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_5288d132a5cc927ae729"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -208,7 +208,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_5288d132a5cc927ae729), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -243,11 +243,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso attrs[0].val.cooperative = 1; config.attrs = attrs; config.numAttrs = 1; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_5288d132a5cc927ae729), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_5288d132a5cc927ae729 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_5288d132a5cc927ae729::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_kernel.cu index 6f32a637fe4..23d6de47619 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp_kernel.cu @@ -162,7 +162,7 @@ __device__ __forceinline__ float mul_f32_rn_noftz(float a, float b) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_5288d132a5cc927ae729(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu index 12eaa156426..ddd04d853e9 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); -namespace cake_binding_kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -108,11 +108,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_y, Tensor config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 33792; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu similarity index 98% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu index 67c8801961a..42e5fc6d9ba 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp_kernel.cu @@ -139,7 +139,7 @@ __device__ __forceinline__ float mul_f32_rn_noftz(float a, float b) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) +kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_binding.cu index 2aad3710643..4dc3f427302 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); -namespace cake_binding_kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -114,7 +114,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso static thread_local uint64_t prepared_devices = 0; const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -147,11 +147,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso attrs[0].val.cooperative = 1; config.attrs = attrs; config.numAttrs = 1; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_kernel.cu index ee4e4b81ac0..6931ec4a87f 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp_kernel.cu @@ -149,7 +149,7 @@ __device__ __forceinline__ float mul_f32_rn_noftz(float a, float b) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) +kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu similarity index 90% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu index c883a6e7339..cab2097c041 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_7485529a1755e5bd3036, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_7485529a1755e5bd3036(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_7485529a1755e5bd3036 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_7485529a1755e5bd3036"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -208,7 +208,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_7485529a1755e5bd3036), cudaFuncAttributeMaxDynamicSharedMemorySize, 82944), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp), cudaFuncAttributeMaxDynamicSharedMemorySize, 82944), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -243,11 +243,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso attrs[0].val.cooperative = 1; config.attrs = attrs; config.numAttrs = 1; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_7485529a1755e5bd3036), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_7485529a1755e5bd3036 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_7485529a1755e5bd3036::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu index 68c667c4333..9437701d148 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp_kernel.cu @@ -162,7 +162,7 @@ __device__ __forceinline__ float mul_f32_rn_noftz(float a, float b) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_7485529a1755e5bd3036(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ dw, float* __restrict__ partial, __nv_bfloat16* __restrict__ hist, int* __restrict__ hist_row_off, float* __restrict__ dpre_tail, int* __restrict__ tail_row_off, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_843c010db78c79107fd3_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_843c010db78c79107fd3_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_binding.cu index ac117208a59..1c9e1a58eee 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_843c010db78c79107fd3_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_843c010db78c79107fd3, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_843c010db78c79107fd3(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_843c010db78c79107fd3 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_843c010db78c79107fd3"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -99,7 +99,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso static thread_local uint64_t prepared_devices = 0; const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_843c010db78c79107fd3), cudaFuncAttributeMaxDynamicSharedMemorySize, 74752), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3), cudaFuncAttributeMaxDynamicSharedMemorySize, 74752), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -124,11 +124,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 74752; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_843c010db78c79107fd3), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_843c010db78c79107fd3 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_843c010db78c79107fd3::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_843c010db78c79107fd3_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_843c010db78c79107fd3_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_kernel.cu index 8ddcc94b538..b54616b3f2f 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_843c010db78c79107fd3_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_843c010db78c79107fd3_kernel.cu @@ -182,7 +182,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_843c010db78c79107fd3(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) +kernel_cake_causal_conv1d_det_train_843c010db78c79107fd3(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_binding.cu index df3bd069694..bcbaef3d864 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); -namespace cake_binding_kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -103,11 +103,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_y, Tensor config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 33920; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_kernel.cu similarity index 98% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_kernel.cu index 854a4d2e713..c46295ed425 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8665d0756b8240ecb4d0_kernel.cu @@ -165,7 +165,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) +kernel_cake_causal_conv1d_det_train_8665d0756b8240ecb4d0(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu similarity index 85% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu index 29bf0b1b08c..9048ab4a2c1 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_binding.cu @@ -14,20 +14,20 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5, for the device source that defines it. Do not edit. #include #include "cake_ffi_checks.h" #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5(float* __restrict__ partial, float* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5(float* __restrict__ partial, float* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec); -namespace cake_binding_kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -80,11 +80,11 @@ void Run(TensorView a_partial, TensorView a_dw, int64_t a_C, int64_t a_main_ctas attrs[0].val.programmaticStreamSerializationAllowed = 1; config.attrs = attrs; config.numAttrs = 1; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu similarity index 97% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu index 007bee71361..44c2965c85a 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5_kernel.cu @@ -61,7 +61,7 @@ __device__ __forceinline__ int make_warp_uniform(int x) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5(float* __restrict__ partial, float* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec) +kernel_cake_causal_conv1d_det_train_86b5ecfbd90aca4b52e5(float* __restrict__ partial, float* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_binding.cu index f3f8e1fc7f3..8f5539ef622 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -96,7 +96,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso static thread_local uint64_t prepared_devices = 0; const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -120,11 +120,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 50176; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_kernel.cu index b7952763b98..be4a41a035f 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_8b9b8db0833fd64a578f_kernel.cu @@ -187,7 +187,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) +kernel_cake_causal_conv1d_det_train_8b9b8db0833fd64a578f(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_binding.cu index 78bd938a809..cdd43975a44 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); -namespace cake_binding_kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -94,7 +94,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso static thread_local uint64_t prepared_devices = 0; const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -117,11 +117,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 50176; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_kernel.cu index ff206112f2c..a97775775f2 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9a07bcaa0223632640ee_kernel.cu @@ -187,7 +187,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) +kernel_cake_causal_conv1d_det_train_9a07bcaa0223632640ee(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_binding.cu index 4e6fb0cb456..e550696b710 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -188,7 +188,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -213,11 +213,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 66560; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_kernel.cu index 64e835a6232..7b78ef1eed3 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9bfa65707a49df6a48ce_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_9bfa65707a49df6a48ce(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu similarity index 87% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu index c49a5a9c098..f679e4fc9b3 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -106,11 +106,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_y, Tensor config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 46208; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu index 3bd6eeb651d..4e4ca845dba 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e_kernel.cu @@ -165,7 +165,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) +kernel_cake_causal_conv1d_det_train_9c1e7ee9f37e3102660e(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_binding.cu index f9bda28dd45..0b52e68def8 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -190,7 +190,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa), cudaFuncAttributeMaxDynamicSharedMemorySize, 50176), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -216,11 +216,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 50176; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_kernel.cu index 6b96f3b91f8..8e38336ad8e 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_9c8310a1d3fa9b8ccffa(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu index e25ad6f3351..4ae0a22b92a 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_a681165b3e81204c73ff, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_a681165b3e81204c73ff(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_a681165b3e81204c73ff { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_a681165b3e81204c73ff"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -193,7 +193,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_a681165b3e81204c73ff), cudaFuncAttributeMaxDynamicSharedMemorySize, 107520), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff), cudaFuncAttributeMaxDynamicSharedMemorySize, 107520), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -220,11 +220,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 107520; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_a681165b3e81204c73ff), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_a681165b3e81204c73ff +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_a681165b3e81204c73ff::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu index 9bb034b5713..c3a6c8bf660 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_a681165b3e81204c73ff_kernel.cu @@ -195,7 +195,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_a681165b3e81204c73ff(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_a681165b3e81204c73ff(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_binding.cu index 2d2bb798ac2..460b1f8d3b0 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -188,7 +188,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c), cudaFuncAttributeMaxDynamicSharedMemorySize, 82944), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c), cudaFuncAttributeMaxDynamicSharedMemorySize, 82944), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -213,11 +213,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 82944; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_kernel.cu index d63fc65e982..cc30594f2da 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_aeff6106e2822f313a0c_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_aeff6106e2822f313a0c(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_binding.cu index 3a53357c944..db559898ad6 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -188,7 +188,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -213,11 +213,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 66560; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_kernel.cu index 15c2b2c27a6..ac1701d7584 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_binding.cu index 20b6316d766..80d698775ad 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -190,7 +190,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -216,11 +216,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 66560; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_kernel.cu index 0c9d8616cfd..c47d5b9956e 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_c8027a64ba3cc38333bf_kernel.cu @@ -195,7 +195,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_c8027a64ba3cc38333bf(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_binding.cu index f8c650d90c9..e6821818c74 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -99,7 +99,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso static thread_local uint64_t prepared_devices = 0; const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4), cudaFuncAttributeMaxDynamicSharedMemorySize, 54272), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4), cudaFuncAttributeMaxDynamicSharedMemorySize, 54272), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -124,11 +124,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 54272; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_kernel.cu index 98a61c27ac0..4ae4e32011a 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4_kernel.cu @@ -187,7 +187,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) +kernel_cake_causal_conv1d_det_train_cc7a117aee9c7ae829c4(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, void* __restrict__ b, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu index d8db9e9b86c..a07f8837e2d 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_binding.cu @@ -14,20 +14,20 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e, for the device source that defines it. Do not edit. #include #include "cake_ffi_checks.h" #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e(float* __restrict__ partial, void* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e(float* __restrict__ partial, void* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -83,11 +83,11 @@ void Run(TensorView a_partial, TensorView a_dw, int64_t a_C, int64_t a_main_ctas attrs[0].val.programmaticStreamSerializationAllowed = 1; config.attrs = attrs; config.numAttrs = 1; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu similarity index 97% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu index 828cad1243d..775d27b8187 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ce49146b9fef24adb47e_kernel.cu @@ -60,7 +60,7 @@ __device__ __forceinline__ int make_warp_uniform(int x) { extern "C" { __global__ __launch_bounds__(128) void -kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e(float* __restrict__ partial, void* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags) +kernel_cake_causal_conv1d_det_train_ce49146b9fef24adb47e(float* __restrict__ partial, void* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_binding.cu similarity index 91% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_binding.cu index 7f54c46fc7b..2a21ac4d7a8 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6, for the device source that defines it. Do not edit. #include #include #include @@ -23,13 +23,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map); -namespace cake_binding_kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -188,7 +188,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -213,11 +213,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 66560; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_kernel.cu index 08228de3f23..717d7ee8d78 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6_kernel.cu @@ -200,7 +200,7 @@ __device__ __forceinline__ unsigned int __as_u32(int v) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) +kernel_cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, const __grid_constant__ CUtensorMap x_map, const __grid_constant__ CUtensorMap dy_map) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_binding.cu similarity index 87% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_binding.cu index d71fe11314d..6a8f9257324 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, void* __restrict__ b, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, void* __restrict__ b, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -110,11 +110,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_y, Tensor config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 33920; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_kernel.cu similarity index 98% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_kernel.cu index 070f5277e02..11f543ab02a 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_d829c95e1d610bc14b23_kernel.cu @@ -165,7 +165,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, void* __restrict__ b, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) +kernel_cake_causal_conv1d_det_train_d829c95e1d610bc14b23(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, void* __restrict__ b, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu index dac3b93fd8f..903de2398f5 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_binding.cu @@ -14,20 +14,20 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2, for the device source that defines it. Do not edit. #include #include "cake_ffi_checks.h" #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2(float* __restrict__ partial, void* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2(float* __restrict__ partial, void* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -83,11 +83,11 @@ void Run(TensorView a_partial, TensorView a_dw, int64_t a_C, int64_t a_main_ctas attrs[0].val.programmaticStreamSerializationAllowed = 1; config.attrs = attrs; config.numAttrs = 1; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu similarity index 97% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu index 6b3bf7360b0..d998ab555c0 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_ea8f91038a9dea452ff2_kernel.cu @@ -61,7 +61,7 @@ __device__ __forceinline__ int make_warp_uniform(int x) { extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2(float* __restrict__ partial, void* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags) +kernel_cake_causal_conv1d_det_train_ea8f91038a9dea452ff2(float* __restrict__ partial, void* __restrict__ dw, int C, int main_ctas, int n_ch_tiles, int ct_vec, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_binding.cu index d7da5f25dbd..8583643d4f5 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_f1545adfd2bb28044419, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_f1545adfd2bb28044419(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags); -namespace cake_binding_kernel_cake_causal_conv1d_train_f1545adfd2bb28044419 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_f1545adfd2bb28044419"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -96,7 +96,7 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso static thread_local uint64_t prepared_devices = 0; const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; if (device_bit == 0 || (prepared_devices & device_bit) == 0) { - CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_train_f1545adfd2bb28044419), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419), cudaFuncAttributeMaxDynamicSharedMemorySize, 66560), "cudaFuncSetAttribute"); prepared_devices |= device_bit; } } @@ -120,11 +120,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_dy, Tenso config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 66560; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_f1545adfd2bb28044419), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_f1545adfd2bb28044419 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_f1545adfd2bb28044419::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_kernel.cu index b221081b172..69389668936 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f1545adfd2bb28044419_kernel.cu @@ -182,7 +182,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_f1545adfd2bb28044419(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) +kernel_cake_causal_conv1d_det_train_f1545adfd2bb28044419(__nv_bfloat16* __restrict__ x, void* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ dy, __nv_bfloat16* __restrict__ dx, float* __restrict__ partial, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles, int flags) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu similarity index 86% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu index 5d72928b3c5..64cc71e0f5c 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_binding.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_train_f5497330bfd286211643, for the device source that defines it. Do not edit. +// Generated by Cake: the tvm-ffi host binding of kernel_cake_causal_conv1d_det_train_f5497330bfd286211643, for the device source that defines it. Do not edit. #include #include @@ -22,13 +22,13 @@ #include -extern "C" __global__ void kernel_cake_causal_conv1d_train_f5497330bfd286211643(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); +extern "C" __global__ void kernel_cake_causal_conv1d_det_train_f5497330bfd286211643(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles); -namespace cake_binding_kernel_cake_causal_conv1d_train_f5497330bfd286211643 { +namespace cake_binding_kernel_cake_causal_conv1d_det_train_f5497330bfd286211643 { using tvm::ffi::TensorView; -constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_train_f5497330bfd286211643"; +constexpr const char* KERNEL_NAME = "kernel_cake_causal_conv1d_det_train_f5497330bfd286211643"; void CheckCuda(cudaError_t status, const char* what) { TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) @@ -103,11 +103,11 @@ void Run(TensorView a_x, TensorView a_w, TensorView a_cu, TensorView a_y, Tensor config.blockDim = dim3(128, 1, 1); config.dynamicSmemBytes = 46208; config.stream = stream; - CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_train_f5497330bfd286211643), kernel_args), + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_causal_conv1d_det_train_f5497330bfd286211643), kernel_args), "cudaLaunchKernelExC"); } -} // namespace cake_binding_kernel_cake_causal_conv1d_train_f5497330bfd286211643 +} // namespace cake_binding_kernel_cake_causal_conv1d_det_train_f5497330bfd286211643 -TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_train_f5497330bfd286211643::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_causal_conv1d_det_train_f5497330bfd286211643::Run); TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu similarity index 99% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu index c1b4417cc18..bf2e1fc0f77 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu +++ b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_causal_conv1d_det_train_f5497330bfd286211643_kernel.cu @@ -165,7 +165,7 @@ __device__ __forceinline__ void cp_async_bulk_gmem2smem( extern "C" { __global__ __launch_bounds__(128, LAUNCH_MIN_BLOCKS) void -kernel_cake_causal_conv1d_train_f5497330bfd286211643(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) +kernel_cake_causal_conv1d_det_train_f5497330bfd286211643(__nv_bfloat16* __restrict__ x, float* __restrict__ w, int* __restrict__ cu, __nv_bfloat16* __restrict__ y, int* __restrict__ work, int S, int C, int num_docs, int search_steps, int ct_vec, int n_sub, int n_ch_tiles) { const int tid = threadIdx.x; const int warp = make_warp_uniform(tid / 32); diff --git a/flashinfer/cake_causal_conv1d_train_backend/csrc/cake_ffi_checks.h b/flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_ffi_checks.h similarity index 100% rename from flashinfer/cake_causal_conv1d_train_backend/csrc/cake_ffi_checks.h rename to flashinfer/cake_causal_conv1d_det_train_backend/csrc/cake_ffi_checks.h diff --git a/flashinfer/cake_causal_conv1d_det_train_backend/kcp.py b/flashinfer/cake_causal_conv1d_det_train_backend/kcp.py new file mode 100644 index 00000000000..393898bcd84 --- /dev/null +++ b/flashinfer/cake_causal_conv1d_det_train_backend/kcp.py @@ -0,0 +1,1870 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Context-parallel (KCP) halo forward / backward of the Cake packed width-4 +depthwise causal convolution + SiLU (FlashInfer issue #6919): the rank's +fragments of the packed documents are convolved with the three document tokens +before each continuing fragment, received from the ranks that hold them by one +``all_to_all_single`` per direction, and the backward sends the first rows' +derivative back to the predecessor fragments, which fold it into their last +rows' ``dx`` before the one BF16 rounding. ``cake_causal_conv1d_det_kcp(x_local, +weight, ctx)`` is the autograd entry; the stable, ``@flashinfer_api``-decorated +facade is :mod:`flashinfer.cake_causal_conv1d_det_train`. This module imports +only ``os``, torch, torch.distributed, tvm_ffi and its own package. + +Registry contract (``kcp_jit``, the KCP records beside the zero-history +registry ``cake_jit``, whose tables stay untouched): kinds ``kcp_fwd`` (one +record per arch: co-resident grid, 4 CTAs per SM, 128 threads x 8 channels), +``kcp_bwd`` (the fallback record and the ``_st8`` record selected by +``channels_at_least`` = the zero-history rule; ``workspace_rule`` +``partial_rows`` with 4 taps) and ``kcp_bwd_head`` (one record; its grid is +``max(1, n_frag x ceil(C / (threads x vec)))`` computed here and launched +through the record's ``arg_plan`` / ``logical`` / ``ffi_entry`` and ``load``); +logical tensor names are ``cake_jit.LOGICAL_TENSORS`` plus ``hist``, +``hist_row_off``, ``dpre_tail``, ``tail_row_off``, ``dpre_head``, +``head_row_off`` (``x_map`` / ``dy_map`` map to ``x`` / ``dy``); the scalars are +``cake_jit.PLAN_SCALARS`` (the KCP programs take no ``flags``). The KCP path +keeps the package's pre-round-6 contract: FP32 taps only (``_weight``'s default +dtype set), no bias, and the legacy channel tile rule ``cake_jit.plan(..., +legacy=True)`` (``ct_vec = min(C / 8, 128)``; its kernels pick the grouped path +from ``ct_vec`` on the device and were validated and measured under that rule), +so C = 1280 still plans (128, 1, 2) here while the training routes tile it as +(32, 4, 5). + +Numerics against the zero-history kernels of this package: the forward runs +the zero-history forward's instructions with history rows, so ``y`` is bitwise +equal to the no-CP ``y`` of the same documents. The backward takes the SiLU +derivative on the FP32 preactivation (the issue's requirement), whereas the +shipped zero-history backward differentiates the BF16-rounded preactivation, +so ``dx`` is not the no-CP ``dx`` bit for bit: the two differ by that +rounding's effect on ``dpre`` plus one BF16 output rounding. Measured on +GB300 (sm_103a) over the single-GPU test rows (``S x C`` up to 12000 x 8192, +random BF16 inputs, FP32 taps of scale 0.5): max ``|dx - dx_noCP|`` 0.03125 +(one BF16 ulp at magnitude 4-8), 0 elements beyond ``1e-2 + 1e-2 |dx_noCP|``, +and the same maximum between the two policies' FP32 oracles with no kernel +involved; each kernel stays within one BF16 ulp of its own oracle. ``dweight`` +is the rank's fixed-order FP32 partial; it is reproducible across launches and +within ``5e-2 + 1e-4 |dw|`` of an FP64 oracle, not bitwise the no-CP partial. + +Shipped host path (this revision; +``HOST_TRIM_DEFAULT = 1``): the context-derived host work of the eager halo +call -- the validated ``cu`` / ``hist_index`` / row-offset tables and halo +splits, the plan scalars, the grids, the record names and the loaded modules +with their positional argument plans -- is computed once per context and +reused, and the kernels launch through a raw-stream FFI context of the current +stream (no ``torch.cuda.Stream`` object per launch). ``_HOST_CACHE`` is keyed +on the context object's identity (the entry pins the context) and, inside it, +on ``(kind, S, C, device, cu_seqlens_local_cpu)`` -- host integers of the +context and the shapes of the call, never a tensor's contents; it holds at most +``HOST_CACHE_MAX = 64`` contexts (the oldest dropped) and +:func:`host_cache_clear` empties it; ``_FFI_DEVICES`` caches the FFI device per +device index. The operand checks, the gather, the exchange, its wait and the +kernel launches run as before with the same arguments on the same stream, so +no bit changes; a replayed graph (``graph="replay"``) contains no host work. +``_set_test_knobs({"host_trim": 0})`` (with ``CAKE_KCP_TEST_KNOBS=1``) selects +the previous host path for the A/B record. + +Checkpoint recompute without the halo exchange (FlashInfer issues #7815 / +#7818): with ``recompute_handle=`` (:class:`RecomputeHandleDetKcp`, or any object +with an ``entries`` dict) the original forward exchanges as always and stores +its history pool ``hist`` (the saved set's own tensor, key ``"conv"``) in the +handle; torch's recompute of that forward inside the backward binds it and +launches the same forward kernel on it -- no gather, no collective, no wait -- +so ``y`` and the saved set ``(x_local, weight, hist)`` are bitwise those of the +exchanging recompute. The backward is unchanged (it still exchanges the head +rows). The forward exchange is a collective: every rank of ``ctx.group`` passes +a handle for the same call or none, created at the same checkpoint nesting +level (no check is made; a mismatch hangs or corrupts data). + +Several convolutions of one layer (FlashInfer issue #7815): ``cake_causal_conv1d_det_multi_kcp(xs, weights, ctx)`` runs +the convolutions of ``n`` tensors that share the context (a layer's q, k and v) with one halo exchange per direction +instead of ``n``: the tensors' send rows travel side by side in one ``all_to_all_single`` and are copied out into +each tensor's own contiguous pool after the wait, and each tensor then runs the single call's kernels with the +single call's arguments, so outputs and gradients (``dweight`` included) are bitwise one single call per tensor. +One recompute handle serves the whole call (its ``"conv"`` entry holds the ``n`` pools). The multi path always +uses the trimmed host path's per-context cache (the ``host_trim`` test knob selects the single call's path only). +""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +import torch +import torch.distributed as dist +import tvm_ffi + +from . import ( + _cu_seqlens, + _device_arch, + _device_index, + _output, + _require_tensor, + _scalars, + _tokens, + _weight, +) +from . import kcp_jit as _loader +from .cake_jit import WIDTH, launch_grid, partial_rows, plan +from .kcp_jit import load, route_modules, run_module, selected_modules + +HALO = WIDTH - 1 # history rows, tail rows and head rows per fragment +HEAD_SLOTS = HALO # send slots per head row (a head token is consumed by at most three predecessor fragments) +KCP_FWD = "kcp_fwd" +KCP_BWD = "kcp_bwd" +KCP_BWD_HEAD = "kcp_bwd_head" +TEST_KNOBS_ENV = "CAKE_KCP_TEST_KNOBS" # the KCP route's private test-hook gate (the variable the KDA host module reads too) +GRAPH_MODES = ( + "off", + "replay", +) # graph replay: keyed CUDA-graph capture / replay of the halo calls; public keyword graph= since an earlier revision +GRAPH_DEFAULT = ( + "off" # shipped: the eager path; the public keyword graph="replay" opts in per call +) + + +def _ceil_div(numerator: int, denominator: int) -> int: + return -(-numerator // denominator) + + +def _table( + ctx: Any, + name: str, + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """One of the context's halo tables, checked against the kernels' ABI shape.""" + + table = _require_tensor(ctx.halo_tables[name], f"ctx.halo_tables[{name!r}]") + if tuple(table.shape) != shape or table.dtype != dtype: + raise ValueError( + f"ctx.halo_tables[{name!r}] must be {dtype} {shape}, got {table.dtype} {tuple(table.shape)}" + ) + if table.device != device or not table.is_contiguous(): + raise ValueError( + f"ctx.halo_tables[{name!r}] must be a contiguous tensor on {device}" + ) + return table + + +def _pool( + rows: int, C: int, *, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + """A fresh row pool; a pool without rows still holds one row so that the kernel has a valid pointer (no table + entry names it).""" + + return torch.empty((max(rows, 1), C), dtype=dtype, device=device) + + +def _exchange( + send: torch.Tensor, + *, + recv_rows: int, + send_splits: tuple[int, ...], + recv_splits: tuple[int, ...], + ctx: Any, +) -> tuple[torch.Tensor, Optional[Any]]: + """One ``all_to_all_single`` of ``send`` (rows in destination-rank order) into a fresh ``[recv_rows, C]`` tensor, + issued asynchronously on the current stream; with one rank the send buffer is the received pool (no collective). + Returns ``(recv, work)``; the caller waits ``work`` at the first consumer.""" + + if int(ctx.world_size) == 1: + return send, None + recv = torch.empty((recv_rows, send.shape[1]), dtype=send.dtype, device=send.device) + work = dist.all_to_all_single( + recv, + send, + output_split_sizes=list(recv_splits), + input_split_sizes=list(send_splits), + group=ctx.group, + async_op=True, + ) + return recv, work + + +def _launch_with_grid( + name: str, values: dict[str, Any], *, grid: tuple[int, int, int], device_index: int +) -> None: + """Launch record ``name`` with logical ``values`` and an explicit grid (the head program's grid depends on the + fragment and channel counts, not on the device); otherwise :func:`cake_jit.run_module`'s argument plan.""" + + record = _loader.MODULES[name] + logical = record["logical"] + grid_values = {"grid_x": grid[0], "grid_y": grid[1], "grid_z": grid[2]} + args = [] + for kind, key in record["arg_plan"]: + if kind == "grid": + args.append(grid_values[key]) + continue + logical_name = logical[key] + if logical_name not in values: + raise ValueError( + f"module {name} needs {logical_name!r} (kernel parameter {key!r})" + ) + args.append(values[logical_name]) + module = load(name, _loader.device_arch(device_index)) + with tvm_ffi.use_torch_stream(): + getattr(module, record["ffi_entry"])(*args) + + +def _wait_ready(ctx: Any) -> None: + """``ctx.wait_ready()`` of a context that has one (a context prepared on another stream orders its table uploads + on the current stream; a duck-typed context without it needs nothing).""" + wait_ready = getattr(ctx, "wait_ready", None) + if wait_ready is not None: + wait_ready() + + +def _context(ctx: Any, *, device: torch.device, T_r: int) -> tuple[torch.Tensor, int]: + """The context's local ``cu_seqlens`` (int32 device ``[n_frag + 1]``) and fragment count, checked against the + tokens the call received.""" + + if int(ctx.T_r) != T_r: + raise ValueError( + f"x_local holds {T_r} tokens but ctx was built for T_r={int(ctx.T_r)}" + ) + if int(ctx.conv_kernel_size) != WIDTH: + raise ValueError( + f"ctx was built for a width-{int(ctx.conv_kernel_size)} convolution; this kernel is width {WIDTH}" + ) + if int(ctx.n_frag) == 0: + # A rank without fragments (T_r == 0): nothing is launched, the collectives get empty splits. + if T_r != 0: + raise ValueError(f"ctx has no fragments but x_local holds {T_r} tokens") + return _require_tensor(ctx.cu_seqlens_local, "ctx.cu_seqlens_local"), 0 + cu, n_frag = _cu_seqlens(ctx.cu_seqlens_local, device=device) + if n_frag != int(ctx.n_frag): + raise ValueError( + f"ctx.cu_seqlens_local names {n_frag} fragments but ctx.n_frag is {int(ctx.n_frag)}" + ) + return cu, n_frag + + +# ---------------------------------------------------------------------------------------------------------------- +# Checkpoint recompute without the halo exchange: the recompute handle (``recompute_handle=``). +# +# The phase comes from ``torch._C._current_graph_task_id()`` (checked once: present and returning an int; without +# it a call with a handle raises): -1 outside the autograd engine (an original forward), the running graph task +# inside a backward (torch's checkpoint recompute, reentrant or not). Rules of a call with a handle: +# * no "conv" entry -> an original forward (also one inside an outer checkpoint's recompute, nested regions): +# gather + exchange as always, then store ``hist`` (no copy) and its metadata; +# * an entry inside the backward -> the recompute: bind the stored ``hist``, validate it against the call, issue +# no collective; the entry stays (``retain_graph`` / a second backward replays it again); +# * an entry outside the backward -> raise (a handle reused across micro-steps or for a second convolution). +# Graph replay (keyword or knob) and a forward inside a CUDA-graph capture are refused with a handle, in the +# autograd Function itself (``_handle_refusals``), so ``CakeCausalConv1dDetFunctionKcp.apply`` refuses like the wrapper. +# The ranks are not checked against each other: the forward exchange is a collective, so every rank of the group +# must take the same branch (all with a handle at the same nesting level, or all without). + +HANDLE_KEY = "conv" # this package's entry in a recompute handle + +_GRAPH_TASK_ID_BINDING = getattr(torch._C, "_current_graph_task_id", None) +_GRAPH_TASK_ID_USABLE: Optional[bool] = None # None = not checked yet + + +class RecomputeHandleDetKcp: + """The boundary tensors of one checkpointed call, carried from its original forward to torch's recompute of that + forward, so the recompute runs without cross-rank communication. + + Create one handle per convolution call and micro-step, outside the outermost checkpoint region, and pass it as an + argument of the checkpointed function, so the recompute sees the same object:: + + def layer(x, weight, handle): + return cake_causal_conv1d_det_kcp(x, weight, ctx, recompute_handle=handle) + + y = torch.utils.checkpoint.checkpoint(layer, x, w, RecomputeHandleDetKcp(), use_reentrant=False) + + ``entries`` maps an operation key to that operation's entry: ``"conv"`` is this package's (the forward's history + pool ``hist``, the very tensor the saved set holds); other packages that accept a handle with an ``entries`` dict + use their own keys, so the handle of a convolution may also serve the recurrence of the same layer. A handle + created inside a checkpoint region is a new, empty one in that region's recompute, which then exchanges again + (correct, not communication-free). The entries live as long as the handle; :meth:`nbytes` is the retained + size. + + The forward halo exchange is a collective. Every rank of ``ctx.group`` must pass a handle for the same call, or + no rank passes one, and every rank creates it at the same checkpoint nesting level. A mismatch pairs one rank's + recompute exchange with another rank's backward exchange: the ranks hang or exchange wrong rows. No check is + made. Never share one handle between two convolution calls (also not in nested regions). Never call + :meth:`clear` on a handle, or reuse it, while a backward that needs its entry is pending.""" + + __slots__ = ("entries",) + + def __init__(self) -> None: + self.entries: dict[str, dict[str, Any]] = {} + + def clear(self) -> None: + """Drop every entry (the retained tensors go with the last reference).""" + self.entries.clear() + + def nbytes(self) -> int: + """Bytes of the distinct tensors the entries hold.""" + seen: dict[int, int] = {} + for entry in self.entries.values(): + for t in entry.values(): + for x in t if isinstance(t, (tuple, list)) else (t,): + if isinstance(x, torch.Tensor): + seen[id(x)] = x.numel() * x.element_size() + return sum(seen.values()) + + +def _kcp_in_backward() -> bool: + """Whether this thread runs inside the autograd engine's backward pass (torch's checkpoint recompute runs there, + an original forward does not). The one place the recompute handle reads the phase. Raises without the torch + binding (checked once).""" + global _GRAPH_TASK_ID_USABLE + if _GRAPH_TASK_ID_USABLE is None: + ok = False + if _GRAPH_TASK_ID_BINDING is not None: + try: + ok = isinstance(_GRAPH_TASK_ID_BINDING(), int) + except Exception: # noqa: BLE001 - an unusable binding disables the option + ok = False + _GRAPH_TASK_ID_USABLE = ok + if not _GRAPH_TASK_ID_USABLE: + raise RuntimeError( + "recompute_handle needs torch._C._current_graph_task_id to tell the recompute from the original " + "forward; this torch does not provide it" + ) + return int(_GRAPH_TASK_ID_BINDING()) != -1 + + +def _handle_refusals(graph: Optional[str], handle: Any) -> None: + """The refusals of a call with a handle, shared by the wrapper and the autograd Function: graph replay (keyword + or test knob), a forward inside a CUDA-graph capture, a handle without an ``entries`` dict.""" + mode = _graph_mode(graph) + if mode != "off": + raise ValueError( + f"recompute_handle needs graph='off' (keyword and test knob), got {mode!r}" + ) + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise ValueError( + "recompute_handle is refused inside a CUDA-graph capture (a replay would overwrite the stored rows)" + ) + if not isinstance(getattr(handle, "entries", None), dict): + raise TypeError( + f"recompute_handle must carry an 'entries' dict (RecomputeHandleDetKcp), got {type(handle).__name__}" + ) + + +def _handle_lookup(handle: Any) -> Optional[dict[str, Any]]: + """The handle's ``"conv"`` entry when this forward is the recompute of the forward that stored it, ``None`` for an + original forward (it exchanges and stores one). An entry found outside the backward raises.""" + entries = getattr(handle, "entries", None) + if not isinstance(entries, dict): + raise TypeError( + f"recompute_handle must carry an 'entries' dict (RecomputeHandleDetKcp), got {type(handle).__name__}" + ) + in_backward = _kcp_in_backward() + entry = entries.get(HANDLE_KEY) + if entry is None: + return None + if not in_backward: + raise RuntimeError( + "recompute_handle: the handle already holds a 'conv' entry outside the backward -- a handle serves one " + "convolution call of one micro-step; create a new RecomputeHandleDetKcp per call and micro-step" + ) + return entry + + +def _handle_hist( + entry: dict[str, Any], x_local: torch.Tensor, ctx: Any +) -> torch.Tensor: + """The stored ``hist`` of ``entry``, checked against this recompute: the same context, token count and channel + count, and the stored shape / dtype / device (BF16, ``x_local``'s device).""" + hist = entry.get("hist") + if not isinstance(hist, torch.Tensor): + raise ValueError("recompute_handle: the 'conv' entry holds no hist tensor") + if entry.get("ctx") is not ctx: + raise ValueError( + "recompute_handle: the 'conv' entry belongs to another KcpContext" + ) + S = int(x_local.shape[-2]) if x_local.ndim >= 2 else -1 + C = int(x_local.shape[-1]) if x_local.ndim >= 1 else -1 + if entry.get("key") != (S, C): + raise ValueError( + f"recompute_handle: the 'conv' entry was stored for (T_r, C) = {entry.get('key')}, " + f"this recompute has {(S, C)}" + ) + if ( + tuple(hist.shape) != entry.get("shape") + or hist.dtype != entry.get("dtype") + or hist.device != entry.get("device") + or hist.dtype != torch.bfloat16 + or hist.device != x_local.device + or hist.ndim != 2 + or int(hist.shape[1]) != C + ): + raise ValueError( + f"recompute_handle: the stored hist ({hist.dtype} {tuple(hist.shape)} on {hist.device}) does not match " + f"its entry ({entry.get('dtype')} {entry.get('shape')} on {entry.get('device')}) or this call " + f"(BF16 [n_hist_rows, {C}] on {x_local.device})" + ) + return hist + + +# ---------------------------------------------------------------------------------------------------------------- +# CUDA-graph replay of a call and the test knobs. +# +# The public keyword ``graph="replay"`` of :func:`cake_causal_conv1d_det_kcp` and of the forward / backward below (shipped +# ``GRAPH_DEFAULT = "off"``) -- or, for the A/B record, the private ``_set_test_knobs({"graph": "replay"})`` (accepted +# only with ``CAKE_KCP_TEST_KNOBS=1``; a keyword contradicting an active knob is a ``ValueError``) -- makes +# :func:`cake_causal_conv1d_det_forward_kcp` / :func:`cake_causal_conv1d_det_backward_kcp` run their first call of a +# validity key (kind, the context object, its world size, and the shape / stride / dtype / data pointer of every +# operand and caller-owned output) eagerly, capture the second call into a ``torch.cuda.CUDAGraph`` (replayed once) +# and replay the graph on every later call of the key: the halo gather, the ``all_to_all_single``, its wait and the +# kernels with the same arguments in the same order (bitwise the eager result) and no host work between them. Any +# other key runs eagerly (silent fallback); a rank without tokens is never captured. A replayed call returns the +# graph's static tensors (detached): the next replay of the key overwrites them, so a framework must keep its inputs +# pointer-stable and consume the outputs before the next call of the same layer. This module keeps its own cache +# (no import of the KDA host module); ``GRAPH_STATS`` counts the events and ``LAST_CALL[kind]`` says how the last +# forward / backward ran (keyword or knob). Every captured graph keeps its fresh tensors in a private pool (one per +# cached key; at most GRAPH_CACHE_MAX keys, the oldest dropped first); ``graph_cache_clear()`` releases them and must run +# before ``torch.distributed.destroy_process_group()`` (a graph holding the NCCL collective blocks ``ncclCommDestroy``). + +# Trimmed host path; shipped default: +# 1 = the trimmed host path of the eager halo call (same launches, bitwise; the module docstring's "Shipped host path"), +# 0 = the previous host path, selectable through _set_test_knobs for the A/B record. +HOST_TRIMS = (0, 1) +HOST_TRIM_DEFAULT = ( + 1 # shipped; 0 = the previous path, via _set_test_knobs({"host_trim": 0}) only +) +_TEST_KNOBS: dict[str, Any] = {"graph": "off", "host_trim": HOST_TRIM_DEFAULT} +GRAPH_STATS: dict[str, int] = {"eager": 0, "capture": 0, "replay": 0, "fallback": 0} +GRAPH_CACHE_MAX = 64 +LAST_CALL: dict[str, dict[str, Any]] = {} + +# -- the trimmed host path (host_trim=1): the per-context host state of the eager halo calls. ``_HOST_CACHE[id(ctx)]`` holds the +# context (pinned, so its id cannot be reused) and one entry per ``(kind, S, C, device, cu_seqlens_local_cpu)``: the +# validated context tables and tensors (``cu``, ``hist_row_off``, ...), the plan scalars, the record names, the loaded +# modules with their positional argument plans and grids. Everything in an entry derives from host integers of the +# context and the shapes of the call; no tensor content is read. At most HOST_CACHE_MAX contexts (the oldest dropped). +HOST_CACHE_MAX = 64 +_HOST_CACHE: dict[int, tuple[Any, dict[tuple, dict[str, Any]]]] = {} +_FFI_DEVICES: dict[int, Any] = {} + + +def _host_cache_entry(ctx: Any, key: tuple) -> dict[str, Any]: + """The (possibly empty, to be filled) host-cache entry of ``ctx`` under ``key``.""" + hit = _HOST_CACHE.get(id(ctx)) + if hit is None or hit[0] is not ctx: + if len(_HOST_CACHE) >= HOST_CACHE_MAX: + del _HOST_CACHE[next(iter(_HOST_CACHE))] + hit = _HOST_CACHE[id(ctx)] = (ctx, {}) + entry = hit[1].get(key) + if entry is None: + entry = hit[1][key] = {} + return entry + + +def host_cache_clear() -> int: + """Drop the trimmed host path's per-context cache; returns the number of contexts dropped.""" + n = len(_HOST_CACHE) + _HOST_CACHE.clear() + return n + + +def _ffi_device(device_index: int) -> Any: + dev = _FFI_DEVICES.get(int(device_index)) + if dev is None: + dev = _FFI_DEVICES[int(device_index)] = tvm_ffi.device( + f"cuda:{int(device_index)}" + ) + return dev + + +def _raw_stream(device_index: int) -> int: + get = getattr(torch._C, "_cuda_getCurrentRawStream", None) + if get is not None: + return int(get(int(device_index))) + return int(torch.cuda.current_stream(int(device_index)).cuda_stream) + + +def _launch_plan( + name: str, grid: tuple[int, int, int], device_index: int +) -> tuple[Any, str, tuple[tuple[str | None, int], ...]]: + """``(module, ffi entry, plan)`` of record ``name`` for the trimmed launch: ``plan`` lists ``(None, grid value)`` + per grid argument and ``(logical name, 0)`` per kernel parameter, in ``run`` order (the same arguments + :func:`run_module` / :func:`_launch_with_grid` build per call).""" + record = _loader.MODULES[name] + logical = record["logical"] + grid_values = { + "grid_x": int(grid[0]), + "grid_y": int(grid[1]), + "grid_z": int(grid[2]), + } + plan: list[tuple[str | None, int]] = [] + for kind, key in record["arg_plan"]: + if kind == "grid": + plan.append((None, grid_values[key])) + else: + if key not in logical: + raise ValueError( + f"module {name}: kernel parameter {key!r} has no logical name" + ) + plan.append((logical[key], 0)) + module = load(name, _loader.device_arch(device_index)) + return module, record["ffi_entry"], tuple(plan) + + +def _launch_trim( + info: tuple[Any, str, tuple[tuple[str | None, int], ...]], + values: dict[str, Any], + device_index: int, +) -> None: + """Launch a cached :func:`_launch_plan` with the per-call ``values`` on the device's current torch stream (raw-stream + FFI context: the same stream ``tvm_ffi.use_torch_stream`` would set, without building a torch.cuda.Stream).""" + module, entry, plan = info + args = [] + for logical_name, grid_value in plan: + if logical_name is None: + args.append(grid_value) + else: + try: + args.append(values[logical_name]) + except KeyError: + raise ValueError( + f"the launch needs logical argument {logical_name!r}" + ) from None + with tvm_ffi.use_raw_stream(_ffi_device(device_index), _raw_stream(device_index)): + getattr(module, entry)(*args) + + +def _set_test_knobs(knobs: Optional[dict]) -> dict: + """Set the module's private test knobs (``graph``: ``"off"`` shipped | ``"replay"``; ``host_trim``: ``1`` shipped = + the trimmed host path of the eager call | ``0`` the previous host path); a knob + the dict does not name keeps its shipped default; ``None`` restores the defaults. Refused without + ``CAKE_KCP_TEST_KNOBS=1`` in the environment. Returns the active knobs.""" + if knobs is None: + _TEST_KNOBS["graph"] = "off" + _TEST_KNOBS["host_trim"] = HOST_TRIM_DEFAULT + return dict(_TEST_KNOBS) + if os.environ.get(TEST_KNOBS_ENV) != "1": + raise ValueError( + f"_set_test_knobs is a private test hook: it is accepted only with {TEST_KNOBS_ENV}=1 in the environment" + ) + if not isinstance(knobs, dict): + raise TypeError(f"_set_test_knobs expects a dict, got {type(knobs).__name__}") + unknown = sorted(set(knobs) - {"graph", "host_trim"}) + if unknown: + raise ValueError( + f"_set_test_knobs: unknown knob(s) {unknown}; known: ['graph', 'host_trim']" + ) + graph = str(knobs.get("graph", "off")) + if graph not in GRAPH_MODES: + raise ValueError( + f"_set_test_knobs: graph must be one of {GRAPH_MODES}, got {graph!r}" + ) + host_trim = knobs.get("host_trim", HOST_TRIM_DEFAULT) + if isinstance(host_trim, bool) or host_trim not in HOST_TRIMS: + raise ValueError( + f"_set_test_knobs: host_trim must be one of {HOST_TRIMS}, got {host_trim!r}" + ) + _TEST_KNOBS["graph"] = graph + _TEST_KNOBS["host_trim"] = int(host_trim) + return dict(_TEST_KNOBS) + + +def _graph_mode(graph: Optional[str]) -> str: + """The effective graph mode of a call: the public keyword when given, else the module's test knob (shipped + ``"off"``); validated against ``GRAPH_MODES``; a keyword contradicting an active knob is a ``ValueError``.""" + knob = _TEST_KNOBS["graph"] + if graph is None: + return knob + mode = str(graph) + if mode not in GRAPH_MODES: + raise ValueError(f"graph must be one of {GRAPH_MODES}, got {graph!r}") + if knob != "off" and mode != knob: + raise ValueError( + f"graph given twice with different values: keyword {mode!r}, _set_test_knobs {knob!r}" + ) + return mode + + +class _GraphEntry: + """One validity key: its graph (after the capture) and the static results.""" + + __slots__ = ("key", "pins", "seen", "graph", "result", "disabled", "digest") + + def __init__(self, key: tuple, pins: tuple) -> None: + self.key, self.pins = ( + key, + pins, + ) # the pins keep the keyed objects alive so their ids cannot be reused + self.seen = 0 + self.graph: Optional[torch.cuda.CUDAGraph] = None + self.result: Optional[tuple] = None + self.disabled: Optional[str] = None + self.digest = f"{key[0]}:{hash(key) & 0xFFFFFFFFFFFF:012x}" + + +_GRAPHS: dict[tuple, _GraphEntry] = {} + + +def _sig(t: Any) -> tuple: + return ( + (t.data_ptr(), tuple(t.shape), tuple(t.stride()), str(t.dtype), t.device.index) + if isinstance(t, torch.Tensor) + else ("none",) + ) + + +def _graph_entry( + kind: str, ctx: Any, tensors: tuple, mode: Optional[str] = None +) -> Optional[_GraphEntry]: + """The cache entry of this call's validity key, or ``None`` when the call is never captured (mode off, no tokens); + ``mode`` = the effective graph mode (:func:`_graph_mode`), ``None`` = the knob.""" + if ( + (_TEST_KNOBS["graph"] if mode is None else mode) == "off" + or int(getattr(ctx, "T_r", 0)) == 0 + or int(getattr(ctx, "n_frag", 0)) == 0 + ): + return None + key = (kind, id(ctx), int(ctx.world_size), tuple(_sig(t) for t in tensors)) + e = _GRAPHS.get(key) + if e is None: + while len(_GRAPHS) >= GRAPH_CACHE_MAX: + del _GRAPHS[next(iter(_GRAPHS))] + e = _GRAPHS[key] = _GraphEntry(key, (ctx,)) + return e + + +def graph_cache_clear() -> int: + """Drop every cached graph with its static buffers and private memory pool; returns the number of entries dropped. + + A graph that captured a NCCL collective (the halo ``all_to_all_single`` here) must be destroyed before the process group + it used is destroyed (``ncclCommDestroy`` waits for graph-captured work; ``torch.distributed.destroy_process_group()`` + hangs while such a graph is alive). Call this before destroying the process group when ``graph="replay"`` (keyword or test knob) was used, and + whenever the operands of a key are gone for good to release its pool. The knob and ``GRAPH_STATS`` are left as they are.""" + n = len(_GRAPHS) + _GRAPHS.clear() + return n + + +def _graph_run( + kind: str, entry: Optional[_GraphEntry], body: Any, mode: Optional[str] = None +) -> tuple: + """Run ``body()`` under the graph mode (keyword or knob): replay the key's graph, or capture this call (the key's second), or run eagerly + (the key's first call, a key whose capture failed, or a call that is never captured). Returns the results.""" + if entry is not None and entry.graph is not None: + entry.graph.replay() + GRAPH_STATS["replay"] += 1 + LAST_CALL[kind] = { + "graph": "replay", + "graph_event": "replay", + "graph_key": entry.digest, + "graph_fallback": None, + } + return entry.result + if entry is not None and entry.seen > 0 and entry.disabled is None: + g = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(g, capture_error_mode="thread_local"): + result = body() + except Exception as exc: # noqa: BLE001 - the failure is recorded; the call runs eagerly below + entry.disabled = f"{type(exc).__name__}: {str(exc)[:200]}" + GRAPH_STATS["fallback"] += 1 + else: + g.replay() + entry.graph, entry.result = g, result + entry.seen += 1 + GRAPH_STATS["capture"] += 1 + LAST_CALL[kind] = { + "graph": "replay", + "graph_event": "capture", + "graph_key": entry.digest, + "graph_fallback": None, + } + return result + result = body() + GRAPH_STATS["eager"] += 1 + if entry is not None: + entry.seen += 1 + LAST_CALL[kind] = { + "graph": _TEST_KNOBS["graph"] if mode is None else mode, + "graph_event": "eager", + "graph_key": None if entry is None else entry.digest, + "graph_fallback": None if entry is None else entry.disabled, + } + return result + + +def cake_causal_conv1d_det_forward_kcp( + x_local: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + *, + out: Optional[torch.Tensor] = None, + graph: Optional[str] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r"""The KCP forward: ``y`` of the rank's fragments, and the history pool ``hist`` it used (saved for the backward). + + ``x_local`` BF16 ``[T_r, C]`` (or ``[1, T_r, C]``) contiguous, the rank's fragments back to back in the order of + ``ctx.cu_seqlens_local``; ``weight`` FP32 ``[C, 4]``; ``ctx`` a ``KcpContext`` (its halo plan ``ctx.halo``, the + device tables ``ctx.halo_tables`` and the process group ``ctx.group``). Every rank of the group calls this + function once per layer, in the same order: the forward halo exchange is a collective (a rank without tokens + joins it with empty splits and returns empty outputs). ``hist`` is a fresh ``[n_hist_rows, C]`` BF16 tensor + (``[1, C]`` when the rank receives no row). ``graph`` (``None`` = the shipped ``GRAPH_DEFAULT`` ``"off"`` | + ``"replay"``): keyed CUDA-graph capture / replay of this call (the graph-replay option; see the section comment + above: pointer-stable operands, private pools, eager fallback, :func:`graph_cache_clear`). + """ + + mode = _graph_mode(graph) + if mode == "off": + return _kcp_forward(x_local, weight, ctx, out=out) + entry = _graph_entry("fwd", ctx, (x_local, weight, out), mode) + y, hist = _graph_run( + "fwd", entry, lambda: _kcp_forward(x_local, weight, ctx, out=out), mode + ) + return (y if out is not None else y.detach()), hist + + +def _kcp_forward( + x_local: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + *, + out: Optional[torch.Tensor] = None, + _hist: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """The forward's host flow (eagerly or inside a capture): see :func:`cake_causal_conv1d_det_forward_kcp`. ``_hist``: + the history pool a recompute handle stored (checked by :func:`_handle_hist`): no gather, no exchange, no wait; + the kernel launch is the same.""" + + x = _require_tensor(x_local, "x_local") + device = x.device + # tables of a context prepared on another stream: order them on this call's stream, on every call (the trimmed + # path's host cache is not per stream; after the first call on a stream this is one dict lookup) + _wait_ready(ctx) + if _TEST_KNOBS["host_trim"] == 1: + return _kcp_forward_trim(x, weight, ctx, out=out, device=device, _hist=_hist) + arch = _device_arch(device) + S, C, x_flat = _tokens(x, "x_local", device=device) + w = _weight(weight, C, device=device) + cu, n_frag = _context(ctx, device=device, T_r=S) + halo = ctx.halo + y = _output(out, "out", shape=tuple(x.shape), dtype=torch.bfloat16, device=device) + if _hist is None: + hist_index = _require_tensor( + ctx.halo_tables["hist_index"], "ctx.halo_tables['hist_index']" + ) + send = ( + x_flat.index_select(0, hist_index) + if hist_index.numel() + else x_flat.new_empty((0, C)) + ) + hist, work = _exchange( + send, + recv_rows=int(halo.n_hist_rows), + send_splits=tuple(halo.fwd_send_rows), + recv_splits=tuple(halo.fwd_recv_rows), + ctx=ctx, + ) + else: + hist, work = _hist, None + if S == 0: + if work is not None: + work.wait() + return y, hist + hist_row_off = _table( + ctx, "hist_row_off", shape=(n_frag, HALO), dtype=torch.int32, device=device + ) + (name,) = route_modules(arch, KCP_FWD) + device_index = _device_index(device) + record = _loader.MODULES[name] + scalars = plan(S, C, n_frag, geometry=record["grid_rule"]["geometry"], legacy=True) + if scalars["n_ch_tiles"] > launch_grid(record, device_index=device_index)[0]: + raise ValueError( + f"C={C} spans more channel tiles than the forward's grid on this device holds" + ) + values: dict[str, Any] = { + "x": x_flat, + "w": w, + "cu": cu, + "y": y.view(S, C), + "hist": hist + if hist.shape[0] + else _pool(0, C, dtype=torch.bfloat16, device=device), + "hist_row_off": hist_row_off, + **_scalars(scalars), + } + if work is not None: + work.wait() + run_module(name, values, device_index=device_index) + return y, hist + + +def _fwd_host( + x: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + *, + out: Optional[torch.Tensor], + device: torch.device, +) -> tuple[int, int, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: + """The trimmed forward's operand checks, output and host-cache entry: ``(S, C, x_flat, w, y, entry)``. The entry + is cached per ``(ctx, S, C, device, cu_seqlens_local_cpu)``: the validated ``cu`` / ``hist_index`` / + ``hist_row_off`` tables, the halo splits, the record, its plan scalars and grid, the loaded module with its + argument plan.""" + + S, C, x_flat = _tokens(x, "x_local", device=device) + w = _weight(weight, C, device=device) + y = _output(out, "out", shape=tuple(x.shape), dtype=torch.bfloat16, device=device) + e = _host_cache_entry( + ctx, ("fwd", S, C, str(device), tuple(ctx.cu_seqlens_local_cpu)) + ) + if not e: + arch = _device_arch(device) + cu, n_frag = _context(ctx, device=device, T_r=S) + halo = ctx.halo + e.update( + cu=cu, + n_frag=n_frag, + hist_index=_require_tensor( + ctx.halo_tables["hist_index"], "ctx.halo_tables['hist_index']" + ), + recv_rows=int(halo.n_hist_rows), + send_splits=tuple(halo.fwd_send_rows), + recv_splits=tuple(halo.fwd_recv_rows), + ) + if S != 0: + e["hist_row_off"] = _table( + ctx, + "hist_row_off", + shape=(n_frag, HALO), + dtype=torch.int32, + device=device, + ) + (name,) = route_modules(arch, KCP_FWD) + device_index = _device_index(device) + record = _loader.MODULES[name] + scalars = plan( + S, C, n_frag, geometry=record["grid_rule"]["geometry"], legacy=True + ) + grid = launch_grid(record, device_index=device_index) + if scalars["n_ch_tiles"] > grid[0]: + raise ValueError( + f"C={C} spans more channel tiles than the forward's grid on this device holds" + ) + e["scalars"] = _scalars(scalars) + e["device_index"] = device_index + e["launch"] = _launch_plan(name, grid, device_index) + return S, C, x_flat, w, y, e + + +def _fwd_gather(x_flat: torch.Tensor, C: int, e: dict[str, Any]) -> torch.Tensor: + """The forward's send rows: the context's ``hist_index`` rows of ``x_flat`` in destination-rank order.""" + hist_index = e["hist_index"] + return ( + x_flat.index_select(0, hist_index) + if hist_index.numel() + else x_flat.new_empty((0, C)) + ) + + +def _fwd_launch( + e: dict[str, Any], + x_flat: torch.Tensor, + w: torch.Tensor, + y: torch.Tensor, + hist: torch.Tensor, + work: Optional[Any], + *, + S: int, + C: int, + device: torch.device, +) -> None: + """Wait ``work`` (when given) and launch the forward kernel of a call with ``S > 0`` on ``hist``.""" + values: dict[str, Any] = { + "x": x_flat, + "w": w, + "cu": e["cu"], + "y": y.view(S, C), + "hist": hist + if hist.shape[0] + else _pool(0, C, dtype=torch.bfloat16, device=device), + "hist_row_off": e["hist_row_off"], + **e["scalars"], + } + if work is not None: + work.wait() + _launch_trim(e["launch"], values, e["device_index"]) + + +def _kcp_forward_trim( + x: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + *, + out: Optional[torch.Tensor], + device: torch.device, + _hist: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Trimmed host path: :func:`_kcp_forward` with the context-derived host work cached per ``(ctx, S, C, device, + cu_seqlens_local_cpu)`` (:func:`_fwd_host`). The operand checks, the gather, the exchange, the wait and the + kernel launch are the untrimmed call's, in the same order on the same stream.""" + + S, C, x_flat, w, y, e = _fwd_host(x, weight, ctx, out=out, device=device) + if _hist is None: + hist, work = _exchange( + _fwd_gather(x_flat, C, e), + recv_rows=e["recv_rows"], + send_splits=e["send_splits"], + recv_splits=e["recv_splits"], + ctx=ctx, + ) + else: + hist, work = _hist, None + if S == 0: + if work is not None: + work.wait() + return y, hist + _fwd_launch(e, x_flat, w, y, hist, work, S=S, C=C, device=device) + return y, hist + + +def cake_causal_conv1d_det_backward_kcp( + x_local: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + dy: torch.Tensor, + hist: torch.Tensor, + *, + dx: Optional[torch.Tensor] = None, + dweight: Optional[torch.Tensor] = None, + graph: Optional[str] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r"""The KCP backward: ``dx`` of the rank's fragments (the halo contributions of the successor fragments folded in + before the one BF16 rounding) and the rank's FP32 ``dweight`` partial (fixed order; its sum over the ranks is the + framework's). + + ``hist`` is the forward's pool. The head program writes the derivative rows of every continuing fragment's first + tokens into the backward exchange's send layout, the exchange delivers them to the predecessors as ``dpre_tail``, + and the main program folds them. Every rank calls this function once per layer, in the same order. ``graph`` as + in the forward (the backward's key holds the forward's ``hist`` pointer, so its capture follows the forward's by + one call). + """ + + mode = _graph_mode(graph) + if mode == "off": + return _kcp_backward(x_local, weight, ctx, dy, hist, dx=dx, dweight=dweight) + entry = _graph_entry("bwd", ctx, (x_local, weight, dy, hist, dx, dweight), mode) + dx_out, dw_out = _graph_run( + "bwd", + entry, + lambda: _kcp_backward(x_local, weight, ctx, dy, hist, dx=dx, dweight=dweight), + mode, + ) + return (dx_out if dx is not None else dx_out.detach()), ( + dw_out if dweight is not None else dw_out.detach() + ) + + +def _kcp_backward( + x_local: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + dy: torch.Tensor, + hist: torch.Tensor, + *, + dx: Optional[torch.Tensor] = None, + dweight: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """The backward's host flow (eagerly or inside a capture): see :func:`cake_causal_conv1d_det_backward_kcp`.""" + + x = _require_tensor(x_local, "x_local") + device = x.device + # as in the forward: every call, before the host-cache lookup of the trimmed path + _wait_ready(ctx) + if _TEST_KNOBS["host_trim"] == 1: + return _kcp_backward_trim( + x, weight, ctx, dy, hist, dx=dx, dweight=dweight, device=device + ) + arch = _device_arch(device) + S, C, x_flat = _tokens(x, "x_local", device=device) + w = _weight(weight, C, device=device) + _, _, dy_flat = _tokens(dy, "dy", device=device, shape=tuple(x.shape)) + cu, n_frag = _context(ctx, device=device, T_r=S) + halo = ctx.halo + hist_t = _require_tensor(hist, "hist") + if ( + hist_t.ndim != 2 + or int(hist_t.shape[1]) != C + or hist_t.dtype != torch.bfloat16 + or not hist_t.is_contiguous() + ): + raise ValueError( + f"hist must be the forward's contiguous BF16 [n_hist_rows, {C}] pool, got {hist_t.dtype} {tuple(hist_t.shape)}" + ) + if hist_t.shape[0] == 0: + # A valid pointer for the kernels; no table entry names a row of it. + hist_t = _pool(0, C, dtype=torch.bfloat16, device=device) + dx_out = _output( + dx, "dx", shape=tuple(x.shape), dtype=torch.bfloat16, device=device + ) + dw_out = _output( + dweight, "dweight", shape=(C, WIDTH), dtype=torch.float32, device=device + ) + device_index = _device_index(device) + n_head = int(halo.n_head_rows) + dpre_head = _pool(n_head, C, dtype=torch.float32, device=device) + if S > 0 and n_head > 0: + hist_row_off = _table( + ctx, "hist_row_off", shape=(n_frag, HALO), dtype=torch.int32, device=device + ) + head_row_off = _table( + ctx, + "head_row_off", + shape=(n_frag, HALO, HEAD_SLOTS), + dtype=torch.int32, + device=device, + ) + (head_name,) = route_modules(arch, KCP_BWD_HEAD) + geometry = _loader.MODULES[head_name]["grid_rule"]["geometry"] + channels_per_cta = int(geometry["threads"]) * int(geometry["vec"]) + grid = (max(1, n_frag * _ceil_div(C, channels_per_cta)), 1, 1) + _launch_with_grid( + head_name, + { + "x": x_flat, + "w": w, + "cu": cu, + "dy": dy_flat, + "hist": hist_t, + "hist_row_off": hist_row_off, + "dpre_head": dpre_head, + "head_row_off": head_row_off, + "S": S, + "C": C, + "num_docs": n_frag, + }, + grid=grid, + device_index=device_index, + ) + dpre_tail, work = _exchange( + dpre_head[:n_head], + recv_rows=int(halo.n_tail_rows), + send_splits=tuple(halo.bwd_send_rows), + recv_splits=tuple(halo.bwd_recv_rows), + ctx=ctx, + ) + if S == 0: + if work is not None: + work.wait() + dw_out.zero_() + return dx_out, dw_out + hist_row_off = _table( + ctx, "hist_row_off", shape=(n_frag, HALO), dtype=torch.int32, device=device + ) + tail_row_off = _table( + ctx, "tail_row_off", shape=(n_frag, HALO), dtype=torch.int32, device=device + ) + names = selected_modules(arch, KCP_BWD, C=C, S=S) + record = _loader.MODULES[names[0]] + scalars = plan(S, C, n_frag, geometry=record["grid_rule"]["geometry"], legacy=True) + grid_x = launch_grid(record, device_index=device_index)[0] + if scalars["n_ch_tiles"] > grid_x: + raise ValueError( + f"C={C} spans more channel tiles than the backward's grid on this device holds" + ) + rows = partial_rows(grid_x, scalars["n_ch_tiles"]) + partial = torch.empty((rows, C, WIDTH), dtype=torch.float32, device=device) + values: dict[str, Any] = { + "x": x_flat, + "w": w, + "cu": cu, + "dy": dy_flat, + "dx": dx_out.view(S, C), + "dw": dw_out, + "partial": partial, + "hist": hist_t, + "hist_row_off": hist_row_off, + "dpre_tail": dpre_tail + if dpre_tail.shape[0] + else _pool(0, C, dtype=torch.float32, device=device), + "tail_row_off": tail_row_off, + **_scalars(scalars), + } + if work is not None: + work.wait() + for name in names: + run_module(name, values, device_index=device_index) + return dx_out, dw_out + + +def _bwd_host( + x: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + dy: torch.Tensor, + hist: torch.Tensor, + *, + dx: Optional[torch.Tensor], + dweight: Optional[torch.Tensor], + device: torch.device, +) -> tuple[Any, ...]: + """The trimmed backward's operand checks, outputs and host-cache entry: ``(S, C, x_flat, w, dy_flat, hist_t, + dx_out, dw_out, entry)``. The entry is cached per ``(ctx, S, C, device, cu_seqlens_local_cpu)``: the tables, the + head grid, the backward records, plan scalars, partial rows, the loaded modules with their argument plans.""" + + S, C, x_flat = _tokens(x, "x_local", device=device) + w = _weight(weight, C, device=device) + _, _, dy_flat = _tokens(dy, "dy", device=device, shape=tuple(x.shape)) + hist_t = _require_tensor(hist, "hist") + if ( + hist_t.ndim != 2 + or int(hist_t.shape[1]) != C + or hist_t.dtype != torch.bfloat16 + or not hist_t.is_contiguous() + ): + raise ValueError( + f"hist must be the forward's contiguous BF16 [n_hist_rows, {C}] pool, got {hist_t.dtype} {tuple(hist_t.shape)}" + ) + if hist_t.shape[0] == 0: + hist_t = _pool(0, C, dtype=torch.bfloat16, device=device) + dx_out = _output( + dx, "dx", shape=tuple(x.shape), dtype=torch.bfloat16, device=device + ) + dw_out = _output( + dweight, "dweight", shape=(C, WIDTH), dtype=torch.float32, device=device + ) + e = _host_cache_entry( + ctx, ("bwd", S, C, str(device), tuple(ctx.cu_seqlens_local_cpu)) + ) + if not e: + arch = _device_arch(device) + cu, n_frag = _context(ctx, device=device, T_r=S) + halo = ctx.halo + device_index = _device_index(device) + n_head = int(halo.n_head_rows) + e.update( + cu=cu, + n_frag=n_frag, + device_index=device_index, + n_head=n_head, + tail_rows=int(halo.n_tail_rows), + send_splits=tuple(halo.bwd_send_rows), + recv_splits=tuple(halo.bwd_recv_rows), + head=None, + ) + if S > 0 and n_head > 0: + hist_row_off = _table( + ctx, + "hist_row_off", + shape=(n_frag, HALO), + dtype=torch.int32, + device=device, + ) + head_row_off = _table( + ctx, + "head_row_off", + shape=(n_frag, HALO, HEAD_SLOTS), + dtype=torch.int32, + device=device, + ) + (head_name,) = route_modules(arch, KCP_BWD_HEAD) + geometry = _loader.MODULES[head_name]["grid_rule"]["geometry"] + channels_per_cta = int(geometry["threads"]) * int(geometry["vec"]) + grid = (max(1, n_frag * _ceil_div(C, channels_per_cta)), 1, 1) + e["head"] = ( + _launch_plan(head_name, grid, device_index), + hist_row_off, + head_row_off, + ) + if S != 0: + e["hist_row_off"] = _table( + ctx, + "hist_row_off", + shape=(n_frag, HALO), + dtype=torch.int32, + device=device, + ) + e["tail_row_off"] = _table( + ctx, + "tail_row_off", + shape=(n_frag, HALO), + dtype=torch.int32, + device=device, + ) + names = selected_modules(arch, KCP_BWD, C=C, S=S) + record = _loader.MODULES[names[0]] + scalars = plan( + S, C, n_frag, geometry=record["grid_rule"]["geometry"], legacy=True + ) + grid_x = launch_grid(record, device_index=device_index)[0] + if scalars["n_ch_tiles"] > grid_x: + raise ValueError( + f"C={C} spans more channel tiles than the backward's grid on this device holds" + ) + e["rows"] = partial_rows(grid_x, scalars["n_ch_tiles"]) + e["scalars"] = _scalars(scalars) + e["launches"] = tuple( + _launch_plan( + name, + launch_grid(_loader.MODULES[name], device_index=device_index), + device_index, + ) + for name in names + ) + return S, C, x_flat, w, dy_flat, hist_t, dx_out, dw_out, e + + +def _bwd_head( + e: dict[str, Any], + x_flat: torch.Tensor, + w: torch.Tensor, + dy_flat: torch.Tensor, + hist_t: torch.Tensor, + *, + S: int, + C: int, + device: torch.device, +) -> torch.Tensor: + """The head program of the backward (when the rank has head rows): the fresh FP32 ``dpre_head`` pool (at least one + row) whose first ``n_head`` rows are the backward exchange's send rows.""" + dpre_head = _pool(e["n_head"], C, dtype=torch.float32, device=device) + if e["head"] is not None: + info, hist_row_off, head_row_off = e["head"] + _launch_trim( + info, + { + "x": x_flat, + "w": w, + "cu": e["cu"], + "dy": dy_flat, + "hist": hist_t, + "hist_row_off": hist_row_off, + "dpre_head": dpre_head, + "head_row_off": head_row_off, + "S": S, + "C": C, + "num_docs": e["n_frag"], + }, + e["device_index"], + ) + return dpre_head + + +def _bwd_main( + e: dict[str, Any], + x_flat: torch.Tensor, + w: torch.Tensor, + dy_flat: torch.Tensor, + hist_t: torch.Tensor, + dx_out: torch.Tensor, + dw_out: torch.Tensor, + dpre_tail: torch.Tensor, + work: Optional[Any], + *, + S: int, + C: int, + device: torch.device, +) -> None: + """Wait ``work`` (when given) and launch the main backward programs of a call with ``S > 0`` (``dx`` with the + received ``dpre_tail`` rows folded in, and the rank's fixed-order FP32 ``dweight`` partial).""" + partial = torch.empty((e["rows"], C, WIDTH), dtype=torch.float32, device=device) + values: dict[str, Any] = { + "x": x_flat, + "w": w, + "cu": e["cu"], + "dy": dy_flat, + "dx": dx_out.view(S, C), + "dw": dw_out, + "partial": partial, + "hist": hist_t, + "hist_row_off": e["hist_row_off"], + "dpre_tail": dpre_tail + if dpre_tail.shape[0] + else _pool(0, C, dtype=torch.float32, device=device), + "tail_row_off": e["tail_row_off"], + **e["scalars"], + } + if work is not None: + work.wait() + for info in e["launches"]: + _launch_trim(info, values, e["device_index"]) + + +def _kcp_backward_trim( + x: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + dy: torch.Tensor, + hist: torch.Tensor, + *, + dx: Optional[torch.Tensor], + dweight: Optional[torch.Tensor], + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Trimmed host path: :func:`_kcp_backward` with the context-derived host work cached (:func:`_bwd_host`); the + operand checks, allocations, exchange, wait and launches are the untrimmed call's.""" + + S, C, x_flat, w, dy_flat, hist_t, dx_out, dw_out, e = _bwd_host( + x, weight, ctx, dy, hist, dx=dx, dweight=dweight, device=device + ) + dpre_head = _bwd_head(e, x_flat, w, dy_flat, hist_t, S=S, C=C, device=device) + dpre_tail, work = _exchange( + dpre_head[: e["n_head"]], + recv_rows=e["tail_rows"], + send_splits=e["send_splits"], + recv_splits=e["recv_splits"], + ctx=ctx, + ) + if S == 0: + if work is not None: + work.wait() + dw_out.zero_() + return dx_out, dw_out + _bwd_main( + e, + x_flat, + w, + dy_flat, + hist_t, + dx_out, + dw_out, + dpre_tail, + work, + S=S, + C=C, + device=device, + ) + return dx_out, dw_out + + +class CakeCausalConv1dDetFunctionKcp(torch.autograd.Function): + """Autograd wrapper: the forward saves ``(x_local, weight)`` and the history pool ``hist`` (no second exchange in + the backward); the backward runs the head program, the backward halo exchange and the main backward program. + ``handle``: a recompute handle (see :class:`RecomputeHandleDetKcp`): the original forward stores ``hist`` in it, the + recompute binds it and exchanges nothing; the saved set is the same three tensors either way. A handle is + refused with graph replay (keyword or knob) and inside a CUDA-graph capture.""" + + @staticmethod + def forward(fn_ctx, x_local, weight, ctx, graph=None, handle=None): + if handle is None: + y, hist = cake_causal_conv1d_det_forward_kcp( + x_local, weight, ctx, graph=graph + ) + else: + _handle_refusals(graph, handle) + entry = _handle_lookup(handle) + if entry is None: + y, hist = _kcp_forward(x_local, weight, ctx) + handle.entries[HANDLE_KEY] = { + "hist": hist, + "shape": tuple(hist.shape), + "dtype": hist.dtype, + "device": hist.device, + "ctx": ctx, + "key": (int(x_local.shape[-2]), int(x_local.shape[-1])), + } + LAST_CALL["recompute_handle"] = {"fwd": "capture"} + else: + hist = _handle_hist(entry, x_local, ctx) + y, hist = _kcp_forward(x_local, weight, ctx, _hist=hist) + LAST_CALL["recompute_handle"] = {"fwd": "replay"} + fn_ctx.save_for_backward(x_local, weight, hist) + fn_ctx.kcp = ctx + fn_ctx.graph = graph # the backward runs under the same graph mode + return y + + @staticmethod + def backward(fn_ctx, grad_y): + x_local, weight, hist = fn_ctx.saved_tensors + if grad_y is None: + grad_y = torch.zeros_like(x_local) + dy = grad_y.to(torch.bfloat16).contiguous() + dx, dweight = cake_causal_conv1d_det_backward_kcp( + x_local, weight, fn_ctx.kcp, dy, hist, graph=fn_ctx.graph + ) + return dx, dweight, None, None, None + + +def cake_causal_conv1d_det_kcp( + x_local: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + *, + graph: Optional[str] = None, + recompute_handle: Optional[Any] = None, +) -> torch.Tensor: + """Autograd-enabled KCP packed width-4 depthwise causal convolution + SiLU of the rank's fragments. + + Parameters + ---------- + x_local : torch.Tensor + BF16 ``[T_r, C]`` (or ``[1, T_r, C]``) contiguous activations of the rank's fragments, back to back in the + order of ``ctx.cu_seqlens_local``. + weight : torch.Tensor + FP32 ``[C, 4]`` contiguous taps; ``weight[c, 3]`` multiplies the current token. + ctx : KcpContext + The rank's context (``flashinfer.experimental.cake_kda_train.build_kcp_context``): fragment table, halo + plan, process group. Every rank of the group calls this function once per layer, in the same order. + graph : str, optional + ``None`` (the shipped default ``GRAPH_DEFAULT`` = ``"off"``, the eager path) or ``"replay"``: keyed CUDA-graph + capture / replay of this call. The first call of a validity key (this + context, its world size, and the shape / stride / dtype / data pointer of ``x_local`` and ``weight``) runs + eagerly, the second is captured, every later one replays the graph (bitwise the eager result). Requirements: + pointer-stable operands and the same ``KcpContext`` per layer, outputs consumed before the next call of the + same layer (a replay overwrites the graph's static tensors); any other key runs eagerly (``GRAPH_STATS``, + ``LAST_CALL``). Each cached key pins a private memory pool (at most ``GRAPH_CACHE_MAX`` keys); call + :func:`graph_cache_clear` before ``torch.distributed.destroy_process_group()``. + recompute_handle : RecomputeHandleDetKcp, optional + ``None`` (default): unchanged. A handle (any object with an ``entries`` dict): the checkpoint recompute of + this call exchanges nothing -- the original forward stores its history pool under ``"conv"`` and torch's + recompute inside the backward binds it (bitwise the exchanging recompute); the backward still exchanges. + One handle per convolution call and micro-step, created outside the outermost checkpoint region; a forward + outside the backward that finds an entry raises. Refused with ``graph="replay"`` (keyword or knob) and + inside a CUDA-graph capture. Retains ``hist`` (``recompute_handle.nbytes()``) while the handle lives. + The forward halo exchange is a collective: every rank of ``ctx.group`` passes a handle for the same call or + none, created at the same checkpoint nesting level. A mismatch pairs one rank's recompute exchange with + another rank's backward exchange and hangs or corrupts data; no check is made. Never share one handle + between two convolution calls (also not in nested regions), and never ``clear()`` or reuse a handle while a + backward that needs it is pending. + + Returns + ------- + torch.Tensor + BF16 ``y`` of ``x_local``'s shape, with the three document tokens before every continuing fragment as + history (zero history at a document's start). The backward returns BF16 ``dx`` (the successor fragments' + contributions folded in before the one rounding) and the rank's FP32 ``dweight`` partial. + """ + + _graph_mode(graph) # validate on the caller's frame + if recompute_handle is not None: + _handle_refusals(graph, recompute_handle) + # Always five inputs: the backward returns five gradients. + return CakeCausalConv1dDetFunctionKcp.apply( + x_local, weight, ctx, graph, recompute_handle + ) + + +# ---------------------------------------------------------------------------------------------------------------- +# Several convolutions of one layer with one halo exchange per direction (FlashInfer issue #7815): the q / k / v +# convolutions of a layer share the context, so their history rows (forward) and head rows (backward) travel in ONE +# ``all_to_all_single`` instead of one per tensor. The send rows of the tensors are packed side by side +# (``[n_rows, sum C_i]``, the context's splits unchanged), and after the wait each tensor's rows are copied out into +# a fresh contiguous pool -- the forward's ``hist`` ``[n_hist_rows, C_i]`` (the saved tensor) and the backward's +# ``dpre_tail`` ``[n_tail_rows, C_i]`` -- before that tensor's kernels launch with the single call's arguments. The +# exchange and the unpacking are copies, and each tensor keeps its own kernels, partial workspace and fixed +# reduction order, so every output and gradient (``dweight`` included) is bitwise that of one single call per tensor. +# One rank (no collective): the per-tensor send rows are the pools, exactly as in the single call. + + +def _multi_operands( + xs: Any, weights: Any, outs: Any, name: str = "xs" +) -> tuple[list[torch.Tensor], list[torch.Tensor], list[Optional[torch.Tensor]]]: + """``xs`` / ``weights`` (/ ``outs``) as lists of one length >= 1, on one device.""" + if isinstance(xs, torch.Tensor) or isinstance(weights, torch.Tensor): + raise TypeError( + f"{name} and weights must be sequences of tensors (one per convolution), not a tensor" + ) + xs, weights = list(xs), list(weights) + if not xs or len(xs) != len(weights): + raise ValueError( + f"{name} and weights must hold the same number (>= 1) of tensors, got {len(xs)} and {len(weights)}" + ) + if outs is None: + outs = [None] * len(xs) + else: + outs = list(outs) + if len(outs) != len(xs): + raise ValueError( + f"outs must hold one tensor per convolution ({len(xs)}), got {len(outs)}" + ) + x0 = _require_tensor(xs[0], f"{name}[0]") + for i, x in enumerate(xs): + if _require_tensor(x, f"{name}[{i}]").device != x0.device: + raise ValueError( + f"{name}[{i}] is on {x.device}, {name}[0] on {x0.device}: one device" + ) + return xs, weights, outs + + +def _unpack(recv: torch.Tensor, widths: list[int]) -> list[torch.Tensor]: + """The tensors' contiguous column blocks of a packed receive pool (one copy each).""" + blocks, off = [], 0 + for C in widths: + blocks.append(recv[:, off : off + C].contiguous()) + off += C + return blocks + + +def cake_causal_conv1d_det_multi_forward_kcp( + xs: Any, + weights: Any, + ctx: Any, + *, + outs: Optional[Any] = None, + graph: Optional[str] = None, +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + r"""The KCP forward of several convolutions of one layer with one halo exchange: ``(ys, hists)``, bitwise + ``cake_causal_conv1d_det_forward_kcp(xs[i], weights[i], ctx)`` per tensor. + + ``xs`` BF16 ``[T_r, C_i]`` (or ``[1, T_r, C_i]``) contiguous tensors of the rank's fragments (the channel counts + may differ); ``weights`` FP32 ``[C_i, 4]``; ``outs`` optional preallocated outputs. Every rank of ``ctx.group`` + calls this function once per layer, in the same order, with the same number of tensors and the same channel + counts: the one forward exchange is a collective. ``hists[i]`` is a fresh contiguous ``[n_hist_rows, C_i]`` + BF16 pool (the saved tensor of tensor ``i``). ``graph`` as in :func:`cake_causal_conv1d_det_forward_kcp` (the + validity key holds every operand). + """ + + xs, weights, outs = _multi_operands(xs, weights, outs) + mode = _graph_mode(graph) + if mode == "off": + return _kcp_forward_multi(xs, weights, ctx, outs=outs) + entry = _graph_entry("fwd_multi", ctx, (*xs, *weights, *outs), mode) + ys, hists = _graph_run( + "fwd_multi", + entry, + lambda: _kcp_forward_multi(xs, weights, ctx, outs=outs), + mode, + ) + ys = tuple( + y if o is not None else y.detach() for y, o in zip(ys, outs, strict=True) + ) + return ys, hists + + +def _kcp_forward_multi( + xs: list[torch.Tensor], + weights: list[torch.Tensor], + ctx: Any, + *, + outs: list[Optional[torch.Tensor]], + _hists: Optional[tuple[torch.Tensor, ...]] = None, +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + """The multi forward's host flow (eagerly or inside a capture): per-tensor operand checks, the packed gather, one + exchange, one wait, the per-tensor unpack copies and the single call's forward kernel per tensor. ``_hists``: + the pools a recompute handle stored (checked by :func:`_handle_hists`): no gather, no exchange, no wait.""" + + device = _require_tensor(xs[0], "xs[0]").device + _wait_ready(ctx) + calls = [ + _fwd_host(x, w, ctx, out=o, device=device) + for x, w, o in zip(xs, weights, outs, strict=True) + ] + S = calls[0][0] + widths = [c[1] for c in calls] + work = None + if _hists is not None: + hists = list(_hists) + elif int(ctx.world_size) == 1: + # one rank: each tensor's send rows are its pool (the single call's identity exchange) + hists = [_fwd_gather(c[2], c[1], c[5]) for c in calls] + else: + e0 = calls[0][5] + sends = [_fwd_gather(c[2], c[1], c[5]) for c in calls] + send = torch.cat(sends, dim=1) if len(sends) > 1 else sends[0] + recv, work = _exchange( + send, + recv_rows=e0["recv_rows"], + send_splits=e0["send_splits"], + recv_splits=e0["recv_splits"], + ctx=ctx, + ) + if work is not None: + work.wait() + hists = _unpack(recv, widths) if len(calls) > 1 else [recv] + ys = tuple(c[4] for c in calls) + if S == 0: + return ys, tuple(hists) + for (S_i, C, x_flat, w, y, e), hist in zip(calls, hists, strict=True): + _fwd_launch(e, x_flat, w, y, hist, None, S=S_i, C=C, device=device) + return ys, tuple(hists) + + +def cake_causal_conv1d_det_multi_backward_kcp( + xs: Any, + weights: Any, + ctx: Any, + dys: Any, + hists: Any, + *, + dxs: Optional[Any] = None, + dweights: Optional[Any] = None, + graph: Optional[str] = None, +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + r"""The KCP backward of several convolutions of one layer with one halo exchange: ``(dxs, dweights)``, bitwise + ``cake_causal_conv1d_det_backward_kcp(xs[i], weights[i], ctx, dys[i], hists[i])`` per tensor. + + The head programs of every tensor run first, their send rows travel packed in one exchange, and after the wait + each tensor's received rows are copied out and folded by that tensor's main backward programs (own partial + workspace, fixed order). ``hists`` are the forward's pools. Every rank calls this function once per layer, in + the same order. ``graph`` as in :func:`cake_causal_conv1d_det_backward_kcp`. + """ + + xs, weights, _ = _multi_operands(xs, weights, None) + n = len(xs) + dys, hists = list(dys), list(hists) + dxs = [None] * n if dxs is None else list(dxs) + dweights = [None] * n if dweights is None else list(dweights) + if not (len(dys) == len(hists) == len(dxs) == len(dweights) == n): + raise ValueError( + f"dys, hists, dxs and dweights must hold one entry per convolution ({n})" + ) + mode = _graph_mode(graph) + if mode == "off": + return _kcp_backward_multi(xs, weights, ctx, dys, hists, dxs, dweights) + entry = _graph_entry( + "bwd_multi", ctx, (*xs, *weights, *dys, *hists, *dxs, *dweights), mode + ) + dx_out, dw_out = _graph_run( + "bwd_multi", + entry, + lambda: _kcp_backward_multi(xs, weights, ctx, dys, hists, dxs, dweights), + mode, + ) + return ( + tuple( + d if o is not None else d.detach() for d, o in zip(dx_out, dxs, strict=True) + ), + tuple( + d if o is not None else d.detach() + for d, o in zip(dw_out, dweights, strict=True) + ), + ) + + +def _kcp_backward_multi( + xs: list[torch.Tensor], + weights: list[torch.Tensor], + ctx: Any, + dys: list[torch.Tensor], + hists: list[torch.Tensor], + dxs: list[Optional[torch.Tensor]], + dweights: list[Optional[torch.Tensor]], +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + """The multi backward's host flow (eagerly or inside a capture): see + :func:`cake_causal_conv1d_det_multi_backward_kcp`.""" + + device = _require_tensor(xs[0], "xs[0]").device + _wait_ready(ctx) + calls = [ + _bwd_host(x, w, ctx, dy, h, dx=dx, dweight=dw, device=device) + for x, w, dy, h, dx, dw in zip( + xs, weights, dys, hists, dxs, dweights, strict=True + ) + ] + heads = [ + _bwd_head(c[8], c[2], c[3], c[4], c[5], S=c[0], C=c[1], device=device) + for c in calls + ] + e0 = calls[0][8] + n_head = e0["n_head"] + work = None + if int(ctx.world_size) == 1: + tails = [h[:n_head] for h in heads] + else: + send = ( + torch.cat([h[:n_head] for h in heads], dim=1) + if len(heads) > 1 + else heads[0][:n_head] + ) + recv, work = _exchange( + send, + recv_rows=e0["tail_rows"], + send_splits=e0["send_splits"], + recv_splits=e0["recv_splits"], + ctx=ctx, + ) + if work is not None: + work.wait() + tails = _unpack(recv, [c[1] for c in calls]) if len(calls) > 1 else [recv] + dx_out = tuple(c[6] for c in calls) + dw_out = tuple(c[7] for c in calls) + if calls[0][0] == 0: + for dw in dw_out: + dw.zero_() + return dx_out, dw_out + for (S, C, x_flat, w, dy_flat, hist_t, dx, dw, e), tail in zip( + calls, tails, strict=True + ): + _bwd_main( + e, x_flat, w, dy_flat, hist_t, dx, dw, tail, None, S=S, C=C, device=device + ) + return dx_out, dw_out + + +def _handle_hists( + entry: dict[str, Any], xs: list[torch.Tensor], ctx: Any +) -> tuple[torch.Tensor, ...]: + """The stored pools of a multi call's ``"conv"`` entry, each checked like :func:`_handle_hist` against its tensor + of this recompute (an entry of a single call, or of another tensor count, raises).""" + pools, keys, shapes = entry.get("hists"), entry.get("keys"), entry.get("shapes") + n = len(xs) + if not all(isinstance(t, tuple) and len(t) == n for t in (pools, keys, shapes)): + raise ValueError( + f"recompute_handle: the 'conv' entry holds no {n}-tensor multi-call pools (an entry of " + f"cake_causal_conv1d_det_kcp or of another tensor count); one handle serves one call" + ) + if entry.get("ctx") is not ctx: + raise ValueError( + "recompute_handle: the 'conv' entry belongs to another KcpContext" + ) + return tuple( + _handle_hist( + { + "hist": hist, + "ctx": ctx, + "key": key, + "shape": shape, + "dtype": entry.get("dtype"), + "device": entry.get("device"), + }, + x, + ctx, + ) + for hist, x, key, shape in zip(pools, xs, keys, shapes, strict=True) + ) + + +class CakeCausalConv1dDetMultiFunctionKcp(torch.autograd.Function): + """Autograd wrapper of several convolutions with one halo exchange per direction: ``apply(ctx, graph, handle, + n, *xs, *weights)`` returns the ``n`` outputs; the forward saves ``(*xs, *weights, *hists)``; the backward runs + :func:`cake_causal_conv1d_det_multi_backward_kcp`. ``handle``: one recompute handle for the call, whose ``"conv"`` + entry holds the ``n`` pools (the saved set's own tensors): the recompute binds them and exchanges nothing. + Refused with graph replay (keyword or knob) and inside a CUDA-graph capture, like the single call.""" + + @staticmethod + def forward(fn_ctx, ctx, graph, handle, n, *tensors): + xs, weights = list(tensors[:n]), list(tensors[n:]) + if handle is None: + ys, hists = cake_causal_conv1d_det_multi_forward_kcp( + xs, weights, ctx, graph=graph + ) + else: + _handle_refusals(graph, handle) + entry = _handle_lookup(handle) + outs = [None] * n + if entry is None: + ys, hists = _kcp_forward_multi(xs, weights, ctx, outs=outs) + handle.entries[HANDLE_KEY] = { + "hists": hists, + "shapes": tuple(tuple(h.shape) for h in hists), + "dtype": torch.bfloat16, + "device": hists[0].device, + "ctx": ctx, + "keys": tuple((int(x.shape[-2]), int(x.shape[-1])) for x in xs), + } + LAST_CALL["recompute_handle"] = {"fwd": "capture", "n": n} + else: + hists = _handle_hists(entry, xs, ctx) + ys, hists = _kcp_forward_multi( + xs, weights, ctx, outs=outs, _hists=hists + ) + LAST_CALL["recompute_handle"] = {"fwd": "replay", "n": n} + fn_ctx.save_for_backward(*xs, *weights, *hists) + fn_ctx.kcp = ctx + fn_ctx.graph = graph + fn_ctx.n = n + return tuple(ys) + + @staticmethod + def backward(fn_ctx, *grad_ys): + n = fn_ctx.n + saved = fn_ctx.saved_tensors + xs, weights, hists = saved[:n], saved[n : 2 * n], saved[2 * n :] + dys = [ + (torch.zeros_like(x) if g is None else g).to(torch.bfloat16).contiguous() + for g, x in zip(grad_ys, xs, strict=True) + ] + dxs, dweights = cake_causal_conv1d_det_multi_backward_kcp( + xs, weights, fn_ctx.kcp, dys, hists, graph=fn_ctx.graph + ) + return (None, None, None, None, *dxs, *dweights) + + +def cake_causal_conv1d_det_multi_kcp( + xs: Any, + weights: Any, + ctx: Any, + *, + graph: Optional[str] = None, + recompute_handle: Optional[Any] = None, +) -> tuple[torch.Tensor, ...]: + """Autograd-enabled KCP convolutions of several tensors of one layer (for example its q, k and v) with one halo + exchange per direction. + + Bitwise ``tuple(cake_causal_conv1d_det_kcp(x, w, ctx) for x, w in zip(xs, weights))`` -- outputs and every gradient, + ``dweight`` included (each tensor keeps its own kernels and fixed-order partial) -- with one collective in the + forward and one in the backward instead of one per tensor: fewer global synchronization points per layer. + + Parameters + ---------- + xs : sequence of torch.Tensor + ``n >= 1`` BF16 ``[T_r, C_i]`` (or ``[1, T_r, C_i]``) contiguous activations of the rank's fragments, in the + order of ``ctx.cu_seqlens_local``; the channel counts may differ. + weights : sequence of torch.Tensor + FP32 ``[C_i, 4]`` contiguous taps, one per tensor. + ctx : KcpContext + The rank's context, as in :func:`cake_causal_conv1d_det_kcp`. Every rank of the group calls this function once + per layer, in the same order, with the same tensor count and channel counts. The ranks are not compared + (that would cost a collective): a different count, channel count or order hangs the exchange or silently + corrupts the history rows, as a mismatch of the single entry's per-rank contract does. + graph : str, optional + As in :func:`cake_causal_conv1d_det_kcp` (keyed CUDA-graph capture / replay; the key holds every operand). + recompute_handle : RecomputeHandleDetKcp, optional + ``None`` (default): unchanged. One handle for the whole call: the original forward stores the ``n`` history + pools under ``"conv"``, torch's recompute inside the backward binds them and exchanges nothing; the backward + still exchanges once. Same rules as :func:`cake_causal_conv1d_det_kcp` (one handle per call and micro-step, + created outside the outermost checkpoint region; refused with ``graph="replay"`` and inside a CUDA-graph + capture; every rank passes one or none). The handle may also carry the ``"kda"`` entry of the layer's + recurrence; a handle's ``"conv"`` entry serves either one multi call or one single call, never both. + + Returns + ------- + tuple of torch.Tensor + The ``n`` BF16 outputs, each of its input's shape. + """ + + _graph_mode(graph) + xs, weights, _ = _multi_operands(xs, weights, None) + if recompute_handle is not None: + _handle_refusals(graph, recompute_handle) + return CakeCausalConv1dDetMultiFunctionKcp.apply( + ctx, graph, recompute_handle, len(xs), *xs, *weights + ) + + +__all__ = [ + "GRAPH_MODES", + "GRAPH_DEFAULT", + "GRAPH_STATS", + "HANDLE_KEY", + "RecomputeHandleDetKcp", + "LAST_CALL", + "graph_cache_clear", + "CakeCausalConv1dDetFunctionKcp", + "cake_causal_conv1d_det_kcp", + "cake_causal_conv1d_det_backward_kcp", + "cake_causal_conv1d_det_forward_kcp", + "CakeCausalConv1dDetMultiFunctionKcp", + "cake_causal_conv1d_det_multi_kcp", + "cake_causal_conv1d_det_multi_backward_kcp", + "cake_causal_conv1d_det_multi_forward_kcp", +] + +# The names before the det rename, bound to the same objects (not in ``__all__``). +CakeCausalConv1dKcpFunction = CakeCausalConv1dDetFunctionKcp +CakeCausalConv1dKcpMultiFunction = CakeCausalConv1dDetMultiFunctionKcp +KcpRecomputeHandle = RecomputeHandleDetKcp +cake_causal_conv1d_kcp = cake_causal_conv1d_det_kcp +cake_causal_conv1d_kcp_backward = cake_causal_conv1d_det_backward_kcp +cake_causal_conv1d_kcp_forward = cake_causal_conv1d_det_forward_kcp +cake_causal_conv1d_kcp_multi = cake_causal_conv1d_det_multi_kcp +cake_causal_conv1d_kcp_multi_backward = cake_causal_conv1d_det_multi_backward_kcp +cake_causal_conv1d_kcp_multi_forward = cake_causal_conv1d_det_multi_forward_kcp diff --git a/flashinfer/cake_causal_conv1d_det_train_backend/kcp_jit.py b/flashinfer/cake_causal_conv1d_det_train_backend/kcp_jit.py new file mode 100644 index 00000000000..89cf4d1f934 --- /dev/null +++ b/flashinfer/cake_causal_conv1d_det_train_backend/kcp_jit.py @@ -0,0 +1,821 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Host plan and loader glue of the Cake context-parallel (KCP) packed width-4 +depthwise causal convolution + SiLU training kernels (SM100 / SM103 / SM107): +the halo-aware programs the KCP host flow (:mod:`.kcp`) launches, recorded +beside the zero-history registry :mod:`.cake_jit`, whose tables this module +never touches. + +Every module below is one mechanically exported production build: a CUDA +device translation unit plus a tvm-ffi binding compiled together for exactly +one architecture. ``MODULES`` and ``ROUTES`` carry the host plan only, in the +vocabulary of :mod:`.cake_jit`: the kernel's ``library`` stem, the +architectures it is delivered for, its launch record, grid rule, workspace +rule, selection rule and the kernel-parameter -> logical-argument map. What to +compile -- sources, digests, nvcc flags -- is recorded once in ``build.json`` +next to this file (its entries sit beside the zero-history entries); +:mod:`.cake_build` serves the library ``_`` from a prebuilt +``aot/`` file when a builder compiled it and otherwise through FlashInfer's JIT +from the same sources and flags, verifying every source digest first. + +The routes of an architecture are the three KCP kinds: ``kcp_fwd`` (one +program: ``y`` and the fragments' history rows ``hist`` from ``x``, ``weight``, +``cu_seqlens`` and the context's history tables), ``kcp_bwd`` (the fallback +program and the big-channel ``_st8`` program, selected by the channel count +exactly as the zero-history backward route selects its programs: a record's +``select_rule`` ``channels_at_least`` owns ``C >= channels[arch]``, the record +without a rule is the fallback; ``dx``, the deterministic FP32 ``dweight`` and +the predecessor fragments' tail derivatives ``dpre_tail`` in one cooperative +launch) and ``kcp_bwd_head`` (one program: the head-row derivatives +``dpre_head`` every consumer fragment sends back; its grid is +``max(1, n_frag x ceil(C / (threads x vec)))``, computed by :mod:`.kcp` from +the record's ``grid_rule.geometry`` and launched through the record's +``arg_plan`` / ``logical`` / ``ffi_entry`` and :func:`load`). + +The host plan reuses :mod:`.cake_jit`'s functions over a record (:func:`plan`, +:func:`launch_grid`, :func:`partial_rows`, :func:`partial_bytes`, +:func:`workspace_layout`, :func:`channel_split`) unchanged; the functions bound +to the tables (:func:`route_modules`, :func:`selected_modules`, +:func:`workspace_record`, :func:`library_name`, :func:`load`, +:func:`run_module`) are spelled here over ``MODULES`` and ``ROUTES``. The +logical tensor names are :mod:`.cake_jit`'s plus the KCP tables ``hist``, +``hist_row_off``, ``dpre_tail``, ``tail_row_off``, ``dpre_head`` and +``head_row_off``; the two ``CUtensorMap`` parameters ``x_map`` / ``dy_map`` of +the TMA backward programs map onto ``x`` / ``dy`` (the binding encodes the +descriptors from the tensors at every launch); the scalars are +:mod:`.cake_jit`'s. + +The program name ``cake_causal_conv1d_train_<20 hex>`` is the first twenty hex +digits of the SHA-256 of the compact JSON ``{"build": , +"module": }`` of the kernel the program was rendered +from: the forward and the head program with their kernel files' default +keywords; the backward with ``vec=8, copy="tma2d"`` on sm_100a and sm_103a, +with ``vec=8, min_blocks=2`` on sm_107a, and the big-channel backward with +``vec=8, copy="tma2d", stage_tokens=8, min_blocks=2`` -- the builds of the +zero-history backward, which the KCP kernels extend. +""" + +from __future__ import annotations + +import functools +from typing import Any, Sequence + +import tvm_ffi + +from . import cake_build +from . import cake_jit as _zero_history +from .cake_jit import ( # noqa: F401 (re-exported: the host plan over a record is the zero-history one) + ARCH_BY_CAPABILITY, + ARCHES, + CHANNEL_MULTIPLE, + INT32_MAX, + LOGICAL_SCALARS, + SELECT_KINDS, + WIDTH, + arch_for_capability, + channel_split, + launch_grid, + partial_bytes, + partial_rows, + plan, + sm_count, + workspace_layout, +) + +# The derived leaves of every record (``kernel_symbol``, ``launch.block``, +# ``launch.cluster``, ``launch.cooperative``, ``launch.dynamic_smem_bytes``, +# ``launch.use_pdl``, ``arg_plan`` and ``logical``) are rewritten by the +# exporter from the rendered program; every other field is host plan. +MODULES: dict[str, dict[str, Any]] = { + "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "y"], + ["buffer", "hist"], + ["buffer", "hist_row_off"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # The grid is the co-resident CTA count, ``ctas_per_sm x SM count``, + # as for the zero-history forward: CTA ``b`` owns channel tile + # ``b % n_ch_tiles`` and an equal contiguous share of the token axis. + # ``ctas_per_sm`` is the measured co-residency of the program (the + # occupancy query's answer, 4 CTAs per SM at 112 registers on GB300; + # sm_100a / sm_107a are taken as the same launch bound); ``geometry`` + # is the kernel's build geometry the host plan derives its scalars + # from. Besides ``y`` the forward writes ``hist`` (BF16 history rows + # ``[n_hist_rows, C]``) at the rows ``hist_row_off`` (int32 + # ``[n_frag, 3]``, ``-1`` = none) names. + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 4, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 33792, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "hist": "hist", + "hist_row_off": "hist_row_off", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "search_steps": "search_steps", + "w": "w", + "x": "x", + "y": "y", + }, + "role": "kcp_fwd", + "workspace_rule": None, + }, + "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp": { + "arches": [ + "sm_100a", + "sm_103a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "dw"], + ["buffer", "partial"], + ["buffer", "hist"], + ["buffer", "hist_row_off"], + ["buffer", "dpre_tail"], + ["buffer", "tail_row_off"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # Cooperative (a grid-wide barrier precedes the fixed-order dweight + # reduction): the grid is exactly the co-resident CTA count, 3 CTAs + # per SM (168 registers at 128 threads, the launch bound), the + # occupancy query's answer on B200 and GB300 for this build. The + # fallback program of the route (no ``select_rule``). Besides ``dx`` + # and ``dw`` it writes the predecessors' tail derivatives + # ``dpre_tail`` (FP32 ``[n_tail_rows, C]``) at the rows + # ``tail_row_off`` (int32 ``[n_frag, 3]``, ``-1`` = none) names and + # reads the forward's ``hist`` rows through ``hist_row_off``. + "grid_rule": { + "cooperative": True, + "ctas_per_sm": 3, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": True, + "dynamic_smem_bytes": 50176, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dpre_tail": "dpre_tail", + "dw": "dw", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "hist": "hist", + "hist_row_off": "hist_row_off", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "tail_row_off": "tail_row_off", + "w": "w", + "x": "x", + "x_map": "x", + }, + "role": "kcp_bwd", + # The FP32 ``partial`` workspace: ``[ceil(grid / n_ch_tiles), C, taps]`` + # rows of per-CTA dweight partials, reduced in fixed order after the + # grid barrier (:func:`partial_rows`, :func:`partial_bytes`). + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + "cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp": { + "arches": [ + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "dw"], + ["buffer", "partial"], + ["buffer", "hist"], + ["buffer", "hist_row_off"], + ["buffer", "dpre_tail"], + ["buffer", "tail_row_off"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # The sm_107a fallback backward (``vec=8, min_blocks=2``: two CTAs per + # SM, bulk copies, no tensor maps), cooperative like every backward + # program; the grid is its co-resident CTA count. + "grid_rule": { + "cooperative": True, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": True, + "dynamic_smem_bytes": 50176, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dpre_tail": "dpre_tail", + "dw": "dw", + "dx": "dx", + "dy": "dy", + "hist": "hist", + "hist_row_off": "hist_row_off", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "tail_row_off": "tail_row_off", + "w": "w", + "x": "x", + }, + "role": "kcp_bwd", + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "dx"], + ["buffer", "dw"], + ["buffer", "partial"], + ["buffer", "hist"], + ["buffer", "hist_row_off"], + ["buffer", "dpre_tail"], + ["buffer", "tail_row_off"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["parameter", "search_steps"], + ["parameter", "ct_vec"], + ["parameter", "n_sub"], + ["parameter", "n_ch_tiles"], + ["buffer", "x_map"], + ["buffer", "dy_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # The big-channel backward (``stage_tokens=8``, eight-row TMA boxes per + # stage, two CTAs per SM): it runs when ``C >= select_rule.channels + # [arch]``, else the architecture's fallback record runs -- the + # zero-history backward's thresholds, which these programs' geometry + # measured. Cooperative; the grid is its co-resident CTA count. + "grid_rule": { + "cooperative": True, + "ctas_per_sm": 2, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": True, + "dynamic_smem_bytes": 82944, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp", + "logical": { + "C": "C", + "S": "S", + "ct_vec": "ct_vec", + "cu": "cu", + "dpre_tail": "dpre_tail", + "dw": "dw", + "dx": "dx", + "dy": "dy", + "dy_map": "dy", + "hist": "hist", + "hist_row_off": "hist_row_off", + "n_ch_tiles": "n_ch_tiles", + "n_sub": "n_sub", + "num_docs": "num_docs", + "partial": "partial", + "search_steps": "search_steps", + "tail_row_off": "tail_row_off", + "w": "w", + "x": "x", + "x_map": "x", + }, + "role": "kcp_bwd", + "select_rule": { + "channels": { + "sm_100a": 512, + "sm_103a": 512, + "sm_107a": 1024, + }, + "kind": "channels_at_least", + }, + "workspace_rule": { + "kind": "partial_rows", + "taps": 4, + }, + }, + "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp": { + "arches": [ + "sm_100a", + "sm_103a", + "sm_107a", + ], + "arg_plan": [ + ["buffer", "x"], + ["buffer", "w"], + ["buffer", "cu"], + ["buffer", "dy"], + ["buffer", "hist"], + ["buffer", "hist_row_off"], + ["buffer", "dpre_head"], + ["buffer", "head_row_off"], + ["parameter", "S"], + ["parameter", "C"], + ["parameter", "num_docs"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + # The head-row derivative program: ``dpre_head`` (FP32 ``[n_head_rows, + # C]``) at the rows ``head_row_off`` (int32 ``[n_frag, 3, 3]``, ``-1`` + # = none) names, from the fragment's first three tokens, the history + # rows and ``dy``. Its grid is ``max(1, n_frag x ceil(C / (threads x + # vec)))`` -- a work-item count, computed by the host flow from + # ``geometry`` (the grid strides over the items) -- not the co-resident + # count; ``ctas_per_sm`` records the measured co-residency (5 CTAs per + # SM at 96 registers and no shared memory, the occupancy query's + # answer on GB300) for a device-bound launch. + "grid_rule": { + "cooperative": False, + "ctas_per_sm": 5, + "geometry": { + "threads": 128, + "vec": 8, + }, + "kind": "co_resident", + }, + "kernel": "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp", + "kernel_symbol": "kernel_cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp", + "launch": { + "block": (128, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 0, + "use_pdl": False, + }, + "library": "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp", + "logical": { + "C": "C", + "S": "S", + "cu": "cu", + "dpre_head": "dpre_head", + "dy": "dy", + "head_row_off": "head_row_off", + "hist": "hist", + "hist_row_off": "hist_row_off", + "num_docs": "num_docs", + "w": "w", + "x": "x", + }, + "role": "kcp_bwd_head", + "workspace_rule": None, + }, +} + +# architecture -> KCP kind -> the records the route may launch, in launch +# order. The forward and the head route carry one program; the backward +# route carries the architecture's fallback program first and the +# big-channel ``_st8`` program second (:func:`selected_modules` launches one). +ROUTES: dict[str, dict[str, list[str]]] = { + "sm_100a": { + "kcp_bwd": [ + "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp", + "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp", + ], + "kcp_bwd_head": [ + "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp", + ], + "kcp_fwd": [ + "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp", + ], + }, + "sm_103a": { + "kcp_bwd": [ + "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp", + "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp", + ], + "kcp_bwd_head": [ + "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp", + ], + "kcp_fwd": [ + "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp", + ], + }, + "sm_107a": { + "kcp_bwd": [ + "cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp", + "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp", + ], + "kcp_bwd_head": [ + "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp", + ], + "kcp_fwd": [ + "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp", + ], + }, +} + +# The KCP kinds (the zero-history kinds are ``cake_jit.KINDS``). +KINDS = ("kcp_fwd", "kcp_bwd", "kcp_bwd_head") +# The history / halo rows every fragment carries (the convolution width minus one). +HALO = WIDTH - 1 +# Logical argument names the KCP host flow binds: the zero-history vocabulary +# (``cake_jit.LOGICAL_TENSORS``, spelled out so the registry stays a literal) +# plus the KCP tables. ``hist``: BF16 history rows ``[n_hist_rows, C]``; +# ``hist_row_off``: int32 ``[n_frag, 3]`` (``-1`` = none); ``dpre_tail``: FP32 +# ``[n_tail_rows, C]`` tail derivatives with ``tail_row_off`` int32 +# ``[n_frag, 3]``; ``dpre_head``: FP32 ``[n_head_rows, C]`` head derivatives +# with ``head_row_off`` int32 ``[n_frag, 3, 3]``. The scalars are the +# zero-history scalars (``cake_jit.LOGICAL_SCALARS``). +LOGICAL_TENSORS = ( + "x", + "w", + "cu", + "y", + "dy", + "dx", + "dw", + "partial", + "hist", + "hist_row_off", + "dpre_tail", + "tail_row_off", + "dpre_head", + "head_row_off", +) + + +def route_applies(*, device_capability: Sequence[int], kind: str) -> bool: + """Whether the verified export owns one KCP training call on a device.""" + + arch = arch_for_capability(device_capability) + if arch is None: + return False + return kind in ROUTES.get(arch, {}) + + +def route_modules(arch: str, kind: str) -> tuple[str, ...]: + """Ordered module names of one KCP route on ``arch`` (every program the route may launch). + + The forward and the head route launch their one program; the backward + launches the one :func:`selected_modules` picks for the call's channel + count. + """ + + if kind not in KINDS: + raise ValueError(f"kind must be one of {KINDS}, got {kind!r}") + names = ROUTES.get(arch, {}).get(kind) + if not names: + raise ValueError( + f"unsupported Cake causal conv1d KCP training route {kind!r} on {arch!r}; " + f"exported architectures: {ARCHES}" + ) + return tuple(names) + + +def select_rule(record: dict[str, Any]) -> dict[str, Any] | None: + """The selection rule of one module record, ``None`` when the record carries none. + + A backward record without a rule (no ``select_rule`` key, or ``None``) is + its route's fallback program. ``{"kind": "channels_at_least", "channels": + {arch: C_min, ...}}`` owns the calls whose channel count is at least + ``channels[arch]`` on the architectures it names and never hits on the + others. Forward and head records carry no rule. + """ + + return record.get("select_rule") + + +def selected_modules( + arch: str, kind: str, *, C: int, S: int | None = None +) -> tuple[str, ...]: + """Ordered module names one KCP call launches on ``arch`` for ``C`` channels. + + The forward and the head route launch every program of their route + (:func:`route_modules`, one each). The backward route carries the + architecture's fallback program and the big-channel program and launches + exactly one: the record whose :func:`select_rule` admits ``C`` on ``arch`` + when one does, else the route's single fallback record (the one without a + rule) -- the rule of the zero-history backward. The choice is a pure + function of the architecture and the channel count (``S`` is accepted and + unused, and the host never reads the data), so a given ``(S, C, device)`` + always runs the same program and ``dweight`` stays bitwise reproducible. + Several admitting records, or zero or several fallbacks, are registry + errors (``RuntimeError``); an unknown rule kind is a ``ValueError``. + """ + + names = route_modules(arch, kind) + if kind != "kcp_bwd": + return names + positive: list[str] = [] + fallback: list[str] = [] + for name in names: + rule = select_rule(MODULES[name]) + if rule is None: + fallback.append(name) + elif _zero_history._admits(rule, arch=arch, C=C): + positive.append(name) + if len(fallback) != 1: + raise RuntimeError( + f"the exported KCP backward route on {arch} declares {len(fallback)} fallback " + "programs (records without a select_rule); expected exactly one" + ) + if len(positive) > 1: + raise RuntimeError( + f"the exported KCP backward route on {arch} selects {len(positive)} programs " + f"for C={int(C)} ({positive}); expected at most one" + ) + if positive: + return (positive[0],) + return (fallback[0],) + + +def workspace_owners(arch: str) -> tuple[dict[str, Any], ...]: + """The KCP backward module records of ``arch`` that own a ``partial`` workspace, in route order.""" + + owners = tuple( + MODULES[name] + for name in route_modules(arch, "kcp_bwd") + if MODULES[name]["workspace_rule"] is not None + ) + if not owners: + raise RuntimeError( + f"the exported KCP backward route on {arch} declares no workspace owner; " + "expected at least one" + ) + return owners + + +def workspace_record(arch: str) -> dict[str, Any]: + """The capacity owner of the KCP backward ``partial`` workspace on ``arch``. + + As for the zero-history route: the owners share one build geometry, so + their rows and bytes grow with ``ctas_per_sm`` alone and the record with + the largest ``ctas_per_sm`` (route order breaks ties) bounds them all. + """ + + owners = workspace_owners(arch) + geometries = { + tuple(sorted(owner["grid_rule"]["geometry"].items())) for owner in owners + } + if len(geometries) != 1: + raise RuntimeError( + f"the KCP backward programs on {arch} declare {len(geometries)} build " + "geometries; one workspace serves them only when they share one" + ) + return max(owners, key=lambda owner: int(owner["grid_rule"]["ctas_per_sm"])) + + +def head_grid(record: dict[str, Any], *, n_frag: int, C: int) -> tuple[int, int, int]: + """The head program's grid: ``max(1, n_frag x ceil(C / (threads x vec)))`` work items over the record's geometry.""" + + geometry = record["grid_rule"]["geometry"] + channels_per_cta = int(geometry["threads"]) * int(geometry["vec"]) + if channels_per_cta <= 0: + raise ValueError( + "the head record's geometry must cover at least one channel per CTA" + ) + items = int(n_frag) * (-(-int(C) // channels_per_cta)) + return (max(1, items), 1, 1) + + +def device_arch(device_index: int) -> str: + """Exported architecture of one CUDA device (``RuntimeError`` when it is not an exported target).""" + + return _zero_history.device_arch(device_index) + + +def library_name(name: str, arch: str) -> str: + """``build.json`` library of KCP module record ``name`` built for ``arch``: ``_``.""" + + record = MODULES[name] + if arch not in record["arches"]: + raise RuntimeError( + f"Cake causal conv1d KCP training module {name} is not delivered for {arch}; " + f"its architectures: {record['arches']}" + ) + return f"{record['library']}_{arch}" + + +@functools.cache +def load(name: str, arch: str) -> Any: + """The loaded tvm-ffi module of KCP record ``name`` built for ``arch``. + + The prebuilt ``aot/`` library when a builder compiled it, else FlashInfer's + JIT build of the same sources and flags (``cake_build.load_library``). + """ + + return cake_build.load_library(library_name(name, arch)) + + +def specs_for_arch(arch: str) -> list[Any]: + """JIT specifications of every ``arch`` library of the package's ``build.json`` without a prebuilt file + (the zero-history and the KCP entries alike; for ahead-of-time warm-up through FlashInfer's + ``build_jit_specs``).""" + + return cake_build.jit_specs(arch) + + +def run_module( + name: str, + values: dict[str, Any], + *, + device_index: int, + grid: tuple[int, int, int] | None = None, +) -> None: + """Launch one exported KCP module with logical argument ``values``. + + ``values`` maps the logical names in ``LOGICAL_TENSORS`` / + ``LOGICAL_SCALARS`` to tensors and Python scalars; the module record maps + its kernel parameter names onto that vocabulary. The launch grid is the + module's :func:`launch_grid` on ``device_index`` (the co-resident count) + unless ``grid`` gives it (the head program's work-item grid, + :func:`head_grid`), and the kernel runs on the current torch stream. + """ + + record = MODULES[name] + logical = record["logical"] + if grid is None: + grid = launch_grid(record, device_index=device_index) + grid_values = { + "grid_x": int(grid[0]), + "grid_y": int(grid[1]), + "grid_z": int(grid[2]), + } + args = [] + for kind, key in record["arg_plan"]: + if kind == "grid": + args.append(grid_values[key]) + continue + logical_name = logical[key] + if logical_name not in values: + raise ValueError( + f"module {name} needs {logical_name!r} (kernel parameter {key!r})" + ) + args.append(values[logical_name]) + module = load(name, device_arch(device_index)) + with tvm_ffi.use_torch_stream(): + getattr(module, record["ffi_entry"])(*args) + + +__all__ = [ + "ARCHES", + "ARCH_BY_CAPABILITY", + "CHANNEL_MULTIPLE", + "HALO", + "INT32_MAX", + "KINDS", + "LOGICAL_SCALARS", + "LOGICAL_TENSORS", + "MODULES", + "ROUTES", + "SELECT_KINDS", + "WIDTH", + "arch_for_capability", + "channel_split", + "device_arch", + "head_grid", + "launch_grid", + "library_name", + "load", + "partial_bytes", + "partial_rows", + "plan", + "route_applies", + "route_modules", + "run_module", + "select_rule", + "selected_modules", + "sm_count", + "specs_for_arch", + "workspace_layout", + "workspace_owners", + "workspace_record", +] diff --git a/flashinfer/cake_causal_conv1d_det_train_kcp.py b/flashinfer/cake_causal_conv1d_det_train_kcp.py new file mode 100644 index 00000000000..d9d6c17e880 --- /dev/null +++ b/flashinfer/cake_causal_conv1d_det_train_kcp.py @@ -0,0 +1,171 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Deterministic context-parallel (KCP) packed width-4 depthwise causal convolution + SiLU training +(SM100 / SM103 / SM107), the Cake backend. ``flashinfer.cake_causal_conv1d_det_train`` keeps the +names these entry points had before (``cake_causal_conv1d_kcp``, ``cake_causal_conv1d_kcp_multi``, +``KcpRecomputeHandle``) as compatibility aliases of the objects defined here. + +Every rank convolves its fragments of the packed batch; the first tokens of a fragment that continues a +document read the history rows of the preceding fragment, exchanged through the halo collective of the +context-parallel context of the KDA step (``flashinfer.cake_kda_det_train_kcp.build_context_det_kcp``). +The per-rank launch sequences, the autograd Functions and the kernels of the ``kcp_jit`` registry live in +``flashinfer.cake_causal_conv1d_det_train_backend.kcp``; the single-rank entry points are in +``flashinfer.cake_causal_conv1d_det_train``:: + + from flashinfer.cake_causal_conv1d_det_train_kcp import ( + cake_causal_conv1d_det_kcp, + cake_causal_conv1d_det_multi_kcp, + ) + + # this rank's fragments, with the history rows of the preceding fragments exchanged + # through the KCP context (FP32 taps, no bias) + y_local = cake_causal_conv1d_det_kcp(x_local, weight, ctx) + + # a layer's q / k / v convolutions with one halo exchange per direction + # (bitwise three single calls) + q, k, v = cake_causal_conv1d_det_multi_kcp((xq, xk, xv), (wq, wk, wv), ctx) +""" + +from __future__ import annotations + +from typing import Any, Optional, Sequence + +import torch + +from .api_logging import flashinfer_api + + +@flashinfer_api +def cake_causal_conv1d_det_kcp( + x_local: torch.Tensor, + weight: torch.Tensor, + ctx: Any, + *, + recompute_handle: Optional[Any] = None, +) -> torch.Tensor: + """Autograd-enabled context-parallel packed width-4 depthwise causal convolution + SiLU of this rank's fragments. + + The first tokens of a fragment that continues a document read the history rows of the preceding fragment + (exchanged through the context's halo collective) instead of zeros, so the result equals the single-rank + :func:`flashinfer.cake_causal_conv1d_det_train.cake_causal_conv1d_det` of the whole batch on this rank's tokens. The KCP path takes FP32 taps and no + bias (the package's pre-round-6 contract). + + Parameters + ---------- + x_local : torch.Tensor + BF16 ``[S_local, C]`` or ``[1, S_local, C]`` contiguous packed activations of this rank's fragments, in the + partition's order; ``C`` is a multiple of 8. + weight : torch.Tensor + FP32 ``[C, 4]`` contiguous taps. + ctx : ContextDetKcp + From ``flashinfer.cake_kda_det_train_kcp.build_context_det_kcp`` (the same context the KDA step of the rank uses). + recompute_handle : RecomputeHandleDetKcp, optional + Keyword-only, default ``None`` (unchanged behavior). A handle (any object with an ``entries`` dict, such as + ``flashinfer.cake_causal_conv1d_det_train_kcp.RecomputeHandleDetKcp``) makes the checkpoint recompute of + this call run without the forward halo exchange: the original forward exchanges as always and stores its + history rows (the saved set's own tensor, key ``"conv"``) in the handle; torch's recompute of the forward + inside the backward reads them back, so ``y`` and the gradients are bitwise those of the same checkpoint + without a handle. The backward still exchanges. Create one handle per convolution call and micro-step, + outside the outermost checkpoint region, and pass it as an argument of the checkpointed function. A + forward outside the backward that finds an entry raises (a handle reused across micro-steps). The forward + halo exchange is a collective: every rank of ``ctx.group`` passes a handle for the same call or none, created + at the same checkpoint nesting level. A mismatch pairs one rank's recompute exchange with another rank's + backward exchange and hangs or corrupts data; no check is made. Never share one handle between two + convolution calls (also not in nested regions), and never ``clear()`` or reuse a handle while a backward + that needs it is pending. + + Returns + ------- + torch.Tensor + BF16 ``y`` of ``x_local``'s shape. The backward runs the head program, the backward halo exchange and the + fused ``dx`` + deterministic FP32 ``dweight`` kernel of this rank's fragments. + """ + + from .cake_causal_conv1d_det_train_backend import kcp as _kcp # noqa: PLC0415 + + return _kcp.cake_causal_conv1d_det_kcp( + x_local, weight, ctx, recompute_handle=recompute_handle + ) + + +@flashinfer_api +def cake_causal_conv1d_det_multi_kcp( + xs: Sequence[torch.Tensor], + weights: Sequence[torch.Tensor], + ctx: Any, + *, + recompute_handle: Optional[Any] = None, +) -> tuple[torch.Tensor, ...]: + """Autograd-enabled context-parallel convolutions of several tensors of one layer (for example its q, k and v) + with one halo exchange per direction. + + The result is bitwise ``tuple(cake_causal_conv1d_det_kcp(x, w, ctx) for x, w in zip(xs, weights))``: outputs and + every gradient, ``dweight`` included (each tensor keeps its own kernels and fixed-order reduction). The forward + issues one halo collective instead of one per tensor (the history rows of all tensors travel side by side and are + copied out after the wait), and so does the backward, so a layer has fewer global synchronization points. + + Parameters + ---------- + xs : sequence of torch.Tensor + ``n >= 1`` BF16 ``[S_local, C_i]`` or ``[1, S_local, C_i]`` contiguous packed activations of this rank's + fragments, in the partition's order; each ``C_i`` is a multiple of 8 (they may differ). + weights : sequence of torch.Tensor + FP32 ``[C_i, 4]`` contiguous taps, one per tensor. + ctx : ContextDetKcp + As in :func:`cake_causal_conv1d_det_kcp`. Every rank of the group calls this function once per layer, in the + same order, with the same tensor count and channel counts. The ranks are not compared (that would cost a + collective): a different count, channel count or order hangs the exchange or silently corrupts the history + rows, as a mismatch of the single entry's per-rank contract does. + recompute_handle : RecomputeHandleDetKcp, optional + Keyword-only, default ``None`` (unchanged behavior). One handle for the whole call: the original forward + stores the ``n`` history pools under ``"conv"`` and torch's recompute inside the backward reads them back, + so the recompute issues no collective; the backward still exchanges once. The rules of + :func:`cake_causal_conv1d_det_kcp` apply (one handle per call and micro-step, created outside the outermost + checkpoint region; every rank passes one or none). The same handle may carry the layer's ``"kda"`` entry; + its ``"conv"`` entry serves one multi call or one single call, never both. + + Returns + ------- + tuple of torch.Tensor + The ``n`` BF16 outputs, each of its input's shape. + """ + + from .cake_causal_conv1d_det_train_backend import kcp as _kcp # noqa: PLC0415 + + return _kcp.cake_causal_conv1d_det_multi_kcp( + xs, weights, ctx, recompute_handle=recompute_handle + ) + + +# ``RecomputeHandleDetKcp`` (the KCP route's checkpoint recompute handle) resolves on first access, so importing this +# module does not import the KCP host module (torch.distributed, the KCP loader), like the entry points above. +def __getattr__(name: str) -> Any: + if name == "RecomputeHandleDetKcp": + from .cake_causal_conv1d_det_train_backend.kcp import RecomputeHandleDetKcp # noqa: PLC0415 + + return RecomputeHandleDetKcp + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted(set(globals()) | {"RecomputeHandleDetKcp"}) + + +__all__ = [ + "RecomputeHandleDetKcp", # noqa: F822 - resolved by __getattr__ + "cake_causal_conv1d_det_kcp", + "cake_causal_conv1d_det_multi_kcp", +] diff --git a/flashinfer/cake_causal_conv1d_train.py b/flashinfer/cake_causal_conv1d_train.py index 70710164bd7..b84805fe269 100644 --- a/flashinfer/cake_causal_conv1d_train.py +++ b/flashinfer/cake_causal_conv1d_train.py @@ -1,456 +1,7 @@ -""" -Copyright (c) 2026 by FlashInfer team. +# Copyright (c) 2026 by FlashInfer team. Licensed under Apache-2.0. +"""Compatibility import for flashinfer.cake_causal_conv1d_det_train (the same module object).""" -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at +from importlib import import_module +import sys - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. - -Cake packed width-4 depthwise causal convolution + SiLU training forward / -backward (SM100 / SM103 / SM107). - -The input is one packed BF16 ``[S, C]`` (or ``[1, S, C]``) activation whose -documents ``cu_seqlens`` (``int32 [num_docs + 1]`` device prefix sums) -delimit; every document is convolved on its own, with zero history before its -first token. With taps ``w[C, 4]`` (``w[c, 3]`` multiplies the current token) -and an optional per-channel bias ``b[C]``, the forward is, in FP32 with one -BF16 rounding, - - pre[t, c] = sum_k w[c, k] * x[t - 3 + k, c] + b[c] # x of another document is 0 - y[t, c] = bf16(pre * sigmoid(pre)) - -The taps and the bias are FP32 or BF16, in any combination (BF16 is widened to -FP32 on load; the arithmetic is FP32 either way); the bias is added last, after -the four taps, which is flash-linear-attention's order. The backward -recomputes ``pre`` in FP32 from ``x`` and evaluates the SiLU derivative in -FP32, - - dpre[t, c] = dy[t, c] * s * (1 + p * (1 - s)), s = sigmoid(p) - dx[t, c] = bf16(sum_k w[c, k] * dpre[t + 3 - k, c]) # dpre of another document is 0 - dweight[c, k] = sum_t dpre[t, c] * x[t - 3 + k, c] # fixed order - dbias[c] = sum_t dpre[t, c] # with a bias - -at ``p`` chosen by ``silu_grad_preactivation``: ``"bf16"`` (default), ``p = -bf16(pre)``, the rounded value the forward stores and the boundary -flash-linear-attention's backward uses (``dweight`` lands within about 3e-5 -rel-L2 of the exact FP64 gradient); ``"fp32"``, ``p = pre`` unrounded, with -the final ``dweight`` reduction in FP64 (within about 2e-7; a slower program). -``dx`` is within one BF16 ulp of exact under both. The backward is two -ordinary launches on one stream: a persistent main program writes ``dx`` and -one FP32 partial row of ``dweight`` (and ``dbias``) per CTA into a workspace, -and a reduce program launched programmatically dependent on it sums the rows -in a fixed order (no floating-point atomics), so the gradients are bitwise -reproducible for a given ``(S, C, device)`` and route. Neither launch is -cooperative: the backward keeps its speed while kernels of other streams -(NCCL collectives) are resident, which a whole-device cooperative grid had to -wait for. Nothing is saved from the forward but ``x``, ``w``, ``b`` and -``cu_seqlens``. - -The forward is one ordinary launch: the co-resident CTAs of each channel tile -take the tile's tokens in chunks from int32 counters on the device (the CTAs -that finish first take the tail), so ``y`` is the same bits whatever CTA -computes a chunk. The counters live in a forward workspace of -``cake_causal_conv1d_train_forward_workspace_bytes(S, C, device)`` bytes that is -zero before its first launch and that every launch leaves zero, so one buffer -serves every later forward call of one stream, CUDA Graph replays included. A -caller may own it (``workspace=``); otherwise the package keeps one zeroed -buffer per (device, current stream), and a call inside a CUDA Graph capture -allocates a zeroed buffer of its own. Two launches that may run at once need -two buffers. - -Inputs and outputs are contiguous (``C`` a multiple of 8); ``dweight`` and -``dbias`` are FP32 or BF16 (the parameter's dtype when not preallocated). -``S``, ``C`` and the document count are runtime scalars: they may change -between calls, the host never reads ``cu_seqlens`` (no synchronization), ``S -== 0`` returns empty outputs without launching, and the calls are CUDA Graph -capturable once the libraries are built. - -A call's arguments (bias, BF16 taps or gradient, the SiLU' policy) select its -route, and within a route the channel count selects one of the architecture's -programs through a host-side rule of the backend's registry (never the data), -so a given ``(S, C, device)`` and route always run the same program. The -workspace size is the maximum over every backward program of every route, so -one buffer of ``cake_causal_conv1d_train_backward_workspace_bytes(S, C, -device)`` serves whichever program a call selects. Before capturing a CUDA -Graph, build every library of the device's architecture (the backend's -``cake_jit.specs_for_arch``) so that no JIT build happens inside the capture; -the captured graph holds the program selected for its shape and route. - -Usage:: - - y = cake_causal_conv1d_train_forward(x, weight, cu_seqlens) - dx, dweight = cake_causal_conv1d_train_backward(x, weight, cu_seqlens, dy) - - # a bias (FP32 or BF16, like the taps) adds a third output - y = cake_causal_conv1d_train_forward(x, weight, cu_seqlens, bias=bias) - dx, dweight, dbias = cake_causal_conv1d_train_backward(x, weight, cu_seqlens, dy, bias=bias) - - # SiLU' at the unrounded FP32 preactivation for a ~2e-7 rel-L2 weight gradient - dx, dweight = cake_causal_conv1d_train_backward( - x, weight, cu_seqlens, dy, silu_grad_preactivation="fp32" - ) - - # autograd (saves x, weight, bias and cu_seqlens only) - y = cake_causal_conv1d(x, weight, cu_seqlens, bias) - - # caller-owned backward workspace (per-CTA FP32 partials, no zeroing needed) - ws = cake_causal_conv1d_train_backward_workspace(S, C, x.device) - dx, dweight = cake_causal_conv1d_train_backward(x, weight, cu_seqlens, dy, workspace=ws) - - # caller-owned forward workspace (zeroed int32 token-chunk counters, left zero - # by every launch: keep one per stream and reuse it for every later call) - fws = cake_causal_conv1d_train_forward_workspace(S, C, x.device) - y = cake_causal_conv1d_train_forward(x, weight, cu_seqlens, workspace=fws) - - # context-parallel training: this rank's fragments of the packed batch, with the - # history rows of the preceding fragments exchanged through the KCP context - # (FP32 taps, no bias) - y_local = cake_causal_conv1d_kcp(x_local, weight, ctx) - - # a layer's q / k / v convolutions with one halo exchange per direction - # (bitwise three single calls) - q, k, v = cake_causal_conv1d_kcp_multi((xq, xk, xv), (wq, wk, wv), ctx) -""" - -from __future__ import annotations - -from typing import Any, Optional, Sequence - -import torch - -from . import cake_causal_conv1d_train_backend as _backend -from .api_logging import flashinfer_api -from .cake_causal_conv1d_train_backend import ( # noqa: F401 - ARCHES, - SUPPORTED_COMPUTE_CAPABILITIES, - CakeCausalConv1dFunction, - cake_causal_conv1d_train_backward_workspace, - cake_causal_conv1d_train_backward_workspace_bytes, - cake_causal_conv1d_train_forward_workspace, - cake_causal_conv1d_train_forward_workspace_bytes, - is_cake_causal_conv1d_train_supported, -) - -# The three stable entry points are real ``def``s here (public-API tooling reads the -# decorated signatures of this module); their bodies live in the self-contained -# backend package next to the generated CUDA sources. The signatures mirror the -# backend's exactly (a test asserts it). - - -@flashinfer_api -def cake_causal_conv1d_train_forward( - x: torch.Tensor, - weight: torch.Tensor, - cu_seqlens: torch.Tensor, - *, - bias: Optional[torch.Tensor] = None, - out: Optional[torch.Tensor] = None, - workspace: Optional[torch.Tensor] = None, -) -> torch.Tensor: - r"""Packed width-4 depthwise causal convolution + SiLU, training forward. - - Parameters - ---------- - x : torch.Tensor - BF16 ``[S, C]`` or ``[1, S, C]`` contiguous packed activations of every - document; ``C`` is a multiple of 8. - weight : torch.Tensor - FP32 or BF16 ``[C, 4]`` contiguous taps; ``weight[c, 3]`` multiplies the - current token, ``weight[c, 0]`` the token three positions back. Any - other width is rejected. - cu_seqlens : torch.Tensor - ``int32 [num_docs + 1]`` contiguous device prefix sums of the document - lengths (``cu_seqlens[0] == 0``, ``cu_seqlens[-1] == S``; a repeated - value is an empty document). The host reads its metadata only: no - element is copied to the host and no stream is synchronized. - bias : torch.Tensor, optional - FP32 or BF16 ``[C]`` contiguous per-channel bias, added after the four - taps; its dtype is independent of ``weight``'s. - out : torch.Tensor, optional - Preallocated BF16 output of ``x``'s shape, contiguous. - workspace : torch.Tensor, optional - Caller-owned ``int32`` buffer from - :func:`cake_causal_conv1d_train_forward_workspace` (the token-chunk - counters: zero before its first launch, left zero by every launch, so - one buffer serves every later call of one stream, CUDA Graph replays - included; two launches that may run at once need two buffers). Its - size (:func:`cake_causal_conv1d_train_forward_workspace_bytes`) is the - most over every forward route. When omitted the package uses a zeroed - buffer it keeps per (device, current stream), or a zeroed buffer of the - call's own inside a CUDA Graph capture. - - Returns - ------- - torch.Tensor - BF16 ``y`` of ``x``'s shape: ``y[t, c] = bf16(silu(sum_k weight[c, k] * - x[t - 3 + k, c] + bias[c]))`` with zero history before each document's - first token. ``S == 0`` returns the empty output without launching. - """ - - return _backend.cake_causal_conv1d_train_forward( - x, weight, cu_seqlens, bias=bias, out=out, workspace=workspace - ) - - -@flashinfer_api -def cake_causal_conv1d_train_backward( - x: torch.Tensor, - weight: torch.Tensor, - cu_seqlens: torch.Tensor, - dy: torch.Tensor, - *, - bias: Optional[torch.Tensor] = None, - silu_grad_preactivation: str = "bf16", - workspace: Optional[torch.Tensor] = None, - dx: Optional[torch.Tensor] = None, - dweight: Optional[torch.Tensor] = None, - dbias: Optional[torch.Tensor] = None, -) -> tuple[torch.Tensor, ...]: - r"""Packed width-4 depthwise causal convolution + SiLU, training backward. - - Recomputes the FP32 preactivation from ``x``, evaluates the SiLU derivative - at it -- rounded to BF16 (``silu_grad_preactivation="bf16"``, the default) - or unrounded (``"fp32"``) -- and produces ``dx``, the weight gradient and, - with a bias, the bias gradient in two ordinary launches on the current - stream: the main program writes ``dx`` and per-CTA partial rows of the - gradients, and a reduce program launched programmatically dependent on it - sums the rows in a fixed order (no cooperative grid, so kernels resident - on other streams never stall the call). - - The call's arguments select the route (bias, BF16 among ``weight`` / - ``dweight``, the SiLU' policy) and the channel count selects the route's - program through a host-side rule of the backend's registry (never the - data), so a given ``(S, C, device)`` and route always run the same program. - Before capturing a CUDA Graph, build every library of the device's - architecture (the backend's ``cake_jit.specs_for_arch``) so that no JIT - build happens inside the capture; the captured graph holds the program - selected for its shape and route. - - Parameters - ---------- - x, weight, cu_seqlens, bias - The forward's inputs (same contracts as - :func:`cake_causal_conv1d_train_forward`); nothing else is saved. - dy : torch.Tensor - BF16 gradient of ``y``, ``x``'s shape, contiguous. - silu_grad_preactivation : {"bf16", "fp32"} - ``"bf16"``: SiLU' at the BF16-rounded recomputed preactivation (the - boundary flash-linear-attention's backward uses; ``dweight`` within - about 3e-5 rel-L2 of the exact FP64 gradient). ``"fp32"``: SiLU' at - the unrounded FP32 preactivation with an FP64 final reduction (about - 2e-7; a slower program). ``dx`` is within one BF16 ulp of exact under - both. - workspace : torch.Tensor, optional - Caller-owned ``uint8`` buffer from - :func:`cake_causal_conv1d_train_backward_workspace` (the per-CTA FP32 - ``dweight`` / ``dbias`` partials; fully written by the kernel, so it - needs no zeroing); allocated when omitted. Its size - (:func:`cake_causal_conv1d_train_backward_workspace_bytes`) is the - maximum over every backward program of every route, so one buffer - serves whichever program a call selects. - dx, dweight, dbias : torch.Tensor, optional - Preallocated outputs: BF16 of ``x``'s shape, ``[C, 4]`` and ``[C]`` - (each FP32 or BF16; the parameter's dtype when omitted), contiguous. - ``dbias`` needs ``bias``. - - Returns - ------- - ``(dx, dweight)``, with a third element ``dbias`` only when ``bias`` is - given: BF16 ``dx`` of ``x``'s shape, ``dx[t, c] = bf16(sum_k weight[c, k] * - dpre[t + 3 - k, c])`` over the tokens of the same document, the ``[C, 4]`` - ``dweight[c, k] = sum_t dpre[t, c] * x[t - 3 + k, c]`` and the ``[C]`` - ``dbias[c] = sum_t dpre[t, c]``, each reduced in a fixed order, so they are - bitwise reproducible for a given ``(S, C, device)`` and route. ``S == 0`` - returns the empty ``dx`` and zero gradients without launching. - """ - - return _backend.cake_causal_conv1d_train_backward( - x, - weight, - cu_seqlens, - dy, - bias=bias, - silu_grad_preactivation=silu_grad_preactivation, - workspace=workspace, - dx=dx, - dweight=dweight, - dbias=dbias, - ) - - -@flashinfer_api -def cake_causal_conv1d( - x: torch.Tensor, - weight: torch.Tensor, - cu_seqlens: torch.Tensor, - bias: Optional[torch.Tensor] = None, - *, - silu_grad_preactivation: str = "bf16", -) -> torch.Tensor: - """Autograd-enabled packed width-4 depthwise causal convolution + SiLU. - - Parameters - ---------- - x : torch.Tensor - BF16 ``[S, C]`` or ``[1, S, C]`` contiguous packed activations on a - supported CUDA device. - weight : torch.Tensor - FP32 or BF16 ``[C, 4]`` contiguous taps. - cu_seqlens : torch.Tensor - ``int32 [num_docs + 1]`` device prefix sums of the document lengths. - bias : torch.Tensor, optional - FP32 or BF16 ``[C]`` contiguous per-channel bias. - silu_grad_preactivation : {"bf16", "fp32"} - The weight gradient's SiLU' policy (see - :func:`cake_causal_conv1d_train_backward`). - - Returns - ------- - torch.Tensor - BF16 ``y`` of ``x``'s shape. The backward saves only ``x``, ``weight``, - ``bias`` and ``cu_seqlens`` and launches the fused ``dx`` + - deterministic ``dweight`` (+ ``dbias``) kernel of - :func:`cake_causal_conv1d_train_backward`; the parameters' gradients - come in the parameters' dtypes. - """ - - return _backend.cake_causal_conv1d( - x, weight, cu_seqlens, bias, silu_grad_preactivation=silu_grad_preactivation - ) - - -@flashinfer_api -def cake_causal_conv1d_kcp( - x_local: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - *, - recompute_handle: Optional[Any] = None, -) -> torch.Tensor: - """Autograd-enabled context-parallel packed width-4 depthwise causal convolution + SiLU of this rank's fragments. - - The first tokens of a fragment that continues a document read the history rows of the preceding fragment - (exchanged through the context's halo collective) instead of zeros, so the result equals the single-rank - :func:`cake_causal_conv1d` of the whole batch on this rank's tokens. The KCP path takes FP32 taps and no - bias (the package's pre-round-6 contract). - - Parameters - ---------- - x_local : torch.Tensor - BF16 ``[S_local, C]`` or ``[1, S_local, C]`` contiguous packed activations of this rank's fragments, in the - partition's order; ``C`` is a multiple of 8. - weight : torch.Tensor - FP32 ``[C, 4]`` contiguous taps. - ctx : KcpContext - From ``flashinfer.kda_train.build_kcp_context`` (the same context the KDA step of the rank uses). - recompute_handle : KcpRecomputeHandle, optional - Keyword-only, default ``None`` (unchanged behavior). A handle (any object with an ``entries`` dict, such as - ``flashinfer.cake_causal_conv1d_train.KcpRecomputeHandle``) makes the checkpoint recompute of - this call run without the forward halo exchange: the original forward exchanges as always and stores its - history rows (the saved set's own tensor, key ``"conv"``) in the handle; torch's recompute of the forward - inside the backward reads them back, so ``y`` and the gradients are bitwise those of the same checkpoint - without a handle. The backward still exchanges. Create one handle per convolution call and micro-step, - outside the outermost checkpoint region, and pass it as an argument of the checkpointed function. A - forward outside the backward that finds an entry raises (a handle reused across micro-steps). The forward - halo exchange is a collective: every rank of ``ctx.group`` passes a handle for the same call or none, created - at the same checkpoint nesting level. A mismatch pairs one rank's recompute exchange with another rank's - backward exchange and hangs or corrupts data; no check is made. Never share one handle between two - convolution calls (also not in nested regions), and never ``clear()`` or reuse a handle while a backward - that needs it is pending. - - Returns - ------- - torch.Tensor - BF16 ``y`` of ``x_local``'s shape. The backward runs the head program, the backward halo exchange and the - fused ``dx`` + deterministic FP32 ``dweight`` kernel of this rank's fragments. - """ - - from .cake_causal_conv1d_train_backend import kcp as _kcp - - return _kcp.cake_causal_conv1d_kcp( - x_local, weight, ctx, recompute_handle=recompute_handle - ) - - -@flashinfer_api -def cake_causal_conv1d_kcp_multi( - xs: Sequence[torch.Tensor], - weights: Sequence[torch.Tensor], - ctx: Any, - *, - recompute_handle: Optional[Any] = None, -) -> tuple[torch.Tensor, ...]: - """Autograd-enabled context-parallel convolutions of several tensors of one layer (for example its q, k and v) - with one halo exchange per direction. - - The result is bitwise ``tuple(cake_causal_conv1d_kcp(x, w, ctx) for x, w in zip(xs, weights))``: outputs and - every gradient, ``dweight`` included (each tensor keeps its own kernels and fixed-order reduction). The forward - issues one halo collective instead of one per tensor (the history rows of all tensors travel side by side and are - copied out after the wait), and so does the backward, so a layer has fewer global synchronization points. - - Parameters - ---------- - xs : sequence of torch.Tensor - ``n >= 1`` BF16 ``[S_local, C_i]`` or ``[1, S_local, C_i]`` contiguous packed activations of this rank's - fragments, in the partition's order; each ``C_i`` is a multiple of 8 (they may differ). - weights : sequence of torch.Tensor - FP32 ``[C_i, 4]`` contiguous taps, one per tensor. - ctx : KcpContext - As in :func:`cake_causal_conv1d_kcp`. Every rank of the group calls this function once per layer, in the - same order, with the same tensor count and channel counts. The ranks are not compared (that would cost a - collective): a different count, channel count or order hangs the exchange or silently corrupts the history - rows, as a mismatch of the single entry's per-rank contract does. - recompute_handle : KcpRecomputeHandle, optional - Keyword-only, default ``None`` (unchanged behavior). One handle for the whole call: the original forward - stores the ``n`` history pools under ``"conv"`` and torch's recompute inside the backward reads them back, - so the recompute issues no collective; the backward still exchanges once. The rules of - :func:`cake_causal_conv1d_kcp` apply (one handle per call and micro-step, created outside the outermost - checkpoint region; every rank passes one or none). The same handle may carry the layer's ``"kda"`` entry; - its ``"conv"`` entry serves one multi call or one single call, never both. - - Returns - ------- - tuple of torch.Tensor - The ``n`` BF16 outputs, each of its input's shape. - """ - - from .cake_causal_conv1d_train_backend import kcp as _kcp - - return _kcp.cake_causal_conv1d_kcp_multi( - xs, weights, ctx, recompute_handle=recompute_handle - ) - - -def __getattr__(name: str) -> Any: - # ``KcpRecomputeHandle`` (the KCP route's checkpoint recompute handle) resolves on first access, so importing the - # facade does not import the KCP host module (torch.distributed, the KCP loader), like the entry point above. - if name == "KcpRecomputeHandle": - from .cake_causal_conv1d_train_backend.kcp import KcpRecomputeHandle - - return KcpRecomputeHandle - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = [ - "ARCHES", - "SUPPORTED_COMPUTE_CAPABILITIES", - "CakeCausalConv1dFunction", - "KcpRecomputeHandle", # noqa: F822 - resolved by __getattr__ - "cake_causal_conv1d", - "cake_causal_conv1d_kcp", - "cake_causal_conv1d_kcp_multi", - "cake_causal_conv1d_train_backward", - "cake_causal_conv1d_train_backward_workspace", - "cake_causal_conv1d_train_backward_workspace_bytes", - "cake_causal_conv1d_train_forward", - "cake_causal_conv1d_train_forward_workspace", - "cake_causal_conv1d_train_forward_workspace_bytes", - "is_cake_causal_conv1d_train_supported", -] +sys.modules[__name__] = import_module("flashinfer.cake_causal_conv1d_det_train") diff --git a/flashinfer/cake_causal_conv1d_train_backend/__init__.py b/flashinfer/cake_causal_conv1d_train_backend/__init__.py index 31143b02340..e932bc99f55 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/__init__.py +++ b/flashinfer/cake_causal_conv1d_train_backend/__init__.py @@ -1,873 +1,15 @@ -""" -Copyright (c) 2026 by FlashInfer team. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. +# Copyright (c) 2026 by FlashInfer team. Licensed under Apache-2.0. +"""Compatibility import for flashinfer.cake_causal_conv1d_det_train_backend. -Backend of :mod:`flashinfer.cake_causal_conv1d_train`: the Cake packed width-4 -depthwise causal convolution + SiLU training forward / backward (SM100 / SM103 / -SM107) as one self-contained package (host API bodies, the ``cake_jit`` registry, -the pinned ``cake_build`` loader and the generated CUDA sources under ``csrc/``). -The stable, ``@flashinfer_api``-decorated entry points and the full user-facing -documentation live in :mod:`flashinfer.cake_causal_conv1d_train`; import that -module. - -Host design (the knobs are keywords; no environment variable): - -- Taps ``weight[C, 4]`` are FP32 or BF16; an optional per-channel ``bias[C]`` - is FP32 or BF16, independently of the taps. The kernels widen BF16 to FP32 - on load and compute in FP32; the bias is added LAST, after the four taps - (FLA's order). ``dweight`` / ``dbias`` default to the parameter's dtype and - may be preallocated as FP32 or BF16. -- ``silu_grad_preactivation``: ``"bf16"`` (default, the shipped policy: SiLU' - at the BF16-rounded preactivation, the boundary FLA's backward uses) or - ``"fp32"`` (SiLU' at the unrounded FP32 preactivation and an FP64 final - reduction of ``dweight``: a weight gradient within ~2e-7 rel-L2 of exact). -- A call's arguments select its route (``cake_jit.forward_route`` / - ``backward_route``): the default routes run the shipped programs unchanged; - the other routes' programs read the ``flags`` bit mask (``cake_jit.FLAG_*``, - derived here from the actual tensors: taps, the ``dweight`` out tensor, the - bias, the ``dbias`` out tensor) for the BF16 / FP32 load and store forms. - An architecture that lacks a route refuses the call with a ``ValueError`` - naming the route. -- The backward workspace is sized for every backward program of the device's - architecture (the bias programs' partials carry a fifth, dbias, slab), so one - buffer serves every route, token count and CUDA Graph. -- The forward (the dynamic token-chunk programs) takes its token chunks from - int32 counters, two per channel tile, in a forward workspace that is zero - before its first launch and that every launch leaves zero: a caller-owned - buffer (``workspace=``, from :func:`cake_causal_conv1d_train_forward_workspace`) - or, when omitted, a zeroed buffer this module keeps per (device, current - stream), grown as needed and never zeroed again; inside a CUDA Graph capture - an omitted workspace is a zeroed buffer of the call's own (its zeroing is - captured with the launch), so a graph never shares counters with eager - calls or with another graph. +An attribute of this package resolves on the det package at first access, and the submodules ``cake_build``, +``cake_jit``, ``kcp`` and ``kcp_jit`` are the det package's module objects. Importing this package imports +nothing else: the KCP host module loads only when ``kcp`` is imported. """ -from __future__ import annotations - -import functools -from typing import Optional - -import torch - -from . import cake_jit as _loader -from .cake_jit import ( - ARCH_BY_CAPABILITY, - ARCHES, - CHANNEL_MULTIPLE, - FLAG_B_BF16, - FLAG_DB_BF16, - FLAG_DW_BF16, - FLAG_W_BF16, - INT32_MAX, - REQUIRED_KINDS, - WIDTH, - arch_for_capability, - backward_route, - counter_words, - forward_route, - forward_workspace_words, - launch_grid, - partial_bytes, - partial_rows, - plan, - route_applies, - route_modules, - run_module, - run_modules, - selected_modules, - workspace_layout, - workspace_owners, -) - -SUPPORTED_COMPUTE_CAPABILITIES: tuple[tuple[int, int], ...] = tuple( - sorted(ARCH_BY_CAPABILITY) -) -# The dtypes of the taps, the bias and their gradients (FP32 or BF16, any combination). -PARAMETER_DTYPES = (torch.float32, torch.bfloat16) -# The two SiLU'-evaluation policies of the weight gradient (``silu_grad_preactivation``). -SILU_GRAD_PREACTIVATIONS = ("bf16", "fp32") -# The forward workspaces of the calls that pass none: (device index, CUDA stream -# handle) -> a zeroed int32 buffer, grown as needed. The launches of one stream -# run in order and each leaves its counters zero, so the stream's buffer serves -# every later forward call on it. -_FORWARD_WORKSPACES: dict[tuple[int, int], torch.Tensor] = {} - - -def _device_index(device: torch.device) -> int: - return device.index if device.index is not None else torch.cuda.current_device() - - -@functools.cache -def _arch_of_index(device_index: int) -> Optional[str]: - """The exported architecture of one CUDA device index, queried once per process.""" - - return arch_for_capability(tuple(torch.cuda.get_device_capability(device_index))) - - -def _device_arch(device: torch.device) -> str: - if device.type != "cuda": - raise ValueError( - f"Cake causal conv1d training requires CUDA tensors, got {device}" - ) - arch = _arch_of_index(_device_index(device)) - if arch is None: - raise ValueError( - "Cake causal conv1d training supports compute capability " - f"{sorted(ARCH_BY_CAPABILITY)} only, got " - f"{tuple(torch.cuda.get_device_capability(device))}" - ) - return arch - - -def is_cake_causal_conv1d_train_supported(device: torch.device) -> bool: - """Whether ``device`` has the exported routes every architecture ships (``fwd`` and ``bwd``). - - The optional routes (bias, BF16 taps / gradients, FP32 SiLU') are checked - per call: a call that needs one the architecture lacks raises a - ``ValueError`` naming it. - """ - - device = torch.device(device) - if device.type != "cuda": - return False - capability = tuple(torch.cuda.get_device_capability(device)) - return all( - route_applies(device_capability=capability, kind=kind) - for kind in REQUIRED_KINDS - ) - - -def _require_tensor(tensor: object, name: str) -> torch.Tensor: - if not isinstance(tensor, torch.Tensor): - raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") - return tensor - - -def _check_placement( - tensor: torch.Tensor, - name: str, - *, - device: torch.device, - dtype: torch.dtype | tuple[torch.dtype, ...], -) -> None: - """On ``device``, of one of the ``dtype``\\ s, contiguous (``ValueError`` otherwise).""" - - dtypes = (dtype,) if isinstance(dtype, torch.dtype) else tuple(dtype) - if tensor.device != device: - raise ValueError(f"{name} must live on {device}, got {tensor.device}") - if tensor.dtype not in dtypes: - wanted = " or ".join(str(d) for d in dtypes) - raise ValueError(f"{name} must be {wanted}, got {tensor.dtype}") - if not tensor.is_contiguous(): - raise ValueError(f"{name} must be contiguous") - - -def _tokens( - tensor: object, - name: str, - *, - device: torch.device, - shape: Optional[tuple[int, ...]] = None, -) -> tuple[int, int, torch.Tensor]: - """Validate a BF16 ``[S, C]`` or ``[1, S, C]`` token tensor; return ``(S, C, [S, C] view)``. - - ``shape`` pins the exact shape another token tensor of the call has (``dy`` - against ``x``). Both kernels read the packed rows as 16-byte vectors of eight - channels, so ``C`` is a multiple of eight; ``S`` and ``C`` are 32-bit kernel - scalars (``S x C`` itself may exceed 2^31: the kernels address with 64-bit - offsets). - """ - - t = _require_tensor(tensor, name) - if t.ndim == 3 and t.shape[0] == 1: - S, C = int(t.shape[1]), int(t.shape[2]) - elif t.ndim == 2: - S, C = int(t.shape[0]), int(t.shape[1]) - else: - raise ValueError( - f"{name} must be a [S, C] or [1, S, C] tensor, got shape {tuple(t.shape)}" - ) - if shape is not None and tuple(t.shape) != tuple(shape): - raise ValueError(f"{name} must have shape {tuple(shape)}, got {tuple(t.shape)}") - _check_placement(t, name, device=device, dtype=torch.bfloat16) - if C <= 0 or C % CHANNEL_MULTIPLE: - raise ValueError( - f"{name} must have a channel count C that is a positive multiple of " - f"{CHANNEL_MULTIPLE}, got C={C}" - ) - if S > INT32_MAX or C > INT32_MAX: - raise ValueError( - f"{name}: S={S} and C={C} must each fit a 32-bit kernel scalar" - ) - return S, C, t.view(S, C) - - -def _weight( - weight: object, - C: int, - *, - device: torch.device, - name: str = "weight", - dtypes: tuple[torch.dtype, ...] = (torch.float32,), -) -> torch.Tensor: - """Validate the ``[C, 4]`` taps of one of ``dtypes`` (the training API passes FP32 and BF16; the KCP path - keeps the default, FP32 only).""" - - w = _require_tensor(weight, name) - if w.ndim != 2 or int(w.shape[1]) != WIDTH: - raise ValueError( - f"{name} must be a [C, {WIDTH}] tensor of a width-{WIDTH} convolution " - f"(tap {WIDTH - 1} multiplies the current token), got shape {tuple(w.shape)}" - ) - if int(w.shape[0]) != C: - raise ValueError( - f"{name} must have one row per channel ([{C}, {WIDTH}]), got shape {tuple(w.shape)}" - ) - _check_placement(w, name, device=device, dtype=dtypes) - return w - - -def _bias(bias: object, C: int, *, device: torch.device) -> torch.Tensor: - """Validate the per-channel ``[C]`` bias (FP32 or BF16, independently of the taps).""" - - b = _require_tensor(bias, "bias") - if b.ndim != 1 or int(b.shape[0]) != C: - raise ValueError( - f"bias must be a [C] tensor with one value per channel ([{C}]), got shape {tuple(b.shape)}" - ) - _check_placement(b, "bias", device=device, dtype=PARAMETER_DTYPES) - return b - - -def _cu_seqlens( - cu_seqlens: object, *, device: torch.device -) -> tuple[torch.Tensor, int]: - """Validate the ``int32 [num_docs + 1]`` device prefix sums; return ``(cu_seqlens, num_docs)``. - - The host reads the tensor's metadata only: no element is copied to the host - and no stream is synchronized, so the call is CUDA Graph capturable. - """ - - cu = _require_tensor(cu_seqlens, "cu_seqlens") - if cu.ndim != 1: - raise ValueError( - f"cu_seqlens must be a 1-D [num_docs + 1] tensor, got shape {tuple(cu.shape)}" - ) - _check_placement(cu, "cu_seqlens", device=device, dtype=torch.int32) - num_docs = int(cu.numel()) - 1 - if num_docs < 1: - raise ValueError( - "cu_seqlens must hold at least two prefix sums (one document), got " - f"{int(cu.numel())} elements" - ) - return cu, num_docs - - -def _output( - tensor: Optional[object], - name: str, - *, - shape: tuple[int, ...], - dtype: torch.dtype, - device: torch.device, - dtypes: Optional[tuple[torch.dtype, ...]] = None, -) -> torch.Tensor: - """A caller-provided output (exact shape, contiguous, of ``dtype`` or of one of ``dtypes`` when given) or a - fresh ``dtype`` one.""" - - if tensor is None: - return torch.empty(shape, dtype=dtype, device=device) - out = _require_tensor(tensor, name) - if tuple(out.shape) != tuple(shape): - raise ValueError( - f"{name} must have shape {tuple(shape)}, got {tuple(out.shape)}" - ) - _check_placement(out, name, device=device, dtype=dtypes or dtype) - return out - - -def _silu_grad_preactivation(value: object) -> str: - if value not in SILU_GRAD_PREACTIVATIONS: - raise ValueError( - f"silu_grad_preactivation must be one of {SILU_GRAD_PREACTIVATIONS}, got {value!r}" - ) - return str(value) - - -def _scalars(scalars: dict[str, int]) -> dict[str, int]: - """The launch scalars :func:`cake_jit.plan` produced (``flags`` is set beside them by the caller).""" - - return {key: int(scalars[key]) for key in _loader.PLAN_SCALARS} - - -def _flags( - w: torch.Tensor, - dw: Optional[torch.Tensor] = None, - b: Optional[torch.Tensor] = None, - db: Optional[torch.Tensor] = None, -) -> int: - """The ``flags`` bit mask of a ``dtype_flags`` program from the ACTUAL tensors of the call.""" - - flags = 0 - for tensor, bit in ( - (w, FLAG_W_BF16), - (dw, FLAG_DW_BF16), - (b, FLAG_B_BF16), - (db, FLAG_DB_BF16), - ): - if tensor is not None and tensor.dtype == torch.bfloat16: - flags |= bit - return flags - - -def _route(arch: str, kind: str, *, needs: str, C: int, S: int) -> tuple[str, ...]: - """The programs ``kind`` launches on ``arch`` for the call (``ValueError`` naming the route when the - architecture's export lacks it).""" - - if kind not in _loader.ROUTES.get(arch, {}): - raise ValueError( - f"{needs} needs the {kind!r} route, which this package does not export for " - f"{arch}; its routes: {sorted(_loader.ROUTES.get(arch, {}))}" - ) - return selected_modules(arch, kind, C=C, S=S) - +from importlib import import_module -def _launch( - record: dict, S: int, C: int, num_docs: int, *, device_index: int -) -> tuple[dict[str, int], tuple[int, int, int]]: - """The kernel scalars and the grid of one launch of ``record`` (the host plan of ``cake_jit``). - CTA ``b`` owns channel tile ``b % n_ch_tiles``, so every channel tile needs - a CTA of the launched program's grid on the device: a ``C`` beyond ``grid - x threads x vec`` channels is refused rather than left partly unwritten. - The caller launches this grid (``run_module`` / ``run_modules`` ``grid``), - so a call plans it once. - """ - - scalars = plan( - S, - C, - num_docs, - geometry=record["grid_rule"]["geometry"], - role=record["role"], - ) - grid = launch_grid(record, device_index=device_index, C=C, S=S) - if scalars["n_ch_tiles"] > grid[0]: - raise ValueError( - f"C={C} spans {scalars['n_ch_tiles']} channel tiles but the launch grid on this " - f"device is {grid[0]} CTAs; every channel tile needs a CTA" - ) - return scalars, grid - - -def cake_causal_conv1d_train_backward_workspace_bytes( - S: int, C: int, device: torch.device -) -> int: - """Bytes of the caller-owned backward workspace for ``C`` channels on ``device``. - - The workspace holds the FP32 ``[rows, C, 4]`` ``dweight`` partials (and, for - the bias routes, a ``[rows, C]`` dbias slab after them) of the backward's - main program, one row per ``n_ch_tiles`` CTAs of a program's launch grid - (``rows = ceil(grid / n_ch_tiles)``, ``grid = grid_mult x ctas_per_sm x SM - count``: the bytes grow with the program's over-decomposition factor - ``grid_mult``; a call whose grid drops CTAs to whole channel-tile ranks, - ``cake_jit.main_grid``, writes fewer rows), sized by the largest such - buffer over every backward program of every backward route of the - architecture: the call's arguments select the route and the channel count - the program, every program writes the leading rows of its own grid inside - this buffer, and its reduce - program sums exactly those. So one buffer serves every route (the bias - routes' fifth slab makes it 25 % larger than the dweight partials alone, - for every caller) and, as it depends on ``C`` and the device only, every - token count; the main program fully writes the rows the reduce program - sums, so the buffer needs no zeroing between launches or CUDA Graph - replays. - """ - - device = torch.device(device) - S, C = int(S), int(C) - if S < 0: - raise ValueError("S must be non-negative") - if C <= 0 or C % CHANNEL_MULTIPLE: - raise ValueError( - f"C must be a positive multiple of {CHANNEL_MULTIPLE}, got {C}" - ) - arch = _device_arch(device) - device_index = _device_index(device) - return max( - int(workspace_layout(owner, C, device_index=device_index)["total_bytes"]) - for owner in workspace_owners(arch) - ) - - -def cake_causal_conv1d_train_backward_workspace( - S: int, C: int, device: torch.device -) -> torch.Tensor: - """Allocate a backward workspace (see ``..._workspace_bytes``; no zeroing is needed).""" - - nbytes = cake_causal_conv1d_train_backward_workspace_bytes(S, C, device) - return torch.empty(nbytes, dtype=torch.uint8, device=torch.device(device)) - - -def cake_causal_conv1d_train_forward_workspace_bytes( - S: int, C: int, device: torch.device -) -> int: - """Bytes of the caller-owned forward workspace for ``C`` channels on ``device``. - - The forward's programs take their token chunks from int32 counters, two - per channel tile (the next-chunk counter and the done-CTA counter of the - tile), sized by the most over every forward program of every forward route - of the architecture (``cake_jit.forward_workspace_words``), so one buffer - serves every route; the size depends on ``C`` only (``S`` is checked and - otherwise unused, as for the backward helper). The buffer must be zero - before its first launch; every launch leaves it zero, so one buffer serves - every later forward call of one stream, CUDA Graph replays included, - without zeroing. Two launches that may run at once (on two streams, or - two graphs replayed concurrently) need two buffers. - """ - - device = torch.device(device) - S, C = int(S), int(C) - if S < 0: - raise ValueError("S must be non-negative") - if C <= 0 or C % CHANNEL_MULTIPLE: - raise ValueError( - f"C must be a positive multiple of {CHANNEL_MULTIPLE}, got {C}" - ) - return 4 * forward_workspace_words(_device_arch(device), C) - - -def cake_causal_conv1d_train_forward_workspace( - S: int, C: int, device: torch.device -) -> torch.Tensor: - """Allocate a zeroed int32 forward workspace (see ``..._workspace_bytes``); keep it for every later forward - call of one stream.""" - - nbytes = cake_causal_conv1d_train_forward_workspace_bytes(S, C, device) - return torch.zeros(nbytes // 4, dtype=torch.int32, device=torch.device(device)) - - -def _forward_work( - workspace: object, words: int, *, device: torch.device, device_index: int -) -> Optional[torch.Tensor]: - """The ``work`` counters of one forward launch needing ``words`` int32 words. - - A caller-owned ``workspace`` (1-D int32, contiguous, on ``device``, at least - ``words`` words, zero before its first launch) is used as is. Without one: - inside a CUDA Graph capture a fresh zeroed buffer (its zeroing is captured - with the launch, so the graph shares no counters with eager calls or other - graphs); else the zeroed buffer this module keeps for the device and the - current stream, replaced by a larger zeroed one when it is too small (the - stream-ordered allocator frees the old one after the launches queued on - the stream). ``None`` when the program takes no counters. - """ - - if workspace is not None: - ws = _require_tensor(workspace, "workspace") - if ws.ndim != 1: - raise ValueError( - f"workspace must be a 1-D int32 CUDA tensor, got shape {tuple(ws.shape)}" - ) - _check_placement(ws, "workspace", device=device, dtype=torch.int32) - if ws.numel() < words: - raise ValueError( - f"workspace holds {ws.numel()} int32 words; this call needs at least " - f"{words}; allocate cake_causal_conv1d_train_forward_workspace(S, C, device)" - ) - return ws - if words == 0: - return None - if torch.cuda.is_current_stream_capturing(): - return torch.zeros(words, dtype=torch.int32, device=device) - key = (device_index, int(torch.cuda.current_stream(device_index).cuda_stream)) - held = _FORWARD_WORKSPACES.get(key) - if held is None or held.numel() < words: - held = torch.zeros(words, dtype=torch.int32, device=device) - _FORWARD_WORKSPACES[key] = held - return held - - -def _partial_view( - workspace: object, record: dict, *, C: int, rows: int, device: torch.device -) -> torch.Tensor: - """The FP32 ``[rows, C, taps]`` capacity view of this launch's partial rows inside ``workspace``. - - ``taps`` is the record's (4, or 5 for a bias program whose kernel lays the - ``[rows, C, 4]`` tap slab out first and the ``[rows, C]`` dbias slab after - it); the view only sizes the kernel's buffer. - """ - - ws = _require_tensor(workspace, "workspace") - if ws.ndim != 1: - raise ValueError( - f"workspace must be a 1-D uint8 CUDA tensor, got shape {tuple(ws.shape)}" - ) - _check_placement(ws, "workspace", device=device, dtype=torch.uint8) - if ws.data_ptr() % 16 != 0: - raise ValueError("workspace must be 16-byte aligned") - nbytes = partial_bytes(record, C, rows) - if ws.numel() < nbytes: - raise ValueError( - f"workspace holds {ws.numel()} bytes; this call needs at least " - f"{nbytes}; allocate " - "cake_causal_conv1d_train_backward_workspace_bytes(S, C, device) bytes" - ) - taps = int(record["workspace_rule"]["taps"]) - return ws[:nbytes].view(torch.float32).view(rows, C, taps) - - -def cake_causal_conv1d_train_forward( - x: torch.Tensor, - weight: torch.Tensor, - cu_seqlens: torch.Tensor, - *, - bias: Optional[torch.Tensor] = None, - out: Optional[torch.Tensor] = None, - workspace: Optional[torch.Tensor] = None, -) -> torch.Tensor: - r"""Packed width-4 depthwise causal convolution + SiLU, training forward. - - Parameters - ---------- - x : torch.Tensor - BF16 ``[S, C]`` or ``[1, S, C]`` contiguous packed activations of every - document; ``C`` is a multiple of 8. - weight : torch.Tensor - FP32 or BF16 ``[C, 4]`` contiguous taps; ``weight[c, 3]`` multiplies the - current token, ``weight[c, 0]`` the token three positions back. - cu_seqlens : torch.Tensor - ``int32 [num_docs + 1]`` contiguous device prefix sums of the document - lengths (``cu_seqlens[0] == 0``, ``cu_seqlens[-1] == S``; a repeated - value is an empty document). The host reads its metadata only. - bias : torch.Tensor, optional - FP32 or BF16 ``[C]`` contiguous per-channel bias (its dtype is - independent of ``weight``'s), added after the four taps. - out : torch.Tensor, optional - Preallocated BF16 output of ``x``'s shape, contiguous. - workspace : torch.Tensor, optional - Caller-owned int32 buffer from - :func:`cake_causal_conv1d_train_forward_workspace` (the token-chunk - counters, zero before its first launch; every launch leaves it zero, - so one buffer serves every later call of one stream). Two launches - that may run at once need two buffers. When omitted the package - uses a zeroed buffer it keeps per (device, current stream), or a - fresh zeroed one inside a CUDA Graph capture. - - Returns - ------- - torch.Tensor - BF16 ``y`` of ``x``'s shape: ``y[t, c] = bf16(silu(sum_k weight[c, k] * - x[t - 3 + k, c] + bias[c]))`` with zero history before each document's - first token. ``S == 0`` returns the empty output without launching. - """ - - x = _require_tensor(x, "x") - device = x.device - arch = _device_arch(device) - S, C, x_flat = _tokens(x, "x", device=device) - w = _weight(weight, C, device=device, dtypes=PARAMETER_DTYPES) - b = None if bias is None else _bias(bias, C, device=device) - cu, num_docs = _cu_seqlens(cu_seqlens, device=device) - y = _output(out, "out", shape=tuple(x.shape), dtype=torch.bfloat16, device=device) - if S == 0: - return y - kind = forward_route(bias=b is not None, bf16_taps=w.dtype == torch.bfloat16) - names = _route(arch, kind, needs="this forward (bias / BF16 taps)", C=C, S=S) - device_index = _device_index(device) - scalars, grid = _launch( - _loader.MODULES[names[0]], S, C, num_docs, device_index=device_index - ) - values: dict[str, object] = { - "x": x_flat, - "w": w, - "cu": cu, - "y": y.view(S, C), - **_scalars(scalars), - "flags": _flags(w, b=b), - } - if b is not None: - values["b"] = b - words = max( - counter_words(_loader.MODULES[name], scalars["n_ch_tiles"]) for name in names - ) - work = _forward_work(workspace, words, device=device, device_index=device_index) - if work is not None: - values["work"] = work - for name in names: - run_module( - name, - values, - device_index=device_index, - grid=grid if name == names[0] else None, - ) - return y - - -def cake_causal_conv1d_train_backward( - x: torch.Tensor, - weight: torch.Tensor, - cu_seqlens: torch.Tensor, - dy: torch.Tensor, - *, - bias: Optional[torch.Tensor] = None, - silu_grad_preactivation: str = "bf16", - workspace: Optional[torch.Tensor] = None, - dx: Optional[torch.Tensor] = None, - dweight: Optional[torch.Tensor] = None, - dbias: Optional[torch.Tensor] = None, -) -> tuple[torch.Tensor, ...]: - r"""Packed width-4 depthwise causal convolution + SiLU, training backward. - - Parameters - ---------- - x, weight, cu_seqlens, bias - The forward's inputs (same contracts as - :func:`cake_causal_conv1d_train_forward`); nothing else is saved. - dy : torch.Tensor - BF16 gradient of ``y``, ``x``'s shape, contiguous. - silu_grad_preactivation : {"bf16", "fp32"} - Where the SiLU derivative of the weight gradient is evaluated: - ``"bf16"`` at the BF16-rounded recomputed preactivation (the shipped - policy and FLA's boundary; ``dweight`` within ~3e-5 rel-L2 of exact), - ``"fp32"`` at the unrounded FP32 preactivation with an FP64 final - reduction (~2e-7). ``dx`` is within one BF16 ulp of exact either way. - workspace : torch.Tensor, optional - Caller-owned ``uint8`` buffer from - :func:`cake_causal_conv1d_train_backward_workspace`; allocated when - omitted. - dx, dweight, dbias : torch.Tensor, optional - Preallocated outputs: BF16 of ``x``'s shape, ``[C, 4]`` and ``[C]`` - (FP32 or BF16 each; the parameter's dtype when omitted), contiguous. - ``dbias`` needs ``bias``. - - Returns - ------- - ``(dx, dweight)``, or ``(dx, dweight, dbias)`` when ``bias`` is given: BF16 - ``dx`` of ``x``'s shape and the ``[C, 4]`` weight gradient (and ``[C]`` bias - gradient), each accumulated in a fixed order, so they are bitwise - reproducible for a given ``(S, C, device)`` and route. ``S == 0`` returns - the empty ``dx`` and zero gradients without launching. - - The call's arguments select the route (``cake_jit.backward_route``: - ``bias``, the policy, BF16 among ``weight`` / ``dweight``), and within a - route the channel count selects the program through the registry's - host-side rule (``cake_jit.selected_modules``; the data is never read), so a - given ``(S, C, device)`` and route always run the same program. The - ``workspace`` is sized for every program of every route - (:func:`cake_causal_conv1d_train_backward_workspace_bytes`); the selected - program uses the leading rows of its own grid. - - A call is two ordinary launches on the current stream: the selected - program writes ``dx`` and one FP32 partial row of the gradients per CTA, - then the route's reduce program, launched programmatically dependent on - it (the launch latency overlaps the main kernel's tail; the dependency is - a stream-order edge, so the pair captures into a CUDA Graph), sums the - rows per channel in a fixed order into ``dweight`` (and ``dbias``) in the - requested dtypes. Neither launch is cooperative, so the backward keeps - its speed while kernels of other streams (NCCL collectives) are resident: - a cooperative grid could not start until they finished. Build every - library of the device's architecture (``cake_jit.specs_for_arch``) before - capturing a CUDA Graph: the capture holds the programs selected for its - shape and route, and no JIT build may happen inside it. - """ - - policy = _silu_grad_preactivation(silu_grad_preactivation) - x = _require_tensor(x, "x") - device = x.device - arch = _device_arch(device) - S, C, x_flat = _tokens(x, "x", device=device) - w = _weight(weight, C, device=device, dtypes=PARAMETER_DTYPES) - b = None if bias is None else _bias(bias, C, device=device) - cu, num_docs = _cu_seqlens(cu_seqlens, device=device) - _, _, dy_flat = _tokens(dy, "dy", device=device, shape=tuple(x.shape)) - dx_out = _output( - dx, "dx", shape=tuple(x.shape), dtype=torch.bfloat16, device=device - ) - dw_out = _output( - dweight, - "dweight", - shape=(C, WIDTH), - dtype=w.dtype, - device=device, - dtypes=PARAMETER_DTYPES, - ) - if b is None: - if dbias is not None: - raise ValueError( - "dbias is the bias gradient's output; pass bias to produce it" - ) - db_out = None - else: - db_out = _output( - dbias, - "dbias", - shape=(C,), - dtype=b.dtype, - device=device, - dtypes=PARAMETER_DTYPES, - ) - if S == 0: - dw_out.zero_() - if db_out is None: - return dx_out, dw_out - db_out.zero_() - return dx_out, dw_out, db_out - kind = backward_route( - bias=b is not None, - fp32_grad=policy == "fp32", - bf16_taps_or_dweight=torch.bfloat16 in (w.dtype, dw_out.dtype), - ) - # Within the route the program is a pure function of the architecture and C - # (S never enters): its own grid sizes the partial rows this call writes - # inside the workspace. - names = _route( - arch, - kind, - needs=f"this backward (bias / BF16 taps or dweight / silu_grad_preactivation={policy!r})", - C=C, - S=S, +def __getattr__(name): + return getattr( + import_module("flashinfer.cake_causal_conv1d_det_train_backend"), name ) - record = _loader.MODULES[names[0]] - if record["workspace_rule"] is None: - raise RuntimeError( - f"the backward program {names[0]} selected for C={C} on {arch} owns no " - "partial workspace" - ) - device_index = _device_index(device) - scalars, grid = _launch(record, S, C, num_docs, device_index=device_index) - grid_x = grid[0] - if workspace is None: - workspace = cake_causal_conv1d_train_backward_workspace(S, C, device) - rows = partial_rows(grid_x, scalars["n_ch_tiles"]) - partial = _partial_view(workspace, record, C=C, rows=rows, device=device) - values: dict[str, object] = { - "x": x_flat, - "w": w, - "cu": cu, - "dy": dy_flat, - "dx": dx_out.view(S, C), - "dw": dw_out, - "partial": partial, - **_scalars(scalars), - "flags": _flags(w, dw_out, b, db_out), - "main_ctas": grid_x, - } - if b is not None: - values["b"] = b - values["db"] = db_out - # The main program (an ordinary launch) writes dx and the partial rows; the - # route's reduce program, programmatically dependent on it through its - # binding's launch attribute, sums into dweight (and dbias) on the same - # stream the rows of each channel tile that a grid of main_ctas CTAs wrote. - # Both are issued back to back in one stream context, so the reduce is - # queued while the main program runs. - run_modules( - ((names[0], grid), (record["reduce"], None)), - values, - device_index=device_index, - ) - if db_out is None: - return dx_out, dw_out - return dx_out, dw_out, db_out - - -class CakeCausalConv1dFunction(torch.autograd.Function): - """Autograd wrapper saving only ``(x, weight, cu_seqlens, bias)`` for the backward. - - ``forward(x, weight, cu_seqlens, bias=None, silu_grad_preactivation="bf16")`` - returns ``y``; the backward runs the fused ``dx`` + deterministic - ``dweight`` (+ ``dbias``) pair of launches of :func:`cake_causal_conv1d_train_backward` - on the incoming gradient (materialized as a contiguous BF16 tensor when - autograd hands over another layout or dtype) and returns the gradients of - ``weight`` and ``bias`` in the parameters' dtypes. - """ - - @staticmethod - def forward(ctx, x, weight, cu_seqlens, bias=None, silu_grad_preactivation="bf16"): - y = cake_causal_conv1d_train_forward(x, weight, cu_seqlens, bias=bias) - ctx.save_for_backward(x, weight, cu_seqlens, bias) - ctx.silu_grad_preactivation = _silu_grad_preactivation(silu_grad_preactivation) - return y - - @staticmethod - def backward(ctx, grad_y): - x, weight, cu_seqlens, bias = ctx.saved_tensors - if grad_y is None: - grad_y = torch.zeros_like(x) - dy = grad_y.to(torch.bfloat16).contiguous() - grads = cake_causal_conv1d_train_backward( - x, - weight, - cu_seqlens, - dy, - bias=bias, - silu_grad_preactivation=ctx.silu_grad_preactivation, - ) - dbias = grads[2] if bias is not None else None - return grads[0], grads[1], None, dbias, None - - -def cake_causal_conv1d( - x: torch.Tensor, - weight: torch.Tensor, - cu_seqlens: torch.Tensor, - bias: Optional[torch.Tensor] = None, - *, - silu_grad_preactivation: str = "bf16", -) -> torch.Tensor: - """Autograd-enabled packed width-4 depthwise causal convolution + SiLU. - - Parameters - ---------- - x : torch.Tensor - BF16 ``[S, C]`` or ``[1, S, C]`` contiguous packed activations on a - supported CUDA device. - weight : torch.Tensor - FP32 or BF16 ``[C, 4]`` contiguous taps. - cu_seqlens : torch.Tensor - ``int32 [num_docs + 1]`` device prefix sums of the document lengths. - bias : torch.Tensor, optional - FP32 or BF16 ``[C]`` contiguous per-channel bias. - silu_grad_preactivation : {"bf16", "fp32"} - The weight gradient's SiLU' policy (see - :func:`cake_causal_conv1d_train_backward`). - - Returns - ------- - torch.Tensor - BF16 ``y`` of ``x``'s shape. The backward saves only ``x``, ``weight``, - ``cu_seqlens`` and ``bias`` and returns BF16 ``dx`` and the ``dweight`` - (and ``dbias``) in the parameters' dtypes. - """ - - return CakeCausalConv1dFunction.apply( - x, weight, cu_seqlens, bias, silu_grad_preactivation - ) - - -__all__ = [ - "ARCHES", - "PARAMETER_DTYPES", - "SILU_GRAD_PREACTIVATIONS", - "SUPPORTED_COMPUTE_CAPABILITIES", - "CakeCausalConv1dFunction", - "cake_causal_conv1d", - "cake_causal_conv1d_train_backward", - "cake_causal_conv1d_train_backward_workspace", - "cake_causal_conv1d_train_backward_workspace_bytes", - "cake_causal_conv1d_train_forward", - "cake_causal_conv1d_train_forward_workspace", - "cake_causal_conv1d_train_forward_workspace_bytes", - "is_cake_causal_conv1d_train_supported", -] diff --git a/flashinfer/cake_causal_conv1d_train_backend/build.json b/flashinfer/cake_causal_conv1d_train_backend/build.json deleted file mode 100644 index d7bdde40dcd..00000000000 --- a/flashinfer/cake_causal_conv1d_train_backend/build.json +++ /dev/null @@ -1,1745 +0,0 @@ -{ - "arch_flags": { - "sm_100a": [ - "-gencode=arch=compute_100a,code=sm_100a" - ], - "sm_103a": [ - "-gencode=arch=compute_103a,code=sm_103a" - ], - "sm_107a": [ - "-gencode=arch=compute_107a,code=sm_107a" - ] - }, - "common": { - "include_dirs": [ - "csrc" - ], - "link": { - "cudart": "shared", - "host_linker": "c++", - "tvm_ffi_symbols": "resolved at load; do not pass --no-undefined" - }, - "link_flags": [ - "-shared", - "-lcudart", - "-lcuda" - ], - "nvcc_flags": [ - "-std=c++17", - "-O3", - "-DNDEBUG", - "--compiler-options=-fPIC", - "--expt-relaxed-constexpr", - "-Xfatbin=-compress-all", - "--compress-mode=size" - ], - "nvcc_flags_if": [ - { - "cuda": ">=12.8", - "flags": [ - "-static-global-template-stub=false" - ] - } - ], - "requires": { - "cuda_toolkit": { - "min": "12.8" - }, - "tvm_ffi": { - "include_dirs": [ - "tvm_ffi.libinfo.find_include_path()", - "tvm_ffi.libinfo.find_dlpack_include_path()" - ], - "python": "apache-tvm-ffi>=0.1.11,<0.2" - } - } - }, - "generator": { - "cake_revision": "e312167e6b4f6d8b4193e218da74fa5b17d8ec14", - "cuda": "13.4.92", - "rendered_at": "2026-10-28T23:53:55Z" - }, - "libraries": [ - { - "arch": "sm_100a", - "closure_sha256": "8181c92a596365c2897a8c664aebe7303b014b1acf186d891a0afc0326306b4c", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_08880ca12e4064bc8cdb_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_08880ca12e4064bc8cdb_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu": "67284293eceb5b55b77677446c96788b62455b98219f02fbabcd1c9d131514a0", - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu": "4d0ee67d1f8825ed025a2fb8575ce8a9a433fd40a30761c59906070b316adc89" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu", - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "108533bd02e09412046532109569c4abcd83825b8207ca632012d3bec41a9cc0", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_08880ca12e4064bc8cdb_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_08880ca12e4064bc8cdb_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu": "67284293eceb5b55b77677446c96788b62455b98219f02fbabcd1c9d131514a0", - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu": "4d0ee67d1f8825ed025a2fb8575ce8a9a433fd40a30761c59906070b316adc89" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu", - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "2d3fdaf6846b5e341b87d13f46ac90754846d6b3d61643e20cd297cc3f74925e", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_08880ca12e4064bc8cdb_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_08880ca12e4064bc8cdb_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu": "67284293eceb5b55b77677446c96788b62455b98219f02fbabcd1c9d131514a0", - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu": "4d0ee67d1f8825ed025a2fb8575ce8a9a433fd40a30761c59906070b316adc89" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_kernel.cu", - "csrc/cake_causal_conv1d_train_08880ca12e4064bc8cdb_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "54f6decd120fad6e7d1b24e51de24ab3f098b7c8e9735fda58a9b4000a38c870", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_149ca6f53934ba630033_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_149ca6f53934ba630033_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_149ca6f53934ba630033_binding.cu": "f29ca7671e13f70bdd14f60bcae481a2a464fa3501121ffb05c491e1ee24cc17", - "csrc/cake_causal_conv1d_train_149ca6f53934ba630033_kernel.cu": "cbbf522578e8b1beb38a3d2295d33e3496735955e528fd7b5057e3d8e76ace11" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_149ca6f53934ba630033_kernel.cu", - "csrc/cake_causal_conv1d_train_149ca6f53934ba630033_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "65ee6ba2b57ce40f1313981aa21a78dd4d47a8013bfbafcaace2e1f860b9ffe4", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_149ca6f53934ba630033_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_149ca6f53934ba630033_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_149ca6f53934ba630033_binding.cu": "f29ca7671e13f70bdd14f60bcae481a2a464fa3501121ffb05c491e1ee24cc17", - "csrc/cake_causal_conv1d_train_149ca6f53934ba630033_kernel.cu": "cbbf522578e8b1beb38a3d2295d33e3496735955e528fd7b5057e3d8e76ace11" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_149ca6f53934ba630033_kernel.cu", - "csrc/cake_causal_conv1d_train_149ca6f53934ba630033_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "5058cdf3822987e7302d4b9a74683b549cb1c2c8261f477a23590c66571d7ac1", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_170d89481a616ae977b3_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_170d89481a616ae977b3_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu": "435d120901907f2fceab47cd0e337a76cb9d03d82f297bff3fcbb401a8b18d2c", - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu": "b8a2c62c900a2377fb783b6e406cc118981d53c26f2c06a835dd4feaed400a06" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu", - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "c84bb1c47efa01513e48055d6b0657c1f6a57e6dbe4d5bcfad8d86b4ef73e6be", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_170d89481a616ae977b3_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_170d89481a616ae977b3_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu": "435d120901907f2fceab47cd0e337a76cb9d03d82f297bff3fcbb401a8b18d2c", - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu": "b8a2c62c900a2377fb783b6e406cc118981d53c26f2c06a835dd4feaed400a06" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu", - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "ec7437f0ed53c98d1d4bb3c32c92b70b2045ce6a61d384cac126909cce5003df", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_170d89481a616ae977b3_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_170d89481a616ae977b3_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu": "435d120901907f2fceab47cd0e337a76cb9d03d82f297bff3fcbb401a8b18d2c", - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu": "b8a2c62c900a2377fb783b6e406cc118981d53c26f2c06a835dd4feaed400a06" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_kernel.cu", - "csrc/cake_causal_conv1d_train_170d89481a616ae977b3_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "60deeacc98b1199b7615dc8f7e30bedf13bf1b6cf014471751027b42b3e076df", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_binding.cu": "f65f1b41c883ad74cafa66682f999f8d71a8d36aa4abb9f3c17020ae9992c40c", - "csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_kernel.cu": "7ac796d3b4c83de20872985503b03806575b4078ecc1c753b0221120197d0ac1" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_kernel.cu", - "csrc/cake_causal_conv1d_train_176ebb3e14f7ad9c94dc_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "053e4624381ef6055fac3af9599e1aa87504607dcc4216d91c3ee348474e824b", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_17cd42fb05079bbbf840_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_17cd42fb05079bbbf840_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu": "d3bf4559f22c2838efbe1a577109579f4cc37ea38df458393174aad7e4f6c6f5", - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu": "481412ca78511d5e1331ec1bd1c6eb2a76c3593e2fe35b8c1e97074a5376ad30" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu", - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "1f4646c653522f1a20a294aa89e1fa1e7972874d0dc48704d5d18710d4dcfaf6", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_17cd42fb05079bbbf840_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_17cd42fb05079bbbf840_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu": "d3bf4559f22c2838efbe1a577109579f4cc37ea38df458393174aad7e4f6c6f5", - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu": "481412ca78511d5e1331ec1bd1c6eb2a76c3593e2fe35b8c1e97074a5376ad30" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu", - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "6f2cf0fec3e5bbc4af75d083182549c0aa94a348aca19a6d1cf2dfb94107a35d", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_17cd42fb05079bbbf840_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_17cd42fb05079bbbf840_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu": "d3bf4559f22c2838efbe1a577109579f4cc37ea38df458393174aad7e4f6c6f5", - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu": "481412ca78511d5e1331ec1bd1c6eb2a76c3593e2fe35b8c1e97074a5376ad30" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_kernel.cu", - "csrc/cake_causal_conv1d_train_17cd42fb05079bbbf840_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "6efc7a5782dc80c79739199191dfce03f91f679481d8908a31d7eb84442dc612", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_1c21c181af94e1635851_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_1c21c181af94e1635851_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu": "2aeebdb7c7a04264b3c8714a1a1f6545c404f6d054b876c32a2e0a3de168fc7e", - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu": "bef0112c63a819bb68e23549a0df9880deac5f5af29542b65178911422695163" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu", - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "7461573cbefd0c9e04cef7e6c57e21de2230bb15053ee547286f29f5e48bd9e8", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_1c21c181af94e1635851_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_1c21c181af94e1635851_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu": "2aeebdb7c7a04264b3c8714a1a1f6545c404f6d054b876c32a2e0a3de168fc7e", - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu": "bef0112c63a819bb68e23549a0df9880deac5f5af29542b65178911422695163" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu", - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "7b7cfbe41609073a7c8d77a65910e146166b4731ea63221635a714b8a1d4057a", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_1c21c181af94e1635851_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_1c21c181af94e1635851_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu": "2aeebdb7c7a04264b3c8714a1a1f6545c404f6d054b876c32a2e0a3de168fc7e", - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu": "bef0112c63a819bb68e23549a0df9880deac5f5af29542b65178911422695163" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_kernel.cu", - "csrc/cake_causal_conv1d_train_1c21c181af94e1635851_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "3e9ef841ae714dc92d207cad59edbba5d3b11cd90d5def9378388f185c204699", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_2675cc1cb4efc0937378_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_2675cc1cb4efc0937378_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu": "c9d0334998b390bf9cc161bfb241cbc3b01208576870491e43891440e2f16803", - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu": "308257ea1704331616d2b7de6b1c466d255e323ffc3ccee6daa1d6447e293fd2" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu", - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "814862315a63bc178076b1134f31926eaf04db274145cb214a040dc6e6275393", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_2675cc1cb4efc0937378_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_2675cc1cb4efc0937378_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu": "c9d0334998b390bf9cc161bfb241cbc3b01208576870491e43891440e2f16803", - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu": "308257ea1704331616d2b7de6b1c466d255e323ffc3ccee6daa1d6447e293fd2" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu", - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "b35535a1eeb2e8cb7cb30c849171b31697ff1c9353733e5a26f10bf897b38fb3", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_2675cc1cb4efc0937378_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_2675cc1cb4efc0937378_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu": "c9d0334998b390bf9cc161bfb241cbc3b01208576870491e43891440e2f16803", - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu": "308257ea1704331616d2b7de6b1c466d255e323ffc3ccee6daa1d6447e293fd2" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_kernel.cu", - "csrc/cake_causal_conv1d_train_2675cc1cb4efc0937378_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "a25a6da9bd422604e197eafea9efcd916d0c559220f1907afb80baca7fa5f014", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_2d91e5f926238453c78c_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_2d91e5f926238453c78c_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_binding.cu": "f80a80924651042130b06ff2486ceb5ef11a61d47a729d82fde44703d84eda6b", - "csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_kernel.cu": "718d5b6e8f2805029dc3e3fe43ecc17b1f345b0f35a32f2488176da834dcaad1" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_kernel.cu", - "csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "f463663ff7ad96564082ce750c1957165b6183efba5c3b13abf9382a0f5f82f0", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_2d91e5f926238453c78c_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_2d91e5f926238453c78c_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_binding.cu": "f80a80924651042130b06ff2486ceb5ef11a61d47a729d82fde44703d84eda6b", - "csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_kernel.cu": "718d5b6e8f2805029dc3e3fe43ecc17b1f345b0f35a32f2488176da834dcaad1" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_kernel.cu", - "csrc/cake_causal_conv1d_train_2d91e5f926238453c78c_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "f99ed2feffd16ca4666272100c154799f98e582c0cdbc3b9c7293d4ec1351b37", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_37e428406a0cc4a8ca1b_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu": "1a45290d23054318c081c5b5e1e53cc79f669aca168235b119bcf2642f91aab5", - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu": "0307d4d3504f98a4b5c419b19d2ce0b4df165fa140dc396a6bce8ae95bb661d8" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu", - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "47ce5a06f2d73e06c01665db468a6931ce01f05fde800a9a7061892bb6b5a8b5", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_37e428406a0cc4a8ca1b_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu": "1a45290d23054318c081c5b5e1e53cc79f669aca168235b119bcf2642f91aab5", - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu": "0307d4d3504f98a4b5c419b19d2ce0b4df165fa140dc396a6bce8ae95bb661d8" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu", - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "eb419544183707e40b5852a8c1374c3ce64f962f1d3b62b38ecd963866a21d42", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_37e428406a0cc4a8ca1b_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu": "1a45290d23054318c081c5b5e1e53cc79f669aca168235b119bcf2642f91aab5", - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu": "0307d4d3504f98a4b5c419b19d2ce0b4df165fa140dc396a6bce8ae95bb661d8" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_kernel.cu", - "csrc/cake_causal_conv1d_train_37e428406a0cc4a8ca1b_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "49de5e7d081380a138c6ff8c82da33424212842496742509dc969fa43e0880f5", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu": "4cc9c5164d4c13b39f1087e51eb9d3a6867126bdf57ea67e35375a5728cf6488", - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu": "c5d3594389858ec2b556412aa03f3829f863f612a0804c8caa4dbc40730f3279" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu", - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "aa75a5bfc53fea38f2e73197ac8543858fd51e32770d32cf94b4cebf0be433a2", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu": "4cc9c5164d4c13b39f1087e51eb9d3a6867126bdf57ea67e35375a5728cf6488", - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu": "c5d3594389858ec2b556412aa03f3829f863f612a0804c8caa4dbc40730f3279" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu", - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "baf8c23e20017e9e145d6c6069f9ae841c4980850b10b45781e305af2cd7de64", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu": "4cc9c5164d4c13b39f1087e51eb9d3a6867126bdf57ea67e35375a5728cf6488", - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu": "c5d3594389858ec2b556412aa03f3829f863f612a0804c8caa4dbc40730f3279" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_kernel.cu", - "csrc/cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "5d3fd779c1d689a9f53d8a1ba55f31df92460c55424a7cd9db8eb85b3d36c5d0", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_517710d1f0988cf570a9_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_517710d1f0988cf570a9_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu": "bf6c46ebbe72ac0ede2d98fea21e7e73b329754e9be25c27ca76ff5583323472", - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu": "9fa84c4aa29dae3aa755cf1596d27a6694d846f9c0301833222c47707d2f14d8" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu", - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "52494b7ff96f30cff8671798d7773db2ac8b16836026f927b2890d11925d5295", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_517710d1f0988cf570a9_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_517710d1f0988cf570a9_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu": "bf6c46ebbe72ac0ede2d98fea21e7e73b329754e9be25c27ca76ff5583323472", - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu": "9fa84c4aa29dae3aa755cf1596d27a6694d846f9c0301833222c47707d2f14d8" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu", - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "49f681e78af3151e25759735de2b33752c59cf58e3ec4c04b0d1b62d60a9d825", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_517710d1f0988cf570a9_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_517710d1f0988cf570a9_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu": "bf6c46ebbe72ac0ede2d98fea21e7e73b329754e9be25c27ca76ff5583323472", - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu": "9fa84c4aa29dae3aa755cf1596d27a6694d846f9c0301833222c47707d2f14d8" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_kernel.cu", - "csrc/cake_causal_conv1d_train_517710d1f0988cf570a9_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "c187d4e484108e3451d5997c74881e67c0ddf42bed070cc466ad971ebdf4111c", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_5288d132a5cc927ae729_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_5288d132a5cc927ae729_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_binding.cu": "e9d1f51f3fa33ef96560a5590a4a38f6e0ef1a479cb57c935b2b53ce3853e62e", - "csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_kernel.cu": "c06d37a5fa4667d6e40d817114f5796fce09db3046085667b62e55247f8923d8" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_kernel.cu", - "csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "2b0e305e5e88d3fc64d508ccbb96ce3d4a742a4869a68f57efa6b09acdd4e548", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_5288d132a5cc927ae729_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_5288d132a5cc927ae729_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_binding.cu": "e9d1f51f3fa33ef96560a5590a4a38f6e0ef1a479cb57c935b2b53ce3853e62e", - "csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_kernel.cu": "c06d37a5fa4667d6e40d817114f5796fce09db3046085667b62e55247f8923d8" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_kernel.cu", - "csrc/cake_causal_conv1d_train_5288d132a5cc927ae729_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "f0427a3c00fbfc0edb6650a4a7fa3b2924b64b4a0889987e39cdec0cf454c795", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_675940a7a4fc0dba4b53_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_675940a7a4fc0dba4b53_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu": "97ccb20269af554aa4e1317d2e0fe4c152020843672cabf717a02eef279bb99a", - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu": "5e66fbcee4ba701174c6a46a6fa34cb6871d583df301567b00469e65ce6f84ed" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu", - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "2844a4a678354529feae15c785c62c936f7cf69ccbaf0e96b1df961523792782", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_675940a7a4fc0dba4b53_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_675940a7a4fc0dba4b53_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu": "97ccb20269af554aa4e1317d2e0fe4c152020843672cabf717a02eef279bb99a", - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu": "5e66fbcee4ba701174c6a46a6fa34cb6871d583df301567b00469e65ce6f84ed" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu", - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "ae72b27ad4f9ebb0bc37e0ee5d8c54b10563d287927113bbad65c7f7f1978ef2", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_675940a7a4fc0dba4b53_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_675940a7a4fc0dba4b53_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu": "97ccb20269af554aa4e1317d2e0fe4c152020843672cabf717a02eef279bb99a", - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu": "5e66fbcee4ba701174c6a46a6fa34cb6871d583df301567b00469e65ce6f84ed" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_kernel.cu", - "csrc/cake_causal_conv1d_train_675940a7a4fc0dba4b53_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "1d8db0e374474def19a949290afc97ed16901601248972dfaf1d59b7b61d17e3", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_binding.cu": "c1332ab0469a57049c78246a1beef72ea071926cdb9d43edcd43872c88299c04", - "csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_kernel.cu": "bbf0d05a8c9756d50f5b6776c375f5844057d5c39703835bddac5ebe09cc51b0" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_kernel.cu", - "csrc/cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "678b6680e1b2e2a9abee3effdcc465507a507d8f1f6a8e935c057ef980c47c3b", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_7485529a1755e5bd3036_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_7485529a1755e5bd3036_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu": "1c13d9c6dc0599cb8772834dd5d272c3ac0194540fd9ad020f33166ead91e80c", - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu": "de02c69950ec7e2ab628bb61fb21991cc8d6e485ea6064eb535c4fdcb807628b" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu", - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "0cdc90f376954df81b7f340f2e799721363be99cabfec85887b75d1b32342fef", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_7485529a1755e5bd3036_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_7485529a1755e5bd3036_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu": "1c13d9c6dc0599cb8772834dd5d272c3ac0194540fd9ad020f33166ead91e80c", - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu": "de02c69950ec7e2ab628bb61fb21991cc8d6e485ea6064eb535c4fdcb807628b" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu", - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "6c44bdb9195fa1cddd588a25915f5229008f90044a4288521471c93fb1dd2ae3", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_7485529a1755e5bd3036_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_7485529a1755e5bd3036_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu": "1c13d9c6dc0599cb8772834dd5d272c3ac0194540fd9ad020f33166ead91e80c", - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu": "de02c69950ec7e2ab628bb61fb21991cc8d6e485ea6064eb535c4fdcb807628b" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_kernel.cu", - "csrc/cake_causal_conv1d_train_7485529a1755e5bd3036_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "9e38e97cdb075fd063f2931a6966049b4529be94037bf5ebe120fdda8063e5a9", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_843c010db78c79107fd3_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_843c010db78c79107fd3_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_843c010db78c79107fd3_binding.cu": "e5fdfa7df6c4e7e7961d05a90985292010f34a1aaa03df36d11aa91dbf838cd4", - "csrc/cake_causal_conv1d_train_843c010db78c79107fd3_kernel.cu": "6261351de0aa7bd088a3e89ad51c129e4de5b184f84368e5a6921c522ab3ee47" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_843c010db78c79107fd3_kernel.cu", - "csrc/cake_causal_conv1d_train_843c010db78c79107fd3_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "3ff781b769466a09f8814a2ac563506d70305ccc0bccb79066f05087d79edd10", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_8665d0756b8240ecb4d0_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_8665d0756b8240ecb4d0_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_binding.cu": "96f38e3f0725cd5a927ae69c89e7295527c8b5e6ee0506213ec6ea2745768b21", - "csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_kernel.cu": "572e7a38873d13d00569a1ca2848fabe6c42ad51e1d39b0af8a5757dc4c17d6f" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_kernel.cu", - "csrc/cake_causal_conv1d_train_8665d0756b8240ecb4d0_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "e2850030a9d1dc35bf97f2969827dceeb240eedbd93d08872979b6fa3e63a353", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu": "d834bd5676f5dc1736e57fa6c6c8b81615bee8358e2af4e3db36c66ff452447f", - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu": "c884bc1733cba848127e880739b961a0d67e151455741b4c33b6ddb3d15af4f9" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu", - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "85800e792fd9cc59e1149e3c8de0f7f2b9bf12089d3644acdfbecd0e6d4e96cb", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu": "d834bd5676f5dc1736e57fa6c6c8b81615bee8358e2af4e3db36c66ff452447f", - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu": "c884bc1733cba848127e880739b961a0d67e151455741b4c33b6ddb3d15af4f9" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu", - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "102187354f9c906ed07363ee740318276aaba57bea476fe14aea7c3ab0b79aee", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu": "d834bd5676f5dc1736e57fa6c6c8b81615bee8358e2af4e3db36c66ff452447f", - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu": "c884bc1733cba848127e880739b961a0d67e151455741b4c33b6ddb3d15af4f9" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_kernel.cu", - "csrc/cake_causal_conv1d_train_86b5ecfbd90aca4b52e5_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "508bddd849e34783a344eebc78785e2cc3702a06dd0f9ffeff31d727d054adeb", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_8b9b8db0833fd64a578f_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_8b9b8db0833fd64a578f_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_binding.cu": "d870571d36def0243936e30b1f7e7d4597fa436f5ee8b54fc4ea10e0f8082ce1", - "csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_kernel.cu": "2191a6094149ac627da55bf4801bbf718a9aa2d29d78907eb6248f99e971cbf6" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_kernel.cu", - "csrc/cake_causal_conv1d_train_8b9b8db0833fd64a578f_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "0a8843c95b0955863752bf9e93863cf52d78ffb89a27775f44268949754fea1b", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_9a07bcaa0223632640ee_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_9a07bcaa0223632640ee_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_binding.cu": "a4edb9897649bbb8d3ed21a4e268d01038a796e2bf182458268c8032053cb06b", - "csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_kernel.cu": "0a70cc039335820214213a99a68ae7df476f4c78eb5dc9e373b267e5e9c9519d" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_kernel.cu", - "csrc/cake_causal_conv1d_train_9a07bcaa0223632640ee_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "2445464fe4821b88950b70f73796dc2e51d02de4bf7584fb5bfbd5cf9a0d6225", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_9bfa65707a49df6a48ce_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_9bfa65707a49df6a48ce_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_binding.cu": "07a044f779dce018c6d0a17de64f7929730680d3fac58aff0745de049210e8ab", - "csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_kernel.cu": "475b6d8e558899b74fd0ce6773e222026305c8b09da196326c8b11930a5bc94a" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_kernel.cu", - "csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "219c0b8c819d5c0cedddc75bd3ae3384d2bf382e7bbe2d8b89c74303acb87f61", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_9bfa65707a49df6a48ce_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_9bfa65707a49df6a48ce_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_binding.cu": "07a044f779dce018c6d0a17de64f7929730680d3fac58aff0745de049210e8ab", - "csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_kernel.cu": "475b6d8e558899b74fd0ce6773e222026305c8b09da196326c8b11930a5bc94a" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_kernel.cu", - "csrc/cake_causal_conv1d_train_9bfa65707a49df6a48ce_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "a75ce44514ab91c1325a1f65d9e25f2e28302793b1b386aadf749127107620dc", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_9c1e7ee9f37e3102660e_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu": "16022ddfe0a39a5cee234327aa02560846527cdae6470da1c4f3925f7a4d6c6a", - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu": "6dd6ca72cc297264918ebcded7b4da238ad50da2bd2220a373c4cd319de6fc2e" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu", - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "d0e0af6b5cbeff21a16e04743997d8618d09040c146f3e4f647f6953907b1704", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_9c1e7ee9f37e3102660e_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu": "16022ddfe0a39a5cee234327aa02560846527cdae6470da1c4f3925f7a4d6c6a", - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu": "6dd6ca72cc297264918ebcded7b4da238ad50da2bd2220a373c4cd319de6fc2e" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu", - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "38365bfcb2ca021590048dcd0223db2307ad4ba4954cb8a0b2bf888b026993c0", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_9c1e7ee9f37e3102660e_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu": "16022ddfe0a39a5cee234327aa02560846527cdae6470da1c4f3925f7a4d6c6a", - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu": "6dd6ca72cc297264918ebcded7b4da238ad50da2bd2220a373c4cd319de6fc2e" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_kernel.cu", - "csrc/cake_causal_conv1d_train_9c1e7ee9f37e3102660e_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "547e8df2dc1ad4eec5d05b3789e0cdddd576ec7ab905fe4018b115d76cda75a8", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_binding.cu": "5dda0dfaeaee03e090de91649f3a82f1745447518da0a1847cbbccbd12447d1d", - "csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_kernel.cu": "40715f1ab1cd7a19b0594836aaa3145995c6a73fcf88ad86207e328a3acf6201" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_kernel.cu", - "csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "73230c64792fd3feab479a477b8fb536d979551e083b5a2451285e0b1fd7a0bc", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_binding.cu": "5dda0dfaeaee03e090de91649f3a82f1745447518da0a1847cbbccbd12447d1d", - "csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_kernel.cu": "40715f1ab1cd7a19b0594836aaa3145995c6a73fcf88ad86207e328a3acf6201" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_kernel.cu", - "csrc/cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "c429d92d958067d22070285e0bb04e26729ec066dd5393aea19600658fac4fb5", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_a681165b3e81204c73ff_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_a681165b3e81204c73ff_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu": "dcc25fc7dc7977412aebc9eceeb05791ac62c09f2f9aaf91bbb826a52e04d129", - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu": "08b2bf8d332c3627108f3e23fc884f25b7c62a60a87ea769135aae58af4baa17" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu", - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "132cf41ddac61304b98d4297810189a1ea4ea5da2e2fbf8c37cd3a7d575ddaaf", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_a681165b3e81204c73ff_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_a681165b3e81204c73ff_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu": "dcc25fc7dc7977412aebc9eceeb05791ac62c09f2f9aaf91bbb826a52e04d129", - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu": "08b2bf8d332c3627108f3e23fc884f25b7c62a60a87ea769135aae58af4baa17" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu", - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "bdd63be1ae5d8ebaee46935c7b136a7746f64d9d4c90a5ca0ed38d9cfaa66d54", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_a681165b3e81204c73ff_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_a681165b3e81204c73ff_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu": "dcc25fc7dc7977412aebc9eceeb05791ac62c09f2f9aaf91bbb826a52e04d129", - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu": "08b2bf8d332c3627108f3e23fc884f25b7c62a60a87ea769135aae58af4baa17" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_kernel.cu", - "csrc/cake_causal_conv1d_train_a681165b3e81204c73ff_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "54d39f1ad2e921cb232dbceab99d6493f0e93f7c264b1172050d63ac62323eaa", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_aeff6106e2822f313a0c_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_aeff6106e2822f313a0c_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_binding.cu": "bc2e652ba15bcca5ffd062d3c1d07c41bc99f9a554b7de2ebd5b42d8f6bcd1d3", - "csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_kernel.cu": "933fe17525b8b65b233d536a44c7c7920dfa59ed1b3c86ebfbdb14a179787f91" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_kernel.cu", - "csrc/cake_causal_conv1d_train_aeff6106e2822f313a0c_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "add12ff82ed9b399181f7f0e539a8df8ba5008549c686cba2ba36893d822e97f", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_binding.cu": "caa435bb8ad6d24ede06affdce2860b40ed4464eff3af45301dd65eb87f5cc38", - "csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_kernel.cu": "c0096a1402f2846c012ccd5e712bd2cc00845937123cdf207d46d882c0190365" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_kernel.cu", - "csrc/cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "024bd8c131848f484e4d460b049fef85b152cef80c68bad595879c1db2a2e593", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_c8027a64ba3cc38333bf_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_c8027a64ba3cc38333bf_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_binding.cu": "a7e841896f193a6f800ecb2bf966a7b02d6aff36ee92715f506ffdfc4b8d6129", - "csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_kernel.cu": "2b545605102d0e80cd1460a25163ee0897f2794410dfced4b97079c7a15d0d1b" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_kernel.cu", - "csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "9c546c26fef41c9257685a1544c3baa807580af64e348edd1ab2072b1972d287", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_c8027a64ba3cc38333bf_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_c8027a64ba3cc38333bf_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_binding.cu": "a7e841896f193a6f800ecb2bf966a7b02d6aff36ee92715f506ffdfc4b8d6129", - "csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_kernel.cu": "2b545605102d0e80cd1460a25163ee0897f2794410dfced4b97079c7a15d0d1b" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_kernel.cu", - "csrc/cake_causal_conv1d_train_c8027a64ba3cc38333bf_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "2b5e5a9be043f2312bae56192262cc7fbd1b65bbc2393deedf364ac60239c413", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_cc7a117aee9c7ae829c4_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_binding.cu": "bf01229de59dadd4c312f41881ba6d6daa174cb04ac7d80efdd48144a7204063", - "csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_kernel.cu": "50973295bd5a4b690a4b8543b426efa244ba851f4de9eab5871434e8843799d2" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_kernel.cu", - "csrc/cake_causal_conv1d_train_cc7a117aee9c7ae829c4_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "8966cdd034af1149ac7f77918ead7979562c686bcbb427dd61a888fa0a7a76a8", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_ce49146b9fef24adb47e_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_ce49146b9fef24adb47e_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu": "1f62015fd9bdbe4239f818d41d7d63f817433657d1e6d608d537c0490f7468c1", - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu": "b6161558774fe9e64319ce9ef5e8340018eea3f9be40084e40338f2d38017043" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu", - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "c3e23d9f57cf0a56e415fdd2d41ca78663f398c3a0b854b60a084ce63c3d8f61", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_ce49146b9fef24adb47e_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_ce49146b9fef24adb47e_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu": "1f62015fd9bdbe4239f818d41d7d63f817433657d1e6d608d537c0490f7468c1", - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu": "b6161558774fe9e64319ce9ef5e8340018eea3f9be40084e40338f2d38017043" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu", - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "529ccdf7ed1d36b0f70499608251d6d45c7214490f97e7ac81811ae0e7b07c0c", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_ce49146b9fef24adb47e_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_ce49146b9fef24adb47e_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu": "1f62015fd9bdbe4239f818d41d7d63f817433657d1e6d608d537c0490f7468c1", - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu": "b6161558774fe9e64319ce9ef5e8340018eea3f9be40084e40338f2d38017043" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_kernel.cu", - "csrc/cake_causal_conv1d_train_ce49146b9fef24adb47e_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "78074c7e46d23e79de2f1b20e433138901de514ab1dbd87b15ef2c0055ffd10d", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_binding.cu": "693298c9ac202cda916fbab120b60f8072641306fa7371e70a87c5dc66518136", - "csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_kernel.cu": "5017c08ee8b031ef93b6812cf9e3fd531409fab6d038f50b472c30c04f6ed399" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_kernel.cu", - "csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "34d1876115d48309962be09f4255493b6cd8be5bb963255e9e5235666c9d2f16", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_binding.cu": "693298c9ac202cda916fbab120b60f8072641306fa7371e70a87c5dc66518136", - "csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_kernel.cu": "5017c08ee8b031ef93b6812cf9e3fd531409fab6d038f50b472c30c04f6ed399" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_kernel.cu", - "csrc/cake_causal_conv1d_train_d46036aa2bc5f90bb0f6_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "b55011f9e8aa50a7accdd0ac22027afddb75b59abe5b599f277475274f9a6a9f", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_d829c95e1d610bc14b23_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_d829c95e1d610bc14b23_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_binding.cu": "649f613d7d60edee643447d1b1f58908d223d6f401becbd50b7bbf61970595b2", - "csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_kernel.cu": "d24b2209852b8b77f0fd17955900bfea2af252652eb859aabb819f768a7a49e6" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_kernel.cu", - "csrc/cake_causal_conv1d_train_d829c95e1d610bc14b23_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "4b04122f4e0d4a96498a962dfb1d08b27808f17601579132c324857e743ab9c7", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_ea8f91038a9dea452ff2_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_ea8f91038a9dea452ff2_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu": "24819a226390735c7572d8117c7de4a15df61b7018504fa0fb2d27ec9cac7b3c", - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu": "a716f616ddafbfffb9444d5c73df2d1fcdc35ece82a5df4fed53aabaf460a5e9" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu", - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "fd5c15993e118e9314f949e3f68dd2df65fb0cec748e50f5d8001aeeb80e2769", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_ea8f91038a9dea452ff2_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_ea8f91038a9dea452ff2_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu": "24819a226390735c7572d8117c7de4a15df61b7018504fa0fb2d27ec9cac7b3c", - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu": "a716f616ddafbfffb9444d5c73df2d1fcdc35ece82a5df4fed53aabaf460a5e9" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu", - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "ee62f17f672e4376bf9b6b60c909965d98b4c241617ddc4fcf6ed0287e9350de", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_ea8f91038a9dea452ff2_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_ea8f91038a9dea452ff2_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu": "24819a226390735c7572d8117c7de4a15df61b7018504fa0fb2d27ec9cac7b3c", - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu": "a716f616ddafbfffb9444d5c73df2d1fcdc35ece82a5df4fed53aabaf460a5e9" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_kernel.cu", - "csrc/cake_causal_conv1d_train_ea8f91038a9dea452ff2_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "2d2eebb1150aad41863783f5f6a1e17290d05d6b2a98fa40d0f5f64d50a3a099", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_f1545adfd2bb28044419_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_f1545adfd2bb28044419_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_binding.cu": "aaadd4e9534872f4aac9ca4c09c9371c10ae961e878675fbf7b97898cd5a85a4", - "csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_kernel.cu": "1bb2163d414690b69c7228326a9e669f6ccb15ca75c303a660ec2234da9a406d" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_kernel.cu", - "csrc/cake_causal_conv1d_train_f1545adfd2bb28044419_binding.cu" - ] - }, - { - "arch": "sm_100a", - "closure_sha256": "a39a126fba96dfa257cf39e76042d0d6e31e1e97def40aae3dcbbf625fc740c5", - "entry": "run", - "gencode": "-gencode=arch=compute_100a,code=sm_100a", - "min_cuda": "12.8", - "name": "cake_causal_conv1d_train_f5497330bfd286211643_sm_100a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_f5497330bfd286211643_sm_100a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu": "d418cab7219dad4fa830268414ccef295430adc465bc6222f953dc0d62fb1cd8", - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu": "3dc9bd451ab2f7cb6fdf18cb6689baffd41a05be03cc6c21bd91609ea36ae567" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu", - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu" - ] - }, - { - "arch": "sm_103a", - "closure_sha256": "601cfe47c9d2873220961be127f9c85f2cacb76725b76879722de47de81011e9", - "entry": "run", - "gencode": "-gencode=arch=compute_103a,code=sm_103a", - "min_cuda": "12.9", - "name": "cake_causal_conv1d_train_f5497330bfd286211643_sm_103a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_f5497330bfd286211643_sm_103a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu": "d418cab7219dad4fa830268414ccef295430adc465bc6222f953dc0d62fb1cd8", - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu": "3dc9bd451ab2f7cb6fdf18cb6689baffd41a05be03cc6c21bd91609ea36ae567" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu", - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu" - ] - }, - { - "arch": "sm_107a", - "closure_sha256": "faa149f2ad044411e8b43c5abc8bb20f3f2a68f13163a014f398e3b31ed44c09", - "entry": "run", - "gencode": "-gencode=arch=compute_107a,code=sm_107a", - "min_cuda": "13.4", - "name": "cake_causal_conv1d_train_f5497330bfd286211643_sm_107a", - "nvcc_flags": [ - "--prec-div=true", - "--prec-sqrt=true", - "--ftz=false" - ], - "nvcc_flags_if": [], - "path": "aot/cake_causal_conv1d_train_f5497330bfd286211643_sm_107a.so", - "sha256": { - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu": "d418cab7219dad4fa830268414ccef295430adc465bc6222f953dc0d62fb1cd8", - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu": "3dc9bd451ab2f7cb6fdf18cb6689baffd41a05be03cc6c21bd91609ea36ae567" - }, - "sources": [ - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_kernel.cu", - "csrc/cake_causal_conv1d_train_f5497330bfd286211643_binding.cu" - ] - } - ], - "package": "flashinfer.cake_causal_conv1d_train_backend", - "schema": "cake-build/1" -} diff --git a/flashinfer/cake_causal_conv1d_train_backend/cake_build.py b/flashinfer/cake_causal_conv1d_train_backend/cake_build.py index 00756137eb1..a8108e9202e 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/cake_build.py +++ b/flashinfer/cake_causal_conv1d_train_backend/cake_build.py @@ -1,324 +1,9 @@ -""" -Copyright (c) 2026 by FlashInfer team. +# Copyright (c) 2026 by FlashInfer team. Licensed under Apache-2.0. +"""Compatibility import for flashinfer.cake_causal_conv1d_det_train_backend.cake_build (the same module object).""" -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at +from importlib import import_module +import sys - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. - -Library loader of one Cake-generated FlashInfer package. - -``build.json`` (schema ``cake-build/1``) next to this file is the only record -of what the package compiles: one entry per ``[_]_`` -library with its package-relative CUDA sources, nvcc flags, toolkit-conditional -flags and digests. A library is served from ``aot/.so`` when a builder -(a vendoring consumer's build rules, or FlashInfer's -``scripts/check_cake_build_manifests.py``) has compiled it from the manifest, -and otherwise through FlashInfer's JIT from the same sources and flags. - -The prebuilt path never imports ``flashinfer``, so the package stays usable in -a process without a FlashInfer installation. A builder that fills ``aot/`` -also writes the sidecar ``aot/build_info.json`` (schema ``cake-build-info/1``: -``{"schema", "cuda": "13.0.88", "cudart_version": 13000, "arches": [...]}``, -``arches`` naming every architecture the builder compiled). Every prebuilt -library exports ``build_cuda_version()``, the ``CUDART_VERSION`` it was -compiled against, and :func:`load_library` refuses a library whose value -disagrees with the sidecar. - -Public surface (``__all__``): :func:`manifest`, :func:`libraries_by_name`, -:func:`prebuilt_path`, :func:`library_path`, :func:`load_library`, -:func:`jit_spec`, :func:`jit_specs` and :func:`record_cuda_version`. - -The Cake exporter writes this file byte-identically into every package from -its loader template; do not edit it in place. -""" - -from __future__ import annotations - -import functools -import hashlib -import json -import re -from pathlib import Path -from typing import Any - -__all__ = [ - "jit_spec", - "jit_specs", - "libraries_by_name", - "library_path", - "load_library", - "manifest", - "prebuilt_path", - "record_cuda_version", -] - -_PACKAGE = Path(__file__).parent -_MANIFEST_SCHEMA = "cake-build/1" -_BUILD_INFO_SCHEMA = "cake-build-info/1" -_MANIFEST_PATH = _PACKAGE / "build.json" -_BUILD_INFO_PATH = _PACKAGE / "aot" / "build_info.json" -_CUDA_CONSTRAINT = re.compile(r"(==|>=)([0-9]+)\.([0-9]+)") - - -@functools.cache -def manifest() -> dict[str, Any]: - """The parsed ``build.json`` of this package (read on first use, not at import). - - One cached document shared by every caller: read it, do not modify it. - """ - try: - text = _MANIFEST_PATH.read_text(encoding="utf-8") - except FileNotFoundError as exc: - raise RuntimeError( - f"{_MANIFEST_PATH} is missing: the package was not exported with its {_MANIFEST_SCHEMA} manifest", - ) from exc - document = json.loads(text) - schema = document.get("schema") - if schema != _MANIFEST_SCHEMA: - raise RuntimeError( - f"{_MANIFEST_PATH} has schema {schema!r}; this loader reads {_MANIFEST_SCHEMA!r}", - ) - return document - - -@functools.cache -def libraries_by_name() -> dict[str, dict[str, Any]]: - """``name -> entry`` over the ``libraries`` of ``build.json``.""" - index: dict[str, dict[str, Any]] = {} - for entry in manifest()["libraries"]: - if entry["name"] in index: - raise RuntimeError(f"{_MANIFEST_PATH} lists {entry['name']!r} twice") - index[entry["name"]] = entry - return index - - -def _entry(name: str) -> dict[str, Any]: - try: - return libraries_by_name()[name] - except KeyError: - raise KeyError(f"{name!r} is not a library of {_MANIFEST_PATH}") from None - - -def prebuilt_path(name: str) -> Path: - """Where a builder places the compiled library of ``name`` (``aot/.so``).""" - return _PACKAGE / _entry(name)["path"] - - -def _cuda_constraint_matches(constraint: str, cuda: tuple[int, int]) -> bool: - """Does toolkit ``cuda`` (major, minor) satisfy ``"==X.Y"`` or ``">=X.Y"``?""" - match = _CUDA_CONSTRAINT.fullmatch(constraint) - if match is None: - raise ValueError( - f"unsupported CUDA constraint {constraint!r}; expected '==X.Y' or '>=X.Y'", - ) - operator, major, minor = match.groups() - version = (int(major), int(minor)) - if operator == "==": - return tuple(cuda) == version - return tuple(cuda) >= version - - -def _conditional_flags(items: list[Any], cuda: tuple[int, int]) -> list[str]: - """The flags of every ``nvcc_flags_if`` item whose ``cuda`` constraint matches.""" - flags: list[str] = [] - for item in items: - if _cuda_constraint_matches(item["cuda"], cuda): - flags.extend(item["flags"]) - return flags - - -def _jit_toolchain(name: str) -> tuple[Any, Any]: - """``(gen_jit_spec, get_cuda_version)`` of the local FlashInfer, or the error.""" - try: - from flashinfer.jit.core import gen_jit_spec - from flashinfer.jit.cpp_ext import get_cuda_version - except ImportError as exc: - raise RuntimeError( - f"{name} is not prebuilt ({prebuilt_path(name)}) and flashinfer's JIT is not available", - ) from exc - return gen_jit_spec, get_cuda_version - - -def _local_cuda_version(get_cuda_version: Any) -> tuple[int, int]: - version = get_cuda_version() - return (int(version.major), int(version.minor)) - - -def _version(text: str) -> tuple[int, int]: - """``(major, minor)`` of a dotted version such as ``"13.0"`` or ``"13.0.88"``.""" - parts = str(text).split(".") - if len(parts) < 2 or not all(part.isdigit() for part in parts[:2]): - raise RuntimeError(f"malformed CUDA version {text!r}") - return (int(parts[0]), int(parts[1])) - - -def _toolkit_builds(entry: dict[str, Any], cuda: tuple[int, int]) -> bool: - """Does toolkit ``cuda`` reach the entry's ``min_cuda`` floor?""" - return cuda >= _version(entry["min_cuda"]) - - -def _verified_sources(entry: dict[str, Any]) -> list[Path]: - """The entry's sources as package paths, each checked against its digest.""" - sources: list[Path] = [] - for relative in entry["sources"]: - path = _PACKAGE / relative - digest = hashlib.sha256(path.read_bytes()).hexdigest() - expected = entry["sha256"][relative] - if digest != expected: - raise RuntimeError( - f"{path} does not match its build.json digest for {entry['name']}: {digest} != {expected}", - ) - sources.append(path) - return sources - - -def jit_spec(name: str) -> Any: - """FlashInfer's ``JitSpecNvcc`` of ``name``, built from the manifest entry. - - Sources are digest-verified here and only here, so a modified or partially - delivered source tree fails closed before anything compiles. The spec name - carries the closure digest: a changed closure never reuses a stale JIT cache - entry. Flags are ``arch_flags[arch]`` + the entry's ``nvcc_flags`` + the - ``nvcc_flags_if`` items matching the local toolkit; FlashInfer's JIT adds - its own defaults (``-O3``, ``-fPIC``, ...), so ``common.nvcc_flags`` are not - repeated here. A toolkit below the entry's ``min_cuda`` raises. - """ - entry = _entry(name) - gen_jit_spec, get_cuda_version = _jit_toolchain(name) - document = manifest() - cuda = _local_cuda_version(get_cuda_version) - if not _toolkit_builds(entry, cuda): - raise RuntimeError( - f"{name} needs CUDA {entry['min_cuda']} or newer; the local toolkit is {cuda[0]}.{cuda[1]}", - ) - flags = [ - *document["arch_flags"][entry["arch"]], - *entry["nvcc_flags"], - *_conditional_flags(entry["nvcc_flags_if"], cuda), - ] - include_dirs = [_PACKAGE / d for d in document["common"]["include_dirs"]] - return gen_jit_spec( - f"{name}_{entry['closure_sha256'][:20]}", - _verified_sources(entry), - extra_cuda_cflags=flags, - extra_ldflags=list(document["common"]["link_flags"]), - extra_include_paths=include_dirs, - use_fast_math=False, - ) - - -def _names_of(arch: str) -> list[str]: - index = libraries_by_name() - return [name for name in index if index[name]["arch"] == arch] - - -def jit_specs(arch: str) -> list[Any]: - """JIT specifications of the ``arch`` libraries still to build locally. - - The entries without a prebuilt file that the local toolkit can build: the - list feeds ``flashinfer.jit.build_jit_specs`` warm-ups, so an entry whose - ``min_cuda`` exceeds the local toolkit is left out (the manifest check - script's "skip min_cuda > toolkit" rule), while :func:`jit_spec` on that - entry by name still raises. When every ``arch`` library is prebuilt the - list is empty and ``flashinfer`` is not imported. - """ - pending = [name for name in _names_of(arch) if not prebuilt_path(name).is_file()] - if not pending: - return [] - _, get_cuda_version = _jit_toolchain(pending[0]) - cuda = _local_cuda_version(get_cuda_version) - return [jit_spec(name) for name in pending if _toolkit_builds(_entry(name), cuda)] - - -def library_path(name: str) -> Path: - """The prebuilt library of ``name`` when present, else its JIT cache path.""" - path = prebuilt_path(name) - if path.is_file(): - return path - return jit_spec(name).get_library_path() - - -@functools.cache -def _build_info() -> dict[str, Any]: - """The builder-written ``aot/build_info.json`` sidecar (prebuilt path only).""" - try: - text = _BUILD_INFO_PATH.read_text(encoding="utf-8") - except FileNotFoundError as exc: - raise RuntimeError( - f"{_BUILD_INFO_PATH} is missing although prebuilt libraries exist under {_PACKAGE / 'aot'}: the builder must write the {_BUILD_INFO_SCHEMA} sidecar", - ) from exc - info = json.loads(text) - schema = info.get("schema") - if schema != _BUILD_INFO_SCHEMA: - raise RuntimeError( - f"{_BUILD_INFO_PATH} has schema {schema!r}; this loader reads {_BUILD_INFO_SCHEMA!r}", - ) - return info - - -def _prebuilt_info(path: Path, arch: str) -> dict[str, Any]: - """The sidecar, which must record a build of ``arch`` for the prebuilt ``path``.""" - info = _build_info() - if arch not in info["arches"]: - raise RuntimeError( - f"{path} exists but {_BUILD_INFO_PATH} records no {arch} build (arches {info['arches']})", - ) - return info - - -def load_library(name: str) -> Any: - """The loaded tvm-ffi module of ``name``: the prebuilt library, else a JIT build.""" - entry = _entry(name) - path = _PACKAGE / entry["path"] - if not path.is_file(): - return jit_spec(name).build_and_load() - info = _prebuilt_info(path, entry["arch"]) - import tvm_ffi - - module = tvm_ffi.load_module(str(path)) - built = int(module.build_cuda_version()) - if built != int(info["cudart_version"]): - raise RuntimeError( - f"{path} was compiled against CUDART {built} but {_BUILD_INFO_PATH} records {info['cudart_version']}", - ) - return module - - -def record_cuda_version(arch: str) -> tuple[int, int]: - """``(major, minor)`` of the toolkit that builds the ``arch`` libraries. - - Prebuilt path, taken as soon as any library of the manifest exists under - ``aot/``: the sidecar ``aot/build_info.json`` is mandatory (missing -> - ``RuntimeError``; no local nvcc is consulted) and its ``cuda`` is the - answer whenever its ``arches`` list ``arch`` -- the builder compiled that - architecture, whether or not every one of its libraries is present. An - ``arch`` library under ``aot/`` that the sidecar does not list raises, as - :func:`load_library` does. - - JIT path, when nothing is prebuilt or the builder did not compile - ``arch``: the local toolkit of FlashInfer's JIT (``RuntimeError`` in a - process without a FlashInfer installation). Call it for the architecture - you are about to load. - """ - names = _names_of(arch) - if not names: - raise KeyError(f"{arch!r} has no libraries in {_MANIFEST_PATH}") - index = libraries_by_name() - prebuilt = [name for name in index if (_PACKAGE / index[name]["path"]).is_file()] - if prebuilt: - of_arch = [name for name in prebuilt if index[name]["arch"] == arch] - if of_arch: - return _version(_prebuilt_info(prebuilt_path(of_arch[0]), arch)["cuda"]) - info = _build_info() - if arch in info["arches"]: - return _version(info["cuda"]) - _, get_cuda_version = _jit_toolchain(names[0]) - return _local_cuda_version(get_cuda_version) +sys.modules[__name__] = import_module( + "flashinfer.cake_causal_conv1d_det_train_backend.cake_build" +) diff --git a/flashinfer/cake_causal_conv1d_train_backend/cake_jit.py b/flashinfer/cake_causal_conv1d_train_backend/cake_jit.py index 6c1e76a0b83..7f8569cf8d0 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/cake_jit.py +++ b/flashinfer/cake_causal_conv1d_train_backend/cake_jit.py @@ -1,3670 +1,9 @@ -""" -Copyright (c) 2026 by FlashInfer team. +# Copyright (c) 2026 by FlashInfer team. Licensed under Apache-2.0. +"""Compatibility import for flashinfer.cake_causal_conv1d_det_train_backend.cake_jit (the same module object).""" -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at +from importlib import import_module +import sys - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. - -Host plan and loader glue of the Cake packed width-4 depthwise causal -convolution + SiLU training kernels (SM100 / SM103 / SM107). - -Every module below is one mechanically exported production build: a CUDA -device translation unit plus a tvm-ffi binding compiled together for exactly -one architecture. ``MODULES`` and ``ROUTES`` carry the host plan only: the -kernel's ``library`` stem, the architectures it is delivered for, its launch -record, grid rule, workspace rule and the kernel-parameter -> logical-argument -map. What to compile -- sources, digests, nvcc flags -- is recorded once in -``build.json`` next to this file; :mod:`.cake_build` serves the library -``_`` from a prebuilt ``aot/`` file when a builder compiled it -and otherwise through FlashInfer's JIT from the same sources and flags, -verifying every source digest before the first build. - -A route lists every program one public call may select. The forward routes -produce ``y`` from ``x``, ``weight`` and ``cu_seqlens``: ``fwd`` (FP32 taps), -``fwd_dt`` (BF16 taps) and ``fwd_bias`` (a per-channel bias, FP32 or BF16 taps -and bias). The backward routes produce ``dx`` and the deterministic -``dweight`` (and ``dbias``): ``bwd`` (FP32 taps and -gradient, SiLU' at the BF16-rounded preactivation), ``bwd_dt`` (BF16 taps and -/ or a BF16 ``dweight``), ``bwd_bias``, ``bwd_fp32g`` (SiLU' at the unrounded -FP32 preactivation for the weight gradient) and ``bwd_bias_fp32g``; -:func:`forward_route` / :func:`backward_route` map a call's arguments to its -route and the ``flags`` bit mask (``FLAG_*``) tells a route's programs which -of ``w``, ``dw``, ``b``, ``db`` are BF16. A route, forward or backward, may -carry several programs of its kernel and selects exactly one per call, -chosen by the channel count (and, for one record, the token count): a -record's ``select_rule`` (``channels_at_least``: ``C >= channels[arch]``; -``channels_tokens_at_least``: that and ``S >= tokens[arch]``) names the -calls it owns, of several owners the one of the highest channel threshold -runs, the record without a rule is the route's fallback, and -:func:`selected_modules` applies that rule (a pure function of the -architecture, the route, ``C`` and ``S``; the host never reads the data). -Token and channel counts are runtime scalars, so a changed problem size never -recompiles anything; building every library of the architecture -(:func:`specs_for_arch`) before a CUDA Graph capture keeps the JIT out of the -capture. - -A forward call is one ordinary launch of the program its route selects, a -dynamic token-chunk forward (round 9): the co-resident CTAs of a channel tile -take the tile's token chunks in order from a device counter, so the CTAs that -finish first take the tail instead of an equal share fixed at launch. Each -forward route ships the program of one eight-token stage per sub-tile and -chunk for every channel count on sm_100a and sm_103a and, on sm_107a, below -a per-route channel count (the fallback); from that count (C = 256 on -``fwd``, 512 on ``fwd_dt``, 1024 on ``fwd_bias``) sm_107a runs the program -of two stages per chunk. The counters -are the forward workspace ``work`` (``workspace_rule`` ``tile_counters``: -:func:`counter_words`, two int32 words per channel tile), zero before the first -launch; the last CTA of each tile resets them, so the buffer is zero after -every launch and every CUDA Graph replay and serves every later launch of one -stream. Two launches that may run at once (two streams, two graphs replayed -concurrently) need two buffers. - -A backward call is two ordinary launches on one stream. The selected program -(persistent, ``grid_mult x ctas_per_sm x SM count`` CTAs) writes ``dx`` and -one FP32 partial row of ``dweight`` (and, on the bias routes, of ``dbias``) -per CTA into the workspace and signals its dependents; then the route's reduce -program (``record["reduce"]``), launched programmatically dependent on it -(``launch.use_pdl``: its binding sets the programmatic-stream-serialization -attribute, and the kernel waits for the main grid before its first load), sums -the rows per channel in a fixed order and stores ``dw`` (and ``db``) in the -dtypes ``flags`` names. Until round 6 the backward was one cooperative launch -whose grid-wide barrier preceded the same reduction in-kernel. A cooperative -grid cannot start while any kernel is resident on another stream -- NCCL -collectives are, for most of a training backward (FlashInfer issue #7038) -- -whereas an ordinary launch whose CTAs find no free slot runs them in a later -wave. ``grid_mult`` (``grid_rule["grid_mult"][arch]``) over-decomposes the -token axis so that such a CTA trails by ``1 / grid_mult`` of a stream rather -than a whole one; the partial rows grow with it. The reduction order is -unchanged, so ``dweight`` is deterministic for a given ``(S, C, device)`` and -route and, at ``grid_mult`` 1, bitwise what the cooperative form produced. -Each backward route has its own reduce program: the default one, and builds -of it with ``dtype_flags`` (the ``dw`` store form), ``bias`` (the fifth, dbias, -slab and ``db``) and ``fp64`` (the ``_fp32g`` routes' FP64 sums). - -The host plan is the three functions :func:`plan` (the kernel scalars, with -the channel tile rule :func:`channel_split` of the record's role: the forward -and the backward rules differ where no multiple of 32 divides the channel -vectors), :func:`launch_grid` (the grid: -``co_resident`` records launch ``ctas_per_sm x SM count`` CTAs, ``occupancy`` -records ``grid_mult x ctas_per_sm x SM count``, less the CTAs -:func:`main_grid` drops to whole channel-tile ranks on the rows it names, -``channels_per_cta`` records ``ceil(C / channels_per_cta)``), -:func:`partial_rows` (the backward -workspace rows) and :func:`counter_words` (the forward's ``work`` words); -everything else derives from them. The public workspace is -sized by the backward program of the largest ``grid_mult x ctas_per_sm x -taps`` over every backward route (:func:`workspace_record`, the capacity -owner; the bias programs' partials carry a fifth, dbias, slab); each program -writes the leading rows of its own grid and its reduce program sums exactly -those. The forward workspace is sized by the forward program of the most -counter words over every forward route (:func:`forward_workspace_words`). - -The program name ``cake_causal_conv1d_train_<20 hex>`` is the first twenty hex -digits of the SHA-256 of the compact JSON ``{"build": , -"module": }`` of the kernel the program was rendered -from; the default reduce program is built with its kernel file's default -keywords and the forward programs with ``fast=True, sched="dynamic", -chunk_stages=1, pair_math=True`` (the dynamic token-chunk schedule of one -eight-token stage per sub-tile and chunk, the f32x2 pair math, the -branch-free in-document stage; ``y`` is bitwise the default build's) plus -the route's keywords, the sm_107a programs the channel count selects with -``chunk_stages=2`` instead (two eight-token stages per sub-tile and chunk); -every backward program carries ``revision=6`` (the -build revision of the split that retired the cooperative form; the exporter -hashes it so the split programs replaced the cooperative ones under new names); -the default route's fallback is built with ``vec=8, copy="tma2d", -stage_pair=True`` (eight channels, 16-byte vectors, per thread; 2-D TMA row -copies; each eight-token copy stage consumed as two four-token halves, the -dweight fold buffer in the stage ring) on sm_100a and sm_103a and with -``min_blocks=2`` (two CTAs per SM) beside those on sm_107a, where for -256 <= C < 512 from S = 262144 tokens the build of ``vec=8, min_blocks=2`` -(four-token copy stages, no ``stage_pair``) runs instead, and the backwards -the ``channels_at_least`` rules select, builds of ``vec=8, copy="tma2d", -stage_tokens=8, min_blocks=2`` (eight-token stages, two CTAs per SM): as is -on sm_100a from C = 512 and on sm_103a from 768, with ``fold_overlap=True, -fast_roll=True`` (the dweight fold buffer over the stage slots, the stage body -as two rolled passes of four tokens) on sm_103a for 512 <= C < 768 and on -sm_107a for 512 <= C < 1024, as is on sm_107a for 1024 <= C < 2048, and -with ``fast_roll=True, roll_prefetch=True`` (each rolled pass loads the -next pass's rows before it computes) on sm_107a from 2048. The programs of -the six other routes are the same bodies built with ``dtype_flags=True`` (the -``flags`` scalar; untyped tap / bias pointers) plus ``bias=True`` for the -bias routes and ``grad_pre="fp32"`` for the ``_fp32g`` routes, their backward -fallbacks without ``stage_pair`` (``vec=8, copy="tma2d"`` on sm_100a and -sm_103a, ``vec=8, min_blocks=2`` on sm_107a); each record's comment names its -build. -""" - -from __future__ import annotations - -import functools -from typing import Any, Sequence - -import tvm_ffi - -from . import cake_build - -# The derived leaves of every record (``kernel_symbol``, ``launch.block``, -# ``launch.cluster``, ``launch.cooperative``, ``launch.dynamic_smem_bytes``, -# ``launch.use_pdl``, ``arg_plan`` and ``logical``) are rewritten by the -# exporter from the rendered program; every other field is host plan. -MODULES: dict[str, dict[str, Any]] = { - # Route ``fwd`` (round 9): the forward's dynamic token-chunk program with the - # f32x2 pair math and the branch-free in-document stage, built with ``fast=True, - # sched="dynamic", chunk_stages=1, pair_math=True``. The CTAs of a channel tile - # take chunks of ``n_sub x 8`` tokens from a device counter instead of equal - # contiguous shares, so the CTAs of the SMs that draw the most bandwidth take - # the tail; ``y`` is bitwise the static program's. The counters are the - # ``work`` buffer (``workspace_rule``: int32, two words per channel tile: the - # next-chunk counter and the done-CTA counter), zero before the first launch; - # the last CTA of each tile resets them, so ``work`` is zero again after every - # launch and every CUDA Graph replay, and one buffer serves every later launch - # of one stream (two launches that may overlap need two buffers). Where a CTA - # holds several sub-tiles, a chunk's three history rows come from the stage's - # slot (the halo rows its copy brings), not from global memory. It is the - # route's fallback (no ``select_rule``): every channel count on sm_100a and - # sm_103a, and C < 256 on sm_107a. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_f5497330bfd286211643": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "y"], - ["buffer", "work"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # The grid is the co-resident CTA count, ``ctas_per_sm x SM count``: - # CTA ``b`` owns channel tile ``b % n_ch_tiles`` and, with the other - # CTAs of that tile, claims the tile's token chunks in order from - # ``work`` (a CTA's chunks rise; its first chunk past ``S`` ends it). - # Every channel tile needs a CTA (the host refuses a ``C`` of more - # tiles than the grid). ``ctas_per_sm`` 4 is the co-residency of the - # program (128 registers at 128 threads, no stack) on B200 (sm_100a), - # GB300 (sm_103a) and R200 (sm_107a). ``geometry`` is the kernel's - # build geometry the host plan derives its scalars from. - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 4, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_f5497330bfd286211643", - "kernel_symbol": "kernel_cake_causal_conv1d_train_f5497330bfd286211643", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 46208, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_f5497330bfd286211643", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "search_steps": "search_steps", - "w": "w", - "work": "work", - "x": "x", - "y": "y", - }, - "role": "fwd", - "workspace_rule": { - "kind": "tile_counters", - "words_per_tile": 2, - }, - }, - "cake_causal_conv1d_train_9bfa65707a49df6a48ce": { - "arches": [ - "sm_100a", - "sm_103a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # An ordinary launch of ``grid_mult x ctas_per_sm x SM count`` CTAs: - # ``ctas_per_sm`` is the measured co-residency, ``grid_mult`` (per - # architecture) splits the token streams that many times finer so a - # CTA another stream's kernel displaces trails by ``1 / grid_mult`` of - # a stream instead of a whole one. Each CTA writes one partial row; - # ``reduce`` sums them. The grid is a function of the device alone. - # The ``stage_pair`` build: each 8-token copy stage is consumed as two - # 4-token halves, the dweight fold buffer in the stage ring. - # ``ctas_per_sm``: 3 CTAs per SM (168 registers at 128 threads, the - # launch bound; ncu launch__occupancy_limit_registers), the occupancy - # query's answer on B200 and GB300. - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 3, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_9bfa65707a49df6a48ce", - "kernel_symbol": "kernel_cake_causal_conv1d_train_9bfa65707a49df6a48ce", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 66560, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_9bfa65707a49df6a48ce", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "role": "bwd", - # The FP32 ``partial`` workspace: ``[ceil(grid / n_ch_tiles), C, taps]`` - # rows (row ``b // n_ch_tiles`` belongs to CTA ``b``) the kernel fully - # writes and ``reduce`` sums in a fixed order; nothing is zeroed by the - # host. - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - "cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # An ordinary launch of ``grid_mult x ctas_per_sm x SM count`` CTAs: - # ``ctas_per_sm`` is the measured co-residency, ``grid_mult`` (per - # architecture) splits the token streams that many times finer so a - # CTA another stream's kernel displaces trails by ``1 / grid_mult`` of - # a stream instead of a whole one. Each CTA writes one partial row; - # ``reduce`` sums them. The grid is a function of the device alone. - # The ``copy="tma2d", stage_pair`` build of the sm_100a / sm_103a - # fallback at ``min_blocks`` 2: 2-D TMA stage copies, each 8-token - # stage consumed as two 4-token halves, the dweight fold buffer in the - # stage ring. ``ctas_per_sm``: 2 CTAs per SM, the program's - # ``min_blocks`` 2 launch bound (ncu launch__occupancy_limit_registers), - # the occupancy query's answer on R200; the ``min_blocks`` 3 build of - # the other two architectures measured 10 % slower at C = 128 on its - # 212 SMs. - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c", - "kernel_symbol": "kernel_cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 66560, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "role": "bwd", - # The FP32 ``partial`` workspace: ``[ceil(grid / n_ch_tiles), C, taps]`` - # rows (row ``b // n_ch_tiles`` belongs to CTA ``b``) the kernel fully - # writes and ``reduce`` sums in a fixed order; nothing is zeroed by the - # host. - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - "cake_causal_conv1d_train_9a07bcaa0223632640ee": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # An ordinary launch of ``grid_mult x ctas_per_sm x SM count`` CTAs: - # ``ctas_per_sm`` is the measured co-residency, ``grid_mult`` (per - # architecture) splits the token streams that many times finer so a - # CTA another stream's kernel displaces trails by ``1 / grid_mult`` of - # a stream instead of a whole one. Each CTA writes one partial row; - # ``reduce`` sums them. The grid is a function of the device alone. - # ``ctas_per_sm``: 2 CTAs per SM, the program's ``min_blocks`` 2 launch - # bound (234 registers at 128 threads; ncu - # launch__occupancy_limit_registers), the occupancy query's answer on - # R200; the ``min_blocks`` 3 build of the other two architectures - # measured 10 % slower at C = 128 on its 212 SMs. - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_9a07bcaa0223632640ee", - "kernel_symbol": "kernel_cake_causal_conv1d_train_9a07bcaa0223632640ee", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 50176, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_9a07bcaa0223632640ee", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - }, - "reduce": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "role": "bwd", - # The four-token build of the sm_107a fallback runs for 256 <= C < 512 - # from S = 262144 tokens; the ``stage_pair`` fallback runs every other - # C < 512 (on those rows its gain was not distinguishable from the - # spread between identical copies of one build). - "select_rule": { - "channels": { - "sm_107a": 256, - }, - "kind": "channels_tokens_at_least", - "tokens": { - "sm_107a": 262144, - }, - }, - # The FP32 ``partial`` workspace: ``[ceil(grid / n_ch_tiles), C, taps]`` - # rows (row ``b // n_ch_tiles`` belongs to CTA ``b``) the kernel fully - # writes and ``reduce`` sums in a fixed order; nothing is zeroed by the - # host. - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Rendered by the exporter from the family adapter's scaffold: the backward - # with stage_tokens 8, copy tma2d, min_blocks 2 (2 CTAs per SM); it runs when - # C >= select_rule.channels[arch] and no routed record of a higher threshold - # admits C; below every threshold the fallback record (no select_rule) runs. - "cake_causal_conv1d_train_37e428406a0cc4a8ca1b": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_37e428406a0cc4a8ca1b", - "kernel_symbol": "kernel_cake_causal_conv1d_train_37e428406a0cc4a8ca1b", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 82944, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_37e428406a0cc4a8ca1b", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "role": "bwd", - "select_rule": { - "channels": { - "sm_100a": 512, - "sm_103a": 768, - "sm_107a": 1024, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Rendered by the exporter from the family adapter's scaffold: the backward - # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, fold_overlap, - # fast_roll - # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no - # routed record of a higher threshold admits C; below every threshold the - # fallback record (no select_rule) runs. - "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6": { - "arches": [ - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_103a": 1, - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6", - "kernel_symbol": "kernel_cake_causal_conv1d_train_d46036aa2bc5f90bb0f6", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 66560, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "role": "bwd", - "select_rule": { - "channels": { - "sm_103a": 512, - "sm_107a": 512, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Rendered by the exporter from the family adapter's scaffold: the backward - # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, fast_roll, - # roll_prefetch - # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no - # routed record of a higher threshold admits C; below every threshold the - # fallback record (no select_rule) runs. - "cake_causal_conv1d_train_aeff6106e2822f313a0c": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_aeff6106e2822f313a0c", - "kernel_symbol": "kernel_cake_causal_conv1d_train_aeff6106e2822f313a0c", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 82944, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_aeff6106e2822f313a0c", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "role": "bwd", - "select_rule": { - "channels": { - "sm_107a": 2048, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Route ``fwd_dt`` (round 9): the dynamic token-chunk forward with - # ``dtype_flags`` (the ``flags`` scalar selects the tap load form: bit 0 set - # reads BF16 taps), built with ``dtype_flags=True, fast=True, - # sched="dynamic", chunk_stages=1, pair_math=True``; ``work`` as the ``fwd`` - # record's. The route's fallback: every channel count on sm_100a and - # sm_103a, and C < 512 on sm_107a. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_9c1e7ee9f37e3102660e": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "y"], - ["buffer", "work"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 4, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_9c1e7ee9f37e3102660e", - "kernel_symbol": "kernel_cake_causal_conv1d_train_9c1e7ee9f37e3102660e", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 46208, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_9c1e7ee9f37e3102660e", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "search_steps": "search_steps", - "w": "w", - "work": "work", - "x": "x", - "y": "y", - }, - "role": "fwd", - "workspace_rule": { - "kind": "tile_counters", - "words_per_tile": 2, - }, - }, - # Route ``fwd_bias`` (round 9): the dynamic token-chunk forward with a - # per-channel bias ``b`` added after the four taps (FLA's order) and - # ``dtype_flags`` (bit 0: BF16 taps, bit 2: BF16 bias), built with ``bias=True, - # dtype_flags=True, fast=True, sched="dynamic", chunk_stages=1, - # pair_math=True``; ``work`` as the ``fwd`` record's. The route's fallback: - # every channel count on sm_100a and sm_103a, and C < 1024 on sm_107a. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_08880ca12e4064bc8cdb": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "y"], - ["buffer", "b"], - ["buffer", "work"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 4, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_08880ca12e4064bc8cdb", - "kernel_symbol": "kernel_cake_causal_conv1d_train_08880ca12e4064bc8cdb", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 46208, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_08880ca12e4064bc8cdb", - "logical": { - "C": "C", - "S": "S", - "b": "b", - "ct_vec": "ct_vec", - "cu": "cu", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "search_steps": "search_steps", - "w": "w", - "work": "work", - "x": "x", - "y": "y", - }, - "role": "fwd", - "workspace_rule": { - "kind": "tile_counters", - "words_per_tile": 2, - }, - }, - # Route ``fwd`` on sm_107a from C = 256 (round 9): the dynamic token-chunk - # forward of two eight-token stages per sub-tile and chunk, built with - # ``fast=True, sched="dynamic", chunk_stages=2, pair_math=True``. On R200 the - # one-stage program loses to the shipped forward from C = 2048 (one sub-tile - # per CTA over several channel tiles) and wins at C = 128, where two stages - # lose; this program runs from the smallest measured C where it beats the - # one-stage program. ``work`` as the ``fwd`` record's. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_8665d0756b8240ecb4d0": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "y"], - ["buffer", "work"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 4, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_8665d0756b8240ecb4d0", - "kernel_symbol": "kernel_cake_causal_conv1d_train_8665d0756b8240ecb4d0", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 33920, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_8665d0756b8240ecb4d0", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "search_steps": "search_steps", - "w": "w", - "work": "work", - "x": "x", - "y": "y", - }, - "role": "fwd", - "select_rule": { - "channels": { - "sm_107a": 256, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "tile_counters", - "words_per_tile": 2, - }, - }, - # Route ``fwd_dt`` on sm_107a from C = 512 (round 9): the two-stage dynamic - # token-chunk forward with ``dtype_flags``, built with ``dtype_flags=True, - # fast=True, sched="dynamic", chunk_stages=2, pair_math=True``; ``work`` as - # the ``fwd`` record's. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_176ebb3e14f7ad9c94dc": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "y"], - ["buffer", "work"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 4, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_176ebb3e14f7ad9c94dc", - "kernel_symbol": "kernel_cake_causal_conv1d_train_176ebb3e14f7ad9c94dc", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 33920, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_176ebb3e14f7ad9c94dc", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "search_steps": "search_steps", - "w": "w", - "work": "work", - "x": "x", - "y": "y", - }, - "role": "fwd", - "select_rule": { - "channels": { - "sm_107a": 512, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "tile_counters", - "words_per_tile": 2, - }, - }, - # Route ``fwd_bias`` on sm_107a from C = 1024 (round 9): the two-stage dynamic - # token-chunk forward with the per-channel bias, built with ``bias=True, - # dtype_flags=True, fast=True, sched="dynamic", chunk_stages=2, - # pair_math=True``; ``work`` as the ``fwd`` record's. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_d829c95e1d610bc14b23": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "y"], - ["buffer", "b"], - ["buffer", "work"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 4, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_d829c95e1d610bc14b23", - "kernel_symbol": "kernel_cake_causal_conv1d_train_d829c95e1d610bc14b23", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 33920, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_d829c95e1d610bc14b23", - "logical": { - "C": "C", - "S": "S", - "b": "b", - "ct_vec": "ct_vec", - "cu": "cu", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "search_steps": "search_steps", - "w": "w", - "work": "work", - "x": "x", - "y": "y", - }, - "role": "fwd", - "select_rule": { - "channels": { - "sm_107a": 1024, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "tile_counters", - "words_per_tile": 2, - }, - }, - # Route ``bwd_dt`` (hand-written fallback, sm_100a / sm_103a): the fallback backward body - # ``vec=8, copy="tma2d"`` with ``dtype_flags=True`` (bit 0: BF16 taps; bit 1, the BF16 - # ``dw`` store, is read by its reduce program); SiLU' at the BF16-rounded preactivation - # as the default route. Like every backward program since the split it is an ordinary - # launch that writes ``dx`` and its partial rows; ``reduce`` names the program that sums - # them (the ``bwd_dt`` reduce: ``dtype_flags`` for the ``dw`` store form). - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa": { - "arches": [ - "sm_100a", - "sm_103a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 3, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa", - "kernel_symbol": "kernel_cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 50176, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_ea8f91038a9dea452ff2", - "role": "bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Route ``bwd_dt`` (hand-written fallback, sm_107a): the sm_107a fallback body ``vec=8, - # min_blocks=2`` with ``dtype_flags=True``. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_8b9b8db0833fd64a578f": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_8b9b8db0833fd64a578f", - "kernel_symbol": "kernel_cake_causal_conv1d_train_8b9b8db0833fd64a578f", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 50176, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_8b9b8db0833fd64a578f", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - }, - "reduce": "cake_causal_conv1d_train_ea8f91038a9dea452ff2", - "role": "bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Route ``bwd_bias`` (hand-written fallback, sm_100a / sm_103a): the fallback body with - # ``bias=True`` (reads ``b``; the ``partial`` workspace grows a fifth slab, ``taps`` 5, - # of dbias partials that its reduce program sums into ``db``) and ``dtype_flags=True`` - # (input bits 0 and 2; the store bits 1 and 3 are the reduce program's), ``min_blocks=2`` - # (the bias accumulators do not fit the 168-register cap of three CTAs per SM). - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_2d91e5f926238453c78c": { - "arches": [ - "sm_100a", - "sm_103a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["buffer", "b"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_2d91e5f926238453c78c", - "kernel_symbol": "kernel_cake_causal_conv1d_train_2d91e5f926238453c78c", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 54272, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_2d91e5f926238453c78c", - "logical": { - "C": "C", - "S": "S", - "b": "b", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_1c21c181af94e1635851", - "role": "bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 5, - }, - }, - # Route ``bwd_bias`` (hand-written fallback, sm_107a): the sm_107a fallback body with - # ``bias=True, dtype_flags=True``. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_cc7a117aee9c7ae829c4": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["buffer", "b"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_cc7a117aee9c7ae829c4", - "kernel_symbol": "kernel_cake_causal_conv1d_train_cc7a117aee9c7ae829c4", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 54272, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_cc7a117aee9c7ae829c4", - "logical": { - "C": "C", - "S": "S", - "b": "b", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - }, - "reduce": "cake_causal_conv1d_train_1c21c181af94e1635851", - "role": "bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 5, - }, - }, - # Route ``bwd_fp32g`` (hand-written fallback, sm_100a / sm_103a): the fallback body with - # ``grad_pre="fp32"`` (SiLU' at the unrounded FP32 preactivation; its reduce program - # sums the partial rows in FP64 with one rounding), ``fold_every=4, fold_kahan=True`` - # and ``dtype_flags=True``, built ``min_blocks=2``: two CTAs per SM (the FP32-SiLU' - # body spills at the 168-register cap of three), so ``ctas_per_sm`` 2. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_c8027a64ba3cc38333bf": { - "arches": [ - "sm_100a", - "sm_103a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_c8027a64ba3cc38333bf", - "kernel_symbol": "kernel_cake_causal_conv1d_train_c8027a64ba3cc38333bf", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 66560, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_c8027a64ba3cc38333bf", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_ce49146b9fef24adb47e", - "role": "bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Route ``bwd_fp32g`` (hand-written fallback, sm_107a): the sm_107a fallback body with - # ``grad_pre="fp32", dtype_flags=True``. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_f1545adfd2bb28044419": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_f1545adfd2bb28044419", - "kernel_symbol": "kernel_cake_causal_conv1d_train_f1545adfd2bb28044419", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 66560, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_f1545adfd2bb28044419", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - }, - "reduce": "cake_causal_conv1d_train_ce49146b9fef24adb47e", - "role": "bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Route ``bwd_bias_fp32g`` (hand-written fallback, sm_100a / sm_103a): the fallback body - # with ``bias=True, grad_pre="fp32", dtype_flags=True``. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_149ca6f53934ba630033": { - "arches": [ - "sm_100a", - "sm_103a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["buffer", "b"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_149ca6f53934ba630033", - "kernel_symbol": "kernel_cake_causal_conv1d_train_149ca6f53934ba630033", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 74752, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_149ca6f53934ba630033", - "logical": { - "C": "C", - "S": "S", - "b": "b", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea", - "role": "bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 5, - }, - }, - # Route ``bwd_bias_fp32g`` (hand-written fallback, sm_107a): the sm_107a fallback body - # with ``bias=True, grad_pre="fp32", dtype_flags=True``. - # ``dynamic_smem_bytes`` 0 is the placeholder of a program whose shared memory follows - # the kernel (the exporter's ``KERNEL_SMEM_PROGRAMS``): the first render writes the - # kernel's figure, the kernel symbol, the argument plan and the logical names it keeps. - "cake_causal_conv1d_train_843c010db78c79107fd3": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["buffer", "b"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_843c010db78c79107fd3", - "kernel_symbol": "kernel_cake_causal_conv1d_train_843c010db78c79107fd3", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 74752, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_843c010db78c79107fd3", - "logical": { - "C": "C", - "S": "S", - "b": "b", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - }, - "reduce": "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea", - "role": "bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 5, - }, - }, - # Rendered by the exporter from the family adapter's scaffold: the backward - # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, dtype_flags, - # fold_layout thread - # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no - # routed record of a higher threshold admits C; below every threshold the - # fallback record (no select_rule) runs. - "cake_causal_conv1d_train_2675cc1cb4efc0937378": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_2675cc1cb4efc0937378", - "kernel_symbol": "kernel_cake_causal_conv1d_train_2675cc1cb4efc0937378", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 82944, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_2675cc1cb4efc0937378", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_ea8f91038a9dea452ff2", - "role": "bwd", - "select_rule": { - "channels": { - "sm_100a": 512, - "sm_103a": 512, - "sm_107a": 1024, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Rendered by the exporter from the family adapter's scaffold: the backward - # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, bias, dtype_flags - # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no - # routed record of a higher threshold admits C; below every threshold the - # fallback record (no select_rule) runs. - "cake_causal_conv1d_train_517710d1f0988cf570a9": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["buffer", "b"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_517710d1f0988cf570a9", - "kernel_symbol": "kernel_cake_causal_conv1d_train_517710d1f0988cf570a9", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 87040, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_517710d1f0988cf570a9", - "logical": { - "C": "C", - "S": "S", - "b": "b", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_1c21c181af94e1635851", - "role": "bwd", - "select_rule": { - "channels": { - "sm_100a": 512, - "sm_103a": 512, - "sm_107a": 1024, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "partial_rows", - "taps": 5, - }, - }, - # Rendered by the exporter from the family adapter's scaffold: the backward - # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, grad_pre fp32, - # fold_every 4, fold_kahan, dtype_flags - # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no - # routed record of a higher threshold admits C; below every threshold the - # fallback record (no select_rule) runs. - "cake_causal_conv1d_train_170d89481a616ae977b3": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_170d89481a616ae977b3", - "kernel_symbol": "kernel_cake_causal_conv1d_train_170d89481a616ae977b3", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 99328, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_170d89481a616ae977b3", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_ce49146b9fef24adb47e", - "role": "bwd", - "select_rule": { - "channels": { - "sm_100a": 512, - "sm_103a": 512, - "sm_107a": 1024, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - # Rendered by the exporter from the family adapter's scaffold: the backward - # with vec 8, copy tma2d, stage_tokens 8, min_blocks 2, bias, grad_pre fp32, - # fold_every 4, fold_kahan, dtype_flags - # (2 CTAs per SM); it runs when C >= select_rule.channels[arch] and no - # routed record of a higher threshold admits C; below every threshold the - # fallback record (no select_rule) runs. - "cake_causal_conv1d_train_a681165b3e81204c73ff": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "partial"], - ["buffer", "b"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["parameter", "flags"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "grid_mult": { - "sm_100a": 1, - "sm_103a": 1, - "sm_107a": 1, - }, - "kind": "occupancy", - }, - "kernel": "cake_causal_conv1d_train_a681165b3e81204c73ff", - "kernel_symbol": "kernel_cake_causal_conv1d_train_a681165b3e81204c73ff", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 107520, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_a681165b3e81204c73ff", - "logical": { - "C": "C", - "S": "S", - "b": "b", - "ct_vec": "ct_vec", - "cu": "cu", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "flags": "flags", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "x_map": "x", - }, - "reduce": "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea", - "role": "bwd", - "select_rule": { - "channels": { - "sm_100a": 512, - "sm_103a": 512, - "sm_107a": 1024, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "partial_rows", - "taps": 5, - }, - }, - # The reduce program every backward route shares: the second launch of a - # backward call. Launched programmatically dependent on the main program - # (``use_pdl``: the binding sets the programmatic-stream-serialization - # launch attribute; the kernel's first statement waits for the main grid, - # so the launch latency overlaps the main kernel's tail and an ordinary - # launch stays correct), it sums the FP32 ``[C, 4]`` partial rows of the - # workspace per channel: channel ``c`` lies in tile ``t = c // (vec x - # ct_vec)``, whose ``ceil((main_ctas - t) / n_ch_tiles)`` rows the main - # grid wrote, in the order the cooperative backward used in-kernel -- one - # warp per channel, lane sums over rows ``lane, lane + 32, ...`` then the - # xor tree 16, 8, 4, 2, 1 -- so ``dweight`` is bitwise what one launch - # produced. Built with its kernel file's default keywords, one program - # for the three architectures. - "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "partial"], - ["buffer", "dw"], - ["parameter", "C"], - ["parameter", "main_ctas"], - ["parameter", "n_ch_tiles"], - ["parameter", "ct_vec"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # One warp per channel, four channels per 128-thread CTA: the grid is - # ``ceil(C / channels_per_cta)`` CTAs, a function of the call's ``C``. - "grid_rule": { - "channels_per_cta": 4, - "cooperative": False, - "kind": "channels_per_cta", - }, - "kernel": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "kernel_symbol": "kernel_cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 2048, - "use_pdl": True, - }, - "library": "cake_causal_conv1d_train_86b5ecfbd90aca4b52e5", - "logical": { - "C": "C", - "ct_vec": "ct_vec", - "dw": "dw", - "main_ctas": "main_ctas", - "n_ch_tiles": "n_ch_tiles", - "partial": "partial", - }, - "role": "reduce", - "workspace_rule": None, - }, - # The reduce program of the ``bwd_dt`` route (hand-written; the render fills the name, - # symbol, argument plan and logical map): the default reduce built with ``dtype_flags=True``; it - # stores ``dw`` as BF16 when ``flags`` bit 1 says so; - # otherwise the default reduce's order and CTA shape (one warp per channel). - "cake_causal_conv1d_train_ea8f91038a9dea452ff2": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "partial"], - ["buffer", "dw"], - ["parameter", "C"], - ["parameter", "main_ctas"], - ["parameter", "n_ch_tiles"], - ["parameter", "ct_vec"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "channels_per_cta": 4, - "cooperative": False, - "kind": "channels_per_cta", - }, - "kernel": "cake_causal_conv1d_train_ea8f91038a9dea452ff2", - "kernel_symbol": "kernel_cake_causal_conv1d_train_ea8f91038a9dea452ff2", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 2048, - "use_pdl": True, - }, - "library": "cake_causal_conv1d_train_ea8f91038a9dea452ff2", - "logical": { - "C": "C", - "ct_vec": "ct_vec", - "dw": "dw", - "flags": "flags", - "main_ctas": "main_ctas", - "n_ch_tiles": "n_ch_tiles", - "partial": "partial", - }, - "role": "reduce", - "workspace_rule": None, - }, - # The reduce program of the ``bwd_bias`` route (hand-written; the render fills the name, - # symbol, argument plan and logical map): the default reduce built with ``bias=True, dtype_flags=True``; it - # sums the fifth, dbias, slab of the partial rows into ``db`` as well; - # stores ``dw`` / ``db`` as BF16 when ``flags`` bit 1 / 3 says so; - # otherwise the default reduce's order and CTA shape (one warp per channel). - "cake_causal_conv1d_train_1c21c181af94e1635851": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "partial"], - ["buffer", "dw"], - ["buffer", "db"], - ["parameter", "C"], - ["parameter", "main_ctas"], - ["parameter", "n_ch_tiles"], - ["parameter", "ct_vec"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "channels_per_cta": 4, - "cooperative": False, - "kind": "channels_per_cta", - }, - "kernel": "cake_causal_conv1d_train_1c21c181af94e1635851", - "kernel_symbol": "kernel_cake_causal_conv1d_train_1c21c181af94e1635851", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 2560, - "use_pdl": True, - }, - "library": "cake_causal_conv1d_train_1c21c181af94e1635851", - "logical": { - "C": "C", - "ct_vec": "ct_vec", - "db": "db", - "dw": "dw", - "flags": "flags", - "main_ctas": "main_ctas", - "n_ch_tiles": "n_ch_tiles", - "partial": "partial", - }, - "role": "reduce", - "workspace_rule": None, - }, - # The reduce program of the ``bwd_fp32g`` route (hand-written; the render fills the name, - # symbol, argument plan and logical map): the default reduce built with ``fp64=True, dtype_flags=True``; it - # sums in FP64 (lane sums and the xor tree) with one final rounding; - # stores ``dw`` as BF16 when ``flags`` bit 1 says so; - # otherwise the default reduce's order and CTA shape (one warp per channel). - "cake_causal_conv1d_train_ce49146b9fef24adb47e": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "partial"], - ["buffer", "dw"], - ["parameter", "C"], - ["parameter", "main_ctas"], - ["parameter", "n_ch_tiles"], - ["parameter", "ct_vec"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "channels_per_cta": 4, - "cooperative": False, - "kind": "channels_per_cta", - }, - "kernel": "cake_causal_conv1d_train_ce49146b9fef24adb47e", - "kernel_symbol": "kernel_cake_causal_conv1d_train_ce49146b9fef24adb47e", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 4096, - "use_pdl": True, - }, - "library": "cake_causal_conv1d_train_ce49146b9fef24adb47e", - "logical": { - "C": "C", - "ct_vec": "ct_vec", - "dw": "dw", - "flags": "flags", - "main_ctas": "main_ctas", - "n_ch_tiles": "n_ch_tiles", - "partial": "partial", - }, - "role": "reduce", - "workspace_rule": None, - }, - # The reduce program of the ``bwd_bias_fp32g`` route (hand-written; the render fills the name, - # symbol, argument plan and logical map): the default reduce built with ``bias=True, fp64=True, dtype_flags=True``; it - # sums the fifth, dbias, slab of the partial rows into ``db`` as well; - # sums in FP64 (lane sums and the xor tree) with one final rounding; - # stores ``dw`` / ``db`` as BF16 when ``flags`` bit 1 / 3 says so; - # otherwise the default reduce's order and CTA shape (one warp per channel). - "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "partial"], - ["buffer", "dw"], - ["buffer", "db"], - ["parameter", "C"], - ["parameter", "main_ctas"], - ["parameter", "n_ch_tiles"], - ["parameter", "ct_vec"], - ["parameter", "flags"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - "grid_rule": { - "channels_per_cta": 4, - "cooperative": False, - "kind": "channels_per_cta", - }, - "kernel": "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea", - "kernel_symbol": "kernel_cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 5120, - "use_pdl": True, - }, - "library": "cake_causal_conv1d_train_4d291f4b5c2cce6ef1ea", - "logical": { - "C": "C", - "ct_vec": "ct_vec", - "db": "db", - "dw": "dw", - "flags": "flags", - "main_ctas": "main_ctas", - "n_ch_tiles": "n_ch_tiles", - "partial": "partial", - }, - "role": "reduce", - "workspace_rule": None, - }, -} - -# Architecture -> route -> the records a call of the route selects from (the -# fallback first). The default routes ``fwd`` / ``bwd`` are the shipped -# programs; the six round-6 routes each hold their hand-written fallback -# (above) and, once rendered, the routed programs the exporter scaffolds from -# it; the forward routes' sm_107a two-stage programs are hand-written records. -ROUTES: dict[str, dict[str, list[str]]] = { - "sm_100a": { - "bwd": [ - "cake_causal_conv1d_train_9bfa65707a49df6a48ce", - "cake_causal_conv1d_train_37e428406a0cc4a8ca1b", - ], - "bwd_bias": [ - "cake_causal_conv1d_train_2d91e5f926238453c78c", - "cake_causal_conv1d_train_517710d1f0988cf570a9", - ], - "bwd_bias_fp32g": [ - "cake_causal_conv1d_train_149ca6f53934ba630033", - "cake_causal_conv1d_train_a681165b3e81204c73ff", - ], - "bwd_dt": [ - "cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa", - "cake_causal_conv1d_train_2675cc1cb4efc0937378", - ], - "bwd_fp32g": [ - "cake_causal_conv1d_train_c8027a64ba3cc38333bf", - "cake_causal_conv1d_train_170d89481a616ae977b3", - ], - "fwd": [ - "cake_causal_conv1d_train_f5497330bfd286211643", - ], - "fwd_bias": [ - "cake_causal_conv1d_train_08880ca12e4064bc8cdb", - ], - "fwd_dt": [ - "cake_causal_conv1d_train_9c1e7ee9f37e3102660e", - ], - }, - "sm_103a": { - "bwd": [ - "cake_causal_conv1d_train_9bfa65707a49df6a48ce", - "cake_causal_conv1d_train_37e428406a0cc4a8ca1b", - "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6", - ], - "bwd_bias": [ - "cake_causal_conv1d_train_2d91e5f926238453c78c", - "cake_causal_conv1d_train_517710d1f0988cf570a9", - ], - "bwd_bias_fp32g": [ - "cake_causal_conv1d_train_149ca6f53934ba630033", - "cake_causal_conv1d_train_a681165b3e81204c73ff", - ], - "bwd_dt": [ - "cake_causal_conv1d_train_9c8310a1d3fa9b8ccffa", - "cake_causal_conv1d_train_2675cc1cb4efc0937378", - ], - "bwd_fp32g": [ - "cake_causal_conv1d_train_c8027a64ba3cc38333bf", - "cake_causal_conv1d_train_170d89481a616ae977b3", - ], - "fwd": [ - "cake_causal_conv1d_train_f5497330bfd286211643", - ], - "fwd_bias": [ - "cake_causal_conv1d_train_08880ca12e4064bc8cdb", - ], - "fwd_dt": [ - "cake_causal_conv1d_train_9c1e7ee9f37e3102660e", - ], - }, - "sm_107a": { - "bwd": [ - "cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c", - "cake_causal_conv1d_train_37e428406a0cc4a8ca1b", - "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6", - "cake_causal_conv1d_train_aeff6106e2822f313a0c", - "cake_causal_conv1d_train_9a07bcaa0223632640ee", - ], - "bwd_bias": [ - "cake_causal_conv1d_train_cc7a117aee9c7ae829c4", - "cake_causal_conv1d_train_517710d1f0988cf570a9", - ], - "bwd_bias_fp32g": [ - "cake_causal_conv1d_train_843c010db78c79107fd3", - "cake_causal_conv1d_train_a681165b3e81204c73ff", - ], - "bwd_dt": [ - "cake_causal_conv1d_train_8b9b8db0833fd64a578f", - "cake_causal_conv1d_train_2675cc1cb4efc0937378", - ], - "bwd_fp32g": [ - "cake_causal_conv1d_train_f1545adfd2bb28044419", - "cake_causal_conv1d_train_170d89481a616ae977b3", - ], - "fwd": [ - "cake_causal_conv1d_train_f5497330bfd286211643", - "cake_causal_conv1d_train_8665d0756b8240ecb4d0", - ], - "fwd_bias": [ - "cake_causal_conv1d_train_08880ca12e4064bc8cdb", - "cake_causal_conv1d_train_d829c95e1d610bc14b23", - ], - "fwd_dt": [ - "cake_causal_conv1d_train_9c1e7ee9f37e3102660e", - "cake_causal_conv1d_train_176ebb3e14f7ad9c94dc", - ], - }, -} - -# One verified build per architecture: B200 (``sm_100a``), B300 (``sm_103a``) -# and R200 (``sm_107a``). Every architecture is compiled with its exact -# ``-gencode`` flag set (``build.json`` ``arch_flags``); there is no family -# fallback. -ARCHES = ("sm_100a", "sm_103a", "sm_107a") -ARCH_BY_CAPABILITY: dict[tuple[int, int], str] = { - (10, 0): "sm_100a", - (10, 3): "sm_103a", - (10, 7): "sm_107a", -} -# The routes (``ROUTES[arch]`` keys) a public call may resolve to. ``fwd`` / -# ``bwd`` are the default routes (FP32 taps, no bias, SiLU' at the BF16-rounded -# preactivation); ``_dt`` routes take BF16 taps and / or store a BF16 gradient -# (the ``flags`` scalar says which), ``_bias`` routes add the per-channel bias -# and its gradient, ``_fp32g`` routes evaluate SiLU' at the unrounded FP32 -# preactivation for the weight gradient. :func:`forward_route` and -# :func:`backward_route` are the host's map from a call's arguments to a route. -KINDS = ( - "fwd", - "fwd_dt", - "fwd_bias", - "bwd", - "bwd_dt", - "bwd_bias", - "bwd_fp32g", - "bwd_bias_fp32g", -) -# The routes every exported architecture ships: the base forward and backward. -# The other routes are optional per architecture; a call that needs one the -# architecture lacks is refused with a ``ValueError`` naming the route. -REQUIRED_KINDS = ("fwd", "bwd") -# The role of a module record: a forward program, a backward program (the -# first launch of a backward call) or the reduce program a backward record -# names in ``record["reduce"]`` (the second launch). -ROLES = ("fwd", "bwd", "reduce") -# The grid-rule vocabulary (``record["grid_rule"]["kind"]``, see -# :func:`launch_grid`): ``co_resident`` launches exactly ``ctas_per_sm x SM -# count`` CTAs, the co-resident count (the forward records and the -# cooperative KCP programs of :mod:`.kcp_jit`, which share this host plan); -# ``occupancy`` launches ``grid_mult[arch] x ctas_per_sm x SM count`` CTAs -# normally (the backward programs since the split; ``grid_mult`` names every -# architecture of its record); ``channels_per_cta`` launches ``ceil(C / -# channels_per_cta)`` CTAs (the reduce programs). -GRID_KINDS = ("co_resident", "occupancy", "channels_per_cta") -# The selection-rule vocabulary of a forward or backward record -# (``record["select_rule"]["kind"]``); of the records whose rule admits a call -# the one of the highest channel threshold runs, and a record without a rule is -# its route's fallback. ``channels_tokens_at_least`` adds a token floor to the -# channel threshold. A reduce record never carries one. -SELECT_KINDS = ("channels_at_least", "channels_tokens_at_least") -# The ``flags`` bit mask of the ``dtype_flags`` programs (every program of the -# six round-6 routes): the host derives it from the tap tensor, the ACTUAL -# ``dweight`` out tensor, the bias tensor and the ACTUAL ``dbias`` out tensor; -# the kernel reads it once in its prologue (load form) and once in its epilogue -# (store form). The default routes' programs take no ``flags``. -FLAG_W_BF16 = 1 # ``w`` holds BF16 taps (else FP32) -FLAG_DW_BF16 = 2 # ``dw`` is stored as BF16 (else FP32) -FLAG_B_BF16 = 4 # ``b`` holds BF16 (else FP32); bias programs -FLAG_DB_BF16 = 8 # ``db`` is stored as BF16 (else FP32); bias programs -# The convolution width: ``weight`` and ``dweight`` are ``[C, WIDTH]`` and tap -# ``WIDTH - 1`` multiplies the current token. -WIDTH = 4 -# A forward thread and a backward thread each own one 16-byte vector of -# eight BF16 channels, so ``C`` is a multiple of eight; the ``vec`` of every -# record's geometry divides this. -CHANNEL_MULTIPLE = 8 -# The kernels take ``S``, ``C`` and the plan scalars as 32-bit integers. -INT32_MAX = 2**31 - 1 -# Logical argument names the public API binds; every module record maps its -# kernel parameter names onto this vocabulary (``record["logical"]``). -LOGICAL_TENSORS = ( - "x", - "w", - "cu", - "y", - "dy", - "dx", - "dw", - "partial", - "b", - "db", - "work", +sys.modules[__name__] = import_module( + "flashinfer.cake_causal_conv1d_det_train_backend.cake_jit" ) -# The workspace-rule vocabulary (``record["workspace_rule"]["kind"]``): a -# backward program's ``partial_rows`` (its FP32 partial rows, :func:`partial_bytes`) -# and a forward program's ``tile_counters`` (its int32 ``work`` counters, -# ``words_per_tile`` per channel tile, :func:`counter_words`); a reduce -# program has none. -WORKSPACE_KINDS = ("partial_rows", "tile_counters") -# The scalars :func:`plan` produces, in kernel order; ``flags`` (the dtype bit -# mask of the ``dtype_flags`` programs) and ``main_ctas`` (the CTA count of the -# selected backward program's launch, ``grid_x``, from which the reduce program -# derives the partial rows of each channel tile, ``ceil((main_ctas - tile) / -# n_ch_tiles)``) are set by the host beside them. -PLAN_SCALARS = ( - "S", - "C", - "num_docs", - "search_steps", - "ct_vec", - "n_sub", - "n_ch_tiles", -) -LOGICAL_SCALARS = (*PLAN_SCALARS, "flags", "main_ctas") - - -def arch_for_capability(device_capability: Sequence[int]) -> str | None: - """Name the exported architecture for one device capability, if any.""" - - values = [int(value) for value in device_capability] - if len(values) != 2: - return None - return ARCH_BY_CAPABILITY.get((values[0], values[1])) - - -def route_applies(*, device_capability: Sequence[int], kind: str) -> bool: - """Whether the verified export owns one training call on a device.""" - - arch = arch_for_capability(device_capability) - if arch is None: - return False - return kind in ROUTES.get(arch, {}) - - -def route_role(kind: str) -> str: - """``"fwd"`` or ``"bwd"``: the kernel a route's programs are built from (``ValueError`` for an unknown route).""" - - if kind not in KINDS: - raise ValueError(f"kind must be one of {KINDS}, got {kind!r}") - return "fwd" if kind.startswith("fwd") else "bwd" - - -def forward_route(*, bias: bool, bf16_taps: bool) -> str: - """The forward route of a call: ``fwd_bias`` with a bias, else ``fwd_dt`` for BF16 taps, else ``fwd``. - - A bias program always has a bias (no runtime switch: adding ``+0.0`` would - flip the sign of ``-0.0`` outputs), so ``bias=None`` never selects one; the - bias programs read FP32 or BF16 taps and bias through ``flags``. - """ - - if bias: - return "fwd_bias" - return "fwd_dt" if bf16_taps else "fwd" - - -def backward_route(*, bias: bool, fp32_grad: bool, bf16_taps_or_dweight: bool) -> str: - """The backward route of a call. - - ``bias`` selects ``bwd_bias`` / ``bwd_bias_fp32g`` (by ``fp32_grad``, the - ``silu_grad_preactivation="fp32"`` policy); without a bias ``fp32_grad`` - selects ``bwd_fp32g``; else BF16 taps or a BF16 ``dweight`` out tensor - select ``bwd_dt``; else the default ``bwd``. Every route but ``bwd`` - reads the tap / gradient dtypes from ``flags``. - """ - - if bias: - return "bwd_bias_fp32g" if fp32_grad else "bwd_bias" - if fp32_grad: - return "bwd_fp32g" - return "bwd_dt" if bf16_taps_or_dweight else "bwd" - - -def route_modules(arch: str, kind: str) -> tuple[str, ...]: - """Ordered module names of one route on ``arch`` (every program the route may launch). - - A call launches the one :func:`selected_modules` picks for its channel - count. - """ - - route_role(kind) - names = ROUTES.get(arch, {}).get(kind) - if not names: - raise ValueError( - f"unsupported Cake causal conv1d training route {kind!r} on {arch!r}; " - f"exported architectures: {ARCHES}" - ) - return tuple(names) - - -def select_rule(record: dict[str, Any]) -> dict[str, Any] | None: - """The selection rule of one module record, ``None`` when the record carries none. - - A record without a rule (no ``select_rule`` key, or ``None``) is its - route's fallback program. ``{"kind": "channels_at_least", "channels": - {arch: C_min, ...}}`` owns the calls whose channel count is at least - ``channels[arch]`` on the architectures it names and never hits on the - others. ``{"kind": "channels_tokens_at_least", "channels": {arch: C_min, - ...}, "tokens": {arch: S_min, ...}}`` owns those calls only when the token - count is also at least ``tokens[arch]``. When several records own a call, - the one of the highest channel threshold runs. - """ - - return record.get("select_rule") - - -def _threshold(rule: dict[str, Any], *, arch: str) -> int | None: - """The channel count from which a selection rule owns the calls on ``arch``, ``None`` where it never does - (``ValueError`` for an unknown kind).""" - - kind = rule["kind"] - if kind in SELECT_KINDS: - channels = rule["channels"] - return int(channels[arch]) if arch in channels else None - raise ValueError( - f"unknown Cake causal conv1d training selection rule: {kind!r} " - f"(known: {SELECT_KINDS})" - ) - - -def _tokens(rule: dict[str, Any], *, arch: str) -> int | None: - """The token count from which a selection rule owns the calls on ``arch`` (its token floor), ``None`` when the - rule has none there (``ValueError`` for an unknown kind).""" - - _threshold(rule, arch=arch) - if rule["kind"] != "channels_tokens_at_least": - return None - tokens = rule["tokens"] - return int(tokens[arch]) if arch in tokens else None - - -def _admits(rule: dict[str, Any], *, arch: str, C: int, S: int | None = None) -> bool: - """Whether a selection rule owns ``C`` channels (and ``S`` tokens) on ``arch``. - - ``ValueError`` for an unknown kind, or when the rule's token floor decides - and ``S`` is ``None``. - """ - - threshold = _threshold(rule, arch=arch) - if threshold is None or int(C) < threshold: - return False - floor = _tokens(rule, arch=arch) - if floor is None: - return True - if S is None: - raise ValueError( - f"the Cake causal conv1d training selection rule {rule!r} on {arch} " - f"for C={int(C)} runs from S {floor}; the selection needs S" - ) - return int(S) >= floor - - -def selected_modules( - arch: str, kind: str, *, C: int, S: int | None = None -) -> tuple[str, ...]: - """Ordered module names one public call launches on ``arch`` for ``C`` channels and ``S`` tokens. - - A route may carry several programs of its kernel and launches exactly one: - of the records whose :func:`select_rule` admits the call on ``arch`` the - one of the highest channel threshold, else the route's single fallback - record (the one without a rule). The choice is a pure function of the - architecture, the route, ``C`` and ``S`` -- ``S`` enters only through a - rule's token floor, and the host never reads the data -- so a given ``(S, - C, device)`` always runs the same program of a route and ``dweight`` - stays bitwise reproducible. Two records of one channel threshold that - both reach ``C``, or zero or several fallbacks, are registry errors - (``RuntimeError``); an unknown rule kind, or ``S`` of ``None`` where a - token floor decides, is a ``ValueError``. - """ - - names = route_modules(arch, kind) - candidates: list[tuple[int, str]] = [] - fallback: list[str] = [] - for name in names: - rule = select_rule(MODULES[name]) - if rule is None: - fallback.append(name) - continue - threshold = _threshold(rule, arch=arch) - if threshold is not None and int(C) >= threshold: - candidates.append((threshold, name)) - if len(fallback) != 1: - raise RuntimeError( - f"the exported route {kind!r} on {arch} declares {len(fallback)} " - "fallback programs (records without a select_rule); expected exactly one" - ) - candidates.sort(reverse=True) - if len(candidates) > 1 and candidates[0][0] == candidates[1][0]: - tied = [name for threshold, name in candidates if threshold == candidates[0][0]] - raise RuntimeError( - f"the exported route {kind!r} on {arch} selects {len(tied)} " - f"programs for C={int(C)} at one threshold {candidates[0][0]} ({tied}); " - "thresholds must be distinct" - ) - for _threshold_c, name in candidates: - if _admits(select_rule(MODULES[name]), arch=arch, C=C, S=S): - return (name,) - return (fallback[0],) - - -def backward_kinds(arch: str) -> tuple[str, ...]: - """The backward routes ``arch`` exports, in :data:`KINDS` order.""" - - return tuple( - kind - for kind in KINDS - if route_role(kind) == "bwd" and ROUTES.get(arch, {}).get(kind) - ) - - -def workspace_owners(arch: str) -> tuple[dict[str, Any], ...]: - """The backward module records of ``arch`` that own a ``partial`` workspace, over every backward route in - :data:`KINDS` order, each route in launch order. - - Every backward program declares a ``workspace_rule``; an architecture - without one is a registry error (``RuntimeError``). - """ - - owners = tuple( - MODULES[name] - for kind in backward_kinds(arch) - for name in route_modules(arch, kind) - if MODULES[name]["workspace_rule"] is not None - ) - if not owners: - raise RuntimeError( - f"the exported backward routes on {arch} declare no workspace owner; " - "expected at least one" - ) - return owners - - -def workspace_record(arch: str) -> dict[str, Any]: - """The capacity owner of the backward ``partial`` workspace on ``arch``. - - The public workspace serves every backward program of every backward route - (:func:`workspace_owners`), each of which writes the leading rows of its - own grid. The owners share one build geometry, so for every ``C`` their - rows are ``ceil(grid_mult x ctas_per_sm x SM count / n_ch_tiles)`` with - one ``n_ch_tiles`` and their bytes ``4 x taps x C x rows``: the record - with the largest :func:`grid_per_sm` ``x taps`` (the bias programs' fifth - slab counts; :data:`KINDS` then route order break ties) bounds them all. - Owners of different geometries are a registry error (``RuntimeError``). - """ - - owners = workspace_owners(arch) - geometries = { - tuple(sorted(owner["grid_rule"]["geometry"].items())) for owner in owners - } - if len(geometries) != 1: - raise RuntimeError( - f"the backward programs on {arch} declare {len(geometries)} build " - "geometries; one workspace serves them only when they share one" - ) - return max( - owners, - key=lambda owner: ( - grid_per_sm(owner, arch) * int(owner["workspace_rule"]["taps"]) - ), - ) - - -def grid_mult(record: dict[str, Any], arch: str) -> int: - """The over-decomposition factor of an ``occupancy`` record on ``arch`` (1 under every other grid rule). - - ``grid_rule["grid_mult"]`` maps each architecture of the record to the - factor its grid multiplies the co-resident CTA count by; the factor per - architecture is the measured trade of tail latency behind a resident - side kernel against partial-row traffic. - """ - - rule = record["grid_rule"] - if rule["kind"] != "occupancy": - return 1 - factor = int(rule["grid_mult"][arch]) - if factor < 1: - raise ValueError( - f"module {record['kernel']}: grid_mult[{arch!r}] must be positive, got {factor}" - ) - return factor - - -def grid_per_sm(record: dict[str, Any], arch: str) -> int: - """CTAs per SM of a ``co_resident`` or ``occupancy`` record's grid on ``arch``: - ``grid_mult x ctas_per_sm`` (the grid is this times the SM count).""" - - return grid_mult(record, arch) * int(record["grid_rule"]["ctas_per_sm"]) - - -def _ceil_div(numerator: int, denominator: int) -> int: - if denominator <= 0: - raise ValueError("denominator must be positive") - return -(-int(numerator) // int(denominator)) - - -# The roles whose programs plan their channel tiles with :func:`channel_split` -# (a reduce program runs on its backward record's plan): each has its own rule. -SPLIT_ROLES = ("fwd", "bwd") - - -def channel_split( - C: int, - *, - geometry: dict[str, int], - role: str | None = None, - legacy: bool = False, -) -> tuple[int, int, int]: - """``(ct_vec, n_sub, n_ch_tiles)`` of ``C`` channels under a record's build geometry and role. - - ``ct_vec`` channel vectors per channel tile (one thread each), ``n_sub = - threads // ct_vec`` token sub-tiles per CTA and ``n_ch_tiles = ceil(C / vec - / ct_vec)`` channel tiles. ``role`` is the record's, ``"fwd"`` or - ``"bwd"`` (a ``ValueError`` without one unless ``legacy``): the two rules - agree except in the fallback below. With ``cv = C / vec`` channel - vectors: ``ct_vec = cv`` when ``cv <= threads``; above that, the multiple - of 32 dividing ``cv`` (at most ``threads``) with the highest lane use - ``ct_vec x (threads // ct_vec)``, the larger ``ct_vec`` on a tie. So C = - 1280 (160 vectors) runs five 32-vector tiles with four token sub-tiles - each, (32, 4, 5), and C = 1536 (64, 2, 3) -- every lane busy -- instead of - two 128-vector tiles at 62 % / 75 % lane use; C <= 1024 and 2048 / 4096 / - 8192 are unchanged. When no multiple of 32 divides ``cv``, the role's - fallback picks the ``ct_vec`` of the highest lane use over the launch, - ``cv / (n_ch_tiles x ct_vec) x (ct_vec x (threads // ct_vec) / threads)`` - rounded to 9 decimals (the larger on a tie), among every multiple of 32 at - most ``threads`` (the last tile partial) and, for the forward, every - divisor of ``cv`` from 40 to ``threads`` (whole tiles). At 128 x 8 the - backward runs C = 1160 as (32, 4, 5), 1344 as (64, 2, 3) and 1600 as (32, - 4, 7), the forward 1160 as (32, 4, 5), 1344 as (42, 3, 4) and 1600 as (40, - 3, 5); both run 1920 as (128, 1, 2). The per element arithmetic does not - depend on the tiling, so ``y`` and ``dx`` are bitwise the legacy plan's; - ``dweight``'s fixed summation order follows the tiling, so its bits differ - from the legacy plan's (both are reproducible). ``legacy=True`` is the - rule the package shipped before round 6, ``ct_vec = min(cv, threads)``, - for either role; the KCP path keeps it (its kernels pick the grouped path - from ``ct_vec`` on the device and are measured under this rule). Each - role's rule is kept in sync with its copies (the kernel generator of - that role and the test plan of that kind). - """ - - return _channel_split( - int(C), int(geometry["vec"]), int(geometry["threads"]), role, bool(legacy) - ) - - -@functools.lru_cache(maxsize=4096) -def _channel_split( - C: int, vec: int, threads: int, role: str | None, legacy: bool -) -> tuple[int, int, int]: - """:func:`channel_split` of plain arguments, memoized: a model's channel - counts are few, so every call after a channel count's first reuses its - plan (a backward call plans its row up to twice).""" - - if (role is None and not legacy) or (role is not None and role not in SPLIT_ROLES): - raise ValueError( - f"role must be one of {SPLIT_ROLES} (or legacy=True), got {role!r}" - ) - if C <= 0 or C % vec: - raise ValueError(f"C={C} must be a positive multiple of {vec}") - channel_vectors = C // vec - if legacy or channel_vectors <= threads: - ct_vec = min(channel_vectors, threads) - else: - candidates = [ - width - for width in range(32, threads + 1, 32) - if channel_vectors % width == 0 - ] - ct_vec = ( - max(candidates, key=lambda width: (width * (threads // width), width)) - if candidates - else _fallback_split(channel_vectors, threads, role) - ) - return ct_vec, threads // ct_vec, _ceil_div(channel_vectors, ct_vec) - - -def _fallback_split(channel_vectors: int, threads: int, role: str) -> int: - """``ct_vec`` of :func:`channel_split` above ``threads`` vectors with no 32-multiple divisor. - - The candidates are every multiple of 32 at most ``threads`` (the last tile - partial: its copies and stores stop at channel ``C``) and, for the - forward, every divisor of ``channel_vectors`` from 40 to ``threads`` (whole - tiles); the one of the highest :func:`_lane_use`, rounded to 9 decimals, - wins, the larger on a tie. No tile is narrower than a warp's 32 vectors - and no whole forward tile narrower than 40. - """ - - candidates = set(range(32, threads + 1, 32)) - if role == "fwd": - candidates |= { - width for width in range(40, threads + 1) if channel_vectors % width == 0 - } - return max( - candidates, - key=lambda width: ( - round(_lane_use(channel_vectors, width, threads), 9), - width, - ), - ) - - -def _lane_use(channel_vectors: int, width: int, threads: int) -> float: - """The busy share of a CTA's lanes over ``ceil(channel_vectors / width)`` tiles of ``width``.""" - - return ( - channel_vectors - / (_ceil_div(channel_vectors, width) * width) - * (width * (threads // width) / threads) - ) - - -def plan( - S: int, - C: int, - num_docs: int, - *, - geometry: dict[str, int], - role: str | None = None, - legacy: bool = False, -) -> dict[str, int]: - """The kernel's launch scalars for one problem under a record's build geometry and role. - - ``geometry`` is ``record["grid_rule"]["geometry"]``: ``threads`` per CTA - and ``vec`` BF16 channels per thread; ``role`` is ``record["role"]`` - (``"fwd"`` or ``"bwd"``). The scalars are ``S``, ``C``, ``num_docs``, - :func:`channel_split`'s ``ct_vec`` / ``n_sub`` / ``n_ch_tiles`` (``role`` - and ``legacy`` are passed on) and ``search_steps``, the number of halvings - of the device-side binary search over ``cu_seqlens`` (``max(1, ceil(log2 - num_docs))``). Every value is a pure function of ``(S, C, num_docs)``, the - geometry and the role: the host never reads ``cu_seqlens``, and the kernel - derives its token ranges from ``S`` and the grid. The ``flags`` scalar of - the ``dtype_flags`` programs is not a plan value: the host derives it from - the tensors of the call. - """ - - S, C, num_docs = int(S), int(C), int(num_docs) - if S < 0 or num_docs < 0: - raise ValueError(f"S and num_docs must be non-negative, got {(S, num_docs)}") - ct_vec, n_sub, n_ch_tiles = channel_split( - C, geometry=geometry, role=role, legacy=legacy - ) - return { - "S": S, - "C": C, - "num_docs": num_docs, - "search_steps": max(1, (num_docs - 1).bit_length()), - "ct_vec": ct_vec, - "n_sub": n_sub, - "n_ch_tiles": n_ch_tiles, - } - - -# Whole-rank trim of the backward's single-wave main grid (:func:`main_grid`). -# A channel tile ``ct`` of a ``grid``-CTA launch gets ``ceil((grid - ct) / -# n_ch_tiles)`` CTAs, so when ``n_ch_tiles`` does not divide the grid two pairs -# of neighbouring tiles stream other token ranges at any instant; a row that is -# not 64-byte aligned shares a cache line (and a 32-byte ``dx`` sector) between -# them, which beyond L2 is then fetched twice. Dropping ``grid % n_ch_tiles`` -# CTAs gives every tile the same CTA count. ``TILE_ALIGNED_GRID[arch]``: the -# co-resident grid the rule was measured at (GB300's 304 CTAs; 0: none), on the -# routes of ``TILE_ALIGNED_ROUTES``, from ``TILE_ALIGNED_MIN_S`` tokens, with a -# last tile of at least ``TILE_ALIGNED_MIN_LAST[ct_vec]`` vectors and at most -# ``TILE_ALIGNED_MAX_TILES[ct_vec]`` tiles (a ``ct_vec`` without an entry: -# never). ``TILE_ALIGNED_CLASSES[arch]`` holds the other devices' measured -# classes, ``(co_resident, ct_vec, n_ch_tiles) -> (min S, min last-tile -# vectors)`` on the routes of ``TILE_ALIGNED_CLASS_ROUTES[arch]`` (B200's 296 -# CTAs; R200 none). Every other launch keeps its grid. ``dx`` does not depend -# on the grid; ``dweight`` / ``dbias`` sum other partial rows, bitwise -# reproducible for the trimmed grid. The rule is kept in sync with its copy in -# the backward kernel generator. -TILE_ALIGNED_GRID: dict[str, int] = {"sm_100a": 0, "sm_103a": 304, "sm_107a": 0} -TILE_ALIGNED_ROUTES: tuple[str, ...] = ("bwd",) -TILE_ALIGNED_MIN_S = 98304 -TILE_ALIGNED_MIN_LAST: dict[int, int] = {32: 15, 64: 49} -TILE_ALIGNED_MAX_TILES: dict[int, int] = {32: 27, 64: 9} -TILE_ALIGNED_CLASSES: dict[str, dict[tuple[int, int, int], tuple[int, int]]] = { - "sm_100a": { - (296, 32, 5): (147456, 27), - (296, 32, 7): (65536, 31), - (296, 64, 3): (147456, 59), - }, - "sm_103a": {}, - "sm_107a": {}, -} -TILE_ALIGNED_CLASS_ROUTES: dict[str, tuple[str, ...]] = { - "sm_100a": ("bwd", "bwd_dt", "bwd_fp32g"), - "sm_103a": (), - "sm_107a": (), -} - - -def main_grid( - co_resident: int, - *, - arch: str, - grid_mult: int = 1, - C: int | None = None, - n_ch_tiles: int = 1, - S: int | None = None, - ct_vec: int | None = None, - route: str | None = None, -) -> int: - """The backward main launch's CTA count: ``grid_mult x co_resident``. - - ``co_resident`` is the program's co-resident CTA count on the device - (``ctas_per_sm x SM count``). A single-wave launch (``grid_mult`` 1) of - a row the caller names (``C``, ``S``, the plan's ``n_ch_tiles`` and - ``ct_vec``, the ``route``) is trimmed to whole channel-tile ranks: - :func:`tile_aligned_grid` on a route of ``TILE_ALIGNED_ROUTES`` when - ``co_resident`` is ``TILE_ALIGNED_GRID[arch]``, else - :func:`class_aligned_grid` on a route of - ``TILE_ALIGNED_CLASS_ROUTES[arch]``. A caller that names no row keeps the - untrimmed grid, the most CTAs any row launches. - """ - - if type(grid_mult) is not int or grid_mult < 1: - raise ValueError(f"grid_mult must be a positive int, got {grid_mult!r}") - grid = int(co_resident) * grid_mult - named = None not in (C, S, ct_vec) - if ( - grid_mult == 1 - and co_resident == TILE_ALIGNED_GRID[arch] - and route in TILE_ALIGNED_ROUTES - and named - ): - grid = tile_aligned_grid(grid, C, n_ch_tiles, S=S, ct_vec=ct_vec) - elif grid_mult == 1 and route in TILE_ALIGNED_CLASS_ROUTES[arch] and named: - grid = class_aligned_grid(grid, C, n_ch_tiles, arch=arch, S=S, ct_vec=ct_vec) - return grid - - -def tile_aligned_grid( - grid: int, C: int, n_ch_tiles: int, *, S: int, ct_vec: int, vec: int = 8 -) -> int: - """``grid`` less ``grid % n_ch_tiles`` CTAs when ``C % 32`` (a row of - ``x`` / ``dy`` / ``dx`` is not 64-byte aligned), ``1 < n_ch_tiles <= - TILE_ALIGNED_MAX_TILES[ct_vec]``, ``grid % n_ch_tiles``, ``grid > - n_ch_tiles``, ``S >= TILE_ALIGNED_MIN_S`` and the last channel tile holds - at least ``TILE_ALIGNED_MIN_LAST[ct_vec]`` vectors (``C / vec - - (n_ch_tiles - 1) x ct_vec``); else ``grid``.""" - - last = C // vec - (n_ch_tiles - 1) * ct_vec - if ( - 1 < n_ch_tiles <= TILE_ALIGNED_MAX_TILES.get(ct_vec, 0) - and C % 32 - and grid % n_ch_tiles - and grid > n_ch_tiles - and S >= TILE_ALIGNED_MIN_S - and last >= TILE_ALIGNED_MIN_LAST.get(ct_vec, ct_vec + 1) - ): - return grid - grid % n_ch_tiles - return grid - - -def class_aligned_grid( - grid: int, - C: int, - n_ch_tiles: int, - *, - arch: str, - S: int, - ct_vec: int, - vec: int = 8, -) -> int: - """``grid`` less ``grid % n_ch_tiles`` CTAs when ``C % 32``, ``grid % - n_ch_tiles`` and ``(grid, ct_vec, n_ch_tiles)`` is a class of - ``TILE_ALIGNED_CLASSES[arch]`` whose bounds the row meets (``S`` and the - last tile's ``C / vec - (n_ch_tiles - 1) x ct_vec`` vectors at least the - class's); else ``grid``.""" - - bound = TILE_ALIGNED_CLASSES[arch].get((grid, ct_vec, n_ch_tiles)) - if bound is None or not C % 32 or not grid % n_ch_tiles or bound[0] > S: - return grid - if C // vec - (n_ch_tiles - 1) * ct_vec < bound[1]: - return grid - return grid - grid % n_ch_tiles - - -def row_may_trim( - record: dict[str, Any], arch: str, co_resident: int, factor: int, C: int, S: int -) -> bool: - """Whether :func:`main_grid` may trim ``record``'s launch for the row. - - Every trim needs a single wave (``factor`` 1), ``C % 32``, the record on a - route of the rule and ``S`` at least the rule's smallest token bound for - ``co_resident``; a call failing any of these launches the untrimmed grid - without planning the row's channel tiles (no host cost beyond the - untrimmed grid's).""" - - if factor != 1 or not C % 32: - return False - library = record["library"] - routes = ROUTES[arch] - if ( - co_resident == TILE_ALIGNED_GRID[arch] - and S >= TILE_ALIGNED_MIN_S - and any(library in routes.get(kind, ()) for kind in TILE_ALIGNED_ROUTES) - ): - return True - return any( - grid == co_resident and bound[0] <= S - for (grid, _, _), bound in TILE_ALIGNED_CLASSES[arch].items() - ) and any( - library in routes.get(kind, ()) for kind in TILE_ALIGNED_CLASS_ROUTES[arch] - ) - - -def record_route(record: dict[str, Any], arch: str) -> str: - """The route of ``ROUTES[arch]`` that lists ``record`` (``ValueError`` if none).""" - - return _library_route(record["library"], arch) - - -@functools.cache -def _library_route(library: str, arch: str) -> str: - for kind, names in ROUTES[arch].items(): - if library in names: - return kind - raise ValueError(f"module {library} is on no route of {arch}") - - -@functools.cache -def sm_count(device_index: int) -> int: - """Streaming multiprocessor count of one CUDA device (immutable per process).""" - - import torch - - return int(torch.cuda.get_device_properties(device_index).multi_processor_count) - - -def launch_grid( - record: dict[str, Any], - *, - device_index: int, - C: int | None = None, - S: int | None = None, -) -> tuple[int, int, int]: - """Launch grid of one module on ``device_index`` under its recorded grid rule. - - ``co_resident`` modules launch exactly ``ctas_per_sm x SM count`` CTAs, - the co-resident count (a cooperative KCP backward never exceeds it). - ``occupancy`` modules launch :func:`main_grid` of that count: - ``grid_mult[arch] x ctas_per_sm x SM count`` CTAs normally - (:func:`grid_per_sm`): ``ctas_per_sm`` co-resident CTAs per SM and - ``grid_mult`` times finer token streams, so a CTA that finds no free slot - beside another stream's kernel trails by a short stream in a later wave. - Given the call's ``C`` and ``S``, a single-wave launch of the rows - :func:`main_grid` names (its route, the record's plan) drops CTAs to whole - channel-tile ranks (the row is planned only where :func:`row_may_trim` - holds). Either grid is a pure function of ``(S, C, device)`` and the - route, so is the deterministic ``dweight`` order, and the untrimmed grid - (``C`` or ``S`` not given) is the most CTAs any call launches, so one - workspace serves every token count. - ``channels_per_cta`` modules (the reduce programs) launch ``ceil(C / - channels_per_cta)`` CTAs for the call's ``C`` channels. - """ - - rule = record["grid_rule"] - kind = rule["kind"] - if kind in ("co_resident", "occupancy"): - arch = device_arch(device_index) - grid_x = grid_per_sm(record, arch) * sm_count(device_index) - if kind == "occupancy" and C is not None and S is not None and C % 32: - co_resident = int(rule["ctas_per_sm"]) * sm_count(device_index) - factor = grid_mult(record, arch) - if row_may_trim(record, arch, co_resident, factor, int(C), int(S)): - ct_vec, _, n_ch_tiles = channel_split( - C, geometry=rule["geometry"], role=record["role"] - ) - grid_x = main_grid( - co_resident, - arch=arch, - grid_mult=factor, - C=int(C), - n_ch_tiles=n_ch_tiles, - S=int(S), - ct_vec=ct_vec, - route=record_route(record, arch), - ) - elif kind == "channels_per_cta": - if C is None: - raise ValueError( - f"the {kind} grid rule of module {record['kernel']} needs the channel count C" - ) - grid_x = _ceil_div(C, int(rule["channels_per_cta"])) - else: - raise ValueError( - f"unknown Cake causal conv1d training grid rule: {kind!r} (known: {GRID_KINDS})" - ) - return (max(1, grid_x), 1, 1) - - -def partial_rows(grid_x: int, n_ch_tiles: int) -> int: - """Rows of the backward ``partial`` workspace for a grid: ``ceil(grid_x / n_ch_tiles)``. - - CTA ``b`` of the grid owns channel tile ``b % n_ch_tiles`` and partial row - ``b // n_ch_tiles``, so tile ``t`` holds ``ceil((grid_x - t) / n_ch_tiles)`` - written rows; the reduce program derives exactly those from ``main_ctas`` - (= ``grid_x``) and ``n_ch_tiles`` and sums nothing else. - """ - - return _ceil_div(grid_x, n_ch_tiles) - - -def partial_bytes(record: dict[str, Any], C: int, rows: int) -> int: - """Bytes of ``rows`` FP32 ``[C, taps]`` partial rows of the module's workspace rule. - - ``taps`` is 4 for the dweight partials; a bias program's rule says 5: its - kernel lays the ``[rows, C, 4]`` tap slab out first and the ``[rows, C]`` - dbias slab after it (the host view is capacity only). - """ - - rule = record["workspace_rule"] - if rule is None or rule["kind"] != "partial_rows": - raise ValueError(f"module {record['kernel_symbol']} owns no backward workspace") - return 4 * int(rule["taps"]) * int(C) * int(rows) - - -def workspace_layout( - record: dict[str, Any], C: int, *, device_index: int -) -> dict[str, int]: - """Byte layout of the caller-owned backward workspace of one module for ``C`` channels. - - The workspace is one CUDA ``uint8`` buffer holding the FP32 ``[rows, C, taps]`` - ``dweight`` (and, ``taps`` 5, dbias) partials, ``rows = partial_rows(grid, - n_ch_tiles)`` for the device's launch grid (``grid_mult x ctas_per_sm x SM - count``: the bytes grow with the record's ``grid_mult``); it depends on - ``C`` and the device only, so one buffer serves every token count. The - kernel fully writes the rows its reduce program sums, so nothing is zeroed - between launches or CUDA Graph replays. - """ - - grid_x = launch_grid(record, device_index=device_index)[0] - _, _, n_ch_tiles = channel_split( - C, geometry=record["grid_rule"]["geometry"], role=record["role"] - ) - rows = partial_rows(grid_x, n_ch_tiles) - return { - "grid": grid_x, - "rows": rows, - "taps": int(record["workspace_rule"]["taps"]), - "total_bytes": partial_bytes(record, C, rows), - } - - -def counter_words(record: dict[str, Any], n_ch_tiles: int) -> int: - """Int32 words of a forward record's ``work`` counters for a launch of ``n_ch_tiles`` channel tiles. - - A ``tile_counters`` rule holds ``words_per_tile`` words per channel tile - (the dynamic forward's next-chunk counter and done-CTA counter); a record - without a workspace rule takes no ``work`` (0 words). - """ - - rule = record["workspace_rule"] - if rule is None: - return 0 - if rule["kind"] != "tile_counters": - raise ValueError( - f"module {record['kernel']}: workspace rule {rule['kind']!r} holds no forward counters" - ) - if n_ch_tiles <= 0: - raise ValueError(f"n_ch_tiles must be positive, got {n_ch_tiles}") - return int(rule["words_per_tile"]) * int(n_ch_tiles) - - -def forward_kinds(arch: str) -> tuple[str, ...]: - """The forward routes ``arch`` exports, in :data:`KINDS` order.""" - - return tuple( - kind - for kind in KINDS - if route_role(kind) == "fwd" and ROUTES.get(arch, {}).get(kind) - ) - - -def forward_workspace_words(arch: str, C: int) -> int: - """Int32 words of the forward workspace for ``C`` channels on ``arch``. - - The largest :func:`counter_words` over every forward program of every - forward route of the architecture (each planned with its own geometry and - role), so one buffer serves every forward route; it depends on ``C`` only - (never on ``S`` or the device's SM count). 0 when no forward program of - the architecture takes ``work``. - """ - - words = 0 - for kind in forward_kinds(arch): - for name in route_modules(arch, kind): - record = MODULES[name] - if record["workspace_rule"] is None: - continue - _, _, n_ch_tiles = channel_split( - C, geometry=record["grid_rule"]["geometry"], role=record["role"] - ) - words = max(words, counter_words(record, n_ch_tiles)) - return words - - -@functools.cache -def device_arch(device_index: int) -> str: - """Exported architecture of one CUDA device (``RuntimeError`` when it is not an exported target).""" - - import torch - - capability = tuple(torch.cuda.get_device_capability(device_index)) - arch = arch_for_capability(capability) - if arch is None: - raise RuntimeError( - f"device {device_index} (capability {capability}) is not an exported Cake causal conv1d training target" - ) - return arch - - -def library_name(name: str, arch: str) -> str: - """``build.json`` library of module record ``name`` built for ``arch``: ``_``.""" - - record = MODULES[name] - if arch not in record["arches"]: - raise RuntimeError( - f"Cake causal conv1d training module {name} is not delivered for {arch}; " - f"its architectures: {record['arches']}" - ) - return f"{record['library']}_{arch}" - - -@functools.cache -def load(name: str, arch: str) -> Any: - """The loaded tvm-ffi module of record ``name`` built for ``arch``. - - The prebuilt ``aot/`` library when a builder compiled it, else FlashInfer's - JIT build of the same sources and flags (``cake_build.load_library``). - """ - - return cake_build.load_library(library_name(name, arch)) - - -def specs_for_arch(arch: str) -> list[Any]: - """JIT specifications of every ``arch`` library without a prebuilt file (for - ahead-of-time warm-up through FlashInfer's ``build_jit_specs``).""" - - return cake_build.jit_specs(arch) - - -_GRID_AXES = {"grid_x": 0, "grid_y": 1, "grid_z": 2} - - -@functools.cache -def _arg_sources(name: str) -> tuple[tuple[bool, Any], ...]: - """``(grid, source)`` per entry of ``MODULES[name]["arg_plan"]``, in - order: a grid slot (``source`` is the axis, 0 / 1 / 2 for ``grid_x`` / - ``grid_y`` / ``grid_z``) or the logical name the kernel parameter binds - (the record's ``logical`` map); the registry is constant, so - :func:`run_modules` reads it once per module.""" - - record = MODULES[name] - logical = record["logical"] - return tuple( - (True, _GRID_AXES[key]) if kind == "grid" else (False, logical[key]) - for kind, key in record["arg_plan"] - ) - - -def _missing_argument(name: str, values: dict[str, Any]) -> ValueError: - """The error of a launch of module ``name`` whose ``values`` lack a - logical name it binds (the first such name, with its kernel parameter).""" - - record = MODULES[name] - logical, key = next( - (record["logical"][key], key) - for kind, key in record["arg_plan"] - if kind != "grid" and record["logical"][key] not in values - ) - return ValueError(f"module {name} needs {logical!r} (kernel parameter {key!r})") - - -@functools.cache -def _ffi_entry(name: str, device_index: int) -> Any: - """The FFI entry of module ``name`` loaded for CUDA device - ``device_index``'s architecture (:func:`load`), looked up once.""" - - module = load(name, device_arch(device_index)) - return getattr(module, MODULES[name]["ffi_entry"]) - - -@functools.cache -def _raw_stream_getter() -> Any: - """``torch._C._cuda_getCurrentRawStream``: the handle of a device's - current torch stream as an integer, without constructing a ``Stream`` - object; or the ``torch.cuda`` form on a torch build without it. Bound on - first use: this module imports torch lazily.""" - - import torch - - getter = getattr(torch._C, "_cuda_getCurrentRawStream", None) - if getter is not None: - return getter - - def fallback(device_index: int) -> int: - return torch.cuda.current_stream(device_index).cuda_stream - - return fallback - - -_FFI_DEVICES: dict[int, Any] = {} - - -def _stream_context(device_index: int) -> Any: - """The tvm-ffi environment-stream context of torch's current stream on - CUDA device ``device_index`` (a CUDA graph's capture stream during - capture): the bindings launch on that stream. The FFI device handle is - kept per index, so entering the context costs a fraction of - ``tvm_ffi.use_torch_stream()``, which builds a ``torch.cuda.Stream`` and - parses a device string on every entry.""" - - device = _FFI_DEVICES.get(device_index) - if device is None: - device = _FFI_DEVICES[device_index] = tvm_ffi.device(f"cuda:{device_index}") - return tvm_ffi.use_raw_stream(device, _raw_stream_getter()(device_index)) - - -def run_modules( - launches: Sequence[tuple[str, tuple[int, int, int] | None]], - values: dict[str, Any], - *, - device_index: int, -) -> None: - """Launch exported modules in order, back to back, with one set of logical - argument ``values``. - - ``launches`` holds ``(name, grid)`` pairs, ``grid`` as in - :func:`run_module`. All of them are issued inside one stream context of - torch's current stream on ``device_index``, each bound (grid and - arguments) right before it is issued, so a programmatically dependent - launch (the backward's reduce program) reaches the stream right behind - the program it depends on, while that one runs, and the first launch - waits for no later one's binding. ``ValueError`` names the first logical - name a module binds that ``values`` lacks. - """ - - with _stream_context(device_index): - for name, grid in launches: - if grid is None: - grid = launch_grid( - MODULES[name], - device_index=device_index, - C=values.get("C"), - S=values.get("S"), - ) - sources = _arg_sources(name) - try: - args = [ - grid[source] if is_grid else values[source] - for is_grid, source in sources - ] - except KeyError: - raise _missing_argument(name, values) from None - _ffi_entry(name, device_index)(*args) - - -def run_module( - name: str, - values: dict[str, Any], - *, - device_index: int, - grid: tuple[int, int, int] | None = None, -) -> None: - """Launch one exported module with logical argument ``values``. - - ``values`` maps the logical names in ``LOGICAL_TENSORS`` / - ``LOGICAL_SCALARS`` to tensors and Python scalars; the module record maps - its kernel parameter names onto that vocabulary. The launch grid is - ``grid`` when the caller planned it (the module's :func:`launch_grid` for - the call, as the public entry points pass it), else the module's - :func:`launch_grid` on ``device_index`` for the call's ``C`` and ``S``, - and the kernel runs on the current torch stream; a - module whose launch record enables programmatic dependent launch - (``use_pdl``) overlaps its launch with the preceding kernel of that stream - and waits for it in-kernel. - """ - - run_modules(((name, grid),), values, device_index=device_index) - - -__all__ = [ - "ARCHES", - "ARCH_BY_CAPABILITY", - "CHANNEL_MULTIPLE", - "FLAG_B_BF16", - "FLAG_DB_BF16", - "FLAG_DW_BF16", - "FLAG_W_BF16", - "GRID_KINDS", - "INT32_MAX", - "KINDS", - "LOGICAL_SCALARS", - "LOGICAL_TENSORS", - "MODULES", - "PLAN_SCALARS", - "REQUIRED_KINDS", - "ROLES", - "ROUTES", - "SELECT_KINDS", - "TILE_ALIGNED_CLASSES", - "TILE_ALIGNED_CLASS_ROUTES", - "TILE_ALIGNED_GRID", - "TILE_ALIGNED_MAX_TILES", - "TILE_ALIGNED_MIN_LAST", - "TILE_ALIGNED_MIN_S", - "TILE_ALIGNED_ROUTES", - "WIDTH", - "WORKSPACE_KINDS", - "arch_for_capability", - "backward_kinds", - "backward_route", - "channel_split", - "class_aligned_grid", - "counter_words", - "device_arch", - "forward_kinds", - "forward_route", - "forward_workspace_words", - "grid_mult", - "grid_per_sm", - "launch_grid", - "library_name", - "load", - "main_grid", - "partial_bytes", - "partial_rows", - "plan", - "record_route", - "route_applies", - "route_modules", - "route_role", - "row_may_trim", - "run_module", - "run_modules", - "select_rule", - "selected_modules", - "sm_count", - "specs_for_arch", - "tile_aligned_grid", - "workspace_layout", - "workspace_owners", - "workspace_record", -] diff --git a/flashinfer/cake_causal_conv1d_train_backend/kcp.py b/flashinfer/cake_causal_conv1d_train_backend/kcp.py index 849a2844f91..902fda1b928 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/kcp.py +++ b/flashinfer/cake_causal_conv1d_train_backend/kcp.py @@ -1,1857 +1,9 @@ -""" -Copyright (c) 2026 by FlashInfer team. +# Copyright (c) 2026 by FlashInfer team. Licensed under Apache-2.0. +"""Compatibility import for flashinfer.cake_causal_conv1d_det_train_backend.kcp (the same module object).""" -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at +from importlib import import_module +import sys - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. - -Context-parallel (KCP) halo forward / backward of the Cake packed width-4 -depthwise causal convolution + SiLU (FlashInfer issue #6919): the rank's -fragments of the packed documents are convolved with the three document tokens -before each continuing fragment, received from the ranks that hold them by one -``all_to_all_single`` per direction, and the backward sends the first rows' -derivative back to the predecessor fragments, which fold it into their last -rows' ``dx`` before the one BF16 rounding. ``cake_causal_conv1d_kcp(x_local, -weight, ctx)`` is the autograd entry; the stable, ``@flashinfer_api``-decorated -facade is :mod:`flashinfer.cake_causal_conv1d_train`. This module imports -only ``os``, torch, torch.distributed, tvm_ffi and its own package. - -Registry contract (``kcp_jit``, the KCP records beside the zero-history -registry ``cake_jit``, whose tables stay untouched): kinds ``kcp_fwd`` (one -record per arch: co-resident grid, 4 CTAs per SM, 128 threads x 8 channels), -``kcp_bwd`` (the fallback record and the ``_st8`` record selected by -``channels_at_least`` = the zero-history rule; ``workspace_rule`` -``partial_rows`` with 4 taps) and ``kcp_bwd_head`` (one record; its grid is -``max(1, n_frag x ceil(C / (threads x vec)))`` computed here and launched -through the record's ``arg_plan`` / ``logical`` / ``ffi_entry`` and ``load``); -logical tensor names are ``cake_jit.LOGICAL_TENSORS`` plus ``hist``, -``hist_row_off``, ``dpre_tail``, ``tail_row_off``, ``dpre_head``, -``head_row_off`` (``x_map`` / ``dy_map`` map to ``x`` / ``dy``); the scalars are -``cake_jit.PLAN_SCALARS`` (the KCP programs take no ``flags``). The KCP path -keeps the package's pre-round-6 contract: FP32 taps only (``_weight``'s default -dtype set), no bias, and the legacy channel tile rule ``cake_jit.plan(..., -legacy=True)`` (``ct_vec = min(C / 8, 128)``; its kernels pick the grouped path -from ``ct_vec`` on the device and were validated and measured under that rule), -so C = 1280 still plans (128, 1, 2) here while the training routes tile it as -(32, 4, 5). - -Numerics against the zero-history kernels of this package: the forward runs -the zero-history forward's instructions with history rows, so ``y`` is bitwise -equal to the no-CP ``y`` of the same documents. The backward takes the SiLU -derivative on the FP32 preactivation (the issue's requirement), whereas the -shipped zero-history backward differentiates the BF16-rounded preactivation, -so ``dx`` is not the no-CP ``dx`` bit for bit: the two differ by that -rounding's effect on ``dpre`` plus one BF16 output rounding. Measured on -GB300 (sm_103a) over the single-GPU test rows (``S x C`` up to 12000 x 8192, -random BF16 inputs, FP32 taps of scale 0.5): max ``|dx - dx_noCP|`` 0.03125 -(one BF16 ulp at magnitude 4-8), 0 elements beyond ``1e-2 + 1e-2 |dx_noCP|``, -and the same maximum between the two policies' FP32 oracles with no kernel -involved; each kernel stays within one BF16 ulp of its own oracle. ``dweight`` -is the rank's fixed-order FP32 partial; it is reproducible across launches and -within ``5e-2 + 1e-4 |dw|`` of an FP64 oracle, not bitwise the no-CP partial. - -Shipped host path (this revision; -``HOST_TRIM_DEFAULT = 1``): the context-derived host work of the eager halo -call -- the validated ``cu`` / ``hist_index`` / row-offset tables and halo -splits, the plan scalars, the grids, the record names and the loaded modules -with their positional argument plans -- is computed once per context and -reused, and the kernels launch through a raw-stream FFI context of the current -stream (no ``torch.cuda.Stream`` object per launch). ``_HOST_CACHE`` is keyed -on the context object's identity (the entry pins the context) and, inside it, -on ``(kind, S, C, device, cu_seqlens_local_cpu)`` -- host integers of the -context and the shapes of the call, never a tensor's contents; it holds at most -``HOST_CACHE_MAX = 64`` contexts (the oldest dropped) and -:func:`host_cache_clear` empties it; ``_FFI_DEVICES`` caches the FFI device per -device index. The operand checks, the gather, the exchange, its wait and the -kernel launches run as before with the same arguments on the same stream, so -no bit changes; a replayed graph (``graph="replay"``) contains no host work. -``_set_test_knobs({"host_trim": 0})`` (with ``CAKE_KCP_TEST_KNOBS=1``) selects -the previous host path for the A/B record. - -Checkpoint recompute without the halo exchange (FlashInfer issues #7815 / -#7818): with ``recompute_handle=`` (:class:`KcpRecomputeHandle`, or any object -with an ``entries`` dict) the original forward exchanges as always and stores -its history pool ``hist`` (the saved set's own tensor, key ``"conv"``) in the -handle; torch's recompute of that forward inside the backward binds it and -launches the same forward kernel on it -- no gather, no collective, no wait -- -so ``y`` and the saved set ``(x_local, weight, hist)`` are bitwise those of the -exchanging recompute. The backward is unchanged (it still exchanges the head -rows). The forward exchange is a collective: every rank of ``ctx.group`` passes -a handle for the same call or none, created at the same checkpoint nesting -level (no check is made; a mismatch hangs or corrupts data). - -Several convolutions of one layer (FlashInfer issue #7815): ``cake_causal_conv1d_kcp_multi(xs, weights, ctx)`` runs -the convolutions of ``n`` tensors that share the context (a layer's q, k and v) with one halo exchange per direction -instead of ``n``: the tensors' send rows travel side by side in one ``all_to_all_single`` and are copied out into -each tensor's own contiguous pool after the wait, and each tensor then runs the single call's kernels with the -single call's arguments, so outputs and gradients (``dweight`` included) are bitwise one single call per tensor. -One recompute handle serves the whole call (its ``"conv"`` entry holds the ``n`` pools). The multi path always -uses the trimmed host path's per-context cache (the ``host_trim`` test knob selects the single call's path only). -""" - -from __future__ import annotations - -import os -from typing import Any, Optional - -import torch -import torch.distributed as dist -import tvm_ffi - -from . import ( - _cu_seqlens, - _device_arch, - _device_index, - _output, - _require_tensor, - _scalars, - _tokens, - _weight, -) -from . import kcp_jit as _loader -from .cake_jit import WIDTH, launch_grid, partial_rows, plan -from .kcp_jit import load, route_modules, run_module, selected_modules - -HALO = WIDTH - 1 # history rows, tail rows and head rows per fragment -HEAD_SLOTS = HALO # send slots per head row (a head token is consumed by at most three predecessor fragments) -KCP_FWD = "kcp_fwd" -KCP_BWD = "kcp_bwd" -KCP_BWD_HEAD = "kcp_bwd_head" -TEST_KNOBS_ENV = "CAKE_KCP_TEST_KNOBS" # the KCP route's private test-hook gate (the variable the KDA host module reads too) -GRAPH_MODES = ( - "off", - "replay", -) # graph replay: keyed CUDA-graph capture / replay of the halo calls; public keyword graph= since an earlier revision -GRAPH_DEFAULT = ( - "off" # shipped: the eager path; the public keyword graph="replay" opts in per call -) - - -def _ceil_div(numerator: int, denominator: int) -> int: - return -(-numerator // denominator) - - -def _table( - ctx: Any, - name: str, - *, - shape: tuple[int, ...], - dtype: torch.dtype, - device: torch.device, -) -> torch.Tensor: - """One of the context's halo tables, checked against the kernels' ABI shape.""" - - table = _require_tensor(ctx.halo_tables[name], f"ctx.halo_tables[{name!r}]") - if tuple(table.shape) != shape or table.dtype != dtype: - raise ValueError( - f"ctx.halo_tables[{name!r}] must be {dtype} {shape}, got {table.dtype} {tuple(table.shape)}" - ) - if table.device != device or not table.is_contiguous(): - raise ValueError( - f"ctx.halo_tables[{name!r}] must be a contiguous tensor on {device}" - ) - return table - - -def _pool( - rows: int, C: int, *, dtype: torch.dtype, device: torch.device -) -> torch.Tensor: - """A fresh row pool; a pool without rows still holds one row so that the kernel has a valid pointer (no table - entry names it).""" - - return torch.empty((max(rows, 1), C), dtype=dtype, device=device) - - -def _exchange( - send: torch.Tensor, - *, - recv_rows: int, - send_splits: tuple[int, ...], - recv_splits: tuple[int, ...], - ctx: Any, -) -> tuple[torch.Tensor, Optional[Any]]: - """One ``all_to_all_single`` of ``send`` (rows in destination-rank order) into a fresh ``[recv_rows, C]`` tensor, - issued asynchronously on the current stream; with one rank the send buffer is the received pool (no collective). - Returns ``(recv, work)``; the caller waits ``work`` at the first consumer.""" - - if int(ctx.world_size) == 1: - return send, None - recv = torch.empty((recv_rows, send.shape[1]), dtype=send.dtype, device=send.device) - work = dist.all_to_all_single( - recv, - send, - output_split_sizes=list(recv_splits), - input_split_sizes=list(send_splits), - group=ctx.group, - async_op=True, - ) - return recv, work - - -def _launch_with_grid( - name: str, values: dict[str, Any], *, grid: tuple[int, int, int], device_index: int -) -> None: - """Launch record ``name`` with logical ``values`` and an explicit grid (the head program's grid depends on the - fragment and channel counts, not on the device); otherwise :func:`cake_jit.run_module`'s argument plan.""" - - record = _loader.MODULES[name] - logical = record["logical"] - grid_values = {"grid_x": grid[0], "grid_y": grid[1], "grid_z": grid[2]} - args = [] - for kind, key in record["arg_plan"]: - if kind == "grid": - args.append(grid_values[key]) - continue - logical_name = logical[key] - if logical_name not in values: - raise ValueError( - f"module {name} needs {logical_name!r} (kernel parameter {key!r})" - ) - args.append(values[logical_name]) - module = load(name, _loader.device_arch(device_index)) - with tvm_ffi.use_torch_stream(): - getattr(module, record["ffi_entry"])(*args) - - -def _wait_ready(ctx: Any) -> None: - """``ctx.wait_ready()`` of a context that has one (a context prepared on another stream orders its table uploads - on the current stream; a duck-typed context without it needs nothing).""" - wait_ready = getattr(ctx, "wait_ready", None) - if wait_ready is not None: - wait_ready() - - -def _context(ctx: Any, *, device: torch.device, T_r: int) -> tuple[torch.Tensor, int]: - """The context's local ``cu_seqlens`` (int32 device ``[n_frag + 1]``) and fragment count, checked against the - tokens the call received.""" - - if int(ctx.T_r) != T_r: - raise ValueError( - f"x_local holds {T_r} tokens but ctx was built for T_r={int(ctx.T_r)}" - ) - if int(ctx.conv_kernel_size) != WIDTH: - raise ValueError( - f"ctx was built for a width-{int(ctx.conv_kernel_size)} convolution; this kernel is width {WIDTH}" - ) - if int(ctx.n_frag) == 0: - # A rank without fragments (T_r == 0): nothing is launched, the collectives get empty splits. - if T_r != 0: - raise ValueError(f"ctx has no fragments but x_local holds {T_r} tokens") - return _require_tensor(ctx.cu_seqlens_local, "ctx.cu_seqlens_local"), 0 - cu, n_frag = _cu_seqlens(ctx.cu_seqlens_local, device=device) - if n_frag != int(ctx.n_frag): - raise ValueError( - f"ctx.cu_seqlens_local names {n_frag} fragments but ctx.n_frag is {int(ctx.n_frag)}" - ) - return cu, n_frag - - -# ---------------------------------------------------------------------------------------------------------------- -# Checkpoint recompute without the halo exchange: the recompute handle (``recompute_handle=``). -# -# The phase comes from ``torch._C._current_graph_task_id()`` (checked once: present and returning an int; without -# it a call with a handle raises): -1 outside the autograd engine (an original forward), the running graph task -# inside a backward (torch's checkpoint recompute, reentrant or not). Rules of a call with a handle: -# * no "conv" entry -> an original forward (also one inside an outer checkpoint's recompute, nested regions): -# gather + exchange as always, then store ``hist`` (no copy) and its metadata; -# * an entry inside the backward -> the recompute: bind the stored ``hist``, validate it against the call, issue -# no collective; the entry stays (``retain_graph`` / a second backward replays it again); -# * an entry outside the backward -> raise (a handle reused across micro-steps or for a second convolution). -# Graph replay (keyword or knob) and a forward inside a CUDA-graph capture are refused with a handle, in the -# autograd Function itself (``_handle_refusals``), so ``CakeCausalConv1dKcpFunction.apply`` refuses like the wrapper. -# The ranks are not checked against each other: the forward exchange is a collective, so every rank of the group -# must take the same branch (all with a handle at the same nesting level, or all without). - -HANDLE_KEY = "conv" # this package's entry in a recompute handle - -_GRAPH_TASK_ID_BINDING = getattr(torch._C, "_current_graph_task_id", None) -_GRAPH_TASK_ID_USABLE: Optional[bool] = None # None = not checked yet - - -class KcpRecomputeHandle: - """The boundary tensors of one checkpointed call, carried from its original forward to torch's recompute of that - forward, so the recompute runs without cross-rank communication. - - Create one handle per convolution call and micro-step, outside the outermost checkpoint region, and pass it as an - argument of the checkpointed function, so the recompute sees the same object:: - - def layer(x, weight, handle): - return cake_causal_conv1d_kcp(x, weight, ctx, recompute_handle=handle) - - y = torch.utils.checkpoint.checkpoint(layer, x, w, KcpRecomputeHandle(), use_reentrant=False) - - ``entries`` maps an operation key to that operation's entry: ``"conv"`` is this package's (the forward's history - pool ``hist``, the very tensor the saved set holds); other packages that accept a handle with an ``entries`` dict - use their own keys, so the handle of a convolution may also serve the recurrence of the same layer. A handle - created inside a checkpoint region is a new, empty one in that region's recompute, which then exchanges again - (correct, not communication-free). The entries live as long as the handle; :meth:`nbytes` is the retained - size. - - The forward halo exchange is a collective. Every rank of ``ctx.group`` must pass a handle for the same call, or - no rank passes one, and every rank creates it at the same checkpoint nesting level. A mismatch pairs one rank's - recompute exchange with another rank's backward exchange: the ranks hang or exchange wrong rows. No check is - made. Never share one handle between two convolution calls (also not in nested regions). Never call - :meth:`clear` on a handle, or reuse it, while a backward that needs its entry is pending.""" - - __slots__ = ("entries",) - - def __init__(self) -> None: - self.entries: dict[str, dict[str, Any]] = {} - - def clear(self) -> None: - """Drop every entry (the retained tensors go with the last reference).""" - self.entries.clear() - - def nbytes(self) -> int: - """Bytes of the distinct tensors the entries hold.""" - seen: dict[int, int] = {} - for entry in self.entries.values(): - for t in entry.values(): - for x in t if isinstance(t, (tuple, list)) else (t,): - if isinstance(x, torch.Tensor): - seen[id(x)] = x.numel() * x.element_size() - return sum(seen.values()) - - -def _kcp_in_backward() -> bool: - """Whether this thread runs inside the autograd engine's backward pass (torch's checkpoint recompute runs there, - an original forward does not). The one place the recompute handle reads the phase. Raises without the torch - binding (checked once).""" - global _GRAPH_TASK_ID_USABLE - if _GRAPH_TASK_ID_USABLE is None: - ok = False - if _GRAPH_TASK_ID_BINDING is not None: - try: - ok = isinstance(_GRAPH_TASK_ID_BINDING(), int) - except Exception: # noqa: BLE001 - an unusable binding disables the option - ok = False - _GRAPH_TASK_ID_USABLE = ok - if not _GRAPH_TASK_ID_USABLE: - raise RuntimeError( - "recompute_handle needs torch._C._current_graph_task_id to tell the recompute from the original " - "forward; this torch does not provide it" - ) - return int(_GRAPH_TASK_ID_BINDING()) != -1 - - -def _handle_refusals(graph: Optional[str], handle: Any) -> None: - """The refusals of a call with a handle, shared by the wrapper and the autograd Function: graph replay (keyword - or test knob), a forward inside a CUDA-graph capture, a handle without an ``entries`` dict.""" - mode = _graph_mode(graph) - if mode != "off": - raise ValueError( - f"recompute_handle needs graph='off' (keyword and test knob), got {mode!r}" - ) - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - raise ValueError( - "recompute_handle is refused inside a CUDA-graph capture (a replay would overwrite the stored rows)" - ) - if not isinstance(getattr(handle, "entries", None), dict): - raise TypeError( - f"recompute_handle must carry an 'entries' dict (KcpRecomputeHandle), got {type(handle).__name__}" - ) - - -def _handle_lookup(handle: Any) -> Optional[dict[str, Any]]: - """The handle's ``"conv"`` entry when this forward is the recompute of the forward that stored it, ``None`` for an - original forward (it exchanges and stores one). An entry found outside the backward raises.""" - entries = getattr(handle, "entries", None) - if not isinstance(entries, dict): - raise TypeError( - f"recompute_handle must carry an 'entries' dict (KcpRecomputeHandle), got {type(handle).__name__}" - ) - in_backward = _kcp_in_backward() - entry = entries.get(HANDLE_KEY) - if entry is None: - return None - if not in_backward: - raise RuntimeError( - "recompute_handle: the handle already holds a 'conv' entry outside the backward -- a handle serves one " - "convolution call of one micro-step; create a new KcpRecomputeHandle per call and micro-step" - ) - return entry - - -def _handle_hist( - entry: dict[str, Any], x_local: torch.Tensor, ctx: Any -) -> torch.Tensor: - """The stored ``hist`` of ``entry``, checked against this recompute: the same context, token count and channel - count, and the stored shape / dtype / device (BF16, ``x_local``'s device).""" - hist = entry.get("hist") - if not isinstance(hist, torch.Tensor): - raise ValueError("recompute_handle: the 'conv' entry holds no hist tensor") - if entry.get("ctx") is not ctx: - raise ValueError( - "recompute_handle: the 'conv' entry belongs to another KcpContext" - ) - S = int(x_local.shape[-2]) if x_local.ndim >= 2 else -1 - C = int(x_local.shape[-1]) if x_local.ndim >= 1 else -1 - if entry.get("key") != (S, C): - raise ValueError( - f"recompute_handle: the 'conv' entry was stored for (T_r, C) = {entry.get('key')}, " - f"this recompute has {(S, C)}" - ) - if ( - tuple(hist.shape) != entry.get("shape") - or hist.dtype != entry.get("dtype") - or hist.device != entry.get("device") - or hist.dtype != torch.bfloat16 - or hist.device != x_local.device - or hist.ndim != 2 - or int(hist.shape[1]) != C - ): - raise ValueError( - f"recompute_handle: the stored hist ({hist.dtype} {tuple(hist.shape)} on {hist.device}) does not match " - f"its entry ({entry.get('dtype')} {entry.get('shape')} on {entry.get('device')}) or this call " - f"(BF16 [n_hist_rows, {C}] on {x_local.device})" - ) - return hist - - -# ---------------------------------------------------------------------------------------------------------------- -# CUDA-graph replay of a call and the test knobs. -# -# The public keyword ``graph="replay"`` of :func:`cake_causal_conv1d_kcp` and of the forward / backward below (shipped -# ``GRAPH_DEFAULT = "off"``) -- or, for the A/B record, the private ``_set_test_knobs({"graph": "replay"})`` (accepted -# only with ``CAKE_KCP_TEST_KNOBS=1``; a keyword contradicting an active knob is a ``ValueError``) -- makes -# :func:`cake_causal_conv1d_kcp_forward` / :func:`cake_causal_conv1d_kcp_backward` run their first call of a -# validity key (kind, the context object, its world size, and the shape / stride / dtype / data pointer of every -# operand and caller-owned output) eagerly, capture the second call into a ``torch.cuda.CUDAGraph`` (replayed once) -# and replay the graph on every later call of the key: the halo gather, the ``all_to_all_single``, its wait and the -# kernels with the same arguments in the same order (bitwise the eager result) and no host work between them. Any -# other key runs eagerly (silent fallback); a rank without tokens is never captured. A replayed call returns the -# graph's static tensors (detached): the next replay of the key overwrites them, so a framework must keep its inputs -# pointer-stable and consume the outputs before the next call of the same layer. This module keeps its own cache -# (no import of the KDA host module); ``GRAPH_STATS`` counts the events and ``LAST_CALL[kind]`` says how the last -# forward / backward ran (keyword or knob). Every captured graph keeps its fresh tensors in a private pool (one per -# cached key; at most GRAPH_CACHE_MAX keys, the oldest dropped first); ``graph_cache_clear()`` releases them and must run -# before ``torch.distributed.destroy_process_group()`` (a graph holding the NCCL collective blocks ``ncclCommDestroy``). - -# Trimmed host path; shipped default: -# 1 = the trimmed host path of the eager halo call (same launches, bitwise; the module docstring's "Shipped host path"), -# 0 = the previous host path, selectable through _set_test_knobs for the A/B record. -HOST_TRIMS = (0, 1) -HOST_TRIM_DEFAULT = ( - 1 # shipped; 0 = the previous path, via _set_test_knobs({"host_trim": 0}) only +sys.modules[__name__] = import_module( + "flashinfer.cake_causal_conv1d_det_train_backend.kcp" ) -_TEST_KNOBS: dict[str, Any] = {"graph": "off", "host_trim": HOST_TRIM_DEFAULT} -GRAPH_STATS: dict[str, int] = {"eager": 0, "capture": 0, "replay": 0, "fallback": 0} -GRAPH_CACHE_MAX = 64 -LAST_CALL: dict[str, dict[str, Any]] = {} - -# -- the trimmed host path (host_trim=1): the per-context host state of the eager halo calls. ``_HOST_CACHE[id(ctx)]`` holds the -# context (pinned, so its id cannot be reused) and one entry per ``(kind, S, C, device, cu_seqlens_local_cpu)``: the -# validated context tables and tensors (``cu``, ``hist_row_off``, ...), the plan scalars, the record names, the loaded -# modules with their positional argument plans and grids. Everything in an entry derives from host integers of the -# context and the shapes of the call; no tensor content is read. At most HOST_CACHE_MAX contexts (the oldest dropped). -HOST_CACHE_MAX = 64 -_HOST_CACHE: dict[int, tuple[Any, dict[tuple, dict[str, Any]]]] = {} -_FFI_DEVICES: dict[int, Any] = {} - - -def _host_cache_entry(ctx: Any, key: tuple) -> dict[str, Any]: - """The (possibly empty, to be filled) host-cache entry of ``ctx`` under ``key``.""" - hit = _HOST_CACHE.get(id(ctx)) - if hit is None or hit[0] is not ctx: - if len(_HOST_CACHE) >= HOST_CACHE_MAX: - del _HOST_CACHE[next(iter(_HOST_CACHE))] - hit = _HOST_CACHE[id(ctx)] = (ctx, {}) - entry = hit[1].get(key) - if entry is None: - entry = hit[1][key] = {} - return entry - - -def host_cache_clear() -> int: - """Drop the trimmed host path's per-context cache; returns the number of contexts dropped.""" - n = len(_HOST_CACHE) - _HOST_CACHE.clear() - return n - - -def _ffi_device(device_index: int) -> Any: - dev = _FFI_DEVICES.get(int(device_index)) - if dev is None: - dev = _FFI_DEVICES[int(device_index)] = tvm_ffi.device( - f"cuda:{int(device_index)}" - ) - return dev - - -def _raw_stream(device_index: int) -> int: - get = getattr(torch._C, "_cuda_getCurrentRawStream", None) - if get is not None: - return int(get(int(device_index))) - return int(torch.cuda.current_stream(int(device_index)).cuda_stream) - - -def _launch_plan( - name: str, grid: tuple[int, int, int], device_index: int -) -> tuple[Any, str, tuple[tuple[str | None, int], ...]]: - """``(module, ffi entry, plan)`` of record ``name`` for the trimmed launch: ``plan`` lists ``(None, grid value)`` - per grid argument and ``(logical name, 0)`` per kernel parameter, in ``run`` order (the same arguments - :func:`run_module` / :func:`_launch_with_grid` build per call).""" - record = _loader.MODULES[name] - logical = record["logical"] - grid_values = { - "grid_x": int(grid[0]), - "grid_y": int(grid[1]), - "grid_z": int(grid[2]), - } - plan: list[tuple[str | None, int]] = [] - for kind, key in record["arg_plan"]: - if kind == "grid": - plan.append((None, grid_values[key])) - else: - if key not in logical: - raise ValueError( - f"module {name}: kernel parameter {key!r} has no logical name" - ) - plan.append((logical[key], 0)) - module = load(name, _loader.device_arch(device_index)) - return module, record["ffi_entry"], tuple(plan) - - -def _launch_trim( - info: tuple[Any, str, tuple[tuple[str | None, int], ...]], - values: dict[str, Any], - device_index: int, -) -> None: - """Launch a cached :func:`_launch_plan` with the per-call ``values`` on the device's current torch stream (raw-stream - FFI context: the same stream ``tvm_ffi.use_torch_stream`` would set, without building a torch.cuda.Stream).""" - module, entry, plan = info - args = [] - for logical_name, grid_value in plan: - if logical_name is None: - args.append(grid_value) - else: - try: - args.append(values[logical_name]) - except KeyError: - raise ValueError( - f"the launch needs logical argument {logical_name!r}" - ) from None - with tvm_ffi.use_raw_stream(_ffi_device(device_index), _raw_stream(device_index)): - getattr(module, entry)(*args) - - -def _set_test_knobs(knobs: Optional[dict]) -> dict: - """Set the module's private test knobs (``graph``: ``"off"`` shipped | ``"replay"``; ``host_trim``: ``1`` shipped = - the trimmed host path of the eager call | ``0`` the previous host path); a knob - the dict does not name keeps its shipped default; ``None`` restores the defaults. Refused without - ``CAKE_KCP_TEST_KNOBS=1`` in the environment. Returns the active knobs.""" - if knobs is None: - _TEST_KNOBS["graph"] = "off" - _TEST_KNOBS["host_trim"] = HOST_TRIM_DEFAULT - return dict(_TEST_KNOBS) - if os.environ.get(TEST_KNOBS_ENV) != "1": - raise ValueError( - f"_set_test_knobs is a private test hook: it is accepted only with {TEST_KNOBS_ENV}=1 in the environment" - ) - if not isinstance(knobs, dict): - raise TypeError(f"_set_test_knobs expects a dict, got {type(knobs).__name__}") - unknown = sorted(set(knobs) - {"graph", "host_trim"}) - if unknown: - raise ValueError( - f"_set_test_knobs: unknown knob(s) {unknown}; known: ['graph', 'host_trim']" - ) - graph = str(knobs.get("graph", "off")) - if graph not in GRAPH_MODES: - raise ValueError( - f"_set_test_knobs: graph must be one of {GRAPH_MODES}, got {graph!r}" - ) - host_trim = knobs.get("host_trim", HOST_TRIM_DEFAULT) - if isinstance(host_trim, bool) or host_trim not in HOST_TRIMS: - raise ValueError( - f"_set_test_knobs: host_trim must be one of {HOST_TRIMS}, got {host_trim!r}" - ) - _TEST_KNOBS["graph"] = graph - _TEST_KNOBS["host_trim"] = int(host_trim) - return dict(_TEST_KNOBS) - - -def _graph_mode(graph: Optional[str]) -> str: - """The effective graph mode of a call: the public keyword when given, else the module's test knob (shipped - ``"off"``); validated against ``GRAPH_MODES``; a keyword contradicting an active knob is a ``ValueError``.""" - knob = _TEST_KNOBS["graph"] - if graph is None: - return knob - mode = str(graph) - if mode not in GRAPH_MODES: - raise ValueError(f"graph must be one of {GRAPH_MODES}, got {graph!r}") - if knob != "off" and mode != knob: - raise ValueError( - f"graph given twice with different values: keyword {mode!r}, _set_test_knobs {knob!r}" - ) - return mode - - -class _GraphEntry: - """One validity key: its graph (after the capture) and the static results.""" - - __slots__ = ("key", "pins", "seen", "graph", "result", "disabled", "digest") - - def __init__(self, key: tuple, pins: tuple) -> None: - self.key, self.pins = ( - key, - pins, - ) # the pins keep the keyed objects alive so their ids cannot be reused - self.seen = 0 - self.graph: Optional[torch.cuda.CUDAGraph] = None - self.result: Optional[tuple] = None - self.disabled: Optional[str] = None - self.digest = f"{key[0]}:{hash(key) & 0xFFFFFFFFFFFF:012x}" - - -_GRAPHS: dict[tuple, _GraphEntry] = {} - - -def _sig(t: Any) -> tuple: - return ( - (t.data_ptr(), tuple(t.shape), tuple(t.stride()), str(t.dtype), t.device.index) - if isinstance(t, torch.Tensor) - else ("none",) - ) - - -def _graph_entry( - kind: str, ctx: Any, tensors: tuple, mode: Optional[str] = None -) -> Optional[_GraphEntry]: - """The cache entry of this call's validity key, or ``None`` when the call is never captured (mode off, no tokens); - ``mode`` = the effective graph mode (:func:`_graph_mode`), ``None`` = the knob.""" - if ( - (_TEST_KNOBS["graph"] if mode is None else mode) == "off" - or int(getattr(ctx, "T_r", 0)) == 0 - or int(getattr(ctx, "n_frag", 0)) == 0 - ): - return None - key = (kind, id(ctx), int(ctx.world_size), tuple(_sig(t) for t in tensors)) - e = _GRAPHS.get(key) - if e is None: - while len(_GRAPHS) >= GRAPH_CACHE_MAX: - del _GRAPHS[next(iter(_GRAPHS))] - e = _GRAPHS[key] = _GraphEntry(key, (ctx,)) - return e - - -def graph_cache_clear() -> int: - """Drop every cached graph with its static buffers and private memory pool; returns the number of entries dropped. - - A graph that captured a NCCL collective (the halo ``all_to_all_single`` here) must be destroyed before the process group - it used is destroyed (``ncclCommDestroy`` waits for graph-captured work; ``torch.distributed.destroy_process_group()`` - hangs while such a graph is alive). Call this before destroying the process group when ``graph="replay"`` (keyword or test knob) was used, and - whenever the operands of a key are gone for good to release its pool. The knob and ``GRAPH_STATS`` are left as they are.""" - n = len(_GRAPHS) - _GRAPHS.clear() - return n - - -def _graph_run( - kind: str, entry: Optional[_GraphEntry], body: Any, mode: Optional[str] = None -) -> tuple: - """Run ``body()`` under the graph mode (keyword or knob): replay the key's graph, or capture this call (the key's second), or run eagerly - (the key's first call, a key whose capture failed, or a call that is never captured). Returns the results.""" - if entry is not None and entry.graph is not None: - entry.graph.replay() - GRAPH_STATS["replay"] += 1 - LAST_CALL[kind] = { - "graph": "replay", - "graph_event": "replay", - "graph_key": entry.digest, - "graph_fallback": None, - } - return entry.result - if entry is not None and entry.seen > 0 and entry.disabled is None: - g = torch.cuda.CUDAGraph() - try: - with torch.cuda.graph(g, capture_error_mode="thread_local"): - result = body() - except Exception as exc: # noqa: BLE001 - the failure is recorded; the call runs eagerly below - entry.disabled = f"{type(exc).__name__}: {str(exc)[:200]}" - GRAPH_STATS["fallback"] += 1 - else: - g.replay() - entry.graph, entry.result = g, result - entry.seen += 1 - GRAPH_STATS["capture"] += 1 - LAST_CALL[kind] = { - "graph": "replay", - "graph_event": "capture", - "graph_key": entry.digest, - "graph_fallback": None, - } - return result - result = body() - GRAPH_STATS["eager"] += 1 - if entry is not None: - entry.seen += 1 - LAST_CALL[kind] = { - "graph": _TEST_KNOBS["graph"] if mode is None else mode, - "graph_event": "eager", - "graph_key": None if entry is None else entry.digest, - "graph_fallback": None if entry is None else entry.disabled, - } - return result - - -def cake_causal_conv1d_kcp_forward( - x_local: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - *, - out: Optional[torch.Tensor] = None, - graph: Optional[str] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r"""The KCP forward: ``y`` of the rank's fragments, and the history pool ``hist`` it used (saved for the backward). - - ``x_local`` BF16 ``[T_r, C]`` (or ``[1, T_r, C]``) contiguous, the rank's fragments back to back in the order of - ``ctx.cu_seqlens_local``; ``weight`` FP32 ``[C, 4]``; ``ctx`` a ``KcpContext`` (its halo plan ``ctx.halo``, the - device tables ``ctx.halo_tables`` and the process group ``ctx.group``). Every rank of the group calls this - function once per layer, in the same order: the forward halo exchange is a collective (a rank without tokens - joins it with empty splits and returns empty outputs). ``hist`` is a fresh ``[n_hist_rows, C]`` BF16 tensor - (``[1, C]`` when the rank receives no row). ``graph`` (``None`` = the shipped ``GRAPH_DEFAULT`` ``"off"`` | - ``"replay"``): keyed CUDA-graph capture / replay of this call (the graph-replay option; see the section comment - above: pointer-stable operands, private pools, eager fallback, :func:`graph_cache_clear`). - """ - - mode = _graph_mode(graph) - if mode == "off": - return _kcp_forward(x_local, weight, ctx, out=out) - entry = _graph_entry("fwd", ctx, (x_local, weight, out), mode) - y, hist = _graph_run( - "fwd", entry, lambda: _kcp_forward(x_local, weight, ctx, out=out), mode - ) - return (y if out is not None else y.detach()), hist - - -def _kcp_forward( - x_local: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - *, - out: Optional[torch.Tensor] = None, - _hist: Optional[torch.Tensor] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """The forward's host flow (eagerly or inside a capture): see :func:`cake_causal_conv1d_kcp_forward`. ``_hist``: - the history pool a recompute handle stored (checked by :func:`_handle_hist`): no gather, no exchange, no wait; - the kernel launch is the same.""" - - x = _require_tensor(x_local, "x_local") - device = x.device - # tables of a context prepared on another stream: order them on this call's stream, on every call (the trimmed - # path's host cache is not per stream; after the first call on a stream this is one dict lookup) - _wait_ready(ctx) - if _TEST_KNOBS["host_trim"] == 1: - return _kcp_forward_trim(x, weight, ctx, out=out, device=device, _hist=_hist) - arch = _device_arch(device) - S, C, x_flat = _tokens(x, "x_local", device=device) - w = _weight(weight, C, device=device) - cu, n_frag = _context(ctx, device=device, T_r=S) - halo = ctx.halo - y = _output(out, "out", shape=tuple(x.shape), dtype=torch.bfloat16, device=device) - if _hist is None: - hist_index = _require_tensor( - ctx.halo_tables["hist_index"], "ctx.halo_tables['hist_index']" - ) - send = ( - x_flat.index_select(0, hist_index) - if hist_index.numel() - else x_flat.new_empty((0, C)) - ) - hist, work = _exchange( - send, - recv_rows=int(halo.n_hist_rows), - send_splits=tuple(halo.fwd_send_rows), - recv_splits=tuple(halo.fwd_recv_rows), - ctx=ctx, - ) - else: - hist, work = _hist, None - if S == 0: - if work is not None: - work.wait() - return y, hist - hist_row_off = _table( - ctx, "hist_row_off", shape=(n_frag, HALO), dtype=torch.int32, device=device - ) - (name,) = route_modules(arch, KCP_FWD) - device_index = _device_index(device) - record = _loader.MODULES[name] - scalars = plan(S, C, n_frag, geometry=record["grid_rule"]["geometry"], legacy=True) - if scalars["n_ch_tiles"] > launch_grid(record, device_index=device_index)[0]: - raise ValueError( - f"C={C} spans more channel tiles than the forward's grid on this device holds" - ) - values: dict[str, Any] = { - "x": x_flat, - "w": w, - "cu": cu, - "y": y.view(S, C), - "hist": hist - if hist.shape[0] - else _pool(0, C, dtype=torch.bfloat16, device=device), - "hist_row_off": hist_row_off, - **_scalars(scalars), - } - if work is not None: - work.wait() - run_module(name, values, device_index=device_index) - return y, hist - - -def _fwd_host( - x: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - *, - out: Optional[torch.Tensor], - device: torch.device, -) -> tuple[int, int, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: - """The trimmed forward's operand checks, output and host-cache entry: ``(S, C, x_flat, w, y, entry)``. The entry - is cached per ``(ctx, S, C, device, cu_seqlens_local_cpu)``: the validated ``cu`` / ``hist_index`` / - ``hist_row_off`` tables, the halo splits, the record, its plan scalars and grid, the loaded module with its - argument plan.""" - - S, C, x_flat = _tokens(x, "x_local", device=device) - w = _weight(weight, C, device=device) - y = _output(out, "out", shape=tuple(x.shape), dtype=torch.bfloat16, device=device) - e = _host_cache_entry( - ctx, ("fwd", S, C, str(device), tuple(ctx.cu_seqlens_local_cpu)) - ) - if not e: - arch = _device_arch(device) - cu, n_frag = _context(ctx, device=device, T_r=S) - halo = ctx.halo - e.update( - cu=cu, - n_frag=n_frag, - hist_index=_require_tensor( - ctx.halo_tables["hist_index"], "ctx.halo_tables['hist_index']" - ), - recv_rows=int(halo.n_hist_rows), - send_splits=tuple(halo.fwd_send_rows), - recv_splits=tuple(halo.fwd_recv_rows), - ) - if S != 0: - e["hist_row_off"] = _table( - ctx, - "hist_row_off", - shape=(n_frag, HALO), - dtype=torch.int32, - device=device, - ) - (name,) = route_modules(arch, KCP_FWD) - device_index = _device_index(device) - record = _loader.MODULES[name] - scalars = plan( - S, C, n_frag, geometry=record["grid_rule"]["geometry"], legacy=True - ) - grid = launch_grid(record, device_index=device_index) - if scalars["n_ch_tiles"] > grid[0]: - raise ValueError( - f"C={C} spans more channel tiles than the forward's grid on this device holds" - ) - e["scalars"] = _scalars(scalars) - e["device_index"] = device_index - e["launch"] = _launch_plan(name, grid, device_index) - return S, C, x_flat, w, y, e - - -def _fwd_gather(x_flat: torch.Tensor, C: int, e: dict[str, Any]) -> torch.Tensor: - """The forward's send rows: the context's ``hist_index`` rows of ``x_flat`` in destination-rank order.""" - hist_index = e["hist_index"] - return ( - x_flat.index_select(0, hist_index) - if hist_index.numel() - else x_flat.new_empty((0, C)) - ) - - -def _fwd_launch( - e: dict[str, Any], - x_flat: torch.Tensor, - w: torch.Tensor, - y: torch.Tensor, - hist: torch.Tensor, - work: Optional[Any], - *, - S: int, - C: int, - device: torch.device, -) -> None: - """Wait ``work`` (when given) and launch the forward kernel of a call with ``S > 0`` on ``hist``.""" - values: dict[str, Any] = { - "x": x_flat, - "w": w, - "cu": e["cu"], - "y": y.view(S, C), - "hist": hist - if hist.shape[0] - else _pool(0, C, dtype=torch.bfloat16, device=device), - "hist_row_off": e["hist_row_off"], - **e["scalars"], - } - if work is not None: - work.wait() - _launch_trim(e["launch"], values, e["device_index"]) - - -def _kcp_forward_trim( - x: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - *, - out: Optional[torch.Tensor], - device: torch.device, - _hist: Optional[torch.Tensor] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Trimmed host path: :func:`_kcp_forward` with the context-derived host work cached per ``(ctx, S, C, device, - cu_seqlens_local_cpu)`` (:func:`_fwd_host`). The operand checks, the gather, the exchange, the wait and the - kernel launch are the untrimmed call's, in the same order on the same stream.""" - - S, C, x_flat, w, y, e = _fwd_host(x, weight, ctx, out=out, device=device) - if _hist is None: - hist, work = _exchange( - _fwd_gather(x_flat, C, e), - recv_rows=e["recv_rows"], - send_splits=e["send_splits"], - recv_splits=e["recv_splits"], - ctx=ctx, - ) - else: - hist, work = _hist, None - if S == 0: - if work is not None: - work.wait() - return y, hist - _fwd_launch(e, x_flat, w, y, hist, work, S=S, C=C, device=device) - return y, hist - - -def cake_causal_conv1d_kcp_backward( - x_local: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - dy: torch.Tensor, - hist: torch.Tensor, - *, - dx: Optional[torch.Tensor] = None, - dweight: Optional[torch.Tensor] = None, - graph: Optional[str] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r"""The KCP backward: ``dx`` of the rank's fragments (the halo contributions of the successor fragments folded in - before the one BF16 rounding) and the rank's FP32 ``dweight`` partial (fixed order; its sum over the ranks is the - framework's). - - ``hist`` is the forward's pool. The head program writes the derivative rows of every continuing fragment's first - tokens into the backward exchange's send layout, the exchange delivers them to the predecessors as ``dpre_tail``, - and the main program folds them. Every rank calls this function once per layer, in the same order. ``graph`` as - in the forward (the backward's key holds the forward's ``hist`` pointer, so its capture follows the forward's by - one call). - """ - - mode = _graph_mode(graph) - if mode == "off": - return _kcp_backward(x_local, weight, ctx, dy, hist, dx=dx, dweight=dweight) - entry = _graph_entry("bwd", ctx, (x_local, weight, dy, hist, dx, dweight), mode) - dx_out, dw_out = _graph_run( - "bwd", - entry, - lambda: _kcp_backward(x_local, weight, ctx, dy, hist, dx=dx, dweight=dweight), - mode, - ) - return (dx_out if dx is not None else dx_out.detach()), ( - dw_out if dweight is not None else dw_out.detach() - ) - - -def _kcp_backward( - x_local: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - dy: torch.Tensor, - hist: torch.Tensor, - *, - dx: Optional[torch.Tensor] = None, - dweight: Optional[torch.Tensor] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """The backward's host flow (eagerly or inside a capture): see :func:`cake_causal_conv1d_kcp_backward`.""" - - x = _require_tensor(x_local, "x_local") - device = x.device - # as in the forward: every call, before the host-cache lookup of the trimmed path - _wait_ready(ctx) - if _TEST_KNOBS["host_trim"] == 1: - return _kcp_backward_trim( - x, weight, ctx, dy, hist, dx=dx, dweight=dweight, device=device - ) - arch = _device_arch(device) - S, C, x_flat = _tokens(x, "x_local", device=device) - w = _weight(weight, C, device=device) - _, _, dy_flat = _tokens(dy, "dy", device=device, shape=tuple(x.shape)) - cu, n_frag = _context(ctx, device=device, T_r=S) - halo = ctx.halo - hist_t = _require_tensor(hist, "hist") - if ( - hist_t.ndim != 2 - or int(hist_t.shape[1]) != C - or hist_t.dtype != torch.bfloat16 - or not hist_t.is_contiguous() - ): - raise ValueError( - f"hist must be the forward's contiguous BF16 [n_hist_rows, {C}] pool, got {hist_t.dtype} {tuple(hist_t.shape)}" - ) - if hist_t.shape[0] == 0: - # A valid pointer for the kernels; no table entry names a row of it. - hist_t = _pool(0, C, dtype=torch.bfloat16, device=device) - dx_out = _output( - dx, "dx", shape=tuple(x.shape), dtype=torch.bfloat16, device=device - ) - dw_out = _output( - dweight, "dweight", shape=(C, WIDTH), dtype=torch.float32, device=device - ) - device_index = _device_index(device) - n_head = int(halo.n_head_rows) - dpre_head = _pool(n_head, C, dtype=torch.float32, device=device) - if S > 0 and n_head > 0: - hist_row_off = _table( - ctx, "hist_row_off", shape=(n_frag, HALO), dtype=torch.int32, device=device - ) - head_row_off = _table( - ctx, - "head_row_off", - shape=(n_frag, HALO, HEAD_SLOTS), - dtype=torch.int32, - device=device, - ) - (head_name,) = route_modules(arch, KCP_BWD_HEAD) - geometry = _loader.MODULES[head_name]["grid_rule"]["geometry"] - channels_per_cta = int(geometry["threads"]) * int(geometry["vec"]) - grid = (max(1, n_frag * _ceil_div(C, channels_per_cta)), 1, 1) - _launch_with_grid( - head_name, - { - "x": x_flat, - "w": w, - "cu": cu, - "dy": dy_flat, - "hist": hist_t, - "hist_row_off": hist_row_off, - "dpre_head": dpre_head, - "head_row_off": head_row_off, - "S": S, - "C": C, - "num_docs": n_frag, - }, - grid=grid, - device_index=device_index, - ) - dpre_tail, work = _exchange( - dpre_head[:n_head], - recv_rows=int(halo.n_tail_rows), - send_splits=tuple(halo.bwd_send_rows), - recv_splits=tuple(halo.bwd_recv_rows), - ctx=ctx, - ) - if S == 0: - if work is not None: - work.wait() - dw_out.zero_() - return dx_out, dw_out - hist_row_off = _table( - ctx, "hist_row_off", shape=(n_frag, HALO), dtype=torch.int32, device=device - ) - tail_row_off = _table( - ctx, "tail_row_off", shape=(n_frag, HALO), dtype=torch.int32, device=device - ) - names = selected_modules(arch, KCP_BWD, C=C, S=S) - record = _loader.MODULES[names[0]] - scalars = plan(S, C, n_frag, geometry=record["grid_rule"]["geometry"], legacy=True) - grid_x = launch_grid(record, device_index=device_index)[0] - if scalars["n_ch_tiles"] > grid_x: - raise ValueError( - f"C={C} spans more channel tiles than the backward's grid on this device holds" - ) - rows = partial_rows(grid_x, scalars["n_ch_tiles"]) - partial = torch.empty((rows, C, WIDTH), dtype=torch.float32, device=device) - values: dict[str, Any] = { - "x": x_flat, - "w": w, - "cu": cu, - "dy": dy_flat, - "dx": dx_out.view(S, C), - "dw": dw_out, - "partial": partial, - "hist": hist_t, - "hist_row_off": hist_row_off, - "dpre_tail": dpre_tail - if dpre_tail.shape[0] - else _pool(0, C, dtype=torch.float32, device=device), - "tail_row_off": tail_row_off, - **_scalars(scalars), - } - if work is not None: - work.wait() - for name in names: - run_module(name, values, device_index=device_index) - return dx_out, dw_out - - -def _bwd_host( - x: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - dy: torch.Tensor, - hist: torch.Tensor, - *, - dx: Optional[torch.Tensor], - dweight: Optional[torch.Tensor], - device: torch.device, -) -> tuple[Any, ...]: - """The trimmed backward's operand checks, outputs and host-cache entry: ``(S, C, x_flat, w, dy_flat, hist_t, - dx_out, dw_out, entry)``. The entry is cached per ``(ctx, S, C, device, cu_seqlens_local_cpu)``: the tables, the - head grid, the backward records, plan scalars, partial rows, the loaded modules with their argument plans.""" - - S, C, x_flat = _tokens(x, "x_local", device=device) - w = _weight(weight, C, device=device) - _, _, dy_flat = _tokens(dy, "dy", device=device, shape=tuple(x.shape)) - hist_t = _require_tensor(hist, "hist") - if ( - hist_t.ndim != 2 - or int(hist_t.shape[1]) != C - or hist_t.dtype != torch.bfloat16 - or not hist_t.is_contiguous() - ): - raise ValueError( - f"hist must be the forward's contiguous BF16 [n_hist_rows, {C}] pool, got {hist_t.dtype} {tuple(hist_t.shape)}" - ) - if hist_t.shape[0] == 0: - hist_t = _pool(0, C, dtype=torch.bfloat16, device=device) - dx_out = _output( - dx, "dx", shape=tuple(x.shape), dtype=torch.bfloat16, device=device - ) - dw_out = _output( - dweight, "dweight", shape=(C, WIDTH), dtype=torch.float32, device=device - ) - e = _host_cache_entry( - ctx, ("bwd", S, C, str(device), tuple(ctx.cu_seqlens_local_cpu)) - ) - if not e: - arch = _device_arch(device) - cu, n_frag = _context(ctx, device=device, T_r=S) - halo = ctx.halo - device_index = _device_index(device) - n_head = int(halo.n_head_rows) - e.update( - cu=cu, - n_frag=n_frag, - device_index=device_index, - n_head=n_head, - tail_rows=int(halo.n_tail_rows), - send_splits=tuple(halo.bwd_send_rows), - recv_splits=tuple(halo.bwd_recv_rows), - head=None, - ) - if S > 0 and n_head > 0: - hist_row_off = _table( - ctx, - "hist_row_off", - shape=(n_frag, HALO), - dtype=torch.int32, - device=device, - ) - head_row_off = _table( - ctx, - "head_row_off", - shape=(n_frag, HALO, HEAD_SLOTS), - dtype=torch.int32, - device=device, - ) - (head_name,) = route_modules(arch, KCP_BWD_HEAD) - geometry = _loader.MODULES[head_name]["grid_rule"]["geometry"] - channels_per_cta = int(geometry["threads"]) * int(geometry["vec"]) - grid = (max(1, n_frag * _ceil_div(C, channels_per_cta)), 1, 1) - e["head"] = ( - _launch_plan(head_name, grid, device_index), - hist_row_off, - head_row_off, - ) - if S != 0: - e["hist_row_off"] = _table( - ctx, - "hist_row_off", - shape=(n_frag, HALO), - dtype=torch.int32, - device=device, - ) - e["tail_row_off"] = _table( - ctx, - "tail_row_off", - shape=(n_frag, HALO), - dtype=torch.int32, - device=device, - ) - names = selected_modules(arch, KCP_BWD, C=C, S=S) - record = _loader.MODULES[names[0]] - scalars = plan( - S, C, n_frag, geometry=record["grid_rule"]["geometry"], legacy=True - ) - grid_x = launch_grid(record, device_index=device_index)[0] - if scalars["n_ch_tiles"] > grid_x: - raise ValueError( - f"C={C} spans more channel tiles than the backward's grid on this device holds" - ) - e["rows"] = partial_rows(grid_x, scalars["n_ch_tiles"]) - e["scalars"] = _scalars(scalars) - e["launches"] = tuple( - _launch_plan( - name, - launch_grid(_loader.MODULES[name], device_index=device_index), - device_index, - ) - for name in names - ) - return S, C, x_flat, w, dy_flat, hist_t, dx_out, dw_out, e - - -def _bwd_head( - e: dict[str, Any], - x_flat: torch.Tensor, - w: torch.Tensor, - dy_flat: torch.Tensor, - hist_t: torch.Tensor, - *, - S: int, - C: int, - device: torch.device, -) -> torch.Tensor: - """The head program of the backward (when the rank has head rows): the fresh FP32 ``dpre_head`` pool (at least one - row) whose first ``n_head`` rows are the backward exchange's send rows.""" - dpre_head = _pool(e["n_head"], C, dtype=torch.float32, device=device) - if e["head"] is not None: - info, hist_row_off, head_row_off = e["head"] - _launch_trim( - info, - { - "x": x_flat, - "w": w, - "cu": e["cu"], - "dy": dy_flat, - "hist": hist_t, - "hist_row_off": hist_row_off, - "dpre_head": dpre_head, - "head_row_off": head_row_off, - "S": S, - "C": C, - "num_docs": e["n_frag"], - }, - e["device_index"], - ) - return dpre_head - - -def _bwd_main( - e: dict[str, Any], - x_flat: torch.Tensor, - w: torch.Tensor, - dy_flat: torch.Tensor, - hist_t: torch.Tensor, - dx_out: torch.Tensor, - dw_out: torch.Tensor, - dpre_tail: torch.Tensor, - work: Optional[Any], - *, - S: int, - C: int, - device: torch.device, -) -> None: - """Wait ``work`` (when given) and launch the main backward programs of a call with ``S > 0`` (``dx`` with the - received ``dpre_tail`` rows folded in, and the rank's fixed-order FP32 ``dweight`` partial).""" - partial = torch.empty((e["rows"], C, WIDTH), dtype=torch.float32, device=device) - values: dict[str, Any] = { - "x": x_flat, - "w": w, - "cu": e["cu"], - "dy": dy_flat, - "dx": dx_out.view(S, C), - "dw": dw_out, - "partial": partial, - "hist": hist_t, - "hist_row_off": e["hist_row_off"], - "dpre_tail": dpre_tail - if dpre_tail.shape[0] - else _pool(0, C, dtype=torch.float32, device=device), - "tail_row_off": e["tail_row_off"], - **e["scalars"], - } - if work is not None: - work.wait() - for info in e["launches"]: - _launch_trim(info, values, e["device_index"]) - - -def _kcp_backward_trim( - x: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - dy: torch.Tensor, - hist: torch.Tensor, - *, - dx: Optional[torch.Tensor], - dweight: Optional[torch.Tensor], - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: - """Trimmed host path: :func:`_kcp_backward` with the context-derived host work cached (:func:`_bwd_host`); the - operand checks, allocations, exchange, wait and launches are the untrimmed call's.""" - - S, C, x_flat, w, dy_flat, hist_t, dx_out, dw_out, e = _bwd_host( - x, weight, ctx, dy, hist, dx=dx, dweight=dweight, device=device - ) - dpre_head = _bwd_head(e, x_flat, w, dy_flat, hist_t, S=S, C=C, device=device) - dpre_tail, work = _exchange( - dpre_head[: e["n_head"]], - recv_rows=e["tail_rows"], - send_splits=e["send_splits"], - recv_splits=e["recv_splits"], - ctx=ctx, - ) - if S == 0: - if work is not None: - work.wait() - dw_out.zero_() - return dx_out, dw_out - _bwd_main( - e, - x_flat, - w, - dy_flat, - hist_t, - dx_out, - dw_out, - dpre_tail, - work, - S=S, - C=C, - device=device, - ) - return dx_out, dw_out - - -class CakeCausalConv1dKcpFunction(torch.autograd.Function): - """Autograd wrapper: the forward saves ``(x_local, weight)`` and the history pool ``hist`` (no second exchange in - the backward); the backward runs the head program, the backward halo exchange and the main backward program. - ``handle``: a recompute handle (see :class:`KcpRecomputeHandle`): the original forward stores ``hist`` in it, the - recompute binds it and exchanges nothing; the saved set is the same three tensors either way. A handle is - refused with graph replay (keyword or knob) and inside a CUDA-graph capture.""" - - @staticmethod - def forward(fn_ctx, x_local, weight, ctx, graph=None, handle=None): - if handle is None: - y, hist = cake_causal_conv1d_kcp_forward(x_local, weight, ctx, graph=graph) - else: - _handle_refusals(graph, handle) - entry = _handle_lookup(handle) - if entry is None: - y, hist = _kcp_forward(x_local, weight, ctx) - handle.entries[HANDLE_KEY] = { - "hist": hist, - "shape": tuple(hist.shape), - "dtype": hist.dtype, - "device": hist.device, - "ctx": ctx, - "key": (int(x_local.shape[-2]), int(x_local.shape[-1])), - } - LAST_CALL["recompute_handle"] = {"fwd": "capture"} - else: - hist = _handle_hist(entry, x_local, ctx) - y, hist = _kcp_forward(x_local, weight, ctx, _hist=hist) - LAST_CALL["recompute_handle"] = {"fwd": "replay"} - fn_ctx.save_for_backward(x_local, weight, hist) - fn_ctx.kcp = ctx - fn_ctx.graph = graph # the backward runs under the same graph mode - return y - - @staticmethod - def backward(fn_ctx, grad_y): - x_local, weight, hist = fn_ctx.saved_tensors - if grad_y is None: - grad_y = torch.zeros_like(x_local) - dy = grad_y.to(torch.bfloat16).contiguous() - dx, dweight = cake_causal_conv1d_kcp_backward( - x_local, weight, fn_ctx.kcp, dy, hist, graph=fn_ctx.graph - ) - return dx, dweight, None, None, None - - -def cake_causal_conv1d_kcp( - x_local: torch.Tensor, - weight: torch.Tensor, - ctx: Any, - *, - graph: Optional[str] = None, - recompute_handle: Optional[Any] = None, -) -> torch.Tensor: - """Autograd-enabled KCP packed width-4 depthwise causal convolution + SiLU of the rank's fragments. - - Parameters - ---------- - x_local : torch.Tensor - BF16 ``[T_r, C]`` (or ``[1, T_r, C]``) contiguous activations of the rank's fragments, back to back in the - order of ``ctx.cu_seqlens_local``. - weight : torch.Tensor - FP32 ``[C, 4]`` contiguous taps; ``weight[c, 3]`` multiplies the current token. - ctx : KcpContext - The rank's context (``flashinfer.experimental.cake_kda_train.build_kcp_context``): fragment table, halo - plan, process group. Every rank of the group calls this function once per layer, in the same order. - graph : str, optional - ``None`` (the shipped default ``GRAPH_DEFAULT`` = ``"off"``, the eager path) or ``"replay"``: keyed CUDA-graph - capture / replay of this call. The first call of a validity key (this - context, its world size, and the shape / stride / dtype / data pointer of ``x_local`` and ``weight``) runs - eagerly, the second is captured, every later one replays the graph (bitwise the eager result). Requirements: - pointer-stable operands and the same ``KcpContext`` per layer, outputs consumed before the next call of the - same layer (a replay overwrites the graph's static tensors); any other key runs eagerly (``GRAPH_STATS``, - ``LAST_CALL``). Each cached key pins a private memory pool (at most ``GRAPH_CACHE_MAX`` keys); call - :func:`graph_cache_clear` before ``torch.distributed.destroy_process_group()``. - recompute_handle : KcpRecomputeHandle, optional - ``None`` (default): unchanged. A handle (any object with an ``entries`` dict): the checkpoint recompute of - this call exchanges nothing -- the original forward stores its history pool under ``"conv"`` and torch's - recompute inside the backward binds it (bitwise the exchanging recompute); the backward still exchanges. - One handle per convolution call and micro-step, created outside the outermost checkpoint region; a forward - outside the backward that finds an entry raises. Refused with ``graph="replay"`` (keyword or knob) and - inside a CUDA-graph capture. Retains ``hist`` (``recompute_handle.nbytes()``) while the handle lives. - The forward halo exchange is a collective: every rank of ``ctx.group`` passes a handle for the same call or - none, created at the same checkpoint nesting level. A mismatch pairs one rank's recompute exchange with - another rank's backward exchange and hangs or corrupts data; no check is made. Never share one handle - between two convolution calls (also not in nested regions), and never ``clear()`` or reuse a handle while a - backward that needs it is pending. - - Returns - ------- - torch.Tensor - BF16 ``y`` of ``x_local``'s shape, with the three document tokens before every continuing fragment as - history (zero history at a document's start). The backward returns BF16 ``dx`` (the successor fragments' - contributions folded in before the one rounding) and the rank's FP32 ``dweight`` partial. - """ - - _graph_mode(graph) # validate on the caller's frame - if recompute_handle is not None: - _handle_refusals(graph, recompute_handle) - # Always five inputs: the backward returns five gradients. - return CakeCausalConv1dKcpFunction.apply( - x_local, weight, ctx, graph, recompute_handle - ) - - -# ---------------------------------------------------------------------------------------------------------------- -# Several convolutions of one layer with one halo exchange per direction (FlashInfer issue #7815): the q / k / v -# convolutions of a layer share the context, so their history rows (forward) and head rows (backward) travel in ONE -# ``all_to_all_single`` instead of one per tensor. The send rows of the tensors are packed side by side -# (``[n_rows, sum C_i]``, the context's splits unchanged), and after the wait each tensor's rows are copied out into -# a fresh contiguous pool -- the forward's ``hist`` ``[n_hist_rows, C_i]`` (the saved tensor) and the backward's -# ``dpre_tail`` ``[n_tail_rows, C_i]`` -- before that tensor's kernels launch with the single call's arguments. The -# exchange and the unpacking are copies, and each tensor keeps its own kernels, partial workspace and fixed -# reduction order, so every output and gradient (``dweight`` included) is bitwise that of one single call per tensor. -# One rank (no collective): the per-tensor send rows are the pools, exactly as in the single call. - - -def _multi_operands( - xs: Any, weights: Any, outs: Any, name: str = "xs" -) -> tuple[list[torch.Tensor], list[torch.Tensor], list[Optional[torch.Tensor]]]: - """``xs`` / ``weights`` (/ ``outs``) as lists of one length >= 1, on one device.""" - if isinstance(xs, torch.Tensor) or isinstance(weights, torch.Tensor): - raise TypeError( - f"{name} and weights must be sequences of tensors (one per convolution), not a tensor" - ) - xs, weights = list(xs), list(weights) - if not xs or len(xs) != len(weights): - raise ValueError( - f"{name} and weights must hold the same number (>= 1) of tensors, got {len(xs)} and {len(weights)}" - ) - if outs is None: - outs = [None] * len(xs) - else: - outs = list(outs) - if len(outs) != len(xs): - raise ValueError( - f"outs must hold one tensor per convolution ({len(xs)}), got {len(outs)}" - ) - x0 = _require_tensor(xs[0], f"{name}[0]") - for i, x in enumerate(xs): - if _require_tensor(x, f"{name}[{i}]").device != x0.device: - raise ValueError( - f"{name}[{i}] is on {x.device}, {name}[0] on {x0.device}: one device" - ) - return xs, weights, outs - - -def _unpack(recv: torch.Tensor, widths: list[int]) -> list[torch.Tensor]: - """The tensors' contiguous column blocks of a packed receive pool (one copy each).""" - blocks, off = [], 0 - for C in widths: - blocks.append(recv[:, off : off + C].contiguous()) - off += C - return blocks - - -def cake_causal_conv1d_kcp_multi_forward( - xs: Any, - weights: Any, - ctx: Any, - *, - outs: Optional[Any] = None, - graph: Optional[str] = None, -) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: - r"""The KCP forward of several convolutions of one layer with one halo exchange: ``(ys, hists)``, bitwise - ``cake_causal_conv1d_kcp_forward(xs[i], weights[i], ctx)`` per tensor. - - ``xs`` BF16 ``[T_r, C_i]`` (or ``[1, T_r, C_i]``) contiguous tensors of the rank's fragments (the channel counts - may differ); ``weights`` FP32 ``[C_i, 4]``; ``outs`` optional preallocated outputs. Every rank of ``ctx.group`` - calls this function once per layer, in the same order, with the same number of tensors and the same channel - counts: the one forward exchange is a collective. ``hists[i]`` is a fresh contiguous ``[n_hist_rows, C_i]`` - BF16 pool (the saved tensor of tensor ``i``). ``graph`` as in :func:`cake_causal_conv1d_kcp_forward` (the - validity key holds every operand). - """ - - xs, weights, outs = _multi_operands(xs, weights, outs) - mode = _graph_mode(graph) - if mode == "off": - return _kcp_forward_multi(xs, weights, ctx, outs=outs) - entry = _graph_entry("fwd_multi", ctx, (*xs, *weights, *outs), mode) - ys, hists = _graph_run( - "fwd_multi", - entry, - lambda: _kcp_forward_multi(xs, weights, ctx, outs=outs), - mode, - ) - ys = tuple( - y if o is not None else y.detach() for y, o in zip(ys, outs, strict=True) - ) - return ys, hists - - -def _kcp_forward_multi( - xs: list[torch.Tensor], - weights: list[torch.Tensor], - ctx: Any, - *, - outs: list[Optional[torch.Tensor]], - _hists: Optional[tuple[torch.Tensor, ...]] = None, -) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: - """The multi forward's host flow (eagerly or inside a capture): per-tensor operand checks, the packed gather, one - exchange, one wait, the per-tensor unpack copies and the single call's forward kernel per tensor. ``_hists``: - the pools a recompute handle stored (checked by :func:`_handle_hists`): no gather, no exchange, no wait.""" - - device = _require_tensor(xs[0], "xs[0]").device - _wait_ready(ctx) - calls = [ - _fwd_host(x, w, ctx, out=o, device=device) - for x, w, o in zip(xs, weights, outs, strict=True) - ] - S = calls[0][0] - widths = [c[1] for c in calls] - work = None - if _hists is not None: - hists = list(_hists) - elif int(ctx.world_size) == 1: - # one rank: each tensor's send rows are its pool (the single call's identity exchange) - hists = [_fwd_gather(c[2], c[1], c[5]) for c in calls] - else: - e0 = calls[0][5] - sends = [_fwd_gather(c[2], c[1], c[5]) for c in calls] - send = torch.cat(sends, dim=1) if len(sends) > 1 else sends[0] - recv, work = _exchange( - send, - recv_rows=e0["recv_rows"], - send_splits=e0["send_splits"], - recv_splits=e0["recv_splits"], - ctx=ctx, - ) - if work is not None: - work.wait() - hists = _unpack(recv, widths) if len(calls) > 1 else [recv] - ys = tuple(c[4] for c in calls) - if S == 0: - return ys, tuple(hists) - for (S_i, C, x_flat, w, y, e), hist in zip(calls, hists, strict=True): - _fwd_launch(e, x_flat, w, y, hist, None, S=S_i, C=C, device=device) - return ys, tuple(hists) - - -def cake_causal_conv1d_kcp_multi_backward( - xs: Any, - weights: Any, - ctx: Any, - dys: Any, - hists: Any, - *, - dxs: Optional[Any] = None, - dweights: Optional[Any] = None, - graph: Optional[str] = None, -) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: - r"""The KCP backward of several convolutions of one layer with one halo exchange: ``(dxs, dweights)``, bitwise - ``cake_causal_conv1d_kcp_backward(xs[i], weights[i], ctx, dys[i], hists[i])`` per tensor. - - The head programs of every tensor run first, their send rows travel packed in one exchange, and after the wait - each tensor's received rows are copied out and folded by that tensor's main backward programs (own partial - workspace, fixed order). ``hists`` are the forward's pools. Every rank calls this function once per layer, in - the same order. ``graph`` as in :func:`cake_causal_conv1d_kcp_backward`. - """ - - xs, weights, _ = _multi_operands(xs, weights, None) - n = len(xs) - dys, hists = list(dys), list(hists) - dxs = [None] * n if dxs is None else list(dxs) - dweights = [None] * n if dweights is None else list(dweights) - if not (len(dys) == len(hists) == len(dxs) == len(dweights) == n): - raise ValueError( - f"dys, hists, dxs and dweights must hold one entry per convolution ({n})" - ) - mode = _graph_mode(graph) - if mode == "off": - return _kcp_backward_multi(xs, weights, ctx, dys, hists, dxs, dweights) - entry = _graph_entry( - "bwd_multi", ctx, (*xs, *weights, *dys, *hists, *dxs, *dweights), mode - ) - dx_out, dw_out = _graph_run( - "bwd_multi", - entry, - lambda: _kcp_backward_multi(xs, weights, ctx, dys, hists, dxs, dweights), - mode, - ) - return ( - tuple( - d if o is not None else d.detach() for d, o in zip(dx_out, dxs, strict=True) - ), - tuple( - d if o is not None else d.detach() - for d, o in zip(dw_out, dweights, strict=True) - ), - ) - - -def _kcp_backward_multi( - xs: list[torch.Tensor], - weights: list[torch.Tensor], - ctx: Any, - dys: list[torch.Tensor], - hists: list[torch.Tensor], - dxs: list[Optional[torch.Tensor]], - dweights: list[Optional[torch.Tensor]], -) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: - """The multi backward's host flow (eagerly or inside a capture): see - :func:`cake_causal_conv1d_kcp_multi_backward`.""" - - device = _require_tensor(xs[0], "xs[0]").device - _wait_ready(ctx) - calls = [ - _bwd_host(x, w, ctx, dy, h, dx=dx, dweight=dw, device=device) - for x, w, dy, h, dx, dw in zip( - xs, weights, dys, hists, dxs, dweights, strict=True - ) - ] - heads = [ - _bwd_head(c[8], c[2], c[3], c[4], c[5], S=c[0], C=c[1], device=device) - for c in calls - ] - e0 = calls[0][8] - n_head = e0["n_head"] - work = None - if int(ctx.world_size) == 1: - tails = [h[:n_head] for h in heads] - else: - send = ( - torch.cat([h[:n_head] for h in heads], dim=1) - if len(heads) > 1 - else heads[0][:n_head] - ) - recv, work = _exchange( - send, - recv_rows=e0["tail_rows"], - send_splits=e0["send_splits"], - recv_splits=e0["recv_splits"], - ctx=ctx, - ) - if work is not None: - work.wait() - tails = _unpack(recv, [c[1] for c in calls]) if len(calls) > 1 else [recv] - dx_out = tuple(c[6] for c in calls) - dw_out = tuple(c[7] for c in calls) - if calls[0][0] == 0: - for dw in dw_out: - dw.zero_() - return dx_out, dw_out - for (S, C, x_flat, w, dy_flat, hist_t, dx, dw, e), tail in zip( - calls, tails, strict=True - ): - _bwd_main( - e, x_flat, w, dy_flat, hist_t, dx, dw, tail, None, S=S, C=C, device=device - ) - return dx_out, dw_out - - -def _handle_hists( - entry: dict[str, Any], xs: list[torch.Tensor], ctx: Any -) -> tuple[torch.Tensor, ...]: - """The stored pools of a multi call's ``"conv"`` entry, each checked like :func:`_handle_hist` against its tensor - of this recompute (an entry of a single call, or of another tensor count, raises).""" - pools, keys, shapes = entry.get("hists"), entry.get("keys"), entry.get("shapes") - n = len(xs) - if not all(isinstance(t, tuple) and len(t) == n for t in (pools, keys, shapes)): - raise ValueError( - f"recompute_handle: the 'conv' entry holds no {n}-tensor multi-call pools (an entry of " - f"cake_causal_conv1d_kcp or of another tensor count); one handle serves one call" - ) - if entry.get("ctx") is not ctx: - raise ValueError( - "recompute_handle: the 'conv' entry belongs to another KcpContext" - ) - return tuple( - _handle_hist( - { - "hist": hist, - "ctx": ctx, - "key": key, - "shape": shape, - "dtype": entry.get("dtype"), - "device": entry.get("device"), - }, - x, - ctx, - ) - for hist, x, key, shape in zip(pools, xs, keys, shapes, strict=True) - ) - - -class CakeCausalConv1dKcpMultiFunction(torch.autograd.Function): - """Autograd wrapper of several convolutions with one halo exchange per direction: ``apply(ctx, graph, handle, - n, *xs, *weights)`` returns the ``n`` outputs; the forward saves ``(*xs, *weights, *hists)``; the backward runs - :func:`cake_causal_conv1d_kcp_multi_backward`. ``handle``: one recompute handle for the call, whose ``"conv"`` - entry holds the ``n`` pools (the saved set's own tensors): the recompute binds them and exchanges nothing. - Refused with graph replay (keyword or knob) and inside a CUDA-graph capture, like the single call.""" - - @staticmethod - def forward(fn_ctx, ctx, graph, handle, n, *tensors): - xs, weights = list(tensors[:n]), list(tensors[n:]) - if handle is None: - ys, hists = cake_causal_conv1d_kcp_multi_forward( - xs, weights, ctx, graph=graph - ) - else: - _handle_refusals(graph, handle) - entry = _handle_lookup(handle) - outs = [None] * n - if entry is None: - ys, hists = _kcp_forward_multi(xs, weights, ctx, outs=outs) - handle.entries[HANDLE_KEY] = { - "hists": hists, - "shapes": tuple(tuple(h.shape) for h in hists), - "dtype": torch.bfloat16, - "device": hists[0].device, - "ctx": ctx, - "keys": tuple((int(x.shape[-2]), int(x.shape[-1])) for x in xs), - } - LAST_CALL["recompute_handle"] = {"fwd": "capture", "n": n} - else: - hists = _handle_hists(entry, xs, ctx) - ys, hists = _kcp_forward_multi( - xs, weights, ctx, outs=outs, _hists=hists - ) - LAST_CALL["recompute_handle"] = {"fwd": "replay", "n": n} - fn_ctx.save_for_backward(*xs, *weights, *hists) - fn_ctx.kcp = ctx - fn_ctx.graph = graph - fn_ctx.n = n - return tuple(ys) - - @staticmethod - def backward(fn_ctx, *grad_ys): - n = fn_ctx.n - saved = fn_ctx.saved_tensors - xs, weights, hists = saved[:n], saved[n : 2 * n], saved[2 * n :] - dys = [ - (torch.zeros_like(x) if g is None else g).to(torch.bfloat16).contiguous() - for g, x in zip(grad_ys, xs, strict=True) - ] - dxs, dweights = cake_causal_conv1d_kcp_multi_backward( - xs, weights, fn_ctx.kcp, dys, hists, graph=fn_ctx.graph - ) - return (None, None, None, None, *dxs, *dweights) - - -def cake_causal_conv1d_kcp_multi( - xs: Any, - weights: Any, - ctx: Any, - *, - graph: Optional[str] = None, - recompute_handle: Optional[Any] = None, -) -> tuple[torch.Tensor, ...]: - """Autograd-enabled KCP convolutions of several tensors of one layer (for example its q, k and v) with one halo - exchange per direction. - - Bitwise ``tuple(cake_causal_conv1d_kcp(x, w, ctx) for x, w in zip(xs, weights))`` -- outputs and every gradient, - ``dweight`` included (each tensor keeps its own kernels and fixed-order partial) -- with one collective in the - forward and one in the backward instead of one per tensor: fewer global synchronization points per layer. - - Parameters - ---------- - xs : sequence of torch.Tensor - ``n >= 1`` BF16 ``[T_r, C_i]`` (or ``[1, T_r, C_i]``) contiguous activations of the rank's fragments, in the - order of ``ctx.cu_seqlens_local``; the channel counts may differ. - weights : sequence of torch.Tensor - FP32 ``[C_i, 4]`` contiguous taps, one per tensor. - ctx : KcpContext - The rank's context, as in :func:`cake_causal_conv1d_kcp`. Every rank of the group calls this function once - per layer, in the same order, with the same tensor count and channel counts. The ranks are not compared - (that would cost a collective): a different count, channel count or order hangs the exchange or silently - corrupts the history rows, as a mismatch of the single entry's per-rank contract does. - graph : str, optional - As in :func:`cake_causal_conv1d_kcp` (keyed CUDA-graph capture / replay; the key holds every operand). - recompute_handle : KcpRecomputeHandle, optional - ``None`` (default): unchanged. One handle for the whole call: the original forward stores the ``n`` history - pools under ``"conv"``, torch's recompute inside the backward binds them and exchanges nothing; the backward - still exchanges once. Same rules as :func:`cake_causal_conv1d_kcp` (one handle per call and micro-step, - created outside the outermost checkpoint region; refused with ``graph="replay"`` and inside a CUDA-graph - capture; every rank passes one or none). The handle may also carry the ``"kda"`` entry of the layer's - recurrence; a handle's ``"conv"`` entry serves either one multi call or one single call, never both. - - Returns - ------- - tuple of torch.Tensor - The ``n`` BF16 outputs, each of its input's shape. - """ - - _graph_mode(graph) - xs, weights, _ = _multi_operands(xs, weights, None) - if recompute_handle is not None: - _handle_refusals(graph, recompute_handle) - return CakeCausalConv1dKcpMultiFunction.apply( - ctx, graph, recompute_handle, len(xs), *xs, *weights - ) - - -__all__ = [ - "GRAPH_MODES", - "GRAPH_DEFAULT", - "GRAPH_STATS", - "HANDLE_KEY", - "KcpRecomputeHandle", - "LAST_CALL", - "graph_cache_clear", - "CakeCausalConv1dKcpFunction", - "cake_causal_conv1d_kcp", - "cake_causal_conv1d_kcp_backward", - "cake_causal_conv1d_kcp_forward", - "CakeCausalConv1dKcpMultiFunction", - "cake_causal_conv1d_kcp_multi", - "cake_causal_conv1d_kcp_multi_backward", - "cake_causal_conv1d_kcp_multi_forward", -] diff --git a/flashinfer/cake_causal_conv1d_train_backend/kcp_jit.py b/flashinfer/cake_causal_conv1d_train_backend/kcp_jit.py index 84edb189ac5..2a7aac41048 100644 --- a/flashinfer/cake_causal_conv1d_train_backend/kcp_jit.py +++ b/flashinfer/cake_causal_conv1d_train_backend/kcp_jit.py @@ -1,821 +1,9 @@ -""" -Copyright (c) 2026 by FlashInfer team. +# Copyright (c) 2026 by FlashInfer team. Licensed under Apache-2.0. +"""Compatibility import for flashinfer.cake_causal_conv1d_det_train_backend.kcp_jit (the same module object).""" -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at +from importlib import import_module +import sys - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. - -Host plan and loader glue of the Cake context-parallel (KCP) packed width-4 -depthwise causal convolution + SiLU training kernels (SM100 / SM103 / SM107): -the halo-aware programs the KCP host flow (:mod:`.kcp`) launches, recorded -beside the zero-history registry :mod:`.cake_jit`, whose tables this module -never touches. - -Every module below is one mechanically exported production build: a CUDA -device translation unit plus a tvm-ffi binding compiled together for exactly -one architecture. ``MODULES`` and ``ROUTES`` carry the host plan only, in the -vocabulary of :mod:`.cake_jit`: the kernel's ``library`` stem, the -architectures it is delivered for, its launch record, grid rule, workspace -rule, selection rule and the kernel-parameter -> logical-argument map. What to -compile -- sources, digests, nvcc flags -- is recorded once in ``build.json`` -next to this file (its entries sit beside the zero-history entries); -:mod:`.cake_build` serves the library ``_`` from a prebuilt -``aot/`` file when a builder compiled it and otherwise through FlashInfer's JIT -from the same sources and flags, verifying every source digest first. - -The routes of an architecture are the three KCP kinds: ``kcp_fwd`` (one -program: ``y`` and the fragments' history rows ``hist`` from ``x``, ``weight``, -``cu_seqlens`` and the context's history tables), ``kcp_bwd`` (the fallback -program and the big-channel ``_st8`` program, selected by the channel count -exactly as the zero-history backward route selects its programs: a record's -``select_rule`` ``channels_at_least`` owns ``C >= channels[arch]``, the record -without a rule is the fallback; ``dx``, the deterministic FP32 ``dweight`` and -the predecessor fragments' tail derivatives ``dpre_tail`` in one cooperative -launch) and ``kcp_bwd_head`` (one program: the head-row derivatives -``dpre_head`` every consumer fragment sends back; its grid is -``max(1, n_frag x ceil(C / (threads x vec)))``, computed by :mod:`.kcp` from -the record's ``grid_rule.geometry`` and launched through the record's -``arg_plan`` / ``logical`` / ``ffi_entry`` and :func:`load`). - -The host plan reuses :mod:`.cake_jit`'s functions over a record (:func:`plan`, -:func:`launch_grid`, :func:`partial_rows`, :func:`partial_bytes`, -:func:`workspace_layout`, :func:`channel_split`) unchanged; the functions bound -to the tables (:func:`route_modules`, :func:`selected_modules`, -:func:`workspace_record`, :func:`library_name`, :func:`load`, -:func:`run_module`) are spelled here over ``MODULES`` and ``ROUTES``. The -logical tensor names are :mod:`.cake_jit`'s plus the KCP tables ``hist``, -``hist_row_off``, ``dpre_tail``, ``tail_row_off``, ``dpre_head`` and -``head_row_off``; the two ``CUtensorMap`` parameters ``x_map`` / ``dy_map`` of -the TMA backward programs map onto ``x`` / ``dy`` (the binding encodes the -descriptors from the tensors at every launch); the scalars are -:mod:`.cake_jit`'s. - -The program name ``cake_causal_conv1d_train_<20 hex>`` is the first twenty hex -digits of the SHA-256 of the compact JSON ``{"build": , -"module": }`` of the kernel the program was rendered -from: the forward and the head program with their kernel files' default -keywords; the backward with ``vec=8, copy="tma2d"`` on sm_100a and sm_103a, -with ``vec=8, min_blocks=2`` on sm_107a, and the big-channel backward with -``vec=8, copy="tma2d", stage_tokens=8, min_blocks=2`` -- the builds of the -zero-history backward, which the KCP kernels extend. -""" - -from __future__ import annotations - -import functools -from typing import Any, Sequence - -import tvm_ffi - -from . import cake_build -from . import cake_jit as _zero_history -from .cake_jit import ( # noqa: F401 (re-exported: the host plan over a record is the zero-history one) - ARCH_BY_CAPABILITY, - ARCHES, - CHANNEL_MULTIPLE, - INT32_MAX, - LOGICAL_SCALARS, - SELECT_KINDS, - WIDTH, - arch_for_capability, - channel_split, - launch_grid, - partial_bytes, - partial_rows, - plan, - sm_count, - workspace_layout, +sys.modules[__name__] = import_module( + "flashinfer.cake_causal_conv1d_det_train_backend.kcp_jit" ) - -# The derived leaves of every record (``kernel_symbol``, ``launch.block``, -# ``launch.cluster``, ``launch.cooperative``, ``launch.dynamic_smem_bytes``, -# ``launch.use_pdl``, ``arg_plan`` and ``logical``) are rewritten by the -# exporter from the rendered program; every other field is host plan. -MODULES: dict[str, dict[str, Any]] = { - "cake_causal_conv1d_train_675940a7a4fc0dba4b53": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "y"], - ["buffer", "hist"], - ["buffer", "hist_row_off"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # The grid is the co-resident CTA count, ``ctas_per_sm x SM count``, - # as for the zero-history forward: CTA ``b`` owns channel tile - # ``b % n_ch_tiles`` and an equal contiguous share of the token axis. - # ``ctas_per_sm`` is the measured co-residency of the program (the - # occupancy query's answer, 4 CTAs per SM at 112 registers on GB300; - # sm_100a / sm_107a are taken as the same launch bound); ``geometry`` - # is the kernel's build geometry the host plan derives its scalars - # from. Besides ``y`` the forward writes ``hist`` (BF16 history rows - # ``[n_hist_rows, C]``) at the rows ``hist_row_off`` (int32 - # ``[n_frag, 3]``, ``-1`` = none) names. - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 4, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_675940a7a4fc0dba4b53", - "kernel_symbol": "kernel_cake_causal_conv1d_train_675940a7a4fc0dba4b53", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 33792, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_675940a7a4fc0dba4b53", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "hist": "hist", - "hist_row_off": "hist_row_off", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "search_steps": "search_steps", - "w": "w", - "x": "x", - "y": "y", - }, - "role": "kcp_fwd", - "workspace_rule": None, - }, - "cake_causal_conv1d_train_5288d132a5cc927ae729": { - "arches": [ - "sm_100a", - "sm_103a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "dw"], - ["buffer", "partial"], - ["buffer", "hist"], - ["buffer", "hist_row_off"], - ["buffer", "dpre_tail"], - ["buffer", "tail_row_off"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # Cooperative (a grid-wide barrier precedes the fixed-order dweight - # reduction): the grid is exactly the co-resident CTA count, 3 CTAs - # per SM (168 registers at 128 threads, the launch bound), the - # occupancy query's answer on B200 and GB300 for this build. The - # fallback program of the route (no ``select_rule``). Besides ``dx`` - # and ``dw`` it writes the predecessors' tail derivatives - # ``dpre_tail`` (FP32 ``[n_tail_rows, C]``) at the rows - # ``tail_row_off`` (int32 ``[n_frag, 3]``, ``-1`` = none) names and - # reads the forward's ``hist`` rows through ``hist_row_off``. - "grid_rule": { - "cooperative": True, - "ctas_per_sm": 3, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_5288d132a5cc927ae729", - "kernel_symbol": "kernel_cake_causal_conv1d_train_5288d132a5cc927ae729", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": True, - "dynamic_smem_bytes": 50176, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_5288d132a5cc927ae729", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dpre_tail": "dpre_tail", - "dw": "dw", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "hist": "hist", - "hist_row_off": "hist_row_off", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "tail_row_off": "tail_row_off", - "w": "w", - "x": "x", - "x_map": "x", - }, - "role": "kcp_bwd", - # The FP32 ``partial`` workspace: ``[ceil(grid / n_ch_tiles), C, taps]`` - # rows of per-CTA dweight partials, reduced in fixed order after the - # grid barrier (:func:`partial_rows`, :func:`partial_bytes`). - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - "cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039": { - "arches": [ - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "dw"], - ["buffer", "partial"], - ["buffer", "hist"], - ["buffer", "hist_row_off"], - ["buffer", "dpre_tail"], - ["buffer", "tail_row_off"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # The sm_107a fallback backward (``vec=8, min_blocks=2``: two CTAs per - # SM, bulk copies, no tensor maps), cooperative like every backward - # program; the grid is its co-resident CTA count. - "grid_rule": { - "cooperative": True, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039", - "kernel_symbol": "kernel_cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": True, - "dynamic_smem_bytes": 50176, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dpre_tail": "dpre_tail", - "dw": "dw", - "dx": "dx", - "dy": "dy", - "hist": "hist", - "hist_row_off": "hist_row_off", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "tail_row_off": "tail_row_off", - "w": "w", - "x": "x", - }, - "role": "kcp_bwd", - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - "cake_causal_conv1d_train_7485529a1755e5bd3036": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "dx"], - ["buffer", "dw"], - ["buffer", "partial"], - ["buffer", "hist"], - ["buffer", "hist_row_off"], - ["buffer", "dpre_tail"], - ["buffer", "tail_row_off"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["parameter", "search_steps"], - ["parameter", "ct_vec"], - ["parameter", "n_sub"], - ["parameter", "n_ch_tiles"], - ["buffer", "x_map"], - ["buffer", "dy_map"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # The big-channel backward (``stage_tokens=8``, eight-row TMA boxes per - # stage, two CTAs per SM): it runs when ``C >= select_rule.channels - # [arch]``, else the architecture's fallback record runs -- the - # zero-history backward's thresholds, which these programs' geometry - # measured. Cooperative; the grid is its co-resident CTA count. - "grid_rule": { - "cooperative": True, - "ctas_per_sm": 2, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_7485529a1755e5bd3036", - "kernel_symbol": "kernel_cake_causal_conv1d_train_7485529a1755e5bd3036", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": True, - "dynamic_smem_bytes": 82944, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_7485529a1755e5bd3036", - "logical": { - "C": "C", - "S": "S", - "ct_vec": "ct_vec", - "cu": "cu", - "dpre_tail": "dpre_tail", - "dw": "dw", - "dx": "dx", - "dy": "dy", - "dy_map": "dy", - "hist": "hist", - "hist_row_off": "hist_row_off", - "n_ch_tiles": "n_ch_tiles", - "n_sub": "n_sub", - "num_docs": "num_docs", - "partial": "partial", - "search_steps": "search_steps", - "tail_row_off": "tail_row_off", - "w": "w", - "x": "x", - "x_map": "x", - }, - "role": "kcp_bwd", - "select_rule": { - "channels": { - "sm_100a": 512, - "sm_103a": 512, - "sm_107a": 1024, - }, - "kind": "channels_at_least", - }, - "workspace_rule": { - "kind": "partial_rows", - "taps": 4, - }, - }, - "cake_causal_conv1d_train_17cd42fb05079bbbf840": { - "arches": [ - "sm_100a", - "sm_103a", - "sm_107a", - ], - "arg_plan": [ - ["buffer", "x"], - ["buffer", "w"], - ["buffer", "cu"], - ["buffer", "dy"], - ["buffer", "hist"], - ["buffer", "hist_row_off"], - ["buffer", "dpre_head"], - ["buffer", "head_row_off"], - ["parameter", "S"], - ["parameter", "C"], - ["parameter", "num_docs"], - ["grid", "grid_x"], - ["grid", "grid_y"], - ["grid", "grid_z"], - ], - "ffi_entry": "run", - # The head-row derivative program: ``dpre_head`` (FP32 ``[n_head_rows, - # C]``) at the rows ``head_row_off`` (int32 ``[n_frag, 3, 3]``, ``-1`` - # = none) names, from the fragment's first three tokens, the history - # rows and ``dy``. Its grid is ``max(1, n_frag x ceil(C / (threads x - # vec)))`` -- a work-item count, computed by the host flow from - # ``geometry`` (the grid strides over the items) -- not the co-resident - # count; ``ctas_per_sm`` records the measured co-residency (5 CTAs per - # SM at 96 registers and no shared memory, the occupancy query's - # answer on GB300) for a device-bound launch. - "grid_rule": { - "cooperative": False, - "ctas_per_sm": 5, - "geometry": { - "threads": 128, - "vec": 8, - }, - "kind": "co_resident", - }, - "kernel": "cake_causal_conv1d_train_17cd42fb05079bbbf840", - "kernel_symbol": "kernel_cake_causal_conv1d_train_17cd42fb05079bbbf840", - "launch": { - "block": (128, 1, 1), - "cluster": (1, 1, 1), - "cooperative": False, - "dynamic_smem_bytes": 0, - "use_pdl": False, - }, - "library": "cake_causal_conv1d_train_17cd42fb05079bbbf840", - "logical": { - "C": "C", - "S": "S", - "cu": "cu", - "dpre_head": "dpre_head", - "dy": "dy", - "head_row_off": "head_row_off", - "hist": "hist", - "hist_row_off": "hist_row_off", - "num_docs": "num_docs", - "w": "w", - "x": "x", - }, - "role": "kcp_bwd_head", - "workspace_rule": None, - }, -} - -# architecture -> KCP kind -> the records the route may launch, in launch -# order. The forward and the head route carry one program; the backward -# route carries the architecture's fallback program first and the -# big-channel ``_st8`` program second (:func:`selected_modules` launches one). -ROUTES: dict[str, dict[str, list[str]]] = { - "sm_100a": { - "kcp_bwd": [ - "cake_causal_conv1d_train_5288d132a5cc927ae729", - "cake_causal_conv1d_train_7485529a1755e5bd3036", - ], - "kcp_bwd_head": [ - "cake_causal_conv1d_train_17cd42fb05079bbbf840", - ], - "kcp_fwd": [ - "cake_causal_conv1d_train_675940a7a4fc0dba4b53", - ], - }, - "sm_103a": { - "kcp_bwd": [ - "cake_causal_conv1d_train_5288d132a5cc927ae729", - "cake_causal_conv1d_train_7485529a1755e5bd3036", - ], - "kcp_bwd_head": [ - "cake_causal_conv1d_train_17cd42fb05079bbbf840", - ], - "kcp_fwd": [ - "cake_causal_conv1d_train_675940a7a4fc0dba4b53", - ], - }, - "sm_107a": { - "kcp_bwd": [ - "cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039", - "cake_causal_conv1d_train_7485529a1755e5bd3036", - ], - "kcp_bwd_head": [ - "cake_causal_conv1d_train_17cd42fb05079bbbf840", - ], - "kcp_fwd": [ - "cake_causal_conv1d_train_675940a7a4fc0dba4b53", - ], - }, -} - -# The KCP kinds (the zero-history kinds are ``cake_jit.KINDS``). -KINDS = ("kcp_fwd", "kcp_bwd", "kcp_bwd_head") -# The history / halo rows every fragment carries (the convolution width minus one). -HALO = WIDTH - 1 -# Logical argument names the KCP host flow binds: the zero-history vocabulary -# (``cake_jit.LOGICAL_TENSORS``, spelled out so the registry stays a literal) -# plus the KCP tables. ``hist``: BF16 history rows ``[n_hist_rows, C]``; -# ``hist_row_off``: int32 ``[n_frag, 3]`` (``-1`` = none); ``dpre_tail``: FP32 -# ``[n_tail_rows, C]`` tail derivatives with ``tail_row_off`` int32 -# ``[n_frag, 3]``; ``dpre_head``: FP32 ``[n_head_rows, C]`` head derivatives -# with ``head_row_off`` int32 ``[n_frag, 3, 3]``. The scalars are the -# zero-history scalars (``cake_jit.LOGICAL_SCALARS``). -LOGICAL_TENSORS = ( - "x", - "w", - "cu", - "y", - "dy", - "dx", - "dw", - "partial", - "hist", - "hist_row_off", - "dpre_tail", - "tail_row_off", - "dpre_head", - "head_row_off", -) - - -def route_applies(*, device_capability: Sequence[int], kind: str) -> bool: - """Whether the verified export owns one KCP training call on a device.""" - - arch = arch_for_capability(device_capability) - if arch is None: - return False - return kind in ROUTES.get(arch, {}) - - -def route_modules(arch: str, kind: str) -> tuple[str, ...]: - """Ordered module names of one KCP route on ``arch`` (every program the route may launch). - - The forward and the head route launch their one program; the backward - launches the one :func:`selected_modules` picks for the call's channel - count. - """ - - if kind not in KINDS: - raise ValueError(f"kind must be one of {KINDS}, got {kind!r}") - names = ROUTES.get(arch, {}).get(kind) - if not names: - raise ValueError( - f"unsupported Cake causal conv1d KCP training route {kind!r} on {arch!r}; " - f"exported architectures: {ARCHES}" - ) - return tuple(names) - - -def select_rule(record: dict[str, Any]) -> dict[str, Any] | None: - """The selection rule of one module record, ``None`` when the record carries none. - - A backward record without a rule (no ``select_rule`` key, or ``None``) is - its route's fallback program. ``{"kind": "channels_at_least", "channels": - {arch: C_min, ...}}`` owns the calls whose channel count is at least - ``channels[arch]`` on the architectures it names and never hits on the - others. Forward and head records carry no rule. - """ - - return record.get("select_rule") - - -def selected_modules( - arch: str, kind: str, *, C: int, S: int | None = None -) -> tuple[str, ...]: - """Ordered module names one KCP call launches on ``arch`` for ``C`` channels. - - The forward and the head route launch every program of their route - (:func:`route_modules`, one each). The backward route carries the - architecture's fallback program and the big-channel program and launches - exactly one: the record whose :func:`select_rule` admits ``C`` on ``arch`` - when one does, else the route's single fallback record (the one without a - rule) -- the rule of the zero-history backward. The choice is a pure - function of the architecture and the channel count (``S`` is accepted and - unused, and the host never reads the data), so a given ``(S, C, device)`` - always runs the same program and ``dweight`` stays bitwise reproducible. - Several admitting records, or zero or several fallbacks, are registry - errors (``RuntimeError``); an unknown rule kind is a ``ValueError``. - """ - - names = route_modules(arch, kind) - if kind != "kcp_bwd": - return names - positive: list[str] = [] - fallback: list[str] = [] - for name in names: - rule = select_rule(MODULES[name]) - if rule is None: - fallback.append(name) - elif _zero_history._admits(rule, arch=arch, C=C): - positive.append(name) - if len(fallback) != 1: - raise RuntimeError( - f"the exported KCP backward route on {arch} declares {len(fallback)} fallback " - "programs (records without a select_rule); expected exactly one" - ) - if len(positive) > 1: - raise RuntimeError( - f"the exported KCP backward route on {arch} selects {len(positive)} programs " - f"for C={int(C)} ({positive}); expected at most one" - ) - if positive: - return (positive[0],) - return (fallback[0],) - - -def workspace_owners(arch: str) -> tuple[dict[str, Any], ...]: - """The KCP backward module records of ``arch`` that own a ``partial`` workspace, in route order.""" - - owners = tuple( - MODULES[name] - for name in route_modules(arch, "kcp_bwd") - if MODULES[name]["workspace_rule"] is not None - ) - if not owners: - raise RuntimeError( - f"the exported KCP backward route on {arch} declares no workspace owner; " - "expected at least one" - ) - return owners - - -def workspace_record(arch: str) -> dict[str, Any]: - """The capacity owner of the KCP backward ``partial`` workspace on ``arch``. - - As for the zero-history route: the owners share one build geometry, so - their rows and bytes grow with ``ctas_per_sm`` alone and the record with - the largest ``ctas_per_sm`` (route order breaks ties) bounds them all. - """ - - owners = workspace_owners(arch) - geometries = { - tuple(sorted(owner["grid_rule"]["geometry"].items())) for owner in owners - } - if len(geometries) != 1: - raise RuntimeError( - f"the KCP backward programs on {arch} declare {len(geometries)} build " - "geometries; one workspace serves them only when they share one" - ) - return max(owners, key=lambda owner: int(owner["grid_rule"]["ctas_per_sm"])) - - -def head_grid(record: dict[str, Any], *, n_frag: int, C: int) -> tuple[int, int, int]: - """The head program's grid: ``max(1, n_frag x ceil(C / (threads x vec)))`` work items over the record's geometry.""" - - geometry = record["grid_rule"]["geometry"] - channels_per_cta = int(geometry["threads"]) * int(geometry["vec"]) - if channels_per_cta <= 0: - raise ValueError( - "the head record's geometry must cover at least one channel per CTA" - ) - items = int(n_frag) * (-(-int(C) // channels_per_cta)) - return (max(1, items), 1, 1) - - -def device_arch(device_index: int) -> str: - """Exported architecture of one CUDA device (``RuntimeError`` when it is not an exported target).""" - - return _zero_history.device_arch(device_index) - - -def library_name(name: str, arch: str) -> str: - """``build.json`` library of KCP module record ``name`` built for ``arch``: ``_``.""" - - record = MODULES[name] - if arch not in record["arches"]: - raise RuntimeError( - f"Cake causal conv1d KCP training module {name} is not delivered for {arch}; " - f"its architectures: {record['arches']}" - ) - return f"{record['library']}_{arch}" - - -@functools.cache -def load(name: str, arch: str) -> Any: - """The loaded tvm-ffi module of KCP record ``name`` built for ``arch``. - - The prebuilt ``aot/`` library when a builder compiled it, else FlashInfer's - JIT build of the same sources and flags (``cake_build.load_library``). - """ - - return cake_build.load_library(library_name(name, arch)) - - -def specs_for_arch(arch: str) -> list[Any]: - """JIT specifications of every ``arch`` library of the package's ``build.json`` without a prebuilt file - (the zero-history and the KCP entries alike; for ahead-of-time warm-up through FlashInfer's - ``build_jit_specs``).""" - - return cake_build.jit_specs(arch) - - -def run_module( - name: str, - values: dict[str, Any], - *, - device_index: int, - grid: tuple[int, int, int] | None = None, -) -> None: - """Launch one exported KCP module with logical argument ``values``. - - ``values`` maps the logical names in ``LOGICAL_TENSORS`` / - ``LOGICAL_SCALARS`` to tensors and Python scalars; the module record maps - its kernel parameter names onto that vocabulary. The launch grid is the - module's :func:`launch_grid` on ``device_index`` (the co-resident count) - unless ``grid`` gives it (the head program's work-item grid, - :func:`head_grid`), and the kernel runs on the current torch stream. - """ - - record = MODULES[name] - logical = record["logical"] - if grid is None: - grid = launch_grid(record, device_index=device_index) - grid_values = { - "grid_x": int(grid[0]), - "grid_y": int(grid[1]), - "grid_z": int(grid[2]), - } - args = [] - for kind, key in record["arg_plan"]: - if kind == "grid": - args.append(grid_values[key]) - continue - logical_name = logical[key] - if logical_name not in values: - raise ValueError( - f"module {name} needs {logical_name!r} (kernel parameter {key!r})" - ) - args.append(values[logical_name]) - module = load(name, device_arch(device_index)) - with tvm_ffi.use_torch_stream(): - getattr(module, record["ffi_entry"])(*args) - - -__all__ = [ - "ARCHES", - "ARCH_BY_CAPABILITY", - "CHANNEL_MULTIPLE", - "HALO", - "INT32_MAX", - "KINDS", - "LOGICAL_SCALARS", - "LOGICAL_TENSORS", - "MODULES", - "ROUTES", - "SELECT_KINDS", - "WIDTH", - "arch_for_capability", - "channel_split", - "device_arch", - "head_grid", - "launch_grid", - "library_name", - "load", - "partial_bytes", - "partial_rows", - "plan", - "route_applies", - "route_modules", - "run_module", - "select_rule", - "selected_modules", - "sm_count", - "specs_for_arch", - "workspace_layout", - "workspace_owners", - "workspace_record", -] diff --git a/pyproject.toml b/pyproject.toml index 8a2a13cc846..37d347efbc4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,7 +86,7 @@ exclude = [ # Cake-exported packages (flashinfer-ai/flashinfer#6553): the generated csrc tree plus the # build manifest (build.json) their cake_build loader reads; prebuilt aot/ libraries never ship. "flashinfer.cake_rmsnorm_det_train_backend" = ["csrc/**", "build.json"] -"flashinfer.cake_causal_conv1d_train_backend" = ["csrc/**", "build.json"] +"flashinfer.cake_causal_conv1d_det_train_backend" = ["csrc/**", "build.json"] "flashinfer.cake_ulysses_backend" = ["csrc/**", "build.json"] "flashinfer.cake_fused_qk_rope_append_backend" = ["csrc/**", "build.json"] "flashinfer.cake_nvfp4_svdquant_gemm_backend" = ["csrc/**", "build.json"] diff --git a/scripts/check_cake_build_manifests.py b/scripts/check_cake_build_manifests.py index 58244223762..40b60be603b 100755 --- a/scripts/check_cake_build_manifests.py +++ b/scripts/check_cake_build_manifests.py @@ -303,6 +303,14 @@ def refs_shipped_modules(package_dir: Path, loader: str) -> set[str]: return generic_library_refs(module_literals(package_dir / loader)["MODULES"]) +def refs_with_kcp(package_dir: Path, loader: str) -> set[str]: + """A package whose ``build.json`` also lists the programs of its context-parallel registry ``kcp_jit.py``: the + generic walk over ``loader`` and the ``MODULES`` records of ``kcp_jit.py``.""" + return refs_generic(package_dir, loader) | refs_shipped_modules( + package_dir, "kcp_jit.py" + ) + + def refs_mok_bf16(package_dir: Path, loader: str) -> set[str]: """``sources.json``: role -> {library, arches, arguments}.""" registry = json.loads((package_dir / loader).read_text()) @@ -335,10 +343,10 @@ class Package: refs_generic, ), Package( - "flashinfer.cake_causal_conv1d_train_backend", - "flashinfer/cake_causal_conv1d_train_backend", + "flashinfer.cake_causal_conv1d_det_train_backend", + "flashinfer/cake_causal_conv1d_det_train_backend", "cake_jit.py", - refs_generic, + refs_with_kcp, ), Package( "flashinfer.cake_gated_act_mxfp8_backend", diff --git a/tests/cpu/test_cake_causal_conv1d_train_kcp_build.py b/tests/cpu/test_cake_causal_conv1d_det_train_kcp_build.py similarity index 88% rename from tests/cpu/test_cake_causal_conv1d_train_kcp_build.py rename to tests/cpu/test_cake_causal_conv1d_det_train_kcp_build.py index 5e4313d38cc..e9e6f9b479d 100644 --- a/tests/cpu/test_cake_causal_conv1d_train_kcp_build.py +++ b/tests/cpu/test_cake_causal_conv1d_det_train_kcp_build.py @@ -28,9 +28,9 @@ from tests.test_helpers.cake_package import cake_package -cake_build = cake_package("cake_causal_conv1d_train", "cake_build") -cake_jit = cake_package("cake_causal_conv1d_train", "cake_jit") -kcp_jit = cake_package("cake_causal_conv1d_train", "kcp_jit") +cake_build = cake_package("cake_causal_conv1d_det_train", "cake_build") +cake_jit = cake_package("cake_causal_conv1d_det_train", "cake_jit") +kcp_jit = cake_package("cake_causal_conv1d_det_train", "kcp_jit") PACKAGE_DIR = Path(cake_build.__file__).parent ARCHES = ("sm_100a", "sm_103a", "sm_107a") @@ -70,7 +70,7 @@ # program name -> (module, build keywords, arches, role, plan, ctas_per_sm) PROGRAMS = { - "cake_causal_conv1d_train_675940a7a4fc0dba4b53": ( + "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp": ( FWD, {}, ARCHES, @@ -78,7 +78,7 @@ FWD_PLAN, 4, ), - "cake_causal_conv1d_train_5288d132a5cc927ae729": ( + "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp": ( BWD, {"vec": 8, "copy": "tma2d"}, ("sm_100a", "sm_103a"), @@ -86,7 +86,7 @@ BWD_TMA_PLAN, 3, ), - "cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039": ( + "cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp": ( BWD, {"vec": 8, "min_blocks": 2}, ("sm_107a",), @@ -94,7 +94,7 @@ BWD_PLAN, 2, ), - "cake_causal_conv1d_train_7485529a1755e5bd3036": ( + "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp": ( BWD, {"vec": 8, "copy": "tma2d", "stage_tokens": 8, "min_blocks": 2}, ARCHES, @@ -102,7 +102,7 @@ BWD_TMA_PLAN, 2, ), - "cake_causal_conv1d_train_17cd42fb05079bbbf840": ( + "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp": ( HEAD, {}, ARCHES, @@ -119,7 +119,7 @@ def _identity(module: str, kwargs: dict) -> str: digest = hashlib.sha256( json.dumps(identity, sort_keys=True, separators=(",", ":")).encode() ).hexdigest() - return f"cake_causal_conv1d_train_{digest[:20]}" + return f"cake_causal_conv1d_det_train_{digest[:20]}_kcp" @pytest.mark.cpu_only @@ -171,16 +171,18 @@ def test_routes_and_selection() -> None: for arch in ARCHES: routes = kcp_jit.ROUTES[arch] assert set(routes) == set(kcp_jit.KINDS) - assert routes["kcp_fwd"] == ["cake_causal_conv1d_train_675940a7a4fc0dba4b53"] + assert routes["kcp_fwd"] == [ + "cake_causal_conv1d_det_train_675940a7a4fc0dba4b53_kcp" + ] assert routes["kcp_bwd_head"] == [ - "cake_causal_conv1d_train_17cd42fb05079bbbf840" + "cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp" ] fallback = ( - "cake_causal_conv1d_train_6ea6d5fc9e7fa3f2d039" + "cake_causal_conv1d_det_train_6ea6d5fc9e7fa3f2d039_kcp" if arch == "sm_107a" - else "cake_causal_conv1d_train_5288d132a5cc927ae729" + else "cake_causal_conv1d_det_train_5288d132a5cc927ae729_kcp" ) - big = "cake_causal_conv1d_train_7485529a1755e5bd3036" + big = "cake_causal_conv1d_det_train_7485529a1755e5bd3036_kcp" assert routes["kcp_bwd"] == [fallback, big] assert kcp_jit.route_modules(arch, "kcp_bwd") == (fallback, big) assert kcp_jit.select_rule(kcp_jit.MODULES[big]) == { @@ -197,7 +199,7 @@ def test_routes_and_selection() -> None: routes["kcp_fwd"][0], ) assert kcp_jit.workspace_record(arch)["role"] == "kcp_bwd" - head = kcp_jit.MODULES["cake_causal_conv1d_train_17cd42fb05079bbbf840"] + head = kcp_jit.MODULES["cake_causal_conv1d_det_train_17cd42fb05079bbbf840_kcp"] assert kcp_jit.head_grid(head, n_frag=3, C=4096) == (3 * 4, 1, 1) assert kcp_jit.head_grid(head, n_frag=0, C=4096) == (1, 1, 1) diff --git a/tests/cpu/test_cake_causal_conv1d_train_kcp_knobs.py b/tests/cpu/test_cake_causal_conv1d_det_train_kcp_knobs.py similarity index 88% rename from tests/cpu/test_cake_causal_conv1d_train_kcp_knobs.py rename to tests/cpu/test_cake_causal_conv1d_det_train_kcp_knobs.py index 1da0c2a5cc4..b388a17e428 100644 --- a/tests/cpu/test_cake_causal_conv1d_train_kcp_knobs.py +++ b/tests/cpu/test_cake_causal_conv1d_det_train_kcp_knobs.py @@ -26,7 +26,7 @@ import pytest import torch -from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp +from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp @pytest.mark.cpu_only @@ -174,9 +174,9 @@ def test_graph_public_keyword_mode_and_key(monkeypatch) -> None: w = torch.empty(16, 4, dtype=torch.float32) ctx = types.SimpleNamespace(T_r=8, n_frag=1, world_size=1) for call in ( # validated on the caller's frame, before any device work - lambda: conv_kcp.cake_causal_conv1d_kcp(x, w, ctx, graph="capture"), - lambda: conv_kcp.cake_causal_conv1d_kcp_forward(x, w, ctx, graph="capture"), - lambda: conv_kcp.cake_causal_conv1d_kcp_backward( + lambda: conv_kcp.cake_causal_conv1d_det_kcp(x, w, ctx, graph="capture"), + lambda: conv_kcp.cake_causal_conv1d_det_forward_kcp(x, w, ctx, graph="capture"), + lambda: conv_kcp.cake_causal_conv1d_det_backward_kcp( x, w, ctx, x, x, graph="capture" ), ): @@ -271,8 +271,8 @@ def test_host_trim_ffi_device_and_cpu_refusal(monkeypatch) -> None: @pytest.mark.cpu_only def test_recompute_handle_object() -> None: - h = conv_kcp.KcpRecomputeHandle() - assert conv_kcp.KcpRecomputeHandle.__slots__ == ("entries",) + h = conv_kcp.RecomputeHandleDetKcp() + assert conv_kcp.RecomputeHandleDetKcp.__slots__ == ("entries",) assert h.entries == {} and h.nbytes() == 0 with pytest.raises(AttributeError): h.other = 1 @@ -283,7 +283,7 @@ def test_recompute_handle_object() -> None: h.clear() assert h.entries == {} and h.nbytes() == 0 assert conv_kcp.HANDLE_KEY == "conv" - assert {"KcpRecomputeHandle", "HANDLE_KEY"} <= set(conv_kcp.__all__) + assert {"RecomputeHandleDetKcp", "HANDLE_KEY"} <= set(conv_kcp.__all__) @pytest.mark.cpu_only @@ -301,18 +301,18 @@ def forward(fctx, t): @staticmethod def backward(fctx, g): seen["in_backward"] = conv_kcp._kcp_in_backward() - seen["empty"] = conv_kcp._handle_lookup(conv_kcp.KcpRecomputeHandle()) + seen["empty"] = conv_kcp._handle_lookup(conv_kcp.RecomputeHandleDetKcp()) seen["entry"] = conv_kcp._handle_lookup(full) return g - full = conv_kcp.KcpRecomputeHandle() + full = conv_kcp.RecomputeHandleDetKcp() full.entries["conv"] = {"hist": torch.zeros(1, 8, dtype=torch.bfloat16)} t = torch.zeros(4, requires_grad=True) Probe.apply(t).sum().backward() assert seen["in_backward"] is True assert seen["empty"] is None assert seen["entry"] is full.entries["conv"] and "conv" in full.entries - assert conv_kcp._handle_lookup(conv_kcp.KcpRecomputeHandle()) is None + assert conv_kcp._handle_lookup(conv_kcp.RecomputeHandleDetKcp()) is None with pytest.raises(RuntimeError, match="micro-step"): conv_kcp._handle_lookup(full) with pytest.raises(TypeError, match="entries"): @@ -326,7 +326,7 @@ def test_recompute_handle_needs_the_graph_task_binding(monkeypatch) -> None: monkeypatch.setattr(conv_kcp, "_GRAPH_TASK_ID_BINDING", None) monkeypatch.setattr(conv_kcp, "_GRAPH_TASK_ID_USABLE", None) with pytest.raises(RuntimeError, match="_current_graph_task_id"): - conv_kcp._handle_lookup(conv_kcp.KcpRecomputeHandle()) + conv_kcp._handle_lookup(conv_kcp.RecomputeHandleDetKcp()) with pytest.raises(RuntimeError, match="_current_graph_task_id"): conv_kcp._kcp_in_backward() @@ -362,24 +362,28 @@ def test_recompute_handle_hist_validation() -> None: @pytest.mark.cpu_only def test_facade_reexports_the_handle_lazily() -> None: - """``flashinfer.cake_causal_conv1d_train.KcpRecomputeHandle`` is the backend's class; importing the facade does not - import the KCP host modules (``kcp`` imports torch.distributed; ``kcp_jit``) -- checked in a fresh interpreter; - the facade's KCP entry carries the keyword-only ``recompute_handle=None``.""" + """``flashinfer.cake_causal_conv1d_det_train_kcp.RecomputeHandleDetKcp`` is the backend's class; importing the KCP + facade (or the det facade, which keeps the previous KCP names) does not import the KCP host modules (``kcp`` imports + torch.distributed; ``kcp_jit``) -- checked in a fresh interpreter; the facade's KCP entry carries the keyword-only + ``recompute_handle=None``, and the previous names are the same objects.""" import subprocess import sys code = ( "import inspect, sys\n" - "import flashinfer.cake_causal_conv1d_train as f\n" - "kcp = 'flashinfer.cake_causal_conv1d_train_backend.kcp'\n" - "jit = 'flashinfer.cake_causal_conv1d_train_backend.kcp_jit'\n" + "import flashinfer.cake_causal_conv1d_det_train_kcp as f\n" + "import flashinfer.cake_causal_conv1d_det_train as d\n" + "kcp = 'flashinfer.cake_causal_conv1d_det_train_backend.kcp'\n" + "jit = 'flashinfer.cake_causal_conv1d_det_train_backend.kcp_jit'\n" "assert kcp not in sys.modules and jit not in sys.modules, sorted(m for m in sys.modules if 'kcp' in m)\n" - "p = inspect.signature(f.cake_causal_conv1d_kcp).parameters['recompute_handle']\n" + "p = inspect.signature(f.cake_causal_conv1d_det_kcp).parameters['recompute_handle']\n" "assert p.kind is inspect.Parameter.KEYWORD_ONLY and p.default is None\n" + "assert d.cake_causal_conv1d_kcp is f.cake_causal_conv1d_det_kcp\n" "assert kcp not in sys.modules\n" - "from flashinfer.cake_causal_conv1d_train import KcpRecomputeHandle\n" - "from flashinfer.cake_causal_conv1d_train_backend.kcp import KcpRecomputeHandle as B\n" - "assert KcpRecomputeHandle is B and 'KcpRecomputeHandle' in f.__all__\n" + "from flashinfer.cake_causal_conv1d_det_train_kcp import RecomputeHandleDetKcp\n" + "from flashinfer.cake_causal_conv1d_det_train_backend.kcp import RecomputeHandleDetKcp as B\n" + "from flashinfer.cake_causal_conv1d_det_train_backend.kcp import KcpRecomputeHandle as C\n" + "assert RecomputeHandleDetKcp is B is C is d.KcpRecomputeHandle and 'RecomputeHandleDetKcp' in f.__all__\n" "print('FACADE_IMPORT_OK')\n" ) r = subprocess.run( diff --git a/tests/norm/test_cake_causal_conv1d_train.py b/tests/norm/test_cake_causal_conv1d_det_train.py similarity index 92% rename from tests/norm/test_cake_causal_conv1d_train.py rename to tests/norm/test_cake_causal_conv1d_det_train.py index 3a322d404a4..8f044007e7b 100644 --- a/tests/norm/test_cake_causal_conv1d_train.py +++ b/tests/norm/test_cake_causal_conv1d_det_train.py @@ -63,8 +63,8 @@ through ``tests.test_helpers.cake_package``, and imports nothing else from flashinfer, so a build that vendors the package under another import path runs it unchanged after pointing that table entry at its package. The -``flashinfer.cake_causal_conv1d_train`` facade and the ``flashinfer`` -top-level exposure are covered by ``test_cake_causal_conv1d_train_api.py``. +``flashinfer.cake_causal_conv1d_det_train`` facade and the ``flashinfer`` +top-level exposure are covered by ``test_cake_causal_conv1d_det_train_api.py``. """ from __future__ import annotations @@ -80,26 +80,28 @@ from tests.test_helpers.cake_package import cake_package -pkg = cake_package("cake_causal_conv1d_train") -cake_build = cake_package("cake_causal_conv1d_train", "cake_build") -loader = cake_package("cake_causal_conv1d_train", "cake_jit") -kcp_loader = cake_package("cake_causal_conv1d_train", "kcp_jit") -cake_causal_conv1d = pkg.cake_causal_conv1d -cake_causal_conv1d_train_backward = pkg.cake_causal_conv1d_train_backward -cake_causal_conv1d_train_backward_workspace = ( - pkg.cake_causal_conv1d_train_backward_workspace +pkg = cake_package("cake_causal_conv1d_det_train") +cake_build = cake_package("cake_causal_conv1d_det_train", "cake_build") +loader = cake_package("cake_causal_conv1d_det_train", "cake_jit") +kcp_loader = cake_package("cake_causal_conv1d_det_train", "kcp_jit") +cake_causal_conv1d_det = pkg.cake_causal_conv1d_det +cake_causal_conv1d_det_train_backward = pkg.cake_causal_conv1d_det_train_backward +cake_causal_conv1d_det_train_backward_workspace = ( + pkg.cake_causal_conv1d_det_train_backward_workspace ) -cake_causal_conv1d_train_backward_workspace_bytes = ( - pkg.cake_causal_conv1d_train_backward_workspace_bytes +cake_causal_conv1d_det_train_backward_workspace_bytes = ( + pkg.cake_causal_conv1d_det_train_backward_workspace_bytes ) -cake_causal_conv1d_train_forward = pkg.cake_causal_conv1d_train_forward -cake_causal_conv1d_train_forward_workspace = ( - pkg.cake_causal_conv1d_train_forward_workspace +cake_causal_conv1d_det_train_forward = pkg.cake_causal_conv1d_det_train_forward +cake_causal_conv1d_det_train_forward_workspace = ( + pkg.cake_causal_conv1d_det_train_forward_workspace ) -cake_causal_conv1d_train_forward_workspace_bytes = ( - pkg.cake_causal_conv1d_train_forward_workspace_bytes +cake_causal_conv1d_det_train_forward_workspace_bytes = ( + pkg.cake_causal_conv1d_det_train_forward_workspace_bytes +) +is_cake_causal_conv1d_det_train_supported = ( + pkg.is_cake_causal_conv1d_det_train_supported ) -is_cake_causal_conv1d_train_supported = pkg.is_cake_causal_conv1d_train_supported SUPPORTED_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) WIDTH = 4 @@ -245,13 +247,13 @@ def _supported() -> bool: return ( torch.cuda.is_available() and _capability() in SUPPORTED_CAPABILITIES - and is_cake_causal_conv1d_train_supported(_device()) + and is_cake_causal_conv1d_det_train_supported(_device()) ) pytestmark = pytest.mark.skipif( not _supported(), - reason="cake_causal_conv1d_train needs an SM100 / SM103 / SM107 device with exported routes", + reason="cake_causal_conv1d_det_train needs an SM100 / SM103 / SM107 device with exported routes", ) @@ -578,7 +580,7 @@ def _check( label or f"S={S} C={C} docs={docs} w={w.dtype} b={None if b is None else b.dtype} {policy}" ) - y = cake_causal_conv1d_train_forward(x, w, cu, bias=b) + y = cake_causal_conv1d_det_train_forward(x, w, cu, bias=b) assert ( y.dtype == torch.bfloat16 and tuple(y.shape) == tuple(x.shape) @@ -597,7 +599,7 @@ def _check( if dbias_dtype is None else torch.empty(C, dtype=dbias_dtype, device=_device()) ) - grads = cake_causal_conv1d_train_backward( + grads = cake_causal_conv1d_det_train_backward( x, w, cu, @@ -653,7 +655,7 @@ def _route_inputs(S: int, C: int, docs, kind: str, *, key: int = 0): def _backward(x, w, cu, dy, b, policy, **kwargs): """``(dx, dweight, dbias_or_None)`` of one backward call.""" - grads = cake_causal_conv1d_train_backward( + grads = cake_causal_conv1d_det_train_backward( x, w, cu, dy, bias=b, silu_grad_preactivation=policy, **kwargs ) return grads[0], grads[1], (grads[2] if b is not None else None) @@ -856,15 +858,15 @@ def test_main_grid_literal_table() -> None: # the kernel module's rules-test rows on both sides of the token and channel # boundaries; C = 512 is the first channel-selected program of every arch. BWD_FALLBACK = { - "sm_100a": "cake_causal_conv1d_train_9bfa65707a49df6a48ce", - "sm_103a": "cake_causal_conv1d_train_9bfa65707a49df6a48ce", - "sm_107a": "cake_causal_conv1d_train_b4f7ff5a3dd9ddd2f41c", + "sm_100a": "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce", + "sm_103a": "cake_causal_conv1d_det_train_9bfa65707a49df6a48ce", + "sm_107a": "cake_causal_conv1d_det_train_b4f7ff5a3dd9ddd2f41c", } -BWD_R200_ST4 = "cake_causal_conv1d_train_9a07bcaa0223632640ee" +BWD_R200_ST4 = "cake_causal_conv1d_det_train_9a07bcaa0223632640ee" BWD_AT_512 = { - "sm_100a": "cake_causal_conv1d_train_37e428406a0cc4a8ca1b", - "sm_103a": "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6", - "sm_107a": "cake_causal_conv1d_train_d46036aa2bc5f90bb0f6", + "sm_100a": "cake_causal_conv1d_det_train_37e428406a0cc4a8ca1b", + "sm_103a": "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6", + "sm_107a": "cake_causal_conv1d_det_train_d46036aa2bc5f90bb0f6", } BWD_SELECTION_TOKENS = (250340, 262143, 262144, 300007) BWD_SELECTION_CHANNELS = (128, 248, 255, 256, 384, 504, 511, 512) @@ -928,7 +930,7 @@ def record_launches(launches, values, *, device_index): launched.extend((name, dict(values), grid) for name, grid in launches) monkeypatch.setattr(pkg, "run_modules", record_launches) - cake_causal_conv1d_train_backward(x, w, cu, dy) + cake_causal_conv1d_det_train_backward(x, w, cu, dy) main = loader.selected_modules(arch, "bwd", C=C, S=S)[0] record = loader.MODULES[main] assert [name for name, _, _ in launched] == [main, record["reduce"]] @@ -1082,10 +1084,12 @@ def test_batched_layout_matches_the_packed_one(case) -> None: S, C, docs = case x, w, cu, dy = _inputs(S, C, docs) - y = cake_causal_conv1d_train_forward(x, w, cu) - dx, dw = cake_causal_conv1d_train_backward(x, w, cu, dy) - y3 = cake_causal_conv1d_train_forward(x.unsqueeze(0), w, cu) - dx3, dw3 = cake_causal_conv1d_train_backward(x.unsqueeze(0), w, cu, dy.unsqueeze(0)) + y = cake_causal_conv1d_det_train_forward(x, w, cu) + dx, dw = cake_causal_conv1d_det_train_backward(x, w, cu, dy) + y3 = cake_causal_conv1d_det_train_forward(x.unsqueeze(0), w, cu) + dx3, dw3 = cake_causal_conv1d_det_train_backward( + x.unsqueeze(0), w, cu, dy.unsqueeze(0) + ) assert tuple(y3.shape) == (1, S, C) and tuple(dx3.shape) == (1, S, C) assert torch.equal(y3[0], y) and torch.equal(dx3[0], dx) and torch.equal(dw3, dw) @@ -1097,16 +1101,16 @@ def test_document_isolation() -> None: x, w, cu, dy = _inputs(S, C, docs) lens = doc_lengths(S, docs) starts = [0, *torch.tensor(lens).cumsum(0).tolist()] - y = cake_causal_conv1d_train_forward(x, w, cu) - dx, _dw = cake_causal_conv1d_train_backward(x, w, cu, dy) + y = cake_causal_conv1d_det_train_forward(x, w, cu) + dx, _dw = cake_causal_conv1d_det_train_backward(x, w, cu, dy) # Perturb every token of document 2 (the 257-token one) in x and dy. a, b = starts[2], starts[3] x2 = x.clone() x2[a:b] = (x2[a:b].float() * -1.5 + 0.25).to(torch.bfloat16) dy2 = dy.clone() dy2[a:b] = (dy2[a:b].float() * 2.0).to(torch.bfloat16) - y2 = cake_causal_conv1d_train_forward(x2, w, cu) - dx2, _dw2 = cake_causal_conv1d_train_backward(x2, w, cu, dy2) + y2 = cake_causal_conv1d_det_train_forward(x2, w, cu) + dx2, _dw2 = cake_causal_conv1d_det_train_backward(x2, w, cu, dy2) assert not torch.equal(y2[a:b], y[a:b]), "the perturbed document did not change" outside = torch.ones(S, dtype=torch.bool, device=_device()) outside[a:b] = False @@ -1121,19 +1125,19 @@ def test_document_isolation() -> None: def test_dweight_is_bitwise_identical_over_three_launches(case) -> None: S, C, docs = case x, w, cu, dy = _inputs(S, C, docs) - workspace = cake_causal_conv1d_train_backward_workspace(S, C, _device()) + workspace = cake_causal_conv1d_det_train_backward_workspace(S, C, _device()) results = [] for _repeat in range(3): dx = torch.full((S, C), float("nan"), dtype=torch.bfloat16, device=_device()) dw = torch.full((C, WIDTH), float("nan"), dtype=torch.float32, device=_device()) # Poison the workspace: the kernel must rewrite every partial row it reduces. workspace.view(torch.float32).fill_(float("nan")) - cake_causal_conv1d_train_backward( + cake_causal_conv1d_det_train_backward( x, w, cu, dy, workspace=workspace, dx=dx, dweight=dw ) torch.cuda.synchronize() results.append((dx.clone(), dw.clone())) - fresh_dx, fresh_dw = cake_causal_conv1d_train_backward(x, w, cu, dy) + fresh_dx, fresh_dw = cake_causal_conv1d_det_train_backward(x, w, cu, dy) for dx, dw in results[1:]: assert torch.equal(dx, results[0][0]) assert torch.equal(dw, results[0][1]) @@ -1169,7 +1173,7 @@ def test_backward_programs_switch_on_channel_count(kind) -> None: assert len(selected) == 1, f"S={S} C={C}: {selected}" launched.add(selected[0]) x, w, cu, dy, b, policy = _route_inputs(S, C, docs, kind) - workspace = cake_causal_conv1d_train_backward_workspace(S, C, _device()) + workspace = cake_causal_conv1d_det_train_backward_workspace(S, C, _device()) # Poison the whole buffer (sized for the largest partial buffer over every route): the selected # program must rewrite every partial row (and dbias row) it reduces. workspace.view(torch.float32).fill_(float("nan")) @@ -1215,7 +1219,7 @@ def test_forward_programs_switch_on_channel_count(kind) -> None: launched.add(selected[0]) x, w, cu, _, b, _ = _route_inputs(S, C, docs, FORWARD_ROUTE_SETTINGS[kind]) label = f"{kind} S={S} C={C} docs={docs} program={selected[0]}" - y = cake_causal_conv1d_train_forward(x, w, cu, bias=b) + y = cake_causal_conv1d_det_train_forward(x, w, cu, bias=b) _assert_within( y, reference_forward(x, w, cu, b), @@ -1268,7 +1272,7 @@ def test_routed_rows_match_fallback_dx_bitwise(case) -> None: pytest.skip(f"S={S} C={C} runs the fallback program {fallback} on {arch}") device = _device() x, w, cu, dy = _inputs(S, C, docs) - dx, dw = cake_causal_conv1d_train_backward(x, w, cu, dy) + dx, dw = cake_causal_conv1d_det_train_backward(x, w, cu, dy) record = loader.MODULES[fallback] scalars = loader.plan( S, @@ -1311,13 +1315,13 @@ def test_routed_rows_match_fallback_dx_bitwise(case) -> None: def test_preallocated_outputs_are_written_in_place() -> None: S, C, docs = 1000, 512, "perf" x, w, cu, dy = _inputs(S, C, docs) - y_ref = cake_causal_conv1d_train_forward(x, w, cu) - dx_ref, dw_ref = cake_causal_conv1d_train_backward(x, w, cu, dy) + y_ref = cake_causal_conv1d_det_train_forward(x, w, cu) + dx_ref, dw_ref = cake_causal_conv1d_det_train_backward(x, w, cu, dy) out = torch.full((S, C), float("nan"), dtype=torch.bfloat16, device=_device()) dx = torch.full((S, C), float("nan"), dtype=torch.bfloat16, device=_device()) dw = torch.full((C, WIDTH), float("nan"), dtype=torch.float32, device=_device()) - assert cake_causal_conv1d_train_forward(x, w, cu, out=out) is out - got = cake_causal_conv1d_train_backward(x, w, cu, dy, dx=dx, dweight=dw) + assert cake_causal_conv1d_det_train_forward(x, w, cu, out=out) is out + got = cake_causal_conv1d_det_train_backward(x, w, cu, dy, dx=dx, dweight=dw) assert len(got) == 2 and got[0] is dx and got[1] is dw assert ( torch.equal(out, y_ref) and torch.equal(dx, dx_ref) and torch.equal(dw, dw_ref) @@ -1332,7 +1336,7 @@ def test_preallocated_bias_outputs_are_written_in_place() -> None: dx = torch.full((S, C), float("nan"), dtype=torch.bfloat16, device=_device()) dw = torch.full((C, WIDTH), float("nan"), dtype=torch.float32, device=_device()) db = torch.full((C,), float("nan"), dtype=torch.float32, device=_device()) - got = cake_causal_conv1d_train_backward( + got = cake_causal_conv1d_det_train_backward( x, w, cu, dy, bias=b, dx=dx, dweight=dw, dbias=db ) assert len(got) == 3 and got[0] is dx and got[1] is dw and got[2] is db @@ -1349,13 +1353,13 @@ def _capture_and_replay(cases, kind: str, *, key: int) -> int: captured = [] for S, C, docs in cases: x, w, cu, dy, b, policy = _route_inputs(S, C, docs, kind, key=key) - eager_y = cake_causal_conv1d_train_forward(x, w, cu, bias=b) + eager_y = cake_causal_conv1d_det_train_forward(x, w, cu, bias=b) eager = (eager_y, *_backward(x, w, cu, dy, b, policy)) y = torch.empty_like(x) dx = torch.empty_like(x) dw = torch.empty(C, WIDTH, dtype=w.dtype, device=device) db = None if b is None else torch.empty(C, dtype=b.dtype, device=device) - workspace = cake_causal_conv1d_train_backward_workspace(S, C, device) + workspace = cake_causal_conv1d_det_train_backward_workspace(S, C, device) def step( x=x, @@ -1370,8 +1374,8 @@ def step( db=db, ws=workspace, ): - cake_causal_conv1d_train_forward(x, w, cu, bias=b, out=y) - cake_causal_conv1d_train_backward( + cake_causal_conv1d_det_train_forward(x, w, cu, bias=b, out=y) + cake_causal_conv1d_det_train_backward( x, w, cu, @@ -1457,7 +1461,7 @@ def record_launches(launches, values, *, device_index): # The backward issues its pair through run_modules, the forward its program through run_module. monkeypatch.setattr(pkg, "run_module", record_launch) monkeypatch.setattr(pkg, "run_modules", record_launches) - cake_causal_conv1d_train_backward(x, w, cu, dy) + cake_causal_conv1d_det_train_backward(x, w, cu, dy) arch = _arch() main = loader.selected_modules(arch, "bwd", C=C, S=S)[0] record = loader.MODULES[main] @@ -1486,7 +1490,7 @@ def record_launches(launches, values, *, device_index): with pytest.raises(ValueError, match="needs the channel count"): loader.launch_grid(reduce, device_index=device.index) launched.clear() - cake_causal_conv1d_train_forward(x, w, cu) + cake_causal_conv1d_det_train_forward(x, w, cu) assert [name for name, _, _, _ in launched] == list( loader.selected_modules(arch, "fwd", C=C, S=S) ) @@ -1537,7 +1541,7 @@ def test_backward_keeps_its_speed_beside_a_resident_side_kernel( device = _device() S, C, docs = 261393, 512, "perf" x, w, cu, dy = _inputs(S, C, docs) - workspace = cake_causal_conv1d_train_backward_workspace(S, C, device) + workspace = cake_causal_conv1d_det_train_backward_workspace(S, C, device) dx = torch.empty_like(x) dw = torch.empty(C, WIDTH, dtype=torch.float32, device=device) side = torch.cuda.Stream(device=device) @@ -1553,7 +1557,7 @@ def timed(with_side_kernel: bool) -> float: ) # one 1-thread CTA, resident for milliseconds torch.cuda._sleep(1_000_000) # let it become resident before the backward start.record() - cake_causal_conv1d_train_backward( + cake_causal_conv1d_det_train_backward( x, w, cu, dy, workspace=workspace, dx=dx, dweight=dw ) end.record() @@ -1585,15 +1589,15 @@ def test_zero_tokens_return_empty_outputs() -> None: w = torch.zeros(C, WIDTH, dtype=torch.float32, device=_device()) b = torch.zeros(C, dtype=torch.bfloat16, device=_device()) cu = torch.zeros(2, dtype=torch.int32, device=_device()) - y = cake_causal_conv1d_train_forward(x, w, cu) + y = cake_causal_conv1d_det_train_forward(x, w, cu) assert tuple(y.shape) == (0, C) and y.dtype == torch.bfloat16 - dx, dw = cake_causal_conv1d_train_backward(x, w, cu, x) + dx, dw = cake_causal_conv1d_det_train_backward(x, w, cu, x) assert tuple(dx.shape) == (0, C) assert torch.equal(dw, torch.zeros_like(dw)) # With a bias the empty call returns three zero gradients in the parameters' dtypes, without a launch. - y = cake_causal_conv1d_train_forward(x, w.to(torch.bfloat16), cu, bias=b) + y = cake_causal_conv1d_det_train_forward(x, w.to(torch.bfloat16), cu, bias=b) assert tuple(y.shape) == (0, C) - dx, dw, db = cake_causal_conv1d_train_backward( + dx, dw, db = cake_causal_conv1d_det_train_backward( x, w, cu, x, bias=b, silu_grad_preactivation="fp32" ) assert tuple(dx.shape) == (0, C) and dw.dtype == torch.float32 @@ -1605,7 +1609,7 @@ def test_zero_tokens_return_empty_outputs() -> None: @pytest.mark.parametrize("kind", tuple(ROUTE_SETTINGS)) def test_autograd_end_to_end(kind) -> None: - """``cake_causal_conv1d`` returns the forward's bits and its backward the fused kernel's gradients, in the + """``cake_causal_conv1d_det`` returns the forward's bits and its backward the fused kernel's gradients, in the parameters' dtypes, under every route; a non-contiguous incoming gradient is materialized.""" _require_route(kind) @@ -1614,9 +1618,9 @@ def test_autograd_end_to_end(kind) -> None: x = x0.detach().clone().requires_grad_(True) w = w0.detach().clone().requires_grad_(True) b = None if b0 is None else b0.detach().clone().requires_grad_(True) - y = cake_causal_conv1d(x, w, cu, b, silu_grad_preactivation=policy) + y = cake_causal_conv1d_det(x, w, cu, b, silu_grad_preactivation=policy) assert torch.equal( - y.detach(), cake_causal_conv1d_train_forward(x0, w0, cu, bias=b0) + y.detach(), cake_causal_conv1d_det_train_forward(x0, w0, cu, bias=b0) ) y.backward(dy) dx_ref, dw_ref, db_ref = _backward(x0, w0, cu, dy, b0, policy) @@ -1627,7 +1631,7 @@ def test_autograd_end_to_end(kind) -> None: x = x0.detach().clone().requires_grad_(True) w = w0.detach().clone().requires_grad_(True) b = None if b0 is None else b0.detach().clone().requires_grad_(True) - cake_causal_conv1d(x, w, cu, b, silu_grad_preactivation=policy).backward( + cake_causal_conv1d_det(x, w, cu, b, silu_grad_preactivation=policy).backward( dy.t().contiguous().t() ) assert torch.equal(x.grad, dx_ref) and torch.equal(w.grad, dw_ref) @@ -1762,7 +1766,7 @@ def test_nan_poisoned_workspace_per_route(kind) -> None: _require_route(kind) for S, C, docs in ROUTE_CASES: x, w, cu, dy, b, policy = _route_inputs(S, C, docs, kind) - workspace = cake_causal_conv1d_train_backward_workspace(S, C, _device()) + workspace = cake_causal_conv1d_det_train_backward_workspace(S, C, _device()) workspace.view(torch.float32).fill_(float("nan")) _check( S, C, docs, x=x, w=w, cu=cu, dy=dy, b=b, policy=policy, workspace=workspace @@ -1805,8 +1809,8 @@ def test_dt_route_with_f32_taps_matches_default_route_bitwise(case) -> None: S, C, docs = case arch, device = _arch(), _device() x, w, cu, dy = _inputs(S, C, docs) - dx, dw = cake_causal_conv1d_train_backward(x, w, cu, dy) - dx_dt, dw_dt = cake_causal_conv1d_train_backward( + dx, dw = cake_causal_conv1d_det_train_backward(x, w, cu, dy) + dx_dt, dw_dt = cake_causal_conv1d_det_train_backward( x, w, cu, dy, dweight=torch.empty(C, WIDTH, dtype=BF16, device=device) ) assert dw_dt.dtype == BF16 @@ -1902,8 +1906,8 @@ def capacity_key(owner): ] assert list(owners) == expected_owners for C in (128, 136, 1280, 8192): - small = cake_causal_conv1d_train_backward_workspace_bytes(1, C, device) - large = cake_causal_conv1d_train_backward_workspace_bytes(262401, C, device) + small = cake_causal_conv1d_det_train_backward_workspace_bytes(1, C, device) + large = cake_causal_conv1d_det_train_backward_workspace_bytes(262401, C, device) assert small == large > 0, "the workspace is independent of S" layout = loader.workspace_layout(record, C, device_index=device.index) _, _, n_ch_tiles = loader.channel_split( @@ -1951,9 +1955,10 @@ def capacity_key(owner): loader.launch_grid(selected, device_index=device.index)[0], n_ch_tiles, ) - assert loader.partial_bytes( - selected, C, rows - ) <= cake_causal_conv1d_train_backward_workspace_bytes(1, C, device), ( + capacity = cake_causal_conv1d_det_train_backward_workspace_bytes( + 1, C, device + ) + assert loader.partial_bytes(selected, C, rows) <= capacity, ( kind, C, selected["library"], @@ -1976,16 +1981,16 @@ def capacity_key(owner): threads // (136 // vec), 1, ) - workspace = cake_causal_conv1d_train_backward_workspace(65536, 512, device) + workspace = cake_causal_conv1d_det_train_backward_workspace(65536, 512, device) assert ( workspace.dtype == torch.uint8 and workspace.numel() - == cake_causal_conv1d_train_backward_workspace_bytes(1, 512, device) + == cake_causal_conv1d_det_train_backward_workspace_bytes(1, 512, device) ) with pytest.raises(ValueError): - cake_causal_conv1d_train_backward_workspace_bytes(-1, 512, device) + cake_causal_conv1d_det_train_backward_workspace_bytes(-1, 512, device) with pytest.raises(ValueError): - cake_causal_conv1d_train_backward_workspace_bytes(1, 130, device) + cake_causal_conv1d_det_train_backward_workspace_bytes(1, 130, device) # Forward workspace rows: whole-row and multi-sub-tile tiles, the forward's fallback tiles (1344, 1600), a @@ -2032,17 +2037,17 @@ def test_forward_workspace_bytes() -> None: for C in (8, 128, 136, 1024, 1160, 1280, 1344, 1600, 2560, 3200, 8192): tiles = _forward_tiles(C) assert loader.forward_workspace_words(arch, C) == 2 * tiles, C - small = cake_causal_conv1d_train_forward_workspace_bytes(1, C, device) - large = cake_causal_conv1d_train_forward_workspace_bytes(262401, C, device) + small = cake_causal_conv1d_det_train_forward_workspace_bytes(1, C, device) + large = cake_causal_conv1d_det_train_forward_workspace_bytes(262401, C, device) assert small == large == 4 * 2 * tiles, C - workspace = cake_causal_conv1d_train_forward_workspace(65536, 2560, device) + workspace = cake_causal_conv1d_det_train_forward_workspace(65536, 2560, device) assert workspace.dtype == torch.int32 and workspace.device == device assert workspace.numel() == 2 * _forward_tiles(2560) and workspace.is_contiguous() assert int(torch.count_nonzero(workspace)) == 0 with pytest.raises(ValueError): - cake_causal_conv1d_train_forward_workspace_bytes(-1, 512, device) + cake_causal_conv1d_det_train_forward_workspace_bytes(-1, 512, device) with pytest.raises(ValueError): - cake_causal_conv1d_train_forward_workspace_bytes(1, 130, device) + cake_causal_conv1d_det_train_forward_workspace_bytes(1, 130, device) def test_forward_reuses_one_workspace_over_calls_and_routes() -> None: @@ -2050,11 +2055,11 @@ def test_forward_reuses_one_workspace_over_calls_and_routes() -> None: leaves the counters zero, and ``y`` is within the reference and bitwise the call without a workspace.""" device = _device() - workspace = cake_causal_conv1d_train_forward_workspace( + workspace = cake_causal_conv1d_det_train_forward_workspace( 1, max( FORWARD_WORKSPACE_CASES, - key=lambda case: cake_causal_conv1d_train_forward_workspace_bytes( + key=lambda case: cake_causal_conv1d_det_train_forward_workspace_bytes( 1, case[1], device ), )[1], @@ -2075,13 +2080,14 @@ def test_forward_reuses_one_workspace_over_calls_and_routes() -> None: for repeat in range(3): for kind, S, C, x, w, cu, b, expected in problems: label = f"repeat {repeat} {kind} S={S} C={C}" - y = cake_causal_conv1d_train_forward(x, w, cu, bias=b, workspace=workspace) + y = cake_causal_conv1d_det_train_forward( + x, w, cu, bias=b, workspace=workspace + ) torch.cuda.synchronize(device) assert int(torch.count_nonzero(workspace)) == 0, label _assert_within(y, expected, atol=Y_ATOL, rtol=Y_RTOL, label=f"y {label}") - assert torch.equal(y, cake_causal_conv1d_train_forward(x, w, cu, bias=b)), ( - label - ) + again = cake_causal_conv1d_det_train_forward(x, w, cu, bias=b) + assert torch.equal(y, again), label def test_forward_on_two_streams_with_two_workspaces() -> None: @@ -2091,12 +2097,12 @@ def test_forward_on_two_streams_with_two_workspaces() -> None: device = _device() S, C, docs = 131072, 1280, "perf" x, w, cu, _ = _inputs(S, C, docs, key=17) - expected = cake_causal_conv1d_train_forward(x, w, cu) + expected = cake_causal_conv1d_det_train_forward(x, w, cu) torch.cuda.synchronize(device) streams = [torch.cuda.Stream(device=device) for _ in range(2)] assert streams[0].cuda_stream != streams[1].cuda_stream workspaces = [ - cake_causal_conv1d_train_forward_workspace(S, C, device) for _ in streams + cake_causal_conv1d_det_train_forward_workspace(S, C, device) for _ in streams ] calls = 6 for owned in (True, False): @@ -2106,7 +2112,7 @@ def test_forward_on_two_streams_with_two_workspaces() -> None: for i in range(calls): for s, stream in enumerate(streams): with torch.cuda.stream(stream): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, @@ -2148,18 +2154,18 @@ def test_forward_workspace_cache_per_device_and_stream() -> None: stream.wait_stream(torch.cuda.current_stream(device)) other.wait_stream(torch.cuda.current_stream(device)) with torch.cuda.stream(stream): - y1 = cake_causal_conv1d_train_forward(*small) + y1 = cake_causal_conv1d_det_train_forward(*small) first = cache[key] assert first.dtype == torch.int32 and first.numel() == 2 * _forward_tiles(512) - y2 = cake_causal_conv1d_train_forward(*small) + y2 = cake_causal_conv1d_det_train_forward(*small) assert cache[key] is first - y3 = cake_causal_conv1d_train_forward(*large) + y3 = cake_causal_conv1d_det_train_forward(*large) grown = cache[key] assert grown is not first and grown.numel() == 2 * _forward_tiles(2560) - y4 = cake_causal_conv1d_train_forward(*small) + y4 = cake_causal_conv1d_det_train_forward(*small) assert cache[key] is grown with torch.cuda.stream(other): - y5 = cake_causal_conv1d_train_forward(*small) + y5 = cake_causal_conv1d_det_train_forward(*small) torch.cuda.synchronize(device) assert cache[other_key].data_ptr() != cache[key].data_ptr() assert int(torch.count_nonzero(cache[key])) == 0 @@ -2184,26 +2190,26 @@ def test_forward_cuda_graph_with_and_without_workspace() -> None: _require_route(settings) S, C, docs = 4097, 1344, (1000, 0, 3, 3094) x, w, cu, _, b, _ = _route_inputs(S, C, docs, settings, key=23) - eager = cake_causal_conv1d_train_forward(x, w, cu, bias=b) - workspace = cake_causal_conv1d_train_forward_workspace(S, C, device) + eager = cake_causal_conv1d_det_train_forward(x, w, cu, bias=b) + workspace = cake_causal_conv1d_det_train_forward_workspace(S, C, device) y_owned, y_own_buffer = torch.empty_like(x), torch.empty_like(x) side = torch.cuda.Stream(device=device) side.wait_stream(torch.cuda.current_stream(device)) with torch.cuda.stream(side): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, bias=b, out=y_owned, workspace=workspace ) - cake_causal_conv1d_train_forward(x, w, cu, bias=b, out=y_own_buffer) + cake_causal_conv1d_det_train_forward(x, w, cu, bias=b, out=y_own_buffer) torch.cuda.current_stream(device).wait_stream(side) torch.cuda.synchronize(device) before = {key: held.data_ptr() for key, held in cache.items()} owned, own_buffer = torch.cuda.CUDAGraph(), torch.cuda.CUDAGraph() with torch.cuda.graph(owned, stream=side): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, bias=b, out=y_owned, workspace=workspace ) with torch.cuda.graph(own_buffer, stream=side): - cake_causal_conv1d_train_forward(x, w, cu, bias=b, out=y_own_buffer) + cake_causal_conv1d_det_train_forward(x, w, cu, bias=b, out=y_own_buffer) assert {key: held.data_ptr() for key, held in cache.items()} == before, kind for replay in range(GRAPH_REPLAYS): y_owned.fill_(float("nan")) @@ -2232,34 +2238,34 @@ def test_forward_cuda_graph_with_and_without_workspace() -> None: def test_forward_workspace_validation() -> None: S, C = 4097, 1280 x, w, cu, _ = _inputs(S, C, (1000, 0, 3, 3094)) - words = cake_causal_conv1d_train_forward_workspace_bytes(S, C, _device()) // 4 + words = cake_causal_conv1d_det_train_forward_workspace_bytes(S, C, _device()) // 4 with pytest.raises(ValueError, match="workspace holds"): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, workspace=torch.zeros(words - 1, dtype=torch.int32, device=_device()), ) with pytest.raises(ValueError, match="torch.int32"): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, workspace=torch.zeros(4 * words, dtype=torch.uint8, device=_device()), ) with pytest.raises(ValueError, match="1-D"): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, workspace=torch.zeros(1, words, dtype=torch.int32, device=_device()), ) with pytest.raises(ValueError, match="must live on"): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, workspace=torch.zeros(words, dtype=torch.int32) ) with pytest.raises(TypeError): - cake_causal_conv1d_train_forward(x, w, cu, workspace=[0] * words) + cake_causal_conv1d_det_train_forward(x, w, cu, workspace=[0] * words) def test_argument_validation() -> None: @@ -2267,41 +2273,41 @@ def test_argument_validation() -> None: x, w, cu, dy = _inputs(S, C, None) b = _bias_input(C, F32) with pytest.raises(ValueError, match="multiple of 8"): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x[:, :130].contiguous(), w[:130].contiguous(), cu ) with pytest.raises(ValueError, match="width-4"): - cake_causal_conv1d_train_forward(x, w[:, :3].contiguous(), cu) + cake_causal_conv1d_det_train_forward(x, w[:, :3].contiguous(), cu) with pytest.raises(ValueError, match="width-4"): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, torch.cat([w, w[:, :1]], dim=1).contiguous(), cu ) with pytest.raises(ValueError, match="one row per channel"): - cake_causal_conv1d_train_forward(x, w[: C // 2].contiguous(), cu) + cake_causal_conv1d_det_train_forward(x, w[: C // 2].contiguous(), cu) with pytest.raises(ValueError, match="torch.int32"): - cake_causal_conv1d_train_forward(x, w, cu.long()) + cake_causal_conv1d_det_train_forward(x, w, cu.long()) with pytest.raises(ValueError, match="cu_seqlens"): - cake_causal_conv1d_train_forward(x, w, cu[:1]) + cake_causal_conv1d_det_train_forward(x, w, cu[:1]) with pytest.raises(ValueError, match="shape"): - cake_causal_conv1d_train_backward(x, w, cu, dy[: S - 1]) + cake_causal_conv1d_det_train_backward(x, w, cu, dy[: S - 1]) with pytest.raises(ValueError, match="contiguous"): - cake_causal_conv1d_train_forward(x.t().contiguous().t(), w, cu) + cake_causal_conv1d_det_train_forward(x.t().contiguous().t(), w, cu) with pytest.raises(ValueError, match="torch.bfloat16"): - cake_causal_conv1d_train_forward(x.float(), w, cu) + cake_causal_conv1d_det_train_forward(x.float(), w, cu) # Per buffer: the taps are FP32 or BF16 (the binding no longer checks the untyped tap / bias pointers, # so the host is the only guard), as are the bias and the gradient outs; the bias is one value per channel. with pytest.raises( ValueError, match="weight must be torch.float32 or torch.bfloat16" ): - cake_causal_conv1d_train_forward(x, w.half(), cu) + cake_causal_conv1d_det_train_forward(x, w.half(), cu) with pytest.raises( ValueError, match="weight must be torch.float32 or torch.bfloat16" ): - cake_causal_conv1d_train_backward(x, w.half(), cu, dy) + cake_causal_conv1d_det_train_backward(x, w.half(), cu, dy) with pytest.raises( ValueError, match="dweight must be torch.float32 or torch.bfloat16" ): - cake_causal_conv1d_train_backward( + cake_causal_conv1d_det_train_backward( x, w, cu, @@ -2311,21 +2317,21 @@ def test_argument_validation() -> None: with pytest.raises( ValueError, match="bias must be torch.float32 or torch.bfloat16" ): - cake_causal_conv1d_train_forward(x, w, cu, bias=b.half()) + cake_causal_conv1d_det_train_forward(x, w, cu, bias=b.half()) with pytest.raises( ValueError, match="bias must be torch.float32 or torch.bfloat16" ): - cake_causal_conv1d_train_backward(x, w, cu, dy, bias=b.half()) + cake_causal_conv1d_det_train_backward(x, w, cu, dy, bias=b.half()) with pytest.raises(ValueError, match=r"bias must be a \[C\]"): - cake_causal_conv1d_train_forward(x, w, cu, bias=b[: C // 2].contiguous()) + cake_causal_conv1d_det_train_forward(x, w, cu, bias=b[: C // 2].contiguous()) with pytest.raises(ValueError, match=r"bias must be a \[C\]"): - cake_causal_conv1d_train_forward(x, w, cu, bias=w[:, :1].contiguous()) + cake_causal_conv1d_det_train_forward(x, w, cu, bias=w[:, :1].contiguous()) with pytest.raises(ValueError, match="contiguous"): - cake_causal_conv1d_train_forward(x, w, cu, bias=w[:, 0]) + cake_causal_conv1d_det_train_forward(x, w, cu, bias=w[:, 0]) with pytest.raises( ValueError, match="dbias must be torch.float32 or torch.bfloat16" ): - cake_causal_conv1d_train_backward( + cake_causal_conv1d_det_train_backward( x, w, cu, @@ -2334,21 +2340,23 @@ def test_argument_validation() -> None: dbias=torch.empty(C, dtype=torch.float16, device=_device()), ) with pytest.raises(ValueError, match="dbias must have shape"): - cake_causal_conv1d_train_backward( + cake_causal_conv1d_det_train_backward( x, w, cu, dy, bias=b, dbias=torch.empty(C + 8, dtype=F32, device=_device()) ) with pytest.raises(ValueError, match="dbias"): - cake_causal_conv1d_train_backward(x, w, cu, dy, dbias=torch.empty_like(b)) + cake_causal_conv1d_det_train_backward(x, w, cu, dy, dbias=torch.empty_like(b)) with pytest.raises(ValueError, match="silu_grad_preactivation"): - cake_causal_conv1d_train_backward(x, w, cu, dy, silu_grad_preactivation="fp16") + cake_causal_conv1d_det_train_backward( + x, w, cu, dy, silu_grad_preactivation="fp16" + ) with pytest.raises(ValueError, match="silu_grad_preactivation"): - cake_causal_conv1d(x, w, cu, silu_grad_preactivation="bf32") + cake_causal_conv1d_det(x, w, cu, silu_grad_preactivation="bf32") with pytest.raises(ValueError, match="shape"): - cake_causal_conv1d_train_forward( + cake_causal_conv1d_det_train_forward( x, w, cu, out=torch.empty(S + 1, C, dtype=torch.bfloat16, device=_device()) ) with pytest.raises(ValueError, match="contiguous"): - cake_causal_conv1d_train_backward( + cake_causal_conv1d_det_train_backward( x, w, cu, @@ -2356,13 +2364,13 @@ def test_argument_validation() -> None: dx=torch.empty(C, S, dtype=torch.bfloat16, device=_device()).t(), ) with pytest.raises(ValueError, match="workspace holds"): - cake_causal_conv1d_train_backward( + cake_causal_conv1d_det_train_backward( x, w, cu, dy, workspace=torch.empty(16, dtype=torch.uint8, device=_device()) ) with pytest.raises(ValueError, match="CUDA"): - cake_causal_conv1d_train_forward(x.cpu(), w.cpu(), cu.cpu()) + cake_causal_conv1d_det_train_forward(x.cpu(), w.cpu(), cu.cpu()) with pytest.raises(TypeError): - cake_causal_conv1d_train_forward(x, w, [0, S]) + cake_causal_conv1d_det_train_forward(x, w, [0, S]) def test_every_route_ships_on_every_architecture() -> None: @@ -2396,7 +2404,7 @@ def test_registry_records_reference_the_build_manifest() -> None: assert manifest["schema"] == "cake-build/1" # The manifest names the package by its FlashInfer import path, also in a copy # vendored under another name (``pkg.__name__`` is then the copy's). - assert manifest["package"] == "flashinfer.cake_causal_conv1d_train_backend" + assert manifest["package"] == "flashinfer.cake_causal_conv1d_det_train_backend" libraries = cake_build.libraries_by_name() package = Path(cake_build.__file__).parent referenced = set() diff --git a/tests/norm/test_cake_causal_conv1d_det_train_aliases.py b/tests/norm/test_cake_causal_conv1d_det_train_aliases.py new file mode 100644 index 00000000000..024c6f86990 --- /dev/null +++ b/tests/norm/test_cake_causal_conv1d_det_train_aliases.py @@ -0,0 +1,221 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Compatibility names of the Cake causal conv1d + SiLU training package. + +Within each namespace (the ``flashinfer.cake_causal_conv1d_det_train`` facade, its +backend package and the ``flashinfer`` top level) every compatibility name is the +same object as its det name and stays out of ``__all__``. The module +``flashinfer.cake_causal_conv1d_train`` is the facade's module object, and the +submodules of ``flashinfer.cake_causal_conv1d_train_backend`` are the det +package's module objects. Both old module paths are also attributes of the +``flashinfer`` package after a plain ``import flashinfer``. Importing either facade +or either package, or reaching it as such an attribute, loads no KCP host module +under either name (checked in fresh interpreters by module file, not by name). +""" + +from __future__ import annotations + +import importlib +import subprocess +import sys +from pathlib import Path + +import pytest + +import flashinfer +import flashinfer.cake_causal_conv1d_det_train as facade +import flashinfer.cake_causal_conv1d_det_train_backend as backend + +pytestmark = pytest.mark.cpu_only + +OLD_FACADE = "flashinfer.cake_causal_conv1d_train" +NEW_FACADE = "flashinfer.cake_causal_conv1d_det_train" +OLD_PACKAGE = "flashinfer.cake_causal_conv1d_train_backend" +NEW_PACKAGE = "flashinfer.cake_causal_conv1d_det_train_backend" +SUBMODULES = ("cake_build", "cake_jit", "kcp", "kcp_jit") + +# (compatibility name, det name) of the facade and the backend package. +PAIRS = ( + ("CakeCausalConv1dFunction", "CakeCausalConv1dDetFunction"), + ("cake_causal_conv1d", "cake_causal_conv1d_det"), + ("cake_causal_conv1d_train_forward", "cake_causal_conv1d_det_train_forward"), + ("cake_causal_conv1d_train_backward", "cake_causal_conv1d_det_train_backward"), + ( + "cake_causal_conv1d_train_forward_workspace", + "cake_causal_conv1d_det_train_forward_workspace", + ), + ( + "cake_causal_conv1d_train_forward_workspace_bytes", + "cake_causal_conv1d_det_train_forward_workspace_bytes", + ), + ( + "cake_causal_conv1d_train_backward_workspace", + "cake_causal_conv1d_det_train_backward_workspace", + ), + ( + "cake_causal_conv1d_train_backward_workspace_bytes", + "cake_causal_conv1d_det_train_backward_workspace_bytes", + ), + ( + "is_cake_causal_conv1d_train_supported", + "is_cake_causal_conv1d_det_train_supported", + ), +) +NAMES = dict(PAIRS) +# The names the ``flashinfer`` top level exports. +TOP_LEVEL_NAMES = ( + "CakeCausalConv1dFunction", + "cake_causal_conv1d", + "cake_causal_conv1d_train_forward", + "cake_causal_conv1d_train_backward", + "cake_causal_conv1d_train_forward_workspace_bytes", + "cake_causal_conv1d_train_backward_workspace_bytes", +) +NAMESPACES = {"facade": facade, "backend": backend} + + +@pytest.mark.parametrize("namespace", sorted(NAMESPACES)) +@pytest.mark.parametrize("old", sorted(NAMES)) +def test_compatibility_name_is_the_det_object(namespace: str, old: str) -> None: + module = NAMESPACES[namespace] + new = NAMES[old] + assert getattr(module, old) is getattr(module, new), (namespace, old, new) + assert new in module.__all__, (namespace, new) + assert old not in module.__all__, (namespace, old) + + +@pytest.mark.parametrize("old", TOP_LEVEL_NAMES) +def test_top_level_compatibility_name_is_the_det_object(old: str) -> None: + new = NAMES[old] + assert getattr(flashinfer, old) is getattr(flashinfer, new), (old, new) + assert getattr(flashinfer, new) is getattr(facade, new), new + + +def test_det_entry_points_are_the_decorated_defs() -> None: + # One ``flashinfer_api`` wrapper per entry point: the compatibility names add none. + for old in ( + "cake_causal_conv1d", + "cake_causal_conv1d_train_forward", + "cake_causal_conv1d_train_backward", + ): + assert getattr(facade, old).__name__ == NAMES[old], old + det_class = backend.CakeCausalConv1dDetFunction + assert det_class.__name__ == "CakeCausalConv1dDetFunction" + + +def test_old_facade_is_the_det_module() -> None: + assert importlib.import_module(OLD_FACADE) is facade + assert sys.modules[OLD_FACADE] is sys.modules[NEW_FACADE] + + +@pytest.mark.parametrize("submodule", SUBMODULES) +def test_old_package_submodule_is_the_det_module(submodule: str) -> None: + old = importlib.import_module(f"{OLD_PACKAGE}.{submodule}") + new = importlib.import_module(f"{NEW_PACKAGE}.{submodule}") + assert old is new, submodule + assert sys.modules[f"{OLD_PACKAGE}.{submodule}"] is new, submodule + assert Path(new.__file__).parent == Path(backend.__file__).parent, submodule + + +def test_old_package_forwards_to_the_det_package() -> None: + old = importlib.import_module(OLD_PACKAGE) + for name in (*NAMES, *NAMES.values(), "ARCHES", "_FORWARD_WORKSPACES"): + assert getattr(old, name) is getattr(backend, name), name + assert old.__all__ == backend.__all__ + # The compatibility package holds Python shims only: no csrc, manifest or libraries. + old_dir = Path(old.__file__).parent + files = sorted(p.name for p in old_dir.iterdir() if p.name != "__pycache__") + assert files == sorted(["__init__.py", *(f"{m}.py" for m in SUBMODULES)]) + + +# The old module paths reached as attributes of ``flashinfer`` with no explicit import. +_ATTRIBUTE_ACCESS = """ +import sys + +import flashinfer + +old_facade = flashinfer.{old_facade} +assert old_facade is flashinfer.{new_facade} is sys.modules["flashinfer.{new_facade}"] +old_package = flashinfer.{old_package} +assert old_package is sys.modules["flashinfer.{old_package}"] +new_package = flashinfer.{new_package} +assert old_package.{old_name} is new_package.{new_name} +""" + +# Every import form of the facade and the package, each in a fresh interpreter. +IMPORTS = { + "flashinfer_attributes": _ATTRIBUTE_ACCESS.format( + old_facade=OLD_FACADE.rpartition(".")[2], + new_facade=NEW_FACADE.rpartition(".")[2], + old_package=OLD_PACKAGE.rpartition(".")[2], + new_package=NEW_PACKAGE.rpartition(".")[2], + old_name=PAIRS[2][0], + new_name=PAIRS[2][1], + ), + "flashinfer": "import flashinfer", + "old_facade": f"import {OLD_FACADE}", + "new_facade": f"import {NEW_FACADE}", + "old_package": f"import {OLD_PACKAGE} as p; p.{PAIRS[2][0]}", + "new_package": f"import {NEW_PACKAGE}", + "old_package_names": f"from {OLD_PACKAGE} import {PAIRS[2][1]}, cake_jit", +} + +_LAZY_CHECK = """ +import sys +from pathlib import Path + +import {new_package} as package + +kcp_files = {{(Path(package.__file__).parent / f).resolve() for f in ("kcp.py", "kcp_jit.py")}} +kcp_names = [p + "." + s for p in ({old_package!r}, {new_package!r}) for s in ("kcp", "kcp_jit")] +loaded = sorted( + name + for name, module in list(sys.modules.items()) + if name in kcp_names + or (getattr(module, "__file__", None) and Path(module.__file__).resolve() in kcp_files) +) +assert not loaded, loaded +print("KCP_NOT_LOADED") +from {old_facade} import KcpRecomputeHandle as a +from {new_facade} import KcpRecomputeHandle as b +from {old_package}.kcp import KcpRecomputeHandle as c +assert a is b is c +assert sys.modules[{old_package!r} + ".kcp"] is sys.modules[{new_package!r} + ".kcp"] +print("KCP_ALIASED") +""" + + +@pytest.mark.parametrize("form", sorted(IMPORTS)) +def test_imports_do_not_load_kcp(form: str) -> None: + code = ( + IMPORTS[form] + + "\n" + + _LAZY_CHECK.format( + old_facade=OLD_FACADE, + new_facade=NEW_FACADE, + old_package=OLD_PACKAGE, + new_package=NEW_PACKAGE, + ) + ) + r = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, timeout=600 + ) + assert r.returncode == 0, (form, r.stdout, r.stderr) + assert "KCP_NOT_LOADED" in r.stdout and "KCP_ALIASED" in r.stdout, ( + form, + r.stdout, + r.stderr, + ) diff --git a/tests/norm/test_cake_causal_conv1d_train_api.py b/tests/norm/test_cake_causal_conv1d_det_train_api.py similarity index 57% rename from tests/norm/test_cake_causal_conv1d_train_api.py rename to tests/norm/test_cake_causal_conv1d_det_train_api.py index efc5d9e3911..484739bafb8 100644 --- a/tests/norm/test_cake_causal_conv1d_train_api.py +++ b/tests/norm/test_cake_causal_conv1d_det_train_api.py @@ -15,12 +15,12 @@ Tests of the flashinfer public surface of the Cake packed width-4 depthwise causal convolution + SiLU training kernels (SM100 / SM103 / SM107): the -``flashinfer.cake_causal_conv1d_train`` facade and its re-export from the +``flashinfer.cake_causal_conv1d_det_train`` facade and its re-export from the ``flashinfer`` top-level package. The facade's three entry points are -pass-throughs to ``flashinfer.cake_causal_conv1d_train_backend``; this module +pass-throughs to ``flashinfer.cake_causal_conv1d_det_train_backend``; this module checks that identity (same objects at the top level, the backend's exact signatures, same bits as the backend) and leaves the numerical contract to -``test_cake_causal_conv1d_train.py``, which imports only the backend package. +``test_cake_causal_conv1d_det_train.py``, which imports only the backend package. """ from __future__ import annotations @@ -31,19 +31,19 @@ import torch import flashinfer -import flashinfer.cake_causal_conv1d_train -import flashinfer.cake_causal_conv1d_train_backend as backend -from flashinfer.cake_causal_conv1d_train import ( - CakeCausalConv1dFunction, - cake_causal_conv1d, - cake_causal_conv1d_train_backward, - cake_causal_conv1d_train_backward_workspace_bytes, - cake_causal_conv1d_train_forward, - cake_causal_conv1d_train_forward_workspace, - cake_causal_conv1d_train_forward_workspace_bytes, - is_cake_causal_conv1d_train_supported, +import flashinfer.cake_causal_conv1d_det_train +import flashinfer.cake_causal_conv1d_det_train_backend as backend +from flashinfer.cake_causal_conv1d_det_train import ( + CakeCausalConv1dDetFunction, + cake_causal_conv1d_det, + cake_causal_conv1d_det_train_backward, + cake_causal_conv1d_det_train_backward_workspace_bytes, + cake_causal_conv1d_det_train_forward, + cake_causal_conv1d_det_train_forward_workspace, + cake_causal_conv1d_det_train_forward_workspace_bytes, + is_cake_causal_conv1d_det_train_supported, ) -from flashinfer.cake_causal_conv1d_train_backend import cake_jit as loader +from flashinfer.cake_causal_conv1d_det_train_backend import cake_jit as loader from flashinfer.jit.env import has_flashinfer_jit_cache from flashinfer.utils import get_compute_capability @@ -59,13 +59,13 @@ def _supported() -> bool: return ( torch.cuda.is_available() and get_compute_capability(_device()) in SUPPORTED_CAPABILITIES - and is_cake_causal_conv1d_train_supported(_device()) + and is_cake_causal_conv1d_det_train_supported(_device()) ) pytestmark = pytest.mark.skipif( not _supported(), - reason="cake_causal_conv1d_train needs an SM100 / SM103 / SM107 device with exported routes", + reason="cake_causal_conv1d_det_train needs an SM100 / SM103 / SM107 device with exported routes", ) @@ -78,28 +78,29 @@ def warmup_jit(): def test_function_signature_is_exported() -> None: - assert flashinfer.cake_causal_conv1d is cake_causal_conv1d + assert flashinfer.cake_causal_conv1d_det is cake_causal_conv1d_det assert ( - flashinfer.cake_causal_conv1d_train_forward is cake_causal_conv1d_train_forward + flashinfer.cake_causal_conv1d_det_train_forward + is cake_causal_conv1d_det_train_forward ) assert ( - flashinfer.cake_causal_conv1d_train_backward - is cake_causal_conv1d_train_backward + flashinfer.cake_causal_conv1d_det_train_backward + is cake_causal_conv1d_det_train_backward ) assert ( - flashinfer.cake_causal_conv1d_train_backward_workspace_bytes - is cake_causal_conv1d_train_backward_workspace_bytes + flashinfer.cake_causal_conv1d_det_train_backward_workspace_bytes + is cake_causal_conv1d_det_train_backward_workspace_bytes ) assert ( - flashinfer.cake_causal_conv1d_train_forward_workspace_bytes - is cake_causal_conv1d_train_forward_workspace_bytes - is backend.cake_causal_conv1d_train_forward_workspace_bytes + flashinfer.cake_causal_conv1d_det_train_forward_workspace_bytes + is cake_causal_conv1d_det_train_forward_workspace_bytes + is backend.cake_causal_conv1d_det_train_forward_workspace_bytes ) assert ( - cake_causal_conv1d_train_forward_workspace - is backend.cake_causal_conv1d_train_forward_workspace + cake_causal_conv1d_det_train_forward_workspace + is backend.cake_causal_conv1d_det_train_forward_workspace ) - assert flashinfer.CakeCausalConv1dFunction is CakeCausalConv1dFunction + assert flashinfer.CakeCausalConv1dDetFunction is CakeCausalConv1dDetFunction def test_facade_signatures_mirror_the_backend() -> None: @@ -107,19 +108,21 @@ def test_facade_signatures_mirror_the_backend() -> None: spelled identically on the ``@flashinfer_api`` facade, which public-API tooling reads.""" for name in ( - "cake_causal_conv1d_train_forward", - "cake_causal_conv1d_train_backward", - "cake_causal_conv1d", + "cake_causal_conv1d_det_train_forward", + "cake_causal_conv1d_det_train_backward", + "cake_causal_conv1d_det", ): - facade = inspect.signature(getattr(flashinfer.cake_causal_conv1d_train, name)) + facade = inspect.signature( + getattr(flashinfer.cake_causal_conv1d_det_train, name) + ) backend_signature = inspect.signature(getattr(backend, name)) assert facade == backend_signature, (name, facade, backend_signature) - backward = inspect.signature(cake_causal_conv1d_train_backward).parameters + backward = inspect.signature(cake_causal_conv1d_det_train_backward).parameters assert backward["bias"].default is None assert backward["silu_grad_preactivation"].default == "bf16" assert backward["dbias"].default is None assert backward["bias"].kind is inspect.Parameter.KEYWORD_ONLY - autograd = inspect.signature(cake_causal_conv1d).parameters + autograd = inspect.signature(cake_causal_conv1d_det).parameters assert autograd["bias"].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD assert autograd["silu_grad_preactivation"].kind is inspect.Parameter.KEYWORD_ONLY assert backend.SILU_GRAD_PREACTIVATIONS == ("bf16", "fp32") @@ -141,16 +144,16 @@ def test_facade_entry_points_match_the_backend() -> None: C, 4, dtype=torch.float32, device=_device(), generator=generator ) cu = torch.tensor([0, 369, 659, 868, S], dtype=torch.int32, device=_device()) - y = cake_causal_conv1d_train_forward(x, w, cu) - assert torch.equal(y, backend.cake_causal_conv1d_train_forward(x, w, cu)) - workspace = cake_causal_conv1d_train_forward_workspace(S, C, _device()) + y = cake_causal_conv1d_det_train_forward(x, w, cu) + assert torch.equal(y, backend.cake_causal_conv1d_det_train_forward(x, w, cu)) + workspace = cake_causal_conv1d_det_train_forward_workspace(S, C, _device()) assert torch.equal( - y, cake_causal_conv1d_train_forward(x, w, cu, workspace=workspace) + y, cake_causal_conv1d_det_train_forward(x, w, cu, workspace=workspace) ) assert int(torch.count_nonzero(workspace)) == 0 - dx, dw = cake_causal_conv1d_train_backward(x, w, cu, dy) - dx_backend, dw_backend = backend.cake_causal_conv1d_train_backward(x, w, cu, dy) + dx, dw = cake_causal_conv1d_det_train_backward(x, w, cu, dy) + dx_backend, dw_backend = backend.cake_causal_conv1d_det_train_backward(x, w, cu, dy) assert torch.equal(dx, dx_backend) and torch.equal(dw, dw_backend) assert torch.equal( - cake_causal_conv1d(x, w, cu), backend.cake_causal_conv1d(x, w, cu) + cake_causal_conv1d_det(x, w, cu), backend.cake_causal_conv1d_det(x, w, cu) ) diff --git a/tests/test_cake_causal_conv1d_train_kcp.py b/tests/test_cake_causal_conv1d_det_train_kcp.py similarity index 85% rename from tests/test_cake_causal_conv1d_train_kcp.py rename to tests/test_cake_causal_conv1d_det_train_kcp.py index 2f620a30ae4..f7c2751b616 100644 --- a/tests/test_cake_causal_conv1d_train_kcp.py +++ b/tests/test_cake_causal_conv1d_det_train_kcp.py @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. -Single-GPU tests of ``cake_causal_conv1d_kcp`` (FlashInfer issue #6919): one process, world size 1, no process +Single-GPU tests of ``cake_causal_conv1d_det_kcp`` (FlashInfer issue #6919): one process, world size 1, no process group. The rank holds every fragment of a partitioned problem (the halo exchange of one rank is the identity), so the history rows, the backward's head rows and the tail fold all run through the host flow. Criteria: @@ -41,15 +41,15 @@ import pytest import torch -from flashinfer.cake_causal_conv1d_train import ( - cake_causal_conv1d_train_backward, - cake_causal_conv1d_train_forward, - is_cake_causal_conv1d_train_supported, +from flashinfer.cake_causal_conv1d_det_train import ( + cake_causal_conv1d_det_train_backward, + cake_causal_conv1d_det_train_forward, + is_cake_causal_conv1d_det_train_supported, ) -from flashinfer.cake_causal_conv1d_train_backend.kcp import ( - cake_causal_conv1d_kcp, - cake_causal_conv1d_kcp_backward, - cake_causal_conv1d_kcp_forward, +from flashinfer.cake_causal_conv1d_det_train_backend.kcp import ( + cake_causal_conv1d_det_kcp, + cake_causal_conv1d_det_backward_kcp, + cake_causal_conv1d_det_forward_kcp, ) WIDTH = 4 @@ -63,8 +63,8 @@ pytestmark = pytest.mark.skipif( not torch.cuda.is_available() - or not is_cake_causal_conv1d_train_supported(torch.device("cuda")), - reason="cake_causal_conv1d_train needs a supported CUDA device (SM100 / SM103 / SM107)", + or not is_cake_causal_conv1d_det_train_supported(torch.device("cuda")), + reason="cake_causal_conv1d_det_train needs a supported CUDA device (SM100 / SM103 / SM107)", ) # Rows: (document lengths, fragment lengths over the packed sequence); each fragment is cut at document @@ -344,8 +344,8 @@ def _bits(t): def test_forward_matches_unpartitioned_bitwise(name, lens, cuts, C): device = torch.device("cuda") p = _problem(lens, cuts, C, device) - y_local, hist = cake_causal_conv1d_kcp_forward(p.x_local, p.w, p.ctx) - y_global = cake_causal_conv1d_train_forward(p.x, p.w, p.cu) + y_local, hist = cake_causal_conv1d_det_forward_kcp(p.x_local, p.w, p.ctx) + y_global = cake_causal_conv1d_det_train_forward(p.x, p.w, p.cu) torch.cuda.synchronize() assert hist.shape[0] == max(1, p.ctx.halo.n_hist_rows) assert not y_local.isnan().any() @@ -366,8 +366,8 @@ def test_backward_matches_reference_and_folds_tails_bitwise( ): device = torch.device("cuda") p = _problem(lens, cuts, C, device) - _, hist = cake_causal_conv1d_kcp_forward(p.x_local, p.w, p.ctx) - dx_local, dw_local = cake_causal_conv1d_kcp_backward( + _, hist = cake_causal_conv1d_det_forward_kcp(p.x_local, p.w, p.ctx) + dx_local, dw_local = cake_causal_conv1d_det_backward_kcp( p.x_local, p.w, p.ctx, p.dy_local, hist ) torch.cuda.synchronize() @@ -382,7 +382,7 @@ def test_backward_matches_reference_and_folds_tails_bitwise( ) # (b) dx against the shipped zero-history backward (policy A) on the unpartitioned documents: the elementwise # clause for a context-parallel result against the no-CP kernel; the measured difference is recorded. - dx_nocp, _ = cake_causal_conv1d_train_backward(p.x, p.w, p.cu, p.dy) + dx_nocp, _ = cake_causal_conv1d_det_train_backward(p.x, p.w, p.cu, p.dy) stats = _assert_within( dx_local, dx_nocp[p.l2g], @@ -395,8 +395,8 @@ def test_backward_matches_reference_and_folds_tails_bitwise( # (c) The unpartitioned problem through the same host flow (one fragment per document, no halo rows): the fold # of the received dpre rows reproduces its dx bit for bit. whole = _problem(lens, tuple(lens), C, device) - _, hist_w = cake_causal_conv1d_kcp_forward(whole.x_local, whole.w, whole.ctx) - dx_whole, _ = cake_causal_conv1d_kcp_backward( + _, hist_w = cake_causal_conv1d_det_forward_kcp(whole.x_local, whole.w, whole.ctx) + dx_whole, _ = cake_causal_conv1d_det_backward_kcp( whole.x_local, whole.w, whole.ctx, whole.dy_local, hist_w ) torch.cuda.synchronize() @@ -414,11 +414,11 @@ def test_autograd_function_and_repeat_are_bitwise(name, lens, cuts, C, record_pr for _ in range(2): x = p.x_local.clone().requires_grad_(True) w = p.w.clone().requires_grad_(True) - y = cake_causal_conv1d_kcp(x, w, p.ctx) + y = cake_causal_conv1d_det_kcp(x, w, p.ctx) y.backward(p.dy_local) runs.append((y.detach(), x.grad, w.grad)) - y_direct, hist = cake_causal_conv1d_kcp_forward(p.x_local, p.w, p.ctx) - dx_direct, dw_direct = cake_causal_conv1d_kcp_backward( + y_direct, hist = cake_causal_conv1d_det_forward_kcp(p.x_local, p.w, p.ctx) + dx_direct, dw_direct = cake_causal_conv1d_det_backward_kcp( p.x_local, p.w, p.ctx, p.dy_local, hist ) torch.cuda.synchronize() @@ -441,7 +441,7 @@ def test_autograd_function_and_repeat_are_bitwise(name, lens, cuts, C, record_pr _assert_within(dw1, dw_ref, atol=DW_ATOL, rtol=DW_RTOL, label="dw vs FP64 oracle") # (b) dx against the shipped zero-history backward (policy A) on the unpartitioned documents under the no-CP # clause; the two policies' difference is recorded for the report (it is not a bitwise relation). - dx_nocp, _ = cake_causal_conv1d_train_backward(p.x, p.w, p.cu, p.dy) + dx_nocp, _ = cake_causal_conv1d_det_train_backward(p.x, p.w, p.cu, p.dy) stats = _assert_within( dx1, dx_nocp[p.l2g], atol=NOCP_ATOL, rtol=NOCP_RTOL, label="dx vs no-CP kernel" ) @@ -455,8 +455,8 @@ def test_empty_rank_returns_empty_outputs(): ctx, S, _, _ = _context((0,), (0,), device) x = torch.empty(0, C, dtype=torch.bfloat16, device=device) w = torch.randn(C, WIDTH, device=device) - y, hist = cake_causal_conv1d_kcp_forward(x, w, ctx) - dx, dw = cake_causal_conv1d_kcp_backward(x, w, ctx, torch.empty_like(x), hist) + y, hist = cake_causal_conv1d_det_forward_kcp(x, w, ctx) + dx, dw = cake_causal_conv1d_det_backward_kcp(x, w, ctx, torch.empty_like(x), hist) assert y.shape == (0, C) and dx.shape == (0, C) assert torch.equal(dw, torch.zeros_like(dw)) @@ -465,7 +465,7 @@ def test_rejects_mismatched_context(): device = torch.device("cuda") p = _problem((5000, 1003), (1, 2, 4997, 1003), 512, device) with pytest.raises(ValueError): - cake_causal_conv1d_kcp_forward(p.x_local[:-1].contiguous(), p.w, p.ctx) + cake_causal_conv1d_det_forward_kcp(p.x_local[:-1].contiguous(), p.w, p.ctx) def test_kcp_keeps_fp32_taps_and_no_bias(): @@ -475,9 +475,9 @@ def test_kcp_keeps_fp32_taps_and_no_bias(): device = torch.device("cuda") p = _problem((5000, 1003), (1, 2, 4997, 1003), 512, device) with pytest.raises(ValueError, match="weight must be torch.float32"): - cake_causal_conv1d_kcp_forward(p.x_local, p.w.to(torch.bfloat16), p.ctx) + cake_causal_conv1d_det_forward_kcp(p.x_local, p.w.to(torch.bfloat16), p.ctx) with pytest.raises(ValueError, match="weight must be torch.float32"): - cake_causal_conv1d_kcp_backward( + cake_causal_conv1d_det_backward_kcp( p.x_local, p.w.to(torch.bfloat16), p.ctx, @@ -486,9 +486,9 @@ def test_kcp_keeps_fp32_taps_and_no_bias(): ) bias = torch.zeros(p.C, device=device) with pytest.raises(TypeError): - cake_causal_conv1d_kcp_forward(p.x_local, p.w, p.ctx, bias=bias) + cake_causal_conv1d_det_forward_kcp(p.x_local, p.w, p.ctx, bias=bias) with pytest.raises(TypeError): - cake_causal_conv1d_kcp(p.x_local, p.w, p.ctx, bias) + cake_causal_conv1d_det_kcp(p.x_local, p.w, p.ctx, bias) def test_kcp_plans_with_the_legacy_tile_rule(monkeypatch): @@ -496,8 +496,8 @@ def test_kcp_plans_with_the_legacy_tile_rule(monkeypatch): n_ch_tiles) = (128, 1, 2)`` (the training routes tile it as (32, 4, 5)), and the forward's ``y`` stays bitwise the unpartitioned forward's.""" - from flashinfer.cake_causal_conv1d_train_backend import cake_jit - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer.cake_causal_conv1d_det_train_backend import cake_jit + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp geometry = {"threads": 128, "vec": 8} legacy = cake_jit.plan(5000, 1280, 2, geometry=geometry, legacy=True) @@ -514,10 +514,12 @@ def recording_plan(*args, **kwargs): monkeypatch.setattr(conv_kcp, "plan", recording_plan) device = torch.device("cuda") p = _problem((5000, 1003), (1, 2, 4997, 1003), 1280, device) - y, hist = cake_causal_conv1d_kcp_forward(p.x_local, p.w, p.ctx) - dx, dw = cake_causal_conv1d_kcp_backward(p.x_local, p.w, p.ctx, p.dy_local, hist) + y, hist = cake_causal_conv1d_det_forward_kcp(p.x_local, p.w, p.ctx) + dx, dw = cake_causal_conv1d_det_backward_kcp( + p.x_local, p.w, p.ctx, p.dy_local, hist + ) assert len(calls) == 2 and all(call.get("legacy") is True for call in calls), calls - y_full = cake_causal_conv1d_train_forward(p.x, p.w, p.cu) + y_full = cake_causal_conv1d_det_train_forward(p.x, p.w, p.cu) assert torch.equal(y, y_full[p.l2g]) assert torch.isfinite(dw).all() and dx.shape == p.x_local.shape @@ -529,14 +531,14 @@ def test_graph_replay_knob_is_bitwise(monkeypatch, form) -> None: the second is captured and replayed once, later calls replay the graph; ``y`` / ``hist`` / ``dx`` / ``dweight`` are bitwise the knob-off results on every call. The backward's key holds the forward's ``hist`` pointer, so its capture follows the forward's by one call. Another input pointer is another key (eager, still bitwise).""" - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp device = torch.device("cuda") p = _problem((5000, 1003), (1, 2, 4997, 1003), 512, device) x, w, ctx = p.x_local, p.w, p.ctx dy = torch.randn(x.shape, device=device).to(torch.bfloat16) - y0, hist0 = cake_causal_conv1d_kcp_forward(x, w, ctx) - dx0, dw0 = cake_causal_conv1d_kcp_backward(x, w, ctx, dy, hist0) + y0, hist0 = cake_causal_conv1d_det_forward_kcp(x, w, ctx) + dx0, dw0 = cake_causal_conv1d_det_backward_kcp(x, w, ctx, dy, hist0) kw = {} if form == "knob": monkeypatch.setenv(conv_kcp.TEST_KNOBS_ENV, "1") @@ -546,10 +548,10 @@ def test_graph_replay_knob_is_bitwise(monkeypatch, form) -> None: try: events = [] for _ in range(4): - y, hist = cake_causal_conv1d_kcp_forward(x, w, ctx, **kw) + y, hist = cake_causal_conv1d_det_forward_kcp(x, w, ctx, **kw) events.append(conv_kcp.LAST_CALL["fwd"]["graph_event"]) assert torch.equal(y, y0) and torch.equal(hist, hist0) - dx, dw = cake_causal_conv1d_kcp_backward(x, w, ctx, dy, hist, **kw) + dx, dw = cake_causal_conv1d_det_backward_kcp(x, w, ctx, dy, hist, **kw) events.append(conv_kcp.LAST_CALL["bwd"]["graph_event"]) assert torch.equal(dx, dx0) and torch.equal(dw, dw0) assert events == [ @@ -563,14 +565,14 @@ def test_graph_replay_knob_is_bitwise(monkeypatch, form) -> None: "replay", ], events assert conv_kcp.GRAPH_STATS["fallback"] == 0 - y, hist = cake_causal_conv1d_kcp_forward(x.clone(), w, ctx, **kw) + y, hist = cake_causal_conv1d_det_forward_kcp(x.clone(), w, ctx, **kw) assert conv_kcp.LAST_CALL["fwd"]["graph_event"] == "eager" assert torch.equal(y, y0) and torch.equal(hist, hist0) if ( form == "keyword" ): # the autograd entry carries the keyword to the forward and to the backward xg = x.clone().requires_grad_(True) - y = conv_kcp.cake_causal_conv1d_kcp(xg, w, ctx, graph="replay") + y = conv_kcp.cake_causal_conv1d_det_kcp(xg, w, ctx, graph="replay") assert conv_kcp.LAST_CALL["fwd"]["graph"] == "replay" assert torch.equal(y.detach(), y0) y.backward(dy) @@ -579,7 +581,7 @@ def test_graph_replay_knob_is_bitwise(monkeypatch, form) -> None: finally: conv_kcp._set_test_knobs(None) conv_kcp._GRAPHS.clear() - y, hist = cake_causal_conv1d_kcp_forward(x, w, ctx) + y, hist = cake_causal_conv1d_det_forward_kcp(x, w, ctx) assert torch.equal(y, y0) and torch.equal(hist, hist0) @@ -587,7 +589,7 @@ def test_graph_replay_knob_is_bitwise(monkeypatch, form) -> None: # Recompute handle (issues #7815 / #7818): torch's checkpoint recompute binds the history pool the original forward # stored instead of gathering and exchanging it again. One rank here (the exchange of one rank is the identity, so # the module's ``_exchange`` calls are counted instead of collectives: one in the original forward, one in the -# backward, none in a recompute with a handle); the NCCL rows are in test_cake_causal_conv1d_train_kcp_dist.py. +# backward, none in a recompute with a handle); the NCCL rows are in test_cake_causal_conv1d_det_train_kcp_dist.py. class _ExchangeCounter: @@ -626,19 +628,19 @@ def table(self) -> dict[str, int]: def _ckpt_run(p, *, handle=None, reentrant=False, ckpt=True, counter=None, twice=False): - """``y`` / ``dx`` / ``dweight`` of ``cake_causal_conv1d_kcp`` (the facade) on the problem ``p``, directly or under a + """``y`` / ``dx`` / ``dweight`` of ``cake_causal_conv1d_det_kcp`` (the facade) on the problem ``p``, directly or under a checkpoint, with ``recompute_handle=handle``.""" import contextlib from torch.utils.checkpoint import checkpoint - from flashinfer import cake_causal_conv1d_train as facade + from flashinfer import cake_causal_conv1d_det_train_kcp as facade x = p.x_local.clone().requires_grad_(True) w = p.w.clone().requires_grad_(True) def fn(x_, w_, h_): - return facade.cake_causal_conv1d_kcp(x_, w_, p.ctx, recompute_handle=h_) + return facade.cake_causal_conv1d_det_kcp(x_, w_, p.ctx, recompute_handle=h_) if not ckpt: y = fn(x, w, handle) @@ -675,12 +677,12 @@ def test_recompute_handle_is_bitwise_and_exchange_free( ``y`` / ``dx`` / ``dweight`` are bitwise the plain checkpoint's, the direct pass's and the explicit forward / backward calls' (``recompute_handle=None`` is the unchanged path); the handle holds exactly the history pool, the saved set's own tensor.""" - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp device = torch.device("cuda") p = _problem(lens, cuts, C, device) - y_e, hist_e = cake_causal_conv1d_kcp_forward(p.x_local, p.w, p.ctx) - dx_e, dw_e = cake_causal_conv1d_kcp_backward( + y_e, hist_e = cake_causal_conv1d_det_forward_kcp(p.x_local, p.w, p.ctx) + dx_e, dw_e = cake_causal_conv1d_det_backward_kcp( p.x_local, p.w, p.ctx, p.dy_local, hist_e ) explicit = {"y": y_e, "dx": dx_e, "dw": dw_e} @@ -688,7 +690,7 @@ def test_recompute_handle_is_bitwise_and_exchange_free( plain_cc = _ExchangeCounter(monkeypatch, conv_kcp) plain = _ckpt_run(p, reentrant=reentrant, counter=plain_cc) monkeypatch.undo() - handle = conv_kcp.KcpRecomputeHandle() + handle = conv_kcp.RecomputeHandleDetKcp() cc = _ExchangeCounter(monkeypatch, conv_kcp) conv_kcp.LAST_CALL.pop("recompute_handle", None) got = _ckpt_run(p, handle=handle, reentrant=reentrant, counter=cc) @@ -726,14 +728,16 @@ def test_recompute_handle_retain_graph_and_nested_checkpoints(monkeypatch) -> No from torch.utils.checkpoint import checkpoint - from flashinfer import cake_causal_conv1d_train as facade - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer import cake_causal_conv1d_det_train_kcp as facade + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp device = torch.device("cuda") p = _problem((5000, 1003), (1, 2, 4997, 1003), 512, device) direct = _ckpt_run(p, ckpt=False) cc = _ExchangeCounter(monkeypatch, conv_kcp) - twice = _ckpt_run(p, handle=conv_kcp.KcpRecomputeHandle(), counter=cc, twice=True) + twice = _ckpt_run( + p, handle=conv_kcp.RecomputeHandleDetKcp(), counter=cc, twice=True + ) assert cc.table() == {"orig": 1, "recompute": 0, "bwd": 2} assert torch.equal(_bits(twice["dx"]), _bits(direct["dx"] + direct["dx"])) assert torch.equal(_bits(twice["dw"]), _bits(direct["dw"] + direct["dw"])) @@ -745,10 +749,10 @@ def nested(inside: bool): w = p.w.clone().requires_grad_(True) def inner(x_, w_, h_): - return facade.cake_causal_conv1d_kcp(x_, w_, p.ctx, recompute_handle=h_) + return facade.cake_causal_conv1d_det_kcp(x_, w_, p.ctx, recompute_handle=h_) def outer(x_, w_, h_): - hh = conv_kcp.KcpRecomputeHandle() if inside else h_ + hh = conv_kcp.RecomputeHandleDetKcp() if inside else h_ yi = checkpoint(inner, x_, w_, hh, use_reentrant=False) return yi * torch.ones_like( yi @@ -758,7 +762,7 @@ def outer(x_, w_, h_): outer, x, w, - conv_kcp.KcpRecomputeHandle(), + conv_kcp.RecomputeHandleDetKcp(), use_reentrant=False, context_fn=lambda: (contextlib.nullcontext(), counter.recompute_cm()), ) @@ -782,55 +786,63 @@ def test_recompute_handle_misuse_raises() -> None: lost its entry before the recompute makes it an original forward again (exchanges, still bitwise).""" from torch.utils.checkpoint import checkpoint - from flashinfer import cake_causal_conv1d_train as facade - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer import cake_causal_conv1d_det_train_kcp as facade + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp device = torch.device("cuda") p = _problem((5000, 1003), (1, 2, 4997, 1003), 512, device) - handle = conv_kcp.KcpRecomputeHandle() - facade.cake_causal_conv1d_kcp(p.x_local, p.w, p.ctx, recompute_handle=handle) + handle = conv_kcp.RecomputeHandleDetKcp() + facade.cake_causal_conv1d_det_kcp(p.x_local, p.w, p.ctx, recompute_handle=handle) assert conv_kcp.LAST_CALL["recompute_handle"] == {"fwd": "capture"} with pytest.raises(RuntimeError, match="per call and micro-step"): - facade.cake_causal_conv1d_kcp(p.x_local, p.w, p.ctx, recompute_handle=handle) + facade.cake_causal_conv1d_det_kcp( + p.x_local, p.w, p.ctx, recompute_handle=handle + ) with pytest.raises(ValueError, match="graph='off'"): - conv_kcp.cake_causal_conv1d_kcp( + conv_kcp.cake_causal_conv1d_det_kcp( p.x_local, p.w, p.ctx, graph="replay", - recompute_handle=conv_kcp.KcpRecomputeHandle(), + recompute_handle=conv_kcp.RecomputeHandleDetKcp(), ) with pytest.raises(TypeError, match="entries"): - facade.cake_causal_conv1d_kcp(p.x_local, p.w, p.ctx, recompute_handle=object()) + facade.cake_causal_conv1d_det_kcp( + p.x_local, p.w, p.ctx, recompute_handle=object() + ) graph = torch.cuda.CUDAGraph() with pytest.raises(ValueError, match="capture"), torch.cuda.graph(graph): - conv_kcp.cake_causal_conv1d_kcp( - p.x_local, p.w, p.ctx, recompute_handle=conv_kcp.KcpRecomputeHandle() + conv_kcp.cake_causal_conv1d_det_kcp( + p.x_local, p.w, p.ctx, recompute_handle=conv_kcp.RecomputeHandleDetKcp() ) del graph # the public autograd Function refuses like the wrapper (graph keyword, capture, no entries dict) - fresh = conv_kcp.KcpRecomputeHandle() + fresh = conv_kcp.RecomputeHandleDetKcp() with pytest.raises(ValueError, match="graph='off'"): - conv_kcp.CakeCausalConv1dKcpFunction.apply( + conv_kcp.CakeCausalConv1dDetFunctionKcp.apply( p.x_local, p.w, p.ctx, "replay", fresh ) with pytest.raises(TypeError, match="entries"): - conv_kcp.CakeCausalConv1dKcpFunction.apply( + conv_kcp.CakeCausalConv1dDetFunctionKcp.apply( p.x_local, p.w, p.ctx, None, object() ) graph = torch.cuda.CUDAGraph() with pytest.raises(ValueError, match="capture"), torch.cuda.graph(graph): - conv_kcp.CakeCausalConv1dKcpFunction.apply(p.x_local, p.w, p.ctx, None, fresh) + conv_kcp.CakeCausalConv1dDetFunctionKcp.apply( + p.x_local, p.w, p.ctx, None, fresh + ) del graph assert not fresh.entries, "a refused call stored an entry" # a recompute whose entry belongs to another context / problem raises inside the backward other = _problem((5000, 1003), (1, 2, 4997, 1003), 512, device) for tamper in ("ctx", "key"): - h = conv_kcp.KcpRecomputeHandle() + h = conv_kcp.RecomputeHandleDetKcp() x = p.x_local.clone().requires_grad_(True) def fn(x_, h_): - return facade.cake_causal_conv1d_kcp(x_, p.w, p.ctx, recompute_handle=h_) + return facade.cake_causal_conv1d_det_kcp( + x_, p.w, p.ctx, recompute_handle=h_ + ) y = checkpoint(fn, x, h, use_reentrant=False) if tamper == "ctx": @@ -841,11 +853,11 @@ def fn(x_, h_): y.backward(p.dy_local) # the entry dropped before the recompute: the recompute is an original forward (exchanges and stores) direct = _ckpt_run(p, ckpt=False) - h = conv_kcp.KcpRecomputeHandle() + h = conv_kcp.RecomputeHandleDetKcp() x = p.x_local.clone().requires_grad_(True) w = p.w.clone().requires_grad_(True) y = checkpoint( - lambda x_, w_, h_: facade.cake_causal_conv1d_kcp( + lambda x_, w_, h_: facade.cake_causal_conv1d_det_kcp( x_, w_, p.ctx, recompute_handle=h_ ), x, @@ -861,21 +873,21 @@ def fn(x_, h_): def test_recompute_handle_knob_graph_is_refused(monkeypatch) -> None: - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp device = torch.device("cuda") p = _problem((65,), (64, 1), 256, device) monkeypatch.setenv(conv_kcp.TEST_KNOBS_ENV, "1") conv_kcp._set_test_knobs({"graph": "replay"}) - handle = conv_kcp.KcpRecomputeHandle() + handle = conv_kcp.RecomputeHandleDetKcp() try: with pytest.raises(ValueError, match="graph='off'"): - conv_kcp.cake_causal_conv1d_kcp( + conv_kcp.cake_causal_conv1d_det_kcp( p.x_local, p.w, p.ctx, recompute_handle=handle ) # the public autograd Function refuses the knob too with pytest.raises(ValueError, match="graph='off'"): - conv_kcp.CakeCausalConv1dKcpFunction.apply( + conv_kcp.CakeCausalConv1dDetFunctionKcp.apply( p.x_local, p.w, p.ctx, None, handle ) finally: @@ -888,7 +900,7 @@ def test_recompute_handle_knob_graph_is_refused(monkeypatch) -> None: def test_recompute_handle_both_host_paths(monkeypatch, host_trim) -> None: """The replay binds the stored pool on the trimmed and on the previous host path alike: bitwise, and no exchange in the recompute (the plain checkpoint, the positive control, exchanges once there).""" - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp device = torch.device("cuda") p = _problem((1, 2, 3, 57), (3, 30, 30), 384, device) @@ -902,7 +914,7 @@ def test_recompute_handle_both_host_paths(monkeypatch, host_trim) -> None: # the next counter wraps the module's own exchange, not the first counter monkeypatch.setattr(conv_kcp, "_exchange", exchange) cc = _ExchangeCounter(monkeypatch, conv_kcp) - got = _ckpt_run(p, handle=conv_kcp.KcpRecomputeHandle(), counter=cc) + got = _ckpt_run(p, handle=conv_kcp.RecomputeHandleDetKcp(), counter=cc) assert conv_kcp.LAST_CALL["recompute_handle"] == {"fwd": "replay"} finally: conv_kcp._set_test_knobs(None) diff --git a/tests/test_cake_causal_conv1d_train_kcp_dist.py b/tests/test_cake_causal_conv1d_det_train_kcp_dist.py similarity index 93% rename from tests/test_cake_causal_conv1d_train_kcp_dist.py rename to tests/test_cake_causal_conv1d_det_train_kcp_dist.py index cf2030aec63..78c454cf8b4 100644 --- a/tests/test_cake_causal_conv1d_train_kcp_dist.py +++ b/tests/test_cake_causal_conv1d_det_train_kcp_dist.py @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. -Multi-GPU rows of ``cake_causal_conv1d_kcp`` under ``torch.utils.checkpoint`` with and without a recompute handle +Multi-GPU rows of ``cake_causal_conv1d_det_kcp`` under ``torch.utils.checkpoint`` with and without a recompute handle (FlashInfer issues #7815 / #7818): ``W`` ranks over NCCL, each holding the fragments of one partition of a packed batch (contiguous, zigzag, per-document zigzag, ragged tables with a zero-token rank, one-token fragments and a fragment that reads its history from three predecessor ranks). Per row and rank: @@ -203,8 +203,9 @@ def _run_rank( ) -> dict: from torch.utils.checkpoint import checkpoint - from flashinfer import cake_causal_conv1d_train as conv - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer import cake_causal_conv1d_det_train as conv + from flashinfer import cake_causal_conv1d_det_train_kcp as kcp_api + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp from flashinfer.experimental.cake_kda_train import cake_kcp_plan as plan lengths, _w, layout, C = row @@ -233,7 +234,7 @@ def leaves(): ), w.detach().clone().requires_grad_(True) def fn(x_, w_, h_): - return conv.cake_causal_conv1d_kcp(x_, w_, ctx, recompute_handle=h_) + return kcp_api.cake_causal_conv1d_det_kcp(x_, w_, ctx, recompute_handle=h_) def run(mode: str, handle=None, *, twice: bool = False, reentrant: bool = False): """mode: direct | ckpt | nested_out (handle outside both regions) | nested_in (handle made in the outer).""" @@ -259,7 +260,7 @@ def run(mode: str, handle=None, *, twice: bool = False, reentrant: bool = False) def outer(x_, w_, h_): # the product with ones saves a tensor of the outer region (so the outer recompute runs) and # changes no bit of y or of its gradient - hh = conv_kcp.KcpRecomputeHandle() if mode == "nested_in" else h_ + hh = conv_kcp.RecomputeHandleDetKcp() if mode == "nested_in" else h_ yi = checkpoint(fn, x_, w_, hh, use_reentrant=False) return yi * torch.ones_like(yi) @@ -300,10 +301,10 @@ def eq(a, b) -> bool: info: dict[str, object] = {"n_hist_rows": n_hist, "T_r": int(x.shape[0])} direct, c_direct = run("direct") plain, c_plain = run("ckpt") - hd = conv_kcp.KcpRecomputeHandle() + hd = conv_kcp.RecomputeHandleDetKcp() withh, c_h = run("ckpt", hd) plain_re, c_plain_re = run("ckpt", reentrant=True) - hd_re = conv_kcp.KcpRecomputeHandle() + hd_re = conv_kcp.RecomputeHandleDetKcp() withh_re, c_h_re = run("ckpt", hd_re, reentrant=True) info["counts"] = { "direct": c_direct, @@ -314,7 +315,7 @@ def eq(a, b) -> bool: } # anchor: y is the single-rank zero-history convolution of the whole pack on this rank's tokens cu_dev = torch.tensor(cu, dtype=torch.int32, device=device) - y_single = conv.cake_causal_conv1d(x_all.to(device), w, cu_dev) + y_single = conv.cake_causal_conv1d_det(x_all.to(device), w, cu_dev) flags["y_eq_single_rank"] = same(direct["y"], _local(part, rank, cu, y_single)) flags["plain_eq_direct"] = eq(plain, direct) flags["handle_eq_plain"] = eq(withh, plain) @@ -340,7 +341,7 @@ def eq(a, b) -> bool: and hd.entries["conv"]["hist"].shape == (n_hist, C) ) # retain_graph: the entry is replayed again, no collective; exactly twice the gradients - hd2 = conv_kcp.KcpRecomputeHandle() + hd2 = conv_kcp.RecomputeHandleDetKcp() twice, c_twice = run("ckpt", hd2, twice=True) info["counts"]["handle_twice"] = c_twice flags["twice_accumulates"] = ( @@ -349,8 +350,8 @@ def eq(a, b) -> bool: and c_twice["calls"] == {"orig": 1, "recompute": 0, "bwd": 2} ) # nested checkpoints - nest_out, c_nest_out = run("nested_out", conv_kcp.KcpRecomputeHandle()) - nest_in, c_nest_in = run("nested_in", conv_kcp.KcpRecomputeHandle()) + nest_out, c_nest_out = run("nested_out", conv_kcp.RecomputeHandleDetKcp()) + nest_in, c_nest_in = run("nested_in", conv_kcp.RecomputeHandleDetKcp()) info["counts"]["nested_out"], info["counts"]["nested_in"] = c_nest_out, c_nest_in flags["nested_out_bitwise"] = eq(nest_out, direct) flags["nested_in_bitwise"] = eq(nest_in, direct) @@ -369,7 +370,7 @@ def eq(a, b) -> bool: } # retained bytes after a checkpointed forward (after the warm-up above): the history pool with a handle, 0 without retained = {} - for name, h in (("plain", None), ("handle", conv_kcp.KcpRecomputeHandle())): + for name, h in (("plain", None), ("handle", conv_kcp.RecomputeHandleDetKcp())): xl, wl = leaves() torch.cuda.synchronize(device) gc.collect() @@ -423,11 +424,11 @@ def _free_port() -> int: def _supported() -> bool: if not torch.cuda.is_available(): return False - from flashinfer.cake_causal_conv1d_train import ( - is_cake_causal_conv1d_train_supported, + from flashinfer.cake_causal_conv1d_det_train import ( + is_cake_causal_conv1d_det_train_supported, ) - return is_cake_causal_conv1d_train_supported(torch.device("cuda")) + return is_cake_causal_conv1d_det_train_supported(torch.device("cuda")) @pytest.mark.parametrize("row", ROWS, ids=_row_id) @@ -451,7 +452,7 @@ def test_recompute_handle_rows(row) -> None: assert report["pass"], report return if not _supported(): - pytest.skip("cake_causal_conv1d_train needs a supported CUDA device") + pytest.skip("cake_causal_conv1d_det_train needs a supported CUDA device") if torch.cuda.device_count() < W: pytest.skip(f"{W} CUDA devices required, {torch.cuda.device_count()} visible") with tempfile.TemporaryDirectory(prefix="cake_conv_kcp_dist_") as tmp: diff --git a/tests/test_cake_causal_conv1d_train_kcp_multi.py b/tests/test_cake_causal_conv1d_det_train_kcp_multi.py similarity index 78% rename from tests/test_cake_causal_conv1d_train_kcp_multi.py rename to tests/test_cake_causal_conv1d_det_train_kcp_multi.py index e60d143fd09..63e4f4e81a1 100644 --- a/tests/test_cake_causal_conv1d_train_kcp_multi.py +++ b/tests/test_cake_causal_conv1d_det_train_kcp_multi.py @@ -13,13 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. -Single-GPU tests of ``cake_causal_conv1d_kcp_multi`` (FlashInfer issue #7815): the convolutions of several tensors of +Single-GPU tests of ``cake_causal_conv1d_det_multi_kcp`` (FlashInfer issue #7815): the convolutions of several tensors of one layer (q / k / v) with one halo exchange per direction. One process, world size 1 (the rank holds every fragment of a partitioned problem, so history rows, head rows and the tail fold all run; one rank issues no collective). The -criterion is bitwise equality with one ``cake_causal_conv1d_kcp`` call per tensor: ``y``, ``hist``, ``dx`` and +criterion is bitwise equality with one ``cake_causal_conv1d_det_kcp`` call per tensor: ``y``, ``hist``, ``dx`` and ``dweight`` -- through the explicit forward / backward functions, the autograd entry, the facade, a checkpoint with and without a recompute handle, and graph replay. The NCCL rows (one collective per direction) are in -``test_cake_causal_conv1d_train_kcp_multi_dist.py``. +``test_cake_causal_conv1d_det_train_kcp_multi_dist.py``. """ from __future__ import annotations @@ -29,13 +29,15 @@ import pytest import torch -from flashinfer.cake_causal_conv1d_train import is_cake_causal_conv1d_train_supported -from tests.test_cake_causal_conv1d_train_kcp import ROWS, WIDTH, _context +from flashinfer.cake_causal_conv1d_det_train import ( + is_cake_causal_conv1d_det_train_supported, +) +from tests.test_cake_causal_conv1d_det_train_kcp import ROWS, WIDTH, _context pytestmark = pytest.mark.skipif( not torch.cuda.is_available() - or not is_cake_causal_conv1d_train_supported(torch.device("cuda")), - reason="cake_causal_conv1d_train needs a supported CUDA device (SM100 / SM103 / SM107)", + or not is_cake_causal_conv1d_det_train_supported(torch.device("cuda")), + reason="cake_causal_conv1d_det_train needs a supported CUDA device (SM100 / SM103 / SM107)", ) # The single-call rows, each with three tensors of the row's C, plus rows whose tensors have different channel counts @@ -56,7 +58,7 @@ def _same(a, b) -> bool: def _kcp(): - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp return conv_kcp @@ -76,8 +78,8 @@ def _problem(lens, cuts, Cs, device, seed=7815): def _singles_explicit(conv_kcp, ctx, xs, ws, dys): out = [] for x, w, dy in zip(xs, ws, dys, strict=True): - y, hist = conv_kcp.cake_causal_conv1d_kcp_forward(x, w, ctx) - dx, dw = conv_kcp.cake_causal_conv1d_kcp_backward(x, w, ctx, dy, hist) + y, hist = conv_kcp.cake_causal_conv1d_det_forward_kcp(x, w, ctx) + dx, dw = conv_kcp.cake_causal_conv1d_det_backward_kcp(x, w, ctx, dy, hist) out.append({"y": y, "hist": hist, "dx": dx, "dw": dw}) return out @@ -88,8 +90,10 @@ def test_multi_forward_backward_bitwise_single_calls(name, lens, cuts, Cs) -> No device = torch.device("cuda") ctx, xs, ws, dys = _problem(lens, cuts, Cs, device) ref = _singles_explicit(conv_kcp, ctx, xs, ws, dys) - ys, hists = conv_kcp.cake_causal_conv1d_kcp_multi_forward(xs, ws, ctx) - dxs, dws = conv_kcp.cake_causal_conv1d_kcp_multi_backward(xs, ws, ctx, dys, hists) + ys, hists = conv_kcp.cake_causal_conv1d_det_multi_forward_kcp(xs, ws, ctx) + dxs, dws = conv_kcp.cake_causal_conv1d_det_multi_backward_kcp( + xs, ws, ctx, dys, hists + ) torch.cuda.synchronize() assert len(ys) == len(hists) == len(dxs) == len(dws) == len(Cs) for i, r in enumerate(ref): @@ -99,10 +103,10 @@ def test_multi_forward_backward_bitwise_single_calls(name, lens, cuts, Cs) -> No assert _same(dws[i], r["dw"]), (name, i, "dweight") # caller-owned outputs are written in place, same bits outs = [torch.empty_like(x) for x in xs] - ys2, _ = conv_kcp.cake_causal_conv1d_kcp_multi_forward(xs, ws, ctx, outs=outs) + ys2, _ = conv_kcp.cake_causal_conv1d_det_multi_forward_kcp(xs, ws, ctx, outs=outs) dx_o = [torch.empty_like(x) for x in xs] dw_o = [torch.empty_like(w) for w in ws] - dxs2, dws2 = conv_kcp.cake_causal_conv1d_kcp_multi_backward( + dxs2, dws2 = conv_kcp.cake_causal_conv1d_det_multi_backward_kcp( xs, ws, ctx, dys, hists, dxs=dx_o, dweights=dw_o ) for i, r in enumerate(ref): @@ -117,7 +121,7 @@ def test_multi_autograd_and_facade_bitwise_single_calls( name, lens, cuts, Cs, batched ) -> None: """The autograd entry and the facade against one autograd single call per tensor (and ``[1, T_r, C]`` inputs).""" - from flashinfer import cake_causal_conv1d_train as facade + from flashinfer import cake_causal_conv1d_det_train_kcp as facade device = torch.device("cuda") ctx, xs, ws, dys = _problem(lens, cuts, Cs, device) @@ -129,10 +133,10 @@ def run(multi: bool): xl = [x.clone().requires_grad_(True) for x in xs] wl = [w.clone().requires_grad_(True) for w in ws] if multi: - ys = facade.cake_causal_conv1d_kcp_multi(xl, wl, ctx) + ys = facade.cake_causal_conv1d_det_multi_kcp(xl, wl, ctx) else: ys = [ - facade.cake_causal_conv1d_kcp(x, w, ctx) + facade.cake_causal_conv1d_det_kcp(x, w, ctx) for x, w in zip(xl, wl, strict=True) ] torch.autograd.backward(list(ys), list(dys)) @@ -174,7 +178,7 @@ def test_multi_recompute_handle_bitwise_single_handles( """A checkpoint of the multi call with ONE handle equals a checkpoint of the single calls with one handle each, the plain checkpoints and the direct pass, bitwise; the recompute replays the entry ("replay"); the handle holds the n pools (the saved set's own tensors, the forward's ``hist`` bits) and retains their bytes.""" - from flashinfer import cake_causal_conv1d_train as facade + from flashinfer import cake_causal_conv1d_det_train_kcp as facade conv_kcp = _kcp() device = torch.device("cuda") @@ -188,11 +192,11 @@ def run(multi: bool, handles, *, ckpt=True): def fn(*args): a_x, a_w, a_h = args[:n], args[n : 2 * n], args[2 * n :] if multi: - return facade.cake_causal_conv1d_kcp_multi( + return facade.cake_causal_conv1d_det_multi_kcp( a_x, a_w, ctx, recompute_handle=a_h[0] ) return tuple( - facade.cake_causal_conv1d_kcp(x, w, ctx, recompute_handle=h) + facade.cake_causal_conv1d_det_kcp(x, w, ctx, recompute_handle=h) for x, w, h in zip(a_x, a_w, a_h, strict=True) ) @@ -205,10 +209,10 @@ def fn(*args): ] direct = run(False, [None] * n, ckpt=False) - singles_h = [conv_kcp.KcpRecomputeHandle() for _ in range(n)] + singles_h = [conv_kcp.RecomputeHandleDetKcp() for _ in range(n)] single_handles = run(False, singles_h) plain_multi = run(True, [None]) - h = conv_kcp.KcpRecomputeHandle() + h = conv_kcp.RecomputeHandleDetKcp() conv_kcp.LAST_CALL.pop("recompute_handle", None) multi_handle = run(True, [h]) assert conv_kcp.LAST_CALL["recompute_handle"] == {"fwd": "replay", "n": n} @@ -233,7 +237,7 @@ def test_multi_recompute_handle_retain_graph_and_nested() -> None: outside nested regions serves both recomputes; the multi handle may also carry another package's entry.""" from torch.utils.checkpoint import checkpoint - from flashinfer import cake_causal_conv1d_train as facade + from flashinfer import cake_causal_conv1d_det_train_kcp as facade conv_kcp = _kcp() device = torch.device("cuda") @@ -242,7 +246,7 @@ def test_multi_recompute_handle_retain_graph_and_nested() -> None: ) def fn(xq, xk, xv, wq, wk, wv, h): - return facade.cake_causal_conv1d_kcp_multi( + return facade.cake_causal_conv1d_det_multi_kcp( (xq, xk, xv), (wq, wk, wv), ctx, recompute_handle=h ) @@ -256,7 +260,7 @@ def leaves(): torch.autograd.backward(list(ys), list(dys)) direct = [t.grad for t in a] a = leaves() - h = conv_kcp.KcpRecomputeHandle() + h = conv_kcp.RecomputeHandleDetKcp() h.entries["kda"] = {"marker": torch.zeros(1, device=device)} ys = checkpoint(fn, *a, h, use_reentrant=False) torch.autograd.backward(list(ys), list(dys), retain_graph=True) @@ -270,7 +274,7 @@ def outer(*args): return tuple(y * torch.ones_like(y) for y in yi) a = leaves() - h = conv_kcp.KcpRecomputeHandle() + h = conv_kcp.RecomputeHandleDetKcp() ys = checkpoint(outer, *a, h, use_reentrant=False) conv_kcp.LAST_CALL.pop("recompute_handle", None) torch.autograd.backward(list(ys), list(dys)) @@ -280,26 +284,32 @@ def outer(*args): def test_multi_recompute_handle_misuse_raises() -> None: - from flashinfer import cake_causal_conv1d_train as facade + from flashinfer import cake_causal_conv1d_det_train_kcp as facade conv_kcp = _kcp() device = torch.device("cuda") ctx, xs, ws, _ = _problem((5000, 1003), (1, 2, 4997, 1003), (512, 512, 512), device) # a handle used outside the backward a second time (reuse across micro-steps) - h = conv_kcp.KcpRecomputeHandle() - facade.cake_causal_conv1d_kcp_multi(xs, ws, ctx, recompute_handle=h) + h = conv_kcp.RecomputeHandleDetKcp() + facade.cake_causal_conv1d_det_multi_kcp(xs, ws, ctx, recompute_handle=h) with pytest.raises(RuntimeError, match="recompute_handle"): - facade.cake_causal_conv1d_kcp_multi(xs, ws, ctx, recompute_handle=h) + facade.cake_causal_conv1d_det_multi_kcp(xs, ws, ctx, recompute_handle=h) # graph replay and a non-handle are refused with pytest.raises(ValueError, match="recompute_handle"): - conv_kcp.cake_causal_conv1d_kcp_multi( - xs, ws, ctx, graph="replay", recompute_handle=conv_kcp.KcpRecomputeHandle() + conv_kcp.cake_causal_conv1d_det_multi_kcp( + xs, + ws, + ctx, + graph="replay", + recompute_handle=conv_kcp.RecomputeHandleDetKcp(), ) with pytest.raises(TypeError, match="entries"): - conv_kcp.cake_causal_conv1d_kcp_multi(xs, ws, ctx, recompute_handle=object()) + conv_kcp.cake_causal_conv1d_det_multi_kcp( + xs, ws, ctx, recompute_handle=object() + ) # entries of a single call / another tensor count are not a multi entry, and the reverse - single = conv_kcp.KcpRecomputeHandle() - facade.cake_causal_conv1d_kcp(xs[0], ws[0], ctx, recompute_handle=single) + single = conv_kcp.RecomputeHandleDetKcp() + facade.cake_causal_conv1d_det_kcp(xs[0], ws[0], ctx, recompute_handle=single) with pytest.raises(ValueError, match="multi-call"): conv_kcp._handle_hists(single.entries["conv"], xs, ctx) with pytest.raises(ValueError, match="multi-call"): @@ -313,13 +323,13 @@ def test_multi_recompute_handle_misuse_raises() -> None: assert all(a is b for a, b in zip(bound, h.entries["conv"]["hists"], strict=True)) # operand shape errors with pytest.raises(TypeError, match="sequences"): - conv_kcp.cake_causal_conv1d_kcp_multi(xs[0], ws[0], ctx) + conv_kcp.cake_causal_conv1d_det_multi_kcp(xs[0], ws[0], ctx) with pytest.raises(ValueError, match="same number"): - conv_kcp.cake_causal_conv1d_kcp_multi(xs, ws[:2], ctx) + conv_kcp.cake_causal_conv1d_det_multi_kcp(xs, ws[:2], ctx) with pytest.raises(ValueError, match="same number"): - conv_kcp.cake_causal_conv1d_kcp_multi([], [], ctx) + conv_kcp.cake_causal_conv1d_det_multi_kcp([], [], ctx) with pytest.raises(ValueError): - conv_kcp.cake_causal_conv1d_kcp_multi_forward( + conv_kcp.cake_causal_conv1d_det_multi_forward_kcp( [xs[0], xs[1][:-1].contiguous()], ws[:2], ctx ) @@ -331,8 +341,8 @@ def test_multi_empty_rank_returns_empty_outputs() -> None: Cs = (512, 256, 512) xs = [torch.empty(0, C, dtype=torch.bfloat16, device=device) for C in Cs] ws = [torch.randn(C, WIDTH, device=device) for C in Cs] - ys, hists = conv_kcp.cake_causal_conv1d_kcp_multi_forward(xs, ws, ctx) - dxs, dws = conv_kcp.cake_causal_conv1d_kcp_multi_backward( + ys, hists = conv_kcp.cake_causal_conv1d_det_multi_forward_kcp(xs, ws, ctx) + dxs, dws = conv_kcp.cake_causal_conv1d_det_multi_backward_kcp( xs, ws, ctx, [torch.empty_like(x) for x in xs], hists ) for y, dx, dw, C in zip(ys, dxs, dws, Cs, strict=True): @@ -353,11 +363,11 @@ def test_multi_graph_replay_is_bitwise() -> None: try: events = [] for _ in range(4): - ys, hists = conv_kcp.cake_causal_conv1d_kcp_multi_forward( + ys, hists = conv_kcp.cake_causal_conv1d_det_multi_forward_kcp( xs, ws, ctx, graph="replay" ) events.append(conv_kcp.LAST_CALL["fwd_multi"]["graph_event"]) - dxs, dws = conv_kcp.cake_causal_conv1d_kcp_multi_backward( + dxs, dws = conv_kcp.cake_causal_conv1d_det_multi_backward_kcp( xs, ws, ctx, dys, hists, graph="replay" ) events.append(conv_kcp.LAST_CALL["bwd_multi"]["graph_event"]) @@ -376,7 +386,7 @@ def test_multi_graph_replay_is_bitwise() -> None: ], events # the autograd entry carries the keyword to the forward and the backward xl = [x.clone().requires_grad_(True) for x in xs] - ys = conv_kcp.cake_causal_conv1d_kcp_multi(xl, ws, ctx, graph="replay") + ys = conv_kcp.cake_causal_conv1d_det_multi_kcp(xl, ws, ctx, graph="replay") assert conv_kcp.LAST_CALL["fwd_multi"]["graph"] == "replay" torch.autograd.backward(list(ys), list(dys)) assert conv_kcp.LAST_CALL["bwd_multi"]["graph"] == "replay" @@ -398,27 +408,30 @@ def test_multi_one_rank_issues_no_exchange(monkeypatch) -> None: monkeypatch.setattr( conv_kcp, "_exchange", lambda *a, **kw: calls.append(1) or original(*a, **kw) ) - ys, hists = conv_kcp.cake_causal_conv1d_kcp_multi_forward(xs, ws, ctx) - conv_kcp.cake_causal_conv1d_kcp_multi_backward(xs, ws, ctx, dys, hists) + ys, hists = conv_kcp.cake_causal_conv1d_det_multi_forward_kcp(xs, ws, ctx) + conv_kcp.cake_causal_conv1d_det_multi_backward_kcp(xs, ws, ctx, dys, hists) assert calls == [] def test_facade_multi_entry_is_lazy_and_keyword_only() -> None: - """The facade's ``cake_causal_conv1d_kcp_multi`` is exported, takes ``recompute_handle`` keyword-only (default - ``None``) and importing the facade still does not import the KCP host modules (fresh interpreter).""" + """The KCP facade's ``cake_causal_conv1d_det_multi_kcp`` is exported, takes ``recompute_handle`` keyword-only + (default ``None``), is the det facade's previous ``cake_causal_conv1d_kcp_multi``, and importing either facade + still does not import the KCP host modules (fresh interpreter).""" import subprocess import sys code = ( "import inspect, sys\n" - "import flashinfer.cake_causal_conv1d_train as f\n" - "kcp = 'flashinfer.cake_causal_conv1d_train_backend.kcp'\n" + "import flashinfer.cake_causal_conv1d_det_train_kcp as f\n" + "import flashinfer.cake_causal_conv1d_det_train as d\n" + "kcp = 'flashinfer.cake_causal_conv1d_det_train_backend.kcp'\n" "assert kcp not in sys.modules\n" - "s = inspect.signature(f.cake_causal_conv1d_kcp_multi).parameters\n" + "s = inspect.signature(f.cake_causal_conv1d_det_multi_kcp).parameters\n" "assert list(s) == ['xs', 'weights', 'ctx', 'recompute_handle'], list(s)\n" "p = s['recompute_handle']\n" "assert p.kind is inspect.Parameter.KEYWORD_ONLY and p.default is None\n" - "assert 'cake_causal_conv1d_kcp_multi' in f.__all__ and kcp not in sys.modules\n" + "assert d.cake_causal_conv1d_kcp_multi is f.cake_causal_conv1d_det_multi_kcp\n" + "assert 'cake_causal_conv1d_det_multi_kcp' in f.__all__ and kcp not in sys.modules\n" "print('FACADE_MULTI_OK')\n" ) r = subprocess.run( diff --git a/tests/test_cake_causal_conv1d_train_kcp_multi_dist.py b/tests/test_cake_causal_conv1d_det_train_kcp_multi_dist.py similarity index 90% rename from tests/test_cake_causal_conv1d_train_kcp_multi_dist.py rename to tests/test_cake_causal_conv1d_det_train_kcp_multi_dist.py index 0add6214c28..4c8e4d1ca1e 100644 --- a/tests/test_cake_causal_conv1d_train_kcp_multi_dist.py +++ b/tests/test_cake_causal_conv1d_det_train_kcp_multi_dist.py @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. -Multi-GPU rows of ``cake_causal_conv1d_kcp_multi`` (FlashInfer issue #7815): the q / k / v convolutions of a layer -with one halo exchange per direction, against one ``cake_causal_conv1d_kcp`` call per tensor. ``W`` ranks over NCCL, +Multi-GPU rows of ``cake_causal_conv1d_det_multi_kcp`` (FlashInfer issue #7815): the q / k / v convolutions of a layer +with one halo exchange per direction, against one ``cake_causal_conv1d_det_kcp`` call per tensor. ``W`` ranks over NCCL, each holding the fragments of one partition of a packed batch (contiguous, zigzag, per-document zigzag, ragged tables with a zero-token rank, one-token fragments and a fragment that reads its history from three predecessor ranks). Per row and rank: @@ -28,7 +28,7 @@ - the multi handle retains exactly the n history pools (``nbytes`` and the allocator's bytes after the forward); - ``graph="replay"`` of the multi forward / backward is bitwise the eager single calls on every call. -Launch forms as ``test_cake_causal_conv1d_train_kcp_dist.py``: ``torch.multiprocessing`` (``pytest``: skipped when +Launch forms as ``test_cake_causal_conv1d_det_train_kcp_dist.py``: ``torch.multiprocessing`` (``pytest``: skipped when fewer than ``W`` CUDA devices are visible) and one process per Slurm task (``SLURM_NTASKS == W``; the rendezvous from ``MASTER_ADDR`` / ``MASTER_PORT``; rows of another ``W`` are skipped under a multi-task step). """ @@ -46,7 +46,7 @@ import torch.distributed as dist import torch.multiprocessing as mp -from tests.test_cake_causal_conv1d_train_kcp_dist import ( +from tests.test_cake_causal_conv1d_det_train_kcp_dist import ( _free_port, _local, _partition, @@ -148,8 +148,9 @@ def _report(W: int, row, flags: dict, info: dict, result_dir, rank: int, tag: st def _run_rank(rank: int, W: int, row, port: int, result_dir, local_id: int) -> dict: from torch.utils.checkpoint import checkpoint - from flashinfer import cake_causal_conv1d_train as conv - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer import cake_causal_conv1d_det_train as conv + from flashinfer import cake_causal_conv1d_det_train_kcp as kcp_api + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp Cs = row[3] n = len(Cs) @@ -165,11 +166,11 @@ def layer(multi: bool, c): def fn(*args): a_x, a_w, a_h = args[:n], args[n : 2 * n], args[2 * n :] if multi: - return conv.cake_causal_conv1d_kcp_multi( + return kcp_api.cake_causal_conv1d_det_multi_kcp( a_x, a_w, c, recompute_handle=a_h[0] ) return tuple( - conv.cake_causal_conv1d_kcp(x, w, c, recompute_handle=h) + kcp_api.cake_causal_conv1d_det_kcp(x, w, c, recompute_handle=h) for x, w, h in zip(a_x, a_w, a_h, strict=True) ) @@ -198,7 +199,7 @@ def run( def outer(*args): hh = ( - [conv_kcp.KcpRecomputeHandle()] + [conv_kcp.RecomputeHandleDetKcp()] if mode == "nested_in" else list(args[2 * n :]) ) @@ -233,19 +234,23 @@ def table(o, r, b): info: dict[str, object] = {"n_hist_rows": n_hist, "T_r": int(xs[0].shape[0])} s_direct, cs_direct = run(False, "direct") s_plain, cs_plain = run(False, "ckpt") - s_hds = [conv_kcp.KcpRecomputeHandle() for _ in range(n)] + s_hds = [conv_kcp.RecomputeHandleDetKcp() for _ in range(n)] s_handle, cs_handle = run(False, "ckpt", s_hds) m_direct, cm_direct = run(True, "direct") m_plain, cm_plain = run(True, "ckpt") - m_hd = conv_kcp.KcpRecomputeHandle() + m_hd = conv_kcp.RecomputeHandleDetKcp() m_handle, cm_handle = run(True, "ckpt", [m_hd]) m_plain_re, cm_plain_re = run(True, "ckpt", reentrant=True) m_handle_re, cm_handle_re = run( - True, "ckpt", [conv_kcp.KcpRecomputeHandle()], reentrant=True + True, "ckpt", [conv_kcp.RecomputeHandleDetKcp()], reentrant=True ) - m_twice, cm_twice = run(True, "ckpt", [conv_kcp.KcpRecomputeHandle()], twice=True) - m_nest_out, cm_nest_out = run(True, "nested_out", [conv_kcp.KcpRecomputeHandle()]) - m_nest_in, cm_nest_in = run(True, "nested_in", [conv_kcp.KcpRecomputeHandle()]) + m_twice, cm_twice = run( + True, "ckpt", [conv_kcp.RecomputeHandleDetKcp()], twice=True + ) + m_nest_out, cm_nest_out = run( + True, "nested_out", [conv_kcp.RecomputeHandleDetKcp()] + ) + m_nest_in, cm_nest_in = run(True, "nested_in", [conv_kcp.RecomputeHandleDetKcp()]) info["counts"] = { "single_direct": cs_direct, "single_plain": cs_plain, @@ -268,7 +273,7 @@ def table(o, r, b): part, rank, cu, - conv.cake_causal_conv1d(xs_all[i].to(device), ws[i], cu_dev), + conv.cake_causal_conv1d_det(xs_all[i].to(device), ws[i], cu_dev), ), ) for i in range(n) @@ -317,7 +322,7 @@ def table(o, r, b): ) # retained bytes after a checkpointed forward: the n history pools with a handle, 0 without retained = {} - for name, h in (("plain", None), ("handle", conv_kcp.KcpRecomputeHandle())): + for name, h in (("plain", None), ("handle", conv_kcp.RecomputeHandleDetKcp())): xl, wl = leaves() torch.cuda.synchronize(device) gc.collect() @@ -338,7 +343,7 @@ def table(o, r, b): gctx = gplan.build_kcp_context(gpart, dist.group.WORLD, rank, device) g_single, cg_single = run(False, "direct", c=gctx) g_multi, cg_multi = run(True, "direct", c=gctx) - g_handle, cg_handle = run(True, "ckpt", [conv_kcp.KcpRecomputeHandle()], c=gctx) + g_handle, cg_handle = run(True, "ckpt", [conv_kcp.RecomputeHandleDetKcp()], c=gctx) info["counts"]["gdn_ctx"] = { "single_direct": cg_single, "multi_direct": cg_multi, @@ -361,22 +366,22 @@ def _run_rank_graph( ) -> dict: """``graph="replay"`` of the multi forward / backward: eager, capture, replay -- bitwise the eager single calls on every call; a rank without tokens is never captured (it joins each collective eagerly).""" - from flashinfer.cake_causal_conv1d_train_backend import kcp as conv_kcp + from flashinfer.cake_causal_conv1d_det_train_backend import kcp as conv_kcp device, part, cu, ctx, xs_all, xs, ws, dys = _setup(rank, W, row, port, local_id) ref = [] for x, w, dy in zip(xs, ws, dys, strict=True): - y, hist = conv_kcp.cake_causal_conv1d_kcp_forward(x, w, ctx) - dx, dw = conv_kcp.cake_causal_conv1d_kcp_backward(x, w, ctx, dy, hist) + y, hist = conv_kcp.cake_causal_conv1d_det_forward_kcp(x, w, ctx) + dx, dw = conv_kcp.cake_causal_conv1d_det_backward_kcp(x, w, ctx, dy, hist) ref.append((y, hist, dx, dw)) events, ok = [], True try: for _ in range(4): - ys, hists = conv_kcp.cake_causal_conv1d_kcp_multi_forward( + ys, hists = conv_kcp.cake_causal_conv1d_det_multi_forward_kcp( xs, ws, ctx, graph="replay" ) events.append(conv_kcp.LAST_CALL["fwd_multi"]["graph_event"]) - dxs, dws = conv_kcp.cake_causal_conv1d_kcp_multi_backward( + dxs, dws = conv_kcp.cake_causal_conv1d_det_multi_backward_kcp( xs, ws, ctx, dys, hists, graph="replay" ) events.append(conv_kcp.LAST_CALL["bwd_multi"]["graph_event"]) @@ -440,7 +445,7 @@ def _launch(row, graph: bool) -> None: assert report["pass"], report return if not _supported(): - pytest.skip("cake_causal_conv1d_train needs a supported CUDA device") + pytest.skip("cake_causal_conv1d_det_train needs a supported CUDA device") if torch.cuda.device_count() < W: pytest.skip(f"{W} CUDA devices required, {torch.cuda.device_count()} visible") with tempfile.TemporaryDirectory(prefix="cake_conv_kcp_multi_dist_") as tmp: diff --git a/tests/test_helpers/cake_package.py b/tests/test_helpers/cake_package.py index 5051152b048..f038e5573fa 100644 --- a/tests/test_helpers/cake_package.py +++ b/tests/test_helpers/cake_package.py @@ -15,7 +15,7 @@ # suite / package name -> import path of the package PACKAGES: dict[str, str] = { "cake_rmsnorm_det_train": "flashinfer.cake_rmsnorm_det_train_backend", - "cake_causal_conv1d_train": "flashinfer.cake_causal_conv1d_train_backend", + "cake_causal_conv1d_det_train": "flashinfer.cake_causal_conv1d_det_train_backend", "cake_gated_act_mxfp8": "flashinfer.cake_gated_act_mxfp8_backend", "cake_rmsnorm_silu_gate_det_train": "flashinfer.cake_rmsnorm_silu_gate_det_train_backend", "cake_ulysses": "flashinfer.cake_ulysses_backend", @@ -48,6 +48,8 @@ "cake_gdn_det_train": "flashinfer.experimental.cake_gdn_det_train", "cake_gqa_train": "flashinfer.experimental.cake_gqa_train", } +# Compatibility key: suites that name the causal conv1d package ``cake_causal_conv1d_train`` load the det package. +PACKAGES["cake_causal_conv1d_train"] = PACKAGES["cake_causal_conv1d_det_train"] def cake_package(name: str, submodule: str | None = None) -> ModuleType: