Fix Qwen 3.5 35B RL gibberish output issue and show end-to-end RL convergence

src/maxtext/configs/post_train/rl.yml

@@ -213,6 +213,7 @@ math_verify_num_procs: null
 # ====== Special tokens/templates for GSM8K reasoning ======
 reasoning_start_token: '<reasoning>'
 reasoning_end_token: '</reasoning>'
+reasoning_start_token_in_prompt: false
 solution_start_token: '<answer>'
 solution_end_token: '</answer>'
 data_template_path: 'maxtext/examples/chat_templates/gsm8k_rl.json'

src/maxtext/configs/post_train/rl_gsm8k_qwen35_35b_v5p64.yml

@@ -0,0 +1,159 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+base_config: "rl.yml"
+
+# ====== Model ======
+model_name: qwen3.5-35b-a3b
+tokenizer_path: Qwen/Qwen3.5-35B-A3B
+tokenizer_type: huggingface
+scan_layers: true
+
+# ====== Hardware ======
+# v5p-64 = v5p-8 (4 chips) * 8 nodes
+# Inference: 4 nodes
+# Training: 4 nodes
+trainer_devices_fraction: 0.5
+sampler_devices_fraction: 0.5
+chips_per_vm: 4
+use_pathways: true
+allow_split_physical_axes: true
+
+# Rollout: TP = 4 (across 4 chips).
+use_standalone_converter: false
+rollout_data_parallelism: -1
+rollout_tensor_parallelism: 4
+rollout_expert_parallelism: 1
+
+# Use the MaxText vLLM adapter so actor and rollout share the same Qwen3.5
+# parameter layout. These overrides were previously supplied only by the
+# experiment JobSet.
+vllm_hf_overrides:
+  architectures: ["MaxTextForCausalLM"]
+vllm_additional_config:
+  maxtext_config:
+    model_name: qwen3.5-35b-a3b
+    model_call_mode: inference
+    attention: vllm_rpa
+    allow_split_physical_axes: true
+    log_config: false
+    weight_dtype: bfloat16
+    prefuse_moe_weights: true
+
+# One checkpoint restore plus an in-memory clone avoids reading/converting the
+# same model twice during startup.
+load_checkpoint_only_once: true
+
+# ====== GRPO ======
+rl:
+  num_generations: 4
+  num_iterations: 1
+  grpo_beta: 0.08
+  grpo_epsilon: 0.2
+  loss_algo: "grpo"
+  loss_agg_mode: "sequence-mean-token-mean"
+  use_agentic_rollout: false
+
+# ====== Training Schedule ======
+# GSM8K train yields 934 full batches at batch_size=8 (7,472 examples, with
+# drop_remainder=True). One epoch covers the full training split once.
+batch_size: 8
+num_batches: 934
+num_epoch: 1
+train_fraction: 1.0
+learning_rate_schedule_steps: 934
+
+train_micro_batch_size: 1
+rollout_micro_batch_size: 8
+
+learning_rate: 3e-6
+warmup_steps_fraction: 0.1
+adam_b1: 0.9
+adam_b2: 0.99
+adam_weight_decay: 0.1
+gradient_clipping_threshold: 0.1
+
+log_period: 20
+eval_interval: 100
+
+# ====== Evaluation ======
+# 20 * 32 = 640 held-out GSM8K prompts per eval pass.
+num_test_batches: 20
+eval_batch_size: 32
+# Qwen3.5 is a thinking model; greedy decoding can fall into endless token
+# loops. Match the model's published sampling defaults.
+eval_sampling_strategy: "standard"
+generation_configs:
+  standard:
+    eval_temperature: 1.0
+    eval_top_k: 20
+    eval_top_p: 0.95
+num_eval_passes: 1
+eval_mode: "pass_at_1"
+
+# ====== Rollout / Generation ======
+max_prefill_predict_length: 256
+max_target_length: 1024
+kv_cache_buffer: 256
+
+decode_sampling_temperature: 1.0
+decode_sampling_top_k: 20
+decode_sampling_nucleus_p: 0.95
+
+hbm_utilization_vllm: 0.60
+swap_space_vllm_gb: 2
+max_num_seqs: 32
+max_num_batched_tokens: 16384
+async_scheduling: false
+enable_dp_attention: false
+# The MaxText Qwen3.5 adapter does not save and restore GDN recurrent state in
+# vLLM's block-addressed prefix cache. Reusing only the attention prefix would
+# pair it with unrelated recurrent state and corrupt generation.
+enable_prefix_caching: false
+
+# ====== Checkpointing ======
+enable_checkpointing: true
+async_checkpointing: false
+checkpoint_period: 250
+max_num_checkpoints_to_keep: 4
+
+# ====== Dataset ======
+dataset_name: "openai/gsm8k"
+eval_dataset_name: "openai/gsm8k"
+train_split: "train"
+eval_split: "test"
+hf_subset: "main"
+data_template_path: "maxtext/examples/chat_templates/qwen35_math_rl.json"
+reasoning_start_token: "<think>"
+reasoning_end_token: "</think>"
+reasoning_start_token_in_prompt: true
+
+# Qwen3.5's chat template opens its native <think> block in the prompt. The
+# model closes it, then emits the task-specific answer block requested above.
+stop_strings: ["</answer>"]
+
+# ====== Reward ======
+reward_exact_answer: 1.0
+reward_white_space_format_match: 1.0
+reward_exact_format_match: 0.1
+reward_partial_format_match: 0.0
+reward_ratio_guess_to_answer_high: 0.0
+reward_ratio_guess_to_answer_low: 0.0
+penalty_incorrect_format: 0.0
+penalty_incorrect_answer: 0.0
+
+math_verify_timeout: 120
+math_verify_num_procs: null
+
+debug: false

