JAXBench: Fix Megablox GMM no-op exploit: use realistic non-underflowing inputs

[charleshong3]

The 11p_Megablox_GMM baseline scaled inputs by limit=1/(M*K) ~= 7.45e-9 (uniform), so each output element was ~2e-13 and underflowed toward zero in bf16. The reference output was therefore ~0 everywhere, which let a no-op kernel (returning zeros, or skipping the grouped matmul entirely) trivially pass np.allclose(atol=1e-2) and report an enormous, meaningless speedup (observed in practice: a degenerate kernel timing ~0.1 ms / ~27x).

This replaces the inputs with small-normal weights/activations (~0.02 scale, output max ~0.14 / mean ~0.02 -- bf16-representable, no K=4096 overflow) and a simulated top-k router with mild popularity bias (router_bias_scale=0.15) so group sizes are realistic and non-uniform. max_expert_size is a static jit arg (added to CONFIG so harnesses that read static_argnums lower the reference correctly). With these inputs a no-op kernel now fails correctness.

Scope: this fixes Megablox's input-underflow mechanism specifically. A separate and distinct no-op class -- KernelBench-derived baselines (18k-50k) that initialize weights/bias to jnp.zeros, yielding an all-zero reference -- is not addressed here and is left to a follow-up.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.