diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py
index ae117bc8a0..171dd100a0 100644
--- a/src/maxtext/models/qwen3.py
+++ b/src/maxtext/models/qwen3.py
@@ -363,7 +363,7 @@ def scan_body(h, args):
 
   o = o.astype(initial_dtype)
 
-  return o, (final_h if initial_state is not None else None)
+  return o, (final_h.astype(compute_dtype) if initial_state is not None else None)
 
 
 def jax_ar_gated_delta_rule(
@@ -970,6 +970,7 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h):
           next_conv_state = next_conv_state[:orig_cache_batch]
           next_recurrent_state = next_recurrent_state[:orig_cache_batch]
 
+    next_recurrent_state = next_recurrent_state.astype(cfg.dtype)
     if model_mode != MODEL_MODE_TRAIN and active_cache is not None:
       active_cache.update_gdn_states(next_recurrent_state, next_conv_state)  # pyrefly: ignore[bad-argument-type]
 
diff --git a/tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh b/tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh
index 36629d7ee2..792c607581 100644
--- a/tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh
+++ b/tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh
@@ -47,7 +47,7 @@ if [ ! -f "${GOLDEN_LOGITS_DISK_LOCATION}" ]; then
   gcloud storage cp ${GOLDEN_LOGITS_PATH} ${GOLDEN_LOGITS_DISK_LOCATION}
 fi
 
-python3 -m tests.utils.forward_pass_logit_checker ${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}/base.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=forward_logits_check load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=true attention=dot_product per_device_batch_size=1 model_name=${MODEL_NAME} max_prefill_predict_length=4 max_target_length=4 async_checkpointing=false sparse_matmul=True ici_fsdp_parallelism=1 ici_expert_parallelism=-1 weight_dtype=float32 dtype=float32 activations_in_float32=true matmul_precision=highest float32_logits=true float32_qk_product=true --golden_logits_path=${GOLDEN_LOGITS_DISK_LOCATION} --atol=1.5 --rtol=1.5 --max_kl_div=0.2
+python3 -m tests.utils.forward_pass_logit_checker ${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}/base.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=forward_logits_check load_parameters_path=${SCANNED_CKPT_PATH} scan_layers=true attention=dot_product per_device_batch_size=1 model_name=${MODEL_NAME} max_prefill_predict_length=4 max_target_length=4 async_checkpointing=false sparse_matmul=True megablox=True ici_fsdp_parallelism=-1 ici_expert_parallelism=1 weight_dtype=float32 dtype=float32 activations_in_float32=true matmul_precision=highest float32_logits=true float32_qk_product=true --golden_logits_path=${GOLDEN_LOGITS_DISK_LOCATION} --atol=1.5 --rtol=1.5 --max_kl_div=0.2
 
 # Run pre-training
 python3 -m maxtext.trainers.pre_train.train ${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}/base.yml base_output_directory=${BASE_OUTPUT_PATH} run_name=q35_pre_training model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_type=grain grain_file_type=tfrecord dataset_path=${DATASET_PATH} enable_checkpointing=false attention=flash sparse_matmul=True megablox=True dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=1 steps=5 max_target_length=1024