perf(sparsecore): enable dense gather reduce for v6e via FP32 intermediate buffer

tests/kernels/dense_gather_reduce_test.py

@@ -258,11 +258,9 @@ def _fake_tpu_info(self, num_lanes, num_cores=1, num_subcores=1):
 
     # (num_lanes, dtype, reduce_group_size, expected_is_compatible)
     @parameterized.named_parameters(
-        # v6e SparseCore (num_lanes=8). bf16 packing=2: 8//8//2 = 0 -> the
-        # Qwen3-30B-A3B crash -> must fall back.
-        ("v6e_bf16_topk8_degenerate", 8, jnp.bfloat16, 8, False),
-        # Same v6e lanes but f32 (packing=1): 8//8//1 = 1 -> kernel is valid,
-        # must NOT be blocked just because it is v6e.
+        # v6e SparseCore (num_lanes=8), bf16, topk=8: enabled via FP32 output buffer (packing=1)
+        ("v6e_bf16_topk8_fp32_buffer", 8, jnp.bfloat16, 8, True),
+        # Same v6e lanes but f32 (packing=1): 8//8//1 = 1 -> kernel is valid.
         ("v6e_f32_topk8_ok", 8, jnp.float32, 8, True),
         # v6e lanes, bf16, smaller group: 8//4//2 = 1 -> valid.
         ("v6e_bf16_topk4_ok", 8, jnp.bfloat16, 4, True),

tests/kernels/experimental/batched_rpa/test_batched_rpa_tuned_params.py

@@ -17,9 +17,40 @@
 import jax.numpy as jnp
 import pytest
 
+import tpu_inference.kernels.experimental.batched_rpa.tuned_params as tp_module
 from tpu_inference.kernels.experimental.batched_rpa import configs
 from tpu_inference.kernels.experimental.batched_rpa.tuned_params import (
-    TunableParams, TuningKey, get_tuned_params)
+    TunableParams, TuningKey, get_block_sizes_override, get_tuned_params)
+
+
+@pytest.mark.parametrize("case", ["decode", "prefill", "mixed"])
+def test_get_block_sizes_override(case):
+    env_name = f"BATCHED_RPA_{case.upper()}_BLOCK_SIZES"
+    with mock.patch.object(tp_module.envs, env_name,
+                           [2048, 512, 1536, 2, 3]):
+        assert get_block_sizes_override(case, 128) == configs.BlockSizes(
+            bq_sz=2048,
+            bq_c_sz=512,
+            bkv_sz=1536,
+            batch_size=2,
+            n_buffer=3,
+        )
+
+
+@pytest.mark.parametrize(
+    "values,error",
+    [
+        ([1, 1, 128, 8], "exactly five"),
+        ([1, 1, 0, 8, 3], "must be positive"),
+        ([2048, 768, 1536, 2, 3], "divide bq_sz"),
+        ([2048, 512, 1500, 2, 3], "multiple of page_size"),
+    ],
+)
+def test_get_block_sizes_override_validation(values, error):
+    with mock.patch.object(tp_module.envs,
+                           "BATCHED_RPA_MIXED_BLOCK_SIZES", values):
+        with pytest.raises(ValueError, match=error):
+            get_block_sizes_override("mixed", 128)
 
 # ---------------------------------------------------------------------------
 # Fixtures
@@ -124,6 +155,49 @@ def test_get_tuned_params_populated_mapping_hit():
     mock_calc.assert_not_called()
 
 
+def test_decode_bkv_size_env_overrides_only_tuned_kv_tile():
+    import tpu_inference.kernels.experimental.batched_rpa.tuned_params as tp_module
+
+    sliding_model_config = configs.ModelConfigs(
+        num_q_heads=16,
+        num_kv_heads=2,
+        head_dim=128,
+        mask_value=-1e9,
+        sliding_window=1024,
+    )
+    key = TuningKey.from_config(sliding_model_config,
+                                _SERVE_CONFIG,
+                                case='decode')
+    global_key = TuningKey.from_config(_MODEL_CONFIG,
+                                       _SERVE_CONFIG,
+                                       case='decode')
+    mapping = {key: _TUNABLE, global_key: _TUNABLE}
+    tpu_info = mock.Mock(mxu_column_size=128)
+
+    with mock.patch.dict(tp_module.tuned_params_mapping, mapping, clear=True), \
+         mock.patch.object(tp_module.envs,
+                           "BATCHED_RPA_DECODE_SLIDING_BKV_SIZE", 1024), \
+         mock.patch.object(tp_module.pltpu,
+                           "get_tpu_info", return_value=tpu_info):
+        result = get_tuned_params(sliding_model_config,
+                                  _SERVE_CONFIG,
+                                  vmem_limit_bytes=1 << 28,
+                                  case='decode')
+        global_result = get_tuned_params(_MODEL_CONFIG,
+                                         _SERVE_CONFIG,
+                                         vmem_limit_bytes=1 << 28,
+                                         case='decode')
+
+    assert result == configs.BlockSizes(
+        bq_sz=_BLOCK_SIZES.bq_sz,
+        bq_c_sz=_BLOCK_SIZES.bq_c_sz,
+        bkv_sz=1024,
+        batch_size=_BLOCK_SIZES.batch_size,
+        n_buffer=_BLOCK_SIZES.n_buffer,
+    )
+    assert global_result == _BLOCK_SIZES
+
+
 # ---------------------------------------------------------------------------
 # TunableParams.__ge__ and __le__
 # ---------------------------------------------------------------------------

tests/layers/common/test_fused_moe_gmm.py

