b12x fused MoE: TrellisMX (procedural NPK) source format with an in-tree P8 runtime

CLAUDE.md

@@ -513,6 +513,63 @@ kernel-source SHA256 wipes and lazily rebuilds the module. Reference usage:
 `flashinfer/quantization/kernels/nvfp4_quantize.py`. Disable with
 `FLASHINFER_CUTE_DSL_DISABLE_CACHE=1`.
 
+The experimental b12x backend uses
+`flashinfer/experimental/b12x/_lib/compiler.py` and stores compiled objects in
+`$B12X_COMPILE_CACHE_DIR`, or `$XDG_CACHE_HOME/b12x/compile` (defaulting to
+`~/.cache/b12x/compile`). Its disk identity includes the callable's transitive
+local source files, compiler support files, resolved specialization, effective
+code-generation options, toolchain, and device architecture. AST import tracing
+includes nested and conditional imports, inherited callables, decorators, and
+callable defaults. External packages participate through toolchain identity.
+Raw environment variables and diagnostic output destinations do not participate;
+environment-selected code generation must be represented by resolved kernel
+parameters or compiler options. Computed imports in implementation modules
+declare local targets in `__b12x_source_dependencies__`. Shared helpers live in
+leaf modules under `_lib/intrinsics/`, `_lib/utils/`, `_lib/dense_gemm/`, and
+`moe/_shared/kernels/`; consumers import the defining module to keep dependency
+closures narrow. Callable source discovery is cached by function or class,
+including the owner class for bound methods, and cleared by
+`clear_compile_cache()`. Source and compiler configuration are process-immutable:
+restart after editing code or changing compiler defaults. Warm in-memory hits
+and prepared CUDA graph replay do not read or parse source files.
+
+b12x autotuning selection keys include the actual compile keys of every eligible
+candidate, including previously losing candidates. Preparation describes their
+programs without lowering or launching kernels and reuses those descriptions for
+compilation. Changes to relevant source, specialization, compiler options, or
+toolchain require a fresh race. Unrelated kernels retain their saved selections.
+The tuning contract's schema versions and `B12X_TUNING_CACHE_VERSION` remain
+explicit controls for changes to tuning semantics and manual retuning.
+
+Experimental b12x preparation requests can opt into `FinalistPolicy` to retain
+measured configurations by tuning-parameter family, independently of the
+standalone winner. The policy has a separate selection-cache identity. Sharded
+finalist requests pass every rank's contribution, including empty shards, to
+`PreparationJob.advance(tuning=...)` in rank order; ordinary requests pass one
+consolidated winner per key. Callers
+prepare alternative plans and capture complete computations as
+`ContextCandidate` objects, then call `PreparationSession.race_contexts` before
+freezing the session. The returned contextual winner does not replace standalone
+selections. Input reset and production occur outside the timed replay; producer
+work that belongs in the measurement must be captured inside the graph. Cold
+cache eviction follows input production. Multi-rank contexts agree on candidate
+catalogs, rank membership, sample schedules, producer counts, and cache policy
+before replay, and score each sample by its slowest rank. Trial graph cleanup
+callbacks run after synchronization; callers release losing plans before final
+capture. Contextual races require an unfrozen session with preparation drained.
+
+Every completed b12x tuning race appends a self-describing record to the
+`ledger_v1` table of the race-progress database beside the selection cache:
+component, encoded query, effort, coverage, measurement protocol, and for every
+candidate its assignment, lowered config, program keys, the lowest
+autotune effort that admits it, and per-round timings, plus the batches raced,
+rejected candidates, and wall seconds per phase. Records are append-only, so
+retuning accumulates replications. `python -m b12x.tools.tuning_ledger
+[--root DIR] [--records] [--replications]` merges tensor-parallel shard records
+into runs and reports default regret, pruning regret (the default tier's best
+candidate against an audit winner), runner-up margins, round-to-round noise,
+and replication agreement.
+
 ### Dispatch Macros
 
 Handle combinatorial parameter spaces:

benchmarks/README.md

@@ -18,6 +18,71 @@ the sample count to stay within the rotating-state capacity and set
 iterations fit. `--dry-run-iters` and `--repeat-iters` request fixed iteration
 counts; they are not duration targets.
 
