Fix: resolve config mismatch in vllm_decode.py

[chiajunglien]

Description

Fixes a initialization failure when running vllm_decode on a single-host TPU (v6e-8).

Cause of the bug

In PR #4068, the weight_dtype was changed from a hardcoded 'bfloat16' string to config.weight_dtype. However, config.weight_dtype is resolved as a raw JAX DType object within MaxText. When this object is passed, OmegaConf fails to parse it, throwing the error:
MAXTEXT CONFIG ERROR: Value 'Float32DType' is not a supported primitive type

Solution

Convert the JAX DType object to its string representation by using config.weight_dtype.name before passing it into the vLLM config.

Tests

python3 -m maxtext.inference.vllm_decode \
    model_name=${MODEL_NAME} \
    load_parameters_path=${UNSCANNED_CKPT_PATH} \
    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
    hbm_utilization_vllm=0.6 \
    prompt="Suggest some famous landmarks in London." \
    use_chat_template=True scan_layers=false

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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!