kernel-fun 0.2.0.dev1__py3-none-any.whl

kernel_fun/__init__.py

@@ -0,0 +1,67 @@
+"""kernel-fun — CuTe/Triton kernels for linear-attention ops on Blackwell.
+
+Each op family is a subpackage exporting a drop-in for its flash-linear-attention
+counterpart, with the same signature and the same return contract:
+
+    from kernel_fun.kda import chunk_kda          # the KDA chunk kernel (CuTe + Triton)
+    from kernel_fun.cconv import causal_conv1d    # the KDA short conv + silu (Triton)
+    # gdn, gnorm to follow into the same shell
+
+A call the package does not implement — wrong architecture, wrong shape, a flag we have
+never seen, a CUDA graph capture of a shape that has not run eagerly yet — is forwarded to
+fla verbatim, so installing this can change how fast a model trains but not what it
+computes beyond kernel-level rounding.
+
+Nothing is imported eagerly: `import kernel_fun` must stay cheap on a machine with no GPU,
+so torch, triton and the CuTe DSL are pulled in by the family that needs them.
+
+Two repos, and changes go to different ones: this package is developed at
+github.com/allenai/kernel-fun (the release artifact — and the only route into
+OLMo-core), while the kernels are researched in github.com/allenai/kernel-fun-dev (the ladder,
+`kernels/<family>/ideas/*`) and cross over by `tools/vendor.py`. A `kernels/...` path in a
+comment here is a path in the LADDER. See the pkg repo's README, "Two repos".
+
+Two environment switches, read per call and documented in the README:
+    KERNEL_FUN_DISABLE=1        forward everything to fla (per family: _KDA_, _CCONV_, ...)
+    KERNEL_FUN_DEBUG=1          log why a call fell back
+"""
+
+from __future__ import annotations
+
+__all__ = ["__version__", "versions"]
+
+__version__ = "0.2.0.dev1"
+
+
+def versions() -> dict[str, str]:
+    """Everything that decides what this package computes and how fast.
+
+    Worth logging once per training run: two of these (the CuTe DSL and cuda-python) are
+    not pinned by any requirement and arrive with the base image, and no bench row in the
+    research repo records them either. When a run is slower than the last one, this is the
+    first thing to diff.
+    """
+    import importlib
+
+    out: dict[str, str] = {"kernel_fun": __version__}
+    for name, mod in (
+        ("torch", "torch"),
+        ("triton", "triton"),
+        ("fla", "fla"),
+        ("cutlass", "cutlass"),
+        ("cuda-python", "cuda.bindings"),
+    ):
+        try:
+            m = importlib.import_module(mod)
+            out[name] = str(getattr(m, "__version__", "unknown"))
+        except Exception:
+            out[name] = "not installed"
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            cap = torch.cuda.get_device_capability()
+            out["device"] = f"{torch.cuda.get_device_name()} sm{cap[0]}{cap[1]}"
+    except Exception:
+        pass
+    return out

kernel_fun/_common/__init__.py

@@ -0,0 +1,11 @@
+"""Machinery every op family in this package shares.
+
+Deliberately small, and deliberately not a place for op logic: the call cache, the support
+predicate, the fla compatibility probe, and once-per-process logging. Those four are what
+the kda / gdn / gnorm chains all needed independently — and each had its own copy in the
+research tree, which is how one leak fix ended up written for gnorm and missing from the
+other two for two weeks.
+
+Nothing here imports torch or cutlass at module scope beyond what it must, so importing
+`kernel_fun` on a machine without a GPU stays cheap.
+"""

kernel_fun/_common/cache.py