@@ -0,0 +1,56 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax.numpy as jnp
+import numpy as np
+
+from tpu_inference.layers.common.fused_moe_gmm import (
+    _invert_permutation, _tokens_from_sorted_assignments)
+
+
+def test_sorted_assignments_map_to_source_tokens_without_gather():
+    rng = np.random.default_rng(1234)
+
+    for num_tokens, topk in ((1, 1), (16, 2), (32, 4), (256, 8)):
+        num_assignments = num_tokens * topk
+        sorted_assignments = jnp.asarray(
+            rng.permutation(num_assignments), dtype=jnp.int32)
+
+        token_indices = jnp.arange(num_tokens, dtype=jnp.int32).repeat(topk)
+        expected = token_indices[sorted_assignments]
+        actual = _tokens_from_sorted_assignments(sorted_assignments, topk)
+
+        np.testing.assert_array_equal(actual, expected)
+
+
+def test_invert_permutation_matches_argsort():
+    rng = np.random.default_rng(1234)
+
+    for size in (1, 16, 256, 1024):
+        permutation = jnp.asarray(rng.permutation(size), dtype=jnp.int32)
+
+        actual = _invert_permutation(permutation)
+        expected = jnp.argsort(permutation)
+
+        np.testing.assert_array_equal(actual, expected)
+
+
+def test_invert_permutation_restores_original_order():
+    permutation = jnp.array([1, 3, 0, 2], dtype=jnp.int32)
+    original = jnp.array([10, 20, 30, 40], dtype=jnp.int32)
+    reordered = original[permutation]
+
+    restored = reordered[_invert_permutation(permutation)]
+
+    np.testing.assert_array_equal(restored, original)

tests/layers/jax/quantization/test_unquantized.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from types import SimpleNamespace
+
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -27,9 +29,11 @@
 from tpu_inference.layers.common.sharding import ShardingAxisName
 from tpu_inference.layers.jax.linear import (JaxEinsum, JaxLinear,
                                              JaxMergedColumnParallelLinear)
+from tpu_inference.layers.jax.moe.moe import JaxRoutedExperts
 from tpu_inference.layers.jax.quantization import QuantizeMethodBase
 from tpu_inference.layers.jax.quantization.unquantized import (
-    UnquantizedConfig, UnquantizedMergedLinearMethod)
+    UnquantizedConfig, UnquantizedFusedMoEMethod,
+    UnquantizedMergedLinearMethod)
 
 
 @pytest.fixture
@@ -205,6 +209,79 @@ def test_load_is_deferred_until_all_projections_arrive(self, rngs):
 
 class TestUnquantizedJaxMoe:
 
+    @pytest.mark.parametrize(
+        "axis_names,shape,use_ep,expected",
+        [
+            (("data", "model"), {
+                "data": 1,
+                "model": 8
+            }, True, (("model", None, None), ("model", None, None))),
+            (("expert", "model"), {
+                "expert": 4,
+                "model": 2
+            }, True, ((('expert', 'model'), None, None),
+                      (('expert', 'model'), None, None))),
+            (("expert", "model"), {
+                "expert": 1,
+                "model": 8
+            }, False, ((None, None, "model"), (None, "model", None))),
+        ],
+    )
+    def test_routed_expert_weight_shardings(self, axis_names, shape, use_ep,
+                                            expected):
+        mesh = SimpleNamespace(axis_names=axis_names, shape=shape)
+        assert JaxRoutedExperts._get_weight_shardings(mesh,
+                                                      use_ep) == expected
+
+    @pytest.mark.parametrize("backend,transpose", [
+        (MoEBackend.FUSED_MOE, True),
+        (MoEBackend.GMM_EP, False),
+        (MoEBackend.GMM_TP, False),
+    ])
+    def test_routed_expert_loader_uses_backend_layout(self, backend,
+                                                      transpose):
+        """Fused and GMM kernels consume different expert-weight layouts."""
+        layer = JaxRoutedExperts.__new__(JaxRoutedExperts)
+        layer.prefix = "experts"
+        layer.moe_backend = backend
+        for name, shape in (
+            ("kernel_gating_EDF", (2, 3, 4)),
+            ("kernel_up_proj_EDF", (2, 3, 4)),
+            ("kernel_down_proj_EFD", (2, 4, 3)),
+        ):
+            param = nnx.Param(jnp.zeros(shape))
+            param.set_metadata(_weights_to_load=[None, None])
+            setattr(layer, name, param)
+
+        checkpoint_weight = torch.arange(12).reshape(4, 3)
+        loaded = layer._load_weights(
+            [("experts.0.gate_proj.weight", checkpoint_weight)])
+
+        assert loaded == set()
+        staged = layer.kernel_gating_EDF._weights_to_load[0]
+        expected = (checkpoint_weight.numpy().T
+                    if transpose else checkpoint_weight.numpy())
+        np.testing.assert_array_equal(staged[0], expected)
+
+    def test_fused_postprocessing_waits_for_all_expert_weights(self):
+        """Streaming load must not fuse and delete partially loaded params."""
+        complete = nnx.Param(jnp.zeros((1, 2, 3)))
+        complete.set_metadata(_weights_to_load=[jnp.zeros((1, 2, 3))])
+        incomplete = nnx.Param(jnp.zeros((1, 2, 3)))
+        incomplete.set_metadata(_weights_to_load=[None])
+        layer = SimpleNamespace(
+            moe_backend=MoEBackend.FUSED_MOE,
+            kernel_gating_EDF=complete,
+            kernel_up_proj_EDF=incomplete,
+            kernel_down_proj_EFD=complete,
+        )
+        method = UnquantizedFusedMoEMethod.__new__(
+            UnquantizedFusedMoEMethod)
+
+        assert method.process_weights_after_loading(layer) is False
+        assert hasattr(layer, "kernel_gating_EDF")
+        assert hasattr(layer, "kernel_up_proj_EDF")
+
     @pytest.fixture
     def mesh(self):
         devices = jax.devices()