Opt-in fast inference: batched encoder precompute + CUDA-graph decode by Tiger14n · Pull Request #120 · OliBomby/Mapperatorinator

Tiger14n · 2026-07-12T13:11:14Z

Adds two opt-in speedups for autoregressive generation. Both default off; when disabled the generate path is byte-for-byte unchanged.

fast_encoder_precompute — The encoder prefill dominates per-window cost. This precomputes all window encoder outputs in batched passes (16-wide), then feeds them into the existing generate loop so each window skips the encoder.

compiled_decode — Replaces the per-step decode with a hand-rolled CUDA-graph capture/replay of the decoder forward (not torch.compile), over a bucketed static KV cache. Auto-enables fast_encoder_precompute.

Measured decode speedups (compiled vs. uncompiled, tok/s)

GPU (arch)	precision	uncompiled	compiled	speedup
RTX 3080 (Ampere)	fp32	55	153	~2.8×
RTX 3080 (Ampere)	bf16	54	302	~5.6×
RTX 3080 (Ampere)	fp16	57	308	~5.4×
GTX 1050 Ti (Pascal)	fp32	36	77	~2.2×

Also confirmed on a GTX 1050 Ti, so compiled decode helps low-end / Colab-class GPUs too. Match precision to the GPU: Ampere+/Turing (incl. T4/Colab) prefer bf16/fp16; older Pascal cards should use fp32.

Super timing

A separate super_timing_compiled flag routes the super-timing generator through the compiled path. Kept as its own toggle because super timing's batched-parallel path can be faster on GPUs with spare compute/VRAM — benchmark before enabling.

Correctness / safety

Opt-in and backward-compatible; no behavior change when off.
Quality-equivalent, not bit-identical: batched encoder matmuls use different GPU kernels than one-at-a-time, so encoder hidden states differ by ~fp-noise (comparable to bf16-vs-fp32 drift). Greedy tracks the original path closely; only sampled runs diverge, at the same seed.
Dynamic-RoPE models (e.g. v30): their decoder runs a data-dependent RoPE frequency update that can't be graph-captured. Since it's a proven no-op within the decode range (the rotary cache covers max_target_positions, which also caps generation), it's neutralized before capture — verified cos/sin bit-identical across the whole range. So compiled works on v30 too.
No truncation vs. the standard path: the static cache is bucketed small for speed, but a window that fills its bucket without EOS retries once on a full-size cache, so generation matches the non-compiled max_length cap at any lookback/lookahead.
Handles CFG, the monotonic-timeshift constraint, and lookback/lookahead EOS identically to model.generate.

Enable with fast_encoder_precompute=true and/or compiled_decode=true.

Files

added osuT5/osuT5/inference/compiled_decode.py — CUDA-graph decode loop
added osuT5/osuT5/inference/fast_generate.py — batched encoder precompute + fast generate
config.py — 3 flags: fast_encoder_precompute, compiled_decode, super_timing_compiled
osuT5/osuT5/inference/cache_utils.py — optional max_cache_len on get_cache
osuT5/osuT5/inference/processor.py — wiring for the fast/compiled path
osuT5/osuT5/inference/super_timing_generator.py — super_timing_compiled path

🤖 Generated with Claude Code

…code Two opt-in speedups for autoregressive generation, both default off; the generate path is unchanged when disabled. - fast_encoder_precompute: precompute all window encoder outputs in batched passes, then feed them into the generate loop so each window skips the ~228ms encoder prefill. - compiled_decode: hand-rolled CUDA-graph capture/replay of the decoder step over a bucketed static cache (auto-enables fast_encoder_precompute). Handles CFG, the monotonic-timeshift constraint, and lookback/lookahead EOS tokens like model.generate; retries a window on a larger cache if it would truncate, and neutralizes dynamic-RoPE models' no-op frequency update so they capture. - super_timing_compiled: route super timing through the compiled path (own flag because its batched-parallel path can be faster on some GPUs). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

OliBomby

This is great, it's way faster than before. I'd love to merge this, but I think it needs some changes to reduce some code duplication and make it more accessible.

OliBomby · 2026-07-12T14:21:12Z

+)
+from ..event import ContextType, EventType
+
+ENC_BATCH = 16  # windows per encoder forward pass


We have a max batch size config parameter. Would prefer to use that instead of hardcoding so this is configurable.

OliBomby · 2026-07-12T14:26:06Z

                predicted_tokens = result[0, max_len:].cpu()
                self.add_predicted_tokens_to_context(context, predicted_tokens, frame_time, trim_lookback, trim_lookahead)

