chore: import upstream snapshot with attribution

tokenspeed-kernel/python/tokenspeed_kernel/ops/transform/__init__.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+# Backend registration (side-effect imports)
+import tokenspeed_kernel.ops.transform.faster_hadamard_transform  # noqa: F401
+import tokenspeed_kernel.ops.transform.triton  # noqa: F401
+import torch
+from tokenspeed_kernel.profiling import ShapeCapture, kernel_scope
+from tokenspeed_kernel.selection import select_kernel
+from tokenspeed_kernel.signature import dense_tensor_format, format_signature
+
+__all__ = ["hadamard_transform"]
+
+
+# ===-----------------------------------------------------------------------===#
+# Hadamard Transform Kernels
+# ===-----------------------------------------------------------------------===#
+
+
+def hadamard_transform(
+    x: torch.Tensor,
+    *,
+    scale: float = 1.0,
+    override: str | None = None,
+    solution: str | None = None,
+) -> torch.Tensor:
+    """Apply a Hadamard transform along the last dimension.
+
+    Args:
+        x: Input tensor. Registered implementations currently accept dense CUDA
+            tensors. The Triton implementation supports last dimension 128.
+        scale: Multiplicative output scale applied after the transform.
+        override: Optional exact kernel override name.
+        solution: Optional registered solution to force through normal selection.
+
+    Returns:
+        Tensor with the same shape and dtype as x.
+
+    """
+    if x.dim() == 0:
+        raise ValueError("hadamard_transform requires at least one dimension")
+
+    traits = {
+        "last_dim": x.shape[-1],
+    }
+    signature = format_signature(x=dense_tensor_format(x.dtype))
+    kernel = select_kernel(
+        "transform",
+        "hadamard_transform",
+        signature,
+        traits=traits,
+        solution=solution,
+        override=override,
+    )
+
+    shape_params = {
+        "shape": tuple(x.shape),
+        "last_dim": x.shape[-1],
+    }
+    ShapeCapture.get().record(
+        "transform",
+        "hadamard_transform",
+        kernel.name,
+        x.dtype,
+        shape_params,
+    )
+
+    with kernel_scope(
+        "transform",
+        "hadamard_transform",
+        x.dtype,
+        kernel_name=kernel.name,
+        **shape_params,
+    ):
+        return kernel(x, scale=scale)

tokenspeed-kernel/python/tokenspeed_kernel/ops/transform/faster_hadamard_transform/__init__.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel.platform import CapabilityRequirement, current_platform
+from tokenspeed_kernel.registry import Priority, register_kernel
+from tokenspeed_kernel.signature import format_signatures
+
+platform = current_platform()
+
+
+if platform.is_nvidia:
+    from fast_hadamard_transform import hadamard_transform
+
+    @register_kernel(
+        "transform",
+        "hadamard_transform",
+        name="fast_hadamard_transform",
+        solution="fast_hadamard_transform",
+        capability=CapabilityRequirement(vendors=frozenset({"nvidia"})),
+        signatures=format_signatures("x", "dense", {torch.bfloat16, torch.float16}),
+        priority=Priority.PERFORMANT,
+    )
+    def fast_hadamard_transform(
+        x: torch.Tensor,
+        *,
+        scale: float = 1.0,
+    ) -> torch.Tensor:
+        return hadamard_transform(x, scale=scale)
+
+
+__all__ = ["fast_hadamard_transform"]

