zerogpu-hacking/sensenova-u1-5-aoti

AOTInductor packages for SenseNova-U1.5-8B-MoT

Compiled artifacts for the per-denoising-step hot path of sensenova/SenseNova-U1.5-8B-MoT, for ZeroGPU Spaces via spaces.zero.torch.aoti.

Measured 1.26x end-to-end on a 28-step 2048x2048 text-to-image run (33.25 s -> 26.42 s), and the compiled path is closer to an fp32 reference than the eager bf16 path it replaces.

What is compiled

One Qwen3DecoderLayer image-generation path, reused by all 42 layers.

language_model.model is 94.9% of wall time, and the 42 identical decoder layers are essentially all of it: 13.358 ms/layer x 42 = 561 ms against a 538 ms measured trunk call. The ViT (2.5%) and fm_head (0.7%) cap out at ~3% combined and are not worth compiling.

The layers' gen path reads its preallocated flash KV buffers as attributes on the cache object, which torch.export cannot trace. The vendored modeling_qwen3.py therefore adds a pure-tensor restatement of that math - gen_attn_core / gen_block_core / Qwen3GenBlock / Qwen3DecoderLayer.gen_core - taking the prefix K/V as explicit tensor arguments and concatenating the current segment instead of copying it into a preallocated buffer.

That restructure was verified to be a no-op on its own: bit-identical images and 32.98 s -> 32.77 s (0.994x) against the original in-place path.

Compiled signature:

gen_core(hidden_states[B,S,C], indexes[3,S], prefix_k[B,P,H_kv,D], prefix_v[B,P,H_kv,D]) -> [B,S,C]

S (image tokens) and P (prefix length) are dynamic, so one package serves every T2I aspect bucket (S = 3888..4096, all six verified) and any prompt length. No static-shape fallback was needed.

Layout

<key>/package/metadata.json
<key>/package/submodules/layers/package.pt2

key = <dtype>/torch<X.Y>/sm<cc>/<shape>. Published: bf16/torch2.8/sm120/dynamic.

A package is valid for exactly one key, and a mismatched one segfaults rather than raising, so the loader refuses anything that is not an exact match and falls back to eager. Weights are not in the package (1.1 MB total): each layer's own parameters and buffers are bound to the one shared graph. Binding must use named_parameters() + named_buffers(), not state_dict() - the two rotary inv_freq buffers are registered non-persistent, so they are missing from state_dict() but are real graph constants.

Measured

RTX PRO 6000 Blackwell Server Edition MIG 2g.48gb (sm120), torch 2.8.0+cu128, bf16, SDPA, 2048x2048, S=4096, P=265.

	eager	AOTI	speedup
one decoder layer	13.358 ms	10.928 ms	1.222x
x42 (trunk call)	561 ms	459 ms
28-step 2048x2048	33.25 s	26.42 s	1.259x
per step	1188 ms	944 ms
peak memory	34.71 GiB	34.62 GiB

Compile takes 31 s. max_autotune was measured to add nothing (cuBLAS already wins the large GEMMs) and is not used. Per-thread package load is ~30 ms - the ContextVar cache in LazyAOTIModel is per-thread, but the cost is negligible and needs no dedicated executor thread.

The win is fusion of the non-GEMM work, not better GEMMs. By kernel time one layer is ~57% large GEMM and ~14% attention; the remaining ~29% (RMSNorms, the rotary recomputation, the chunk/transpose/cat reshaping around q/k/v, and the GQA repeat_interleave) is what inductor collapses.

Numerical equivalence

The compiled path is more accurate than the eager path, not less. Against an fp32 reference built from the same weights and inputs:

layer	eager bf16 vs fp32	AOTI bf16 vs fp32	ratio
0	1.160e-2	8.848e-3	0.76
20	9.220e-3	6.066e-3	0.66

Both are ~1e-2 RMS-relative from fp32: that is the model's own bf16 noise floor at this width. The eager-vs-AOTI disagreement (1.1-1.4e-2 RMS-relative per layer) is the sum of two independent roundings of that size, not a defect.

Consequently the final image does change - 81% of pixels, mean 3.3/255 - but this generation is not stable to any change in accumulation order. Control: swapping between two equally correct SDPA kernels (FLASH vs cuDNN), with no compiler involved, already moves 73% of pixels (mean 1.8/255). Eager vs eager on the same backend is bit-identical, so the baseline itself is deterministic. Composition, colour, and subject are preserved; fine foreground detail is reshuffled.

Attention backends on sm120 (torch 2.8), for reference

S_q=4096, S_kv=4144, H 32/8, D=128, bf16, non-causal:

path	ms
as shipped (repeat_interleave + default SDPA)	1.792
FLASH + repeat_interleave	1.630
FLASH + enable_gqa	1.621
cuDNN + enable_gqa	1.429
memory-efficient	4.158

torch 2.8's SDPA already has a working flash backend on sm120, so installing flash_attn has very little to buy: attention is 1.79 ms of a 13.36 ms layer, and the best possible backend swap saves 0.36 ms/layer = 2.6% end-to-end.

Use

Build recipe and loader: zerogpu-hacking/sensenova-u1-5-aoti-compiler (sn_aoti.py, INTEGRATION.md). Consumer side is one call at startup:

import sn_aoti
sn_aoti.maybe_load(model)   # no-op unless SN_AOTI=1 and the key matches exactly