[misc]: harden the rotary cache hit path

.buildkite/pipeline.yml

@@ -481,6 +481,7 @@ steps:
                 - "fastvideo/layers/**"
                 - "fastvideo/worker/**"
                 - "fastvideo/entrypoints/**"
+                - "fastvideo/performance/**"
                 - "fastvideo/tests/performance/**"
                 - ".buildkite/performance-benchmarks/**"
                 - "pyproject.toml"

.github/scripts/gate_full_suite.sh

@@ -0,0 +1,133 @@
+#!/usr/bin/env bash
+# Gate the expensive Buildkite full suite on the cheap GitHub checks.
+#
+# Polls the workflow runs for the PR head commit and only exits 0 once the
+# watched cheap workflows (pre-commit, docs build) have succeeded, so the
+# 'ready' label cannot burn ~20 GPU lanes on a head that a cheap check has
+# already doomed.
+#
+# Semantics:
+#   - watched run completed with a bad conclusion  -> exit 1 (fail CLOSED:
+#     no full suite; the next push re-arms via the 'synchronize' trigger)
+#   - watched run cancelled                         -> still pending: the docs
+#     workflow's repo-global 'pages' concurrency group cancels runs superseded
+#     by unrelated pushes, so 'cancelled' is not a verdict on this PR
+#   - watched runs pending                          -> poll until done
+#   - docs run absent                               -> not applicable after a
+#     short grace period ('Deploy Documentation' is path-filtered on PRs)
+#   - pre-commit run absent                         -> keep polling: pre-commit
+#     is never path-filtered, so its absence is always anomalous
+#   - 'ready' label removed while waiting           -> exit 1 (fail CLOSED:
+#     un-labeling is a deliberate maintainer action)
+#   - GitHub API unreachable or timeout             -> exit 0 (fail OPEN,
+#     loud warning: never brick CI on a GitHub outage)
+#
+# Required env: PR_SHA (PR head commit), PR_NUMBER, GITHUB_REPOSITORY, GH_TOKEN.
+set -euo pipefail
+
+: "${PR_SHA:?PR_SHA (PR head commit) is required}"
+: "${PR_NUMBER:?PR_NUMBER (pull request number) is required}"
+: "${GITHUB_REPOSITORY:?GITHUB_REPOSITORY is required}"
+
+# Workflow-level `name:` values that must be green before the full suite
+# may start. "Deploy Documentation" is path-filtered on PRs, so its run may
+# legitimately never exist; pre-commit always runs, so it must appear.
+WATCHED_NAMES='["pre-commit", "Deploy Documentation"]'
+WATCHED_REGEX='^(pre-commit|Deploy Documentation)$'
+POLL_SECS="${POLL_SECS:-20}"
+GRACE_SECS="${GRACE_SECS:-60}"
+MAX_WAIT_SECS="${MAX_WAIT_SECS:-1500}"
+
+# Bound each API call so a hung connection hits the 3-strike fail-open path
+# instead of pinning the loop until the job timeout (which would fail closed
+# on exactly the GitHub-outage case this script is meant to survive).
+if command -v timeout >/dev/null 2>&1; then
+  gh_api() { timeout 30 gh api "$@"; }
+else
+  gh_api() { gh api "$@"; }  # macOS dev boxes; CI always has coreutils timeout
+fi
+
+# The workflow checked the label before starting the gate, but the wait can
+# last ~25 min: re-check once before any exit 0 and fail closed if 'ready'
+# was removed in the meantime. An API error here proceeds (the label was
+# present when the gate started; never brick CI on an outage).
+recheck_ready_label() {
+  local pr_json
+  if pr_json=$(gh_api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" 2>/dev/null); then
+    if ! jq -e '[.labels[]?.name] | index("ready")' <<<"$pr_json" >/dev/null 2>&1; then
+      echo "::error::PR #${PR_NUMBER} no longer has the 'ready' label —" \
+        "NOT triggering the Buildkite full suite. Re-add the label to re-arm."
+      exit 1
+    fi
+  else
+    echo "::warning::Could not re-check the 'ready' label on PR #${PR_NUMBER}; proceeding (it was present when the gate started)."
+  fi
+}
+
+start=$(date +%s)
+api_fails=0
+missing=""
+
+while true; do
+  elapsed=$(( $(date +%s) - start ))
+
+  if runs_json=$(gh_api "repos/${GITHUB_REPOSITORY}/actions/runs?head_sha=${PR_SHA}&per_page=100" 2>/dev/null) \
+     && state=$(jq --arg re "$WATCHED_REGEX" '
+          [.workflow_runs[]? | select(.name // "" | test($re))]
+          | group_by(.name) | map(max_by(.id))
+          | map({name, status, conclusion})' <<<"$runs_json" 2>/dev/null); then
+    api_fails=0
+    echo "t+${elapsed}s watched checks: $(jq -c . <<<"$state")"
+
+    failed=$(jq -r '[.[] | select(.status == "completed"
+          and (.conclusion | IN("success", "skipped", "neutral", "cancelled") | not))]
+        | map(.name) | join(", ")' <<<"$state")
+    if [ -n "$failed" ]; then
+      echo "::error::Cheap check(s) failed on ${PR_SHA}: ${failed}." \
+        "NOT triggering the Buildkite full suite. Push a fix (the 'ready'" \
+        "label re-arms on every push), or re-run the failed check and then" \
+        "re-run this workflow."
+      exit 1
+    fi
+
+    # 'cancelled' counts as pending: wait for a re-run to reach a real verdict
+    # (bounded by MAX_WAIT, then the fail-open below).
+    pending=$(jq '[.[] | select(.status != "completed" or .conclusion == "cancelled")] | length' <<<"$state")
+    missing=$(jq -r --argjson watched "$WATCHED_NAMES" '($watched - map(.name)) | join(", ")' <<<"$state")
+    if [ "$pending" -eq 0 ]; then
+      if [ -z "$missing" ]; then
+        recheck_ready_label
+        echo "All watched cheap checks are green — full suite may proceed."
+        exit 0
+      fi
+      case "$missing" in
+        *pre-commit*)
+          echo "pre-commit run not found for ${PR_SHA} yet; waiting (pre-commit is never path-filtered, so its absence is anomalous)."
+          ;;
+        *)
+          if [ "$elapsed" -ge "$GRACE_SECS" ]; then
+            recheck_ready_label
+            echo "::warning::Watched run(s) never appeared for ${PR_SHA}: ${missing} (path-filtered, likely not applicable). Proceeding on the checks that did run."
+            exit 0
+          fi
+          echo "Waiting up to ${GRACE_SECS}s grace for path-filtered run(s) to appear: ${missing}."
+          ;;
+      esac
+    fi
+  else
+    api_fails=$(( api_fails + 1 ))
+    echo "::warning::GitHub API error querying workflow runs for ${PR_SHA} (attempt ${api_fails}/3)."
+    if [ "$api_fails" -ge 3 ]; then
+      recheck_ready_label
+      echo "::warning::FAILING OPEN: cannot query GitHub check status — triggering the full suite WITHOUT the cheap-check gate."
+      exit 0
+    fi
+  fi
+
+  if [ "$elapsed" -ge "$MAX_WAIT_SECS" ]; then
+    recheck_ready_label
+    echo "::warning::FAILING OPEN: watched checks still pending after $(( MAX_WAIT_SECS / 60 )) min${missing:+ (never appeared: ${missing})} — triggering the full suite anyway."
+    exit 0
+  fi
+  sleep "$POLL_SECS"
+done

.github/scripts/test_gate_full_suite.sh

@@ -0,0 +1,122 @@
+#!/usr/bin/env bash
+# Self-test for gate_full_suite.sh using a mocked `gh`. No network, runs on
+# any dev box: bash .github/scripts/test_gate_full_suite.sh
+set -u
+here=$(cd "$(dirname "$0")" && pwd)
+tmp=$(mktemp -d)
+trap 'rm -rf "$tmp"' EXIT
+
+# Mock gh. Asserts the exact endpoint (including head_sha) it is called
+# with — an endpoint typo in the gate script fails the test rather than
+# silently serving canned data. On the runs endpoint it serves
+# $MOCK_DIR/response_<call#>.json, sticking on the highest existing file,
+# and exits 1 if none exist (simulates a GitHub API outage). On the pulls
+# endpoint it serves $MOCK_DIR/pr.json, defaulting to a 'ready'-labeled PR.
+cat > "$tmp/gh" <<'EOF'
+#!/usr/bin/env bash
+if [ "${1:-}" != "api" ]; then
+  echo "unexpected gh invocation: $*" >> "$MOCK_DIR/endpoint_error"
+  exit 2
+fi
+case "${2:-}" in
+  "repos/o/r/actions/runs?head_sha=deadbeef&per_page=100")
+    n=$(( $(cat "$MOCK_DIR/count" 2>/dev/null || echo 0) + 1 ))
+    echo "$n" > "$MOCK_DIR/count"
+    while [ "$n" -gt 0 ]; do
+      if [ -f "$MOCK_DIR/response_$n.json" ]; then
+        cat "$MOCK_DIR/response_$n.json"
+        exit 0
+      fi
+      n=$(( n - 1 ))
+    done
+    echo "api outage" >&2
+    exit 1
+    ;;
+  "repos/o/r/pulls/42")
+    if [ -f "$MOCK_DIR/pr.json" ]; then
+      cat "$MOCK_DIR/pr.json"
+    else
+      echo '{"labels": [{"name": "ready"}]}'
+    fi
+    ;;
+  *)
+    echo "unexpected gh endpoint: $2" >> "$MOCK_DIR/endpoint_error"
+    exit 2
+    ;;
+esac
+EOF
+chmod +x "$tmp/gh"
+
+PC_OK='{"name": "pre-commit", "id": 1, "status": "completed", "conclusion": "success"}'
+PC_BAD='{"name": "pre-commit", "id": 1, "status": "completed", "conclusion": "failure"}'
+PC_PENDING='{"name": "pre-commit", "id": 1, "status": "in_progress", "conclusion": null}'
+DOCS_OK='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "success"}'
+DOCS_BAD='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "failure"}'
+DOCS_CANCELLED='{"name": "Deploy Documentation", "id": 2, "status": "completed", "conclusion": "cancelled"}'
+OTHER='{"name": "Trigger Full Suite", "id": 3, "status": "in_progress", "conclusion": null}'
+NULL_NAME='{"name": null, "id": 4, "status": "completed", "conclusion": "failure"}'
+PC_OK_RERUN='{"name": "pre-commit", "id": 5, "status": "completed", "conclusion": "success"}'
+
+fails=0
+want_log=""  # optional: expect() also greps out.log for this regex, then resets
+pr_json=""   # optional: served for the pulls (label re-check) endpoint, then resets
+raw_body=""  # optional: serve responses verbatim instead of wrapping in workflow_runs
+expect() { # <name> <expected-exit> <response json>...
+  local name=$1 want=$2 dir i=1
+  shift 2
+  dir=$(mktemp -d "$tmp/test_XXXXXX")
+  for body in "$@"; do
+    if [ -n "$raw_body" ]; then
+      printf '%s' "$body" > "$dir/response_$i.json"
+    else
+      printf '{"workflow_runs": [%s]}' "$body" > "$dir/response_$i.json"
+    fi
+    i=$(( i + 1 ))
+  done
+  [ -n "$pr_json" ] && printf '%s' "$pr_json" > "$dir/pr.json"
+  ( export PATH="$tmp:$PATH" MOCK_DIR="$dir" PR_SHA=deadbeef PR_NUMBER=42 \
+      GITHUB_REPOSITORY=o/r POLL_SECS=0 GRACE_SECS=1 MAX_WAIT_SECS=3
+    bash "$here/gate_full_suite.sh" > "$dir/out.log" 2>&1 )
+  local rc=$?
+  if [ "$rc" -ne "$want" ]; then
+    echo "FAIL: $name (exit $rc, want $want)"
+    cat "$dir/out.log"
+    fails=1
+  elif [ -f "$dir/endpoint_error" ]; then
+    echo "FAIL: $name (mock gh got an unexpected call)"
+    cat "$dir/endpoint_error"
+    fails=1
+  elif [ -n "$want_log" ] && ! grep -Eq "$want_log" "$dir/out.log"; then
+    echo "FAIL: $name (log does not match: $want_log)"
+    cat "$dir/out.log"
+    fails=1
+  else
+    echo "ok: $name"
+  fi
+  want_log="" pr_json="" raw_body=""
+}
+
+expect "both green -> proceed" 0 "$PC_OK, $DOCS_OK, $OTHER, $NULL_NAME"
+expect "docs build failed -> blocked" 1 "$PC_OK, $DOCS_BAD"
+expect "pre-commit failed -> blocked" 1 "$PC_BAD"
+expect "pending then green -> proceed" 0 "$PC_PENDING" "$PC_OK, $DOCS_OK"
+want_log="never appeared.*Deploy Documentation"
+expect "docs run absent (path-filtered) -> proceed after grace" 0 "$PC_OK"
+expect "API outage -> fail open" 0
+want_log="FAILING OPEN"
+expect "pending past MAX_WAIT -> fail open" 0 "$PC_PENDING"
+want_log="FAILING OPEN"
+expect "unrelated runs only -> no grace, fail open at MAX_WAIT" 0 "$OTHER"
+expect "cancelled docs then green -> proceed" 0 \
+  "$PC_OK, $DOCS_CANCELLED" "$PC_OK, $DOCS_OK"
+want_log="FAILING OPEN"
+expect "cancelled docs forever -> fail open at MAX_WAIT" 0 "$PC_OK, $DOCS_CANCELLED"
+want_log="FAILING OPEN"
+expect "pre-commit absent -> no grace, fail open at MAX_WAIT" 0 "$DOCS_OK"
+expect "duplicate run names -> latest wins" 0 "$PC_BAD, $PC_OK_RERUN, $DOCS_OK"
+raw_body=1
+expect "garbage response body -> fail open" 0 "this is not json"
+pr_json='{"labels": [{"name": "other"}]}'
+expect "ready label removed mid-gate -> blocked" 1 "$PC_OK, $DOCS_OK"
+
+exit "$fails"

