luminal_python: suppress torch.export _guards_fn instead of disabling auto-dynamic shapes

Replaces the WIP `automatic_dynamic_shapes = False` workaround (commit
3a3cd049) with a targeted monkey-patch of `torch.export._unlift.
_ok_to_generate_guards_fn`. That function already supports a call-stack
opt-out (used by executorch / modai / on_device_ai / torchao); we extend
it with a "luminal" check so torch.export skips inserting the
`_guards_fn` submodule whenever luminal is the embedder.

Why the previous workaround was costly: with `automatic_dynamic_shapes
= False`, the bench loop's `compiled(input_ids, cache_position=tensor([k]))`
recompiles once per `cache_position` *value*, i.e. one full luminal
compile per generated token. gemma3-4b smoke = ~2 hr CPU + 200 GB host
RSS. The L NameError it was working around fires during
aot_autograd's fx.Interpreter trace of a re-exported GraphModule that
contains the L-referencing `_guards_fn` body — a dead-end for any
non-dynamo consumer of the exported graph.

Skipping `_guards_fn` generation at the source restores the
compile-once-run-many behaviour of dynamic-shape promotion: dynamo
promotes the varying dim to a SymInt on the second compile and reuses
the same compiled graph for all subsequent values.

The monkey-patch is scoped to luminal's call stack — other consumers
of `torch.export` in the same Python process see unmodified behaviour.

Verified via a multi-shape compile smoke (`compiled(rand(4,8))` then
`compiled(rand(5,8))`): no L NameError. The remaining downstream
`SymInt` input passthrough is handled by `_specialize_sym_scalar` in
pt2.py and is unrelated to this fix.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

.github/workflows/cuda-clippy.yml

@@ -0,0 +1,30 @@
+name: CUDA Clippy
+
+on:
+    push:
+        branches: ["main"]
+    pull_request:
+        branches: ["main"]
+    workflow_dispatch:
+
+jobs:
+    cuda_clippy:
+        name: CUDA Clippy
+        runs-on: cuda_t4_runner
+        container:
+            image: ghcr.io/luminal-ai/luminal-docker:cuda
+            options: --gpus all
+        timeout-minutes: 20
+
+        steps:
+            - uses: actions/checkout@v6
+            - name: Mark workspace as safe for git
+              run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
+            - uses: actions/setup-python@v5
+              with:
+                  python-version: "3.11"
+            - name: Update Rust toolchain
+              run: rustup update
+            - uses: pre-commit/action@v3.0.1
+              with:
+                  extra_args: cargo-clippy --all-files

.github/workflows/fmt.yml

@@ -0,0 +1,23 @@
+name: Fmt
+
+on:
+    push:
+        branches: ["main"]
+    pull_request:
+        branches: ["main"]
+    workflow_dispatch:
+
+jobs:
+    fmt:
+        name: Fmt
+        runs-on: ubuntu-latest
+        timeout-minutes: 20
+
+        steps:
+            - uses: actions/checkout@v6
+            - uses: actions/setup-python@v5
+              with:
+                  python-version: "3.11"
+            - uses: pre-commit/action@v3.0.1
+              with:
+                  extra_args: cargo-fmt --all-files

.github/workflows/lint.yml

@@ -1,86 +0,0 @@
-name: Lint
-
-on:
-    push:
-        branches: ["main"]
-    pull_request:
-        branches: ["main"]
-    workflow_dispatch:
-
-env:
-    CARGO_TERM_COLOR: always
-
-jobs:
-    ruff:
-        name: Ruff
-        runs-on: ubuntu-latest
-        timeout-minutes: 20
-
-        steps:
-            - uses: actions/checkout@v6
-            - uses: actions/setup-python@v5
-              with:
-                  python-version: "3.11"
-            - uses: pre-commit/action@v3.0.1
-              with:
-                  extra_args: ruff-check --all-files
-
-    ruff_format:
-        name: Ruff Format
-        runs-on: ubuntu-latest
-        timeout-minutes: 20
-
-        steps:
-            - uses: actions/checkout@v6
-            - uses: actions/setup-python@v5
-              with:
-                  python-version: "3.11"
-            - uses: pre-commit/action@v3.0.1
-              with:
-                  extra_args: ruff-format --all-files
-
-    clippy:
-        name: Clippy
-        runs-on: ubuntu-latest
-        timeout-minutes: 20
-
-        steps:
-            - uses: actions/checkout@v6
-            - uses: actions/setup-python@v5
-              with:
-                  python-version: "3.11"
-            - name: Update Rust toolchain
-              run: rustup update
-            - uses: pre-commit/action@v3.0.1
-              with:
-                  extra_args: cargo-clippy --all-files
-
-    metal_clippy:
-        name: Metal Clippy
-        runs-on: macos-14
-        timeout-minutes: 30
-
-        steps:
-            - uses: actions/checkout@v6
-            - uses: actions/setup-python@v5
-              with:
-                  python-version: "3.11"
-            - name: Update Rust toolchain
-              run: rustup update
-            - uses: pre-commit/action@v3.0.1
-              with:
-                  extra_args: --hook-stage manual cargo-clippy-metal --all-files
-
-    fmt:
-        name: Fmt
-        runs-on: ubuntu-latest
-        timeout-minutes: 20
-
-        steps:
-            - uses: actions/checkout@v6
-            - uses: actions/setup-python@v5
-              with:
-                  python-version: "3.11"
-            - uses: pre-commit/action@v3.0.1
-              with:
-                  extra_args: cargo-fmt --all-files

.github/workflows/metal-clippy.yml

@@ -0,0 +1,25 @@
+name: Metal Clippy
+
+on:
+    push:
+        branches: ["main"]
+    pull_request:
+        branches: ["main"]
+    workflow_dispatch:
+
+jobs:
+    metal_clippy:
+        name: Metal Clippy
+        runs-on: macos-14
+        timeout-minutes: 30
+
+        steps:
+            - uses: actions/checkout@v6
+            - uses: actions/setup-python@v5
+              with:
+                  python-version: "3.11"
+            - name: Update Rust toolchain
+              run: rustup update
+            - uses: pre-commit/action@v3.0.1
+              with:
+                  extra_args: --hook-stage manual cargo-clippy-metal --all-files

.github/workflows/modal-examples.yml

@@ -3,15 +3,18 @@ name: Modal Examples
 on:
     push:
         branches: ["main"]
-    pull_request:
+    pull_request_target:
         branches: ["main"]
-        types: [opened, synchronize, reopened, ready_for_review]
+        types: [labeled, synchronize]
     workflow_dispatch:
 
 jobs:
     modal_example:
-        # Keep the draft check PR-specific so push/manual runs still execute.
-        if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}
+        if: >-
+            github.event_name == 'push'
+            || github.event_name == 'workflow_dispatch'
+            || (github.event_name == 'pull_request_target'
+                && contains(github.event.pull_request.labels.*.name, 'modal-ready'))
         name: "${{ matrix.example }} (Modal ${{ matrix.gpu.type }})"
         runs-on: ubuntu-latest
         environment: Modal
@@ -27,6 +30,8 @@ jobs:
 
         steps:
             - uses: actions/checkout@v6
+              with:
+                  ref: ${{ github.event.pull_request.head.sha || github.sha }}
             - name: Set up Python
               uses: actions/setup-python@v5
               with:

.github/workflows/ruff-format.yml

@@ -0,0 +1,23 @@
+name: Ruff Format
+
+on:
+    push:
+        branches: ["main"]
+    pull_request:
+        branches: ["main"]
+    workflow_dispatch:
+
+jobs:
+    ruff_format:
+        name: Ruff Format
+        runs-on: ubuntu-latest
+        timeout-minutes: 20
+
+        steps:
+            - uses: actions/checkout@v6
+            - uses: actions/setup-python@v5
+              with:
+                  python-version: "3.11"
+            - uses: pre-commit/action@v3.0.1
+              with:
+                  extra_args: ruff-format --all-files

.github/workflows/ruff.yml

@@ -0,0 +1,23 @@
+name: Ruff
+
+on:
+    push:
+        branches: ["main"]
+    pull_request:
+        branches: ["main"]
+    workflow_dispatch:
+
+jobs:
+    ruff:
+        name: Ruff
+        runs-on: ubuntu-latest
+        timeout-minutes: 20
+
+        steps:
+            - uses: actions/checkout@v6
+            - uses: actions/setup-python@v5
+              with:
+                  python-version: "3.11"
+            - uses: pre-commit/action@v3.0.1
+              with:
+                  extra_args: ruff-check --all-files

.github/workflows/test-core.yml

@@ -0,0 +1,24 @@
+name: Test Core
+
+on:
+    push:
+        branches: ["main"]
+    pull_request:
+        branches: ["main"]
+    workflow_dispatch:
+
+env:
+    CARGO_TERM_COLOR: always
+
+jobs:
+    core_unit_test:
+        name: Core Unit Tests
+        runs-on: ubuntu-latest
+        container:
+            image: ghcr.io/luminal-ai/luminal-docker:cpu
+        timeout-minutes: 20
+
+        steps:
+            - uses: actions/checkout@v6
+            - name: Run tests
+              run: cargo test --workspace --exclude luminal_cuda_lite --exclude luminal_metal --exclude luminal_bench --verbose

.github/workflows/test-cuda.yml

@@ -3,46 +3,35 @@ name: Test CUDA
 on:
     push:
         branches: ["main"]
-    pull_request:
+    pull_request_target:
         branches: ["main"]
+        types: [labeled, synchronize]
     workflow_dispatch:
 
-env:
-    CARGO_TERM_COLOR: always
-
 jobs:
-    cuda_clippy:
-        name: Cuda Clippy
-        runs-on: cuda_t4_runner
-        container:
-            image: ghcr.io/luminal-ai/luminal-docker:cuda
-            options: --gpus all
+    cuda_unit_test:
+        if: >-
+            github.event_name == 'push'
+            || github.event_name == 'workflow_dispatch'
+            || (github.event_name == 'pull_request_target'
+                && contains(github.event.pull_request.labels.*.name, 'modal-ready'))
+        name: Cuda Unit Tests
+        runs-on: ubuntu-latest
+        environment: Modal
         timeout-minutes: 30
 
         steps:
             - uses: actions/checkout@v6
-            - name: Mark workspace as a safe git directory
-              run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
-            - uses: actions/setup-python@v5
+              with:
+                  ref: ${{ github.event.pull_request.head.sha || github.sha }}
+            - name: Set up Python
+              uses: actions/setup-python@v5
               with:
                   python-version: "3.11"
-            - uses: pre-commit/action@v3.0.1
-              with:
-                  extra_args: --hook-stage manual cargo-clippy-cuda-lite --all-files
-
-    cuda_unit_test:
-        name: Cuda Unit Tests
-        runs-on: cuda_t4_runner
-        container:
-            image: ghcr.io/luminal-ai/luminal-docker:cuda
-            options: --gpus all
-        timeout-minutes: 30
-
-        steps:
-            - uses: actions/checkout@v6
-            - name: Detect GPU compute capability
-              run: |
-                  CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d '.')
-                  echo "CUDA_COMPUTE_CAP=${CAP}" >> "$GITHUB_ENV"
-            - name: Run CUDA crate tests
-              run: cargo test -p luminal_cuda_lite --verbose -- --test-threads=1
+            - name: Install Modal
+              run: pip install modal
+            - name: Run CUDA tests on Modal
+              env:
+                  MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+                  MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+              run: modal run ci/modal_cargo_test.py

.github/workflows/test-metal.yml

@@ -0,0 +1,19 @@
+name: Test Metal
+
+on:
+    push:
+        branches: ["main"]
+    pull_request:
+        branches: ["main"]
+    workflow_dispatch:
+
+jobs:
+    metal_unit_test:
+        name: Metal Unit Tests
+        runs-on: macos-14
+        timeout-minutes: 30
+
+        steps:
+            - uses: actions/checkout@v6
+            - name: Run Metal crate tests
+              run: rustup update; cargo test -p luminal_metal --verbose -- --test-threads=1

.github/workflows/test-python-cuda.yml

@@ -1,56 +1,20 @@
-name: Test
+name: Test Python CUDA
 
 on:
     push:
         branches: ["main"]
-    pull_request:
+    pull_request_target:
         branches: ["main"]
+        types: [labeled, synchronize]
     workflow_dispatch:
 
-env:
-    CARGO_TERM_COLOR: always
-
 jobs:
-    core_unit_test:
-        name: Core Unit Tests
-        runs-on: ubuntu-latest
-        container:
-            image: ghcr.io/luminal-ai/luminal-docker:cpu
-        timeout-minutes: 20
-
-        steps:
-            - uses: actions/checkout@v6
-            - name: Run tests
-              run: cargo test --workspace --exclude luminal_cuda_lite --exclude luminal_metal --exclude luminal_bench --verbose
-    metal_unit_test:
-        name: Metal Unit Tests
-        runs-on: macos-14
-        timeout-minutes: 30
-
-        steps:
-            - uses: actions/checkout@v6
-            - name: Run Metal crate tests
-              run: rustup update; cargo test -p luminal_metal --verbose -- --test-threads=1
-    python_native_tests:
-        name: Python Native Tests
-        runs-on: ubuntu-latest
-        container:
-            image: ghcr.io/luminal-ai/luminal-docker:cpu
-        timeout-minutes: 45
-        defaults:
-            run:
-                working-directory: crates/luminal_python
-
-        steps:
-            - uses: actions/checkout@v6
-            - name: Update Rust toolchain
-              run: rustup update
-            - name: Build maturin extension
-              run: uv run maturin develop --manifest-path rust/Cargo.toml
-            - name: Run pytest
-              run: uv run pytest tests/test_hlir_ops.py tests/test_unary.py -v -m "not slow"
-
     python_cuda_tests:
+        if: >-
+            github.event_name == 'push'
+            || github.event_name == 'workflow_dispatch'
+            || (github.event_name == 'pull_request_target'
+                && contains(github.event.pull_request.labels.*.name, 'modal-ready'))
         name: Python CUDA Tests
         runs-on: ubuntu-latest
         environment: Modal
@@ -61,6 +25,8 @@ jobs:
 
         steps:
             - uses: actions/checkout@v6
+              with:
+                  ref: ${{ github.event.pull_request.head.sha || github.sha }}
             - name: Set up Python
               uses: actions/setup-python@v5
               with:

.github/workflows/test-python-native.yml

@@ -0,0 +1,28 @@
+name: Test Python Native
+
+on:
+    push:
+        branches: ["main"]
+    pull_request:
+        branches: ["main"]
+    workflow_dispatch:
+
+jobs:
+    python_native_tests:
+        name: Python Native Tests
+        runs-on: ubuntu-latest
+        container:
+            image: ghcr.io/luminal-ai/luminal-docker:cpu
+        timeout-minutes: 45
+        defaults:
+            run:
+                working-directory: crates/luminal_python
+
+        steps:
+            - uses: actions/checkout@v6
+            - name: Update Rust toolchain
+              run: rustup update
+            - name: Build maturin extension
+              run: uv run maturin develop --manifest-path rust/Cargo.toml
+            - name: Run pytest
+              run: uv run pytest tests/test_hlir_ops.py tests/test_unary.py -v -m "not slow"

.gitignore

@@ -37,3 +37,9 @@ __pycache__/
 dist/
 build/
 uv.lock
+
+# TTFT benchmark SQLite database (per-machine state)
+benchmarks/ttft/bench.db
+benchmarks/ttft/bench.db-journal
+benchmarks/ttft/bench.db-wal
+benchmarks/ttft/bench.db-shm

Cargo.toml

@@ -32,6 +32,7 @@ pretty-duration = "0.1.1"
 anyhow = "1.0"
 graphviz-rust = { version = "0.9", default-features = false}
 lru = "0.16.2"
+rayon = "1.10"
 
 [workspace.package]
 edition = "2024"

README.md

@@ -1,10 +1,10 @@
-<img href="luminal.com" alt="Screenshot 2025-08-14 at 9 18 54 PM" src="https://github.com/user-attachments/assets/c5832634-55d5-45b7-ba65-6efe36afce4a" />
+<img href="luminal.com" alt="Screenshot 2025-08-14 at 9 18 54 PM" src="https://github.com/luminal-ai/luminal/blob/main/docs/logo/inference_at_the_speed_of_light.png" />
 
 <h3 align="center">
 Luminal is a high-performance general-purpose inference compiler.
 </h3>
 
-[![CI Status](https://img.shields.io/github/actions/workflow/status/jafioti/luminal/test.yml?style=for-the-badge&logo=github-actions&logoColor=white&branch=main)](https://github.com/jafioti/luminal/actions)
+[![CI Status](https://img.shields.io/github/actions/workflow/status/luminal-ai/luminal/test-core.yml?style=for-the-badge&logo=github-actions&logoColor=white&branch=main)](https://github.com/luminal-ai/luminal/actions)
 [![Docs](https://img.shields.io/badge/Documentation-green?style=for-the-badge&color=0D9373)](https://docs.luminalai.com)
 [![Current Crates.io Version](https://img.shields.io/crates/v/luminal.svg?style=for-the-badge&logo=rust)](https://crates.io/crates/luminal)
 [![discord](https://dcbadge.limes.pink/api/server/APjuwHAbGy)](https://discord.gg/APjuwHAbGy)

benchmarks/ttft/bench_python_baseline.py

@@ -0,0 +1,117 @@
+"""Pure HuggingFace/PyTorch TTFT + TPOT bench. Prints a JSON line on stdout.
+
+Measures:
+  TTFT — sum of single-token forward-pass durations over the prompt, using
+         a StaticCache. Methodology matches bench_python_luminal.py and the
+         rust path so the cross-path comparison is apples-to-apples.
+  TPOT — average time per output token during KV-cache greedy decode.
+"""
+
+import argparse
+import json
+import statistics
+import time
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.cache_utils import StaticCache
+
+from bench_utils import encode_prompt, measure_tpot, static_cache_config
+
+DEFAULT_MODEL = "NousResearch/Meta-Llama-3-8B-Instruct"
+DEFAULT_PROMPT = "Explain what a neural network is in a paragraph."
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default=DEFAULT_MODEL)
+    ap.add_argument("--prompt", default=DEFAULT_PROMPT)
+    ap.add_argument("--warmups", type=int, default=1)
+    ap.add_argument("--iters", type=int, default=3)
+    ap.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    ap.add_argument("--decode-tokens", type=int, default=50,
+                    help="Number of tokens to generate for TPOT measurement (0 = skip).")
+    ap.add_argument("--max-cache-len", type=int, default=256,
+                    help="StaticCache max sequence length.")
+    args = ap.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    input_ids = encode_prompt(tokenizer, args.prompt, device)
+    prompt_tokens = int(input_ids.shape[-1])
+
+    config = AutoConfig.from_pretrained(args.model)
+    config._attn_implementation = "eager"
+
+    model = (
+        AutoModelForCausalLM.from_pretrained(args.model, config=config, torch_dtype=dtype)
+        .eval()
+        .to(device)
+    )
+
+    single_token = torch.zeros(1, 1, dtype=torch.long, device=device)
+
+    cache_config = static_cache_config(config)
+
+    def make_cache():
+        return StaticCache(
+            config=cache_config,
+            max_batch_size=1,
+            max_cache_len=args.max_cache_len,
+            device=device,
+            dtype=dtype,
+        )
+
+    def measure_ttft() -> float:
+        """Sum of per-token forward-pass durations over prompt_tokens steps."""
+        kv = make_cache()
+        # Eager init at position 0 to satisfy StaticCache.lazy_initialization.
+        with torch.no_grad():
+            model(single_token, past_key_values=kv,
+                  cache_position=torch.tensor([0], device=device))
+        total_ms = 0.0
+        for pos in range(1, prompt_tokens):
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            with torch.no_grad():
+                model(single_token, past_key_values=kv,
+                      cache_position=torch.tensor([pos], device=device))
+                if device.type == "cuda":
+                    torch.cuda.synchronize()
+            total_ms += (time.perf_counter() - t0) * 1000.0
+        return total_ms
+
+    for _ in range(args.warmups):
+        measure_ttft()
+
+    ttft_samples_ms = [measure_ttft() for _ in range(args.iters)]
+
+    result = {
+        "path": "python_baseline",
+        "model": args.model,
+        "device": str(device),
+        "dtype": args.dtype,
+        "prompt_tokens": prompt_tokens,
+        "iters": args.iters,
+        "ttft_ms": statistics.median(ttft_samples_ms),
+        "ttft_ms_mean": sum(ttft_samples_ms) / len(ttft_samples_ms),
+        "ttft_ms_samples": ttft_samples_ms,
+        "note": "sequential per-token, StaticCache KV cache",
+    }
+
+    if args.decode_tokens > 0:
+        tpot_samples_ms = measure_tpot(model, input_ids, device, args.decode_tokens)
+        tpot_ms = sum(tpot_samples_ms) / len(tpot_samples_ms)
+        result["decode_tokens"] = args.decode_tokens
+        result["tpot_ms"] = tpot_ms
+        result["tpot_ms_samples"] = tpot_samples_ms
+        result["throughput_tps"] = 1000.0 / tpot_ms
+
+    print("BENCH_RESULT " + json.dumps(result))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ttft/bench_python_luminal.py

@@ -0,0 +1,196 @@
+"""Python -> Luminal TTFT + TPOT bench via torch.compile(backend=luminal_backend).
+
+Methodology mirrors examples/llama (the Rust path):
+  - One eager prefill step initialises the StaticCache (required by transformers'
+    StaticCache.lazy_initialization) before compilation.
+  - TTFT: run one forward pass per prompt token sequentially, each advancing
+    cache_position by 1; sum durations.
+  - TPOT: run --decode-tokens more single-token passes; average durations.
+  - StaticCache pre-allocates K/V buffers up to max_cache_len; no growing allocation.
+
+Prints a BENCH_RESULT JSON line on stdout.
+"""
+
+import argparse
+import gc
+import json
+import statistics
+import time
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.cache_utils import StaticCache
+
+from bench_utils import encode_prompt, static_cache_config
+from luminal import luminal_backend
+
+DEFAULT_MODEL = "NousResearch/Meta-Llama-3-8B-Instruct"
+DEFAULT_PROMPT = "Explain what a neural network is in a paragraph."
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default=DEFAULT_MODEL)
+    ap.add_argument("--prompt", default=DEFAULT_PROMPT)
+    ap.add_argument("--warmups", type=int, default=1)
+    ap.add_argument("--iters", type=int, default=3)
+    ap.add_argument(
+        "--search-iters",
+        type=int,
+        default=500,
+        help="Egraph search iterations (matches examples/llama default of 500).",
+    )
+    ap.add_argument(
+        "--decode-tokens",
+        type=int,
+        default=50,
+        help="Tokens to generate for TPOT measurement (0 = skip TPOT).",
+    )
+    ap.add_argument(
+        "--max-cache-len",
+        type=int,
+        default=256,
+        help="StaticCache max sequence length.",
+    )
+    ap.add_argument(
+        "--dtype",
+        default="float32",
+        choices=["float32", "bfloat16", "float16"],
+        help="Torch dtype for model + StaticCache.",
+    )
+    args = ap.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    input_ids = encode_prompt(tokenizer, args.prompt, device)
+    prompt_tokens = int(input_ids.shape[-1])
+
+    config = AutoConfig.from_pretrained(args.model)
+    config._attn_implementation = "eager"
+
+    model = (
+        AutoModelForCausalLM.from_pretrained(args.model, config=config, torch_dtype=dtype)
+        .eval()
+        .to(device)
+    )
+
+    single_token = torch.zeros(1, 1, dtype=torch.long, device=device)
+
+    cache_config = static_cache_config(config)
+
+    def make_cache():
+        return StaticCache(
+            config=cache_config,
+            max_batch_size=1,
+            max_cache_len=args.max_cache_len,
+            device=device,
+            dtype=dtype,
+        )
+
+    # Step 0: run ONE eager prefill to initialise the cache tensors and call
+    # mark_static_address (required by transformers' StaticCache before compile).
+    cache = make_cache()
+    with torch.no_grad():
+        model(single_token, past_key_values=cache, cache_position=torch.tensor([0], device=device))
+
+    # Compile for a single-token input — same graph is reused for every step.
+    # Compilation happens on the first call after the eager init above.
+    t0 = time.perf_counter()
+    compiled = torch.compile(
+        model,
+        backend=luminal_backend,
+        options={"search_iterations": args.search_iters},
+    )
+    cache_position = torch.tensor([1], dtype=torch.long, device=device)
+    with torch.no_grad():
+        compiled(single_token, past_key_values=cache, cache_position=cache_position)
+        if device.type == "cuda":
+            torch.cuda.synchronize()
+    compile_ms = (time.perf_counter() - t0) * 1000.0
+
+    gc.collect()
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+
+    def one_step(pos: int, kv_cache):
+        cache_pos = torch.tensor([pos], dtype=torch.long, device=device)
+        with torch.no_grad():
+            compiled(single_token, past_key_values=kv_cache, cache_position=cache_pos)
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+
+    def measure_ttft():
+        """Sum of per-token forward-pass durations over prompt_tokens steps.
+
+        Uses a fresh cache so each TTFT measurement is independent.
+        """
+        kv = make_cache()
+        # Eager init for this fresh cache (required before compiled can run on it).
+        with torch.no_grad():
+            model(single_token, past_key_values=kv, cache_position=torch.tensor([0], device=device))
+        total_ms = 0.0
+        # Step 0 was the eager init above; measure from step 1 to prompt_tokens.
+        for pos in range(1, prompt_tokens):
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            one_step(pos, kv)
+            total_ms += (time.perf_counter() - t0) * 1000.0
+        return total_ms
+
+    def measure_tpot(n, start_pos: int):
+        """Average single-token forward-pass duration over n decode steps."""
+        kv = make_cache()
+        # Eager init
+        with torch.no_grad():
+            model(single_token, past_key_values=kv, cache_position=torch.tensor([0], device=device))
+        # One warmup step.
+        one_step(1, kv)
+        step_times_ms = []
+        for i in range(n):
+            pos = start_pos + i
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            one_step(pos, kv)
+            step_times_ms.append((time.perf_counter() - t0) * 1000.0)
+        return step_times_ms
+
+    # Warmups before timing TTFT (all run after compilation is complete).
+    for _ in range(args.warmups):
+        measure_ttft()
+
+    ttft_samples_ms = [measure_ttft() for _ in range(args.iters)]
+
+    tpot_ms_samples = []
+    if args.decode_tokens > 0:
+        tpot_ms_samples = measure_tpot(args.decode_tokens, start_pos=prompt_tokens)
+
+    tpot_ms = sum(tpot_ms_samples) / len(tpot_ms_samples) if tpot_ms_samples else None
+    throughput_tps = (1000.0 / tpot_ms) if tpot_ms else None
+
+    result = {
+        "path": "python_luminal",
+        "model": args.model,
+        "device": str(device),
+        "dtype": args.dtype,
+        "prompt_tokens": prompt_tokens,
+        "iters": args.iters,
+        "ttft_ms": statistics.median(ttft_samples_ms),
+        "ttft_ms_mean": sum(ttft_samples_ms) / len(ttft_samples_ms),
+        "ttft_ms_samples": ttft_samples_ms,
+        "compile_ms": compile_ms,
+        "search_iters": args.search_iters,
+        "decode_tokens": args.decode_tokens if args.decode_tokens > 0 else None,
+        "tpot_ms": tpot_ms,
+        "tpot_ms_samples": tpot_ms_samples,
+        "throughput_tps": throughput_tps,
+        "note": "sequential per-token, StaticCache KV cache",
+    }
+    print("BENCH_RESULT " + json.dumps(result))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ttft/bench_python_torch_compile.py

@@ -0,0 +1,138 @@
+"""Vanilla torch.compile TTFT + TPOT bench. Prints a JSON line on stdout.
+
+Uses the default inductor backend (torch.compile without a custom backend).
+TTFT uses sequential per-token prefill with a StaticCache so the methodology
+matches bench_python_baseline.py, bench_python_luminal.py, and the rust path.
+"""
+
+import argparse
+import json
+import statistics
+import time
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.cache_utils import StaticCache
+
+from bench_utils import encode_prompt, measure_tpot, static_cache_config
+
+DEFAULT_MODEL = "NousResearch/Meta-Llama-3-8B-Instruct"
+DEFAULT_PROMPT = "Explain what a neural network is in a paragraph."
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default=DEFAULT_MODEL)
+    ap.add_argument("--prompt", default=DEFAULT_PROMPT)
+    ap.add_argument("--warmups", type=int, default=1)
+    ap.add_argument("--iters", type=int, default=3)
+    ap.add_argument("--dtype", default="float32", choices=["float32", "bfloat16", "float16"])
+    ap.add_argument(
+        "--decode-tokens", type=int, default=50,
+        help="Number of tokens to generate for TPOT measurement (0 = skip).",
+    )
+    ap.add_argument("--max-cache-len", type=int, default=256,
+                    help="StaticCache max sequence length.")
+    args = ap.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    input_ids = encode_prompt(tokenizer, args.prompt, device)
+    prompt_tokens = int(input_ids.shape[-1])
+
+    config = AutoConfig.from_pretrained(args.model)
+    config._attn_implementation = "eager"
+
+    model = (
+        AutoModelForCausalLM.from_pretrained(args.model, config=config, torch_dtype=dtype)
+        .eval()
+        .to(device)
+    )
+
+    single_token = torch.zeros(1, 1, dtype=torch.long, device=device)
+
+    cache_config = static_cache_config(config)
+
+    def make_cache():
+        return StaticCache(
+            config=cache_config,
+            max_batch_size=1,
+            max_cache_len=args.max_cache_len,
+            device=device,
+            dtype=dtype,
+        )
+
+    # Eager init on the uncompiled model so the StaticCache buffers get
+    # registered (mark_static_address) before torch.compile traces them.
+    init_cache = make_cache()
+    with torch.no_grad():
+        model(single_token, past_key_values=init_cache,
+              cache_position=torch.tensor([0], device=device))
+
+    compiled = torch.compile(model)
+
+    # First compiled call triggers JIT compilation; time it as compile_ms.
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    with torch.no_grad():
+        compiled(single_token, past_key_values=init_cache,
+                 cache_position=torch.tensor([1], device=device))
+        if device.type == "cuda":
+            torch.cuda.synchronize()
+    compile_ms = (time.perf_counter() - t0) * 1000.0
+
+    def measure_ttft() -> float:
+        """Sum of per-token compiled-forward durations over prompt_tokens steps."""
+        kv = make_cache()
+        # Fresh cache needs eager init via the uncompiled model first.
+        with torch.no_grad():
+            model(single_token, past_key_values=kv,
+                  cache_position=torch.tensor([0], device=device))
+        total_ms = 0.0
+        for pos in range(1, prompt_tokens):
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            with torch.no_grad():
+                compiled(single_token, past_key_values=kv,
+                         cache_position=torch.tensor([pos], device=device))
+                if device.type == "cuda":
+                    torch.cuda.synchronize()
+            total_ms += (time.perf_counter() - t0) * 1000.0
+        return total_ms
+
+    for _ in range(args.warmups):
+        measure_ttft()
+
+    ttft_samples_ms = [measure_ttft() for _ in range(args.iters)]
+
+    result = {
+        "path": "python_torch_compile",
+        "model": args.model,
+        "device": str(device),
+        "dtype": args.dtype,
+        "prompt_tokens": prompt_tokens,
+        "iters": args.iters,
+        "ttft_ms": statistics.median(ttft_samples_ms),
+        "ttft_ms_mean": sum(ttft_samples_ms) / len(ttft_samples_ms),
+        "ttft_ms_samples": ttft_samples_ms,
+        "compile_ms": compile_ms,
+        "note": "sequential per-token, StaticCache KV cache (torch.compile inductor)",
+    }
+
+    if args.decode_tokens > 0:
+        tpot_samples_ms = measure_tpot(compiled, input_ids, device, args.decode_tokens)
+        tpot_ms = sum(tpot_samples_ms) / len(tpot_samples_ms)
+        result["decode_tokens"] = args.decode_tokens
+        result["tpot_ms"] = tpot_ms
+        result["tpot_ms_samples"] = tpot_samples_ms
+        result["throughput_tps"] = 1000.0 / tpot_ms
+
+    print("BENCH_RESULT " + json.dumps(result))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ttft/bench_utils.py

@@ -0,0 +1,94 @@
+"""Shared helpers for the Python benchmark scripts."""
+
+import time
+
+import torch
+
+
+class _CfgWithoutKvShared:
+    """Wrapper that hides `num_kv_shared_layers` from a HF config.
+
+    transformers 5.6 has a bug in StaticCache.__init__:
+        if hasattr(config, "num_kv_shared_layers"):
+            layer_types = layer_types[: -config.num_kv_shared_layers]
+    For configs where the attribute is 0 (e.g. Gemma-4), `[:-0]` returns an
+    empty list, leaving StaticCache with zero layer slots, and the LM's
+    first `past_key_values.update(..., layer_idx=0)` raises IndexError.
+
+    This wrapper makes `hasattr(...)` return False so the bad branch never
+    fires. Used via `static_cache_config(config)` below.
+    """
+    __slots__ = ("_inner",)
+
+    def __init__(self, inner):
+        object.__setattr__(self, "_inner", inner)
+
+    def __getattr__(self, name):
+        if name == "num_kv_shared_layers":
+            raise AttributeError(name)
+        return getattr(self._inner, name)
+
+    def get_text_config(self, *args, **kwargs):
+        return _CfgWithoutKvShared(self._inner.get_text_config(*args, **kwargs))
+
+
+def static_cache_config(config):
+    """Return a config suitable for `StaticCache(config=..., ...)`.
+
+    Two normalizations:
+      1. Multimodal wrappers (Gemma4ForConditionalGeneration, ...) nest the
+         actual LM config under `.text_config`. Pass that, not the wrapper,
+         so layer/head counts match the inner LM.
+      2. If the resulting config has `num_kv_shared_layers == 0`, wrap it to
+         hide the attribute (works around the transformers 5.6 slice bug).
+    """
+    cfg = getattr(config, "text_config", config)
+    if getattr(cfg, "num_kv_shared_layers", None) == 0:
+        cfg = _CfgWithoutKvShared(cfg)
+    return cfg
+
+
+def encode_prompt(tokenizer, prompt: str, device):
+    """Tokenize prompt using chat template if available, falling back to raw tokenization."""
+    messages = [{"role": "user", "content": prompt}]
+    try:
+        encoded = tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True, return_tensors="pt"
+        )
+    except (ValueError, AttributeError):
+        encoded = tokenizer(prompt, return_tensors="pt")
+    if hasattr(encoded, "input_ids"):
+        return encoded.input_ids.to(device)
+    if isinstance(encoded, dict):
+        return encoded["input_ids"].to(device)
+    return encoded.to(device)
+
+
+def measure_tpot(model, input_ids, device, decode_tokens: int) -> list[float]:
+    """Prefill once with KV cache, then time each subsequent single-token decode step."""
+    with torch.no_grad():
+        out = model(input_ids, use_cache=True)
+        if device.type == "cuda":
+            torch.cuda.synchronize()
+        past = out.past_key_values
+        next_id = out.logits[:, -1:].argmax(-1)
+
+        out = model(next_id, past_key_values=past, use_cache=True)
+        if device.type == "cuda":
+            torch.cuda.synchronize()
+        past = out.past_key_values
+        next_id = out.logits[:, -1:].argmax(-1)
+
+        step_times_ms = []
+        for _ in range(decode_tokens):
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            out = model(next_id, past_key_values=past, use_cache=True)
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+            step_times_ms.append((time.perf_counter() - t0) * 1000.0)
+            past = out.past_key_values
+            next_id = out.logits[:, -1:].argmax(-1)
+
+    return step_times_ms

benchmarks/ttft/benchmarks.toml

@@ -0,0 +1,92 @@
+[ur_test]
+models = ["llama-8b", "qwen3-4b", "gemma3-4b", "gemma4-moe", "qwen3-moe"]
+# 3-point sweep (low/mid/high). The previous list [5, 10, 20, 50, 100, 500]
+# spent ~62 extra minutes on s=5/s=20/s=50 with little additional information.
+search_sweep_iters = [10, 100, 500]
+
+[configs.llama-8b]
+model = "NousResearch/Meta-Llama-3-8B-Instruct"
+rust_package = "llama"
+search_iters = 500
+iters = 10
+warmups = 2
+decode_tokens = 50
+# On-disk weights are bf16-majority. fp32 upcast doubled python_luminal's
+# egglog Search peak past the 525 GB unified pool and triggered SIGKILLs on
+# gemma3-4b (and same risk here). bf16 matches rust's load path.
+dtype = "bfloat16"
+
+[configs.as_fast_as_possible]
+prompt = "The"
+search_iters = 1
+iters = 1
+warmups = 0
+decode_tokens = 5
+
+[configs.qwen3-4b]
+model = "Qwen/Qwen3-4B"
+rust_package = "qwen"
+search_iters = 50
+iters = 10
+warmups = 2
+decode_tokens = 20
+# bf16-majority on-disk; see llama-8b note.
+dtype = "bfloat16"
+
+[configs.gemma3-4b]
+model = "unsloth/gemma-3-4b-it"
+rust_package = "gemma"
+search_iters = 50
+iters = 10
+warmups = 2
+decode_tokens = 20
+# bf16-majority on-disk; see llama-8b note.
+dtype = "bfloat16"
+
+[configs.gemma4-moe]
+model = "google/gemma-4-26B-A4B"
+rust_package = "gemma4_moe"
+search_iters = 50
+iters = 10
+warmups = 2
+decode_tokens = 20
+# 26B params at fp32 = 104 GB → OOM on a 94 GB GPU. Use bf16 (matches the
+# on-disk safetensors dtype) so the python paths can actually load.
+dtype = "bfloat16"
+
+[configs.qwen3-moe]
+model = "Qwen/Qwen3-30B-A3B"
+rust_package = "qwen3_moe"
+search_iters = 50
+iters = 10
+warmups = 2
+decode_tokens = 20
+# 30B params at fp32 = 120 GB → OOM. See gemma4-moe note.
+dtype = "bfloat16"
+
+[configs.llama-8b-const]
+model = "NousResearch/Meta-Llama-3-8B-Instruct"
+rust_package = "llama"
+prompt = "We the People of the United States, in Order to form a more perfect Union, establish Justice, insure domestic Tranquility, provide for the common defence, promote the general Welfare, and secure the Blessings of Liberty to ourselves and our Posterity, do ordain and establish this Constitution for the United States of America."
+search_iters = 500
+iters = 10
+warmups = 2
+decode_tokens = 20
+
+[configs.qwen3-4b-const]
+model = "Qwen/Qwen3-4B"
+rust_package = "qwen"
+prompt = "We the People of the United States, in Order to form a more perfect Union, establish Justice, insure domestic Tranquility, provide for the common defence, promote the general Welfare, and secure the Blessings of Liberty to ourselves and our Posterity, do ordain and establish this Constitution for the United States of America."
+search_iters = 50
+iters = 10
+warmups = 2
+decode_tokens = 20
+
+[configs.gemma3-4b-const]
+model = "unsloth/gemma-3-4b-it"
+rust_package = "gemma"
+prompt = "We the People of the United States, in Order to form a more perfect Union, establish Justice, insure domestic Tranquility, provide for the common defence, promote the general Welfare, and secure the Blessings of Liberty to ourselves and our Posterity, do ordain and establish this Constitution for the United States of America."
+search_iters = 50
+iters = 10
+warmups = 2
+decode_tokens = 20

benchmarks/ttft/dashboard.html

@@ -0,0 +1,610 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>Luminal · Benchmark Dashboard</title>
+<link rel="preconnect" href="https://fonts.googleapis.com">
+<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+<link href="https://fonts.googleapis.com/css2?family=Geist:wght@300;400;500;600&family=Geist+Mono:wght@300;400;500&display=swap" rel="stylesheet">
+<script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
+<style>
+*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
+html { -webkit-font-smoothing: antialiased; scroll-behavior: smooth; }
+
+body {
+  font-family: 'Geist', system-ui, sans-serif;
+  background: #030712;
+  color: #d7d8d9;
+  min-height: 100vh;
+  line-height: 1.5;
+}
+
+/* ── NAV ── */
+nav {
+  position: sticky;
+  top: 0;
+  z-index: 50;
+  height: 56px;
+  background: rgba(8, 15, 17, 0.92);
+  backdrop-filter: blur(8px);
+  -webkit-backdrop-filter: blur(8px);
+  border-bottom: 1px solid #2d3335;
+  display: flex;
+  align-items: center;
+  padding: 0 24px;
+  gap: 0;
+}
+.nav-brand {
+  display: flex;
+  align-items: center;
+  gap: 8px;
+  font-family: 'Geist Mono', monospace;
+  font-size: 14px;
+  font-weight: 500;
+  letter-spacing: 0.05em;
+  color: #2faa6e;
+  text-decoration: none;
+}
+.nav-dot {
+  width: 6px;
+  height: 6px;
+  background: #2faa6e;
+  border-radius: 50%;
+  flex-shrink: 0;
+  animation: pulse-glow 2s ease-in-out infinite;
+}
+.nav-sep {
+  color: #2d3335;
+  margin: 0 14px;
+  font-size: 18px;
+  font-weight: 300;
+}
+.nav-page {
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.1em;
+  text-transform: uppercase;
+  color: #7e8385;
+}
+
+@keyframes pulse-glow {
+  0%, 100% { opacity: 1; }
+  50%       { opacity: 0.35; }
+}
+
+/* ── MAIN ── */
+main {
+  max-width: 1200px;
+  margin: 0 auto;
+  padding: 40px 24px 80px;
+}
+
+/* ── PAGE HEADER ── */
+.page-header {
+  margin-bottom: 40px;
+  padding-bottom: 32px;
+  border-bottom: 1px solid #1c2225;
+}
+.page-eyebrow {
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.1em;
+  text-transform: uppercase;
+  color: #2faa6e;
+  margin-bottom: 10px;
+}
+.page-title {
+  font-size: 30px;
+  font-weight: 500;
+  letter-spacing: -0.025em;
+  color: #d7d8d9;
+  margin-bottom: 10px;
+}
+.page-meta {
+  font-size: 14px;
+  color: #7e8385;
+  display: flex;
+  align-items: center;
+  gap: 0;
+  flex-wrap: wrap;
+}
+.meta-sep {
+  font-family: 'Geist Mono', monospace;
+  color: #2d3335;
+  margin: 0 10px;
+}
+.meta-val {
+  font-family: 'Geist Mono', monospace;
+  font-size: 13px;
+  color: #5b5f61;
+}
+
+/* ── LEGEND STRIP ── */
+.legend-strip {
+  display: flex;
+  flex-wrap: wrap;
+  gap: 6px;
+  margin-bottom: 32px;
+}
+.legend-pill {
+  display: flex;
+  align-items: center;
+  gap: 6px;
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.04em;
+  color: #a1a4a5;
+  background: #141b1d;
+  border: 1px solid #2d3335;
+  border-radius: 2px;
+  padding: 4px 10px;
+}
+.legend-swatch {
+  width: 8px;
+  height: 8px;
+  border-radius: 50%;
+  flex-shrink: 0;
+}
+
+/* ── SECTIONS ── */
+section { margin-bottom: 48px; }
+.section-header {
+  display: flex;
+  align-items: baseline;
+  gap: 10px;
+  margin-bottom: 16px;
+  padding-bottom: 12px;
+  border-bottom: 1px solid #1c2225;
+  flex-wrap: wrap;
+}
+.section-eyebrow {
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.1em;
+  text-transform: uppercase;
+  color: #404647;
+}
+.section-title {
+  font-size: 18px;
+  font-weight: 500;
+  color: #d7d8d9;
+  letter-spacing: -0.01em;
+}
+.section-title .unit {
+  color: #7e8385;
+  font-weight: 400;
+}
+.section-tag {
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.04em;
+  text-transform: uppercase;
+  color: #2faa6e;
+  background: #162322;
+  border: 1px solid #1c372e;
+  padding: 2px 8px;
+  border-radius: 2px;
+  margin-left: auto;
+}
+
+/* ── CHART GRID ── */
+.chart-grid {
+  display: grid;
+  gap: 10px;
+}
+.chart-card {
+  background: #141b1d;
+  border: 1px solid #2d3335;
+  border-radius: 2px;
+  overflow: hidden;
+  transition: border-color 150ms;
+  min-width: 0;
+}
+.chart-card:hover { border-color: #404647; }
+.chart-card-header {
+  padding: 10px 14px 0;
+  display: flex;
+  align-items: center;
+}
+.model-tag {
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.06em;
+  text-transform: uppercase;
+  color: #7e8385;
+}
+
+/* ── FOOTER ── */
+footer {
+  max-width: 1200px;
+  margin: 0 auto;
+  padding: 20px 24px;
+  border-top: 1px solid #1c2225;
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.04em;
+  color: #404647;
+  display: flex;
+  justify-content: space-between;
+  flex-wrap: wrap;
+  gap: 8px;
+}
+
+.section-divider {
+  border: none;
+  border-top: 1px solid #1c2225;
+  margin: 8px 0 40px;
+}
+.sweep-hint {
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.04em;
+  color: #404647;
+  margin-bottom: 12px;
+}
+
+@media (max-width: 768px) {
+  .chart-grid { grid-template-columns: 1fr !important; }
+  .page-title { font-size: 22px; }
+}
+</style>
+</head>
+<body>
+
+<nav>
+  <a class="nav-brand" href="https://luminal.com">
+    <span class="nav-dot"></span>luminal
+  </a>
+  <span class="nav-sep">/</span>
+  <span class="nav-page">benchmarks</span>
+</nav>
+
+<main>
+
+<header class="page-header">
+  <p class="page-eyebrow">performance · time-series</p>
+  <h1 class="page-title">Benchmark Dashboard</h1>
+  <div class="page-meta">
+    <span>Last updated</span>
+    <span class="meta-sep">·</span>
+    <span class="meta-val">May 01, 2026 · 18:56</span>
+    <span class="meta-sep">·</span>
+    <span class="meta-val">1 run in history</span>
+  </div>
+</header>
+
+<div class="legend-strip">
+  <div class="legend-pill"><span class="legend-swatch" style="background:#5b5f61"></span>HF Baseline</div><div class="legend-pill"><span class="legend-swatch" style="background:#3b82f6"></span>torch.compile</div><div class="legend-pill"><span class="legend-swatch" style="background:#a855f7"></span>luminal backend</div><div class="legend-pill"><span class="legend-swatch" style="background:#e8855a"></span>Rust (luminal)</div>
+</div>
+
+
+<section>
+  <div class="section-header">
+    <span class="section-eyebrow">metric</span>
+    <h2 class="section-title">TTFT <span class="unit">over time</span></h2>
+    <span class="section-tag">Time to first token (ms)</span>
+  </div>
+  <div class="chart-grid" style="grid-template-columns: repeat(4, 1fr)">
+<div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">llama-8b</span>
+  </div>
+  <div id="c_ttft_ms_llama_8b"></div>
+  <script>
+    Plotly.newPlot("c_ttft_ms_llama_8b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [705.9654394979589], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [307.66548847896047], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [461.48114453535527], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [1026.86], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 48, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-4b</span>
+  </div>
+  <div id="c_ttft_ms_qwen3_4b"></div>
+  <script>
+    Plotly.newPlot("c_ttft_ms_qwen3_4b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [869.2860195587855], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [298.27259748708457], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [485.3892414830625], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [398.58], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma3-4b</span>
+  </div>
+  <div id="c_ttft_ms_gemma3_4b"></div>
+  <script>
+    Plotly.newPlot("c_ttft_ms_gemma3_4b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [951.1196144158021], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [300.9451600664761], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [404.43], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma4-moe</span>
+  </div>
+  <div id="c_ttft_ms_gemma4_moe"></div>
+  <script>
+    Plotly.newPlot("c_ttft_ms_gemma4_moe", [{"x": ["2026-05-01T18-56-26-996695"], "y": [837.3980740143452], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [245.510076492792], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-moe</span>
+  </div>
+  <div id="c_ttft_ms_qwen3_moe"></div>
+  <script>
+    Plotly.newPlot("c_ttft_ms_qwen3_moe", [{"x": ["2026-05-01T18-56-26-996695"], "y": [1565.540504961973], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [460.077923577046], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [21002.791983017232], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [662.07], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div>
+  </div>
+</section>
+<section>
+  <div class="section-header">
+    <span class="section-eyebrow">metric</span>
+    <h2 class="section-title">TPOT <span class="unit">over time</span></h2>
+    <span class="section-tag">Time per output token (ms)</span>
+  </div>
+  <div class="chart-grid" style="grid-template-columns: repeat(4, 1fr)">
+<div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">llama-8b</span>
+  </div>
+  <div id="c_tpot_ms_llama_8b"></div>
+  <script>
+    Plotly.newPlot("c_tpot_ms_llama_8b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [34.15271903970279], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [171.7862353892997], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [23.078908618772402], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [51.64], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 48, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-4b</span>
+  </div>
+  <div id="c_tpot_ms_qwen3_4b"></div>
+  <script>
+    Plotly.newPlot("c_tpot_ms_qwen3_4b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [47.71483448566869], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [468.56868775503244], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [26.90318431414198], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [40.62], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma3-4b</span>
+  </div>
+  <div id="c_tpot_ms_gemma3_4b"></div>
+  <script>
+    Plotly.newPlot("c_tpot_ms_gemma3_4b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [52.498737201676704], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [2197.426627812092], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [38.99], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma4-moe</span>
+  </div>
+  <div id="c_tpot_ms_gemma4_moe"></div>
+  <script>
+    Plotly.newPlot("c_tpot_ms_gemma4_moe", [{"x": ["2026-05-01T18-56-26-996695"], "y": [83.64427039632574], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [654.9649795080768], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-moe</span>
+  </div>
+  <div id="c_tpot_ms_qwen3_moe"></div>
+  <script>
+    Plotly.newPlot("c_tpot_ms_qwen3_moe", [{"x": ["2026-05-01T18-56-26-996695"], "y": [84.527321747737], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "HF Baseline", "line": {"color": "#5b5f61", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>HF Baseline</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [753.0061075551203], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [1166.8824461026816], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [60.08], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} ms<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " ms", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div>
+  </div>
+</section>
+<section>
+  <div class="section-header">
+    <span class="section-eyebrow">metric</span>
+    <h2 class="section-title">Time to Search <span class="unit">over time</span></h2>
+    <span class="section-tag">Search time (sec)</span>
+  </div>
+  <div class="chart-grid" style="grid-template-columns: repeat(4, 1fr)">
+<div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">llama-8b</span>
+  </div>
+  <div id="c_compile_ms_llama_8b"></div>
+  <script>
+    Plotly.newPlot("c_compile_ms_llama_8b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [18.760145067994017], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [95.96263545705006], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [84.45343], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": true, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 48, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " sec", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-4b</span>
+  </div>
+  <div id="c_compile_ms_qwen3_4b"></div>
+  <script>
+    Plotly.newPlot("c_compile_ms_qwen3_4b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [4.680963660997804], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [45.345814052037895], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [19.92977], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " sec", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma3-4b</span>
+  </div>
+  <div id="c_compile_ms_gemma3_4b"></div>
+  <script>
+    Plotly.newPlot("c_compile_ms_gemma3_4b", [{"x": ["2026-05-01T18-56-26-996695"], "y": [26.649526304972824], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [156.84164], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " sec", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma4-moe</span>
+  </div>
+  <div id="c_compile_ms_gemma4_moe"></div>
+  <script>
+    Plotly.newPlot("c_compile_ms_gemma4_moe", [{"x": ["2026-05-01T18-56-26-996695"], "y": [38.81582092499593], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " sec", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-moe</span>
+  </div>
+  <div id="c_compile_ms_qwen3_moe"></div>
+  <script>
+    Plotly.newPlot("c_compile_ms_qwen3_moe", [{"x": ["2026-05-01T18-56-26-996695"], "y": [8.341281775035895], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "torch.compile", "line": {"color": "#3b82f6", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>torch.compile</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [111.70731823903043], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "luminal backend", "line": {"color": "#a855f7", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>luminal backend</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}, {"x": ["2026-05-01T18-56-26-996695"], "y": [80.83241000000001], "customdata": [["b2bd91f5", "2026-05-01T18:56:26.996695"]], "type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "line": {"color": "#e8855a", "width": 2}, "marker": {"size": 7, "symbol": "circle"}, "connectgaps": false, "showlegend": false, "hovertemplate": "<b>Rust (luminal)</b><br>%{customdata[1]}<br>%{y:.1f} sec<br><span style='color:#7e8385'>commit %{customdata[0]}</span><extra></extra>"}], {"plot_bgcolor": "#0d1416", "paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"}, "margin": {"t": 16, "b": 16, "l": 52, "r": 12}, "height": 280, "xaxis": {"type": "category", "categoryorder": "array", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "tickangle": -30, "automargin": true, "zeroline": false, "categoryarray": ["2026-05-01T18-56-26-996695"], "tickvals": ["2026-05-01T18-56-26-996695"], "ticktext": ["May 01 \u00b7 18:56"]}, "yaxis": {"rangemode": "tozero", "color": "#5b5f61", "gridcolor": "#1c2225", "linecolor": "#2d3335", "tickfont": {"size": 11, "family": "Geist Mono, monospace"}, "ticksuffix": " sec", "zeroline": false}, "legend": {"orientation": "h", "y": -0.28, "x": 0, "font": {"size": 11, "color": "#a1a4a5"}, "bgcolor": "rgba(0,0,0,0)", "visible": false}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}},
+      {responsive: true, displayModeBar: false});
+  </script>
+</div>
+  </div>
+</section>
+<hr class='section-divider'>
+<section>
+  <div class="section-header">
+    <span class="section-eyebrow">sweep · 3d</span>
+    <h2 class="section-title">TTFT <span class="unit">vs search budget · over time</span></h2>
+    <span class="section-tag">1 run</span>
+  </div>
+  <p class="sweep-hint">Drag to rotate · scroll to zoom · each curve = one run</p>
+  <div class="chart-grid" style="grid-template-columns: repeat(4, 1fr)">
+<div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">llama-8b</span>
+  </div>
+  <div id="sw_ttft_ms_llama_8b"></div>
+  <script>
+    Plotly.newPlot("sw_ttft_ms_llama_8b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [470.7036415056791, 460.72837291285396, 472.43661794345826], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [751.03, 1038.34, 453.16], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " ms", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-4b</span>
+  </div>
+  <div id="sw_ttft_ms_qwen3_4b"></div>
+  <script>
+    Plotly.newPlot("sw_ttft_ms_qwen3_4b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [465.02652901108377, 465.9317950136028, 495.75577257201076], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [398.44, 390.08, 559.29], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " ms", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma3-4b</span>
+  </div>
+  <div id="sw_ttft_ms_gemma3_4b"></div>
+  <script>
+    Plotly.newPlot("sw_ttft_ms_gemma3_4b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [388.19, 436.49, 386.13], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " ms", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-moe</span>
+  </div>
+  <div id="sw_ttft_ms_qwen3_moe"></div>
+  <script>
+    Plotly.newPlot("sw_ttft_ms_qwen3_moe", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [21002.663500519702, 21018.686580006033, 21034.366824431345], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [656.7, 540.37, 542.34], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " ms", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div>
+  </div>
+</section>
+<section>
+  <div class="section-header">
+    <span class="section-eyebrow">sweep · 3d</span>
+    <h2 class="section-title">TPOT <span class="unit">vs search budget · over time</span></h2>
+    <span class="section-tag">1 run</span>
+  </div>
+  <p class="sweep-hint">Drag to rotate · scroll to zoom · each curve = one run</p>
+  <div class="chart-grid" style="grid-template-columns: repeat(4, 1fr)">
+<div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">llama-8b</span>
+  </div>
+  <div id="sw_tpot_ms_llama_8b"></div>
+  <script>
+    Plotly.newPlot("sw_tpot_ms_llama_8b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [23.540849717101082, 23.101884137140587, 23.610779400914907], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [38.2, 51.92, 24.09], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " ms", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-4b</span>
+  </div>
+  <div id="sw_tpot_ms_qwen3_4b"></div>
+  <script>
+    Plotly.newPlot("sw_tpot_ms_qwen3_4b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [25.875402649398893, 25.884080055402592, 27.492373346467502], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [40.64, 39.98, 55.37], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " ms", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma3-4b</span>
+  </div>
+  <div id="sw_tpot_ms_gemma3_4b"></div>
+  <script>
+    Plotly.newPlot("sw_tpot_ms_gemma3_4b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [37.47, 41.95, 37.25], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " ms", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-moe</span>
+  </div>
+  <div id="sw_tpot_ms_qwen3_moe"></div>
+  <script>
+    Plotly.newPlot("sw_tpot_ms_qwen3_moe", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [1166.6714247548953, 1167.2746865515364, 1168.7990181031637], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [59.6, 48.79, 48.88], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} ms<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " ms", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div>
+  </div>
+</section>
+<section>
+  <div class="section-header">
+    <span class="section-eyebrow">sweep · 3d</span>
+    <h2 class="section-title">Time to Search <span class="unit">vs search budget · over time</span></h2>
+    <span class="section-tag">1 run</span>
+  </div>
+  <p class="sweep-hint">Drag to rotate · scroll to zoom · each curve = one run</p>
+  <div class="chart-grid" style="grid-template-columns: repeat(4, 1fr)">
+<div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">llama-8b</span>
+  </div>
+  <div id="sw_compile_ms_llama_8b"></div>
+  <script>
+    Plotly.newPlot("sw_compile_ms_llama_8b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [28.428826077957638, 43.57440591201885, 95.52432684396626], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} sec<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [15.14307, 30.12727, 84.87889], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} sec<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "sec", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " sec", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-4b</span>
+  </div>
+  <div id="sw_compile_ms_qwen3_4b"></div>
+  <script>
+    Plotly.newPlot("sw_compile_ms_qwen3_4b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [37.92102829599753, 54.08867314597592, 118.29659596900456], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} sec<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [12.448030000000001, 27.06796, 81.89342], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} sec<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "sec", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " sec", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">gemma3-4b</span>
+  </div>
+  <div id="sw_compile_ms_gemma3_4b"></div>
+  <script>
+    Plotly.newPlot("sw_compile_ms_gemma3_4b", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [102.18644, 186.34269, 498.48983000000004], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} sec<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "sec", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " sec", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div><div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">qwen3-moe</span>
+  </div>
+  <div id="sw_compile_ms_qwen3_moe"></div>
+  <script>
+    Plotly.newPlot("sw_compile_ms_qwen3_moe", [{"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [93.47603664599592, 132.266081985028, 298.05094401398674], "name": "luminal backend", "legendgroup": "python_luminal", "showlegend": true, "line": {"color": "#a855f7", "width": 5}, "marker": {"color": "#a855f7", "size": 4}, "hovertemplate": "<b>luminal backend</b><br>s=%{x} iters<br>%{z:.1f} sec<br>May 01 \u00b7 b2bd91f5<extra></extra>"}, {"type": "scatter3d", "mode": "lines+markers", "x": [10, 100, 500], "y": ["May 01", "May 01", "May 01"], "z": [25.48138, 47.5342, 134.79345], "name": "Rust (luminal)", "legendgroup": "rust", "showlegend": true, "line": {"color": "#e8855a", "width": 5}, "marker": {"color": "#e8855a", "size": 4}, "hovertemplate": "<b>Rust (luminal)</b><br>s=%{x} iters<br>%{z:.1f} sec<br>May 01 \u00b7 b2bd91f5<extra></extra>"}], {"paper_bgcolor": "#141b1d", "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11}, "height": 420, "margin": {"t": 20, "b": 0, "l": 0, "r": 0}, "legend": {"orientation": "h", "y": -0.05, "x": 0, "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"}, "bgcolor": "rgba(0,0,0,0)"}, "hoverlabel": {"bgcolor": "#1c2225", "bordercolor": "#2d3335", "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"}}, "scene": {"bgcolor": "#0d1416", "xaxis": {"title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}}, "type": "log", "tickvals": [5, 10, 20, 50, 100, 500], "ticktext": ["5", "10", "20", "50", "100", "500"], "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335", "zerolinecolor": "#2d3335"}, "yaxis": {"title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}}, "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "zaxis": {"title": {"text": "sec", "font": {"size": 10, "color": "#7e8385"}}, "rangemode": "tozero", "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"}, "ticksuffix": " sec", "gridcolor": "#1c2225", "linecolor": "#2d3335"}, "camera": {"eye": {"x": 1.6, "y": -1.6, "z": 0.9}}}},
+      {responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]});
+  </script>
+</div>
+  </div>
+</section>
+
+</main>
+
+<footer>
+  <span>luminal · benchmark dashboard</span>
+  <span>generated May 01, 2026 · 18:56</span>
+</footer>
+
+</body>
+</html>

benchmarks/ttft/db.py

@@ -0,0 +1,242 @@
+"""SQLite persistence for TTFT/TPOT benchmark runs.
+
+Two tables:
+  runs    — one row per orchestrator invocation
+  results — many rows per run, one per (path, config) combination
+
+`results` carries every field that today's BENCH_RESULT JSON record carries.
+Per-iteration sample arrays (`ttft_ms_samples`, `tpot_ms_samples`) are kept as
+JSON TEXT — they're archival, no consumer aggregates over them.
+
+The default DB path is benchmarks/ttft/bench.db (gitignored). Schema is
+created lazily on first connect.
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+from pathlib import Path
+from typing import Any, Iterable
+
+BENCH_DIR = Path(__file__).resolve().parent
+DEFAULT_DB_PATH = BENCH_DIR / "bench.db"
+
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS runs (
+  run_id        TEXT PRIMARY KEY,
+  timestamp     TEXT NOT NULL,
+  git_commit    TEXT,
+  git_branch    TEXT,
+  gpu_name      TEXT,
+  gpu_driver    TEXT,
+  gpu_vram_mb   INTEGER,
+  cuda_version  TEXT,
+  mode          TEXT NOT NULL  -- 'single' | 'all-configs' | 'search-sweep' | 'ur-test' | 'ur-test-fast'
+);
+
+CREATE TABLE IF NOT EXISTS results (
+  id              INTEGER PRIMARY KEY AUTOINCREMENT,
+  run_id          TEXT NOT NULL REFERENCES runs(run_id) ON DELETE CASCADE,
+  path            TEXT NOT NULL,
+  model           TEXT NOT NULL,
+  model_key       TEXT,
+  config          TEXT NOT NULL,
+  device          TEXT,
+  dtype           TEXT,
+  prompt_tokens   INTEGER,
+  iters           INTEGER,
+  decode_tokens   INTEGER,
+  search_iters    INTEGER,
+  ttft_ms         REAL,
+  ttft_ms_mean    REAL,
+  tpot_ms         REAL,
+  throughput_tps  REAL,
+  compile_ms      REAL,
+  note            TEXT,
+  error           TEXT,
+  ttft_ms_samples TEXT,
+  tpot_ms_samples TEXT,
+  created_at      TEXT NOT NULL DEFAULT (datetime('now'))
+);
+
+CREATE INDEX IF NOT EXISTS idx_results_run    ON results(run_id);
+CREATE INDEX IF NOT EXISTS idx_results_path   ON results(path);
+CREATE INDEX IF NOT EXISTS idx_results_config ON results(config);
+CREATE INDEX IF NOT EXISTS idx_results_modelk ON results(model_key);
+"""
+
+
+# Columns that map 1:1 from a BENCH_RESULT record dict into `results`.
+_SCALAR_RESULT_COLS = (
+    "path", "model", "model_key", "config",
+    "device", "dtype",
+    "prompt_tokens", "iters", "decode_tokens", "search_iters",
+    "ttft_ms", "ttft_ms_mean", "tpot_ms", "throughput_tps", "compile_ms",
+    "note", "error",
+)
+_SAMPLE_COLS = ("ttft_ms_samples", "tpot_ms_samples")
+_ALL_RESULT_COLS = ("run_id",) + _SCALAR_RESULT_COLS + _SAMPLE_COLS
+
+
+def connect(path: str | Path = DEFAULT_DB_PATH) -> sqlite3.Connection:
+    """Open (or create) the bench DB and ensure the schema exists."""
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(p)
+    conn.row_factory = sqlite3.Row
+    conn.execute("PRAGMA foreign_keys = ON")
+    conn.executescript(_SCHEMA)
+    return conn
+
+
+def insert_run(
+    conn: sqlite3.Connection,
+    *,
+    run_id: str,
+    timestamp: str,
+    mode: str,
+    git_commit: str | None = None,
+    git_branch: str | None = None,
+    gpu_name: str | None = None,
+    gpu_driver: str | None = None,
+    gpu_vram_mb: int | None = None,
+    cuda_version: str | None = None,
+    if_exists: str = "ignore",
+) -> str:
+    """Insert a run row.  if_exists='ignore' (default) leaves an existing
+    row untouched; 'replace' overwrites."""
+    verb = {"ignore": "INSERT OR IGNORE", "replace": "INSERT OR REPLACE"}[if_exists]
+    conn.execute(
+        f"""{verb} INTO runs
+            (run_id, timestamp, git_commit, git_branch,
+             gpu_name, gpu_driver, gpu_vram_mb, cuda_version, mode)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+        (run_id, timestamp, git_commit, git_branch,
+         gpu_name, gpu_driver, gpu_vram_mb, cuda_version, mode),
+    )
+    return run_id
+
+
+def insert_result(conn: sqlite3.Connection, run_id: str, record: dict[str, Any]) -> int:
+    """Insert one BENCH_RESULT-shaped record under the given run_id."""
+    values = [run_id]
+    for col in _SCALAR_RESULT_COLS:
+        values.append(record.get(col))
+    for col in _SAMPLE_COLS:
+        v = record.get(col)
+        values.append(json.dumps(v) if v is not None else None)
+    placeholders = ", ".join(["?"] * len(_ALL_RESULT_COLS))
+    cols = ", ".join(_ALL_RESULT_COLS)
+    cur = conn.execute(
+        f"INSERT INTO results ({cols}) VALUES ({placeholders})",
+        values,
+    )
+    return cur.lastrowid
+
+
+def insert_results(conn: sqlite3.Connection, run_id: str, records: Iterable[dict[str, Any]]) -> int:
+    """Bulk-insert; returns count."""
+    n = 0
+    for r in records:
+        insert_result(conn, run_id, r)
+        n += 1
+    return n
+
+
+def latest_run_id(conn: sqlite3.Connection) -> str | None:
+    row = conn.execute(
+        "SELECT run_id FROM runs ORDER BY timestamp DESC, run_id DESC LIMIT 1"
+    ).fetchone()
+    return row["run_id"] if row else None
+
+
+def load_run(conn: sqlite3.Connection, run_id: str) -> dict[str, Any] | None:
+    row = conn.execute("SELECT * FROM runs WHERE run_id = ?", (run_id,)).fetchone()
+    return dict(row) if row else None
+
+
+def load_runs(conn: sqlite3.Connection) -> list[dict[str, Any]]:
+    """All runs, oldest → newest."""
+    rows = conn.execute(
+        "SELECT * FROM runs ORDER BY timestamp ASC, run_id ASC"
+    ).fetchall()
+    return [dict(r) for r in rows]
+
+
+def _row_to_record(row: sqlite3.Row) -> dict[str, Any]:
+    """Convert a results row into a BENCH_RESULT-shaped dict, stripping NULLs
+    so consumers see the same shape they did with JSON."""
+    out: dict[str, Any] = {}
+    for col in _SCALAR_RESULT_COLS:
+        v = row[col]
+        if v is not None:
+            out[col] = v
+    for col in _SAMPLE_COLS:
+        v = row[col]
+        if v is not None:
+            out[col] = json.loads(v)
+    return out
+
+
+def load_results(conn: sqlite3.Connection, run_id: str) -> list[dict[str, Any]]:
+    """All results for one run, in insertion order."""
+    rows = conn.execute(
+        "SELECT * FROM results WHERE run_id = ? ORDER BY id ASC", (run_id,)
+    ).fetchall()
+    return [_row_to_record(r) for r in rows]
+
+
+def load_history(conn: sqlite3.Connection) -> list[dict[str, Any]]:
+    """Mirror the legacy gen_dashboard.load_history() shape:
+    [{"meta": {...}, "results": [...], "sweep": [...]}], sorted oldest→newest.
+    Splits results vs sweep by config-startswith('s=')."""
+    out = []
+    for run in load_runs(conn):
+        run_id = run["run_id"]
+        meta = {
+            "run_id":      run_id,
+            "timestamp":   run["timestamp"],
+            "git_commit":  run["git_commit"] or "?",
+            "git_branch":  run["git_branch"] or "?",
+        }
+        if run["gpu_name"] is not None:
+            meta["gpu_name"] = run["gpu_name"]
+        if run["gpu_driver"] is not None:
+            meta["gpu_driver"] = run["gpu_driver"]
+        if run["gpu_vram_mb"] is not None:
+            meta["gpu_vram_mb"] = run["gpu_vram_mb"]
+        if run["cuda_version"] is not None:
+            meta["cuda_version"] = run["cuda_version"]
+
+        records = load_results(conn, run_id)
+        comparison, sweep = [], []
+        for r in records:
+            (sweep if r.get("config", "").startswith("s=") else comparison).append(r)
+        out.append({"meta": meta, "results": comparison, "sweep": sweep})
+    return out
+
+
+# ── self-test ────────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    # In-memory smoke test: round-trip one record.
+    conn = sqlite3.connect(":memory:")
+    conn.row_factory = sqlite3.Row
+    conn.executescript(_SCHEMA)
+    insert_run(conn, run_id="test", timestamp="2026-04-27T00:00:00", mode="single")
+    insert_result(conn, "test", {
+        "path": "rust",
+        "model": "test-model",
+        "config": "default",
+        "ttft_ms": 12.34,
+        "ttft_ms_samples": [12.0, 12.5, 12.3],
+        "search_iters": 500,
+    })
+    [row] = load_results(conn, "test")
+    assert row["path"] == "rust", row
+    assert row["ttft_ms"] == 12.34, row
+    assert row["ttft_ms_samples"] == [12.0, 12.5, 12.3], row
+    assert latest_run_id(conn) == "test"
+    print("db.py smoke test ok")

benchmarks/ttft/gen_dashboard.py

@@ -0,0 +1,832 @@
+"""Time-series benchmark dashboard generator.
+
+Reads every run from the SQLite DB (benchmarks/ttft/bench.db) and produces a
+single standalone HTML file with Plotly.js charts styled to match luminal.com.
+
+Layout:
+  TTFT over time  →  one chart per model, lines = execution paths
+  TPOT over time  →  same
+
+Usage:
+  python3 benchmarks/ttft/gen_dashboard.py [--db PATH] [--out FILE]
+"""
+
+import argparse
+import json
+from datetime import datetime
+from pathlib import Path
+
+import db
+
+BENCH_DIR = Path(__file__).resolve().parent
+
+# Path colours – kept distinct against the dark green Luminal accent
+PATH_COLORS = {
+    "python_baseline":      "#5b5f61",  # muted slate
+    "python_torch_compile": "#3b82f6",  # blue (luminal accent palette)
+    "python_luminal":       "#a855f7",  # purple (luminal accent palette)
+    "rust":                 "#e8855a",  # warm orange – Rust brand feel
+}
+PATH_LABELS = {
+    "python_baseline":      "HF Baseline",
+    "python_torch_compile": "torch.compile",
+    "python_luminal":       "luminal backend",
+    "rust":                 "Rust (luminal)",
+}
+PATH_ORDER = ["python_baseline", "python_torch_compile", "python_luminal", "rust"]
+
+# (key, short label, y-axis label, scale, axis ticksuffix)
+# scale is applied to raw value before plotting (e.g. ms → sec via 0.001).
+METRICS = [
+    ("ttft_ms",    "TTFT",            "Time to first token (ms)",   1.0,   " ms"),
+    ("tpot_ms",    "TPOT",            "Time per output token (ms)", 1.0,   " ms"),
+    ("compile_ms", "Time to Search",  "Search time (sec)",          0.001, " sec"),
+]
+
+
+# ── data loading ─────────────────────────────────────────────────────────────
+
+def load_history(db_path: Path) -> list[dict]:
+    """Return [{"meta", "results", "sweep"}, …] from the bench DB,
+    oldest→newest. Same shape the legacy JSON loader returned."""
+    if not Path(db_path).exists():
+        return []
+    conn = db.connect(db_path)
+    return db.load_history(conn)
+
+
+def build_series(runs: list[dict]) -> tuple[dict, list[str], list[str]]:
+    """Returns (data, run_ids, run_labels).
+
+    - data[model][path][metric] = [(run_id, value, commit, ts), ...]
+      `run_id` is the categorical x value; `ts` is kept for tooltip formatting.
+    - run_ids: chronological list of every run that appears in the comparison data.
+    - run_labels: parallel to run_ids; "MMM DD · HH:MM" for nice axis ticks.
+
+    The categorical x-axis (one column per run_id) replaces the previous
+    `type: date` axis. With multiple runs on the same day, the date axis
+    silently stacked them on one column; the category axis spaces them
+    evenly so each run is visually distinct.
+    """
+    data: dict = {}
+    seen_run_ids: list[str] = []
+    seen_ts: dict[str, str] = {}
+
+    for run in runs:
+        run_id = run["meta"]["run_id"]
+        ts = run["meta"]["timestamp"]
+        commit = run["meta"].get("git_commit", "?")
+        had_data = False
+        for r in run["results"]:
+            if r.get("error") or r.get("ttft_ms") is None:
+                continue
+            model = r.get("config", r.get("model", "unknown"))
+            path = r.get("path", "unknown")
+            data.setdefault(model, {}).setdefault(path, {})
+            for metric, _, _, scale, _ in METRICS:
+                val = r.get(metric)
+                if val is not None:
+                    data[model][path].setdefault(metric, []).append(
+                        (run_id, val * scale, commit, ts)
+                    )
+                    had_data = True
+        if had_data and run_id not in seen_ts:
+            seen_run_ids.append(run_id)
+            seen_ts[run_id] = ts
+
+    run_ids = sorted(seen_run_ids, key=lambda rid: seen_ts.get(rid, rid))
+    run_labels = []
+    for rid in run_ids:
+        ts = seen_ts.get(rid, rid)
+        try:
+            run_labels.append(datetime.fromisoformat(ts).strftime("%b %d · %H:%M"))
+        except ValueError:
+            run_labels.append(rid[:16].replace("T", " "))
+    return data, run_ids, run_labels
+
+
+def build_sweep_series(runs: list[dict]) -> tuple[dict, list[str]]:
+    """Collect sweep records from ALL runs for 3D charting.
+
+    Returns:
+      data[model_key][path][metric][run_id] = {
+          "label":  str,              # short date label for Y axis
+          "commit": str,
+          "points": [(iters, ms), …]  # sorted by iters
+      }
+      run_ids: list[str] in chronological order (oldest → newest)
+    """
+    data: dict = {}
+    run_ids: list[str] = []
+
+    for run in runs:
+        if not run.get("sweep"):
+            continue
+        run_id = run["meta"]["run_id"]
+        commit = run["meta"].get("git_commit", "?")
+        try:
+            label = datetime.fromisoformat(run["meta"]["timestamp"]).strftime("%b %d")
+        except ValueError:
+            label = run_id[:10]
+        if run_id not in run_ids:
+            run_ids.append(run_id)
+
+        for r in run["sweep"]:
+            if r.get("error"):
+                continue
+            n = r.get("search_iters")
+            if n is None:
+                cfg = r.get("config", "")
+                if cfg.startswith("s="):
+                    try:
+                        n = int(cfg[2:])
+                    except ValueError:
+                        continue
+            if n is None:
+                continue
+            model_key = r.get("model_key", "unknown")
+            path = r.get("path", "unknown")
+            for metric, _, _, scale, _ in METRICS:
+                val = r.get(metric)
+                if val is None:
+                    continue
+                (data
+                    .setdefault(model_key, {})
+                    .setdefault(path, {})
+                    .setdefault(metric, {})
+                    .setdefault(run_id, {"label": label, "commit": commit, "points": []})
+                    ["points"].append((n, val * scale)))
+
+    # Sort points within each run by search_iters
+    for mk in data:
+        for path in data[mk]:
+            for metric in data[mk][path]:
+                for run_id in data[mk][path][metric]:
+                    data[mk][path][metric][run_id]["points"].sort(key=lambda x: x[0])
+
+    return data, run_ids
+
+
+# ── chart building ────────────────────────────────────────────────────────────
+
+def _traces_json(path_data: dict, metric: str, show_legend: bool, unit: str = " ms") -> str:
+    traces = []
+    for path in PATH_ORDER:
+        if path not in path_data or metric not in path_data[path]:
+            continue
+        pts = path_data[path][metric]
+        # pts: list of (run_id, val, commit, ts)
+        trace = {
+            "x": [p[0] for p in pts],
+            "y": [p[1] for p in pts],
+            "customdata": [[p[2], p[3]] for p in pts],
+            "type": "scatter",
+            "mode": "lines+markers",
+            "name": PATH_LABELS.get(path, path),
+            "line": {"color": PATH_COLORS.get(path, "#aaa"), "width": 2},
+            "marker": {"size": 7, "symbol": "circle"},
+            "connectgaps": False,
+            "showlegend": show_legend,
+            "hovertemplate": (
+                f"<b>{PATH_LABELS.get(path, path)}</b><br>"
+                "%{customdata[1]}<br>"
+                f"%{{y:.1f}}{unit}<br>"
+                "<span style='color:#7e8385'>commit %{customdata[0]}</span>"
+                "<extra></extra>"
+            ),
+        }
+        traces.append(trace)
+    return json.dumps(traces)
+
+
+_CHART_LAYOUT = {
+    "plot_bgcolor":  "#0d1416",
+    "paper_bgcolor": "#141b1d",
+    "font":          {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9"},
+    "margin":        {"t": 16, "b": 48, "l": 52, "r": 12},
+    "height":        280,
+    "xaxis": {
+        # Categorical: one column per run, evenly spaced. Same-day runs
+        # used to collapse on a date axis; this keeps every run distinct.
+        "type":          "category",
+        "categoryorder": "array",  # categoryarray injected per chart
+        "color":         "#5b5f61",
+        "gridcolor":     "#1c2225",
+        "linecolor":     "#2d3335",
+        "tickfont":      {"size": 11, "family": "Geist Mono, monospace"},
+        "tickangle":     -30,
+        "automargin":    True,
+        "zeroline":      False,
+    },
+    "yaxis": {
+        "rangemode": "tozero",
+        "color":     "#5b5f61",
+        "gridcolor": "#1c2225",
+        "linecolor": "#2d3335",
+        "tickfont":  {"size": 11, "family": "Geist Mono, monospace"},
+        "ticksuffix": " ms",
+        "zeroline":  False,
+    },
+    "legend": {
+        "orientation": "h",
+        "y": -0.28,
+        "x": 0,
+        "font": {"size": 11, "color": "#a1a4a5"},
+        "bgcolor": "rgba(0,0,0,0)",
+    },
+    "hoverlabel": {
+        "bgcolor":    "#1c2225",
+        "bordercolor":"#2d3335",
+        "font":       {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"},
+    },
+}
+
+
+def _chart_card(div_id: str, model: str, traces_json: str, show_legend: bool,
+                run_ids: list[str], run_labels: list[str], unit: str = " ms") -> str:
+    layout = dict(_CHART_LAYOUT)
+    xaxis = {
+        **layout["xaxis"],
+        "categoryarray": run_ids,
+        "tickvals":      run_ids,
+        "ticktext":      run_labels,
+    }
+    layout = {**layout,
+              "xaxis": xaxis,
+              "yaxis": {**layout["yaxis"], "ticksuffix": unit}}
+    if not show_legend:
+        layout = {**layout, "legend": {**layout["legend"], "visible": False},
+                  "margin": {**layout["margin"], "b": 16}}
+    return f"""<div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">{model}</span>
+  </div>
+  <div id="{div_id}"></div>
+  <script>
+    Plotly.newPlot("{div_id}", {traces_json}, {json.dumps(layout)},
+      {{responsive: true, displayModeBar: false}});
+  </script>
+</div>"""
+
+
+def _sweep_3d_traces_json(model_data: dict, metric: str, run_ids: list[str], unit: str = " ms") -> str:
+    """One scatter3d trace per (path, run) — same colour per path, stacked by run on Y."""
+    traces = []
+    path_legend_shown: set[str] = set()
+
+    for run_id in run_ids:
+        for path in PATH_ORDER:
+            run_map = model_data.get(path, {}).get(metric, {})
+            if run_id not in run_map:
+                continue
+            entry = run_map[run_id]
+            pts = entry["points"]
+            label = entry["label"]
+            commit = entry["commit"]
+            color = PATH_COLORS.get(path, "#aaa")
+            show_legend = path not in path_legend_shown
+            path_legend_shown.add(path)
+
+            traces.append({
+                "type": "scatter3d",
+                "mode": "lines+markers",
+                "x": [p[0] for p in pts],   # search iters
+                "y": [label] * len(pts),     # run label (categorical)
+                "z": [p[1] for p in pts],    # value (already scaled by build_sweep_series)
+                "name": PATH_LABELS.get(path, path),
+                "legendgroup": path,
+                "showlegend": show_legend,
+                "line":   {"color": color, "width": 5},
+                "marker": {"color": color, "size": 4},
+                "hovertemplate": (
+                    f"<b>{PATH_LABELS.get(path, path)}</b><br>"
+                    f"s=%{{x}} iters<br>%{{z:.1f}}{unit}<br>"
+                    f"{label} · {commit}"
+                    "<extra></extra>"
+                ),
+            })
+
+    # Cross-run wire lines: for each path, connect same-budget points across
+    # runs. Makes regressions at a fixed search budget visible as a kink in the
+    # wireframe. Dashed + thinner than the per-run curves; legendgroup matches
+    # the path so toggling one toggles both.
+    for path in PATH_ORDER:
+        metric_runs = model_data.get(path, {}).get(metric, {})
+        if len(metric_runs) < 2:
+            continue
+        color = PATH_COLORS.get(path, "#aaa")
+        # by_budget[iters] -> list of (run_label, value) in chronological order
+        by_budget: dict = {}
+        for run_id in run_ids:
+            if run_id not in metric_runs:
+                continue
+            entry = metric_runs[run_id]
+            for iters, val in entry["points"]:
+                by_budget.setdefault(iters, []).append((entry["label"], val))
+        for budget, items in sorted(by_budget.items()):
+            if len(items) < 2:
+                continue
+            traces.append({
+                "type": "scatter3d",
+                "mode": "lines",
+                "x": [budget] * len(items),
+                "y": [it[0] for it in items],
+                "z": [it[1] for it in items],
+                "legendgroup": path,
+                "showlegend": False,
+                "line": {"color": color, "width": 2, "dash": "dash"},
+                "hovertemplate": (
+                    f"<b>{PATH_LABELS.get(path, path)} @ s={budget}</b><br>"
+                    f"%{{y}}: %{{z:.1f}}{unit}"
+                    "<extra></extra>"
+                ),
+            })
+    return json.dumps(traces)
+
+
+_SWEEP_3D_LAYOUT = {
+    "paper_bgcolor": "#141b1d",
+    "font": {"family": "Geist, system-ui, sans-serif", "color": "#d7d8d9", "size": 11},
+    "height": 420,
+    "margin": {"t": 20, "b": 0, "l": 0, "r": 0},
+    "legend": {
+        "orientation": "h",
+        "y": -0.05,
+        "x": 0,
+        "font": {"size": 11, "color": "#a1a4a5", "family": "Geist Mono, monospace"},
+        "bgcolor": "rgba(0,0,0,0)",
+    },
+    "hoverlabel": {
+        "bgcolor": "#1c2225",
+        "bordercolor": "#2d3335",
+        "font": {"size": 12, "color": "#d7d8d9", "family": "Geist Mono, monospace"},
+    },
+    "scene": {
+        "bgcolor": "#0d1416",
+        "xaxis": {
+            "title": {"text": "search iters", "font": {"size": 10, "color": "#7e8385"}},
+            "type": "log",
+            "tickvals": [5, 10, 20, 50, 100, 500],
+            "ticktext": ["5", "10", "20", "50", "100", "500"],
+            "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"},
+            "gridcolor": "#1c2225",
+            "linecolor": "#2d3335",
+            "zerolinecolor": "#2d3335",
+        },
+        "yaxis": {
+            "title": {"text": "run", "font": {"size": 10, "color": "#7e8385"}},
+            "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"},
+            "gridcolor": "#1c2225",
+            "linecolor": "#2d3335",
+        },
+        "zaxis": {
+            "title": {"text": "ms", "font": {"size": 10, "color": "#7e8385"}},
+            "rangemode": "tozero",
+            "tickfont": {"size": 10, "family": "Geist Mono, monospace", "color": "#5b5f61"},
+            "ticksuffix": " ms",
+            "gridcolor": "#1c2225",
+            "linecolor": "#2d3335",
+        },
+        "camera": {
+            "eye": {"x": 1.6, "y": -1.6, "z": 0.9},
+        },
+    },
+}
+
+
+def _sweep_3d_card(div_id: str, model: str, traces_json: str, unit: str = " ms") -> str:
+    layout = {**_SWEEP_3D_LAYOUT,
+              "scene": {**_SWEEP_3D_LAYOUT["scene"],
+                        "zaxis": {**_SWEEP_3D_LAYOUT["scene"]["zaxis"],
+                                  "title": {**_SWEEP_3D_LAYOUT["scene"]["zaxis"]["title"],
+                                            "text": unit.strip()},
+                                  "ticksuffix": unit}}}
+    return f"""<div class="chart-card">
+  <div class="chart-card-header">
+    <span class="model-tag">{model}</span>
+  </div>
+  <div id="{div_id}"></div>
+  <script>
+    Plotly.newPlot("{div_id}", {traces_json}, {json.dumps(layout)},
+      {{responsive: true, displayModeBar: true, displaylogo: false,
+        modeBarButtonsToRemove: ["toImage","sendDataToCloud"]}});
+  </script>
+</div>"""
+
+
+# ── HTML assembly ─────────────────────────────────────────────────────────────
+
+def build_html(runs: list[dict], data: dict,
+               run_ids: list[str], run_labels: list[str],
+               sweep_data: dict | None = None,
+               sweep_run_ids: list[str] | None = None) -> str:
+    # Preserve insertion order of models as seen across runs
+    models = list(dict.fromkeys(
+        r["config"]
+        for run in runs
+        for r in run["results"]
+        if not r.get("config", "").startswith("s=") and not r.get("error")
+    ))
+
+    last_ts = ""
+    if runs:
+        raw = runs[-1]["meta"]["timestamp"]
+        try:
+            last_ts = datetime.fromisoformat(raw).strftime("%b %d, %Y · %H:%M")
+        except ValueError:
+            last_ts = raw[:16].replace("T", " ")
+
+    n_runs = len(runs)
+
+    sections_html = ""
+    for metric_key, metric_label, ylabel, _scale, unit in METRICS:
+        active_models = [
+            m for m in models
+            if any(metric_key in data.get(m, {}).get(p, {}) for p in PATH_ORDER)
+        ]
+        if not active_models:
+            continue
+
+        cards_html = ""
+        first = True
+        for model in active_models:
+            path_data = data.get(model, {})
+            div_id = f"c_{metric_key}_{model.replace('-','_').replace('.','_')}"
+            traces = _traces_json(path_data, metric_key, show_legend=first, unit=unit)
+            cards_html += _chart_card(div_id, model, traces, show_legend=first,
+                                      run_ids=run_ids, run_labels=run_labels, unit=unit)
+            first = False
+
+        n = len(active_models)
+        # Clamp columns so charts don't get too narrow; wrap at 4
+        cols = min(n, 4)
+        sections_html += f"""
+<section>
+  <div class="section-header">
+    <span class="section-eyebrow">metric</span>
+    <h2 class="section-title">{metric_label} <span class="unit">over time</span></h2>
+    <span class="section-tag">{ylabel}</span>
+  </div>
+  <div class="chart-grid" style="grid-template-columns: repeat({cols}, 1fr)">
+{cards_html}
+  </div>
+</section>"""
+
+    # ── sweep sections (3D) ──────────────────────────────────────────────────
+    sweep_sections_html = ""
+    if sweep_data and sweep_run_ids:
+        sweep_models = list(sweep_data.keys())
+        for metric_key, metric_label, ylabel, _scale, unit in METRICS:
+            active = [
+                m for m in sweep_models
+                if any(
+                    run_id in sweep_data[m].get(p, {}).get(metric_key, {})
+                    for p in PATH_ORDER
+                    for run_id in sweep_run_ids
+                )
+            ]
+            if not active:
+                continue
+            cards_html = ""
+            for model in active:
+                div_id = f"sw_{metric_key}_{model.replace('-','_').replace('.','_')}"
+                traces = _sweep_3d_traces_json(sweep_data[model], metric_key, sweep_run_ids, unit=unit)
+                cards_html += _sweep_3d_card(div_id, model, traces, unit=unit)
+            cols = min(len(active), 4)
+            run_count = len(sweep_run_ids)
+            sweep_sections_html += f"""
+<section>
+  <div class="section-header">
+    <span class="section-eyebrow">sweep · 3d</span>
+    <h2 class="section-title">{metric_label} <span class="unit">vs search budget · over time</span></h2>
+    <span class="section-tag">{run_count} run{"s" if run_count != 1 else ""}</span>
+  </div>
+  <p class="sweep-hint">Drag to rotate · scroll to zoom · each curve = one run</p>
+  <div class="chart-grid" style="grid-template-columns: repeat({cols}, 1fr)">
+{cards_html}
+  </div>
+</section>"""
+
+    return f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>Luminal · Benchmark Dashboard</title>
+<link rel="preconnect" href="https://fonts.googleapis.com">
+<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+<link href="https://fonts.googleapis.com/css2?family=Geist:wght@300;400;500;600&family=Geist+Mono:wght@300;400;500&display=swap" rel="stylesheet">
+<script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
+<style>
+*, *::before, *::after {{ box-sizing: border-box; margin: 0; padding: 0; }}
+html {{ -webkit-font-smoothing: antialiased; scroll-behavior: smooth; }}
+
+body {{
+  font-family: 'Geist', system-ui, sans-serif;
+  background: #030712;
+  color: #d7d8d9;
+  min-height: 100vh;
+  line-height: 1.5;
+}}
+
+/* ── NAV ── */
+nav {{
+  position: sticky;
+  top: 0;
+  z-index: 50;
+  height: 56px;
+  background: rgba(8, 15, 17, 0.92);
+  backdrop-filter: blur(8px);
+  -webkit-backdrop-filter: blur(8px);
+  border-bottom: 1px solid #2d3335;
+  display: flex;
+  align-items: center;
+  padding: 0 24px;
+  gap: 0;
+}}
+.nav-brand {{
+  display: flex;
+  align-items: center;
+  gap: 8px;
+  font-family: 'Geist Mono', monospace;
+  font-size: 14px;
+  font-weight: 500;
+  letter-spacing: 0.05em;
+  color: #2faa6e;
+  text-decoration: none;
+}}
+.nav-dot {{
+  width: 6px;
+  height: 6px;
+  background: #2faa6e;
+  border-radius: 50%;
+  flex-shrink: 0;
+  animation: pulse-glow 2s ease-in-out infinite;
+}}
+.nav-sep {{
+  color: #2d3335;
+  margin: 0 14px;
+  font-size: 18px;
+  font-weight: 300;
+}}
+.nav-page {{
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.1em;
+  text-transform: uppercase;
+  color: #7e8385;
+}}
+
+@keyframes pulse-glow {{
+  0%, 100% {{ opacity: 1; }}
+  50%       {{ opacity: 0.35; }}
+}}
+
+/* ── MAIN ── */
+main {{
+  max-width: 1200px;
+  margin: 0 auto;
+  padding: 40px 24px 80px;
+}}
+
+/* ── PAGE HEADER ── */
+.page-header {{
+  margin-bottom: 40px;
+  padding-bottom: 32px;
+  border-bottom: 1px solid #1c2225;
+}}
+.page-eyebrow {{
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.1em;
+  text-transform: uppercase;
+  color: #2faa6e;
+  margin-bottom: 10px;
+}}
+.page-title {{
+  font-size: 30px;
+  font-weight: 500;
+  letter-spacing: -0.025em;
+  color: #d7d8d9;
+  margin-bottom: 10px;
+}}
+.page-meta {{
+  font-size: 14px;
+  color: #7e8385;
+  display: flex;
+  align-items: center;
+  gap: 0;
+  flex-wrap: wrap;
+}}
+.meta-sep {{
+  font-family: 'Geist Mono', monospace;
+  color: #2d3335;
+  margin: 0 10px;
+}}
+.meta-val {{
+  font-family: 'Geist Mono', monospace;
+  font-size: 13px;
+  color: #5b5f61;
+}}
+
+/* ── LEGEND STRIP ── */
+.legend-strip {{
+  display: flex;
+  flex-wrap: wrap;
+  gap: 6px;
+  margin-bottom: 32px;
+}}
+.legend-pill {{
+  display: flex;
+  align-items: center;
+  gap: 6px;
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.04em;
+  color: #a1a4a5;
+  background: #141b1d;
+  border: 1px solid #2d3335;
+  border-radius: 2px;
+  padding: 4px 10px;
+}}
+.legend-swatch {{
+  width: 8px;
+  height: 8px;
+  border-radius: 50%;
+  flex-shrink: 0;
+}}
+
+/* ── SECTIONS ── */
+section {{ margin-bottom: 48px; }}
+.section-header {{
+  display: flex;
+  align-items: baseline;
+  gap: 10px;
+  margin-bottom: 16px;
+  padding-bottom: 12px;
+  border-bottom: 1px solid #1c2225;
+  flex-wrap: wrap;
+}}
+.section-eyebrow {{
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.1em;
+  text-transform: uppercase;
+  color: #404647;
+}}
+.section-title {{
+  font-size: 18px;
+  font-weight: 500;
+  color: #d7d8d9;
+  letter-spacing: -0.01em;
+}}
+.section-title .unit {{
+  color: #7e8385;
+  font-weight: 400;
+}}
+.section-tag {{
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.04em;
+  text-transform: uppercase;
+  color: #2faa6e;
+  background: #162322;
+  border: 1px solid #1c372e;
+  padding: 2px 8px;
+  border-radius: 2px;
+  margin-left: auto;
+}}
+
+/* ── CHART GRID ── */
+.chart-grid {{
+  display: grid;
+  gap: 10px;
+}}
+.chart-card {{
+  background: #141b1d;
+  border: 1px solid #2d3335;
+  border-radius: 2px;
+  overflow: hidden;
+  transition: border-color 150ms;
+  min-width: 0;
+}}
+.chart-card:hover {{ border-color: #404647; }}
+.chart-card-header {{
+  padding: 10px 14px 0;
+  display: flex;
+  align-items: center;
+}}
+.model-tag {{
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.06em;
+  text-transform: uppercase;
+  color: #7e8385;
+}}
+
+/* ── FOOTER ── */
+footer {{
+  max-width: 1200px;
+  margin: 0 auto;
+  padding: 20px 24px;
+  border-top: 1px solid #1c2225;
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.04em;
+  color: #404647;
+  display: flex;
+  justify-content: space-between;
+  flex-wrap: wrap;
+  gap: 8px;
+}}
+
+.section-divider {{
+  border: none;
+  border-top: 1px solid #1c2225;
+  margin: 8px 0 40px;
+}}
+.sweep-hint {{
+  font-family: 'Geist Mono', monospace;
+  font-size: 11px;
+  letter-spacing: 0.04em;
+  color: #404647;
+  margin-bottom: 12px;
+}}
+
+@media (max-width: 768px) {{
+  .chart-grid {{ grid-template-columns: 1fr !important; }}
+  .page-title {{ font-size: 22px; }}
+}}
+</style>
+</head>
+<body>
+
+<nav>
+  <a class="nav-brand" href="https://luminal.com">
+    <span class="nav-dot"></span>luminal
+  </a>
+  <span class="nav-sep">/</span>
+  <span class="nav-page">benchmarks</span>
+</nav>
+
+<main>
+
+<header class="page-header">
+  <p class="page-eyebrow">performance · time-series</p>
+  <h1 class="page-title">Benchmark Dashboard</h1>
+  <div class="page-meta">
+    <span>Last updated</span>
+    <span class="meta-sep">·</span>
+    <span class="meta-val">{last_ts}</span>
+    <span class="meta-sep">·</span>
+    <span class="meta-val">{n_runs} run{"s" if n_runs != 1 else ""} in history</span>
+  </div>
+</header>
+
+<div class="legend-strip">
+  {"".join(
+      f'<div class="legend-pill"><span class="legend-swatch" style="background:{PATH_COLORS[p]}"></span>{PATH_LABELS[p]}</div>'
+      for p in PATH_ORDER
+  )}
+</div>
+
+{sections_html}
+{"<hr class='section-divider'>" + sweep_sections_html if sweep_sections_html else ""}
+
+</main>
+
+<footer>
+  <span>luminal · benchmark dashboard</span>
+  <span>generated {last_ts}</span>
+</footer>
+
+</body>
+</html>
+"""
+
+
+# ── entry point ───────────────────────────────────────────────────────────────
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--db", default=str(db.DEFAULT_DB_PATH),
+                    help=f"SQLite bench DB (default: {db.DEFAULT_DB_PATH})")
+    ap.add_argument("--out", default=str(BENCH_DIR / "dashboard.html"),
+                    help="Output HTML file")
+    args = ap.parse_args()
+
+    runs = load_history(Path(args.db))
+    if not runs:
+        print(f"No runs found in {args.db}. Run --ur-test (or backfill) first.")
+        return
+
+    data, run_ids, run_labels = build_series(runs)
+    sweep_data, sweep_run_ids = build_sweep_series(runs)
+    html = build_html(runs, data, run_ids, run_labels, sweep_data, sweep_run_ids)
+    Path(args.out).write_text(html)
+
+    print(f"wrote {args.out}  ({len(runs)} runs, {sum(len(v) for v in data.values())} model×path series)")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ttft/gen_report.py

@@ -0,0 +1,349 @@
+#!/usr/bin/env python3
+"""Generate a standalone HTML benchmark report from a single benchmark run.
+
+Usage:
+    python3 gen_report.py [--db PATH] [--run RUN_ID] [--out report.html] [--title "..."]
+
+Sections are split out of a single run automatically:
+  - per-model_key, "comparison" (configs not matching s=N)  →  grouped bar chart
+  - per-model_key, "sweep" (configs matching s=N)           →  line chart (log X)
+For runs without model_key (e.g. single-config runs), one section per detected
+shape is produced instead.
+"""
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+
+import db
+
+PATH_ORDER = ["python_baseline", "python_torch_compile", "python_luminal", "rust"]
+PATH_LABELS = {
+    "python_baseline":      "HF Baseline",
+    "python_torch_compile": "torch.compile",
+    "python_luminal":       "luminal backend",
+    "rust":                 "Rust (luminal)",
+}
+PATH_COLORS = {
+    "python_baseline":      "#888888",
+    "python_torch_compile": "#5ab552",
+    "python_luminal":       "#4c9ed9",
+    "rust":                 "#d97a4c",
+}
+
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+def _fmt(v, decimals=1, suffix=""):
+    return f"{v:.{decimals}f}{suffix}" if v is not None else "—"
+
+def _section_title(path: Path) -> str:
+    stem = path.stem.replace("_", " ").replace("-", " ")
+    return stem.title()
+
+def _is_sweep(configs: list[str]) -> bool:
+    return bool(configs) and all(re.fullmatch(r"s=\d+", c) for c in configs)
+
+def _group_by_config(results: list[dict]) -> dict[str, dict[str, dict]]:
+    """Return {config: {path: result_dict}}."""
+    out: dict[str, dict[str, dict]] = {}
+    for r in results:
+        cfg = r.get("config", "default")
+        out.setdefault(cfg, {})[r["path"]] = r
+    return out
+
+
+# ── chart builders (return Plotly figure dicts) ───────────────────────────────
+
+def _bar_figure(by_config: dict, metric: str, title: str,
+                scale: float = 1.0, unit: str = "ms") -> dict:
+    configs = list(by_config.keys())
+    traces = []
+    for path in PATH_ORDER:
+        ys, texts = [], []
+        for cfg in configs:
+            r = by_config[cfg].get(path)
+            raw = r.get(metric) if r and not r.get("error") else None
+            v = raw * scale if raw is not None else None
+            ys.append(v if v is not None else 0)
+            texts.append(f"{v:.1f} {unit}" if v is not None else "n/a")
+        if any(y > 0 for y in ys):
+            traces.append({
+                "type": "bar",
+                "name": PATH_LABELS.get(path, path),
+                "x": configs,
+                "y": ys,
+                "text": texts,
+                "textposition": "outside",
+                "marker": {"color": PATH_COLORS.get(path, "#aaaaaa")},
+                "hovertemplate": "%{x}<br>" + PATH_LABELS.get(path, path)
+                                 + f": %{{y:.1f}} {unit}<extra></extra>",
+            })
+    return {
+        "data": traces,
+        "layout": {
+            "title": title,
+            "yaxis": {"title": unit, "rangemode": "tozero"},
+            "barmode": "group",
+            "legend": {"orientation": "h", "y": -0.2},
+            "margin": {"t": 50, "b": 80},
+            "plot_bgcolor": "#fafafa",
+            "paper_bgcolor": "#ffffff",
+        },
+    }
+
+
+def _line_figure(by_config: dict, metric: str, title: str,
+                 scale: float = 1.0, unit: str = "ms") -> dict:
+    """Line chart for sweep data. Config names are 's=N'; X = N (log scale)."""
+    def _iter(cfg):
+        m = re.fullmatch(r"s=(\d+)", cfg)
+        return int(m.group(1)) if m else 0
+
+    configs_sorted = sorted(by_config.keys(), key=_iter)
+    xs = [_iter(c) for c in configs_sorted]
+
+    paths_present = {p for cfg in by_config.values() for p in cfg}
+    traces = []
+    for path in PATH_ORDER:
+        if path not in paths_present:
+            continue
+        ys = []
+        for cfg in configs_sorted:
+            r = by_config[cfg].get(path)
+            raw = r.get(metric) if r and not r.get("error") else None
+            ys.append(raw * scale if raw is not None else None)
+        if any(y is not None for y in ys):
+            traces.append({
+                "type": "scatter",
+                "mode": "lines+markers",
+                "name": PATH_LABELS.get(path, path),
+                "x": xs,
+                "y": ys,
+                "marker": {"size": 8, "color": PATH_COLORS.get(path, "#aaaaaa")},
+                "line": {"color": PATH_COLORS.get(path, "#aaaaaa"), "width": 2},
+                "hovertemplate": "iters=%{x}<br>" + PATH_LABELS.get(path, path)
+                                 + f": %{{y:.1f}} {unit}<extra></extra>",
+            })
+    return {
+        "data": traces,
+        "layout": {
+            "title": title,
+            "xaxis": {"title": "Search iterations", "type": "log",
+                      "tickvals": xs, "ticktext": [str(x) for x in xs]},
+            "yaxis": {"title": unit, "rangemode": "tozero"},
+            "legend": {"orientation": "h", "y": -0.25},
+            "margin": {"t": 50, "b": 90},
+            "plot_bgcolor": "#fafafa",
+            "paper_bgcolor": "#ffffff",
+        },
+    }
+
+
+# ── table builder ─────────────────────────────────────────────────────────────
+
+def _table_html(results: list[dict]) -> str:
+    rows = []
+    for r in sorted(results, key=lambda r: (r.get("config", ""), PATH_ORDER.index(r["path"]) if r["path"] in PATH_ORDER else 99)):
+        error = r.get("error")
+        style = ' style="background:#fff0f0"' if error else ""
+        path_label = PATH_LABELS.get(r["path"], r["path"])
+        cfg = r.get("config", "—")
+        ttft = _fmt(r.get("ttft_ms"), 1, " ms")
+        tpot = _fmt(r.get("tpot_ms"), 1, " ms")
+        tput = _fmt(r.get("throughput_tps"), 1, " tok/s")
+        comp = _fmt(r.get("compile_ms"), 0, " ms") if r.get("compile_ms") else "—"
+        ptok = str(r.get("prompt_tokens", "—"))
+        note = (r.get("error") or r.get("note") or "")[:90]
+        note_style = ' style="color:#c00"' if error else ' style="color:#777"'
+        rows.append(
+            f'<tr{style}>'
+            f'<td>{path_label}</td><td>{cfg}</td>'
+            f'<td>{ttft}</td><td>{tpot}</td><td>{tput}</td>'
+            f'<td>{comp}</td><td>{ptok}</td>'
+            f'<td{note_style}>{note}</td>'
+            f'</tr>'
+        )
+    return (
+        '<table>'
+        '<thead><tr>'
+        '<th>Path</th><th>Config</th>'
+        '<th>TTFT</th><th>TPOT</th><th>Throughput</th>'
+        '<th>Compile</th><th>Prompt tokens</th><th>Note</th>'
+        '</tr></thead>'
+        '<tbody>' + "\n".join(rows) + '</tbody>'
+        '</table>'
+    )
+
+
+# ── section builder ───────────────────────────────────────────────────────────
+
+def _section_html(sec_id: str, title: str, results: list[dict], fig_counter: list) -> str:
+    by_config = _group_by_config(results)
+    configs = list(by_config.keys())
+    sweep = _is_sweep(configs)
+
+    models = list(dict.fromkeys(r.get("model", "") for r in results if r.get("model")))
+    model_str = ", ".join(models) if models else "—"
+    prompt_tokens = list(dict.fromkeys(r.get("prompt_tokens") for r in results if r.get("prompt_tokens")))
+    tok_str = "/".join(str(t) for t in prompt_tokens) + " prompt tokens" if prompt_tokens else ""
+
+    builder = _line_figure if sweep else _bar_figure
+    ttft_fig = builder(by_config, "ttft_ms", "TTFT")
+    has_tpot = any(r.get("tpot_ms") is not None for r in results if not r.get("error"))
+    tpot_fig = builder(by_config, "tpot_ms", "TPOT") if has_tpot else None
+    has_compile = any(r.get("compile_ms") is not None and r.get("compile_ms") > 0
+                      for r in results if not r.get("error"))
+    compile_fig = (builder(by_config, "compile_ms", "Time to Search",
+                           scale=0.001, unit="sec")
+                   if has_compile else None)
+
+    def chart_div(fig):
+        n = fig_counter[0]
+        fig_counter[0] += 1
+        return (
+            f'<div id="fig{n}" class="chart"></div>'
+            f'<script>Plotly.newPlot("fig{n}", {json.dumps(fig["data"])}, {json.dumps(fig["layout"])}, {{responsive:true}});</script>'
+        )
+
+    charts_html = f'<div class="charts-row">{chart_div(ttft_fig)}'
+    if tpot_fig:
+        charts_html += chart_div(tpot_fig)
+    if compile_fig:
+        charts_html += chart_div(compile_fig)
+    charts_html += '</div>'
+
+    return f"""
+<section id="{sec_id}">
+  <h2>{title}</h2>
+  <p class="meta">{model_str}{" · " + tok_str if tok_str else ""} · {len(results)} results</p>
+  {charts_html}
+  {_table_html(results)}
+</section>
+"""
+
+
+# ── full page ─────────────────────────────────────────────────────────────────
+
+CSS = """
+* { box-sizing: border-box; margin: 0; padding: 0; }
+body { font-family: system-ui, sans-serif; background: #f0f2f5; color: #222; }
+header { background: #1a1a2e; color: #fff; padding: 1rem 2rem;
+         position: sticky; top: 0; z-index: 100; display: flex;
+         align-items: center; gap: 2rem; }
+header h1 { font-size: 1.2rem; white-space: nowrap; }
+nav a { color: #a0c4ff; text-decoration: none; font-size: 0.9rem;
+        padding: 0.3rem 0.7rem; border-radius: 4px; white-space: nowrap; }
+nav a:hover { background: rgba(255,255,255,0.15); }
+main { max-width: 1400px; margin: 0 auto; padding: 2rem; display: flex;
+       flex-direction: column; gap: 2.5rem; }
+section { background: #fff; border-radius: 8px; padding: 1.5rem 2rem;
+          box-shadow: 0 1px 4px rgba(0,0,0,.08); }
+h2 { font-size: 1.3rem; margin-bottom: 0.4rem; }
+.meta { color: #666; font-size: 0.85rem; margin-bottom: 1.2rem; }
+.charts-row { display: flex; gap: 1.5rem; flex-wrap: wrap; margin-bottom: 1.5rem; }
+.chart { flex: 1; min-width: 340px; height: 360px; }
+table { width: 100%; border-collapse: collapse; font-size: 0.82rem; }
+thead tr { background: #f5f5f5; }
+th, td { padding: 0.45rem 0.7rem; text-align: left;
+         border-bottom: 1px solid #e8e8e8; }
+th { font-weight: 600; white-space: nowrap; }
+tr:last-child td { border-bottom: none; }
+tr:hover { background: #fafafa; }
+"""
+
+def _build_html(sections: list[tuple[str, str, list[dict]]], title: str) -> str:
+    nav_links = "".join(f'<a href="#{sid}">{stitle}</a>' for sid, stitle, _ in sections)
+    fig_counter = [0]
+    body = "".join(_section_html(sid, stitle, results, fig_counter)
+                   for sid, stitle, results in sections)
+    return f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>{title}</title>
+<script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
+<style>{CSS}</style>
+</head>
+<body>
+<header>
+  <h1>{title}</h1>
+  <nav>{nav_links}</nav>
+</header>
+<main>{body}</main>
+</body>
+</html>"""
+
+
+# ── CLI ───────────────────────────────────────────────────────────────────────
+
+def _sections_for_run(results: list[dict]) -> list[tuple[str, str, list[dict]]]:
+    """Split a single run's results into (sec_id, title, records) sections.
+
+    Splits first by model_key (NULL → 'results'), then within each by
+    sweep-vs-comparison based on config 's=N' shape."""
+    by_key: dict[str | None, list[dict]] = {}
+    for r in results:
+        by_key.setdefault(r.get("model_key"), []).append(r)
+
+    sections: list[tuple[str, str, list[dict]]] = []
+    for key, recs in by_key.items():
+        comp, sweep = [], []
+        for r in recs:
+            (sweep if str(r.get("config", "")).startswith("s=") else comp).append(r)
+        prefix = (key or "results").replace("-", "_").replace(".", "_")
+        title_prefix = key or "Results"
+        if comp:
+            sections.append((f"{prefix}_comparison",
+                             f"{title_prefix} comparison".strip().title(),
+                             comp))
+        if sweep:
+            sections.append((f"{prefix}_sweep",
+                             f"{title_prefix} sweep".strip().title(),
+                             sweep))
+    return sections
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--db", default=str(db.DEFAULT_DB_PATH),
+                    help=f"SQLite bench DB (default: {db.DEFAULT_DB_PATH})")
+    ap.add_argument("--run", default=None,
+                    help="Run ID to render (default: latest run in DB)")
+    ap.add_argument("--out", default=None,
+                    help="Output HTML path (default: report.html in benchmarks/ttft/)")
+    ap.add_argument("--title", default="Luminal TTFT Benchmark Report",
+                    help="Page title and heading")
+    args = ap.parse_args()
+
+    if not Path(args.db).exists():
+        print(f"DB not found: {args.db}", file=sys.stderr)
+        sys.exit(1)
+
+    conn = db.connect(args.db)
+    run_id = args.run or db.latest_run_id(conn)
+    if run_id is None:
+        print(f"No runs in {args.db}", file=sys.stderr)
+        sys.exit(1)
+
+    results = db.load_results(conn, run_id)
+    if not results:
+        print(f"No results for run {run_id}", file=sys.stderr)
+        sys.exit(1)
+
+    sections = _sections_for_run(results)
+    if not sections:
+        print(f"No section data for run {run_id}", file=sys.stderr)
+        sys.exit(1)
+
+    out = Path(args.out) if args.out else Path(__file__).parent / "report.html"
+    html = _build_html(sections, f"{args.title} — {run_id}")
+    out.write_text(html)
+    print(f"wrote {out}  (run {run_id}, {len(sections)} sections, {len(results)} results)")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ttft/report.html

@@ -0,0 +1,148 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>Luminal TTFT Benchmark Report — 2026-05-01T18-56-26-996695</title>
+<script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
+<style>
+* { box-sizing: border-box; margin: 0; padding: 0; }
+body { font-family: system-ui, sans-serif; background: #f0f2f5; color: #222; }
+header { background: #1a1a2e; color: #fff; padding: 1rem 2rem;
+         position: sticky; top: 0; z-index: 100; display: flex;
+         align-items: center; gap: 2rem; }
+header h1 { font-size: 1.2rem; white-space: nowrap; }
+nav a { color: #a0c4ff; text-decoration: none; font-size: 0.9rem;
+        padding: 0.3rem 0.7rem; border-radius: 4px; white-space: nowrap; }
+nav a:hover { background: rgba(255,255,255,0.15); }
+main { max-width: 1400px; margin: 0 auto; padding: 2rem; display: flex;
+       flex-direction: column; gap: 2.5rem; }
+section { background: #fff; border-radius: 8px; padding: 1.5rem 2rem;
+          box-shadow: 0 1px 4px rgba(0,0,0,.08); }
+h2 { font-size: 1.3rem; margin-bottom: 0.4rem; }
+.meta { color: #666; font-size: 0.85rem; margin-bottom: 1.2rem; }
+.charts-row { display: flex; gap: 1.5rem; flex-wrap: wrap; margin-bottom: 1.5rem; }
+.chart { flex: 1; min-width: 340px; height: 360px; }
+table { width: 100%; border-collapse: collapse; font-size: 0.82rem; }
+thead tr { background: #f5f5f5; }
+th, td { padding: 0.45rem 0.7rem; text-align: left;
+         border-bottom: 1px solid #e8e8e8; }
+th { font-weight: 600; white-space: nowrap; }
+tr:last-child td { border-bottom: none; }
+tr:hover { background: #fafafa; }
+</style>
+</head>
+<body>
+<header>
+  <h1>Luminal TTFT Benchmark Report — 2026-05-01T18-56-26-996695</h1>
+  <nav><a href="#llama_8b_comparison">Llama-8B Comparison</a><a href="#llama_8b_sweep">Llama-8B Sweep</a><a href="#qwen3_4b_comparison">Qwen3-4B Comparison</a><a href="#qwen3_4b_sweep">Qwen3-4B Sweep</a><a href="#gemma3_4b_comparison">Gemma3-4B Comparison</a><a href="#gemma3_4b_sweep">Gemma3-4B Sweep</a><a href="#gemma4_moe_comparison">Gemma4-Moe Comparison</a><a href="#gemma4_moe_sweep">Gemma4-Moe Sweep</a><a href="#qwen3_moe_comparison">Qwen3-Moe Comparison</a><a href="#qwen3_moe_sweep">Qwen3-Moe Sweep</a></nav>
+</header>
+<main>
+<section id="llama_8b_comparison">
+  <h2>Llama-8B Comparison</h2>
+  <p class="meta">NousResearch/Meta-Llama-3-8B-Instruct · 21 prompt tokens · 4 results</p>
+  <div class="charts-row"><div id="fig0" class="chart"></div><script>Plotly.newPlot("fig0", [{"type": "bar", "name": "HF Baseline", "x": ["llama-8b"], "y": [705.9654394979589], "text": ["706.0 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["llama-8b"], "y": [307.66548847896047], "text": ["307.7 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["llama-8b"], "y": [461.48114453535527], "text": ["461.5 ms"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["llama-8b"], "y": [1026.86], "text": ["1026.9 ms"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig1" class="chart"></div><script>Plotly.newPlot("fig1", [{"type": "bar", "name": "HF Baseline", "x": ["llama-8b"], "y": [34.15271903970279], "text": ["34.2 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["llama-8b"], "y": [171.7862353892997], "text": ["171.8 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["llama-8b"], "y": [23.078908618772402], "text": ["23.1 ms"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["llama-8b"], "y": [51.64], "text": ["51.6 ms"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig2" class="chart"></div><script>Plotly.newPlot("fig2", [{"type": "bar", "name": "torch.compile", "x": ["llama-8b"], "y": [18.760145067994017], "text": ["18.8 sec"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} sec<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["llama-8b"], "y": [95.96263545705006], "text": ["96.0 sec"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} sec<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["llama-8b"], "y": [84.45343], "text": ["84.5 sec"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "yaxis": {"title": "sec", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr><td>HF Baseline</td><td>llama-8b</td><td>706.0 ms</td><td>34.2 ms</td><td>29.3 tok/s</td><td>—</td><td>21</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>torch.compile</td><td>llama-8b</td><td>307.7 ms</td><td>171.8 ms</td><td>5.8 tok/s</td><td>18760 ms</td><td>21</td><td style="color:#777">sequential per-token, StaticCache KV cache (torch.compile inductor)</td></tr>
+<tr><td>luminal backend</td><td>llama-8b</td><td>461.5 ms</td><td>23.1 ms</td><td>43.3 tok/s</td><td>95963 ms</td><td>21</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>llama-8b</td><td>1026.9 ms</td><td>51.6 ms</td><td>19.4 tok/s</td><td>84453 ms</td><td>21</td><td style="color:#777">sum of per-token prefill durations</td></tr></tbody></table>
+</section>
+
+<section id="llama_8b_sweep">
+  <h2>Llama-8B Sweep</h2>
+  <p class="meta">NousResearch/Meta-Llama-3-8B-Instruct · 21 prompt tokens · 6 results</p>
+  <div class="charts-row"><div id="fig3" class="chart"></div><script>Plotly.newPlot("fig3", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [470.7036415056791, 460.72837291285396, 472.43661794345826], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [751.03, 1038.34, 453.16], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig4" class="chart"></div><script>Plotly.newPlot("fig4", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [23.540849717101082, 23.101884137140587, 23.610779400914907], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [38.2, 51.92, 24.09], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig5" class="chart"></div><script>Plotly.newPlot("fig5", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [28.428826077957638, 43.57440591201885, 95.52432684396626], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} sec<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [15.14307, 30.12727, 84.87889], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "sec", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr><td>luminal backend</td><td>s=10</td><td>470.7 ms</td><td>23.5 ms</td><td>42.5 tok/s</td><td>28429 ms</td><td>21</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=10</td><td>751.0 ms</td><td>38.2 ms</td><td>26.2 tok/s</td><td>15143 ms</td><td>21</td><td style="color:#777">sum of per-token prefill durations</td></tr>
+<tr><td>luminal backend</td><td>s=100</td><td>460.7 ms</td><td>23.1 ms</td><td>43.3 tok/s</td><td>43574 ms</td><td>21</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=100</td><td>1038.3 ms</td><td>51.9 ms</td><td>19.3 tok/s</td><td>30127 ms</td><td>21</td><td style="color:#777">sum of per-token prefill durations</td></tr>
+<tr><td>luminal backend</td><td>s=500</td><td>472.4 ms</td><td>23.6 ms</td><td>42.4 tok/s</td><td>95524 ms</td><td>21</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=500</td><td>453.2 ms</td><td>24.1 ms</td><td>41.5 tok/s</td><td>84879 ms</td><td>21</td><td style="color:#777">sum of per-token prefill durations</td></tr></tbody></table>
+</section>
+
+<section id="qwen3_4b_comparison">
+  <h2>Qwen3-4B Comparison</h2>
+  <p class="meta">Qwen/Qwen3-4B · 19/11 prompt tokens · 4 results</p>
+  <div class="charts-row"><div id="fig6" class="chart"></div><script>Plotly.newPlot("fig6", [{"type": "bar", "name": "HF Baseline", "x": ["qwen3-4b"], "y": [869.2860195587855], "text": ["869.3 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["qwen3-4b"], "y": [298.27259748708457], "text": ["298.3 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["qwen3-4b"], "y": [485.3892414830625], "text": ["485.4 ms"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["qwen3-4b"], "y": [398.58], "text": ["398.6 ms"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig7" class="chart"></div><script>Plotly.newPlot("fig7", [{"type": "bar", "name": "HF Baseline", "x": ["qwen3-4b"], "y": [47.71483448566869], "text": ["47.7 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["qwen3-4b"], "y": [468.56868775503244], "text": ["468.6 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["qwen3-4b"], "y": [26.90318431414198], "text": ["26.9 ms"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["qwen3-4b"], "y": [40.62], "text": ["40.6 ms"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig8" class="chart"></div><script>Plotly.newPlot("fig8", [{"type": "bar", "name": "torch.compile", "x": ["qwen3-4b"], "y": [4.680963660997804], "text": ["4.7 sec"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} sec<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["qwen3-4b"], "y": [45.345814052037895], "text": ["45.3 sec"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} sec<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["qwen3-4b"], "y": [19.92977], "text": ["19.9 sec"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "yaxis": {"title": "sec", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr><td>HF Baseline</td><td>qwen3-4b</td><td>869.3 ms</td><td>47.7 ms</td><td>21.0 tok/s</td><td>—</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>torch.compile</td><td>qwen3-4b</td><td>298.3 ms</td><td>468.6 ms</td><td>2.1 tok/s</td><td>4681 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache (torch.compile inductor)</td></tr>
+<tr><td>luminal backend</td><td>qwen3-4b</td><td>485.4 ms</td><td>26.9 ms</td><td>37.2 tok/s</td><td>45346 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>qwen3-4b</td><td>398.6 ms</td><td>40.6 ms</td><td>24.6 tok/s</td><td>19930 ms</td><td>11</td><td style="color:#777">sum of per-token prefill durations</td></tr></tbody></table>
+</section>
+
+<section id="qwen3_4b_sweep">
+  <h2>Qwen3-4B Sweep</h2>
+  <p class="meta">Qwen/Qwen3-4B · 19/11 prompt tokens · 6 results</p>
+  <div class="charts-row"><div id="fig9" class="chart"></div><script>Plotly.newPlot("fig9", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [465.02652901108377, 465.9317950136028, 495.75577257201076], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [398.44, 390.08, 559.29], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig10" class="chart"></div><script>Plotly.newPlot("fig10", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [25.875402649398893, 25.884080055402592, 27.492373346467502], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [40.64, 39.98, 55.37], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig11" class="chart"></div><script>Plotly.newPlot("fig11", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [37.92102829599753, 54.08867314597592, 118.29659596900456], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} sec<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [12.448030000000001, 27.06796, 81.89342], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "sec", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr><td>luminal backend</td><td>s=10</td><td>465.0 ms</td><td>25.9 ms</td><td>38.6 tok/s</td><td>37921 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=10</td><td>398.4 ms</td><td>40.6 ms</td><td>24.6 tok/s</td><td>12448 ms</td><td>11</td><td style="color:#777">sum of per-token prefill durations</td></tr>
+<tr><td>luminal backend</td><td>s=100</td><td>465.9 ms</td><td>25.9 ms</td><td>38.6 tok/s</td><td>54089 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=100</td><td>390.1 ms</td><td>40.0 ms</td><td>25.0 tok/s</td><td>27068 ms</td><td>11</td><td style="color:#777">sum of per-token prefill durations</td></tr>
+<tr><td>luminal backend</td><td>s=500</td><td>495.8 ms</td><td>27.5 ms</td><td>36.4 tok/s</td><td>118297 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=500</td><td>559.3 ms</td><td>55.4 ms</td><td>18.1 tok/s</td><td>81893 ms</td><td>11</td><td style="color:#777">sum of per-token prefill durations</td></tr></tbody></table>
+</section>
+
+<section id="gemma3_4b_comparison">
+  <h2>Gemma3-4B Comparison</h2>
+  <p class="meta">unsloth/gemma-3-4b-it · 19/11 prompt tokens · 4 results</p>
+  <div class="charts-row"><div id="fig12" class="chart"></div><script>Plotly.newPlot("fig12", [{"type": "bar", "name": "HF Baseline", "x": ["gemma3-4b"], "y": [951.1196144158021], "text": ["951.1 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["gemma3-4b"], "y": [300.9451600664761], "text": ["300.9 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["gemma3-4b"], "y": [404.43], "text": ["404.4 ms"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig13" class="chart"></div><script>Plotly.newPlot("fig13", [{"type": "bar", "name": "HF Baseline", "x": ["gemma3-4b"], "y": [52.498737201676704], "text": ["52.5 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["gemma3-4b"], "y": [2197.426627812092], "text": ["2197.4 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["gemma3-4b"], "y": [38.99], "text": ["39.0 ms"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig14" class="chart"></div><script>Plotly.newPlot("fig14", [{"type": "bar", "name": "torch.compile", "x": ["gemma3-4b"], "y": [26.649526304972824], "text": ["26.6 sec"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} sec<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["gemma3-4b"], "y": [156.84164], "text": ["156.8 sec"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "yaxis": {"title": "sec", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr><td>HF Baseline</td><td>gemma3-4b</td><td>951.1 ms</td><td>52.5 ms</td><td>19.0 tok/s</td><td>—</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>torch.compile</td><td>gemma3-4b</td><td>300.9 ms</td><td>2197.4 ms</td><td>0.5 tok/s</td><td>26650 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache (torch.compile inductor)</td></tr>
+<tr style="background:#fff0f0"><td>luminal backend</td><td>gemma3-4b</td><td>—</td><td>—</td><td>—</td><td>—</td><td>—</td><td style="color:#c00">bench_python_luminal.py failed with code 1</td></tr>
+<tr><td>Rust (luminal)</td><td>gemma3-4b</td><td>404.4 ms</td><td>39.0 ms</td><td>25.6 tok/s</td><td>156842 ms</td><td>11</td><td style="color:#777">sum of per-token prefill durations</td></tr></tbody></table>
+</section>
+
+<section id="gemma3_4b_sweep">
+  <h2>Gemma3-4B Sweep</h2>
+  <p class="meta">unsloth/gemma-3-4b-it · 11 prompt tokens · 6 results</p>
+  <div class="charts-row"><div id="fig15" class="chart"></div><script>Plotly.newPlot("fig15", [{"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [388.19, 436.49, 386.13], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig16" class="chart"></div><script>Plotly.newPlot("fig16", [{"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [37.47, 41.95, 37.25], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig17" class="chart"></div><script>Plotly.newPlot("fig17", [{"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [102.18644, 186.34269, 498.48983000000004], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "sec", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr style="background:#fff0f0"><td>luminal backend</td><td>s=10</td><td>—</td><td>—</td><td>—</td><td>—</td><td>—</td><td style="color:#c00">bench_python_luminal.py failed with code 1</td></tr>
+<tr><td>Rust (luminal)</td><td>s=10</td><td>388.2 ms</td><td>37.5 ms</td><td>26.7 tok/s</td><td>102186 ms</td><td>11</td><td style="color:#777">sum of per-token prefill durations</td></tr>
+<tr style="background:#fff0f0"><td>luminal backend</td><td>s=100</td><td>—</td><td>—</td><td>—</td><td>—</td><td>—</td><td style="color:#c00">bench_python_luminal.py failed with code 1</td></tr>
+<tr><td>Rust (luminal)</td><td>s=100</td><td>436.5 ms</td><td>42.0 ms</td><td>23.8 tok/s</td><td>186343 ms</td><td>11</td><td style="color:#777">sum of per-token prefill durations</td></tr>
+<tr style="background:#fff0f0"><td>luminal backend</td><td>s=500</td><td>—</td><td>—</td><td>—</td><td>—</td><td>—</td><td style="color:#c00">bench_python_luminal.py failed with code 1</td></tr>
+<tr><td>Rust (luminal)</td><td>s=500</td><td>386.1 ms</td><td>37.2 ms</td><td>26.8 tok/s</td><td>498490 ms</td><td>11</td><td style="color:#777">sum of per-token prefill durations</td></tr></tbody></table>
+</section>
+
+<section id="gemma4_moe_comparison">
+  <h2>Gemma4-Moe Comparison</h2>
+  <p class="meta">google/gemma-4-26B-A4B · 11 prompt tokens · 4 results</p>
+  <div class="charts-row"><div id="fig18" class="chart"></div><script>Plotly.newPlot("fig18", [{"type": "bar", "name": "HF Baseline", "x": ["gemma4-moe"], "y": [837.3980740143452], "text": ["837.4 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["gemma4-moe"], "y": [245.510076492792], "text": ["245.5 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig19" class="chart"></div><script>Plotly.newPlot("fig19", [{"type": "bar", "name": "HF Baseline", "x": ["gemma4-moe"], "y": [83.64427039632574], "text": ["83.6 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["gemma4-moe"], "y": [654.9649795080768], "text": ["655.0 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig20" class="chart"></div><script>Plotly.newPlot("fig20", [{"type": "bar", "name": "torch.compile", "x": ["gemma4-moe"], "y": [38.81582092499593], "text": ["38.8 sec"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "yaxis": {"title": "sec", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr><td>HF Baseline</td><td>gemma4-moe</td><td>837.4 ms</td><td>83.6 ms</td><td>12.0 tok/s</td><td>—</td><td>11</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>torch.compile</td><td>gemma4-moe</td><td>245.5 ms</td><td>655.0 ms</td><td>1.5 tok/s</td><td>38816 ms</td><td>11</td><td style="color:#777">sequential per-token, StaticCache KV cache (torch.compile inductor)</td></tr>
+<tr style="background:#fff0f0"><td>luminal backend</td><td>gemma4-moe</td><td>—</td><td>—</td><td>—</td><td>—</td><td>—</td><td style="color:#c00">bench_python_luminal.py failed with code -9</td></tr>
+<tr style="background:#fff0f0"><td>Rust (luminal)</td><td>gemma4-moe</td><td>—</td><td>—</td><td>—</td><td>—</td><td>—</td><td style="color:#c00">rust bench failed with code -9</td></tr></tbody></table>
+</section>
+
+<section id="gemma4_moe_sweep">
+  <h2>Gemma4-Moe Sweep</h2>
+  <p class="meta">google/gemma-4-26B-A4B · 2 results</p>
+  <div class="charts-row"><div id="fig21" class="chart"></div><script>Plotly.newPlot("fig21", [], {"title": "TTFT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10], "ticktext": ["10"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr style="background:#fff0f0"><td>luminal backend</td><td>s=10</td><td>—</td><td>—</td><td>—</td><td>—</td><td>—</td><td style="color:#c00">bench_python_luminal.py failed with code -9</td></tr>
+<tr style="background:#fff0f0"><td>Rust (luminal)</td><td>s=10</td><td>—</td><td>—</td><td>—</td><td>—</td><td>—</td><td style="color:#c00">rust bench failed with code -9</td></tr></tbody></table>
+</section>
+
+<section id="qwen3_moe_comparison">
+  <h2>Qwen3-Moe Comparison</h2>
+  <p class="meta">Qwen/Qwen3-30B-A3B · 19 prompt tokens · 4 results</p>
+  <div class="charts-row"><div id="fig22" class="chart"></div><script>Plotly.newPlot("fig22", [{"type": "bar", "name": "HF Baseline", "x": ["qwen3-moe"], "y": [1565.540504961973], "text": ["1565.5 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["qwen3-moe"], "y": [460.077923577046], "text": ["460.1 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["qwen3-moe"], "y": [21002.791983017232], "text": ["21002.8 ms"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["qwen3-moe"], "y": [662.07], "text": ["662.1 ms"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig23" class="chart"></div><script>Plotly.newPlot("fig23", [{"type": "bar", "name": "HF Baseline", "x": ["qwen3-moe"], "y": [84.527321747737], "text": ["84.5 ms"], "textposition": "outside", "marker": {"color": "#888888"}, "hovertemplate": "%{x}<br>HF Baseline: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "torch.compile", "x": ["qwen3-moe"], "y": [753.0061075551203], "text": ["753.0 ms"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["qwen3-moe"], "y": [1166.8824461026816], "text": ["1166.9 ms"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["qwen3-moe"], "y": [60.08], "text": ["60.1 ms"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "yaxis": {"title": "ms", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig24" class="chart"></div><script>Plotly.newPlot("fig24", [{"type": "bar", "name": "torch.compile", "x": ["qwen3-moe"], "y": [8.341281775035895], "text": ["8.3 sec"], "textposition": "outside", "marker": {"color": "#5ab552"}, "hovertemplate": "%{x}<br>torch.compile: %{y:.1f} sec<extra></extra>"}, {"type": "bar", "name": "luminal backend", "x": ["qwen3-moe"], "y": [111.70731823903043], "text": ["111.7 sec"], "textposition": "outside", "marker": {"color": "#4c9ed9"}, "hovertemplate": "%{x}<br>luminal backend: %{y:.1f} sec<extra></extra>"}, {"type": "bar", "name": "Rust (luminal)", "x": ["qwen3-moe"], "y": [80.83241000000001], "text": ["80.8 sec"], "textposition": "outside", "marker": {"color": "#d97a4c"}, "hovertemplate": "%{x}<br>Rust (luminal): %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "yaxis": {"title": "sec", "rangemode": "tozero"}, "barmode": "group", "legend": {"orientation": "h", "y": -0.2}, "margin": {"t": 50, "b": 80}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr><td>HF Baseline</td><td>qwen3-moe</td><td>1565.5 ms</td><td>84.5 ms</td><td>11.8 tok/s</td><td>—</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>torch.compile</td><td>qwen3-moe</td><td>460.1 ms</td><td>753.0 ms</td><td>1.3 tok/s</td><td>8341 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache (torch.compile inductor)</td></tr>
+<tr><td>luminal backend</td><td>qwen3-moe</td><td>21002.8 ms</td><td>1166.9 ms</td><td>0.9 tok/s</td><td>111707 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>qwen3-moe</td><td>662.1 ms</td><td>60.1 ms</td><td>16.6 tok/s</td><td>80832 ms</td><td>—</td><td style="color:#777">sum of per-token prefill durations</td></tr></tbody></table>
+</section>
+
+<section id="qwen3_moe_sweep">
+  <h2>Qwen3-Moe Sweep</h2>
+  <p class="meta">Qwen/Qwen3-30B-A3B · 19 prompt tokens · 6 results</p>
+  <div class="charts-row"><div id="fig25" class="chart"></div><script>Plotly.newPlot("fig25", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [21002.663500519702, 21018.686580006033, 21034.366824431345], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [656.7, 540.37, 542.34], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TTFT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig26" class="chart"></div><script>Plotly.newPlot("fig26", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [1166.6714247548953, 1167.2746865515364, 1168.7990181031637], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} ms<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [59.6, 48.79, 48.88], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} ms<extra></extra>"}], {"title": "TPOT", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "ms", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script><div id="fig27" class="chart"></div><script>Plotly.newPlot("fig27", [{"type": "scatter", "mode": "lines+markers", "name": "luminal backend", "x": [10, 100, 500], "y": [93.47603664599592, 132.266081985028, 298.05094401398674], "marker": {"size": 8, "color": "#4c9ed9"}, "line": {"color": "#4c9ed9", "width": 2}, "hovertemplate": "iters=%{x}<br>luminal backend: %{y:.1f} sec<extra></extra>"}, {"type": "scatter", "mode": "lines+markers", "name": "Rust (luminal)", "x": [10, 100, 500], "y": [25.48138, 47.5342, 134.79345], "marker": {"size": 8, "color": "#d97a4c"}, "line": {"color": "#d97a4c", "width": 2}, "hovertemplate": "iters=%{x}<br>Rust (luminal): %{y:.1f} sec<extra></extra>"}], {"title": "Time to Search", "xaxis": {"title": "Search iterations", "type": "log", "tickvals": [10, 100, 500], "ticktext": ["10", "100", "500"]}, "yaxis": {"title": "sec", "rangemode": "tozero"}, "legend": {"orientation": "h", "y": -0.25}, "margin": {"t": 50, "b": 90}, "plot_bgcolor": "#fafafa", "paper_bgcolor": "#ffffff"}, {responsive:true});</script></div>
+  <table><thead><tr><th>Path</th><th>Config</th><th>TTFT</th><th>TPOT</th><th>Throughput</th><th>Compile</th><th>Prompt tokens</th><th>Note</th></tr></thead><tbody><tr><td>luminal backend</td><td>s=10</td><td>21002.7 ms</td><td>1166.7 ms</td><td>0.9 tok/s</td><td>93476 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=10</td><td>656.7 ms</td><td>59.6 ms</td><td>16.8 tok/s</td><td>25481 ms</td><td>—</td><td style="color:#777">sum of per-token prefill durations</td></tr>
+<tr><td>luminal backend</td><td>s=100</td><td>21018.7 ms</td><td>1167.3 ms</td><td>0.9 tok/s</td><td>132266 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=100</td><td>540.4 ms</td><td>48.8 ms</td><td>20.5 tok/s</td><td>47534 ms</td><td>—</td><td style="color:#777">sum of per-token prefill durations</td></tr>
+<tr><td>luminal backend</td><td>s=500</td><td>21034.4 ms</td><td>1168.8 ms</td><td>0.9 tok/s</td><td>298051 ms</td><td>19</td><td style="color:#777">sequential per-token, StaticCache KV cache</td></tr>
+<tr><td>Rust (luminal)</td><td>s=500</td><td>542.3 ms</td><td>48.9 ms</td><td>20.5 tok/s</td><td>134793 ms</td><td>—</td><td style="color:#777">sum of per-token prefill durations</td></tr></tbody></table>
+</section>
+</main>
+</body>
+</html>

benchmarks/ttft/run.py

@@ -0,0 +1,683 @@
+"""TTFT + TPOT benchmark orchestrator.
+
+Runs four paths in isolated subprocesses:
+  1. python_baseline       — HuggingFace / PyTorch eager on CUDA
+  2. python_torch_compile  — torch.compile(model) inductor backend
+  3. python_luminal        — torch.compile(model, backend=luminal_backend)
+  4. rust                  — examples/<package> binary (luminal_cuda_lite)
+
+Use --config to select a named configuration, or --all-configs to run every
+entry in CONFIGS. All output is written to the SQLite bench DB
+(benchmarks/ttft/bench.db); the TUI / dashboard / report read from there.
+
+Notes on comparability:
+  - python_baseline: single chunked forward for TTFT; KV-cache decode for TPOT.
+  - python_torch_compile: inductor, same chunked prefill as baseline; first
+    call triggers JIT compilation (recorded separately as compile_ms).
+  - python_luminal: sequential per-token prefill with StaticCache; TPOT via
+    autoregressive decode steps.
+  - rust: sequential per-token prefill; TTFT = sum of prefill step durations.
+Steady-state execution only — compile / egraph-search time excluded from TTFT but
+recorded separately as compile_ms for all paths that support it.
+"""
+
+import argparse
+import datetime
+import json
+import os
+import re
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+try:
+    import tomllib
+except ImportError:
+    try:
+        import tomli as tomllib  # type: ignore[no-redef]
+    except ImportError:
+        raise ImportError("Python 3.11+ or 'pip install tomli' required to load benchmarks.toml")
+
+import db
+
+BENCH_DIR = Path(__file__).resolve().parent
+REPO_ROOT = BENCH_DIR.parent.parent
+
+DEFAULT_PROMPT = "Explain what a neural network is in a paragraph."
+DEFAULT_MODEL = "NousResearch/Meta-Llama-3-8B-Instruct"
+
+_CONFIG_PATH = BENCH_DIR / "benchmarks.toml"
+with open(_CONFIG_PATH, "rb") as _f:
+    _BENCH_CONFIG = tomllib.load(_f)
+
+# Named benchmark configurations.  Each entry overrides any subset of the
+# CLI defaults; explicit CLI flags always take precedence over the config.
+CONFIGS: dict = _BENCH_CONFIG["configs"]
+UR_TEST_MODELS: list = _BENCH_CONFIG["ur_test"]["models"]
+SEARCH_SWEEP_ITERS: list = _BENCH_CONFIG["ur_test"]["search_sweep_iters"]
+
+SWEEP_CONFIG_PREFIX = "s="
+
+BENCH_LINE = re.compile(r"^BENCH_RESULT (.*)$", re.MULTILINE)
+RUST_TTFT_LINE = re.compile(r"TTFT:\s*([0-9]+\.?[0-9]*)\s*ms")
+RUST_TPOT_LINE = re.compile(r"TPOT:\s*([0-9]+\.?[0-9]*)\s*ms")
+RUST_COMPILE_LINE = re.compile(r"COMPILE:\s*([0-9]+\.?[0-9]*)\s*ms")
+RUST_PROMPT_LINE = re.compile(r"Prompt:\s*(\d+)\s*tokens")
+
+
+def _stream(proc, tee_prefix):
+    """Drain subprocess stdout, tee-ing to our stdout line-by-line. Returns full stdout."""
+    buf = []
+    assert proc.stdout is not None
+    for line in proc.stdout:
+        buf.append(line)
+        sys.stdout.write(f"[{tee_prefix}] {line}")
+        sys.stdout.flush()
+    proc.wait()
+    return "".join(buf)
+
+
+_MEM_LOG_PATH = os.environ.get("BENCH_MEM_LOG", "/tmp/bench_mem_snapshots.log")
+
+
+def _snapshot_memory(label: str) -> None:
+    """Append a host+GPU memory snapshot to BENCH_MEM_LOG. Cheap, never raises."""
+    try:
+        ts = datetime.datetime.now().isoformat(timespec="seconds")
+        meminfo_keys = ("MemTotal", "MemFree", "MemAvailable", "Cached", "Slab", "SReclaimable")
+        meminfo = {}
+        with open("/proc/meminfo") as f:
+            for line in f:
+                k, _, rest = line.partition(":")
+                if k in meminfo_keys:
+                    meminfo[k] = rest.strip().split()[0]  # kB
+        try:
+            gpu = subprocess.check_output(
+                ["nvidia-smi", "--query-gpu=memory.used,memory.free,memory.total",
+                 "--format=csv,noheader,nounits"],
+                stderr=subprocess.DEVNULL, text=True, timeout=5,
+            ).strip().splitlines()[0]
+        except Exception:
+            gpu = "n/a"
+        parent_rss = "?"
+        try:
+            with open(f"/proc/{os.getpid()}/status") as f:
+                for line in f:
+                    if line.startswith("VmRSS:"):
+                        parent_rss = line.split()[1]
+                        break
+        except Exception:
+            pass
+        host_str = " ".join(f"{k}={meminfo.get(k, '?')}kB" for k in meminfo_keys)
+        with open(_MEM_LOG_PATH, "a") as f:
+            f.write(f"{ts}  [{label}]  parent_rss={parent_rss}kB  {host_str}  gpu(used,free,total MiB)={gpu}\n")
+    except Exception as e:
+        sys.stderr.write(f"[mem-snapshot warn] {e}\n")
+
+
+def _cargo_env():
+    """Return env dict with ~/.cargo/bin prepended to PATH."""
+    cargo_bin = str(Path.home() / ".cargo" / "bin")
+    path = os.environ.get("PATH", "")
+    if cargo_bin not in path:
+        path = f"{cargo_bin}:{path}"
+    return {**os.environ, "PATH": path}
+
+
+def run_rust(_prompt, package="llama", env_vars=None):
+    print(f"\n=== Running: rust (examples/{package}) ===", flush=True)
+    cmd = ["cargo", "run", "--release", "-p", package]
+    env = _cargo_env()
+    if env_vars:
+        env.update(env_vars)
+    proc = subprocess.Popen(
+        cmd,
+        cwd=REPO_ROOT,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+        env=env,
+    )
+    output = _stream(proc, "rust")
+    if proc.returncode != 0:
+        raise RuntimeError(f"rust bench failed with code {proc.returncode}")
+    m = RUST_TTFT_LINE.search(output)
+    if not m:
+        raise RuntimeError("could not find 'TTFT: X ms' in rust stdout")
+    ttft_ms = float(m.group(1))
+    result = {
+        "path": "rust",
+        "model": DEFAULT_MODEL,
+        "ttft_ms": ttft_ms,
+        "note": "sum of per-token prefill durations",
+    }
+    m_compile = RUST_COMPILE_LINE.search(output)
+    if m_compile:
+        result["compile_ms"] = float(m_compile.group(1))
+    m_tpot = RUST_TPOT_LINE.search(output)
+    if m_tpot:
+        tpot_ms = float(m_tpot.group(1))
+        result["tpot_ms"] = tpot_ms
+        result["throughput_tps"] = 1000.0 / tpot_ms
+    m_prompt = RUST_PROMPT_LINE.search(output)
+    if m_prompt:
+        result["prompt_tokens"] = int(m_prompt.group(1))
+    return result
+
+
+def run_python_script(name, extra_args):
+    script = BENCH_DIR / name
+    print(f"\n=== Running: {script.name} ===", flush=True)
+    cmd = [sys.executable, str(script), *extra_args]
+    proc = subprocess.Popen(
+        cmd,
+        cwd=REPO_ROOT,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+        env={**os.environ},
+    )
+    output = _stream(proc, script.stem)
+    if proc.returncode != 0:
+        raise RuntimeError(f"{script.name} failed with code {proc.returncode}")
+    m = BENCH_LINE.search(output)
+    if not m:
+        raise RuntimeError(f"no BENCH_RESULT line in {script.name} output")
+    return json.loads(m.group(1))
+
+
+PATH_ORDER = ["python_baseline", "python_torch_compile", "python_luminal", "rust"]
+PATH_LABELS = {
+    "python_baseline": "Python\n(HF baseline)",
+    "python_torch_compile": "Python\n(torch.compile)",
+    "python_luminal": "Python → Rust\n(luminal_backend)",
+    "rust": "Rust\n(examples/llama)",
+}
+PATH_COLORS = {
+    "python_baseline": "#888888",
+    "python_torch_compile": "#5ab552",
+    "python_luminal": "#4c9ed9",
+    "rust": "#d97a4c",
+}
+
+
+def run_one_config(config_name, settings, global_skip, inter_path_cooldown=0):
+    """Run all four paths for one config. Returns list of result dicts tagged with 'config'."""
+    model = settings["model"]
+    rust_package = settings["rust_package"]
+    prompt = settings["prompt"]
+    iters = settings["iters"]
+    warmups = settings["warmups"]
+    decode_tokens = settings["decode_tokens"]
+    search_iters = settings["search_iters"]
+    dtype = settings.get("dtype", "float32")
+    skip = set(global_skip) | set(settings.get("skip", []))
+
+    common_py = [
+        "--model", model,
+        "--prompt", prompt,
+        "--iters", str(iters),
+        "--warmups", str(warmups),
+        "--decode-tokens", str(decode_tokens),
+        "--dtype", dtype,
+    ]
+    luminal_py = common_py + ["--search-iters", str(search_iters)]
+
+    rust_env = {"SEARCH_GRAPHS": str(search_iters), "PROMPT": prompt, "ITERS": str(iters)}
+
+    results = []
+    first_path = True
+    for path, fn in [
+        ("python_baseline",      lambda: run_python_script("bench_python_baseline.py", common_py)),
+        ("python_torch_compile", lambda: run_python_script("bench_python_torch_compile.py", common_py)),
+        ("python_luminal",       lambda: run_python_script("bench_python_luminal.py", luminal_py)),
+        ("rust",                 lambda: run_rust(prompt, package=rust_package, env_vars=rust_env)),
+    ]:
+        if path in skip:
+            continue
+        if not first_path and inter_path_cooldown > 0:
+            print(f"  [cooldown {inter_path_cooldown}s]", flush=True)
+            time.sleep(inter_path_cooldown)
+        first_path = False
+        _snapshot_memory(f"{config_name}/{path} BEFORE")
+        try:
+            r = fn()
+            r["config"] = config_name
+            r["model"] = model  # ensure correct model is always tagged
+            if path in ("python_luminal", "rust"):
+                r["search_iters"] = search_iters
+            results.append(r)
+        except Exception as e:
+            print(f"\n[WARN] {config_name}/{path} failed: {e}", flush=True)
+            results.append({
+                "path": path,
+                "config": config_name,
+                "model": model,
+                "error": str(e),
+                "ttft_ms": None,
+            })
+        _snapshot_memory(f"{config_name}/{path} AFTER")
+    return results
+
+
+def plot(results, out_path):
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    # Group by config so each config gets its own subplot column.
+    configs_seen: list[str] = []
+    by_config: dict[str, dict] = {}
+    for r in results:
+        cfg = r.get("config", "default")
+        if cfg not in by_config:
+            configs_seen.append(cfg)
+            by_config[cfg] = {}
+        by_config[cfg][r["path"]] = r
+
+    has_tpot = any(
+        r.get("tpot_ms") is not None
+        for r in results
+        if not r.get("error")
+    )
+    nrows = 2 if has_tpot else 1
+    ncols = len(configs_seen)
+    fig, axes = plt.subplots(nrows, ncols, figsize=(6 * ncols, 4.5 * nrows), squeeze=False)
+
+    for col, cfg in enumerate(configs_seen):
+        by_path = by_config[cfg]
+        present = [p for p in PATH_ORDER if p in by_path]
+
+        def _bar(ax, title, ylabel, key):
+            raw = [by_path[p].get(key) for p in present]
+            ys = [v if v is not None else 0.0 for v in raw]
+            cs = [PATH_COLORS.get(p, "#aaaaaa") if raw[i] is not None else "#cccccc"
+                  for i, p in enumerate(present)]
+            xs = [PATH_LABELS.get(p, p) for p in present]
+            bars = ax.bar(xs, ys, color=cs)
+            ax.set_ylabel(ylabel)
+            ax.set_title(f"{title} — {cfg}")
+            ax.grid(axis="y", alpha=0.3)
+            for b, v in zip(bars, raw):
+                if v is not None:
+                    ax.text(b.get_x() + b.get_width() / 2, v, f"{v:.0f} ms",
+                            ha="center", va="bottom", fontsize=9)
+
+        _bar(axes[0][col], "TTFT", "Time to first token (ms)", "ttft_ms")
+        if has_tpot:
+            _bar(axes[1][col], "TPOT", "Time per output token (ms)", "tpot_ms")
+
+    fig.tight_layout()
+    fig.savefig(out_path, dpi=150)
+    print(f"wrote {out_path}")
+
+
+def run_ur_test(args, conn, run_id):
+    """The ur-test: all 4 paths at default budget + full search sweep, for each model.
+
+    Inserts each result into the DB as it is produced so a mid-run crash still
+    leaves partial data behind.
+    """
+    all_results = []
+
+    for model_idx, model_key in enumerate(UR_TEST_MODELS):
+        s = _settings_for_config(model_key, args)
+
+        if model_idx > 0:
+            print(f"\n  [cooldown 30s between models]", flush=True)
+            time.sleep(30)
+
+        # ── Phase 1: comparison — all 4 paths at the model's default search budget ──
+        print(f"\n{'='*60}\nUR-TEST COMPARISON: {model_key}\n{'='*60}", flush=True)
+        comp_results = run_one_config(model_key, s, args.skip, inter_path_cooldown=20)
+        for r in comp_results:
+            r["model_key"] = model_key
+            db.insert_result(conn, run_id, r)
+        conn.commit()
+        all_results.extend(comp_results)
+
+        # ── Phase 2: search sweep — python_luminal + rust across all budgets ──
+        if args.no_sweep:
+            continue
+        print(f"\n{'='*60}\nUR-TEST SWEEP: {model_key}\n{'='*60}", flush=True)
+        sweep_skip_base = set(args.skip) | {"python_baseline", "python_torch_compile"}
+        # Memory peak in egglog Search grows monotonically with search-iters.
+        # If a path SIGKILLs (-9) at budget N, every higher budget will too —
+        # skip it to avoid wasting another ~hour per model on guaranteed OOMs.
+        oom_paths: set[str] = set()
+        for n in SEARCH_SWEEP_ITERS:
+            print(f"  [cooldown 20s before s={n}]", flush=True)
+            time.sleep(20)
+            sweep_skip = list(sweep_skip_base | oom_paths)
+            if oom_paths:
+                print(f"  [skip-on-prior-OOM] {sorted(oom_paths)} OOM'd at lower budget; skipping at s={n}", flush=True)
+            sweep_s = {**s, "search_iters": n}
+            results_n = run_one_config(f"s={n}", sweep_s, sweep_skip, inter_path_cooldown=20)
+            for r in results_n:
+                r["model_key"] = model_key  # preserve ur-test model identity for dashboard
+                db.insert_result(conn, run_id, r)
+                if "code -9" in (r.get("error") or ""):
+                    oom_paths.add(r["path"])
+            conn.commit()
+            all_results.extend(results_n)
+
+    print("\nGenerate report with:")
+    print(f"  python3 benchmarks/ttft/gen_report.py --db benchmarks/ttft/bench.db --run {run_id} \\")
+    print("    --out benchmarks/ttft/report.html")
+    print("\nGenerate dashboard with:")
+    print("  python3 benchmarks/ttft/gen_dashboard.py --out benchmarks/ttft/dashboard.html")
+
+    return all_results
+
+
+def _git_info():
+    """Return (short_commit, branch) from the repo, or ('unknown', 'unknown') if unavailable."""
+    try:
+        commit = subprocess.check_output(
+            ["git", "rev-parse", "--short", "HEAD"],
+            cwd=REPO_ROOT, stderr=subprocess.DEVNULL, text=True,
+        ).strip()
+        branch = subprocess.check_output(
+            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
+            cwd=REPO_ROOT, stderr=subprocess.DEVNULL, text=True,
+        ).strip()
+        return commit, branch
+    except Exception:
+        return "unknown", "unknown"
+
+
+def _gpu_info() -> dict:
+    """Return GPU metadata from nvidia-smi, or empty dict if unavailable."""
+    try:
+        out = subprocess.check_output(
+            [
+                "nvidia-smi",
+                "--query-gpu=name,driver_version,memory.total",
+                "--format=csv,noheader,nounits",
+            ],
+            stderr=subprocess.DEVNULL,
+            text=True,
+        ).strip()
+        if not out:
+            return {}
+        parts = [p.strip() for p in out.splitlines()[0].split(",")]
+        if len(parts) < 3:
+            return {}
+        return {
+            "gpu_name": parts[0],
+            "gpu_driver": parts[1],
+            "gpu_vram_mb": int(parts[2]),
+        }
+    except Exception:
+        return {}
+
+
+def _cuda_version() -> str:
+    """Return CUDA version string from nvidia-smi, or 'unknown'."""
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query", "--display=COMPUTE"],
+            stderr=subprocess.DEVNULL,
+            text=True,
+        )
+        for line in out.splitlines():
+            if "CUDA Version" in line:
+                return line.split(":")[-1].strip()
+    except Exception:
+        pass
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi"], stderr=subprocess.DEVNULL, text=True
+        )
+        import re as _re
+        m = _re.search(r"CUDA Version:\s*([\d.]+)", out)
+        if m:
+            return m.group(1)
+    except Exception:
+        pass
+    return "unknown"
+
+
+def _record_run(conn, mode):
+    """Insert a `runs` row capturing this orchestrator invocation. Returns run_id.
+
+    Uses microsecond resolution in the run_id so two invocations within the
+    same wallclock second never collide on the runs PRIMARY KEY (insert_run
+    defaults to OR IGNORE, which would otherwise silently merge them and
+    corrupt history). Microseconds also let the dashboard plot back-to-back
+    runs at distinct x-positions instead of stacking them on one date label.
+    """
+    now = datetime.datetime.now()
+    run_id = now.strftime("%Y-%m-%dT%H-%M-%S-%f")
+    commit, branch = _git_info()
+    db.insert_run(
+        conn,
+        run_id=run_id,
+        timestamp=now.isoformat(),
+        mode=mode,
+        git_commit=commit,
+        git_branch=branch,
+        cuda_version=_cuda_version(),
+        **_gpu_info(),
+    )
+    conn.commit()
+    return run_id
+
+
+def _settings_from_args(args):
+    """Build a settings dict from parsed CLI args."""
+    return {
+        "model": args.model,
+        "rust_package": args.rust_package,
+        "prompt": args.prompt,
+        "iters": args.iters,
+        "warmups": args.warmups,
+        "decode_tokens": args.decode_tokens,
+        "search_iters": args.search_iters,
+        "dtype": args.dtype,
+        "skip": [],
+    }
+
+
+def _settings_for_config(config_name, args):
+    """Merge CONFIGS[config_name] over CLI arg defaults."""
+    cfg = CONFIGS[config_name]
+    return {
+        "model":        cfg.get("model",        args.model),
+        "rust_package": cfg.get("rust_package", args.rust_package),
+        "prompt":       cfg.get("prompt",       args.prompt),
+        "iters":        cfg.get("iters",        args.iters),
+        "warmups":      cfg.get("warmups",      args.warmups),
+        "decode_tokens":cfg.get("decode_tokens",args.decode_tokens),
+        "search_iters": cfg.get("search_iters", args.search_iters),
+        "dtype":        cfg.get("dtype",        args.dtype),
+        "skip":         cfg.get("skip",         []),
+    }
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--config",
+        choices=list(CONFIGS),
+        default=None,
+        help="Named benchmark configuration. Sets parameter defaults; explicit flags override.",
+    )
+    ap.add_argument(
+        "--all-configs",
+        action="store_true",
+        dest="all_configs",
+        help="Run every entry in CONFIGS into a single run_id in the DB.",
+    )
+    ap.add_argument(
+        "--search-sweep",
+        action="store_true",
+        dest="search_sweep",
+        help=(
+            "Run python_luminal + rust across all SEARCH_SWEEP_ITERS budgets "
+            f"({SEARCH_SWEEP_ITERS}). Uses --config (default: llama-8b) as the base settings."
+        ),
+    )
+    ap.add_argument(
+        "--skip-configs",
+        nargs="*",
+        default=[],
+        choices=list(CONFIGS),
+        dest="skip_configs",
+        metavar="CONFIG",
+        help="Config names to exclude when using --all-configs.",
+    )
+    ap.add_argument(
+        "--no-sweep",
+        action="store_true",
+        dest="no_sweep",
+        help=(
+            "With --ur-test: skip the search-budget sweep phase and only run "
+            "the 4-path comparison for each model. ~1.5 hr instead of ~5 hr."
+        ),
+    )
+    ap.add_argument("--model", default=DEFAULT_MODEL)
+    ap.add_argument("--rust-package", default="llama", dest="rust_package",
+                    help="Cargo package name for the rust bench (examples/<name>).")
+    ap.add_argument("--prompt", default=DEFAULT_PROMPT)
+    ap.add_argument("--iters", type=int, default=3)
+    ap.add_argument("--warmups", type=int, default=1)
+    ap.add_argument("--skip", nargs="*", default=[],
+                    choices=["rust", "python_luminal", "python_baseline", "python_torch_compile"])
+    ap.add_argument("--out", default=str(BENCH_DIR / "ttft.png"))
+    ap.add_argument("--db", default=str(db.DEFAULT_DB_PATH),
+                    help="SQLite database file (default: benchmarks/ttft/bench.db).")
+    ap.add_argument("--run", default=None, dest="run",
+                    help="With --render-only: run_id to render (default: latest).")
+    ap.add_argument(
+        "--decode-tokens", type=int, default=50,
+        help="Tokens to generate for TPOT measurement (0 = skip TPOT).",
+    )
+    ap.add_argument(
+        "--search-iters", type=int, default=500,
+        help="Egraph search iterations for the python_luminal path.",
+    )
+    ap.add_argument(
+        "--dtype", default="float32",
+        choices=["float32", "bfloat16", "float16"],
+        help="Torch dtype for the python paths. Configs may override per-model.",
+    )
+    ap.add_argument(
+        "--render-only", action="store_true",
+        help="Skip running benches; render an existing run from the DB. "
+             "Use --run RUN_ID to pick a specific run, otherwise the latest is used.",
+    )
+    ap.add_argument(
+        "--ur-test", action="store_true", dest="ur_test",
+        help=(
+            f"The mega-test: run all 4 paths at default budget + full search sweep "
+            f"({SEARCH_SWEEP_ITERS}) for each of {UR_TEST_MODELS}."
+        ),
+    )
+
+    # Pre-parse to apply named config as argparse defaults so explicit CLI
+    # flags still override them.
+    pre, _ = ap.parse_known_args()
+    if pre.config and not (pre.all_configs or getattr(pre, "search_sweep", False)):
+        cfg = CONFIGS[pre.config]
+        ap.set_defaults(**{k: v for k, v in cfg.items() if k not in ("skip",)})
+    args = ap.parse_args()
+    if pre.config and not args.all_configs and not args.search_sweep:
+        for path in CONFIGS[pre.config].get("skip", []):
+            if path not in args.skip:
+                args.skip.append(path)
+
+    conn = db.connect(args.db)
+
+    if args.render_only:
+        run_id = args.run or db.latest_run_id(conn)
+        if run_id is None:
+            sys.exit(f"--render-only: no runs found in {args.db}")
+        results = db.load_results(conn, run_id)
+        if not results:
+            sys.exit(f"--render-only: no results found for run {run_id} in {args.db}")
+        print(f"rendering run {run_id} ({len(results)} results)")
+    else:
+        mode = (
+            ("ur-test-fast" if args.no_sweep else "ur-test") if args.ur_test
+            else "search-sweep" if args.search_sweep
+            else "all-configs"  if args.all_configs
+            else "single"
+        )
+        run_id = _record_run(conn, mode)
+        print(f"run_id: {run_id}  →  {args.db}")
+
+        if args.ur_test:
+            results = run_ur_test(args, conn, run_id)
+        elif args.search_sweep:
+            results = []
+            # Base settings come from --config (default: llama-8b) or bare CLI args.
+            base = (
+                _settings_for_config(args.config, args)
+                if args.config
+                else _settings_for_config("llama-8b", args)
+            )
+            sweep_skip = set(args.skip) | {"python_baseline", "python_torch_compile"}
+            for i, n in enumerate(SEARCH_SWEEP_ITERS):
+                if i > 0:
+                    print(f"  [cooldown 20s — letting CUDA free previous model memory]", flush=True)
+                    time.sleep(20)
+                print(f"\n{'='*60}\nSEARCH SWEEP: s={n}\n{'='*60}", flush=True)
+                s = {**base, "search_iters": n}
+                rs = run_one_config(f"s={n}", s, list(sweep_skip))
+                for r in rs:
+                    db.insert_result(conn, run_id, r)
+                conn.commit()
+                results.extend(rs)
+        elif args.all_configs:
+            results = []
+            for config_name in CONFIGS:
+                if config_name in args.skip_configs:
+                    continue
+                print(f"\n{'='*60}\nCONFIG: {config_name}\n{'='*60}", flush=True)
+                settings = _settings_for_config(config_name, args)
+                rs = run_one_config(config_name, settings, args.skip)
+                for r in rs:
+                    db.insert_result(conn, run_id, r)
+                conn.commit()
+                results.extend(rs)
+        else:
+            config_name = args.config or "default"
+            settings = (
+                _settings_for_config(args.config, args)
+                if args.config
+                else _settings_from_args(args)
+            )
+            results = run_one_config(config_name, settings, args.skip)
+            for r in results:
+                db.insert_result(conn, run_id, r)
+            conn.commit()
+
+    # Summary
+    configs_in_results = list(dict.fromkeys(r.get("config", "default") for r in results))
+    for cfg in configs_in_results:
+        group = [r for r in results if r.get("config", "default") == cfg]
+        print(f"\nSummary ({cfg}):")
+        for r in group:
+            if r.get("error"):
+                print(f"  {r['path']:>22}:  FAILED — {r['error']}")
+                continue
+            if r.get("ttft_ms") is None:
+                print(f"  {r['path']:>22}:  no data")
+                continue
+            compile_ms = r.get("compile_ms")
+            compile_str = f"  compile {compile_ms:.0f} ms" if compile_ms is not None else ""
+            tpot = r.get("tpot_ms")
+            tput = r.get("throughput_tps")
+            tpot_str = f"  TPOT {tpot:.2f} ms  ({tput:.1f} tok/s)" if tpot is not None else ""
+            print(f"  {r['path']:>22}:  TTFT {r['ttft_ms']:.2f} ms{compile_str}{tpot_str}")
+
+    plot(results, args.out)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/ttft/run.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+# TTFT benchmark entrypoint. Runs via uv against the luminal_python venv.
+set -e
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+REPO_ROOT="$( cd "$SCRIPT_DIR/../.." && pwd )"
+cd "$REPO_ROOT/crates/luminal_python"
+exec uv run python "$SCRIPT_DIR/run.py" "$@"

ci/modal_cargo_test.py

@@ -0,0 +1,68 @@
+import modal
+import subprocess
+import os
+
+gpu_type = os.environ.get("GPU_TYPE", "T4")
+CUDARC_CUDA_VERSION = "12080"
+
+app = modal.App("luminal-ci-cargo-test")
+
+WORKDIR = "/workspace/luminal"
+
+cuda_image = (
+    modal.Image.from_registry("nvcr.io/nvidia/pytorch:25.03-py3")
+    .apt_install("protobuf-compiler")
+    .run_commands(
+        "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y",
+    )
+    .env(
+        {
+            "PATH": "/root/.cargo/bin:$PATH",
+            "CUDARC_CUDA_VERSION": CUDARC_CUDA_VERSION,
+        }
+    )
+    .add_local_dir(".", remote_path=WORKDIR, copy=True)
+)
+
+
+@app.function(
+    image=cuda_image,
+    gpu=gpu_type,
+    timeout=1800,  # 30 minutes
+)
+def run_cargo_test():
+    """Run cargo test for luminal_cuda_lite on a Modal GPU."""
+    subprocess.run(["nvidia-smi"], check=True)
+
+    # Detect GPU compute capability
+    result = subprocess.run(
+        ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    compute_cap = result.stdout.strip().replace(".", "")
+
+    subprocess.run(
+        [
+            "cargo",
+            "test",
+            "-p",
+            "luminal_cuda_lite",
+            "--verbose",
+            "--",
+            "--test-threads=1",
+        ],
+        cwd=WORKDIR,
+        env={
+            **os.environ,
+            "CUDARC_CUDA_VERSION": CUDARC_CUDA_VERSION,
+            "CUDA_COMPUTE_CAP": compute_cap,
+        },
+        check=True,
+    )
+
+
+@app.local_entrypoint()
+def main():
+    run_cargo_test.remote()

crates/luminal_bench/examples/debug_ops.rs

@@ -106,13 +106,13 @@ impl Case {
         let out = match self {
             Case::Mul => {
                 let x = cx.tensor(size);
-                x.clone() * x
+                x * x
             }
             Case::Sigmoid => cx.tensor(size).sigmoid(),
             Case::Tanh => cx.tensor(size).tanh(),
             Case::GeluInner => {
                 let x = cx.tensor(size);
-                (0.797_884_560_8_f32 * x.clone() * (1. + 0.044_715_f32 * x.clone() * x)).tanh()
+                (0.797_884_6_f32 * x * (1. + 0.044_715_f32 * x * x)).tanh()
             }
             Case::Gelu => cx.tensor(size).gelu(),
             Case::LayerNorm => {
@@ -447,10 +447,10 @@ where
         if let Some(ref backend) = backend_analysis {
             print_lowering_analysis(backend);
         }
-    } else if !args.inspect_ops.is_empty() {
-        if let Some(ref backend) = backend_analysis {
-            print_lowering_analysis(backend);
-        }
+    } else if !args.inspect_ops.is_empty()
+        && let Some(ref backend) = backend_analysis
+    {
+        print_lowering_analysis(backend);
     }
 
     // Trace facts for explicit variables.

crates/luminal_cuda_lite/src/dyn_backend.rs

@@ -0,0 +1,75 @@
+//! [`DynBackend`] implementation for the CUDA lite runtime.
+
+use luminal::dtype::DType;
+use luminal::dyn_backend::{BackendCompileArgs, DynBackend, compile_backend};
+use luminal::prelude::*;
+
+use crate::cudarc::driver::CudaContext;
+use crate::runtime::CudaRuntime;
+
+/// [`DynBackend`] wrapper for [`CudaRuntime`].
+pub struct CudaLiteDynBackend {
+    pub runtime: CudaRuntime,
+}
+
+impl DynBackend for CudaLiteDynBackend {
+    fn name(&self) -> &str {
+        "cuda_lite"
+    }
+    fn device_type(&self) -> &str {
+        "cuda"
+    }
+
+    fn set_data_bytes(&mut self, node: NodeIndex, bytes: Vec<u8>, _dtype: DType) {
+        self.runtime.set_data(node, bytes);
+    }
+    fn set_data_f32(&mut self, node: NodeIndex, data: Vec<f32>) {
+        self.runtime.set_data(node, data);
+    }
+    fn get_output_f32(&self, node: NodeIndex) -> Vec<f32> {
+        self.runtime.get_f32(node)
+    }
+    fn get_output_i32(&self, node: NodeIndex) -> Vec<i32> {
+        self.runtime.get_i32(node)
+    }
+    fn get_output_bool(&self, node: NodeIndex) -> Vec<bool> {
+        self.runtime.get_bool(node)
+    }
+    fn execute(&mut self, dyn_map: &FxHashMap<char, usize>) {
+        self.runtime.execute(dyn_map);
+    }
+
+    fn supports_device_ptrs(&self) -> bool {
+        true
+    }
+    unsafe fn set_device_ptr(&mut self, node: NodeIndex, ptr: u64, n: usize) {
+        unsafe { self.runtime.set_device_ptr(node, ptr, n) }
+    }
+    unsafe fn set_output_device_ptr(&mut self, node: NodeIndex, ptr: u64, n: usize) {
+        unsafe { self.runtime.set_output_device_ptr(node, ptr, n) }
+    }
+    fn output_is_zero_copy(&self, node: NodeIndex) -> bool {
+        self.runtime.output_is_zero_copy(node)
+    }
+    unsafe fn copy_output_to_device_ptr(&self, node: NodeIndex, ptr: u64, n: usize) {
+        unsafe { self.runtime.copy_output_to_device_ptr(node, ptr, n) }
+    }
+}
+
+pub fn cuda_lite_factory(
+    graph: &mut Graph,
+    args: BackendCompileArgs,
+) -> Result<Box<dyn DynBackend>, String> {
+    let cuda_ctx = CudaContext::new(0).map_err(|e| format!("CUDA init failed: {e}"))?;
+    let stream = cuda_ctx.default_stream();
+    compile_backend::<CudaRuntime>(
+        graph,
+        args,
+        || Ok(CudaRuntime::initialize(stream)),
+        |rt, node, bytes, _dtype| {
+            rt.set_data(node, bytes);
+        },
+        Some(&|rt, node, ptr, n| unsafe { rt.set_device_ptr(node, ptr, n) }),
+        |rt| Box::new(CudaLiteDynBackend { runtime: rt }),
+    )
+}

crates/luminal_cuda_lite/src/host/cublaslt/mod.rs

@@ -32,6 +32,7 @@ use crate::{
         driver::{CudaSlice, CudaStream, DevicePtr},
     },
     host::{HostOp, cublas::parse_cublas_op},
+    try_create_cublaslt,
 };
 
 #[derive(Debug)]
@@ -248,6 +249,19 @@ fn dtype_to_cuda_types(dtype: DType) -> (cudaDataType, cublasComputeType_t, cuda
     }
 }
 
+impl CuBlasLt {
+    fn get_cublaslt(&self, stream: &Arc<CudaStream>) -> anyhow::Result<Arc<CudaBlasLT>> {
+        if let Some(cublaslt) = self.cublaslt.get() {
+            return Ok(cublaslt.clone());
+        }
+        let created = try_create_cublaslt(stream.clone()).map_err(|message| {
+            anyhow::anyhow!("cuBLASLt unavailable on this machine: {message}")
+        })?;
+        let _ = self.cublaslt.set(created.clone());
+        Ok(created)
+    }
+}
+
 impl HostOp for CuBlasLt {
     fn execute(
         &self,
@@ -324,9 +338,7 @@ impl HostOp for CuBlasLt {
         )
         .entered();
 
-        let cublaslt = self
-            .cublaslt
-            .get_or_init(|| Arc::new(CudaBlasLT::new(stream.clone()).unwrap()));
+        let cublaslt = self.get_cublaslt(stream)?;
 
         let mut matmul_desc: cublasLtMatmulDesc_t = std::ptr::null_mut();
         let mut a_desc: cublasLtMatrixLayout_t = std::ptr::null_mut();
@@ -461,7 +473,8 @@ impl HostOp for CuBlasLt {
             cublasLtMatmulDescDestroy(matmul_desc);
         }
 
-        stream.synchronize()?;
+        // No stream.synchronize() here — CUDA stream ordering guarantees
+        // sequential execution. The runtime syncs once at the end of execute().
         Ok(())
     }
 

crates/luminal_cuda_lite/src/host/moe/glumoe_rewrite.egg

@@ -1,128 +1,213 @@
-; GLUMoE: Match the expert computation subgraph of a Gated MoE (SwiGLU variant).
+; GLUMoE: Match the expert computation subgraph of a gated MoE.
 ;
-; This matches the pattern produced by QwenMoE::forward() starting from the
-; expert gathers through to the final weighted sum, and replaces it with a
-; fused GLUMoE HostOp.
+; One fused op supports two activation modes:
+;   mode=0: Qwen-style SwiGLU (silu(gate) * up)
+;   mode=1: Gemma-style GELU (gate * sigmoid(1.595769 * gate * (1 + 0.044715 * gate^2)))
 ;
-; Inputs extracted:
-;   ?x            - input activations [s, H] F32
-;   ?topk_idx     - top-k expert indices [s, k] Int (from argsort+slice)
-;   ?topk_vals    - top-k routing values [s, k] F32 (from gather on softmax)
-;   ?gate_up_w    - stacked gate+up expert weights [E, intermediate*2, H] BF16
-;   ?down_w       - stacked down expert weights [E, H, intermediate] BF16
-;
-; The pattern captures:
-;   1. Gate-up expert gather (Iota, Mul, Cast, Iota, Cast, Add, Cast, Gather)
-;   2. Cast BF16→F32 of gathered gate-up weights
-;   3. Gate-up batched matmul (Mul + SumReduce)
-;   4. Gate/Up split via Iota+Gather (slice semantics)
-;   5. SwiGLU: silu(gate) * up
-;   6. Down expert gather (same pattern as gate-up)
-;   7. Cast BF16→F32 of gathered down weights
-;   8. Down batched matmul (Mul + SumReduce)
-;   9. Weighted sum: (down_out * topk_values) summed over k
-;
-; Variables with ? prefix are egglog pattern variables.
-; We use wildcards (?_xxx) for shapes/strides we don't extract.
+; To keep matching fast, we stage through marker states:
+;   1) Shared gate-up matmul marker
+;   2) Activation marker (separate swiglu / gemma_gelu paths)
+;   3) Down matmul marker (separate swiglu / gemma_gelu paths)
+;   4) Final GLUMoE fusion (separate swiglu / gemma_gelu rules)
+
+(datatype*
+    (GLUMoEGateUpState
+        (MkGLUMoEGateUpState Expression Expression Expression IR IR IR)
+    )
+    (GLUMoESwiGLUState
+        (MkGLUMoESwiGLUState GLUMoEGateUpState)
+    )
+    (GLUMoEGemmaGELUState
+        (MkGLUMoEGemmaGELUState GLUMoEGateUpState)
+    )
+    (GLUMoESwiGLUDownState
+        (MkGLUMoESwiGLUDownState Expression Expression Expression GLUMoESwiGLUState IR IR)
+    )
+    (GLUMoEGemmaDownState
+        (MkGLUMoEGemmaDownState Expression Expression Expression GLUMoEGemmaGELUState IR IR)
+    )
+)
+
+(function glumoe_gate_up (IR) GLUMoEGateUpState :merge new)
+(function glumoe_swiglu (IR) GLUMoESwiGLUState :merge new)
+(function glumoe_gemma_gelu (IR) GLUMoEGemmaGELUState :merge new)
+(function glumoe_swiglu_down (IR) GLUMoESwiGLUDownState :merge new)
+(function glumoe_gemma_down (IR) GLUMoEGemmaDownState :merge new)
 
 (rule
     (
         ; ===== Gate-up expert gather =====
-        ; t51: Iota for base index (expert_idx * io_gu)
         (= ?gu_iota_base (Op (Iota ?gu_io ?gu_iota_base_range) (INil)))
-        ; t52: Mul topk_indices * io → base offsets [s, k]
         (= ?gu_mul_base (Op (Mul ?gu_mul_base_shape ?gu_mul_base_a_stride ?gu_mul_base_b_stride ?gu_mul_base_out_stride) (ICons ?topk_idx (ICons ?gu_iota_base (INil)))))
-        ; t53: Cast to F32
-        (= ?gu_cast_base (Op (Cast ?gu_cast_base_size (F32)) (ICons ?gu_mul_base (INil))))
-        ; t54: Iota for within-expert index
         (= ?gu_iota_within (Op (Iota (MIter) ?gu_iota_within_range) (INil)))
-        ; t55: Cast within to F32
-        (= ?gu_cast_within (Op (Cast ?gu_cast_within_size (F32)) (ICons ?gu_iota_within (INil))))
-        ; t56: Add base + within → flat gather indices
-        (= ?gu_add_idx (Op (Add ?gu_add_shape ?gu_add_a_stride ?gu_add_b_stride ?gu_add_out_stride) (ICons ?gu_cast_base (ICons ?gu_cast_within (INil)))))
-        ; t57: Cast to Int
-        (= ?gu_cast_idx (Op (Cast ?gu_cast_idx_size (Int)) (ICons ?gu_add_idx (INil))))
-        ; t58: Gather gate_up weights
-        (= ?gu_gathered (Op (Gather ?gu_gather_idx_shape ?gu_gather_idx_stride ?gu_gather_data_shape ?gu_gather_data_stride) (ICons ?gu_cast_idx (ICons ?gate_up_w (INil)))))
+        (= ?gu_add_idx (Op (Add ?gu_add_shape ?gu_add_a_stride ?gu_add_b_stride ?gu_add_out_stride) (ICons ?gu_mul_base (ICons ?gu_iota_within (INil)))))
+        (= ?gu_gathered (Op (Gather ?gu_gather_idx_shape ?gu_gather_idx_stride ?gu_gather_data_shape ?gu_gather_data_stride) (ICons ?gu_add_idx (ICons ?gate_up_w (INil)))))
 
         ; ===== Cast BF16→F32 =====
-        ; t59: Cast gathered gate_up to F32
         (= ?gu_f32 (Op (Cast ?gu_f32_size (F32)) (ICons ?gu_gathered (INil))))
 
         ; ===== Gate-up batched matmul =====
-        ; t60: Mul x * gathered_gu (broadcast multiply)
         (= ?gu_matmul_mul (Op (Mul ?gu_matmul_mul_shape ?gu_matmul_a_stride ?gu_matmul_b_stride ?gu_matmul_mul_out_stride) (ICons ?x (ICons ?gu_f32 (INil)))))
-        ; t61: SumReduce over K dimension
         (= ?gu_matmul (Op (Sum ?gu_matmul_out_shape ?gu_matmul_k ?gu_matmul_in_stride ?gu_matmul_k_stride ?gu_matmul_out_stride) (ICons ?gu_matmul_mul (INil))))
+    )
+    (
+        (set (glumoe_gate_up ?gu_matmul)
+            (MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_iota_within_range ?x ?topk_idx ?gate_up_w))
+    )
+    :name "GLUMoE gate-up matmul marker"
+)
+
+; ===== SwiGLU activation marker =====
+(rule
+    (
+        (= ?gate_up_state (glumoe_gate_up ?gu_matmul))
+        (= ?gate_up_state (MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_within_range ?x ?topk_idx ?gate_up_w))
 
-        ; ===== Up slice via Iota+Gather =====
-        ; t62: Iota with complex expression (slicing the "up" half)
         (= ?up_iota (Op (Iota ?up_iota_expr ?up_iota_range) (INil)))
-        ; t63: Gather to select up portion from matmul result
         (= ?up_slice (Op (Gather ?up_gather_idx_shape ?up_gather_idx_stride ?up_gather_data_shape ?up_gather_data_stride) (ICons ?up_iota (ICons ?gu_matmul (INil)))))
 
-        ; ===== SwiGLU: silu(gate) * up =====
-        ; t64: Constant(-1)
         (= ?neg1 (Op (Constant -1.000000) (INil)))
-        ; t65: gate * -1
         (= ?neg_gate (Op (Mul ?silu_shape1 ?silu_a_stride1 ?silu_b_stride1 ?silu_out_stride1) (ICons ?gu_matmul (ICons ?neg1 (INil)))))
-        ; t66: Constant(log2e)
         (= ?log2e (Op (Constant 1.442695) (INil)))
-        ; t67: neg_gate * log2e
         (= ?scaled (Op (Mul ?silu_shape2 ?silu_a_stride2 ?silu_b_stride2 ?silu_out_stride2) (ICons ?neg_gate (ICons ?log2e (INil)))))
-        ; t68: exp2
         (= ?exp2_val (Op (Exp2 ?silu_shape3 ?silu_in_stride3 ?silu_out_stride3) (ICons ?scaled (INil))))
-        ; t69: Constant(1)
         (= ?one (Op (Constant 1.000000) (INil)))
-        ; t70: exp2 + 1
         (= ?plus1 (Op (Add ?silu_shape4 ?silu_a_stride4 ?silu_b_stride4 ?silu_out_stride4) (ICons ?exp2_val (ICons ?one (INil)))))
-        ; t71: recip
         (= ?sigmoid (Op (Recip ?silu_shape5 ?silu_in_stride5 ?silu_out_stride5) (ICons ?plus1 (INil))))
-        ; t72: gate * sigmoid(gate) = silu(gate)
         (= ?silu_out (Op (Mul ?silu_shape6 ?silu_a_stride6 ?silu_b_stride6 ?silu_out_stride6) (ICons ?gu_matmul (ICons ?sigmoid (INil)))))
-        ; t73: silu(gate) * up
         (= ?swiglu_out (Op (Mul ?swiglu_shape ?swiglu_a_stride ?swiglu_b_stride ?swiglu_out_stride) (ICons ?silu_out (ICons ?up_slice (INil)))))
+    )
+    (
+        (set (glumoe_swiglu ?swiglu_out) (MkGLUMoESwiGLUState ?gate_up_state))
+    )
+    :name "GLUMoE swiglu marker"
+)
+
+; ===== Gemma GELU activation marker =====
+(rule
+    (
+        (= ?gate_up_state (glumoe_gate_up ?gu_matmul))
+        (= ?gate_up_state (MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_within_range ?x ?topk_idx ?gate_up_w))
+
+        (= ?up_iota (Op (Iota ?up_iota_expr ?up_iota_range) (INil)))
+        (= ?up_slice (Op (Gather ?up_gather_idx_shape ?up_gather_idx_stride ?up_gather_data_shape ?up_gather_data_stride) (ICons ?up_iota (ICons ?gu_matmul (INil)))))
+
+        (= ?gelu_coeff_inner (Op (Constant 0.044715) (INil)))
+        (= ?gelu_inner_scaled (Op (Mul ?gelu_inner_scaled_shape ?gelu_inner_scaled_a_stride ?gelu_inner_scaled_b_stride ?gelu_inner_scaled_out_stride) (ICons ?gu_matmul (ICons ?gelu_coeff_inner (INil)))))
+        (= ?gelu_inner_quad (Op (Mul ?gelu_inner_quad_shape ?gelu_inner_quad_a_stride ?gelu_inner_quad_b_stride ?gelu_inner_quad_out_stride) (ICons ?gelu_inner_scaled (ICons ?gu_matmul (INil)))))
+        (= ?gelu_one (Op (Constant 1.000000) (INil)))
+        (= ?gelu_poly (Op (Add ?gelu_poly_shape ?gelu_poly_a_stride ?gelu_poly_b_stride ?gelu_poly_out_stride) (ICons ?gelu_inner_quad (ICons ?gelu_one (INil)))))
+        (= ?gelu_coeff_outer (Op (Constant 1.595769) (INil)))
+        (= ?gelu_outer_scaled (Op (Mul ?gelu_outer_scaled_shape ?gelu_outer_scaled_a_stride ?gelu_outer_scaled_b_stride ?gelu_outer_scaled_out_stride) (ICons ?gu_matmul (ICons ?gelu_coeff_outer (INil)))))
+        (= ?gelu_scaled (Op (Mul ?gelu_scaled_shape ?gelu_scaled_a_stride ?gelu_scaled_b_stride ?gelu_scaled_out_stride) (ICons ?gelu_outer_scaled (ICons ?gelu_poly (INil)))))
+        (= ?neg1 (Op (Constant -1.000000) (INil)))
+        (= ?gelu_neg (Op (Mul ?gelu_neg_shape ?gelu_neg_a_stride ?gelu_neg_b_stride ?gelu_neg_out_stride) (ICons ?gelu_scaled (ICons ?neg1 (INil)))))
+        (= ?log2e (Op (Constant 1.442695) (INil)))
+        (= ?gelu_exp_scaled (Op (Mul ?gelu_exp_scaled_shape ?gelu_exp_scaled_a_stride ?gelu_exp_scaled_b_stride ?gelu_exp_scaled_out_stride) (ICons ?gelu_neg (ICons ?log2e (INil)))))
+        (= ?gelu_exp2_val (Op (Exp2 ?gelu_exp_shape ?gelu_exp_in_stride ?gelu_exp_out_stride) (ICons ?gelu_exp_scaled (INil))))
+        (= ?gelu_plus1 (Op (Add ?gelu_plus1_shape ?gelu_plus1_a_stride ?gelu_plus1_b_stride ?gelu_plus1_out_stride) (ICons ?gelu_exp2_val (ICons ?gelu_one (INil)))))
+        (= ?gelu_sigmoid (Op (Recip ?gelu_sigmoid_shape ?gelu_sigmoid_in_stride ?gelu_sigmoid_out_stride) (ICons ?gelu_plus1 (INil))))
+        (= ?gelu_out (Op (Mul ?gelu_out_shape ?gelu_out_a_stride ?gelu_out_b_stride ?gelu_out_out_stride) (ICons ?gu_matmul (ICons ?gelu_sigmoid (INil)))))
+        (= ?gemma_out (Op (Mul ?geglu_shape ?geglu_a_stride ?geglu_b_stride ?geglu_out_stride) (ICons ?gelu_out (ICons ?up_slice (INil)))))
+    )
+    (
+        (set (glumoe_gemma_gelu ?gemma_out) (MkGLUMoEGemmaGELUState ?gate_up_state))
+    )
+    :name "GLUMoE gemma gelu marker"
+)
+
+; ===== SwiGLU down marker =====
+(rule
+    (
+        (= ?swiglu_state (glumoe_swiglu ?swiglu_out))
+        (= ?swiglu_state (MkGLUMoESwiGLUState ?gate_up_state))
 
-        ; ===== Down expert gather =====
-        ; t74: Iota for base index (expert_idx * io_down)
         (= ?dn_iota_base (Op (Iota ?dn_io ?dn_iota_base_range) (INil)))
-        ; t75: Mul topk_indices * io_down
         (= ?dn_mul_base (Op (Mul ?dn_mul_base_shape ?dn_mul_base_a_stride ?dn_mul_base_b_stride ?dn_mul_base_out_stride) (ICons ?topk_idx (ICons ?dn_iota_base (INil)))))
-        ; t76: Cast to F32
-        (= ?dn_cast_base (Op (Cast ?dn_cast_base_size (F32)) (ICons ?dn_mul_base (INil))))
-        ; t77: Iota for within-expert index
         (= ?dn_iota_within (Op (Iota (MIter) ?dn_iota_within_range) (INil)))
-        ; t78: Cast within to F32
-        (= ?dn_cast_within (Op (Cast ?dn_cast_within_size (F32)) (ICons ?dn_iota_within (INil))))
-        ; t79: Add base + within
-        (= ?dn_add_idx (Op (Add ?dn_add_shape ?dn_add_a_stride ?dn_add_b_stride ?dn_add_out_stride) (ICons ?dn_cast_base (ICons ?dn_cast_within (INil)))))
-        ; t80: Cast to Int
-        (= ?dn_cast_idx (Op (Cast ?dn_cast_idx_size (Int)) (ICons ?dn_add_idx (INil))))
-        ; t81: Gather down weights
-        (= ?dn_gathered (Op (Gather ?dn_gather_idx_shape ?dn_gather_idx_stride ?dn_gather_data_shape ?dn_gather_data_stride) (ICons ?dn_cast_idx (ICons ?down_w (INil)))))
-
-        ; ===== Cast BF16→F32 =====
-        ; t82: Cast gathered down to F32
+        (= ?dn_add_idx (Op (Add ?dn_add_shape ?dn_add_a_stride ?dn_add_b_stride ?dn_add_out_stride) (ICons ?dn_mul_base (ICons ?dn_iota_within (INil)))))
+        (= ?dn_gathered (Op (Gather ?dn_gather_idx_shape ?dn_gather_idx_stride ?dn_gather_data_shape ?dn_gather_data_stride) (ICons ?dn_add_idx (ICons ?down_w (INil)))))
         (= ?dn_f32 (Op (Cast ?dn_f32_size (F32)) (ICons ?dn_gathered (INil))))
-
-        ; ===== Down batched matmul =====
-        ; t83: Mul swiglu_out * gathered_down (broadcast multiply)
         (= ?dn_matmul_mul (Op (Mul ?dn_matmul_mul_shape ?dn_matmul_a_stride ?dn_matmul_b_stride ?dn_matmul_mul_out_stride) (ICons ?swiglu_out (ICons ?dn_f32 (INil)))))
-        ; t84: SumReduce
         (= ?dn_matmul (Op (Sum ?dn_matmul_out_shape ?dn_matmul_k ?dn_matmul_in_stride ?dn_matmul_k_stride ?dn_matmul_out_stride) (ICons ?dn_matmul_mul (INil))))
+    )
+    (
+        (set (glumoe_swiglu_down ?dn_matmul)
+            (MkGLUMoESwiGLUDownState ?dn_io ?dn_matmul_k ?dn_iota_within_range ?swiglu_state ?topk_idx ?down_w))
+    )
+    :name "GLUMoE swiglu down marker"
+)
+
+; ===== Gemma GELU down marker =====
+(rule
+    (
+        (= ?gemma_state (glumoe_gemma_gelu ?gemma_out))
+        (= ?gemma_state (MkGLUMoEGemmaGELUState ?gate_up_state))
+
+        (= ?dn_iota_base (Op (Iota ?dn_io ?dn_iota_base_range) (INil)))
+        (= ?dn_mul_base (Op (Mul ?dn_mul_base_shape ?dn_mul_base_a_stride ?dn_mul_base_b_stride ?dn_mul_base_out_stride) (ICons ?topk_idx (ICons ?dn_iota_base (INil)))))
+        (= ?dn_iota_within (Op (Iota (MIter) ?dn_iota_within_range) (INil)))
+        (= ?dn_add_idx (Op (Add ?dn_add_shape ?dn_add_a_stride ?dn_add_b_stride ?dn_add_out_stride) (ICons ?dn_mul_base (ICons ?dn_iota_within (INil)))))
+        (= ?dn_gathered (Op (Gather ?dn_gather_idx_shape ?dn_gather_idx_stride ?dn_gather_data_shape ?dn_gather_data_stride) (ICons ?dn_add_idx (ICons ?down_w (INil)))))
+        (= ?dn_f32 (Op (Cast ?dn_f32_size (F32)) (ICons ?dn_gathered (INil))))
+        (= ?dn_matmul_mul (Op (Mul ?dn_matmul_mul_shape ?dn_matmul_a_stride ?dn_matmul_b_stride ?dn_matmul_mul_out_stride) (ICons ?gemma_out (ICons ?dn_f32 (INil)))))
+        (= ?dn_matmul (Op (Sum ?dn_matmul_out_shape ?dn_matmul_k ?dn_matmul_in_stride ?dn_matmul_k_stride ?dn_matmul_out_stride) (ICons ?dn_matmul_mul (INil))))
+    )
+    (
+        (set (glumoe_gemma_down ?dn_matmul)
+            (MkGLUMoEGemmaDownState ?dn_io ?dn_matmul_k ?dn_iota_within_range ?gemma_state ?topk_idx ?down_w))
+    )
+    :name "GLUMoE gemma down marker"
+)
+
+; ===== Final fusion: mode 0 (SwiGLU) =====
+(rule
+    (
+        (= ?down_state (glumoe_swiglu_down ?dn_matmul))
+        (= ?down_state (MkGLUMoESwiGLUDownState ?dn_io ?dn_matmul_k ?dn_within_range ?swiglu_state ?topk_idx ?down_w))
+        (= ?swiglu_state (MkGLUMoESwiGLUState ?gate_up_state))
+        (= ?gate_up_state (MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_within_range ?x ?topk_idx ?gate_up_w))
 
-        ; ===== Weighted sum over k experts =====
-        ; t85: Mul down_out * topk_values
         (= ?weighted (Op (Mul ?weighted_shape ?weighted_a_stride ?weighted_b_stride ?weighted_out_stride) (ICons ?dn_matmul (ICons ?topk_vals (INil)))))
-        ; t86: SumReduce over k dimension → [s, H]
         (= ?output (Op (Sum ?output_shape ?output_k ?output_in_stride ?output_k_stride ?output_out_stride) (ICons ?weighted (INil))))
     )
     (
         (let ?glumoe (Op (GLUMoE
             ?gu_io ?dn_io ?gu_matmul_k ?dn_matmul_k ?output_k
-            ?gu_iota_within_range ?dn_iota_within_range)
-            (ICons ?x (ICons ?topk_idx (ICons ?topk_vals (ICons ?gate_up_w (ICons ?down_w (INil))))))))
+            ?gu_within_range ?dn_within_range (MNum 0))
+            (ICons ?x (ICons ?topk_idx (ICons ?topk_vals (ICons ?gate_up_w (ICons ?down_w (ICons ?topk_vals (INil)))))))))
         (union ?output ?glumoe)
     )
-    :name "GLUMoE fused expert computation"
+    :name "GLUMoE fused expert computation (swiglu)"
+)
+
+; ===== Final fusion: mode 1 (Gemma GELU) =====
+(rule
+    (
+        (= ?down_state (glumoe_gemma_down ?dn_matmul))
+        (= ?down_state (MkGLUMoEGemmaDownState ?dn_io ?dn_matmul_k ?dn_within_range ?gemma_state ?topk_idx ?down_w))
+        (= ?gemma_state (MkGLUMoEGemmaGELUState ?gate_up_state))
+        (= ?gate_up_state (MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_within_range ?x ?topk_idx ?gate_up_w))
+
+        ; Gemma expert weights: topk_weights = normed_topk * per_expert_scale.gather(topk_idx)
+        (= ?per_expert_vals (Op (Gather ?scale_gather_idx_shape ?scale_gather_idx_stride ?scale_gather_data_shape ?scale_gather_data_stride) (ICons ?topk_idx (ICons ?per_expert_scale (INil)))))
+        (= ?topk_row_offsets (Op (Iota ?topk_row_offsets_expr ?topk_row_offsets_range) (INil)))
+        (= ?topk_flat_idx (Op (Add ?topk_flat_idx_shape ?topk_flat_idx_a_stride ?topk_flat_idx_b_stride ?topk_flat_idx_out_stride) (ICons ?topk_row_offsets (ICons ?topk_idx (INil)))))
+        (= ?topk_vals (Op (Gather ?topk_vals_gather_idx_shape ?topk_vals_gather_idx_stride ?topk_vals_gather_data_shape ?topk_vals_gather_data_stride) (ICons ?topk_flat_idx (ICons ?routing_weights (INil)))))
+        (= ?topk_norm (Op (Sum ?topk_norm_shape ?output_k ?topk_norm_in_stride ?topk_norm_k_stride ?topk_norm_out_stride) (ICons ?topk_vals (INil))))
+        (= ?topk_norm_factor (Op (Recip ?topk_norm_recip_shape ?topk_norm_recip_in_stride ?topk_norm_recip_out_stride) (ICons ?topk_norm (INil))))
+        (= ?normed_topk (Op (Mul ?normed_topk_shape ?normed_topk_a_stride ?normed_topk_b_stride ?normed_topk_out_stride) (ICons ?topk_vals (ICons ?topk_norm_factor (INil)))))
+        (= ?expert_weights (Op (Mul ?expert_weights_shape ?expert_weights_a_stride ?expert_weights_b_stride ?expert_weights_out_stride) (ICons ?normed_topk (ICons ?per_expert_vals (INil)))))
+
+        (= ?weighted (Op (Mul ?weighted_shape ?weighted_a_stride ?weighted_b_stride ?weighted_out_stride) (ICons ?dn_matmul (ICons ?expert_weights (INil)))))
+        (= ?output (Op (Sum ?output_shape ?output_k ?output_in_stride ?output_k_stride ?output_out_stride) (ICons ?weighted (INil))))
+    )
+    (
+        (let ?glumoe (Op (GLUMoE
+            ?gu_io ?dn_io ?gu_matmul_k ?dn_matmul_k ?output_k
+            ?gu_within_range ?dn_within_range (MNum 1))
+            (ICons ?x (ICons ?topk_idx (ICons ?topk_vals (ICons ?gate_up_w (ICons ?down_w (ICons ?per_expert_scale (INil)))))))))
+        (union ?output ?glumoe)
+    )
+    :name "GLUMoE fused expert computation (gemma_gelu)"
 )

crates/luminal_cuda_lite/src/host/moe/mod.rs

@@ -33,14 +33,15 @@ use crate::{
         },
     },
     host::HostOp,
+    try_create_cublaslt,
 };
 
 const WORKSPACE_SIZE: usize = 32 * 1024 * 1024; // 32 MiB
 
 /// Fused GLU-MoE HostOp matched via egglog pattern.
 ///
-/// Replaces the expert computation subgraph (expert gathers + matmuls + SwiGLU
-/// + weighted sum) with an efficient cuBLASLt implementation.
+/// Replaces the expert computation subgraph (expert gathers + matmuls + gated
+/// activation + weighted sum) with an efficient cuBLASLt implementation.
 ///
 /// Inputs (graph edges, in order):
 ///   0: x              [seq, hidden]                        F32
@@ -48,9 +49,13 @@ const WORKSPACE_SIZE: usize = 32 * 1024 * 1024; // 32 MiB
 ///   2: topk_values    [seq, k]                             F32
 ///   3: gate_up_w      [E, gate_up_dim, hidden]             BF16
 ///   4: down_w         [E, hidden, intermediate]             BF16
+///   5: mode_aux
+///      - SwiGLU: ignored (rewriter wires `topk_values` again)
+///      - GemmaGELU: per_expert_scale [E]                   F32
 ///
 /// Output: [seq, hidden] F32
 pub struct GLUMoE {
+    pub(crate) mode: GLUMoEMode,
     /// Product of gate_up weight dimensions per expert (gate_up_dim * hidden) used for gather stride
     gu_io: Expression,
     /// Product of down weight dimensions per expert (hidden * intermediate) used for gather stride
@@ -69,9 +74,35 @@ pub struct GLUMoE {
     module: OnceLock<(Arc<CudaModule>, CudaFunction, CudaFunction)>,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum GLUMoEMode {
+    SwiGLU,
+    GemmaGELU,
+}
+
+impl GLUMoEMode {
+    fn from_mode_id(mode_id: usize) -> Self {
+        match mode_id {
+            0 => Self::SwiGLU,
+            1 => Self::GemmaGELU,
+            other => {
+                panic!("Unknown GLUMoE mode id: {other}");
+            }
+        }
+    }
+
+    fn activation_kernel_mode(self) -> i32 {
+        match self {
+            Self::SwiGLU => 0,
+            Self::GemmaGELU => 1,
+        }
+    }
+}
+
 impl Default for GLUMoE {
     fn default() -> Self {
         Self {
+            mode: GLUMoEMode::SwiGLU,
             gu_io: Expression::default(),
             dn_io: Expression::default(),
             gu_matmul_k: Expression::default(),
@@ -88,6 +119,7 @@ impl Default for GLUMoE {
 impl std::fmt::Debug for GLUMoE {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("GLUMoE")
+            .field("mode", &self.mode)
             .field("gu_io", &self.gu_io)
             .field("dn_io", &self.dn_io)
             .field("gu_matmul_k", &self.gu_matmul_k)
@@ -100,6 +132,7 @@ impl std::fmt::Debug for GLUMoE {
 impl Clone for GLUMoE {
     fn clone(&self) -> Self {
         Self {
+            mode: self.mode,
             gu_io: self.gu_io,
             dn_io: self.dn_io,
             gu_matmul_k: self.gu_matmul_k,
@@ -114,9 +147,15 @@ impl Clone for GLUMoE {
 }
 
 impl GLUMoE {
-    fn get_cublaslt(&self, stream: &Arc<CudaStream>) -> &Arc<CudaBlasLT> {
-        self.cublaslt
-            .get_or_init(|| Arc::new(CudaBlasLT::new(stream.clone()).unwrap()))
+    fn get_cublaslt(&self, stream: &Arc<CudaStream>) -> anyhow::Result<Arc<CudaBlasLT>> {
+        if let Some(cublaslt) = self.cublaslt.get() {
+            return Ok(cublaslt.clone());
+        }
+        let created = try_create_cublaslt(stream.clone()).map_err(|message| {
+            anyhow::anyhow!("cuBLASLt unavailable on this machine: {message}")
+        })?;
+        let _ = self.cublaslt.set(created.clone());
+        Ok(created)
     }
 
     fn get_kernels(
@@ -134,23 +173,34 @@ extern "C" __global__ void f32_to_bf16(unsigned long long in_ptr, unsigned long
     if (i < n) out[i] = __float2bfloat16(in_[i]);
 }
 
-extern "C" __global__ void swiglu_bf16(unsigned long long gate_up_ptr, unsigned long long out_ptr, int intermediate) {
+extern "C" __global__ void glu_activation_bf16(
+    unsigned long long gate_up_ptr,
+    unsigned long long out_ptr,
+    int intermediate,
+    int mode
+) {
     const __nv_bfloat16* gate_up = (const __nv_bfloat16*)gate_up_ptr;
     __nv_bfloat16* out = (__nv_bfloat16*)out_ptr;
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i < intermediate) {
         float gate = __bfloat162float(gate_up[i]);
         float up   = __bfloat162float(gate_up[i + intermediate]);
-        float silu = gate / (1.0f + expf(-gate));
-        out[i] = __float2bfloat16(silu * up);
+        float activated;
+        if (mode == 0) {
+            activated = gate / (1.0f + expf(-gate));
+        } else {
+            float scaled = 1.5957691216f * gate * (1.0f + 0.044715f * gate * gate);
+            activated = gate / (1.0f + expf(-scaled));
+        }
+        out[i] = __float2bfloat16(activated * up);
     }
 }
 "#;
             let ptx = compile_module_image_for_current_device(stream.context(), src).unwrap();
             let module = stream.context().load_module(ptx).unwrap();
             let f32_to_bf16 = module.load_function("f32_to_bf16").unwrap();
-            let swiglu = module.load_function("swiglu_bf16").unwrap();
-            (module, f32_to_bf16, swiglu)
+            let activation = module.load_function("glu_activation_bf16").unwrap();
+            (module, f32_to_bf16, activation)
         })
     }
 }
@@ -168,12 +218,27 @@ impl EgglogOp for GLUMoE {
                 ("output_k", EXPRESSION),
                 ("gu_within_range", EXPRESSION),
                 ("dn_within_range", EXPRESSION),
+                ("mode", EXPRESSION),
             ],
         )
     }
 
+    fn rewrites(&self) -> Vec<Rule> {
+        vec![Rule::raw(
+            "(rule
+                (
+                    (= ?e (Op (GLUMoE ?gu_io ?dn_io ?gu_matmul_k ?dn_matmul_k ?output_k ?gu_within_range ?dn_within_range ?mode) ?inputs))
+                )
+                (
+                    (set (dtype ?e) (F32))
+                )
+                :ruleset dtype_prop
+            )",
+        )]
+    }
+
     fn n_inputs(&self) -> usize {
-        5
+        6
     }
 
     fn early_rewrites(&self) -> Vec<Rule> {
@@ -195,8 +260,14 @@ impl EgglogOp for GLUMoE {
         let output_k = extract_expr(egraph, kind_children[4], expr_cache).unwrap();
         let gu_within_range = extract_expr(egraph, kind_children[5], expr_cache).unwrap();
         let dn_within_range = extract_expr(egraph, kind_children[6], expr_cache).unwrap();
+        let mode_expr = extract_expr(egraph, kind_children[7], expr_cache).unwrap();
+        let mode_id = mode_expr
+            .to_usize()
+            .unwrap_or_else(|| panic!("GLUMoE mode must be static, got expression: {mode_expr}"));
+        let mode = GLUMoEMode::from_mode_id(mode_id);
 
         let extracted = GLUMoE {
+            mode,
             gu_io,
             dn_io,
             gu_matmul_k,
@@ -209,7 +280,7 @@ impl EgglogOp for GLUMoE {
         };
 
         let op = LLIROp::new::<dyn HostOp>(Box::new(extracted) as Box<dyn HostOp>);
-        // Return the 5 IR inputs: x, topk_idx, topk_vals, gate_up_w, down_w
+        // Return the 6 IR inputs: x, topk_idx, topk_values, gate_up_w, down_w, mode_aux
         (op, input_enodes)
     }
 
@@ -230,9 +301,9 @@ impl HostOp for GLUMoE {
         // Resolve dimensions
         let hidden = self.gu_matmul_k.exec(dyn_map).unwrap();
         let intermediate = self.dn_matmul_k.exec(dyn_map).unwrap();
-        let top_k = self.output_k.exec(dyn_map).unwrap();
+        let top_k_expected = self.output_k.exec(dyn_map).unwrap();
         let gate_up_dim = self.gu_io.exec(dyn_map).unwrap() / hidden; // gate_up_dim = gu_io / hidden
-        let _num_experts = self.gu_within_range.exec(dyn_map).unwrap() / (gate_up_dim * hidden);
+        let num_experts = self.gu_within_range.exec(dyn_map).unwrap() / (gate_up_dim * hidden);
 
         // Derive seq from x buffer size: x is [seq, hidden] F32 → seq = len / (hidden * 4)
         let x_buf = buffers[&inputs[0]];
@@ -243,6 +314,7 @@ impl HostOp for GLUMoE {
         let topk_vals_buf = buffers[&inputs[2]]; // [seq, k] F32
         let gate_up_buf = buffers[&inputs[3]]; // [E, gate_up_dim, hidden] BF16
         let down_buf = buffers[&inputs[4]]; // [E, hidden, intermediate] BF16
+        let mode_aux_buf = buffers[&inputs[5]];
         let output_buf = buffers[&self_node]; // [seq, hidden] F32
 
         // Get raw device pointer addresses
@@ -251,14 +323,59 @@ impl HostOp for GLUMoE {
         let down_ptr = buf_ptr(down_buf, stream);
         let output_ptr = buf_ptr(output_buf, stream);
 
-        let cublaslt = self.get_cublaslt(stream);
-        let (_, f32_to_bf16_fn, swiglu_fn) = self.get_kernels(stream);
+        let cublaslt = self.get_cublaslt(stream)?;
+        let (_, f32_to_bf16_fn, activation_fn) = self.get_kernels(stream);
 
-        // Read topk indices and values from GPU
+        // Read top-k routing values from GPU
         let topk_idx_host: Vec<u8> = stream.clone_dtoh(topk_idx_buf)?;
         let topk_idx_i32: &[i32] = bytemuck::cast_slice(&topk_idx_host);
         let topk_vals_host: Vec<u8> = stream.clone_dtoh(topk_vals_buf)?;
         let topk_vals_f32: &[f32] = bytemuck::cast_slice(&topk_vals_host);
+        let idx_k = topk_idx_i32
+            .len()
+            .checked_div(seq)
+            .unwrap_or(top_k_expected);
+        let val_k = topk_vals_f32
+            .len()
+            .checked_div(seq)
+            .unwrap_or(top_k_expected);
+        let top_k = idx_k.min(val_k);
+        if seq > 0 && top_k == 0 {
+            return Ok(());
+        }
+
+        // Mode-dependent expert weights used for the final reduction:
+        // - SwiGLU: direct topk values
+        // - GemmaGELU: normalize topk values and scale by per-expert factors
+        let mut expert_weights_storage: Vec<f32> = Vec::new();
+        let expert_weights_f32: &[f32] = match self.mode {
+            GLUMoEMode::SwiGLU => topk_vals_f32,
+            GLUMoEMode::GemmaGELU => {
+                let per_expert_scale_host: Vec<u8> = stream.clone_dtoh(mode_aux_buf)?;
+                let per_expert_scale_f32: &[f32] = bytemuck::cast_slice(&per_expert_scale_host);
+                debug_assert!(per_expert_scale_f32.len() >= num_experts);
+                expert_weights_storage.resize(seq * top_k, 0.0);
+                for t in 0..seq {
+                    let base = t * top_k;
+                    let vals = &topk_vals_f32[base..base + top_k];
+                    let norm = vals.iter().copied().sum::<f32>();
+                    let inv_norm = if norm != 0.0 { norm.recip() } else { 0.0 };
+                    for i in 0..top_k {
+                        let expert_idx = topk_idx_i32[base + i] as usize;
+                        if expert_idx >= per_expert_scale_f32.len() {
+                            anyhow::bail!(
+                                "GLUMoE Gemma mode expert index {} out of bounds {}",
+                                expert_idx,
+                                per_expert_scale_f32.len()
+                            );
+                        }
+                        let scale = per_expert_scale_f32[expert_idx];
+                        expert_weights_storage[base + i] = vals[i] * inv_norm * scale;
+                    }
+                }
+                &expert_weights_storage
+            }
+        };
 
         // Allocate temp buffers
         let x_bf16_buf = unsafe { stream.alloc::<u8>(seq * hidden * 2)? }; // BF16
@@ -291,22 +408,10 @@ impl HostOp for GLUMoE {
         let gu_stride = (gate_up_dim * hidden * 2) as u64; // bytes per expert gate_up (BF16)
         let down_stride = (hidden * intermediate * 2) as u64; // bytes per expert down (BF16)
 
-        // Normalize top-k values per token (norm_topk_prob=true)
-        let mut normalized_vals = topk_vals_f32.to_vec();
-        for t in 0..seq {
-            let row = &mut normalized_vals[t * top_k..(t + 1) * top_k];
-            let sum: f32 = row.iter().sum();
-            if sum > 0.0 {
-                for v in row.iter_mut() {
-                    *v /= sum;
-                }
-            }
-        }
-
         for t in 0..seq {
             let x_t_ptr = xbf16_ptr + (t * hidden * 2) as u64; // BF16
             let expert_indices = &topk_idx_i32[t * top_k..(t + 1) * top_k];
-            let weights = &normalized_vals[t * top_k..(t + 1) * top_k];
+            let weights = &expert_weights_f32[t * top_k..(t + 1) * top_k];
 
             for (i, (&expert_idx, &weight)) in expert_indices.iter().zip(weights.iter()).enumerate()
             {
@@ -316,7 +421,7 @@ impl HostOp for GLUMoE {
                 let expert_gu_ptr = gate_up_ptr + expert_idx as u64 * gu_stride;
                 cublas_matmul(
                     stream,
-                    cublaslt,
+                    &cublaslt,
                     ws_ptr,
                     gate_up_dim as u64,
                     1,
@@ -335,17 +440,19 @@ impl HostOp for GLUMoE {
                     0.0f32,
                 )?;
 
-                // b. SwiGLU kernel (BF16 → BF16)
+                // b. Mode-specific gated activation (BF16 → BF16)
                 let moe_int = intermediate as i32;
-                let swiglu_blocks = (moe_int as u32).div_ceil(256);
+                let activation_mode = self.mode.activation_kernel_mode();
+                let activation_blocks = (moe_int as u32).div_ceil(256);
                 unsafe {
                     stream
-                        .launch_builder(swiglu_fn)
+                        .launch_builder(activation_fn)
                         .arg(&gu_out_ptr)
                         .arg(&hid_ptr)
                         .arg(&moe_int)
+                        .arg(&activation_mode)
                         .launch(LaunchConfig {
-                            grid_dim: (swiglu_blocks, 1, 1),
+                            grid_dim: (activation_blocks, 1, 1),
                             block_dim: (256, 1, 1),
                             shared_mem_bytes: 0,
                         })?;
@@ -358,7 +465,7 @@ impl HostOp for GLUMoE {
                 let beta = if i == 0 { 0.0f32 } else { 1.0f32 };
                 cublas_matmul_mixed(
                     stream,
-                    cublaslt,
+                    &cublaslt,
                     ws_ptr,
                     hidden as u64,
                     1,

crates/luminal_cuda_lite/src/kernel/cuda_graph.rs

@@ -653,4 +653,53 @@ mod tests {
         }
         assert_close(&rt.get_f32(output), &expected, 1e-2, 1e-2);
     }
+
+    /// Test that CUDA graphs produce correct results when dynamic dimensions
+    /// change incrementally across many executions (simulating a decode loop
+    /// where position offset increments each step).
+    #[test]
+    fn test_cuda_graph_incremental_dim_changes() {
+        let Some(stream) = get_cuda_stream() else {
+            return;
+        };
+        let mut cx = Graph::default();
+        let a = cx.tensor('s');
+        let b = cx.tensor('s');
+        let c = ((a + b) * a).output();
+
+        let initial_size = 128;
+        cx.set_dim('s', initial_size);
+        let mut rt = CudaRuntime::initialize(stream);
+        let data_a = random_f32_vec(initial_size, 42, -0.5, 0.5);
+        let data_b = random_f32_vec(initial_size, 43, -0.5, 0.5);
+        rt.set_data(a, data_a.clone());
+        rt.set_data(b, data_b.clone());
+        cx.build_search_space::<CudaRuntime>();
+        rt = cx.search(rt, 5);
+
+        // Initial execution
+        rt.execute(&cx.dyn_map);
+        let eps = dtype_epsilon(luminal::dtype::DType::F32);
+        let tol = eps * TOLERANCE_SAFETY_FACTOR;
+        let expected: Vec<f32> = data_a
+            .iter()
+            .zip(&data_b)
+            .map(|(a, b)| (a + b) * a)
+            .collect();
+        assert_close(&rt.get_f32(c), &expected, tol, tol);
+
+        // Incrementally change the dynamic dimension 10 times,
+        // simulating decode steps where position offset grows.
+        for step in 1..=10usize {
+            let size = initial_size + step;
+            cx.set_dim('s', size);
+            let da = random_f32_vec(size, 100 + step as u64, -0.5, 0.5);
+            let db = random_f32_vec(size, 200 + step as u64, -0.5, 0.5);
+            rt.set_data(a, da.clone());
+            rt.set_data(b, db.clone());
+            rt.execute(&cx.dyn_map);
+            let expected: Vec<f32> = da.iter().zip(&db).map(|(a, b)| (a + b) * a).collect();
+            assert_close(&rt.get_f32(c), &expected, tol, tol);
+        }
+    }
 }

crates/luminal_cuda_lite/src/kernel/fusion/fused_ops.rs

@@ -0,0 +1,451 @@
+// =========================================================================
+// Fused elementwise op variants used inside FusionStart/FusionEnd regions.
+//
+// Each `FusedX` struct mirrors its un-fused `KernelX` sibling field-for-field
+// and serves a single purpose: give the egglog rules a distinct sort to
+// rewrite into so a pair-fuse rule's RHS can never re-match its own LHS
+// pattern. Cascade prevention by typing.
+//
+// `compile()` is a *fallback* path. The fast path collapses each FE-rooted
+// region into one CUDA kernel inside `region_codegen` and FusedX/FS/FE
+// never reach kernel_to_host's compile loop. But extraction can produce
+// LLIR shapes the detector doesn't sweep into a region, so each FusedX's
+// standalone `compile()` falls back to emitting the same kernel its
+// un-fused KernelX sibling would — correct, just one launch per op.
+// =========================================================================
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::{
+    egglog_utils::{
+        api::{Rule, SortDef, sort},
+        base::{DTYPE, ELIST, OP_KIND},
+        extract_dtype, extract_expr_list,
+    },
+    op::*,
+    prelude::*,
+};
+
+use crate::{
+    compile_module_image_for_current_device, cuda_dtype,
+    kernel::KernelOp,
+    kernel::hlir::{dtype_includes, generate_dyn_dims_defines},
+};
+
+pub type Ops = (
+    FusedSin,
+    FusedSqrt,
+    FusedExp,
+    FusedExp2,
+    FusedLog2,
+    FusedRecip,
+    FusedAdd,
+    FusedMul,
+);
+
+// Standard `compile()` return tuple (matches the trait signature).
+type CompileOut = (
+    CudaFunction,
+    Arc<CudaModule>,
+    String,
+    (Expression, Expression, Expression),
+    (Expression, Expression, Expression),
+    Expression,
+    FxHashMap<char, CudaSlice<u8>>,
+);
+
+// =========================================================================
+// Fallback kernel templates — used when a FusedX op reaches
+// `kernel_to_host` standalone (region detection missed it). Same CUDA as
+// the matching un-fused KernelX would emit, parameterised by the per-op
+// body expression. The fast path goes through `region_codegen`.
+// =========================================================================
+
+#[allow(clippy::too_many_arguments)]
+fn compile_unary_fallback(
+    stream: &Arc<CudaStream>,
+    compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    kernel_name: &str,
+    body_expr: &str, // CUDA expression on `in[{in_idx}]`, e.g. "sinf(in[{in_idx}])"
+    shape: &[Expression],
+    in_strides: &[Expression],
+    out_strides: &[Expression],
+    dtype: DType,
+) -> CompileOut {
+    let vars = shape
+        .iter()
+        .flat_map(|e| e.dyn_vars())
+        .chain(in_strides.iter().flat_map(|e| e.dyn_vars()))
+        .chain(out_strides.iter().flat_map(|e| e.dyn_vars()))
+        .collect::<FxHashSet<_>>();
+    let cuda_ty = cuda_dtype(dtype);
+    let includes = dtype_includes(&[dtype]);
+    let (dyn_defines, _sorted_dims) = generate_dyn_dims_defines(&vars);
+    let dyn_dims_param = if vars.is_empty() {
+        ""
+    } else {
+        ", const int* dyn_dims"
+    };
+    let n_elements = shape.iter().copied().product::<Expression>().to_kernel();
+    let out_idx = flatten_strides(shape, out_strides).to_kernel();
+    let in_idx = flatten_strides(shape, in_strides).to_kernel();
+    let body = body_expr.replace("{in_idx}", &in_idx);
+    let kernel = format!(
+        "{includes}\n{dyn_defines}\nextern \"C\" {{\n\
+         \x20   __global__ void {kernel_name}({cuda_ty} *out, const {cuda_ty} *in{dyn_dims_param}) {{\n\
+         \x20       long long const_z = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n\
+         \x20       if (const_z >= {n_elements}) return;\n\
+         \x20       out[{out_idx}] = {body};\n\
+         \x20   }}\n}}"
+    );
+    let (module, func) = if let Some((m, f)) = compile_cache.get(&kernel) {
+        (m.clone(), f.clone())
+    } else {
+        let ptx = compile_module_image_for_current_device(stream.context(), &kernel).unwrap();
+        let module = stream.context().load_module(ptx).unwrap();
+        let func = module.load_function(kernel_name).unwrap();
+        compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+        (module, func)
+    };
+    let out_size = shape.iter().copied().product::<Expression>();
+    (
+        func,
+        module,
+        kernel,
+        (out_size.ceil_div(256), 1.into(), 1.into()),
+        (out_size.min(256), 1.into(), 1.into()),
+        0.into(),
+        FxHashMap::default(),
+    )
+}
+
+#[allow(clippy::too_many_arguments)]
+fn compile_binary_fallback(
+    stream: &Arc<CudaStream>,
+    compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    kernel_name: &str,
+    op_str: &str, // CUDA infix operator, e.g. "+", "*"
+    out_shape: &[Expression],
+    a_stride: &[Expression],
+    b_stride: &[Expression],
+    out_stride: &[Expression],
+    dtype: DType,
+) -> CompileOut {
+    let vars = out_shape
+        .iter()
+        .flat_map(|e| e.dyn_vars())
+        .chain(a_stride.iter().flat_map(|e| e.dyn_vars()))
+        .chain(b_stride.iter().flat_map(|e| e.dyn_vars()))
+        .chain(out_stride.iter().flat_map(|e| e.dyn_vars()))
+        .collect::<FxHashSet<_>>();
+    let cuda_ty = cuda_dtype(dtype);
+    let includes = dtype_includes(&[dtype, dtype]);
+    let (dyn_defines, _sorted_dims) = generate_dyn_dims_defines(&vars);
+    let dyn_dims_param = if vars.is_empty() {
+        ""
+    } else {
+        ", const int* dyn_dims"
+    };
+    let n_elements = out_shape
+        .iter()
+        .copied()
+        .product::<Expression>()
+        .to_kernel();
+    let out_idx = flatten_strides(out_shape, out_stride).to_kernel();
+    let a_idx = flatten_strides(out_shape, a_stride).to_kernel();
+    let b_idx = flatten_strides(out_shape, b_stride).to_kernel();
+    let kernel = format!(
+        "{includes}\n{dyn_defines}\nextern \"C\" {{\n\
+         \x20   __global__ void {kernel_name}({cuda_ty} *C, const {cuda_ty} *A, const {cuda_ty} *B{dyn_dims_param}) {{\n\
+         \x20       long long const_z = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n\
+         \x20       if (const_z >= {n_elements}) return;\n\
+         \x20       C[{out_idx}] = A[{a_idx}] {op_str} B[{b_idx}];\n\
+         \x20   }}\n}}"
+    );
+    let (module, func) = if let Some((m, f)) = compile_cache.get(&kernel) {
+        (m.clone(), f.clone())
+    } else {
+        let ptx = compile_module_image_for_current_device(stream.context(), &kernel).unwrap();
+        let module = stream.context().load_module(ptx).unwrap();
+        let func = module.load_function(kernel_name).unwrap();
+        compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+        (module, func)
+    };
+    let out_size = out_shape.iter().copied().product::<Expression>();
+    (
+        func,
+        module,
+        kernel,
+        (out_size.ceil_div(256), 1.into(), 1.into()),
+        (out_size.min(256), 1.into(), 1.into()),
+        0.into(),
+        FxHashMap::default(),
+    )
+}
+
+/// Generate `pub struct $Name { … unary fields … }` plus its `EgglogOp` and
+/// `KernelOp` impls. `$kernel_name` names the CUDA function (and the cache
+/// key); `$body` is the per-op CUDA expression, e.g. `"sinf(in[{in_idx}])"`.
+macro_rules! impl_fused_unary {
+    ($Name:ident, $sort:literal, $kernel_name:literal, $body:literal) => {
+        #[derive(Default, Debug, Clone)]
+        pub struct $Name {
+            pub(crate) shape: Vec<Expression>,
+            pub(crate) in_strides: Vec<Expression>,
+            pub(crate) out_strides: Vec<Expression>,
+            pub(crate) dtype: DType,
+        }
+
+        impl EgglogOp for $Name {
+            fn sort(&self) -> SortDef {
+                sort(
+                    OP_KIND,
+                    $sort,
+                    &[
+                        ("shape", ELIST),
+                        ("strides", ELIST),
+                        ("out_strides", ELIST),
+                        ("dtype", DTYPE),
+                    ],
+                )
+            }
+            fn n_inputs(&self) -> usize {
+                1
+            }
+            fn rewrites(&self) -> Vec<Rule> {
+                Vec::new()
+            }
+            fn cleanup(&self) -> bool {
+                false
+            }
+            fn extract<'a>(
+                &'a self,
+                egraph: &'a SerializedEGraph,
+                kind_children: &[&'a ENodeId],
+                input_enodes: Vec<&'a ENodeId>,
+                list_cache: &mut FxHashMap<&'a ENodeId, Vec<Expression>>,
+                expr_cache: &mut FxHashMap<&'a ENodeId, Expression>,
+            ) -> (LLIROp, Vec<&'a ENodeId>) {
+                (
+                    LLIROp::new::<dyn KernelOp>(Box::new(Self {
+                        shape: extract_expr_list(egraph, kind_children[0], list_cache, expr_cache)
+                            .unwrap(),
+                        in_strides: extract_expr_list(
+                            egraph,
+                            kind_children[1],
+                            list_cache,
+                            expr_cache,
+                        )
+                        .unwrap(),
+                        out_strides: extract_expr_list(
+                            egraph,
+                            kind_children[2],
+                            list_cache,
+                            expr_cache,
+                        )
+                        .unwrap(),
+                        dtype: extract_dtype(egraph, kind_children[3]),
+                    })),
+                    input_enodes,
+                )
+            }
+        }
+
+        impl KernelOp for $Name {
+            fn compile(
+                &self,
+                stream: &Arc<CudaStream>,
+                compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+            ) -> CompileOut {
+                compile_unary_fallback(
+                    stream,
+                    compile_cache,
+                    $kernel_name,
+                    $body,
+                    &self.shape,
+                    &self.in_strides,
+                    &self.out_strides,
+                    self.dtype,
+                )
+            }
+            fn output_size(&self) -> Expression {
+                self.shape.iter().copied().product()
+            }
+            fn output_bytes(&self) -> Expression {
+                (self.output_size() * self.dtype.bits()).ceil_div(8)
+            }
+            fn bytes_loaded(&self) -> Expression {
+                self.output_bytes()
+            }
+            fn bytes_stored(&self) -> Expression {
+                self.output_bytes()
+            }
+            fn flops(&self) -> Expression {
+                self.shape.iter().copied().product()
+            }
+            fn output_dtype(&self) -> DType {
+                self.dtype
+            }
+            fn kernel_name(&self) -> &'static str {
+                $sort
+            }
+        }
+    };
+}
+
+/// As `impl_fused_unary!` but for binary ops: 5-field sort signature
+/// (shape + per-input strides + out_stride + dtype), n_inputs = 2.
+/// `$op_str` is the CUDA infix operator, e.g. `"+"`, `"*"`.
+macro_rules! impl_fused_binary {
+    ($Name:ident, $sort:literal, $kernel_name:literal, $op_str:literal) => {
+        #[derive(Default, Debug, Clone)]
+        pub struct $Name {
+            pub(crate) out_shape: Vec<Expression>,
+            pub(crate) a_stride: Vec<Expression>,
+            pub(crate) b_stride: Vec<Expression>,
+            pub(crate) out_stride: Vec<Expression>,
+            pub(crate) dtype: DType,
+        }
+
+        impl EgglogOp for $Name {
+            fn sort(&self) -> SortDef {
+                sort(
+                    OP_KIND,
+                    $sort,
+                    &[
+                        ("shape", ELIST),
+                        ("a_strides", ELIST),
+                        ("b_strides", ELIST),
+                        ("out_strides", ELIST),
+                        ("dtype", DTYPE),
+                    ],
+                )
+            }
+            fn n_inputs(&self) -> usize {
+                2
+            }
+            fn rewrites(&self) -> Vec<Rule> {
+                Vec::new()
+            }
+            fn cleanup(&self) -> bool {
+                false
+            }
+            fn extract<'a>(
+                &'a self,
+                egraph: &'a SerializedEGraph,
+                kind_children: &[&'a ENodeId],
+                input_enodes: Vec<&'a ENodeId>,
+                list_cache: &mut FxHashMap<&'a ENodeId, Vec<Expression>>,
+                expr_cache: &mut FxHashMap<&'a ENodeId, Expression>,
+            ) -> (LLIROp, Vec<&'a ENodeId>) {
+                (
+                    LLIROp::new::<dyn KernelOp>(Box::new(Self {
+                        out_shape: extract_expr_list(
+                            egraph,
+                            kind_children[0],
+                            list_cache,
+                            expr_cache,
+                        )
+                        .unwrap(),
+                        a_stride: extract_expr_list(
+                            egraph,
+                            kind_children[1],
+                            list_cache,
+                            expr_cache,
+                        )
+                        .unwrap(),
+                        b_stride: extract_expr_list(
+                            egraph,
+                            kind_children[2],
+                            list_cache,
+                            expr_cache,
+                        )
+                        .unwrap(),
+                        out_stride: extract_expr_list(
+                            egraph,
+                            kind_children[3],
+                            list_cache,
+                            expr_cache,
+                        )
+                        .unwrap(),
+                        dtype: extract_dtype(egraph, kind_children[4]),
+                    })),
+                    input_enodes,
+                )
+            }
+        }
+
+        impl KernelOp for $Name {
+            fn compile(
+                &self,
+                stream: &Arc<CudaStream>,
+                compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+            ) -> CompileOut {
+                compile_binary_fallback(
+                    stream,
+                    compile_cache,
+                    $kernel_name,
+                    $op_str,
+                    &self.out_shape,
+                    &self.a_stride,
+                    &self.b_stride,
+                    &self.out_stride,
+                    self.dtype,
+                )
+            }
+            fn output_size(&self) -> Expression {
+                self.out_shape.iter().copied().product()
+            }
+            fn output_bytes(&self) -> Expression {
+                (self.output_size() * self.dtype.bits()).ceil_div(8)
+            }
+            fn bytes_loaded(&self) -> Expression {
+                let bytes = (self.output_size() * self.dtype.bits()).ceil_div(8);
+                bytes + bytes
+            }
+            fn bytes_stored(&self) -> Expression {
+                self.output_bytes()
+            }
+            fn flops(&self) -> Expression {
+                self.out_shape.iter().copied().product()
+            }
+            fn output_dtype(&self) -> DType {
+                self.dtype
+            }
+            fn kernel_name(&self) -> &'static str {
+                $sort
+            }
+        }
+    };
+}
+
+impl_fused_unary!(FusedSin, "FusedSin", "fused_sin_k", "sinf(in[{in_idx}])");
+impl_fused_unary!(
+    FusedSqrt,
+    "FusedSqrt",
+    "fused_sqrt_k",
+    "sqrtf(in[{in_idx}])"
+);
+impl_fused_unary!(FusedExp, "FusedExp", "fused_exp_k", "expf(in[{in_idx}])");
+impl_fused_unary!(
+    FusedExp2,
+    "FusedExp2",
+    "fused_exp2_k",
+    "exp2f(in[{in_idx}])"
+);
+impl_fused_unary!(
+    FusedLog2,
+    "FusedLog2",
+    "fused_log2_k",
+    "log2f(in[{in_idx}])"
+);
+impl_fused_unary!(
+    FusedRecip,
+    "FusedRecip",
+    "fused_recip_k",
+    "1.0f / in[{in_idx}]"
+);
+
+impl_fused_binary!(FusedAdd, "FusedAdd", "fused_add_k", "+");
+impl_fused_binary!(FusedMul, "FusedMul", "fused_mul_k", "*");

crates/luminal_cuda_lite/src/kernel/fusion/markers.rs

@@ -0,0 +1,490 @@
+// =========================================================================
+// Fusion boundary markers — FusionStart and FusionEnd.
+//
+// Tag-like LLIR ops that bracket a region of elementwise ops destined to
+// be emitted as a single CUDA kernel:
+//   - N FusionStart nodes per region (one per FS leaf — distinct external
+//     reads),
+//   - exactly 1 FusionEnd per region.
+//
+// `FusionEnd::rewrites()` carries the seven rule families that build and
+// extend regions (pair-fuse / grow / merge); the actual single-kernel
+// codegen lives in `region_codegen`. Like FusedX, both markers'
+// `compile()` is `unreachable!()` — region codegen folds them away
+// before kernel_to_host's compile loop reaches an interior node.
+// =========================================================================
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::{
+    egglog_utils::{
+        api::{Rule, SortDef, sort},
+        base::{DTYPE, ELIST, OP_KIND},
+        extract_dtype, extract_expr_list,
+    },
+    op::*,
+    prelude::*,
+};
+
+use crate::{
+    compile_module_image_for_current_device, cuda_dtype,
+    kernel::KernelOp,
+    kernel::hlir::{dtype_includes, generate_dyn_dims_defines},
+};
+
+/// Identity-memcpy kernel used as a *fallback* when a FusionStart or
+/// FusionEnd reaches `kernel_to_host`'s compile loop standalone (i.e.,
+/// region detection didn't sweep it into a `CompileUnit::Region`). The
+/// fast path is region collapse, but model-fuzz extraction sometimes
+/// produces LLIR shapes the detector doesn't catch; this keeps
+/// execution correct in those cases.
+#[allow(clippy::type_complexity)]
+fn compile_identity_kernel(
+    stream: &Arc<CudaStream>,
+    compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    kernel_name: &str,
+    shape: &[Expression],
+    strides: &[Expression],
+    dtype: DType,
+) -> CompileOut {
+    let vars = shape
+        .iter()
+        .flat_map(|e| e.dyn_vars())
+        .chain(strides.iter().flat_map(|e| e.dyn_vars()))
+        .collect::<FxHashSet<_>>();
+    let cuda_ty = cuda_dtype(dtype);
+    let includes = dtype_includes(&[dtype]);
+    let (dyn_defines, _sorted_dims) = generate_dyn_dims_defines(&vars);
+    let dyn_dims_param = if vars.is_empty() {
+        ""
+    } else {
+        ", const int* dyn_dims"
+    };
+    let n_elements = shape.iter().copied().product::<Expression>().to_kernel();
+    let idx = flatten_strides(shape, strides).to_kernel();
+    let kernel = format!(
+        "{includes}\n{dyn_defines}\nextern \"C\" {{\n\
+         \x20   __global__ void {kernel_name}({cuda_ty} *out, const {cuda_ty} *in{dyn_dims_param}) {{\n\
+         \x20       long long const_z = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n\
+         \x20       if (const_z >= {n_elements}) return;\n\
+         \x20       out[{idx}] = in[{idx}];\n\
+         \x20   }}\n}}"
+    );
+    let (module, func) = if let Some((m, f)) = compile_cache.get(&kernel) {
+        (m.clone(), f.clone())
+    } else {
+        let ptx = compile_module_image_for_current_device(stream.context(), &kernel).unwrap();
+        let module = stream.context().load_module(ptx).unwrap();
+        let func = module.load_function(kernel_name).unwrap();
+        compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+        (module, func)
+    };
+    let out_size = shape.iter().copied().product::<Expression>();
+    (
+        func,
+        module,
+        kernel,
+        (out_size.ceil_div(256), 1.into(), 1.into()),
+        (out_size.min(256), 1.into(), 1.into()),
+        0.into(),
+        FxHashMap::default(),
+    )
+}
+
+pub type Ops = (FusionStart, FusionEnd);
+
+type CompileOut = (
+    CudaFunction,
+    Arc<CudaModule>,
+    String,
+    (Expression, Expression, Expression),
+    (Expression, Expression, Expression),
+    Expression,
+    FxHashMap<char, CudaSlice<u8>>,
+);
+
+// =========================================================================
+// FusionStart
+// =========================================================================
+
+#[derive(Default, Debug, Clone)]
+pub struct FusionStart {
+    pub(crate) shape: Vec<Expression>,
+    pub(crate) strides: Vec<Expression>,
+    pub(crate) dtype: DType,
+}
+
+impl EgglogOp for FusionStart {
+    fn sort(&self) -> SortDef {
+        sort(
+            OP_KIND,
+            "FusionStart",
+            &[("shape", ELIST), ("strides", ELIST), ("dtype", DTYPE)],
+        )
+    }
+    fn n_inputs(&self) -> usize {
+        1
+    }
+    fn rewrites(&self) -> Vec<Rule> {
+        // No idempotence rule. `FusionStart(FusionStart(x)) ≡ FusionStart(x)`
+        // would unify nested markers and create eclass cycles via the
+        // pair-fuse rules; without it, occasional re-firings produce extra
+        // semantically-correct identity layers, bounded by the run schedule.
+        Vec::new()
+    }
+    fn cleanup(&self) -> bool {
+        false
+    }
+    fn extract<'a>(
+        &'a self,
+        egraph: &'a SerializedEGraph,
+        kind_children: &[&'a ENodeId],
+        input_enodes: Vec<&'a ENodeId>,
+        list_cache: &mut FxHashMap<&'a ENodeId, Vec<Expression>>,
+        expr_cache: &mut FxHashMap<&'a ENodeId, Expression>,
+    ) -> (LLIROp, Vec<&'a ENodeId>) {
+        (
+            LLIROp::new::<dyn KernelOp>(Box::new(Self {
+                shape: extract_expr_list(egraph, kind_children[0], list_cache, expr_cache).unwrap(),
+                strides: extract_expr_list(egraph, kind_children[1], list_cache, expr_cache)
+                    .unwrap(),
+                dtype: extract_dtype(egraph, kind_children[2]),
+            })),
+            input_enodes,
+        )
+    }
+}
+
+impl KernelOp for FusionStart {
+    fn compile(
+        &self,
+        stream: &Arc<CudaStream>,
+        compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    ) -> CompileOut {
+        compile_identity_kernel(
+            stream,
+            compile_cache,
+            "fusion_start_k",
+            &self.shape,
+            &self.strides,
+            self.dtype,
+        )
+    }
+    fn output_size(&self) -> Expression {
+        self.shape.iter().copied().product()
+    }
+    fn output_bytes(&self) -> Expression {
+        (self.output_size() * self.dtype.bits()).ceil_div(8)
+    }
+    fn output_dtype(&self) -> DType {
+        self.dtype
+    }
+    fn kernel_name(&self) -> &'static str {
+        "FusionStart"
+    }
+}
+
+// =========================================================================
+// FusionEnd
+// =========================================================================
+
+#[derive(Default, Debug, Clone)]
+pub struct FusionEnd {
+    pub(crate) shape: Vec<Expression>,
+    pub(crate) strides: Vec<Expression>,
+    pub(crate) dtype: DType,
+}
+
+impl EgglogOp for FusionEnd {
+    fn sort(&self) -> SortDef {
+        sort(
+            OP_KIND,
+            "FusionEnd",
+            &[("shape", ELIST), ("strides", ELIST), ("dtype", DTYPE)],
+        )
+    }
+    fn n_inputs(&self) -> usize {
+        1
+    }
+
+    fn rewrites(&self) -> Vec<Rule> {
+        // Ablation switch: with `LUMINAL_DISABLE_BINARY_FUSION=1` set, do
+        // not register any fusion rules. The e-graph never sees the FS/FE
+        // bracketed alternative, extraction always picks the un-fused
+        // form, and the runtime path matches main with no fusion at all.
+        // Used to A/B fusion's runtime impact on a single binary.
+        if std::env::var("LUMINAL_DISABLE_BINARY_FUSION").is_ok() {
+            return Vec::new();
+        }
+        // Seven rule families build and extend FE-bracketed regions. Each
+        // pair-fuse rule's LHS pattern matches *un-fused* `KernelX` ops; the
+        // RHS produces `FusedX` variants in a different egglog sort, so the
+        // rule's own output cannot re-match its LHS — cascade is prevented
+        // by typing rather than by a discriminator field.
+        //
+        // Stride compatibility is expressed by reusing variable names: a
+        // unary inside a region matches `(KernelU ?shape ?s ?s ?dt)` (in =
+        // out, no transpose); a binary feeding a downstream op binds the
+        // binary's out-stride to the downstream op's in-stride along the
+        // connecting side.
+        let mut rules = Vec::new();
+
+        // (KernelX kind, FusedX kind)
+        let unaries: &[(&str, &str)] = &[
+            ("KernelSin", "FusedSin"),
+            ("KernelSqrt", "FusedSqrt"),
+            ("KernelExp", "FusedExp"),
+            ("KernelExp2", "FusedExp2"),
+            ("KernelLog2", "FusedLog2"),
+            ("KernelRecip", "FusedRecip"),
+        ];
+        // (KernelX kind, FusedX kind, rule-name label)
+        let binaries: &[(&str, &str, &str)] = &[
+            ("KernelAdd", "FusedAdd", "Add"),
+            ("KernelMul", "FusedMul", "Mul"),
+        ];
+
+        // 1. Pair-fuse U → U: U2(U1(x)) → FE(FU2(FU1(FS(x)))).
+        for (ki1, fi1) in unaries {
+            for (ko2, fo2) in unaries {
+                rules.push(Rule::raw(format!(
+                    "(rule (
+                        (= ?u1 (Op ({ki1} ?shape ?s ?s ?dt) (ICons ?x (INil))))
+                        (= ?u2 (Op ({ko2} ?shape ?s ?s ?dt) (ICons ?u1 (INil))))
+                     ) (
+                        (let ?fs (Op (FusionStart ?shape ?s ?dt) (ICons ?x (INil))))
+                        (let ?fu1 (Op ({fi1} ?shape ?s ?s ?dt) (ICons ?fs (INil))))
+                        (let ?fu2 (Op ({fo2} ?shape ?s ?s ?dt) (ICons ?fu1 (INil))))
+                        (let ?fe (Op (FusionEnd ?shape ?s ?dt) (ICons ?fu2 (INil))))
+                        (union ?u2 ?fe)
+                     ) :name \"pair-fuse-U-U-{ki1}-{ko2}\")"
+                )));
+            }
+        }
+
+        // 2. Pair-fuse B → U: U(B(a, b)) → FE(FU(FB(FS(a), FS(b)))).
+        for (kb, fb, lb) in binaries {
+            for (ku, fu) in unaries {
+                rules.push(Rule::raw(format!(
+                    "(rule (
+                        (= ?bin (Op ({kb} ?shape ?a_s ?b_s ?o_s ?dt)
+                                     (ICons ?a (ICons ?b (INil)))))
+                        (= ?u (Op ({ku} ?shape ?o_s ?o_s ?dt) (ICons ?bin (INil))))
+                     ) (
+                        (let ?fs_a (Op (FusionStart ?shape ?a_s ?dt) (ICons ?a (INil))))
+                        (let ?fs_b (Op (FusionStart ?shape ?b_s ?dt) (ICons ?b (INil))))
+                        (let ?fbin (Op ({fb} ?shape ?a_s ?b_s ?o_s ?dt)
+                                       (ICons ?fs_a (ICons ?fs_b (INil)))))
+                        (let ?fu (Op ({fu} ?shape ?o_s ?o_s ?dt) (ICons ?fbin (INil))))
+                        (let ?fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fu (INil))))
+                        (union ?u ?fe)
+                     ) :name \"pair-fuse-B-U-{lb}-{ku}\")"
+                )));
+            }
+        }
+
+        // 3. Pair-fuse U → B (lhs / rhs): unary feeds binary's A or B input.
+        //    LHS:  B(U(a), b) → FE(FB(FU(FS(a)), FS(b))).
+        //    RHS:  B(a, U(b)) → FE(FB(FS(a), FU(FS(b)))).
+        for (ku, fu) in unaries {
+            for (kb, fb, lb) in binaries {
+                rules.push(Rule::raw(format!(
+                    "(rule (
+                        (= ?u (Op ({ku} ?shape ?u_s ?u_s ?dt) (ICons ?a (INil))))
+                        (= ?bin (Op ({kb} ?shape ?u_s ?b_s ?o_s ?dt)
+                                     (ICons ?u (ICons ?b (INil)))))
+                     ) (
+                        (let ?fs_a (Op (FusionStart ?shape ?u_s ?dt) (ICons ?a (INil))))
+                        (let ?fs_b (Op (FusionStart ?shape ?b_s ?dt) (ICons ?b (INil))))
+                        (let ?fu (Op ({fu} ?shape ?u_s ?u_s ?dt) (ICons ?fs_a (INil))))
+                        (let ?fbin (Op ({fb} ?shape ?u_s ?b_s ?o_s ?dt)
+                                       (ICons ?fu (ICons ?fs_b (INil)))))
+                        (let ?fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fbin (INil))))
+                        (union ?bin ?fe)
+                     ) :name \"pair-fuse-U-B-lhs-{ku}-{lb}\")"
+                )));
+                rules.push(Rule::raw(format!(
+                    "(rule (
+                        (= ?u (Op ({ku} ?shape ?u_s ?u_s ?dt) (ICons ?b (INil))))
+                        (= ?bin (Op ({kb} ?shape ?a_s ?u_s ?o_s ?dt)
+                                     (ICons ?a (ICons ?u (INil)))))
+                     ) (
+                        (let ?fs_a (Op (FusionStart ?shape ?a_s ?dt) (ICons ?a (INil))))
+                        (let ?fs_b (Op (FusionStart ?shape ?u_s ?dt) (ICons ?b (INil))))
+                        (let ?fu (Op ({fu} ?shape ?u_s ?u_s ?dt) (ICons ?fs_b (INil))))
+                        (let ?fbin (Op ({fb} ?shape ?a_s ?u_s ?o_s ?dt)
+                                       (ICons ?fs_a (ICons ?fu (INil)))))
+                        (let ?fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fbin (INil))))
+                        (union ?bin ?fe)
+                     ) :name \"pair-fuse-U-B-rhs-{ku}-{lb}\")"
+                )));
+            }
+        }
+
+        // 4. Pair-fuse B → B (lhs / rhs): inner binary feeds outer's A or B.
+        for (kbi, fbi, lbi) in binaries {
+            for (kbo, fbo, lbo) in binaries {
+                rules.push(Rule::raw(format!(
+                    "(rule (
+                        (= ?bi (Op ({kbi} ?shape ?ai_s ?bi_s ?oi_s ?dt)
+                                    (ICons ?a (ICons ?b (INil)))))
+                        (= ?bo (Op ({kbo} ?shape ?oi_s ?co_s ?oo_s ?dt)
+                                    (ICons ?bi (ICons ?c (INil)))))
+                     ) (
+                        (let ?fs_a (Op (FusionStart ?shape ?ai_s ?dt) (ICons ?a (INil))))
+                        (let ?fs_b (Op (FusionStart ?shape ?bi_s ?dt) (ICons ?b (INil))))
+                        (let ?fs_c (Op (FusionStart ?shape ?co_s ?dt) (ICons ?c (INil))))
+                        (let ?fbi (Op ({fbi} ?shape ?ai_s ?bi_s ?oi_s ?dt)
+                                       (ICons ?fs_a (ICons ?fs_b (INil)))))
+                        (let ?fbo (Op ({fbo} ?shape ?oi_s ?co_s ?oo_s ?dt)
+                                       (ICons ?fbi (ICons ?fs_c (INil)))))
+                        (let ?fe (Op (FusionEnd ?shape ?oo_s ?dt) (ICons ?fbo (INil))))
+                        (union ?bo ?fe)
+                     ) :name \"pair-fuse-B-B-lhs-{lbi}-{lbo}\")"
+                )));
+                rules.push(Rule::raw(format!(
+                    "(rule (
+                        (= ?bi (Op ({kbi} ?shape ?ai_s ?bi_s ?oi_s ?dt)
+                                    (ICons ?a (ICons ?b (INil)))))
+                        (= ?bo (Op ({kbo} ?shape ?co_s ?oi_s ?oo_s ?dt)
+                                    (ICons ?c (ICons ?bi (INil)))))
+                     ) (
+                        (let ?fs_a (Op (FusionStart ?shape ?ai_s ?dt) (ICons ?a (INil))))
+                        (let ?fs_b (Op (FusionStart ?shape ?bi_s ?dt) (ICons ?b (INil))))
+                        (let ?fs_c (Op (FusionStart ?shape ?co_s ?dt) (ICons ?c (INil))))
+                        (let ?fbi (Op ({fbi} ?shape ?ai_s ?bi_s ?oi_s ?dt)
+                                       (ICons ?fs_a (ICons ?fs_b (INil)))))
+                        (let ?fbo (Op ({fbo} ?shape ?co_s ?oi_s ?oo_s ?dt)
+                                       (ICons ?fs_c (ICons ?fbi (INil)))))
+                        (let ?fe (Op (FusionEnd ?shape ?oo_s ?dt) (ICons ?fbo (INil))))
+                        (union ?bo ?fe)
+                     ) :name \"pair-fuse-B-B-rhs-{lbi}-{lbo}\")"
+                )));
+            }
+        }
+
+        // 5. Grow FE → U: U(FE(inner)) → FE(FU(inner)). No new FS.
+        for (ku, fu) in unaries {
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?fe (Op (FusionEnd ?shape ?s ?dt) (ICons ?inner (INil))))
+                    (= ?u (Op ({ku} ?shape ?s ?s ?dt) (ICons ?fe (INil))))
+                 ) (
+                    (let ?fu (Op ({fu} ?shape ?s ?s ?dt) (ICons ?inner (INil))))
+                    (let ?new_fe (Op (FusionEnd ?shape ?s ?dt) (ICons ?fu (INil))))
+                    (union ?u ?new_fe)
+                 ) :name \"grow-FE-U-{ku}\")"
+            )));
+        }
+
+        // 6. Grow FE → B (lhs / rhs): one input is the FE, the other external.
+        for (kb, fb, lb) in binaries {
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?fe (Op (FusionEnd ?shape ?a_s ?dt) (ICons ?inner_a (INil))))
+                    (= ?bin (Op ({kb} ?shape ?a_s ?b_s ?o_s ?dt)
+                                 (ICons ?fe (ICons ?b (INil)))))
+                 ) (
+                    (let ?fs_b (Op (FusionStart ?shape ?b_s ?dt) (ICons ?b (INil))))
+                    (let ?fbin (Op ({fb} ?shape ?a_s ?b_s ?o_s ?dt)
+                                   (ICons ?inner_a (ICons ?fs_b (INil)))))
+                    (let ?new_fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fbin (INil))))
+                    (union ?bin ?new_fe)
+                 ) :name \"grow-FE-B-lhs-{lb}\")"
+            )));
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?fe (Op (FusionEnd ?shape ?b_s ?dt) (ICons ?inner_b (INil))))
+                    (= ?bin (Op ({kb} ?shape ?a_s ?b_s ?o_s ?dt)
+                                 (ICons ?a (ICons ?fe (INil)))))
+                 ) (
+                    (let ?fs_a (Op (FusionStart ?shape ?a_s ?dt) (ICons ?a (INil))))
+                    (let ?fbin (Op ({fb} ?shape ?a_s ?b_s ?o_s ?dt)
+                                   (ICons ?fs_a (ICons ?inner_b (INil)))))
+                    (let ?new_fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fbin (INil))))
+                    (union ?bin ?new_fe)
+                 ) :name \"grow-FE-B-rhs-{lb}\")"
+            )));
+        }
+
+        // 7. Merge two FEs at a binary: B(FE(ia), FE(ib)) → FE(FB(ia, ib)).
+        //    Both inners reused, no new FS — shared external tensors with
+        //    upstream FSes stay at one FS.
+        for (kb, fb, lb) in binaries {
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?fe_a (Op (FusionEnd ?shape ?a_s ?dt) (ICons ?inner_a (INil))))
+                    (= ?fe_b (Op (FusionEnd ?shape ?b_s ?dt) (ICons ?inner_b (INil))))
+                    (= ?bin (Op ({kb} ?shape ?a_s ?b_s ?o_s ?dt)
+                                 (ICons ?fe_a (ICons ?fe_b (INil)))))
+                 ) (
+                    (let ?fbin (Op ({fb} ?shape ?a_s ?b_s ?o_s ?dt)
+                                   (ICons ?inner_a (ICons ?inner_b (INil)))))
+                    (let ?new_fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fbin (INil))))
+                    (union ?bin ?new_fe)
+                 ) :name \"merge-FE-FE-{lb}\")"
+            )));
+        }
+
+        // No dissolve rule (`FS(FE(x)) → x`): unioning FS's eclass with FE's
+        // inner eclass creates self-referential eclasses after grow rules
+        // extend the downstream region, and extraction then panics with
+        // `Cycle(NodeIndex(_))`. Grow rules already compose adjacent regions
+        // correctly without dissolve.
+
+        rules
+    }
+
+    fn cleanup(&self) -> bool {
+        false
+    }
+
+    fn extract<'a>(
+        &'a self,
+        egraph: &'a SerializedEGraph,
+        kind_children: &[&'a ENodeId],
+        input_enodes: Vec<&'a ENodeId>,
+        list_cache: &mut FxHashMap<&'a ENodeId, Vec<Expression>>,
+        expr_cache: &mut FxHashMap<&'a ENodeId, Expression>,
+    ) -> (LLIROp, Vec<&'a ENodeId>) {
+        (
+            LLIROp::new::<dyn KernelOp>(Box::new(Self {
+                shape: extract_expr_list(egraph, kind_children[0], list_cache, expr_cache).unwrap(),
+                strides: extract_expr_list(egraph, kind_children[1], list_cache, expr_cache)
+                    .unwrap(),
+                dtype: extract_dtype(egraph, kind_children[2]),
+            })),
+            input_enodes,
+        )
+    }
+}
+
+impl KernelOp for FusionEnd {
+    fn compile(
+        &self,
+        stream: &Arc<CudaStream>,
+        compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    ) -> CompileOut {
+        compile_identity_kernel(
+            stream,
+            compile_cache,
+            "fusion_end_k",
+            &self.shape,
+            &self.strides,
+            self.dtype,
+        )
+    }
+    fn output_size(&self) -> Expression {
+        self.shape.iter().copied().product()
+    }
+    fn output_bytes(&self) -> Expression {
+        (self.output_size() * self.dtype.bits()).ceil_div(8)
+    }
+    fn output_dtype(&self) -> DType {
+        self.dtype
+    }
+    fn kernel_name(&self) -> &'static str {
+        "FusionEnd"
+    }
+}

crates/luminal_cuda_lite/src/kernel/fusion/mod.rs

@@ -0,0 +1,26 @@
+//! Binary-inclusive elementwise kernel fusion.
+//!
+//! - `markers` — `FusionStart` / `FusionEnd` ops + the seven egglog rule
+//!   families that build and extend FE-bracketed regions.
+//! - `fused_ops` — eight `FusedX` op variants (interior to a region) so
+//!   pair-fuse rules' RHS sit in a different egglog sort than their LHS,
+//!   blocking cascade by typing.
+//! - `region_codegen` — `kernel_to_host` calls into here to collapse each
+//!   FE-rooted region into a single CUDA kernel at compile time.
+//!
+//! The LLIR keeps `FusionStart` / `FusedX` / `FusionEnd` nodes after
+//! extraction; `region_codegen` is the only place that walks them.
+
+pub mod fused_ops;
+pub mod markers;
+pub mod region_codegen;
+
+pub use fused_ops::{
+    FusedAdd, FusedExp, FusedExp2, FusedLog2, FusedMul, FusedRecip, FusedSin, FusedSqrt,
+};
+pub use markers::{FusionEnd, FusionStart};
+
+/// All fusion-related op types that the egglog runtime needs to know about
+/// (markers + interior FusedX variants). Combined into a flat tuple for the
+/// `Ops` registry in `kernel::mod`.
+pub type Ops = (markers::Ops, fused_ops::Ops);

crates/luminal_cuda_lite/src/kernel/fusion/region_codegen.rs

@@ -0,0 +1,479 @@
+// =========================================================================
+// Region codegen for FusionStart / FusionEnd-bracketed fused regions.
+//
+// PR1 left FusedX / FusionStart / FusionEnd nodes in the post-extraction
+// LLIR, each compiling to its own standalone CUDA kernel. PR2 collapses
+// every FusionEnd-rooted region into ONE fused CUDA kernel at codegen
+// time — without rewriting the LLIR.
+//
+// Pipeline:
+//   `kernel_to_host` builds a Vec<CompileUnit> from the topo order:
+//     - CompileUnit::Single(node)  — un-fused KernelX, compiled as before.
+//     - CompileUnit::Region(rgn)   — one FE + its interior FusedX DAG +
+//                                    its FS leaves. Compiled here as a
+//                                    single CUDA kernel that reads from
+//                                    the region's external inputs once,
+//                                    chains all FusedX bodies through
+//                                    register-resident locals, and writes
+//                                    the FE's output.
+//
+// The CompiledKernel for a Region is keyed on the FE node and stores
+// `inputs = external producer NodeIndices` (one per interior FusionStart),
+// so the existing buffer-pointer wiring in to_host.rs picks up the right
+// device pointers at execute time. Interior FusedX / FusionStart nodes
+// never enter the kernels Vec — they have no buffers, no launches.
+// =========================================================================
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::{
+    graph::LLIRGraph,
+    prelude::{
+        petgraph::{Direction, algo::toposort, visit::EdgeRef},
+        *,
+    },
+};
+
+use as_any::Downcast;
+
+use crate::{
+    compile_module_image_for_current_device, cuda_dtype,
+    kernel::KernelOp,
+    kernel::fusion::markers::{FusionEnd, FusionStart},
+    kernel::hlir::{dtype_includes, generate_dyn_dims_defines},
+};
+
+// =========================================================================
+// Compile units — what `kernel_to_host` iterates over instead of nodes.
+// =========================================================================
+
+#[derive(Debug, Clone)]
+pub(crate) struct RegionUnit {
+    /// The FusionEnd node that anchors this region.
+    pub fe_node: NodeIndex,
+    /// Interior FusedX nodes, in topological order (predecessors before
+    /// consumers). Used to emit register-binding statements in dependency
+    /// order in the fused CUDA kernel body.
+    pub fusedx_topo: Vec<NodeIndex>,
+    /// FusionStart nodes that bound the region's leaves. One per external
+    /// read site — duplicates (different FS LLIR nodes wrapping the same
+    /// upstream tensor) are kept separate so each read uses its own
+    /// strides; the host launch passes the same device pointer twice.
+    pub fs_nodes: Vec<NodeIndex>,
+    /// External producer NodeIndices, one per `fs_nodes` entry in the same
+    /// order. Becomes the `inputs` field of the FE's `CompiledKernel`, and
+    /// the kernel function's `in0`, `in1`, ... parameters in that order.
+    pub external_inputs: Vec<NodeIndex>,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) enum CompileUnit {
+    Single(NodeIndex),
+    Region(RegionUnit),
+}
+
+// =========================================================================
+// Region detection.
+// =========================================================================
+
+/// Group a sub-DAG's topo order into compile units. Each FusionEnd node
+/// becomes the root of a `CompileUnit::Region`; the region's interior
+/// FusedX and FusionStart nodes are absorbed into that region and removed
+/// from the per-node iteration. Anything else is wrapped in
+/// `CompileUnit::Single`.
+/// Globally-absorbed FS / FE markers — the set of marker nodes that any
+/// `FusionEnd` in the LLIR walks back to during region detection. A
+/// marker is "absorbed" iff some FE in the LLIR can reach it by walking
+/// incoming edges through `FusionEnd` / `FusedX` nodes, stopping at
+/// `FusionStart` leaves.
+///
+/// This is computed once over the full LLIR rather than per-convex-
+/// subgraph, because `partition_marked_convex` may put a shared FS leaf
+/// (one whose e-graph congruence-deduplicated it across multiple
+/// regions) into a different subgraph than the FE that absorbs it.
+/// Without this global view, `build_compile_units` running on the FS's
+/// subgraph would not see any FE walking back to the FS, would emit the
+/// FS as `CompileUnit::Single`, and the markers' identity-memcpy
+/// fallback would compile and launch — pure overhead at runtime.
+pub(crate) fn globally_absorbed_markers(llir_graph: &LLIRGraph) -> FxHashSet<NodeIndex> {
+    let name_of = |idx: NodeIndex| -> Option<&'static str> {
+        llir_graph
+            .node_weight(idx)
+            .and_then(|op| op.to_dialect::<dyn KernelOp>().map(|k| k.kernel_name()))
+    };
+
+    let mut absorbed: FxHashSet<NodeIndex> = FxHashSet::default();
+    for fe in llir_graph.node_indices() {
+        if name_of(fe) != Some("FusionEnd") {
+            continue;
+        }
+        let mut visited: FxHashSet<NodeIndex> = FxHashSet::default();
+        let mut stack: Vec<NodeIndex> = vec![fe];
+        visited.insert(fe);
+        while let Some(cur) = stack.pop() {
+            for pred in llir_graph.neighbors_directed(cur, Direction::Incoming) {
+                if !visited.insert(pred) {
+                    continue;
+                }
+                match name_of(pred) {
+                    Some("FusionStart") => {
+                        absorbed.insert(pred);
+                    }
+                    Some("FusionEnd") => {
+                        absorbed.insert(pred);
+                        stack.push(pred);
+                    }
+                    Some(other) if other.starts_with("Fused") => {
+                        absorbed.insert(pred);
+                        stack.push(pred);
+                    }
+                    _ => {}
+                }
+            }
+        }
+    }
+    absorbed
+}
+
+pub(crate) fn build_compile_units(
+    topo_order: &[NodeIndex],
+    llir_graph: &LLIRGraph,
+    globally_absorbed: &FxHashSet<NodeIndex>,
+) -> Vec<CompileUnit> {
+    let name_of = |idx: NodeIndex| -> Option<&'static str> {
+        llir_graph
+            .node_weight(idx)
+            .and_then(|op| op.to_dialect::<dyn KernelOp>().map(|k| k.kernel_name()))
+    };
+
+    // First pass: every FusionEnd in the subgraph anchors a region; gather
+    // the region's interior + FS leaves by walking incoming edges
+    // backward, stopping at FusionStart (a leaf — its predecessor is the
+    // external producer, outside the region).
+    let mut absorbed: FxHashSet<NodeIndex> = FxHashSet::default();
+    let mut regions: FxHashMap<NodeIndex, RegionUnit> = FxHashMap::default();
+
+    for &node in topo_order {
+        if name_of(node) != Some("FusionEnd") {
+            continue;
+        }
+
+        let mut interior: Vec<NodeIndex> = Vec::new();
+        let mut fs_nodes: Vec<NodeIndex> = Vec::new();
+        let mut visited: FxHashSet<NodeIndex> = FxHashSet::default();
+        let mut stack: Vec<NodeIndex> = Vec::new();
+        stack.push(node);
+        visited.insert(node);
+
+        while let Some(cur) = stack.pop() {
+            for pred in llir_graph.neighbors_directed(cur, Direction::Incoming) {
+                if !visited.insert(pred) {
+                    continue;
+                }
+                match name_of(pred) {
+                    Some("FusionStart") => {
+                        fs_nodes.push(pred);
+                        // Don't recurse past FS — its predecessor is
+                        // external (outside the region).
+                    }
+                    Some("FusionEnd") => {
+                        // A nested FE inside a region. Under the current
+                        // rule design these are cascade artifacts — treat
+                        // them as transparent (walk through) rather than
+                        // as a separate region. The outer region absorbs
+                        // them. They do not become CompileUnit::Region
+                        // anchors because their eclass is already the
+                        // outer region's.
+                        absorbed.insert(pred);
+                        stack.push(pred);
+                    }
+                    Some(other) if other.starts_with("Fused") => {
+                        interior.push(pred);
+                        stack.push(pred);
+                    }
+                    _ => {
+                        // Non-marker, non-FusedX predecessor inside what
+                        // we thought was a region. Shouldn't happen with
+                        // the current rules; treat conservatively: do
+                        // not absorb — let the kernel_to_host single
+                        // path handle it. This means the region is
+                        // malformed and we likely should not have a
+                        // region at all. Caller will see incomplete
+                        // interior; the safer thing is to fall back.
+                    }
+                }
+            }
+        }
+
+        // Topological order on the interior + FS nodes (so the kernel
+        // emits `let v = ...;` lines after their inputs are bound). We
+        // use the parent graph's toposort filtered to in-region nodes.
+        let mut region_set: FxHashSet<NodeIndex> = FxHashSet::default();
+        region_set.extend(interior.iter().copied());
+        region_set.extend(fs_nodes.iter().copied());
+        let topo = toposort(llir_graph, None).expect("LLIR cycle in region detection");
+        let interior_topo: Vec<NodeIndex> = topo
+            .iter()
+            .copied()
+            .filter(|n| region_set.contains(n) && interior.contains(n))
+            .collect();
+        let fs_topo: Vec<NodeIndex> = topo
+            .iter()
+            .copied()
+            .filter(|n| region_set.contains(n) && fs_nodes.contains(n))
+            .collect();
+
+        // External producer for each FS leaf, in the same order.
+        let external_inputs: Vec<NodeIndex> = fs_topo
+            .iter()
+            .map(|&fs| {
+                llir_graph
+                    .neighbors_directed(fs, Direction::Incoming)
+                    .next()
+                    .expect("FusionStart with no predecessor")
+            })
+            .collect();
+
+        absorbed.extend(interior_topo.iter().copied());
+        absorbed.extend(fs_topo.iter().copied());
+
+        regions.insert(
+            node,
+            RegionUnit {
+                fe_node: node,
+                fusedx_topo: interior_topo,
+                fs_nodes: fs_topo,
+                external_inputs,
+            },
+        );
+    }
+
+    // Second pass: emit compile units in original topo order, replacing
+    // FE nodes with their RegionUnit and skipping anything absorbed —
+    // either by a region in *this* subgraph (`absorbed`) or by any
+    // region anywhere in the LLIR (`globally_absorbed`). Skipping the
+    // latter prevents the identity-memcpy fallback from firing on
+    // shared FS markers whose consumers live in other convex subgraphs:
+    // those FSes are absorbed by some other region, and the consuming
+    // region reads from FS's external producer, so the FS never needs
+    // its own kernel.
+    let mut units: Vec<CompileUnit> = Vec::new();
+    for &node in topo_order {
+        if let Some(region) = regions.remove(&node) {
+            units.push(CompileUnit::Region(region));
+        } else if absorbed.contains(&node) || globally_absorbed.contains(&node) {
+            continue;
+        } else {
+            units.push(CompileUnit::Single(node));
+        }
+    }
+    units
+}
+
+// =========================================================================
+// Per-FusedX body templates.
+//
+// Each entry takes the names of the local variables holding the op's
+// inputs and returns a CUDA expression evaluating to the op's output
+// (a register-resident value, no buffer involved).
+// =========================================================================
+
+fn fused_body(name: &str, locals: &[&str]) -> String {
+    match name {
+        "FusedSin" => format!("sinf({})", locals[0]),
+        "FusedSqrt" => format!("sqrtf({})", locals[0]),
+        "FusedExp" => format!("expf({})", locals[0]),
+        "FusedExp2" => format!("exp2f({})", locals[0]),
+        "FusedLog2" => format!("log2f({})", locals[0]),
+        "FusedRecip" => format!("1.0f / {}", locals[0]),
+        "FusedAdd" => format!("{} + {}", locals[0], locals[1]),
+        "FusedMul" => format!("{} * {}", locals[0], locals[1]),
+        other => panic!("region_codegen: unknown FusedX op {other}"),
+    }
+}
+
+// =========================================================================
+// Region compilation — emit one CUDA kernel for the whole region.
+// =========================================================================
+
+#[allow(clippy::type_complexity)]
+pub(crate) struct CompiledRegion {
+    pub function: CudaFunction,
+    pub module: Arc<CudaModule>,
+    pub kernel_str: String,
+    pub grid: (Expression, Expression, Expression),
+    pub block: (Expression, Expression, Expression),
+    pub shared_mem: Expression,
+    pub constants: FxHashMap<char, CudaSlice<u8>>,
+}
+
+#[allow(clippy::type_complexity)]
+pub(crate) fn compile_region(
+    region: &RegionUnit,
+    llir_graph: &LLIRGraph,
+    stream: &Arc<CudaStream>,
+    compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+) -> CompiledRegion {
+    // Resolve FE: shape, strides (for the write), dtype.
+    let fe_op = llir_graph[region.fe_node]
+        .to_dialect::<dyn KernelOp>()
+        .expect("FE node must be a KernelOp");
+    let fe_struct: &FusionEnd = (***fe_op)
+        .downcast_ref::<FusionEnd>()
+        .expect("region root must be FusionEnd");
+    let out_shape: &[Expression] = &fe_struct.shape;
+    let out_strides: &[Expression] = &fe_struct.strides;
+    let dtype: DType = fe_struct.dtype;
+
+    // Aggregate all dynamic vars used anywhere in the region (FS strides,
+    // FE strides, FusedX shape — all FusedX share `out_shape`, but their
+    // own strides are likewise relevant for any future stride-affine ops).
+    let mut all_vars: FxHashSet<char> = FxHashSet::default();
+    all_vars.extend(out_shape.iter().flat_map(|e| e.dyn_vars()));
+    all_vars.extend(out_strides.iter().flat_map(|e| e.dyn_vars()));
+    for &fs_idx in &region.fs_nodes {
+        let fs_op = llir_graph[fs_idx].to_dialect::<dyn KernelOp>().unwrap();
+        let fs_struct: &FusionStart = (***fs_op).downcast_ref::<FusionStart>().unwrap();
+        all_vars.extend(fs_struct.strides.iter().flat_map(|e| e.dyn_vars()));
+    }
+
+    let cuda_ty = cuda_dtype(dtype);
+    let includes = dtype_includes(&[dtype]);
+    let (dyn_defines, _sorted_dims) = generate_dyn_dims_defines(&all_vars);
+    let dyn_dims_param = if all_vars.is_empty() {
+        ""
+    } else {
+        ", const int* dyn_dims"
+    };
+
+    let n_elements = out_shape
+        .iter()
+        .copied()
+        .product::<Expression>()
+        .to_kernel();
+
+    // Build kernel signature: out, then one input per FS leaf in
+    // `region.fs_nodes` order. The `external_inputs` list (parallel to
+    // `fs_nodes`) is what the host wires into the launch params.
+    let mut signature_params: Vec<String> = vec![format!("{cuda_ty} *out")];
+    for i in 0..region.fs_nodes.len() {
+        signature_params.push(format!("const {cuda_ty} *in{i}"));
+    }
+    let signature = signature_params.join(", ");
+
+    // Body: read FS leaves, then walk FusedX in topo order emitting a
+    // local per op, then write FE output. Every node gets a local keyed
+    // by a position-in-region index so the kernel string is invariant
+    // under NodeIndex churn (each `egglog_to_llir` reissues NodeIndexes,
+    // so naming locals by `n.index()` would invalidate the kernel
+    // string cache on every search candidate). Indices: FS leaves get
+    // 0..fs_nodes.len(), FusedX get fs_nodes.len()..(+ fusedx_topo.len()).
+    let mut local_idx_map: FxHashMap<NodeIndex, usize> = FxHashMap::default();
+    for (i, &fs_idx) in region.fs_nodes.iter().enumerate() {
+        local_idx_map.insert(fs_idx, i);
+    }
+    let fs_count = region.fs_nodes.len();
+    for (i, &op_idx) in region.fusedx_topo.iter().enumerate() {
+        local_idx_map.insert(op_idx, fs_count + i);
+    }
+    let local_name = |n: NodeIndex| format!("v_{}", local_idx_map[&n]);
+
+    let mut body = String::new();
+    body.push_str(&format!(
+        "        long long const_z = (long long)blockIdx.x * blockDim.x + threadIdx.x;\n\
+         \x20       if (const_z >= {n_elements}) return;\n"
+    ));
+
+    // FS leaves: each reads from its corresponding `in_i` parameter using
+    // its own strides.
+    for (i, &fs_idx) in region.fs_nodes.iter().enumerate() {
+        let fs_op = llir_graph[fs_idx].to_dialect::<dyn KernelOp>().unwrap();
+        let fs_struct: &FusionStart = (***fs_op).downcast_ref::<FusionStart>().unwrap();
+        let read_idx = flatten_strides(out_shape, &fs_struct.strides).to_kernel();
+        body.push_str(&format!(
+            "        {cuda_ty} {name} = in{i}[{read_idx}];\n",
+            name = local_name(fs_idx),
+        ));
+    }
+
+    // FusedX ops in topo order. Each looks up its predecessor locals
+    // (in incoming-edge id order to match the original op's input
+    // arity / position).
+    for &op_idx in &region.fusedx_topo {
+        let op_ref = llir_graph[op_idx].to_dialect::<dyn KernelOp>().unwrap();
+        let op_name = op_ref.kernel_name();
+
+        let mut input_locals: Vec<String> = llir_graph
+            .edges_directed(op_idx, Direction::Incoming)
+            .map(|e| (e.id(), e.source()))
+            .collect::<Vec<_>>()
+            .into_iter()
+            .map(|(_, src)| local_name(src))
+            .collect();
+        // Sort by edge id like the rest of the codegen does for stable
+        // input ordering.
+        let mut edges: Vec<(_, NodeIndex)> = llir_graph
+            .edges_directed(op_idx, Direction::Incoming)
+            .map(|e| (e.id(), e.source()))
+            .collect();
+        edges.sort_by_key(|(eid, _)| *eid);
+        input_locals = edges.into_iter().map(|(_, src)| local_name(src)).collect();
+        let inputs_ref: Vec<&str> = input_locals.iter().map(|s| s.as_str()).collect();
+
+        let expr = fused_body(op_name, &inputs_ref);
+        body.push_str(&format!(
+            "        {cuda_ty} {name} = {expr};\n",
+            name = local_name(op_idx),
+        ));
+    }
+
+    // FE write: pick the FusedX feeding FE (its single incoming edge in
+    // the region — a FusedX or, in degenerate single-FS regions which
+    // shouldn't arise, an FS).
+    let fe_input: NodeIndex = llir_graph
+        .neighbors_directed(region.fe_node, Direction::Incoming)
+        .next()
+        .expect("FusionEnd with no predecessor");
+    let fe_input_local = local_name(fe_input);
+    let write_idx = flatten_strides(out_shape, out_strides).to_kernel();
+    body.push_str(&format!("        out[{write_idx}] = {fe_input_local};\n"));
+
+    let kernel = format!(
+        "{includes}\n\
+         {dyn_defines}\n\
+         extern \"C\" {{\n\
+         \x20   __global__ void fused_region_k({signature}{dyn_dims_param}) {{\n\
+         {body}\
+         \x20   }}\n\
+         }}"
+    );
+
+    let (module, function) = if let Some((m, f)) = compile_cache.get(&kernel) {
+        (m.clone(), f.clone())
+    } else {
+        let ptx = compile_module_image_for_current_device(stream.context(), &kernel)
+            .expect("region kernel PTX compile failed");
+        let module = stream
+            .context()
+            .load_module(ptx)
+            .expect("module load failed");
+        let function = module
+            .load_function("fused_region_k")
+            .expect("region kernel function not found");
+        compile_cache.insert(kernel.clone(), (module.clone(), function.clone()));
+        (module, function)
+    };
+
+    let out_size = out_shape.iter().copied().product::<Expression>();
+
+    CompiledRegion {
+        function,
+        module,
+        kernel_str: kernel,
+        grid: (out_size.ceil_div(256), 1.into(), 1.into()),
+        block: (out_size.min(256), 1.into(), 1.into()),
+        shared_mem: 0.into(),
+        constants: FxHashMap::default(),
+    }
+}

crates/luminal_cuda_lite/src/kernel/hlir.rs

@@ -69,7 +69,7 @@ pub type Ops = (
 
 /// Build a rewrite that matches an HLIR op, reads dtype(s) from the given source fields,
 /// and unions with a kernel op that has the same fields plus the dtype(s) appended.
-fn kernel_rewrite<H: Default + EgglogOp, L: Default + EgglogOp>() -> Rule {
+pub fn kernel_rewrite<H: Default + EgglogOp, L: Default + EgglogOp>() -> Rule {
     let hlir = H::default().sort();
     let llir = L::default().sort();
     let (mut args, hlir_kind_term) = hlir.new_call();
@@ -415,8 +415,12 @@ extern \"C\" {{
         long long iters = {iters};
 
         {dtype} partial = 0;
+        {dtype} comp = 0;   // Kahan compensation
         for (long long i = tid; i < iters; i += THREADS_PER_BLOCK) {{
-            partial += in_data[in_start + {iter_stride_of_i}];
+            {dtype} y = in_data[in_start + {iter_stride_of_i}] - comp;
+            {dtype} t = partial + y;
+            comp = (t - partial) - y;
+            partial = t;
         }}
 
         #pragma unroll
@@ -630,8 +634,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(), // No per-module constants needed
         )
@@ -793,8 +797,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )
@@ -986,12 +990,13 @@ extern \"C\" {{
             compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
             (module, func)
         };
+        let out_size = self.out_shape.iter().copied().product::<Expression>();
         (
             func,
             module,
             kernel,
-            (self.out_shape.iter().copied().product(), 1.into(), 1.into()),
-            (1.into(), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )
@@ -1195,7 +1200,25 @@ impl KernelOp for KernelScatter {
 
         // Single-kernel scatter: copy dest→output then scatter src→output[indexes]
         // Launched as 1 block of 1024 threads with __syncthreads() barrier.
-        // Uses float4 vectorized copy (4x throughput) for the copy phase.
+        // Uses float4 vectorized copy (16 bytes per op) for the copy phase.
+        //
+        // The number of dtype elements that fit in a float4 (16 bytes) depends
+        // on the element size. Computing `n_vec = n_dest / 4` would only be
+        // correct for 4-byte dtypes — for bf16 it walks 2× past the end of
+        // `out`, producing CUDA_ERROR_ILLEGAL_ADDRESS once the OOB region
+        // happens to land on an unmapped page.
+        let elements_per_vec: usize = match self.dtype {
+            DType::F64 => 2,
+            DType::F32 | DType::Int => 4,
+            DType::F16 | DType::Bf16 | DType::I16 | DType::U16 => 8,
+            DType::Bool
+            | DType::I8
+            | DType::U8
+            | DType::F8UE8M0
+            | DType::F8E4M3
+            | DType::F8E5M2 => 16,
+            other => panic!("Unsupported dtype for scatter vectorization: {other:?}"),
+        };
         let n_src_elements = self
             .index_shape
             .iter()
@@ -1220,15 +1243,17 @@ extern \"C\" {{
         int tid = threadIdx.x;
         long long n_dest = {n_dest_elements};
         long long n_src = {n_src_elements};
-        // Phase 1: vectorized copy dest → output (float4 = 4 elements per op)
-        long long n_vec = n_dest / 4;
+        // Phase 1: vectorized copy dest → output (float4 = 16 bytes / iter,
+        // i.e. {elements_per_vec} {dtype} elements). n_vec is sized so the
+        // total bytes covered (`n_vec * 16`) never exceed `n_dest * sizeof({dtype})`.
+        long long n_vec = n_dest / {elements_per_vec};
         float4 *out4 = (float4 *)out;
         const float4 *dest4 = (const float4 *)dest;
         for (long long i = tid; i < n_vec; i += blockDim.x) {{
             out4[i] = dest4[i];
         }}
-        // Handle remaining elements
-        long long remainder_start = n_vec * 4;
+        // Handle remaining elements (the dtype-tail past the last full float4).
+        long long remainder_start = n_vec * {elements_per_vec};
         for (long long i = remainder_start + tid; i < n_dest; i += blockDim.x) {{
             out[i] = dest[i];
         }}
@@ -1611,8 +1636,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )
@@ -1765,8 +1790,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )
@@ -1919,8 +1944,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )
@@ -2055,7 +2080,7 @@ extern \"C\" {{
     __global__ void recip_k({dtype} *out, const {dtype} *in{dyn_dims_param}) {{
         long long const_z = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         if (const_z >= {n_elements}) return;
-        out[{out_idx}] = 1.0f / in[{in_idx}];
+        out[{out_idx}] = ({dtype})1.0f / in[{in_idx}];
     }}
 }}"
         );
@@ -2073,8 +2098,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )
@@ -2227,8 +2252,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )
@@ -2388,8 +2413,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )
@@ -2563,8 +2588,8 @@ extern \"C\" {{
             func,
             module,
             kernel,
-            (out_size.ceil_div(128), 1.into(), 1.into()),
-            (out_size.min(128), 1.into(), 1.into()),
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
             0.into(),
             FxHashMap::default(),
         )

crates/luminal_cuda_lite/src/kernel/mod.rs

@@ -10,12 +10,13 @@ use luminal_tracing::schema::{
 use uuid::Uuid;
 
 pub mod cuda_graph;
+pub mod fusion;
 pub mod hlir;
 pub mod other_ops;
 
 pub use cuda_graph::*;
 
-pub type Ops = (hlir::Ops, other_ops::Ops);
+pub type Ops = (hlir::Ops, other_ops::Ops, fusion::Ops);
 
 /// Build a mapping from interned string IDs to their string values for a given sequence.
 fn build_interned_strings(trace: &schema::Trace) -> std::collections::HashMap<(u32, u64), String> {

crates/luminal_cuda_lite/src/kernel/other_ops.rs

@@ -3,14 +3,14 @@ use std::sync::Arc;
 use crate::{
     compile_module_image_for_current_device, cuda_dtype,
     kernel::KernelOp,
-    kernel::hlir::{dtype_includes, generate_dyn_dims_defines},
+    kernel::hlir::{dtype_includes, generate_dyn_dims_defines, kernel_rewrite},
 };
 use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
 use itertools::Itertools;
 use luminal::{
     egglog_utils::{
         api::{Rule, SortDef, sort},
-        base::{DTYPE, ELIST, EXPRESSION, OP_KIND},
+        base::{DTYPE, ELIST, EXPRESSION, OP_KIND, STRING},
         extract_dtype, extract_expr, extract_expr_list,
     },
     op::*,
@@ -22,6 +22,9 @@ pub type Ops = (
     KernelBatchMatVec,
     KernelBatchMatMul,
     KernelScatterNoCopy,
+    KernelSoftmax,
+    KernelExp,
+    KernelSigmoid,
 );
 
 #[derive(Default, Debug, Clone)]
@@ -1151,6 +1154,7 @@ impl EgglogOp for KernelSoftmax {
                 ("out_strides", ELIST),
                 ("reduce_dim", EXPRESSION),
                 ("reduce_stride", EXPRESSION),
+                ("dtype", DTYPE),
             ],
         )
     }
@@ -1160,8 +1164,24 @@ impl EgglogOp for KernelSoftmax {
     }
 
     fn rewrites(&self) -> Vec<Rule> {
-        // No rewrite rules yet - this op is not in the Ops tuple.
-        vec![]
+        vec![
+            kernel_rewrite::<luminal::hlir::Softmax, Self>(),
+            // Also add a direct rewrite that assumes F32 dtype, in case dtype
+            // propagation hasn't reached the Softmax node yet.
+            Rule::raw(
+                "(rule
+                (
+                    (= ?sm (Op (Softmax ?shape ?in_strides ?out_strides ?reduce_dim ?reduce_stride) ?inputs))
+                )
+                (
+                    (let ?ksm (Op (KernelSoftmax ?shape ?in_strides ?out_strides ?reduce_dim ?reduce_stride (F32)) ?inputs))
+                    (union ?sm ?ksm)
+                    (set (dtype ?ksm) (F32))
+                )
+                :name \"softmax-to-kernel-f32\"
+            )",
+            ),
+        ]
     }
 
     fn cleanup(&self) -> bool {
@@ -1176,16 +1196,21 @@ impl EgglogOp for KernelSoftmax {
         list_cache: &mut FxHashMap<&'a ENodeId, Vec<Expression>>,
         expr_cache: &mut FxHashMap<&'a ENodeId, Expression>,
     ) -> (LLIROp, Vec<&'a ENodeId>) {
+        let out_shape =
+            extract_expr_list(egraph, kind_children[0], list_cache, expr_cache).unwrap();
+        let in_stride =
+            extract_expr_list(egraph, kind_children[1], list_cache, expr_cache).unwrap();
+        let out_stride =
+            extract_expr_list(egraph, kind_children[2], list_cache, expr_cache).unwrap();
+        let reduce_dim = extract_expr(egraph, kind_children[3], expr_cache).unwrap();
+        let reduce_stride = extract_expr(egraph, kind_children[4], expr_cache).unwrap();
         (
             LLIROp::new::<dyn KernelOp>(Box::new(Self {
-                out_shape: extract_expr_list(egraph, kind_children[0], list_cache, expr_cache)
-                    .unwrap(),
-                in_stride: extract_expr_list(egraph, kind_children[1], list_cache, expr_cache)
-                    .unwrap(),
-                out_stride: extract_expr_list(egraph, kind_children[2], list_cache, expr_cache)
-                    .unwrap(),
-                reduce_dim: extract_expr(egraph, kind_children[3], expr_cache).unwrap(),
-                reduce_stride: extract_expr(egraph, kind_children[4], expr_cache).unwrap(),
+                out_shape,
+                in_stride,
+                out_stride,
+                reduce_dim,
+                reduce_stride,
             })),
             input_enodes,
         )
@@ -1374,3 +1399,370 @@ extern \"C\" {{
         "Softmax"
     }
 }
+
+// KernelExp: native exp (uses expf instead of exp2f * constant)
+// Single-kernel alternative to the 3-kernel Constant+Mul+Exp2 path.
+// Improves numerical precision by avoiding the truncated log2(e) constant.
+
+#[derive(Default, Debug, Clone)]
+pub struct KernelExp {
+    shape: Vec<Expression>,
+    in_strides: Vec<Expression>,
+    out_strides: Vec<Expression>,
+    dtype: DType,
+}
+
+impl EgglogOp for KernelExp {
+    fn sort(&self) -> SortDef {
+        sort(
+            OP_KIND,
+            "KernelExp",
+            &[
+                ("shape", ELIST),
+                ("strides", ELIST),
+                ("out_strides", ELIST),
+                ("dtype", DTYPE),
+            ],
+        )
+    }
+
+    fn n_inputs(&self) -> usize {
+        1
+    }
+
+    fn rewrites(&self) -> Vec<Rule> {
+        vec![
+            // Match Exp2(Mul(x, log2e_constant)) directly.
+            // This matches the pattern created by frontend exp() = (self * (1/ln(2))).exp2()
+            Rule::raw(
+                "(rule
+                (
+                    (= ?mul (Op (Mul ?shape ?x_stride ?const_stride ?inter_stride) (ICons ?x (ICons ?exp_const (INil)))))
+                    (= ?exp2 (Op (Exp2 ?shape ?inter_stride ?out_stride) (ICons ?mul (INil))))
+                    (= ?dt (dtype ?x))
+                    (= ?cv (Op (Constant ?val) (INil)))
+                    (= ?exp_const ?cv)
+                    (> ?val 1.44)
+                    (< ?val 1.45)
+                )
+                (
+                    (let ?kexp (Op (KernelExp ?shape ?x_stride ?out_stride ?dt) (ICons ?x (INil))))
+                    (union ?exp2 ?kexp)
+                    (set (dtype ?kexp) ?dt)
+                )
+                :name \"direct-exp-fusion\"
+            )",
+            ),
+        ]
+    }
+
+    fn cleanup(&self) -> bool {
+        false
+    }
+
+    fn extract<'a>(
+        &'a self,
+        egraph: &'a SerializedEGraph,
+        kind_children: &[&'a ENodeId],
+        input_enodes: Vec<&'a ENodeId>,
+        list_cache: &mut FxHashMap<&'a ENodeId, Vec<Expression>>,
+        expr_cache: &mut FxHashMap<&'a ENodeId, Expression>,
+    ) -> (LLIROp, Vec<&'a ENodeId>) {
+        (
+            LLIROp::new::<dyn KernelOp>(Box::new(Self {
+                shape: extract_expr_list(egraph, kind_children[0], list_cache, expr_cache).unwrap(),
+                in_strides: extract_expr_list(egraph, kind_children[1], list_cache, expr_cache)
+                    .unwrap(),
+                out_strides: extract_expr_list(egraph, kind_children[2], list_cache, expr_cache)
+                    .unwrap(),
+                dtype: extract_dtype(egraph, kind_children[3]),
+            })),
+            input_enodes,
+        )
+    }
+}
+
+impl KernelOp for KernelExp {
+    fn compile(
+        &self,
+        stream: &Arc<CudaStream>,
+        compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    ) -> (
+        CudaFunction,
+        Arc<CudaModule>,
+        String,
+        (Expression, Expression, Expression),
+        (Expression, Expression, Expression),
+        Expression,
+        FxHashMap<char, CudaSlice<u8>>,
+    ) {
+        let vars = self
+            .shape
+            .iter()
+            .flat_map(|e| e.dyn_vars())
+            .chain(self.in_strides.iter().flat_map(|e| e.dyn_vars()))
+            .chain(self.out_strides.iter().flat_map(|e| e.dyn_vars()))
+            .collect::<FxHashSet<_>>();
+        let dtype = cuda_dtype(self.dtype);
+        let includes = dtype_includes(&[self.dtype]);
+        let (dyn_defines, _sorted_dims) = generate_dyn_dims_defines(&vars);
+        let dyn_dims_param = if vars.is_empty() {
+            ""
+        } else {
+            ", const int* dyn_dims"
+        };
+        let n_elements = self
+            .shape
+            .iter()
+            .copied()
+            .product::<Expression>()
+            .to_kernel();
+        let out_idx = flatten_strides(&self.shape, &self.out_strides).to_kernel();
+        let in_idx = flatten_strides(&self.shape, &self.in_strides).to_kernel();
+        let kernel = format!(
+            "{includes}
+{dyn_defines}
+extern \"C\" {{
+    __global__ void exp_k({dtype} *out, const {dtype} *in{dyn_dims_param}) {{
+        long long const_z = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+        if (const_z >= {n_elements}) return;
+        out[{out_idx}] = expf(in[{in_idx}]);
+    }}
+}}"
+        );
+        let (module, func) = if let Some((module, func)) = compile_cache.get(&kernel) {
+            (module.clone(), func.clone())
+        } else {
+            let ptx = compile_module_image_for_current_device(stream.context(), &kernel).unwrap();
+            let module = stream.context().load_module(ptx).unwrap();
+            let func = module.load_function("exp_k").unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+        let out_size = self.shape.iter().copied().product::<Expression>();
+        (
+            func,
+            module,
+            kernel,
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
+            0.into(),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        self.shape.iter().copied().product()
+    }
+
+    fn output_bytes(&self) -> Expression {
+        (self.output_size() * self.dtype.bits()).ceil_div(8)
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        self.output_bytes()
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_bytes()
+    }
+
+    fn flops(&self) -> Expression {
+        self.shape.iter().copied().product()
+    }
+
+    fn output_dtype(&self) -> DType {
+        self.dtype
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "Exp"
+    }
+}
+
+// KernelSigmoid: fused sigmoid = 1/(1+exp(-x))
+// Single-kernel alternative to the 5-kernel Neg+Exp+Const+Add+Recip path.
+
+#[derive(Default, Debug, Clone)]
+pub struct KernelSigmoid {
+    shape: Vec<Expression>,
+    in_strides: Vec<Expression>,
+    out_strides: Vec<Expression>,
+    dtype: DType,
+}
+
+impl EgglogOp for KernelSigmoid {
+    fn sort(&self) -> SortDef {
+        sort(
+            OP_KIND,
+            "KernelSigmoid",
+            &[
+                ("shape", ELIST),
+                ("strides", ELIST),
+                ("out_strides", ELIST),
+                ("dtype", DTYPE),
+            ],
+        )
+    }
+
+    fn n_inputs(&self) -> usize {
+        1
+    }
+
+    fn rewrites(&self) -> Vec<Rule> {
+        vec![
+            // Match the HLIR pattern directly: Recip(Add(Exp2(Mul(Mul(x, -1), log2e)), 1))
+            Rule::raw(
+                "(rule
+                (
+                    (= ?neg1 (Op (Constant ?nv) (INil)))
+                    (< ?nv -0.99)
+                    (> ?nv -1.01)
+                    (= ?neg_x (Op (Mul ?shape ?x_stride ?neg_stride ?neg_out_stride) (ICons ?x (ICons ?neg1 (INil)))))
+                    (= ?log2e (Op (Constant ?lv) (INil)))
+                    (> ?lv 1.44)
+                    (< ?lv 1.45)
+                    (= ?scaled (Op (Mul ?shape ?neg_out_stride ?log2e_stride ?scaled_stride) (ICons ?neg_x (ICons ?log2e (INil)))))
+                    (= ?exp2 (Op (Exp2 ?shape ?scaled_stride ?exp_stride) (ICons ?scaled (INil))))
+                    (= ?one (Op (Constant ?ov) (INil)))
+                    (> ?ov 0.99)
+                    (< ?ov 1.01)
+                    (= ?plus_one (Op (Add ?shape ?exp_stride ?one_stride ?add_stride) (ICons ?exp2 (ICons ?one (INil)))))
+                    (= ?sig_out (Op (Recip ?shape ?add_stride ?out_stride) (ICons ?plus_one (INil))))
+                    (= ?dt (dtype ?x))
+                )
+                (
+                    (let ?ksig (Op (KernelSigmoid ?shape ?x_stride ?out_stride ?dt) (ICons ?x (INil))))
+                    (union ?sig_out ?ksig)
+                    (set (dtype ?ksig) ?dt)
+                )
+                :name \"direct-sigmoid-fusion\"
+            )",
+            ),
+        ]
+    }
+
+    fn cleanup(&self) -> bool {
+        false
+    }
+
+    fn extract<'a>(
+        &'a self,
+        egraph: &'a SerializedEGraph,
+        kind_children: &[&'a ENodeId],
+        input_enodes: Vec<&'a ENodeId>,
+        list_cache: &mut FxHashMap<&'a ENodeId, Vec<Expression>>,
+        expr_cache: &mut FxHashMap<&'a ENodeId, Expression>,
+    ) -> (LLIROp, Vec<&'a ENodeId>) {
+        (
+            LLIROp::new::<dyn KernelOp>(Box::new(Self {
+                shape: extract_expr_list(egraph, kind_children[0], list_cache, expr_cache).unwrap(),
+                in_strides: extract_expr_list(egraph, kind_children[1], list_cache, expr_cache)
+                    .unwrap(),
+                out_strides: extract_expr_list(egraph, kind_children[2], list_cache, expr_cache)
+                    .unwrap(),
+                dtype: extract_dtype(egraph, kind_children[3]),
+            })),
+            input_enodes,
+        )
+    }
+}
+
+impl KernelOp for KernelSigmoid {
+    fn compile(
+        &self,
+        stream: &Arc<CudaStream>,
+        compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    ) -> (
+        CudaFunction,
+        Arc<CudaModule>,
+        String,
+        (Expression, Expression, Expression),
+        (Expression, Expression, Expression),
+        Expression,
+        FxHashMap<char, CudaSlice<u8>>,
+    ) {
+        let vars = self
+            .shape
+            .iter()
+            .flat_map(|e| e.dyn_vars())
+            .chain(self.in_strides.iter().flat_map(|e| e.dyn_vars()))
+            .chain(self.out_strides.iter().flat_map(|e| e.dyn_vars()))
+            .collect::<FxHashSet<_>>();
+        let dtype = cuda_dtype(self.dtype);
+        let includes = dtype_includes(&[self.dtype]);
+        let (dyn_defines, _sorted_dims) = generate_dyn_dims_defines(&vars);
+        let dyn_dims_param = if vars.is_empty() {
+            ""
+        } else {
+            ", const int* dyn_dims"
+        };
+        let n_elements = self
+            .shape
+            .iter()
+            .copied()
+            .product::<Expression>()
+            .to_kernel();
+        let out_idx = flatten_strides(&self.shape, &self.out_strides).to_kernel();
+        let in_idx = flatten_strides(&self.shape, &self.in_strides).to_kernel();
+        let kernel = format!(
+            "{includes}
+{dyn_defines}
+extern \"C\" {{
+    __global__ void sigmoid_k({dtype} *out, const {dtype} *in{dyn_dims_param}) {{
+        long long const_z = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+        if (const_z >= {n_elements}) return;
+        out[{out_idx}] = 1.0f / (1.0f + expf(-in[{in_idx}]));
+    }}
+}}"
+        );
+        let (module, func) = if let Some((module, func)) = compile_cache.get(&kernel) {
+            (module.clone(), func.clone())
+        } else {
+            let ptx = compile_module_image_for_current_device(stream.context(), &kernel).unwrap();
+            let module = stream.context().load_module(ptx).unwrap();
+            let func = module.load_function("sigmoid_k").unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+        let out_size = self.shape.iter().copied().product::<Expression>();
+        (
+            func,
+            module,
+            kernel,
+            (out_size.ceil_div(256), 1.into(), 1.into()),
+            (out_size.min(256), 1.into(), 1.into()),
+            0.into(),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        self.shape.iter().copied().product()
+    }
+
+    fn output_bytes(&self) -> Expression {
+        (self.output_size() * self.dtype.bits()).ceil_div(8)
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        self.output_bytes()
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_bytes()
+    }
+
+    fn flops(&self) -> Expression {
+        // neg + exp + add + recip = ~4 ops per element
+        self.shape.iter().copied().product::<Expression>() * 4
+    }
+
+    fn output_dtype(&self) -> DType {
+        self.dtype
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "Sigmoid"
+    }
+}

crates/luminal_cuda_lite/src/kernel/to_host.rs

@@ -7,7 +7,8 @@ use std::cell::RefCell;
 use std::sync::Arc;
 
 use cudarc::driver::{
-    CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, sys::CUgraphNode,
+    CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr,
+    sys::{CUgraphNode, CUresult, cuLaunchKernel},
 };
 use itertools::Itertools;
 use luminal::{
@@ -26,6 +27,7 @@ use crate::{
     kernel::{
         CudaFunctionExt, CudaGraphExecHandle, CudaGraphHandle, KernelOp, create_cuda_event,
         destroy_cuda_event,
+        fusion::region_codegen::{self, CompileUnit},
         hlir::{clear_global_dyn_dims, get_global_dyn_dims, set_global_dyn_dims},
     },
     runtime::partition_marked_convex,
@@ -274,6 +276,14 @@ impl CudaGraphOp {
         buffers: &FxHashMap<NodeIndex, &CudaSlice<u8>>,
         dyn_map: &FxHashMap<char, usize>,
     ) -> anyhow::Result<()> {
+        // Debug path: launch each kernel sequentially with sync between, so the
+        // failing kernel surfaces instead of the generic "CudaGraph" panic.
+        // Enable via `LUMINAL_DEBUG_SEQ=1`. Slow — only for diagnosing
+        // CUDA_ERROR_ILLEGAL_ADDRESS / NaN / wrong-output bugs in graph batching.
+        if std::env::var("LUMINAL_DEBUG_SEQ").is_ok() {
+            return self.execute_sequential_for_debug(stream, buffers, dyn_map);
+        }
+
         let mut state = self.state.borrow_mut();
         let _span = span!(Level::TRACE, "cuda_graph", kernels = state.kernels.len()).entered();
 
@@ -302,8 +312,10 @@ impl CudaGraphOp {
                 kernel.internal_bufs = kernel.kernel_op.allocate_internal_buffers(stream, dyn_map);
             }
         }
-        // Force full rebuild when dims change (debug: testing if update_kernel_node is the issue)
-        if dyn_map_changed || needs_internal_realloc {
+        // Only force full rebuild when internal buffer sizes change.
+        // Dim-only changes (e.g. position offset `p` incrementing each decode step) are
+        // handled by updating the dyn_dims device buffer + kernel node params in-place.
+        if needs_internal_realloc {
             state.cuda_graph = None;
             state.cuda_graph_exec = None;
             state.node_to_graph_node.clear();
@@ -444,6 +456,152 @@ impl CudaGraphOp {
         Ok(())
     }
 
+    /// Diagnostic path for kernel-level errors that surface as a generic
+    /// `CUDA_ERROR_ILLEGAL_ADDRESS` panic from the batched cuda_graph_exec
+    /// launch. Bypasses CUDA-graph batching entirely: builds params per
+    /// kernel and launches each via `cuLaunchKernel`, syncing afterwards so
+    /// the offending kernel reports itself instead of being hidden inside
+    /// the graph's atomic launch.
+    ///
+    /// Enabled via `LUMINAL_DEBUG_SEQ=1`. ~10–100× slower than the graph
+    /// path; not for production.
+    fn execute_sequential_for_debug(
+        &self,
+        stream: &Arc<CudaStream>,
+        buffers: &FxHashMap<NodeIndex, &CudaSlice<u8>>,
+        dyn_map: &FxHashMap<char, usize>,
+    ) -> anyhow::Result<()> {
+        let mut state = self.state.borrow_mut();
+        let num_kernels = state.kernels.len();
+
+        // Allocate dyn_dims_buffer if needed and copy current values.
+        if !self.dyn_dims_order.is_empty() && state.dyn_dims_buffer.is_none() {
+            state.dyn_dims_buffer = Some(stream.alloc_zeros::<i32>(self.dyn_dims_order.len())?);
+        }
+        if !self.dyn_dims_order.is_empty() {
+            let values: Vec<i32> = self
+                .dyn_dims_order
+                .iter()
+                .map(|d| dyn_map.get(d).copied().unwrap_or(0) as i32)
+                .collect();
+            if let Some(buf) = state.dyn_dims_buffer.as_mut() {
+                stream.memcpy_htod(&values, buf)?;
+            }
+        }
+        let dyn_dims_ptr = state
+            .dyn_dims_buffer
+            .as_ref()
+            .map(|buf| buf.device_ptr(stream).0)
+            .unwrap_or(0);
+
+        // Collect buffer pointers (mirrors the graph path).
+        let mut buffer_ptrs: FxHashMap<NodeIndex, u64> = FxHashMap::default();
+        for &node in &self.buffer_nodes {
+            if let Some(buf) = buffers.get(&node) {
+                buffer_ptrs.insert(node, buf.device_ptr(stream).0);
+            }
+        }
+        for kernel in state.kernels.iter() {
+            if let Some(input_idx) = kernel.kernel_op.output_aliases_input()
+                && let Some(&input_ptr) = buffer_ptrs.get(&kernel.inputs[input_idx])
+            {
+                buffer_ptrs.insert(kernel.node, input_ptr);
+            }
+        }
+
+        // Allocate internal buffers + run pre_execute for every kernel up front.
+        for idx in 0..num_kernels {
+            let kernel = &mut state.kernels[idx];
+            if kernel.internal_bufs.is_empty() {
+                kernel.internal_bufs = kernel.kernel_op.allocate_internal_buffers(stream, dyn_map);
+            }
+            kernel.kernel_op.pre_execute(
+                stream,
+                &mut kernel.internal_bufs,
+                &mut kernel.constants,
+                &buffer_ptrs,
+                dyn_map,
+            );
+        }
+
+        let cu_stream = stream.cu_stream();
+
+        for idx in 0..num_kernels {
+            let kernel = &state.kernels[idx];
+            let kernel_name = kernel.kernel_op.kernel_name();
+            let node = kernel.node;
+
+            let grid = (
+                kernel.grid.0.exec(dyn_map).unwrap() as u32,
+                kernel.grid.1.exec(dyn_map).unwrap() as u32,
+                kernel.grid.2.exec(dyn_map).unwrap() as u32,
+            );
+            let block = (
+                kernel.block.0.exec(dyn_map).unwrap() as u32,
+                kernel.block.1.exec(dyn_map).unwrap() as u32,
+                kernel.block.2.exec(dyn_map).unwrap() as u32,
+            );
+            let shared_mem = kernel.shared_mem.exec(dyn_map).unwrap() as u32;
+
+            let output_ptr = buffer_ptrs.get(&node).copied().unwrap_or(0);
+            let input_ptrs: Vec<u64> = kernel
+                .inputs
+                .iter()
+                .map(|inp| buffer_ptrs.get(inp).copied().unwrap_or(0))
+                .collect();
+
+            let param_values = kernel.kernel_op.build_params(
+                stream,
+                output_ptr,
+                &input_ptrs,
+                &kernel.internal_bufs,
+                dyn_dims_ptr,
+            );
+            let mut params = UnifiedKernelParams::new(param_values);
+            let cu_func = unsafe { kernel.function.raw_function() };
+
+            let result = unsafe {
+                cuLaunchKernel(
+                    cu_func,
+                    grid.0,
+                    grid.1,
+                    grid.2,
+                    block.0,
+                    block.1,
+                    block.2,
+                    shared_mem,
+                    cu_stream,
+                    params.as_cuda_params(),
+                    std::ptr::null_mut(),
+                )
+            };
+            if result != CUresult::CUDA_SUCCESS {
+                eprintln!(
+                    "[seq-debug] kernel #{idx}/{num_kernels} '{kernel_name}' \
+                     node={node:?} grid={grid:?} block={block:?} \
+                     output_ptr={output_ptr:#x} inputs={input_ptrs:#x?} \
+                     LAUNCH FAILED: {result:?}"
+                );
+                anyhow::bail!(
+                    "kernel #{idx} '{kernel_name}' (node {node:?}) launch failed: {result:?}"
+                );
+            }
+            if let Err(e) = stream.synchronize() {
+                eprintln!(
+                    "[seq-debug] kernel #{idx}/{num_kernels} '{kernel_name}' \
+                     node={node:?} grid={grid:?} block={block:?} \
+                     output_ptr={output_ptr:#x} inputs={input_ptrs:#x?} \
+                     SYNC FAILED: {e}"
+                );
+                anyhow::bail!(
+                    "kernel #{idx} '{kernel_name}' (node {node:?}) sync failed: {e}"
+                );
+            }
+        }
+
+        Ok(())
+    }
+
     /// Build the CUDA graph from compiled kernels.
     fn build_graph(
         &self,
@@ -653,6 +811,11 @@ pub fn kernel_to_host(
     }
 
     let kernel_subgraphs = partition_marked_convex(llir_graph, &kernel_ops_in_graph).unwrap();
+    // Compute the set of FS / FE / FusedX nodes globally absorbed by some
+    // FusionEnd in the LLIR. Used by `build_compile_units` to suppress the
+    // identity-memcpy fallback for shared FS leaves whose consumers live
+    // in a different convex subgraph than the FS itself.
+    let globally_absorbed = region_codegen::globally_absorbed_markers(llir_graph);
 
     // Track which kernel node belongs to which CudaGraphOp (for later edge creation)
     let mut kernel_to_cuda_graph: FxHashMap<NodeIndex, NodeIndex> = FxHashMap::default();
@@ -687,45 +850,98 @@ pub fn kernel_to_host(
             set_global_dyn_dims(global_dyn_dims.clone());
         }
 
-        // Compile all kernels with global ordering for correct dyn_dims indices
-        let mut kernels = Vec::with_capacity(topo_order.len());
-        for kernel_node_idx in &topo_order {
-            let kernel_op_ref = llir_graph[*kernel_node_idx]
-                .to_dialect::<dyn KernelOp>()
-                .unwrap();
+        // Group the topo order into compile units: each FusionEnd-rooted
+        // region collapses to a single CompileUnit::Region (one fused
+        // CUDA kernel for the whole DAG); everything else stays as
+        // CompileUnit::Single (the existing per-op compile path).
+        let compile_units =
+            region_codegen::build_compile_units(&topo_order, llir_graph, &globally_absorbed);
 
-            let (kernel_function, _, _kernel_str, grid, block, shared_mem, constants) =
-                kernel_op_ref.compile(cuda_stream, kernel_cache);
+        // Compile all units with global ordering for correct dyn_dims indices
+        let mut kernels = Vec::with_capacity(compile_units.len());
+        for unit in &compile_units {
+            match unit {
+                CompileUnit::Single(kernel_node_idx) => {
+                    let kernel_op_ref = llir_graph[*kernel_node_idx]
+                        .to_dialect::<dyn KernelOp>()
+                        .unwrap();
 
-            // Collect inputs from graph edges
-            let mut inputs: Vec<NodeIndex> = llir_graph
-                .edges_directed(*kernel_node_idx, Direction::Incoming)
-                .sorted_by_key(|e| e.id())
-                .map(|e| e.source())
-                .collect_vec();
+                    let (kernel_function, _, _kernel_str, grid, block, shared_mem, constants) =
+                        kernel_op_ref.compile(cuda_stream, kernel_cache);
 
-            // Collect buffer nodes and sizes
-            // Only add kernel nodes with non-zero output size (MegakernelOps have size 0)
-            let output_size = kernel_op_ref.output_size();
-            if output_size.exec(&FxHashMap::default()).unwrap_or(1) != 0 {
-                all_buffer_nodes.insert(*kernel_node_idx);
-                all_buffer_sizes.insert(*kernel_node_idx, output_size);
+                    // Collect inputs from graph edges
+                    let inputs: Vec<NodeIndex> = llir_graph
+                        .edges_directed(*kernel_node_idx, Direction::Incoming)
+                        .sorted_by_key(|e| e.id())
+                        .map(|e| e.source())
+                        .collect_vec();
+
+                    // Collect buffer nodes and sizes
+                    // Only add kernel nodes with non-zero output size (MegakernelOps have size 0)
+                    let output_size = kernel_op_ref.output_size();
+                    if output_size.exec(&FxHashMap::default()).unwrap_or(1) != 0 {
+                        all_buffer_nodes.insert(*kernel_node_idx);
+                        all_buffer_sizes.insert(*kernel_node_idx, output_size);
+                    }
+                    all_buffer_nodes.extend(inputs.iter().copied());
+
+                    let kernel_op: Arc<Box<dyn KernelOp>> = Arc::clone(kernel_op_ref);
+
+                    kernels.push(CompiledKernel::new(
+                        *kernel_node_idx,
+                        kernel_function,
+                        grid,
+                        block,
+                        shared_mem,
+                        inputs,
+                        kernel_op.clone(),
+                        constants,
+                        kernel_op.kernel_name(),
+                    ));
+                }
+                CompileUnit::Region(region) => {
+                    // Generate one fused CUDA kernel for the whole region.
+                    let compiled = region_codegen::compile_region(
+                        region,
+                        llir_graph,
+                        cuda_stream,
+                        kernel_cache,
+                    );
+
+                    // The region's CompiledKernel is keyed on the FE node
+                    // (so FE provides trait methods like output_size /
+                    // build_params) but its `inputs` are the external
+                    // producers, not FE's literal LLIR predecessors —
+                    // those are interior FusedX nodes that don't exist
+                    // as buffer-bearing nodes from the host's view.
+                    let fe_op_ref = llir_graph[region.fe_node]
+                        .to_dialect::<dyn KernelOp>()
+                        .unwrap();
+
+                    let inputs: Vec<NodeIndex> = region.external_inputs.clone();
+
+                    let output_size = fe_op_ref.output_size();
+                    if output_size.exec(&FxHashMap::default()).unwrap_or(1) != 0 {
+                        all_buffer_nodes.insert(region.fe_node);
+                        all_buffer_sizes.insert(region.fe_node, output_size);
+                    }
+                    all_buffer_nodes.extend(inputs.iter().copied());
+
+                    let kernel_op: Arc<Box<dyn KernelOp>> = Arc::clone(fe_op_ref);
+
+                    kernels.push(CompiledKernel::new(
+                        region.fe_node,
+                        compiled.function,
+                        compiled.grid,
+                        compiled.block,
+                        compiled.shared_mem,
+                        inputs,
+                        kernel_op,
+                        compiled.constants,
+                        "FusedRegion",
+                    ));
+                }
             }
-            all_buffer_nodes.extend(inputs.iter().copied());
-
-            let kernel_op: Arc<Box<dyn KernelOp>> = Arc::clone(kernel_op_ref);
-
-            kernels.push(CompiledKernel::new(
-                *kernel_node_idx,
-                kernel_function,
-                grid,
-                block,
-                shared_mem,
-                inputs,
-                kernel_op.clone(),
-                constants,
-                kernel_op.kernel_name(),
-            ));
         }
 
         // Get the possibly-extended global ordering (kernels may have discovered new dims)
@@ -818,22 +1034,41 @@ pub fn kernel_to_host(
         }
     }
 
-    // Add collected edges (deduplicate), skipping back-edges to preserve DAG property
+    // Add each cross-CudaGraphOp dep edge iff it would carry new ordering
+    // information without closing a cycle. The previous topo-position gate
+    // ("skip when src_pos >= dst_pos") was too coarse: it dropped edges
+    // whose src happened to land later in the toposort than their dst even
+    // when no path dst→src actually existed, leaving consumers free to run
+    // before the producer wrote their input buffer (wrong outputs); and it
+    // also added edges that were already implied by an existing src→dst
+    // path (extra serialization, no new info).
     let edges_to_add: FxHashSet<(NodeIndex, NodeIndex)> = edges_to_add.into_iter().collect();
-    let topo = toposort(&*llir_graph, None).unwrap();
-    let mut topo_pos: FxHashMap<NodeIndex, usize> = FxHashMap::default();
-    for (i, n) in topo.iter().enumerate() {
-        topo_pos.insert(*n, i);
-    }
+    use petgraph::algo::has_path_connecting;
     for (src, dst) in edges_to_add {
-        // Only add forward edges (src before dst in topo order) to avoid creating cycles
-        let src_pos = topo_pos.get(&src).copied().unwrap_or(usize::MAX);
-        let dst_pos = topo_pos.get(&dst).copied().unwrap_or(usize::MAX);
-        if src_pos >= dst_pos {
-            continue; // Skip back-edges
+        if has_path_connecting(&*llir_graph, src, dst, None) {
+            continue; // already ordered src→dst by some path; edge redundant
         }
-        if !llir_graph.edges_connecting(src, dst).any(|_| true) {
-            llir_graph.add_edge(src, dst, ());
+        if has_path_connecting(&*llir_graph, dst, src, None) {
+            continue; // adding src→dst would close a cycle
+        }
+        llir_graph.add_edge(src, dst, ());
+    }
+
+    // Strip fully-absorbed marker nodes (FusionStart, nested FusionEnd,
+    // FusedX) from the LLIR. Region codegen has already folded them into
+    // a single fused CUDA function anchored at each region's root
+    // FusionEnd; the absorbed nodes have no consumers outside the region
+    // and never need their own buffers. Removing them keeps later
+    // per-execute walks (e.g., `allocate_intermediate_buffers`) from
+    // chewing through dead nodes every decode token.
+    //
+    // Root FusionEnd nodes are NOT in `globally_absorbed` (they were the
+    // walks' starting points), so we keep them — they're the kernel
+    // anchor for the region's compiled kernel.
+    for node in globally_absorbed {
+        // Defensive: only remove if the node still exists.
+        if llir_graph.node_weight(node).is_some() {
+            llir_graph.remove_node(node);
         }
     }
 }

crates/luminal_cuda_lite/src/lib.rs

@@ -1,6 +1,6 @@
+pub mod dyn_backend;
 pub mod host;
 pub mod kernel;
-pub mod logical;
 pub mod runtime;
 use std::{
     ffi::{CStr, CString},
@@ -10,6 +10,8 @@ use std::{
 
 pub use cudarc;
 
+use cudarc::{cublaslt::CudaBlasLT, driver::CudaStream};
+
 #[cfg(test)]
 mod tests;
 
@@ -138,6 +140,25 @@ fn cuda_driver_diagnostics() -> (Option<i32>, Option<i32>) {
     (driver_version, None)
 }
 
+pub(crate) fn try_create_cublaslt(
+    stream: Arc<CudaStream>,
+) -> std::result::Result<Arc<CudaBlasLT>, String> {
+    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| CudaBlasLT::new(stream))) {
+        Ok(Ok(handle)) => Ok(Arc::new(handle)),
+        Ok(Err(err)) => Err(err.to_string()),
+        Err(payload) => {
+            let message = if let Some(message) = payload.downcast_ref::<String>() {
+                message.clone()
+            } else if let Some(message) = payload.downcast_ref::<&str>() {
+                message.to_string()
+            } else {
+                "cuBLASLt initialization panicked".to_string()
+            };
+            Err(message)
+        }
+    }
+}
+
 fn cuda_nvrtc_compile_options(target_arch: &str) -> Vec<String> {
     let mut options = cuda_nvrtc_include_paths()
         .into_iter()
@@ -187,9 +208,9 @@ fn get_cubin(program: nvrtc_sys::nvrtcProgram) -> Result<Vec<u8>, NvrtcError> {
     }
 
     let mut cubin = Vec::with_capacity(cubin_size);
-    cubin.resize(cubin_size, 0);
-    unsafe { nvrtc_sys::nvrtcGetCUBIN(program, cubin.as_mut_ptr()) }.result()?;
-    Ok(cubin.into_iter().map(|byte| byte as u8).collect())
+    cubin.resize(cubin_size, 0u8);
+    unsafe { nvrtc_sys::nvrtcGetCUBIN(program, cubin.as_mut_ptr() as *mut _) }.result()?;
+    Ok(cubin)
 }
 
 pub(crate) fn compile_module_image_for_current_device<S: AsRef<str>>(

crates/luminal_cuda_lite/src/logical.rs

@@ -1,71 +0,0 @@
-use std::fmt::Debug;
-
-use luminal::{
-    egglog_utils::api::{Rule, SortDef},
-    hlir::unary_sort,
-    op::EgglogOp,
-};
-
-pub type Ops = (Exp, Sigmoid);
-
-#[derive(Debug, Default)]
-pub struct Exp;
-impl EgglogOp for Exp {
-    fn sort(&self) -> SortDef {
-        unary_sort("Exp")
-    }
-
-    fn cleanup(&self) -> bool {
-        true
-    }
-
-    fn rewrites(&self) -> Vec<Rule> {
-        vec![Rule::raw(
-            "(rule
-            (
-                (= ?exp_const (Op (Constant 1.442695) (INil)))
-                (= ?mul (Op (Mul ?shape ?x_stride ?const_stride ?intermediate_stride) (ICons ?x (ICons ?exp_const (INil)))))
-                (= ?exp2 (Op (Exp2 ?shape ?intermediate_stride ?out_stride) (ICons ?mul (INil))))
-                (= ?dt (dtype ?x))
-            )
-            (
-                (let ?exp (Op (Exp ?shape ?x_stride ?out_stride) (ICons ?x (INil))))
-                (union ?exp2 ?exp)
-                (set (dtype ?exp) ?dt)
-            )
-        )",
-        )]
-    }
-}
-
-#[derive(Default, Debug, Clone)]
-pub struct Sigmoid;
-impl EgglogOp for Sigmoid {
-    fn sort(&self) -> SortDef {
-        unary_sort("Sigmoid")
-    }
-
-    fn cleanup(&self) -> bool {
-        true
-    }
-
-    fn rewrites(&self) -> Vec<Rule> {
-        vec![Rule::raw("(rule
-            (
-                (= ?neg1 (Op (Constant -1.0) (INil)))
-                (= ?neg_input (Op (Mul ?input_range ?input_stride ?const_stride ?intermediate_stride) (ICons ?input (ICons ?neg1 (INil)))))
-                (= ?exp (Op (Exp ?input_range ?intermediate_stride ?exp_stride) (ICons ?neg_input (INil))))
-                (= ?one (Op (Constant 1.0) (INil)))
-                (= ?plus_one (Op (Add ?input_range ?exp_stride ?const_stride ?plus_one_stride) (ICons ?exp (ICons ?one (INil)))))
-                (= ?sig_out (Op (Recip ?input_range ?plus_one_stride ?out_stride) (ICons ?plus_one (INil))))
-                (= ?dt (dtype ?input))
-            )
-            (
-                (let ?sig (Op (Sigmoid ?input_range ?input_stride ?out_stride) (ICons ?input (INil))))
-                (union ?sig_out ?sig)
-                (set (dtype ?sig) ?dt)
-            )
-            :name \"sigmoid\"
-        )")]
-    }
-}

crates/luminal_cuda_lite/src/runtime.rs

@@ -119,6 +119,18 @@ pub struct CudaRuntime {
     active_bucket: usize,
     /// Bucket definitions per dimension (empty = single-bucket mode)
     dim_buckets: FxHashMap<char, Vec<DimBucket>>,
+
+    /// Non-owning CudaSlice wrappers for external device pointers.
+    /// ManuallyDrop prevents cuMemFree — the external allocator (e.g. PyTorch) owns the memory.
+    external_buffers: FxHashMap<NodeIndex, std::mem::ManuallyDrop<CudaSlice<u8>>>,
+
+    /// Pending output pointer registrations: HLIR output id -> (device_ptr, n_bytes)
+    /// Set by python before execute(), consumed at start of execute()
+    output_ptr_registrations: FxHashMap<NodeIndex, (u64, usize)>,
+
+    /// Non-owning CudaSlice views of external output pointers, keyed by LLIR data node
+    /// ManuallyDrop prevents cuMemFree -- Pytorch owns the memory
+    external_output_buffers: FxHashMap<NodeIndex, std::mem::ManuallyDrop<CudaSlice<u8>>>,
 }
 
 impl CudaRuntime {
@@ -199,6 +211,48 @@ impl CudaRuntime {
         self.changed_hlir.insert(id);
     }
 
+    /// Set an external CUDA device pointer as input data. Zero-copy.
+    /// The caller must ensure the pointer remains valid for the runtime's lifetime.
+    ///
+    /// # Safety
+    /// The device pointer must point to a valid CUDA allocation on the same device
+    /// as this runtime's stream, with at least `n_bytes` bytes available.
+    pub unsafe fn set_device_ptr(&mut self, id: impl ToId, device_ptr: u64, n_bytes: usize) {
+        debug_assert!(device_ptr != 0, "set_device_ptr called with null pointer");
+        let id = id.to_id();
+        // Create CudaSlice view via cudarc's upgrade_device_ptr.
+        // ManuallyDrop prevents cuMemFree on drop (external allocator owns this memory).
+        let slice = unsafe {
+            self.cuda_stream
+                .upgrade_device_ptr::<u8>(device_ptr, n_bytes)
+        };
+        self.external_buffers
+            .insert(id, std::mem::ManuallyDrop::new(slice));
+        self.hlir_buffers.insert(id, CudaInput::Ptr(device_ptr));
+        self.changed_hlir.insert(id);
+    }
+
+    /// Register an external device pointer for an output tensor (zero-copy output).
+    /// The pointer is stored lazily — resolution to LLIR nodes happens in execute().
+    ///
+    /// # Safety
+    /// The device pointer must point to a valid CUDA allocation with at least `n_bytes` bytes,
+    /// and must remain valid through the next execute() call.
+    pub unsafe fn set_output_device_ptr(&mut self, id: impl ToId, device_ptr: u64, n_bytes: usize) {
+        debug_assert!(
+            device_ptr != 0,
+            "set_output_device_ptr called with null pointer"
+        );
+        self.output_ptr_registrations
+            .insert(id.to_id(), (device_ptr, n_bytes));
+    }
+
+    pub fn output_is_zero_copy(&self, id: impl ToId) -> bool {
+        let producer = self.find_producer_node(id);
+        let data_node = self.follow_aliases(producer);
+        self.external_output_buffers.contains_key(&data_node)
+    }
+
     /// Find the LLIR producing node for an output tensor.
     fn find_producer_node(&self, id: impl ToId) -> NodeIndex {
         let id = id.to_id();
@@ -281,12 +335,15 @@ impl CudaRuntime {
                 .expect("Cannot find input tensor in runtime!")
             {
                 CudaInput::Buffer(buf) => self.cuda_stream.clone_dtoh(buf).unwrap(),
-                CudaInput::Ptr(p) => {
-                    // Raw pointer — need size from cached_buffer_ptrs or error
-                    panic!(
-                        "Cannot read raw pointer input (ptr=0x{:x}) — use Buffer variant",
-                        p
-                    );
+                CudaInput::Ptr(_) => {
+                    // External device pointer — use the CudaSlice view from external_buffers
+                    if let Some(ext) = self.external_buffers.get(hlir_node) {
+                        self.cuda_stream.clone_dtoh(&**ext).unwrap()
+                    } else {
+                        panic!(
+                            "Cannot read raw pointer input — no external_buffers entry for node"
+                        );
+                    }
                 }
             }
         } else {
@@ -302,6 +359,101 @@ impl CudaRuntime {
         }
     }
 
+    /// Resolve the device-side CudaSlice for an output tensor without copying to host.
+    /// Used by copy_output_to_device_ptr for DtoD transfers.
+    fn resolve_output_slice(&self, id: impl ToId) -> &CudaSlice<u8> {
+        let data_id = self.resolve_data_node(id);
+        let bucket = self.active();
+        if let Some(hlir_node) = bucket.llir_to_hlir.get(&data_id) {
+            match self
+                .hlir_buffers
+                .get(hlir_node)
+                .expect("Cannot find input tensor in runtime!")
+            {
+                CudaInput::Buffer(buf) => buf,
+                CudaInput::Ptr(_) => self
+                    .external_buffers
+                    .get(hlir_node)
+                    .map(|ext| &**ext)
+                    .expect("Cannot read raw pointer input — no external_buffers entry for node"),
+            }
+        } else {
+            bucket
+                .buffers
+                .get(&data_id)
+                .expect("Cannot find tensor in runtime!")
+        }
+    }
+
+    /// Copy output tensor data to an external CUDA device pointer (DtoD).
+    /// Much faster than get_f32 + HtoD for CUDA-to-CUDA workflows.
+    ///
+    /// # Safety
+    /// The dest_ptr must be a valid CUDA device allocation with at least n_bytes available.
+    pub unsafe fn copy_output_to_device_ptr(&self, id: impl ToId, dest_ptr: u64, n_bytes: usize) {
+        debug_assert!(
+            dest_ptr != 0,
+            "copy_output_to_device_ptr called with null pointer"
+        );
+        let src_slice = self.resolve_output_slice(id);
+        let src_ptr = src_slice.device_ptr(&self.cuda_stream).0;
+        let copy_bytes = n_bytes.min(src_slice.len());
+        unsafe {
+            cudarc::driver::result::memcpy_dtod_async(
+                dest_ptr,
+                src_ptr,
+                copy_bytes,
+                self.cuda_stream.cu_stream(),
+            )
+            .expect("cuMemcpyDtoDAsync failed");
+        }
+        self.cuda_stream.synchronize().unwrap();
+    }
+
+    /// Resolve pending output pointer registrations into external_output_buffers.
+    /// Called at the start of execute(), after buffer allocation and HLIR sync.
+    fn apply_output_ptr_registrations(&mut self) {
+        // clear stale external output buffers from previous execution
+        self.external_output_buffers.clear();
+
+        if self.output_ptr_registrations.is_empty() {
+            return;
+        }
+
+        // Collect registrations to avoid borrow conflict (drain borrows self mutably,
+        // but find_producer_node/follow_aliases need &self).
+
+        let registrations: Vec<_> = self.output_ptr_registrations.drain().collect();
+
+        for (hlir_id, (device_ptr, n_bytes)) in registrations {
+            // Resolve HLIR output id -> LLIR producer -> follow aliases -> data node
+            let producer = self.find_producer_node(hlir_id);
+            let data_node = self.follow_aliases(producer);
+
+            // If data_node is an HLIR input (aliased output), skip — can't substitute
+            if self.compiled_buckets[self.active_bucket]
+                .llir_to_hlir
+                .contains_key(&data_node)
+            {
+                continue;
+            }
+
+            // Create non-owning CudaSlice view of PyTorch's buffer
+            let slice = unsafe {
+                self.cuda_stream
+                    .upgrade_device_ptr::<u8>(device_ptr, n_bytes)
+            };
+
+            self.external_output_buffers
+                .insert(data_node, std::mem::ManuallyDrop::new(slice));
+
+            // Update cached_buffer_ptrs so CudaGraphOp picks up the new pointer
+            self.compiled_buckets[self.active_bucket]
+                .cached_buffer_ptrs
+                .insert(data_node, device_ptr);
+        }
+    }
+
     pub fn get_f32(&self, id: impl ToId) -> Vec<f32> {
         let bytes = self.get_output_data(id);
         let bytes = bytes.leak();
@@ -512,6 +664,22 @@ impl CudaRuntime {
             if bucket.llir_graph[node].to_op::<Input>().is_some() {
                 continue;
             }
+            // Skip fusion marker / interior nodes. Region codegen folds
+            // FusionStart / FusionEnd / FusedX into a single CUDA function
+            // anchored at the FusionEnd; these marker nodes never need a
+            // device buffer of their own at runtime, so walking them here
+            // each step (with `p` incrementing every decode token) is
+            // pure overhead. Skipping them recovers ~2 ms / token on
+            // llama with fusion enabled.
+            if let Some(op) = bucket.llir_graph[node].to_dialect::<dyn KernelOp>() {
+                let kn = op.kernel_name();
+                if kn == "FusionStart" || kn.starts_with("Fused") {
+                    continue;
+                }
+                // Note: we deliberately keep "FusionEnd" because it is the
+                // anchor for the region's compiled kernel and DOES need a
+                // buffer for the region's output.
+            }
             let needed_bytes =
                 if let Some(op) = bucket.llir_graph[node].to_dialect::<dyn KernelOp>() {
                     let out_bytes = op.output_bytes();
@@ -684,7 +852,7 @@ fn format_duration_precise(d: &std::time::Duration) -> String {
 }
 
 impl Runtime for CudaRuntime {
-    type Ops = (crate::logical::Ops, crate::kernel::Ops, crate::host::Ops);
+    type Ops = (crate::kernel::Ops, crate::host::Ops);
     type CompileArg = Arc<CudaStream>;
     type ExecReturn = ();
     type ProfileMetric = Duration;
@@ -702,9 +870,16 @@ impl Runtime for CudaRuntime {
             compiled_buckets: vec![CompiledBucket::new()],
             active_bucket: 0,
             dim_buckets: FxHashMap::default(),
+            output_ptr_registrations: FxHashMap::default(),
+            external_output_buffers: FxHashMap::default(),
+            external_buffers: FxHashMap::default(),
         }
     }
 
+    fn aggregate_profile_metrics(metrics: &[Self::ProfileMetric]) -> Self::ProfileMetric {
+        metrics.iter().copied().sum()
+    }
+
     #[tracing::instrument(skip_all)]
     fn load_llir(&mut self, llir_graph: &LLIRGraph) {
         // Sync before clearing old data to ensure all operations complete
@@ -737,15 +912,13 @@ impl Runtime for CudaRuntime {
         }
     }
 
-    fn allocate_dummy_input(&mut self, node_index: usize, num_elements: usize) {
-        // Use small non-zero values (ones) instead of zeros so that NaN-producing
-        // graph variants are detected during profiling. Zero inputs often hide
-        // numerical issues that appear with real data.
-        let host_data = vec![1.0f32; num_elements];
-        let buf = self
-            .cuda_stream
-            .clone_htod(bytemuck::cast_slice::<f32, u8>(&host_data))
-            .unwrap();
+    fn allocate_dummy_input(&mut self, node_index: usize, num_bytes: usize) {
+        // Boundary scratch buffers are sized in raw bytes and may represent
+        // non-float tensors such as gather/scatter indices. Initialize with zero
+        // bytes so integer boundaries stay in-range and the raw allocation size
+        // matches the requested tensor storage.
+        let host_data = vec![0u8; num_bytes];
+        let buf = self.cuda_stream.clone_htod(&host_data).unwrap();
         let id = NodeIndex::new(node_index);
         self.hlir_buffers.insert(id, CudaInput::Buffer(buf));
         self.changed_hlir.insert(id);
@@ -923,6 +1096,9 @@ impl Runtime for CudaRuntime {
         // Ensure all CUDA graphs are built (handles first execute and any missing graphs)
         self.prebuild_graphs(dyn_map);
 
+        // Resolve external output pointer registrations (zero-copy output path)
+        self.apply_output_ptr_registrations();
+
         let total_start = std::time::Instant::now();
         let bucket = &self.compiled_buckets[self.active_bucket];
 
@@ -932,16 +1108,32 @@ impl Runtime for CudaRuntime {
 
             // Build buffer map for the HostOp interface
             let mut buffer_map: FxHashMap<NodeIndex, &CudaSlice<u8>> = FxHashMap::default();
-            // Add output buffer
-            if let Some(buf) = bucket.buffers.get(&exec_op.output) {
+
+            // Add output buffer -- prefer external output pointer if registered (zero copy)
+            if let Some(ext) = self.external_output_buffers.get(&exec_op.output) {
+                buffer_map.insert(exec_op.output, &**ext);
+            } else if let Some(buf) = bucket.buffers.get(&exec_op.output) {
                 buffer_map.insert(exec_op.output, buf);
             }
             // Add input buffers (prefer HLIR weight buffers over intermediate placeholders)
             for inp in exec_op.inputs.iter() {
-                if let Some(hlir_node) = bucket.llir_to_hlir.get(inp)
-                    && let Some(CudaInput::Buffer(buf)) = self.hlir_buffers.get(hlir_node)
-                {
-                    buffer_map.insert(*inp, buf);
+                if let Some(hlir_node) = bucket.llir_to_hlir.get(inp) {
+                    match self.hlir_buffers.get(hlir_node) {
+                        Some(CudaInput::Buffer(buf)) => {
+                            buffer_map.insert(*inp, buf);
+                        }
+                        Some(CudaInput::Ptr(_)) => {
+                            if let Some(ext) = self.external_buffers.get(hlir_node) {
+                                buffer_map.insert(*inp, &**ext);
+                            }
+                        }
+                        None => {}
+                    }
+                    if !buffer_map.contains_key(inp)
+                        && let Some(buf) = bucket.buffers.get(inp)
+                    {
+                        buffer_map.insert(*inp, buf);
+                    }
                 } else if let Some(buf) = bucket.buffers.get(inp) {
                     buffer_map.insert(*inp, buf);
                 }
@@ -950,27 +1142,47 @@ impl Runtime for CudaRuntime {
             let extra_nodes = exec_op.internal.extra_buffer_nodes();
             for extra_node in extra_nodes {
                 if let Entry::Vacant(e) = buffer_map.entry(extra_node) {
-                    if let Some(buf) = bucket.buffers.get(&extra_node) {
-                        e.insert(buf);
-                    } else if let Some(hlir_node) = bucket.llir_to_hlir.get(&extra_node)
-                        && let Some(CudaInput::Buffer(buf)) = self.hlir_buffers.get(hlir_node)
-                    {
+                    if let Some(ext) = self.external_output_buffers.get(&extra_node) {
+                        e.insert(&**ext);
+                    } else if let Some(buf) = bucket.buffers.get(&extra_node) {
                         e.insert(buf);
+                    } else if let Some(hlir_node) = bucket.llir_to_hlir.get(&extra_node) {
+                        match self.hlir_buffers.get(hlir_node) {
+                            Some(CudaInput::Buffer(buf)) => {
+                                e.insert(buf);
+                            }
+                            Some(CudaInput::Ptr(_)) => {
+                                if let Some(ext) = self.external_buffers.get(hlir_node) {
+                                    e.insert(&**ext);
+                                }
+                            }
+                            None => {}
+                        }
                     }
                 }
             }
             // Resolve output aliases
             for (&alias_node, &alias_target) in &bucket.output_alias_map {
-                if let std::collections::hash_map::Entry::Occupied(mut e) =
-                    buffer_map.entry(alias_node)
-                {
-                    if let Some(hlir_node) = bucket.llir_to_hlir.get(&alias_target)
-                        && let Some(CudaInput::Buffer(buf)) = self.hlir_buffers.get(hlir_node)
-                    {
-                        e.insert(buf);
-                    } else if let Some(buf) = bucket.buffers.get(&alias_target) {
-                        e.insert(buf);
-                    }
+                if !buffer_map.contains_key(&alias_node) {
+                    continue;
+                }
+                // Try HLIR buffer first (includes external device pointers)
+                let resolved: Option<&CudaSlice<u8>> =
+                    if let Some(hlir_node) = bucket.llir_to_hlir.get(&alias_target) {
+                        match self.hlir_buffers.get(hlir_node) {
+                            Some(CudaInput::Buffer(buf)) => Some(buf),
+                            Some(CudaInput::Ptr(_)) => {
+                                self.external_buffers.get(hlir_node).map(|ext| &**ext)
+                            }
+                            None => None,
+                        }
+                    } else {
+                        None
+                    };
+                if let Some(buf) = resolved {
+                    buffer_map.insert(alias_node, buf);
+                } else if let Some(buf) = bucket.buffers.get(&alias_target) {
+                    buffer_map.insert(alias_node, buf);
                 }
             }
             let _span = span!(
@@ -1017,11 +1229,6 @@ impl Runtime for CudaRuntime {
             }
         }
 
-        // Final sync to ensure all operations completed successfully
-        self.cuda_stream
-            .synchronize()
-            .expect("Final sync failed in execute");
-
         // Consume input buffers
         if self.profiling {
             return;
@@ -1074,6 +1281,7 @@ impl Runtime for CudaRuntime {
 
         for hlir_node in to_consume {
             self.hlir_buffers.remove(&hlir_node);
+            self.external_buffers.remove(&hlir_node);
             let bucket = &mut self.compiled_buckets[self.active_bucket];
             if let Some(llir_node) = bucket.hlir_to_llir.get(&hlir_node) {
                 bucket.cached_buffer_ptrs.remove(llir_node);

crates/luminal_cuda_lite/src/tests/bucket_tests.rs

@@ -41,7 +41,7 @@ fn test_bucket_dispatch_simple() {
     rt.set_data(a, vec![1.0f32; 4]);
 
     let mut rng = SmallRng::seed_from_u64(42);
-    rt = cx.search_rng(rt, 5, &mut rng);
+    rt = cx.search_options(rt, SearchOptions::new(5), &mut rng);
 
     // Test bucket 1: s=1
     cx.set_dim('s', 1);
@@ -85,7 +85,7 @@ fn test_bucket_matmul_dynamic() {
     rt.set_data(b_tensor, b_data.clone());
 
     let mut rng = SmallRng::seed_from_u64(42);
-    rt = cx.search_rng(rt, 5, &mut rng);
+    rt = cx.search_options(rt, SearchOptions::new(5), &mut rng);
 
     // Execute at s=1
     cx.set_dim('s', 1);
@@ -140,7 +140,7 @@ fn test_bucket_results_match_unbucketed() {
     let input_data = random_f32_vec(12, seed, -1.0, 1.0);
     rt1.set_data(a1, input_data.clone());
     let mut rng1 = SmallRng::seed_from_u64(seed);
-    rt1 = cx1.search_rng(rt1, 5, &mut rng1);
+    rt1 = cx1.search_options(rt1, SearchOptions::new(5), &mut rng1);
     rt1.set_data(a1, input_data.clone());
     rt1.execute(&cx1.dyn_map);
     let result_unbucketed = rt1.get_f32(b1);
@@ -153,7 +153,7 @@ fn test_bucket_results_match_unbucketed() {
     let mut rt2 = CudaRuntime::initialize(stream.clone());
     rt2.set_data(a2, input_data.clone());
     let mut rng2 = SmallRng::seed_from_u64(seed);
-    rt2 = cx2.search_rng(rt2, 5, &mut rng2);
+    rt2 = cx2.search_options(rt2, SearchOptions::new(5), &mut rng2);
     rt2.set_data(a2, input_data.clone());
     rt2.execute(&cx2.dyn_map);
     let result_bucketed = rt2.get_f32(b2);
@@ -179,7 +179,7 @@ fn test_bucket_out_of_range_panics() {
     cx.set_dim('s', 1);
     rt.set_data(a, vec![1.0f32; 4]);
     let mut rng = SmallRng::seed_from_u64(42);
-    rt = cx.search_rng(rt, 3, &mut rng);
+    rt = cx.search_options(rt, SearchOptions::new(3), &mut rng);
 
     // s=10 is outside all buckets — should panic
     cx.set_dim('s', 10);
@@ -204,7 +204,7 @@ fn test_bucket_no_buckets_backward_compat() {
     let input_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
     rt.set_data(a, input_data.clone());
     let mut rng = SmallRng::seed_from_u64(42);
-    rt = cx.search_rng(rt, 3, &mut rng);
+    rt = cx.search_options(rt, SearchOptions::new(3), &mut rng);
 
     rt.set_data(a, input_data.clone());
     rt.execute(&cx.dyn_map);
@@ -249,7 +249,7 @@ fn test_bucket_switch_preserves_weights() {
     rt.set_data(b_tensor, b_data.clone());
 
     let mut rng = SmallRng::seed_from_u64(42);
-    rt = cx.search_rng(rt, 5, &mut rng);
+    rt = cx.search_options(rt, SearchOptions::new(5), &mut rng);
 
     // Execute with bucket 1 (s=1)
     cx.set_dim('s', 1);
@@ -305,7 +305,7 @@ fn test_bucket_multiple_executions_same_bucket() {
     cx.set_dim('s', 1);
     rt.set_data(a, vec![1.0f32; 4]);
     let mut rng = SmallRng::seed_from_u64(42);
-    rt = cx.search_rng(rt, 3, &mut rng);
+    rt = cx.search_options(rt, SearchOptions::new(3), &mut rng);
 
     // Execute at different sizes within the same bucket
     for s in [1, 2, 4, 8] {

crates/luminal_cuda_lite/src/tests/consumed_buffer_tests.rs

@@ -301,9 +301,8 @@ fn test_scatter_kv_cache_roundtrip() {
 }
 
 /// Test scatter with TWO cache buffers and dual outputs (closer to llama K+V pattern).
-/// Also verifies graph_break interaction.
 #[test]
-fn test_scatter_dual_cache_with_graph_break() {
+fn test_scatter_dual_cache() {
     let ctx = CudaContext::new(0).unwrap();
     ctx.bind_to_thread().unwrap();
     let stream = ctx.default_stream();
@@ -348,7 +347,7 @@ fn test_scatter_dual_cache_with_graph_break() {
     // Use seeded search for deterministic scatter variant selection.
     // Seed 0 reliably selects Scatter (not ScatterNoCopy) for both caches.
     let mut rng = rand::rngs::SmallRng::seed_from_u64(0);
-    rt = cx.search_rng(rt, 5, &mut rng);
+    rt = cx.search_options(rt, SearchOptions::new(5), &mut rng);
 
     // Print selected variants
     for node in rt.llir_graph().node_weights() {

crates/luminal_cuda_lite/src/tests/fusion.rs

@@ -0,0 +1,986 @@
+use luminal::egglog_utils::{egglog_to_llir, random_initial_choice};
+use luminal::prelude::*;
+
+use crate::kernel::KernelOp;
+use crate::runtime::CudaRuntime;
+use crate::tests::utilities::{
+    TOLERANCE_SAFETY_FACTOR, dtype_epsilon, random_f32_vec, test_binary_cuda, test_unary_cuda,
+};
+
+#[test]
+fn test_two_unary_ops_fuse() {
+    // Marker form: `a.sin().sqrt()` should fuse into a region with FusedSin
+    // and FusedSqrt under one FusionEnd (per pair-fuse U→U).
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let _b = a.sin().sqrt().output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedSin", "FusedSqrt"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 1 && r.end_count == 1),
+        "expected a marker region of {expected:?} with 1 FusionStart, got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_stride_mismatch_prevents_fusion() {
+    // A permute between sin and sqrt gives sqrt a non-contiguous view of sin's
+    // contiguous output, so sqrt's in_strides != its out_strides and the
+    // non-linear `?s ?s` match in the pair-fuse U→U rule can't fire.
+    let mut cx = Graph::new();
+    let a = cx.tensor((3, 4));
+    let _b = a.sin().permute((1, 0)).sqrt().output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    for r in &regions {
+        let has_sin = r.internal_ops_sorted.iter().any(|n| n == "FusedSin");
+        let has_sqrt = r.internal_ops_sorted.iter().any(|n| n == "FusedSqrt");
+        assert!(
+            !(has_sin && has_sqrt),
+            "permute between sin and sqrt must prevent them sharing a fused region, \
+             but found: {r:#?}"
+        );
+    }
+}
+
+#[test]
+fn test_reduction_prevents_unary_fusion() {
+    // A reduction between two unaries is not elementwise, so pair-fuse U→U
+    // (which only matches adjacent elementwise pairs) must not fire across
+    // the reduction.
+    let mut cx = Graph::new();
+    let a = cx.tensor((4, 4));
+    let _b = a.sin().sum(1).sqrt().output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    for r in &regions {
+        let has_sin = r.internal_ops_sorted.iter().any(|n| n == "FusedSin");
+        let has_sqrt = r.internal_ops_sorted.iter().any(|n| n == "FusedSqrt");
+        assert!(
+            !(has_sin && has_sqrt),
+            "reduction between sin and sqrt must prevent them sharing a fused region, \
+             but found: {r:#?}"
+        );
+    }
+}
+
+#[test]
+fn test_unary_fusion_preserves_output() {
+    // End-to-end numerical check: sqrt(sin(x)) must produce the same values
+    // whether or not the fusion rule fired. Runs on GPU when available;
+    // silently no-ops otherwise via get_cuda_stream().
+    let seed = 0xC0FFEEu64;
+    let gen_lambda = |n, s| random_f32_vec(n, s, 0.0, 1.0);
+    test_unary_cuda::<f32>(
+        8,
+        |a| a.sin().sqrt(),
+        |a| a.sin().unwrap().sqrt().unwrap(),
+        gen_lambda,
+        seed,
+    );
+}
+
+#[test]
+fn test_three_unary_ops_fuse() {
+    // A chain of 3 pure-elementwise unaries with matching strides should be
+    // reachable as a single marker region containing all three FusedX ops.
+    let mut cx = Graph::new();
+    let a = cx.tensor(16);
+    let _b = a.sin().sqrt().exp2().output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedSin", "FusedSqrt", "FusedExp2"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 1 && r.end_count == 1),
+        "expected a marker region of {expected:?} with 1 FusionStart, got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_four_unary_ops_fuse() {
+    // 4-op chain should collapse into a single marker region containing all
+    // four FusedX ops (one pair-fuse + repeated grow-FE→U firings).
+    let mut cx = Graph::new();
+    let a = cx.tensor(16);
+    let _b = a.sin().sqrt().exp2().log2().output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedSin", "FusedSqrt", "FusedExp2", "FusedLog2"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 1 && r.end_count == 1),
+        "expected a marker region of {expected:?} with 1 FusionStart, got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_three_unary_chain_preserves_output() {
+    // End-to-end numerical check for a 3-op chain.
+    // Uses sin→sqrt→sin because candle lacks exp2/log2 and this still exercises
+    // a 3-link chain. The structural tests above cover the distinct-ops shape.
+    let seed = 0xBEEFu64;
+    let gen_lambda = |n, s| random_f32_vec(n, s, 0.0, 1.0);
+    test_unary_cuda::<f32>(
+        16,
+        |a| a.sin().sqrt().sin(),
+        |a| a.sin().unwrap().sqrt().unwrap().sin().unwrap(),
+        gen_lambda,
+        seed,
+    );
+}
+
+/// Isolated per-kernel microbenchmark: time two unfused kernels
+/// (`sqrt_k` then `recip_k`) vs one fused kernel (`fused_k` that does
+/// `1.0f / sqrtf(x)` in a single launch) on a fixed-size input, using
+/// CUDA events for device-side timing.
+///
+/// Ignored by default — run with
+/// `cargo test -p luminal_cuda_lite -- --ignored bench_fused_vs_unfused_sqrt_recip --nocapture`.
+#[test]
+#[ignore]
+fn bench_fused_vs_unfused_sqrt_recip() {
+    use crate::compile_module_image_for_current_device;
+    use cudarc::driver::{CudaContext, LaunchConfig, PushKernelArg};
+
+    const N: usize = 1 << 20; // 1M elements
+    const WARMUP: usize = 100;
+    const TRIALS: usize = 2000;
+
+    let ctx = match CudaContext::new(0) {
+        Ok(c) => c,
+        Err(_) => return, // no GPU available, skip
+    };
+    ctx.bind_to_thread().unwrap();
+    let stream = ctx.default_stream();
+
+    // Prepare input (values in (0, 1] so sqrt/recip are well-defined).
+    let host_input: Vec<f32> = (0..N).map(|i| (i as f32 + 1.0) / (N as f32)).collect();
+    let d_in = stream.clone_htod(&host_input).unwrap();
+    let mut d_scratch = stream.alloc_zeros::<f32>(N).unwrap();
+    let mut d_out = stream.alloc_zeros::<f32>(N).unwrap();
+
+    let compile = |src: &str, name: &str| {
+        let ptx = compile_module_image_for_current_device(stream.context(), src).unwrap();
+        let module = stream.context().load_module(ptx).unwrap();
+        module.load_function(name).unwrap()
+    };
+
+    let sqrt_k = compile(
+        r#"
+extern "C" __global__ void sqrt_k(float* out, const float* in, long long n) {
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    out[i] = sqrtf(in[i]);
+}
+"#,
+        "sqrt_k",
+    );
+    let recip_k = compile(
+        r#"
+extern "C" __global__ void recip_k(float* out, const float* in, long long n) {
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    out[i] = 1.0f / in[i];
+}
+"#,
+        "recip_k",
+    );
+    let fused_k = compile(
+        r#"
+extern "C" __global__ void fused_k(float* out, const float* in, long long n) {
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    float v = in[i];
+    v = sqrtf(v);
+    v = 1.0f / v;
+    out[i] = v;
+}
+"#,
+        "fused_k",
+    );
+
+    let cfg = LaunchConfig::for_num_elems(N as u32);
+    let n_arg: i64 = N as i64;
+
+    let launch_unfused = |d_out: &mut cudarc::driver::CudaSlice<f32>,
+                          d_scratch: &mut cudarc::driver::CudaSlice<f32>| {
+        let mut b = stream.launch_builder(&sqrt_k);
+        b.arg(&mut *d_scratch).arg(&d_in).arg(&n_arg);
+        unsafe { b.launch(cfg) }.unwrap();
+        let mut b = stream.launch_builder(&recip_k);
+        b.arg(d_out).arg(&*d_scratch).arg(&n_arg);
+        unsafe { b.launch(cfg) }.unwrap();
+    };
+    let launch_fused = |d_out: &mut cudarc::driver::CudaSlice<f32>| {
+        let mut b = stream.launch_builder(&fused_k);
+        b.arg(d_out).arg(&d_in).arg(&n_arg);
+        unsafe { b.launch(cfg) }.unwrap();
+    };
+
+    // Warmup
+    for _ in 0..WARMUP {
+        launch_unfused(&mut d_out, &mut d_scratch);
+        launch_fused(&mut d_out);
+    }
+    stream.synchronize().unwrap();
+
+    let start = ctx.new_event(None).unwrap();
+    let end = ctx.new_event(None).unwrap();
+
+    // Time unfused
+    start.record(&stream).unwrap();
+    for _ in 0..TRIALS {
+        launch_unfused(&mut d_out, &mut d_scratch);
+    }
+    end.record(&stream).unwrap();
+    end.synchronize().unwrap();
+    let unfused_total_ms = start.elapsed_ms(&end).unwrap();
+
+    // Time fused
+    start.record(&stream).unwrap();
+    for _ in 0..TRIALS {
+        launch_fused(&mut d_out);
+    }
+    end.record(&stream).unwrap();
+    end.synchronize().unwrap();
+    let fused_total_ms = start.elapsed_ms(&end).unwrap();
+
+    let unfused_us = unfused_total_ms as f64 * 1_000.0 / TRIALS as f64;
+    let fused_us = fused_total_ms as f64 * 1_000.0 / TRIALS as f64;
+    let speedup = unfused_us / fused_us;
+
+    println!(
+        "\n[fusion microbench, N={N}, trials={TRIALS}]\n\
+         unfused (sqrt_k; recip_k):  {unfused_us:8.3} us/iter ({unfused_total_ms:.2} ms total)\n\
+         fused   (sqrtf; 1.0f/):     {fused_us:8.3} us/iter ({fused_total_ms:.2} ms total)\n\
+         speedup: {speedup:.2}x"
+    );
+}
+
+// =========================================================================
+// Binary-inclusive fusion tests (marker-based FusionStart / FusionEnd scheme).
+//
+// Detects fused regions by walking backward from each `FusionEnd`-tagged LLIR
+// node through `Direction::Incoming` edges until a `FusionStart` is reached.
+// The walker stops at FusionStarts (they mark the external-input boundary of
+// the region). A region's summary is: the sorted set of internal op names,
+// the count of distinct FusionStart nodes reached, and the count of FusionEnd
+// nodes (invariant: always 1 per region).
+// =========================================================================
+
+/// A single fused region extracted from the LLIR graph after egglog.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct FusedRegion {
+    /// Sorted internal op `kernel_name()`s, excluding the `FusionStart` /
+    /// `FusionEnd` markers. Sorted so DAG traversal order doesn't produce
+    /// spurious "distinct" regions.
+    internal_ops_sorted: Vec<String>,
+    /// Number of distinct `FusionStart` nodes reached by the walk. Per design
+    /// this equals the number of distinct external input tensors.
+    start_count: usize,
+    /// Number of `FusionEnd` nodes in the region. Per design this is always 1.
+    end_count: usize,
+}
+
+/// Helper: collect every distinct fused region reachable across many random
+/// extractions of the search space.
+fn extract_all_fused_regions(cx: &mut Graph) -> Vec<FusedRegion> {
+    cx.build_search_space::<CudaRuntime>();
+    let egraph = cx.egraph().expect("egraph not built");
+    let ops = cx.egglog_ops().expect("ops not built");
+    let custom_ops = &cx.custom_ops;
+
+    let mut seen: Vec<FusedRegion> = Vec::new();
+    // 200 samples: the random extractor picks one e-node per e-class per
+    // call, and the fully-fused diamond form lives in an e-class with
+    // many equivalent forms. 50 was flaky; 200 is reliably stable and
+    // each sample is cheap (~100 µs).
+    for _ in 0..200 {
+        let choices = random_initial_choice(egraph, &mut rand::rng());
+        let mut list_cache = Default::default();
+        let mut expr_cache = Default::default();
+        let llir = egglog_to_llir(
+            egraph,
+            choices,
+            ops,
+            custom_ops,
+            &mut list_cache,
+            &mut expr_cache,
+            None,
+        );
+
+        let name_of = |idx: NodeIndex| -> Option<String> {
+            llir.node_weight(idx).and_then(|op| {
+                op.to_dialect::<dyn KernelOp>()
+                    .map(|k| k.kernel_name().to_string())
+            })
+        };
+
+        let end_nodes: Vec<NodeIndex> = llir
+            .node_indices()
+            .filter(|&idx| name_of(idx).as_deref() == Some("FusionEnd"))
+            .collect();
+
+        for end in end_nodes {
+            let mut internal: Vec<String> = Vec::new();
+            // Count distinct external input *tensors*, not distinct FusionStart
+            // node indices. Egglog rule firings can emit multiple FusionStart
+            // enodes that all wrap the same source tensor (e.g. when the same
+            // `a` is consumed at two sites inside the fused region, each
+            // pair-fuse / grow firing mints its own FusionStart). Those are
+            // logically one FusionStart per the design invariant
+            // ("N = number of distinct external input tensors").
+            let mut start_sources: FxHashSet<NodeIndex> = FxHashSet::default();
+            let mut visited: FxHashSet<NodeIndex> = FxHashSet::default();
+            visited.insert(end);
+            let mut stack = vec![end];
+
+            // Resolve chains of nested FusionStart wrappers (cascade artifact)
+            // to the real external source. A FusionStart whose incoming neighbor
+            // is itself a FusionStart — or a FusionEnd whose region is fully
+            // inside ours — is a cascade layer, not a new external tensor.
+            let resolve_source = |mut n: NodeIndex| -> NodeIndex {
+                loop {
+                    match name_of(n).as_deref() {
+                        Some("FusionStart") | Some("FusionEnd") => {
+                            let mut inc = llir.neighbors_directed(n, petgraph::Direction::Incoming);
+                            match inc.next() {
+                                Some(p) => n = p,
+                                None => return n,
+                            }
+                        }
+                        _ => return n,
+                    }
+                }
+            };
+
+            while let Some(node) = stack.pop() {
+                for pred in llir.neighbors_directed(node, petgraph::Direction::Incoming) {
+                    if !visited.insert(pred) {
+                        continue;
+                    }
+                    match name_of(pred).as_deref() {
+                        Some("FusionStart") => {
+                            // If this FS's predecessor is itself a FE (or a
+                            // chain of FS/FE wrappers that eventually hits a
+                            // non-marker op inside the region), the FS is a
+                            // cascade artifact, not a real external boundary.
+                            // Walk past it and its upstream FE into the same
+                            // region. Otherwise treat the predecessor as the
+                            // external source tensor — which may be a KernelOp
+                            // *or* a non-KernelOp (HLIR loadable) node, so we
+                            // can't gate counting on `name_of` being `Some`.
+                            let mut inc =
+                                llir.neighbors_directed(pred, petgraph::Direction::Incoming);
+                            match inc.next() {
+                                Some(src_node)
+                                    if name_of(src_node).as_deref() == Some("FusionEnd") =>
+                                {
+                                    // Merge adjacent regions — treat the FS/FE
+                                    // pair as internal; walk past the upstream
+                                    // FE into its region.
+                                    visited.insert(src_node);
+                                    stack.push(src_node);
+                                }
+                                Some(src_node) => {
+                                    start_sources.insert(resolve_source(src_node));
+                                }
+                                None => {
+                                    // FS with no predecessor — degenerate.
+                                }
+                            }
+                        }
+                        Some("FusionEnd") => {
+                            // Transparent: inner FusionEnds are cascade-wart
+                            // artifacts from grow rules re-firing and creating
+                            // nested `FE(Op(FE(...)))` wrappers. They don't
+                            // represent real work or a real boundary — walk
+                            // past them and do not count them as internal ops.
+                            stack.push(pred);
+                        }
+                        Some(other) => {
+                            internal.push(other.to_string());
+                            stack.push(pred);
+                        }
+                        None => {
+                            // Non-KernelOp predecessor (shouldn't appear inside a
+                            // fused region under the design). Stop walking this path.
+                        }
+                    }
+                }
+            }
+
+            internal.sort();
+            // Skip singleton regions: every elementwise op has a seeded
+            // `FE(Op(FS(...)))` form, so random extraction will surface
+            // many one-op regions that are equivalent to not fusing. We
+            // only care about regions that represent real multi-op fusion.
+            if internal.len() < 2 {
+                continue;
+            }
+            let region = FusedRegion {
+                internal_ops_sorted: internal,
+                start_count: start_sources.len(),
+                end_count: 1,
+            };
+            if !seen.contains(&region) {
+                seen.push(region);
+            }
+        }
+    }
+    seen
+}
+
+fn sorted_names(items: &[&str]) -> Vec<String> {
+    let mut v: Vec<String> = items.iter().map(|s| (*s).to_string()).collect();
+    v.sort();
+    v
+}
+
+// ---- Structural tests: the expected fused shape is reachable ----
+
+#[test]
+fn test_single_binary_does_not_fuse_alone() {
+    // A lone elementwise op gets a seeded singleton region by design; we
+    // filter singletons out in `extract_all_fused_regions`. What this test
+    // asserts is that no *multi-op* region appears for a standalone binary
+    // — nothing to grow into.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let _c = (a + b).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    assert!(
+        regions.is_empty(),
+        "a solo binary op should not form a multi-op fused region, but got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_chain_of_binaries_fuses() {
+    // `(a + b) * c`: three external inputs collapse into one region with
+    // internal [Add, Mul] and 3 FusionStarts.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let c = cx.tensor(8);
+    let _d = ((a + b) * c).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedMul"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 3),
+        "expected a fused region of {expected:?} with 3 FusionStarts, got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_binary_then_unary_fuses() {
+    // `sin(a + b)`: binary feeds a unary inside one fused region.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let _c = (a + b).sin().output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedSin"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 2),
+        "expected a fused region of {expected:?} with 2 FusionStarts, got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_unary_then_binary_fuses() {
+    // `sin(a) + b`: unary feeds a binary inside one fused region.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let _c = (a.sin() + b).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedSin"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 2),
+        "expected a fused region of {expected:?} with 2 FusionStarts, got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_diamond_dag_fuses() {
+    // The canonical diamond-DAG example agreed with the user:
+    //   t = a + b; u = exp2(t); v = sin(t); w = u * a; out = w + v
+    // `a` is reused (feeds outer Add and Mul) and `t` is reused (feeds Exp2 and
+    // Sin). Expected: one fused region with internal ops [Add, Add, Exp2, Mul,
+    // Sin], 2 FusionStarts (distinct tensors a, b), 1 FusionEnd.
+    // We use exp2 rather than exp because the frontend's exp() desugars to
+    // Mul(x, LOG2E).exp2(), which would add a constant input and a Mul op and
+    // obscure the diamond topology this test is checking.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let t = a + b;
+    let u = t.exp2();
+    let v = t.sin();
+    let w = u * a;
+    let _out = (w + v).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedAdd", "FusedExp2", "FusedMul", "FusedSin"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 2 && r.end_count == 1),
+        "expected diamond DAG to fuse into one region with ops {expected:?}, \
+         2 FusionStarts, 1 FusionEnd. Got: {regions:#?}"
+    );
+}
+
+// ---- Negative tests: fusion must NOT happen across these blockers ----
+
+#[test]
+fn test_reduction_blocks_binary_fusion() {
+    // A reduction between a binary and anything downstream is not elementwise,
+    // so Add and SumReduce must never appear in the same fused region.
+    let mut cx = Graph::new();
+    let a = cx.tensor((4, 4));
+    let b = cx.tensor((4, 4));
+    let _c = (a + b).sum(1).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    for r in &regions {
+        let has_add = r.internal_ops_sorted.iter().any(|n| n == "FusedAdd");
+        let has_sum = r.internal_ops_sorted.iter().any(|n| n == "SumReduce");
+        assert!(
+            !(has_add && has_sum),
+            "FusedAdd and SumReduce must not share a fused region, but got: {r:#?}"
+        );
+    }
+}
+
+#[test]
+fn test_stride_mismatch_blocks_binary_fusion() {
+    // A permute gives `b` a non-contiguous view whose strides do not match `a`'s,
+    // so the binary fusion rule's stride-compatibility check must prevent the
+    // Add from being absorbed into any fused region.
+    let mut cx = Graph::new();
+    let a = cx.tensor((3, 4));
+    let b = cx.tensor((4, 3));
+    let _c = (a + b.permute((1, 0))).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    for r in &regions {
+        assert!(
+            !r.internal_ops_sorted.iter().any(|n| n == "FusedAdd"),
+            "permuted binary must not fuse into a region, but found: {r:#?}"
+        );
+    }
+}
+
+// ---- Numerical parity tests: fused output matches candle reference ----
+
+#[test]
+fn test_simple_binary_fusion_preserves_output() {
+    // End-to-end numerical check: `a + b` on GPU matches candle's add across
+    // all reachable genomes (fused or unfused) via test_binary_cuda's fuzzer.
+    let seed = 0xADDBEEFu64;
+    let eps = dtype_epsilon(luminal::dtype::DType::F32);
+    let tol = eps * TOLERANCE_SAFETY_FACTOR;
+    test_binary_cuda::<f32>(
+        16,
+        16,
+        |a, b| a + b,
+        |a, b| (a + b).unwrap(),
+        |n, s| random_f32_vec(n, s, 0.0, 1.0),
+        |n, s| random_f32_vec(n, s, 0.0, 1.0),
+        seed,
+        tol,
+        tol,
+    );
+}
+
+#[test]
+fn test_diamond_dag_preserves_output() {
+    // Numerical parity for the diamond DAG: `(exp(a+b) * a) + sin(a+b)`
+    // matches candle's equivalent across fused and unfused genomes.
+    // Inputs are drawn from [-1, 1] so exp() doesn't overflow.
+    let seed = 0xD1A_0D1Au64;
+    let eps = dtype_epsilon(luminal::dtype::DType::F32);
+    // Five-op chain with exp + sin: allow ~5x safety to absorb accumulated
+    // rounding vs candle's kernels.
+    let tol = eps * TOLERANCE_SAFETY_FACTOR * 5.0;
+    test_binary_cuda::<f32>(
+        16,
+        16,
+        |a, b| {
+            let t = a + b;
+            let u = t.exp();
+            let v = t.sin();
+            let w = u * a;
+            w + v
+        },
+        |a, b| {
+            let t = (&a + &b).unwrap();
+            let u = t.exp().unwrap();
+            let v = t.sin().unwrap();
+            let w = (&u * &a).unwrap();
+            (&w + &v).unwrap()
+        },
+        |n, s| random_f32_vec(n, s, -1.0, 1.0),
+        |n, s| random_f32_vec(n, s, -1.0, 1.0),
+        seed,
+        tol,
+        tol,
+    );
+}
+
+// ---- Marker invariant tests ----
+
+#[test]
+fn test_fused_region_has_exactly_one_end() {
+    // Design invariant: a fused region always has exactly one FusionEnd.
+    // Uses the diamond DAG so there's real fan-in/out inside the region.
+    // See test_diamond_dag_fuses for why we use exp2 directly.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let t = a + b;
+    let u = t.exp2();
+    let v = t.sin();
+    let w = u * a;
+    let _out = (w + v).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedAdd", "FusedExp2", "FusedMul", "FusedSin"]);
+    let full = regions
+        .iter()
+        .find(|r| r.internal_ops_sorted == expected)
+        .expect("expected at least one extraction to produce the full 5-op diamond region");
+    assert_eq!(
+        full.end_count, 1,
+        "fused region must have exactly one FusionEnd, got {}",
+        full.end_count
+    );
+}
+
+#[test]
+fn test_fused_region_starts_match_distinct_external_tensors() {
+    // Design invariant: FusionStart count == number of distinct external input
+    // tensors, NOT number of edges crossing the boundary. In the diamond DAG
+    // `a` is consumed inside the region by two ops (outer Add + Mul), so a
+    // per-edge counting scheme would give 3; the correct per-distinct-tensor
+    // count is 2 ({a, b}).
+    // See test_diamond_dag_fuses for why we use exp2 directly.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let t = a + b;
+    let u = t.exp2();
+    let v = t.sin();
+    let w = u * a;
+    let _out = (w + v).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedAdd", "FusedExp2", "FusedMul", "FusedSin"]);
+    // Multiple 5-op extractions are reachable: the merge-FE-FE rule fires
+    // across paths that may have minted distinct FS enodes for the shared
+    // tensor `a` at separate sites. The design invariant is that *some*
+    // extraction collapses those into the deduped form (one FS per distinct
+    // tensor → 2 FS for {a, b}); we don't require every random sample to.
+    let matching: Vec<&FusedRegion> = regions
+        .iter()
+        .filter(|r| r.internal_ops_sorted == expected)
+        .collect();
+    assert!(
+        !matching.is_empty(),
+        "expected at least one extraction to produce the full 5-op diamond region, \
+         got: {regions:#?}"
+    );
+    assert!(
+        matching
+            .iter()
+            .any(|r| r.start_count == 2 && r.end_count == 1),
+        "expected at least one 5-op diamond extraction with FusionStart count == 2 \
+         (one per distinct external tensor) and FusionEnd count == 1; got: {matching:#?}"
+    );
+}
+
+// ---- Targeted rule-family tests (one per family / orientation) ----
+//
+// The structural and diamond tests above hit several rule families at once.
+// These narrow tests pin each rule family / orientation independently so a
+// regression in one rule shows up as a single failing test rather than a
+// confusing diamond mismatch.
+
+#[test]
+fn test_pair_fuse_unary_unary_marker_form() {
+    // Pair-fuse U→U: `a.sin().sqrt()` should be reachable as a marker-bracketed
+    // region containing FusedSin and FusedSqrt (with one FusionStart for `a`).
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let _b = a.sin().sqrt().output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedSin", "FusedSqrt"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 1 && r.end_count == 1),
+        "expected marker region of {expected:?} with 1 FusionStart, got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_pair_fuse_unary_to_binary_rhs() {
+    // Pair-fuse U→B (RHS variant): `a + b.sin()`. The unary is on the
+    // binary's B input, so the rule's RHS-orientation version is what fires.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let _c = (a + b.sin()).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedSin"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 2),
+        "expected a fused region of {expected:?} with 2 FusionStarts (RHS-side unary), \
+         got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_pair_fuse_binary_to_binary_rhs() {
+    // Pair-fuse B→B (RHS variant): `c * (a + b)`. The inner binary feeds the
+    // outer binary's B input, exercising the mirror direction of the rule
+    // covered by test_chain_of_binaries_fuses.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let c = cx.tensor(8);
+    let _d = (c * (a + b)).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedMul"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 3),
+        "expected a fused region of {expected:?} with 3 FusionStarts (RHS-side inner binary), \
+         got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_grow_fe_to_binary_rhs() {
+    // Grow FE→B (RHS variant): `c + (a.sin() + b)`. Once the inner
+    // `a.sin() + b` is fused, the outer `+ c` consumes that FE on its B input
+    // (because we wrote `c + (...)` — `c` is on LHS, FE on RHS), exercising
+    // grow-FE-B-rhs to absorb the outer Add into the same region.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let c = cx.tensor(8);
+    let _d = (c + (a.sin() + b)).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedAdd", "FusedSin"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 3),
+        "expected a 3-op fused region of {expected:?} with 3 FusionStarts (grow into RHS), \
+         got: {regions:#?}"
+    );
+}
+
+#[test]
+fn test_merge_two_regions_at_outer_binary() {
+    // Merge: `(sin(a) + b) + (sqrt(c) + d)`. Each side independently pair-fuses
+    // U→B on its own (the unary gives the inner Add a fusion partner that
+    // doesn't pull in the outer Add), so both sides become FEs. The outer Add
+    // then fires merge-FE-FE-Add to collapse them into a single region.
+    // Without the unaries, `(a+b) + (c+d)` would only ever pair-fuse one
+    // inner Add at a time with the outer Add — merge wouldn't have two FEs to
+    // combine because the inner Adds never become singleton FEs on their own.
+    let mut cx = Graph::new();
+    let a = cx.tensor(8);
+    let b = cx.tensor(8);
+    let c = cx.tensor(8);
+    let d = cx.tensor(8);
+    let _e = ((a.sin() + b) + (c.sqrt() + d)).output();
+
+    let regions = extract_all_fused_regions(&mut cx);
+    let expected = sorted_names(&["FusedAdd", "FusedAdd", "FusedAdd", "FusedSin", "FusedSqrt"]);
+    assert!(
+        regions
+            .iter()
+            .any(|r| r.internal_ops_sorted == expected && r.start_count == 4),
+        "expected a 5-op merged region (two pair-fused sides combined at outer Add) with \
+         4 FusionStarts, got: {regions:#?}"
+    );
+}
+
+/// Microbench: time three unfused kernels (`add_k` → `sin_k` → `sqrt_k`)
+/// vs one fused kernel (`(a + b).sin().sqrt()` in a single launch) on a
+/// fixed-size input, using CUDA events for device-side timing. Mirrors
+/// the existing sqrt→recip bench but on the binary-inclusive 3-op DAG
+/// PR2's region codegen targets.
+///
+/// Ignored by default — run with
+/// `cargo test -p luminal_cuda_lite -- --ignored bench_fused_region_vs_unfused_3op --nocapture`.
+#[test]
+#[ignore]
+fn bench_fused_region_vs_unfused_3op() {
+    use crate::compile_module_image_for_current_device;
+    use cudarc::driver::{CudaContext, LaunchConfig, PushKernelArg};
+
+    const N: usize = 1 << 20; // 1M elements
+    const WARMUP: usize = 100;
+    const TRIALS: usize = 2000;
+
+    let ctx = match CudaContext::new(0) {
+        Ok(c) => c,
+        Err(_) => return, // no GPU available, skip
+    };
+    ctx.bind_to_thread().unwrap();
+    let stream = ctx.default_stream();
+
+    // Inputs in (0, 1] keep `sin` < 1 and `sqrt` well-defined post-add.
+    let host_a: Vec<f32> = (0..N)
+        .map(|i| (i as f32 + 1.0) / (N as f32) * 0.5)
+        .collect();
+    let host_b: Vec<f32> = (0..N)
+        .map(|i| (i as f32 + 1.0) / (N as f32) * 0.5)
+        .collect();
+    let d_a = stream.clone_htod(&host_a).unwrap();
+    let d_b = stream.clone_htod(&host_b).unwrap();
+    let mut d_scratch1 = stream.alloc_zeros::<f32>(N).unwrap();
+    let mut d_scratch2 = stream.alloc_zeros::<f32>(N).unwrap();
+    let mut d_out = stream.alloc_zeros::<f32>(N).unwrap();
+
+    let compile = |src: &str, name: &str| {
+        let ptx = compile_module_image_for_current_device(stream.context(), src).unwrap();
+        let module = stream.context().load_module(ptx).unwrap();
+        module.load_function(name).unwrap()
+    };
+
+    let add_k = compile(
+        r#"
+extern "C" __global__ void add_k(float* out, const float* a, const float* b, long long n) {
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    out[i] = a[i] + b[i];
+}
+"#,
+        "add_k",
+    );
+    let sin_k = compile(
+        r#"
+extern "C" __global__ void sin_k(float* out, const float* in, long long n) {
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    out[i] = sinf(in[i]);
+}
+"#,
+        "sin_k",
+    );
+    let sqrt_k = compile(
+        r#"
+extern "C" __global__ void sqrt_k(float* out, const float* in, long long n) {
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    out[i] = sqrtf(in[i]);
+}
+"#,
+        "sqrt_k",
+    );
+    let fused_k = compile(
+        r#"
+extern "C" __global__ void fused_k(float* out, const float* a, const float* b, long long n) {
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    float v = a[i] + b[i];
+    v = sinf(v);
+    v = sqrtf(v);
+    out[i] = v;
+}
+"#,
+        "fused_k",
+    );
+
+    let cfg = LaunchConfig::for_num_elems(N as u32);
+    let n_arg: i64 = N as i64;
+
+    let launch_unfused =
+        |d_out: &mut cudarc::driver::CudaSlice<f32>,
+         d_scratch1: &mut cudarc::driver::CudaSlice<f32>,
+         d_scratch2: &mut cudarc::driver::CudaSlice<f32>| {
+            let mut b = stream.launch_builder(&add_k);
+            b.arg(&mut *d_scratch1).arg(&d_a).arg(&d_b).arg(&n_arg);
+            unsafe { b.launch(cfg) }.unwrap();
+            let mut b = stream.launch_builder(&sin_k);
+            b.arg(&mut *d_scratch2).arg(&*d_scratch1).arg(&n_arg);
+            unsafe { b.launch(cfg) }.unwrap();
+            let mut b = stream.launch_builder(&sqrt_k);
+            b.arg(d_out).arg(&*d_scratch2).arg(&n_arg);
+            unsafe { b.launch(cfg) }.unwrap();
+        };
+    let launch_fused = |d_out: &mut cudarc::driver::CudaSlice<f32>| {
+        let mut b = stream.launch_builder(&fused_k);
+        b.arg(d_out).arg(&d_a).arg(&d_b).arg(&n_arg);
+        unsafe { b.launch(cfg) }.unwrap();
+    };
+
+    // Warmup
+    for _ in 0..WARMUP {
+        launch_unfused(&mut d_out, &mut d_scratch1, &mut d_scratch2);
+        launch_fused(&mut d_out);
+    }
+    stream.synchronize().unwrap();
+
+    // Host-side wall-clock timing: synchronize before/after each batch so the
+    // measured interval covers exactly the GPU work for `TRIALS` iterations.
+    // (CUDA event-based timing is the more precise option in principle, but
+    // `event.elapsed_ms` on this driver/cudarc combo errors with
+    // CUDA_ERROR_INVALID_HANDLE — see bench_fused_vs_unfused_sqrt_recip
+    // above which fails the same way. Wall-clock is reliable here.)
+    let unfused_start = std::time::Instant::now();
+    for _ in 0..TRIALS {
+        launch_unfused(&mut d_out, &mut d_scratch1, &mut d_scratch2);
+    }
+    stream.synchronize().unwrap();
+    let unfused_total_ms = unfused_start.elapsed().as_secs_f64() * 1_000.0;
+
+    let fused_start = std::time::Instant::now();
+    for _ in 0..TRIALS {
+        launch_fused(&mut d_out);
+    }
+    stream.synchronize().unwrap();
+    let fused_total_ms = fused_start.elapsed().as_secs_f64() * 1_000.0;
+
+    let unfused_us = unfused_total_ms * 1_000.0 / TRIALS as f64;
+    let fused_us = fused_total_ms * 1_000.0 / TRIALS as f64;
+    let speedup = unfused_us / fused_us;
+
+    println!(
+        "\n[fusion microbench, (a+b).sin().sqrt(), N={N}, trials={TRIALS}]\n\
+         unfused (add_k; sin_k; sqrt_k): {unfused_us:8.3} us/iter ({unfused_total_ms:.2} ms total)\n\
+         fused   (one kernel):           {fused_us:8.3} us/iter ({fused_total_ms:.2} ms total)\n\
+         speedup: {speedup:.2}x"
+    );
+}

crates/luminal_cuda_lite/src/tests/mod.rs

@@ -5,10 +5,14 @@ mod bucket_tests;
 #[cfg(test)]
 mod consumed_buffer_tests;
 #[cfg(test)]
+mod fusion;
+#[cfg(test)]
 mod model_fuzz;
 #[cfg(test)]
 mod op_functional_tests;
 #[cfg(test)]
 mod performance_tests;
 #[cfg(test)]
+mod qwen3_moe_rewrite;
+#[cfg(test)]
 mod transformer;

crates/luminal_cuda_lite/src/tests/qwen3_moe_rewrite.rs

@@ -0,0 +1,314 @@
+use half::bf16;
+use luminal::{dtype::DType, prelude::*, shape::Expression};
+
+use super::utilities::{assert_close, get_cuda_stream, random_f32_vec};
+use crate::{
+    host::{
+        HostOp,
+        moe::{GLUMoE, GLUMoEMode},
+    },
+    runtime::CudaRuntime,
+};
+
+const SEQ: usize = 2;
+const HIDDEN: usize = 16;
+const NUM_EXPERTS: usize = 8;
+const TOP_K: usize = 2;
+const MOE_INTERMEDIATE: usize = 6;
+const RMS_NORM_EPS: f32 = 1e-6;
+
+struct QwenMoeGraph {
+    graph: Graph,
+    x: GraphTensor,
+    router: GraphTensor,
+    gate_up_weights: GraphTensor,
+    down_weights: GraphTensor,
+    output: GraphTensor,
+}
+
+struct GemmaMoeGraph {
+    graph: Graph,
+    router_input: GraphTensor,
+    expert_input: GraphTensor,
+    router_scale: GraphTensor,
+    router_proj: GraphTensor,
+    per_expert_scale: GraphTensor,
+    gate_up_weights: GraphTensor,
+    down_weights: GraphTensor,
+    output: GraphTensor,
+}
+
+fn build_qwen_moe_graph() -> QwenMoeGraph {
+    let mut cx = Graph::default();
+    let x = cx.tensor(('s', HIDDEN));
+    let router = cx.tensor((NUM_EXPERTS, HIDDEN));
+    let gate_up_weights = cx
+        .tensor((NUM_EXPERTS, MOE_INTERMEDIATE * 2, HIDDEN))
+        .as_dtype(DType::Bf16);
+    let down_weights = cx
+        .tensor((NUM_EXPERTS, HIDDEN, MOE_INTERMEDIATE))
+        .as_dtype(DType::Bf16);
+
+    let n = x.dims().len();
+    let e_dim = *router.dims().first().unwrap();
+    let k_expr = Expression::from(TOP_K);
+
+    let routing_weights = x.matmul(router.t()).softmax(n - 1);
+    let top_k_indices = routing_weights.topk_indexes(TOP_K, n - 1);
+
+    let row_offsets = x
+        .graph()
+        .iota(Expression::from('z') / k_expr * e_dim, top_k_indices.dims());
+    let routing_flat_idx = row_offsets + top_k_indices;
+    let top_k_values = routing_weights.gather(routing_flat_idx);
+
+    let gate_up_gathered = gather_experts(x, top_k_indices, gate_up_weights).cast(DType::F32);
+    let x_exp = x.expand_dim(n - 1, TOP_K).unsqueeze(n);
+    let gate_up_out = x_exp.matmul(gate_up_gathered.transpose(2, 3)).squeeze(n);
+    let gate = gate_up_out.slice((.., .., ..MOE_INTERMEDIATE));
+    let up = gate_up_out.slice((.., .., MOE_INTERMEDIATE..));
+    let hidden = gate.silu() * up;
+
+    let down_gathered = gather_experts(x, top_k_indices, down_weights).cast(DType::F32);
+    let down_out = hidden
+        .unsqueeze(2)
+        .matmul(down_gathered.transpose(2, 3))
+        .squeeze(2);
+    let output = (down_out * top_k_values.unsqueeze(top_k_values.dims().len()))
+        .sum(n - 1)
+        .output();
+
+    QwenMoeGraph {
+        graph: cx,
+        x,
+        router,
+        gate_up_weights,
+        down_weights,
+        output,
+    }
+}
+
+fn build_gemma_moe_graph() -> GemmaMoeGraph {
+    let mut cx = Graph::default();
+    let router_input = cx.tensor(('s', HIDDEN));
+    let expert_input = cx.tensor(('s', HIDDEN));
+    let router_scale = cx.tensor(HIDDEN);
+    let router_proj = cx.tensor((NUM_EXPERTS, HIDDEN));
+    let per_expert_scale = cx.tensor(NUM_EXPERTS);
+    let gate_up_weights = cx
+        .tensor((NUM_EXPERTS, MOE_INTERMEDIATE * 2, HIDDEN))
+        .as_dtype(DType::Bf16);
+    let down_weights = cx
+        .tensor((NUM_EXPERTS, HIDDEN, MOE_INTERMEDIATE))
+        .as_dtype(DType::Bf16);
+
+    let n = router_input.dims().len();
+    let e_dim = *router_proj.dims().first().unwrap();
+    let k_expr = Expression::from(TOP_K);
+
+    let router_hidden = router_input.std_norm(n - 1, RMS_NORM_EPS)
+        * router_scale.expand_lhs(&router_input.dims()[..n - 1])
+        * (HIDDEN as f32).sqrt().recip();
+    let routing_weights = router_hidden.matmul(router_proj.t()).softmax(n - 1);
+
+    let top_k_indices = routing_weights.topk_indexes(TOP_K, n - 1);
+    let row_offsets = router_input
+        .graph()
+        .iota(Expression::from('z') / k_expr * e_dim, top_k_indices.dims());
+    let routing_flat_idx = row_offsets + top_k_indices;
+    let top_k_values = routing_weights.gather(routing_flat_idx);
+    let top_k_norm = top_k_values.sum(n - 1).expand_dim(n - 1, TOP_K);
+    let top_k_weights = (top_k_values / top_k_norm) * per_expert_scale.gather(top_k_indices);
+
+    let gate_up_gathered =
+        gather_experts(expert_input, top_k_indices, gate_up_weights).cast(DType::F32);
+    let x_exp = expert_input.expand_dim(n - 1, TOP_K).unsqueeze(n);
+    let gate_up_out = x_exp.matmul(gate_up_gathered.transpose(2, 3)).squeeze(n);
+    let gate = gate_up_out.slice((.., .., ..MOE_INTERMEDIATE));
+    let up = gate_up_out.slice((.., .., MOE_INTERMEDIATE..));
+    let hidden = gemma_gelu(gate) * up;
+
+    let down_gathered = gather_experts(expert_input, top_k_indices, down_weights).cast(DType::F32);
+    let down_out = hidden
+        .unsqueeze(2)
+        .matmul(down_gathered.transpose(2, 3))
+        .squeeze(2);
+    let output = (down_out * top_k_weights.unsqueeze(top_k_weights.dims().len()))
+        .sum(n - 1)
+        .output();
+
+    GemmaMoeGraph {
+        graph: cx,
+        router_input,
+        expert_input,
+        router_scale,
+        router_proj,
+        per_expert_scale,
+        gate_up_weights,
+        down_weights,
+        output,
+    }
+}
+
+fn gather_experts(
+    graph_source: GraphTensor,
+    top_k_indices: GraphTensor,
+    weights: GraphTensor,
+) -> GraphTensor {
+    let (_, d1, d2) = weights.dims3();
+    let io = d1 * d2;
+    let base = top_k_indices * io;
+    let within = graph_source.graph().iota(Expression::from('z'), (d1, d2));
+    let n_base = base.dims().len();
+    let exp_base = base.expand_dim(n_base, d1).expand_dim(n_base + 1, d2);
+    let mut exp_within = within;
+    for (axis, dim) in base.dims().iter().enumerate() {
+        exp_within = exp_within.expand_dim(axis, *dim);
+    }
+    let expert_flat_idx = exp_base + exp_within;
+    weights.gather(expert_flat_idx)
+}
+
+#[allow(clippy::excessive_precision)]
+fn gemma_gelu(x: GraphTensor) -> GraphTensor {
+    let scaled = 1.5957691216 * x * (1. + 0.044715 * x * x);
+    x * scaled.sigmoid()
+}
+
+fn glumoe_modes(rt: &CudaRuntime) -> Vec<GLUMoEMode> {
+    rt.llir_graph()
+        .node_weights()
+        .filter_map(|node| {
+            let op = node.to_dialect::<dyn HostOp>()?;
+            op.as_any()
+                .downcast_ref::<GLUMoE>()
+                .map(|glumoe| glumoe.mode)
+        })
+        .collect()
+}
+
+fn run_qwen_moe(use_glumoe: bool) -> (Vec<f32>, Vec<GLUMoEMode>) {
+    let Some(stream) = get_cuda_stream() else {
+        return (vec![], vec![]);
+    };
+
+    let mut model = build_qwen_moe_graph();
+    model.graph.set_dim('s', SEQ);
+    if use_glumoe {
+        model.graph.build_search_space::<CudaRuntime>();
+    } else {
+        model
+            .graph
+            .build_search_space_exclude_ops::<CudaRuntime, GLUMoE>();
+    }
+
+    let x_data = random_f32_vec(SEQ * HIDDEN, 11, -0.15, 0.15);
+    let router_data = random_f32_vec(NUM_EXPERTS * HIDDEN, 12, -0.2, 0.2);
+    let gate_up_data = random_f32_vec(NUM_EXPERTS * MOE_INTERMEDIATE * 2 * HIDDEN, 13, -0.1, 0.1)
+        .into_iter()
+        .map(bf16::from_f32)
+        .collect::<Vec<_>>();
+    let down_data = random_f32_vec(NUM_EXPERTS * HIDDEN * MOE_INTERMEDIATE, 14, -0.1, 0.1)
+        .into_iter()
+        .map(bf16::from_f32)
+        .collect::<Vec<_>>();
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.set_data(model.x, x_data);
+    rt.set_data(model.router, router_data);
+    rt.set_data(model.gate_up_weights, gate_up_data);
+    rt.set_data(model.down_weights, down_data);
+    rt = model.graph.search(rt, 10);
+    rt.execute(&model.graph.dyn_map);
+
+    (rt.get_f32(model.output.id), glumoe_modes(&rt))
+}
+
+fn run_gemma_moe(use_glumoe: bool) -> (Vec<f32>, Vec<GLUMoEMode>) {
+    let Some(stream) = get_cuda_stream() else {
+        return (vec![], vec![]);
+    };
+
+    let mut model = build_gemma_moe_graph();
+    model.graph.set_dim('s', SEQ);
+    if use_glumoe {
+        model.graph.build_search_space::<CudaRuntime>();
+    } else {
+        model
+            .graph
+            .build_search_space_exclude_ops::<CudaRuntime, GLUMoE>();
+    }
+
+    let router_input_data = random_f32_vec(SEQ * HIDDEN, 21, -0.15, 0.15);
+    let expert_input_data = random_f32_vec(SEQ * HIDDEN, 22, -0.15, 0.15);
+    let router_scale_data = random_f32_vec(HIDDEN, 23, 0.7, 1.3);
+    let router_proj_data = random_f32_vec(NUM_EXPERTS * HIDDEN, 24, -0.2, 0.2);
+    let per_expert_scale_data = random_f32_vec(NUM_EXPERTS, 25, 0.5, 1.5);
+    let gate_up_data = random_f32_vec(NUM_EXPERTS * MOE_INTERMEDIATE * 2 * HIDDEN, 26, -0.1, 0.1)
+        .into_iter()
+        .map(bf16::from_f32)
+        .collect::<Vec<_>>();
+    let down_data = random_f32_vec(NUM_EXPERTS * HIDDEN * MOE_INTERMEDIATE, 27, -0.1, 0.1)
+        .into_iter()
+        .map(bf16::from_f32)
+        .collect::<Vec<_>>();
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.set_data(model.router_input, router_input_data);
+    rt.set_data(model.expert_input, expert_input_data);
+    rt.set_data(model.router_scale, router_scale_data);
+    rt.set_data(model.router_proj, router_proj_data);
+    rt.set_data(model.per_expert_scale, per_expert_scale_data);
+    rt.set_data(model.gate_up_weights, gate_up_data);
+    rt.set_data(model.down_weights, down_data);
+    rt = model.graph.search(rt, 10);
+    rt.execute(&model.graph.dyn_map);
+
+    (rt.get_f32(model.output.id), glumoe_modes(&rt))
+}
+
+#[test]
+fn test_glumoe_matches_qwen_swiglu_pattern() {
+    let (_result, modes) = run_qwen_moe(true);
+    if modes.is_empty() {
+        return;
+    }
+
+    assert_eq!(modes, vec![GLUMoEMode::SwiGLU]);
+}
+
+#[test]
+fn test_glumoe_matches_gemma_gelu_pattern() {
+    let (_result, modes) = run_gemma_moe(true);
+    if modes.is_empty() {
+        return;
+    }
+
+    assert_eq!(modes, vec![GLUMoEMode::GemmaGELU]);
+}
+
+#[test]
+fn test_glumoe_swiglu_matches_unfused_output() {
+    let (expected, baseline_modes) = run_qwen_moe(false);
+    if expected.is_empty() {
+        return;
+    }
+    assert!(baseline_modes.is_empty());
+
+    let (actual, fused_modes) = run_qwen_moe(true);
+    assert_eq!(fused_modes, vec![GLUMoEMode::SwiGLU]);
+    assert_close(&actual, &expected, 3e-2, 3e-2);
+}
+
+#[test]
+fn test_glumoe_gemma_gelu_matches_unfused_output() {
+    let (expected, baseline_modes) = run_gemma_moe(false);
+    if expected.is_empty() {
+        return;
+    }
+    assert!(baseline_modes.is_empty());
+
+    let (actual, fused_modes) = run_gemma_moe(true);
+    assert_eq!(fused_modes, vec![GLUMoEMode::GemmaGELU]);
+    assert_close(&actual, &expected, 3e-2, 3e-2);
+}

crates/luminal_cuda_lite/src/tests/transformer.rs

@@ -300,7 +300,7 @@ fn test_mini_transformer_two_layers() {
     let input = cx.tensor((SEQ, HIDDEN));
     let layer1 = MiniTransformerLayer::init(&mut cx);
     let layer2 = MiniTransformerLayer::init(&mut cx);
-    let x = layer1.forward(input).graph_break();
+    let x = layer1.forward(input);
     let out = layer2.forward(x).output();
 
     cx.build_search_space::<CudaRuntime>();
@@ -508,3 +508,32 @@ fn test_swiglu_mlp_cuda() {
 
     assert_close(&result, &expected, 1e-3, 1e-3);
 }
+
+/// Body=1, trips=3 chain of scalar Muls plus a residual back to the
+/// chain's initial value. Auto-rolling sees this as a state-carrying loop
+/// with state at input position 0; the rolled HLIR must round-trip through
+/// egglog (rolled body Mul + LoopStart/LoopInput/LoopEnd markers) and
+/// `unroll_loops_in_llir` must reconstruct the flat 3-mul chain plus
+/// rewire the residual edge to reference the chain's initial input
+/// (outside the body) — not a per-iter clone.
+#[test]
+fn test_rolled_chained_scalar_muls() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+    let mut cx = Graph::default();
+    let x = cx.tensor((1, 4, 32));
+    let chained = ((x * 2.0_f32) * 3.0_f32) * 5.0_f32;
+    let out = (chained + x).output();
+
+    cx.build_search_space::<CudaRuntime>();
+    let mut rt = CudaRuntime::initialize(stream);
+    let x_data = random_f32_vec(4 * 32, 101, -0.5, 0.5);
+    rt.set_data(x, x_data.clone());
+    rt = cx.search(rt, 3);
+    rt.execute(&cx.dyn_map);
+
+    let result = rt.get_f32(out);
+    let expected: Vec<f32> = x_data.iter().map(|v| v * 2.0 * 3.0 * 5.0 + v).collect();
+    assert_close(&result, &expected, 1e-5, 1e-5);
+}

crates/luminal_cuda_lite/src/tests/utilities.rs

@@ -468,7 +468,7 @@ pub fn fuzz_genomes<T: TestDType>(
 
             let mut list_cache = FxHashMap::default();
             let mut expr_cache = FxHashMap::default();
-            let llir_graph = egglog_to_llir(
+            let mut llir_graph = egglog_to_llir(
                 egraph,
                 genome.clone(),
                 ops,
@@ -477,6 +477,12 @@ pub fn fuzz_genomes<T: TestDType>(
                 &mut expr_cache,
                 None,
             );
+            // Same finalization as `Graph::search` performs on the chosen
+            // best LLIR: collapse the rolled body's loop markers into a
+            // fully-unrolled LLIR. The runtime cannot execute LoopStart /
+            // LoopEnd / LoopInput / LoopOutput markers — they exist only as
+            // a search-time scaffold the auto-roll prepass introduces.
+            unroll_loops_in_llir(&mut llir_graph);
 
             let mut rt = CudaRuntime::initialize(stream.clone());
             rt.load_llir(&llir_graph);

crates/luminal_metal/src/dyn_backend.rs

@@ -0,0 +1,48 @@
+//! [`DynBackend`] implementation for the Metal runtime.
+
+use luminal::dtype::DType;
+use luminal::dyn_backend::{bytes_to_native_data, compile_backend, BackendCompileArgs, DynBackend};
+use luminal::prelude::*;
+
+use crate::runtime::MetalRuntime;
+
+/// [`DynBackend`] wrapper for [`MetalRuntime`].
+pub struct MetalDynBackend {
+    pub runtime: MetalRuntime,
+}
+
+impl DynBackend for MetalDynBackend {
+    fn name(&self) -> &str {
+        "metal"
+    }
+
+    fn set_data_bytes(&mut self, node: NodeIndex, bytes: Vec<u8>, dtype: DType) {
+        self.runtime
+            .set_data(node, bytes_to_native_data(bytes, dtype));
+    }
+    fn set_data_f32(&mut self, node: NodeIndex, data: Vec<f32>) {
+        self.runtime.set_data(node, data);
+    }
+    fn get_output_f32(&self, node: NodeIndex) -> Vec<f32> {
+        self.runtime.get_f32(node)
+    }
+    fn execute(&mut self, dyn_map: &FxHashMap<char, usize>) {
+        self.runtime.execute(dyn_map);
+    }
+}
+
+pub fn metal_factory(
+    graph: &mut Graph,
+    args: BackendCompileArgs,
+) -> Result<Box<dyn DynBackend>, String> {
+    compile_backend::<MetalRuntime>(
+        graph,
+        args,
+        || Ok(MetalRuntime::initialize(())),
+        |rt, node, bytes, dtype| {
+            rt.set_data(node, bytes_to_native_data(bytes, dtype));
+        },
+        None,
+        |rt| Box::new(MetalDynBackend { runtime: rt }),
+    )
+}

crates/luminal_metal/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod dyn_backend;
 pub mod kernel;
 pub mod runtime;
 

crates/luminal_metal/src/runtime.rs

@@ -234,6 +234,10 @@ impl Runtime for MetalRuntime {
         }
     }
 
+    fn aggregate_profile_metrics(metrics: &[Self::ProfileMetric]) -> Self::ProfileMetric {
+        metrics.iter().copied().sum()
+    }
+
     #[tracing::instrument(skip_all)]
     fn load_llir(&mut self, llir_graph: &LLIRGraph) {
         self.pipelines.clear();

crates/luminal_python/CLAUDE.md

@@ -24,7 +24,7 @@ consult before writing new egglog rules, CUDA kernels, or optimizer passes.
 ## Testing Best Practices
 
 ### Overview
-The luminal_python crate provides a bridge between PyTorch models and the luminal library via ONNX. Tests should verify this integration end-to-end by testing the actual user workflow: PyTorch model → torch.compile → luminal backend.
+The luminal_python crate provides a bridge between PyTorch models and the luminal library via the PT2 Export pipeline. Tests should verify this integration end-to-end by testing the actual user workflow: PyTorch model → torch.compile → luminal backend.
 
 ### Test Pattern (CORRECT)
 
@@ -67,11 +67,11 @@ class AddTestModel(torch.nn.Module):
 
 ### What NOT to Do
 
-**❌ DO NOT create ONNX files directly in tests:**
+**❌ DO NOT create pt2 files directly in tests:**
 ```python
 # WRONG - bypasses the PyTorch integration
-model_path = create_onnx_model(...)
-graph_result = luminal.process_onnx(model_path, backend='native')
+model_path = create_pt2_model(...)
+graph_result = luminal.process_pt(model_path, backend='native')
 ```
 
 **✓ DO create PyTorch models and use torch.compile:**
@@ -83,16 +83,16 @@ model_compiled = torch.compile(model, backend=luminal_backend)
 
 ### Rationale
 
-- **End-to-end testing**: Tests verify the complete PyTorch → ONNX → luminal pipeline
+- **End-to-end testing**: Tests verify the complete PyTorch → Pt2 → luminal pipeline
 - **User-facing API**: Tests use the same API that users will use (torch.compile)
 - **Correctness**: Comparing compiled vs original PyTorch output ensures correctness
 - **Maintainability**: Consistent pattern across all tests makes the codebase easier to understand
-- **Simplicity**: No manual ONNX file creation, no tempfile cleanup, no numpy comparisons
+- **Simplicity**: No manual Pt2 file creation, no tempfile cleanup, no numpy comparisons
 
 ### Special Cases
 
 **Testing constants:**
-Use inline tensor literals in the forward method - PyTorch exports these as ONNX Constant nodes:
+Use inline tensor literals in the forward method - these are exported as constant tensors:
 ```python
 def forward(self, x: torch.Tensor) -> torch.Tensor:
     constant = torch.tensor([1.0, 2.0, 3.0])
@@ -100,14 +100,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 ```
 
 **Testing type casts:**
-Use `.to(dtype)` method - PyTorch exports these as ONNX Cast nodes:
+Use `.to(dtype)` method - these are exported as type cast operations:
 ```python
 def forward(self, x: torch.Tensor) -> torch.Tensor:
     return x.to(torch.float32)
 ```
 
 **Testing complex operations:**
-Chain operations naturally in PyTorch - ONNX export handles the conversion:
+Chain operations naturally in PyTorch - the export pipeline handles the conversion:
 ```python
 def forward(self, x: torch.Tensor) -> torch.Tensor:
     transposed = x.transpose(0, 1)

crates/luminal_python/LessonsLearned.md

@@ -340,7 +340,7 @@ with matching shape tracker dimensions.
 
 ---
 
-## Bug: TopK values wrong on CUDA (gather_elements with sliced non-contiguous indices)
+## 2026-03-05 — TopK Values Wrong on CUDA (gather_elements with sliced non-contiguous indices)
 
 1. **Symptom**: `test_topk_values` failed on CUDA — rows 0-1 were correct but rows 2+ returned
    the value at column 0 of each row (all three top-k positions got the same value).
@@ -748,3 +748,305 @@ method rather than string-matching on Debug output. Additionally, when diagnosin
 candidates rejected" during search, check whether the rejection is from actual float NaN
 or from dtype misinterpretation — the key diagnostic is whether the NaN pattern is
 identical across all attempts (dtype issue) vs varying (actual numerical issue).
+
+## 2026-04-22 — Benchmark python_luminal Path: NativeRuntime Panic on CUDA Weights
+
+### What the symptom was
+
+Running `benchmarks/ttft/run.py` with the `python_luminal` path panicked deep in Rust:
+
+```
+thread panicked at src/hlir.rs:2239:40: no entry found for key
+```
+
+The panic occurred in `NativeRuntime::execute` when the `Output` node tried to read its
+predecessor's buffer from `self.buffers` — and the buffer wasn't there.
+
+### What the actual root cause was
+
+The luminal Python wheel was built without `--features cuda` (plain `maturin build --release`).
+This means `_cuda_lite_factory_capsule` is not compiled into the `.so` file. In `main.py`,
+`_detect_factory_capsule` catches the resulting `ImportError` and **silently** falls back to
+`_native_factory_capsule` (NativeRuntime / CPU runtime).
+
+The benchmark model (`LlamaForCausalLM.from_pretrained(...).to("cuda")`) has all weights as
+CUDA device pointers. `BackendCompileArgs.device_ptrs` is populated with these GPU pointers.
+NativeRuntime has no mechanism to handle GPU-resident weight data — the `device_ptrs` map is
+simply ignored. After search completes (it can search because it uses dummy CPU data during
+profiling), the first real `execute()` call processes the graph:
+
+1. `Input` nodes are skipped (their buffers should be pre-populated by `set_input_from_ptr`)
+2. Weight `Input` nodes were set via `set_input_device_ptr` — but NativeRuntime's
+   `set_input_device_ptr` likely no-ops or stores garbage, leaving those buffers empty
+3. The `Output` node looks up its predecessor's buffer → key not found → panic
+
+### Why it was hard to find
+
+1. **Silent fallback**: `_detect_factory_capsule` catches `ImportError` without logging a
+   warning. Nothing in stdout indicates you're running on CPU when the model is on GPU.
+2. **Search succeeds**: The e-graph search runs to completion (searches 1 group, 1 chunk in
+   ~15s) because it uses 1.0f32 dummy data that doesn't need GPU. The failure only occurs at
+   first real execution.
+3. **Misleading error site**: `hlir.rs:2239` is in NativeRuntime's buffer-copy loop for Output
+   nodes — it gives no indication that the root cause is a missing CUDA feature flag at build time.
+4. **Backtrace required**: Without `RUST_BACKTRACE=1`, only the panic message is visible;
+   the `NativeRuntime` frame that reveals the CPU fallback is hidden.
+
+### The fix
+
+Rebuild the wheel with CUDA support:
+```bash
+maturin build --release --features cuda
+pip install target/wheels/luminal_python-*.whl --force-reinstall
+```
+
+Or via the test runner: `./run_tests_cuda.sh` uses `maturin develop --features cuda -r`.
+
+Consider adding an explicit warning or error in `_detect_factory_capsule` when CUDA inputs are
+detected but no CUDA factory is available:
+
+```python
+if device.type == "cuda":
+    try:
+        from .luminal import _cuda_lite_factory_capsule
+        return _cuda_lite_factory_capsule()
+    except ImportError:
+        import warnings
+        warnings.warn(
+            "CUDA inputs detected but luminal was built without --features cuda. "
+            "Falling back to NativeRuntime (CPU) — this will likely panic at runtime.",
+            RuntimeWarning,
+            stacklevel=3,
+        )
+```
+
+### The regression test
+
+`test_hf_llama3_8b_instruct_1layer` in `tests/test_llama3.py` — tests the exact architecture
+from the benchmark (Meta-Llama-3-8B-Instruct, 4096 hidden, 32 attn heads, 8 KV heads) with
+1 layer and random weights. This test passes with `--features cuda` and panics without it.
+
+### General principle
+
+**When a feature gate silently changes the runtime backend, assert that the selected backend
+is compatible with the input device.** A CUDA tensor flowing into a CPU-only runtime is always
+a programming error, not a graceful degradation. The failure should surface at factory
+selection time (with a clear error message), not deep in a Rust buffer-copy loop.
+
+---
+
+## 2026-03-25 — KernelExp/KernelSigmoid: Fused CUDA Kernels for Precision
+
+1. **Symptom**: `test_hf_llama3_full` (16-layer Llama-3.2-1B) had ~1e-4 max diff vs PyTorch.
+2. **Root cause**: `exp(x)` was computed as `exp2(x * 1.442695)` — the constant truncated by `{:.6}` format + extra multiply adds rounding. Sigmoid was 5 separate kernels. SumReduce had naive accumulation.
+3. **Why hard**: Per-operation error was ~1e-7 but compounded over 16 layers × ~25 extra materializations. The egglog `Exp` rewrite depends on exact constant format matching.
+4. **Fix**: Added `KernelExp` (uses `expf()`), `KernelSigmoid` (uses `1/(1+expf(-x))`), and Kahan summation in SumReduce. Each uses both `kernel_rewrite` and a direct egglog pattern match with range checks (e.g., `(> ?val 1.44) (< ?val 1.45)`) to bypass constant format dependency.
+5. **Principle**: When decomposed CUDA kernel chains cause precision loss, add fused kernels via `kernel_rewrite`. For robustness, add BOTH the logical-op rewrite path AND a direct HLIR pattern match — the constant format in egglog can be fragile.
+
+---
+
+## 2026-04-23 — NativeRuntime Multi-Call Panic: Input Buffers Cleared After Each Run
+
+1. **Symptom**: The compiled model panicked with `hlir.rs:XXXX: no entry found for key` on the second call. First call succeeded; subsequent calls failed.
+2. **Root cause**: `NativeRuntime::execute` in `src/hlir.rs` called `self.buffers.retain(|k, _| output_nodes.contains(k))` after each run to free intermediate buffers. This correctly pruned temporary buffers but also pruned the Input-node buffers that hold model weights — so on the second call, the weight tensors were gone.
+3. **Why hard**: The bug never manifested in the test suite because every test called the compiled model exactly once per compile. The issue only appeared when running a bench loop that called the model multiple times. The panic location (deep in buffer lookup) gave no indication that the root cause was in the buffer retention policy.
+4. **Fix**: Changed the retain predicate to keep both `Output` and `Input` nodes:
+   ```rust
+   let keep_nodes = graph.node_indices()
+       .filter(|n| is::<Output> || is::<Input>)
+       .collect();
+   self.buffers.retain(|k, _| keep_nodes.contains(k));
+   ```
+5. **Principle**: When buffer lifetime policies are changed to free memory after a run, always verify that *persistent* state (model weights stored in Input nodes) is excluded from the cleanup sweep. A test that compiles + calls once per test function will never catch a multi-call regression — add a dedicated multi-call test for any compiled runtime.
+
+---
+
+## 2026-04-23 — PT2 USER_INPUT_MUTATION Outputs Confuse Dynamo Caller
+
+1. **Symptom**: With `StaticCache`, the compiled model returned `[1]` (cumulative_length update) instead of `[1, vocab_size]` logits. The wrong tensor was silently mapped to the output variable.
+2. **Root cause**: When `torch.export` encounters in-place mutations to input tensors (KV cache updates via `index_copy_`), it lifts them as `USER_INPUT_MUTATION` output specs, placed *before* the actual `USER_OUTPUT` logits in `ep.graph_signature.output_specs`. The compiled model returned all outputs; dynamo mapped index 0 (the mutation) to the first return value.
+3. **Why hard**: The output shape `[1]` from `cumulative_length` looked like a valid (though wrong) output. No error was raised — just wrong logits. Required inspecting `ep.graph_signature.output_specs` and understanding the ordering convention for different `OutputKind` values.
+4. **Fix**: In `pt2_backend`, parse `output_specs` to build a `mutation_mappings` list and `user_output_indices`. Wrap the compiled model to: (a) copy mutation outputs back into the corresponding input tensors, and (b) return only the `USER_OUTPUT` tensors.
+5. **Principle**: After `torch.export(...).run_decompositions()`, always inspect `ep.graph_signature.output_specs` when the model has in-place operations (KV cache, BN running stats). The output ordering is: mutations first, then actual outputs — and the caller only expects actual outputs.
+
+---
+
+## 2026-04-23 — CUDA Version Mismatch: torch+cuXXX Must Match System Driver
+
+1. **Symptom**: `torch.cuda.is_available()` returned `False` despite `nvidia-smi` showing a GPU. Warning: "CUDA initialization: The NVIDIA driver on your system is too old (found version 12080)."
+2. **Root cause**: `torch==2.11.0+cu130` requires CUDA 13.0 which needs driver >= 575. The system has driver 570 (CUDA 12.8 max). The mismatch caused silent CPU fallback — no error, just False from `is_available()`.
+3. **Why hard**: The bench appeared to start successfully (model loaded, compilation ran) but produced no results because it was running an 8B model on CPU. Zero output with exit code 0 looked like a hang or silent crash.
+4. **Fix**: Installed `torch==2.11.0+cu128` from `https://download.pytorch.org/whl/cu128`. CUDA 12.8 matches driver 570. Also needed matching `torchvision==0.26.0+cu128` and the `nvidia-cusparselt-cu12` runtime library.
+5. **Principle**: Before running any CUDA-dependent bench or test, verify `torch.cuda.is_available()` returns `True`. Check `nvidia-smi` CUDA Version field against the `+cuXXX` suffix in `torch.__version__` — they must match (CUDA runtime ≤ driver's max supported version). Never assume CPU fallback "works" for large model benchmarks.
+
+---
+
+## 2026-04-26 — Loop unroll-union rules silently disabled in full egglog stage
+
+1. **Symptom**: Python `test_llama_transformer_block` (CUDA backend) produced output ~1e-2 off from PyTorch (atol=1e-4) on the `loop_rolling` branch. All component tests (RMSNorm, attention, SwiGLU, RoPE) passed. The diff pattern was suspicious: row 0 of the (1,4,32) output matched exactly, rows 1–3 differed slightly. Disabling rolling fixed it.
+2. **Root cause**: The auto-roll prepass folds three sequential scalar muls in PyTorch's `pow(2)` decomposition (`exp2(log2(x) * 0.693 * 2.0 * 1.442)` — the last constant is `log2(e)`). The kernel `direct-exp-fusion` egglog rule rewrites `Mul(?x, log2_e_const) → Exp2(...)` into `KernelExp(?x)` (single `expf()` instead of separate exp2f + multiply by truncated log2(e)). Without rolling, this fusion fires and the float chain stays stable; with rolling the fusion can't see through the `LoopStart`/`LoopEnd` markers, so the chain stays as `KernelMul → KernelExp2`, and the truncated `log2(e)` constant accumulates ~1e-7 error per layer that compounds into ~1e-2 over the full block.
+
+   The unroll-union rules I'd added (`Mul`/`Add`/etc. binary-op rules that union a rolled body with its fully-unrolled equivalent) were registered only in `EgglogOp::early_rewrites()`, not `rewrites()`. The egglog driver feeds `early_rewrites` only into the early-stage program and `rewrites` only into the full-stage program. So the unrolled chain materialised in the early egraph, the early→full extract picked the (cheaper) rolled form, the unrolled chain was lost, and `direct-exp-fusion` (which runs in the full stage) had nothing to match against.
+3. **Why hard**: The post-unroll LLIR for the rolled vs un-rolled paths *looked* nearly identical when scanned visually — both had the Log2 → Mul × 3 → Exp2 chain. The diff was 2 extra Muls vs no-rolling, and the actual semantic gap was visible only in op-name counts: WITH-rolling had 3 `KernelExp2` and 0 `KernelExp`, WITHOUT-rolling had 1 `KernelExp2` and 2 `KernelExp`. Tracking the missing fusion to the early/full ruleset split required reading the egglog driver carefully and noticing that `OpTextParts` builds `early_rewrites` and `full_rewrites` from disjoint method calls.
+4. **Fix**: Register `binary_op_unroll_rules` in BOTH `early_rewrites()` (so fusion patterns like GLUMoE can match before the early-stage extract, which is what fixed `test_glumoe_gemma_gelu_matches_unfused_output` earlier in the session) AND `rewrites()` (so kernel-level rewrites like `direct-exp-fusion` can match in the full stage on the unrolled chain). One block per binary op (`Add`, `Mul`, `Mod`, `LessThan`).
+5. **Principle**: When egglog has multiple stages (early/full) with disjoint rule sets, any rewrite that materialises new HLIR/IR enodes (rather than just lowering to LLIR) needs to fire in BOTH stages if downstream rewrites in BOTH stages might want to see the new structure. Putting "preparatory" rewrites only in `early_rewrites` means their effect is lost across the early→full handoff. The narrow rule of thumb: if your rule's outputs are intended to enable matches by other rules, audit which stages those other rules run in and register accordingly.
+
+---
+
+## 2026-04-26 — `unroll_loops_in_llir` panicked on iteration-invariant body producers
+
+1. **Symptom**: Modal CI/CD job for the gemma example panicked at `src/graph.rs:1867` with `no entry found for key`. The line is `clone_map[i - 1][&body_producer]` inside `unroll_loops_in_llir`'s `resolve_src` closure — `body_producer` (the LoopEnd's incoming source for that slot) wasn't a key in the per-iteration clone map. cuda_lite/python tests didn't repro: only triggered by the specific genome and graph shapes that gemma's longer search settles on.
+2. **Root cause**: `body_nodes` is computed by walking *forward* from each LoopStart/LoopInput/LoopInputStatic outgoing edge, stopping at markers and `Output` ops. Some egglog-extracted LLIRs land a `body_producer` that isn't reachable via that forward walk — i.e., its only ancestors are non-marker (a constant, an external input, or an op whose chain was congruence-merged off the marker chain by rules like `LoopInputStatic inline`). Semantically this is a degenerate "iteration-invariant body": every iter computes the same value, so the loop's state never changes. The per-iter clone path needed a fallback for that case.
+3. **Why hard**: cuda_lite and python tests don't generate genomes that produce this shape, so local runs always pass. The forward-walk-only definition of `body_nodes` is *almost* always right — only specific extraction shapes from longer searches expose the gap. Test-driven debugging has limited reach when the failure mode depends on a search trajectory the local fuzzers don't explore.
+4. **Fix**: in `unroll_loops_in_llir::resolve_src`, when the LoopStart-resolved `body_producer` isn't in `body_nodes`, return `body_producer` itself for iter > 0 instead of indexing `clone_map[i - 1]`. The body op didn't depend on the loop variable, so every iter > 0 carries the same value forward — using `body_producer` directly is semantically correct. Mirrored the same `unwrap_or(body_producer)` fallback in the post-loop substitution map (`marker_post_sub` for LoopEnd / LoopOutputSelect). Added a backward-walk-from-end-markers backfill in `collapse_loops_to_first_iter` so its body-node iteration also covers these nodes (it doesn't have a clone_map, but does need to rewire body ops' incoming edges before deleting markers).
+5. **Principle**: When a graph-walk-derived set is used as a hashmap key requirement, every code path that *could* produce a key outside that set needs a graceful fallback — not just a defensive `expect`. For loop unrolling specifically, the rule is: `body_nodes` is the set of "ops that participate in per-iter computation"; ops on the LoopEnd's path that *don't* participate (iteration-invariant) are still legitimate, and need a "no clone, share across iters" path through `resolve_src` and `marker_post_sub`. Forward-walk-only `body_nodes` is correct only when extraction never produces iteration-invariant body producers — and in an egglog-driven search, that's not a guarantee you can make.
+
+---
+
+## 2026-04-26 — Iteration-invariant state slots are a first-class concept, not a defensive fallback
+
+1. **Symptom + fix recap**: gemma Modal CI panicked at `clone_map[i-1][&body_producer]` because some state slots' `body_producer` (LoopEnd's incoming) isn't in `body_nodes` (forward walk from input markers). The first commit pair (16de9638 / 93fb02c4) caught this with `.unwrap_or(body_producer)` — which works but reads as "defensive, unclear *why* this case exists."
+2. **What's actually happening**: extracted LLIR from gemma legitimately puts a `KernelConstant` at LoopEnd's incoming for some state slots. e.g. for one slot of gemma's body=104 trips=5 rolling: `initial = KernelConstant 1.442695` (log2 e), `body_producer = same node`. For another: `body_producer = KernelConstant 9.21034` (ln 10000, RoPE's frequency base after `Log2 * ln(2)` simplification). egglog's kernel-level rewrites legitimately union body-slot eclasses with these constants when the body chain provably reduces to them. The state really is iteration-invariant — every iter sees the same value.
+3. **Why "defensive fallback" framing is misleading**: it implies the LLIR is broken. It isn't. The forward-walk-only `body_nodes` definition just doesn't cover this case, because the case requires no per-iter cloning at all. A *node not reachable from any loop input marker has no input-marker ancestor*, so by construction its value doesn't depend on the loop's per-iter state.
+4. **Cleaner formulation**: name the concept. Compute an `iteration_invariant_slots: HashSet<LoopStart>` set at the same time `start_meta` is built, with the rule `body_producer ∉ body_nodes ⇒ iteration_invariant`. `resolve_src` and `marker_post_sub` then have explicit branches: if the slot is invariant, use `body_producer` directly; otherwise the standard per-iter clone lookup. The behavior is the same as the `unwrap_or` band-aid, but the code now documents that this is a real, sound case the unroll handles correctly — not a panic suppressor.
+5. **Principle**: when an `unwrap_or` papers over a case that turns out to be semantically valid, the right cleanup isn't to keep the `unwrap_or` and add a comment — it's to name the case. Hoist the predicate into a set or enum and branch on it explicitly. The compiler then enforces that every consumer of the per-iter cloning machinery has an opinion on iteration-invariant slots, instead of silently relying on a `Map::get` returning `None` at the right moment.
+
+---
+
+## 2026-04-30 — `translate_grouped_mm` casted the full expert weight to F32, OOMing search on Qwen3-MoE
+
+### What the symptom was
+
+`benchmarks/ttft/run.py --config qwen3-moe` crashed every search-profile attempt with:
+```
+crates/luminal_cuda_lite/src/runtime.rs:711: called `Result::unwrap()` on an `Err` value:
+  DriverError(CUDA_ERROR_OUT_OF_MEMORY, "out of memory")
+```
+The DB shows this had been failing every run for ~2 weeks. The rust `examples/qwen3_moe` ran fine end-to-end. python_baseline / python_torch_compile / qwen3-4b were all fine — only python_luminal × qwen3-moe failed.
+
+### What the actual root cause was
+
+`translate_grouped_mm` in `crates/luminal_python/rust/src/translator/tensor.rs` was lowering HF's `_grouped_mm(input, weight, offs)` op to a *full-broadcast* batched matmul plus a group-mask:
+
+```rust
+let weight_f = weight.cast(DType::F32);                  // [G=128, K, N] cast → 1.5 GB / layer
+let input_batched = input_f.expand_dim(0, g);
+let all_out = input_batched.matmul(weight_f);            // [G, S, N]
+let mask = ... (g_arange == expert_id).cast(F32);
+let out = (all_out * mask.expand_dim(2, n)).sum(0);      // mask + sum over G
+```
+
+The full `[G, K, N]` F32 cast intermediate is 1.5 GB / layer for gate-up and 0.6 GB / layer for down on Qwen3-30B-A3B. With 60 GB of persistent bf16 weights already on a 97 GB GPU, the search-time profiler ran out of memory allocating those casts.
+
+By contrast, `examples/qwen3_moe`'s `gather_experts` gathers only the top-K active experts per token first, then casts that small `[s, k, d1, d2]` slice (~100 MB / layer). The GLUMoE host op (`crates/luminal_cuda_lite/src/host/moe/glumoe_rewrite.egg`) is also wired to this gather pattern.
+
+### Why it was hard to find
+
+1. **Code path was reasonable in isolation**: at small scale (`test_grouped_mm_fallback`: g=2, K=8, N=16) the broadcast version was fine — the F32 cast was only 1 KB, and search profiling never noticed.
+2. **The error reported "out of memory" but the rest of the system looked healthy**: 60 GB weights + 37 GB headroom looks like plenty until you realise 48 layers × 2.1 GB cast intermediates per layer doesn't fit, even after loop rolling.
+3. **The DB's `code 1` failures looked the same as a Python exception** — the actual panic site (`runtime.rs:711:64` `stream.alloc_zeros(needed_bytes).unwrap()`) had to be recovered from a tmux scrollback because the orchestrator's stdout was already torn down by the time we looked.
+
+### The fix
+
+Rewrote `translate_grouped_mm` to gather first, matmul second:
+
+```rust
+// expert_id[m] = first g s.t. m < offs[g], clamped to [0, G-1]
+let expert_id = ge_boundary.sum(0).minimum_f32(g_max_f).cast(DType::Int);
+
+// flat_idx = expert_id * (K*N) + iota('z', (K, N))   — same shape as
+// rust qwen3_moe's `gather_experts`
+let flat_idx = (expert_id * (k * n))
+    .expand_dim(1, k).expand_dim(2, n)
+    + self.graph.iota(Expression::from('z'), (k, n)).expand_dim(0, s);
+
+let weight_gathered = weight.gather(flat_idx);            // [S, K, N], bf16
+let result = input.cast(F32).unsqueeze(1)
+    .matmul(weight_gathered.cast(F32))                    // [S, 1, N]
+    .squeeze(1);
+```
+
+Two important details:
+
+1. **Clamp `expert_id` to `[0, G-1]`**: at search time, dummy data fills `offs` with all-1s (`make_ones_bytes` in `compile_backend`). For S>1 that pushes `expert_id` to G (boundary count = G), which is one past the last valid expert and OOBs the gather. HF's own grouped-MM forward also clamps for the same reason (invalid expert IDs from EP).
+2. **Don't cast the full weight**: the cast moved from before the batched-matmul (over `[G, K, N]`) to after the gather (over `[S, K, N]`). 16× shrink at prefill (S=top_k=8 vs G=128).
+
+### Result
+
+`search-iters=1` end-to-end works on Qwen3-30B-A3B: `BENCH_RESULT … "ttft_ms": 9350.5, "tpot_ms": 1166.7`. The OOM is gone.
+
+`search-iters>=5` still crashes — but with a *different*, downstream `CUDA_ERROR_ILLEGAL_ADDRESS` during execution after search completes. That looks like the same family as the 2026-03-07 / 2026-03-09 egglog-extractor non-determinism bugs (some mutation during search picks a kernel/rewrite combo that's broken at this scale). It's a separate investigation — the gather-based lowering is correct in isolation (`test_grouped_mm_fallback` passes; a synthetic `g=128, S=8, K=2048, N=1536` bf16 test passes with max-diff ~2.4e-4).
+
+### General principle
+
+**When lowering an op that takes a per-row index over a large parameter, gather first and cast second — never cast the full parameter to F32 just because your matmul kernel is F32-only.** A "broadcast over G + mask" pattern is mathematically equivalent to "gather per-row" but materialises a G× larger intermediate — fine for tests, ruinous on real MoE checkpoints. When in doubt, mirror the rust example's pattern: the egglog fusion rules (GLUMoE here) are written to recognise the gather form, not the broadcast-and-mask form.
+
+Also: search-time dummy-1 inputs are not the same shape as runtime inputs. Anything you compute from a runtime tensor (cumsum offsets, routing indices, mask boundaries) needs to remain in-bounds for the dummy. Clamp index-producing chains as a matter of course, not just when the math says you "should" — `make_ones_bytes` is a hostile witness.
+
+---
+
+## 2026-05-01 — `KernelScatter` float4 vectorization wrote 2× past end of buffer for bf16/f16 KV cache
+
+### What the symptom was
+
+After the `translate_grouped_mm` gather rewrite (above) cleared the OOM, the qwen3-moe bench progressed past search but panicked during execution roughly 40% of the time:
+```
+crates/luminal_cuda_lite/src/runtime.rs:1204:
+  CUDA execute error in "CudaGraph":
+    DriverError(CUDA_ERROR_ILLEGAL_ADDRESS, "an illegal memory access was encountered")
+```
+qwen3-4b (dense) was unaffected; the bf16 KV cache in HF `StaticCache` was the only path triggering it. The rust `examples/qwen3_moe` ran fine because it uses an F32 KV cache.
+
+### What the actual root cause was
+
+`KernelScatter::compile` in `crates/luminal_cuda_lite/src/kernel/hlir.rs` emitted a hand-written CUDA copy phase that vectorised through `float4` (16-byte) reads/writes:
+
+```cuda
+long long n_vec = n_dest / 4;          // ← assumes 4-byte dtype
+float4 *out4 = (float4 *)out;
+const float4 *dest4 = (const float4 *)dest;
+for (long long i = tid; i < n_vec; i += blockDim.x) {
+    out4[i] = dest4[i];                 // ← writes 16 B per iteration
+}
+long long remainder_start = n_vec * 4;  // ← also assumes 4 elem/vec
+```
+
+For `dtype=F32` (4 bytes), `n_vec * 16 = n_dest * 4` bytes — exactly fills the buffer. For `dtype=Bf16` (2 bytes), `n_vec * 16 = (n_dest/4) * 16 = n_dest * 4` bytes, which is **2× the actual buffer size of `n_dest * 2` bytes**. The write walks half the buffer past the end of `out` (and reads past `dest`).
+
+Whether that produced an `ILLEGAL_ADDRESS` depended on whether the OOB region happened to land on an unmapped page. For different search outcomes, the surrounding allocator state differed → ~60% it was silent corruption, ~40% it crashed the CUDA context. That probabilistic mix is why the bug had been hidden — no test exercised a bf16 scatter (every existing scatter test uses F32 by default), and the rust example uses F32 KV cache so it was never seen there either.
+
+### Why it was hard to find
+
+1. **Probabilistic, but search-determinate**: the rewrite from HLIR `Scatter` → `KernelScatter` always fires (it's the only non-NoCopy path), so the kernel is always present. The crash depends on memory layout, which depends on which other kernels the search picked. Made it look like an egglog-mutation issue rather than a kernel-correctness issue.
+2. **Existing test coverage was F32-only**: `test_scatter_execution_correctness` (in `tests/consumed_buffer_tests.rs`) explicitly tries 50 random extractions to cover both `Scatter` and `ScatterNoCopy`, but always with `cx.tensor(5)` which defaults to F32. The bug would never surface there.
+3. **The panic message hid the kernel name**: it surfaced as a generic `"CudaGraph"` host-op panic — the cuda_graph_exec batches all kernels into one atomic launch, so the failing kernel disappears into the batch. To localize it I had to add a `LUMINAL_DEBUG_SEQ` env var to `CudaGraphOp::execute_internal` that bypasses graph batching and launches each kernel via `cuLaunchKernel` with a sync afterwards, surfacing kernel name + node + grid/block/pointers when one fails.
+
+### The fix
+
+Parameterise `n_vec` and the remainder-loop start by the number of dtype elements that fit in 16 bytes:
+
+```rust
+let elements_per_vec: usize = match self.dtype {
+    DType::F64 => 2,
+    DType::F32 | DType::Int => 4,
+    DType::F16 | DType::Bf16 | DType::I16 | DType::U16 => 8,
+    DType::Bool | DType::I8 | DType::U8
+        | DType::F8UE8M0 | DType::F8E4M3 | DType::F8E5M2 => 16,
+    other => panic!("Unsupported dtype for scatter vectorization: {other:?}"),
+};
+```
+and substitute `{elements_per_vec}` into the kernel template (both the `n_vec` calc and `remainder_start`). For F32 / Int the generated code is byte-for-byte identical to before, so existing F32 tests are unaffected; for any other dtype the byte coverage now exactly equals `n_dest * sizeof(dtype)` as intended.
+
+### Result
+
+Before fix: 3/5 success at iters=10 (probabilistic).
+After fix: 5/5 at iters=10, 3/3 at iters=50. All 206 HLIR tests still pass. TTFT/TPOT identical (~9.35s / ~1.17s).
+
+### General principle
+
+**Hand-rolled CUDA vectorisation with a fixed-width type (`float4`, `float2`, `int4`, …) is almost always specialised to one element size.** When the same kernel template is parameterised by `dtype`, every byte-count expression has to be too. The cheapest correct form is "elements per vector load" computed from the dtype's byte size — never hardcode `/4`.
+
+Also: **F32 is not a representative test dtype for kernels with vector loads.** When a kernel is written generic-over-dtype, the test matrix needs to actually exercise the dtypes (bf16, f16, bool) where the vector-element-count differs. A `test_scatter_bf16` would have caught this years before the qwen3-moe bench did. Same trap likely exists wherever else `float4` is cast over a `{dtype} *` template.
+
+Diagnostic also added: `LUMINAL_DEBUG_SEQ=1` on the python_luminal path will now bypass `CudaGraphOp` batching at execute time, launching each kernel sequentially with a sync afterwards. If a future ILLEGAL_ADDRESS hides inside a batched graph again, this surfaces the kernel name and node index immediately.
+

crates/luminal_python/modal_pytest_runner.py

@@ -186,7 +186,7 @@ class TestRunner:
         env = os.environ.copy()
         existing = env.get("PYTHONPATH")
         env["PYTHONPATH"] = f"{SRC_PATH}:{existing}" if existing else SRC_PATH
-        env["LUMINAL_BACKEND"] = "cuda"
+        env["LUMINAL_TEST_DEVICE"] = "cuda"
         env["UV_PROJECT_ENVIRONMENT"] = VENV_PATH
         env["MATURIN_PEP517_ARGS"] = "--features cuda --profile release"
         env["CUDARC_CUDA_VERSION"] = CUDARC_CUDA_VERSION

crates/luminal_python/pyproject.toml

@@ -7,8 +7,6 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy>=2.0.2",
     "torch>=2.10.0",
-    "onnx",
-    "onnxscript",
     "safetensors",
 ]
 
@@ -47,6 +45,6 @@ dev = [
     "pytest-randomly>=4.0.1",
     "transformers>=4.40.0",
     "diffusers>=0.35.0",
-    "onnxsim",
     "modal>=1.3.5",
+    "matplotlib>=3.8",
 ]

crates/luminal_python/run_all_tests.sh

@@ -16,13 +16,9 @@ rm -rf rust/target/wheels rust/target/debug rust/target/release
 uv run maturin develop --manifest-path rust/Cargo.toml
 
 echo ""
-echo "--- 1a: Native + ONNX ---"
+echo "--- 1a: Native backend tests ---"
 uv run pytest $NATIVE_TESTS -v
 
-echo ""
-echo "--- 1b: Native + PT2 ---"
-LUMINAL_EXPORT_MODE=pt2 uv run pytest $NATIVE_TESTS -v
-
 # ── Phase 2: CUDA Backend ───────────────────────────────────
 
 echo ""
@@ -31,12 +27,8 @@ rm -rf rust/target/wheels rust/target/debug rust/target/release
 uv run maturin develop --manifest-path rust/Cargo.toml --features cuda -r
 
 echo ""
-echo "--- 2a: CUDA + ONNX ---"
-RUST_BACKTRACE=1 LUMINAL_BACKEND=cuda uv run pytest $CUDA_TESTS -m "not slow" -v
-
-echo ""
-echo "--- 2b: CUDA + PT2 ---"
-RUST_BACKTRACE=1 LUMINAL_BACKEND=cuda LUMINAL_EXPORT_MODE=pt2 uv run pytest $CUDA_TESTS -m "not slow" -v
+echo "--- 2a: CUDA ---"
+RUST_BACKTRACE=1 LUMINAL_TEST_DEVICE=cuda uv run pytest $CUDA_TESTS -m "not slow" -v
 
 echo ""
 echo "=========================================="

crates/luminal_python/run_test_fx.sh

@@ -1,20 +0,0 @@
-#!/bin/bash
-set -e
-
-echo "=== Luminal Python Test Runner (PT2 Export Mode) ==="
-echo ""
-
-# Force clean rebuild of Rust extension
-echo "Step 1: Cleaning previous builds..."
-rm -rf rust/target/wheels rust/target/debug rust/target/release
-
-# Rebuild in development mode (faster compilation)
-echo "Step 2: Building Rust extension..."
-uv run maturin develop --manifest-path rust/Cargo.toml
-
-# Run pytest with PT2 export mode
-echo "Step 3: Running pytest with PT2 export mode..."
-LUMINAL_EXPORT_MODE=pt2 uv run pytest tests/test_hlir_ops.py tests/test_unary.py -v
-
-echo ""
-echo "=== Tests Complete ==="

crates/luminal_python/run_tests_cuda.sh

@@ -14,7 +14,7 @@ uv run maturin develop --manifest-path rust/Cargo.toml --features cuda -r
 
 # Run pytest with CUDA backend
 echo "Step 3: Running pytest with CUDA backend..."
-RUST_BACKTRACE=1 LUMINAL_BACKEND=cuda uv run pytest tests/test_hlir_ops.py tests/test_unary.py tests/test_llama3.py -m "not slow" -v
+RUST_BACKTRACE=1 LUMINAL_TEST_DEVICE=cuda uv run pytest tests/test_llama3.py tests/test_hlir_ops.py tests/test_unary.py -v
 
 echo ""
 echo "=== Tests Complete ==="

crates/luminal_python/run_tests_cuda_fx.sh

@@ -1,19 +0,0 @@
-#!/bin/bash
-set -e
-
-echo "=== Luminal Python Test Runner (CUDA + PT2 Export Mode) ==="
-echo ""
-
-# Force clean rebuild of Rust extension
-echo "Step 1: Cleaning previous builds..."
-rm -rf rust/target/wheels rust/target/debug rust/target/release
-
-# Rebuild in development mode (faster compilation)
-echo "Step 2: Building Rust extension..."
-uv run maturin develop --manifest-path rust/Cargo.toml --features cuda -r
-
-# Run pytest with CUDA backend and PT2 export mode
-echo "Step 3: Running pytest with CUDA backend + PT2 export mode..."
-RUST_BACKTRACE=1 LUMINAL_BACKEND=cuda LUMINAL_EXPORT_MODE=pt2 uv run pytest tests/test_hlir_ops.py tests/test_unary.py tests/test_llama3.py -m "not slow" -v
-echo ""
-echo "=== Tests Complete ==="

crates/luminal_python/rust/Cargo.toml

@@ -12,8 +12,6 @@ path = "src/lib.rs"
 cuda = ["dep:luminal_cuda_lite"]
 
 [dependencies]
-onnx-protobuf = "0.2"
-protobuf = "~3.4"
 rustc-hash = "2.1.1"
 luminal = {path= "../../.."}
 luminal_cuda_lite = {path="../../luminal_cuda_lite", optional = true}

crates/luminal_python/rust/src/compiled_graph.rs

@@ -1,423 +1,134 @@
 use luminal::{
-    prelude::{
-        tracing::{Level, span, trace},
-        *,
-    },
+    dyn_backend::{BackendCompileArgs, BackendFactory, DynBackend},
+    prelude::*,
     shape::Expression,
     visualization::ToDot,
 };
-use onnx_protobuf::{GraphProto, ModelProto};
 use pyo3::prelude::*;
-use std::{
-    collections::{HashMap, HashSet},
-    path::Path,
-};
+use std::collections::HashMap;
 
-#[cfg(feature = "cuda")]
-use crate::util::transpose_weight_data;
-use crate::{
-    dispatch::process_onnx_nodes,
-    runtime::*,
-    util::{
-        DimParamMap, get_shape_for_onnx_value, get_shape_for_onnx_value_expr,
-        load_all_tensor_floats, load_initializer_as_f32,
-    },
-};
+use crate::typed_data::TypedData;
+
+/// Maps symbolic dimension parameter names (e.g. "seq_len") to luminal Expression variable chars.
+pub type DimParamMap = HashMap<String, char>;
+
+/// Convert luminal DType to PT2 dtype integer code (for python interop)
+/// Types without a direct Pytorch equivalent map to the closest safe representation
+fn luminal_dtype_to_pt2_code(dtype: DType) -> u32 {
+    match dtype {
+        DType::U8 => 1,
+        DType::I8 => 2,
+        DType::I16 => 3,
+        DType::Int => 4, // i32
+        DType::U16 => 4, // u16 -> i32 (Pytorch has no u16 in older versions)
+        DType::F16 => 6,
+        DType::F32 | DType::TF32 => 7,
+        DType::F64 => 8,
+        DType::Bool => 12,
+        DType::Bf16 => 13,
+        _ => panic!("luminal_dtype_to_pt2_code: unsupported dtype {:?}", dtype),
+    }
+}
+
+/// Common intermediate result from translating a model graph.
+pub struct GraphTranslation {
+    pub graph: Graph,
+    pub tensor_ids: HashMap<String, NodeIndex>,
+    pub input_names: Vec<String>,
+    pub output_names: Vec<String>,
+    pub output_shape_exprs: Vec<Vec<Expression>>,
+    pub output_dtypes: Vec<DType>,
+    pub input_shape_exprs: Vec<Vec<Expression>>,
+    pub dim_param_map: DimParamMap,
+}
+
+/// Pre-loaded weight data from any model format (dtype-aware).
+pub struct WeightData {
+    /// (Input node label, typed data) for weights and constants.
+    pub weights: Vec<(String, TypedData)>,
+    /// label → element count for ALL Input nodes (for CUDA dummy data sizing).
+    pub tensor_sizes: HashMap<String, usize>,
+    /// label → (device_ptr, n_bytes) for zero-copy CUDA weight sharing.
+    pub device_ptrs: HashMap<String, (u64, usize)>,
+}
 
 #[pyclass(unsendable)]
 pub struct CompiledGraph {
     pub graph: Graph,
-    pub runtime: RuntimeBackend,
+    pub runtime: Box<dyn DynBackend>,
     pub tensor_ids: HashMap<String, NodeIndex>,
+    /// Cached label → NodeIndex map for O(1) lookups in set_weight_* methods.
+    label_map: HashMap<String, NodeIndex>,
     pub input_names: Vec<String>,
     pub output_names: Vec<String>,
     pub output_shapes: Vec<Vec<usize>>,
     pub output_shape_exprs: Vec<Vec<Expression>>,
+    pub output_dtypes: Vec<DType>,
     pub input_shape_exprs: Vec<Vec<Expression>>,
     pub dim_param_map: DimParamMap,
 }
 
 impl CompiledGraph {
+    /// Compilation pipeline for PT2/FX graphs.
+    ///
+    /// Takes a `GraphTranslation` (produced by `translate_pt2`) and `WeightData`,
+    /// builds the backend via the global registry, loads weights, and
+    /// returns a ready-to-execute `CompiledGraph`.
     pub fn parse_graph(
-        model: ModelProto,
-        model_directory: &Path,
-        backend: &str,
+        translation: GraphTranslation,
+        weight_data: WeightData,
+        factory: BackendFactory,
+        search_iters: usize,
     ) -> Result<CompiledGraph, String> {
-        let _span = span!(Level::TRACE, "Onnx Graphing Parsing").entered();
-        let onnx_graph = &model.graph;
-        let mut cx = Graph::new();
-        // We will need to track the tensors we allocate so we can match up inputs and outputs in the graph
-        let mut tensors: HashMap<String, GraphTensor> = HashMap::new();
+        let GraphTranslation {
+            mut graph,
+            tensor_ids,
+            input_names,
+            output_names,
+            output_shape_exprs,
+            output_dtypes,
+            input_shape_exprs,
+            dim_param_map,
+        } = translation;
 
-        // Dynamic dimension tracking
-        let mut dim_param_map: DimParamMap = HashMap::new();
-        let mut next_char = 'a';
-
-        // This is the name of all of the tensors we will need to fill in parameters for
-        let initializer_names: HashSet<&str> = onnx_graph
-            .initializer
-            .iter()
-            .map(|t| t.name.as_str())
-            .collect();
-
-        // Input is an overloaded term in Onnx, it both means the inputs into the model, like the next token
-        // and the parameters of the layers, for this we don't want any of the parameters
-        // Input here is in the straightforward meaning, those tensors you feed into the network for a
-        // forward passd
-        let input_names: Vec<String> = onnx_graph
-            .input
-            .iter()
-            .filter(|inp| !initializer_names.contains(inp.name.as_str()))
-            .map(|inp| inp.name.clone())
-            .collect();
-
-        // Create "holding" tensors for the input
-        // this way they can be considered in the graph computation, and later as we do mutiple runs we can target them and swap out the values
-        // in them and not need to recompile the network
-        for input in &onnx_graph.input {
-            // Use expression-aware shape parsing to detect DimParam (dynamic dims)
-            let shape_exprs =
-                get_shape_for_onnx_value_expr(input, &mut dim_param_map, &mut next_char);
-            if shape_exprs.is_empty() {
-                // Fall back to concrete parsing (initializer shapes don't have DimParam)
-                let shape = get_shape_for_onnx_value(input);
-                if shape.is_empty() {
-                    trace!("Input {} skipped because it is empty", input.name.clone());
-                    continue;
-                }
-                let tensor = cx.named_tensor(input.name.clone(), shape);
-                trace!("Input {} added to tensors", input.name.clone());
-                tensors.insert(input.name.clone(), tensor);
-                continue;
-            }
-            // Always F32: Python runtime always sends float32 data via .float().numpy()
-            let tensor = cx.named_tensor(input.name.clone(), shape_exprs);
-            trace!("Input {} added to tensors", input.name.clone());
-            tensors.insert(input.name.clone(), tensor);
-        }
-
-        for init in &onnx_graph.initializer {
-            if !tensors.contains_key(&init.name) {
-                let mut shape: Vec<usize> = init.dims.iter().map(|&d| d as usize).collect();
-                // Scalar (0-dim) tensors have empty dims; represent as [1] in luminal
-                if shape.is_empty() {
-                    shape = vec![1];
-                }
-                let tensor = cx.named_tensor(init.name.clone(), shape);
-                tensors.insert(init.name.clone(), tensor);
-            }
-        }
-
-        let mut weight_data = Vec::new();
-
-        let mut known_values: HashMap<String, Vec<f32>> = HashMap::new();
-
-        for init in &onnx_graph.initializer {
-            let n_elements: usize = init
-                .dims
+        // Build compile args from WeightData (convert TypedData -> raw bytes + dtype)
+        let compile_args = BackendCompileArgs {
+            search_iters,
+            weights: weight_data
+                .weights
                 .iter()
-                .map(|&d| d as usize)
-                .product::<usize>()
-                .max(1);
-            // MAGIC_NUMBER:
-            if n_elements <= 32 {
-                if let Some(floats) = load_initializer_as_f32(init) {
-                    known_values.insert(init.name.clone(), floats);
-                } else {
-                    // Questions
-                    // Should this be fatal
-                    // Should this be a print or a log
-                    panic!("Unable to initializer values for {:?}", init.name);
-                }
-            }
-        }
-        // Shape expressions map for propagating symbolic shape values through
-        // Shape→Gather→Unsqueeze→Concat chains in dynamic ONNX graphs
-        let mut shape_exprs: HashMap<String, Vec<Expression>> = HashMap::new();
-
-        // Process computation nodes (Constant nodes add to weight_data)
-        process_onnx_nodes(
-            &onnx_graph.node,
-            &mut tensors,
-            &mut cx,
-            &mut weight_data,
-            &mut known_values,
-            &mut shape_exprs,
-        )
-        .map_err(|e| format!("process_onnx_nodes failed: {}", e))?;
-
-        // Mark weight/constant tensors as persistent so their buffers survive
-        // execute()'s input consumption. User inputs (like input_ids) are NOT persisted
-        // since they are re-set via set_input() before each execution.
-        for (name, gt) in &tensors {
-            if !input_names.contains(name) {
-                gt.persist();
-            }
-        }
-
-        let has_dynamic = !dim_param_map.is_empty();
-
-        // Mark graph outputs (must happen before build_search_space)
-        let mut output_names = Vec::new();
-        let mut output_shapes = Vec::new();
-        let mut output_shape_exprs = Vec::new();
-        for output_vi in &onnx_graph.output {
-            if let Some(&gt) = tensors.get(&output_vi.name) {
-                // Force contiguous if the shape tracker is a non-contiguous view
-                // (e.g. a view-only slice that changed dims without a gather).
-                // Without this, get_f32 returns the full underlying buffer.
-                let gt = if gt.shape != gt.shape.contiguous() {
-                    let contiguous = gt * 1.0;
-                    tensors.insert(output_vi.name.clone(), contiguous);
-                    contiguous
-                } else {
-                    gt
-                };
-                gt.output();
-                let dims = gt.dims();
-
-                // Store Expression-based shapes for dynamic resolution
-                output_shape_exprs.push(dims.clone());
-
-                // For concrete output shapes, resolve now; for dynamic, use placeholder
-                let shape: Vec<usize> = dims.iter().map(|d| d.to_usize().unwrap_or(1)).collect();
-                if shape.is_empty() {
-                    return Err(format!(
-                        "Output tensor '{}' has no shape information in the ONNX model",
-                        output_vi.name
-                    ));
-                }
-                output_names.push(output_vi.name.clone());
-                output_shapes.push(shape);
-            }
-        }
-        // If we have dynamic dims, set initial values in the graph's dyn_map
-        // based on the concrete shapes from the example input used during export
-        if has_dynamic {
-            for input in &onnx_graph.input {
-                if initializer_names.contains(input.name.as_str()) {
-                    continue;
-                }
-                let concrete_shape = get_shape_for_onnx_value(input);
-                let expr_shape =
-                    get_shape_for_onnx_value_expr(input, &mut dim_param_map, &mut next_char);
-                for (expr, concrete) in expr_shape.iter().zip(concrete_shape.iter()) {
-                    if expr.to_usize().is_none() {
-                        // This is a symbolic dim — set initial value in dyn_map
-                        // Extract the char variable from the expression
-                        if let Some(ch) = dim_param_map
-                            .values()
-                            .find(|&&ch| Expression::from(ch) == *expr)
-                        {
-                            cx.set_dim(*ch, *concrete);
-                        }
-                    }
-                }
-            }
-        }
-        // Extract weight data from initializers (handles inline + external storage)
-        // Batch load reads each external file only once instead of per-tensor
-        for (name, floats) in load_all_tensor_floats(&onnx_graph.initializer, model_directory) {
-            if let Some(f) = floats {
-                weight_data.push((name, f));
-            }
-        }
-
-        // Collect tensor name -> NodeIndex mapping
-        let tensor_ids: HashMap<String, NodeIndex> = tensors
-            .iter()
-            .map(|(name, gt)| (name.clone(), gt.id))
-            .collect();
-
-        // Track which tensor names are Input nodes (includes those created during process_onnx_nodes)
-        let input_tensor_names: HashSet<String> = tensors.keys().cloned().collect();
-
-        let rt = match backend {
-            #[cfg(feature = "cuda")]
-            "cuda" => CompiledGraph::build_cuda_backend(
-                onnx_graph,
-                model_directory,
-                &mut tensors,
-                &mut weight_data,
-                &mut cx,
-                &input_tensor_names,
-            )?,
-            "native" => CompiledGraph::build_native_backend(
-                onnx_graph,
-                model_directory,
-                &mut tensors,
-                &mut weight_data,
-                &mut cx,
-                &input_tensor_names,
-            )?,
-            _ => {
-                #[cfg(feature = "cuda")]
-                {
-                    return Err(format!(
-                        "Invalid backend '{}'. Must be 'native' or 'cuda'",
-                        backend
-                    ));
-                }
-                #[cfg(not(feature = "cuda"))]
-                {
-                    if backend == "cuda" {
-                        return Err(
-                            "CUDA backend requested, but this luminal extension was built without the `cuda` feature. Rebuild with `maturin develop --features cuda -r` or use backend='native'."
-                                .to_string(),
-                        );
-                    }
-                    return Err(format!(
-                        "Invalid backend '{}'. This build only supports 'native'. Rebuild with the `cuda` feature to enable 'cuda'.",
-                        backend
-                    ));
-                }
-            }
+                .map(|(label, td)| (label.clone(), td.bytes.clone(), td.dtype))
+                .collect(),
+            tensor_sizes: weight_data.tensor_sizes,
+            device_ptrs: weight_data.device_ptrs,
         };
 
-        // Build input_shape_exprs for user inputs (needed for auto-dim detection)
-        let input_shape_exprs: Vec<Vec<Expression>> = input_names
+        // Create backend via the factory directly
+        let rt =
+            luminal::dyn_backend::compile_backend_from_factory(factory, &mut graph, compile_args)?;
+
+        // Resolve concrete output shapes from expressions
+        let output_shapes: Vec<Vec<usize>> = output_shape_exprs
             .iter()
-            .map(|name| {
-                if let Some(&gt) = tensors.get(name) {
-                    gt.dims()
-                } else {
-                    vec![]
-                }
-            })
+            .map(|exprs| exprs.iter().map(|e| e.to_usize().unwrap_or(1)).collect())
             .collect();
 
+        let label_map = luminal::dyn_backend::build_label_map(&graph);
+
         Ok(CompiledGraph {
-            graph: cx,
+            graph,
             runtime: rt,
             tensor_ids,
+            label_map,
             input_names,
             output_names,
             output_shapes,
             output_shape_exprs,
+            output_dtypes,
             input_shape_exprs,
             dim_param_map,
         })
     }
-
-    #[cfg(feature = "cuda")]
-    fn build_cuda_backend(
-        onnx_graph: &protobuf::MessageField<GraphProto>,
-        model_directory: &Path,
-        tensors: &mut HashMap<String, GraphTensor>,
-        weight_data: &mut Vec<(String, Vec<f32>)>,
-        context: &mut Graph,
-        input_tensor_names: &HashSet<String>,
-    ) -> Result<RuntimeBackend, String> {
-        let compute_n_elements = |name: &str| -> usize {
-            if let Some(vi) = onnx_graph.input.iter().find(|i| i.name == name) {
-                let shape = get_shape_for_onnx_value(vi);
-                shape.iter().product::<usize>()
-            } else if let Some(init) = onnx_graph.initializer.iter().find(|i| i.name == name) {
-                init.dims.iter().map(|&d| d as usize).product::<usize>()
-            } else if let Some((_, data)) = weight_data.iter().find(|(n, _)| n == name) {
-                data.len()
-            } else {
-                0
-            }
-        };
-
-        // CUDA: Two-phase - set data BEFORE search for profiling
-        let (mut cuda_rt, _stream) = prepare_cuda(context)?;
-
-        // Set dummy data for ALL input tensors using small non-zero values (ones).
-        // IMPORTANT: Must use 1.0, NOT 0.0. Zero inputs cause NaN in many ops:
-        //   - fmod(0, 0) = NaN  (Mod)
-        //   - recip(0) = inf → weight * inf = NaN  (Div)
-        //   - log(0) = -inf  (Pow)
-        //   - chain ops with zero produce NaN  (Erf)
-        // The search's has_nan_outputs check then rejects ALL candidates, causing
-        // "Failed to find viable genome" errors. See LessonsLearned.md entry #1.
-        // Note: torch.compile passes model weights as additional ONNX inputs (not
-        // initializers), so these dummy values also cover weight tensors.
-        for (name, gt) in &mut *tensors {
-            if !input_tensor_names.contains(name) {
-                continue;
-            }
-            let n_elements = compute_n_elements(name);
-            if n_elements > 0 {
-                cuda_rt.set_data(gt.id, vec![1.0f32; n_elements]);
-            }
-        }
-
-        // Overwrite with real initializer data (for accurate profiling)
-        // Batch load reads each external file only once
-        let init_data = load_all_tensor_floats(&onnx_graph.initializer, model_directory);
-        for (i, (name, floats_opt)) in init_data.iter().enumerate() {
-            let floats = match floats_opt {
-                Some(f) => f,
-                None => continue,
-            };
-            if let Some(gt) = tensors.get(name) {
-                cuda_rt.set_data(gt.id, floats.clone());
-            }
-            let kn_name = format!("{}_kn", name);
-            if let Some(gt_kn) = tensors.get(&kn_name) {
-                let dims: Vec<usize> = onnx_graph.initializer[i]
-                    .dims
-                    .iter()
-                    .map(|&d| d as usize)
-                    .collect();
-                if dims.len() == 2 {
-                    let transposed = transpose_weight_data(floats, dims[0], dims[1]);
-                    cuda_rt.set_data(gt_kn.id, transposed);
-                }
-            }
-        }
-
-        // Load constant node data
-        for (name, floats) in weight_data {
-            if let Some(gt) = tensors.get(name) {
-                cuda_rt.set_data(gt.id, floats.clone());
-            }
-        }
-
-        // Now finalize (search with profiling, data is available)
-        let cuda_rt = finalize_cuda(context, cuda_rt);
-
-        Ok(cuda_rt)
-    }
-
-    fn build_native_backend(
-        onnx_graph: &protobuf::MessageField<GraphProto>,
-        model_directory: &Path,
-        tensors: &mut HashMap<String, GraphTensor>,
-        weight_data: &mut Vec<(String, Vec<f32>)>,
-        context: &mut Graph,
-        _input_tensor_names: &HashSet<String>,
-    ) -> Result<RuntimeBackend, String> {
-        let mut rt = initialize_native(context)?;
-        context.search(NativeRuntime::default(), 1);
-
-        // Set initializer data - these MUST exist after optimization (they're weights)
-        // Skip _kn variants - they might be optimized away
-        // Batch load reads each external file only once
-        for (name, floats_opt) in load_all_tensor_floats(&onnx_graph.initializer, model_directory) {
-            let floats = match floats_opt {
-                Some(f) => f,
-                None => continue,
-            };
-            if let Some(gt) = tensors.get(&name) {
-                rt.set_data(gt.id, floats);
-            }
-        }
-
-        // Load constant node data, but skip _kn transposed variants
-        for (name, floats) in weight_data {
-            // Skip _kn transposed variants - might be optimized away
-            if name.ends_with("_kn") {
-                continue;
-            }
-            if let Some(gt) = tensors.get(name) {
-                rt.set_data(gt.id, floats.clone());
-            }
-        }
-        Ok(rt)
-    }
 }
 
 #[pymethods]
@@ -428,6 +139,24 @@ impl CompiledGraph {
         self.input_names.clone()
     }
 
+    /// Get the PT2 dtype codes for all inputs (in order of input_names).
+    #[getter]
+    fn input_dtypes(&self) -> Vec<u32> {
+        self.input_names
+            .iter()
+            .map(|name| {
+                if let Some(&node_id) = self.tensor_ids.get(name)
+                    && let Some(input) = (*self.graph.graph[node_id])
+                        .as_any()
+                        .downcast_ref::<luminal::hlir::Input>()
+                {
+                    return luminal_dtype_to_pt2_code(input.dtype);
+                }
+                7 // default to f32
+            })
+            .collect()
+    }
+
     /// Get the list of output tensor names.
     #[getter]
     fn output_names(&self) -> Vec<String> {
@@ -446,12 +175,24 @@ impl CompiledGraph {
         self.tensor_ids.keys().cloned().collect()
     }
 
-    /// Get the name of the active backend (native or cuda).
+    /// Get the name of the active backend.
     #[getter]
-    fn backend(&self) -> &'static str {
+    fn backend(&self) -> &str {
         self.runtime.name()
     }
 
+    /// The device type this backend operates on (e.g. "cpu", "cuda").
+    #[getter]
+    fn device_type(&self) -> &str {
+        self.runtime.device_type()
+    }
+
+    /// Whether the active backend supports device pointer operations (zero-copy GPU I/O).
+    #[getter]
+    fn supports_device_ptrs(&self) -> bool {
+        self.runtime.supports_device_ptrs()
+    }
+
     /// Whether this graph has dynamic (symbolic) dimensions.
     #[getter]
     fn has_dynamic_dims(&self) -> bool {
@@ -516,12 +257,136 @@ impl CompiledGraph {
         Ok(result)
     }
 
-    /// Set input tensor data by name.
+    /// Set input tensor data by name (f32, for backward compatibility).
     fn set_input(&mut self, name: &str, data: Vec<f32>) -> PyResult<()> {
         let node_id = self.tensor_ids.get(name).ok_or_else(|| {
             PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!("Unknown input tensor: {}", name))
         })?;
-        self.runtime.set_data(*node_id, data);
+        self.runtime.set_data_f32(*node_id, data);
+        Ok(())
+    }
+
+    /// Set input tensor data from a CPU host memory pointer (dtype-aware).
+    /// The pointer must point to contiguous data. `n_bytes` is the total byte count.
+    /// `dtype_code` uses PT2 numbering (7=f32, 6=f16, 13=bf16, etc.).
+    /// Converts source format to luminal's native format (e.g., i64→i32, f64→f32).
+    fn set_input_from_ptr(
+        &mut self,
+        name: &str,
+        ptr: u64,
+        n_bytes: usize,
+        dtype_code: u32,
+    ) -> PyResult<()> {
+        debug_assert!(ptr != 0, "set_input_from_ptr called with null pointer");
+        let node_id = self.tensor_ids.get(name).ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!("Unknown input tensor: {}", name))
+        })?;
+        let raw_bytes = unsafe { std::slice::from_raw_parts(ptr as *const u8, n_bytes).to_vec() };
+        let typed = TypedData::from_pytorch_bytes(raw_bytes, dtype_code);
+        self.runtime
+            .set_data_bytes(*node_id, typed.bytes, typed.dtype);
+        Ok(())
+    }
+
+    /// Set input from a device pointer. Zero-copy on device.
+    /// The pointer must be a valid device allocation with at least n_bytes bytes.
+    /// Requires a GPU backend (e.g. CUDA).
+    fn set_input_device_ptr(
+        &mut self,
+        name: &str,
+        device_ptr: u64,
+        n_bytes: usize,
+    ) -> PyResult<()> {
+        if !self.runtime.supports_device_ptrs() {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "set_input_device_ptr requires a GPU backend",
+            ));
+        }
+        let node_id = self.tensor_ids.get(name).ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!("Unknown input tensor: {}", name))
+        })?;
+        unsafe { self.runtime.set_device_ptr(*node_id, device_ptr, n_bytes) };
+        Ok(())
+    }
+
+    /// Set a weight from a device pointer (e.g. "fc1.weight"). Zero-copy on device.
+    /// Requires a GPU backend.
+    fn set_weight_device_ptr(
+        &mut self,
+        label: &str,
+        device_ptr: u64,
+        n_bytes: usize,
+    ) -> PyResult<()> {
+        if !self.runtime.supports_device_ptrs() {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "set_weight_device_ptr requires a GPU backend",
+            ));
+        }
+        let &node_id = self.label_map.get(label).ok_or_else(|| {
+            pyo3::exceptions::PyKeyError::new_err(format!("No Input node with label: {}", label))
+        })?;
+        unsafe { self.runtime.set_device_ptr(node_id, device_ptr, n_bytes) };
+        Ok(())
+    }
+
+    /// Register an external device pointer for an output tensor (zero-copy output).
+    /// Call before run() — the runtime will write kernel results directly into this buffer.
+    /// For aliased outputs (in-place ops), falls back to DtoD copy; check output_is_zero_copy() after run().
+    /// Requires a GPU backend.
+    fn set_output_device_ptr(
+        &mut self,
+        name: &str,
+        device_ptr: u64,
+        n_bytes: usize,
+    ) -> PyResult<()> {
+        if !self.runtime.supports_device_ptrs() {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "set_output_device_ptr requires a GPU backend",
+            ));
+        }
+        let node_id = self.tensor_ids.get(name).ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                "Unknown output tensor: {}",
+                name
+            ))
+        })?;
+        unsafe {
+            self.runtime
+                .set_output_device_ptr(*node_id, device_ptr, n_bytes)
+        };
+        Ok(())
+    }
+
+    /// Check whether an output tensor was zero-copied (written directly to the registered pointer).
+    /// Returns false for aliased outputs that need a fallback DtoD copy, or if no GPU backend.
+    /// Must be called after run().
+    fn output_is_zero_copy(&self, name: &str) -> PyResult<bool> {
+        let node_id = self.tensor_ids.get(name).ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                "Unknown output tensor: {}",
+                name
+            ))
+        })?;
+        Ok(self.runtime.output_is_zero_copy(*node_id))
+    }
+
+    /// Set a weight tensor from a CPU host pointer, matching by Input node label (dtype-aware).
+    /// `n_bytes` is the total byte count. `dtype_code` uses PT2 numbering (7=f32, 6=f16, 13=bf16, etc.).
+    fn set_weight_from_ptr(
+        &mut self,
+        label: &str,
+        ptr: u64,
+        n_bytes: usize,
+        dtype_code: u32,
+    ) -> PyResult<()> {
+        debug_assert!(ptr != 0, "set_weight_from_ptr called with null pointer");
+        let &node_id = self.label_map.get(label).ok_or_else(|| {
+            pyo3::exceptions::PyKeyError::new_err(format!("No Input node with label: {}", label))
+        })?;
+        let bytes = unsafe { std::slice::from_raw_parts(ptr as *const u8, n_bytes).to_vec() };
+        let typed = TypedData::from_pytorch_bytes(bytes, dtype_code);
+        self.runtime
+            .set_data_bytes(node_id, typed.bytes, typed.dtype);
         Ok(())
     }
 
@@ -537,7 +402,16 @@ impl CompiledGraph {
         })
     }
 
-    /// Get output tensor data by name.
+    /// Get the PT2 dtype codes for all outputs (in order).
+    #[getter]
+    fn output_dtypes(&self) -> Vec<u32> {
+        self.output_dtypes
+            .iter()
+            .map(|d| luminal_dtype_to_pt2_code(*d))
+            .collect()
+    }
+
+    /// Get output tensor data by name as f32 (copies to host).
     fn get_output(&self, name: &str) -> PyResult<Vec<f32>> {
         let node_id = self.tensor_ids.get(name).ok_or_else(|| {
             PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
@@ -545,6 +419,50 @@ impl CompiledGraph {
                 name
             ))
         })?;
-        Ok(self.runtime.get_f32(*node_id))
+        Ok(self.runtime.get_output_f32(*node_id))
+    }
+
+    /// Get output tensor data by name as i32 (copies to host).
+    fn get_output_i32(&self, name: &str) -> PyResult<Vec<i32>> {
+        let node_id = self.tensor_ids.get(name).ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                "Unknown output tensor: {}",
+                name
+            ))
+        })?;
+        Ok(self.runtime.get_output_i32(*node_id))
+    }
+
+    /// Get output tensor data by name as bool (copies to host).
+    fn get_output_bool(&self, name: &str) -> PyResult<Vec<bool>> {
+        let node_id = self.tensor_ids.get(name).ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                "Unknown output tensor: {}",
+                name
+            ))
+        })?;
+        Ok(self.runtime.get_output_bool(*node_id))
+    }
+
+    /// Copy output tensor data directly to a device pointer (DtoD).
+    /// Avoids the DtoH + HtoD round-trip of get_output() + .to(device).
+    /// Requires a GPU backend.
+    fn copy_output_to_device_ptr(&self, name: &str, dest_ptr: u64, n_bytes: usize) -> PyResult<()> {
+        if !self.runtime.supports_device_ptrs() {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "copy_output_to_device_ptr requires a GPU backend",
+            ));
+        }
+        let node_id = self.tensor_ids.get(name).ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                "Unknown output tensor: {}",
+                name
+            ))
+        })?;
+        unsafe {
+            self.runtime
+                .copy_output_to_device_ptr(*node_id, dest_ptr, n_bytes)
+        };
+        Ok(())
     }
 }

crates/luminal_python/rust/src/dispatch.rs

@@ -1,248 +0,0 @@
-use std::collections::HashMap;
-
-use luminal::{prelude::*, shape::Expression};
-use onnx_protobuf::NodeProto;
-
-use crate::ops_parse::*;
-
-pub fn process_onnx_nodes(
-    nodes: &[NodeProto],
-    tensors: &mut HashMap<String, GraphTensor>,
-    cx: &mut Graph,
-    weight_data: &mut Vec<(String, Vec<f32>)>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-    shape_exprs: &mut HashMap<String, Vec<Expression>>,
-) -> Result<(), String> {
-    for node in nodes {
-        match node.op_type.as_str() {
-            "Add" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Add",
-                |a, b| a + b,
-                shape_exprs,
-                known_values,
-            )?,
-            "Mod" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Mod",
-                |a, b| a % b,
-                shape_exprs,
-                known_values,
-            )?,
-            "Sub" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Sub",
-                |a, b| a - b,
-                shape_exprs,
-                known_values,
-            )?,
-            "Mul" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Mul",
-                |a, b| a * b,
-                shape_exprs,
-                known_values,
-            )?,
-            "Div" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Div",
-                |a, b| a / b,
-                shape_exprs,
-                known_values,
-            )?,
-            "Sqrt" => parse_unary_op(node, tensors, "Sqrt", |a| a.sqrt())?,
-            "Transpose" => parse_transpose_node(node, tensors)?,
-            "Concat" => parse_concat_node(node, tensors, shape_exprs, known_values)?,
-            "Floor" => parse_floor_node(node, tensors)?,
-            "Ceil" => parse_ceil_node(node, tensors)?,
-            "Sin" => parse_unary_op(node, tensors, "Sin", |a| a.sin())?,
-            "Neg" => parse_unary_op(node, tensors, "Neg", |a| -a)?,
-            "Cos" => parse_unary_op(node, tensors, "Cos", |a| a.cos())?,
-            "Pow" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Pow",
-                |a, b| a.pow(b),
-                shape_exprs,
-                known_values,
-            )?,
-            "Sigmoid" => parse_unary_op(node, tensors, "Sigmoid", |a| a.sigmoid())?,
-            "Tanh" => parse_unary_op(node, tensors, "Tanh", |a| a.tanh())?,
-            "Relu" => parse_unary_op(node, tensors, "Relu", |a| a.relu())?,
-            "Softmax" => parse_softmax_node(node, tensors)?,
-            "Abs" => parse_unary_op(node, tensors, "Abs", |a| a.abs())?,
-            "Reciprocal" => parse_unary_op(node, tensors, "Reciprocal", |a| a.reciprocal())?,
-            "Clip" => parse_clip_node(node, tensors, known_values)?,
-            "Equal" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Equal",
-                |a, b| a.eq(b),
-                shape_exprs,
-                known_values,
-            )?,
-            "Where" => parse_where_node(node, tensors)?,
-            "Constant" => {
-                parse_constant_node(node, tensors, cx, weight_data, known_values, shape_exprs)?
-            }
-            "ConstantOfShape" => {
-                parse_constant_of_shape(node, tensors, cx, weight_data, known_values, shape_exprs)?
-            }
-            "Cast" => parse_cast_node(node, tensors, weight_data, known_values, shape_exprs)?,
-            "MatMul" => parse_matmul_node(node, tensors)?,
-            "Reshape" => parse_reshape_node(node, tensors, known_values, shape_exprs)?,
-            "Shape" => parse_shape_node(node, tensors, cx, weight_data, known_values, shape_exprs)?,
-            "Gather" => {
-                parse_gather_node(node, tensors, cx, weight_data, known_values, shape_exprs)?
-            }
-            "GatherND" => parse_gathernd_node(node, tensors, cx, weight_data, known_values)?,
-            "Less" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Less",
-                |a, b| a.lt(b),
-                shape_exprs,
-                known_values,
-            )?,
-            "Greater" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Greater",
-                |a, b| b.lt(a),
-                shape_exprs,
-                known_values,
-            )?,
-            "LessOrEqual" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "LessOrEqual",
-                |a, b| a.le(b),
-                shape_exprs,
-                known_values,
-            )?,
-            "GreaterOrEqual" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "GreaterOrEqual",
-                |a, b| a.ge(b),
-                shape_exprs,
-                known_values,
-            )?,
-            "Not" => parse_not_node(node, tensors)?,
-            "And" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "And",
-                |a, b| a.cast(DType::F32) * b.cast(DType::F32),
-                shape_exprs,
-                known_values,
-            )?,
-            "Or" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Or",
-                |a, b| (a.cast(DType::F32) + b.cast(DType::F32)).minimum_f32(1.0),
-                shape_exprs,
-                known_values,
-            )?,
-            "Xor" => parse_binary_broadcast_op(
-                node,
-                tensors,
-                "Xor",
-                |a, b| a.ne(b),
-                shape_exprs,
-                known_values,
-            )?,
-            "Min" => parse_variadic_broadcast_op(
-                node,
-                tensors,
-                "Min",
-                |a, b| a.minimum(b),
-                shape_exprs,
-                known_values,
-            )?,
-            "Max" => parse_variadic_broadcast_op(
-                node,
-                tensors,
-                "Max",
-                |a, b| a.maximum(b),
-                shape_exprs,
-                known_values,
-            )?,
-            "Identity" => parse_identity(node, tensors, known_values, shape_exprs)?,
-            "Unsqueeze" => parse_unsqueeze_node(node, tensors, known_values, shape_exprs)?,
-            "Squeeze" => parse_squeeze_node(node, tensors, known_values, shape_exprs)?,
-            "ReduceSum" => parse_reduce_op(
-                node,
-                tensors,
-                known_values,
-                "ReduceSum",
-                |t, axes| t.sum(axes),
-                |flat, _n| flat.sum(1),
-            )?,
-            "ReduceMax" => parse_reduce_op(
-                node,
-                tensors,
-                known_values,
-                "ReduceMax",
-                |t, axes| t.max(axes),
-                |flat, _n| flat.max(1),
-            )?,
-            "ReduceMin" => parse_reduce_op(
-                node,
-                tensors,
-                known_values,
-                "ReduceMin",
-                |t, axes| t.min(axes),
-                |flat, _n| flat.min(1),
-            )?,
-            "ReduceMean" => parse_reduce_op(
-                node,
-                tensors,
-                known_values,
-                "ReduceMean",
-                |t, axes| t.mean(axes),
-                |flat, n| flat.sum(1) / n as f32,
-            )?,
-            "Trilu" => parse_trilu_node(node, tensors, cx, known_values)?,
-            "GatherElements" => parse_gather_elements_node(node, tensors)?,
-            "ScatterElements" => parse_scatter_elements_node(node, tensors)?,
-            "ScatterND" => parse_scatter_nd_node(node, tensors)?,
-            "Expand" => parse_expand_node(node, tensors, known_values, shape_exprs)?,
-            "IsNaN" => parse_unary_op(node, tensors, "IsNaN", |a| a.ne(a))?,
-            "LayerNormalization" => parse_layernorm_node(node, tensors)?,
-            "Gemm" => parse_gemm_node(node, tensors)?,
-            "Erf" => parse_erf_node(node, tensors)?,
-            "Slice" => parse_slice_node(node, tensors, known_values, shape_exprs)?,
-            "Split" => parse_split_node(node, tensors, known_values)?,
-            "TopK" => parse_topk_node(node, tensors, known_values)?,
-            "OneHot" => parse_onehot_node(node, tensors, known_values)?,
-            "Range" => parse_range_node(node, tensors, cx, weight_data, known_values, shape_exprs)?,
-            "CumSum" => parse_cumsum_node(node, tensors, known_values)?,
-            "Gelu" => parse_unary_op(node, tensors, "Gelu", |a| a.gelu())?,
-            "Conv" => parse_conv_node(node, tensors)?,
-            "Pad" => parse_pad_node(node, tensors, known_values)?,
-            "Resize" => parse_resize_node(node, tensors, known_values)?,
-            "Tile" => parse_tile_node(node, tensors, known_values)?,
-            "ReduceL2" => parse_reduce_op(
-                node,
-                tensors,
-                known_values,
-                "ReduceL2",
-                |t, axes| (t * t).sum(axes).sqrt(),
-                |flat, _n| (flat * flat).sum(1).sqrt(),
-            )?,
-            "GroupNormalization" => parse_group_norm_node(node, tensors)?,
-            _ => {
-                panic!("Missing Node {}", node.op_type)
-            }
-        }
-    }
-
-    Ok(())
-}

crates/luminal_python/rust/src/lib.rs

@@ -1,8 +1,5 @@
 mod compiled_graph;
-mod dispatch;
-mod ops_parse;
-mod runtime;
-mod util;
+pub mod typed_data;
 
 // PT2 modules
 mod pt2_compiled_model;
@@ -12,82 +9,42 @@ mod pt2_util;
 mod translator;
 
 use compiled_graph::CompiledGraph;
-use onnx_protobuf::ModelProto;
-use protobuf::Message;
-use pt2_compiled_model::compile_pt2;
+use pt2_compiled_model::process_pt2;
 use pyo3::prelude::*;
-use std::fs;
-use std::path::Path;
-
-fn validate_backend(backend: &str) -> PyResult<()> {
-    match backend {
-        "native" => Ok(()),
-        #[cfg(feature = "cuda")]
-        "cuda" => Ok(()),
-        #[cfg(not(feature = "cuda"))]
-        "cuda" => Err(pyo3::exceptions::PyValueError::new_err(
-            "CUDA backend requested, but this luminal extension was built without the `cuda` feature. Rebuild with `maturin develop --features cuda -r` or use backend='native'.",
-        )),
-        _ => {
-            #[cfg(feature = "cuda")]
-            {
-                Err(pyo3::exceptions::PyValueError::new_err(format!(
-                    "Invalid backend '{}'. Must be 'native' or 'cuda'",
-                    backend
-                )))
-            }
-            #[cfg(not(feature = "cuda"))]
-            {
-                Err(pyo3::exceptions::PyValueError::new_err(format!(
-                    "Invalid backend '{}'. This build only supports 'native'. Rebuild with the `cuda` feature to enable 'cuda'.",
-                    backend
-                )))
-            }
-        }
-    }
-}
-
-#[pyfunction]
-#[pyo3(signature = (path, backend="native"))]
-fn process_onnx(path: &str, backend: &str) -> PyResult<CompiledGraph> {
-    validate_backend(backend)?;
-
-    parse_onnx(path, backend).map_err(pyo3::exceptions::PyRuntimeError::new_err)
-}
-
-fn parse_onnx(path: &str, backend: &str) -> Result<CompiledGraph, String> {
-    let data = fs::read(path).map_err(|e| format!("Failed to read file: {}", e))?;
-    let model_directory = Path::new(path).parent().unwrap_or(Path::new("."));
-    let model = ModelProto::parse_from_bytes(&data)
-        .map_err(|e| format!("Failed to parse Onnx Model: {}", e))?;
-
-    let opset_version = model
-        .opset_import
-        .iter()
-        .find(|entry| entry.domain.is_empty())
-        .map(|entry| entry.version);
-
-    match opset_version {
-        Some(20) => {}
-        Some(v) => {
-            return Err(format!(
-                "Unsupported ONNX opset version {v}. Only opset 20 is supported."
-            ));
-        }
-        None => {
-            return Err(
-                "No ONNX opset version found in model. Only opset 20 is supported.".to_string(),
-            );
-        }
-    }
-
-    CompiledGraph::parse_graph(model, model_directory, backend)
-}
+use pyo3::types::PyCapsule;
 
 #[pymodule]
 fn luminal(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(process_onnx, m)?)?;
-    m.add_function(wrap_pyfunction!(compile_pt2, m)?)?;
+    m.add_function(wrap_pyfunction!(process_pt2, m)?)?;
     m.add_class::<CompiledGraph>()?;
+    m.add_function(wrap_pyfunction!(_native_factory_capsule, m)?)?;
+    #[cfg(feature = "cuda")]
+    m.add_function(wrap_pyfunction!(_cuda_lite_factory_capsule, m)?)?;
     Ok(())
 }
+
+// ---------------------------------------------------------------------------
+// Factory capsule helpers
+// ---------------------------------------------------------------------------
+
+/// Wrapper to put a function pointer into a PyCapsule.
+#[allow(dead_code)]
+struct FnPtrWrapper(pub *const std::ffi::c_void);
+unsafe impl Send for FnPtrWrapper {}
+
+/// PyCapsule wrapping the native (CPU) backend factory.
+#[pyfunction]
+fn _native_factory_capsule<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyCapsule>> {
+    let fptr = ::luminal::dyn_backend::native_factory as *const std::ffi::c_void;
+    let name = ::luminal::dyn_backend::BACKEND_FACTORY_CAPSULE_NAME.to_owned();
+    PyCapsule::new(py, FnPtrWrapper(fptr), Some(name))
+}
+
+/// PyCapsule wrapping the cuda_lite backend factory.
+#[cfg(feature = "cuda")]
+#[pyfunction]
+fn _cuda_lite_factory_capsule<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyCapsule>> {
+    let fptr = luminal_cuda_lite::dyn_backend::cuda_lite_factory as *const std::ffi::c_void;
+    let name = ::luminal::dyn_backend::BACKEND_FACTORY_CAPSULE_NAME.to_owned();
+    PyCapsule::new(py, FnPtrWrapper(fptr), Some(name))
+}

crates/luminal_python/rust/src/ops_parse/binary.rs

@@ -1,187 +0,0 @@
-use std::collections::HashMap;
-
-use luminal::{
-    prelude::{tracing::trace, *},
-    shape::Expression,
-};
-use onnx_protobuf::NodeProto;
-
-use crate::util::{broadcast_to_expr, compute_broadcast_shape_expr};
-
-/// Handle Where node: conditional select — output[i] = condition[i] ? x[i] : y[i]
-///
-/// ONNX Where uses numpy-style broadcasting across all three inputs.
-pub fn parse_where_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    assert!(node.input.len() == 3, "Where should have 3 inputs");
-    let condition = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Where: missing condition tensor '{}'", node.input[0]))?;
-    let x = *tensors
-        .get(&node.input[1])
-        .ok_or_else(|| format!("Where: missing X tensor '{}'", node.input[1]))?;
-    let y = *tensors
-        .get(&node.input[2])
-        .ok_or_else(|| format!("Where: missing Y tensor '{}'", node.input[2]))?;
-
-    let output_name = &node.output[0];
-
-    // ONNX Where broadcasts all 3 inputs to a common shape
-    let bc_shape = compute_broadcast_shape_expr(
-        &condition.dims(),
-        &compute_broadcast_shape_expr(&x.dims(), &y.dims()),
-    );
-    let condition = broadcast_to_expr(condition, &bc_shape);
-    let x = broadcast_to_expr(x, &bc_shape);
-    let y = broadcast_to_expr(y, &bc_shape);
-
-    let result = x.cond(condition, y);
-    tensors.insert(output_name.clone(), result);
-    Ok(())
-}
-
-pub fn parse_binary_broadcast_op(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    op_name: &str,
-    op: impl Fn(GraphTensor, GraphTensor) -> GraphTensor,
-    shape_exprs: &mut HashMap<String, Vec<Expression>>,
-    known_values: &HashMap<String, Vec<f32>>,
-) -> Result<(), String> {
-    trace!("Starting parse: {} Node", op_name);
-    assert!(
-        node.input.len() == 2,
-        "{} should have 2 inputs, got {}",
-        op_name,
-        node.input.len()
-    );
-    assert!(
-        node.output.len() == 1,
-        "{} should have 1 output, got {}",
-        op_name,
-        node.output.len()
-    );
-    // Shape-only path: if any input is shape-only (not in tensors), do Expression arithmetic
-    let a_missing = !tensors.contains_key(&node.input[0]);
-    let b_missing = !tensors.contains_key(&node.input[1]);
-    if a_missing || b_missing {
-        // At least one input is shape-only. Do shape_exprs arithmetic and return.
-        let se_a = shape_exprs.get(&node.input[0]).cloned().or_else(|| {
-            known_values
-                .get(&node.input[0])
-                .map(|kv| kv.iter().map(|&v| Expression::from(v as usize)).collect())
-        });
-        let se_b = shape_exprs.get(&node.input[1]).cloned().or_else(|| {
-            known_values
-                .get(&node.input[1])
-                .map(|kv| kv.iter().map(|&v| Expression::from(v as usize)).collect())
-        });
-        if let (Some(se_a), Some(se_b)) = (se_a, se_b)
-            && se_a.len() == 1
-            && se_b.len() == 1
-        {
-            let result_expr = match op_name {
-                "Add" => Some(se_a[0] + se_b[0]),
-                "Sub" => Some(se_a[0] - se_b[0]),
-                "Mul" => Some(se_a[0] * se_b[0]),
-                "Div" => Some(se_a[0] / se_b[0]),
-                _ => None,
-            };
-            if let Some(expr) = result_expr {
-                shape_exprs.insert(node.output[0].clone(), vec![expr]);
-            }
-        }
-        trace!("Finished parse: {} Node (shape-only)", op_name);
-        return Ok(());
-    }
-
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("{}: missing input '{}'", op_name, node.input[0]))?;
-    let b = *tensors
-        .get(&node.input[1])
-        .ok_or_else(|| format!("{}: missing input '{}'", op_name, node.input[1]))?;
-    let broadcast_shape = compute_broadcast_shape_expr(&a.dims(), &b.dims());
-    let a_bc = broadcast_to_expr(a, &broadcast_shape);
-    let b_bc = broadcast_to_expr(b, &broadcast_shape);
-    let result = op(a_bc, b_bc);
-    tensors.insert(node.output[0].clone(), result);
-
-    // Propagate shape_exprs for scalar shape arithmetic (e.g., Add(1, seq_len))
-    // At least one input must be in shape_exprs; the other can come from known_values.
-    let has_shape_expr =
-        shape_exprs.contains_key(&node.input[0]) || shape_exprs.contains_key(&node.input[1]);
-    if has_shape_expr {
-        let se_a = shape_exprs.get(&node.input[0]).cloned().or_else(|| {
-            known_values
-                .get(&node.input[0])
-                .map(|kv| kv.iter().map(|&v| Expression::from(v as usize)).collect())
-        });
-        let se_b = shape_exprs.get(&node.input[1]).cloned().or_else(|| {
-            known_values
-                .get(&node.input[1])
-                .map(|kv| kv.iter().map(|&v| Expression::from(v as usize)).collect())
-        });
-        if let (Some(se_a), Some(se_b)) = (se_a, se_b)
-            && se_a.len() == 1
-            && se_b.len() == 1
-        {
-            let result_expr = match op_name {
-                "Add" => Some(se_a[0] + se_b[0]),
-                "Sub" => Some(se_a[0] - se_b[0]),
-                "Mul" => Some(se_a[0] * se_b[0]),
-                "Div" => Some(se_a[0] / se_b[0]),
-                _ => None,
-            };
-            if let Some(expr) = result_expr {
-                shape_exprs.insert(node.output[0].clone(), vec![expr]);
-            }
-        }
-    }
-
-    trace!("Finished parse: {} Node", op_name);
-    Ok(())
-}
-
-pub fn parse_variadic_broadcast_op(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    op_name: &str,
-    op: impl Fn(GraphTensor, GraphTensor) -> GraphTensor,
-    _shape_exprs: &mut HashMap<String, Vec<Expression>>,
-    _known_values: &HashMap<String, Vec<f32>>,
-) -> Result<(), String> {
-    trace!("Starting parse: {} Node", op_name);
-    assert!(
-        node.input.len() >= 2,
-        "{} needs at least two inputs, got {}",
-        op_name,
-        node.input.len()
-    );
-    assert!(
-        node.output.len() == 1,
-        "{} nodes only have one output, got {}",
-        op_name,
-        node.output.len()
-    );
-
-    let mut result = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("{}: missing input tensor '{}'", op_name, node.input[0]))?;
-
-    for input_name in &node.input[1..] {
-        let rhs = *tensors
-            .get(input_name)
-            .ok_or_else(|| format!("{}: missing input tensor '{}'", op_name, input_name))?;
-        let broadcast_shape = compute_broadcast_shape_expr(&result.dims(), &rhs.dims());
-        let lhs_bc = broadcast_to_expr(result, &broadcast_shape);
-        let rhs_bc = broadcast_to_expr(rhs, &broadcast_shape);
-        result = op(lhs_bc, rhs_bc);
-    }
-
-    tensors.insert(node.output[0].clone(), result);
-    trace!("Finished parse: {} Node", op_name);
-    Ok(())
-}

crates/luminal_python/rust/src/ops_parse/convolution.rs

@@ -1,194 +0,0 @@
-use std::collections::HashMap;
-
-use luminal::{
-    prelude::{tracing::trace, *},
-    shape::Expression,
-};
-use onnx_protobuf::NodeProto;
-
-use crate::util::get_int_attr;
-
-/// Get an integer-list attribute from a node, with a default value applied per element.
-fn get_ints_attr(node: &NodeProto, name: &str, default_elem: i64, spatial: usize) -> Vec<usize> {
-    for attr in &node.attribute {
-        if attr.name == name {
-            return attr.ints.iter().map(|&v| v as usize).collect();
-        }
-    }
-    vec![default_elem as usize; spatial]
-}
-
-/// Parse an ONNX Conv node.
-///
-/// Supports N-dimensional convolution (1D, 2D, 3D) with group=1.
-/// Uses the unfold-based approach from `luminal_nn::ConvND`.
-///
-/// Input layout: [batch, C_in, spatial...]
-/// Weight layout: [C_out, C_in/group, kernel...]
-/// Optional bias: [C_out]
-pub fn parse_conv_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Starting parse: Conv Node");
-
-    assert!(
-        node.input.len() >= 2,
-        "Conv needs at least 2 inputs (X, W), got {}",
-        node.input.len()
-    );
-
-    let x = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Conv: missing input X '{}'", node.input[0]))?;
-    let w = *tensors
-        .get(&node.input[1])
-        .ok_or_else(|| format!("Conv: missing weight W '{}'", node.input[1]))?;
-    let bias = if node.input.len() > 2 && !node.input[2].is_empty() {
-        Some(
-            *tensors
-                .get(&node.input[2])
-                .ok_or_else(|| format!("Conv: missing bias B '{}'", node.input[2]))?,
-        )
-    } else {
-        None
-    };
-
-    let x_dims = x.dims();
-    let w_dims = w.dims();
-    let rank = x_dims.len();
-    assert!(
-        rank >= 3,
-        "Conv: input must be at least 3D (batch, channels, spatial...), got {rank}D"
-    );
-
-    let spatial = rank - 2; // number of spatial dimensions
-
-    // Parse attributes
-    let kernel_shape = get_ints_attr(node, "kernel_shape", 1, spatial);
-    let strides = get_ints_attr(node, "strides", 1, spatial);
-    let dilations = get_ints_attr(node, "dilations", 1, spatial);
-    let group = get_int_attr(node, "group", 1) as usize;
-
-    // Parse pads: ONNX format is [begin_0, begin_1, ..., end_0, end_1, ...]
-    let pads_flat = get_ints_attr(node, "pads", 0, 2 * spatial);
-    let mut pads_begin = vec![0usize; spatial];
-    let mut pads_end = vec![0usize; spatial];
-    if pads_flat.len() == 2 * spatial {
-        pads_begin[..spatial].copy_from_slice(&pads_flat[..spatial]);
-        pads_end[..spatial].copy_from_slice(&pads_flat[spatial..(spatial + spatial)]);
-    }
-
-    assert_eq!(
-        group, 1,
-        "Conv: only group=1 is currently supported, got {group}"
-    );
-
-    // Get channel dimensions
-    let ch_out = w_dims[0]
-        .to_usize()
-        .ok_or("Conv: weight C_out must be concrete")?;
-    let ch_in = x_dims[1]
-        .to_usize()
-        .ok_or("Conv: input C_in must be concrete")?;
-
-    let kernel_product: usize = kernel_shape.iter().product();
-
-    // Reshape weight from ONNX [C_out, C_in, *kernel] to [C_out, C_in * kernel_product]
-    let w_reshaped = {
-        let mut wt = w;
-        wt.shape = ShapeTracker::new(vec![ch_out, ch_in * kernel_product]);
-        wt
-    };
-
-    // Pad spatial dimensions
-    let mut padding: Vec<(Expression, Expression)> = vec![(0.into(), 0.into()); rank];
-    for i in 0..spatial {
-        let axis = 2 + i; // batch=0, channel=1, spatial starts at 2
-        padding[axis] = (
-            Expression::from(pads_begin[i]),
-            Expression::from(pads_end[i]),
-        );
-    }
-    let padded = x.pad(padding, 0.0);
-
-    // Build unfold parameters (ones for batch/channel, actual for spatial)
-    let mut kernel_full = vec![1usize; rank];
-    let mut stride_full = vec![1usize; rank];
-    let mut dilation_full = vec![1usize; rank];
-    for i in 0..spatial {
-        let axis = 2 + i;
-        kernel_full[axis] = kernel_shape[i];
-        stride_full[axis] = strides[i];
-        dilation_full[axis] = dilations[i];
-    }
-
-    let unfolded = padded.unfold(kernel_full, stride_full, dilation_full);
-    // unfolded shape: [win_N, win_C, win_spatial..., k_batch=1, k_chan=1, k_spatial...]
-    //                  (2*rank dimensions total)
-
-    // Step 1: Permute to [N, win_spatial..., C_in, k_batch, k_chan, k_spatial...]
-    // This groups: batch | output spatial | channel+kernel (for merging)
-    let mut perm: Vec<usize> = Vec::with_capacity(2 * rank);
-    perm.push(0); // win_N (batch)
-    perm.extend(2..2 + spatial); // win_spatial dims
-    perm.push(1); // win_C (= C_in)
-    perm.extend(rank..2 * rank); // all kernel dims: k_batch=1, k_chan=1, k_spatial...
-    let permuted = unfolded.permute(perm);
-
-    // Step 2: Capture output spatial dimensions (win_spatial sizes)
-    let output_spatial_dims: Vec<Expression> = permuted.dims()[1..1 + spatial].to_vec();
-
-    // Step 3: Merge all channel+kernel dims into one (C_in * kernel_product)
-    // From index (1+spatial) to end there are (1 + 2 + spatial) dims to merge
-    let mut patches = permuted;
-    let target_before_spatial_merge = 2 + spatial; // [N, spatial..., merged_patch]
-    while patches.dims().len() > target_before_spatial_merge {
-        let last = patches.dims().len();
-        patches = patches.merge_dims(last - 2, last - 1);
-    }
-    // patches: [N, spatial_0, ..., spatial_{s-1}, C_in * kernel_product]
-
-    // Step 4: Merge spatial dims into one
-    for _ in 1..spatial {
-        patches = patches.merge_dims(1, 2);
-    }
-    // patches: [N, spatial_product, C_in * kernel_product]
-
-    // Step 5: Matmul with weight
-    let mut out = patches.matmul(w_reshaped.permute((1, 0)));
-    // out: [N, spatial_product, C_out]
-
-    // Step 6: Restore spatial dimensions via split_dims
-    // Split from innermost spatial dim first (reverse order, skip outermost)
-    for i in (1..spatial).rev() {
-        out = out.split_dims(1, output_spatial_dims[i]);
-    }
-    // out: [N, spatial_0, spatial_1, ..., spatial_{s-1}, C_out]
-
-    // Step 7: Move C_out from last position to position 1 (after batch)
-    let mut final_order: Vec<usize> = Vec::with_capacity(2 + spatial);
-    final_order.push(0); // batch
-    final_order.push(1 + spatial); // C_out
-    final_order.extend(1..1 + spatial); // spatial dims
-    out = out.permute(final_order);
-    // out: [N, C_out, spatial_0, ..., spatial_{s-1}]
-
-    // Add bias if present: bias shape [C_out], broadcast to [1, C_out, 1, 1, ...]
-    if let Some(b) = bias {
-        let mut bias_expanded = b;
-        // Expand to [1, C_out, 1, 1, ...]
-        bias_expanded = bias_expanded.expand_dim(0, 1); // batch dim
-        for i in 0..spatial {
-            let out_dims = out.dims();
-            let spatial_size = out_dims[2 + i];
-            bias_expanded = bias_expanded.expand_dim(2 + i, spatial_size);
-        }
-        out += bias_expanded;
-    }
-
-    tensors.insert(node.output[0].clone(), out);
-
-    trace!("Finished parse: Conv Node");
-    Ok(())
-}

crates/luminal_python/rust/src/ops_parse/matmul.rs

@@ -1,70 +0,0 @@
-use std::collections::HashMap;
-
-use luminal::prelude::{tracing::trace, *};
-use onnx_protobuf::NodeProto;
-
-use crate::util::{broadcast_to_expr, get_float_attr, get_int_attr};
-
-pub fn parse_matmul_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Started parse: MatMul Node");
-    assert!(node.input.len() == 2, "MatMul should have exactly 2 inputs");
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("MatMul: missing input tensor '{}'", node.input[0]))?;
-    let b = *tensors
-        .get(&node.input[1])
-        .ok_or_else(|| format!("MatMul: missing input tensor '{}'", node.input[1]))?;
-
-    //TODO: enforce some kind of check here that they are broadcastable
-    let result = a.matmul(b);
-    let output_name = &node.output[0];
-    tensors.insert(output_name.clone(), result);
-    trace!("Finished parse: MatMul Node");
-    Ok(())
-}
-
-/// Handle Gemm node: Y = alpha * (transA ? A.T : A) @ (transB ? B.T : B) + beta * C
-///
-/// Attributes: transA (default 0), transB (default 0), alpha (default 1.0), beta (default 1.0)
-/// Input C (bias) is optional.
-pub fn parse_gemm_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Started parse: Gemm Node");
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Gemm: missing input A '{}'", node.input[0]))?;
-    let b = *tensors
-        .get(&node.input[1])
-        .ok_or_else(|| format!("Gemm: missing input B '{}'", node.input[1]))?;
-
-    let trans_a = get_int_attr(node, "transA", 0) != 0;
-    let trans_b = get_int_attr(node, "transB", 0) != 0;
-    let alpha = get_float_attr(node, "alpha", 1.0);
-    let beta = get_float_attr(node, "beta", 1.0);
-
-    let a_mat = if trans_a { a.permute(vec![1, 0]) } else { a };
-    let b_mat = if trans_b { b.permute(vec![1, 0]) } else { b };
-
-    let mut result = a_mat.matmul(b_mat);
-    if alpha != 1.0 {
-        result *= alpha;
-    }
-
-    if node.input.len() > 2 && !node.input[2].is_empty() {
-        let c = *tensors
-            .get(&node.input[2])
-            .ok_or_else(|| format!("Gemm: missing bias C '{}'", node.input[2]))?;
-        let c_scaled = if beta != 1.0 { c * beta } else { c };
-        let result_shape = result.dims();
-        result += broadcast_to_expr(c_scaled, &result_shape);
-    }
-
-    tensors.insert(node.output[0].clone(), result);
-    trace!("Finished parse: Gemm Node");
-    Ok(())
-}

crates/luminal_python/rust/src/ops_parse/mod.rs

@@ -1,15 +0,0 @@
-pub mod binary;
-pub mod convolution;
-pub mod matmul;
-pub mod movement;
-pub mod reduction;
-pub mod tensor;
-pub mod unary;
-
-pub use binary::*;
-pub use convolution::*;
-pub use matmul::*;
-pub use movement::*;
-pub use reduction::*;
-pub use tensor::*;
-pub use unary::*;

crates/luminal_python/rust/src/ops_parse/reduction.rs

@@ -1,172 +0,0 @@
-use std::collections::HashMap;
-
-use luminal::prelude::{tracing::trace, *};
-use onnx_protobuf::NodeProto;
-
-use crate::util::get_int_attr;
-
-/// Handle TopK node: return the top-k values and indices along an axis.
-///
-/// output[0] = values (F32), output[1] = indices (Int, can be empty/unused).
-/// For largest=true (default): uses topk_indexes + gather_elements.
-/// For largest=false: uses argsort(ascending).slice_along(..k) + gather_elements.
-/// Indices output is stored as-is (Int dtype); downstream Cast handles F32 conversion.
-/// The "sorted" attribute is ignored — output is always sorted.
-pub fn parse_topk_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-) -> Result<(), String> {
-    let x = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("TopK: missing input '{}'", node.input[0]))?;
-    let k = known_values
-        .get(&node.input[1])
-        .ok_or("TopK: k must be constant")?[0] as usize;
-
-    let rank = x.dims().len() as i64;
-    let raw_axis = get_int_attr(node, "axis", -1);
-    let axis = if raw_axis < 0 {
-        (raw_axis + rank) as usize
-    } else {
-        raw_axis as usize
-    };
-
-    let largest = get_int_attr(node, "largest", 1) != 0;
-
-    // Compute full argsort, then gather all sorted values, then slice both to top-k.
-    // This avoids passing a non-contiguous sliced index tensor into gather_elements,
-    // which triggers a CUDA kernel bug when data and index sizes differ along the axis.
-    let full_argsort = x.argsort(axis, largest);
-    let indices = full_argsort.slice_along(..k, axis);
-    let values = x.gather_elements(full_argsort, axis).slice_along(..k, axis);
-
-    // ONNX output[0] = values, output[1] = indices
-    if !node.output[0].is_empty() {
-        tensors.insert(node.output[0].clone(), values);
-    }
-    if node.output.len() > 1 && !node.output[1].is_empty() {
-        // Force materialization of Int indices; downstream Cast(INT64→FLOAT) handles the
-        // F32 conversion via the *1.0 workaround in parse_cast_node.
-        tensors.insert(node.output[1].clone(), indices * 1.0);
-    }
-    Ok(())
-}
-
-pub fn parse_reduce_op(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-    op_name: &str,
-    reduce_op: impl Fn(GraphTensor, Vec<usize>) -> GraphTensor,
-    all_axes_op: impl Fn(GraphTensor, usize) -> GraphTensor,
-) -> Result<(), String> {
-    trace!("Starting parse: {} Node", op_name);
-    assert!(
-        !node.input.is_empty(),
-        "{} should have at least 1 input",
-        op_name
-    );
-    assert!(
-        node.output.len() == 1,
-        "{} should have exactly 1 output",
-        op_name
-    );
-
-    let input = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("{}: missing input tensor '{}'", op_name, node.input[0]))?;
-
-    let keepdims = get_int_attr(node, "keepdims", 1) != 0;
-    let noop_with_empty_axes = get_int_attr(node, "noop_with_empty_axes", 0) != 0;
-
-    let ndim = input.dims().len();
-
-    // Resolve axes from second input (opset 13+) or from attribute (opset 11)
-    let raw_axes: Vec<i64> = if node.input.len() > 1 && !node.input[1].is_empty() {
-        let axes_vals = known_values.get(&node.input[1]).ok_or_else(|| {
-            format!(
-                "{}: axes input '{}' must be a known constant",
-                op_name, node.input[1]
-            )
-        })?;
-        axes_vals.iter().map(|&v| v as i64).collect()
-    } else if let Some(attr) = node.attribute.iter().find(|a| a.name == "axes") {
-        attr.ints.clone()
-    } else {
-        vec![]
-    };
-
-    let output_name = &node.output[0];
-
-    // Handle empty axes: noop or reduce all
-    let raw_axes: Vec<i64> = if raw_axes.is_empty() {
-        if noop_with_empty_axes {
-            tensors.insert(output_name.clone(), input);
-            trace!("Finished parse: {} Node (noop)", op_name);
-            return Ok(());
-        } else {
-            (0..ndim as i64).collect()
-        }
-    } else {
-        raw_axes
-    };
-
-    // Normalize negative axes and convert to usize
-    let mut normalized_axes: Vec<usize> = raw_axes
-        .iter()
-        .map(|&a| {
-            if a < 0 {
-                (ndim as i64 + a) as usize
-            } else {
-                a as usize
-            }
-        })
-        .collect();
-    normalized_axes.sort();
-    normalized_axes.dedup();
-
-    // Save original sorted axes for keepdims unsqueeze bookkeeping
-    let sorted_axes = normalized_axes.clone();
-
-    let input_dims = input.dims();
-
-    if normalized_axes.len() == ndim {
-        // All-axes reduction: flatten to [1, N] and reduce axis 1 → [1].
-        // luminal's Expression::product() returns 0 for empty iterators, so a reduce
-        // producing a 0-dim tensor causes CUDA to launch with grid (0,1,1), which is
-        // invalid. Using [1, N] → reduce(1) → [1] avoids this entirely.
-        let total: usize = input_dims
-            .iter()
-            .map(|d| d.to_usize().expect("reduce: dim must be concrete"))
-            .product();
-        let mut flat = input;
-        flat.shape = ShapeTracker::new(vec![1, total]);
-        let mut result = all_axes_op(flat, total);
-
-        if keepdims {
-            // Insert (ndim-1) additional size-1 dims to produce [1]*ndim
-            for i in 1..ndim {
-                result = result.unsqueeze(i);
-            }
-        }
-
-        tensors.insert(output_name.clone(), result);
-        trace!("Finished parse: {} Node (all-axes)", op_name);
-        return Ok(());
-    }
-
-    // Partial reduction: luminal's ToAxes API handles axis shifting internally
-    let mut result = reduce_op(input, normalized_axes);
-
-    // Re-insert size-1 dims at original positions (ascending order keeps positions correct)
-    if keepdims {
-        for &axis in &sorted_axes {
-            result = result.unsqueeze(axis);
-        }
-    }
-
-    tensors.insert(output_name.clone(), result);
-    trace!("Finished parse: {} Node", op_name);
-    Ok(())
-}

crates/luminal_python/rust/src/ops_parse/tensor.rs

@@ -1,453 +0,0 @@
-use std::collections::HashMap;
-
-use luminal::{
-    prelude::{tracing::trace, *},
-    shape::Expression,
-};
-use onnx_protobuf::NodeProto;
-
-use crate::util::{broadcast_to_expr, get_int_attr};
-
-/// Handle Constant node: creates a tensor from embedded data in the node attributes.
-///
-/// Supports FLOAT, INT64, INT32, and FLOAT64 data types (all converted to f32).
-/// The resulting tensor is registered as a known constant for downstream folding.
-pub fn parse_constant_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    cx: &mut Graph,
-    weight_data: &mut Vec<(String, Vec<f32>)>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-    shape_exprs: &mut HashMap<String, Vec<Expression>>,
-) -> Result<(), String> {
-    trace!("Starting parse: Constant Node");
-    assert!(
-        node.output.len() == 1,
-        "Constant should have exactly one output"
-    );
-
-    // Find the "value" attribute (type TENSOR)
-    let value_attr = node
-        .attribute
-        .iter()
-        .find(|a| a.name == "value")
-        .ok_or_else(|| "Constant node missing 'value' attribute".to_string())?;
-
-    let tensor_proto = value_attr
-        .t
-        .as_ref()
-        .ok_or_else(|| "Constant 'value' attribute has no TensorProto".to_string())?;
-
-    // Determine shape: empty dims = scalar = [1] for luminal
-    let shape: Vec<usize> = if tensor_proto.dims.is_empty() {
-        vec![1]
-    } else {
-        tensor_proto.dims.iter().map(|&d| d as usize).collect()
-    };
-
-    // Extract float data based on data_type
-    let floats: Vec<f32> = match tensor_proto.data_type {
-        1 => {
-            // FLOAT (f32)
-            if !tensor_proto.float_data.is_empty() {
-                tensor_proto.float_data.clone()
-            } else {
-                tensor_proto
-                    .raw_data
-                    .chunks_exact(4)
-                    .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
-                    .collect()
-            }
-        }
-        6 => {
-            // INT32
-            if !tensor_proto.int32_data.is_empty() {
-                tensor_proto.int32_data.iter().map(|&v| v as f32).collect()
-            } else {
-                tensor_proto
-                    .raw_data
-                    .chunks_exact(4)
-                    .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f32)
-                    .collect()
-            }
-        }
-        7 => {
-            // INT64
-            if !tensor_proto.int64_data.is_empty() {
-                tensor_proto.int64_data.iter().map(|&v| v as f32).collect()
-            } else {
-                tensor_proto
-                    .raw_data
-                    .chunks_exact(8)
-                    .map(|c| {
-                        i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) as f32
-                    })
-                    .collect()
-            }
-        }
-        dt => return Err(format!("Constant node: unsupported data_type {}", dt)),
-    };
-
-    let output_name = &node.output[0];
-    let tensor = cx.named_tensor(output_name.clone(), shape);
-    tensors.insert(output_name.clone(), tensor);
-    known_values.insert(output_name.clone(), floats.clone());
-    // Also propagate as concrete shape_exprs for downstream shape computation chains
-    shape_exprs.insert(
-        output_name.clone(),
-        floats
-            .iter()
-            .map(|&v| Expression::from(v as usize))
-            .collect(),
-    );
-    weight_data.push((output_name.clone(), floats));
-
-    trace!("Finished parse: Constant Node");
-    Ok(())
-}
-
-/// Handle Shape node: extract the shape of the input tensor as a 1D constant.
-///
-/// For static shapes, stores as known_values. For dynamic shapes (containing
-/// Expression variables), stores in shape_exprs for downstream shape computation chains.
-pub fn parse_shape_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    cx: &mut Graph,
-    weight_data: &mut Vec<(String, Vec<f32>)>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-    shape_exprs: &mut HashMap<String, Vec<Expression>>,
-) -> Result<(), String> {
-    trace!("Started parse: Shape");
-    assert!(node.input.len() == 1, "Shape should have exactly 1 input");
-    let input = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Shape: missing input tensor '{}'", node.input[0]))?;
-
-    let all_dims = input.dims();
-
-    // Handle start/end attributes (ONNX Shape opset 15+: extract a slice of dims)
-    let start = get_int_attr(node, "start", 0) as usize;
-    let end_attr = get_int_attr(node, "end", all_dims.len() as i64);
-    let end = if end_attr < 0 {
-        (all_dims.len() as i64 + end_attr) as usize
-    } else {
-        (end_attr as usize).min(all_dims.len())
-    };
-    let dims: Vec<Expression> = all_dims[start..end].to_vec();
-
-    let output_name = &node.output[0];
-
-    // Always store in shape_exprs (supports both concrete and symbolic dims)
-    shape_exprs.insert(output_name.clone(), dims.clone());
-
-    // For concrete dims, also store in known_values for backward compat
-    let all_concrete = dims.iter().all(|d| d.to_usize().is_some());
-    let shape_values: Vec<f32> = dims
-        .iter()
-        .map(|d| d.to_usize().unwrap_or(1) as f32)
-        .collect();
-
-    if all_concrete {
-        // Concrete shape: create tensor + known_values + weight_data
-        let tensor = cx.named_tensor(output_name.clone(), vec![shape_values.len()]);
-        tensors.insert(output_name.clone(), tensor);
-        known_values.insert(output_name.clone(), shape_values.clone());
-        weight_data.push((output_name.clone(), shape_values));
-    }
-    // For symbolic shapes, don't create a tensor — it's shape-only
-
-    trace!("Finished parse: Shape");
-    Ok(())
-}
-
-/// Handle ConstantOfShape node: creates a tensor of a given shape filled with a constant value.
-///
-/// The shape is taken from the input tensor (which must be a known constant).
-/// The fill value comes from the "value" attribute (default 0.0).
-pub fn parse_constant_of_shape(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    cx: &mut Graph,
-    weight_data: &mut Vec<(String, Vec<f32>)>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-    shape_exprs: &mut HashMap<String, Vec<Expression>>,
-) -> Result<(), String> {
-    trace!("Starting parse: ConstantOfShape Node");
-    assert!(
-        node.input.len() == 1,
-        "ConstantOfShape should have exactly one input (shape)"
-    );
-    assert!(
-        node.output.len() == 1,
-        "ConstantOfShape should have exactly one output"
-    );
-
-    // Extract fill value from "value" attribute (TensorProto scalar), default 0.0
-    let fill_value: f32 = node
-        .attribute
-        .iter()
-        .find(|a| a.name == "value")
-        .and_then(|attr| attr.t.as_ref())
-        .map(|tp| {
-            if !tp.float_data.is_empty() {
-                tp.float_data[0]
-            } else if !tp.int32_data.is_empty() {
-                tp.int32_data[0] as f32
-            } else if !tp.raw_data.is_empty() {
-                match tp.data_type {
-                    1 => f32::from_le_bytes([
-                        tp.raw_data[0],
-                        tp.raw_data[1],
-                        tp.raw_data[2],
-                        tp.raw_data[3],
-                    ]),
-                    6 => i32::from_le_bytes([
-                        tp.raw_data[0],
-                        tp.raw_data[1],
-                        tp.raw_data[2],
-                        tp.raw_data[3],
-                    ]) as f32,
-                    7 => i64::from_le_bytes([
-                        tp.raw_data[0],
-                        tp.raw_data[1],
-                        tp.raw_data[2],
-                        tp.raw_data[3],
-                        tp.raw_data[4],
-                        tp.raw_data[5],
-                        tp.raw_data[6],
-                        tp.raw_data[7],
-                    ]) as f32,
-                    _ => 0.0,
-                }
-            } else {
-                0.0
-            }
-        })
-        .unwrap_or(0.0);
-
-    let output_name = &node.output[0];
-
-    // Try shape_exprs first (for dynamic shapes), then known_values
-    if let Some(se) = shape_exprs.get(&node.input[0]) {
-        let shape: Vec<Expression> = se.clone();
-
-        // Check if all dims are concrete
-        if let Some(concrete) = shape
-            .iter()
-            .map(|e| e.to_usize())
-            .collect::<Option<Vec<usize>>>()
-        {
-            // Fully concrete: create named tensor with weight data
-            let numel: usize = concrete.iter().product();
-            let floats: Vec<f32> = vec![fill_value; numel];
-            let tensor = cx.named_tensor(output_name.clone(), concrete);
-            tensors.insert(output_name.clone(), tensor);
-            known_values.insert(output_name.clone(), floats.clone());
-            weight_data.push((output_name.clone(), floats));
-        } else {
-            // Dynamic shape: create scalar constant and broadcast to symbolic shape.
-            // The scalar always has concrete data (1 element), and the shape is
-            // resolved at runtime via ShapeTracker/dyn_map. Broadcast uses stride-0
-            // expansion, so only 1 float is needed in the backing buffer.
-            let scalar = cx.constant_float(fill_value);
-            let result = broadcast_to_expr(scalar, se);
-            tensors.insert(output_name.clone(), result);
-        }
-    } else {
-        let shape_values = known_values.get(&node.input[0]).ok_or_else(|| {
-            format!(
-                "ConstantOfShape: shape input '{}' must be a known constant or shape_expr",
-                node.input[0]
-            )
-        })?;
-        let shape: Vec<usize> = shape_values.iter().map(|&v| v as usize).collect();
-        let numel: usize = shape.iter().product();
-        let floats: Vec<f32> = vec![fill_value; numel];
-
-        let tensor = cx.named_tensor(output_name.clone(), shape);
-        tensors.insert(output_name.clone(), tensor);
-        known_values.insert(output_name.clone(), floats.clone());
-        weight_data.push((output_name.clone(), floats));
-    }
-
-    trace!("Finished parse: ConstantOfShape Node");
-    Ok(())
-}
-
-/// Handle Identity node: output is a direct alias of the input tensor.
-///
-/// Propagates known constant values for downstream constant folding.
-pub fn parse_identity(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-    shape_exprs: &mut HashMap<String, Vec<Expression>>,
-) -> Result<(), String> {
-    trace!("Starting parse: Identity Node");
-    assert!(node.input.len() == 1, "Identity should only have one input");
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Identity: missing input tensor '{}'", node.input[0]))?;
-
-    assert!(
-        node.output.len() == 1,
-        "Identity should only have a single output"
-    );
-
-    let output_name = &node.output[0];
-
-    // Force materialization using Expression-aware broadcast
-    let dims = a.dims();
-    let one = a.graph().constant_float(1.0);
-    let one_expanded = broadcast_to_expr(one, &dims);
-    let result = a * one_expanded;
-    tensors.insert(output_name.clone(), result);
-
-    // Propagate known values
-    if let Some(vals) = known_values.get(&node.input[0]).cloned() {
-        known_values.insert(output_name.clone(), vals);
-    }
-    // Propagate shape_exprs
-    if let Some(se) = shape_exprs.get(&node.input[0]).cloned() {
-        shape_exprs.insert(output_name.clone(), se);
-    }
-
-    trace!("Finished parse: Identity Node");
-    Ok(())
-}
-
-/// Handle Range node: creates a 1D tensor [start, start+delta, start+2*delta, ...] up to limit.
-///
-/// Used by dynamo ONNX export for generating position indices (arange).
-/// Supports Expression-based limits for dynamic sequence lengths.
-pub fn parse_range_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    cx: &mut Graph,
-    weight_data: &mut Vec<(String, Vec<f32>)>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-    shape_exprs: &mut HashMap<String, Vec<Expression>>,
-) -> Result<(), String> {
-    trace!("Starting parse: Range Node");
-    assert!(
-        node.input.len() == 3,
-        "Range needs 3 inputs: start, limit, delta"
-    );
-
-    let output_name = &node.output[0];
-
-    // Try to get concrete values from known_values first
-    let start_val = known_values
-        .get(&node.input[0])
-        .and_then(|v| v.first().copied());
-    let limit_val = known_values
-        .get(&node.input[1])
-        .and_then(|v| v.first().copied());
-    let delta_val = known_values
-        .get(&node.input[2])
-        .and_then(|v| v.first().copied());
-
-    // Also check shape_exprs for symbolic limit
-    let limit_expr = shape_exprs
-        .get(&node.input[1])
-        .and_then(|v| v.first().cloned());
-
-    let start = start_val.unwrap_or(0.0);
-    let delta = delta_val.unwrap_or(1.0);
-
-    if start == 0.0 && delta == 1.0 {
-        // Simple arange case — most common for position indices
-        if let Some(expr) = limit_expr {
-            // Dynamic limit: create arange with symbolic length
-            let tensor = cx.arange(expr);
-            // Cast to F32 (luminal arange returns Int dtype)
-            let result = tensor.cast(DType::F32);
-            tensors.insert(output_name.clone(), result);
-            shape_exprs.insert(output_name.clone(), vec![expr]);
-        } else if let Some(limit) = limit_val {
-            let n = limit as usize;
-            let floats: Vec<f32> = (0..n).map(|i| i as f32).collect();
-            let tensor = cx.named_tensor(output_name.clone(), vec![n]);
-            tensors.insert(output_name.clone(), tensor);
-            known_values.insert(output_name.clone(), floats.clone());
-            weight_data.push((output_name.clone(), floats));
-        } else {
-            return Err("Range: limit must be known or symbolic".to_string());
-        }
-    } else if let (Some(s), Some(l), Some(d)) = (start_val, limit_val, delta_val) {
-        // Fully concrete range
-        let mut floats = Vec::new();
-        let mut v = s;
-        while (d > 0.0 && v < l) || (d < 0.0 && v > l) {
-            floats.push(v);
-            v += d;
-        }
-        let tensor = cx.named_tensor(output_name.clone(), vec![floats.len()]);
-        tensors.insert(output_name.clone(), tensor);
-        known_values.insert(output_name.clone(), floats.clone());
-        weight_data.push((output_name.clone(), floats));
-    } else {
-        return Err("Range: cannot handle non-trivial dynamic ranges yet".to_string());
-    }
-
-    trace!("Finished parse: Range Node");
-    Ok(())
-}
-
-/// Handle CumSum node: cumulative sum along an axis.
-///
-/// For the simple case of axis=0 on a 1D tensor [0, 1, 2, ...] (position indices),
-/// the cumsum is equivalent to [0, 1, 3, 6, ...]. For dynamic ONNX graphs,
-/// this is typically used for position_ids computation.
-pub fn parse_cumsum_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-) -> Result<(), String> {
-    trace!("Starting parse: CumSum Node");
-    assert!(node.input.len() >= 2, "CumSum needs at least 2 inputs");
-
-    let input = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("CumSum: missing input '{}'", node.input[0]))?;
-
-    let axis_val = known_values
-        .get(&node.input[1])
-        .and_then(|v| v.first().copied())
-        .unwrap_or(0.0) as i64;
-
-    let dims = input.dims();
-    let ndim = dims.len();
-    let _axis = if axis_val < 0 {
-        (ndim as i64 + axis_val) as usize
-    } else {
-        axis_val as usize
-    };
-
-    // For constant folding
-    if let Some(vals) = known_values.get(&node.input[0]).cloned() {
-        let output_name = &node.output[0];
-        let mut cumsum = vals.clone();
-        // Simple 1D cumsum
-        if ndim == 1 {
-            for i in 1..cumsum.len() {
-                cumsum[i] += cumsum[i - 1];
-            }
-        }
-        known_values.insert(output_name.clone(), cumsum);
-        // Just alias the tensor (same shape)
-        tensors.insert(output_name.clone(), input);
-        trace!("Finished parse: CumSum Node (constant folded)");
-        return Ok(());
-    }
-
-    // For dynamic: cumsum is hard to express in luminal primitives.
-    // For the specific pattern used in Llama position_ids (cumsum of ones = arange),
-    // we just pass through since arange is already handled by Range node.
-    let output_name = &node.output[0];
-    tensors.insert(output_name.clone(), input);
-
-    trace!("Finished parse: CumSum Node");
-    Ok(())
-}

crates/luminal_python/rust/src/ops_parse/unary.rs

@@ -1,440 +0,0 @@
-use std::collections::HashMap;
-
-use luminal::{
-    prelude::{tracing::trace, *},
-    shape::Expression,
-};
-use onnx_protobuf::NodeProto;
-
-use crate::util::{broadcast_to_expr, get_float_attr, get_int_attr};
-
-/// Handle Softmax node: output = softmax(input[0], axis)
-///
-/// ONNX axis attribute defaults to -1 (last dimension, opset 13+).
-/// Negative axis is normalized against the input rank.
-pub fn parse_softmax_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Starting parse: Softmax Node");
-    assert!(
-        node.input.len() == 1,
-        "Softmax nodes need to have one input, {} where present",
-        node.input.len()
-    );
-    assert!(
-        node.output.len() == 1,
-        "Softmax nodes only have one output, {} where present",
-        node.output.len(),
-    );
-    let output_name = &node.output[0];
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Softmax: missing input tensor '{}'", node.input[0]))?;
-
-    let ndim = a.dims().len();
-    let raw_axis = get_int_attr(node, "axis", -1);
-    let axis = if raw_axis < 0 {
-        (ndim as i64 + raw_axis) as usize
-    } else {
-        raw_axis as usize
-    };
-
-    let result = a.softmax(axis);
-    tensors.insert(output_name.clone(), result);
-    trace!("Finished parse: Softmax Node");
-    Ok(())
-}
-
-/// Handle Not node: logical NOT — output = 1.0 - input[0]
-pub fn parse_not_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Starting parse: Not Node");
-    assert!(
-        node.input.len() == 1,
-        "Not nodes need to have one input {} where present",
-        node.input.len()
-    );
-    assert!(
-        node.output.len() == 1,
-        "Not nodes only have one output, {} where present",
-        node.output.len()
-    );
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Not: missing input tensor '{}'", node.input[0]))?;
-    let a_f32 = a.cast(DType::F32);
-    let result = 1.0_f32 - a_f32;
-    tensors.insert(node.output[0].clone(), result);
-    trace!("Finished parse: Not Node");
-    Ok(())
-}
-
-/// Handle Clip node: output = clip(input[0], min, max)
-///
-/// Equivalent to torch.clamp. min and max are optional tensor inputs
-/// (typically constants) residing in known_values.
-pub fn parse_clip_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    known_values: &HashMap<String, Vec<f32>>,
-) -> Result<(), String> {
-    trace!("Starting parse: Clip Node");
-    let output_name = &node.output[0];
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Clip: missing input tensor '{}'", node.input[0]))?;
-
-    // input[1] = min (optional), input[2] = max (optional)
-    let min_name = node.input.get(1).map(String::as_str).unwrap_or("");
-    let max_name = node.input.get(2).map(String::as_str).unwrap_or("");
-
-    let min_val = if min_name.is_empty() {
-        None
-    } else {
-        known_values.get(min_name).map(|v| v[0])
-    };
-    let max_val = if max_name.is_empty() {
-        None
-    } else {
-        known_values.get(max_name).map(|v| v[0])
-    };
-
-    let result = match (min_val, max_val) {
-        (Some(lo), Some(hi)) => a.clip(lo, hi),
-        (Some(lo), None) => a.maximum_f32(lo),
-        (None, Some(hi)) => a.minimum_f32(hi),
-        (None, None) => a,
-    };
-
-    tensors.insert(output_name.clone(), result);
-    trace!("Finished parse: Clip Node");
-    Ok(())
-}
-
-/// Handle Floor node: output = floor(input[0])
-///
-/// Implemented as: trunc(x) - (x < trunc(x) ? 1 : 0)
-/// where trunc is truncation toward zero via cast to Int then back to F32.
-/// This correctly handles negative non-integer values (e.g. floor(-1.5) = -2).
-pub fn parse_floor_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Starting parse: Floor Node");
-    assert!(
-        node.input.len() == 1,
-        "Floor nodes need to have one input {} where present",
-        node.input.len()
-    );
-    assert!(
-        node.output.len() == 1,
-        "Floor nodes only have one output, {} where present",
-        node.output.len(),
-    );
-    let output_name = &node.output[0];
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Floor: missing input tensor '{}'", node.input[0]))?;
-
-    // trunc(x): truncation toward zero
-    let trunc = a.cast(DType::Int).cast(DType::F32);
-    // For negative non-integers, x < trunc(x), so subtract 1
-    // Cast lt result (Bool) to F32 before arithmetic
-    let adjustment = a.lt(trunc).cast(DType::F32);
-    let result = trunc - adjustment;
-    tensors.insert(output_name.clone(), result);
-    trace!("Finished parse: Floor Node");
-
-    Ok(())
-}
-
-/// Handle Ceil node: output = ceil(input[0])
-///
-/// Implemented as: trunc(x) + (x > trunc(x) ? 1 : 0)
-/// where trunc is truncation toward zero via cast to Int then back to F32.
-/// This correctly handles positive non-integer values (e.g. ceil(1.5) = 2).
-pub fn parse_ceil_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Starting parse: Ceil Node");
-    assert!(
-        node.input.len() == 1,
-        "Ceil nodes need to have one input {} where present",
-        node.input.len()
-    );
-    assert!(
-        node.output.len() == 1,
-        "Ceil nodes only have one output, {} where present",
-        node.output.len(),
-    );
-    let output_name = &node.output[0];
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Ceil: missing input tensor '{}'", node.input[0]))?;
-
-    // trunc(x): truncation toward zero
-    let trunc = a.cast(DType::Int).cast(DType::F32);
-    // For positive non-integers, x > trunc(x), so add 1
-    let adjustment = a.gt(trunc).cast(DType::F32);
-    let result = trunc + adjustment;
-    tensors.insert(output_name.clone(), result);
-    trace!("Finished parse: Ceil Node");
-
-    Ok(())
-}
-
-pub fn parse_cast_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    weight_data: &mut Vec<(String, Vec<f32>)>,
-    known_values: &mut HashMap<String, Vec<f32>>,
-    shape_exprs: &mut HashMap<String, Vec<Expression>>,
-) -> Result<(), String> {
-    trace!("Starting parse: Cast Node");
-    assert!(node.input.len() == 1, "Cast should have exactly 1 input");
-    let input = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("Cast: missing input tensor '{}'", node.input[0]))?;
-
-    // ONNX data type enum → luminal DType
-    let to = get_int_attr(node, "to", 1);
-    let dtype = match to {
-        1 => DType::F32,     // FLOAT
-        10 => DType::F16,    // FLOAT16
-        16 => DType::Bf16,   // BFLOAT16
-        6 | 7 => DType::Int, // INT32, INT64
-        9 => DType::F32,     // BOOL → treat as F32 (0.0/1.0)
-        11 => DType::F32,    // DOUBLE → F32 (downcast)
-        _ => DType::F32,     // fallback
-    };
-
-    let cast_result = input.cast(dtype);
-    let output_name = &node.output[0];
-
-    let result = if cast_result.id == input.id {
-        input
-    } else {
-        cast_result
-    };
-
-    tensors.insert(output_name.clone(), result);
-
-    // Propagate known values (cast is a no-op for our f32 storage)
-    if let Some(vals) = known_values.get(&node.input[0]).cloned() {
-        let folded = if to == 9 {
-            vals.iter()
-                .map(|&v| if v != 0.0 { 1.0 } else { 0.0 })
-                .collect()
-        } else if to == 6 || to == 7 {
-            vals.iter().map(|&v| (v as i64) as f32).collect()
-        } else {
-            vals
-        };
-        known_values.insert(output_name.clone(), folded.clone());
-        weight_data.push((output_name.clone(), folded));
-    }
-    // Propagate shape_exprs
-    if let Some(se) = shape_exprs.get(&node.input[0]).cloned() {
-        shape_exprs.insert(output_name.clone(), se);
-    }
-
-    trace!("Finished parse: Cast Node");
-    Ok(())
-}
-
-pub fn parse_unary_op(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-    op_name: &str,
-    op: impl Fn(GraphTensor) -> GraphTensor,
-) -> Result<(), String> {
-    trace!("Starting parse: {} Node", op_name);
-    assert!(
-        node.input.len() == 1,
-        "{} should have 1 input, got {}",
-        op_name,
-        node.input.len()
-    );
-    assert!(
-        node.output.len() == 1,
-        "{} should have 1 output, got {}",
-        op_name,
-        node.output.len()
-    );
-    let a = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("{}: missing input tensor '{}'", op_name, node.input[0]))?;
-    let result = op(a);
-    tensors.insert(node.output[0].clone(), result);
-    trace!("Finished parse: {} Node", op_name);
-    Ok(())
-}
-
-/// Handle Erf node: output = erf(input[0])
-///
-/// Uses the Abramowitz & Stegun 7.1.26 polynomial approximation (max error < 1.5e-7):
-///   For x ≥ 0: erf(x) ≈ 1 - (a1·t + a2·t² + a3·t³ + a4·t⁴ + a5·t⁵) · exp(-x²)
-///   where t = 1 / (1 + 0.3275911·x)
-///     a1 =  0.254829592
-///     a2 = -0.284496736
-///     a3 =  1.421413741
-///     a4 = -1.453152027
-///     a5 =  1.061405429
-/// Extended to all x via odd symmetry: erf(-x) = -erf(x).
-pub fn parse_erf_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    parse_unary_op(node, tensors, "Erf", |x| {
-        let a = x.abs();
-        let t = (1.0_f32 + 0.3275911_f32 * a).reciprocal();
-        // Horner evaluation of a1*t + a2*t² + a3*t³ + a4*t⁴ + a5*t⁵
-        // poly = t*(a1 + t*(a2 + t*(a3 + t*(a4 + a5*t))))
-        let h = t * 1.061_405_4_f32 - 1.453_152_1_f32; // a4 + a5*t
-        let h = t * h + 1.421_413_8_f32;
-        let h = t * h - 0.284_496_72_f32;
-        let h = t * h + 0.254_829_6_f32;
-        let poly = t * h;
-        let erf_abs = 1.0_f32 - poly * (-a * a).exp();
-        x.sign() * erf_abs
-    })
-}
-
-/// Handle LayerNormalization node (opset 17).
-///
-/// Inputs: X (required), scale (required), bias (optional)
-/// Attributes: axis (default -1), epsilon (default 1e-5)
-/// Normalizes over axes [axis, axis+1, ..., rank-1], then applies scale and bias.
-/// Only output 0 (the normalized result) is wired; outputs 1/2 (mean, inv_std_var)
-/// are training-only and not supported for inference.
-pub fn parse_layernorm_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Starting parse: LayerNormalization Node");
-    let input = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("LayerNorm: missing input '{}'", node.input[0]))?;
-    let scale = *tensors
-        .get(&node.input[1])
-        .ok_or_else(|| format!("LayerNorm: missing scale '{}'", node.input[1]))?;
-
-    let ndim = input.dims().len();
-    let axis_raw = get_int_attr(node, "axis", -1);
-    let axis = if axis_raw < 0 {
-        (ndim as i64 + axis_raw) as usize
-    } else {
-        axis_raw as usize
-    };
-    let epsilon = get_float_attr(node, "epsilon", 1e-5);
-    let axes: Vec<usize> = (axis..ndim).collect();
-
-    let mut result = input.layer_norm(axes, epsilon);
-
-    // Apply scale (broadcast to input shape using Expression-aware broadcast)
-    let input_shape = input.dims();
-    result *= broadcast_to_expr(scale, &input_shape);
-
-    // Apply optional bias
-    if node.input.len() > 2 && !node.input[2].is_empty() {
-        let bias = *tensors
-            .get(&node.input[2])
-            .ok_or_else(|| format!("LayerNorm: missing bias '{}'", node.input[2]))?;
-        result += broadcast_to_expr(bias, &input_shape);
-    }
-
-    tensors.insert(node.output[0].clone(), result);
-    trace!("Finished parse: LayerNormalization Node");
-    Ok(())
-}
-
-/// Handle GroupNormalization node (opset 18).
-///
-/// Inputs: X [N, C, spatial...], scale [num_groups], bias [num_groups]
-/// Attributes: num_groups (required), epsilon (default 1e-5)
-///
-/// Normalizes over channels-per-group and spatial dims, then applies per-group scale/bias.
-/// Decomposed into: reshape [N, G, C/G, spatial...] -> layer_norm over [C/G, spatial...] ->
-/// reshape back to [N, C, spatial...] -> scale + bias (broadcast).
-pub fn parse_group_norm_node(
-    node: &NodeProto,
-    tensors: &mut HashMap<String, GraphTensor>,
-) -> Result<(), String> {
-    trace!("Starting parse: GroupNormalization Node");
-
-    assert!(
-        node.input.len() >= 3,
-        "GroupNormalization needs 3 inputs (X, scale, bias), got {}",
-        node.input.len()
-    );
-
-    let x = *tensors
-        .get(&node.input[0])
-        .ok_or_else(|| format!("GroupNorm: missing input X '{}'", node.input[0]))?;
-    let scale = *tensors
-        .get(&node.input[1])
-        .ok_or_else(|| format!("GroupNorm: missing scale '{}'", node.input[1]))?;
-    let bias = *tensors
-        .get(&node.input[2])
-        .ok_or_else(|| format!("GroupNorm: missing bias '{}'", node.input[2]))?;
-
-    let x_dims = x.dims();
-    let ndim = x_dims.len();
-    assert!(
-        ndim >= 3,
-        "GroupNorm: input must be at least 3D [N, C, spatial...], got {ndim}D"
-    );
-
-    let num_groups = get_int_attr(node, "num_groups", 1) as usize;
-    let epsilon = get_float_attr(node, "epsilon", 1e-5);
-
-    let n = x_dims[0]
-        .to_usize()
-        .expect("GroupNorm: batch must be concrete");
-    let c = x_dims[1]
-        .to_usize()
-        .expect("GroupNorm: channels must be concrete");
-    assert_eq!(
-        c % num_groups,
-        0,
-        "GroupNorm: channels {c} must be divisible by num_groups {num_groups}"
-    );
-    let cpg = c / num_groups; // channels per group
-
-    // Reshape X from [N, C, spatial...] to [N, G, C/G, spatial...]
-    let spatial_dims: Vec<Expression> = x_dims[2..].to_vec();
-    let mut reshaped = x;
-    let mut new_shape = vec![n, num_groups, cpg];
-    for d in &spatial_dims {
-        new_shape.push(
-            d.to_usize()
-                .expect("GroupNorm: spatial dims must be concrete"),
-        );
-    }
-    reshaped.shape = ShapeTracker::new(new_shape.clone());
-
-    // Normalize over axes [2, 3, ..., ndim] (C/G + spatial dims)
-    let norm_axes: Vec<usize> = (2..new_shape.len()).collect();
-    let mut normed = reshaped.layer_norm(norm_axes, epsilon);
-
-    // Reshape back to [N, C, spatial...]
-    let mut orig_shape = vec![n, c];
-    for d in &spatial_dims {
-        orig_shape.push(d.to_usize().unwrap());
-    }
-    normed *= 1.0;
-    normed.shape = ShapeTracker::new(orig_shape.clone());
-
-    // Apply scale and bias (both shape [C], broadcast to [N, C, spatial...])
-    let target_shape: Vec<Expression> = orig_shape.iter().map(|&d| Expression::from(d)).collect();
-    let result =
-        normed * broadcast_to_expr(scale, &target_shape) + broadcast_to_expr(bias, &target_shape);
-
-    tensors.insert(node.output[0].clone(), result);
-    trace!("Finished parse: GroupNormalization Node");
-    Ok(())
-}

crates/luminal_python/rust/src/pt2_compiled_model.rs

@@ -1,19 +1,18 @@
-use luminal::graph::Graph as LuminalGraph;
+use luminal::dyn_backend::BackendFactory;
+use luminal::prelude::tracing::warn;
 use luminal::prelude::*;
 use pyo3::prelude::*;
+use pyo3::types::{PyCapsule, PyCapsuleMethods};
 use std::collections::HashMap;
 
-#[cfg(feature = "cuda")]
-use luminal_cuda_lite::cudarc::driver::CudaContext;
-#[cfg(feature = "cuda")]
-use luminal_cuda_lite::runtime::CudaRuntime;
-
-use crate::compiled_graph::CompiledGraph;
-use crate::pt2_parser;
+use crate::compiled_graph::{CompiledGraph, DimParamMap, GraphTranslation, WeightData};
 use crate::pt2_schema;
-use crate::runtime::RuntimeBackend;
 use crate::translator;
-use crate::util::DimParamMap;
+use crate::typed_data::TypedData;
+use crate::{pt2_parser, pt2_util};
+
+/// Pre-loaded weight/constant data paired with tensor sizes.
+type PreloadResult = (Vec<(String, TypedData)>, HashMap<String, usize>);
 
 fn resolve_dim_sizes(
     sizes: &[pt2_schema::DimSize],
@@ -39,32 +38,89 @@ fn resolve_dim_sizes(
 }
 
 #[pyfunction]
-pub fn compile_pt2(
+#[pyo3(signature = (pt2_path, weights_path, search_iters, factory_capsule, weight_device_ptrs=None))]
+pub fn process_pt2(
     pt2_path: &str,
     weights_path: &str,
-    backend: &str,
     search_iters: usize,
+    factory_capsule: &Bound<'_, PyCapsule>,
+    weight_device_ptrs: Option<HashMap<String, (u64, usize)>>,
 ) -> PyResult<CompiledGraph> {
-    compile_pt2_inner(pt2_path, weights_path, backend, search_iters)
-        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e:#}")))
+    let factory: BackendFactory = {
+        let expected = ::luminal::dyn_backend::BACKEND_FACTORY_CAPSULE_NAME;
+        match factory_capsule.name()? {
+            Some(name) => {
+                // SAFETY: the &CStr is used immediately (for a byte-wise
+                // comparison) and never stored; the capsule is borrowed for
+                // the duration of this function, so the name pointer stays
+                // valid for as long as we read it here.
+                let actual = unsafe { name.as_cstr() };
+                if actual != expected {
+                    return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                        "factory_capsule has wrong name: expected {:?}, got {:?}",
+                        expected, actual,
+                    )));
+                }
+            }
+            None => {
+                return Err(pyo3::exceptions::PyValueError::new_err(
+                    "factory_capsule has no name; expected \"luminal.backend_factory\"",
+                ));
+            }
+        }
+        let wrapper_ptr = factory_capsule
+            .pointer_checked(Some(expected))
+            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?
+            .as_ptr() as *const *const std::ffi::c_void;
+        let fn_ptr = unsafe { *wrapper_ptr };
+        if fn_ptr.is_null() {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "factory_capsule inner function pointer is null",
+            ));
+        }
+        unsafe { std::mem::transmute(fn_ptr) }
+    };
+    compile_pt2(
+        pt2_path,
+        weights_path,
+        search_iters,
+        weight_device_ptrs.unwrap_or_default(),
+        factory,
+    )
+    .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e:#}")))
 }
 
-fn compile_pt2_inner(
+fn compile_pt2(
     pt2_path: &str,
     weights_path: &str,
-    backend: &str,
     search_iters: usize,
+    weight_device_ptrs: HashMap<String, (u64, usize)>,
+    factory: BackendFactory,
 ) -> anyhow::Result<CompiledGraph> {
+    let (translation, mut weights) = translate_pt2(pt2_path, weights_path)?;
+    weights.device_ptrs = weight_device_ptrs;
+
+    CompiledGraph::parse_graph(translation, weights, factory, search_iters)
+        .map_err(|e| anyhow::anyhow!(e))
+}
+
+/// Translate a PT2 exported model into a format-neutral GraphTranslation + WeightData.
+pub fn translate_pt2(
+    pt2_path: &str,
+    weights_path: &str,
+) -> anyhow::Result<(GraphTranslation, WeightData)> {
     let parsed = pt2_parser::parse_pt2(pt2_path)?;
     let translated = translator::translate(&parsed)?;
     let mut graph = translated.graph;
 
+    // Set initial dynamic dim values from symbol ranges
     for (sym_name, c) in &translated.sym_map.sym_to_char {
         if let Some(rc) = translated.sym_map.ranges.get(sym_name) {
             graph.set_dim(*c, rc.min_val as usize);
         }
     }
 
+    // Compute shape expressions and dtypes from PT2 tensor metadata
     let output_shape_exprs: Vec<Vec<Expression>> = translated
         .output_ids
         .iter()
@@ -76,6 +132,17 @@ fn compile_pt2_inner(
         })
         .collect();
 
+    let output_dtypes: Vec<DType> = translated
+        .output_ids
+        .iter()
+        .map(|(name, _id)| {
+            parsed
+                .tensor_meta(name)
+                .map(|meta| pt2_util::torch_dtype_int_to_luminal(meta.dtype))
+                .unwrap_or(DType::F32)
+        })
+        .collect();
+
     let input_names: Vec<String> = translated
         .user_input_ids
         .iter()
@@ -98,45 +165,6 @@ fn compile_pt2_inner(
         })
         .collect();
 
-    let user_input_sizes: Vec<(NodeIndex, usize)> = translated
-        .user_input_ids
-        .iter()
-        .map(|(name, id)| {
-            let meta = parsed.tensor_meta(name);
-            let n_elements = meta
-                .map(|m| {
-                    m.sizes
-                        .iter()
-                        .map(|s| s.hint().unwrap_or(1) as usize)
-                        .product()
-                })
-                .unwrap_or(1);
-            (*id, n_elements)
-        })
-        .collect();
-
-    let runtime = match backend {
-        "cpu" | "native" => {
-            graph.build_search_space::<NativeRuntime>();
-            let mut rt = graph.search(NativeRuntime::default(), search_iters);
-            if !weights_path.is_empty() {
-                load_safetensors_native(&mut rt, &graph, weights_path)?;
-            }
-            load_constants_native(&mut rt, &graph, &parsed)?;
-            RuntimeBackend::Native(rt)
-        }
-        "cuda" | "gpu" => init_cuda_runtime(
-            &mut graph,
-            weights_path,
-            &parsed,
-            &user_input_sizes,
-            search_iters,
-        )?,
-        other => {
-            anyhow::bail!("Unknown backend: {other}. Use 'cpu' or 'cuda'.");
-        }
-    };
-
     // Build tensor_ids from user inputs and outputs
     let mut tensor_ids: HashMap<String, NodeIndex> = HashMap::new();
     for (name, id) in &translated.user_input_ids {
@@ -146,80 +174,91 @@ fn compile_pt2_inner(
         tensor_ids.insert(name.clone(), *id);
     }
 
-    // Resolve concrete output shapes
-    let output_shapes: Vec<Vec<usize>> = output_shape_exprs
-        .iter()
-        .map(|exprs| exprs.iter().map(|e| e.to_usize().unwrap_or(1)).collect())
-        .collect();
+    // Pre-load weights and compute tensor sizes for CUDA dummy data
+    let mut weights: Vec<(String, TypedData)> = Vec::new();
+    let mut tensor_sizes: HashMap<String, usize> = HashMap::new();
+
+    // Load safetensors weights
+    if !weights_path.is_empty() {
+        let (st_weights, st_sizes) = preload_safetensors(&graph, weights_path)?;
+        weights.extend(st_weights);
+        tensor_sizes.extend(st_sizes);
+    }
+
+    // Load PT2 constants from ZIP archive
+    let (const_weights, const_sizes) = preload_constants(&graph, &parsed)?;
+    weights.extend(const_weights);
+    tensor_sizes.extend(const_sizes);
+
+    // Add tensor sizes from PT2 metadata for parameters/buffers not in safetensors
+    // (covers case when weights are loaded via device pointers after compilation)
+    for input_kind in parsed.classify_inputs() {
+        let (graph_name, original_name) = match &input_kind {
+            pt2_parser::InputKind::Parameter {
+                graph_name,
+                original_name,
+            } => (graph_name.as_str(), original_name.as_str()),
+            pt2_parser::InputKind::Buffer {
+                graph_name,
+                original_name,
+            } => (graph_name.as_str(), original_name.as_str()),
+            pt2_parser::InputKind::UserInput { .. } => continue,
+        };
+        // Always use authoritative sizes from model.json tensor_meta,
+        // even if preload_constants inserted a different (possibly stripped) size.
+        if let Some(meta) = parsed.tensor_meta(graph_name) {
+            let n: usize = meta
+                .sizes
+                .iter()
+                .map(|s| s.hint().unwrap_or(1) as usize)
+                .product();
+            tensor_sizes.insert(original_name.to_string(), n);
+        }
+    }
+
+    // Add user input sizes
+    for (name, _id) in &translated.user_input_ids {
+        if !tensor_sizes.contains_key(name)
+            && let Some(meta) = parsed.tensor_meta(name)
+        {
+            let n: usize = meta
+                .sizes
+                .iter()
+                .map(|s| s.hint().unwrap_or(1) as usize)
+                .product();
+            tensor_sizes.insert(name.clone(), n);
+        }
+    }
 
-    // Build dim_param_map from sym_map
     let dim_param_map: DimParamMap = translated.sym_map.sym_to_char;
 
-    Ok(CompiledGraph {
+    let translation = GraphTranslation {
         graph,
-        runtime,
         tensor_ids,
         input_names,
         output_names,
-        output_shapes,
+        output_dtypes,
         output_shape_exprs,
         input_shape_exprs,
         dim_param_map,
-    })
-}
+    };
 
-#[cfg(feature = "cuda")]
-fn init_cuda_runtime(
-    graph: &mut LuminalGraph,
-    weights_path: &str,
-    parsed: &pt2_parser::ParsedPT2,
-    user_input_sizes: &[(NodeIndex, usize)],
-    search_iters: usize,
-) -> anyhow::Result<RuntimeBackend> {
-    let cuda_ctx =
-        CudaContext::new(0).map_err(|e| anyhow::anyhow!("CUDA context init failed: {e}"))?;
-    let stream = cuda_ctx.default_stream();
+    let weight_data = WeightData {
+        weights,
+        tensor_sizes,
+        device_ptrs: HashMap::new(),
+    };
 
-    graph.build_search_space::<CudaRuntime>();
-    let mut rt = CudaRuntime::initialize(stream);
-
-    // Phase 1: Set ALL input nodes to safe dummy data (1.0) for search profiling.
-    // Real weights/constants may contain -inf (e.g. causal attention mask) which
-    // produce NaN in intermediate computations (e.g. -inf - (-inf) = NaN in softmax
-    // decomposition), causing the search's has_nan_outputs check to reject ALL
-    // candidates. We load real data only AFTER the search completes.
-    set_all_inputs_dummy_cuda(&mut rt, graph, weights_path, parsed, user_input_sizes)?;
-
-    let mut rt = graph.search(rt, search_iters);
-
-    if !weights_path.is_empty() {
-        load_safetensors_cuda(&mut rt, graph, weights_path)?;
-    }
-    load_constants_cuda(&mut rt, graph, parsed)?;
-
-    Ok(RuntimeBackend::Cuda(Box::new(rt)))
-}
-
-#[cfg(not(feature = "cuda"))]
-fn init_cuda_runtime(
-    _graph: &mut LuminalGraph,
-    _weights_path: &str,
-    _parsed: &pt2_parser::ParsedPT2,
-    _user_input_sizes: &[(NodeIndex, usize)],
-    _search_iters: usize,
-) -> anyhow::Result<RuntimeBackend> {
-    anyhow::bail!("CUDA support not compiled. Rebuild with --features cuda")
+    Ok((translation, weight_data))
 }
 
 // ---------------------------------------------------------------------------
-// Weight loading
+// Weight pre-loading helpers
 // ---------------------------------------------------------------------------
 
-fn load_safetensors_impl(
-    cx: &LuminalGraph,
-    file_path: &str,
-    mut set_data: impl FnMut(NodeIndex, Vec<f32>),
-) -> anyhow::Result<()> {
+/// Pre-load all safetensors weights that match Input nodes in the graph.
+/// Returns (weight data, tensor sizes for all tensors in the file).
+fn preload_safetensors(graph: &Graph, file_path: &str) -> anyhow::Result<PreloadResult> {
     use memmap2::MmapOptions;
     use safetensors::SafeTensors;
     use std::fs::File;
@@ -229,95 +268,75 @@ fn load_safetensors_impl(
     let st = SafeTensors::deserialize(&mmap)
         .map_err(|e| anyhow::anyhow!("SafeTensors deserialize error: {e}"))?;
 
-    for node in cx.graph.node_indices() {
-        if let Some(input) = (*cx.graph[node])
+    let mut weights = Vec::new();
+    let mut sizes = HashMap::new();
+
+    // Get sizes for ALL tensors in the file (for dummy data allocation)
+    for (name, info) in st.tensors() {
+        let n: usize = info.shape().iter().product();
+        sizes.insert(name.to_string(), n);
+    }
+
+    // Load weight data for Input nodes that match safetensors tensor names
+    for node_id in graph.graph.node_indices() {
+        if let Some(input) = (*graph.graph[node_id])
             .as_any()
             .downcast_ref::<luminal::hlir::Input>()
             && let Ok(tensor) = st.tensor(&input.label)
         {
-            let f32s = bytes_to_f32(tensor.data(), safetensors_dtype_to_pt2(tensor.dtype()));
-            set_data(node, f32s);
+            let types = bytes_to_typed(tensor.data(), safetensors_dtype_to_pt2(tensor.dtype()));
+            weights.push((input.label.clone(), types));
         }
     }
 
-    Ok(())
+    Ok((weights, sizes))
 }
 
-fn load_safetensors_native(
-    rt: &mut NativeRuntime,
-    cx: &LuminalGraph,
-    file_path: &str,
-) -> anyhow::Result<()> {
-    load_safetensors_impl(cx, file_path, |node, data| rt.set_data(node, data))
-}
-
-#[cfg(feature = "cuda")]
-fn load_safetensors_cuda(
-    rt: &mut CudaRuntime,
-    cx: &LuminalGraph,
-    file_path: &str,
-) -> anyhow::Result<()> {
-    load_safetensors_impl(cx, file_path, |node, data| rt.set_data(node, data))
-}
-
-/// Set ALL input nodes to dummy 1.0 data for safe CUDA search profiling.
-#[cfg(feature = "cuda")]
-fn set_all_inputs_dummy_cuda(
-    rt: &mut CudaRuntime,
-    cx: &LuminalGraph,
-    weights_path: &str,
+/// Pre-load all PT2 constants from the ZIP archive.
+/// Returns (constant data, tensor sizes for all constants).
+fn preload_constants(
+    _graph: &Graph,
     parsed: &pt2_parser::ParsedPT2,
-    user_input_sizes: &[(NodeIndex, usize)],
-) -> anyhow::Result<()> {
-    use memmap2::MmapOptions;
-    use safetensors::SafeTensors;
-    use std::fs::File;
+) -> anyhow::Result<PreloadResult> {
+    let constants_config = match &parsed.constants_config {
+        Some(c) => c,
+        None => return Ok((Vec::new(), HashMap::new())),
+    };
 
-    let mut label_sizes: HashMap<String, usize> = HashMap::new();
+    let mut weights = Vec::new();
+    let mut sizes = HashMap::new();
 
-    if !weights_path.is_empty() {
-        let f = File::open(weights_path)?;
-        let mmap = unsafe { MmapOptions::new().map(&f)? };
-        let st = SafeTensors::deserialize(&mmap)
-            .map_err(|e| anyhow::anyhow!("SafeTensors deserialize error: {e}"))?;
-        for (name, info) in st.tensors() {
-            let n: usize = info.shape().iter().product();
-            label_sizes.insert(name.to_string(), n);
-        }
-    }
+    for (name, entry) in &constants_config.config {
+        let n: usize = entry
+            .tensor_meta
+            .sizes
+            .iter()
+            .map(|s| s.hint().unwrap_or(1) as usize)
+            .product();
+        sizes.insert(name.clone(), n);
 
-    if let Some(cc) = &parsed.constants_config {
-        for (name, entry) in &cc.config {
-            let n: usize = entry
-                .tensor_meta
-                .sizes
-                .iter()
-                .map(|s| s.hint().unwrap_or(1) as usize)
-                .product();
-            label_sizes.insert(name.clone(), n);
-        }
-    }
-
-    for node_id in cx.graph.node_indices() {
-        if let Some(input) = (*cx.graph[node_id])
-            .as_any()
-            .downcast_ref::<luminal::hlir::Input>()
-        {
-            if let Some(&n) = label_sizes.get(&input.label) {
-                if n > 0 {
-                    rt.set_data(node_id, vec![1.0f32; n]);
-                }
+        let raw_bytes = match pt2_parser::read_constant_bytes(
+            &parsed.pt2_path,
+            &parsed.archive_prefix,
+            entry,
+        ) {
+            Ok(b) => b,
+            Err(e) => {
+                warn!("failed to load constant '{}': {:#}", name, e);
+                continue;
             }
-        }
+        };
+        let typed_data = bytes_to_typed(&raw_bytes, entry.tensor_meta.dtype);
+        weights.push((name.clone(), typed_data));
     }
 
-    for &(id, n_elements) in user_input_sizes {
-        rt.set_data(id, vec![1.0f32; n_elements]);
-    }
-
-    Ok(())
+    Ok((weights, sizes))
 }
 
+// ---------------------------------------------------------------------------
+// Byte conversion helpers
+// ---------------------------------------------------------------------------
+
 /// Convert safetensors Dtype to PT2 dtype number.
 fn safetensors_dtype_to_pt2(dtype: safetensors::Dtype) -> u32 {
     match dtype {
@@ -335,106 +354,52 @@ fn safetensors_dtype_to_pt2(dtype: safetensors::Dtype) -> u32 {
     }
 }
 
-/// Convert raw bytes to f32 using PT2 dtype numbering.
-fn bytes_to_f32(bytes: &[u8], dtype: u32) -> Vec<f32> {
+/// Convert raw bytes to TypedData using PT2 dtype numbering.
+/// Preserves native byte format for types luminal supports directly (f32, f16, bf16, i32, bool, u8, i8).
+/// Converts i64/f64/i16 to the closest luminal-native representation.
+fn bytes_to_typed(bytes: &[u8], dtype: u32) -> TypedData {
     match dtype {
-        7 => bytes
-            .chunks_exact(4)
-            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
-            .collect(),
-        6 => bytes
-            .chunks_exact(2)
-            .map(|b| half::f16::from_le_bytes([b[0], b[1]]).to_f32())
-            .collect(),
-        13 => bytes
-            .chunks_exact(2)
-            .map(|b| half::bf16::from_le_bytes([b[0], b[1]]).to_f32())
-            .collect(),
-        8 => bytes
-            .chunks_exact(8)
-            .map(|b| f64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as f32)
-            .collect(),
-        5 => bytes
-            .chunks_exact(8)
-            .map(|b| i64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as f32)
-            .collect(),
-        4 => bytes
-            .chunks_exact(4)
-            .map(|b| i32::from_le_bytes([b[0], b[1], b[2], b[3]]) as f32)
-            .collect(),
-        3 => bytes
-            .chunks_exact(2)
-            .map(|b| i16::from_le_bytes([b[0], b[1]]) as f32)
-            .collect(),
-        2 => bytes.iter().map(|&b| (b as i8) as f32).collect(),
-        1 => bytes.iter().map(|&b| b as f32).collect(),
-        12 => bytes
-            .iter()
-            .map(|&b| if b != 0 { 1.0 } else { 0.0 })
-            .collect(),
+        // Types that map directly — preserve raw bytes
+        7 => TypedData::from_raw(bytes.to_vec(), DType::F32),
+        6 => TypedData::from_raw(bytes.to_vec(), DType::F16),
+        13 => TypedData::from_raw(bytes.to_vec(), DType::Bf16),
+        4 => TypedData::from_raw(bytes.to_vec(), DType::Int), // i32
+        1 => TypedData::from_raw(bytes.to_vec(), DType::U8),
+        2 => TypedData::from_raw(bytes.to_vec(), DType::I8),
+        12 => TypedData::from_raw(bytes.to_vec(), DType::Bool),
+
+        // i64 → i32 (truncate, matching luminal's Int type)
+        5 => {
+            let i32s: Vec<i32> = bytes
+                .chunks_exact(8)
+                .map(|b| {
+                    i64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as i32
+                })
+                .collect();
+            TypedData::from_i32_vec(i32s)
+        }
+        // f64 → f32 (downcast, luminal has no F64 in practice for most ops)
+        8 => {
+            let f32s: Vec<f32> = bytes
+                .chunks_exact(8)
+                .map(|b| {
+                    f64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as f32
+                })
+                .collect();
+            TypedData::from_f32_vec(f32s)
+        }
+        // i16 → i32 (widen to luminal's Int)
+        3 => {
+            let i32s: Vec<i32> = bytes
+                .chunks_exact(2)
+                .map(|b| i16::from_le_bytes([b[0], b[1]]) as i32)
+                .collect();
+            TypedData::from_i32_vec(i32s)
+        }
         _ => {
-            eprintln!("[luminal] Warning: unrecognized dtype {dtype}, interpreting as f32");
-            bytes
-                .chunks_exact(4)
-                .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
-                .collect()
+            let luminal_dtype = pt2_util::torch_dtype_int_to_luminal(dtype);
+            warn!("Unrecognized dtype {dtype}, interpreting as {luminal_dtype:?}");
+            TypedData::from_raw(bytes.to_vec(), luminal_dtype)
         }
     }
 }
-
-fn load_constants_impl(
-    cx: &LuminalGraph,
-    parsed: &pt2_parser::ParsedPT2,
-    mut set_data: impl FnMut(NodeIndex, Vec<f32>),
-) -> anyhow::Result<()> {
-    let constants_config = match &parsed.constants_config {
-        Some(c) => c,
-        None => return Ok(()),
-    };
-
-    for (name, entry) in &constants_config.config {
-        let raw_bytes = match pt2_parser::read_constant_bytes(
-            &parsed.pt2_path,
-            &parsed.archive_prefix,
-            entry,
-        ) {
-            Ok(b) => b,
-            Err(e) => {
-                eprintln!(
-                    "[luminal] Warning: failed to load constant '{}': {:#}",
-                    name, e
-                );
-                continue;
-            }
-        };
-        let f32_data = bytes_to_f32(&raw_bytes, entry.tensor_meta.dtype);
-
-        for node_id in cx.graph.node_indices() {
-            if let Some(input) = (*cx.graph[node_id])
-                .as_any()
-                .downcast_ref::<luminal::hlir::Input>()
-                && input.label == *name
-            {
-                set_data(node_id, f32_data.clone());
-            }
-        }
-    }
-    Ok(())
-}
-
-fn load_constants_native(
-    rt: &mut NativeRuntime,
-    cx: &LuminalGraph,
-    parsed: &pt2_parser::ParsedPT2,
-) -> anyhow::Result<()> {
-    load_constants_impl(cx, parsed, |node, data| rt.set_data(node, data))
-}
-
-#[cfg(feature = "cuda")]
-fn load_constants_cuda(
-    rt: &mut CudaRuntime,
-    cx: &LuminalGraph,
-    parsed: &pt2_parser::ParsedPT2,
-) -> anyhow::Result<()> {
-    load_constants_impl(cx, parsed, |node, data| rt.set_data(node, data))
-}

crates/luminal_python/rust/src/pt2_parser.rs

@@ -160,7 +160,31 @@ pub fn parse_pt2(path: &str) -> Result<ParsedPT2> {
     let file = File::open(path).with_context(|| format!("Failed to open PT2 file: {path}"))?;
     let mut archive = ZipArchive::new(file).context("Failed to read PT2 ZIP archive")?;
 
-    // Determine archive prefix from the first entry
+    // Torch >= 2.6 uses a flat archive with no prefix directory; detect by presence of the
+    // well-known root-level file.  Older torch used a prefix (e.g. "archive/models/model.json").
+    let is_new_format = archive
+        .file_names()
+        .any(|n| n == "serialized_exported_program.json");
+
+    if is_new_format {
+        let program: ExportedProgram = {
+            let mut entry = archive.by_name("serialized_exported_program.json")?;
+            let mut buf = String::new();
+            entry.read_to_string(&mut buf)?;
+            serde_json::from_str(&buf)
+                .context("Failed to parse serialized_exported_program.json")?
+        };
+        // Tensor constants live in serialized_constants.pt; Python extracts them
+        // and loads them post-compile via set_weight_from_ptr.
+        return Ok(ParsedPT2 {
+            program,
+            constants_config: None,
+            archive_prefix: String::new(),
+            pt2_path: path.to_string(),
+        });
+    }
+
+    // Old prefix-based format.
     let archive_prefix = {
         let first = archive
             .file_names()

crates/luminal_python/rust/src/pt2_schema.rs

@@ -77,6 +77,7 @@ pub enum Argument {
     SymInts(SymIntsArg),
     SymInt(SymIntArg),
     Expr(ExprArg),
+    #[allow(dead_code)]
     ScalarType(ScalarTypeArg),
     Tensors(TensorsArg),
     OptionalTensors(OptionalTensorsArg),
@@ -168,6 +169,7 @@ pub struct NoneArg {
 }
 
 #[derive(Debug, Clone, Deserialize)]
+#[allow(dead_code)]
 pub struct ScalarTypeArg {
     pub as_scalar_type: u32,
 }
@@ -224,6 +226,7 @@ impl Argument {
         }
     }
 
+    #[allow(dead_code)]
     pub fn as_scalar_type(&self) -> Option<u32> {
         match self {
             Argument::ScalarType(s) => Some(s.as_scalar_type),

crates/luminal_python/rust/src/pt2_util.rs

@@ -16,6 +16,7 @@ pub enum ReductionOp {
     Mean,
     Max,
     Min,
+    Prod,
 }
 
 /// Normalize a potentially negative dimension index.

crates/luminal_python/rust/src/runtime.rs

@@ -1,89 +0,0 @@
-use luminal::prelude::*;
-#[cfg(feature = "cuda")]
-use luminal_cuda_lite::cudarc::driver::{CudaContext, CudaStream};
-#[cfg(feature = "cuda")]
-use luminal_cuda_lite::runtime::CudaRuntime;
-use rustc_hash::FxHashMap;
-#[cfg(feature = "cuda")]
-use std::sync::Arc;
-
-/// Enum wrapper for runtime backends allowing runtime selection.
-pub enum RuntimeBackend {
-    Native(NativeRuntime),
-    #[cfg(feature = "cuda")]
-    Cuda(Box<CudaRuntime>),
-}
-
-impl RuntimeBackend {
-    /// Set input data for a tensor node.
-    pub fn set_data(&mut self, node: NodeIndex, data: Vec<f32>) {
-        match self {
-            RuntimeBackend::Native(rt) => rt.set_data(node, data),
-            #[cfg(feature = "cuda")]
-            RuntimeBackend::Cuda(rt) => rt.set_data(node, data),
-        }
-    }
-
-    /// Execute the compiled graph.
-    pub fn execute(&mut self, dyn_map: &FxHashMap<char, usize>) {
-        match self {
-            RuntimeBackend::Native(rt) => rt.execute(dyn_map),
-            #[cfg(feature = "cuda")]
-            RuntimeBackend::Cuda(rt) => rt.execute(dyn_map),
-        }
-    }
-
-    /// Get output data from a tensor node.
-    pub fn get_f32(&self, node: NodeIndex) -> Vec<f32> {
-        match self {
-            RuntimeBackend::Native(rt) => rt.get_f32(node).to_vec(),
-            #[cfg(feature = "cuda")]
-            RuntimeBackend::Cuda(rt) => rt.get_f32(node),
-        }
-    }
-
-    /// Get the name of the active backend.
-    pub fn name(&self) -> &'static str {
-        match self {
-            RuntimeBackend::Native(_) => "native",
-            #[cfg(feature = "cuda")]
-            RuntimeBackend::Cuda(_) => "cuda",
-        }
-    }
-}
-
-// ============================================================================
-// Two-phase initialization for CUDA (required because profiling executes graph)
-// ============================================================================
-
-/// Prepare CUDA runtime: build search space and create runtime, but don't search yet.
-/// Returns the unoptimized runtime that can have data set on it.
-///
-/// Use this with `finalize_cuda` for proper CUDA initialization:
-/// 1. Call `prepare_cuda` to get the runtime
-/// 2. Set data on the runtime using `rt.set_data(node_id, data)`
-/// 3. Call `finalize_cuda` to run profiling with data available
-#[cfg(feature = "cuda")]
-pub fn prepare_cuda(context: &mut Graph) -> Result<(CudaRuntime, Arc<CudaStream>), String> {
-    let cuda_ctx =
-        CudaContext::new(0).map_err(|e| format!("Failed to init CUDA context: {}", e))?;
-    let stream = cuda_ctx.default_stream();
-    context.build_search_space::<CudaRuntime>();
-    let rt = CudaRuntime::initialize(stream.clone());
-    Ok((rt, stream))
-}
-
-/// Finalize CUDA runtime: run search with data already set.
-#[cfg(feature = "cuda")]
-pub fn finalize_cuda(context: &mut Graph, rt: CudaRuntime) -> RuntimeBackend {
-    let optimized_rt = context.search(rt, 10);
-    RuntimeBackend::Cuda(Box::new(optimized_rt))
-}
-
-/// Initialize a native (CPU) runtime using single-phase approach.
-/// NativeRuntime validates Input nodes, so we must search first, then set data.
-pub fn initialize_native(context: &mut Graph) -> Result<RuntimeBackend, String> {
-    context.build_search_space::<NativeRuntime>();
-    let rt = context.search(NativeRuntime::default(), 10);
-    Ok(RuntimeBackend::Native(rt))
-}

crates/luminal_python/rust/src/translator/binary.rs

@@ -12,6 +12,7 @@ impl<'a> Translator<'a> {
         let arg1 = &node.inputs[1].arg;
         if let Some(name) = arg1.as_tensor_name() {
             let b = self.get_tensor(name)?;
+            let (a, b) = ensure_same_dtype(a, b);
             let (a, b) = broadcast_binary(a, b);
             Ok(match op {
                 BinaryOp::Add => a + b,

crates/luminal_python/rust/src/translator/conv.rs

@@ -0,0 +1,407 @@
+use anyhow::Result;
+use luminal::prelude::*;
+
+use crate::pt2_schema::*;
+
+use super::Translator;
+
+const CONV_INPUT_ARG: usize = 0;
+const CONV_WEIGHT_ARG: usize = 1;
+const CONV_BIAS_ARG: usize = 2;
+const CONV_STRIDE_ARG: usize = 3;
+const CONV_PADDING_ARG: usize = 4;
+const CONV_DILATION_ARG: usize = 5;
+const CONV_GROUPS_ARG: usize = 6;
+
+const CONVOLUTION_TRANSPOSED_ARG: usize = 6;
+const CONVOLUTION_OUTPUT_PADDING_ARG: usize = 7;
+const CONVOLUTION_GROUPS_ARG: usize = 8;
+
+impl<'a> Translator<'a> {
+    /// Translate aten.conv{1,2,3}d.default and aten.convolution.default.
+    ///
+    /// The PT2 export may omit defaulted trailing arguments entirely. In practice this means
+    /// conv{N}d.default can show up as just `(input, weight)` for the no-bias, stride=1,
+    /// padding=0, dilation=1, groups=1 case.
+    pub(crate) fn translate_conv(&mut self, node: &Node) -> Result<GraphTensor> {
+        let input = self.get_input_tensor(node, CONV_INPUT_ARG)?;
+        let weight = self.get_input_tensor(node, CONV_WEIGHT_ARG)?;
+        let bias = self.get_input_tensor(node, CONV_BIAS_ARG).ok();
+
+        let x_dims = input.dims();
+        let w_dims = weight.dims();
+        let rank = x_dims.len();
+        let spatial = rank - 2;
+
+        let stride = self
+            .get_ints_arg(node, CONV_STRIDE_ARG)
+            .unwrap_or_else(|_| vec![1; spatial]);
+        let padding = self
+            .get_ints_arg(node, CONV_PADDING_ARG)
+            .unwrap_or_else(|_| vec![0; spatial]);
+        let mut dilation = self
+            .get_ints_arg(node, CONV_DILATION_ARG)
+            .unwrap_or_else(|_| vec![1; spatial]);
+        let groups = if node.target == "torch.ops.aten.convolution.default" {
+            let transposed = self
+                .get_bool_arg(node, CONVOLUTION_TRANSPOSED_ARG)
+                .unwrap_or(false);
+            anyhow::ensure!(
+                !transposed,
+                "conv: ConvTranspose / transposed=true is not supported yet"
+            );
+            let output_padding = self
+                .get_ints_arg(node, CONVOLUTION_OUTPUT_PADDING_ARG)
+                .unwrap_or_else(|_| vec![0; spatial]);
+            anyhow::ensure!(
+                output_padding.iter().all(|&v| v == 0),
+                "conv: output_padding is not supported for non-transposed convolution"
+            );
+            self.get_int_arg(node, CONVOLUTION_GROUPS_ARG).unwrap_or(1) as usize
+        } else {
+            self.get_int_arg(node, CONV_GROUPS_ARG).unwrap_or(1) as usize
+        };
+        if dilation.len() != spatial {
+            dilation = vec![1; spatial];
+        }
+
+        let ch_out = w_dims[0]
+            .to_usize()
+            .ok_or_else(|| anyhow::anyhow!("conv: weight C_out must be concrete"))?;
+        let ch_in = x_dims[1]
+            .to_usize()
+            .ok_or_else(|| anyhow::anyhow!("conv: input C_in must be concrete"))?;
+        anyhow::ensure!(
+            stride.len() == spatial && padding.len() == spatial && dilation.len() == spatial,
+            "conv: stride/padding/dilation rank must match spatial rank {spatial}"
+        );
+        anyhow::ensure!(
+            groups > 0 && ch_in % groups == 0 && ch_out % groups == 0,
+            "conv: invalid group configuration (C_in={ch_in}, C_out={ch_out}, groups={groups})"
+        );
+        let ch_per_group = ch_in / groups;
+
+        let kernel_shape: Vec<usize> = w_dims[2..]
+            .iter()
+            .map(|d| {
+                d.to_usize()
+                    .ok_or_else(|| anyhow::anyhow!("conv: kernel dims must be concrete"))
+            })
+            .collect::<Result<_>>()?;
+        let kernel_product: usize = kernel_shape.iter().product();
+
+        // ATen uses symmetric padding (same begin/end)
+        let stride_u: Vec<usize> = stride.iter().map(|&v| v as usize).collect();
+        let padding_u: Vec<usize> = padding.iter().map(|&v| v as usize).collect();
+        let dilation_u: Vec<usize> = dilation.iter().map(|&v| v as usize).collect();
+
+        let mut out = if groups > 1 {
+            let group_out = ch_out / groups;
+
+            if ch_per_group == 1 {
+                // Depthwise (including channel multiplier > 1): avoid per-channel slicing.
+                depthwise_conv(
+                    input,
+                    weight,
+                    &kernel_shape,
+                    &stride_u,
+                    &dilation_u,
+                    &padding_u,
+                    &padding_u,
+                    ch_in,
+                    group_out,
+                    kernel_product,
+                    spatial,
+                )
+            } else {
+                // General grouped: pre-pad full input then slice per group
+                let padded_input = {
+                    let mut pad_spec: Vec<(Expression, Expression)> =
+                        vec![(0.into(), 0.into()); 2 + spatial];
+                    for i in 0..spatial {
+                        pad_spec[2 + i] = (padding_u[i].into(), padding_u[i].into());
+                    }
+                    input.pad(pad_spec, 0.0)
+                };
+
+                let no_pad = vec![0usize; spatial];
+                let mut group_outputs = Vec::with_capacity(groups);
+                for g in 0..groups {
+                    let x_g = slice_channel_group(padded_input, g, ch_per_group, spatial);
+                    let w_g =
+                        slice_weight_group(weight, g, group_out, ch_per_group * kernel_product);
+                    group_outputs.push(conv_unfold(
+                        x_g,
+                        w_g,
+                        &kernel_shape,
+                        &stride_u,
+                        &dilation_u,
+                        &no_pad,
+                        &no_pad,
+                        ch_per_group,
+                        group_out,
+                        spatial,
+                    ));
+                }
+
+                let mut result = group_outputs[0];
+                for g_out in &group_outputs[1..] {
+                    result = result.concat_along(*g_out, 1);
+                }
+                result
+            }
+        } else {
+            let mut w_flat = weight;
+            w_flat.shape = ShapeTracker::new_with_element_bits(
+                vec![ch_out, ch_in * kernel_product],
+                weight.dtype.bits(),
+            );
+
+            conv_unfold(
+                input,
+                w_flat,
+                &kernel_shape,
+                &stride_u,
+                &dilation_u,
+                &padding_u,
+                &padding_u,
+                ch_in,
+                ch_out,
+                spatial,
+            )
+        };
+
+        if let Some(b) = bias {
+            let out_dims = out.dims();
+            let mut b_expanded = b.expand_dim(0, 1);
+            for i in 0..spatial {
+                b_expanded = b_expanded.expand_dim(2 + i, out_dims[2 + i]);
+            }
+            out += b_expanded;
+        }
+
+        Ok(out)
+    }
+}
+
+/// Slice input channels for one group.
+/// Caller must pre-pad `x` so no additional padding is applied to the slice.
+fn slice_channel_group(
+    x: GraphTensor,
+    g: usize,
+    ch_per_group: usize,
+    spatial: usize,
+) -> GraphTensor {
+    let start = g * ch_per_group;
+    let end = start + ch_per_group;
+    let dims = x.dims();
+    let rank = 2 + spatial;
+    let mut slices: Vec<(Expression, Expression)> = Vec::with_capacity(rank);
+    slices.push((0.into(), dims[0]));
+    slices.push((start.into(), end.into()));
+    for dim in dims.iter().take(rank).skip(2) {
+        slices.push((0.into(), *dim));
+    }
+    x.slice(slices)
+}
+
+/// Slice and flatten weight for one group.
+fn slice_weight_group(
+    w: GraphTensor,
+    g: usize,
+    group_out: usize,
+    flat_inner: usize,
+) -> GraphTensor {
+    let start = g * group_out;
+    let end = start + group_out;
+    let w_dims = w.dims();
+    let mut slices: Vec<(Expression, Expression)> = Vec::with_capacity(w_dims.len());
+    slices.push((start.into(), end.into()));
+    for dim in w_dims.iter().skip(1) {
+        slices.push((0.into(), *dim));
+    }
+    // Materialize through Add: binary op outputs are contiguous in Luminal, which makes the
+    // following flatten safe for the sliced weight buffer.
+    let w_sliced = w.slice(slices) + 0.0;
+    let mut w_flat = w_sliced;
+    w_flat.shape =
+        ShapeTracker::new_with_element_bits(vec![group_out, flat_inner], w_sliced.dtype.bits());
+    w_flat
+}
+
+/// Core unfold-based convolution for a single group.
+///
+/// `x`: [batch, ch_in, spatial...]
+/// `w_flat`: [ch_out, ch_in * kernel_product] (already reshaped)
+/// Returns: [batch, ch_out, out_spatial...]
+#[allow(clippy::too_many_arguments)]
+fn conv_unfold(
+    x: GraphTensor,
+    w_flat: GraphTensor,
+    kernel_shape: &[usize],
+    strides: &[usize],
+    dilations: &[usize],
+    pads_begin: &[usize],
+    pads_end: &[usize],
+    _ch_in: usize,
+    _ch_out: usize,
+    spatial: usize,
+) -> GraphTensor {
+    let rank = 2 + spatial;
+
+    // Pad spatial dimensions (skip if all padding is zero)
+    let needs_pad = pads_begin.iter().any(|&p| p > 0) || pads_end.iter().any(|&p| p > 0);
+    let padded = if needs_pad {
+        let mut padding: Vec<(Expression, Expression)> = vec![(0.into(), 0.into()); rank];
+        for i in 0..spatial {
+            padding[2 + i] = (pads_begin[i].into(), pads_end[i].into());
+        }
+        x.pad(padding, 0.0)
+    } else {
+        x
+    };
+
+    // Build full-rank unfold parameters (1 for batch/channel, actual for spatial)
+    let mut kernel_full = vec![1usize; rank];
+    let mut stride_full = vec![1usize; rank];
+    let mut dilation_full = vec![1usize; rank];
+    kernel_full[2..(spatial + 2)].copy_from_slice(&kernel_shape[..spatial]);
+    stride_full[2..(spatial + 2)].copy_from_slice(&strides[..spatial]);
+    dilation_full[2..(spatial + 2)].copy_from_slice(&dilations[..spatial]);
+
+    let unfolded = padded.unfold(kernel_full, stride_full, dilation_full);
+    // Shape: [win_N, win_C, win_spatial..., k_N=1, k_C=1, k_spatial...]
+
+    // Permute to [N, win_spatial..., C_in, k_N, k_C, k_spatial...]
+    let mut perm: Vec<usize> = Vec::with_capacity(2 * rank);
+    perm.push(0);
+    perm.extend(2..2 + spatial);
+    perm.push(1);
+    perm.extend(rank..2 * rank);
+    let permuted = unfolded.permute(perm);
+
+    let output_spatial_dims: Vec<Expression> = permuted.dims()[1..1 + spatial].to_vec();
+
+    // Merge all channel+kernel dims into [N, spatial..., ch_in * kernel_product]
+    let mut patches = permuted;
+    let target = 2 + spatial;
+    while patches.dims().len() > target {
+        let last = patches.dims().len();
+        patches = patches.merge_dims(last - 2, last - 1);
+    }
+
+    // Merge spatial dims into one
+    for _ in 1..spatial {
+        patches = patches.merge_dims(1, 2);
+    }
+    // patches: [N, spatial_product, ch_in * kernel_product]
+
+    let mut out = patches.matmul(w_flat.permute((1, 0)));
+    // out: [N, spatial_product, ch_out]
+
+    // Restore spatial dimensions
+    for i in (1..spatial).rev() {
+        out = out.split_dims(1, output_spatial_dims[i]);
+    }
+
+    // Move ch_out from last to position 1: [N, ch_out, spatial...]
+    let mut final_order: Vec<usize> = Vec::with_capacity(2 + spatial);
+    final_order.push(0);
+    final_order.push(1 + spatial);
+    final_order.extend(1..1 + spatial);
+    out.permute(final_order)
+}
+
+/// Depthwise convolution: groups == in_channels, ch_per_group == 1.
+///
+/// Processes all channels simultaneously using element-wise multiply + reduce,
+/// avoiding per-channel input slicing which can cause index-expression bugs in luminal.
+///
+/// out[n, c, oh, ow] = sum_k patches[n, c, oh, ow, k] * weight[c, k]
+#[allow(clippy::too_many_arguments)]
+fn depthwise_conv(
+    x: GraphTensor,
+    w: GraphTensor, // [C, 1, *kernel]
+    kernel_shape: &[usize],
+    strides: &[usize],
+    dilations: &[usize],
+    pads_begin: &[usize],
+    pads_end: &[usize],
+    ch: usize,
+    group_out: usize,
+    kernel_product: usize,
+    spatial: usize,
+) -> GraphTensor {
+    let rank = 2 + spatial;
+
+    let needs_pad = pads_begin.iter().any(|&p| p > 0) || pads_end.iter().any(|&p| p > 0);
+    let padded = if needs_pad {
+        let mut padding: Vec<(Expression, Expression)> = vec![(0.into(), 0.into()); rank];
+        for i in 0..spatial {
+            padding[2 + i] = (pads_begin[i].into(), pads_end[i].into());
+        }
+        x.pad(padding, 0.0)
+    } else {
+        x
+    };
+
+    // Unfold the full [N, C, H+2p, W+2p] with kernel [1, 1, kH, kW]
+    let mut kernel_full = vec![1usize; rank];
+    let mut stride_full = vec![1usize; rank];
+    let mut dilation_full = vec![1usize; rank];
+    kernel_full[2..(spatial + 2)].copy_from_slice(&kernel_shape[..spatial]);
+    stride_full[2..(spatial + 2)].copy_from_slice(&strides[..spatial]);
+    dilation_full[2..(spatial + 2)].copy_from_slice(&dilations[..spatial]);
+
+    let unfolded = padded.unfold(kernel_full, stride_full, dilation_full);
+    // Shape: [N, C, out_H, out_W, 1, 1, kH, kW]
+
+    // Permute to [N, C, out_spatial..., k_all...]
+    let mut perm: Vec<usize> = Vec::with_capacity(2 * rank);
+    perm.push(0); // N
+    perm.push(1); // C
+    perm.extend(2..2 + spatial); // win_spatial
+    perm.extend(rank..2 * rank); // all kernel dims
+    let permuted = unfolded.permute(perm);
+
+    let out_spatial_dims: Vec<Expression> = permuted.dims()[2..2 + spatial].to_vec();
+
+    // Merge all kernel dims (including 1-size k_N, k_C) into kernel_product
+    let target = 3 + spatial; // [N, C, spatial..., K]
+    let mut patches = permuted;
+    while patches.dims().len() > target {
+        let last = patches.dims().len();
+        patches = patches.merge_dims(last - 2, last - 1);
+    }
+    // patches: [N, C, out_H, ..., out_W, kernel_product]
+
+    // Merge spatial into one: [N, C, out_spatial_product, kernel_product]
+    for _ in 1..spatial {
+        patches = patches.merge_dims(2, 3);
+    }
+
+    // Weight [C * group_out, 1, *kernel] -> [C, group_out, kernel_product]
+    let mut w_flat = w;
+    w_flat.shape =
+        ShapeTracker::new_with_element_bits(vec![ch, group_out, kernel_product], w.dtype.bits());
+
+    // patches: [N, C, out_spatial_product, kernel_product]
+    // Expand to [N, C, group_out, out_spatial_product, kernel_product]
+    let patches = patches.expand_dim(2, group_out);
+
+    // Expand weight for broadcast: [1, C, group_out, out_spatial_product, kernel_product]
+    let w_expanded = w_flat.expand_dim(0, 1).expand_dim(3, patches.dims()[3]);
+
+    // Element-wise multiply and sum over kernel dim
+    let product = patches * w_expanded;
+    let mut out = product.sum(vec![4]).merge_dims(1, 2);
+    // out: [N, C * group_out, out_spatial_product]
+
+    // Restore spatial dimensions
+    for i in (1..spatial).rev() {
+        out = out.split_dims(2, out_spatial_dims[i]);
+    }
+    // out: [N, C, out_spatial_0, ..., out_spatial_{s-1}]
+
+    out
+}

crates/luminal_python/rust/src/translator/dispatch.rs

@@ -51,6 +51,7 @@ impl<'a> Translator<'a> {
             "torch.ops.aten.sub.Scalar" => self.translate_binary_scalar_op(node, BinaryOp::Sub)?,
             "torch.ops.aten.div.Tensor" => self.translate_binary_op(node, BinaryOp::Div)?,
             "torch.ops.aten.div.Scalar" => self.translate_binary_scalar_op(node, BinaryOp::Div)?,
+            "torch.ops.aten.div.Tensor_mode" => self.translate_div_tensor_mode(node)?,
 
             // Unary ops
             "torch.ops.aten.neg.default" => self.translate_unary_op(node, |a| a * (-1.0))?,
@@ -66,74 +67,75 @@ impl<'a> Translator<'a> {
             }
             "torch.ops.aten.sigmoid.default" => self.translate_unary_op(node, |a| a.sigmoid())?,
             "torch.ops.aten.relu.default" => self.translate_unary_op(node, |a| a.relu())?,
-            "torch.ops.aten.silu.default" => self.translate_unary_op(node, |a| a.swish())?,
             "torch.ops.aten.tanh.default" => self.translate_unary_op(node, |a| a.tanh())?,
             "torch.ops.aten.abs.default" => self.translate_unary_op(node, |a| a.abs())?,
             "torch.ops.aten.log.default" => self.translate_unary_op(node, |a| a.log())?,
+            "torch.ops.aten.log2.default" => self.translate_unary_op(node, |a| a.log2())?,
+            "torch.ops.aten.exp2.default" => self.translate_unary_op(node, |a| a.exp2())?,
+            "torch.ops.aten.sign.default" => self.translate_sign(node)?,
+            "torch.ops.aten.bitwise_not.default" => self.translate_bitwise_not(node)?,
 
             // Cast
             "torch.ops.aten._to_copy.default" => self.translate_to_copy(node)?,
-            "torch.ops.aten.to.dtype" => self.translate_to_dtype(node)?,
-            "torch.ops.aten.to.dtype_layout" => self.translate_to_dtype_layout(node)?,
 
-            // No-op pass-throughs
-            "torch.ops.aten.alias.default"
-            | "torch.ops.aten.detach_.default"
-            | "torch.ops.aten.lift_fresh_copy.default" => self.get_input_tensor(node, 0)?,
-            "torch.ops.aten.dropout.default" => self.get_input_tensor(node, 0)?,
+            // No-op
+            "torch.ops.aten.alias.default" => self.get_input_tensor(node, 0)?,
 
             // Shape ops
-            "torch.ops.aten.view.default"
-            | "torch.ops.aten.reshape.default"
-            | "torch.ops.aten._unsafe_view.default" => self.translate_reshape(node)?,
+            "torch.ops.aten.view.default" => self.translate_reshape(node)?,
             "torch.ops.aten.permute.default" => self.translate_permute(node)?,
-            "torch.ops.aten.transpose.int" => self.translate_transpose(node)?,
-            "torch.ops.aten.t.default" => {
-                let a = self.get_input_tensor(node, 0)?;
-                a.t()
-            }
             "torch.ops.aten.unsqueeze.default" => {
                 let a = self.get_input_tensor(node, 0)?;
                 let dim = self.get_int_arg(node, 1)?;
                 let dim = normalize_dim(dim, a.shape.len() + 1);
                 a.unsqueeze(dim)
             }
-            "torch.ops.aten.squeeze.dim" | "torch.ops.aten.squeeze.default" => {
+            "torch.ops.aten.squeeze.dims" => {
                 let a = self.get_input_tensor(node, 0)?;
-                if node.inputs.len() > 1 {
-                    let dim = self.get_int_arg(node, 1)?;
-                    let dim = normalize_dim(dim, a.shape.len());
-                    a.squeeze(dim)
-                } else {
-                    let mut result = a;
-                    let dims = a.shape.dims;
-                    let mut offset = 0;
-                    for (i, d) in dims.iter().enumerate() {
-                        if d.to_usize() == Some(1) {
-                            result = result.squeeze(i - offset);
-                            offset += 1;
-                        }
+                let dims = self.get_ints_arg(node, 1)?;
+                let ndim = a.shape.len();
+                let mut sorted_dims: Vec<usize> =
+                    dims.iter().map(|&d| normalize_dim(d, ndim)).collect();
+                sorted_dims.sort();
+                let mut result = a;
+                let mut offset = 0;
+                for d in sorted_dims {
+                    if result.shape.dims[d - offset].to_usize() == Some(1) {
+                        result = result.squeeze(d - offset);
+                        offset += 1;
                     }
-                    result
                 }
+                result
             }
             "torch.ops.aten.expand.default" => self.translate_expand(node)?,
-            "torch.ops.aten.contiguous.default" | "torch.ops.aten.clone.default" => {
+            "torch.ops.aten.clone.default" => {
                 let a = self.get_input_tensor(node, 0)?;
                 if !a.shape.is_contiguous() { a + 0.0 } else { a }
             }
+            "torch.ops.aten.argsort.default" => self.translate_argsort(node)?,
 
             // Matmul
-            "torch.ops.aten.mm.default"
-            | "torch.ops.aten.bmm.default"
-            | "torch.ops.aten.matmul.default" => {
+            "torch.ops.aten.mm.default" | "torch.ops.aten.bmm.default" => {
                 let a = self.get_input_tensor(node, 0)?;
                 let b = self.get_input_tensor(node, 1)?;
+                let (a, b) = ensure_same_dtype(a, b);
                 a.matmul(b)
             }
 
-            // Linear
-            "torch.ops.aten.linear.default" => self.translate_linear(node)?,
+            // addmm: beta*input + alpha*(mat1 @ mat2)
+            "torch.ops.aten.addmm.default" => {
+                let input = self.get_input_tensor(node, 0)?;
+                let mat1 = self.get_input_tensor(node, 1)?;
+                let mat2 = self.get_input_tensor(node, 2)?;
+                let beta = self.get_float_arg(node, 3).unwrap_or(1.0) as f32;
+                let alpha = self.get_float_arg(node, 4).unwrap_or(1.0) as f32;
+                let mm = mat1.matmul(mat2);
+                let (input, mm) = broadcast_binary(input, mm);
+                input * beta + mm * alpha
+            }
+
+            // Convolution
+            "torch.ops.aten.convolution.default" => self.translate_conv(node)?,
 
             // Reduction ops
             "torch.ops.aten.sum.dim_IntList" => self.translate_reduction(node, ReductionOp::Sum)?,
@@ -142,16 +144,14 @@ impl<'a> Translator<'a> {
 
             // Slice/index ops
             "torch.ops.aten.slice.Tensor" => self.translate_slice(node)?,
-            "torch.ops.aten.select.int" => self.translate_select(node)?,
             "torch.ops.aten.cat.default" => self.translate_cat(node)?,
-            "torch.ops.aten.index_select.default" => self.translate_index_select(node)?,
             "torch.ops.aten.index.Tensor" => self.translate_index_tensor(node)?,
 
             // Embedding
             "torch.ops.aten.embedding.default" => self.translate_embedding(node)?,
 
             // Softmax
-            "torch.ops.aten._softmax.default" | "torch.ops.aten.softmax.int" => {
+            "torch.ops.aten._softmax.default" => {
                 let a = self.get_input_tensor(node, 0)?;
                 let dim = self.get_int_arg(node, 1)?;
                 let dim = normalize_dim(dim, a.shape.len());
@@ -159,11 +159,12 @@ impl<'a> Translator<'a> {
             }
 
             // LayerNorm
-            "torch.ops.aten.layer_norm.default" => self.translate_layer_norm(node)?,
+            "torch.ops.aten.native_layer_norm.default" => self.translate_layer_norm(node)?,
 
             // Where
             "torch.ops.aten.where.self" => self.translate_where(node)?,
             "torch.ops.aten.where.ScalarOther" => self.translate_where_scalar_other(node)?,
+            "torch.ops.aten.masked_fill.Scalar" => self.translate_masked_fill_scalar(node)?,
 
             // Pow
             "torch.ops.aten.pow.Tensor_Scalar" => {
@@ -179,23 +180,33 @@ impl<'a> Translator<'a> {
             }
 
             // Creation ops
-            "torch.ops.aten.arange.default" | "torch.ops.aten.arange.start" => {
-                self.translate_arange(node)?
-            }
+            "torch.ops.aten.arange.start_step" => self.translate_arange(node)?,
             "torch.ops.aten.full.default" => self.translate_full(node)?,
-            "torch.ops.aten.zeros.default" | "torch.ops.aten.zeros_like.default" => {
-                self.translate_zeros(node)?
-            }
-            "torch.ops.aten.ones.default" | "torch.ops.aten.ones_like.default" => {
-                self.translate_ones(node)?
-            }
-            "torch.ops.aten.new_ones.default" => self.translate_new_ones(node)?,
+            "torch.ops.aten.full_like.default" => self.translate_full_like(node)?,
+            "torch.ops.aten.empty_permuted.default"
+            | "torch.ops.aten.empty.memory_format" => self.translate_empty(node)?,
+            "torch.ops.aten.histc.default" => self.translate_histc(node)?,
 
+            // Grouped matmul (MoE expert dispatch).
+            // aten._grouped_mm is the native op; transformers::grouped_mm_fallback
+            // is a Python-implemented custom_op (transformers/integrations/moe.py)
+            // used by HF MoE when _grouped_mm isn't available for the activation
+            // dtype. Both have identical (input, weight, offs) signature; route
+            // both through the same batched-matmul + group-mask lowering.
+            "torch.ops.aten._grouped_mm.default"
+            | "torch.ops.transformers.grouped_mm_fallback.default" => {
+                self.translate_grouped_mm(node)?
+            }
+            "torch.ops.aten.scalar_tensor.default" => {
+                let val = self.get_float_arg(node, 0)? as f32;
+                self.graph.constant_float(val)
+            }
             // Scalar comparisons
             "torch.ops.aten.gt.Scalar" => self.translate_scalar_comparison(node, |a, s| a.gt(s))?,
             "torch.ops.aten.lt.Scalar" => self.translate_scalar_comparison(node, |a, s| a.lt(s))?,
             "torch.ops.aten.ge.Scalar" => self.translate_scalar_comparison(node, |a, s| a.ge(s))?,
             "torch.ops.aten.le.Scalar" => self.translate_scalar_comparison(node, |a, s| a.le(s))?,
+            "torch.ops.aten.eq.Scalar" => self.translate_scalar_comparison(node, |a, s| a.eq(s))?,
 
             // Tensor comparisons
             "torch.ops.aten.ne.Scalar" => {
@@ -222,7 +233,7 @@ impl<'a> Translator<'a> {
                 let (a, b) = broadcast_binary(a, b);
                 a.le(b)
             }
-            "torch.ops.aten.__and__.Tensor" | "torch.ops.aten.logical_and.default" => {
+            "torch.ops.aten.bitwise_and.Tensor" | "torch.ops.aten.logical_and.default" => {
                 let a = self.get_input_tensor(node, 0)?;
                 let b = self.get_input_tensor(node, 1)?;
                 let (a, b) = broadcast_binary(a, b);
@@ -230,7 +241,11 @@ impl<'a> Translator<'a> {
                 let b = b.cast(DType::F32);
                 (a * b).cast(DType::Bool)
             }
-            "torch.ops.aten.logical_or.default" => {
+            "torch.ops.aten.bitwise_or.Tensor" | "torch.ops.aten.logical_or.default" => {
+                // Both arms use the same bool-OR lowering. Gemma-4's sliding+full
+                // attention mask fusion emits bitwise_or on boolean tensors; the
+                // integer semantics of bitwise_or aren't exercised by any op in
+                // the test suite, so we rely on inputs being boolean-typed.
                 let a = self.get_input_tensor(node, 0)?;
                 let b = self.get_input_tensor(node, 1)?;
                 let (a, b) = broadcast_binary(a, b);
@@ -248,9 +263,7 @@ impl<'a> Translator<'a> {
             }
 
             // Clamp
-            "torch.ops.aten.clamp.default" | "torch.ops.aten.clamp_min.default" => {
-                self.translate_clamp(node)?
-            }
+            "torch.ops.aten.clamp.default" => self.translate_clamp(node)?,
 
             // Cumsum
             "torch.ops.aten.cumsum.default" => {
@@ -265,9 +278,6 @@ impl<'a> Translator<'a> {
                 a.cumsum(dim)
             }
 
-            // Diff
-            "torch.ops.aten.diff.default" => self.translate_diff(node)?,
-
             // Floor / Ceil / Erf (approximations)
             "torch.ops.aten.floor.default" => {
                 let a = self.get_input_tensor(node, 0)?;
@@ -287,24 +297,40 @@ impl<'a> Translator<'a> {
             }
             "torch.ops.aten.erf.default" => {
                 let a = self.get_input_tensor(node, 0)?;
-                // Abramowitz & Stegun approximation 7.1.28 (max error ~1.5e-7)
-                // erf(x) = sign(x) * (1 - poly(t) * exp(-x^2))
-                // where t = 1/(1 + 0.3275911*|x|), poly in Horner form
-                let ax = a.abs();
-                let x2 = a * a;
-                let t = (ax * 0.3275911_f32 + 1.0).reciprocal();
-                // Horner: t*(a1 + t*(a2 + t*(a3 + t*(a4 + t*a5))))
-                let poly = t
-                    * (t * (t
-                        * (t * (t * 1.061_405_4_f32 + (-1.453_152_1_f32)) + 1.421_413_8_f32)
-                        + (-0.284_496_72_f32))
-                        + 0.254_829_6_f32);
-                let result_abs =
-                    self.graph.constant_float(1.0).expand_rhs(a.shape) - poly * (x2 * (-1.0)).exp();
-                // sign(x) = 2*(x >= 0) - 1
-                let zero = self.graph.constant_float(0.0).expand_rhs(a.shape);
-                let sign = a.ge(zero).cast(DType::F32) * 2.0 - 1.0;
-                result_abs * sign
+                self.erf_approx(a)
+            }
+            "torch.ops.aten.gelu.default" => {
+                let a_in = self.get_input_tensor(node, 0)?;
+                // PyTorch's gelu has a kwarg `approximate` (default "none").
+                // "none"  → 0.5 * x * (1 + erf(x / sqrt(2)))           (exact)
+                // "tanh"  → 0.5 * x * (1 + tanh(c * (x + 0.044715*x^3)))
+                //          where c = sqrt(2/pi) ≈ 0.7978845608
+                // Gemma family uses approximate="tanh" but lowering may emit
+                // either form; honour whatever the FX graph carries.
+                let approximate = node.inputs.iter().find_map(|input| {
+                    if input.name == "approximate"
+                        && let Argument::Other(val) = &input.arg
+                    {
+                        return val.as_str().map(|s| s.to_string());
+                    }
+                    None
+                });
+                // Promote to F32 around the constants/comparisons (same reason
+                // as clamp/erf — luminal binary ops assert matching dtypes).
+                let orig = a_in.dtype;
+                let a = if orig == DType::F32 { a_in } else { a_in.cast(DType::F32) };
+                let half = self.graph.constant_float(0.5).expand_rhs(a.shape);
+                let one = self.graph.constant_float(1.0).expand_rhs(a.shape);
+                let result = if approximate.as_deref() == Some("tanh") {
+                    let x2 = a * a;
+                    let inner = a * (x2 * 0.044715_f32 + 1.0) * 0.797_884_56_f32;
+                    half * a * (one + inner.tanh())
+                } else {
+                    let scaled = a * 0.707_106_77_f32; // 1 / sqrt(2)
+                    let erf_val = self.erf_approx(scaled);
+                    half * a * (one + erf_val)
+                };
+                if orig == DType::F32 { result } else { result.cast(orig) }
             }
             "torch.ops.aten.isnan.default" => {
                 let a = self.get_input_tensor(node, 0)?;
@@ -352,45 +378,12 @@ impl<'a> Translator<'a> {
                 let (a, b) = broadcast_binary(a, b);
                 a.gt(b)
             }
-            "torch.ops.aten.ne.Tensor" => {
-                let a = self.get_input_tensor(node, 0)?;
-                let b = self.get_input_tensor(node, 1)?;
-                let (a, b) = ensure_same_dtype(a, b);
-                let (a, b) = broadcast_binary(a, b);
-                a.ne(b)
-            }
 
-            // Reductions without dim arg (full reduce)
-            // Flatten to [1, N] and reduce axis 1 to avoid multi-step HLIR
-            // that CUDA can't schedule (grid (0,1,1) invalid launch).
-            "torch.ops.aten.sum.default" => {
-                let a = self.get_input_tensor(node, 0)?;
-                let total = concrete_numel(&a)?;
-                let mut flat = a;
-                flat.shape = ShapeTracker::new(vec![1, total]);
-                flat.sum(vec![1])
-            }
-            "torch.ops.aten.mean.default" => {
-                let a = self.get_input_tensor(node, 0)?;
-                let total = concrete_numel(&a)?;
-                let mut flat = a;
-                flat.shape = ShapeTracker::new(vec![1, total]);
-                flat.sum(vec![1]) / total as f32
-            }
-            "torch.ops.aten.max.default" => {
-                let a = self.get_input_tensor(node, 0)?;
-                let total = concrete_numel(&a)?;
-                let mut flat = a;
-                flat.shape = ShapeTracker::new(vec![1, total]);
-                flat.max(vec![1])
-            }
-            "torch.ops.aten.min.default" => {
-                let a = self.get_input_tensor(node, 0)?;
-                let total = concrete_numel(&a)?;
-                let mut flat = a;
-                flat.shape = ShapeTracker::new(vec![1, total]);
-                flat.min(vec![1])
-            }
+            // Full-reduce variants (no dim arg) — handled by translate_reduction fallback
+            "torch.ops.aten.sum.default" => self.translate_reduction(node, ReductionOp::Sum)?,
+            "torch.ops.aten.mean.default" => self.translate_reduction(node, ReductionOp::Mean)?,
+            "torch.ops.aten.max.default" => self.translate_reduction(node, ReductionOp::Max)?,
+            "torch.ops.aten.min.default" => self.translate_reduction(node, ReductionOp::Min)?,
             "torch.ops.aten.amin.default" => self.translate_reduction(node, ReductionOp::Min)?,
 
             // Gather (axis-aware)
@@ -398,7 +391,13 @@ impl<'a> Translator<'a> {
 
             // Scatter ops
             "torch.ops.aten.scatter.src" => self.translate_scatter_src(node)?,
-            "torch.ops.aten.index_put_.default" => self.translate_index_put(node)?,
+            "torch.ops.aten.scatter.value" => self.translate_scatter_value(node)?,
+            "torch.ops.aten.index_put_.default" | "torch.ops.aten.index_put.default" => {
+                self.translate_index_put(node)?
+            }
+
+            // Integer routing math
+            "torch.ops.aten.floor_divide.default" => self.translate_floor_divide(node)?,
 
             // Triangular
             "torch.ops.aten.tril.default" => self.translate_tril(node)?,
@@ -410,13 +409,14 @@ impl<'a> Translator<'a> {
                 return Ok(());
             }
 
-            // Split
-            "torch.ops.aten.split.Tensor" | "torch.ops.aten.split_with_sizes.default" => {
-                self.translate_split(node)?
+            // Sort — handles its own output storage, returns early
+            "torch.ops.aten.sort.default" => {
+                self.translate_sort(node)?;
+                return Ok(());
             }
 
-            // One-hot
-            "torch.ops.aten.one_hot.default" => self.translate_one_hot(node)?,
+            // Split
+            "torch.ops.aten.split_with_sizes.default" => self.translate_split_with_sizes(node)?,
 
             // Fmod
             "torch.ops.aten.fmod.Tensor" => {
@@ -425,12 +425,8 @@ impl<'a> Translator<'a> {
                 let (a, b) = broadcast_binary(a, b);
                 a % b
             }
-            "torch.ops.aten.fmod.Scalar" | "torch.ops.aten.remainder.Scalar" => {
-                let a = self.get_input_tensor(node, 0)?;
-                let val = self.get_float_arg(node, 1)? as f32;
-                let b = self.graph.constant_float(val).expand_rhs(a.shape);
-                a % b
-            }
+            // Prod reduction
+            "torch.ops.aten.prod.dim_int" => self.translate_reduction(node, ReductionOp::Prod)?,
 
             other => {
                 bail!("Unsupported ATen op: {other}");
@@ -444,15 +440,6 @@ impl<'a> Translator<'a> {
     }
 }
 
-/// Compute total element count, returning an error if any dimension is symbolic.
-fn concrete_numel(a: &GraphTensor) -> Result<usize> {
-    a.dims().iter().try_fold(1usize, |acc, d| {
-        d.to_usize().map(|v| acc * v).ok_or_else(|| {
-            anyhow::anyhow!("Full reduction requires concrete dimensions, got symbolic dim")
-        })
-    })
-}
-
 impl<'a> Translator<'a> {
     fn translate_scalar_comparison(
         &mut self,

crates/luminal_python/rust/src/translator/matmul.rs

@@ -1,23 +0,0 @@
-use anyhow::Result;
-use luminal::prelude::*;
-
-use crate::pt2_schema::*;
-use crate::pt2_util::broadcast_binary;
-
-use super::Translator;
-
-impl<'a> Translator<'a> {
-    pub(crate) fn translate_linear(&mut self, node: &Node) -> Result<GraphTensor> {
-        let input = self.get_input_tensor(node, 0)?;
-        let weight = self.get_input_tensor(node, 1)?;
-        let result = input.matmul(weight.t());
-
-        if node.inputs.len() > 2
-            && let Ok(bias) = self.get_input_tensor(node, 2)
-        {
-            let (result, bias) = broadcast_binary(result, bias);
-            return Ok(result + bias);
-        }
-        Ok(result)
-    }
-}

crates/luminal_python/rust/src/translator/mod.rs

@@ -3,8 +3,8 @@
 //! Walks the parsed PT2 graph and constructs an equivalent Luminal computation graph.
 
 mod binary;
+mod conv;
 mod dispatch;
-mod matmul;
 mod movement;
 mod reduction;
 mod tensor;
@@ -18,6 +18,7 @@ use luminal::prelude::*;
 
 use crate::pt2_parser::{InputKind, ParsedPT2, SymDimMap};
 use crate::pt2_schema::*;
+use crate::pt2_util;
 
 /// Result of translating a PT2 graph to a Luminal graph.
 pub struct TranslatedGraph {
@@ -67,6 +68,9 @@ impl<'a> Translator<'a> {
     fn translate_graph(&mut self) -> Result<()> {
         self.create_inputs()?;
 
+        // Per-block partitioning is now handled automatically by the upstream
+        // loop-rolling prepass; this translator no longer needs to insert
+        // manual graph breaks at RMSNorm boundaries.
         let nodes = &self.parsed.program.graph_module.graph.nodes;
         for (i, node) in nodes.iter().enumerate() {
             self.translate_node(node)
@@ -76,7 +80,13 @@ impl<'a> Translator<'a> {
         let output_names = self.parsed.output_names();
         for name in &output_names {
             let tensor = self.get_tensor(name)?;
-            let tensor = tensor + 0.0;
+            let tensor = if tensor.dtype == DType::Bool {
+                tensor.cast(DType::Int).cast(DType::Bool)
+            } else if tensor.dtype == DType::Int {
+                tensor
+            } else {
+                tensor + 0.0
+            };
             tensor.output();
             self.output_ids.push((name.clone(), tensor.id));
         }
@@ -97,7 +107,12 @@ impl<'a> Translator<'a> {
                         .tensor_meta(graph_name)
                         .with_context(|| format!("Missing tensor meta for param {graph_name}"))?;
                     let shape = self.tensor_meta_to_shape(meta)?;
-                    let tensor = self.graph.named_tensor(original_name, shape);
+                    let dtype = pt2_util::torch_dtype_int_to_luminal(meta.dtype);
+                    let tensor = self
+                        .graph
+                        .named_tensor(original_name, shape)
+                        .as_dtype(dtype);
+                    tensor.persist();
                     self.tensors.insert(graph_name.clone(), tensor);
                 }
                 InputKind::Buffer {
@@ -109,7 +124,12 @@ impl<'a> Translator<'a> {
                         .tensor_meta(graph_name)
                         .with_context(|| format!("Missing tensor meta for buffer {graph_name}"))?;
                     let shape = self.tensor_meta_to_shape(meta)?;
-                    let tensor = self.graph.named_tensor(original_name, shape);
+                    let dtype = pt2_util::torch_dtype_int_to_luminal(meta.dtype);
+                    let tensor = self
+                        .graph
+                        .named_tensor(original_name, shape)
+                        .as_dtype(dtype);
+                    tensor.persist();
                     self.tensors.insert(graph_name.clone(), tensor);
                 }
                 InputKind::UserInput { graph_name } => {
@@ -118,7 +138,8 @@ impl<'a> Translator<'a> {
                         .tensor_meta(graph_name)
                         .with_context(|| format!("Missing tensor meta for input {graph_name}"))?;
                     let shape = self.tensor_meta_to_shape(meta)?;
-                    let tensor = self.graph.named_tensor(graph_name, shape);
+                    let dtype = pt2_util::torch_dtype_int_to_luminal(meta.dtype);
+                    let tensor = self.graph.named_tensor(graph_name, shape).as_dtype(dtype);
                     self.user_input_ids.push((graph_name.clone(), tensor.id));
                     self.tensors.insert(graph_name.clone(), tensor);
                 }
@@ -138,7 +159,6 @@ impl<'a> Translator<'a> {
 
     // --- Helper methods ---
 
-    /// Look up tensor metadata by name, checking subgraph extras first.
     pub(crate) fn tensor_meta(&self, name: &str) -> Option<&TensorMeta> {
         self.extra_tensor_values
             .get(name)
@@ -319,3 +339,4 @@ impl<'a> Translator<'a> {
         None
     }
 }
+

crates/luminal_python/rust/src/translator/movement.rs

@@ -6,6 +6,11 @@ use crate::pt2_util::*;
 
 use super::Translator;
 
+const SCATTER_INPUT_ARG: usize = 0;
+const SCATTER_DIM_ARG: usize = 1;
+const SCATTER_INDEX_ARG: usize = 2;
+const SCATTER_VALUE_ARG: usize = 3;
+
 impl<'a> Translator<'a> {
     pub(crate) fn translate_reshape(&mut self, node: &Node) -> Result<GraphTensor> {
         let a = self.get_input_tensor(node, 0)?;
@@ -49,15 +54,6 @@ impl<'a> Translator<'a> {
         Ok(a.permute(axes))
     }
 
-    pub(crate) fn translate_transpose(&mut self, node: &Node) -> Result<GraphTensor> {
-        let a = self.get_input_tensor(node, 0)?;
-        let dim0 = self.get_int_arg(node, 1)?;
-        let dim1 = self.get_int_arg(node, 2)?;
-        let dim0 = normalize_dim(dim0, a.shape.len());
-        let dim1 = normalize_dim(dim1, a.shape.len());
-        Ok(a.transpose(dim0, dim1))
-    }
-
     pub(crate) fn translate_expand(&mut self, node: &Node) -> Result<GraphTensor> {
         let mut a = self.get_input_tensor(node, 0)?;
         let neg1_expr = Expression::from(-1i32);
@@ -124,20 +120,6 @@ impl<'a> Translator<'a> {
         Ok(a.slice_along(start..end, dim))
     }
 
-    pub(crate) fn translate_select(&mut self, node: &Node) -> Result<GraphTensor> {
-        let a = self.get_input_tensor(node, 0)?;
-        let dim = self.get_int_arg(node, 1)?;
-        let dim = normalize_dim(dim, a.shape.len());
-        let index = self.get_int_arg(node, 2)?;
-        let index = if index < 0 {
-            bail!("Negative select index not yet supported");
-        } else {
-            index as usize
-        };
-
-        Ok(a.slice_along(index..index + 1, dim).squeeze(dim))
-    }
-
     pub(crate) fn translate_cat(&mut self, node: &Node) -> Result<GraphTensor> {
         let tensors: Vec<GraphTensor> = if let Some(names) = node.inputs[0].arg.as_tensors() {
             names
@@ -184,31 +166,6 @@ impl<'a> Translator<'a> {
         Ok(result)
     }
 
-    pub(crate) fn translate_index_select(&mut self, node: &Node) -> Result<GraphTensor> {
-        let a = self.get_input_tensor(node, 0)?;
-        let dim = self.get_int_arg(node, 1)?;
-        let dim = normalize_dim(dim, a.shape.len());
-        let indices = self.get_input_tensor(node, 2)?.cast(DType::Int);
-        let src_dims = a.shape.dims;
-        let idx_len = indices.shape.dims[0];
-
-        // Reshape 1D indices [K] → [1,..,K,..,1] with K at position `dim`
-        let mut idx = indices;
-        for _ in 0..dim {
-            idx = idx.unsqueeze(0);
-        }
-        for _ in (dim + 1)..src_dims.len() {
-            idx = idx.expand_dim(idx.shape.len(), Expression::from(1usize));
-        }
-
-        // Expand to output shape: src_dims with dim replaced by idx_len
-        let mut target: Vec<Expression> = src_dims.to_vec();
-        target[dim] = idx_len;
-        idx.shape.expand(target);
-
-        Ok(a.gather_elements(idx, dim))
-    }
-
     pub(crate) fn translate_embedding(&mut self, node: &Node) -> Result<GraphTensor> {
         let weight = self.get_input_tensor(node, 0)?;
         let indices = self.get_input_tensor(node, 1)?;
@@ -407,32 +364,130 @@ impl<'a> Translator<'a> {
         Ok(a.scatter_elements(indices.cast(DType::Int), src, dim))
     }
 
-    pub(crate) fn translate_index_put(&mut self, node: &Node) -> Result<GraphTensor> {
-        let a = self.get_input_tensor(node, 0)?;
-        let index_names = node.inputs[1]
-            .arg
-            .as_tensors()
-            .context("index_put: indices not as_tensors")?;
-        let values = self.get_input_tensor(node, 2)?;
-
-        if index_names.len() == 1 {
-            let indices = self.get_tensor(&index_names[0].name)?.cast(DType::Int);
-            // scatter_nd expects indices of shape [batch, K] where K = number of index dims.
-            // PT2's index_put gives 1D indices [batch]; reshape to [batch, 1].
-            let indices = if indices.shape.len() == 1 {
-                indices.expand_dim(1, Expression::from(1usize))
-            } else {
-                indices
-            };
-            Ok(a.scatter_nd(indices, values))
+    pub(crate) fn translate_scatter_value(&mut self, node: &Node) -> Result<GraphTensor> {
+        let a = self.get_input_tensor(node, SCATTER_INPUT_ARG)?;
+        let dim = self.get_int_arg(node, SCATTER_DIM_ARG)?;
+        let dim = normalize_dim(dim, a.shape.len());
+        let indices = self.get_input_tensor(node, SCATTER_INDEX_ARG)?;
+        let value_arg = &node
+            .inputs
+            .get(SCATTER_VALUE_ARG)
+            .context("scatter.value missing value input")?
+            .arg;
+        let value = if let Some(b) = value_arg.as_bool() {
+            self.graph.constant(if b { 1 } else { 0 }).cast(a.dtype)
+        } else if let Some(i) = value_arg.as_int() {
+            self.graph.constant(i).cast(a.dtype)
+        } else if let Some(f) = value_arg.as_float() {
+            self.graph.constant_float(f as f32).cast(a.dtype)
         } else {
-            bail!("index_put with multiple index tensors not yet supported");
+            bail!("scatter.value: unsupported scalar argument {:?}", value_arg);
         }
+        .expand_rhs(indices.shape);
+        Ok(a.scatter_elements(indices.cast(DType::Int), value, dim))
     }
 
-    pub(crate) fn translate_split(&mut self, node: &Node) -> Result<GraphTensor> {
+    pub(crate) fn translate_index_put(&mut self, node: &Node) -> Result<GraphTensor> {
         let a = self.get_input_tensor(node, 0)?;
-        let split_size = self.get_int_arg(node, 1)? as usize;
+        let values = self.get_input_tensor(node, 2)?;
+
+        // --- all-tensor indices: bool-mask blend or scatter_nd ---
+        if let Some(index_names) = node.inputs[1].arg.as_tensors() {
+            if index_names.len() == 1 {
+                let idx_tensor = self.get_tensor(&index_names[0].name)?;
+
+                // Boolean-mask index_put: when the only index is a Bool tensor whose
+                // shape matches the data tensor, PyTorch semantics are
+                //   data[mask] = value   ↔   where(mask, value, data)
+                // NOT a scatter into positions. Casting the Bool mask to Int and
+                // feeding it to scatter_nd would reinterpret True/False as row
+                // indices 1/0 and silently corrupt the data. Reproducer:
+                //   x = arange(16).reshape(4, 4); mask = zeros(4, 4, dtype=bool)
+                //   y = x.clone(); y[mask] = 99   # eager: y == x (no-op)
+                // Pre-fix the compiled graph wrote 99 to row 0; this branch
+                // ensures the bool-mask path lowers to a where-blend instead.
+                if idx_tensor.dtype == DType::Bool && idx_tensor.shape.dims == a.shape.dims {
+                    let mask_f = idx_tensor.cast(a.dtype);
+                    let values_b = values.cast(a.dtype).expand_rhs(a.shape);
+                    // Implements where(mask, value, a) as
+                    //     a*(1 - mask) + value*mask
+                    // — works without a dedicated cond op for any numeric dtype.
+                    let one = self
+                        .graph
+                        .constant_float(1.0)
+                        .cast(a.dtype)
+                        .expand_rhs(a.shape);
+                    return Ok(a * (one - mask_f) + values_b * mask_f);
+                }
+
+                // Integer-index scatter: index_put with indices=[idx_tensor] writes
+                // into dim 0 of `a` at every position named in idx_tensor (flattened),
+                // broadcasting values across the trailing dims of `a`. Always pad
+                // a trailing size-1 dim so rank-1 and rank-N cases share a path.
+                let indices = idx_tensor.cast(DType::Int);
+                let new_last = indices.shape.len();
+                let indices = indices.expand_dim(new_last, Expression::from(1usize));
+                return Ok(a.scatter_nd(indices, values));
+            }
+            bail!("index_put with multiple all-tensor indices not yet supported");
+        }
+
+        // --- optional-tensor indices: [None, arange_tensor, None, ...] ---
+        // Each None means "all of that dimension"; one tensor means "index into that dim".
+        // StaticCache uses this for KV updates: cache[:, :, position, :] = new_value.
+        if let Some(opt_tensors) = node.inputs[1].arg.as_optional_tensors() {
+            use crate::pt2_schema::OptionalTensorEntry;
+            let mut first_non_none_dim = 0usize;
+            let mut idx_name: Option<String> = None;
+            let mut non_none_count = 0usize;
+
+            for (i, entry) in opt_tensors.iter().enumerate() {
+                if let OptionalTensorEntry::Tensor(t) = entry {
+                    if idx_name.is_none() {
+                        first_non_none_dim = i;
+                    }
+                    idx_name = Some(t.as_tensor.name.clone());
+                    non_none_count += 1;
+                }
+            }
+
+            if non_none_count != 1 {
+                bail!(
+                    "index_put with optional tensors: only single non-None index supported \
+                     (got {non_none_count})"
+                );
+            }
+
+            let mut indices = self.get_tensor(&idx_name.unwrap())?.cast(DType::Int);
+
+            // Expand 1-D indices [P] to values.shape for scatter_elements:
+            // Build [1, ..., 1, P, 1, ..., 1] with P at first_non_none_dim, then broadcast.
+            let rank = a.shape.len();
+            // Insert singleton dims before first_non_none_dim
+            for i in 0..first_non_none_dim {
+                indices = indices.expand_dim(i, Expression::from(1usize));
+            }
+            // Insert singleton dims after first_non_none_dim
+            let current_rank = indices.shape.len();
+            for j in current_rank..rank {
+                indices = indices.expand_dim(j, Expression::from(1usize));
+            }
+            // Broadcast singletons to values shape
+            let values_shape: Vec<Expression> = values.shape.dims[..rank].to_vec();
+            indices.shape.expand(values_shape);
+
+            return Ok(a.scatter_elements(indices, values, first_non_none_dim));
+        }
+
+        bail!(
+            "index_put: unsupported indices format: {:?}",
+            node.inputs[1].arg
+        )
+    }
+
+    pub(crate) fn translate_split_with_sizes(&mut self, node: &Node) -> Result<GraphTensor> {
+        let a = self.get_input_tensor(node, 0)?;
+        let sizes = self.get_ints_arg(node, 1)?;
         let dim = if node.inputs.len() > 2 {
             self.get_int_arg(node, 2).unwrap_or(0)
         } else {
@@ -440,35 +495,32 @@ impl<'a> Translator<'a> {
         };
         let dim = normalize_dim(dim, a.shape.len());
 
-        let dim_size = a.shape.dims[dim];
-        if let Some(total) = dim_size.to_usize() {
-            // Collect output names from as_tensors (multi-output) or as_tensor (single)
-            let output_names: Vec<String> = node
-                .outputs
-                .first()
-                .and_then(|o| o.as_tensors.as_ref())
-                .map(|ts| ts.iter().map(|t| t.name.clone()).collect())
-                .unwrap_or_else(|| {
-                    node.outputs
-                        .iter()
-                        .filter_map(|o| o.as_tensor.as_ref().map(|t| t.name.clone()))
-                        .collect()
-                });
+        let output_names: Vec<String> = node
+            .outputs
+            .first()
+            .and_then(|o| o.as_tensors.as_ref())
+            .map(|ts| ts.iter().map(|t| t.name.clone()).collect())
+            .unwrap_or_else(|| {
+                node.outputs
+                    .iter()
+                    .filter_map(|o| o.as_tensor.as_ref().map(|t| t.name.clone()))
+                    .collect()
+            });
 
-            // Store each chunk under its output name
-            for (i, out_name) in output_names.iter().enumerate() {
-                let start = i * split_size;
-                let end = ((i + 1) * split_size).min(total);
-                if start < total {
-                    let chunk = a.slice_along(start..end, dim);
-                    self.tensors.insert(out_name.clone(), chunk);
-                }
+        let mut offset = 0usize;
+        let mut first_chunk = None;
+        for (i, &size) in sizes.iter().enumerate() {
+            let size = size as usize;
+            let chunk = a.slice_along(offset..offset + size, dim);
+            if let Some(name) = output_names.get(i) {
+                self.tensors.insert(name.clone(), chunk);
             }
-
-            // Return the first chunk
-            Ok(a.slice_along(0..split_size.min(total), dim))
-        } else {
-            Ok(a.slice_along(0..split_size, dim))
+            if i == 0 {
+                first_chunk = Some(chunk);
+            }
+            offset += size;
         }
+
+        first_chunk.ok_or_else(|| anyhow::anyhow!("split_with_sizes: empty sizes list"))
     }
 }

crates/luminal_python/rust/src/translator/reduction.rs

@@ -6,6 +6,15 @@ use crate::pt2_util::*;
 
 use super::Translator;
 
+/// Compute total element count, returning an error if any dimension is symbolic.
+fn concrete_numel(a: &GraphTensor) -> Result<usize> {
+    a.dims().iter().try_fold(1usize, |acc, d| {
+        d.to_usize().map(|v| acc * v).ok_or_else(|| {
+            anyhow::anyhow!("Full reduction requires concrete dimensions, got symbolic dim")
+        })
+    })
+}
+
 impl<'a> Translator<'a> {
     pub(crate) fn translate_reduction(
         &mut self,
@@ -13,21 +22,42 @@ impl<'a> Translator<'a> {
         op: ReductionOp,
     ) -> Result<GraphTensor> {
         let a = self.get_input_tensor(node, 0)?;
-        let dims = self.get_ints_arg(node, 1)?;
-        let keepdim = if node.inputs.len() > 2 {
-            self.get_bool_arg(node, 2).unwrap_or(false)
-        } else {
-            false
-        };
 
-        let ndim = a.shape.len();
-        let axes: Vec<usize> = dims.iter().map(|&d| normalize_dim(d, ndim)).collect();
+        // Try to get dims arg; if missing or empty, fall back to full reduce
+        let dims_result = self.get_ints_arg(node, 1);
+        let (axes, keepdim) = match dims_result {
+            Ok(ref dims) if !dims.is_empty() => {
+                let ndim = a.shape.len();
+                let axes: Vec<usize> = dims.iter().map(|&d| normalize_dim(d, ndim)).collect();
+                let keepdim = if node.inputs.len() > 2 {
+                    self.get_bool_arg(node, 2).unwrap_or(false)
+                } else {
+                    false
+                };
+                (axes, keepdim)
+            }
+            _ => {
+                // Full reduce: flatten to [1, N] and reduce axis 1
+                let total = concrete_numel(&a)?;
+                let mut flat = a;
+                flat.shape = ShapeTracker::new(vec![1, total]);
+                let result = match op {
+                    ReductionOp::Sum => flat.sum(vec![1]),
+                    ReductionOp::Mean => flat.sum(vec![1]) / total as f32,
+                    ReductionOp::Max => flat.max(vec![1]),
+                    ReductionOp::Min => flat.min(vec![1]),
+                    ReductionOp::Prod => flat.prod(vec![1]),
+                };
+                return Ok(result);
+            }
+        };
 
         let mut result = match op {
             ReductionOp::Sum => a.sum(axes.clone()),
             ReductionOp::Mean => a.mean(axes.clone()),
             ReductionOp::Max => a.max(axes.clone()),
             ReductionOp::Min => a.min(axes.clone()),
+            ReductionOp::Prod => a.prod(axes.clone()),
         };
 
         if keepdim {

crates/luminal_python/rust/src/translator/tensor.rs

@@ -6,6 +6,27 @@ use crate::pt2_util::*;
 
 use super::Translator;
 
+const FULL_SHAPE_ARG: usize = 0;
+const FULL_VALUE_ARG: usize = 1;
+
+const FULL_LIKE_INPUT_ARG: usize = 0;
+const FULL_LIKE_VALUE_ARG: usize = 1;
+
+const TOPK_INPUT_ARG: usize = 0;
+const TOPK_K_ARG: usize = 1;
+const TOPK_DIM_ARG: usize = 2;
+
+const SORT_INPUT_ARG: usize = 0;
+const SORT_DIM_ARG: usize = 1;
+const SORT_DESCENDING_ARG: usize = 2;
+
+const WHERE_COND_ARG: usize = 0;
+const WHERE_X_ARG: usize = 1;
+const WHERE_OTHER_ARG: usize = 2;
+
+const TRIANGULAR_INPUT_ARG: usize = 0;
+const TRIANGULAR_DIAGONAL_ARG: usize = 1;
+
 impl<'a> Translator<'a> {
     pub(crate) fn translate_arange(&mut self, node: &Node) -> Result<GraphTensor> {
         let positional_args: Vec<Expression> = node
@@ -18,31 +39,124 @@ impl<'a> Translator<'a> {
         match positional_args.len() {
             0 => anyhow::bail!("arange: no positional args found"),
             1 => Ok(self.graph.arange(positional_args[0])),
-            _ => Ok(self
+            2 => Ok(self
                 .graph
                 .arange_options(positional_args[0], positional_args[1], 1)),
+            _ => Ok(self.graph.arange_options(
+                positional_args[0],
+                positional_args[1],
+                positional_args[2],
+            )),
         }
     }
 
     pub(crate) fn translate_full(&mut self, node: &Node) -> Result<GraphTensor> {
+        let shape = self.get_exprs_arg(node, FULL_SHAPE_ARG)?;
+        // fill_value can be float, int, or bool after decomposition
+        let val = if let Ok(f) = self.get_float_arg(node, FULL_VALUE_ARG) {
+            f as f32
+        } else if let Ok(b) = self.get_bool_arg(node, FULL_VALUE_ARG) {
+            if b { 1.0 } else { 0.0 }
+        } else {
+            anyhow::bail!(
+                "full: unsupported fill value type: {:?}",
+                node.inputs.get(FULL_VALUE_ARG)
+            );
+        };
+        let dtype = self.output_meta_dtype(node)?;
+        let value = self.graph.constant_float(val).cast(dtype);
+        Ok(if shape.is_empty() {
+            value
+        } else {
+            value.expand_rhs(shape)
+        })
+    }
+
+    /// Translate `aten.histc.default(input, bins, min, max)` → `Tensor[bins]`.
+    ///
+    /// Counts how many input elements fall in each of `bins` equal-width
+    /// buckets over `[min, max]`. PyTorch's histc accepts only 1D input;
+    /// HF MoE forwards emit it on flattened expert-assignment tensors to
+    /// produce per-expert token counts (one_hot + sum, essentially).
+    ///
+    /// Implementation: arange over bins, broadcast to [G, N], element-wise
+    /// `(lower <= input < upper)` into a F32 mask, sum over the input axis.
+    /// The right edge of the last bin is technically inclusive in PyTorch;
+    /// we treat it as exclusive — for the typical MoE use (integer expert
+    /// IDs in `[0, num_experts)`), no input ever equals `max` so this is
+    /// indistinguishable.
+    pub(crate) fn translate_histc(&mut self, node: &Node) -> Result<GraphTensor> {
+        let input = self.get_input_tensor(node, 0)?;
+        let bins = self.get_int_arg(node, 1)? as usize;
+        let min_val = self.get_float_arg(node, 2)? as f32;
+        let max_val = self.get_float_arg(node, 3)? as f32;
+
+        anyhow::ensure!(
+            input.shape.len() == 1,
+            "histc: only 1D input supported (got {}D)",
+            input.shape.len()
+        );
+        let n = input.shape.dims[0];
+        let g = Expression::from(bins);
+
+        let input_f = input.cast(DType::F32);
+        let step = (max_val - min_val) / bins as f32;
+
+        // Per-bin lower edges: arange(bins) * step + min.
+        let bin_idx = self.graph.arange(g).cast(DType::F32);
+        let lower_1d = bin_idx * step + min_val;
+        let upper_1d = lower_1d + step;
+
+        // Broadcast to [G, N] and produce the boolean mask.
+        let input_b = input_f.expand_dim(0, g);
+        let lower = lower_1d.expand_dim(1, n);
+        let upper = upper_1d.expand_dim(1, n);
+
+        let in_lower = input_b.ge(lower).cast(DType::F32);
+        let in_upper = input_b.lt(upper).cast(DType::F32);
+        let mask = in_lower * in_upper;
+
+        Ok(mask.sum(1))
+    }
+
+    /// Translate `aten.empty_permuted.default(size, physical_layout, **kwargs)`
+    /// → zero-filled tensor of shape `size`.
+    ///
+    /// PyTorch's `empty_permuted` allocates uninitialized memory with a given
+    /// stride permutation; downstream code typically overwrites every element
+    /// before reading. Luminal's tensor abstraction doesn't expose strides, so
+    /// the physical_layout hint is irrelevant — we just emit a zero tensor of
+    /// the requested shape and dtype. (Same approach works for `aten.empty`
+    /// variants when they show up.)
+    pub(crate) fn translate_empty(&mut self, node: &Node) -> Result<GraphTensor> {
         let shape = self.get_exprs_arg(node, 0)?;
-        let val = self.get_float_arg(node, 1)? as f32;
-        Ok(self.graph.constant_float(val).expand_rhs(shape))
+        let dtype = self.output_meta_dtype(node)?;
+        let value = self.graph.constant_float(0.0).cast(dtype);
+        Ok(if shape.is_empty() {
+            value
+        } else {
+            value.expand_rhs(shape)
+        })
     }
 
-    pub(crate) fn translate_zeros(&mut self, node: &Node) -> Result<GraphTensor> {
-        self.translate_constant_fill(node, 0.0)
+    pub(crate) fn translate_full_like(&mut self, node: &Node) -> Result<GraphTensor> {
+        let reference = self.get_input_tensor(node, FULL_LIKE_INPUT_ARG)?;
+        let val = if let Ok(f) = self.get_float_arg(node, FULL_LIKE_VALUE_ARG) {
+            f as f32
+        } else if let Ok(b) = self.get_bool_arg(node, FULL_LIKE_VALUE_ARG) {
+            if b { 1.0 } else { 0.0 }
+        } else {
+            anyhow::bail!(
+                "full_like: unsupported fill value type: {:?}",
+                node.inputs.get(FULL_LIKE_VALUE_ARG)
+            );
+        };
+        let dtype = self.output_meta_dtype(node)?;
+        let value = self.graph.constant_float(val).cast(dtype);
+        Ok(value.expand_rhs(reference.shape))
     }
 
-    pub(crate) fn translate_ones(&mut self, node: &Node) -> Result<GraphTensor> {
-        self.translate_constant_fill(node, 1.0)
-    }
-
-    pub(crate) fn translate_new_ones(&mut self, node: &Node) -> Result<GraphTensor> {
-        self.translate_constant_fill(node, 1.0)
-    }
-
-    fn translate_constant_fill(&mut self, node: &Node, val: f32) -> Result<GraphTensor> {
+    fn output_meta_dtype(&self, node: &Node) -> Result<DType> {
         let output_name = node
             .outputs
             .first()
@@ -51,32 +165,127 @@ impl<'a> Translator<'a> {
             .unwrap_or_default();
         let meta = self
             .tensor_meta(&output_name)
-            .context("Missing tensor meta for constant fill output")?;
-        let shape = self.tensor_meta_to_shape(meta)?;
-        if shape.is_empty() {
-            Ok(self.graph.constant_float(val))
-        } else {
-            Ok(self.graph.constant_float(val).expand_rhs(shape))
-        }
+            .context("Missing tensor meta for output dtype")?;
+        Ok(torch_dtype_int_to_luminal(meta.dtype))
+    }
+
+    /// Translate `aten._grouped_mm.default(input, weight, offs)` → `Tensor[S, N]`.
+    ///
+    /// Grouped matmul: `input` is `[S, K]` (tokens sorted by expert), `weight` is
+    /// `[G, K, N]` (per-expert weights), `offs` is `[G]` cumulative token counts.
+    /// Output `[S, N]` where token m (in group g s.t. `offs[g-1] <= m < offs[g]`)
+    /// is multiplied by `weight[g]`.
+    ///
+    /// Implementation: for each token m we (a) compute its expert id from offs,
+    /// (b) gather only that expert's `[K, N]` slice from weight, and (c) do a
+    /// single per-token matmul. The gather pattern mirrors the rust qwen3_moe
+    /// example's `gather_experts`, which the GLUMoE host-op fusion in
+    /// `luminal_cuda_lite` is designed to recognise.
+    ///
+    /// Why not the straightforward `[G, S, K] @ [G, K, N] → [G, S, N]` + mask:
+    /// it forces a full F32 cast of the entire `[G, K, N]` weight tensor as
+    /// search-time intermediate, which OOMs on real MoE checkpoints
+    /// (Qwen3-30B-A3B: 1.5 GB / layer × 48 layers for gate-up alone). Gathering
+    /// first keeps the F32 cast on `[S, K, N]` instead — for prefill (S = top_k)
+    /// that is a 16× shrink (G=128, top_k=8).
+    ///
+    /// `offs` flows through as a runtime tensor — the routing decision is computed
+    /// at execution time by the gate network and the same compiled graph handles
+    /// any routing pattern without recompilation.
+    pub(crate) fn translate_grouped_mm(&mut self, node: &Node) -> Result<GraphTensor> {
+        let input = self.get_input_tensor(node, 0)?;
+        let weight = self.get_input_tensor(node, 1)?;
+        let offs = self.get_input_tensor(node, 2)?;
+
+        anyhow::ensure!(
+            input.shape.len() == 2,
+            "_grouped_mm: input must be 2D, got {}D",
+            input.shape.len()
+        );
+        anyhow::ensure!(
+            weight.shape.len() == 3,
+            "_grouped_mm: weight must be 3D, got {}D",
+            weight.shape.len()
+        );
+        anyhow::ensure!(
+            offs.shape.len() == 1,
+            "_grouped_mm: offs must be 1D, got {}D",
+            offs.shape.len()
+        );
+
+        let s = input.shape.dims[0];
+        let g = weight.shape.dims[0];
+        let k = weight.shape.dims[1];
+        let n = weight.shape.dims[2];
+
+        // expert_id[m] = number of g s.t. m >= offs[g]
+        //              = first g s.t. m < offs[g], i.e. the expert assigned to m.
+        // Clamp to [0, G-1] before using as gather index. Matches HF MoE's
+        // `expert_ids.clamp(0, num_experts-1)` for invalid IDs from EP, AND
+        // protects search-time profiling: dummy-1 input bytes give offs=[1,…,1],
+        // which makes `m >= offs[g]` true for m≥1 and pushes expert_id to G,
+        // out of bounds for the weight gather. Clamping keeps the gather safe.
+        let g_max_f = (g
+            .to_usize()
+            .context("_grouped_mm: G (num_experts) must be concrete")?
+            as f32)
+            - 1.0;
+        let offs_f = offs.cast(DType::F32);
+        let s_arange_f = self.graph.arange(s).cast(DType::F32);
+        let ge_boundary = s_arange_f
+            .expand_dim(0, g)
+            .ge(offs_f.expand_dim(1, s))
+            .cast(DType::F32);
+        let expert_id = ge_boundary
+            .sum(0)
+            .minimum_f32(g_max_f)
+            .cast(DType::Int); // [S] Int
+
+        // Flat gather index into weight (treated as a length-G*K*N 1D buffer):
+        //   flat[m, k_, n_] = expert_id[m] * (K*N) + k_ * N + n_
+        // Encoded as `Mul(expert_id, Iota(io_const)) + Iota(MIter, K*N)` so the
+        // resulting Gather matches the GLUMoE / gather-experts egglog patterns.
+        let io = k * n;
+        let base = expert_id * io;
+        let within = self.graph.iota(Expression::from('z'), (k, n));
+        let exp_base = base.expand_dim(1, k).expand_dim(2, n);
+        let exp_within = within.expand_dim(0, s);
+        let flat_idx = exp_base + exp_within;
+
+        // Gather → [S, K, N]. Preserves weight's native dtype (bf16 stays bf16).
+        let weight_gathered = weight.gather(flat_idx);
+
+        // Cast for matmul — now on the small gathered slice, not the full weight.
+        let input_f = input.cast(DType::F32);
+        let weight_f = weight_gathered.cast(DType::F32);
+
+        // Per-token matmul: [S, 1, K] @ [S, K, N] → [S, 1, N] → [S, N].
+        let result = input_f.unsqueeze(1).matmul(weight_f).squeeze(1);
+
+        Ok(result.cast(input.dtype))
     }
 
     pub(crate) fn translate_where(&mut self, node: &Node) -> Result<GraphTensor> {
         let cond = self.get_input_tensor(node, 0)?;
         let x = self.get_input_tensor(node, 1)?;
         let y = self.get_input_tensor(node, 2)?;
+        // Ensure x and y have the same dtype
+        let (x, y) = ensure_same_dtype(x, y);
         // Broadcast all three tensors to a common shape first
         let (cond_b, x_b) = broadcast_binary(cond, x);
         let (cond_bc, y_b) = broadcast_binary(cond_b, y);
         let (x_bc, y_bc) = broadcast_binary(x_b, y_b);
         let c = cond_bc.cast(DType::F32);
+        let x_f = x_bc.cast(DType::F32);
+        let y_f = y_bc.cast(DType::F32);
         let one = self.graph.constant_float(1.0).expand_rhs(c.shape);
-        Ok(c * x_bc + (one - c) * y_bc)
+        Ok(c * x_f + (one - c) * y_f)
     }
 
     pub(crate) fn translate_where_scalar_other(&mut self, node: &Node) -> Result<GraphTensor> {
-        let cond = self.get_input_tensor(node, 0)?;
-        let x = self.get_input_tensor(node, 1)?;
-        let other_val = self.get_float_arg(node, 2)? as f32;
+        let cond = self.get_input_tensor(node, WHERE_COND_ARG)?;
+        let x = self.get_input_tensor(node, WHERE_X_ARG)?;
+        let other_val = self.get_float_arg(node, WHERE_OTHER_ARG)? as f32;
         // Broadcast cond and x to a common shape
         let (cond_b, x_b) = broadcast_binary(cond, x);
         let c = cond_b.cast(DType::F32);
@@ -85,33 +294,6 @@ impl<'a> Translator<'a> {
         Ok(c * x_b + (one - c) * other)
     }
 
-    pub(crate) fn translate_diff(&mut self, node: &Node) -> Result<GraphTensor> {
-        let input = self.get_input_tensor(node, 0)?;
-        let dim = if node.inputs.len() > 2 {
-            self.get_int_arg(node, 2).unwrap_or(-1)
-        } else {
-            -1
-        };
-        let dim = normalize_dim(dim, input.shape.len());
-
-        let prepend = if node.inputs.len() > 3 {
-            self.get_input_tensor(node, 3).ok()
-        } else {
-            None
-        };
-
-        let x = if let Some(prep) = prepend {
-            prep.concat_along(input, dim)
-        } else {
-            input
-        };
-
-        let dim_size = x.shape.dims[dim];
-        let front = x.slice_along(Expression::from(1)..dim_size, dim);
-        let back = x.slice_along(Expression::from(0)..dim_size - 1, dim);
-        Ok(front - back)
-    }
-
     pub(crate) fn translate_tril(&mut self, node: &Node) -> Result<GraphTensor> {
         self.translate_triangular(node, false)
     }
@@ -121,9 +303,9 @@ impl<'a> Translator<'a> {
     }
 
     fn translate_triangular(&mut self, node: &Node, upper: bool) -> Result<GraphTensor> {
-        let a = self.get_input_tensor(node, 0)?;
-        let diagonal = if node.inputs.len() > 1 {
-            self.get_int_arg(node, 1).unwrap_or(0) as i32
+        let a = self.get_input_tensor(node, TRIANGULAR_INPUT_ARG)?;
+        let diagonal = if node.inputs.len() > TRIANGULAR_DIAGONAL_ARG {
+            self.get_int_arg(node, TRIANGULAR_DIAGONAL_ARG).unwrap_or(0) as i32
         } else {
             0
         };
@@ -154,10 +336,10 @@ impl<'a> Translator<'a> {
     }
 
     pub(crate) fn translate_topk(&mut self, node: &Node) -> Result<()> {
-        let a = self.get_input_tensor(node, 0)?;
-        let k = self.get_int_arg(node, 1)? as usize;
-        let dim = if node.inputs.len() > 2 {
-            self.get_int_arg(node, 2).unwrap_or(-1)
+        let a = self.get_input_tensor(node, TOPK_INPUT_ARG)?;
+        let k = self.get_int_arg(node, TOPK_K_ARG)? as usize;
+        let dim = if node.inputs.len() > TOPK_DIM_ARG {
+            self.get_int_arg(node, TOPK_DIM_ARG).unwrap_or(-1)
         } else {
             -1
         };
@@ -177,13 +359,10 @@ impl<'a> Translator<'a> {
                 None
             };
 
-        // Use full argsort then slice, rather than topk_indexes/topk_values directly.
-        // This avoids a CUDA gather kernel bug when data and index shapes differ
-        // along the gather axis (topk_indexes returns a sliced tensor).
-        let full_argsort = a.argsort(dim, true);
+        // Build top-k outputs from a full stable argsort, then slice to k.
+        let full_argsort = a.stable_argsort(dim, true);
 
-        // Only build each branch when its output is consumed.
-        // Dead nodes in the graph can confuse the CUDA optimizer.
+        // Only build the outputs that are consumed.
         if let Some(val_name) = values_name
             && !val_name.is_empty()
         {
@@ -191,8 +370,7 @@ impl<'a> Translator<'a> {
             self.tensors.insert(val_name, values);
         }
         if let Some(idx_name) = indices_name {
-            // Materialize Int indices as F32 with `* 1.0` to force a contiguous copy.
-            // Without this, CUDA can't correctly read the sliced Int view.
+            // Materialize the sliced indices through a copy before storing them.
             let indices = full_argsort.slice_along(..k, dim) * 1.0;
             self.tensors.insert(idx_name, indices);
         }
@@ -200,19 +378,49 @@ impl<'a> Translator<'a> {
         Ok(())
     }
 
-    pub(crate) fn translate_one_hot(&mut self, node: &Node) -> Result<GraphTensor> {
-        let a = self.get_input_tensor(node, 0)?;
-        let num_classes = self.get_int_arg(node, 1)? as usize;
-        // one_hot: output[..., i] = 1 if input[...] == i else 0
-        let a_int = a.cast(DType::Int);
-        let classes = self.graph.arange(num_classes);
-        // Expand a to [..., 1] and classes to [..., num_classes]
-        let a_expanded = a_int.expand_dim(a.shape.len(), num_classes);
-        let mut classes_expanded = classes;
-        for d in a.shape.dims.iter().rev() {
-            classes_expanded = classes_expanded.expand_dim(0, *d);
+    pub(crate) fn translate_sort(&mut self, node: &Node) -> Result<()> {
+        let a = self.get_input_tensor(node, SORT_INPUT_ARG)?;
+        let dim = if node.inputs.len() > SORT_DIM_ARG {
+            self.get_int_arg(node, SORT_DIM_ARG).unwrap_or(-1)
+        } else {
+            -1
+        };
+        let descending = if node.inputs.len() > SORT_DESCENDING_ARG {
+            self.get_bool_arg(node, SORT_DESCENDING_ARG)
+                .unwrap_or(false)
+        } else {
+            false
+        };
+        let dim = normalize_dim(dim, a.shape.len());
+
+        // Determine output names (sort returns (values, indices))
+        let values_name = node
+            .outputs
+            .first()
+            .and_then(|o| o.as_tensor.as_ref().map(|t| t.name.clone()));
+        let indices_name =
+            if let Some(ts) = node.outputs.first().and_then(|o| o.as_tensors.as_ref()) {
+                ts.get(1).map(|t| t.name.clone())
+            } else if node.outputs.len() > 1 {
+                node.outputs[1].as_tensor.as_ref().map(|t| t.name.clone())
+            } else {
+                None
+            };
+
+        let full_argsort = a.stable_argsort(dim, descending);
+
+        if let Some(val_name) = values_name
+            && !val_name.is_empty()
+        {
+            let values = a.gather_elements(full_argsort, dim);
+            self.tensors.insert(val_name, values);
         }
-        Ok(a_expanded.eq(classes_expanded).cast(DType::Int))
+        if let Some(idx_name) = indices_name {
+            let indices = full_argsort * 1.0;
+            self.tensors.insert(idx_name, indices);
+        }
+
+        Ok(())
     }
 
     pub(crate) fn translate_wrap_set_grad(&mut self, node: &Node) -> Result<()> {

crates/luminal_python/rust/src/translator/unary.rs

@@ -6,7 +6,38 @@ use crate::pt2_util::{broadcast_binary, torch_dtype_int_to_luminal};
 
 use super::Translator;
 
+const ARGSORT_INPUT_ARG: usize = 0;
+const ARGSORT_DIM_ARG: usize = 1;
+const ARGSORT_DESCENDING_ARG: usize = 2;
+
+const MASKED_FILL_INPUT_ARG: usize = 0;
+const MASKED_FILL_MASK_ARG: usize = 1;
+const MASKED_FILL_VALUE_ARG: usize = 2;
+
+const FLOOR_DIVIDE_INPUT_ARG: usize = 0;
+const FLOOR_DIVIDE_OTHER_ARG: usize = 1;
+
+const DIV_MODE_INPUT_ARG: usize = 0;
+const DIV_MODE_OTHER_ARG: usize = 1;
+
 impl<'a> Translator<'a> {
+    pub(crate) fn translate_argsort(&mut self, node: &Node) -> Result<GraphTensor> {
+        let a = self.get_input_tensor(node, ARGSORT_INPUT_ARG)?;
+        let dim = if node.inputs.len() > ARGSORT_DIM_ARG {
+            self.get_int_arg(node, ARGSORT_DIM_ARG).unwrap_or(-1)
+        } else {
+            -1
+        };
+        let descending = if node.inputs.len() > ARGSORT_DESCENDING_ARG {
+            self.get_bool_arg(node, ARGSORT_DESCENDING_ARG)
+                .unwrap_or(false)
+        } else {
+            false
+        };
+        let dim = crate::pt2_util::normalize_dim(dim, a.shape.len());
+        Ok(a.stable_argsort(dim, descending))
+    }
+
     pub(crate) fn translate_unary_op(
         &mut self,
         node: &Node,
@@ -17,43 +48,17 @@ impl<'a> Translator<'a> {
     }
 
     pub(crate) fn translate_to_copy(&mut self, node: &Node) -> Result<GraphTensor> {
-        let a = self.get_input_tensor(node, 0)?;
-        for input in &node.inputs {
-            if input.name == "dtype"
-                && let Some(dtype_int) = input.arg.as_int()
-            {
-                let dtype = torch_dtype_int_to_luminal(dtype_int as u32);
-                return Ok(a.cast(dtype));
-            }
-        }
-        Ok(a)
-    }
-
-    pub(crate) fn translate_to_dtype(&mut self, node: &Node) -> Result<GraphTensor> {
-        let a = self.get_input_tensor(node, 0)?;
-        if let Some(dtype_int) = node.inputs.get(1).and_then(|i| i.arg.as_scalar_type()) {
-            let dtype = torch_dtype_int_to_luminal(dtype_int);
-            Ok(a.cast(dtype))
-        } else if let Some(dtype_int) = node.inputs.get(1).and_then(|i| i.arg.as_int()) {
-            let dtype = torch_dtype_int_to_luminal(dtype_int as u32);
-            Ok(a.cast(dtype))
-        } else {
-            Ok(a)
-        }
-    }
-
-    pub(crate) fn translate_to_dtype_layout(&mut self, node: &Node) -> Result<GraphTensor> {
         let a = self.get_input_tensor(node, 0)?;
         for input in &node.inputs {
             if input.name == "dtype" {
-                if let Some(dtype_int) = input.arg.as_scalar_type() {
-                    let dtype = torch_dtype_int_to_luminal(dtype_int);
-                    return Ok(a.cast(dtype));
-                }
                 if let Some(dtype_int) = input.arg.as_int() {
                     let dtype = torch_dtype_int_to_luminal(dtype_int as u32);
                     return Ok(a.cast(dtype));
                 }
+                if let Some(dtype_int) = input.arg.as_scalar_type() {
+                    let dtype = torch_dtype_int_to_luminal(dtype_int);
+                    return Ok(a.cast(dtype));
+                }
             }
         }
         Ok(a)
@@ -90,6 +95,155 @@ impl<'a> Translator<'a> {
         Ok(result)
     }
 
+    pub(crate) fn translate_sign(&mut self, node: &Node) -> Result<GraphTensor> {
+        let a = self.get_input_tensor(node, 0)?;
+        let zero = self
+            .graph
+            .constant_float(0.0)
+            .cast(a.dtype)
+            .expand_rhs(a.shape);
+        let pos = a.gt(zero).cast(DType::Int);
+        let neg = a.lt(zero).cast(DType::Int);
+        let signed = pos - neg;
+        Ok(if a.dtype == DType::Int {
+            signed
+        } else {
+            signed.cast(a.dtype)
+        })
+    }
+
+    pub(crate) fn translate_bitwise_not(&mut self, node: &Node) -> Result<GraphTensor> {
+        let a = self.get_input_tensor(node, 0)?;
+        Ok(match a.dtype {
+            DType::Bool => {
+                let one = self
+                    .graph
+                    .constant_float(1.0)
+                    .cast(DType::Int)
+                    .expand_rhs(a.shape);
+                (one - a.cast(DType::Int)).cast(DType::Bool)
+            }
+            DType::Int => (a + 1) * -1.0,
+            other => {
+                anyhow::bail!("bitwise_not only supports Bool/Int routing tensors, got {other:?}")
+            }
+        })
+    }
+
+    pub(crate) fn translate_masked_fill_scalar(&mut self, node: &Node) -> Result<GraphTensor> {
+        let input = self.get_input_tensor(node, MASKED_FILL_INPUT_ARG)?;
+        let mask = self.get_input_tensor(node, MASKED_FILL_MASK_ARG)?;
+        let fill = self.get_float_arg(node, MASKED_FILL_VALUE_ARG)? as f32;
+        let (input, mask) = broadcast_binary(input, mask);
+        let work_dtype = if input.dtype == DType::Bool {
+            DType::Int
+        } else {
+            input.dtype
+        };
+        let input_work = if input.dtype == DType::Bool {
+            input.cast(DType::Int)
+        } else {
+            input
+        };
+        let mask_work = mask.cast(work_dtype);
+        let fill_work = self
+            .graph
+            .constant_float(fill)
+            .cast(work_dtype)
+            .expand_rhs(input_work.shape);
+        let one = self
+            .graph
+            .constant_float(1.0)
+            .cast(work_dtype)
+            .expand_rhs(input_work.shape);
+        let result = mask_work * fill_work + (one - mask_work) * input_work;
+        Ok(if input.dtype == DType::Bool {
+            result.cast(DType::Bool)
+        } else {
+            result
+        })
+    }
+
+    pub(crate) fn translate_floor_divide(&mut self, node: &Node) -> Result<GraphTensor> {
+        let a = self.get_input_tensor(node, FLOOR_DIVIDE_INPUT_ARG)?;
+        let b = if let Some(name) = node
+            .inputs
+            .get(FLOOR_DIVIDE_OTHER_ARG)
+            .and_then(|i| i.arg.as_tensor_name())
+        {
+            self.get_tensor(name)?
+        } else {
+            let scalar = self.get_float_arg(node, FLOOR_DIVIDE_OTHER_ARG)? as f32;
+            self.graph
+                .constant_float(scalar)
+                .cast(a.dtype)
+                .expand_rhs(a.shape)
+        };
+        let (a, b) = crate::pt2_util::ensure_same_dtype(a, b);
+        let (a, b) = broadcast_binary(a, b);
+        let quotient = a.cast(DType::F32) / b.cast(DType::F32);
+        let trunc = quotient.cast(DType::Int).cast(DType::F32);
+        let adjust = quotient.lt(trunc).cast(DType::F32);
+        let floored = trunc - adjust;
+        Ok(if a.dtype == DType::Int {
+            floored.cast(DType::Int)
+        } else {
+            floored.cast(a.dtype)
+        })
+    }
+
+    pub(crate) fn translate_div_tensor_mode(&mut self, node: &Node) -> Result<GraphTensor> {
+        let a = self.get_input_tensor(node, DIV_MODE_INPUT_ARG)?;
+        let b = if let Some(name) = node
+            .inputs
+            .get(DIV_MODE_OTHER_ARG)
+            .and_then(|i| i.arg.as_tensor_name())
+        {
+            self.get_tensor(name)?
+        } else {
+            let scalar = self.get_float_arg(node, DIV_MODE_OTHER_ARG)? as f32;
+            self.graph
+                .constant_float(scalar)
+                .cast(a.dtype)
+                .expand_rhs(a.shape)
+        };
+        let (a, b) = crate::pt2_util::ensure_same_dtype(a, b);
+        let (a, b) = broadcast_binary(a, b);
+
+        // Check rounding_mode kwarg
+        let rounding_mode = node.inputs.iter().find_map(|input| {
+            if input.name == "rounding_mode"
+                && let Argument::Other(val) = &input.arg
+            {
+                return val.as_str().map(|s| s.to_string());
+            }
+            None
+        });
+
+        let quotient = a.cast(DType::F32) / b.cast(DType::F32);
+        match rounding_mode.as_deref() {
+            Some("floor") => {
+                let trunc = quotient.cast(DType::Int).cast(DType::F32);
+                let adjust = quotient.lt(trunc).cast(DType::F32);
+                let floored = trunc - adjust;
+                Ok(if a.dtype == DType::Int {
+                    floored.cast(DType::Int)
+                } else {
+                    floored.cast(a.dtype)
+                })
+            }
+            Some("trunc") => Ok(if a.dtype == DType::Int {
+                quotient.cast(DType::Int)
+            } else {
+                quotient.cast(DType::Int).cast(a.dtype)
+            }),
+            _ => {
+                // No rounding mode — regular division
+                Ok(quotient.cast(a.dtype))
+            }
+        }
+    }
+
     pub(crate) fn translate_clamp(&mut self, node: &Node) -> Result<GraphTensor> {
         let a = self.get_input_tensor(node, 0)?;
         let min_val = if node.inputs.len() > 1 {
@@ -103,13 +257,54 @@ impl<'a> Translator<'a> {
             None
         };
 
-        let mut result = a;
+        // maximum_f32 / minimum_f32 internally use `.lt(F32 scalar)`, which
+        // asserts matching tensor dtypes. Without this, clamp on an Int tensor
+        // (e.g. Qwen3-MoE routes `cache_position.clamp(...)` through here)
+        // panics inside luminal core. Promote to F32 around the bounds check
+        // and cast back at the end.
+        let original_dtype = a.dtype;
+        let needs_promote = original_dtype != DType::F32;
+        let mut result = if needs_promote { a.cast(DType::F32) } else { a };
         if let Some(min) = min_val {
             result = result.maximum_f32(min);
         }
         if let Some(max) = max_val {
             result = result.minimum_f32(max);
         }
+        if needs_promote {
+            result = result.cast(original_dtype);
+        }
         Ok(result)
     }
+
+    /// Compute `erf(a)` via the Abramowitz & Stegun 7.1.28 approximation
+    /// (max error ~1.5e-7). Shared by `aten.erf.default` and the exact
+    /// `aten.gelu.default` (which is `0.5 * x * (1 + erf(x / sqrt(2)))`).
+    ///
+    /// erf(x) = sign(x) * (1 - poly(t) * exp(-x^2))
+    /// where t = 1/(1 + 0.3275911*|x|), poly is degree 5 in Horner form.
+    ///
+    /// Promotes the input to F32 internally (the approximation constants are
+    /// F32 anyway, and luminal's binary ops assert matching dtypes — running
+    /// this on Bf16 input directly trips the assertion at `a.ge(zero)`).
+    /// Restores the original dtype on return.
+    pub(crate) fn erf_approx(&mut self, a: GraphTensor) -> GraphTensor {
+        let orig = a.dtype;
+        let a = if orig == DType::F32 { a } else { a.cast(DType::F32) };
+        let ax = a.abs();
+        let x2 = a * a;
+        let t = (ax * 0.3275911_f32 + 1.0).reciprocal();
+        let poly = t
+            * (t * (t
+                * (t * (t * 1.061_405_4_f32 + (-1.453_152_1_f32)) + 1.421_413_8_f32)
+                + (-0.284_496_72_f32))
+                + 0.254_829_6_f32);
+        let result_abs =
+            self.graph.constant_float(1.0).expand_rhs(a.shape) - poly * (x2 * (-1.0)).exp();
+        // sign(x) = 2*(x >= 0) - 1
+        let zero = self.graph.constant_float(0.0).expand_rhs(a.shape);
+        let sign = a.ge(zero).cast(DType::F32) * 2.0 - 1.0;
+        let result = result_abs * sign;
+        if orig == DType::F32 { result } else { result.cast(orig) }
+    }
 }

crates/luminal_python/rust/src/typed_data.rs

@@ -0,0 +1,352 @@
+//! Dtype-aware buffer type for the luminal_python bridge.
+//!
+//! `TypedData` wraps raw bytes with a `DType` tag, enabling multi-dtype data flow
+//! through the PT2 path without forcing everything to f32.
+
+use luminal::hlir::NativeData;
+use luminal::prelude::tracing::warn;
+use luminal::prelude::*;
+
+/// A dtype-tagged byte buffer. All weight, constant, and input data flows through this type.
+#[derive(Clone, Debug)]
+pub struct TypedData {
+    pub bytes: Vec<u8>,
+    pub dtype: DType,
+}
+
+impl TypedData {
+    /// Wrap raw bytes with a dtype tag. Caller must ensure bytes are correctly formatted.
+    pub fn from_raw(bytes: Vec<u8>, dtype: DType) -> Self {
+        Self { bytes, dtype }
+    }
+
+    /// Number of bytes in the buffer
+    pub fn n_bytes(&self) -> usize {
+        self.bytes.len()
+    }
+
+    /// Number of logical elements (for byte-aligned dtypes)
+    pub fn n_elements(&self) -> usize {
+        let bits = self.dtype.bits();
+        if bits >= 8 {
+            self.bytes.len() / (bits / 8)
+        } else {
+            // sub-byte types: multiple elements per byte
+            self.bytes.len() * (8 / bits)
+        }
+    }
+
+    /// Read element at `idx` as f64 (used by From<TypedData> for NativeData fallback).
+    fn as_f64(&self, idx: usize) -> f64 {
+        match self.dtype {
+            DType::F32 => {
+                let start = idx * 4;
+                f32::from_le_bytes([
+                    self.bytes[start],
+                    self.bytes[start + 1],
+                    self.bytes[start + 2],
+                    self.bytes[start + 3],
+                ]) as f64
+            }
+            DType::F64 => {
+                let start = idx * 8;
+                f64::from_le_bytes([
+                    self.bytes[start],
+                    self.bytes[start + 1],
+                    self.bytes[start + 2],
+                    self.bytes[start + 3],
+                    self.bytes[start + 4],
+                    self.bytes[start + 5],
+                    self.bytes[start + 6],
+                    self.bytes[start + 7],
+                ])
+            }
+            DType::F16 => {
+                let start = idx * 2;
+                half::f16::from_le_bytes([self.bytes[start], self.bytes[start + 1]]).to_f64()
+            }
+            DType::Bf16 => {
+                let start = idx * 2;
+                half::bf16::from_le_bytes([self.bytes[start], self.bytes[start + 1]]).to_f64()
+            }
+            DType::Int => {
+                let start = idx * 4;
+                i32::from_le_bytes([
+                    self.bytes[start],
+                    self.bytes[start + 1],
+                    self.bytes[start + 2],
+                    self.bytes[start + 3],
+                ]) as f64
+            }
+            DType::I8 => self.bytes[idx] as i8 as f64,
+            DType::U8 => self.bytes[idx] as f64,
+            DType::I16 | DType::U16 => {
+                let start = idx * 2;
+                let val = i16::from_le_bytes([self.bytes[start], self.bytes[start + 1]]);
+                if self.dtype == DType::U16 {
+                    val as u16 as f64
+                } else {
+                    val as f64
+                }
+            }
+            DType::Bool => {
+                if self.bytes[idx] != 0 {
+                    1.0
+                } else {
+                    0.0
+                }
+            }
+            _ => panic!("as_f64 not supported for {:?}", self.dtype),
+        }
+    }
+    // -- Constructors from typed Vecs --
+
+    pub fn from_f32_vec(data: Vec<f32>) -> Self {
+        let bytes = unsafe {
+            std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 4).to_vec()
+        };
+        Self {
+            bytes,
+            dtype: DType::F32,
+        }
+    }
+
+    pub fn from_f16_vec(data: Vec<half::f16>) -> Self {
+        let bytes = unsafe {
+            std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 2).to_vec()
+        };
+        Self {
+            bytes,
+            dtype: DType::F16,
+        }
+    }
+
+    pub fn from_bf16_vec(data: Vec<half::bf16>) -> Self {
+        let bytes = unsafe {
+            std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 2).to_vec()
+        };
+        Self {
+            bytes,
+            dtype: DType::Bf16,
+        }
+    }
+
+    pub fn from_i32_vec(data: Vec<i32>) -> Self {
+        let bytes = unsafe {
+            std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 4).to_vec()
+        };
+        Self {
+            bytes,
+            dtype: DType::Int,
+        }
+    }
+
+    pub fn from_bool_vec(data: Vec<bool>) -> Self {
+        let bytes: Vec<u8> = data.iter().map(|&b| b as u8).collect();
+        Self {
+            bytes,
+            dtype: DType::Bool,
+        }
+    }
+
+    /// Convert raw bytes from a PyTorch tensor (identified by PT2 dtype code) to TypedData
+    /// in luminal's native format. Handles widening/narrowing conversions for types where
+    /// PyTorch's byte layout differs from luminal's:
+    /// - i64 → i32, f64 → f32 (luminal has no 64-bit types)
+    /// - i16 → i32, u8 → i32, i8 → i32 (luminal maps all integer types to i32 for PT2)
+    pub fn from_pytorch_bytes(bytes: Vec<u8>, dtype_code: u32) -> Self {
+        match dtype_code {
+            // Types that map directly — preserve raw bytes
+            7 => Self::from_raw(bytes, DType::F32),
+            6 => Self::from_raw(bytes, DType::F16),
+            13 => Self::from_raw(bytes, DType::Bf16),
+            4 => Self::from_raw(bytes, DType::Int), // i32
+            12 => Self::from_raw(bytes, DType::Bool),
+            // i64 → i32 (truncate)
+            5 => {
+                let i32s: Vec<i32> = bytes
+                    .chunks_exact(8)
+                    .map(|b| {
+                        i64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as i32
+                    })
+                    .collect();
+                Self::from_i32_vec(i32s)
+            }
+            // f64 → f32 (downcast)
+            8 => {
+                let f32s: Vec<f32> = bytes
+                    .chunks_exact(8)
+                    .map(|b| {
+                        f64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as f32
+                    })
+                    .collect();
+                Self::from_f32_vec(f32s)
+            }
+            // i16 → i32 (widen)
+            3 => {
+                let i32s: Vec<i32> = bytes
+                    .chunks_exact(2)
+                    .map(|b| i16::from_le_bytes([b[0], b[1]]) as i32)
+                    .collect();
+                Self::from_i32_vec(i32s)
+            }
+            // u8 → i32 (widen)
+            1 => {
+                let i32s: Vec<i32> = bytes.iter().map(|&b| b as i32).collect();
+                Self::from_i32_vec(i32s)
+            }
+            // i8 → i32 (widen, signed)
+            2 => {
+                let i32s: Vec<i32> = bytes.iter().map(|&b| (b as i8) as i32).collect();
+                Self::from_i32_vec(i32s)
+            }
+            // Unknown: best-effort pass-through as f32
+            _ => {
+                warn!("Unrecognized pytorch dtype code {dtype_code}, interpreting as f32");
+                Self::from_raw(bytes, DType::F32)
+            }
+        }
+    }
+
+    /// Create an n-element buffer of "safe" dummy values (1.0 for floats, 1 for ints, true for bool).
+    /// IMPORTANT: Must use 1, NOT 0. Zero inputs cause NaN in many ops (fmod, recip, log, etc.).
+    pub fn ones(n_elements: usize, dtype: DType) -> Self {
+        match dtype {
+            DType::F32 | DType::TF32 => Self::from_f32_vec(vec![1.0f32; n_elements]),
+            DType::F64 => {
+                let data = vec![1.0f64; n_elements];
+                let bytes = unsafe {
+                    std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 8).to_vec()
+                };
+                Self {
+                    bytes,
+                    dtype: DType::F64,
+                }
+            }
+            DType::F16 => Self::from_f16_vec(vec![half::f16::from_f32(1.0); n_elements]),
+            DType::Bf16 => Self::from_bf16_vec(vec![half::bf16::from_f32(1.0); n_elements]),
+            DType::Int => Self::from_i32_vec(vec![1i32; n_elements]),
+            DType::I8 => Self::from_raw(vec![1u8; n_elements], DType::I8),
+            DType::U8 => Self::from_raw(vec![1u8; n_elements], DType::U8),
+            DType::I16 => {
+                let data = vec![1i16; n_elements];
+                let bytes = unsafe {
+                    std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 2).to_vec()
+                };
+                Self {
+                    bytes,
+                    dtype: DType::I16,
+                }
+            }
+            DType::U16 => {
+                let data = vec![1u16; n_elements];
+                let bytes = unsafe {
+                    std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 2).to_vec()
+                };
+                Self {
+                    bytes,
+                    dtype: DType::U16,
+                }
+            }
+            DType::Bool => Self::from_bool_vec(vec![true; n_elements]),
+            _ => panic!("TypedData::ones not supported for {:?}", dtype),
+        }
+    }
+}
+
+/// Convert TypedData to NativeData for the native runtime.
+impl From<TypedData> for NativeData {
+    fn from(td: TypedData) -> Self {
+        match td.dtype {
+            DType::F32 | DType::TF32 => {
+                let data: Vec<f32> = td
+                    .bytes
+                    .chunks_exact(4)
+                    .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
+                    .collect();
+                NativeData::F32(data)
+            }
+            DType::F64 => {
+                // Downcast f64 -> f32 for native runtime (which only has F32 variant for floats > 32-bit)
+                let data: Vec<f32> = td
+                    .bytes
+                    .chunks_exact(8)
+                    .map(|b| {
+                        f64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as f32
+                    })
+                    .collect();
+                NativeData::F32(data)
+            }
+            DType::F16 => {
+                let data: Vec<half::f16> = td
+                    .bytes
+                    .chunks_exact(2)
+                    .map(|b| half::f16::from_le_bytes([b[0], b[1]]))
+                    .collect();
+                NativeData::F16(data)
+            }
+            DType::Bf16 => {
+                let data: Vec<half::bf16> = td
+                    .bytes
+                    .chunks_exact(2)
+                    .map(|b| half::bf16::from_le_bytes([b[0], b[1]]))
+                    .collect();
+                NativeData::Bf16(data)
+            }
+            DType::Int => {
+                let data: Vec<i32> = td
+                    .bytes
+                    .chunks_exact(4)
+                    .map(|b| i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
+                    .collect();
+                NativeData::Int(data)
+            }
+            DType::Bool => {
+                let data: Vec<bool> = td.bytes.iter().map(|&b| b != 0).collect();
+                NativeData::Bool(data)
+            }
+            // Integer types that map to NativeData::Int
+            DType::I8 => {
+                let data: Vec<i32> = td.bytes.iter().map(|&b| b as i8 as i32).collect();
+                NativeData::Int(data)
+            }
+            DType::U8 => {
+                let data: Vec<i32> = td.bytes.iter().map(|&b| b as i32).collect();
+                NativeData::Int(data)
+            }
+            DType::I16 => {
+                let data: Vec<i32> = td
+                    .bytes
+                    .chunks_exact(2)
+                    .map(|b| i16::from_le_bytes([b[0], b[1]]) as i32)
+                    .collect();
+                NativeData::Int(data)
+            }
+            DType::U16 => {
+                let data: Vec<i32> = td
+                    .bytes
+                    .chunks_exact(2)
+                    .map(|b| u16::from_le_bytes([b[0], b[1]]) as i32)
+                    .collect();
+                NativeData::Int(data)
+            }
+            // Sub-byte and F8 types: store as raw f32 for native runtime (best effort)
+            _ => {
+                // For exotic types, the native runtime can't handle them natively.
+                // Store as f32 with element-wise conversion.
+                let data: Vec<f32> = (0..td.n_elements()).map(|i| td.as_f64(i) as f32).collect();
+                NativeData::F32(data)
+            }
+        }
+    }
+}
+
+/// Convert &TypedData to NativeData (clone the bytes).
+impl From<&TypedData> for NativeData {
+    fn from(td: &TypedData) -> Self {
+        td.clone().into()
+    }
+}
+
+// CUDA runtime conversion is implemented via ToCudaInput in runtime.rs
+// (behind the `cuda` feature gate) since it depends on cudarc types.

crates/luminal_python/rust/src/util.rs

@@ -1,477 +0,0 @@
-use std::{collections::HashMap, fs, path::Path};
-
-use luminal::{prelude::GraphTensor, shape::Expression};
-use onnx_protobuf::NodeProto;
-
-/// Maps ONNX dim_param names (e.g. "seq_len") to luminal Expression variable chars ('a'..'w').
-pub type DimParamMap = HashMap<String, char>;
-
-// Given a Value from the Onnx proto return its tensor Shape, if it exists
-// Note: some times pytorch will create tensors with a 0 shape
-// we might want to handle, 0 shape and No shape as seperate ideas
-pub fn get_shape_for_onnx_value(value: &onnx_protobuf::ValueInfoProto) -> Vec<usize> {
-    if let Some(type_proto) = value.type_.as_ref()
-        && let Some(onnx_protobuf::type_proto::Value::TensorType(tensor)) = &type_proto.value
-        && let Some(shape) = tensor.shape.as_ref()
-    {
-        // Scalar (0-dim) tensors have an empty dim list; represent as [1] in luminal
-        if shape.dim.is_empty() {
-            return vec![1];
-        }
-        return shape
-            .dim
-            .iter()
-            .map(|dimension| {
-                if let Some(onnx_protobuf::tensor_shape_proto::dimension::Value::DimValue(v)) =
-                    &dimension.value
-                {
-                    *v as usize
-                } else {
-                    1
-                }
-            })
-            .collect();
-    }
-
-    vec![]
-}
-
-/// Like `get_shape_for_onnx_value`, but returns `Vec<Expression>` with symbolic vars for DimParam dims.
-/// Allocates new variable chars in `dim_param_map` for unseen dim_param names.
-/// `next_char` is updated to the next available char after allocation.
-pub fn get_shape_for_onnx_value_expr(
-    value: &onnx_protobuf::ValueInfoProto,
-    dim_param_map: &mut DimParamMap,
-    next_char: &mut char,
-) -> Vec<Expression> {
-    if let Some(type_proto) = value.type_.as_ref()
-        && let Some(onnx_protobuf::type_proto::Value::TensorType(tensor)) = &type_proto.value
-        && let Some(shape) = tensor.shape.as_ref()
-    {
-        if shape.dim.is_empty() {
-            return vec![Expression::from(1usize)];
-        }
-        return shape
-            .dim
-            .iter()
-            .map(|dimension| match &dimension.value {
-                Some(onnx_protobuf::tensor_shape_proto::dimension::Value::DimValue(v)) => {
-                    Expression::from(*v as usize)
-                }
-                Some(onnx_protobuf::tensor_shape_proto::dimension::Value::DimParam(name)) => {
-                    let ch = *dim_param_map.entry(name.clone()).or_insert_with(|| {
-                        let c = *next_char;
-                        *next_char = (c as u8 + 1) as char;
-                        c
-                    });
-                    Expression::from(ch)
-                }
-                _ => Expression::from(1usize),
-            })
-            .collect();
-    }
-
-    vec![]
-}
-
-/// Compute the broadcast output shape for two tensors using Expressions (numpy rules).
-pub fn compute_broadcast_shape_expr(a: &[Expression], b: &[Expression]) -> Vec<Expression> {
-    let max_rank = a.len().max(b.len());
-    let mut result = Vec::with_capacity(max_rank);
-
-    for i in 0..max_rank {
-        let a_dim = if i < max_rank - a.len() {
-            Expression::from(1usize)
-        } else {
-            a[i - (max_rank - a.len())]
-        };
-        let b_dim = if i < max_rank - b.len() {
-            Expression::from(1usize)
-        } else {
-            b[i - (max_rank - b.len())]
-        };
-
-        // If both are concrete, use max. If one is 1, use the other.
-        // Otherwise, assume they match (same symbolic dim).
-        let dim = match (a_dim.to_usize(), b_dim.to_usize()) {
-            (Some(a_val), Some(b_val)) => Expression::from(a_val.max(b_val)),
-            (Some(1), _) => b_dim,
-            (_, Some(1)) => a_dim,
-            _ => a_dim, // Both symbolic — assume compatible
-        };
-        result.push(dim);
-    }
-    result
-}
-
-/// Broadcast a tensor's shape to match a target Expression shape (numpy-style broadcasting).
-/// Left-pads with size-1 dims, then expands dims that are 1 to match target.
-pub fn broadcast_to_expr(mut tensor: GraphTensor, target_shape: &[Expression]) -> GraphTensor {
-    let src_dims = tensor.dims();
-    let src_len = src_dims.len();
-    let tgt_len = target_shape.len();
-
-    if src_len == tgt_len {
-        tensor.shape.expand(target_shape.to_vec());
-        return tensor;
-    }
-
-    // Left-pad with size-1 dims
-    for _ in 0..(tgt_len - src_len) {
-        tensor = tensor.expand_dim(0, 1);
-    }
-
-    tensor.shape.expand(target_shape.to_vec());
-    tensor
-}
-
-/// Convert inline data from a TensorProto to f32, based on data_type.
-/// Returns None if the tensor has no inline data (e.g. external storage).
-fn convert_inline_data(init: &onnx_protobuf::TensorProto) -> Option<Vec<f32>> {
-    match init.data_type {
-        1 => {
-            // FLOAT
-            if !init.float_data.is_empty() {
-                return Some(init.float_data.clone());
-            }
-            if !init.raw_data.is_empty() {
-                return Some(parse_raw_bytes_as_f32(&init.raw_data, 1));
-            }
-        }
-        7 => {
-            // INT64
-            if !init.int64_data.is_empty() {
-                return Some(init.int64_data.iter().map(|&v| v as f32).collect());
-            }
-            if !init.raw_data.is_empty() {
-                return Some(parse_raw_bytes_as_f32(&init.raw_data, 7));
-            }
-        }
-        6 => {
-            // INT32
-            if !init.int32_data.is_empty() {
-                return Some(init.int32_data.iter().map(|&v| v as f32).collect());
-            }
-            if !init.raw_data.is_empty() {
-                return Some(parse_raw_bytes_as_f32(&init.raw_data, 6));
-            }
-        }
-        9 => {
-            // BOOL
-            if !init.raw_data.is_empty() {
-                return Some(parse_raw_bytes_as_f32(&init.raw_data, 9));
-            }
-            if !init.int32_data.is_empty() {
-                return Some(
-                    init.int32_data
-                        .iter()
-                        .map(|&v| if v != 0 { 1.0 } else { 0.0 })
-                        .collect(),
-                );
-            }
-        }
-        _ => {
-            // Fallback: try float_data or interpret raw_data as F32
-            if !init.float_data.is_empty() {
-                return Some(init.float_data.clone());
-            }
-            if !init.raw_data.is_empty() {
-                return Some(parse_raw_bytes_as_f32(&init.raw_data, 1));
-            }
-        }
-    }
-    None
-}
-
-/// Parse a raw byte slice as f32 values, respecting the ONNX data_type.
-fn parse_raw_bytes_as_f32(bytes: &[u8], data_type: i32) -> Vec<f32> {
-    match data_type {
-        1 => bytes
-            .chunks_exact(4)
-            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
-            .collect(),
-        7 => bytes
-            .chunks_exact(8)
-            .map(|c| i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]]) as f32)
-            .collect(),
-        6 => bytes
-            .chunks_exact(4)
-            .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f32)
-            .collect(),
-        9 => bytes
-            .iter()
-            .map(|&b| if b != 0 { 1.0 } else { 0.0 })
-            .collect(),
-        _ => bytes
-            .chunks_exact(4)
-            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
-            .collect(),
-    }
-}
-
-/// Load float data from a TensorProto, handling inline (float_data/raw_data) and external storage.
-/// Prefer `load_all_tensor_floats` for batch loading (avoids redundant file reads).
-#[allow(dead_code)]
-pub fn load_tensor_floats(init: &onnx_protobuf::TensorProto, model_dir: &Path) -> Option<Vec<f32>> {
-    // Try inline data first
-    if let Some(floats) = convert_inline_data(init) {
-        return Some(floats);
-    }
-    // Try external data (data_location == EXTERNAL = 1)
-    if !init.external_data.is_empty() {
-        let mut location: Option<&str> = None;
-        let mut offset: u64 = 0;
-        let mut length: Option<u64> = None;
-        for entry in &init.external_data {
-            match entry.key.as_str() {
-                "location" => location = Some(&entry.value),
-                "offset" => offset = entry.value.parse().unwrap_or(0),
-                "length" => length = entry.value.parse().ok(),
-                _ => {}
-            }
-        }
-        if let Some(loc) = location {
-            let ext_path = model_dir.join(loc);
-            match fs::read(&ext_path) {
-                Ok(file_data) => {
-                    let start = offset as usize;
-                    let end = match length {
-                        Some(len) => start + len as usize,
-                        None => file_data.len(),
-                    };
-                    if end > file_data.len() {
-                        return None;
-                    }
-                    return Some(parse_raw_bytes_as_f32(
-                        &file_data[start..end],
-                        init.data_type,
-                    ));
-                }
-                Err(_) => {
-                    return None;
-                }
-            }
-        }
-    }
-    None
-}
-
-/// Batch-load float data from multiple TensorProtos, reading each external file only once.
-/// Returns results in the same order as `inits`, with `None` for tensors that couldn't be loaded.
-pub fn load_all_tensor_floats(
-    inits: &[onnx_protobuf::TensorProto],
-    model_dir: &Path,
-) -> Vec<(String, Option<Vec<f32>>)> {
-    let mut results: Vec<(String, Option<Vec<f32>>)> = Vec::with_capacity(inits.len());
-
-    // Pending external data entries: (result_index, offset, length, data_type)
-    // grouped by file location
-    type ExternalEntry = (usize, u64, Option<u64>, i32);
-    let mut external_pending: HashMap<String, Vec<ExternalEntry>> = HashMap::new();
-
-    for (i, init) in inits.iter().enumerate() {
-        // Try inline data first
-        if let Some(floats) = convert_inline_data(init) {
-            results.push((init.name.clone(), Some(floats)));
-            continue;
-        }
-
-        // Check for external data
-        if !init.external_data.is_empty() {
-            let mut location: Option<String> = None;
-            let mut offset: u64 = 0;
-            let mut length: Option<u64> = None;
-            for entry in &init.external_data {
-                match entry.key.as_str() {
-                    "location" => location = Some(entry.value.clone()),
-                    "offset" => offset = entry.value.parse().unwrap_or(0),
-                    "length" => length = entry.value.parse().ok(),
-                    _ => {}
-                }
-            }
-            if let Some(loc) = location {
-                // Push placeholder, will fill in later
-                results.push((init.name.clone(), None));
-                external_pending
-                    .entry(loc)
-                    .or_default()
-                    .push((i, offset, length, init.data_type));
-                continue;
-            }
-        }
-
-        results.push((init.name.clone(), None));
-    }
-
-    // Read each external file once and extract all tensor slices
-    for (loc, entries) in &external_pending {
-        let ext_path = model_dir.join(loc);
-        let file_data = match fs::read(&ext_path) {
-            Ok(data) => data,
-            Err(_) => continue, // results already have None
-        };
-        for &(idx, offset, length, data_type) in entries {
-            let start = offset as usize;
-            let end = match length {
-                Some(len) => start + len as usize,
-                None => file_data.len(),
-            };
-            if end > file_data.len() {
-                continue;
-            }
-            results[idx].1 = Some(parse_raw_bytes_as_f32(&file_data[start..end], data_type));
-        }
-    }
-
-    results
-}
-
-/// Load initializer data as f32 values, handling multiple ONNX data types.
-/// Used to seed known_values with small constant initializers for constant folding.
-pub fn load_initializer_as_f32(init: &onnx_protobuf::TensorProto) -> Option<Vec<f32>> {
-    match init.data_type {
-        1 => {
-            // FLOAT
-            if !init.float_data.is_empty() {
-                Some(init.float_data.clone())
-            } else if !init.raw_data.is_empty() {
-                Some(
-                    init.raw_data
-                        .chunks_exact(4)
-                        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
-                        .collect(),
-                )
-            } else {
-                None
-            }
-        }
-        7 => {
-            // INT64
-            if !init.int64_data.is_empty() {
-                Some(init.int64_data.iter().map(|&v| v as f32).collect())
-            } else if !init.raw_data.is_empty() {
-                Some(
-                    init.raw_data
-                        .chunks_exact(8)
-                        .map(|c| {
-                            i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]])
-                                as f32
-                        })
-                        .collect(),
-                )
-            } else {
-                None
-            }
-        }
-        6 => {
-            // INT32
-            if !init.int32_data.is_empty() {
-                Some(init.int32_data.iter().map(|&v| v as f32).collect())
-            } else if !init.raw_data.is_empty() {
-                Some(
-                    init.raw_data
-                        .chunks_exact(4)
-                        .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f32)
-                        .collect(),
-                )
-            } else {
-                None
-            }
-        }
-        16 => {
-            // BFLOAT16 — 2 bytes per element, upper 16 bits of f32
-            if !init.raw_data.is_empty() {
-                Some(
-                    init.raw_data
-                        .chunks_exact(2)
-                        .map(|c| {
-                            let bits = u16::from_le_bytes([c[0], c[1]]);
-                            f32::from_bits((bits as u32) << 16)
-                        })
-                        .collect(),
-                )
-            } else {
-                None
-            }
-        }
-        9 => {
-            // BOOL — 1 byte per element, 0 → 0.0, non-zero → 1.0
-            if !init.raw_data.is_empty() {
-                Some(
-                    init.raw_data
-                        .iter()
-                        .map(|&b| if b != 0 { 1.0 } else { 0.0 })
-                        .collect(),
-                )
-            } else if !init.int32_data.is_empty() {
-                Some(
-                    init.int32_data
-                        .iter()
-                        .map(|&v| if v != 0 { 1.0 } else { 0.0 })
-                        .collect(),
-                )
-            } else {
-                None
-            }
-        }
-        11 => {
-            // FLOAT64
-            if !init.raw_data.is_empty() {
-                Some(
-                    init.raw_data
-                        .chunks_exact(8)
-                        .map(|c| {
-                            f64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]])
-                                as f32
-                        })
-                        .collect(),
-                )
-            } else {
-                None
-            }
-        }
-        _ => None,
-    }
-}
-
-/// Transpose weight data from [rows, cols] to [cols, rows] row-major layout
-#[cfg(feature = "cuda")]
-pub fn transpose_weight_data(data: &[f32], rows: usize, cols: usize) -> Vec<f32> {
-    let mut transposed = vec![0.0f32; rows * cols];
-    for r in 0..rows {
-        for c in 0..cols {
-            transposed[c * rows + r] = data[r * cols + c];
-        }
-    }
-    transposed
-}
-
-/// Get an integer attribute from a node, with a default value
-pub fn get_int_attr(node: &NodeProto, name: &str, default: i64) -> i64 {
-    for attr in &node.attribute {
-        if attr.name == name {
-            return attr.i;
-        }
-    }
-    default
-}
-
-/// Get a string attribute from a node, with a default value
-pub fn get_str_attr(node: &NodeProto, name: &str, default: &str) -> String {
-    for attr in &node.attribute {
-        if attr.name == name {
-            return String::from_utf8_lossy(&attr.s).into_owned();
-        }
-    }
-    default.to_string()
-}
-
-/// Get a float attribute from a node, with a default value
-pub fn get_float_attr(node: &NodeProto, name: &str, default: f32) -> f32 {
-    for attr in &node.attribute {
-        if attr.name == name {
-            return attr.f;
-        }
-    }
-    default
-}

crates/luminal_python/src/luminal/__init__.py

@@ -1,18 +1,66 @@
 """Luminal Python bindings - PyTorch backend using Luminal."""
 
 # Import Python components
+# Register DynamicCache pytree serialization once at import time
+import torch.export._unlift as _torch_export_unlift
+
+from .cache_utils import _register_cache_serialization
 from .compiled_model import CompiledModel
-from .main import luminal_backend
 
 # Import Rust extension components (built by maturin)
-# These are available directly in the package namespace
-from .luminal import process_onnx, CompiledGraph, compile_pt2
+from .luminal import CompiledGraph, process_pt2
+from .main import luminal_backend, register_backend
+
+_register_cache_serialization()
+
+# ---------------------------------------------------------------------------
+# Suppress torch.export's `_guards_fn` insertion when luminal is on the stack.
+#
+# When `torch._dynamo.config.automatic_dynamic_shapes=True` (the default) and
+# a model is called with shapes that vary across calls, dynamo promotes the
+# changing dim to a SymInt and re-traces. During the re-trace, torch.export's
+# `_unlift_exported_program_lifted_states` (in `torch/export/_unlift.py`)
+# generates a `_guards_fn` submodule whose body closes over `L` — dynamo's
+# locals namespace. When aot_autograd later evaluates the resulting
+# GraphModule via fx.Interpreter, the closure's free `L` reference doesn't
+# resolve and we get
+#       NameError: name 'L' is not defined
+# (gemma3 + StaticCache reproduces this deterministically).
+#
+# torch.export's own opt-out — `_ok_to_generate_guards_fn` — already walks
+# the call stack for filename patterns to suppress guard generation for
+# specific embedders (executorch, modai, on_device_ai, torchao). Add
+# "luminal" to the same suppression set by monkey-patching the function.
+# Net effect: torch.export never inserts `_guards_fn`, so re-tracing
+# succeeds, dynamic-shape compile-once-run-many works, and StaticCache
+# decode loops compile in ~one shot instead of per-token.
+# ---------------------------------------------------------------------------
+
+_orig_ok_to_generate_guards_fn = _torch_export_unlift._ok_to_generate_guards_fn
+
+
+def _luminal_aware_ok_to_generate_guards_fn() -> bool:
+    """Return False whenever luminal is anywhere in the call stack."""
+    import inspect
+
+    frame = inspect.currentframe()
+    try:
+        while frame is not None:
+            if "luminal" in frame.f_code.co_filename:
+                return False
+            frame = frame.f_back
+    finally:
+        del frame  # avoid reference cycle
+    return _orig_ok_to_generate_guards_fn()
+
+
+_torch_export_unlift._ok_to_generate_guards_fn = _luminal_aware_ok_to_generate_guards_fn
 
 # Re-export everything for clean package interface
 __all__ = [
     "CompiledModel",
     "luminal_backend",
-    "process_onnx",
+    "register_backend",
     "CompiledGraph",
-    "compile_pt2",
+    "process_pt2",
 ]

crates/luminal_python/src/luminal/compiled_model.py

@@ -4,21 +4,45 @@ from typing import List
 
 import torch
 
+from .dtype_util import code_to_torch_dtype
+from .dtype_util import torch_dtype_code as _torch_dtype_code
+
 
 class CompiledModel:
     """Wrapper around CompiledGraph that handles PyTorch tensor conversion."""
 
-    def __init__(self, graph_result):
+    def __init__(
+        self, graph_result, weight_refs=None, input_names=None, user_indices=None
+    ):
         """Initialize with a compiled CompiledGraph from Rust.
 
         Args:
-            graph_result: The CompiledGraph from luminal_python.process_onnx() or compile_pt2()
+            graph_result: The CompiledGraph from luminal_python.process_pt2()
+            weight_refs: List of PyTorch tensors to keep alive (prevents GC of shared weights)
+            input_names: Override for user input names. If None, uses graph_result.input_names.
+            user_indices: When torch.compile lifts model parameters into extra args,
+                this tells __call__ which arg positions are actual user inputs.
+                None means all args are user inputs (PT2 path).
         """
         self._graph = graph_result
-        self._input_names = graph_result.input_names
+        self._input_names = input_names or graph_result.input_names
         self._output_names = graph_result.output_names
         self._output_shapes = graph_result.output_shapes
         self._has_dynamic_dims = getattr(graph_result, "has_dynamic_dims", False)
+        self._weight_refs = weight_refs or []
+        self._user_indices = user_indices
+        self._is_gpu = getattr(graph_result, "device_type", "cpu") != "cpu"
+        self._supports_device_ptrs = getattr(
+            graph_result, "supports_device_ptrs", False
+        )
+        # Expected input dtypes from graph (used to convert user inputs)
+        input_dtype_codes = graph_result.input_dtypes
+        self._input_dtypes = [
+            code_to_torch_dtype(input_dtype_codes[i])
+            if i < len(input_dtype_codes)
+            else torch.float32
+            for i in range(len(self._input_names))
+        ]
 
     def set_dim(self, param_name: str, value: int) -> None:
         """Set a dynamic dimension value by its param name."""
@@ -36,49 +60,139 @@ class CompiledModel:
         """Execute the compiled model with PyTorch tensor inputs.
 
         Args:
-            *inputs: PyTorch tensors matching the model's input signature
+            *inputs: PyTorch tensors. When torch.compile lifts model parameters,
+                this includes both weights and user inputs. user_indices filters
+                to just the user inputs.
 
         Returns:
             Tuple of PyTorch tensors containing the model outputs
         """
-        if len(inputs) != len(self._input_names):
-            raise ValueError(
-                f"Expected {len(self._input_names)} inputs, got {len(inputs)}"
-            )
+        # Extract user inputs (torch.compile may pass lifted weights as extra args)
+        if self._user_indices is not None:
+            user_inputs = [inputs[i] for i in self._user_indices]
+        else:
+            if len(inputs) != len(self._input_names):
+                raise ValueError(
+                    f"Expected {len(self._input_names)} inputs, got {len(inputs)}"
+                )
+            user_inputs = inputs
 
         input_device = inputs[0].device if inputs else torch.device("cpu")
 
         # Auto-detect dynamic dims from input shapes
         if self._has_dynamic_dims:
-            input_shapes = [list(t.shape) for t in inputs]
+            input_shapes = [list(t.shape) for t in user_inputs]
             self._graph.auto_set_dims_from_input_shapes(input_shapes)
 
-        # Set input data
-        for name, tensor in zip(self._input_names, inputs):
-            # Convert to contiguous float32 numpy array (move to CPU first for CUDA tensors)
-            arr = tensor.detach().cpu().contiguous().float().numpy()
-            data = arr.flatten().tolist()
-            self._graph.set_input(name, data)
+        # Set user input data via pointer.
+        # Convert to the graph's expected dtype so bytes match the Input node's dtype tag.
+        # For CUDA inputs, keep references alive so the caching allocator doesn't
+        # recycle GPU memory before run() reads the pointers.
+        _input_refs = []
+        for name, tensor, expected_dtype in zip(
+            self._input_names, user_inputs, self._input_dtypes
+        ):
+            if self._supports_device_ptrs and tensor.is_cuda:
+                t = tensor.detach().contiguous().to(expected_dtype)
+                n_bytes = t.numel() * t.element_size()
+                self._graph.set_input_device_ptr(name, t.data_ptr(), n_bytes)
+                _input_refs.append(t)
+            else:
+                t = tensor.detach().cpu().contiguous().to(expected_dtype)
+                n_bytes = t.numel() * t.element_size()
+                dtype_code = _torch_dtype_code(t.dtype)
+                self._graph.set_input_from_ptr(name, t.data_ptr(), n_bytes, dtype_code)
 
-        # Run the graph
-        self._graph.run()
-
-        # Get output shapes — resolve dynamically if needed
+        # Resolve output shapes before run() (needed for pre-allocation).
         if self._has_dynamic_dims:
             output_shapes = self._graph.resolve_output_shapes()
         else:
             output_shapes = self._output_shapes
 
-        # Get outputs and convert back to PyTorch tensors on the same device as inputs
-        outputs = []
-        for name, shape in zip(self._output_names, output_shapes):
-            data = self._graph.get_output(name)
-            tensor = (
-                torch.tensor(data, dtype=torch.float32)
-                .reshape(tuple(shape))
-                .to(input_device)
-            )
-            outputs.append(tensor)
+        output_dtype_codes = self._graph.output_dtypes
+
+        # CUDA zero-copy path: pre-allocate output tensors and register their device
+        # pointers so the final kernel writes directly into PyTorch's buffer.
+        _use_zero_copy = self._supports_device_ptrs
+        output_tensors = []
+        if _use_zero_copy:
+            for i, (name, shape) in enumerate(zip(self._output_names, output_shapes)):
+                out_dtype = (
+                    code_to_torch_dtype(output_dtype_codes[i])
+                    if i < len(output_dtype_codes)
+                    else torch.float32
+                )
+                out = torch.empty(shape, dtype=out_dtype, device=input_device)
+                if out_dtype.is_floating_point:
+                    self._graph.set_output_device_ptr(
+                        name, out.data_ptr(), out.numel() * out.element_size()
+                    )
+                output_tensors.append(out)
+
+        # Run the graph
+        self._graph.run()
+
+        # Collect outputs
+        if _use_zero_copy:
+            outputs = []
+            for i, (name, shape) in enumerate(zip(self._output_names, output_shapes)):
+                out_dtype = (
+                    code_to_torch_dtype(output_dtype_codes[i])
+                    if i < len(output_dtype_codes)
+                    else torch.float32
+                )
+                out = output_tensors[i]
+                if out_dtype.is_floating_point:
+                    if not self._graph.output_is_zero_copy(name):
+                        self._graph.copy_output_to_device_ptr(
+                            name, out.data_ptr(), out.numel() * out.element_size()
+                        )
+                elif out_dtype == torch.int32:
+                    data = self._graph.get_output_i32(name)
+                    out = (
+                        torch.tensor(data, dtype=torch.int32)
+                        .reshape(tuple(shape))
+                        .to(input_device)
+                    )
+                elif out_dtype == torch.bool:
+                    data = self._graph.get_output_bool(name)
+                    out = (
+                        torch.tensor(data, dtype=torch.bool)
+                        .reshape(tuple(shape))
+                        .to(input_device)
+                    )
+                else:
+                    data = self._graph.get_output(name)
+                    out = (
+                        torch.tensor(data, dtype=torch.float32)
+                        .reshape(tuple(shape))
+                        .to(out_dtype)
+                        .to(input_device)
+                    )
+                outputs.append(out)
+        else:
+            # Native path: retrieve as f32, then convert to target dtype if needed.
+            outputs = []
+            for i, (name, shape) in enumerate(zip(self._output_names, output_shapes)):
+                out_dtype = (
+                    code_to_torch_dtype(output_dtype_codes[i])
+                    if i < len(output_dtype_codes)
+                    else torch.float32
+                )
+                if out_dtype == torch.int32:
+                    data = self._graph.get_output_i32(name)
+                    out = torch.tensor(data, dtype=torch.int32).reshape(tuple(shape))
+                elif out_dtype == torch.bool:
+                    data = self._graph.get_output_bool(name)
+                    out = torch.tensor(data, dtype=torch.bool).reshape(tuple(shape))
+                else:
+                    data = self._graph.get_output(name)
+                    out = (
+                        torch.tensor(data, dtype=torch.float32)
+                        .reshape(tuple(shape))
+                        .to(out_dtype)
+                    )
+                out = out.to(input_device)
+                outputs.append(out)
 
-        # Return as a tuple (TorchDynamo expects tuple return from backend callables)
         return tuple(outputs)

crates/luminal_python/src/luminal/dtype_util.py

@@ -0,0 +1,28 @@
+"""Shared dtype utility functions for the luminal Python Bridge"""
+
+import torch
+
+_TORCH_DTYPE_TO_CODE = {
+    torch.uint8: 1,
+    torch.int8: 2,
+    torch.int16: 3,
+    torch.int32: 4,
+    torch.int64: 5,
+    torch.float16: 6,
+    torch.float32: 7,
+    torch.float64: 8,
+    torch.bool: 12,
+    torch.bfloat16: 13,
+}
+
+_CODE_TO_TORCH_DTYPE = {v: k for k, v in _TORCH_DTYPE_TO_CODE.items()}
+
+
+def torch_dtype_code(dtype):
+    """Map torch.dtype to PT2 dtype integer code."""
+    return _TORCH_DTYPE_TO_CODE.get(dtype, 7)  # default to f32
+
+
+def code_to_torch_dtype(code):
+    """Map PT2 dtype integer code to torch.dtype."""
+    return _CODE_TO_TORCH_DTYPE.get(code, torch.float32)

crates/luminal_python/src/luminal/main.py

@@ -1,68 +1,127 @@
-import os
-import tempfile
-
 import torch
 import torch._dynamo
 
-import luminal
+from .dtype_util import torch_dtype_code as _torch_dtype_code
 
-from .cache_utils import _register_cache_serialization
-from .compiled_model import CompiledModel
+
+# ---------------------------------------------------------------------------
+# Shared helpers (used by PT2 path and compiled_model)
+# ---------------------------------------------------------------------------
+
+
+def _detect_factory_capsule(example_inputs):
+    """Pick the best built-in factory capsule based on input device.
+
+    Walks example_inputs for the first Tensor to read .device from. With
+    dynamic=True, dynamo may pass SymInt/SymFloat alongside Tensors and those
+    don't have a .device attribute — falling back to CPU on a SymInt-only call
+    would silently route to the wrong backend, so prefer the first Tensor."""
+    device = torch.device("cpu")
+    for v in example_inputs or ():
+        if isinstance(v, torch.Tensor):
+            device = v.device
+            break
+    if device.type == "cuda":
+        try:
+            from .luminal import _cuda_lite_factory_capsule
+
+            return _cuda_lite_factory_capsule()
+        except ImportError:
+            pass
+    from .luminal import _native_factory_capsule
+
+    return _native_factory_capsule()
+
+
+def _collect_weight_pointers(weights):
+    """Partition weight tensors into CUDA device pointers and CPU host pointers.
+
+    Preserves native dtype — no forced conversion to float32.
+
+    Args:
+        weights: dict of name -> torch.Tensor
+
+    Returns:
+        (keep_alive, device_ptrs, cpu_ptrs) where:
+        - keep_alive: list[Tensor] to prevent GC of shared weight memory
+        - device_ptrs: {name: (device_ptr, n_bytes)}
+        - cpu_ptrs: {name: (host_ptr, n_bytes, dtype_code)}
+    """
+    keep_alive = []
+    device_ptrs = {}
+    cpu_ptrs = {}
+    for name, tensor in weights.items():
+        t = tensor.detach().contiguous()
+        n_bytes = t.numel() * t.element_size()
+        if t.is_cuda:
+            keep_alive.append(t)
+            device_ptrs[name] = (t.data_ptr(), n_bytes)
+        else:
+            t = t.cpu() if t.is_cuda else t
+            keep_alive.append(t)
+            cpu_ptrs[name] = (t.data_ptr(), n_bytes, _torch_dtype_code(t.dtype))
+    return keep_alive, device_ptrs, cpu_ptrs
+
+
+def _load_cpu_weights(compiled_graph, cpu_weights):
+    """Load CPU weight data into a compiled graph after Rust compilation."""
+    for name, (ptr, n_bytes, dtype_code) in cpu_weights.items():
+        compiled_graph.set_weight_from_ptr(name, ptr, n_bytes, dtype_code)
+
+
+# ---------------------------------------------------------------------------
+# Backend registration
+# ---------------------------------------------------------------------------
+
+
+def register_backend(factory_capsule):
+    """Wrap a backend factory PyCapsule into a torch.compile-compatible callable.
+
+    Args:
+        factory_capsule: PyCapsule wrapping a BackendFactory fn pointer.
+
+    Returns:
+        A callable(gm, example_inputs, options=None) suitable for torch.compile.
+    """
+
+    def backend(gm, example_inputs, options=None):
+        return _compile_pt2(gm, example_inputs, factory_capsule, options=options)
+
+    return backend
+
+
+# ---------------------------------------------------------------------------
+# torch.compile backend entry point (auto-detecting)
+# ---------------------------------------------------------------------------
 
 
 def luminal_backend(gm, example_inputs, options=None):
-    """Luminal torch.compile backend.
+    """Auto-detecting torch.compile backend.
 
-    Usage:
-        torch.compile(model, backend=luminal_backend)
-        torch.compile(model, backend=luminal_backend, options={"export_mode": "pt2"})
+    Picks cuda_lite if inputs are on CUDA (and cuda feature is compiled in),
+    native otherwise.
 
-    Options:
-        export_mode: "onnx" (default) or "pt2"
-        opset: ONNX opset version (default 20)
+    For external backends, use register_backend with the backend's factory capsule.
     """
-    options = options or {}
-
-    # Env var override
-    env_mode = os.getenv("LUMINAL_EXPORT_MODE", "").lower()
-    export_mode = (
-        env_mode if env_mode in ("pt2", "onnx") else options.get("export_mode", "onnx")
-    )
-    opset = options.get("opset", 20)
-
-    _register_cache_serialization()
-    device = example_inputs[0].device if example_inputs else torch.device("cpu")
-    backend = "cuda" if device.type == "cuda" else "native"
-
-    if export_mode == "pt2":
-        return _compile_pt2(gm, example_inputs, backend)
-    return _compile_onnx(gm, example_inputs, backend, opset=opset)
+    capsule = _detect_factory_capsule(example_inputs)
+    return _compile_pt2(gm, example_inputs, capsule, options=options)
 
 
-def _compile_onnx(gm, example_inputs, backend, opset=20):
-    """ONNX compilation path."""
-    tmp = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False)
-    tmp_path = tmp.name
-    tmp.close()
-    _ = gm.eval()
-    try:
-        _ = torch.onnx.export(
-            gm,
-            tuple(example_inputs),
-            tmp_path,
-            opset_version=opset,
-            input_names=[f"input_{i}" for i in range(len(example_inputs))],
-        )
-
-        result = luminal.process_onnx(tmp_path, backend)
-    finally:
-        os.unlink(tmp_path)
-    compiled = CompiledModel(result)
-    return compiled
+# ---------------------------------------------------------------------------
+# PT2 compilation path (delegates to pt2 module)
+# ---------------------------------------------------------------------------
 
 
-def _compile_pt2(gm, example_inputs, backend):
+def _compile_pt2(gm, example_inputs, factory_capsule, options=None):
     """PT2/torch.export path — delegates to pt2.pt2_backend."""
     from .pt2 import pt2_backend
 
-    return pt2_backend(gm, example_inputs, backend=backend)
+    search_iterations = None
+    if options is not None:
+        search_iterations = options.get("search_iterations")
+    return pt2_backend(
+        gm,
+        example_inputs,
+        factory=factory_capsule,
+        search_iterations=search_iterations,
+    )

crates/luminal_python/src/luminal/pt2.py

@@ -11,11 +11,87 @@ import shutil
 import tempfile
 
 import torch
-from safetensors.torch import save_file
 
-from .cache_utils import _register_cache_serialization
 from .compiled_model import CompiledModel
-from .luminal import compile_pt2 as _compile_pt2_rust
+from .luminal import process_pt2
+from .main import _collect_weight_pointers, _detect_factory_capsule, _load_cpu_weights
+
+# ---------------------------------------------------------------------------
+# DynamicCache <> pytree registration
+#
+# Without this, torch.export.export raises when handed an HF model that
+# returns CausalLMOutputWithPast(past_key_values=DynamicCache(...)), which
+# is every model with use_cache=True. The registration mirrors the one in
+# transformers.integrations.executorch.register_dynamic_cache_export_support
+# — same dict-based flatten (key_cache / value_cache lists), same replay via
+# cache.update(k, v, idx), and the matching torch.fx._pytree spec for FX
+# graphs. Done at module import so both entry points (pt2_backend via
+# torch.compile and the direct compile() call) get it for free.
+# ---------------------------------------------------------------------------
+
+
+def _get_cache_dict(cache):
+    """Flatten a DynamicCache to a dict of parallel key/value lists."""
+    return {
+        "key_cache": [layer.keys for layer in cache.layers if layer.keys is not None],
+        "value_cache": [
+            layer.values for layer in cache.layers if layer.values is not None
+        ],
+    }
+
+
+def _flatten_dynamic_cache(cache):
+    return torch.utils._pytree._dict_flatten(_get_cache_dict(cache))
+
+
+def _flatten_with_keys_dynamic_cache(cache):
+    return torch.utils._pytree._dict_flatten_with_keys(_get_cache_dict(cache))
+
+
+def _unflatten_dynamic_cache(values, context):
+    from transformers.cache_utils import DynamicCache
+
+    dictionary = torch.utils._pytree._dict_unflatten(values, context)
+    cache = DynamicCache()
+    key_list = dictionary.get("key_cache", [])
+    value_list = dictionary.get("value_cache", [])
+    for idx in range(max(len(key_list), len(value_list))):
+        k = key_list[idx] if idx < len(key_list) else None
+        v = value_list[idx] if idx < len(value_list) else None
+        cache.update(k, v, idx)
+    return cache
+
+
+def _register_cache_serialization():
+    """Register DynamicCache with both torch.utils._pytree and torch.fx._pytree.
+
+    Idempotent: a second call is a no-op. Silently skipped if transformers is
+    not installed.
+    """
+    try:
+        from transformers.cache_utils import DynamicCache
+    except ImportError:
+        return
+
+    if DynamicCache in torch.utils._pytree.SUPPORTED_NODES:
+        return
+
+    torch.utils._pytree.register_pytree_node(
+        DynamicCache,
+        _flatten_dynamic_cache,
+        _unflatten_dynamic_cache,
+        serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
+        flatten_with_keys_fn=_flatten_with_keys_dynamic_cache,
+    )
+    torch.fx._pytree.register_pytree_flatten_spec(
+        DynamicCache,
+        lambda cache, spec: torch.fx._pytree._dict_flatten_spec(
+            _get_cache_dict(cache), spec
+        ),
+    )
+
+
+_register_cache_serialization()
 
 
 # ---------------------------------------------------------------------------
@@ -34,37 +110,112 @@ def _export_kwargs():
     return kwargs
 
 
-def _save_and_compile(ep, backend, search_iterations):
-    """Save ExportedProgram + weights to temp files, compile via Rust, return CompiledModel."""
-    tmpdir = tempfile.mkdtemp(prefix="luminal_")
+def _extract_pt2_constants(pt2_path):
+    """Extract tensor constants from the new flat PT2 format (torch >= 2.6).
+
+    In the new format, inline constants (e.g. ``torch.tensor([1., 2.])``) are
+    stored in ``serialized_constants.pt`` rather than individual ZIP entries.
+    The Rust parser skips them (returns constants_config=None); this function
+    reads them back and returns a cpu_ptrs dict ready for _load_cpu_weights.
+
+    Returns (keep_alive, cpu_ptrs) — keep_alive must stay alive until after
+    _load_cpu_weights returns (set_weight_from_ptr copies the bytes).
+    """
+    import io
+    import zipfile
+
+    from .dtype_util import torch_dtype_code as _torch_dtype_code
+
     try:
-        pt2_path = os.path.join(tmpdir, "model.pt2")
-        weights_path = os.path.join(tmpdir, "weights.safetensors")
+        with zipfile.ZipFile(pt2_path) as z:
+            if "serialized_constants.pt" not in z.namelist():
+                return [], {}
+            data = z.read("serialized_constants.pt")
+    except Exception:
+        return [], {}
 
-        torch.export.save(ep, pt2_path)
+    constants = torch.load(io.BytesIO(data), weights_only=False)
+    if not constants:
+        return [], {}
 
-        state_dict = {k: v.float().clone() for k, v in ep.state_dict.items()}
-        if state_dict:
-            save_file(state_dict, weights_path)
+    keep_alive = []
+    cpu_ptrs = {}
+    for name, tensor in constants.items():
+        t = tensor.detach().cpu().contiguous()
+        keep_alive.append(t)
+        n_bytes = t.numel() * t.element_size()
+        cpu_ptrs[name] = (t.data_ptr(), n_bytes, _torch_dtype_code(t.dtype))
+    return keep_alive, cpu_ptrs
+
+
+def _save_and_compile(ep_or_path, factory, search_iterations, original_weights=None):
+    """Compile a PT2 model via Rust, return CompiledModel.
+
+    Args:
+        ep_or_path: Either an ExportedProgram (will be saved to a temp file) or
+            a path to an already-saved .pt2 file.
+        factory: PyCapsule wrapping the BackendFactory to use.
+        original_weights: Optional dict mapping state_dict key -> original PyTorch tensor.
+            When provided, device pointers are taken from these tensors instead of
+            ep.state_dict (which torch.export may have cloned), enabling true zero-copy
+            sharing with the original model's GPU memory.
+    """
+    owns_tmpdir = not isinstance(ep_or_path, str)
+    tmpdir = tempfile.mkdtemp(prefix="luminal_") if owns_tmpdir else None
+    try:
+        if owns_tmpdir:
+            pt2_path = os.path.join(tmpdir, "model.pt2")
+            torch.export.save(ep_or_path, pt2_path)
+            weight_source = (
+                original_weights if original_weights else ep_or_path.state_dict
+            )
         else:
-            weights_path = ""
+            pt2_path = ep_or_path
+            weight_source = original_weights or {}
 
-        compiled = _compile_pt2_rust(pt2_path, weights_path, backend, search_iterations)
-        return CompiledModel(compiled)
+        # Collect weight pointers for Rust (avoids duplicate GPU buffer allocation)
+        keep_alive, weight_device_ptrs, cpu_weights = _collect_weight_pointers(
+            weight_source
+        )
+
+        # Compile with device pointers — search uses actual weight memory (zero-copy)
+        compiled = process_pt2(
+            pt2_path, "", search_iterations, factory, weight_device_ptrs
+        )
+
+        # Load CPU weights; also load inline tensor constants from the new flat
+        # PT2 format (torch >= 2.6 stores them in serialized_constants.pt).
+        const_keep_alive, const_cpu_weights = _extract_pt2_constants(pt2_path)
+        cpu_weights.update(const_cpu_weights)
+        _load_cpu_weights(compiled, cpu_weights)
+        del const_keep_alive  # bytes were copied by set_weight_from_ptr
+
+        return CompiledModel(compiled, weight_refs=keep_alive)
     finally:
-        shutil.rmtree(tmpdir, ignore_errors=True)
+        if owns_tmpdir and tmpdir:
+            shutil.rmtree(tmpdir, ignore_errors=True)
 
 
 def _reinternalize_lifted_params(gm, example_inputs):
     """Re-internalize lifted params as buffers so torch.export sees them as model state.
 
     torch.compile lifts model parameters out of the module and passes them as
-    extra elements in example_inputs.  The Rust PT2 compiler expects weights in
-    the .pt2 state dict, not as runtime inputs.  This function reverses the
+    extra elements in example_inputs. The Rust PT2 compiler may expect weights in
+    the .pt2 state dict, not as runtime inputs. This function reverses the
     lifting by registering them as buffers and replacing the placeholder nodes
     with get_attr nodes.
 
-    Returns (gm, user_inputs) where user_inputs contains only the real inputs.
+    SymInt/SymFloat/SymBool values in example_inputs are rejected by
+    torch.export.export as user inputs ("Unsupported input type
+    <class 'torch.SymInt'>"). We don't restructure the graph for this — we
+    specialize the *value* to its concrete hint (a plain int/float/bool), which
+    torch.export accepts. The placeholder stays in place; the traced graph
+    proceeds as if dynamo had specialized this dim. Invisible to callers of
+    `torch.compile(..., backend=luminal_backend)`.
+
+    Returns (gm, user_inputs, original_weights) where:
+      - user_inputs contains only real inputs (Tensors and concrete scalars)
+      - original_weights maps buffer name -> original tensor (for zero-copy device pointers)
     """
     buffer_indices = []
     user_indices = []
@@ -80,12 +231,15 @@ def _reinternalize_lifted_params(gm, example_inputs):
                 user_indices.append(placeholder_idx)
             placeholder_idx += 1
 
+    original_weights = {}
     if buffer_nodes:
         for i, node in enumerate(buffer_nodes):
             attr_name = f"_luminal_param_{i}"
-            gm.register_buffer(
-                attr_name, example_inputs[buffer_indices[i]].detach().clone()
-            )
+            # Keep a reference to the original tensor for zero-copy device pointers.
+            # torch.export.export may clone the registered buffer, so we bypass
+            # the EP's state_dict and use the originals directly.
+            original_weights[attr_name] = example_inputs[buffer_indices[i]]
+            gm.register_buffer(attr_name, example_inputs[buffer_indices[i]].detach())
             with gm.graph.inserting_before(node):
                 new_node = gm.graph.create_node("get_attr", attr_name)
                 new_node.meta = node.meta.copy()
@@ -94,12 +248,47 @@ def _reinternalize_lifted_params(gm, example_inputs):
         gm.graph.lint()
         gm.recompile()
 
-    user_inputs = (
+    raw_user_inputs = (
         [example_inputs[i] for i in user_indices]
         if user_indices
         else list(example_inputs)
     )
-    return gm, user_inputs
+    user_inputs = [
+        _specialize_sym_scalar(v) if _is_sym_scalar(v) else v
+        for v in raw_user_inputs
+    ]
+    return gm, user_inputs, original_weights
+
+
+def _is_sym_scalar(val) -> bool:
+    """True for torch SymInt/SymFloat/SymBool — anything torch.export's fakify
+    rejects as a user input. Plain int/float/bool are fine; only the symbolic
+    wrappers need specialization."""
+    if val is None:
+        return False
+    if isinstance(val, torch.Tensor):
+        return False
+    return type(val).__name__ in ("SymInt", "SymFloat", "SymBool") or isinstance(
+        val, (torch.SymInt, torch.SymFloat, torch.SymBool)
+    )
+
+
+def _specialize_sym_scalar(val):
+    """Resolve a SymInt/SymFloat/SymBool to its concrete hint. Falls back to
+    str(val) -> primitive parse if the SymNode hint is missing."""
+    try:
+        if isinstance(val, torch.SymBool):
+            return bool(val)
+        if isinstance(val, torch.SymFloat):
+            return float(val)
+        return int(val)
+    except Exception:
+        # SymNodes without a hint — try parsing the str form as a last resort.
+        s = str(val)
+        try:
+            return int(s)
+        except ValueError:
+            return float(s)
 
 
 # ---------------------------------------------------------------------------
@@ -111,7 +300,7 @@ def compile(
     model,
     example_input,
     search_iterations=25,
-    backend=None,
+    factory=None,
     export_kwargs=None,
     dynamic_dim=None,
 ):
@@ -121,22 +310,18 @@ def compile(
         model: A PyTorch nn.Module.
         example_input: Example input tensor(s) for tracing.
         search_iterations: Number of optimization search iterations.
-        backend: "cpu" or "cuda". Auto-detected if None.
+        factory: PyCapsule wrapping a BackendFactory. Auto-detected if None.
         export_kwargs: Extra kwargs passed to torch.export.export.
         dynamic_dim: Which input dimension to make dynamic.
 
     Returns:
         A CompiledModel callable.
     """
-    _register_cache_serialization()
-
     if dynamic_dim is None:
         dynamic_dim = "auto"
 
-    if backend is None:
-        backend = os.environ.get("LUMINAL_BACKEND", None)
-        if backend is None:
-            backend = "cuda" if torch.cuda.is_available() else "cpu"
+    if factory is None:
+        factory = _detect_factory_capsule([example_input])
 
     kwargs = export_kwargs or {}
     extra = _export_kwargs()
@@ -170,6 +355,7 @@ def compile(
                     dynamic_shapes=dynamic_shapes,
                     **extra,
                 )
+                ep = ep.run_decompositions()
                 break
             except Exception:
                 continue
@@ -182,20 +368,97 @@ def compile(
             dynamic_shapes=None,
             **extra,
         )
+        ep = ep.run_decompositions()
 
-    return _save_and_compile(ep, backend, search_iterations)
+    return _save_and_compile(ep, factory, search_iterations)
 
 
-def pt2_backend(gm, example_inputs, backend=None):
+def pt2_backend(gm, example_inputs, factory=None, search_iterations=None):
     """torch.compile backend using PT2 pipeline.
 
-    Usage: torch.compile(model, backend=luminal.pt2.pt2_backend)
+    Usage: torch.compile(model, backend=luminal.register_backend(capsule))
     """
-    _register_cache_serialization()
-    if backend is None:
-        device = example_inputs[0].device if example_inputs else torch.device("cpu")
-        backend = "cuda" if device.type == "cuda" else "cpu"
+    import gc
+
+    if factory is None:
+        factory = _detect_factory_capsule(example_inputs)
+    if search_iterations is None:
+        search_iterations = 10
+
     gm = gm.eval()
-    gm, user_inputs = _reinternalize_lifted_params(gm, example_inputs)
+    gm, user_inputs, original_weights = _reinternalize_lifted_params(gm, example_inputs)
+
     ep = torch.export.export(gm, tuple(user_inputs), **_export_kwargs())
-    return _save_and_compile(ep, backend, 10)
+    ep = ep.run_decompositions()
+
+    # Detect USER_INPUT_MUTATION outputs (e.g., in-place KV cache updates).
+    # These must be written back to the original input tensors after each call.
+    # Only USER_OUTPUT results are returned to the torch.compile caller.
+    try:
+        from torch.export.graph_signature import OutputKind
+
+        mutation_mappings = []  # list of (compiled_output_idx, user_input_idx)
+        user_output_indices = []
+        for i, spec in enumerate(ep.graph_signature.output_specs):
+            if spec.kind == OutputKind.USER_INPUT_MUTATION:
+                # target is 'args_N' — index into user_inputs
+                try:
+                    arg_idx = int(spec.target.split("_")[1])
+                    mutation_mappings.append((i, arg_idx))
+                except (ValueError, IndexError):
+                    user_output_indices.append(i)
+            else:
+                user_output_indices.append(i)
+    except ImportError:
+        mutation_mappings = []
+        user_output_indices = None  # unknown; return all outputs
+
+    # When using shared memory (original_weights), strip large weight buffers from
+    # the EP before saving.  The Rust side uses device pointers for these weights,
+    # not the .pt2 file data, so serializing them is pure IO waste (~32 GB for 8B
+    # models).  Replacing with tiny CPU scalars shrinks the .pt2 to < 1 MB.
+    if original_weights:
+        for key in list(ep._state_dict.keys()):
+            if key in original_weights:
+                orig = ep._state_dict[key]
+                ep._state_dict[key] = torch.zeros(1, dtype=orig.dtype, device="cpu")
+                del orig
+
+    # Save the exported program to disk, then free it and the traced graph module
+    # BEFORE Rust compilation. torch.export clones the state_dict internally, so
+    # holding ep alive during compilation would double the weight memory on GPU.
+    tmpdir = tempfile.mkdtemp(prefix="luminal_")
+    pt2_path = os.path.join(tmpdir, "model.pt2")
+    torch.export.save(ep, pt2_path)
+
+    del ep, gm
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    try:
+        result = _save_and_compile(
+            pt2_path, factory, search_iterations, original_weights=original_weights
+        )
+    finally:
+        shutil.rmtree(tmpdir, ignore_errors=True)
+
+    # Wrap the compiled model to handle USER_INPUT_MUTATION: write updated tensors
+    # back into the original input buffers and return only USER_OUTPUT tensors.
+    if mutation_mappings:
+        _compiled = result
+        _mut = mutation_mappings
+        _usr = user_output_indices
+
+        def _mutation_wrapper(*inputs):
+            outputs = _compiled(*inputs)
+            for out_idx, inp_idx in _mut:
+                if inp_idx < len(inputs) and out_idx < len(outputs):
+                    inputs[inp_idx].copy_(outputs[out_idx])
+            if _usr is not None:
+                return tuple(outputs[i] for i in _usr if i < len(outputs))
+            return outputs
+
+        return _mutation_wrapper
+
+    return result

crates/luminal_python/tests/_test_kimi_k25.py

@@ -1,194 +0,0 @@
-"""Kimi-K2.5 / DeepseekV3 model integration tests.
-
-Tests the DeepseekV3 text backbone (MoE + MLA attention with LoRA-compressed KV,
-SwiGLU, YaRN RoPE) through the PyTorch -> ONNX -> luminal pipeline.
-
-The model code requires trust_remote_code=True and uses custom HF modules from
-moonshotai/Kimi-K2.5. Since torch.compile cannot trace the MoE routing (it uses
-.numpy() and tensor indexing incompatible with dynamo), tests use manual ONNX
-export + onnxsim simplification + luminal.process_onnx.
-"""
-
-import os
-import tempfile
-import warnings
-
-import onnx
-import onnxsim
-import pytest
-import torch
-
-warnings.filterwarnings("ignore")
-
-
-def _get_deepseek_v3_classes():
-    """Import DeepseekV3Config and DeepseekV3ForCausalLM from the Kimi-K2.5 HF repo."""
-    import importlib
-
-    from transformers import AutoConfig
-
-    config = AutoConfig.from_pretrained("moonshotai/Kimi-K2.5", trust_remote_code=True)
-    tc = config.text_config
-    DeepseekV3Config = type(tc)
-    pkg = DeepseekV3Config.__module__.rsplit(".", 1)[0]
-    modeling_mod = importlib.import_module(f"{pkg}.modeling_deepseek")
-    return DeepseekV3Config, modeling_mod.DeepseekV3ForCausalLM
-
-
-def _make_deepseek_v3_config(
-    DeepseekV3Config,
-    hidden_size: int = 64,
-    num_attention_heads: int = 4,
-    num_key_value_heads: int = 4,
-    num_hidden_layers: int = 1,
-    intermediate_size: int = 128,
-    vocab_size: int = 256,
-    kv_lora_rank: int = 16,
-    q_lora_rank: int = 32,
-    qk_nope_head_dim: int = 8,
-    qk_rope_head_dim: int = 8,
-    v_head_dim: int = 8,
-    n_routed_experts: int = 4,
-    num_experts_per_tok: int = 2,
-    n_shared_experts: int = 1,
-    moe_intermediate_size: int = 32,
-    first_k_dense_replace: int = 1,
-):
-    """Create a small DeepseekV3Config for testing."""
-    config = DeepseekV3Config(
-        hidden_size=hidden_size,
-        num_attention_heads=num_attention_heads,
-        num_key_value_heads=num_key_value_heads,
-        num_hidden_layers=num_hidden_layers,
-        intermediate_size=intermediate_size,
-        vocab_size=vocab_size,
-        max_position_embeddings=128,
-        kv_lora_rank=kv_lora_rank,
-        q_lora_rank=q_lora_rank,
-        qk_nope_head_dim=qk_nope_head_dim,
-        qk_rope_head_dim=qk_rope_head_dim,
-        v_head_dim=v_head_dim,
-        n_routed_experts=n_routed_experts,
-        num_experts_per_tok=num_experts_per_tok,
-        n_shared_experts=n_shared_experts,
-        moe_intermediate_size=moe_intermediate_size,
-        first_k_dense_replace=first_k_dense_replace,
-        use_cache=False,
-        n_group=1,
-        topk_group=1,
-        topk_method="noaux_tc",
-        scoring_func="sigmoid",
-        rope_scaling={
-            "type": "yarn",
-            "rope_type": "yarn",
-            "factor": 4.0,
-            "original_max_position_embeddings": 32,
-            "beta_fast": 32.0,
-            "beta_slow": 1.0,
-            "mscale": 1.0,
-            "mscale_all_dim": 1.0,
-            "rope_theta": 10000.0,
-        },
-        rope_theta=10000.0,
-    )
-    config._attn_implementation = "eager"
-    return config
-
-
-def _export_and_simplify(model, input_ids):
-    """Export model to ONNX and simplify with onnxsim to constant-fold shape chains."""
-    tmp = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False)
-    tmp_path = tmp.name
-    tmp.close()
-    try:
-        torch.onnx.export(
-            model,
-            (input_ids,),
-            tmp_path,
-            opset_version=20,
-            input_names=["input_ids"],
-            output_names=["logits"],
-            dynamo=False,
-        )
-        m = onnx.load(tmp_path)
-        m_sim, check = onnxsim.simplify(m)
-        assert check, "onnxsim simplification failed"
-        onnx.save(m_sim, tmp_path)
-        return tmp_path
-    except Exception:
-        os.unlink(tmp_path)
-        raise
-
-
-def _run_deepseek_v3_test(config, DeepseekV3ForCausalLM, backend: str, atol: float):
-    """Export DeepseekV3 to ONNX, simplify, run through luminal, compare."""
-    import luminal
-
-    model = DeepseekV3ForCausalLM(config).eval()
-    input_ids = torch.tensor([[1, 2, 3, 4]])
-
-    onnx_path = _export_and_simplify(model, input_ids)
-    try:
-        graph = luminal.process_onnx(onnx_path, backend)
-        graph.set_input("input_ids", [1.0, 2.0, 3.0, 4.0])
-        graph.run()
-        logits_data = graph.get_output("logits")
-        logits = torch.tensor(logits_data, dtype=torch.float32).reshape(
-            1, 4, config.vocab_size
-        )
-    finally:
-        os.unlink(onnx_path)
-
-    with torch.no_grad():
-        ref = model(input_ids)
-
-    assert torch.allclose(logits, ref.logits, atol=atol), (
-        f"max_diff={torch.max(torch.abs(logits - ref.logits)).item():.2e}"
-    )
-
-
-# ========== Tests ==========
-
-
-def test_deepseek_v3_tiny_dense():
-    """Tiny DeepseekV3 with dense MLP (no MoE): 64 hidden, 1 layer, MLA attention."""
-    DeepseekV3Config, DeepseekV3ForCausalLM = _get_deepseek_v3_classes()
-    config = _make_deepseek_v3_config(
-        DeepseekV3Config,
-        first_k_dense_replace=1,  # all layers use dense MLP
-    )
-    backend = os.environ.get("LUMINAL_BACKEND", "native")
-    _run_deepseek_v3_test(config, DeepseekV3ForCausalLM, backend, atol=1e-5)
-
-
-@pytest.mark.xfail(reason="MoE routing uses Int/F32 mixed ops not yet supported")
-def test_deepseek_v3_tiny_moe():
-    """Tiny DeepseekV3 with MoE: 64 hidden, 1 layer, 4 routed experts + 1 shared."""
-    DeepseekV3Config, DeepseekV3ForCausalLM = _get_deepseek_v3_classes()
-    config = _make_deepseek_v3_config(
-        DeepseekV3Config,
-        first_k_dense_replace=0,  # all layers use MoE
-    )
-    backend = os.environ.get("LUMINAL_BACKEND", "native")
-    _run_deepseek_v3_test(config, DeepseekV3ForCausalLM, backend, atol=1e-5)
-
-
-def test_deepseek_v3_small_dense():
-    """Small DeepseekV3 with dense MLP: 256 hidden, 1 layer."""
-    DeepseekV3Config, DeepseekV3ForCausalLM = _get_deepseek_v3_classes()
-    config = _make_deepseek_v3_config(
-        DeepseekV3Config,
-        hidden_size=256,
-        num_attention_heads=8,
-        num_key_value_heads=8,
-        intermediate_size=512,
-        vocab_size=1024,
-        kv_lora_rank=32,
-        q_lora_rank=64,
-        qk_nope_head_dim=16,
-        qk_rope_head_dim=16,
-        v_head_dim=16,
-        first_k_dense_replace=1,
-    )
-    backend = os.environ.get("LUMINAL_BACKEND", "native")
-    _run_deepseek_v3_test(config, DeepseekV3ForCausalLM, backend, atol=1e-4)

crates/luminal_python/tests/_test_qwen3.py

@@ -1,7 +1,7 @@
 """Qwen3-8B HuggingFace model integration tests.
 
 Tests progressively larger HuggingFace Qwen3ForCausalLM configs through the
-PyTorch -> ONNX -> luminal pipeline via torch.compile. Qwen3 shares the same
+PyTorch -> PT2 -> luminal pipeline via torch.compile. Qwen3 shares the same
 architecture family as Llama (GQA, RoPE, SwiGLU MLP, RMSNorm).
 """
 
@@ -10,7 +10,6 @@ import torch._dynamo
 
 from luminal import luminal_backend
 
-
 # ========== HuggingFace Qwen3ForCausalLM Tests ==========
 
 
@@ -56,12 +55,12 @@ def _run_hf_qwen3_test(config, device: torch.device, atol: float):
 def test_hf_qwen3_tiny(device: torch.device):
     """HuggingFace Qwen3ForCausalLM -- tiny (64 hidden, 1 layer, ~70K params)."""
     config = _make_qwen3_config(
-        hidden_size=64,
-        num_attention_heads=4,
-        num_key_value_heads=2,
+        hidden_size=32,
+        num_attention_heads=2,
+        num_key_value_heads=1,
         num_hidden_layers=1,
-        intermediate_size=128,
-        vocab_size=256,
+        intermediate_size=64,
+        vocab_size=128,
     )
     _run_hf_qwen3_test(config, device, atol=1e-5)
 
@@ -161,167 +160,6 @@ def test_hf_qwen3_decode_loop_static(device: torch.device):
         tokens.append(next_token)
 
 
-def test_hf_qwen3_decode_loop_dynamic():
-    """Decode loop with dynamic shapes -- compile once, run with varying seq_len.
-
-    Bypasses torch.compile to use luminal's dynamic dim support directly.
-    Exports ONNX once with dynamic_axes, then calls set_dim/set_input/run/get_output.
-    """
-    import os
-    import tempfile
-
-    from transformers import Qwen3Config, Qwen3ForCausalLM
-
-    import luminal
-
-    config = Qwen3Config(
-        hidden_size=64,
-        num_attention_heads=4,
-        num_key_value_heads=2,
-        num_hidden_layers=1,
-        intermediate_size=128,
-        vocab_size=256,
-        max_position_embeddings=128,
-        use_cache=False,
-        attn_implementation="eager",
-    )
-    model = Qwen3ForCausalLM(config).eval()
-
-    # Export ONNX once with dynamic seq_len
-    dummy = torch.tensor([[1, 2, 3, 4]])
-    tmp = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False)
-    tmp_path = tmp.name
-    tmp.close()
-    try:
-        torch.onnx.export(
-            model,
-            (dummy,),
-            tmp_path,
-            opset_version=20,
-            input_names=["input_ids"],
-            output_names=["logits"],
-            dynamic_axes={"input_ids": {1: "seq_len"}, "logits": {1: "seq_len"}},
-        )
-
-        graph = luminal.process_onnx(tmp_path, "native")
-    finally:
-        os.unlink(tmp_path)
-
-    assert graph.has_dynamic_dims, "Graph should have dynamic dims"
-    assert "seq_len" in graph.dim_params, f"Expected 'seq_len' in {graph.dim_params}"
-
-    tokens = [1, 2, 3, 4]
-    for step in range(3):
-        seq_len = len(tokens)
-        graph.set_dim("seq_len", seq_len)
-
-        # Set input as float (luminal works with f32 internally)
-        graph.set_input("input_ids", [float(t) for t in tokens])
-        graph.run()
-
-        # Get output and reshape using resolved shapes
-        output_shapes = graph.resolve_output_shapes()
-        logits_data = graph.get_output("logits")
-        logits = torch.tensor(logits_data, dtype=torch.float32).reshape(
-            output_shapes[0]
-        )
-
-        # Compare against PyTorch reference
-        input_ids = torch.tensor([tokens])
-        with torch.no_grad():
-            ref = model(input_ids)
-
-        assert torch.allclose(logits, ref.logits, atol=1e-4), (
-            f"step {step}: max_diff={torch.max(torch.abs(logits - ref.logits)).item():.2e}"
-        )
-
-        next_token = ref.logits[0, -1, :].argmax().item()
-        tokens.append(next_token)
-
-
-def test_hf_qwen3_8b_decode_loop_dynamic():
-    """Decode loop with dynamic shapes on real Qwen3-8B -- compile once, run with varying seq_len.
-
-    Full 8B model with pretrained weights, ONNX exported once with dynamic_axes
-    for seq_len, then decoded autoregressively without recompilation.
-    """
-    import os
-    import tempfile
-
-    from transformers import AutoConfig, AutoTokenizer, Qwen3ForCausalLM
-
-    import luminal
-
-    backend = os.environ.get("LUMINAL_BACKEND", "cuda")
-
-    config = AutoConfig.from_pretrained("Qwen/Qwen3-8B")
-    config.use_cache = False
-    config._attn_implementation = "eager"
-    print("Loaded config")
-    model = Qwen3ForCausalLM.from_pretrained(
-        "Qwen/Qwen3-8B",
-        config=config,
-        torch_dtype=torch.float32,
-    ).eval()
-    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
-    print("Loaded Model")
-
-    # Export ONNX once with dynamic seq_len
-    dummy = torch.tensor([[1, 2, 3, 4]])
-    tmp = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False)
-    tmp_path = tmp.name
-    tmp.close()
-
-    try:
-        torch.onnx.export(
-            model,
-            (dummy,),
-            tmp_path,
-            opset_version=20,
-            input_names=["input_ids"],
-            output_names=["logits"],
-            dynamic_axes={"input_ids": {1: "seq_len"}, "logits": {1: "seq_len"}},
-        )
-        print("Exported onnx")
-        graph = luminal.process_onnx(tmp_path, backend)
-    finally:
-        os.unlink(tmp_path)
-    print("Exported Model")
-    assert graph.has_dynamic_dims, "Graph should have dynamic dims"
-    assert "seq_len" in graph.dim_params, f"Expected 'seq_len' in {graph.dim_params}"
-
-    prompt = "The capital of france is"
-    tokens = tokenizer.encode(prompt)
-    print(f"Prompt: '{prompt}' -> {len(tokens)} tokens: {tokens}")
-
-    num_generate = 3
-    for step in range(num_generate):
-        seq_len = len(tokens)
-        graph.set_dim("seq_len", seq_len)
-
-        graph.set_input("input_ids", [float(t) for t in tokens])
-        graph.run()
-
-        output_shapes = graph.resolve_output_shapes()
-        logits_data = graph.get_output("logits")
-        logits = torch.tensor(logits_data, dtype=torch.float32).reshape(
-            output_shapes[0]
-        )
-
-        # Compare against PyTorch reference
-        input_ids = torch.tensor([tokens])
-        with torch.no_grad():
-            ref = model(input_ids)
-
-        assert torch.allclose(logits, ref.logits, atol=1e-3), (
-            f"step {step}: max_diff={torch.max(torch.abs(logits - ref.logits)).item():.2e}"
-        )
-
-        next_token = ref.logits[0, -1, :].argmax().item()
-        tokens.append(next_token)
-        print(f"Step {step}: '{tokenizer.decode(tokens)}'")
-
-
 def test_hf_qwen3_8b_full(device: torch.device):
     """HuggingFace Qwen3ForCausalLM -- full Qwen3-8B with real pretrained weights.
 

crates/luminal_python/tests/_test_qwen_image.py

@@ -1,426 +0,0 @@
-"""Qwen-Image diffusion model integration tests.
-
-Tests the QwenImageTransformer2DModel (MMDiT denoiser) and AutoencoderKLQwenImage (VAE)
-through the PyTorch -> ONNX -> luminal pipeline.
-
-The transformer uses complex-valued RoPE (torch.view_as_complex) which isn't ONNX-exportable,
-so tests use a wrapper that pre-computes RoPE as real-valued cos/sin and replaces the
-attention processor with a real-valued equivalent.
-
-The VAE uses Conv3d, which is supported via the N-dimensional unfold-based conv parser.
-"""
-
-import os
-import tempfile
-import warnings
-
-import onnx
-import onnxsim
-import pytest
-import torch
-import torch.nn as nn
-
-warnings.filterwarnings("ignore")
-
-
-# ============================================================================
-# Transformer helpers
-# ============================================================================
-
-
-def _apply_rope_real(x, cos, sin):
-    """Apply RoPE using real-valued cos/sin. x: [B, S, H, D], cos/sin: [S, D/2]."""
-    d = x.shape[-1]
-    x1 = x[..., : d // 2]
-    x2 = x[..., d // 2 :]
-    cos = cos.unsqueeze(0).unsqueeze(2)  # [1, S, 1, D/2]
-    sin = sin.unsqueeze(0).unsqueeze(2)
-    rotated_x1 = x1 * cos - x2 * sin
-    rotated_x2 = x2 * cos + x1 * sin
-    return torch.cat([rotated_x1, rotated_x2], dim=-1)
-
-
-class RealRoPEAttnProcessor:
-    """Attention processor that uses real-valued RoPE for ONNX compatibility.
-
-    Replaces the default QwenDoubleStreamAttnProcessor2_0 which uses
-    torch.view_as_complex (not ONNX-exportable).
-    """
-
-    def __call__(
-        self,
-        attn,
-        hidden_states,
-        encoder_hidden_states=None,
-        encoder_hidden_states_mask=None,
-        attention_mask=None,
-        image_rotary_emb=None,
-    ):
-        seq_txt = encoder_hidden_states.shape[1]
-
-        img_query = attn.to_q(hidden_states)
-        img_key = attn.to_k(hidden_states)
-        img_value = attn.to_v(hidden_states)
-
-        txt_query = attn.add_q_proj(encoder_hidden_states)
-        txt_key = attn.add_k_proj(encoder_hidden_states)
-        txt_value = attn.add_v_proj(encoder_hidden_states)
-
-        img_query = img_query.unflatten(-1, (attn.heads, -1))
-        img_key = img_key.unflatten(-1, (attn.heads, -1))
-        img_value = img_value.unflatten(-1, (attn.heads, -1))
-
-        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
-        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
-        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
-
-        if attn.norm_q is not None:
-            img_query = attn.norm_q(img_query)
-        if attn.norm_k is not None:
-            img_key = attn.norm_k(img_key)
-        if attn.norm_added_q is not None:
-            txt_query = attn.norm_added_q(txt_query)
-        if attn.norm_added_k is not None:
-            txt_key = attn.norm_added_k(txt_key)
-
-        if image_rotary_emb is not None:
-            img_cos, img_sin, txt_cos, txt_sin = image_rotary_emb
-            img_query = _apply_rope_real(img_query, img_cos, img_sin)
-            img_key = _apply_rope_real(img_key, img_cos, img_sin)
-            txt_query = _apply_rope_real(txt_query, txt_cos, txt_sin)
-            txt_key = _apply_rope_real(txt_key, txt_cos, txt_sin)
-
-        joint_query = torch.cat([txt_query, img_query], dim=1)
-        joint_key = torch.cat([txt_key, img_key], dim=1)
-        joint_value = torch.cat([txt_value, img_value], dim=1)
-
-        joint_query = joint_query.transpose(1, 2)
-        joint_key = joint_key.transpose(1, 2)
-        joint_value = joint_value.transpose(1, 2)
-        joint_hidden = torch.nn.functional.scaled_dot_product_attention(
-            joint_query, joint_key, joint_value, dropout_p=0.0, is_causal=False
-        )
-        joint_hidden = joint_hidden.transpose(1, 2)
-        joint_hidden = joint_hidden.flatten(2, 3)
-
-        txt_attn = joint_hidden[:, :seq_txt, :]
-        img_attn = joint_hidden[:, seq_txt:, :]
-
-        img_attn = attn.to_out[0](img_attn.contiguous())
-        if len(attn.to_out) > 1:
-            img_attn = attn.to_out[1](img_attn)
-        txt_attn = attn.to_add_out(txt_attn.contiguous())
-
-        return img_attn, txt_attn
-
-
-class TransformerONNXWrapper(nn.Module):
-    """Wraps QwenImageTransformer2DModel for ONNX export.
-
-    Pre-computes complex RoPE frequencies as real cos/sin buffers and replaces
-    the attention processors with ONNX-friendly real-valued versions.
-    """
-
-    def __init__(self, model, img_shapes, txt_seq_len):
-        super().__init__()
-        self.model = model
-
-        for block in self.model.transformer_blocks:
-            block.attn.set_processor(RealRoPEAttnProcessor())
-
-        with torch.no_grad():
-            img_freqs, txt_freqs = model.pos_embed(
-                img_shapes, max_txt_seq_len=txt_seq_len
-            )
-        self.register_buffer("img_cos", img_freqs.real.float().contiguous())
-        self.register_buffer("img_sin", img_freqs.imag.float().contiguous())
-        self.register_buffer("txt_cos", txt_freqs.real.float().contiguous())
-        self.register_buffer("txt_sin", txt_freqs.imag.float().contiguous())
-
-    def forward(self, hidden_states, encoder_hidden_states, timestep):
-        hidden_states = self.model.img_in(hidden_states)
-        timestep = timestep.to(hidden_states.dtype)
-
-        encoder_hidden_states = self.model.txt_norm(encoder_hidden_states)
-        encoder_hidden_states = self.model.txt_in(encoder_hidden_states)
-
-        temb = self.model.time_text_embed(timestep, hidden_states)
-
-        rope = (self.img_cos, self.img_sin, self.txt_cos, self.txt_sin)
-
-        for block in self.model.transformer_blocks:
-            encoder_hidden_states, hidden_states = block(
-                hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                encoder_hidden_states_mask=None,
-                temb=temb,
-                image_rotary_emb=rope,
-            )
-
-        hidden_states = self.model.norm_out(hidden_states, temb)
-        output = self.model.proj_out(hidden_states)
-        return output
-
-
-def _make_tiny_transformer_config():
-    """Tiny transformer config: ~100K params, 1 layer."""
-    return dict(
-        patch_size=2,
-        in_channels=4,
-        out_channels=4,
-        num_layers=1,
-        attention_head_dim=16,
-        num_attention_heads=4,
-        joint_attention_dim=64,
-        axes_dims_rope=(4, 6, 6),
-    )
-
-
-def _make_small_transformer_config():
-    """Small transformer config: ~1M params, 2 layers."""
-    return dict(
-        patch_size=2,
-        in_channels=16,
-        out_channels=16,
-        num_layers=2,
-        attention_head_dim=32,
-        num_attention_heads=8,
-        joint_attention_dim=256,
-        axes_dims_rope=(8, 12, 12),
-    )
-
-
-def _make_medium_transformer_config():
-    """Medium transformer config: ~39M params, 4 layers."""
-    return dict(
-        patch_size=2,
-        in_channels=32,
-        out_channels=32,
-        num_layers=4,
-        attention_head_dim=64,
-        num_attention_heads=8,
-        joint_attention_dim=512,
-        axes_dims_rope=(8, 28, 28),
-    )
-
-
-def _run_transformer_test(config, atol):
-    """Compile transformer with luminal backend, compare to PyTorch reference."""
-    from diffusers.models import QwenImageTransformer2DModel
-
-    from luminal import luminal_backend
-
-    model = QwenImageTransformer2DModel(**config).eval()
-    img_seq_len = 4
-    txt_seq_len = 3
-
-    wrapper = TransformerONNXWrapper(model, [(1, 2, 2)], txt_seq_len).eval()
-    wrapper_compiled = torch.compile(wrapper, backend=luminal_backend)
-
-    hidden = torch.randn(1, img_seq_len, config["in_channels"])
-    encoder_hs = torch.randn(1, txt_seq_len, config["joint_attention_dim"])
-    timestep = torch.tensor([1.0])
-
-    with torch.no_grad():
-        ref = wrapper(hidden, encoder_hs, timestep)
-        out = wrapper_compiled(hidden, encoder_hs, timestep)
-
-    assert torch.allclose(out, ref, atol=atol), (
-        f"max_diff={torch.max(torch.abs(out - ref)).item():.2e}"
-    )
-
-
-# ============================================================================
-# VAE helpers
-# ============================================================================
-
-
-class _OnnxFriendlyUpsample(nn.Module):
-    """Replaces nn.Upsample with repeat_interleave for ONNX compatibility."""
-
-    def __init__(self, scale_factor):
-        super().__init__()
-        if isinstance(scale_factor, (tuple, list)):
-            self.scale_factors = [int(s) for s in scale_factor]
-        else:
-            sf = int(scale_factor)
-            self.scale_factors = [sf]
-
-    def forward(self, x):
-        for dim_offset, sf in enumerate(self.scale_factors):
-            if sf > 1:
-                x = x.repeat_interleave(sf, dim=2 + dim_offset)
-        return x
-
-
-def _make_tiny_vae_config():
-    """Tiny VAE config for testing."""
-    return dict(
-        base_dim=8,
-        z_dim=4,
-        dim_mult=[1, 2],
-        num_res_blocks=1,
-        attn_scales=[],
-        temperal_downsample=[False],
-        dropout=0.0,
-        input_channels=3,
-    )
-
-
-def _make_medium_vae_config():
-    """Medium VAE config: base_dim=32, z_dim=8."""
-    return dict(
-        base_dim=32,
-        z_dim=8,
-        dim_mult=[1, 2, 4],
-        num_res_blocks=2,
-        attn_scales=[],
-        temperal_downsample=[False, True],
-        dropout=0.0,
-        input_channels=3,
-    )
-
-
-def _prepare_vae_for_onnx(vae):
-    """Replace non-ONNX-exportable modules in the VAE."""
-    import diffusers.models.autoencoders.autoencoder_kl_qwenimage as vae_mod
-
-    def _replace(module):
-        for name, child in module.named_children():
-            if isinstance(child, vae_mod.QwenImageUpsample):
-                setattr(module, name, _OnnxFriendlyUpsample(child.scale_factor))
-            else:
-                _replace(child)
-
-    _replace(vae)
-    return vae
-
-
-class _VAEDecoderWrapper(nn.Module):
-    def __init__(self, vae):
-        super().__init__()
-        self.vae = vae
-
-    def forward(self, z):
-        return self.vae.decode(z).sample
-
-
-def _export_and_simplify(wrapper, inputs, input_names, output_names):
-    """Export model to ONNX and simplify with onnxsim."""
-    tmp = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False)
-    tmp_path = tmp.name
-    tmp.close()
-    try:
-        torch.onnx.export(
-            wrapper,
-            inputs,
-            tmp_path,
-            opset_version=20,
-            input_names=input_names,
-            output_names=output_names,
-            dynamo=False,
-        )
-        m = onnx.load(tmp_path)
-        m_sim, check = onnxsim.simplify(m)
-        assert check, "onnxsim simplification failed"
-        onnx.save(m_sim, tmp_path)
-        return tmp_path
-    except Exception:
-        os.unlink(tmp_path)
-        raise
-
-
-def _run_vae_test(config, atol):
-    """Export VAE decoder to ONNX, run through luminal, compare."""
-    from diffusers import AutoencoderKLQwenImage
-
-    import luminal
-
-    backend = os.environ.get("LUMINAL_BACKEND", "native")
-    vae = AutoencoderKLQwenImage(**config).eval()
-    vae = _prepare_vae_for_onnx(vae)
-
-    wrapper = _VAEDecoderWrapper(vae).eval()
-    latents = torch.randn(1, config["z_dim"], 1, 4, 4)
-
-    with torch.no_grad():
-        ref = wrapper(latents)
-
-    onnx_path = _export_and_simplify(wrapper, (latents,), ["latents"], ["output"])
-    try:
-        graph = luminal.process_onnx(onnx_path, backend)
-        graph.set_input("latents", latents.flatten().tolist())
-        graph.run()
-        out_data = graph.get_output("output")
-        out = torch.tensor(out_data, dtype=torch.float32).reshape(ref.shape)
-    finally:
-        os.unlink(onnx_path)
-
-    assert torch.allclose(out, ref, atol=atol), (
-        f"max_diff={torch.max(torch.abs(out - ref)).item():.2e}"
-    )
-
-
-# ============================================================================
-# Tests
-# ============================================================================
-
-
-def test_qwen_image_transformer_tiny():
-    """Tiny QwenImage transformer: 1 layer, 4 heads, dim=64."""
-    _run_transformer_test(_make_tiny_transformer_config(), atol=1e-4)
-
-
-def test_qwen_image_transformer_small():
-    """Small QwenImage transformer: 2 layers, 8 heads, dim=256."""
-    _run_transformer_test(_make_small_transformer_config(), atol=1e-4)
-
-
-def test_qwen_image_transformer_medium():
-    """Medium QwenImage transformer: 4 layers, 8 heads, dim=512."""
-    _run_transformer_test(_make_medium_transformer_config(), atol=1e-4)
-
-
-def test_qwen_image_transformer_full():
-    """Full QwenImage transformer (production defaults)."""
-    from diffusers.models import QwenImageTransformer2DModel
-
-    from luminal import luminal_backend
-
-    model = QwenImageTransformer2DModel().eval()
-    config = {k: v for k, v in dict(model.config).items() if not k.startswith("_")}
-
-    wrapper = TransformerONNXWrapper(model, [(1, 2, 2)], txt_seq_len=3).eval()
-    wrapper_compiled = torch.compile(wrapper, backend=luminal_backend)
-
-    hidden = torch.randn(1, 4, config["in_channels"])
-    encoder_hs = torch.randn(1, 3, config["joint_attention_dim"])
-    timestep = torch.tensor([1.0])
-
-    with torch.no_grad():
-        ref = wrapper(hidden, encoder_hs, timestep)
-        out = wrapper_compiled(hidden, encoder_hs, timestep)
-
-    assert torch.allclose(out, ref, atol=1e-4), (
-        f"max_diff={torch.max(torch.abs(out - ref)).item():.2e}"
-    )
-
-
-def test_qwen_image_vae_decoder_tiny():
-    """Tiny QwenImage VAE decoder: base_dim=8, z_dim=4."""
-    _run_vae_test(_make_tiny_vae_config(), atol=1e-3)
-
-
-def test_qwen_image_vae_decoder_medium():
-    """Medium QwenImage VAE decoder: base_dim=32, z_dim=8."""
-    _run_vae_test(_make_medium_vae_config(), atol=1e-3)
-
-
-@pytest.mark.skip(reason="Full production VAE -- expected to be slow/OOM")
-def test_qwen_image_vae_decoder_full():
-    """Full QwenImage VAE decoder (production defaults)."""
-    from diffusers import AutoencoderKLQwenImage
-
-    config = dict(AutoencoderKLQwenImage().config)
-    config = {k: v for k, v in config.items() if not k.startswith("_")}
-    _run_vae_test(config, atol=1e-3)