+## Prepared DeepSeek V4.1 MLA on SM120/SM121
+
+The experimental [V4.1 MLA benchmark](experimental/b12x/benchmark_deepseek_v41_mla.py)
+runs rank-local Flare TP4 attention shapes on one selected GPU:
+
+```sh
+python benchmarks/experimental/b12x/benchmark_compressed_sparse_mla.py \
+  --preset flare-dsv41-tp4 \
+  --model-config /path/to/DeepSeek-V4.1-Flash-1017/config.json \
+  --device 4 --output-csv /tmp/v41-mla.csv
+```
+
+The profile uses BF16 arithmetic, 16 query heads, a 1M-token context capacity,
+8192-column indexed page tables, and a 512-token live context. It covers SWA,
+C2, C1, and draft SWA through prepared plans and captured CUDA graph replay.
+`--verify-shapes BxWxP` keeps the per-request query ceiling `W` separate from
+the physical packed query rows `P`. A 4096-row prefill uses 128 suffix rows
+for C1 attention. `--describe` prints the resolved contracts without importing
+GPU libraries; use `--preset flare-dsv41-tp4 --help` for all profile options.
+
+Queries, cache values, sinks, and selections are synthetic. The indexed cache
+view has 895489 rows, requiring about 33 GB for C1; ordinary PyTorch backing
+differs from Flare's virtual-memory residency. `--cache-pages` explicitly
+overrides that descriptor for smaller fixtures. Timing includes attention and
+logical-to-physical page mapping, excluding projections, index selection,
+cache writes, and communication. L2 eviction runs outside timed replay by
+default; `--no-flush-l2` disables it. CSV results retain actual layouts,
+selected configurations, numerical errors, cache policy, and raw timings.
+
+## Prepared BF16 router projection and top-k
+
+The experimental [router benchmark](experimental/b12x/benchmark_router_projection.py)
+compares tuned b12x projection plus top-k with a CuTe DSL single-launch prototype.
+Both round projection logits to BF16, break ties by descending expert ID, and
+return selected logits and normalized softmax weights in FP32. The fused kernel
+uses caller-owned logit scratch and completion counters; it is not selected by
+the serving runtime. Four warps cooperate on each expert by default, with eight
+warps per CTA. This SIMT kernel targets small decode batches; wider batches can
+be slower than the split path.
+
+```sh
+python -m benchmarks.experimental.b12x.benchmark_router_projection \
+  --device 10 --checkpoint /path/to/Qwen3.8-Flash-Next-FP4-Local \
+  --layers 0 23 47 mtp --rows 1 2 3 4 5 6 --output /tmp/router.json
+```
+
+Activations are synthetic. Without `--checkpoint`, weights are synthetic too.
+The default geometry is 2560 inputs, 512 experts, and top-10 routing. Correctness
+checks require exact expert IDs against BF16 projection and stable sorting,
+check selected logits and weights, and replay changed inputs and all-zero ties
+with poisoned scratch and fixed allocations. A failed check prevents timing.
+The reference disables reduced-precision BF16 partial reductions for its matmul,
+then restores the setting before executing either measured implementation.
+The preparation session remains alive and frozen through measurement.
+
+Timings cover the complete operation and counter recycling in captured graphs,
+using balanced rounds from `b12x.testing.benchmark.measure_calls`. The default
+128 MiB L2 eviction runs outside timing; `--eviction-mib 0` measures warm-cache
+replay. JSON retains per-layer/per-shape timings, raw samples, the selected split
+projection configuration, and the fused `--rows-per-tile`, `--warps`, and
+`--warps-per-expert` geometry.
+Use `--implementations split` for the control alone or `--check-only` to validate
+without timing. These component timings do not measure expert overlap or model
+throughput.
+
 ## Overview
 
 This framework provides tools to:

benchmarks/experimental/b12x/autotune_dense_fp4_tile_k.py

@@ -27,9 +27,10 @@
     make_nvfp4_operand,
 )
 from benchmarks.experimental.b12x.common import make_l2_flush_fn, resolve_l2_flush_bytes
-from b12x._lib.dense_gemm import _select_default_dense_gemm_plan, dense_gemm
-from b12x._lib.dense_gemm import DenseGemmKernel
-from b12x._lib.utils import get_num_sm
+from b12x._lib.dense_gemm.policy import _select_default_dense_gemm_plan
+from b12x._lib.dense_gemm.execution import dense_gemm
+from b12x._lib.dense_gemm.kernel import DenseGemmKernel
+from b12x._lib.utils.device import get_num_sm
 
 
 @dataclass(frozen=True)

benchmarks/experimental/b12x/autotune_dense_mxfp8_plan.py

@@ -30,12 +30,10 @@
     cosine_similarity,
     make_mxfp8_operand,
 )
-from b12x._lib.dense_gemm import (
-    DenseGemmKernel,
-    _select_default_dense_gemm_plan,
-    dense_gemm,
-)
-from b12x._lib.utils import get_num_sm
+from b12x._lib.dense_gemm.kernel import DenseGemmKernel
+from b12x._lib.dense_gemm.policy import _select_default_dense_gemm_plan
+from b12x._lib.dense_gemm.execution import dense_gemm
+from b12x._lib.utils.device import get_num_sm
 
 
 BASELINE_COSINE_THRESHOLD = 0.9999

benchmarks/experimental/b12x/autotune_dense_mxfp8_split_k.py

@@ -28,12 +28,12 @@
     make_mxfp8_operand,
 )
 from benchmarks.experimental.b12x.common import make_l2_flush_fn, resolve_l2_flush_bytes