@@ -0,0 +1,99 @@
+"""The CuTe call cache: marshal once, poke the pointer, launch.
+
+Every CuTe kernel in this package drives its compiled object through the same three
+mechanisms, and they live here once rather than copy-pasted per kernel (which is how they
+exist in the research tree, and how the keepalive bug below survived in four places at
+once):
+
+  cute_view        — a torch tensor as a cute.Tensor with the mode order cute wants and the
+                     batch-ish modes marked dynamic, so one compile serves every B/T.
+  retarget         — rewrite the descriptor's data pointer for this call. Marshaling costs
+                     ~0.07ms per kernel; a ctypes word-write costs nothing.
+  release_keepalives — drop the DLPack reference to the tensor that exported the view.
+
+The last one is not an optimization, it is a leak fix: `from_dlpack` requires the consumer
+to keep the exporter alive, so a cached view pins the FIRST call's tensor for the life of
+the process. Measured on the kda chain at B8/T1024/H8/HV8/K128/V256: 778 MiB pinned across
+its four cache entries, ~24 GiB at prod8192 shapes, never freed — which is exactly the
+condition that pushes the caching allocator into device-synchronizing cudaMalloc retries.
+gnorm hit the same bug at +3.5 GiB/rank on the 1.4b ladder.
+
+IMPORTANT for anything cached here: outputs are allocated fresh per call and retargeted,
+never reused. A cache-owned output silently aliases across calls, which a bench that makes
+one call per iteration cannot see.
+"""
+
+from __future__ import annotations
+
+import ctypes
+
+import torch
+from cutlass.cute.runtime import from_dlpack
+
+
+def cute_view(t: torch.Tensor, perm: tuple[int, ...], dyn_modes: tuple[int, ...]):
+    """A cute.Tensor view of `t`, permuted to `perm`, with `dyn_modes` dynamic.
+
+    cute wants modes ordered outermost-to-innermost and checks that against the strides.
+    `tt.dim_order()` is not usable: it places size-1 modes by contiguity heuristics, so e.g.
+    (512,128,2,1):(256,1,128,131072) comes back as (0,2,3,1) and cute rejects it. Sorting
+    the permuted tensor's strides has the mirror problem — size-1 modes tie with their
+    neighbour and land on the wrong side. Take the order from the UNPERMUTED tensor, where
+    descending stride is unambiguous, and map it through perm.
+
+    detach: inside an autograd.Function the incoming leaves still carry requires_grad and
+    dlpack refuses to export those. Every gradient path here is hand-written, so only the
+    storage is wanted.
+    """
+    t = t.detach()
+    base_order = sorted(range(t.dim()), key=lambda i: -t.stride(i))
+    new_of_old = {old: new for new, old in enumerate(perm)}
+    stride_order = tuple(new_of_old[d] for d in base_order)
+    ct = from_dlpack(t.permute(*perm), assumed_align=16)
+    for m in dyn_modes:
+        ct = ct.mark_compact_shape_dynamic(mode=m, stride_order=stride_order)
+    return ct
+
+
+def retarget(ct, t: torch.Tensor) -> None:
+    """Point a cached view at this call's tensor. One 64-bit store."""
+    ptr = t.data_ptr()
+    assert ptr & 15 == 0, "kernel views assume 16B alignment"
+    ctypes.c_uint64.from_address(ct.__c_pointers__()[0]).value = ptr
+
+
+def release_keepalives(*cts) -> None:
+    """Drop cached views' DLPack references to the tensors that exported them.
+
+    Call this after `cute.compile` and before the entry lands in a call cache, on EVERY
+    view including the outputs' — a cached entry needs only the C memref descriptor, which
+    the view owns, since `retarget` rewrites the data word before every launch.
+
+    Populate the pointer cache first, then null the two reference fields. Those names are
+    cutlass internals: fail loudly if an upgrade renames them rather than silently going
+    back to pinning gigabytes.
+    """
+    for ct in cts:
+        ct.__c_pointers__()  # materialize _c_pointers_cache before the source goes away
+        d = ct.__dict__
+        assert "_dlpack_data" in d and "_dltensor_wrapper" in d, (
+            "cutlass.cute.runtime._Tensor changed its keepalive fields; re-derive "
+            "release_keepalives against the new version (see kernel_fun._common.cache)"
+        )
+        d["_dlpack_data"] = None
+        d["_dltensor_wrapper"] = None
+
+
+def out_specs(*tensors: torch.Tensor) -> tuple:
+    """(shape, dtype) per output — what a cache entry keeps instead of the tensors."""
+    return tuple((tuple(t.shape), t.dtype) for t in tensors)
+
+
+def alloc_outs(specs: tuple, device) -> tuple:
+    return tuple(torch.empty(shape, device=device, dtype=dtype) for shape, dtype in specs)
+
+
+def stream_handle() -> int:
+    """The current stream's raw handle — part of every call key, so a stream switch makes a
+    new entry rather than launching a stale one."""
+    return torch.cuda.current_stream().cuda_stream

