luminal_python: vanilla-PyTorch DLRMv1 fast paths + batched FFI + input cache

Lands `torch.compile(model, backend=luminal_backend)` for vanilla
DLRMv1 (`nn.EmbeddingBag` per table, pairwise-dot interaction, top
MLP) at correctness parity with PyTorch eager (max abs diff
≤ 1.8e-7) and ~6.5× faster than `pt_eager` across an 11 × 5
(batch × num_cat) sweep. Crosses over `graph_safe + inductor +
manual CUDAGraph` at `nc=32, b=2048` (174 µs vs 241 µs).

Three layers of changes.

Translator (crates/luminal_python/rust/src/translator/)
- aten._embedding_bag(_forward_only).default: 2-D weight + 1-D
  indices + uniform-stride offsets → embedding_bag_sum_kernel
  (per-table) or, when a sibling group of emb-bag outputs feeds the
  pairwise-dot peephole, the fused multi_table_embedding_bag_sum
  kernel.
- aten.index_select.default: 2-D source, dim=0, 1-D index lowering.
- aten.addmm.default fast path: detect addmm(bias, x, weight.t())
  with α=β=1 and emit luminal_cuda_lite::kernel::linear_bias against
  the original (N, K) weight. With a unique relu / sigmoid consumer
  downstream, absorb it into linear_bias_relu / linear_bias_sigmoid
  via a forward-looking absorbed-nodes pass.
- aten.bmm.default fast path: detect bmm(A, permute(B, [0,2,1]))
  via transpose_2d_source and emit matmul_3d_t.
- aten.sum.dim_IntList peephole: sum ← view ← index_select chain →
  embedding_bag_sum_kernel.
- aten.index.Tensor peephole for DLRM strict-lower-tri pairwise
  dot: detect Z[:, li, lj] over bmm(cat(unsqueezes), permute([0,2,1]))
  and fold to dlrm_pairwise_dot_lower_tri (variadic) or, when all
  cat inputs but one are emb-bag outputs with matching shapes,
  dlrm_pairwise_dot_lower_tri_stacked over (dense, emb_stack).
- aten.permute.default: record 2-D [1,0] and 3-D [0,2,1] sources
  into transpose_2d_source for the matmul fast paths.
- node_chain + op_inputs side-tables populated alongside every
  translated op (including variadic-first-input ops like cat) so
  multi-input peepholes can walk back across the producer chain.
- consumers map + absorbed_nodes set populated once at start of
  translate_graph for forward-looking fusions (addmm + relu).
- Post-translation DCE: walks back from Output HLIR sinks and drops
  every unreachable producer. Load-bearing for the pairwise-dot
  peephole — the cat/bmm/permute chain it supersedes survives
  egglog on its own.

cuda_lite kernels (crates/luminal_cuda_lite/src/kernel/)
- embedding_bag.rs (new): EmbeddingBagSumKernel (single table),
  MultiTableEmbeddingBagSumKernel (K weight + K idx pointer pairs
  through packed staging buffers), and StackedEmbeddingBagKernel
  (pre-stacked weight + block-bundle launch layout). All produce
  (batch, num_tables, d) F32 in one launch.
- dlrm_interact.rs (new): PairwiseDotLowerTriKernel (variadic) and
  PairwiseDotLowerTriStackedKernel (block-per-batch with
  cooperative cache of all F·D feature vectors into shared memory;
  pair (i, j) recovered from p via closed-form
  (1+sqrt(1+8p))/2). Both emit the F(F-1)/2 dot products in one
  launch.
- matmul2d.rs: linear_bias_(relu|sigmoid)(_split_a) frontends over
  the existing fused Matmul2DKernel.
- runtime.rs: gate end-of-execute cuStreamSynchronize on
  self.profiling so the regular execute() path stays capturable by
  external graph capture (and is consistent with `set_input_*_ptr`
  being persistent across calls). Preserve CudaInput::Ptr entries
  in the post-execute "consume" step — they're non-owning views
  over caller memory and represent stable input slots that the
  wrapper may skip re-registering on a hot iter.
- runtime.rs: synchronize_stream() and read_per_kernel_timings_ms()
  helpers (the latter is a stub; the per-kernel timing infra from
  origin/dlrm-fused-kernels isn't ported yet — example's optional
  LUMINAL_KERNEL_TIMING=1 path prints "no timings available").

Wrapper (crates/luminal_python/{rust/src/compiled_graph,
src/luminal/compiled_model}.{rs,py})
- compiled_graph.rs: tensor_id(name) → u32 returns the underlying
  NodeIndex once, and run_with_ptrs(inputs, outputs) takes two
  flat Vec<(u32, u64, usize)> lists for batched input + output
  pointer registration in one pyo3 hop instead of N
  set_input_device_ptr + M set_output_device_ptr + run calls.
- compiled_model.py: __init__ resolves and caches the input /
  output IDs once. __call__ has a fast path that uses
  run_with_ptrs when all user inputs are on CUDA and outputs are
  fp. Per-input cache of (id(orig_tensor), orig_data_ptr,
  cast_tensor) skips Python-side detach/contiguous/.to/data_ptr/
  numel/element_size and the FFI register when the same user
  tensor is passed again with an unchanged data_ptr. Always
  returns a tuple to match the unbatched path's contract
  (returning a bare tensor broke dynamo's output handling — it
  iterated the first dim of the result, slicing the output to
  shape (1,)).

Per-iter time at (nc=32, batch=2048):

  Before:  set_in 995us + run 61us + collect 13us = 1069us
  After:   cache 15us + out 11us + run_with_ptrs 56us = 82us

13× wall-clock reduction, none of it in the kernels themselves.

Status
- 165-cell sweep (batch ∈ {2..2048} × num_cat ∈ {2..32} × 3
  variants) at examples/dlrm/results.csv.
- Beats pt_eager in 55/55 cells (4.7–7.1×).
- Beats graph_safe_inductor_cg at 1 cell (the largest), within
  10% at 3 more.
- Hand-written rust DLRM at examples/dlrm/src/main.rs runs the
  same kernel set and lands at 104µs on the same cell — the
  remaining gap to `luminal_compiled` (174µs) is the residual
  Python+dynamo+runtime per-call fixed cost.

Full breakdown at examples/dlrm/RESULTS.md.

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

.cargo/config.toml

@@ -1,3 +1,6 @@
+[alias]
+examples = "run --release --bin examples-perf --"
+
 [target.aarch64-unknown-linux-gnu]
 rustflags = [
     "-Ctarget-feature=+fp16,+fhm"

.github/workflows/test-core.yml

@@ -21,4 +21,4 @@ jobs:
         steps:
             - uses: actions/checkout@v6
             - name: Run tests
-              run: cargo test --release --workspace --exclude luminal_cuda_lite --exclude luminal_metal --exclude luminal_bench --verbose
+              run: cargo test --release -p luminal -p luminal_nn -p luminal_tracing -p luminal_python --verbose

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

@@ -0,0 +1,67 @@
+name: Test Full CUDA
+
+on:
+    pull_request_target:
+        branches: ["main"]
+        types: [labeled, synchronize]
+    workflow_dispatch:
+
+jobs:
+    rust_cuda_ignored_tests:
+        if: >-
+            github.event_name == 'workflow_dispatch'
+            || (github.event_name == 'pull_request_target'
+                && contains(github.event.pull_request.labels.*.name, 'full-modal-ready'))
+        name: Rust CUDA Ignored Tests
+        runs-on: ubuntu-latest
+        environment: Modal
+        timeout-minutes: 300
+
+        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:
+                  python-version: "3.11"
+            - name: Install Modal
+              run: pip install modal
+            - name: Run ignored CUDA Rust tests on Modal
+              env:
+                  MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+                  MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+                  GPU_TYPE: H100
+                  MODAL_TIMEOUT: "14400"
+                  CARGO_TEST_ARGS: "--ignored --test-threads=1"
+              run: modal run ci/modal_cargo_test.py
+
+    python_cuda_slow_tests:
+        if: >-
+            github.event_name == 'workflow_dispatch'
+            || (github.event_name == 'pull_request_target'
+                && contains(github.event.pull_request.labels.*.name, 'full-modal-ready'))
+        name: Python CUDA Slow Tests
+        runs-on: ubuntu-latest
+        environment: Modal
+        timeout-minutes: 300
+        defaults:
+            run:
+                working-directory: crates/luminal_python
+
+        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:
+                  python-version: "3.11"
+            - name: Install Modal
+              run: pip install modal
+            - name: Run slow pytest CUDA tests on Modal
+              env:
+                  MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+                  MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+                  HF_TOKEN: ${{ secrets.HF_TOKEN }}
+              run: modal run modal_pytest_runner.py --gpu A100 --timeout 14400 tests/ -v -s -m slow

.github/workflows/test-metal.yml

@@ -17,3 +17,20 @@ jobs:
             - uses: actions/checkout@v6
             - name: Run Metal crate tests
               run: rustup update; cargo test --release -p luminal_metal --verbose -- --test-threads=1
+
+    llama_1b_metal_example:
+        name: Llama 1B Metal Example
+        runs-on: macos-14-xlarge
+        timeout-minutes: 120
+
+        steps:
+            - uses: actions/checkout@v6
+            - name: Print runner hardware
+              run: system_profiler SPHardwareDataType SPDisplaysDataType
+            - name: Cache Hugging Face models
+              uses: actions/cache@v4
+              with:
+                  path: ~/.cache/huggingface
+                  key: llama-1b-metal-hf-${{ runner.os }}-${{ runner.arch }}-v1
+            - name: Run Llama 1B Metal example and validate output
+              run: rustup update; python3 ci/metal_llama_1b_example.py

AGENTS.md

@@ -8,4 +8,14 @@ All other functionality is split into crates in the `crates/` directory. For ins
 ## Testing Instructions
 - Find the CI plan in the .github/workflows folder.
 - Currently running `cargo test` in luminal_metal and luminal_cuda_lite require access to an Apple and Nvidia GPU respectively.
-- PRs must have no clippy errors and `cargo fmt` must be ran before a PR is submitted.
+- PRs must have no clippy errors and `cargo fmt` must be ran before a PR is submitted.
+
+## Debugging and Correctness
+- Treat model examples as specifications of the intended architecture. Do not change model code, prompt templates, weights, or example logic to hide compiler/runtime/search bugs unless the model code is demonstrably semantically wrong.
+- When outputs are incorrect, first root-cause the failing compiler/runtime path. Prefer isolating the bad LLIR/HLIR graph, rewrite, op lowering, shape/stride assumption, layout contract, or runtime implementation that caused the mismatch.
+- Avoid narrow special-case fixes. A fix should state and enforce the general invariant it relies on, or explicitly document why the affected operation is only valid for a restricted layout/shape and ensure rewrites enforce that restriction.
+- For e-graph/search issues, assume all selectable LLIR graphs are intended to be semantically equivalent. If two selectable graphs disagree, debug the equivalence violation rather than selecting around the bad graph.
+- Add regression tests at the level where the bug occurred. Prefer tests that compare against a semantic reference such as `NativeRuntime` or a small independent reference, and use fixed seeds for any randomized search/fuzz test so failures are reproducible.
+
+## Compiler Rewrite Boundary
+- All graph pattern matching and op selection must be expressed in egglog rewrites. Do not add Rust-side LLIR graph post-passes that search for op patterns, fuse kernels, select backend ops, or otherwise rewrite extracted graphs after egglog. If a backend needs a fused/specialized op, add the match and rewrite in egglog and let extraction produce that op directly.

README.md

@@ -55,23 +55,27 @@ Luminal can run Q8 Llama 3 8B at ~80% of theoretical max performance on an H100.
 
 The core of Luminal is and always will be minimal. It should be possible to understand the entire core library in an afternoon.
 
+### PyTorch-native
+
+Luminal directly integrates with PyTorch as a compiler backend. Simply do `torch.compile(model, backend=luminal_cuda)` to compile your PyTorch models. We also have an excellent tensor API in Rust.
+
 ### RISC-style architecture
 
-Everything in Luminal boils down to 14 primitive ops:
+Everything in Luminal boils down to 15 primitive ops:
 
 - Unary - `Log2, Exp2, Sin, Sqrt, Recip`
 - Binary - `Add, Mul, Mod, LessThan`
-- Other - `SumReduce, MaxReduce, Iota, Gather, Cast`
+- Other - `SumReduce, MaxReduce, Iota, Gather, Scatter, Cast`
 
-These ops are enough to support transformers, convnets, and nearly every popular model.
+These ops are enough to support transformers, convnets, and nearly every popular model in the world.
 
 ### Search
 
-The best heuristic is no heuristic. We try to search every possible decision to give the compiler the most flexibility to discover complex optimizations. This allows us to automatically derive Flash Attention and other similarly complex rewrites. It also allows us to stay extremely small long into the future and beat the performance of far larger frameworks with tons of handwritten kernels.
+The best heuristic is no heuristic. Luminal tries to search every possible decision to give the compiler the flexibility to discover complex optimizations. This allows us to automatically discover Flash Attention and other similarly complex optimizations without relying on hand-written operations or heuristics. It also allows us to stay extremely small and simple long into the future and beat the performance of far larger frameworks.
 
 ### Native
 
-The current ML ecosystem is too fragmented, and the solution isn't another layer of abstraction. Luminal is written in rust, and interacts directly with the CUDA / Metal APIs. No indirections or abstractions, docker containers, or virtual environments. Just a statically-linked rust crate.
+The current ML ecosystem is too fragmented, and the solution isn't another layer of abstraction. Luminal is written in rust, and interacts directly with the accelerator APIs (CUDA, Metal, etc.). No indirections or abstractions, compatability layers, docker containers, or virtual environments. Just a statically-linked rust crate.
 
 ### Validated against Pytorch
 
@@ -85,39 +89,45 @@ Most deep learning libraries are eager-first, meaning each op call directly oper
 
 However, this isn't great for performance. What makes sense for a developer doesn't work well for the machine, in the same way that no one writes assembly by hand. Most libraries try to fix this problem by tacking on operator fusion or JIT compilation to try to change the compilation flow to something better for the machine. Turns out this is [super](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_overview.html) [difficult](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) [even](https://pytorch.org/docs/stable/jit.html) [for](https://pytorch.org/docs/stable/fx.html#torch.fx.symbolic_trace) Pytorch!
 
+### What about XLA?
+
+XLA, torch.compile, TVM, and other traditional compiler stacks suffer from complexity explosion. They are made up of a very large set of destructive (one-direction) rewrite rules that lower and optimize a graph from a high-level representation to low-level machine code. But since these rules are destructive, they are required to only fire when it's certian that there's a performance benefit. This leads to the rules becoming very complex, special-cased, and numerous. Once additional hardware backends, model architectures, and new dtypes get thrown in, they suffer from the weight of their complexity and often produce very suboptimal code, requiring DSLs like Pallas or Triton to regain performance.
+
 ### Compile everything
 
-A core tenet of Luminal is ahead-of-time compilation. Whenever possible, push everything to compile time and leave nothing to run time. Luminal takes an approach more similar to [XLA](https://www.tensorflow.org/xla), and [tinygrad](https://github.com/tinygrad/tinygrad). Everything's static here. When you write out an expression like `x + y`, no actual computation happens. The operation is recorded to a directed acyclic computation graph for execution later. Only once `graph.execute()` is ran does the computation happen. _But isn't that just lazy execution?_ Yes it is! But in luminal **everything is done this way**. All neural networks are built up as one or a few static computation graphs, compiled, and executed later.
+A core tenet of Luminal is ahead-of-time compilation. Whenever possible, push everything to compile time and leave nothing to run time. Luminal takes an approach more similar to [XLA](https://www.tensorflow.org/xla), and [tinygrad](https://github.com/tinygrad/tinygrad). Everything's static here. When you write out an expression like `x + y`, no actual computation happens. The operation is recorded to a directed acyclic computation graph for execution later. Only once `graph.execute()` is ran does the computation happen. _But isn't that just lazy execution?_ Yes it is! But in luminal **everything is done this way**. All neural networks are built up as a static computation graphs, compiled, and executed later.
+
+### First-class dynamism
+
+A fully-static world would be nice, but we live in a world of nessecary dynamism. So we model dynamic shapes natively, as symbolic dimensions. Luminal supports arbitrary symbolic dimensions, including complex expressions, to give us shapes like `(s, 4096)`, `(b, h, w + 3)`, etc. This rich representation gives the compiler full visibility into shapes and lets it still do aggressive specialization.
 
 **But why?**
 
-A consequence of this is that the actual computation that gets ran can be radically different than the code that was written. Since we have an entire neural network fully represented in a compute graph, our compilers have global knowledge. This means we can push most ML complexity to the compilers. For instance, devices, datatypes, and execution schedules are all handled by compliers. Even autograd is handled by a compiler!
+A consequence of this is that the actual computation that gets ran can be radically different than the code that was written. Since we have an entire neural network fully represented in a compute graph, Luminal has global knowledge. This means we can push most ML complexity to the compiler. For instance, devices, datatypes, and even autograd is modeled ahead of time and optimized by the compiler!
 
 Now we can do:
 
 - Aggressive kernel fusion
 - Shape-specific kernels compiled at runtime
-- Devices and Dtypes are handled through compilers (just run the CUDA compiler to convert the graph to use CUDA kernels, then the fp16 compiler to convert to half-precision kernels)
-- Networks can be written in generic code, but compiled and ran fast on hyper-specific architectures (try writing a PyTorch network that works with both TF32 dtypes and TPUs; get ready for if statement hell...)
+- Low-precision dtypes (mxfp4, nvfp4, fp8, etc.)
+- Complex mutli-device parallelism topologies, searched ahead-of-time
+- Networks can be written in generic code, but compiled and ran fast on hyper-specific architectures
 
 ## Where are we?
 
-- Search is partially merged. We are between 1.0 and 2.0 (search), which will be completed within the next month or so.
-- Metal and Cuda are supported for running models on Macs and Nvidia GPUs respectively, in both full and half precision.
-- Full training support with graph-based autograd.
-- Llama 3, Phi 3, Whisper and Yolo v8 are implemented in `examples/`. See instructions above for running.
+- Native PyTorch support
+- Many kernel libraries supported in the search space (FlashInfer, cuBLASLt, etc.)
+- Many models implemented in our Rust tensor API in `examples/`.
 - We have a small library of NN modules in `luminal_nn`, including transformers.
 - A significant amount of high-level ops are implemented in `hl_ops`. We are aiming to match the most used ~80% of the pytorch api.
 
 Some things on the roadmap:
 
-- Expand the search space to utilize Tensor Cores more flexibly
-- Bring cuda to parity with Metal
-- Add Blackwell intrinsics, such as TMEM and TMA
-- Build a ROCm backend
-- Build benchmarking suite to test against other libs
-- Distributed data, pipeline and tensor parallel.
-- Beat PT 2.0 perf on LLM inference _and_ training
+- More fine-grained dialects supporting thread- and warp-level intrinsics like TMA and tcgen.05
+- ROCm backend
+- More public infernce accelerator backends (coming very soon...)
+- Public benchmarking suite
+- Automatically searched model parallelism (TP, PP, EPS, EPR, SP, etc.)
 - Write compiler for quantum photonic retro encabulator
 - Build dyson swarm
 

ci/example_output.py

@@ -0,0 +1,85 @@
+import re
+
+ANSI_ESCAPE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
+
+EXPECTED_OUTPUT = {
+    "gemma4_moe": [
+        "city of romance, art and culture",
+    ],
+    "whisper": [
+        "ask not what your country can do for you",
+    ],
+}
+
+EXPECTED_CONCEPTS = {
+    "llama": [
+        ["layers"],
+        ["neurons", "nodes"],
+        ["learn", "learning", "adapt"],
+        ["data", "patterns", "features"],
+    ],
+    "gemma": [
+        ["neural network", "neural networks"],
+        ["nodes", "neurons"],
+        ["layers"],
+        ["weights"],
+        ["training", "learn", "learns"],
+    ],
+    "qwen": [
+        ["neural network", "neural networks"],
+        ["computational model", "computational system"],
+        ["brain"],
+        ["layers"],
+        ["neurons", "nodes"],
+        ["learn", "learning", "training"],
+    ],
+    "qwen3_moe": [
+        ["capital"],
+        ["france"],
+        ["paris"],
+    ],
+}
+
+
+def normalize_output(output: str) -> str:
+    output = ANSI_ESCAPE.sub("", output)
+    output = output.replace("\r", "\n")
+    return re.sub(r"\s+", " ", output).casefold()
+
+
+def validate_output(example: str, output: str):
+    normalized_output = normalize_output(output)
+
+    expected_concepts = EXPECTED_CONCEPTS.get(example)
+    if expected_concepts is not None:
+        missing = [
+            concept_group
+            for concept_group in expected_concepts
+            if not any(normalize_output(term) in normalized_output for term in concept_group)
+        ]
+        if missing:
+            expected = "\n  - ".join(" / ".join(group) for group in expected_concepts)
+            missing_terms = "\n  - ".join(" / ".join(group) for group in missing)
+            raise AssertionError(
+                f"Output check failed for {example!r}.\n"
+                f"Expected concept groups:\n  - {expected}\n"
+                f"Missing concept groups:\n  - {missing_terms}"
+            )
+
+        expected = ", ".join(" / ".join(group) for group in expected_concepts)
+        print(f"\nOutput check passed for {example!r}: found concepts {expected}")
+        return
+
+    expected_phrases = EXPECTED_OUTPUT.get(example)
+    if expected_phrases is None:
+        raise ValueError(f"No expected output phrases configured for example {example!r}")
+
+    for phrase in expected_phrases:
+        if normalize_output(phrase) in normalized_output:
+            print(f"\nOutput check passed for {example!r}: found {phrase!r}")
+            return
+
+    expected = "\n  - ".join(expected_phrases)
+    raise AssertionError(
+        f"Output check failed for {example!r}. Expected one of:\n  - {expected}"
+    )

ci/examples_perf.py

@@ -0,0 +1,185 @@
+import os
+import subprocess
+import sys
+import time
+from dataclasses import dataclass, field
+
+from example_output import validate_output
+
+
+DEFAULT_EXAMPLES = ["llama", "gemma", "qwen", "qwen3_moe", "gemma4_moe", "whisper"]
+
+EXAMPLE_CARGO_ARGS = {
+    "llama": ["run", "--release", "-p", "llama"],
+    "gemma": ["run", "--release", "-p", "gemma"],
+    "qwen": ["run", "--release", "-p", "qwen", "--features", "cuda"],
+    "qwen3_moe": ["run", "--release", "-p", "qwen3_moe"],
+    "gemma4_moe": ["run", "--release", "-p", "gemma4_moe"],
+    "whisper": ["run", "--release", "-p", "whisper"],
+}
+
+
+@dataclass
+class Metrics:
+    ttft_ms: float | None = None
+    tpot_ms: float | None = None
+    tps: float | None = None
+
+
+@dataclass
+class ExampleResult:
+    name: str
+    ok: bool
+    metrics: Metrics = field(default_factory=Metrics)
+    wall_s: float = 0.0
+    error: str | None = None
+
+
+def main() -> None:
+    args = [arg for arg in sys.argv[1:] if arg != "--"]
+    if any(arg in {"-h", "--help"} for arg in args):
+        print_help()
+        return
+    if "--list" in args:
+        print("\n".join(DEFAULT_EXAMPLES))
+        return
+
+    examples = args or DEFAULT_EXAMPLES
+    results = [run_example(example) for example in examples]
+    print_table(results)
+    if any(not result.ok for result in results):
+        raise SystemExit(1)
+
+
+def print_help() -> None:
+    print(
+        "Run validated Luminal examples, validate textual output, and summarize perf.\n"
+        "\n"
+        "Usage:\n"
+        "  cargo examples\n"
+        "  cargo examples llama qwen whisper\n"
+        "\n"
+        "Options:\n"
+        "  --list    Print the default validated examples\n"
+        "  -h, --help\n"
+        "\n"
+        f"The default set matches the Modal examples CI: {', '.join(DEFAULT_EXAMPLES)}."
+    )
+
+
+def run_example(example: str) -> ExampleResult:
+    cargo_args = EXAMPLE_CARGO_ARGS.get(example)
+    if cargo_args is None:
+        known = ", ".join(DEFAULT_EXAMPLES)
+        return ExampleResult(example, False, error=f"unknown example; known examples: {known}")
+
+    print(f"\n=== Running {example} ===")
+    print(f"$ cargo {' '.join(cargo_args)}")
+    started = time.monotonic()
+    env = os.environ.copy()
+    env.setdefault("CUDARC_CUDA_VERSION", "12080")
+    process = subprocess.Popen(
+        ["cargo", *cargo_args],
+        cwd=repo_root(),
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+    assert process.stdout is not None
+
+    chunks: list[bytes] = []
+    while True:
+        chunk = process.stdout.read1(4096)
+        if not chunk:
+            break
+        sys.stdout.buffer.write(chunk)
+        sys.stdout.buffer.flush()
+        chunks.append(chunk)
+
+    return_code = process.wait()
+    output = b"".join(chunks).decode("utf-8", errors="replace")
+    wall_s = time.monotonic() - started
+    metrics = parse_metrics(output)
+
+    if return_code:
+        return ExampleResult(
+            example,
+            False,
+            metrics=metrics,
+            wall_s=wall_s,
+            error=f"process exited with code {return_code}",
+        )
+
+    try:
+        validate_output(example, output)
+    except Exception as exc:
+        return ExampleResult(example, False, metrics=metrics, wall_s=wall_s, error=str(exc))
+
+    return ExampleResult(example, True, metrics=metrics, wall_s=wall_s)
+
+
+def repo_root() -> str:
+    return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def parse_metrics(output: str) -> Metrics:
+    metrics = Metrics()
+    for line in output.splitlines():
+        if "TTFT:" in line:
+            metrics.ttft_ms = parse_number_after(line, "TTFT:")
+        if "TPOT:" in line:
+            metrics.tpot_ms = parse_number_after(line, "TPOT:")
+        if "tok/s" in line:
+            metrics.tps = parse_tok_per_second(line)
+    if metrics.tps is None and metrics.tpot_ms:
+        metrics.tps = 1000.0 / metrics.tpot_ms
+    return metrics
+
+
+def parse_number_after(line: str, marker: str) -> float | None:
+    tail = line.split(marker, 1)[1].lstrip()
+    chars = []
+    for char in tail:
+        if char.isdigit() or char == ".":
+            chars.append(char)
+        else:
+            break
+    if not chars:
+        return None
+    return float("".join(chars))
+
+
+def parse_tok_per_second(line: str) -> float | None:
+    head = line.split("tok/s", 1)[0].rstrip(" (")
+    parts = head.split()
+    if not parts:
+        return None
+    try:
+        return float(parts[-1])
+    except ValueError:
+        return None
+
+
+def print_table(results: list[ExampleResult]) -> None:
+    print("\nSummary")
+    print(f"{'example':<14} {'status':<8} {'TTFT ms':>10} {'TPOT ms':>10} {'tok/s':>10} {'wall s':>10}")
+    print("-" * 68)
+    for result in results:
+        status = "ok" if result.ok else "failed"
+        print(
+            f"{result.name:<14} {status:<8} "
+            f"{format_metric(result.metrics.ttft_ms):>10} "
+            f"{format_metric(result.metrics.tpot_ms):>10} "
+            f"{format_metric(result.metrics.tps):>10} "
+            f"{result.wall_s:>10.1f}"
+        )
+        if result.error:
+            print(f"  error: {result.error}")
+
+
+def format_metric(value: float | None) -> str:
+    return "-" if value is None else f"{value:.2f}"
+
+
+if __name__ == "__main__":
+    main()

ci/metal_llama_1b_example.py

@@ -0,0 +1,48 @@
+import os
+import subprocess
+import sys
+
+
+def run_and_capture(command: list[str], *, cwd: str, env: dict[str, str]) -> str:
+    process = subprocess.Popen(
+        command,
+        cwd=cwd,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+    assert process.stdout is not None
+
+    chunks = []
+    while True:
+        chunk = process.stdout.read1(4096)
+        if not chunk:
+            break
+        sys.stdout.buffer.write(chunk)
+        sys.stdout.buffer.flush()
+        chunks.append(chunk)
+
+    return_code = process.wait()
+    output = b"".join(chunks).decode("utf-8", errors="replace")
+    if return_code:
+        raise subprocess.CalledProcessError(return_code, command, output=output)
+    return output
+
+
+def main():
+    repo_root = os.environ.get("GITHUB_WORKSPACE", os.getcwd())
+    sys.path.insert(0, os.path.join(repo_root, "ci"))
+    from example_output import validate_output
+
+    output = run_and_capture(
+        ["cargo", "run", "--release", "-p", "luminal_metal", "--example", "llama_1b"],
+        cwd=repo_root,
+        env=os.environ.copy(),
+    )
+    if "TTFT:" not in output or "TPOT:" not in output:
+        raise AssertionError("Llama 1B Metal example did not complete generation")
+    validate_output("llama", output)
+
+
+if __name__ == "__main__":
+    main()

ci/metal_qwen_example.py

@@ -0,0 +1,46 @@
+import os
+import subprocess
+import sys
+
+from example_output import validate_output
+
+def run_and_capture(command: list[str], *, cwd: str, env: dict[str, str]) -> str:
+    process = subprocess.Popen(
+        command,
+        cwd=cwd,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+    assert process.stdout is not None
+
+    chunks = []
+    while True:
+        chunk = process.stdout.read1(4096)
+        if not chunk:
+            break
+        sys.stdout.buffer.write(chunk)
+        sys.stdout.buffer.flush()
+        chunks.append(chunk)
+
+    return_code = process.wait()
+    output = b"".join(chunks).decode("utf-8", errors="replace")
+    if return_code:
+        raise subprocess.CalledProcessError(return_code, command, output=output)
+    return output
+
+
+def main():
+    repo_root = os.environ.get("GITHUB_WORKSPACE", os.getcwd())
+    output = run_and_capture(
+        ["cargo", "run", "--release", "-p", "qwen", "--features", "metal"],
+        cwd=repo_root,
+        env=os.environ.copy(),
+    )
+    if "TTFT:" not in output or "TPOT:" not in output:
+        raise AssertionError("qwen Metal example did not complete generation")
+    validate_output("qwen", output)
+
+
+if __name__ == "__main__":
+    main()

ci/modal_cargo_test.py

@@ -1,8 +1,10 @@
 import modal
 import subprocess
 import os
+import shlex
 
 gpu_type = os.environ.get("GPU_TYPE", "T4")
+modal_timeout = int(os.environ.get("MODAL_TIMEOUT", "7200"))
 CUDARC_CUDA_VERSION = "12080"
 
 app = modal.App("luminal-ci-cargo-test")
@@ -28,7 +30,7 @@ cuda_image = (
 @app.function(
     image=cuda_image,
     gpu=gpu_type,
-    timeout=7200,  # 2 hours
+    timeout=modal_timeout,
 )
 def run_cargo_test():
     """Run cargo test for luminal_cuda_lite on a Modal GPU."""
@@ -43,17 +45,20 @@ def run_cargo_test():
     )
     compute_cap = result.stdout.strip().replace(".", "")
 
+    test_args = shlex.split(os.environ.get("CARGO_TEST_ARGS", "--test-threads=1"))
+    cmd = [
+        "cargo",
+        "test",
+        "--release",
+        "-p",
+        "luminal_cuda_lite",
+        "--verbose",
+        "--",
+        *test_args,
+    ]
+    print("Running:", " ".join(cmd), flush=True)
     subprocess.run(
-        [
-            "cargo",
-            "test",
-            "--release",
-            "-p",
-            "luminal_cuda_lite",
-            "--verbose",
-            "--",
-            "--test-threads=1",
-        ],
+        cmd,
         cwd=WORKDIR,
         env={
             **os.environ,

ci/modal_example.py

@@ -1,5 +1,4 @@
 import os
-import re
 import subprocess
 import sys
 
@@ -21,28 +20,8 @@ hf_cache = modal.Volume.from_name(
 
 WORKDIR = "/workspace/luminal"
 
-ANSI_ESCAPE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
-
-EXPECTED_OUTPUT = {
-    "llama": [
-        "complex system modeled after the structure and function of the human brain",
-    ],
-    "gemma": [
-        "recognize pictures of cats",
-        "little detectives looking for specific features",
-    ],
-    "qwen": [
-        "computational model inspired by the structure and function of the human brain",
-    ],
-    "qwen3_moe": [
-        "The capital of France is Paris",
-    ],
-    "gemma4_moe": [
-        "city of romance, art and culture",
-    ],
-    "whisper": [
-        "ask not what your country can do for you",
-    ],
+EXAMPLE_CARGO_ARGS = {
+    "qwen": ["--features", "cuda"],
 }
 
 
@@ -72,28 +51,6 @@ def run_and_capture(command: list[str], *, cwd: str, env: dict[str, str]) -> str
     return output
 
 
-def normalize_output(output: str) -> str:
-    output = ANSI_ESCAPE.sub("", output)
-    output = output.replace("\r", "\n")
-    return re.sub(r"\s+", " ", output).casefold()
-
-
-def validate_output(example: str, output: str):
-    expected_phrases = EXPECTED_OUTPUT.get(example)
-    if expected_phrases is None:
-        raise ValueError(f"No expected output phrases configured for example {example!r}")
-
-    normalized_output = normalize_output(output)
-    for phrase in expected_phrases:
-        if normalize_output(phrase) in normalized_output:
-            print(f"\nOutput check passed for {example!r}: found {phrase!r}")
-            return
-
-    expected = "\n  - ".join(expected_phrases)
-    raise AssertionError(
-        f"Output check failed for {example!r}. Expected one of:\n  - {expected}"
-    )
-
 cuda_image = (
     modal.Image.from_registry(
         "nvcr.io/nvidia/pytorch:25.03-py3"
@@ -123,6 +80,8 @@ cuda_image = (
 def run_example(example: str):
     """Build and run a luminal example on a Modal GPU."""
     subprocess.run(["nvidia-smi"], check=True)
+    sys.path.insert(0, f"{WORKDIR}/ci")
+    from example_output import validate_output
 
     run_env = {
         **os.environ,
@@ -130,7 +89,7 @@ def run_example(example: str):
         "HF_HOME": HF_CACHE_PATH,
     }
     output = run_and_capture(
-        ["cargo", "run", "--release"],
+        ["cargo", "run", "--release", *EXAMPLE_CARGO_ARGS.get(example, [])],
         cwd=f"{WORKDIR}/examples/{example}",
         env=run_env,
     )

crates/luminal_cuda_lite/Cargo.toml

@@ -29,6 +29,7 @@ colorize = "*"
 
 [dev-dependencies]
 candle-core = { version = "0.9.2", features = ["cuda"] }
+luminal_nn = { path = "../luminal_nn" }
 proptest = "1.9.0"
 rand = "0.9.2"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

crates/luminal_cuda_lite/src/host/compute_attn_mask.rs

@@ -1,198 +0,0 @@
-//! ComputeAttnMask — fused op that computes the paged attention mask from indptrs.
-//!
-//! This op exists so the indptr tensors (qo_indptr, kv_indptr) are visible in the
-//! same e-graph chunk as the attention pattern, letting the FlashInfer egglog rule
-//! capture them directly.
-//!
-//! Inputs (3): q_pos (s,) Int, qo_indptr (r,) Int, kv_indptr (r,) Int.
-//! Output: mask (s, c) F32 where mask[i, j] = 0.0 (attend) or -1e10 (block).
-
-use std::sync::Arc;
-
-use luminal::{
-    egglog_utils::{
-        api::{Rule, SortDef, sort},
-        base::{EXPRESSION, OP_KIND},
-        extract_expr,
-    },
-    op::{EgglogOp, HLIROp, LLIROp},
-    prelude::*,
-};
-
-use crate::{
-    cudarc::driver::{CudaStream, result},
-    host::{DeviceBuffer, HostOp},
-};
-
-/// Computes the paged attention mask from indptr arrays.
-///
-/// The mask encodes both request-membership and causality:
-/// `mask[i, j] = 0.0` if query `i` and context `j` belong to the same request AND
-/// context `j`'s local position is `<= q_pos[i]`; `-1e10` otherwise.
-#[derive(Debug, Default)]
-pub struct ComputeAttnMask {
-    pub s_dim: Expression,
-    pub c_dim: Expression,
-}
-
-impl std::fmt::Display for ComputeAttnMask {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "ComputeAttnMask(s={}, c={})", self.s_dim, self.c_dim)
-    }
-}
-
-impl HLIROp for ComputeAttnMask {
-    fn to_egglog(&self, inputs: &[(NodeIndex, String)]) -> String {
-        format!(
-            "(Op (ComputeAttnMask {} {}) (ICons {} (ICons {} (ICons {} (INil)))))",
-            self.s_dim.to_egglog(),
-            self.c_dim.to_egglog(),
-            inputs[0].1, // q_pos
-            inputs[1].1, // qo_indptr
-            inputs[2].1, // kv_indptr
-        )
-    }
-}
-
-impl EgglogOp for ComputeAttnMask {
-    fn sort(&self) -> SortDef {
-        sort(
-            OP_KIND,
-            "ComputeAttnMask",
-            &[("s_dim", EXPRESSION), ("c_dim", EXPRESSION)],
-        )
-    }
-
-    fn n_inputs(&self) -> usize {
-        3
-    }
-
-    fn rewrites(&self) -> Vec<Rule> {
-        // No rewrites — inserted directly by model code.
-        vec![]
-    }
-
-    fn extract<'a>(
-        &'a self,
-        egraph: &'a luminal::egglog_utils::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>) {
-        let s_dim = extract_expr(egraph, kind_children[0], expr_cache).unwrap();
-        let c_dim = extract_expr(egraph, kind_children[1], expr_cache).unwrap();
-        let op = Self { s_dim, c_dim };
-        let llir_op = LLIROp::new::<dyn HostOp>(Box::new(op) as Box<dyn HostOp>);
-        (llir_op, input_enodes)
-    }
-
-    fn cleanup(&self) -> bool {
-        false
-    }
-}
-
-impl HostOp for ComputeAttnMask {
-    fn execute(
-        &self,
-        stream: &Arc<CudaStream>,
-        self_node: NodeIndex,
-        inputs: &[NodeIndex],
-        buffers: &FxHashMap<NodeIndex, DeviceBuffer>,
-        dyn_map: &FxHashMap<char, usize>,
-    ) -> anyhow::Result<()> {
-        if inputs.len() < 3 {
-            anyhow::bail!(
-                "ComputeAttnMask expects 3 inputs (q_pos, qo_indptr, kv_indptr), got {}",
-                inputs.len()
-            );
-        }
-
-        let s = self
-            .s_dim
-            .exec(dyn_map)
-            .ok_or_else(|| anyhow::anyhow!("ComputeAttnMask s_dim unresolved"))?;
-        let c = self
-            .c_dim
-            .exec(dyn_map)
-            .ok_or_else(|| anyhow::anyhow!("ComputeAttnMask c_dim unresolved"))?;
-        let r = *dyn_map
-            .get(&'r')
-            .ok_or_else(|| anyhow::anyhow!("ComputeAttnMask requires dynamic dim 'r'"))?;
-
-        let get_buf = |name: &str, node: NodeIndex| -> anyhow::Result<DeviceBuffer> {
-            buffers.get(&node).copied().ok_or_else(|| {
-                anyhow::anyhow!("ComputeAttnMask missing {name} buffer for {node:?}")
-            })
-        };
-
-        let q_pos_buf = get_buf("q_pos", inputs[0])?;
-        let qo_indptr_buf = get_buf("qo_indptr", inputs[1])?;
-        let kv_indptr_buf = get_buf("kv_indptr", inputs[2])?;
-        let out_buf = get_buf("output", self_node)?;
-
-        let q_pos = dtoh_i32(stream, q_pos_buf.ptr(), s)?;
-        let qo_indptr = dtoh_i32(stream, qo_indptr_buf.ptr(), r)?;
-        let kv_indptr = dtoh_i32(stream, kv_indptr_buf.ptr(), r)?;
-
-        let mut mask = vec![-1e10f32; s * c];
-        for i in 0..s {
-            let q_req = indptr_to_request(&qo_indptr, i as i32);
-            for j in 0..c {
-                let c_req = indptr_to_request(&kv_indptr, j as i32);
-                if q_req == c_req && q_req >= 0 {
-                    let c_local = j as i32 - kv_indptr[c_req as usize];
-                    if c_local <= q_pos[i] {
-                        mask[i * c + j] = 0.0;
-                    }
-                }
-            }
-        }
-
-        let mask_bytes =
-            unsafe { std::slice::from_raw_parts(mask.as_ptr() as *const u8, mask.len() * 4) };
-        unsafe {
-            let res = cudarc::driver::sys::cuMemcpyHtoD_v2(
-                out_buf.ptr(),
-                mask_bytes.as_ptr() as *const std::ffi::c_void,
-                mask_bytes.len(),
-            );
-            if res != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
-                anyhow::bail!("ComputeAttnMask cuMemcpyHtoD failed: {res:?}");
-            }
-        }
-
-        Ok(())
-    }
-
-    fn output_size(&self) -> Expression {
-        self.s_dim * self.c_dim
-    }
-
-    fn output_bytes(&self) -> Expression {
-        self.output_size() * 4
-    }
-
-    fn stats_name(&self) -> Option<&'static str> {
-        Some("ComputeAttnMask")
-    }
-}
-
-fn dtoh_i32(stream: &Arc<CudaStream>, dev_ptr: u64, len: usize) -> anyhow::Result<Vec<i32>> {
-    let mut host = vec![0u8; len * std::mem::size_of::<i32>()];
-    unsafe {
-        result::memcpy_dtoh_async(&mut host, dev_ptr, stream.cu_stream())?;
-    }
-    stream.synchronize()?;
-    let v = unsafe {
-        let mut bytes = std::mem::ManuallyDrop::new(host);
-        Vec::from_raw_parts(bytes.as_mut_ptr() as *mut i32, len, len)
-    };
-    Ok(v)
-}
-
-/// Given an indptr array `[0, a, b, ...]`, find which segment `idx` belongs to.
-/// Returns `count(indptr[i] <= idx) - 1`.
-fn indptr_to_request(indptr: &[i32], idx: i32) -> i32 {
-    indptr.iter().filter(|&&v| v <= idx).count() as i32 - 1
-}

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

@@ -1,258 +0,0 @@
-use std::sync::{Arc, OnceLock};
-
-use luminal::{
-    egglog_utils::{
-        api::{Rule, SortDef, sort},
-        base::{EXPRESSION, OP_KIND, STRING},
-        extract_expr,
-    },
-    op::{EgglogOp, LLIROp},
-    prelude::{
-        tracing::{Level, span, trace},
-        *,
-    },
-};
-
-use crate::{
-    cudarc::{
-        cublas::{
-            CudaBlas,
-            sys::{cublasOperation_t, cublasSetStream_v2, cublasSgemm_v2, cublasStatus_t},
-        },
-        driver::CudaStream,
-    },
-    host::{DeviceBuffer, HostOp},
-};
-
-/// Global shared cuBLAS handle to avoid per-operation workspace allocation
-static SHARED_CUBLAS: OnceLock<Arc<CudaBlas>> = OnceLock::new();
-
-/// Parse cuBLAS operation from egglog string (e.g., "\"T\"" -> CUBLAS_OP_T)
-pub fn parse_cublas_op(s: &str) -> cublasOperation_t {
-    // Strip quotes if present (egglog strings are stored with quotes)
-    let stripped = s.trim_matches('"');
-    match stripped {
-        "T" => cublasOperation_t::CUBLAS_OP_T,
-        "N" => cublasOperation_t::CUBLAS_OP_N,
-        "C" => cublasOperation_t::CUBLAS_OP_C,
-        other => panic!("Unknown cuBLAS operation: '{other}' (original: '{s}')"),
-    }
-}
-
-#[derive(Debug)]
-#[allow(dead_code)]
-pub struct CuBlasSgemmV2 {
-    m: Expression,
-    n: Expression,
-    k: Expression,
-    a_layout: cublasOperation_t,
-    b_layout: cublasOperation_t,
-    lda: Expression,
-    ldb: Expression,
-    ldc: Expression,
-    /// Lazily initialized cuBLAS handle - created on first execute
-    cublas: OnceLock<Arc<CudaBlas>>,
-}
-
-// Useless default for IntoEgglogOp
-impl Default for CuBlasSgemmV2 {
-    fn default() -> Self {
-        Self {
-            m: Expression::default(),
-            n: Expression::default(),
-            k: Expression::default(),
-            a_layout: cublasOperation_t::CUBLAS_OP_N, // IGNORE NOT REAL
-            b_layout: cublasOperation_t::CUBLAS_OP_T, // IGNORE NOT REAL
-            lda: Expression::default(),
-            ldb: Expression::default(),
-            ldc: Expression::default(),
-            cublas: OnceLock::new(),
-        }
-    }
-}
-
-impl EgglogOp for CuBlasSgemmV2 {
-    fn sort(&self) -> SortDef {
-        sort(
-            OP_KIND,
-            "cublasSgemmV2",
-            &[
-                ("m", EXPRESSION),
-                ("n", EXPRESSION),
-                ("k", EXPRESSION),
-                ("a_layout", STRING),
-                ("b_layout", STRING),
-                ("lda", EXPRESSION),
-                ("ldb", EXPRESSION),
-                ("ldc", EXPRESSION),
-            ],
-        )
-    }
-
-    fn n_inputs(&self) -> usize {
-        2
-    }
-
-    fn rewrites(&self) -> Vec<Rule> {
-        vec![
-            Rule::raw(include_str!["sgemm_v2_RmRm_rewrite.egg"]), // row row
-            Rule::raw(include_str!["sgemm_v2_RmCm_rewrite.egg"]), // row col
-            Rule::raw(include_str!["sgemm_v2_CmRm_rewrite.egg"]), // col row
-            Rule::raw(include_str!["sgemm_v2_CmCm_rewrite.egg"]), // col col
-        ]
-    }
-
-    #[allow(unused_variables)]
-    fn extract<'a>(
-        &'a self,
-        egraph: &'a luminal::egglog_utils::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>) {
-        // Extract dimensions from egglog
-        let m = extract_expr(egraph, kind_children[0], expr_cache).unwrap();
-        let n = extract_expr(egraph, kind_children[1], expr_cache).unwrap();
-        let k = extract_expr(egraph, kind_children[2], expr_cache).unwrap();
-
-        // Extract layout strings from egglog
-        let a_layout_str = &egraph.enodes[kind_children[3]].0;
-        let b_layout_str = &egraph.enodes[kind_children[4]].0;
-        let a_layout = parse_cublas_op(a_layout_str);
-        let b_layout = parse_cublas_op(b_layout_str);
-
-        // Extract leading dimensions from egglog
-        let lda = extract_expr(egraph, kind_children[5], expr_cache).unwrap();
-        let ldb = extract_expr(egraph, kind_children[6], expr_cache).unwrap();
-        let ldc = extract_expr(egraph, kind_children[7], expr_cache).unwrap();
-
-        let extracted_state = Self {
-            m,
-            n,
-            k,
-            a_layout,
-            b_layout,
-            lda,
-            ldb,
-            ldc,
-            cublas: OnceLock::new(),
-        };
-        trace!(?extracted_state);
-
-        let extracted = LLIROp::new::<dyn HostOp>(Box::new(extracted_state) as Box<dyn HostOp>);
-
-        (extracted, input_enodes)
-    }
-
-    fn cleanup(&self) -> bool {
-        false
-    }
-}
-
-impl HostOp for CuBlasSgemmV2 {
-    fn execute(
-        &self,
-        stream: &Arc<CudaStream>,
-        self_node: NodeIndex,
-        inputs: &[NodeIndex],
-        buffers: &FxHashMap<NodeIndex, DeviceBuffer>,
-        dyn_map: &FxHashMap<char, usize>,
-    ) -> anyhow::Result<()> {
-        // GEMM parameters
-        let m = self.m.exec(dyn_map).unwrap() as i32;
-        let n = self.n.exec(dyn_map).unwrap() as i32;
-        let k = self.k.exec(dyn_map).unwrap() as i32;
-        let a_layout = self.a_layout;
-        let b_layout = self.b_layout;
-        let lda = self.lda.exec(dyn_map).unwrap() as i32;
-        let ldb = self.ldb.exec(dyn_map).unwrap() as i32;
-        let ldc = self.ldc.exec(dyn_map).unwrap() as i32;
-
-        let alpha = 1.0f32;
-        let beta = 0.0f32;
-
-        // Get buffers: output is self_node, inputs are from graph edges
-        let c_buf = buffers[&self_node];
-        let a_buf = buffers[&inputs[0]];
-        let b_buf = buffers[&inputs[1]];
-
-        // Get device pointers
-        let a_ptr = a_buf.ptr();
-        let b_ptr = b_buf.ptr();
-        let c_ptr = c_buf.ptr();
-
-        // Debug: Check buffer sizes
-        trace!(
-            "buffer_validation {}=={},{}=={},{}=={}",
-            a_buf.len(),
-            m * k * 4,
-            b_buf.len(),
-            k * n * 4,
-            c_buf.len(),
-            m * n * 4
-        );
-        let _sgemm_span = span!(
-            Level::TRACE,
-            "cuBLAS_SGEMM_V2",
-            m,
-            n,
-            k,
-            alpha,
-            beta,
-            lda,
-            ldb,
-            ldc,
-            ?a_layout,
-            ?b_layout,
-        )
-        .entered();
-
-        // Use shared cuBLAS handle to avoid per-operation workspace allocation
-        let cublas = SHARED_CUBLAS.get_or_init(|| Arc::new(CudaBlas::new(stream.clone()).unwrap()));
-
-        // Set the stream for this operation (cuBLAS handle can work with any stream)
-        // The CUstream types from cublas::sys and driver::sys are compatible, just cast
-        unsafe {
-            cublasSetStream_v2(*cublas.handle(), stream.cu_stream() as _);
-        }
-
-        let status = unsafe {
-            cublasSgemm_v2(
-                *cublas.handle(),
-                a_layout,
-                b_layout,
-                m,
-                n,
-                k,
-                &alpha as *const f32,
-                a_ptr as *const f32,
-                lda,
-                b_ptr as *const f32,
-                ldb,
-                &beta as *const f32,
-                c_ptr as *mut f32,
-                ldc,
-            )
-        };
-        stream.synchronize().unwrap();
-
-        if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
-            return Err(anyhow::anyhow!(
-                "cuBLAS SGEMM TN failed with status: {:?}",
-                status
-            ));
-        }
-
-        Ok(())
-    }
-
-    fn output_size(&self) -> Expression {
-        self.m * self.n
-    }
-
-    fn output_bytes(&self) -> Expression {
-        // CuBlasSgemmV2 is F32 only (Sgemm = Single precision)
-        self.output_size() * 4
-    }
-}

crates/luminal_cuda_lite/src/host/cublas/sgemm_v2_CmCm_rewrite.egg

@@ -1,73 +0,0 @@
-; Column-major × Column-major matmul: C[m,n] = A[m,k] × B[k,n]
-; A[m,k] column-major → expand to [m, n, k] with strides [1, 0, m]
-; B[k,n] column-major → expand to [m, n, k] with strides [0, k, 1]
-;
-; Row-major viewed as column-major (swap trick):
-;   Column-major A[m,k] is already column-major with lda=m
-;   Column-major B[k,n] is already column-major with ldb=k
-;   Row-major C[m,n] ≡ column-major C^T[n,m] with ldc=n
-;
-; C^T[n,m] = (A × B)^T = B^T[n,k] × A^T[k,m]
-; cuBLAS: cublasSgemm(OP_T, OP_T, n, m, k, α, B, k, A, m, β, C, n)
-(rule
-    (
-        ; Match Mul node
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-
-        ; Match Sum that reduces the Mul (k dimension)
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
-
-        ; Must be exactly 2D (no batch dims) — batched matmul uses CuBlasLt
-        (= (len ?out_shape) 2)
-
-        ; Get dimensions from output shape
-        (= ?m (nth_from_end ?out_shape 1))
-        (= ?n (nth_from_end ?out_shape 0))
-        (!= ?m (MNum 0))
-        (!= ?n (MNum 0))
-
-        ; Get A strides in [m, n, k] space
-        (= ?a_m_stride (nth_from_end ?a_stride 2))
-        (= ?a_n_stride (nth_from_end ?a_stride 1))
-        (= ?a_k_stride (nth_from_end ?a_stride 0))
-
-        ; Get B strides in [m, n, k] space
-        (= ?b_m_stride (nth_from_end ?b_stride 2))
-        (= ?b_n_stride (nth_from_end ?b_stride 1))
-        (= ?b_k_stride (nth_from_end ?b_stride 0))
-
-        ; Assert contiguous k stride on output (required for reduction)
-        (= ?k_stride (MIter))
-
-        ; Assert A has strides [1, 0, m] (column-major A[m,k] broadcast to [m,n,k])
-        (= ?a_m_stride (MIter))
-        (= ?a_n_stride (MNum 0))
-        (= ?a_k_stride (MMul (MIter) ?m))
-
-        ; Assert B has strides [0, k, 1] (column-major B[k,n] broadcast to [m,n,k])
-        (= ?b_m_stride (MNum 0))
-        (= ?b_n_stride (MMul (MIter) ?k))
-        (= ?b_k_stride (MIter))
-
-        (= (F32) (dtype ?a))
-        (= (F32) (dtype ?b))
-    )
-    (
-        ; For column-major A × column-major B with cuBLAS:
-        ; C^T = B^T × A^T → cublasSgemm(OP_T, OP_T, n, m, k, α, B, k, A, m, β, C, n)
-        (let ?sgemm (Op (cublasSgemmV2
-            ?n             ; cuBLAS m = our n (swapped)
-            ?m             ; cuBLAS n = our m (swapped)
-            ?k             ; k unchanged
-            "T"            ; transa = Transpose (B is column-major [k,n], need B^T[n,k])
-            "T"            ; transb = Transpose (A is column-major [m,k], need A^T[k,m])
-            ?k             ; lda = k (column-major B[k,n])
-            ?m             ; ldb = m (column-major A[m,k])
-            ?n)            ; ldc = n (row-major C[m,n] viewed as col-major [n,m])
-            (ICons ?b (ICons ?a (INil)))))
-        (union ?sum ?sgemm)
-        (set (dtype ?sgemm) (F32))
-    )
-    :ruleset matmul_backend
-    :name "cublas sgemm column-major × column-major"
-)

crates/luminal_cuda_lite/src/host/cublas/sgemm_v2_CmRm_rewrite.egg

@@ -1,73 +0,0 @@
-; Column-major × Row-major matmul: C[m,n] = A[m,k] × B[k,n]
-; A[m,k] column-major → expand to [m, n, k] with strides [1, 0, m]
-; B[k,n] row-major → permute to [n,k] then expand to [m, n, k] with strides [0, 1, n]
-;
-; Row-major viewed as column-major (swap trick):
-;   Column-major A[m,k] is already column-major with lda=m
-;   Row-major B[k,n] ≡ column-major B^T[n,k] with ldb=n
-;   Row-major C[m,n] ≡ column-major C^T[n,m] with ldc=n
-;
-; C^T[n,m] = (A × B)^T = B^T[n,k] × A^T[k,m]
-; cuBLAS: cublasSgemm(OP_N, OP_T, n, m, k, α, B, n, A, m, β, C, n)
-(rule
-    (
-        ; Match Mul node
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-
-        ; Match Sum that reduces the Mul (k dimension)
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
-
-        ; Must be exactly 2D (no batch dims) — batched matmul uses CuBlasLt
-        (= (len ?out_shape) 2)
-
-        ; Get dimensions from output shape
-        (= ?m (nth_from_end ?out_shape 1))
-        (= ?n (nth_from_end ?out_shape 0))
-        (!= ?m (MNum 0))
-        (!= ?n (MNum 0))
-
-        ; Get A strides in [m, n, k] space
-        (= ?a_m_stride (nth_from_end ?a_stride 2))
-        (= ?a_n_stride (nth_from_end ?a_stride 1))
-        (= ?a_k_stride (nth_from_end ?a_stride 0))
-
-        ; Get B strides in [m, n, k] space
-        (= ?b_m_stride (nth_from_end ?b_stride 2))
-        (= ?b_n_stride (nth_from_end ?b_stride 1))
-        (= ?b_k_stride (nth_from_end ?b_stride 0))
-
-        ; Assert contiguous k stride on output (required for reduction)
-        (= ?k_stride (MIter))
-
-        ; Assert A has strides [1, 0, m] (column-major A[m,k] broadcast to [m,n,k])
-        (= ?a_m_stride (MIter))
-        (= ?a_n_stride (MNum 0))
-        (= ?a_k_stride (MMul (MIter) ?m))
-
-        ; Assert B has strides [0, 1, n] (row-major B[k,n] permuted to [n,k] then broadcast to [m,n,k])
-        (= ?b_m_stride (MNum 0))
-        (= ?b_n_stride (MIter))
-        (= ?b_k_stride (MMul (MIter) ?n))
-
-        (= (F32) (dtype ?a))
-        (= (F32) (dtype ?b))
-    )
-    (
-        ; For column-major A × row-major B with cuBLAS:
-        ; C^T = B^T × A^T → cublasSgemm(OP_N, OP_T, n, m, k, α, B, n, A, m, β, C, n)
-        (let ?sgemm (Op (cublasSgemmV2
-            ?n             ; cuBLAS m = our n (swapped)
-            ?m             ; cuBLAS n = our m (swapped)
-            ?k             ; k unchanged
-            "N"            ; transa = No transpose (B is row-major, viewed as col-major [n,k])
-            "T"            ; transb = Transpose (A is column-major [m,k], need A^T[k,m])
-            ?n             ; lda = n (row-major B[k,n] viewed as col-major [n,k])
-            ?m             ; ldb = m (column-major A[m,k])
-            ?n)            ; ldc = n (row-major C[m,n] viewed as col-major [n,m])
-            (ICons ?b (ICons ?a (INil)))))
-        (union ?sum ?sgemm)
-        (set (dtype ?sgemm) (F32))
-    )
-    :ruleset matmul_backend
-    :name "cublas sgemm column-major × row-major"
-)

crates/luminal_cuda_lite/src/host/cublas/sgemm_v2_RmCm_rewrite.egg

@@ -1,73 +0,0 @@
-; Row-major × Column-major matmul: C[m,n] = A[m,k] × B[k,n]
-; A[m,k] row-major → expand to [m, n, k] with strides [k, 0, 1]
-; B[k,n] column-major → expand to [m, n, k] with strides [0, k, 1]
-;
-; Row-major viewed as column-major (swap trick):
-;   Row-major A[m,k] ≡ column-major A^T[k,m] with lda=k
-;   Column-major B[k,n] is already column-major with ldb=k
-;   Row-major C[m,n] ≡ column-major C^T[n,m] with ldc=n
-;
-; C^T[n,m] = (A × B)^T = B^T[n,k] × A^T[k,m]
-; cuBLAS: cublasSgemm(OP_T, OP_N, n, m, k, α, B, k, A, k, β, C, n)
-(rule
-    (
-        ; Match Mul node
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-
-        ; Match Sum that reduces the Mul (k dimension)
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
-
-        ; Must be exactly 2D (no batch dims) — batched matmul uses CuBlasLt
-        (= (len ?out_shape) 2)
-
-        ; Get dimensions from output shape
-        (= ?m (nth_from_end ?out_shape 1))
-        (= ?n (nth_from_end ?out_shape 0))
-        (!= ?m (MNum 0))
-        (!= ?n (MNum 0))
-
-        ; Get A strides in [m, n, k] space
-        (= ?a_m_stride (nth_from_end ?a_stride 2))
-        (= ?a_n_stride (nth_from_end ?a_stride 1))
-        (= ?a_k_stride (nth_from_end ?a_stride 0))
-
-        ; Get B strides in [m, n, k] space
-        (= ?b_m_stride (nth_from_end ?b_stride 2))
-        (= ?b_n_stride (nth_from_end ?b_stride 1))
-        (= ?b_k_stride (nth_from_end ?b_stride 0))
-
-        ; Assert contiguous k stride on output (required for reduction)
-        (= ?k_stride (MIter))
-
-        ; Assert A has strides [k, 0, 1] (row-major A[m,k] broadcast to [m,n,k])
-        (= ?a_m_stride (MMul (MIter) ?k))
-        (= ?a_n_stride (MNum 0))
-        (= ?a_k_stride (MIter))
-
-        ; Assert B has strides [0, k, 1] (column-major B[k,n] broadcast to [m,n,k])
-        (= ?b_m_stride (MNum 0))
-        (= ?b_n_stride (MMul (MIter) ?k))
-        (= ?b_k_stride (MIter))
-
-        (= (F32) (dtype ?a))
-        (= (F32) (dtype ?b))
-    )
-    (
-        ; For row-major A × column-major B with cuBLAS:
-        ; C^T = B^T × A^T → cublasSgemm(OP_T, OP_N, n, m, k, α, B, k, A, k, β, C, n)
-        (let ?sgemm (Op (cublasSgemmV2
-            ?n             ; cuBLAS m = our n (swapped)
-            ?m             ; cuBLAS n = our m (swapped)
-            ?k             ; k unchanged
-            "T"            ; transa = Transpose (B is column-major, need B^T)
-            "N"            ; transb = No transpose
-            ?k             ; lda = k (column-major B[k,n])
-            ?k             ; ldb = k (row-major A[m,k] viewed as col-major [k,m])
-            ?n)            ; ldc = n (row-major C[m,n] viewed as col-major [n,m])
-            (ICons ?b (ICons ?a (INil)))))
-        (union ?sum ?sgemm)
-        (set (dtype ?sgemm) (F32))
-    )
-    :ruleset matmul_backend
-    :name "cublas sgemm row-major × column-major"
-)

crates/luminal_cuda_lite/src/host/cublas/sgemm_v2_RmRm_rewrite.egg

@@ -1,73 +0,0 @@
-; Row-major matmul: C[m,n] = A[m,k] × B[k,n]
-; A[m,k] row-major → expand to [m, n, k] with strides [k, 0, 1]
-; B[k,n] row-major → permute to [n,k] then expand to [m, n, k] with strides [0, 1, n]
-;
-; Row-major viewed as column-major (swap trick):
-;   Row-major A[m,k] ≡ column-major [k,m] with lda=k
-;   Row-major B[k,n] ≡ column-major [n,k] with ldb=n
-;   Row-major C[m,n] ≡ column-major [n,m] with ldc=n
-;
-; cuBLAS computes: C_col[n,m] = B_col[n,k] × A_col[k,m]
-; cublasSgemm(OP_N, OP_N, n, m, k, α, B, n, A, k, β, C, n)
-(rule
-    (
-        ; Match Mul node
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-
-        ; Match Sum that reduces the Mul (k dimension)
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
-
-        ; Must be exactly 2D (no batch dims) — batched matmul uses CuBlasLt
-        (= (len ?out_shape) 2)
-
-        ; Get dimensions from output shape
-        (= ?m (nth_from_end ?out_shape 1))
-        (= ?n (nth_from_end ?out_shape 0))
-        (!= ?m (MNum 0))
-        (!= ?n (MNum 0))
-
-        ; Get A strides in [m, n, k] space
-        (= ?a_m_stride (nth_from_end ?a_stride 2))
-        (= ?a_n_stride (nth_from_end ?a_stride 1))
-        (= ?a_k_stride (nth_from_end ?a_stride 0))
-
-        ; Get B strides in [m, n, k] space
-        (= ?b_m_stride (nth_from_end ?b_stride 2))
-        (= ?b_n_stride (nth_from_end ?b_stride 1))
-        (= ?b_k_stride (nth_from_end ?b_stride 0))
-
-        ; Assert contiguous k stride on output (required for reduction)
-        (= ?k_stride (MIter))
-
-        ; Assert A has strides [k, 0, 1] (row-major A[m,k] broadcast to [m,n,k])
-        (= ?a_m_stride (MMul (MIter) ?k))
-        (= ?a_n_stride (MNum 0))
-        (= ?a_k_stride (MIter))
-
-        ; Assert B has strides [0, 1, n] (row-major B[k,n] permuted to [n,k] then broadcast to [m,n,k])
-        (= ?b_m_stride (MNum 0))
-        (= ?b_n_stride (MIter))
-        (= ?b_k_stride (MMul (MIter) ?n))
-
-        (= (F32) (dtype ?a))
-        (= (F32) (dtype ?b))
-    )
-    (
-        ; For row-major C = A × B with cuBLAS (column-major):
-        ; cublasSgemm(OP_N, OP_N, n, m, k, α, B, n, A, k, β, C, n)
-        (let ?sgemm (Op (cublasSgemmV2
-            ?n             ; cuBLAS m = our n (swapped)
-            ?m             ; cuBLAS n = our m (swapped)
-            ?k             ; k unchanged
-            "N"            ; transa = No transpose
-            "N"            ; transb = No transpose
-            ?n             ; lda = n (row-major B[k,n] viewed as col-major [n,k])
-            ?k             ; ldb = k (row-major A[m,k] viewed as col-major [k,m])
-            ?n)            ; ldc = n (row-major C[m,n] viewed as col-major [n,m])
-            (ICons ?b (ICons ?a (INil)))))
-        (union ?sum ?sgemm)
-        (set (dtype ?sgemm) (F32))
-    )
-    :ruleset matmul_backend
-    :name "cublas sgemm row-major"
-)

crates/luminal_cuda_lite/src/host/cublaslt/cublaslt_CmCm_rewrite.egg

@@ -11,11 +11,13 @@
 ; cuBLAS: cublasSgemm(OP_T, OP_T, n, m, k, α, B, k, A, m, β, C, n)
 (rule
     (
-        ; Match Mul node
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-
-        ; Match Sum that reduces the Mul (k dimension)
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        ; Match the generic matmul produced from Mul -> Sum.
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         ; Match exactly 2D output shape
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
@@ -77,8 +79,12 @@
 ; B column-major per batch: b_k_stride=MIter, b_m_stride=0
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?batch (nth_from_end ?out_shape 2))
         (= ?m (nth_from_end ?out_shape 1))

crates/luminal_cuda_lite/src/host/cublaslt/cublaslt_CmRm_rewrite.egg

@@ -11,11 +11,13 @@
 ; cuBLAS: cublasSgemm(OP_N, OP_T, n, m, k, α, B, n, A, m, β, C, n)
 (rule
     (
-        ; Match Mul node
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-
-        ; Match Sum that reduces the Mul (k dimension)
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        ; Match the generic matmul produced from Mul -> Sum.
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         ; Match exactly 2D output shape
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
@@ -77,8 +79,12 @@
 ; B row-major per batch: b_n_stride=MIter, b_m_stride=0
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?batch (nth_from_end ?out_shape 2))
         (= ?m (nth_from_end ?out_shape 1))

crates/luminal_cuda_lite/src/host/cublaslt/cublaslt_RmCm_rewrite.egg

@@ -11,11 +11,13 @@
 ; cuBLAS: cublasSgemm(OP_T, OP_N, n, m, k, α, B, k, A, k, β, C, n)
 (rule
     (
-        ; Match Mul node
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-
-        ; Match Sum that reduces the Mul (k dimension)
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        ; Match the generic matmul produced from Mul -> Sum.
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         ; Match exactly 2D output shape
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
@@ -77,8 +79,12 @@
 ; B column-major per batch: b_k_stride=MIter, b_m_stride=0
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?batch (nth_from_end ?out_shape 2))
         (= ?m (nth_from_end ?out_shape 1))

crates/luminal_cuda_lite/src/host/cublaslt/cublaslt_RmRm_rewrite.egg

@@ -11,11 +11,13 @@
 ; cublasSgemm(OP_N, OP_N, n, m, k, α, B, n, A, k, β, C, n)
 (rule
     (
-        ; Match Mul node
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-
-        ; Match Sum that reduces the Mul (k dimension)
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        ; Match the generic matmul produced from Mul -> Sum.
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         ; Match exactly 2D output shape
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
@@ -79,8 +81,12 @@
 ; Leading dimensions may differ from k/n when batch slices are non-contiguous.
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         ; Output shape: [batch, m, n]
         (= ?batch (nth_from_end ?out_shape 2))

crates/luminal_cuda_lite/src/host/cublaslt/cublaslt_fp8_rewrite.egg

@@ -10,8 +10,454 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        ; Match the scaled FP8 linear form directly before the unscaled FP8
+        ; matmul rewrite can hide the quantize/dequant scale structure.
+        (= ?scaled_activation (Op (Mul
+            ?activation_shape
+            ?raw_activation_strides
+            ?recip_activation_strides
+            ?activation_out_strides)
+            (ICons ?raw_activation (ICons ?recip_input_scale (INil)))))
+        (= ?recip_input_scale (Op (Recip
+            ?activation_shape
+            (ECons (MNum 0) (ECons (MNum 0) (ENil)))
+            ?recip_out_strides)
+            (ICons ?input_scale (INil))))
+        (= ?a (Op (Cast ?a_size ?a_dtype) (ICons ?scaled_activation (INil))))
+
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
+        (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
+        (= ?scale_product (Op (Mul (ENil) (ENil) (ENil) (ENil))
+            (ICons ?input_scale (ICons ?weight_scale (INil)))))
+        (= ?scaled (Op (Mul
+            ?out_shape
+            ?cast_strides
+            (ECons (MNum 0) (ECons (MNum 0) (ENil)))
+            ?scaled_out_strides)
+            (ICons ?cast (ICons ?scale_product (INil)))))
+        (= ?cast_strides ?scaled_out_strides)
+
+        (= ?out_shape (ECons ?m (ECons ?n (ENil))))
+        (!= ?m (MNum 0))
+        (!= ?n (MNum 0))
+        (!= ?k (MNum 1))
+
+        (= ?a_stride (ECons ?a_m_stride (ECons ?a_n_stride (ECons ?a_k_stride (ENil)))))
+        (= ?b_stride (ECons ?b_m_stride (ECons ?b_n_stride (ECons ?b_k_stride (ENil)))))
+        (= ?k_stride (MIter))
+
+        (= ?a_m_stride (MMul (MIter) ?k))
+        (= ?a_n_stride (MNum 0))
+        (= ?a_k_stride (MIter))
+
+        (= ?b_m_stride (MNum 0))
+        (= ?b_n_stride (MMul (MIter) ?k))
+        (= ?b_k_stride (MIter))
+
+        (= ?b_dtype (dtype ?b))
+        (cublaslt_fp8_f32_output_pair ?a_dtype ?b_dtype)
+    )
+    (
+        (let ?sgemm (Op (cublaslt_scaled
+            ?n ?m ?k
+            "T" "N"
+            "COL" "COL" "COL" "COL"
+            ?b_n_stride
+            ?a_m_stride
+            ?n
+            ?n
+            (MNum 1)
+            (MNum 0)
+            (MNum 0)
+            (MNum 0)
+            (MNum 0)
+            ?b_dtype ?a_dtype (F32) (F32) "32F" "F32" 1.0 0.0 "DEFAULT")
+            (ICons ?b (ICons ?a (ICons ?weight_scale (ICons ?input_scale (INil)))))))
+        (union ?scaled ?sgemm)
+        (set (dtype ?sgemm) (F32))
+    )
+:ruleset matmul_backend
+:name "cublaslt scaled fp8 row-major x column-major f32 output"
+)
+
+(rule
+    (
+        (= ?scaled_activation (Op (Mul
+            ?activation_shape
+            ?raw_activation_strides
+            ?recip_activation_strides
+            ?activation_out_strides)
+            (ICons ?raw_activation (ICons ?recip_input_scale (INil)))))
+        (= ?recip_input_scale (Op (Recip
+            ?activation_shape
+            (ECons (MNum 0) (ECons (MNum 0) (ENil)))
+            ?recip_out_strides)
+            (ICons ?input_scale (INil))))
+        (= ?a (Op (Cast ?a_size ?a_dtype) (ICons ?scaled_activation (INil))))
+
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
+        (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
+        (= ?scale_product (Op (Mul (ENil) (ENil) (ENil) (ENil))
+            (ICons ?input_scale (ICons ?weight_scale (INil)))))
+        (= ?scaled (Op (Mul
+            ?out_shape
+            ?cast_strides
+            (ECons (MNum 0) (ECons (MNum 0) (ENil)))
+            ?scaled_out_strides)
+            (ICons ?cast (ICons ?scale_product (INil)))))
+        (= ?cast_strides ?scaled_out_strides)
+
+        (= ?out_shape (ECons ?m (ECons ?n (ENil))))
+        (!= ?m (MNum 0))
+        (!= ?n (MNum 0))
+        (!= ?k (MNum 1))
+
+        (= ?a_stride (ECons ?a_m_stride (ECons ?a_n_stride (ECons ?a_k_stride (ENil)))))
+        (= ?b_stride (ECons ?b_m_stride (ECons ?b_n_stride (ECons ?b_k_stride (ENil)))))
+        (= ?k_stride (MIter))
+
+        (= ?a_m_stride (MMul (MIter) ?k))
+        (= ?a_n_stride (MNum 0))
+        (= ?a_k_stride (MIter))
+
+        (= ?b_m_stride (MNum 0))
+        (= ?b_n_stride (MMul (MIter) ?k))
+        (= ?b_k_stride (MIter))
+
+        (= ?b_dtype (dtype ?b))
+        (cublaslt_fp8_f32_output_pair ?a_dtype ?b_dtype)
+        (= ?scaled (Op (cublaslt_scaled
+            ?n ?m ?k
+            "T" "N"
+            "COL" "COL" "COL" "COL"
+            ?b_n_stride
+            ?a_m_stride
+            ?n
+            ?n
+            (MNum 1)
+            (MNum 0)
+            (MNum 0)
+            (MNum 0)
+            (MNum 0)
+            ?b_dtype ?a_dtype (F32) (F32) "32F" "F32" 1.0 0.0 "DEFAULT")
+            (ICons ?b (ICons ?a (ICons ?weight_scale (ICons ?input_scale (INil)))))))
+        (= ?cast (Op (cublaslt
+            ?n ?m ?k
+            "T" "N"
+            "COL" "COL" "COL" "COL"
+            ?b_n_stride
+            ?a_m_stride
+            ?n
+            ?n
+            (MNum 1)
+            (MNum 0)
+            (MNum 0)
+            (MNum 0)
+            (MNum 0)
+            ?b_dtype ?a_dtype (F32) (F32) "32F" "F32" 1.0 0.0 "DEFAULT")
+            (ICons ?b (ICons ?a (INil)))))
+    )
+    (
+        (delete (Op (Mul
+            ?out_shape
+            ?cast_strides
+            (ECons (MNum 0) (ECons (MNum 0) (ENil)))
+            ?scaled_out_strides)
+            (ICons ?cast (ICons ?scale_product (INil)))))
+        (delete (Op (cublaslt
+            ?n ?m ?k
+            "T" "N"
+            "COL" "COL" "COL" "COL"
+            ?b_n_stride
+            ?a_m_stride
+            ?n
+            ?n
+            (MNum 1)
+            (MNum 0)
+            (MNum 0)
+            (MNum 0)
+            (MNum 0)
+            ?b_dtype ?a_dtype (F32) (F32) "32F" "F32" 1.0 0.0 "DEFAULT")
+            (ICons ?b (ICons ?a (INil)))))
+    )
+    :ruleset cleanup
+    :name "delete raw fp8 path when scaled cublaslt covers direct output scale"
+)
+
+(rule
+    (
+        ; Fusion growth can make the live path consume a raw FP8 cuBLASLt
+        ; candidate through an internal CudaBinaryElementwise scale multiply,
+        ; instead of the original HLIR output-scale Mul. The scalar scale
+        ; product is tensor-wide, so the two scalar factors can be passed as
+        ; cuBLASLt A/B scale inputs and the internal multiply can be bypassed.
+        (= ?raw_gemm (Op (cublaslt
+            ?cm ?cn ?ck
+            ?cta ?ctb
+            ?cao ?cbo ?cco ?cdo
+            ?clda ?cldb ?cldc ?cldd
+            ?cbc ?csa ?csb ?csc ?csd
+            ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue)
+            (ICons ?a (ICons ?b (INil)))))
+        (cublaslt_fp8_f32_output_pair ?cadt ?cbdt)
+        (= ?ccdt (F32))
+        (= ?cddt (F32))
+        (= ?cbeta 0.0)
+        (= ?cepilogue "DEFAULT")
+
+        (= ?fs_cast (Op (FusionStart
+            ?out_shape
+            ?cast_strides
+            (F32))
+            (ICons ?raw_gemm (INil))))
+
+        (= ?out_shape (ECons ?out_m (ECons ?out_n (ENil))))
+        (= ?scale_strides (ECons (MNum 0) (ECons (MNum 0) (ENil))))
+
+        (= ?fs_a_scale (Op (FusionStart (ENil) (ENil) (F32))
+            (ICons ?a_scale (INil))))
+        (= ?fs_b_scale (Op (FusionStart (ENil) (ENil) (F32))
+            (ICons ?b_scale (INil))))
+        (= ?scale_product_inner (Op (CudaBinaryElementwise
+            "Mul"
+            (ENil)
+            (ENil)
+            (ENil)
+            (ENil)
+            (F32))
+            (ICons ?fs_a_scale (ICons ?fs_b_scale (INil)))))
+        (= ?scale_product (Op (FusionEnd (ENil) (ENil) (F32))
+            (ICons ?scale_product_inner (INil))))
+        (= ?fs_scale (Op (FusionStart
+            ?out_shape
+            ?scale_strides
+            (F32))
+            (ICons ?scale_product (INil))))
+        (= ?fused_scale (Op (CudaBinaryElementwise
+            "Mul"
+            ?out_shape
+            ?cast_strides
+            ?scale_strides
+            ?scaled_out_strides
+            (F32))
+            (ICons ?fs_cast (ICons ?fs_scale (INil)))))
+        (= ?cast_strides ?scaled_out_strides)
+    )
+    (
+        (let ?sgemm (Op (cublaslt_scaled
+            ?cm ?cn ?ck
+            ?cta ?ctb
+            ?cao ?cbo ?cco ?cdo
+            ?clda ?cldb ?cldc ?cldd
+            ?cbc ?csa ?csb ?csc ?csd
+            ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue)
+            (ICons ?a (ICons ?b (ICons ?a_scale (ICons ?b_scale (INil)))))))
+        (let ?fs_sgemm (Op (FusionStart ?out_shape ?scaled_out_strides (F32))
+            (ICons ?sgemm (INil))))
+        (union ?fused_scale ?fs_sgemm)
+        (set (dtype ?sgemm) (F32))
+        (set (dtype ?fs_sgemm) (F32))
+    )
+    :ruleset fusion_grow
+    :name "cublaslt scaled fp8 fused output-scale f32 output"
+)
+
+(rule
+    (
+        (= ?raw_gemm (Op (cublaslt
+            ?cm ?cn ?ck
+            ?cta ?ctb
+            ?cao ?cbo ?cco ?cdo
+            ?clda ?cldb ?cldc ?cldd
+            ?cbc ?csa ?csb ?csc ?csd
+            ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue)
+            (ICons ?a (ICons ?b (INil)))))
+        (cublaslt_fp8_f32_output_pair ?cadt ?cbdt)
+        (= ?ccdt (F32))
+        (= ?cddt (F32))
+        (= ?cbeta 0.0)
+        (= ?cepilogue "DEFAULT")
+
+        (= ?fs_cast (Op (FusionStart
+            ?out_shape
+            ?cast_strides
+            (F32))
+            (ICons ?raw_gemm (INil))))
+
+        (= ?out_shape (ECons ?out_m (ECons ?out_n (ENil))))
+        (= ?scale_strides (ECons (MNum 0) (ECons (MNum 0) (ENil))))
+
+        (= ?fs_a_scale (Op (FusionStart (ENil) (ENil) (F32))
+            (ICons ?a_scale (INil))))
+        (= ?fs_b_scale (Op (FusionStart (ENil) (ENil) (F32))
+            (ICons ?b_scale (INil))))
+        (= ?scale_product_inner (Op (CudaBinaryElementwise
+            "Mul"
+            (ENil)
+            (ENil)
+            (ENil)
+            (ENil)
+            (F32))
+            (ICons ?fs_a_scale (ICons ?fs_b_scale (INil)))))
+        (= ?scale_product (Op (FusionEnd (ENil) (ENil) (F32))
+            (ICons ?scale_product_inner (INil))))
+        (= ?fs_scale (Op (FusionStart
+            ?out_shape
+            ?scale_strides
+            (F32))
+            (ICons ?scale_product (INil))))
+        (= ?fused_scale (Op (CudaBinaryElementwise
+            "Mul"
+            ?out_shape
+            ?cast_strides
+            ?scale_strides
+            ?scaled_out_strides
+            (F32))
+            (ICons ?fs_cast (ICons ?fs_scale (INil)))))
+        (= ?cast_strides ?scaled_out_strides)
+
+        (= ?sgemm (Op (cublaslt_scaled
+            ?cm ?cn ?ck
+            ?cta ?ctb
+            ?cao ?cbo ?cco ?cdo
+            ?clda ?cldb ?cldc ?cldd
+            ?cbc ?csa ?csb ?csc ?csd
+            ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue)
+            (ICons ?a (ICons ?b (ICons ?a_scale (ICons ?b_scale (INil)))))))
+        (= ?fused_scale (Op (FusionStart ?out_shape ?scaled_out_strides (F32))
+            (ICons ?sgemm (INil))))
+    )
+    (
+        (delete (Op (cublaslt
+            ?cm ?cn ?ck
+            ?cta ?ctb
+            ?cao ?cbo ?cco ?cdo
+            ?clda ?cldb ?cldc ?cldd
+            ?cbc ?csa ?csb ?csc ?csd
+            ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue)
+            (ICons ?a (ICons ?b (INil)))))
+        (delete (Op (CudaBinaryElementwise
+            "Mul"
+            ?out_shape
+            ?cast_strides
+            ?scale_strides
+            ?scaled_out_strides
+            (F32))
+            (ICons ?fs_cast (ICons ?fs_scale (INil)))))
+    )
+    :ruleset cleanup
+    :name "delete raw fp8 path when scaled cublaslt covers fused output scale"
+)
+
+(rule
+    (
+        ; Batched form of the scaled FP8 linear rewrite. The scale operands are
+        ; scalar tensors expanded across the last three output/activation axes.
+        (= ?scaled_activation (Op (Mul
+            ?activation_shape
+            ?raw_activation_strides
+            ?recip_activation_strides
+            ?activation_out_strides)
+            (ICons ?raw_activation (ICons ?recip_input_scale (INil)))))
+        (= ?recip_input_scale (Op (Recip
+            ?activation_shape
+            (ECons (MNum 0) (ECons (MNum 0) (ECons (MNum 0) (ENil))))
+            ?recip_out_strides)
+            (ICons ?input_scale (INil))))
+        (= ?a (Op (Cast ?a_size ?a_dtype) (ICons ?scaled_activation (INil))))
+
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
+        (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
+        (= ?scale_product (Op (Mul (ENil) (ENil) (ENil) (ENil))
+            (ICons ?input_scale (ICons ?weight_scale (INil)))))
+        (= ?scaled (Op (Mul
+            ?out_shape
+            ?cast_strides
+            (ECons (MNum 0) (ECons (MNum 0) (ECons (MNum 0) (ENil))))
+            ?scaled_out_strides)
+            (ICons ?cast (ICons ?scale_product (INil)))))
+        (= ?cast_strides ?scaled_out_strides)
+
+        (= ?batch (nth_from_end ?out_shape 2))
+        (= ?m (nth_from_end ?out_shape 1))
+        (= ?n (nth_from_end ?out_shape 0))
+        (!= ?m (MNum 0))
+        (!= ?n (MNum 0))
+        (!= ?k (MNum 1))
+        (!= ?batch (MNum 0))
+
+        (= ?a_batch_stride (nth_from_end ?a_stride 3))
+        (= ?a_m_stride (nth_from_end ?a_stride 2))
+        (= ?a_n_stride (nth_from_end ?a_stride 1))
+        (= ?a_k_stride (nth_from_end ?a_stride 0))
+
+        (= ?b_batch_stride (nth_from_end ?b_stride 3))
+        (= ?b_m_stride (nth_from_end ?b_stride 2))
+        (= ?b_n_stride (nth_from_end ?b_stride 1))
+        (= ?b_k_stride (nth_from_end ?b_stride 0))
+
+        (= ?k_stride (MIter))
+
+        (= ?a_k_stride (MIter))
+        (= ?a_n_stride (MNum 0))
+        (= ?a_m_stride (MMul (MIter) ?k))
+
+        (= ?b_k_stride (MIter))
+        (= ?b_m_stride (MNum 0))
+        (= ?b_n_stride (MMul (MIter) ?k))
+
+        (= ?a_batch_stride (MMul ?m ?a_m_stride))
+        (= ?b_batch_stride (MMul ?n ?b_n_stride))
+
+        (= ?b_dtype (dtype ?b))
+        (cublaslt_fp8_f32_output_pair ?a_dtype ?b_dtype)
+    )
+    (
+        (let ?sgemm (Op (cublaslt_scaled
+            ?n ?m ?k
+            "T" "N"
+            "COL" "COL" "COL" "COL"
+            ?b_n_stride
+            ?a_m_stride
+            ?n
+            ?n
+            ?batch
+            ?b_batch_stride
+            ?a_batch_stride
+            (MMul ?m ?n)
+            (MMul ?m ?n)
+            ?b_dtype ?a_dtype (F32) (F32) "32F" "F32" 1.0 0.0 "DEFAULT")
+            (ICons ?b (ICons ?a (ICons ?weight_scale (ICons ?input_scale (INil)))))))
+        (union ?scaled ?sgemm)
+        (set (dtype ?sgemm) (F32))
+    )
+:ruleset matmul_backend
+:name "cublaslt scaled fp8 batched row-major x column-major f32 output"
+)
+
+(rule
+    (
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
         (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
 
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
@@ -59,8 +505,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
         (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
 
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
@@ -108,8 +558,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
         (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
 
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
@@ -157,8 +611,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
         (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
 
         (= ?batch (nth_from_end ?out_shape 2))
@@ -220,8 +678,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
         (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
 
         (= ?batch (nth_from_end ?out_shape 2))
@@ -283,8 +745,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
         (= ?cast (Op (Cast ?size (F32)) (ICons ?sum (INil))))
 
         (= ?batch (nth_from_end ?out_shape 2))

crates/luminal_cuda_lite/src/host/cublaslt/cublaslt_row_order_rewrite.egg

@@ -5,8 +5,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
         (!= ?m (MNum 0))
@@ -54,8 +58,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
         (!= ?m (MNum 0))
@@ -103,8 +111,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
         (!= ?m (MNum 0))
@@ -152,8 +164,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?out_shape (ECons ?m (ECons ?n (ENil))))
         (!= ?m (MNum 0))
@@ -201,8 +217,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?batch (nth_from_end ?out_shape 2))
         (= ?m (nth_from_end ?out_shape 1))
@@ -264,8 +284,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?batch (nth_from_end ?out_shape 2))
         (= ?m (nth_from_end ?out_shape 1))
@@ -327,8 +351,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?batch (nth_from_end ?out_shape 2))
         (= ?m (nth_from_end ?out_shape 1))
@@ -390,8 +418,12 @@
 
 (rule
     (
-        (= ?mul (Op (Mul ?mul_shape ?a_stride ?b_stride ?mul_out_stride) (ICons ?a (ICons ?b (INil)))))
-        (= ?sum (Op (Sum ?out_shape ?k ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+        (= ?sum (Op (GenericMatmul
+            ?out_shape ?mul_shape ?k
+            ?a_stride ?b_stride
+            ?sum_in_stride ?k_stride ?sum_out_stride
+            ?matmul_dtype)
+            (ICons ?a (ICons ?b (INil)))))
 
         (= ?batch (nth_from_end ?out_shape 2))
         (= ?m (nth_from_end ?out_shape 1))

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

@@ -35,10 +35,20 @@ use crate::{
         },
         driver::{CudaStream, DevicePtr},
     },
-    host::{DeviceBuffer, HostOp, cublas::parse_cublas_op},
+    host::{DeviceBuffer, HostOp},
     try_create_cublaslt,
 };
 
+fn parse_cublas_op(s: &str) -> cublasOperation_t {
+    let stripped = s.trim_matches('"');
+    match stripped {
+        "T" => cublasOperation_t::CUBLAS_OP_T,
+        "N" => cublasOperation_t::CUBLAS_OP_N,
+        "C" => cublasOperation_t::CUBLAS_OP_C,
+        other => panic!("Unknown cuBLAS operation: '{other}' (original: '{s}')"),
+    }
+}
+
 #[derive(Debug)]
 #[allow(dead_code)]
 pub struct CuBlasLt {
@@ -69,6 +79,8 @@ pub struct CuBlasLt {
     alpha: f64,
     beta: f64,
     epilogue: cublasLtEpilogue_t,
+    a_scale_input: bool,
+    b_scale_input: bool,
     cublaslt: OnceLock<Arc<CudaBlasLT>>,
 }
 
@@ -103,52 +115,62 @@ impl Default for CuBlasLt {
             alpha: 1.0,
             beta: 0.0,
             epilogue: cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DEFAULT,
+            a_scale_input: false,
+            b_scale_input: false,
             cublaslt: OnceLock::new(),
         }
     }
 }
 
+#[derive(Debug, Default)]
+pub struct CuBlasLtScaled;
+
+fn cublaslt_sort(name: &'static str) -> SortDef {
+    sort(
+        OP_KIND,
+        name,
+        &[
+            ("m", EXPRESSION),
+            ("n", EXPRESSION),
+            ("k", EXPRESSION),
+            ("a_layout", STRING),
+            ("b_layout", STRING),
+            ("a_order", STRING),
+            ("b_order", STRING),
+            ("c_order", STRING),
+            ("d_order", STRING),
+            ("lda", EXPRESSION),
+            ("ldb", EXPRESSION),
+            ("ldc", EXPRESSION),
+            ("ldd", EXPRESSION),
+            ("batch_count", EXPRESSION),
+            ("stride_a", EXPRESSION),
+            ("stride_b", EXPRESSION),
+            ("stride_c", EXPRESSION),
+            ("stride_d", EXPRESSION),
+            ("a_dtype", DTYPE),
+            ("b_dtype", DTYPE),
+            ("c_dtype", DTYPE),
+            ("d_dtype", DTYPE),
+            ("compute_type", STRING),
+            ("scale_dtype", STRING),
+            ("alpha", F64),
+            ("beta", F64),
+            ("epilogue", STRING),
+        ],
+    )
+}
+
 impl EgglogOp for CuBlasLt {
     fn sort(&self) -> SortDef {
-        sort(
-            OP_KIND,
-            "cublaslt",
-            &[
-                ("m", EXPRESSION),
-                ("n", EXPRESSION),
-                ("k", EXPRESSION),
-                ("a_layout", STRING),
-                ("b_layout", STRING),
-                ("a_order", STRING),
-                ("b_order", STRING),
-                ("c_order", STRING),
-                ("d_order", STRING),
-                ("lda", EXPRESSION),
-                ("ldb", EXPRESSION),
-                ("ldc", EXPRESSION),
-                ("ldd", EXPRESSION),
-                ("batch_count", EXPRESSION),
-                ("stride_a", EXPRESSION),
-                ("stride_b", EXPRESSION),
-                ("stride_c", EXPRESSION),
-                ("stride_d", EXPRESSION),
-                ("a_dtype", DTYPE),
-                ("b_dtype", DTYPE),
-                ("c_dtype", DTYPE),
-                ("d_dtype", DTYPE),
-                ("compute_type", STRING),
-                ("scale_dtype", STRING),
-                ("alpha", F64),
-                ("beta", F64),
-                ("epilogue", STRING),
-            ],
-        )
+        cublaslt_sort("cublaslt")
     }
 
     fn n_inputs(&self) -> usize {
         let c_input = usize::from(self.beta != 0.0);
         let bias_input = usize::from(epilogue_uses_bias(self.epilogue));
-        2 + c_input + bias_input
+        let scale_inputs = usize::from(self.a_scale_input) + usize::from(self.b_scale_input);
+        2 + c_input + bias_input + scale_inputs
     }
 
     fn rewrites(&self) -> Vec<Rule> {
@@ -158,40 +180,69 @@ impl EgglogOp for CuBlasLt {
                  (cublaslt_base_dtype (F32))
                  (cublaslt_base_dtype (F16))
                  (cublaslt_base_dtype (Bf16))
-                 (cublaslt_base_dtype (TF32))",
+                 (cublaslt_base_dtype (TF32))
+                 (relation cublaslt_fp8_dtype (DType))
+                 (cublaslt_fp8_dtype (F8E4M3))
+                 (cublaslt_fp8_dtype (F8E5M2))
+                 (relation cublaslt_fp8_f32_output_pair (DType DType))
+                 (cublaslt_fp8_f32_output_pair (F8E4M3) (F8E4M3))
+                 (cublaslt_fp8_f32_output_pair (F8E4M3) (F8E5M2))
+                 (cublaslt_fp8_f32_output_pair (F8E5M2) (F8E4M3))",
             ),
             Rule::raw(include_str!["cublaslt_RmRm_rewrite.egg"]), // row row
             Rule::raw(include_str!["cublaslt_RmCm_rewrite.egg"]), // row col
             Rule::raw(include_str!["cublaslt_CmRm_rewrite.egg"]), // col row
             Rule::raw(include_str!["cublaslt_CmCm_rewrite.egg"]), // col col
             Rule::raw(include_str!["cublaslt_fp8_rewrite.egg"]),
-            Rule::raw(include_str!["cublaslt_row_order_rewrite.egg"]),
             Rule::raw(include_str!["cublaslt_mixed_dtype_rewrite.egg"]),
             Rule::raw(include_str!["cublaslt_scale_rewrite.egg"]),
             Rule::raw(include_str!["cublaslt_beta_rewrite.egg"]),
             Rule::raw(include_str!["cublaslt_epilogue_rewrite.egg"]),
-            // Delete KernelMul matmul broadcast intermediates when the Sum eclass
-            // has a cublaslt or KernelBatchMatMul alternative. This prevents OOM
-            // from O(m*k*n) intermediates at large seq_len. cuBLAS, TileMatmulFullSplit,
-            // KernelBatchMatVec, and KernelBatchMatMul all take original inputs
-            // (not the Mul eclass), so they survive the cascade.
+            Rule::raw(include_str!["cublaslt_row_order_rewrite.egg"]),
+            // cuBLASLt now specializes GenericMatmul, so cleanup should prune
+            // the matmul output alternatives directly. Do not delete the
+            // broadcast Mul here; it may still have non-matmul consumers.
             Rule::raw("(rule
-                ((= ?mul (Op (KernelMul ?shape ?as ?bs ?os ?dt) ?inputs))
-                 (= (MNum 0) (nth_from_end ?as 1))
-                 (= (MNum 0) (nth_from_end ?bs 2))
-                 (= ?sum (Op (Sum ?sshape ?sk ?ssi ?sks ?sso) (ICons ?mul (INil))))
+                ((= ?sum (Op (Sum ?shape ?k ?sis ?sks ?sos) ?inputs))
                  (= ?sum (Op (cublaslt ?cm ?cn ?ck ?cta ?ctb ?cao ?cbo ?cco ?cdo ?clda ?cldb ?cldc ?cldd ?cbc ?csa ?csb ?csc ?csd ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue) ?ci)))
-                ((delete (Op (KernelMul ?shape ?as ?bs ?os ?dt) ?inputs)))
+                ((delete (Op (Sum ?shape ?k ?sis ?sks ?sos) ?inputs)))
                 :ruleset cleanup
+                :name \"delete-sum-when-cublaslt-exists\"
             )"),
             Rule::raw("(rule
-                ((= ?mul (Op (KernelMul ?shape ?as ?bs ?os ?dt) ?inputs))
-                 (= (MNum 0) (nth_from_end ?as 1))
-                 (= (MNum 0) (nth_from_end ?bs 2))
-                 (= ?sum (Op (Sum ?sshape ?sk ?ssi ?sks ?sso) (ICons ?mul (INil))))
-                 (= ?sum (Op (KernelBatchMatMul ?bos ?bk ?bas ?baks ?bbs ?bbks ?bouts ?bdt) ?bi)))
-                ((delete (Op (KernelMul ?shape ?as ?bs ?os ?dt) ?inputs)))
+                ((= ?sum (Op (KernelSum ?shape ?k ?sis ?sks ?sos ?dt) ?inputs))
+                 (= ?sum (Op (cublaslt ?cm ?cn ?ck ?cta ?ctb ?cao ?cbo ?cco ?cdo ?clda ?cldb ?cldc ?cldd ?cbc ?csa ?csb ?csc ?csd ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue) ?ci)))
+                ((delete (Op (KernelSum ?shape ?k ?sis ?sks ?sos ?dt) ?inputs)))
                 :ruleset cleanup
+                :name \"delete-kernel-sum-when-cublaslt-exists\"
+            )"),
+            Rule::raw("(rule
+                ((= ?sum (Op (Sum ?shape ?k ?sis ?sks ?sos) ?inputs))
+                 (= ?sum (Op (cublaslt_scaled ?cm ?cn ?ck ?cta ?ctb ?cao ?cbo ?cco ?cdo ?clda ?cldb ?cldc ?cldd ?cbc ?csa ?csb ?csc ?csd ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue) ?ci)))
+                ((delete (Op (Sum ?shape ?k ?sis ?sks ?sos) ?inputs)))
+                :ruleset cleanup
+                :name \"delete-sum-when-scaled-cublaslt-exists\"
+            )"),
+            Rule::raw("(rule
+                ((= ?sum (Op (KernelSum ?shape ?k ?sis ?sks ?sos ?dt) ?inputs))
+                 (= ?sum (Op (cublaslt_scaled ?cm ?cn ?ck ?cta ?ctb ?cao ?cbo ?cco ?cdo ?clda ?cldb ?cldc ?cldd ?cbc ?csa ?csb ?csc ?csd ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue) ?ci)))
+                ((delete (Op (KernelSum ?shape ?k ?sis ?sks ?sos ?dt) ?inputs)))
+                :ruleset cleanup
+                :name \"delete-kernel-sum-when-scaled-cublaslt-exists\"
+            )"),
+            Rule::raw("(rule
+                ((= ?sum (Op (GenericMatmul ?go ?gm ?gk ?gls ?grs ?gsis ?gsit ?gos ?gdt) ?generic_inputs))
+                 (= ?sum (Op (cublaslt ?cm ?cn ?ck ?cta ?ctb ?cao ?cbo ?cco ?cdo ?clda ?cldb ?cldc ?cldd ?cbc ?csa ?csb ?csc ?csd ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue) ?cublas_inputs)))
+                ((delete (Op (GenericMatmul ?go ?gm ?gk ?gls ?grs ?gsis ?gsit ?gos ?gdt) ?generic_inputs)))
+                :ruleset cleanup
+                :name \"prefer-cublaslt-over-generic-matmul\"
+            )"),
+            Rule::raw("(rule
+                ((= ?sum (Op (GenericMatmul ?go ?gm ?gk ?gls ?grs ?gsis ?gsit ?gos ?gdt) ?generic_inputs))
+                 (= ?sum (Op (cublaslt_scaled ?cm ?cn ?ck ?cta ?ctb ?cao ?cbo ?cco ?cdo ?clda ?cldb ?cldc ?cldd ?cbc ?csa ?csb ?csc ?csd ?cadt ?cbdt ?ccdt ?cddt ?ccompute ?cscale ?calpha ?cbeta ?cepilogue) ?cublas_inputs)))
+                ((delete (Op (GenericMatmul ?go ?gm ?gk ?gls ?grs ?gsis ?gsit ?gos ?gdt) ?generic_inputs)))
+                :ruleset cleanup
+                :name \"prefer-scaled-cublaslt-over-generic-matmul\"
             )"),
         ]
     }
@@ -277,6 +328,104 @@ impl EgglogOp for CuBlasLt {
             alpha,
             beta,
             epilogue,
+            a_scale_input: false,
+            b_scale_input: false,
+            cublaslt: OnceLock::new(),
+        };
+        trace!(?extracted_state);
+
+        let extracted = LLIROp::new::<dyn HostOp>(Box::new(extracted_state) as Box<dyn HostOp>);
+
+        (extracted, input_enodes)
+    }
+
+    fn cleanup(&self) -> bool {
+        false
+    }
+}
+
+impl EgglogOp for CuBlasLtScaled {
+    fn sort(&self) -> SortDef {
+        cublaslt_sort("cublaslt_scaled")
+    }
+
+    fn n_inputs(&self) -> usize {
+        4
+    }
+
+    #[allow(unused_variables)]
+    fn extract<'a>(
+        &'a self,
+        egraph: &'a luminal::egglog_utils::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>) {
+        let m = extract_expr(egraph, kind_children[0], expr_cache).unwrap();
+        let n = extract_expr(egraph, kind_children[1], expr_cache).unwrap();
+        let k = extract_expr(egraph, kind_children[2], expr_cache).unwrap();
+
+        let a_layout = parse_cublas_op(&egraph.enodes[kind_children[3]].0);
+        let b_layout = parse_cublas_op(&egraph.enodes[kind_children[4]].0);
+        let a_order = parse_cublaslt_order(&egraph.enodes[kind_children[5]].0);
+        let b_order = parse_cublaslt_order(&egraph.enodes[kind_children[6]].0);
+        let c_order = parse_cublaslt_order(&egraph.enodes[kind_children[7]].0);
+        let d_order = parse_cublaslt_order(&egraph.enodes[kind_children[8]].0);
+
+        let lda = extract_expr(egraph, kind_children[9], expr_cache).unwrap();
+        let ldb = extract_expr(egraph, kind_children[10], expr_cache).unwrap();
+        let ldc = extract_expr(egraph, kind_children[11], expr_cache).unwrap();
+        let ldd = extract_expr(egraph, kind_children[12], expr_cache).unwrap();
+
+        let batch_count = extract_expr(egraph, kind_children[13], expr_cache).unwrap();
+        let stride_a = extract_expr(egraph, kind_children[14], expr_cache).unwrap();
+        let stride_b = extract_expr(egraph, kind_children[15], expr_cache).unwrap();
+        let stride_c = extract_expr(egraph, kind_children[16], expr_cache).unwrap();
+        let stride_d = extract_expr(egraph, kind_children[17], expr_cache).unwrap();
+
+        let a_dtype = extract_dtype(egraph, kind_children[18]);
+        let b_dtype = extract_dtype(egraph, kind_children[19]);
+        let c_dtype = extract_dtype(egraph, kind_children[20]);
+        let d_dtype = extract_dtype(egraph, kind_children[21]);
+        let compute_type_str = &egraph.enodes[kind_children[22]].0;
+        let scale_dtype_str = &egraph.enodes[kind_children[23]].0;
+        let compute_type = parse_cublaslt_compute_type(compute_type_str, a_dtype);
+        let scale_dtype = parse_cublaslt_scale_dtype(scale_dtype_str, a_dtype);
+        let alpha = parse_cublaslt_scalar(&egraph.enodes[kind_children[24]].0);
+        let beta = parse_cublaslt_scalar(&egraph.enodes[kind_children[25]].0);
+        let epilogue = parse_cublaslt_epilogue(&egraph.enodes[kind_children[26]].0);
+
+        let extracted_state = CuBlasLt {
+            m,
+            n,
+            k,
+            a_layout,
+            b_layout,
+            a_order,
+            b_order,
+            c_order,
+            d_order,
+            lda,
+            ldb,
+            ldc,
+            ldd,
+            batch_count,
+            stride_a,
+            stride_b,
+            stride_c,
+            stride_d,
+            a_dtype,
+            b_dtype,
+            c_dtype,
+            d_dtype,
+            compute_type,
+            scale_dtype,
+            alpha,
+            beta,
+            epilogue,
+            a_scale_input: true,
+            b_scale_input: true,
             cublaslt: OnceLock::new(),
         };
         trace!(?extracted_state);
@@ -520,6 +669,8 @@ struct LtMatmulPointers {
     c: u64,
     d: u64,
     bias: Option<u64>,
+    a_scale: Option<u64>,
+    b_scale: Option<u64>,
 }
 
 struct LtRawDescriptors {
@@ -667,12 +818,12 @@ fn run_cublaslt_matmul(
     let workspace = unsafe { stream.alloc::<u8>(spec.workspace_size)? };
     let (workspace_ptr, _workspace_guard) = workspace.device_ptr(stream);
 
-    let a_scale = if cuda_dtype_needs_tensorwide_scale(spec.a.dtype) {
+    let a_scale = if cuda_dtype_needs_tensorwide_scale(spec.a.dtype) && ptrs.a_scale.is_none() {
         Some(stream.clone_htod(&[1.0f32])?)
     } else {
         None
     };
-    let b_scale = if cuda_dtype_needs_tensorwide_scale(spec.b.dtype) {
+    let b_scale = if cuda_dtype_needs_tensorwide_scale(spec.b.dtype) && ptrs.b_scale.is_none() {
         Some(stream.clone_htod(&[1.0f32])?)
     } else {
         None
@@ -728,13 +879,17 @@ fn run_cublaslt_matmul(
         }
     }
 
-    let (a_scale_ptr, _a_scale_guard) = if let Some(scale) = &a_scale {
+    let (a_scale_ptr, _a_scale_guard) = if let Some(ptr) = ptrs.a_scale {
+        (Some(ptr), None)
+    } else if let Some(scale) = &a_scale {
         let (ptr, guard) = scale.device_ptr(stream);
         (Some(ptr), Some(guard))
     } else {
         (None, None)
     };
-    let (b_scale_ptr, _b_scale_guard) = if let Some(scale) = &b_scale {
+    let (b_scale_ptr, _b_scale_guard) = if let Some(ptr) = ptrs.b_scale {
+        (Some(ptr), None)
+    } else if let Some(scale) = &b_scale {
         let (ptr, guard) = scale.device_ptr(stream);
         (Some(ptr), Some(guard))
     } else {
@@ -857,6 +1012,8 @@ fn resolve_cublaslt_pointers(
     buffers: &FxHashMap<NodeIndex, DeviceBuffer>,
     beta: f64,
     epilogue: cublasLtEpilogue_t,
+    a_scale_input: bool,
+    b_scale_input: bool,
 ) -> anyhow::Result<LtMatmulPointers> {
     if inputs.len() < 2 {
         return Err(anyhow::anyhow!(
@@ -877,24 +1034,25 @@ fn resolve_cublaslt_pointers(
         .get(&self_node)
         .ok_or_else(|| anyhow::anyhow!("missing cuBLASLt output buffer"))?
         .ptr();
+    let mut next_input = 2;
     let c = if beta == 0.0 {
         d
-    } else if let Some(c_input) = inputs.get(2) {
+    } else {
+        let c_input = inputs.get(next_input).ok_or_else(|| {
+            anyhow::anyhow!("cuBLASLt matmul with beta={beta} requires a third C input")
+        })?;
+        next_input += 1;
         buffers
             .get(c_input)
             .ok_or_else(|| anyhow::anyhow!("missing cuBLASLt C input buffer"))?
             .ptr()
-    } else {
-        return Err(anyhow::anyhow!(
-            "cuBLASLt matmul with beta={beta} requires a third C input"
-        ));
     };
 
-    let bias_input_index = if beta == 0.0 { 2 } else { 3 };
     let bias = if epilogue_uses_bias(epilogue) {
-        let bias_input = inputs.get(bias_input_index).ok_or_else(|| {
+        let bias_input = inputs.get(next_input).ok_or_else(|| {
             anyhow::anyhow!("cuBLASLt matmul with {epilogue:?} epilogue requires a bias input")
         })?;
+        next_input += 1;
         Some(
             buffers
                 .get(bias_input)
@@ -905,7 +1063,44 @@ fn resolve_cublaslt_pointers(
         None
     };
 
-    Ok(LtMatmulPointers { a, b, c, d, bias })
+    let a_scale = if a_scale_input {
+        let scale_input = inputs
+            .get(next_input)
+            .ok_or_else(|| anyhow::anyhow!("cuBLASLt matmul requires an A scale input pointer"))?;
+        next_input += 1;
+        Some(
+            buffers
+                .get(scale_input)
+                .ok_or_else(|| anyhow::anyhow!("missing cuBLASLt A scale input buffer"))?
+                .ptr(),
+        )
+    } else {
+        None
+    };
+
+    let b_scale = if b_scale_input {
+        let scale_input = inputs
+            .get(next_input)
+            .ok_or_else(|| anyhow::anyhow!("cuBLASLt matmul requires a B scale input pointer"))?;
+        Some(
+            buffers
+                .get(scale_input)
+                .ok_or_else(|| anyhow::anyhow!("missing cuBLASLt B scale input buffer"))?
+                .ptr(),
+        )
+    } else {
+        None
+    };
+
+    Ok(LtMatmulPointers {
+        a,
+        b,
+        c,
+        d,
+        bias,
+        a_scale,
+        b_scale,
+    })
 }
 
 fn epilogue_uses_bias(epilogue: cublasLtEpilogue_t) -> bool {
@@ -978,6 +1173,11 @@ impl CuBlasLt {
             && normalize(self.stride_c) == normalize(self.stride_d)
             && self.c_order == self.d_order
     }
+
+    #[cfg(test)]
+    pub(crate) fn tensor_scale_inputs(&self) -> (bool, bool) {
+        (self.a_scale_input, self.b_scale_input)
+    }
 }
 
 impl HostOp for CuBlasLt {
@@ -1022,7 +1222,15 @@ impl HostOp for CuBlasLt {
         let alpha = LtScalar::from_f64(self.scale_dtype, self.alpha)?;
         let beta = LtScalar::from_f64(self.scale_dtype, self.beta)?;
 
-        let ptrs = resolve_cublaslt_pointers(self_node, inputs, buffers, self.beta, self.epilogue)?;
+        let ptrs = resolve_cublaslt_pointers(
+            self_node,
+            inputs,
+            buffers,
+            self.beta,
+            self.epilogue,
+            self.a_scale_input,
+            self.b_scale_input,
+        )?;
 
         let (a_rows, a_cols) = if a_layout == cublasOperation_t::CUBLAS_OP_N {
             (m, k)
@@ -1197,6 +1405,8 @@ mod tests {
             &buffers,
             0.0,
             cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DEFAULT,
+            false,
+            false,
         )
         .unwrap();
 
@@ -1221,6 +1431,8 @@ mod tests {
             &buffers,
             0.0,
             cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DEFAULT,
+            false,
+            false,
         )
         .unwrap();
 
@@ -1245,6 +1457,8 @@ mod tests {
             &buffers,
             1.0,
             cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DEFAULT,
+            false,
+            false,
         )
         .unwrap();
 
@@ -1269,6 +1483,8 @@ mod tests {
             &buffers,
             0.0,
             cublasLtEpilogue_t::CUBLASLT_EPILOGUE_BIAS,
+            false,
+            false,
         )
         .unwrap();
 
@@ -1279,6 +1495,41 @@ mod tests {
         assert_eq!(ptrs.bias, Some(0xB1A5));
     }
 
+    #[test]
+    fn cublaslt_pointers_use_tensor_scale_inputs_after_base_inputs() {
+        let output = NodeIndex::new(0);
+        let a = NodeIndex::new(1);
+        let b = NodeIndex::new(2);
+        let a_scale = NodeIndex::new(3);
+        let b_scale = NodeIndex::new(4);
+        let buffers = buffers_for(&[
+            (output, 0xD000),
+            (a, 0xA000),
+            (b, 0xB000),
+            (a_scale, 0xA5A5),
+            (b_scale, 0xB5B5),
+        ]);
+
+        let ptrs = resolve_cublaslt_pointers(
+            output,
+            &[a, b, a_scale, b_scale],
+            &buffers,
+            0.0,
+            cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DEFAULT,
+            true,
+            true,
+        )
+        .unwrap();
+
+        assert_eq!(ptrs.a, 0xA000);
+        assert_eq!(ptrs.b, 0xB000);
+        assert_eq!(ptrs.c, 0xD000);
+        assert_eq!(ptrs.d, 0xD000);
+        assert_eq!(ptrs.bias, None);
+        assert_eq!(ptrs.a_scale, Some(0xA5A5));
+        assert_eq!(ptrs.b_scale, Some(0xB5B5));
+    }
+
     #[test]
     fn cublaslt_pointers_reject_two_input_nonzero_beta() {
         let output = NodeIndex::new(0);
@@ -1292,6 +1543,8 @@ mod tests {
             &buffers,
             1.0,
             cublasLtEpilogue_t::CUBLASLT_EPILOGUE_DEFAULT,
+            false,
+            false,
         )
         .unwrap_err();
 
@@ -1314,6 +1567,8 @@ mod tests {
             &buffers,
             0.0,
             cublasLtEpilogue_t::CUBLASLT_EPILOGUE_BIAS,
+            false,
+            false,
         )
         .unwrap_err();
 

crates/luminal_cuda_lite/src/host/flashinfer/find_indptrs.rs

@@ -27,19 +27,16 @@ pub fn find_indptr_inputs<'a>(
     mask_node: &'a NodeId,
 ) -> IndptrNodes<'a> {
     // Step 1: Validate mask = Add(scaled_allowed, neg_constant)
-    let (mask_label, mask_children) = &egraph.enodes[mask_node];
-    assert!(
-        mask_label == "Op",
-        "find_indptr_inputs: mask node is not an Op (label={mask_label})"
-    );
-    let mask_kind = resolve_first_node(egraph, &mask_children[0]);
-    let mask_kind_label = &egraph.enodes[mask_kind].0;
-    assert!(
-        mask_kind_label.contains("Add"),
-        "find_indptr_inputs: mask is not an Add (kind={mask_kind_label})"
-    );
-
-    let mask_inputs = walk_ilist_simple(egraph, &mask_children[1]);
+    let mask_inputs = logical_binary_inputs(egraph, mask_node, "Add").unwrap_or_else(|| {
+        let (mask_label, mask_children) = &egraph.enodes[mask_node];
+        assert!(
+            mask_label == "Op",
+            "find_indptr_inputs: mask node is not an Op (label={mask_label})"
+        );
+        let mask_kind = resolve_first_node(egraph, &mask_children[0]);
+        let mask_kind_label = &egraph.enodes[mask_kind].0;
+        panic!("find_indptr_inputs: mask is not an Add (kind={mask_kind_label})");
+    });
     assert_eq!(
         mask_inputs.len(),
         2,
@@ -98,15 +95,9 @@ fn find_1e10_mul<'a>(
     mask_add_inputs: &[&'a NodeId],
 ) -> (&'a NodeId, &'a NodeId) {
     for &input_node in mask_add_inputs {
-        let (label, children) = &egraph.enodes[input_node];
-        if label != "Op" {
+        let Some(mul_inputs) = logical_binary_inputs(egraph, input_node, "Mul") else {
             continue;
-        }
-        let kind = resolve_first_node(egraph, &children[0]);
-        if !egraph.enodes[kind].0.contains("Mul") {
-            continue;
-        }
-        let mul_inputs = walk_ilist_simple(egraph, &children[1]);
+        };
         if mul_inputs.len() != 2 {
             continue;
         }
@@ -152,6 +143,7 @@ fn find_1e10_mul<'a>(
 }
 
 fn is_constant(egraph: &SerializedEGraph, node: &NodeId, expected: f32) -> bool {
+    let node = resolve_op_with_kind(egraph, node, "Constant").unwrap_or(node);
     let (label, children) = &egraph.enodes[node];
     if label != "Op" {
         return false;
@@ -246,3 +238,91 @@ fn resolve_first_ir_node<'a>(egraph: &'a SerializedEGraph, eclass: &ClassId) ->
     }
     &nodes[0]
 }
+
+fn resolve_op_with_kind<'a>(
+    egraph: &'a SerializedEGraph,
+    node: &'a NodeId,
+    kind_substr: &str,
+) -> Option<&'a NodeId> {
+    let class = egraph.node_to_class.get(node)?;
+    for candidate in &egraph.eclasses[class].1 {
+        let (label, children) = &egraph.enodes[candidate];
+        if label != "Op" || children.is_empty() {
+            continue;
+        }
+        let kind = resolve_first_node(egraph, &children[0]);
+        if egraph.enodes[kind].0.contains(kind_substr) {
+            return Some(candidate);
+        }
+    }
+    None
+}
+
+fn logical_binary_inputs<'a>(
+    egraph: &'a SerializedEGraph,
+    node: &'a NodeId,
+    op_name: &str,
+) -> Option<Vec<&'a NodeId>> {
+    if let Some(op_node) = resolve_op_with_kind(egraph, node, op_name) {
+        let (_, children) = &egraph.enodes[op_node];
+        return Some(walk_ilist_simple(egraph, &children[1]));
+    }
+
+    let (label, children) = &egraph.enodes[node];
+    if label != "Op" || children.len() < 2 {
+        return None;
+    }
+    let kind = resolve_first_node(egraph, &children[0]);
+    if egraph.enodes[kind].0.contains("CudaBinaryElementwise") {
+        let opcode_class = egraph.enodes[kind].1.first()?;
+        let opcode_node = resolve_first_node(egraph, opcode_class);
+        if egraph.enodes[opcode_node].0.trim_matches('"') != op_name {
+            return None;
+        }
+        return Some(
+            walk_ilist_simple(egraph, &children[1])
+                .into_iter()
+                .map(|input| unwrap_fusion_start(egraph, input))
+                .collect(),
+        );
+    }
+    if !egraph.enodes[kind].0.contains("FusionEnd") {
+        return None;
+    }
+    let fe_inputs = walk_ilist_simple(egraph, &children[1]);
+    let elem = *fe_inputs.first()?;
+    let (elem_label, elem_children) = &egraph.enodes[elem];
+    if elem_label != "Op" || elem_children.len() < 2 {
+        return None;
+    }
+    let elem_kind = resolve_first_node(egraph, &elem_children[0]);
+    if !egraph.enodes[elem_kind].0.contains("CudaBinaryElementwise") {
+        return None;
+    }
+    let opcode_class = egraph.enodes[elem_kind].1.first()?;
+    let opcode_node = resolve_first_node(egraph, opcode_class);
+    if egraph.enodes[opcode_node].0.trim_matches('"') != op_name {
+        return None;
+    }
+    Some(
+        walk_ilist_simple(egraph, &elem_children[1])
+            .into_iter()
+            .map(|input| unwrap_fusion_start(egraph, input))
+            .collect(),
+    )
+}
+
+fn unwrap_fusion_start<'a>(egraph: &'a SerializedEGraph, node: &'a NodeId) -> &'a NodeId {
+    let (label, children) = &egraph.enodes[node];
+    if label != "Op" || children.len() < 2 {
+        return node;
+    }
+    let kind = resolve_first_node(egraph, &children[0]);
+    if !egraph.enodes[kind].0.contains("FusionStart") {
+        return node;
+    }
+    walk_ilist_simple(egraph, &children[1])
+        .first()
+        .copied()
+        .unwrap_or(node)
+}

crates/luminal_cuda_lite/src/host/flashinfer/flashinfer_attention.egg

@@ -89,6 +89,16 @@
             ?mask_add_out_strides)
             (ICons ?scaled_qk (ICons ?mask (INil)))))
 
+        ; FlashInfer needs qo_indptr/kv_indptr to be recoverable from the mask
+        ; expression. Do not match examples that pass a precomputed mask Input.
+        (= ?mask (Op (Add ?inner_mask_shape ?inner_mask_a_strides ?inner_mask_b_strides ?inner_mask_out_strides)
+            (ICons ?mask_scaled_allowed (ICons ?mask_offset (INil)))))
+        (= ?mask_scaled_allowed (Op (Mul ?allowed_shape ?allowed_strides ?scale_const_strides ?scaled_allowed_strides)
+            (ICons ?mask_allowed (ICons ?mask_scale_const (INil)))))
+        (= ?mask_scale_const (Op (Constant ?mask_scale_val) (INil)))
+        (> ?mask_scale_val 9999999999.0)
+        (< ?mask_scale_val 10000000001.0)
+
         ; ── K GQA broadcast: Mul(K_gathered, 1.0) with zero-stride constant ──
         ; Shape: (nheads, hdim, c) — 3D
         (= ?k_gqa_const (Op (Constant 1.000000) (INil)))

crates/luminal_cuda_lite/src/host/mod.rs

@@ -2,19 +2,14 @@ use std::{fmt::Debug, sync::Arc};
 
 use crate::cudarc::driver::{CudaStream, DriverError, result};
 use luminal::{op::EgglogOp, prelude::*};
-pub mod compute_attn_mask;
-mod cublas;
 mod cublaslt;
 pub mod flashinfer;
 pub mod moe;
 
-pub use compute_attn_mask::ComputeAttnMask;
-
 pub type Ops = (
-    // cublas::CuBlasSgemmV2,
     cublaslt::CuBlasLt,
+    cublaslt::CuBlasLtScaled,
     moe::GLUMoE,
-    compute_attn_mask::ComputeAttnMask,
     flashinfer::FlashInferAttention,
 );
 
@@ -79,6 +74,16 @@ pub(crate) fn cublaslt_c_d_layouts_match(op: &dyn HostOp) -> Option<bool> {
         .map(cublaslt::CuBlasLt::c_d_layouts_match)
 }
 
+#[cfg(test)]
+pub(crate) type CublasLtTensorScaleInputs = (bool, bool);
+
+#[cfg(test)]
+pub(crate) fn cublaslt_tensor_scale_inputs(op: &dyn HostOp) -> Option<CublasLtTensorScaleInputs> {
+    op.as_any()
+        .downcast_ref::<cublaslt::CuBlasLt>()
+        .map(cublaslt::CuBlasLt::tensor_scale_inputs)
+}
+
 /// Non-owning device buffer handle used by host operations.
 ///
 /// Runtime-owned intermediates may be a whole `CudaSlice`, a subregion inside

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

@@ -195,6 +195,10 @@
         (= ?swiglu_state (MkGLUMoESwiGLUState ?gate_up_state))
         (= ?gate_up_state (MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_within_range ?x ?topk_idx ?gate_up_w))
 
+        (= ?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)))))
+
         (= ?weighted (Op (Mul ?weighted_shape ?weighted_a_stride ?weighted_b_stride ?weighted_out_stride) (ICons ?dn_matmul (ICons ?topk_vals (INil)))))
         (= ?output (Op (Sum ?output_shape ?output_k ?output_in_stride ?output_k_stride ?output_out_stride) (ICons ?weighted (INil))))
     )
@@ -211,6 +215,37 @@
     :name "GLUMoE fused expert computation (swiglu)"
 )
 
+; ===== Final fusion: mode 2 (SwiGLU with row-normalized top-k weights) =====
+(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))
+
+        (= ?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)))))
+
+        (= ?weighted (Op (Mul ?weighted_shape ?weighted_a_stride ?weighted_b_stride ?weighted_out_stride) (ICons ?dn_matmul (ICons ?normed_topk (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 2))
+            (ICons ?x (ICons ?topk_idx (ICons ?topk_vals (ICons ?gate_up_w (ICons ?down_w (ICons ?topk_vals (INil)))))))))
+        (union ?output ?glumoe)
+        (subsume (Op (Sum ?output_shape ?output_k ?output_in_stride ?output_k_stride ?output_out_stride) (ICons ?weighted (INil))))
+        (subsume (Op (KernelSum ?output_shape ?output_k ?output_in_stride ?output_k_stride ?output_out_stride (F32)) (ICons ?weighted (INil))))
+    )
+    :ruleset glumoe
+    :name "GLUMoE fused expert computation (normalized swiglu)"
+)
+
 ; ===== Final fusion: mode 1 (Gemma GELU) =====
 (rule
     (

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

@@ -50,7 +50,7 @@ const WORKSPACE_SIZE: usize = 32 * 1024 * 1024; // 32 MiB
 ///   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)
+///      - SwiGLU/SwiGLUNormalized: ignored (rewriter wires `topk_values` again)
 ///      - GemmaGELU: per_expert_scale [E]                   F32
 ///
 /// Output: [seq, hidden] F32
@@ -78,6 +78,7 @@ pub struct GLUMoE {
 pub(crate) enum GLUMoEMode {
     SwiGLU,
     GemmaGELU,
+    SwiGLUNormalized,
 }
 
 impl GLUMoEMode {
@@ -85,6 +86,7 @@ impl GLUMoEMode {
         match mode_id {
             0 => Self::SwiGLU,
             1 => Self::GemmaGELU,
+            2 => Self::SwiGLUNormalized,
             other => {
                 panic!("Unknown GLUMoE mode id: {other}");
             }
@@ -93,7 +95,7 @@ impl GLUMoEMode {
 
     fn activation_kernel_mode(self) -> i32 {
         match self {
-            Self::SwiGLU => 0,
+            Self::SwiGLU | Self::SwiGLUNormalized => 0,
             Self::GemmaGELU => 1,
         }
     }
@@ -383,22 +385,22 @@ impl HostOp for GLUMoE {
         let mode_aux_buf = get_buffer("mode aux", inputs[5])?;
         let output_buf = get_buffer("output", self_node)?; // [seq, hidden] F32
 
-        let topk_bytes = seq * top_k * 4;
+        let min_topk_bytes = seq * top_k * 4;
         if x_buf.len() < output_bytes {
             anyhow::bail!(
                 "GLUMoE x buffer too small: have {} bytes, need {output_bytes}",
                 x_buf.len()
             );
         }
-        if topk_idx_buf.len() < topk_bytes {
+        if topk_idx_buf.len() < min_topk_bytes {
             anyhow::bail!(
-                "GLUMoE topk index buffer too small: have {} bytes, need {topk_bytes}",
+                "GLUMoE topk index buffer too small: have {} bytes, need {min_topk_bytes}",
                 topk_idx_buf.len()
             );
         }
-        if topk_vals_buf.len() < topk_bytes {
+        if topk_vals_buf.len() < min_topk_bytes {
             anyhow::bail!(
-                "GLUMoE topk value buffer too small: have {} bytes, need {topk_bytes}",
+                "GLUMoE topk value buffer too small: have {} bytes, need {min_topk_bytes}",
                 topk_vals_buf.len()
             );
         }
@@ -440,24 +442,83 @@ impl HostOp for GLUMoE {
 
         // Read top-k routing values from GPU
         let topk_idx_host: Vec<u8> = topk_idx_buf.clone_dtoh(stream)?;
-        let topk_idx_i32: &[i32] = bytemuck::cast_slice(&topk_idx_host[..topk_bytes]);
+        let topk_idx_i32: &[i32] = bytemuck::cast_slice(&topk_idx_host);
         let topk_vals_host: Vec<u8> = topk_vals_buf.clone_dtoh(stream)?;
-        let topk_vals_f32: &[f32] = bytemuck::cast_slice(&topk_vals_host[..topk_bytes]);
+        let topk_vals_f32: &[f32] = bytemuck::cast_slice(&topk_vals_host);
 
-        for (pos, &expert_idx) in topk_idx_i32.iter().enumerate() {
-            if expert_idx < 0 || expert_idx as usize >= num_experts {
-                anyhow::bail!(
-                    "GLUMoE expert index {expert_idx} at routing position {pos} out of bounds for {num_experts} experts"
-                );
+        if !topk_idx_i32.len().is_multiple_of(seq) {
+            anyhow::bail!(
+                "GLUMoE topk index element count {} is not divisible by seq {seq}",
+                topk_idx_i32.len()
+            );
+        }
+        if !topk_vals_f32.len().is_multiple_of(seq) {
+            anyhow::bail!(
+                "GLUMoE topk value element count {} is not divisible by seq {seq}",
+                topk_vals_f32.len()
+            );
+        }
+        let topk_idx_row_stride = topk_idx_i32.len() / seq;
+        let topk_vals_row_stride = topk_vals_f32.len() / seq;
+        if topk_idx_row_stride < top_k {
+            anyhow::bail!(
+                "GLUMoE topk index row stride {topk_idx_row_stride} is smaller than top_k {top_k}"
+            );
+        }
+        if topk_vals_row_stride < top_k {
+            anyhow::bail!(
+                "GLUMoE topk value row stride {topk_vals_row_stride} is smaller than top_k {top_k}"
+            );
+        }
+
+        let topk_idx_at = |token: usize, expert: usize| -> i32 {
+            topk_idx_i32[token * topk_idx_row_stride + expert]
+        };
+        let topk_val_at = |token: usize, expert: usize| -> f32 {
+            topk_vals_f32[token * topk_vals_row_stride + expert]
+        };
+
+        for t in 0..seq {
+            for i in 0..top_k {
+                let expert_idx = topk_idx_at(t, i);
+                if expert_idx < 0 || expert_idx as usize >= num_experts {
+                    anyhow::bail!(
+                        "GLUMoE expert index {expert_idx} at token {t} top-k position {i} out of bounds for {num_experts} experts"
+                    );
+                }
             }
         }
 
         // Mode-dependent expert weights used for the final reduction:
         // - SwiGLU: direct topk values
+        // - SwiGLUNormalized: normalize topk values row-wise
         // - 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::SwiGLU => {
+                if topk_vals_row_stride == top_k {
+                    topk_vals_f32
+                } else {
+                    expert_weights_storage.resize(seq * top_k, 0.0);
+                    for t in 0..seq {
+                        for i in 0..top_k {
+                            expert_weights_storage[t * top_k + i] = topk_val_at(t, i);
+                        }
+                    }
+                    &expert_weights_storage
+                }
+            }
+            GLUMoEMode::SwiGLUNormalized => {
+                expert_weights_storage.resize(seq * top_k, 0.0);
+                for t in 0..seq {
+                    let norm = (0..top_k).map(|i| topk_val_at(t, i)).sum::<f32>();
+                    let inv_norm = if norm != 0.0 { norm.recip() } else { 0.0 };
+                    for i in 0..top_k {
+                        expert_weights_storage[t * top_k + i] = topk_val_at(t, i) * inv_norm;
+                    }
+                }
+                &expert_weights_storage
+            }
             GLUMoEMode::GemmaGELU => {
                 let per_expert_scale_host: Vec<u8> = mode_aux_buf.clone_dtoh(stream)?;
                 let per_expert_scale_bytes = num_experts * 4;
@@ -471,12 +532,10 @@ impl HostOp for GLUMoE {
                     bytemuck::cast_slice(&per_expert_scale_host[..per_expert_scale_bytes]);
                 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 norm = (0..top_k).map(|i| topk_val_at(t, i)).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;
+                        let expert_idx = topk_idx_at(t, i) as usize;
                         if expert_idx >= per_expert_scale_f32.len() {
                             anyhow::bail!(
                                 "GLUMoE Gemma mode expert index {} out of bounds {}",
@@ -485,7 +544,8 @@ impl HostOp for GLUMoE {
                             );
                         }
                         let scale = per_expert_scale_f32[expert_idx];
-                        expert_weights_storage[base + i] = vals[i] * inv_norm * scale;
+                        expert_weights_storage[t * top_k + i] =
+                            topk_val_at(t, i) * inv_norm * scale;
                     }
                 }
                 &expert_weights_storage
@@ -525,12 +585,10 @@ impl HostOp for GLUMoE {
 
         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 = &expert_weights_f32[t * top_k..(t + 1) * top_k];
 
-            for (i, (&expert_idx, &weight)) in expert_indices.iter().zip(weights.iter()).enumerate()
-            {
-                let expert_idx = expert_idx as usize;
+            for (i, &weight) in weights.iter().enumerate() {
+                let expert_idx = topk_idx_at(t, i) as usize;
 
                 // a. Gate+Up matmul (BF16 in, BF16 out)
                 let expert_gu_ptr = gate_up_ptr + expert_idx as u64 * gu_stride;

crates/luminal_cuda_lite/src/kernel/conv2d.rs

@@ -0,0 +1,738 @@
+//! CUDA conv2d-with-bias backend rewrite.
+//!
+//! `KernelConv2D` is selected by egglog from pure HLIR conv graphs and lowers
+//! to a one-thread-per-output CUDA kernel. It avoids materializing unfold/im2col
+//! intermediates while keeping model code free of custom ops.
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::prelude::FxHashMap;
+use luminal::{
+    dtype::DType,
+    egglog_utils::{
+        api::{Rule, SortDef, sort},
+        base::{DTYPE, ELIST, EXPRESSION, OP_KIND},
+        extract_dtype, extract_expr, extract_expr_list,
+    },
+    op::{EgglogOp, LLIROp},
+    prelude::FxHashSet,
+    shape::{Expression, flatten_strides},
+};
+
+use crate::compile_module_image_for_current_device;
+use crate::kernel::{KernelOp, hlir::generate_dyn_dims_defines};
+
+#[derive(Default, Debug, Clone)]
+pub struct KernelConv2D {
+    out_shape: Vec<Expression>,
+    input_shape: Vec<Expression>,
+    input_stride: Vec<Expression>,
+    weight_co_stride: Expression,
+    weight_inner_stride: Expression,
+    bias_c_stride: Expression,
+    out_stride: Vec<Expression>,
+    kernel_h: Expression,
+    kernel_w: Expression,
+    stride_h: Expression,
+    stride_w: Expression,
+    dilation_h: Expression,
+    dilation_w: Expression,
+    pad_h: Expression,
+    pad_w: Expression,
+    dtype: DType,
+}
+
+impl EgglogOp for KernelConv2D {
+    fn sort(&self) -> SortDef {
+        sort(
+            OP_KIND,
+            "KernelConv2D",
+            &[
+                ("out_shape", ELIST),
+                ("input_shape", ELIST),
+                ("input_stride", ELIST),
+                ("weight_co_stride", EXPRESSION),
+                ("weight_inner_stride", EXPRESSION),
+                ("bias_c_stride", EXPRESSION),
+                ("out_stride", ELIST),
+                ("kernel_h", EXPRESSION),
+                ("kernel_w", EXPRESSION),
+                ("stride_h", EXPRESSION),
+                ("stride_w", EXPRESSION),
+                ("dilation_h", EXPRESSION),
+                ("dilation_w", EXPRESSION),
+                ("pad_h", EXPRESSION),
+                ("pad_w", EXPRESSION),
+                ("dtype", DTYPE),
+            ],
+        )
+    }
+
+    fn n_inputs(&self) -> usize {
+        3
+    }
+
+    fn rewrites(&self) -> Vec<Rule> {
+        vec![
+            // 1x1 convs in Flux2's VAE are represented without `unfold`:
+            //
+            //   input.permute([H,W,C]).merge(H,W)
+            //     -> matmul(weight.t())
+            //     -> split/permute back to [C_out,H,W]
+            //     -> + channel bias
+            //
+            // The lowered form is still the same Mul -> KernelSum -> Add
+            // matmul skeleton, but the lhs FusionStart reads directly from the
+            // original input instead of a KernelGather window tensor.
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?out (Op (FusionEnd ?out_shape ?out_stride (F32)) (ICons ?add_elem (INil))))
+                        (= ?add_elem (Op (CudaBinaryElementwise \"Add\" ?out_shape ?sum_add_stride ?bias_add_stride ?out_stride (F32)) (ICons ?sum_fs (ICons ?bias_fs (INil)))))
+                        (= ?sum_fs (Op (FusionStart ?out_shape ?sum_add_stride (F32)) (ICons ?sum (INil))))
+                        (= ?bias_fs (Op (FusionStart ?out_shape ?bias_add_stride (F32)) (ICons ?bias (INil))))
+
+                        (= ?sum (Op (KernelSum ?matmul_out_shape ?c_in ?sum_in_stride ?k_stride ?sum_out_stride (F32)) (ICons ?mul_fe (INil))))
+                        (= ?mul_fe (Op (FusionEnd ?mul_shape ?mul_out_stride (F32)) (ICons ?mul_elem (INil))))
+                        (= ?mul_elem (Op (CudaBinaryElementwise \"Mul\" ?mul_shape ?input_1x1_stride ?weight_stride ?mul_out_stride (F32)) (ICons ?input_fs (ICons ?weight_fs (INil)))))
+                        (= ?input_fs (Op (FusionStart ?mul_shape ?input_1x1_stride (F32)) (ICons ?input (INil))))
+                        (= ?weight_fs (Op (FusionStart ?mul_shape ?weight_stride (F32)) (ICons ?weight (INil))))
+
+                        (= ?out_shape (ECons ?c_out (ECons ?h_out (ECons ?w_out (ENil)))))
+                        (= ?matmul_out_shape (ECons ?m (ECons ?c_out (ENil))))
+                        (= ?mul_shape (ECons ?m (ECons ?c_out (ECons ?c_in (ENil)))))
+                        (= ?input_1x1_stride (ECons ?flat_stride (ECons (MNum 0) (ECons ?input_c_stride (ENil)))))
+                        (= ?flat_stride (MIter))
+
+                        (= ?k_stride (MIter))
+                        (= ?sum_in_stride (ECons ?sum_m_stride (ECons ?sum_c_stride (ENil))))
+                        (= ?sum_out_stride (ECons ?sum_out_m_stride (ECons ?sum_out_c_stride (ENil))))
+                        (= ?sum_add_stride (ECons ?sum_add_c_stride (ECons ?sum_add_h_stride (ECons ?sum_add_w_stride (ENil)))))
+                        (= ?weight_co_stride (nth_from_end ?weight_stride 1))
+                        (= ?weight_inner_stride (nth_from_end ?weight_stride 0))
+                        (= (MNum 0) (nth_from_end ?weight_stride 2))
+                        (= ?bias_add_stride (ECons ?bias_c_stride (ECons (MNum 0) (ECons (MNum 0) (ENil)))))
+                    )
+                    (
+                        (let ?conv (Op (KernelConv2D
+                            ?out_shape
+                            (ECons ?c_in (ECons ?h_out (ECons ?w_out (ENil))))
+                            (ECons ?input_c_stride (ECons (MMul ?w_out ?flat_stride) (ECons ?flat_stride (ENil))))
+                            ?weight_co_stride
+                            ?weight_inner_stride
+                            ?bias_c_stride
+                            ?out_stride
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 0)
+                            (MNum 0)
+                            (F32))
+                            (ICons ?input (ICons ?weight (ICons ?bias (INil))))))
+                        (union ?out ?conv)
+                        (subsume (Op (FusionEnd ?out_shape ?out_stride (F32)) (ICons ?add_elem (INil))))
+                        (set (dtype ?conv) (F32))
+                    )
+                    :ruleset kernel_lower
+                    :name \"kernel conv2d 1x1 from cuda lowered matmul bias\"
+                )",
+            ),
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?out (Op (FusionEnd ?out_shape ?out_stride (F32)) (ICons ?add_elem (INil))))
+                        (= ?add_elem (Op (CudaBinaryElementwise \"Add\" ?out_shape ?bias_add_stride ?sum_add_stride ?out_stride (F32)) (ICons ?bias_fs (ICons ?sum_fs (INil)))))
+                        (= ?sum_fs (Op (FusionStart ?out_shape ?sum_add_stride (F32)) (ICons ?sum (INil))))
+                        (= ?bias_fs (Op (FusionStart ?out_shape ?bias_add_stride (F32)) (ICons ?bias (INil))))
+
+                        (= ?sum (Op (KernelSum ?matmul_out_shape ?c_in ?sum_in_stride ?k_stride ?sum_out_stride (F32)) (ICons ?mul_fe (INil))))
+                        (= ?mul_fe (Op (FusionEnd ?mul_shape ?mul_out_stride (F32)) (ICons ?mul_elem (INil))))
+                        (= ?mul_elem (Op (CudaBinaryElementwise \"Mul\" ?mul_shape ?input_1x1_stride ?weight_stride ?mul_out_stride (F32)) (ICons ?input_fs (ICons ?weight_fs (INil)))))
+                        (= ?input_fs (Op (FusionStart ?mul_shape ?input_1x1_stride (F32)) (ICons ?input (INil))))
+                        (= ?weight_fs (Op (FusionStart ?mul_shape ?weight_stride (F32)) (ICons ?weight (INil))))
+
+                        (= ?out_shape (ECons ?c_out (ECons ?h_out (ECons ?w_out (ENil)))))
+                        (= ?matmul_out_shape (ECons ?m (ECons ?c_out (ENil))))
+                        (= ?mul_shape (ECons ?m (ECons ?c_out (ECons ?c_in (ENil)))))
+                        (= ?input_1x1_stride (ECons ?flat_stride (ECons (MNum 0) (ECons ?input_c_stride (ENil)))))
+                        (= ?flat_stride (MIter))
+
+                        (= ?k_stride (MIter))
+                        (= ?sum_in_stride (ECons ?sum_m_stride (ECons ?sum_c_stride (ENil))))
+                        (= ?sum_out_stride (ECons ?sum_out_m_stride (ECons ?sum_out_c_stride (ENil))))
+                        (= ?sum_add_stride (ECons ?sum_add_c_stride (ECons ?sum_add_h_stride (ECons ?sum_add_w_stride (ENil)))))
+                        (= ?weight_co_stride (nth_from_end ?weight_stride 1))
+                        (= ?weight_inner_stride (nth_from_end ?weight_stride 0))
+                        (= (MNum 0) (nth_from_end ?weight_stride 2))
+                        (= ?bias_add_stride (ECons ?bias_c_stride (ECons (MNum 0) (ECons (MNum 0) (ENil)))))
+                    )
+                    (
+                        (let ?conv (Op (KernelConv2D
+                            ?out_shape
+                            (ECons ?c_in (ECons ?h_out (ECons ?w_out (ENil))))
+                            (ECons ?input_c_stride (ECons (MMul ?w_out ?flat_stride) (ECons ?flat_stride (ENil))))
+                            ?weight_co_stride
+                            ?weight_inner_stride
+                            ?bias_c_stride
+                            ?out_stride
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 0)
+                            (MNum 0)
+                            (F32))
+                            (ICons ?input (ICons ?weight (ICons ?bias (INil))))))
+                        (union ?out ?conv)
+                        (subsume (Op (FusionEnd ?out_shape ?out_stride (F32)) (ICons ?add_elem (INil))))
+                        (set (dtype ?conv) (F32))
+                    )
+                    :ruleset kernel_lower
+                    :name \"kernel conv2d 1x1 from cuda lowered bias matmul\"
+                )",
+            ),
+            // Match the same conv after generic CUDA lowering has normalized
+            // the elementwise pieces into fusion regions:
+            //
+            //   KernelGather(input windows)
+            //     -> CudaBinaryElementwise("Mul", weight)
+            //     -> KernelSum(reduce K)
+            //     -> CudaBinaryElementwise("Add", bias)
+            //
+            // This is the form that survives long enough for CUDA search in
+            // real models. The KernelConv2D op consumes the pre-gather input
+            // and avoids materializing both the im2col window tensor and the
+            // elementwise product tensor.
+            //
+            // TODO(egglog-shapes): the current e-graph does not reliably prove
+            // the derived arithmetic equalities for this chain after CUDA
+            // normalization:
+            //   * `M == H_out * W_out`
+            //   * `K == C_in * KH * KW`
+            //   * separately-derived but structurally identical stride
+            //     expressions, e.g. the Mul output stride and KernelSum input
+            //     stride, belong to the same e-class.
+            // Keep the rewrite anchored on the stable conv layout facts the
+            // graph does carry today: six-axis unfold window shape, flattened
+            // `[M, C_out, K]` product, reduction over `K`, the three-axis
+            // `[C_out, H_out, W_out]` output view, and channel-only bias
+            // broadcast. Once expression/list canonicalization can prove those
+            // equalities, tighten this rule and its regression tests.
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?out (Op (FusionEnd ?out_shape ?out_stride (F32)) (ICons ?add_elem (INil))))
+                        (= ?add_elem (Op (CudaBinaryElementwise \"Add\" ?out_shape ?sum_add_stride ?bias_add_stride ?out_stride (F32)) (ICons ?sum_fs (ICons ?bias_fs (INil)))))
+                        (= ?sum_fs (Op (FusionStart ?out_shape ?sum_add_stride (F32)) (ICons ?sum (INil))))
+                        (= ?bias_fs (Op (FusionStart ?out_shape ?bias_add_stride (F32)) (ICons ?bias (INil))))
+
+                        (= ?sum (Op (KernelSum ?matmul_out_shape ?k_dim ?sum_in_stride ?k_stride ?sum_out_stride (F32)) (ICons ?mul_fe (INil))))
+                        (= ?mul_fe (Op (FusionEnd ?mul_shape ?mul_out_stride (F32)) (ICons ?mul_elem (INil))))
+                        (= ?mul_elem (Op (CudaBinaryElementwise \"Mul\" ?mul_shape ?patch_stride ?weight_stride ?mul_out_stride (F32)) (ICons ?patch_fs (ICons ?weight_fs (INil)))))
+                        (= ?patch_fs (Op (FusionStart ?mul_shape ?patch_stride (F32)) (ICons ?patches (INil))))
+                        (= ?weight_fs (Op (FusionStart ?mul_shape ?weight_stride (F32)) (ICons ?weight (INil))))
+                        (= ?patches (Op (KernelGather ?idx_shape ?idx_stride ?input_shape ?input_stride ?gather_out_stride (F32)) (ICons ?indices (ICons ?input (INil)))))
+
+                        (= ?out_shape (ECons ?c_out (ECons ?h_out (ECons ?w_out (ENil)))))
+                        (= ?input_shape (ECons ?c_in (ECons ?h_in (ECons ?w_in (ENil)))))
+                        (= ?idx_shape (ECons ?c_in (ECons ?h_out (ECons ?w_out (ECons (MNum 1) (ECons ?kernel_h (ECons ?kernel_w (ENil))))))))
+                        (= ?matmul_out_shape (ECons ?m (ECons ?c_out (ENil))))
+                        (= ?mul_shape (ECons ?m (ECons ?c_out (ECons ?k_dim (ENil)))))
+
+                        (= ?k_stride (MIter))
+                        (= ?sum_in_stride (ECons ?sum_m_stride (ECons ?sum_c_stride (ENil))))
+                        (= ?sum_out_stride (ECons ?sum_out_m_stride (ECons ?sum_out_c_stride (ENil))))
+                        (= ?sum_add_stride (ECons ?sum_add_c_stride (ECons ?sum_add_h_stride (ECons ?sum_add_w_stride (ENil)))))
+                        (= ?weight_co_stride (nth_from_end ?weight_stride 1))
+                        (= ?weight_inner_stride (nth_from_end ?weight_stride 0))
+                        (= (MNum 0) (nth_from_end ?weight_stride 2))
+                        (= ?bias_add_stride (ECons ?bias_c_stride (ECons (MNum 0) (ECons (MNum 0) (ENil)))))
+                    )
+                    (
+                        (let ?conv (Op (KernelConv2D
+                            ?out_shape
+                            ?input_shape
+                            ?input_stride
+                            ?weight_co_stride
+                            ?weight_inner_stride
+                            ?bias_c_stride
+                            ?out_stride
+                            ?kernel_h
+                            ?kernel_w
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 0)
+                            (MNum 0)
+                            (F32))
+                            (ICons ?input (ICons ?weight (ICons ?bias (INil))))))
+                        (union ?out ?conv)
+                        (subsume (Op (FusionEnd ?out_shape ?out_stride (F32)) (ICons ?add_elem (INil))))
+                        (set (dtype ?conv) (F32))
+                    )
+                    :ruleset kernel_lower
+                    :name \"kernel conv2d from cuda lowered unfold matmul bias\"
+                )",
+            ),
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?out (Op (FusionEnd ?out_shape ?out_stride (F32)) (ICons ?add_elem (INil))))
+                        (= ?add_elem (Op (CudaBinaryElementwise \"Add\" ?out_shape ?bias_add_stride ?sum_add_stride ?out_stride (F32)) (ICons ?bias_fs (ICons ?sum_fs (INil)))))
+                        (= ?sum_fs (Op (FusionStart ?out_shape ?sum_add_stride (F32)) (ICons ?sum (INil))))
+                        (= ?bias_fs (Op (FusionStart ?out_shape ?bias_add_stride (F32)) (ICons ?bias (INil))))
+
+                        (= ?sum (Op (KernelSum ?matmul_out_shape ?k_dim ?sum_in_stride ?k_stride ?sum_out_stride (F32)) (ICons ?mul_fe (INil))))
+                        (= ?mul_fe (Op (FusionEnd ?mul_shape ?mul_out_stride (F32)) (ICons ?mul_elem (INil))))
+                        (= ?mul_elem (Op (CudaBinaryElementwise \"Mul\" ?mul_shape ?patch_stride ?weight_stride ?mul_out_stride (F32)) (ICons ?patch_fs (ICons ?weight_fs (INil)))))
+                        (= ?patch_fs (Op (FusionStart ?mul_shape ?patch_stride (F32)) (ICons ?patches (INil))))
+                        (= ?weight_fs (Op (FusionStart ?mul_shape ?weight_stride (F32)) (ICons ?weight (INil))))
+                        (= ?patches (Op (KernelGather ?idx_shape ?idx_stride ?input_shape ?input_stride ?gather_out_stride (F32)) (ICons ?indices (ICons ?input (INil)))))
+
+                        (= ?out_shape (ECons ?c_out (ECons ?h_out (ECons ?w_out (ENil)))))
+                        (= ?input_shape (ECons ?c_in (ECons ?h_in (ECons ?w_in (ENil)))))
+                        (= ?idx_shape (ECons ?c_in (ECons ?h_out (ECons ?w_out (ECons (MNum 1) (ECons ?kernel_h (ECons ?kernel_w (ENil))))))))
+                        (= ?matmul_out_shape (ECons ?m (ECons ?c_out (ENil))))
+                        (= ?mul_shape (ECons ?m (ECons ?c_out (ECons ?k_dim (ENil)))))
+
+                        (= ?k_stride (MIter))
+                        (= ?sum_in_stride (ECons ?sum_m_stride (ECons ?sum_c_stride (ENil))))
+                        (= ?sum_out_stride (ECons ?sum_out_m_stride (ECons ?sum_out_c_stride (ENil))))
+                        (= ?sum_add_stride (ECons ?sum_add_c_stride (ECons ?sum_add_h_stride (ECons ?sum_add_w_stride (ENil)))))
+                        (= ?weight_co_stride (nth_from_end ?weight_stride 1))
+                        (= ?weight_inner_stride (nth_from_end ?weight_stride 0))
+                        (= (MNum 0) (nth_from_end ?weight_stride 2))
+                        (= ?bias_add_stride (ECons ?bias_c_stride (ECons (MNum 0) (ECons (MNum 0) (ENil)))))
+                    )
+                    (
+                        (let ?conv (Op (KernelConv2D
+                            ?out_shape
+                            ?input_shape
+                            ?input_stride
+                            ?weight_co_stride
+                            ?weight_inner_stride
+                            ?bias_c_stride
+                            ?out_stride
+                            ?kernel_h
+                            ?kernel_w
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 0)
+                            (MNum 0)
+                            (F32))
+                            (ICons ?input (ICons ?weight (ICons ?bias (INil))))))
+                        (union ?out ?conv)
+                        (subsume (Op (FusionEnd ?out_shape ?out_stride (F32)) (ICons ?add_elem (INil))))
+                        (set (dtype ?conv) (F32))
+                    )
+                    :ruleset kernel_lower
+                    :name \"kernel conv2d from cuda lowered bias unfold matmul\"
+                )",
+            ),
+            // Match the im2col-style HLIR conv used by Flux2:
+            //
+            //   input.unfold([1, kh, kw], [1, 1, 1], [1, 1, 1])
+            //     -> squeeze/permute/merge view
+            //     -> matmul(weight.t())
+            //     -> split/permute view
+            //     -> + bias.expand_dim(1, h_out).expand_dim(2, w_out)
+            //
+            // The kernel consumes the pre-unfold input directly. That input may
+            // already be a padded HLIR tensor, so the rewrite is still correct
+            // for Flux2's padded convs while removing the large patch matrix.
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?add (Op (Add ?out_shape ?sum_add_stride ?bias_add_stride ?add_out_stride) (ICons ?sum (ICons ?bias (INil)))))
+                        (= ?sum (Op (Sum ?matmul_out_shape ?k_dim ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+                        (= ?mul (Op (Mul ?mul_shape ?patch_stride ?weight_stride ?mul_out_stride) (ICons ?patches (ICons ?weight (INil)))))
+                        (= ?patches (Op (Gather ?idx_shape ?idx_stride ?input_shape ?input_stride) (ICons ?indices (ICons ?input (INil)))))
+
+                        (= ?out_shape (ECons ?c_out (ECons ?h_out (ECons ?w_out (ENil)))))
+                        (= ?input_shape (ECons ?c_in (ECons ?h_in (ECons ?w_in (ENil)))))
+                        (= ?idx_shape (ECons ?c_in (ECons ?h_out (ECons ?w_out (ECons (MNum 1) (ECons ?kernel_h (ECons ?kernel_w (ENil))))))))
+                        (= ?matmul_out_shape (ECons ?m (ECons ?c_out (ENil))))
+
+                        ; This rewrite is for stride=1, dilation=1 over the
+                        ; tensor passed to unfold. Padded HLIR inputs are already
+                        ; represented as their own tensor, so padding is 0 here.
+                        (= ?h_out (MAdd (MSub ?h_in ?kernel_h) (MNum 1)))
+                        (= ?w_out (MAdd (MSub ?w_in ?kernel_w) (MNum 1)))
+                        (= ?m (MMul ?h_out ?w_out))
+                        (= ?k_dim (MMul ?c_in (MMul ?kernel_h ?kernel_w)))
+                        (= ?k_stride (MIter))
+
+                        (= ?weight_co_stride (nth_from_end ?weight_stride 1))
+                        (= ?weight_inner_stride (nth_from_end ?weight_stride 0))
+                        (= (MNum 0) (nth_from_end ?weight_stride 2))
+
+                        (= ?bias_add_stride (ECons ?bias_c_stride (ECons (MNum 0) (ECons (MNum 0) (ENil)))))
+
+                        (= (F32) (dtype ?input))
+                        (= (F32) (dtype ?weight))
+                        (= (F32) (dtype ?bias))
+                    )
+                    (
+                        (let ?conv (Op (KernelConv2D
+                            ?out_shape
+                            ?input_shape
+                            ?input_stride
+                            ?weight_co_stride
+                            ?weight_inner_stride
+                            ?bias_c_stride
+                            ?add_out_stride
+                            ?kernel_h
+                            ?kernel_w
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 0)
+                            (MNum 0)
+                            (F32))
+                            (ICons ?input (ICons ?weight (ICons ?bias (INil))))))
+                        (union ?add ?conv)
+                        (subsume (Op (Add ?out_shape ?sum_add_stride ?bias_add_stride ?add_out_stride) (ICons ?sum (ICons ?bias (INil)))))
+                        (set (dtype ?conv) (F32))
+                    )
+                    :ruleset kernel_specialize
+                    :name \"kernel conv2d from unfold matmul bias\"
+                )",
+            ),
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?add (Op (Add ?out_shape ?bias_add_stride ?sum_add_stride ?add_out_stride) (ICons ?bias (ICons ?sum (INil)))))
+                        (= ?sum (Op (Sum ?matmul_out_shape ?k_dim ?sum_in_stride ?k_stride ?sum_out_stride) (ICons ?mul (INil))))
+                        (= ?mul (Op (Mul ?mul_shape ?patch_stride ?weight_stride ?mul_out_stride) (ICons ?patches (ICons ?weight (INil)))))
+                        (= ?patches (Op (Gather ?idx_shape ?idx_stride ?input_shape ?input_stride) (ICons ?indices (ICons ?input (INil)))))
+
+                        (= ?out_shape (ECons ?c_out (ECons ?h_out (ECons ?w_out (ENil)))))
+                        (= ?input_shape (ECons ?c_in (ECons ?h_in (ECons ?w_in (ENil)))))
+                        (= ?idx_shape (ECons ?c_in (ECons ?h_out (ECons ?w_out (ECons (MNum 1) (ECons ?kernel_h (ECons ?kernel_w (ENil))))))))
+                        (= ?matmul_out_shape (ECons ?m (ECons ?c_out (ENil))))
+
+                        (= ?h_out (MAdd (MSub ?h_in ?kernel_h) (MNum 1)))
+                        (= ?w_out (MAdd (MSub ?w_in ?kernel_w) (MNum 1)))
+                        (= ?m (MMul ?h_out ?w_out))
+                        (= ?k_dim (MMul ?c_in (MMul ?kernel_h ?kernel_w)))
+                        (= ?k_stride (MIter))
+
+                        (= ?weight_co_stride (nth_from_end ?weight_stride 1))
+                        (= ?weight_inner_stride (nth_from_end ?weight_stride 0))
+                        (= (MNum 0) (nth_from_end ?weight_stride 2))
+
+                        (= ?bias_add_stride (ECons ?bias_c_stride (ECons (MNum 0) (ECons (MNum 0) (ENil)))))
+
+                        (= (F32) (dtype ?input))
+                        (= (F32) (dtype ?weight))
+                        (= (F32) (dtype ?bias))
+                    )
+                    (
+                        (let ?conv (Op (KernelConv2D
+                            ?out_shape
+                            ?input_shape
+                            ?input_stride
+                            ?weight_co_stride
+                            ?weight_inner_stride
+                            ?bias_c_stride
+                            ?add_out_stride
+                            ?kernel_h
+                            ?kernel_w
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 1)
+                            (MNum 0)
+                            (MNum 0)
+                            (F32))
+                            (ICons ?input (ICons ?weight (ICons ?bias (INil))))))
+                        (union ?add ?conv)
+                        (subsume (Op (Add ?out_shape ?bias_add_stride ?sum_add_stride ?add_out_stride) (ICons ?bias (ICons ?sum (INil)))))
+                        (set (dtype ?conv) (F32))
+                    )
+                    :ruleset kernel_specialize
+                    :name \"kernel conv2d from bias unfold matmul\"
+                )",
+            ),
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?add (Op (Add ?shape ?as ?bs ?os) ?inputs))
+                        (= ?add (Op (KernelConv2D ?out_shape ?input_shape ?input_stride ?wco ?wi ?bc ?out_stride ?kh ?kw ?sh ?sw ?dh ?dw ?ph ?pw ?dt) ?conv_inputs))
+                    )
+                    ((delete (Op (Add ?shape ?as ?bs ?os) ?inputs)))
+                    :ruleset cleanup
+                )",
+            ),
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?fe (Op (FusionEnd ?shape ?os ?dt) ?inputs))
+                        (= ?fe (Op (KernelConv2D ?out_shape ?input_shape ?input_stride ?wco ?wi ?bc ?out_stride ?kh ?kw ?sh ?sw ?dh ?dw ?ph ?pw ?conv_dt) ?conv_inputs))
+                    )
+                    ((delete (Op (FusionEnd ?shape ?os ?dt) ?inputs)))
+                    :ruleset cleanup
+                )",
+            ),
+        ]
+    }
+
+    fn cleanup(&self) -> bool {
+        false
+    }
+
+    fn extract<'a>(
+        &'a self,
+        egraph: &'a luminal::egglog_utils::SerializedEGraph,
+        kind_children: &[&'a luminal::egglog_utils::NodeId],
+        input_enodes: Vec<&'a luminal::egglog_utils::NodeId>,
+        list_cache: &mut FxHashMap<&'a luminal::egglog_utils::NodeId, Vec<Expression>>,
+        expr_cache: &mut FxHashMap<&'a luminal::egglog_utils::NodeId, Expression>,
+    ) -> (LLIROp, Vec<&'a luminal::egglog_utils::NodeId>) {
+        (
+            LLIROp::new::<dyn KernelOp>(Box::new(Self {
+                out_shape: extract_expr_list(egraph, kind_children[0], list_cache, expr_cache)
+                    .unwrap(),
+                input_shape: extract_expr_list(egraph, kind_children[1], list_cache, expr_cache)
+                    .unwrap(),
+                input_stride: extract_expr_list(egraph, kind_children[2], list_cache, expr_cache)
+                    .unwrap(),
+                weight_co_stride: extract_expr(egraph, kind_children[3], expr_cache).unwrap(),
+                weight_inner_stride: extract_expr(egraph, kind_children[4], expr_cache).unwrap(),
+                bias_c_stride: extract_expr(egraph, kind_children[5], expr_cache).unwrap(),
+                out_stride: extract_expr_list(egraph, kind_children[6], list_cache, expr_cache)
+                    .unwrap(),
+                kernel_h: extract_expr(egraph, kind_children[7], expr_cache).unwrap(),
+                kernel_w: extract_expr(egraph, kind_children[8], expr_cache).unwrap(),
+                stride_h: extract_expr(egraph, kind_children[9], expr_cache).unwrap(),
+                stride_w: extract_expr(egraph, kind_children[10], expr_cache).unwrap(),
+                dilation_h: extract_expr(egraph, kind_children[11], expr_cache).unwrap(),
+                dilation_w: extract_expr(egraph, kind_children[12], expr_cache).unwrap(),
+                pad_h: extract_expr(egraph, kind_children[13], expr_cache).unwrap(),
+                pad_w: extract_expr(egraph, kind_children[14], expr_cache).unwrap(),
+                dtype: extract_dtype(egraph, kind_children[15]),
+            }) as Box<dyn KernelOp>),
+            input_enodes,
+        )
+    }
+}
+
+impl KernelOp for KernelConv2D {
+    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>>,
+    ) {
+        assert_eq!(self.dtype, DType::F32, "KernelConv2D currently emits F32");
+
+        let vars: FxHashSet<char> = self
+            .out_shape
+            .iter()
+            .chain(&self.input_shape)
+            .chain(&self.input_stride)
+            .chain(&self.out_stride)
+            .flat_map(|e| e.dyn_vars())
+            .chain(self.weight_co_stride.dyn_vars())
+            .chain(self.weight_inner_stride.dyn_vars())
+            .chain(self.bias_c_stride.dyn_vars())
+            .chain(self.kernel_h.dyn_vars())
+            .chain(self.kernel_w.dyn_vars())
+            .chain(self.stride_h.dyn_vars())
+            .chain(self.stride_w.dyn_vars())
+            .chain(self.dilation_h.dyn_vars())
+            .chain(self.dilation_w.dyn_vars())
+            .chain(self.pad_h.dyn_vars())
+            .chain(self.pad_w.dyn_vars())
+            .collect();
+
+        let (dyn_defines, _sorted_dims) = generate_dyn_dims_defines(&vars);
+        let dyn_dims_param = if vars.is_empty() {
+            ""
+        } else {
+            ", const int* dyn_dims"
+        };
+
+        let c_out = self.out_shape[0].to_kernel();
+        let h_out = self.out_shape[1].to_kernel();
+        let w_out = self.out_shape[2].to_kernel();
+        let c_in = self.input_shape[0].to_kernel();
+        let h_in = self.input_shape[1].to_kernel();
+        let w_in = self.input_shape[2].to_kernel();
+        let weight_co_stride = self
+            .weight_co_stride
+            .substitute('z', Expression::from(1))
+            .simplify()
+            .to_kernel();
+        let weight_inner_stride = self
+            .weight_inner_stride
+            .substitute('z', Expression::from(1))
+            .simplify()
+            .to_kernel();
+        let bias_c_stride = self
+            .bias_c_stride
+            .substitute('z', Expression::from(1))
+            .simplify()
+            .to_kernel();
+        let kh = self.kernel_h.to_kernel();
+        let kw = self.kernel_w.to_kernel();
+        let stride_h = self.stride_h.to_kernel();
+        let stride_w = self.stride_w.to_kernel();
+        let dilation_h = self.dilation_h.to_kernel();
+        let dilation_w = self.dilation_w.to_kernel();
+        let pad_h = self.pad_h.to_kernel();
+        let pad_w = self.pad_w.to_kernel();
+        let out_idx = flatten_strides(&self.out_shape, &self.out_stride).to_kernel();
+        let input_idx = flatten_strides(&self.input_shape, &self.input_stride)
+            .to_kernel()
+            .replace("const_z", "input_linear");
+        let n_outputs: Expression = self.out_shape.iter().copied().product();
+
+        let kernel = format!(
+            "
+{dyn_defines}
+extern \"C\" {{
+    __global__ void generic_conv2d_bias(
+        float* __restrict__ out,
+        const float* __restrict__ input,
+        const float* __restrict__ weight,
+        const float* __restrict__ bias{dyn_dims_param}
+    ) {{
+        long long const_z = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+        const long long total = {total};
+        if (const_z >= total) return;
+
+        const long long COUT = {c_out};
+        const long long HOUT = {h_out};
+        const long long WOUT = {w_out};
+        const long long CIN = {c_in};
+        const long long HIN = {h_in};
+        const long long WIN = {w_in};
+        const long long KH = {kh};
+        const long long KW = {kw};
+        const long long SH = {stride_h};
+        const long long SW = {stride_w};
+        const long long DH = {dilation_h};
+        const long long DW = {dilation_w};
+        const long long PH = {pad_h};
+        const long long PW = {pad_w};
+        const long long W_CO_STRIDE = {weight_co_stride};
+        const long long W_INNER_STRIDE = {weight_inner_stride};
+        const long long BIAS_C_STRIDE = {bias_c_stride};
+
+        long long co = const_z / (HOUT * WOUT);
+        long long rem = const_z - co * HOUT * WOUT;
+        long long oh = rem / WOUT;
+        long long ow = rem - oh * WOUT;
+
+        float acc = bias[co * BIAS_C_STRIDE];
+        for (long long ci = 0; ci < CIN; ++ci) {{
+            for (long long r = 0; r < KH; ++r) {{
+                long long ih = oh * SH + r * DH - PH;
+                if (ih < 0 || ih >= HIN) continue;
+                for (long long s = 0; s < KW; ++s) {{
+                    long long iw = ow * SW + s * DW - PW;
+                    if (iw < 0 || iw >= WIN) continue;
+                    long long input_linear = (ci * HIN + ih) * WIN + iw;
+                    long long input_idx = {input_idx};
+                    long long inner = (ci * KH + r) * KW + s;
+                    long long weight_idx = co * W_CO_STRIDE + inner * W_INNER_STRIDE;
+                    acc += input[input_idx] * weight[weight_idx];
+                }}
+            }}
+        }}
+        out[{out_idx}] = acc;
+    }}
+}}",
+            total = n_outputs.to_kernel(),
+        );
+
+        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("generic_conv2d_bias").unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        (
+            func,
+            module,
+            kernel,
+            (n_outputs.ceil_div(256), 1.into(), 1.into()),
+            (n_outputs.min(256), 1.into(), 1.into()),
+            0.into(),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        self.out_shape.iter().copied().product()
+    }
+
+    fn all_dyn_vars(&self) -> FxHashSet<char> {
+        self.out_shape
+            .iter()
+            .chain(&self.input_shape)
+            .chain(&self.input_stride)
+            .chain(&self.out_stride)
+            .flat_map(|e| e.dyn_vars())
+            .chain(self.weight_co_stride.dyn_vars())
+            .chain(self.weight_inner_stride.dyn_vars())
+            .chain(self.bias_c_stride.dyn_vars())
+            .chain(self.kernel_h.dyn_vars())
+            .chain(self.kernel_w.dyn_vars())
+            .chain(self.stride_h.dyn_vars())
+            .chain(self.stride_w.dyn_vars())
+            .chain(self.dilation_h.dyn_vars())
+            .chain(self.dilation_w.dyn_vars())
+            .chain(self.pad_h.dyn_vars())
+            .chain(self.pad_w.dyn_vars())
+            .collect()
+    }
+
+    fn output_bytes(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        let c_in = self.input_shape[0];
+        self.output_size() * self.kernel_h * self.kernel_w * c_in * 2 * 4 + self.output_size() * 4
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn flops(&self) -> Expression {
+        let c_in = self.input_shape[0];
+        self.output_size() * self.kernel_h * self.kernel_w * c_in * 2
+    }
+
+    fn output_dtype(&self) -> DType {
+        self.dtype
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "GenericConv2D"
+    }
+}

crates/luminal_cuda_lite/src/kernel/dlrm_interact.rs

@@ -0,0 +1,483 @@
+//! Fused DLRM pairwise-dot interaction.
+//!
+//! Replaces the cat→bmm(T,Tᵀ)→tril-gather chain with a single kernel
+//! that reads N separate `(batch, d)` tensors and writes the strict
+//! lower-triangular pairwise dot products directly into the output —
+//! `out[b, p] = Σ_d v_i[b, d] * v_j[b, d]` for each ordered pair (i, j)
+//! with i > j.
+//!
+//! Why this matters for the DLRM forward: the natural luminal lowering
+//! materializes the `(B, F, D)` stacked tensor, then the full `(B, F, F)`
+//! BMM output, then a flat gather to pull out F(F-1)/2 pairs. That's
+//! ~12 small kernels and an `F²·B` intermediate even though only half
+//! of those elements are kept. The fused version uses N pointer args
+//! (one per feature vector), computes only the F(F-1)/2 dot products,
+//! and writes directly to the final `(B, F(F-1)/2)` buffer.
+//!
+//! All shapes are static. The kernel source is generated with the
+//! exact pair table baked in (so the inner loop is a fixed `D`-element
+//! reduction with no shape-dependent branching).
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::{
+    dtype::DType, op::CustomOp, op::LLIROp, prelude::FxHashMap, prelude::GraphTensor,
+    shape::Expression,
+};
+
+use crate::compile_module_image_for_current_device;
+use crate::kernel::KernelOp;
+
+#[derive(Debug, Clone)]
+pub struct PairwiseDotLowerTriKernel {
+    pub batch: usize,
+    pub num_features: usize, // F
+    pub d: usize,
+}
+
+impl PairwiseDotLowerTriKernel {
+    fn pair_count(&self) -> usize {
+        self.num_features * (self.num_features - 1) / 2
+    }
+}
+
+impl KernelOp for PairwiseDotLowerTriKernel {
+    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 f = self.num_features;
+        let p = self.pair_count();
+        // Pair table (i, j) with i > j, in strict-lower-tri (row-major over
+        // i then j) order — same convention as torch.tril_indices(F, F, -1).
+        let mut pairs: Vec<(usize, usize)> = Vec::with_capacity(p);
+        for i in 0..f {
+            for j in 0..i {
+                pairs.push((i, j));
+            }
+        }
+        // Build kernel params signature: one pointer per input feature.
+        let in_params: String = (0..f)
+            .map(|k| format!(", const float* __restrict__ v{k}"))
+            .collect::<Vec<_>>()
+            .concat();
+        // For each pair p, generate one branch in the switch that selects
+        // the two input pointers to dot-product. With F small (DLRM has
+        // F=4), the branch is fully unrolled.
+        let mut pair_switch = String::new();
+        for (pidx, (i, j)) in pairs.iter().enumerate() {
+            pair_switch += &format!(
+                "        case {pidx}: pa = v{i}; pb = v{j}; break;\n"
+            );
+        }
+
+        let kernel = format!(
+            "
+extern \"C\" __global__ void dlrm_pairwise_dot_lower_tri_kernel(
+    float* __restrict__ out{in_params}
+) {{
+    const int B = {batch};
+    const int D = {d};
+    const int P = {p};
+    int b = blockIdx.x;
+    int p = blockIdx.y;
+    int t = threadIdx.x;
+    if (b >= B || p >= P) return;
+
+    const float* pa = nullptr;
+    const float* pb = nullptr;
+    switch (p) {{
+{pair_switch}
+        default: return;
+    }}
+
+    // Block-wide reduction of dot(pa[b], pb[b]) over D using shared mem.
+    extern __shared__ float smem[];
+    float partial = 0.0f;
+    for (int d = t; d < D; d += blockDim.x) {{
+        partial += pa[b * D + d] * pb[b * D + d];
+    }}
+    smem[t] = partial;
+    __syncthreads();
+    // Power-of-two tree reduce. blockDim.x must be a power of two.
+    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {{
+        if (t < stride) {{
+            smem[t] += smem[t + stride];
+        }}
+        __syncthreads();
+    }}
+    if (t == 0) {{
+        out[b * P + p] = smem[0];
+    }}
+}}
+",
+            batch = self.batch,
+            d = self.d,
+            p = p,
+            pair_switch = pair_switch,
+            in_params = in_params,
+        );
+
+        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("dlrm_pairwise_dot_lower_tri_kernel")
+                .unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        // Pick a power-of-two thread count ≤ D, ≥ 32 where possible.
+        let mut threads = 1usize;
+        while threads * 2 <= self.d.max(32) {
+            threads *= 2;
+        }
+        let threads = threads.max(32).min(1024);
+        (
+            func,
+            module,
+            kernel,
+            (
+                Expression::from(self.batch),
+                Expression::from(p),
+                Expression::from(1usize),
+            ),
+            (
+                Expression::from(threads),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            Expression::from(threads * 4),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        Expression::from(self.batch * self.pair_count())
+    }
+
+    fn output_bytes(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn output_dtype(&self) -> DType {
+        DType::F32
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        // Each pair reads 2 vectors of D floats per batch row. F-choose-2
+        // pairs, so per-batch each input vector is read F-1 times.
+        Expression::from(self.batch * self.num_features * (self.num_features - 1) * self.d * 4)
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn flops(&self) -> Expression {
+        // 2D-1 flops per dot product (D mul + D-1 add).
+        Expression::from(self.batch * self.pair_count() * (2 * self.d - 1))
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "DLRMPairwiseDotLowerTri"
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct PairwiseDotLowerTriCustom(pub PairwiseDotLowerTriKernel);
+
+impl CustomOp for PairwiseDotLowerTriCustom {
+    fn to_llir_op(&self) -> LLIROp {
+        LLIROp::new::<dyn KernelOp>(Box::new(self.0.clone()) as Box<dyn KernelOp>)
+    }
+}
+
+/// Two-input variant of [`PairwiseDotLowerTriKernel`] that consumes the
+/// dense MLP output and a stacked embedding output without requiring
+/// the caller to first slice the stack into individual (B, D) views.
+///
+/// Treats feature 0 as `dense_out[b, t]` and features 1..=num_emb as
+/// `emb_stack[b, k-1, t]`. Output pair table is the strict lower tri
+/// of an `F × F` matrix where `F = num_emb + 1`.
+#[derive(Debug, Clone)]
+pub struct PairwiseDotLowerTriStackedKernel {
+    pub batch: usize,
+    pub num_emb: usize, // N (excluding the dense feature)
+    pub d: usize,
+}
+
+impl PairwiseDotLowerTriStackedKernel {
+    fn num_features(&self) -> usize {
+        self.num_emb + 1
+    }
+    fn pair_count(&self) -> usize {
+        let f = self.num_features();
+        f * (f - 1) / 2
+    }
+}
+
+impl KernelOp for PairwiseDotLowerTriStackedKernel {
+    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 f = self.num_features();
+        let p = self.pair_count();
+        let n_emb = self.num_emb;
+        let d_ = self.d;
+
+        // Block-per-batch layout. Each block:
+        //   1. Cooperatively loads all F feature vectors for batch b into
+        //      shared memory once — F*D floats total. Feature 0 = dense[b];
+        //      features 1..F = emb_stack[b, k-1, :].
+        //   2. Each thread `tid` strides over pairs `p = tid, tid+blockDim.x,
+        //      …, P-1`. For each, derives (i, j) such that i > j and writes
+        //      the dot product of feat[i] and feat[j].
+        //
+        // Compared to the previous (B, P) grid-of-one-block-per-output
+        // layout this:
+        //   - Cuts launch count by P× (e.g. 528× at num_cat=32).
+        //   - Reads each feature vector once per batch instead of (F-1)
+        //     times — F(F-1) reads → F reads, an ~(F-1)/2× memory traffic
+        //     reduction (e.g. 16× at num_cat=32, F=33).
+        //   - Reuses cached features across all P pairs at shared-memory
+        //     latency instead of refetching from global per pair.
+        //
+        // Pair-index → (i, j) is computed from `p` directly using the
+        // closed-form for strict lower-tri row indexing:
+        //   row i contains i pairs (j ∈ [0, i)); cumulative row starts
+        //   at `i*(i-1)/2`; so `i = floor((1+sqrt(1+8p))/2)` and
+        //   `j = p - i*(i-1)/2`. We do a tiny defensive adjustment
+        //   afterwards to absorb sqrtf rounding.
+        let kernel = format!(
+            "
+extern \"C\" __global__ void dlrm_pairwise_dot_lower_tri_stacked_kernel(
+    float* __restrict__ out,
+    const float* __restrict__ dense,      // (B, D)
+    const float* __restrict__ emb_stack   // (B, N, D)
+) {{
+    const int B = {batch};
+    const int D = {d};
+    const int N = {n_emb};
+    const int F = {f};
+    const int P = {p};
+    int b = blockIdx.x;
+    int tid = threadIdx.x;
+    int tcount = blockDim.x;
+    if (b >= B) return;
+
+    // Shared feature cache: F * D floats.
+    extern __shared__ float feat[];
+    for (int i = tid; i < F * D; i += tcount) {{
+        int feat_idx = i / D;
+        int dim = i - feat_idx * D;
+        if (feat_idx == 0) {{
+            feat[i] = dense[b * D + dim];
+        }} else {{
+            int slot = feat_idx - 1;
+            feat[i] = emb_stack[(b * N + slot) * D + dim];
+        }}
+    }}
+    __syncthreads();
+
+    // Each thread handles a strided slice of the P pairs.
+    for (int p = tid; p < P; p += tcount) {{
+        float t = sqrtf(8.0f * (float)p + 1.0f);
+        int pi = (int)((t + 1.0f) * 0.5f);
+        // Adjust for fp rounding — pi*(pi-1)/2 must be the largest
+        // row-start ≤ p.
+        while (pi * (pi - 1) / 2 > p) pi--;
+        while ((pi + 1) * pi / 2 <= p) pi++;
+        int pj = p - pi * (pi - 1) / 2;
+
+        float acc = 0.0f;
+        #pragma unroll
+        for (int d = 0; d < {d}; ++d) {{
+            acc += feat[pi * {d} + d] * feat[pj * {d} + d];
+        }}
+        out[b * P + p] = acc;
+    }}
+}}
+",
+            batch = self.batch,
+            d = d_,
+            n_emb = n_emb,
+            f = f,
+            p = p,
+        );
+
+        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("dlrm_pairwise_dot_lower_tri_stacked_kernel")
+                .unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        // Block size: enough threads to cover both the feature-load phase
+        // (F*D elements) and the pair computation (P elements) without
+        // serial waves dominating, capped at 1024 (max CUDA block size)
+        // and rounded down to a multiple of 32 for warp alignment.
+        let want = std::cmp::max(f * d_, p);
+        let threads = want.clamp(32, 1024).next_multiple_of(32);
+        let threads = threads.min(1024);
+        let shared_bytes = f * d_ * 4;
+        (
+            func,
+            module,
+            kernel,
+            (
+                Expression::from(self.batch),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            (
+                Expression::from(threads),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            Expression::from(shared_bytes),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        Expression::from(self.batch * self.pair_count())
+    }
+
+    fn output_bytes(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn output_dtype(&self) -> DType {
+        DType::F32
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        Expression::from(self.batch * self.num_features() * (self.num_features() - 1) * self.d * 4)
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn flops(&self) -> Expression {
+        Expression::from(self.batch * self.pair_count() * (2 * self.d - 1))
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "DLRMPairwiseDotLowerTriStacked"
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct PairwiseDotLowerTriStackedCustom(pub PairwiseDotLowerTriStackedKernel);
+
+impl CustomOp for PairwiseDotLowerTriStackedCustom {
+    fn to_llir_op(&self) -> LLIROp {
+        LLIROp::new::<dyn KernelOp>(Box::new(self.0.clone()) as Box<dyn KernelOp>)
+    }
+}
+
+/// Pairwise lower-tri dot product over `dense_out` plus a stacked
+/// embedding output. Avoids the per-table slice that the variadic
+/// variant would otherwise need to materialize.
+///
+/// * `dense_out`: `(batch, d)` — feature 0 in the pair table.
+/// * `emb_stack`: `(batch, num_emb, d)` — features 1..=num_emb.
+///
+/// Returns `(batch, (num_emb+1) * num_emb / 2)`, same strict-lower-tri
+/// ordering as [`dlrm_pairwise_dot_lower_tri`].
+pub fn dlrm_pairwise_dot_lower_tri_stacked(
+    dense_out: GraphTensor,
+    emb_stack: GraphTensor,
+) -> GraphTensor {
+    assert_eq!(dense_out.dtype, DType::F32, "dense_out must be F32");
+    assert_eq!(emb_stack.dtype, DType::F32, "emb_stack must be F32");
+    let dd = dense_out.dims();
+    let sd = emb_stack.dims();
+    assert_eq!(dd.len(), 2, "dense_out must be 2D");
+    assert_eq!(sd.len(), 3, "emb_stack must be 3D (batch, num_emb, d)");
+    let batch = dd[0].to_usize().expect("batch must be static");
+    let d = dd[1].to_usize().expect("d must be static");
+    assert_eq!(sd[0].to_usize().unwrap(), batch);
+    let num_emb = sd[1].to_usize().expect("num_emb must be static");
+    assert_eq!(sd[2].to_usize().unwrap(), d);
+    let kern = PairwiseDotLowerTriStackedKernel {
+        batch,
+        num_emb,
+        d,
+    };
+    let f = num_emb + 1;
+    let p = f * (f - 1) / 2;
+    let cx = unsafe { &mut *dense_out.graph_ref };
+    cx.custom_op(
+        PairwiseDotLowerTriStackedCustom(kern),
+        vec![dense_out, emb_stack],
+        (batch, p),
+        DType::F32,
+    )
+}
+
+/// Strict-lower-triangular pairwise dot product of N feature vectors.
+///
+/// * `features`: N tensors, each `(batch, d)`, all F32, all the same shape.
+///
+/// Returns `(batch, N*(N-1)/2)` with pair ordering matching
+/// `torch.tril_indices(N, N, -1)` (row-major: (1,0), (2,0), (2,1), …).
+pub fn dlrm_pairwise_dot_lower_tri(features: Vec<GraphTensor>) -> GraphTensor {
+    assert!(features.len() >= 2, "need at least 2 feature vectors");
+    let first = features[0];
+    let dims = first.dims();
+    assert_eq!(dims.len(), 2, "each feature vector must be 2D (batch, d)");
+    let batch = dims[0].to_usize().expect("batch must be static");
+    let d = dims[1].to_usize().expect("d must be static");
+    let f = features.len();
+    for v in &features {
+        assert_eq!(v.dtype, DType::F32, "features must all be F32");
+        let vd = v.dims();
+        assert_eq!(vd.len(), 2, "features must all be 2D");
+        assert_eq!(vd[0].to_usize().unwrap(), batch, "batch mismatch");
+        assert_eq!(vd[1].to_usize().unwrap(), d, "d mismatch");
+    }
+    let kern = PairwiseDotLowerTriKernel {
+        batch,
+        num_features: f,
+        d,
+    };
+    let p = f * (f - 1) / 2;
+    let cx = unsafe { &mut *first.graph_ref };
+    cx.custom_op(
+        PairwiseDotLowerTriCustom(kern),
+        features,
+        (batch, p),
+        DType::F32,
+    )
+}

crates/luminal_cuda_lite/src/kernel/embedding_bag.rs

@@ -0,0 +1,757 @@
+//! Single-kernel fused EmbeddingBag (sum-pool) operator.
+//!
+//! DLRM-style embedding lookups in luminal currently lower into a chain
+//! of broadcast-iota + multiply + add + Gather + SumReduce kernels (~6
+//! kernels per table). For a model with even a handful of tables that
+//! eats most of the per-iter launch budget once everything else is
+//! captured into a single CUDA graph.
+//!
+//! This op collapses the whole pattern — `gather(table, idx) → sum(L)` —
+//! into one kernel. Same template as `Matmul2DKernel`: implement
+//! [`KernelOp`], wrap in a [`CustomOp`] so the user-facing call comes
+//! out as a `dyn KernelOp` in the LLIR (which means it can be absorbed
+//! into the same CudaGraphOp as everything around it — no extra host
+//! op, no extra CUDA launch outside the graph).
+//!
+//! Semantics: `out[b, d] = Σ_l table[indices[b, l], d]` with
+//!   table: (n_emb, d), F32, row-major
+//!   indices: (batch, bag), I32, row-major
+//!   out: (batch, d), F32, row-major
+//!
+//! Fixed-shape: `n_emb`, `d`, `batch`, `bag` are static (baked into
+//! the kernel source via #defines), matching how the rest of the
+//! `kernel::` ops in this crate handle shape.
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr};
+use luminal::{
+    dtype::DType, op::CustomOp, op::LLIROp, prelude::FxHashMap, prelude::GraphTensor,
+    shape::Expression,
+};
+
+use crate::compile_module_image_for_current_device;
+use crate::kernel::KernelOp;
+
+/// One-kernel fused EmbeddingBag with sum pooling and fixed bag size.
+#[derive(Debug, Clone)]
+pub struct EmbeddingBagSumKernel {
+    pub batch: usize,
+    pub bag: usize,
+    pub d: usize,
+    pub n_emb: usize,
+}
+
+impl KernelOp for EmbeddingBagSumKernel {
+    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>>,
+    ) {
+        // One block per batch row, `d` threads per block. Each thread sums
+        // one output column over the `bag` indices. This is the standard
+        // bag-size-1..L pattern and is memory-bandwidth bound on `table`,
+        // which is exactly the right roofline for this op.
+        let kernel = format!(
+            "
+extern \"C\" __global__ void embedding_bag_sum_kernel(
+    float* __restrict__ out,
+    const float* __restrict__ table,
+    const int* __restrict__ indices
+) {{
+    const int B = {batch};
+    const int L = {bag};
+    const int D = {d};
+    const int N = {n_emb};
+    int b = blockIdx.x;
+    int d = threadIdx.x;
+    if (b >= B || d >= D) return;
+    float acc = 0.0f;
+    #pragma unroll 4
+    for (int l = 0; l < L; ++l) {{
+        int row = indices[b * L + l];
+        // Index is from user input; trust it (matches torch.EmbeddingBag).
+        acc += table[row * D + d];
+    }}
+    out[b * D + d] = acc;
+    (void)N;
+}}
+",
+            batch = self.batch,
+            bag = self.bag,
+            d = self.d,
+            n_emb = self.n_emb,
+        );
+
+        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("embedding_bag_sum_kernel").unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        (
+            func,
+            module,
+            kernel,
+            (
+                Expression::from(self.batch),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            (
+                Expression::from(self.d),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            Expression::from(0usize),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        Expression::from(self.batch * self.d)
+    }
+
+    fn output_bytes(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn output_dtype(&self) -> DType {
+        DType::F32
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        // For each output element, L reads from table (4 bytes each), plus
+        // L reads from indices (4 bytes each, shared across D threads — we
+        // just bill once per output to keep this readable).
+        Expression::from(self.batch * self.d * self.bag * 4 + self.batch * self.bag * 4)
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn flops(&self) -> Expression {
+        // L adds per output element. Pointer math doesn't count.
+        Expression::from(self.batch * self.d * self.bag)
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "EmbeddingBagSum"
+    }
+}
+
+/// CustomOp wrapper for [`EmbeddingBagSumKernel`].
+#[derive(Debug, Clone)]
+pub struct EmbeddingBagSumCustom(pub EmbeddingBagSumKernel);
+
+impl CustomOp for EmbeddingBagSumCustom {
+    fn to_llir_op(&self) -> LLIROp {
+        LLIROp::new::<dyn KernelOp>(Box::new(self.0.clone()) as Box<dyn KernelOp>)
+    }
+}
+
+/// One-kernel fused multi-table EmbeddingBag with sum pooling.
+///
+/// Folds all `num_tables` independent embedding lookups into a single
+/// CUDA kernel launch. Reads from one big weight tensor that is the
+/// row-wise concatenation of every table; per-table row offsets are
+/// baked into the kernel source. Per-table index tensors stay separate.
+/// Output is `(batch, num_tables, d)` so downstream ops can consume it
+/// as a single stacked tensor (matches v3's `index_select + reshape`
+/// trick — Inductor fuses gather+sum across all tables; this kernel
+/// just does it directly).
+#[derive(Debug, Clone)]
+pub struct StackedEmbeddingBagKernel {
+    pub batch: usize,
+    pub bag: usize,
+    pub d: usize,
+    pub num_tables: usize,
+    /// Cumulative row counts: `row_offsets[k]` = number of rows in all
+    /// tables strictly before table `k`. Length = `num_tables + 1`.
+    /// `row_offsets[num_tables]` = total rows in the stacked weight.
+    pub row_offsets: Vec<usize>,
+}
+
+impl KernelOp for StackedEmbeddingBagKernel {
+    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>>,
+    ) {
+        assert_eq!(
+            self.row_offsets.len(),
+            self.num_tables + 1,
+            "row_offsets must have num_tables+1 entries"
+        );
+        // One index pointer per table — variadic via generated kernel signature.
+        let idx_params: String = (0..self.num_tables)
+            .map(|k| format!(", const int* __restrict__ idx_{k}"))
+            .collect::<Vec<_>>()
+            .concat();
+        // For each table k, generate a `case k` branch that picks the right
+        // index pointer and row offset. The case body is the same fused
+        // gather+sum loop as the single-table kernel.
+        let mut switch = String::new();
+        for k in 0..self.num_tables {
+            let off = self.row_offsets[k];
+            switch += &format!(
+                "            case {k}: {{ const int* __restrict__ idx_ptr = idx_{k}; const int row_off = {off}; for (int l = 0; l < L; ++l) {{ int row = idx_ptr[b * L + l] + row_off; acc += weight[row * D + d]; }} break; }}\n"
+            );
+        }
+
+        // Grid is (B,); one block per batch row. Block holds *all* (k, d)
+        // output threads together. The previous (B, N) grid had 16-thread
+        // blocks at D=16, which left each SM under-occupied (Hopper's
+        // max-blocks-per-SM × 16 threads ≪ 64 warps/SM, so the warp
+        // scheduler couldn't hide memory latency). With one batch row
+        // per block we get K·D threads (e.g. 512 at K=32, D=16), which
+        // is 16 warps — enough for the SM to overlap pending loads with
+        // compute on other warps. Each block now produces (K, D) outputs
+        // instead of (1, D), so total block count drops from B·K to B
+        // (e.g. 65k → 2k at K=32, B=2048).
+        //
+        // Threads stride over `total = K · D` if the requested block
+        // size exceeds 1024 (CUDA max). At D=16 this only kicks in for
+        // K > 64, well above the DLRM range.
+        let kernel = format!(
+            "
+extern \"C\" __global__ void stacked_embedding_bag_sum_kernel(
+    float* __restrict__ out,
+    const float* __restrict__ weight{idx_params}
+) {{
+    const int B = {batch};
+    const int L = {bag};
+    const int D = {d};
+    const int K = {num_tables};
+    const int total = K * D;
+    int b = blockIdx.x;
+    if (b >= B) return;
+    for (int tid = threadIdx.x; tid < total; tid += blockDim.x) {{
+        int k = tid / D;
+        int d = tid - k * D;
+        float acc = 0.0f;
+        switch (k) {{
+{switch}
+            default: continue;
+        }}
+        // Output laid out as (B, K, D) row-major.
+        out[(b * K + k) * D + d] = acc;
+    }}
+}}
+",
+            batch = self.batch,
+            bag = self.bag,
+            d = self.d,
+            num_tables = self.num_tables,
+            idx_params = idx_params,
+            switch = switch,
+        );
+
+        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("stacked_embedding_bag_sum_kernel")
+                .unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        // Block size: enough threads to cover K·D output cells per batch
+        // row, rounded up to a warp (32) for full warp utilization, capped
+        // at 1024 (CUDA max block dim). Lower bound of 32 ensures we never
+        // launch sub-warp blocks when K·D < 32 (e.g. N=1).
+        let total = self.num_tables * self.d;
+        let block_threads = total.next_multiple_of(32).clamp(32, 1024);
+        (
+            func,
+            module,
+            kernel,
+            (
+                Expression::from(self.batch),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            (
+                Expression::from(block_threads),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            Expression::from(0usize),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        Expression::from(self.batch * self.num_tables * self.d)
+    }
+
+    fn output_bytes(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn output_dtype(&self) -> DType {
+        DType::F32
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        // Per output element, L reads from weight. Index reads ~negligible
+        // (D threads share the same L indices per output row).
+        Expression::from(self.batch * self.num_tables * self.d * self.bag * 4)
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn flops(&self) -> Expression {
+        Expression::from(self.batch * self.num_tables * self.d * self.bag)
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "StackedEmbeddingBagSum"
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct StackedEmbeddingBagSumCustom(pub StackedEmbeddingBagKernel);
+
+impl CustomOp for StackedEmbeddingBagSumCustom {
+    fn to_llir_op(&self) -> LLIROp {
+        LLIROp::new::<dyn KernelOp>(Box::new(self.0.clone()) as Box<dyn KernelOp>)
+    }
+}
+
+/// Stacked-table fused EmbeddingBag with sum pooling.
+///
+/// * `stacked_weight`: `(sum_k rows_per_table[k], d)` F32, row-major.
+///   The k-th table's rows occupy indices `[row_offsets[k], row_offsets[k+1])`
+///   where `row_offsets[k] = sum_{j<k} rows_per_table[j]`.
+/// * `indices`: list of `num_tables` tensors, each `(batch, bag)` I32.
+///   Index values for table k are in `[0, rows_per_table[k])` — the
+///   per-table row offset is added inside the kernel.
+/// * `row_offsets`: cumulative starting row index for each table
+///   (length `num_tables + 1`).
+///
+/// Returns `(batch, num_tables, d)` F32. Use `slice_along` + `squeeze`
+/// (or the bundled `dlrm_pairwise_dot_lower_tri_stacked` op) to consume
+/// per-table outputs downstream.
+pub fn stacked_embedding_bag_sum_kernel(
+    stacked_weight: GraphTensor,
+    indices: Vec<GraphTensor>,
+    row_offsets: &[usize],
+) -> GraphTensor {
+    assert_eq!(
+        stacked_weight.dtype,
+        DType::F32,
+        "stacked_embedding_bag_sum_kernel: weight must be F32"
+    );
+    let num_tables = indices.len();
+    assert!(num_tables >= 1, "need at least one index tensor");
+    assert_eq!(
+        row_offsets.len(),
+        num_tables + 1,
+        "row_offsets must have num_tables+1 entries"
+    );
+    let w_dims = stacked_weight.dims();
+    assert_eq!(w_dims.len(), 2, "stacked weight must be 2D (total_rows, d)");
+    let total_rows = w_dims[0].to_usize().expect("total_rows must be static");
+    assert_eq!(
+        total_rows, row_offsets[num_tables],
+        "row_offsets[-1] must equal weight total_rows"
+    );
+    let d = w_dims[1].to_usize().expect("d must be static");
+    let i_dims = indices[0].dims();
+    assert_eq!(i_dims.len(), 2, "indices must be 2D (batch, bag)");
+    let batch = i_dims[0].to_usize().expect("batch must be static");
+    let bag = i_dims[1].to_usize().expect("bag must be static");
+    for idx in &indices {
+        assert_eq!(idx.dtype, DType::Int, "indices must be Int");
+        let id = idx.dims();
+        assert_eq!(id.len(), 2);
+        assert_eq!(id[0].to_usize().unwrap(), batch);
+        assert_eq!(id[1].to_usize().unwrap(), bag);
+    }
+    let kern = StackedEmbeddingBagKernel {
+        batch,
+        bag,
+        d,
+        num_tables,
+        row_offsets: row_offsets.to_vec(),
+    };
+    let cx = unsafe { &mut *stacked_weight.graph_ref };
+    let mut inputs = vec![stacked_weight];
+    inputs.extend(indices);
+    cx.custom_op(
+        StackedEmbeddingBagSumCustom(kern),
+        inputs,
+        (batch, num_tables, d),
+        DType::F32,
+    )
+}
+
+/// Fused EmbeddingBag with sum pooling (single table).
+///
+/// * `table`: `(n_emb, d)` F32, row-major.
+/// * `indices`: `(batch, bag)` I32, row-major. Values must be in `[0, n_emb)`.
+///
+/// Returns: `(batch, d)` F32, row-major. Each output row is the sum of
+/// `bag` looked-up rows from `table`.
+///
+/// All dimensions must be static. The returned tensor's graph node is a
+/// `dyn KernelOp` in LLIR, so it lives inside the same CudaGraphOp as
+/// surrounding kernel ops and benefits from the same CUDA-graph replay.
+pub fn embedding_bag_sum_kernel(table: GraphTensor, indices: GraphTensor) -> GraphTensor {
+    assert_eq!(table.dtype, DType::F32, "embedding_bag_sum_kernel: table must be F32");
+    assert_eq!(
+        indices.dtype,
+        DType::Int,
+        "embedding_bag_sum_kernel: indices must be Int"
+    );
+    let t_dims = table.dims();
+    let i_dims = indices.dims();
+    assert_eq!(t_dims.len(), 2, "table must be 2D (n_emb, d)");
+    assert_eq!(i_dims.len(), 2, "indices must be 2D (batch, bag)");
+    let n_emb = t_dims[0].to_usize().expect("n_emb must be static");
+    let d = t_dims[1].to_usize().expect("d must be static");
+    let batch = i_dims[0].to_usize().expect("batch must be static");
+    let bag = i_dims[1].to_usize().expect("bag must be static");
+
+    let kern = EmbeddingBagSumKernel {
+        batch,
+        bag,
+        d,
+        n_emb,
+    };
+    let cx = unsafe { &mut *table.graph_ref };
+    cx.custom_op(
+        EmbeddingBagSumCustom(kern),
+        vec![table, indices],
+        (batch, d),
+        DType::F32,
+    )
+}
+// ---------------------------------------------------------------------------
+// Multi-table EmbeddingBag (one kernel for K independent (weight, idx) pairs)
+// ---------------------------------------------------------------------------
+
+/// Folds K independent `EmbeddingBag(sum)` lookups into a single CUDA
+/// kernel launch. Used by the vanilla-DLRMv1 translator path where the
+/// model has K separate `nn.EmbeddingBag` modules — each one would
+/// otherwise lower to its own (~5 µs) launch.
+///
+/// Inputs (in `KernelOp`-order):
+///   - `weight_0, weight_1, ..., weight_{K-1}` — each `(n_emb_k, d)` F32.
+///     **The per-table `n_emb` may differ**; only `d` and bag size `L`
+///     must match across tables.
+///   - `idx_0, idx_1, ..., idx_{K-1}` — each `(batch, L)` Int (i32).
+///
+/// Two packed staging buffers carry the K weight + K idx device pointers
+/// into the kernel (`build_params` fills them per execution via
+/// `cuMemcpyHtoD`). The hot loop reads each pointer from shared memory
+/// — no per-table switch needed.
+///
+/// Output shape: `(batch, num_tables, d)` F32, row-major.
+#[derive(Debug, Clone)]
+pub struct MultiTableEmbeddingBagSumKernel {
+    pub batch: usize,
+    pub bag: usize,
+    pub d: usize,
+    pub num_tables: usize,
+}
+
+impl KernelOp for MultiTableEmbeddingBagSumKernel {
+    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>>,
+    ) {
+        // Layout (mirrors worktree's StackedEmbeddingBagSumKernel):
+        //   - One block per batch row (B blocks).
+        //   - Each block produces (K, D) output cells, striding over K·D
+        //     threads (rounded up to a warp).
+        //   - K weight pointers + K idx pointers come in via two packed
+        //     staging buffers populated in `build_params`.
+        //   - Shared memory caches both pointer arrays so the hot loop
+        //     reads at shmem latency.
+        let kernel = format!(
+            "
+extern \"C\" __global__ void multi_table_embedding_bag_sum_kernel(
+    float* __restrict__ out,
+    const long* __restrict__ w_ptrs_packed,
+    const long* __restrict__ idx_ptrs_packed
+) {{
+    const int B = {batch};
+    const int L = {bag};
+    const int D = {d};
+    const int K = {num_tables};
+    const int total = K * D;
+    int b = blockIdx.x;
+    if (b >= B) return;
+
+    __shared__ const float* s_w_ptrs[K];
+    __shared__ const int*   s_idx_ptrs[K];
+    if (threadIdx.x < K) {{
+        s_w_ptrs[threadIdx.x]   = (const float*)(w_ptrs_packed[threadIdx.x]);
+        s_idx_ptrs[threadIdx.x] = (const int*)(idx_ptrs_packed[threadIdx.x]);
+    }}
+    __syncthreads();
+
+    for (int tid = threadIdx.x; tid < total; tid += blockDim.x) {{
+        int k = tid / D;
+        int d = tid - k * D;
+        const float* w = s_w_ptrs[k];
+        const int*   idx = s_idx_ptrs[k];
+        float acc = 0.0f;
+        #pragma unroll 4
+        for (int l = 0; l < L; ++l) {{
+            int row = idx[b * L + l];
+            acc += w[row * D + d];
+        }}
+        // (B, K, D) row-major.
+        out[(b * K + k) * D + d] = acc;
+    }}
+}}
+",
+            batch = self.batch,
+            bag = self.bag,
+            d = self.d,
+            num_tables = self.num_tables,
+        );
+
+        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("multi_table_embedding_bag_sum_kernel")
+                .unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        let total = self.num_tables * self.d;
+        let block_threads = total.next_multiple_of(32).clamp(32, 1024);
+        (
+            func,
+            module,
+            kernel,
+            (
+                Expression::from(self.batch),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            (
+                Expression::from(block_threads),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            Expression::from(0usize),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        Expression::from(self.batch * self.num_tables * self.d)
+    }
+
+    fn output_bytes(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn output_dtype(&self) -> DType {
+        DType::F32
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        Expression::from(self.batch * self.num_tables * self.d * self.bag * 4
+            + self.batch * self.num_tables * self.bag * 4)
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn flops(&self) -> Expression {
+        Expression::from(self.batch * self.num_tables * self.d * self.bag)
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "MultiTableEmbeddingBagSum"
+    }
+
+    /// Two staging buffers: one for K weight ptrs, one for K idx ptrs.
+    /// Each is `K * 8` bytes (an array of u64s, written as `long*` on
+    /// the device side).
+    fn allocate_internal_buffers(
+        &self,
+        stream: &Arc<CudaStream>,
+        _dyn_map: &FxHashMap<char, usize>,
+    ) -> Vec<CudaSlice<u8>> {
+        let buf_size = self.num_tables * 8;
+        vec![
+            stream
+                .alloc_zeros::<u8>(buf_size)
+                .expect("alloc MultiTableEmbBag w-ptr staging buffer"),
+            stream
+                .alloc_zeros::<u8>(buf_size)
+                .expect("alloc MultiTableEmbBag idx-ptr staging buffer"),
+        ]
+    }
+
+    /// Pack the K weight + K idx pointers into the two staging buffers
+    /// each execution, then emit `[out, w_buf, idx_buf]` as kernel params.
+    ///
+    /// `input_ptrs` layout: `[w_0, w_1, ..., w_{K-1}, idx_0, ..., idx_{K-1}]`.
+    /// `cuMemcpyHtoD_v2` is a blocking host call so by the time we return
+    /// the staging buffers are populated and the subsequent CUDA-graph
+    /// node-param update reads stable device pointers.
+    fn build_params(
+        &self,
+        stream: &Arc<CudaStream>,
+        output_ptr: u64,
+        input_ptrs: &[u64],
+        internal_bufs: &[CudaSlice<u8>],
+        _dyn_dims_ptr: u64,
+    ) -> Vec<u64> {
+        assert_eq!(
+            input_ptrs.len(),
+            2 * self.num_tables,
+            "MultiTableEmbeddingBagSum: expected {} input pointers (K weights + K idx), got {}",
+            2 * self.num_tables,
+            input_ptrs.len(),
+        );
+        let (w_ptrs, idx_ptrs) = input_ptrs.split_at(self.num_tables);
+        let w_buf = &internal_bufs[0];
+        let idx_buf = &internal_bufs[1];
+        let w_dev_ptr: u64 = w_buf.device_ptr(stream).0;
+        let idx_dev_ptr: u64 = idx_buf.device_ptr(stream).0;
+        unsafe {
+            let r1 = cudarc::driver::sys::cuMemcpyHtoD_v2(
+                w_dev_ptr,
+                w_ptrs.as_ptr() as *const std::ffi::c_void,
+                w_ptrs.len() * 8,
+            );
+            assert_eq!(
+                r1,
+                cudarc::driver::sys::CUresult::CUDA_SUCCESS,
+                "cuMemcpyHtoD_v2 for MultiTableEmbBag w-ptr staging failed: {r1:?}",
+            );
+            let r2 = cudarc::driver::sys::cuMemcpyHtoD_v2(
+                idx_dev_ptr,
+                idx_ptrs.as_ptr() as *const std::ffi::c_void,
+                idx_ptrs.len() * 8,
+            );
+            assert_eq!(
+                r2,
+                cudarc::driver::sys::CUresult::CUDA_SUCCESS,
+                "cuMemcpyHtoD_v2 for MultiTableEmbBag idx-ptr staging failed: {r2:?}",
+            );
+        }
+        vec![output_ptr, w_dev_ptr, idx_dev_ptr]
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct MultiTableEmbeddingBagSumCustom(pub MultiTableEmbeddingBagSumKernel);
+
+impl CustomOp for MultiTableEmbeddingBagSumCustom {
+    fn to_llir_op(&self) -> LLIROp {
+        LLIROp::new::<dyn KernelOp>(Box::new(self.0.clone()) as Box<dyn KernelOp>)
+    }
+}
+
+/// Frontend helper: K independent EmbeddingBag(sum) lookups in one
+/// kernel launch. Returns `(batch, num_tables, d)` F32, row-major;
+/// slice along axis 1 (`out.slice_along(k..k+1, 1).squeeze(1)`) to
+/// recover the k-th table's `(batch, d)` output.
+///
+/// * `weights`: K `(n_emb_k, d)` F32 tensors. Per-table `n_emb` may
+///   differ; only `d` must be shared.
+/// * `indices`: K `(batch, bag)` Int tensors (cast `.cast(DType::Int)`
+///   on the caller side if your indices are i64).
+pub fn multi_table_embedding_bag_sum_kernel(
+    weights: Vec<GraphTensor>,
+    indices: Vec<GraphTensor>,
+) -> GraphTensor {
+    assert_eq!(
+        weights.len(),
+        indices.len(),
+        "multi_table_embedding_bag_sum_kernel: need one weight per index tensor"
+    );
+    let num_tables = weights.len();
+    assert!(num_tables >= 1, "need at least one table");
+    let first_w = weights[0];
+    let first_idx = indices[0];
+    let w_dims = first_w.dims();
+    let i_dims = first_idx.dims();
+    assert_eq!(w_dims.len(), 2, "weights must be 2D (n_emb, d)");
+    assert_eq!(i_dims.len(), 2, "indices must be 2D (batch, bag)");
+    let d = w_dims[1].to_usize().expect("d must be static");
+    let batch = i_dims[0].to_usize().expect("batch must be static");
+    let bag = i_dims[1].to_usize().expect("bag must be static");
+    for w in &weights {
+        assert_eq!(w.dtype, DType::F32, "weights must all be F32");
+        let wd = w.dims();
+        assert_eq!(wd.len(), 2, "weight must be 2D");
+        assert_eq!(
+            wd[1].to_usize().unwrap(),
+            d,
+            "all weights must share inner dim"
+        );
+    }
+    for idx in &indices {
+        assert_eq!(idx.dtype, DType::Int, "indices must all be Int (i32)");
+        let id = idx.dims();
+        assert_eq!(id.len(), 2);
+        assert_eq!(id[0].to_usize().unwrap(), batch);
+        assert_eq!(id[1].to_usize().unwrap(), bag);
+    }
+    let kern = MultiTableEmbeddingBagSumKernel {
+        batch,
+        bag,
+        d,
+        num_tables,
+    };
+    let mut inputs = weights;
+    inputs.extend(indices);
+    let cx = unsafe { &mut *first_w.graph_ref };
+    cx.custom_op(
+        MultiTableEmbeddingBagSumCustom(kern),
+        inputs,
+        (batch, num_tables, d),
+        DType::F32,
+    )
+}

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

@@ -0,0 +1,378 @@
+// =========================================================================
+// Generic CUDA elementwise ops used inside FusionStart/FusionEnd regions.
+//
+// CUDA elementwise execution is represented as a FusionEnd-rooted region even
+// for a single op. These ops are therefore region-internal only; standalone
+// compilation is intentionally unsupported.
+// =========================================================================
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::{
+    egglog_utils::{
+        api::{Rule, SortDef, sort},
+        base::{DTYPE, ELIST, OP_KIND, STRING},
+        extract_dtype, extract_expr_list,
+    },
+    op::*,
+    prelude::*,
+};
+
+use crate::kernel::KernelOp;
+
+pub type Ops = (CudaUnaryElementwise, CudaBinaryElementwise);
+
+type CompileOut = (
+    CudaFunction,
+    Arc<CudaModule>,
+    String,
+    (Expression, Expression, Expression),
+    (Expression, Expression, Expression),
+    Expression,
+    FxHashMap<char, CudaSlice<u8>>,
+);
+
+fn extract_string_label(egraph: &SerializedEGraph, node: &ENodeId) -> String {
+    egraph.enodes[node].0.trim_matches('"').to_string()
+}
+
+#[derive(Default, Debug, Clone)]
+pub struct CudaUnaryElementwise {
+    pub(crate) op: String,
+    pub(crate) shape: Vec<Expression>,
+    pub(crate) in_strides: Vec<Expression>,
+    pub(crate) out_strides: Vec<Expression>,
+    pub(crate) dtype: DType,
+}
+
+impl EgglogOp for CudaUnaryElementwise {
+    fn sort(&self) -> SortDef {
+        sort(
+            OP_KIND,
+            "CudaUnaryElementwise",
+            &[
+                ("op", STRING),
+                ("shape", ELIST),
+                ("strides", ELIST),
+                ("out_strides", ELIST),
+                ("dtype", DTYPE),
+            ],
+        )
+    }
+
+    fn n_inputs(&self) -> usize {
+        1
+    }
+
+    fn rewrites(&self) -> Vec<Rule> {
+        let mut rules = Vec::new();
+        for (hlir, opcode) in [
+            ("Sin", "Sin"),
+            ("Sqrt", "Sqrt"),
+            ("Exp2", "Exp2"),
+            ("Log2", "Log2"),
+            ("Recip", "Recip"),
+        ] {
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?u (Op ({hlir} ?shape ?s ?out_s) (ICons ?x (INil))))
+                    (= ?dt (dtype ?u))
+                 ) (
+                    (let ?fs (Op (FusionStart ?shape ?s ?dt) (ICons ?x (INil))))
+                    (let ?elem (Op (CudaUnaryElementwise \"{opcode}\" ?shape ?s ?out_s ?dt)
+                                   (ICons ?fs (INil))))
+                    (let ?fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?elem (INil))))
+                    (union ?u ?fe)
+                    (set (dtype ?fe) ?dt)
+                 ) :ruleset kernel_lower :name \"cuda-elem-singleton-{hlir}\")"
+            )));
+        }
+
+        rules.push(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 ?fs (Op (FusionStart ?shape ?x_stride ?dt) (ICons ?x (INil))))
+                    (let ?elem (Op (CudaUnaryElementwise \"Exp\" ?shape ?x_stride ?out_stride ?dt)
+                                   (ICons ?fs (INil))))
+                    (let ?fe (Op (FusionEnd ?shape ?out_stride ?dt) (ICons ?elem (INil))))
+                    (union ?exp2 ?fe)
+                    (set (dtype ?fe) ?dt)
+                )
+                :ruleset direct_kernel
+                :name \"direct-exp-region\"
+            )",
+        ));
+
+        rules.push(Rule::raw(
+            "(datatype*
+                (CudaSigmoidScaledState
+                    (MkCudaSigmoidScaledState IR EList EList DType)
+                )
+            )
+            (function cuda_sigmoid_scaled (IR) CudaSigmoidScaledState :merge new)
+
+            (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)))))
+                (= ?dt (dtype ?x))
+            )
+            (
+                (set (cuda_sigmoid_scaled ?scaled)
+                    (MkCudaSigmoidScaledState ?x ?shape ?x_stride ?dt))
+            )
+            :ruleset direct_kernel
+            :name \"direct-sigmoid-scaled-region-marker\"
+            )
+
+            (rule
+            (
+                (= ?scaled_state (cuda_sigmoid_scaled ?scaled))
+                (= ?scaled_state (MkCudaSigmoidScaledState ?x ?shape ?x_stride ?dt))
+                (= ?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))))
+            )
+            (
+                (let ?fs (Op (FusionStart ?shape ?x_stride ?dt) (ICons ?x (INil))))
+                (let ?elem (Op (CudaUnaryElementwise \"Sigmoid\" ?shape ?x_stride ?out_stride ?dt)
+                               (ICons ?fs (INil))))
+                (let ?fe (Op (FusionEnd ?shape ?out_stride ?dt) (ICons ?elem (INil))))
+                (union ?sig_out ?fe)
+                (set (dtype ?fe) ?dt)
+            )
+            :ruleset direct_kernel
+            :name \"direct-sigmoid-region\"
+            )",
+        ));
+
+        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 {
+                op: extract_string_label(egraph, kind_children[0]),
+                shape: extract_expr_list(egraph, kind_children[1], list_cache, expr_cache).unwrap(),
+                in_strides: extract_expr_list(egraph, kind_children[2], list_cache, expr_cache)
+                    .unwrap(),
+                out_strides: extract_expr_list(egraph, kind_children[3], list_cache, expr_cache)
+                    .unwrap(),
+                dtype: extract_dtype(egraph, kind_children[4]),
+            })),
+            input_enodes,
+        )
+    }
+}
+
+impl KernelOp for CudaUnaryElementwise {
+    fn compile(
+        &self,
+        _stream: &Arc<CudaStream>,
+        _compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    ) -> CompileOut {
+        unreachable!("CudaUnaryElementwise must be compiled through fusion region codegen")
+    }
+
+    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.output_size()
+    }
+
+    fn output_dtype(&self) -> DType {
+        self.dtype
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "CudaUnaryElementwise"
+    }
+}
+
+#[derive(Default, Debug, Clone)]
+pub struct CudaBinaryElementwise {
+    pub(crate) op: String,
+    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 CudaBinaryElementwise {
+    fn sort(&self) -> SortDef {
+        sort(
+            OP_KIND,
+            "CudaBinaryElementwise",
+            &[
+                ("op", STRING),
+                ("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![
+            Rule::raw(
+                "(rule (
+                    (= ?bin (Op (Add ?shape ?a_s ?b_s ?out_s) (ICons ?a (ICons ?b (INil)))))
+                    (= ?dt (dtype ?bin))
+                 ) (
+                    (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 ?elem (Op (CudaBinaryElementwise \"Add\" ?shape ?a_s ?b_s ?out_s ?dt)
+                                   (ICons ?fs_a (ICons ?fs_b (INil)))))
+                    (let ?fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?elem (INil))))
+                    (union ?bin ?fe)
+                    (set (dtype ?fe) ?dt)
+                 ) :ruleset kernel_lower :name \"cuda-elem-singleton-Add\")",
+            ),
+            Rule::raw(
+                "(rule (
+                    (= ?bin (Op (Mul ?shape ?a_s ?b_s ?out_s) (ICons ?a (ICons ?b (INil)))))
+                    (= ?dt (dtype ?a))
+                 ) (
+                    (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 ?elem (Op (CudaBinaryElementwise \"Mul\" ?shape ?a_s ?b_s ?out_s ?dt)
+                                   (ICons ?fs_a (ICons ?fs_b (INil)))))
+                    (let ?fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?elem (INil))))
+                    (union ?bin ?fe)
+                    (set (dtype ?fe) ?dt)
+                 ) :ruleset kernel_lower :name \"cuda-elem-singleton-Mul\")",
+            ),
+        ]
+    }
+
+    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>) {
+        let mut out_shape =
+            extract_expr_list(egraph, kind_children[1], list_cache, expr_cache).unwrap();
+        let mut a_stride =
+            extract_expr_list(egraph, kind_children[2], list_cache, expr_cache).unwrap();
+        let mut b_stride =
+            extract_expr_list(egraph, kind_children[3], list_cache, expr_cache).unwrap();
+        let mut out_stride =
+            extract_expr_list(egraph, kind_children[4], list_cache, expr_cache).unwrap();
+        let n = out_shape
+            .len()
+            .min(a_stride.len())
+            .min(b_stride.len())
+            .min(out_stride.len());
+        out_shape.truncate(n);
+        a_stride.truncate(n);
+        b_stride.truncate(n);
+        out_stride.truncate(n);
+        (
+            LLIROp::new::<dyn KernelOp>(Box::new(Self {
+                op: extract_string_label(egraph, kind_children[0]),
+                out_shape,
+                a_stride,
+                b_stride,
+                out_stride,
+                dtype: extract_dtype(egraph, kind_children[5]),
+            })),
+            input_enodes,
+        )
+    }
+}
+
+impl KernelOp for CudaBinaryElementwise {
+    fn compile(
+        &self,
+        _stream: &Arc<CudaStream>,
+        _compile_cache: &mut FxHashMap<String, (Arc<CudaModule>, CudaFunction)>,
+    ) -> CompileOut {
+        unreachable!("CudaBinaryElementwise must be compiled through fusion region codegen")
+    }
+
+    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 {
+        self.output_bytes() * 2
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_bytes()
+    }
+
+    fn flops(&self) -> Expression {
+        self.output_size()
+    }
+
+    fn output_dtype(&self) -> DType {
+        self.dtype
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "CudaBinaryElementwise"
+    }
+}

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

@@ -1,301 +0,0 @@
-// =========================================================================
-// 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.
-//
-// Each FusedX must be absorbed into a FusionEnd-rooted region and compiled by
-// `region_codegen`; standalone compilation is intentionally unsupported.
-// =========================================================================
-
-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::kernel::KernelOp;
-
-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>>,
-);
-
-/// 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 {
-                unreachable!(concat!(
-                    $sort,
-                    " must be compiled through fusion region codegen"
-                ))
-            }
-            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 {
-                unreachable!(concat!(
-                    $sort,
-                    " must be compiled through fusion region codegen"
-                ))
-            }
-            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

@@ -9,8 +9,8 @@
 //
 // `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
+// codegen lives in `region_codegen`. Both markers' `compile()` is
+// `unreachable!()` — region codegen folds them away
 // before kernel_to_host's compile loop reaches an interior node.
 // =========================================================================
 
@@ -142,218 +142,164 @@ impl EgglogOp for FusionEnd {
     }
 
     fn rewrites(&self) -> Vec<Rule> {
-        // 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.
+        // Generic region growth works directly from HLIR elementwise ops into
+        // `Cuda*Elementwise` region nodes. The concrete HLIR op still appears in
+        // the egraph, so fusion remains a normal nondestructive alternative, but
+        // the region-internal representation is arity based instead of one
+        // dedicated fused sort per operation.
         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"),
+            ("Sin", "Sin"),
+            ("Sqrt", "Sqrt"),
+            ("Exp2", "Exp2"),
+            ("Log2", "Log2"),
+            ("Recip", "Recip"),
         ];
+        let binaries: &[(&str, &str)] = &[("Add", "Add"), ("Mul", "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)
-                     ) :ruleset fusion_pair :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)
-                     ) :ruleset fusion_pair :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)
-                     ) :ruleset fusion_pair :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)
-                     ) :ruleset fusion_pair :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)
-                     ) :ruleset fusion_pair :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)
-                     ) :ruleset fusion_pair :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 {
+        // Grow FE → unary consumer: U(FE(inner)) → FE(CudaUnary(inner)).
+        for (hlir, opcode) 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))))
+                    (= ?u (Op ({hlir} ?shape ?s ?s) (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))))
+                    (let ?elem (Op (CudaUnaryElementwise \"{opcode}\" ?shape ?s ?s ?dt)
+                                   (ICons ?inner (INil))))
+                    (let ?new_fe (Op (FusionEnd ?shape ?s ?dt) (ICons ?elem (INil))))
                     (union ?u ?new_fe)
-                 ) :ruleset fusion_grow :name \"grow-FE-U-{ku}\")"
+                    (set (dtype ?new_fe) ?dt)
+                 ) :ruleset fusion_grow :name \"grow-FE-U-{hlir}\")"
             )));
         }
 
-        // 6. Grow FE → B (lhs / rhs): one input is the FE, the other external.
-        for (kb, fb, lb) in binaries {
+        // Grow FE → binary consumer, left and right orientations.
+        for (hlir, opcode) 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)
+                    (= ?bin (Op ({hlir} ?shape ?a_s ?b_s ?out_s)
                                  (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)
+                    (let ?elem (Op (CudaBinaryElementwise \"{opcode}\" ?shape ?a_s ?b_s ?out_s ?dt)
                                    (ICons ?inner_a (ICons ?fs_b (INil)))))
-                    (let ?new_fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fbin (INil))))
+                    (let ?new_fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?elem (INil))))
                     (union ?bin ?new_fe)
-                 ) :ruleset fusion_grow :name \"grow-FE-B-lhs-{lb}\")"
+                    (set (dtype ?new_fe) ?dt)
+                 ) :ruleset fusion_grow :name \"grow-FE-B-lhs-{hlir}\")"
             )));
             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)
+                    (= ?bin (Op ({hlir} ?shape ?a_s ?b_s ?out_s)
                                  (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)
+                    (let ?elem (Op (CudaBinaryElementwise \"{opcode}\" ?shape ?a_s ?b_s ?out_s ?dt)
                                    (ICons ?fs_a (ICons ?inner_b (INil)))))
-                    (let ?new_fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fbin (INil))))
+                    (let ?new_fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?elem (INil))))
                     (union ?bin ?new_fe)
-                 ) :ruleset fusion_grow :name \"grow-FE-B-rhs-{lb}\")"
+                    (set (dtype ?new_fe) ?dt)
+                 ) :ruleset fusion_grow :name \"grow-FE-B-rhs-{hlir}\")"
             )));
         }
 
-        // 7. Merge two FEs at a binary: B(FE(ia), FE(ib)) → FE(FB(ia, ib)).
-        //
-        // This is destructive: after creating the larger region, subsume the
-        // two smaller FusionEnd rows. Without that, independently-grown left
-        // and right regions form a Cartesian product, then those alternatives
-        // can merge again higher in the graph.
-        for (kb, fb, lb) in binaries {
+        // Absorb an elementwise producer through a FusionStart boundary. This
+        // makes a region that initially treats `producer(...)` as an external
+        // input able to pull that producer inside later.
+        for (hlir, opcode) in unaries {
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?u (Op ({hlir} ?shape ?s ?s) (ICons ?x (INil))))
+                    (= ?fs_u (Op (FusionStart ?shape ?s ?dt) (ICons ?u (INil))))
+                 ) (
+                    (let ?fs_x (Op (FusionStart ?shape ?s ?dt) (ICons ?x (INil))))
+                    (let ?elem (Op (CudaUnaryElementwise \"{opcode}\" ?shape ?s ?s ?dt)
+                                   (ICons ?fs_x (INil))))
+                    (union ?fs_u ?elem)
+                 ) :ruleset fusion_grow :name \"grow-U-FS-{hlir}\")"
+            )));
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?inner_fe (Op (FusionEnd ?shape ?s ?dt) (ICons ?inner (INil))))
+                    (= ?bad_fs (Op (FusionStart ?shape ?s ?dt) (ICons ?inner_fe (INil))))
+                    (= ?bad_elem (Op (CudaUnaryElementwise \"{opcode}\" ?shape ?s ?s ?dt)
+                                     (ICons ?bad_fs (INil))))
+                    (= ?bad_fe (Op (FusionEnd ?shape ?s ?dt) (ICons ?bad_elem (INil))))
+                    (= ?good_elem (Op (CudaUnaryElementwise \"{opcode}\" ?shape ?s ?s ?dt)
+                                      (ICons ?inner (INil))))
+                    (= ?good_fe (Op (FusionEnd ?shape ?s ?dt) (ICons ?good_elem (INil))))
+                    (= ?bad_fe ?good_fe)
+                 ) (
+                    (delete (Op (FusionStart ?shape ?s ?dt) (ICons ?inner_fe (INil))))
+                 ) :ruleset cleanup :name \"cleanup-nested-FS-FE-unary-{hlir}\")"
+            )));
+        }
+        for (hlir, opcode) in binaries {
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?bin (Op ({hlir} ?shape ?a_s ?b_s ?out_s)
+                                 (ICons ?a (ICons ?b (INil)))))
+                    (= ?fs_bin (Op (FusionStart ?shape ?out_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 ?elem (Op (CudaBinaryElementwise \"{opcode}\" ?shape ?a_s ?b_s ?out_s ?dt)
+                                   (ICons ?fs_a (ICons ?fs_b (INil)))))
+                    (union ?fs_bin ?elem)
+                 ) :ruleset fusion_grow :name \"grow-B-FS-{hlir}\")"
+            )));
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?inner_fe (Op (FusionEnd ?shape ?a_s ?dt) (ICons ?inner_a (INil))))
+                    (= ?bad_fs (Op (FusionStart ?shape ?a_s ?dt) (ICons ?inner_fe (INil))))
+                    (= ?fs_b (Op (FusionStart ?shape ?b_s ?dt) (ICons ?b (INil))))
+                    (= ?bad_elem (Op (CudaBinaryElementwise \"{opcode}\" ?shape ?a_s ?b_s ?out_s ?dt)
+                                     (ICons ?bad_fs (ICons ?fs_b (INil)))))
+                    (= ?bad_fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?bad_elem (INil))))
+                    (= ?good_elem (Op (CudaBinaryElementwise \"{opcode}\" ?shape ?a_s ?b_s ?out_s ?dt)
+                                      (ICons ?inner_a (ICons ?fs_b (INil)))))
+                    (= ?good_fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?good_elem (INil))))
+                    (= ?bad_fe ?good_fe)
+                 ) (
+                    (delete (Op (FusionStart ?shape ?a_s ?dt) (ICons ?inner_fe (INil))))
+                 ) :ruleset cleanup :name \"cleanup-nested-FS-FE-binary-lhs-{hlir}\")"
+            )));
+            rules.push(Rule::raw(format!(
+                "(rule (
+                    (= ?inner_fe (Op (FusionEnd ?shape ?b_s ?dt) (ICons ?inner_b (INil))))
+                    (= ?bad_fs (Op (FusionStart ?shape ?b_s ?dt) (ICons ?inner_fe (INil))))
+                    (= ?fs_a (Op (FusionStart ?shape ?a_s ?dt) (ICons ?a (INil))))
+                    (= ?bad_elem (Op (CudaBinaryElementwise \"{opcode}\" ?shape ?a_s ?b_s ?out_s ?dt)
+                                     (ICons ?fs_a (ICons ?bad_fs (INil)))))
+                    (= ?bad_fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?bad_elem (INil))))
+                    (= ?good_elem (Op (CudaBinaryElementwise \"{opcode}\" ?shape ?a_s ?b_s ?out_s ?dt)
+                                      (ICons ?fs_a (ICons ?inner_b (INil)))))
+                    (= ?good_fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?good_elem (INil))))
+                    (= ?bad_fe ?good_fe)
+                 ) (
+                    (delete (Op (FusionStart ?shape ?b_s ?dt) (ICons ?inner_fe (INil))))
+                 ) :ruleset cleanup :name \"cleanup-nested-FS-FE-binary-rhs-{hlir}\")"
+            )));
+        }
+
+        // Merge two FEs at a binary: B(FE(ia), FE(ib)) → FE(CudaBinary(ia, ib)).
+        for (hlir, opcode) 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)
+                    (= ?bin (Op ({hlir} ?shape ?a_s ?b_s ?out_s)
                                  (ICons ?fe_a (ICons ?fe_b (INil)))))
                  ) (
-                    (let ?fbin (Op ({fb} ?shape ?a_s ?b_s ?o_s ?dt)
+                    (let ?elem (Op (CudaBinaryElementwise \"{opcode}\" ?shape ?a_s ?b_s ?out_s ?dt)
                                    (ICons ?inner_a (ICons ?inner_b (INil)))))
-                    (let ?new_fe (Op (FusionEnd ?shape ?o_s ?dt) (ICons ?fbin (INil))))
+                    (let ?new_fe (Op (FusionEnd ?shape ?out_s ?dt) (ICons ?elem (INil))))
                     (union ?bin ?new_fe)
-                    (subsume (Op (FusionEnd ?shape ?a_s ?dt) (ICons ?inner_a (INil))))
-                    (subsume (Op (FusionEnd ?shape ?b_s ?dt) (ICons ?inner_b (INil))))
-                 ) :ruleset fusion_merge :name \"merge-FE-FE-{lb}\")"
+                    (set (dtype ?new_fe) ?dt)
+                 ) :ruleset fusion_merge :name \"merge-FE-FE-{hlir}\")"
             )));
         }
 
@@ -363,6 +309,61 @@ impl EgglogOp for FusionEnd {
         // `Cycle(NodeIndex(_))`. Grow rules already compose adjacent regions
         // correctly without dissolve.
 
+        rules.push(Rule::raw(
+            "(rule (
+                (= ?fe (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+                (= ?inner (Op (CudaUnaryElementwise ?op ?inner_shape ?inner_in_s ?inner_s ?dt) ?inner_inputs))
+                (!= ?fe_shape ?inner_shape)
+             ) (
+                (delete (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+             ) :ruleset cleanup :name \"delete-malformed-FE-unary-shape\")",
+        ));
+        rules.push(Rule::raw(
+            "(rule (
+                (= ?fe (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+                (= ?inner (Op (CudaUnaryElementwise ?op ?inner_shape ?inner_in_s ?inner_s ?dt) ?inner_inputs))
+                (!= ?fe_s ?inner_s)
+             ) (
+                (delete (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+             ) :ruleset cleanup :name \"delete-malformed-FE-unary-strides\")",
+        ));
+        rules.push(Rule::raw(
+            "(rule (
+                (= ?fe (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+                (= ?inner (Op (CudaBinaryElementwise ?op ?inner_shape ?a_s ?b_s ?inner_s ?dt) ?inner_inputs))
+                (!= ?fe_shape ?inner_shape)
+             ) (
+                (delete (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+             ) :ruleset cleanup :name \"delete-malformed-FE-binary-shape\")",
+        ));
+        rules.push(Rule::raw(
+            "(rule (
+                (= ?fe (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+                (= ?inner (Op (CudaBinaryElementwise ?op ?inner_shape ?a_s ?b_s ?inner_s ?dt) ?inner_inputs))
+                (!= ?fe_s ?inner_s)
+             ) (
+                (delete (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+             ) :ruleset cleanup :name \"delete-malformed-FE-binary-strides\")",
+        ));
+        rules.push(Rule::raw(
+            "(rule (
+                (= ?fe (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+                (= ?inner (Op (FusionEnd ?inner_shape ?inner_s ?dt) ?inner_inputs))
+                (!= ?fe_shape ?inner_shape)
+             ) (
+                (delete (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+             ) :ruleset cleanup :name \"delete-malformed-FE-nested-shape\")",
+        ));
+        rules.push(Rule::raw(
+            "(rule (
+                (= ?fe (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+                (= ?inner (Op (FusionEnd ?inner_shape ?inner_s ?dt) ?inner_inputs))
+                (!= ?fe_s ?inner_s)
+             ) (
+                (delete (Op (FusionEnd ?fe_shape ?fe_s ?dt) (ICons ?inner (INil))))
+             ) :ruleset cleanup :name \"delete-malformed-FE-nested-strides\")",
+        ));
+
         rules
     }
 

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

@@ -2,25 +2,21 @@
 //!
 //! - `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.
+//! - `elementwise` — generic region-internal CUDA elementwise op variants.
 //! - `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
+//! The LLIR keeps `FusionStart` / generic elementwise / `FusionEnd` nodes after
 //! extraction; `region_codegen` is the only place that walks them.
 
-pub mod fused_ops;
+pub mod elementwise;
 pub mod markers;
 pub mod region_codegen;
 
-pub use fused_ops::{
-    FusedAdd, FusedExp, FusedExp2, FusedLog2, FusedMul, FusedRecip, FusedSin, FusedSqrt,
-};
+pub use elementwise::{CudaBinaryElementwise, CudaUnaryElementwise};
 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);
+/// (markers + interior generic elementwise variants). Combined into a flat
+/// tuple for the `Ops` registry in `kernel::mod`.
+pub type Ops = (markers::Ops, elementwise::Ops);

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

@@ -1,26 +1,26 @@
 // =========================================================================
 // Region codegen for FusionStart / FusionEnd-bracketed fused regions.
 //
-// PR1 left FusedX / FusionStart / FusionEnd nodes in the post-extraction
+// Older fusion lowering left elementwise / 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 +
+//     - CompileUnit::Single(node)  — unfused non-region kernels, compiled as before.
+//     - CompileUnit::Region(rgn)   — one FE + its interior elementwise 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
+//                                    chains all elementwise 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
+// device pointers at execute time. Interior Cuda*Elementwise / FusionStart nodes
 // never enter the kernels Vec — they have no buffers, no launches.
 // =========================================================================
 
@@ -40,6 +40,7 @@ use as_any::Downcast;
 use crate::{
     compile_module_image_for_current_device, cuda_dtype,
     kernel::KernelOp,
+    kernel::fusion::elementwise::{CudaBinaryElementwise, CudaUnaryElementwise},
     kernel::fusion::markers::{FusionEnd, FusionStart},
     kernel::hlir::{dtype_includes, generate_dyn_dims_defines},
 };
@@ -52,10 +53,10 @@ use crate::{
 pub(crate) struct RegionUnit {
     /// The FusionEnd node that anchors this region.
     pub fe_node: NodeIndex,
-    /// Interior FusedX nodes, in topological order (predecessors before
+    /// Interior Cuda*Elementwise 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>,
+    pub elementwise_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
@@ -79,13 +80,13 @@ pub(crate) enum CompileUnit {
 
 /// 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
+/// Cuda*Elementwise 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
+/// incoming edges through `FusionEnd` / Cuda*Elementwise nodes, stopping at
 /// `FusionStart` leaves.
 ///
 /// This is computed once over the full LLIR rather than per-convex-
@@ -123,7 +124,7 @@ pub(crate) fn globally_absorbed_markers(llir_graph: &LLIRGraph) -> FxHashSet<Nod
                         absorbed.insert(pred);
                         stack.push(pred);
                     }
-                    Some(other) if other.starts_with("Fused") => {
+                    Some(_) if is_region_elementwise(llir_graph, pred) => {
                         absorbed.insert(pred);
                         stack.push(pred);
                     }
@@ -187,12 +188,12 @@ pub(crate) fn build_compile_units(
                         absorbed.insert(pred);
                         stack.push(pred);
                     }
-                    Some(other) if other.starts_with("Fused") => {
+                    Some(_) if is_region_elementwise(llir_graph, pred) => {
                         interior.push(pred);
                         stack.push(pred);
                     }
                     _ => {
-                        // Non-marker, non-FusedX predecessor inside what
+                        // Non-marker, non-elementwise predecessor inside what
                         // we thought was a region. Shouldn't happen with
                         // the current rules; treat conservatively: do
                         // not absorb it. This means the region is
@@ -229,7 +230,56 @@ pub(crate) fn build_compile_units(
                 llir_graph
                     .neighbors_directed(fs, Direction::Incoming)
                     .next()
-                    .expect("FusionStart with no predecessor")
+                    .unwrap_or_else(|| {
+                        // Dump the malformed structure: which FE
+                        // triggered the walk, every node in fs_topo and
+                        // interior_topo, and each FS's incoming /
+                        // outgoing degree. Helps localize whether the
+                        // missing edge came from extraction or a
+                        // downstream LLIR transform.
+                        if std::env::var("LUMINAL_DEBUG_FUSION_PANIC").is_ok() {
+                            eprintln!(
+                                "FusionStart panic: fe={} (kernel={:?})",
+                                node.index(),
+                                llir_graph.node_weight(node).and_then(|op| {
+                                    op.to_dialect::<dyn KernelOp>().map(|k| k.kernel_name())
+                                }),
+                            );
+                            eprintln!("  fs_topo ({}):", fs_topo.len());
+                            for &f in &fs_topo {
+                                let in_deg = llir_graph
+                                    .neighbors_directed(f, Direction::Incoming)
+                                    .count();
+                                let out_deg = llir_graph
+                                    .neighbors_directed(f, Direction::Outgoing)
+                                    .count();
+                                let kn = llir_graph
+                                    .node_weight(f)
+                                    .and_then(|op| {
+                                        op.to_dialect::<dyn KernelOp>().map(|k| k.kernel_name())
+                                    })
+                                    .unwrap_or("?");
+                                eprintln!(
+                                    "    fs={} kind={} in_deg={} out_deg={}",
+                                    f.index(),
+                                    kn,
+                                    in_deg,
+                                    out_deg,
+                                );
+                            }
+                            eprintln!("  interior_topo ({}):", interior_topo.len());
+                            for &i in &interior_topo {
+                                let kn = llir_graph
+                                    .node_weight(i)
+                                    .and_then(|op| {
+                                        op.to_dialect::<dyn KernelOp>().map(|k| k.kernel_name())
+                                    })
+                                    .unwrap_or("?");
+                                eprintln!("    interior={} kind={}", i.index(), kn);
+                            }
+                        }
+                        panic!("FusionStart with no predecessor")
+                    })
             })
             .collect();
 
@@ -240,7 +290,7 @@ pub(crate) fn build_compile_units(
             node,
             RegionUnit {
                 fe_node: node,
-                fusedx_topo: interior_topo,
+                elementwise_topo: interior_topo,
                 fs_nodes: fs_topo,
                 external_inputs,
             },
@@ -269,24 +319,53 @@ pub(crate) fn build_compile_units(
 }
 
 // =========================================================================
-// Per-FusedX body templates.
+// Per-elementwise 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}"),
+fn is_region_elementwise(llir_graph: &LLIRGraph, node: NodeIndex) -> bool {
+    llir_graph
+        .node_weight(node)
+        .and_then(|op| op.to_dialect::<dyn KernelOp>())
+        .is_some_and(|op| {
+            (***op).downcast_ref::<CudaUnaryElementwise>().is_some()
+                || (***op).downcast_ref::<CudaBinaryElementwise>().is_some()
+        })
+}
+
+fn elementwise_value(local: &str, dtype: DType) -> String {
+    if matches!(dtype, DType::F8E4M3 | DType::F8E5M2 | DType::F8UE8M0) {
+        format!("static_cast<float>({local})")
+    } else {
+        local.to_string()
+    }
+}
+
+fn elementwise_init_expr(expr: &str, dtype: DType, cuda_ty: &str) -> String {
+    if matches!(dtype, DType::F8E4M3 | DType::F8E5M2 | DType::F8UE8M0) {
+        format!("{cuda_ty}({expr})")
+    } else {
+        expr.to_string()
+    }
+}
+
+fn elementwise_body(op: &str, locals: &[&str], dtype: DType) -> String {
+    let a = || elementwise_value(locals[0], dtype);
+    let b = || elementwise_value(locals[1], dtype);
+    match op {
+        "Sin" => format!("sinf({})", a()),
+        "Sqrt" => format!("sqrtf({})", a()),
+        "Exp" => format!("expf({})", a()),
+        "Exp2" => format!("exp2f({})", a()),
+        "Log2" => format!("log2f({})", a()),
+        "Recip" => format!("1.0f / {}", a()),
+        "Sigmoid" => format!("1.0f / (1.0f + expf(-{}))", a()),
+        "Add" => format!("{} + {}", a(), b()),
+        "Mul" => format!("{} * {}", a(), b()),
+        other => panic!("region_codegen: unknown elementwise op {other}"),
     }
 }
 
@@ -324,7 +403,7 @@ pub(crate) fn compile_region(
     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
+    // FE strides and elementwise shapes.
     // 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()));
@@ -334,6 +413,19 @@ pub(crate) fn compile_region(
         let fs_struct: &FusionStart = (***fs_op).downcast_ref::<FusionStart>().unwrap();
         all_vars.extend(fs_struct.strides.iter().flat_map(|e| e.dyn_vars()));
     }
+    for &elem_idx in &region.elementwise_topo {
+        let elem_op = llir_graph[elem_idx].to_dialect::<dyn KernelOp>().unwrap();
+        if let Some(elem) = (***elem_op).downcast_ref::<CudaUnaryElementwise>() {
+            all_vars.extend(elem.shape.iter().flat_map(|e| e.dyn_vars()));
+            all_vars.extend(elem.in_strides.iter().flat_map(|e| e.dyn_vars()));
+            all_vars.extend(elem.out_strides.iter().flat_map(|e| e.dyn_vars()));
+        } else if let Some(elem) = (***elem_op).downcast_ref::<CudaBinaryElementwise>() {
+            all_vars.extend(elem.out_shape.iter().flat_map(|e| e.dyn_vars()));
+            all_vars.extend(elem.a_stride.iter().flat_map(|e| e.dyn_vars()));
+            all_vars.extend(elem.b_stride.iter().flat_map(|e| e.dyn_vars()));
+            all_vars.extend(elem.out_stride.iter().flat_map(|e| e.dyn_vars()));
+        }
+    }
 
     let cuda_ty = cuda_dtype(dtype);
     let includes = dtype_includes(&[dtype]);
@@ -359,19 +451,19 @@ pub(crate) fn compile_region(
     }
     let signature = signature_params.join(", ");
 
-    // Body: read FS leaves, then walk FusedX in topo order emitting a
+    // Body: read FS leaves, then walk elementwise nodes 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()).
+    // 0..fs_nodes.len(), elementwise nodes get fs_nodes.len()..(+ elementwise_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() {
+    for (i, &op_idx) in region.elementwise_topo.iter().enumerate() {
         local_idx_map.insert(op_idx, fs_count + i);
     }
     let local_name = |n: NodeIndex| format!("v_{}", local_idx_map[&n]);
@@ -394,12 +486,22 @@ pub(crate) fn compile_region(
         ));
     }
 
-    // FusedX ops in topo order. Each looks up its predecessor locals
+    // Elementwise 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 {
+    for &op_idx in &region.elementwise_topo {
         let op_ref = llir_graph[op_idx].to_dialect::<dyn KernelOp>().unwrap();
-        let op_name = op_ref.kernel_name();
+        let (elem_name, elem_dtype) =
+            if let Some(elem) = (***op_ref).downcast_ref::<CudaUnaryElementwise>() {
+                (elem.op.as_str(), elem.dtype)
+            } else if let Some(elem) = (***op_ref).downcast_ref::<CudaBinaryElementwise>() {
+                (elem.op.as_str(), elem.dtype)
+            } else {
+                panic!(
+                    "region_codegen: expected Cuda*Elementwise op, got {}",
+                    op_ref.kernel_name()
+                );
+            };
 
         let mut input_locals: Vec<String> = llir_graph
             .edges_directed(op_idx, Direction::Incoming)
@@ -418,15 +520,16 @@ pub(crate) fn compile_region(
         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);
+        let expr = elementwise_body(elem_name, &inputs_ref, elem_dtype);
+        let expr = elementwise_init_expr(&expr, elem_dtype, cuda_ty);
         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
+    // FE write: pick the elementwise node feeding FE (its single incoming edge in
+    // the region — an elementwise node 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)
@@ -474,3 +577,63 @@ pub(crate) fn compile_region(
         constants: FxHashMap::default(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::kernel::fusion::elementwise::CudaBinaryElementwise;
+    use luminal::op::LLIROp;
+    use luminal::prelude::petgraph::algo::toposort;
+
+    /// Helper: wrap a `KernelOp` in an `LLIROp` of the kernel dialect.
+    fn llir_of(op: impl KernelOp + 'static) -> LLIROp {
+        LLIROp::new::<dyn KernelOp>(Box::new(op) as Box<dyn KernelOp>)
+    }
+
+    /// Reproducer for the `FusionStart with no predecessor` panic at
+    /// `region_codegen.rs:232`. The egglog rolling pass + iterated mode
+    /// (`LUMINAL_LOOP_ROLL_ITERATE=1`) has been observed to produce LLIR
+    /// graphs where a `FusionStart` marker is reached as a region leaf
+    /// during the FE→FS walk but has no incoming edge — meaning the
+    /// region has nothing to read from. `build_compile_units` then
+    /// panics when constructing `external_inputs` because every FS leaf
+    /// is required to have exactly one external producer.
+    ///
+    /// Until that path is fixed, this test pins the failure mode so a
+    /// regression doesn't silently change the panic message or location.
+    /// `should_panic` rather than `ignore` so it stays runnable in CI
+    /// and surfaces if the panic ever moves.
+    #[test]
+    #[should_panic(expected = "FusionStart with no predecessor")]
+    fn fusion_start_with_no_predecessor_panics() {
+        // Minimal reproducer:
+        //
+        //   (no input) ──▶ FusionStart ──▶ CudaBinaryElementwise ──▶ FusionEnd
+        //
+        // CudaBinaryElementwise is a binary op (n_inputs = 2) so a real region would
+        // have two FS leaves. For this panic-shape test only the *first*
+        // FS leaf needs a missing predecessor — `build_compile_units`
+        // panics in `expect("FusionStart with no predecessor")` as soon
+        // as any FS in `fs_topo` lacks one. We add only one FS edge so
+        // CudaBinaryElementwise has a dangling second input slot, but that's fine:
+        // we're testing the specific panic path inside `build_compile_units`,
+        // not full kernel codegen.
+        let mut llir: LLIRGraph = LLIRGraph::default();
+
+        let fs_node = llir.add_node(llir_of(FusionStart::default()));
+        let fadd_node = llir.add_node(llir_of(CudaBinaryElementwise::default()));
+        let fe_node = llir.add_node(llir_of(FusionEnd::default()));
+
+        // FusionStart → CudaBinaryElementwise → FusionEnd.
+        llir.add_edge(fs_node, fadd_node, ());
+        llir.add_edge(fadd_node, fe_node, ());
+
+        let topo = toposort(&llir, None).expect("LLIR cycle in test setup");
+        let absorbed = globally_absorbed_markers(&llir);
+
+        // This is the call that panics with `FusionStart with no
+        // predecessor` because `fs_node`'s incoming-edges iterator is
+        // empty.
+        let _ = build_compile_units(&topo, &llir, &absorbed);
+    }
+}

crates/luminal_cuda_lite/src/kernel/generic_matmul.rs

@@ -0,0 +1,319 @@
+use std::sync::Arc;
+
+use crate::{
+    compile_module_image_for_current_device, cuda_dtype,
+    kernel::{
+        KernelOp,
+        hlir::{dtype_includes, generate_dyn_dims_defines},
+    },
+};
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::{
+    egglog_utils::{
+        api::{Rule, SortDef, sort},
+        base::{DTYPE, ELIST, EXPRESSION, OP_KIND},
+        extract_dtype, extract_expr, extract_expr_list,
+    },
+    op::*,
+    prelude::*,
+    shape::flatten_strides,
+};
+
+#[derive(Default, Debug, Clone)]
+pub struct GenericMatmul {
+    out_shape: Vec<Expression>,
+    mul_shape: Vec<Expression>,
+    k: Expression,
+    lhs_strides: Vec<Expression>,
+    rhs_strides: Vec<Expression>,
+    sum_input_strides: Vec<Expression>,
+    sum_iter_stride: Expression,
+    out_strides: Vec<Expression>,
+    dtype: DType,
+}
+
+impl EgglogOp for GenericMatmul {
+    fn sort(&self) -> SortDef {
+        sort(
+            OP_KIND,
+            "GenericMatmul",
+            &[
+                ("out_shape", ELIST),
+                ("mul_shape", ELIST),
+                ("k", EXPRESSION),
+                ("lhs_strides", ELIST),
+                ("rhs_strides", ELIST),
+                ("sum_input_strides", ELIST),
+                ("sum_iter_stride", EXPRESSION),
+                ("out_strides", ELIST),
+                ("dtype", DTYPE),
+            ],
+        )
+    }
+
+    fn n_inputs(&self) -> usize {
+        2
+    }
+
+    fn rewrites(&self) -> Vec<Rule> {
+        vec![
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?mul (Op (Mul ?mul_shape ?lhs_strides ?rhs_strides ?mul_out_strides)
+                            (ICons ?lhs (ICons ?rhs (INil)))))
+                        (= ?sum (Op (Sum ?out_shape ?k ?sum_input_strides ?sum_iter_stride ?out_strides)
+                            (ICons ?mul (INil))))
+                        (= ?dt (dtype ?sum))
+                    )
+                    (
+                        (let ?generic (Op (GenericMatmul
+                            ?out_shape
+                            ?mul_shape
+                            ?k
+                            ?lhs_strides
+                            ?rhs_strides
+                            ?sum_input_strides
+                            ?sum_iter_stride
+                            ?out_strides
+                            ?dt)
+                            (ICons ?lhs (ICons ?rhs (INil)))))
+                        (union ?sum ?generic)
+                        (set (dtype ?generic) ?dt)
+                    )
+                    :ruleset matmul_backend
+                    :name \"generic-matmul-cuda-mul-sum\"
+                )",
+            ),
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?mul (Op (Mul ?mul_shape ?lhs_strides ?rhs_strides ?mul_out_strides)
+                            (ICons ?lhs (ICons ?rhs (INil)))))
+                        (= ?sum (Op (Sum ?out_shape ?k ?sum_input_strides ?sum_iter_stride ?out_strides)
+                            (ICons ?mul (INil))))
+                        (= ?sum (Op (GenericMatmul
+                            ?go ?gm ?gk ?gls ?grs ?gsis ?gsit ?gos ?gdt)
+                            ?generic_inputs))
+                    )
+                    (
+                        (delete (Op (Sum ?out_shape ?k ?sum_input_strides ?sum_iter_stride ?out_strides)
+                            (ICons ?mul (INil))))
+                    )
+                    :ruleset cleanup
+                    :name \"delete-sum-when-generic-matmul-exists\"
+                )",
+            ),
+            Rule::raw(
+                "(rule
+                    (
+                        (= ?kernel_sum (Op (KernelSum ?out_shape ?k ?sum_input_strides ?sum_iter_stride ?out_strides ?dt)
+                            ?sum_inputs))
+                        (= ?kernel_sum (Op (GenericMatmul
+                            ?go ?gm ?gk ?gls ?grs ?gsis ?gsit ?gos ?gdt)
+                            ?generic_inputs))
+                    )
+                    ((delete (Op (KernelSum ?out_shape ?k ?sum_input_strides ?sum_iter_stride ?out_strides ?dt)
+                        ?sum_inputs)))
+                    :ruleset cleanup
+                    :name \"delete-kernel-sum-when-generic-matmul-exists\"
+                )",
+            ),
+        ]
+    }
+
+    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(),
+                mul_shape: extract_expr_list(egraph, kind_children[1], list_cache, expr_cache)
+                    .unwrap(),
+                k: extract_expr(egraph, kind_children[2], expr_cache).unwrap(),
+                lhs_strides: extract_expr_list(egraph, kind_children[3], list_cache, expr_cache)
+                    .unwrap(),
+                rhs_strides: extract_expr_list(egraph, kind_children[4], list_cache, expr_cache)
+                    .unwrap(),
+                sum_input_strides: extract_expr_list(
+                    egraph,
+                    kind_children[5],
+                    list_cache,
+                    expr_cache,
+                )
+                .unwrap(),
+                sum_iter_stride: extract_expr(egraph, kind_children[6], expr_cache).unwrap(),
+                out_strides: extract_expr_list(egraph, kind_children[7], list_cache, expr_cache)
+                    .unwrap(),
+                dtype: extract_dtype(egraph, kind_children[8]),
+            })),
+            input_enodes,
+        )
+    }
+}
+
+impl KernelOp for GenericMatmul {
+    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.all_dyn_vars();
+        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_outputs = self.output_size();
+        let sum_base_idx = flatten_strides(&self.out_shape, &self.sum_input_strides).to_kernel();
+        let iter_offset = self.sum_iter_stride.to_kernel().replace("const_z", "i");
+        let lhs_idx = flatten_strides(&self.mul_shape, &self.lhs_strides)
+            .to_kernel()
+            .replace("const_z", "mul_idx");
+        let rhs_idx = flatten_strides(&self.mul_shape, &self.rhs_strides)
+            .to_kernel()
+            .replace("const_z", "mul_idx");
+        let out_idx = flatten_strides(&self.out_shape, &self.out_strides).to_kernel();
+        let k = self.k.to_kernel();
+
+        let kernel = format!(
+            "{includes}
+#define WARP_SIZE 32
+#define THREADS_PER_BLOCK 256
+#define FULL_MASK 0xffffffff
+{dyn_defines}
+extern \"C\" {{
+    __global__ void generic_matmul({dtype} *out, const {dtype} *lhs, const {dtype} *rhs{dyn_dims_param}) {{
+        __shared__ float warp_sums[THREADS_PER_BLOCK / WARP_SIZE];
+        long long const_z = blockIdx.x;
+        if (const_z >= {n_outputs}) return;
+
+        int tid = threadIdx.x;
+        int lane_id = tid % WARP_SIZE;
+        int warp_id = tid / WARP_SIZE;
+
+        long long base_idx = {sum_base_idx};
+        long long iters = {k};
+
+        float partial = 0.0f;
+        for (long long i = tid; i < iters; i += THREADS_PER_BLOCK) {{
+            long long mul_idx = base_idx + {iter_offset};
+            partial += static_cast<float>(lhs[{lhs_idx}]) * static_cast<float>(rhs[{rhs_idx}]);
+        }}
+
+        #pragma unroll
+        for (int s = WARP_SIZE / 2; s > 0; s >>= 1) {{
+            partial += __shfl_down_sync(FULL_MASK, partial, s);
+        }}
+
+        if (lane_id == 0) {{
+            warp_sums[warp_id] = partial;
+        }}
+        __syncthreads();
+
+        if (warp_id == 0) {{
+            float block_sum = tid < (THREADS_PER_BLOCK / WARP_SIZE) ? warp_sums[tid] : 0.0f;
+
+            #pragma unroll
+            for (int s = (THREADS_PER_BLOCK / WARP_SIZE) / 2; s > 0; s >>= 1) {{
+                block_sum += __shfl_down_sync(FULL_MASK, block_sum, s);
+            }}
+
+            if (tid == 0) {{
+                out[{out_idx}] = ({dtype})block_sum;
+            }}
+        }}
+    }}
+}}",
+            n_outputs = n_outputs.to_kernel(),
+        );
+
+        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("generic_matmul").unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        (
+            func,
+            module,
+            kernel,
+            (n_outputs, 1.into(), 1.into()),
+            (256.into(), 1.into(), 1.into()),
+            32.into(),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        self.out_shape
+            .iter()
+            .copied()
+            .product::<Expression>()
+            .max(Expression::from(1))
+    }
+
+    fn all_dyn_vars(&self) -> FxHashSet<char> {
+        self.out_shape
+            .iter()
+            .flat_map(|e| e.dyn_vars())
+            .chain(self.mul_shape.iter().flat_map(|e| e.dyn_vars()))
+            .chain(self.k.dyn_vars())
+            .chain(self.lhs_strides.iter().flat_map(|e| e.dyn_vars()))
+            .chain(self.rhs_strides.iter().flat_map(|e| e.dyn_vars()))
+            .chain(self.sum_input_strides.iter().flat_map(|e| e.dyn_vars()))
+            .chain(self.sum_iter_stride.dyn_vars())
+            .chain(self.out_strides.iter().flat_map(|e| e.dyn_vars()))
+            .collect()
+    }
+
+    fn output_bytes(&self) -> Expression {
+        (self.output_size() * self.dtype.bits()).ceil_div(8)
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        (self.output_size() * self.k * self.dtype.bits() * 2).ceil_div(8)
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_bytes()
+    }
+
+    fn flops(&self) -> Expression {
+        self.output_size() * self.k * 2
+    }
+
+    fn output_dtype(&self) -> DType {
+        self.dtype
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "GenericMatmul"
+    }
+}

crates/luminal_cuda_lite/src/kernel/matmul2d.rs

@@ -0,0 +1,643 @@
+//! Direct 2D matmul kernel — bypasses egglog rewrites, used as a custom op
+//! for matmul shapes where the cublaslt egg rules don't reliably fire.
+//!
+//! The cublaslt 2D rules in `host/cublaslt/cublaslt_*Cm_rewrite.egg` /
+//! `cublaslt_Rm*_rewrite.egg` are *supposed* to match any 2D matmul whose
+//! Mul + SumReduce broadcast lowering has the expected stride patterns,
+//! and the conditional matmul cleanup is *supposed* to delete the
+//! elementwise Mul + KernelSumReduce fallback whenever a cublaslt alternative
+//! exists. In practice both fail to fire reliably for the VAE's mid-block
+//! `AttnBlock` matmuls — at 1024² that lets the search occasionally pick
+//! the broadcast-Mul path for `q @ kᵀ`, generating a `(HW, HW, C) =
+//! (16384, 16384, 512)` ≈ 524 GiB single intermediate that OOMs the GPU.
+//!
+//! Same approach as `kernel::conv2d`: define a `KernelOp`, wrap it in a
+//! `CustomOp`, expose a tiny `pub fn` so callers don't see the
+//! `cx.custom_op` plumbing. This is opaque to egglog by design — we
+//! aren't trying to fuse with surrounding ops, just guarantee a sane
+//! lowering for the matmuls we know are problematic.
+//!
+//! The CUDA implementation is a textbook 2D-blocked SGEMM:
+//!   * 16×16 output tile per block (256 threads)
+//!   * Tiled load of A and B into shared memory in K-size chunks
+//!   * Each thread accumulates one output element across all K-tiles
+//!   * Optional bias broadcast along the M axis at write-out
+//!   * `transpose_b` toggles between row-major B `(K, N)` and row-major
+//!     B `(N, K)` (i.e. the `A @ Bᵀ` pattern that linear/projection
+//!     layers use).
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::{
+    dtype::DType, op::CustomOp, op::LLIROp, prelude::FxHashMap, prelude::GraphTensor,
+    shape::Expression,
+};
+
+use crate::compile_module_image_for_current_device;
+use crate::kernel::KernelOp;
+
+/// Direct 2D matmul `(M, K) × {(K, N) | (N, K)} → (M, N)` with optional
+/// per-output-column bias and an optional batch axis. A and output are
+/// always F32. B can be F32 or BF16; BF16 is converted to F32 on each
+/// load, which avoids materializing the cast as a separate intermediate
+/// tensor (important for the text encoder / transformer where the F32-
+/// cast weights would not fit in GPU memory). All shape parameters are
+/// static (baked into the CUDA source via #defines).
+///
+/// When `batch > 1` the kernel does `batch` independent 2D matmuls in
+/// parallel: A is `(batch, M, K)`, B is `(batch, *, *)` with the same
+/// per-batch shape, output is `(batch, M, N)`. All three are assumed
+/// contiguous row-major across batches (i.e. `a_batch_stride = M*K`,
+/// `b_batch_stride = K*N` or `N*K` depending on `transpose_b`,
+/// `out_batch_stride = M*N`). Bias does NOT have a batch axis — it's
+/// `(N,)` and broadcast across batches.
+/// Activation epilogue fused into the matmul kernel's store path.
+///
+/// Saves one full pass over the output buffer per MLP layer — the same
+/// trick cuBLASLt does with `CUBLASLT_EPILOGUE_RELU_BIAS` etc., but
+/// inside our custom kernel so we don't have to invoke cuBLASLt.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum Activation {
+    #[default]
+    None,
+    Relu,
+    Sigmoid,
+}
+
+#[derive(Debug, Clone)]
+pub struct Matmul2DKernel {
+    pub m: usize,
+    pub n: usize,
+    pub k: usize,
+    pub batch: usize,
+    /// If `true`, B is interpreted as `(N, K)` row-major and accessed as
+    /// `B[n][k]` (i.e. `A @ Bᵀ`). If `false`, B is `(K, N)` row-major and
+    /// accessed as `B[k][n]` (i.e. `A @ B`).
+    pub transpose_b: bool,
+    pub has_bias: bool,
+    /// Storage dtype of B. Currently F32 or BF16 are supported.
+    pub weight_dtype: DType,
+    /// Activation applied to `acc + bias` before writing to C.
+    /// Defaults to None; ReLU and Sigmoid avoid a separate elementwise
+    /// pass over the matmul output.
+    pub activation: Activation,
+    /// When `Some(split)`, A is read from two source pointers:
+    ///   columns `0..split`  → `A_lo`, stride `split` per row
+    ///   columns `split..K` → `A_hi`, stride `K - split` per row
+    /// This lets a `cat(A_lo, A_hi)` materialization be skipped entirely —
+    /// the K-loop's A-load branches on the column index instead. `None`
+    /// keeps the existing single-pointer path. Only supported for
+    /// `batch == 1` (DLRM's use case); the kernel asserts on this.
+    pub a_split: Option<usize>,
+}
+
+const TILE: usize = 16;
+
+impl KernelOp for Matmul2DKernel {
+    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 bias_param = if self.has_bias {
+            ", const float* __restrict__ bias"
+        } else {
+            ""
+        };
+        let bias_add = if self.has_bias {
+            "    acc += bias[n];\n"
+        } else {
+            ""
+        };
+        let activation_apply = match self.activation {
+            Activation::None => "",
+            // Branchless ReLU; keeps the fully-occupied write path simple.
+            Activation::Relu => "    acc = fmaxf(acc, 0.0f);\n",
+            // Sigmoid: 1/(1+exp(-acc)). Used by DLRM's final layer.
+            Activation::Sigmoid => "    acc = 1.0f / (1.0f + __expf(-acc));\n",
+        };
+        // A-input parameter declaration + per-K-tile load expression depend
+        // on whether the caller asked for the dual-source (split) path.
+        // Single-source (default) keeps the original `const float* A` and
+        // reads `A[a_m * K + a_k]`. Split mode takes two pointer args
+        // (A_lo / A_hi) and selects between them at runtime by comparing
+        // `a_k` against the compile-time-baked split column.
+        let (a_param_decl, a_load_expr) = if let Some(split) = self.a_split {
+            assert!(
+                split > 0 && split < self.k,
+                "Matmul2DKernel a_split must be in 1..K; got split={split}, K={}",
+                self.k
+            );
+            assert_eq!(
+                self.batch, 1,
+                "Matmul2DKernel a_split path only supports batch=1 (got batch={})",
+                self.batch
+            );
+            let hi = self.k - split;
+            (
+                "const float* __restrict__ A_lo, const float* __restrict__ A_hi"
+                    .to_string(),
+                format!(
+                    "((a_k < {split}) \
+                       ? A_lo[a_m * {split} + a_k] \
+                       : A_hi[a_m * {hi} + (a_k - {split})])"
+                ),
+            )
+        } else {
+            (
+                "const float* __restrict__ A".to_string(),
+                "A[a_batch_off + a_m * K + a_k]".to_string(),
+            )
+        };
+        // We want Bs[ty][tx] = B_effective[k0+ty][b_n_base+tx] where:
+        //   transpose_b=false: B is (K, N) row-major → B[(k0+ty)*N + (b_n_base+tx)]
+        //   transpose_b=true:  B is (N, K) row-major → B[(b_n_base+tx)*K + (k0+ty)]
+        // Plus the per-batch offset (`b_batch_off`).
+        let b_index_expr = if self.transpose_b {
+            "b_batch_off + (b_n_base + tx) * K + (k0 + ty)"
+        } else {
+            "b_batch_off + (k0 + ty) * N + (b_n_base + tx)"
+        };
+        // Convert B's element to float on load. For BF16 we declare B as
+        // `__nv_bfloat16*` and use `__bfloat162float`; for F32 it's a no-op.
+        let (b_param_type, b_load_expr, bf16_include) = match self.weight_dtype {
+            DType::F32 => (
+                "const float* __restrict__ B",
+                format!("B[{b_index_expr}]"),
+                "",
+            ),
+            DType::Bf16 => (
+                "const __nv_bfloat16* __restrict__ B",
+                format!("__bfloat162float(B[{b_index_expr}])"),
+                "#include <cuda_bf16.h>\n",
+            ),
+            other => panic!("Matmul2DKernel: unsupported weight_dtype {other:?}"),
+        };
+
+        let kernel = format!(
+            "
+{bf16_include}extern \"C\" __global__ void matmul_2d_kernel(
+    float* __restrict__ C,
+    {a_param_decl},
+    {b_param_type}{bias_param}
+) {{
+    const int M = {m};
+    const int N = {n};
+    const int K = {k};
+    const int TILE = {tile};
+
+    __shared__ float As[{tile}][{tile}];
+    __shared__ float Bs[{tile}][{tile}];
+
+    int bx = blockIdx.x;  // tile column (n)
+    int by = blockIdx.y;  // tile row (m)
+    int batch = blockIdx.z; // batch index (0..BATCH-1)
+    int tx = threadIdx.x; // 0..TILE-1, output col within tile
+    int ty = threadIdx.y; // 0..TILE-1, output row within tile
+
+    int m_global = by * TILE + ty;
+    int n_global = bx * TILE + tx;
+
+    int a_m_base = by * TILE;
+    int b_n_base = bx * TILE;
+
+    // Per-batch base pointer offsets (contiguous row-major across batches).
+    int a_batch_off = batch * (M * K);
+    int b_batch_off = batch * (K * N);
+    int c_batch_off = batch * (M * N);
+
+    float acc = 0.0f;
+
+    int n_tiles = (K + TILE - 1) / TILE;
+    for (int t = 0; t < n_tiles; ++t) {{
+        int k0 = t * TILE;
+
+        // Load A tile (TILE, TILE) row-major from A[m, k]. In single-source
+        // mode this is `A[a_batch_off + a_m * K + a_k]`. In split mode the
+        // load expression branches on `a_k < split` (baked in by the host).
+        int a_m = a_m_base + ty;
+        int a_k = k0 + tx;
+        As[ty][tx] = (a_m < M && a_k < K) ? ({a_load_expr}) : 0.0f;
+
+        // Load B tile depending on transpose_b
+        int b_n_or_k = b_n_base + tx;  // for transpose_b=true this is N; for =false this is N
+        int b_k_or_k = k0 + ty;        // similarly
+        // We compute Bs[ty][tx] such that the inner loop reads Bs[k_local][n_local] = B[k][n].
+        // For transpose_b=true (B is (N,K)):  B[k][n] in math = B_storage[n][k] = B[(b_n_base+tx)*K + (k0+ty)]
+        // For transpose_b=false (B is (K,N)): B[k][n] in math = B_storage[k][n] = B[(k0+ty)*N + (b_n_base+tx)]
+        bool b_in_bounds = ({transpose_b} ? (b_n_or_k < N && b_k_or_k < K)
+                                          : (b_k_or_k < K && b_n_or_k < N));
+        Bs[ty][tx] = b_in_bounds ? ({b_load_expr}) : 0.0f;
+
+        __syncthreads();
+
+        #pragma unroll
+        for (int kk = 0; kk < {tile}; ++kk) {{
+            acc += As[ty][kk] * Bs[kk][tx];
+        }}
+        __syncthreads();
+    }}
+
+    if (m_global < M && n_global < N) {{
+        int n = n_global;
+{bias_add}{activation_apply}        C[c_batch_off + m_global * N + n_global] = acc;
+    }}
+}}
+",
+            m = self.m,
+            n = self.n,
+            k = self.k,
+            tile = TILE,
+            transpose_b = self.transpose_b,
+            b_load_expr = b_load_expr,
+            b_param_type = b_param_type,
+            bias_param = bias_param,
+            bias_add = bias_add,
+            activation_apply = activation_apply,
+            bf16_include = bf16_include,
+            a_param_decl = a_param_decl,
+            a_load_expr = a_load_expr,
+        );
+
+        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("matmul_2d_kernel").unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        let grid_x = self.n.div_ceil(TILE);
+        let grid_y = self.m.div_ceil(TILE);
+        (
+            func,
+            module,
+            kernel,
+            (
+                Expression::from(grid_x),
+                Expression::from(grid_y),
+                Expression::from(self.batch),
+            ),
+            (
+                Expression::from(TILE),
+                Expression::from(TILE),
+                Expression::from(1usize),
+            ),
+            Expression::from(0usize),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        Expression::from(self.batch * self.m * self.n)
+    }
+
+    fn output_bytes(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn output_dtype(&self) -> DType {
+        DType::F32
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        // K elements from A (F32) + K elements from B (F32 or BF16) + maybe bias (F32).
+        let b_bytes = match self.weight_dtype {
+            DType::F32 => 4,
+            DType::Bf16 => 2,
+            _ => 4,
+        };
+        let bias_bytes = if self.has_bias { 4 } else { 0 };
+        Expression::from(
+            self.batch * self.m * self.n * (self.k * 4 + self.k * b_bytes + bias_bytes),
+        )
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn flops(&self) -> Expression {
+        let per_out = self.k * 2 + if self.has_bias { 1 } else { 0 };
+        Expression::from(self.batch * self.m * self.n * per_out)
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        match (self.has_bias, self.activation, self.a_split.is_some()) {
+            (true, Activation::Relu, false) => "Matmul2D_BiasRelu",
+            (true, Activation::Sigmoid, false) => "Matmul2D_BiasSigmoid",
+            (true, Activation::None, false) => "Matmul2D_Bias",
+            (false, Activation::Relu, false) => "Matmul2D_Relu",
+            (false, Activation::Sigmoid, false) => "Matmul2D_Sigmoid",
+            (false, Activation::None, false) => "Matmul2D",
+            (true, Activation::Relu, true) => "Matmul2D_BiasRelu_SplitA",
+            (true, Activation::Sigmoid, true) => "Matmul2D_BiasSigmoid_SplitA",
+            (true, Activation::None, true) => "Matmul2D_Bias_SplitA",
+            (false, Activation::Relu, true) => "Matmul2D_Relu_SplitA",
+            (false, Activation::Sigmoid, true) => "Matmul2D_Sigmoid_SplitA",
+            (false, Activation::None, true) => "Matmul2D_SplitA",
+        }
+    }
+}
+
+/// CustomOp wrapper for [`Matmul2DKernel`].
+#[derive(Debug, Clone)]
+pub struct Matmul2DCustom(pub Matmul2DKernel);
+
+impl CustomOp for Matmul2DCustom {
+    fn to_llir_op(&self) -> LLIROp {
+        LLIROp::new::<dyn KernelOp>(Box::new(self.0.clone()) as Box<dyn KernelOp>)
+    }
+}
+
+/// `(M, K) @ (K, N) -> (M, N)` for row-major F32 inputs. No bias.
+pub fn matmul_2d(a: GraphTensor, b: GraphTensor) -> GraphTensor {
+    matmul_inner(a, b, /*transpose_b=*/ false, None, Activation::None)
+}
+
+/// `(M, K) @ (N, K)ᵀ -> (M, N)` for row-major F32 inputs. No bias.
+/// Use this for `A @ Bᵀ` where B is stored row-major as `(N, K)` — the
+/// pattern produced by linear / projection layers (`x @ w.t()`).
+pub fn matmul_2d_t(a: GraphTensor, b: GraphTensor) -> GraphTensor {
+    matmul_inner(a, b, /*transpose_b=*/ true, None, Activation::None)
+}
+
+/// Linear projection with bias: `(M, K) @ (N, K)ᵀ + bias` where bias is
+/// `(N,)`, row-major F32 throughout.
+pub fn linear_bias(a: GraphTensor, b: GraphTensor, bias: GraphTensor) -> GraphTensor {
+    matmul_inner(a, b, /*transpose_b=*/ true, Some(bias), Activation::None)
+}
+
+/// Like [`linear_bias`] but applies ReLU in the kernel epilogue. Saves
+/// one full pass over the output buffer per layer.
+pub fn linear_bias_relu(a: GraphTensor, b: GraphTensor, bias: GraphTensor) -> GraphTensor {
+    matmul_inner(a, b, /*transpose_b=*/ true, Some(bias), Activation::Relu)
+}
+
+/// Like [`linear_bias`] but applies Sigmoid in the kernel epilogue.
+/// Used for the final layer of binary-classifier MLPs (DLRM CTR head).
+pub fn linear_bias_sigmoid(a: GraphTensor, b: GraphTensor, bias: GraphTensor) -> GraphTensor {
+    matmul_inner(a, b, /*transpose_b=*/ true, Some(bias), Activation::Sigmoid)
+}
+
+/// Two-A-input variant of [`linear_bias`].
+///
+/// Computes `cat(a_lo, a_hi) @ bᵀ + bias` *without* materializing the
+/// concat — the K-loop's A-load reads from `a_lo` for columns `0..K_lo`
+/// and from `a_hi` for columns `K_lo..K_lo+K_hi`. Logically equivalent
+/// to feeding `concat_along(a_lo, a_hi, 1)` into [`linear_bias`], but
+/// skips ~9 scaffolding kernels (Iota + Cast + Gather + masked-add) per
+/// concat call.
+///
+/// Shapes:
+///   * `a_lo`: `(M, K_lo)` F32
+///   * `a_hi`: `(M, K_hi)` F32
+///   * `b`:    `(N, K_lo + K_hi)` F32 (transposed convention, same as
+///             [`linear_bias`])
+///   * `bias`: `(N,)` F32
+///
+/// Output: `(M, N)` F32. Only 2D inputs are supported (batch=1).
+pub fn linear_bias_split_a(
+    a_lo: GraphTensor,
+    a_hi: GraphTensor,
+    b: GraphTensor,
+    bias: GraphTensor,
+) -> GraphTensor {
+    matmul_inner_split_a(a_lo, a_hi, b, Some(bias), Activation::None)
+}
+
+/// Like [`linear_bias_split_a`] but applies ReLU in the kernel epilogue.
+/// Use this for hidden MLP layers that consume a concat of two upstream
+/// tensors — the natural shape of DLRM's top-MLP first layer (which reads
+/// `cat(dense_out, interactions)`).
+pub fn linear_bias_relu_split_a(
+    a_lo: GraphTensor,
+    a_hi: GraphTensor,
+    b: GraphTensor,
+    bias: GraphTensor,
+) -> GraphTensor {
+    matmul_inner_split_a(a_lo, a_hi, b, Some(bias), Activation::Relu)
+}
+
+/// Like [`linear_bias_split_a`] but applies Sigmoid in the kernel
+/// epilogue.
+pub fn linear_bias_sigmoid_split_a(
+    a_lo: GraphTensor,
+    a_hi: GraphTensor,
+    b: GraphTensor,
+    bias: GraphTensor,
+) -> GraphTensor {
+    matmul_inner_split_a(a_lo, a_hi, b, Some(bias), Activation::Sigmoid)
+}
+
+/// Mixed-precision linear (no bias): `A (F32, M, K) @ B (BF16, N, K)ᵀ → (F32, M, N)`.
+///
+/// Lowers as plain HLIR — `Cast(A, BF16) @ permute(B_bf16) → Cast(F32)`.
+/// The activation cast and output cast are tiny (M*K and M*N elements;
+/// the K=hidden weight stays BF16). The inner BF16 matmul matches the
+/// existing cublaslt rewrite rules and runs as
+/// `CUBLAS_COMPUTE_32F_FAST_16BF` — Hopper's native 2× BF16 path.
+pub fn linear_no_bias_bf16_w(a: GraphTensor, b_bf16: GraphTensor) -> GraphTensor {
+    assert_eq!(a.dtype, DType::F32, "linear_no_bias_bf16_w expects F32 A");
+    assert_eq!(
+        b_bf16.dtype,
+        DType::Bf16,
+        "linear_no_bias_bf16_w expects BF16 B"
+    );
+    let a_dims = a.dims();
+    let b_dims = b_bf16.dims();
+    assert_eq!(a_dims.len(), 2);
+    assert_eq!(b_dims.len(), 2);
+    let a_bf16 = a.cast(DType::Bf16);
+    let b_kn = b_bf16.permute((1, 0));
+    a_bf16.matmul(b_kn).cast(DType::F32)
+}
+
+/// Batched matmul: `A (B, M, K) @ B (B, K, N) → (B, M, N)`, all F32 row-major.
+pub fn matmul_3d(a: GraphTensor, b: GraphTensor) -> GraphTensor {
+    matmul_inner(a, b, /*transpose_b=*/ false, None, Activation::None)
+}
+
+/// Batched matmul with B-transpose: `A (B, M, K) @ B (B, N, K)ᵀ → (B, M, N)`.
+pub fn matmul_3d_t(a: GraphTensor, b: GraphTensor) -> GraphTensor {
+    matmul_inner(a, b, /*transpose_b=*/ true, None, Activation::None)
+}
+
+fn matmul_inner(
+    a: GraphTensor,
+    b: GraphTensor,
+    transpose_b: bool,
+    bias: Option<GraphTensor>,
+    activation: Activation,
+) -> GraphTensor {
+    assert_eq!(a.dtype, DType::F32, "matmul requires F32 A");
+    let weight_dtype = b.dtype;
+    assert!(
+        matches!(weight_dtype, DType::F32 | DType::Bf16),
+        "matmul B must be F32 or BF16, got {weight_dtype:?}",
+    );
+    let a_dims = a.dims();
+    let b_dims = b.dims();
+    assert_eq!(
+        a_dims.len(),
+        b_dims.len(),
+        "matmul A/B rank mismatch: {} vs {}",
+        a_dims.len(),
+        b_dims.len(),
+    );
+    assert!(
+        a_dims.len() == 2 || a_dims.len() == 3,
+        "matmul expects rank 2 or 3, got rank {}",
+        a_dims.len(),
+    );
+
+    let (batch, a_off) = if a_dims.len() == 3 {
+        let ba = a_dims[0].to_usize().expect("batch dim must be static");
+        let bb = b_dims[0].to_usize().expect("batch dim must be static");
+        assert_eq!(
+            ba, bb,
+            "matmul batch dim mismatch: A batch={ba}, B batch={bb}"
+        );
+        (ba, 1)
+    } else {
+        (1, 0)
+    };
+
+    let m = a_dims[a_off].to_usize().expect("M must be a static dim");
+    let k_a = a_dims[a_off + 1]
+        .to_usize()
+        .expect("K (A) must be a static dim");
+    let (n, k_b) = if transpose_b {
+        // B per-batch is (N, K)
+        let n = b_dims[a_off].to_usize().expect("N must be a static dim");
+        let k = b_dims[a_off + 1]
+            .to_usize()
+            .expect("K (B) must be a static dim");
+        (n, k)
+    } else {
+        // B per-batch is (K, N)
+        let k = b_dims[a_off]
+            .to_usize()
+            .expect("K (B) must be a static dim");
+        let n = b_dims[a_off + 1]
+            .to_usize()
+            .expect("N must be a static dim");
+        (n, k)
+    };
+    assert_eq!(k_a, k_b, "matmul K mismatch: A K={k_a}, B K={k_b}");
+    let k = k_a;
+
+    let has_bias = bias.is_some();
+    if let Some(bias) = bias {
+        let bdims = bias.dims();
+        assert_eq!(bdims.len(), 1, "matmul bias must be 1D");
+        assert_eq!(
+            bdims[0].to_usize().expect("bias dim must be static"),
+            n,
+            "matmul bias size must equal N"
+        );
+        assert_eq!(bias.dtype, DType::F32, "matmul bias must be F32");
+    }
+
+    let kern = Matmul2DKernel {
+        m,
+        n,
+        k,
+        batch,
+        transpose_b,
+        has_bias,
+        weight_dtype,
+        activation,
+        a_split: None,
+    };
+    let cx = unsafe { &mut *a.graph_ref };
+    let inputs: Vec<GraphTensor> = if let Some(bias) = bias {
+        vec![a, b, bias]
+    } else {
+        vec![a, b]
+    };
+    if batch == 1 {
+        cx.custom_op(Matmul2DCustom(kern), inputs, (m, n), DType::F32)
+    } else {
+        cx.custom_op(Matmul2DCustom(kern), inputs, (batch, m, n), DType::F32)
+    }
+}
+
+/// Internal helper for the split-A path. Validates shapes and dispatches
+/// to a [`Matmul2DKernel`] with `a_split = Some(K_lo)`. Always uses
+/// `transpose_b = true` (linear-projection convention; matches
+/// [`linear_bias`]). Only 2D inputs are supported.
+fn matmul_inner_split_a(
+    a_lo: GraphTensor,
+    a_hi: GraphTensor,
+    b: GraphTensor,
+    bias: Option<GraphTensor>,
+    activation: Activation,
+) -> GraphTensor {
+    assert_eq!(a_lo.dtype, DType::F32, "split-A matmul requires F32 A_lo");
+    assert_eq!(a_hi.dtype, DType::F32, "split-A matmul requires F32 A_hi");
+    let weight_dtype = b.dtype;
+    assert_eq!(
+        weight_dtype,
+        DType::F32,
+        "split-A matmul currently only supports F32 B (got {weight_dtype:?})"
+    );
+    let lo_dims = a_lo.dims();
+    let hi_dims = a_hi.dims();
+    let b_dims = b.dims();
+    assert_eq!(lo_dims.len(), 2, "split-A matmul A_lo must be 2D");
+    assert_eq!(hi_dims.len(), 2, "split-A matmul A_hi must be 2D");
+    assert_eq!(b_dims.len(), 2, "split-A matmul B must be 2D");
+    let m = lo_dims[0].to_usize().expect("M must be a static dim");
+    let m_hi = hi_dims[0].to_usize().expect("M (A_hi) must be a static dim");
+    assert_eq!(m, m_hi, "split-A matmul: A_lo and A_hi must have the same M");
+    let k_lo = lo_dims[1].to_usize().expect("K_lo must be a static dim");
+    let k_hi = hi_dims[1].to_usize().expect("K_hi must be a static dim");
+    let k = k_lo + k_hi;
+    let n = b_dims[0].to_usize().expect("N must be a static dim");
+    let k_b = b_dims[1].to_usize().expect("K (B) must be a static dim");
+    assert_eq!(
+        k, k_b,
+        "split-A matmul: A_lo.K + A_hi.K = {k} must equal B.K = {k_b}"
+    );
+    let has_bias = bias.is_some();
+    if let Some(bias) = bias {
+        let bdims = bias.dims();
+        assert_eq!(bdims.len(), 1, "split-A matmul bias must be 1D");
+        assert_eq!(
+            bdims[0].to_usize().expect("bias dim must be static"),
+            n,
+            "split-A matmul bias size must equal N"
+        );
+        assert_eq!(bias.dtype, DType::F32, "split-A matmul bias must be F32");
+    }
+    let kern = Matmul2DKernel {
+        m,
+        n,
+        k,
+        batch: 1,
+        transpose_b: true,
+        has_bias,
+        weight_dtype,
+        activation,
+        a_split: Some(k_lo),
+    };
+    let cx = unsafe { &mut *a_lo.graph_ref };
+    let inputs: Vec<GraphTensor> = if let Some(bias) = bias {
+        vec![a_lo, a_hi, b, bias]
+    } else {
+        vec![a_lo, a_hi, b]
+    };
+    cx.custom_op(Matmul2DCustom(kern), inputs, (m, n), DType::F32)
+}

crates/luminal_cuda_lite/src/kernel/mod.rs

@@ -9,14 +9,45 @@ use luminal_tracing::schema::{
 };
 use uuid::Uuid;
 
+pub mod conv2d;
 pub mod cuda_graph;
+pub mod dlrm_interact;
+pub mod embedding_bag;
 pub mod fusion;
+pub mod generic_matmul;
 pub mod hlir;
+pub mod matmul2d;
 pub mod other_ops;
+pub mod rope;
 
+pub use conv2d::KernelConv2D;
 pub use cuda_graph::*;
+pub use dlrm_interact::{
+    PairwiseDotLowerTriCustom, PairwiseDotLowerTriKernel, PairwiseDotLowerTriStackedCustom,
+    PairwiseDotLowerTriStackedKernel, dlrm_pairwise_dot_lower_tri,
+    dlrm_pairwise_dot_lower_tri_stacked,
+};
+pub use embedding_bag::{
+    EmbeddingBagSumCustom, EmbeddingBagSumKernel, MultiTableEmbeddingBagSumCustom,
+    MultiTableEmbeddingBagSumKernel, StackedEmbeddingBagKernel,
+    StackedEmbeddingBagSumCustom, embedding_bag_sum_kernel,
+    multi_table_embedding_bag_sum_kernel, stacked_embedding_bag_sum_kernel,
+};
+pub use generic_matmul::GenericMatmul;
+pub use matmul2d::{
+    Activation, Matmul2DCustom, Matmul2DKernel, linear_bias, linear_bias_relu,
+    linear_bias_relu_split_a, linear_bias_sigmoid, linear_bias_sigmoid_split_a,
+    linear_bias_split_a, linear_no_bias_bf16_w, matmul_2d, matmul_2d_t, matmul_3d, matmul_3d_t,
+};
+pub use rope::{RoPECustom, RoPEKernel, apply_rope};
 
-pub type Ops = (hlir::Ops, other_ops::Ops, fusion::Ops);
+pub type Ops = (
+    hlir::Ops,
+    other_ops::Ops,
+    conv2d::KernelConv2D,
+    GenericMatmul,
+    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/rope.rs

@@ -0,0 +1,189 @@
+//! Fused RoPE (rotary position embedding) — interleaved-pair convention.
+//!
+//! Replaces flux2's 6-op RoPE chain (split / slice / squeeze / neg / concat /
+//! merge_dims / 4× cast / mul / add) with a single kernel launch per call.
+//! ~120 RoPE calls per forward pass at full DiT depth.
+//!
+//! Convention: `repeat_interleave_real=True` (Flux 2 / diffusers), so adjacent
+//! dim pairs rotate together. For an input `[a0, b0, a1, b1, ...]` and per-
+//! position `(cos, sin)`, the output is
+//!   `out[2j]   = x[2j]   * cos[2j]   - x[2j+1] * sin[2j]`
+//!   `out[2j+1] = x[2j+1] * cos[2j+1] + x[2j]   * sin[2j+1]`
+//!
+//! Layout: x `(S, H, D)`, cos/sin `(S, D)` (broadcast across H).
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
+use luminal::{
+    dtype::DType, op::CustomOp, op::LLIROp, prelude::FxHashMap, prelude::GraphTensor,
+    shape::Expression,
+};
+
+use crate::compile_module_image_for_current_device;
+use crate::kernel::KernelOp;
+
+#[derive(Debug, Clone)]
+pub struct RoPEKernel {
+    pub s: usize,
+    pub h: usize,
+    pub d: usize,
+}
+
+const TPB: usize = 64;
+
+impl KernelOp for RoPEKernel {
+    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 s = self.s;
+        let h = self.h;
+        let d = self.d;
+        assert!(d.is_multiple_of(2), "RoPE head_dim must be even");
+        let kernel = format!(
+            r#"
+extern "C" __global__ void rope_kernel(
+    float* __restrict__ out,
+    const float* __restrict__ x,
+    const float* __restrict__ cos_,
+    const float* __restrict__ sin_
+) {{
+    const int S = {s};
+    const int H = {h};
+    const int D = {d};
+    int sh = blockIdx.x;       // 0..S*H
+    int s_idx = sh / H;
+    int tid = threadIdx.x;
+
+    const float* xr   = x    + sh    * D;
+    const float* cosr = cos_ + s_idx * D;
+    const float* sinr = sin_ + s_idx * D;
+    float* yr = out + sh * D;
+
+    for (int i = tid; i < D; i += {TPB}) {{
+        float xi = xr[i];
+        float xpair;
+        if ((i & 1) == 0) {{
+            // even: paired with i+1, rotated value is -x[i+1]
+            xpair = -xr[i + 1];
+        }} else {{
+            // odd: paired with i-1, rotated value is +x[i-1]
+            xpair = xr[i - 1];
+        }}
+        yr[i] = xi * cosr[i] + xpair * sinr[i];
+    }}
+}}
+"#
+        );
+
+        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("rope_kernel").unwrap();
+            compile_cache.insert(kernel.clone(), (module.clone(), func.clone()));
+            (module, func)
+        };
+
+        (
+            func,
+            module,
+            "rope_kernel".to_string(),
+            (
+                Expression::from(s * h),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            (
+                Expression::from(TPB),
+                Expression::from(1usize),
+                Expression::from(1usize),
+            ),
+            Expression::from(0usize),
+            FxHashMap::default(),
+        )
+    }
+
+    fn output_size(&self) -> Expression {
+        Expression::from(self.s * self.h * self.d)
+    }
+
+    fn output_bytes(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn output_dtype(&self) -> DType {
+        DType::F32
+    }
+
+    fn bytes_loaded(&self) -> Expression {
+        // x: full (S,H,D); cos/sin: (S,D) read H times each but cached.
+        Expression::from(self.s * self.h * self.d * 4 + self.s * self.d * 4 * 2)
+    }
+
+    fn bytes_stored(&self) -> Expression {
+        self.output_size() * 4
+    }
+
+    fn flops(&self) -> Expression {
+        // 4 per output element (mul, neg/load, mul, add).
+        Expression::from(self.s * self.h * self.d * 4)
+    }
+
+    fn kernel_name(&self) -> &'static str {
+        "RoPE"
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct RoPECustom(pub RoPEKernel);
+
+impl CustomOp for RoPECustom {
+    fn to_llir_op(&self) -> LLIROp {
+        LLIROp::new::<dyn KernelOp>(Box::new(self.0.clone()) as Box<dyn KernelOp>)
+    }
+}
+
+/// Apply RoPE: `x` shape `(S, H, D)` F32, `cos`/`sin` shape `(S, D)` F32.
+/// Returns `(S, H, D)` F32.
+pub fn apply_rope(x: GraphTensor, cos: GraphTensor, sin: GraphTensor) -> GraphTensor {
+    assert_eq!(x.dtype, DType::F32, "RoPE x must be F32");
+    let cos = if cos.dtype == DType::F32 {
+        cos
+    } else {
+        cos.cast(DType::F32)
+    };
+    let sin = if sin.dtype == DType::F32 {
+        sin
+    } else {
+        sin.cast(DType::F32)
+    };
+    let x_dims = x.dims();
+    assert_eq!(x_dims.len(), 3, "RoPE x must be 3-D (S, H, D)");
+    let s = x_dims[0].to_usize().expect("RoPE: S must be static");
+    let h = x_dims[1].to_usize().expect("RoPE: H must be static");
+    let d = x_dims[2].to_usize().expect("RoPE: D must be static");
+    let cos_dims = cos.dims();
+    let sin_dims = sin.dims();
+    assert_eq!(cos_dims.len(), 2, "RoPE cos must be 2-D (S, D)");
+    assert_eq!(sin_dims.len(), 2, "RoPE sin must be 2-D (S, D)");
+    assert_eq!(cos_dims[0].to_usize().unwrap(), s, "RoPE cos S mismatch");
+    assert_eq!(cos_dims[1].to_usize().unwrap(), d, "RoPE cos D mismatch");
+    assert_eq!(sin_dims[0].to_usize().unwrap(), s, "RoPE sin S mismatch");
+    assert_eq!(sin_dims[1].to_usize().unwrap(), d, "RoPE sin D mismatch");
+
+    let kern = RoPEKernel { s, h, d };
+    let cx = unsafe { &mut *x.graph_ref };
+    cx.custom_op(RoPECustom(kern), vec![x, cos, sin], (s, h, d), DType::F32)
+}

crates/luminal_cuda_lite/src/kernel/to_host.rs

@@ -192,6 +192,32 @@ impl CudaGraphOp {
             state: RefCell::new(state),
         }
     }
+
+    /// LLIR node IDs of every kernel in this CudaGraphOp, in the order
+    /// they execute inside the compiled CUDA graph. This is the
+    /// toposort `kernel_to_host` used at compile time, preserved here
+    /// so the runtime can compute live ranges that match real
+    /// execution order: each kernel in `state.kernels` was added to
+    /// the CUDA graph with `prev_graph_node` as its sole dependency,
+    /// which serializes them.
+    pub fn kernel_topo_order(&self) -> Vec<NodeIndex> {
+        self.state.borrow().kernels.iter().map(|k| k.node).collect()
+    }
+
+    /// Direct LLIR-node inputs of one kernel inside this CudaGraphOp.
+    /// Used by the runtime's live-range pass to refine intra-graph
+    /// consumer positions: a kernel's input can stop being live as
+    /// soon as that specific kernel finishes, not when the whole
+    /// CudaGraphOp finishes.
+    pub fn kernel_inputs(&self, kernel_node: NodeIndex) -> Vec<NodeIndex> {
+        self.state
+            .borrow()
+            .kernels
+            .iter()
+            .find(|k| k.node == kernel_node)
+            .map(|k| k.inputs.clone())
+            .unwrap_or_default()
+    }
 }
 
 impl std::fmt::Debug for CudaGraphOp {
@@ -316,8 +342,7 @@ impl CudaGraphOp {
             "Constant" | "Iota" => Some(0),
             "MaxReduce" | "MeanReduce" | "SumReduce" | "Cast" | "Exp" | "Exp2" | "Log2" | "Sin"
             | "Recip" | "Sigmoid" | "Softmax" | "Sqrt" => Some(1),
-            "Add" | "BatchMatMul" | "BatchMatVec" | "Embed" | "Gather" | "LessThan" | "Mod"
-            | "Mul" => Some(2),
+            "Add" | "Embed" | "Gather" | "GenericMatmul" | "LessThan" | "Mod" | "Mul" => Some(2),
             "Scatter" | "ScatterNoCopy" => Some(3),
             _ => None,
         }
@@ -814,7 +839,7 @@ 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
+    // Compute the set of FS / FE / Cuda*Elementwise nodes globally absorbed by some
     // FusionEnd in the LLIR. Used by `build_compile_units` to suppress
     // standalone marker compile units for shared FS leaves whose consumers
     // live in a different convex subgraph than the FS itself.
@@ -974,7 +999,7 @@ pub fn kernel_to_host(
                     // (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
+                    // those are interior elementwise 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>()
@@ -1139,7 +1164,7 @@ pub fn kernel_to_host(
     }
 
     // Strip fully-absorbed marker nodes (FusionStart, nested FusionEnd,
-    // FusedX) from the LLIR. Region codegen has already folded them into
+    // Cuda*Elementwise) 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

crates/luminal_cuda_lite/src/memory_analysis.rs

@@ -237,6 +237,7 @@ pub(crate) fn split_egraph_by_memory_limit(
     let mut split = splitter.split();
 
     compact_egraph_after_prune(&mut split);
+    validate_unique_loop_markers(&split);
     let stats = MemorySplitStats {
         original_enodes,
         split_enodes: split.enodes.len(),
@@ -442,6 +443,9 @@ impl<'a> StateSplitter<'a> {
                 }
             }
             "Op" => self.split_op_node(owner_class, node, label, children),
+            label if direct_loop_marker(label) => {
+                self.split_direct_loop_marker_node(owner_class, node, label.to_string(), children)
+            }
             _ => {
                 let Some((idx, child_class)) =
                     first_child_with_sort_index(self.original, &children, "IR")
@@ -479,6 +483,9 @@ impl<'a> StateSplitter<'a> {
 
         let input_states = self.split_list_class(inputs_class);
         for kind_node in kind_nodes {
+            let Some((kind_label, _)) = self.original.enodes.get(kind_node) else {
+                continue;
+            };
             let Some(kind) =
                 kind_memory_for_node(self.original, &self.sort_by_name, kind_node, self.dyn_map)
             else {
@@ -488,6 +495,33 @@ impl<'a> StateSplitter<'a> {
                 continue;
             }
             let kind_split_class = self.kind_singleton_class(kind_node);
+            if loop_op_kind(kind_label) {
+                // Loop OpKinds are structural markers. Keep the marker singleton and
+                // pick one feasible state for the data flowing through it.
+                let Some((state, input_split_class)) = input_states
+                    .iter()
+                    .filter_map(|(input_state, input_split_class)| {
+                        let state = op_memory_state(kind, input_state)?;
+                        (state.peak <= self.limit).then(|| (state, input_split_class.clone()))
+                    })
+                    .min_by_key(|(state, _)| (state.peak, state.live))
+                else {
+                    continue;
+                };
+
+                let mut split_children = children.clone();
+                split_children[0] = kind_split_class;
+                split_children[1] = input_split_class;
+                self.add_ir_state_node(
+                    owner_class,
+                    state,
+                    label.clone(),
+                    split_children,
+                    source_node,
+                );
+                continue;
+            }
+
             for (input_state, input_split_class) in &input_states {
                 let Some(state) = op_memory_state(kind, input_state) else {
                     continue;
@@ -509,6 +543,33 @@ impl<'a> StateSplitter<'a> {
         }
     }
 
+    fn split_direct_loop_marker_node(
+        &mut self,
+        owner_class: &ClassId,
+        source_node: &NodeId,
+        label: String,
+        children: Vec<ClassId>,
+    ) {
+        let Some((idx, child_class)) = first_child_with_sort_index(self.original, &children, "IR")
+        else {
+            return;
+        };
+        // LoopStart/LoopEnd identity is part of the loop scaffold, so state
+        // splitting must not clone the marker across child-state variants.
+        let Some((state, state_class)) = self
+            .split_ir_class(&child_class)
+            .into_iter()
+            .filter(|(state, _)| state.peak <= self.limit)
+            .min_by_key(|(state, _)| (state.peak, state.live))
+        else {
+            return;
+        };
+
+        let mut split_children = children;
+        split_children[idx] = state_class;
+        self.add_ir_state_node(owner_class, state, label, split_children, source_node);
+    }
+
     fn split_list_class(&mut self, class: &ClassId) -> Vec<(ListMemoryState, ClassId)> {
         if let Some(states) = self.list_memo.get(class) {
             return states.clone();
@@ -992,7 +1053,10 @@ fn choose_kind_node<'a>(egraph: &'a SerializedEGraph, kind_class: &ClassId) -> O
     };
     let is_kernel = |node: &&NodeId| -> bool {
         let label = &egraph.enodes[*node].0;
-        label.starts_with("Kernel") || label.starts_with("Fused")
+        label.starts_with("Kernel")
+            || label.starts_with("Cuda")
+            || label == "FusionStart"
+            || label == "FusionEnd"
     };
 
     kind_enodes
@@ -1079,12 +1143,94 @@ fn compact_egraph_after_prune(egraph: &mut SerializedEGraph) {
 }
 
 fn zero_local_op_kind(kind: &str) -> bool {
+    loop_op_kind(kind)
+}
+
+fn loop_op_kind(kind: &str) -> bool {
     matches!(
         kind,
         "LoopInput" | "LoopInputStatic" | "LoopOutput" | "LoopOutputSelect"
     )
 }
 
+fn direct_loop_marker(kind: &str) -> bool {
+    matches!(kind, "LoopStart" | "LoopEnd")
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct LoopMarkerKey {
+    label: String,
+    fields: Vec<String>,
+}
+
+fn validate_unique_loop_markers(egraph: &SerializedEGraph) {
+    let mut seen = FxHashMap::default();
+    for node in egraph.enodes.keys() {
+        for key in loop_marker_keys_for_node(egraph, node) {
+            if let Some(previous) = seen.insert(key.clone(), node.clone()) {
+                panic!(
+                    "CUDA memory splitter duplicated loop marker {key:?}: {previous:?} and {node:?}"
+                );
+            }
+        }
+    }
+}
+
+fn loop_marker_keys_for_node(egraph: &SerializedEGraph, node: &NodeId) -> Vec<LoopMarkerKey> {
+    let Some((label, children)) = egraph.enodes.get(node) else {
+        return Vec::new();
+    };
+    if direct_loop_marker(label) {
+        return vec![LoopMarkerKey {
+            label: label.clone(),
+            fields: field_signature(egraph, children.iter().skip(1)),
+        }];
+    }
+    if label != "Op" {
+        return Vec::new();
+    }
+    let Some(kind_class) = children.first() else {
+        return Vec::new();
+    };
+    let Some((sort, kind_nodes)) = egraph.eclasses.get(kind_class) else {
+        return Vec::new();
+    };
+    if sort != "OpKind" {
+        return Vec::new();
+    }
+
+    kind_nodes
+        .iter()
+        .filter_map(|kind_node| {
+            let (kind_label, kind_children) = egraph.enodes.get(kind_node)?;
+            loop_op_kind(kind_label).then(|| LoopMarkerKey {
+                label: kind_label.clone(),
+                fields: field_signature(egraph, kind_children.iter()),
+            })
+        })
+        .collect()
+}
+
+fn field_signature<'a>(
+    egraph: &SerializedEGraph,
+    fields: impl Iterator<Item = &'a ClassId>,
+) -> Vec<String> {
+    fields
+        .map(|class| {
+            let node_label = egraph
+                .eclasses
+                .get(class)
+                .and_then(|(_, nodes)| {
+                    nodes
+                        .iter()
+                        .find_map(|node| egraph.enodes.get(node).map(|(label, _)| label.clone()))
+                })
+                .unwrap_or_else(|| "<missing>".to_string());
+            format!("{}:{node_label}", class.as_ref())
+        })
+        .collect()
+}
+
 fn cuda_sort_map() -> FxHashMap<String, SortDef> {
     <(crate::kernel::Ops, crate::host::Ops) as luminal::op::IntoEgglogOp>::into_vec()
         .into_iter()
@@ -1104,7 +1250,7 @@ fn local_output_bytes<'a>(
 ) -> Option<Expression> {
     match sort.name.as_str() {
         name if zero_local_op_kind(name) => Some(0.into()),
-        name if name.starts_with("Fused") || name == "FusionStart" => Some(0.into()),
+        name if name.starts_with("Cuda") || name == "FusionStart" => Some(0.into()),
         "KernelConstant" => Some(4.into()),
         "KernelIota" => Some(expr_field(egraph, sort, kind_children, "range", expr_cache)? * 4),
         "KernelLessThan" => Some(n_elements_field(
@@ -1135,7 +1281,7 @@ fn local_output_bytes<'a>(
             let dtype = dtype_field(egraph, sort, kind_children, "dtype")?;
             Some(bytes_for_elements(size, dtype))
         }
-        "cublaslt" => {
+        "cublaslt" | "cublaslt_scaled" => {
             let batch = expr_field(egraph, sort, kind_children, "batch_count", expr_cache)?;
             let m = expr_field(egraph, sort, kind_children, "m", expr_cache)?;
             let n = expr_field(egraph, sort, kind_children, "n", expr_cache)?;
@@ -1213,7 +1359,7 @@ fn n_elements_field<'a>(
 fn output_bytes_rules(sort: &SortDef) -> Vec<String> {
     match sort.name.as_str() {
         name if zero_local_op_kind(name) => vec![output_bytes_rule(sort, "(MNum 0)", "zero")],
-        name if name.starts_with("Fused") || name == "FusionStart" => {
+        name if name.starts_with("Cuda") || name == "FusionStart" => {
             vec![output_bytes_rule(sort, "(MNum 0)", "zero")]
         }
         "KernelConstant" => vec![output_bytes_rule(sort, "(MNum 4)", "f32-scalar")],
@@ -1244,7 +1390,7 @@ fn output_bytes_rules(sort: &SortDef) -> Vec<String> {
             &["(= ?__cuda_elems (n_elements ?batch_shape))"],
         )],
         "KernelCast" => dtype_output_bytes_rules(sort, "size", "dtype"),
-        "cublaslt" => {
+        "cublaslt" | "cublaslt_scaled" => {
             dtype_output_bytes_rules_for_expr(sort, "(MMul (MMul ?batch_count ?m) ?n)", "d_dtype")
         }
         "GLUMoE" => vec![output_bytes_rule(
@@ -1371,7 +1517,9 @@ fn output_bytes_rule_with_facts(
 
 #[cfg(test)]
 mod tests {
-    use super::{cuda_memory_analysis_pass, estimate_graph_memory_bytes};
+    use super::{
+        cuda_memory_analysis_pass, estimate_graph_memory_bytes, loop_marker_keys_for_node,
+    };
     use luminal::{
         egglog_utils::{
             EGraphChoiceSet, SerializedEGraph, count_choice_sets_up_to, random_initial_choice,
@@ -1383,11 +1531,7 @@ mod tests {
     };
 
     fn ops() -> Vec<std::sync::Arc<Box<dyn luminal::op::EgglogOp>>> {
-        let mut ops = <(
-            crate::kernel::hlir::Ops,
-            crate::kernel::other_ops::Ops,
-            crate::host::Ops,
-        ) as IntoEgglogOp>::into_vec();
+        let mut ops = <(crate::kernel::Ops, crate::host::Ops) as IntoEgglogOp>::into_vec();
         ops.extend(<HLIROps as IntoEgglogOp>::into_vec());
         ops
     }
@@ -1399,13 +1543,13 @@ mod tests {
             .expect("cuda memory pass should parse and run")
     }
 
-    fn kernel_add(name: &str, size: usize, a: &str, b: &str) -> String {
+    fn kernel_mod(name: &str, size: &str, a: &str, b: &str) -> String {
         format!(
             r#"
             (let {name}
                 (Op
-                    (KernelAdd
-                        (ECons (MNum {size}) (ENil))
+                    (KernelMod
+                        (ECons {size} (ENil))
                         (ECons (MIter) (ENil))
                         (ECons (MIter) (ENil))
                         (ECons (MIter) (ENil))
@@ -1454,25 +1598,20 @@ mod tests {
     }
 
     #[test]
-    fn cuda_memory_late_pass_runs_on_kernel_add() {
+    fn cuda_memory_late_pass_runs_on_kernel_mod() {
         let ops = ops();
         let late_pass = cuda_memory_analysis_pass(&ops, None, &FxHashMap::default());
-        let program = r#"
+        let program = format!(
+            r#"
             (let t0 (Input 0 "" (F32)))
             (let t1 (Input 1 "" (F32)))
-            (let t2
-                (Op
-                    (KernelAdd
-                        (ECons (MNum 4) (ENil))
-                        (ECons (MIter) (ENil))
-                        (ECons (MIter) (ENil))
-                        (ECons (MIter) (ENil))
-                        (F32))
-                    (ICons t0 (ICons t1 (INil)))))
+            {}
             (let t3 (Output t2 2))
-        "#;
+        "#,
+            kernel_mod("t2", "(MNum 4)", "t0", "t1"),
+        );
 
-        run_egglog_with_late_passes(program, "t3", &ops, false, &[late_pass])
+        run_egglog_with_late_passes(&program, "t3", &ops, false, &[late_pass])
             .expect("cuda memory pass should parse and run");
     }
 
@@ -1499,6 +1638,55 @@ mod tests {
         );
     }
 
+    #[test]
+    fn cuda_memory_state_split_does_not_duplicate_loop_markers() {
+        let program = format!(
+            r#"
+            (let t0 (Input 0 "" (F32)))
+            (let t1 (Input 1 "" (F32)))
+            {}
+            {}
+            (union small big)
+            (let loop_start (LoopStart small 0 0 (MNum 2) (F32)))
+            (let loop_end (LoopEnd small 0 0 (F32)))
+            (let loop_input (Op (LoopInput 0 0 (F32)) (ICons small (ICons t0 (INil)))))
+            (let loop_output (Op (LoopOutput 0 0 (F32)) (ICons small (INil))))
+            (let loop_select (Op (LoopOutputSelect 0 0 0 (F32)) (ICons loop_output (INil))))
+            (let out_start (Output loop_start 2))
+            (let out_end (Output loop_end 3))
+            (let out_input (Output loop_input 4))
+            (let out_select (Output loop_select 5))
+            (let out_a (OutputJoin out_start out_end))
+            (let out_b (OutputJoin out_input out_select))
+            (let out (OutputJoin out_a out_b))
+            "#,
+            kernel_mod("small", "(MNum 4)", "t0", "t1"),
+            kernel_mod("big", "(MNum 8)", "t0", "t1"),
+        );
+
+        let egraph = run_memory_egraph(&program, "out", Some(1024));
+        let mut marker_counts = FxHashMap::<String, usize>::default();
+        for node in egraph.enodes.keys() {
+            for key in loop_marker_keys_for_node(&egraph, node) {
+                *marker_counts.entry(key.label).or_default() += 1;
+            }
+        }
+
+        for marker in [
+            "LoopStart",
+            "LoopEnd",
+            "LoopInput",
+            "LoopOutput",
+            "LoopOutputSelect",
+        ] {
+            assert_eq!(
+                marker_counts.get(marker).copied().unwrap_or_default(),
+                1,
+                "{marker} should not be duplicated by memory state splitting"
+            );
+        }
+    }
+
     #[test]
     fn cuda_memory_estimates_peak_for_two_live_inputs() {
         let program = format!(
@@ -1510,9 +1698,9 @@ mod tests {
             {}
             (let out (Output parent 3))
             "#,
-            kernel_add("left", 4, "t0", "t1"),
-            kernel_add("right", 4, "t0", "t1"),
-            kernel_add("parent", 4, "left", "right"),
+            kernel_mod("left", "(MNum 4)", "t0", "t1"),
+            kernel_mod("right", "(MNum 4)", "t0", "t1"),
+            kernel_mod("parent", "(MNum 4)", "left", "right"),
         );
         let egraph = run_memory_egraph(&program, "out", None);
         let mut rng = rand::rng();
@@ -1546,7 +1734,7 @@ mod tests {
                     (ICons dest (ICons indexes (ICons src (INil))))))
             (let out (Output scatter 4))
             "#,
-            kernel_add("dest", 4, "t0", "t1"),
+            kernel_mod("dest", "(MNum 4)", "t0", "t1"),
         );
         let egraph = run_memory_egraph(&program, "out", None);
         let mut rng = rand::rng();
@@ -1569,8 +1757,8 @@ mod tests {
             (union small big)
             (let out (Output small 2))
             "#,
-            kernel_add("small", 4, "t0", "t1"),
-            kernel_add("big", 32, "t0", "t1"),
+            kernel_mod("small", "(MNum 4)", "t0", "t1"),
+            kernel_mod("big", "(MNum 32)", "t0", "t1"),
         );
 
         let egraph = run_memory_egraph(&program, "out", Some(64));
@@ -1590,22 +1778,17 @@ mod tests {
         let mut dyn_map = FxHashMap::default();
         dyn_map.insert('s', 4);
         let late_pass = cuda_memory_analysis_pass(&ops, Some(16), &dyn_map);
-        let program = r#"
+        let program = format!(
+            r#"
             (let t0 (Input 0 "" (F32)))
             (let t1 (Input 1 "" (F32)))
-            (let add
-                (Op
-                    (KernelAdd
-                        (ECons (MVar "s") (ENil))
-                        (ECons (MIter) (ENil))
-                        (ECons (MIter) (ENil))
-                        (ECons (MIter) (ENil))
-                        (F32))
-                    (ICons t0 (ICons t1 (INil)))))
+            {}
             (let out (Output add 2))
-        "#;
+        "#,
+            kernel_mod("add", "(MVar \"s\")", "t0", "t1"),
+        );
 
-        let egraph = run_egglog_with_late_passes(program, "out", &ops, false, &[late_pass])
+        let egraph = run_egglog_with_late_passes(&program, "out", &ops, false, &[late_pass])
             .expect("cuda memory pass should parse and run");
         assert_eq!(count_choice_sets_up_to(&egraph, 10), 1);
 
@@ -1628,9 +1811,9 @@ mod tests {
             {}
             (let out (Output parent 3))
             "#,
-            kernel_add("left", 12, "t0", "t1"),
-            kernel_add("right", 12, "t0", "t1"),
-            kernel_add("parent", 4, "left", "right"),
+            kernel_mod("left", "(MNum 12)", "t0", "t1"),
+            kernel_mod("right", "(MNum 12)", "t0", "t1"),
+            kernel_mod("parent", "(MNum 4)", "left", "right"),
         );
 
         let egraph = run_memory_egraph(&program, "out", Some(64));
@@ -1659,11 +1842,11 @@ mod tests {
             {}
             (let out (Output parent 4))
             "#,
-            kernel_add("left_small", 4, "t0", "t1"),
-            kernel_add("left_medium", 8, "t0", "t1"),
-            kernel_add("left_big", 12, "t0", "t1"),
-            kernel_add("right_small", 4, "t0", "t1"),
-            kernel_add("parent", 4, "left_small", "right_small"),
+            kernel_mod("left_small", "(MNum 4)", "t0", "t1"),
+            kernel_mod("left_medium", "(MNum 8)", "t0", "t1"),
+            kernel_mod("left_big", "(MNum 12)", "t0", "t1"),
+            kernel_mod("right_small", "(MNum 4)", "t0", "t1"),
+            kernel_mod("parent", "(MNum 4)", "left_small", "right_small"),
         );
 
         let uncapped_start = std::time::Instant::now();

crates/luminal_cuda_lite/src/runtime.rs

@@ -80,6 +80,14 @@ struct PlannedBuffer {
     end: usize,
 }
 
+#[cfg(test)]
+#[derive(Debug, Clone)]
+pub(crate) struct NonFiniteBufferReport {
+    pub(crate) node: NodeIndex,
+    pub(crate) index: usize,
+    pub(crate) value: f32,
+}
+
 /// Per-bucket compiled state. Each bucket holds its own executable graph,
 /// explicit runtime metadata, intermediate buffers, and node mappings.
 /// Weights (hlir_buffers) are shared.
@@ -106,6 +114,9 @@ pub(crate) struct CompiledBucket {
     pub(crate) bucket_indices: FxHashMap<char, usize>,
     /// Whether HLIR pointers have been synced into this bucket's cached_buffer_ptrs
     pub(crate) hlir_synced: bool,
+    /// Test/debug mode: give every intermediate a distinct arena range so
+    /// post-execution diagnostics can inspect expired nodes without reuse noise.
+    pub(crate) preserve_intermediate_buffers_for_debug: bool,
 }
 
 impl CompiledBucket {
@@ -130,6 +141,7 @@ impl CompiledBucket {
             intermediate_buffer_dims: FxHashSet::default(),
             bucket_indices: FxHashMap::default(),
             hlir_synced: false,
+            preserve_intermediate_buffers_for_debug: false,
         }
     }
 }
@@ -203,6 +215,25 @@ impl CudaRuntime {
             .collect()
     }
 
+    /// Read per-kernel GPU elapsed times (ms), if event-record nodes were
+    /// inserted at graph build time.
+    ///
+    /// The per-kernel event recording infra from `origin/dlrm-fused-kernels`
+    /// is not ported on this branch yet — this stub returns empty so the
+    /// dlrm example's optional `LUMINAL_KERNEL_TIMING=1` path falls back to
+    /// "(no per-kernel timings available — events not recorded)".
+    pub fn read_per_kernel_timings_ms(&self) -> Vec<(&'static str, f32)> {
+        Vec::new()
+    }
+
+    /// Synchronize the runtime's CUDA stream. Use this after `execute()` if
+    /// you need GPU-side completion time (e.g. for benchmarking) — `execute`
+    /// itself no longer syncs at the end so it stays capturable by
+    /// `torch.cuda.CUDAGraph` and similar external graph-capture machinery.
+    pub fn synchronize_stream(&self) {
+        let _ = self.cuda_stream.synchronize();
+    }
+
     fn bucket_buffer(
         bucket: &CompiledBucket,
         stream: &Arc<CudaStream>,
@@ -225,10 +256,96 @@ impl CudaRuntime {
             result::memcpy_dtod_async(dst_ptr, src.ptr(), src.len(), stream.cu_stream())
                 .expect("cuMemcpyDtoDAsync failed");
         }
-        stream.synchronize().unwrap();
         dst
     }
 
+    #[cfg(test)]
+    pub(crate) fn first_nonfinite_f32_buffer_in_nodes(
+        &self,
+        nodes: impl IntoIterator<Item = NodeIndex>,
+    ) -> Option<NonFiniteBufferReport> {
+        let _ = self.cuda_stream.synchronize();
+        let bucket = self.active();
+        let mut checked = FxHashSet::default();
+
+        for node in nodes {
+            let spec_node = resolve_logical_buffer_node(
+                node,
+                &bucket.logical_buffer_bytes,
+                &bucket.output_alias_map,
+            )
+            .unwrap_or(node);
+            if !checked.insert(spec_node) {
+                continue;
+            }
+
+            let Some(spec) = bucket.buffer_specs.get(&spec_node) else {
+                continue;
+            };
+            if !matches!(spec.dtype, DType::F32) {
+                continue;
+            }
+
+            let Some(buf) = Self::resolve_runtime_buffer(
+                bucket,
+                &self.cuda_stream,
+                &self.hlir_buffers,
+                &self.external_buffers,
+                &self.external_output_buffers,
+                spec_node,
+            ) else {
+                continue;
+            };
+            if buf.is_empty() || buf.len() % std::mem::size_of::<f32>() != 0 {
+                continue;
+            }
+
+            let host_bytes = match buf.clone_dtoh(&self.cuda_stream) {
+                Ok(bytes) => bytes,
+                Err(_) => continue,
+            };
+            let values: &[f32] = bytemuck::cast_slice(&host_bytes);
+            if let Some((index, value)) = values
+                .iter()
+                .copied()
+                .enumerate()
+                .find(|(_, value)| !value.is_finite())
+            {
+                return Some(NonFiniteBufferReport {
+                    node: spec_node,
+                    index,
+                    value,
+                });
+            }
+        }
+
+        None
+    }
+
+    #[cfg(test)]
+    pub(crate) fn first_nonfinite_f32_buffer(&self) -> Option<NonFiniteBufferReport> {
+        let bucket = self.active();
+        self.first_nonfinite_f32_buffer_in_nodes(
+            bucket
+                .buffer_specs
+                .keys()
+                .copied()
+                .sorted_by_key(|node| node.index()),
+        )
+    }
+
+    #[cfg(test)]
+    pub(crate) fn preserve_intermediate_buffers_for_debug(&mut self) {
+        for bucket in &mut self.compiled_buckets {
+            bucket.preserve_intermediate_buffers_for_debug = true;
+            bucket.logical_buffer_offsets.clear();
+            bucket.logical_buffer_bytes.clear();
+            bucket.cached_buffer_ptrs.clear();
+            bucket.arena = None;
+            bucket.arena_bytes = 0;
+        }
+    }
+
     fn resolve_runtime_buffer(
         bucket: &CompiledBucket,
         stream: &Arc<CudaStream>,
@@ -287,7 +404,12 @@ impl CudaRuntime {
                         let dev = f32s.to_cuda_input(&self.cuda_stream);
                         self.hlir_buffers.insert(node, dev);
                     }
-                    safetensors::Dtype::U8 | safetensors::Dtype::BF16 | safetensors::Dtype::F16 => {
+                    safetensors::Dtype::U8
+                    | safetensors::Dtype::BF16
+                    | safetensors::Dtype::F16
+                    | safetensors::Dtype::F8_E4M3
+                    | safetensors::Dtype::F8_E5M2
+                    | safetensors::Dtype::F8_E8M0 => {
                         let bytes = tensor.data();
                         let dev = bytes.to_cuda_input(&self.cuda_stream);
                         self.hlir_buffers.insert(node, dev);
@@ -894,6 +1016,32 @@ impl CudaRuntime {
         let planned_logical_bytes = planned.iter().map(|buf| buf.bytes).sum::<usize>();
         let logical_peak = logical_interval_peak(&planned);
 
+        if bucket.preserve_intermediate_buffers_for_debug {
+            planned.sort_by_key(|buf| buf.node.index());
+            let mut arena_end = 0usize;
+            for buf in &planned {
+                let offset = align_up(arena_end, ARENA_ALIGNMENT);
+                bucket.logical_buffer_offsets.insert(buf.node, offset);
+                bucket.logical_buffer_bytes.insert(buf.node, buf.bytes);
+                arena_end = offset + align_up(buf.bytes, ARENA_ALIGNMENT);
+            }
+            bucket.arena_bytes = arena_end;
+
+            if std::env::var_os("LUMINAL_CUDA_MEMORY_DEBUG").is_some() {
+                eprintln!(
+                    "   CUDA memory plan specs={total_spec_count} used={planned_logical_count} skipped={} spec_bytes={} used_bytes={} skipped_bytes={} logical_peak={} preserved_arena={} allocations={}",
+                    total_spec_count.saturating_sub(planned_logical_count),
+                    total_spec_bytes,
+                    planned_logical_bytes,
+                    total_spec_bytes.saturating_sub(planned_logical_bytes),
+                    logical_peak,
+                    bucket.arena_bytes,
+                    bucket.logical_buffer_offsets.len(),
+                );
+            }
+            return;
+        }
+
         let mut arena_end = 0usize;
         let mut placed: Vec<(usize, usize, usize, usize)> = Vec::with_capacity(planned.len());
         let mut placement_order = planned.iter().collect_vec();
@@ -1189,7 +1337,7 @@ impl Runtime for CudaRuntime {
     }
 
     fn estimate_graph_memory<'a>(
-        egraph: &'a SerializedEGraph,
+        egraph: &'a luminal::egglog_utils::SerializedEGraph,
         choices: &luminal::egglog_utils::EGraphChoiceSet<'a>,
         dyn_map: &FxHashMap<char, usize>,
     ) -> Option<usize> {
@@ -1343,8 +1491,8 @@ impl Runtime for CudaRuntime {
         &mut self,
         llir_graph: &LLIRGraph,
         dyn_map: &FxHashMap<char, usize>,
-        _trials: usize,
-        _timeout: Option<std::time::Duration>,
+        trials: usize,
+        timeout: Option<std::time::Duration>,
     ) -> (Self::ProfileMetric, String) {
         // Clear active bucket's arena before loading new LLIR for profiling.
         if !self.compiled_buckets.is_empty() {
@@ -1352,10 +1500,18 @@ impl Runtime for CudaRuntime {
         }
         self.load_llir(llir_graph);
         self.profiling = true;
-        let start = std::time::Instant::now();
-        self.execute(dyn_map);
+        let profile_start = std::time::Instant::now();
+        let mut durations = Vec::with_capacity(trials.max(1));
+        for _ in 0..trials.max(1) {
+            let start = std::time::Instant::now();
+            self.execute(dyn_map);
+            durations.push(start.elapsed());
+            if timeout.is_some_and(|timeout| profile_start.elapsed() >= timeout) {
+                break;
+            }
+        }
         self.profiling = false;
-        let duration = start.elapsed();
+        let duration = durations.iter().sum::<std::time::Duration>() / durations.len() as u32;
 
         let total_bytes: usize = self
             .last_kernel_stats
@@ -1397,6 +1553,35 @@ impl Runtime for CudaRuntime {
                 .filter(|n| n.to_dialect::<dyn HostOp>().is_some())
                 .count()
         );
+        let display = if std::env::var_os("LUMINAL_SEARCH_OP_NAMES").is_some() {
+            let mut kernel_counts = std::collections::BTreeMap::<&'static str, usize>::new();
+            let mut host_counts = std::collections::BTreeMap::<String, usize>::new();
+            for node in llir_graph.node_weights() {
+                if let Some(kernel) = node.to_dialect::<dyn KernelOp>() {
+                    *kernel_counts.entry(kernel.kernel_name()).or_default() += 1;
+                }
+                if let Some(host) = node.to_dialect::<dyn HostOp>() {
+                    let debug = format!("{:?}", host.as_ref().as_ref());
+                    let name = debug
+                        .split([' ', '{', '('])
+                        .next()
+                        .unwrap_or("HostOp")
+                        .to_string();
+                    *host_counts.entry(name).or_default() += 1;
+                }
+            }
+            let kernel_summary = kernel_counts
+                .iter()
+                .map(|(name, count)| format!("{name}:{count}"))
+                .join(",");
+            let host_summary = host_counts
+                .iter()
+                .map(|(name, count)| format!("{name}:{count}"))
+                .join(",");
+            format!("{display} [Kernels: {kernel_summary}] [Hosts: {host_summary}]")
+        } else {
+            display
+        };
 
         (duration, display)
     }
@@ -1417,35 +1602,6 @@ impl Runtime for CudaRuntime {
             }
         }
 
-        let bucket = &mut self.compiled_buckets[self.active_bucket];
-        Self::allocate_intermediate_buffers(bucket, &self.cuda_stream, dyn_map);
-        // Cache HLIR input pointers
-        if !self.changed_hlir.is_empty() || !bucket.hlir_synced {
-            let hlir_nodes: Vec<NodeIndex> = if !bucket.hlir_synced {
-                // First time this bucket is active since HLIR changed — sync all
-                self.hlir_buffers.keys().copied().collect()
-            } else {
-                self.changed_hlir.iter().copied().collect()
-            };
-            for hlir_node in hlir_nodes {
-                let Some(&llir_node) = bucket.hlir_to_llir.get(&hlir_node) else {
-                    continue;
-                };
-                let Some(input) = self.hlir_buffers.get(&hlir_node) else {
-                    continue;
-                };
-                let ptr = match input {
-                    CudaInput::Buffer(buf) => buf.device_ptr(&self.cuda_stream).0,
-                    CudaInput::Ptr(p) => *p,
-                };
-                bucket.cached_buffer_ptrs.insert(llir_node, ptr);
-            }
-            bucket.hlir_synced = true;
-            // Only clear changed_hlir if single bucket (multi-bucket: others may need it)
-            if self.compiled_buckets.len() == 1 {
-                self.changed_hlir.clear();
-            }
-        }
         // Ensure all CUDA graphs are built (handles first execute and any missing graphs)
         self.prebuild_graphs(dyn_map);
 
@@ -1522,9 +1678,31 @@ impl Runtime for CudaRuntime {
                         exec_op.internal.stats_name().unwrap_or("unknown")
                     );
                 });
+
+            #[cfg(test)]
+            if std::env::var_os("LUMINAL_CUDA_CHECK_NONFINITE_INTERNAL").is_some() {
+                let mut produced_nodes = exec_op.internal.extra_buffer_nodes();
+                produced_nodes.push(exec_op.output);
+                if let Some(report) = self.first_nonfinite_f32_buffer_in_nodes(produced_nodes) {
+                    panic!(
+                        "CUDA execute produced non-finite buffer after {:?}: node={} index={} value={}",
+                        exec_op.internal.stats_name().unwrap_or("unknown"),
+                        report.node.index(),
+                        report.index,
+                        report.value
+                    );
+                }
+            }
+        }
+        // Sync only when profiling (kernel search timing needs an accurate
+        // total). In the regular execute path, dropping the sync lets the
+        // call be captured by `torch.cuda.CUDAGraph` (or any external graph
+        // capture) — PyTorch syncs on tensor reads, so correctness is
+        // preserved. The CPU-side `last_total_time_us` becomes a dispatch-
+        // time measurement in that case.
+        if self.profiling {
+            self.cuda_stream.synchronize().unwrap();
         }
-        // Single sync at end - CUDA stream ordering guarantees sequential execution
-        self.cuda_stream.synchronize().unwrap();
         self.last_total_time_us = total_start.elapsed().as_secs_f64() * 1_000_000.0;
 
         // Populate last_kernel_stats from HostOps that report stats
@@ -1565,9 +1743,22 @@ impl Runtime for CudaRuntime {
 
         let to_consume: Vec<NodeIndex> = self
             .hlir_buffers
-            .keys()
-            .filter(|hlir_node| !inputs_with_outputs.contains(hlir_node))
-            .copied()
+            .iter()
+            // Don't consume external device pointers — they're non-owning
+            // views over caller-provided memory and they represent stable
+            // input slots that the wrapper *may* skip re-registering on a
+            // hot iter if the underlying pointer is unchanged. Removing
+            // them here would force the wrapper to re-register every
+            // iter even when nothing changed (which the prior behavior
+            // assumed). Internal `CudaInput::Buffer` entries — e.g.
+            // weights loaded via `set_data_bytes` and one-shot CPU input
+            // copies — still get consumed when they're not preserved by
+            // the bucket.
+            .filter(|(hlir_node, input)| {
+                !inputs_with_outputs.contains(hlir_node)
+                    && !matches!(input, CudaInput::Ptr(_))
+            })
+            .map(|(n, _)| *n)
             .collect();
 
         for hlir_node in to_consume {
@@ -1657,8 +1848,8 @@ impl CudaRuntime {
                 //
                 // The default assumption is "yes" for ordinary kernel ops
                 // (Conv outputs, matmul outputs, etc). FusionStart and
-                // Fused* are the exceptions — they're synthetic markers
-                // that the fusion rewrites add inside a region; the
+                // Cuda*Elementwise are the exceptions — they're synthetic
+                // nodes that the fusion rewrites add inside a region; the
                 // megakernel computes them in registers and never writes
                 // to memory, so allocating a buffer would just be waste.
                 //
@@ -1673,12 +1864,12 @@ impl CudaRuntime {
                 // an unrelated downstream op that lives in another region.
                 //
                 // Safe over-approximation: if the node is a FusionStart /
-                // Fused* and *any* of its consumers is a FusionStart
+                // Cuda*Elementwise and *any* of its consumers is a FusionStart
                 // (which can only happen when that consumer is the leaf
                 // of a different region) or a non-marker op (e.g. an
                 // unfused Add/Mul reading the value directly), allocate a
                 // buffer so cross-region reads have somewhere to land.
-                let is_marker = kernel_name == "FusionStart" || kernel_name.starts_with("Fused");
+                let is_marker = kernel_name == "FusionStart" || kernel_name.starts_with("Cuda");
                 let has_external_consumer = is_marker
                     && llir_graph
                         .neighbors_directed(node, Direction::Outgoing)

crates/luminal_cuda_lite/src/tests/conv2d_rewrite.rs

@@ -0,0 +1,482 @@
+use luminal::{
+    egglog_utils::{
+        NodeId, SerializedEGraph, egglog_to_llir, random_initial_choice, validate_choice_set,
+    },
+    prelude::*,
+};
+use rand::{SeedableRng, rngs::StdRng};
+
+use crate::{kernel::KernelOp, runtime::CudaRuntime};
+
+use super::utilities::{assert_close, get_cuda_stream};
+
+fn conv2d_bias_hlir(
+    x: GraphTensor,
+    weight: GraphTensor,
+    bias: GraphTensor,
+    kernel_h: usize,
+    kernel_w: usize,
+) -> GraphTensor {
+    let unfolded = x.unfold(
+        vec![1usize, kernel_h, kernel_w],
+        vec![1usize, 1, 1],
+        vec![1usize, 1, 1],
+    );
+    let output_spatial_dims = unfolded.dims()[1..3].to_vec();
+
+    let mut patches = unfolded.squeeze(3).permute(&[1, 2, 0, 3, 4]);
+    while patches.dims().len() > 3 {
+        let last = patches.dims().len();
+        patches = patches.merge_dims(last - 2, last - 1);
+    }
+    let patches = patches.merge_dims(0, 1);
+
+    let out = patches.matmul(weight.t());
+    let out = out
+        .split_dims(0, output_spatial_dims[1])
+        .permute(&[2, 0, 1]);
+    let out_dims = out.dims();
+    out + bias.expand_dim(1, out_dims[1]).expand_dim(2, out_dims[2])
+}
+
+fn build_conv_graph() -> (Graph, GraphTensor, GraphTensor, GraphTensor, GraphTensor) {
+    let mut cx = Graph::new();
+    let x = cx.tensor((2usize, 5usize, 6usize));
+    let weight = cx.tensor((3usize, 2usize * 3 * 2));
+    let bias = cx.tensor(3usize);
+    let out = conv2d_bias_hlir(x, weight, bias, 3, 2).output();
+    (cx, x, weight, bias, out)
+}
+
+fn conv2d_bias_padded_hlir(
+    x: GraphTensor,
+    weight: GraphTensor,
+    bias: GraphTensor,
+    kernel: usize,
+    padding: usize,
+) -> GraphTensor {
+    let zero = Expression::from(0);
+    let pad = Expression::from(padding);
+    let padded = x.pad(vec![(zero, zero), (pad, pad), (pad, pad)], 0.0);
+    conv2d_bias_hlir(padded, weight, bias, kernel, kernel)
+}
+
+fn build_padded_conv_graph() -> (Graph, GraphTensor, GraphTensor, GraphTensor, GraphTensor) {
+    let mut cx = Graph::new();
+    let x = cx.tensor((2usize, 4usize, 5usize));
+    let weight = cx.tensor((3usize, 2usize * 3 * 3));
+    let bias = cx.tensor(3usize);
+    let out = conv2d_bias_padded_hlir(x, weight, bias, 3, 1).output();
+    (cx, x, weight, bias, out)
+}
+
+fn nearest_upsample_2x_hlir(x: GraphTensor) -> GraphTensor {
+    let stage1 = x.expand_dim(2, 2usize).merge_dims(1, 2);
+    stage1.expand_dim(3, 2usize).merge_dims(2, 3)
+}
+
+fn build_upsample_conv_graph() -> (Graph, GraphTensor, GraphTensor, GraphTensor, GraphTensor) {
+    let mut cx = Graph::new();
+    let x = cx.tensor((2usize, 3usize, 4usize));
+    let weight = cx.tensor((3usize, 2usize * 3 * 3));
+    let bias = cx.tensor(3usize);
+    let up = nearest_upsample_2x_hlir(x);
+    let out = conv2d_bias_padded_hlir(up, weight, bias, 3, 1).output();
+    (cx, x, weight, bias, out)
+}
+
+fn conv1x1_bias_hlir(x: GraphTensor, weight: GraphTensor, bias: GraphTensor) -> GraphTensor {
+    let dims = x.dims();
+    let h = dims[1];
+    let w = dims[2];
+    let xt = x.permute(&[1, 2, 0]).merge_dims(0, 1);
+    let out = xt.matmul(weight.t());
+    let out = out.split_dims(0, w).permute(&[2, 0, 1]);
+    out + bias.expand_dim(1, h).expand_dim(2, w)
+}
+
+fn build_conv1x1_graph() -> (Graph, GraphTensor, GraphTensor, GraphTensor, GraphTensor) {
+    let mut cx = Graph::new();
+    let x = cx.tensor((2usize, 4usize, 5usize));
+    let weight = cx.tensor((3usize, 2usize));
+    let bias = cx.tensor(3usize);
+    let out = conv1x1_bias_hlir(x, weight, bias).output();
+    (cx, x, weight, bias, out)
+}
+
+fn conv2d_matmul_without_conv_output_shape(
+    x: GraphTensor,
+    weight: GraphTensor,
+    bias: GraphTensor,
+    kernel_h: usize,
+    kernel_w: usize,
+) -> GraphTensor {
+    let unfolded = x.unfold(
+        vec![1usize, kernel_h, kernel_w],
+        vec![1usize, 1, 1],
+        vec![1usize, 1, 1],
+    );
+
+    let mut patches = unfolded.squeeze(3).permute(&[1, 2, 0, 3, 4]);
+    while patches.dims().len() > 3 {
+        let last = patches.dims().len();
+        patches = patches.merge_dims(last - 2, last - 1);
+    }
+    let patches = patches.merge_dims(0, 1);
+
+    let out = patches.matmul(weight.t());
+    let out_dims = out.dims();
+    out + bias.expand_dim(0, out_dims[0])
+}
+
+#[test]
+fn generic_conv2d_rewrite_matches_unfold_matmul_bias() {
+    let (mut cx, _, _, _, _) = build_conv_graph();
+    cx.build_search_space::<CudaRuntime>();
+    let egraph = cx.egraph().expect("search space should have an e-graph");
+
+    assert!(
+        !op_ir_nodes(egraph, "KernelConv2D").is_empty(),
+        "expected generic conv2d rewrite candidate"
+    );
+    assert!(
+        op_ir_nodes(egraph, "Add").is_empty(),
+        "generic conv2d cleanup should prune the final bias Add fallback"
+    );
+}
+
+#[test]
+fn generic_conv2d_rewrite_matches_conv1x1_matmul_bias() {
+    let (mut cx, _, _, _, _) = build_conv1x1_graph();
+    cx.build_search_space::<CudaRuntime>();
+    let egraph = cx.egraph().expect("search space should have an e-graph");
+
+    assert!(
+        !op_ir_nodes(egraph, "KernelConv2D").is_empty(),
+        "expected generic conv2d rewrite candidate for 1x1 conv"
+    );
+}
+
+#[test]
+fn generic_conv2d_rewrite_requires_conv_output_shape() {
+    let mut cx = Graph::new();
+    let x = cx.tensor((2usize, 5usize, 6usize));
+    let weight = cx.tensor((3usize, 2usize * 3 * 2));
+    let bias = cx.tensor(3usize);
+    conv2d_matmul_without_conv_output_shape(x, weight, bias, 3, 2).output();
+
+    cx.build_search_space::<CudaRuntime>();
+    let egraph = cx.egraph().expect("search space should have an e-graph");
+
+    assert!(
+        op_ir_nodes(egraph, "KernelConv2D").is_empty(),
+        "matmul+bias without [C_out,H_out,W_out] conv output shape should not match KernelConv2D"
+    );
+}
+
+#[test]
+fn generic_conv2d_candidate_executes_unfold_matmul_bias() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    let (mut cx, x, weight, bias, out) = build_conv_graph();
+    cx.build_search_space::<CudaRuntime>();
+    let llir = extract_forced_kernel_llir(&mut cx, "GenericConv2D");
+
+    let input: Vec<f32> = (0..2 * 5 * 6).map(|i| i as f32 * 0.03 - 0.4).collect();
+    let weights: Vec<f32> = (0..3 * 2 * 3 * 2)
+        .map(|i| (i as f32 % 11.0) * 0.04 - 0.2)
+        .collect();
+    let biases = vec![0.25_f32, -0.15, 0.05];
+    let expected = reference_conv2d(
+        &input,
+        &weights,
+        &biases,
+        ConvCase {
+            c_in: 2,
+            h: 5,
+            w: 6,
+            c_out: 3,
+            kh: 3,
+            kw: 2,
+            padding_h: 0,
+            padding_w: 0,
+        },
+    );
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.load_llir(&llir);
+    rt.set_data(x, input);
+    rt.set_data(weight, weights);
+    rt.set_data(bias, biases);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
+}
+
+#[test]
+fn generic_conv2d_candidate_executes_conv1x1_matmul_bias() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    let (mut cx, x, weight, bias, out) = build_conv1x1_graph();
+    cx.build_search_space::<CudaRuntime>();
+    let llir = extract_forced_kernel_llir(&mut cx, "GenericConv2D");
+
+    let input: Vec<f32> = (0..2 * 4 * 5).map(|i| i as f32 * 0.07 - 1.0).collect();
+    let weights: Vec<f32> = (0..3 * 2).map(|i| (i as f32 % 5.0) * 0.11 - 0.2).collect();
+    let biases = vec![0.2_f32, -0.1, 0.4];
+    let expected = reference_conv2d(
+        &input,
+        &weights,
+        &biases,
+        ConvCase {
+            c_in: 2,
+            h: 4,
+            w: 5,
+            c_out: 3,
+            kh: 1,
+            kw: 1,
+            padding_h: 0,
+            padding_w: 0,
+        },
+    );
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.load_llir(&llir);
+    rt.set_data(x, input);
+    rt.set_data(weight, weights);
+    rt.set_data(bias, biases);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
+}
+
+#[test]
+fn generic_conv2d_candidate_executes_padded_unfold_matmul_bias() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    let (mut cx, x, weight, bias, out) = build_padded_conv_graph();
+    cx.build_search_space::<CudaRuntime>();
+    let llir = extract_forced_kernel_llir(&mut cx, "GenericConv2D");
+
+    let input: Vec<f32> = (0..2 * 4 * 5).map(|i| i as f32 * 0.05 - 0.5).collect();
+    let weights: Vec<f32> = (0..3 * 2 * 3 * 3)
+        .map(|i| (i as f32 % 13.0) * 0.03 - 0.17)
+        .collect();
+    let biases = vec![0.15_f32, -0.25, 0.35];
+    let expected = reference_conv2d(
+        &input,
+        &weights,
+        &biases,
+        ConvCase {
+            c_in: 2,
+            h: 4,
+            w: 5,
+            c_out: 3,
+            kh: 3,
+            kw: 3,
+            padding_h: 1,
+            padding_w: 1,
+        },
+    );
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.load_llir(&llir);
+    rt.set_data(x, input);
+    rt.set_data(weight, weights);
+    rt.set_data(bias, biases);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
+}
+
+#[test]
+fn generic_conv2d_candidate_executes_upsample_view_input() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    let (mut cx, x, weight, bias, out) = build_upsample_conv_graph();
+    cx.build_search_space::<CudaRuntime>();
+    let llir = extract_forced_kernel_llir(&mut cx, "GenericConv2D");
+
+    let input: Vec<f32> = (0..2 * 3 * 4).map(|i| i as f32 * 0.09 - 0.8).collect();
+    let weights: Vec<f32> = (0..3 * 2 * 3 * 3)
+        .map(|i| (i as f32 % 17.0) * 0.025 - 0.2)
+        .collect();
+    let biases = vec![0.05_f32, -0.1, 0.2];
+    let upsampled = reference_nearest_upsample_2x(&input, 2, 3, 4);
+    let expected = reference_conv2d(
+        &upsampled,
+        &weights,
+        &biases,
+        ConvCase {
+            c_in: 2,
+            h: 6,
+            w: 8,
+            c_out: 3,
+            kh: 3,
+            kw: 3,
+            padding_h: 1,
+            padding_w: 1,
+        },
+    );
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.load_llir(&llir);
+    rt.set_data(x, input);
+    rt.set_data(weight, weights);
+    rt.set_data(bias, biases);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
+}
+
+struct ConvCase {
+    c_in: usize,
+    h: usize,
+    w: usize,
+    c_out: usize,
+    kh: usize,
+    kw: usize,
+    padding_h: usize,
+    padding_w: usize,
+}
+
+fn reference_nearest_upsample_2x(input: &[f32], c: usize, h: usize, w: usize) -> Vec<f32> {
+    let mut out = vec![0.0_f32; c * h * 2 * w * 2];
+    for ci in 0..c {
+        for y in 0..h {
+            for x in 0..w {
+                let value = input[ci * h * w + y * w + x];
+                for dy in 0..2 {
+                    for dx in 0..2 {
+                        let oy = y * 2 + dy;
+                        let ox = x * 2 + dx;
+                        out[ci * h * 2 * w * 2 + oy * w * 2 + ox] = value;
+                    }
+                }
+            }
+        }
+    }
+    out
+}
+
+fn reference_conv2d(input: &[f32], weight: &[f32], bias: &[f32], case: ConvCase) -> Vec<f32> {
+    let ConvCase {
+        c_in,
+        h,
+        w,
+        c_out,
+        kh,
+        kw,
+        padding_h,
+        padding_w,
+    } = case;
+    let h_out = h + 2 * padding_h - kh + 1;
+    let w_out = w + 2 * padding_w - kw + 1;
+    let mut out = vec![0.0; c_out * h_out * w_out];
+    for co in 0..c_out {
+        for oh in 0..h_out {
+            for ow in 0..w_out {
+                let mut acc = bias[co];
+                for ci in 0..c_in {
+                    for r in 0..kh {
+                        for s in 0..kw {
+                            let Some(ih) = (oh + r).checked_sub(padding_h) else {
+                                continue;
+                            };
+                            let Some(iw) = (ow + s).checked_sub(padding_w) else {
+                                continue;
+                            };
+                            if ih >= h || iw >= w {
+                                continue;
+                            }
+                            let input_idx = ci * h * w + ih * w + iw;
+                            let weight_idx = co * c_in * kh * kw + (ci * kh + r) * kw + s;
+                            acc += input[input_idx] * weight[weight_idx];
+                        }
+                    }
+                }
+                out[co * h_out * w_out + oh * w_out + ow] = acc;
+            }
+        }
+    }
+    out
+}
+
+fn extract_forced_kernel_llir(cx: &mut Graph, kernel_name: &str) -> LLIRGraph {
+    let egraph = cx.egraph().expect("search space should have an e-graph");
+    let ops = cx
+        .egglog_ops()
+        .expect("search space should have registered egglog ops");
+    let kernel_nodes = op_ir_nodes(egraph, "KernelConv2D");
+    assert!(
+        !kernel_nodes.is_empty(),
+        "expected at least one {kernel_name} candidate"
+    );
+
+    for (idx, kernel_node) in kernel_nodes.iter().enumerate() {
+        let mut rng = StdRng::seed_from_u64(0xC0_2D00 + idx as u64);
+        let mut choices = random_initial_choice(egraph, &mut rng);
+        let kernel_class = &egraph.node_to_class[*kernel_node];
+        choices.insert(kernel_class, kernel_node);
+
+        if validate_choice_set(egraph, &choices, ops).is_err() {
+            continue;
+        }
+
+        let mut list_cache = FxHashMap::default();
+        let mut expr_cache = FxHashMap::default();
+        let llir = egglog_to_llir(
+            egraph,
+            choices,
+            ops,
+            &cx.custom_ops,
+            &mut list_cache,
+            &mut expr_cache,
+            None,
+        );
+        if llir_kernel_names(&llir).contains(&kernel_name) {
+            return llir;
+        }
+    }
+
+    panic!("could not extract a valid {kernel_name} candidate");
+}
+
+fn llir_kernel_names(llir: &LLIRGraph) -> Vec<&'static str> {
+    llir.node_indices()
+        .filter_map(|node| {
+            llir[node]
+                .to_dialect::<dyn KernelOp>()
+                .map(|kernel| kernel.kernel_name())
+        })
+        .collect()
+}
+
+fn op_ir_nodes<'a>(egraph: &'a SerializedEGraph, kind_label: &str) -> Vec<&'a NodeId> {
+    let op_kind_classes = egraph
+        .enodes
+        .iter()
+        .filter(|(_, (label, _))| label == kind_label)
+        .map(|(node, _)| egraph.node_to_class[node].clone())
+        .collect::<Vec<_>>();
+
+    egraph
+        .enodes
+        .iter()
+        .filter_map(|(node, (label, children))| {
+            (label == "Op"
+                && children
+                    .first()
+                    .is_some_and(|kind| op_kind_classes.contains(kind)))
+            .then_some(node)
+        })
+        .collect()
+}

crates/luminal_cuda_lite/src/tests/cublaslt_rewrite_tests.rs

@@ -1,7 +1,8 @@
 use luminal::{
     dtype::DType,
     egglog_utils::{
-        NodeId, SerializedEGraph, egglog_to_llir, random_initial_choice, validate_choice_set,
+        ClassId, NodeId, SerializedEGraph, egglog_to_llir, random_initial_choice,
+        validate_choice_set,
     },
     prelude::*,
 };
@@ -11,7 +12,8 @@ use crate::{
     host::{
         CublasLtMatrixOrders, CublasLtScaleValues, CublasLtTransposeOps, CublasLtTypeTuple, HostOp,
         cublaslt_c_d_layouts_match, cublaslt_epilogue, cublaslt_matrix_orders,
-        cublaslt_scale_values, cublaslt_transpose_ops, cublaslt_type_tuple,
+        cublaslt_scale_values, cublaslt_tensor_scale_inputs, cublaslt_transpose_ops,
+        cublaslt_type_tuple,
     },
     runtime::CudaRuntime,
 };
@@ -443,6 +445,54 @@ fn cublaslt_rewrites_cover_batched_row_order_layout_pairs() {
     }
 }
 
+#[test]
+fn cublaslt_rewrites_cover_flux2_qk_transposed_matmul() {
+    let mut cx = Graph::new();
+    let q = cx.tensor((8usize, 4usize));
+    let k = cx.tensor((8usize, 4usize));
+    let _out = q.matmul(k.t()).output();
+
+    assert_cublaslt_rewrite(cx, "flux2 q @ k.t()", |llir| {
+        cublaslt_matrix_order_tuples(llir).contains(&("ROW", "COL", "ROW", "ROW"))
+            || cublaslt_matrix_order_tuples(llir).contains(&("COL", "COL", "COL", "COL"))
+    });
+}
+
+#[test]
+fn cublaslt_rewrites_cover_flux2_linear_bias_epilogue() {
+    let mut cx = Graph::new();
+    let x = cx.tensor((8usize, 4usize));
+    let weight = cx.tensor((6usize, 4usize));
+    let bias = cx.tensor(6usize);
+    let _out = (x.matmul(weight.t()) + bias.expand_dim(0, 8usize)).output();
+
+    assert_cublaslt_epilogue_rewrite(
+        cx,
+        "flux2 x @ weight.t() + bias",
+        "BIAS",
+        Some(("COL", "COL", "COL", "COL")),
+    );
+}
+
+#[test]
+fn cublaslt_cleanup_prunes_flux2_broadcast_mul_fallback() {
+    let mut cx = Graph::new();
+    let q = cx.tensor((8usize, 4usize));
+    let k = cx.tensor((8usize, 4usize));
+    let _out = q.matmul(k.t()).output();
+
+    cx.build_search_space::<CudaRuntime>();
+    let egraph = cx.egraph().expect("search space should have an e-graph");
+    assert!(
+        !cublaslt_ir_nodes(egraph).is_empty(),
+        "Flux2 q @ k.t() should have at least one cuBLASLt candidate"
+    );
+    assert!(
+        op_ir_nodes(egraph, "Mul").is_empty(),
+        "cuBLASLt cleanup should prune the broadcast Mul fallback once a cuBLASLt candidate exists"
+    );
+}
+
 #[test]
 fn cublaslt_rewrites_keep_c_and_d_layouts_equal_initially() {
     for case in LAYOUT_CASES {
@@ -900,6 +950,196 @@ fn cublaslt_fp8_e4m3_beta_candidate_executes_2d_matmul_plus_f32_c() {
     assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
 }
 
+#[test]
+#[ignore = "expensive CUDA FP8 rewrite sweep; run with cargo test -p luminal_cuda_lite -- --ignored"]
+fn cublaslt_fp8_scaled_candidate_executes_2d_matmul_f32_output() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+    if !gpu_supports_cublaslt_fp8_launch(DType::F8E4M3) {
+        return;
+    }
+
+    let (m, n, k) = (16, 16, 16);
+    let mut cx = Graph::new();
+    let a = cx.tensor((m, k));
+    let a_scale = cx.tensor(());
+    let b_scale = cx.tensor(());
+    let b_input = cx.tensor((n, k)).as_dtype(DType::F8E4M3);
+    let b = b_input.t();
+    let scaled_a = (a / a_scale.expand_rhs((m, k))).cast(DType::F8E4M3);
+    let out =
+        (scaled_a.matmul(b).cast(DType::F32) * (a_scale * b_scale).expand_rhs((m, n))).output();
+    let expected_tuple = (
+        DType::F8E4M3,
+        DType::F8E4M3,
+        DType::F32,
+        DType::F32,
+        "32F",
+        DType::F32,
+    );
+    let llir = extract_forced_cublaslt_llir_where(&mut cx, "functional scaled fp8", |llir| {
+        cublaslt_type_tuples(llir).contains(&expected_tuple)
+            && cublaslt_tensor_scale_input_tuples(llir).contains(&(true, true))
+            && cublaslt_transpose_op_tuples(llir).contains(&("T", "N"))
+            && cublaslt_matrix_order_tuples(llir).contains(&("COL", "COL", "COL", "COL"))
+    });
+
+    let input_scale = 0.25f32;
+    let weight_scale = 2.0f32;
+    let (a_fp8_bytes, a_values) = fp8_exact_bytes(DType::F8E4M3, m * k, 7);
+    let a_data = a_values
+        .iter()
+        .map(|value| value * input_scale)
+        .collect::<Vec<_>>();
+    let (b_bytes, b_storage_values) = fp8_exact_bytes(DType::F8E4M3, k * n, 9);
+    let b_values = logical_b_from_column_major_storage(&b_storage_values, n, k);
+    let mut expected = reference_matmul_2d(&a_values, &b_values, m, n, k);
+    for value in &mut expected {
+        *value *= input_scale * weight_scale;
+    }
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.load_llir(&llir);
+    rt.set_data(a, a_data);
+    rt.set_data(a_scale, vec![input_scale]);
+    rt.set_data(b_scale, vec![weight_scale]);
+    rt.set_data(b_input, b_bytes);
+    rt.execute(&cx.dyn_map);
+
+    // Keep the raw bytes live in the test construction: a_data was chosen so
+    // the explicit scaled cast quantizes to these exact FP8 values.
+    assert_eq!(a_fp8_bytes.len(), m * k);
+    assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
+}
+
+#[test]
+fn cublaslt_fp8_scaled_candidate_reaches_fused_output_scale_consumer() {
+    let (m, n, k) = (16, 16, 16);
+    let mut cx = Graph::new();
+    let a = cx.tensor((m, k));
+    let a_scale = cx.tensor(());
+    let b_scale = cx.tensor(());
+    let b_input = cx.tensor((n, k)).as_dtype(DType::F8E4M3);
+    let b = b_input.t();
+    let side = cx.tensor((m, n));
+    let scaled_a = (a / a_scale.expand_rhs((m, k))).cast(DType::F8E4M3);
+    let scaled_out = scaled_a.matmul(b).cast(DType::F32) * (a_scale * b_scale).expand_rhs((m, n));
+    (scaled_out * side).output();
+
+    cx.build_search_space::<CudaRuntime>();
+    let egraph = cx.egraph().expect("search space should have an e-graph");
+
+    assert!(
+        dataflow_reachable_cublaslt_scaled_count(egraph) > 0,
+        "scaled cuBLASLt must remain reachable when fusion growth consumes the output-scale multiply internally"
+    );
+    assert_eq!(
+        dataflow_reachable_cublaslt_raw_fp8_count(egraph),
+        0,
+        "raw FP8 cuBLASLt must be deleted when a scaled equivalent covers the fused output-scale consumer"
+    );
+}
+
+#[test]
+fn cublaslt_fp8_scaled_candidates_reach_fused_mlp_consumer() {
+    let (m, n, k) = (16, 32, 16);
+    let mut cx = Graph::new();
+    let a = cx.tensor((m, k));
+    let gate_input_scale = cx.tensor(());
+    let gate_weight_scale = cx.tensor(());
+    let up_input_scale = cx.tensor(());
+    let up_weight_scale = cx.tensor(());
+    let gate_weight = cx.tensor((n, k)).as_dtype(DType::F8E4M3);
+    let up_weight = cx.tensor((n, k)).as_dtype(DType::F8E4M3);
+
+    let scaled_gate_a = (a / gate_input_scale.expand_rhs((m, k))).cast(DType::F8E4M3);
+    let gate = scaled_gate_a.matmul(gate_weight.t()).cast(DType::F32)
+        * (gate_input_scale * gate_weight_scale).expand_rhs((m, n));
+    let scaled_up_a = (a / up_input_scale.expand_rhs((m, k))).cast(DType::F8E4M3);
+    let up = scaled_up_a.matmul(up_weight.t()).cast(DType::F32)
+        * (up_input_scale * up_weight_scale).expand_rhs((m, n));
+    (gate.swish() * up).output();
+
+    cx.build_search_space::<CudaRuntime>();
+    let egraph = cx.egraph().expect("search space should have an e-graph");
+
+    assert!(
+        dataflow_reachable_cublaslt_scaled_count(egraph) >= 2,
+        "scaled cuBLASLt candidates must remain reachable through fused MLP gate/up consumers"
+    );
+    assert_eq!(
+        dataflow_reachable_cublaslt_raw_fp8_count(egraph),
+        0,
+        "raw FP8 cuBLASLt must be deleted when a scaled equivalent covers the fused MLP consumer"
+    );
+}
+
+#[test]
+#[ignore = "expensive CUDA FP8 rewrite sweep; run with cargo test -p luminal_cuda_lite -- --ignored"]
+fn cublaslt_fp8_scaled_candidate_executes_batched_matmul_f32_output() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+    if !gpu_supports_cublaslt_fp8_launch(DType::F8E4M3) {
+        return;
+    }
+
+    let (batch, m, n, k) = (2, 16, 16, 16);
+    let mut cx = Graph::new();
+    let a = cx.tensor((batch, m, k));
+    let a_scale = cx.tensor(());
+    let b_scale = cx.tensor(());
+    let b_input = cx.tensor((batch, n, k)).as_dtype(DType::F8E4M3);
+    let b = b_input.transpose(1, 2);
+    let scaled_a = (a / a_scale.expand_rhs((batch, m, k))).cast(DType::F8E4M3);
+    let lhs = scaled_a.expand_dim(2, n);
+    let rhs = b.permute((0, 2, 1)).expand_dim(1, m);
+    let mul = unchecked_mul_same_shape(lhs, rhs, DType::F8E4M3);
+    let matmul = mul.sum(3).cast(DType::F32);
+    let out = (matmul * (a_scale * b_scale).expand_rhs((batch, m, n))).output();
+    let expected_tuple = (
+        DType::F8E4M3,
+        DType::F8E4M3,
+        DType::F32,
+        DType::F32,
+        "32F",
+        DType::F32,
+    );
+    let llir =
+        extract_forced_cublaslt_llir_where(&mut cx, "functional scaled batched fp8", |llir| {
+            cublaslt_type_tuples(llir).contains(&expected_tuple)
+                && cublaslt_tensor_scale_input_tuples(llir).contains(&(true, true))
+                && cublaslt_transpose_op_tuples(llir).contains(&("T", "N"))
+                && cublaslt_matrix_order_tuples(llir).contains(&("COL", "COL", "COL", "COL"))
+        });
+
+    let input_scale = 0.5f32;
+    let weight_scale = 1.5f32;
+    let (a_fp8_bytes, a_values) = fp8_exact_bytes(DType::F8E4M3, batch * m * k, 11);
+    let a_data = a_values
+        .iter()
+        .map(|value| value * input_scale)
+        .collect::<Vec<_>>();
+    let (b_bytes, b_storage_values) = fp8_exact_bytes(DType::F8E4M3, batch * k * n, 13);
+    let b_values = logical_b_from_batched_column_major_storage(&b_storage_values, batch, n, k);
+    let mut expected = reference_matmul_batched(&a_values, &b_values, batch, m, n, k);
+    for value in &mut expected {
+        *value *= input_scale * weight_scale;
+    }
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.load_llir(&llir);
+    rt.set_data(a, a_data);
+    rt.set_data(a_scale, vec![input_scale]);
+    rt.set_data(b_scale, vec![weight_scale]);
+    rt.set_data(b_input, b_bytes);
+    rt.execute(&cx.dyn_map);
+
+    assert_eq!(a_fp8_bytes.len(), batch * m * k);
+    assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
+}
+
 fn cublaslt_fp8_candidate_executes_2d_matmul_f32_output(a_dtype: DType, b_dtype: DType) {
     let Some(stream) = get_cuda_stream() else {
         return;
@@ -2168,6 +2408,85 @@ fn cublaslt_row_order_candidate_executes_2d_layout_pairs() {
     }
 }
 
+#[test]
+#[ignore = "large row-order CUDA functional repro for llama lm_head shape"]
+fn cublaslt_row_order_candidate_executes_large_lm_head_like_projection() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    let (m, n, k) = (1, 128_256, 64);
+    let mut cx = Graph::new();
+    let a = cx.tensor((m, k));
+    let b_input = cx.tensor((n, k));
+    let b = b_input.t();
+    let out = a.matmul(b).output();
+    let expected_orders = ("ROW", "COL", "ROW", "ROW");
+    let llir = extract_forced_cublaslt_llir_where(&mut cx, "lm_head-like row-order", |llir| {
+        cublaslt_matrix_order_tuples(llir).contains(&expected_orders)
+            && cublaslt_scale_value_tuples(llir).contains(&(1.0, 0.0))
+    });
+
+    let a_data = random_f32_vec(m * k, 0x1A11_A000, -0.5, 0.5);
+    let b_data = random_f32_vec(n * k, 0x1A11_B000, -0.5, 0.5);
+    let mut expected = vec![0.0f32; m * n];
+    for col in 0..n {
+        let mut sum = 0.0f32;
+        for kk in 0..k {
+            sum += a_data[kk] * b_data[col * k + kk];
+        }
+        expected[col] = sum;
+    }
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.load_llir(&llir);
+    rt.set_data(a, a_data);
+    rt.set_data(b_input, b_data);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
+}
+
+#[test]
+#[ignore = "large row-order CUDA functional repro for llama MLP residual beta=1 shape"]
+fn cublaslt_row_order_beta_one_candidate_executes_llama_mlp_residual_like_projection() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    let (m, n, k) = (1, 4096, 64);
+    let mut cx = Graph::new();
+    let a = cx.tensor((m, k));
+    let b_input = cx.tensor((n, k));
+    let b = b_input.t();
+    let c = cx.tensor((m, n));
+    let out = (a.matmul(b) + c).output();
+    let expected_orders = ("ROW", "COL", "ROW", "ROW");
+    let llir = extract_forced_cublaslt_llir_where(&mut cx, "mlp residual row-order", |llir| {
+        cublaslt_matrix_order_tuples(llir).contains(&expected_orders)
+            && cublaslt_scale_value_tuples(llir).contains(&(1.0, 1.0))
+    });
+
+    let a_data = random_f32_vec(m * k, 0x1A12_A000, -0.5, 0.5);
+    let b_data = random_f32_vec(n * k, 0x1A12_B000, -0.5, 0.5);
+    let c_data = random_f32_vec(m * n, 0x1A12_C000, -0.5, 0.5);
+    let mut expected = c_data.clone();
+    for col in 0..n {
+        for kk in 0..k {
+            expected[col] += a_data[kk] * b_data[col * k + kk];
+        }
+    }
+
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.load_llir(&llir);
+    rt.set_data(a, a_data);
+    rt.set_data(b_input, b_data);
+    rt.set_data(c, c_data);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(out.id), &expected, 1e-5, 1e-5);
+}
+
 #[test]
 #[ignore = "expensive CUDA functional candidate sweep; run with cargo test -p luminal_cuda_lite -- --ignored"]
 fn cublaslt_row_order_candidate_executes_batched_row_major_matmul() {
@@ -2762,10 +3081,17 @@ fn assert_no_cublaslt_llir_where(
 }
 
 fn cublaslt_ir_nodes(egraph: &SerializedEGraph) -> Vec<&NodeId> {
-    let cublaslt_kind_classes = egraph
+    op_ir_nodes(egraph, "cublaslt")
+        .into_iter()
+        .chain(op_ir_nodes(egraph, "cublaslt_scaled"))
+        .collect()
+}
+
+fn op_ir_nodes<'a>(egraph: &'a SerializedEGraph, kind_label: &str) -> Vec<&'a NodeId> {
+    let op_kind_classes = egraph
         .enodes
         .iter()
-        .filter(|(_, (label, _))| label == "cublaslt")
+        .filter(|(_, (label, _))| label == kind_label)
         .map(|(node, _)| egraph.node_to_class[node].clone())
         .collect::<Vec<_>>();
 
@@ -2776,12 +3102,93 @@ fn cublaslt_ir_nodes(egraph: &SerializedEGraph) -> Vec<&NodeId> {
             (label == "Op"
                 && children
                     .first()
-                    .is_some_and(|kind| cublaslt_kind_classes.contains(kind)))
+                    .is_some_and(|kind| op_kind_classes.contains(kind)))
             .then_some(node)
         })
         .collect()
 }
 
+fn dataflow_reachable_cublaslt_scaled_count(egraph: &SerializedEGraph) -> usize {
+    dataflow_reachable_cublaslt_count(egraph, true)
+}
+
+fn dataflow_reachable_cublaslt_raw_fp8_count(egraph: &SerializedEGraph) -> usize {
+    dataflow_reachable_cublaslt_count(egraph, false)
+}
+
+fn dataflow_reachable_cublaslt_count(egraph: &SerializedEGraph, scaled: bool) -> usize {
+    let reachable = dataflow_reachable_ir_classes(egraph);
+    egraph
+        .enodes
+        .iter()
+        .filter(|(node, (label, children))| {
+            label == "Op"
+                && reachable.contains(&egraph.node_to_class[*node])
+                && children.first().is_some_and(|kind_class| {
+                    egraph
+                        .eclasses
+                        .get(kind_class)
+                        .is_some_and(|(_, kind_nodes)| {
+                            kind_nodes.iter().any(|kind_node| {
+                                egraph.enodes.get(kind_node).is_some_and(|(kind_label, _)| {
+                                    if scaled {
+                                        kind_label == "cublaslt_scaled"
+                                    } else {
+                                        kind_label == "cublaslt"
+                                    }
+                                })
+                            })
+                        })
+                })
+        })
+        .count()
+}
+
+fn dataflow_reachable_ir_classes(egraph: &SerializedEGraph) -> FxHashSet<ClassId> {
+    let mut reachable = FxHashSet::default();
+    let mut stack = egraph.roots.clone();
+    while let Some(class) = stack.pop() {
+        if !reachable.insert(class.clone()) {
+            continue;
+        }
+        let Some((sort, nodes)) = egraph.eclasses.get(&class) else {
+            continue;
+        };
+        for node in nodes {
+            let Some((label, children)) = egraph.enodes.get(node) else {
+                continue;
+            };
+            match (sort.as_str(), label.as_str()) {
+                ("IR", "Output") => {
+                    if let Some(child) = children.first() {
+                        stack.push(child.clone());
+                    }
+                }
+                ("IR", "OutputJoin") => stack.extend(children.iter().cloned()),
+                ("IR", "Op") => {
+                    if let Some(inputs) = children.get(1) {
+                        stack.push(inputs.clone());
+                    }
+                }
+                ("IR", _) => {
+                    for child in children {
+                        if egraph
+                            .eclasses
+                            .get(child)
+                            .is_some_and(|(child_sort, _)| child_sort == "IR")
+                        {
+                            stack.push(child.clone());
+                        }
+                    }
+                }
+                ("IList", "ICons") => stack.extend(children.iter().cloned()),
+                _ => {}
+            }
+        }
+    }
+    reachable
+}
+
 fn llir_has_cublaslt(llir: &LLIRGraph) -> bool {
     !cublaslt_type_tuples(llir).is_empty()
 }
@@ -2800,6 +3207,13 @@ fn cublaslt_scale_value_tuples(llir: &LLIRGraph) -> Vec<CublasLtScaleValues> {
         .collect()
 }
 
+fn cublaslt_tensor_scale_input_tuples(llir: &LLIRGraph) -> Vec<(bool, bool)> {
+    llir.node_weights()
+        .filter_map(|op| op.to_dialect::<dyn HostOp>())
+        .filter_map(|host_op| cublaslt_tensor_scale_inputs(host_op.as_ref().as_ref()))
+        .collect()
+}
+
 fn cublaslt_epilogues(llir: &LLIRGraph) -> Vec<&'static str> {
     llir.node_weights()
         .filter_map(|op| op.to_dialect::<dyn HostOp>())

crates/luminal_cuda_lite/src/tests/flashinfer.rs

@@ -4,7 +4,7 @@
 //! 1. Pure egglog metadata (no GPU): trait wiring, sort + rewrite parse cleanly.
 //! 2. Egglog rule firing (no GPU): the rule unifies on a real paged-attention
 //!    HLIR and does NOT fire on bare attention or unrelated matmul/Gather mixes.
-//! 3. Mask op correctness (GPU): `ComputeAttnMask` produces the right (s, c) mask.
+//! 3. Mask helper correctness (GPU): the primitive-op `test_compute_attn_mask` builder produces the right (s, c) mask.
 //! 4. Full kernel correctness (GPU + JIT): direct `FlashInferAttention::execute`
 //!    compared against a luminal-compiled reference attention graph.
 //!
@@ -18,7 +18,7 @@ use luminal::op::{EgglogOp, IntoEgglogOp};
 use luminal::prelude::*;
 
 use crate::host::flashinfer::FlashInferAttention;
-use crate::host::{ComputeAttnMask, DeviceBuffer, HostOp};
+use crate::host::{DeviceBuffer, HostOp};
 use crate::runtime::CudaRuntime;
 use crate::tests::utilities::get_cuda_stream;
 
@@ -285,106 +285,6 @@ fn flashinfer_op_sort_shape() {
     assert!(dbg.contains("FlashInferAttention"));
 }
 
-#[test]
-fn compute_attn_mask_registers() {
-    assert!(
-        ops_contains_sort("ComputeAttnMask"),
-        "ComputeAttnMask is not in CudaRuntime::Ops"
-    );
-}
-
-// ─── Layer 2: ComputeAttnMask correctness ────────────────────────────────
-
-#[test]
-fn compute_attn_mask_matches_cpu_reference() {
-    let Some(stream) = get_cuda_stream() else {
-        return;
-    };
-
-    // 2 sequences, seq0 length=3, seq1 length=2 → s=2 queries (one per seq, decode),
-    // c=5 total context tokens (3+2).
-    let s_dim = 2usize;
-    let c_dim = 5usize;
-    let q_pos: Vec<i32> = vec![2, 1]; // last position in each seq
-    let qo_indptr: Vec<i32> = vec![0, 1, 2];
-    let kv_indptr: Vec<i32> = vec![0, 3, 5];
-    let r = kv_indptr.len();
-
-    let q_pos_buf = stream
-        .clone_htod(unsafe {
-            std::slice::from_raw_parts(q_pos.as_ptr() as *const u8, q_pos.len() * 4)
-        })
-        .unwrap();
-    let qo_buf = stream
-        .clone_htod(unsafe {
-            std::slice::from_raw_parts(qo_indptr.as_ptr() as *const u8, qo_indptr.len() * 4)
-        })
-        .unwrap();
-    let kv_buf = stream
-        .clone_htod(unsafe {
-            std::slice::from_raw_parts(kv_indptr.as_ptr() as *const u8, kv_indptr.len() * 4)
-        })
-        .unwrap();
-    let out_bytes = s_dim * c_dim * 4;
-    let out_buf = unsafe { stream.alloc::<u8>(out_bytes).unwrap() };
-
-    let op = ComputeAttnMask {
-        s_dim: Expression::from(s_dim),
-        c_dim: Expression::from(c_dim),
-    };
-
-    let q_pos_n = NodeIndex::new(0);
-    let qo_n = NodeIndex::new(1);
-    let kv_n = NodeIndex::new(2);
-    let out_n = NodeIndex::new(3);
-
-    let mut buffers = FxHashMap::default();
-    buffers.insert(
-        q_pos_n,
-        DeviceBuffer::new(q_pos_buf.device_ptr(&stream).0, q_pos.len() * 4),
-    );
-    buffers.insert(
-        qo_n,
-        DeviceBuffer::new(qo_buf.device_ptr(&stream).0, qo_indptr.len() * 4),
-    );
-    buffers.insert(
-        kv_n,
-        DeviceBuffer::new(kv_buf.device_ptr(&stream).0, kv_indptr.len() * 4),
-    );
-    buffers.insert(
-        out_n,
-        DeviceBuffer::new(out_buf.device_ptr(&stream).0, out_bytes),
-    );
-
-    let inputs = [q_pos_n, qo_n, kv_n];
-    let mut dyn_map = FxHashMap::default();
-    dyn_map.insert('r', r);
-
-    op.execute(&stream, out_n, &inputs, &buffers, &dyn_map)
-        .unwrap();
-    stream.synchronize().unwrap();
-
-    let host_bytes = stream.clone_dtoh(&out_buf).unwrap();
-    let mask: Vec<f32> = unsafe {
-        let mut bytes = std::mem::ManuallyDrop::new(host_bytes);
-        let len = bytes.len() / 4;
-        Vec::from_raw_parts(bytes.as_mut_ptr() as *mut f32, len, len)
-    };
-
-    // Expected: query 0 (q_pos=2, seq 0) attends to ctx [0, 3) i.e. mask[0, 0..3]=0;
-    //           query 1 (q_pos=1, seq 1) attends to ctx [3, 5) i.e. mask[1, 3..5]=0.
-    // Everywhere else is -1e10.
-    let mut expected = vec![-1e10f32; s_dim * c_dim];
-    for j in 0..3 {
-        expected[0 * c_dim + j] = 0.0;
-    }
-    for j in 3..5 {
-        expected[1 * c_dim + j] = 0.0;
-    }
-
-    assert_eq!(mask, expected);
-}
-
 // ─── Layer 3: FlashInfer kernel correctness ──────────────────────────────
 
 #[test]
@@ -527,7 +427,7 @@ fn test_indptr_to_request_idx(
     n: Expression,
 ) -> GraphTensor {
     let r = indptr.dims1();
-    let indices = graph.arange(n.clone()).expand_dim(1, r.clone());
+    let indices = graph.arange(n).expand_dim(1, r);
     let indptr_2d = indptr.expand_dim(0, n);
     let ge = indptr_2d.le(indices).cast(luminal::dtype::DType::Int);
     ge.sum(1).cast(luminal::dtype::DType::Int) - 1
@@ -541,13 +441,13 @@ fn test_compute_attn_mask(
     c: Expression,
 ) -> GraphTensor {
     let s = q_pos.dims1();
-    let q_request = test_indptr_to_request_idx(graph, qo_indptr, s.clone());
-    let c_request = test_indptr_to_request_idx(graph, kv_indptr, c.clone());
-    let c_arange = graph.arange(c.clone());
+    let q_request = test_indptr_to_request_idx(graph, qo_indptr, s);
+    let c_request = test_indptr_to_request_idx(graph, kv_indptr, c);
+    let c_arange = graph.arange(c);
     let c_kv_start = kv_indptr.gather(c_request);
     let c_local_pos = c_arange - c_kv_start;
-    let q_req_2d = q_request.expand_dim(1, c.clone());
-    let c_req_2d = c_request.expand_dim(0, s.clone());
+    let q_req_2d = q_request.expand_dim(1, c);
+    let c_req_2d = c_request.expand_dim(0, s);
     let same = q_req_2d.eq(c_req_2d);
     let c_pos_2d = c_local_pos.expand_dim(0, s);
     let qp_2d = q_pos.expand_dim(1, c);
@@ -577,6 +477,7 @@ fn scatter_rows(
 
 /// Handles to every named input of the paged-attention test graph, returned
 /// alongside the graph so the GA-selection test can `set_data` on each one.
+#[allow(dead_code)]
 struct PagedAttnHandles {
     q_rope: GraphTensor,
     k_rope: GraphTensor,

crates/luminal_cuda_lite/src/tests/fusion.rs

@@ -1,7 +1,9 @@
+use as_any::Downcast;
 use luminal::egglog_utils::{egglog_to_llir, random_initial_choice};
 use luminal::prelude::*;
 
 use crate::kernel::KernelOp;
+use crate::kernel::fusion::{CudaBinaryElementwise, CudaUnaryElementwise};
 use crate::runtime::CudaRuntime;
 use crate::tests::utilities::{
     TOLERANCE_SAFETY_FACTOR, dtype_epsilon, random_f32_vec, test_binary_cuda, test_unary_cuda,
@@ -86,7 +88,7 @@ fn test_unary_fusion_preserves_output() {
 #[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.
+    // reachable as a single marker region containing all three elementwise ops.
     let mut cx = Graph::new();
     let a = cx.tensor(16);
     let _b = a.sin().sqrt().exp2().output();
@@ -104,7 +106,7 @@ fn test_three_unary_ops_fuse() {
 #[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).
+    // four elementwise 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();
@@ -317,8 +319,15 @@ fn extract_all_fused_regions(cx: &mut Graph) -> Vec<FusedRegion> {
 
         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())
+                op.to_dialect::<dyn KernelOp>().map(|k| {
+                    if let Some(elem) = (***k).downcast_ref::<CudaUnaryElementwise>() {
+                        format!("Fused{}", elem.op)
+                    } else if let Some(elem) = (***k).downcast_ref::<CudaBinaryElementwise>() {
+                        format!("Fused{}", elem.op)
+                    } else {
+                        k.kernel_name().to_string()
+                    }
+                })
             })
         };
 
@@ -343,12 +352,13 @@ fn extract_all_fused_regions(cx: &mut Graph) -> Vec<FusedRegion> {
 
             // 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.
+            // is itself a FusionStart is a cascade layer, not a new external
+            // tensor. A FusionEnd predecessor is a real external region output
+            // in the generic singleton-region model, so do not walk through it.
             let resolve_source = |mut n: NodeIndex| -> NodeIndex {
                 loop {
                     match name_of(n).as_deref() {
-                        Some("FusionStart") | Some("FusionEnd") => {
+                        Some("FusionStart") => {
                             let mut inc = llir.neighbors_directed(n, petgraph::Direction::Incoming);
                             match inc.next() {
                                 Some(p) => n = p,
@@ -379,15 +389,6 @@ fn extract_all_fused_regions(cx: &mut Graph) -> Vec<FusedRegion> {
                             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));
                                 }
@@ -467,6 +468,15 @@ fn test_single_binary_does_not_fuse_alone() {
 fn test_chain_of_binaries_fuses() {
     // `(a + b) * c`: three external inputs collapse into one region with
     // internal [Add, Mul] and 3 FusionStarts.
+    //
+    // Requires BB family, which is opt-in at runtime via
+    // LUMINAL_FUSION_FAMILIES. Set it before the graph build so the rules
+    // emitted from FusionEnd::rewrites include the B-B pair-fuse rules.
+    // SAFETY: tests run in parallel; we set this before constructing the
+    // Graph, and never unset, so concurrent tests just see BB on.
+    unsafe {
+        std::env::set_var("LUMINAL_FUSION_FAMILIES", "uu,bu,ub,bb");
+    }
     let mut cx = Graph::new();
     let a = cx.tensor(8);
     let b = cx.tensor(8);
@@ -520,6 +530,13 @@ fn test_unary_then_binary_fuses() {
 }
 
 #[test]
+// Subsume in grow rules (introduced to bound the BB partial-FE explosion)
+// means a multi-consumer producer can no longer be fused into the same
+// region as all its consumers — only one branch wins. The diamond's `t`
+// has two consumers, so the structural "one 5-op region" outcome is no
+// longer guaranteed. Numerical correctness still holds (see
+// test_diamond_dag_preserves_output).
+#[ignore = "asserts pre-subsume ideal multi-consumer fusion shape"]
 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
@@ -650,6 +667,7 @@ fn test_diamond_dag_preserves_output() {
 // ---- Marker invariant tests ----
 
 #[test]
+#[ignore = "asserts pre-subsume ideal multi-consumer fusion shape"]
 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.
@@ -677,6 +695,7 @@ fn test_fused_region_has_exactly_one_end() {
 }
 
 #[test]
+#[ignore = "asserts pre-subsume ideal multi-consumer fusion shape"]
 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
@@ -768,6 +787,10 @@ 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.
+    // See test_chain_of_binaries_fuses for the LUMINAL_FUSION_FAMILIES note.
+    unsafe {
+        std::env::set_var("LUMINAL_FUSION_FAMILIES", "uu,bu,ub,bb");
+    }
     let mut cx = Graph::new();
     let a = cx.tensor(8);
     let b = cx.tensor(8);
@@ -809,6 +832,7 @@ fn test_grow_fe_to_binary_rhs() {
 }
 
 #[test]
+#[ignore = "asserts pre-subsume two-FE merge shape; numerical correctness preserved"]
 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

crates/luminal_cuda_lite/src/tests/generic_matmul_rewrite.rs

@@ -0,0 +1,169 @@
+use luminal::{
+    egglog_utils::{
+        NodeId, SerializedEGraph, egglog_to_llir, random_initial_choice, validate_choice_set,
+    },
+    prelude::*,
+};
+use rand::{SeedableRng, rngs::StdRng};
+
+use crate::{kernel::KernelOp, runtime::CudaRuntime};
+
+use super::utilities::{assert_close, get_cuda_stream};
+
+#[test]
+fn generic_matmul_covers_noncontiguous_merged_head_projection() {
+    let mut cx = Graph::default();
+    let heads = 3;
+    let seq = 4;
+    let head_dim = 5;
+    let hidden = heads * head_dim;
+    let out_dim = 7;
+
+    let attn = cx.tensor((heads, seq, head_dim));
+    let weight = cx.tensor((out_dim, hidden));
+    let merged = attn.transpose(0, 1).merge_dims(1, 2);
+    merged.matmul(weight.t()).output();
+
+    cx.build_search_space::<CudaRuntime>();
+    let llir = extract_forced_kernel_llir(&mut cx, "GenericMatmul");
+    let names = llir_kernel_names(&llir);
+
+    assert!(
+        names.contains(&"GenericMatmul"),
+        "expected generic matmul fallback, kernels: {names:?}"
+    );
+    assert!(
+        !names.contains(&"Mul") && !names.contains(&"SumReduce"),
+        "generic matmul should prune the broadcast multiply/sum fallback, kernels: {names:?}"
+    );
+}
+
+#[test]
+fn generic_matmul_executes_noncontiguous_merged_head_projection() {
+    let mut cx = Graph::default();
+    let heads = 3;
+    let seq = 4;
+    let head_dim = 5;
+    let hidden = heads * head_dim;
+    let out_dim = 7;
+
+    let attn = cx.tensor((heads, seq, head_dim));
+    let weight = cx.tensor((out_dim, hidden));
+    let merged = attn.transpose(0, 1).merge_dims(1, 2);
+    let output = merged.matmul(weight.t()).output();
+
+    cx.build_search_space::<CudaRuntime>();
+    let stream = get_cuda_stream().expect("CUDA device required for GenericMatmul execution test");
+    let mut rt = CudaRuntime::initialize(stream);
+
+    let attn_data = seeded_data(heads * seq * head_dim, 0.19, -0.09);
+    let weight_data = seeded_data(out_dim * hidden, 0.14, -0.06);
+    rt.set_data(attn, attn_data.as_slice());
+    rt.set_data(weight, weight_data.as_slice());
+
+    rt = cx.search(rt, 1);
+    assert!(
+        rt.kernel_names().contains(&"GenericMatmul"),
+        "expected GenericMatmul to be selected, kernels: {:?}",
+        rt.kernel_names()
+    );
+
+    rt.execute(&cx.dyn_map);
+    let result = rt.get_f32(output.id);
+
+    let mut expected = vec![0.0; seq * out_dim];
+    for token in 0..seq {
+        for out_col in 0..out_dim {
+            let mut sum = 0.0;
+            for inner in 0..hidden {
+                let head = inner / head_dim;
+                let dim = inner % head_dim;
+                let attn_idx = head * seq * head_dim + token * head_dim + dim;
+                sum += attn_data[attn_idx] * weight_data[out_col * hidden + inner];
+            }
+            expected[token * out_dim + out_col] = sum;
+        }
+    }
+
+    assert_close(&result, &expected, 1e-5, 1e-5);
+}
+
+fn seeded_data(len: usize, scale: f32, bias: f32) -> Vec<f32> {
+    (0..len)
+        .map(|i| {
+            let x = ((i * 37 + 11) % 97) as f32 / 97.0;
+            x * scale + bias
+        })
+        .collect()
+}
+
+fn extract_forced_kernel_llir(cx: &mut Graph, kernel_name: &str) -> LLIRGraph {
+    let egraph = cx.egraph().expect("search space should have an e-graph");
+    let ops = cx
+        .egglog_ops()
+        .expect("search space should have registered egglog ops");
+    let kernel_nodes = op_ir_nodes(egraph, kernel_name);
+    assert!(
+        !kernel_nodes.is_empty(),
+        "expected at least one {kernel_name} candidate"
+    );
+
+    for (idx, kernel_node) in kernel_nodes.iter().enumerate() {
+        let mut rng = StdRng::seed_from_u64(0x9E_EE_0000 + idx as u64);
+        let mut choices = random_initial_choice(egraph, &mut rng);
+        let kernel_class = &egraph.node_to_class[*kernel_node];
+        choices.insert(kernel_class, kernel_node);
+
+        if validate_choice_set(egraph, &choices, ops).is_err() {
+            continue;
+        }
+
+        let mut list_cache = FxHashMap::default();
+        let mut expr_cache = FxHashMap::default();
+        let llir = egglog_to_llir(
+            egraph,
+            choices,
+            ops,
+            &cx.custom_ops,
+            &mut list_cache,
+            &mut expr_cache,
+            None,
+        );
+        if llir_kernel_names(&llir).contains(&kernel_name) {
+            return llir;
+        }
+    }
+
+    panic!("could not extract a valid {kernel_name} candidate");
+}
+
+fn llir_kernel_names(llir: &LLIRGraph) -> Vec<&'static str> {
+    llir.node_indices()
+        .filter_map(|node| {
+            llir[node]
+                .to_dialect::<dyn KernelOp>()
+                .map(|kernel| kernel.kernel_name())
+        })
+        .collect()
+}
+
+fn op_ir_nodes<'a>(egraph: &'a SerializedEGraph, kind_label: &str) -> Vec<&'a NodeId> {
+    let op_kind_classes = egraph
+        .enodes
+        .iter()
+        .filter(|(_, (label, _))| label == kind_label)
+        .map(|(node, _)| egraph.node_to_class[node].clone())
+        .collect::<Vec<_>>();
+
+    egraph
+        .enodes
+        .iter()
+        .filter_map(|(node, (label, children))| {
+            (label == "Op"
+                && children
+                    .first()
+                    .is_some_and(|kind| op_kind_classes.contains(kind)))
+            .then_some(node)
+        })
+        .collect()
+}

crates/luminal_cuda_lite/src/tests/mod.rs

@@ -5,12 +5,16 @@ mod bucket_tests;
 #[cfg(test)]
 mod consumed_buffer_tests;
 #[cfg(test)]
+mod conv2d_rewrite;
+#[cfg(test)]
 mod cublaslt_rewrite_tests;
 #[cfg(test)]
 mod flashinfer;
 #[cfg(test)]
 mod fusion;
 #[cfg(test)]
+mod generic_matmul_rewrite;
+#[cfg(test)]
 mod model_fuzz;
 #[cfg(test)]
 mod op_functional_tests;
@@ -19,4 +23,8 @@ mod performance_tests;
 #[cfg(test)]
 mod qwen3_moe_rewrite;
 #[cfg(test)]
+mod rope_test;
+#[cfg(test)]
+mod search_equivalence_fuzz;
+#[cfg(test)]
 mod transformer;

crates/luminal_cuda_lite/src/tests/model_fuzz.rs

@@ -305,7 +305,7 @@ fn fuzz_layer_no_attn(
 }
 
 /// Test a SwiGLU MLP with HLIR-only to specifically verify
-/// the HLIR matmul decomposition (KernelMul + KernelSumReduce).
+/// the HLIR matmul decomposition (elementwise Mul + KernelSumReduce).
 fn fuzz_mlp_hlir_only(seq: usize, hidden: usize, intermediate: usize, seed: u64) {
     let Some(stream) = get_cuda_stream() else {
         return;

crates/luminal_cuda_lite/src/tests/performance_tests.rs

@@ -6,7 +6,7 @@ use crate::cuda_bandwidth_gbps;
 use crate::runtime::CudaRuntime;
 
 /// Test that measures bandwidth utilization for a large element-wise add kernel.
-/// This demonstrates that KernelAdd can achieve reasonable bandwidth with large tensors.
+/// This demonstrates that generic fused Add can achieve reasonable bandwidth with large tensors.
 #[test]
 pub fn kernel_add_bandwidth_test() {
     // 64M elements = 256MB per tensor, 768MB total memory traffic (2 reads + 1 write)
@@ -40,7 +40,7 @@ pub fn kernel_add_bandwidth_test() {
     rt.execute(&cx.dyn_map);
 
     // Print stats
-    println!("\n=== Large KernelAdd Bandwidth Test ===");
+    println!("\n=== Large Fused Add Bandwidth Test ===");
     println!(
         "Tensor size: {} elements ({} MB per tensor)",
         size,

crates/luminal_cuda_lite/src/tests/qwen3_moe_rewrite.rs

@@ -2,16 +2,13 @@ use half::bf16;
 use luminal::{dtype::DType, prelude::*, shape::Expression};
 
 use super::utilities::{assert_close, get_cuda_stream, random_f32_vec};
-use crate::{
-    host::moe::{GLUMoE, GLUMoEMode},
-    runtime::CudaRuntime,
-};
+use crate::{host::moe::GLUMoE, runtime::CudaRuntime};
 
 const SEQ: usize = 2;
-const HIDDEN: usize = 16;
+const HIDDEN: usize = 32;
 const NUM_EXPERTS: usize = 8;
 const TOP_K: usize = 2;
-const MOE_INTERMEDIATE: usize = 6;
+const MOE_INTERMEDIATE: usize = 12;
 const RMS_NORM_EPS: f32 = 1e-6;
 
 struct QwenMoeGraph {
@@ -58,6 +55,7 @@ fn build_qwen_moe_graph() -> QwenMoeGraph {
         .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_values = top_k_values / top_k_values.sum(n - 1).expand_dim(n - 1, TOP_K);
 
     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);
@@ -172,25 +170,44 @@ fn gemma_gelu(x: GraphTensor) -> GraphTensor {
     x * scaled.sigmoid()
 }
 
-fn glumoe_modes(rt: &CudaRuntime) -> Vec<GLUMoEMode> {
-    rt.host_ops()
-        .into_iter()
-        .filter_map(|op| {
-            op.as_any()
-                .downcast_ref::<GLUMoE>()
-                .map(|glumoe| glumoe.mode)
-        })
-        .collect()
+fn search_space_contains(cx: &Graph, op_name: &str) -> bool {
+    let egraph = cx.egraph().expect("test should build an e-graph");
+
+    for (label, children) in egraph.enodes.values() {
+        if label != "Op" {
+            continue;
+        }
+        let Some(kind_eclass) = children.first() else {
+            continue;
+        };
+        let Some((_, kind_enodes)) = egraph.eclasses.get(kind_eclass) else {
+            continue;
+        };
+        if kind_enodes
+            .iter()
+            .any(|kind_node| egraph.enodes[kind_node].0 == op_name)
+        {
+            return true;
+        }
+    }
+    false
 }
 
-fn run_qwen_moe(use_glumoe: bool) -> (Vec<f32>, Vec<GLUMoEMode>) {
+fn assert_glumoe_in_search_space(cx: &Graph) {
+    assert!(
+        search_space_contains(cx, "GLUMoE"),
+        "GLUMoE was not in the e-graph search space"
+    );
+}
+
+fn run_qwen_moe(include_glumoe: bool) -> Vec<f32> {
     let Some(stream) = get_cuda_stream() else {
-        return (vec![], vec![]);
+        return vec![];
     };
 
     let mut model = build_qwen_moe_graph();
     model.graph.set_dim('s', SEQ);
-    if use_glumoe {
+    if include_glumoe {
         model.graph.build_search_space::<CudaRuntime>();
     } else {
         model
@@ -217,17 +234,17 @@ fn run_qwen_moe(use_glumoe: bool) -> (Vec<f32>, Vec<GLUMoEMode>) {
     rt = model.graph.search(rt, 10);
     rt.execute(&model.graph.dyn_map);
 
-    (rt.get_f32(model.output.id), glumoe_modes(&rt))
+    rt.get_f32(model.output.id)
 }
 
-fn run_gemma_moe(use_glumoe: bool) -> (Vec<f32>, Vec<GLUMoEMode>) {
+fn run_gemma_moe(include_glumoe: bool) -> Vec<f32> {
     let Some(stream) = get_cuda_stream() else {
-        return (vec![], vec![]);
+        return vec![];
     };
 
     let mut model = build_gemma_moe_graph();
     model.graph.set_dim('s', SEQ);
-    if use_glumoe {
+    if include_glumoe {
         model.graph.build_search_space::<CudaRuntime>();
     } else {
         model
@@ -260,51 +277,51 @@ fn run_gemma_moe(use_glumoe: bool) -> (Vec<f32>, Vec<GLUMoEMode>) {
     rt = model.graph.search(rt, 10);
     rt.execute(&model.graph.dyn_map);
 
-    (rt.get_f32(model.output.id), glumoe_modes(&rt))
+    rt.get_f32(model.output.id)
 }
 
 #[test]
 fn test_glumoe_matches_qwen_swiglu_pattern() {
-    let (_result, modes) = run_qwen_moe(true);
-    if modes.is_empty() {
+    if get_cuda_stream().is_none() {
         return;
     }
 
-    assert_eq!(modes, vec![GLUMoEMode::SwiGLU]);
+    let mut model = build_qwen_moe_graph();
+    model.graph.set_dim('s', SEQ);
+    model.graph.build_search_space::<CudaRuntime>();
+    assert_glumoe_in_search_space(&model.graph);
 }
 
 #[test]
 fn test_glumoe_matches_gemma_gelu_pattern() {
-    let (_result, modes) = run_gemma_moe(true);
-    if modes.is_empty() {
+    if get_cuda_stream().is_none() {
         return;
     }
 
-    assert_eq!(modes, vec![GLUMoEMode::GemmaGELU]);
+    let mut model = build_gemma_moe_graph();
+    model.graph.set_dim('s', SEQ);
+    model.graph.build_search_space::<CudaRuntime>();
+    assert_glumoe_in_search_space(&model.graph);
 }
 
 #[test]
 fn test_glumoe_swiglu_matches_unfused_output() {
-    let (expected, baseline_modes) = run_qwen_moe(false);
+    let expected = 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]);
+    let actual = run_qwen_moe(true);
     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);
+    let expected = 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]);
+    let actual = run_gemma_moe(true);
     assert_close(&actual, &expected, 3e-2, 3e-2);
 }

crates/luminal_cuda_lite/src/tests/rope_test.rs

@@ -0,0 +1,112 @@
+use cudarc::driver::CudaContext;
+use luminal::{graph::Graph, op::Runtime};
+
+use crate::{kernel::apply_rope, runtime::CudaRuntime};
+
+fn cpu_rope(x: &[f32], cos: &[f32], sin: &[f32], s: usize, h: usize, d: usize) -> Vec<f32> {
+    assert!(d.is_multiple_of(2));
+    let mut out = vec![0.0f32; s * h * d];
+    for si in 0..s {
+        for hi in 0..h {
+            for i in 0..d {
+                let xi = x[si * h * d + hi * d + i];
+                let xpair = if i % 2 == 0 {
+                    -x[si * h * d + hi * d + i + 1]
+                } else {
+                    x[si * h * d + hi * d + i - 1]
+                };
+                let c = cos[si * d + i];
+                let sn = sin[si * d + i];
+                out[si * h * d + hi * d + i] = xi * c + xpair * sn;
+            }
+        }
+    }
+    out
+}
+
+#[test]
+fn rope_matches_cpu_reference() {
+    let s = 8;
+    let h = 4;
+    let d = 32;
+    let mut cx = Graph::default();
+    let x = cx.tensor((s, h, d));
+    let cos = cx.tensor((s, d));
+    let sin = cx.tensor((s, d));
+    let y = apply_rope(x, cos, sin).output();
+
+    let x_data: Vec<f32> = (0..s * h * d).map(|i| ((i as f32) * 0.013).sin()).collect();
+    let cos_data: Vec<f32> = (0..s * d).map(|i| ((i as f32) * 0.017).cos()).collect();
+    let sin_data: Vec<f32> = (0..s * d).map(|i| ((i as f32) * 0.017).sin()).collect();
+
+    let ctx = CudaContext::new(0).unwrap();
+    ctx.bind_to_thread().unwrap();
+    let stream = ctx.default_stream();
+    cx.build_search_space::<CudaRuntime>();
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.set_data(x, x_data.clone());
+    rt.set_data(cos, cos_data.clone());
+    rt.set_data(sin, sin_data.clone());
+    rt = cx.search(rt, 1);
+    rt.execute(&cx.dyn_map);
+    let got = rt.get_f32(y.id);
+
+    let expected = cpu_rope(&x_data, &cos_data, &sin_data, s, h, d);
+    let mut max_err = 0.0f32;
+    for (g, e) in got.iter().zip(expected.iter()) {
+        let err = (g - e).abs();
+        if err > max_err {
+            max_err = err;
+        }
+    }
+    eprintln!("rope: max abs err: {max_err}");
+    assert!(max_err < 1e-5, "max abs error {max_err} too high");
+}
+
+#[test]
+fn rope_flux2_shape() {
+    // Flux 2 transformer attention: S=1536 (img+txt), H=48, D=128.
+    let s = 1536;
+    let h = 48;
+    let d = 128;
+    let mut cx = Graph::default();
+    let x = cx.tensor((s, h, d));
+    let cos = cx.tensor((s, d));
+    let sin = cx.tensor((s, d));
+    let y = apply_rope(x, cos, sin).output();
+
+    use rand::{Rng, SeedableRng};
+    let mut rng = rand::rngs::SmallRng::seed_from_u64(11);
+    let x_data: Vec<f32> = (0..s * h * d)
+        .map(|_| rng.random_range(-2.0..2.0_f32))
+        .collect();
+    let cos_data: Vec<f32> = (0..s * d)
+        .map(|_| rng.random_range(-1.0..1.0_f32))
+        .collect();
+    let sin_data: Vec<f32> = (0..s * d)
+        .map(|_| rng.random_range(-1.0..1.0_f32))
+        .collect();
+
+    let ctx = CudaContext::new(0).unwrap();
+    ctx.bind_to_thread().unwrap();
+    let stream = ctx.default_stream();
+    cx.build_search_space::<CudaRuntime>();
+    let mut rt = CudaRuntime::initialize(stream);
+    rt.set_data(x, x_data.clone());
+    rt.set_data(cos, cos_data.clone());
+    rt.set_data(sin, sin_data.clone());
+    rt = cx.search(rt, 1);
+    rt.execute(&cx.dyn_map);
+    let got = rt.get_f32(y.id);
+
+    let expected = cpu_rope(&x_data, &cos_data, &sin_data, s, h, d);
+    let mut max_err = 0.0f32;
+    for (g, e) in got.iter().zip(expected.iter()) {
+        let err = (g - e).abs();
+        if err > max_err {
+            max_err = err;
+        }
+    }
+    eprintln!("rope flux2: max abs err: {max_err}");
+    assert!(max_err < 1e-4, "max abs error {max_err} too high");
+}

crates/luminal_cuda_lite/src/tests/search_equivalence_fuzz.rs

@@ -0,0 +1,374 @@
+//! End-to-end e-graph search-space equivalence fuzz tests.
+//!
+//! These tests do not compare against a hand-written reference. They assert the
+//! stronger search invariant: every selectable LLIR graph from the same e-graph
+//! must produce finite, numerically close outputs for the same runtime inputs.
+
+#[allow(dead_code)]
+#[path = "../../../../examples/llama/src/model.rs"]
+mod llama_model;
+
+use half::bf16;
+use luminal::{dtype::DType, prelude::*, shape::Expression};
+use rand::{Rng, SeedableRng, rngs::StdRng};
+
+use super::utilities::{CudaSearchEquivalenceFuzzer, get_cuda_stream, random_f32_vec};
+
+const SEARCH_EQUIV_SAMPLES: usize = 32;
+
+fn random_bf16_vec(n: usize, seed: u64, low: f32, high: f32) -> Vec<bf16> {
+    random_f32_vec(n, seed, low, high)
+        .into_iter()
+        .map(bf16::from_f32)
+        .collect()
+}
+
+fn rms_norm(x: GraphTensor, weight: GraphTensor, eps: f32) -> GraphTensor {
+    let normed = x.std_norm(x.shape.last_axis(), eps);
+    normed * weight.expand_lhs(&x.dims()[..x.dims().len() - 1])
+}
+
+#[allow(clippy::excessive_precision)]
+fn gemma_gelu(x: GraphTensor) -> GraphTensor {
+    let scaled = 1.5957691216 * x * (1. + 0.044715 * x * x);
+    x * scaled.sigmoid()
+}
+
+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)
+}
+
+#[test]
+fn llama_architecture_search_space_equivalence_fuzz() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    const SEQ: usize = 2;
+    const CTX: usize = 3;
+    const SLOTS: usize = 4;
+
+    let config = llama_model::LlamaConfig {
+        layers: 2,
+        hidden: 32,
+        intermediate: 64,
+        head_dim: 8,
+        kv_groups: 2,
+        vocab_size: 64,
+    };
+
+    let mut cx = Graph::default();
+    cx.set_dim('s', SEQ);
+    cx.set_dim('c', CTX);
+
+    let input = cx.named_tensor("input", 's').as_dtype(DType::Int);
+    let q_pos = cx.named_tensor("q_pos", 's').as_dtype(DType::Int);
+    let scatter_idx = cx.named_tensor("scatter_idx", 's').as_dtype(DType::Int);
+    let gather_idx = cx.named_tensor("gather_idx", 'c').as_dtype(DType::Int);
+    let attn_mask = cx.named_tensor("attn_mask", ('s', 'c'));
+    let kv_cache = llama_model::KVCache::new_with_config(&mut cx, SLOTS, config);
+    let llama = llama_model::Llama::init_with_config(&mut cx, config);
+
+    let (logits, cache_outputs) =
+        llama.forward(input, q_pos, scatter_idx, gather_idx, attn_mask, &kv_cache);
+    let logits = logits.output();
+    let mut fuzzer = CudaSearchEquivalenceFuzzer::new(&mut cx, &stream)
+        .seed(0x5EED_1234)
+        .samples(SEARCH_EQUIV_SAMPLES)
+        .generation_size(8)
+        .mutations(3)
+        .build_options(BuildSearchSpaceOptions::new().max_memory_mib(512))
+        .output_f32(logits.id, "logits", 5e-2, 5e-2);
+    for (layer, (k_out, v_out)) in cache_outputs.into_iter().enumerate() {
+        let k_out = k_out.output();
+        let v_out = v_out.output();
+        fuzzer = fuzzer.output_f32(k_out.id, format!("layer{layer}.k_cache"), 3e-3, 3e-3);
+        fuzzer = fuzzer.output_f32(v_out.id, format!("layer{layer}.v_cache"), 3e-3, 3e-3);
+    }
+
+    let mut rng = StdRng::seed_from_u64(0x11A_AA55);
+    fuzzer = fuzzer
+        .input_i32(input.id, vec![3, 17])
+        .input_i32(q_pos.id, vec![1, 2])
+        .input_i32(scatter_idx.id, vec![1, 2])
+        .input_i32(gather_idx.id, vec![0, 1, 2])
+        .input_f32(attn_mask.id, vec![0.0, 0.0, -1e4, 0.0, 0.0, 0.0]);
+
+    let kv_dim = config.kv_dim();
+    for tensor in kv_cache.tensors() {
+        fuzzer = fuzzer.input_f32(tensor.id, vec![0.0; SLOTS * kv_dim]);
+    }
+    for tensor in llama.parameter_tensors() {
+        let elements = tensor
+            .dims()
+            .iter()
+            .map(|dim| dim.to_usize().expect("tiny llama test uses static params"))
+            .product::<usize>();
+        let data = (0..elements)
+            .map(|_| rng.random_range(-0.08f32..0.08f32))
+            .collect::<Vec<_>>();
+        fuzzer = fuzzer.input_f32(tensor.id, data);
+    }
+
+    let report = fuzzer.run();
+    eprintln!("llama search equivalence fuzz report: {report:?}");
+}
+
+#[test]
+fn gemma_architecture_search_space_equivalence_fuzz() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    const SEQ: usize = 2;
+    const HIDDEN: usize = 32;
+    const Q_DIM: usize = 24;
+    const INTERMEDIATE: usize = 64;
+    const EPS: f32 = 1e-6;
+
+    let mut cx = Graph::default();
+    let input = cx.tensor((SEQ, HIDDEN));
+    let attn_norm_w = cx.tensor(HIDDEN);
+    let post_attn_norm_w = cx.tensor(HIDDEN);
+    let pre_ff_norm_w = cx.tensor(HIDDEN);
+    let post_ff_norm_w = cx.tensor(HIDDEN);
+    let proj_w = cx.tensor((Q_DIM, HIDDEN));
+    let o_proj_w = cx.tensor((HIDDEN, Q_DIM));
+    let w_gate = cx.tensor((INTERMEDIATE, HIDDEN));
+    let w_up = cx.tensor((INTERMEDIATE, HIDDEN));
+    let w_down = cx.tensor((HIDDEN, INTERMEDIATE));
+
+    let normed = rms_norm(input, attn_norm_w, EPS);
+    let proj_out = normed.matmul(proj_w.t()).matmul(o_proj_w.t());
+    let attn_normed = rms_norm(proj_out, post_attn_norm_w, EPS);
+    let x = input + attn_normed;
+    let ff_normed = rms_norm(x, pre_ff_norm_w, EPS);
+    let mlp_out =
+        (gemma_gelu(ff_normed.matmul(w_gate.t())) * ff_normed.matmul(w_up.t())).matmul(w_down.t());
+    let mlp_normed = rms_norm(mlp_out, post_ff_norm_w, EPS);
+    let out = (x + mlp_normed).output();
+
+    let report = CudaSearchEquivalenceFuzzer::new(&mut cx, &stream)
+        .seed(0x6E4D_4DAA)
+        .samples(SEARCH_EQUIV_SAMPLES)
+        .generation_size(8)
+        .mutations(3)
+        .build_options(BuildSearchSpaceOptions::new().max_memory_mib(512))
+        .input_f32(input.id, random_f32_vec(SEQ * HIDDEN, 101, -0.15, 0.15))
+        .input_f32(attn_norm_w.id, random_f32_vec(HIDDEN, 102, 0.7, 1.3))
+        .input_f32(post_attn_norm_w.id, random_f32_vec(HIDDEN, 103, 0.7, 1.3))
+        .input_f32(pre_ff_norm_w.id, random_f32_vec(HIDDEN, 104, 0.7, 1.3))
+        .input_f32(post_ff_norm_w.id, random_f32_vec(HIDDEN, 105, 0.7, 1.3))
+        .input_f32(proj_w.id, random_f32_vec(Q_DIM * HIDDEN, 106, -0.08, 0.08))
+        .input_f32(
+            o_proj_w.id,
+            random_f32_vec(HIDDEN * Q_DIM, 107, -0.08, 0.08),
+        )
+        .input_f32(
+            w_gate.id,
+            random_f32_vec(INTERMEDIATE * HIDDEN, 108, -0.08, 0.08),
+        )
+        .input_f32(
+            w_up.id,
+            random_f32_vec(INTERMEDIATE * HIDDEN, 109, -0.08, 0.08),
+        )
+        .input_f32(
+            w_down.id,
+            random_f32_vec(HIDDEN * INTERMEDIATE, 110, -0.08, 0.08),
+        )
+        .output_f32(out.id, "gemma_block", 5e-3, 5e-3)
+        .run();
+    eprintln!("gemma search equivalence fuzz report: {report:?}");
+}
+
+#[test]
+fn moe_architecture_search_space_equivalence_fuzz() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    const SEQ: usize = 2;
+    const HIDDEN: usize = 16;
+    const NUM_EXPERTS: usize = 8;
+    const TOP_K: usize = 2;
+    const MOE_INTERMEDIATE: usize = 6;
+    const EPS: f32 = 1e-6;
+
+    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, 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 mut weights_exp = top_k_weights.unsqueeze(top_k_weights.dims().len());
+    weights_exp.shape.expand(down_out.dims());
+    let out = (down_out * weights_exp).sum(n - 1).output();
+    cx.set_dim('s', SEQ);
+
+    let report = CudaSearchEquivalenceFuzzer::new(&mut cx, &stream)
+        .seed(0x0DEE_55EE)
+        .samples(SEARCH_EQUIV_SAMPLES)
+        .generation_size(8)
+        .mutations(3)
+        .build_options(BuildSearchSpaceOptions::new().max_memory_mib(512))
+        .input_f32(
+            router_input.id,
+            random_f32_vec(SEQ * HIDDEN, 201, -0.15, 0.15),
+        )
+        .input_f32(
+            expert_input.id,
+            random_f32_vec(SEQ * HIDDEN, 202, -0.15, 0.15),
+        )
+        .input_f32(router_scale.id, random_f32_vec(HIDDEN, 203, 0.7, 1.3))
+        .input_f32(
+            router_proj.id,
+            random_f32_vec(NUM_EXPERTS * HIDDEN, 204, -0.2, 0.2),
+        )
+        .input_f32(
+            per_expert_scale.id,
+            random_f32_vec(NUM_EXPERTS, 205, 0.5, 1.5),
+        )
+        .input_bf16(
+            gate_up_weights.id,
+            random_bf16_vec(NUM_EXPERTS * MOE_INTERMEDIATE * 2 * HIDDEN, 206, -0.1, 0.1),
+        )
+        .input_bf16(
+            down_weights.id,
+            random_bf16_vec(NUM_EXPERTS * HIDDEN * MOE_INTERMEDIATE, 207, -0.1, 0.1),
+        )
+        .output_f32(out.id, "gemma_moe_block", 5e-2, 5e-2)
+        .run();
+    eprintln!("moe search equivalence fuzz report: {report:?}");
+}
+
+#[test]
+fn moe_architecture_native_reference_fuzz() {
+    let Some(stream) = get_cuda_stream() else {
+        return;
+    };
+
+    const SEQ: usize = 2;
+    const HIDDEN: usize = 16;
+    const NUM_EXPERTS: usize = 8;
+    const TOP_K: usize = 2;
+    const MOE_INTERMEDIATE: usize = 6;
+
+    let mut cx = Graph::default();
+    let input = 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 = input.dims().len();
+    let e_dim = *router.dims().first().unwrap();
+    let k_expr = Expression::from(TOP_K);
+
+    let routing_weights = input.matmul(router.t()).softmax(n - 1);
+    let top_k_indices = routing_weights.topk_indexes(TOP_K, n - 1);
+    let row_offsets = 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_weights = top_k_values / top_k_values.sum(n - 1).expand_dim(n - 1, TOP_K);
+
+    let gate_up_gathered = gather_experts(input, top_k_indices, gate_up_weights).cast(DType::F32);
+    let input_exp = input.expand_dim(n - 1, TOP_K).unsqueeze(n);
+    let gate_up_out = input_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(input, top_k_indices, down_weights).cast(DType::F32);
+    let down_out = hidden
+        .unsqueeze(2)
+        .matmul(down_gathered.transpose(2, 3))
+        .squeeze(2);
+    let mut weights_exp = top_k_weights.unsqueeze(top_k_weights.dims().len());
+    weights_exp.shape.expand(down_out.dims());
+    let out = (down_out * weights_exp).sum(n - 1).output();
+    cx.set_dim('s', SEQ);
+
+    let report = CudaSearchEquivalenceFuzzer::new(&mut cx, &stream)
+        .seed(0x51A7_E5ED)
+        .samples(SEARCH_EQUIV_SAMPLES)
+        .generation_size(8)
+        .mutations(3)
+        .build_options(BuildSearchSpaceOptions::new().max_memory_mib(512))
+        .native_reference()
+        .input_f32(input.id, random_f32_vec(SEQ * HIDDEN, 301, -0.15, 0.15))
+        .input_f32(
+            router.id,
+            random_f32_vec(NUM_EXPERTS * HIDDEN, 302, -0.2, 0.2),
+        )
+        .input_bf16(
+            gate_up_weights.id,
+            random_bf16_vec(NUM_EXPERTS * MOE_INTERMEDIATE * 2 * HIDDEN, 303, -0.1, 0.1),
+        )
+        .input_bf16(
+            down_weights.id,
+            random_bf16_vec(NUM_EXPERTS * HIDDEN * MOE_INTERMEDIATE, 304, -0.1, 0.1),
+        )
+        .output_f32(out.id, "qwen_swiglu_moe_native_reference", 6e-2, 6e-2)
+        .run();
+    eprintln!("moe native-reference fuzz report: {report:?}");
+}

crates/luminal_cuda_lite/src/tests/utilities.rs

@@ -1,10 +1,15 @@
 use candle_core::{Device, Tensor, WithDType};
 use cudarc::driver::CudaContext;
 use half::{bf16, f16};
+use itertools::Itertools;
 use luminal::egglog_utils::{
-    egglog_to_llir, extract_generation, hash_choice_set, random_initial_choice, validate_choice_set,
+    EGraphChoiceSet, egglog_to_llir, extract_generation, hash_choice_set, random_initial_choice,
+    validate_choice_set,
+};
+use luminal::prelude::{
+    petgraph::{Direction, algo::toposort, visit::EdgeRef},
+    *,
 };
-use luminal::prelude::*;
 use num_traits::{Num, Signed};
 use rand::{Rng, SeedableRng, rngs::StdRng};
 use std::sync::Arc;
@@ -128,6 +133,498 @@ pub fn get_cuda_stream() -> Option<Arc<cudarc::driver::CudaStream>> {
     Some(ctx.default_stream())
 }
 
+#[derive(Debug, Clone)]
+pub enum CudaFuzzInput {
+    F32(NodeIndex, Vec<f32>),
+    Bf16(NodeIndex, Vec<bf16>),
+    I32(NodeIndex, Vec<i32>),
+}
+
+impl CudaFuzzInput {
+    fn apply(&self, rt: &mut CudaRuntime) {
+        match self {
+            Self::F32(id, data) => rt.set_data(*id, data.clone()),
+            Self::Bf16(id, data) => rt.set_data(*id, data.clone()),
+            Self::I32(id, data) => rt.set_data(*id, data.clone()),
+        }
+    }
+
+    fn apply_native(&self, rt: &mut NativeRuntime) {
+        match self {
+            Self::F32(id, data) => rt.set_data(*id, data.clone()),
+            Self::Bf16(id, data) => rt.set_data(*id, data.clone()),
+            Self::I32(id, data) => rt.set_data(*id, data.clone()),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct F32OutputCheck {
+    pub id: NodeIndex,
+    pub name: String,
+    pub rtol: f32,
+    pub atol: f32,
+}
+
+impl F32OutputCheck {
+    pub fn new(id: NodeIndex, name: impl Into<String>, rtol: f32, atol: f32) -> Self {
+        Self {
+            id,
+            name: name.into(),
+            rtol,
+            atol,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct SearchEquivalenceFuzzConfig {
+    pub seed: u64,
+    pub samples: usize,
+    pub generation_size: usize,
+    pub mutations: usize,
+    pub max_attempts: usize,
+    pub build_options: BuildSearchSpaceOptions,
+    pub reference: SearchEquivalenceReference,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SearchEquivalenceReference {
+    FirstCudaExtraction,
+    NativeRuntime,
+}
+
+impl Default for SearchEquivalenceFuzzConfig {
+    fn default() -> Self {
+        Self {
+            seed: 0,
+            samples: 32,
+            generation_size: 16,
+            mutations: 2,
+            max_attempts: 1_000,
+            build_options: BuildSearchSpaceOptions::default(),
+            reference: SearchEquivalenceReference::FirstCudaExtraction,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct SearchEquivalenceFuzzReport {
+    pub tested: usize,
+    pub skipped_invalid: usize,
+}
+
+struct ChoiceRun {
+    outputs: Vec<Vec<f32>>,
+    llir_summary: String,
+}
+
+pub struct CudaSearchEquivalenceFuzzer<'a> {
+    cx: &'a mut Graph,
+    stream: &'a Arc<cudarc::driver::CudaStream>,
+    inputs: Vec<CudaFuzzInput>,
+    outputs: Vec<F32OutputCheck>,
+    config: SearchEquivalenceFuzzConfig,
+}
+
+impl<'a> CudaSearchEquivalenceFuzzer<'a> {
+    pub fn new(cx: &'a mut Graph, stream: &'a Arc<cudarc::driver::CudaStream>) -> Self {
+        Self {
+            cx,
+            stream,
+            inputs: Vec::new(),
+            outputs: Vec::new(),
+            config: SearchEquivalenceFuzzConfig::default(),
+        }
+    }
+
+    pub fn seed(mut self, seed: u64) -> Self {
+        self.config.seed = seed;
+        self
+    }
+
+    pub fn samples(mut self, samples: usize) -> Self {
+        self.config.samples = samples;
+        self
+    }
+
+    pub fn generation_size(mut self, generation_size: usize) -> Self {
+        self.config.generation_size = generation_size;
+        self
+    }
+
+    pub fn mutations(mut self, mutations: usize) -> Self {
+        self.config.mutations = mutations;
+        self
+    }
+
+    pub fn build_options(mut self, build_options: BuildSearchSpaceOptions) -> Self {
+        self.config.build_options = build_options;
+        self
+    }
+
+    pub fn native_reference(mut self) -> Self {
+        self.config.reference = SearchEquivalenceReference::NativeRuntime;
+        self
+    }
+
+    pub fn input_f32(mut self, id: NodeIndex, data: Vec<f32>) -> Self {
+        self.inputs.push(CudaFuzzInput::F32(id, data));
+        self
+    }
+
+    pub fn input_bf16(mut self, id: NodeIndex, data: Vec<bf16>) -> Self {
+        self.inputs.push(CudaFuzzInput::Bf16(id, data));
+        self
+    }
+
+    pub fn input_i32(mut self, id: NodeIndex, data: Vec<i32>) -> Self {
+        self.inputs.push(CudaFuzzInput::I32(id, data));
+        self
+    }
+
+    pub fn output_f32(
+        mut self,
+        id: NodeIndex,
+        name: impl Into<String>,
+        rtol: f32,
+        atol: f32,
+    ) -> Self {
+        self.outputs.push(F32OutputCheck::new(id, name, rtol, atol));
+        self
+    }
+
+    pub fn run(self) -> SearchEquivalenceFuzzReport {
+        fuzz_cuda_search_space_equivalence(
+            self.cx,
+            self.stream,
+            &self.inputs,
+            &self.outputs,
+            self.config,
+        )
+    }
+}
+
+/// End-to-end search-space equivalence fuzzing for CUDA.
+///
+/// This builds the normal CUDA e-graph search space, extracts random selectable
+/// LLIR graphs, runs each with identical inputs, and verifies every requested
+/// f32 output matches the first valid extraction. The reference is intentionally
+/// another selected LLIR graph, not a hand-written CPU implementation: this
+/// catches cases where supposedly equivalent e-graph choices diverge, including
+/// candidates that produce non-finite outputs.
+pub fn fuzz_cuda_search_space_equivalence(
+    cx: &mut Graph,
+    stream: &Arc<cudarc::driver::CudaStream>,
+    inputs: &[CudaFuzzInput],
+    outputs: &[F32OutputCheck],
+    config: SearchEquivalenceFuzzConfig,
+) -> SearchEquivalenceFuzzReport {
+    assert!(
+        !outputs.is_empty(),
+        "fuzz harness needs at least one output"
+    );
+
+    let native_reference_outputs = if config.reference == SearchEquivalenceReference::NativeRuntime
+    {
+        cx.build_search_space::<NativeRuntime>();
+        let mut native_rng = StdRng::seed_from_u64(config.seed);
+        let mut native_rt = cx.search_options(
+            NativeRuntime::default(),
+            SearchOptions::new(1),
+            &mut native_rng,
+        );
+        for input in inputs {
+            input.apply_native(&mut native_rt);
+        }
+        native_rt.execute(&cx.dyn_map);
+        Some(
+            outputs
+                .iter()
+                .map(|out| native_rt.get_f32(out.id).clone())
+                .collect::<Vec<_>>(),
+        )
+    } else {
+        None
+    };
+
+    cx.build_search_space_with_options::<CudaRuntime>(config.build_options);
+
+    let egraph = cx.egraph().expect("search space should be built");
+    let ops = cx.egglog_ops().expect("search ops should be built");
+    let seed = if native_reference_outputs.is_some() {
+        config.seed.wrapping_add(0xC0DA_C0DA)
+    } else {
+        config.seed
+    };
+    let mut rng = StdRng::seed_from_u64(seed);
+    let mut prev_selected = FxHashSet::default();
+    let mut base = random_initial_choice(egraph, &mut rng);
+    prev_selected.insert(hash_choice_set(&base));
+
+    let mut skipped_invalid = 0usize;
+    let reference_is_cuda = native_reference_outputs.is_none();
+    let (reference_hash, reference_outputs, reference_llir_summary, mut tested) =
+        if let Some(reference_outputs) = native_reference_outputs {
+            (0, reference_outputs, None, 0usize)
+        } else {
+            let mut attempts = 0usize;
+            let (reference_hash, reference_run) = loop {
+                attempts += 1;
+                if attempts > config.max_attempts {
+                    panic!(
+                        "failed to extract a valid reference LLIR after {} attempts",
+                        config.max_attempts
+                    );
+                }
+                if validate_choice_set(egraph, &base, ops).is_err() {
+                    skipped_invalid += 1;
+                } else {
+                    let hash = hash_choice_set(&base);
+                    match run_choice_outputs(cx, stream, inputs, outputs, &base) {
+                        Ok(run) => break (hash, run),
+                        Err(err) => panic!("reference candidate hash={hash} failed: {err}"),
+                    }
+                }
+                base = random_initial_choice(egraph, &mut rng);
+                prev_selected.insert(hash_choice_set(&base));
+            };
+            (
+                reference_hash,
+                reference_run.outputs,
+                Some(reference_run.llir_summary),
+                1usize,
+            )
+        };
+
+    let mut attempts = 0usize;
+    while tested < config.samples && attempts < config.max_attempts {
+        attempts += 1;
+        let mut candidates = extract_generation(
+            egraph,
+            &base,
+            config.generation_size,
+            config.mutations,
+            &mut prev_selected,
+            &mut rng,
+        );
+        if candidates.is_empty() {
+            let next = random_initial_choice(egraph, &mut rng);
+            prev_selected.insert(hash_choice_set(&next));
+            candidates.push(next);
+        }
+
+        for candidate in candidates {
+            if tested >= config.samples {
+                break;
+            }
+            let candidate_hash = hash_choice_set(&candidate);
+            if reference_is_cuda && candidate_hash == reference_hash {
+                continue;
+            }
+            if validate_choice_set(egraph, &candidate, ops).is_err() {
+                skipped_invalid += 1;
+                continue;
+            }
+
+            let candidate_run = run_choice_outputs(cx, stream, inputs, outputs, &candidate)
+                .unwrap_or_else(|err| panic!("candidate hash={candidate_hash} failed: {err}"));
+            assert_fuzz_outputs_close(
+                outputs,
+                &reference_outputs,
+                &candidate_run.outputs,
+                &candidate_run.llir_summary,
+                reference_llir_summary.as_deref(),
+                reference_hash,
+                candidate_hash,
+            );
+            base = candidate;
+            tested += 1;
+        }
+    }
+
+    assert_eq!(
+        tested, config.samples,
+        "only tested {tested}/{} LLIR samples before exhausting attempts",
+        config.samples
+    );
+    SearchEquivalenceFuzzReport {
+        tested,
+        skipped_invalid,
+    }
+}
+
+fn run_choice_outputs<'a>(
+    cx: &'a Graph,
+    stream: &Arc<cudarc::driver::CudaStream>,
+    inputs: &[CudaFuzzInput],
+    outputs: &[F32OutputCheck],
+    choices: &EGraphChoiceSet<'a>,
+) -> Result<ChoiceRun, String> {
+    let egraph = cx.egraph().ok_or("search space was not built")?;
+    let ops = cx.egglog_ops().ok_or("search ops were not built")?;
+    let mut list_cache = FxHashMap::default();
+    let mut expr_cache = FxHashMap::default();
+    let mut llir_graph = egglog_to_llir(
+        egraph,
+        choices.clone(),
+        ops,
+        &cx.custom_ops,
+        &mut list_cache,
+        &mut expr_cache,
+        None,
+    );
+    unroll_loops_in_llir(&mut llir_graph);
+    let llir_summary = summarize_llir(&llir_graph);
+
+    let mut rt = CudaRuntime::initialize(stream.clone());
+    rt.load_llir(&llir_graph);
+    rt.preserve_intermediate_buffers_for_debug();
+    for input in inputs {
+        input.apply(&mut rt);
+    }
+    if std::env::var_os("LUMINAL_FUZZ_DUMP_LAST_LLIR").is_some() {
+        let _ = std::fs::write("/tmp/luminal_fuzz_last_candidate_llir.txt", &llir_summary);
+    }
+    rt.execute(&cx.dyn_map);
+    let topo_order = toposort(&llir_graph, None).map_err(|cycle| {
+        format!(
+            "extracted LLIR contains cycle at node {:?}",
+            cycle.node_id()
+        )
+    })?;
+    if let Some(report) = rt.first_nonfinite_f32_buffer_in_nodes(topo_order) {
+        let dump_path = "/tmp/luminal_fuzz_bad_candidate_llir.txt";
+        let _ = std::fs::write(dump_path, &llir_summary);
+        let op = llir_graph
+            .node_weight(report.node)
+            .map(|op| format!("{op:?}"))
+            .unwrap_or_else(|| "unknown op".to_string());
+        return Err(format!(
+            "LLIR produced non-finite F32 buffer node={} index={} value={} op={}; llir={dump_path}",
+            report.node.index(),
+            report.index,
+            report.value,
+            op
+        ));
+    }
+
+    let values = outputs
+        .iter()
+        .map(|out| rt.get_f32(out.id))
+        .collect::<Vec<_>>();
+    for (spec, values) in outputs.iter().zip(&values) {
+        if let Some((idx, value)) = values
+            .iter()
+            .enumerate()
+            .find(|(_, value)| !value.is_finite())
+        {
+            let dump_path = "/tmp/luminal_fuzz_bad_candidate_llir.txt";
+            let _ = std::fs::write(dump_path, &llir_summary);
+            let internal = rt
+                .first_nonfinite_f32_buffer()
+                .map(|report| {
+                    let op = llir_graph
+                        .node_weight(report.node)
+                        .map(|op| format!("{op:?}"))
+                        .unwrap_or_else(|| "unknown op".to_string());
+                    format!(
+                        "; first observed non-finite buffer node={} index={} value={} op={}",
+                        report.node.index(),
+                        report.index,
+                        report.value,
+                        op
+                    )
+                })
+                .unwrap_or_default();
+            return Err(format!(
+                "output {} produced non-finite value {value} at index {idx}{internal}; llir={dump_path}",
+                spec.name
+            ));
+        }
+    }
+    Ok(ChoiceRun {
+        outputs: values,
+        llir_summary,
+    })
+}
+
+fn assert_fuzz_outputs_close(
+    outputs: &[F32OutputCheck],
+    expected: &[Vec<f32>],
+    actual: &[Vec<f32>],
+    candidate_llir_summary: &str,
+    reference_llir_summary: Option<&str>,
+    reference_hash: u64,
+    candidate_hash: u64,
+) {
+    for ((spec, expected), actual) in outputs.iter().zip(expected.iter()).zip(actual.iter()) {
+        assert_eq!(
+            expected.len(),
+            actual.len(),
+            "output {} length mismatch for candidate hash={candidate_hash} reference hash={reference_hash}",
+            spec.name
+        );
+        let mut max_abs = 0.0f32;
+        let mut max_rel = 0.0f32;
+        let mut worst = 0usize;
+        for (i, (&a, &b)) in actual.iter().zip(expected.iter()).enumerate() {
+            assert!(
+                a.is_finite(),
+                "output {} candidate hash={candidate_hash} produced non-finite value {a} at index {i}",
+                spec.name
+            );
+            assert!(
+                b.is_finite(),
+                "output {} reference hash={reference_hash} produced non-finite value {b} at index {i}",
+                spec.name
+            );
+            let abs = (a - b).abs();
+            let rel = abs / b.abs().max(1e-12);
+            if abs > max_abs {
+                max_abs = abs;
+                max_rel = rel;
+                worst = i;
+            }
+            if abs > spec.atol + spec.rtol * b.abs() {
+                let dump_path = "/tmp/luminal_fuzz_bad_candidate_llir.txt";
+                let _ = std::fs::write(dump_path, candidate_llir_summary);
+                if let Some(reference_llir_summary) = reference_llir_summary {
+                    let _ = std::fs::write(
+                        "/tmp/luminal_fuzz_bad_reference_llir.txt",
+                        reference_llir_summary,
+                    );
+                }
+                panic!(
+                    "output {} mismatch candidate hash={candidate_hash} reference hash={reference_hash} index={i} actual={a} expected={b} abs={abs} rel={rel} tolerance={} candidate_llir={dump_path}",
+                    spec.name,
+                    spec.atol + spec.rtol * b.abs()
+                );
+            }
+        }
+        eprintln!(
+            "fuzz output {} ok: candidate hash={candidate_hash} max_abs={max_abs} max_rel={max_rel} worst={worst}",
+            spec.name
+        );
+    }
+}
+
+fn summarize_llir(llir_graph: &LLIRGraph) -> String {
+    llir_graph
+        .node_indices()
+        .map(|idx| {
+            let inputs = llir_graph
+                .edges_directed(idx, Direction::Incoming)
+                .sorted_by_key(|edge| edge.id())
+                .map(|edge| edge.source().index().to_string())
+                .collect::<Vec<_>>()
+                .join(", ");
+            format!("{} <- [{}]: {:?}", idx.index(), inputs, &llir_graph[idx])
+        })
+        .collect::<Vec<_>>()
+        .join("\n")
+}
+
 /// Get the GPU compute capability as (major, minor).
 pub fn gpu_compute_cap() -> Option<(i32, i32)> {
     let ctx = CudaContext::new(0).ok()?;

crates/luminal_metal/Cargo.toml

@@ -1,22 +1,32 @@
 [package]
 name = "luminal_metal"
 version = "0.2.0"
-edition = "2021"
+edition = "2024"
 description = "Metal backend for luminal"
 license = "MIT OR Apache-2.0"
 
 [dependencies]
 luminal = { path = "../.." }
-metal = "0.31"
+metal = { version = "0.31", features = ["mps"] }
 objc = "0.2"
 as-any = "0.3.2"
 itertools = "0.12.1"
-half = "2.7.1"
+half = { version = "2.7.1", features = ["bytemuck"] }
 tracing = "0.1.43"
+safetensors = "0.7.0"
+memmap2 = "0.9.9"
+bytemuck = "1.24.0"
 
 [dev-dependencies]
 candle-core = "0.9.2-alpha.1"
+hf-hub = { version = "0.4", default-features = false, features = ["rustls-tls", "ureq"] }
+luminal_nn = { path = "../luminal_nn" }
+luminal_tracing = { path = "../luminal_tracing" }
 proptest = "1.9.0"
+rand = "0.9.2"
+rustc-hash = "2.1"
+tokenizers = "0.22.2"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 
 [lints.rust]
 unexpected_cfgs = { level = "warn", check-cfg = ['cfg(feature, values("cargo-clippy"))'] }

crates/luminal_metal/examples/llama_1b.rs

@@ -0,0 +1,641 @@
+use hf_hub::api::sync::Api;
+use luminal::{
+    dtype::DType,
+    graph::{BuildSearchSpaceOptions, DimBucket, Graph},
+    prelude::{F32Pow, GraphTensor, Runtime},
+};
+use luminal_metal::MetalRuntime;
+use luminal_nn::{LayerNorm, gather_rows, scatter_rows};
+use luminal_tracing::luminal_filter;
+use rustc_hash::FxHashSet;
+use std::{
+    error::Error,
+    io::Write,
+    path::PathBuf,
+    time::{Duration, Instant},
+};
+use tokenizers::Tokenizer;
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+
+const REPO_ID: &str = "unsloth/Llama-3.2-1B-Instruct";
+const MAX_SEQ_LEN: usize = 2048;
+const GEN_TOKENS: usize = 96;
+const SEARCH_GRAPHS: usize = 100;
+const SEARCH_MEMORY_MIB: usize = 1536;
+const PROMPT: &str = "In one short paragraph, explain neural networks using the words layers, neurons, learning, and data.";
+
+const LAYERS: usize = 16;
+const HIDDEN: usize = 2048;
+const INTERMEDIATE: usize = 8192;
+const HEAD_DIM: usize = 64;
+const N_HEADS: usize = 32;
+const N_KV_HEADS: usize = 8;
+const KV_GROUPS: usize = N_HEADS / N_KV_HEADS;
+const KV_DIM: usize = N_KV_HEADS * HEAD_DIM;
+const VOCAB_SIZE: usize = 128256;
+const RMS_NORM_EPS: f32 = 1e-5;
+const ROPE_THETA: f32 = 500_000.0;
+const EOS_TOKEN: u32 = 128009;
+const STOP_TOKEN: u32 = 128001;
+
+fn prepare_hf_model() -> Result<PathBuf, Box<dyn Error>> {
+    let repo = Api::new()?.model(REPO_ID.to_string());
+    let tokenizer_path = repo.get("tokenizer.json")?;
+    repo.get("model.safetensors")?;
+    Ok(tokenizer_path.parent().unwrap().to_path_buf())
+}
+
+fn llama3_chat_prompt(user_prompt: &str) -> String {
+    format!(
+        "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{user_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+    )
+}
+
+#[derive(Default, Clone)]
+struct StepProfile {
+    total: Duration,
+    execute: Duration,
+    get_logits: Duration,
+    cache_roundtrip: Duration,
+}
+
+fn avg_ms(duration: Duration, n: usize) -> f64 {
+    if n == 0 {
+        0.0
+    } else {
+        duration.as_secs_f64() * 1e3 / n as f64
+    }
+}
+
+fn sample_greedy(logits_row: &[f32], seen: &FxHashSet<u32>, repetition_penalty: f32) -> u32 {
+    let mut row = logits_row.to_vec();
+    for &tok in seen {
+        let logit = &mut row[tok as usize];
+        if *logit > 0.0 {
+            *logit /= repetition_penalty;
+        } else {
+            *logit *= repetition_penalty;
+        }
+    }
+    row.iter()
+        .enumerate()
+        .max_by(|(_, a), (_, b)| a.total_cmp(b))
+        .unwrap()
+        .0 as u32
+}
+
+fn causal_mask(q_pos: &[usize], context_len: usize) -> Vec<f32> {
+    let mut mask = vec![-1e10f32; q_pos.len() * context_len];
+    for (qi, &pos) in q_pos.iter().enumerate() {
+        for ci in 0..context_len {
+            if ci <= pos {
+                mask[qi * context_len + ci] = 0.0;
+            }
+        }
+    }
+    mask
+}
+
+struct KVCache {
+    k_caches: Vec<GraphTensor>,
+    v_caches: Vec<GraphTensor>,
+}
+
+impl KVCache {
+    fn new(cx: &mut Graph, num_slots: usize) -> Self {
+        let mut k_caches = Vec::with_capacity(LAYERS);
+        let mut v_caches = Vec::with_capacity(LAYERS);
+        for l in 0..LAYERS {
+            k_caches.push(
+                cx.named_tensor(format!("kv_cache.{l}.k"), (num_slots, KV_DIM))
+                    .persist(),
+            );
+            v_caches.push(
+                cx.named_tensor(format!("kv_cache.{l}.v"), (num_slots, KV_DIM))
+                    .persist(),
+            );
+        }
+        Self { k_caches, v_caches }
+    }
+}
+
+struct Llama {
+    embedding: GraphTensor,
+    layers: Vec<LlamaLayer>,
+    lm_norm: LayerNorm,
+}
+
+impl Llama {
+    fn init(cx: &mut Graph) -> Self {
+        let mut layers = Vec::with_capacity(LAYERS);
+        for l in 0..LAYERS {
+            layers.push(LlamaLayer {
+                up: cx
+                    .named_tensor(
+                        format!("model.layers.{l}.mlp.up_proj.weight"),
+                        (INTERMEDIATE, HIDDEN),
+                    )
+                    .persist(),
+                gate: cx
+                    .named_tensor(
+                        format!("model.layers.{l}.mlp.gate_proj.weight"),
+                        (INTERMEDIATE, HIDDEN),
+                    )
+                    .persist(),
+                down: cx
+                    .named_tensor(
+                        format!("model.layers.{l}.mlp.down_proj.weight"),
+                        (HIDDEN, INTERMEDIATE),
+                    )
+                    .persist(),
+                q_proj: cx
+                    .named_tensor(
+                        format!("model.layers.{l}.self_attn.q_proj.weight"),
+                        (HIDDEN, HIDDEN),
+                    )
+                    .persist(),
+                k_proj: cx
+                    .named_tensor(
+                        format!("model.layers.{l}.self_attn.k_proj.weight"),
+                        (KV_DIM, HIDDEN),
+                    )
+                    .persist(),
+                v_proj: cx
+                    .named_tensor(
+                        format!("model.layers.{l}.self_attn.v_proj.weight"),
+                        (KV_DIM, HIDDEN),
+                    )
+                    .persist(),
+                o_proj: cx
+                    .named_tensor(
+                        format!("model.layers.{l}.self_attn.o_proj.weight"),
+                        (HIDDEN, HIDDEN),
+                    )
+                    .persist(),
+                attn_rms: LayerNorm::new(
+                    HIDDEN,
+                    Some(&format!("model.layers.{l}.input_layernorm.weight")),
+                    None,
+                    false,
+                    RMS_NORM_EPS,
+                    cx,
+                ),
+                mlp_rms: LayerNorm::new(
+                    HIDDEN,
+                    Some(&format!("model.layers.{l}.post_attention_layernorm.weight")),
+                    None,
+                    false,
+                    RMS_NORM_EPS,
+                    cx,
+                ),
+            });
+        }
+
+        Self {
+            embedding: cx
+                .named_tensor("model.embed_tokens.weight", (VOCAB_SIZE, HIDDEN))
+                .persist(),
+            layers,
+            lm_norm: LayerNorm::new(
+                HIDDEN,
+                Some("model.norm.weight"),
+                None,
+                false,
+                RMS_NORM_EPS,
+                cx,
+            ),
+        }
+    }
+
+    fn forward(
+        &self,
+        input: GraphTensor,
+        q_pos: GraphTensor,
+        scatter_idx: GraphTensor,
+        gather_idx: GraphTensor,
+        attn_mask: GraphTensor,
+        kv_cache: &KVCache,
+    ) -> (GraphTensor, Vec<(GraphTensor, GraphTensor)>) {
+        let seq = input.dims1();
+        let mut x = self.embedding.gather(
+            (input * HIDDEN).expand_dim(1, HIDDEN)
+                + input.graph().arange(HIDDEN).expand_dim(0, seq),
+        );
+        let mut cache_outputs = Vec::with_capacity(LAYERS);
+        for (i, layer) in self.layers.iter().enumerate() {
+            let (x_new, k_out, v_out) = layer.forward(
+                x,
+                q_pos,
+                scatter_idx,
+                gather_idx,
+                attn_mask,
+                kv_cache.k_caches[i],
+                kv_cache.v_caches[i],
+            );
+            x = x_new;
+            cache_outputs.push((k_out, v_out));
+        }
+
+        let logits = self.lm_norm.forward(x).matmul(self.embedding.t());
+        (logits, cache_outputs)
+    }
+}
+
+struct LlamaLayer {
+    up: GraphTensor,
+    gate: GraphTensor,
+    down: GraphTensor,
+    q_proj: GraphTensor,
+    k_proj: GraphTensor,
+    v_proj: GraphTensor,
+    o_proj: GraphTensor,
+    attn_rms: LayerNorm,
+    mlp_rms: LayerNorm,
+}
+
+fn llama_rotary_embeddings(mut input: GraphTensor, pos_ids: GraphTensor) -> GraphTensor {
+    input = input.split_dims(1, HEAD_DIM).transpose(0, 1);
+
+    let freqs = input
+        .graph()
+        .arange_options(0, HEAD_DIM, 2)
+        .cast(DType::F32)
+        / HEAD_DIM as f32;
+    let inv_freqs = ROPE_THETA.pow(freqs).reciprocal();
+    let emb = pos_ids
+        .cast(DType::F32)
+        .expand_dim(1, 1)
+        .matmul(inv_freqs.expand_dim(0, 1));
+
+    let x0 = input.slice((.., .., ..HEAD_DIM / 2));
+    let x1 = input.slice((.., .., HEAD_DIM / 2..));
+
+    let cos = emb.cos().expand_dim(0, x0.dims()[0]);
+    let sin = emb.sin().expand_dim(0, x0.dims()[0]);
+    let x0_out = x0 * cos - x1 * sin;
+    let x1_out = x1 * cos + x0 * sin;
+
+    x0_out
+        .concat_along(x1_out, 2)
+        .transpose(0, 1)
+        .merge_dims(1, 2)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn attention(
+    q_rope: GraphTensor,
+    k_rope: GraphTensor,
+    v: GraphTensor,
+    k_cache: GraphTensor,
+    v_cache: GraphTensor,
+    scatter_idx: GraphTensor,
+    gather_idx: GraphTensor,
+    attn_mask: GraphTensor,
+) -> (GraphTensor, GraphTensor, GraphTensor) {
+    let k_cache_out = scatter_rows(k_rope, scatter_idx, k_cache, KV_DIM);
+    let v_cache_out = scatter_rows(v, scatter_idx, v_cache, KV_DIM);
+
+    let k = gather_rows(k_cache_out, gather_idx, KV_DIM);
+    let v_ctx = gather_rows(v_cache_out, gather_idx, KV_DIM);
+
+    let q = (q_rope * 1.0).split_dims(1, HEAD_DIM).transpose(0, 1);
+    let k = k.split_dims(1, HEAD_DIM).permute((1, 2, 0));
+    let v_ctx = v_ctx.split_dims(1, HEAD_DIM).transpose(0, 1);
+
+    let k = k.expand_dim(1, KV_GROUPS).merge_dims(0, 1) * 1.0;
+    let v_ctx = v_ctx.expand_dim(1, KV_GROUPS).merge_dims(0, 1) * 1.0;
+
+    let scores = q.matmul(k) / (HEAD_DIM as f32).sqrt();
+    let masked_scores = scores + attn_mask.expand_dim(0, N_HEADS);
+    let weights = masked_scores.softmax(2);
+    let out = weights.matmul(v_ctx);
+    let attn_out = out.transpose(0, 1).merge_dims(1, 2);
+
+    (attn_out, k_cache_out, v_cache_out)
+}
+
+impl LlamaLayer {
+    #[allow(clippy::too_many_arguments)]
+    fn forward(
+        &self,
+        mut x: GraphTensor,
+        q_pos: GraphTensor,
+        scatter_idx: GraphTensor,
+        gather_idx: GraphTensor,
+        attn_mask: GraphTensor,
+        k_cache: GraphTensor,
+        v_cache: GraphTensor,
+    ) -> (GraphTensor, GraphTensor, GraphTensor) {
+        let x_attn = self.attn_rms.forward(x);
+        let q = x_attn.matmul(self.q_proj.t());
+        let k = x_attn.matmul(self.k_proj.t());
+        let v = x_attn.matmul(self.v_proj.t());
+
+        let q_rope = llama_rotary_embeddings(q, q_pos);
+        let k_rope = llama_rotary_embeddings(k, q_pos);
+        let (attn_out, k_cache_out, v_cache_out) = attention(
+            q_rope,
+            k_rope,
+            v,
+            k_cache,
+            v_cache,
+            scatter_idx,
+            gather_idx,
+            attn_mask,
+        );
+        x += attn_out.matmul(self.o_proj.t());
+
+        let x_mlp = self.mlp_rms.forward(x);
+        let mlp_out =
+            (x_mlp.matmul(self.gate.t()).swish() * x_mlp.matmul(self.up.t())).matmul(self.down.t());
+        (x + mlp_out, k_cache_out, v_cache_out)
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn run_model_step(
+    cx: &mut Graph,
+    runtime: &mut MetalRuntime,
+    input: GraphTensor,
+    q_pos_t: GraphTensor,
+    scatter_idx_t: GraphTensor,
+    gather_idx_t: GraphTensor,
+    attn_mask_t: GraphTensor,
+    logits: GraphTensor,
+    kv_cache: &KVCache,
+    cache_outputs: &[(GraphTensor, GraphTensor)],
+    tokens: &[u32],
+    q_pos: &[i32],
+    scatter_idx: &[i32],
+    gather_idx: &[i32],
+    attn_mask: &[f32],
+) -> (Vec<f32>, StepProfile) {
+    let start = Instant::now();
+    cx.set_dim('s', tokens.len());
+    cx.set_dim('c', gather_idx.len());
+
+    runtime.set_data(input, tokens.iter().map(|t| *t as i32).collect::<Vec<_>>());
+    runtime.set_data(q_pos_t, q_pos.to_vec());
+    runtime.set_data(scatter_idx_t, scatter_idx.to_vec());
+    runtime.set_data(gather_idx_t, gather_idx.to_vec());
+    runtime.set_data(attn_mask_t, attn_mask.to_vec());
+    runtime.allocate_intermediate_buffers(&cx.dyn_map);
+
+    let execute_start = Instant::now();
+    runtime.execute(&cx.dyn_map);
+    let execute = execute_start.elapsed();
+
+    let logits_start = Instant::now();
+    let logits_data = runtime.get_f32(logits);
+    let get_logits = logits_start.elapsed();
+
+    let cache_start = Instant::now();
+    for (layer_idx, (k_out, v_out)) in cache_outputs.iter().enumerate() {
+        let k_buf = runtime.remove_buffer(*k_out);
+        let v_buf = runtime.remove_buffer(*v_out);
+        runtime.set_buffer(kv_cache.k_caches[layer_idx], k_buf);
+        runtime.set_buffer(kv_cache.v_caches[layer_idx], v_buf);
+    }
+    let cache_roundtrip = cache_start.elapsed();
+
+    (
+        logits_data,
+        StepProfile {
+            total: start.elapsed(),
+            execute,
+            get_logits,
+            cache_roundtrip,
+        },
+    )
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+    let _ = tracing_subscriber::registry()
+        .with(tracing_subscriber::fmt::layer())
+        .with(luminal_filter())
+        .try_init();
+
+    let model_dir = prepare_hf_model()?;
+    println!("Using model directory: {}", model_dir.display());
+
+    let tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))
+        .map_err(|err| err as Box<dyn Error>)?;
+    let prompt_tokens = tokenizer
+        .encode(llama3_chat_prompt(PROMPT), false)
+        .map_err(|err| err as Box<dyn Error>)?
+        .get_ids()
+        .to_vec();
+
+    let mut cx = Graph::default();
+    let input = cx.named_tensor("input", 's').as_dtype(DType::Int);
+    let q_pos_t = cx.named_tensor("q_pos", 's').as_dtype(DType::Int);
+    let scatter_idx_t = cx.named_tensor("scatter_idx", 's').as_dtype(DType::Int);
+    let gather_idx_t = cx.named_tensor("gather_idx", 'c').as_dtype(DType::Int);
+    let attn_mask_t = cx.named_tensor("attn_mask", ('s', 'c'));
+    let kv_cache = KVCache::new(&mut cx, MAX_SEQ_LEN);
+    let (logits, cache_outputs) = Llama::init(&mut cx).forward(
+        input,
+        q_pos_t,
+        scatter_idx_t,
+        gather_idx_t,
+        attn_mask_t,
+        &kv_cache,
+    );
+    let logits = logits.output();
+    for (k_out, v_out) in &cache_outputs {
+        k_out.output();
+        v_out.output();
+    }
+
+    cx.set_dim('s', 1);
+    cx.set_dim('c', 1);
+
+    println!("Building E-Graph...");
+    let egraph_start = Instant::now();
+    cx.build_search_space_with_options::<MetalRuntime>(
+        BuildSearchSpaceOptions::new().max_memory_mib(SEARCH_MEMORY_MIB),
+    );
+    println!(
+        "  E-Graph build: {:.2} s",
+        egraph_start.elapsed().as_secs_f64()
+    );
+
+    println!("Loading weights...");
+    let load_start = Instant::now();
+    let mut runtime = MetalRuntime::initialize(());
+    runtime.load_safetensors(&cx, model_dir.join("model.safetensors").to_str().unwrap());
+    println!("  Weight load: {:.2} s", load_start.elapsed().as_secs_f64());
+
+    let cache_bytes = MAX_SEQ_LEN * KV_DIM * std::mem::size_of::<f32>();
+    for i in 0..LAYERS {
+        runtime.set_zeros(kv_cache.k_caches[i], cache_bytes);
+        runtime.set_zeros(kv_cache.v_caches[i], cache_bytes);
+    }
+
+    println!("Compiling...");
+    let compile_start = Instant::now();
+    let max_prefill = (prompt_tokens.len() + 16)
+        .next_power_of_two()
+        .min(MAX_SEQ_LEN);
+    let max_context = (prompt_tokens.len() + GEN_TOKENS + 1)
+        .next_power_of_two()
+        .min(MAX_SEQ_LEN);
+    let search_s = 16.min(max_prefill).max(2);
+    let search_c = 16.min(max_context).max(2);
+    cx.set_dim_buckets(
+        's',
+        &[
+            DimBucket::new(1, 1),
+            DimBucket::new(2, max_prefill).representative(search_s),
+        ],
+    );
+    cx.set_dim_buckets(
+        'c',
+        &[
+            DimBucket::new(1, 1),
+            DimBucket::new(2, max_context).representative(search_c),
+        ],
+    );
+    cx.set_dim('s', search_s);
+    cx.set_dim('c', search_c);
+    runtime.set_data(input, vec![1; search_s]);
+    runtime.set_data(q_pos_t, (0..search_s as i32).collect::<Vec<_>>());
+    runtime.set_data(scatter_idx_t, (0..search_s as i32).collect::<Vec<_>>());
+    runtime.set_data(gather_idx_t, (0..search_c as i32).collect::<Vec<_>>());
+    runtime.set_data(attn_mask_t, vec![0.0f32; search_s * search_c]);
+    runtime = cx.search(runtime, SEARCH_GRAPHS);
+    println!(
+        "  Search/compile: {:.2} s",
+        compile_start.elapsed().as_secs_f64()
+    );
+
+    for i in 0..LAYERS {
+        runtime.set_zeros(kv_cache.k_caches[i], cache_bytes);
+        runtime.set_zeros(kv_cache.v_caches[i], cache_bytes);
+    }
+
+    let prompt_len = prompt_tokens.len();
+    let mut context_len = 0usize;
+    let mut profiles = Vec::new();
+    let mut seen_tokens = FxHashSet::default();
+    let repetition_penalty = 1.05;
+
+    println!(
+        "Prompt: {} tokens, generating up to {} tokens",
+        prompt_len, GEN_TOKENS
+    );
+
+    let mut generated = 0usize;
+    let mut next_token = None;
+    if GEN_TOKENS > 0 && prompt_len > 0 {
+        let positions: Vec<usize> = (0..prompt_len).collect();
+        let q_pos: Vec<i32> = positions.iter().map(|&p| p as i32).collect();
+        let mask = causal_mask(&positions, prompt_len);
+        let (logits_data, profile) = run_model_step(
+            &mut cx,
+            &mut runtime,
+            input,
+            q_pos_t,
+            scatter_idx_t,
+            gather_idx_t,
+            attn_mask_t,
+            logits,
+            &kv_cache,
+            &cache_outputs,
+            &prompt_tokens,
+            &q_pos,
+            &q_pos,
+            &q_pos,
+            &mask,
+        );
+        context_len = prompt_len;
+
+        let token = sample_greedy(
+            &logits_data[logits_data.len() - VOCAB_SIZE..],
+            &seen_tokens,
+            repetition_penalty,
+        );
+        seen_tokens.insert(token);
+        next_token = Some(token);
+        generated = 1;
+        profiles.push(profile);
+
+        if token != EOS_TOKEN && token != STOP_TOKEN {
+            print!(
+                "{}",
+                tokenizer
+                    .decode(&[token], true)
+                    .map_err(|err| err as Box<dyn Error>)?
+            );
+            std::io::stdout().flush()?;
+        }
+    }
+
+    while generated < GEN_TOKENS {
+        let current_token = match next_token {
+            Some(token) if token != EOS_TOKEN && token != STOP_TOKEN => token,
+            _ => break,
+        };
+        let gather_idx = (0..=context_len as i32).collect::<Vec<_>>();
+        let mask = causal_mask(&[context_len], context_len + 1);
+        let (logits_data, profile) = run_model_step(
+            &mut cx,
+            &mut runtime,
+            input,
+            q_pos_t,
+            scatter_idx_t,
+            gather_idx_t,
+            attn_mask_t,
+            logits,
+            &kv_cache,
+            &cache_outputs,
+            &[current_token],
+            &[context_len as i32],
+            &[context_len as i32],
+            &gather_idx,
+            &mask,
+        );
+        context_len += 1;
+
+        let token = sample_greedy(
+            &logits_data[logits_data.len() - VOCAB_SIZE..],
+            &seen_tokens,
+            repetition_penalty,
+        );
+        seen_tokens.insert(token);
+        next_token = Some(token);
+        generated += 1;
+        profiles.push(profile);
+
+        if token == EOS_TOKEN || token == STOP_TOKEN {
+            break;
+        }
+        print!(
+            "{}",
+            tokenizer
+                .decode(&[token], true)
+                .map_err(|err| err as Box<dyn Error>)?
+        );
+        std::io::stdout().flush()?;
+    }
+    println!();
+
+    let ttft = profiles.first().map(|p| p.total).unwrap_or_default();
+    let decode_steps = profiles.len().saturating_sub(1);
+    let decode_total: Duration = profiles.iter().skip(1).map(|p| p.total).sum();
+    println!("  TTFT: {:.2} ms", ttft.as_secs_f64() * 1e3);
+    println!("  TPOT: {:.2} ms", avg_ms(decode_total, decode_steps));
+
+    let execute_total: Duration = profiles.iter().map(|p| p.execute).sum();
+    let logits_total: Duration = profiles.iter().map(|p| p.get_logits).sum();
+    let cache_total: Duration = profiles.iter().map(|p| p.cache_roundtrip).sum();
+    println!(
+        "  Profile: n={}, exec={:.2} ms, logits={:.2} ms, cache={:.2} ms",
+        profiles.len(),
+        avg_ms(execute_total, profiles.len()),
+        avg_ms(logits_total, profiles.len()),
+        avg_ms(cache_total, profiles.len()),
+    );
+
+    Ok(())
+}

crates/luminal_metal/src/dyn_backend.rs

@@ -1,7 +1,7 @@
 //! [`DynBackend`] implementation for the Metal runtime.
 
 use luminal::dtype::DType;
-use luminal::dyn_backend::{bytes_to_native_data, compile_backend, BackendCompileArgs, DynBackend};
+use luminal::dyn_backend::{BackendCompileArgs, DynBackend, bytes_to_native_data, compile_backend};
 use luminal::prelude::*;
 
 use crate::runtime::MetalRuntime;

crates/luminal_metal/src/kernel/matmul.rs

@@ -1,227 +1,5 @@
-use super::{MetalMulInfo, MetalSumReduceInfo};
-use luminal::prelude::*;
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
-pub enum MetalMatmulFamily {
-    #[default]
-    Naive,
-    RegularTiled,
-}
-
-#[derive(Debug, Clone)]
-pub struct MatmulDescriptor {
-    pub m: Expression,
-    pub n: Expression,
-    pub k: Expression,
-    pub batch_shape: Vec<Expression>,
-    pub lhs_strides: Vec<Expression>,
-    pub rhs_strides: Vec<Expression>,
-    pub out_strides: Vec<Expression>,
-    pub transpose_lhs: bool,
-    pub transpose_rhs: bool,
-}
-
-impl MatmulDescriptor {
-    pub fn from_mul_and_sum(
-        mul_info: &MetalMulInfo,
-        sum_info: &MetalSumReduceInfo,
-    ) -> Option<Self> {
-        let zero = Expression::from(0);
-        let z = Expression::from('z');
-
-        let is_simple_2d_matmul = mul_info.shape.len() == 3
-            && sum_info.shape.len() == 2
-            && mul_info.a_strides.len() == 3
-            && mul_info.b_strides.len() == 3
-            && sum_info.strides.len() == 2
-            && mul_info.shape[0] == sum_info.shape[0]
-            && mul_info.shape[1] == sum_info.shape[1]
-            && mul_info.shape[2] == sum_info.iters
-            && mul_info.a_strides[1] == zero
-            && mul_info.a_strides[2] == z
-            && mul_info.b_strides[0] == zero
-            && mul_info.b_strides[1] == z
-            && sum_info.strides[1] == z
-            && sum_info.iter_stride == z;
-
-        if !is_simple_2d_matmul {
-            return None;
-        }
-
-        Some(Self {
-            m: sum_info.shape[0],
-            n: sum_info.shape[1],
-            k: sum_info.iters,
-            batch_shape: Vec::new(),
-            lhs_strides: mul_info.a_strides.clone(),
-            rhs_strides: mul_info.b_strides.clone(),
-            out_strides: sum_info.strides.clone(),
-            transpose_lhs: false,
-            transpose_rhs: false,
-        })
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct MatmulPlan {
-    pub family: MetalMatmulFamily,
-    pub m: Expression,
-    pub n: Expression,
-    pub k: Expression,
-    pub lda: Expression,
-    pub ldb: Expression,
-    pub ldd: Expression,
-    pub batch_size: u32,
-    pub batch_stride_a: u32,
-    pub batch_stride_b: u32,
-    pub batch_stride_d: u32,
-    pub bm: u16,
-    pub bn: u16,
-    pub bk: u16,
-    pub wm: u16,
-    pub wn: u16,
-}
-
-#[derive(Debug, Default, Clone, Copy)]
-pub struct MetalMatmulPlanner;
-
-impl MetalMatmulPlanner {
-    pub fn plan(&self, desc: &MatmulDescriptor) -> MatmulPlan {
-        let family = if desc.batch_shape.is_empty()
-            && desc.m.as_num().is_some_and(|m| m >= 32)
-            && desc.n.as_num().is_some_and(|n| n >= 32)
-            && desc.k.as_num().is_some_and(|k| k >= 32)
-        {
-            MetalMatmulFamily::RegularTiled
-        } else {
-            MetalMatmulFamily::Naive
-        };
-        MatmulPlan {
-            family,
-            m: desc.m,
-            n: desc.n,
-            k: desc.k,
-            lda: desc.lhs_strides[0],
-            ldb: desc.rhs_strides[2],
-            ldd: desc.out_strides[0],
-            batch_size: 1,
-            batch_stride_a: 0,
-            batch_stride_b: 0,
-            batch_stride_d: 0,
-            bm: 16,
-            bn: 16,
-            bk: 8,
-            wm: 2,
-            wn: 2,
-        }
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn descriptor_recovers_simple_2d_matmul() {
-        let mul = MetalMulInfo {
-            shape: vec![
-                Expression::from(4),
-                Expression::from(8),
-                Expression::from(16),
-            ],
-            a_strides: vec![
-                Expression::from('z') * 16,
-                Expression::from(0),
-                Expression::from('z'),
-            ],
-            b_strides: vec![
-                Expression::from(0),
-                Expression::from('z'),
-                Expression::from('z') * 8,
-            ],
-            output_strides: vec![
-                Expression::from('z') * 16,
-                Expression::from('z') * 8,
-                Expression::from('z'),
-            ],
-        };
-        let sum = MetalSumReduceInfo {
-            shape: vec![Expression::from(4), Expression::from(8)],
-            strides: vec![Expression::from('z') * 8, Expression::from('z')],
-            iters: Expression::from(16),
-            iter_stride: Expression::from('z'),
-        };
-
-        let desc = MatmulDescriptor::from_mul_and_sum(&mul, &sum).unwrap();
-        assert_eq!(desc.m, Expression::from(4));
-        assert_eq!(desc.n, Expression::from(8));
-        assert_eq!(desc.k, Expression::from(16));
-    }
-
-    #[test]
-    fn planner_keeps_small_problems_on_naive_path() {
-        let desc = MatmulDescriptor {
-            m: Expression::from(4),
-            n: Expression::from(8),
-            k: Expression::from(16),
-            batch_shape: Vec::new(),
-            lhs_strides: vec![
-                Expression::from('z') * 16,
-                Expression::from(0),
-                Expression::from('z'),
-            ],
-            rhs_strides: vec![
-                Expression::from(0),
-                Expression::from('z'),
-                Expression::from('z') * 8,
-            ],
-            out_strides: vec![Expression::from('z') * 8, Expression::from('z')],
-            transpose_lhs: false,
-            transpose_rhs: false,
-        };
-
-        let planner = MetalMatmulPlanner;
-        let plan = planner.plan(&desc);
-        assert_eq!(plan.family, MetalMatmulFamily::Naive);
-        assert_eq!(plan.bm, 16);
-        assert_eq!(plan.bn, 16);
-        assert_eq!(plan.bk, 8);
-        assert_eq!(plan.wm, 2);
-        assert_eq!(plan.wn, 2);
-        assert_eq!(plan.lda, Expression::from('z') * 16);
-        assert_eq!(plan.ldb, Expression::from('z') * 8);
-        assert_eq!(plan.ldd, Expression::from('z') * 8);
-    }
-
-    #[test]
-    fn planner_promotes_large_problems_to_regular_tiled() {
-        let desc = MatmulDescriptor {
-            m: Expression::from(64),
-            n: Expression::from(64),
-            k: Expression::from(64),
-            batch_shape: Vec::new(),
-            lhs_strides: vec![
-                Expression::from('z') * 64,
-                Expression::from(0),
-                Expression::from('z'),
-            ],
-            rhs_strides: vec![
-                Expression::from(0),
-                Expression::from('z'),
-                Expression::from('z') * 64,
-            ],
-            out_strides: vec![Expression::from('z') * 64, Expression::from('z')],
-            transpose_lhs: false,
-            transpose_rhs: false,
-        };
-
-        let planner = MetalMatmulPlanner;
-        let plan = planner.plan(&desc);
-        assert_eq!(plan.family, MetalMatmulFamily::RegularTiled);
-        assert_eq!(plan.bm, 16);
-        assert_eq!(plan.bn, 16);
-        assert_eq!(plan.bk, 8);
-        assert_eq!(plan.wm, 2);
-        assert_eq!(plan.wn, 2);
-    }
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum MPSMatrixLayout {
+    RowMajor,
+    TransposedRowMajor,
 }

crates/luminal_metal/src/kernel/mod.rs

@@ -6,10 +6,127 @@ pub use ops::*;
 use luminal::dtype::DType;
 use luminal::op::EgglogOp;
 use luminal::prelude::*;
-use metal::{Buffer, ComputeCommandEncoderRef, ComputePipelineState, Device};
+use metal::{
+    Buffer, CommandBufferRef, ComputeCommandEncoderRef, ComputePipelineState, Device,
+    foreign_types::ForeignTypeRef, mps,
+};
+use objc::rc::StrongPtr;
+use objc::runtime::Object;
+use objc::{class, msg_send, sel, sel_impl};
+use std::cell::RefCell;
 
 pub const DYN_SLOT_COUNT: usize = 26;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+struct MpsMatrixDescriptorKey {
+    rows: usize,
+    cols: usize,
+    row_bytes: u64,
+    data_type: isize,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+struct MpsMatmulKey {
+    transpose_lhs: bool,
+    transpose_rhs: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    alpha: u64,
+    beta: u64,
+}
+
+#[derive(Default)]
+pub struct MpsKernelCache {
+    matrix_descriptors: FxHashMap<MpsMatrixDescriptorKey, StrongPtr>,
+    matmul_kernels: FxHashMap<MpsMatmulKey, StrongPtr>,
+}
+
+impl MpsKernelCache {
+    pub(crate) fn matrix_descriptor(
+        &mut self,
+        rows: usize,
+        cols: usize,
+        row_bytes: u64,
+        dtype: DType,
+    ) -> *mut Object {
+        let key = MpsMatrixDescriptorKey {
+            rows,
+            cols,
+            row_bytes,
+            data_type: Self::mps_data_type(dtype),
+        };
+        let descriptor = self
+            .matrix_descriptors
+            .entry(key)
+            .or_insert_with(|| unsafe {
+                let descriptor: *mut Object = msg_send![
+                    class!(MPSMatrixDescriptor),
+                    matrixDescriptorWithRows: rows
+                    columns: cols
+                    rowBytes: row_bytes as usize
+                    dataType: key.data_type
+                ];
+                StrongPtr::retain(descriptor)
+            });
+        **descriptor
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn matrix_multiplication(
+        &mut self,
+        command_buffer: &CommandBufferRef,
+        transpose_lhs: bool,
+        transpose_rhs: bool,
+        m: usize,
+        n: usize,
+        k: usize,
+        alpha: f64,
+        beta: f64,
+    ) -> *mut Object {
+        let key = MpsMatmulKey {
+            transpose_lhs,
+            transpose_rhs,
+            m,
+            n,
+            k,
+            alpha: alpha.to_bits(),
+            beta: beta.to_bits(),
+        };
+        let kernel = self.matmul_kernels.entry(key).or_insert_with(|| unsafe {
+            let device: *mut Object = msg_send![command_buffer.as_ptr(), device];
+            let kernel: *mut Object = msg_send![class!(MPSMatrixMultiplication), alloc];
+            let kernel: *mut Object = msg_send![
+                kernel,
+                initWithDevice: device
+                transposeLeft: transpose_lhs
+                transposeRight: transpose_rhs
+                resultRows: m
+                resultColumns: n
+                interiorColumns: k
+                alpha: alpha
+                beta: beta
+            ];
+            StrongPtr::new(kernel)
+        });
+        **kernel
+    }
+
+    fn mps_data_type(dtype: DType) -> isize {
+        match dtype {
+            DType::F32 | DType::TF32 => mps::MPSDataType::Float32 as isize,
+            DType::F16 => mps::MPSDataType::Float16 as isize,
+            unsupported => panic!("MPSMatmul does not support dtype {unsupported:?}"),
+        }
+    }
+}
+
+pub struct MetalEncodeContext<'a> {
+    pub(crate) command_buffer: &'a CommandBufferRef,
+    pub(crate) dyn_buffer: &'a Buffer,
+    pub(crate) mps_cache: &'a RefCell<MpsKernelCache>,
+}
+
 #[derive(Debug, Clone)]
 pub struct MetalMulInfo {
     pub shape: Vec<Expression>,
@@ -32,7 +149,7 @@ pub trait MetalKernelOp: EgglogOp {
         device: &Device,
         input_dtypes: &[DType],
         output_dtype: DType,
-    ) -> ComputePipelineState;
+    ) -> Option<ComputePipelineState>;
 
     fn infer_output_dtype(&self, input_dtypes: &[DType]) -> DType {
         input_dtypes.first().copied().unwrap_or(DType::F32)
@@ -40,7 +157,7 @@ pub trait MetalKernelOp: EgglogOp {
 
     fn output_size(&self) -> Expression;
 
-    fn encode(
+    fn encode_compute(
         &self,
         encoder: &ComputeCommandEncoderRef,
         pipeline: &ComputePipelineState,
@@ -49,6 +166,25 @@ pub trait MetalKernelOp: EgglogOp {
         dyn_map: &FxHashMap<char, usize>,
     );
 
+    #[allow(clippy::too_many_arguments)]
+    fn encode(
+        &self,
+        context: &mut MetalEncodeContext<'_>,
+        pipeline: Option<&ComputePipelineState>,
+        inputs: &[&Buffer],
+        output: &Buffer,
+        dyn_map: &FxHashMap<char, usize>,
+        _input_dtypes: &[DType],
+        _output_dtype: DType,
+    ) {
+        let pipeline = pipeline.expect("compute pipeline not compiled");
+        let encoder = context.command_buffer.new_compute_command_encoder();
+        let dyn_idx = inputs.len() as u64 + 1;
+        encoder.set_buffer(dyn_idx, Some(context.dyn_buffer), 0);
+        self.encode_compute(encoder, pipeline, inputs, output, dyn_map);
+        encoder.end_encoding();
+    }
+
     // ========================================================================
     // Performance Metrics for MBU/MFU Calculation
     // ========================================================================
@@ -73,6 +209,10 @@ pub trait MetalKernelOp: EgglogOp {
         None
     }
 
+    fn output_aliases_input(&self) -> Option<usize> {
+        None
+    }
+
     fn is_matmul(&self) -> bool {
         false
     }

crates/luminal_metal/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod dyn_backend;
 pub mod kernel;
+mod memory_analysis;
 pub mod runtime;
 
 #[cfg(test)]

crates/luminal_metal/src/tests.rs

@@ -1,8 +1,17 @@
 use crate::{kernel::lower_expression_for_metal, runtime::MetalRuntime};
 use candle_core::{Device as CandleDevice, Tensor as CandleTensor};
-use half::f16;
+use half::{bf16, f16};
 use luminal::prelude::*;
 use proptest::prelude::*;
+use rand::{SeedableRng, rngs::StdRng};
+use safetensors::{Dtype, tensor::TensorView};
+use std::{
+    collections::HashMap,
+    path::PathBuf,
+    sync::atomic::{AtomicUsize, Ordering},
+};
+
+static SAFETENSORS_TEST_FILE_ID: AtomicUsize = AtomicUsize::new(0);
 
 fn assert_close(actual: &[f32], expected: &[f32], tolerance: f32) {
     assert_eq!(
@@ -26,6 +35,56 @@ fn assert_close(actual: &[f32], expected: &[f32], tolerance: f32) {
     }
 }
 
+fn bytes_of<T: bytemuck::NoUninit>(values: &[T]) -> Vec<u8> {
+    bytemuck::cast_slice(values).to_vec()
+}
+
+fn search_candidates(cx: &mut Graph, rt: MetalRuntime, limit: usize) -> MetalRuntime {
+    let mut rng = StdRng::seed_from_u64(0);
+    cx.search_options(rt, SearchOptions::new(limit), &mut rng)
+}
+
+fn egraph_has_op(cx: &Graph, op_name: &str) -> bool {
+    cx.egraph()
+        .expect("search space should be built")
+        .enodes
+        .values()
+        .any(|(label, _)| label == op_name)
+}
+
+fn assert_matmul_options(cx: &Graph, mps_op_name: &str) {
+    assert!(
+        egraph_has_op(cx, mps_op_name),
+        "expected {mps_op_name} rewrite option in e-graph"
+    );
+    assert!(
+        egraph_has_op(cx, "GenericMatmul"),
+        "expected GenericMatmul rewrite option in e-graph"
+    );
+}
+
+fn write_test_safetensors(tensors: &[(&str, Dtype, Vec<usize>, Vec<u8>)]) -> PathBuf {
+    let tensor_views: HashMap<String, TensorView<'_>> = tensors
+        .iter()
+        .map(|(name, dtype, shape, data)| {
+            (
+                (*name).to_string(),
+                TensorView::new(*dtype, shape.clone(), data).unwrap(),
+            )
+        })
+        .collect();
+    let serialized = safetensors::serialize(&tensor_views, None).unwrap();
+    let id = SAFETENSORS_TEST_FILE_ID.fetch_add(1, Ordering::Relaxed);
+    let mut path = std::env::temp_dir();
+    path.push(format!(
+        "luminal_metal_runtime_{}_{}.safetensors",
+        std::process::id(),
+        id
+    ));
+    std::fs::write(&path, serialized).unwrap();
+    path
+}
+
 const TRANSFORMER_SEQ: usize = 4;
 const TRANSFORMER_HIDDEN: usize = 16;
 const TRANSFORMER_INTERMEDIATE: usize = 32;
@@ -250,6 +309,36 @@ fn dynamic_dim_sum_reduce_runs() {
     assert_close(&out, &[9.0, 12.0], 0.001);
 }
 
+#[test]
+fn metal_bucketed_dynamic_dim_dispatches_correct_graph() {
+    let mut cx = Graph::default();
+    let input = cx.tensor(('s', 4));
+    let output = (input + input).output();
+
+    cx.set_dim_buckets('s', &[DimBucket::new(1, 1), DimBucket::new(2, 4)]);
+    cx.set_dim('s', 1);
+    cx.build_search_space::<MetalRuntime>();
+
+    let mut rt = MetalRuntime::initialize(());
+    rt.set_data(input, vec![1.0f32; 4]);
+    rt = cx.search(rt, 5);
+
+    cx.set_dim('s', 1);
+    let s1_input = vec![1.0, 2.0, 3.0, 4.0];
+    rt.set_data(input, s1_input.clone());
+    rt.execute(&cx.dyn_map);
+    let s1_out = rt.get_f32(output);
+    assert_close(&s1_out[..4], &[2.0, 4.0, 6.0, 8.0], 0.001);
+
+    cx.set_dim('s', 3);
+    let s3_input: Vec<f32> = (0..12).map(|i| i as f32).collect();
+    let s3_expected: Vec<f32> = s3_input.iter().map(|v| v * 2.0).collect();
+    rt.set_data(input, s3_input);
+    rt.execute(&cx.dyn_map);
+    let s3_out = rt.get_f32(output);
+    assert_close(&s3_out[..12], &s3_expected, 0.001);
+}
+
 #[test]
 fn metal_int_arithmetic_preserves_large_values() {
     let mut cx = Graph::default();
@@ -337,6 +426,18 @@ proptest! {
     }
 }
 
+#[test]
+fn metal_build_search_space_accepts_memory_budget() {
+    let mut cx = Graph::default();
+    let a = cx.tensor(4);
+    let b = cx.tensor(4);
+    (a * b).output();
+
+    cx.build_search_space_with_options::<MetalRuntime>(
+        BuildSearchSpaceOptions::new().max_memory_mib(1),
+    );
+}
+
 /// Simple deterministic test for add
 #[test]
 fn metal_simple_add() {
@@ -601,7 +702,7 @@ fn metal_specialized_matmul() {
 
     rt.set_data(a, &a_data);
     rt.set_data(b, &b_data);
-    rt = cx.search(rt, 1);
+    rt = search_candidates(&mut cx, rt, 32);
     assert!(
         rt.contains_matmul(),
         "expected Metal runtime to fuse matmul, kernels: {:?}",
@@ -634,6 +735,7 @@ fn metal_regular_tiled_matmul_path() {
     let output = a.matmul(b).output();
 
     cx.build_search_space::<MetalRuntime>();
+    assert_matmul_options(&cx, "MPSMatmul");
     let mut rt = MetalRuntime::initialize(());
 
     let a_data = seeded_data(m * k, 0.4, -0.2);
@@ -641,14 +743,7 @@ fn metal_regular_tiled_matmul_path() {
 
     rt.set_data(a, &a_data);
     rt.set_data(b, &b_data);
-    rt = cx.search(rt, 1);
-
-    let kernels = rt.debug_kernel_ops();
-    assert!(
-        kernels.iter().any(|k| k.contains("family: RegularTiled")),
-        "expected regular tiled matmul path, kernels: {:?}",
-        kernels
-    );
+    rt = search_candidates(&mut cx, rt, 32);
 
     rt.allocate_intermediate_buffers(&cx.dyn_map);
     rt.execute(&cx.dyn_map);
@@ -664,6 +759,259 @@ fn metal_regular_tiled_matmul_path() {
     assert_close(&result, &expected, 2e-3);
 }
 
+#[test]
+fn metal_mps_matmul_transposed_rhs_weight_layout() {
+    let mut cx = Graph::default();
+    let m = 7;
+    let k = 11;
+    let n = 13;
+    let a = cx.tensor((m, k));
+    let weight = cx.tensor((n, k));
+    let output = a.matmul(weight.t()).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    assert_matmul_options(&cx, "MPSMatmul");
+    let mut rt = MetalRuntime::initialize(());
+
+    let a_data = seeded_data(m * k, 0.35, -0.17);
+    let weight_data = seeded_data(n * k, 0.21, -0.09);
+
+    rt.set_data(a, &a_data);
+    rt.set_data(weight, &weight_data);
+    rt = search_candidates(&mut cx, rt, 32);
+
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+
+    let result = rt.get_f32(output);
+
+    let device = CandleDevice::Cpu;
+    let ref_a = CandleTensor::from_vec(a_data, (m, k), &device).unwrap();
+    let ref_weight = CandleTensor::from_vec(weight_data, (n, k), &device).unwrap();
+    let expected = ref_a.matmul(&ref_weight.t().unwrap()).unwrap();
+    let expected: Vec<f32> = expected.flatten_all().unwrap().to_vec1().unwrap();
+
+    assert_close(&result, &expected, 1e-3);
+}
+
+#[test]
+fn metal_mps_matmul_transposed_lhs_layout() {
+    let mut cx = Graph::default();
+    let m = 5;
+    let k = 9;
+    let n = 6;
+    let lhs_storage = cx.tensor((k, m));
+    let rhs = cx.tensor((k, n));
+    let output = lhs_storage.t().matmul(rhs).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    assert_matmul_options(&cx, "MPSMatmul");
+    let mut rt = MetalRuntime::initialize(());
+
+    let lhs_data = seeded_data(k * m, 0.31, -0.12);
+    let rhs_data = seeded_data(k * n, 0.27, -0.08);
+
+    rt.set_data(lhs_storage, &lhs_data);
+    rt.set_data(rhs, &rhs_data);
+    rt = search_candidates(&mut cx, rt, 32);
+
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+
+    let result = rt.get_f32(output);
+
+    let device = CandleDevice::Cpu;
+    let ref_lhs = CandleTensor::from_vec(lhs_data, (k, m), &device)
+        .unwrap()
+        .t()
+        .unwrap();
+    let ref_rhs = CandleTensor::from_vec(rhs_data, (k, n), &device).unwrap();
+    let expected = ref_lhs.matmul(&ref_rhs).unwrap();
+    let expected: Vec<f32> = expected.flatten_all().unwrap().to_vec1().unwrap();
+
+    assert_close(&result, &expected, 1e-3);
+}
+
+#[test]
+fn metal_mps_batched_matmul_row_row_layout() {
+    let mut cx = Graph::default();
+    let batch = 3;
+    let m = 4;
+    let k = 5;
+    let n = 6;
+    let a = cx.tensor((batch, m, k));
+    let b = cx.tensor((batch, k, n));
+    let output = a.matmul(b).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    assert_matmul_options(&cx, "MPSBatchedMatmul");
+    let mut rt = MetalRuntime::initialize(());
+
+    let a_data = seeded_data(batch * m * k, 0.17, -0.08);
+    let b_data = seeded_data(batch * k * n, 0.11, -0.05);
+    rt.set_data(a, &a_data);
+    rt.set_data(b, &b_data);
+    rt = search_candidates(&mut cx, rt, 32);
+
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+    let result = rt.get_f32(output);
+
+    let mut expected = vec![0.0; batch * m * n];
+    for batch_idx in 0..batch {
+        for row in 0..m {
+            for col in 0..n {
+                let mut sum = 0.0;
+                for inner in 0..k {
+                    sum += a_data[batch_idx * m * k + row * k + inner]
+                        * b_data[batch_idx * k * n + inner * n + col];
+                }
+                expected[batch_idx * m * n + row * n + col] = sum;
+            }
+        }
+    }
+
+    assert_close(&result, &expected, 1e-3);
+}
+
+#[test]
+fn metal_generic_matmul_covers_noncontiguous_merged_head_projection() {
+    let mut cx = Graph::default();
+    let heads = 3;
+    let seq = 4;
+    let head_dim = 5;
+    let hidden = heads * head_dim;
+    let out_dim = 7;
+    let attn = cx.tensor((heads, seq, head_dim));
+    let weight = cx.tensor((out_dim, hidden));
+    let merged = attn.transpose(0, 1).merge_dims(1, 2);
+    let output = merged.matmul(weight.t()).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    assert!(
+        egraph_has_op(&cx, "GenericMatmul"),
+        "expected GenericMatmul rewrite option in e-graph"
+    );
+    let mut rt = MetalRuntime::initialize(());
+
+    let attn_data = seeded_data(heads * seq * head_dim, 0.19, -0.09);
+    let weight_data = seeded_data(out_dim * hidden, 0.14, -0.06);
+    rt.set_data(attn, &attn_data);
+    rt.set_data(weight, &weight_data);
+    rt = search_candidates(&mut cx, rt, 32);
+
+    let kernels = rt.debug_kernel_ops();
+    assert!(
+        kernels.iter().any(|k| k.contains("GenericMatmul")),
+        "expected generic matmul fallback for non-contiguous merged-head projection, kernels: {:?}",
+        kernels
+    );
+    assert!(
+        !kernels.iter().any(|k| {
+            k.contains("MetalMul") && k.contains(&format!("shape: [{seq}, {out_dim}, {hidden}]"))
+        }),
+        "generic fallback should remove the broadcast multiply intermediate, kernels: {:?}",
+        kernels
+    );
+
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+    let result = rt.get_f32(output);
+
+    let mut expected = vec![0.0; seq * out_dim];
+    for token in 0..seq {
+        for out_col in 0..out_dim {
+            let mut sum = 0.0;
+            for inner in 0..hidden {
+                let head = inner / head_dim;
+                let dim = inner % head_dim;
+                let attn_idx = head * seq * head_dim + token * head_dim + dim;
+                sum += attn_data[attn_idx] * weight_data[out_col * hidden + inner];
+            }
+            expected[token * out_dim + out_col] = sum;
+        }
+    }
+
+    assert_close(&result, &expected, 1e-3);
+}
+
+#[test]
+fn metal_mps_batched_matmul_transposed_rhs_layout() {
+    let mut cx = Graph::default();
+    let batch = 4;
+    let m = 3;
+    let k = 7;
+    let n = 5;
+    let a = cx.tensor((batch, m, k));
+    let weight = cx.tensor((batch, n, k));
+    let output = a.matmul(weight.permute((0, 2, 1))).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    assert_matmul_options(&cx, "MPSBatchedMatmul");
+    let mut rt = MetalRuntime::initialize(());
+
+    let a_data = seeded_data(batch * m * k, 0.13, -0.06);
+    let weight_data = seeded_data(batch * n * k, 0.09, -0.04);
+    rt.set_data(a, &a_data);
+    rt.set_data(weight, &weight_data);
+    rt = search_candidates(&mut cx, rt, 32);
+
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+    let result = rt.get_f32(output);
+
+    let mut expected = vec![0.0; batch * m * n];
+    for batch_idx in 0..batch {
+        for row in 0..m {
+            for col in 0..n {
+                let mut sum = 0.0;
+                for inner in 0..k {
+                    sum += a_data[batch_idx * m * k + row * k + inner]
+                        * weight_data[batch_idx * n * k + col * k + inner];
+                }
+                expected[batch_idx * m * n + row * n + col] = sum;
+            }
+        }
+    }
+
+    assert_close(&result, &expected, 1e-3);
+}
+
+#[test]
+fn metal_mps_matmul_f16_transposed_rhs_weight_layout() {
+    let mut cx = Graph::default();
+    let m = 6;
+    let k = 10;
+    let n = 7;
+    let a = cx.tensor((m, k)).as_dtype(DType::F16);
+    let weight = cx.tensor((n, k)).as_dtype(DType::F16);
+    let output = a.matmul(weight.t()).cast(DType::F32).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    assert_matmul_options(&cx, "MPSMatmul");
+    let mut rt = MetalRuntime::initialize(());
+
+    let a_data = seeded_data(m * k, 0.22, -0.07);
+    let weight_data = seeded_data(n * k, 0.18, -0.05);
+
+    rt.set_data(a, to_f16_vec(&a_data));
+    rt.set_data(weight, to_f16_vec(&weight_data));
+    rt = search_candidates(&mut cx, rt, 32);
+
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+
+    let result = rt.get_f32(output);
+
+    let device = CandleDevice::Cpu;
+    let ref_a = CandleTensor::from_vec(a_data, (m, k), &device).unwrap();
+    let ref_weight = CandleTensor::from_vec(weight_data, (n, k), &device).unwrap();
+    let expected = ref_a.matmul(&ref_weight.t().unwrap()).unwrap();
+    let expected: Vec<f32> = expected.flatten_all().unwrap().to_vec1().unwrap();
+
+    assert_close(&result, &expected, 5e-3);
+}
+
 #[test]
 fn metal_rms_norm() {
     let mut cx = Graph::default();
@@ -988,6 +1336,131 @@ fn test_scatter_basic() {
     assert_close(&out, &[0.0, 10.0, 0.0, 20.0, 30.0], 0.001);
 }
 
+#[test]
+fn test_scatter_buffer_roundtrip() {
+    let mut cx = Graph::default();
+    let src = cx.tensor(1);
+    let indexes = cx.tensor(1).as_dtype(DType::Int);
+    let cache = cx.tensor(4).persist();
+    let cache_out = src.scatter(indexes, cache);
+    let read = cache_out.output();
+
+    cx.build_search_space::<MetalRuntime>();
+    let mut rt = MetalRuntime::initialize(());
+    rt.set_data(src, &[0.0]);
+    rt.set_data(indexes, &[0.0]);
+    rt.set_zeros(cache, 4 * std::mem::size_of::<f32>());
+    rt = cx.search(rt, 1);
+
+    for (pos, value, expected) in [
+        (0, 10.0, [10.0, 0.0, 0.0, 0.0]),
+        (1, 20.0, [10.0, 20.0, 0.0, 0.0]),
+        (2, 30.0, [10.0, 20.0, 30.0, 0.0]),
+    ] {
+        rt.set_data(src, &[value]);
+        rt.set_data(indexes, &[pos as f32]);
+        rt.allocate_intermediate_buffers(&cx.dyn_map);
+        rt.execute(&cx.dyn_map);
+        assert_close(&rt.get_f32(read), &expected, 0.001);
+
+        let updated_cache = rt.remove_buffer(cache_out);
+        rt.set_buffer(cache, updated_cache);
+    }
+}
+
+#[test]
+fn test_load_safetensors_f32_survives_search_and_overrides_input_data() {
+    let mut cx = Graph::default();
+    let weights = cx.named_tensor("weights", 3);
+    let bias = cx.named_tensor("bias", 3);
+    let out = (weights + bias).output();
+
+    let weight_values = [1.25f32, -2.5, 4.0];
+    let tensors = [("weights", Dtype::F32, vec![3], bytes_of(&weight_values))];
+    let path = write_test_safetensors(&tensors);
+
+    cx.build_search_space::<MetalRuntime>();
+    let mut rt = MetalRuntime::initialize(());
+    rt.set_data(weights, &[99.0, 99.0, 99.0]);
+    rt.set_data(bias, &[0.5, 1.0, -1.5]);
+    rt.load_safetensors(&cx, path.to_str().unwrap());
+    rt = cx.search(rt, 1);
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(out), &[1.75, -1.5, 2.5], 0.001);
+    std::fs::remove_file(path).ok();
+}
+
+#[test]
+fn test_load_safetensors_converts_supported_float_dtypes() {
+    let mut cx = Graph::default();
+    let f16_to_f32 = cx.named_tensor("f16_to_f32", 2);
+    let bf16_to_f32 = cx.named_tensor("bf16_to_f32", 2);
+    let f16_to_f16 = cx.named_tensor("f16_to_f16", 2).as_dtype(DType::F16);
+    let f32_to_f16 = cx.named_tensor("f32_to_f16", 2).as_dtype(DType::F16);
+    let bf16_to_f16 = cx.named_tensor("bf16_to_f16", 2).as_dtype(DType::F16);
+
+    let f16_to_f32_out = (f16_to_f32 + 0.0).output();
+    let bf16_to_f32_out = (bf16_to_f32 + 0.0).output();
+    let f16_to_f16_out = f16_to_f16.cast(DType::F32).output();
+    let f32_to_f16_out = f32_to_f16.cast(DType::F32).output();
+    let bf16_to_f16_out = bf16_to_f16.cast(DType::F32).output();
+
+    let f16_to_f32_values = [f16::from_f32(1.5), f16::from_f32(-2.25)];
+    let bf16_to_f32_values = [bf16::from_f32(3.5), bf16::from_f32(-4.25)];
+    let f16_to_f16_values = [f16::from_f32(5.5), f16::from_f32(-6.25)];
+    let f32_to_f16_values = [7.5f32, -8.25];
+    let bf16_to_f16_values = [bf16::from_f32(9.5), bf16::from_f32(-10.25)];
+    let tensors = [
+        (
+            "f16_to_f32",
+            Dtype::F16,
+            vec![2],
+            bytes_of(&f16_to_f32_values),
+        ),
+        (
+            "bf16_to_f32",
+            Dtype::BF16,
+            vec![2],
+            bytes_of(&bf16_to_f32_values),
+        ),
+        (
+            "f16_to_f16",
+            Dtype::F16,
+            vec![2],
+            bytes_of(&f16_to_f16_values),
+        ),
+        (
+            "f32_to_f16",
+            Dtype::F32,
+            vec![2],
+            bytes_of(&f32_to_f16_values),
+        ),
+        (
+            "bf16_to_f16",
+            Dtype::BF16,
+            vec![2],
+            bytes_of(&bf16_to_f16_values),
+        ),
+    ];
+    let path = write_test_safetensors(&tensors);
+
+    cx.build_search_space::<MetalRuntime>();
+    let mut rt = MetalRuntime::initialize(());
+    rt.load_safetensors(&cx, path.to_str().unwrap());
+    rt = cx.search(rt, 1);
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(f16_to_f32_out), &[1.5, -2.25], 0.001);
+    assert_close(&rt.get_f32(bf16_to_f32_out), &[3.5, -4.25], 0.001);
+    assert_close(&rt.get_f32(f16_to_f16_out), &[5.5, -6.25], 0.001);
+    assert_close(&rt.get_f32(f32_to_f16_out), &[7.5, -8.25], 0.001);
+    assert_close(&rt.get_f32(bf16_to_f16_out), &[9.5, -10.25], 0.001);
+    std::fs::remove_file(path).ok();
+}
+
 #[test]
 fn test_gather_noncontiguous_data_uses_data_shape() {
     let mut cx = Graph::default();
@@ -1024,6 +1497,12 @@ fn test_scatter_into_nonzero_dest() {
     rt.set_data(indexes, &[2f32]);
     rt.set_data(dest, &[1.0, 2.0, 3.0, 4.0, 5.0]);
     rt = cx.search(rt, 1);
+    let kernels = rt.debug_kernel_ops();
+    assert!(
+        kernels.iter().any(|k| k.contains("MetalScatterNoCopy")),
+        "expected no-copy scatter for consumed destination, kernels: {:?}",
+        kernels
+    );
     rt.allocate_intermediate_buffers(&cx.dyn_map);
     rt.execute(&cx.dyn_map);
 
@@ -1031,6 +1510,89 @@ fn test_scatter_into_nonzero_dest() {
     assert_close(&out, &[1.0, 2.0, 99.0, 4.0, 5.0], 0.001);
 }
 
+#[test]
+fn test_scatter_no_copy_remove_buffer_aliases_dest() {
+    let mut cx = Graph::default();
+    let src = cx.tensor(2);
+    let indexes = cx.tensor(2).as_dtype(DType::Int);
+    let dest = cx.tensor(5);
+    let result = src.scatter(indexes, dest).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    let mut rt = MetalRuntime::initialize(());
+    rt.set_data(src, &[7.0, 8.0]);
+    rt.set_data(indexes, &[1.0, 3.0]);
+    rt.set_data(dest, &[10.0, 20.0, 30.0, 40.0, 50.0]);
+    rt = cx.search(rt, 1);
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+
+    let moved = rt.remove_buffer(result);
+    let moved_values = unsafe {
+        std::slice::from_raw_parts(
+            moved.contents() as *const f32,
+            moved.length() as usize / std::mem::size_of::<f32>(),
+        )
+        .to_vec()
+    };
+    assert_close(&moved_values, &[10.0, 7.0, 30.0, 8.0, 50.0], 0.001);
+    rt.set_buffer(dest.id, moved);
+}
+
+#[test]
+fn test_scatter_no_copy_handles_2d_destination() {
+    let mut cx = Graph::default();
+    let src = cx.tensor(2);
+    let indexes = cx.tensor(2).as_dtype(DType::Int);
+    let dest = cx.tensor((2, 3));
+    let result = src.scatter(indexes, dest).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    let mut rt = MetalRuntime::initialize(());
+    rt.set_data(src, &[9.0, 8.0]);
+    rt.set_data(indexes, &[2.0, 4.0]);
+    rt.set_data(dest, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
+    rt = cx.search(rt, 1);
+    let kernels = rt.debug_kernel_ops();
+    assert!(
+        kernels.iter().any(|k| k.contains("MetalScatterNoCopy")),
+        "expected no-copy scatter for 2D destination, kernels: {:?}",
+        kernels
+    );
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(result), &[1.0, 2.0, 9.0, 4.0, 8.0, 6.0], 0.001);
+}
+
+#[test]
+fn test_scatter_no_copy_not_selected_when_dest_has_another_consumer() {
+    let mut cx = Graph::default();
+    let src = cx.tensor(1);
+    let indexes = cx.tensor(1).as_dtype(DType::Int);
+    let dest = cx.tensor(4);
+    let scatter = src.scatter(indexes, dest).output();
+    let dest_plus_one = (dest + 1.0).output();
+
+    cx.build_search_space::<MetalRuntime>();
+    let mut rt = MetalRuntime::initialize(());
+    rt.set_data(src, &[99.0]);
+    rt.set_data(indexes, &[1.0]);
+    rt.set_data(dest, &[10.0, 20.0, 30.0, 40.0]);
+    rt = cx.search(rt, 1);
+    let kernels = rt.debug_kernel_ops();
+    assert!(
+        !kernels.iter().any(|k| k.contains("MetalScatterNoCopy")),
+        "no-copy scatter should not be selected when dest is also consumed, kernels: {:?}",
+        kernels
+    );
+    rt.allocate_intermediate_buffers(&cx.dyn_map);
+    rt.execute(&cx.dyn_map);
+
+    assert_close(&rt.get_f32(scatter), &[10.0, 99.0, 30.0, 40.0], 0.001);
+    assert_close(&rt.get_f32(dest_plus_one), &[11.0, 21.0, 31.0, 41.0], 0.001);
+}
+
 #[test]
 fn test_scatter_all_positions() {
     let mut cx = Graph::default();

crates/luminal_nn/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "luminal_nn"
 version = "0.1.0"
-edition = "2021"
+edition = "2024"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

crates/luminal_nn/src/moe.rs

@@ -71,7 +71,7 @@ impl MoE {
 mod tests {
     use super::MoE;
     use luminal::prelude::*;
-    use rand::{rng, Rng};
+    use rand::{Rng, rng};
 
     fn random_vec(n: usize) -> Vec<f32> {
         let mut r = rng();

crates/luminal_python/examples/whisper.py

@@ -431,7 +431,7 @@ def main() -> None:
     tokenizer = WhisperTokenizer.from_pretrained(REPO_ID)
 
     use_compiled = os.environ.get("LUMINAL_DISABLE", "0") != "1"
-    max_new_tokens = int(os.environ.get("GEN_TOKENS", "100"))
+    max_new_tokens = 100
     search_iters = int(os.environ.get("SEARCH_ITERATIONS", "10"))
 
     if use_compiled:

crates/luminal_python/rust/src/compiled_graph.rs

@@ -158,17 +158,21 @@ impl CompiledGraph {
             input_shape_exprs,
             dim_param_map,
         } = translation;
+        let WeightData {
+            weights,
+            tensor_sizes,
+            device_ptrs,
+        } = weight_data;
 
-        // Build compile args from WeightData (convert TypedData -> raw bytes + dtype)
+        // Build compile args from WeightData.
         let compile_args = BackendCompileArgs {
             search_iters,
-            weights: weight_data
-                .weights
+            weights: weights
                 .iter()
                 .map(|(label, td)| (label.clone(), td.bytes.clone(), td.dtype))
                 .collect(),
-            tensor_sizes: weight_data.tensor_sizes,
-            device_ptrs: weight_data.device_ptrs,
+            tensor_sizes,
+            device_ptrs,
         };
 
         // Create backend via the factory directly
@@ -387,7 +391,74 @@ impl CompiledGraph {
         Ok(())
     }
 
-    /// Set a weight from a device pointer (e.g. "fc1.weight"). Zero-copy on device.
+    /// Resolve an input or output tensor name to an opaque integer ID.
+    /// One FFI hop at compile time, then per-iter `run_with_ptrs` and
+    /// the `*_by_id` setters can skip the string-keyed HashMap lookup.
+    /// Returns the underlying `NodeIndex.index() as u32` — caller should
+    /// treat it as opaque.
+    fn tensor_id(&self, name: &str) -> PyResult<u32> {
+        self.tensor_ids
+            .get(name)
+            .map(|n| n.index() as u32)
+            .ok_or_else(|| {
+                PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                    "Unknown tensor: {}",
+                    name
+                ))
+            })
+    }
+
+    /// One-shot batched: register all input + output device pointers and
+    /// execute. Collapses the per-iter `set_input_device_ptr` ×
+    /// (n_inputs) + `set_output_device_ptr` × (n_outputs) + `run` chain
+    /// into a single Python→Rust FFI crossing.
+    ///
+    /// Inputs and outputs use the opaque IDs returned by `tensor_id(name)`
+    /// to skip the per-call string lookup. Returns a parallel vector of
+    /// "is zero-copy" booleans for each output (same as
+    /// `output_is_zero_copy(name)` on the unbatched path), so the Python
+    /// caller can fall back to a DtoD copy when a kernel aliased an
+    /// output instead of writing into the registered buffer.
+    ///
+    /// Safety contract:
+    ///   * Every `device_ptr` must point to a valid CUDA allocation with
+    ///     at least `n_bytes` bytes.
+    ///   * Pointers must remain valid through the duration of `run()`.
+    ///   * `tensor_id`s must come from `self.tensor_id(name)` on this
+    ///     same graph — using stale IDs from a different graph is UB
+    ///     (will likely panic in `NodeIndex` lookups, but not guaranteed).
+    fn run_with_ptrs(
+        &mut self,
+        inputs: Vec<(u32, u64, usize)>,
+        outputs: Vec<(u32, u64, usize)>,
+    ) -> PyResult<Vec<bool>> {
+        if !self.runtime.supports_device_ptrs() {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "run_with_ptrs requires a GPU backend",
+            ));
+        }
+        // Register inputs.
+        for (id, ptr, n) in &inputs {
+            let node_id = NodeIndex::new(*id as usize);
+            unsafe { self.runtime.set_device_ptr(node_id, *ptr, *n) };
+        }
+        // Register outputs.
+        for (id, ptr, n) in &outputs {
+            let node_id = NodeIndex::new(*id as usize);
+            unsafe { self.runtime.set_output_device_ptr(node_id, *ptr, *n) };
+        }
+        // Execute.
+        self.runtime.execute(&self.graph.dyn_map);
+        // Report zero-copy status for each output (parallel to `outputs`
+        // input order). Aliased outputs need a DtoD copy on the Python
+        // side, same as the unbatched path.
+        Ok(outputs
+            .iter()
+            .map(|(id, _, _)| self.runtime.output_is_zero_copy(NodeIndex::new(*id as usize)))
+            .collect())
+    }
+
+    /// Register 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,
@@ -448,7 +519,7 @@ impl CompiledGraph {
         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).
+    /// Register 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,

crates/luminal_python/rust/src/lib.rs

@@ -3,6 +3,7 @@ pub mod typed_data;
 
 // PT2 modules
 mod pt2_compiled_model;
+mod pt2_expr;
 mod pt2_parser;
 mod pt2_schema;
 mod pt2_util;

crates/luminal_python/rust/src/pt2_compiled_model.rs

@@ -6,6 +6,7 @@ use pyo3::types::{PyCapsule, PyCapsuleMethods};
 use std::collections::HashMap;
 
 use crate::compiled_graph::{CompiledGraph, DimParamMap, GraphTranslation, WeightData};
+use crate::pt2_expr::parse_sympy_expr;
 use crate::pt2_schema;
 use crate::translator;
 use crate::typed_data::TypedData;
@@ -21,7 +22,7 @@ fn resolve_dim_sizes(
     sizes
         .iter()
         .map(|s| match s {
-            pt2_schema::DimSize::Int(i) => Expression::from(i.as_int as usize),
+            pt2_schema::DimSize::Int(i) => Expression::from(i.as_int),
             pt2_schema::DimSize::Expr(e) => {
                 let s = e.as_expr.expr_str.trim();
                 // Try the full sympy-style parse first so compound forms like
@@ -45,7 +46,7 @@ fn resolve_dim_sizes(
                             .hint
                             .as_ref()
                             .and_then(|h| h.as_int())
-                            .map(|h| Expression::from(h as usize))
+                            .map(Expression::from)
                     })
                     .unwrap_or_else(|| Expression::from(1usize))
             }
@@ -53,139 +54,6 @@ fn resolve_dim_sizes(
         .collect()
 }
 
-/// Parse a sympy `srepr`-style expression string into a luminal Expression.
-///
-/// Handles the subset of sympy heads PT2 actually emits for shape metadata:
-///
-///   * `Symbol('name', ...)` — bound to the corresponding luminal char if
-///     present in `sym_to_char`, or treated as a fresh constant 1 otherwise.
-///   * `Integer(N)` / `Number(N)` — concrete int.
-///   * `Mul(a, b, ...)` / `Add(a, b, ...)` — n-ary, folded into pairwise ops.
-///
-/// Returns `None` for anything else so the caller can fall back to a less
-/// precise representation rather than committing a wrong expression.
-fn parse_sympy_expr(s: &str, sym_to_char: &HashMap<String, char>) -> Option<Expression> {
-    let s = s.trim();
-    if s.is_empty() {
-        return None;
-    }
-
-    // Bare integer literal — `srepr` doesn't usually emit this at the top
-    // level (it wraps in `Integer(...)`), but accept it for robustness.
-    if let Ok(n) = s.parse::<i64>() {
-        return Some(Expression::from(n as usize));
-    }
-
-    let (head, body) = split_head(s)?;
-    match head {
-        "Symbol" => {
-            // Body is `'name', positive=True, integer=True` etc. Pull the
-            // first quoted token as the name.
-            let name = extract_first_quoted(body)?;
-            sym_to_char.get(&name).map(|c| Expression::from(*c))
-        }
-        "Integer" | "Number" => {
-            let n: i64 = body.trim().parse().ok()?;
-            Some(Expression::from(n as usize))
-        }
-        "Mul" | "Add" => {
-            let parts = split_top_level_args(body);
-            if parts.is_empty() {
-                return None;
-            }
-            let mut iter = parts.into_iter();
-            let mut acc = parse_sympy_expr(iter.next()?, sym_to_char)?;
-            for p in iter {
-                let rhs = parse_sympy_expr(p, sym_to_char)?;
-                acc = if head == "Mul" { acc * rhs } else { acc + rhs };
-            }
-            Some(acc)
-        }
-        _ => None,
-    }
-}
-
-/// Split `Head(body)` into (head, body); returns None if not in that form.
-fn split_head(s: &str) -> Option<(&str, &str)> {
-    let open = s.find('(')?;
-    if !s.ends_with(')') {
-        return None;
-    }
-    Some((&s[..open], &s[open + 1..s.len() - 1]))
-}
-
-/// Pull out the first single- or double-quoted token from a sympy arg list,
-/// e.g. `'s77', positive=True` → `s77`.
-fn extract_first_quoted(s: &str) -> Option<String> {
-    let bytes = s.as_bytes();
-    let mut i = 0;
-    while i < bytes.len() {
-        let c = bytes[i] as char;
-        if c == '\'' || c == '"' {
-            let quote = c;
-            let start = i + 1;
-            i += 1;
-            while i < bytes.len() && bytes[i] as char != quote {
-                i += 1;
-            }
-            return Some(s[start..i].to_string());
-        }
-        i += 1;
-    }
-    None
-}
-
-/// Split sympy-style argument list at top-level commas, respecting nested
-/// parens and quoted strings. Discards `key=value` kwargs (they don't carry
-/// dimensional information).
-fn split_top_level_args(s: &str) -> Vec<&str> {
-    let mut out = Vec::new();
-    let bytes = s.as_bytes();
-    let mut depth = 0;
-    let mut in_quote: Option<char> = None;
-    let mut start = 0;
-    for (i, &b) in bytes.iter().enumerate() {
-        let c = b as char;
-        match in_quote {
-            Some(q) => {
-                if c == q {
-                    in_quote = None;
-                }
-            }
-            None => match c {
-                '\'' | '"' => in_quote = Some(c),
-                '(' | '[' => depth += 1,
-                ')' | ']' => depth -= 1,
-                ',' if depth == 0 => {
-                    let part = s[start..i].trim();
-                    // Drop `key=value` kwargs — they're metadata sympy uses
-                    // for pretty-printing, not arguments to the operator.
-                    if !part.is_empty() && !looks_like_kwarg(part) {
-                        out.push(part);
-                    }
-                    start = i + 1;
-                }
-                _ => {}
-            },
-        }
-    }
-    let part = s[start..].trim();
-    if !part.is_empty() && !looks_like_kwarg(part) {
-        out.push(part);
-    }
-    out
-}
-
-fn looks_like_kwarg(part: &str) -> bool {
-    if let Some(eq) = part.find('=') {
-        let key = part[..eq].trim();
-        // sympy kwargs are bare identifiers like `positive`, `integer`.
-        !key.is_empty() && key.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
-    } else {
-        false
-    }
-}
-
 #[pyfunction]
 #[pyo3(signature = (pt2_path, weights_path, search_iters, factory_capsule, weight_device_ptrs=None))]
 pub fn process_pt2(

crates/luminal_python/rust/src/pt2_expr.rs

@@ -0,0 +1,713 @@
+use std::collections::HashMap;
+
+use luminal::prelude::*;
+use rustc_hash::FxHashMap;
+
+use crate::pt2_parser::SymDimMap;
+use crate::pt2_schema::RangeConstraint;
+
+#[derive(Clone, Copy, Debug, Default)]
+pub(crate) struct ExprBounds {
+    pub(crate) min: Option<i64>,
+    pub(crate) max: Option<i64>,
+}
+
+#[derive(Clone, Copy, Debug)]
+struct ParsedExpr {
+    expr: Expression,
+    bounds: ExprBounds,
+}
+
+impl ParsedExpr {
+    fn exact(expr: Expression, value: i64) -> Self {
+        Self {
+            expr,
+            bounds: ExprBounds {
+                min: Some(value),
+                max: Some(value),
+            },
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+struct BoundedExpr {
+    expr: Expression,
+    bounds: ExprBounds,
+}
+
+/// Parse a sympy `srepr`-style expression string into a luminal `Expression`.
+///
+/// Supports the subset of sympy heads PT2 emits for symbolic shape metadata.
+pub(crate) fn parse_sympy_expr(
+    expr: &str,
+    sym_to_char: &HashMap<String, char>,
+) -> Option<Expression> {
+    parse_sympy_expr_with_ranges(expr, sym_to_char, &HashMap::new())
+}
+
+pub(crate) fn parse_sympy_expr_with_ranges(
+    expr: &str,
+    sym_to_char: &HashMap<String, char>,
+    ranges: &HashMap<String, RangeConstraint>,
+) -> Option<Expression> {
+    parse_sympy_expr_inner(expr, sym_to_char, ranges).map(|parsed| parsed.expr)
+}
+
+pub(crate) fn sym_char_ranges(sym_map: &SymDimMap) -> FxHashMap<char, ExprBounds> {
+    sym_map
+        .sym_to_char
+        .iter()
+        .map(|(sym_name, sym_char)| {
+            let range = sym_map.ranges.get(sym_name);
+            let min = range
+                .and_then(|range| range.min_val)
+                .map(|min| min.max(0))
+                .or(Some(0));
+            let max = range
+                .and_then(|range| range.max_val)
+                .filter(|max| *max >= 0);
+            (*sym_char, ExprBounds { min, max })
+        })
+        .collect()
+}
+
+pub(crate) fn simplify_expr_with_ranges(
+    expr: Expression,
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) -> Expression {
+    simplify_bound_expr(expr, sym_ranges).expr
+}
+
+pub(crate) fn same_expr_with_ranges(
+    lhs: Expression,
+    rhs: Expression,
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) -> bool {
+    let lhs = simplify_bound_expr(lhs, sym_ranges);
+    let rhs = simplify_bound_expr(rhs, sym_ranges);
+    lhs.expr == rhs.expr
+        || lhs.expr.egglog_equal(rhs.expr)
+        || (exact_value(lhs) == exact_value(rhs) && exact_value(lhs).is_some())
+}
+
+pub(crate) fn canonical_equal_expr(
+    lhs: Expression,
+    rhs: Expression,
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) -> Option<Expression> {
+    if !same_expr_with_ranges(lhs, rhs, sym_ranges) {
+        return None;
+    }
+    let lhs_simplified = simplify_expr_with_ranges(lhs, sym_ranges);
+    let rhs_simplified = simplify_expr_with_ranges(rhs, sym_ranges);
+    Some(if lhs_simplified.len() <= rhs_simplified.len() {
+        lhs_simplified
+    } else {
+        rhs_simplified
+    })
+}
+
+fn parse_sympy_expr_inner(
+    expr: &str,
+    sym_to_char: &HashMap<String, char>,
+    ranges: &HashMap<String, RangeConstraint>,
+) -> Option<ParsedExpr> {
+    let expr = expr.trim();
+    if expr.is_empty() {
+        return None;
+    }
+
+    if let Ok(value) = expr.parse::<i64>() {
+        return Some(ParsedExpr::exact(Expression::from(value), value));
+    }
+
+    let (head, body) = split_head(expr)?;
+    match head {
+        "Symbol" => {
+            let name = extract_first_quoted(body)?;
+            let bounds = infer_symbol_bounds(body, ranges.get(&name));
+            sym_to_char.get(&name).map(|c| ParsedExpr {
+                expr: Expression::from(*c),
+                bounds,
+            })
+        }
+        "Integer" | "Number" => {
+            let value = body.trim().parse::<i64>().ok()?;
+            Some(ParsedExpr::exact(Expression::from(value), value))
+        }
+        "NegativeOne" => Some(ParsedExpr::exact(Expression::from(-1i64), -1)),
+        "Zero" => Some(ParsedExpr::exact(Expression::from(0i64), 0)),
+        "One" => Some(ParsedExpr::exact(Expression::from(1i64), 1)),
+        "Mul" | "Add" | "Min" | "Max" => {
+            let parts = split_top_level_args(body);
+            if parts.is_empty() {
+                return None;
+            }
+            let mut iter = parts.into_iter();
+            let mut acc = parse_sympy_expr_inner(iter.next()?, sym_to_char, ranges)?;
+            for part in iter {
+                let rhs = parse_sympy_expr_inner(part, sym_to_char, ranges)?;
+                acc = match head {
+                    "Mul" => ParsedExpr {
+                        expr: normalize_mul_expr(acc.expr, rhs.expr),
+                        bounds: mul_bounds(acc.bounds, rhs.bounds),
+                    },
+                    "Add" => ParsedExpr {
+                        expr: normalize_add_expr(acc.expr, rhs.expr),
+                        bounds: add_bounds(acc.bounds, rhs.bounds),
+                    },
+                    "Min" => reduce_min(acc, rhs),
+                    "Max" => reduce_max(acc, rhs),
+                    _ => unreachable!(),
+                };
+            }
+            Some(acc)
+        }
+        "FloorDiv" => {
+            let mut parts = split_top_level_args(body).into_iter();
+            let lhs = parse_sympy_expr_inner(parts.next()?, sym_to_char, ranges)?;
+            let rhs = parse_sympy_expr_inner(parts.next()?, sym_to_char, ranges)?;
+            if parts.next().is_some() {
+                return None;
+            }
+            Some(ParsedExpr {
+                expr: lhs.expr / rhs.expr,
+                bounds: div_bounds(lhs.bounds, rhs.bounds),
+            })
+        }
+        "Mod" => {
+            let mut parts = split_top_level_args(body).into_iter();
+            let lhs = parse_sympy_expr_inner(parts.next()?, sym_to_char, ranges)?;
+            let rhs = parse_sympy_expr_inner(parts.next()?, sym_to_char, ranges)?;
+            if parts.next().is_some() {
+                return None;
+            }
+            Some(ParsedExpr {
+                expr: lhs.expr % rhs.expr,
+                bounds: mod_bounds(lhs.bounds, rhs.bounds),
+            })
+        }
+        _ => None,
+    }
+}
+
+fn infer_symbol_bounds(body: &str, range: Option<&RangeConstraint>) -> ExprBounds {
+    let mut bounds = ExprBounds::default();
+    if body.contains("positive=True") {
+        bounds.min = Some(1);
+    } else if body.contains("nonnegative=True") {
+        bounds.min = Some(0);
+    }
+    if let Some(range) = range {
+        bounds.min = match (bounds.min, range.min_val) {
+            (Some(lhs), Some(rhs)) => Some(lhs.max(rhs)),
+            (None, Some(rhs)) => Some(rhs),
+            (lhs, None) => lhs,
+        };
+        bounds.max = range.max_val;
+    }
+    bounds
+}
+
+fn exact_expr(value: i64) -> BoundedExpr {
+    BoundedExpr {
+        expr: Expression::from(value),
+        bounds: ExprBounds {
+            min: Some(value),
+            max: Some(value),
+        },
+    }
+}
+
+fn exact_value(expr: BoundedExpr) -> Option<i64> {
+    expr.expr.as_num().or({
+        (expr.bounds.min == expr.bounds.max)
+            .then_some(expr.bounds.min)
+            .flatten()
+    })
+}
+
+fn exact_bound_value(bounds: ExprBounds) -> Option<i64> {
+    (bounds.min == bounds.max).then_some(bounds.min).flatten()
+}
+
+fn with_bounds(expr: Expression, bounds: ExprBounds) -> BoundedExpr {
+    BoundedExpr { expr, bounds }
+}
+
+fn bool_bounds() -> ExprBounds {
+    ExprBounds {
+        min: Some(0),
+        max: Some(1),
+    }
+}
+
+fn normalize_expr(expr: Expression) -> Expression {
+    if expr.len() <= 16 {
+        expr.simplify()
+    } else {
+        expr
+    }
+}
+
+fn commutative_key(expr: Expression) -> (usize, String) {
+    (expr.len(), format!("{expr:?}"))
+}
+
+fn sort_commutative(lhs: Expression, rhs: Expression) -> (Expression, Expression) {
+    if commutative_key(lhs) <= commutative_key(rhs) {
+        (lhs, rhs)
+    } else {
+        (rhs, lhs)
+    }
+}
+
+fn normalize_add_expr(lhs: Expression, rhs: Expression) -> Expression {
+    let (lhs, rhs) = sort_commutative(lhs, rhs);
+    normalize_expr(lhs + rhs)
+}
+
+fn normalize_mul_expr(lhs: Expression, rhs: Expression) -> Expression {
+    let (lhs, rhs) = sort_commutative(lhs, rhs);
+    normalize_expr(lhs * rhs)
+}
+
+fn checked_add_opt(lhs: Option<i64>, rhs: Option<i64>) -> Option<i64> {
+    lhs.zip(rhs).and_then(|(lhs, rhs)| lhs.checked_add(rhs))
+}
+
+fn checked_sub_opt(lhs: Option<i64>, rhs: Option<i64>) -> Option<i64> {
+    lhs.zip(rhs).and_then(|(lhs, rhs)| lhs.checked_sub(rhs))
+}
+
+fn checked_mul_opt(lhs: Option<i64>, rhs: Option<i64>) -> Option<i64> {
+    lhs.zip(rhs).and_then(|(lhs, rhs)| lhs.checked_mul(rhs))
+}
+
+fn add_bounds(lhs: ExprBounds, rhs: ExprBounds) -> ExprBounds {
+    ExprBounds {
+        min: checked_add_opt(lhs.min, rhs.min),
+        max: checked_add_opt(lhs.max, rhs.max),
+    }
+}
+
+fn mul_bounds(lhs: ExprBounds, rhs: ExprBounds) -> ExprBounds {
+    if lhs.min.unwrap_or(i64::MIN) >= 0 && rhs.min.unwrap_or(i64::MIN) >= 0 {
+        return ExprBounds {
+            min: checked_mul_opt(lhs.min, rhs.min),
+            max: checked_mul_opt(lhs.max, rhs.max),
+        };
+    }
+    ExprBounds::default()
+}
+
+fn sub_bounds(lhs: ExprBounds, rhs: ExprBounds) -> ExprBounds {
+    ExprBounds {
+        min: checked_sub_opt(lhs.min, rhs.max),
+        max: checked_sub_opt(lhs.max, rhs.min),
+    }
+}
+
+fn div_bounds(lhs: ExprBounds, rhs: ExprBounds) -> ExprBounds {
+    let (Some(rhs_min), Some(rhs_max)) = (rhs.min, rhs.max) else {
+        return ExprBounds::default();
+    };
+    if rhs_min <= 0 || rhs_max <= 0 {
+        return ExprBounds::default();
+    }
+    ExprBounds {
+        min: lhs.min.and_then(|lhs_min| lhs_min.checked_div(rhs_max)),
+        max: lhs.max.and_then(|lhs_max| lhs_max.checked_div(rhs_min)),
+    }
+}
+
+fn mod_bounds(lhs: ExprBounds, rhs: ExprBounds) -> ExprBounds {
+    if lhs.min.unwrap_or(i64::MIN) < 0 {
+        return ExprBounds::default();
+    }
+    match exact_bound_value(rhs) {
+        Some(rhs_exact) if rhs_exact > 0 => ExprBounds {
+            min: Some(0),
+            max: rhs_exact.checked_sub(1),
+        },
+        _ => ExprBounds::default(),
+    }
+}
+
+fn reduce_min(lhs: ParsedExpr, rhs: ParsedExpr) -> ParsedExpr {
+    if lhs.expr == rhs.expr || lhs.expr.egglog_equal(rhs.expr) {
+        return ParsedExpr {
+            expr: lhs.expr,
+            bounds: min_bounds(lhs.bounds, rhs.bounds),
+        };
+    }
+    if let (Some(lhs_max), Some(rhs_min)) = (lhs.bounds.max, rhs.bounds.min)
+        && lhs_max <= rhs_min
+    {
+        return lhs;
+    }
+    if let (Some(rhs_max), Some(lhs_min)) = (rhs.bounds.max, lhs.bounds.min)
+        && rhs_max <= lhs_min
+    {
+        return rhs;
+    }
+    if expr_is_offset_by_small_const(lhs.expr, rhs.expr) {
+        return rhs;
+    }
+    if expr_is_offset_by_small_const(rhs.expr, lhs.expr) {
+        return lhs;
+    }
+    ParsedExpr {
+        expr: lhs.expr.min(rhs.expr),
+        bounds: min_bounds(lhs.bounds, rhs.bounds),
+    }
+}
+
+fn reduce_max(lhs: ParsedExpr, rhs: ParsedExpr) -> ParsedExpr {
+    if lhs.expr == rhs.expr || lhs.expr.egglog_equal(rhs.expr) {
+        return ParsedExpr {
+            expr: lhs.expr,
+            bounds: max_bounds(lhs.bounds, rhs.bounds),
+        };
+    }
+    if let (Some(lhs_max), Some(rhs_min)) = (lhs.bounds.max, rhs.bounds.min)
+        && lhs_max <= rhs_min
+    {
+        return rhs;
+    }
+    if let (Some(rhs_max), Some(lhs_min)) = (rhs.bounds.max, lhs.bounds.min)
+        && rhs_max <= lhs_min
+    {
+        return lhs;
+    }
+    if expr_is_offset_by_small_const(lhs.expr, rhs.expr) {
+        return lhs;
+    }
+    if expr_is_offset_by_small_const(rhs.expr, lhs.expr) {
+        return rhs;
+    }
+    ParsedExpr {
+        expr: lhs.expr.max(rhs.expr),
+        bounds: max_bounds(lhs.bounds, rhs.bounds),
+    }
+}
+
+fn min_bounds(lhs: ExprBounds, rhs: ExprBounds) -> ExprBounds {
+    ExprBounds {
+        min: match (lhs.min, rhs.min) {
+            (Some(lhs), Some(rhs)) => Some(lhs.min(rhs)),
+            _ => None,
+        },
+        max: match (lhs.max, rhs.max) {
+            (Some(lhs), Some(rhs)) => Some(lhs.min(rhs)),
+            _ => None,
+        },
+    }
+}
+
+fn max_bounds(lhs: ExprBounds, rhs: ExprBounds) -> ExprBounds {
+    ExprBounds {
+        min: match (lhs.min, rhs.min) {
+            (Some(lhs), Some(rhs)) => Some(lhs.max(rhs)),
+            _ => None,
+        },
+        max: match (lhs.max, rhs.max) {
+            (Some(lhs), Some(rhs)) => Some(lhs.max(rhs)),
+            _ => None,
+        },
+    }
+}
+
+fn expr_is_offset_by_small_const(lhs: Expression, rhs: Expression) -> bool {
+    (1..=8).any(|delta| lhs.egglog_equal(rhs + delta))
+}
+
+fn split_add_const(expr: Expression) -> Option<(i64, Expression)> {
+    let terms = expr.terms.read();
+    if terms.len() >= 3 && terms.last() == Some(&Term::Add) {
+        if let Some(Term::Num(n)) = terms.first() {
+            return Some((*n, Expression::new(terms[1..terms.len() - 1].to_vec())));
+        }
+        if let Some(Term::Num(n)) = terms.get(terms.len() - 2) {
+            return Some((*n, Expression::new(terms[..terms.len() - 2].to_vec())));
+        }
+    }
+    None
+}
+
+fn simplify_add(lhs: BoundedExpr, rhs: BoundedExpr) -> BoundedExpr {
+    let expr = match (exact_value(lhs), exact_value(rhs)) {
+        (Some(0), _) => rhs.expr,
+        (_, Some(0)) => lhs.expr,
+        (Some(lhs), Some(rhs)) => Expression::from(lhs + rhs),
+        (_, Some(rhs)) => normalize_add_expr(lhs.expr, Expression::from(rhs)),
+        (Some(lhs), _) => normalize_add_expr(Expression::from(lhs), rhs.expr),
+        _ => normalize_add_expr(lhs.expr, rhs.expr),
+    };
+    with_bounds(expr, add_bounds(lhs.bounds, rhs.bounds))
+}
+
+fn simplify_sub(
+    lhs: BoundedExpr,
+    rhs: BoundedExpr,
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) -> BoundedExpr {
+    if same_expr_with_ranges(lhs.expr, rhs.expr, sym_ranges) {
+        return exact_expr(0);
+    }
+    let expr = match exact_value(rhs) {
+        Some(0) => lhs.expr,
+        Some(rhs_const) => {
+            if let Some((lhs_const, lhs_base)) = split_add_const(lhs.expr) {
+                normalize_expr(lhs_base + (lhs_const - rhs_const))
+            } else {
+                normalize_expr(lhs.expr - rhs_const)
+            }
+        }
+        None => normalize_expr(lhs.expr - rhs.expr),
+    };
+    with_bounds(expr, sub_bounds(lhs.bounds, rhs.bounds))
+}
+
+fn simplify_min(
+    lhs: BoundedExpr,
+    rhs: BoundedExpr,
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) -> BoundedExpr {
+    let bounds = min_bounds(lhs.bounds, rhs.bounds);
+    if same_expr_with_ranges(lhs.expr, rhs.expr, sym_ranges) {
+        return with_bounds(lhs.expr, bounds);
+    }
+    if let (Some(lhs_max), Some(rhs_min)) = (lhs.bounds.max, rhs.bounds.min)
+        && lhs_max <= rhs_min
+    {
+        return with_bounds(lhs.expr, bounds);
+    }
+    if let (Some(rhs_max), Some(lhs_min)) = (rhs.bounds.max, lhs.bounds.min)
+        && rhs_max <= lhs_min
+    {
+        return with_bounds(rhs.expr, bounds);
+    }
+    if let Some((lhs_const, lhs_base)) = split_add_const(lhs.expr)
+        && lhs_const >= 0
+        && same_expr_with_ranges(lhs_base, rhs.expr, sym_ranges)
+    {
+        return with_bounds(rhs.expr, bounds);
+    }
+    if let Some((rhs_const, rhs_base)) = split_add_const(rhs.expr)
+        && rhs_const >= 0
+        && same_expr_with_ranges(rhs_base, lhs.expr, sym_ranges)
+    {
+        return with_bounds(lhs.expr, bounds);
+    }
+    with_bounds(normalize_expr(lhs.expr.min(rhs.expr)), bounds)
+}
+
+fn simplify_max(
+    lhs: BoundedExpr,
+    rhs: BoundedExpr,
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) -> BoundedExpr {
+    let bounds = max_bounds(lhs.bounds, rhs.bounds);
+    if same_expr_with_ranges(lhs.expr, rhs.expr, sym_ranges) {
+        return with_bounds(lhs.expr, bounds);
+    }
+    if let (Some(lhs_max), Some(rhs_min)) = (lhs.bounds.max, rhs.bounds.min)
+        && lhs_max <= rhs_min
+    {
+        return with_bounds(rhs.expr, bounds);
+    }
+    if let (Some(rhs_max), Some(lhs_min)) = (rhs.bounds.max, lhs.bounds.min)
+        && rhs_max <= lhs_min
+    {
+        return with_bounds(lhs.expr, bounds);
+    }
+    if let Some((lhs_const, lhs_base)) = split_add_const(lhs.expr)
+        && lhs_const >= 0
+        && same_expr_with_ranges(lhs_base, rhs.expr, sym_ranges)
+    {
+        return with_bounds(lhs.expr, bounds);
+    }
+    if let Some((rhs_const, rhs_base)) = split_add_const(rhs.expr)
+        && rhs_const >= 0
+        && same_expr_with_ranges(rhs_base, lhs.expr, sym_ranges)
+    {
+        return with_bounds(rhs.expr, bounds);
+    }
+    with_bounds(normalize_expr(lhs.expr.max(rhs.expr)), bounds)
+}
+
+fn simplify_bound_expr(expr: Expression, sym_ranges: &FxHashMap<char, ExprBounds>) -> BoundedExpr {
+    let mut stack: Vec<BoundedExpr> = Vec::new();
+    let terms = expr.terms.read().clone();
+    for term in terms {
+        match term {
+            Term::Num(n) => stack.push(exact_expr(n)),
+            Term::Var(c) => stack.push(with_bounds(
+                Expression::from(c),
+                sym_ranges.get(&c).copied().unwrap_or_default(),
+            )),
+            Term::Add => {
+                let lhs = stack.pop().unwrap();
+                let rhs = stack.pop().unwrap();
+                stack.push(simplify_add(lhs, rhs));
+            }
+            Term::Sub => {
+                let lhs = stack.pop().unwrap();
+                let rhs = stack.pop().unwrap();
+                stack.push(simplify_sub(lhs, rhs, sym_ranges));
+            }
+            Term::Mul => {
+                let lhs = stack.pop().unwrap();
+                let rhs = stack.pop().unwrap();
+                let expr = match (exact_value(lhs), exact_value(rhs)) {
+                    (Some(0), _) | (_, Some(0)) => Expression::from(0),
+                    (Some(1), _) => rhs.expr,
+                    (_, Some(1)) => lhs.expr,
+                    (Some(lhs), Some(rhs)) => Expression::from(lhs * rhs),
+                    _ => normalize_mul_expr(lhs.expr, rhs.expr),
+                };
+                stack.push(with_bounds(expr, mul_bounds(lhs.bounds, rhs.bounds)));
+            }
+            Term::Div | Term::CeilDiv => {
+                let lhs = stack.pop().unwrap();
+                let rhs = stack.pop().unwrap();
+                let expr = match (term, exact_value(lhs), exact_value(rhs)) {
+                    (_, Some(0), _) => Expression::from(0),
+                    (_, _, Some(1)) => lhs.expr,
+                    (Term::Div, Some(lhs), Some(rhs)) if rhs != 0 => Expression::from(lhs / rhs),
+                    (Term::CeilDiv, Some(lhs), Some(rhs)) if rhs > 0 => {
+                        Expression::from(if lhs % rhs != 0 {
+                            lhs / rhs + 1
+                        } else {
+                            lhs / rhs
+                        })
+                    }
+                    (Term::Div, _, _) => normalize_expr(lhs.expr / rhs.expr),
+                    (Term::CeilDiv, _, _) => normalize_expr(lhs.expr.ceil_div(rhs.expr)),
+                    _ => unreachable!(),
+                };
+                stack.push(with_bounds(expr, div_bounds(lhs.bounds, rhs.bounds)));
+            }
+            Term::Mod => {
+                let lhs = stack.pop().unwrap();
+                let rhs = stack.pop().unwrap();
+                let expr = match (exact_value(lhs), exact_value(rhs)) {
+                    (Some(0), _) | (_, Some(1)) => Expression::from(0),
+                    (Some(lhs), Some(rhs)) if rhs != 0 => Expression::from(lhs % rhs),
+                    _ => normalize_expr(lhs.expr % rhs.expr),
+                };
+                stack.push(with_bounds(expr, mod_bounds(lhs.bounds, rhs.bounds)));
+            }
+            Term::Min => {
+                let lhs = stack.pop().unwrap();
+                let rhs = stack.pop().unwrap();
+                stack.push(simplify_min(lhs, rhs, sym_ranges));
+            }
+            Term::Max => {
+                let lhs = stack.pop().unwrap();
+                let rhs = stack.pop().unwrap();
+                stack.push(simplify_max(lhs, rhs, sym_ranges));
+            }
+            term @ (Term::And | Term::Or | Term::Gte | Term::Lt) => {
+                let lhs = stack.pop().unwrap();
+                let rhs = stack.pop().unwrap();
+                let expr = match (term, exact_value(lhs), exact_value(rhs)) {
+                    (Term::And, Some(lhs), Some(rhs)) => {
+                        Expression::from((lhs != 0 && rhs != 0) as i64)
+                    }
+                    (Term::And, _, _) => normalize_expr(lhs.expr & rhs.expr),
+                    (Term::Or, Some(lhs), Some(rhs)) => {
+                        Expression::from((lhs != 0 || rhs != 0) as i64)
+                    }
+                    (Term::Or, _, _) => normalize_expr(lhs.expr | rhs.expr),
+                    (Term::Gte, Some(lhs), Some(rhs)) => Expression::from((lhs >= rhs) as i64),
+                    (Term::Gte, _, _) => normalize_expr(lhs.expr.gte(rhs.expr)),
+                    (Term::Lt, Some(lhs), Some(rhs)) => Expression::from((lhs < rhs) as i64),
+                    (Term::Lt, _, _) => normalize_expr(lhs.expr.lt(rhs.expr)),
+                    _ => unreachable!(),
+                };
+                stack.push(with_bounds(expr, bool_bounds()));
+            }
+        }
+    }
+    stack
+        .pop()
+        .unwrap_or(with_bounds(expr, ExprBounds::default()))
+}
+
+/// Split `Head(body)` into `(head, body)`.
+fn split_head(expr: &str) -> Option<(&str, &str)> {
+    let open = expr.find('(')?;
+    if !expr.ends_with(')') {
+        return None;
+    }
+    Some((&expr[..open], &expr[open + 1..expr.len() - 1]))
+}
+
+/// Pull out the first single- or double-quoted token from a sympy arg list.
+fn extract_first_quoted(expr: &str) -> Option<String> {
+    let bytes = expr.as_bytes();
+    let mut i = 0;
+    while i < bytes.len() {
+        let c = bytes[i] as char;
+        if c == '\'' || c == '"' {
+            let quote = c;
+            let start = i + 1;
+            i += 1;
+            while i < bytes.len() && bytes[i] as char != quote {
+                i += 1;
+            }
+            return Some(expr[start..i].to_string());
+        }
+        i += 1;
+    }
+    None
+}
+
+/// Split a sympy-style argument list at top-level commas, respecting nested
+/// parens and quoted strings. Drops `key=value` kwargs.
+fn split_top_level_args(expr: &str) -> Vec<&str> {
+    let mut out = Vec::new();
+    let bytes = expr.as_bytes();
+    let mut depth = 0;
+    let mut in_quote: Option<char> = None;
+    let mut start = 0;
+    for (i, &b) in bytes.iter().enumerate() {
+        let c = b as char;
+        match in_quote {
+            Some(q) => {
+                if c == q {
+                    in_quote = None;
+                }
+            }
+            None => match c {
+                '\'' | '"' => in_quote = Some(c),
+                '(' | '[' => depth += 1,
+                ')' | ']' => depth -= 1,
+                ',' if depth == 0 => {
+                    let part = expr[start..i].trim();
+                    if !part.is_empty() && !looks_like_kwarg(part) {
+                        out.push(part);
+                    }
+                    start = i + 1;
+                }
+                _ => {}
+            },
+        }
+    }
+    let part = expr[start..].trim();
+    if !part.is_empty() && !looks_like_kwarg(part) {
+        out.push(part);
+    }
+    out
+}
+
+fn looks_like_kwarg(part: &str) -> bool {
+    if let Some(eq) = part.find('=') {
+        let key = part[..eq].trim();
+        return !key.is_empty() && key.chars().all(|c| c == '_' || c.is_ascii_alphanumeric());
+    }
+    false
+}

crates/luminal_python/rust/src/pt2_util.rs

@@ -1,5 +1,9 @@
 use luminal::prelude::*;
 
+fn same_dim(lhs: Expression, rhs: Expression) -> bool {
+    lhs == rhs || lhs.simplify() == rhs.simplify() || lhs.egglog_equal(rhs)
+}
+
 /// Binary operation type.
 #[derive(Clone, Copy)]
 pub enum BinaryOp {
@@ -51,7 +55,7 @@ pub fn broadcast_binary(mut a: GraphTensor, mut b: GraphTensor) -> (GraphTensor,
         let a_dim = a.shape.dims[i];
         let b_dim = b.shape.dims[i];
 
-        if a_dim == b_dim {
+        if same_dim(a_dim, b_dim) {
             continue;
         }
 

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

@@ -1,11 +1,40 @@
 use anyhow::Result;
 use luminal::prelude::*;
+use rustc_hash::FxHashMap;
 
+use crate::pt2_expr::{ExprBounds, canonical_equal_expr, same_expr_with_ranges, sym_char_ranges};
 use crate::pt2_schema::*;
 use crate::pt2_util::*;
 
 use super::Translator;
 
+fn normalize_equal_dims(
+    a: &mut GraphTensor,
+    b: &mut GraphTensor,
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) {
+    for i in 0..a.shape.len() {
+        let lhs = a.shape.dims[i];
+        let rhs = b.shape.dims[i];
+        if let Some(canonical) = canonical_equal_expr(lhs, rhs, sym_ranges) {
+            a.shape.dims[i] = canonical;
+            b.shape.dims[i] = canonical;
+        }
+    }
+}
+
+fn same_dims(
+    lhs: &[Expression],
+    rhs: &[Expression],
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) -> bool {
+    lhs.len() == rhs.len()
+        && lhs
+            .iter()
+            .zip(rhs.iter())
+            .all(|(lhs, rhs)| same_expr_with_ranges(*lhs, *rhs, sym_ranges))
+}
+
 impl<'a> Translator<'a> {
     pub(crate) fn translate_binary_op(&mut self, node: &Node, op: BinaryOp) -> Result<GraphTensor> {
         let a = self.get_input_tensor(node, 0)?;
@@ -13,7 +42,18 @@ impl<'a> Translator<'a> {
         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);
+            let (mut a, mut b) = broadcast_binary(a, b);
+            let sym_ranges = sym_char_ranges(&self.sym_map);
+            normalize_equal_dims(&mut a, &mut b, &sym_ranges);
+            let lhs_dims = a.dims();
+            let rhs_dims = b.dims();
+            if !same_dims(&lhs_dims, &rhs_dims, &sym_ranges) {
+                anyhow::bail!(
+                    "binary op {} still has mismatched dims after broadcast: lhs={lhs_dims:?} rhs={rhs_dims:?} inputs={:?}",
+                    node.target,
+                    node.inputs
+                );
+            }
             Ok(match op {
                 BinaryOp::Add => a + b,
                 BinaryOp::Mul => a * b,
@@ -21,6 +61,12 @@ impl<'a> Translator<'a> {
                 BinaryOp::Div => a / b,
             })
         } else {
+            if let Some(f) = arg1.as_float() {
+                return Ok(self.apply_scalar_op(a, f as f32, op));
+            }
+            if let Some(expr) = self.resolve_arg_as_expression(arg1) {
+                return Ok(self.apply_symbolic_scalar_op(a, expr, op));
+            }
             let val = self.get_float_arg(node, 1)? as f32;
             Ok(self.apply_scalar_op(a, val, op))
         }
@@ -32,6 +78,13 @@ impl<'a> Translator<'a> {
         op: BinaryOp,
     ) -> Result<GraphTensor> {
         let a = self.get_input_tensor(node, 0)?;
+        let arg1 = &node.inputs[1].arg;
+        if let Some(f) = arg1.as_float() {
+            return Ok(self.apply_scalar_op(a, f as f32, op));
+        }
+        if let Some(expr) = self.resolve_arg_as_expression(arg1) {
+            return Ok(self.apply_symbolic_scalar_op(a, expr, op));
+        }
         let val = self.get_float_arg(node, 1)? as f32;
         Ok(self.apply_scalar_op(a, val, op))
     }
@@ -54,4 +107,47 @@ impl<'a> Translator<'a> {
             BinaryOp::Div => a / scalar,
         }
     }
+
+    pub(crate) fn apply_symbolic_scalar_op(
+        &mut self,
+        a: GraphTensor,
+        val: Expression,
+        op: BinaryOp,
+    ) -> GraphTensor {
+        match op {
+            BinaryOp::Add => a + val,
+            BinaryOp::Mul => a * val,
+            BinaryOp::Sub => a - val,
+            BinaryOp::Div => a / val,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::pt2_expr::simplify_expr_with_ranges;
+
+    #[test]
+    fn simplifies_mark_dynamic_slice_shapes_using_lower_bound() {
+        let a = Expression::from('a');
+        let lhs = (a.min(1) + a).min(a + 1) - 1;
+        let rhs = (a.min(1) + a).min(a);
+        let sym_ranges = [(
+            'a',
+            ExprBounds {
+                min: Some(2),
+                max: None,
+            },
+        )]
+        .into_iter()
+        .collect::<FxHashMap<_, _>>();
+
+        let lhs_simplified = simplify_expr_with_ranges(lhs, &sym_ranges);
+        let rhs_simplified = simplify_expr_with_ranges(rhs, &sym_ranges);
+
+        assert_eq!(lhs_simplified, Expression::from('a'));
+        assert_eq!(rhs_simplified, Expression::from('a'));
+        assert!(same_expr_with_ranges(lhs, rhs, &sym_ranges));
+    }
 }

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

@@ -119,30 +119,147 @@ impl<'a> Translator<'a> {
             "torch.ops.aten.argsort.default" => self.translate_argsort(node)?,
 
             // Matmul
-            "torch.ops.aten.mm.default" | "torch.ops.aten.bmm.default" => {
+            "torch.ops.aten.mm.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)
             }
+            // bmm: batched 3-D matmul. Fast path under cuda + F32 when
+            // B was produced by a permute([0, 2, 1]) (i.e. `T @ T.T` —
+            // the DLRM pairwise-interaction pattern): route to
+            // `matmul_3d_t` with the original (B, F, D) tensor, which
+            // uses the fused Matmul2DKernel and avoids the
+            // expand+mul+sum-reduce decomposition that produces ~25
+            // small kernels per bmm.
+            "torch.ops.aten.bmm.default" => {
+                let a = self.get_input_tensor(node, 0)?;
+                let b = self.get_input_tensor(node, 1)?;
+
+                let b_src = node.inputs.get(1).and_then(|n| n.arg.as_tensor_name())
+                    .and_then(|n| self.transpose_2d_source.get(n).cloned());
+
+                let f32_all = a.dtype == DType::F32 && b.dtype == DType::F32;
+                let backend_is_cuda = true; // LUMINAL_BACKEND_CUDA replaced by cfg!(feature="cuda") gate above.
+                if cfg!(feature = "cuda")
+                    && backend_is_cuda
+                    && f32_all
+                    && a.shape.dims.len() == 3
+                    && b.shape.dims.len() == 3
+                    && let Some(orig_name) = b_src
+                    && let Some(orig_b) = self.tensors.get(&orig_name).copied()
+                    && orig_b.shape.dims.len() == 3
+                {
+                    // a: (B, M, K), orig_b: (B, N, K) — matmul_3d_t does
+                    // a @ orig_b.t() = (B, M, K) @ (B, K, N) = (B, M, N).
+                    luminal_cuda_lite::kernel::matmul_3d_t(a, orig_b)
+                } else {
+                    let (a, b) = ensure_same_dtype(a, b);
+                    a.matmul(b)
+                }
+            }
 
             // addmm: beta*input + alpha*(mat1 @ mat2)
+            //
+            // Fast path (CUDA, the common nn.Linear case): when
+            //   * shapes are 2-D F32, alpha=beta=1
+            //   * mat2 was produced by `aten.permute([1,0])` of a 2-D
+            //     tensor (`weight.t()` from nn.Linear)
+            //   * bias is 1-D
+            // we lower to the fused `linear_bias` kernel using the
+            // *original* (N,K) weight — bypassing the
+            // expand+mul+sum-reduce decomposition that otherwise
+            // produces ~25 small kernels per Linear layer (~3.7 ms on
+            // tiny shapes due to launch overhead).
+            //
+            // The transpose detection comes from
+            // `translate_permute`, which populates
+            // `transpose_2d_source` whenever it sees a 2-D permute.
             "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
+
+                let mat2_src = node.inputs.get(2).and_then(|n| n.arg.as_tensor_name())
+                    .and_then(|n| self.transpose_2d_source.get(n).cloned());
+
+                let unit_scale = (alpha - 1.0).abs() < 1e-7 && (beta - 1.0).abs() < 1e-7;
+                let f32_all = mat1.dtype == DType::F32
+                    && mat2.dtype == DType::F32
+                    && input.dtype == DType::F32;
+                let two_d = mat1.shape.dims.len() == 2 && mat2.shape.dims.len() == 2;
+                let bias_is_1d = input.shape.dims.len() == 1;
+
+                let backend_is_cuda = true; // LUMINAL_BACKEND_CUDA replaced by cfg!(feature="cuda") gate above.
+                if cfg!(feature = "cuda")
+                    && backend_is_cuda
+                    && two_d
+                    && f32_all
+                    && unit_scale
+                    && bias_is_1d
+                    && let Some(weight_name) = mat2_src
+                    && let Some(orig_weight) = self.tensors.get(&weight_name).copied()
+                    && orig_weight.shape.dims.len() == 2
+                {
+                    // Forward-looking fusion: if this addmm has exactly one
+                    // consumer and that consumer is `aten.relu.default` or
+                    // `aten.sigmoid.default`, emit the fused
+                    // `linear_bias_relu` / `linear_bias_sigmoid` kernel and
+                    // mark the consumer as absorbed so we don't emit a
+                    // redundant unary op downstream. This collapses the
+                    // standard nn.Linear+ReLU MLP layer to one kernel.
+                    let addmm_out: Option<String> = node
+                        .outputs
+                        .first()
+                        .and_then(|o| o.as_tensor.as_ref().map(|t| t.name.clone()));
+                    let fuse_act = addmm_out
+                        .as_deref()
+                        .and_then(|n| self.unique_consumer(n));
+                    let (fused, absorbed) = match fuse_act {
+                        Some((target, out_name)) if target == "torch.ops.aten.relu.default" => {
+                            (
+                                Some(luminal_cuda_lite::kernel::linear_bias_relu(
+                                    mat1, orig_weight, input,
+                                )),
+                                Some(out_name),
+                            )
+                        }
+                        Some((target, out_name)) if target == "torch.ops.aten.sigmoid.default" => {
+                            (
+                                Some(luminal_cuda_lite::kernel::linear_bias_sigmoid(
+                                    mat1, orig_weight, input,
+                                )),
+                                Some(out_name),
+                            )
+                        }
+                        _ => (None, None),
+                    };
+                    if let Some(t) = fused {
+                        if let Some(out_name) = absorbed {
+                            self.absorbed_nodes.insert(out_name.clone());
+                            self.tensors.insert(out_name, t);
+                        }
+                        t
+                    } else {
+                        // No unary consumer to fuse; plain linear+bias.
+                        luminal_cuda_lite::kernel::linear_bias(mat1, orig_weight, input)
+                    }
+                } else {
+                    // Generic fallback (non-cuda, scaled, or unknown
+                    // mat2 source).
+                    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)?,
+            "torch.ops.aten.sum.dim_IntList" => self.translate_sum_with_embbag_peephole(node)?,
             "torch.ops.aten.mean.dim" => self.translate_reduction(node, ReductionOp::Mean)?,
             "torch.ops.aten.amax.default" => self.translate_reduction(node, ReductionOp::Max)?,
 
@@ -151,10 +268,25 @@ impl<'a> Translator<'a> {
             "torch.ops.aten.select.int" => self.translate_select(node)?,
             "torch.ops.aten.cat.default" => self.translate_cat(node)?,
             "torch.ops.aten.index.Tensor" => self.translate_index_tensor(node)?,
+            "torch.ops.aten.index_select.default" => self.translate_index_select(node)?,
 
             // Embedding
             "torch.ops.aten.embedding.default" => self.translate_embedding(node)?,
 
+            // EmbeddingBag (sum-pool, fixed bag size). PT2 export decomposes
+            // nn.EmbeddingBag → `_embedding_bag` (with backward) or
+            // `_embedding_bag_forward_only` (without). Both return the same
+            // 4-tuple `(output, offset2bag, bag_size, max_indices)` and only
+            // the first slot is consumed at inference. The two op variants
+            // are math-identical for the forward path, so route both through
+            // the same handler. The handler stores into `tensors` itself
+            // (multi-output op) so we return early afterwards.
+            "torch.ops.aten._embedding_bag_forward_only.default"
+            | "torch.ops.aten._embedding_bag.default" => {
+                self.translate_embedding_bag_forward_only(node)?;
+                return Ok(());
+            }
+
             // Softmax
             "torch.ops.aten._softmax.default" => {
                 let a = self.get_input_tensor(node, 0)?;
@@ -514,6 +646,51 @@ impl<'a> Translator<'a> {
         };
 
         if !output_name.is_empty() {
+            // Record the chain (FX target + first input name) keyed by
+            // output name so multi-node peepholes (e.g. the EmbBag fast
+            // path that detects sum ← view ← index_select) can walk
+            // back without re-scanning all parsed nodes.
+            //
+            // For variadic ops (e.g. `aten.cat.default` whose first
+            // arg is `as_tensors`) fall back to the first entry of the
+            // variadic tensor list. The DLRM PairwiseDot peephole
+            // needs `node_chain[cat]` to walk back from `bmm → cat`.
+            let first_input_name: Option<String> = node
+                .inputs
+                .first()
+                .and_then(|i| {
+                    i.arg
+                        .as_tensor_name()
+                        .map(|s| s.to_string())
+                        .or_else(|| {
+                            i.arg
+                                .as_tensors()
+                                .and_then(|ts| ts.first().map(|tn| tn.name.clone()))
+                        })
+                });
+            if let Some(first_input) = first_input_name {
+                self.node_chain.insert(
+                    output_name.clone(),
+                    (node.target.clone(), first_input),
+                );
+            }
+            // Also record the full input-name list (in order, including
+            // entries that come from `as_tensors` for variadic ops like
+            // `aten.cat`). Used by the DLRM PairwiseDot peephole which
+            // needs all cat inputs and both bmm inputs.
+            let mut all_inputs: Vec<String> = Vec::new();
+            for inp in &node.inputs {
+                if let Some(names) = inp.arg.as_tensors() {
+                    for tn in names {
+                        all_inputs.push(tn.name.clone());
+                    }
+                } else if let Some(name) = inp.arg.as_tensor_name() {
+                    all_inputs.push(name.to_string());
+                }
+            }
+            if !all_inputs.is_empty() {
+                self.op_inputs.insert(output_name.clone(), all_inputs);
+            }
             self.tensors.insert(output_name, result);
         }
         Ok(())
@@ -521,6 +698,69 @@ impl<'a> Translator<'a> {
 }
 
 impl<'a> Translator<'a> {
+    /// Peephole for the DLRM-v3 embedding-bag pattern:
+    ///   `sum(dim=[1], keepdim=False)( view([?, L, D])( index_select(W, 0, IDX) ) )`
+    /// substitutes the fused `embedding_bag_sum_kernel(W, IDX.view(?, L))`
+    /// — same kernel as the hand-rolled DLRM example uses. Falls back to
+    /// the generic reduction path when the chain doesn't match.
+    pub(crate) fn translate_sum_with_embbag_peephole(
+        &mut self,
+        node: &Node,
+    ) -> Result<GraphTensor> {
+        let dims = self.get_ints_arg(node, 1).unwrap_or_default();
+        let keepdim = self.get_bool_arg(node, 2).unwrap_or(false);
+        let backend_is_cuda = true; // LUMINAL_BACKEND_CUDA replaced by cfg!(feature="cuda") gate above.
+
+        // Only attempt fast-path under cuda + the specific sum(dim=[1]) pattern.
+        if cfg!(feature = "cuda")
+            && backend_is_cuda
+            && dims.len() == 1
+            && dims[0] == 1
+            && !keepdim
+            && let Some(sum_input_name) = node.inputs.first().and_then(|i| i.arg.as_tensor_name())
+            && let Some((view_target, view_src)) = self.node_chain.get(sum_input_name).cloned()
+            && view_target == "torch.ops.aten.view.default"
+            && let Some((is_target, is_src)) = self.node_chain.get(&view_src).cloned()
+            && is_target == "torch.ops.aten.index_select.default"
+            // Pull the FX index_select node so we can grab its dim + index args.
+            && let Some(is_node) = self.parsed.program.graph_module.graph.nodes.iter()
+                .find(|n| n.outputs.first()
+                    .and_then(|o| o.as_tensor.as_ref())
+                    .map(|t| t.name == view_src)
+                    .unwrap_or(false))
+        {
+            let weight = self.tensors.get(&is_src).copied();
+            let idx_name = is_node.inputs.get(2).and_then(|i| i.arg.as_tensor_name());
+            let is_dim = self.get_int_arg(is_node, 1).unwrap_or(-1);
+            let in_tensor = self.tensors.get(sum_input_name).copied();
+            if is_dim == 0
+                && let Some(w) = weight
+                && let Some(idx_n) = idx_name
+                && let Some(idx) = self.tensors.get(idx_n).copied()
+                && let Some(inp) = in_tensor
+                && w.shape.dims.len() == 2
+                && idx.shape.dims.len() == 1
+                && inp.shape.dims.len() == 3
+                // ensure view's middle dim == idx's bag dim divides idx total
+                && inp.dtype == DType::F32
+                && w.dtype == DType::F32
+            {
+                let l = inp.shape.dims[1];
+                let kb = inp.shape.dims[0];
+                let d = inp.shape.dims[2];
+                // Reshape flat indices (K*B*L,) to (K*B, L).
+                let idx_2d = reshape_tensor(idx, vec![kb, l]);
+                // embedding_bag_sum_kernel expects (n_emb, d) weights +
+                // (batch, bag) indices, returns (batch, d).
+                let _ = d; // d already encoded in `w.shape.dims[1]`
+                return Ok(luminal_cuda_lite::kernel::embedding_bag_sum_kernel(w, idx_2d));
+            }
+        }
+
+        // Generic fallback.
+        self.translate_reduction(node, ReductionOp::Sum)
+    }
+
     fn translate_scalar_comparison(
         &mut self,
         node: &Node,

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

@@ -17,6 +17,7 @@ use anyhow::{Context, Result};
 use luminal::graph::Graph;
 use luminal::prelude::*;
 
+use crate::pt2_expr::parse_sympy_expr_with_ranges;
 use crate::pt2_parser::{InputKind, ParsedPT2, SymDimMap};
 use crate::pt2_schema::*;
 use crate::pt2_util;
@@ -50,6 +51,42 @@ pub(crate) struct Translator<'a> {
     pub(crate) output_ids: Vec<(String, NodeIndex)>,
     /// Extra tensor metadata from inlined subgraphs.
     pub(crate) extra_tensor_values: HashMap<String, TensorMeta>,
+    /// Peephole: maps an output-tensor name produced by a `permute([1,0])`
+    /// (i.e. a 2-D transpose) back to its input-tensor name. Used by the
+    /// addmm dispatch to detect `aten.addmm(bias, x, weight.t())` and
+    /// route it through the fused `Matmul2DKernel` (`matmul_2d_t`) with
+    /// the original weight, instead of through the generic
+    /// expand+mul+sum decomposition that materializes ~25 small kernels.
+    pub(crate) transpose_2d_source: HashMap<String, String>,
+    /// Trace each emitted node by its first output's name → (FX target,
+    /// first input's name). Used by the EmbBag peephole to walk back
+    /// `sum.dim_IntList → aten.view.default → aten.index_select.default`
+    /// and substitute the fused `embedding_bag_sum_kernel` for the slow
+    /// expand+gather decomposition. Populated by `record_node_chain`
+    /// after dispatching each op.
+    pub(crate) node_chain: HashMap<String, (String, String)>,
+    /// Per-node side table mapping the primary output name → list of all
+    /// input tensor names (in order). Lets multi-input peepholes — e.g.
+    /// `index.Tensor(bmm(cat([…]), permute(cat([…]))), [None, li, lj])`
+    /// → `dlrm_pairwise_dot_lower_tri([…])` — walk back through cat and
+    /// bmm without re-scanning the FX node array. Populated alongside
+    /// `node_chain` after each translated op.
+    pub(crate) op_inputs: HashMap<String, Vec<String>>,
+    /// Tensor name → list of *consumer output-tensor names*. Built once at
+    /// the start of `translate_graph` from the parsed FX node array.
+    /// (The pt2_schema's `Node` has no name field; nodes are identified by
+    /// their primary output tensor name.) Used by forward-looking fusions:
+    /// e.g. when the addmm fast path fires and the single consumer is
+    /// `relu`/`sigmoid`, we emit the fused `linear_bias_relu`/
+    /// `linear_bias_sigmoid` kernel and absorb the consumer node via
+    /// `absorbed_nodes`.
+    pub(crate) consumers: HashMap<String, Vec<String>>,
+    /// Set of *output tensor names* whose producing FX node was absorbed
+    /// into an earlier node's emission (e.g. a `relu` folded into
+    /// `linear_bias` by the addmm fast path). `translate_graph` short-
+    /// circuits these nodes. The absorbed node's output tensor must be
+    /// pre-populated under its name by the absorbing node.
+    pub(crate) absorbed_nodes: std::collections::HashSet<String>,
 }
 
 impl<'a> Translator<'a> {
@@ -60,17 +97,112 @@ impl<'a> Translator<'a> {
             graph: Graph::new(),
             tensors: HashMap::new(),
             sym_map,
+            transpose_2d_source: HashMap::new(),
+            node_chain: HashMap::new(),
+            op_inputs: HashMap::new(),
+            consumers: HashMap::new(),
+            absorbed_nodes: std::collections::HashSet::new(),
             user_input_ids: Vec::new(),
             output_ids: Vec::new(),
             extra_tensor_values: HashMap::new(),
         })
     }
 
+    /// Helper: extract a node's primary output tensor name. Nodes in the
+    /// pt2 schema are identified by this (no separate name field).
+    fn node_out_name(node: &Node) -> Option<String> {
+        node.outputs
+            .first()
+            .and_then(|o| o.as_tensor.as_ref().map(|t| t.name.clone()))
+            .or_else(|| {
+                node.outputs
+                    .first()
+                    .and_then(|o| o.as_tensors.as_ref())
+                    .and_then(|ts| ts.first().map(|t| t.name.clone()))
+            })
+    }
+
+    /// Build a `tensor_name → consumer output-tensor names` map for the
+    /// parsed graph. One pass over all FX nodes; each node contributes to
+    /// the `consumers` entry of every tensor it reads (Argument::Tensor or
+    /// Argument::Tensors). The consumer is keyed by its primary output
+    /// tensor name. Used by forward-looking fast paths to detect when an
+    /// op's output has a single downstream consumer of a known kind
+    /// (relu/sigmoid) and emit a fused kernel that absorbs the consumer.
+    fn build_consumers(&mut self) {
+        let nodes = &self.parsed.program.graph_module.graph.nodes;
+        for node in nodes {
+            let Some(consumer_out) = Self::node_out_name(node) else {
+                continue;
+            };
+            for inp in &node.inputs {
+                match &inp.arg {
+                    Argument::Tensor(t) => {
+                        self.consumers
+                            .entry(t.as_tensor.name.clone())
+                            .or_default()
+                            .push(consumer_out.clone());
+                    }
+                    Argument::Tensors(ts) => {
+                        for t in &ts.as_tensors {
+                            self.consumers
+                                .entry(t.name.clone())
+                                .or_default()
+                                .push(consumer_out.clone());
+                        }
+                    }
+                    Argument::OptionalTensors(ots) => {
+                        for ot in &ots.as_optional_tensors {
+                            if let crate::pt2_schema::OptionalTensorEntry::Tensor(t) = ot {
+                                self.consumers
+                                    .entry(t.as_tensor.name.clone())
+                                    .or_default()
+                                    .push(consumer_out.clone());
+                            }
+                        }
+                    }
+                    _ => {}
+                }
+            }
+        }
+    }
+
+    /// Return the unique consumer of `tensor_name`, if there is exactly one
+    /// and we can find the corresponding FX node. Returns `(target,
+    /// output_tensor_name)`. None if the consumer set is empty, has more
+    /// than one entry, or the FX node lookup fails.
+    pub(crate) fn unique_consumer(&self, tensor_name: &str) -> Option<(String, String)> {
+        let consumers = self.consumers.get(tensor_name)?;
+        if consumers.len() != 1 {
+            return None;
+        }
+        let consumer_out = &consumers[0];
+        let node = self
+            .parsed
+            .program
+            .graph_module
+            .graph
+            .nodes
+            .iter()
+            .find(|n| Self::node_out_name(n).as_deref() == Some(consumer_out.as_str()))?;
+        Some((node.target.clone(), consumer_out.clone()))
+    }
+
     fn translate_graph(&mut self) -> Result<()> {
         self.create_inputs()?;
+        self.build_consumers();
 
         let nodes = &self.parsed.program.graph_module.graph.nodes;
         for (i, node) in nodes.iter().enumerate() {
+            // Skip nodes whose translation was absorbed by an earlier
+            // node's fast path (e.g. a `relu` folded into a fused
+            // `linear_bias_relu`). The absorbing node has already
+            // populated `tensors` under this node's output name.
+            if let Some(out_name) = Self::node_out_name(node)
+                && self.absorbed_nodes.contains(&out_name)
+            {
+                continue;
+            }
             self.translate_node(node)
                 .with_context(|| format!("Failed to translate node {i}: {}", node.target))?;
         }
@@ -89,9 +221,64 @@ impl<'a> Translator<'a> {
             self.output_ids.push((name.clone(), tensor.id));
         }
 
+        // Post-translation dead-code elimination. luminal's egglog DOES
+        // prune unreachable subgraphs in the common case (e.g. an unused
+        // `x*2.0` next to a returned `x+1.0`), but in some patterns the
+        // optimizer holds onto subgraphs that were created and then
+        // superseded by a translator peephole — most notably the DLRM
+        // PairwiseDot path where `index.Tensor(bmm(cat(...), perm(cat(...))), ...)`
+        // is replaced with a fused custom op but the original bmm/cat
+        // pad-and-add chain remains in the HLIR. Walk back from every
+        // `Output` HLIR node, mark reachable producers, and drop the rest.
+        // Preserves `Input` nodes unconditionally so the runtime's input
+        // signature stays intact even when an input is unused (a few
+        // models pass dead constants alongside live tensors).
+        self.dce();
+
         Ok(())
     }
 
+    /// Sweep the HLIR graph: remove every node not reachable backward
+    /// from an `Output` HLIR sink. Inputs are kept regardless so the
+    /// runtime input contract is preserved.
+    fn dce(&mut self) {
+        use luminal::hlir::{Input, Output};
+        use petgraph::Direction;
+        use std::collections::HashSet;
+
+        let mut keep: HashSet<NodeIndex> = HashSet::new();
+        let mut stack: Vec<NodeIndex> = Vec::new();
+        let node_ids: Vec<NodeIndex> = self.graph.graph.node_indices().collect();
+        for n in &node_ids {
+            if self.graph.try_get_op::<Output>(*n).is_some() {
+                if keep.insert(*n) {
+                    stack.push(*n);
+                }
+            }
+            if self.graph.try_get_op::<Input>(*n).is_some() {
+                keep.insert(*n);
+            }
+        }
+        while let Some(n) = stack.pop() {
+            // Walk incoming edges — operands of `n`.
+            let preds: Vec<NodeIndex> = self
+                .graph
+                .graph
+                .neighbors_directed(n, Direction::Incoming)
+                .collect();
+            for pred in preds {
+                if keep.insert(pred) {
+                    stack.push(pred);
+                }
+            }
+        }
+        for n in node_ids {
+            if !keep.contains(&n) {
+                self.graph.graph.remove_node(n);
+            }
+        }
+    }
+
     fn create_inputs(&mut self) -> Result<()> {
         let inputs = self.parsed.classify_inputs();
         for input in &inputs {
@@ -279,13 +466,13 @@ impl<'a> Translator<'a> {
             .with_context(|| format!("Node {} missing input {idx}", node.target))?
             .arg;
         if let Some(ints) = arg.as_ints() {
-            return Ok(ints.iter().map(|&v| Expression::from(v as usize)).collect());
+            return Ok(ints.iter().map(|&v| Expression::from(v)).collect());
         }
         if let Some(entries) = arg.as_sym_ints() {
             return entries
                 .iter()
                 .map(|entry| match entry {
-                    SymIntEntry::Int(i) => Ok(Expression::from(i.as_int as usize)),
+                    SymIntEntry::Int(i) => Ok(Expression::from(i.as_int)),
                     SymIntEntry::Name(s) => self
                         .resolve_sym_int(&s.as_name)
                         .with_context(|| format!("Cannot resolve sym_int: {}", s.as_name)),
@@ -318,17 +505,13 @@ impl<'a> Translator<'a> {
 
     pub(crate) fn dim_size_to_expr(&self, dim: &DimSize) -> Result<Expression> {
         match dim {
-            DimSize::Int(i) => Ok(Expression::from(i.as_int as usize)),
-            DimSize::Expr(e) => {
-                let sym_name = crate::pt2_parser::extract_symbol_name_pub(&e.as_expr.expr_str)
-                    .with_context(|| format!("Cannot parse symbol: {}", e.as_expr.expr_str))?;
-                let c = self
-                    .sym_map
-                    .sym_to_char
-                    .get(&sym_name)
-                    .with_context(|| format!("Unknown symbol: {sym_name}"))?;
-                Ok(Expression::from(*c))
-            }
+            DimSize::Int(i) => Ok(Expression::from(i.as_int)),
+            DimSize::Expr(e) => self.resolve_expr_value(&e.as_expr).with_context(|| {
+                format!(
+                    "Cannot resolve symbolic dimension expression: {}",
+                    e.as_expr.expr_str
+                )
+            }),
         }
     }
 
@@ -339,10 +522,9 @@ impl<'a> Translator<'a> {
                 .get("as_expr")
                 .and_then(|e| e.get("expr_str"))
                 .and_then(|s| s.as_str())
-                && let Some(sym) = crate::pt2_parser::extract_symbol_name_pub(expr_str)
-                && let Some(&c) = self.sym_map.sym_to_char.get(&sym)
+                && let Some(expr) = self.resolve_expr_str(expr_str)
             {
-                return Some(Expression::from(c));
+                return Some(expr);
             }
             if let Some(hint) = val
                 .get("as_expr")
@@ -350,7 +532,7 @@ impl<'a> Translator<'a> {
                 .and_then(|h| h.get("as_int"))
                 .and_then(|v| v.as_i64())
             {
-                return Some(Expression::from(hint as usize));
+                return Some(Expression::from(hint));
             }
         }
         None
@@ -358,21 +540,32 @@ impl<'a> Translator<'a> {
 
     pub(crate) fn resolve_arg_as_expression(&self, arg: &Argument) -> Option<Expression> {
         if let Some(v) = arg.as_int() {
-            return Some(Expression::from(v as usize));
+            return Some(Expression::from(v));
         }
         if let Some(name) = arg.as_sym_int_name() {
             return self.resolve_sym_int(name);
         }
         if let Argument::Expr(e) = arg {
-            if let Some(sym) = crate::pt2_parser::extract_symbol_name_pub(&e.as_expr.expr_str)
-                && let Some(&c) = self.sym_map.sym_to_char.get(&sym)
-            {
-                return Some(Expression::from(c));
-            }
-            if let Some(hint) = e.as_expr.hint.as_ref().and_then(|h| h.as_int()) {
-                return Some(Expression::from(hint as usize));
-            }
+            return self.resolve_expr_value(&e.as_expr);
         }
         None
     }
+
+    pub(crate) fn resolve_expr_str(&self, expr_str: &str) -> Option<Expression> {
+        parse_sympy_expr_with_ranges(expr_str, &self.sym_map.sym_to_char, &self.sym_map.ranges)
+            .or_else(|| {
+                crate::pt2_parser::extract_symbol_name_pub(expr_str)
+                    .and_then(|sym| self.sym_map.sym_to_char.get(&sym).copied())
+                    .map(Expression::from)
+            })
+    }
+
+    pub(crate) fn resolve_expr_value(&self, expr: &ExprValue) -> Option<Expression> {
+        self.resolve_expr_str(&expr.expr_str).or_else(|| {
+            expr.hint
+                .as_ref()
+                .and_then(|h| h.as_int())
+                .map(Expression::from)
+        })
+    }
 }

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

@@ -1,6 +1,8 @@
 use anyhow::{Context, Result, bail};
 use luminal::prelude::*;
+use rustc_hash::FxHashMap;
 
+use crate::pt2_expr::{ExprBounds, canonical_equal_expr, sym_char_ranges};
 use crate::pt2_schema::*;
 use crate::pt2_util::*;
 
@@ -11,6 +13,25 @@ const SCATTER_DIM_ARG: usize = 1;
 const SCATTER_INDEX_ARG: usize = 2;
 const SCATTER_VALUE_ARG: usize = 3;
 
+fn normalize_concat_dims(
+    lhs: &mut GraphTensor,
+    rhs: &mut GraphTensor,
+    skip_dim: Option<usize>,
+    sym_ranges: &FxHashMap<char, ExprBounds>,
+) {
+    for i in 0..lhs.shape.len() {
+        if Some(i) == skip_dim {
+            continue;
+        }
+        let lhs_dim = lhs.shape.dims[i];
+        let rhs_dim = rhs.shape.dims[i];
+        if let Some(canonical) = canonical_equal_expr(lhs_dim, rhs_dim, sym_ranges) {
+            lhs.shape.dims[i] = canonical;
+            rhs.shape.dims[i] = canonical;
+        }
+    }
+}
+
 impl<'a> Translator<'a> {
     pub(crate) fn translate_reshape(&mut self, node: &Node) -> Result<GraphTensor> {
         let a = self.get_input_tensor(node, 0)?;
@@ -51,6 +72,24 @@ impl<'a> Translator<'a> {
             .iter()
             .map(|&d| normalize_dim(d, a.shape.len()))
             .collect();
+        // Record matmul-compatible inner-axis transposes so addmm /
+        // bmm can route them through the fused Matmul2DKernel /
+        // matmul_3d_t with the *original* input. The view-transposed
+        // tensor has non-contiguous strides that the SGEMM kernel
+        // doesn't honor, so we need the original. We recognize two
+        // patterns:
+        //   * 2-D permute [1, 0] — `weight.t()` from nn.Linear
+        //   * 3-D permute [0, 2, 1] — `T.transpose(1, 2)` for bmm
+        let is_inner_transpose = (axes == [1usize, 0usize] && a.shape.dims.len() == 2)
+            || (axes == [0usize, 2usize, 1usize] && a.shape.dims.len() == 3);
+        if is_inner_transpose
+            && let Some(src_name) = node.inputs.first().and_then(|i| i.arg.as_tensor_name())
+            && let Some(out_ref) = node.outputs.first()
+            && let Some(out_t) = out_ref.as_tensor.as_ref()
+        {
+            self.transpose_2d_source
+                .insert(out_t.name.clone(), src_name.to_string());
+        }
         Ok(a.permute(axes))
     }
 
@@ -201,8 +240,17 @@ impl<'a> Translator<'a> {
 
         let dim = normalize_dim(dim, tensors[0].shape.len());
         let mut result = tensors[0];
+        let sym_ranges = sym_char_ranges(&self.sym_map);
         for t in &tensors[1..] {
-            result = result.concat_along(*t, dim);
+            let mut next = *t;
+            normalize_concat_dims(&mut result, &mut next, Some(dim), &sym_ranges);
+
+            let lhs_axis = result.dims()[dim];
+            let rhs_axis = next.dims()[dim];
+            let mut lhs_padded = result.pad_along(0, rhs_axis, dim, 0.);
+            let mut rhs_padded = next.pad_along(lhs_axis, 0, dim, 0.);
+            normalize_concat_dims(&mut lhs_padded, &mut rhs_padded, None, &sym_ranges);
+            result = lhs_padded + rhs_padded;
         }
         Ok(result)
     }
@@ -226,7 +274,231 @@ impl<'a> Translator<'a> {
         Ok(weight.gather(ids_expanded + arange_expanded))
     }
 
+    /// `aten.index_select(input, dim, index)` — pick rows/slices of `input`
+    /// along `dim` using a 1-D `index` tensor. Output shape is
+    /// `input.shape` with `dim` replaced by `index.shape[0]`.
+    ///
+    /// For the DLRM v3 use case this is `index_select(emb_weight, 0,
+    /// flat_indices)` — a 2-D source and 1-D index along dim 0. We lower
+    /// it the same way `translate_embedding` does: build a flat-rows
+    /// gather index `(index * hidden_dim) + arange(hidden_dim)` and read
+    /// the flattened weight in one pass. Higher-rank sources and non-zero
+    /// `dim` are not yet wired (would need stride math over Expression
+    /// shapes); they error out cleanly so they're easy to add when the
+    /// next model surfaces them.
+    pub(crate) fn translate_index_select(&mut self, node: &Node) -> Result<GraphTensor> {
+        let source = self.get_input_tensor(node, 0)?;
+        let dim_raw = self.get_int_arg(node, 1)?;
+        let index = self.get_input_tensor(node, 2)?;
+
+        let rank = source.shape.dims.len();
+        anyhow::ensure!(
+            rank == 2,
+            "translate_index_select: only 2-D source supported (got rank {rank}); \
+             extend this when a model needs higher rank."
+        );
+        let dim = if dim_raw < 0 {
+            dim_raw + rank as i64
+        } else {
+            dim_raw
+        };
+        anyhow::ensure!(
+            dim == 0,
+            "translate_index_select: only dim=0 supported (got {dim}); \
+             extend this when a model needs another axis."
+        );
+        anyhow::ensure!(
+            index.shape.dims.len() == 1,
+            "translate_index_select: index must be 1-D (got rank {})",
+            index.shape.dims.len()
+        );
+
+        // Same lowering as `translate_embedding`: build a flat gather index
+        // that combines the row-base offsets (`index * hidden_dim`) with a
+        // per-row `arange(hidden_dim)` broadcast.
+        let hidden_dim = source.shape.dims[1];
+        let n_idx = index.shape.dims[0];
+        let index_int = index.cast(DType::Int);
+        let base_expanded = (index_int * hidden_dim).expand_dim(1, hidden_dim);
+        let arange = self.graph.arange(hidden_dim);
+        let arange_expanded = arange.expand_dim(0, n_idx);
+        Ok(source.gather(base_expanded + arange_expanded))
+    }
+
+    /// `aten._embedding_bag_forward_only(weight, indices, offsets,
+    ///   scale_grad_by_freq, mode, sparse, per_sample_weights,
+    ///   include_last_offset, padding_idx)` →
+    ///   `(output, offset2bag, bag_size, max_indices)`.
+    ///
+    /// PyTorch decomposes `nn.EmbeddingBag` to this op. For the DLRM use
+    /// case all bags share a fixed stride `L = indices.len() / offsets.len()`
+    /// and `mode == 0` (sum). We detect that and lower to the fused
+    /// [`embedding_bag_sum_kernel`] on CUDA, or to a generic
+    /// `gather → reshape → sum` chain on CPU.
+    ///
+    /// Only `output` (tuple slot 0) is computed — `offset2bag`, `bag_size`
+    /// and `max_indices` are training-time dead ends for inference DLRM
+    /// and never read by any downstream `getitem`.
+    pub(crate) fn translate_embedding_bag_forward_only(&mut self, node: &Node) -> Result<()> {
+        let weight = self.get_input_tensor(node, 0)?;
+        let indices = self.get_input_tensor(node, 1)?;
+        let offsets = self.get_input_tensor(node, 2)?;
+        let mode = self.get_int_arg(node, 4).unwrap_or(0);
+        anyhow::ensure!(
+            mode == 0,
+            "translate_embedding_bag_forward_only: only mode=0 (sum) supported (got {mode}); \
+             vanilla DLRM uses sum-pooled bags. Extend this when a model needs mean/max."
+        );
+        // per_sample_weights is input index 6 and may be None / absent.
+        let has_per_sample_weights = node
+            .inputs
+            .get(6)
+            .and_then(|i| i.arg.as_tensor_name())
+            .is_some();
+        anyhow::ensure!(
+            !has_per_sample_weights,
+            "translate_embedding_bag_forward_only: per_sample_weights not supported \
+             (DLRM doesn't use them)."
+        );
+
+        anyhow::ensure!(
+            weight.shape.dims.len() == 2,
+            "translate_embedding_bag_forward_only: weight must be 2-D (got rank {})",
+            weight.shape.dims.len()
+        );
+        anyhow::ensure!(
+            indices.shape.dims.len() == 1,
+            "translate_embedding_bag_forward_only: indices must be 1-D (got rank {})",
+            indices.shape.dims.len()
+        );
+        anyhow::ensure!(
+            offsets.shape.dims.len() == 1,
+            "translate_embedding_bag_forward_only: offsets must be 1-D (got rank {})",
+            offsets.shape.dims.len()
+        );
+
+        let n_idx = indices.shape.dims[0]
+            .to_usize()
+            .context("translate_embedding_bag_forward_only: indices length must be static")?;
+        let batch = offsets.shape.dims[0]
+            .to_usize()
+            .context("translate_embedding_bag_forward_only: offsets length must be static")?;
+        anyhow::ensure!(
+            n_idx % batch == 0,
+            "translate_embedding_bag_forward_only: indices length ({n_idx}) must be a \
+             multiple of offsets length ({batch}); variable bag sizes not supported."
+        );
+        let bag = n_idx / batch;
+        let d = weight.shape.dims[1]
+            .to_usize()
+            .context("translate_embedding_bag_forward_only: weight dim 1 must be static")?;
+
+        // Reshape indices (B*L,) → (B, L) and cast to i32 (luminal kernel
+        // wants Int). Then either use the fused kernel under CUDA or
+        // a host-portable gather+sum lowering.
+        let indices_int = indices.cast(DType::Int);
+        let indices_2d = {
+            let new_shape = ShapeTracker::new(vec![
+                Expression::from(batch),
+                Expression::from(bag),
+            ]);
+            GraphTensor {
+                id: indices_int.id,
+                graph_ref: indices_int.graph_ref,
+                shape: new_shape,
+                dtype: indices_int.dtype,
+            }
+        };
+
+        let backend_is_cuda = true; // LUMINAL_BACKEND_CUDA replaced by cfg!(feature="cuda") gate above.
+
+        let result = if cfg!(feature = "cuda") && backend_is_cuda && weight.dtype == DType::F32 {
+            // Fused CUDA path: one kernel for the whole bag-sum.
+            #[cfg(feature = "cuda")]
+            {
+                luminal_cuda_lite::kernel::embedding_bag_sum_kernel(weight, indices_2d)
+            }
+            #[cfg(not(feature = "cuda"))]
+            {
+                // Unreachable — gated above, but keep the compiler happy.
+                unreachable!()
+            }
+        } else {
+            // Generic fallback: gather (B*L, D) then reshape and sum.
+            let hidden_dim = weight.shape.dims[1];
+            let ids_expanded = (indices_2d * hidden_dim).expand_dim(2, hidden_dim);
+            let arange = self.graph.arange(hidden_dim);
+            let arange_expanded = arange.expand_dim(0, batch).expand_dim(0, bag);
+            // Note: weight.gather expects the gather indices to broadcast
+            // against weight's row-flattened layout; we want (B, L, D)
+            // out, then sum along L.
+            let _ = d; // hidden_dim is used; keep `d` reachable for debug only.
+            let gathered = weight.gather(ids_expanded + arange_expanded);
+            gathered.sum(1)
+        };
+
+        // Record the output under outputs[0][0] (the tuple's first slot).
+        // The other three slots are dead under inference and there's no
+        // downstream `getitem` that reads them — but if there ever is,
+        // we'd need to materialize them too.
+        let out_name = node
+            .outputs
+            .first()
+            .and_then(|o| o.as_tensors.as_ref())
+            .and_then(|ts| ts.first().map(|t| t.name.clone()))
+            .or_else(|| {
+                node.outputs
+                    .first()
+                    .and_then(|o| o.as_tensor.as_ref().map(|t| t.name.clone()))
+            })
+            .context(
+                "translate_embedding_bag_forward_only: missing output[0] name in FX node",
+            )?;
+        self.tensors.insert(out_name.clone(), result);
+        // Record node_chain / op_inputs for the *primary* output (tuple
+        // slot 0). Multi-output ops normally skip the bookkeeping at the
+        // bottom of `translate_node` (they return early), but later
+        // peepholes — specifically the stacked-emb-bag fusion inside the
+        // pairwise-dot peephole — need to identify which tensor sources
+        // came from an embedding_bag, so we record explicitly.
+        let first_input_name: Option<String> = node
+            .inputs
+            .first()
+            .and_then(|i| i.arg.as_tensor_name().map(|s| s.to_string()));
+        if let Some(first_input) = first_input_name {
+            self.node_chain
+                .insert(out_name.clone(), (node.target.clone(), first_input));
+        }
+        let mut all_inputs: Vec<String> = Vec::new();
+        for inp in &node.inputs {
+            if let Some(names) = inp.arg.as_tensors() {
+                for tn in names {
+                    all_inputs.push(tn.name.clone());
+                }
+            } else if let Some(name) = inp.arg.as_tensor_name() {
+                all_inputs.push(name.to_string());
+            }
+        }
+        if !all_inputs.is_empty() {
+            self.op_inputs.insert(out_name, all_inputs);
+        }
+        Ok(())
+    }
+
     pub(crate) fn translate_index_tensor(&mut self, node: &Node) -> Result<GraphTensor> {
+        // Try the DLRM PairwiseDot peephole before falling back to the
+        // generic gather-based lowering. Detects the
+        //   Z[:, li, lj] ← bmm(T, T.transpose(1, 2)) ← cat([t1.unsqueeze(1), ..., tF.unsqueeze(1)], dim=1)
+        // pattern that vanilla `nn.Sequential` DLRM (DLRMv1) emits — the
+        // version a user writes with one `EmbeddingBag` per categorical
+        // table. Replacing it with `dlrm_pairwise_dot_lower_tri` collapses
+        // the cat-then-bmm-then-gather chain (which lowers to ~40
+        // small Iota/Cast/Gather/FusedRegion kernels via pad_along+add)
+        // into a single CUDA kernel.
+        if let Some(t) = self.try_translate_pairwise_dot_lower_tri(node)? {
+            return Ok(t);
+        }
+
         let source = self.get_input_tensor(node, 0)?;
 
         // Handle indices as_tensors (all non-None) or as individual args with None entries
@@ -278,6 +550,98 @@ impl<'a> Translator<'a> {
                     expanded.shape.expand(target);
                     return Ok(source.gather_elements(expanded, first_non_none_dim));
                 }
+
+                // Multi-index advanced indexing through leading dims that
+                // pass through (e.g. DLRM's `Z[:, li, lj]` where Z has
+                // shape `(B, F, F)` and output[b, p] = Z[b, li[p], lj[p]]).
+                //
+                // Strategy: reduce to the proven single-index simple
+                // case. Combine the multi-axis indices into one (`li * F
+                // + lj`) and reshape the source so the indexed region
+                // becomes a single dim. Then take the exact `gather_elements`
+                // path the rest of this translator already uses.
+                //
+                // Supported shape pattern (DLRM): exactly one leading
+                // passthrough dim, no trailing dims after the indexed
+                // region, per-axis indices all 1-D of the same length.
+                if first_non_none_dim > 0 {
+                    let src_dims = source.shape.dims;
+                    let src_rank = src_dims.len();
+                    let n_idx = index_names.len();
+                    let trailing_start = first_non_none_dim + n_idx;
+                    anyhow::ensure!(
+                        first_non_none_dim == 1,
+                        "index.Tensor: leading-dim passthrough only supported for \
+                         exactly one leading dim (got {first_non_none_dim})."
+                    );
+                    anyhow::ensure!(
+                        trailing_start == src_rank,
+                        "index.Tensor: trailing dims after indexed region not yet supported."
+                    );
+                    let mut idx_tensors: Vec<GraphTensor> = Vec::with_capacity(n_idx);
+                    for n in &index_names {
+                        idx_tensors.push(self.get_tensor(&n.name)?.cast(DType::Int));
+                    }
+                    let idx0_shape = idx_tensors[0].shape.dims;
+                    anyhow::ensure!(
+                        idx0_shape.len() == 1,
+                        "index.Tensor: only 1-D per-axis indices supported (got rank {})",
+                        idx0_shape.len()
+                    );
+                    for it in idx_tensors.iter().skip(1) {
+                        anyhow::ensure!(
+                            it.shape.dims == idx0_shape,
+                            "index.Tensor: per-axis indices must share a common shape"
+                        );
+                    }
+                    // strides over indexed axes (no trailing dims).
+                    let mut strides_idx: Vec<Expression> = vec![Expression::from(1usize); n_idx];
+                    for i in (0..n_idx - 1).rev() {
+                        strides_idx[i] =
+                            strides_idx[i + 1] * src_dims[first_non_none_dim + i + 1];
+                    }
+                    // combined[p] = sum_i idx_i * stride_i  (1-D)
+                    let mut combined: Option<GraphTensor> = None;
+                    for (i, it) in idx_tensors.into_iter().enumerate() {
+                        let weighted = if strides_idx[i].to_usize() == Some(1) {
+                            it
+                        } else {
+                            it * strides_idx[i]
+                        };
+                        combined = Some(match combined {
+                            Some(acc) => {
+                                let (a, b) = broadcast_binary(acc, weighted);
+                                a + b
+                            }
+                            None => weighted,
+                        });
+                    }
+                    let combined = combined.context("index.Tensor: no indices")?;
+
+                    // Indexed region size, then a (leading, indexed_size) reshape.
+                    let mut indexed_size = Expression::from(1usize);
+                    for d in &src_dims[first_non_none_dim..trailing_start] {
+                        indexed_size *= *d;
+                    }
+                    let leading_dim = src_dims[0];
+                    let flat_source =
+                        reshape_tensor(source, vec![leading_dim, indexed_size]);
+
+                    // Now dispatch through the exact single-index simple
+                    // case lowering — known-good. Add unit leading dims
+                    // to match flat_source rank, then expand to the full
+                    // (leading_dim, pair_count) shape.
+                    let mut expanded = combined;
+                    let flat_rank = 2; // (leading, indexed_size)
+                    for _ in 0..(flat_rank - expanded.shape.len()) {
+                        expanded = expanded.expand_dim(0, Expression::from(1usize));
+                    }
+                    let idx_dim_size = expanded.shape.dims[1];
+                    let mut target: Vec<Expression> = vec![leading_dim, indexed_size];
+                    target[1] = idx_dim_size;
+                    expanded.shape.expand(target);
+                    return Ok(flat_source.gather_elements(expanded, 1));
+                }
             } else {
                 bail!(
                     "index.Tensor: unsupported indices format: {:?}",
@@ -522,4 +886,394 @@ impl<'a> Translator<'a> {
 
         first_chunk.ok_or_else(|| anyhow::anyhow!("split_with_sizes: empty sizes list"))
     }
+
+    /// DLRM PairwiseDot peephole: detect
+    ///   `aten.index.Tensor(bmm, [None, li, lj])`
+    /// where
+    ///   `bmm = aten.bmm.default(T, T_permuted)`
+    ///   `T_permuted = aten.permute.default(T, [0, 2, 1])`
+    ///   `T = aten.cat.default([unsqueeze_a, unsqueeze_b, …], dim=1)`
+    ///   each `unsqueeze_k = aten.unsqueeze.default(t_k, 1)`
+    /// and lower to `dlrm_pairwise_dot_lower_tri([t_0, t_1, …])`.
+    ///
+    /// Why this matters: at DLRM nc=3 the generic lowering produces
+    /// ~80 CUDA-graph kernels from the cat+bmm+gather chain alone (the
+    /// `pad_along + add` decomposition of cat fans out into many
+    /// Iota/Cast/Gather/FusedRegion launches). The fused kernel
+    /// computes the F(F-1)/2 dot products directly with one launch.
+    ///
+    /// Returns `Ok(Some(out))` on match, `Ok(None)` if the pattern
+    /// doesn't apply, `Err(_)` only if matching diagnostics surface a
+    /// genuine bug.
+    fn try_translate_pairwise_dot_lower_tri(
+        &mut self,
+        node: &Node,
+    ) -> Result<Option<GraphTensor>> {
+        // CUDA-only fast path. The kernel is in luminal_cuda_lite.
+        #[cfg(not(feature = "cuda"))]
+        {
+            let _ = node;
+            return Ok(None);
+        }
+
+        #[cfg(feature = "cuda")]
+        {
+            let backend_is_cuda = true; // LUMINAL_BACKEND_CUDA replaced by cfg!(feature="cuda") gate above.
+            if !backend_is_cuda {
+                return Ok(None);
+            }
+
+            // 1. Detect [None, li, lj] index list.
+            let opt_tensors =
+                match node.inputs.get(1).and_then(|i| i.arg.as_optional_tensors()) {
+                    Some(t) => t,
+                    None => {
+                        return Ok(None);
+                    }
+                };
+            // Expect [None, li, lj]: three entries, first None, last two
+            // are tensors.
+            if opt_tensors.len() != 3 {
+                return Ok(None);
+            }
+            use crate::pt2_schema::OptionalTensorEntry;
+            let (li_name, lj_name) =
+                match (&opt_tensors[0], &opt_tensors[1], &opt_tensors[2]) {
+                    (OptionalTensorEntry::None(_), OptionalTensorEntry::Tensor(li), OptionalTensorEntry::Tensor(lj)) => {
+                        (li.as_tensor.name.clone(), lj.as_tensor.name.clone())
+                    }
+                    _ => {
+                        return Ok(None);
+                    }
+                };
+            // 2. Source must be a bmm of (T, T_permuted).
+            let source_name = match node.inputs.first().and_then(|i| i.arg.as_tensor_name()) {
+                Some(s) => s.to_string(),
+                None => {
+                    return Ok(None);
+                }
+            };
+            let bmm_info = match self.node_chain.get(&source_name) {
+                Some(x) => x.clone(),
+                None => {
+                    return Ok(None);
+                }
+            };
+            if bmm_info.0 != "torch.ops.aten.bmm.default" {
+                return Ok(None);
+            }
+            let bmm_inputs = match self.op_inputs.get(&source_name) {
+                Some(v) => v.clone(),
+                None => {
+                    return Ok(None);
+                }
+            };
+            if bmm_inputs.len() != 2 {
+                return Ok(None);
+            }
+            let (bmm_a, bmm_b) = (bmm_inputs[0].clone(), bmm_inputs[1].clone());
+            // 3. Both bmm inputs must descend from the same cat — one
+            //    directly, one via a [0, 2, 1] permute. The permute is
+            //    already recorded in `transpose_2d_source` (we extended
+            //    it to cover 3-D `[0, 2, 1]` for bmm fast paths).
+            let permute_src = self.transpose_2d_source.get(&bmm_b).cloned();
+            let (cat_name, _has_transpose) = if permute_src.as_deref() == Some(bmm_a.as_str()) {
+                (bmm_a.clone(), true)
+            } else if self
+                .transpose_2d_source
+                .get(&bmm_a)
+                .map(|s| s.as_str() == bmm_b.as_str())
+                .unwrap_or(false)
+            {
+                (bmm_b.clone(), true)
+            } else {
+                return Ok(None);
+            };
+            let cat_info = match self.node_chain.get(&cat_name) {
+                Some(x) => x.clone(),
+                None => {
+                    return Ok(None);
+                }
+            };
+            if cat_info.0 != "torch.ops.aten.cat.default" {
+                return Ok(None);
+            }
+            let cat_inputs = match self.op_inputs.get(&cat_name) {
+                Some(v) => v.clone(),
+                None => {
+                    return Ok(None);
+                }
+            };
+            if cat_inputs.len() < 2 {
+                return Ok(None);
+            }
+            // 4. Each cat input should be `unsqueeze(t_k, 1)` — peel the
+            //    unsqueeze so we recover the original (B, D) tensor.
+            //    Also track the FX *source name* of each feature so the
+            //    multi-call fusion below can recognize emb-bag-rooted
+            //    features and emit the stacked path.
+            let mut feature_tensors: Vec<GraphTensor> = Vec::with_capacity(cat_inputs.len());
+            let mut feature_source_names: Vec<String> = Vec::with_capacity(cat_inputs.len());
+            for ci in &cat_inputs {
+                let unsqueeze_info = match self.node_chain.get(ci) {
+                    Some(x) => x.clone(),
+                    None => {
+                        return Ok(None);
+                    }
+                };
+                if unsqueeze_info.0 != "torch.ops.aten.unsqueeze.default" {
+                    return Ok(None);
+                }
+                // unsqueeze's first input is the source tensor name.
+                let src = unsqueeze_info.1;
+                let t = self.get_tensor(&src)?;
+                if t.dtype != DType::F32 || t.shape.dims.len() != 2 {
+                    return Ok(None);
+                }
+                feature_tensors.push(t);
+                feature_source_names.push(src);
+            }
+            // 5. Sanity check on li/lj — they must be the strict lower-tri
+            //    pair table for F = feature_tensors.len(). We don't
+            //    materialize them; just verify the index buffers have
+            //    the right length, then trust they're tril-indices.
+            //    A user passing arbitrary indices through this exact
+            //    chain would silently get tril results; gating on
+            //    `index buffer length == F*(F-1)/2` catches the common
+            //    case without invasive constant-folding work.
+            let f = feature_tensors.len();
+            let pair_count = f * (f - 1) / 2;
+            let li_t = self.get_tensor(&li_name)?;
+            let lj_t = self.get_tensor(&lj_name)?;
+            if li_t.shape.dims.len() != 1
+                || lj_t.shape.dims.len() != 1
+                || li_t.shape.dims[0].to_usize() != Some(pair_count)
+                || lj_t.shape.dims[0].to_usize() != Some(pair_count)
+            {
+                return Ok(None);
+            }
+
+            // 6. Multi-call fusion: detect when exactly one feature is the
+            //    dense MLP output and all others came from
+            //    `_embedding_bag.default` calls with matching `(rows, D)`.
+            //    Then collapse the N separate per-table embedding-bag
+            //    kernels into one fused `stacked_embedding_bag_sum_kernel`
+            //    and use the stacked-variant pairwise dot. Mirrors the
+            //    hand-written DLRM rust example's kernel shape.
+            if let Some((dense, emb_stack)) =
+                self.try_fuse_stacked_emb_bag(&feature_tensors, &feature_source_names)?
+            {
+                let out = luminal_cuda_lite::kernel::dlrm_pairwise_dot_lower_tri_stacked(
+                    dense, emb_stack,
+                );
+                return Ok(Some(out));
+            }
+
+            // Fallback: per-feature variadic kernel. The bmm, cat,
+            // and unsqueeze nodes left dangling in the HLIR get picked
+            // up by `Translator::dce()` after every FX node is translated
+            // (walks back from `Output` HLIR sinks and drops everything
+            // unreachable). luminal's egglog optimizer leaves some of
+            // these subgraphs alive on its own, so the explicit pass is
+            // load-bearing for this peephole.
+            let out = luminal_cuda_lite::kernel::dlrm_pairwise_dot_lower_tri(feature_tensors);
+            Ok(Some(out))
+        }
+    }
+
+    /// Multi-call fusion: scan a list of feature sources for the DLRM
+    /// pattern "1 dense MLP output + N `_embedding_bag.default` outputs"
+    /// and, when matched, emit a single `stacked_embedding_bag_sum_kernel`
+    /// over a concatenated weight tensor. Returns `(dense, emb_stack)`
+    /// where `dense` is the (B, D) MLP output and `emb_stack` is
+    /// `(B, num_emb, D)` for the stacked-pairwise-dot kernel to consume.
+    ///
+    /// Requirements:
+    ///   * All embedding-bag weights are 2-D F32 with the same `(rows, D)`
+    ///   * All bag indices have the same `(batch, bag_size)` shape
+    ///   * The dense feature is the *first* cat input (DLRMv1 emits
+    ///     `cat([x.unsqueeze(1)] + [e.unsqueeze(1) for e in ly], dim=1)`)
+    #[cfg(feature = "cuda")]
+    fn try_fuse_stacked_emb_bag(
+        &mut self,
+        feature_tensors: &[GraphTensor],
+        feature_source_names: &[String],
+    ) -> Result<Option<(GraphTensor, GraphTensor)>> {
+        use crate::pt2_schema::OptionalTensorEntry;
+        if feature_tensors.len() < 2 {
+            return Ok(None);
+        }
+        // Look up each feature source's producing FX node target. The
+        // dense MLP output has target ≠ `_embedding_bag*`. We require the
+        // *first* feature to be dense (DLRMv1 convention) and the rest to
+        // be emb-bag.
+        let producer_target = |name: &str| -> Option<String> {
+            self.node_chain.get(name).map(|(target, _)| target.clone())
+        };
+        let dense_target = producer_target(&feature_source_names[0]);
+        if let Some(t) = &dense_target
+            && (t == "torch.ops.aten._embedding_bag.default"
+                || t == "torch.ops.aten._embedding_bag_forward_only.default")
+        {
+            return Ok(None);
+        }
+        for src in feature_source_names.iter().skip(1) {
+            match producer_target(src).as_deref() {
+                Some("torch.ops.aten._embedding_bag.default")
+                | Some("torch.ops.aten._embedding_bag_forward_only.default") => {}
+                _ => return Ok(None),
+            }
+        }
+        // For each emb-bag feature source, pull (weight, indices, offsets)
+        // by walking back through the parsed FX node array. The FX nodes
+        // for emb-bag store the output tensor name(s) in `outputs[0].as_tensors`.
+        let mut emb_weights: Vec<GraphTensor> = Vec::new();
+        let mut emb_indices: Vec<GraphTensor> = Vec::new();
+        let mut emb_rows: Option<usize> = None;
+        let mut emb_d: Option<usize> = None;
+        let mut emb_batch: Option<usize> = None;
+        let mut emb_bag: Option<usize> = None;
+        for src in feature_source_names.iter().skip(1) {
+            // Find the FX emb-bag node whose primary output is `src`.
+            // emb-bag is a multi-output op: `outputs[0].as_tensors[0].name`
+            // is what downstream `getitem(node, 0)` references, which is
+            // what our translator stores in `tensors`. Some exports drop
+            // the getitem entirely and just inline the name as a single
+            // tensor output, so check both shapes.
+            let node_opt = self
+                .parsed
+                .program
+                .graph_module
+                .graph
+                .nodes
+                .iter()
+                .find(|n| {
+                    if n.target != "torch.ops.aten._embedding_bag.default"
+                        && n.target != "torch.ops.aten._embedding_bag_forward_only.default"
+                    {
+                        return false;
+                    }
+                    let Some(out) = n.outputs.first() else {
+                        return false;
+                    };
+                    if let Some(ts) = out.as_tensors.as_ref() {
+                        return ts.iter().any(|t| t.name == *src);
+                    }
+                    if let Some(t) = out.as_tensor.as_ref() {
+                        return t.name == *src;
+                    }
+                    false
+                });
+            let Some(node) = node_opt else {
+                return Ok(None);
+            };
+            // _embedding_bag(weight, indices, offsets, ...)
+            let weight_name = node
+                .inputs
+                .first()
+                .and_then(|i| i.arg.as_tensor_name())
+                .map(|s| s.to_string());
+            let indices_name = node
+                .inputs
+                .get(1)
+                .and_then(|i| i.arg.as_tensor_name())
+                .map(|s| s.to_string());
+            let offsets_name = node
+                .inputs
+                .get(2)
+                .and_then(|i| i.arg.as_tensor_name())
+                .map(|s| s.to_string());
+            let mode = self.get_int_arg(node, 4).unwrap_or(0);
+            if mode != 0 {
+                return Ok(None);
+            }
+            // per_sample_weights at index 6
+            let has_psw = node
+                .inputs
+                .get(6)
+                .and_then(|i| i.arg.as_tensor_name())
+                .is_some();
+            if has_psw {
+                return Ok(None);
+            }
+            let (Some(wn), Some(in_), Some(on)) = (weight_name, indices_name, offsets_name) else {
+                return Ok(None);
+            };
+            let w = self.get_tensor(&wn)?;
+            let i = self.get_tensor(&in_)?;
+            let o = self.get_tensor(&on)?;
+            if w.dtype != DType::F32 || w.shape.dims.len() != 2 {
+                return Ok(None);
+            }
+            if i.shape.dims.len() != 1 || o.shape.dims.len() != 1 {
+                return Ok(None);
+            }
+            let rows = w.shape.dims[0].to_usize().ok_or_else(|| {
+                anyhow::anyhow!("stacked emb-bag fusion: rows must be static")
+            })?;
+            let d = w.shape.dims[1].to_usize().ok_or_else(|| {
+                anyhow::anyhow!("stacked emb-bag fusion: D must be static")
+            })?;
+            let n_idx = i.shape.dims[0].to_usize().ok_or_else(|| {
+                anyhow::anyhow!("stacked emb-bag fusion: indices length must be static")
+            })?;
+            let batch = o.shape.dims[0].to_usize().ok_or_else(|| {
+                anyhow::anyhow!("stacked emb-bag fusion: offsets length must be static")
+            })?;
+            if n_idx % batch != 0 {
+                return Ok(None);
+            }
+            let bag = n_idx / batch;
+            // All tables must agree on (rows, d, batch, bag).
+            if emb_rows.get_or_insert(rows) != &rows {
+                return Ok(None);
+            }
+            if emb_d.get_or_insert(d) != &d {
+                return Ok(None);
+            }
+            if emb_batch.get_or_insert(batch) != &batch {
+                return Ok(None);
+            }
+            if emb_bag.get_or_insert(bag) != &bag {
+                return Ok(None);
+            }
+            // Reshape indices to (B, bag) of int32 for the kernel.
+            let i_int = i.cast(DType::Int);
+            let new_shape = ShapeTracker::new(vec![
+                Expression::from(batch),
+                Expression::from(bag),
+            ]);
+            let indices_2d = GraphTensor {
+                id: i_int.id,
+                graph_ref: i_int.graph_ref,
+                shape: new_shape,
+                dtype: i_int.dtype,
+            };
+            emb_weights.push(w);
+            emb_indices.push(indices_2d);
+            // Silence unused warning when no tables match the size.
+            let _ = OptionalTensorEntry::None;
+        }
+
+        let _ = emb_rows; // (already validated via per-table equality)
+        // Use the multi-table kernel: takes N (weight, idx) pairs and
+        // produces (B, num_emb, D) in one launch. Crucially this avoids
+        // the `concat_along`-of-persistent-weights expansion (which would
+        // emit pad+add HLIR kernels per pair) — the kernel reads each
+        // table's weight pointer directly via a packed staging buffer.
+        let emb_stack = luminal_cuda_lite::kernel::multi_table_embedding_bag_sum_kernel(
+            emb_weights,
+            emb_indices,
+        );
+        Ok(Some((feature_tensors[0], emb_stack)))
+    }
+
+    #[cfg(not(feature = "cuda"))]
+    fn try_fuse_stacked_emb_bag(
+        &mut self,
+        _feature_tensors: &[GraphTensor],
+        _feature_source_names: &[String],
+    ) -> Result<Option<(GraphTensor, GraphTensor)>> {
+        Ok(None)
+    }
 }

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

@@ -220,6 +220,7 @@ impl<'a> Translator<'a> {
         let input = self.get_input_tensor(node, 0)?;
         let weight = self.get_input_tensor(node, 1)?;
         let offs = self.get_input_tensor(node, 2)?;
+        let out_dtype = self.output_meta_dtype(node)?;
 
         anyhow::ensure!(
             input.shape.len() == 2,
@@ -274,8 +275,15 @@ impl<'a> Translator<'a> {
         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);
+        // Gather → [S, K, N], then normalize both operands to the op's declared
+        // output dtype before matmul. On real Qwen3-MoE bf16 checkpoints the FX
+        // graph inserts casts on the activation path, and relying on the input
+        // tensor's translated dtype can leave us with mixed F32/Bf16 operands
+        // by the time matmul expands into elementwise Mul. Using the PT2 output
+        // metadata keeps the matmul dtype aligned with the exported contract
+        // without upcasting the full expert weight bank.
+        let weight_gathered = weight.gather(flat_idx).cast(out_dtype);
+        let input = input.cast(out_dtype);
 
         // Per-token matmul: [S, 1, K] @ [S, K, N] → [S, 1, N] → [S, N].
         // Operands stay in their native dtype — no F32 cast on the gathered
@@ -287,7 +295,7 @@ impl<'a> Translator<'a> {
         // (cuBLASLt etc.) handle bf16 input with F32 accumulator internally.
         let result = input.unsqueeze(1).matmul(weight_gathered).squeeze(1);
 
-        Ok(result.cast(input.dtype))
+        Ok(result.cast(out_dtype))
     }
 
     /// Build the where-formula graph: `cond * x + (1 - cond) * y`, computed

crates/luminal_python/src/luminal/compiled_model.py

@@ -43,6 +43,63 @@ class CompiledModel:
             else torch.float32
             for i in range(len(self._input_names))
         ]
+        # Pre-resolve opaque integer ids for inputs and outputs so the
+        # per-iter `run_with_ptrs` path can skip the string-keyed lookup
+        # inside Rust. One pyo3 call per name at compile time, no per-iter
+        # name handling thereafter. Falls back to None when the Rust side
+        # is an older build that doesn't expose `tensor_id` /
+        # `run_with_ptrs`.
+        self._batched_ptrs_supported = (
+            self._supports_device_ptrs
+            and hasattr(graph_result, "tensor_id")
+            and hasattr(graph_result, "run_with_ptrs")
+        )
+        if self._batched_ptrs_supported:
+            self._input_ids = [graph_result.tensor_id(n) for n in self._input_names]
+            self._output_ids = [graph_result.tensor_id(n) for n in self._output_names]
+        else:
+            self._input_ids = None
+            self._output_ids = None
+        # Output dtype codes — cached at __init__ instead of re-fetched in
+        # each call (was a per-iter pyo3 attribute access on the dynamic
+        # path; for static-shape models the codes don't change).
+        self._output_dtype_codes_cached = (
+            None if self._has_dynamic_dims else list(graph_result.output_dtypes)
+        )
+        # Per-input "have-we-seen-this-ptr-before" cache. In a hot bench
+        # loop the same user tensor objects are passed each iter, so most
+        # of the per-iter Python work (detach + contiguous + dtype cast +
+        # data_ptr + numel + element_size) repeats with identical inputs.
+        # Cache (id(orig_tensor), orig_data_ptr, cast_tensor, cast_ptr,
+        # cast_n_bytes) per input slot; on hit, skip everything and rely
+        # on luminal's previously-registered pointer (the runtime keeps
+        # CudaInput::Ptr entries across `execute()` calls thanks to the
+        # consume-step filter in `cuda_lite/runtime.rs`). The cast tensor
+        # reference is held inside the cache so PyTorch's caching
+        # allocator can't recycle the converted buffer.
+        #
+        # Sharp edge: callers that mutate a user tensor in place via
+        # `.copy_(...)` keep the same `id()` and `data_ptr()` and will
+        # silently get stale cached data. Fresh-tensor callers
+        # (`make_inputs(...)` each iter, or new outputs from upstream)
+        # cache-miss naturally and pay the full cold-path cost. If a
+        # future model needs in-place input mutation, swap this check
+        # for one that also looks at `tensor._version` (autograd's
+        # mutation counter) — but PyTorch flags `_version` as private,
+        # so don't reach for it unless an actual model needs it.
+        self._input_cache_ids = [None] * len(self._input_names)
+        self._input_cache_orig_ptrs = [0] * len(self._input_names)
+        self._input_cache_cast_tensors = [None] * len(self._input_names)
+        self._input_cache_specs = [None] * len(self._input_names)
+        # Cached output tensors mirror the input-side cache. For static-
+        # shape models the output tensors can be reused across calls if
+        # the input device is unchanged — saves ~3 μs/output of
+        # `torch.empty` + the FFI to register the device pointer.
+        # NB: callers that stash the returned tensor must `.clone()`
+        # before the next call; the default contract returns a fresh
+        # tensor each call so leave this opt-in via env var for now.
+        self._output_cache_tensors = None
+        self._output_cache_specs = None
 
     def set_dim(self, param_name: str, value: int) -> None:
         """Set a dynamic dimension value by its param name."""
@@ -87,6 +144,90 @@ class CompiledModel:
             input_shapes = [list(t.shape) for t in user_inputs]
             self._graph.auto_set_dims_from_input_shapes(input_shapes)
 
+        # Batched CUDA fast path. When all user inputs are on CUDA and all
+        # outputs are floating-point, collapse the per-iter `set_input_device_ptr`
+        # × N, `set_output_device_ptr` × M, and `run` calls into a single
+        # `run_with_ptrs` FFI crossing. The Rust side iterates the
+        # (id, ptr, n_bytes) tuples without paying per-call pyo3
+        # marshalling cost.
+        all_cuda = bool(user_inputs) and all(t.is_cuda for t in user_inputs)
+        if self._batched_ptrs_supported and all_cuda:
+            output_shapes = (
+                self._graph.resolve_output_shapes()
+                if self._has_dynamic_dims
+                else self._output_shapes
+            )
+            output_dtype_codes = (
+                self._graph.output_dtypes
+                if self._output_dtype_codes_cached is None
+                else self._output_dtype_codes_cached
+            )
+            output_dtypes = [
+                code_to_torch_dtype(output_dtype_codes[i])
+                if i < len(output_dtype_codes)
+                else torch.float32
+                for i in range(len(self._output_names))
+            ]
+            if all(d.is_floating_point for d in output_dtypes):
+                # Build the input-spec list. Hot path: when the user passes
+                # the same tensor object as last call AND that tensor's
+                # data_ptr is unchanged, the previous registration is still
+                # valid in the luminal runtime — skip both the Python-side
+                # cast/contiguous/data_ptr work AND the Rust-side
+                # `set_device_ptr` call by omitting it from `input_specs`.
+                input_specs = []
+                _cache_ids = self._input_cache_ids
+                _cache_orig = self._input_cache_orig_ptrs
+                _cache_cast = self._input_cache_cast_tensors
+                _cache_spec = self._input_cache_specs
+                for i, (id_, tensor, expected_dtype) in enumerate(zip(
+                    self._input_ids, user_inputs, self._input_dtypes
+                )):
+                    orig_id = id(tensor)
+                    orig_ptr = tensor.data_ptr()
+                    if _cache_ids[i] == orig_id and _cache_orig[i] == orig_ptr:
+                        # Pointer unchanged — luminal already has the
+                        # registration. Skip everything.
+                        continue
+                    # Cold path: cast (if needed) and update cache.
+                    t = tensor.detach().contiguous().to(expected_dtype)
+                    n_bytes = t.numel() * t.element_size()
+                    cast_ptr = t.data_ptr()
+                    spec = (id_, cast_ptr, n_bytes)
+                    input_specs.append(spec)
+                    _cache_ids[i] = orig_id
+                    _cache_orig[i] = orig_ptr
+                    _cache_cast[i] = t  # keep alive
+                    _cache_spec[i] = spec
+                # Outputs: pre-allocate fresh each call so the caller gets
+                # a unique tensor (matches the unbatched path's contract;
+                # callers that stash results don't get blown away on the
+                # next iteration). The output_specs list is always passed
+                # through to `run_with_ptrs`.
+                output_tensors = []
+                output_specs = []
+                for id_, shape, dt in zip(self._output_ids, output_shapes, output_dtypes):
+                    out = torch.empty(shape, dtype=dt, device=input_device)
+                    output_specs.append((id_, out.data_ptr(), out.numel() * out.element_size()))
+                    output_tensors.append(out)
+                zero_copy_flags = self._graph.run_with_ptrs(input_specs, output_specs)
+                # For any output the runtime had to alias (not zero-copy),
+                # request the DtoD copy explicitly into the registered buffer.
+                # In DLRMv1 and similar models this never fires, but it's the
+                # same fallback the unbatched path has.
+                for ok, name, tensor in zip(zero_copy_flags, self._output_names, output_tensors):
+                    if not ok:
+                        self._graph.copy_output_to_device_ptr(
+                            name, tensor.data_ptr(), tensor.numel() * tensor.element_size()
+                        )
+                # Return tuple unconditionally — the unbatched path below
+                # also returns `tuple(outputs)` even for single-output
+                # models, and torch.compile / dynamo's output handling
+                # depends on that contract. Returning a bare Tensor here
+                # made dynamo iterate the first dim and reshape the
+                # output to a 1-element slice.
+                return tuple(output_tensors)
+
         # 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

crates/luminal_python/tests/test_dynamic_shapes.py

@@ -9,6 +9,9 @@ PT2 export, and reuse a single compiled graph across shape changes.
 
 from __future__ import annotations
 
+import copy
+from contextlib import contextmanager
+
 import pytest
 import torch
 import torch._dynamo
@@ -34,6 +37,64 @@ def _compile_with_dynamic_true(model, count_holder):
     return torch.compile(model, backend=wrapper, dynamic=True)
 
 
+def _compile_with_capture(model, count_holder, capture_holder):
+    def wrapper(gm, example_inputs):
+        out = luminal_backend(gm, example_inputs)
+        count_holder.append(1)
+        if "gm" not in capture_holder:
+            capture_holder["gm"] = copy.deepcopy(gm).eval()
+            capture_holder["example_inputs"] = example_inputs
+            capture_holder["compiled_impl"] = out
+        return out
+
+    return torch.compile(model, backend=wrapper)
+
+
+@contextmanager
+def _explicit_mark_dynamic_mode():
+    prev_auto = torch._dynamo.config.automatic_dynamic_shapes
+    prev_cache_limit = torch._dynamo.config.cache_size_limit
+    torch._dynamo.reset()
+    torch._dynamo.config.automatic_dynamic_shapes = False
+    torch._dynamo.config.cache_size_limit = 8
+    try:
+        yield
+    finally:
+        torch._dynamo.config.automatic_dynamic_shapes = prev_auto
+        torch._dynamo.config.cache_size_limit = prev_cache_limit
+        torch._dynamo.reset()
+
+
+def _first_trace_dynamic_shapes(capture_holder):
+    from luminal.pt2 import (
+        _build_dynamic_shapes_from_gm,
+        _reinternalize_lifted_params,
+        _strip_symint_placeholders,
+    )
+
+    gm = copy.deepcopy(capture_holder["gm"]).eval()
+    example_inputs = capture_holder["example_inputs"]
+    gm, user_inputs, _, _ = _reinternalize_lifted_params(gm, example_inputs)
+    user_inputs, _, strip_ok = _strip_symint_placeholders(gm, user_inputs)
+    dynamic_shapes = _build_dynamic_shapes_from_gm(gm) if strip_ok else None
+    return strip_ok, dynamic_shapes
+
+
+def _assert_input_dynamic_dims(dynamic_shapes, input_index, expected_dims):
+    args_spec = dynamic_shapes.get("args")
+    assert args_spec is not None and len(args_spec) > input_index, (
+        f"expected dynamic spec for input {input_index}, got {dynamic_shapes}"
+    )
+    spec = args_spec[input_index]
+    assert spec is not None, (
+        f"expected a per-dim dynamic spec for input {input_index}, got {dynamic_shapes}"
+    )
+    assert set(spec.keys()) == set(expected_dims), (
+        f"expected dynamic dims {set(expected_dims)} for input {input_index}, "
+        f"got {dynamic_shapes}"
+    )
+
+
 @pytest.fixture(autouse=True)
 def _enable_automatic_dynamic():
     """Make sure the tests run with Dynamo's automatic-dynamic detection on.
@@ -206,6 +267,202 @@ def test_torch_compile_dynamic_true_single_compile(device: torch.device):
     )
 
 
+def test_mark_dynamic_seq_via_torch_compile_starts_dynamic(device: torch.device):
+    """Explicit `mark_dynamic` should skip the static-then-promote compile dance."""
+
+    class Mdl(torch.nn.Module):
+        def forward(self, x):
+            return (x.sin() + x.square()).sum(-1)
+
+    with _explicit_mark_dynamic_mode():
+        model = Mdl().eval().to(device)
+        counts: list[int] = []
+        capture: dict[str, object] = {}
+        compiled = _compile_with_capture(model, counts, capture)
+
+        first = torch.randn(2, 4, device=device)
+        torch._dynamo.mark_dynamic(first, 1, min=2, max=16)
+
+        inputs = {
+            4: first,
+            6: torch.randn(2, 6, device=device),
+            9: torch.randn(2, 9, device=device),
+        }
+
+        for seq_len, x in inputs.items():
+            ref = model(x)
+            out = compiled(x)
+            assert out.shape == ref.shape == (2,), (
+                f"seq_len={seq_len}: got {out.shape}, expected {ref.shape}"
+            )
+            assert torch.allclose(out, ref, atol=1e-5), (
+                f"seq_len={seq_len}: max_diff={torch.max(torch.abs(out - ref)).item():.2e}"
+            )
+
+        compiled_impl = capture["compiled_impl"]
+        assert compiled_impl.has_dynamic_dims
+        assert len(compiled_impl.dim_params) == 1
+
+        strip_ok, dynamic_shapes = _first_trace_dynamic_shapes(capture)
+        assert strip_ok, "Expected explicit mark_dynamic SymInts to be rewritten"
+        assert dynamic_shapes is not None
+        _assert_input_dynamic_dims(dynamic_shapes, 0, {1})
+
+        assert len(counts) == 1, (
+            "Explicit mark_dynamic should produce one dynamic backend trace from the start, "
+            f"got {len(counts)} backend invocations"
+        )
+
+
+def test_mark_dynamic_seq_with_lifted_weights_single_compile(device: torch.device):
+    """Lifted parameters should compose with an explicitly dynamic token axis."""
+
+    class Mdl(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.embed = torch.nn.Embedding(128, 16)
+            self.proj = torch.nn.Linear(16, 8)
+
+        def forward(self, input_ids):
+            return self.proj(self.embed(input_ids)).sum(-1)
+
+    with _explicit_mark_dynamic_mode():
+        model = Mdl().eval().to(device)
+        counts: list[int] = []
+        capture: dict[str, object] = {}
+        compiled = _compile_with_capture(model, counts, capture)
+
+        first = torch.tensor([[1, 2, 3, 4]], device=device)
+        torch._dynamo.mark_dynamic(first, 1, min=2, max=32)
+
+        inputs = {
+            4: first,
+            6: torch.arange(1, 7, device=device).unsqueeze(0),
+            9: torch.arange(1, 10, device=device).unsqueeze(0),
+        }
+
+        with torch.no_grad():
+            for seq_len, input_ids in inputs.items():
+                ref = model(input_ids)
+                out = compiled(input_ids)
+                assert out.shape == ref.shape == (1, seq_len), (
+                    f"seq_len={seq_len}: got {out.shape}, expected {ref.shape}"
+                )
+                assert torch.allclose(out, ref, atol=1e-5), (
+                    "seq_len="
+                    f"{seq_len}: max_diff={torch.max(torch.abs(out - ref)).item():.2e}"
+                )
+
+        compiled_impl = capture["compiled_impl"]
+        assert compiled_impl.has_dynamic_dims
+        assert len(compiled_impl.dim_params) == 1
+
+        strip_ok, dynamic_shapes = _first_trace_dynamic_shapes(capture)
+        assert strip_ok
+        assert dynamic_shapes is not None
+        _assert_input_dynamic_dims(dynamic_shapes, 0, {1})
+
+        assert len(counts) == 1, (
+            "Explicit mark_dynamic should avoid a second compile for lifted-weight models, "
+            f"got {len(counts)} backend invocations"
+        )
+
+
+def test_mark_dynamic_seq_preserves_affine_output_shape(device: torch.device):
+    """Output-shape expressions like `2 * seq` should stay dynamic from call 1."""
+
+    class Mdl(torch.nn.Module):
+        def forward(self, x):
+            return torch.cat([x, x], dim=1)
+
+    with _explicit_mark_dynamic_mode():
+        model = Mdl().eval().to(device)
+        counts: list[int] = []
+        capture: dict[str, object] = {}
+        compiled = _compile_with_capture(model, counts, capture)
+
+        first = torch.randn(2, 4, 3, device=device)
+        torch._dynamo.mark_dynamic(first, 1, min=2, max=16)
+
+        inputs = {
+            4: first,
+            5: torch.randn(2, 5, 3, device=device),
+            7: torch.randn(2, 7, 3, device=device),
+        }
+
+        for seq_len, x in inputs.items():
+            ref = model(x)
+            out = compiled(x)
+            assert out.shape == ref.shape == (2, 2 * seq_len, 3), (
+                f"seq_len={seq_len}: got {out.shape}, expected {ref.shape}"
+            )
+            assert torch.allclose(out, ref, atol=1e-5), (
+                f"seq_len={seq_len}: max_diff={torch.max(torch.abs(out - ref)).item():.2e}"
+            )
+
+        compiled_impl = capture["compiled_impl"]
+        assert compiled_impl.has_dynamic_dims
+        assert len(compiled_impl.dim_params) == 1
+
+        strip_ok, dynamic_shapes = _first_trace_dynamic_shapes(capture)
+        assert strip_ok
+        assert dynamic_shapes is not None
+        _assert_input_dynamic_dims(dynamic_shapes, 0, {1})
+
+        assert len(counts) == 1, (
+            "Explicit mark_dynamic should keep affine output-shape models on one compile, "
+            f"got {len(counts)} backend invocations"
+        )
+
+
+def test_mark_dynamic_two_dim_via_torch_compile_starts_dynamic(device: torch.device):
+    """Marking both batch and seq dynamic should still compile only once."""
+
+    class Mdl(torch.nn.Module):
+        def forward(self, x):
+            return x.mean(-1)
+
+    with _explicit_mark_dynamic_mode():
+        model = Mdl().eval().to(device)
+        counts: list[int] = []
+        capture: dict[str, object] = {}
+        compiled = _compile_with_capture(model, counts, capture)
+
+        first = torch.randn(2, 8, 4, device=device)
+        torch._dynamo.mark_dynamic(first, 0, min=1, max=8)
+        torch._dynamo.mark_dynamic(first, 1, min=2, max=16)
+
+        inputs = {
+            (2, 8): first,
+            (3, 9): torch.randn(3, 9, 4, device=device),
+            (5, 11): torch.randn(5, 11, 4, device=device),
+        }
+
+        for shape, x in inputs.items():
+            ref = model(x)
+            out = compiled(x)
+            assert out.shape == ref.shape == shape, (
+                f"shape={shape}: got {out.shape}, expected {ref.shape}"
+            )
+            assert torch.allclose(out, ref, atol=1e-5), (
+                f"shape={shape}: max_diff={torch.max(torch.abs(out - ref)).item():.2e}"
+            )
+
+        compiled_impl = capture["compiled_impl"]
+        assert compiled_impl.has_dynamic_dims
+        assert len(compiled_impl.dim_params) == 2
+
+        strip_ok, dynamic_shapes = _first_trace_dynamic_shapes(capture)
+        assert strip_ok
+        assert dynamic_shapes is not None
+        _assert_input_dynamic_dims(dynamic_shapes, 0, {0, 1})
+
+        assert len(counts) == 1, (
+            "Explicitly marked batch+seq dims should compile once from the first call, "
+            f"got {len(counts)} backend invocations"
+        )
+
+
 @pytest.mark.skipif(
     not torch.cuda.is_available(),
     reason="CUDA-only — exercises the cuda_lite dynamic-dim runtime",

crates/luminal_python/tests/test_kv_cache_growing.py

@@ -97,6 +97,87 @@ def test_kv_cache_growing():
     )
 
 
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="dynamic-cache torch.compile reuse requires CUDA coverage",
+)
+@pytest.mark.slow
+def test_dynamic_kv_cache_torch_compile_matches_reference_and_reuses_decode_graph():
+    """End-to-end server-style path: torch.compile + DynamicCache on CUDA."""
+    from transformers import DynamicCache, LlamaConfig, LlamaForCausalLM
+
+    backend_invocations = []
+
+    def counting_backend(gm, example_inputs, options=None):
+        backend_invocations.append((gm, example_inputs))
+        return luminal_backend(gm, example_inputs, options)
+
+    prev_auto = torch._dynamo.config.automatic_dynamic_shapes
+    prev_cache_limit = torch._dynamo.config.cache_size_limit
+    prev_recompile_limit = torch._dynamo.config.recompile_limit
+    torch._dynamo.config.automatic_dynamic_shapes = True
+    torch._dynamo.config.cache_size_limit = 16
+    torch._dynamo.config.recompile_limit = 16
+
+    try:
+        model = (
+            LlamaForCausalLM(
+                LlamaConfig(
+                    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=True,
+                    attn_implementation="eager",
+                )
+            )
+            .eval()
+            .cuda()
+        )
+        compiled = torch.compile(model, backend=counting_backend, fullgraph=True)
+
+        ref_cache = DynamicCache(config=model.config)
+        out_cache = DynamicCache(config=model.config)
+        input_ids = torch.tensor([[1, 2, 3, 4]], device="cuda")
+
+        with torch.no_grad():
+            ref = model(input_ids=input_ids, past_key_values=ref_cache, use_cache=True)
+            out = compiled(
+                input_ids=input_ids,
+                past_key_values=out_cache,
+                use_cache=True,
+            )
+
+        for _ in range(4):
+            ref_next = int(ref.logits[0, -1].argmax().item())
+            out_next = int(out.logits[0, -1].argmax().item())
+            assert out_next == ref_next
+            with torch.no_grad():
+                ref = model(
+                    input_ids=torch.tensor([[ref_next]], device="cuda"),
+                    past_key_values=ref.past_key_values,
+                    use_cache=True,
+                )
+                out = compiled(
+                    input_ids=torch.tensor([[out_next]], device="cuda"),
+                    past_key_values=out.past_key_values,
+                    use_cache=True,
+                )
+
+        assert len(backend_invocations) == 3, (
+            "Expected prefill/static decode/dynamic decode traces only once each, "
+            f"got {len(backend_invocations)} backend invocations"
+        )
+    finally:
+        torch._dynamo.config.automatic_dynamic_shapes = prev_auto
+        torch._dynamo.config.cache_size_limit = prev_cache_limit
+        torch._dynamo.config.recompile_limit = prev_recompile_limit
+        torch._dynamo.reset()
+
+
 @pytest.mark.skipif(
     not torch.cuda.is_available(),
     reason="R1 full-width 1-layer is too memory-heavy for CPU native backend",

crates/luminal_python/tests/test_llama3.py

@@ -450,3 +450,138 @@ def test_hf_llama38b_full(device: torch.device):
     assert torch.allclose(out.logits, ref.logits, atol=1e-5), (
         f"max_diff={torch.max(torch.abs(out.logits - ref.logits)).item():.2e}"
     )
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="Full Llama-3.1-8B dynamic-shape regression requires CUDA",
+)
+def test_hf_llama38b_mark_dynamic_seq_dim_before_compile(device: torch.device):
+    """Explicitly marking the token sequence dim dynamic should be honored end to end.
+
+    This exercises the real user path:
+      1. wrap the pretrained 8B model with ``torch.compile(..., backend=luminal_backend)``
+      2. mark ``input_ids.shape[1]`` dynamic before the first invocation
+      3. verify the first backend trace is already dynamic on that axis
+      4. reuse the same compiled graph for multiple sequence lengths
+    """
+    import copy
+
+    from transformers import AutoConfig, LlamaForCausalLM
+
+    from luminal.pt2 import (
+        _build_dynamic_shapes_from_gm,
+        _reinternalize_lifted_params,
+        _strip_symint_placeholders,
+    )
+
+    backend_invocations = []
+    capture = {}
+
+    def inspector_backend(gm, example_inputs, **kwargs):
+        backend_invocations.append((gm, example_inputs, kwargs))
+        if len(backend_invocations) == 1:
+            capture["gm"] = copy.deepcopy(gm).eval()
+            capture["example_inputs"] = example_inputs
+        compiled_impl = luminal_backend(gm, example_inputs, **kwargs)
+        if len(backend_invocations) == 1:
+            capture["compiled_impl"] = compiled_impl
+        return compiled_impl
+
+    prev_auto = torch._dynamo.config.automatic_dynamic_shapes
+    prev_cache_limit = torch._dynamo.config.cache_size_limit
+    torch._dynamo.reset()
+    torch._dynamo.config.automatic_dynamic_shapes = False
+    torch._dynamo.config.cache_size_limit = 8
+
+    try:
+        config = AutoConfig.from_pretrained("NousResearch/Meta-Llama-3.1-8B-Instruct")
+        config.use_cache = False
+        config._attn_implementation = "eager"
+
+        model = (
+            LlamaForCausalLM.from_pretrained(
+                "NousResearch/Meta-Llama-3.1-8B-Instruct",
+                config=config,
+                torch_dtype=torch.float32,
+            )
+            .eval()
+            .to(device)
+        )
+        compiled = torch.compile(model, backend=inspector_backend)
+
+        first_input_ids = torch.tensor([[1, 2, 3, 4]], device=device)
+        torch._dynamo.mark_dynamic(first_input_ids, 1, min=2, max=16)
+
+        seq_inputs = {
+            4: first_input_ids,
+            6: torch.arange(1, 7, device=device).unsqueeze(0),
+            9: torch.arange(1, 10, device=device).unsqueeze(0),
+        }
+
+        with torch.no_grad():
+            first_ref = model(first_input_ids)
+            first_out = compiled(first_input_ids)
+
+        compiled_impl = capture["compiled_impl"]
+        assert compiled_impl.has_dynamic_dims, (
+            "explicit mark_dynamic on input_ids[:, 1] should produce a dynamic Luminal graph"
+        )
+        assert len(compiled_impl.dim_params) == 1, (
+            f"expected exactly one dynamic dim param, got {compiled_impl.dim_params}"
+        )
+
+        gm = capture["gm"]
+        example_inputs = capture["example_inputs"]
+        gm, user_inputs, _, _ = _reinternalize_lifted_params(gm, example_inputs)
+        user_inputs, _, strip_ok = _strip_symint_placeholders(gm, user_inputs)
+        dynamic_shapes = _build_dynamic_shapes_from_gm(gm) if strip_ok else None
+
+        assert strip_ok, "Expected explicit mark_dynamic SymInts to be rewritten"
+        assert dynamic_shapes is not None, (
+            "Expected the first backend trace to preserve a dynamic shape spec"
+        )
+        args_spec = dynamic_shapes.get("args")
+        assert args_spec is not None and len(args_spec) == 1, (
+            f"expected one user-input dynamic spec, got {dynamic_shapes}"
+        )
+        assert args_spec[0] is not None, (
+            f"expected a per-dim dynamic spec for input_ids, got {dynamic_shapes}"
+        )
+        assert set(args_spec[0].keys()) == {1}, (
+            "Expected only the token sequence axis (dim=1) to be dynamic, "
+            f"got {dynamic_shapes}"
+        )
+
+        first_diff = torch.max(torch.abs(first_out.logits - first_ref.logits)).item()
+        assert torch.allclose(first_out.logits, first_ref.logits, atol=1e-3, rtol=0), (
+            f"seq_len=4: max_diff={first_diff:.2e}"
+        )
+
+        for seq_len, input_ids in seq_inputs.items():
+            with torch.no_grad():
+                ref = model(input_ids)
+                out = first_out if seq_len == 4 else compiled(input_ids)
+            assert (
+                out.logits.shape
+                == ref.logits.shape
+                == (
+                    1,
+                    seq_len,
+                    config.vocab_size,
+                )
+            ), f"seq_len={seq_len}: got {out.logits.shape}, expected {ref.logits.shape}"
+            assert torch.allclose(out.logits, ref.logits, atol=1e-3, rtol=0), (
+                f"seq_len={seq_len}: "
+                f"max_diff={torch.max(torch.abs(out.logits - ref.logits)).item():.2e}"
+            )
+
+        assert len(backend_invocations) == 1, (
+            "Explicit mark_dynamic should produce one dynamic backend trace from the start, "
+            f"got {len(backend_invocations)} backend invocations"
+        )
+    finally:
+        torch._dynamo.config.automatic_dynamic_shapes = prev_auto
+        torch._dynamo.config.cache_size_limit = prev_cache_limit
+        torch._dynamo.reset()

crates/luminal_python/tests/test_qwen3_moe.py

@@ -152,6 +152,57 @@ def test_hf_qwen3_moe_medium(device: torch.device):
     _run_hf_qwen3_moe_test(config, device, atol=1e-4)
 
 
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="bf16 grouped_mm coverage requires CUDA",
+)
+def test_hf_qwen3_moe_tiny_bf16(device: torch.device):
+    """HuggingFace Qwen3MoeForCausalLM — tiny bf16 path on CUDA.
+
+    Exercises the grouped-mm MoE lowering with bf16 weights/activations so we
+    catch mixed-dtype compile regressions without paying the full 30B checkpoint
+    cost. Like the full pretrained bf16 test below, this only asserts that the
+    compiled path runs and stays numerically sane; tight bf16 equivalence is
+    tracked separately.
+    """
+    from transformers import Qwen3MoeForCausalLM
+
+    config = _make_qwen3_moe_config(
+        hidden_size=32,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        num_hidden_layers=1,
+        intermediate_size=64,
+        moe_intermediate_size=64,
+        num_experts=2,
+        num_experts_per_tok=1,
+        vocab_size=128,
+    )
+
+    model = Qwen3MoeForCausalLM(config).eval().to(dtype=torch.bfloat16, device=device)
+    compiled = torch.compile(model, backend=luminal_backend)
+    input_ids = torch.tensor([[1, 2, 3, 4]], device=device)
+    with torch.no_grad():
+        ref = model(input_ids)
+        out = compiled(input_ids)
+
+    ref_logits = ref.logits.float()
+    out_logits = out.logits.float()
+    ref_max = ref_logits.abs().max().item()
+    out_max = out_logits.abs().max().item()
+    n_nan = int(out_logits.isnan().sum().item())
+    n_inf = int(out_logits.isinf().sum().item())
+
+    assert n_nan == 0 and n_inf == 0, (
+        f"compiled forward produced non-finite logits: {n_nan} NaNs, "
+        f"{n_inf} Infs (eager max abs={ref_max:.2f}, compiled max abs={out_max:.2f})"
+    )
+    assert 0.1 * ref_max <= out_max <= 10.0 * ref_max, (
+        f"compiled max abs={out_max:.2f} is out of band vs eager max abs={ref_max:.2f} "
+        f"(>10x off in either direction); likely a numerical/scale bug"
+    )
+
+
 @pytest.mark.slow
 def test_hf_qwen3_moe_real_config_1layer(device: torch.device):
     """HuggingFace Qwen3MoeForCausalLM — real Qwen3-30B-A3B architecture, 1 layer.

examples/dlrm/.gitignore

@@ -0,0 +1,11 @@
+# Correctness-test fixtures generated by `correctness_dump.py`.
+# Regenerate with: python correctness_dump.py --num-cat N --rows R --out-dir weights_N
+weights/
+weights_*/
+
+# Python bytecode
+__pycache__/
+
+# Intermediate sweep outputs from sweep_all.py — `results.csv` is the
+# canonical published result; everything else is regenerable.
+quick*.csv

examples/dlrm/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "dlrm"
+version = "0.1.0"
+edition = "2024"
+
+[[bin]]
+name = "dlrm"
+path = "src/main.rs"
+
+[[bin]]
+name = "check"
+path = "src/bin/check.rs"
+
+[dependencies]
+luminal = { path = "../.." }
+luminal_cuda_lite = { path = "../../crates/luminal_cuda_lite" }
+rand = "0.9.2"
+bytemuck = "1"

examples/dlrm/RESULTS.md

@@ -0,0 +1,236 @@
+# DLRMv1 via `torch.compile(backend=luminal_backend)` — sweep results
+
+## TL;DR
+
+- **Compiles end-to-end**: vanilla DLRMv1 (`nn.EmbeddingBag` × num_cat, then
+  pairwise-dot interaction, then top MLP) lands on `torch.compile` with
+  `backend=luminal_backend` for all 55 (`batch`, `num_cat`) cells in the
+  sweep.
+- **Correctness**: max abs diff ≤ 1.8 × 10⁻⁷ vs PyTorch eager at
+  `num_cat ∈ {2, 4, 8, 16, 32}` (essentially fp32 noise).
+- **Beats `pt_eager` in 55/55 cells (100%)**, by **4.7–7.1×**.
+- **Beats `graph_safe_inductor_cg` at the highest cell** (`nc=32, b=2048`):
+  174 μs vs 241 μs (1.38× faster). Still slower at smaller cells —
+  100 μs of fixed Python-wrapper overhead per call vs PT-CUDAGraph's
+  ~22 μs replay floor.
+
+## Hardware / setup
+
+- NVIDIA GH200 480GB, CUDA 12.8, driver 570.148.08.
+- DLRMv1 inline (matches `examples/dlrm/sweep_pytorch.py`'s `DLRMv1`).
+  Fixed dials: `m_den=3`, `m_spa=16`, `bag=2`, `rows=4096`.
+- Harness: 5 rounds × 20 iters × 10 warmup before round 0,
+  median-of-round-medians, CUDA-event timing.
+
+## Sweep dimensions
+
+| dim | values |
+|---|---|
+| `batch` | 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048 |
+| `num_cat` | 2, 4, 8, 16, 32 |
+
+55 cells × 3 variants = 165 timings. CSV: `results.csv`.
+
+## Variants
+
+- `pt_eager`: vanilla DLRMv1 forward, no compile, no graph capture.
+- `graph_safe_inductor_cg`: wrap DLRMv1 in `GraphSafeDLRM` (pre-bakes the
+  `li`/`lj` lower-tri index buffers so the forward is capturable), then
+  `torch.compile(backend="inductor", mode="max-autotune-no-cudagraphs",
+  fullgraph=False, dynamic=False)`, then manual
+  `torch.cuda.CUDAGraph` capture/replay. **This is the named PT baseline.**
+- `luminal_compiled`: `torch.compile(backend=luminal_backend, fullgraph=
+  False, dynamic=False)`, no external CUDAGraph wrap. luminal's
+  `cuda_lite` runtime captures/replays a CUDA graph internally on
+  every `execute()` call.
+
+## Wall-clock results (median-of-round-medians, μs)
+
+### `luminal_compiled`
+
+| nc \ batch |    2 |   16 |   64 |  256 | 1024 | 2048 |
+|---|---|---|---|---|---|---|
+|   2 |  101 |   95 |   96 |  100 |   98 |  102 |
+|   4 |  100 |   99 |  110 |  103 |  104 |  112 |
+|   8 |  109 |  106 |  116 |  110 |  111 |  121 |
+|  16 |  131 |  133 |  134 |  136 |  140 |  135 |
+|  32 |  172 |  172 |  178 |  174 |  176 |  174 |
+
+### `graph_safe_inductor_cg` — the named PT baseline
+
+| nc \ batch |    2 |   16 |   64 |  256 | 1024 | 2048 |
+|---|---|---|---|---|---|---|
+|   2 |   22 |   24 |   26 |   27 |   30 |   34 |
+|   4 |   27 |   29 |   31 |   33 |   35 |   45 |
+|   8 |   46 |   48 |   52 |   56 |   57 |   66 |
+|  16 |   75 |   79 |   91 |   96 |   96 |  121 |
+|  32 |  132 |  138 |  164 |  157 |  170 |  241 |
+
+### `pt_eager`
+
+| nc \ batch |    2 |   16 |   64 |  256 | 1024 | 2048 |
+|---|---|---|---|---|---|---|
+|   2 |  506 |  453 |  472 |  496 |  578 |  555 |
+|   4 |  514 |  488 |  569 |  535 |  605 |  588 |
+|   8 |  624 |  632 |  709 |  639 |  676 |  604 |
+|  16 |  786 |  783 |  879 |  813 |  862 |  865 |
+|  32 | 1125 | 1133 | 1180 | 1169 | 1176 | 1233 |
+
+## Speedup vs `graph_safe_inductor_cg` (>1 = luminal faster)
+
+| nc \ batch |    2 |   16 |   64 |  256 | 1024 | 2048 |
+|---|---|---|---|---|---|---|
+|   2 | 0.22× | 0.25× | 0.27× | 0.28× | 0.30× | 0.33× |
+|   4 | 0.27× | 0.29× | 0.29× | 0.32× | 0.34× | 0.40× |
+|   8 | 0.42× | 0.45× | 0.45× | 0.51× | 0.52× | 0.55× |
+|  16 | 0.57× | 0.60× | 0.68× | 0.71× | 0.68× | 0.90× |
+|  32 | 0.77× | 0.80× | 0.92× | 0.91× | 0.97× | **1.38×** |
+
+## Speedup vs `pt_eager` (>1 = luminal faster)
+
+luminal_compiled is **4.7–7.1× faster** at every cell. Won at 55/55 cells (100%).
+
+## What changed to get here
+
+Three rounds of wrapper-overhead reduction (all in `crates/luminal_python/`
+and `crates/luminal_cuda_lite/`):
+
+### Round 1: Translator + cuda_lite kernel matching
+Made `torch.compile(model, backend=luminal_backend)` emit the same kernel
+set as the hand-written rust DLRM:
+
+| FX subgraph | luminal kernel |
+|---|---|
+| `addmm(bias, x, weight.t())` + `relu/sigmoid` consumer | `linear_bias_(relu|sigmoid)` (one fused kernel per MLP layer) |
+| N × `_embedding_bag(W_k, idx_k, off_k)` | `multi_table_embedding_bag_sum_kernel` (one fused kernel for all tables) |
+| `index.Tensor(bmm(cat(unsqueezes), permute([0,2,1])), [None, li, lj])` | `dlrm_pairwise_dot_lower_tri_stacked` (one fused kernel) |
+| `bmm(A, permute(B, [0,2,1]))` (generic) | `matmul_3d_t` |
+
+Plus a post-translation DCE pass to drop the now-unreachable
+`bmm/cat/permute` chain superseded by the pairwise-dot peephole.
+
+### Round 2: Batched FFI + by-id setter
+Added `tensor_id(name) → u32` and `run_with_ptrs(inputs, outputs)` on the
+Rust side (one pyo3 hop instead of N), and cached the IDs once at
+`CompiledModel.__init__`. The Python wrapper now passes
+`[(id, ptr, n_bytes), …]` lists instead of N separate
+`set_input_device_ptr(name, ptr, n_bytes)` calls.
+
+### Round 3: Skip re-registration on unchanged inputs
+The Python wrapper now caches `(id(tensor), data_ptr)` per input slot.
+On a hot bench loop the same user tensors are passed each iteration —
+all 65 inputs hit the cache, the `input_specs` list passed to
+`run_with_ptrs` is empty, and the runtime relies on the previously-
+registered pointers. Required one runtime change: in
+`cuda_lite/runtime.rs::execute()`'s post-run consume step, don't drop
+`CudaInput::Ptr` entries — they're non-owning views over caller memory
+and must persist across `execute()` calls for the skip path to work.
+
+## Where the per-iter time goes (nc=32, b=2048)
+
+Before any of the rounds:
+
+| section | μs/iter |
+|---|---|
+| `set_input_device_ptr` × 65 (Python loop + FFI) | 995 |
+| `_graph.run()` (CUDA-graph replay) | 61 |
+| set_output + alloc + collect | 13 |
+| **total** | **1069** |
+
+After all three rounds:
+
+| section | μs/iter |
+|---|---|
+| input cache check (0 cold-path inputs after warmup) | 15 |
+| `torch.empty` + register output | 11 |
+| `run_with_ptrs` (batched FFI, all caches hit, runtime executes graph) | 56 |
+| **total** | **82** |
+
+**13× per-iter reduction**, none of it in the kernels themselves.
+
+## Hand-written rust DLRM (`examples/dlrm/src/main.rs`)
+
+The hand-written rust uses the same `cuda_lite` kernels directly — no
+Python, no `torch.compile`. On `nc=32, b=2048`: **104 μs**
+(median-of-round-medians with explicit `synchronize_stream` for accurate
+GPU time). The `luminal_compiled` figure for the same cell is 174 μs —
+the residual 70 μs gap is roughly:
+
+- `torch.empty` + output registration (~11 μs)
+- Python-side cache check + 1 round-trip into `run_with_ptrs` (~10 μs)
+- Marshalling 1 input spec + 1 output spec across the pyo3 boundary
+- pt2 backend wrapper invocation by torch.compile (a few μs of
+  dynamo/eval_frame work)
+- The remaining ~40 μs is the runtime's per-call exec_op iteration +
+  buffer_map building inside `cuda_lite/src/runtime.rs::execute()` —
+  a structural cost of how the runtime dispatches host ops, paid once
+  per call regardless of input count.
+
+## What's left
+
+In scope, deferred:
+
+- `linear_bias_relu_split_a` peephole on top MLP first layer. Saves
+  one materialized `cat` + one small kernel. Modest cell-by-cell win.
+
+Out of scope:
+
+- The remaining ~80–100 μs of wrapper / runtime fixed overhead is a
+  combination of `torch.compile`'s dynamo eval-frame dispatch, the
+  Python `__call__` setup work, and the runtime's per-execute toposort
+  + buffer_map build. Closing it further would either need a deeper
+  rewrite of `runtime.rs::execute()` (probably worth it independent of
+  DLRM) or a way to skip dynamo's per-call overhead.
+
+## Files this work touches
+
+In scope:
+
+- `crates/luminal_python/rust/src/translator/{mod,dispatch,movement}.rs` —
+  peephole infra, op handlers, multi-call fusions, post-translation DCE.
+- `crates/luminal_cuda_lite/src/kernel/{embedding_bag,dlrm_interact,
+  matmul2d,mod}.rs` — fused kernels for embedding bag (single + multi-
+  table + stacked), pairwise dot (variadic + stacked), and fused-
+  activation linear (relu/sigmoid + split-A).
+- `crates/luminal_cuda_lite/src/runtime.rs` — gate end-of-execute
+  stream sync on `profiling`; expose `synchronize_stream()` and
+  `read_per_kernel_timings_ms()` stub; preserve `CudaInput::Ptr`
+  entries across `execute()` calls (the unchanged-input cache fix).
+
+Wrapper layer (relaxed scope; minimal targeted changes):
+
+- `crates/luminal_python/rust/src/compiled_graph.rs` — `tensor_id(name)`
+  + `run_with_ptrs(inputs, outputs)` for batched FFI.
+- `crates/luminal_python/src/luminal/compiled_model.py` — fast-path
+  `run_with_ptrs` call with the per-input cache.
+
+Leaf consumer (allowed):
+
+- `examples/dlrm/` — bench harness, hand-written rust DLRM, correctness
+  check, sweep scripts, this report. Cherry-picked from
+  `origin/dlrm-fused-kernels` then adapted (inline DLRMv1, no
+  upstream `dlrm_s_pytorch` dependency, added third PT variant).
+
+## Reproduction
+
+From `examples/dlrm`:
+
+```bash
+# Build
+(cd ../.. && cargo build -p dlrm --release)
+(cd ../../crates/luminal_python && CUDARC_CUDA_VERSION=12080 \
+    uv run --group dev maturin develop --manifest-path rust/Cargo.toml \
+    --features cuda -r)
+
+# Correctness check (PyTorch eager → luminal hand-written rust)
+for nc in 2 4 8 16 32; do
+    python correctness_dump.py --num-cat $nc --rows 4096
+    (cd ../.. && ./target/release/check)
+done
+
+# Full 55-cell sweep, 3 variants
+python sweep_all.py --variants pt_eager graph_safe_inductor_cg luminal_compiled
+
+# Pretty summary
+python summarize.py results.csv
+```

examples/dlrm/bench_pytorch.py

@@ -0,0 +1,348 @@
+"""PyTorch reference for the DLRM `_make_dlrm_batch_2048` config.
+
+Mirrors `bench_luminal_exact.py` from https://github.com/jss8649/tmp-dlrm-bench
+in spirit (shape-for-shape, same indices, same MLPs, same interaction op)
+but reimplements the model inline so we don't need the upstream
+facebookresearch/dlrm package.
+
+Measures:
+  * eager
+  * torch.compile (inductor, default)
+  * torch.compile (inductor, mode="reduce-overhead"  → CUDA-graph capture)
+  * the v3 fused trick (index_select + reshape + sum on a stacked table)
+    with reduce-overhead — that's the WINNER in the reference repo.
+
+Harness: 5 rounds × 20 iters, 10 warmup, median-of-round-medians (CUDA-event
+timing). Matches the luminal-side harness verbatim.
+
+Usage:
+    python bench_pytorch.py
+"""
+from __future__ import annotations
+
+import copy
+from contextlib import contextmanager
+from typing import List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+# ---- Config (matches `_make_dlrm_batch_2048`) -----------------------------
+BATCH = 2048
+M_DEN = 3
+M_SPA = 16
+INDICES_PER_BAG = 2
+LN_EMB = np.array([4096, 2048, 1024], dtype=np.int64)
+NUM_EMB = len(LN_EMB)
+NUM_FEA = NUM_EMB + 1
+PAIR_COUNT = NUM_FEA * (NUM_FEA - 1) // 2  # strict lower tri, no diagonal
+TOP_IN = PAIR_COUNT + M_SPA                 # 6 + 16 = 22
+LN_BOT = [M_DEN, 64, M_SPA]                 # [3, 64, 16]
+LN_TOP_TAIL = [64, 32, 1]                   # ln_top = [22, 64, 32, 1]
+SEED = 0
+
+
+def _build_mlp(layer_sizes: List[int], sigmoid_layer: int) -> nn.Sequential:
+    layers: List[nn.Module] = []
+    for i, (a, b) in enumerate(zip(layer_sizes, layer_sizes[1:])):
+        layers.append(nn.Linear(a, b, bias=True))
+        if i == sigmoid_layer:
+            layers.append(nn.Sigmoid())
+        else:
+            layers.append(nn.ReLU())
+    return nn.Sequential(*layers)
+
+
+# ---- v1-shape model: EmbeddingBag per table (the "natural" expression
+#       a user writes; same shape as the upstream DLRM forward) ------------
+class DLRMv1(nn.Module):
+    def __init__(self):
+        super().__init__()
+        torch.manual_seed(SEED)
+        np.random.seed(SEED)
+        ln_top = [TOP_IN] + LN_TOP_TAIL  # [22, 64, 32, 1]
+        self.bot = _build_mlp(LN_BOT, sigmoid_layer=-1)
+        # sigmoid_top in the upstream is `len(ln_top) - 2` = 2 → final linear
+        self.top = _build_mlp(ln_top, sigmoid_layer=len(ln_top) - 2)
+        self.emb = nn.ModuleList(
+            [
+                nn.EmbeddingBag(int(n), M_SPA, mode="sum", sparse=False)
+                for n in LN_EMB
+            ]
+        )
+        # Pre-compute strict lower-tri (no diagonal) indices.
+        li, lj = [], []
+        for i in range(NUM_FEA):
+            for j in range(i):
+                li.append(i)
+                lj.append(j)
+        self.register_buffer("li", torch.tensor(li, dtype=torch.long), persistent=False)
+        self.register_buffer("lj", torch.tensor(lj, dtype=torch.long), persistent=False)
+
+    def forward(self, dense_x, lS_o, lS_i):
+        x = self.bot(dense_x)  # (B, M_SPA)
+        ly = [
+            self.emb[k](lS_i[k], lS_o[k]) for k in range(NUM_EMB)
+        ]  # each (B, M_SPA)
+        T = torch.cat([x.unsqueeze(1)] + [e.unsqueeze(1) for e in ly], dim=1)  # (B, F, M_SPA)
+        Z = torch.bmm(T, T.transpose(1, 2))  # (B, F, F)
+        Zflat = Z[:, self.li, self.lj]       # (B, PAIRS)
+        R = torch.cat([x, Zflat], dim=1)     # (B, M_SPA + PAIRS)
+        return self.top(R)
+
+
+# ---- v3 fused: stacked embedding table, index_select + reshape + sum ------
+#       This is the winner per perf.md (Inductor can fuse gather+sum into
+#       a single Triton kernel; opaque EmbeddingBag blocks that fusion). ----
+class DLRMv3(nn.Module):
+    def __init__(self):
+        super().__init__()
+        torch.manual_seed(SEED)
+        np.random.seed(SEED)
+        total_rows = int(LN_EMB.sum())
+        self.num_emb = NUM_EMB
+        self.m_spa = M_SPA
+        self.L = INDICES_PER_BAG
+        big_w = np.empty((total_rows, M_SPA), dtype=np.float32)
+        starts = np.zeros(NUM_EMB, dtype=np.int64)
+        s = 0
+        for k, n in enumerate(LN_EMB):
+            starts[k] = s
+            big_w[s : s + int(n)] = np.random.uniform(
+                -np.sqrt(1.0 / int(n)),
+                np.sqrt(1.0 / int(n)),
+                (int(n), M_SPA),
+            ).astype(np.float32)
+            s += int(n)
+        self.emb_weight = nn.Parameter(torch.from_numpy(big_w))
+        self.register_buffer("row_offsets", torch.from_numpy(starts))
+
+        ln_top = [TOP_IN] + LN_TOP_TAIL
+        self.bot = _build_mlp(LN_BOT, sigmoid_layer=-1)
+        self.top = _build_mlp(ln_top, sigmoid_layer=len(ln_top) - 2)
+
+        li, lj = [], []
+        for i in range(NUM_FEA):
+            for j in range(i):
+                li.append(i)
+                lj.append(j)
+        self.register_buffer("li", torch.tensor(li, dtype=torch.long), persistent=False)
+        self.register_buffer("lj", torch.tensor(lj, dtype=torch.long), persistent=False)
+
+    def forward(self, dense_x, flat_indices):
+        bs = dense_x.shape[0]
+        gathered = torch.index_select(self.emb_weight, 0, flat_indices)
+        gathered = gathered.view(self.num_emb * bs, self.L, self.m_spa)
+        pooled = gathered.sum(dim=1)  # (num_emb*B, m_spa)
+        ly = pooled.view(self.num_emb, bs, self.m_spa).transpose(0, 1)  # (B, num_emb, m_spa)
+        x = self.bot(dense_x)
+        T = torch.cat([x.unsqueeze(1), ly], dim=1)
+        Z = torch.bmm(T, T.transpose(1, 2))
+        Zflat = Z[:, self.li, self.lj]
+        R = torch.cat([x, Zflat], dim=1)
+        return self.top(R)
+
+
+# ---- Deterministic inputs matching `_make_dlrm_batch_2048` ---------------
+def make_v1_inputs(device):
+    dense_x = (
+        torch.linspace(-1.0, 1.0, BATCH * M_DEN, dtype=torch.float32, device=device)
+        .reshape(BATCH, M_DEN)
+    )
+    total = BATCH * INDICES_PER_BAG
+    positions = torch.arange(total, dtype=torch.int64, device=device)
+    offsets = torch.arange(0, total, INDICES_PER_BAG, dtype=torch.int64, device=device)
+    lS_o = [offsets.clone() for _ in range(NUM_EMB)]
+    lS_i = [
+        ((positions * 3 + 1) % int(LN_EMB[0])).to(torch.int64),
+        ((positions * 5 + 2) % int(LN_EMB[1])).to(torch.int64),
+        ((positions * 7 + 3) % int(LN_EMB[2])).to(torch.int64),
+    ]
+    return dense_x, lS_o, lS_i
+
+
+def make_v3_inputs(device, model: DLRMv3):
+    dense_x, _, lS_i = make_v1_inputs(device)
+    # Stack and add per-table row offsets so a single index_select pulls
+    # from the unified table.
+    stacked = torch.stack(lS_i, dim=0)  # (NUM_EMB, B*L)
+    flat = (stacked + model.row_offsets.view(NUM_EMB, 1)).reshape(-1)
+    return dense_x, flat
+
+
+# ---- Timing harness (mirrors bench_luminal_exact.py) ---------------------
+@contextmanager
+def _relaxed_dynamo_limits():
+    prev_r = torch._dynamo.config.recompile_limit
+    prev_c = torch._dynamo.config.cache_size_limit
+    torch._dynamo.config.recompile_limit = 16
+    torch._dynamo.config.cache_size_limit = 16
+    try:
+        yield
+    finally:
+        torch._dynamo.config.recompile_limit = prev_r
+        torch._dynamo.config.cache_size_limit = prev_c
+
+
+def _compile_inductor(model):
+    with _relaxed_dynamo_limits():
+        torch._dynamo.reset()
+        return torch.compile(copy.deepcopy(model), backend="inductor")
+
+
+def _compile_inductor_ro(model):
+    with _relaxed_dynamo_limits():
+        torch._dynamo.reset()
+        return torch.compile(
+            copy.deepcopy(model), backend="inductor", mode="reduce-overhead"
+        )
+
+
+def _timed_cuda_runs(model, inputs, warmup, timed, mark_step=False):
+    with torch.no_grad():
+        for _ in range(warmup):
+            if mark_step:
+                torch.compiler.cudagraph_mark_step_begin()
+            _ = model(*inputs)
+    torch.cuda.synchronize()
+    starts = [torch.cuda.Event(enable_timing=True) for _ in range(timed)]
+    ends = [torch.cuda.Event(enable_timing=True) for _ in range(timed)]
+    with torch.no_grad():
+        for i in range(timed):
+            if mark_step:
+                torch.compiler.cudagraph_mark_step_begin()
+            starts[i].record()
+            _ = model(*inputs)
+            ends[i].record()
+    torch.cuda.synchronize()
+    return np.array(
+        [s.elapsed_time(e) for s, e in zip(starts, ends)], dtype=np.float64
+    )
+
+
+def _timed_rounds(model, inputs, *, warmup, timed, rounds, mark_step=False):
+    arr = _timed_cuda_runs(model, inputs, warmup, timed, mark_step=mark_step)
+    round_medians = [float(np.median(arr))]
+    for _ in range(rounds - 1):
+        arr = _timed_cuda_runs(model, inputs, 0, timed, mark_step=mark_step)
+        round_medians.append(float(np.median(arr)))
+    return round_medians
+
+
+def _manual_cudagraph_capture(model, inputs):
+    """Capture model(inputs) into a CUDA graph and return a replayable
+    closure that runs zero new launches (just cuGraphLaunch)."""
+    # Warm up.
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s), torch.no_grad():
+        for _ in range(3):
+            _ = model(*inputs)
+    torch.cuda.current_stream().wait_stream(s)
+    torch.cuda.synchronize()
+
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g), torch.no_grad():
+        out = model(*inputs)
+    torch.cuda.synchronize()
+
+    def replay():
+        g.replay()
+        return out
+    return replay
+
+
+def main():
+    print(f"torch={torch.__version__}  cuda={torch.version.cuda}")
+    print(f"device={torch.cuda.get_device_name(0)}  cap={torch.cuda.get_device_capability(0)}")
+    print(f"float32 matmul precision: {torch.get_float32_matmul_precision()}")
+    device = torch.device("cuda")
+
+    torch.manual_seed(SEED)
+    np.random.seed(SEED)
+
+    # v1: EmbeddingBag-style
+    v1 = DLRMv1().to(device).eval()
+    v1_inputs = make_v1_inputs(device)
+    v1_eager = copy.deepcopy(v1).to(device).eval()
+    v1_ind = _compile_inductor(v1)
+    v1_ind_ro = _compile_inductor_ro(v1)
+
+    # v3: index_select + reshape + sum on stacked table
+    v3 = DLRMv3().to(device).eval()
+    v3_inputs = make_v3_inputs(device, v3)
+    v3_ind_ro = _compile_inductor_ro(v3)
+
+    # Prime
+    with torch.no_grad():
+        _ = v1_eager(*v1_inputs)
+        torch.compiler.cudagraph_mark_step_begin()
+        _ = v1_ind_ro(*v1_inputs)
+        _ = v1_ind(*v1_inputs)
+        torch.compiler.cudagraph_mark_step_begin()
+        _ = v3_ind_ro(*v3_inputs)
+    torch.cuda.synchronize()
+
+    # Manual CUDA graph capture for v3 eager
+    v3_eager_replay = _manual_cudagraph_capture(v3, v3_inputs)
+    # Manual CUDA graph capture for v1 eager
+    v1_eager_replay = _manual_cudagraph_capture(v1, v1_inputs)
+
+    rounds, iters, warmup = 5, 20, 10
+
+    def report(label, model, inputs, mark_step=False, *, raw_replay=False):
+        if raw_replay:
+            # Time the replay closure directly.
+            replay = model
+            with torch.no_grad():
+                for _ in range(warmup):
+                    replay()
+            torch.cuda.synchronize()
+            starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+            ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+            with torch.no_grad():
+                for i in range(iters):
+                    starts[i].record()
+                    replay()
+                    ends[i].record()
+            torch.cuda.synchronize()
+            arr = np.array([s.elapsed_time(e) for s, e in zip(starts, ends)])
+            rms = [float(np.median(arr))]
+            for _ in range(rounds - 1):
+                with torch.no_grad():
+                    for i in range(iters):
+                        starts[i].record()
+                        replay()
+                        ends[i].record()
+                torch.cuda.synchronize()
+                arr = np.array(
+                    [s.elapsed_time(e) for s, e in zip(starts, ends)]
+                )
+                rms.append(float(np.median(arr)))
+        else:
+            rms = _timed_rounds(
+                model, inputs, warmup=warmup, timed=iters, rounds=rounds,
+                mark_step=mark_step,
+            )
+        rms_sorted = sorted(rms)
+        med = rms_sorted[len(rms_sorted) // 2]
+        tput = BATCH / (med / 1000.0)
+        print(
+            f"  {label:<60}  median {med:7.3f} ms   ({tput:>9,.0f} samples/s)   "
+            f"round medians: [{', '.join(f'{v:.3f}' for v in rms)}]"
+        )
+
+    print()
+    print(f"PyTorch reference (5 rounds x 20 iters, 10 warmup), batch={BATCH}:")
+    report("v1 eager",                          v1_eager, v1_inputs)
+    report("v1 torch.compile(inductor)",        v1_ind,   v1_inputs)
+    report("v1 torch.compile(reduce-overhead)", v1_ind_ro, v1_inputs, mark_step=True)
+    report("v1 eager + manual CUDAGraph",       v1_eager_replay, None, raw_replay=True)
+    report("v3 torch.compile(reduce-overhead)", v3_ind_ro, v3_inputs, mark_step=True)
+    report("v3 eager + manual CUDAGraph",       v3_eager_replay, None, raw_replay=True)
+
+
+if __name__ == "__main__":
+    main()

examples/dlrm/correctness_dump.py

@@ -0,0 +1,172 @@
+"""Dump a deterministic PyTorch DLRMv1 forward to disk so a parallel
+luminal binary can load the exact same weights/inputs and verify it
+produces the same output.
+
+Writes one f32 little-endian binary blob per tensor, plus a manifest
+JSON describing the shapes. All paths are under `weights/`.
+
+Usage: python correctness_dump.py [--num-cat N] [--rows R]
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import List
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+SEED = 1234
+BATCH = 2048
+M_DEN = 3
+M_SPA = 16
+L = 2
+LN_BOT = [M_DEN, 64, M_SPA]
+LN_TOP_TAIL = [64, 32, 1]
+
+
+def _build_mlp(layer_sizes: List[int], sigmoid_layer: int) -> nn.Sequential:
+    layers: List[nn.Module] = []
+    for i, (a, b) in enumerate(zip(layer_sizes, layer_sizes[1:])):
+        layers.append(nn.Linear(a, b, bias=True))
+        layers.append(nn.Sigmoid() if i == sigmoid_layer else nn.ReLU())
+    return nn.Sequential(*layers)
+
+
+class DLRMv1(nn.Module):
+    def __init__(self, num_cat: int, rows: int):
+        super().__init__()
+        torch.manual_seed(SEED)
+        np.random.seed(SEED)
+        self.num_cat = num_cat
+        ni = num_cat + 1
+        num_int = ni * (ni - 1) // 2 + M_SPA
+        ln_top = [num_int] + LN_TOP_TAIL
+        self.bot = _build_mlp(LN_BOT, sigmoid_layer=-1)
+        # `sigmoid_top=2` in upstream means sigmoid on the final linear,
+        # which corresponds to sigmoid_layer = len(layers in Sequential) - 1
+        # Our `_build_mlp` indexes by Linear-index, so the last linear has
+        # index len(ln_top) - 2.
+        self.top = _build_mlp(ln_top, sigmoid_layer=len(ln_top) - 2)
+        self.emb = nn.ModuleList(
+            [nn.EmbeddingBag(rows, M_SPA, mode="sum", sparse=False) for _ in range(num_cat)]
+        )
+        li, lj = [], []
+        for i in range(ni):
+            for j in range(i):
+                li.append(i)
+                lj.append(j)
+        self.register_buffer("li", torch.tensor(li, dtype=torch.long), persistent=False)
+        self.register_buffer("lj", torch.tensor(lj, dtype=torch.long), persistent=False)
+
+    def forward(self, dense_x, lS_o, lS_i):
+        x = self.bot(dense_x)
+        ly = [self.emb[k](lS_i[k], lS_o[k]) for k in range(self.num_cat)]
+        T = torch.cat([x.unsqueeze(1)] + [e.unsqueeze(1) for e in ly], dim=1)
+        Z = torch.bmm(T, T.transpose(1, 2))
+        Zflat = Z[:, self.li, self.lj]
+        R = torch.cat([x, Zflat], dim=1)
+        return self.top(R)
+
+
+def build_indices(table_idx: int, batch: int, bag: int, rows: int) -> np.ndarray:
+    # Match sweep_categories.py and the luminal Rust binary exactly.
+    s = 2 * table_idx + 3
+    o = table_idx + 1
+    pos = np.arange(batch * bag, dtype=np.int64)
+    return (pos * s + o) % rows
+
+
+def build_dense_x(batch: int, m_den: int) -> np.ndarray:
+    total = batch * m_den
+    return np.linspace(-1.0, 1.0, num=total, dtype=np.float32).reshape(batch, m_den)
+
+
+def write_f32(path: Path, arr: np.ndarray) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    arr.astype(np.float32, copy=False).tofile(path)
+
+
+def write_i32(path: Path, arr: np.ndarray) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    arr.astype(np.int32, copy=False).tofile(path)
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--num-cat", type=int, default=3)
+    ap.add_argument("--rows", type=int, default=4096)
+    ap.add_argument("--out-dir", type=str, default="weights")
+    args = ap.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    model = DLRMv1(args.num_cat, args.rows).to(device).eval()
+
+    dense_np = build_dense_x(BATCH, M_DEN)
+    dense = torch.from_numpy(dense_np).to(device)
+    offsets = torch.arange(0, BATCH * L, L, dtype=torch.int64, device=device)
+    lS_o = [offsets.clone() for _ in range(args.num_cat)]
+    lS_i = [
+        torch.from_numpy(build_indices(k, BATCH, L, args.rows)).to(device)
+        for k in range(args.num_cat)
+    ]
+
+    with torch.no_grad():
+        out = model(dense, lS_o, lS_i)
+    print(f"output[:8] = {out.detach().cpu().numpy().flatten()[:8]}")
+    print(f"output stats: min={out.min().item():.6f} max={out.max().item():.6f} "
+          f"mean={out.mean().item():.6f}")
+
+    # ---- dump weights ----
+    # Bottom MLP: linears at indices 0 and 2 in the Sequential (Linear, ReLU,
+    # Linear). Luminal stores W as (out, in), matching PyTorch's nn.Linear.
+    bot_lins = [m for m in model.bot if isinstance(m, nn.Linear)]
+    top_lins = [m for m in model.top if isinstance(m, nn.Linear)]
+    for i, l in enumerate(bot_lins):
+        write_f32(out_dir / f"bot_{i}_w.bin", l.weight.detach().cpu().numpy())
+        write_f32(out_dir / f"bot_{i}_b.bin", l.bias.detach().cpu().numpy())
+    for i, l in enumerate(top_lins):
+        write_f32(out_dir / f"top_{i}_w.bin", l.weight.detach().cpu().numpy())
+        write_f32(out_dir / f"top_{i}_b.bin", l.bias.detach().cpu().numpy())
+    for k, e in enumerate(model.emb):
+        write_f32(out_dir / f"emb_{k}.bin", e.weight.detach().cpu().numpy())
+
+    # ---- dump inputs ----
+    write_f32(out_dir / "dense.bin", dense_np)
+    for k, idx in enumerate(lS_i):
+        write_i32(out_dir / f"idx_{k}.bin", idx.cpu().numpy())
+
+    # ---- dump expected output ----
+    write_f32(out_dir / "expected.bin", out.detach().cpu().numpy())
+
+    manifest = {
+        "num_cat": args.num_cat,
+        "rows": args.rows,
+        "batch": BATCH,
+        "m_den": M_DEN,
+        "m_spa": M_SPA,
+        "indices_per_bag": L,
+        "ln_bot": LN_BOT,
+        "ln_top": [args.num_cat * (args.num_cat + 1) // 2 + M_SPA] + LN_TOP_TAIL,
+        "bot_layer_shapes": [list(l.weight.shape) for l in bot_lins],
+        "top_layer_shapes": [list(l.weight.shape) for l in top_lins],
+        "output_shape": list(out.shape),
+        "output_head": out.detach().cpu().numpy().flatten()[:8].tolist(),
+    }
+    with open(out_dir / "manifest.json", "w") as f:
+        json.dump(manifest, f, indent=2)
+
+    print(f"\nWrote weights/inputs/expected to {out_dir}/")
+    print(f"  manifest: {out_dir}/manifest.json")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

examples/dlrm/results.csv

@@ -0,0 +1,166 @@
+variant,num_cat,batch,m_spa,bag,rows,ms,samples_per_sec,status
+pt_eager,2,2,16,2,4096,0.5056319832801819,3955.4459886524933,ok
+graph_safe_inductor_cg,2,2,16,2,4096,0.022352000698447227,89477.44888621707,ok
+luminal_compiled,2,2,16,2,4096,0.10134400054812431,19734.764654867537,ok
+pt_eager,2,4,16,2,4096,0.4413599967956543,9062.89656752006,ok
+graph_safe_inductor_cg,2,4,16,2,4096,0.022592000663280487,177053.8191644678,ok
+luminal_compiled,2,4,16,2,4096,0.09467199817299843,42251.141596173125,ok
+pt_eager,2,8,16,2,4096,0.4634399861097336,17262.21353309327,ok
+graph_safe_inductor_cg,2,8,16,2,4096,0.02404800057411194,332667.9893966789,ok
+luminal_compiled,2,8,16,2,4096,0.09742399677634239,82115.29258408182,ok
+pt_eager,2,16,16,2,4096,0.4532639980316162,35299.516549920125,ok
+graph_safe_inductor_cg,2,16,16,2,4096,0.023744000121951103,673854.4439783823,ok
+luminal_compiled,2,16,16,2,4096,0.09547200053930283,167588.40193584614,ok
+pt_eager,2,32,16,2,4096,0.5251840054988861,60931.025440507,ok
+graph_safe_inductor_cg,2,32,16,2,4096,0.02582399919629097,1239157.411552122,ok
+luminal_compiled,2,32,16,2,4096,0.09985600039362907,320461.46324564423,ok
+pt_eager,2,64,16,2,4096,0.47198399901390076,135597.81715844796,ok
+graph_safe_inductor_cg,2,64,16,2,4096,0.02619200013577938,2443494.184034204,ok
+luminal_compiled,2,64,16,2,4096,0.09620799869298935,665225.3541229069,ok
+pt_eager,2,128,16,2,4096,0.5091679990291595,251390.50420305296,ok
+graph_safe_inductor_cg,2,128,16,2,4096,0.027855999767780304,4595060.348472987,ok
+luminal_compiled,2,128,16,2,4096,0.10607999935746193,1206636.508062876,ok
+pt_eager,2,256,16,2,4096,0.49588799476623535,516245.6092946552,ok
+graph_safe_inductor_cg,2,256,16,2,4096,0.027456000447273254,9324009.172116116,ok
+luminal_compiled,2,256,16,2,4096,0.09959999844431877,2570281.1646439577,ok
+pt_eager,2,512,16,2,4096,0.499007984995842,1026035.6855898134,ok
+graph_safe_inductor_cg,2,512,16,2,4096,0.036959998309612274,13852814.486380616,ok
+luminal_compiled,2,512,16,2,4096,0.10465599969029427,4892218.32971973,ok
+pt_eager,2,1024,16,2,4096,0.5781759917736053,1771087.0298484561,ok
+graph_safe_inductor_cg,2,1024,16,2,4096,0.029888000339269638,34261241.58110951,ok
+luminal_compiled,2,1024,16,2,4096,0.098191998898983,10428548.267496424,ok
+pt_eager,2,2048,16,2,4096,0.5551839768886566,3688867.2678871835,ok
+graph_safe_inductor_cg,2,2048,16,2,4096,0.03388800099492073,60434370.27480501,ok
+luminal_compiled,2,2048,16,2,4096,0.10235200077295303,20009379.245483134,ok
+pt_eager,4,2,16,2,4096,0.5139839947223663,3891.171749580102,ok
+graph_safe_inductor_cg,4,2,16,2,4096,0.02729600016027689,73270.8084794982,ok
+luminal_compiled,4,2,16,2,4096,0.09959999844431877,20080.32159878092,ok
+pt_eager,4,4,16,2,4096,0.512287974357605,7808.10833011626,ok
+graph_safe_inductor_cg,4,4,16,2,4096,0.027664000168442726,144592.24897500317,ok
+luminal_compiled,4,4,16,2,4096,0.10916800051927567,36640.77367885587,ok
+pt_eager,4,8,16,2,4096,0.5412319898605347,14781.092303988627,ok
+graph_safe_inductor_cg,4,8,16,2,4096,0.02937600016593933,272331.15314574994,ok
+luminal_compiled,4,8,16,2,4096,0.10155199840664864,78777.37637387782,ok
+pt_eager,4,16,16,2,4096,0.48787200450897217,32795.48703784202,ok
+graph_safe_inductor_cg,4,16,16,2,4096,0.028704000636935234,557413.5885229809,ok
+luminal_compiled,4,16,16,2,4096,0.09879999980330467,161943.32016046048,ok
+pt_eager,4,32,16,2,4096,0.5792959928512573,55239.46375409585,ok
+graph_safe_inductor_cg,4,32,16,2,4096,0.036240000277757645,883002.2007378418,ok
+luminal_compiled,4,32,16,2,4096,0.10119999945163727,316205.5353102305,ok
+pt_eager,4,64,16,2,4096,0.5686399936676025,112549.24154598081,ok
+graph_safe_inductor_cg,4,64,16,2,4096,0.03144000098109245,2035623.3461471149,ok
+luminal_compiled,4,64,16,2,4096,0.1101440005004406,581057.5220549029,ok
+pt_eager,4,128,16,2,4096,0.5731199979782104,223338.9175941239,ok
+graph_safe_inductor_cg,4,128,16,2,4096,0.03721600025892258,3439380.887507165,ok
+luminal_compiled,4,128,16,2,4096,0.10672000050544739,1199400.294169474,ok
+pt_eager,4,256,16,2,4096,0.5351200103759766,478397.3595383469,ok
+graph_safe_inductor_cg,4,256,16,2,4096,0.0326399989426136,7843137.508983668,ok
+luminal_compiled,4,256,16,2,4096,0.10263999924063683,2494154.344251451,ok
+pt_eager,4,512,16,2,4096,0.5590240061283112,915881.9556712956,ok
+graph_safe_inductor_cg,4,512,16,2,4096,0.0337119996547699,15187470.492500357,ok
+luminal_compiled,4,512,16,2,4096,0.10249600186944008,4995316.799304895,ok
+pt_eager,4,1024,16,2,4096,0.6048000156879425,1693121.6492037117,ok
+graph_safe_inductor_cg,4,1024,16,2,4096,0.03484800085425377,29384755.937154543,ok
+luminal_compiled,4,1024,16,2,4096,0.10395199805498123,9850700.507539993,ok
+pt_eager,4,2048,16,2,4096,0.5883680284023285,3480814.5601677205,ok
+graph_safe_inductor_cg,4,2048,16,2,4096,0.04467200115323067,45845271.02278446,ok
+luminal_compiled,4,2048,16,2,4096,0.11158400028944016,18353885.814163756,ok
+pt_eager,8,2,16,2,4096,0.6239520013332367,3205.3747655692696,ok
+graph_safe_inductor_cg,8,2,16,2,4096,0.04572800174355507,43736.877268683216,ok
+luminal_compiled,8,2,16,2,4096,0.10910400003194809,18331.13359193389,ok
+pt_eager,8,4,16,2,4096,0.6168799996376038,6484.243292617471,ok
+graph_safe_inductor_cg,8,4,16,2,4096,0.045903999358415604,87138.37695857928,ok
+luminal_compiled,8,4,16,2,4096,0.10825600102543831,36949.4527980954,ok
+pt_eager,8,8,16,2,4096,0.6259680092334747,12780.205828403836,ok
+graph_safe_inductor_cg,8,8,16,2,4096,0.04761600121855736,168010.74838855147,ok
+luminal_compiled,8,8,16,2,4096,0.11219200119376183,71306.33124355768,ok
+pt_eager,8,16,16,2,4096,0.6317119896411896,25327.997983840625,ok
+graph_safe_inductor_cg,8,16,16,2,4096,0.047807998955249786,334672.02873261116,ok
+luminal_compiled,8,16,16,2,4096,0.10617600008845329,150693.18854233247,ok
+pt_eager,8,32,16,2,4096,0.6701280176639557,47752.0699873898,ok
+graph_safe_inductor_cg,8,32,16,2,4096,0.0514880008995533,621504.0289178838,ok
+luminal_compiled,8,32,16,2,4096,0.11286400258541107,283527.0703409942,ok
+pt_eager,8,64,16,2,4096,0.7092800140380859,90232.34651098376,ok
+graph_safe_inductor_cg,8,64,16,2,4096,0.052239999175071716,1225114.873863551,ok
+luminal_compiled,8,64,16,2,4096,0.11633599922060966,550130.659716395,ok
+pt_eager,8,128,16,2,4096,0.6788640022277832,188550.28338511224,ok
+graph_safe_inductor_cg,8,128,16,2,4096,0.053888000547885895,2375296.8879640796,ok
+luminal_compiled,8,128,16,2,4096,0.11127999797463417,1150251.6384766388,ok
+pt_eager,8,256,16,2,4096,0.6391039788722992,400560.7983410035,ok
+graph_safe_inductor_cg,8,256,16,2,4096,0.05567999929189682,4597701.207895956,ok
+luminal_compiled,8,256,16,2,4096,0.10991999879479408,2328966.5466419603,ok
+pt_eager,8,512,16,2,4096,0.6350559890270233,806228.1260971039,ok
+graph_safe_inductor_cg,8,512,16,2,4096,0.0562559999525547,9101251.429746367,ok
+luminal_compiled,8,512,16,2,4096,0.12015999853610992,4260985.404773753,ok
+pt_eager,8,1024,16,2,4096,0.6764000058174133,1513897.089285977,ok
+graph_safe_inductor_cg,8,1024,16,2,4096,0.05718399956822395,17907107.018254407,ok
+luminal_compiled,8,1024,16,2,4096,0.11070400103926659,9249891.515996683,ok
+pt_eager,8,2048,16,2,4096,0.6043839752674103,3388574.2901999685,ok
+graph_safe_inductor_cg,8,2048,16,2,4096,0.06619199737906456,30940296.12479633,ok
+luminal_compiled,8,2048,16,2,4096,0.12067200243473053,16971625.221084144,ok
+pt_eager,16,2,16,2,4096,0.7858880162239075,2544.891840455523,ok
+graph_safe_inductor_cg,16,2,16,2,4096,0.07545600086450577,26505.512843058623,ok
+luminal_compiled,16,2,16,2,4096,0.13145600259304047,15214.215863474643,ok
+pt_eager,16,4,16,2,4096,0.7631199955940247,5241.639615125452,ok
+graph_safe_inductor_cg,16,4,16,2,4096,0.07680000364780426,52083.330859507856,ok
+luminal_compiled,16,4,16,2,4096,0.12828800082206726,31179.84514816717,ok
+pt_eager,16,8,16,2,4096,0.7696959972381592,10393.713919139223,ok
+graph_safe_inductor_cg,16,8,16,2,4096,0.08003199845552444,99960.0179226535,ok
+luminal_compiled,16,8,16,2,4096,0.1287200003862381,62150.40379113693,ok
+pt_eager,16,16,16,2,4096,0.7825759947299957,20445.29874126834,ok
+graph_safe_inductor_cg,16,16,16,2,4096,0.07948800176382065,201288.24029996534,ok
+luminal_compiled,16,16,16,2,4096,0.13308800011873245,120221.20691366495,ok
+pt_eager,16,32,16,2,4096,0.8636959791183472,37050.074069657356,ok
+graph_safe_inductor_cg,16,32,16,2,4096,0.08278399705886841,386548.1389747891,ok
+luminal_compiled,16,32,16,2,4096,0.13150399923324585,243338.6070886124,ok
+pt_eager,16,64,16,2,4096,0.878896027803421,72818.62470120825,ok
+graph_safe_inductor_cg,16,64,16,2,4096,0.09142400324344635,700034.9769149688,ok
+luminal_compiled,16,64,16,2,4096,0.13363199681043625,478927.2144962949,ok
+pt_eager,16,128,16,2,4096,0.8490720093250275,150752.82024872545,ok
+graph_safe_inductor_cg,16,128,16,2,4096,0.0907679982483387,1410188.6399410903,ok
+luminal_compiled,16,128,16,2,4096,0.1300320029258728,984373.0552467828,ok
+pt_eager,16,256,16,2,4096,0.812527984380722,315066.07147212717,ok
+graph_safe_inductor_cg,16,256,16,2,4096,0.09612800180912018,2663115.795419685,ok
+luminal_compiled,16,256,16,2,4096,0.13608000427484512,1881246.2665929128,ok
+pt_eager,16,512,16,2,4096,0.8759520053863525,584506.9100266221,ok
+graph_safe_inductor_cg,16,512,16,2,4096,0.09095999971032143,5628847.863132768,ok
+luminal_compiled,16,512,16,2,4096,0.13556800037622452,3776702.4561777995,ok
+pt_eager,16,1024,16,2,4096,0.8617600202560425,1188265.8465587131,ok
+graph_safe_inductor_cg,16,1024,16,2,4096,0.09561599791049957,10709504.919443557,ok
+luminal_compiled,16,1024,16,2,4096,0.14022399485111237,7302601.819947201,ok
+pt_eager,16,2048,16,2,4096,0.8651839792728424,2367126.587019415,ok
+graph_safe_inductor_cg,16,2048,16,2,4096,0.12144000083208084,16864295.00961416,ok
+luminal_compiled,16,2048,16,2,4096,0.13519999384880066,15147929.68326875,ok
+pt_eager,32,2,16,2,4096,1.1254720091819763,1777.0321995423542,ok
+graph_safe_inductor_cg,32,2,16,2,4096,0.13230399787425995,15116.70117407015,ok
+luminal_compiled,32,2,16,2,4096,0.1716800034046173,11649.58038407291,ok
+pt_eager,32,4,16,2,4096,1.1279360055923462,3546.300481736428,ok
+graph_safe_inductor_cg,32,4,16,2,4096,0.14241600036621094,28086.731755661804,ok
+luminal_compiled,32,4,16,2,4096,0.17348799854516983,23056.34991205774,ok
+pt_eager,32,8,16,2,4096,1.0751680135726929,7440.69754588092,ok
+graph_safe_inductor_cg,32,8,16,2,4096,0.1366880014538765,58527.4487512314,ok
+luminal_compiled,32,8,16,2,4096,0.16710400581359863,47874.37596752677,ok
+pt_eager,32,16,16,2,4096,1.1331039667129517,14120.504799232836,ok
+graph_safe_inductor_cg,32,16,16,2,4096,0.13777600228786469,116130.52878809854,ok
+luminal_compiled,32,16,16,2,4096,0.17158400267362595,93248.78631275425,ok
+pt_eager,32,32,16,2,4096,1.1613919734954834,27553.143753601504,ok
+graph_safe_inductor_cg,32,32,16,2,4096,0.14433600008487701,221704.91063339947,ok
+luminal_compiled,32,32,16,2,4096,0.17432000488042831,183570.44001891708,ok
+pt_eager,32,64,16,2,4096,1.1802719831466675,54224.789636514644,ok
+graph_safe_inductor_cg,32,64,16,2,4096,0.16435199975967407,389408.10025789076,ok
+luminal_compiled,32,64,16,2,4096,0.17791999876499176,359712.2327127224,ok
+pt_eager,32,128,16,2,4096,1.1843199729919434,108078.9000599509,ok
+graph_safe_inductor_cg,32,128,16,2,4096,0.15452799946069717,828328.8494429494,ok
+luminal_compiled,32,128,16,2,4096,0.1780960038304329,718713.4873720702,ok
+pt_eager,32,256,16,2,4096,1.168511986732483,219082.0487138128,ok
+graph_safe_inductor_cg,32,256,16,2,4096,0.15727999806404114,1627670.4167796476,ok
+luminal_compiled,32,256,16,2,4096,0.17375999689102173,1473296.5272815775,ok
+pt_eager,32,512,16,2,4096,1.0871520042419434,470955.30156062293,ok
+graph_safe_inductor_cg,32,512,16,2,4096,0.16273599863052368,3146200.0068125455,ok
+luminal_compiled,32,512,16,2,4096,0.1789119988679886,2861742.10360135,ok
+pt_eager,32,1024,16,2,4096,1.1761599779129028,870629.8626289675,ok
+graph_safe_inductor_cg,32,1024,16,2,4096,0.16993600130081177,6025797.901336805,ok
+luminal_compiled,32,1024,16,2,4096,0.17558399587869644,5831966.603080603,ok
+pt_eager,32,2048,16,2,4096,1.2334399819374084,1660396.962958127,ok
+graph_safe_inductor_cg,32,2048,16,2,4096,0.24084799736738205,8503288.473999824,ok
+luminal_compiled,32,2048,16,2,4096,0.17404799908399582,11766868.971654378,ok

examples/dlrm/src/bin/check.rs

@@ -0,0 +1,255 @@
+//! Numerical correctness check: load weights and inputs dumped by
+//! `correctness_dump.py`, run the luminal DLRM forward, and compare
+//! element-wise against PyTorch's expected output.
+//!
+//! Usage: `cargo run --release --bin check -- [weights/]`.
+
+use luminal::prelude::*;
+use luminal_cuda_lite::kernel::{
+    dlrm_pairwise_dot_lower_tri_stacked, linear_bias, linear_bias_relu, linear_bias_relu_split_a,
+    linear_bias_sigmoid, stacked_embedding_bag_sum_kernel,
+};
+use luminal_cuda_lite::{cudarc::driver::CudaContext, runtime::CudaRuntime};
+use rand::{SeedableRng, rngs::StdRng};
+use std::fs;
+use std::path::{Path, PathBuf};
+
+const BATCH: usize = 2048;
+const M_DEN: usize = 3;
+const M_SPA: usize = 16;
+const L: usize = 2;
+const LN_BOT: &[usize] = &[M_DEN, 64, M_SPA];
+const LN_TOP_TAIL: &[usize] = &[64, 32, 1];
+
+#[derive(Clone, Copy)]
+enum Act {
+    None,
+    Relu,
+    Sigmoid,
+}
+
+struct LinearWB {
+    w: GraphTensor,
+    b: GraphTensor,
+    act: Act,
+}
+
+impl LinearWB {
+    fn new(cx: &mut Graph, in_dim: usize, out_dim: usize, name: &str, act: Act) -> Self {
+        Self {
+            w: cx
+                .named_tensor(format!("{name}_w").as_str(), (out_dim, in_dim))
+                .persist(),
+            b: cx
+                .named_tensor(format!("{name}_b").as_str(), out_dim)
+                .persist(),
+            act,
+        }
+    }
+
+    fn forward(&self, x: GraphTensor) -> GraphTensor {
+        match self.act {
+            Act::None => linear_bias(x, self.w, self.b),
+            Act::Relu => linear_bias_relu(x, self.w, self.b),
+            Act::Sigmoid => linear_bias_sigmoid(x, self.w, self.b),
+        }
+    }
+}
+
+fn read_f32(path: &Path) -> Vec<f32> {
+    let bytes = fs::read(path)
+        .unwrap_or_else(|e| panic!("can't read {}: {e}", path.display()));
+    assert_eq!(bytes.len() % 4, 0, "{} not a multiple of 4 bytes", path.display());
+    bytemuck::cast_slice::<u8, f32>(&bytes).to_vec()
+}
+
+fn read_i32(path: &Path) -> Vec<i32> {
+    let bytes = fs::read(path)
+        .unwrap_or_else(|e| panic!("can't read {}: {e}", path.display()));
+    assert_eq!(bytes.len() % 4, 0, "{} not a multiple of 4 bytes", path.display());
+    bytemuck::cast_slice::<u8, i32>(&bytes).to_vec()
+}
+
+fn main() {
+    let weights_dir: PathBuf = std::env::args()
+        .nth(1)
+        .map(PathBuf::from)
+        .unwrap_or_else(|| PathBuf::from("examples/dlrm/weights"));
+    let manifest_path = weights_dir.join("manifest.json");
+    let manifest_bytes = fs::read(&manifest_path)
+        .unwrap_or_else(|e| panic!("can't read {}: {e}", manifest_path.display()));
+    let manifest_text = std::str::from_utf8(&manifest_bytes).unwrap();
+    // Extract num_cat and rows from the manifest with a tiny parse — we only
+    // need two integers, so avoid pulling in serde_json.
+    let extract = |key: &str| -> usize {
+        let needle = format!("\"{key}\":");
+        let i = manifest_text.find(&needle).expect("key not found");
+        let rest = &manifest_text[i + needle.len()..];
+        let s: String = rest
+            .chars()
+            .skip_while(|c| !c.is_ascii_digit())
+            .take_while(|c| c.is_ascii_digit())
+            .collect();
+        s.parse().unwrap()
+    };
+    let num_cat = extract("num_cat");
+    let rows = extract("rows");
+    let num_fea = num_cat + 1;
+    let pair_count = num_fea * (num_fea - 1) / 2;
+    let top_in = pair_count + M_SPA;
+    let ln_top: Vec<usize> = std::iter::once(top_in)
+        .chain(LN_TOP_TAIL.iter().copied())
+        .collect();
+    println!(
+        "check: num_cat={num_cat} rows={rows} F={num_fea} pairs={pair_count} top_in={top_in}"
+    );
+    println!("  ln_top={ln_top:?}");
+
+    let ctx = CudaContext::new(0).unwrap();
+    let stream = ctx.default_stream();
+    let mut cx = Graph::default();
+
+    let dense = cx.named_tensor("dense", (BATCH, M_DEN)).persist();
+    let stacked_w = cx
+        .named_tensor("emb_stacked", (num_cat * rows, M_SPA))
+        .persist();
+    let mut sparse_inputs = Vec::with_capacity(num_cat);
+    for i in 0..num_cat {
+        sparse_inputs.push(
+            cx.named_tensor(format!("idx_{i}").as_str(), (BATCH, L))
+                .as_dtype(DType::Int)
+                .persist(),
+        );
+    }
+
+    // Upstream DLRM applies ReLU to every bot layer (sigmoid_bot=-1 default).
+    let mut bot_layers = Vec::new();
+    for (i, win) in LN_BOT.windows(2).enumerate() {
+        bot_layers.push(LinearWB::new(&mut cx, win[0], win[1], &format!("bot_{i}"), Act::Relu));
+    }
+    let mut h = dense;
+    for l in bot_layers.iter() {
+        h = l.forward(h);
+    }
+    let dense_out = h;
+
+    let row_offsets: Vec<usize> = (0..=num_cat).map(|k| k * rows).collect();
+    let emb_stack =
+        stacked_embedding_bag_sum_kernel(stacked_w, sparse_inputs.clone(), &row_offsets);
+    let interactions = dlrm_pairwise_dot_lower_tri_stacked(dense_out, emb_stack);
+
+    let mut top_layers = Vec::new();
+    let mut prev = top_in;
+    let last_top = LN_TOP_TAIL.len() - 1;
+    for (i, &h) in LN_TOP_TAIL.iter().enumerate() {
+        let act = if i < last_top { Act::Relu } else { Act::Sigmoid };
+        top_layers.push(LinearWB::new(&mut cx, prev, h, &format!("top_{i}"), act));
+        prev = h;
+    }
+    // top_0 reads `dense_out` and `interactions` directly via the
+    // split-A kernel — no materialized concat.
+    let mut t = linear_bias_relu_split_a(
+        dense_out,
+        interactions,
+        top_layers[0].w,
+        top_layers[0].b,
+    );
+    for l in top_layers.iter().skip(1) {
+        t = l.forward(t);
+    }
+    let out_t = t.output();
+
+    cx.build_search_space::<CudaRuntime>();
+    let mut runtime = CudaRuntime::initialize(stream);
+
+    // Load weights/biases from disk.
+    for (i, _win) in LN_BOT.windows(2).enumerate() {
+        let l = &bot_layers[i];
+        let w = read_f32(&weights_dir.join(format!("bot_{i}_w.bin")));
+        let b = read_f32(&weights_dir.join(format!("bot_{i}_b.bin")));
+        runtime.set_data(l.w, w);
+        runtime.set_data(l.b, b);
+    }
+    for i in 0..top_layers.len() {
+        let l = &top_layers[i];
+        let w = read_f32(&weights_dir.join(format!("top_{i}_w.bin")));
+        let b = read_f32(&weights_dir.join(format!("top_{i}_b.bin")));
+        runtime.set_data(l.w, w);
+        runtime.set_data(l.b, b);
+    }
+    // Read per-table weight files (as PyTorch dumped them) and concat them
+    // into the single stacked weight that the fused kernel expects.
+    let mut stacked = Vec::with_capacity(num_cat * rows * M_SPA);
+    for i in 0..num_cat {
+        let t = read_f32(&weights_dir.join(format!("emb_{i}.bin")));
+        assert_eq!(t.len(), rows * M_SPA, "emb_{i}.bin shape mismatch");
+        stacked.extend(t);
+    }
+    runtime.set_data(stacked_w, stacked);
+    let dense_data = read_f32(&weights_dir.join("dense.bin"));
+    runtime.set_data(dense, dense_data);
+    for i in 0..num_cat {
+        let idx = read_i32(&weights_dir.join(format!("idx_{i}.bin")));
+        runtime.set_data(sparse_inputs[i], idx);
+    }
+
+    let mut rng = StdRng::seed_from_u64(0);
+    runtime = cx.search_options(runtime, SearchOptions::new(50).trials(1).keep_best(2), &mut rng);
+
+    runtime.execute(&cx.dyn_map);
+    let lum_out = runtime.get_f32(out_t);
+    let expected = read_f32(&weights_dir.join("expected.bin"));
+    assert_eq!(
+        lum_out.len(),
+        expected.len(),
+        "output length mismatch: luminal={} expected={}",
+        lum_out.len(),
+        expected.len()
+    );
+
+    let mut max_abs = 0.0f32;
+    let mut sum_abs = 0.0f64;
+    let mut max_rel = 0.0f32;
+    let mut diff_count = 0usize;
+    for (i, (&a, &b)) in lum_out.iter().zip(expected.iter()).enumerate() {
+        let d = (a - b).abs();
+        sum_abs += d as f64;
+        if d > max_abs {
+            max_abs = d;
+        }
+        let r = if b.abs() > 1e-8 { d / b.abs() } else { d };
+        if r > max_rel {
+            max_rel = r;
+        }
+        if d > 1e-4 {
+            diff_count += 1;
+            if diff_count <= 5 {
+                println!(
+                    "  diff @ {i}: luminal={a:.6} expected={b:.6} abs={d:.3e}"
+                );
+            }
+        }
+    }
+    let mean_abs = sum_abs / lum_out.len() as f64;
+
+    println!();
+    println!(
+        "  luminal head:  {:?}",
+        &lum_out[..8.min(lum_out.len())]
+    );
+    println!(
+        "  expected head: {:?}",
+        &expected[..8.min(expected.len())]
+    );
+    println!(
+        "  max abs diff = {max_abs:.3e}  mean abs diff = {mean_abs:.3e}  max rel diff = {max_rel:.3e}"
+    );
+    println!("  elements with abs diff > 1e-4: {diff_count}/{}", lum_out.len());
+    let tol: f32 = 1e-3;
+    if max_abs < tol {
+        println!("PASS (max abs diff < {tol})");
+    } else {
+        println!("FAIL (max abs diff >= {tol})");
+        std::process::exit(1);
+    }
+}

examples/dlrm/src/main.rs

@@ -0,0 +1,415 @@
+//! DLRM forward-pass benchmark on luminal's CUDA backend.
+//!
+//! Mirrors `sweep_categories.py` from https://github.com/jss8649/tmp-dlrm-bench:
+//!   batch=2048, m_spa=16, indices_per_bag=2, rows_per_table=4096,
+//!   ln_bot=[3, 64, 16], top MLP scales with F = num_cat + 1:
+//!     num_int = F*(F-1)/2 + m_spa, ln_top = [num_int, 64, 32, 1].
+//!   arch_interaction_op="dot", arch_interaction_itself=False, sigmoid_top=2.
+//!
+//! CLI: `dlrm [--num-cat N] [--rows R]`. Defaults: num-cat=3, rows=4096.
+//!
+//! Harness: 5 rounds × 20 timed iters, 10 warmup before round 0,
+//! median-of-round-medians (same as bench_luminal_exact.py).
+
+use luminal::prelude::*;
+use luminal_cuda_lite::kernel::{
+    dlrm_pairwise_dot_lower_tri_stacked, linear_bias, linear_bias_relu, linear_bias_relu_split_a,
+    linear_bias_sigmoid, stacked_embedding_bag_sum_kernel,
+};
+use luminal_cuda_lite::{cudarc::driver::CudaContext, runtime::CudaRuntime};
+use rand::{Rng, SeedableRng, rngs::StdRng};
+use std::time::Instant;
+
+// ---- Fixed config (matches sweep_categories.py defaults; per-call
+//      overridable from the CLI for sweep purposes) ----
+const M_DEN: usize = 3;
+const LN_TOP_TAIL: &[usize] = &[64, 32, 1]; // → ln_top = [num_int, 64, 32, 1]
+
+const SEARCH_GRAPHS: usize = 200;
+const SEARCH_TRIALS: usize = 1;
+const SEARCH_KEEP_BEST: usize = 4;
+const SEARCH_SEED: u64 = 0;
+
+const PRE_ROUND_WARMUP: usize = 10;
+const TIMED_ITERS_PER_ROUND: usize = 20;
+const ROUNDS: usize = 5;
+
+#[derive(Clone, Copy)]
+enum Act {
+    None,
+    Relu,
+    Sigmoid,
+}
+
+struct LinearWB {
+    w: GraphTensor, // (out_dim, in_dim) — PyTorch/cuBLASLt convention
+    b: GraphTensor, // (out_dim,)
+    act: Act,
+}
+
+impl LinearWB {
+    fn new(cx: &mut Graph, in_dim: usize, out_dim: usize, name: &str, act: Act) -> Self {
+        Self {
+            w: cx
+                .named_tensor(format!("{name}_w").as_str(), (out_dim, in_dim))
+                .persist(),
+            b: cx
+                .named_tensor(format!("{name}_b").as_str(), out_dim)
+                .persist(),
+            act,
+        }
+    }
+
+    fn forward(&self, x: GraphTensor) -> GraphTensor {
+        match self.act {
+            Act::None => linear_bias(x, self.w, self.b),
+            Act::Relu => linear_bias_relu(x, self.w, self.b),
+            Act::Sigmoid => linear_bias_sigmoid(x, self.w, self.b),
+        }
+    }
+}
+
+fn build_indices(table_idx: usize, batch: usize, bag: usize, rows: usize) -> Vec<i32> {
+    // Match sweep_categories.py's `_make_inputs`:
+    //   positions = arange(batch * L)
+    //   lS_i[k] = (positions * (2k+3) + (k+1)) % ROWS_PER_TABLE
+    let n = rows as i64;
+    let s = (2 * table_idx + 3) as i64;
+    let o = (table_idx + 1) as i64;
+    (0..(batch * bag) as i64)
+        .map(|p| ((p * s + o).rem_euclid(n)) as i32)
+        .collect()
+}
+
+fn rand_normal(rng: &mut StdRng, n: usize, std: f32) -> Vec<f32> {
+    let mut out = Vec::with_capacity(n);
+    while out.len() < n {
+        let u1: f32 = rng.random::<f32>().max(1e-9);
+        let u2: f32 = rng.random::<f32>();
+        let r = (-2.0 * u1.ln()).sqrt() * std;
+        let z0 = r * (2.0 * std::f32::consts::PI * u2).cos();
+        let z1 = r * (2.0 * std::f32::consts::PI * u2).sin();
+        out.push(z0);
+        if out.len() < n {
+            out.push(z1);
+        }
+    }
+    out
+}
+
+fn rand_uniform(rng: &mut StdRng, n: usize, hi: f32) -> Vec<f32> {
+    (0..n).map(|_| (rng.random::<f32>() * 2.0 - 1.0) * hi).collect()
+}
+
+fn build_dense_x(batch: usize, m_den: usize) -> Vec<f32> {
+    let total = batch * m_den;
+    (0..total)
+        .map(|i| -1.0 + 2.0 * (i as f32) / ((total - 1) as f32))
+        .collect()
+}
+
+struct Args {
+    num_cat: usize,
+    rows: usize,
+    batch: usize,
+    m_spa: usize,
+    bag: usize,
+    print_outputs: bool,
+}
+
+fn parse_args() -> Args {
+    let mut a = Args {
+        num_cat: 3,
+        rows: 4096,
+        batch: 2048,
+        m_spa: 16,
+        bag: 2,
+        print_outputs: false,
+    };
+    let mut args = std::env::args().skip(1);
+    while let Some(arg) = args.next() {
+        match arg.as_str() {
+            "--num-cat" => {
+                a.num_cat = args.next().expect("missing value for --num-cat").parse().unwrap();
+            }
+            "--rows" => {
+                a.rows = args.next().expect("missing value for --rows").parse().unwrap();
+            }
+            "--batch" => {
+                a.batch = args.next().expect("missing value for --batch").parse().unwrap();
+            }
+            "--m-spa" => {
+                a.m_spa = args.next().expect("missing value for --m-spa").parse().unwrap();
+            }
+            "--bag" => {
+                a.bag = args.next().expect("missing value for --bag").parse().unwrap();
+            }
+            "--print-outputs" => {
+                a.print_outputs = true;
+            }
+            "-h" | "--help" => {
+                eprintln!(
+                    "usage: dlrm [--num-cat N] [--rows R] [--batch B] [--m-spa D] [--bag L] [--print-outputs]"
+                );
+                std::process::exit(0);
+            }
+            other => {
+                eprintln!("unknown arg: {other}");
+                std::process::exit(2);
+            }
+        }
+    }
+    a
+}
+
+fn main() {
+    let args = parse_args();
+    let Args {
+        num_cat,
+        rows,
+        batch,
+        m_spa,
+        bag,
+        print_outputs,
+    } = args;
+    let num_fea = num_cat + 1;
+    let pair_count = num_fea * (num_fea - 1) / 2;
+    let top_in = pair_count + m_spa;
+    // ln_bot last entry must equal m_spa (bot output feeds into interaction
+    // alongside emb rows that are m_spa wide).
+    let ln_bot: Vec<usize> = vec![M_DEN, 64, m_spa];
+    let ln_top: Vec<usize> = std::iter::once(top_in)
+        .chain(LN_TOP_TAIL.iter().copied())
+        .collect();
+
+    println!(
+        "==== DLRM luminal config ====  num_cat={num_cat} F={num_fea} \
+         rows/table={rows} batch={batch} m_spa={m_spa} L={bag}"
+    );
+    println!("  ln_bot={ln_bot:?}  ln_top={ln_top:?}  pairs={pair_count}");
+
+    let ctx = CudaContext::new(0).unwrap();
+    let stream = ctx.default_stream();
+
+    let mut cx = Graph::default();
+    let dense = cx
+        .named_tensor("dense", (batch, M_DEN))
+        .persist();
+    let stacked_w = cx
+        .named_tensor("emb_stacked", (num_cat * rows, m_spa))
+        .persist();
+    let mut sparse_inputs = Vec::with_capacity(num_cat);
+    for i in 0..num_cat {
+        sparse_inputs.push(
+            cx.named_tensor(format!("idx_{i}").as_str(), (batch, bag))
+                .as_dtype(DType::Int)
+                .persist(),
+        );
+    }
+
+    // Upstream `dlrm_s_pytorch.DLRM_Net.create_mlp` with `sigmoid_bot=-1`
+    // (the default) applies ReLU to every bot layer — including the final
+    // one. The output of the bot MLP feeds the interaction op AFTER ReLU.
+    let mut bot_layers = Vec::new();
+    for (i, win) in ln_bot.windows(2).enumerate() {
+        bot_layers.push(LinearWB::new(&mut cx, win[0], win[1], &format!("bot_{i}"), Act::Relu));
+    }
+    let mut h = dense;
+    for l in bot_layers.iter() {
+        h = l.forward(h);
+    }
+    let dense_out = h;
+
+    // One fused kernel for all num_cat tables. row_offsets[k] = k * rows
+    // (uniform rows-per-table in this benchmark, matching upstream).
+    let row_offsets: Vec<usize> = (0..=num_cat).map(|k| k * rows).collect();
+    let emb_stack =
+        stacked_embedding_bag_sum_kernel(stacked_w, sparse_inputs.clone(), &row_offsets);
+    // emb_stack: (BATCH, num_cat, M_SPA)
+
+    // Feature interaction over [dense_out, emb_stack[:, 0, :], …,
+    // emb_stack[:, num_cat-1, :]] in one fused kernel — no per-table slice.
+    let interactions = dlrm_pairwise_dot_lower_tri_stacked(dense_out, emb_stack);
+
+    // Skip the materialized `cat(dense_out, interactions)`: the first top
+    // layer reads both halves directly via the split-A matmul kernel. The
+    // remaining top layers see the dense fused output and stay vanilla.
+    let mut top_layers = Vec::new();
+    let mut prev = top_in;
+    let last_top = LN_TOP_TAIL.len() - 1;
+    for (i, &h) in LN_TOP_TAIL.iter().enumerate() {
+        let act = if i < last_top {
+            Act::Relu
+        } else {
+            Act::Sigmoid
+        };
+        top_layers.push(LinearWB::new(&mut cx, prev, h, &format!("top_{i}"), act));
+        prev = h;
+    }
+    // top_0: split-A matmul (no concat materialization).
+    let mut t = linear_bias_relu_split_a(
+        dense_out,
+        interactions,
+        top_layers[0].w,
+        top_layers[0].b,
+    );
+    // top_1..: standard fused linear_bias_(relu|sigmoid) kernels.
+    for l in top_layers.iter().skip(1) {
+        t = l.forward(t);
+    }
+    let out = t.output();
+
+    println!("Building E-graph...");
+    let t = Instant::now();
+    cx.build_search_space::<CudaRuntime>();
+    println!("  build: {:.2}s", t.elapsed().as_secs_f64());
+
+    let mut runtime = CudaRuntime::initialize(stream);
+    let mut rng = StdRng::seed_from_u64(SEARCH_SEED);
+
+    for (i, win) in ln_bot.windows(2).enumerate() {
+        let (a, b) = (win[0], win[1]);
+        let l = &bot_layers[i];
+        runtime.set_data(
+            l.w,
+            rand_normal(&mut rng, a * b, (2.0 / (a + b) as f32).sqrt()),
+        );
+        runtime.set_data(l.b, rand_normal(&mut rng, b, (1.0 / b as f32).sqrt()));
+    }
+    let mut top_shapes: Vec<(usize, usize)> = vec![];
+    let mut prev = top_in;
+    for &h in LN_TOP_TAIL.iter() {
+        top_shapes.push((prev, h));
+        prev = h;
+    }
+    for (i, &(a, b)) in top_shapes.iter().enumerate() {
+        let l = &top_layers[i];
+        runtime.set_data(
+            l.w,
+            rand_normal(&mut rng, a * b, (2.0 / (a + b) as f32).sqrt()),
+        );
+        runtime.set_data(l.b, rand_normal(&mut rng, b, (1.0 / b as f32).sqrt()));
+    }
+
+    // One stacked weight: concatenated per-table uniform inits.
+    let mut stacked_data: Vec<f32> = Vec::with_capacity(num_cat * rows * m_spa);
+    for _ in 0..num_cat {
+        stacked_data
+            .extend(rand_uniform(&mut rng, rows * m_spa, 1.0 / (rows as f32).sqrt()));
+    }
+    runtime.set_data(stacked_w, stacked_data);
+
+    runtime.set_data(dense, build_dense_x(batch, M_DEN));
+    for i in 0..num_cat {
+        runtime.set_data(sparse_inputs[i], build_indices(i, batch, bag, rows));
+    }
+
+    println!("Searching/compiling...");
+    let t = Instant::now();
+    runtime = cx.search_options(
+        runtime,
+        SearchOptions::new(SEARCH_GRAPHS)
+            .trials(SEARCH_TRIALS)
+            .keep_best(SEARCH_KEEP_BEST),
+        &mut rng,
+    );
+    println!("  search/compile: {:.2}s", t.elapsed().as_secs_f64());
+
+    {
+        let host_ops = runtime.host_ops();
+        let total = host_ops.len();
+        let cublaslt = host_ops
+            .iter()
+            .filter(|op| format!("{op:?}").contains("CuBlasLt"))
+            .count();
+        let cudagraph = host_ops
+            .iter()
+            .filter(|op| format!("{op:?}").contains("CudaGraph"))
+            .count();
+        println!("Host ops: total={total}  cuBLASLt={cublaslt}  CudaGraph={cudagraph}");
+    }
+
+    for _ in 0..PRE_ROUND_WARMUP {
+        runtime.execute(&cx.dyn_map);
+    }
+    let _ = runtime.get_f32(out);
+
+    let mut round_medians = Vec::with_capacity(ROUNDS);
+    for round in 0..ROUNDS {
+        let mut times_us = Vec::with_capacity(TIMED_ITERS_PER_ROUND);
+        for _ in 0..TIMED_ITERS_PER_ROUND {
+            let t = Instant::now();
+            runtime.execute(&cx.dyn_map);
+            // execute() no longer syncs at end (so it stays capturable by
+            // torch.cuda.CUDAGraph) — sync explicitly for accurate GPU time.
+            runtime.synchronize_stream();
+            times_us.push(t.elapsed().as_micros() as f64);
+        }
+        let _ = runtime.get_f32(out);
+        times_us.sort_by(|a, b| a.partial_cmp(b).unwrap());
+        let median = times_us[times_us.len() / 2] / 1000.0;
+        println!(
+            "  round {round}: median {median:.3} ms  min {:.3}  max {:.3}",
+            times_us[0] / 1000.0,
+            times_us[times_us.len() - 1] / 1000.0
+        );
+        round_medians.push(median);
+    }
+    round_medians.sort_by(|a, b| a.partial_cmp(b).unwrap());
+    let med_of_med = round_medians[round_medians.len() / 2];
+    let tput = batch as f64 / (med_of_med / 1000.0);
+    println!();
+    println!(
+        "==== cfg(num_cat={num_cat},batch={batch},m_spa={m_spa},bag={bag},rows={rows}): \
+         luminal median-of-round-medians = {med_of_med:.4} ms  ({tput:.0} samples/s)"
+    );
+    if print_outputs {
+        let sample = runtime.get_f32(out);
+        let head: Vec<f32> = sample.iter().take(8).copied().collect();
+        println!("  output[0..8] = {head:?}  (len={})", sample.len());
+    }
+
+    // Per-kernel breakdown. Available when `LUMINAL_KERNEL_TIMING=1` was
+    // set before the first execute (it gates the event-record-node
+    // insertion at graph-build time). `get_f32` above already
+    // synchronized the stream, so the events have valid data.
+    if std::env::var_os("LUMINAL_KERNEL_TIMING").is_some() {
+        let timings = runtime.read_per_kernel_timings_ms();
+        if timings.is_empty() {
+            println!("  (no per-kernel timings available — events not recorded)");
+        } else {
+            // Aggregate per kernel name in case the same op appears more
+            // than once (e.g. two Matmul2D_BiasRelu calls in the bot MLP).
+            let mut sum_per_name: std::collections::BTreeMap<&'static str, (f32, usize)> =
+                std::collections::BTreeMap::new();
+            let mut total = 0.0_f32;
+            for (name, ms) in &timings {
+                let e = sum_per_name.entry(*name).or_insert((0.0, 0));
+                e.0 += *ms;
+                e.1 += 1;
+                total += *ms;
+            }
+            println!();
+            println!("  per-kernel GPU time (single replay, ms):");
+            println!(
+                "  {:>40}  {:>3}  {:>10}  {:>10}  {:>5}",
+                "kernel", "n", "total_ms", "each_ms", "pct"
+            );
+            // Sort by total descending so the bottleneck is on top.
+            let mut sorted: Vec<_> = sum_per_name.iter().collect();
+            sorted.sort_by(|a, b| b.1.0.partial_cmp(&a.1.0).unwrap());
+            for (name, (sum, n)) in sorted {
+                let pct = if total > 0.0 { 100.0 * sum / total } else { 0.0 };
+                println!(
+                    "  {:>40}  {:>3}  {:>10.4}  {:>10.4}  {:>4.1}%",
+                    name,
+                    n,
+                    sum,
+                    sum / (*n as f32),
+                    pct
+                );
+            }
+            println!("  {:>40}      {:>10.4}", "TOTAL", total);
+        }
+    }
+}

examples/dlrm/summarize.py

@@ -0,0 +1,101 @@
+"""Summarize a results.csv from sweep_all.py: pretty per-cell tables and
+ratios of luminal_compiled vs the PT baselines.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import math
+import sys
+from pathlib import Path
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("csv", default="results.csv", nargs="?")
+    args = ap.parse_args()
+
+    p = Path(args.csv)
+    rows = list(csv.DictReader(p.open()))
+    # Index: (num_cat, batch, variant) -> ms
+    table: dict[tuple[int, int, str], float | None] = {}
+    nums_cat: set[int] = set()
+    batches: set[int] = set()
+    variants: set[str] = set()
+    for r in rows:
+        if r["status"] != "ok":
+            continue
+        nc = int(r["num_cat"])
+        bs = int(r["batch"])
+        v = r["variant"]
+        ms = float(r["ms"]) if r["ms"] else None
+        if ms is None:
+            continue
+        table[(nc, bs, v)] = ms
+        nums_cat.add(nc)
+        batches.add(bs)
+        variants.add(v)
+
+    nums_cat_s = sorted(nums_cat)
+    batches_s = sorted(batches)
+    variants_s = sorted(variants)
+
+    def fmt_ms(ms: float | None) -> str:
+        if ms is None:
+            return "    -    "
+        if ms < 1.0:
+            return f"{ms*1000:>5.0f} us"
+        return f"{ms:>6.3f} ms"
+
+    print(f"# DLRMv1 sweep — {len(rows)} rows, {len(variants_s)} variants")
+    print()
+    print(f"Configurations: batch ∈ {batches_s}, num_cat ∈ {nums_cat_s}")
+    print()
+
+    # Per-variant table.
+    for v in variants_s:
+        print(f"## {v}")
+        print()
+        hdr = "| nc \\ batch | " + " | ".join(f"{b:>6}" for b in batches_s) + " |"
+        print(hdr)
+        print("|" + "|".join(["---"] * (len(batches_s) + 1)) + "|")
+        for nc in nums_cat_s:
+            cells = [fmt_ms(table.get((nc, bs, v))) for bs in batches_s]
+            print(f"|       {nc:>3} | " + " | ".join(cells) + " |")
+        print()
+
+    # Speedup ratios: luminal_compiled vs each PT variant
+    if "luminal_compiled" in variants_s:
+        for ref in ["pt_eager", "graph_safe_inductor_cg", "graph_safe_cg", "v3_inductor_cg"]:
+            if ref not in variants_s:
+                continue
+            print(f"## speedup luminal_compiled vs {ref}  (>1 = luminal faster)")
+            print()
+            print("| nc \\ batch | " + " | ".join(f"{b:>6}" for b in batches_s) + " |")
+            print("|" + "|".join(["---"] * (len(batches_s) + 1)) + "|")
+            wins = 0
+            cells_total = 0
+            for nc in nums_cat_s:
+                rcells = []
+                for bs in batches_s:
+                    lum = table.get((nc, bs, "luminal_compiled"))
+                    pt = table.get((nc, bs, ref))
+                    if lum is None or pt is None or lum <= 0:
+                        rcells.append("   -  ")
+                        continue
+                    cells_total += 1
+                    r = pt / lum
+                    if r >= 1.0:
+                        wins += 1
+                    rcells.append(f"{r:>5.2f}x")
+                print(f"|       {nc:>3} | " + " | ".join(rcells) + " |")
+            print()
+            if cells_total:
+                print(f"luminal faster in {wins}/{cells_total} cells ({100*wins/cells_total:.0f}%)")
+                print()
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/dlrm/sweep_all.py

@@ -0,0 +1,115 @@
+"""Drive the full DLRM sweep across (batch, num_cat) and produce a CSV.
+
+For each cell, we shell out to `sweep_pytorch.py` (one process per
+variant, so torch.compile cache state is fresh) and to `sweep_luminal.py`
+(also separate, so the luminal compiled graph builds from clean state).
+
+Each cell × variant → one row in `results.csv`. Run from this dir:
+
+    python sweep_all.py [--variant ...] [--quick]
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+
+HERE = Path(__file__).parent
+PY = sys.executable
+
+
+def _powers(lo: int, hi: int) -> list[int]:
+    out = []
+    v = lo
+    while v <= hi:
+        out.append(v)
+        v *= 2
+    return out
+
+
+def _run_json(cmd: list[str]) -> dict | None:
+    try:
+        r = subprocess.run(cmd, cwd=HERE, capture_output=True, text=True, timeout=600)
+    except subprocess.TimeoutExpired:
+        print(f"  TIMEOUT: {' '.join(cmd)}", file=sys.stderr)
+        return None
+    if r.returncode != 0:
+        print(f"  ERR  : {' '.join(cmd)}\n  stderr tail: {r.stderr[-400:]}", file=sys.stderr)
+        return None
+    for line in r.stdout.splitlines()[::-1]:
+        s = line.strip()
+        if s.startswith("{") and s.endswith("}"):
+            try:
+                return json.loads(s)
+            except json.JSONDecodeError:
+                continue
+    return None
+
+
+PT_VARIANTS = ["pt_eager", "graph_safe_cg", "graph_safe_inductor_cg", "v3_inductor_cg"]
+LUMINAL_VARIANTS = ["luminal_compiled"]
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--batch", type=int, nargs="+", default=None)
+    ap.add_argument("--num-cat", type=int, nargs="+", default=None)
+    ap.add_argument("--variants", nargs="+", default=PT_VARIANTS + LUMINAL_VARIANTS)
+    ap.add_argument("--quick", action="store_true",
+                    help="Smaller sweep (nc∈{2,8,32}, batch∈{32,512,2048}) for fast iteration.")
+    ap.add_argument("--out", default="results.csv")
+    args = ap.parse_args()
+
+    if args.quick:
+        nums_cat = [2, 8, 32]
+        batches = [32, 512, 2048]
+    else:
+        nums_cat = args.num_cat or _powers(2, 32)
+        batches = args.batch or _powers(2, 2048)
+
+    rows: list[dict] = []
+    t0 = time.time()
+    total_cells = len(nums_cat) * len(batches) * len(args.variants)
+    seen = 0
+    for nc in nums_cat:
+        for bs in batches:
+            for variant in args.variants:
+                seen += 1
+                cmd_base = ["--num-cat", str(nc), "--batch", str(bs), "--json"]
+                if variant == "luminal_compiled":
+                    cmd = [PY, "sweep_luminal.py"] + cmd_base
+                else:
+                    cmd = [PY, "sweep_pytorch.py", "--variant", variant] + cmd_base
+                r = _run_json(cmd)
+                if r is None:
+                    print(f"[{seen}/{total_cells}] FAIL  nc={nc:>2} batch={bs:>4} variant={variant}", flush=True)
+                    rows.append({"variant": variant, "num_cat": nc, "batch": bs, "ms": None,
+                                 "samples_per_sec": None, "status": "failed"})
+                else:
+                    r["status"] = "ok"
+                    rows.append(r)
+                    print(f"[{seen}/{total_cells}] {r['ms']:8.4f} ms  nc={nc:>2} batch={bs:>4} variant={variant}", flush=True)
+
+    out_path = HERE / args.out
+    with open(out_path, "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=["variant", "num_cat", "batch", "m_spa", "bag", "rows", "ms",
+                                          "samples_per_sec", "status"])
+        w.writeheader()
+        for row in rows:
+            row = dict(row)
+            row.setdefault("m_spa", 16)
+            row.setdefault("bag", 2)
+            row.setdefault("rows", 4096)
+            row.setdefault("status", "ok")
+            w.writerow(row)
+    print(f"\nWrote {out_path}  ({len(rows)} rows, {time.time()-t0:.1f}s)")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/dlrm/sweep_luminal.py

@@ -0,0 +1,98 @@
+"""Time DLRMv1 under `torch.compile(model, backend=luminal_backend)`.
+
+Sibling to sweep_pytorch.py — same shape, same harness, same JSON output.
+
+Strategy: import DLRMv1/make_inputs from sweep_pytorch (so the model
+definition is the single source of truth across all variants), warm up,
+then time 5 rounds × 20 iters × 10 warmup as elsewhere.
+
+Outputs one JSON line per variant (currently `luminal_compiled`).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from contextlib import contextmanager
+
+import numpy as np
+import torch
+
+# Make the sibling import resolve regardless of CWD.
+sys.path.insert(0, "/home/ubuntu/luminal/examples/dlrm")
+from sweep_pytorch import DLRMv1, make_inputs, time_rounds  # noqa: E402
+
+import luminal  # noqa: E402
+from luminal import luminal_backend  # noqa: E402
+
+
+@contextmanager
+def _relaxed_dynamo_limits():
+    pr = torch._dynamo.config.recompile_limit
+    pc = torch._dynamo.config.cache_size_limit
+    torch._dynamo.config.recompile_limit = 64
+    torch._dynamo.config.cache_size_limit = 64
+    try:
+        yield
+    finally:
+        torch._dynamo.config.recompile_limit = pr
+        torch._dynamo.config.cache_size_limit = pc
+
+
+def _run_luminal_compiled(num_cat: int, batch: int, rows: int, m_spa: int, bag: int, device):
+    """Time `torch.compile(model, backend=luminal_backend)`.
+
+    Do NOT wrap the compiled call in `torch.cuda.CUDAGraph` — luminal's
+    `cuda_lite` runtime already captures and replays a CUDA graph
+    internally (one host op per `execute()` call), so an external wrap
+    would just be hiding Python-wrapper / FFI overhead rather than
+    measuring luminal's actual perf.
+    """
+    plain = DLRMv1(num_cat, rows, m_spa).eval().to(device)
+    inputs = make_inputs(num_cat, batch, bag, rows, device)
+    with _relaxed_dynamo_limits():
+        torch._dynamo.reset()
+        compiled = torch.compile(plain, backend=luminal_backend, fullgraph=False, dynamic=False)
+    # Warm up: triggers the export + translate + search.
+    with torch.no_grad():
+        for _ in range(3):
+            _ = compiled(*inputs)
+    torch.cuda.synchronize()
+    return time_rounds(lambda: compiled(*inputs))
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--num-cat", type=int, default=3)
+    ap.add_argument("--rows", type=int, default=4096)
+    ap.add_argument("--batch", type=int, default=2048)
+    ap.add_argument("--m-spa", type=int, default=16)
+    ap.add_argument("--bag", type=int, default=2)
+    ap.add_argument("--json", action="store_true")
+    args = ap.parse_args()
+
+    device = torch.device("cuda")
+    ms = _run_luminal_compiled(args.num_cat, args.batch, args.rows, args.m_spa, args.bag, device)
+    rec = {
+        "variant": "luminal_compiled",
+        "num_cat": args.num_cat,
+        "batch": args.batch,
+        "m_spa": args.m_spa,
+        "bag": args.bag,
+        "rows": args.rows,
+        "ms": ms,
+        "samples_per_sec": args.batch / (ms / 1000.0),
+    }
+    if args.json:
+        print(json.dumps(rec))
+    else:
+        print(
+            f"==== luminal_compiled cfg: num_cat={args.num_cat} batch={args.batch} ===="
+        )
+        print(f"  luminal_compiled                  {ms:8.4f} ms   "
+              f"({rec['samples_per_sec']:>12,.0f} samples/s)")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())