src/maxtext/configs/types.py

@@ -2439,6 +2439,10 @@ class RLSpecialTokens(BaseModel):
 
   reasoning_start_token: str = Field("<reasoning>", description="Token to mark the beginning of a reasoning section.")
   reasoning_end_token: str = Field("</reasoning>", description="Token to mark the end of a reasoning section.")
+  reasoning_start_token_in_prompt: bool = Field(
+      False,
+      description="Whether the chat template prefilled the reasoning start token, so it is absent from the completion.",
+  )
   solution_start_token: str = Field("<answer>", description="Token to mark the beginning of a solution section.")
   solution_end_token: str = Field("</answer>", description="Token to mark the end of a solution section.")
 
@@ -3942,12 +3946,14 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
 class RLConfig(
     LogitsAndLoss,
     Engram,
+    ManifoldConstrainedHyperConnections,
     RematAndOffload,
     Attention,
     Llama4Attention,
     LayoutAndSharding,
     InferenceLayout,
     InferenceGeneral,
+    PrefixCaching,
     Decoding,
     IciParallelism,
     DcnParallelism,

src/maxtext/examples/chat_templates/qwen35_math_rl.json

@@ -0,0 +1,4 @@
+{
+  "SYSTEM_PROMPT": "Solve the problem step by step. After your reasoning, place only the final numerical answer between {solution_start_token} and {solution_end_token}.",
+  "TEMPLATE": "{system_prompt}\n\n{question}"
+}

src/maxtext/integration/vllm/_hybrid_cache.py

@@ -0,0 +1,38 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Hybrid cache-layout helpers for MaxText's vLLM adapter."""
+
+import math
+from typing import Any
+
+
+def build_qwen_gdn_cache_layout(cfg: Any, torch_module: Any):
+  """Returns the shapes, dtypes, and unpadded bytes for a Qwen GDN cache."""
+  key_dim = cfg.gdn_key_head_dim * cfg.gdn_num_key_heads
+  value_dim = cfg.gdn_value_head_dim * cfg.gdn_num_value_heads
+  conv_dim = key_dim * 2 + value_dim
+
+  shapes = (
+      (cfg.gdn_conv_kernel_dim - 1, conv_dim),
+      (cfg.gdn_num_value_heads, cfg.gdn_key_head_dim, cfg.gdn_value_head_dim),
+  )
+  # This is the TPU Inference / upstream vLLM contract regardless of model
+  # weight or attention-KV dtype.
+  dtypes = (torch_module.bfloat16, torch_module.float32)
+  page_size_bytes = sum(
+      math.prod(shape) * torch_module.empty((), dtype=dtype).element_size()
+      for shape, dtype in zip(shapes, dtypes, strict=True)
+  )
+  return shapes, dtypes, page_size_bytes

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -23,6 +23,7 @@
 from jax.sharding import Mesh
 from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE
 from maxtext.configs import pyconfig
+from maxtext.integration.vllm._hybrid_cache import build_qwen_gdn_cache_layout
 from maxtext.utils import lora_utils
 from maxtext.utils import max_logging
 from maxtext.utils import model_creation_utils
@@ -378,7 +379,6 @@ def patch_kv_cache_manager():
     from tpu_inference.runner.kv_cache_manager import KVCacheManager
     from vllm.v1.kv_cache_interface import MambaSpec
     import torch
-    import numpy as np
   except ImportError as e:
     # Gracefully handle missing imports in standard JAX environments (e.g. unit tests on CPU)
     max_logging.log(f"Skipping KVCacheManager patch (tpu_inference or dependencies not installed): {e}")
@@ -414,31 +414,12 @@ def patched_get_kv_cache_spec(self):
     if decoder_block_str in ("qwen3_next", "qwen3_5"):
       interval = cfg.inhomogeneous_layer_cycle_interval
 
-      num_v_heads = cfg.gdn_num_value_heads
-      num_k_heads = cfg.gdn_num_key_heads
-      head_k_dim = cfg.gdn_key_head_dim
-      head_v_dim = cfg.gdn_value_head_dim
-      conv_kernel_size = cfg.gdn_conv_kernel_dim
-
-      key_dim = head_k_dim * num_k_heads
-      value_dim = head_v_dim * num_v_heads
-      conv_dim = key_dim * 2 + value_dim
-
-      conv_state_shape = (conv_kernel_size - 1, conv_dim)
-      recurrent_state_shape = (num_v_heads, head_k_dim, head_v_dim)
-
-      mamba_shapes = (conv_state_shape, recurrent_state_shape)
-
-      torch_dtype = torch.bfloat16
-      if str(cfg.dtype) == "float32":
-        torch_dtype = torch.float32
-      elif str(cfg.dtype) == "float16":
-        torch_dtype = torch.float16
-      mamba_dtypes = (torch_dtype, torch_dtype)
-
-      # Calculate unpadded mamba page size
-      dtype_size = 4 if torch_dtype == torch.float32 else 2
-      unpadded_mamba_page_size = sum(int(np.prod(shape)) * dtype_size for shape in mamba_shapes)
+      # Qwen GDN keeps its short convolution history in BF16, but recurrence is
+      # accumulated and persisted in FP32. Declaring both caches as the model
+      # dtype silently quantizes the recurrent state after every generated token.
+      mamba_shapes, mamba_dtypes, unpadded_mamba_page_size = build_qwen_gdn_cache_layout(
+          cfg, torch
+      )
 
       # Calculate attn_page_size_bytes
       from tpu_inference.layers.common.sharding import ShardingAxisName