Compat matrix: current-release rows¶

This page records rows of the TorchLens compat matrix that were measured against releases newer than the pinned CI environment (python 3.10, torch 2.13.0, transformers 5.18.0). Each row is a capture plus validation run, executed by a test in tests/real_model/test_compat_matrix_rows.py, or a named export check. A row says exactly what ran, on which versions and which kind of host; a row that did not run says so and never shows as passing.

The broader support statement (torch 2.1 to 2.12+ for eager capture, the torch version table, the limitations catalog) lives in limitations.md and the torch capability rows of tl.report.doctor().

Rows measured 2026-11-03¶

Router-logit outputs¶

Three small mixture-of-experts causal LMs (Mixtral, Qwen2-MoE and OLMoE architectures; two layers, four experts, top-2 routing, constructed weights, no download) run with output_router_logits=True under both the eager and the sdpa attention implementations. Each trace returns the exact model output: the logits and every per-layer router-logit tensor are bit-equal to the direct call, each router-logit tensor is an output op of the trace, and forward replay validation (tl.validate(..., scope="forward")) passes.

DTensor tensor parallelism¶

Two processes form a CPU gloo device mesh. A two-linear MLP is sharded over it through torch.distributed.tensor.parallel.parallelize_module (column-wise then row-wise), and the sharded forward runs (the parameters are DTensors). On every rank, capture refuses at entry with the typed tl.errors.DistributedCaptureUnsupportedError, whose findings name the dtensor and tensor_parallel kinds. This refusal is the supported behaviour today: capturing below the DTensor dispatch layer would record rank-local shims instead of the real ops. A capture design for sharded modules is planned but not built; until it is, capture a rank-local dense copy of the module.

A model loaded by transformers with distributed_config=transformers.DistributedConfig(tp_plan="auto") is the row's second half. A tiny random-weight Llama checkpoint, pinned by revision and digest in the test artifact registry, loads sharded on each rank (the attention projections are DTensors), its logits match the same checkpoint loaded unsharded, and capture refuses with the same typed error and kinds. Transformers applies a tp_plan only when the rank layout is in the process environment (RANK, WORLD_SIZE, LOCAL_RANK); without it the checkpoint loads unsharded and there is no tensor parallelism to test. Transformers refuses a tp_plan on any machine with an MPS device, so this half does not run on macOS. On macOS, gloo needs the loopback interface for a local mesh (GLOO_SOCKET_IFNAME=lo0; the test sets it when unset), else the shard scatter waits out its timeout.

flash_attention_4¶

