Fix matmul precision and vmem in test scripts

[shuningjin]

Description

1. fix vmem for gpt oss test

• Problem: gpt oss test has RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem in megablox gmm operation during forward pass.
• Fix b/461483388
• Solution: export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=81920'

2. fix matmul precision in multiple tests

• Problem:matmul_precision=float32 used by multiple test scripts was working previously. It now gives error Input should be 'default', 'high' or 'highest' [type=enum, input_value='float32', input_type=str] (see deepseek run: cmd, log). This is due to recent commit.
• Solution: adding float32 to enum value now. Note it is equivalent to matmul_precision=highest, reference

Tests

run forward pass for gpt-oss on v5p-8 (xpk), with export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=81920' and matmul_precision=float32

• cmd: https://paste.googleplex.com/6504516713316352
• log: https://cloudlogging.app.goo.gl/wrjVTB5HPji9uJ8B6

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.