VllmSampler: mapped weight sync into tpu-inference's vLLM (torchax) model

[wenxindongwork]

Summary

VllmSampler.update_params: mapped weight sync into tpu-inference's vLLM (torchax) model implementation (MODEL_IMPL_TYPE=vllm), which is the path native Qwen3.5 runs on.

On that path the runner keeps its weights in a flat {name: jax.Array} dict in an internal, TP/MoE-backend-dependent layout, so writing mapped trainer weights into transformer_state directly is not enough. update_params now:

1. runs the usual preprocess_src_state + transfer_state_with_mappings, but against the runner model's canonical weight specs (VllmModelWrapper.canonical_weight_specs(state): vLLM's TP=1 parameter names/shapes) with reshard_fn=None;
2. hands the resulting canonical arrays to VllmModelWrapper.load_canonical_weights(weights, state), which re-runs tpu-inference's own layout processing and resharding in place.

The nnx path is unchanged. The torchax path is detected by isinstance(runner.state, dict) plus the presence of load_canonical_weights on the runner model.

Also: _vllm_config now merges the caller's additional_config["sharding"]["sharding_strategy"] (e.g. tpu-inference's attn_dp_size) instead of overwriting it, so production sharding settings can be passed through the sampler.

Testing

End to end on v7x-8 with MaxText Qwen3.5-35B-A3B (bf16) → update_params → tpu-inference torchax model started from random weights: all 613 weight tensors bit-identical to the HF-loaded model, greedy generations 4/4 identical to the in-process HF-weight reference. update_params takes ~150 s on the first call (JIT of the preprocess and processing), HBM returns to 28.6 GiB/device.

Also validated with the production Qwen3.5 sampler configuration passed through VllmConfig (TP=8, enable_expert_parallel, enable_dp_attention=True with additional_config={"sharding": {"sharding_strategy": {"attn_dp_size": 4}}}, prefix caching, async scheduling, chunked prefill): 4/4 identical generations after the sync.

Companion PRs: vllm-project/tpu-inference#3477 (load_canonical_weights), AI-Hypercomputer/maxtext#5034 (Qwen3.5 mapping with preprocess_src_state).