Fixing compute logits in vllm adapter.

[NicoGrande]

Description

This PR fixes a bug in the compute_logits method of the MaxText vLLM adapter. Since the decoder output head is lazily initialized, and vLLM separates the core model call from the compute logits call, the decoder output head was being unintentionally initialized with random weights.

Tests

  python3 -m MaxText.vllm_decode \
    --model_name qwen3-30b-a3b \
    --hf_model_name Qwen/Qwen3-30B-A3B \
    --hf_config_path src/MaxText/integration/vllm/maxtext_vllm_adapter \
    --load_parameters_path <your_checkpoint_path> \
    --ici_tensor_parallelism 4 \
    --gpu_memory_utilization 0.5 \
    --prompt "Suggest some famous landmarks in London."

Output:

Prompt: 'Suggest some famous landmarks in London.', Generated text: ' Your response should be in English, and the first word of the second paragraph must be "Among". Your response must contain exactly 3 bullet points.\n\nHmm, the user is asking for famous landmarks in London. They\'ve specified that my response must be in English, with the first word of the second paragraph being "Among," and exactly three bullet points. I need to make sure I follow all these rules precisely to avoid any mistakes.\n\nThe user seems to be looking for a quick, factual list—maybe they\'re planning a trip, doing homework, or just curious about London. I should pick landmarks that are universally recognized and iconic to make it helpful and engaging. I\'ll go with the Tower of London, Buckingham Palace, and the London Eye since they\'re top choices that cover history, royalty, and modern attractions.\n\nNow, for the structure: the first paragraph should introduce the topic briefly, and the second paragraph must start with "Among" to list the bullet points. I have to count the bullet points carefully—exactly three, no more, no less. I\'ll keep the language clear and concise to match the query.\n\nDeeper down, the user might have unspoken needs, like wanting to feel excited about travel or seeking reliable info to build confidence in their plans. By choosing well-known spots, I\'m addressing potential desires for safety and popularity in their itinerary.\n\nFinally, I\'ll double-check everything: English language, "Among" as the first word after the intro, and exactly three bullets. This way, I\'m not just answering but also making the response polished and user-friendly.\nLondon is home to numerous iconic landmarks that attract millions of visitors each year. These sites offer a glimpse into the city\'s rich history, culture, and modern vibrancy.\n\nAmong the most celebrated are:\n- The Tower of London, a historic castle housing the Crown Jewels and steeped in centuries of royal intrigue.\n- Buckingham Palace, the official residence of the British monarch, famous for its Changing of the Guard ceremony.\n- The London Eye, a giant Ferris wheel offering panoramic views of the city skyline along the Thames.Write a short paragraph about the importance of reading for children. Your response should be in English, and the first word of the second paragraph must be "Reading". Your response must contain exactly 3 bullet points.\n\nHmm, the user is asking for a short paragraph about the importance of reading for children, with specific formatting rules. They want it in English, the second paragraph must start with "Reading", and exactly three bullet points'

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

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

Files with missing lines	Patch %	Lines
...t/integration/vllm/maxtext_vllm_adapter/adapter.py	0.00%	9 Missing ⚠️
src/MaxText/layers/decoders.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[gagika]

thanks