+    def generate_sequential_fast(


I kinda dislike code quality wise how there is separate methods for generate_sequential and generate_sequential_fast, and model_generate and model_generate_fast. I think the idea to precompute encoder features is good enough to just make that the default, so it's good to just replace/update the old methods instead of making copies.

OliBomby · 2026-07-12T14:33:10Z

+def model_generate_fast(model, tokenizer, model_kwargs, generate_kwargs,
+                        encoder_outputs: BaseModelOutput):
+    """Drop-in for server.model_generate that takes precomputed encoder_outputs.
+


I'd like the fast mode to be available in the InferenceServer too. Can you turn the model_generate in server.py into the fast mode? That way it will also be fast in the web-ui and calc_fid scripts. It's alright if there is no opt-out.

The inference server might need another endpoint for pre-computing encoder windows for clients, since the clients won't own the model.

OliBomby · 2026-07-12T14:37:35Z

+
+
+@torch.no_grad()
+def model_generate_compiled(model, tokenizer, model_kwargs, generate_kwargs,


Can this one also be used by the InferenceServer? Would like to see this speed available there for fast batched inference in the Web UI.

Also can you educate me on how this is different from using torch.compile on the HF model generate loop?

Also can you educate me on how this is different from using torch.compile on the HF model generate loop?

torch.compile fuses kernels in the forward and can CUDA graph it, but it leaves the HF's eager generate loop and its per token Python/CPU overhead and host syncs in place, which is the suspected bottleneck for a small decoder. This PR captures the decode step into a CUDA graph directly and replaces the loop to remove that overhead entirely, while handling the custom model, CFG, logits processors and other constraints that make torch.compile(generate) fail here. it's also more portable since CUDA graphs are a runtime feature with no compiler/Triton dependency, it should give Windows users the speedups without having issues with the finicky torch compile support on Windows (though i haven't tested Windows myself).

Also, torch.compile was the first thing attempted but had zero gains on the plain generate loop since the loop overhead dominates. your question had me attempt a quick test, torch.compile on top of the CUDA graph decode had 30-45% gains on top reaching ~400 tok/s, but only if you compile just the fixed-shape decode step. compiling the whole model recompiles on every prompt length and collapses to ~23 tok/s. since it needs a warm-up compile per cache bucket it only amortizes in a long-lived process like the InferenceServer, and it pulls Triton back in, so i'd keep it as a follow-up in a separate PR. i've only measured speed so far, not verified output equivalence.

OliBomby · 2026-07-12T14:54:50Z

+    fast_encoder_precompute: bool = False  # Precompute all encoder outputs in a batched pass before the decode loop (major speedup, bit-compatible)
+    compiled_decode: bool = False  # CUDA-graph decode loop (auto-enables fast_encoder_precompute); ~2-6x faster decode (fp16/bf16 > fp32), quality-equivalent
+    super_timing_compiled: bool = False  # Use the compiled CUDA-graph decode path for super timing instead of the batched-parallel path. Separate from compiled_decode because super timing's parallel path may be faster on some GPUs; benchmark before enabling.


Add these to the config yaml files too.

I think there is very little reason to not enable these options by default. I think let fast encoder precompute be on by default with no way to turn off, and let compiled be on by default for at least V30 and V32, where its confirmed working.
Also BF16 precision should be on by default, and automatically turned off if incompatible GPU is detected.

OliBomby requested changes Jul 12, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Opt-in fast inference: batched encoder precompute + CUDA-graph decode#120

Opt-in fast inference: batched encoder precompute + CUDA-graph decode#120
Tiger14n wants to merge 1 commit into
OliBomby:mainfrom
Tiger14n:feat/compiled-decode

Tiger14n commented Jul 12, 2026

Uh oh!

OliBomby left a comment

Uh oh!

OliBomby Jul 12, 2026

Uh oh!

OliBomby Jul 12, 2026

Uh oh!

OliBomby Jul 12, 2026

Uh oh!

OliBomby Jul 12, 2026

Uh oh!

OliBomby Jul 12, 2026

Uh oh!

Tiger14n Jul 12, 2026

Uh oh!

OliBomby Jul 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants



		@torch.no_grad()
		def model_generate_compiled(model, tokenizer, model_kwargs, generate_kwargs,

Uh oh!

Conversation

Tiger14n commented Jul 12, 2026

Measured decode speedups (compiled vs. uncompiled, tok/s)

Super timing

Correctness / safety

Files

Uh oh!

OliBomby left a comment

Choose a reason for hiding this comment

Uh oh!

OliBomby Jul 12, 2026

Choose a reason for hiding this comment

Uh oh!

OliBomby Jul 12, 2026

Choose a reason for hiding this comment

Uh oh!

OliBomby Jul 12, 2026

Choose a reason for hiding this comment

Uh oh!

OliBomby Jul 12, 2026

Choose a reason for hiding this comment

Uh oh!

OliBomby Jul 12, 2026

Choose a reason for hiding this comment

Uh oh!

Tiger14n Jul 12, 2026

Choose a reason for hiding this comment

Uh oh!

OliBomby Jul 12, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants