Fix Attention static-attribute error

[hengtaoguo]

Description

The Attention module was initializing query_norm and key_norm as None at the end of __init__, which made them static attributes in Flax NNX. When the code conditionally tried to assign RMSNorm or Qwen3NextRMSNorm module instances to these attributes, NNX raised a ValueError because you cannot assign module objects to static attributes.

• This is breaking our Gemma3/Qwen3 XLML tests and decoding utils for these two model families.

This PR fixes the issue: query_norm and key_norm are now properly initialized as either NNX modules or None based on the configuration.

Fixes: b/458142671

Tests

Qwen3-4b:

# Command
python -m MaxText.decode MaxText/configs/base.yml model_name=qwen3-4b tokenizer_path=Qwen/Qwen3-4B load_parameters_path=gs://hengtaoguo-maxtext-logs/checkpoints/qwen3-4b/unscanned/2025-10-27-02-00/0/items per_device_batch_size=1 run_name=ht_test max_prefill_predict_length=16 max_target_length=32 steps=1 async_checkpointing=false scan_layers=false use_multimodal=false prompt='Paris is the' hf_access_token=<hf_token>

# Result
Input `Paris is the` -> ` capital of France, and the capital of France is the seat of government. Therefore,`

Gemma3-4b

# Command
python -m MaxText.decode MaxText/configs/base.yml model_name=gemma3-4b tokenizer_path=assets/tokenizer.gemma3 load_parameters_path=gs://maxtext-gemma/unified/gemma3/4b/unscanned/2025-08-09-01-17/0/items per_device_batch_size=1 run_name=ht_test max_prefill_predict_length=272 max_target_length=300 steps=1 async_checkpointing=false scan_layers=false use_multimodal=true prompt=\'Describe\ image\ \<start_of_image\>\' image_path=\'src/MaxText/test_assets/test_image.jpg\' attention=\'dot_product\'

# Result
Input `<start_of_turn>user
Describe image <start_of_image><end_of_turn>
<start_of_turn>model
` -> `Here's a description of the image:

**Overall Impression:**

The image is a bright, expansive cityscape view of Seattle, Washington,`

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.

[RissyRan]

LGTM

[parambole]

PTAL: #2604