[RFC] [Qwen-Image] WS1 Single-GPU Kernel List

[zhangj1an]

Motivation

It is time to support multimodal train-infer consistency! This task is open to the community for contribution and I will actively review it.

Our 1st model: Qwen-Image

To start with, we will work on Qwen-Image first, it has standard MMDiT structure, with no MoE. Also it is supported by vllm-omni and sglang-omni. Later we can work on bigger models. For your reference, this is the architecture breakdown of Qwen-Image, also a detailed walkthrough of the MMDiT module.

Here is a list of single card ops that needs to be written, open to community for contribution.

Kernel table

Status: 🙋 open → ⏳ in progress → 👀 in review → ✅ merged

To claim: put your handle in the GitHub column and open a PR against this file.

Kernel	What it does	fwd / bwd	WS2	GitHub	PR	Status
latent_pack_unpack	Patchify / unpatchify permutation and truncation rule	fwd + bwd	yes	@nodeeeeee		🙋
latent_normalize	(z − mean) * std_inv and its inverse, 16 per-channel constants	fwd + bwd	yes	@JLiu4Coding		⏳
latent_proj_gemm	img_in 64→3072 and proj_out 3072→64, FP32 throughout	fwd + bwd	no	@haoruilee		⏳
txt_in_rmsnorm_linear	Text-side RMSNorm(3584) + Linear 3584→3072	fwd + bwd	no			🙋
timestep_embed_mlp	Sinusoidal (256 ch, scale 1000) + 2-layer MLP → 3072	fwd + bwd	no			🙋
adaln_modulation	chunk(3) into shift/scale/gate, apply norm(x)*(1+scale)+shift	fwd + bwd	yes			🙋
adaln_gate_residual	Gated residual write-back x + gate * sublayer(h)	fwd + bwd	yes	@BruceLoveDecimal		⏳
qkv_bias_gemm	3072 → 3×3072 + bias, one weight set per stream	fwd + bwd	yes (TP col)			🙋
qk_rmsnorm	Per-head RMSNorm on Q/K, head_dim = 128	fwd + bwd	no	NPU: @erfgss		🙋
multi_axis_rope	Axis split [16, 56, 56], text on the grid diagonal	fwd + bwd	yes (SP)	NPU: @erfgss		🙋
joint_attn_qk_gemm	QKᵀ / √d, seq = T_txt + T_img	fwd + bwd	yes			🙋
joint_attn_softmax	Online softmax over the key sequence, fixed tiling	fwd + bwd	yes			🙋
joint_attn_av_gemm	A · V	fwd + bwd	yes			🙋
joint_attention_fused	Fusion of the three above, must match their composition byte for byte	fwd + bwd	yes (Ulysses / Ring)			🙋
attn_out_bias_gemm	Attention output projection 3072→3072 + bias	fwd + bwd	yes (TP row)			🙋
mlp_up_gemm_gelu	3072→12288 + bias + tanh-approximate GELU	fwd + bwd	yes (TP col)			🙋
mlp_down_gemm	12288→3072 + bias	fwd + bwd	yes (TP row)			🙋
flow_sigma_table	Sigma schedule: dynamic shift + terminal stretch	fwd only	yes (must match across ranks)			🙋
flow_sde_step_logp	SDE step prev_mean and logp, FP32 throughout	fwd + bwd (dlogp)	yes			🙋

How to contribute

Please refer to #204 as an example to add a new ops. You must include python, triton, cuda implementation, also binding + registry gating + clean fallback + tests + benchmark + validation environment.

1. Fixed reduction tree, frozen and documented. Addition order never changes with batch size, token count, SM count, or any runtime condition.
2. FP32 accumulators everywhere.
3. No Split-K, Stream-K, or atomic partial accumulation.
4. No fast-math, no __expf / __logf, no TF32, no compiler reassociation.
5. One FP32 → BF16 cast per kernel, at the final output only.
6. Batch invariance. A row's output bytes do not depend on who it is batched with or
   how it is tiled. This is the core acceptance test and the precondition for WS2.
7. Trace records: reduction order, accumulator precision, Split-K / Stream-K state, TF32 state, kernel fingerprint, selected backend.

CUDA is the bit-level reference backend. Other backends must match it byte for byte on the same dataset, or register an explicit tolerance profile with the hardware difference written down. No silent fallback.

Acceptance

Every PR checks off:

• Forward matches the FP32 CPU reference byte for byte.
• Backward matches the fixed FP32 reference; multi-path gradient sums use a fixed order.
• Batch invariance: batch size, launch geometry, and tiling do not change a row's bytes.
• Shapes cover {1024², 1328², 1664×928}.
• Trace records reduction order, accumulator precision, Split-K / Stream-K state, TF32 state, kernel fingerprint.
• FP32 CPU reference and bit-equality harness delivered.
• The kernel's own pinned items have explicit test cases.

WS1 exit criteria

• All 19 rows ✅.
• One MMDiT block assembled from WS1 kernels: forward within the tolerance contract,
  base.grad is None, LoRA dA / dB finite and non-zero.
• Per-layer error curve across 60 layers produced, for locating first divergence.
• Single-step logp reproducible on fixed inputs and byte-equal to the rollout side.

Then WS2 opens: build <op>_parallel for every kernel marked yes, with byte equality against the WS1 version as the acceptance criterion.

CC list

if you are interested just ping below!
@Dnoob @Zhifu-Liu @bitborne @Billy1900 @icenfly @JLiu4Coding @haoruilee @erfgss @nodeeeeee @Chen-BUPT @ryankert01 @BruceLoveDecimal @Flink-ddd

[BruceLoveDecimal]

hi, can i take adaln_gate_residual as a try?

[Flink-ddd]

Awesome! I'll recommend this task not just to existing contributors, but also to external developers looking to contribute to our project.

[JLiu4Coding]

Hi @zhangj1an , I am so exciited for our new plan. Can I take the latent_normalize?

[haoruilee]

Hi @zhangj1an, can I take latent_proj_gemm?

[nodeeeeee]

Hi @zhangj1an, can I take latent_pack_unpack?

[erfgss]

Hi @zhangj1an @Flink-ddd ，can I take qk_rmsnorm and multi_axis_rope in Ascend A2/A3,？