[perf] Add fine-grained FA4 VSA kernels

2026-08-17 17:39:20 +00:00
227 changed files with 52722 additions and 23454 deletions
@@ -1,149 +0,0 @@
-name: macOS MLX Smoke
-
-on:
-  pull_request:
-    branches: [main]
-    paths:
-      - ".github/workflows/ci-macos-mlx.yml"
-      - "fastvideo/mlx_runtime/**"
-      - "fastvideo/tests/mlx/**"
-      - "fastvideo/tests/platforms/test_mps_vsa_error.py"
-      - "fastvideo/platforms/mps.py"
-      - "fastvideo/platforms/__init__.py"
-      - "fastvideo/__init__.py"
-      - "examples/inference/basic/mlx_*.py"
-      - "fastvideo/benchmarks/mlx_*.py"
-      - "pyproject.toml"
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-concurrency:
-  group: macos-mlx-${{ github.ref }}
-  cancel-in-progress: true
-
-jobs:
-  mlx-smoke:
-    if: github.event_name == 'workflow_dispatch' || github.event.pull_request.draft != true
-    runs-on: macos-15
-    timeout-minutes: 25
-    env:
-      FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA
-      TOKENIZERS_PARALLELISM: "false"
-      MASTER_ADDR: localhost
-      MASTER_PORT: "29513"
-    steps:
-      - uses: actions/checkout@v4
-
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
-          cache: pip
-
-      - uses: astral-sh/setup-uv@v3
-
-      - name: Install lightweight MLX smoke dependencies
-        run: |
-          uv pip install --system \
-            --index-url https://download.pytorch.org/whl/cpu \
-            torch==2.11.0 torchvision torchaudio
-          uv pip install --system \
-            pytest numpy scipy pillow imageio einops cloudpickle filelock \
-            PyYAML diffusers huggingface_hub remote-pdb safetensors loguru mlx \
-            "ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0"
-
-      - name: Show Apple runtime
-        run: |
-          python - <<'PY'
-          import platform
-          import mlx.core as mx
-          import torch
-
-          print("machine:", platform.machine())
-          print("processor:", platform.processor())
-          print("mlx default device:", mx.default_device())
-          memory_size = mx.metal.device_info().get("memory_size") if mx.metal.is_available() else "metal unavailable"
-          print("mlx memory_size:", memory_size)
-          print("torch:", torch.__version__)
-          print("torch mps available:", torch.backends.mps.is_available())
-          PY
-
-      - name: Run MLX smoke tests
-        run: |
-          python -m pytest \
-            fastvideo/tests/mlx/test_dmd_sampling.py \
-            fastvideo/tests/mlx/test_memory_limits.py \
-            fastvideo/tests/mlx/test_quant_capability.py \
-            fastvideo/tests/mlx/test_mlx_dit_parity.py \
-            fastvideo/tests/mlx/test_mlx_compile_parity.py \
-            fastvideo/tests/mlx/test_mlx_checkpoint.py \
-            fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \
-            fastvideo/tests/mlx/test_taehv_decode.py \
-            fastvideo/tests/mlx/test_frame_upsample.py \
-            fastvideo/tests/mlx/test_mlx_fast_spatial.py \
-            fastvideo/tests/mlx/test_mlx_refine.py \
-            fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \
-            fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \
-            fastvideo/tests/mlx/test_wan22_sample.py \
-            fastvideo/tests/mlx/test_windowed_attention.py \
-            fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \
-            fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \
-            fastvideo/tests/platforms/test_mps_vsa_error.py \
-            -q
-
-  # Same tests on MLX's CPU backend. Hosted macOS runners are scarce and
-  # slower to schedule; this Linux job gives fast PR signal on the identical
-  # graph (the parity tests were designed to be backend-agnostic), while the
-  # macOS job above stays the source of truth for Metal behavior.
-  mlx-smoke-linux-cpu:
-    if: github.event_name == 'workflow_dispatch' || github.event.pull_request.draft != true
-    runs-on: ubuntu-latest
-    timeout-minutes: 20
-    env:
-      FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA
-      TOKENIZERS_PARALLELISM: "false"
-      MASTER_ADDR: localhost
-      MASTER_PORT: "29513"
-    steps:
-      - uses: actions/checkout@v4
-
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
-          cache: pip
-
-      - uses: astral-sh/setup-uv@v3
-
-      - name: Install lightweight MLX smoke dependencies (CPU backend)
-        run: |
-          uv pip install --system \
-            --index-url https://download.pytorch.org/whl/cpu \
-            torch==2.11.0 torchvision torchaudio
-          uv pip install --system \
-            pytest numpy scipy pillow imageio einops cloudpickle filelock \
-            PyYAML diffusers huggingface_hub remote-pdb safetensors loguru "mlx[cpu]" \
-            "ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0"
-
-      - name: Run MLX smoke tests (CPU backend)
-        run: |
-          python -m pytest \
-            fastvideo/tests/mlx/test_dmd_sampling.py \
-            fastvideo/tests/mlx/test_memory_limits.py \
-            fastvideo/tests/mlx/test_quant_capability.py \
-            fastvideo/tests/mlx/test_mlx_dit_parity.py \
-            fastvideo/tests/mlx/test_mlx_compile_parity.py \
-            fastvideo/tests/mlx/test_mlx_checkpoint.py \
-            fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \
-            fastvideo/tests/mlx/test_taehv_decode.py \
-            fastvideo/tests/mlx/test_frame_upsample.py \
-            fastvideo/tests/mlx/test_mlx_fast_spatial.py \
-            fastvideo/tests/mlx/test_mlx_refine.py \
-            fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \
-            fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \
-            fastvideo/tests/mlx/test_wan22_sample.py \
-            fastvideo/tests/mlx/test_windowed_attention.py \
-            fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \
-            fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \
-            fastvideo/tests/platforms/test_mps_vsa_error.py \
-            -q
@@ -6,7 +6,6 @@ on:
    paths:
      - 'docs/**'
      - 'examples/**'
-      - 'scripts/inference/**'
      - 'mkdocs.yml'
      - 'requirements-mkdocs.in'
      - 'requirements-mkdocs.txt'
@@ -17,7 +16,6 @@ on:
    paths:
      - 'docs/**'
      - 'examples/**'
-      - 'scripts/inference/**'
      - 'mkdocs.yml'
      - 'requirements-mkdocs.in'
      - 'requirements-mkdocs.txt'
@@ -9,7 +9,6 @@
 **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.**

 ## NEWS
- `2026/08/19`: FastVideo now supports MLX on Apple Silicon with [FastMetal-QAD](https://huggingface.co/collections/FastVideo/fastmetal), a family of 1.3B, 5B, and 14B models optimized for Mac—follow the [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/) and read the [Blog](https://haoailab.com/blogs/fastmetal/).
 - `2026/06/23`: Release FastWan-QAD: 5s of Video generated in 1.8s E2E. See the [FastWan-QAD models](https://huggingface.co/FastVideo/FastWan-QAD-FP8-1.3B), [Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/), and [blog](https://haoailab.com/blogs/fastwan-qad/).
 - `2026/03/17`: Release demo: Into the Dreamverse: Vibe Directing in FastVideo, check out the [Blog](https://haoailab.com/blogs/dreamverse/).
 - `2026/03/13`: Release demo: Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU, check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/).
@@ -63,11 +62,6 @@ UV_TORCH_BACKEND=cu126 uv pip install fastvideo
 Use `UV_TORCH_BACKEND=cu130` on CUDA 13. Apple silicon users should follow the
 [MPS installation guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).

-> **On an Apple Silicon Mac?** FastVideo runs FastWan text-to-video natively
-> through an MLX runtime — a 5-second 480p clip generated locally, no cloud,
-> no discrete GPU. Install with `uv pip install -e '.[mlx]'` and follow the
-> [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
-
 Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.

 > **On an NVIDIA DGX Spark (GB10 / ARM64 + CUDA 13)?** There's no prebuilt ARM wheel for the FastVideo CUDA kernel, so it's an editable from-source install (`UV_TORCH_BACKEND=cu130 uv pip install -e .`, which compiles that kernel for you) rather than `UV_TORCH_BACKEND=cu130 uv pip install fastvideo`. A compatible prebuilt ARM64 FlashAttention wheel is available separately. Follow the [DGX Spark install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/spark/).
@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=14c377950125c70b7a9dabf9c561fca53715ac7d
+ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083

 # Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
 # select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
@@ -1,52 +0,0 @@
-{
-  "recipes": [
-    {
-      "id": "fastwan21-t2v",
-      "task": "Text to video",
-      "label": "FastWan2.1 1.3B (distilled + VSA)",
-      "model": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
-      "source": "scripts/inference/inference_wan_VSA_DMD_1_3B.yaml",
-      "command": "FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN fastvideo generate --config scripts/inference/inference_wan_VSA_DMD_1_3B.yaml"
-    },
-    {
-      "id": "wan22-t2v",
-      "task": "Text to video",
-      "label": "Wan2.2 A14B",
-      "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-      "source": "examples/inference/basic/basic_wan2_2.py",
-      "command": "python examples/inference/basic/basic_wan2_2.py"
-    },
-    {
-      "id": "wan21-i2v",
-      "task": "Image to video",
-      "label": "Wan2.1 14B 480P",
-      "model": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
-      "source": "scripts/inference/inference_wan_i2v.yaml",
-      "command": "fastvideo generate --config scripts/inference/inference_wan_i2v.yaml"
-    },
-    {
-      "id": "turbowan22-i2v",
-      "task": "Image to video",
-      "label": "TurboWan2.2 A14B",
-      "model": "loayrashid/TurboWan2.2-I2V-A14B-Diffusers",
-      "source": "examples/inference/basic/basic_turbodiffusion_i2v.py",
-      "command": "python examples/inference/basic/basic_turbodiffusion_i2v.py"
-    },
-    {
-      "id": "wan22-ti2v",
-      "task": "Text or image to video",
-      "label": "Wan2.2 TI2V 5B",
-      "model": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
-      "source": "examples/inference/basic/basic_wan2_2_ti2v.py",
-      "command": "python examples/inference/basic/basic_wan2_2_ti2v.py"
-    },
-    {
-      "id": "matrix-game-2",
-      "task": "Interactive world",
-      "label": "Matrix Game 2.0",
-      "model": "FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers",
-      "source": "examples/inference/basic/basic_matrixgame2.py",
-      "command": "python examples/inference/basic/basic_matrixgame2.py"
-    }
-  ]
-}
@@ -1,60 +0,0 @@
-(() => {
-  let recipesPromise;
-
-  const loadRecipes = (url) => {
-    recipesPromise ||= fetch(url).then((response) => {
-      if (!response.ok) throw new Error(`HTTP ${response.status}`);
-      return response.json();
-    });
-    return recipesPromise;
-  };
-
-  const init = () => {
-    document.querySelectorAll("[data-cookbook]").forEach(async (root) => {
-      if (root.dataset.initialized) return;
-      root.dataset.initialized = "true";
-
-      const select = root.querySelector("[data-cookbook-recipe]");
-      const model = root.querySelector("[data-cookbook-model]");
-      const source = root.querySelector("[data-cookbook-source]");
-      const command = root.querySelector("[data-cookbook-command]");
-      const status = root.querySelector("[data-cookbook-status]");
-
-      try {
-        const { recipes } = await loadRecipes(root.dataset.recipes);
-        const byId = new Map(recipes.map((recipe) => [recipe.id, recipe]));
-        const groups = new Map();
-
-        select.replaceChildren();
-        recipes.forEach((recipe) => {
-          if (!groups.has(recipe.task)) {
-            const group = document.createElement("optgroup");
-            group.label = recipe.task;
-            groups.set(recipe.task, group);
-            select.append(group);
-          }
-          groups.get(recipe.task).append(new Option(recipe.label, recipe.id));
-        });
-
-        const render = () => {
-          const recipe = byId.get(select.value);
-          model.textContent = recipe.model;
-          source.textContent = recipe.source;
-          source.href = `https://github.com/hao-ai-lab/FastVideo/blob/main/${recipe.source}`;
-          command.textContent = recipe.command;
-          status.textContent = `${recipe.label} selected.`;
-        };
-
-        select.addEventListener("change", render);
-        select.disabled = false;
-        render();
-      } catch (error) {
-        status.textContent = "Recipes could not be loaded. Use the examples link below.";
-        console.error("Failed to load FastVideo cookbook recipes", error);
-      }
-    });
-  };
-
-  if (window.document$) window.document$.subscribe(init);
-  else document.addEventListener("DOMContentLoaded", init);
-})();
@@ -42,46 +42,6 @@ img {
    margin: 0 auto;
 }

-.cookbook-picker {
-    padding: 1rem;
-    border: 0.05rem solid var(--md-default-fg-color--lightest);
-    border-radius: 0.2rem;
-}
-
-.cookbook-picker select {
-    width: 100%;
-    padding: 0.6rem;
-    color: var(--md-default-fg-color);
-    background: var(--md-default-bg-color);
-    border: 0.05rem solid var(--md-default-fg-color--lighter);
-    border-radius: 0.2rem;
-}
-
-.cookbook-picker dl {
-    display: grid;
-    grid-template-columns: max-content 1fr;
-    gap: 0.25rem 1rem;
-}
-
-.cookbook-picker dt {
-    font-weight: 700;
-}
-
-.cookbook-picker dd {
-    margin: 0;
-    min-width: 0;
-    overflow-wrap: anywhere;
-}
-
-.cookbook-picker__status {
-    position: absolute;
-    width: 1px;
-    height: 1px;
-    overflow: hidden;
-    clip: rect(0, 0, 0, 0);
-    white-space: nowrap;
-}
-
 .md-typeset .copy-page-button.md-button {
    float: right;
    margin: 0 0 1rem 1rem;
@@ -1,44 +0,0 @@
-# Inference Cookbook
-
-Choose a complete recipe maintained in the FastVideo repository. Each command
-runs its checked-in source directly, so coupled model, GPU, offload, and
-attention settings do not drift into unsupported combinations.
-
-The commands expect a local clone:
-
-```bash
-git clone https://github.com/hao-ai-lab/FastVideo.git
-cd FastVideo
-```
-
-<div class="cookbook-picker" data-cookbook data-recipes="../assets/cookbook-recipes.json">
-  <label for="cookbook-recipe"><strong>Recipe</strong></label>
-  <select id="cookbook-recipe" data-cookbook-recipe disabled>
-    <option>Loading recipes…</option>
-  </select>
-  <dl>
-    <dt>Model</dt>
-    <dd data-cookbook-model>Loading…</dd>
-    <dt>Source</dt>
-    <dd><a data-cookbook-source href="../inference/examples/basic/">Browse maintained examples</a></dd>
-  </dl>
-  <pre><code class="language-bash" data-cookbook-command>Loading…</code></pre>
-  <p class="cookbook-picker__status" role="status" aria-live="polite" data-cookbook-status></p>
-  <noscript>
-    JavaScript is needed for the recipe picker. Browse the
-    <a href="../inference/examples/examples_inference_index/">inference examples</a>
-    instead.
-  </noscript>
-</div>
-
-## Customize a recipe
-
-Start from the checked-in source, then change only the settings your model
-supports:
-
- [Configuration](../inference/configuration.md) covers the Python and CLI
-  config surfaces.
- [Optimizations](../inference/optimizations.md) covers attention backends,
-  compilation, and memory tradeoffs.
- [Support matrix](../inference/support_matrix.md) lists supported models and
-  optimizations.
@@ -1,128 +0,0 @@
-# Fast mode (RIFE) — Apple Silicon
-
-`--fast` makes local generation ~2.7× faster by **generating fewer frames and
-interpolating the rest** with an Apple-Silicon-native RIFE model, instead of
-denoising every frame. Video-diffusion denoise is dominated by self-attention,
-which is O(tokens²); halving the frames cuts the token count ~2× and the denoise
-compute ~3.7×, so the wall-clock drops far more than 2×. RIFE (which estimates
-its own optical flow — no motion vectors needed) fills the dropped frames back
-in for ~1.4 s, and a light unsharp pass counters its softening.
-
-Measured on the 1.3B INT8 QAD model (fox, 480×832×81, M4): generate 41 + RIFE→81
-runs in ~35 s of denoise vs ~90 s full, at reconstruction MS-SSIM **0.97**.
-Reproduce with `python -m fastvideo.benchmarks.eval_metalfx_rife --mode int8`.
-
-> **Note:** Apple's *MetalFX* frame interpolation is **not** usable here — it
-> requires game-engine motion vectors + depth, which diffusion output lacks. We
-> use the video-native **`rife-mlx`** model instead (Metal-backed, torch-free).
-
-## Install
-
-```bash
-uv pip install -e ".[mlx]"   # RIFE ships vendored; this only needs MLX
-```
-
-## Use
-
-```bash
-python examples/inference/basic/mlx_wan_prompt_to_video.py \
-  --mlx-checkpoint <FastWan2.1-T2V-1.3B-INT8-QAD> \
-  --prompt "A red fox trotting through a snowy pine forest at golden hour, cinematic" \
-  --num-frames 81 --fast \
-  --output-path video_samples/fox_fast.mp4
-```
-
-`--num-frames` stays the *target* length; fast mode generates the smallest
-VAE-aligned keyframe count that RIFE can interpolate to that target.
-
-| Flag | Default | Meaning |
-|---|---|---|
-| `--fast` / `--no-fast` | off | enable fast mode |
-| `--fast-factor` | 2 | generate 1/factor of the frames (2 = half) |
-| `--fast-sharpen` | 0.6 | light unsharp strength to counter RIFE softness (0 disables) |
-
-Fast mode composes with everything else (`--mlx-quantization int8`,
-`--mlx-compile`, TAEHV vs `--decode-backend wan-vae`). Keep `--fast-factor` at 2
-for quality — larger temporal gaps are where RIFE starts inventing motion.
-
-## Spatial fast mode (`--fast-spatial`)
-
-The spatial twin of `--fast`: instead of dropping frames, drop pixels. Denoise
-*and decode* at `height/width // fast-spatial-scale`, then resample the decoded
-frames up to the requested size. Self-attention is O(tokens²), so halving each
-spatial axis cuts the token count 4× and the denoise time far more than that —
-measured on the 1.3B INT8 QAD model at 480×832×81, M4 Max: **86.1 s → 10.3 s**
-of denoise. It composes with `--fast`; both together run the same clip in
-**4.5 s** of denoise.
-
-```bash
-python examples/inference/basic/mlx_wan_prompt_to_video.py \
-  --prompt "A red fox trotting through a snowy pine forest at golden hour, cinematic" \
-  --height 480 --width 832 --num-frames 81 --fast-spatial \
-  --output-path video_samples/fox_fast_spatial.mp4
-```
-
-| Flag | Default | Meaning |
-|---|---|---|
-| `--fast-spatial` / `--no-fast-spatial` | off | enable spatial fast mode |
-| `--fast-spatial-scale` | 2 | denoise at 1/scale of each spatial axis |
-| `--fast-spatial-upsample-mode` | `lanczos` | pixel interpolation kernel (`lanczos`, `cubic`, `bilinear`, `nearest`) |
-| `--fast-spatial-sharpen` | 0.4 | light unsharp strength to counter resampling softness (0 disables) |
-
-### The upsample must happen in pixel space
-
-This is the one thing to get right. The obvious implementation — bilinearly
-upsample the finished latents and decode at the target size — **does not work**,
-and produces a distinctive failure: correct composition and silhouette under a
-smeared, hazy veil, with ringing along strong edges.
-
-A Wan latent cell is a *learned code* for an 8×8 (Wan2.1) or 16×16 (Wan2.2)
-pixel block, not a low-pass sample of the image. The average of two adjacent
-codes is not the code of the averaged blocks; it is a vector the decoder was
-never trained on. Measured on Wan2.1-1.3B at 480×832, a 2× bilinear latent
-upsample destroys **62%** of the latent's high-frequency energy while leaving
-its overall magnitude intact — exactly the signature of that veil. At Wan2.2-5B
-the same operation degrades to black or noise.
-
-Decoded RGB frames have no such problem: an image *is* a sampled 2-D signal, so
-Lanczos interpolation is the operation it was defined for. The result is soft —
-it carries stage-1's real detail budget and no more — but clean and coherent.
-
-`--refine` gets away with a latent-space upsample only because a second DMD pass
-re-denoises the hand-off; spatial fast mode passes the latent straight to the
-decoder, so it cannot.
-
-## Refine (`--refine`) stage-2 timesteps
-
-`--refine` hands stage 1 to stage 2 as `(1 - sigma) * upsampled + sigma * noise`,
-where `sigma` comes from the *first* stage-2 timestep. FastWan's DMD grid opens
-at `t=1000`, which is `sigma == 1` exactly — so a stage-2 grid that starts there
-weights the stage-1 result at zero and refine silently degrades into a plain
-full-resolution run at twice the cost.
-
-Left unset, `--refine-dmd-denoising-steps` now derives the stage-2 grid from the
-stage-1 one with leading full-noise steps dropped (`1000,757,522` → `757,522`).
-That keeps the pass on timesteps the distilled student was trained on while
-letting stage-1 structure through: hand-off `sigma = 0.757`, stage-1 weight
-`0.243`. Passing a grid that starts at full noise is now an error rather than a
-silently wasted pass.
-
-The run prints the resolved hand-off so it is visible:
-
-```
-[refine] stage-2 hand-off sigma=0.7568 (stage-1 weight 0.2432)
-```
-
-There is a trade-off in choosing that grid. Later start = more of the draft
-survives, but fewer stage-2 steps. On Wan2.1 the default `757,522` gives weight
-0.243 with two steps; `--refine-dmd-denoising-steps 522` gives weight 0.478 with
-one. `--refine-sigma` decouples the noise level from the timestep entirely — it
-logs a warning, because the DiT is then told a timestep that does not match the
-noise it receives.
-
-**Wan2.2-5B has a lower ceiling.** Its warped schedule maps `1000,757,522` to
-sigmas `1.000, 0.940, 0.845`, so the best available stage-1 weight is **0.060**
-(vs 0.243 at 1.3B). Un-warped (`--no-warp`) the same grid gives `1.000, 0.757,
-0.522` and a weight of 0.243 — but warping is what matches the FastVideo
-sampling schedule, so turning it off changes the timesteps the distilled student
-sees. Which is better at 5B is unresolved and needs a run on real 5B weights.
@@ -76,10 +76,6 @@ surfaces:
      lora_target_modules: "Legacy LoRA configuration surface pending dedicated component API."
      output_type: "Legacy output formatting surface pending GenerationResult cleanup."
      VSA_sparsity: "Model-specific inference optimization not yet represented in the typed public schema."
-      VSA_tile_size: "VSA-H3 tile geometry request; model-specific optimization not yet represented in the typed public schema."
-      vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema."
-      vae_parallel_encode: "MiniMax-H3 sequence-parallel reference VAE encode opt-in; model-specific optimization not yet represented in the typed public schema."
-      vae_parallel_decode_strategy: "Chunk-transport collective for vae_parallel_decode; model-specific optimization not yet represented in the typed public schema."
      attention_backend: "Process-wide default attention-backend request applied per component at load time; kernel-selection knob not yet represented in the typed public schema."
      moba_config_path: "Model-specific MoBA optimization surface not yet represented in the typed public schema."
      master_port: "Executor/bootstrap compatibility field; not part of the canonical inference schema."
@@ -2,7 +2,6 @@
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/docs/source/generate_examples.py

 import itertools
-import json
 import os
 import re
 from dataclasses import dataclass, field
@@ -20,40 +19,6 @@ GENERATED_DOC_PREFIXES = (
    "training/examples/",
    "distillation/examples/",
 )
-COOKBOOK_DATA = ROOT_DIR / "docs/assets/cookbook-recipes.json"
-COOKBOOK_SOURCE_ROOTS = (
-    ROOT_DIR / "examples/inference",
-    ROOT_DIR / "scripts/inference",
-)
-
-
-def validate_cookbook() -> None:
-    """Keep cookbook entries tied to checked-in runnable sources."""
-    recipes = json.loads(COOKBOOK_DATA.read_text(encoding="utf-8")).get("recipes")
-    if not isinstance(recipes, list) or not recipes:
-        raise ValueError(f"{COOKBOOK_DATA}: recipes must be a non-empty list")
-
-    seen: set[str] = set()
-    for recipe in recipes:
-        required = ("id", "task", "label", "model", "source", "command")
-        missing = {key for key in required if not recipe.get(key)}
-        if missing:
-            raise ValueError(f"Cookbook recipe is missing: {', '.join(sorted(missing))}")
-        if recipe["id"] in seen:
-            raise ValueError(f"Duplicate cookbook recipe id: {recipe['id']}")
-        seen.add(recipe["id"])
-
-        source = (ROOT_DIR / recipe["source"]).resolve()
-        if not any(source.is_relative_to(root.resolve()) for root in COOKBOOK_SOURCE_ROOTS):
-            raise ValueError(f"Cookbook source is outside an approved directory: {recipe['source']}")
-        if not source.is_file():
-            raise ValueError(f"Cookbook source does not exist: {recipe['source']}")
-
-        source_text = source.read_text(encoding="utf-8")
-        if recipe["model"] not in source_text:
-            raise ValueError(f"Cookbook model is not present in {recipe['source']}: {recipe['model']}")
-        if recipe["source"] not in recipe["command"]:
-            raise ValueError(f"Cookbook command does not invoke its source: {recipe['id']}")


 def fix_case(text: str) -> str:
@@ -571,7 +536,6 @@ def on_pre_build(config, **kwargs):
    MkDocs hook to generate examples before building the documentation.
    This function is called automatically by MkDocs' native hook system.
    """
-    validate_cookbook()
    print("Generating example documentation...")
    generate_examples(generate_main_index=True)
    print("Example documentation generated successfully!")
@@ -585,7 +549,6 @@ def on_page_context(context, page, **kwargs):


 if __name__ == "__main__":
-    validate_cookbook()
    print("Generating example documentation...")
    generate_examples(generate_main_index=True)
    print("Example documentation generated successfully!")
@@ -65,7 +65,6 @@ uv pip install flash-attn --no-build-isolation -v

 ## Next Steps

- [Quick Start](quick_start.md) - Generate your first video
- [Inference Cookbook](../cookbook/index.md) - Choose a maintained recipe
+- [Quick Start Guide](quick_start.md) - Get started with your first video generation
 - [Configuration](../inference/configuration.md) - Learn about configuration options
- [Examples](../inference/examples/examples_inference_index.md) - Explore scripts and notebooks
+- [Examples](../inference/examples/examples_inference_index.md) - Explore example scripts and notebooks
@@ -49,12 +49,10 @@ brew install ffmpeg

 ### Installation

-FastWan's native Apple Silicon runtime requires the `mlx` extra.
-
 #### With uv (recommended)

 ```bash
-uv pip install "fastvideo[mlx]"
+uv pip install fastvideo
 ```

 #### With Conda environment (alternative)
@@ -62,7 +60,7 @@ uv pip install "fastvideo[mlx]"
 `uv` works inside an active conda env too, so prefer `uv pip` for the actual install:

 ```bash
-uv pip install "fastvideo[mlx]"
+uv pip install fastvideo
 ```

 ### Installation from Source
@@ -78,13 +76,13 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 Basic installation:

 ```bash
-uv pip install -e ".[mlx]"
+uv pip install -e .
 ```

 Alternative with Conda environment:

 ```bash
-uv pip install -e ".[mlx]"
+uv pip install -e .
 ```

 ## Development Environment Setup
@@ -23,21 +23,61 @@ Also optionally install flash-attn:
 uv pip install flash-attn --no-build-isolation -v
 ```

-## Choose a maintained recipe
+## Basic Usage

-The cookbook selects complete, checked-in recipes instead of mixing model,
-parallelism, offload, and attention settings independently.
+### Text-to-Video Generation

-[Open the inference cookbook](../cookbook/index.md){ .md-button .md-button--primary }
+```python
+from fastvideo import VideoGenerator

-!!! tip "Need more control?"
-    Start from a maintained recipe, then use the
-    [configuration](../inference/configuration.md) and
-    [optimization](../inference/optimizations.md) guides for supported changes.
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+### Image-to-Video Generation
+
+```python
+from fastvideo import VideoGenerator, SamplingParam
+
+def main():
+    # Create the generator
+    model_name = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    generator = VideoGenerator.from_pretrained(model_name, num_gpus=1)
+
+    # Set up parameters with an initial image
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.num_frames = 107
+
+    # Generate video based on the image
+    prompt = "A photograph coming to life with gentle movement"
+    generator.generate_video(prompt, sampling_param=sampling_param,
+                             output_path="my_videos/",
+                             save_video=True)
+
+if __name__ == '__main__':
+    main()
+```

 ## Next Steps

- [Inference Cookbook](../cookbook/index.md) - Choose a maintained recipe
 - [Installation Guide](installation.md) - Detailed installation instructions
 - [Configuration](../inference/configuration.md) - Learn about configuration options
 - [Examples](../inference/examples/examples_inference_index.md) - Explore more
@@ -178,17 +178,6 @@ optimizations: absence means **untested**, not incompatible.
 | Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ |

-## Apple Silicon native runtime
-
-| Release path | Model | Mode | Validated hardware | Status |
-| --- | --- | --- | --- | --- |
-| MLX FastWan T2V | FastWan-QAD-INT8-1.3B `[release model ID pending]` | 480x832, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 36 GB unified-memory class, MLX 0.31.2 | Release candidate; requires release-owner visual sign-off |
-
-This is a text-to-video-only source-install release. It is validated on the
-hardware listed above; MLX allocator caps are not evidence of support for a
-physical 16 GB Mac. See [Apple Silicon FastWan](../getting_started/installation/mps.md)
-for the supported command and release gates.
-
 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.

 ***Lucy Edit Dev uses a non-commercial model license. FastVideo support is
@@ -33,14 +33,6 @@ For the typed config/request path added during the inference API refactor:
 python examples/inference/basic/basic_dmd_new_api.py
 ```

-For the few-step (4-step, DMD2-distilled) MiniMax-H3 preview, generating synchronized video and audio, optionally with block-sparse VSA attention:
-```
-python examples/inference/basic/basic_fasth3.py --prompt "your prompt" [--vsa-sparsity 0.9]
-```
-The default checkpoint `FastVideo/FastVideo-Minimax-FastH3-Preview-v0.1` is private on the Hub while its license review completes; until it flips public, pass `--model-path` with a local snapshot of the release.
-
-On Blackwell (sm_100) GPUs with a `fastvideo-kernel` build that carries the sm_100a block-sparse extension, `--vsa-kernel sm100a` routes the tile-64 attention forwards through the CUDA kernel instead of Triton (it sets `FASTVIDEO_VSA_SM100A=1` before the pipeline boots); if the extension or the arch is missing, the run warns once and falls back to Triton.
-
 ## Basic Walkthrough

 All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines!
@@ -1,180 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Few-step video+audio generation with the DMD2-distilled MiniMax H3 preview.
-
-FastVideo/FastVideo-Minimax-FastH3-Preview-v0.1 is a 4-step distillation of
-MiniMaxAI/MiniMax-H3 (data-free DMD2): it walks a 4-step grid on the release
-sampler's shift-12 schedule instead of the base model's 50 steps, generating
-synchronized video and audio in one pipeline call.
-
-The student was trained with block-sparse video attention (VSA, 64-token
-tiles) and its checkpoint carries the trained sparse-gate parameters
-(``attn.to_gate_compress``), so this script always runs the VSA-H3 attention
-backend. At the default ``--vsa-sparsity 0.0`` the attention math is exactly
-dense (every tile is selected); raise the sparsity for additional speedup.
-"""
-
-from __future__ import annotations
-
-import argparse
-import os
-from pathlib import Path
-
-from fastvideo import VideoGenerator
-from fastvideo.api import (
-    CompileConfig,
-    EngineConfig,
-    GenerationRequest,
-    GeneratorConfig,
-    OffloadConfig,
-    OutputConfig,
-    ParallelismConfig,
-    PipelineSelection,
-    SamplingConfig,
-)
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--model-path", default="FastVideo/FastVideo-Minimax-FastH3-Preview-v0.1")
-    # The HF repo is private while the MiniMax H3 Community License review
-    # completes; until it flips public, pass --model-path with a local
-    # snapshot of the release instead (e.g. the team export at
-    # /mnt/lustre/vlm-wlsaidhi/fastvideo/exports/FastVideo-Minimax-FastH3-Preview-v0.1).
-    parser.add_argument("--prompt", required=True)
-    parser.add_argument("--output", default="outputs/fasth3")
-    parser.add_argument("--height", type=int, default=768)
-    parser.add_argument("--width", type=int, default=1344)
-    parser.add_argument("--num-frames", type=int, default=124)
-    # num_inference_steps counts sigma-GRID POINTS, matching the base model's
-    # convention ("50 steps" = a 50-point grid = 49 transformer forwards). The
-    # student's distilled 4-step grid is 4 FORWARDS, i.e. a 5-point grid
-    # (t = 1000, 750, 500, 250 -> 0 on the shift-12 schedule) — so the correct
-    # default here is 5. Other grids are off-distribution.
-    parser.add_argument("--steps",
-                        type=int,
-                        default=5,
-                        help="num_inference_steps = sigma-grid points; N points run N-1 denoising "
-                        "forwards. 5 (default) is the distilled 4-forward grid")
-    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--num-gpus", type=int, default=4)
-    parser.add_argument("--vsa-sparsity",
-                        type=float,
-                        default=0.0,
-                        help="Run-level VSA sparsity in [0, 1). 0.0 (default) selects every tile, which is "
-                        "exactly dense attention; the student was trained at 0.9")
-    # 64 is the trained contract: the student was TRAINED with 64-token
-    # (4,4,4) tiles, and its to_gate_compress gates were learned against
-    # pooling at that granularity — keep 64 unless you are ablating.
-    parser.add_argument("--vsa-tile-size",
-                        type=int,
-                        choices=(64, 256),
-                        default=64,
-                        help="VSA-H3 tile size in tokens; 64 (default) is what the student was trained "
-                        "with and runs the native Triton block-sparse path, 256 is the FA4-CuTe-capable "
-                        "geometry for ablations")
-    parser.add_argument("--vsa-kernel",
-                        choices=("triton", "sm100a"),
-                        default="triton",
-                        help="Block-sparse kernel for the tile-64 attention forward: triton (default, "
-                        "portable fwd+bwd) or sm100a — the opt-in Blackwell CUDA forward "
-                        "(fastvideo_kernel.block_sparse_attn_sm100a). sm100a needs an sm_100 GPU and a "
-                        "fastvideo-kernel build that carries the extension; if a precondition fails at "
-                        "run time the attention layer logs one warning and falls back to Triton. Only "
-                        "meaningful with --vsa-tile-size 64")
-    parser.add_argument("--torch-compile", action="store_true", help="torch.compile the DiT transformer path")
-    parser.add_argument("--compile-mode",
-                        default=None,
-                        help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
-    parser.add_argument("--repeats",
-                        type=int,
-                        default=1,
-                        help="generate N times; with --torch-compile the first run pays "
-                        "compilation, so steady-state is the last repeat")
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    output_dir = Path(args.output)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    if args.vsa_kernel == "sm100a":
-        # The attention backend reads FASTVIDEO_VSA_SM100A per forward; set it
-        # before the pipeline boots so spawned GPU workers inherit it. The
-        # kernel is forward-only and inference runs under no-grad, so every
-        # denoising forward qualifies for the CUDA route.
-        os.environ["FASTVIDEO_VSA_SM100A"] = "1"
-
-    # Boot-time run configuration, folded into FastVideoArgs (the same route
-    # examples/inference/basic/basic_minimax_h3_t2v.py uses for sparsity):
-    # - attention_backend: the checkpoint carries trained to_gate_compress
-    #   gates, which only exist under the VSA-H3 backend — a dense-backend
-    #   load would reject them as unexpected weights. Layers that do not
-    #   support VSA-H3 (e.g. the token refiner) fall back to flash attention.
-    # - VSA_tile_size: forwarded even at sparsity 0.0 because the gate-compress
-    #   branch pools per tile, and the gates were trained at 64 tokens/tile.
-    experimental: dict[str, object] = {
-        "attention_backend": "VIDEO_SPARSE_ATTN_H3",
-        "VSA_tile_size": args.vsa_tile_size,
-    }
-    if args.vsa_sparsity > 0.0:
-        experimental["VSA_sparsity"] = args.vsa_sparsity
-
-    generator = VideoGenerator.from_config(
-        GeneratorConfig(
-            model_path=args.model_path,
-            pipeline=PipelineSelection(experimental=experimental),
-            engine=EngineConfig(
-                num_gpus=args.num_gpus,
-                use_fsdp_inference=args.num_gpus > 1,
-                parallelism=ParallelismConfig(tp_size=1, sp_size=args.num_gpus),
-                offload=OffloadConfig(
-                    dit=False,
-                    dit_layerwise=False,
-                    text_encoder=True,
-                    vae=True,
-                    pin_cpu_memory=False,
-                ),
-                compile=CompileConfig(
-                    enabled=args.torch_compile,
-                    mode=args.compile_mode,
-                ),
-            ),
-        ))
-    try:
-        request = GenerationRequest(
-            prompt=args.prompt,
-            negative_prompt="",
-            sampling=SamplingConfig(
-                height=args.height,
-                width=args.width,
-                num_frames=args.num_frames,
-                fps=24,
-                num_inference_steps=args.steps,
-                # the base model is guidance-distilled; the student inherits it
-                guidance_scale=1.0,
-                batch_cfg=False,
-                seed=args.seed,
-            ),
-            output=OutputConfig(
-                output_path=str(output_dir / "fasth3.mp4"),
-                save_video=True,
-                return_frames=False,
-            ),
-        )
-        result = generator.generate(request)
-        print(f"Output written to: {result.video_path}")
-        if result.generation_time is not None:
-            # machine-readable: benchmark harnesses parse this line to separate
-            # generation from model-load time (last occurrence = steady state)
-            print(f"Generation time: {result.generation_time:.2f}s")
-        for _ in range(args.repeats - 1):
-            result = generator.generate(request)
-            if result.generation_time is not None:
-                print(f"Generation time: {result.generation_time:.2f}s")
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
@@ -27,8 +27,7 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3")
    # Rank-reduced AdaLN checkpoint (-39% params, -23 GiB VRAM): pass
    #   --model-path noctuashap/MiniMax-H3-pruned-r16
-    # (or a local dir produced by
-    #   scripts/checkpoint_conversion/convert_minimax_h3_adaln_rank.py).
+    # (or a local dir converted with tools/minimax_h3/fit_adaln_basis.py).
    # adaln_rank is read from the checkpoint config; no other flags needed.
    # Rank-reduced checkpoints are inference-only: training needs the
    # full-rank release.
@@ -27,8 +27,7 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3")
    # Rank-reduced AdaLN checkpoint (-39% params, -23 GiB VRAM): pass
    #   --model-path noctuashap/MiniMax-H3-pruned-r16
-    # (or a local dir produced by
-    #   scripts/checkpoint_conversion/convert_minimax_h3_adaln_rank.py).
+    # (or a local dir converted with tools/minimax_h3/fit_adaln_basis.py).
    # adaln_rank is read from the checkpoint config; no other flags needed.
    # Rank-reduced checkpoints are inference-only: training needs the
    # full-rank release.
@@ -24,8 +24,7 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3")
    # Rank-reduced AdaLN checkpoint (-39% params, -23 GiB VRAM): pass
    #   --model-path noctuashap/MiniMax-H3-pruned-r16
-    # (or a local dir produced by
-    #   scripts/checkpoint_conversion/convert_minimax_h3_adaln_rank.py).
+    # (or a local dir converted with tools/minimax_h3/fit_adaln_basis.py).
    # adaln_rank is read from the checkpoint config; no other flags needed.
    # Rank-reduced checkpoints are inference-only: training needs the
    # full-rank release.
@@ -1,46 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Tiny MLX RIFE frame-interpolation smoke test."""
-
-from __future__ import annotations
-
-import argparse
-import time
-
-import numpy as np
-
-from fastvideo.mlx_runtime.rife_interp import interpolate, load_model
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(
-        description="MLX RIFE 4.25 frame interpolation smoke test."
-    )
-    parser.add_argument(
-        "--self-test",
-        action="store_true",
-        help="Run a tiny two-frame interpolation test.",
-    )
-    args = parser.parse_args()
-    if not args.self_test:
-        raise SystemExit("Nothing to do; pass --self-test")
-
-    frame0 = np.zeros((64, 96, 3), dtype=np.uint8)
-    frame1 = np.zeros((64, 96, 3), dtype=np.uint8)
-    frame1[:, :, 0] = 255
-    start = time.perf_counter()
-    model = load_model()
-    load_s = time.perf_counter() - start
-    start = time.perf_counter()
-    frames = interpolate([frame0, frame1], factor=2, model=model)
-    interp_s = time.perf_counter() - start
-    assert len(frames) == 3
-    assert frames[1].shape == frame0.shape
-    assert frames[1].dtype == np.uint8
-    print(
-        "MLX RIFE self-test passed: "
-        f"load_s={load_s:.3f} interp_s={interp_s:.3f} shape={frames[1].shape}"
-    )
-
-
-if __name__ == "__main__":
-    main()
@@ -1,510 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""End-to-end Wan2.2-TI2V-5B generation on Apple Silicon (MLX DiT + MLX TAEHV).
-
-Pipeline: torch/MPS UMT5 encode (shared with 1.3B) → MLXWan22DiT 3-step DMD
-(warped schedule, flow_shift=5) → MLX TAEHV decode (taew2_2.pth). Fully MLX
-on the heavy DiT + decode path.
-
-    PYTHONPATH=$PWD python examples/inference/basic/mlx_wan22_generate.py \
-      --prompt "A red fox trotting through a snowy pine forest at golden hour" \
-      --output-path video_samples/demo_5b/fox_5b_mlx.mp4
-
-Decoder backends: ``taehv`` (default, MLX, ~seconds), ``taehv-torch`` (parity),
-``wan-vae`` (full AutoencoderKLWan on MPS, slow).
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import time
-from pathlib import Path
-
-import numpy as np
-
-from fastvideo.mlx_runtime.fast_spatial import DEFAULT_FAST_SPATIAL_SHARPEN
-from fastvideo.mlx_runtime.frame_upsample import DEFAULT_PIXEL_UPSAMPLE_MODE, PIXEL_UPSAMPLE_MODES
-from fastvideo.mlx_runtime.memory import cleanup_mlx
-from fastvideo.mlx_runtime.prompt_cache import (
-    fingerprint_digest,
-    load_prompt_cache,
-    save_prompt_cache,
-    text_encoder_fingerprint,
-)
-from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count
-
-FASTWAN21_MODEL_ID = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
-FASTWAN22_MODEL_ID = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
-DEFAULT_HEIGHT = 448
-DEFAULT_WIDTH = 832
-DEFAULT_NUM_FRAMES = 121
-
-def _resolve_model_paths(
-    *,
-    text_encoder_root: Path | None,
-    dit_checkpoint: Path | None,
-    dit_config: Path | None,
-    vae_root: Path | None,
-    mlx_checkpoint: Path | None,
-    decode_backend: str,
-) -> tuple[Path, Path | None, Path | None, Path | None]:
-    """Download only the missing assets required by the selected Wan2.2 path."""
-    from huggingface_hub import snapshot_download
-
-    if text_encoder_root is None:
-        text_encoder_root = Path(snapshot_download(
-            FASTWAN21_MODEL_ID,
-            allow_patterns=["tokenizer/*", "text_encoder/*"],
-        ))
-    if mlx_checkpoint is None and (dit_checkpoint is None or dit_config is None):
-        patterns = []
-        if dit_checkpoint is None:
-            patterns.append("transformer/diffusion_pytorch_model.safetensors")
-        if dit_config is None:
-            patterns.append("transformer/config.json")
-        model_root = Path(snapshot_download(FASTWAN22_MODEL_ID, allow_patterns=patterns))
-        dit_checkpoint = dit_checkpoint or model_root / "transformer/diffusion_pytorch_model.safetensors"
-        dit_config = dit_config or model_root / "transformer/config.json"
-    if decode_backend == "wan-vae" and vae_root is None:
-        model_root = Path(snapshot_download(FASTWAN22_MODEL_ID, allow_patterns=["vae/*"]))
-        vae_root = model_root / "vae"
-    return text_encoder_root, dit_checkpoint, dit_config, vae_root
-
-
-def _prompt_cache_fingerprint(
-    *,
-    prompt: str,
-    prompt_used: str,
-    enhance_prompt: bool,
-    enhance_prompt_backend: str,
-    text_encoder_root: Path,
-    max_sequence_length: int,
-    dtype: str,
-) -> dict[str, object]:
-    return {
-        "prompt": prompt,
-        "prompt_used": prompt_used,
-        "enhance_prompt": enhance_prompt,
-        "enhance_prompt_backend": enhance_prompt_backend,
-        "text_encoder": text_encoder_fingerprint(text_encoder_root),
-        "max_sequence_length": max_sequence_length,
-        "dtype": dtype,
-    }
-
-
-def _default_prompt_cache_path(fingerprint: dict[str, object]) -> Path:
-    """Content-addressed default cache file for a prompt fingerprint.
-
-    The Wan2.1 entrypoint caches prompt embeddings by default; this one only
-    did so when handed an explicit ``--prompt-embeds-cache`` path, so every 5B
-    run paid a full UMT5 encode (~45s on an M4 Max) even for a repeat prompt.
-    The fingerprint already covers everything that changes the embedding, so
-    hash it for the filename.
-    """
-    digest = fingerprint_digest(fingerprint)[:32]
-    return Path.home() / ".cache" / "fastvideo" / "prompt_embeds" / f"wan22_{digest}.npy"
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(
-        description="MLX Wan2.2-5B T2V (encode → DiT DMD → TAEHV/VAE decode)"
-    )
-    parser.add_argument(
-        "--prompt",
-        default="A red fox trotting through a snowy pine forest at golden hour, cinematic",
-    )
-    parser.add_argument(
-        "--output-path",
-        type=Path,
-        default=Path("video_samples/demo_5b/fox_5b_mlx.mp4"),
-    )
-    parser.add_argument(
-        "--text-encoder-root",
-        type=Path,
-        default=None,
-        help="Root with text_encoder/ + tokenizer/",
-    )
-    parser.add_argument(
-        "--prompt-embeds-cache",
-        type=Path,
-        default=None,
-        help="Explicit .npy UMT5 embedding cache file. Overrides the automatic "
-        "content-addressed cache (--prompt-cache).",
-    )
-    parser.add_argument(
-        "--prompt-cache",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Cache prompt embeddings under ~/.cache/fastvideo/prompt_embeds so "
-        "repeat runs skip the text encoder entirely. Default: on.",
-    )
-    parser.add_argument(
-        "--text-encoder-device",
-        choices=("auto", "cpu", "mps"),
-        default="cpu",
-        help="Device for UMT5 encoding. CPU is safest beside the 5B MLX DiT.",
-    )
-    parser.add_argument(
-        "--enhance-prompt",
-        action="store_true",
-        help="Apply deterministic local cinematic prompt enrichment before UMT5.",
-    )
-    parser.add_argument(
-        "--enhance-prompt-backend",
-        choices=("template",),
-        default="template",
-        help="Prompt enrichment backend.",
-    )
-    parser.add_argument(
-        "--dit-checkpoint",
-        type=Path,
-        default=None,
-    )
-    parser.add_argument("--dit-config", type=Path, default=None)
-    parser.add_argument(
-        "--mlx-checkpoint",
-        type=Path,
-        default=None,
-        help="Pre-quantized MLX DiT checkpoint directory. Rewrapped with Wan2.2 per-token conditioning.",
-    )
-    parser.add_argument("--vae-root", type=Path, default=None)
-    parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT)
-    parser.add_argument("--width", type=int, default=DEFAULT_WIDTH)
-    parser.add_argument(
-        "--num-frames",
-        type=int,
-        default=DEFAULT_NUM_FRAMES,
-        help="Pixel frames (121 at 24fps = 5.04 seconds)",
-    )
-    parser.add_argument("--seed", type=int, default=1234)
-    parser.add_argument("--renoise-seed", type=int, default=0)
-    parser.add_argument("--fps", type=int, default=24)
-    parser.add_argument("--flow-shift", type=float, default=5.0)
-    parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
-    parser.add_argument(
-        "--no-warp",
-        action="store_true",
-        help="Disable schedule warping (debug only).",
-    )
-    parser.add_argument(
-        "--fast",
-        action="store_true",
-        help="Generate fewer frames then RIFE-interpolate to --num-frames.",
-    )
-    parser.add_argument("--fast-factor", type=int, default=2)
-    parser.add_argument("--fast-sharpen", type=float, default=0.6)
-    parser.add_argument(
-        "--fast-spatial",
-        action="store_true",
-        help="Denoise and decode at reduced spatial resolution, then resample "
-        "the decoded frames up to the target size.",
-    )
-    parser.add_argument("--fast-spatial-scale", type=int, default=2)
-    parser.add_argument(
-        "--fast-spatial-upsample-mode",
-        choices=PIXEL_UPSAMPLE_MODES,
-        default=DEFAULT_PIXEL_UPSAMPLE_MODE,
-    )
-    parser.add_argument("--fast-spatial-sharpen", type=float, default=DEFAULT_FAST_SPATIAL_SHARPEN)
-    parser.add_argument(
-        "--refine",
-        action="store_true",
-        help="Two-pass DMD: coarse denoise, upsample/re-noise, full-res denoise.",
-    )
-    parser.add_argument("--refine-scale", type=int, default=2)
-    parser.add_argument(
-        "--refine-upsample-mode",
-        choices=("bilinear", "nearest"),
-        default="bilinear",
-    )
-    parser.add_argument("--no-refine-add-noise", action="store_true")
-    parser.add_argument(
-        "--decode-backend",
-        choices=("taehv", "taehv-torch", "wan-vae"),
-        default="taehv",
-    )
-    parser.add_argument("--save-latents", type=Path, default=None)
-    parser.add_argument("--metrics-json", type=Path, default=None,
-                        help="Write measured run metadata as JSON for reports or galleries.")
-    parser.add_argument(
-        "--compile",
-        action="store_true",
-        help="Compile the DiT forward with mx.compile; fallback to eager on failure.",
-    )
-    args = parser.parse_args()
-
-    if args.fast_factor < 2:
-        parser.error("--fast-factor must be at least 2")
-    # --fast-spatial used to be rejected here because it upsampled the completed
-    # 48-channel latent, which is out of distribution for the decoder and gave
-    # black or noisy video. The upsample now runs on decoded frames, so the
-    # latent never leaves the grid it was denoised on and the mode is usable.
-    if args.refine and args.fast_spatial:
-        print("[wan22] --refine takes precedence over --fast-spatial")
-    args.text_encoder_root, args.dit_checkpoint, args.dit_config, args.vae_root = _resolve_model_paths(
-        text_encoder_root=args.text_encoder_root,
-        dit_checkpoint=args.dit_checkpoint,
-        dit_config=args.dit_config,
-        vae_root=args.vae_root,
-        mlx_checkpoint=args.mlx_checkpoint,
-        decode_backend=args.decode_backend,
-    )
-    target_frames = args.num_frames
-    if args.fast:
-        args.num_frames = aligned_keyframe_count(target_frames, args.fast_factor)
-        print(
-            f"[wan22 fast] generating {args.num_frames} frames, "
-            f"RIFE {args.fast_factor}x -> {target_frames}"
-        )
-
-    import mlx.core as mx
-    import torch
-
-    from examples.inference.basic.mlx_wan_prompt_to_video import (
-        _postprocess_video,
-        encode_prompt,
-        make_rotary_embeddings,
-    )
-    from fastvideo.mlx_runtime.fast_spatial import plan_fast_spatial
-    from fastvideo.mlx_runtime.refine import (
-        default_refine_timesteps,
-        plan_refine_resolutions,
-        prepare_refine_latents,
-    )
-    from fastvideo.mlx_runtime.wan22 import (
-        mlx_wan22_dit_from_diffusers_safetensors,
-        mlx_wan22_dit_from_mlx_checkpoint,
-    )
-    from fastvideo.mlx_runtime.wan22_sample import build_wan22_dmd_schedule, sample_wan22_dmd
-    from fastvideo.mlx_runtime.wan_vae import decode_latents_to_video
-
-    if args.mlx_checkpoint is not None:
-        config = json.loads((args.mlx_checkpoint / "mlx_dit.json").read_text())["config"]
-    else:
-        config = json.loads(args.dit_config.read_text())
-    patch_size = tuple(config.get("patch_size", (1, 2, 2)))
-    if args.refine:
-        active_plan = plan_refine_resolutions(
-            height=args.height, width=args.width, num_frames=args.num_frames,
-            spatial_scale=args.refine_scale, vae_spatial_compression=16,
-            vae_temporal_compression=4, patch_size=patch_size, enabled=True,
-        )
-        spatial_mode = "refine"
-    elif args.fast_spatial:
-        fast_spatial_plan = plan_fast_spatial(
-            height=args.height, width=args.width, num_frames=args.num_frames,
-            spatial_scale=args.fast_spatial_scale, vae_spatial_compression=16,
-            vae_temporal_compression=4, patch_size=patch_size,
-            upsample_mode=args.fast_spatial_upsample_mode,
-            sharpen=args.fast_spatial_sharpen, enabled=True,
-        )
-        active_plan = fast_spatial_plan.plan
-        spatial_mode = "fast_spatial"
-    else:
-        active_plan = plan_refine_resolutions(
-            height=args.height, width=args.width, num_frames=args.num_frames,
-            spatial_scale=1, vae_spatial_compression=16, vae_temporal_compression=4,
-            patch_size=patch_size, enabled=False,
-        )
-        spatial_mode = "off"
-    lat_h, lat_w = active_plan.stage1_latent_height, active_plan.stage1_latent_width
-    lat_t = active_plan.latent_frames
-    in_ch = int(config["in_channels"])
-    print(f"[5B] latent {in_ch}x{lat_t}x{lat_h}x{lat_w}", flush=True)
-
-    total_start = time.perf_counter()
-    prompt_for_encode = args.prompt
-    enhance_backend = None
-    enhance_elapsed_s = 0.0
-    if args.enhance_prompt:
-        from fastvideo.mlx_runtime.prompt_enhance import enhance_prompt
-
-        enhancement = enhance_prompt(args.prompt, backend=args.enhance_prompt_backend)
-        prompt_for_encode = enhancement.enhanced
-        enhance_backend = enhancement.backend
-        enhance_elapsed_s = enhancement.elapsed_s
-        print(f"[enhance] backend={enhance_backend} in {enhance_elapsed_s:.2f}s", flush=True)
-        print(f"[enhance] prompt: {prompt_for_encode}", flush=True)
-
-    t0 = time.perf_counter()
-    prompt_cache_fingerprint = _prompt_cache_fingerprint(
-        prompt=args.prompt,
-        prompt_used=prompt_for_encode,
-        enhance_prompt=args.enhance_prompt,
-        enhance_prompt_backend=args.enhance_prompt_backend,
-        text_encoder_root=args.text_encoder_root,
-        max_sequence_length=512,
-        dtype="fp16",
-    )
-    prompt_cache_path = args.prompt_embeds_cache
-    if prompt_cache_path is None and args.prompt_cache:
-        prompt_cache_path = _default_prompt_cache_path(prompt_cache_fingerprint)
-    cached_embeds = load_prompt_cache(
-        prompt_cache_path,
-        prompt_cache_fingerprint,
-    )
-    if cached_embeds is not None:
-        embeds = torch.from_numpy(cached_embeds).contiguous()
-    else:
-        embeds = encode_prompt(
-            model_root=args.text_encoder_root,
-            prompt=prompt_for_encode,
-            max_sequence_length=512,
-            device_arg=args.text_encoder_device,
-            dtype_arg="fp16",
-        )
-        save_prompt_cache(
-            prompt_cache_path,
-            embeds.cpu().numpy(),
-            prompt_cache_fingerprint,
-        )
-    ehs = mx.array(embeds.numpy()).astype(mx.float16)
-    prompt_encode_s = time.perf_counter() - t0
-    print(f"[5B] prompt encoded {tuple(ehs.shape)} in {prompt_encode_s:.1f}s", flush=True)
-
-    t1 = time.perf_counter()
-    if args.mlx_checkpoint is not None:
-        dit = mlx_wan22_dit_from_mlx_checkpoint(
-            args.mlx_checkpoint,
-            compile=args.compile,
-        )
-    else:
-        dit = mlx_wan22_dit_from_diffusers_safetensors(
-            args.dit_checkpoint,
-            args.dit_config,
-            dtype="fp16",
-            compile=args.compile,
-        )
-    dit_load_s = time.perf_counter() - t1
-    print(f"[5B] DiT loaded in {dit_load_s:.1f}s", flush=True)
-
-    freqs = make_rotary_embeddings(config, latent_frames=lat_t, latent_height=lat_h, latent_width=lat_w)
-    gen = torch.Generator().manual_seed(args.seed)
-    noise = mx.array(
-        torch.randn(1, in_ch, lat_t, lat_h, lat_w, generator=gen, dtype=torch.float32).numpy()).astype(mx.float16)
-
-    steps = [int(s) for s in args.dmd_denoising_steps.split(",") if s.strip()]
-    t2 = time.perf_counter()
-    mx.reset_peak_memory()
-    latents = sample_wan22_dmd(
-        dit,
-        ehs,
-        noise,
-        freqs,
-        dmd_denoising_steps=steps,
-        flow_shift=args.flow_shift,
-        warp_denoising_step=not args.no_warp,
-        seed=args.renoise_seed,
-    )
-    if spatial_mode == "refine":
-        schedule, warped_steps = build_wan22_dmd_schedule(
-            steps, flow_shift=args.flow_shift, warp_denoising_step=not args.no_warp,
-        )
-        # The grid opens at sigma == 1, where the hand-off
-        # `(1 - sigma) * upsampled + sigma * noise` weights stage 1 at zero and
-        # refine silently becomes a plain full-res run. Drop the leading
-        # full-noise steps so stage 1 actually reaches stage 2.
-        stage2_warped = default_refine_timesteps(schedule, warped_steps)
-        stage2_steps = steps[len(warped_steps) - len(stage2_warped):]
-        sigma = schedule.sigma_for(stage2_warped[0])
-        print(f"[5B refine] stage-2 steps={stage2_steps} sigma={sigma:.4f} "
-              f"(stage-1 weight {1.0 - sigma:.4f})", flush=True)
-        latents = prepare_refine_latents(
-            latents, scale=args.refine_scale, sigma=sigma,
-            add_noise_flag=not args.no_refine_add_noise,
-            upsample_mode=args.refine_upsample_mode, seed=args.renoise_seed + 1,
-        )
-        freqs_stage2 = make_rotary_embeddings(
-            config, latent_frames=lat_t,
-            latent_height=active_plan.stage2_latent_height,
-            latent_width=active_plan.stage2_latent_width,
-        )
-        latents = sample_wan22_dmd(
-            dit, ehs, latents, freqs_stage2, dmd_denoising_steps=stage2_steps,
-            flow_shift=args.flow_shift, warp_denoising_step=not args.no_warp,
-            seed=args.renoise_seed + 2,
-        )
-    # spatial_mode == "fast_spatial" leaves the latents on the stage-1 grid;
-    # the resample happens after decode, in _postprocess_video.
-    denoise_s = time.perf_counter() - t2
-    peak = mx.get_peak_memory() / (1024**3)
-    print(f"[5B] denoise {len(steps)} steps in {denoise_s:.1f}s, peak {peak:.2f} GiB", flush=True)
-
-    latents_np = np.array(latents.astype(mx.float32))
-    if args.save_latents is not None:
-        args.save_latents.parent.mkdir(parents=True, exist_ok=True)
-        np.savez(args.save_latents, latents=latents_np, prompt=args.prompt, seed=args.seed)
-        print(f"[5B] wrote latents {args.save_latents}", flush=True)
-
-    if spatial_mode == "refine":
-        del freqs_stage2
-    del dit, latents, ehs, noise, freqs
-    cleanup_mlx()
-
-    metrics = decode_latents_to_video(
-        latents_np,
-        args.output_path,
-        fps=args.fps,
-        backend=args.decode_backend,
-        vae_dir=args.vae_root if args.decode_backend == "wan-vae" else None,
-        z_dim=in_ch,
-    )
-    # One h264 round-trip for both post-decode passes (see _postprocess_video).
-    rife_s = 0.0
-    rife_request = ({
-        "factor": args.fast_factor,
-        "target_frames": target_frames,
-        "sharpen": args.fast_sharpen,
-    } if args.fast else None)
-    spatial_request = fast_spatial_plan if spatial_mode == "fast_spatial" else None
-    if rife_request is not None or spatial_request is not None:
-        rife_start = time.perf_counter()
-        _postprocess_video(
-            video_path=args.output_path, fps=args.fps,
-            rife=rife_request, spatial=spatial_request,
-        )
-        rife_s = time.perf_counter() - rife_start
-    print(f"[5B] decoded via {metrics['backend']} in {metrics['decode_s']:.1f}s → {args.output_path}", flush=True)
-    summary = {
-        "output_path": str(args.output_path.resolve()),
-        "prompt": args.prompt,
-        "prompt_used": prompt_for_encode,
-        "enhance_prompt": args.enhance_prompt,
-        "enhance_backend": enhance_backend,
-        "enhance_elapsed_s": round(enhance_elapsed_s, 3),
-        "height": args.height,
-        "width": args.width,
-        "fps": args.fps,
-        "target_frames": target_frames,
-        "generated_frames": args.num_frames,
-        "seed": args.seed,
-        "renoise_seed": args.renoise_seed,
-        "dmd_denoising_steps": steps,
-        "flow_shift": args.flow_shift,
-        "warp": not args.no_warp,
-        "spatial_mode": spatial_mode,
-        "fast": args.fast,
-        "fast_factor": args.fast_factor if args.fast else None,
-        "fast_spatial_scale": args.fast_spatial_scale if args.fast_spatial else None,
-        "refine_scale": args.refine_scale if args.refine else None,
-        "decode_backend": args.decode_backend,
-        "prompt_encode_s": round(prompt_encode_s, 3),
-        "dit_load_s": round(dit_load_s, 3),
-        "denoise_s": round(denoise_s, 3),
-        "decode_s": round(metrics["decode_s"], 3),
-        "rife_s": round(rife_s, 3),
-        "wall_total_s": round(time.perf_counter() - total_start, 3),
-        "peak_gib": round(peak, 3),
-        "latent_shape": [in_ch, lat_t, lat_h, lat_w],
-        "stage2_latent_shape": [in_ch, lat_t, active_plan.stage2_latent_height, active_plan.stage2_latent_width],
-        "mlx_checkpoint": str(args.mlx_checkpoint.resolve()) if args.mlx_checkpoint else None,
-    }
-    if args.metrics_json is not None:
-        args.metrics_json.parent.mkdir(parents=True, exist_ok=True)
-        args.metrics_json.write_text(json.dumps(summary, indent=2) + "\n")
-        print(f"[5B] wrote metrics {args.metrics_json}", flush=True)
-    print(json.dumps(summary, indent=2), flush=True)
-
-
-if __name__ == "__main__":
-    main()
@@ -1,130 +0,0 @@
-"""Compare Wan VAE and TAEHV decode on saved FastWan latents."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import time
-from pathlib import Path
-
-import numpy as np
-
-from examples.inference.basic.mlx_wan_prompt_to_video import DEFAULT_MODEL_ROOT, decode_latents_to_video
-
-
-def _torch_mps_memory() -> dict[str, int | None]:
-    """
-    Report current and recommended memory usage for the MPS backend.
-
-    Returns:
-        dict[str, int | None]: Memory metrics in bytes, or `None` values when
-            PyTorch or MPS is unavailable.
-    """
-    try:
-        import torch
-    except ImportError:
-        return {
-            "current_allocated_bytes": None,
-            "driver_allocated_bytes": None,
-            "recommended_max_bytes": None,
-        }
-    if not torch.backends.mps.is_available():
-        return {
-            "current_allocated_bytes": None,
-            "driver_allocated_bytes": None,
-            "recommended_max_bytes": None,
-        }
-    return {
-        "current_allocated_bytes": int(torch.mps.current_allocated_memory()),
-        "driver_allocated_bytes": int(torch.mps.driver_allocated_memory()),
-        "recommended_max_bytes": int(torch.mps.recommended_max_memory()),
-    }
-
-
-def _parse_backends(raw: str) -> list[str]:
-    """
-    Parse and validate a comma-separated list of decoding backends.
-
-    Parameters:
-        raw (str): Comma-separated backend names.
-
-    Returns:
-        list[str]: Trimmed, supported backend names in input order.
-
-    Raises:
-        ValueError: If the input contains an unsupported backend.
-    """
-    backends = [backend.strip() for backend in raw.split(",") if backend.strip()]
-    allowed = {"wan-vae", "taehv"}
-    unknown = sorted(set(backends) - allowed)
-    if unknown:
-        raise ValueError(f"Unsupported decode backends: {unknown}")
-    return backends
-
-
-def main() -> None:
-    """
-    Benchmark selected Wan latent decoding backends and record their performance metrics.
-
-    Loads the specified latent array, decodes it with each selected backend, exports the
-    results as MP4 files, and writes per-backend timing and Torch MPS memory metrics to
-    `metrics.json`.
-    """
-    parser = argparse.ArgumentParser(description="Benchmark decode backends on saved Wan/FastWan latents.")
-    parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
-    parser.add_argument("--latents-path", type=Path, required=True)
-    parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_decode_benchmark"))
-    parser.add_argument("--backends", default="wan-vae,taehv")
-    parser.add_argument("--fps", type=int, default=16)
-    parser.add_argument("--torch-device", default="auto")
-    parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
-    parser.add_argument("--taehv-source-path", type=Path, default=None)
-    parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
-    parser.add_argument("--taehv-parallel", action="store_true")
-    args = parser.parse_args()
-
-    latents = np.load(args.latents_path)
-    args.output_dir.mkdir(parents=True, exist_ok=True)
-    rows = []
-
-    for backend in _parse_backends(args.backends):
-        print(f"=== Decode backend: {backend} ===")
-        before = _torch_mps_memory()
-        start = time.perf_counter()
-        output_path = args.output_dir / f"{args.latents_path.stem}_{backend}.mp4"
-        decode_latents_to_video(
-            model_root=args.model_root,
-            latents_np=latents,
-            output_path=output_path,
-            fps=args.fps,
-            device_arg=args.torch_device,
-            dtype_arg=args.torch_dtype,
-            backend=backend,
-            taehv_source_path=args.taehv_source_path,
-            taehv_checkpoint_path=args.taehv_checkpoint_path,
-            taehv_parallel=args.taehv_parallel,
-        )
-        elapsed = time.perf_counter() - start
-        after = _torch_mps_memory()
-        metrics = {
-            "backend": backend,
-            "latents_path": str(args.latents_path),
-            "latents_shape": list(latents.shape),
-            "decode_export_s": elapsed,
-            "torch_mps_current_before_bytes": before["current_allocated_bytes"],
-            "torch_mps_current_after_bytes": after["current_allocated_bytes"],
-            "torch_mps_driver_before_bytes": before["driver_allocated_bytes"],
-            "torch_mps_driver_after_bytes": after["driver_allocated_bytes"],
-            "torch_mps_recommended_max_bytes": after["recommended_max_bytes"],
-            "output_path": str(output_path),
-        }
-        rows.append(metrics)
-        print(json.dumps(metrics, indent=2))
-
-    metrics_path = args.output_dir / "metrics.json"
-    metrics_path.write_text(json.dumps(rows, indent=2))
-    print(f"Wrote decode metrics to: {metrics_path}")
-
-
-if __name__ == "__main__":
-    main()
@@ -1,365 +0,0 @@
-"""Benchmark MLX FastWan quantization modes with one shared prompt encode."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import time
-from pathlib import Path
-from typing import cast
-
-import numpy as np
-
-from examples.inference.basic.mlx_wan_prompt_to_video import (
-    DEFAULT_MODEL_ROOT,
-    decode_latents_to_video,
-    encode_prompt,
-    make_rotary_embeddings,
-)
-from fastvideo.mlx_runtime.memory import cleanup_mlx
-
-
-def _parse_modes(raw: str) -> list[str]:
-    """
-    Parse and validate a comma-separated list of quantization modes.
-
-    Parameters:
-        raw (str): Comma-separated mode names.
-
-    Returns:
-        list[str]: Normalized, whitespace-trimmed mode names.
-
-    Raises:
-        ValueError: If any mode is unsupported.
-    """
-    modes = [mode.strip() for mode in raw.split(",") if mode.strip()]
-    allowed = {"none", "int8", "int4", "mxfp8", "mxfp4", "nvfp4"}
-    unknown = sorted(set(modes) - allowed)
-    if unknown:
-        raise ValueError(f"Unsupported modes: {unknown}")
-    return modes
-
-
-def _latent_delta_metrics(candidate: np.ndarray, baseline: np.ndarray) -> dict[str, float]:
-    """
-    Compare candidate and baseline latent arrays using error and signal-quality metrics.
-
-    Parameters:
-        candidate (np.ndarray): Latent array to evaluate.
-        baseline (np.ndarray): Reference latent array for comparison.
-
-    Returns:
-        dict[str, float]: Mean squared error, mean absolute error, maximum absolute
-        error, and signal-to-noise ratio in decibels between the arrays.
-    """
-    diff = candidate.astype(np.float32) - baseline.astype(np.float32)
-    mse = float(np.mean(np.square(diff)))
-    mae = float(np.mean(np.abs(diff)))
-    max_abs = float(np.max(np.abs(diff)))
-    signal = float(np.mean(np.square(baseline.astype(np.float32))))
-    return {
-        "latent_mse_vs_fp16": mse,
-        "latent_mae_vs_fp16": mae,
-        "latent_max_abs_vs_fp16": max_abs,
-        "latent_snr_db_vs_fp16": float(10.0 * np.log10(signal / mse)) if mse > 0 else float("inf"),
-    }
-
-
-def _torch_mps_memory() -> dict[str, int | None]:
-    """
-    Report PyTorch MPS memory statistics when PyTorch MPS is available.
-
-    Returns:
-        dict[str, int | None]: A mapping of MPS memory metric names to byte counts, or `None` values when PyTorch or MPS is unavailable.
-    """
-    try:
-        import torch
-    except ImportError:
-        return {
-            "torch_mps_current_allocated_bytes": None,
-            "torch_mps_driver_allocated_bytes": None,
-            "torch_mps_recommended_max_bytes": None,
-        }
-    if not torch.backends.mps.is_available():
-        return {
-            "torch_mps_current_allocated_bytes": None,
-            "torch_mps_driver_allocated_bytes": None,
-            "torch_mps_recommended_max_bytes": None,
-        }
-    return {
-        "torch_mps_current_allocated_bytes": int(torch.mps.current_allocated_memory()),
-        "torch_mps_driver_allocated_bytes": int(torch.mps.driver_allocated_memory()),
-        "torch_mps_recommended_max_bytes": int(torch.mps.recommended_max_memory()),
-    }
-
-
-def _decode_with_metrics(*, args, latents: np.ndarray, output_path: Path) -> dict[str, float | int | None | str]:
-    """
-    Decode latents to a video and collect export timing and PyTorch MPS memory metrics.
-
-    Parameters:
-        args: Configuration values for decoding and video export.
-        latents (np.ndarray): Latent representation to decode.
-        output_path (Path): Destination path for the exported video.
-
-    Returns:
-        dict[str, float | int | None | str]: Video export duration and PyTorch MPS memory measurements.
-    """
-    before = _torch_mps_memory()
-    decode_start = time.perf_counter()
-    decode_latents_to_video(
-        model_root=args.model_root,
-        latents_np=latents,
-        output_path=output_path,
-        fps=args.fps,
-        device_arg=args.torch_device,
-        dtype_arg=args.torch_dtype,
-        backend=args.decode_backend,
-        taehv_source_path=args.taehv_source_path,
-        taehv_checkpoint_path=args.taehv_checkpoint_path,
-        taehv_parallel=args.taehv_parallel,
-    )
-    decode_time = time.perf_counter() - decode_start
-    after = _torch_mps_memory()
-    return {
-        "decode_export_s": decode_time,
-        "decode_torch_mps_current_before_bytes": before["torch_mps_current_allocated_bytes"],
-        "decode_torch_mps_current_after_bytes": after["torch_mps_current_allocated_bytes"],
-        "decode_torch_mps_driver_before_bytes": before["torch_mps_driver_allocated_bytes"],
-        "decode_torch_mps_driver_after_bytes": after["torch_mps_driver_allocated_bytes"],
-        "decode_torch_mps_recommended_max_bytes": after["torch_mps_recommended_max_bytes"],
-    }
-
-
-def _run_one_mode(
-    *,
-    mode: str,
-    args,
-    config: dict,
-    checkpoint_path: Path,
-    config_path: Path,
-    prompt_embeds,
-    freqs_cis,
-):
-    """
-    Run denoising for one quantization mode and collect performance and memory metrics.
-
-    Parameters:
-        mode (str): Quantization mode to benchmark.
-        args: Benchmark configuration, including dtype, dimensions, seed, scheduler, and denoising settings.
-        config (dict): Model configuration containing the input channel count.
-        checkpoint_path (Path): Path to the transformer checkpoint.
-        config_path (Path): Path to the transformer configuration.
-        prompt_embeds: Encoded prompt embeddings shared across benchmark modes.
-        freqs_cis: Rotary positional embeddings used during denoising.
-
-    Returns:
-        dict: The mode name, generated latent array, and metrics for model loading,
-        denoising, step timing, and MLX memory usage.
-    """
-    import mlx.core as mx
-    import torch
-
-    from fastvideo.benchmarks.mlx_fastwan_bench import denoise_dmd_on_device
-    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-    from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors
-    from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step
-
-    mx_dtype = mx.float16 if args.mlx_dtype == "fp16" else mx.float32
-    quantization = None if mode == "none" else mode
-    latent_frames = (args.num_frames - 1) // 4 + 1
-    latent_height = args.height // 8
-    latent_width = args.width // 8
-
-    load_start = time.perf_counter()
-    mx.clear_cache()
-    mx.reset_peak_memory()
-    dit = mlx_dit_from_diffusers_safetensors(
-        checkpoint_path,
-        config_path,
-        dtype=args.mlx_dtype,
-        quantization=quantization,
-    )
-    load_time = time.perf_counter() - load_start
-    load_peak_memory = mx.get_peak_memory()
-
-    scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift)
-    schedule = MLXDMDSchedule.from_torch_scheduler(scheduler)
-    timesteps = [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()]
-    # Same torch generator sequence as the original host-round-trip loop
-    # (initial latents first, then one re-noise draw per intermediate step),
-    # so every mode still shares identical stochasticity.
-    generator = torch.Generator(device="cpu").manual_seed(args.seed)
-    latents_seed = torch.randn(
-        (1, int(config["in_channels"]), latent_frames, latent_height, latent_width),
-        generator=generator,
-        dtype=torch.float32,
-    ).numpy()
-    renoise_by_step = [
-        torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy()
-        for _ in range(max(0, len(timesteps) - 1))
-    ]
-    latents = mx.array(latents_seed).astype(mx_dtype)
-    encoder_hidden_states = mx.array(prompt_embeds.numpy()).astype(mx_dtype)
-
-    denoise_start = time.perf_counter()
-    mx.reset_peak_memory()
-    latents_np, step_times = denoise_dmd_on_device(
-        mx=mx,
-        dit=dit,
-        latents=latents,
-        encoder_hidden_states=encoder_hidden_states,
-        freqs_cis=freqs_cis,
-        timesteps=timesteps,
-        renoise_by_step=renoise_by_step,
-        schedule=schedule,
-        dmd_step=dmd_step,
-        mx_dtype=mx_dtype,
-    )
-    denoise_time = time.perf_counter() - denoise_start
-    denoise_peak_memory = mx.get_peak_memory()
-    active_memory = mx.get_active_memory()
-    return {
-        "mode": mode,
-        "latents": latents_np,
-        "metrics": {
-            "mlx_dit_load_s": load_time,
-            "mlx_denoise_s": denoise_time,
-            "mlx_denoise_first_step_s": step_times[0] if step_times else None,
-            "mlx_load_peak_bytes": int(load_peak_memory),
-            "mlx_denoise_peak_bytes": int(denoise_peak_memory),
-            "mlx_active_after_denoise_bytes": int(active_memory),
-        },
-    }
-
-
-def main() -> None:
-    """
-    Run the MLX FastWan quantization benchmark for the selected modes and write latency, memory, output, and latent-difference metrics to the output directory.
-    """
-    parser = argparse.ArgumentParser(description="Benchmark MLX FastWan quantization modes.")
-    parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
-    parser.add_argument("--prompt", default="A snow leopard walks across a windy mountain ridge.")
-    parser.add_argument("--height", type=int, default=192)
-    parser.add_argument("--width", type=int, default=320)
-    parser.add_argument("--num-frames", type=int, default=17)
-    parser.add_argument("--dmd-denoising-steps", default="1000,757,522")
-    parser.add_argument("--flow-shift", type=float, default=8.0)
-    parser.add_argument("--max-sequence-length", type=int, default=256)
-    parser.add_argument("--seed", type=int, default=1024)
-    parser.add_argument("--fps", type=int, default=16)
-    parser.add_argument("--torch-device", default="auto")
-    parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
-    parser.add_argument("--mlx-dtype", choices=("fp16", "fp32"), default="fp16")
-    parser.add_argument("--modes", default="none,int8,int4,mxfp8,mxfp4,nvfp4")
-    parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_quant_benchmark"))
-    parser.add_argument("--decode-backend", choices=("none", "wan-vae", "taehv"), default="taehv")
-    parser.add_argument("--taehv-source-path", type=Path, default=None)
-    parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
-    parser.add_argument("--taehv-parallel", action="store_true")
-    args = parser.parse_args()
-
-    import mlx.core as mx
-    import torch
-
-    mx.random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    args.output_dir.mkdir(parents=True, exist_ok=True)
-
-    config_path = args.model_root / "transformer/config.json"
-    checkpoint_path = args.model_root / "transformer/diffusion_pytorch_model.safetensors"
-    config = json.loads(config_path.read_text())
-    latent_frames = (args.num_frames - 1) // 4 + 1
-    latent_height = args.height // 8
-    latent_width = args.width // 8
-
-    prompt_start = time.perf_counter()
-    prompt_embeds = encode_prompt(
-        model_root=args.model_root,
-        prompt=args.prompt,
-        max_sequence_length=args.max_sequence_length,
-        device_arg=args.torch_device,
-        dtype_arg=args.torch_dtype,
-    )
-    prompt_time = time.perf_counter() - prompt_start
-    freqs_cis = make_rotary_embeddings(
-        config,
-        latent_frames=latent_frames,
-        latent_height=latent_height,
-        latent_width=latent_width,
-    )
-
-    from fastvideo.mlx_runtime.fastwan import UnsupportedMLXQuantizationError
-
-    baseline_latents = None
-    rows = []
-    for mode in _parse_modes(args.modes):
-        print(f"=== MLX quant mode: {mode} ===")
-        mode_start = time.perf_counter()
-        try:
-            result = _run_one_mode(
-                mode=mode,
-                args=args,
-                config=config,
-                checkpoint_path=checkpoint_path,
-                config_path=config_path,
-                prompt_embeds=prompt_embeds,
-                freqs_cis=freqs_cis,
-            )
-        except UnsupportedMLXQuantizationError as exc:
-            print(f"skipping mode (unsupported by this MLX build): {exc}")
-            rows.append({"mode": mode, "status": "unsupported_by_mlx", "error": str(exc)})
-            continue
-        cleanup_mlx(mx)
-        latents = result["latents"]
-        if baseline_latents is None:
-            baseline_latents = latents
-        latent_path = args.output_dir / f"latents_{mode}.npy"
-        np.save(latent_path, latents)
-
-        decode_time = 0.0
-        decode_metrics = {}
-        output_path = None
-        if args.decode_backend != "none":
-            output_path = args.output_dir / f"video_{mode}_{args.decode_backend}_{args.height}x{args.width}x{args.num_frames}.mp4"
-            decode_metrics = _decode_with_metrics(args=args, latents=latents, output_path=output_path)
-            decode_time = cast(float, decode_metrics["decode_export_s"])
-
-        mode_total = time.perf_counter() - mode_start
-        mlx_denoise_peak_bytes = int(result["metrics"]["mlx_denoise_peak_bytes"])
-        mlx_active_bytes = int(result["metrics"]["mlx_active_after_denoise_bytes"])
-        metrics = {
-            "mode": mode,
-            "status": "ok",
-            "prompt_encode_shared_s": prompt_time,
-            "height": args.height,
-            "width": args.width,
-            "num_frames": args.num_frames,
-            "decode_backend": args.decode_backend,
-            "decode_export_s": decode_time,
-            "mode_total_excluding_shared_prompt_s": mode_total,
-            "mode_total_including_shared_prompt_s": mode_total + prompt_time,
-            "latents_path": str(latent_path),
-            "output_path": str(output_path) if output_path else None,
-            "mlx_denoise_peak_gib": mlx_denoise_peak_bytes / (1024**3),
-            "mlx_active_after_denoise_gib": mlx_active_bytes / (1024**3),
-            "mlx_dit_peak_under_16gb": mlx_denoise_peak_bytes < 16 * 1024**3,
-            "mlx_dit_active_under_16gb": mlx_active_bytes < 16 * 1024**3,
-            "mac_16gb_status": (
-                "dit_memory_fits_16gb_measured_decode_separately"
-                if mlx_denoise_peak_bytes < 16 * 1024**3 else "dit_memory_exceeds_16gb"
-            ),
-            **result["metrics"],
-            **decode_metrics,
-            **_latent_delta_metrics(latents, baseline_latents),
-        }
-        rows.append(metrics)
-        print(json.dumps(metrics, indent=2))
-
-    metrics_path = args.output_dir / "metrics.json"
-    metrics_path.write_text(json.dumps(rows, indent=2))
-    print(f"Wrote benchmark metrics to: {metrics_path}")
-
-
-if __name__ == "__main__":
-    main()
@@ -1,104 +0,0 @@
-"""Compare generated MP4s against a reference MP4 with simple pixel metrics."""
-
-from __future__ import annotations
-
-import argparse
-import json
-from pathlib import Path
-
-import numpy as np
-
-
-def _read_video(path: Path) -> np.ndarray:
-    """
-    Read all frames from a video file as an RGB NumPy array.
-
-    Parameters:
-        path (Path): Path to the video file.
-
-    Returns:
-        np.ndarray: Video frames stacked along the first axis.
-
-    Raises:
-        ValueError: If the video contains no readable frames.
-    """
-    import cv2
-
-    cap = cv2.VideoCapture(str(path))
-    frames = []
-    try:
-        while True:
-            ok, frame_bgr = cap.read()
-            if not ok:
-                break
-            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
-            frames.append(frame_rgb)
-    finally:
-        cap.release()
-    if not frames:
-        raise ValueError(f"No frames read from {path}")
-    return np.stack(frames, axis=0)
-
-
-def _metrics(candidate: np.ndarray, reference: np.ndarray) -> dict[str, float | int | list[int]]:
-    """
-    Compute pixel-level comparison metrics between candidate and reference video frames.
-
-    Parameters:
-        candidate (np.ndarray): Candidate video frames in frame, height, width, and channel order.
-        reference (np.ndarray): Reference video frames with the same shape as the candidate.
-
-    Returns:
-        dict[str, float | int | list[int]]: Frame dimensions and pixel comparison metrics, including MSE, MAE, maximum absolute difference, and PSNR in decibels.
-
-    Raises:
-        ValueError: If the candidate and reference arrays have different shapes.
-    """
-    if candidate.shape != reference.shape:
-        raise ValueError(f"Shape mismatch: candidate={candidate.shape}, reference={reference.shape}")
-    candidate_f = candidate.astype(np.float32)
-    reference_f = reference.astype(np.float32)
-    diff = candidate_f - reference_f
-    mse = float(np.mean(np.square(diff)))
-    mae = float(np.mean(np.abs(diff)))
-    max_abs = float(np.max(np.abs(diff)))
-    psnr = float(20.0 * np.log10(255.0 / np.sqrt(mse))) if mse > 0 else float("inf")
-    return {
-        "frames": int(candidate.shape[0]),
-        "height": int(candidate.shape[1]),
-        "width": int(candidate.shape[2]),
-        "channels": int(candidate.shape[3]),
-        "mse_vs_reference": mse,
-        "mae_vs_reference": mae,
-        "max_abs_vs_reference": max_abs,
-        "psnr_db_vs_reference": psnr,
-    }
-
-
-def main() -> None:
-    """Compare candidate MP4 videos with a reference and write pixel-level metrics to a JSON file."""
-    parser = argparse.ArgumentParser(description="Compare MP4s against a reference MP4.")
-    parser.add_argument("--reference", type=Path, required=True)
-    parser.add_argument("--candidates", type=Path, nargs="+", required=True)
-    parser.add_argument("--metrics-json", type=Path, required=True)
-    args = parser.parse_args()
-
-    reference = _read_video(args.reference)
-    rows = []
-    for candidate_path in args.candidates:
-        candidate = _read_video(candidate_path)
-        row = {
-            "reference_path": str(args.reference),
-            "candidate_path": str(candidate_path),
-            **_metrics(candidate, reference),
-        }
-        rows.append(row)
-        print(json.dumps(row, indent=2))
-
-    args.metrics_json.parent.mkdir(parents=True, exist_ok=True)
-    args.metrics_json.write_text(json.dumps(rows, indent=2))
-    print(f"Wrote video quality metrics to: {args.metrics_json}")
-
-
-if __name__ == "__main__":
-    main()
@@ -318,38 +318,6 @@ if(BUILD_CXX_KERNELS)

    # Combined FastVideo Extension
    # Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace
-    # ---------------------------------------------------------------------------
-    # VSA block-sparse attention forward, Blackwell (sm_100a) only.
-    #
-    # NOTE the "a" suffix: -arch=sm_100a is NOT enough -- it emits a plain sm_100 target and
-    # ptxas rejects every tcgen05 / setmaxnreg instruction. The explicit gencode spelling
-    # below is required, and matches the 10.0a entry in TORCH_CUDA_ARCH_LIST.
-    #
-    # Built for 64-token sparse blocks -- FastVideo's default (4,4,4) tiling, so no
-    # tile-size change and no top-k granularity change is needed. VSA_BHSD
-    # selects [B, H, S, D]. Both are compile-time; the Python is_supported() checks incoming
-    # tensors against them so callers fall back to Triton rather than getting a wrong answer.
-    # Read the ENVIRONMENT as well as the cache variable. When TORCH_CUDA_ARCH_LIST is
-    # exported (build.sh, and `pip install` with it set) the branch above only prints it --
-    # the cmake variable stays empty, so testing that alone silently skips the kernel and
-    # leaves a build that succeeds with the op missing.
-    set(ENABLE_VSA_SM100A OFF)
-    set(_VSA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}")
-    if(NOT _VSA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST})
-        set(_VSA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}")
-    endif()
-    if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)")
-        set(ENABLE_VSA_SM100A ON)
-    endif()
-    if(ENABLE_VSA_SM100A)
-        message(STATUS "fastvideo-kernel: building block_sparse_sm100a (Blackwell, 64- and 128-token blocks)")
-        list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_sm100a.cu
-                                      csrc/attention/block_sparse_blk128_sm100a.cu)
-        set_source_files_properties(csrc/attention/block_sparse_sm100a.cu
-                                    csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES
-            COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-DVSA_BHSD=true")
-    endif()
-
    Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI
        ${EXTENSION_SOURCES}
    )
@@ -365,14 +333,10 @@ if(BUILD_CXX_KERNELS)

    # Build compile definitions list
    set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops)
-    if(ENABLE_VSA_SM100A)
-        list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A)
-    endif()
    if(ENABLE_TK_KERNELS)
        list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE)
    endif()

-
    target_compile_definitions(fastvideo_kernel_ops PRIVATE ${COMPILE_DEFS})

    target_compile_options(fastvideo_kernel_ops PRIVATE
@@ -17,7 +17,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
 | Kernels | Where | Used when |
 |---|---|---|
 | Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils | `python/fastvideo_kernel/triton_kernels/` | automatic fallback when the matching C++ op is absent (`ops.py`, `turbodiffusion_ops.py`) |
-| FA4 CuTe-DSL block-sparse forward/backward (VSA-128/256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
+| FA4 CuTe-DSL block-sparse (VSA-256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
 | VMoBA `moba_attn_varlen` | `vmoba.py` | wraps flash-attn varlen |

 ## What gets built where, and when
@@ -64,49 +64,42 @@ cd fastvideo-kernel
 ./build.sh --rocm
 ```

-### Optional: FA4 CuTe block-sparse backend (VSA-128/256 fastpath)
+### Optional: FA4 CuTe block-sparse backend

-The VSA-128/256 fastpaths (tile volume 128 or 256, on NVIDIA Blackwell / sm_100) route to the
-FlashAttention-4 CuTe-DSL block-sparse kernel exposed as `flash_attn.cute`. This is
-an **optional** dependency: it is imported lazily, and `video_sparse_attn`
-transparently falls back to the Triton backend when it is absent (so the package is
-fully usable without it).
+The NVIDIA Blackwell / sm_100 VSA fastpath routes to FastVideo's in-tree
+FlashAttention-4 CuTe-DSL source fork, exposed as `flash_attn.cute`. This is an
+**optional** dependency: it is imported lazily, and `video_sparse_attn`
+transparently falls back to the Triton backend when it is absent (so the package
+is fully usable without it).

-The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`
-and the public/private forward-backward bridges in `flash_attn.cute.interface`) are provided upstream by
-[Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to
-commit `14c377950125c70b7a9dabf9c561fca53715ac7d`, the revision FastVideo pins as
-the `flash-attn-4` source in the repo-root `pyproject.toml`; other revisions may
-have incompatible block-sparse forward/backward interfaces.
-
-Install it under its distribution name so its own runtime stack resolves with it.
-Do **not** pre-install `nvidia-cutlass-dsl` by hand: this revision pins
-`nvidia-cutlass-dsl==4.6.0.dev0` exactly, and a hand-installed 4.5.x floor either
-gets silently upgraded or, if something else holds it back, leaves the CuTe
-kernels broken.
+The fork lives in `fastvideo-kernel/fa4`, retains the `flash-attn-4` distribution
+name, and records its exact upstream revision and refresh procedure in
+`fa4/UPSTREAM.md`. Install it editable while developing the kernels:

 ```bash
-pip install torchvision
-pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@14c377950125c70b7a9dabf9c561fca53715ac7d#subdirectory=flash_attn/cute"
+uv pip install --no-deps --editable fastvideo-kernel/fa4
 ```

-That resolves `nvidia-cutlass-dsl` to 4.6.0.dev0 and `quack-kernels` to 0.5.3, a
-combination this revision works with. A mismatched CuTe DSL only surfaces when the
-kernel JIT-compiles, so the error points at CuTe internals rather than at the
-install:
+The production VSA-256 entrypoint keeps its existing logical Q256/KV256 mask.
+When the CuTe backend is enabled, an optional shape override expands that mask
+to the finer physical schedules without changing selected token pairs:

-| Error on first VSA-128/256 CuTe call | Cause |
-|---|---|
-| `TypeError: fmax() missing 1 required positional argument: 'b'` | `nvidia-cutlass-dsl` 4.5.x |
-| `AttributeError: module 'cutlass.cute.core' has no attribute 'ThrMma'` | `quack-kernels` older than 0.5.1 |
-| `ImportError: cannot import name 'alloc_reserved_mbarrier'` | `quack-kernels` 0.6.2 or newer |
+```bash
+FASTVIDEO_VSA_CUTEDSL=1 FASTVIDEO_VSA_FA4_BLOCK_SHAPE=128x64 ...
+# At the original Q256 API boundary, logical Q64 children are safely paired
+# onto the optimized physical Q128/KV64 schedule:
+FASTVIDEO_VSA_CUTEDSL=1 FASTVIDEO_VSA_FA4_BLOCK_SHAPE=64x64 ...
+```

-An environment whose `flash_attn.cute` came from a prebuilt flash-attn wheel rather
-than from this pin hits the first row; that is what the overlay step in
-`docker/Dockerfile` works around.
+The default `FASTVIDEO_VSA_FA4_BLOCK_SHAPE=256x256` preserves the historical
+route (its logical KV256 edges are implemented with physical KV128 tiles).
+Direct callers with arbitrary, distinct Q64 masks still use the native
+Q64/KV64 kernel. The Q128/KV64 path enables score/probability double buffering
+by default; set `FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER=0` only to run its retained
+legacy schedule for comparison or rollback.

-The CuTe kernels JIT-compile on first use. Forward and backward are verified on
-Blackwell (sm_100) against `tests/test_vsa128_*.py` and `tests/test_vsa256_*.py`.
+The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against
+`tests/test_fa4_vsa_block_shapes.py` and `tests/test_vsa256_forward*.py`.

 ## Usage

@@ -158,19 +151,37 @@ statistics and bitwise-compatible `dV`.

 ### VSA (block-sparse) TFLOPs

-After building/installing `fastvideo-kernel`, run:
+For the GB200 FA4 comparison, install the in-tree source fork and run the
+issue-#4554-derived harness. Its default `exact256` mask mode preserves the
+same selected token pairs across Q×KV block shapes, reports sparse-aware
+TFLOP/s, MFU against the 2.5-PFLOP/s dense-BF16 GB200 peak, and efficiency
+relative to the raw FA4 256×256 path:
+
+```bash
+uv pip install --no-deps --editable fastvideo-kernel/fa4
+CUDA_VISIBLE_DEVICES=0 CUTE_DSL_ENABLE_TVM_FFI=1 \
+  uv run --no-sync python fastvideo-kernel/benchmarks/bench_vsa_blackwell.py \
+  --seq_lens 32768 --sparsities dense 90 \
+  --block_shapes 256x256 128x64 64x64
+```
+
+To exercise every GPU in a four-GB200 tray, launch one independent replica per
+device. Each rank writes its own suffixed JSON (for example,
+`/tmp/fa4_tray.rank0.json`):
+
+```bash
+CUTE_DSL_ENABLE_TVM_FFI=1 uv run --no-sync torchrun --standalone --nproc-per-node=4 \
+  fastvideo-kernel/benchmarks/bench_vsa_blackwell.py \
+  --seq_lens 32768 --sparsities dense 90 \
+  --block_shapes 256x256 128x64 64x64 --out /tmp/fa4_tray.json
+```
+
+The older generic benchmark below measures FastVideo's 64-token block-sparse
+wrapper; it does not exercise the FA4 fine-grained source fork:

 ```bash
 cd fastvideo-kernel
 python benchmarks/bench_vsa.py --batch_size 1 --num_heads 16 --head_dim 128 --q_seq_lens 49152 --topk 64
-
-# VSA-256 FA4 CuTe forward/backward on Blackwell
-python benchmarks/bench_vsa.py --block_size 256 --use_cute \
-  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 20
-
-# VSA-128 FA4 CuTe forward/backward on Blackwell
-python benchmarks/bench_vsa.py --block_size 128 --use_cute \
-  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 40
 ```

 ### TurboDiffusion Kernels
@@ -2,9 +2,8 @@
 """
 Benchmark VSA *wrapper* performance (forward + backward) and report TFLOPs.

-This script benchmarks the autograd-enabled wrappers:
-  - 64-token TK/Triton: fastvideo_kernel.block_sparse_attn.block_sparse_attn
-  - 128/256-token Triton/CuTe: fastvideo_kernel.block_sparse_attn_256
+This script benchmarks the autograd-enabled wrapper:
+  - fastvideo_kernel.block_sparse_attn.block_sparse_attn

 So measured time includes wrapper overhead (map->index conversion, dispatch) plus kernel time.
 """
@@ -24,6 +23,9 @@ try:
 except Exception as e:  # pragma: no cover
    raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e

+BLOCK_M = 64
+BLOCK_N = 64
+

 def set_seed(seed: int = 42) -> None:
    random.seed(seed)
@@ -39,11 +41,7 @@ def parse_arguments() -> argparse.Namespace:
    p.add_argument("--num_heads", type=int, default=12)
    p.add_argument("--head_dim", type=int, default=128, choices=[64, 128])
    p.add_argument("--topk", type=int, default=None, help="KV blocks per Q block (default: ~90%% sparsity)")
-    p.add_argument("--q_seq_lens",
-                   type=int,
-                   nargs="+",
-                   default=[49152],
-                   help="Q sequence lengths (must be divisible by --block_size)")
+    p.add_argument("--q_seq_lens", type=int, nargs="+", default=[49152], help="Q sequence lengths (must be /64)")
    p.add_argument("--kv_seq_lens",
                   type=int,
                   nargs="+",
@@ -53,13 +51,9 @@ def parse_arguments() -> argparse.Namespace:
    p.add_argument("--rep", type=int, default=20)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"])
-    p.add_argument("--block_size", type=int, default=64, choices=[64, 128, 256])
    p.add_argument("--force_triton",
                   action="store_true",
                   help="Force wrapper to use Triton path (if supported by shapes).")
-    p.add_argument("--use_cute",
-                   action="store_true",
-                   help="Use the optional FA4 CuTe forward/backward path (requires --block_size 128 or 256).")
    return p.parse_args()


@@ -90,38 +84,18 @@ def bench_ms(fn: Callable[[], object], warmup: int, rep: int) -> float:
    return do_bench(fn, warmup=warmup, rep=rep, quantiles=None)


-def _configure_backend(args: argparse.Namespace) -> None:
-    if args.use_cute and args.block_size not in (128, 256):
-        raise ValueError("--use_cute requires --block_size 128 or 256")
-    if args.use_cute and args.force_triton:
-        raise ValueError("--use_cute and --force_triton are mutually exclusive")
-
-    if args.force_triton:
-        os.environ.pop("FASTVIDEO_VSA_CUTEDSL", None)
-        os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
-    elif args.use_cute:
-        os.environ.pop("FASTVIDEO_VSA_TRITON", None)
-        os.environ.pop("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", None)
-        os.environ["FASTVIDEO_VSA_CUTEDSL"] = "1"
-
-
 def main() -> None:
    args = parse_arguments()
    set_seed(args.seed)
-    _configure_backend(args)

    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16

+    if args.force_triton:
+        os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
+
    from fastvideo_kernel.block_sparse_attn import block_sparse_attn
-    from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128, block_sparse_attn_256

    bs, h, d = args.batch_size, args.num_heads, args.head_dim
-    block_size = args.block_size
-    attention = {
-        64: block_sparse_attn,
-        128: block_sparse_attn_128,
-        256: block_sparse_attn_256,
-    }[block_size]
    kv_seq_lens = args.kv_seq_lens
    if kv_seq_lens is None:
        kv_seq_lens = args.q_seq_lens
@@ -131,22 +105,20 @@ def main() -> None:
    print("VSA Block-Sparse Attention Benchmark (WRAPPER)")
    print(f"device: {torch.cuda.get_device_name(0)}")
    print(f"batch={bs}, heads={h}, head_dim={d}, dtype={args.dtype}")
-    print(f"block_size={block_size}")
+    print(f"BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}")
    print("NOTE: timings include wrapper overhead (map->index + dispatch).")
-    if args.use_cute:
-        print("dispatch: FA4 CuTe")
-    elif args.force_triton:
+    if args.force_triton:
        print("dispatch: forced Triton (FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1)")
    else:
        print("dispatch: SM90 if available, else Triton")

    for q_len, kv_len in zip(args.q_seq_lens, kv_seq_lens):
-        if q_len % block_size != 0 or kv_len % block_size != 0:
-            print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by {block_size}")
+        if q_len % BLOCK_M != 0 or kv_len % BLOCK_N != 0:
+            print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by 64")
            continue

-        num_q_blocks = q_len // block_size
-        num_kv_blocks = kv_len // block_size
+        num_q_blocks = q_len // BLOCK_M
+        num_kv_blocks = kv_len // BLOCK_N
        topk = args.topk if args.topk is not None else max(1, num_kv_blocks // 10)
        topk = min(topk, num_kv_blocks)

@@ -157,11 +129,11 @@ def main() -> None:
        q, k, v = create_qkv(bs, h, q_len, kv_len, d, dtype)
        block_map = make_block_map(bs, h, num_q_blocks, num_kv_blocks, topk)

-        # Variable block sizes: default full logical blocks.
-        variable_block_sizes = torch.full((num_kv_blocks, ), block_size, dtype=torch.int32, device="cuda")
+        # Variable block sizes: default full blocks (64 tokens per KV block)
+        variable_block_sizes = torch.full((num_kv_blocks, ), BLOCK_N, dtype=torch.int32, device="cuda")

        def _fwd():
-            return attention(q, k, v, block_map, variable_block_sizes)
+            return block_sparse_attn(q, k, v, block_map, variable_block_sizes)

        fwd_ms = bench_ms(_fwd, warmup=args.warmup, rep=args.rep)

@@ -170,7 +142,7 @@ def main() -> None:
        q_ = q.detach().requires_grad_(True)
        k_ = k.detach().requires_grad_(True)
        v_ = v.detach().requires_grad_(True)
-        o_, _aux_ = attention(q_, k_, v_, block_map, variable_block_sizes)
+        o_, _aux_ = block_sparse_attn(q_, k_, v_, block_map, variable_block_sizes)
        og = torch.randn_like(o_)
        loss = (o_ * og).sum()

@@ -184,7 +156,7 @@ def main() -> None:
            rep=max(5, args.rep // 2),
        )

-        flops = flops_sparse_attention(bs, h, d, q_len, topk, block_size)
+        flops = flops_sparse_attention(bs, h, d, q_len, topk, BLOCK_N)
        fwd_tflops = flops / fwd_ms * 1e-12 * 1e3
        # Rough backward multiplier (attention backward typically ~2-3x forward)
        bwd_tflops = (2.5 * flops) / bwd_ms * 1e-12 * 1e3
@@ -0,0 +1,984 @@
+#!/usr/bin/env python3
+"""Block-sparse attention micro-benchmark: flashinfer ``vsa_blackwell`` and
+FA4 CuTe-DSL BSA, on identical block masks where the backends can represent
+them (GB200 / SM100).
+
+This harness originated in FlashInfer issue #4554 and its benchmark gist:
+  https://github.com/flashinfer-ai/flashinfer/issues/4554
+  https://gist.github.com/SolitaryThinker/90a1d1447929fc38dc509c1852e76532
+The imported baseline is pinned to immutable gist revision
+``e15ac9066f23ef3690e33e1cc1fdac45b4b9099f``.
+
+Measures achieved sparse-aware algorithmic TFLOP/s over a sequence-length x
+sparsity x (Q block, KV block) grid. Every successful arm is checked against
+the same fp32 masked-SDPA reference for that cell. FlashInfer's plan()/run()
+split is timed separately (VSA masks are data-dependent and change every
+layer x step, so mask-per-call deployments pay both).
+
+The default ``--mask_mode exact256`` samples one Q256/KV256 mask and expands
+it into each finer block shape. Thus all shapes select exactly the same token
+pairs and their output and relative-efficiency comparisons are apples-to-
+apples. ``--mask_mode native`` samples independently at each requested shape.
+
+Arms:
+  flashinfer  — BlockSparseAttentionWrapper(backend="vsa_blackwell"),
+                R=C=128. It runs only when the requested logical mask is
+                exactly representable by 128-token tiles.
+  cutedsl256  — original FastVideo 256x256 wrapper arm from the pinned gist;
+                used as a measured relative-efficiency reference.
+  fa4_wrapper — FastVideo's direct fine-grained BSHD wrapper, including
+                block-map conversion and variable-block-size mask plumbing.
+  fa4         — direct ``flash_attn.cute.interface._flash_attn_fwd`` call.
+                Sparse Q and KV block sizes are configured independently.
+
+For the stock 256x256 FA4 path, a logical KV256 edge expands into two physical
+KV128 edges, matching ``block_sparse_attn_256``. The 128x64 and 64x64 paths
+call FA4 with physical tiles (128, 64) and (64, 64), respectively. Mask-index
+construction and output allocation are outside the raw FA4 timing region.
+The direct ``fa4_wrapper`` 64x64 row is therefore the native Q64/KV64 kernel;
+it is not the public VSA-256 ``FASTVIDEO_VSA_FA4_BLOCK_SHAPE=64x64`` adapter,
+which safely pairs identical Q64 children onto physical Q128/KV64 and is
+measured by selecting that environment variable on the ``cutedsl256`` arm.
+
+The vendored FA4 source is a flat ``flash_attn.cute`` package. It can be used
+without installation by putting that directory on ``PYTHONPATH``; this script
+detects the flat layout and mounts it under the expected import namespace:
+
+  PYTHONPATH="$PWD/fastvideo-kernel/fa4${PYTHONPATH:+:$PYTHONPATH}" \\
+    uv run --no-sync python \\
+    fastvideo-kernel/benchmarks/bench_vsa_blackwell.py --quick
+
+MFU is reported separately from TFLOP/s. Its default denominator is the
+official 2,500 TFLOP/s dense BF16 tensor-core peak for one GB200 GPU and is
+explicitly overrideable with ``--peak_bf16_tflops``. Relative efficiency
+against measured 256x256 baselines does not depend on that nominal peak.
+
+Examples:
+  python bench_vsa_blackwell.py
+  python bench_vsa_blackwell.py --quick
+  python bench_vsa_blackwell.py --seq_lens 32768 --sparsities dense 90 --block_shapes 256x256 128x64 64x64
+  python bench_vsa_blackwell.py --mask_mode native --block_shapes 128x64 64x64
+  torchrun --standalone --nproc-per-node=4 bench_vsa_blackwell.py --quick --out /tmp/fa4_tray.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import importlib.util
+import inspect
+import json
+import math
+import os
+from pathlib import Path
+import random
+import subprocess
+import sys
+import time
+import traceback
+import types
+
+os.environ.setdefault("FASTVIDEO_VSA_CUTEDSL", "1")
+
+import numpy as np
+import torch
+
+FI_BLOCK = 128  # flashinfer vsa_blackwell R=C
+DEFAULT_BLOCK_SHAPES = ((256, 256), (128, 64), (64, 64))
+DEFAULT_GB200_BF16_DENSE_TFLOPS = 2500.0
+GB200_DATASHEET_URL = (
+    "https://dam-cdn.nvd.orangelogic.com/AssetLink/"
+    "y441155802qub41q118b2852i557jem5.pdf"
+)
+ISSUE_URL = "https://github.com/flashinfer-ai/flashinfer/issues/4554"
+GIST_URL = "https://gist.github.com/SolitaryThinker/90a1d1447929fc38dc509c1852e76532"
+GIST_REVISION = "e15ac9066f23ef3690e33e1cc1fdac45b4b9099f"
+
+KEEP_FRAC = {
+    "dense": 1.0,
+    "50": 0.5,
+    "60": 0.4,
+    "70": 0.3,
+    "75": 0.25,
+    "80": 0.2,
+    "87.5": 0.125,
+    "90": 0.1,
+}
+
+
+class UnsupportedConfig(ValueError):
+    """The backend cannot exactly represent the requested logical mask."""
+
+
+def parse_block_shape(value: str) -> tuple[int, int]:
+    parts = value.lower().replace(",", "x").split("x")
+    if len(parts) != 2:
+        raise argparse.ArgumentTypeError(f"expected QxKV, got {value!r}")
+    try:
+        q_block, kv_block = (int(part) for part in parts)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError(f"expected integer QxKV, got {value!r}") from exc
+    if q_block <= 0 or kv_block <= 0:
+        raise argparse.ArgumentTypeError("Q and KV block sizes must be positive")
+    return q_block, kv_block
+
+
+def flops_sparse_attention(
+    bs: int,
+    d: int,
+    selected_edges: int,
+    q_block: int,
+    kv_block: int,
+) -> float:
+    """QK^T + PV FLOPs over selected token pairs only.
+
+    ``selected_edges`` already includes the head dimension. Each selected
+    logical edge covers ``q_block * kv_block`` token pairs, and each of QK^T
+    and PV costs two FLOPs per head-dimension element.
+    """
+    return 4.0 * bs * d * selected_edges * q_block * kv_block
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def make_logical_mask(h: int, nq: int, nkv: int, keep: int, seed: int) -> torch.Tensor:
+    """[H, NQ, NKV] bool with one diagonal anchor and random other edges.
+
+    CPU generator so the pattern is device-independent and fully seeded.
+    """
+    if not 1 <= keep <= nkv:
+        raise ValueError(f"keep must be in [1, {nkv}], got {keep}")
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    scores = torch.rand(h, nq, nkv, generator=g)
+    q_rows = torch.arange(nq)
+    # For equal sequence lengths this is the KV block containing the first
+    # token of the Q block. It keeps every softmax row nonempty even when
+    # Q and KV block sizes differ.
+    diagonal_anchor = torch.div(q_rows * nkv, nq, rounding_mode="floor").clamp_max(nkv - 1)
+    scores[:, q_rows, diagonal_anchor] = 2.0  # rand() < 1, so this always wins topk
+    idx = torch.topk(scores, keep, dim=-1).indices
+    m = torch.zeros(h, nq, nkv, dtype=torch.bool)
+    m.scatter_(-1, idx, True)
+    assert int(m.sum(-1).min()) == keep and int(m.sum(-1).max()) == keep
+    return m.cuda()
+
+
+def expand_exact256_mask(base_mask: torch.Tensor, q_block: int, kv_block: int) -> torch.Tensor:
+    """Expand a Q256/KV256 mask without changing its selected token pairs."""
+    if 256 % q_block or 256 % kv_block:
+        raise ValueError(f"QxKV block shape {q_block}x{kv_block} must divide 256")
+    return base_mask.repeat_interleave(256 // q_block, dim=1).repeat_interleave(
+        256 // kv_block, dim=2
+    )
+
+
+def ref_masked_sdpa_fp32(q, k, v, logical_mask, q_block, kv_block, scale, qchunk=2048):
+    """fp32 SDPA with the expanded boolean token mask, chunked over q.
+
+    q,k,v: [1,S,H,D] bf16. Returns [S,H,D] fp32.
+    TF32 is disabled in main().
+    """
+    _, seq_len, heads, head_dim = q.shape
+    if seq_len % q_block or k.shape[1] % kv_block:
+        raise ValueError("reference requires sequence lengths divisible by their logical block sizes")
+    qchunk = max(q_block, qchunk // q_block * q_block)
+    q32 = q[0].permute(1, 0, 2).float()  # [H,S,D]
+    k32 = k[0].permute(1, 0, 2).float()
+    v32 = v[0].permute(1, 0, 2).float()
+    out = torch.empty(heads, seq_len, head_dim, dtype=torch.float32, device=q.device)
+    for i in range(0, seq_len, qchunk):
+        j = min(i + qchunk, seq_len)
+        scores = torch.matmul(q32[:, i:j], k32.transpose(-1, -2)).mul_(scale)
+        mask_rows = logical_mask[:, i // q_block:j // q_block]
+        mask_tokens = mask_rows.repeat_interleave(q_block, dim=1)
+        mask_tokens = mask_tokens.repeat_interleave(kv_block, dim=2)  # [H,c,S]
+        scores.masked_fill_(~mask_tokens, float("-inf"))
+        out[:, i:j] = torch.matmul(torch.softmax(scores, dim=-1), v32)
+        del scores, mask_tokens
+    return out.permute(1, 0, 2).contiguous()  # [S,H,D]
+
+
+def arm_cutedsl(q, k, v, logical_mask, q_block, kv_block):
+    """FA4-lineage CuTe-DSL 256 path. Returns (bench_fn, out [S,H,D], extra).
+
+    Timed as the full wrapper call (mask expansion + map->index included) —
+    that is its per-call hot path in a mask-per-call deployment.
+    """
+    if (q_block, kv_block) != (256, 256):
+        raise UnsupportedConfig("cutedsl256 is defined only for logical Q256/KV256")
+    _mount_flat_fa4_from_pythonpath()
+    from fastvideo_kernel import block_sparse_attn_256 as wrapper_module
+
+    block_sparse_attn_256_bshd = wrapper_module.block_sparse_attn_256_bshd
+
+    nkv = logical_mask.shape[-1]
+    vbs = torch.full((nkv, ), kv_block, dtype=torch.int32, device=q.device)
+    mask = logical_mask.unsqueeze(0)  # [1,H,NQ,NKV]
+
+    def fn():
+        return block_sparse_attn_256_bshd(q, k, v, mask, vbs)
+
+    out, _lse = fn()
+    interface = importlib.import_module("flash_attn.cute.interface")
+    return fn, out[0], {
+        "fa4_interface": str(Path(interface.__file__).resolve()),
+        "fastvideo_wrapper": str(Path(wrapper_module.__file__).resolve()),
+        "fa4_import_mode": _FA4_IMPORT_MODE,
+        "route_semantics": "public_vsa256",
+        "requested_vsa256_block_shape": os.environ.get(
+            "FASTVIDEO_VSA_FA4_BLOCK_SHAPE", "256x256"
+        ),
+    }
+
+
+def arm_fa4_wrapper(q, k, v, logical_mask, q_block, kv_block):
+    """FastVideo's direct fine-grained BSHD wrapper.
+
+    This arm deliberately includes boolean-map conversion and the VBS
+    ``mask_mod``/aux-tensor plumbing in the timed region. Fixed-length cells
+    use fully valid KV blocks. In particular, its 64x64 row stays on native
+    Q64/KV64 rather than entering the public VSA-256 coalescing adapter.
+    """
+    if (q_block, kv_block) not in ((128, 64), (64, 64)):
+        raise UnsupportedConfig("fa4_wrapper supports Q128/KV64 and Q64/KV64")
+
+    _mount_flat_fa4_from_pythonpath()
+    from fastvideo_kernel import block_sparse_attn_cute_fwd as wrapper_module
+
+    block_sparse_attn_cute_fwd_bshd = wrapper_module.block_sparse_attn_cute_fwd_bshd
+
+    nkv = logical_mask.shape[-1]
+    vbs = torch.full((nkv,), kv_block, dtype=torch.int32, device=q.device)
+    mask = logical_mask.unsqueeze(0)
+
+    def fn():
+        return block_sparse_attn_cute_fwd_bshd(q, k, v, mask, vbs)
+
+    out, _lse = fn()
+    interface = importlib.import_module("flash_attn.cute.interface")
+    return fn, out[0], {
+        "fa4_interface": str(Path(interface.__file__).resolve()),
+        "fastvideo_wrapper": str(Path(wrapper_module.__file__).resolve()),
+        "fa4_import_mode": _FA4_IMPORT_MODE,
+        "route_semantics": "direct_native_fine_grained",
+        "physical_tile": [q_block, kv_block],
+    }
+
+
+def arm_flashinfer(q, k, v, logical_mask, q_block, kv_block):
+    """flashinfer vsa_blackwell (blk128). Returns (bench_fn, out, extra).
+
+    run() timed alone (flashinfer's documented plan-once/run-many model);
+    steady-state plan() wall time recorded separately per cell.
+    """
+    if q_block % FI_BLOCK or kv_block % FI_BLOCK:
+        raise UnsupportedConfig(
+            "flashinfer vsa_blackwell R=C=128 cannot exactly represent this logical mask"
+        )
+    from flashinfer.sparse import BlockSparseAttentionWrapper
+
+    if "block_mask" not in inspect.signature(BlockSparseAttentionWrapper.plan).parameters:
+        raise UnsupportedConfig(
+            "installed FlashInfer lacks the per-head block_mask planning API used by "
+            "issue #4554 (the issue used flashinfer 0.6.16.post2)"
+        )
+
+    _, seq_len, heads, head_dim = q.shape
+    mask128 = logical_mask.repeat_interleave(q_block // FI_BLOCK, dim=1).repeat_interleave(
+        kv_block // FI_BLOCK, dim=2
+    )
+    workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
+    wrapper = BlockSparseAttentionWrapper(workspace, backend="vsa_blackwell")
+
+    def do_plan():
+        wrapper.plan(
+            None,
+            None,
+            M=seq_len,
+            N=seq_len,
+            R=FI_BLOCK,
+            C=FI_BLOCK,
+            num_qo_heads=heads,
+            num_kv_heads=heads,
+            head_dim=head_dim,
+            block_mask=mask128,
+            q_data_type=torch.bfloat16,
+            o_data_type=torch.bfloat16,
+        )
+
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    do_plan()
+    torch.cuda.synchronize()
+    plan_first = (time.perf_counter() - t0) * 1e3  # includes JIT on first cell
+    plan_steady = []
+    for _ in range(3):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        do_plan()
+        torch.cuda.synchronize()
+        plan_steady.append((time.perf_counter() - t0) * 1e3)
+    qn, kn, vn = q[0], k[0], v[0]  # NHD views of the same storage
+
+    def fn():
+        return wrapper.run(qn, kn, vn)
+
+    out = fn()
+    return fn, out, {
+        "plan_ms_first": round(plan_first, 2),
+        "plan_ms": round(min(plan_steady), 2),
+        "keep_alive": wrapper,
+    }
+
+
+_FA4_IMPORT_MODE = "normal"
+
+
+def _flat_fa4_pythonpath() -> Path | None:
+    """Find FastVideo's flat ``flash_attn.cute`` source on PYTHONPATH."""
+    for entry in os.environ.get("PYTHONPATH", "").split(os.pathsep):
+        if not entry:
+            continue
+        candidate = Path(entry).resolve()
+        if all((candidate / name).is_file() for name in ("UPSTREAM.md", "interface.py", "block_sparsity.py")):
+            return candidate
+    return None
+
+
+def _mount_flat_fa4_from_pythonpath() -> None:
+    """Mount the flat vendored tree as ``flash_attn.cute`` without installing.
+
+    A regular upstream checkout already has ``flash_attn/cute`` and needs no
+    special handling. FastVideo intentionally vendors only that subpackage,
+    so its source root itself is the package directory.
+    """
+    global _FA4_IMPORT_MODE
+    source = _flat_fa4_pythonpath()
+    if source is None:
+        return
+
+    try:
+        current = importlib.util.find_spec("flash_attn.cute.interface")
+    except (ImportError, AttributeError, ValueError):
+        current = None
+    if current is not None and current.origin is not None and Path(current.origin).resolve().parent == source:
+        _FA4_IMPORT_MODE = "local-pythonpath"
+        return
+
+    # Load the package under its real namespace so the vendored source's
+    # absolute ``flash_attn.cute.*`` imports resolve back to the same tree.
+    try:
+        import flash_attn
+    except ModuleNotFoundError as exc:
+        if exc.name != "flash_attn":
+            raise
+        # The fork is a standalone ``flash_attn.cute`` distribution. Create
+        # its namespace parent when FA2's top-level package is not installed.
+        flash_attn = types.ModuleType("flash_attn")
+        flash_attn.__path__ = []
+        sys.modules["flash_attn"] = flash_attn
+
+    for module_name in tuple(sys.modules):
+        if module_name == "flash_attn.cute" or module_name.startswith("flash_attn.cute."):
+            del sys.modules[module_name]
+    spec = importlib.util.spec_from_file_location(
+        "flash_attn.cute",
+        source / "__init__.py",
+        submodule_search_locations=[str(source)],
+    )
+    if spec is None or spec.loader is None:
+        raise ImportError(f"could not create an import spec for local FA4 source {source}")
+    package = importlib.util.module_from_spec(spec)
+    sys.modules["flash_attn.cute"] = package
+    setattr(flash_attn, "cute", package)
+    spec.loader.exec_module(package)
+    _FA4_IMPORT_MODE = "flat-pythonpath"
+
+
+def _load_fa4_symbols():
+    _mount_flat_fa4_from_pythonpath()
+    block_sparsity = importlib.import_module("flash_attn.cute.block_sparsity")
+    interface = importlib.import_module("flash_attn.cute.interface")
+    return block_sparsity.BlockSparseTensorsTorch, interface._flash_attn_fwd, interface
+
+
+def _fa4_physical_tiles(q_block: int, kv_block: int) -> tuple[int, int]:
+    """Map logical sparse blocks to FA4's physical MMA tiles.
+
+    KV256 is the historical logical VSA block and expands to KV128. Smaller
+    requested KV blocks remain physical. Q256 is two Q128 stages; Q128 and
+    Q64 select the single-stage paths being optimized.
+    """
+    tile_m = min(q_block, 128)
+    tile_n = min(kv_block, 128)
+    if q_block % tile_m or kv_block % tile_n:
+        raise UnsupportedConfig(
+            f"logical QxKV={q_block}x{kv_block} is not divisible by physical "
+            f"FA4 tile {tile_m}x{tile_n}"
+        )
+    return tile_m, tile_n
+
+
+def arm_fa4(q, k, v, logical_mask, q_block, kv_block):
+    """Direct FA4 block-sparse forward; mask preparation is not timed."""
+    BlockSparseTensorsTorch, _flash_attn_fwd, interface = _load_fa4_symbols()
+    tile_m, tile_n = _fa4_physical_tiles(q_block, kv_block)
+    kv_expansion = kv_block // tile_n
+    physical_mask = logical_mask.repeat_interleave(kv_expansion, dim=-1).contiguous()
+    heads, nq, nkv_physical = physical_mask.shape
+    selected_per_row = int(physical_mask[0, 0].sum().item())
+    if selected_per_row <= 0 or not bool((physical_mask.sum(dim=-1) == selected_per_row).all()):
+        raise ValueError("direct FA4 arm requires the same positive selected-block count per row")
+
+    physical_indices = torch.arange(nkv_physical, dtype=torch.int32, device=q.device)
+    physical_indices = physical_indices.view(1, 1, -1).expand_as(physical_mask)
+    full_idx = physical_indices.masked_select(physical_mask).view(1, heads, nq, selected_per_row)
+    full_cnt = torch.full((1, heads, nq), selected_per_row, dtype=torch.int32, device=q.device)
+    # Every selected logical block is fully valid in this fixed-length benchmark.
+    # Keep a compact empty partial-block representation to exercise FA4's full path.
+    mask_cnt = torch.zeros((1, heads, nq), dtype=torch.int32, device=q.device)
+    mask_idx = torch.zeros((1, heads, nq, 1), dtype=torch.int32, device=q.device)
+    sparse_tensors = BlockSparseTensorsTorch(
+        mask_block_cnt=mask_cnt,
+        mask_block_idx=mask_idx,
+        full_block_cnt=full_cnt,
+        full_block_idx=full_idx.contiguous(),
+        block_size=(q_block, tile_n),
+    )
+
+    batch, seq_len, _, _ = q.shape
+    out_buffer = torch.empty_like(q)
+    lse_buffer = torch.empty((batch, heads, seq_len), dtype=torch.float32, device=q.device)
+    scale = 1.0 / math.sqrt(q.shape[-1])
+
+    def fn():
+        return _flash_attn_fwd(
+            q,
+            k,
+            v,
+            out=out_buffer,
+            lse=lse_buffer,
+            softmax_scale=scale,
+            tile_mn=(tile_m, tile_n),
+            num_splits=1,
+            pack_gqa=False,
+            block_sparse_tensors=sparse_tensors,
+            causal=False,
+            return_lse=True,
+        )[:2]
+
+    out, _lse = fn()
+    return fn, out[0], {
+        "fa4_interface": str(Path(interface.__file__).resolve()),
+        "fa4_import_mode": _FA4_IMPORT_MODE,
+        "fa4_tile_m": tile_m,
+        "fa4_tile_n": tile_n,
+        "fa4_sparse_q_block": q_block,
+        "fa4_sparse_kv_block": tile_n,
+        "fa4_kv_expansion": kv_expansion,
+        "keep_alive": (sparse_tensors, out_buffer, lse_buffer),
+    }
+
+
+ARMS = {
+    "flashinfer": arm_flashinfer,
+    "cutedsl256": arm_cutedsl,
+    "fa4_wrapper": arm_fa4_wrapper,
+    "fa4": arm_fa4,
+}
+
+
+def detect_arms(requested):
+    if requested != ["auto"]:
+        return requested
+    arms = []
+    try:
+        from flashinfer.sparse import BlockSparseAttentionWrapper  # noqa: F401
+        if "block_mask" not in inspect.signature(BlockSparseAttentionWrapper.plan).parameters:
+            raise RuntimeError(
+                "per-head block_mask planning API unavailable; issue #4554 used "
+                "flashinfer 0.6.16.post2"
+            )
+        arms.append("flashinfer")
+    except Exception as exc:  # noqa: BLE001
+        print(f"# arm flashinfer unavailable: {exc!r}", flush=True)
+    fa4_available = False
+    try:
+        _load_fa4_symbols()
+        fa4_available = True
+    except Exception as exc:  # noqa: BLE001
+        print(f"# arm fa4 unavailable: {exc!r}", flush=True)
+    if fa4_available:
+        try:
+            from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_256_bshd  # noqa: F401
+
+            arms.append("cutedsl256")
+        except Exception as exc:  # noqa: BLE001
+            print(f"# arm cutedsl256 unavailable (optional): {exc!r}", flush=True)
+        try:
+            from fastvideo_kernel.block_sparse_attn_cute_fwd import (  # noqa: F401
+                block_sparse_attn_cute_fwd_bshd,
+            )
+
+            arms.append("fa4_wrapper")
+        except Exception as exc:  # noqa: BLE001
+            print(f"# arm fa4_wrapper unavailable (optional): {exc!r}", flush=True)
+        arms.append("fa4")
+    if not arms:
+        raise RuntimeError(
+            "no benchmark arm importable — need flashinfer-python, fastvideo_kernel, and/or FA4"
+        )
+    return arms
+
+
+def print_table(rows, arms):
+    header = "| seq | sparsity | Q blk | KV blk | keep/NKV |"
+    separator = "|---|---|---|---|---|"
+    for arm in arms:
+        header += f" {arm} ms | {arm} TFLOP/s | {arm} MFU % |"
+        separator += "---|---|---|"
+        if arm == "flashinfer":
+            header += " plan ms |"
+            separator += "---|"
+        if arm == "fa4":
+            header += " vs raw 256 % | vs cutedsl256 % |"
+            separator += "---|---|"
+        if arm == "fa4_wrapper":
+            header += " vs cutedsl256 % |"
+            separator += "---|"
+    print("\n" + header + "\n" + separator)
+    for row in rows:
+        line = (
+            f"| {row['seq_len']} | {row['sparsity']} | {row['q_block']} | {row['kv_block']} | "
+            f"{row['keep_kv_blocks']}/{row['nkv_blocks']} |"
+        )
+        for arm in arms:
+            cell = row.get(arm, {})
+            if cell.get("status") == "ok":
+                line += f" {cell['fwd_ms']:.3f} | {cell['tflops']:.1f} | {cell['mfu_pct']:.1f} |"
+                if arm == "flashinfer":
+                    line += f" {cell.get('plan_ms', float('nan')):.2f} |"
+                if arm == "fa4":
+                    raw_relative = cell.get("vs_fa4_256_pct", float("nan"))
+                    wrapper_relative = cell.get("vs_cutedsl256_pct", float("nan"))
+                    line += f" {raw_relative:.1f} | {wrapper_relative:.1f} |"
+                if arm == "fa4_wrapper":
+                    wrapper_relative = cell.get("vs_cutedsl256_pct", float("nan"))
+                    line += f" {wrapper_relative:.1f} |"
+            else:
+                status = cell.get("status", "FAILED")
+                line += f" {status} | — | — |"
+                if arm == "flashinfer":
+                    line += " — |"
+                if arm == "fa4":
+                    line += " — | — |"
+                if arm == "fa4_wrapper":
+                    line += " — |"
+        print(line)
+    print()
+
+
+def add_relative_efficiency(cell_rows):
+    """Attach denominator-free efficiency relative to the 256x256 baselines."""
+    baseline = next(
+        (row for row in cell_rows if (row["q_block"], row["kv_block"]) == (256, 256)),
+        None,
+    )
+    if baseline is None:
+        return
+    raw_tflops = baseline.get("fa4", {}).get("tflops")
+    wrapper_tflops = baseline.get("cutedsl256", {}).get("tflops")
+    for row in cell_rows:
+        fa4_result = row.get("fa4", {})
+        if fa4_result.get("status") == "ok":
+            if raw_tflops:
+                fa4_result["vs_fa4_256_pct"] = round(
+                    100.0 * fa4_result["tflops"] / raw_tflops,
+                    2,
+                )
+            if wrapper_tflops:
+                fa4_result["vs_cutedsl256_pct"] = round(
+                    100.0 * fa4_result["tflops"] / wrapper_tflops,
+                    2,
+                )
+        fine_wrapper_result = row.get("fa4_wrapper", {})
+        if fine_wrapper_result.get("status") == "ok" and wrapper_tflops:
+            fine_wrapper_result["vs_cutedsl256_pct"] = round(
+                100.0 * fine_wrapper_result["tflops"] / wrapper_tflops,
+                2,
+            )
+
+
+def main():
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size > 1:
+        if local_rank >= torch.cuda.device_count():
+            raise RuntimeError(
+                f"LOCAL_RANK={local_rank} exceeds visible CUDA device count={torch.cuda.device_count()}"
+            )
+        torch.cuda.set_device(local_rank)
+
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--seq_lens", type=int, nargs="+", default=[4096, 8192, 16384, 32768, 49152, 65536])
+    parser.add_argument("--sparsities", type=str, nargs="+", default=list(KEEP_FRAC), choices=list(KEEP_FRAC))
+    parser.add_argument("--arms", type=str, nargs="+", default=["auto"], choices=["auto", *ARMS])
+    parser.add_argument(
+        "--block_shapes",
+        type=parse_block_shape,
+        nargs="+",
+        default=list(DEFAULT_BLOCK_SHAPES),
+        metavar="QxKV",
+        help="independent logical QxKV sparse block pairs (default: 256x256 128x64 64x64)",
+    )
+    parser.add_argument(
+        "--mask_mode",
+        choices=("exact256", "native"),
+        default="exact256",
+        help=(
+            "exact256 expands one shared Q256/KV256 mask into every shape; "
+            "native samples each shape independently (default: exact256)"
+        ),
+    )
+    parser.add_argument(
+        "--peak_bf16_tflops",
+        type=float,
+        default=DEFAULT_GB200_BF16_DENSE_TFLOPS,
+        help=(
+            "per-GPU dense BF16 tensor-core peak used only for MFU %% "
+            f"(default: official GB200 {DEFAULT_GB200_BF16_DENSE_TFLOPS:.0f} TFLOP/s)"
+        ),
+    )
+    parser.add_argument(
+        "--quick",
+        action="store_true",
+        help="sanity grid: {8192,32768} x {dense,90} x requested block shapes",
+    )
+    parser.add_argument("--out", default="bsa_bench_results.json")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--rep", type=int, default=20)
+    parser.add_argument("--num_heads", type=int, default=12)
+    parser.add_argument("--head_dim", type=int, default=128)
+    args = parser.parse_args()
+    if args.quick:
+        args.seq_lens, args.sparsities = [8192, 32768], ["dense", "90"]
+    if "auto" in args.arms and args.arms != ["auto"]:
+        parser.error("--arms auto cannot be combined with explicit arms")
+    if args.peak_bf16_tflops <= 0:
+        parser.error("--peak_bf16_tflops must be positive")
+    args.block_shapes = list(dict.fromkeys(args.block_shapes))
+    for seq_len in args.seq_lens:
+        if args.mask_mode == "exact256" and seq_len % 256:
+            parser.error(f"exact256 mask mode requires sequence length divisible by 256; got {seq_len}")
+        for q_block, kv_block in args.block_shapes:
+            if seq_len % q_block or seq_len % kv_block:
+                parser.error(
+                    f"sequence length {seq_len} must be divisible by QxKV block shape "
+                    f"{q_block}x{kv_block}"
+                )
+            if args.mask_mode == "exact256" and (256 % q_block or 256 % kv_block):
+                parser.error(
+                    "exact256 mask mode requires both block sizes to divide 256; "
+                    f"got {q_block}x{kv_block}"
+                )
+
+    from triton.testing import do_bench
+
+    # The fp32 reference must be true fp32, not TF32.
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    arms = detect_arms(args.arms)
+    batch, heads, head_dim = 1, args.num_heads, args.head_dim
+    scale = 1.0 / math.sqrt(head_dim)
+
+    meta = {
+        "device": torch.cuda.get_device_name(),
+        "cuda_device_index": torch.cuda.current_device(),
+        "rank": rank,
+        "local_rank": local_rank,
+        "world_size": world_size,
+        "capability": list(torch.cuda.get_device_capability()),
+        "torch": torch.__version__,
+        "cuda": torch.version.cuda,
+        "seed": args.seed,
+        "warmup": args.warmup,
+        "rep": args.rep,
+        "batch": batch,
+        "heads": heads,
+        "head_dim": head_dim,
+        "dtype": "bfloat16",
+        "arms": arms,
+        "block_shapes": [list(shape) for shape in args.block_shapes],
+        "mask_mode": args.mask_mode,
+        "mask_base_block": 256 if args.mask_mode == "exact256" else None,
+        "peak_bf16_dense_tflops": args.peak_bf16_tflops,
+        "peak_note": (
+            "official single-GB200 dense BF16 tensor-core peak; override for clock-derived diagnostics"
+        ),
+        "peak_source": GB200_DATASHEET_URL,
+        "flop_formula": "4*bs*d*selected_logical_edges*q_block*kv_block",
+        "mfu_formula": "100*sparse_aware_algorithmic_tflops/peak_bf16_dense_tflops",
+        "correctness_gate": "finite and max_abs_vs_fp32 <= max(1e-2, 8*bf16_rounding_floor)",
+        "timing_scope": {
+            "flashinfer": "run only; plan reported separately",
+            "cutedsl256": "full FastVideo wrapper including mask expansion and map-to-index",
+            "fa4_wrapper": (
+                "FastVideo direct fine-grained BSHD wrapper including map-to-index and VBS mask "
+                "plumbing; 64x64 is native Q64/KV64, not the public VSA-256 coalescing adapter"
+            ),
+            "fa4": "raw _flash_attn_fwd dispatch; mask/index preparation and output allocation excluded",
+        },
+        "provenance": {
+            "issue": ISSUE_URL,
+            "gist": GIST_URL,
+            "gist_revision": GIST_REVISION,
+        },
+        "fa4_env": {
+            name: os.environ.get(name)
+            for name in (
+                "FASTVIDEO_FA4_VSA_DUAL_STREAM",
+                "FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER",
+                "FASTVIDEO_VSA_FA4_BLOCK_SHAPE",
+            )
+        },
+    }
+    try:
+        meta["driver"] = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], text=True
+        ).split()[0]
+    except Exception:  # noqa: BLE001
+        meta["driver"] = None
+    for module in ("flashinfer", "triton", "fastvideo_kernel"):
+        try:
+            meta[module] = __import__(module).__version__
+        except Exception:  # noqa: BLE001
+            pass
+    if any(arm in arms for arm in ("cutedsl256", "fa4_wrapper", "fa4")):
+        try:
+            _, _, fa4_interface = _load_fa4_symbols()
+            meta["fa4_interface"] = str(Path(fa4_interface.__file__).resolve())
+            meta["fa4_import_mode"] = _FA4_IMPORT_MODE
+        except Exception as exc:  # noqa: BLE001
+            meta["fa4_import_error"] = f"{type(exc).__name__}: {exc}"
+    print("META " + json.dumps(meta), flush=True)
+
+    rows = []
+    for seq_len in args.seq_lens:
+        # Same q,k,v for every sparsity cell of this seq len (isolates sparsity).
+        set_seed(args.seed)
+        q = torch.randn(batch, seq_len, heads, head_dim, dtype=torch.bfloat16, device="cuda")
+        k = torch.randn(batch, seq_len, heads, head_dim, dtype=torch.bfloat16, device="cuda")
+        v = torch.randn(batch, seq_len, heads, head_dim, dtype=torch.bfloat16, device="cuda")
+
+        for label in args.sparsities:
+            fraction = KEEP_FRAC[label]
+            base_mask_seed = args.seed + seq_len + int(fraction * 1000)
+            if args.mask_mode == "exact256":
+                base_nblocks_256 = seq_len // 256
+                base_keep_256 = (
+                    base_nblocks_256
+                    if label == "dense"
+                    else max(1, round(fraction * base_nblocks_256))
+                )
+                base_mask_256 = make_logical_mask(
+                    heads,
+                    base_nblocks_256,
+                    base_nblocks_256,
+                    base_keep_256,
+                    base_mask_seed,
+                )
+            else:
+                base_nblocks_256 = None
+                base_keep_256 = None
+                base_mask_256 = None
+            exact_ref = (
+                ref_masked_sdpa_fp32(
+                    q,
+                    k,
+                    v,
+                    base_mask_256,
+                    256,
+                    256,
+                    scale,
+                )
+                if base_mask_256 is not None
+                else None
+            )
+            cell_rows = []
+            fa4_shape_outputs = {}
+            for q_block, kv_block in args.block_shapes:
+                nq_blocks, nkv_blocks = seq_len // q_block, seq_len // kv_block
+                if args.mask_mode == "exact256":
+                    if base_mask_256 is None:  # pragma: no cover - guarded above
+                        raise AssertionError("exact256 mode requires a base mask")
+                    keep = base_keep_256 * (256 // kv_block)
+                    mask_seed = base_mask_seed
+                    logical_mask = expand_exact256_mask(base_mask_256, q_block, kv_block)
+                else:
+                    keep = nkv_blocks if label == "dense" else max(1, round(fraction * nkv_blocks))
+                    mask_seed = (
+                        base_mask_seed
+                        + q_block * 1_000_003
+                        + kv_block * 9_176
+                    )
+                    logical_mask = make_logical_mask(
+                        heads,
+                        nq_blocks,
+                        nkv_blocks,
+                        keep,
+                        mask_seed,
+                    )
+                selected_edges = heads * nq_blocks * keep
+                flops = flops_sparse_attention(
+                    batch,
+                    head_dim,
+                    selected_edges,
+                    q_block,
+                    kv_block,
+                )
+                row = {
+                    "seq_len": seq_len,
+                    "sparsity": label,
+                    "mask_mode": args.mask_mode,
+                    "base_nblocks_256": base_nblocks_256,
+                    "base_keep_256": base_keep_256,
+                    "base_mask_seed": base_mask_seed,
+                    "q_block": q_block,
+                    "kv_block": kv_block,
+                    "nq_blocks": nq_blocks,
+                    "nkv_blocks": nkv_blocks,
+                    "keep_kv_blocks": keep,
+                    "selected_logical_edges": selected_edges,
+                    "actual_sparsity": round(1.0 - keep / nkv_blocks, 4),
+                    "mask_seed": mask_seed,
+                    "flops": flops,
+                }
+
+                ref = (
+                    exact_ref
+                    if exact_ref is not None
+                    else ref_masked_sdpa_fp32(
+                        q,
+                        k,
+                        v,
+                        logical_mask,
+                        q_block,
+                        kv_block,
+                        scale,
+                    )
+                )
+                # Irreducible bf16 output-rounding floor: even a bit-exact kernel
+                # emitting bf16 cannot beat this vs the fp32 reference.
+                row["bf16_floor_max_abs"] = float(
+                    (ref - ref.to(torch.bfloat16).float()).abs().max()
+                )
+                row["correctness_max_abs_limit"] = max(
+                    1e-2,
+                    8.0 * row["bf16_floor_max_abs"],
+                )
+
+                outputs = {}
+                for name in arms:
+                    result = {}
+                    try:
+                        fn, out, extra = ARMS[name](
+                            q,
+                            k,
+                            v,
+                            logical_mask,
+                            q_block,
+                            kv_block,
+                        )
+                        torch.cuda.synchronize()
+                        result["max_abs_vs_fp32"] = float((out.float() - ref).abs().max())
+                        if not bool(torch.isfinite(out).all()):
+                            raise RuntimeError(f"{name} produced NaN or Inf")
+                        if result["max_abs_vs_fp32"] > row["correctness_max_abs_limit"]:
+                            raise RuntimeError(
+                                f"{name} max_abs_vs_fp32={result['max_abs_vs_fp32']:.6g} "
+                                f"exceeds limit={row['correctness_max_abs_limit']:.6g}"
+                            )
+                        outputs[name] = out
+                        elapsed_ms = float(
+                            do_bench(fn, warmup=args.warmup, rep=args.rep, quantiles=None)
+                        )
+                        achieved_tflops = flops / elapsed_ms * 1e-9
+                        result["fwd_ms"] = round(elapsed_ms, 4)
+                        result["tflops"] = round(achieved_tflops, 2)
+                        result["mfu_pct"] = round(
+                            100.0 * achieved_tflops / args.peak_bf16_tflops,
+                            2,
+                        )
+                        for key, value in extra.items():
+                            if key != "keep_alive":
+                                result[key] = value
+                        result["status"] = "ok"
+                        del fn, extra
+                    except UnsupportedConfig as exc:
+                        result["status"] = "SKIPPED"
+                        result["error"] = str(exc)
+                    except Exception as exc:  # noqa: BLE001
+                        result["status"] = "FAILED"
+                        result["error"] = f"{type(exc).__name__}: {exc}"[:500]
+                        traceback.print_exc()
+                    row[name] = result
+                    torch.cuda.empty_cache()
+
+                output_names = list(outputs)
+                if len(output_names) > 1:
+                    row["cross_max_abs"] = {}
+                    for lhs_idx, lhs in enumerate(output_names):
+                        for rhs in output_names[lhs_idx + 1:]:
+                            row["cross_max_abs"][f"{lhs}_vs_{rhs}"] = float(
+                                (outputs[lhs].float() - outputs[rhs].float()).abs().max()
+                            )
+                if args.mask_mode == "exact256" and "fa4" in outputs:
+                    fa4_shape_outputs[(q_block, kv_block)] = outputs["fa4"]
+                outputs.clear()
+                del ref, logical_mask
+                torch.cuda.empty_cache()
+                cell_rows.append(row)
+
+            if args.mask_mode == "exact256":
+                baseline_output = fa4_shape_outputs.get((256, 256))
+                if baseline_output is not None:
+                    for row in cell_rows:
+                        shape = (row["q_block"], row["kv_block"])
+                        target_output = fa4_shape_outputs.get(shape)
+                        if target_output is not None:
+                            row["fa4"]["max_abs_vs_fa4_256"] = float(
+                                (target_output.float() - baseline_output.float()).abs().max()
+                            )
+                fa4_shape_outputs.clear()
+            add_relative_efficiency(cell_rows)
+            for row in cell_rows:
+                print("ROW " + json.dumps(row), flush=True)
+            rows.extend(cell_rows)
+            del base_mask_256, exact_ref
+        del q, k, v
+        torch.cuda.empty_cache()
+
+    output_path = Path(args.out)
+    if world_size > 1:
+        output_path = output_path.with_name(
+            f"{output_path.stem}.rank{rank}{output_path.suffix}"
+        )
+    with output_path.open("w") as output_file:
+        json.dump({"meta": meta, "rows": rows}, output_file, indent=2)
+    print_table(rows, arms)
+    print(f"wrote {output_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA required")
+    sys.exit(main())
@@ -1,7 +0,0 @@
-// block_sparse_blk128_sm100a.cu -- the 128-token-block instantiation of the torch binding.
-//
-// Same source as block_sparse_sm100a.cu with VSA_BLK128 set: the kernel and launch land in
-// namespace vsa_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the
-// exported entry point becomes block_sparse_sm100a_blk128_fwd.
-#define VSA_BLK128 true
-#include "block_sparse_sm100a.cu"
@@ -1,201 +0,0 @@
-#ifndef BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH
-#define BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH
-
-// Launch surface for the sm_100a VSA block-sparse FMHA forward.
-//
-// Everything a caller needs: a POD argument struct, a predicate saying whether this build can
-// run those arguments, and one launch entry point. The benchmark in
-// block_sparse_bench_sm100a.cu and the torch binding both go through here, so there is
-// one tensormap construction and one launch configuration rather than two that can drift.
-//
-// Two compile-time knobs select the four builds:
-//   VSA_BLK128  false -> 64-token sparse blocks, true -> 128-token
-//   VSA_BHSD    false -> [token][head][dim] (BSHD), true -> [batch][head][token][dim] (BHSD)
-
-#include "block_sparse_kernel_sm100a.cuh"
-
-namespace VSA_NAMESPACE {
-
-struct BlockSparseVsaArgs {
-  const __nv_bfloat16* q;
-  const __nv_bfloat16* k;
-  const __nv_bfloat16* v;      // natural layout; only blk128 reads it (blk64 still needs v_t)
-  const __nv_bfloat16* v_t;    // unused: kept so the bench's V_T buffer still binds
-  __nv_bfloat16* o;
-  float* lse;                  // [batch, num_heads, seqlen] fp32, or nullptr
-
-  const int* q2k_idx;              // [batch*num_heads*num_blocks, max_kv] int32
-  const int* q2k_num;              // [batch*num_heads*num_blocks] int32
-  const int* variable_block_sizes; // [num_blocks] int32, valid tokens per block
-
-  int batch;
-  int num_heads;
-  int seqlen;
-  int head_dim;
-  int num_blocks;
-  int max_kv;
-  float sm_scale;
-};
-
-// cudaSuccess iff this build can run `a`. Deliberately conservative: the caller is expected
-// to fall back to its own implementation rather than get a wrong answer.
-__host__ inline cudaError_t block_sparse_supported(const BlockSparseVsaArgs& a) {
-  if (a.head_dim != HEAD_DIM) return cudaErrorInvalidValue;      // compile-time in the kernel
-  if (a.num_blocks % 2 != 0) return cudaErrorInvalidValue;       // a CTA owns an adjacent pair
-  if (a.seqlen != a.num_blocks * BLOCK) return cudaErrorInvalidValue;
-  if (a.max_kv < 1 || a.num_blocks < 1) return cudaErrorInvalidValue;
-  if (a.q == nullptr || a.k == nullptr || a.o == nullptr) return cudaErrorInvalidValue;
-  if (a.q2k_idx == nullptr || a.q2k_num == nullptr) return cudaErrorInvalidValue;
-  // FastVideo always supplies this; without it padded keys would be attended as real zeros.
-  if (a.variable_block_sizes == nullptr) return cudaErrorInvalidValue;
-  // V is read MN-major at BOTH block sizes now, so no pre-transposed V_T is ever needed.
-  if (a.v == nullptr) return cudaErrorInvalidValue;
-  return cudaSuccess;
-}
-
-__host__ inline cudaError_t launch_block_sparse_sm100a(const BlockSparseVsaArgs& a,
-                                                           cudaStream_t stream) {
-  const cudaError_t sup = block_sparse_supported(a);
-  if (sup != cudaSuccess) return sup;
-
-  const int B = a.batch, H = a.num_heads, S = a.seqlen, hd = a.head_dim;
-  const int num_blocks = a.num_blocks, max_kv = a.max_kv;
-  const long tq = (long)B * S;
-  const int packed_mtiles_per_seq = num_blocks / 2;
-  const int total_work = B * H * packed_mtiles_per_seq;
-  constexpr bool BHSD = VSA_BHSD;
-
-  CUtensorMap tq_, tk_, tvt_, tv_, to_;
-  {
-    uint64_t gd[4] = { (uint64_t)SUB_COLS_BF16, BHSD ? (uint64_t)((long)B * H) : (uint64_t)H,
-                       BHSD ? (uint64_t)S : (uint64_t)tq, (uint64_t)Q_SUBTILES };
-    uint64_t gs[3] = { BHSD ? (uint64_t)((long)S * hd) * 2u : (uint64_t)hd * 2u,
-                       BHSD ? (uint64_t)hd * 2u : (uint64_t)((long)H * hd) * 2u,
-                       (uint64_t)SUB_COLS_BF16 * 2u };
-    uint32_t bd[4] = { (uint32_t)SUB_COLS_BF16, 1u, (uint32_t)M_TILE, (uint32_t)Q_SUBTILES };
-    uint32_t es[4] = { 1u, 1u, 1u, 1u };
-    if (cuTensorMapEncodeTiled(&tq_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4,
-                               const_cast<__nv_bfloat16*>(a.q), gd, gs, bd, es,
-                               CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
-                               CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
-                               CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
-      return cudaErrorInvalidValue;
-    if (cuTensorMapEncodeTiled(&to_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, a.o, gd, gs, bd, es,
-                               CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
-                               CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
-                               CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
-      return cudaErrorInvalidValue;
-  }
-  {
-    uint64_t gd[4] = { (uint64_t)SUB_COLS_BF16,
-                       BHSD ? (uint64_t)S : (uint64_t)tq,
-                       BHSD ? (uint64_t)(hd / SUB_COLS_BF16)
-                            : (uint64_t)((long)H * hd / SUB_COLS_BF16),
-                       (uint64_t)((long)B * H) };
-    uint64_t gs[3] = { BHSD ? (uint64_t)hd * 2u : (uint64_t)((long)H * hd) * 2u,
-                       (uint64_t)SUB_COLS_BF16 * 2u,
-                       (uint64_t)((long)S * hd) * 2u };
-    uint32_t bd[4] = { (uint32_t)SUB_COLS_BF16, (uint32_t)BLOCK,
-                       BLK128 ? (uint32_t)K_SUBTILES : 1u, 1u };
-    uint32_t es[4] = { 1u, 1u, 1u, 1u };
-    if (cuTensorMapEncodeTiled(&tk_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, BHSD ? 4 : 3,
-                               const_cast<__nv_bfloat16*>(a.k), gd, gs, bd, es,
-                               CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
-                               CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
-                               CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
-      return cudaErrorInvalidValue;
-    // V map is byte-for-byte the K map over a.v: MN-major V needs no transpose (blk128).
-    const __nv_bfloat16* vbase = a.v ? a.v : a.k;
-    if (cuTensorMapEncodeTiled(&tv_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, BHSD ? 4 : 3,
-                               const_cast<__nv_bfloat16*>(vbase), gd, gs, bd, es,
-                               CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
-                               CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
-                               CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
-      return cudaErrorInvalidValue;
-  }
-  // V_T map: blk64 only. Unused at blk128 but must still be a valid tensormap to pass by value.
-  {
-    const __nv_bfloat16* vt = a.v_t ? a.v_t : a.k;
-    if constexpr (BLK128) {
-      uint64_t gd[3] = { (uint64_t)SUB_COLS_BF16, (uint64_t)((long)H * hd),
-                         (uint64_t)((long)tq / SUB_COLS_BF16) };
-      uint64_t gs[2] = { (uint64_t)tq * 2u, (uint64_t)SUB_COLS_BF16 * 2u };
-      uint32_t bd[3] = { (uint32_t)SUB_COLS_BF16, (uint32_t)hd, (uint32_t)V_SUBTILES };
-      uint32_t es[3] = { 1u, 1u, 1u };
-      if (cuTensorMapEncodeTiled(&tvt_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3,
-                                 const_cast<__nv_bfloat16*>(vt), gd, gs, bd, es,
-                                 CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
-                                 CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
-                                 CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) != CUDA_SUCCESS)
-        return cudaErrorInvalidValue;
-    } else {
-      if (make_tma_2d_tiled(&tvt_, const_cast<__nv_bfloat16*>(vt), (long)H * hd, (int)tq, hd,
-                            SUB_COLS_BF16, 2, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
-                            CU_TENSOR_MAP_SWIZZLE_128B) != cudaSuccess)
-        return cudaErrorInvalidValue;
-    }
-  }
-
-  const size_t smem =
-        (size_t)2 * Q_TILE_BYTES + NUM_KV_STAGES * KV_RING_SLOT_BYTES
-      + (size_t)2 * M_TILE * HEAD_DIM * sizeof(__nv_bfloat16)
-      + (2 * NUM_KV_STAGES + 22) * 8
-      + (size_t)CLC_STAGES * (2 * 8 + 16) + 16
-      + 8
-      + (size_t)2 * STAT_REGIONS * STATS * sizeof(float)
-      + 256;
-
-#ifndef VSA_NAMED_BAR
-#define VSA_NAMED_BAR false
-#endif
-#ifndef VSA_THROTTLE
-#define VSA_THROTTLE false
-#endif
-#ifndef VSA_USE_CLC
-#define VSA_USE_CLC true
-#endif
-  constexpr bool FULL_NAMED_BAR = VSA_NAMED_BAR, EX2_EMU = true, SPLIT_P = true,
-                 SOFTMAX_THROTTLE = VSA_THROTTLE, USE_CLC = VSA_USE_CLC,
-                 Q_RASTER = true, MHA = true;
-  auto kfn = &fmha_context_bf16_gen_kernel<32, FULL_NAMED_BAR, EX2_EMU, SPLIT_P,
-                                           SOFTMAX_THROTTLE, USE_CLC, Q_RASTER, MHA,
-                                           /*RESCALE_THRESHOLD=*/8, /*BHSD=*/VSA_BHSD>;
-  cudaError_t e = cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
-  if (e != cudaSuccess) return e;
-
-  const unsigned long long magic0 = make_magic((unsigned)(H * packed_mtiles_per_seq));
-  const unsigned long long magic1 = make_magic((unsigned)H);
-  const unsigned long long magic2 = make_magic((unsigned)packed_mtiles_per_seq);
-  const float scale_log2 = a.sm_scale * (float)M_LOG2E;
-
-  int numSM = 0;
-  e = cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, 0);
-  if (e != cudaSuccess) return e;
-  const int num_ctas = USE_CLC ? total_work : (total_work < numSM ? total_work : numSM);
-  dim3 grid(num_ctas, 1, 1), block(N_WARPS * 32, 1, 1);
-
-  if (USE_CLC) {
-    cudaLaunchConfig_t cfg = {};
-    cfg.gridDim = grid; cfg.blockDim = block; cfg.dynamicSmemBytes = smem; cfg.stream = stream;
-    cudaLaunchAttribute cfgAttr[1];
-    cfgAttr[0].id = cudaLaunchAttributeClusterDimension;
-    cfgAttr[0].val.clusterDim.x = 1; cfgAttr[0].val.clusterDim.y = 1;
-    cfgAttr[0].val.clusterDim.z = 1;
-    cfg.attrs = cfgAttr; cfg.numAttrs = 1;
-    return cudaLaunchKernelEx(&cfg, kfn, tq_, tk_, tvt_, tv_, to_, S, H, scale_log2, B,
-                              num_blocks, packed_mtiles_per_seq, max_kv, magic0, magic1, magic2,
-                              a.q2k_idx, a.q2k_num, a.variable_block_sizes, a.lse);
-  }
-  kfn<<<grid, block, smem, stream>>>(tq_, tk_, tvt_, tv_, to_, S, H, scale_log2, B, num_blocks,
-                                     packed_mtiles_per_seq, max_kv, magic0, magic1, magic2,
-                                     a.q2k_idx, a.q2k_num, a.variable_block_sizes, a.lse);
-  return cudaGetLastError();
-}
-
-}  // namespace VSA_NAMESPACE
-
-// Callers (the bench, the torch binding) keep using unqualified names; each translation unit
-// only ever sees the one configuration its VSA_BLK128 selected.
-using namespace VSA_NAMESPACE;
-
-#endif  // BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH
@@ -1,114 +0,0 @@
-// block_sparse_sm100a.cu -- torch binding for the sm_100a VSA block-sparse FMHA forward.
-//
-// Forward only: returns (out, lse) so FastVideo's existing Triton backward keeps working
-// unchanged. lse is exactly the M tensor triton_block_sparse_attn_forward writes --
-// max(qk * qk_scale) + log2(l), [B, H, S] fp32 -- which is what lets
-// block_sparse_attn_backward_triton run against our forward untouched.
-//
-// The build is fixed at compile time by two flags, so one extension carries one configuration:
-//   VSA_BLK128  false -> 64-token sparse blocks, true -> 128-token
-//   VSA_BHSD    false -> [B, S, H, D], true -> [B, H, S, D]
-#include <torch/extension.h>
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-
-#include "block_sparse_launch_sm100a.cuh"
-
-namespace {
-
-void check_qkv(const torch::Tensor& t, const char* name, int64_t B, int64_t H, int64_t S,
-               int64_t D) {
-  TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
-  TORCH_CHECK(t.scalar_type() == at::kBFloat16, name, " must be bfloat16, got ", t.scalar_type());
-  TORCH_CHECK(t.dim() == 4, name, " must be 4-D, got ", t.dim(), " dims");
-  TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
-  if (VSA_BHSD) {
-    TORCH_CHECK(t.size(0) == B && t.size(1) == H && t.size(2) == S && t.size(3) == D, name,
-                " has shape ", t.sizes(), ", expected [", B, ",", H, ",", S, ",", D, "]");
-  } else {
-    TORCH_CHECK(t.size(0) == B && t.size(1) == S && t.size(2) == H && t.size(3) == D, name,
-                " has shape ", t.sizes(), ", expected [", B, ",", S, ",", H, ",", D, "]");
-  }
-}
-
-void check_index(const torch::Tensor& t, const char* name) {
-  TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
-  TORCH_CHECK(t.scalar_type() == at::kInt, name, " must be int32, got ", t.scalar_type());
-  TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
-}
-
-}  // namespace
-
-// The exported symbol carries the block size: block_sparse_sm100a_fwd is the 64-token build,
-// block_sparse_sm100a_blk128_fwd the 128-token one (block_sparse_blk128_sm100a.cu re-includes
-// this file with VSA_BLK128 set). The python backend picks by the metadata's block size.
-#if VSA_BLK128
-#define BLOCK_SPARSE_SM100A_FWD block_sparse_sm100a_blk128_fwd
-#else
-#define BLOCK_SPARSE_SM100A_FWD block_sparse_sm100a_fwd
-#endif
-
-// Returns {out} or {out, lse}. Layout of out matches the inputs.
-std::vector<torch::Tensor> BLOCK_SPARSE_SM100A_FWD(torch::Tensor q, torch::Tensor k,
-                                                       torch::Tensor v,
-                                                       c10::optional<torch::Tensor> v_t,
-                                                       torch::Tensor q2k_idx,
-                                                       torch::Tensor q2k_num,
-                                                       torch::Tensor variable_block_sizes,
-                                                       double sm_scale, bool need_lse) {
-  const at::cuda::OptionalCUDAGuard guard(device_of(q));
-
-  const int64_t B = q.size(0);
-  const int64_t H = VSA_BHSD ? q.size(1) : q.size(2);
-  const int64_t S = VSA_BHSD ? q.size(2) : q.size(1);
-  const int64_t D = q.size(3);
-
-  check_qkv(q, "q", B, H, S, D);
-  check_qkv(k, "k", B, H, S, D);
-  check_qkv(v, "v", B, H, S, D);
-  check_index(q2k_idx, "q2k_idx");
-  check_index(q2k_num, "q2k_num");
-  check_index(variable_block_sizes, "variable_block_sizes");
-
-  const int64_t num_blocks = variable_block_sizes.numel();
-  const int64_t max_kv = q2k_idx.size(-1);
-  TORCH_CHECK(S == num_blocks * BLOCK, "seqlen ", S, " must equal num_blocks (", num_blocks,
-              ") * ", BLOCK, "; FastVideo pads the sequence up to whole blocks");
-
-  auto out = torch::empty_like(q);
-  torch::Tensor lse;
-  if (need_lse) lse = torch::empty({B, H, S}, q.options().dtype(torch::kFloat32));
-
-  BlockSparseVsaArgs a{};
-  a.q = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr());
-  a.k = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr());
-  a.v = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr());
-  a.v_t = v_t.has_value() ? reinterpret_cast<const __nv_bfloat16*>(v_t->data_ptr()) : nullptr;
-  a.o = reinterpret_cast<__nv_bfloat16*>(out.data_ptr());
-  a.lse = need_lse ? lse.data_ptr<float>() : nullptr;
-  a.q2k_idx = q2k_idx.data_ptr<int>();
-  a.q2k_num = q2k_num.data_ptr<int>();
-  a.variable_block_sizes = variable_block_sizes.data_ptr<int>();
-  a.batch = (int)B;
-  a.num_heads = (int)H;
-  a.seqlen = (int)S;
-  a.head_dim = (int)D;
-  a.num_blocks = (int)num_blocks;
-  a.max_kv = (int)max_kv;
-  a.sm_scale = (float)sm_scale;
-
-  // Report an unsupported regime loudly rather than returning plausible-looking wrong values.
-  TORCH_CHECK(block_sparse_supported(a) == cudaSuccess,
-              "block_sparse_sm100a: unsupported configuration -- requires head_dim==",
-              HEAD_DIM, ", an even num_blocks, seqlen == num_blocks*", BLOCK,
-              ", and a variable_block_sizes tensor. Got head_dim=", D, " num_blocks=",
-              num_blocks, " seqlen=", S);
-
-  const cudaError_t err = launch_block_sparse_sm100a(a, at::cuda::getCurrentCUDAStream());
-  TORCH_CHECK(err == cudaSuccess,
-              "block_sparse_sm100a launch failed: ", cudaGetErrorString(err));
-
-  if (need_lse) return {out, lse};
-  return {out};
-}
@@ -1,877 +0,0 @@
-// primitives.cuh -- device primitives for the sm_100a VSA block-sparse attention
-// forward: tcgen05 (alloc / mma / ld / st / commit / wait / fence), TMA load / store /
-// tensormap, mbarrier, cluster launch control, setmaxnreg, fast math, and the FMHA helpers.
-//
-// Generated and pruned to what the kernel reaches -- do not edit by hand.
-#pragma once
-#include <cstdint>
-#include <cstdio>
-#include <cstdlib>
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_runtime.h>
-#include <cmath>
-#include <cassert>
-#include <cstring>
-#include <vector_types.h>
-
-#ifndef CUDA_CHECK
-#define CUDA_CHECK(stmt) do {                                                 \
-    cudaError_t _e = (stmt);                                                  \
-    if (_e != cudaSuccess) {                                                  \
-      fprintf(stderr, "CUDA error %s:%d: %s -> %s\n",                         \
-              __FILE__, __LINE__, #stmt, cudaGetErrorString(_e));             \
-      std::exit(1);                                                           \
-    }                                                                         \
-  } while (0)
-#endif
-
-__device__ __forceinline__
-uint64_t mbarrier_arrive(uint32_t mbar_smem) {
-  uint64_t state;
-  asm volatile("mbarrier.arrive.shared::cta.b64 %0, [%1];\n"
-               : "=l"(state) : "r"(mbar_smem) : "memory");
-  return state;
-}
-
-__device__ __forceinline__
-void mbarrier_arrive_nostate(uint32_t mbar_smem) {
-  asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n"
-               :: "r"(mbar_smem) : "memory");
-}
-
-__device__ __forceinline__
-void mbarrier_arrive_cluster_default(uint32_t cluster_smem_addr) {
-  asm volatile("mbarrier.arrive.shared::cluster.b64 _, [%0];\n"
-               :: "r"(cluster_smem_addr) : "memory");
-}
-
-__device__ __forceinline__
-void mbarrier_arrive_expect_tx(uint32_t mbar_smem, uint32_t expected_bytes) {
-  asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;\n"
-               :: "r"(mbar_smem), "r"(expected_bytes) : "memory");
-}
-
-__device__ __forceinline__
-void mbarrier_wait_parity_suspend(uint32_t mbar_smem, uint32_t phase_parity) {
-  asm volatile(
-    "{\n"
-    ".reg .pred P1;\n"
-    "LAB_WAIT:\n"
-    "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, 10000000;\n"
-    "@!P1 bra.uni LAB_WAIT;\n"
-    "}\n"
-    :: "r"(mbar_smem), "r"(phase_parity) : "memory");
-}
-
-__device__ __forceinline__
-void mbarrier_wait_parity(uint32_t mbar_smem, uint32_t phase_parity) {
-  asm volatile(
-    "{\n"
-    ".reg .pred P1;\n"
-    "LAB_WAIT_HOT:\n"
-    "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
-    "@!P1 bra.uni LAB_WAIT_HOT;\n"
-    "}\n"
-    :: "r"(mbar_smem), "r"(phase_parity) : "memory");
-}
-
-__device__ __forceinline__
-void fence_proxy_async_shared_cta() {
-  asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
-}
-
-__device__ __forceinline__
-void fence_proxy_async_shared() {
-  asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory");
-}
-
-__device__ __forceinline__ void clc_try_cancel_async(
-    uint32_t smem_dst, uint32_t mbar_smem) {
-  asm volatile(
-    "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.b128"
-    " [%0], [%1];\n"
-    :: "r"(smem_dst), "r"(mbar_smem) : "memory");
-}
-
-__device__ __forceinline__ void clc_load_response(
-    uint32_t smem_slot, uint32_t& r0, uint32_t& r1,
-    uint32_t& r2, uint32_t& r3) {
-  asm volatile("ld.shared::cta.v4.b32 {%0, %1, %2, %3}, [%4];\n"
-               : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
-               : "r"(smem_slot));
-}
-
-template <int NUM_STAGES>
-struct MbarrierPhaseTracker {
-
-  uint32_t phase[NUM_STAGES];
-  int idx;
-
-  __device__ __forceinline__
-  void init() {
-    for (int i = 0; i < NUM_STAGES; ++i) phase[i] = 0;
-    idx = 0;
-  }
-
-  __device__ __forceinline__
-  uint32_t current_phase() const { return phase[idx]; }
-
-  __device__ __forceinline__
-  void advance() {
-    phase[idx] ^= 1u;
-    idx = (idx + 1) % NUM_STAGES;
-  }
-
-  __device__ __forceinline__
-  int stage() const { return idx; }
-};
-
-template <int NUM_STAGES>
-struct PhaseTracker {
-  int stage;
-  uint32_t phase;
-
-  __device__ __forceinline__
-  PhaseTracker() : stage(0), phase(0) {}
-
-  __device__ __forceinline__
-  void advance() {
-    stage++;
-    if (stage == NUM_STAGES) {
-      stage = 0;
-      phase ^= 1;
-    }
-  }
-
-  __device__ __forceinline__
-  int get_stage() const { return stage; }
-
-  __device__ __forceinline__
-  uint32_t get_phase() const { return phase; }
-};
-
-template <int NUM_STAGES>
-struct EmptyPhaseTracker {
-  int stage;
-  uint32_t phase;
-
-  __device__ __forceinline__
-  EmptyPhaseTracker() : stage(0), phase(1) {}
-
-  __device__ __forceinline__
-  void advance() {
-    stage++;
-    if (stage == NUM_STAGES) {
-      stage = 0;
-      phase ^= 1;
-    }
-  }
-
-  __device__ __forceinline__
-  int get_stage() const { return stage; }
-
-  __device__ __forceinline__
-  uint32_t get_phase() const { return phase; }
-};
-
-template <int STAGES>
-__device__ __forceinline__
-void advance_stage_phase(int& stage, uint32_t& phase) {
-  ++stage;
-  if (stage == STAGES) {
-    stage = 0;
-    phase ^= 1u;
-  }
-}
-
-static constexpr uint32_t SM100_CLC_PEER_MASK = 0xFEFFFFFF;
-
-struct ClcTileInfo {
-  int  m_tile;
-  int  n_tile;
-  bool valid;
-};
-
-enum class ClcRasterOrder { AlongN, AlongM };
-
-__device__ __forceinline__
-void clc_arrive_expect_tx_cta(uint32_t clc_full_local_addr, uint32_t tx_bytes) {
-  if ((threadIdx.x & 31) == 0) {
-    mbarrier_arrive_expect_tx(clc_full_local_addr, tx_bytes);
-  }
-}
-
-__device__ __forceinline__
-void clc_consumer_release(uint32_t clc_empty_local_addr) {
-  uint32_t peer0_addr = clc_empty_local_addr & SM100_CLC_PEER_MASK;
-  mbarrier_arrive_cluster_default(peer0_addr);
-}
-
-__device__ __forceinline__
-void clc_consumer_release_cta(uint32_t clc_empty_local_addr) {
-  mbarrier_arrive_nostate(clc_empty_local_addr);
-}
-
-template <int CLUSTER_SHAPE_M, int CLUSTER_SHAPE_N, ClcRasterOrder ORDER>
-__device__ __forceinline__
-ClcTileInfo clc_parse_response(uint32_t resp_smem_addr) {
-  uint32_t d0, d1, d2, d3;
-  fence_proxy_async_shared_cta();
-  clc_load_response(resp_smem_addr, d0, d1, d2, d3);
-  const int  ctaid_x = static_cast<int>(d0);
-  const int  ctaid_y = static_cast<int>(d1 & 0xFFFFu);
-  const bool valid   = (d2 & 1u) != 0u;
-  (void)d3;
-
-  ClcTileInfo info;
-  info.valid = valid;
-  if constexpr (ORDER == ClcRasterOrder::AlongN) {
-    info.m_tile = ctaid_y / CLUSTER_SHAPE_M;
-    info.n_tile = ctaid_x / CLUSTER_SHAPE_N;
-  } else {
-    info.m_tile = ctaid_x / CLUSTER_SHAPE_M;
-    info.n_tile = ctaid_y / CLUSTER_SHAPE_N;
-  }
-  return info;
-}
-
-template <int CLUSTER_SHAPE_M, int CLUSTER_SHAPE_N, ClcRasterOrder ORDER,
-          int CTA_GROUP = 2, bool SUSPEND = false>
-__device__ __forceinline__
-ClcTileInfo clc_fetch_next_tile(
-    uint64_t* clc_full_bar, uint64_t* clc_empty_bar, uint32_t* clc_response,
-    int clc_cons_stage, uint32_t clc_cons_phase, bool do_release) {
-  uint32_t full_addr = static_cast<uint32_t>(
-      __cvta_generic_to_shared(&clc_full_bar[clc_cons_stage]));
-  if constexpr (SUSPEND) mbarrier_wait_parity_suspend(full_addr, clc_cons_phase);
-  else                   mbarrier_wait_parity(full_addr, clc_cons_phase);
-  uint32_t resp_addr = static_cast<uint32_t>(
-      __cvta_generic_to_shared(&clc_response[clc_cons_stage * 4]));
-  ClcTileInfo t = clc_parse_response<
-      CLUSTER_SHAPE_M, CLUSTER_SHAPE_N, ORDER>(resp_addr);
-  if (do_release) {
-    uint32_t empty_local = static_cast<uint32_t>(
-        __cvta_generic_to_shared(&clc_empty_bar[clc_cons_stage]));
-    if constexpr (CTA_GROUP == 1) {
-      clc_consumer_release_cta(empty_local);
-    } else {
-      clc_consumer_release(empty_local);
-    }
-  }
-  return t;
-}
-
-template <int STAGES = 2>
-__device__ __forceinline__
-void clc_fetch_next_tile_advance(int& clc_cons_stage,
-                                 uint32_t& clc_cons_phase) {
-  advance_stage_phase<STAGES>(clc_cons_stage, clc_cons_phase);
-}
-
-__device__ __forceinline__ unsigned fdiv(unsigned n, unsigned long long pk) {
-  unsigned M = (unsigned)pk;
-  if (M == 0u) return n;
-  return __umulhi(n, M) >> (unsigned)(pk >> 32);
-}
-__host__ inline unsigned long long make_magic(unsigned d) {
-  if (d <= 1u) return 0ULL;
-  unsigned l = 0; while ((1u << (l + 1)) <= d) ++l;
-  unsigned p = 31u + l;
-  unsigned long long m = ((1ull << p) + (unsigned long long)d - 1ull) / d;
-  return (m & 0xffffffffULL) | ((unsigned long long)(p - 32u) << 32);
-}
-
-template <int BARRIER_ID>
-__device__ __forceinline__
-void bar_sync(uint32_t thread_count) {
-  static_assert(BARRIER_ID >= 0 && BARRIER_ID <= 15,
-                "bar.sync: BARRIER_ID must be in [0, 15]");
-  asm volatile("bar.sync %0, %1;\n"
-               :: "n"(BARRIER_ID), "r"(thread_count) : "memory");
-}
-
-template <int BARRIER_ID>
-__device__ __forceinline__
-void bar_arrive(uint32_t thread_count) {
-  static_assert(BARRIER_ID >= 0 && BARRIER_ID <= 15,
-                "bar.arrive: BARRIER_ID must be in [0, 15]");
-  asm volatile("bar.arrive %0, %1;\n"
-               :: "n"(BARRIER_ID), "r"(thread_count) : "memory");
-}
-
-__device__ __forceinline__ void full_bar_arrive(int m_tile, int band) {
-  switch (1 + m_tile * 4 + band) {
-    case 1: bar_arrive<1>(64); break;
-    case 2: bar_arrive<2>(64); break;
-    case 3: bar_arrive<3>(64); break;
-    case 4: bar_arrive<4>(64); break;
-    case 5: bar_arrive<5>(64); break;
-    case 6: bar_arrive<6>(64); break;
-    case 7: bar_arrive<7>(64); break;
-    case 8: bar_arrive<8>(64); break;
-  }
-}
-__device__ __forceinline__ void full_bar_wait(int m_tile, int band) {
-  switch (1 + m_tile * 4 + band) {
-    case 1: bar_sync<1>(64); break;
-    case 2: bar_sync<2>(64); break;
-    case 3: bar_sync<3>(64); break;
-    case 4: bar_sync<4>(64); break;
-    case 5: bar_sync<5>(64); break;
-    case 6: bar_sync<6>(64); break;
-    case 7: bar_sync<7>(64); break;
-    case 8: bar_sync<8>(64); break;
-  }
-}
-
-template <bool IS_CAUSAL, int K_TILE>
-__device__ __forceinline__ void mask_s_row_r2p(float* scores, int k_offset, int q_pos, int seqlen_k) {
-  int n_keep = seqlen_k - k_offset;
-  if constexpr (IS_CAUSAL) {
-    const int causal = q_pos - k_offset + 1;
-    n_keep = n_keep < causal ? n_keep : causal;
-  }
-  #pragma unroll
-  for (int s = 0; s < K_TILE / 32; ++s) {
-    int m = (s + 1) * 32 - n_keep;
-    m = m < 0 ? 0 : (m > 32 ? 32 : m);
-    const uint32_t keep = (m >= 32) ? 0u : (0xFFFFFFFFu >> m);
-    #pragma unroll
-    for (int i = 0; i < 32; ++i)
-      if (!(keep & (1u << i))) scores[s * 32 + i] = -INFINITY;
-  }
-}
-
-template <int CTA_GROUP>
-__device__ __forceinline__ void tcgen05_alloc(uint32_t smem_dst_ptr,
-                                              uint32_t n_cols) {
-  static_assert(CTA_GROUP == 1 || CTA_GROUP == 2,
-                "tcgen05_alloc: CTA_GROUP must be 1 or 2");
-  if constexpr (CTA_GROUP == 1) {
-    asm volatile(
-      "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n"
-      :: "r"(smem_dst_ptr), "r"(n_cols));
-  } else {
-    asm volatile(
-      "tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;\n"
-      :: "r"(smem_dst_ptr), "r"(n_cols));
-  }
-}
-
-__device__ __forceinline__ void tcgen05_st_32x32b_x16(
-    uint32_t tmem_addr, const uint32_t (&r)[16]) {
-  asm volatile("tcgen05.st.sync.aligned.32x32b.x16.b32 "
-               "[%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,"
-               "%11,%12,%13,%14,%15,%16};\n"
-               :: "r"(tmem_addr),
-                  "r"(r[0]),"r"(r[1]),"r"(r[2]),"r"(r[3]),
-                  "r"(r[4]),"r"(r[5]),"r"(r[6]),"r"(r[7]),
-                  "r"(r[8]),"r"(r[9]),"r"(r[10]),"r"(r[11]),
-                  "r"(r[12]),"r"(r[13]),"r"(r[14]),"r"(r[15]));
-}
-
-__device__ __forceinline__ void tcgen05_st_32x32b_x32(
-    uint32_t tmem_addr, const uint32_t (&r)[32]) {
-  asm volatile("tcgen05.st.sync.aligned.32x32b.x32.b32 "
-               "[%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,"
-               "%11,%12,%13,%14,%15,%16,%17,%18,%19,%20,"
-               "%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,"
-               "%31,%32};\n"
-               :: "r"(tmem_addr),
-                  "r"(r[0]),"r"(r[1]),"r"(r[2]),"r"(r[3]),
-                  "r"(r[4]),"r"(r[5]),"r"(r[6]),"r"(r[7]),
-                  "r"(r[8]),"r"(r[9]),"r"(r[10]),"r"(r[11]),
-                  "r"(r[12]),"r"(r[13]),"r"(r[14]),"r"(r[15]),
-                  "r"(r[16]),"r"(r[17]),"r"(r[18]),"r"(r[19]),
-                  "r"(r[20]),"r"(r[21]),"r"(r[22]),"r"(r[23]),
-                  "r"(r[24]),"r"(r[25]),"r"(r[26]),"r"(r[27]),
-                  "r"(r[28]),"r"(r[29]),"r"(r[30]),"r"(r[31]));
-}
-
-__device__ __forceinline__ void tcgen05_commit1_lead(uint32_t lead, uint32_t mbar_smem_addr) {
-  asm volatile(
-    "{\n\t"
-    ".reg .pred q;\n\t"
-    "setp.ne.b32 q, %0, 0;\n\t"
-    "@q tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%1];\n\t"
-    "}\n"
-    :: "r"(lead), "r"(mbar_smem_addr));
-}
-
-__device__ __forceinline__ void tcgen05_wait_st() {
-  asm volatile("tcgen05.wait::st.sync.aligned;\n" ::: "memory");
-}
-
-__device__ __forceinline__ void tcgen05_fence_before_thread_sync() {
-  asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory");
-}
-
-__device__ __forceinline__
-void tma_load_2d(uint32_t smem_dst, const void* tensormap_ptr,
-                 uint32_t mbar_smem, int coord_x, int coord_y) {
-  asm volatile(
-    "cp.async.bulk.tensor.2d.shared::cluster.global.tile"
-    ".mbarrier::complete_tx::bytes"
-    " [%0], [%1, {%3, %4}], [%2];\n"
-    :: "r"(smem_dst), "l"(tensormap_ptr),
-       "r"(mbar_smem), "r"(coord_x), "r"(coord_y)
-    : "memory");
-}
-
-__device__ __forceinline__
-void tma_load_3d(uint32_t smem_dst, const void* tensormap_ptr,
-                 uint32_t mbar_smem, int c0, int c1, int c2) {
-  asm volatile(
-    "cp.async.bulk.tensor.3d.shared::cluster.global.tile"
-    ".mbarrier::complete_tx::bytes"
-    " [%0], [%1, {%3, %4, %5}], [%2];\n"
-    :: "r"(smem_dst), "l"(tensormap_ptr),
-       "r"(mbar_smem), "r"(c0), "r"(c1), "r"(c2)
-    : "memory");
-}
-
-__device__ __forceinline__
-void tma_load_4d(uint32_t smem_dst, const void* tensormap_ptr,
-                 uint32_t mbar_smem, int c0, int c1, int c2, int c3) {
-  asm volatile(
-    "cp.async.bulk.tensor.4d.shared::cluster.global.tile"
-    ".mbarrier::complete_tx::bytes"
-    " [%0], [%1, {%3, %4, %5, %6}], [%2];\n"
-    :: "r"(smem_dst), "l"(tensormap_ptr),
-       "r"(mbar_smem), "r"(c0), "r"(c1), "r"(c2), "r"(c3)
-    : "memory");
-}
-
-template <int CTA_GROUP>
-__device__ __forceinline__ void tcgen05_dealloc(uint32_t tmem_addr,
-                                                uint32_t n_cols) {
-  static_assert(CTA_GROUP == 1 || CTA_GROUP == 2,
-                "tcgen05_dealloc: CTA_GROUP must be 1 or 2");
-  if constexpr (CTA_GROUP == 1) {
-    asm volatile(
-      "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n"
-      :: "r"(tmem_addr), "r"(n_cols));
-  } else {
-    asm volatile(
-      "tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;\n"
-      :: "r"(tmem_addr), "r"(n_cols));
-  }
-}
-
-__device__ __forceinline__
-void tma_store_2d(const void* tensormap_ptr, int coord_x, int coord_y,
-                  uint32_t smem_src) {
-  asm volatile(
-    "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group"
-    " [%0, {%1, %2}], [%3];\n"
-    :: "l"(tensormap_ptr), "r"(coord_x), "r"(coord_y),
-       "r"(smem_src)
-    : "memory");
-}
-
-__device__ __forceinline__
-void tma_store_3d(const void* tensormap_ptr, int c0, int c1, int c2,
-                  uint32_t smem_src) {
-  asm volatile(
-    "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group"
-    " [%0, {%1, %2, %3}], [%4];\n"
-    :: "l"(tensormap_ptr), "r"(c0), "r"(c1), "r"(c2), "r"(smem_src)
-    : "memory");
-}
-
-__device__ __forceinline__
-void tma_store_4d(const void* tensormap_ptr, int c0, int c1, int c2, int c3,
-                  uint32_t smem_src) {
-  asm volatile(
-    "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group"
-    " [%0, {%1, %2, %3, %4}], [%5];\n"
-    :: "l"(tensormap_ptr), "r"(c0), "r"(c1), "r"(c2), "r"(c3), "r"(smem_src)
-    : "memory");
-}
-
-inline cudaError_t make_tma_2d_tiled(
-    CUtensorMap* out,
-    const void* ptr, int rows, int cols, int box_rows, int box_cols,
-    int elem_bytes, CUtensorMapDataType dtype,
-    CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_128B,
-    CUtensorMapL2promotion l2 = CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
-    CUtensorMapFloatOOBfill oob = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) {
-  uint64_t globalDim[2]     = { (uint64_t)cols, (uint64_t)rows };
-  uint64_t globalStrides[1] = { (uint64_t)cols * (uint64_t)elem_bytes };
-  uint32_t boxDim[2]        = { (uint32_t)box_cols, (uint32_t)box_rows };
-  uint32_t elemStrides[2]   = { 1u, 1u };
-
-  CUresult r = cuTensorMapEncodeTiled(
-      out, dtype, 2,
-      const_cast<void*>(ptr), globalDim, globalStrides,
-      boxDim, elemStrides,
-      CU_TENSOR_MAP_INTERLEAVE_NONE,
-      swizzle, l2, oob);
-  return (r == CUDA_SUCCESS) ? cudaSuccess : cudaErrorInvalidValue;
-}
-
-__device__ __forceinline__
-void cp_async_bulk_commit_group() {
-  asm volatile("cp.async.bulk.commit_group;\n" ::: "memory");
-}
-
-template <int N>
-__device__ __forceinline__
-void cp_async_bulk_wait_group_read() {
-  asm volatile("cp.async.bulk.wait_group.read %0;\n" :: "n"(N) : "memory");
-}
-
-__device__ __forceinline__
-void mbarrier_init(uint32_t mbar_smem, uint32_t arrive_count) {
-  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n"
-               :: "r"(mbar_smem), "r"(arrive_count) : "memory");
-}
-
-template <int CTA_GROUP>
-__device__ __forceinline__ void tcgen05_relinquish_alloc_permit() {
-  static_assert(CTA_GROUP == 1 || CTA_GROUP == 2,
-                "tcgen05_relinquish_alloc_permit: CTA_GROUP must be 1 or 2");
-  if constexpr (CTA_GROUP == 1) {
-    asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::);
-  } else {
-    asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;\n" ::);
-  }
-}
-
-__device__ __forceinline__
-void fence_mbarrier_init_release_cluster() {
-  asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory");
-}
-
-__device__ __forceinline__ void tcgen05_mma_f16_ss_lead(uint32_t lead,
-    uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b, uint32_t idesc,
-    bool enable_input_d) {
-  asm volatile(
-    "{\n\t"
-    ".reg .pred p, q;\n\t"
-    "setp.ne.b32 q, %0, 0;\n\t"
-    "setp.ne.b32 p, %5, 0;\n\t"
-    "@q tcgen05.mma.cta_group::1.kind::f16 [%1], %2, %3, %4, {%6, %7, %8, %9}, p;\n\t"
-    "}\n"
-    :: "r"(lead), "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(idesc),
-       "r"(enable_input_d ? 1u : 0u), "r"(0u), "r"(0u), "r"(0u), "r"(0u));
-}
-
-__device__ __forceinline__ void tcgen05_mma_f16_ts_1sm_lead(uint32_t lead,
-    uint32_t tmem_c, uint32_t tmem_a, uint64_t desc_b, uint32_t idesc,
-    bool enable_input_d) {
-  asm volatile(
-    "{\n\t"
-    ".reg .pred p, q;\n\t"
-    "setp.ne.b32 q, %0, 0;\n\t"
-    "setp.ne.b32 p, %5, 0;\n\t"
-    "@q tcgen05.mma.cta_group::1.kind::f16 [%1], [%2], %3, %4, {%6, %7, %8, %9}, p;\n\t"
-    "}\n"
-    :: "r"(lead), "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(idesc),
-       "r"(enable_input_d ? 1u : 0u), "r"(0u), "r"(0u), "r"(0u), "r"(0u));
-}
-
-enum class SmemSwizzleBlackwell : uint32_t {
-  None = 0,
-  B128_32atom = 1,
-  B128 = 2,
-  B64  = 4,
-  B32  = 6,
-};
-
-__device__ __host__ __forceinline__ uint64_t build_smem_desc_blackwell(
-    uint32_t smem_addr,
-    uint32_t stride_byte_offset,
-    uint32_t leading_byte_offset,
-    SmemSwizzleBlackwell swizzle = SmemSwizzleBlackwell::B128,
-    uint32_t base_offset = 0) {
-  uint64_t d = 0;
-  d |= static_cast<uint64_t>((smem_addr >> 4) & 0x3FFF);
-  d |= static_cast<uint64_t>((leading_byte_offset >> 4) & 0x3FFF) << 16;
-  d |= static_cast<uint64_t>((stride_byte_offset >> 4) & 0x3FFF) << 32;
-  d |= static_cast<uint64_t>(1) << 46;
-  d |= (static_cast<uint64_t>(base_offset) & 0x7) << 49;
-  d |= static_cast<uint64_t>(static_cast<uint32_t>(swizzle) & 0x7) << 61;
-  return d;
-}
-
-__device__ __forceinline__
-uint32_t elect_one_sync() {
-  uint32_t elected;
-  asm volatile(
-    "{\n\t"
-    ".reg .pred p;\n\t"
-    "elect.sync %0|p, 0xffffffff;\n\t"
-    "selp.b32 %0, 1, 0, p;\n\t"
-    "}\n"
-    : "=r"(elected));
-  return elected;
-}
-
-__device__ __forceinline__
-uint32_t elect_one_sync(uint32_t membermask) {
-  uint32_t elected;
-  asm volatile(
-    "{\n\t"
-    ".reg .pred p;\n\t"
-    "elect.sync %0|p, %1;\n\t"
-    "selp.b32 %0, 1, 0, p;\n\t"
-    "}\n"
-    : "=r"(elected) : "r"(membermask));
-  return elected;
-}
-
-template <int N>
-__device__ __forceinline__
-void setmaxnreg_dec() {
-  static_assert(N >= 24 && N <= 256 && N % 8 == 0,
-                "setmaxnreg_dec: N must be in [24, 256] and a multiple of 8");
-  asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n"
-               :: "n"(N) : "memory");
-}
-
-template <int N>
-__device__ __forceinline__
-void setmaxnreg_inc() {
-  static_assert(N >= 24 && N <= 256 && N % 8 == 0,
-                "setmaxnreg_inc: N must be in [24, 256] and a multiple of 8");
-  asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n"
-               :: "n"(N) : "memory");
-}
-
-__device__ __forceinline__
-uint32_t cvt_f32x2_to_bf16x2(float a, float b) {
-  uint32_t r;
-  asm volatile("cvt.rn.bf16x2.f32 %0, %2, %1;\n"
-               : "=r"(r) : "f"(a), "f"(b));
-  return r;
-}
-
-namespace {
-__device__ __forceinline__ uint64_t f32x2_bits(float2 v) {
-  uint64_t b; __builtin_memcpy(&b, &v, 8); return b;
-}
-__device__ __forceinline__ float2 f32x2_make(uint64_t b) {
-  float2 v; __builtin_memcpy(&v, &b, 8); return v;
-}
-}
-
-__device__ __forceinline__ float2 fmul2(float2 a, float2 b) {
-  uint64_t d;
-  asm volatile("mul.f32x2 %0, %1, %2;\n" : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b)));
-  return f32x2_make(d);
-}
-
-__device__ __forceinline__ float2 fadd2(float2 a, float2 b) {
-  uint64_t d;
-  asm volatile("add.f32x2 %0, %1, %2;\n" : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b)));
-  return f32x2_make(d);
-}
-
-__device__ __forceinline__ float2 ffma2(float2 a, float2 b, float2 c) {
-  uint64_t d;
-  asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n"
-               : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b)), "l"(f32x2_bits(c)));
-  return f32x2_make(d);
-}
-
-__device__ __forceinline__ float2 f32x2_splat(float s) { return make_float2(s, s); }
-
-__device__ __forceinline__ float ex2_approx_f32(float z) {
-  float d;
-  asm volatile("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(d) : "f"(z));
-  return d;
-}
-
-__device__ __forceinline__ float2 ex2_emu_f32x2(float x, float y) {
-  uint32_t ox, oy;
-  asm volatile(
-    "{\n\t"
-    ".reg .f32 f1,f2,f3,f4,f5,f6,f7;\n\t"
-    ".reg .b64 l1,l2,l3,l4,l5,l6,l7,l8,l9,l10;\n\t"
-    ".reg .s32 r1,r2,r3,r4,r5,r6,r7,r8;\n\t"
-    "max.f32 f1, %2, 0fC2FE0000;\n\t"
-    "max.f32 f2, %3, 0fC2FE0000;\n\t"
-    "mov.b64 l1, {f1, f2};\n\t"
-    "mov.f32 f3, 0f4B400000;\n\t"
-    "mov.b64 l2, {f3, f3};\n\t"
-    "add.rm.f32x2 l7, l1, l2;\n\t"
-    "sub.rn.f32x2 l8, l7, l2;\n\t"
-    "sub.rn.f32x2 l9, l1, l8;\n\t"
-    "mov.f32 f7, 0f3D9DF09D;\n\t"
-    "mov.b64 l6, {f7, f7};\n\t"
-    "mov.f32 f6, 0f3E6906A4;\n\t"
-    "mov.b64 l5, {f6, f6};\n\t"
-    "mov.f32 f5, 0f3F31F519;\n\t"
-    "mov.b64 l4, {f5, f5};\n\t"
-    "mov.f32 f4, 0f3F800000;\n\t"
-    "mov.b64 l3, {f4, f4};\n\t"
-    "fma.rn.f32x2 l10, l9, l6, l5;\n\t"
-    "fma.rn.f32x2 l10, l10, l9, l4;\n\t"
-    "fma.rn.f32x2 l10, l10, l9, l3;\n\t"
-    "mov.b64 {r1, r2}, l7;\n\t"
-    "mov.b64 {r3, r4}, l10;\n\t"
-    "shl.b32 r5, r1, 23;\n\t"
-    "add.s32 r7, r5, r3;\n\t"
-    "shl.b32 r6, r2, 23;\n\t"
-    "add.s32 r8, r6, r4;\n\t"
-    "mov.b32 %0, r7;\n\t"
-    "mov.b32 %1, r8;\n\t"
-    "}\n"
-    : "=r"(ox), "=r"(oy) : "f"(x), "f"(y));
-  float2 r; __builtin_memcpy(&r.x, &ox, 4); __builtin_memcpy(&r.y, &oy, 4);
-  return r;
-}
-
-__device__ __forceinline__ float rcp_approx_ftz_f32(float x) {
-  float d;
-  asm("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(d) : "f"(x));
-  return d;
-}
-
-__device__ __forceinline__ uint32_t make_idesc_table44(
-    int M, int N,
-    uint32_t dtype, uint32_t atype, uint32_t btype,
-    bool transpose_a = false, bool transpose_b = false,
-    bool negate_a = false, bool negate_b = false) {
-  uint32_t idesc = 0;
-  idesc |= (dtype & 0x3) << 4;
-  idesc |= (atype & 0x7) << 7;
-  idesc |= (btype & 0x7) << 10;
-  idesc |= (negate_a ? 1u : 0u) << 13;
-  idesc |= (negate_b ? 1u : 0u) << 14;
-  idesc |= (transpose_a ? 1u : 0u) << 15;
-  idesc |= (transpose_b ? 1u : 0u) << 16;
-  idesc |= ((static_cast<uint32_t>(N) >> 3) & 0x3F) << 17;
-  idesc |= ((static_cast<uint32_t>(M) >> 4) & 0x1F) << 24;
-  return idesc;
-}
-
-__device__ __forceinline__ uint32_t make_idesc_bf16_f32(
-    int M, int N, bool ta = false, bool tb = false) {
-  return make_idesc_table44(M, N,  1,
-                              1,  1, ta, tb);
-}
-
-__device__ __forceinline__ void tcgen05_ld_32x32b_x16(
-    uint32_t tmem_addr, uint32_t (&r)[16]) {
-  asm volatile("tcgen05.ld.sync.aligned.32x32b.x16.b32 "
-               "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,"
-               "%10,%11,%12,%13,%14,%15}, [%16];\n"
-               : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]),
-                 "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]),
-                 "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]),
-                 "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15])
-               : "r"(tmem_addr));
-}
-
-__device__ __forceinline__ void tcgen05_ld_32x32b_x32(
-    uint32_t tmem_addr, uint32_t (&r)[32]) {
-  asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 "
-               "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,"
-               "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,"
-               "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,"
-               "%30,%31}, [%32];\n"
-               : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]),
-                 "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]),
-                 "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]),
-                 "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15]),
-                 "=r"(r[16]),"=r"(r[17]),"=r"(r[18]),"=r"(r[19]),
-                 "=r"(r[20]),"=r"(r[21]),"=r"(r[22]),"=r"(r[23]),
-                 "=r"(r[24]),"=r"(r[25]),"=r"(r[26]),"=r"(r[27]),
-                 "=r"(r[28]),"=r"(r[29]),"=r"(r[30]),"=r"(r[31])
-               : "r"(tmem_addr));
-}
-
-__device__ __forceinline__ void tcgen05_ld_32x32b_x64(
-    uint32_t tmem_addr, uint32_t (&r)[64]) {
-  asm volatile("tcgen05.ld.sync.aligned.32x32b.x64.b32 "
-               "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,"
-               "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,"
-               "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,"
-               "%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,"
-               "%40,%41,%42,%43,%44,%45,%46,%47,%48,%49,"
-               "%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,"
-               "%60,%61,%62,%63}, [%64];\n"
-               : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]),
-                 "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]),
-                 "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]),
-                 "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15]),
-                 "=r"(r[16]),"=r"(r[17]),"=r"(r[18]),"=r"(r[19]),
-                 "=r"(r[20]),"=r"(r[21]),"=r"(r[22]),"=r"(r[23]),
-                 "=r"(r[24]),"=r"(r[25]),"=r"(r[26]),"=r"(r[27]),
-                 "=r"(r[28]),"=r"(r[29]),"=r"(r[30]),"=r"(r[31]),
-                 "=r"(r[32]),"=r"(r[33]),"=r"(r[34]),"=r"(r[35]),
-                 "=r"(r[36]),"=r"(r[37]),"=r"(r[38]),"=r"(r[39]),
-                 "=r"(r[40]),"=r"(r[41]),"=r"(r[42]),"=r"(r[43]),
-                 "=r"(r[44]),"=r"(r[45]),"=r"(r[46]),"=r"(r[47]),
-                 "=r"(r[48]),"=r"(r[49]),"=r"(r[50]),"=r"(r[51]),
-                 "=r"(r[52]),"=r"(r[53]),"=r"(r[54]),"=r"(r[55]),
-                 "=r"(r[56]),"=r"(r[57]),"=r"(r[58]),"=r"(r[59]),
-                 "=r"(r[60]),"=r"(r[61]),"=r"(r[62]),"=r"(r[63])
-               : "r"(tmem_addr));
-}
-
-__device__ __forceinline__ void tcgen05_ld_32x32b_x128(
-    uint32_t tmem_addr, uint32_t (&r)[128]) {
-  asm volatile("tcgen05.ld.sync.aligned.32x32b.x128.b32 "
-               "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,"
-               "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19,"
-               "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,"
-               "%30,%31,%32,%33,%34,%35,%36,%37,%38,%39,"
-               "%40,%41,%42,%43,%44,%45,%46,%47,%48,%49,"
-               "%50,%51,%52,%53,%54,%55,%56,%57,%58,%59,"
-               "%60,%61,%62,%63,%64,%65,%66,%67,%68,%69,"
-               "%70,%71,%72,%73,%74,%75,%76,%77,%78,%79,"
-               "%80,%81,%82,%83,%84,%85,%86,%87,%88,%89,"
-               "%90,%91,%92,%93,%94,%95,%96,%97,%98,%99,"
-               "%100,%101,%102,%103,%104,%105,%106,%107,%108,%109,"
-               "%110,%111,%112,%113,%114,%115,%116,%117,%118,%119,"
-               "%120,%121,%122,%123,%124,%125,%126,%127}, [%128];\n"
-               : "=r"(r[  0]),"=r"(r[  1]),"=r"(r[  2]),"=r"(r[  3]),
-                 "=r"(r[  4]),"=r"(r[  5]),"=r"(r[  6]),"=r"(r[  7]),
-                 "=r"(r[  8]),"=r"(r[  9]),"=r"(r[ 10]),"=r"(r[ 11]),
-                 "=r"(r[ 12]),"=r"(r[ 13]),"=r"(r[ 14]),"=r"(r[ 15]),
-                 "=r"(r[ 16]),"=r"(r[ 17]),"=r"(r[ 18]),"=r"(r[ 19]),
-                 "=r"(r[ 20]),"=r"(r[ 21]),"=r"(r[ 22]),"=r"(r[ 23]),
-                 "=r"(r[ 24]),"=r"(r[ 25]),"=r"(r[ 26]),"=r"(r[ 27]),
-                 "=r"(r[ 28]),"=r"(r[ 29]),"=r"(r[ 30]),"=r"(r[ 31]),
-                 "=r"(r[ 32]),"=r"(r[ 33]),"=r"(r[ 34]),"=r"(r[ 35]),
-                 "=r"(r[ 36]),"=r"(r[ 37]),"=r"(r[ 38]),"=r"(r[ 39]),
-                 "=r"(r[ 40]),"=r"(r[ 41]),"=r"(r[ 42]),"=r"(r[ 43]),
-                 "=r"(r[ 44]),"=r"(r[ 45]),"=r"(r[ 46]),"=r"(r[ 47]),
-                 "=r"(r[ 48]),"=r"(r[ 49]),"=r"(r[ 50]),"=r"(r[ 51]),
-                 "=r"(r[ 52]),"=r"(r[ 53]),"=r"(r[ 54]),"=r"(r[ 55]),
-                 "=r"(r[ 56]),"=r"(r[ 57]),"=r"(r[ 58]),"=r"(r[ 59]),
-                 "=r"(r[ 60]),"=r"(r[ 61]),"=r"(r[ 62]),"=r"(r[ 63]),
-                 "=r"(r[ 64]),"=r"(r[ 65]),"=r"(r[ 66]),"=r"(r[ 67]),
-                 "=r"(r[ 68]),"=r"(r[ 69]),"=r"(r[ 70]),"=r"(r[ 71]),
-                 "=r"(r[ 72]),"=r"(r[ 73]),"=r"(r[ 74]),"=r"(r[ 75]),
-                 "=r"(r[ 76]),"=r"(r[ 77]),"=r"(r[ 78]),"=r"(r[ 79]),
-                 "=r"(r[ 80]),"=r"(r[ 81]),"=r"(r[ 82]),"=r"(r[ 83]),
-                 "=r"(r[ 84]),"=r"(r[ 85]),"=r"(r[ 86]),"=r"(r[ 87]),
-                 "=r"(r[ 88]),"=r"(r[ 89]),"=r"(r[ 90]),"=r"(r[ 91]),
-                 "=r"(r[ 92]),"=r"(r[ 93]),"=r"(r[ 94]),"=r"(r[ 95]),
-                 "=r"(r[ 96]),"=r"(r[ 97]),"=r"(r[ 98]),"=r"(r[ 99]),
-                 "=r"(r[100]),"=r"(r[101]),"=r"(r[102]),"=r"(r[103]),
-                 "=r"(r[104]),"=r"(r[105]),"=r"(r[106]),"=r"(r[107]),
-                 "=r"(r[108]),"=r"(r[109]),"=r"(r[110]),"=r"(r[111]),
-                 "=r"(r[112]),"=r"(r[113]),"=r"(r[114]),"=r"(r[115]),
-                 "=r"(r[116]),"=r"(r[117]),"=r"(r[118]),"=r"(r[119]),
-                 "=r"(r[120]),"=r"(r[121]),"=r"(r[122]),"=r"(r[123]),
-                 "=r"(r[124]),"=r"(r[125]),"=r"(r[126]),"=r"(r[127])
-               : "r"(tmem_addr));
-}
-
-__device__ __forceinline__
-uint32_t smem_ptr_u32(const void* ptr) {
-  return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
-}
-
-__device__ __forceinline__
-void sts_f32(uint32_t smem_addr, float val) {
-  asm volatile("st.shared.f32 [%0], %1;" :: "r"(smem_addr), "f"(val) : "memory");
-}
-
@@ -28,31 +28,10 @@ void register_rms_norm(pybind11::module_ &);
 void register_layer_norm(pybind11::module_ &);
 void register_gemm(pybind11::module_ &);

-#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A
-extern std::vector<torch::Tensor> block_sparse_sm100a_fwd(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<torch::Tensor> v_t,
-    torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes,
-    double sm_scale, bool need_lse);
-extern std::vector<torch::Tensor> block_sparse_sm100a_blk128_fwd(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<torch::Tensor> v_t,
-    torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes,
-    double sm_scale, bool need_lse);
-#endif
-
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.doc() = "FastVideo CUDA Kernels";

-#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A
-    m.def("block_sparse_sm100a_fwd",
-          torch::wrap_pybind_function(block_sparse_sm100a_fwd),
-          "VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a)");
-    m.def("block_sparse_sm100a_blk128_fwd",
-          torch::wrap_pybind_function(block_sparse_sm100a_blk128_fwd),
-          "VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a)");
-#endif
-
 #ifdef TK_COMPILE_ST_ATTN
-
    m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention (Hopper)");
 #endif

@@ -0,0 +1,4 @@
+[flake8]
+max-line-length = 100
+# W503: line break before binary operator
+ignore = E731, E741, F841, W503
@@ -0,0 +1,8 @@
+Tri Dao
+Jay Shah
+Ted Zadouri
+Markus Hoehnerbach
+Vijay Thakkar
+Timmy Liu
+Driss Guessous
+Reuben Stern
@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,6 @@
+include UPSTREAM.md
+global-exclude *.egg-info/*
+prune flash_attn_4.egg-info
+prune flash_attn.egg-info
+prune build
+prune dist
@@ -0,0 +1,33 @@
+# FlashAttention-4 (CuTeDSL)
+
+FlashAttention-4 is a CuTeDSL-based implementation of FlashAttention for Hopper and Blackwell GPUs.
+
+## Installation
+
+```sh
+pip install flash-attn-4
+```
+
+If you're on CUDA 13, install with the `cu13` extra for best performance:
+
+```sh
+pip install "flash-attn-4[cu13]"
+```
+
+## Usage
+
+```python
+from flash_attn.cute import flash_attn_func, flash_attn_varlen_func
+
+out = flash_attn_func(q, k, v, causal=True)
+```
+
+## Development
+
+```sh
+git clone https://github.com/Dao-AILab/flash-attention.git
+cd flash-attention
+pip install -e "flash_attn/cute[dev]"       # CUDA 12.x
+pip install -e "flash_attn/cute[dev,cu13]"  # CUDA 13.x (e.g. B200)
+pytest tests/cute/
+```
@@ -0,0 +1,37 @@
+# FastVideo FA4 source fork
+
+This directory vendors the `flash_attn.cute` Python package from
+[Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention) at
+commit [`82d6441eec5d4dfec120153db2c0145ae855a083`](https://github.com/Dao-AILab/flash-attention/commit/82d6441eec5d4dfec120153db2c0145ae855a083).
+The upstream source path is `flash_attn/cute/`, and its BSD-3-Clause license is
+preserved in `LICENSE`.
+
+FastVideo keeps this as a first-class source fork because its video sparse
+attention kernels need FA4 scheduling and tile-shape changes that are developed
+and benchmarked together with `fastvideo-kernel`. The distribution intentionally
+retains the upstream name `flash-attn-4` and import namespace `flash_attn.cute`.
+
+The forward-kernel delta currently contains FastVideo's single-Q-block,
+two-KV-stream Q128/KV64 schedule, four-slot score/probability double buffering,
+and FP32 online-softmax merge. The Q64/KV64 primitives were adapted from
+upstream's guarded smaller-tile commit
+[`526c18d25bcbc7fc7d6740ab3c7c84ed2d42cb0b`](https://github.com/Dao-AILab/flash-attention/commit/526c18d25bcbc7fc7d6740ab3c7c84ed2d42cb0b),
+then extended with the same two-stream sparse traversal and merge. These are
+forward-only specializations; the existing upstream paths remain the fallback
+outside their explicit SM100/SM110 dtype, shape, and metadata gates.
+
+Install the working tree without resolving dependencies:
+
+```bash
+uv pip install --no-deps --editable fastvideo-kernel/fa4
+```
+
+The initial performance comparison came from FlashInfer issue
+[#4554](https://github.com/flashinfer-ai/flashinfer/issues/4554) and its
+[benchmark harness](https://gist.github.com/SolitaryThinker/90a1d1447929fc38dc509c1852e76532)
+at gist revision `e15ac9066f23ef3690e33e1cc1fdac45b4b9099f`.
+
+When refreshing from upstream, copy only `flash_attn/cute/` from an explicitly
+checked-out commit, retain this file and FastVideo-specific changes, update the
+commit and package version above, and rerun the FA4 correctness and GB200
+performance gates before accepting the refresh.
@@ -0,0 +1,18 @@
+"""Flash Attention CUTE (CUDA Template Engine) implementation."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("flash-attn-4")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+from .interface import (
+    flash_attn_func,
+    flash_attn_varlen_func,
+)
+
+__all__ = [
+    "flash_attn_func",
+    "flash_attn_varlen_func",
+]
@@ -0,0 +1,103 @@
+# Copyright (c) 2025, Tri Dao.
+from typing import Type, Callable, Optional
+
+import cutlass
+import cutlass.cute as cute
+
+
+def get_smem_layout_atom(dtype: Type[cutlass.Numeric], k_dim: int) -> cute.ComposedLayout:
+    dtype_byte = cutlass.const_expr(dtype.width // 8)
+    bytes_per_row = cutlass.const_expr(k_dim * dtype_byte)
+    smem_k_block_size = (
+        cutlass.const_expr(
+            128
+            if bytes_per_row % 128 == 0
+            else (64 if bytes_per_row % 64 == 0 else (32 if bytes_per_row % 32 == 0 else 16))
+        )
+        // dtype_byte
+    )
+    swizzle_bits = (
+        4
+        if smem_k_block_size == 128
+        else (3 if smem_k_block_size == 64 else (2 if smem_k_block_size == 32 else 1))
+    )
+    swizzle_base = 2 if dtype_byte == 4 else (3 if dtype_byte == 2 else 4)
+    return cute.make_composed_layout(
+        cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base),
+        0,
+        cute.make_ordered_layout(
+            (8 if cutlass.const_expr(k_dim % 32 == 0) else 16, smem_k_block_size), order=(1, 0)
+        ),
+    )
+
+
+@cute.jit
+def gemm(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    tCsA: cute.Tensor,
+    tCsB: cute.Tensor,
+    smem_thr_copy_A: cute.TiledCopy,
+    smem_thr_copy_B: cute.TiledCopy,
+    hook_fn: Optional[Callable] = None,
+    A_in_regs: cutlass.Constexpr[bool] = False,
+    B_in_regs: cutlass.Constexpr[bool] = False,
+    swap_AB: cutlass.Constexpr[bool] = False,
+) -> None:
+    if cutlass.const_expr(swap_AB):
+        gemm(
+            tiled_mma,
+            acc,
+            tCrB,
+            tCrA,
+            tCsB,
+            tCsA,
+            smem_thr_copy_B,
+            smem_thr_copy_A,
+            hook_fn,
+            A_in_regs=B_in_regs,
+            B_in_regs=A_in_regs,
+            swap_AB=False,
+        )
+    else:
+        tCrA_copy_view = smem_thr_copy_A.retile(tCrA)
+        tCrB_copy_view = smem_thr_copy_B.retile(tCrB)
+        if cutlass.const_expr(not A_in_regs):
+            cute.copy(smem_thr_copy_A, tCsA[None, None, 0], tCrA_copy_view[None, None, 0])
+        if cutlass.const_expr(not B_in_regs):
+            cute.copy(smem_thr_copy_B, tCsB[None, None, 0], tCrB_copy_view[None, None, 0])
+        for k in cutlass.range_constexpr(cute.size(tCsA.shape[2])):
+            if k < cute.size(tCsA.shape[2]) - 1:
+                if cutlass.const_expr(not A_in_regs):
+                    cute.copy(
+                        smem_thr_copy_A, tCsA[None, None, k + 1], tCrA_copy_view[None, None, k + 1]
+                    )
+                if cutlass.const_expr(not B_in_regs):
+                    cute.copy(
+                        smem_thr_copy_B, tCsB[None, None, k + 1], tCrB_copy_view[None, None, k + 1]
+                    )
+            cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc)
+            if cutlass.const_expr(k == 0 and hook_fn is not None):
+                hook_fn()
+
+
+@cute.jit
+def gemm_rs(
+    tiled_mma: cute.TiledMma,
+    acc: cute.Tensor,
+    tCrA: cute.Tensor,
+    tCrB: cute.Tensor,
+    tCsB: cute.Tensor,
+    smem_thr_copy_B: cute.TiledCopy,
+    hook_fn: Optional[Callable] = None,
+) -> None:
+    tCrB_copy_view = smem_thr_copy_B.retile(tCrB)
+    cute.copy(smem_thr_copy_B, tCsB[None, None, 0], tCrB_copy_view[None, None, 0])
+    for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])):
+        if cutlass.const_expr(k < cute.size(tCrA.shape[2]) - 1):
+            cute.copy(smem_thr_copy_B, tCsB[None, None, k + 1], tCrB_copy_view[None, None, k + 1])
+        cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc)
+        if cutlass.const_expr(k == 0 and hook_fn is not None):
+            hook_fn()
@@ -0,0 +1,71 @@
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm
+
+
+@dsl_user_op
+def ld_acquire(lock_ptr: cute.Pointer, *, loc=None, ip=None) -> cutlass.Int32:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    state = llvm.inline_asm(
+        T.i32(),
+        [lock_ptr_i64],
+        "ld.global.acquire.gpu.b32 $0, [$1];",
+        "=r,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+    return cutlass.Int32(state)
+
+
+@dsl_user_op
+def red_relaxed(
+    lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None
+) -> None:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [lock_ptr_i64, Int32(val).ir_value(loc=loc, ip=ip)],
+        "red.relaxed.gpu.global.add.s32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def red_release(
+    lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None
+) -> None:
+    lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [lock_ptr_i64, Int32(val).ir_value(loc=loc, ip=ip)],
+        "red.release.gpu.global.add.s32 [$0], $1;",
+        "l,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@cute.jit
+def wait_eq(lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: Int32) -> None:
+    flag_ptr = lock_ptr + flag_offset
+    if thread_idx == 0:
+        read_val = Int32(0)
+        while read_val != val:
+            read_val = ld_acquire(flag_ptr)
+
+
+@cute.jit
+def arrive_inc(
+    lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: cutlass.Constexpr[Int32]
+) -> None:
+    flag_ptr = lock_ptr + flag_offset
+    if thread_idx == 0:
+        red_release(flag_ptr, val)
+        # red_relaxed(flag_ptr, val)
@@ -0,0 +1,243 @@
+"""Shared benchmark utilities: attention_ref, cuDNN helpers, flops calculation."""
+
+import math
+import torch
+
+try:
+    import cudnn
+except ImportError:
+    cudnn = None
+
+
+# ── FLOPS calculation ────────────────────────────────────────────────────────
+
+
+def flops(
+    batch,
+    nheads,
+    seqlen_q,
+    seqlen_k,
+    headdim,
+    headdim_v,
+    causal=False,
+    window_size=(None, None),
+    has_qv=False,
+):
+    if causal:
+        avg_seqlen = (max(0, seqlen_k - seqlen_q) + seqlen_k) / 2
+    else:
+        if window_size == (None, None):
+            avg_seqlen = seqlen_k
+        else:
+            row_idx = torch.arange(seqlen_q, device="cuda")
+            col_left = (
+                torch.maximum(row_idx + seqlen_k - seqlen_q - window_size[0], torch.tensor(0))
+                if window_size[0] is not None
+                else torch.zeros_like(row_idx)
+            )
+            col_right = (
+                torch.minimum(
+                    row_idx + seqlen_k - seqlen_q + window_size[1], torch.tensor(seqlen_k - 1)
+                )
+                if window_size[1] is not None
+                else torch.full_like(row_idx, seqlen_k - 1)
+            )
+            avg_seqlen = (col_right - col_left + 1).float().mean().item()
+    eff_headdim = headdim + headdim_v if has_qv else headdim
+    return batch * nheads * 2 * seqlen_q * avg_seqlen * (eff_headdim + headdim_v)
+
+
+# ── Bandwidth calculation ────────────────────────────────────────────────────
+
+
+def bandwidth_fwd_bytes(
+    batch,
+    nheads,
+    nheads_kv,
+    seqlen_q,
+    seqlen_k,
+    headdim,
+    headdim_v,
+    dtype_bytes=2,
+    has_qv=False,
+    shared_kv=False,
+):
+    """HBM traffic for one attention pass: read Q,K,V + write O."""
+    q = batch * nheads * seqlen_q * headdim
+    qv = batch * nheads * seqlen_q * headdim_v if has_qv else 0
+    k = batch * nheads_kv * seqlen_k * headdim
+    v = batch * nheads_kv * seqlen_k * headdim_v if not shared_kv else 0
+    o = batch * nheads * seqlen_q * headdim_v
+    return (q + qv + k + v + o) * dtype_bytes
+
+
+def bandwidth_bwd_bytes(
+    batch, nheads, nheads_kv, seqlen_q, seqlen_k, headdim, headdim_v, dtype_bytes=2
+):
+    """HBM traffic for one attention pass: read Q,K,V,dO + write dQ,dK,dV."""
+    q = batch * nheads * seqlen_q * headdim
+    k = batch * nheads_kv * seqlen_k * headdim
+    v = batch * nheads_kv * seqlen_k * headdim_v
+    do = batch * nheads * seqlen_q * headdim_v
+    dq = q
+    dk = k
+    dv = v
+    return (q + k + v + do + dq + dk + dv) * dtype_bytes
+
+
+# ── Reference attention ─────────────────────────────────────────────────────
+
+_attention_ref_mask_cache = {}
+
+
+def attention_ref(q, k, v, causal=False):
+    """Standard attention reference implementation.
+
+    Args:
+        q, k, v: (batch, seqlen, nheads, headdim) tensors.
+        causal: whether to apply causal mask.
+    """
+    softmax_scale = 1.0 / math.sqrt(q.shape[-1])
+    scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k)
+    if causal:
+        if scores.shape[-2] not in _attention_ref_mask_cache:
+            mask = torch.tril(
+                torch.ones(scores.shape[-2:], device=scores.device, dtype=torch.bool), diagonal=0
+            )
+            _attention_ref_mask_cache[scores.shape[-2]] = mask
+        else:
+            mask = _attention_ref_mask_cache[scores.shape[-2]]
+        scores = scores.masked_fill(mask, float("-inf"))
+    attn = torch.softmax(scores, dim=-1)
+    return torch.einsum("bhts,bshd->bthd", attn, v)
+
+
+# ── cuDNN graph helpers ─────────────────────────────────────────────────────
+
+_TORCH_TO_CUDNN_DTYPE = {
+    torch.float16: "HALF",
+    torch.bfloat16: "BFLOAT16",
+    torch.float32: "FLOAT",
+    torch.int32: "INT32",
+    torch.int64: "INT64",
+}
+
+
+def _build_cudnn_graph(io_dtype, tensors, build_fn):
+    """Build a cuDNN graph.  Returns (graph, variant_pack, workspace)."""
+    assert cudnn is not None, "cuDNN is not available"
+    cudnn_dtype = getattr(cudnn.data_type, _TORCH_TO_CUDNN_DTYPE[io_dtype])
+    graph = cudnn.pygraph(
+        io_data_type=cudnn_dtype,
+        intermediate_data_type=cudnn.data_type.FLOAT,
+        compute_data_type=cudnn.data_type.FLOAT,
+    )
+    graph_tensors = {name: graph.tensor_like(t.detach()) for name, t in tensors.items()}
+    variant_pack = build_fn(graph, graph_tensors)
+    graph.validate()
+    graph.build_operation_graph()
+    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
+    graph.check_support()
+    graph.build_plans()
+    workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8)
+    return graph, variant_pack, workspace
+
+
+def cudnn_fwd_setup(q, k, v, causal=False, window_size_left=None):
+    """Build a cuDNN forward SDPA graph.
+
+    Args:
+        q, k, v: (batch, nheads, seqlen, headdim) tensors (cuDNN layout).
+        causal: whether to apply causal mask.
+        window_size_left: sliding window size (None for no window).
+
+    Returns:
+        (fwd_fn, o_gpu, stats_gpu) where fwd_fn is a zero-arg callable.
+    """
+    b, nheads, seqlen_q, headdim = q.shape
+    headdim_v = v.shape[-1]
+    o_gpu = torch.empty(b, nheads, seqlen_q, headdim_v, dtype=q.dtype, device=q.device)
+    stats_gpu = torch.empty(b, nheads, seqlen_q, 1, dtype=torch.float32, device=q.device)
+
+    def build(graph, gt):
+        o, stats = graph.sdpa(
+            name="sdpa",
+            q=gt["q"],
+            k=gt["k"],
+            v=gt["v"],
+            is_inference=False,
+            attn_scale=1.0 / math.sqrt(headdim),
+            use_causal_mask=causal or window_size_left is not None,
+            sliding_window_length=window_size_left
+            if window_size_left is not None and not causal
+            else None,
+        )
+        o.set_output(True).set_dim(o_gpu.shape).set_stride(o_gpu.stride())
+        stats.set_output(True).set_data_type(cudnn.data_type.FLOAT)
+        return {gt["q"]: q, gt["k"]: k, gt["v"]: v, o: o_gpu, stats: stats_gpu}
+
+    graph, variant_pack, workspace = _build_cudnn_graph(q.dtype, {"q": q, "k": k, "v": v}, build)
+
+    def fwd_fn():
+        graph.execute(variant_pack, workspace)
+        return o_gpu
+
+    return fwd_fn, o_gpu, stats_gpu
+
+
+def cudnn_bwd_setup(q, k, v, o, g, lse, causal=False, window_size_left=None):
+    """Build a cuDNN backward SDPA graph.
+
+    Args:
+        q, k, v, o, g, lse: (batch, nheads, seqlen, dim) tensors (cuDNN layout).
+        causal: whether to apply causal mask.
+        window_size_left: sliding window size (None for no window).
+
+    Returns:
+        bwd_fn: zero-arg callable that returns (dq, dk, dv).
+    """
+    headdim = q.shape[-1]
+    dq_gpu, dk_gpu, dv_gpu = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+
+    def build(graph, gt):
+        dq, dk, dv = graph.sdpa_backward(
+            name="sdpa_backward",
+            q=gt["q"],
+            k=gt["k"],
+            v=gt["v"],
+            o=gt["o"],
+            dO=gt["g"],
+            stats=gt["lse"],
+            attn_scale=1.0 / math.sqrt(headdim),
+            use_causal_mask=causal or window_size_left is not None,
+            sliding_window_length=window_size_left
+            if window_size_left is not None and not causal
+            else None,
+            use_deterministic_algorithm=False,
+        )
+        dq.set_output(True).set_dim(dq_gpu.shape).set_stride(dq_gpu.stride())
+        dk.set_output(True).set_dim(dk_gpu.shape).set_stride(dk_gpu.stride())
+        dv.set_output(True).set_dim(dv_gpu.shape).set_stride(dv_gpu.stride())
+        return {
+            gt["q"]: q,
+            gt["k"]: k,
+            gt["v"]: v,
+            gt["o"]: o,
+            gt["g"]: g,
+            gt["lse"]: lse,
+            dq: dq_gpu,
+            dk: dk_gpu,
+            dv: dv_gpu,
+        }
+
+    graph, variant_pack, workspace = _build_cudnn_graph(
+        q.dtype,
+        {"q": q, "k": k, "v": v, "o": o, "g": g, "lse": lse},
+        build,
+    )
+
+    def bwd_fn():
+        graph.execute(variant_pack, workspace)
+        return dq_gpu, dk_gpu, dv_gpu
+
+    return bwd_fn
@@ -0,0 +1,268 @@
+# Copyright (c) 2023, Tri Dao.
+"""Useful functions for writing test code."""
+
+import torch
+import torch.utils.benchmark as benchmark
+
+
+def benchmark_forward(
+    fn, *inputs, repeats=10, desc="", verbose=True, amp=False, amp_dtype=torch.float16, **kwinputs
+):
+    """Use Pytorch Benchmark on the forward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Forward pass")
+
+    def amp_wrapper(*inputs, **kwinputs):
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            fn(*inputs, **kwinputs)
+
+    t = benchmark.Timer(
+        stmt="fn_amp(*inputs, **kwinputs)",
+        globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_backward(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the backward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Backward pass")
+    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+        y = fn(*inputs, **kwinputs)
+        if type(y) is tuple:
+            y = y[0]
+    if grad is None:
+        grad = torch.randn_like(y)
+    else:
+        if grad.shape != y.shape:
+            raise RuntimeError("Grad shape does not match output shape")
+
+    def f(*inputs, y, grad):
+        # Set .grad to None to avoid extra operation of gradient accumulation
+        for x in inputs:
+            if isinstance(x, torch.Tensor):
+                x.grad = None
+        y.backward(grad, retain_graph=True)
+
+    t = benchmark.Timer(
+        stmt="f(*inputs, y=y, grad=grad)",
+        globals={"f": f, "inputs": inputs, "y": y, "grad": grad},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_combined(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Forward + Backward pass")
+    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+        y = fn(*inputs, **kwinputs)
+        if type(y) is tuple:
+            y = y[0]
+    if grad is None:
+        grad = torch.randn_like(y)
+    else:
+        if grad.shape != y.shape:
+            raise RuntimeError("Grad shape does not match output shape")
+
+    def f(grad, *inputs, **kwinputs):
+        for x in inputs:
+            if isinstance(x, torch.Tensor):
+                x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            y = fn(*inputs, **kwinputs)
+            if type(y) is tuple:
+                y = y[0]
+        y.backward(grad, retain_graph=True)
+
+    t = benchmark.Timer(
+        stmt="f(grad, *inputs, **kwinputs)",
+        globals={"f": f, "fn": fn, "inputs": inputs, "grad": grad, "kwinputs": kwinputs},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_fwd_bwd(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    return (
+        benchmark_forward(
+            fn,
+            *inputs,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_backward(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+    )
+
+
+def benchmark_all(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    return (
+        benchmark_forward(
+            fn,
+            *inputs,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_backward(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_combined(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+    )
+
+
+def pytorch_profiler(
+    fn,
+    *inputs,
+    trace_filename=None,
+    backward=False,
+    amp=False,
+    amp_dtype=torch.float16,
+    cpu=False,
+    verbose=True,
+    **kwinputs,
+):
+    """Wrap benchmark functions in Pytorch profiler to see CUDA information."""
+    if backward:
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+            g = torch.randn_like(out)
+    for _ in range(30):  # Warm up
+        if backward:
+            for x in inputs:
+                if isinstance(x, torch.Tensor):
+                    x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+        # Backward should be done outside autocast
+        if backward:
+            out.backward(g, retain_graph=True)
+    activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [
+        torch.profiler.ProfilerActivity.CUDA
+    ]
+    with torch.profiler.profile(
+        activities=activities,
+        record_shapes=True,
+        # profile_memory=True,
+        with_stack=True,
+    ) as prof:
+        if backward:
+            for x in inputs:
+                if isinstance(x, torch.Tensor):
+                    x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+        if backward:
+            out.backward(g, retain_graph=True)
+    if verbose:
+        # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50))
+        print(prof.key_averages().table(row_limit=50))
+    if trace_filename is not None:
+        prof.export_chrome_trace(trace_filename)
+
+
+def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs):
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    fn(*inputs, **kwinputs)
+    torch.cuda.synchronize()
+    mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000)
+    if verbose:
+        print(f"{desc} max memory: {mem}GB")
+    torch.cuda.empty_cache()
+    return mem
@@ -0,0 +1,434 @@
+# Benchmark FP8 attention for FA4 (CuTe-DSL) on SM100.
+#
+# Run (recommended):
+#   python -m flash_attn.cute.benchmark_flash_attention_fp8
+#
+# Notes:
+# - This is intended to be used while bringing up FP8 support for SM100.
+# - FP8 correctness depends on descales + max-offset scaling being implemented in the SM100 kernel.
+#   This script optionally checks output vs a BF16 PyTorch baseline on dequantized FP8 inputs.
+#
+# Adapted from: `hopper/benchmark_flash_attention_fp8.py`
+
+from __future__ import annotations
+
+import argparse
+import inspect
+import math
+import time
+from typing import Iterable
+
+import torch
+from einops import rearrange
+
+from flash_attn.cute.benchmark import benchmark_forward
+from flash_attn.cute.interface import _flash_attn_fwd as flash_attn_cute_fwd
+
+try:
+    import cudnn
+except ImportError:
+    cudnn = None
+
+
+def _torch_float8_dtype(name: str) -> torch.dtype:
+    if name in ("fp8", "fp8_e4m3", "fp8_e4m3fn"):
+        return torch.float8_e4m3fn
+    if name in ("fp8_e5m2", "fp8_e5m2fn"):
+        return torch.float8_e5m2
+    raise ValueError(f"Unsupported fp8 dtype name: {name}")
+
+
+def _parse_int_list(csv: str) -> list[int]:
+    out: list[int] = []
+    for part in csv.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        out.append(int(part))
+    return out
+
+
+def attention_pytorch(qkv: torch.Tensor, causal: bool) -> torch.Tensor:
+    """
+    qkv: (batch, seqlen, 3, nheads, headdim)
+    out: (batch, seqlen, nheads, headdim)
+    """
+    batch_size, seqlen, _, nheads, d = qkv.shape
+    q, k, v = qkv.unbind(dim=2)
+    q = rearrange(q, "b t h d -> (b h) t d")
+    k = rearrange(k, "b s h d -> (b h) d s")
+    softmax_scale = 1.0 / math.sqrt(d)
+    scores = torch.empty(batch_size * nheads, seqlen, seqlen, dtype=qkv.dtype, device=qkv.device)
+    scores = rearrange(
+        torch.baddbmm(scores, q, k, beta=0, alpha=softmax_scale), "(b h) t s -> b h t s", h=nheads
+    )
+    if causal:
+        causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)
+        scores = scores + causal_mask.to(dtype=scores.dtype)
+    attention = torch.softmax(scores, dim=-1)
+    output = torch.einsum("bhts,bshd->bthd", attention, v)
+    return output.to(dtype=qkv.dtype)
+
+
+def flops(batch: int, seqlen: int, headdim: int, nheads: int, causal: bool) -> int:
+    # Matches the hopper benchmark’s convention.
+    return 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1)
+
+
+def efficiency(flop: int, seconds: float) -> float:
+    return (flop / seconds / 1e12) if not math.isnan(seconds) else 0.0
+
+
+def time_fwd(fn, *args, repeats: int, **kwargs) -> float:
+    time.sleep(1)  # reduce residual throttling effects between benchmarks
+    _, m = benchmark_forward(fn, *args, repeats=repeats, verbose=False, **kwargs)
+    return float(m.mean)
+
+
+def convert_to_cudnn_type(torch_type):
+    if torch_type == torch.float16:
+        return cudnn.data_type.HALF
+    if torch_type == torch.bfloat16:
+        return cudnn.data_type.BFLOAT16
+    if torch_type == torch.float32:
+        return cudnn.data_type.FLOAT
+    if torch_type == torch.int32:
+        return cudnn.data_type.INT32
+    if torch_type == torch.int64:
+        return cudnn.data_type.INT64
+    if torch_type == torch.float8_e4m3fn:
+        return cudnn.data_type.FP8_E4M3
+    if torch_type == torch.float8_e5m2:
+        return cudnn.data_type.FP8_E5M2
+    raise ValueError("Unsupported tensor data type.")
+
+
+def cudnn_sdpa_fp8_setup(qkv: torch.Tensor, seqlen_q: int, seqlen_k: int, causal: bool):
+    """Minimal cudnn.fp8 sdpa runner (optional)."""
+    assert cudnn is not None, "cudnn python bindings not available"
+    b, _, _, nheads, headdim = qkv.shape
+    o_gpu = torch.zeros(b, seqlen_q, nheads, headdim, dtype=qkv.dtype, device=qkv.device)
+    o_gpu_transposed = torch.as_strided(
+        o_gpu,
+        [b, nheads, seqlen_q, headdim],
+        [nheads * seqlen_q * headdim, headdim, nheads * headdim, 1],
+    )
+    amax_s_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device)
+    amax_o_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device)
+
+    graph = cudnn.pygraph(
+        io_data_type=convert_to_cudnn_type(qkv.dtype),
+        intermediate_data_type=cudnn.data_type.FLOAT,
+        compute_data_type=cudnn.data_type.FLOAT,
+    )
+
+    new_q = torch.as_strided(
+        qkv,
+        [b, nheads, seqlen_q, headdim],
+        [seqlen_q * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
+        storage_offset=0,
+    )
+    q = graph.tensor(
+        name="Q",
+        dim=list(new_q.shape),
+        stride=list(new_q.stride()),
+        data_type=convert_to_cudnn_type(qkv.dtype),
+    )
+
+    new_k = torch.as_strided(
+        qkv,
+        [b, nheads, seqlen_k, headdim],
+        [seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
+        storage_offset=nheads * headdim,
+    )
+    k = graph.tensor(
+        name="K",
+        dim=list(new_k.shape),
+        stride=list(new_k.stride()),
+        data_type=convert_to_cudnn_type(qkv.dtype),
+    )
+
+    new_v = torch.as_strided(
+        qkv,
+        [b, nheads, seqlen_k, headdim],
+        [seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1],
+        storage_offset=nheads * headdim * 2,
+    )
+    v = graph.tensor(
+        name="V",
+        dim=list(new_v.shape),
+        stride=list(new_v.stride()),
+        data_type=convert_to_cudnn_type(qkv.dtype),
+    )
+
+    def _scale_tensor():
+        return graph.tensor(dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT)
+
+    default_scale_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float32, device="cuda")
+    descale_q = _scale_tensor()
+    descale_k = _scale_tensor()
+    descale_v = _scale_tensor()
+    descale_s = _scale_tensor()
+    scale_s = _scale_tensor()
+    scale_o = _scale_tensor()
+
+    o, _, amax_s, amax_o = graph.sdpa_fp8(
+        q=q,
+        k=k,
+        v=v,
+        descale_q=descale_q,
+        descale_k=descale_k,
+        descale_v=descale_v,
+        descale_s=descale_s,
+        scale_s=scale_s,
+        scale_o=scale_o,
+        is_inference=True,
+        attn_scale=1.0 / math.sqrt(headdim),
+        use_causal_mask=causal,
+        name="sdpa",
+    )
+    o.set_output(True).set_dim(o_gpu_transposed.shape).set_stride(o_gpu_transposed.stride())
+    amax_s.set_output(False).set_dim(amax_s_gpu.shape).set_stride(amax_s_gpu.stride())
+    amax_o.set_output(False).set_dim(amax_o_gpu.shape).set_stride(amax_o_gpu.stride())
+
+    graph.validate()
+    graph.build_operation_graph()
+    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
+    graph.check_support()
+    graph.build_plans()
+
+    variant_pack = {
+        q: new_q,
+        k: new_k,
+        v: new_v,
+        descale_q: default_scale_gpu,
+        descale_k: default_scale_gpu,
+        descale_v: default_scale_gpu,
+        descale_s: default_scale_gpu,
+        scale_s: default_scale_gpu,
+        scale_o: default_scale_gpu,
+        o: o_gpu_transposed,
+        amax_s: amax_s_gpu,
+        amax_o: amax_o_gpu,
+    }
+    workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8)
+
+    def run():
+        graph.execute(variant_pack, workspace)
+        return o_gpu
+
+    return run
+
+
+def _maybe_pass_descales(callable_, **kwargs):
+    sig = inspect.signature(callable_)
+    return {k: v for k, v in kwargs.items() if k in sig.parameters}
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repeats", type=int, default=30)
+    parser.add_argument("--dim", type=int, default=2048)
+    parser.add_argument("--headdims", default="64,128")
+    parser.add_argument("--dtype", default="fp8_e4m3fn")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--check",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Enable correctness checks vs BF16 PyTorch baseline.",
+    )
+    parser.add_argument(
+        "--check-quantization-only",
+        action="store_true",
+        help="Check FP8 kernel vs dequantized-FP8 baseline (quantization error only).",
+    )
+    parser.add_argument("--atol-bf16", type=float, default=0.10)
+    parser.add_argument("--rtol-bf16", type=float, default=0.10)
+    parser.add_argument("--atol-fp8", type=float, default=0.50)
+    parser.add_argument("--rtol-fp8", type=float, default=0.50)
+    parser.add_argument("--run-cudnn", action="store_true")
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    major, minor = torch.cuda.get_device_capability()
+    if major != 10:
+        raise RuntimeError(
+            f"This benchmark is for SM100 (compute capability 10.x). Got {major}.{minor}."
+        )
+
+    torch.manual_seed(args.seed)
+    device = "cuda"
+    fp8_dtype = _torch_float8_dtype(args.dtype)
+    headdim_vals = _parse_int_list(args.headdims)
+    bs_seqlen_vals = [(32, 512), (16, 1024), (8, 2048), (4, 4096), (2, 8192), (1, 16384)]
+
+    methods = ["Pytorch", "FA4-CuTe-BF16", "FA4-CuTe-FP8"] + (
+        ["cuDNN-FP8"] if args.run_cudnn and cudnn is not None else []
+    )
+
+    fp8_failures = []
+
+    for headdim in headdim_vals:
+        for causal in (False, True):
+            for batch, seqlen in bs_seqlen_vals:
+                torch.cuda.empty_cache()
+                nheads = args.dim // headdim
+                if args.dim % headdim != 0:
+                    raise ValueError(f"--dim must be divisible by headdim ({args.dim=} {headdim=})")
+
+                q_bf16 = torch.randn(
+                    batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
+                )
+                k_bf16 = torch.randn(
+                    batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
+                )
+                v_bf16 = torch.randn(
+                    batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16
+                )
+                qkv_bf16 = torch.stack([q_bf16, k_bf16, v_bf16], dim=2)
+
+                times = {}
+                speeds = {}
+
+                out_ref_bf16 = None
+                try:
+                    out_ref_bf16 = attention_pytorch(qkv_bf16, causal=causal)  # warmup / reference
+                    t = time_fwd(attention_pytorch, qkv_bf16, causal=causal, repeats=args.repeats)
+                    times["Pytorch"] = t
+                except RuntimeError as e:
+                    if "out of memory" in str(e).lower():
+                        times["Pytorch"] = float("nan")
+                        out_ref_bf16 = None
+                    else:
+                        raise
+
+                # FA4 / CuTe BF16 baseline
+                try:
+                    softmax_scale = headdim**-0.5
+                    out_fa4_bf16, *_ = flash_attn_cute_fwd(
+                        q_bf16, k_bf16, v_bf16, softmax_scale=softmax_scale, causal=causal
+                    )  # warmup / compile
+                    t = time_fwd(
+                        flash_attn_cute_fwd,
+                        q_bf16,
+                        k_bf16,
+                        v_bf16,
+                        softmax_scale=softmax_scale,
+                        causal=causal,
+                        repeats=args.repeats,
+                    )
+                    times["FA4-CuTe-BF16"] = t
+                    if args.check and out_ref_bf16 is not None:
+                        torch.testing.assert_close(
+                            out_fa4_bf16,
+                            out_ref_bf16,
+                            atol=args.atol_bf16,
+                            rtol=args.rtol_bf16,
+                        )
+                except Exception as e:
+                    # Treat as fatal: BF16 kernel should be usable for basic sanity checking.
+                    raise RuntimeError("FA4-CuTe BF16 baseline failed") from e
+
+                # FA4 / CuTe FP8
+                q_fp8 = q_bf16.to(fp8_dtype)
+                k_fp8 = k_bf16.to(fp8_dtype)
+                v_fp8 = v_bf16.to(fp8_dtype)
+
+                # Placeholder descales (FA3-style: per-(batch, kv_head)).
+                q_descale = torch.ones(batch, nheads, device=device, dtype=torch.float32)
+                k_descale = torch.ones(batch, nheads, device=device, dtype=torch.float32)
+                v_descale = torch.ones(batch, nheads, device=device, dtype=torch.float32)
+
+                # Optional: FP8 reference baseline (dequantized FP8 -> PyTorch) for quantization-error-only checks
+                out_ref_fp8 = None
+                if args.check and args.check_quantization_only:
+                    try:
+                        # Dequantize FP8 inputs back to BF16 (applying descales)
+                        q_ref_fp8 = (q_fp8.to(torch.bfloat16) * q_descale[:, None, :, None]).to(
+                            torch.bfloat16
+                        )
+                        k_ref_fp8 = (k_fp8.to(torch.bfloat16) * k_descale[:, None, :, None]).to(
+                            torch.bfloat16
+                        )
+                        v_ref_fp8 = (v_fp8.to(torch.bfloat16) * v_descale[:, None, :, None]).to(
+                            torch.bfloat16
+                        )
+                        qkv_ref_fp8 = torch.stack([q_ref_fp8, k_ref_fp8, v_ref_fp8], dim=2)
+                        out_ref_fp8 = attention_pytorch(qkv_ref_fp8, causal=causal)
+                    except RuntimeError as e:
+                        if "out of memory" in str(e).lower():
+                            out_ref_fp8 = None
+                        else:
+                            raise
+
+                fa4_kwargs = dict(softmax_scale=softmax_scale, causal=causal)
+                fa4_kwargs.update(
+                    _maybe_pass_descales(
+                        flash_attn_cute_fwd,
+                        q_descale=q_descale,
+                        k_descale=k_descale,
+                        v_descale=v_descale,
+                    )
+                )
+
+                try:
+                    # Warmup/compile (will raise until FP8 is implemented)
+                    out_fa4_fp8, *_ = flash_attn_cute_fwd(q_fp8, k_fp8, v_fp8, **fa4_kwargs)
+                    t = time_fwd(
+                        flash_attn_cute_fwd,
+                        q_fp8,
+                        k_fp8,
+                        v_fp8,
+                        repeats=args.repeats,
+                        **fa4_kwargs,
+                    )
+                    times["FA4-CuTe-FP8"] = t
+                    if args.check:
+                        # Choose baseline: quantization-only (dequantized FP8) or full (BF16)
+                        if args.check_quantization_only:
+                            ref_baseline = out_ref_fp8
+                        else:
+                            ref_baseline = out_ref_bf16
+
+                        if ref_baseline is not None:
+                            torch.testing.assert_close(
+                                out_fa4_fp8,
+                                ref_baseline,
+                                atol=args.atol_fp8,
+                                rtol=args.rtol_fp8,
+                            )
+                except Exception as e:
+                    fp8_failures.append((causal, headdim, batch, seqlen, repr(e)))
+                    times["FA4-CuTe-FP8"] = float("nan")
+
+                if args.run_cudnn and cudnn is not None:
+                    qkv_fp8 = qkv_bf16.to(fp8_dtype)
+                    runner = cudnn_sdpa_fp8_setup(qkv_fp8, seqlen, seqlen, causal=causal)
+                    _ = runner()  # warmup
+                    t = time_fwd(lambda: runner(), repeats=args.repeats)
+                    times["cuDNN-FP8"] = t
+
+                print(f"### causal={causal}, headdim={headdim}, batch={batch}, seqlen={seqlen} ###")
+                for method in methods:
+                    t = times.get(method, float("nan"))
+                    speeds[method] = efficiency(flops(batch, seqlen, headdim, nheads, causal), t)
+                    if math.isnan(t):
+                        print(f"{method} fwd: (skipped)")
+                    else:
+                        print(f"{method} fwd: {speeds[method]:.2f} TFLOPs/s, {t * 1e3:.3f} ms")
+                if math.isnan(times.get("FA4-CuTe-FP8", float("nan"))):
+                    print("FA4-CuTe-FP8 status: FAILED")
+
+    if fp8_failures:
+        print(f"\nFP8 failures: {len(fp8_failures)} (showing first 5)")
+        for causal, headdim, batch, seqlen, err in fp8_failures[:5]:
+            print(f"- causal={causal} headdim={headdim} batch={batch} seqlen={seqlen}: {err}")
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
@@ -0,0 +1,156 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+from typing import Tuple, Optional
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+
+from flash_attn.cute.seqlen_info import SeqlenInfoQK, SeqlenInfoQKNewK
+
+
+@dataclass(frozen=True)
+class BlockInfo:
+    tile_m: cutlass.Constexpr[int]
+    tile_n: cutlass.Constexpr[int]
+    is_causal: cutlass.Constexpr[bool]
+    is_local: cutlass.Constexpr[bool] = False
+    is_split_kv: cutlass.Constexpr[bool] = False
+    window_size_left: Optional[Int32] = None
+    window_size_right: Optional[Int32] = None
+    qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1
+
+    @cute.jit
+    def get_n_block_min_max(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+        split_idx: Int32 = 0,
+        num_splits: Int32 = 1,
+    ) -> Tuple[Int32, Int32]:
+        n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n)
+        if const_expr(self.is_causal or (self.is_local and self.window_size_right is not None)):
+            m_idx_max = (m_block + 1) * self.tile_m
+            if const_expr(self.qhead_per_kvhead_packgqa > 1):
+                m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa)
+            n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            n_idx_right = n_idx if const_expr(self.is_causal) else n_idx + self.window_size_right
+            n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, self.tile_n))
+        n_block_min = 0
+        if const_expr(self.is_local and self.window_size_left is not None):
+            m_idx_min = m_block * self.tile_m
+            if const_expr(self.qhead_per_kvhead_packgqa > 1):
+                m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa
+            n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            n_idx_left = n_idx - self.window_size_left
+            n_block_min = cutlass.max(n_idx_left // self.tile_n, 0)
+        if cutlass.const_expr(self.is_split_kv):
+            num_n_blocks_per_split = (
+                Int32(0)
+                if n_block_max <= n_block_min
+                else (n_block_max - n_block_min + num_splits - 1) // num_splits
+            )
+            n_block_min = n_block_min + split_idx * num_n_blocks_per_split
+            n_block_max = cutlass.min(n_block_min + num_n_blocks_per_split, n_block_max)
+        return n_block_min, n_block_max
+
+    @cute.jit
+    def get_m_block_min_max(self, seqlen_info: SeqlenInfoQK, n_block: Int32) -> Tuple[Int32, Int32]:
+        m_block_max = cute.ceil_div(seqlen_info.seqlen_q, self.tile_m)
+        m_block_min = 0
+        if const_expr(self.is_causal or (self.is_local and self.window_size_right is not None)):
+            n_idx_min = n_block * self.tile_n
+            m_idx = n_idx_min + seqlen_info.seqlen_q - seqlen_info.seqlen_k
+            m_idx_right = m_idx if const_expr(self.is_causal) else m_idx - self.window_size_right
+            m_block_min = max(m_block_min, m_idx_right // self.tile_m)
+        if const_expr(self.is_local and self.window_size_left is not None):
+            n_idx_max = (n_block + 1) * self.tile_n
+            m_idx = n_idx_max + seqlen_info.seqlen_q - seqlen_info.seqlen_k
+            m_idx_left = m_idx + self.window_size_left
+            m_block_max = min(m_block_max, cute.ceil_div(m_idx_left, self.tile_m))
+        return m_block_min, m_block_max
+
+    @cute.jit
+    def get_n_block_k_new_min_max(
+        self,
+        seqlen_info: SeqlenInfoQKNewK,
+        m_block: Int32,
+        split_idx: Int32 = 0,
+        num_splits: Int32 = 1,
+    ) -> Tuple[Int32, Int32]:
+        """Get the block range for new K tokens (append KV).
+
+        First computes the full n_block range via get_n_block_min_max, then maps
+        those blocks into the new-K index space by subtracting seqlen_k_og.
+        """
+        n_block_min, n_block_max = self.get_n_block_min_max(
+            seqlen_info,
+            m_block,
+            split_idx,
+            num_splits,
+        )
+        idx_k_new_min = cutlass.max(n_block_min * self.tile_n - seqlen_info.seqlen_k_og, 0)
+        idx_k_new_max = cutlass.min(
+            n_block_max * self.tile_n - seqlen_info.seqlen_k_og, seqlen_info.seqlen_k_new
+        )
+        n_block_new_min = idx_k_new_min // self.tile_n
+        n_block_new_max = (
+            cute.ceil_div(idx_k_new_max, self.tile_n)
+            if idx_k_new_max > idx_k_new_min
+            else n_block_new_min
+        )
+        return n_block_new_min, n_block_new_max
+
+    @cute.jit
+    def get_n_block_min_causal_local_mask(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+        n_block_min: Int32,
+    ) -> Int32:
+        """If we have separate iterations with causal or local masking at the start, where do we stop"""
+        m_idx_min = m_block * self.tile_m
+        if const_expr(self.qhead_per_kvhead_packgqa > 1):
+            m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa
+        n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+        n_idx_right = (
+            n_idx
+            if const_expr(not self.is_local or self.window_size_right is None)
+            else n_idx + self.window_size_right
+        )
+        return cutlass.max(n_block_min, n_idx_right // self.tile_n)
+
+    @cute.jit
+    def get_n_block_min_before_local_mask(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+        n_block_min: Int32,
+    ) -> Int32:
+        """If we have separate iterations with local masking at the end, where do we stop the non-masked iterations"""
+        if const_expr(not self.is_local or self.window_size_left is None):
+            return n_block_min
+        else:
+            m_idx_max = (m_block + 1) * self.tile_m
+            if const_expr(self.qhead_per_kvhead_packgqa > 1):
+                m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa)
+            n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            n_idx_left = n_idx - self.window_size_left
+            return cutlass.max(n_block_min, cute.ceil_div(n_idx_left, self.tile_n))
+
+    @cute.jit
+    def get_n_block_max_for_m_block(
+        self,
+        seqlen_info: SeqlenInfoQK,
+        m_block: Int32,
+    ) -> Int32:
+        n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n)
+        if const_expr(self.is_causal or self.window_size_right is not None):
+            m_idx_max = (m_block + 1) * self.tile_m
+            if const_expr(self.qhead_per_kvhead_packgqa > 1):
+                m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa)
+            n_idx_right = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q
+            if const_expr(self.window_size_right is not None):
+                n_idx_right += self.window_size_right
+            n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, self.tile_n))
+        return n_block_max
@@ -0,0 +1,676 @@
+"""
+Block-sparsity utilities for FlexAttention
+"""
+
+from typing import Callable, NamedTuple, Tuple
+
+import cutlass.cute as cute
+import torch
+
+from flash_attn.cute.cute_dsl_utils import get_broadcast_dims, to_cute_tensor
+
+
+def ceildiv(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
+class BlockSparseTensors(NamedTuple):
+    mask_block_cnt: cute.Tensor
+    mask_block_idx: cute.Tensor
+    full_block_cnt: cute.Tensor | None = None
+    full_block_idx: cute.Tensor | None = None
+    cu_total_m_blocks: cute.Tensor | None = None
+    cu_block_idx_offsets: cute.Tensor | None = None
+    dq_write_order: cute.Tensor | None = None
+    dq_write_order_full: cute.Tensor | None = None
+
+    def __new_from_mlir_values__(self, values):
+        new_fields = []
+        idx = 0
+        for original in self:
+            if original is None:
+                new_fields.append(None)
+            else:
+                new_fields.append(values[idx])
+                idx += 1
+        return BlockSparseTensors(*new_fields)
+
+
+class BlockSparseTensorsTorch(NamedTuple):
+    mask_block_cnt: torch.Tensor
+    mask_block_idx: torch.Tensor
+    full_block_cnt: torch.Tensor | None = None
+    full_block_idx: torch.Tensor | None = None
+    cu_total_m_blocks: torch.Tensor | None = None
+    cu_block_idx_offsets: torch.Tensor | None = None
+    block_size: tuple[int, int] | None = None
+    dq_write_order: torch.Tensor | None = None
+    dq_write_order_full: torch.Tensor | None = None
+    spt: bool | None = None
+
+
+def _ordered_to_dense_simple(
+    num_blocks: torch.Tensor,
+    indices: torch.Tensor,
+    num_cols: int,
+) -> torch.Tensor:
+    """Convert ordered sparse representation to dense binary matrix.
+
+    Args:
+        num_blocks: [B, H, num_rows] count of valid entries per row
+        indices: [B, H, num_rows, max_entries] column indices (valid entries packed left)
+        num_cols: total number of columns
+
+    Returns:
+        dense: [B, H, num_rows, num_cols] binary int32 matrix
+    """
+    B, H, num_rows, max_entries = indices.shape
+    device = indices.device
+    dense = torch.zeros(B, H, num_rows, num_cols + 1, dtype=torch.int32, device=device)
+    col_range = torch.arange(max_entries, device=device)
+    valid = col_range[None, None, None, :] < num_blocks[:, :, :, None]
+    safe_indices = torch.where(valid, indices.long(), num_cols)
+    row_idx = torch.arange(num_rows, device=device)[None, None, :, None].expand_as(indices)
+    b_idx = torch.arange(B, device=device)[:, None, None, None].expand_as(indices)
+    h_idx = torch.arange(H, device=device)[None, :, None, None].expand_as(indices)
+    dense[b_idx, h_idx, row_idx, safe_indices] = 1
+    return dense[:, :, :, :num_cols]
+
+
+def compute_dq_write_order(
+    fwd_mask_cnt: torch.Tensor,
+    fwd_mask_idx: torch.Tensor,
+    fwd_full_cnt: torch.Tensor | None,
+    fwd_full_idx: torch.Tensor | None,
+    bwd_mask_cnt: torch.Tensor,
+    bwd_mask_idx: torch.Tensor,
+    bwd_full_cnt: torch.Tensor | None,
+    bwd_full_idx: torch.Tensor | None,
+    spt: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Compute dQ write-order metadata for deterministic block-sparse backward.
+
+    For each (n_block, i) in the backward iteration, computes the semaphore
+    lock value: the rank of n_block in the combined (partial + full) sorted
+    contributor list for the target m_block.
+
+    Lock values are assigned in ascending n_block order (or descending if spt=True)
+    to guarantee deadlock-freedom with the CTA scheduling order.
+
+    Args:
+        fwd_mask_cnt: [B, H, num_m_blocks] partial contributor counts per m_block
+        fwd_mask_idx: [B, H, num_m_blocks, max_kv] partial contributor n_block indices (ascending)
+        fwd_full_cnt: [B, H, num_m_blocks] full contributor counts per m_block (optional)
+        fwd_full_idx: [B, H, num_m_blocks, max_kv] full contributor n_block indices (optional)
+        bwd_mask_cnt: [B, H, num_n_blocks] partial iteration counts per n_block
+        bwd_mask_idx: [B, H, num_n_blocks, max_q] partial iteration m_block indices
+        bwd_full_cnt: [B, H, num_n_blocks] full iteration counts per n_block (optional)
+        bwd_full_idx: [B, H, num_n_blocks, max_q] full iteration m_block indices (optional)
+        spt: if True, reverse ordering (highest n_block gets lock_value=0)
+
+    Returns:
+        (dq_write_order, dq_write_order_full): tensors parallel to bwd_mask_idx
+        and bwd_full_idx respectively, containing lock values.
+    """
+    device = fwd_mask_idx.device
+    B, H, num_m, max_kv_partial = fwd_mask_idx.shape
+    _, _, num_n, max_q_partial = bwd_mask_idx.shape
+
+    has_full = fwd_full_cnt is not None and fwd_full_idx is not None
+
+    dense_partial = _ordered_to_dense_simple(fwd_mask_cnt, fwd_mask_idx, num_n)
+    if has_full:
+        dense_full = _ordered_to_dense_simple(fwd_full_cnt, fwd_full_idx, num_n)
+        dense = (dense_partial + dense_full).clamp(max=1)
+    else:
+        dense = dense_partial
+
+    cumsum = dense.cumsum(dim=-1)
+    rank_table = (cumsum - dense).to(torch.int32)
+
+    if spt:
+        total_per_m = cumsum[:, :, :, -1:]
+        rank_table = (total_per_m - 1 - rank_table).to(torch.int32)
+
+    def _gather_write_order(bwd_idx, bwd_cnt):
+        b_i = torch.arange(B, device=device)[:, None, None, None].expand_as(bwd_idx)
+        h_i = torch.arange(H, device=device)[None, :, None, None].expand_as(bwd_idx)
+        n_i = torch.arange(bwd_idx.shape[2], device=device)[None, None, :, None].expand_as(bwd_idx)
+        m_vals = bwd_idx.long().clamp(0, num_m - 1)
+        return rank_table[b_i, h_i, m_vals, n_i].to(torch.int32)
+
+    dq_write_order = _gather_write_order(bwd_mask_idx, bwd_mask_cnt)
+
+    dq_write_order_full = None
+    if has_full and bwd_full_cnt is not None and bwd_full_idx is not None:
+        dq_write_order_full = _gather_write_order(bwd_full_idx, bwd_full_cnt)
+
+    return dq_write_order, dq_write_order_full
+
+
+def compute_dq_write_order_from_block_mask(
+    block_mask,
+    spt: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    (
+        _seq_q,
+        _seq_k,
+        kv_mask_cnt,
+        kv_mask_idx,
+        full_kv_cnt,
+        full_kv_idx,
+        q_mask_cnt,
+        q_mask_idx,
+        full_q_cnt,
+        full_q_idx,
+        *_,
+    ) = block_mask.as_tuple()
+    return compute_dq_write_order(
+        kv_mask_cnt,
+        kv_mask_idx,
+        full_kv_cnt,
+        full_kv_idx,
+        q_mask_cnt,
+        q_mask_idx,
+        full_q_cnt,
+        full_q_idx,
+        spt=spt,
+    )
+
+
+def get_sparse_q_block_size(
+    tensors: BlockSparseTensorsTorch | None,
+    seqlen_q: int,
+) -> int | None:
+    """Return the Q sparse block size, or None when sparsity is unset or ambiguous."""
+    if tensors is None:
+        return None
+    if tensors.block_size is not None:
+        return tensors.block_size[0]
+    num_m_blocks = tensors.mask_block_idx.shape[2]
+    min_block_size = ceildiv(seqlen_q, num_m_blocks)
+    max_block_size = seqlen_q if num_m_blocks == 1 else (seqlen_q - 1) // (num_m_blocks - 1)
+    if min_block_size != max_block_size:
+        return None
+    return min_block_size
+
+
+def _expand_sparsity_tensor(
+    tensor: torch.Tensor,
+    expected_shape: Tuple[int, ...],
+    tensor_name: str,
+    context: str | None,
+    hint: str | Callable[[], str] | None,
+) -> torch.Tensor:
+    """Check if we need to expand the tensor to expected shape, and do so if possible."""
+    needs_expand = tensor.shape != expected_shape
+    if not needs_expand:
+        return tensor
+    can_expand = all(map(lambda cur, tgt: cur == tgt or cur == 1, tensor.shape, expected_shape))
+    if not can_expand:
+        context_clause = f" ({context})" if context else ""
+        resolved_hint = hint() if callable(hint) else hint
+        hint_clause = f" Hint: {resolved_hint}" if resolved_hint else ""
+        raise ValueError(
+            f"{tensor_name}{context_clause} with shape {tensor.shape} cannot be expanded to expected shape {expected_shape}."
+            f"{hint_clause}"
+        )
+    return tensor.expand(*expected_shape)
+
+
+def _check_and_expand_block(
+    name: str,
+    cnt: torch.Tensor | None,
+    idx: torch.Tensor | None,
+    expected_count_shape: Tuple[int, ...],
+    expected_index_shape: Tuple[int, ...],
+    context: str | None,
+    hint: str | Callable[[], str] | None,
+) -> Tuple[torch.Tensor | None, torch.Tensor | None]:
+    if (cnt is None) != (idx is None):
+        raise ValueError(
+            f"{name}_block_cnt and {name}_block_idx must both be provided or both be None"
+        )
+    if cnt is None or idx is None:
+        return None, None
+    if cnt.dtype != torch.int32 or idx.dtype != torch.int32:
+        raise ValueError(f"{name}_block tensors must have dtype torch.int32")
+    if cnt.device != idx.device:
+        raise ValueError(f"{name}_block_cnt and {name}_block_idx must be on the same device")
+    if not cnt.is_cuda or not idx.is_cuda:
+        raise ValueError(f"{name}_block tensors must live on CUDA")
+    expanded_cnt = _expand_sparsity_tensor(
+        cnt, expected_count_shape, f"{name}_block_cnt", context, hint
+    )
+    # [Note] Allow Compact block sparse indices
+    # Allow the last dimension (n_blocks) of idx to be <= expected, since
+    # FA4 only accesses indices 0..cnt-1 per query tile. This enables compact
+    # index tensors that avoid O(N^2) memory at long sequence lengths.
+    if idx.ndim == 4 and idx.shape[3] <= expected_index_shape[3]:
+        expected_index_shape = (*expected_index_shape[:3], idx.shape[3])
+    expanded_idx = _expand_sparsity_tensor(
+        idx, expected_index_shape, f"{name}_block_idx", context, hint
+    )
+    return expanded_cnt, expanded_idx
+
+
+def _check_and_expand_metadata_tensor(
+    name: str,
+    tensor: torch.Tensor | None,
+    expected_shape: Tuple[int, ...],
+    context: str | None,
+    hint: str | Callable[[], str] | None,
+    device: torch.device,
+) -> torch.Tensor | None:
+    if tensor is None:
+        return None
+    if tensor.dtype != torch.int32:
+        raise ValueError(f"{name} must have dtype torch.int32")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on the same device as block sparse tensors")
+    if not tensor.is_cuda:
+        raise ValueError(f"{name} must live on CUDA")
+    return _expand_sparsity_tensor(tensor, expected_shape, name, context, hint)
+
+
+def get_block_sparse_expected_shapes(
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    m_block_size: int,
+    n_block_size: int,
+    q_stage: int,
+) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]:
+    """Return (expected_count_shape, expected_index_shape) for block sparse normalization."""
+    m_block_size_effective = q_stage * m_block_size
+    expected_m_blocks = ceildiv(seqlen_q, m_block_size_effective)
+    expected_n_blocks = ceildiv(seqlen_k, n_block_size)
+    expected_count_shape = (batch_size, num_head, expected_m_blocks)
+    expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks)
+    return expected_count_shape, expected_index_shape
+
+
+def infer_block_sparse_expected_shapes(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    m_block_size: int,
+    n_block_size: int,
+    q_stage: int,
+    context: str,
+    sparse_block_size_q: int | None = None,
+    sparse_block_size_kv: int | None = None,
+) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int], int]:
+    """Infer shapes and scaling for block-sparse tensors.
+
+    Expectations:
+    - mask_block_cnt is (B, H, M) and mask_block_idx is (B, H, M, N).
+    - Batch/head dims may be 1 for broadcast, or match the requested sizes.
+    - sparse_block_size_kv must match tile_n.
+    - sparse_block_size_q must be a multiple of q_stage * tile_m.
+    - If sparse_block_size_q is omitted and seqlen_q/num_m_blocks is ambiguous,
+      the caller must provide block_size to disambiguate. TODO will make this required in a future PR.
+    """
+    base_m_block = q_stage * m_block_size
+    base_n_block = n_block_size
+    if sparse_block_size_kv is None:
+        sparse_block_size_kv = base_n_block
+    if sparse_block_size_kv != base_n_block:
+        raise ValueError(f"Block sparse tensors{context} require BLOCK_SIZE_KV={base_n_block}.")
+    if tensors.mask_block_idx is None:
+        raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.")
+    num_m_blocks = tensors.mask_block_idx.shape[2]
+
+    if sparse_block_size_q is None:
+        sparse_block_size_q = get_sparse_q_block_size(tensors, seqlen_q)
+        if sparse_block_size_q is None and base_m_block != 1:
+            raise ValueError(
+                f"Block sparse tensors{context} require explicit sparse_block_size[0] "
+                f"to disambiguate block size for seqlen_q={seqlen_q} and num_m_blocks={num_m_blocks}."
+            )
+        if sparse_block_size_q is None:
+            sparse_block_size_q = ceildiv(seqlen_q, num_m_blocks)
+
+    if sparse_block_size_q % base_m_block != 0:
+        raise ValueError(
+            f"Block sparse tensors{context} have block size {sparse_block_size_q}, "
+            f"which must be a multiple of {base_m_block}."
+        )
+
+    expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)
+    expected_n_blocks = ceildiv(seqlen_k, sparse_block_size_kv)
+    q_subtile_factor = sparse_block_size_q // base_m_block
+    expected_count_shape = (batch_size, num_head, expected_m_blocks)
+    expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks)
+
+    mask_block_cnt = tensors.mask_block_cnt
+    mask_block_idx = tensors.mask_block_idx
+    if mask_block_cnt is None or mask_block_idx is None:
+        raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.")
+    if mask_block_cnt.ndim != 3 or mask_block_idx.ndim != 4:
+        raise ValueError(
+            f"Block sparse tensors{context} must have shapes (B, H, M) and (B, H, M, N)."
+        )
+    for dim_name, cur, tgt in (
+        ("batch", mask_block_cnt.shape[0], expected_count_shape[0]),
+        ("head", mask_block_cnt.shape[1], expected_count_shape[1]),
+    ):
+        if cur != tgt and cur != 1:
+            raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.")
+    for dim_name, cur, tgt in (
+        ("batch", mask_block_idx.shape[0], expected_index_shape[0]),
+        ("head", mask_block_idx.shape[1], expected_index_shape[1]),
+    ):
+        if cur != tgt and cur != 1:
+            raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.")
+    if mask_block_cnt.shape[2] != mask_block_idx.shape[2]:
+        raise ValueError(f"Block sparse tensors{context} must share the same m-block dimension.")
+    # [Note] Allow Compact block sparse indices: FA4 only accesses indices 0..cnt-1
+    # per query tile, so idx.shape[3] can be <= expected_n_blocks.
+    if mask_block_idx.shape[3] > expected_n_blocks:
+        raise ValueError(
+            f"Block sparse tensors{context} n-block dimension must be <= {expected_n_blocks}."
+        )
+    if expected_m_blocks != num_m_blocks:
+        raise ValueError(
+            f"Block sparse tensors{context} m-block dimension {num_m_blocks} does not match "
+            f"sparse_block_size_q={sparse_block_size_q}. "
+            f"Set BlockSparseTensorsTorch.block_size to match the BlockMask BLOCK_SIZE."
+        )
+    return expected_count_shape, expected_index_shape, q_subtile_factor
+
+
+def get_block_sparse_expected_shapes_bwd(
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    m_block_size: int,
+    n_block_size: int,
+    q_subtile_factor: int,
+) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]:
+    """Return (expected_count_shape, expected_index_shape) for backward block sparse normalization.
+
+    Backward uses Q-direction indexing (transposed from forward), where shapes are
+    indexed by N-blocks first, then M-blocks. The sparse_block_size_q is determined
+    by q_subtile_factor * m_block_size.
+    """
+    sparse_block_size_q = q_subtile_factor * m_block_size
+    expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)
+    expected_n_blocks = ceildiv(seqlen_k, n_block_size)
+    expected_count_shape = (batch_size, num_head, expected_n_blocks)
+    expected_index_shape = (batch_size, num_head, expected_n_blocks, expected_m_blocks)
+    return expected_count_shape, expected_index_shape
+
+
+def normalize_block_sparse_tensors(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    expected_count_shape: Tuple[int, ...],
+    expected_index_shape: Tuple[int, ...],
+    context: str | None = None,
+    hint: str | Callable[[], str] | None = None,
+) -> BlockSparseTensorsTorch:
+    if tensors.mask_block_cnt is None or tensors.mask_block_idx is None:
+        raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.")
+
+    mask_cnt, mask_idx = _check_and_expand_block(
+        "mask",
+        tensors.mask_block_cnt,
+        tensors.mask_block_idx,
+        expected_count_shape,
+        expected_index_shape,
+        context,
+        hint,
+    )
+    if mask_cnt is None or mask_idx is None:
+        raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.")
+
+    full_cnt, full_idx = _check_and_expand_block(
+        "full",
+        tensors.full_block_cnt,
+        tensors.full_block_idx,
+        expected_count_shape,
+        expected_index_shape,
+        context,
+        hint,
+    )
+    if full_cnt is not None and mask_cnt.device != full_cnt.device:
+        raise ValueError("All block sparse tensors must be on the same device")
+
+    dq_write_order = _check_and_expand_metadata_tensor(
+        "dq_write_order",
+        tensors.dq_write_order,
+        tuple(mask_idx.shape),
+        context,
+        hint,
+        mask_cnt.device,
+    )
+    dq_write_order_full = _check_and_expand_metadata_tensor(
+        "dq_write_order_full",
+        tensors.dq_write_order_full,
+        tuple(full_idx.shape) if full_idx is not None else expected_index_shape,
+        context,
+        hint,
+        mask_cnt.device,
+    )
+    spt = tensors.spt
+    if spt is not None and not isinstance(spt, bool):
+        raise ValueError("spt must be a bool when provided")
+    if spt is not None and dq_write_order is None:
+        raise ValueError("spt requires dq_write_order to be provided")
+
+    return BlockSparseTensorsTorch(
+        mask_block_cnt=mask_cnt,
+        mask_block_idx=mask_idx,
+        full_block_cnt=full_cnt,
+        full_block_idx=full_idx,
+        cu_total_m_blocks=tensors.cu_total_m_blocks,
+        cu_block_idx_offsets=tensors.cu_block_idx_offsets,
+        block_size=tensors.block_size,
+        dq_write_order=dq_write_order,
+        dq_write_order_full=dq_write_order_full,
+        spt=spt,
+    )
+
+
+def is_block_sparsity_enabled(tensors: BlockSparseTensorsTorch) -> bool:
+    return any(t is not None for t in (tensors.full_block_cnt, tensors.mask_block_cnt))
+
+
+def get_block_sparse_broadcast_pattern(
+    tensors: BlockSparseTensorsTorch,
+) -> Tuple[Tuple[bool, ...], ...] | None:
+    """Return broadcast pattern for block sparse tensors by checking actual strides.
+
+    Returns a tuple of broadcast patterns (one per tensor) where each pattern
+    is a tuple of bools indicating which dims have stride=0.
+    This is used in compile keys to ensure kernels are recompiled when
+    broadcast patterns change, since CuTe's mark_layout_dynamic() keeps
+    stride=0 as static.
+
+    The tensors should already be expanded/normalized before calling this function.
+
+    Returns None if block sparsity is not enabled.
+    """
+    if not is_block_sparsity_enabled(tensors):
+        return None
+
+    patterns = []
+    for tensor in (
+        tensors.mask_block_cnt,
+        tensors.mask_block_idx,
+        tensors.full_block_cnt,
+        tensors.full_block_idx,
+        tensors.dq_write_order,
+        tensors.dq_write_order_full,
+    ):
+        if tensor is not None:
+            patterns.append(get_broadcast_dims(tensor))
+        else:
+            patterns.append(None)
+    return tuple(patterns)
+
+
+def normalize_block_sparse_config(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    block_size: tuple[int, int],
+    q_stage: int,
+) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None, int]:
+    """Validate the block-sparse config, infer expected shapes, and normalize.
+
+    Handles both fixed-length (3D `[B, H, M]` / 4D `[B, H, M, N]`) and varlen
+    (2D `[H, total_m_blocks]` / `[H, total_n_blocks]`) layouts. Varlen is
+    detected by `tensors.cu_total_m_blocks is not None` and forces
+    `q_subtile_factor == 1` (TODO: potentially remove this restriction).
+    """
+    m_block_size, n_block_size = block_size
+    if tensors.block_size is None:
+        sparse_block_size_q, sparse_block_size_kv = None, n_block_size
+    else:
+        sparse_block_size_q, sparse_block_size_kv = tensors.block_size
+    if sparse_block_size_kv != n_block_size:
+        raise ValueError(
+            f"Block sparsity requires sparse_block_size[1]={n_block_size} to match tile_n."
+        )
+    if tensors.cu_total_m_blocks is not None:
+        base_m_block = q_stage * m_block_size
+        if sparse_block_size_q is not None and sparse_block_size_q != base_m_block:
+            raise ValueError(
+                f"Varlen block sparsity requires sparse_block_size[0]={base_m_block} "
+                f"(= q_stage * tile_m); got {sparse_block_size_q}."
+            )
+        total_m_blocks = tensors.mask_block_cnt.shape[-1]
+        total_n_blocks = tensors.mask_block_idx.shape[-1]
+        expected_count_shape = (num_head, total_m_blocks)
+        expected_index_shape = (num_head, total_n_blocks)
+        q_subtile_factor = 1
+    else:
+        expected_count_shape, expected_index_shape, q_subtile_factor = (
+            infer_block_sparse_expected_shapes(
+                tensors,
+                batch_size=batch_size,
+                num_head=num_head,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                m_block_size=m_block_size,
+                n_block_size=n_block_size,
+                q_stage=q_stage,
+                context="forward",
+                sparse_block_size_q=sparse_block_size_q,
+                sparse_block_size_kv=sparse_block_size_kv,
+            )
+        )
+    normalized_tensors = normalize_block_sparse_tensors(
+        tensors,
+        expected_count_shape=expected_count_shape,
+        expected_index_shape=expected_index_shape,
+    )
+    return (
+        normalized_tensors,
+        get_block_sparse_broadcast_pattern(normalized_tensors),
+        q_subtile_factor,
+    )
+
+
+def normalize_block_sparse_config_bwd(
+    tensors: BlockSparseTensorsTorch,
+    *,
+    batch_size: int,
+    num_head: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    block_size: tuple[int, int],
+    q_subtile_factor: int,
+) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None]:
+    m_block_size, n_block_size = block_size
+    if tensors.block_size is None:
+        sparse_block_size_q, sparse_block_size_kv = q_subtile_factor * m_block_size, n_block_size
+    else:
+        sparse_block_size_q, sparse_block_size_kv = tensors.block_size
+    if sparse_block_size_q != q_subtile_factor * m_block_size:
+        raise ValueError(
+            f"Block sparsity expects sparse_block_size_q={q_subtile_factor * m_block_size} "
+            f"for q_subtile_factor={q_subtile_factor}."
+        )
+    if sparse_block_size_kv != n_block_size:
+        raise ValueError(
+            f"Block sparsity expects sparse_block_size[1]={n_block_size} to match tile_n."
+        )
+    expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes_bwd(
+        batch_size,
+        num_head,
+        seqlen_q,
+        seqlen_k,
+        m_block_size,
+        n_block_size,
+        q_subtile_factor,
+    )
+    normalized_tensors = normalize_block_sparse_tensors(
+        tensors,
+        expected_count_shape=expected_count_shape,
+        expected_index_shape=expected_index_shape,
+        context="_flash_attn_bwd",
+        hint=lambda: (
+            f"Backward expects Q-direction block-sparse tensors (q_mask_cnt/q_mask_idx, "
+            f"and optionally full_q_cnt/full_q_idx). Regenerate the backward BlockMask with "
+            f"BLOCK_SIZE=({q_subtile_factor * m_block_size}, {n_block_size})."
+        ),
+    )
+    return normalized_tensors, get_block_sparse_broadcast_pattern(normalized_tensors)
+
+
+def to_cute_block_sparse_tensors(
+    tensors: BlockSparseTensorsTorch, enable_tvm_ffi: bool = True
+) -> BlockSparseTensors | None:
+    """Convert torch block sparsity tensors to CuTe tensors, optionally for tvm ffi"""
+    if not is_block_sparsity_enabled(tensors):
+        return None
+    mask_block_cnt_tensor, mask_block_idx_tensor = [
+        to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi)
+        for t in (tensors.mask_block_cnt, tensors.mask_block_idx)
+    ]
+    full_block_cnt_tensor, full_block_idx_tensor = [
+        to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi)
+        if t is not None
+        else None
+        for t in (tensors.full_block_cnt, tensors.full_block_idx)
+    ]
+    cu_total_m_blocks_tensor, cu_block_idx_offsets_tensor = [
+        to_cute_tensor(t, assumed_align=4, leading_dim=0, enable_tvm_ffi=enable_tvm_ffi)
+        if t is not None
+        else None
+        for t in (tensors.cu_total_m_blocks, tensors.cu_block_idx_offsets)
+    ]
+    dq_write_order_tensor, dq_write_order_full_tensor = [
+        to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi)
+        if t is not None
+        else None
+        for t in (tensors.dq_write_order, tensors.dq_write_order_full)
+    ]
+
+    return BlockSparseTensors(
+        mask_block_cnt_tensor,
+        mask_block_idx_tensor,
+        full_block_cnt_tensor,
+        full_block_idx_tensor,
+        cu_total_m_blocks_tensor,
+        cu_block_idx_offsets_tensor,
+        dq_write_order_tensor,
+        dq_write_order_full_tensor,
+    )
+
+
+def fast_sampling(mask_mod):
+    """Convenience decorator to mark mask_mod as safe for 5-point fast sampling"""
+    mask_mod.use_fast_sampling = True
+    return mask_mod
@@ -0,0 +1,281 @@
+# Manage Ahead-of-Time (AOT) compiled kernels
+import fcntl
+import hashlib
+import os
+import pickle
+import sys
+import tempfile
+import time
+from functools import lru_cache
+from getpass import getuser
+from pathlib import Path
+from typing import Hashable, TypeAlias
+
+import ctypes
+
+import cutlass
+import cutlass.cute as cute
+import tvm_ffi
+from cutlass.cutlass_dsl import JitCompiledFunction
+from flash_attn.cute.fa_logging import fa_log
+
+# Pre-load cute DSL runtime libraries with RTLD_GLOBAL so that their symbols
+# (e.g. _cudaLibraryLoadData) are visible to .so modules loaded later via dlopen.
+# Upstream cute.runtime.load_module loads these without RTLD_GLOBAL, which causes
+# "undefined symbol" errors when loading cached kernels from disk.
+for _lib_path in cute.runtime.find_runtime_libraries(enable_tvm_ffi=False):
+    if Path(_lib_path).exists():
+        ctypes.CDLL(_lib_path, mode=ctypes.RTLD_GLOBAL)
+
+CompileKeyType: TypeAlias = tuple[Hashable, ...]
+CallableFunction: TypeAlias = JitCompiledFunction | tvm_ffi.Function
+
+# Enable cache via `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1`
+CUTE_DSL_CACHE_ENABLED: bool = os.getenv("FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED", "0") == "1"
+
+
+# Customize cache dir via `FLASH_ATTENTION_CUTE_DSL_CACHE_DIR`, default is
+# `/tmp/${USER}/flash_attention_cute_dsl_cache``
+CUTE_DSL_CACHE_DIR: str | None = os.getenv("FLASH_ATTENTION_CUTE_DSL_CACHE_DIR", None)
+
+
+def get_cache_path() -> Path:
+    if CUTE_DSL_CACHE_DIR is not None:
+        cache_dir = Path(CUTE_DSL_CACHE_DIR)
+    else:
+        cache_dir = Path(tempfile.gettempdir()) / getuser() / "flash_attention_cute_dsl_cache"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+@lru_cache(maxsize=1)
+def _compute_source_fingerprint() -> str:
+    """
+    Hash all CuTe Python sources plus runtime ABI stamps into a short fingerprint.
+
+    The fingerprint changes whenever:
+    - Any .py file under flash_attn/cute is added, removed, renamed, or modified.
+    - The Python minor version changes (e.g. 3.13 -> 3.14).
+    - The cutlass or tvm_ffi package version changes.
+
+    Computed once per process and cached.
+    """
+    cute_root = Path(__file__).resolve().parent
+    h = hashlib.sha256()
+
+    h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode())
+    h.update(f"cutlass={cutlass.__version__}".encode())
+    h.update(f"tvm_ffi={tvm_ffi.__version__}".encode())
+
+    for src in sorted(cute_root.rglob("*.py")):
+        if not src.is_file():
+            continue
+        h.update(src.relative_to(cute_root).as_posix().encode())
+        content = src.read_bytes()
+        h.update(len(content).to_bytes(8, "little"))
+        h.update(content)
+
+    return h.hexdigest()
+
+
+class FileLock:
+    """Context manager for advisory file locks using fcntl.flock.
+
+    Supports exclusive (write) and shared (read) locks.
+    Always blocks with polling until the lock is acquired or timeout is reached.
+
+    Usage:
+        with FileLock(lock_path, exclusive=True, timeout=15, label="abc"):
+            # do work under lock
+    """
+
+    def __init__(
+        self,
+        lock_path: Path,
+        exclusive: bool,
+        timeout: float = 15,
+        label: str = "",
+    ):
+        """
+        Args:
+            lock_path: Path to the lock file on disk.
+            exclusive: True for exclusive (write) lock, False for shared (read) lock.
+            timeout: Max seconds to wait for lock acquisition before raising RuntimeError.
+            label: Optional human-readable label for error messages.
+        """
+        self.lock_path: Path = lock_path
+        self.exclusive: bool = exclusive
+        self.timeout: float = timeout
+        self.label: str = label
+        self._fd: int = -1
+
+    @property
+    def _lock_label(self) -> str:
+        kind = "exclusive" if self.exclusive else "shared"
+        return f"{kind} {self.label}" if self.label else kind
+
+    def __enter__(self) -> "FileLock":
+        open_flags = os.O_WRONLY | os.O_CREAT if self.exclusive else os.O_RDONLY | os.O_CREAT
+        lock_type = fcntl.LOCK_EX if self.exclusive else fcntl.LOCK_SH
+
+        self._fd = os.open(str(self.lock_path), open_flags)
+
+        deadline = time.monotonic() + self.timeout
+        acquired = False
+        while time.monotonic() < deadline:
+            try:
+                fcntl.flock(self._fd, lock_type | fcntl.LOCK_NB)
+                acquired = True
+                break
+            except OSError:
+                time.sleep(0.1)
+        if not acquired:
+            os.close(self._fd)
+            self._fd = None
+            raise RuntimeError(
+                f"Timed out after {self.timeout}s waiting for "
+                f"{self._lock_label} lock: {self.lock_path}"
+            )
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self._fd is not None:
+            fcntl.flock(self._fd, fcntl.LOCK_UN)
+            os.close(self._fd)
+            self._fd = None
+
+
+class JITCache:
+    """
+    In-memory cache for compiled functions.
+    """
+
+    def __init__(self):
+        self.cache: dict[CompileKeyType, CallableFunction] = {}
+
+    def __setitem__(self, key: CompileKeyType, fn: JitCompiledFunction) -> None:
+        self.cache[key] = fn
+
+    def __getitem__(self, key: CompileKeyType) -> CallableFunction:
+        return self.cache[key]
+
+    def __contains__(self, key: CompileKeyType) -> bool:
+        return key in self.cache
+
+    def clear(self) -> None:
+        """
+        Clear in-memory cache of compiled functions
+        """
+        self.cache.clear()
+
+
+class JITPersistentCache(JITCache):
+    """
+    In-memory cache for compiled functions, which is also backed by persistent storage.
+    Use cutedsl ahead-of-time (AOT) compilation, only supporting enable_tvm_ffi=True
+    """
+
+    EXPORT_FUNCTION_PREFIX = "func"
+    LOCK_TIMEOUT_SECONDS = 15
+
+    def __init__(self, cache_path: Path):
+        super().__init__()
+        cache_path.mkdir(parents=True, exist_ok=True)
+        self.cache_path: Path = cache_path
+
+    def __setitem__(self, key: CompileKeyType, fn: JitCompiledFunction) -> None:
+        JITCache.__setitem__(self, key, fn)
+        self._try_export_to_storage(key, fn)
+
+    def __getitem__(self, key: CompileKeyType) -> CallableFunction:
+        # Use __contains__ to try populating in-memory cache with persistent storage
+        self.__contains__(key)
+        return JITCache.__getitem__(self, key)
+
+    def __contains__(self, key: CompileKeyType) -> bool:
+        # Checks in-memory cache first, then tries loading from storage.
+        # When returning True, guarantees the in-memory cache is populated.
+        if JITCache.__contains__(self, key):
+            return True
+        return self._try_load_from_storage(key)
+
+    def _try_load_from_storage(self, key: CompileKeyType) -> bool:
+        """
+        Try to load a function from persistent storage into in-memory cache.
+        Returns True if loaded successfully, False if not found on disk.
+        Holds a shared lock during loading to prevent concurrent writes.
+        """
+        sha256_hex = self._key_to_hash(key)
+        obj_path = self.cache_path / f"{sha256_hex}.o"
+        with FileLock(
+            self._lock_path(sha256_hex),
+            exclusive=False,
+            timeout=self.LOCK_TIMEOUT_SECONDS,
+            label=sha256_hex,
+        ):
+            if obj_path.exists():
+                fa_log(1, f"Loading compiled function from disk: {obj_path}")
+                m = cute.runtime.load_module(str(obj_path), enable_tvm_ffi=True)
+                fn = getattr(m, self.EXPORT_FUNCTION_PREFIX)
+                JITCache.__setitem__(self, key, fn)
+                return True
+            else:
+                fa_log(1, f"Cache miss on disk for key hash {sha256_hex}")
+        return False
+
+    def _try_export_to_storage(self, key: CompileKeyType, fn: JitCompiledFunction) -> None:
+        """Export a compiled function to persistent storage under exclusive lock."""
+        sha256_hex = self._key_to_hash(key)
+        with FileLock(
+            self._lock_path(sha256_hex),
+            exclusive=True,
+            timeout=self.LOCK_TIMEOUT_SECONDS,
+            label=sha256_hex,
+        ):
+            obj_path = self.cache_path / f"{sha256_hex}.o"
+            if obj_path.exists():
+                # Another process already exported.
+                fa_log(1, f"Skipping export, already on disk: {obj_path}")
+                return
+            fa_log(1, f"Exporting compiled function to disk: {obj_path}")
+            fn.export_to_c(
+                object_file_path=str(obj_path),
+                function_name=self.EXPORT_FUNCTION_PREFIX,
+            )
+            fa_log(1, f"Successfully exported compiled function to disk: {obj_path}")
+
+    def _key_to_hash(self, key: CompileKeyType) -> str:
+        return hashlib.sha256(pickle.dumps(key)).hexdigest()
+
+    def _lock_path(self, sha256_hex: str) -> Path:
+        return self.cache_path / f"{sha256_hex}.lock"
+
+    def clear(self) -> None:
+        """
+        Not only clear the in-memory cache. Also purge persistent compilation cache.
+        """
+        fa_log(1, f"Clearing persistent cache at {self.cache_path}")
+        super().clear()
+        for child in self.cache_path.iterdir():
+            child.unlink()
+
+
+def get_jit_cache(name: str | None = None) -> JITCache:
+    """
+    JIT cache factory.
+    `name` is an optional identifier to create subdirectories to manage cache.
+
+    When persistent caching is enabled, artifacts are namespaced under a
+    source fingerprint directory so that code or dependency changes
+    automatically invalidate stale entries.
+    """
+    if CUTE_DSL_CACHE_ENABLED:
+        path = get_cache_path() / _compute_source_fingerprint()
+        if name:
+            path = path / name
+        fa_log(1, f"Creating persistent JIT cache at {path}")
+        return JITPersistentCache(path)
+    else:
+        fa_log(1, "Persistent cache disabled, using in-memory JIT cache")
+        return JITCache()
@@ -0,0 +1,551 @@
+from functools import partial
+from typing import Callable, Optional, Tuple
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass import Boolean, Int8, Int32, const_expr
+
+from flash_attn.cute.block_sparsity import (
+    BlockSparseTensors,
+    BlockSparseTensorsTorch,
+    to_cute_block_sparse_tensors,
+)
+from flash_attn.cute.block_sparse_utils import get_curr_blocksparse_tensors
+from flash_attn.cute.testing import is_fake_mode
+from flash_attn.cute.cute_dsl_utils import (
+    get_aux_tensor_metadata,
+    to_cute_aux_tensor,
+    to_cute_tensor,
+)
+from flash_attn.cute.utils import (
+    get_batch_from_cu_tensor,
+    hash_callable,
+    scalar_to_ssa,
+    ssa_to_scalar,
+)
+from flash_attn.cute.mask import call_mask_mod
+from flash_attn.cute.seqlen_info import SeqlenInfoQK
+from flash_attn.cute.utils import AuxData
+
+
+class BlockSparsityKernel:
+    """Block sparsity kernel for FlexAttention.
+
+    This kernel computes `mask_mod` for every token of each block
+    to determine if an n block is full, masked, or neither.
+
+    Writes block counts and indices to a BlockSparseTensors object.
+
+    When use_fast_sampling=True, uses 5-point sampling (4 corners + center)
+    which is much faster but only suitable for masks where this is sufficient.
+
+    TODO:
+        - optimize mask_mod evaluation
+        - transposed tensors for bwd pass
+    """
+
+    def __init__(
+        self,
+        mask_mod: Callable,
+        tile_mn: Tuple[int, int],
+        compute_full_blocks: bool = True,
+        use_aux_tensors: bool = False,
+        use_fast_sampling: bool = False,
+    ):
+        self.mask_mod = mask_mod
+        self.tile_mn = tile_mn
+        self.compute_full_blocks = compute_full_blocks
+        self.use_aux_tensors = use_aux_tensors
+        self.use_fast_sampling = use_fast_sampling
+
+    @cute.jit
+    def __call__(
+        self,
+        blocksparse_tensors: BlockSparseTensors,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
+    ):
+        mask_cnt, mask_idx, full_cnt, full_idx, mCuTotalMBlocks, mCuBlockIdxOffsets, *_ = (
+            blocksparse_tensors
+        )
+
+        self.is_varlen_q = const_expr(mCuSeqlensQ is not None)
+
+        if const_expr(self.compute_full_blocks):
+            assert full_cnt is not None and full_idx is not None, (
+                "full block tensors must be provided when computing full blocks"
+            )
+        if const_expr(not self.is_varlen_q):
+            batch_size, num_heads, num_m_blocks, _ = mask_idx.shape
+            total_m_blocks = batch_size * num_m_blocks
+        else:
+            assert const_expr(mCuTotalMBlocks is not None), (
+                "mCuTotalMBlocks must be provided when varlen q"
+            )
+            num_heads, total_m_blocks = mask_cnt.shape  # num_m_blocks is total_m_blocks
+            batch_size = mCuSeqlensQ.shape[0] - 1
+
+        if const_expr(self.use_fast_sampling):
+            num_threads = 5
+            self.num_warps = 1
+        else:
+            num_threads = self.tile_mn[0]
+            self.num_warps = (num_threads + 32 - 1) // 32
+
+        if const_expr(not self.is_varlen_q):
+            grid = [num_m_blocks, num_heads, batch_size]
+        else:
+            grid = [total_m_blocks, num_heads, 1]
+
+        self.kernel(
+            blocksparse_tensors,
+            seqlen_q,
+            seqlen_k,
+            batch_size,
+            mCuSeqlensQ,
+            mCuSeqlensK,
+            mSeqUsedQ,
+            mSeqUsedK,
+            mCuTotalMBlocks,
+            mCuBlockIdxOffsets,
+            aux_data,
+        ).launch(grid=grid, block=[num_threads, 1, 1])
+
+    @cute.kernel
+    def kernel(
+        self,
+        blocksparse_tensors: BlockSparseTensors,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        batch_size: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuBlockIdxOffsets: Optional[cute.Tensor] = None,
+        aux_data: AuxData = AuxData(),
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        warp_idx = cute.arch.warp_idx()
+        lane_id = cute.arch.lane_idx()
+
+        ssa = partial(scalar_to_ssa, dtype=Int32)
+
+        @cute.struct
+        class SharedStorage:
+            reduction_buffer_smem: cute.struct.Align[
+                cute.struct.MemRange[cutlass.Int8, 2 * self.num_warps], 1024
+            ]
+
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage, 16)
+
+        reduction_buffer = storage.reduction_buffer_smem.get_tensor(
+            cute.make_layout((self.num_warps, 2))
+        )
+        SeqlenInfoCls = partial(
+            SeqlenInfoQK.create,
+            seqlen_q_static=seqlen_q,
+            seqlen_k_static=seqlen_k,
+            mCuSeqlensQ=mCuSeqlensQ,
+            mCuSeqlensK=mCuSeqlensK,
+            mSeqUsedQ=mSeqUsedQ,
+            mSeqUsedK=mSeqUsedK,
+            mCuTotalMBlocks=mCuTotalMBlocks,
+            mCuBlockIdxOffsets=mCuBlockIdxOffsets,
+            tile_m=self.tile_mn[0],
+            tile_n=self.tile_mn[1],
+        )
+
+        if const_expr(not self.is_varlen_q):
+            m_block, head_idx, batch_idx = cute.arch.block_idx()
+        else:
+            global_m_block, head_idx, _ = cute.arch.block_idx()
+            batch_idx = get_batch_from_cu_tensor(global_m_block, mCuTotalMBlocks)
+            m_block = global_m_block - mCuTotalMBlocks[batch_idx]
+
+        seqlen = SeqlenInfoCls(batch_idx)
+        seqlen_q = seqlen.seqlen_q
+        seqlen_k = seqlen.seqlen_k
+        global_m_block = seqlen.m_block_offset + m_block
+
+        num_n_blocks = (seqlen_k + self.tile_mn[1] - 1) // self.tile_mn[1]
+
+        _, curr_mask_idx, _, curr_full_idx = get_curr_blocksparse_tensors(
+            batch_idx, head_idx, m_block, blocksparse_tensors, seqlen
+        )
+
+        num_mask_blocks = Int32(0)
+        num_full_blocks = Int32(0)
+
+        m_base = m_block * self.tile_mn[0]
+        if const_expr(self.use_fast_sampling):
+            # Loop-invariant per-thread q_idx for the 5 sample points
+            # (tidx 0, 1: top corners; 2, 3: bottom corners; 4: center).
+            q_idx_sample = m_base
+            if tidx == 2 or tidx == 3:
+                q_idx_sample = cutlass.min(m_base + self.tile_mn[0] - 1, seqlen_q - 1)
+            elif tidx == 4:
+                q_idx_sample = m_base + cutlass.min(seqlen_q - m_base, self.tile_mn[0]) // 2
+        else:
+            q_idx_thread = m_base + tidx
+            thread_in_bounds = Boolean(tidx < self.tile_mn[0] and q_idx_thread < seqlen_q)
+
+        for n_block in cutlass.range(num_n_blocks):
+            n_base = n_block * self.tile_mn[1]
+
+            if const_expr(self.use_fast_sampling):
+                # 5-point sampling (4 corners + center). Interior n_blocks
+                # (n_base + tile_n <= seqlen_k) skip the OOB clamp on the right /
+                # center samples.
+                is_interior = (n_base + self.tile_mn[1]) <= seqlen_k
+                n_right = Int32(0)
+                n_mid = Int32(0)
+                if is_interior:
+                    n_right = n_base + self.tile_mn[1] - 1
+                    n_mid = n_base + self.tile_mn[1] // 2
+                else:
+                    n_right = cutlass.min(n_base + self.tile_mn[1] - 1, seqlen_k - 1)
+                    n_mid = n_base + cutlass.min(seqlen_k - n_base, self.tile_mn[1]) // 2
+
+                kv_idx = n_base
+                if tidx == 1 or tidx == 3:
+                    kv_idx = n_right
+                elif tidx == 4:
+                    kv_idx = n_mid
+
+                thread_result = Boolean(False)
+                thread_is_valid = Boolean(False)
+                if tidx < 5:
+                    thread_is_valid = Boolean(True)
+                    thread_result = ssa_to_scalar(
+                        call_mask_mod(
+                            self.mask_mod,
+                            ssa(batch_idx),
+                            ssa(head_idx),
+                            ssa(q_idx_sample),
+                            ssa(kv_idx),
+                            seqlen,
+                            aux_data,
+                        )
+                    )
+
+                has_unmasked = cute.arch.vote_any_sync(thread_result & thread_is_valid)
+                has_masked = cute.arch.vote_any_sync(Boolean(not thread_result) & thread_is_valid)
+
+            else:
+                # Full path. Interior blocks (n_base + tile_n <= seqlen_k) drop the
+                # per-element bound check; the boundary block (at most one) keeps it.
+                thread_has_unmasked = Boolean(False)
+                thread_has_masked = Boolean(False)
+                kv_idx = Int32(0)
+                is_interior = (n_base + self.tile_mn[1]) <= seqlen_k
+
+                if is_interior:
+                    if thread_in_bounds:
+                        for c in cutlass.range(self.tile_mn[1], unroll_full=True):
+                            mask_val = ssa_to_scalar(
+                                call_mask_mod(
+                                    self.mask_mod,
+                                    ssa(batch_idx),
+                                    ssa(head_idx),
+                                    ssa(q_idx_thread),
+                                    ssa(n_base + c),
+                                    seqlen,
+                                    aux_data,
+                                )
+                            )
+                            thread_has_unmasked |= Boolean(mask_val)
+                            thread_has_masked |= Boolean(not mask_val)
+                else:
+                    if thread_in_bounds:
+                        for c in cutlass.range(self.tile_mn[1], unroll_full=True):
+                            kv_idx = n_base + c
+                            if kv_idx < seqlen_k:
+                                mask_val = ssa_to_scalar(
+                                    call_mask_mod(
+                                        self.mask_mod,
+                                        ssa(batch_idx),
+                                        ssa(head_idx),
+                                        ssa(q_idx_thread),
+                                        ssa(kv_idx),
+                                        seqlen,
+                                        aux_data,
+                                    )
+                                )
+                                thread_has_unmasked |= Boolean(mask_val)
+                                thread_has_masked |= Boolean(not mask_val)
+
+                warp_unmasked = cute.arch.vote_any_sync(thread_has_unmasked & thread_in_bounds)
+                warp_masked = cute.arch.vote_any_sync(thread_has_masked & thread_in_bounds)
+                if lane_id == 0:
+                    reduction_buffer[warp_idx, 0] = Int8(1) if warp_unmasked else Int8(0)
+                    reduction_buffer[warp_idx, 1] = Int8(1) if warp_masked else Int8(0)
+                cute.arch.sync_threads()
+
+                # Cross-warp OR via warp 0; thread 0 (lane 0 of warp 0) holds the result.
+                has_unmasked = Boolean(False)
+                has_masked = Boolean(False)
+                if warp_idx == 0:
+                    lane_unmasked = Boolean(False)
+                    lane_masked = Boolean(False)
+                    if lane_id < self.num_warps:
+                        lane_unmasked = reduction_buffer[lane_id, 0] != Int8(0)
+                        lane_masked = reduction_buffer[lane_id, 1] != Int8(0)
+                    has_unmasked = cute.arch.vote_any_sync(lane_unmasked)
+                    has_masked = cute.arch.vote_any_sync(lane_masked)
+
+            # Only thread 0 updates the output arrays (common to both paths)
+            if tidx == 0:
+                # Block classification based on what we found:
+                # - If has_masked and has_unmasked: partial block (needs masking)
+                # - If only has_unmasked: full block (no masking needed)
+                # - If only has_masked: skip this block entirely
+                is_partial = Boolean(has_masked and has_unmasked)
+                is_full = Boolean(has_unmasked and (not has_masked))
+
+                if is_partial:
+                    curr_mask_idx[num_mask_blocks] = n_block
+                    num_mask_blocks += 1
+                elif is_full and const_expr(self.compute_full_blocks):
+                    curr_full_idx[num_full_blocks] = n_block
+                    num_full_blocks += 1
+
+        # Only thread 0 writes back the counts
+        if tidx == 0:
+            mask_cnt, _, full_cnt, *_ = blocksparse_tensors
+            if const_expr(self.is_varlen_q):
+                mask_cnt[head_idx, global_m_block] = num_mask_blocks
+                if const_expr(self.compute_full_blocks):
+                    full_cnt[head_idx, global_m_block] = num_full_blocks
+            else:
+                mask_cnt[batch_idx, head_idx, m_block] = num_mask_blocks
+                if const_expr(self.compute_full_blocks):
+                    full_cnt[batch_idx, head_idx, m_block] = num_full_blocks
+
+
+def compute_block_sparsity(
+    tile_m,
+    tile_n,
+    batch_size,
+    num_heads,
+    seqlen_q,
+    seqlen_k,
+    mask_mod: Callable,
+    aux_tensors: Optional[list],
+    device,
+    aux_scalars: Optional[tuple] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    cu_total_m_blocks: Optional[torch.Tensor] = None,
+    cu_block_idx_offsets: Optional[torch.Tensor] = None,
+    compute_full_blocks: bool = True,
+    use_fast_sampling: bool = False,
+) -> BlockSparseTensorsTorch:
+    """
+    Computes block sparsity for a given `mask_mod`.
+
+    Args:
+        tile_m: The tile size for the m dimension.
+        tile_n: The tile size for the n dimension.
+        batch_size: The batch size.
+        num_heads: The number of heads.
+        seqlen_q: The sequence length for the query.
+        seqlen_k: The sequence length for the key.
+        mask_mod: The `mask_mod` callable to use.
+        aux_tensors: A list of auxiliary tensors.
+        device: The device to use.
+        cu_seqlens_q: Cumulative q sequence lengths for varlen
+        cu_seqlens_k: Cumulative k sequence lengths for varlen
+        seqused_q: Per-batch effective q sequence lengths
+        seqused_k: Per-batch effective k sequence lengths
+        cu_total_m_blocks: Cumulative total m blocks tensor for varlen q
+        cu_block_idx_offsets: Cumulative offsets into the packed mask_block_idx /
+            full_block_idx tensors per batch (== cumsum of M_b * N_b).
+        compute_full_blocks: Whether to compute full blocks. If False, only partially-masked blocks are computed.
+        use_fast_sampling: Whether to use 5-point sampling (4 corners + center). This is much faster, but only suitable for masks where this check is sufficient.
+
+    Returns:
+        BlockSparseTensorsTorch
+    """
+    aux_scalars = tuple(aux_scalars) if aux_scalars else None
+
+    # Check if mask_mod is marked as suitable for 5-point sampling
+    use_fast_sampling = getattr(mask_mod, "use_fast_sampling", use_fast_sampling)
+
+    num_m_blocks = (seqlen_q + tile_m - 1) // tile_m
+    num_n_blocks = (seqlen_k + tile_n - 1) // tile_n
+
+    if cu_seqlens_q is not None:
+        assert cu_total_m_blocks is not None, "total m blocks must be provided when varlen q"
+        total_m_blocks = cu_total_m_blocks[-1].item()
+        if cu_block_idx_offsets is None and (cu_seqlens_k is not None or seqused_k is not None):
+            # Derive cu_block_idx_offsets from per-batch K seqlens.
+            cu_block_idx_offsets_list = [0]
+            for batch_idx in range(batch_size):
+                batch_seqlen_q = cu_seqlens_q[batch_idx + 1].item() - cu_seqlens_q[batch_idx].item()
+                if cu_seqlens_k is not None:
+                    batch_seqlen_k = (
+                        cu_seqlens_k[batch_idx + 1].item() - cu_seqlens_k[batch_idx].item()
+                    )
+                else:
+                    batch_seqlen_k = seqused_k[batch_idx].item()
+                num_m_blocks_batch = (batch_seqlen_q + tile_m - 1) // tile_m
+                num_n_blocks_batch = (batch_seqlen_k + tile_n - 1) // tile_n
+                cu_block_idx_offsets_list.append(
+                    cu_block_idx_offsets_list[-1] + num_m_blocks_batch * num_n_blocks_batch
+                )
+            cu_block_idx_offsets = torch.tensor(
+                cu_block_idx_offsets_list, dtype=torch.int32, device=device
+            )
+        if cu_block_idx_offsets is not None:
+            total_n_blocks = cu_block_idx_offsets[-1].item()
+        else:
+            # Uniform-K varlen-Q: every batch has the same K seqlen.
+            total_n_blocks = total_m_blocks * num_n_blocks
+
+        mask_block_cnt = torch.zeros((num_heads, total_m_blocks), device=device, dtype=torch.int32)
+        mask_block_idx = torch.zeros((num_heads, total_n_blocks), device=device, dtype=torch.int32)
+        full_block_cnt = (
+            torch.zeros((num_heads, total_m_blocks), device=device, dtype=torch.int32)
+            if compute_full_blocks
+            else None
+        )
+        full_block_idx = (
+            torch.zeros((num_heads, total_n_blocks), device=device, dtype=torch.int32)
+            if compute_full_blocks
+            else None
+        )
+    else:
+        total_m_blocks = batch_size * num_m_blocks
+        total_n_blocks = batch_size * num_m_blocks * num_n_blocks
+
+        mask_block_cnt = torch.zeros(
+            (batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32
+        )
+        mask_block_idx = torch.zeros(
+            (batch_size, num_heads, num_m_blocks, num_n_blocks), device=device, dtype=torch.int32
+        )
+        full_block_cnt = (
+            torch.zeros((batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32)
+            if compute_full_blocks
+            else None
+        )
+        full_block_idx = (
+            torch.zeros(
+                (batch_size, num_heads, num_m_blocks, num_n_blocks),
+                device=device,
+                dtype=torch.int32,
+            )
+            if compute_full_blocks
+            else None
+        )
+
+    blocksparse_tensors_torch = BlockSparseTensorsTorch(
+        mask_block_cnt=mask_block_cnt,
+        mask_block_idx=mask_block_idx,
+        full_block_cnt=full_block_cnt,
+        full_block_idx=full_block_idx,
+        cu_total_m_blocks=cu_total_m_blocks,
+        cu_block_idx_offsets=cu_block_idx_offsets,
+        block_size=(tile_m, tile_n),
+    )
+
+    mask_mod_hash = hash_callable(mask_mod)
+    if aux_tensors is not None:
+        aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors)
+    else:
+        aux_tensor_metadata = None
+    aux_scalar_metadata = tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None
+
+    compile_key = (
+        tile_m,
+        tile_n,
+        mask_mod_hash,
+        aux_tensor_metadata,
+        aux_scalar_metadata,
+        compute_full_blocks,
+        cu_seqlens_q is None,
+        cu_seqlens_k is None,
+        seqused_q is None,
+        seqused_k is None,
+        aux_tensors is not None,
+        use_fast_sampling,
+    )
+    if compile_key not in compute_block_sparsity.compile_cache:
+        (
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+        ) = [
+            to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None
+            for t in (
+                cu_seqlens_q,
+                cu_seqlens_k,
+                seqused_q,
+                seqused_k,
+            )
+        ]
+        blocksparse_tensors = to_cute_block_sparse_tensors(
+            blocksparse_tensors_torch, enable_tvm_ffi=True
+        )
+        if aux_tensors is not None:
+            cute_aux_tensors = [to_cute_aux_tensor(buf) for buf in aux_tensors]
+        else:
+            cute_aux_tensors = None
+        kernel = BlockSparsityKernel(
+            mask_mod,
+            tile_mn=(tile_m, tile_n),
+            compute_full_blocks=compute_full_blocks,
+            use_aux_tensors=aux_tensors is not None,
+            use_fast_sampling=use_fast_sampling,
+        )
+
+        compute_block_sparsity.compile_cache[compile_key] = cute.compile(
+            kernel,
+            blocksparse_tensors,
+            seqlen_q,
+            seqlen_k,
+            cu_seqlens_q_tensor,
+            cu_seqlens_k_tensor,
+            seqused_q_tensor,
+            seqused_k_tensor,
+            AuxData(cute_aux_tensors, aux_scalars),
+            options="--enable-tvm-ffi",
+        )
+
+    if not is_fake_mode():
+        compute_block_sparsity.compile_cache[compile_key](
+            (
+                blocksparse_tensors_torch.mask_block_cnt,
+                blocksparse_tensors_torch.mask_block_idx,
+                blocksparse_tensors_torch.full_block_cnt,
+                blocksparse_tensors_torch.full_block_idx,
+                blocksparse_tensors_torch.cu_total_m_blocks,
+                blocksparse_tensors_torch.cu_block_idx_offsets,
+                blocksparse_tensors_torch.dq_write_order,
+                blocksparse_tensors_torch.dq_write_order_full,
+            ),
+            seqlen_q,
+            seqlen_k,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            AuxData(aux_tensors, aux_scalars),
+        )
+
+    return blocksparse_tensors_torch
+
+
+compute_block_sparsity.compile_cache = {}
@@ -0,0 +1,372 @@
+# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+
+import math
+from typing import Optional, Type, Callable
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute.nvgpu import cpasync
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm
+import cutlass.pipeline
+
+
+@dsl_user_op
+def cvt_copy(
+    atom: cute.CopyAtom,
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    *,
+    pred: Optional[cute.Tensor] = None,
+    loc=None,
+    ip=None,
+    **kwargs,
+) -> None:
+    assert isinstance(src.iterator, cute.Pointer) and src.memspace == cute.AddressSpace.rmem
+    if const_expr(src.element_type != dst.element_type):
+        src_cvt = cute.make_fragment_like(src, dst.element_type, loc=loc, ip=ip)
+        src_cvt.store(src.load().to(dst.element_type))
+        src = src_cvt
+    cute.copy(atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
+
+
+@dsl_user_op
+def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+    dst = cute.make_fragment_like(src, src.element_type, loc=loc, ip=ip)
+    cute.autovec_copy(src, dst, loc=loc, ip=ip)
+    return dst
+
+
+@dsl_user_op
+def get_copy_atom(
+    dtype: Type[cutlass.Numeric], num_copy_elems: int, is_async: bool = False, *, loc=None, ip=None
+) -> cute.CopyAtom:
+    num_copy_bits = const_expr(min(128, num_copy_elems * dtype.width))
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    return cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+
+
+@dsl_user_op
+def make_tmem_copy(
+    tmem_copy_atom: cute.CopyAtom, num_wg: int = 1, *, loc=None, ip=None
+) -> cute.CopyAtom:
+    num_dp, num_bits, num_rep, _ = sm100_utils.get_tmem_copy_properties(tmem_copy_atom)
+    assert num_dp == 32
+    assert num_bits == 32
+    tiler_mn = (cute.make_layout((128 * num_rep * num_wg // 32, 32), stride=(32, 1)),)
+    layout_tv = cute.make_layout(
+        ((32, 4, num_wg), (num_rep, 32)), stride=((0, 1, 4 * num_rep), (4, 4 * num_rep * num_wg))
+    )
+    return cute.make_tiled_copy(tmem_copy_atom, layout_tv, tiler_mn)
+
+
+@dsl_user_op
+def copy(
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    *,
+    pred: Optional[cute.Tensor] = None,
+    num_copy_elems: int = 1,
+    is_async: bool = False,
+    loc=None,
+    ip=None,
+    **kwargs,
+) -> None:
+    copy_atom = get_copy_atom(src.element_type, num_copy_elems, is_async)
+    cute.copy(copy_atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
+
+
+def tiled_copy_1d(
+    dtype: Type[cutlass.Numeric], num_threads: int, num_copy_elems: int = 1, is_async: bool = False
+) -> cute.TiledCopy:
+    num_copy_bits = num_copy_elems * dtype.width
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+    thr_layout = cute.make_layout(num_threads)
+    val_layout = cute.make_layout(num_copy_elems)
+    return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
+
+
+def tiled_copy_2d(
+    dtype: Type[cutlass.Numeric], major_mode_size: int, num_threads: int, is_async: bool = False
+) -> cute.TiledCopy:
+    num_copy_bits = math.gcd(major_mode_size, 128 // dtype.width) * dtype.width
+    copy_elems = num_copy_bits // dtype.width
+    copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp()
+    copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+    gmem_threads_per_row = major_mode_size // copy_elems
+    assert num_threads % gmem_threads_per_row == 0
+    thr_layout = cute.make_ordered_layout(
+        (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+        order=(1, 0),
+    )
+    val_layout = cute.make_layout((1, copy_elems))
+    return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout)
+
+
+@dsl_user_op
+def atomic_add_fp32x4(
+    a: Float32, b: Float32, c: Float32, d: Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None
+) -> None:
+    gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
+    # cache_hint = cutlass.Int64(0x12F0000000000000)
+    llvm.inline_asm(
+        None,
+        [
+            gmem_ptr_i64,
+            Float32(a).ir_value(loc=loc, ip=ip),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            Float32(c).ir_value(loc=loc, ip=ip),
+            Float32(d).ir_value(loc=loc, ip=ip),
+        ],
+        # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()],
+        "{\n\t"
+        # ".reg .b128 abcd;\n\t"
+        # "mov.b128 abcd, {$1, $2, $3, $4};\n\t"
+        ".reg .v4 .f32 abcd;\n\t"
+        # "mov.b128 abcd, {$1, $2, $3, $4};\n\t"
+        "mov.f32 abcd.x, $1;\n\t"
+        "mov.f32 abcd.y, $2;\n\t"
+        "mov.f32 abcd.z, $3;\n\t"
+        "mov.f32 abcd.w, $4;\n\t"
+        "red.global.add.v4.f32 [$0], abcd;\n\t"
+        # "red.global.add.L2::cache_hint.v4.f32 [$0], abcd, 0x14F0000000000000;\n\t"
+        "}\n",
+        # "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;",
+        # "red.global.add.L2::cache_hint.f32 [$0], $1, $2;",
+        "l,f,f,f,f",
+        # "l,f,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def set_block_rank(
+    smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None
+) -> Int32:
+    """Map the given smem pointer to the address at another CTA rank in the cluster."""
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
+            "mapa.shared::cluster.u32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def store_shared_remote_fp32x4(
+    a: Float32,
+    b: Float32,
+    c: Float32,
+    d: Float32,
+    smem_ptr: cute.Pointer,
+    mbar_ptr: cute.Pointer,
+    peer_cta_rank_in_cluster: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    remote_smem_ptr_i32 = set_block_rank(
+        smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    remote_mbar_ptr_i32 = set_block_rank(
+        mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    llvm.inline_asm(
+        None,
+        [
+            remote_smem_ptr_i32,
+            remote_mbar_ptr_i32,
+            Float32(a).ir_value(loc=loc, ip=ip),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            Float32(c).ir_value(loc=loc, ip=ip),
+            Float32(d).ir_value(loc=loc, ip=ip),
+        ],
+        "{\n\t"
+        ".reg .v4 .f32 abcd;\n\t"
+        "mov.f32 abcd.x, $2;\n\t"
+        "mov.f32 abcd.y, $3;\n\t"
+        "mov.f32 abcd.z, $4;\n\t"
+        "mov.f32 abcd.w, $5;\n\t"
+        "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 [$0], abcd, [$1];\n\t"
+        "}\n",
+        "r,r,f,f,f,f",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def cpasync_bulk_s2cluster(
+    smem_src_ptr: cute.Pointer,
+    smem_dst_ptr: cute.Pointer,
+    mbar_ptr: cute.Pointer,
+    size: int | Int32,
+    peer_cta_rank_in_cluster: Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    smem_src_ptr_i32 = smem_src_ptr.toint(loc=loc, ip=ip).ir_value()
+    smem_dst_ptr_i32 = set_block_rank(
+        smem_dst_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    mbar_ptr_i32 = set_block_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [
+            smem_dst_ptr_i32,
+            smem_src_ptr_i32,
+            mbar_ptr_i32,
+            Int32(size).ir_value(loc=loc, ip=ip),
+        ],
+        "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [$0], [$1], $3, [$2];",
+        "r,r,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def cpasync_bulk_g2s(
+    gmem_ptr: cute.Pointer,
+    smem_ptr: cute.Pointer,
+    tma_bar_ptr: cute.Pointer,
+    size: int | Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    mbar_ptr_i32 = tma_bar_ptr.toint(loc=loc, ip=ip).ir_value()
+    llvm.inline_asm(
+        None,
+        [gmem_ptr_i64, smem_ptr_i32, mbar_ptr_i32, Int32(size).ir_value()],
+        "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [$1], [$0], $3, [$2];",
+        "l,r,r,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def cpasync_reduce_bulk_add_f32(
+    smem_ptr: cute.Pointer,
+    gmem_ptr: cute.Pointer,
+    store_bytes: int | Int32,
+    *,
+    loc=None,
+    ip=None,
+):
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    # cache_hint = cutlass.Int64(0x14F0000000000000)  # EVICT_LAST
+    llvm.inline_asm(
+        None,
+        [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()],
+        "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;",
+        "l,r,r",
+        # [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value(), cache_hint.ir_value()],
+        # "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [$0], [$1], $2, $3;",
+        # "l,r,r,l",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+def cpasync_bulk_get_copy_fn(
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    single_stage: bool = False,
+    **kwargs,
+) -> Callable:
+    # src_is_smem = const_expr(
+    #     isinstance(src_tensor.iterator, cute.Pointer)
+    #     and src_tensor.memspace == cute.AddressSpace.smem
+    # )
+    group_rank_src = const_expr(cute.rank(src_tensor) - (1 if not single_stage else 0))
+    group_rank_dst = const_expr(cute.rank(dst_tensor) - (1 if not single_stage else 0))
+    # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
+    src = cute.group_modes(src_tensor, 0, group_rank_src)
+    dst = cute.group_modes(dst_tensor, 0, group_rank_dst)
+
+    def copy_bulk(src_idx, dst_idx, **new_kwargs):
+        size = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8)
+        cpasync_bulk_g2s(
+            src[None, src_idx].iterator,
+            dst[None, dst_idx].iterator,
+            size=size,
+            **new_kwargs,
+            **kwargs,
+        )
+
+    def copy_bulk_single_stage(**new_kwargs):
+        size = const_expr(cute.size(src.shape) * src.element_type.width // 8)
+        cpasync_bulk_g2s(src.iterator, dst.iterator, size=size, **new_kwargs, **kwargs)
+
+    return copy_bulk if const_expr(not single_stage) else copy_bulk_single_stage
+
+
+def tma_get_copy_fn(
+    atom: cute.CopyAtom,
+    cta_coord: cute.Coord,
+    cta_layout: cute.Layout,
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    filter_zeros: bool = False,
+    single_stage: bool = False,
+    **kwargs,
+) -> Callable:
+    src_is_smem = const_expr(
+        isinstance(src_tensor.iterator, cute.Pointer)
+        and src_tensor.memspace == cute.AddressSpace.smem
+    )
+    smem_tensor, gmem_tensor = (src_tensor, dst_tensor) if src_is_smem else (dst_tensor, src_tensor)
+    group_rank_smem = const_expr(cute.rank(smem_tensor) - (1 if not single_stage else 0))
+    group_rank_gmem = const_expr(cute.rank(gmem_tensor) - (1 if not single_stage else 0))
+    # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
+    s, g = cpasync.tma_partition(
+        atom,
+        cta_coord,
+        cta_layout,
+        cute.group_modes(smem_tensor, 0, group_rank_smem),
+        cute.group_modes(gmem_tensor, 0, group_rank_gmem),
+    )
+    if const_expr(filter_zeros):
+        s = cute.filter_zeros(s)
+        g = cute.filter_zeros(g)
+    src, dst = (s, g) if src_is_smem else (g, s)
+
+    def copy_tma(src_idx, dst_idx, **new_kwargs):
+        cute.copy(atom, src[None, src_idx], dst[None, dst_idx], **new_kwargs, **kwargs)
+
+    def copy_tma_single_stage(**new_kwargs):
+        cute.copy(atom, src, dst, **new_kwargs, **kwargs)
+
+    return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g
+
+
+def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync):
+    def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
+        copy(
+            src_idx=src_idx,
+            dst_idx=producer_state.index,
+            tma_bar_ptr=pipeline.producer_get_barrier(producer_state),
+            **new_kwargs,
+        )
+
+    return copy_fn
@@ -0,0 +1,151 @@
+"""
+System ptxas replacement for CUTLASS DSL.
+Environment variables:
+    CUTE_DSL_PTXAS_PATH    - Path to ptxas (e.g., /usr/local/cuda/bin/ptxas)
+    CUTE_DSL_PTXAS_VERBOSE - Set to 1 for verbose output
+"""
+
+import os
+import sys
+import re
+import ctypes
+import subprocess
+from pathlib import Path
+
+import cutlass
+
+
+CUTE_DSL_PTXAS_PATH = os.environ.get("CUTE_DSL_PTXAS_PATH", None)
+VERBOSE = os.environ.get("CUTE_DSL_PTXAS_VERBOSE", "0") == "1"
+
+_original_load_cuda_library = None
+_user_wanted_ptx = False  # True if user originally set CUTE_DSL_KEEP_PTX=1
+
+
+def _log(msg):
+    if VERBOSE:
+        print(f"[ptxas] {msg}", file=sys.stderr)
+
+
+def _get_ptx(compiled_func) -> tuple[str, Path] | None:
+    """Find and read PTX file, stripping null bytes."""
+    func_name = getattr(compiled_func, "function_name", None)
+    if not func_name:
+        return None
+
+    dump_dir = os.environ.get("CUTE_DSL_DUMP_DIR", Path.cwd())
+    for ptx_path in Path(dump_dir).glob(f"*{func_name}*.ptx"):
+        content = ptx_path.read_text().rstrip("\x00")
+        if ".entry " in content and content.rstrip().endswith("}"):
+            _log(f"Found PTX: {ptx_path}")
+            return content, ptx_path
+    return None
+
+
+def _compile_ptx(ptx_path: Path, ptx_content: str) -> bytes:
+    """Compile PTX to cubin using system ptxas."""
+    # Extract arch from PTX
+    match = re.search(r"\.target\s+(sm_\d+[a-z]?)", ptx_content)
+    arch = match.group(1) if match else "sm_90a"
+
+    # Write stripped content back if needed
+    if ptx_path.read_text() != ptx_content:
+        ptx_path.write_text(ptx_content)
+
+    # Compile
+    cubin_tmp = ptx_path.with_suffix(".cubin.tmp")
+    try:
+        assert CUTE_DSL_PTXAS_PATH is not None
+        result = subprocess.run(
+            [CUTE_DSL_PTXAS_PATH, f"-arch={arch}", "-O3", "-o", str(cubin_tmp), str(ptx_path)],
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            raise RuntimeError(f"ptxas failed: {result.stderr}")
+
+        cubin_data = cubin_tmp.read_bytes()
+        _log(f"Compiled {ptx_path.name} -> {len(cubin_data)} bytes ({arch})")
+
+        # Save cubin if CUTE_DSL_KEEP_CUBIN is set
+        if os.environ.get("CUTE_DSL_KEEP_CUBIN", "0") == "1":
+            cubin_out = ptx_path.with_suffix(".cubin")
+            cubin_out.write_bytes(cubin_data)
+            _log(f"Saved: {cubin_out}")
+
+        return cubin_data
+    finally:
+        cubin_tmp.unlink(missing_ok=True)
+
+
+def _patched_load_cuda_library(self):
+    """Replacement for _load_cuda_library that uses system ptxas."""
+
+    result = _get_ptx(self)
+    if not result:
+        _log("PTX not found, falling back to embedded ptxas")
+        return _original_load_cuda_library(self)
+
+    ptx_content, ptx_path = result
+
+    try:
+        cubin = _compile_ptx(ptx_path, ptx_content)
+    except Exception as e:
+        _log(f"Compilation failed ({e}), falling back to embedded ptxas")
+        return _original_load_cuda_library(self)
+
+    # Load cubin
+    import cuda.bindings.runtime as cuda_runtime
+
+    err, library = cuda_runtime.cudaLibraryLoadData(cubin, None, None, 0, None, None, 0)
+    if err != cuda_runtime.cudaError_t.cudaSuccess:
+        _log(f"cudaLibraryLoadData failed ({err}), falling back to embedded ptxas")
+        return _original_load_cuda_library(self)
+
+    # Register kernels on all devices
+    _, cuda_load_to_device = self._get_cuda_init_and_load()
+    lib_ptr = ctypes.c_void_p(int(library))
+    dev_id = ctypes.c_int32(0)
+    err_val = ctypes.c_int32(0)
+    args = (ctypes.c_void_p * 3)(
+        ctypes.cast(ctypes.pointer(lib_ptr), ctypes.c_void_p),
+        ctypes.cast(ctypes.pointer(dev_id), ctypes.c_void_p),
+        ctypes.cast(ctypes.pointer(err_val), ctypes.c_void_p),
+    )
+
+    for dev in range(self.num_devices):
+        dev_id.value = dev
+        cuda_load_to_device(args)
+        if err_val.value != 0:
+            _log("cuda_load_to_device failed, falling back to embedded ptxas")
+            return _original_load_cuda_library(self)
+
+    _log(f"Loaded kernel from {ptx_path.name}")
+
+    # Delete PTX if user didn't originally want it kept
+    if not _user_wanted_ptx:
+        ptx_path.unlink(missing_ok=True)
+
+    return [cuda_runtime.cudaLibrary_t(lib_ptr.value)]
+
+
+def patch():
+    """Install system ptxas hook. Call before importing cutlass."""
+    global _original_load_cuda_library, _user_wanted_ptx
+
+    assert CUTE_DSL_PTXAS_PATH is not None
+    if not os.path.isfile(CUTE_DSL_PTXAS_PATH) or not os.access(CUTE_DSL_PTXAS_PATH, os.X_OK):
+        raise RuntimeError(f"ptxas not found: {CUTE_DSL_PTXAS_PATH}")
+
+    # Track if user originally wanted PTX kept
+    _user_wanted_ptx = os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1"
+    # os.environ['CUTE_DSL_KEEP_PTX'] = '1'
+    assert os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1", (
+        "Require CUTE_DSL_KEEP_PTX=1 to use system's ptxas"
+    )
+
+    cls = cutlass.cutlass_dsl.cuda_jit_executor.CudaDialectJitCompiledFunction
+    _original_load_cuda_library = cls._load_cuda_library
+    cls._load_cuda_library = _patched_load_cuda_library
+    _log("Patch applied")
+    return
@@ -0,0 +1,158 @@
+# Copyright (c) 2025, Tri Dao.
+
+from typing import Tuple
+from functools import lru_cache
+
+import torch
+
+try:
+    from triton.tools.disasm import extract
+except ImportError:
+    extract = None
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cutlass_dsl import NumericMeta
+from cutlass.cute.runtime import from_dlpack
+
+StaticTypes = (cutlass.Constexpr, NumericMeta, int, bool, str, float, type(None))
+
+
+load_cubin_module_data_og = cutlass.base_dsl.runtime.cuda.load_cubin_module_data
+cute_compile_og = cute.compile
+
+
+torch2cute_dtype_map = {
+    torch.float16: cutlass.Float16,
+    torch.bfloat16: cutlass.BFloat16,
+    torch.float32: cutlass.Float32,
+    torch.float8_e4m3fn: cutlass.Float8E4M3FN,
+    torch.float8_e5m2: cutlass.Float8E5M2,
+}
+
+
+@lru_cache
+def get_max_active_clusters(cluster_size):
+    return cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size=cluster_size)
+
+
+@lru_cache
+def get_device_capacity(device: torch.device = None) -> Tuple[int, int]:
+    return torch.cuda.get_device_capability(device)
+
+
+def assume_strides_aligned(t):
+    """Assume all strides except the last are divisible by 128 bits.
+
+    Python int strides (e.g., stride=0 from GQA expand) are kept as-is
+    since they're static and don't need alignment assumptions.
+    """
+    divby = 128 // t.element_type.width
+    strides = tuple(s if isinstance(s, int) else cute.assume(s, divby=divby) for s in t.stride[:-1])
+    return (*strides, t.stride[-1])
+
+
+def assume_tensor_aligned(t):
+    """Rebuild a tensor with 128-bit aligned stride assumptions. Passes through None."""
+    if t is None:
+        return None
+    return cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=assume_strides_aligned(t)))
+
+
+def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True):
+    """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1."""
+    if t is None:
+        return None
+    # NOTE: torch 2.9.1 doesn't support fp8 via DLPack but 2.11.0 nightly does
+    # currently export raw bytes as uint8 and tell cutlass correct type
+    # can directly export as fp8 when torch supports it
+    if t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+        tensor = from_dlpack(
+            t.view(torch.uint8).detach(),
+            assumed_align=assumed_align,
+            enable_tvm_ffi=enable_tvm_ffi,
+        )
+        tensor.element_type = (
+            cutlass.Float8E4M3FN if t.dtype == torch.float8_e4m3fn else cutlass.Float8E5M2
+        )
+    else:
+        tensor = from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=enable_tvm_ffi)
+    if fully_dynamic:
+        return tensor.mark_layout_dynamic()
+    if leading_dim == -1:
+        leading_dim = t.ndim - 1
+    return tensor.mark_layout_dynamic(leading_dim=leading_dim)
+
+
+def to_cute_aux_tensor(t, enable_tvm_ffi=True):
+    """Convert torch tensor to cute tensor for TVM FFI, tailored to FlexAttention aux tensors.
+    This allows the user to specify alignment and leading dimension for aux tensors used in
+    custom score_mod callables.
+    """
+    assumed_align: int = getattr(t, "__assumed_align__", None)
+    leading_dim: int = getattr(t, "__leading_dim__", None)
+    fully_dynamic: bool = leading_dim is None
+
+    return to_cute_tensor(
+        t,
+        assumed_align=assumed_align,
+        leading_dim=leading_dim,
+        fully_dynamic=fully_dynamic,
+        enable_tvm_ffi=enable_tvm_ffi,
+    )
+
+
+def get_aux_tensor_metadata(aux_tensors):
+    return tuple(
+        (
+            getattr(t, "__assumed_align__", 0),
+            getattr(t, "__leading_dim__", -1),
+            hasattr(t, "__leading_dim__"),
+        )
+        for t in aux_tensors
+    )
+
+
+def get_broadcast_dims(tensor: torch.Tensor) -> Tuple[bool, ...]:
+    """Return tuple of bools indicating which dims have stride=0 (broadcast).
+
+    This is useful for compile keys since CuTe's mark_layout_dynamic() keeps
+    stride=0 as static, meaning kernels compiled with different broadcast
+    patterns are not interchangeable.
+    """
+    return tuple(s == 0 for s in tensor.stride())
+
+
+# credit: monellz (https://github.com/NVIDIA/cutlass/issues/2658#issuecomment-3630564264)
+def dump_kernel_attributes(compiled_kernel):
+    from cuda.bindings import driver
+    from cutlass.utils import HardwareInfo
+    import torch
+
+    device_id = torch.cuda.current_device()
+    hardware_info = HardwareInfo(device_id=device_id)
+    cubin_data = compiled_kernel.artifacts.CUBIN
+    assert cubin_data is not None, "cubin_data is None, need '--keep-cubin' option when compiling"
+    cuda_library = hardware_info._checkCudaErrors(
+        driver.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0)
+    )
+    kernels = hardware_info._checkCudaErrors(driver.cuLibraryEnumerateKernels(1, cuda_library))
+    kernel = hardware_info._checkCudaErrors(driver.cuKernelGetFunction(kernels[0]))
+    # more metrics: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__EXEC.html#group__CUDA__EXEC_1g5e92a1b0d8d1b82cb00dcfb2de15961b
+    local_size_bytes = hardware_info._checkCudaErrors(
+        driver.cuFuncGetAttribute(
+            driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES,
+            kernel,
+        )
+    )
+    num_regs = hardware_info._checkCudaErrors(
+        driver.cuFuncGetAttribute(
+            driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS,
+            kernel,
+        )
+    )
+
+    print("--- Kernel Info ---")
+    print(f"local_size_bytes: {local_size_bytes}")
+    print(f"num_regs: {num_regs}")
+    print("--- End Kernel Info ---")
@@ -0,0 +1,97 @@
+# Copyright (c) 2025, Tri Dao.
+
+"""Unified FlashAttention logging controlled by a single ``FA_LOG_LEVEL`` env var.
+
+Host-side messages go through Python ``logging`` (logger name ``flash_attn``).
+A default ``StreamHandler`` is attached automatically when ``FA_LOG_LEVEL >= 1``
+so that standalone scripts get output without extra setup; applications that
+configure their own logging can remove or replace it via the standard API.
+
+FA_LOG_LEVEL mapping::
+
+    0  off       nothing logged
+    1  host      host-side summaries only (no kernel printf)
+    2  kernel    host + curated kernel traces
+    3  max       host + all kernel traces (noisy, perf hit)
+
+Set via environment variable::
+
+    FA_LOG_LEVEL=1 python train.py
+
+Device-side ``cute.printf`` calls are compile-time eliminated via
+``cutlass.const_expr`` when the log level is below the callsite threshold,
+so there is zero performance cost when device logging is off.
+Changing the log level after kernel compilation requires a recompile
+(the level participates in the forward compile key).
+"""
+
+import logging
+import os
+import sys
+
+import cutlass.cute as cute
+from cutlass import const_expr
+
+_LOG_LEVEL_NAMES = {"off": 0, "host": 1, "kernel": 2, "max": 3}
+
+
+def _parse_log_level(raw: str) -> int:
+    if raw in _LOG_LEVEL_NAMES:
+        return _LOG_LEVEL_NAMES[raw]
+    try:
+        level = int(raw)
+    except ValueError:
+        return 0
+    return max(0, min(level, 3))
+
+
+_fa_log_level: int = _parse_log_level(os.environ.get("FA_LOG_LEVEL", "0"))
+
+_logger = logging.getLogger("flash_attn")
+_logger.addHandler(logging.NullHandler())
+_default_handler: logging.Handler | None = None
+
+
+def _configure_default_handler() -> None:
+    global _default_handler
+    if _fa_log_level >= 1:
+        if _default_handler is None:
+            _default_handler = logging.StreamHandler(sys.stdout)
+            _default_handler.setFormatter(logging.Formatter("[FA] %(message)s"))
+            _logger.addHandler(_default_handler)
+        _logger.setLevel(logging.DEBUG)
+    else:
+        if _default_handler is not None:
+            _logger.removeHandler(_default_handler)
+            _default_handler = None
+        _logger.setLevel(logging.WARNING)
+
+
+_configure_default_handler()
+
+
+def get_fa_log_level() -> int:
+    return _fa_log_level
+
+
+def set_fa_log_level(level: int | str) -> None:
+    """Set the FA log level programmatically.
+
+    Host logging takes effect immediately.  Device logging changes only
+    affect kernels compiled after this call (new compile-key selection).
+    """
+    global _fa_log_level
+    if isinstance(level, str):
+        level = _parse_log_level(level)
+    _fa_log_level = max(0, min(int(level), 3))
+    _configure_default_handler()
+
+
+def fa_log(level: int, msg: str):
+    if _fa_log_level >= level:
+        _logger.info(msg)
+
+
+def fa_printf(level: int, fmt, *args):
+    if const_expr(_fa_log_level >= level):
+        cute.printf(fmt, *args)
@@ -0,0 +1,21 @@
+# Copyright (c) 2025, Tri Dao.
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32
+
+
+@cute.jit
+def clz(x: Int32) -> Int32:
+    # for i in cutlass.range_constexpr(32):
+    #     if (1 << (31 - i)) & x:
+    #         return Int32(i)
+    # return Int32(32)
+    # Early exit is not supported yet
+    res = Int32(32)
+    done = False
+    for i in cutlass.range(32):
+        if ((1 << (31 - i)) & x) and not done:
+            res = Int32(i)
+            done = True
+    return res
@@ -0,0 +1,587 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_bwd_postprocess_kernel.h
+# from Cutlass C++ to Cute-DSL.
+import math
+from typing import Callable, Optional, Type
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils.hopper_helpers as sm90_utils_basic
+import cutlass.utils.blackwell_helpers as sm100_utils_basic
+from cutlass.cute.nvgpu import cpasync, warp, warpgroup
+from cutlass import Float32, const_expr
+from cutlass.utils import LayoutEnum
+
+from quack import copy_utils
+from quack import layout_utils
+from quack import sm90_utils
+
+from flash_attn.cute import utils
+from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
+from flash_attn.cute import ampere_helpers as sm80_utils
+from flash_attn.cute.seqlen_info import SeqlenInfoQK
+import cutlass.cute.nvgpu.tcgen05 as tcgen05
+from quack.cute_dsl_utils import ParamsBase
+from flash_attn.cute.tile_scheduler import (
+    SingleTileScheduler,
+    SingleTileVarlenScheduler,
+    TileSchedulerArguments,
+)
+
+
+class FlashAttentionBackwardPostprocess:
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        head_dim: int,
+        arch: int,
+        tile_m: int = 128,
+        num_threads: int = 256,
+        AtomLayoutMdQ: int = 1,
+        dQ_swapAB: bool = False,
+        use_2cta_instrs: bool = False,
+        cluster_size: int = 1,  # for varlen offsets
+    ):
+        """
+        :param head_dim: head dimension
+        :type head_dim: int
+        :param tile_m: m block size
+        :type tile_m: int
+        """
+        self.dtype = dtype
+        self.tile_m = tile_m
+        assert arch // 10 in [8, 9, 10, 11, 12], (
+            "Only Ampere (8.x), Hopper (9.x), and Blackwell (10.x, 11.x, 12.x) are supported"
+        )
+        self.arch = arch
+        # padding head_dim to a multiple of 32 as k_block_size
+        hdim_multiple_of = 32
+        self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of)
+        self.check_hdim_oob = head_dim != self.tile_hdim
+        self.num_threads = num_threads
+        self.AtomLayoutMdQ = AtomLayoutMdQ
+        self.dQ_swapAB = dQ_swapAB
+        self.use_2cta_instrs = use_2cta_instrs and arch // 10 in [10, 11] and head_dim != 64
+        self.cluster_size = cluster_size
+
+    @staticmethod
+    def can_implement(dtype, head_dim, tile_m, num_threads) -> bool:
+        """Check if the kernel can be implemented with the given parameters.
+
+        :param dtype: data type
+        :type dtype: cutlass.Numeric
+        :param head_dim: head dimension
+        :type head_dim: int
+        :param tile_m: m block size
+        :type tile_m: int
+
+        :return: True if the kernel can be implemented, False otherwise
+        :rtype: bool
+        """
+        if dtype not in [cutlass.Float16, cutlass.BFloat16]:
+            return False
+        if head_dim % 8 != 0:
+            return False
+        if num_threads % 32 != 0:
+            return False
+        return True
+
+    def _get_tiled_mma(self):
+        if const_expr(self.arch // 10 in [8, 12]):
+            num_mma_warps = self.num_threads // 32
+            atom_layout_dQ = (
+                (self.AtomLayoutMdQ, num_mma_warps // self.AtomLayoutMdQ, 1)
+                if const_expr(not self.dQ_swapAB)
+                else (num_mma_warps // self.AtomLayoutMdQ, self.AtomLayoutMdQ, 1)
+            )
+            tiled_mma = cute.make_tiled_mma(
+                warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)),
+                atom_layout_dQ,
+                permutation_mnk=(atom_layout_dQ[0] * 16, atom_layout_dQ[1] * 16, 16),
+            )
+        elif const_expr(self.arch // 10 == 9):
+            num_wg_mma = self.num_threads // 128
+            atom_layout_dQ = (self.AtomLayoutMdQ, num_wg_mma // self.AtomLayoutMdQ)
+            tiler_mn_dQ = (self.tile_m // atom_layout_dQ[0], self.tile_hdim // atom_layout_dQ[1])
+            tiled_mma = sm90_utils_basic.make_trivial_tiled_mma(
+                self.dtype,
+                self.dtype,
+                warpgroup.OperandMajorMode.K,  # These don't matter, we only care about the accum
+                warpgroup.OperandMajorMode.K,
+                Float32,
+                atom_layout_mnk=(atom_layout_dQ if not self.dQ_swapAB else atom_layout_dQ[::-1])
+                + (1,),
+                tiler_mn=tiler_mn_dQ if not self.dQ_swapAB else tiler_mn_dQ[::-1],
+            )
+        else:
+            cta_group = tcgen05.CtaGroup.ONE
+            tiled_mma = sm100_utils_basic.make_trivial_tiled_mma(
+                self.dtype,
+                tcgen05.OperandMajorMode.MN,  # dS_major_mode
+                tcgen05.OperandMajorMode.MN,  # Kt_major_mode
+                Float32,
+                cta_group,
+                (self.tile_m, self.tile_hdim),
+            )
+        if const_expr(self.arch // 10 in [8, 9, 12]):
+            assert self.num_threads == tiled_mma.size
+        return tiled_mma
+
+    def _setup_attributes(self):
+        # ///////////////////////////////////////////////////////////////////////////////
+        # GMEM Tiled copy:
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Thread layouts for copies
+        universal_copy_bits = 128
+        async_copy_elems_accum = universal_copy_bits // Float32.width
+        atom_async_copy_accum = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            Float32,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        # We don't do bound checking for the gmem -> smem load so we just assert here.
+        assert (self.tile_m * self.tile_hdim // async_copy_elems_accum) % self.num_threads == 0
+        self.g2s_tiled_copy_dQaccum = cute.make_tiled_copy_tv(
+            atom_async_copy_accum,
+            cute.make_layout(self.num_threads),
+            cute.make_layout(async_copy_elems_accum),
+        )
+        num_s2r_copy_elems = 1 if const_expr(self.arch // 10 in [8, 12]) else 4
+        if const_expr(self.arch // 10 in [8, 12]):
+            self.s2r_tiled_copy_dQaccum = copy_utils.tiled_copy_1d(
+                Float32, self.num_threads, num_s2r_copy_elems
+            )
+            self.sdQaccum_layout = cute.make_layout(self.tile_m * self.tile_hdim)
+        elif const_expr(self.arch // 10 == 9):
+            num_threads_per_warp_group = 128
+            num_wg_mma = self.num_threads // 128
+            self.s2r_tiled_copy_dQaccum = cute.make_tiled_copy_tv(
+                cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128),
+                cute.make_layout((num_threads_per_warp_group, num_wg_mma)),  # thr_layout
+                cute.make_layout(128 // Float32.width),  # val_layout
+            )
+            self.sdQaccum_layout = cute.make_layout(
+                (self.tile_m * self.tile_hdim // num_wg_mma, num_wg_mma)
+            )
+        else:
+            self.dQ_reduce_ncol = 32
+            dQaccum_reduce_stage = self.tile_hdim // self.dQ_reduce_ncol
+            assert self.num_threads == 128  # TODO: currently hard-coded
+            self.s2r_tiled_copy_dQaccum = copy_utils.tiled_copy_1d(
+                Float32, self.num_threads, num_s2r_copy_elems
+            )
+            self.sdQaccum_layout = cute.make_layout(
+                (self.tile_m * self.tile_hdim // dQaccum_reduce_stage, dQaccum_reduce_stage)
+            )
+
+        num_copy_elems = 128 // self.dtype.width
+        threads_per_row = math.gcd(128, self.tile_hdim) // num_copy_elems
+        self.gmem_tiled_copy_dQ = copy_utils.tiled_copy_2d(
+            self.dtype, threads_per_row, self.num_threads, num_copy_elems
+        )
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Shared memory layout: dQ
+        # ///////////////////////////////////////////////////////////////////////////////
+        # We can't just use kHeadDim here. E.g. if MMA shape is 64 x 96 but split across 2 WGs,
+        # then setting kBlockKSmem to 32 will cause "Static shape_div failure".
+        # We want to treat it as 64 x 48, so kBlockKSmem should be 16.
+        mma_shape_n = self.tiled_mma.get_tile_size(1)
+        if const_expr(self.arch // 10 in [8, 12]):
+            sdQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, mma_shape_n)
+            self.sdQ_layout = cute.tile_to_shape(
+                sdQ_layout_atom, (self.tile_m, self.tile_hdim), (0, 1)
+            )
+        elif const_expr(self.arch // 10 == 9):
+            wg_d_dQ = num_wg_mma // self.AtomLayoutMdQ
+            self.sdQ_layout = sm90_utils.make_smem_layout(
+                self.dtype,
+                LayoutEnum.ROW_MAJOR,
+                (self.tile_m, self.tile_hdim),
+                major_mode_size=self.tile_hdim // wg_d_dQ,
+            )
+        else:
+            # TODO: this is hard-coded for hdim 128
+            self.sdQ_layout = sm100_utils_basic.make_smem_layout_epi(
+                self.dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_hdim), 1
+            )
+
+    @cute.jit
+    def __call__(
+        self,
+        mdQaccum: cute.Tensor,
+        mdQ: cute.Tensor,
+        scale: cutlass.Float32,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        # Get the data type and check if it is fp16 or bf16
+        if const_expr(mdQ.element_type not in [cutlass.Float16, cutlass.BFloat16]):
+            raise TypeError("Only Float16 or BFloat16 is supported")
+        if const_expr(mdQaccum is not None):
+            if const_expr(mdQaccum.element_type not in [cutlass.Float32]):
+                raise TypeError("dQaccum tensor must be Float32")
+
+        mdQaccum, mdQ = [assume_tensor_aligned(t) for t in (mdQaccum, mdQ)]
+
+        self.tiled_mma = self._get_tiled_mma()
+        self._setup_attributes()
+
+        smem_size = max(
+            cute.size_in_bytes(cutlass.Float32, self.sdQaccum_layout),
+            cute.size_in_bytes(self.dtype, self.sdQ_layout),
+        )
+
+        if const_expr(mCuSeqlensQ is not None):
+            TileScheduler = SingleTileVarlenScheduler
+            num_head = mdQ.shape[1]
+            num_batch = mCuSeqlensQ.shape[0] - 1
+            num_block = cute.ceil_div(mdQ.shape[0], self.tile_m)
+        else:
+            TileScheduler = SingleTileScheduler
+            num_head = mdQ.shape[2]
+            num_batch = mdQ.shape[0]
+            num_block = cute.ceil_div(mdQ.shape[1], self.tile_m)
+
+        tile_sched_args = TileSchedulerArguments(
+            num_block=num_block,
+            num_head=num_head,
+            num_batch=num_batch,
+            num_splits=1,
+            seqlen_k=0,
+            headdim=mdQ.shape[2],
+            headdim_v=0,
+            total_q=mdQ.shape[0],
+            tile_shape_mn=(self.tile_m, 1),
+            mCuSeqlensQ=mCuSeqlensQ,
+            mSeqUsedQ=mSeqUsedQ,
+        )
+
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+
+        # grid_dim: (m_block, num_head, batch_size)
+        self.kernel(
+            mdQaccum,
+            mdQ,
+            mCuSeqlensQ,
+            mSeqUsedQ,
+            scale,
+            self.tiled_mma,
+            self.dQ_swapAB,
+            self.sdQaccum_layout,
+            self.sdQ_layout,
+            self.g2s_tiled_copy_dQaccum,
+            self.s2r_tiled_copy_dQaccum,
+            self.gmem_tiled_copy_dQ,
+            tile_sched_params,
+            TileScheduler,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            smem=smem_size,
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mdQaccum: cute.Tensor,
+        mdQ: cute.Tensor,
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        scale: cutlass.Float32,
+        tiled_mma: cute.TiledMma,
+        dQ_swapAB: cutlass.Constexpr,
+        sdQaccum_layout: cute.Layout,
+        sdQ_layout: cute.ComposedLayout,
+        g2s_tiled_copy_dQaccum: cute.TiledCopy,
+        s2r_tiled_copy_dQaccum: cute.TiledCopy,
+        gmem_tiled_copy_dQ: cute.TiledCopy,
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
+    ):
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Get shared memory buffer
+        # ///////////////////////////////////////////////////////////////////////////////
+        smem = cutlass.utils.SmemAllocator()
+        sdQaccum = smem.allocate_tensor(cutlass.Float32, sdQaccum_layout, byte_alignment=1024)
+        sdQaccum_flat = cute.make_tensor(sdQaccum.iterator, cute.make_layout(cute.size(sdQaccum)))
+        if const_expr(self.arch // 10 in [8, 9, 12]):
+            sdQ = cute.make_tensor(cute.recast_ptr(sdQaccum.iterator, dtype=self.dtype), sdQ_layout)
+        else:
+            # extra stage dimension
+            sdQ = cute.make_tensor(
+                cute.recast_ptr(sdQaccum.iterator, sdQ_layout.inner, dtype=self.dtype),
+                sdQ_layout.outer,
+            )[None, None, 0]
+        sdQt = layout_utils.transpose_view(sdQ)
+
+        # Thread index, block index
+        tidx, _, _ = cute.arch.thread_idx()
+
+        tile_scheduler = TileScheduler.create(tile_sched_params)
+        work_tile = tile_scheduler.initial_work_tile_info()
+
+        m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+
+        if work_tile.is_valid_tile:
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Get the appropriate tiles for this thread block.
+            # ///////////////////////////////////////////////////////////////////////////////
+
+            seqlen = SeqlenInfoQK.create(
+                batch_idx,
+                mdQ.shape[1],
+                0,
+                mCuSeqlensQ=mCuSeqlensQ,
+                mCuSeqlensK=None,
+                mSeqUsedQ=mSeqUsedQ,
+                mSeqUsedK=None,
+                tile_m=self.tile_m * self.cluster_size,
+            )
+            if const_expr(not seqlen.has_cu_seqlens_q):
+                mdQ_cur = mdQ[batch_idx, None, head_idx, None]
+                mdQaccum_cur = mdQaccum[batch_idx, head_idx, None]
+                head_dim = mdQ.shape[3]
+            else:
+                padded_offset_q = seqlen.padded_offset_q
+                mdQ_cur = cute.domain_offset((seqlen.offset_q, 0), mdQ[None, head_idx, None])
+                mdQaccum_cur = cute.domain_offset(
+                    (padded_offset_q * self.tile_hdim,), mdQaccum[head_idx, None]
+                )
+                head_dim = mdQ.shape[2]
+
+                # HACK: Compiler doesn't seem to recognize that padding
+                # by padded_offset_q * self.tile_hdim keeps alignment
+                # since statically divisible by 4
+
+                mdQaccum_cur_ptr = cute.make_ptr(
+                    dtype=mdQaccum_cur.element_type,
+                    value=mdQaccum_cur.iterator.toint(),
+                    mem_space=mdQaccum_cur.iterator.memspace,
+                    assumed_align=mdQaccum.iterator.alignment,
+                )
+                mdQaccum_cur = cute.make_tensor(mdQaccum_cur_ptr, mdQaccum_cur.layout)
+
+            gdQaccum = cute.local_tile(mdQaccum_cur, (self.tile_m * self.tile_hdim,), (m_block,))
+            gdQ = cute.local_tile(mdQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0))
+
+            seqlen_q = seqlen.seqlen_q
+            seqlen_q_rounded = cute.round_up(seqlen_q, self.tile_m)
+
+            if const_expr(self.arch // 10 in [10, 11] and self.use_2cta_instrs):
+                # 2-CTA: remap dQaccum layout into TMEM view before writing sdQ
+                num_reduce_threads = self.num_threads
+                thr_mma_dsk = tiled_mma.get_slice(tidx)
+                dQacc_shape = thr_mma_dsk.partition_shape_C((self.tile_m, self.tile_hdim))
+                tdQtdQ = thr_mma_dsk.make_fragment_C(dQacc_shape)
+                tdQtdQ = cute.make_tensor(tdQtdQ.iterator, tdQtdQ.layout)
+
+                tmem_load_atom = cute.make_copy_atom(
+                    tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.dQ_reduce_ncol)), Float32
+                )
+                tiled_tmem_ld = tcgen05.make_tmem_copy(tmem_load_atom, tdQtdQ)
+                thr_tmem_ld = tiled_tmem_ld.get_slice(tidx)
+
+                cdQ = cute.make_identity_tensor((self.tile_m, self.tile_hdim))
+                tdQcdQ = thr_mma_dsk.partition_C(cdQ)
+                tdQcdQ_tensor = cute.make_tensor(tdQcdQ.iterator, tdQcdQ.layout)
+                tdQrdQ = thr_tmem_ld.partition_D(tdQcdQ_tensor)
+
+                tiled_copy_accum = s2r_tiled_copy_dQaccum
+                g2s_thr_copy = tiled_copy_accum.get_slice(tidx)
+
+                # S -> R
+                tdQrdQ_fp32 = cute.make_rmem_tensor(tdQrdQ.shape, cutlass.Float32)
+                tdQrdQ_s2r = cute.make_tensor(tdQrdQ_fp32.iterator, tdQrdQ_fp32.shape)
+
+                smem_copy_atom = sm100_utils_basic.get_smem_store_op(
+                    LayoutEnum.ROW_MAJOR, self.dtype, cutlass.Float32, tiled_tmem_ld
+                )
+                r2s_tiled_copy = cute.make_tiled_copy(
+                    smem_copy_atom,
+                    layout_tv=tiled_tmem_ld.layout_dst_tv_tiled,
+                    tiler_mn=tiled_tmem_ld.tiler_mn,
+                )
+                tdQsdQ_r2s = thr_tmem_ld.partition_D(thr_mma_dsk.partition_C(sdQ))
+                tdQrdQ_r2s = cute.make_rmem_tensor(tdQsdQ_r2s.shape, self.dtype)
+
+                num_stages = cute.size(tdQrdQ_fp32, mode=[1])
+                stage_stride = self.dQ_reduce_ncol
+                row_groups = 2
+                assert num_stages % row_groups == 0
+                assert num_reduce_threads % row_groups == 0
+                stage_groups = num_stages // row_groups
+                threads_per_row_group = num_reduce_threads // row_groups
+                stage_loads = tuple((row_group, row_group) for row_group in range(row_groups))
+                stage_iters = tuple(
+                    (row_group, row_group * threads_per_row_group)
+                    for row_group in range(row_groups)
+                )
+                s2r_lane = tidx % threads_per_row_group
+                s2r_buf = tidx // threads_per_row_group
+
+                gdQaccum_layout_g2s = cute.make_layout(
+                    shape=(self.tile_m * self.dQ_reduce_ncol, 1), stride=(1, 0)
+                )
+                sdQaccum_g2s = g2s_thr_copy.partition_D(sdQaccum)
+
+                # G -> S
+                for stage_group in cutlass.range_constexpr(stage_groups):
+                    for stage_offset, smem_buf in stage_loads:
+                        stage_idx = stage_group + stage_offset * stage_groups
+                        gdQaccum_stage = cute.local_tile(
+                            gdQaccum,
+                            (self.tile_m * self.dQ_reduce_ncol,),
+                            (stage_idx,),
+                        )
+                        gdQaccum_stage_g2s = cute.make_tensor(
+                            gdQaccum_stage.iterator,
+                            gdQaccum_layout_g2s,
+                        )
+                        tdQgdQ = g2s_thr_copy.partition_S(gdQaccum_stage_g2s)
+                        cute.copy(
+                            g2s_thr_copy,
+                            tdQgdQ[None, None, 0],
+                            sdQaccum_g2s[None, None, smem_buf],
+                        )
+
+                    cute.arch.fence_view_async_shared()
+                    cute.arch.barrier(barrier_id=6, number_of_threads=num_reduce_threads)
+
+                    # S -> R
+                    for stage_offset, lane_offset in stage_iters:
+                        stage_idx = stage_group + stage_offset * stage_groups
+                        s2r_src_tidx = s2r_lane + lane_offset
+                        s2r_thr_copy = tiled_copy_accum.get_slice(s2r_src_tidx)
+                        sdQaccum_src = s2r_thr_copy.partition_S(sdQaccum)[None, None, s2r_buf]
+
+                        tdQrdQ_s2r_cpy = tdQrdQ_s2r[None, stage_idx, None, None]
+                        tdQrdQ_r2s_cpy = cute.make_tensor(
+                            tdQrdQ_s2r_cpy.iterator, cute.make_layout(sdQaccum_src.shape)
+                        )
+                        cute.copy(s2r_thr_copy, sdQaccum_src, tdQrdQ_r2s_cpy)
+                        cute.arch.fence_view_async_shared()
+                        cute.arch.barrier(barrier_id=7, number_of_threads=num_reduce_threads)
+
+                        # R -> S
+                        stage_lo = stage_idx % stage_stride
+                        stage_hi = stage_idx // stage_stride
+                        tdQrdQ_r2s_cpy = cute.make_tensor(
+                            cute.recast_ptr(tdQrdQ_r2s_cpy.iterator),
+                            tdQrdQ_r2s[((None, 0), (stage_lo, stage_hi), 0, 0)].shape,
+                        )
+                        dQ_vec = tdQrdQ_r2s_cpy.load() * scale
+                        tdQrdQ_r2s[((None, 0), (stage_lo, stage_hi), 0, 0)].store(
+                            dQ_vec.to(self.dtype)
+                        )
+
+                # R -> S
+                cute.copy(
+                    r2s_tiled_copy,
+                    tdQrdQ_r2s[None, None, None, 0],
+                    tdQsdQ_r2s[None, None, None, 0],
+                )
+                cute.arch.fence_view_async_shared()
+                cute.arch.barrier(barrier_id=8, number_of_threads=num_reduce_threads)
+            else:
+                # Step 1: load dQaccum from gmem to smem
+                g2s_thr_copy_dQaccum = g2s_tiled_copy_dQaccum.get_slice(tidx)
+                tdQgdQaccum = g2s_thr_copy_dQaccum.partition_S(gdQaccum)
+                tdQsdQaccumg2s = g2s_thr_copy_dQaccum.partition_D(sdQaccum_flat)
+                cute.copy(g2s_tiled_copy_dQaccum, tdQgdQaccum, tdQsdQaccumg2s)
+                cute.arch.cp_async_commit_group()
+                cute.arch.cp_async_wait_group(0)
+                cute.arch.barrier()
+
+                # Step 2: load dQ from smem to rmem
+                s2r_thr_copy_dQaccum = s2r_tiled_copy_dQaccum.get_slice(tidx)
+                tdQsdQaccum = s2r_thr_copy_dQaccum.partition_S(sdQaccum)
+                tile_shape = (self.tile_m, self.tile_hdim)
+                acc = None
+                tiled_copy_t2r = None
+                if const_expr(self.arch // 10 in [8, 9, 12]):
+                    acc_shape = tiled_mma.partition_shape_C(
+                        tile_shape if const_expr(not dQ_swapAB) else tile_shape[::-1]
+                    )
+                    acc = cute.make_rmem_tensor(acc_shape, cutlass.Float32)
+                    assert cute.size(acc) == cute.size(tdQsdQaccum)
+                else:
+                    thr_mma = tiled_mma.get_slice(0)  # 1-CTA
+                    dQacc_shape = tiled_mma.partition_shape_C((self.tile_m, self.tile_hdim))
+                    tdQtdQ = tiled_mma.make_fragment_C(dQacc_shape)
+                    tdQcdQ = thr_mma.partition_C(
+                        cute.make_identity_tensor((self.tile_m, self.tile_hdim))
+                    )
+                    tmem_load_atom = cute.make_copy_atom(
+                        tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.dQ_reduce_ncol)),
+                        Float32,
+                    )
+                    tiled_copy_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tdQtdQ)
+                    thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+                    tdQrdQ_t2r_shape = thr_copy_t2r.partition_D(tdQcdQ).shape
+                    acc = cute.make_rmem_tensor(tdQrdQ_t2r_shape, Float32)
+                tdQrdQaccum = cute.make_tensor(acc.iterator, cute.make_layout(tdQsdQaccum.shape))
+                cute.autovec_copy(tdQsdQaccum, tdQrdQaccum)
+                # Convert tdQrdQaccum from fp32 to fp16/bf16
+                rdQ = cute.make_fragment_like(acc, self.dtype)
+                rdQ.store((acc.load() * scale).to(self.dtype))
+
+                # Step 3: Copy dQ from register to smem
+                cute.arch.barrier()  # make sure all threads have finished loading dQaccum
+                if const_expr(self.arch // 10 in [8, 9, 12]):
+                    copy_atom_r2s_dQ = utils.get_smem_store_atom(
+                        self.arch, self.dtype, transpose=self.dQ_swapAB
+                    )
+                    tiled_copy_r2s_dQ = cute.make_tiled_copy_C(copy_atom_r2s_dQ, tiled_mma)
+                else:
+                    # copy_atom_r2s_dQ = sm100_utils_basic.get_smem_store_op(
+                    #     LayoutEnum.ROW_MAJOR, self.dtype, Float32, tiled_copy_t2r,
+                    # )
+                    # tiled_copy_r2s_dQ = cute.make_tiled_copy_D(copy_atom_r2s_dQ, tiled_copy_t2r)
+                    thr_layout_r2s_dQ = cute.make_layout((self.num_threads, 1))  # 128 threads
+                    val_layout_r2s_dQ = cute.make_layout((1, 128 // self.dtype.width))
+                    copy_atom_r2s_dQ = cute.make_copy_atom(
+                        cute.nvgpu.CopyUniversalOp(),
+                        self.dtype,
+                        num_bits_per_copy=128,
+                    )
+                    tiled_copy_r2s_dQ = cute.make_tiled_copy_tv(
+                        copy_atom_r2s_dQ, thr_layout_r2s_dQ, val_layout_r2s_dQ
+                    )
+                thr_copy_r2s_dQ = tiled_copy_r2s_dQ.get_slice(tidx)
+                cdQ = cute.make_identity_tensor((self.tile_m, self.tile_hdim))
+                if const_expr(self.arch // 10 in [8, 9, 12]):
+                    taccdQrdQ = thr_copy_r2s_dQ.retile(rdQ)
+                else:
+                    taccdQcdQ_shape = thr_copy_r2s_dQ.partition_S(cdQ).shape
+                    taccdQrdQ = cute.make_tensor(rdQ.iterator, taccdQcdQ_shape)
+                taccdQsdQ = thr_copy_r2s_dQ.partition_D(
+                    sdQ if const_expr(not self.dQ_swapAB) else sdQt
+                )
+                cute.copy(thr_copy_r2s_dQ, taccdQrdQ, taccdQsdQ)
+
+            # Step 4: Copy dQ from smem to register to prepare for coalesced write to gmem
+            cute.arch.barrier()  # make sure all smem stores are done
+            gmem_thr_copy_dQ = gmem_tiled_copy_dQ.get_slice(tidx)
+            tdQgdQ = gmem_thr_copy_dQ.partition_S(gdQ)
+            tdQsdQ = gmem_thr_copy_dQ.partition_D(sdQ)
+            tdQrdQ = cute.make_fragment_like(tdQsdQ, self.dtype)
+            # TODO: check OOB when reading from smem if kBlockM isn't evenly tiled
+            cute.autovec_copy(tdQsdQ, tdQrdQ)
+
+            # Step 5: Copy dQ from register to gmem
+            tdQcdQ = gmem_thr_copy_dQ.partition_S(cdQ)
+            tdQpdQ = utils.predicate_k(tdQcdQ, limit=head_dim)
+            for rest_m in cutlass.range(cute.size(tdQrdQ.shape[1]), unroll_full=True):
+                if tdQcdQ[0, rest_m, 0][0] < seqlen_q - m_block * self.tile_m:
+                    cute.copy(
+                        gmem_tiled_copy_dQ,
+                        tdQrdQ[None, rest_m, None],
+                        tdQgdQ[None, rest_m, None],
+                        pred=tdQpdQ[None, rest_m, None],
+                    )
@@ -0,0 +1,465 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_bwd_preprocess_kernel.h
+# from Cutlass C++ to Cute-DSL.
+#
+# Computes D_i = (dO_i * O_i).sum(dim=-1), optionally adjusted for LSE gradient:
+#   D'_i = D_i - dLSE_i
+# This works because in the backward pass:
+#   dS_ij = P_ij * (dP_ij - D_i)                     [standard]
+# When LSE is differentiable, d(loss)/d(S_ij) gets an extra term dLSE_i * P_ij
+# (since d(LSE_i)/d(S_ij) = P_ij), giving:
+#   dS_ij = P_ij * (dP_ij - D_i) + dLSE_i * P_ij
+#         = P_ij * (dP_ij - (D_i - dLSE_i))
+# So the main backward kernel is unchanged; we just replace D with D' = D - dLSE here.
+import math
+import operator
+from functools import partial
+from typing import Callable, Type, Optional
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, const_expr
+from cutlass.cutlass_dsl import Arch, BaseDSL
+
+from quack import copy_utils, layout_utils
+
+from flash_attn.cute import utils
+from flash_attn.cute.seqlen_info import SeqlenInfo
+from quack.cute_dsl_utils import ParamsBase
+from flash_attn.cute.tile_scheduler import (
+    SingleTileScheduler,
+    SingleTileVarlenScheduler,
+    TileSchedulerArguments,
+)
+from flash_attn.cute.pack_gqa import pack_gqa_layout
+
+
+class FlashAttentionBackwardPreprocess:
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        head_dim: int,
+        head_dim_v: int,
+        tile_m: int = 128,
+        num_threads: int = 256,
+        use_padded_offsets: bool = True,
+        nheads_major: bool = False,
+        pack_gqa: bool = False,
+        qhead_per_kvhead: int = 1,
+        nheads_kv: int = 1,
+    ):
+        """
+        All contiguous dimensions must be at least 16 bytes aligned which indicates the head dimension
+        should be a multiple of 8.
+
+        :param head_dim: head dimension
+        :type head_dim: int
+        :param tile_m: m block size
+        :type tile_m: int
+        :param num_threads: number of threads
+        :type num_threads: int
+        """
+        self.use_pdl = BaseDSL._get_dsl().get_arch_enum() >= Arch.sm_90a
+        self.dtype = dtype
+        self.tile_m = tile_m
+        # padding head_dim to a multiple of 32 as k_block_size
+        hdim_multiple_of = 32
+        self.head_dim_padded = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of)
+        self.head_dim_v_padded = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of)
+        self.check_hdim_v_oob = head_dim_v != self.head_dim_v_padded
+        self.num_threads = num_threads
+        self.use_padded_offsets = use_padded_offsets
+        self.nheads_major = nheads_major
+        self.pack_gqa = pack_gqa
+        self.qhead_per_kvhead = qhead_per_kvhead
+        self.nheads_kv = nheads_kv
+
+    @staticmethod
+    def can_implement(dtype, head_dim, tile_m, num_threads) -> bool:
+        """Check if the kernel can be implemented with the given parameters.
+
+        :param dtype: data type
+        :type dtype: cutlass.Numeric
+        :param head_dim: head dimension
+        :type head_dim: int
+        :param tile_m: m block size
+        :type tile_m: int
+        :param num_threads: number of threads
+        :type num_threads: int
+
+        :return: True if the kernel can be implemented, False otherwise
+        :rtype: bool
+        """
+        if dtype not in [cutlass.Float16, cutlass.BFloat16]:
+            return False
+        if head_dim % 8 != 0:
+            return False
+        if num_threads % 32 != 0:
+            return False
+        if num_threads < tile_m:  # For multiplying lse with log2
+            return False
+        return True
+
+    def _setup_attributes(self):
+        # ///////////////////////////////////////////////////////////////////////////////
+        # GMEM Tiled copy:
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Thread layouts for copies
+        # We want kBlockKGmem to be a power of 2 so that when we do the summing,
+        # it's just between threads in the same warp
+        gmem_k_block_size = (
+            128
+            if self.head_dim_v_padded % 128 == 0
+            else (
+                64
+                if self.head_dim_v_padded % 64 == 0
+                else (32 if self.head_dim_v_padded % 32 == 0 else 16)
+            )
+        )
+        num_copy_elems = 128 // self.dtype.width
+        threads_per_row = gmem_k_block_size // num_copy_elems
+        self.gmem_tiled_copy_O = copy_utils.tiled_copy_2d(
+            self.dtype, threads_per_row, self.num_threads, num_copy_elems
+        )
+        universal_copy_bits = 128
+        num_copy_elems_dQaccum = universal_copy_bits // Float32.width
+        assert (
+            self.tile_m * self.head_dim_padded // num_copy_elems_dQaccum
+        ) % self.num_threads == 0
+        self.gmem_tiled_copy_dQaccum = copy_utils.tiled_copy_1d(
+            Float32, self.num_threads, num_copy_elems_dQaccum
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        mO: cute.Tensor,  # (batch, seqlen, nheads, head_dim_v) or (total_q, nheads, head_dim_v)
+        mdO: cute.Tensor,  # same shape as mO
+        mPdPsum: cute.Tensor,  # (batch, nheads, seqlen_padded) or (nheads, total_q_padded)
+        mLSE: Optional[cute.Tensor],  # (batch, nheads, seqlen) or (nheads, total_q)
+        mLSElog2: Optional[cute.Tensor],  # same shape as mPdPsum
+        # (batch, nheads, seqlen_padded * head_dim_v) or (nheads, total_q_padded * head_dim_v)
+        mdQaccum: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],  # (batch + 1,)
+        mSeqUsedQ: Optional[cute.Tensor],  # (batch,)
+        mdLSE: Optional[cute.Tensor],  # (batch, nheads, seqlen) or (nheads, total_q)
+        mRowMax: Optional[cute.Tensor],  # (b, s, n, h) or (t, n, h)
+        mScaleP: Optional[cute.Tensor],  # == mRowMax
+        softmax_scale: Float32,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        # Get the data type and check if it is fp16 or bf16
+        if const_expr(not (mO.element_type == mdO.element_type)):
+            raise TypeError("All tensors must have the same data type")
+        if const_expr(mO.element_type not in [cutlass.Float16, cutlass.BFloat16]):
+            raise TypeError("Only Float16 or BFloat16 is supported")
+        if const_expr(mPdPsum.element_type not in [Float32]):
+            raise TypeError("PdPsum tensor must be Float32")
+        if const_expr(mdQaccum is not None):
+            assert self.nheads_major is False
+            assert self.pack_gqa is False
+            assert self.use_padded_offsets is True
+            if const_expr(mdQaccum.element_type not in [Float32]):
+                raise TypeError("dQaccum tensor must be Float32")
+        if const_expr(mLSE is not None):
+            if const_expr(mLSE.element_type not in [Float32]):
+                raise TypeError("LSE tensor must be Float32")
+        if const_expr(mLSElog2 is not None):
+            if const_expr(mLSElog2.element_type not in [Float32]):
+                raise TypeError("LSElog2 tensor must be Float32")
+        if const_expr(mdLSE is not None):
+            if const_expr(mdLSE.element_type not in [Float32]):
+                raise TypeError("dLSE tensor must be Float32")
+        if const_expr(mScaleP is not None):
+            assert self.nheads_major is True
+            assert self.pack_gqa is True
+            assert mRowMax is not None
+            if const_expr(mScaleP.element_type not in [Float32]):
+                raise TypeError("ScaleP tensor must be Float32")
+            if const_expr(mRowMax.element_type not in [Float32]):
+                raise TypeError("RowMax tensor must be Float32")
+
+        self._setup_attributes()
+
+        # (b, s, h, d)  -> (s, d, h, b)  or
+        # (total, h, d) -> (total, d, h)
+        QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
+        mO, mdO = [
+            cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=QO_layout_transpose))
+            for mX in (mO, mdO)
+        ]
+
+        if const_expr(not self.nheads_major):
+            # (batch, nheads, seqlen) -> (seqlen, nheads, batch) or
+            # (nheads, total_q) -> (total_q, nheads)
+            transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
+        else:
+            # (batch, seqlen, nheads) -> (seqlen, nheads, batch) or
+            # (total_q, nheads) -> (total_q, nheads)
+            transpose = [1, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 1]
+        mPdPsum, mLSE, mLSElog2, mdLSE, mdQaccum = [
+            layout_utils.select(mX, transpose) if mX is not None else None
+            for mX in (mPdPsum, mLSE, mLSElog2, mdLSE, mdQaccum)
+        ]
+
+        # (b, s, n, h) => (s, n, h, b) or
+        # (total, n, h) == (total, n, h)
+        rowmax_layout_transpose = [1, 2, 3, 0] if const_expr(mCuSeqlensQ is None) else [0, 1, 2]
+        if const_expr(mRowMax is not None):
+            mRowMax = layout_utils.select(mRowMax, rowmax_layout_transpose)
+        if const_expr(mScaleP is not None):
+            mScaleP = layout_utils.select(mScaleP, rowmax_layout_transpose)
+
+        # pack gqa
+        if const_expr(self.pack_gqa):
+            mO, mdO, mRowMax, mScaleP = [
+                pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2)
+                if mX is not None
+                else None
+                for mX in (mO, mdO, mRowMax, mScaleP)
+            ]
+            mPdPsum, mLSE, mLSElog2, mdLSE = [
+                pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=1)
+                if mX is not None
+                else None
+                for mX in (mPdPsum, mLSE, mLSElog2, mdLSE)
+            ]
+
+        # mO: (s, d, h, b) or (total, d, h)
+        if const_expr(mCuSeqlensQ is not None):
+            TileScheduler = SingleTileVarlenScheduler
+            num_head = mO.shape[2]
+            num_batch = mCuSeqlensQ.shape[0] - 1
+        else:
+            TileScheduler = SingleTileScheduler
+            num_head = mO.shape[2]
+            num_batch = mO.shape[3]
+
+        tile_sched_args = TileSchedulerArguments(
+            num_block=cute.ceil_div(mO.shape[0], self.tile_m),
+            num_head=num_head,
+            num_batch=num_batch,
+            num_splits=1,
+            seqlen_k=0,
+            headdim=0,
+            headdim_v=mO.shape[1],
+            total_q=cute.size(mO.shape[0])
+            if const_expr(mCuSeqlensQ is not None)
+            else cute.size(mO.shape[0]) * cute.size(mO.shape[3]),
+            tile_shape_mn=(self.tile_m, 1),
+            mCuSeqlensQ=mCuSeqlensQ,
+            mSeqUsedQ=mSeqUsedQ,
+            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
+        )
+
+        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
+        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
+
+        LOG2_E = math.log2(math.e)
+        softmax_scale_log2 = softmax_scale * LOG2_E
+
+        self.kernel(
+            mO,
+            mdO,
+            mPdPsum,
+            mLSE,
+            mLSElog2,
+            mdQaccum,
+            mCuSeqlensQ,
+            mSeqUsedQ,
+            mdLSE,
+            mRowMax,
+            mScaleP,
+            softmax_scale_log2,
+            self.gmem_tiled_copy_O,
+            self.gmem_tiled_copy_dQaccum,
+            tile_sched_params,
+            TileScheduler,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            stream=stream,
+            use_pdl=self.use_pdl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mO: cute.Tensor,
+        mdO: cute.Tensor,
+        mPdPsum: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        mLSElog2: Optional[cute.Tensor],
+        mdQaccum: Optional[cute.Tensor],
+        mCuSeqlensQ: Optional[cute.Tensor],
+        mSeqUsedQ: Optional[cute.Tensor],
+        mdLSE: Optional[cute.Tensor],
+        mRowMax: Optional[cute.Tensor],
+        mScaleP: Optional[cute.Tensor],
+        softmax_scale_log2: Float32,
+        gmem_tiled_copy_O: cute.TiledCopy,
+        gmem_tiled_copy_dQaccum: cute.TiledCopy,
+        tile_sched_params: ParamsBase,
+        TileScheduler: cutlass.Constexpr[Callable],
+    ):
+        # Thread index, block index
+        tidx, _, _ = cute.arch.thread_idx()
+
+        tile_scheduler = TileScheduler.create(tile_sched_params)
+        work_tile = tile_scheduler.initial_work_tile_info()
+        m_block, head_idx, batch_idx, _ = work_tile.tile_idx
+
+        # This kernel is launched with use_pdl=True, so the GPU may start executing it in
+        # "prologue" mode while the previous stream kernel is still running. We must wait
+        # before touching any upstream GMEM outputs (mO, mdO, mLSE); otherwise we risk
+        # reading a partially-written dout, which silently corrupts dpsum = sum(O * dO) and
+        # propagates to dQ/dK via dS = P * (dP - dpsum).
+        if const_expr(self.use_pdl):
+            cute.arch.griddepcontrol_wait()
+
+        if work_tile.is_valid_tile:
+            # ///////////////////////////////////////////////////////////////////////////////
+            # Get the appropriate tiles for this thread block.
+            # ///////////////////////////////////////////////////////////////////////////////
+            seqlen_static = mO.shape[0] if const_expr(not self.pack_gqa) else mO.shape[0][1]
+            seqlen = SeqlenInfo.create(
+                batch_idx, seqlen_static, mCuSeqlensQ, mSeqUsedQ, tile=self.tile_m
+            )
+            # (seqlen, dv)
+            mO_cur, mdO_cur = [
+                seqlen.offset_batch(mX, batch_idx, dim=3)[None, None, head_idx] for mX in (mO, mdO)
+            ]
+            mPdPsum_cur = seqlen.offset_batch(
+                mPdPsum, batch_idx, dim=2, padded=self.use_padded_offsets
+            )[None, head_idx]
+            headdim_v = mO_cur.shape[1]
+            seqlen_q = (
+                seqlen.seqlen
+                if const_expr(not self.pack_gqa)
+                else seqlen.seqlen * self.qhead_per_kvhead
+            )
+            seqlen_q_rounded = cute.round_up(seqlen_q, self.tile_m)
+            seqlen_limit = seqlen_q - m_block * self.tile_m
+
+            lse = None
+            if const_expr(mLSE is not None):
+                mLSE_cur = seqlen.offset_batch(mLSE, batch_idx, dim=2)[None, head_idx]
+                gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,))
+                lse = Float32.inf
+                if tidx < seqlen_limit:
+                    lse = gLSE[tidx]
+
+            blk_shape = (self.tile_m, self.head_dim_v_padded)
+            gO = cute.local_tile(mO_cur, blk_shape, (m_block, 0))
+            gdO = cute.local_tile(mdO_cur, blk_shape, (m_block, 0))
+            gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx)
+            # (CPY_Atom, CPY_M, CPY_K)
+            tOgO = gmem_thr_copy_O.partition_S(gO)
+            tOgdO = gmem_thr_copy_O.partition_S(gdO)
+            cO = cute.make_identity_tensor(blk_shape)
+            tOcO = gmem_thr_copy_O.partition_S(cO)
+            t0OcO = gmem_thr_copy_O.get_slice(0).partition_S(cO)
+            tOpO = None
+            if const_expr(self.check_hdim_v_oob):
+                tOpO = copy_utils.predicate_k(tOcO, limit=headdim_v)
+            # Each copy will use the same predicate
+            copy = partial(copy_utils.copy, pred=tOpO)
+
+            tOrO = cute.make_rmem_tensor_like(tOgO)
+            tOrdO = cute.make_rmem_tensor_like(tOgdO)
+            if const_expr(self.check_hdim_v_oob):
+                tOrO.fill(0.0)
+                tOrdO.fill(0.0)
+            assert tOgO.shape == tOgdO.shape
+            for m in cutlass.range(cute.size(tOrO.shape[1]), unroll_full=True):
+                # Instead of using tOcO, we using t0OcO and subtract the offset from the limit.
+                # This is bc the entries of t0OcO are known at compile time.
+                if t0OcO[0, m, 0][0] < seqlen_limit - tOcO[0][0]:
+                    copy(tOgO[None, m, None], tOrO[None, m, None])
+                    copy(tOgdO[None, m, None], tOrdO[None, m, None])
+            # O and dO loads are done; signal that the next kernel can start.
+            # Correctness is ensured by griddepcontrol_wait() in bwd_sm90 before it reads our outputs.
+            if const_expr(self.use_pdl):
+                cute.arch.griddepcontrol_launch_dependents()
+            # Sum across the "k" dimension
+            pdpsum = (tOrO.load().to(Float32) * tOrdO.load().to(Float32)).reduce(
+                cute.ReductionOp.ADD, init_val=0.0, reduction_profile=(0, None, 1)
+            )
+            threads_per_row = gmem_tiled_copy_O.layout_src_tv_tiled[0].shape[0]
+            assert cute.arch.WARP_SIZE % threads_per_row == 0
+            pdpsum = utils.warp_reduce(pdpsum, operator.add, width=threads_per_row)
+            PdP_sum = cute.make_rmem_tensor(cute.size(tOrO, mode=[1]), Float32)
+            PdP_sum.store(pdpsum)
+
+            # If dLSE is provided, compute D' = D - dLSE (see module docstring for derivation).
+            gdLSE = None
+            if const_expr(mdLSE is not None):
+                mdLSE_cur = seqlen.offset_batch(mdLSE, batch_idx, dim=2)[None, head_idx]
+                gdLSE = cute.local_tile(mdLSE_cur, (self.tile_m,), (m_block,))
+
+            # Write PdPsum from rmem -> gmem
+            gPdPsum = cute.local_tile(mPdPsum_cur, (self.tile_m,), (m_block,))
+            # Only the thread corresponding to column 0 writes out the PdPsum to gmem
+            if tOcO[0, 0, 0][1] == 0:
+                for m in cutlass.range(cute.size(PdP_sum), unroll_full=True):
+                    row = tOcO[0, m, 0][0]
+                    PdPsum_val = 0.0
+                    if row < seqlen_limit:
+                        PdPsum_val = PdP_sum[m]
+                        if const_expr(mdLSE is not None):
+                            PdPsum_val -= gdLSE[row]
+                    gPdPsum[row] = PdPsum_val
+
+            # Clear dQaccum
+            if const_expr(mdQaccum is not None):
+                mdQaccum_cur = seqlen.offset_batch(
+                    mdQaccum,
+                    batch_idx,
+                    dim=2,
+                    padded=self.use_padded_offsets,
+                    multiple=self.head_dim_padded,
+                )[None, head_idx]
+                blkdQaccum_shape = (self.tile_m * self.head_dim_padded,)
+                gdQaccum = cute.local_tile(mdQaccum_cur, blkdQaccum_shape, (m_block,))
+                gmem_thr_copy_dQaccum = gmem_tiled_copy_dQaccum.get_slice(tidx)
+                tdQgdQaccum = gmem_thr_copy_dQaccum.partition_S(gdQaccum)
+                zero = cute.make_rmem_tensor_like(tdQgdQaccum)
+                zero.fill(0.0)
+                cute.copy(gmem_tiled_copy_dQaccum, zero, tdQgdQaccum)
+
+            LOG2_E = math.log2(math.e)
+            lse_log2 = lse * LOG2_E if lse != -Float32.inf else 0.0
+            if const_expr(mLSElog2 is not None):
+                mLSElog2_cur = seqlen.offset_batch(
+                    mLSElog2, batch_idx, dim=2, padded=self.use_padded_offsets
+                )[None, head_idx]
+                gLSElog2 = cute.local_tile(mLSElog2_cur, (self.tile_m,), (m_block,))
+                LOG2_E = math.log2(math.e)
+                if tidx < seqlen_q_rounded - m_block * self.tile_m:
+                    gLSElog2[tidx] = lse_log2
+
+            if const_expr(mRowMax is not None):
+                assert mLSE is not None
+                # (s, n)
+                mRowMax_cur, mScaleP_cur = [
+                    seqlen.offset_batch(mX, batch_idx, dim=3)[None, None, head_idx]
+                    for mX in (mRowMax, mScaleP)
+                ]
+                # (tile_m, n)
+                gRowMax, gScaleP = [
+                    cute.local_tile(mX, (self.tile_m,), (m_block, None))
+                    for mX in (mRowMax_cur, mScaleP_cur)
+                ]
+
+                assert self.tile_m <= self.num_threads
+                if const_expr(self.tile_m == self.num_threads) or tidx < self.tile_m:
+                    for n in cutlass.range(gRowMax.shape[1], unroll=4):
+                        row_max = gRowMax[tidx, n]
+                        scale = 0.0
+                        if row_max != -Float32.inf and lse != -Float32.inf:
+                            scale = softmax_scale_log2 * row_max - lse_log2
+                            scale = cute.math.exp2(scale, fastmath=True)
+                        gScaleP[tidx, n] = scale
@@ -0,0 +1,55 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# SM120 (Blackwell GeForce / DGX Spark) backward pass.
+#
+# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has
+# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses
+# FlashAttentionBackwardSm80 and overrides the SMEM capacity check accordingly.
+
+import cutlass
+import cutlass.utils as utils_basic
+
+from flash_attn.cute.flash_bwd import FlashAttentionBackwardSm80
+
+
+class FlashAttentionBackwardSm120(FlashAttentionBackwardSm80):
+    @staticmethod
+    def can_implement(
+        dtype,
+        head_dim,
+        head_dim_v,
+        m_block_size,
+        n_block_size,
+        num_stages_Q,
+        num_stages_dO,
+        num_threads,
+        is_causal,
+        V_in_regs=False,
+    ) -> bool:
+        """Check if the kernel can be implemented on SM120.
+
+        Same logic as SM80 but uses SM120's shared memory capacity (99 KB).
+        """
+        if dtype not in [cutlass.Float16, cutlass.BFloat16]:
+            return False
+        if head_dim % 8 != 0:
+            return False
+        if head_dim_v % 8 != 0:
+            return False
+        if n_block_size % 16 != 0:
+            return False
+        if num_threads % 32 != 0:
+            return False
+        # Shared memory usage: Q tile + dO tile + K tile + V tile
+        smem_usage_Q = m_block_size * head_dim * num_stages_Q * 2
+        smem_usage_dO = m_block_size * head_dim_v * num_stages_dO * 2
+        smem_usage_K = n_block_size * head_dim * 2
+        smem_usage_V = n_block_size * head_dim_v * 2
+        smem_usage_QV = (
+            (smem_usage_Q + smem_usage_V) if not V_in_regs else max(smem_usage_Q, smem_usage_V)
+        )
+        smem_usage = smem_usage_QV + smem_usage_dO + smem_usage_K
+        # SM120 has 99 KB shared memory (vs 163 KB on SM80)
+        smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120")
+        if smem_usage > smem_capacity:
+            return False
+        return True
@@ -0,0 +1,698 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_combine_kernel.h
+# from Cutlass C++ to Cute-DSL.
+import math
+from typing import Type, Optional
+from functools import partial
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.nvgpu import cpasync
+from cutlass import Float32, Int32, Boolean, const_expr
+
+from flash_attn.cute import utils
+from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
+from flash_attn.cute.seqlen_info import SeqlenInfo
+from cutlass.cute import FastDivmodDivisor
+
+
+class FlashAttentionForwardCombine:
+    def __init__(
+        self,
+        dtype: Type[cutlass.Numeric],
+        dtype_partial: Type[cutlass.Numeric],
+        head_dim: int,
+        tile_m: int = 8,
+        k_block_size: int = 64,
+        log_max_splits: int = 4,
+        num_threads: int = 256,
+        stages: int = 4,
+    ):
+        """
+        Forward combine kernel for split attention computation.
+
+        :param dtype: output data type
+        :param dtype_partial: partial accumulation data type
+        :param head_dim: head dimension
+        :param tile_m: m block size
+        :param k_block_size: k block size
+        :param log_max_splits: log2 of maximum splits
+        :param num_threads: number of threads
+        :param varlen: whether using variable length sequences
+        :param stages: number of pipeline stages
+        """
+        self.dtype = dtype
+        self.dtype_partial = dtype_partial
+        self.head_dim = head_dim
+        self.tile_m = tile_m
+        self.k_block_size = k_block_size
+        self.max_splits = 1 << log_max_splits
+        self.num_threads = num_threads
+        self.is_even_k = head_dim % k_block_size == 0
+        self.stages = stages
+
+    @staticmethod
+    def can_implement(
+        dtype,
+        dtype_partial,
+        head_dim,
+        tile_m,
+        k_block_size,
+        log_max_splits,
+        num_threads,
+    ) -> bool:
+        """Check if the kernel can be implemented with the given parameters."""
+        if dtype not in [cutlass.Float16, cutlass.BFloat16, cutlass.Float32]:
+            return False
+        if dtype_partial not in [cutlass.Float16, cutlass.BFloat16, Float32]:
+            return False
+        if head_dim % 8 != 0:
+            return False
+        if num_threads % 32 != 0:
+            return False
+        if tile_m % 8 != 0:
+            return False
+        max_splits = 1 << log_max_splits
+        if max_splits > 256:
+            return False
+        if (tile_m * max_splits) % num_threads != 0:
+            return False
+        return True
+
+    def _setup_attributes(self):
+        # GMEM copy setup for O partial
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // self.dtype_partial.width
+        assert self.k_block_size % async_copy_elems == 0
+
+        k_block_gmem = (
+            128 if self.k_block_size % 128 == 0 else (64 if self.k_block_size % 64 == 0 else 32)
+        )
+        gmem_threads_per_row = k_block_gmem // async_copy_elems
+        assert self.num_threads % gmem_threads_per_row == 0
+
+        # Async copy atom for O partial load
+        atom_async_copy_partial = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            self.dtype_partial,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        tOpartial_layout = cute.make_ordered_layout(
+            (self.num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        vOpartial_layout = cute.make_layout((1, async_copy_elems))  # 4 vals per load
+        self.gmem_tiled_copy_O_partial = cute.make_tiled_copy_tv(
+            atom_async_copy_partial, tOpartial_layout, vOpartial_layout
+        )
+
+        # GMEM copy setup for final O (use universal copy for store)
+        atom_universal_copy = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            self.dtype,
+            num_bits_per_copy=async_copy_elems * self.dtype.width,
+        )
+        self.gmem_tiled_copy_O = cute.make_tiled_copy_tv(
+            atom_universal_copy,
+            tOpartial_layout,
+            vOpartial_layout,  # 4 vals per store
+        )
+
+        # LSE copy setup with async copy (alignment = 1)
+        lse_copy_bits = Float32.width  # 1 element per copy, width is in bits
+        m_block_smem = (
+            128
+            if self.tile_m % 128 == 0
+            else (
+                64
+                if self.tile_m % 64 == 0
+                else (32 if self.tile_m % 32 == 0 else (16 if self.tile_m % 16 == 0 else 8))
+            )
+        )
+        gmem_threads_per_row_lse = m_block_smem
+        assert self.num_threads % gmem_threads_per_row_lse == 0
+
+        # Async copy atom for LSE load
+        atom_async_copy_lse = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS),
+            Float32,
+            num_bits_per_copy=lse_copy_bits,
+        )
+        tLSE_layout = cute.make_ordered_layout(
+            (self.num_threads // gmem_threads_per_row_lse, gmem_threads_per_row_lse),
+            order=(1, 0),
+        )
+        vLSE_layout = cute.make_layout(1)
+        self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv(
+            atom_async_copy_lse, tLSE_layout, vLSE_layout
+        )
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Shared memory
+        # ///////////////////////////////////////////////////////////////////////////////
+
+        # Shared memory to register copy for LSE
+        self.smem_threads_per_col_lse = self.num_threads // m_block_smem
+        assert 32 % self.smem_threads_per_col_lse == 0  # Must divide warp size
+
+        s2r_layout_atom_lse = cute.make_ordered_layout(
+            (self.smem_threads_per_col_lse, self.num_threads // self.smem_threads_per_col_lse),
+            order=(0, 1),
+        )
+        self.s2r_tiled_copy_LSE = cute.make_tiled_copy_tv(
+            cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32),
+            s2r_layout_atom_lse,
+            cute.make_layout(1),
+        )
+
+        # LSE shared memory layout with swizzling to avoid bank conflicts
+        # This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts
+        if const_expr(m_block_smem == 8):
+            smem_lse_swizzle = cute.make_swizzle(5, 0, 5)
+        elif const_expr(m_block_smem == 16):
+            smem_lse_swizzle = cute.make_swizzle(4, 0, 4)
+        else:
+            smem_lse_swizzle = cute.make_swizzle(3, 2, 3)
+        smem_layout_atom_lse = cute.make_composed_layout(
+            smem_lse_swizzle, 0, cute.make_ordered_layout((8, m_block_smem), order=(1, 0))
+        )
+        self.smem_layout_lse = cute.tile_to_shape(
+            smem_layout_atom_lse, (self.max_splits, self.tile_m), (0, 1)
+        )
+
+        # O partial shared memory layout (simple layout for pipeline stages)
+        self.smem_layout_o = cute.make_ordered_layout(
+            (self.tile_m, self.k_block_size, self.stages), order=(1, 0, 2)
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        mO_partial: cute.Tensor,
+        mLSE_partial: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor] = None,
+        cu_seqlens: Optional[cute.Tensor] = None,
+        seqused: Optional[cute.Tensor] = None,
+        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,
+        varlen_batch_idx: Optional[cute.Tensor] = None,
+        semaphore_to_reset: Optional[cute.Tensor] = None,
+        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
+        stream: cuda.CUstream = None,
+    ):
+        # Type checking
+        if const_expr(not (mO_partial.element_type == self.dtype_partial)):
+            raise TypeError("O partial tensor must match dtype_partial")
+        if const_expr(not (mO.element_type == self.dtype)):
+            raise TypeError("O tensor must match dtype")
+        if const_expr(mLSE_partial.element_type not in [Float32]):
+            raise TypeError("LSE partial tensor must be Float32")
+        if const_expr(mLSE is not None and mLSE.element_type not in [Float32]):
+            raise TypeError("LSE tensor must be Float32")
+
+        # Shape validation - input tensors are in user format, need to be converted to kernel format
+        if const_expr(len(mO_partial.shape) not in [4, 5]):
+            raise ValueError(
+                "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)"
+            )
+        if const_expr(len(mLSE_partial.shape) not in [3, 4]):
+            raise ValueError(
+                "LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)"
+            )
+        if const_expr(len(mO.shape) not in [3, 4]):
+            raise ValueError(
+                "O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)"
+            )
+        if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]):
+            raise ValueError(
+                "LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)"
+            )
+
+        mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)]
+        # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b)
+        # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h)
+        O_partial_layout_transpose = (
+            [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2]
+        )
+        # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h)
+        mO_partial = cute.make_tensor(
+            mO_partial.iterator, cute.select(mO_partial.layout, mode=O_partial_layout_transpose)
+        )
+        O_layout_transpose = [1, 3, 2, 0] if const_expr(cu_seqlens is None) else [0, 2, 1]
+        mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose))
+        # (num_splits, b, seqlen, h) -> (seqlen, num_splits, h, b)
+        # or (num_splits, total_q, h) -> (total_q, num_splits, h)
+        LSE_partial_layout_transpose = [2, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 0, 2]
+        mLSE_partial = cute.make_tensor(
+            mLSE_partial.iterator,
+            cute.select(mLSE_partial.layout, mode=LSE_partial_layout_transpose),
+        )
+        # (b, seqlen, h) -> (seqlen, h, b) or (total_q, h) -> (total_q, h)
+        LSE_layout_transpose = [1, 2, 0] if const_expr(cu_seqlens is None) else [0, 1]
+        mLSE = (
+            cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose))
+            if mLSE is not None
+            else None
+        )
+
+        # Determine if we have variable length sequences
+        varlen = const_expr(cu_seqlens is not None or seqused is not None)
+
+        self._setup_attributes()
+
+        @cute.struct
+        class SharedStorage:
+            sLSE: cute.struct.Align[
+                cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128
+            ]
+            sMaxValidSplit: cute.struct.Align[cute.struct.MemRange[Int32, self.tile_m], 128]
+            sO: cute.struct.Align[
+                cute.struct.MemRange[self.dtype_partial, cute.cosize(self.smem_layout_o)], 128
+            ]
+
+        smem_size = SharedStorage.size_in_bytes()
+
+        # Grid dimensions: (ceil_div(seqlen, m_block), ceil_div(head_dim, k_block), num_head * batch)
+        seqlen = mO_partial.shape[0]
+        num_head = mO_partial.shape[3]
+        batch_size = (
+            mO_partial.shape[4]
+            if const_expr(cu_seqlens is None)
+            else Int32(cu_seqlens.shape[0] - 1)
+        )
+
+        # Create FastDivmodDivisor objects for efficient division
+        seqlen_divmod = FastDivmodDivisor(seqlen)
+        head_divmod = FastDivmodDivisor(num_head)
+
+        grid_dim = (
+            cute.ceil_div(seqlen * num_head, self.tile_m),
+            cute.ceil_div(self.head_dim, self.k_block_size),
+            batch_size,
+        )
+
+        self.kernel(
+            mO_partial,
+            mLSE_partial,
+            mO,
+            mLSE,
+            cu_seqlens,
+            seqused,
+            num_splits_dynamic_ptr,
+            varlen_batch_idx,
+            semaphore_to_reset,
+            SharedStorage,
+            self.smem_layout_lse,
+            self.smem_layout_o,
+            self.gmem_tiled_copy_O_partial,
+            self.gmem_tiled_copy_O,
+            self.gmem_tiled_copy_LSE,
+            self.s2r_tiled_copy_LSE,
+            seqlen_divmod,
+            head_divmod,
+            varlen,
+        ).launch(
+            grid=grid_dim,
+            block=[self.num_threads, 1, 1],
+            smem=smem_size,
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mO_partial: cute.Tensor,
+        mLSE_partial: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        cu_seqlens: Optional[cute.Tensor],
+        seqused: Optional[cute.Tensor],
+        num_splits_dynamic_ptr: Optional[cute.Tensor],
+        varlen_batch_idx: Optional[cute.Tensor],
+        semaphore_to_reset: Optional[cute.Tensor],
+        SharedStorage: cutlass.Constexpr,
+        smem_layout_lse: cute.Layout | cute.ComposedLayout,
+        smem_layout_o: cute.Layout,
+        gmem_tiled_copy_O_partial: cute.TiledCopy,
+        gmem_tiled_copy_O: cute.TiledCopy,
+        gmem_tiled_copy_LSE: cute.TiledCopy,
+        s2r_tiled_copy_LSE: cute.TiledCopy,
+        seqlen_divmod: FastDivmodDivisor,
+        head_divmod: FastDivmodDivisor,
+        varlen: cutlass.Constexpr[bool],
+    ):
+        # Thread and block indices
+        tidx, _, _ = cute.arch.thread_idx()
+        m_block, k_block, maybe_virtual_batch = cute.arch.block_idx()
+
+        # Map virtual batch index to real batch index (for persistent tile schedulers)
+        batch_idx = (
+            varlen_batch_idx[maybe_virtual_batch]
+            if const_expr(varlen_batch_idx is not None)
+            else maybe_virtual_batch
+        )
+
+        # ///////////////////////////////////////////////////////////////////////////////
+        # Get shared memory buffer
+        # ///////////////////////////////////////////////////////////////////////////////
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        sLSE = storage.sLSE.get_tensor(smem_layout_lse)
+        sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,))
+        sO = storage.sO.get_tensor(smem_layout_o)
+
+        # Handle semaphore reset — wait for dependent grids first
+        if const_expr(semaphore_to_reset is not None):
+            if (
+                tidx == 0
+                and m_block == cute.arch.grid_dim()[0] - 1
+                and k_block == cute.arch.grid_dim()[1] - 1
+                and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1
+            ):
+                cute.arch.griddepcontrol_wait()
+                semaphore_to_reset[0] = 0
+
+        # Get number of splits (use maybe_virtual_batch for per-batch-slot splits)
+        num_splits = (
+            num_splits_dynamic_ptr[maybe_virtual_batch]
+            if const_expr(num_splits_dynamic_ptr is not None)
+            else mLSE_partial.shape[1]
+        )
+        # Handle variable length sequences using SeqlenInfo
+        seqlen_info = SeqlenInfo.create(
+            batch_idx=batch_idx,
+            seqlen_static=mO_partial.shape[0],
+            cu_seqlens=cu_seqlens,
+            seqused=seqused,
+            # Don't need to pass in tile size since we won't use offset_padded
+        )
+        seqlen, offset = seqlen_info.seqlen, seqlen_info.offset
+
+        # Extract number of heads (head index will be determined dynamically)
+        num_head = mO_partial.shape[3]
+        max_idx = seqlen * num_head
+
+        # Early exit for single split if dynamic
+        if (const_expr(num_splits_dynamic_ptr is None) or num_splits > 1) and (
+            const_expr(not varlen) or m_block * self.tile_m < max_idx
+        ):
+            # Wait for dependent grids (e.g., the main attention kernel that produces O_partial/LSE_partial)
+            cute.arch.griddepcontrol_wait()
+
+            # ===============================
+            # Step 1: Load LSE_partial from gmem to shared memory
+            # ===============================
+
+            mLSE_partial_cur = seqlen_info.offset_batch(mLSE_partial, batch_idx, dim=3)
+            mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,))
+            gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx)
+            tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE)
+            # Create identity tensor for coordinate tracking
+            cLSE = cute.make_identity_tensor((self.max_splits, self.tile_m))
+            tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE)
+
+            # Load LSE partial values
+            for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True):
+                mi = tLSEcLSE[0, 0, m][1]  # Get m coordinate
+                idx = m_block * self.tile_m + mi
+                if idx < max_idx:
+                    # Calculate actual sequence position and head using FastDivmodDivisor
+                    if const_expr(not varlen):
+                        head_idx, m_idx = divmod(idx, seqlen_divmod)
+                    else:
+                        head_idx = idx // seqlen
+                        m_idx = idx - head_idx * seqlen
+                    mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx]
+                    for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True):
+                        si = tLSEcLSE[0, s, 0][0]  # Get split coordinate
+                        if si < num_splits:
+                            cute.copy(
+                                gmem_thr_copy_LSE,
+                                mLSE_partial_cur_copy[None, si],
+                                tLSEsLSE[None, s, m],
+                            )
+                        else:
+                            tLSEsLSE[None, s, m].fill(-Float32.inf)
+                # Don't need to zero out the rest of the LSEs, as we will not write the output to gmem
+            cute.arch.cp_async_commit_group()
+
+            # ===============================
+            # Step 2: Load O_partial for pipeline stages
+            # ===============================
+
+            gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_slice(tidx)
+            cO = cute.make_identity_tensor((self.tile_m, self.k_block_size))
+            tOcO = gmem_thr_copy_O_partial.partition_D(cO)
+            tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO)
+            mO_partial_cur = seqlen_info.offset_batch(mO_partial, batch_idx, dim=4)
+
+            # Precompute these values to avoid recomputing them in the loop
+            num_rows = const_expr(cute.size(tOcO, mode=[1]))
+            tOmidx = cute.make_rmem_tensor(num_rows, cutlass.Int32)
+            tOhidx = cute.make_rmem_tensor(num_rows, cutlass.Int32)
+            tOrOptr = cute.make_rmem_tensor(num_rows, cutlass.Int64)
+            for m in cutlass.range(num_rows, unroll_full=True):
+                mi = tOcO[0, m, 0][0]  # m coordinate
+                idx = m_block * self.tile_m + mi
+                if const_expr(not varlen):
+                    tOhidx[m], tOmidx[m] = divmod(idx, seqlen_divmod)
+                else:
+                    tOhidx[m] = idx // seqlen
+                    tOmidx[m] = idx - tOhidx[m] * seqlen
+                tOrOptr[m] = utils.elem_pointer(
+                    mO_partial_cur, (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m])
+                ).toint()
+                if idx >= max_idx:
+                    tOhidx[m] = -1
+
+            tOpO = None
+            if const_expr(not self.is_even_k):
+                tOpO = cute.make_rmem_tensor(cute.size(tOcO, mode=[2]), Boolean)
+                for k in cutlass.range(cute.size(tOpO), unroll_full=True):
+                    tOpO[k] = tOcO[0, 0, k][1] < mO_partial.shape[1] - k_block * self.k_block_size
+                # if cute.arch.thread_idx()[0] == 0 and k_block == 1: cute.print_tensor(tOpO)
+
+            load_O_partial = partial(
+                self.load_O_partial,
+                gmem_tiled_copy_O_partial,
+                tOrOptr,
+                tOsO_partial,
+                tOhidx,
+                tOpO,
+                tOcO,
+                mO_partial_cur.layout,
+            )
+
+            # Load first few stages of O_partial
+            for stage in cutlass.range(self.stages - 1, unroll_full=True):
+                if stage < num_splits:
+                    load_O_partial(stage, stage)
+                cute.arch.cp_async_commit_group()
+
+            # ===============================
+            # Step 3: Load and transpose LSE from smem to registers
+            # ===============================
+
+            # Wait for LSE and initial O partial stages to complete
+            cute.arch.cp_async_wait_group(self.stages - 1)
+            cute.arch.sync_threads()
+            # if cute.arch.thread_idx()[0] == 0:
+            #     # cute.print_tensor(sLSE)
+            #     for i in range(64):
+            #         cute.printf("sLSE[%d, 0] = %f", i, sLSE[i, 0])
+            # cute.arch.sync_threads()
+
+            s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_slice(tidx)
+            ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE)
+            ts2rrLSE = cute.make_rmem_tensor_like(ts2rsLSE)
+            cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE)
+
+            # ===============================
+            # Step 4: Compute final LSE along split dimension
+            # ===============================
+
+            lse_sum = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Float32)
+            ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE)
+            # We compute the max valid split for each row to short-circuit the computation later
+            max_valid_split = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Int32)
+            assert cute.size(ts2rrLSE, mode=[0]) == 1
+            # Compute max, scales, and final LSE for each row
+            for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                # Find max LSE value across splits
+                threads_per_col = const_expr(self.smem_threads_per_col_lse)
+                lse_max = cute.arch.warp_reduction_max(
+                    ts2rrLSE[None, None, m]
+                    .load()
+                    .reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0),
+                    threads_in_group=threads_per_col,
+                )
+                # if cute.arch.thread_idx()[0] == 0: cute.printf(lse_max)
+                # Find max valid split index
+                max_valid_idx = -1
+                for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True):
+                    if ts2rrLSE[0, s, m] != -Float32.inf:
+                        max_valid_idx = ts2rcLSE[0, s, 0][0]  # Get split coordinate
+                # if cute.arch.thread_idx()[0] < 32: cute.printf(max_valid_idx)
+                max_valid_split[m] = cute.arch.warp_reduction_max(
+                    max_valid_idx, threads_in_group=threads_per_col
+                )
+                # Compute exp scales and sum
+                lse_max_cur = (
+                    0.0 if lse_max == -Float32.inf else lse_max
+                )  # In case all local LSEs are -inf
+                LOG2_E = math.log2(math.e)
+                lse_sum_cur = 0.0
+                for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True):
+                    scale = cute.math.exp2(
+                        ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E), fastmath=True
+                    )
+                    lse_sum_cur += scale
+                    ts2rrLSE[0, s, m] = scale  # Store scale for later use
+                lse_sum_cur = cute.arch.warp_reduction_sum(
+                    lse_sum_cur, threads_in_group=threads_per_col
+                )
+                lse_sum[m] = cute.math.log(lse_sum_cur, fastmath=True) + lse_max
+                # Normalize scales
+                inv_sum = (
+                    0.0 if (lse_sum_cur == 0.0 or lse_sum_cur != lse_sum_cur) else 1.0 / lse_sum_cur
+                )
+                ts2rrLSE[None, None, m].store(ts2rrLSE[None, None, m].load() * inv_sum)
+            # Store the scales exp(lse - lse_logsum) back to smem
+            cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE)
+
+            # Store max valid split to smem
+            for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
+                    mi = ts2rcLSE[0, 0, m][1]
+                    if mi < self.tile_m:
+                        sMaxValidSplit[mi] = max_valid_split[m]
+
+            # ===============================
+            # Step 5: Store final LSE to gmem
+            # ===============================
+
+            if const_expr(mLSE is not None):
+                if const_expr(cu_seqlens is None):
+                    mLSE_cur = mLSE[None, None, batch_idx]
+                else:
+                    mLSE_cur = cute.domain_offset((offset, 0), mLSE)
+                if k_block == 0:  # Only first k_block writes LSE when mLSE is provided
+                    for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True):
+                        if ts2rcLSE[0, 0, m][0] == 0:  # Only thread responsible for s=0 writes
+                            mi = ts2rcLSE[0, 0, m][1]
+                            idx = m_block * self.tile_m + mi
+                            if idx < max_idx:
+                                if const_expr(not varlen):
+                                    head_idx, m_idx = divmod(idx, seqlen_divmod)
+                                else:
+                                    head_idx = idx // seqlen
+                                    m_idx = idx - head_idx * seqlen
+                                mLSE_cur[m_idx, head_idx] = lse_sum[m]
+
+            # ===============================
+            # Step 6: Read O_partial and accumulate final O
+            # ===============================
+
+            cute.arch.sync_threads()
+
+            # Get max valid split for this thread
+            thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]]
+            for m in cutlass.range(1, cute.size(tOcO, mode=[1]), unroll_full=True):
+                thr_max_valid_split = max(thr_max_valid_split, sMaxValidSplit[tOcO[0, m, 0][0]])
+
+            tOrO_partial = cute.make_rmem_tensor_like(tOsO_partial[None, None, None, 0])
+            tOrO = cute.make_rmem_tensor_like(tOrO_partial, Float32)
+            tOrO.fill(0.0)
+
+            stage_load = self.stages - 1
+            stage_compute = 0
+
+            # Main accumulation loop
+            for s in cutlass.range(thr_max_valid_split + 1, unroll=4):
+                # Get scales for this split
+                scale = cute.make_rmem_tensor(num_rows, Float32)
+                for m in cutlass.range(num_rows, unroll_full=True):
+                    scale[m] = sLSE[s, tOcO[0, m, 0][0]]  # Get scale from smem
+
+                # Load next stage if needed
+                split_to_load = s + self.stages - 1
+                if split_to_load <= thr_max_valid_split:
+                    load_O_partial(split_to_load, stage_load)
+                cute.arch.cp_async_commit_group()
+                stage_load = 0 if stage_load == self.stages - 1 else stage_load + 1
+
+                # Wait for the current stage to be ready
+                cute.arch.cp_async_wait_group(self.stages - 1)
+                # We don't need __syncthreads() because each thread is just reading its own data from smem
+                # Copy from smem to registers
+                cute.autovec_copy(tOsO_partial[None, None, None, stage_compute], tOrO_partial)
+                stage_compute = 0 if stage_compute == self.stages - 1 else stage_compute + 1
+
+                # Accumulate scaled partial results
+                for m in cutlass.range(num_rows, unroll_full=True):
+                    if tOhidx[m] >= 0 and scale[m] > 0.0:
+                        tOrO[None, m, None].store(
+                            tOrO[None, m, None].load()
+                            + scale[m] * tOrO_partial[None, m, None].load().to(Float32)
+                        )
+
+            # ===============================
+            # Step 7: Write final O to gmem
+            # ===============================
+
+            rO = cute.make_rmem_tensor_like(tOrO, self.dtype)
+            rO.store(tOrO.load().to(self.dtype))
+            mO_cur = seqlen_info.offset_batch(mO, batch_idx, dim=3)
+            if const_expr(cu_seqlens is None):
+                mO_cur = mO[None, None, None, batch_idx]
+            else:
+                mO_cur = cute.domain_offset((offset, 0, 0), mO)
+            mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur)
+            elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1]))
+            # mO_cur_copy = cute.tiled_divide(mO_cur, (1, elems_per_store,))
+            gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx)
+            # Write final results
+            for m in cutlass.range(num_rows, unroll_full=True):
+                if tOhidx[m] >= 0:
+                    mO_cur_copy = cute.tiled_divide(
+                        mO_cur[tOmidx[m], None, tOhidx[m]], (elems_per_store,)
+                    )
+                    for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True):
+                        k_idx = tOcO[0, 0, k][1] // elems_per_store
+                        if const_expr(self.is_even_k) or tOpO[k]:
+                            cute.copy(gmem_thr_copy_O, rO[None, m, k], mO_cur_copy[None, k_idx])
+
+    @cute.jit
+    def load_O_partial(
+        self,
+        gmem_tiled_copy_O_partial: cute.TiledCopy,
+        tOrOptr: cute.Tensor,
+        tOsO_partial: cute.Tensor,
+        tOhidx: cute.Tensor,
+        tOpO: Optional[cute.Tensor],
+        tOcO: cute.Tensor,
+        mO_cur_partial_layout: cute.Layout,
+        split: Int32,
+        stage: Int32,
+    ) -> None:
+        elems_per_load = const_expr(cute.size(gmem_tiled_copy_O_partial.layout_tv_tiled[1]))
+        tOsO_partial_cur = tOsO_partial[None, None, None, stage]
+        for m in cutlass.range(cute.size(tOcO, [1]), unroll_full=True):
+            if tOhidx[m] >= 0:
+                o_gmem_ptr = cute.make_ptr(
+                    tOsO_partial.element_type, tOrOptr[m], cute.AddressSpace.gmem, assumed_align=16
+                )
+                mO_partial_cur = cute.make_tensor(
+                    o_gmem_ptr, cute.slice_(mO_cur_partial_layout, (0, None, None, 0))
+                )
+                mO_partial_cur_copy = cute.tiled_divide(mO_partial_cur, (elems_per_load,))
+                for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True):
+                    k_idx = tOcO[0, 0, k][1] // elems_per_load
+                    if const_expr(tOpO is None) or tOpO[k]:
+                        cute.copy(
+                            gmem_tiled_copy_O_partial,
+                            mO_partial_cur_copy[None, k_idx, split],
+                            tOsO_partial_cur[None, m, k],
+                        )
@@ -0,0 +1,59 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+# SM120 (Blackwell GeForce / DGX Spark) forward pass.
+#
+# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has
+# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses
+# FlashAttentionForwardSm80 and overrides the SMEM capacity check accordingly.
+
+import cutlass
+import cutlass.utils as utils_basic
+
+from flash_attn.cute.flash_fwd import FlashAttentionForwardSm80
+
+
+class FlashAttentionForwardSm120(FlashAttentionForwardSm80):
+    # Keep arch = 80 to use CpAsync code paths (no TMA for output).
+    # The compilation target is determined by the GPU at compile time, not this field.
+    arch = 80
+
+    @staticmethod
+    def can_implement(
+        dtype,
+        head_dim,
+        head_dim_v,
+        tile_m,
+        tile_n,
+        num_stages,
+        num_threads,
+        is_causal,
+        Q_in_regs=False,
+    ) -> bool:
+        """Check if the kernel can be implemented on SM120.
+
+        Same logic as SM80 but uses SM120's shared memory capacity (99 KB).
+        """
+        if dtype not in [cutlass.Float16, cutlass.BFloat16]:
+            return False
+        if head_dim % 8 != 0:
+            return False
+        if head_dim_v % 8 != 0:
+            return False
+        if tile_n % 16 != 0:
+            return False
+        if num_threads % 32 != 0:
+            return False
+        # Shared memory usage: Q tile + (K tile + V tile)
+        smem_usage_Q = tile_m * head_dim * 2
+        smem_usage_K = tile_n * head_dim * num_stages * 2
+        smem_usage_V = tile_n * head_dim_v * num_stages * 2
+        smem_usage_QV = (
+            (smem_usage_Q + smem_usage_V) if not Q_in_regs else max(smem_usage_Q, smem_usage_V)
+        )
+        smem_usage = smem_usage_QV + smem_usage_K
+        # SM120 has 99 KB shared memory (vs 163 KB on SM80)
+        smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120")
+        if smem_usage > smem_capacity:
+            return False
+        if (tile_m * 2) % num_threads != 0:
+            return False
+        return True
@@ -0,0 +1,296 @@
+# Copyright (c) 2025, Tri Dao.
+# Ported Cutlass code from C++ to Python:
+# https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/mma_sm100_desc.hpp
+# https://github.com/NVIDIA/cutlass/blob/main/include/cute/atom/mma_traits_sm100.hpp
+
+from enum import IntEnum
+
+import cutlass
+import cutlass.cute as cute
+
+# ---------------------------------------------------------------------------
+# Enumerations that match the HW encodings (values MUST stay identical)
+# ---------------------------------------------------------------------------
+
+
+class Major(IntEnum):  # matrix “layout” in the ISA docs
+    K = 0
+    MN = 1
+
+
+class ScaleIn(IntEnum):  # negate flags
+    One = 0
+    Neg = 1
+
+
+class Saturate(IntEnum):
+    False_ = 0
+    True_ = 1
+
+
+class CFormat(IntEnum):  # 2-bit field (bits 4-5)
+    F16 = 0
+    F32 = 1
+    S32 = 2
+
+
+class F16F32Format(IntEnum):  # 3-bit field (A/B element type)
+    F16 = 0
+    BF16 = 1
+    TF32 = 2
+
+
+class S8Format(IntEnum):
+    UINT8 = 0
+    INT8 = 1
+
+
+class MXF8F6F4Format(IntEnum):
+    E4M3 = 0
+    E5M2 = 1
+    E2M3 = 3
+    E3M2 = 4
+    E2M1 = 5
+
+
+class MaxShift(IntEnum):
+    NoShift = 0
+    MaxShift8 = 1
+    MaxShift16 = 2
+    MaxShift32 = 3
+
+
+# ---------------------------------------------------------------------------
+# CUTLASS-type → encoding helpers
+# ---------------------------------------------------------------------------
+
+
+def to_UMMA_format(cutlass_type) -> int:
+    """
+    Map a CUTLASS scalar class to the 3-bit encoding for Matrix A/B.
+    """
+    if cutlass_type is cutlass.Int8:
+        return S8Format.INT8
+    # Unsigned 8-bit (if available in your CUTLASS build)
+    if cutlass_type is cutlass.Uint8:
+        return S8Format.UINT8
+    # FP-16 / BF-16
+    if cutlass_type is cutlass.Float16:
+        return F16F32Format.F16
+    if cutlass_type is cutlass.BFloat16:
+        return F16F32Format.BF16
+    # TensorFloat-32 (8-bit exponent, 10-bit mantissa packed in 19 bits)
+    if cutlass_type is cutlass.TFloat32:
+        return F16F32Format.TF32
+    # Float-8 / Float-6 / Float-4 – add whenever CUTLASS exposes them
+    if cutlass_type is cutlass.Float8E4M3FN:
+        return MXF8F6F4Format.E4M3
+    if cutlass_type is cutlass.Float8E5M2:
+        return MXF8F6F4Format.E5M2
+    raise TypeError(f"Unsupported CUTLASS scalar type for A/B: {cutlass_type!r}")
+
+
+def to_C_format(cutlass_type) -> int:
+    """
+    Map a CUTLASS scalar class to the 2-bit accumulator encoding.
+    """
+    if cutlass_type is cutlass.Float16:
+        return CFormat.F16
+    if cutlass_type is cutlass.Float32:
+        return CFormat.F32
+    if cutlass_type is cutlass.Int32:
+        return CFormat.S32
+    raise TypeError(f"Unsupported CUTLASS scalar type for accumulator: {cutlass_type!r}")
+
+
+# ---------------------------------------------------------------------------
+# The constructor – accepts only CUTLASS scalar classes
+# ---------------------------------------------------------------------------
+
+
+def make_instr_desc(
+    a_type,  # CUTLASS scalar class, e.g. cutlass.Int8
+    b_type,
+    c_type,
+    M: int,  # 64, 128 or 256
+    N: int,  # 8 … 256 (multiple of 8)
+    a_major: Major,
+    b_major: Major,
+    a_neg: ScaleIn = ScaleIn.One,
+    b_neg: ScaleIn = ScaleIn.One,
+    c_sat: Saturate = Saturate.False_,
+    is_sparse: bool = False,
+    max_shift: MaxShift = MaxShift.NoShift,
+) -> int:
+    """
+    Build the 32-bit instruction descriptor for Blackwell MMA.
+    All matrix/accumulator **types must be CUTLASS scalar classes** –
+    passing integers is forbidden.
+    """
+    # --- encode element formats -------------------------------------------------
+    a_fmt = int(to_UMMA_format(a_type))
+    b_fmt = int(to_UMMA_format(b_type))
+    c_fmt = int(to_C_format(c_type))
+
+    # --- range checks on M/N -----------------------------------------------------
+    if M not in (64, 128, 256):
+        raise ValueError("M must be 64, 128 or 256")
+    if N < 8 or N > 256 or (N & 7):
+        raise ValueError("N must be a multiple of 8 in the range 8…256")
+
+    m_dim = M >> 4  # 5-bit field
+    n_dim = N >> 3  # 6-bit field
+
+    # fmt: off
+    # --- pack the bit-fields -----------------------------------------------------
+    desc = 0
+    desc |= (0                 & 0x3) << 0        # sparse_id2 (always 0 here)
+    desc |= (int(is_sparse)    & 0x1) << 2        # sparse_flag
+    desc |= (int(c_sat)        & 0x1) << 3        # saturate
+    desc |= (c_fmt             & 0x3) << 4        # c_format
+    desc |= (a_fmt             & 0x7) << 7        # a_format
+    desc |= (b_fmt             & 0x7) << 10       # b_format
+    desc |= (int(a_neg)        & 0x1) << 13       # a_negate
+    desc |= (int(b_neg)        & 0x1) << 14       # b_negate
+    desc |= (int(a_major)      & 0x1) << 15       # a_major
+    desc |= (int(b_major)      & 0x1) << 16       # b_major
+    desc |= (n_dim             & 0x3F) << 17      # n_dim (6 bits)
+    desc |= (m_dim             & 0x1F) << 24      # m_dim (5 bits)
+    desc |= (int(max_shift)    & 0x3) << 30       # max_shift (2 bits)
+    # fmt: on
+
+    return desc & 0xFFFF_FFFF  # ensure 32-bit result
+
+
+def mma_op_to_idesc(op: cute.nvgpu.tcgen05.mma.MmaOp):
+    return make_instr_desc(
+        op.a_dtype,
+        op.b_dtype,
+        op.acc_dtype,
+        op.shape_mnk[0],
+        op.shape_mnk[1],
+        Major.K if op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K else Major.MN,
+        Major.K if op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K else Major.MN,
+    )
+
+
+class LayoutType(IntEnum):  # occupies the top-3 bits [61:64)
+    SWIZZLE_NONE = 0  # (a.k.a. “INTERLEAVE” in older docs)
+    SWIZZLE_128B_BASE32B = 1
+    SWIZZLE_128B = 2
+    SWIZZLE_64B = 4
+    SWIZZLE_32B = 6
+    # values 3,5,7 are reserved / illegal for UMMA
+
+
+# ---------------------------------------------------------------------------
+#  Helpers – figure out the SWIZZLE_* family from the tensor layout
+# ---------------------------------------------------------------------------
+
+
+def _layout_type(swizzle: cute.Swizzle) -> LayoutType:
+    B, M, S = swizzle.num_bits, swizzle.num_base, swizzle.num_shift
+
+    if M == 4:  # Swizzle<*,4,3>
+        if S != 3:
+            raise ValueError("Unexpected swizzle shift – want S==3 for M==4")
+        return {
+            0: LayoutType.SWIZZLE_NONE,
+            1: LayoutType.SWIZZLE_32B,
+            2: LayoutType.SWIZZLE_64B,
+            3: LayoutType.SWIZZLE_128B,
+        }[B]  # KeyError ⇒ invalid B→ raise
+    if M == 5:  # Swizzle<2,5,2> (the only legal triple for M==5)
+        if (B, S) != (2, 2):
+            raise ValueError("Only Swizzle<2,5,2> supported for 128B_BASE32B")
+        return LayoutType.SWIZZLE_128B_BASE32B
+
+    # Any other (M,B,S) triple is not a UMMA-legal shared-memory layout
+    raise ValueError("Unsupported swizzle triple for UMMA smem descriptor")
+
+
+def make_smem_desc_base(layout: cute.Layout, swizzle: cute.Swizzle, major: Major) -> int:
+    """
+    Convert a 2-D *shared-memory* Cute layout into the Blackwell 64-bit
+    smem-descriptor, without the smem start address.
+    layout must correspond to layout of an uint128 tensor.
+    """
+    # ------------------------------------------------------------------ meta
+    layout_type = _layout_type(swizzle)  # resolve SWIZZLE_* family
+
+    VERSION = 1  # bits 46–47
+    LBO_MODE = 0  # bit  52
+    BASE_OFFSET = 0  # bits 49–51   (CUTLASS always 0)
+
+    # ---------------------------------------------------------- strides  (units: uint128_t = 16 B)
+    swizzle_atom_mn_size = {
+        LayoutType.SWIZZLE_NONE: 1,
+        LayoutType.SWIZZLE_32B: 2,
+        LayoutType.SWIZZLE_64B: 4,
+        LayoutType.SWIZZLE_128B: 8,
+        LayoutType.SWIZZLE_128B_BASE32B: 8,
+    }[layout_type]
+
+    if major is Major.MN:
+        swizzle_atom_k_size = 4 if layout_type is LayoutType.SWIZZLE_128B_BASE32B else 8
+        canonical_layout = cute.logical_divide(layout, (swizzle_atom_mn_size, swizzle_atom_k_size))
+        if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))):
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected profile failure.")
+        stride_00 = canonical_layout.stride[0][0]
+        if layout_type is not LayoutType.SWIZZLE_NONE and stride_00 != 1:
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.")
+        stride_10 = canonical_layout.stride[1][0]
+        if stride_10 != swizzle_atom_mn_size:
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.")
+        stride_01, stride_11 = canonical_layout.stride[0][1], canonical_layout.stride[1][1]
+        if layout_type is LayoutType.SWIZZLE_NONE:
+            stride_byte_offset, leading_byte_offset = stride_01, stride_11
+        else:
+            stride_byte_offset, leading_byte_offset = stride_11, stride_01
+    else:
+        if layout_type == LayoutType.SWIZZLE_128B_BASE32B:
+            raise ValueError("SWIZZLE_128B_BASE32B is invalid for Major-K")
+        if not cute.size(layout.shape[0]) % 8 == 0:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected MN-size multiple of 8.")
+        canonical_layout = cute.logical_divide(layout, (8, 2))
+        if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))):
+            raise ValueError("Not a canonical UMMA_K Layout: Expected profile failure.")
+        stride_00 = canonical_layout.stride[0][0]
+        if stride_00 != swizzle_atom_mn_size:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.")
+        stride_10 = canonical_layout.stride[1][0]
+        if layout_type is not LayoutType.SWIZZLE_NONE and stride_10 != 1:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.")
+        stride_01 = canonical_layout.stride[0][1]
+        stride_byte_offset, leading_byte_offset = stride_01, stride_10
+
+    # ------------------------------------------------------------------ pack
+    desc = 0
+    # leading_byte_offset_  [16:30)
+    desc |= (leading_byte_offset & 0x3FFF) << 16
+    # stride_byte_offset_   [32:46)
+    desc |= (stride_byte_offset & 0x3FFF) << 32
+    # version_             [46:48)
+    desc |= (VERSION & 0x3) << 46
+    # base_offset_         [49:52)
+    desc |= (BASE_OFFSET & 0x7) << 49
+    # lbo_mode_            [52:53)
+    desc |= (LBO_MODE & 0x1) << 52
+    # layout_type_         [61:64)
+    desc |= (int(layout_type) & 0x7) << 61
+
+    return desc & 0xFFFF_FFFF_FFFF_FFFF  # force 64-bit width
+
+
+def make_smem_desc_start_addr(start_addr: cute.Pointer) -> cutlass.Int32:
+    # 14 bits, remove 4 LSB (bits 0-13 in desc)
+    return (start_addr.toint() & 0x3FFFF) >> 4
+
+
+def smem_desc_base_from_tensor(sA: cute.Tensor, major: Major) -> int:
+    sA_swizzle = sA.iterator.type.swizzle_type
+    return make_smem_desc_base(
+        cute.recast_layout(128, sA.element_type.width, sA.layout[0]),
+        sA_swizzle,
+        major,
+    )
@@ -0,0 +1,68 @@
+# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+
+import enum
+
+
+class NamedBarrierFwd(enum.IntEnum):
+    Epilogue = enum.auto()  # starts from 1 as barrier 0 is reserved for sync_threads()
+    WarpSchedulerWG1 = enum.auto()
+    WarpSchedulerWG2 = enum.auto()
+    WarpSchedulerWG3 = enum.auto()
+    PFull = enum.auto()
+    PEmpty = enum.auto()
+
+
+class NamedBarrierFwdSm100(enum.IntEnum):
+    Epilogue = enum.auto()  # starts from 1 as barrier 0 is reserved for sync_threads()
+    TmemPtr = enum.auto()
+    SoftmaxStatsW0 = enum.auto()
+    SoftmaxStatsW1 = enum.auto()
+    SoftmaxStatsW2 = enum.auto()
+    SoftmaxStatsW3 = enum.auto()
+    SoftmaxStatsW4 = enum.auto()
+    SoftmaxStatsW5 = enum.auto()
+    SoftmaxStatsW6 = enum.auto()
+    SoftmaxStatsW7 = enum.auto()
+    # Reserve one independent SMEM-P visibility barrier per compute slot. A
+    # single shared barrier can mix arrivals from the two softmax warpgroups.
+    SoftmaxSmemP0 = enum.auto()
+    SoftmaxSmemP1 = enum.auto()
+    CorrectionScale = enum.auto()
+
+
+class NamedBarrierBwd(enum.IntEnum):
+    Epilogue = enum.auto()
+    WarpSchedulerWG1 = enum.auto()
+    WarpSchedulerWG2 = enum.auto()
+    WarpSchedulerWG3 = enum.auto()
+    PdS = enum.auto()
+    dQFullWG0 = enum.auto()
+    dQFullWG1 = enum.auto()
+    dQFullWG2 = enum.auto()
+    dQEmptyWG0 = enum.auto()
+    dQEmptyWG1 = enum.auto()
+    dQEmptyWG2 = enum.auto()
+
+
+class NamedBarrierBwdSm100(enum.IntEnum):
+    EpilogueWG1 = enum.auto()
+    EpilogueWG2 = enum.auto()
+    Compute = enum.auto()
+    dQaccReduce = enum.auto()
+    TmemPtr = enum.auto()
+
+
+class NamedBarrierFwdSm100_MLA2CTA(enum.IntEnum):
+    Epilogue = enum.auto()
+    TmemPtr = enum.auto()
+    Cpasync = enum.auto()
+    Softmax = enum.auto()
+    SoftmaxStatsFull = enum.auto()
+    SoftmaxStatsEmpty = enum.auto()
+
+
+class NamedBarrierBwdSm100_MLA2CTA(enum.IntEnum):
+    Epilogue = enum.auto()
+    TmemPtr = enum.auto()
+    Cpasync = enum.auto()
+    Softmax = enum.auto()
@@ -0,0 +1,263 @@
+# Copyright (c) 2025, Tri Dao.
+
+from dataclasses import dataclass
+from typing import Union, Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.nvgpu import cpasync
+
+
+from quack import layout_utils
+import flash_attn.cute.utils as utils
+
+
+def pack_gqa_layout(T, qhead_per_kvhead, nheads_kv, head_idx):
+    """Reshape a tensor to fold qhead_per_kvhead into the seqlen dimension (mode 0).
+
+    The head dimension is at mode ``head_idx``.  Modes before it (1..head_idx-1)
+    are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept
+    as-is (e.g. batch).
+
+    For Q/O tensors (head_idx=2):
+        (seqlen_q, headdim, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...)
+    For LSE tensors (head_idx=1):
+        (seqlen_q, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...)
+    """
+    head_stride = T.stride[head_idx]
+    shape_packed = (
+        (qhead_per_kvhead, T.shape[0]),
+        *[T.shape[i] for i in range(1, head_idx)],
+        nheads_kv,
+        *[T.shape[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    stride_packed = (
+        (head_stride, T.stride[0]),
+        *[T.stride[i] for i in range(1, head_idx)],
+        head_stride * qhead_per_kvhead,
+        *[T.stride[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    return cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed))
+
+
+def make_packgqa_tiled_tma_atom(
+    op: cute.atom.CopyOp,
+    gmem_tensor: cute.Tensor,
+    smem_layout: Union[cute.Layout, cute.ComposedLayout],
+    cta_tiler: Tuple[int, int],
+    qhead_per_kvhead: int,
+    head_idx: int,
+):
+    # This packing and unpacking of the layout is so that we keep the same TMA dimension as usual.
+    # e.g. for (seqlen, d, nheads, b) layout, we still have 4D TMA after packing to
+    # ((nheads, seqlen), d, b).
+    # If we instead pack directly to ((qhead_per_kvhead, seqlen), d, nheads_kv, b) we'd have 5D TMA.
+    # Pack headdim and seqlen dim into 1: (seqlen, d, nheads, b) -> ((nheads, seqlen), d, b)
+    gmem_tensor = layout_utils.select(
+        gmem_tensor, [head_idx, *range(head_idx), *range(head_idx + 1, cute.rank(gmem_tensor))]
+    )
+    gmem_tensor = cute.group_modes(gmem_tensor, 0, 2)
+    assert cta_tiler[0] % qhead_per_kvhead == 0, (
+        "CTA tile size in the seqlen dimension must be divisible by qhead_per_kvhead"
+    )
+    tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(
+        op,
+        gmem_tensor,
+        smem_layout,
+        ((qhead_per_kvhead, cta_tiler[0] // qhead_per_kvhead), cta_tiler[1]),  # No mcast
+    )
+    # Unpack from ((nheads, seqlen), d, b) -> ((qhead_per_kvhead, seqlen), d, nheads_kv, b)
+    T = tma_tensor
+    shape_packed = (
+        (qhead_per_kvhead, T.shape[0][1]),
+        *[T.shape[i] for i in range(1, head_idx)],
+        T.shape[0][0] // qhead_per_kvhead,
+        *[T.shape[i] for i in range(head_idx, len(T.shape))],
+    )
+    stride_packed = (
+        *[T.stride[i] for i in range(head_idx)],
+        T.stride[0][0] * qhead_per_kvhead,
+        *[T.stride[i] for i in range(head_idx, len(T.shape))],
+    )
+    tma_tensor = cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed))
+    return tma_atom, tma_tensor
+
+
+def unpack_gqa_layout(T, qhead_per_kvhead, head_idx):
+    """Reverse of pack_gqa_layout: unfold qhead_per_kvhead from the seqlen dimension (mode 0).
+
+    The head dimension is at mode ``head_idx``.  Modes before it (1..head_idx-1)
+    are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept
+    as-is (e.g. batch).
+
+    For Q/O tensors (head_idx=2):
+        ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...) -> (seqlen_q, headdim, nheads, batch, ...)
+    For LSE tensors (head_idx=1):
+        ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...) -> (seqlen_q, nheads, batch, ...)
+    """
+    seqlen_stride = T.stride[0][1]
+    head_stride = T.stride[0][0]
+    shape_unpacked = (
+        T.shape[0][1],
+        *[T.shape[i] for i in range(1, head_idx)],
+        T.shape[head_idx] * qhead_per_kvhead,
+        *[T.shape[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    stride_unpacked = (
+        seqlen_stride,
+        *[T.stride[i] for i in range(1, head_idx)],
+        head_stride,
+        *[T.stride[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    return cute.make_tensor(T.iterator, cute.make_layout(shape_unpacked, stride=stride_unpacked))
+
+
+@dataclass
+class PackGQA:
+    m_block_size: cutlass.Constexpr[int]
+    head_dim_padded: cutlass.Constexpr[int]
+    check_hdim_oob: cutlass.Constexpr[bool]
+    qhead_per_kvhead: cutlass.Constexpr[bool]
+
+    @cute.jit
+    def compute_ptr(
+        self,
+        tensor: cute.Tensor,
+        cRows: cute.Tensor,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        threads_per_row: cutlass.Constexpr[int],
+        num_threads: cutlass.Constexpr[int],
+    ):
+        num_ptr_per_thread = cute.ceil_div(cute.size(cRows), threads_per_row)
+        tPrPtr = cute.make_rmem_tensor(num_ptr_per_thread, cutlass.Int64)
+        for i in cutlass.range_constexpr(num_ptr_per_thread):
+            row = i * num_threads + cRows[tidx % threads_per_row][0]
+            idx = block * self.m_block_size + row
+            m_idx = idx // self.qhead_per_kvhead
+            h_idx = idx - m_idx * self.qhead_per_kvhead
+            tPrPtr[i] = utils.elem_pointer(tensor, ((h_idx, m_idx),)).toint()
+        return tPrPtr
+
+    @cute.jit
+    def load_Q(
+        self,
+        mQ: cute.Tensor,  # ((qhead_per_kvhead, seqlen_q), headdim)
+        sQ: cute.Tensor,  # (m_block_size, head_dim_padded)
+        gmem_tiled_copy: cute.TiledCopy,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        gmem_thr_copy = gmem_tiled_copy.get_slice(tidx)
+        cQ = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        tQsQ = gmem_thr_copy.partition_D(sQ)
+        tQcQ = gmem_thr_copy.partition_S(cQ)
+        t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ)
+        tQpQ = utils.predicate_k(tQcQ, limit=mQ.shape[1])
+        tQcQ_row = tQcQ[0, None, 0]
+        threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0]
+        assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE"
+        num_threads = gmem_tiled_copy.size
+        tPrQPtr = self.compute_ptr(mQ[None, 0], tQcQ_row, tidx, block, threads_per_row, num_threads)
+        for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])):
+            q_ptr_i64 = utils.shuffle_sync(
+                tPrQPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row
+            )
+            q_gmem_ptr = cute.make_ptr(
+                mQ.element_type, q_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            if (
+                t0QcQ[0, m, 0][0]
+                < seqlen * self.qhead_per_kvhead - block * self.m_block_size - tQcQ_row[0][0]
+            ):
+                mQ_cur = cute.make_tensor(q_gmem_ptr, (self.head_dim_padded,))
+                elems_per_load = cute.size(tQsQ.shape[0][0])
+                mQ_cur_copy = cute.tiled_divide(mQ_cur, (elems_per_load,))
+                for k in cutlass.range_constexpr(cute.size(tQsQ.shape[2])):
+                    ki = tQcQ[0, 0, k][1] // elems_per_load
+                    cute.copy(
+                        gmem_thr_copy,
+                        mQ_cur_copy[None, ki],
+                        tQsQ[None, m, k],
+                        pred=tQpQ[None, m, k] if cutlass.const_expr(self.check_hdim_oob) else None,
+                    )
+            # We don't need to clear the sQ smem tiles since we'll only write out the valid outputs
+
+    @cute.jit
+    def store_LSE(
+        self,
+        mLSE: cute.Tensor,  # (qhead_per_kvhead, seqlen_q)
+        tLSErLSE: cute.Tensor,  # (m_block_size, head_dim_padded)
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        caccO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        taccOcO = thr_mma.partition_C(caccO)
+        taccOcO_row = layout_utils.reshape_acc_to_mn(taccOcO)[None, 0]
+        assert cute.size(tLSErLSE) == cute.size(taccOcO_row)
+        threads_per_row = tiled_mma.tv_layout_C.shape[0][0]
+        assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE"
+        assert cute.size(tLSErLSE) <= threads_per_row
+        num_threads = tiled_mma.size
+        tPrLSEPtr = self.compute_ptr(mLSE, taccOcO_row, tidx, block, threads_per_row, num_threads)
+        for m in cutlass.range_constexpr(cute.size(tLSErLSE)):
+            lse_ptr_i64 = utils.shuffle_sync(
+                tPrLSEPtr[m // threads_per_row],
+                m % threads_per_row,
+                width=threads_per_row,
+            )
+            lse_gmem_ptr = cute.make_ptr(
+                mLSE.element_type, lse_ptr_i64, cute.AddressSpace.gmem, assumed_align=4
+            )
+            row = block * self.m_block_size + taccOcO_row[m][0]
+            # Only the thread corresponding to column 0 writes out the lse to gmem
+            if taccOcO[0][1] == 0 and row < seqlen * self.qhead_per_kvhead:
+                mLSE_copy = cute.make_tensor(lse_gmem_ptr, (1,))
+                mLSE_copy[0] = tLSErLSE[m]
+
+    @cute.jit
+    def store_O(
+        self,
+        mO: cute.Tensor,  # ((qhead_per_kvhead, seqlen_q), headdim)
+        tOrO: cute.Tensor,  # (m_block_size, head_dim_padded) split across threads according to gmem_tiled_copy
+        gmem_tiled_copy: cute.TiledCopy,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        gmem_thr_copy = gmem_tiled_copy.get_slice(tidx)
+        cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        tOcO = gmem_thr_copy.partition_S(cO)
+        t0OcO = gmem_thr_copy.get_slice(0).partition_S(cO)
+        tOpO = utils.predicate_k(tOcO, limit=mO.shape[1])
+        tOcO_row = tOcO[0, None, 0]
+        threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0]
+        assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE"
+        num_threads = gmem_tiled_copy.size
+        tPrOPtr = self.compute_ptr(mO[None, 0], tOcO_row, tidx, block, threads_per_row, num_threads)
+        for m in cutlass.range_constexpr(cute.size(tOrO.shape[1])):
+            o_ptr_i64 = utils.shuffle_sync(
+                tPrOPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row
+            )
+            o_gmem_ptr = cute.make_ptr(
+                mO.element_type, o_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            if (
+                t0OcO[0, m, 0][0]
+                < seqlen * self.qhead_per_kvhead - block * self.m_block_size - tOcO_row[0][0]
+            ):
+                mO_cur = cute.make_tensor(o_gmem_ptr, (self.head_dim_padded,))
+                elems_per_load = cute.size(tOrO.shape[0][0])
+                mO_cur_copy = cute.tiled_divide(mO_cur, (elems_per_load,))
+                for k in cutlass.range_constexpr(cute.size(tOrO.shape[2])):
+                    ki = tOcO[0, 0, k][1] // elems_per_load
+                    cute.copy(
+                        gmem_thr_copy,
+                        tOrO[None, m, k],
+                        mO_cur_copy[None, ki],
+                        pred=tOpO[None, m, k] if cutlass.const_expr(self.check_hdim_oob) else None,
+                    )
@@ -0,0 +1,247 @@
+from typing import Type
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.nvgpu import cpasync
+from cutlass import Int32, const_expr
+
+from flash_attn.cute import utils
+from quack.cute_dsl_utils import ParamsBase
+from cutlass.cute import FastDivmodDivisor
+
+import math
+
+
+@dataclass
+class PagedKVManager(ParamsBase):
+    mPageTable: cute.Tensor
+    mK_paged: cute.Tensor
+    mV_paged: cute.Tensor
+    thread_idx: Int32
+
+    page_size_divmod: FastDivmodDivisor
+    seqlen_k: Int32
+    leftpad_k: Int32
+    n_block_size: Int32
+    num_threads: cutlass.Constexpr[Int32]
+    head_dim_padded: cutlass.Constexpr[Int32]
+    head_dim_v_padded: cutlass.Constexpr[Int32]
+
+    arch: cutlass.Constexpr[Int32]
+    v_gmem_transposed: cutlass.Constexpr[bool]
+
+    gmem_threads_per_row: cutlass.Constexpr[Int32]
+    page_entry_per_thread: Int32
+    async_copy_elems: Int32
+
+    gmem_tiled_copy_KV: cute.TiledCopy
+    gmem_thr_copy_KV: cute.TiledCopy
+    tPrPage: cute.Tensor
+    tPrPageOffset: cute.Tensor
+    tKpK: cute.Tensor
+    tVpV: cute.Tensor
+
+    @staticmethod
+    def create(
+        mPageTable: cute.Tensor,
+        mK_paged: cute.Tensor,
+        mV_paged: cute.Tensor,
+        page_size_divmod: FastDivmodDivisor,
+        bidb: Int32,
+        bidh: Int32,
+        thread_idx: Int32,
+        seqlen_k: Int32,
+        leftpad_k: Int32,
+        n_block_size: cutlass.Constexpr[Int32],
+        head_dim_padded: cutlass.Constexpr[Int32],
+        head_dim_v_padded: cutlass.Constexpr[Int32],
+        num_threads: cutlass.Constexpr[Int32],
+        dtype: Type[cutlass.Numeric],
+        arch: cutlass.Constexpr[int] = 100,
+    ):
+        # SM100 transposes V in gmem to (dv, page_size, num_pages);
+        # SM90 keeps V as (page_size, dv, num_pages), same layout as K.
+        v_gmem_transposed = arch != 90
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // dtype.width
+        dtype_bytes = dtype.width // 8
+        gmem_k_block_size = math.gcd(
+            head_dim_padded,
+            head_dim_v_padded,
+            128 // dtype_bytes,
+        )
+        assert gmem_k_block_size % async_copy_elems == 0
+        gmem_threads_per_row = gmem_k_block_size // async_copy_elems
+        assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0
+        atom_async_copy = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            dtype,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        thr_layout = cute.make_ordered_layout(
+            (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        val_layout = cute.make_layout((1, async_copy_elems))
+        gmem_tiled_copy_KV = cute.make_tiled_copy_tv(atom_async_copy, thr_layout, val_layout)
+        gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx)
+        page_entry_per_thread = n_block_size // num_threads
+
+        tPrPage = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
+        tPrPageOffset = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
+
+        mPageTable = mPageTable[bidb, None]
+        mK_paged = mK_paged[None, None, bidh, None]
+        mV_paged = mV_paged[None, None, bidh, None]
+
+        cK = cute.make_identity_tensor((n_block_size, head_dim_padded))
+        tKcK = gmem_thr_copy_KV.partition_S(cK)
+        tKpK = utils.predicate_k(tKcK, limit=mK_paged.shape[1])
+
+        if const_expr(head_dim_padded == head_dim_v_padded):
+            tVpV = tKpK
+        else:
+            cV = cute.make_identity_tensor((n_block_size, head_dim_v_padded))
+            tVcV = gmem_thr_copy_KV.partition_S(cV)
+            # When V is transposed in gmem, dv is shape[0]; otherwise dv is shape[1] (same as K)
+            V_limit = cute.size(mV_paged.shape[0 if v_gmem_transposed else 1])
+            tVpV = utils.predicate_k(tVcV, limit=V_limit)
+
+        return PagedKVManager(
+            mPageTable,
+            mK_paged,
+            mV_paged,
+            thread_idx,
+            page_size_divmod,
+            seqlen_k,
+            leftpad_k,
+            n_block_size,
+            num_threads,
+            head_dim_padded,
+            head_dim_v_padded,
+            arch,
+            v_gmem_transposed,
+            gmem_threads_per_row,
+            page_entry_per_thread,
+            async_copy_elems,
+            gmem_tiled_copy_KV,
+            gmem_thr_copy_KV,
+            tPrPage,
+            tPrPageOffset,
+            tKpK,
+            tVpV,
+        )
+
+    @cute.jit
+    def load_page_table(self, n_block: Int32):
+        for i in cutlass.range(self.page_entry_per_thread, unroll=1):
+            row = (
+                i * self.num_threads
+                + (self.thread_idx % self.gmem_threads_per_row)
+                * (self.num_threads // self.gmem_threads_per_row)
+                + (self.thread_idx // self.gmem_threads_per_row)
+            )
+            row_idx = n_block * self.n_block_size + row
+
+            page_idx, page_offset = divmod(row_idx + self.leftpad_k, self.page_size_divmod)
+
+            is_valid = (
+                (i + 1) * self.num_threads <= self.n_block_size or row < self.n_block_size
+            ) and row_idx < self.seqlen_k
+            page = self.mPageTable[page_idx] if is_valid else 0
+
+            self.tPrPage[i] = page
+            self.tPrPageOffset[i] = page_offset
+
+    @cute.jit
+    def compute_X_ptr(self, K_or_V: str, d_offset: int = 0):
+        tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64)
+        mX = self.mK_paged if const_expr(K_or_V == "K") else self.mV_paged
+        # K is always (page_size, d, num_pages). V matches K when not transposed,
+        # but is (dv, page_size, num_pages) when transposed (SM100).
+        transposed = const_expr(K_or_V == "V" and self.v_gmem_transposed)
+        for i in cutlass.range(self.page_entry_per_thread, unroll=1):
+            page = self.tPrPage[i]
+            page_offset = self.tPrPageOffset[i]
+            if const_expr(transposed):
+                tPrXPtr[i] = utils.elem_pointer(mX, (d_offset, page_offset, page)).toint()
+            else:
+                tPrXPtr[i] = utils.elem_pointer(mX, (page_offset, d_offset, page)).toint()
+        return tPrXPtr
+
+    @cute.jit
+    def _flatten_smem_sm100(self, sX: cute.Tensor, K_or_V: str):
+        """Flatten SM100 smem ((a,b), cta_split, k) to (a,(b,k)); transpose V to (d,page_size)."""
+        sX_pi = cute.make_tensor(
+            sX.iterator,
+            cute.make_layout(
+                (sX.shape[0][0], (sX.shape[0][1], sX.shape[2])),
+                stride=(sX.stride[0][0], (sX.stride[0][1], sX.stride[2])),
+            ),
+        )
+        if const_expr(K_or_V == "V"):
+            sX_pi = cute.make_tensor(sX_pi.iterator, cute.select(sX_pi.layout, mode=[1, 0]))
+        return sX_pi
+
+    @cute.jit
+    def _copy_row_async(
+        self,
+        tXsX: cute.Tensor,
+        tXcX: cute.Tensor,
+        mX_paged_cur_copy: cute.Tensor,
+        m: Int32,
+        should_load: cute.Tensor,
+    ):
+        """Issue cp.async copies for one row across all k-tiles."""
+        for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
+            ki = tXcX[0, 0, k][1] // self.async_copy_elems
+            mX_paged_cur_copy_ki = mX_paged_cur_copy[None, ki]
+            tXsX_k = tXsX[None, m, k]
+            mX_paged_cur_copy_ki = cute.make_tensor(mX_paged_cur_copy_ki.iterator, tXsX_k.layout)
+            cute.copy(
+                self.gmem_tiled_copy_KV,
+                mX_paged_cur_copy_ki,
+                tXsX_k,
+                pred=should_load,
+            )
+
+    @cute.jit
+    def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str):
+        assert K_or_V in ("K", "V")
+
+        tPrXPtr = self.compute_X_ptr(K_or_V)
+
+        if const_expr(self.arch == 90):
+            # SM90: sX is already stage-sliced by caller (sK[None, None, stage]).
+            # Flatten hierarchical modes to get (n_block_size, head_dim).
+            sX_pi = cute.group_modes(sX, 0, 1)
+            # SM90 does NOT transpose V here (it's transposed via utils.transpose_view before MMA)
+        else:
+            sX_pi = self._flatten_smem_sm100(sX, K_or_V)
+
+        head_dim = self.head_dim_v_padded if const_expr(K_or_V == "V") else self.head_dim_padded
+        cX = cute.make_identity_tensor((self.n_block_size, head_dim))
+        tXsX = self.gmem_thr_copy_KV.partition_D(sX_pi)
+        tXcX = self.gmem_thr_copy_KV.partition_S(cX)
+        tXc0X = self.gmem_thr_copy_KV.get_slice(0).partition_S(cX)
+
+        seqlenk_row_limit = (
+            self.seqlen_k - n_block * self.n_block_size - tXcX[0][0] if n_block >= 0 else 0
+        )
+        for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
+            row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit
+            should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], cute.Boolean)
+            should_load.fill(row_valid)
+
+            x_ptr_i64 = utils.shuffle_sync(
+                tPrXPtr[m // self.gmem_threads_per_row],
+                m % self.gmem_threads_per_row,
+                width=self.gmem_threads_per_row,
+            )
+            x_gmem_ptr = cute.make_ptr(
+                self.mK_paged.element_type, x_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            mX_paged_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
+            mX_paged_cur_copy = cute.tiled_divide(mX_paged_cur, (self.async_copy_elems,))
+            self._copy_row_async(tXsX, tXcX, mX_paged_cur_copy, m, should_load)
@@ -0,0 +1,402 @@
+# Copyright (c) 2025, Tri Dao.
+
+from typing import Optional
+from dataclasses import dataclass
+
+import cutlass.cute as cute
+from cutlass import Boolean, Int32, const_expr
+from cutlass.cutlass_dsl import if_generate, dsl_user_op
+from cutlass.pipeline import PipelineState
+from cutlass.pipeline import PipelineUserType
+from cutlass.pipeline import NamedBarrier as NamedBarrierOg
+from cutlass.pipeline import PipelineAsync as PipelineAsyncOg
+from cutlass.pipeline import PipelineCpAsync as PipelineCpAsyncOg
+from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg
+from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg
+from cutlass.pipeline import PipelineUmmaAsync as PipelineUmmaAsyncOg
+from cutlass.pipeline import PipelineAsyncUmma as PipelineAsyncUmmaOg
+
+
+def _override_create(parent_cls, child_cls):
+    """Create a static factory that constructs parent_cls then re-classes to child_cls."""
+
+    @staticmethod
+    def create(*args, **kwargs):
+        obj = parent_cls.create(*args, **kwargs)
+        # Can't assign to __class__ directly since the dataclass is frozen
+        object.__setattr__(obj, "__class__", child_cls)
+        return obj
+
+    return create
+
+
+def _make_state(index: Int32, phase: Int32) -> PipelineState:
+    """Construct a PipelineState from index and phase (count/stages unused by callers)."""
+    return PipelineState(stages=0, count=Int32(0), index=index, phase=phase)
+
+
+class PipelineStateSimple:
+    """
+    Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer.
+    Use a single Int32 to store both the index and phase bit, then we use divmod to get the
+    index and phase. If stages is a power of 2, divmod turns into bit twiddling.
+    """
+
+    def __init__(self, stages: int, phase_index: Int32):
+        self._stages = stages
+        self._phase_index = phase_index
+
+    def clone(self) -> "PipelineStateSimple":
+        return PipelineStateSimple(self.stages, self._phase_index)
+
+    @property
+    def stages(self) -> int:
+        return self._stages
+
+    @property
+    def index(self) -> Int32:
+        if const_expr(self._stages == 1):
+            return Int32(0)
+        else:
+            return self._phase_index % self._stages
+
+    @property
+    def phase(self) -> Int32:
+        # PTX docs say that the phase parity needs to be 0 or 1, so by right we need to
+        # take modulo 2. But in practice just passing the phase in without modulo works fine.
+        if const_expr(self._stages == 1):
+            return self._phase_index
+        else:
+            return self._phase_index // self._stages
+
+    def advance(self):
+        if const_expr(self._stages == 1):
+            self._phase_index ^= 1
+        else:
+            self._phase_index += 1
+
+    def __extract_mlir_values__(self):
+        phase_index = self._phase_index
+        return [phase_index.ir_value()]
+
+    def __new_from_mlir_values__(self, values):
+        return PipelineStateSimple(self.stages, Int32(values[0]))
+
+
+def make_pipeline_state(type: PipelineUserType, stages: int):
+    """
+    Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1.
+    """
+    if type is PipelineUserType.Producer:
+        return PipelineStateSimple(stages, Int32(stages))
+    elif type is PipelineUserType.Consumer:
+        return PipelineStateSimple(stages, Int32(0))
+    else:
+        assert False, "Error: invalid PipelineUserType specified for make_pipeline_state."
+
+
+# ── Shared helpers ───────────────────────────────────────────────────────────
+
+
+def _call_with_elect_one(parent_method, self, state, elect_one, syncwarp, loc, ip):
+    """Optionally wrap a parent pipeline method call in sync_warp + elect_one."""
+    if const_expr(elect_one):
+        if const_expr(syncwarp):
+            cute.arch.sync_warp()
+        with cute.arch.elect_one():
+            parent_method(self, state, loc=loc, ip=ip)
+    else:
+        parent_method(self, state, loc=loc, ip=ip)
+
+
+# ── Mixin: _w_index / _w_index_phase variants that delegate to parent ───────
+# Each parent class has PipelineState-based methods (producer_acquire, producer_commit,
+# consumer_wait, consumer_release). The _w_index_phase variants just construct a
+# PipelineState from (index, phase) and delegate.
+
+
+class _PipelineIndexPhaseMixin:
+    """Mixin providing _w_index_phase / _w_index methods that delegate to PipelineState-based parents."""
+
+    @dsl_user_op
+    def producer_acquire_w_index_phase(
+        self,
+        index: Int32,
+        phase: Int32,
+        try_acquire_token: Optional[Boolean] = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        state = _make_state(index, phase)
+        # Call the parent's producer_acquire (which takes PipelineState)
+        self.producer_acquire(state, try_acquire_token, loc=loc, ip=ip)
+
+    @dsl_user_op
+    def producer_commit_w_index(self, index: Int32, *, loc=None, ip=None):
+        state = _make_state(index, Int32(0))
+        self.producer_commit(state, loc=loc, ip=ip)
+
+    @dsl_user_op
+    def consumer_wait_w_index_phase(
+        self,
+        index: Int32,
+        phase: Int32,
+        try_wait_token: Optional[Boolean] = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        state = _make_state(index, phase)
+        self.consumer_wait(state, try_wait_token, loc=loc, ip=ip)
+
+    @dsl_user_op
+    def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None):
+        state = _make_state(index, Int32(0))
+        self.consumer_release(state, loc=loc, ip=ip)
+
+
+# ── NamedBarrier ─────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class NamedBarrier(NamedBarrierOg):
+    create = _override_create(NamedBarrierOg, None)  # patched below
+
+    @dsl_user_op
+    def arrive_w_index(self, index: Int32, *, loc=None, ip=None) -> None:
+        """
+        The aligned flavor of arrive is used when all threads in the CTA will execute the
+        same instruction. See PTX documentation.
+        """
+        cute.arch.barrier_arrive(
+            barrier_id=self.barrier_id + index,
+            number_of_threads=self.num_threads,
+            loc=loc,
+            ip=ip,
+        )
+
+    @dsl_user_op
+    def arrive_and_wait_w_index(self, index: Int32, *, loc=None, ip=None) -> None:
+        cute.arch.barrier(
+            barrier_id=self.barrier_id + index,
+            number_of_threads=self.num_threads,
+            loc=loc,
+            ip=ip,
+        )
+
+
+NamedBarrier.create = _override_create(NamedBarrierOg, NamedBarrier)
+
+
+# ── PipelineAsync ────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineAsync(_PipelineIndexPhaseMixin, PipelineAsyncOg):
+    """
+    PipelineAsync with optional elect_one for producer_commit and consumer_release.
+
+    When elect_one_*=True (set at create time), only one elected thread per warp
+    signals the barrier arrive. This is useful when the mask count is set to 1 per warp.
+
+    Args (to create):
+        elect_one_commit: If True, only elected thread signals producer_commit.
+        syncwarp_before_commit: If True (default), issue syncwarp before elect_one.
+        elect_one_release: If True, only elected thread signals consumer_release.
+        syncwarp_before_release: If True (default), issue syncwarp before elect_one.
+            Set syncwarp to False when threads are already converged (e.g. after wgmma wait_group).
+    """
+
+    _elect_one_commit: bool = False
+    _syncwarp_before_commit: bool = True
+    _elect_one_release: bool = False
+    _syncwarp_before_release: bool = True
+
+    @staticmethod
+    def create(
+        *args,
+        elect_one_commit: bool = False,
+        syncwarp_before_commit: bool = True,
+        elect_one_release: bool = False,
+        syncwarp_before_release: bool = True,
+        **kwargs,
+    ):
+        obj = PipelineAsyncOg.create(*args, **kwargs)
+        object.__setattr__(obj, "__class__", PipelineAsync)
+        object.__setattr__(obj, "_elect_one_commit", elect_one_commit)
+        object.__setattr__(obj, "_syncwarp_before_commit", syncwarp_before_commit)
+        object.__setattr__(obj, "_elect_one_release", elect_one_release)
+        object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release)
+        return obj
+
+    @dsl_user_op
+    def producer_commit(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineAsyncOg.producer_commit,
+            self,
+            state,
+            self._elect_one_commit,
+            self._syncwarp_before_commit,
+            loc,
+            ip,
+        )
+
+    @dsl_user_op
+    def consumer_release(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineAsyncOg.consumer_release,
+            self,
+            state,
+            self._elect_one_release,
+            self._syncwarp_before_release,
+            loc,
+            ip,
+        )
+
+    # _w_index variants inherited from _PipelineIndexPhaseMixin, which delegate
+    # to producer_commit / consumer_release above.
+
+
+# ── PipelineCpAsync ──────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineCpAsync(_PipelineIndexPhaseMixin, PipelineCpAsyncOg):
+    _elect_one_release: bool = False
+    _syncwarp_before_release: bool = True
+
+    @staticmethod
+    def create(
+        *args,
+        elect_one_release: bool = False,
+        syncwarp_before_release: bool = True,
+        **kwargs,
+    ):
+        obj = PipelineCpAsyncOg.create(*args, **kwargs)
+        object.__setattr__(obj, "__class__", PipelineCpAsync)
+        object.__setattr__(obj, "_elect_one_release", elect_one_release)
+        object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release)
+        return obj
+
+    @dsl_user_op
+    def consumer_release(self, state: PipelineState, *, loc=None, ip=None):
+        _call_with_elect_one(
+            PipelineCpAsyncOg.consumer_release,
+            self,
+            state,
+            self._elect_one_release,
+            self._syncwarp_before_release,
+            loc,
+            ip,
+        )
+
+    # _w_index variants inherited from _PipelineIndexPhaseMixin.
+
+
+# ── PipelineTmaAsync ────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineTmaAsync(_PipelineIndexPhaseMixin, PipelineTmaAsyncOg):
+    """Override producer_acquire to take in extra_tx_count parameter."""
+
+    @dsl_user_op
+    def producer_acquire(
+        self,
+        state: PipelineState,
+        try_acquire_token: Optional[Boolean] = None,
+        extra_tx_count: int = 0,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks.
+        """
+        if_generate(
+            try_acquire_token is None or try_acquire_token == 0,
+            lambda: self.sync_object_empty.wait(state.index, state.phase, loc=loc, ip=ip),
+            loc=loc,
+            ip=ip,
+        )
+        if const_expr(extra_tx_count == 0):
+            self.sync_object_full.arrive(state.index, self.producer_mask, loc=loc, ip=ip)
+        else:
+            tx_count = self.sync_object_full.tx_count + extra_tx_count
+            self.sync_object_full.arrive_and_expect_tx(state.index, tx_count, loc=loc, ip=ip)
+
+
+PipelineTmaAsync.create = _override_create(PipelineTmaAsyncOg, PipelineTmaAsync)
+
+
+# ── PipelineTmaUmma ─────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineTmaUmma(_PipelineIndexPhaseMixin, PipelineTmaUmmaOg):
+    """Override producer_acquire to take in extra_tx_count parameter."""
+
+    @dsl_user_op
+    def producer_acquire(
+        self,
+        state: PipelineState,
+        try_acquire_token: Optional[Boolean] = None,
+        extra_tx_count: int = 0,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks.
+        """
+        if_generate(
+            try_acquire_token is None or try_acquire_token == 0,
+            lambda: self.sync_object_empty.wait(state.index, state.phase, loc=loc, ip=ip),
+            loc=loc,
+            ip=ip,
+        )
+        if const_expr(extra_tx_count == 0):
+            if_generate(
+                self.is_leader_cta,
+                lambda: self.sync_object_full.arrive(
+                    state.index, self.producer_mask, loc=loc, ip=ip
+                ),
+                loc=loc,
+                ip=ip,
+            )
+        else:
+            tx_count = self.sync_object_full.tx_count + extra_tx_count
+            if_generate(
+                self.is_leader_cta,
+                lambda: self.sync_object_full.arrive_and_expect_tx(
+                    state.index, tx_count, loc=loc, ip=ip
+                ),
+                loc=loc,
+                ip=ip,
+            )
+
+
+PipelineTmaUmma.create = _override_create(PipelineTmaUmmaOg, PipelineTmaUmma)
+
+
+# ── PipelineUmmaAsync ───────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineUmmaAsync(_PipelineIndexPhaseMixin, PipelineUmmaAsyncOg):
+    pass
+
+
+PipelineUmmaAsync.create = _override_create(PipelineUmmaAsyncOg, PipelineUmmaAsync)
+
+
+# ── PipelineAsyncUmma ───────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PipelineAsyncUmma(_PipelineIndexPhaseMixin, PipelineAsyncUmmaOg):
+    pass
+
+
+PipelineAsyncUmma.create = _override_create(PipelineAsyncUmmaOg, PipelineAsyncUmma)
@@ -0,0 +1,68 @@
+[build-system]
+requires = ["setuptools>=75"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flash-attn-4"
+version = "4.0.0b20.dev2+fastvideo.82d6441"
+description = "Flash Attention CUTE (CUDA Template Engine) implementation"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "BSD-3-Clause"
+authors = [
+    {name = "Tri Dao"},
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+
+dependencies = [
+    "nvidia-cutlass-dsl==4.6.0.dev0",
+    "torch",
+    "einops",
+    "typing_extensions",
+    "apache-tvm-ffi>=0.1.5,<0.2",
+    "torch-c-dlpack-ext",
+    "quack-kernels>=0.5.0",
+]
+
+[project.optional-dependencies]
+cu13 = ["nvidia-cutlass-dsl[cu13]==4.6.0.dev0"]
+dev = [
+    "pytest",
+    "pytest-xdist",
+    "ruff",
+]
+
+[project.urls]
+Homepage = "https://github.com/Dao-AILab/flash-attention"
+Repository = "https://github.com/Dao-AILab/flash-attention"
+
+[tool.setuptools]
+packages = ["flash_attn.cute"]
+package-dir = {"flash_attn.cute" = "."}
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+    { index = "pytorch-cu130", marker = "extra == 'cu13'" },
+]
+
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+ignore = [
+    "E731",  # do not assign a lambda expression, use a def
+    "E741",  # Do not use variables named 'I', 'O', or 'l'
+    "F841",  # local variable is assigned to but never used
+    "D102",  # Missing docstring in public methods
+]
@@ -0,0 +1,302 @@
+from typing import Optional
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+
+from quack import copy_utils
+
+"""
+This consolidates all the info related to sequence length. This is so that we can do all
+the gmem reads once at the beginning of each tile, rather than having to repeat these reads
+to compute various things like n_block_min, n_block_max, etc.
+"""
+
+
+@dataclass(frozen=True)
+class SeqlenInfo:
+    offset: Int32
+    offset_padded: Int32
+    seqlen: Int32
+    has_cu_seqlens: cutlass.Constexpr[bool] = False
+
+    @staticmethod
+    def create(
+        batch_idx: Int32,
+        seqlen_static: Int32,
+        cu_seqlens: Optional[cute.Tensor] = None,
+        seqused: Optional[cute.Tensor] = None,
+        tile: cutlass.Constexpr[int] = 128,
+    ):
+        offset = 0 if const_expr(cu_seqlens is None) else cu_seqlens[batch_idx]
+        offset_padded = (
+            0
+            if const_expr(cu_seqlens is None)
+            # Add divby so that the compiler knows the alignment when moving by offset_padded
+            else cute.assume((offset + batch_idx * tile) // tile * tile, divby=tile)
+        )
+        if const_expr(seqused is not None):
+            seqlen = seqused[batch_idx]
+        elif const_expr(cu_seqlens is not None):
+            seqlen = cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]
+        else:
+            seqlen = seqlen_static
+        return SeqlenInfo(offset, offset_padded, seqlen, has_cu_seqlens=cu_seqlens is not None)
+
+    def offset_batch(
+        self,
+        mT: cute.Tensor,
+        batch_idx: Int32,
+        dim: int,
+        padded: cutlass.Constexpr[bool] = False,
+        multiple: int = 1,
+    ) -> cute.Tensor:
+        """Offset a tensor by batch index. batch dim is at position `dim`, seqlen is at dim=0."""
+        if const_expr(not self.has_cu_seqlens):
+            idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mT) - 1 - dim)
+            return mT[idx]
+        else:
+            off = multiple * (self.offset if const_expr(not padded) else self.offset_padded)
+            offset = off if const_expr(cute.rank(mT.shape[0]) == 1) else (0, off)
+            idx = (offset,) + (None,) * (cute.rank(mT) - 1)
+            return cute.domain_offset(idx, mT)
+
+
+@dataclass(frozen=True)
+class SeqlenInfoQK:
+    offset_q: Int32
+    offset_k: Int32
+    padded_offset_q: Int32
+    padded_offset_k: Int32
+    seqlen_q: Int32
+    seqlen_k: Int32
+    m_block_offset: Int32
+    block_idx_offset: Int32
+    num_n_blocks: Int32
+    has_cu_seqlens_q: cutlass.Constexpr[bool]
+    has_cu_seqlens_k: cutlass.Constexpr[bool]
+    has_seqused_q: cutlass.Constexpr[bool]
+    has_seqused_k: cutlass.Constexpr[bool]
+
+    @staticmethod
+    def create(
+        batch_idx: Int32,
+        seqlen_q_static: Int32,
+        seqlen_k_static: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mCuTotalMBlocks: Optional[cute.Tensor] = None,
+        mCuBlockIdxOffsets: Optional[cute.Tensor] = None,
+        tile_m: cutlass.Constexpr[Int32] = 128,
+        tile_n: cutlass.Constexpr[Int32] = 128,
+    ):
+        offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx]
+        offset_k = 0 if const_expr(mCuSeqlensK is None) else mCuSeqlensK[batch_idx]
+        padded_offset_q = (
+            0
+            if const_expr(mCuSeqlensQ is None)
+            else cute.assume((offset_q + batch_idx * tile_m) // tile_m * tile_m, divby=tile_m)
+        )
+        padded_offset_k = (
+            0
+            if const_expr(mCuSeqlensK is None)
+            else cute.assume((offset_k + batch_idx * tile_n) // tile_n * tile_n, divby=tile_n)
+        )
+        if const_expr(mSeqUsedQ is not None):
+            seqlen_q = mSeqUsedQ[batch_idx]
+        else:
+            seqlen_q = (
+                seqlen_q_static
+                if const_expr(mCuSeqlensQ is None)
+                else mCuSeqlensQ[batch_idx + 1] - offset_q
+            )
+        if const_expr(mSeqUsedK is not None):
+            seqlen_k = mSeqUsedK[batch_idx]
+        else:
+            seqlen_k = (
+                seqlen_k_static
+                if const_expr(mCuSeqlensK is None)
+                else mCuSeqlensK[batch_idx + 1] - offset_k
+            )
+        m_block_offset = 0 if const_expr(mCuTotalMBlocks is None) else mCuTotalMBlocks[batch_idx]
+        num_n_blocks = (seqlen_k + tile_n - 1) // tile_n
+        block_idx_offset = (
+            mCuBlockIdxOffsets[batch_idx]
+            if const_expr(mCuBlockIdxOffsets is not None)
+            else m_block_offset * num_n_blocks
+        )
+        return SeqlenInfoQK(
+            offset_q,
+            offset_k,
+            padded_offset_q,
+            padded_offset_k,
+            seqlen_q,
+            seqlen_k,
+            m_block_offset,
+            block_idx_offset,
+            num_n_blocks,
+            has_cu_seqlens_q=mCuSeqlensQ is not None,
+            has_cu_seqlens_k=mCuSeqlensK is not None,
+            has_seqused_q=mSeqUsedQ is not None,
+            has_seqused_k=mSeqUsedK is not None,
+        )
+
+    def offset_batch_Q(
+        self,
+        mQ: cute.Tensor,
+        batch_idx: Int32,
+        dim: int,
+        padded: cutlass.Constexpr[bool] = False,
+        ragged: cutlass.Constexpr[bool] = False,
+    ) -> cute.Tensor:
+        """Seqlen must be the first dimension of mQ"""
+        if const_expr(not ragged):
+            if const_expr(not self.has_cu_seqlens_q):
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
+                return mQ[idx]
+            else:
+                offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q
+                offset_q = offset_q if const_expr(cute.rank(mQ.shape[0]) == 1) else (None, offset_q)
+                idx = (offset_q,) + (None,) * (cute.rank(mQ) - 1)
+                return cute.domain_offset(idx, mQ)
+        else:
+            if const_expr(not self.has_cu_seqlens_q):
+                offset_q = 0
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
+                mQ = mQ[idx]
+            else:
+                offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q
+            if const_expr(cute.rank(mQ.shape[0]) == 1):
+                return copy_utils.offset_ragged_tensor(
+                    mQ, offset_q, self.seqlen_q, ragged_dim=0, ptr_shift=True
+                )
+            else:  # PackGQA
+                assert cute.rank(mQ.shape[0]) == 2
+                # Unpack before calling offset_ragged_tensor, then pack
+                idx = ((None, None),) + (None,) * (cute.rank(mQ) - 1)
+                mQ = mQ[idx]
+                mQ = copy_utils.offset_ragged_tensor(
+                    mQ, offset_q, self.seqlen_q, ragged_dim=1, ptr_shift=True
+                )
+                return cute.group_modes(mQ, 0, 2)
+
+    def offset_batch_K(
+        self,
+        mK: cute.Tensor,
+        batch_idx: Int32,
+        dim: int,
+        padded: cutlass.Constexpr[bool] = False,
+        ragged: cutlass.Constexpr[bool] = False,
+        multiple: int = 1,
+    ) -> cute.Tensor:
+        """Seqlen must be the first dimension of mK"""
+        if const_expr(not ragged):
+            if const_expr(not self.has_cu_seqlens_k):
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
+                return mK[idx]
+            else:
+                offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k
+                offset_k *= multiple
+                idx = (offset_k,) + (None,) * (cute.rank(mK) - 1)
+                return cute.domain_offset(idx, mK)
+        else:
+            if const_expr(not self.has_cu_seqlens_k):
+                offset_k = 0
+                idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
+                mK = mK[idx]
+            else:
+                offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k
+                offset_k *= multiple
+            return copy_utils.offset_ragged_tensor(
+                mK, offset_k, self.seqlen_k, ragged_dim=0, ptr_shift=True
+            )
+
+
+@dataclass(frozen=True)
+class SeqlenInfoQKNewK:
+    """Sequence length info for append-KV with left-padding and new K support.
+
+    Extends SeqlenInfoQK with:
+    - leftpad_k: left padding for K (tokens to skip at the start of the KV cache)
+    - offset_k_new: offset into the new K tensor
+    - seqlen_k_og: original K length (before appending new K), excluding leftpad
+    - seqlen_k_new: length of new K to append
+    - seqlen_k: total K length (seqlen_k_og + seqlen_k_new)
+    - seqlen_rotary: position for rotary embedding computation
+    """
+
+    leftpad_k: Int32
+    offset_q: Int32
+    offset_k: Int32
+    offset_k_new: Int32
+    seqlen_q: Int32
+    seqlen_k_og: Int32
+    seqlen_k_new: Int32
+    seqlen_k: Int32
+    seqlen_rotary: Int32
+
+    @staticmethod
+    def create(
+        batch_idx: Int32,
+        seqlen_q_static: Int32,
+        seqlen_k_static: Int32,
+        shape_K_new_0: Int32,
+        mCuSeqlensQ: Optional[cute.Tensor] = None,
+        mCuSeqlensK: Optional[cute.Tensor] = None,
+        mCuSeqlensKNew: Optional[cute.Tensor] = None,
+        mSeqUsedQ: Optional[cute.Tensor] = None,
+        mSeqUsedK: Optional[cute.Tensor] = None,
+        mLeftpadK: Optional[cute.Tensor] = None,
+        mSeqlensRotary: Optional[cute.Tensor] = None,
+    ):
+        leftpad_k = 0 if const_expr(mLeftpadK is None) else mLeftpadK[batch_idx]
+        offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx]
+        if const_expr(mCuSeqlensK is not None):
+            offset_k = mCuSeqlensK[batch_idx] + leftpad_k
+        else:
+            offset_k = leftpad_k if const_expr(mCuSeqlensQ is not None) else 0
+        offset_k_new = 0 if const_expr(mCuSeqlensKNew is None) else mCuSeqlensKNew[batch_idx]
+        # seqlen_q
+        if const_expr(mSeqUsedQ is not None):
+            seqlen_q = mSeqUsedQ[batch_idx]
+        elif const_expr(mCuSeqlensQ is not None):
+            seqlen_q = mCuSeqlensQ[batch_idx + 1] - mCuSeqlensQ[batch_idx]
+        else:
+            seqlen_q = seqlen_q_static
+        # seqlen_k_og: original K length (excluding leftpad)
+        if const_expr(mSeqUsedK is not None):
+            seqlen_k_og = mSeqUsedK[batch_idx] - leftpad_k
+        elif const_expr(mCuSeqlensK is not None):
+            seqlen_k_og = mCuSeqlensK[batch_idx + 1] - mCuSeqlensK[batch_idx] - leftpad_k
+        else:
+            seqlen_k_og = (
+                seqlen_k_static - leftpad_k
+                if const_expr(mCuSeqlensQ is not None)
+                else seqlen_k_static
+            )
+        # seqlen_k_new
+        if const_expr(mCuSeqlensKNew is None):
+            seqlen_k_new = 0 if const_expr(mCuSeqlensQ is None) else shape_K_new_0
+        else:
+            seqlen_k_new = mCuSeqlensKNew[batch_idx + 1] - mCuSeqlensKNew[batch_idx]
+        seqlen_k = seqlen_k_og if const_expr(mCuSeqlensQ is None) else seqlen_k_og + seqlen_k_new
+
+        # seqlen_rotary: defaults to seqlen_k_og + leftpad_k unless explicitly provided
+        if const_expr(mSeqlensRotary is not None):
+            seqlen_rotary = mSeqlensRotary[batch_idx]
+        else:
+            seqlen_rotary = seqlen_k_og + leftpad_k
+        return SeqlenInfoQKNewK(
+            leftpad_k,
+            offset_q,
+            offset_k,
+            offset_k_new,
+            seqlen_q,
+            seqlen_k_og,
+            seqlen_k_new,
+            seqlen_k,
+            seqlen_rotary,
+        )
@@ -0,0 +1,298 @@
+# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin.
+
+
+"""Fused multi-head attention (FMHA) backward for the SM100 architecture using CUTE DSL.
+
+Constraints:
+* Supported head dimensions: 256 only
+* mma_tiler_mn must be 64,64
+* Batch size must be the same for Q, K, and V tensors
+"""
+
+import cuda.bindings.driver as cuda
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.typing import Int32
+
+from flash_attn.cute.sm100_hd256_2cta_fmha_backward_dqkernel import (
+    BlackwellFusedMultiHeadAttentionBackwardDQKernel,
+)
+from flash_attn.cute.sm100_hd256_2cta_fmha_backward_dkdvkernel import (
+    BlackwellFusedMultiHeadAttentionBackwardDKDVKernel,
+)
+from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
+from flash_attn.cute.utils import AuxData
+
+
+def _as_bshkrd_tensor(
+    tensor: cute.Tensor,
+    h_k: Int32,
+    h_r: Int32,
+    varlen: bool,
+) -> cute.Tensor:
+    """Normalize (B,S,H,D)/(S,H,D) tensors to (B,S,H_k,H_r,D) view."""
+    if cutlass.const_expr(cute.rank(tensor.layout) == 5):
+        return tensor
+    if cutlass.const_expr(cute.rank(tensor.layout) == 4):
+        return cute.make_tensor(
+            tensor.iterator,
+            cute.make_layout(
+                (tensor.shape[0], tensor.shape[1], h_k, h_r, tensor.shape[3]),
+                stride=(
+                    tensor.stride[0],
+                    tensor.stride[1],
+                    tensor.stride[2] * h_r,
+                    tensor.stride[2],
+                    tensor.stride[3],
+                ),
+            ),
+        )
+    assert cutlass.const_expr(cute.rank(tensor.layout) == 3), "Expected rank-3 varlen tensor"
+    assert cutlass.const_expr(varlen), "Rank-3 input is only valid for varlen backward"
+    return cute.make_tensor(
+        tensor.iterator,
+        cute.make_layout(
+            (1, tensor.shape[0], h_k, h_r, tensor.shape[2]),
+            stride=(
+                0,
+                tensor.stride[0],
+                tensor.stride[1] * h_r,
+                tensor.stride[1],
+                tensor.stride[2],
+            ),
+        ),
+    )
+
+
+def _as_shhb_tensor(
+    tensor: cute.Tensor,
+    h_k: Int32,
+    h_r: Int32,
+    b: Int32,
+    varlen: bool,
+) -> cute.Tensor:
+    """Normalize (B,H,S)/(H,S) tensors to (S, ((H_r, H_k), B)) view."""
+    if cutlass.const_expr(cute.rank(tensor.layout) == 3):
+        return cute.make_tensor(
+            tensor.iterator,
+            cute.make_layout(
+                (tensor.shape[2], ((h_r, h_k), tensor.shape[0])),
+                stride=(
+                    tensor.stride[2],
+                    ((tensor.stride[1], tensor.stride[1] * h_r), tensor.stride[0]),
+                ),
+            ),
+        )
+    assert cutlass.const_expr(cute.rank(tensor.layout) == 2), "Expected rank-2 varlen tensor"
+    assert cutlass.const_expr(varlen), "Rank-2 input is only valid for varlen backward"
+    return cute.make_tensor(
+        tensor.iterator,
+        cute.make_layout(
+            (tensor.shape[1], ((h_r, h_k), b)),
+            stride=(
+                tensor.stride[1],
+                ((tensor.stride[0], tensor.stride[0] * h_r), 0),
+            ),
+        ),
+    )
+
+
+class BlackwellFusedMultiHeadAttentionBackward:
+    """FMHA backward class for executing CuTeDSL kernel."""
+
+    def __init__(
+        self,
+        head_dim: int,
+        head_dim_v: int | None = None,
+        is_causal: bool = False,
+        is_local: bool = False,
+        qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+        is_persistent: bool = False,
+        deterministic: bool = False,
+        cluster_size: int = 1,
+        use_2cta_instrs: bool = False,
+        score_mod: cutlass.Constexpr | None = None,
+        score_mod_bwd: cutlass.Constexpr | None = None,
+        mask_mod: cutlass.Constexpr | None = None,
+        has_aux_tensors: cutlass.Constexpr = False,
+        q_subtile_factor: cutlass.Constexpr[int] = 1,
+        tile_m_dq: int = 128,
+        tile_n_dq: int = 128,
+        tile_m_dkdv: int = 128,
+        tile_n_dkdv: int = 64,
+        window_size_left: int | None = None,
+        window_size_right: int | None = None,
+        use_clc_scheduler: bool = False,
+    ):
+        """Initialization."""
+        head_dim_v = head_dim if head_dim_v is None else head_dim_v
+        assert head_dim == 256 and head_dim_v == 256, (
+            "SM100 dedicated backward kernel only supports (head_dim, head_dim_v) = (256, 256)"
+        )
+        assert not is_local, "SM100 backward with head_dim=256 does not support local attention"
+        assert tile_m_dq == 128 and tile_n_dq == 128, (
+            "SM100 dedicated backward kernel only supports tile_m_dq=128 and tile_n_dq=128"
+        )
+        assert tile_m_dkdv == 128 and tile_n_dkdv == 64, (
+            "SM100 dedicated backward kernel only supports tile_m_dkdv=128 and tile_n_dkdv=64"
+        )
+        assert score_mod is None and score_mod_bwd is None and mask_mod is None, (
+            "SM100 backward with head_dim=256 does not support score_mod/mask_mod"
+        )
+        assert not deterministic, (
+            "SM100 backward with head_dim=256 does not support deterministic mode"
+        )
+        assert not has_aux_tensors, "SM100 backward with head_dim=256 does not support aux_tensors"
+        assert cluster_size in (1, 2), (
+            "SM100 backward with head_dim=256 only supports cluster_size in {1, 2}"
+        )
+        assert use_2cta_instrs, "SM100 backward with head_dim=256 requires use_2cta_instrs=True"
+        # q_subtile_factor is accepted for interface parity with FlashAttentionBackwardSm100,
+        # but this dedicated kernel uses fixed internal behavior.
+
+        self.acc_dtype = cutlass.Float32
+        self.is_causal = is_causal
+        self.window_size_left = (
+            None if (window_size_left is None or window_size_left < 0) else window_size_left
+        )
+        self.window_size_right = (
+            None if (window_size_right is None or window_size_right < 0) else window_size_right
+        )
+        self.tile_m_dq = tile_m_dq
+        self.tile_n_dq = tile_n_dq
+        self.tile_m_dkdv = tile_m_dkdv
+        self.tile_n_dkdv = tile_n_dkdv
+        self.use_clc_scheduler = use_clc_scheduler
+
+        self.dq_kernel = BlackwellFusedMultiHeadAttentionBackwardDQKernel(
+            self.acc_dtype,
+            (self.tile_m_dq, self.tile_n_dq, 256),
+            self.is_causal,
+            self.window_size_left,
+            self.window_size_right,
+            False,  # is_persistent
+            False,  # split_head
+            use_clc_scheduler=self.use_clc_scheduler,
+        )
+        self.dkdv_kernel = BlackwellFusedMultiHeadAttentionBackwardDKDVKernel(
+            self.acc_dtype,
+            (self.tile_m_dkdv, self.tile_n_dkdv, 256),
+            self.is_causal,
+            self.window_size_left,
+            self.window_size_right,
+            use_clc_scheduler=self.use_clc_scheduler,
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        Q: cute.Tensor,
+        K: cute.Tensor,
+        V: cute.Tensor,
+        dO: cute.Tensor,
+        lse_log2: cute.Tensor,
+        dpsum: cute.Tensor,
+        dQ_accum: cute.Tensor | None,
+        dK: cute.Tensor,
+        dV: cute.Tensor,
+        scale_softmax: cutlass.Float32,
+        cumulative_s_q: cute.Tensor | None,
+        cumulative_s_k: cute.Tensor | None,
+        seqused_q: cute.Tensor | None = None,
+        seqused_k: cute.Tensor | None = None,
+        window_size_left: Int32 | None = None,
+        window_size_right: Int32 | None = None,
+        dQ_semaphore: cute.Tensor | None = None,
+        dK_semaphore: cute.Tensor | None = None,
+        dV_semaphore: cute.Tensor | None = None,
+        aux_data: AuxData = AuxData(),
+        block_sparse_tensors: cute.Tensor | None = None,
+        stream: cuda.CUstream = None,
+    ):
+        """Host function to launch CuTeDSL kernel."""
+        assert seqused_q is None and seqused_k is None, (
+            "SM100 backward with head_dim=256 does not support seqused_q/seqused_k"
+        )
+        assert window_size_left is None and window_size_right is None, (
+            "SM100 backward with head_dim=256 uses constructor-provided window sizes"
+        )
+        assert dQ_semaphore is None and dK_semaphore is None and dV_semaphore is None, (
+            "SM100 backward with head_dim=256 does not use semaphores"
+        )
+        assert block_sparse_tensors is None, (
+            "SM100 backward with head_dim=256 does not support block sparse tensors"
+        )
+        assert aux_data.tensors is None or len(aux_data.tensors) == 0, (
+            "SM100 backward with head_dim=256 does not support aux_tensors"
+        )
+        assert aux_data.scalars is None or len(aux_data.scalars) == 0, (
+            "SM100 backward with head_dim=256 does not support aux_scalars"
+        )
+        assert dQ_accum is not None, (
+            "SM100 backward with head_dim=256 expects dQ tensor at dQ_accum slot"
+        )
+        dQ = dQ_accum
+        varlen = cumulative_s_q is not None or cumulative_s_k is not None
+        q_rank = cute.rank(Q.layout)
+        k_rank = cute.rank(K.layout)
+        if cutlass.const_expr(q_rank == 5):
+            h_q = Q.shape[2] * Q.shape[3]
+        elif cutlass.const_expr(q_rank == 4):
+            h_q = Q.shape[2]
+        else:
+            h_q = Q.shape[1]
+        if cutlass.const_expr(k_rank == 5):
+            h_k = K.shape[2]
+        elif cutlass.const_expr(k_rank == 4):
+            h_k = K.shape[2]
+        else:
+            h_k = K.shape[1]
+        h_r = h_q // h_k
+        if cutlass.const_expr(cumulative_s_q is not None):
+            b = cumulative_s_q.shape[0] - 1
+        elif cutlass.const_expr(cumulative_s_k is not None):
+            b = cumulative_s_k.shape[0] - 1
+        else:
+            b = Q.shape[0]
+
+        Q, K, V, dQ, dK, dV, dO = [assume_tensor_aligned(t) for t in (Q, K, V, dQ, dK, dV, dO)]
+
+        Q = _as_bshkrd_tensor(Q, h_k, h_r, varlen)
+        K = _as_bshkrd_tensor(K, h_k, 1, varlen)
+        V = _as_bshkrd_tensor(V, h_k, 1, varlen)
+        dQ = _as_bshkrd_tensor(dQ, h_k, h_r, varlen)
+        dK = _as_bshkrd_tensor(dK, h_k, 1, varlen)
+        dV = _as_bshkrd_tensor(dV, h_k, 1, varlen)
+        dO = _as_bshkrd_tensor(dO, h_k, h_r, varlen)
+        scaled_LSE = _as_shhb_tensor(lse_log2, h_k, h_r, b, varlen)
+        sum_OdO = _as_shhb_tensor(dpsum, h_k, h_r, b, varlen)
+
+        # Keep original order: dQ first, then dKdV.
+        self.dq_kernel(
+            Q,
+            K,
+            V,
+            dQ,
+            dO,
+            scaled_LSE,
+            sum_OdO,
+            cumulative_s_q,
+            cumulative_s_k,
+            scale_softmax,
+            stream,
+        )
+        self.dkdv_kernel(
+            Q,
+            K,
+            V,
+            dK,
+            dV,
+            dO,
+            scaled_LSE,
+            sum_OdO,
+            cumulative_s_q,
+            cumulative_s_k,
+            scale_softmax,
+            stream,
+        )
@@ -0,0 +1,402 @@
+"""Search feasible SM90 fwd/bwd attention configs for given (head_dim, head_dim_v).
+
+Enumerates tile sizes, swap modes, atom layouts, and staging options.
+Checks GMMA divisibility, register budget, and shared memory budget.
+
+Usage:
+    python flash_attn/cute/sm90_config_search.py --headdim 128
+    python flash_attn/cute/sm90_config_search.py --mode fwd --headdim 192-128
+    python flash_attn/cute/sm90_config_search.py --mode bwd --headdim 192 --tile-n 64,96
+"""
+
+import math
+
+# H100 hardware limits
+SMEM_LIMIT = 224 * 1024  # 228 KB minus ~3 KB for LSE, dPsum, mbarriers
+REG_LIMITS = {2: 216, 3: 128}  # per-WG budget: 2WG=240-24, 3WG=160-32
+THREADS_PER_WG = 128
+
+
+def _divisors(n):
+    return [d for d in range(1, n + 1) if n % d == 0]
+
+
+def _acc_regs(M, N, num_wg):
+    """Accumulator registers per thread per WG."""
+    return M * N // (num_wg * THREADS_PER_WG)
+
+
+def _check_mma(M, N, num_wg, atom_layout_m, swap_AB):
+    """Check MMA feasibility. Returns regs per WG, or None if infeasible.
+
+    GMMA atom M=64. Swap exchanges (M, N) and atom layout.
+    Requires: M divisible by (atom_layout_m * 64), N by (atom_layout_n * 8).
+    """
+    if swap_AB:
+        M, N = N, M
+        atom_layout_m = num_wg // atom_layout_m
+    atom_layout_n = num_wg // atom_layout_m
+    if M % (atom_layout_m * 64) != 0 or N % (atom_layout_n * 8) != 0:
+        return None
+    return _acc_regs(M, N, num_wg)
+
+
+def _mma_traffic(M_eff, N_eff, K_red, num_wg, wg_n, is_rs=False):
+    """Total SMEM read traffic for one MMA (all WGs combined).
+
+    num_instr = (M_eff / 64) * wg_n instructions total.
+    Each reads A(64, K_red) and B(N_eff/wg_n, K_red) from smem (bf16).
+    """
+    num_instr = (M_eff // 64) * wg_n
+    A_per = 64 * K_red * 2 if not is_rs else 0
+    B_per = (N_eff // wg_n) * K_red * 2
+    return num_instr * (A_per + B_per)
+
+
+# ============================================================================
+# Backward
+# ============================================================================
+
+
+def _check_bwd_config(
+    hdim,
+    hdimv,
+    tile_m,
+    tile_n,
+    num_wg,
+    SdP_swapAB,
+    dKV_swapAB,
+    dQ_swapAB,
+    AtomLayoutMSdP,
+    AtomLayoutNdKV,
+    AtomLayoutMdQ,
+):
+    reg_limit = REG_LIMITS[num_wg]
+
+    # MMA feasibility
+    regs_SdP = _check_mma(tile_m, tile_n, num_wg, AtomLayoutMSdP, SdP_swapAB)
+    regs_dK = _check_mma(tile_n, hdim, num_wg, AtomLayoutNdKV, dKV_swapAB)
+    regs_dV = _check_mma(tile_n, hdimv, num_wg, AtomLayoutNdKV, dKV_swapAB)
+    regs_dQ = _check_mma(tile_m, hdim, num_wg, AtomLayoutMdQ, dQ_swapAB)
+    if any(r is None for r in (regs_SdP, regs_dK, regs_dV, regs_dQ)):
+        return None
+
+    # Peak regs: max(S+dP, dQ) + dK + dV
+    total_regs = max(2 * regs_SdP, regs_dQ) + regs_dK + regs_dV
+    if total_regs > reg_limit:
+        return None
+
+    # SMEM
+    mma_dkv_is_rs = (
+        AtomLayoutMSdP == 1 and AtomLayoutNdKV == num_wg and SdP_swapAB and not dKV_swapAB
+    )
+    Q_stage, PdS_stage = 2, 1
+
+    for dO_stage in (2, 1):
+        sQ = tile_m * hdim * 2 * Q_stage
+        sK = tile_n * hdim * 2
+        sV = tile_n * hdimv * 2
+        sdO = tile_m * hdimv * 2 * dO_stage
+        sPdS = tile_m * tile_n * 2 * PdS_stage
+        sP = sPdS if not mma_dkv_is_rs else 0
+        sdQaccum = tile_m * hdim * 4
+        smem = sQ + sK + sV + sdO + sP + sPdS + sdQaccum
+        if smem <= SMEM_LIMIT:
+            break
+    else:
+        return None
+
+    # SMEM traffic
+    def _swap(a, b, s):
+        return (b, a) if s else (a, b)
+
+    def _wg_n(al_m, s):
+        return al_m if s else num_wg // al_m
+
+    M_s, N_s = _swap(tile_m, tile_n, SdP_swapAB)
+    wn_SdP = _wg_n(AtomLayoutMSdP, SdP_swapAB)
+    traffic_S = _mma_traffic(M_s, N_s, hdim, num_wg, wn_SdP)
+    traffic_dP = _mma_traffic(M_s, N_s, hdimv, num_wg, wn_SdP)
+
+    wn_dKV = _wg_n(AtomLayoutNdKV, dKV_swapAB)
+    M_dv, N_dv = _swap(tile_n, hdimv, dKV_swapAB)
+    traffic_dV = _mma_traffic(M_dv, N_dv, tile_m, num_wg, wn_dKV, is_rs=mma_dkv_is_rs)
+    M_dk, N_dk = _swap(tile_n, hdim, dKV_swapAB)
+    traffic_dK = _mma_traffic(M_dk, N_dk, tile_m, num_wg, wn_dKV, is_rs=mma_dkv_is_rs)
+
+    M_dq, N_dq = _swap(tile_m, hdim, dQ_swapAB)
+    wn_dQ = _wg_n(AtomLayoutMdQ, dQ_swapAB)
+    traffic_dQ = _mma_traffic(M_dq, N_dq, tile_n, num_wg, wn_dQ)
+
+    traffic_P_store = tile_m * tile_n * 2 if not mma_dkv_is_rs else 0
+    traffic_dS_store = tile_m * tile_n * 2
+    traffic_dQ_smem = tile_m * hdim * 4 * 2  # store + TMA load
+
+    smem_traffic = (
+        traffic_S
+        + traffic_dP
+        + traffic_dV
+        + traffic_dK
+        + traffic_dQ
+        + traffic_P_store
+        + traffic_dS_store
+        + traffic_dQ_smem
+    )
+
+    return dict(
+        tile_m=tile_m,
+        tile_n=tile_n,
+        num_wg=num_wg,
+        Q_stage=Q_stage,
+        dO_stage=dO_stage,
+        PdS_stage=PdS_stage,
+        SdP_swapAB=SdP_swapAB,
+        dKV_swapAB=dKV_swapAB,
+        dQ_swapAB=dQ_swapAB,
+        AtomLayoutMSdP=AtomLayoutMSdP,
+        AtomLayoutNdKV=AtomLayoutNdKV,
+        AtomLayoutMdQ=AtomLayoutMdQ,
+        mma_dkv_is_rs=mma_dkv_is_rs,
+        regs_SdP=regs_SdP,
+        regs_dK=regs_dK,
+        regs_dV=regs_dV,
+        regs_dQ=regs_dQ,
+        total_regs=total_regs,
+        reg_limit=reg_limit,
+        smem_bytes=smem,
+        smem_kb=smem / 1024,
+        smem_traffic=smem_traffic,
+        smem_traffic_kb=smem_traffic / 1024,
+        smem_traffic_per_block=smem_traffic / (tile_m * tile_n),
+    )
+
+
+def find_feasible_bwd_configs(
+    head_dim,
+    head_dim_v=None,
+    tile_m_choices=(64, 80, 96, 112, 128),
+    tile_n_choices=(64, 80, 96, 112, 128),
+):
+    if head_dim_v is None:
+        head_dim_v = head_dim
+    hdim = int(math.ceil(head_dim / 32) * 32)
+    hdimv = int(math.ceil(head_dim_v / 32) * 32)
+
+    results = []
+    for num_wg in (2, 3):
+        divs = _divisors(num_wg)
+        for tile_m in tile_m_choices:
+            for tile_n in tile_n_choices:
+                for SdP_swap in (False, True):
+                    if (tile_n if SdP_swap else tile_m) % 64 != 0:
+                        continue
+                    for dKV_swap in (False, True):
+                        if not dKV_swap and tile_n % 64 != 0:
+                            continue
+                        if dKV_swap and (hdim % 64 != 0 or hdimv % 64 != 0):
+                            continue
+                        for dQ_swap in (False, True):
+                            if (hdim if dQ_swap else tile_m) % 64 != 0:
+                                continue
+                            for a1 in divs:
+                                for a2 in divs:
+                                    for a3 in divs:
+                                        cfg = _check_bwd_config(
+                                            hdim,
+                                            hdimv,
+                                            tile_m,
+                                            tile_n,
+                                            num_wg,
+                                            SdP_swap,
+                                            dKV_swap,
+                                            dQ_swap,
+                                            a1,
+                                            a2,
+                                            a3,
+                                        )
+                                        if cfg is not None:
+                                            results.append(cfg)
+
+    results.sort(key=lambda c: (-c["tile_n"], -c["tile_m"], c["smem_traffic_per_block"]))
+    return results
+
+
+def print_bwd_configs(configs, max_results=20):
+    if not configs:
+        print("No feasible configs found!")
+        return
+    n = min(len(configs), max_results)
+    print(f"Found {len(configs)} feasible configs (showing top {n}):\n")
+    hdr = (
+        f"{'wg':>2} {'tm':>3} {'tn':>3}  "
+        f"{'SdP':>3} {'dKV':>3} {'dQ':>3}  "
+        f"{'aSdP':>4} {'adKV':>4} {'adQ':>4}  "
+        f"{'Qs':>2} {'dOs':>3}  "
+        f"{'rS':>3} {'rdK':>3} {'rdV':>3} {'rdQ':>3} {'tot':>4}/{'':<3}  "
+        f"{'smem':>5}  {'traffic':>7}  {'tr/blk':>6}"
+    )
+    print(hdr)
+    print("-" * len(hdr))
+    B = lambda b: "T" if b else "F"
+    for c in configs[:max_results]:
+        print(
+            f"{c['num_wg']:>2} {c['tile_m']:>3} {c['tile_n']:>3}  "
+            f"{B(c['SdP_swapAB']):>3} {B(c['dKV_swapAB']):>3} {B(c['dQ_swapAB']):>3}  "
+            f"{c['AtomLayoutMSdP']:>4} {c['AtomLayoutNdKV']:>4} {c['AtomLayoutMdQ']:>4}  "
+            f"{c['Q_stage']:>2} {c['dO_stage']:>3}  "
+            f"{c['regs_SdP']:>3} {c['regs_dK']:>3} {c['regs_dV']:>3} {c['regs_dQ']:>3} "
+            f"{c['total_regs']:>4}/{c['reg_limit']:<3}  "
+            f"{c['smem_kb']:>4.0f}K  "
+            f"{c['smem_traffic_kb']:>6.0f}K  "
+            f"{c['smem_traffic_per_block']:>6.1f}"
+        )
+
+
+# ============================================================================
+# Forward
+# ============================================================================
+
+
+def _check_fwd_config(hdim, hdimv, tile_n, num_wg, pv_is_rs, overlap_wg):
+    reg_limit = REG_LIMITS[num_wg]
+    tile_m = num_wg * 64
+
+    if tile_n % 8 != 0:
+        return None
+
+    regs_S = _acc_regs(tile_m, tile_n, num_wg)
+    regs_O = _acc_regs(tile_m, hdimv, num_wg)
+    regs_P = regs_S // 2  # bf16 = half of f32
+
+    if overlap_wg:
+        total_regs = regs_S + regs_P + regs_O
+    else:
+        total_regs = regs_S + regs_O
+
+    if total_regs > reg_limit:
+        return None
+
+    # SMEM: 1 stage Q, 2 stages K/V, O overlaps Q, sP if not RS
+    sQ = tile_m * hdim * 2
+    sK = tile_n * hdim * 2 * 2
+    sV = tile_n * hdimv * 2 * 2
+    sO = tile_m * hdimv * 2
+    sP = tile_m * tile_n * 2 if not pv_is_rs else 0
+    smem = max(sQ, sO) + sK + sV + sP
+    if smem > SMEM_LIMIT:
+        return None
+
+    # SMEM traffic: num_instr = num_wg (all WGs in M, wg_n=1)
+    traffic_S = num_wg * (64 * hdim * 2 + tile_n * hdim * 2)
+    A_pv = 64 * tile_n * 2 if not pv_is_rs else 0
+    traffic_O = num_wg * (A_pv + hdimv * tile_n * 2)
+    traffic_P_store = tile_m * tile_n * 2 if not pv_is_rs else 0
+    smem_traffic = traffic_S + traffic_O + traffic_P_store
+
+    return dict(
+        tile_m=tile_m,
+        tile_n=tile_n,
+        num_wg=num_wg,
+        pv_is_rs=pv_is_rs,
+        overlap_wg=overlap_wg,
+        regs_S=regs_S,
+        regs_O=regs_O,
+        regs_P=regs_P,
+        total_regs=total_regs,
+        reg_limit=reg_limit,
+        smem_bytes=smem,
+        smem_kb=smem / 1024,
+        smem_traffic=smem_traffic,
+        smem_traffic_kb=smem_traffic / 1024,
+        smem_traffic_per_block=smem_traffic / (tile_m * tile_n),
+    )
+
+
+def find_feasible_fwd_configs(
+    head_dim, head_dim_v=None, tile_n_choices=(64, 80, 96, 112, 128, 144, 160, 176, 192)
+):
+    if head_dim_v is None:
+        head_dim_v = head_dim
+    hdim = int(math.ceil(head_dim / 32) * 32)
+    hdimv = int(math.ceil(head_dim_v / 32) * 32)
+
+    results = []
+    for num_wg in (2, 3):
+        for tile_n in tile_n_choices:
+            for pv_is_rs in (True, False):
+                for overlap_wg in (True, False):
+                    cfg = _check_fwd_config(hdim, hdimv, tile_n, num_wg, pv_is_rs, overlap_wg)
+                    if cfg is not None:
+                        results.append(cfg)
+
+    results.sort(key=lambda c: (-c["tile_n"], c["smem_traffic_per_block"]))
+    return results
+
+
+def print_fwd_configs(configs, max_results=20):
+    if not configs:
+        print("No feasible configs found!")
+        return
+    n = min(len(configs), max_results)
+    print(f"Found {len(configs)} feasible configs (showing top {n}):\n")
+    hdr = (
+        f"{'wg':>2} {'tm':>3} {'tn':>3}  "
+        f"{'RS':>2} {'olap':>4}  "
+        f"{'rS':>3} {'rP':>3} {'rO':>3} {'tot':>4}/{'':<3}  "
+        f"{'smem':>5}  {'traffic':>7}  {'tr/blk':>6}"
+    )
+    print(hdr)
+    print("-" * len(hdr))
+    B = lambda b: "T" if b else "F"
+    for c in configs[:max_results]:
+        print(
+            f"{c['num_wg']:>2} {c['tile_m']:>3} {c['tile_n']:>3}  "
+            f"{B(c['pv_is_rs']):>2} {B(c['overlap_wg']):>4}  "
+            f"{c['regs_S']:>3} {c['regs_P']:>3} {c['regs_O']:>3} "
+            f"{c['total_regs']:>4}/{c['reg_limit']:<3}  "
+            f"{c['smem_kb']:>4.0f}K  "
+            f"{c['smem_traffic_kb']:>6.0f}K  "
+            f"{c['smem_traffic_per_block']:>6.1f}"
+        )
+
+
+# ============================================================================
+# CLI
+# ============================================================================
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Search feasible SM90 MMA configs")
+    parser.add_argument("--mode", choices=["fwd", "bwd", "both"], default="both")
+    parser.add_argument(
+        "--headdim", type=str, default="128", help="Head dim, or hdim-hdimv (e.g. 192-128)"
+    )
+    parser.add_argument("--tile-m", type=str, default="64,80,96,112,128", help="Bwd tile_m choices")
+    parser.add_argument(
+        "--tile-n",
+        type=str,
+        default=None,
+        help="tile_n choices (default: fwd up to 192, bwd up to 128)",
+    )
+    parser.add_argument("-n", "--num-results", type=int, default=30)
+    args = parser.parse_args()
+
+    parts = args.headdim.split("-")
+    hdim = int(parts[0])
+    hdimv = int(parts[1]) if len(parts) > 1 else hdim
+
+    TN_FWD = "64,80,96,112,128,144,160,176,192"
+    TN_BWD = "64,80,96,112,128"
+
+    if args.mode in ("fwd", "both"):
+        tn = tuple(int(x) for x in (args.tile_n or TN_FWD).split(","))
+        print(f"=== FWD configs: hdim={hdim}, hdimv={hdimv} ===\n")
+        print_fwd_configs(find_feasible_fwd_configs(hdim, hdimv, tn), args.num_results)
+        print()
+
+    if args.mode in ("bwd", "both"):
+        tm = tuple(int(x) for x in args.tile_m.split(","))
+        tn = tuple(int(x) for x in (args.tile_n or TN_BWD).split(","))
+        print(f"=== BWD configs: hdim={hdim}, hdimv={hdimv} ===\n")
+        print_bwd_configs(find_feasible_bwd_configs(hdim, hdimv, tm, tn), args.num_results)
@@ -0,0 +1,740 @@
+# Copyright (c) 2025, Tri Dao.
+
+import math
+import operator
+from typing import Tuple
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Boolean
+
+from quack import layout_utils
+import flash_attn.cute.utils as utils
+from quack.cute_dsl_utils import ParamsBase
+from flash_attn.cute.seqlen_info import SeqlenInfoQK
+from flash_attn.cute.utils import AuxData
+
+
+@cute.jit
+def call_score_mod(
+    score_mod: cutlass.Constexpr,
+    score,
+    batch_idx,
+    head_idx,
+    q_idx,
+    kv_idx,
+    seqlen_info,
+    aux_data: AuxData,
+):
+    aux_tensors = aux_data.tensors if aux_data.tensors is not None else ()
+    # Compatibility shim for pre-aux_scalars score_mod callables.
+    if cutlass.const_expr(aux_data.scalars is not None):
+        return score_mod(
+            score,
+            batch_idx,
+            head_idx,
+            q_idx=q_idx,
+            kv_idx=kv_idx,
+            seqlen_info=seqlen_info,
+            aux_tensors=aux_tensors,
+            aux_scalars=aux_data.scalars,
+        )
+    return score_mod(
+        score,
+        batch_idx,
+        head_idx,
+        q_idx=q_idx,
+        kv_idx=kv_idx,
+        seqlen_info=seqlen_info,
+        aux_tensors=aux_tensors,
+    )
+
+
+@cute.jit
+def call_score_mod_bwd(
+    score_mod_bwd: cutlass.Constexpr,
+    grad,
+    score,
+    batch_idx,
+    head_idx,
+    q_idx,
+    kv_idx,
+    seqlen_info,
+    aux_data: AuxData,
+):
+    aux_tensors = aux_data.tensors if aux_data.tensors is not None else ()
+    # Compatibility shim for pre-aux_scalars score_mod_bwd callables.
+    if cutlass.const_expr(aux_data.scalars is not None):
+        return score_mod_bwd(
+            grad,
+            score,
+            batch_idx,
+            head_idx,
+            q_idx=q_idx,
+            kv_idx=kv_idx,
+            seqlen_info=seqlen_info,
+            aux_tensors=aux_tensors,
+            aux_scalars=aux_data.scalars,
+        )
+    return score_mod_bwd(
+        grad,
+        score,
+        batch_idx,
+        head_idx,
+        q_idx=q_idx,
+        kv_idx=kv_idx,
+        seqlen_info=seqlen_info,
+        aux_tensors=aux_tensors,
+    )
+
+
+@dataclass
+class Softmax(ParamsBase):
+    scale_log2: Float32
+    num_rows: cutlass.Constexpr[int]
+    row_max: cute.Tensor
+    row_sum: cute.Tensor
+    arch: cutlass.Constexpr[int] = 80
+    softmax_scale: Float32 | None = None
+
+    @staticmethod
+    def create(
+        scale_log2: Float32,
+        num_rows: cutlass.Constexpr[int],
+        arch: cutlass.Constexpr[int] = 80,
+        softmax_scale: Float32 | None = None,
+    ):
+        row_max = cute.make_rmem_tensor(num_rows, Float32)
+        row_sum = cute.make_rmem_tensor(num_rows, Float32)
+        return Softmax(scale_log2, num_rows, row_max, row_sum, arch, softmax_scale)
+
+    def reset(self) -> None:
+        self.row_max.fill(-Float32.inf)
+        self.row_sum.fill(0.0)
+
+    def _compute_row_max(
+        self, acc_S_row: cute.TensorSSA, init_val: float | Float32 | None = None
+    ) -> Float32:
+        return utils.fmax_reduce(acc_S_row, init_val, arch=self.arch)
+
+    def _compute_row_sum(
+        self, acc_S_row_exp: cute.TensorSSA, init_val: float | Float32 | None = None
+    ) -> Float32:
+        return utils.fadd_reduce(acc_S_row_exp, init_val, arch=self.arch)
+
+    @cute.jit
+    def online_softmax(
+        self,
+        acc_S: cute.Tensor,
+        is_first: cutlass.Constexpr[bool] = False,
+        check_inf: cutlass.Constexpr[bool] = True,
+    ) -> cute.Tensor:
+        """Apply online softmax and return the row_scale to rescale O.
+
+        :param acc_S: acc_S tensor
+        :type acc_S: cute.Tensor
+        :param is_first: is first n_block
+        :type is_first: cutlass.Constexpr
+        """
+        # Change acc_S to M,N layout view.
+        acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S)
+        row_scale = cute.make_fragment_like(self.row_max, Float32)
+
+        row_max = self.row_max
+        row_sum = self.row_sum
+        scale_log2 = self.scale_log2
+        arch = self.arch
+
+        # Each iteration processes one row of acc_S
+        for r in cutlass.range(cute.size(row_max), unroll_full=True):
+            acc_S_row = acc_S_mn[r, None].load()  # (n_block_size)
+
+            row_max_cur = utils.fmax_reduce(
+                acc_S_row,
+                init_val=row_max[r] if cutlass.const_expr(not is_first) else None,
+                arch=arch,
+            )
+
+            row_max_cur = cute.arch.warp_reduction_max(row_max_cur, threads_in_group=4)
+            # Update row_max before changing row_max_cur to safe value for -inf
+            row_max_prev = row_max[r]
+            row_max[r] = row_max_cur
+
+            if cutlass.const_expr(check_inf):
+                row_max_cur = 0.0 if row_max_cur == -Float32.inf else row_max_cur
+
+            if cutlass.const_expr(is_first):
+                row_max_cur_scaled = row_max_cur * scale_log2
+                acc_S_row_exp = cute.math.exp2(
+                    acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True
+                )
+                acc_S_row_sum = utils.fadd_reduce(acc_S_row_exp, init_val=None, arch=arch)
+                row_scale[r] = 1.0
+            else:
+                row_max_cur_scaled = row_max_cur * scale_log2
+                acc_S_row_exp = cute.math.exp2(
+                    acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True
+                )
+                # row_scale[r] = cute.math.exp2(row_max_prev * self.scale_log2 - row_max_cur_scaled)
+                row_scale[r] = cute.math.exp2(
+                    (row_max_prev - row_max_cur) * scale_log2, fastmath=True
+                )
+                acc_S_row_sum = utils.fadd_reduce(
+                    acc_S_row_exp, init_val=row_sum[r] * row_scale[r], arch=arch
+                )
+
+            row_sum[r] = acc_S_row_sum
+            acc_S_mn[r, None].store(acc_S_row_exp)
+
+        return row_scale
+
+    @cute.jit
+    def finalize(
+        self, final_scale: Float32 = 1.0, sink_val: Float32 | cute.Tensor | None = None
+    ) -> cute.Tensor:
+        """Finalize the online softmax by computing the scale and logsumexp."""
+        if cutlass.const_expr(sink_val is not None and isinstance(sink_val, cute.Tensor)):
+            assert cute.size(sink_val) == cute.size(self.row_sum)
+        row_sum = self.row_sum
+        row_max = self.row_max
+        scale_log2 = self.scale_log2
+
+        # quad reduction for row_sum as we didn't do it during each iteration of online softmax
+        row_sum.store(utils.warp_reduce(row_sum.load(), operator.add, width=4))
+        row_scale = cute.make_fragment_like(row_max, Float32)
+
+        for r in cutlass.range(cute.size(row_sum), unroll_full=True):
+            if cutlass.const_expr(sink_val is not None):
+                sink_val_cur = sink_val if not isinstance(sink_val, cute.Tensor) else sink_val[r]
+                LOG2_E = math.log2(math.e)
+                row_sum[r] += cute.math.exp2(
+                    sink_val_cur * LOG2_E - row_max[r] * scale_log2, fastmath=True
+                )
+
+            # if row_sum is zero or nan, set acc_O_mn_row to 1.0
+            acc_O_mn_row_is_zero_or_nan = row_sum[r] == 0.0 or row_sum[r] != row_sum[r]
+            row_scale[r] = (
+                cute.arch.rcp_approx(row_sum[r] if not acc_O_mn_row_is_zero_or_nan else 1.0)
+            ) * final_scale
+            row_sum_cur = row_sum[r]
+            LN2 = math.log(2.0)
+            row_sum[r] = (
+                (row_max[r] * scale_log2 + cute.math.log2(row_sum_cur, fastmath=True)) * LN2
+                if not acc_O_mn_row_is_zero_or_nan
+                else -Float32.inf
+            )
+        return row_scale
+
+    @cute.jit
+    def rescale_O(self, acc_O: cute.Tensor, row_scale: cute.Tensor) -> None:
+        """Scale each row of acc_O by the given scale tensor.
+        :param acc_O: input tensor
+        :type acc_O: cute.Tensor
+        :param row_scale: row_scale tensor
+        :type row_scale: cute.Tensor
+        """
+        acc_O_mn = layout_utils.reshape_acc_to_mn(acc_O)
+        assert cute.size(row_scale) == cute.size(acc_O_mn, mode=[0])
+        for r in cutlass.range(cute.size(row_scale), unroll_full=True):
+            acc_O_mn[r, None].store(acc_O_mn[r, None].load() * row_scale[r])
+
+
+@dataclass
+class SoftmaxSm100(Softmax):
+    rescale_threshold: cutlass.Constexpr[float] = 0.0
+    max_offset: cutlass.Constexpr[int] = 0
+    row_pair_stride: cutlass.Constexpr[int] = 0
+
+    @staticmethod
+    def create(
+        scale_log2: Float32,
+        rescale_threshold: cutlass.Constexpr[float] = 0.0,
+        softmax_scale: Float32 | None = None,
+        max_offset: cutlass.Constexpr[int] = 0,
+        row_pair_stride: cutlass.Constexpr[int] = 0,
+    ):
+        num_rows = 1
+        arch = 100
+        row_max = cute.make_rmem_tensor(num_rows, Float32)
+        row_sum = cute.make_rmem_tensor(num_rows, Float32)
+        return SoftmaxSm100(
+            scale_log2,
+            num_rows,
+            row_max,
+            row_sum,
+            arch,
+            softmax_scale,
+            rescale_threshold=rescale_threshold,
+            max_offset=max_offset,
+            row_pair_stride=row_pair_stride,
+        )
+
+    @cute.jit
+    def compute_row_max_local(self, acc_S_row: cute.TensorSSA, is_first: Boolean) -> Float32:
+        if cutlass.const_expr(is_first):
+            row_max_new = self._compute_row_max(acc_S_row)
+        else:
+            row_max_old = self.row_max[0]
+            row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old)
+        return row_max_new
+
+    @cute.jit
+    def update_row_max_from_local(
+        self,
+        row_max_new: Float32,
+        is_first: Boolean,
+    ) -> Tuple[Float32, Float32]:
+        if cutlass.const_expr(self.row_pair_stride != 0):
+            row_max_new = utils.fmax(
+                row_max_new,
+                cute.arch.shuffle_sync_bfly(
+                    row_max_new,
+                    offset=self.row_pair_stride,
+                    mask=-1,
+                    mask_and_clamp=31,
+                ),
+            )
+        if cutlass.const_expr(is_first):
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale = 0.0
+        else:
+            row_max_old = self.row_max[0]
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2
+            acc_scale = cute.math.exp2(acc_scale_)
+            if cutlass.const_expr(self.rescale_threshold > 0.0):
+                if acc_scale_ >= -self.rescale_threshold:
+                    row_max_new = row_max_old
+                    row_max_safe = row_max_old
+                    acc_scale = 1.0
+        self.row_max[0] = row_max_new
+        return row_max_safe, acc_scale
+
+    @cute.jit
+    def update_row_max(self, acc_S_row: cute.TensorSSA, is_first: int) -> Tuple[Float32, Float32]:
+        if cutlass.const_expr(is_first):
+            row_max_new = self._compute_row_max(acc_S_row)
+            if cutlass.const_expr(self.row_pair_stride != 0):
+                row_max_new = utils.fmax(
+                    row_max_new,
+                    cute.arch.shuffle_sync_bfly(
+                        row_max_new,
+                        offset=self.row_pair_stride,
+                        mask=-1,
+                        mask_and_clamp=31,
+                    ),
+                )
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale = 0.0
+        else:
+            row_max_old = self.row_max[0]
+            row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old)
+            if cutlass.const_expr(self.row_pair_stride != 0):
+                row_max_new = utils.fmax(
+                    row_max_new,
+                    cute.arch.shuffle_sync_bfly(
+                        row_max_new,
+                        offset=self.row_pair_stride,
+                        mask=-1,
+                        mask_and_clamp=31,
+                    ),
+                )
+            row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
+            acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2
+            acc_scale = cute.math.exp2(acc_scale_, fastmath=True)
+            if cutlass.const_expr(self.rescale_threshold > 0.0):
+                if acc_scale_ >= -self.rescale_threshold:
+                    row_max_new = row_max_old
+                    row_max_safe = row_max_old
+                    acc_scale = 1.0
+        self.row_max[0] = row_max_new
+        return row_max_safe, acc_scale
+
+    def update_row_sum(
+        self, acc_S_row_exp: cute.TensorSSA, row_scale: Float32, is_first: int = False
+    ) -> None:
+        if cutlass.const_expr(self.row_pair_stride != 0):
+            local_sum = self._compute_row_sum(acc_S_row_exp)
+            row_sum_new = local_sum + cute.arch.shuffle_sync_bfly(
+                local_sum,
+                offset=self.row_pair_stride,
+                mask=-1,
+                mask_and_clamp=31,
+            )
+            if cutlass.const_expr(not is_first):
+                row_sum_new = row_sum_new + self.row_sum[0] * row_scale
+            self.row_sum[0] = row_sum_new
+        else:
+            init_val = self.row_sum[0] * row_scale if cutlass.const_expr(not is_first) else None
+            self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=init_val)
+
+    @cute.jit
+    def scale_subtract_rowmax(
+        self,
+        acc_S_row: cute.Tensor,
+        row_max: Float32,
+    ):
+        assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements"
+        row_max_scaled = row_max * self.scale_log2
+        max_offset = Float32(self.max_offset)
+        bias = max_offset - row_max_scaled
+        for i in cutlass.range(0, cute.size(acc_S_row.shape), 2, unroll_full=True):
+            acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2(
+                (acc_S_row[i], acc_S_row[i + 1]),
+                (self.scale_log2, self.scale_log2),
+                (bias, bias),
+            )
+
+    @cute.jit
+    def apply_exp2_convert(
+        self,
+        acc_S_row: cute.Tensor,
+        acc_S_row_converted: cute.Tensor,
+        ex2_emu_freq: cutlass.Constexpr[int] = 0,
+        ex2_emu_res: cutlass.Constexpr[int] = 4,
+        ex2_emu_start_frg: cutlass.Constexpr[int] = 0,
+    ):
+        assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements"
+        frg_tile = 32
+        assert frg_tile % 2 == 0
+        frg_cnt = cute.size(acc_S_row) // frg_tile
+        assert cute.size(acc_S_row) % frg_tile == 0
+        acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile))
+        acc_S_row_converted_frg = cute.logical_divide(
+            acc_S_row_converted, cute.make_layout(frg_tile)
+        )
+        for j in cutlass.range_constexpr(frg_cnt):
+            for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2):
+                # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
+                # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True)
+                if cutlass.const_expr(ex2_emu_freq == 0):
+                    acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
+                    acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True)
+                else:
+                    if cutlass.const_expr(
+                        k % ex2_emu_freq < ex2_emu_freq - ex2_emu_res
+                        or j >= frg_cnt - 1
+                        or j < ex2_emu_start_frg
+                    ):
+                        acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
+                        acc_S_row_frg[k + 1, j] = cute.math.exp2(
+                            acc_S_row_frg[k + 1, j], fastmath=True
+                        )
+                    else:
+                        # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.e2e_asm2(acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j])
+                        acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.ex2_emulation_2(
+                            acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]
+                        )
+            acc_S_row_converted_frg[None, j].store(
+                acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type)
+            )
+
+    @cute.jit
+    def scale_apply_exp2_convert(
+        self,
+        acc_S_row: cute.Tensor,
+        row_max: Float32,
+        acc_S_row_converted: cute.Tensor,
+    ):
+        assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements"
+        minus_row_max_scaled = -row_max * self.scale_log2
+        for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2):
+            acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2(
+                (acc_S_row[i], acc_S_row[i + 1]),
+                (self.scale_log2, self.scale_log2),
+                (minus_row_max_scaled, minus_row_max_scaled),
+            )
+
+        # for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2):
+        #     acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2(
+        #         (acc_S_row[i], acc_S_row[i + 1]),
+        #         (self.scale_log2, self.scale_log2),
+        #         (minus_row_max_scaled, minus_row_max_scaled),
+        #     )
+        #     acc_S_row[i] = cute.math.exp2(acc_S_row[i], fastmath=True)
+        #     acc_S_row[i + 1] = cute.math.exp2(acc_S_row[i + 1], fastmath=True)
+
+        frg_tile = 32
+        assert frg_tile % 2 == 0
+        frg_cnt = cute.size(acc_S_row) // frg_tile
+        assert cute.size(acc_S_row) % frg_tile == 0
+        acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile))
+        acc_S_row_converted_frg = cute.logical_divide(
+            acc_S_row_converted, cute.make_layout(frg_tile)
+        )
+        for j in cutlass.range_constexpr(frg_cnt):
+            for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2):
+                # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = (
+                #     cute.arch.fma_packed_f32x2(
+                #         (acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]),
+                #         (self.scale_log2, self.scale_log2),
+                #         (minus_row_max_scaled, minus_row_max_scaled),
+                #     )
+                # )
+                # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
+                # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True)
+                acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
+                acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True)
+            acc_S_row_converted_frg[None, j].store(
+                acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type)
+            )
+
+
+@cute.jit
+def floor_if_packed(
+    q_idx,
+    qhead_per_kvhead: cutlass.Constexpr[int],
+) -> cute.Tensor:
+    """Convert q_idx to packed format for Pack-GQA."""
+    if cutlass.const_expr(qhead_per_kvhead == 1):
+        return q_idx
+    return q_idx // qhead_per_kvhead
+
+
+@cute.jit
+def apply_score_mod_inner(
+    score_tensor,
+    index_tensor,
+    score_mod: cutlass.Constexpr,
+    batch_idx,
+    head_idx,
+    softmax_scale,
+    vec_size: cutlass.Constexpr,
+    qk_acc_dtype: cutlass.Constexpr,
+    aux_data: AuxData,
+    fastdiv_mods,
+    seqlen_info: SeqlenInfoQK,
+    constant_q_idx: cutlass.Constexpr,
+    qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+    transpose_indices: cutlass.Constexpr[bool] = False,
+):
+    """Shared implementation for applying score modification.
+
+    Args:
+        score_tensor: The scores to modify (acc_S for flash_fwd, tSrS_t2r for sm100)
+        index_tensor: Index positions (tScS for flash_fwd, tScS_t2r for sm100)
+        score_mod: The score modification function to apply
+        batch_idx: Batch index
+        head_idx: Head index
+        softmax_scale: Scale to apply
+        vec_size: Vector size for processing elements
+        qk_acc_dtype: Data type for accumulator
+        aux_tensors: Optional aux_tensors for FlexAttention
+        aux_scalars: Optional runtime scalar captures for FlexAttention
+        fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping
+        seqlen_info: Sequence length info
+        constant_q_idx: If provided, use this constant for all q_idx values
+                        If None, compute q_idx per-element
+        qhead_per_kvhead_packgqa: Pack-GQA replication factor. Divide q_idx by this
+                                  when greater than 1 so score mods see logical heads.
+        transpose_indices: If True, swap q_idx/kv_idx in index_tensor (for bwd kernel where S is transposed)
+    """
+    # Index positions in the index_tensor tuple
+    # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx
+    # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx
+    if cutlass.const_expr(transpose_indices):
+        q_idx_pos = cutlass.const_expr(1)
+        kv_idx_pos = cutlass.const_expr(0)
+    else:
+        q_idx_pos = cutlass.const_expr(0)
+        kv_idx_pos = cutlass.const_expr(1)
+
+    n_vals = cutlass.const_expr(cute.size(score_tensor.shape))
+    score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype)
+    kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    # SSA values for batch (constant across all elements)
+    batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,))
+
+    # Handle q_idx based on whether it's constant
+    q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    # For Pack-GQA with non-constant q_idx, we need per-element head indices
+    # since a thread may process multiple query head indices
+    if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+        head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    for i in cutlass.range(0, n_vals, vec_size, unroll_full=True):
+        for j in cutlass.range(vec_size, unroll_full=True):
+            score_vec[j] = score_tensor[i + j] * softmax_scale
+
+            # Extract head offset from packed q_idx for Pack-GQA
+            if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+                q_idx_packed = index_tensor[i + j][q_idx_pos]
+                # Building up the logical q_head idx: final_q_head = kv_head * qhead_per_kvhead + (q_physical % qhead_per_kvhead)
+                q_idx_logical = q_idx_packed // qhead_per_kvhead
+                head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead
+                head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset
+
+            # If we will do loads we mod, in order to not read OOB
+            if cutlass.const_expr(aux_data.tensors is not None and fastdiv_mods is not None):
+                if cutlass.const_expr(constant_q_idx is None):
+                    seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
+                    q_idx_floored = floor_if_packed(
+                        index_tensor[i + j][q_idx_pos], qhead_per_kvhead
+                    )
+                    _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod)
+                    q_idx_vec[j] = q_idx_wrapped
+                else:
+                    _, seqlen_k_divmod = fastdiv_mods
+
+                _, kv_idx_wrapped = divmod(index_tensor[i + j][kv_idx_pos], seqlen_k_divmod)
+                kv_idx_vec[j] = kv_idx_wrapped
+            else:
+                # No bounds checking - direct indexing
+                if constant_q_idx is None:
+                    q_idx_vec[j] = floor_if_packed(index_tensor[i + j][q_idx_pos], qhead_per_kvhead)
+                kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos]
+
+        # Convert to SSA for score_mod call
+        score_ssa = score_vec.load()
+        kv_idx_ssa = kv_idx_vec.load()
+        if cutlass.const_expr(constant_q_idx is None):
+            q_idx_ssa = q_idx_vec.load()
+        else:
+            # NB we do not apply Pack-GQA division here, as constant_q_idx is assumed to already be logical
+            q_idx_const = constant_q_idx
+            q_idx_ssa = utils.scalar_to_ssa(q_idx_const, cutlass.Int32).broadcast_to((vec_size,))
+
+        # Compute head_idx_ssa: per-element for Pack-GQA with non-constant q_idx, constant otherwise
+        if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+            head_idx_ssa = head_idx_vec.load()
+        else:
+            head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,))
+
+        post_mod_scores = call_score_mod(
+            score_mod,
+            score_ssa,
+            batch_idx_ssa,
+            head_idx_ssa,
+            q_idx_ssa,
+            kv_idx_ssa,
+            seqlen_info,
+            aux_data,
+        )
+
+        # Write back modified scores
+        score_vec.store(post_mod_scores)
+        for j in cutlass.range(vec_size, unroll_full=True):
+            score_tensor[i + j] = score_vec[j]
+
+
+@cute.jit
+def apply_score_mod_bwd_inner(
+    grad_tensor,
+    score_tensor,
+    index_tensor,
+    score_mod_bwd: cutlass.Constexpr,
+    batch_idx,
+    head_idx,
+    softmax_scale,
+    vec_size: cutlass.Constexpr,
+    qk_acc_dtype: cutlass.Constexpr,
+    aux_data: AuxData,
+    fastdiv_mods,
+    seqlen_info,
+    constant_q_idx: cutlass.Constexpr,
+    qhead_per_kvhead: cutlass.Constexpr[int] = 1,
+    transpose_indices: cutlass.Constexpr[bool] = False,
+):
+    """Apply backward score modification (joint graph).
+
+    Args:
+        grad_tensor: in/out: dlogits rewritten in-place with d(scaled_scores)
+        score_tensor: pre-mod scores (unscaled QK tile), scaled by softmax_scale internally
+        index_tensor: Index positions (same as forward)
+        score_mod_bwd: The backward score modification function (joint graph)
+        batch_idx: Batch index
+        head_idx: Head index
+        softmax_scale: Scale to apply to score_tensor
+        vec_size: Vector size for processing elements
+        qk_acc_dtype: Data type for accumulator
+        aux_tensors: Optional aux_tensors for FlexAttention
+        aux_scalars: Optional runtime scalar captures for FlexAttention
+        fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping
+        seqlen_info: Sequence length info
+        constant_q_idx: If provided, use this constant for all q_idx values
+        qhead_per_kvhead: Pack-GQA replication factor
+        transpose_indices: If True, swap q_idx/kv_idx in index_tensor
+    """
+    # Index positions in the index_tensor tuple
+    # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx
+    # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx
+    if cutlass.const_expr(transpose_indices):
+        q_idx_pos = cutlass.const_expr(1)
+        kv_idx_pos = cutlass.const_expr(0)
+    else:
+        q_idx_pos = cutlass.const_expr(0)
+        kv_idx_pos = cutlass.const_expr(1)
+    n_vals = cutlass.const_expr(cute.size(grad_tensor.shape))
+    grad_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype)
+    score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype)
+    kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+    batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,))
+    q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    # For Pack-GQA with non-constant q_idx, we need per-element head indices
+    if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+        head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32)
+
+    for i in cutlass.range(0, n_vals, vec_size, unroll_full=True):
+        for j in cutlass.range(vec_size, unroll_full=True):
+            grad_vec[j] = grad_tensor[i + j]
+            # Scale score so joint graph sees same value as forward score_mod
+            score_vec[j] = score_tensor[i + j] * softmax_scale
+
+            if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+                q_idx_packed = index_tensor[i + j][q_idx_pos]
+                q_idx_logical = q_idx_packed // qhead_per_kvhead
+                head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead
+                head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset
+
+            if cutlass.const_expr(aux_data.tensors is not None and fastdiv_mods is not None):
+                if cutlass.const_expr(constant_q_idx is None):
+                    seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods
+                    q_idx_floored = floor_if_packed(
+                        index_tensor[i + j][q_idx_pos], qhead_per_kvhead
+                    )
+                    _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod)
+                    q_idx_vec[j] = q_idx_wrapped
+                else:
+                    _, seqlen_k_divmod = fastdiv_mods
+
+                _, kv_idx_wrapped = divmod(index_tensor[i + j][kv_idx_pos], seqlen_k_divmod)
+                kv_idx_vec[j] = kv_idx_wrapped
+            else:
+                # No bounds checking - direct indexing
+                if constant_q_idx is None:
+                    q_idx_vec[j] = floor_if_packed(index_tensor[i + j][q_idx_pos], qhead_per_kvhead)
+                kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos]
+
+        grad_ssa = grad_vec.load()
+        score_ssa = score_vec.load()
+        kv_idx_ssa = kv_idx_vec.load()
+
+        if cutlass.const_expr(constant_q_idx is None):
+            q_idx_ssa = q_idx_vec.load()
+        else:
+            q_idx_ssa = utils.scalar_to_ssa(constant_q_idx, cutlass.Int32).broadcast_to((vec_size,))
+
+        if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None):
+            head_idx_ssa = head_idx_vec.load()
+        else:
+            head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,))
+
+        grad_out_ssa = call_score_mod_bwd(
+            score_mod_bwd,
+            grad_ssa,
+            score_ssa,
+            batch_idx_ssa,
+            head_idx_ssa,
+            q_idx_ssa,
+            kv_idx_ssa,
+            seqlen_info,
+            aux_data,
+        )
+
+        grad_vec.store(grad_out_ssa)
+        for j in cutlass.range(vec_size, unroll_full=True):
+            grad_tensor[i + j] = grad_vec[j]
@@ -0,0 +1,487 @@
+import math
+from contextlib import nullcontext
+from functools import wraps
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch._guards import active_fake_mode
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+
+class IndexFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, indices):
+        ctx.save_for_backward(indices)
+        assert input.ndim >= 2
+        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+        second_dim = other_shape.numel()
+        return torch.gather(
+            rearrange(input, "b ... -> b (...)"),
+            0,
+            repeat(indices, "z -> z d", d=second_dim),
+        ).reshape(-1, *other_shape)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (indices,) = ctx.saved_tensors
+        assert grad_output.ndim >= 2
+        other_shape = grad_output.shape[1:]
+        grad_output = rearrange(grad_output, "b ... -> b (...)")
+        grad_input = torch.zeros(
+            [ctx.first_axis_dim, grad_output.shape[1]],
+            device=grad_output.device,
+            dtype=grad_output.dtype,
+        )
+        grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+index_first_axis = IndexFirstAxis.apply
+
+
+class IndexPutFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, values, indices, first_axis_dim):
+        ctx.save_for_backward(indices)
+        assert indices.ndim == 1
+        assert values.ndim >= 2
+        output = torch.zeros(
+            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
+        )
+        output[indices] = values
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (indices,) = ctx.saved_tensors
+        grad_values = grad_output[indices]
+        return grad_values, None, None
+
+
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+def unpad_input(hidden_states, attention_mask, unused_mask=None):
+    all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask
+    seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32)
+    used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    in_fake_mode = active_fake_mode() is not None
+    if not in_fake_mode:
+        indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten()
+        max_seqlen_in_batch = seqlens_in_batch.max().item()
+    else:
+        # torch.nonzero and .item() are not supported in FakeTensorMode
+        batch_size, seqlen = attention_mask.shape
+        indices = torch.arange(batch_size * seqlen, device=hidden_states.device)
+        max_seqlen_in_batch = seqlen
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices),
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+        used_seqlens_in_batch,
+    )
+
+
+def pad_input(hidden_states, indices, batch, seqlen):
+    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)
+
+
+def generate_random_padding_mask(
+    max_seqlen, batch_size, device, mode="random", zero_lengths=False, min_seqlen=None
+):
+    assert mode in ["full", "random", "third"]
+    min_seqlen = min_seqlen if min_seqlen is not None else 0 if zero_lengths else 1
+    if mode == "full":
+        lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32)
+    elif mode == "random":
+        lengths = torch.randint(
+            max(min_seqlen, max_seqlen - 20),
+            max_seqlen + 1,
+            (batch_size, 1),
+            device=device,
+        )
+    else:
+        lengths = torch.randint(
+            max(min_seqlen, max_seqlen // 3),
+            max_seqlen + 1,
+            (batch_size, 1),
+            device=device,
+        )
+
+    if zero_lengths:
+        for i in range(batch_size):
+            if i % 5 == 0:
+                lengths[i] = 0
+        lengths[-1] = 0
+    padding_mask = (
+        repeat(torch.arange(max_seqlen, device=device), "s -> b s", b=batch_size) < lengths
+    )
+    return padding_mask
+
+
+def generate_qkv(
+    q,
+    k,
+    v,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    qv=None,
+    kvpacked=False,
+    qkvpacked=False,
+    query_unused_mask=None,
+    key_unused_mask=None,
+):
+    assert not (kvpacked and qkvpacked)
+    batch_size, seqlen_q, nheads, d = q.shape
+    d_v = v.shape[-1]
+    _, seqlen_k, nheads_k, _ = k.shape
+    assert k.shape == (batch_size, seqlen_k, nheads_k, d)
+    assert v.shape == (batch_size, seqlen_k, nheads_k, d_v)
+    if query_unused_mask is not None or key_unused_mask is not None:
+        assert not kvpacked
+        assert not qkvpacked
+
+    if query_padding_mask is not None:
+        q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, seqused_q = unpad_input(
+            q, query_padding_mask, query_unused_mask
+        )
+        output_pad_fn = lambda output_unpad: pad_input(
+            output_unpad, indices_q, batch_size, seqlen_q
+        )
+        qv_unpad = rearrange(qv, "b s ... -> (b s) ...")[indices_q] if qv is not None else None
+    else:
+        q_unpad = rearrange(q, "b s h d -> (b s) h d")
+        cu_seqlens_q = torch.arange(
+            0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q_unpad.device
+        )
+        seqused_q = None
+        max_seqlen_q = seqlen_q
+        output_pad_fn = lambda output_unpad: rearrange(
+            output_unpad, "(b s) h d -> b s h d", b=batch_size
+        )
+        qv_unpad = rearrange(qv, "b s ... -> (b s) ...") if qv is not None else None
+
+    if key_padding_mask is not None:
+        k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, seqused_k = unpad_input(
+            k, key_padding_mask, key_unused_mask
+        )
+        v_unpad, *_ = unpad_input(v, key_padding_mask, key_unused_mask)
+    else:
+        k_unpad = rearrange(k, "b s h d -> (b s) h d")
+        v_unpad = rearrange(v, "b s h d -> (b s) h d")
+        cu_seqlens_k = torch.arange(
+            0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, device=k_unpad.device
+        )
+        seqused_k = None
+        max_seqlen_k = seqlen_k
+
+    if qkvpacked:
+        assert (query_padding_mask == key_padding_mask).all()
+        assert nheads == nheads_k
+        qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)
+        qkv = torch.stack([q, k, v], dim=2)
+        if query_padding_mask is not None:
+            dqkv_pad_fn = lambda dqkv_unpad: pad_input(dqkv_unpad, indices_q, batch_size, seqlen_q)
+        else:
+            dqkv_pad_fn = lambda dqkv_unpad: rearrange(
+                dqkv_unpad, "(b s) t h d -> b s t h d", b=batch_size
+            )
+        return (
+            qkv_unpad.detach().requires_grad_(),
+            cu_seqlens_q,
+            max_seqlen_q,
+            qkv.detach().requires_grad_(),
+            output_pad_fn,
+            dqkv_pad_fn,
+        )
+    elif kvpacked:
+        kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)
+        kv = torch.stack([k, v], dim=2)
+        dq_pad_fn = output_pad_fn
+        if key_padding_mask is not None:
+            dkv_pad_fn = lambda dkv_unpad: pad_input(dkv_unpad, indices_k, batch_size, seqlen_k)
+        else:
+            dkv_pad_fn = lambda dkv_unpad: rearrange(
+                dkv_unpad, "(b s) t h d -> b s t h d", b=batch_size
+            )
+        return (
+            q_unpad.detach().requires_grad_(),
+            kv_unpad.detach().requires_grad_(),
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            q.detach().requires_grad_(),
+            kv.detach().requires_grad_(),
+            output_pad_fn,
+            dq_pad_fn,
+            dkv_pad_fn,
+        )
+    else:
+        dq_pad_fn = output_pad_fn
+        if key_padding_mask is not None:
+            dk_pad_fn = lambda dk_unpad: pad_input(dk_unpad, indices_k, batch_size, seqlen_k)
+        else:
+            dk_pad_fn = lambda dk_unpad: rearrange(dk_unpad, "(b s) h d -> b s h d", b=batch_size)
+        return (
+            q_unpad.detach().requires_grad_(),
+            k_unpad.detach().requires_grad_(),
+            v_unpad.detach().requires_grad_(),
+            qv_unpad.detach() if qv is not None else None,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            q.detach().requires_grad_(),
+            k.detach().requires_grad_(),
+            v.detach().requires_grad_(),
+            qv.detach() if qv is not None else None,
+            output_pad_fn,
+            dq_pad_fn,
+            dk_pad_fn,
+        )
+
+
+def construct_local_mask(
+    seqlen_q,
+    seqlen_k,
+    window_size=(None, None),
+    sink_token_length=0,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    device=None,
+):
+    row_idx = rearrange(torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1")
+    col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)
+    if key_leftpad is not None:
+        key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1")
+        col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0])
+        col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32)
+    sk = (
+        seqlen_k
+        if key_padding_mask is None
+        else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sq = (
+        seqlen_q
+        if query_padding_mask is None
+        else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    if window_size[0] is None:
+        return col_idx > row_idx + sk - sq + window_size[1]
+    else:
+        sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
+        if window_size[1] is None:
+            local_mask_left = col_idx > sk
+        else:
+            local_mask_left = col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk)
+        return torch.logical_or(
+            local_mask_left,
+            torch.logical_and(
+                col_idx < row_idx + sk - sq - window_size[0], col_idx >= sink_token_length
+            ),
+        )
+
+
+def construct_chunk_mask(
+    seqlen_q,
+    seqlen_k,
+    attention_chunk,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    device=None,
+):
+    row_idx = rearrange(torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1")
+    col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long)
+    if key_leftpad is not None:
+        key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1")
+        col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0])
+        col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32)
+    sk = (
+        seqlen_k
+        if key_padding_mask is None
+        else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sq = (
+        seqlen_q
+        if query_padding_mask is None
+        else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1")
+    )
+    sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
+    col_limit_left_chunk = row_idx + sk - sq - (row_idx + sk - sq) % attention_chunk
+    return torch.logical_or(
+        col_idx < col_limit_left_chunk, col_idx >= col_limit_left_chunk + attention_chunk
+    )
+
+
+def attention_ref(
+    q,
+    k,
+    v,
+    query_padding_mask=None,
+    key_padding_mask=None,
+    key_leftpad=None,
+    attn_bias=None,
+    dropout_p=0.0,
+    dropout_mask=None,
+    causal=False,
+    qv=None,
+    q_descale=None,
+    k_descale=None,
+    v_descale=None,
+    window_size=(None, None),
+    attention_chunk=0,
+    sink_token_length=0,
+    learnable_sink: Optional[torch.Tensor] = None,
+    softcap=0.0,
+    upcast=True,
+    reorder_ops=False,
+    intermediate_dtype=None,
+    return_lse=False,
+    gather_kv_indices=None,
+):
+    assert v is not None
+    has_qk = q is not None and k is not None
+    assert has_qk or qv is not None
+    if causal:
+        window_size = (window_size[0], 0)
+    dtype_og = v.dtype
+    q_shape = q.shape if q is not None else qv.shape
+    if upcast:
+        q, k, v, qv = [t.float() if t is not None else None for t in (q, k, v, qv)]
+    if q_descale is not None:
+        q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q_shape[2] // v.shape[2])
+        q, qv = [(t.float() * q_descale).to(t.dtype) if t is not None else None for t in (q, qv)]
+    if k_descale is not None:
+        k = (k.float() * rearrange(k_descale, "b h -> b 1 h 1")).to(dtype=k.dtype)
+    if v_descale is not None:
+        v = (v.float() * rearrange(v_descale, "b h -> b 1 h 1")).to(dtype=v.dtype)
+    seqlen_q, seqlen_k = q_shape[1], v.shape[1]
+    k, v = [
+        repeat(t, "b s h d -> b s (h g) d", g=q_shape[2] // t.shape[2]) if t is not None else None
+        for t in (k, v)
+    ]
+    d = q_shape[-1]  # == dv for qv
+    dv = v.shape[-1]
+    softmax_scale = 1.0 / math.sqrt(d if qv is None or q is None else d + dv)
+    if has_qk:
+        scores = torch.einsum(
+            "bthd,bshd->bhts",
+            q if reorder_ops else q * softmax_scale,
+            k * softmax_scale if reorder_ops else k,
+        )
+    if qv is not None:
+        qv_scores = torch.einsum(
+            "bthd,bshd->bhts",
+            qv if reorder_ops else qv * softmax_scale,
+            v * softmax_scale if reorder_ops else v,
+        )
+        scores = qv_scores if not has_qk else scores + qv_scores
+    if softcap > 0:
+        scores = torch.tanh(scores / softcap) * softcap
+    if key_padding_mask is not None:
+        scores.masked_fill_(rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf"))
+    local_mask = None
+    if window_size[0] is not None or window_size[1] is not None:
+        local_mask = construct_local_mask(
+            seqlen_q,
+            seqlen_k,
+            window_size,
+            sink_token_length,
+            query_padding_mask,
+            key_padding_mask,
+            key_leftpad=key_leftpad,
+            device=v.device,
+        )
+    if attention_chunk > 0:
+        chunk_mask = construct_chunk_mask(
+            seqlen_q,
+            seqlen_k,
+            attention_chunk,
+            query_padding_mask,
+            key_padding_mask,
+            key_leftpad=key_leftpad,
+            device=v.device,
+        )
+        local_mask = (
+            torch.logical_or(local_mask, chunk_mask) if local_mask is not None else chunk_mask
+        )
+    if gather_kv_indices is not None:
+        batch = q_shape[0]
+        topk_len = gather_kv_indices.shape[2]
+        if topk_len < seqlen_k:
+            topk_index_mask = torch.full(
+                (batch, seqlen_q, seqlen_k), False, device="cuda"
+            ).scatter_(-1, gather_kv_indices, True)
+            scores.masked_fill_(rearrange(~topk_index_mask, "b t s -> b 1 t s"), float("-inf"))
+    if local_mask is not None:
+        scores.masked_fill_(local_mask, float("-inf"))
+    if attn_bias is not None:
+        scores = scores + attn_bias
+    # After all masks are applied, before softmax:
+    # scores shape: [b, h, t, s]
+    lse = torch.logsumexp(scores, dim=-1)  # [b, h, t]
+    if learnable_sink is None:
+        attention = torch.softmax(scores, dim=-1).to(v.dtype)
+    else:
+        scores_fp32 = scores.to(torch.float32)
+        logits_max = torch.amax(scores_fp32, dim=-1, keepdim=True)
+        learnable_sink = rearrange(learnable_sink, "h -> h 1 1")
+        logits_or_sinks_max = torch.maximum(learnable_sink, logits_max)
+        unnormalized_scores = torch.exp(scores_fp32 - logits_or_sinks_max)
+        normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + torch.exp(
+            learnable_sink - logits_or_sinks_max
+        )
+        # LSE with sink: log(Z) = log(normalizer) + max
+        lse = (torch.log(normalizer.squeeze(-1)) + logits_or_sinks_max.squeeze(-1)).to(dtype_og)
+        attention = (unnormalized_scores / normalizer).to(v.dtype)
+    if query_padding_mask is not None:
+        attention = attention.masked_fill(rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0)
+    if key_padding_mask is not None:
+        attention = attention.masked_fill(rearrange(~key_padding_mask, "b s -> b 1 1 s"), 0.0)
+    if local_mask is not None:
+        attention = attention.masked_fill(torch.all(local_mask, dim=-1, keepdim=True), 0.0)
+    dropout_scaling = 1.0 / (1 - dropout_p)
+    if dropout_mask is not None:
+        attention_drop = attention.masked_fill(~dropout_mask, 0.0)
+    else:
+        attention_drop = attention
+    if intermediate_dtype is not None:
+        attention_drop = attention_drop.to(intermediate_dtype).to(attention_drop.dtype)
+    output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling)
+    if query_padding_mask is not None:
+        output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0)
+    if return_lse:
+        return output.to(dtype_og), attention.to(dtype_og), lse.to(dtype_og)
+    return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)
+
+
+def maybe_fake_tensor_mode(fake: bool = True):
+    """
+    One way to populate/pre-compile cache is to use torch fake tensor mode,
+    which does not allocate actual GPU tensors but retains tensor shape/dtype
+    metadata for cute.compile.
+    """
+
+    def decorator(fn):
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            with FakeTensorMode() if fake else nullcontext():
+                return fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def is_fake_mode() -> bool:
+    return active_fake_mode() is not None
@@ -0,0 +1,278 @@
+from typing import Type, Optional
+from dataclasses import dataclass
+import operator
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync
+from cutlass import Int32, Uint32, const_expr, Boolean
+
+from flash_attn.cute import utils
+from flash_attn.cute.utils import warp_reduce
+from quack.cute_dsl_utils import ParamsBase
+
+import math
+
+
+@dataclass
+class CpasyncGatherKVManager(ParamsBase):
+    mIndexTopk: cute.Tensor
+    sBitmask: Optional[cute.Tensor]
+
+    cta_rank_in_cluster: Int32
+    thread_idx: Int32
+    warp_idx: Int32
+
+    topk_length: Int32
+    seqlen_k_limit: Int32
+    tile_n: Int32
+    num_threads: cutlass.Constexpr[Int32]
+    hdim: cutlass.Constexpr[Int32]
+    hdim_v: cutlass.Constexpr[Int32]
+    num_hdimv_splits: cutlass.Constexpr[Int32]
+    cta_group_size: cutlass.Constexpr[Int32]
+
+    gmem_threads_per_row: cutlass.Constexpr[Int32]
+    topk_indices_per_thread: Int32
+    async_copy_elems: Int32
+
+    gmem_tiled_copy_KV: cute.TiledCopy
+    gmem_thr_copy_KV: cute.TiledCopy
+
+    rTopk: cute.Tensor
+    rTopkHalf: cute.Tensor
+    # for bitmask
+    rTopk_NonInterleaved: cute.Tensor
+
+    pipeline_bitmask: Optional[pipeline.PipelineAsync]
+    cpasync_barrier: Optional[pipeline.NamedBarrier]
+
+    disable_bitmask: cutlass.Constexpr[Boolean]
+
+    @staticmethod
+    def create(
+        mIndexTopk: cute.Tensor,
+        cta_rank_in_cluster: Int32,
+        thread_idx: Int32,
+        warp_idx: Int32,
+        topk_length: Int32,
+        seqlen_k_limit: Int32,
+        tile_n: cutlass.Constexpr[Int32],
+        hdim: cutlass.Constexpr[Int32],
+        hdim_v: cutlass.Constexpr[Int32],
+        num_hdimv_splits: cutlass.Constexpr[Int32],
+        num_threads: cutlass.Constexpr[Int32],
+        dtype: Type[cutlass.Numeric],
+        cta_group_size: cutlass.Constexpr[Int32],
+        cpasync_barrier: Optional[pipeline.NamedBarrier] = None,
+        disable_bitmask: cutlass.Constexpr[Boolean] = False,
+        sBitmask: Optional[cute.Tensor] = None,
+        pipeline_bitmask: Optional[pipeline.PipelineAsync] = None,
+    ):
+        assert tile_n % num_threads == 0
+        assert num_threads == 128
+        assert hdim % 64 == 0
+        assert (hdim_v // num_hdimv_splits // cta_group_size) % 64 == 0
+        assert num_threads % cute.arch.WARP_SIZE == 0
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // dtype.width
+        dtype_bytes = dtype.width // 8
+        # assumes hdim is never part of transposed operand
+        gmem_k_block_size = math.gcd(
+            hdim,
+            hdim_v // num_hdimv_splits // cta_group_size,
+            128 // dtype_bytes,
+        )
+        assert gmem_k_block_size % async_copy_elems == 0
+        gmem_threads_per_row = gmem_k_block_size // async_copy_elems
+        assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0
+        atom_async_copy = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            dtype,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        thr_layout = cute.make_ordered_layout(
+            (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        val_layout = cute.make_layout((1, async_copy_elems))
+        gmem_tiled_copy_KV = cute.make_tiled_copy_tv(atom_async_copy, thr_layout, val_layout)
+        gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx)
+        topk_indices_per_thread = tile_n // num_threads
+
+        rTopk = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+        rTopkHalf = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+        rTopk_NonInterleaved = cute.make_rmem_tensor((topk_indices_per_thread,), Int32)
+
+        return CpasyncGatherKVManager(
+            mIndexTopk,
+            sBitmask,
+            cta_rank_in_cluster,
+            thread_idx,
+            warp_idx,
+            topk_length,
+            seqlen_k_limit,
+            tile_n,
+            num_threads,
+            hdim,
+            hdim_v,
+            num_hdimv_splits,
+            cta_group_size,
+            gmem_threads_per_row,
+            topk_indices_per_thread,
+            async_copy_elems,
+            gmem_tiled_copy_KV,
+            gmem_thr_copy_KV,
+            rTopk,
+            rTopkHalf,
+            rTopk_NonInterleaved,
+            pipeline_bitmask,
+            cpasync_barrier,
+            disable_bitmask,
+        )
+
+    @cute.jit
+    def load_index_topk(
+        self,
+        n_block: Int32,
+        transpose: bool,
+    ):
+        entries_per_thread = self.topk_indices_per_thread
+        rTopk = self.rTopk if const_expr(transpose) else self.rTopkHalf
+
+        for i in cutlass.range_constexpr(entries_per_thread):
+            row = (
+                i * self.num_threads
+                + (self.thread_idx % self.gmem_threads_per_row)
+                * (self.num_threads // self.gmem_threads_per_row)
+                + (self.thread_idx // self.gmem_threads_per_row)
+            )
+            # need this if not offset in load_X
+            # if const_expr(not transpose):
+            #     row += self.cta_rank_in_cluster * (self.tile_n//self.cta_group_size)
+            #     row = row % self.tile_n
+            row_idx = n_block * self.tile_n + row
+            rTopk[i] = self.mIndexTopk[row_idx]
+
+            if const_expr(not transpose and not self.disable_bitmask):
+                row_non_interleaved = i * self.num_threads + self.thread_idx
+                row_idx_non_interleaved = n_block * self.tile_n + row_non_interleaved
+                self.rTopk_NonInterleaved[0] = self.mIndexTopk[row_idx_non_interleaved]
+
+    @cute.jit
+    def compute_bitmask(
+        self,
+        producer_state_bitmask,
+    ):
+        assert self.pipeline_bitmask is not None, "pipeline_bitmask not provided"
+        assert self.cpasync_barrier is not None, "cpasync barrier not provided"
+
+        lane_idx = cute.arch.lane_idx()
+        assert cute.size(self.rTopk_NonInterleaved) == 1
+        bitmask = Uint32(0)
+
+        # Step 1. Construct per-thread bitmask
+        topk_idx = self.rTopk_NonInterleaved[0]
+        is_valid = topk_idx >= 0 and topk_idx < self.seqlen_k_limit
+        if is_valid:
+            bitmask = Uint32(1 << lane_idx)
+
+        # Step 2. Warp shuffle bitwise OR = add since indices are exclusive.
+        bitmask = warp_reduce(bitmask, operator.add)
+
+        self.pipeline_bitmask.producer_acquire(producer_state_bitmask)
+        # store to smem and sync threads
+        if lane_idx == 0:
+            self.sBitmask[self.warp_idx, producer_state_bitmask.index] = bitmask
+        self.cpasync_barrier.arrive_and_wait()
+
+        self.pipeline_bitmask.producer_commit(producer_state_bitmask)
+        producer_state_bitmask.advance()
+        return producer_state_bitmask
+
+    @cute.jit
+    def compute_X_ptr(
+        self,
+        mX: cute.Tensor,
+        transpose: bool,
+        d_offset: int = 0,
+    ):
+        entries_per_thread = self.topk_indices_per_thread
+        tPrXPtr = cute.make_rmem_tensor((entries_per_thread,), cutlass.Int64)
+        tPrRowValid = cute.make_rmem_tensor((entries_per_thread,), cutlass.Int32)
+        rTopk = self.rTopk if const_expr(transpose) else self.rTopkHalf
+
+        for i in cutlass.range_constexpr(entries_per_thread):
+            topk_idx = rTopk[i]
+            if const_expr(not self.disable_bitmask):
+                row_valid = topk_idx >= 0 and topk_idx < self.seqlen_k_limit
+                tPrRowValid[i] = row_valid
+            if const_expr(not transpose):
+                tPrXPtr[i] = utils.elem_pointer(mX, (topk_idx, d_offset)).toint()
+            else:
+                tPrXPtr[i] = utils.elem_pointer(mX, (d_offset, topk_idx)).toint()
+
+        return tPrXPtr, tPrRowValid
+
+    @cute.jit
+    def load_X(
+        self,
+        mX: cute.Tensor,
+        sX: cute.Tensor,
+        transpose: bool,
+        K_or_V: str,
+        d_offset: int = 0,
+    ):
+        assert K_or_V in ("K", "V")
+        cta_tile_n = self.tile_n if const_expr(transpose) else self.tile_n // self.cta_group_size
+        head_dim = self.hdim if const_expr(K_or_V == "K") else self.hdim_v // self.num_hdimv_splits
+        if const_expr(transpose):
+            head_dim = head_dim // self.cta_group_size
+        order = (1, 0) if const_expr(transpose) else (0, 1)
+
+        sX_nd_layout = cute.make_ordered_layout((cta_tile_n, head_dim), order=order)
+        sX_nd = cute.composition(sX, sX_nd_layout)
+
+        cX = cute.make_identity_tensor((cta_tile_n, head_dim))
+        tXsX = self.gmem_thr_copy_KV.partition_D(sX_nd)
+        tXcX = self.gmem_thr_copy_KV.partition_S(cX)
+
+        tPrXPtr, tPrRowValid = self.compute_X_ptr(mX, transpose, d_offset)
+
+        if const_expr(not transpose):
+            offset = self.cta_rank_in_cluster * (self.gmem_threads_per_row // self.cta_group_size)
+        else:
+            offset = 0
+
+        for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
+            if const_expr(not self.disable_bitmask):
+                row_valid = utils.shuffle_sync(
+                    tPrRowValid[m // self.gmem_threads_per_row],
+                    (m + offset) % self.gmem_threads_per_row,
+                    width=self.gmem_threads_per_row,
+                )
+                should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], Boolean)
+                should_load.fill(Boolean(row_valid))
+            x_ptr_i64 = utils.shuffle_sync(
+                tPrXPtr[m // self.gmem_threads_per_row],
+                (m + offset) % self.gmem_threads_per_row,
+                width=self.gmem_threads_per_row,
+            )
+            x_gmem_ptr = cute.make_ptr(
+                mX.element_type, x_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            mX_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
+            mX_cur_copy = cute.tiled_divide(mX_cur, (self.async_copy_elems,))
+
+            for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
+                ki = tXcX[0, 0, k][1] // self.async_copy_elems
+                mX_cur_copy_ki = mX_cur_copy[None, ki]
+                tXsX_k = tXsX[None, m, k]
+                mX_cur_copy_ki = cute.make_tensor(mX_cur_copy_ki.iterator, tXsX_k.layout)
+                cute.copy(
+                    self.gmem_tiled_copy_KV,
+                    mX_cur_copy_ki,
+                    tXsX_k,
+                    pred=should_load if const_expr(not self.disable_bitmask) else None,
+                )
@@ -0,0 +1,957 @@
+# Copyright (c) 2025, Tri Dao.
+
+import math
+import hashlib
+import inspect
+import os
+from functools import partial
+from typing import Type, Callable, Optional, Tuple, overload, NamedTuple
+
+import cutlass
+import cutlass.cute as cute
+
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute import FastDivmodDivisor
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import nvvm, llvm
+from cutlass.cute.runtime import from_dlpack
+
+
+import quack.activation
+
+_MIXER_ATTRS = ("__vec_size__",)
+
+
+class AuxData(NamedTuple):
+    tensors: tuple | list | None = None
+    scalars: tuple | None = None
+
+
+# Obtained from sollya:
+# fpminimax(exp(x * log(2.0)), 1, [|1,24...|],[0;1],relative);
+POLY_EX2 = {
+    0: (1.0),
+    1: (
+        1.0,
+        0.922497093677520751953125,
+    ),
+    2: (
+        1.0,
+        0.6657850742340087890625,
+        0.330107033252716064453125,
+    ),
+    3: (
+        1.0,
+        0.695146143436431884765625,
+        0.227564394474029541015625,
+        0.077119089663028717041015625,
+    ),
+    4: (
+        1.0,
+        0.693042695522308349609375,
+        0.2412912547588348388671875,
+        5.2225358784198760986328125e-2,
+        1.3434938155114650726318359375e-2,
+    ),
+    5: (
+        1.0,
+        0.693151414394378662109375,
+        0.24016360938549041748046875,
+        5.5802188813686370849609375e-2,
+        9.01452265679836273193359375e-3,
+        1.86810153536498546600341796875e-3,
+    ),
+}
+
+_fa_clc_enabled: bool = os.environ.get("FA_CLC", "0") == "1"
+_fa_disable_2cta_enabled: bool = os.environ.get("FA_DISABLE_2CTA", "0") == "1"
+
+
+def _is_cuda_12() -> bool:
+    """Check if the CUDA toolkit version is 12.x.
+
+    2CTA forward non-causal has a codegen regression on CUDA 12 that causes
+    ~18% slowdown compared to 1CTA. This is fixed in CUDA 13.x.
+    """
+    try:
+        import torch
+
+        cuda_version = torch.version.cuda
+        if cuda_version is not None:
+            major = cuda_version.split(".")[0]
+            return int(major) == 12
+    except Exception:
+        pass
+    return False
+
+
+_fa_disable_2cta_cuda12: bool = _is_cuda_12()
+
+
+def _get_use_clc_scheduler_default() -> bool:
+    return _fa_clc_enabled
+
+
+def _get_disable_2cta_default(is_fwd: bool = False) -> bool:
+    if is_fwd:
+        return _fa_disable_2cta_enabled or _fa_disable_2cta_cuda12
+    else:
+        return _fa_disable_2cta_enabled
+
+
+def _compute_base_hash(func: Callable) -> str:
+    """Compute hash from source code or bytecode and closure values."""
+    try:
+        data = inspect.getsource(func).encode()
+    except (OSError, TypeError):
+        if hasattr(func, "__code__") and func.__code__ is not None:
+            data = func.__code__.co_code
+        else:
+            data = repr(func).encode()
+
+    hasher = hashlib.sha256(data)
+
+    if hasattr(func, "__closure__") and func.__closure__ is not None:
+        for cell in func.__closure__:
+            hasher.update(repr(cell.cell_contents).encode())
+
+    return hasher.hexdigest()
+
+
+def hash_callable(
+    func: Callable, mixer_attrs: Tuple[str] = _MIXER_ATTRS, set_cute_hash: bool = True
+) -> str:
+    """Hash a callable based on the source code or bytecode and closure values.
+    Fast-path: if the callable (or its __wrapped__ base) has a ``__cute_hash__``
+    attribute, that value is returned immediately as the base hash, then
+    metadata dunders are mixed in to produce the final dict-key hash.
+    set_cute_hash: whether or not to set func.__cute_hash__
+    """
+    # Resolve base hash
+    if hasattr(func, "__cute_hash__"):
+        base_hash = func.__cute_hash__
+    else:
+        # Unwrap decorated functions (e.g., cute.jit wrappers).
+        base_func = getattr(func, "__wrapped__", func)
+
+        if hasattr(base_func, "__cute_hash__"):
+            base_hash = base_func.__cute_hash__
+        else:
+            base_hash = _compute_base_hash(base_func)
+
+            if set_cute_hash:
+                base_func.__cute_hash__ = base_hash
+
+    # Mix in mutable metadata dunders
+    mixer_values = tuple(getattr(func, attr, None) for attr in mixer_attrs)
+
+    if all(v is None for v in mixer_values):
+        return base_hash
+
+    hasher = hashlib.sha256(base_hash.encode())
+
+    for attr, val in zip(mixer_attrs, mixer_values):
+        hasher.update(f"{attr}={val!r}".encode())
+
+    return hasher.hexdigest()
+
+
+def create_softcap_scoremod(softcap_val):
+    @cute.jit
+    def scoremod_premask_fn(
+        acc_S_SSA, batch_idx, head_idx, q_idx, kv_idx, seqlen_info, aux_tensors
+    ):
+        scores = acc_S_SSA / softcap_val
+        return softcap_val * cute.math.tanh(scores, fastmath=True)
+
+    return scoremod_premask_fn
+
+
+def create_softcap_scoremod_bwd(softcap_val):
+    @cute.jit
+    def scoremod_bwd_fn(
+        grad_out_SSA, score_SSA, batch_idx, head_idx, q_idx, kv_idx, seqlen_info, aux_tensors
+    ):
+        scores = score_SSA / softcap_val
+        tanh_scores = cute.math.tanh(scores, fastmath=True)
+        return grad_out_SSA * (1.0 - tanh_scores * tanh_scores)
+
+    return scoremod_bwd_fn
+
+
+LOG2_E = math.log2(math.e)
+
+
+def compute_softmax_scale_log2(softmax_scale, score_mod):
+    """Compute softmax_scale_log2 and adjusted softmax_scale based on whether score_mod is used.
+
+    When score_mod is None, fold the log2(e) factor into softmax_scale_log2 and set softmax_scale
+    to None. When score_mod is present, keep softmax_scale separate so it can be applied before
+    the score_mod, and set softmax_scale_log2 to just the change-of-base constant.
+
+    Returns (softmax_scale_log2, softmax_scale).
+    """
+    if const_expr(score_mod is None):
+        return softmax_scale * LOG2_E, None
+    else:
+        return LOG2_E, softmax_scale
+
+
+def compute_fastdiv_mods(mQ, mK, qhead_per_kvhead, pack_gqa, aux_tensors, mPageTable=None):
+    """Compute FastDivmodDivisor pairs for aux_tensors index computation.
+
+    Returns a (seqlen_q_divmod, seqlen_k_divmod) tuple, or None if aux_tensors is None.
+    """
+    if const_expr(aux_tensors is None):
+        return None
+    seqlen_q = cute.size(mQ.shape[0]) // (qhead_per_kvhead if const_expr(pack_gqa) else 1)
+    seqlen_k = (
+        cute.size(mK.shape[0])
+        if const_expr(mPageTable is None)
+        else mK.shape[0] * mPageTable.shape[1]
+    )
+    return (FastDivmodDivisor(seqlen_q), FastDivmodDivisor(seqlen_k))
+
+
+def convert_from_dlpack(x, leading_dim, alignment=16, divisibility=1) -> cute.Tensor:
+    return (
+        from_dlpack(x, assumed_align=alignment)
+        .mark_layout_dynamic(leading_dim=leading_dim)
+        .mark_compact_shape_dynamic(
+            mode=leading_dim, stride_order=x.dim_order(), divisibility=divisibility
+        )
+    )
+
+
+def convert_from_dlpack_compact_dynamic(
+    x,
+    *,
+    dynamic_modes: tuple[int, ...],
+    alignment: int = 16,
+    stride_order=None,
+    divisibility: int = 1,
+    enable_tvm_ffi: bool = False,
+) -> cute.Tensor:
+    """Convert via DLPack and mark selected compact dimensions as dynamic."""
+    if isinstance(dynamic_modes, int):
+        dynamic_modes = (dynamic_modes,)
+    if stride_order is None:
+        stride_order = x.dim_order()
+    t = (
+        from_dlpack(x, assumed_align=alignment, enable_tvm_ffi=True)
+        if enable_tvm_ffi
+        else from_dlpack(x, assumed_align=alignment)
+    )
+    for mode in dynamic_modes:
+        t = t.mark_compact_shape_dynamic(
+            mode=mode,
+            stride_order=stride_order,
+            divisibility=divisibility,
+        )
+    return t
+
+
+def convert_from_dlpack_leading_static(
+    x, leading_dim, alignment=16, static_modes=None, stride_order=None
+) -> cute.Tensor:
+    if stride_order is None:
+        stride_order = x.dim_order()
+    x_ = from_dlpack(x, assumed_align=alignment)
+    for i in range(x.ndim):
+        if i != leading_dim and (static_modes is None or i not in static_modes):
+            x_ = x_.mark_compact_shape_dynamic(mode=i, stride_order=stride_order)
+    return x_
+
+
+def make_tiled_copy_A(
+    copy_atom: cute.CopyAtom, tiled_mma: cute.TiledMma, swapAB: cutlass.Constexpr[bool] = False
+) -> cute.TiledCopy:
+    if const_expr(swapAB):
+        return cute.make_tiled_copy_B(copy_atom, tiled_mma)
+    else:
+        return cute.make_tiled_copy_A(copy_atom, tiled_mma)
+
+
+def make_tiled_copy_B(
+    copy_atom: cute.CopyAtom, tiled_mma: cute.TiledMma, swapAB: cutlass.Constexpr[bool] = False
+) -> cute.TiledCopy:
+    if const_expr(swapAB):
+        return cute.make_tiled_copy_A(copy_atom, tiled_mma)
+    else:
+        return cute.make_tiled_copy_B(copy_atom, tiled_mma)
+
+
+def mma_make_fragment_A(
+    smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False
+) -> cute.Tensor:
+    if const_expr(swapAB):
+        return mma_make_fragment_B(smem, thr_mma)
+    else:
+        return thr_mma.make_fragment_A(thr_mma.partition_A(smem))
+
+
+def mma_make_fragment_B(
+    smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False
+) -> cute.Tensor:
+    if const_expr(swapAB):
+        return mma_make_fragment_A(smem, thr_mma)
+    else:
+        return thr_mma.make_fragment_B(thr_mma.partition_B(smem))
+
+
+def get_smem_store_atom(
+    arch: cutlass.Constexpr[int], element_type: Type[cute.Numeric], transpose: bool = False
+) -> cute.CopyAtom:
+    if const_expr(arch < 90 or element_type.width != 16):
+        return cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            element_type,
+            num_bits_per_copy=2 * element_type.width,
+        )
+    else:
+        return cute.make_copy_atom(
+            cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=transpose, num_matrices=4),
+            element_type,
+        )
+
+
+@cute.jit
+def warp_reduce(
+    val: cute.TensorSSA | cute.Numeric,
+    op: Callable,
+    width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE,
+) -> cute.TensorSSA | cute.Numeric:
+    if const_expr(isinstance(val, cute.TensorSSA)):
+        res = cute.make_rmem_tensor(val.shape, val.dtype)
+        res.store(val)
+        for i in cutlass.range_constexpr(cute.size(val.shape)):
+            res[i] = warp_reduce(res[i], op, width)
+        return res.load()
+    else:
+        for i in cutlass.range_constexpr(int(math.log2(width))):
+            val = op(val, cute.arch.shuffle_sync_bfly(val, offset=1 << i))
+    return val
+
+
+@dsl_user_op
+def smid(*, loc=None, ip=None) -> Int32:
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [],
+            "mov.u32 $0, %smid;",
+            "=r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def fmax(
+    a: float | Float32, b: float | Float32, c: float | Float32 | None = None, *, loc=None, ip=None
+) -> Float32:
+    return Float32(
+        nvvm.fmax(
+            Float32(a).ir_value(loc=loc, ip=ip),
+            Float32(b).ir_value(loc=loc, ip=ip),
+            c=Float32(c).ir_value(loc=loc, ip=ip) if c is not None else None,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
+@cute.jit
+def fmax_reduce(
+    x: cute.TensorSSA, init_val: float | Float32 | None = None, arch: cutlass.Constexpr[int] = 80
+) -> Float32:
+    if const_expr(arch < 100 or cute.size(x.shape) % 8 != 0):
+        # if const_expr(init_val is None):
+        #     init_val = -cutlass.Float32.if
+        # return x.reduce(cute.ReductionOp.MAX, init_val, 0)
+        res = cute.make_rmem_tensor(x.shape, Float32)
+        res.store(x)
+        # local_max = [res[0], res[1]]
+        # for i in cutlass.range_constexpr(2, cute.size(x.shape), 2):
+        #     local_max[0] = fmax(local_max[0], res[i + 0])
+        #     local_max[1] = fmax(local_max[1], res[i + 1])
+        # local_max[0] = fmax(local_max[0], local_max[1])
+        # return local_max[0] if const_expr(init_val is None) else fmax(local_max[0], init_val)
+        local_max = [res[0], res[1], res[2], res[3]]
+        for i in cutlass.range_constexpr(4, cute.size(x.shape), 4):
+            local_max[0] = fmax(local_max[0], res[i + 0])
+            local_max[1] = fmax(local_max[1], res[i + 1])
+            local_max[2] = fmax(local_max[2], res[i + 2])
+            local_max[3] = fmax(local_max[3], res[i + 3])
+        local_max[0] = fmax(local_max[0], local_max[1])
+        local_max[2] = fmax(local_max[2], local_max[3])
+        local_max[0] = fmax(local_max[0], local_max[2])
+        return local_max[0] if const_expr(init_val is None) else fmax(local_max[0], init_val)
+    else:
+        # [2025-06-15] x.reduce only seems to use 50% 3-input max and 50% 2-input max
+        # We instead force the 3-input max.
+        res = cute.make_rmem_tensor(x.shape, Float32)
+        res.store(x)
+        local_max_0 = (
+            fmax(init_val, res[0], res[1])
+            if const_expr(init_val is not None)
+            else fmax(res[0], res[1])
+        )
+        local_max = [
+            local_max_0,
+            fmax(res[2], res[3]),
+            fmax(res[4], res[5]),
+            fmax(res[6], res[7]),
+        ]
+        for i in cutlass.range_constexpr(8, cute.size(x.shape), 8):
+            local_max[0] = fmax(local_max[0], res[i], res[i + 1])
+            local_max[1] = fmax(local_max[1], res[i + 2], res[i + 3])
+            local_max[2] = fmax(local_max[2], res[i + 4], res[i + 5])
+            local_max[3] = fmax(local_max[3], res[i + 6], res[i + 7])
+        local_max[0] = fmax(local_max[0], local_max[1])
+        return fmax(local_max[0], local_max[2], local_max[3])
+
+
+@cute.jit
+def fadd_reduce(
+    x: cute.TensorSSA, init_val: float | Float32 | None = None, arch: cutlass.Constexpr[int] = 80
+) -> Float32:
+    if const_expr(arch < 100 or cute.size(x.shape) % 8 != 0):
+        if const_expr(init_val is None):
+            init_val = Float32.zero
+        return x.reduce(cute.ReductionOp.ADD, init_val, 0)
+        # res = cute.make_rmem_tensor(x.shape, Float32)
+        # res.store(x)
+        # local_sum = [res[0], res[1], res[2], res[3]]
+        # for i in cutlass.range_constexpr(4, cute.size(x.shape), 4):
+        #     local_sum[0] += res[i + 0]
+        #     local_sum[1] += res[i + 1]
+        #     local_sum[2] += res[i + 2]
+        #     local_sum[3] += res[i + 3]
+        # local_sum[0] += local_sum[1]
+        # local_sum[2] += local_sum[3]
+        # local_sum[0] += local_sum[2]
+        # return local_sum[0] if const_expr(init_val is None) else local_sum[0] + init_val
+    else:
+        res = cute.make_rmem_tensor(x.shape, Float32)
+        res.store(x)
+        local_sum_0 = (
+            cute.arch.add_packed_f32x2((init_val, 0.0), (res[0], res[1]))
+            # cute.arch.add_packed_f32x2((init_val / 2, init_val / 2), (res[0], res[1]))
+            if const_expr(init_val is not None)
+            else (res[0], res[1])
+        )
+        local_sum = [local_sum_0, (res[2], res[3]), (res[4], res[5]), (res[6], res[7])]
+        for i in cutlass.range_constexpr(8, cute.size(x.shape), 8):
+            local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], (res[i + 0], res[i + 1]))
+            local_sum[1] = cute.arch.add_packed_f32x2(local_sum[1], (res[i + 2], res[i + 3]))
+            local_sum[2] = cute.arch.add_packed_f32x2(local_sum[2], (res[i + 4], res[i + 5]))
+            local_sum[3] = cute.arch.add_packed_f32x2(local_sum[3], (res[i + 6], res[i + 7]))
+        local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], local_sum[1])
+        local_sum[2] = cute.arch.add_packed_f32x2(local_sum[2], local_sum[3])
+        local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], local_sum[2])
+        return local_sum[0][0] + local_sum[0][1]
+
+
+@dsl_user_op
+def atomic_add_fp32(a: float | Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> None:
+    # gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value()
+    # # cache_hint = cutlass.Int64(0x12F0000000000000)
+    # llvm.inline_asm(
+    #     None,
+    #     [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip)],
+    #     # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()],
+    #     "red.global.add.f32 [$0], $1;",
+    #     # "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;",
+    #     # "red.global.add.L2::cache_hint.f32 [$0], $1, $2;",
+    #     "l,f",
+    #     # "l,f,l",
+    #     has_side_effects=True,
+    #     is_align_stack=False,
+    #     asm_dialect=llvm.AsmDialect.AD_ATT,
+    # )
+    nvvm.atomicrmw(
+        res=T.f32(), op=nvvm.AtomicOpKind.FADD, ptr=gmem_ptr.llvm_ptr, a=Float32(a).ir_value()
+    )
+
+
+@dsl_user_op
+def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
+    return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
+
+
+@cute.jit
+def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor:
+    # Only compute predicates for the "k" dimension. For the mn dimension, we will use "if"
+    tApA = cute.make_rmem_tensor(
+        cute.make_layout(
+            (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])),
+            stride=(cute.size(tAcA, mode=[2]), 0, 1),
+        ),
+        cutlass.Boolean,
+    )
+    for rest_v in cutlass.range_constexpr(tApA.shape[0]):
+        for rest_k in cutlass.range_constexpr(tApA.shape[2]):
+            tApA[rest_v, 0, rest_k] = cute.elem_less(tAcA[(0, rest_v), 0, rest_k][1], limit)
+    return tApA
+
+
+def canonical_warp_group_idx(sync: bool = True) -> cutlass.Int32:
+    warp_group_idx = cute.arch.thread_idx()[0] // 128
+    if const_expr(sync):
+        warp_group_idx = cute.arch.make_warp_uniform(warp_group_idx)
+    return warp_group_idx
+
+
+# @dsl_user_op
+# def warp_vote_any_lt(a: float | Float32, b: float | Float32, *, loc=None, ip=None) -> cutlass.Boolean:
+#     mask = cutlass.Int32(-1)
+#     return cutlass.Boolean(
+#         llvm.inline_asm(
+#             T.i32(),
+#             [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip), mask.ir_value(loc=loc, ip=ip)],
+#             ".pred p1, p2;\n"
+#             "setp.lt.f32 p1, $1, $2;\n"
+#             "vote.sync.any.pred p2, p1, $3;\n"
+#             "selp.u32 $0, 1, 0, p2;",
+#             # "selp.u32 $0, 1, 0, p1;",
+#             "=r,f,f,r",
+#             has_side_effects=False,
+#             is_align_stack=False,
+#             asm_dialect=llvm.AsmDialect.AD_ATT,
+#         )
+#     )
+
+
+@cute.jit
+def shuffle_sync(
+    value: cute.Numeric,
+    offset: cute.typing.Int,
+    width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE,
+) -> cute.Numeric:
+    assert value.width % 32 == 0, "value type must be a multiple of 32 bits"
+    # 1 -> 0b11111, 2 -> 0b11110, 4 -> 0b11100, 8 -> 0b11000, 16 -> 0b10000, 32 -> 0b00000
+    mask = cute.arch.WARP_SIZE - width
+    clamp = cute.arch.WARP_SIZE - 1
+    mask_and_clamp = mask << 8 | clamp
+    # important: need stride 1 and not 0 for recast_tensor to work
+    val = cute.make_rmem_tensor(cute.make_layout((1,), stride=(1,)), type(value))
+    val[0] = value
+    val_i32 = cute.recast_tensor(val, cutlass.Int32)
+    for i in cutlass.range_constexpr(cute.size(val_i32)):
+        val_i32[i] = cute.arch.shuffle_sync(val_i32[i], offset, mask_and_clamp=mask_and_clamp)
+    return val[0]
+
+
+@dsl_user_op
+def shl_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32:
+    """
+    Left-shift val by shift bits using PTX shl.b32 (sign-agnostic).
+
+    Named ``shl_u32`` (not ``shl_b32``) because python type annotations
+    distinguish signed/unsigned.
+
+    PTX semantics (§9.7.8.8): "Shift amounts greater than the register width N
+    are clamped to N."  So ``shl.b32 d, a, 32`` is well-defined and yields 0.
+
+    This differs from C/C++ and LLVM IR, where shifting by >= the type width is
+    undefined behavior.  CuTeDSL compiles through MLIR -> LLVM IR, so a plain
+    Python-level ``Uint32(x) << Uint32(n)`` inherits LLVM's UB: the optimizer
+    may treat the result as poison and eliminate dependent code.  Inline PTX
+    bypasses the LLVM IR shift entirely — the instruction is emitted verbatim
+    into PTX where clamping makes it safe for all shift amounts.
+    """
+    return cutlass.Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [
+                cutlass.Uint32(val).ir_value(loc=loc, ip=ip),
+                cutlass.Uint32(shift).ir_value(loc=loc, ip=ip),
+            ],
+            "shl.b32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def shr_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32:
+    """
+    Unsigned right-shift val by shift bits using PTX shr.u32 (zero-fills).
+
+    See ``shl_u32`` docstring for why inline PTX is used instead of plain
+    CuTeDSL shift operators (LLVM shift-by-type-width UB).
+    """
+    return cutlass.Uint32(
+        llvm.inline_asm(
+            T.i32(),
+            [
+                cutlass.Uint32(val).ir_value(loc=loc, ip=ip),
+                cutlass.Uint32(shift).ir_value(loc=loc, ip=ip),
+            ],
+            "shr.u32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@cute.jit
+def warp_prefix_sum(val: cutlass.Int32, lane: Optional[cutlass.Int32] = None) -> cutlass.Int32:
+    if const_expr(lane is None):
+        lane = cute.arch.lane_idx()
+    # if cute.arch.thread_idx()[0] >= 128 and cute.arch.thread_idx()[0] < 128 + 32 and cute.arch.block_idx()[0] == 0: cute.printf("tidx = %d, val = %d", cute.arch.thread_idx()[0] % 32, val)
+    for i in cutlass.range_constexpr(int(math.log2(cute.arch.WARP_SIZE))):
+        offset = 1 << i
+        # Very important that we set mask_and_clamp to 0
+        partial_sum = cute.arch.shuffle_sync_up(val, offset=offset, mask_and_clamp=0)
+        if lane >= offset:
+            val += partial_sum
+        # if cute.arch.thread_idx()[0] >= 128 and cute.arch.thread_idx()[0] < 128 + 32 and cute.arch.block_idx()[0] == 0: cute.printf("tidx = %d, partial_sum = %d, val = %d", cute.arch.thread_idx()[0] % 32, partial_sum, val)
+    return val
+
+
+@dsl_user_op
+def cvt_f16x2_f32(
+    a: float | Float32, b: float | Float32, to_dtype: Type, *, loc=None, ip=None
+) -> cutlass.Int32:
+    assert to_dtype in [cutlass.BFloat16, cutlass.Float16], "to_dtype must be BFloat16 or Float16"
+    return cutlass.Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)],
+            f"cvt.rn.{'bf16x2' if to_dtype is cutlass.BFloat16 else 'f16x2'}.f32 $0, $2, $1;",
+            "=r,f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@overload
+def cvt_f16(src: cute.Tensor, dst: cute.Tensor) -> None: ...
+
+
+@overload
+def cvt_f16(src: cute.Tensor, dtype: Type[cute.Numeric]) -> cute.Tensor: ...
+
+
+@cute.jit
+def cvt_f16(src: cute.Tensor, dst_or_dtype):
+    """Convert Float32 tensor to Float16/BFloat16.
+
+    Args:
+        src: Source tensor with Float32 element type
+        dst_or_dtype: Either a destination tensor or a dtype (Float16/BFloat16)
+
+    Returns:
+        None if dst is a tensor, or a new tensor if dtype is provided
+    """
+    if const_expr(isinstance(dst_or_dtype, type)):
+        # dtype variant: create new tensor and call the tensor variant
+        dtype = dst_or_dtype
+        dst = cute.make_rmem_tensor(src.shape, dtype)
+        cvt_f16(src, dst)
+        return dst
+    else:
+        # tensor variant: write to dst
+        dst = dst_or_dtype
+        assert cute.size(dst.shape) == cute.size(src.shape), "dst and src must have the same size"
+        assert cute.size(src.shape) % 2 == 0, "src must have an even number of elements"
+        assert dst.element_type in [cutlass.BFloat16, cutlass.Float16], (
+            "dst must be BFloat16 or Float16"
+        )
+        assert src.element_type is Float32, "src must be Float32"
+        dst_i32 = cute.recast_tensor(dst, cutlass.Int32)
+        assert cute.size(dst_i32.shape) * 2 == cute.size(src.shape)
+        for i in cutlass.range_constexpr(cute.size(dst_i32)):
+            dst_i32[i] = cvt_f16x2_f32(src[2 * i], src[2 * i + 1], dst.element_type)
+
+
+@dsl_user_op
+@cute.jit
+def evaluate_polynomial(x: Float32, poly: Tuple[Float32, ...], *, loc=None, ip=None) -> Float32:
+    deg = len(poly) - 1
+    out = poly[deg]
+    for i in cutlass.range_constexpr(deg - 1, -1, -1):
+        out = out * x + poly[i]
+    return out
+
+
+@dsl_user_op
+@cute.jit
+def evaluate_polynomial_2(
+    x: Float32, y: Float32, poly: Tuple[Float32, ...], *, loc=None, ip=None
+) -> Tuple[Float32, Float32]:
+    deg = len(poly) - 1
+    out = (poly[deg], poly[deg])
+    for i in cutlass.range_constexpr(deg - 1, -1, -1):
+        out = cute.arch.fma_packed_f32x2(out, (x, y), (poly[i], poly[i]))
+    return out
+
+
+@dsl_user_op
+def add_round_down(x: float | Float32, y: float | Float32, *, loc=None, ip=None) -> Float32:
+    # There's probably a way to call llvm or nvvm to do this instead of ptx
+    return cutlass.Float32(
+        llvm.inline_asm(
+            T.f32(),
+            [Float32(x).ir_value(loc=loc, ip=ip), Float32(y).ir_value(loc=loc, ip=ip)],
+            "add.rm.ftz.f32 $0, $1, $2;",
+            "=f,f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def combine_int_frac_ex2(x_rounded: Float32, frac_ex2: Float32, *, loc=None, ip=None) -> Float32:
+    return cutlass.Float32(
+        llvm.inline_asm(
+            T.f32(),
+            [
+                Float32(x_rounded).ir_value(loc=loc, ip=ip),
+                Float32(frac_ex2).ir_value(loc=loc, ip=ip),
+            ],
+            "{\n\t"
+            ".reg .s32 x_rounded_i, frac_ex_i, x_rounded_e, out_i;\n\t"
+            "mov.b32 x_rounded_i, $1;\n\t"
+            "mov.b32 frac_ex_i, $2;\n\t"
+            "shl.b32 x_rounded_e, x_rounded_i, 23;\n\t"
+            # add.u32 generates IMAD instruction and add.s32 generates LEA instruction
+            # IMAD uses the FMA pipeline and LEA uses the ALU pipeline, afaik
+            "add.s32 out_i, x_rounded_e, frac_ex_i;\n\t"
+            "mov.b32 $0, out_i;\n\t"
+            "}\n",
+            "=f,f,f",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def ex2_emulation(x: Float32, *, poly_degree: int = 3, loc=None, ip=None) -> Float32:
+    assert poly_degree in POLY_EX2, f"Polynomial degree {poly_degree} not supported"
+    # We assume x <= 127.0
+    fp32_round_int = float(2**23 + 2**22)
+    x_clamped = cute.arch.fmax(x, -127.0)
+    # We want to round down here, so that the fractional part is in [0, 1)
+    x_rounded = add_round_down(x_clamped, fp32_round_int, loc=loc, ip=ip)
+    # The integer floor of x is now in the last 8 bits of x_rounded
+    # We assume the next 2 ops round to nearest even. The rounding mode is important.
+    x_rounded_back = x_rounded - fp32_round_int
+    x_frac = x_clamped - x_rounded_back
+    x_frac_ex2 = evaluate_polynomial(x_frac, POLY_EX2[poly_degree], loc=loc, ip=ip)
+    return combine_int_frac_ex2(x_rounded, x_frac_ex2, loc=loc, ip=ip)
+
+
+# TODO: check that the ex2_emulation_2 produces the same SASS as the ptx version
+@dsl_user_op
+def ex2_emulation_2(
+    x: Float32, y: Float32, *, poly_degree: int = 3, loc=None, ip=None
+) -> Tuple[Float32, Float32]:
+    # We assume x <= 127.0 and y <= 127.0
+    fp32_round_int = float(2**23 + 2**22)
+    xy_clamped = (cute.arch.fmax(x, -127.0), cute.arch.fmax(y, -127.0))
+    # We want to round down here, so that the fractional part is in [0, 1)
+    xy_rounded = cute.arch.add_packed_f32x2(xy_clamped, (fp32_round_int, fp32_round_int), rnd="rm")
+    # The integer floor of x & y are now in the last 8 bits of xy_rounded
+    # We want the next 2 ops to round to nearest even. The rounding mode is important.
+    xy_rounded_back = quack.activation.sub_packed_f32x2(
+        xy_rounded, (fp32_round_int, fp32_round_int)
+    )
+    xy_frac = quack.activation.sub_packed_f32x2(xy_clamped, xy_rounded_back)
+    xy_frac_ex2 = evaluate_polynomial_2(*xy_frac, POLY_EX2[poly_degree], loc=loc, ip=ip)
+    x_out = combine_int_frac_ex2(xy_rounded[0], xy_frac_ex2[0], loc=loc, ip=ip)
+    y_out = combine_int_frac_ex2(xy_rounded[1], xy_frac_ex2[1], loc=loc, ip=ip)
+    return x_out, y_out
+
+
+@dsl_user_op
+def e2e_asm2(x: Float32, y: Float32, *, loc=None, ip=None) -> Tuple[Float32, Float32]:
+    out_f32x2 = llvm.inline_asm(
+        llvm.StructType.get_literal([T.f32(), T.f32()]),
+        [Float32(x).ir_value(loc=loc, ip=ip), Float32(y, loc=loc, ip=ip).ir_value()],
+        "{\n\t"
+        ".reg .f32 f1, f2, f3, f4, f5, f6, f7;\n\t"
+        ".reg .b64 l1, l2, l3, l4, l5, l6, l7, l8, l9, l10;\n\t"
+        ".reg .s32 r1, r2, r3, r4, r5, r6, r7, r8;\n\t"
+        "max.ftz.f32 f1, $2, 0fC2FE0000;\n\t"
+        "max.ftz.f32 f2, $3, 0fC2FE0000;\n\t"
+        "mov.b64 l1, {f1, f2};\n\t"
+        "mov.f32 f3, 0f4B400000;\n\t"
+        "mov.b64 l2, {f3, f3};\n\t"
+        "add.rm.ftz.f32x2 l7, l1, l2;\n\t"
+        "sub.rn.ftz.f32x2 l8, l7, l2;\n\t"
+        "sub.rn.ftz.f32x2 l9, l1, l8;\n\t"
+        "mov.f32 f7, 0f3D9DF09D;\n\t"
+        "mov.b64 l6, {f7, f7};\n\t"
+        "mov.f32 f6, 0f3E6906A4;\n\t"
+        "mov.b64 l5, {f6, f6};\n\t"
+        "mov.f32 f5, 0f3F31F519;\n\t"
+        "mov.b64 l4, {f5, f5};\n\t"
+        "mov.f32 f4, 0f3F800000;\n\t"
+        "mov.b64 l3, {f4, f4};\n\t"
+        "fma.rn.ftz.f32x2 l10, l9, l6, l5;\n\t"
+        "fma.rn.ftz.f32x2 l10, l10, l9, l4;\n\t"
+        "fma.rn.ftz.f32x2 l10, l10, l9, l3;\n\t"
+        "mov.b64 {r1, r2}, l7;\n\t"
+        "mov.b64 {r3, r4}, l10;\n\t"
+        "shl.b32 r5, r1, 23;\n\t"
+        "add.s32 r7, r5, r3;\n\t"
+        "shl.b32 r6, r2, 23;\n\t"
+        "add.s32 r8, r6, r4;\n\t"
+        "mov.b32 $0, r7;\n\t"
+        "mov.b32 $1, r8;\n\t"
+        "}\n",
+        "=r,=r,f,f",
+        has_side_effects=False,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+    out0 = Float32(llvm.extractvalue(T.f32(), out_f32x2, [0], loc=loc, ip=ip))
+    out1 = Float32(llvm.extractvalue(T.f32(), out_f32x2, [1], loc=loc, ip=ip))
+    return out0, out1
+
+
+@dsl_user_op
+def domain_offset_aligned(
+    coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None
+) -> cute.Tensor:
+    assert isinstance(tensor.iterator, cute.Pointer)
+    # We assume that applying the offset does not change the pointer alignment
+    new_ptr = cute.make_ptr(
+        tensor.element_type,
+        elem_pointer(tensor, coord).toint(),
+        tensor.memspace,
+        assumed_align=tensor.iterator.alignment,
+    )
+    return cute.make_tensor(new_ptr, tensor.layout)
+
+
+@dsl_user_op
+def warp_reduction(
+    val: cute.Numeric, op: Callable, *, threads_in_group: int = 32, loc=None, ip=None
+) -> cute.Numeric:
+    """Warp-wide reduction helper for a custom binary op."""
+    offset = threads_in_group // 2
+    while offset > 0:
+        val = op(
+            val,
+            cute.arch.shuffle_sync_bfly(
+                val, offset=offset, mask=-1, mask_and_clamp=31, loc=loc, ip=ip
+            ),
+        )
+        offset //= 2
+    return val
+
+
+warp_reduction_max = partial(
+    warp_reduction, op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else max(x, y)
+)
+warp_reduction_sum = partial(warp_reduction, op=lambda x, y: x + y)  # noqa: FURB118
+
+
+@dsl_user_op
+def make_cotiled_copy(
+    atom: cute.CopyAtom, atom_layout_tv: cute.Layout, data_layout: cute.Layout, *, loc=None, ip=None
+) -> cute.TiledCopy:
+    """Compatibility wrapper for deprecated CuTeDSL `make_cotiled_copy`."""
+    assert cute.is_static(atom_layout_tv.type), "atom_layout_tv must be static"
+    assert cute.is_static(data_layout.type), "data_layout must be static"
+
+    inv_layout_ = cute.left_inverse(data_layout, loc=loc, ip=ip)
+    inv_data_layout = cute.make_layout(
+        (inv_layout_.shape, (1)), stride=(inv_layout_.stride, (0)), loc=loc, ip=ip
+    )
+    layout_tv_data = cute.composition(inv_data_layout, atom_layout_tv, loc=loc, ip=ip)
+
+    atom_layout_v_to_check = cute.coalesce(
+        cute.make_layout(atom_layout_tv.shape[1], stride=atom_layout_tv.stride[1], loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+    data_layout_v_to_check = cute.coalesce(
+        cute.composition(
+            data_layout,
+            cute.make_layout(
+                layout_tv_data.shape[1], stride=layout_tv_data.stride[1], loc=loc, ip=ip
+            ),
+            loc=loc,
+            ip=ip,
+        ),
+        loc=loc,
+        ip=ip,
+    )
+    assert data_layout_v_to_check == atom_layout_v_to_check, (
+        "the memory pointed to by atom_layout_tv does not exist in the data_layout."
+    )
+
+    flat_data_shape = cute.product_each(data_layout.shape, loc=loc, ip=ip)
+    tiler = tuple(
+        cute.filter(
+            cute.composition(
+                cute.make_layout(
+                    flat_data_shape,
+                    stride=tuple(0 if j != i else 1 for j in range(cute.rank(flat_data_shape))),
+                    loc=loc,
+                    ip=ip,
+                ),
+                layout_tv_data,
+                loc=loc,
+                ip=ip,
+            ),
+            loc=loc,
+            ip=ip,
+        )
+        for i in range(cute.rank(flat_data_shape))
+    )
+    tile2data = cute.composition(
+        cute.make_layout(flat_data_shape, loc=loc, ip=ip), tiler, loc=loc, ip=ip
+    )
+    layout_tv = cute.composition(
+        cute.left_inverse(tile2data, loc=loc, ip=ip), layout_tv_data, loc=loc, ip=ip
+    )
+    return cute.make_tiled_copy(atom, layout_tv, tiler, loc=loc, ip=ip)
+
+
+@cute.jit
+def scalar_to_ssa(a: cute.Numeric, dtype) -> cute.TensorSSA:
+    """Convert a scalar to a cute TensorSSA of shape (1,) and given dtype"""
+    vec = cute.make_rmem_tensor(1, dtype)
+    vec[0] = a
+    return vec.load()
+
+
+def ssa_to_scalar(val):
+    """Could inline but nice for reflecting the above api"""
+    return val[0]
+
+
+@cute.jit
+def get_batch_from_cu_tensor(idx: Int32, cu_tensor: cute.Tensor) -> Int32:
+    """Binary search to determine batch from packed index in a cumulative tensor"""
+    batch_size = cute.size(cu_tensor) - 1
+    lo = Int32(0)
+    hi = batch_size
+
+    while lo < hi:
+        mid = (lo + hi) // 2
+        if cu_tensor[mid + 1] <= idx:
+            lo = mid + 1
+        else:
+            hi = mid
+
+    return lo
@@ -1,4 +1,4 @@
-"""VSA-128/256 block-sparse attention wrappers.
+"""VSA-256 block-sparse attention wrapper.

 The default 256-block path is Triton: it expands the logical 256-block map
 to the existing 64-block Triton kernel via a dense 4x4 expansion per logical
@@ -6,9 +6,13 @@ edge ("route A"), and requires no optional dependencies.

 The FA4 CuTe block-sparse fastpath (intended for Blackwell sm_100+) is
 *opt-in* via ``FASTVIDEO_VSA_CUTEDSL=1``. It routes to
-:mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`, which natively operates
-on 128-token Q/KV blocks (the 256 wrapper expands its logical KV map and
-sizes into that physical representation). The CuTe kernel
+:mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`. By default this wrapper
+expands logical KV256 blocks to the historical physical KV128 tiles. Set
+``FASTVIDEO_VSA_FA4_BLOCK_SHAPE=128x64`` for the fine-grained FA4 schedule.
+The production ``64x64`` selection pairs adjacent Q64 children into physical
+Q128 tiles: their KV lists are identical because this wrapper derives them
+from one original Q256 metadata row. Direct Q64 callers with arbitrary masks
+continue to use the native Q64/KV64 path. The CuTe kernel
 (``flash_attn.cute`` with block-sparsity) is an optional dependency,
 imported lazily only when this fastpath is selected.

@@ -30,12 +34,23 @@ from .block_sparse_attn import block_sparse_attn_triton, _force_triton
 # FA4 CuTe build (``flash_attn.cute``) and make it a hard dependency of the
 # default Triton path.

-_KV_BLOCK_PHYS = 128  # FA4 CuTe BSA forward uses 128-token KV blocks.
+_LOGICAL_BLOCK_SIZE = 256
+_FA4_BLOCK_SHAPES = {
+    # The public name describes the original VSA mask. FA4 internally splits
+    # its logical KV256 edge into two physical KV128 tiles.
+    "256x256": (256, 128),
+    "128x64": (128, 64),
+    # This pairing is safe only at this original-Q256 wrapper boundary: all
+    # four Q64 children inherit exactly the same KV list from their parent.
+    # Keep arbitrary/distinct Q64 maps on block_sparse_attn_cute_fwd's native
+    # Q64/KV64 path, where they are not coalesced.
+    "64x64": (128, 64),
+}
 _KV_BLOCK_TRITON = 64  # Existing Triton path uses 64-token KV blocks.


 def _resolve_backend() -> str:
-    """Pick the backend for the 128/256-block VSA paths.
+    """Pick the backend for the 256-block VSA path.

    Default is Triton (no optional deps). The FA4 CuTe fastpath is opt-in
    via ``FASTVIDEO_VSA_CUTEDSL=1`` and requires the optional FA4 CuTe
@@ -49,51 +64,52 @@ def _resolve_backend() -> str:
    return "triton"


-def _expand_mask_and_sizes_128_to_64(
-    logical_mask_128: torch.Tensor,
-    logical_kv_sizes_128: torch.Tensor,
+def _resolve_fa4_block_shape() -> Tuple[int, int]:
+    requested = os.environ.get("FASTVIDEO_VSA_FA4_BLOCK_SHAPE", "256x256").lower()
+    try:
+        return _FA4_BLOCK_SHAPES[requested]
+    except KeyError as exc:
+        choices = ", ".join(_FA4_BLOCK_SHAPES)
+        raise ValueError(
+            f"unsupported FASTVIDEO_VSA_FA4_BLOCK_SHAPE={requested!r}; choose one of {choices}"
+        ) from exc
+
+
+def _expand_mask_and_sizes_256_for_fa4(
+    logical_mask_256: torch.Tensor,
+    logical_kv_sizes_256: torch.Tensor,
+    q_block_size: int,
+    kv_block_size: int,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Expand a [B, H, Qb128, KVb128] map to 64-token Triton tiles."""
-    expanded_mask = logical_mask_128.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3)
-    sizes_i32 = logical_kv_sizes_128.to(torch.int32)
-    offsets = torch.tensor(
-        [0, _KV_BLOCK_TRITON],
+    """Expand a logical Q256/KV256 map to one supported physical FA4 shape.
+
+    Every child inherits its parent edge. KV valid-token counts are clamped
+    into each child's window, preserving partial and empty tail blocks.
+    """
+    if _LOGICAL_BLOCK_SIZE % q_block_size or _LOGICAL_BLOCK_SIZE % kv_block_size:
+        raise ValueError(
+            f"FA4 QxKV blocks must divide {_LOGICAL_BLOCK_SIZE}, got "
+            f"{q_block_size}x{kv_block_size}"
+        )
+    q_factor = _LOGICAL_BLOCK_SIZE // q_block_size
+    kv_factor = _LOGICAL_BLOCK_SIZE // kv_block_size
+    expanded_mask = logical_mask_256.repeat_interleave(q_factor, dim=2)
+    expanded_mask = expanded_mask.repeat_interleave(kv_factor, dim=3)
+
+    sizes_i32 = logical_kv_sizes_256.to(torch.int32)
+    offsets = torch.arange(
+        kv_factor,
        dtype=torch.int32,
        device=sizes_i32.device,
-    )
+    ) * kv_block_size
    expanded_sizes = torch.clamp(
        sizes_i32[:, None] - offsets[None, :],
        min=0,
-        max=_KV_BLOCK_TRITON,
+        max=kv_block_size,
    ).reshape(-1)
    return expanded_mask, expanded_sizes


-def _expand_mask_and_sizes_256_to_128(
-    logical_mask_256: torch.Tensor,
-    logical_kv_sizes_256: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Expand a [B, H, Qb256, KVb256] map to [B, H, Qb256, KVb128].
-
-    Each logical 256-token KV block splits into two physical 128-token
-    children. Each child inherits the logical mask edge; its valid-token
-    count is the logical count clamped into the child's window.
-    """
-    expanded_mask = logical_mask_256.repeat_interleave(2, dim=3)
-
-    sizes_i32 = logical_kv_sizes_256.to(torch.int32)
-    child0 = torch.clamp(sizes_i32, min=0, max=_KV_BLOCK_PHYS)
-    child1 = torch.clamp(sizes_i32 - _KV_BLOCK_PHYS, min=0, max=_KV_BLOCK_PHYS)
-    expanded_sizes = torch.empty(
-        (sizes_i32.numel() * 2, ),
-        dtype=torch.int32,
-        device=sizes_i32.device,
-    )
-    expanded_sizes[0::2] = child0
-    expanded_sizes[1::2] = child1
-    return expanded_mask, expanded_sizes
-
-
 def _expand_mask_and_sizes_256_to_64(
    logical_mask_256: torch.Tensor,
    logical_kv_sizes_256: torch.Tensor,
@@ -132,63 +148,6 @@ def _triton_via_route_a(
    return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)


-def _triton_via_route_a_128(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    logical_mask_128: torch.Tensor,
-    logical_kv_sizes_128: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    from .triton_kernels.index import map_to_index as triton_map_to_index
-
-    mask_64, sizes_64 = _expand_mask_and_sizes_128_to_64(logical_mask_128, logical_kv_sizes_128)
-    q2k_idx, q2k_num = triton_map_to_index(mask_64.to(torch.bool))
-    return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
-
-
-def block_sparse_attn_128(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    logical_block_map_128: torch.Tensor,
-    logical_variable_block_sizes_128: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """VSA-128 sparse-branch entrypoint for [B, H, S, D] inputs."""
-    if logical_block_map_128.dim() == 3:
-        logical_block_map_128 = logical_block_map_128.unsqueeze(0)
-
-    if _resolve_backend() == "triton":
-        return _triton_via_route_a_128(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
-
-    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd
-    return block_sparse_attn_cute_fwd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
-
-
-def block_sparse_attn_128_bshd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    logical_block_map_128: torch.Tensor,
-    logical_variable_block_sizes_128: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """VSA-128 sparse-branch entrypoint for [B, S, H, D] inputs."""
-    if logical_block_map_128.dim() == 3:
-        logical_block_map_128 = logical_block_map_128.unsqueeze(0)
-
-    if _resolve_backend() == "triton":
-        out_bhsd, aux = _triton_via_route_a_128(
-            q.transpose(1, 2).contiguous(),
-            k.transpose(1, 2).contiguous(),
-            v.transpose(1, 2).contiguous(),
-            logical_block_map_128,
-            logical_variable_block_sizes_128,
-        )
-        return out_bhsd.transpose(1, 2).contiguous(), aux
-
-    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
-    return block_sparse_attn_cute_fwd_bshd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
-
-
 def block_sparse_attn_256(
    q: torch.Tensor,
    k: torch.Tensor,
@@ -203,9 +162,15 @@ def block_sparse_attn_256(
    if _resolve_backend() == "triton":
        return _triton_via_route_a(q, k, v, logical_block_map_256, logical_variable_block_sizes_256)

-    mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(logical_block_map_256, logical_variable_block_sizes_256)
+    q_block_size, kv_block_size = _resolve_fa4_block_shape()
+    physical_mask, physical_sizes = _expand_mask_and_sizes_256_for_fa4(
+        logical_block_map_256,
+        logical_variable_block_sizes_256,
+        q_block_size,
+        kv_block_size,
+    )
    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd
-    return block_sparse_attn_cute_fwd(q, k, v, mask_128, sizes_128)
+    return block_sparse_attn_cute_fwd(q, k, v, physical_mask, physical_sizes)


 def block_sparse_attn_256_bshd(
@@ -233,6 +198,12 @@ def block_sparse_attn_256_bshd(
        )
        return out_bhsd.transpose(1, 2).contiguous(), aux

-    mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(logical_block_map_256, logical_variable_block_sizes_256)
+    q_block_size, kv_block_size = _resolve_fa4_block_shape()
+    physical_mask, physical_sizes = _expand_mask_and_sizes_256_for_fa4(
+        logical_block_map_256,
+        logical_variable_block_sizes_256,
+        q_block_size,
+        kv_block_size,
+    )
    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
-    return block_sparse_attn_cute_fwd_bshd(q, k, v, mask_128, sizes_128)
+    return block_sparse_attn_cute_fwd_bshd(q, k, v, physical_mask, physical_sizes)
@@ -1,18 +1,19 @@
-"""FA4 CuTe-DSL block-sparse attention adapter.
+"""CuTe-DSL block-sparse attention forward kernel.

-This module adapts VSA's ``(block_map, variable_block_sizes)`` inputs into
-FA4's forward and backward ``BlockSparseTensorsTorch`` representations.
-FA4's public ``flash_attn_func`` owns the forward/backward autograd bridge.
+Thin wrapper around `flash_attn.cute.interface._flash_attn_fwd` that adapts
+VSA's `(block_map, variable_block_sizes)` inputs into FA4's
+`BlockSparseTensorsTorch` representation and the per-KV-block validity mask.

 Both [B, H, S, D] (BHSD) and [B, S, H, D] (BSHD) entrypoints are provided.
-The BSHD variant is preferred from VSA-128/256 callers to avoid layout
+The BSHD variant is preferred from VSA-256 callers to avoid layout
 round-trips on the hot path.

 The FA4 CuTe block-sparse kernel (``flash_attn.cute`` with
-``block_sparsity``) is an *optional* dependency: it is imported lazily and
-only exercised when the VSA-128/256 CuTe fastpath is explicitly selected
-(``FASTVIDEO_VSA_CUTEDSL=1``). The default path is Triton and does not require
-it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``.
+``block_sparsity``) is an *optional* dependency: it is imported lazily. The
+legacy VSA-256 route selects it with ``FASTVIDEO_VSA_CUTEDSL=1``; direct
+fine-grained callers can use Q128/KV64 and, on SM100, Q64/KV64. The default
+VSA-256 path is Triton and does not require FA4. The CuTe path also needs
+``nvidia-cutlass-dsl`` and ``quack-kernels``.
 """

 from __future__ import annotations
@@ -22,14 +23,13 @@ from typing import Tuple

 import torch

-_FA4_IMPORT_HINT = ("VSA-128/256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
+_FA4_IMPORT_HINT = ("The CuTe block-sparse fastpath requires a FlashAttention-4 CuTe build that "
                    "provides `flash_attn.cute` with block-sparsity support (plus "
                    "`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional "
-                    "dependency; the default path is Triton. Install the FA4 CuTe "
-                    "build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath.")
+                    "dependency; the default VSA-256 path is Triton. Install the FA4 CuTe "
+                    "build; set FASTVIDEO_VSA_CUTEDSL=1 when selecting the VSA-256 route.")


-@functools.lru_cache(maxsize=1)
 def _load_fa4_cute():
    """Lazily import the optional FA4 CuTe block-sparse symbols.

@@ -39,39 +39,14 @@ def _load_fa4_cute():
    """
    try:
        from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch
-        from flash_attn.cute.interface import (
-            _flash_attn_bwd,
-            _flash_attn_fwd,
-            flash_attn_func,
-        )
+        from flash_attn.cute.interface import _flash_attn_fwd
    except ImportError as exc:  # pragma: no cover - optional dependency
        raise ImportError(_FA4_IMPORT_HINT) from exc
-    return BlockSparseTensorsTorch, flash_attn_func, _flash_attn_fwd, _flash_attn_bwd
+    return BlockSparseTensorsTorch, _flash_attn_fwd


-# FA4's physical Q tile size; KV block size comes from the VSA caller.
-_FA4_Q_BLOCK_SIZE = 128
-
-
-class _SingleQStageLength(int):
-    """Keep the real length while selecting FA4's one-stage Q128 path.
-
-    On sm_100 FA4 derives ``q_stage`` from ``max_seqlen_q > tile_m``. Its
-    kernel supports one 128-token Q stage, but the fixed-length public wrapper
-    does not expose that choice. VSA-128 must select it explicitly; otherwise
-    adjacent logical Q blocks are merged into a 256-token sparse block.
-    """
-
-    def __mul__(self, other):
-        return type(self)(int(self) * int(other))
-
-    def __rmul__(self, other):
-        return type(self)(int(other) * int(self))
-
-    def __gt__(self, other):
-        if int(other) == _FA4_Q_BLOCK_SIZE:
-            return False
-        return int(self) > int(other)
+# Default Q-side tile size; the Q64/KV64 specialization uses a 64-row tile.
+_M_BLOCK_SIZE_DEFAULT = 128


 def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -90,12 +65,24 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    return triton_map_to_index(block_map)


-def _choose_q_sparse_block_size(q_len: int, q_tile_size: int = _FA4_Q_BLOCK_SIZE) -> int:
-    # FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > q_tile_size.
+def _choose_q_sparse_block_size(q_len: int, m_block_size: int = _M_BLOCK_SIZE_DEFAULT) -> int:
+    # FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > m_block_size.
    major, _ = torch.cuda.get_device_capability()
-    if major >= 10 and q_len > q_tile_size:
-        return 2 * q_tile_size
-    return q_tile_size
+    if major >= 10 and q_len > m_block_size:
+        return 2 * m_block_size
+    return m_block_size
+
+
+def _choose_m_block_size(q_block_size: int, kv_block_size: int) -> int:
+    """Choose the physical FA4 Q tile without changing sparse-map semantics."""
+    if q_block_size == 64:
+        if kv_block_size != 64:
+            raise ValueError("the FA4 Q64 specialization requires KV64 blocks")
+        major, _minor = torch.cuda.get_device_capability()
+        if major != 10:
+            raise RuntimeError("the FA4 Q64/KV64 specialization currently requires SM100")
+        return 64
+    return _M_BLOCK_SIZE_DEFAULT


 def _aggregate_q_block_map(
@@ -160,35 +147,30 @@ def _build_vbs_mask_mod(kv_block_size: int):
    return _vbs_mask_mod


-def _build_sparse_tensors(
+def _cute_forward(
+    q_bshd: torch.Tensor,
+    k_bshd: torch.Tensor,
+    v_bshd: torch.Tensor,
    block_map: torch.Tensor,
    variable_block_sizes: torch.Tensor,
    *,
-    q_len: int,
    q_block_size: int,
    kv_block_size: int,
-    need_backward: bool,
-    force_q_sparse_block_size: int | None = None,
-) -> Tuple[object, object | None]:
-    """Build the Q-owned forward and KV-owned backward sparse metadata.
-
-    ``need_backward`` is False on inference-only calls: the backward metadata
-    is a pair of dense ``[B, H, kv_blocks, q_blocks]`` int32 index tensors that
-    FA4 keeps alive on its autograd ctx until backward runs, so building it
-    when nothing requires grad is pure overhead (~80 MiB per call at Wan-14B
-    720p shape).
-    """
-    BlockSparseTensorsTorch, _, _, _ = _load_fa4_cute()
-    if force_q_sparse_block_size is None:
-        q_sparse_candidate = _choose_q_sparse_block_size(q_len)
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Internal: FA4 CuTe BSA fwd with BSHD inputs."""
+    BlockSparseTensorsTorch, _flash_attn_fwd = _load_fa4_cute()
+    m_block_size = _choose_m_block_size(q_block_size, kv_block_size)
+    if q_block_size == m_block_size:
+        # Q128 and Q64 maps may select different KV blocks in adjacent rows.
+        # Preserve their declared granularity instead of unioning them into the
+        # historical two-stage Q256 schedule.
+        q_sparse_block_size = q_block_size
+    else:
+        q_sparse_candidate = _choose_q_sparse_block_size(q_bshd.shape[1], m_block_size)
        q_sparse_block_size = max(
            q_block_size,
            ((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size,
        )
-    else:
-        q_sparse_block_size = force_q_sparse_block_size
-        if q_sparse_block_size < q_block_size or q_sparse_block_size % q_block_size != 0:
-            raise ValueError("force_q_sparse_block_size must be a positive multiple of q_block_size")
    sparse_map = _aggregate_q_block_map(
        block_map,
        q_sparse_block_size=q_sparse_block_size,
@@ -196,166 +178,35 @@ def _build_sparse_tensors(
    )
    kv_full = (variable_block_sizes == kv_block_size).view(1, 1, 1, -1)
    kv_partial = ((variable_block_sizes > 0) & (variable_block_sizes < kv_block_size)).view(1, 1, 1, -1)
+    full_map = sparse_map & kv_full
+    mask_map = sparse_map & kv_partial

-    def from_maps(full_map: torch.Tensor, mask_map: torch.Tensor) -> object:
-        full_block_idx, full_block_cnt = _map_to_index(full_map.contiguous())
-        mask_block_idx, mask_block_cnt = _map_to_index(mask_map.contiguous())
-        return BlockSparseTensorsTorch(
-            full_block_cnt=full_block_cnt.to(torch.int32).contiguous(),
-            full_block_idx=full_block_idx.to(torch.int32).contiguous(),
-            mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(),
-            mask_block_idx=mask_block_idx.to(torch.int32).contiguous(),
-            block_size=(q_sparse_block_size, kv_block_size),
-        )
+    full_block_idx, full_block_cnt = _map_to_index(full_map)
+    mask_block_idx, mask_block_cnt = _map_to_index(mask_map)

-    forward_sparse_tensors = from_maps(
-        sparse_map & kv_full,
-        sparse_map & kv_partial,
+    sparse_tensors = BlockSparseTensorsTorch(
+        full_block_cnt=full_block_cnt.to(torch.int32).contiguous(),
+        full_block_idx=full_block_idx.to(torch.int32).contiguous(),
+        mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(),
+        mask_block_idx=mask_block_idx.to(torch.int32).contiguous(),
+        block_size=(q_sparse_block_size, kv_block_size),
    )

-    if not need_backward:
-        return forward_sparse_tensors, None
-
-    # FA4 backward is KV-owned: for each physical KV tile, list the sparse
-    # query tiles that selected it. Full and partial KV tiles stay separate
-    # so the token-level validity mask only runs for padded tiles.
-    backward_sparse_tensors = from_maps(
-        (sparse_map & kv_full).transpose(2, 3),
-        (sparse_map & kv_partial).transpose(2, 3),
-    )
-    return forward_sparse_tensors, backward_sparse_tensors
-
-
-def _cute_attention_q128_forward(
-    q_bshd: torch.Tensor,
-    k_bshd: torch.Tensor,
-    v_bshd: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-    *,
-    need_backward: bool,
-) -> Tuple[torch.Tensor, torch.Tensor, object | None]:
-    """Run FA4 with one physical Q stage per logical VSA-128 block."""
-    _, _, flash_attn_fwd, _ = _load_fa4_cute()
-    forward_sparse_tensors, backward_sparse_tensors = _build_sparse_tensors(
-        block_map,
-        variable_block_sizes,
-        q_len=q_bshd.shape[1],
-        q_block_size=_FA4_Q_BLOCK_SIZE,
-        kv_block_size=_FA4_Q_BLOCK_SIZE,
-        need_backward=need_backward,
-        force_q_sparse_block_size=_FA4_Q_BLOCK_SIZE,
-    )
-    out, lse = flash_attn_fwd(
+    # _flash_attn_fwd returns (out, lse, p, row_max); keep the first two.
+    out, lse = _flash_attn_fwd(
        q_bshd,
        k_bshd,
        v_bshd,
-        tile_mn=(_FA4_Q_BLOCK_SIZE, _FA4_Q_BLOCK_SIZE),
-        max_seqlen_q=_SingleQStageLength(q_bshd.shape[1]),
-        mask_mod=_build_vbs_mask_mod(_FA4_Q_BLOCK_SIZE),
-        block_sparse_tensors=forward_sparse_tensors,
+        tile_mn=(m_block_size, kv_block_size),
+        mask_mod=_build_vbs_mask_mod(kv_block_size),
+        block_sparse_tensors=sparse_tensors,
        aux_tensors=[variable_block_sizes],
        causal=False,
        return_lse=True,
    )[:2]
-    return out, lse, backward_sparse_tensors
-
-
-class _CuteAttentionQ128(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes):
-        out, lse, backward_sparse_tensors = _cute_attention_q128_forward(
-            q_bshd,
-            k_bshd,
-            v_bshd,
-            block_map,
-            variable_block_sizes,
-            need_backward=True,
-        )
-        ctx.save_for_backward(q_bshd, k_bshd, v_bshd, out, lse, variable_block_sizes)
-        ctx.backward_sparse_tensors = backward_sparse_tensors
-        ctx.mark_non_differentiable(lse)
-        ctx.set_materialize_grads(False)
-        return out, lse
-
-    @staticmethod
-    def backward(ctx, grad_out, grad_lse):
-        del grad_lse
-        q_bshd, k_bshd, v_bshd, out, lse, variable_block_sizes = ctx.saved_tensors
-        if grad_out is None:
-            grad_out = torch.zeros_like(out)
-        _, _, _, flash_attn_bwd = _load_fa4_cute()
-        dq, dk, dv = flash_attn_bwd(
-            q_bshd,
-            k_bshd,
-            v_bshd,
-            out,
-            grad_out.contiguous(),
-            lse,
-            softmax_scale=q_bshd.shape[-1]**-0.5,
-            mask_mod=_build_vbs_mask_mod(_FA4_Q_BLOCK_SIZE),
-            aux_tensors=[variable_block_sizes],
-            block_sparse_tensors=ctx.backward_sparse_tensors,
-        )
-        return dq, dk, dv, None, None
-
-
-def _cute_attention_q128(
-    q_bshd: torch.Tensor,
-    k_bshd: torch.Tensor,
-    v_bshd: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q_bshd, k_bshd, v_bshd))
-    if need_backward:
-        return _CuteAttentionQ128.apply(q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes)
-    out, lse, _ = _cute_attention_q128_forward(
-        q_bshd,
-        k_bshd,
-        v_bshd,
-        block_map,
-        variable_block_sizes,
-        need_backward=False,
-    )
    return out, lse


-def _cute_attention(
-    q_bshd: torch.Tensor,
-    k_bshd: torch.Tensor,
-    v_bshd: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Run FA4's autograd-enabled block-sparse attention with BSHD inputs."""
-    _, flash_attn_func, _, _ = _load_fa4_cute()
-    q_block_size = q_bshd.shape[1] // block_map.shape[2]
-    kv_block_size = k_bshd.shape[1] // block_map.shape[3]
-    if q_block_size == kv_block_size == _FA4_Q_BLOCK_SIZE:
-        return _cute_attention_q128(q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes)
-    need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q_bshd, k_bshd, v_bshd))
-    forward_sparse_tensors, backward_sparse_tensors = _build_sparse_tensors(
-        block_map,
-        variable_block_sizes,
-        q_len=q_bshd.shape[1],
-        q_block_size=q_block_size,
-        kv_block_size=kv_block_size,
-        need_backward=need_backward,
-    )
-    return flash_attn_func(
-        q_bshd,
-        k_bshd,
-        v_bshd,
-        mask_mod=_build_vbs_mask_mod(kv_block_size),
-        aux_tensors=[variable_block_sizes],
-        block_sparse_tensors=forward_sparse_tensors,
-        block_sparse_tensors_bwd=backward_sparse_tensors,
-        return_lse=True,
-    )
-
-
 def block_sparse_attn_cute_fwd(
    q: torch.Tensor,
    k: torch.Tensor,
@@ -363,25 +214,34 @@ def block_sparse_attn_cute_fwd(
    block_map: torch.Tensor,
    variable_block_sizes: torch.Tensor,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Autograd-enabled CuTe block-sparse attention for [B, H, S, D]."""
+    """CuTe forward-only block-sparse attention with [B, H, S, D] inputs."""
    if block_map.dim() == 3:
        block_map = block_map.unsqueeze(0)
+    q_block_size = q.shape[2] // block_map.shape[2]
+    kv_block_size = k.shape[2] // block_map.shape[3]

    q_bshd = q.transpose(1, 2).contiguous()
    k_bshd = k.transpose(1, 2).contiguous()
    v_bshd = v.transpose(1, 2).contiguous()
-    out_bshd, lse = _cute_attention(
+    out_bshd, lse_bshd = _cute_forward(
        q_bshd,
        k_bshd,
        v_bshd,
        block_map,
        variable_block_sizes,
+        q_block_size=q_block_size,
+        kv_block_size=kv_block_size,
    )
    out = out_bshd.transpose(1, 2).contiguous()
-    # FA4 already returns lse as [B, H, S], matching the Triton path's aux
-    # contract, so it needs no transpose. Detach before any further op: the
-    # value is informational and callers never backprop through it.
-    return out, lse.detach()
+    if lse_bshd is None:
+        lse = torch.empty(
+            (q.shape[0], q.shape[1], q.shape[2]),
+            dtype=torch.float32,
+            device=q.device,
+        )
+    else:
+        lse = lse_bshd.transpose(1, 2).contiguous()
+    return out, lse


 def block_sparse_attn_cute_fwd_bshd(
@@ -391,16 +251,27 @@ def block_sparse_attn_cute_fwd_bshd(
    block_map: torch.Tensor,
    variable_block_sizes: torch.Tensor,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Autograd-enabled CuTe block-sparse attention for [B, S, H, D]."""
+    """CuTe forward-only block-sparse attention with [B, S, H, D] inputs."""
    if block_map.dim() == 3:
        block_map = block_map.unsqueeze(0)
+    q_block_size = q.shape[1] // block_map.shape[2]
+    kv_block_size = k.shape[1] // block_map.shape[3]

-    out, lse = _cute_attention(
+    out, lse_bshd = _cute_forward(
        q,
        k,
        v,
        block_map,
        variable_block_sizes,
+        q_block_size=q_block_size,
+        kv_block_size=kv_block_size,
    )
-    # lse is [B, H, S] regardless of the q/k/v layout; see above.
-    return out, lse.detach()
+    if lse_bshd is None:
+        lse = torch.empty(
+            (q.shape[0], q.shape[2], q.shape[1]),
+            dtype=torch.float32,
+            device=q.device,
+        )
+    else:
+        lse = lse_bshd.transpose(1, 2).contiguous()
+    return out, lse
@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""sm_100a (Blackwell) CUDA block-sparse VSA forward.
-
-A third backend behind the same VSA op as the Triton and CuTe-DSL paths. Forward only: it
-returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward``
-writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]`` fp32 -- so
-``block_sparse_attn_backward_triton`` runs against it unchanged.
-
-The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks
-(tile volumes 64 and 128 in ``build_vsa_metadata``); the block size is inferred from the
-tensors and picks the op. Anything else falls back to Triton via ``is_supported``.
-"""
-
-from typing import Tuple
-
-import torch
-
-try:
-    # The pybind symbols live on fastvideo_kernel_ops, NOT on the _C package that contains it.
-    # `import fastvideo_kernel._C as _C` resolves to the namespace package, whose __init__ is
-    # empty, so hasattr() fails on a wheel install and the caller silently falls back with the
-    # kernel built and present.
-    from fastvideo_kernel._C import fastvideo_kernel_ops as _C
-    _FWD_BY_BLOCK = {
-        64: getattr(_C, "block_sparse_sm100a_fwd", None),
-        128: getattr(_C, "block_sparse_sm100a_blk128_fwd", None),
-    }
-    _HAS_VSA_SM100A = any(_FWD_BY_BLOCK.values())
-except ImportError:  # pragma: no cover - extension not built
-    _C = None
-    _FWD_BY_BLOCK = {}
-    _HAS_VSA_SM100A = False
-
-_SM100 = (10, 0)
-HEAD_DIM = 128
-# Must match the -DVSA_BHSD the extension was compiled with (see CMakeLists).
-BHSD = True
-
-
-def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int:
-    num_blocks = variable_block_sizes.numel()
-    seqlen = q.shape[2] if BHSD else q.shape[1]
-    return 0 if num_blocks == 0 or seqlen % num_blocks else seqlen // num_blocks
-
-
-def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
-    """True iff this build can run these tensors; otherwise the caller uses Triton.
-
-    Static facts only -- shapes, dtypes, arch, layout. Deliberately NO reads of tensor
-    contents: the previous ``int(variable_block_sizes.min())`` was a GPU->CPU sync on every
-    call, and the kernel no longer needs it (see below). This predicate must stay cheap
-    enough to sit on a per-layer dispatch path.
-
-    What the kernel accepts (and is tested to handle):
-      * q/k/v: contiguous 4-D bf16 CUDA tensors on an sm_100 device, head_dim 128, laid out
-        as compiled (BHSD here); seqlen == num_blocks * block with an EVEN num_blocks (a CTA
-        owns an adjacent pair of query blocks) and a 64- or 128-token build present.
-      * q2k_num: any per-row counts in [0, max_kv], NON-uniform across rows included. Rows
-        with count 0 produce exactly-zero output rows (and a finite LSE sentinel) rather
-        than attending anywhere -- so no ``.min()`` floor is required of the caller.
-      * q2k_idx: rows only need valid entries (in [0, num_blocks)) BELOW that row's count;
-        padding past the count (e.g. map_to_index's -1 fill) is never dereferenced. max_kv
-        (= q2k_idx.shape[-1]) must be >= 1, which the host launcher re-checks.
-      * variable_block_sizes: per-KV-block valid-token counts in [0, block]; keys at or past
-        a block's count are masked. Integer metadata is converted to int32/contiguous by
-        ``block_sparse_attn_sm100a`` itself, so int64 inputs merely cost a cast.
-    """
-    if not _HAS_VSA_SM100A or not q.is_cuda:
-        return False
-    if torch.cuda.get_device_capability(q.device) != _SM100:
-        return False
-    if q.dtype != torch.bfloat16 or q.dim() != 4 or q.shape[-1] != HEAD_DIM:
-        return False
-    if not q.is_contiguous():
-        return False
-    if _FWD_BY_BLOCK.get(_block_size(q, variable_block_sizes)) is None:
-        return False
-    # A CTA owns an adjacent pair of query blocks.
-    if variable_block_sizes.numel() % 2 != 0:
-        return False
-    # Metadata must be integer-typed so the wrapper's int32 conversion is value-preserving.
-    if not variable_block_sizes.is_cuda or variable_block_sizes.dtype not in (torch.int32, torch.int64):
-        return False
-    return True
-
-
-def block_sparse_attn_sm100a(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    q2k_idx: torch.Tensor,
-    q2k_num: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-    need_lse: bool = True,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Forward pass. Returns ``(out, lse)``; ``out`` has q's layout."""
-    fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)]
-    idx = q2k_idx.to(torch.int32).contiguous()
-    num = q2k_num.to(torch.int32).contiguous()
-    vbs = variable_block_sizes.to(torch.int32).contiguous()
-    sm_scale = 1.0 / (q.shape[-1]**0.5)
-    res = fwd(q.contiguous(), k.contiguous(), v.contiguous(), None,
-              idx, num, vbs, sm_scale, need_lse)
-    return (res[0], res[1]) if need_lse else (res[0], None)
@@ -2,8 +2,6 @@ import math
 import torch
 from .block_sparse_attn import block_sparse_attn
 from .block_sparse_attn_256 import (
-    block_sparse_attn_128,
-    block_sparse_attn_128_bshd,
    block_sparse_attn_256,
    block_sparse_attn_256_bshd,
 )
@@ -76,13 +74,12 @@ def video_sparse_attn(

    Dispatches the sparse branch by ``block_elements = prod(block_size)``:
    - 64  -> existing TK/Triton path (see ``block_sparse_attn_from_indices``).
-    - 128 -> Triton fallback or CuTe FA4 block-sparse attention.
    - 256 -> CuTe FA4 block-sparse attention (see ``block_sparse_attn_256``).

    Backend overrides:
    - ``FASTVIDEO_VSA_TRITON=1`` forces Triton in either path.
    - ``FASTVIDEO_VSA_TK=1`` prefers the sm_90 TK kernel in the 64-block path.
-    - ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 128/256-block paths.
+    - ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 256-block path.
    """
    if isinstance(block_size, int):
        block_size = (block_size, block_size, block_size)
@@ -122,9 +119,8 @@ def video_sparse_attn(
    # Sparse branch (fused Triton topk mask)
    mask = fused_topk_mask(scores, topk)

-    if block_elements in (128, 256):
-        attention = block_sparse_attn_128 if block_elements == 128 else block_sparse_attn_256
-        out_s = attention(q, k, v, mask, variable_block_sizes)[0]
+    if block_elements == 256:
+        out_s = block_sparse_attn_256(q, k, v, mask, variable_block_sizes)[0]
    else:
        out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]

@@ -146,14 +142,14 @@ def video_sparse_attn_bshd(
    """VSA entrypoint for [B, S, H, D] tensors.

    Avoids the BHSD<->BSHD round-trip that ``video_sparse_attn`` performs on
-    the CuTe 128/256-block paths; the 64-block path still expects BHSD and is not
+    the CuTe 256-block path; the 64-block path still expects BHSD and is not
    supported here.
    """
    if isinstance(block_size, int):
        block_size = (block_size, block_size, block_size)
    block_elements = block_size[0] * block_size[1] * block_size[2]
-    if block_elements not in (128, 256):
-        raise ValueError("video_sparse_attn_bshd is only defined for block_elements=128 or 256 "
+    if block_elements != 256:
+        raise ValueError("video_sparse_attn_bshd is only defined for block_elements=256 "
                         f"(got {block_elements}); use video_sparse_attn for the 64-block path.")

    batch, q_seq_len, heads, dim = q.shape
@@ -175,15 +171,19 @@ def video_sparse_attn_bshd(
        raise ValueError(f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
                         f"got {q_variable_block_sizes.numel()}")

-    # Compression branch (BSHD-native: match fused_block_mean's semantics).
-    # Padding values are expected to be zero; gradients are broadcast across
-    # the full padded block, just like the BHSD fused common path.
+    # Compression branch (BSHD-native: mean over the 256-token axis).
+    token_idx = torch.arange(block_elements, device=q.device, dtype=torch.int32)
+    q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1,
+                                                                         1)).view(1, q_num_blocks, block_elements, 1, 1)
+    kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1,
+                                                                        1)).view(1, kv_num_blocks, block_elements, 1, 1)
+
    q_c = q.view(batch, q_num_blocks, block_elements, heads, dim)
    k_c = k.view(batch, kv_num_blocks, block_elements, heads, dim)
    v_c = v.view(batch, kv_num_blocks, block_elements, heads, dim)
-    q_c = (q_c.float().sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
-    k_c = (k_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
-    v_c = (v_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
+    q_c = ((q_c.float() * q_token_valid).sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
+    k_c = ((k_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
+    v_c = ((v_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
    q_ch = q_c.permute(0, 2, 1, 3).contiguous()
    k_ch = k_c.permute(0, 2, 1, 3).contiguous()
    v_ch = v_c.permute(0, 2, 1, 3).contiguous()
@@ -195,15 +195,13 @@ def video_sparse_attn_bshd(

    # Sparse branch (fused Triton topk mask + CuTe BSHD).
    mask = fused_topk_mask(scores, topk)
-    attention = block_sparse_attn_128_bshd if block_elements == 128 else block_sparse_attn_256_bshd
-    out_s, _ = attention(q, k, v, mask, variable_block_sizes)
+    out_s, _ = block_sparse_attn_256_bshd(q, k, v, mask, variable_block_sizes)

-    # Out-of-place: ``out_s`` is the tensor FA4's autograd node saved for its
-    # backward, so mutating it in place invalidates the graph.
-    out_view = out_s.view(batch, q_num_blocks, block_elements, heads, dim)
+    out = out_s
+    out_view = out.view(batch, q_num_blocks, block_elements, heads, dim)
    if compress_attn_weight is not None:
        gate_view = compress_attn_weight.view(batch, q_num_blocks, block_elements, heads, dim)
-        out = out_view + out_c_blk.unsqueeze(2) * gate_view
+        out_view.add_(out_c_blk.unsqueeze(2) * gate_view)
    else:
-        out = out_view + out_c_blk.unsqueeze(2)
-    return out.view(batch, q_seq_len, heads, dim)
+        out_view.add_(out_c_blk.unsqueeze(2))
+    return out
@@ -237,12 +237,7 @@ def _attn_bwd_dkdv(
        # Load m before computing qk to reduce pipeline stall.
        offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
        m = tl.load(M + offs_m)
-        # Recompute logits exactly as the forward does: raw bf16 operands into
-        # the dot, fp32 scale after accumulation. A bf16 pre-scaled K perturbs
-        # the recomputed logits relative to the saved M by an error
-        # proportional to |logit|, which exp2 amplifies into arbitrarily wrong
-        # probabilities at large activations.
-        qkT = tl.dot(k, qT) * (sm_scale * 1.4426950408889634)
+        qkT = tl.dot(k, qT)
        pT = tl.math.exp2(qkT - m[None, :])
        mask = tl.arange(0, BLOCK_N1) < block_size
        pT = tl.where(mask[:, None], pT, 0.0)
@@ -273,7 +268,6 @@ def _attn_bwd_dq(
        do,
        m,
        D,
-        sm_scale,
        # shared by Q/K/V/DO.
        q2k_index,
        q2k_num,
@@ -321,7 +315,7 @@ def _attn_bwd_dq(
        block_sparse_offset = (kv_idx * 2 + half) * step_n * stride_tok
        kT = tl.load(kT_ptrs + block_sparse_offset)
        vT = tl.load(vT_ptrs + block_sparse_offset)
-        qk = tl.dot(q, kT) * (sm_scale * 1.4426950408889634)
+        qk = tl.dot(q, kT)
        p = tl.math.exp2(qk - m)
        offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
        mask = offs_in_block < block_size
@@ -330,7 +324,8 @@ def _attn_bwd_dq(
        dp = tl.dot(do, vT).to(tl.float32)
        ds = p * (dp - Di[:, None])
        ds = ds.to(tl.bfloat16)
-        # Compute dQ (kT is raw; the caller applies sm_scale once at the end).
+        # Compute dQ.
+        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
        dq += tl.dot(ds, tl.trans(kT))
        # Increment pointers.
    return dq
@@ -458,7 +453,6 @@ def _attn_bwd(
        do,
        m,
        D,  #
-        sm_scale,
        q2k_index,
        q2k_num,
        max_kv_blks,
@@ -476,7 +470,7 @@ def _attn_bwd(
    )
    # Write back dQ.
    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq *= sm_scale
+    dq *= LN2
    tl.store(dq_ptrs, dq)


@@ -597,7 +591,6 @@ def _attn_bwd_dq_kernel(
        Q,
        K,
        V,
-        sm_scale,
        DO,  #
        DQ,
        M,
@@ -670,7 +663,6 @@ def _attn_bwd_dq_kernel(
        do,
        m,
        D,
-        sm_scale,
        q2k_index,
        q2k_num,
        max_kv_blks,
@@ -688,7 +680,7 @@ def _attn_bwd_dq_kernel(
    )

    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq_acc *= sm_scale
+    dq_acc *= LN2
    tl.store(dq_ptrs, dq_acc)


@@ -756,11 +748,9 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
    dv = torch.empty_like(v)
    BATCH, N_HEAD = q.shape[:2]
    BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
-    # K stays raw: the backward kernels apply sm_scale in fp32 after the dot,
-    # matching the forward's rounding exactly. (A bf16 pre-scaled K perturbs
-    # the recomputed logits vs the saved M; exp2 turns that into unboundedly
-    # wrong probabilities at large activations.)
+    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
    arg_k = k
+    arg_k = arg_k * (sm_scale * RCP_LN2)
    PRE_BLOCK = 64
    assert Tq % PRE_BLOCK == 0
    pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
@@ -823,7 +813,6 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
        q,
        arg_k,
        v,
-        sm_scale,
        do,
        dq,
        M,
@@ -14,9 +14,7 @@ import math
 import torch

 VSA_TILE_SIZE = (4, 4, 4)
-# 128 is served by the sm_100a CUDA backend (block_sparse_attn_sm100a); 64 and 256 by
-# Triton and the CuTe-DSL path. A volume here only needs a backend that accepts it.
-_SUPPORTED_VSA_BLOCK_VOLUMES = (64, 128, 256)
+_SUPPORTED_VSA_BLOCK_VOLUMES = (64, 256)


 def _canonicalize_device(device: torch.device | str) -> torch.device:
@@ -0,0 +1,526 @@
+"""Correctness coverage for FastVideo's fine-grained FA4 sparse tiles."""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def _require_fa4_sm100(monkeypatch):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    major, _minor = torch.cuda.get_device_capability()
+    if major not in (10, 11):
+        pytest.skip("fine-grained FA4 VSA requires SM100/SM110")
+    pytest.importorskip(
+        "flash_attn.cute.block_sparsity",
+        reason="optional FastVideo FA4 source package is not installed",
+    )
+    monkeypatch.setenv("FASTVIDEO_FA4_VSA_DUAL_STREAM", "1")
+    monkeypatch.setenv("FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER", "1")
+
+
+def _selected_block_map(
+    heads: int,
+    q_blocks: int,
+    kv_blocks: int,
+    selected_blocks: int,
+    device: torch.device,
+) -> torch.Tensor:
+    block_map = torch.zeros(
+        1,
+        heads,
+        q_blocks,
+        kv_blocks,
+        dtype=torch.bool,
+        device=device,
+    )
+    for head in range(heads):
+        for q_block in range(q_blocks):
+            start = (3 * head + 5 * q_block) % kv_blocks
+            indices = [(start + 2 * offset) % kv_blocks for offset in range(selected_blocks)]
+            block_map[0, head, q_block, indices] = True
+    return block_map
+
+
+def _ordered_sparse_tensors(
+    block_map: torch.Tensor,
+    q_block_size: int,
+    kv_block_size: int,
+    masked_blocks: int,
+):
+    from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch
+
+    batch, heads, q_blocks, kv_blocks = block_map.shape
+    selected_blocks = int(block_map.sum(dim=-1).min().item())
+    assert bool((block_map.sum(dim=-1) == selected_blocks).all())
+    if not 0 <= masked_blocks <= selected_blocks:
+        raise ValueError("masked_blocks must be within the selected-block count")
+
+    selected = torch.arange(kv_blocks, dtype=torch.int32, device=block_map.device)
+    selected = selected.view(1, 1, 1, kv_blocks).expand(batch, heads, q_blocks, kv_blocks)
+    selected = selected.masked_select(block_map).view(batch, heads, q_blocks, selected_blocks)
+    selected = selected.sort(dim=-1).values
+
+    mask_count = masked_blocks
+    full_count = selected_blocks - masked_blocks
+    mask_idx = torch.zeros(
+        batch,
+        heads,
+        q_blocks,
+        max(mask_count, 1),
+        dtype=torch.int32,
+        device=block_map.device,
+    )
+    full_idx = torch.zeros(
+        batch,
+        heads,
+        q_blocks,
+        max(full_count, 1),
+        dtype=torch.int32,
+        device=block_map.device,
+    )
+    if mask_count:
+        mask_idx[..., :mask_count] = selected[..., :mask_count]
+    if full_count:
+        full_idx[..., :full_count] = selected[..., mask_count:]
+
+    return BlockSparseTensorsTorch(
+        mask_block_cnt=torch.full(
+            (batch, heads, q_blocks),
+            mask_count,
+            dtype=torch.int32,
+            device=block_map.device,
+        ),
+        mask_block_idx=mask_idx,
+        full_block_cnt=torch.full(
+            (batch, heads, q_blocks),
+            full_count,
+            dtype=torch.int32,
+            device=block_map.device,
+        ),
+        full_block_idx=full_idx,
+        block_size=(q_block_size, kv_block_size),
+    )
+
+
+def _torch_reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    q_block_size: int,
+    kv_block_size: int,
+    variable_block_sizes: torch.Tensor | None = None,
+) -> torch.Tensor:
+    token_mask = block_map.repeat_interleave(q_block_size, dim=2)
+    token_mask = token_mask.repeat_interleave(kv_block_size, dim=3)
+    if variable_block_sizes is not None:
+        kv_positions = torch.arange(k.shape[1], device=k.device)
+        kv_blocks = kv_positions // kv_block_size
+        kv_offsets = kv_positions % kv_block_size
+        valid_tokens = kv_offsets < variable_block_sizes[kv_blocks]
+        token_mask = token_mask & valid_tokens.view(1, 1, 1, -1)
+    q_heads = q.permute(0, 2, 1, 3).float()
+    k_heads = k.permute(0, 2, 1, 3).float()
+    v_heads = v.permute(0, 2, 1, 3).float()
+    scores = torch.matmul(q_heads, k_heads.transpose(-2, -1)) / math.sqrt(q.shape[-1])
+    scores.masked_fill_(~token_mask, float("-inf"))
+    out = torch.matmul(torch.softmax(scores, dim=-1), v_heads)
+    return out.permute(0, 2, 1, 3).contiguous()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize(
+    (
+        "q_block_size",
+        "kv_block_size",
+        "selected_blocks",
+        "masked_blocks",
+        "sp_double_buffer",
+    ),
+    [
+        pytest.param(128, 64, 1, 0, True, id="q128_kv64_one_stream_empty"),
+        pytest.param(128, 64, 2, 1, True, id="q128_kv64_even_mixed"),
+        pytest.param(128, 64, 3, 1, True, id="q128_kv64_odd_mixed"),
+        pytest.param(128, 64, 3, 0, True, id="q128_kv64_odd_full_only"),
+        pytest.param(128, 64, 3, 1, False, id="q128_kv64_legacy_optout"),
+        pytest.param(64, 64, 3, 1, True, id="q64_kv64_odd_mixed"),
+    ],
+)
+def test_fa4_vsa_fine_grained_forward(
+    q_block_size: int,
+    kv_block_size: int,
+    selected_blocks: int,
+    masked_blocks: int,
+    sp_double_buffer: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from flash_attn.cute.interface import _flash_attn_fwd
+
+    monkeypatch.setenv(
+        "FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER",
+        "1" if sp_double_buffer else "0",
+    )
+    torch.manual_seed(2026 + q_block_size + selected_blocks)
+    device = torch.device("cuda")
+    batch, heads, head_dim = 1, 8, 128
+    q_blocks, kv_blocks = 4, 7
+    q_len = q_blocks * q_block_size
+    kv_len = kv_blocks * kv_block_size
+    q = torch.randn(batch, q_len, heads, head_dim, dtype=torch.bfloat16, device=device)
+    k = torch.randn(batch, kv_len, heads, head_dim, dtype=torch.bfloat16, device=device)
+    v = torch.randn_like(k)
+    block_map = _selected_block_map(heads, q_blocks, kv_blocks, selected_blocks, device)
+    sparse_tensors = _ordered_sparse_tensors(
+        block_map,
+        q_block_size,
+        kv_block_size,
+        masked_blocks,
+    )
+
+    out = _flash_attn_fwd(
+        q,
+        k,
+        v,
+        tile_mn=(q_block_size, kv_block_size),
+        block_sparse_tensors=sparse_tensors,
+        causal=False,
+        return_lse=True,
+    )[0]
+    ref = _torch_reference(q, k, v, block_map, q_block_size, kv_block_size)
+
+    assert bool(torch.isfinite(out).all())
+    error = (out.float() - ref).abs()
+    avg_abs = float(error.mean())
+    max_rel = float(error.max() / (ref.abs().mean() + 1e-6))
+    print(
+        f"[fa4 {q_block_size}x{kv_block_size} selected={selected_blocks} "
+        f"masked={masked_blocks} spdb={sp_double_buffer}] "
+        f"avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}"
+    )
+    assert avg_abs < 1e-3
+    assert max_rel < 0.2
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("q_block_size", [128, 64], ids=["q128_kv64", "q64_kv64"])
+def test_fa4_vsa_dual_stream_persistent_phase_transitions(q_block_size: int) -> None:
+    """Exercise even, odd, empty, and one-block tiles in one persistent launch."""
+    if q_block_size == 64 and torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("the Q64/KV64 specialization currently requires SM100")
+
+    from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch
+    from flash_attn.cute.interface import _flash_attn_fwd
+
+    torch.manual_seed(4554 + q_block_size)
+    device = torch.device("cuda")
+    kv_block_size = 64
+    counts = [2, 3, 0, 1, 4]
+    masked_counts = [1, 1, 0, 1, 2]
+    batch, heads, head_dim, kv_blocks = 1, 4, 128, 8
+    q_blocks = len(counts)
+    q = torch.randn(
+        batch,
+        q_blocks * q_block_size,
+        heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    k = torch.randn(
+        batch,
+        kv_blocks * kv_block_size,
+        heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    v = torch.randn_like(k)
+
+    mask_count = torch.empty(batch, heads, q_blocks, dtype=torch.int32, device=device)
+    full_count = torch.empty_like(mask_count)
+    mask_idx = torch.zeros(
+        batch,
+        heads,
+        q_blocks,
+        max(masked_counts),
+        dtype=torch.int32,
+        device=device,
+    )
+    full_idx = torch.zeros(
+        batch,
+        heads,
+        q_blocks,
+        max(count - masked for count, masked in zip(counts, masked_counts)),
+        dtype=torch.int32,
+        device=device,
+    )
+    block_map = torch.zeros(
+        batch,
+        heads,
+        q_blocks,
+        kv_blocks,
+        dtype=torch.bool,
+        device=device,
+    )
+    for head in range(heads):
+        for q_block, (count, masked) in enumerate(zip(counts, masked_counts)):
+            chosen = sorted({(head + 3 * q_block + 2 * offset) % kv_blocks for offset in range(count)})
+            assert len(chosen) == count
+            mask_count[0, head, q_block] = masked
+            full_count[0, head, q_block] = count - masked
+            if masked:
+                mask_idx[0, head, q_block, :masked] = torch.tensor(
+                    chosen[:masked],
+                    dtype=torch.int32,
+                    device=device,
+                )
+            if count > masked:
+                full_idx[0, head, q_block, :count - masked] = torch.tensor(
+                    chosen[masked:],
+                    dtype=torch.int32,
+                    device=device,
+                )
+            if count:
+                block_map[0, head, q_block, chosen] = True
+
+    sparse_tensors = BlockSparseTensorsTorch(
+        mask_block_cnt=mask_count,
+        mask_block_idx=mask_idx,
+        full_block_cnt=full_count,
+        full_block_idx=full_idx,
+        block_size=(q_block_size, kv_block_size),
+    )
+    out, lse = _flash_attn_fwd(
+        q,
+        k,
+        v,
+        tile_mn=(q_block_size, kv_block_size),
+        block_sparse_tensors=sparse_tensors,
+        causal=False,
+        return_lse=True,
+    )[:2]
+    ref = _torch_reference(q, k, v, block_map, q_block_size, kv_block_size)
+
+    nonempty_rows = torch.tensor(
+        [count > 0 for count in counts],
+        dtype=torch.bool,
+        device=device,
+    ).repeat_interleave(q_block_size)
+    error = (out[:, nonempty_rows].float() - ref[:, nonempty_rows]).abs()
+    empty_out = out[:, ~nonempty_rows].float()
+    empty_lse = lse[:, :, ~nonempty_rows]
+    avg_abs = float(error.mean())
+    max_abs = float(error.max())
+    print(
+        f"[fa4 persistent {q_block_size}x{kv_block_size}] "
+        f"avg_abs={avg_abs:.6e}, max_abs={max_abs:.6e}"
+    )
+    assert avg_abs < 1e-3
+    assert max_abs < 1e-2
+    assert float(empty_out.abs().max()) == 0.0
+    assert bool(torch.isneginf(empty_lse).all())
+
+
+@pytest.mark.cuda
+def test_fa4_vsa_q128_kv64_bshd_wrapper_variable_blocks() -> None:
+    """Cover partial and empty KV blocks through FastVideo's VBS mask mod."""
+    from fastvideo_kernel.block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
+
+    torch.manual_seed(5128)
+    device = torch.device("cuda")
+    batch, heads, head_dim = 1, 8, 128
+    q_block_size, kv_block_size = 128, 64
+    q_blocks, kv_blocks, selected_blocks = 4, 7, 3
+    q = torch.randn(
+        batch,
+        q_blocks * q_block_size,
+        heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    k = torch.randn(
+        batch,
+        kv_blocks * kv_block_size,
+        heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    v = torch.randn_like(k)
+    block_map = _selected_block_map(heads, q_blocks, kv_blocks, selected_blocks, device)
+    variable_block_sizes = torch.tensor(
+        [64, 23, 64, 7, 0, 51, 64],
+        dtype=torch.int32,
+        device=device,
+    )
+
+    out, lse = block_sparse_attn_cute_fwd_bshd(
+        q,
+        k,
+        v,
+        block_map,
+        variable_block_sizes,
+    )
+    ref = _torch_reference(
+        q,
+        k,
+        v,
+        block_map,
+        q_block_size,
+        kv_block_size,
+        variable_block_sizes,
+    )
+
+    assert bool(torch.isfinite(out).all())
+    assert bool(torch.isfinite(lse).all())
+    error = (out.float() - ref).abs()
+    avg_abs = float(error.mean())
+    max_rel = float(error.max() / (ref.abs().mean() + 1e-6))
+    print(f"[fa4 wrapper Q128/KV64 VBS] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
+    assert avg_abs < 1e-3
+    assert max_rel < 0.2
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize(
+    ("q_block_size", "kv_block_size"),
+    [
+        pytest.param(128, 64, id="q128_kv64"),
+        pytest.param(64, 64, id="q64_kv64"),
+    ],
+)
+def test_fa4_vsa_fine_grained_bshd_wrapper(
+    q_block_size: int,
+    kv_block_size: int,
+) -> None:
+    """Exercise the FastVideo BSHD entrypoint used by the H3 backend."""
+    if q_block_size == 64 and torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("the Q64/KV64 specialization currently requires SM100")
+
+    from fastvideo_kernel.block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
+
+    torch.manual_seed(4048 + q_block_size)
+    device = torch.device("cuda")
+    batch, heads, head_dim = 1, 8, 128
+    q_blocks, kv_blocks, selected_blocks = 4, 7, 3
+    q = torch.randn(
+        batch,
+        q_blocks * q_block_size,
+        heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    k = torch.randn(
+        batch,
+        kv_blocks * kv_block_size,
+        heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    v = torch.randn_like(k)
+    block_map = _selected_block_map(heads, q_blocks, kv_blocks, selected_blocks, device)
+    variable_block_sizes = torch.full(
+        (kv_blocks,),
+        kv_block_size,
+        dtype=torch.int32,
+        device=device,
+    )
+
+    out, lse = block_sparse_attn_cute_fwd_bshd(
+        q,
+        k,
+        v,
+        block_map,
+        variable_block_sizes,
+    )
+    ref = _torch_reference(q, k, v, block_map, q_block_size, kv_block_size)
+
+    assert bool(torch.isfinite(out).all())
+    assert bool(torch.isfinite(lse).all())
+    error = (out.float() - ref).abs()
+    avg_abs = float(error.mean())
+    max_rel = float(error.max() / (ref.abs().mean() + 1e-6))
+    print(
+        f"[fa4 wrapper {q_block_size}x{kv_block_size}] "
+        f"avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}"
+    )
+    assert avg_abs < 1e-3
+    assert max_rel < 0.2
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("block_shape", ["256x256", "128x64", "64x64"])
+def test_vsa256_bshd_fa4_block_shape_routes(
+    block_shape: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Route one logical Q256/KV256 mask through every FA4 specialization."""
+    if block_shape == "64x64" and torch.cuda.get_device_capability()[0] != 10:
+        pytest.skip("the Q64/KV64 specialization currently requires SM100")
+
+    from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_256_bshd
+
+    monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
+    monkeypatch.setenv("FASTVIDEO_VSA_FA4_BLOCK_SHAPE", block_shape)
+    torch.manual_seed(6256)
+    device = torch.device("cuda")
+    batch, heads, head_dim = 1, 8, 128
+    logical_block_size = 256
+    q_blocks, kv_blocks, selected_blocks = 4, 7, 3
+    q = torch.randn(
+        batch,
+        q_blocks * logical_block_size,
+        heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    k = torch.randn(
+        batch,
+        kv_blocks * logical_block_size,
+        heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    v = torch.randn_like(k)
+    block_map = _selected_block_map(heads, q_blocks, kv_blocks, selected_blocks, device)
+    variable_block_sizes = torch.tensor(
+        [256, 137, 255, 64, 201, 1, 192],
+        dtype=torch.int32,
+        device=device,
+    )
+
+    out, lse = block_sparse_attn_256_bshd(
+        q,
+        k,
+        v,
+        block_map,
+        variable_block_sizes,
+    )
+    ref = _torch_reference(
+        q,
+        k,
+        v,
+        block_map,
+        logical_block_size,
+        logical_block_size,
+        variable_block_sizes,
+    )
+
+    assert bool(torch.isfinite(out).all())
+    assert bool(torch.isfinite(lse).all())
+    error = (out.float() - ref).abs()
+    avg_abs = float(error.mean())
+    max_rel = float(error.max() / (ref.abs().mean() + 1e-6))
+    print(f"[VSA256 route {block_shape}] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
+    assert avg_abs < 1e-3
+    assert max_rel < 0.2
--- a/Show More
+++ b/Show More