Yi30/hunyuan-image3-dit-attn-capture

HunyuanImage-3.0 DiT attention capture — handoff

Real attention inputs (q/k/v + mask) captured from the HunyuanImage-3.0 DiT, for kernel work: choosing/implementing an attention backend that can serve this model. The mask here is structural, not a padding mask, so it is the part that constrains what a kernel can accept.

What's in this bundle

rank0 = tensor-parallel rank 0 of a TP=4 run, so 8 heads per rank (32 total attention heads sharded 4 ways). Heads are the 3rd axis: (batch, seq, heads, dim).

The two dumps

Sequence layout: 1249 prompt tokens + 4096 image latents + 3 trailing specials = 5348. The image span is exactly [1249, 5345); 4096 = 64x64, i.e. 1024/16 downsample.

call0_rank0.npz — prefill (1 pass per request)

All 5348 positions are queried at once, against the full 5348. Mask density 0.7933.

call32_rank0.npz — denoise (7 passes per request)

Note Q and KV lengths differ (4099 vs 5345): only the generated-image tokens (plus 3 specials) are re-queried each step, against the cached prompt KV. Mask density 0.9994.

q/k/v were converted from bf16 to float32 on save. For bit-exact kernel work, note they are float32 here, not the bf16 the model actually feeds.

The mask contract

Verified elementwise against the dumped array — zero mismatches:

mask == tril(i >= j)  OR  (i and j both in [1249, 5345))

A causal base plus one dense rectangular block on the generated-image span.

S = 5348:  causal base          14,303,226 pairs (50.01%)
           image-span block adds 8,386,560 pairs (29.32%)
           span = 4096 tokens = exactly the image latent grid

Consequences for a kernel:

• —`is_causal` alone is insufficient. The model hard-raises if no mask is given (hunyuan_image3_transformer.py:1271, "Hunyuan dense attention requires an attention mask").
• —A 4D bool mask (or equivalent additive bias) is the simplest correct interface.
• —A "causal + full span" decomposition is also exact: the two parts partition query rows, so outputs merge by torch.cat with no log-sum-exp merge.
• —Neither length is block-aligned: 5348 % 128 == 100, and the denoise pair is 4099 / 5345. A padding kernel must keep the mask consistent across the pad.

At denoise only 3 rows are restricted:

row-visibility histogram: {1247: 1, 1248: 1, 1249: 1, 5345: 4096}

The 3 leading rows are special tokens that see only the pre-image text prefix; all 4096 image rows see the full 5345 cached KV.

Using the dumps

import numpy as np

with np.load("out/call0_rank0.npz") as z:
    q, k, v, mask = z["query"], z["key"], z["value"], z["mask"]

# (1, S, H, D) -> (1, H, S, D) for flash/sage-style kernels
q = q.transpose(0, 2, 1, 3).contiguous()
k = k.transpose(0, 2, 1, 3).contiguous()
v = v.transpose(0, 2, 1, 3).contiguous()

# reference (masked SDPA)
import torch
o = torch.nn.functional.scaled_dot_product_attention(
    torch.from_numpy(q), torch.from_numpy(k), torch.from_numpy(v),
    attn_mask=torch.from_numpy(mask), scale=128 ** -0.5,
)

mask is (B, 1, Q, KV) bool, True = may attend. For a no-mask baseline, note that dropping the mask is not a valid approximation here — it changes the text/image attention pattern, not just padding.

Re-capturing

cd <this dir>
PYTHON=/path/to/vllm-omni/.venv/bin/python ./run_capture.sh

run_capture.sh generates a DiT-only deploy config, puts sitecustomize.py on PYTHONPATH so the hook also reaches the spawn-launched diffusion workers, and runs the shared T2I example. Useful env vars:

Re-analyse / re-render:

$PYTHON analyze_capture.py out          # shapes, mask stats
$PYTHON visualize_mask.py out out.png   # PNG (needs PIL, not matplotlib)

HUNYUAN_ATTN_COLLECT_ALL=1 records every call instead of the first 80. Mask summarisation allocates a float32 copy of the S x S bool, so it is only done for rows actually written.

Environment these were captured in

Caveats

• —Captured on CUDA, not XPU. Shapes and the mask contract are model-level and should carry over, but head-per-rank layout depends on the parallel config. At TP=1 you would see 32 heads per rank rather than 8.
• —GQA is already expanded. q_heads == kv_heads == 8 at the kernel because repeat_kv runs in the dense path when allgather_degree == 1. With allgather_degree > 1 KV stays compressed and a GQA-capable kernel is required.
• —Sparsity is not the win. Denoise is 7 of the 8 passes and is 99.94% dense — a causal+span split has almost nothing to skip there. Only prefill (29.3% of pairs in the non-causal block) benefits. Any gain from a quantized/masked kernel here is throughput, not skipped work.