Crate rlx_opscope

Expand description

rlx-opscope — a data-pattern recording harness.

The idea: run existing op graphs with varied synthetic (and, later, real) data and record cheap sketches of every tensor flowing through the ops we care about — matmuls first. Offline we mine those sketches for exploitable structure (sparsity, per-channel outliers, quantization headroom, sequence structure) that justifies a specialized kernel.

The load-bearing trick (see the crate README) is that we do not tap the executor. We rewrite the graph: for each op-site we append reduction / histogram nodes on its inputs and output and mark them as extra graph outputs. The stats are then computed by the backend’s own kernels — so this works identically on CPU, Metal, CUDA, … and costs nothing when not applied.

This module provides:

inject_matmul_stats — the stat-injection graph pass (matmul MVP).
StatConfig / StatSpec — what to record and how to label it.
Dist / gen — synthetic distribution generators.
Recorder — a dependency-free tidy-CSV sink.

Modules§

bytes: Static bytes-moved accounting — a memory-traffic ledger read off the IR.
dataflow: Dataflow repeated-pattern mining. Where motifs walks linear single- consumer op chains, this canonicalizes each node’s input cone (the branching sub-DAG that produces it, up to a depth) into a structural signature and counts recurrences. A sub-DAG that recurs across the graph is a unit of repeated computation — a candidate to be decomposed once and shared, or fused into a specialized op. This is the “track the data flow that can be additionally decomposed, especially repeated patterns” seam.
demo: Small illustrative graph builders (mlp / transformer / moe) shared by the structural/shape demos. Shapes are representative; the structural and shape miners read op kinds + dims, they don’t execute these.
guard: A real specialized kernel behind a runtime guard — the actuation the optimizer recommends. opscope-optimize says “this op’s input repeats over time → memoize”; MemoMatmul is that kernel: a cheap guard (input hash + exact verify) skips the matmul on a cache hit and returns the stored output, falling back to a dense compute on a miss. Correct by construction — the verify guarantees a hit is the identical input, so the output is bit-exact.
kernels: Data-pattern-specialized matmul kernels — the actuation for the exploits the probe finds. Each is a decomposition of the dense matmul that the mined structure makes cheap:
layers: Per-layer decomposition mining + a DADO allocator.
live: Live sampled recorder — the production path. Instead of writing every run to a CSV, it rides a real forward pass, samples 1-in-N runs, and folds each op-site’s element-wise fingerprint into bounded crate::online sketches: a HyperLogLog estimates how many distinct inputs a site has seen (→ temporal recurrence / memoize opportunity) and a reservoir keeps a sample — all at near-zero, constant memory regardless of stream length.
motifs: Structural op-subsequence mining — the “patterns in the sequences of ops” reading. Pure graph analysis, no runtime data: walk single-consumer chains (the ones a fusion pass could collapse) and count recurring op-kind n-grams. A motif that recurs across many layers, weighted by its length, is a fusion-kernel candidate.
online: Cross-cutting — streaming/online sketches for a sampled real-inference recording path. To ride the production forward pass at near-zero overhead you can’t keep every tensor; you keep bounded per-op summaries updated in one pass: a reservoir sample, a fixed-bucket quantile sketch (t-digest-lite), and a HyperLogLog cardinality estimate. All deterministic (hash-seeded, no RNG dependency), so recordings are reproducible.
optimize: Optimization-opportunity miner. Fuses three signals the harness records into one ranked list of what to actually optimize:
parity: Op-level parity checking. After stat-injection / fusion-skip / attention decompose / any IR transform, prove the result is unchanged — not just at the final output, but op by op. expose_all_ops marks every compute node as a graph output, so two structure-preserving graphs can be compared tensor by tensor. Structure-changing transforms (fusion) are compared at the primary output within tolerance.
probe: Tier 2 — host probes on a captured (deep-dumped) tensor, closing the gaps the cheap CSV sketches can’t see: effective rank (→ factored matmul), N:M structured-sparsity feasibility (→ sparse tensor-core GEMM), quantization error at a target bit-width (→ int kernel), and value cardinality (→ LUT/palettized weights). Pure host arithmetic on &[f32].
shapes: Tier 3 — mine the workload rather than the values. Per op we estimate FLOPs, bytes moved, and arithmetic intensity (FLOP/byte), which gives a roofline classification (compute- vs memory-bound → drives fusion strategy) and a histogram of the actual GEMM shapes that occur (→ autotune/specialize kernels for the hot shapes, generate a dispatch table). Pure static analysis over any Graph — no execution.
svd: Small dependency-free truncated SVD — the shared numerical primitive under every tensor decomposition here: Tucker (HOSVD unfolds → SVD each mode), TT (TT-SVD sweeps SVD the sequential reshape), and Monarch (rank-1 SVD of each block). Power iteration with explicit deflation: robust and plenty accurate for a mining tool that only needs the leading r triplets.
timing: MEASURED execution time — the empirical complement to crate::shapes’s analytical roofline. shapes.rs predicts “memory-bound” from arithmetic intensity vs a fixed ridge point; this module measures achieved bandwidth so the claim carries a number, and attributes real wall-clock to op regions.

