Accept inputs shape in Attention/MLA constructor

[bvandermoon]

Description

Accept inputs shape instead of the full inputs in Attention/MLA. This will help since we are about to move Attention initialization to __init__ of the decoder layers. The full inputs are not available at that point.

Example:

This Attention initialization needs to move to __init__ in LlamaDecoderLayer. lnx is generated in __call__, so it won't be available in __init__

Tests

• Base model train gives same perf before/after:

python3 -m MaxText.train MaxText/configs/base.yml  \
    run_name=<run_name> \
    base_output_directory=gs://<gcs_bucket> \
    dataset_type=synthetic \
    steps=10

• Deepseek3-test gives same perf before/after:

python3 -m MaxText.train MaxText/configs/base.yml \
    run_name=bvandermoon-$RANDOM \
    base_output_directory=gs://bvandermoon-multipod-maxtext \
    dataset_type=synthetic \
    steps=10 \
    model_name=deepseek3-test \
    mla_naive_kvcache=False \ # Not needed for train
    max_target_length=256 \
    per_device_batch_size=1 \
    ici_fsdp_parallelism=-1 \
    scan_layers=false \
    weight_dtype=bfloat16 opt_type=sgd

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• 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.

[NuojCheng]

LGTM

[richjames0]

LGTM