Repository navigation
[RFC][Qwen-Image][CUDA/ROCm] WS1/WS2 kernel roadmap, ablation matrix and integration plan #386
Description
Goal and scope
Support Qwen-Image train–infer consistency in RL-Kernel. Given identical weights, prompt embeddings and masks, initial latents, sigma schedule, CFG settings and SDE noise, strict training and rollout must produce byte-equal per-step velocities v_t, latents x_t and log-probabilities logp_t.
- Phase A: Base text-to-image MMDiT, starting from pre-encoded prompts and latents. Single-GPU kernels (WS1), distributed execution (WS2), and full SDE trajectories.
- Phase B: Qwen2.5-VL text encoder, VAE and edit variants.
- Phase C: Training/rollout integration, LoRA gradient validation and end-to-end experiments.
Edit-specific conditioning, additional timestep conditioning, layered RoPE, ControlNet and guidance-distilled variants are outside Phase A. Approximate step caches are excluded from strict mode.
Model contract
qwenimage_arch_fingerprint pins the checkpoint and runtime revisions, graph, weight mapping, dtype policy, scheduler, CFG and SDE arithmetic.
| Component | Contract |
|---|---|
| MMDiT | 60 double-stream blocks; hidden size 3072; 24 heads × 128; BF16 checkpoint |
| Streams | Separate image/text projection, modulation and MLP weights; joint non-causal attention over [text, image] |
| Text input | RMSNorm(3584, eps 1e-6) + Linear 3584→3072; padded prompts require a key mask |
| Image input/output | 16-channel latents; 2×2 packing into 64 values/token; projections 64→3072 and 3072→64 |
| Latent layout | Normalize with (z − mean) * std_inv; pack with permute(0, 2, 4, 1, 3, 5) |
| Timestep | sigma input; 256-channel sinusoid, cos first, internal scale 1000; MLP 256→3072→3072 |
| Q/K preprocessing | Learned-gain per-head RMSNorm, eps 1e-6, before RoPE |
| RoPE | Three axes [16, 56, 56], theta 10000; centred spatial indices; text offset max(H/2, W/2) |
| Block modulation | Parameter-free LayerNorm, eps 1e-6; per-stream SiLU + Linear 3072→18432; two (shift, scale, gate) triples |
| Residual | x + gate * sublayer(h) |
| MLP | 3072→12288→3072, with bias and tanh-approximate GELU |
| Final modulation | Image stream only; SiLU + Linear 3072→6144, (scale, shift) order |
| Sigma schedule | linspace(1, 1/N, N); exponential dynamic shift, 0.5 at 256 tokens and 0.9 at 8192; terminal stretch to 0.02, then append 0 |
| CFG | neg + s * (cond − neg), followed by per-token norm(cond) / norm(comb) over 64 channels |
| RL transition | Flow-GRPO SDE conversion, pinned separately from the released ODE scheduler; step and logp computed in FP32 |
| Reference shapes | 1024², 1328² and 1664×928 |
Architecture
Numerical requirements
The existing WS1 numerical standard applies:
- Fixed reduction trees and FP32 accumulators. No uncontracted Split-K, Stream-K, split-KV or atomic accumulation.
- No TF32, fast math or compiler-dependent reassociation. Cast FP32→BF16 only at the declared kernel output boundary.
- Batch invariance: a logical row’s bytes must not change with batch size, batch position, prompt padding, launch geometry or tiling.
- Fused kernels and distributed execution must reproduce the unfused WS1 result byte for byte.
- Preserve sample, CFG branch, stream, position and timestep identity through every layout change.
- Pin the FP32 sigma table and per-sample noise, including generator device, dtype and consumption order. Use a fixed reduction order for
logp. - Unsupported strict configurations fail closed.
Operator accuracy is checked against an independent FP32 CPU reference: byte equality where required by the arithmetic contract, and declared tolerances for transcendental operations. Strict replay and topology invariance require byte equality.
CUDA and ROCm qualify separately within pinned hardware/runtime profiles. Cross-backend equality is measured; any deviation requires an explicit tolerance profile.
Work items
Claim a row with a platform-prefixed handle and open a PR against test-qwenimage. CUDA and ROCm can be claimed independently; other backends are welcome.
Status follows OPEN → IN PROGRESS → IN REVIEW → MERGED. For claimed rows, the status applies to the listed platform; unlisted platform slots remain open.
All differentiable WS1 operators require forward and backward implementations. Land unfused single-GPU operators before fusion or distribution.
Foundation
| Work item | Scope | Claim | PR | Status |
|---|---|---|---|---|
qwenimage_arch_fingerprint |
Pin model, weight map, scheduler, CFG and SDE contract | OPEN | ||
qwenimage_operator_trace |
Record requested/actual backend, kernel identity, arithmetic policy and fallback state | OPEN | ||
qwenimage_trajectory_artifact |
Record replay inputs, per-step tensor hashes and first drift | OPEN |
WS1: single-GPU kernels
| Work item | Scope | Claim | PR | Status |
|---|---|---|---|---|
latent_pack_unpack |
Patchify/unpatchify | CUDA: @nodeeeeee | #410 | IN REVIEW |
latent_normalize |
Per-channel normalization and inverse | CUDA: @JLiu4Coding | IN PROGRESS | |
latent_proj_gemm |
Image input/output projections | CUDA: @haoruilee | IN PROGRESS | |
txt_in_rmsnorm_linear |
Text input normalization and projection | CUDA: @yanbao1217 | IN PROGRESS | |
timestep_embed_mlp |
Sinusoid and timestep MLP | CUDA: @zzf2026-svg | #478 | IN PROGRESS |
adaln_mod_projection |
Per-stream modulation projections | OPEN | ||
adaln_modulation |
LayerNorm, scale and shift | CUDA: @hongleng | IN PROGRESS | |
adaln_gate_residual |
Gated residual | CUDA: @BruceLoveDecimal | IN PROGRESS | |
qkv_bias_gemm |
Per-stream packed QKV + bias | OPEN | ||
qk_rmsnorm |
Learned-gain Q/K RMSNorm | Ascend: @erfgss | #411 | IN REVIEW |
multi_axis_rope |
Image/text three-axis RoPE | Ascend: @erfgss | #411 | IN REVIEW |
joint_attn_qk_gemm |
Joint QKᵀ / √d |
CUDA: @ryankert01 | IN PROGRESS | |
joint_attn_softmax |
Masked softmax with fixed reduction order | CUDA: @BPbruce | #453 | IN PROGRESS |
joint_attn_av_gemm |
Attention-weighted value accumulation | CUDA: @russli1993 | IN PROGRESS | |
joint_attention_fused |
Byte-equal fusion of the three attention operators | OPEN | ||
attn_out_bias_gemm |
Per-stream attention output projection | CUDA: @yanbao1217 | #459 | IN REVIEW |
mlp_up_gemm_gelu |
MLP up projection + tanh GELU | OPEN | ||
mlp_down_gemm |
MLP down projection | @luoyueyuguang | IN PROGRESS | |
final_adaln_out |
Final image-stream modulation | OPEN | ||
cfg_combine_renorm |
CFG combine and per-token renormalization | @ScottNi0809 | IN PROGRESS | |
flow_sigma_table |
Dynamic shift and terminal stretch; forward only | CUDA: @Txy02 | IN PROGRESS | |
flow_sde_step_logp |
FP32 SDE transition and log-probability | @LG-0927 | IN PROGRESS |
WS2: distributed execution
TP shards attention heads and MLP channels. SP shards image rows and their RoPE positions, with text and prompt metadata replicated. CFG-parallel assigns one branch per rank.
Gather the full latent before CFG combine, scheduler step and logp. SP padding must be masked or avoided.
| Work item | Scope | Claim | PR | Status |
|---|---|---|---|---|
tp_qkv_bias_gemm |
Column-parallel QKV for both streams | OPEN | ||
tp_attn_out_bias_gemm |
Row-parallel output projection, canonical reduction tree, bias added once | OPEN | ||
tp_mlp_gemm_gelu |
Column-parallel up/GELU and row-parallel down | OPEN | ||
tp_replicated_modules |
Verify replicated input/output, timestep, modulation and norm modules | OPEN | ||
sp_shard_gather |
Co-shard image rows and positions; gather in fixed order | OPEN | ||
sp_attention_agrs |
All-gather/shared attention core/reduce-scatter baseline | OPEN | ||
sp_attention_ulysses |
Head redistribution through exact all-to-all | OPEN | ||
sp_attention_ring |
Fixed global block and partial-state merge order | OPEN | ||
sp_rowwise_ops |
Verify rowwise operators on image shards | OPEN | ||
sp_arbitrary_length |
Padding invariance, including 6889 image tokens | OPEN | ||
cfg_parallel_invariance |
Batched, sequential and distributed CFG equivalence | OPEN | ||
dist_sigma_noise_identity |
Identical sigma, timestep and sample noise across ranks | OPEN | ||
dist_sde_step_logp |
Gathered SDE step and logp match WS1 |
OPEN |
Reuse deterministic collectives, the canonical TP FFN, CP attention infrastructure and canonical_cp.py for gradient reduction.
Validation and integration
| Work item | Scope | Claim | PR | Status |
|---|---|---|---|---|
ws1_one_mmdit_block |
Both streams, forward/backward and intermediate comparison | OPEN | ||
ws1_full_mmdit |
Full 60-block model and first-drift localization | OPEN | ||
ws1_trajectory_logp |
Full CFG/SDE trajectory | OPEN | ||
ws2_tp2_tp4_tp8 |
TP equivalence against WS1 | OPEN | ||
ws2_sp2_sp4 |
SP equivalence for AG/RS, Ulysses and Ring | OPEN | ||
ws2_tp_sp_cfg |
TP4×SP2, SP2×CFG2 and TP2×SP2×CFG2 | OPEN | ||
ws2_lora_grad_invariance |
Distributed LoRA gradients match single rank | OPEN | ||
qwenimage_ablation_overlay |
Semantic mismatch probes and invariant controls | OPEN | ||
diffusers_runtime_adapter_qwenimage |
Strict operator binding in pinned Diffusers | OPEN | ||
qwenimage_checkpoint_mapper |
Both stream weight maps and LoRA targets | OPEN | ||
qwenimage_training_provider |
Training context and strict operator injection | OPEN | ||
vllm_omni_runtime_adapter_qwenimage |
Strict binding and actual-backend readback | OPEN | ||
sglang_omni_runtime_adapter_qwenimage |
Strict binding and actual-backend readback | OPEN | ||
qwenimage_runtime_modes |
Strict / audit / auto / off | OPEN | ||
qwen25vl_text_contract |
Phase B: tokenization, template, feature extraction and mask | OPEN | ||
vae_contract |
Phase B: VAE encode/decode and normalization | OPEN | ||
edit_pipeline_contract |
Phase B: edit conditioning and positions | OPEN | ||
phase_b_full_chain |
Phase B: prompt/image to decoded image | OPEN | ||
final_e2e_cuda |
CUDA training/rollout experiment | OPEN | ||
final_e2e_rocm |
ROCm training/rollout experiment | OPEN | ||
final_distributed_run |
Full trajectory under combined topologies | OPEN | ||
cuda_perf_closeout |
CUDA optimization after correctness | OPEN | ||
rocm_perf_closeout |
ROCm optimization after correctness | OPEN |
Acceptance
Operator PRs
- Independent FP32 reference, forward/backward correctness and batch-invariance tests.
- Reference shapes plus a short synthetic case; clear handling of unsupported inputs.
- Both stream weight sets covered where applicable.
- Documented reduction order, precision, cast points and supported geometry.
- Actual-backend trace, reproducible validation command and pinned hardware/runtime environment.
- Benchmark and implementation for the claimed platform. See [WS1][kernels] Batch-invariant logprob (CUDA) #204 for the operator PR pattern.
WS1 and WS2
- All WS1 operators pass, followed by one block, the full model and a complete trajectory.
- Strict training recomputation matches rollout
v_tandlogp_t; complete replays matchx_tat every step. - Frozen base weights have
base.grad is None; LoRAdA / dBare finite and non-zero. - TP2/4/8, SP2/4 and combined topologies match WS1, including LoRA gradients.
- Prompt padding, non-divisible image lengths and CFG execution layout preserve bytes.
- First-drift reports identify the operator, layer, denoising step, logical row and owning rank.
- Strict traces contain no undeclared fallback.
Each platform closes against its own pinned profile.
Experiments
Mismatch probes change one factor at a time: latent layout/normalization, timestep encoding, stream weights, sequence order, RoPE, Q/K preprocessing, modulation, GELU, sigma schedule, CFG, step precision, noise or distributed metadata.
Invariant controls cover repeated execution, batching, padding, trajectory replay, training recomputation, TP/SP/CFG layouts and gradient checkpointing.
Run production/production, strict/production, production/strict and strict/strict training–rollout combinations with matched replay inputs. Report velocity, latent and logp drift, KL, first drift, reward and performance.
For Phase B, include sealed text-to-image and edit fixtures. Latent equality remains the hard gate; decoded-image metrics are secondary. Performance reports separate model, scheduler, communication, text encoder and VAE time.
References
CC: @Dnoob @Zhifu-Liu @bitborne @Billy1900 @icenfly @JLiu4Coding @haoruilee @erfgss @nodeeeeee @Chen-BUPT @ryankert01 @BruceLoveDecimal @Flink-ddd
Activity
BruceLoveDecimal commentedon Sep 5, 2026
hi, can i take adaln_gate_residual as a try?
Flink-ddd commentedon Sep 5, 2026
Awesome! I'll recommend this task not just to existing contributors, but also to external developers looking to contribute to our project.
JLiu4Coding commentedon Sep 6, 2026
Hi @zhangj1an , I am so exciited for our new plan. Can I take the latent_normalize?
haoruilee commentedon Sep 8, 2026
Hi @zhangj1an, can I take latent_proj_gemm?
nodeeeeee commentedon Sep 9, 2026
Hi @zhangj1an, can I take
latent_pack_unpack?erfgss commentedon Sep 12, 2026
Hi @zhangj1an @Flink-ddd ，can I take qk_rmsnorm and multi_axis_rope in Ascend A2/A3,？
BPbruce commentedon Sep 15, 2026
Hi @zhangj1an, I'm interested in this part. Can I take
joint_attn_softmax?32 remaining items
yanbao1217 commentedon Oct 10, 2026
Hi @zhangj1an @Flink-ddd, txt_in_rmsnorm_linear (CUDA claim) is ready for review in #506.
Flink-ddd commentedon Oct 10, 2026
Important Notice: Repository Refactor
Hi everyone! The repository refactor has been merged into main and synced to test-qwenimage. Please merge the latest test-qwenimage from RL-Align/RL-Kernel into your development branch, resolve any conflicts, and update your code and tests to follow the new directory structure. Please rerun the relevant tests after updating. Existing PRs can be updated in place and should continue targeting test-qwenimage, there is no need to open new PRs. Thank you.
fusheng-ji commentedon Oct 11, 2026
Hi @zhangj1an @Flink-ddd, I'd like to claim
adaln_mod_projectionandfinal_adaln_out(CUDA: @fusheng-ji). Both are ready for review as two stacked PRs againsttest-qwenimage:adaln_mod_projection) #537adaln_mod_projection: per-stream SiLU + Linear 3072→18432, with forward and backward.final_adaln_out) #538final_adaln_out: image-stream LayerNorm +(scale, shift)modulation, with forward and backward. It reuses [WS1][Qwen-Image][CUDA] AdaLN modulation projection (adaln_mod_projection) #537's projection, so please review [WS1][Qwen-Image][CUDA] AdaLN modulation projection (adaln_mod_projection) #537 first.Both use fixed-order FP32 reductions, with no Split-K, atomics, TF32 or fast-math. Each is checked against an independent CPU FP32/FP64 reference at the existing tolerances, with raw-byte batch-invariance checks on B200.
Before writing the kernels, I measured existing vLLM, SGLang, Megatron, Liger and flash-attention paths for reuse. For the projection, the batch-invariant GEMMs are a viable alternative, and #537 asks which you prefer.
luoyueyuguang commentedon Oct 11, 2026
Hi @zhangj1an @Flink-ddd,
mlp_down_gemm— PR #499 is ready for review. The branch is rebased onto the currenttest-qwenimage(after the layout refactor,16b8117) and no longer conflicts with it:test-qwenimageis an ancestor of the head commit and GitHub reports the PR asMERGEABLE. Awaiting review.mlp_up_gemm_gelu— I would like to take this operators if it is still open; PR #541 is ready for review.Scope: the MMDiT feed-forward up projection of both streams (
img_mlp.net.0/net.1,txt_mlp.net.0/net.1),3072 -> 12288with bias, fused with the tanh-approximate GELU — forward plus backward (dx/dW/db) — with an independent FP32 CPU reference, the cross-platform Triton backend (ROCm/MUSA slot) and the operator page. The row reports two frozen contracts:mlp-up-gemm-gelu-tree(portable fp32 tree; its pre-activation is byte-equal to the fp32 CPU reference) andmlp-up-gemm-gelu-mma(Hopper TMA+wgmma and Triton; byte-identical to each other). Placement follows the refactored layout:rl_engine/reference/gemm,rl_engine/backends/cuda/gemm,rl_engine/backends/shared/triton/gemm,tests/ops/gemm,benchmarks/operators/gemm.Thanks!
port(attn-out): re-port RL-Align#459 onto the refactored repository l…
BPbruce commentedon Oct 12, 2026
Hi @zhangj1an @Flink-ddd,
PR #453 is ready for review. I've also left some questions in the PR comments:
I'd also like to claim
ws1_one_mmdit_blockunder #386.I plan to start with a PyTorch reference implementation of a complete MMDiT Block, then gradually integrate the WS1 operators as they become available, validating the forward/backward passes and intermediate results.
WIP RL-Align#386: add AdaLN modulation WS1 implementations and tests
Dnoob commentedon Oct 13, 2026
@zhangj1an i'd like to take the CUDA slot of
qk_rmsnormDnoob commentedon Oct 14, 2026
@Flink-ddd @zhangj1an @erfgss Upstream Qwen-Image's QK norm has a learned gain: the checkpoint contains
attn.norm_q.weight(shape[128], eps=1e-6), and diffusers builds it with a weight and appliesbf16(x * rsqrt(mean(x²) + eps)) * weight. The reference in #411 is parameter-free.weight(backward needsdx+dg)?feat(qwenimage): add CUDA AdaLN modulation projection
feat(qwenimage): add CUDA final AdaLN output
port(txt-in): re-port RL-Align#506 onto the refactored repository layout
WIP RL-Align#386: add AdaLN modulation WS1 implementations and tests
JLiu4Coding commentedon Oct 26, 2026
Hi @zhangj1an, I would like to take "adaln_mod_projection" and "qkv_bias_gemm". Grateful if I can be assinged. Also I have made PR #577 as an implemenation of latent_normalize. Your any comment will be appreciated.
test(rope): qualify the complete CUDA package on WSL2