tokenspeed-kernel/python/tokenspeed_kernel/ops/transform/triton/__init__.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+from tokenspeed_kernel.platform import CapabilityRequirement
+from tokenspeed_kernel.registry import Priority, register_kernel
+from tokenspeed_kernel.signature import format_signatures
+
+
+@triton.jit
+def _hadamard_128_kernel(
+    x,
+    out,
+    n_rows: tl.constexpr,
+    scale: tl.constexpr,
+    BLOCK_OUT: tl.constexpr,
+):
+    row = tl.program_id(0)
+    out_block = tl.program_id(1)
+    out_offsets = out_block * BLOCK_OUT + tl.arange(0, BLOCK_OUT)
+    in_offsets = tl.arange(0, 128)
+
+    vals = tl.load(x + row * 128 + in_offsets).to(tl.float32)
+    bits = out_offsets[:, None] & in_offsets[None, :]
+    parity = bits ^ (bits >> 1)
+    parity = parity ^ (parity >> 2)
+    parity = parity ^ (parity >> 4)
+    parity = parity & 1
+    signs = tl.where(parity == 0, 1.0, -1.0)
+    acc = tl.sum(signs * vals[None, :], axis=1) * scale
+
+    tl.store(
+        out + row * 128 + out_offsets,
+        acc,
+        mask=(row < n_rows) & (out_offsets < 128),
+    )
+
+
+@register_kernel(
+    "transform",
+    "hadamard_transform",
+    name="triton_hadamard_transform_128",
+    solution="triton",
+    capability=CapabilityRequirement(vendors=frozenset({"nvidia", "amd"})),
+    signatures=format_signatures(
+        "x",
+        "dense",
+        {torch.bfloat16, torch.float16, torch.float32},
+    ),
+    traits={
+        "last_dim": frozenset({128}),
+    },
+    priority=Priority.PORTABLE,
+)
+def triton_hadamard_transform_128(
+    x: torch.Tensor,
+    *,
+    scale: float = 1.0,
+) -> torch.Tensor:
+    """Apply a length-128 Sylvester Hadamard transform along the last dim."""
+    if x.shape[-1] != 128:
+        raise ValueError(
+            f"triton_hadamard_transform_128 requires last dim 128, got {x.shape[-1]}"
+        )
+    if not x.is_cuda:
+        raise RuntimeError("triton_hadamard_transform_128 requires a CUDA tensor")
+    if x.dtype not in (torch.bfloat16, torch.float16, torch.float32):
+        raise TypeError(
+            f"triton_hadamard_transform_128 does not support dtype {x.dtype}"
+        )
+
+    shape = x.shape
+    x_2d = x.reshape(-1, 128).contiguous()
+    out = torch.empty_like(x_2d)
+    if x_2d.shape[0] == 0:
+        return out.reshape(shape)
+    _hadamard_128_kernel[(x_2d.shape[0], 8)](
+        x_2d,
+        out,
+        n_rows=x_2d.shape[0],
+        scale=float(scale),
+        BLOCK_OUT=16,
+        num_warps=8,
+        num_stages=1,
+    )
+    return out.reshape(shape)
+
+
+__all__ = ["triton_hadamard_transform_128"]

tokenspeed-kernel/python/tokenspeed_kernel/ops/transform/weak_ref.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Non-owning tensor view for CUDA-graph break closures.
+
+``weak_ref_tensor(t)`` returns a tensor aliasing ``t``'s memory WITHOUT owning
+its storage (``at::from_blob`` with a no-op deleter). A breakable-CUDA-graph
+break closure that holds such a view does not pin ``t``'s mempool block, so the
+graph pool can recycle blocks across segments and captures -- graph memory
+drops from sum-over-buckets of break inputs back to ~peak-live. Correctness
+relies on replay order: the captured segment rewrites the aliased address
+before the break reads it, every replay (the same discipline that lets the
+decode graph share one pool across batch sizes).
+
+Adapted from vLLM ``csrc/ops.h`` / SGLang ``sgl-kernel/csrc/memory/
+weak_ref_tensor.cpp``. Compiled lazily via ``torch.utils.cpp_extension``
+(pure C++ against ATen -- no nvcc, no wheel rebuild); falls back to identity
+(strong ref: correct, more memory) if compilation is unavailable.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_CPP_SOURCE = r"""
+#include <ATen/ATen.h>
+#include <vector>
+
+at::Tensor weak_ref_tensor(const at::Tensor& tensor) {
+  TORCH_CHECK(tensor.is_cuda(), "weak_ref_tensor expects a CUDA tensor");
+  void* data_ptr = tensor.data_ptr();
+  std::vector<int64_t> sizes = tensor.sizes().vec();
+  std::vector<int64_t> strides = tensor.strides().vec();
+  auto options = tensor.options();
+  return at::from_blob(data_ptr, sizes, strides, options);
+}
+"""
+
+_module = None
+_load_failed = False
+
+
+def _load():
+    global _module, _load_failed
+    if _module is not None or _load_failed:
+        return _module
+    try:
+        from torch.utils.cpp_extension import load_inline
+
+        _module = load_inline(
+            name="tokenspeed_weak_ref_tensor",
+            cpp_sources=_CPP_SOURCE,
+            functions=["weak_ref_tensor"],
+            with_cuda=False,
+            verbose=False,
+        )
+    except Exception as exc:  # pragma: no cover - host without a C++ toolchain
+        _load_failed = True
+        logger.warning(
+            "weak_ref_tensor extension unavailable (%s: %s); falling back to "
+            "identity (strong refs pin graph-pool blocks -- correct, but "
+            "breakable-graph capture memory scales with the bucket sum).",
+            type(exc).__name__,
+            exc,
+        )
+    return _module
+
+
+def weak_ref_tensor(t: torch.Tensor) -> torch.Tensor:
+    """Return a non-owning view of CUDA tensor ``t`` (alias, same shape/strides).
+
+    Args:
+        t: A CUDA tensor. The caller must guarantee the aliased memory is
+            rewritten (or still valid) whenever the view is consumed -- for
+            breakable-graph closures this holds because the captured segment
+            producing ``t`` replays before the break reads the view.
+
+    Returns:
+        A tensor sharing ``t``'s memory without owning its storage, or ``t``
+        itself if the extension is unavailable (identity fallback).
+    """
+    mod = _load()
+    if mod is None:
+        return t
+    return mod.weak_ref_tensor(t)