Set check_vma true for moe sparse matmul

[NuojCheng]

Description

This PR enables check_vma=True for the sparse_matmul shard_map, and remove some unnecessaary psum, which are backward transpose of jax.lax.pvary, see jax doc.

This PR currently only supports pure FSDP + megablox kernel.

Following tests are outdated.

Performance Improvements

Test Case 1: Deepseek-test(small) + ZeRO1 + GA

• Metric: Step Time
• Result: Reduced from 1457ms to 1435ms.
• MFU Increase: ~1.5%
• Profiles:
  • Before: xprof
  • After: xprof

Test Case 2: Deepseek3-test Model

• Metric: Step Time
• Result: Reduced from 2238ms to 2091ms.
• MFU Increase: ~7%
• Profiles:
  • Before: xprof (Screenshot)
  • After: xprof (Screenshot)

Test Case 3: Deepseek2-16b on v5p-32 Cluster

• Metric: Step Time
• Result: Reduced from 17045ms to 16808ms.
• Improvement: ~1.41% reduction in step time.
• Details:
  • Before:
    • Command
    • xprof
    • Log
  • After:
    • Command
    • xprof
    • Log

Test Case 4: Deepseek3-671b on tpu7x-512

• Metric: MFU
• Improvement: ~17% improvement
• Details:
  • Before:
    • Command
    • xprof
    • Log
  • After:
    • Command
    • xprof
    • Log

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.