[RFC][BAGEL-7B-MoT][CUDA/ROCm] WS1/WS2 kernel roadmap and integration plan #435

Status: Proposed

Target: post-v0.1.0 community roadmap

Related: #386, #434

Upstream checkpoint pinned for this draft: ByteDance-Seed/BAGEL-7B-MoT@5019f57.

1. Motivation

BAGEL is a unified understanding + generation model. One Qwen2.5-7B-shaped decoder holds two weight sets (Mixture-of-Transformers, MoT) that share self-attention: text and ViT tokens use the understanding expert, VAE-latent tokens use the generation expert. Images are generated by rectified flow inside the same decoder, conditioned on a causally prefilled prompt KV cache.

This adds three things RL-Kernel does not cover today:

1. Token-routed dense weights. Every GEMM and norm picks its weight set per token by a fixed index set. No learned router, but a row's bytes must not depend on the modality mix of the pack.
2. A diffusion policy inside an LLM decoder. Latent tokens query the prompt KV non-causally at every step. The policy logprob is the ODE→SDE transition density, as in [RFC][Qwen-Image][CUDA/ROCm] WS1/WS2 kernel roadmap, ablation matrix and integration plan #386, but with BAGEL's own schedule, sign and timestep convention.
3. Three-branch renormalised CFG. Up to three branches, two scales, a timestep window and a norm-ratio rescale. The reference renorm reduces over the whole packed batch.

The strict objective:

Given the same weights, prompt tokens, initial latents, schedule, CFG parameters and SDE noise, training and rollout must execute the same declared arithmetic contract and produce exactly equal per-step velocities v_t and log-probabilities logp_t in strict mode.

• Phase A - text-to-image: generation-path operators, full-trajectory parity, TP / SP / CFG-parallel invariance.
• Phase B - text-output policy, image-conditioned editing, and think-then-generate. Opens after Phase A closeout.

2. Checkpoint fingerprint

Upstream references:

• https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT
• https://github.com/ByteDance-Seed/Bagel
• vllm-omni vllm_omni/diffusion/models/bagel/
• Flow-GRPO: https://arxiv.org/abs/2505.05470

Upstream runtime status

• vllm-omni serves BAGEL (single-stage and Thinker + DiT) with CFG-parallel, Ulysses / Ring SP, step execution and opt-in trajectory output. It is the initial rollout runtime.
• Upstream BAGEL trains with FSDP and supports freeze_und. No VIME / Megatron provider exists; treat trainer binding as new integration work.

Reference hazards

These are plausible implementation errors that raise nothing:

• config.json says max_latent_size: 32; the weights hold a 64x64 table. At 512² both give in-range, different position ids.
• config.json says timestep_shift: 1.0; the pipeline default is 3.0. At N = 50 that is 30 vs 41 guided steps.
• BAGEL runs N-1 steps; the Lance subclass runs N.
• global CFG renorm reduces over the packed batch; per-request combine exists only in step execution.
• Sequential, CFG-parallel, SP and step-batched paths pack rows differently.
• QK norm runs in FP32, then RoPE and Q/K in BF16.
• Initial noise is drawn on CPU in one path and regenerated on device in another.
• Cached prefill needs a bottom-right causal mask; SDPA is_causal=True is top-left aligned.
• Rollout log_prob comes from a pluggable scheduler; there is no in-tree BAGEL SDE.

3. Numerical contract

The WS1 numerical standard remains authoritative: fixed accumulator precision and reduction order, no Split-K / Stream-K / split-KV / atomics without a contracted merge tree, no TF32 or fast math, casts only at declared boundaries, fail closed on unsupported geometry.

BAGEL-specific rules

bagel_arch_fingerprint lands these as a versioned profile. Changing any of them is a new profile.

• Schedule built in FP64, stored and traced in FP32. s comes from the run config, never config.json.
• SDE step and logp in FP32 with sigma = t, dt = t_i - t_{i+1} > 0. The sign is the opposite of [RFC][Qwen-Image][CUDA/ROCm] WS1/WS2 kernel roadmap, ablation matrix and integration plan #386. logp sums over all latent elements in a fixed order.
• t = 1 first step and t -> 0 last step are separate test cases.
• CFG branch set, scales, interval (compared on FP32 t), renorm type and min are policy. Renorm is per request. cfg_text_scale <= 1 is a separate profile.
• Initial and transition noise use an explicit per-sample generator with fixed device, dtype and consumption order.
• Phase A trains only the gen expert (*_moe_gen, vae2llm, llm2vae, time_embedder), by LoRA or full fine-tune. Frozen und operators still need dX for marker rows.

Required invariances

TP8 is unsupported (28 heads). CUDA and ROCm each require exact parity within a pinned profile; cross-platform byte equality is reported, not assumed.

4. Work-item table

Status: OPEN -> IN PROGRESS -> IN REVIEW -> MERGED.

To claim a task, put your handle in the GitHub column and open a PR. A row is complete only when implementation, independent reference, invariance tests and benchmarks against the native path land together.

5. Recommended claim order

1. bagel_arch_fingerprint
2. mot_fused_add_rmsnorm, mot_routed_gemm, mot_qk_norm_rope, mot_swiglu_mlp
3. gen_joint_attention in parallel with latent_embed_io, cfg_combine_renorm, flow_sde_step_logp
4. full_model_chain
5. tp_invariance, sp_invariance, cfg_parallel_invariance, fsdp_grad_invariance
6. Phase B rows

6. Contribution notes

• Do not claim a routed row with one expert. Both weight sets and the modality-mix test are required.
• Do not reuse [RFC][Qwen-Image][CUDA/ROCm] WS1/WS2 kernel roadmap, ablation matrix and integration plan #386 kernels where the contract differs. Test timestep scaling, dt sign and latent layout explicitly.
• Do not use decoded images as the first comparator. Compare velocities, logp and latents first.
• Strict mode fails closed on batch-coupled renorm, a position grid that disagrees with the table, unpinned RNG, and unsupported TP / SP layouts.

Open questions

1. Gen-expert LoRA or full fine-tune?
2. Transport hash-checked rollout KV to the trainer, or rebuild it?
3. Keep per-request global renorm, or switch to channel?
4. Keep the reference BF16 RoPE, or define an FP32 profile?

If you are interested

Claim one row above and link the implementation PR.