tile-kernels 1.0.0__tar.gz

tile_kernels-1.0.0/.editorconfig

@@ -0,0 +1,58 @@
+# https://editorconfig.org/
+
+root = true
+
+[*]
+charset = utf-8
+end_of_line = lf
+indent_style = space
+indent_size = 4
+trim_trailing_whitespace = true
+insert_final_newline = true
+
+[*.{py,pyi}]
+indent_size = 4
+
+[*.{cpp,hpp,cxx,cc,c,h,cu,cuh}]
+indent_size = 4
+
+[*.rs]
+indent_size = 4
+
+[*.go]
+indent_style = tab
+
+[*.{yaml,yml}]
+indent_size = 2
+
+[.clang-{format,tidy}]
+indent_size = 2
+
+[Makefile]
+indent_style = tab
+
+[*.sh]
+indent_size = 4
+
+[*.bat]
+indent_size = 4
+end_of_line = crlf
+
+[*.md]
+indent_size = 2
+x-soft-wrap-text = true
+
+[*.rst]
+indent_size = 4
+x-soft-wrap-text = true
+
+[*.{html,xml,css,scss,js,jsx,ts,tsx,vue}]
+indent_size = 2
+
+[**/test{,s,ing}/**/*.txt]
+trim_trailing_whitespace = false
+insert_final_newline = false
+
+[**/example{,s}/**/*.txt]
+trim_trailing_whitespace = false
+insert_final_newline = false

tile_kernels-1.0.0/.gitignore

@@ -0,0 +1,20 @@
+.DS_Store
+build
+dist
+*.egg-info
+*.pyc
+__pycache__/
+
+# PyCharm Settings
+.idea/
+*.iml
+*.iws
+
+# Python packaging version file
+_version.py
+
+# VSCode Settings
+/.vscode
+
+# workspace
+workspace

tile_kernels-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 DeepSeek
+
+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.

tile_kernels-1.0.0/PKG-INFO