kernel_fun/_common/compat.py

@@ -0,0 +1,168 @@
+"""Are the flash-linear-attention internals we call still the ones we were built against?
+
+This package is a partial replacement: several stages of every chain are fla's own kernels,
+called by private path, and one of them (`chunk_kda_fwd_kernel_inter_solve_fused`) is a raw
+`triton.jit` kernel we launch with our own grid and constexpr list. That is a deliberate
+trade — reusing fla's kernels is why our numbers are comparable to fla's launch for launch —
+but it means an fla upgrade can change our results without changing our code.
+
+So: check at warmup, not at import. A raise here happens before a training run spends any
+time; the same drift discovered at step 40,000 costs a day.
+
+Policy: warn on a version we have not tested, raise on a symbol or signature that moved.
+KERNEL_FUN_FALLBACK=1 downgrades the raise to a warning plus a fallback, for whoever needs
+the cluster running now rather than correct attribution.
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+import os
+from functools import cache
+
+log = logging.getLogger(__name__)
+
+# Versions this package has actually been run against, not a guess at compatibility.
+TESTED_FLA = {"0.5.2"}
+
+# module path -> symbols we import from it. Kept as data so the check and the docs cannot
+# drift apart, and so a new family only has to extend this table.
+FLA_SYMBOLS: dict[str, tuple[str, ...]] = {
+    "fla.ops.kda.chunk_intra": (
+        "chunk_kda_fwd_intra",
+        "chunk_kda_bwd_intra",
+        "chunk_kda_fwd_kernel_inter_solve_fused",
+    ),
+    "fla.ops.kda.chunk_intra_token_parallel": ("chunk_kda_fwd_intra_token_parallel",),
+    "fla.ops.kda.chunk_bwd": ("chunk_kda_bwd_dAv", "chunk_kda_bwd_wy_dqkg_fused"),
+    "fla.ops.kda.wy_fast": ("recompute_w_u_fwd",),
+    "fla.ops.kda.gate": ("kda_gate_chunk_cumsum", "kda_gate_bwd"),
+    "fla.ops.common.chunk_delta_h": (
+        "chunk_gated_delta_rule_fwd_h",
+        "chunk_gated_delta_rule_bwd_dhu",
+    ),
+    "fla.ops.common.gate": ("fused_beta_sigmoid", "fused_beta_sigmoid_bwd"),
+    "fla.ops.gla.chunk": ("chunk_gla_fwd_o_gk",),
+    "fla.ops.utils": ("chunk_local_cumsum",),
+    "fla.ops.utils.constant": ("RCP_LN2",),
+    "fla.ops.utils.cache": ("fla_cache_autotune",),
+    "fla.ops.utils.op": ("exp2",),
+    "fla.utils": (
+        "autotune_cache_kwargs",
+        "check_shared_mem",
+        "input_guard",
+        "autocast_custom_fwd",
+        "autocast_custom_bwd",
+    ),
+    "fla.modules.l2norm": ("l2norm_fwd", "l2norm_bwd"),
+    # cconv: the public entry point is also the fallback, and its parameter list is what
+    # the whitelist in kernel_fun.cconv.ops classifies.
+    "fla.modules.convolution": ("causal_conv1d",),
+}
+
+# Keywords we pass by name. A rename upstream is a TypeError at the worst possible moment
+# otherwise; a reorder of the raw triton kernel's args would be worse — wrong answers.
+FLA_KWARGS: dict[tuple[str, str], tuple[str, ...]] = {
+    ("fla.ops.kda.wy_fast", "recompute_w_u_fwd"): ("q", "k", "v", "beta", "A", "gk"),
+    ("fla.ops.kda.gate", "kda_gate_chunk_cumsum"): (
+        "g", "A_log", "dt_bias", "scale", "chunk_size",
+    ),
+    ("fla.ops.kda.gate", "kda_gate_bwd"): ("g", "A_log", "dt_bias", "dyg"),
+    ("fla.ops.kda.chunk_bwd", "chunk_kda_bwd_dAv"): (
+        "q", "k", "v", "do", "A", "scale", "chunk_size",
+    ),
+    ("fla.ops.utils", "chunk_local_cumsum"): ("g", "chunk_size", "reverse"),
+    ("fla.modules.convolution", "causal_conv1d"): (
+        "x", "weight", "bias", "residual", "initial_state", "output_final_state",
+        "activation", "backend", "cu_seqlens", "cu_seqlens_cpu", "chunk_indices",
+        "cp_context",
+    ),
+}
+
+
+class FlaDriftError(RuntimeError):
+    """An fla internal we depend on moved. See the message for which one."""
+
+
+def _fail(msg: str) -> None:
+    if os.environ.get("KERNEL_FUN_FALLBACK", "0") == "1":
+        log.warning("kernel-fun: %s (KERNEL_FUN_FALLBACK=1: falling back to fla)", msg)
+        return
+    raise FlaDriftError(
+        f"{msg}\n"
+        f"kernel-fun calls flash-linear-attention internals directly; tested against "
+        f"{sorted(TESTED_FLA)}. Set KERNEL_FUN_FALLBACK=1 to degrade to fla's own kernels "
+        f"instead of raising."
+    )
+
+
+@cache
+def check_fla() -> bool:
+    """Verify every fla symbol and keyword this package uses. Cached; call from warmup.
+
+    Returns False when drift was found and KERNEL_FUN_FALLBACK downgraded it, so callers
+    can route to fla wholesale.
+    """
+    import importlib
+
+    import fla
+
+    version = getattr(fla, "__version__", "unknown")
+    if version not in TESTED_FLA:
+        log.warning(
+            "kernel-fun: flash-linear-attention %s is untested here (tested: %s); "
+            "verify numerics before trusting a training run",
+            version, sorted(TESTED_FLA),
+        )
+
+    ok = True
+    for mod_path, names in FLA_SYMBOLS.items():
+        try:
+            mod = importlib.import_module(mod_path)
+        except ImportError as e:
+            _fail(f"cannot import {mod_path}: {e}")
+            ok = False
+            continue
+        for name in names:
+            if not hasattr(mod, name):
+                _fail(f"{mod_path}.{name} is gone in fla {version}")
+                ok = False
+
+    for (mod_path, name), kwargs in FLA_KWARGS.items():
+        try:
+            fn = getattr(importlib.import_module(mod_path), name)
+        except (ImportError, AttributeError):
+            continue  # already reported above
+        params = _parameters(fn)
+        if params is None:
+            continue
+        missing = [kw for kw in kwargs if kw not in params]
+        if missing:
+            _fail(f"{mod_path}.{name} no longer accepts {missing} in fla {version}")
+            ok = False
+    return ok
+
+
+def _parameters(fn) -> set[str] | None:
+    """Parameter names of a python function or a triton.jit kernel.
+
+    A JITFunction is not introspectable by `inspect.signature`; it carries `arg_names`.
+    That list is also ORDERED, which matters for the one kernel we launch ourselves — a
+    reorder upstream would silently feed our arguments to the wrong parameters.
+    """
+    arg_names = getattr(fn, "arg_names", None)
+    if arg_names is not None:
+        return set(arg_names)
+    try:
+        return set(inspect.signature(fn).parameters)
+    except (TypeError, ValueError):  # pragma: no cover - exotic callables
+        return None
+
+
+def kernel_arg_names(mod_path: str, name: str) -> tuple[str, ...]:
+    """Ordered parameter names of a raw triton kernel, for an exact-match assertion."""
+    import importlib
+
+    fn = getattr(importlib.import_module(mod_path), name)
+    return tuple(getattr(fn, "arg_names", ()))

kernel_fun/_common/support.py

@@ -0,0 +1,267 @@
+"""Can this call use our kernels — and if not, exactly why?
+
+Every public entry point in this package answers that question before doing anything, and
+delegates to flash-linear-attention when the answer is no. The reason string is the
+important half: a silent fallback reads as a correct 1.00x result, and the last time
+kernels from this repo landed in training the single most expensive question was "did they
+actually run?".
+
+Probes that touch the driver or import cutlass are cached — the previous port called
+`import cutlass.cute` on every forward. Probe functions are module-level and cached so a
+test can monkeypatch one and clear its cache to exercise the fallback on any GPU.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from functools import cache
+
+import torch
+
+log = logging.getLogger(__name__)
+
+# Serial-scan kernels need a full GPU: a one-CTA-per-(chunk, b*hv) kernel at a 64-CTA grid
+# lost 1.4ms to fla in the gdn ladder. Below this the fla path is genuinely faster, so the
+# gate is a performance decision, not a capability one.
+MIN_CTAS = 256
+
+_logged: set[str] = set()
+_VERSIONS_KEY = "kernel-fun versions"  # sentinel in _logged, so a test's reset re-arms it
+
+
+def log_once(message: str, level: int = logging.INFO) -> None:
+    if message in _logged:
+        return
+    _logged.add(message)
+    log.log(level, message)
+
+
+def log_versions_once() -> None:
+    """Log `versions()` the first time any family engages.
+
+    This belongs HERE, not in the caller's module: every entry point that calls it is
+    already `torch.compiler.disable`d, so the line costs no graph break, whereas the same
+    call from a compiled `forward` splits the block (and, in olmo-core, hit `lru_cache` on
+    a dict return and raised `TypeError: unhashable type: 'dict'`). Two of these versions
+    — the CuTe DSL and cuda-python — are pinned by nothing and ride in with the image,
+    so when a run is slower than the last one this is the first thing to diff.
+    """
+    if _VERSIONS_KEY in _logged:
+        return  # before versions(): it queries device properties, ~14 us, and this is per call
+    _logged.add(_VERSIONS_KEY)
+    from .. import versions
+
+    log.info(f"kernel-fun {versions()}")
+
+
+@cache
+def has_cute() -> bool:
+    """Is the CuTe DSL importable? Cached: the import pulls MLIR and costs seconds.
+
+    Also warns, once, when the DSL's CUDA build visibly disagrees with torch's — see
+    `cute_cuda_mismatch`. A warning and not a gate: the wrong-build DSL may well still
+    compile (a CUDA 12 toolchain on a CUDA 13 driver is a supported combination), and
+    silently routing a working install to fla is the failure this package exists to avoid.
+    """
+    try:
+        import cuda.bindings.driver  # noqa: F401
+        import cutlass  # noqa: F401
+        import cutlass.cute  # noqa: F401
+    except Exception:  # pragma: no cover - environment-dependent
+        return False
+    reason = cute_cuda_mismatch(torch.version.cuda, _installed_dists())
+    if reason is not None:
+        log_once(f"kernel-fun: {reason}", logging.WARNING)
+    return True
+
+
+def _installed_dists() -> frozenset[str]:
+    import importlib.metadata as md
+
+    return frozenset(
+        d.metadata["Name"].lower() for d in md.distributions() if d.metadata["Name"]
+    )
+
+
+def cute_cuda_mismatch(torch_cuda: str | None, installed: frozenset[str]) -> str | None:
+    """Is the installed CuTe DSL the CUDA build torch was built for? Pure, for the test.
+
+    On PyPI `nvidia-cutlass-dsl` >= 4.6 always pulls `nvidia-cutlass-dsl-libs-cu12` and
+    ships the CUDA 13 libraries only through its `[cu13]` extra, so a CUDA 13 torch next to
+    a DSL whose cu13 libs are absent means somebody wrote the bare requirement (this
+    package's own extra is `kernel-fun[cu13]`; OLMo-core's is FA4's `[cu13]`). Only that
+    one direction is detectable: before 4.6 the CUDA 12 libraries live inside the base
+    wheel, and a cu13-only install on a CUDA 12 torch cannot be told from a pre-split one.
+    """
+    if not torch_cuda or "nvidia-cutlass-dsl" not in installed:
+        return None
+    major = torch_cuda.split(".")[0]
+    has_cu12 = "nvidia-cutlass-dsl-libs-cu12" in installed
+    has_cu13 = "nvidia-cutlass-dsl-libs-cu13" in installed
+    if major == "13" and has_cu12 and not has_cu13:
+        return (
+            f"torch is a CUDA {torch_cuda} build but the CuTe DSL installed is the CUDA 12 "
+            f"one (nvidia-cutlass-dsl-libs-cu12 without -cu13). Install "
+            f"nvidia-cutlass-dsl[cu13] — kernel-fun[cu13] does — or expect cute.compile "
+            f"to fail rather than fall back"
+        )
+    return None
+
+
+@cache
+def arch_at_least(device_index: int, major: int) -> bool:
+    """For the Triton-only families: any Blackwell-or-newer datacenter part, or Hopper.
+
+    Those kernels need no tcgen05 and no CuTe DSL, so the sm100 gate below would deny them
+    for no reason; but nobody has timed them below sm100, so callers state their floor.
+    """
+    return torch.cuda.get_device_capability(device_index)[0] >= major
+
+
+@cache
+def arch_ok(device_index: int) -> bool:
+    """sm100 exactly — Blackwell datacenter (B200/B300).
+
+    Not `major >= 10`: sm_120 is consumer Blackwell and has no tcgen05, so the MMA kernels
+    would fail inside cute.compile rather than fall back. Not `major == 10 and minor == 0`
+    either: B300 reports (10, 3).
+    """
+    return torch.cuda.get_device_capability(device_index)[0] == 10
+
+
+def disabled(family: str) -> bool:
+    """KERNEL_FUN_DISABLE=1 kills every family; KERNEL_FUN_<FAMILY>_DISABLE=1 kills one.
+
+    Read per call, never at import: the point of a kill switch is that someone can set it
+    on a run that is already failing, and an import-time read would depend on which module
+    got imported first.
+    """
+    return (
+        os.environ.get("KERNEL_FUN_DISABLE", "0") == "1"
+        or os.environ.get(f"KERNEL_FUN_{family.upper()}_DISABLE", "0") == "1"
+    )
+
+
+def debug() -> bool:
+    return os.environ.get("KERNEL_FUN_DEBUG", "0") == "1"
+
+
+def min_ctas(family: str) -> int:
+    """`family`'s dispatch floor: `MIN_CTAS`, or `KERNEL_FUN_<FAMILY>_MIN_CTAS` if set.
+
+    A performance heuristic, not a capability — the kernels run below it, they just lose to
+    fla where the CuTe scans underfill the GPU — and the crossover belongs to the workload,
+    not to the hardware. So a shape that has been *measured* below the default can opt in
+    from outside, instead of a training script assigning to `MIN_CTAS` and moving whichever
+    per-stage gates happened to bind it at import.
+
+    Only the chain-level gate moves. Every kernel keeps its own floor, so a lowered value
+    changes which stages are ours, never what any of them computes.
+
+    Read per call, like `disabled`. A value that is not a positive integer is ignored with a
+    warning: falling back is this package's job, ending someone's run at step 1 is not.
+    """
+    raw = os.environ.get(f"KERNEL_FUN_{family.upper()}_MIN_CTAS")
+    if raw is None:
+        return MIN_CTAS
+    try:
+        floor = int(raw)
+    except ValueError:
+        floor = 0
+    if floor < 1:
+        log_once(
+            f"kernel-fun {family}: ignoring KERNEL_FUN_{family.upper()}_MIN_CTAS={raw!r} "
+            f"(want a positive integer); using the default floor of {MIN_CTAS} CTAs",
+            logging.WARNING,
+        )
+        return MIN_CTAS
+    return floor
+
+
+def capturing() -> bool:
+    """Is a CUDA graph being captured on this stream?"""
+    try:
+        return torch.cuda.is_current_stream_capturing()
+    except Exception:  # pragma: no cover - older torch without the query
+        return False
+
+
+# Call signatures (per family: shapes, dtypes, device, kernel-selecting flags — see the
+# family's `_warm_sig`) that have completed an eager forward / backward through our kernels.
+_WARM: set[tuple] = set()
+
+
+def mark_warm(sig: tuple, phase: str) -> None:
+    """Record that `sig` just ran `phase` ("fwd" or "bwd") outside any capture."""
+    if not capturing():
+        _WARM.add((*sig, phase))
+
+
+def capture_unsupported_reason(sig: tuple, needs_grad: bool) -> str | None:
+    """Under CUDA graph capture, our kernels run only for a shape that has already run eagerly.
+
+    Capture bakes every pointer the region touches; the call cache's launch-time pointer
+    poke is no different from what cuBLAS does, and with the static buffers a graphed
+    callable owns it replays correctly (verified bit-identical fwd+bwd, 2026-09-08). What
+    CANNOT happen inside a capture is a first call at a shape: `cute.compile`, and the
+    Triton autotuners' trial launches (ours and fla's both synchronize). Both are keyed on
+    shapes and dtypes, never on the stream, so a call-cache miss on a new stream for a warm
+    shape is only `from_dlpack` plus an allocation, and captures fine. Hence the stream is
+    deliberately NOT part of the signature: `make_graphed_callables` warms on one side
+    stream and captures on another, and keying on it would send every graphed call to fla
+    with nothing but a log line to say so.
+
+    The fallback for a cold shape is fla, which under capture has the same autotune
+    problem — but that is fla's failure to report, loudly, not ours to hide.
+    """
+    if not capturing():
+        return None
+    if (*sig, "fwd") not in _WARM:
+        return (
+            "CUDA graph capture of a shape that has not run eagerly yet (run one forward "
+            "at this shape before capturing — the compile and autotune steps cannot be "
+            "captured)"
+        )
+    if needs_grad and (*sig, "bwd") not in _WARM:
+        return (
+            "CUDA graph capture of a shape whose backward has not run eagerly yet (run one "
+            "fwd+bwd at this shape before capturing)"
+        )
+    return None
+
+
+def _basic_unsupported_reason(t: torch.Tensor, family: str) -> str | None:
+    if disabled(family):
+        return f"KERNEL_FUN_{family.upper()}_DISABLE / KERNEL_FUN_DISABLE is set"
+    if not t.is_cuda:
+        return "not a CUDA tensor"
+    # CUDA graph capture is gated per shape, after the shape checks: capture_unsupported_reason.
+    return None
+
+
+def common_unsupported_reason(t: torch.Tensor, family: str) -> str | None:
+    """The gates every CuTe family shares. Family-specific shape checks live with the family."""
+    reason = _basic_unsupported_reason(t, family)
+    if reason is not None:
+        return reason
+    if not arch_ok(t.device.index or 0):
+        cap = torch.cuda.get_device_capability(t.device.index or 0)
+        return f"device capability sm{cap[0]}{cap[1]} is not sm100 (B200/B300)"
+    if not has_cute():
+        return "the CUTLASS CuTe DSL is not installed"
+    return None
+
+
+def triton_unsupported_reason(t: torch.Tensor, family: str, min_major: int = 9) -> str | None:
+    """The same, for a family whose kernels are all Triton: no CuTe, a lower arch floor."""
+    reason = _basic_unsupported_reason(t, family)
+    if reason is not None:
+        return reason
+    if not arch_at_least(t.device.index or 0, min_major):
+        cap = torch.cuda.get_device_capability(t.device.index or 0)
+        return (
+            f"device capability sm{cap[0]}{cap[1]} is below sm{min_major}0 "
+            f"(measured on sm100; untimed below it)"
+        )
+    return None

