Introduce SC dense gather reduce kernel for MoE layers

[BirdsOfAFthr]

Description

This PR introduces a SparseCore Pallas kernel implementation for a dense gather-reduce in MoE.

Details & Context

• Problem: Under pure TP, the gather-reduce step (which collects outputs from active experts and performs the top-k weighted reduction) was executed using a native JAX implementation. This native path is highly unperformant on TPU.
• Solution: Similar to what we did for EP, where we have ragged gather reduce kernel to offload the ops to SC, we introduce dense gather reduce kernel for TP.
• Limitations / Future Work:
  • Currently limited to hardware platforms equipped with TPU SparseCore.
  • Fusing the gather reduce ops into gmm kernel, which is agnostic to sharding.
  • Remove the sequential dependency chain of vector addtions.

Tests

• Verified correctness and NaN-handling behavior against a JAX reference implementation across multiple shapes, dtypes (bfloat16 and float32), and reduction group sizes (top-k).
• Ensured the unit tests gracefully fallback.

Commands to Run:

PYTHONPATH=. .venv/bin/python -m unittest tests.kernels.dense_gather_reduce_test