Fixed e2e test for Qwen3.5-35b-a3b

[AusarYao28]

Description

Fix the failing end-to-end test for Qwen3.5-35b-a3b. Caused by sparse_matmul option issue and mismatch dtype in qwen 3 implementation

FIXES: b/553630181

Tests

xpk workload create --cluster=v5p-128-bodaborg-europe-west4-b --zone=europe-west4-b --project=cloud-tpu-multipod-dev --tpu-type=v5p-128 --num-slices=1 --priority=medium --workload=yaoyuchen-qwen35-35b-e2e --command="bash tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh" --docker-image=gcr.io/cloud-tpu-multipod-dev/yaoyuchen_google_com-runner:latest

Result: https://pantheon.corp.google.com/kubernetes/workload/overview?project=cloud-tpu-multipod-dev&pageState=(%22workload_list_table%22:(%22f%22:%22%255B%257B_22k_22_3A_22Is%2520system%2520object_22_2C_22t_22_3A11_2C_22v_22_3A_22_5C_22False_~*false_5C_22_22_2C_22i_22_3A_22is_system_22%257D_2C%257B_22k_22_3A_22_22_2C_22t_22_3A10_2C_22v_22_3A_22_5C_22yaoyuchen_5C_22_22%257D%255D%22))

Checklist

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request updates the Qwen3 model implementation to cast recurrent states (final_h and next_recurrent_state) to their appropriate target data types (compute_dtype and cfg.dtype). It also modifies an end-to-end test script for Qwen3.5-35b-a3b to disable sparse_matmul during the forward pass logit check. There are no review comments, and I have no feedback to provide.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[AusarYao28]

Test result for updated code: https://pantheon.corp.google.com/kubernetes/service/europe-west4/v5p-128-bodaborg-europe-west4-b/default/yaoyuchen-qwen35-35b-e2e/overview?project=cloud-tpu-multipod-dev

[parambole]

LGTM