attn_implementation="flash_attention_4" needs a CUDA device that the FlashAttention 4 package supports (compute capability 8.0 and newer; transformers refuses it on CPU). FlashAttention 4 (flash_attn.cute) allocates its output with torch.empty and fills it from a CuTe kernel reached through TVM FFI, below the torch function layer TorchLens wraps and without the numpy or dlpack export the dynamic foreign-kernel detector sees, so a plain capture would record the allocation in place of the attention. TorchLens therefore records every autograd.Function of flash_attn.cute.interface (FlashAttnFunc, FlashAttnVarlenFunc) as one autogradfunction boundary op with transform_kind="autograd.Function" and transform_config["exits"] == ["flash_attn.cute"]: its parents are the projected q, k and v, its output the attention output, and it replays by calling the Function's apply again (the treatment numpy and dlpack round trips already get). The op is disclosed once per Function class per capture by the foreign_kernel_boundary warning. On the measured GPU the trace returns the exact logits, and with the completeness witness armed for every capture (torchlens.backends.torch.wrappers.wrap_torch(completeness_witness=True)) the trace is verified (trace.provenance.verified is True, no CaptureGapWarning). Under torch.no_grad() the kernel returns the attention output alone and tl.validate(..., scope="forward") passes with the dispatch census balanced (nothing exempted). With autograd enabled the kernel also returns its log-sum-exp, so each call is two sibling ops sharing one func_call_id (index 0 the bfloat16 attention output, index 1 the float32 log-sum-exp that nothing consumes); the log-sum-exp depends on q and k only. The known-kernel table declares that per-output dependence (ruling: "(b) Declare per-output dependence per kernel in a table."), so the perturbation check excuses the log-sum-exp's v edge by declaration and nothing else: the attention output's q, k and v edges and the log-sum-exp's q and k edges stay perturbation-tested, an undeclared kernel, output or argument position is strict, and test_flash_attention_4_grad_enabled_validates asserts the grad-enabled capture validates. The same model on sdpa or eager attention, same GPU and dtype (bfloat16), captures and validates op by op. tl.report.doctor() reports whether the package is installed as HAS_FLASH_ATTENTION_4_INTERFACE (read from the package's search path; nothing is imported). The row's test asserts the boundary ops, the exact output and the passing validation under torch.no_grad() and with autograd enabled; a companion test pins the verified marking under the armed witness.

flash_attention_2¶

attn_implementation="flash_attention_2" needs a CUDA device and the flash-attn package, whose prebuilt wheels pin a torch release (2.4 to 2.8 as of flash-attn 2.8.3.post2; there is no wheel for torch 2.13 or 2.14, and the PyPI distribution is a source build). The kernel entry flash_attn.flash_attn_interface._flash_attn_forward is a torch custom op, so the torch function layer TorchLens wraps sees it directly: no boundary op, no foreign_kernel_boundary disclosure and no CaptureGapWarning, with the completeness witness armed as well. Each attention call is recorded as four sibling ops sharing one func_call_id (index 0 the bfloat16 attention output that feeds the next op; the float32 log-sum-exp, the dropout mask and the int64 rng state, which nothing consumes), each with the projected q, k and v as parents. On the measured GPU the trace returns the exact logits under torch.no_grad() and with autograd enabled. Forward validation passes (test_flash_attention_2_validates): the unconsumed siblings do not read every parent the call carries (the log-sum-exp and the dropout mask read q and k, the rng state none of q, k and v), and the kernel table declares that per-output dependence for the custom op flash_attn::_flash_attn_forward (ruling: "(b) Declare per-output dependence per kernel in a table."), so the perturbation check excuses exactly those edges. The attention output has no row and keeps every edge tested; another operator under the same name, a parent at any other position or a keyword stays strict. The dated gap flash-attention-2-impl in the real-model suite is closed. With dropout (dropout_p > 0, return_attn_probs=True, train mode) every output replays bit for bit (test_flash_attention_2_dropout_replays_exactly): replay restores the captured CUDA generator state, so the kernel redraws the same dropout mask, and the mask's padding past the sequence lengths, which the kernel leaves undefined, is zeroed where the recorded call returns, in the capture and in the replay alike.

Model Explorer with adapter 0.2.x¶

ai-edge-model-explorer-adapter 0.2.0 (August 2026) replaced the adapter's Python extension with a C ABI shared library. Every ai-edge-model-explorer release the model-explorer extra admits (>=0.1.33,<0.2) depends on the 0.2 adapter line: 0.1.33 needs ~=0.2.0 and 0.1.34 needs ~=0.2.1. TorchLens writes Model Explorer JSON graph collections, which the app loads through its built-in JSON path; the adapter's converters handle other formats and are not on TorchLens' export path.

Measured 2026-11-03 on a Linux x86_64 CPU worker in the pinned CI environment (python 3.10, torch 2.13.0, transformers 5.18.0):

Each export check traces the model (ResNet-18 with random weights on a 1x3x224x224 input; GPT-2 small from its default config with random weights on 8 token ids), writes the file with tl.export.model_explorer, passes tl.export.validate_model_explorer_payload, rebuilds every graph through the vendor's graph_builder dataclasses (no duplicate node ids). The served read-back starts tl.export.model_explorer_serve on the file, checks the health and index routes, and reads the graph back over the server's file route equal to the file. ResNet-18 exports 151 nodes and 158 edges in the execution graph (about 564 KB); GPT-2 small exports 567 nodes and 638 edges (about 1.39 MB).