kernel_fun/cconv/__init__.py

@@ -0,0 +1,24 @@
+"""cconv — the KDA short convolution: width-4 depthwise causal conv + silu.
+
+    from kernel_fun.cconv import causal_conv1d
+
+Same signature and return contract as `fla.modules.convolution.causal_conv1d`; anything this
+package does not implement is forwarded to fla, so the swap is one import line and
+reverting it is one line back.
+
+A KDA layer makes three of these calls per forward (q, k at D = n_heads*128, v at 2x that),
+and the 810m/B300 step trace put them at ~50 ms of a ~520 ms step — the largest non-GEMM
+item in the step, against an ~8-10 ms bandwidth roofline. Measured isolated on B300 at the
+production call (B=16, T=8192, D=2048, W=4, bf16 x / fp32 weight; 2026-09-01, the ladder's
+cconv/001 record 002):
+
+    backward (dx, dw; no forward re-run)   1.597 ms -> 0.344 ms   4.65x   (4.7 TB/s)
+    forward  (silu fused)                  0.307 ms -> 0.190 ms   1.62x   (5.6 TB/s)
+
+The same ratios hold at D=1024 and D=4096. Both kernels are Triton; see ops.py for the
+supported box and _kernels/strip.py for the design.
+"""
+
+from .ops import causal_conv1d, is_supported, warmup
+
+__all__ = ["causal_conv1d", "is_supported", "warmup"]

kernel_fun/cconv/_kernels/__init__.py

@@ -0,0 +1,6 @@
+"""Vendored kernels: copied from the research ladder at the commit in _provenance.py.
+
+These files are edited on arrival (env knobs frozen to the branch that won) and are
+otherwise the measured code. Edit them HERE only for packaging concerns; kernel work
+happens in the ladder, and comes back through tools/vendor.py.
+"""