refactor(cake_causal_conv1d): deterministic KCP names for the context-parallel causal conv1d entry points

benchmarks/bench_cake_causal_conv1d_train.py → 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()

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
 

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

flashinfer/__init__.py

@@ -1,4 +1,4 @@
 """
 Copyright (c) 2023 by FlashInfer team.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -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 @@
 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,
 )