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

src/maxtext/configs/types.py

@@ -3786,7 +3786,10 @@ class RLConfig(
     Engram,
     RematAndOffload,
     Attention,
+    MlaAttention,
+    CompressedAttention,
     Llama4Attention,
+    SplashAttention,
     LayoutAndSharding,
     InferenceLayout,
     InferenceGeneral,
@@ -3795,6 +3798,7 @@ class RLConfig(
     DcnParallelism,
     HardwareAndMesh,
     ModelArchitecture,
+    Qwen3Next,
     MoBa,
     # Positional Embeddings
     PositionalEmbedding,
@@ -3806,6 +3810,7 @@ class RLConfig(
     DeepSeekMoE,
     # General MaxText Configs
     RunInfo,
+    TrainingLoop,
     Checkpointing,
     OrbaxStorage,
     DataTypes,
@@ -3816,6 +3821,8 @@ class RLConfig(
     # Debugging and Profiling
     DevelopmentAndDebugging,
     Profiling,
+    # Multimodal Configs
+    Multimodal,
     # For compatibility with trainer in post_train/rl
     RL,
     RLCluster,
@@ -3837,6 +3844,8 @@ class RLConfig(
   enable_dropout: bool = Field(True, description="Enables dropout in the model.")
   dropout_rate: float = Field(0.0, ge=0.0, le=1.0, description="The dropout rate.")
   init_weights_seed: int = Field(0, description="Seed for model weight initialization.")
+  debug_converter: bool = Field(False, description="Enable converter debug checks.")
+  vllm_load_format: str = Field("dummy", description="vLLM load format (dummy or auto).")
   log_period: int = Field(100, description="Frequency (in steps) to log metrics and flush to Tensorboard.")
   hf_access_token: None | str = Field(None, description="Hugging Face API access token.")
   enable_tunix_perf_metrics: bool = Field(

src/maxtext/integration/vllm/torchax_converter/qwen35_moe.py

@@ -109,8 +109,10 @@ def _unstack_rep(x):
               jnp.concatenate([q_tp_shards[t], k_tp_shards[t], v_tp_shards[t]], axis=0) for t in range(tp_size)
           ]
 
-          self.vllm_state[f"{prefix}.self_attn.qkv_proj.weight"] = jnp.concatenate(tp_interleaved, axis=0)
-          self.vllm_state[f"{prefix}.self_attn.o_proj.weight"] = jnp.transpose(o_layers[rep], (1, 0))
+          self.vllm_state[f"{prefix}.self_attn.qkv_proj.weight"] = jnp.transpose(
+              jnp.concatenate(tp_interleaved, axis=0), (1, 0)
+          )
+          self.vllm_state[f"{prefix}.self_attn.o_proj.weight"] = o_layers[rep]
           self.vllm_state[f"{prefix}.self_attn.q_norm.weight"] = qnorm_layers[rep]
           self.vllm_state[f"{prefix}.self_attn.k_norm.weight"] = knorm_layers[rep]
 
@@ -157,7 +159,9 @@ def _unstack_rep(x):
           qkvz_interleaved = [
               jnp.concatenate([q_shards[s], k_shards[s], v_shards[s], z_shards[s]], axis=0) for s in range(tp_size)
           ]
-          self.vllm_state[f"{prefix}.linear_attn.in_proj_qkvz.weight"] = jnp.concatenate(qkvz_interleaved, axis=0)
+          self.vllm_state[f"{prefix}.linear_attn.in_proj_qkvz.weight"] = jnp.transpose(
+              jnp.concatenate(qkvz_interleaved, axis=0), (1, 0)
+          )
 
           # Extract MaxText GDN BA Layout
           t_m_ba = jnp.transpose(ba_layers[rep], (1, 0))
@@ -171,9 +175,11 @@ def _unstack_rep(x):
           a_shards = jnp.split(a, tp_size, axis=0)
 
           ba_interleaved = [jnp.concatenate([b_shards[s], a_shards[s]], axis=0) for s in range(tp_size)]
-          self.vllm_state[f"{prefix}.linear_attn.in_proj_ba.weight"] = jnp.concatenate(ba_interleaved, axis=0)
+          self.vllm_state[f"{prefix}.linear_attn.in_proj_ba.weight"] = jnp.transpose(
+              jnp.concatenate(ba_interleaved, axis=0), (1, 0)
+          )
 
-          self.vllm_state[f"{prefix}.linear_attn.out_proj.weight"] = jnp.transpose(out_layers[rep], (1, 0))
+          self.vllm_state[f"{prefix}.linear_attn.out_proj.weight"] = out_layers[rep]
           self.vllm_state[f"{prefix}.linear_attn.conv1d.weight"] = jnp.transpose(conv_layers[rep], (2, 1, 0))
           self.vllm_state[f"{prefix}.linear_attn.A_log"] = A_log_layers[rep]
           self.vllm_state[f"{prefix}.linear_attn.dt_bias"] = dt_bias_layers[rep]
@@ -263,8 +269,8 @@ def _convert_moe(self, params):
               axis=1,
           ).reshape(-1, sh_g.shape[1])
 
-          self.vllm_state[f"{p}.mlp.shared_expert.gate_up_proj.weight"] = shared_gate_up
-          self.vllm_state[f"{p}.mlp.shared_expert.down_proj.weight"] = sh_down_layers[rep]
+          self.vllm_state[f"{p}.mlp.shared_expert.gate_up_proj.weight"] = jnp.transpose(shared_gate_up, (1, 0))
+          self.vllm_state[f"{p}.mlp.shared_expert.down_proj.weight"] = jnp.transpose(sh_down_layers[rep], (1, 0))
 
           if "shared_expert_gate" in mlp_block:
             self.vllm_state[f"{p}.mlp.shared_expert_gate.weight"] = sh_gate_router_layers[rep]

src/maxtext/integration/vllm/torchax_converter/validate_converter.py

@@ -360,6 +360,9 @@ def validate_converter(argv) -> None:
   # --- Weight assignment ----------------------------------------------------
   with timer(f"Assigning {len(maxtext_vllm_state)} weights to vLLM model"):
     for key, weight in maxtext_vllm_state.items():
+      if key not in golden_llm_state:
+        logging.warning("Key %s not in golden_llm_state, skipping direct assignment", key)
+        continue
       weight_array = weight.value if hasattr(weight, "value") else weight
       dst_sharding = golden_llm_state[key].sharding
       golden_llm_state[key] = reshard_pytree(weight_array, dst_sharding, donate_input=False, cache_plan=True)