Krea 2 TileMix Attention (FP16/INT8 tiles)

Krea 2 is a 12B single-stream DiT, and at 1–2 megapixels the thing eating your seconds is attention. Not the step count, not the VAE - the joint [text, image] sequence squaring off against itself 28 times. This node swaps that attention for a fused Triton kernel that computes some of the score matrix in INT8 and keeps the rest in FP16, routing tile by tile. You give up a little exactness on tiles that barely matter, and you get the time back.

What it actually patches

Only attn(q, k, v) inside the 28 gated blocks. The projections, QK RMSNorm, 3-axis RoPE, the sigmoid(gate) and the output projection all stay stock. That small blast radius matters - attention packs tend to whack the whole model, and you find out three generations later.

The kernel is bidirectional online-softmax attention: diffusion-shaped, not the causal triangle you'd write for an LLM. Every tile of the score matrix goes to an FP16 or INT8 compute path by a static map, and both paths feed one running softmax state. Two things stay FP16 whatever you set: every tile that reads a text key block or writes a text query row, and a diagonal band of key blocks around each query. The INT8 percentage counts only the image-image tiles left over. Text conditioning is never quantized, and local structure gets exact math because it's the part you can see.

It also skips Krea's repeat_interleave, resolving GQA (48 query heads / 12 KV heads) inside the kernel. Krea's stock path is quantifiably slower than plain SDPA, so the bar here is lower than it looks.

The idea and the pattern names come from the TileMix paper, which targets causal LLM decoding; this is an independent bidirectional reimplementation, Apache-2.0.

The inputs that matter

Wire it between your UNETLoader and the sampler:

UNETLoader → Krea 2 TileMix Attention (FP16/INT8 tiles) → KSampler → VAEDecode

• pattern (band) is the one knob worth thinking about. band keeps a diagonal FP16 window and sends the off-diagonal image tiles to INT8; one sends every image tile to INT8; zero is an FP16-only sanity mode. The other three are research leftovers; you won't beat band.
• percentage (0.5) - how much of the image-image matrix goes INT8. The author's sweet spot with band is 0.5–0.75: up for speed, down if quality drops.
• min_tokens (1024) keeps Krea's original attention below that joint token count. There's little to win on a short sequence.
• smooth_k (True) mean-centres K per head before quantisation. Leave it on - exact for softmax, and measurably better on keys with a shared channel offset.
• qk_fp32_accumulate (False) is the accuracy escape hatch - SDPA-level precision on the FP16 tiles for roughly 25% slower FP16 math.
• strict (False) raises kernel errors instead of falling back quietly. Debugging tool, not a daily setting.
• enabled just passes the model through when off.

One MODEL output, into whatever consumes your loaded model downstream.

Install

ComfyUI Manager, search the pack title, or:

cd ComfyUI/custom_nodes
git clone https://github.com/chinoll/TileMix-comfyui.git
# Triton + CUDA PyTorch are already in ComfyUI; nothing else to install

Restart ComfyUI. No model downloads, no calibration - requirements.txt is just torch>=2.4 and triton>=3.0. The kernel imports lazily, so if Triton is missing the node hides itself and logs a warning instead of killing startup. What you actually need is a working Triton build on an NVIDIA card (SM80+, so 30-series and up; developed on RTX 5090/5060 Ti), running bf16 or fp16.

What speedup to expect

Pure FP16 tiles, no INT8 at all, still bench at ~1.13–1.45× plain bf16 SDPA, because GQA moved inside the kernel. band at 0.5 gets ~1.49× median, ~1.65× at 0.75. Against Krea's real repeat_interleave + SDPA path - itself 0.96× SDPA - that's roughly 1.4–1.7× for the block, and the accuracy table behind it uses adversarial random inputs, the worst case for INT8. Versus SageAttention 1 at the same INT8 budget the claim is 2–3× lower error, structurally: Sage quantizes text tiles, this doesn't.

Where people get burned

You can't tell if it's on. It logs. Look for [Krea2 TileMix] active (…) on the first generation - that's your only confirmation the kernel is live. A dense fallback: line instead means the original attention ran, with a reason attached: an fp32 input, an attention mask present, a token count under min_tokens, or a non-Krea model (SDXL just warns and passes through).

You stacked it with another attention patch. It deliberately bails if something in Krea's attn1_patch chain installs a mask, so don't run SageAttention and this together and expect a clean read on which one helped.

Windows. Triton there is the classic pain point - without a community wheel, the node just doesn't appear. Silent-absent, not broken.

You expected more. Under ~1024 tokens there's almost nothing to route, and 1.5× on the attention block is a rounding error when your workflow is mostly VAE decode and text encoding. A lever for high-res work, not a miracle.