-from b12x._lib.dense_gemm import (
+from b12x._lib.dense_gemm.policy import (
     _select_default_dense_gemm_plan,
     _select_mxfp8_tile_k,
-    dense_gemm,
 )
-from b12x._lib.utils import get_num_sm
+from b12x._lib.dense_gemm.execution import dense_gemm
+from b12x._lib.utils.device import get_num_sm
 
 
 BASELINE_COSINE_THRESHOLD = 0.9999

benchmarks/experimental/b12x/benchmark_block_fp8_gemv.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+"""Small-row serialized block-FP8: GEMV regime vs the dense GEMM plan.
+
+Times fixed ``expected_m`` plans (the decode shape vLLM declares per CUDA graph
+size) through ``blockscaled.mm_block_fp8``, once with the GEMV regime and once
+with it disabled, using the preparation engine's GPU sampler.  ``--cold`` flushes L2 before
+every call, which is what a decode step sees for weights larger than L2.
+
+    python benchmarks/experimental/b12x/benchmark_block_fp8_gemv.py [--cold] [--rows 1,2,4,8]
+        [--shapes 3712x4096,1024x7168]
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+
+import torch
+
+from b12x.gemm import blockscaled
+from b12x.gemm.blockscaled import _block_fp8_gemv
+from b12x.preparation import PreparationSession, PreparedCall
+from b12x.testing.benchmark import measure_calls
+
+DEFAULT_SHAPES = [
+    (n, k)
+    for k in (1536, 4096, 7168)
+    for n in (256, 512, 1024, 2048, 3712, 4096, 6144, 8192, 16384)
+]
+
+
+@contextlib.contextmanager
+def _gemv_enabled(enabled: bool):
+    original = _block_fp8_gemv.supports
+    if not enabled:
+        _block_fp8_gemv.supports = lambda *args: False
+    try:
+        yield
+    finally:
+        _block_fp8_gemv.supports = original
+
+
+def _prepare(lhs, lhs_scale, rhs, rhs_scale, rows):
+    query = blockscaled.FixedBlockscaledQuery(
+        recipe="block_fp8",
+        call_kind="serialized",
+        max_rows=rows,
+        in_features=rhs.shape[1],
+        padded_in_features=rhs.shape[1],
+        out_features=rhs.shape[0],
+        input_dtype="float8_e4m3fn",
+        output_dtype="bfloat16",
+        expected_m=rows,
+    )
+    plan = blockscaled.plan(query)
+
+    def call(state):
+        return PreparedCall(
+            run=lambda: state.run_serialized(
+                lhs,
+                lhs_scale,
+                rhs,
+                rhs_scale,
+                None,
+                ab_dtype="float8_e4m3fn",
+                sf_dtype="float32",
+                c_dtype="bfloat16",
+                sf_vec_size=128,
+                block_fp8=True,
+                stream=None,
+            )
+        )
+
+    session = PreparationSession(device=lhs.device, autotune=False, compile_workers=4)
+    session.__enter__()
+    session.prepare((plan.request(name="bench", prepare_call=call),))
+    return session, plan
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--rows", default="1,2,4,8")
+    parser.add_argument("--shapes", default=None, help="NxK list, e.g. 3712x4096")
+    parser.add_argument("--cold", action="store_true")
+    parser.add_argument("--iters", type=int, default=25)
+    args = parser.parse_args()
+    shapes = (
+        [tuple(int(v) for v in s.split("x")) for s in args.shapes.split(",")]
+        if args.shapes
+        else DEFAULT_SHAPES
+    )
+    rows_list = [int(r) for r in args.rows.split(",")]
+    device = torch.device("cuda")
+    flush = (
+        torch.empty(320 << 20, dtype=torch.uint8, device=device) if args.cold else None
+    )
+    print(
+        f"{'N':>6} {'K':>6} {'M':>2} {'dense_us':>9} {'gemv_us':>8} {'speedup':>7}  eligible"
+    )
+    for n, k in shapes:
+        gen = torch.Generator(device=device).manual_seed(n * 131 + k)
+        rhs = torch.randn((n, k), generator=gen, device=device).to(torch.float8_e4m3fn)
+        rhs_scale = torch.rand(
+            ((n + 127) // 128, k // 128), generator=gen, device=device
+        )
+        for rows in rows_list:
+            lhs = torch.randn((rows, k), generator=gen, device=device).to(
+                torch.float8_e4m3fn
+            )
+            lhs_scale = torch.rand((rows, k // 128), generator=gen, device=device)
+            lhs_ref = lhs.double() * lhs_scale.double().repeat_interleave(128, dim=1)
+            rhs_ref = rhs.double() * rhs_scale.double().repeat_interleave(128, dim=0)[
+                :n
+            ].repeat_interleave(128, dim=1)
+            reference = lhs_ref @ rhs_ref.T
+            ulp = reference.abs().clamp_min(1e-30) * 2.0**-7
+            calls = {}
+            with contextlib.ExitStack() as stack:
+                for label, enabled in (("dense", False), ("gemv", True)):
+                    with _gemv_enabled(enabled):
+                        original_max = _block_fp8_gemv.MAX_OUT_FEATURES
+                        _block_fp8_gemv.MAX_OUT_FEATURES = 1 << 30
+                        try:
+                            session, plan = _prepare(
+                                lhs, lhs_scale, rhs, rhs_scale, rows
+                            )
+                        finally:
+                            _block_fp8_gemv.MAX_OUT_FEATURES = original_max
+                        stack.callback(session.close)
+
+                    def fn(plan=plan):
+                        return blockscaled.mm_block_fp8(
+                            lhs, lhs_scale, rhs, rhs_scale, plan=plan
+                        )
+
+                    output = fn()
+                    error = (output.double() - reference).abs()
+                    assert torch.isfinite(output).all()
+                    if enabled:
+                        assert bool(
+                            (error <= ulp + 1e-6 * reference.abs().max()).all()
+                        ), label
+                    else:
+                        # Dense split-K may round partial sums to BF16.
+                        assert float(error.norm() / reference.norm()) < 4e-3, label
+                    calls[label] = PreparedCall(
+                        run=fn, produce=lambda: None, owners=(plan, session)
+                    )
+                measured = measure_calls(
+                    calls,
+                    samples=args.iters,
+                    eviction=(lambda: None) if flush is None else flush.zero_,
+                )
+                times = measured.latencies_us
+            eligible = _block_fp8_gemv.supports(rows, n, k)
+            print(
+                f"{n:6d} {k:6d} {rows:2d} {times['dense']:9.2f} {times['gemv']:8.2f} "
+                f"{times['dense'] / times['gemv']:7.2f}  {'yes' if eligible else 'no'}",
+                flush=True,
+            )
+
+
+if __name__ == "__main__":
+    main()