Status: Proposed
Target: post-v0.1.0 community roadmap
Related: #386

Reference Qwen3-8B Dense track: #180, #204, #228, #230, #240, #243, #280, #315, #322, #336, #338, #343, #352, #360, #361, #377.

1. Motivation

RL-Kernel has already closed the dense Qwen3-8B train/rollout consistency loop on CUDA and ROCm, including WS1 operator contracts, WS2 distributed execution, VIME/vLLM integration, ablation, and end-to-end experiments.

google/gemma-4-31B-it is the next dense model target. The goal is the same:

Given the same logical tokens, weights, positions, masks, multimodal inputs, and sampling context, training and rollout must execute the same declared arithmetic contract and produce exactly equal selected-token log-probabilities in strict mode.

This RFC covers CUDA and ROCm, and splits the work into two phases:

• Phase A — text backbone: close WS1, WS2, ablation, vLLM/VIME integration, and end-to-end consistency for the 31B text backbone.
• Phase B — multimodal path: add the Gemma 4 vision encoder, vision soft-token construction/injection, bidirectional vision-token mask semantics, and full VLM train/rollout validation.

The split keeps the text-path consistency work unblocked while VIME-side Gemma 4 multimodal support is brought up.

2. Checkpoint fingerprint

The implementation must be pinned to the exact checkpoint/runtime revision used by each validation run. The current google/gemma-4-31B-it architecture relevant to this RFC is:

Vision-side fingerprint:

Upstream references:

• https://huggingface.co/google/gemma-4-31B-it/blob/main/config.json
• https://github.com/huggingface/transformers/tree/main/src/transformers/models/gemma4
• https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/gemma4.py
• https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/gemma4_mm.py

Upstream runtime status

• vLLM: Gemma 4 inference already exists, including the hybrid text attention path and multimodal model path. RL-Kernel therefore needs strict operator binding and runtime provenance, not a new serving implementation.
• VIME: as of 2026-09-17, the current vime_plugins/models/ tree does not contain a Gemma 4 provider. VIME support is therefore treated as new integration work in this RFC.

3. Architecture map

Figure 1 — Gemma-4-31B-it architecture and RL-Kernel consistency boundaries. The text backbone is closed first; the image path joins through deterministic vision soft-token construction and embedding merge.

The central difference from Qwen3-8B is that Gemma 4 does not have one uniform attention layer. Sliding and global layers have different head geometry, KV ownership, RoPE, mask semantics, cache behavior, and K/V projection semantics. They must therefore have separate runtime identities and validation rows.

One decoder layer

Figure 2 — One Gemma 4 decoder layer. WS1 qualifies each arithmetic boundary; WS2 wraps the same boundary with deterministic TP/CP/SP ownership and communication order.

4. Numerical contract

The existing RL-Kernel WS1 numerical standard remains authoritative. Gemma 4 does not get a weaker consistency definition.

For every strict reduction-bearing boundary:

1. accumulator precision is declared and fixed;
2. reduction/addition order is documented and independent of batch size, batch position, packing, launch geometry, SM/CU count, or runtime occupancy;
3. Split-K / Stream-K / split-KV / cross-CTA atomics are forbidden unless the split and merge tree is itself fixed by the contract;
4. TF32 and fast-math reassociation are disabled in strict mode;
5. casts occur only at declared epilogue points;
6. one logical row is invariant to unrelated batch rows, padding mutation, chunking, and rank placement;
7. unsupported strict geometry fails closed rather than silently falling back to a production/native path;
8. runtime provenance records requested backend, actual backend, kernel/implementation ID, schedule ID, split policy, accumulator dtype, layer type, device, dtype, shape fingerprint, and fallback state.

Platform acceptance:

• CUDA: exact train/rollout selected-token logprob equality under the pinned contract.
• ROCm: exact train/rollout selected-token logprob equality under the pinned contract.
• CUDA ↔ ROCm: cross-platform byte equality is tested and reported, but it is not assumed. Any unavoidable platform-level transcendental/compiler difference must use an explicit named comparator profile with provenance.

5. Kernel and 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 not complete until both the implementation and the row-specific validation contract are included.

Legend for reuse:

• Reuse — existing Qwen3 arithmetic kernel can be reused after Gemma shape qualification.
• Extend — existing infrastructure is reusable, but Gemma-specific arithmetic or semantics are required.
• New — a new strict operator or integration boundary is required.

6. Recommended claim order

For contributors, the recommended order is:

1. gemma4_arch_fingerprint
2. gemma4_operator_trace
3. scaled_token_embedding / tied_embedding_lm_head_contract
4. gemma_rmsnorm / qk_rmsnorm_d256_d512
5. sliding_qkv_gemm / global_q_gemm / global_k_eq_v_projection
6. sliding_rope_d256 / global_proportional_p_rope
7. sliding_attention_core / global_attention_core / hybrid_kv_cache_contract
8. ffn_gate_up_gemm / gelu_tanh_and_mul / ffn_down_gemm
9. tied_lm_head_projection / final_logit_softcap / softcapped_selected_logprob
10. ws1_full_text_chain / ws1_hybrid_attention_regression
11. ws2_tp2 → ws2_tp4 → ws2_tp4_cp2 → ws2_tp4_cp2_sp
12. gemma4_ablation_overlay and Gemma-specific ablation rows
13. vllm_runtime_adapter_gemma4 / vime_provider_gemma4_text / vime_checkpoint_mapping_gemma4
14. Phase B multimodal rows
15. final end-to-end experiments and performance closeout

7. Contribution notes

• A work item is not done until its validation contract is landed together with the implementation.
• Sliding and global attention are two different acceptance items. A contributor should not claim “Gemma attention done” by covering only one side.
• TP8 is a separate WS2 acceptance item because the global path has only 4 KV heads and therefore introduces replicated KV-head ownership.
• Correctness comes before optimization. In particular, the CP sliding-window path should first land a correctness-first deterministic contract before any optimized halo-exchange implementation is accepted.
• Phase B multimodal work must preserve the exact same strict logprob contract used by the text-only path.

if you are interested just ping below!