Structs§

Recorder: A dependency-free tidy-CSV sink: one row per sketch element. Columns: run_id,backend,dist,M,K,N,site,role,stat,idx,value. Tidy/long form so scalars (idx=0), per-channel vectors (idx=channel), and histograms (idx=bin) all share one flat schema that pandas/polars can groupby.
ResidualSpec: One residual-Add tap: the recorded Σx² of each input, plus the block index.
StatConfig: Which sketches to append per tapped tensor.
StatSpec: One appended output: where it is in graph.outputs, and what it means.

Enums§

Dist: Synthetic input distributions. Each is chosen to make a different sketch pop, so the miner can be validated against ground truth.

Functions§

inject_activation_fakequant: Rewrite graph so every Op::MatMul’s activation input (input[0]) is per-token int8 fake-quantized — the missing half of an end-to-end W8A8 simulation (combine with int8 weights in the WeightMap). In a fused-attention qwen graph the MatMul nodes are exactly the linear projections (Q/K/V/O/ gate/up/down + lm_head) — attention’s internal QK/AV live inside Op::Attention — so this quantizes precisely the linear-layer activations, as real W8A8 does. Outputs are preserved (same indices); compile with skip_fusion (the inserted nodes would otherwise trip SwiGLU fusion). per_channel picks the scale scope (see fakequant) — false = per-token (deployable W8A8), true = per-channel (the outlier-robust ceiling).
inject_attention_stats: Tap fused Op::Attention nodes: decompose each to its softmax-weights node (via rlx’s own attention_softmax_weights) and append per-query peak mass + per-key received mass as extra outputs. This is how attention concentration / sink keys are recovered from a fused attention op (which otherwise exposes no softmax to tap). Adds a shadow Q·Kᵀ→softmax per site.
inject_infer_stats: Append attention (Op::Softmax) and MoE-routing (Op::TopK) sketches as extra graph outputs — the distinctly inference signals:
inject_matmul_stats: Rewrite graph so that every Op::MatMul / Op::DequantMatMul site’s lhs, rhs, out tensors get sketch outputs appended. The original outputs stay first (same indices), so downstream code that reads output 0 is unaffected. Returns the rewritten graph and the manifest describing each appended output.
inject_matmul_stats_filtered: inject_matmul_stats but only taps sites for which keep(lhs_numel, out_numel) is true. The injected graph’s compile cost grows superlinearly with the number of tapped sites, so whole-model injection (100+ matmuls) can become intractable; a keep predicate lets a caller tap just the tensors it reports (e.g. residual stream + query by their element counts), keeping the graph small enough to compile at depth.
inject_residual_stats: Tap every residual Add(residual, block_out) — record Σx² of both inputs so the analysis can compute ‖block-delta‖/‖residual‖, the block-influence signal (ShortGPT/Gromov): a sublayer whose delta ≈ 0 barely changes the residual stream ⇒ it’s near-identity ⇒ a layer-skip candidate. Compile with skip_fusion so the residual adds aren’t folded into FusedResidualRmsNorm before the tap sees them. Outputs are preserved (same indices).
sample: Generate a rows × cols matrix (row-major) drawn from dist, seeded for reproducibility.
skip_residual_blocks: Rewrite graph to skip the residual blocks at the given execution orders (same ordering as inject_residual_stats): each targeted Add(residual, block_out) is replaced by its residual input (input[0]), so the block’s output is bypassed — and its whole subgraph becomes dead code the compiler drops, actually removing that layer’s weight-stream + compute. The concrete “replace-like-fusing” action for a near-identity layer (layer minimization).
write_site_costs: Write a sidecar mapping each op-site to its FLOPs (deduplicated), consumed by opscope-optimize to weight temporal-recurrence savings by op cost.