@@ -0,0 +1,116 @@
+Metadata-Version: 2.4
+Name: tile_kernels
+Version: 1.0.0
+Summary: Tilelang-based kernels.
+Author-email: Chenhao Xu <xch@deepseek.com>, Xiangwen Wang <xiangwen.wang@deepseek.com>, Huanqi Cao <caohuanqi@deepseek.com>, Rui Tian <tianr22@deepseek.com>, Weilin Zhao <zhaoweilin@deepseek.com>, Kuai Yu <yukuai@deepseek.com>, Chenggang Zhao <chenggangz@deepseek.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/deepseek-ai/TileKernels
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.10
+Requires-Dist: tilelang>=0.1.9
+Provides-Extra: dev
+Requires-Dist: setuptools; extra == "dev"
+Requires-Dist: wheel; extra == "dev"
+Requires-Dist: setuptools-scm>=8; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-xdist; extra == "dev"
+Requires-Dist: pytest-repeat; extra == "dev"
+Dynamic: license-file
+
+# Tile Kernels
+
+Optimized GPU kernels for LLM operations, built with [TileLang](https://github.com/tile-ai/tilelang). TileLang is a domain-specific language for expressing high-performance GPU kernels in Python, featuring easy migration, agile development, and automatic optimization.
+
+Most kernels in this project approach the limit of hardware performance regarding the compute intensity and memory bandwidth. Some of them have already been used in internal training and inference scenarios. However, they do not represent best practices and we are actively working on improving the code quality and documentation.
+
+## Features
+
+- **Gating** — Top-k expert selection and scoring for Mixture of Experts routing
+- **MoE Routing** — Token-to-expert mapping, fused expansion/reduction and weight normalization
+- **Quantization** — Per-token, per-block, and per-channel FP8/FP4/E5M6 casting with fused SwiGLU+quantization ops
+- **Transpose** — Batched transpose operations
+- **Engram** — Engram gating kernels with fused RMSNorm, forward/backward passes and weight gradient reduction
+- **Manifold HyperConnection** — Hyper-connection kernels including Sinkhorn normalization and mix splitting/application
+- **Modeling** — High-level `torch.autograd.Function` wrappers composing low-level kernels into trainable layers (engram gate, mHC pipeline)
+
+## Requirements
+
+- Python 3.10 or higher
+- PyTorch 2.10 or higher
+- TileLang 0.1.9 or higher
+- NVIDIA SM90 or SM100 architecture GPU
+- CUDA Toolkit 13.1 or higher
+
+## Installation
+
+### Install a local development version
+
+```bash
+pip install -e ".[dev]"
+```
+
+### Install a release version
+
+```bash
+pip install tile-kernels
+```
+
+## Testing
+
+Tests using pytest:
+
+### Test single test file
+
+```bash
+pytest tests/transpose/test_transpose.py -n 4 # Correctness only with 4 workers
+pytest tests/transpose/test_transpose.py --run-benchmark # Correctness + Benchmarking
+```
+
+### Pressure test
+
+```bash
+TK_FULL_TEST=1 pytest -n 4 --count 2
+```
+
+## Project Structure
+
+```txt
+tile_kernels/
+├── moe/        # Mixture of Experts routing related kernels
+├── quant/      # FP8/FP4/E5M6 quantization
+├── transpose/  # Batched transpose
+├── engram/     # Engram gating kernels
+├── mhc/        # Manifold HyperConnection kernels
+├── modeling/   # High-level autograd modeling layers (engram, mHC)
+├── torch/      # PyTorch reference implementations
+└── testing/    # Test and benchmark utilities
+```
+
+## Acknowledgement
+
+This project is built on [TileLang](https://github.com/tile-ai/tilelang). Thanks and respect to the developers!
+
+## License
+
+This code repository is released under [the MIT License](LICENSE).
+
+## Citation
+
+```bibtex
+@misc{tilekernels,
+      title={TileKernels},
+      author={Xiangwen Wang, Chenhao Xu, Huanqi Cao, Rui Tian, Weilin Zhao, Kuai Yu and Chenggang Zhao},
+      year={2026},
+      publisher = {GitHub},
+      howpublished = {\url{https://github.com/deepseek-ai/TileKernels}},
+}
+```

tile_kernels-1.0.0/README.md

@@ -0,0 +1,88 @@
+# Tile Kernels
+
+Optimized GPU kernels for LLM operations, built with [TileLang](https://github.com/tile-ai/tilelang). TileLang is a domain-specific language for expressing high-performance GPU kernels in Python, featuring easy migration, agile development, and automatic optimization.
+
+Most kernels in this project approach the limit of hardware performance regarding the compute intensity and memory bandwidth. Some of them have already been used in internal training and inference scenarios. However, they do not represent best practices and we are actively working on improving the code quality and documentation.
+
+## Features
+
+- **Gating** — Top-k expert selection and scoring for Mixture of Experts routing
+- **MoE Routing** — Token-to-expert mapping, fused expansion/reduction and weight normalization
+- **Quantization** — Per-token, per-block, and per-channel FP8/FP4/E5M6 casting with fused SwiGLU+quantization ops
+- **Transpose** — Batched transpose operations
+- **Engram** — Engram gating kernels with fused RMSNorm, forward/backward passes and weight gradient reduction
+- **Manifold HyperConnection** — Hyper-connection kernels including Sinkhorn normalization and mix splitting/application
+- **Modeling** — High-level `torch.autograd.Function` wrappers composing low-level kernels into trainable layers (engram gate, mHC pipeline)
+
+## Requirements
+
+- Python 3.10 or higher
+- PyTorch 2.10 or higher
+- TileLang 0.1.9 or higher
+- NVIDIA SM90 or SM100 architecture GPU
+- CUDA Toolkit 13.1 or higher
+
+## Installation
+
+### Install a local development version
+
+```bash
+pip install -e ".[dev]"
+```
+
+### Install a release version
+
+```bash
+pip install tile-kernels
+```
+
+## Testing
+
+Tests using pytest:
+
+### Test single test file
+
+```bash
+pytest tests/transpose/test_transpose.py -n 4 # Correctness only with 4 workers
+pytest tests/transpose/test_transpose.py --run-benchmark # Correctness + Benchmarking
+```
+
+### Pressure test
+
+```bash
+TK_FULL_TEST=1 pytest -n 4 --count 2
+```
+
+## Project Structure
+
+```txt
+tile_kernels/
+├── moe/        # Mixture of Experts routing related kernels
+├── quant/      # FP8/FP4/E5M6 quantization
+├── transpose/  # Batched transpose
+├── engram/     # Engram gating kernels
+├── mhc/        # Manifold HyperConnection kernels
+├── modeling/   # High-level autograd modeling layers (engram, mHC)
+├── torch/      # PyTorch reference implementations
+└── testing/    # Test and benchmark utilities
+```
+
+## Acknowledgement
+
+This project is built on [TileLang](https://github.com/tile-ai/tilelang). Thanks and respect to the developers!
+
+## License
+
+This code repository is released under [the MIT License](LICENSE).
+
+## Citation
+
+```bibtex
+@misc{tilekernels,
+      title={TileKernels},
+      author={Xiangwen Wang, Chenhao Xu, Huanqi Cao, Rui Tian, Weilin Zhao, Kuai Yu and Chenggang Zhao},
+      year={2026},
+      publisher = {GitHub},
+      howpublished = {\url{https://github.com/deepseek-ai/TileKernels}},
+}
+```

tile_kernels-1.0.0/pyproject.toml

@@ -0,0 +1,58 @@
+[build-system]
+requires = ["setuptools", "wheel", "setuptools-scm>=8"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools_scm]
+version_file = "tile_kernels/_version.py"
+
+[project]
+name = "tile_kernels"
+dynamic = ["version"]
+description = "Tilelang-based kernels."
+readme = "README.md"
+license = "MIT"
+authors = [
+    { name = "Chenhao Xu", email = "xch@deepseek.com" },
+    { name = "Xiangwen Wang", email = "xiangwen.wang@deepseek.com" },
+    { name = "Huanqi Cao", email = "caohuanqi@deepseek.com" },
+    { name = "Rui Tian", email = "tianr22@deepseek.com"},
+    { name = "Weilin Zhao", email = "zhaoweilin@deepseek.com" },
+    { name = "Kuai Yu", email = "yukuai@deepseek.com" },
+    { name = "Chenggang Zhao", email = "chenggangz@deepseek.com"}
+]
+dependencies = [
+    "torch>=2.10",
+    "tilelang>=0.1.9"
+]
+requires-python = ">=3.10"
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.optional-dependencies]
+dev = ["setuptools", "wheel", "setuptools-scm>=8", "pytest", "pytest-xdist", "pytest-repeat"]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["tile_kernels*"]
+
+[project.urls]
+Homepage = "https://github.com/deepseek-ai/TileKernels"
+
+# Linter tools
+
+[tool.ruff]
+line-length = 150
+
+[tool.ruff.lint]
+select = ["Q000"]
+fixable = ["Q000"]
+
+[tool.ruff.lint.flake8-quotes]
+inline-quotes = "single"

tile_kernels-1.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tile_kernels-1.0.0/tests/conftest.py

@@ -0,0 +1,10 @@
+# Root-level conftest
+#
+# Loads the benchmark plugin (CLI options, markers, fixtures).
+# The plugin lives in a file deliberately NOT named conftest.py to
+# avoid pluggy's duplicate-registration error.
+
+pytest_plugins = [
+    'tests.pytest_random_plugin',
+    'tests.pytest_benchmark_plugin',
+]

tile_kernels-1.0.0/tests/engram/test_engram_fused_weight.py

@@ -0,0 +1,56 @@
+import os
+import pytest
+import torch
+
+from tile_kernels.engram import fused_weight
+from tile_kernels.testing.numeric import assert_equal, count_bytes
+from tile_kernels.testing.generator import generate_hidden_sizes
+from tile_kernels.testing.bench import make_param_id
+
+# Disable TileLang prints
+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
+
+
+def generate_test_data(params):
+    hc_mult = params['hc']
+    hidden_size = params['hidden']
+    wh_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    we_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    return (wh_data, we_data)
+
+
+def generate_test_params(is_benchmark: bool) -> list[dict]:
+    return [
+        {'hc': hc, 'hidden': hidden_size}
+        for hc in (4,)
+        for hidden_size in generate_hidden_sizes(128)
+    ]
+
+
+@pytest.mark.parametrize('params', generate_test_params(is_benchmark=False), ids=make_param_id)
+def test_engram_fused_weight(params):
+    wh_data, we_data = generate_test_data(params)
+
+    ref = wh_data.float() * we_data.float()
+    out = fused_weight(wh_data, we_data)
+
+    assert_equal(out, ref)
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize('params', generate_test_params(is_benchmark=True), ids=make_param_id)
+def test_engram_fused_weight_benchmark(benchmark_timer, benchmark_record, params):
+    wh_data, we_data = generate_test_data(params)
+    out = fused_weight(wh_data, we_data)
+
+    t_us = benchmark_timer(lambda: fused_weight(wh_data, we_data))
+
+    num_bytes = count_bytes(wh_data, we_data, out)
+    bandwidth_gbs = num_bytes / t_us / 1e3
+    benchmark_record(
+        kernel='fused_weight',
+        operation='fwd',
+        params=params,
+        time_us=t_us,
+        bandwidth_gbs=bandwidth_gbs,
+    )

tile_kernels-1.0.0/tests/engram/test_engram_gate_bwd.py

@@ -0,0 +1,112 @@
+import os
+import pytest
+import torch
+
+from tile_kernels.engram import engram_gate_bwd
+from tile_kernels.torch.engram import engram_gate_ref
+from tile_kernels.testing.numeric import calc_diff, count_bytes
+from tile_kernels.testing.generator import generate_hidden_sizes, generate_num_tokens
+from tile_kernels.testing.bench import make_param_id
+
+# Disable TileLang prints
+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
+
+
+def generate_test_data(params):
+    num_tokens = params['num_tokens']
+    hc_mult = params['hc']
+    hidden_size = params['hidden']
+    eps = 1e-20
+    clamp_value = 1e-6
+    x_data = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    k_data = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    v_data = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device='cuda')
+    wh_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    we_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    weight_fused = wh_data.float() * we_data.float()
+    grad_out = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    return (x_data, k_data, v_data, wh_data, we_data, weight_fused, grad_out, eps, clamp_value)
+
+
+def generate_test_params(is_benchmark: bool) -> list[dict]:
+    return [
+        {'num_tokens': t, 'hc': hc, 'hidden': hidden_size}
+        for t in generate_num_tokens(is_benchmark=is_benchmark)
+        for hc in (4,)
+        for hidden_size in generate_hidden_sizes(128)
+    ]
+
+
+@pytest.mark.parametrize('params', generate_test_params(is_benchmark=False), ids=make_param_id)
+def test_engram_gate_bwd(params):
+    (x_data, k_data, v_data, wh_data, we_data, weight_fused, grad_out, eps, clamp_value) = generate_test_data(params)
+
+    # Reference: forward with intermediates + autograd backward
+    x_ref = x_data.clone().requires_grad_(True)
+    k_ref = k_data.clone().requires_grad_(True)
+    v_ref = v_data.clone().requires_grad_(True)
+    # Cast to float32 so autograd produces fp32 gradients matching the kernel
+    wh_ref = wh_data.float().requires_grad_(True)
+    we_ref = we_data.float().requires_grad_(True)
+    o_ref, dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref = engram_gate_ref(
+        x_ref, k_ref, v_ref, wh_ref, we_ref, clamp_value, eps, save_for_backward=True,
+    )
+    o_ref.backward(grad_out)
+
+    # Kernel backward using ref intermediates
+    grad_x, grad_k, grad_v, grad_w_partial = engram_gate_bwd(
+        grad_out, x_data, k_data, v_data, weight_fused,
+        dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref, clamp_value,
+    )
+    grad_w_fused = grad_w_partial.sum(0)
+    grad_wh = grad_w_fused * we_data.float()
+    grad_we = grad_w_fused * wh_data.float()
+
+    # Correctness
+    diff_x = calc_diff(grad_x, x_ref.grad)
+    assert diff_x < 1e-8, f'grad_x mismatch: {diff_x:.6e}'
+    diff_k = calc_diff(grad_k, k_ref.grad)
+    assert diff_k < 1e-8, f'grad_k mismatch: {diff_k:.6e}'
+    diff_v = calc_diff(grad_v, v_ref.grad)
+    assert diff_v < 1e-8, f'grad_v mismatch: {diff_v:.6e}'
+    diff_wh = calc_diff(grad_wh, wh_ref.grad)
+    assert diff_wh < 1e-8, f'grad_wh mismatch: {diff_wh:.6e}'
+    diff_we = calc_diff(grad_we, we_ref.grad)
+    assert diff_we < 1e-8, f'grad_we mismatch: {diff_we:.6e}'
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize('params', generate_test_params(is_benchmark=True), ids=make_param_id)
+def test_engram_gate_bwd_benchmark(benchmark_timer, benchmark_record, params):
+    (x_data, k_data, v_data, wh_data, we_data, weight_fused, grad_out, eps, clamp_value) = generate_test_data(params)
+
+    # Forward to get intermediates
+    o_ref, dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref = engram_gate_ref(
+        x_data, k_data, v_data, wh_data, we_data, clamp_value, eps, save_for_backward=True,
+    )
+
+    grad_x, grad_k, grad_v, grad_w_partial = engram_gate_bwd(
+        grad_out, x_data, k_data, v_data, weight_fused,
+        dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref, clamp_value,
+    )
+    grad_w_fused = grad_w_partial.sum(0)
+    grad_wh = grad_w_fused * we_data.float()
+    grad_we = grad_w_fused * wh_data.float()
+
+    func_bwd = lambda: engram_gate_bwd(
+        grad_out, x_data, k_data, v_data, weight_fused,
+        dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref, clamp_value,
+    )
+    t_us = benchmark_timer(func_bwd)
+    num_bytes = count_bytes(
+        grad_out, x_data, k_data, v_data, weight_fused,
+        dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref,
+        grad_x, grad_k, grad_v, grad_wh, grad_we,
+    )
+    benchmark_record(
+        kernel='engram_gate_bwd',
+        operation='bwd',
+        params=params,
+        time_us=t_us,
+        bandwidth_gbs=num_bytes / t_us / 1e3,
+    )

tile_kernels-1.0.0/tests/engram/test_engram_gate_fwd.py

@@ -0,0 +1,108 @@
+import os
+import pytest
+import torch
+
+from tile_kernels.engram import engram_gate_fwd
+from tile_kernels.torch.engram import engram_gate_ref
+from tile_kernels.testing.numeric import assert_equal, calc_diff, count_bytes
+from tile_kernels.testing.generator import generate_hidden_sizes, generate_num_tokens
+from tile_kernels.testing.bench import make_param_id
+
+# Disable TileLang prints
+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
+
+
+def generate_test_data(params):
+    num_tokens = params['num_tokens']
+    hc_mult = params['hc']
+    hidden_size = params['hidden']
+    eps = 1e-20
+    clamp_value = 1e-6
+    x_data = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    k_data = torch.randn(num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    v_data = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device='cuda')
+    wh_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    we_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device='cuda')
+    weight_fused = wh_data.float() * we_data.float()
+    return (x_data, k_data, v_data, wh_data, we_data, weight_fused, eps, clamp_value)
+
+
+def generate_test_params(is_benchmark: bool) -> list[dict]:
+    return [
+        {'num_tokens': t, 'hc': hc, 'hidden': hidden_size}
+        for t in generate_num_tokens(is_benchmark=is_benchmark)
+        for hc in (4,)
+        for hidden_size in generate_hidden_sizes(128)
+    ]
+
+
+@pytest.mark.parametrize('params', generate_test_params(is_benchmark=False), ids=make_param_id)
+def test_engram_gate_fwd(params):
+    (x_data, k_data, v_data, wh_data, we_data, weight_fused, eps, clamp_value) = generate_test_data(params)
+
+    out_ref, dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref = engram_gate_ref(
+        x_data, k_data, v_data, wh_data, we_data, clamp_value, eps, save_for_backward=True,
+    )
+
+    # Correctness: save_for_backward=True
+    out_save, dot, gate_score, rstd_x, rstd_k = engram_gate_fwd(
+        x_data, k_data, v_data, weight_fused, eps, clamp_value, save_for_backward=True,
+    )
+    assert dot is not None and gate_score is not None and rstd_x is not None and rstd_k is not None
+    diff_out = calc_diff(out_save, out_ref)
+    assert diff_out < 2e-10, f'out_save mismatch: {diff_out:.6e}'
+    diff_dot = calc_diff(dot, dot_ref)
+    assert diff_dot < 2e-10, f'dot mismatch: {diff_dot:.6e}'
+    diff_gate = calc_diff(gate_score, gate_score_ref)
+    assert diff_gate < 2e-10, f'gate_score mismatch: {diff_gate:.6e}'
+    diff_rstd_x = calc_diff(rstd_x, rstd_x_ref)
+    assert diff_rstd_x < 2e-10, f'rstd_x mismatch: {diff_rstd_x:.6e}'
+    diff_rstd_k = calc_diff(rstd_k, rstd_k_ref)
+    assert diff_rstd_k < 2e-10, f'rstd_k mismatch: {diff_rstd_k:.6e}'
+
+    # Correctness: save_for_backward=False
+    out_no_save, dot_n, gate_score_n, rstd_x_n, rstd_k_n = engram_gate_fwd(
+        x_data, k_data, v_data, weight_fused, eps, clamp_value, save_for_backward=False,
+    )
+    assert dot_n is None and gate_score_n is None and rstd_x_n is None and rstd_k_n is None
+    diff_out = calc_diff(out_no_save, out_ref)
+    assert diff_out < 2e-10, f'out_no_save mismatch: {diff_out:.6e}'
+    assert_equal(out_no_save, out_save)
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize('params', generate_test_params(is_benchmark=True), ids=make_param_id)
+def test_engram_gate_fwd_benchmark(benchmark_timer, benchmark_record, params):
+    (x_data, k_data, v_data, _, _, weight_fused, eps, clamp_value) = generate_test_data(params)
+
+    # Benchmark save_for_backward=True
+    out_save, dot, gate_score, rstd_x, rstd_k = engram_gate_fwd(
+        x_data, k_data, v_data, weight_fused, eps, clamp_value, save_for_backward=True,
+    )
+    t_save_us = benchmark_timer(lambda: engram_gate_fwd(
+        x_data, k_data, v_data, weight_fused, eps, clamp_value, save_for_backward=True,
+    ))
+    num_bytes_save = count_bytes(x_data, k_data, v_data, weight_fused, out_save, dot, gate_score, rstd_x, rstd_k)
+    benchmark_record(
+        kernel='engram_gate_fwd',
+        operation='fwd',
+        params={**params, 'save': True},
+        time_us=t_save_us,
+        bandwidth_gbs=num_bytes_save / t_save_us / 1e3,
+    )
+
+    # Benchmark save_for_backward=False
+    out_no_save = engram_gate_fwd(
+        x_data, k_data, v_data, weight_fused, eps, clamp_value, save_for_backward=False,
+    )[0]
+    t_no_save_us = benchmark_timer(lambda: engram_gate_fwd(
+        x_data, k_data, v_data, weight_fused, eps, clamp_value, save_for_backward=False,
+    ))
+    num_bytes_no_save = count_bytes(x_data, k_data, v_data, weight_fused, out_no_save)
+    benchmark_record(
+        kernel='engram_gate_fwd',
+        operation='fwd',
+        params={**params, 'save': False},
+        time_us=t_no_save_us,
+        bandwidth_gbs=num_bytes_no_save / t_no_save_us / 1e3,
+    )