.github/workflows/ci-precommit.yml

@@ -47,3 +47,6 @@ jobs:
     - uses: pre-commit/action@v3.0.1
       with:
         extra_args: --all-files --hook-stage manual
+    # After pre-commit so a self-test failure cannot mask lint failures.
+    - name: Full-suite gate self-test
+      run: bash .github/scripts/test_gate_full_suite.sh

.github/workflows/ci-trigger-full-suite.yml

@@ -7,6 +7,7 @@ on:
 permissions:
   contents: read
   pull-requests: read
+  actions: read
 
 concurrency:
   group: full-suite-${{ github.event.pull_request.number }}
@@ -18,6 +19,8 @@ jobs:
       (github.event.action == 'labeled' && github.event.label.name == 'ready')
       || github.event.action == 'synchronize'
     runs-on: ubuntu-latest
+    # Gate below may wait for cheap checks (up to MAX_WAIT_SECS = 25 min).
+    timeout-minutes: 35
     steps:
       - name: Check ready label
         id: check
@@ -49,6 +52,20 @@ jobs:
               "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds/${build_num}/cancel"
           done
 
+      # Checks out the BASE branch (default for pull_request_target), so PR
+      # authors cannot tamper with the gate script.
+      - name: Checkout gate script
+        if: steps.check.outputs.has_ready == 'true'
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+
+      - name: Wait for pre-commit and docs build
+        if: steps.check.outputs.has_ready == 'true'
+        env:
+          GH_TOKEN: ${{ github.token }}
+          PR_SHA: ${{ github.event.pull_request.head.sha }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+        run: bash .github/scripts/gate_full_suite.sh
+
       - name: Trigger Buildkite Full Suite
         if: steps.check.outputs.has_ready == 'true'
         env:

.github/workflows/publish-fastvideo-kernel-hub.yml

@@ -1,95 +0,0 @@
-name: Publish FastVideo Kernel to Hugging Face Hub
-
-# Version-gated like the PyPI flow (publish-kernel.yml): a push to main only
-# publishes when the fastvideo-kernel version actually changes. Use
-# workflow_dispatch as the manual override.
-on:
-  push:
-    branches:
-      - main
-    paths:
-      - "fastvideo-kernel/pyproject.toml"
-  workflow_dispatch:
-    inputs:
-      branch:
-        description: "Hub branch to upload the build"
-        default: ""
-        required: false
-
-# A Hub build takes hours; never race two publishes to the same repo.
-concurrency:
-  group: ${{ github.workflow }}
-  cancel-in-progress: false
-
-jobs:
-  check-version-change:
-    runs-on: ubuntu-latest
-    outputs:
-      version-changed: ${{ steps.check-version.outputs.changed }}
-      new-version: ${{ steps.check-version.outputs.new-version }}
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 2
-
-      - name: Check if version changed
-        id: check-version
-        run: |
-          cd fastvideo-kernel
-          # Get current commit's version from pyproject.toml
-          # Use ^ to match start of line to avoid matching minimum-version
-          NEW_VERSION=$(grep -oP '^version\s*=\s*"\K[^"]+' pyproject.toml)
-          echo "New version: $NEW_VERSION"
-
-          # Get previous version from git history
-          # Note: git show expects path relative to repo root
-          OLD_VERSION=$(git show HEAD~1:fastvideo-kernel/pyproject.toml | grep -oP '^version\s*=\s*"\K[^"]+' || echo "0.0.0")
-          echo "Old version: $OLD_VERSION"
-
-          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
-            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
-            echo "changed=true" >> "$GITHUB_OUTPUT"
-            echo "new-version=$NEW_VERSION" >> "$GITHUB_OUTPUT"
-          else
-            echo "Version did not change"
-            echo "changed=false" >> "$GITHUB_OUTPUT"
-          fi
-
-  build:
-    needs: check-version-change
-    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-
-      - name: Build kernel via HF Jobs and upload to hao-ai-lab/fastvideo-kernel
-        # Pinned to a commit SHA: this third-party action receives HF_TOKEN.
-        uses: huggingface/kernel-builder-job@ee590cee1ff01e922310401f18d5554b8da448ac # main, 2026-06-03
-        with:
-          token: ${{ secrets.HF_TOKEN }}
-          namespace: hao-ai-lab
-          flavor: cpu-xl
-          timeout: "21600"
-          script: |
-            set +x
-            export HF_TOKEN="${{ secrets.HF_TOKEN }}"
-            export GIT_LFS_SKIP_SMUDGE=1
-
-            git clone "${{ github.server_url }}/${{ github.repository }}" FastVideo
-            cd FastVideo
-            git checkout "${{ github.sha }}"
-            git submodule update --init --recursive fastvideo-kernel/include/cutlass fastvideo-kernel/include/tk
-            cd fastvideo-kernel
-
-            UPLOAD_BRANCH="${{ github.event.inputs.branch }}"
-            BRANCH_ARGS=()
-            if [ -n "${UPLOAD_BRANCH}" ]; then
-              BRANCH_ARGS=(--branch "${UPLOAD_BRANCH}")
-            fi
-
-            nix run github:huggingface/kernels#kernel-builder -- build-and-upload \
-              --max-jobs 4 \
-              --cores 8 \
-              --repo-id hao-ai-lab/fastvideo-kernel \
-              "${BRANCH_ARGS[@]}"

docs/contributing/ci_architecture.md

@@ -99,6 +99,13 @@ status.
 Full Suite is also path-filtered. It validates broader behavior before Mergify
 can merge a PR.
 
+A `ready`-labeled PR does not hit Buildkite immediately:
+`ci-trigger-full-suite.yml` first runs `.github/scripts/gate_full_suite.sh`,
+which waits for the cheap Tier-1 checks (pre-commit, docs build) on the PR
+head. A red cheap check blocks the suite (fail closed; the next push re-arms
+it), while a GitHub outage or a >25 min wait lets it run anyway (fail open).
+`/test full` bypasses the gate.
+
 | Buildkite label | `TEST_TYPE` | Main watched paths |
 |---|---|---|
 | SSIM Tests | `ssim` | `fastvideo/**/*.py`, `pyproject.toml`, `docker/Dockerfile` |

fastvideo-kernel/CARD.md

@@ -1,37 +0,0 @@
----
-license: apache-2.0
----
-
-# FastVideo Kernel
-
-Hub-compatible FastVideo CUDA kernels packaged with Hugging Face `kernel-builder`.
-
-## Usage
-
-```python
-from kernels import get_kernel
-
-fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1)
-```
-
-The module exposes:
-
-- `sta_fwd`
-- `block_sparse_fwd`
-- `block_sparse_bwd`
-- `rms_norm`
-- `layer_norm`
-- `int8_quant`
-- `int8_gemm`
-
-These APIs are thin wrappers around the native FastVideo kernels. Higher-level
-fallback logic remains in the regular `fastvideo-kernel` Python package.
-
-## Requirements
-
-This first Hub build targets CUDA Hopper (`sm_90a`) because the native attention
-kernels rely on Hopper/ThunderKittens features.
-
-## Source
-
-Source repository: https://github.com/hao-ai-lab/FastVideo/tree/main/fastvideo-kernel

fastvideo-kernel/README.md

@@ -40,25 +40,6 @@ Notes:
 
 ## Installation
 
-### Hugging Face Kernel Hub
-
-FastVideo kernels are built and published with the Hugging Face
-[`kernels`](https://huggingface.co/docs/kernels/index) project. Downstream code
-loads the precompiled kernels from
-[`hao-ai-lab/fastvideo-kernel`](https://huggingface.co/hao-ai-lab/fastvideo-kernel)
-without building locally:
-
-```python
-from kernels import get_kernel
-
-fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1)
-```
-
-The Hub package exposes the native CUDA entrypoints (`sta_fwd`,
-`block_sparse_fwd`, `block_sparse_bwd`, `rms_norm`, `layer_norm`, `int8_quant`,
-and `int8_gemm`). The full FastVideo Python package still provides the
-high-level fallback logic used during local development.
-
 ### Standard Installation (Local Development)
 This will automatically detect your GPU architecture. If an NVIDIA Hopper (H100/sm_90a) GPU is detected, ThunderKittens kernels will be enabled. Otherwise, they will be skipped, and the package will use Triton fallbacks at runtime.
 

fastvideo-kernel/build.toml

@@ -1,52 +0,0 @@
-[general]
-name = "fastvideo_kernel"
-version = 1
-license = "Apache-2.0"
-backends = ["cuda"]
-upstream = "https://github.com/hao-ai-lab/FastVideo.git"
-source = "https://github.com/hao-ai-lab/FastVideo.git"
-
-[general.hub]
-repo-id = "hao-ai-lab/fastvideo-kernel"
-
-[torch]
-src = [
-  "torch-ext/torch_binding.cpp",
-  "torch-ext/torch_binding.h",
-]
-
-[kernel.fastvideo_kernel]
-backend = "cuda"
-depends = ["torch"]
-include = [
-  "csrc",
-  "csrc/turbodiffusion",
-  "include/cutlass/include",
-  "include/tk/include",
-  "include/tk/prototype",
-]
-src = [
-  "csrc/attention/st_attn_h100.cu",
-  "csrc/attention/block_sparse_h100.cu",
-  "csrc/turbodiffusion/gemm/gemm.cu",
-  "csrc/turbodiffusion/norm/rmsnorm.cu",
-  "csrc/turbodiffusion/norm/layernorm.cu",
-  "csrc/turbodiffusion/quant/quant.cu",
-]
-# Deliberately Hopper-only (sm_90a): the TK attention kernels here require
-# Hopper, unlike the multi-arch PyPI wheels (see .github/workflows/publish-kernel.yml).
-cuda-capabilities = ["9.0a"]
-cuda-flags = [
-  "-DNDEBUG",
-  "-O3",
-  "-std=c++20",
-  "--use_fast_math",
-  "--expt-extended-lambda",
-  "--expt-relaxed-constexpr",
-  "-Xcompiler=-fno-strict-aliasing",
-  "-Xcompiler=-fPIC",
-  "-DTORCH_COMPILE",
-  "-DKITTENS_HOPPER",
-  "-DFASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION",
-  "-Xptxas=--warn-on-spills",
-]

fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu

@@ -71,12 +71,11 @@ void int8_gemm(
 
 }
 
-#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION
 void register_gemm(pybind11::module_ &m) {
     m.def("gemm_cuda", &int8_gemm);
 }
-#endif
 
 
   
 
+

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu

@@ -9,10 +9,10 @@
 #include "norm/layernorm.hpp"
 
 auto layer_norm(
-  at::Tensor const Input,
+  at::Tensor const Input, 
   float eps,
-  std::optional<at::Tensor> W,
-  std::optional<at::Tensor> const B,
+  std::optional<at::Tensor const> W,
+  std::optional<at::Tensor const> const B,
   std::optional<at::Tensor> Output
 ) {
   int64_t const m = Input.size(0);
@@ -93,8 +93,7 @@ auto layer_norm(
   return Output;
 }
 
-#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION
 void register_layer_norm(pybind11::module_ &m) {
     m.def("layer_norm_cuda", &layer_norm);
 }
-#endif
+

fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu

@@ -75,8 +75,6 @@ auto rms_norm(
   return Output;
 }
 
-#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION
 void register_rms_norm(pybind11::module_ &m) {
     m.def("rms_norm_cuda", &rms_norm);
 }
-#endif

fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu

@@ -70,8 +70,6 @@ auto quant(
   return std::make_tuple(Output, Output_S);
 }
 
-#ifndef FASTVIDEO_KERNEL_DISABLE_PYBIND_REGISTRATION
 void register_quant(pybind11::module_ &m) {
     m.def("quant_cuda", &quant);
 }
-#endif

fastvideo-kernel/example.py

@@ -1,11 +0,0 @@
-import torch
-from kernels import get_kernel
-
-
-fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1)
-
-x = torch.randn((16, 1024), device="cuda", dtype=torch.float16)
-weight = torch.ones((1024,), device="cuda", dtype=torch.float16)
-y = fastvideo_kernel.rms_norm(x, 1e-6, weight)
-
-print(y.shape)

fastvideo-kernel/flake.nix

@@ -1,18 +0,0 @@
-{
-  description = "Flake for FastVideo Kernel Hub builds";
-
-  inputs = {
-    self.submodules = true;
-    kernel-builder.url = "github:huggingface/kernels";
-  };
-
-  outputs =
-    {
-      self,
-      kernel-builder,
-    }:
-    kernel-builder.lib.genKernelFlakeOutputs {
-      inherit self;
-      path = ./.;
-    };
-}

fastvideo-kernel/torch-ext/fastvideo_kernel/__init__.py

@@ -1,122 +0,0 @@
-from typing import Optional
-
-import torch
-
-from ._ops import ops
-
-
-def sta_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    out: torch.Tensor,
-    kernel_t_size: int,
-    kernel_h_size: int,
-    kernel_w_size: int,
-    text_length: int,
-    process_text: bool,
-    has_text: bool,
-    kernel_aspect_ratio_flag: int,
-) -> torch.Tensor:
-    return ops.sta_fwd(
-        q,
-        k,
-        v,
-        out,
-        kernel_t_size,
-        kernel_h_size,
-        kernel_w_size,
-        text_length,
-        process_text,
-        has_text,
-        kernel_aspect_ratio_flag,
-    )
-
-
-def block_sparse_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    q2k_block_sparse_index: torch.Tensor,
-    q2k_block_sparse_num: torch.Tensor,
-    kv_block_size: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    out, lse = ops.block_sparse_fwd(
-        q,
-        k,
-        v,
-        q2k_block_sparse_index,
-        q2k_block_sparse_num,
-        kv_block_size,
-    )
-    return out, lse
-
-
-def block_sparse_bwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    out: torch.Tensor,
-    l_vec: torch.Tensor,
-    out_grad: torch.Tensor,
-    k2q_block_sparse_index: torch.Tensor,
-    k2q_block_sparse_num: torch.Tensor,
-    kv_block_size: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    q_grad, k_grad, v_grad = ops.block_sparse_bwd(
-        q,
-        k,
-        v,
-        out,
-        l_vec,
-        out_grad,
-        k2q_block_sparse_index,
-        k2q_block_sparse_num,
-        kv_block_size,
-    )
-    return q_grad, k_grad, v_grad
-
-
-def rms_norm(
-    input: torch.Tensor,
-    eps: float,
-    weight: Optional[torch.Tensor] = None,
-    output: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    return ops.rms_norm_cuda(input, eps, weight, output)
-
-
-def layer_norm(
-    input: torch.Tensor,
-    eps: float,
-    weight: Optional[torch.Tensor] = None,
-    bias: Optional[torch.Tensor] = None,
-    output: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    return ops.layer_norm_cuda(input, eps, weight, bias, output)
-
-
-def int8_quant(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    quantized, scale = ops.quant_cuda(input)
-    return quantized, scale
-
-
-def int8_gemm(
-    a: torch.Tensor,
-    a_scale: torch.Tensor,
-    b: torch.Tensor,
-    b_scale: torch.Tensor,
-    out: torch.Tensor,
-) -> torch.Tensor:
-    return ops.gemm_cuda(a, a_scale, b, b_scale, out)
-
-
-__all__ = [
-    "sta_fwd",
-    "block_sparse_fwd",
-    "block_sparse_bwd",
-    "rms_norm",
-    "layer_norm",
-    "int8_quant",
-    "int8_gemm",
-]

fastvideo-kernel/torch-ext/torch_binding.cpp

@@ -1,99 +0,0 @@
-#include <torch/library.h>
-
-#include "registration.h"
-#include "torch_binding.h"
-
-torch::Tensor rms_norm_cuda(
-    torch::Tensor const& input,
-    double eps,
-    std::optional<torch::Tensor> const& weight,
-    std::optional<torch::Tensor> const& output) {
-  std::optional<at::Tensor> mutable_output = output;
-  auto result = rms_norm(input, static_cast<float>(eps), weight, mutable_output);
-  TORCH_CHECK(result.has_value(), "rms_norm_cuda did not produce an output tensor");
-  return result.value();
-}
-
-torch::Tensor layer_norm_cuda(
-    torch::Tensor const& input,
-    double eps,
-    std::optional<torch::Tensor> const& weight,
-    std::optional<torch::Tensor> const& bias,
-    std::optional<torch::Tensor> const& output) {
-  auto result = layer_norm(input, static_cast<float>(eps), weight, bias, output);
-  TORCH_CHECK(result.has_value(), "layer_norm_cuda did not produce an output tensor");
-  return result.value();
-}
-
-std::tuple<torch::Tensor, torch::Tensor> quant_cuda(torch::Tensor const& input) {
-  std::optional<torch::Tensor> output;
-  std::optional<torch::Tensor> output_scale;
-  auto result = quant(input, output, output_scale);
-  TORCH_CHECK(std::get<0>(result).has_value(), "quant_cuda did not produce an int8 tensor");
-  TORCH_CHECK(std::get<1>(result).has_value(), "quant_cuda did not produce a scale tensor");
-  return std::make_tuple(std::get<0>(result).value(), std::get<1>(result).value());
-}
-
-torch::Tensor gemm_cuda(
-    torch::Tensor const& a,
-    torch::Tensor const& a_scale,
-    torch::Tensor const& b,
-    torch::Tensor const& b_scale,
-    torch::Tensor c) {
-  int8_gemm(a, a_scale, b, b_scale, c);
-  return c;
-}
-
-TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
-  ops.def(
-      // `out` is written in place and returned, so the return must be declared
-      // as the same alias set — otherwise torch.compile can reorder/eliminate
-      // the mutation and silently produce wrong results.
-      "sta_fwd(Tensor q, Tensor k, Tensor v, Tensor(a!) out, int kernel_t_size, int kernel_h_size, "
-      "int kernel_w_size, int text_length, bool process_text, bool has_text, "
-      "int kernel_aspect_ratio_flag) -> Tensor(a!)");
-#if defined(CUDA_KERNEL)
-  ops.impl("sta_fwd", torch::kCUDA, &sta_forward);
-#endif
-
-  ops.def(
-      "block_sparse_fwd(Tensor q, Tensor k, Tensor v, Tensor q2k_block_sparse_index, "
-      "Tensor q2k_block_sparse_num, Tensor kv_block_size) -> Tensor[]");
-#if defined(CUDA_KERNEL)
-  ops.impl("block_sparse_fwd", torch::kCUDA, &block_sparse_attention_forward);
-#endif
-
-  ops.def(
-      "block_sparse_bwd(Tensor q, Tensor k, Tensor v, Tensor out, Tensor l_vec, Tensor out_grad, "
-      "Tensor k2q_block_sparse_index, Tensor k2q_block_sparse_num, Tensor kv_block_size) -> Tensor[]");
-#if defined(CUDA_KERNEL)
-  ops.impl("block_sparse_bwd", torch::kCUDA, &block_sparse_attention_backward);
-#endif
-
-  // When `output` is provided it is written in place and aliased by the
-  // return value; the schema must say so for torch.compile correctness.
-  ops.def("rms_norm_cuda(Tensor input, float eps, Tensor? weight=None, Tensor(a!)? output=None) -> Tensor(a!)");
-#if defined(CUDA_KERNEL)
-  ops.impl("rms_norm_cuda", torch::kCUDA, &rms_norm_cuda);
-#endif
-
-  ops.def(
-      "layer_norm_cuda(Tensor input, float eps, Tensor? weight=None, Tensor? bias=None, "
-      "Tensor(a!)? output=None) -> Tensor(a!)");
-#if defined(CUDA_KERNEL)
-  ops.impl("layer_norm_cuda", torch::kCUDA, &layer_norm_cuda);
-#endif
-
-  ops.def("quant_cuda(Tensor input) -> (Tensor, Tensor)");
-#if defined(CUDA_KERNEL)
-  ops.impl("quant_cuda", torch::kCUDA, &quant_cuda);
-#endif
-
-  // `c` is the accumulator written in place and returned (same alias set).
-  ops.def("gemm_cuda(Tensor a, Tensor a_scale, Tensor b, Tensor b_scale, Tensor(a!) c) -> Tensor(a!)");
-#if defined(CUDA_KERNEL)
-  ops.impl("gemm_cuda", torch::kCUDA, &gemm_cuda);
-#endif
-}
-
-REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

fastvideo-kernel/torch-ext/torch_binding.h

@@ -1,86 +0,0 @@
-#pragma once
-
-#include <optional>
-#include <tuple>
-#include <vector>
-
-#include <torch/torch.h>
-
-torch::Tensor sta_forward(
-    torch::Tensor q,
-    torch::Tensor k,
-    torch::Tensor v,
-    torch::Tensor o,
-    int kernel_t_size,
-    int kernel_h_size,
-    int kernel_w_size,
-    int text_length,
-    bool process_text,
-    bool has_text,
-    int kernel_aspect_ratio_flag);
-
-std::vector<torch::Tensor> block_sparse_attention_forward(
-    torch::Tensor q,
-    torch::Tensor k,
-    torch::Tensor v,
-    torch::Tensor q2k_block_sparse_index,
-    torch::Tensor q2k_block_sparse_num,
-    torch::Tensor kv_block_size);
-
-std::vector<torch::Tensor> block_sparse_attention_backward(
-    torch::Tensor q,
-    torch::Tensor k,
-    torch::Tensor v,
-    torch::Tensor o,
-    torch::Tensor l_vec,
-    torch::Tensor og,
-    torch::Tensor k2q_block_sparse_index,
-    torch::Tensor k2q_block_sparse_num,
-    torch::Tensor kv_block_size);
-
-std::optional<at::Tensor> rms_norm(
-    at::Tensor const& input,
-    float eps,
-    std::optional<at::Tensor> const& weight,
-    std::optional<at::Tensor>& output);
-
-std::optional<at::Tensor> layer_norm(
-    at::Tensor const input,
-    float eps,
-    std::optional<at::Tensor> weight,
-    std::optional<at::Tensor> const bias,
-    std::optional<at::Tensor> output);
-
-std::tuple<std::optional<torch::Tensor>, std::optional<torch::Tensor>> quant(
-    torch::Tensor const& input,
-    std::optional<torch::Tensor>& output,
-    std::optional<torch::Tensor>& output_scale);
-
-void int8_gemm(
-    at::Tensor const& a,
-    at::Tensor const& a_scale,
-    at::Tensor const& b,
-    at::Tensor const& b_scale,
-    torch::Tensor& c);
-
-torch::Tensor rms_norm_cuda(
-    torch::Tensor const& input,
-    double eps,
-    std::optional<torch::Tensor> const& weight,
-    std::optional<torch::Tensor> const& output);
-
-torch::Tensor layer_norm_cuda(
-    torch::Tensor const& input,
-    double eps,
-    std::optional<torch::Tensor> const& weight,
-    std::optional<torch::Tensor> const& bias,
-    std::optional<torch::Tensor> const& output);
-
-std::tuple<torch::Tensor, torch::Tensor> quant_cuda(torch::Tensor const& input);
-
-torch::Tensor gemm_cuda(
-    torch::Tensor const& a,
-    torch::Tensor const& a_scale,
-    torch::Tensor const& b,
-    torch::Tensor const& b_scale,
-    torch::Tensor c);

fastvideo/layers/rotary_embedding.py

@@ -445,6 +445,21 @@ def get_nd_rotary_pos_embed(
     return cos, sin
 
 
+_ROTARY_POS_EMBED_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
+# Bound the table cache so long-running servers / causal models (which vary
+# start_frame per frame) cannot grow it without limit; entries are large float64
+# tensors. Least-recently-used eviction keeps the active resolution(s) hot while
+# capping memory.
+_ROTARY_POS_EMBED_CACHE_MAXSIZE = 16
+
+
+def _hashable(value: Any) -> Any:
+    """Return a hashable view of a scalar or sequence for use in a cache key."""
+    if isinstance(value, list | tuple):
+        return tuple(value)
+    return value
+
+
 def get_rotary_pos_embed(
     rope_sizes,
     hidden_size,
@@ -495,6 +510,31 @@ def get_rotary_pos_embed(
         sp_rank = 0
         sp_world_size = 1
 
+    # Memoize on every output-affecting argument; the table is constant across
+    # denoising steps, so this avoids recomputing the float64 cos/sin tables.
+    cache_key = (
+        _hashable(rope_sizes),
+        tuple(rope_dim_list),
+        rope_theta,
+        _hashable(theta_rescale_factor),
+        _hashable(interpolation_factor),
+        shard_dim,
+        sp_rank,
+        sp_world_size,
+        dtype,
+        start_frame,
+        use_real,
+    )
+    cached = _ROTARY_POS_EMBED_CACHE.get(cache_key)
+    if cached is not None:
+        # Move to most-recently-used position so the active table is not evicted
+        # when several resolutions / buckets share the process (LRU recency).
+        # Pop with a default: a concurrent eviction between the get() above and
+        # here would otherwise raise KeyError on the hit path.
+        if _ROTARY_POS_EMBED_CACHE.pop(cache_key, None) is not None:
+            _ROTARY_POS_EMBED_CACHE[cache_key] = cached
+        return cached
+
     freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
         rope_dim_list,
         rope_sizes,
@@ -508,6 +548,14 @@ def get_rotary_pos_embed(
         start_frame=start_frame,
         use_real=use_real,
     )
+    # The returned tensors are shared cache entries: callers must never mutate
+    # them in place. Note .to(device) is an identity alias when the tensor is
+    # already on the target device (e.g. CPU runs), so it does NOT guarantee a
+    # copy — treat the tables as read-only and copy before any in-place op.
+    # Reached only on a miss, so evict the least-recently-used entry at capacity.
+    if len(_ROTARY_POS_EMBED_CACHE) >= _ROTARY_POS_EMBED_CACHE_MAXSIZE:
+        _ROTARY_POS_EMBED_CACHE.pop(next(iter(_ROTARY_POS_EMBED_CACHE)))
+    _ROTARY_POS_EMBED_CACHE[cache_key] = (freqs_cos, freqs_sin)
     return freqs_cos, freqs_sin
 
 

fastvideo/tests/ops/test_rotary_pos_embed_cache.py

@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the get_rotary_pos_embed memoization cache."""
+
+import pytest
+import torch
+
+from fastvideo.layers.rotary_embedding import (
+    _ROTARY_POS_EMBED_CACHE,
+    _ROTARY_POS_EMBED_CACHE_MAXSIZE,
+    get_rotary_pos_embed,
+)
+
+
+def _rope_dim_list(hidden_size: int, heads_num: int) -> list[int]:
+    """Return the default 3-axis rope_dim_list used by the video DiTs."""
+    d = hidden_size // heads_num
+    return [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
+
+
+def _call(
+    rope_sizes=(21, 30, 52),
+    hidden_size=1536,
+    heads_num=12,
+    rope_dim_list="default",
+    rope_theta=10000.0,
+    dtype=torch.float64,
+    start_frame=0,
+    use_real=True,
+    **kwargs,
+):
+    """Thin wrapper around get_rotary_pos_embed with DiT-like defaults."""
+    if rope_dim_list == "default":
+        rope_dim_list = _rope_dim_list(hidden_size, heads_num)
+    return get_rotary_pos_embed(
+        rope_sizes,
+        hidden_size,
+        heads_num,
+        rope_dim_list,
+        rope_theta,
+        dtype=dtype,
+        start_frame=start_frame,
+        use_real=use_real,
+        **kwargs,
+    )
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    """Isolate every test by clearing the module-level cache around it."""
+    _ROTARY_POS_EMBED_CACHE.clear()
+    yield
+    _ROTARY_POS_EMBED_CACHE.clear()
+
+
+def test_repeated_call_hits_cache():
+    """A second identical call returns the exact same tensor objects."""
+    cos1, sin1 = _call()
+    cos2, sin2 = _call()
+    assert cos1 is cos2 and sin1 is sin2
+    assert len(_ROTARY_POS_EMBED_CACHE) == 1
+
+
+def test_many_identical_calls_keep_single_entry():
+    """Many identical calls never grow the cache beyond one entry."""
+    for _ in range(10):
+        _call()
+    assert len(_ROTARY_POS_EMBED_CACHE) == 1
+
+
+@pytest.mark.parametrize(
+    "rope_sizes,hidden_size,heads_num,dtype,use_real",
+    [
+        ((21, 30, 52), 1536, 12, torch.float64, True),
+        ((21, 45, 80), 5120, 40, torch.float64, True),
+        ((1, 16, 16), 1536, 12, torch.float32, True),
+        ((4, 8, 8), 1536, 12, torch.float64, False),
+    ],
+)
+def test_cached_matches_fresh_recompute(rope_sizes, hidden_size, heads_num,
+                                        dtype, use_real):
+    """Cached tensors are bitwise-equal to a fresh uncached recompute."""
+    cos_cached, sin_cached = _call(rope_sizes=rope_sizes,
+                                   hidden_size=hidden_size,
+                                   heads_num=heads_num,
+                                   dtype=dtype,
+                                   use_real=use_real)
+    _ROTARY_POS_EMBED_CACHE.clear()
+    cos_fresh, sin_fresh = _call(rope_sizes=rope_sizes,
+                                 hidden_size=hidden_size,
+                                 heads_num=heads_num,
+                                 dtype=dtype,
+                                 use_real=use_real)
+    assert torch.equal(cos_cached, cos_fresh)
+    assert torch.equal(sin_cached, sin_fresh)
+
+
+@pytest.mark.parametrize(
+    "kwargs_a,kwargs_b",
+    [
+        ({"rope_sizes": (21, 30, 52)}, {"rope_sizes": (21, 45, 80)}),
+        ({"dtype": torch.float64}, {"dtype": torch.float32}),
+        ({"use_real": True}, {"use_real": False}),
+        ({"start_frame": 0}, {"start_frame": 3}),
+        ({"rope_theta": 10000.0}, {"rope_theta": 5000.0}),
+        ({"shard_dim": 0}, {"shard_dim": 1}),
+    ],
+)
+def test_distinct_args_create_distinct_entries(kwargs_a, kwargs_b):
+    """Any output-affecting argument difference yields a separate cache entry."""
+    _call(**kwargs_a)
+    _call(**kwargs_b)
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_none_rope_dim_list_shares_key_with_equivalent_list():
+    """None rope_dim_list and its derived explicit list map to one entry."""
+    # head_dim must be divisible by 3 for the None branch to stay valid.
+    hidden_size, heads_num = 1536, 16  # head_dim == 96 -> [32, 32, 32]
+    _call(rope_dim_list=None, hidden_size=hidden_size, heads_num=heads_num)
+    before = len(_ROTARY_POS_EMBED_CACHE)
+    _call(rope_dim_list=[32, 32, 32], hidden_size=hidden_size, heads_num=heads_num)
+    assert len(_ROTARY_POS_EMBED_CACHE) == before == 1
+
+
+def test_use_real_controls_last_dim():
+    """use_real=True spans full head_dim; use_real=False spans half."""
+    cos_full, _ = _call(use_real=True)
+    cos_half, _ = _call(use_real=False)
+    assert cos_full.shape[-1] == 128
+    assert cos_half.shape[-1] == 64
+
+
+@pytest.mark.parametrize("rope_sizes", [(1, 1, 1), (1, 30, 52), (21, 1, 1)])
+def test_degenerate_grid_shapes(rope_sizes):
+    """Degenerate single-element axes still produce a correctly sized table."""
+    cos, sin = _call(rope_sizes=rope_sizes)
+    expected = rope_sizes[0] * rope_sizes[1] * rope_sizes[2]
+    assert cos.shape[0] == expected
+    assert sin.shape[0] == expected
+
+
+def test_scalar_and_list_factors_are_hashable_and_distinct():
+    """List-valued rescale factors are hashable and keyed apart from scalars."""
+    _call(theta_rescale_factor=1.0)
+    _call(theta_rescale_factor=[1.0, 1.0, 1.0])
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_caller_device_copy_does_not_corrupt_cache():
+    """The .to()/.float() copy callers perform must not mutate cached tensors."""
+    cos, _ = _call()
+    snapshot = cos.clone()
+    _ = cos.to("cpu").float()
+    cos_again, _ = _call()
+    assert torch.equal(cos_again, snapshot)
+
+
+def test_start_frame_offsets_values():
+    """A non-zero start_frame shifts the temporal positions, changing output."""
+    cos0, _ = _call(start_frame=0)
+    cos3, _ = _call(start_frame=3)
+    assert not torch.equal(cos0, cos3)
+    assert len(_ROTARY_POS_EMBED_CACHE) == 2
+
+
+def test_cache_is_bounded_and_evicts_oldest():
+    """The cache caps at the max size and evicts the oldest entry first."""
+    # Tiny grids keep this lightweight; each start_frame is a distinct key.
+    overshoot = _ROTARY_POS_EMBED_CACHE_MAXSIZE + 4
+    for frame in range(overshoot):
+        _call(rope_sizes=(2, 2, 2), start_frame=frame)
+        assert len(_ROTARY_POS_EMBED_CACHE) <= _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    # The earliest-inserted frames must have been evicted; the latest survive.
+    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}  # start_frame slot
+    assert overshoot - 1 in surviving
+    assert 0 not in surviving
+
+
+def test_cache_hit_refreshes_recency():
+    """Re-accessing an entry protects it from eviction over an untouched one."""
+    _call(rope_sizes=(2, 2, 2), start_frame=0)  # entry we will keep hot
+    for frame in range(1, _ROTARY_POS_EMBED_CACHE_MAXSIZE):
+        _call(rope_sizes=(2, 2, 2), start_frame=frame)
+    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
+    _call(rope_sizes=(2, 2, 2), start_frame=0)   # hit -> frame 0 becomes most recent
+    _call(rope_sizes=(2, 2, 2), start_frame=99)  # miss -> evicts now-oldest (frame 1)
+    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}
+    assert 0 in surviving
+    assert 1 not in surviving