[MaxText][vLLM] Fix weight converter layout transposes for Qwen3.5 MoE

[igorts-git]

Description

Fix projection matrix layout transpositions in Qwen35MaxTextToVLLMConverter and restore missing Qwen3.5 config fields in RLConfig.

Details

1. Weight Converter Transpose Fix (qwen35_moe.py):
   • Transposed linear projection weight matrices (qkv_proj, o_proj, in_proj_qkvz, in_proj_ba, out_proj, shared_expert.gate_up_proj, shared_expert.down_proj) from PyTorch (out, in) convention to vLLM tpu-inference JAX (in, out) layout.
   • Eliminates 190 matrix shape mismatches during weight sync and resolves post-conversion gibberish text generation.
2. Config Field Validation (types.py):
   • Added Qwen3Next, MlaAttention, CompressedAttention, Llama4Attention, SplashAttention, and Multimodal base classes to RLConfig.
   • Resolves Pydantic validation failures for GDN linear attention parameters (gdn_conv_kernel_dim, gdn_key_head_dim, etc.) when running post-training RL workflows.
3. Validation Safety Guard (validate_converter.py):
   • Added defensive key lookup check to skip weight assignment if a converted key is not present in golden_llm_state.

BUG: b/521604343

Tests

• Executed validate_converter.py for qwen3.5-35b-a3b on TPU VM (v5p-8):
  • Status: PASSED (Exit Code 0).
  • Weight Conversion: 573 weights assigned in 1.27s.
  • Text Generation: Prompt 'The capital of France is' produced fluent, coherent output (' Paris.\nThe capital of France is Paris...').
• Unit tests (tests/unit/pyconfig_test.py): PASSED 100% (19/19 passed).

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• 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]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...t/integration/vllm/torchax_converter/qwen35_moe.py	0.00%	7 Missing ⚠️
...ation/vllm/torchax_converter/validate_converter.py	0.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!