perf(gemma4): select decode down weights in output blocks

hf_adapters/hf_gemma4_moe.py

@@ -19,6 +19,8 @@
 Both paths share one device-resident expert-weight set.
 """
 
+from contextlib import nullcontext
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -37,6 +39,70 @@
 
 _MOE_TILE = 32  # Decode gather requires tiles with at least two rows.
 
+# Enabled automatically for supported decode calls. The route assignment and
+# the compiler's indexed-selection layout are a measured pair: either alone
+# regressed this workload. Keep the compiler option scoped to decode calls.
+_DECODE_ROUTE_SCHEDULE = True
+
+
+def _decode_route_schedule_enabled(tokens, top_k):
+    """Pair R8 with its compiler capability; other shapes use ordinary decode."""
+    if not _DECODE_ROUTE_SCHEDULE or tokens != 1 or top_k != 8:
+        return False
+    from torch_spyre._inductor import config
+
+    return (
+        hasattr(config, "indexed_selection_consumer_layout")
+        and config.sencores == 32
+        and not config.ignore_work_division_hints
+        and not config.ignore_wsr_hints
+    )
+
+
+# Independent of gate/up reduction blocking. Enabled at width 1024;
+# smaller widths and intermediate-retention experiments are not shipped.
+_DECODE_DOWN_OUTPUT_PANEL = 1024
+
+
+def _decode_down_output_blocks(activated, down_bank, expert_indices, block_size):
+    """Select output-column blocks before indexing the original expert bank.
+
+    Every output still sums its full reduction dimension in one BMM. Concatenate
+    columns in order; never concatenate weights into a full selected slab.
+    """
+    from torch_spyre._inductor.propagate_hints import spyre_hint
+
+    rows, _, intermediate = activated.shape
+    hidden = down_bank.shape[-1]
+    outputs = []
+    for start in range(0, hidden, block_size):
+        width = min(block_size, hidden - start)
+        selected = down_bank[:, :, start : start + width][expert_indices].reshape(
+            rows, intermediate, width
+        )
+        # The indexed load keeps data columns unsplit on this compiler. H:4
+        # needs a proven distributed load or explicit transfer, not a new hint.
+        with spyre_hint(named_dims=["R", "ONE", "H"], work_div={"R": rows, "H": 1}):
+            outputs.append(torch.bmm(activated, selected))
+    return torch.cat(outputs, dim=-1)
+
+
+def _decode_down_panel(hidden, intermediate, dtypes, route_schedule):
+    """Choose the measured block width only for its supported decode shape."""
+    if _DECODE_DOWN_OUTPUT_PANEL not in (None, 1024):
+        raise ValueError("Unsupported decode block width; expected 1024 or None")
+    if (
+        route_schedule
+        and hidden == 2816
+        and intermediate == 704
+        # Both host formats use SEN169_FP16 device arithmetic/storage.
+        # Gemma's checkpoint uses bfloat16; float32 is a different device path.
+        and dtypes[0] in (torch.float16, torch.bfloat16)
+        and all(dtype == dtypes[0] for dtype in dtypes)
+    ):
+        return _DECODE_DOWN_OUTPUT_PANEL
+    return None
+
 
 def _name_prefill_inputs(x, gate, up, down):
     from torch_spyre._inductor.wsr.propagate_named_dims import (
@@ -97,6 +163,13 @@ def _compiled_moe_loop_region(
     from torch_spyre._inductor.propagate_hints import spyre_hint
 
     T, H = x_expert.shape
+    route_schedule = _decode_route_schedule_enabled(T, top_k)
+    down_panel = _decode_down_panel(
+        H,
+        gate_dev.shape[-1],
+        (x_expert.dtype, gate_dev.dtype, up_dev.dtype, down_dev.dtype),
+        route_schedule,
+    )
     probs = _router_probs(
         x_router,
         router_proj_w,
@@ -122,12 +195,24 @@ def _compiled_moe_loop_region(
         )
         gate = gate_dev[expert_indices].reshape(rows, H, intermediate)
         up = up_dev[expert_indices].reshape(rows, H, intermediate)
-        down = down_dev[expert_indices].reshape(rows, intermediate, H)
+        if down_panel is None:
+            down = down_dev[expert_indices].reshape(rows, intermediate, H)
 
-        gate_out = torch.bmm(inputs, gate)
-        up_out = torch.bmm(inputs, up)
+        if route_schedule:
+            with spyre_hint(named_dims=["R", "ONE", "F"], work_div={"R": 8}):
+                gate_out = torch.bmm(inputs, gate)
+                up_out = torch.bmm(inputs, up)
+        else:
+            gate_out = torch.bmm(inputs, gate)
+            up_out = torch.bmm(inputs, up)
         activated = F.gelu(gate_out, approximate="tanh") * up_out
-        expert_out = torch.bmm(activated, down).reshape(T, top_k, H)
+        if down_panel is not None:
+            expert_out = _decode_down_output_blocks(
+                activated, down_dev, expert_indices, down_panel
+            )
+        else:
+            expert_out = torch.bmm(activated, down)
+        expert_out = expert_out.reshape(T, top_k, H)
 
         # Scale on the H-carrying tensor because bare [T,K] products have no
         # legal layout. The widened source gives the gather a physical stick.
@@ -405,15 +490,25 @@ def forward(
                 hidden_states = self._compiled_prefill_ffn(hidden_states, layer_scalar)
             _reset_named_dims()
         else:
-            hidden_states, key_cache, value_cache = self._compiled_decode(
-                hidden_states,
-                selected_freqs,
-                attn_mask,
-                key_cache,
-                value_cache,
-                cache_index,
-                layer_scalar,
+            # Do not enable the slower route-only configuration on an older
+            # compiler. Both the wrapper and region use this same eligibility.
+            decode_config = (
+                optional_spyre_config_patch({"indexed_selection_consumer_layout": True})
+                if _decode_route_schedule_enabled(
+                    hidden_states.shape[0] * hidden_states.shape[1], self._moe_k
+                )
+                else nullcontext()
             )
+            with decode_config:
+                hidden_states, key_cache, value_cache = self._compiled_decode(
+                    hidden_states,
+                    selected_freqs,
+                    attn_mask,
+                    key_cache,
+                    value_cache,
+                    cache_index,
+                    layer_scalar,
+                )
 
         return hidden_states, key_cache, value_cache
 

tests/cpu/_gemma4_decode_perf_helpers.py

@@ -0,0 +1,88 @@
+# Copyright 2026 The Torch-Spyre Authors.
+#
+# 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
+#
+#     http://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.
+
+"""Load shipped functions for dependency-free CPU geometry tests.
+
+These tests stub only Spyre hints/model imports, not tensor operations. They
+do not substitute for compiled layout, placement or numerical device tests.
+"""
+
+import ast
+import contextlib
+import sys
+from pathlib import Path
+from types import ModuleType
+from unittest.mock import patch
+
+import torch
+
+SOURCE = Path(__file__).resolve().parents[2] / "hf_adapters/hf_gemma4_moe.py"
+
+
+def load_functions(names, **overrides):
+    parsed = ast.parse(SOURCE.read_text())
+    namespace = {
+        "torch": torch,
+        "F": torch.nn.functional,
+        "nullcontext": contextlib.nullcontext,
+    }
+    for node in parsed.body:
+        if isinstance(node, ast.Assign):
+            for name in node.targets:
+                if isinstance(name, ast.Name) and name.id.startswith("_DECODE_"):
+                    namespace[name.id] = ast.literal_eval(node.value)
+    namespace.update(overrides)
+    nodes = [
+        n
+        for n in ast.walk(parsed)
+        if isinstance(n, ast.FunctionDef) and n.name in names
+    ]
+    if len(nodes) != len(names):
+        raise AssertionError(f"Expected exactly the requested functions: {names}")
+    nodes += [
+        n
+        for n in parsed.body
+        if isinstance(n, ast.FunctionDef)
+        and n.name in {"_decode_route_schedule_enabled", "_decode_down_panel"}
+        and n.name not in names
+    ]
+    exec(
+        compile(ast.Module(body=nodes, type_ignores=[]), str(SOURCE), "exec"), namespace
+    )
+    return namespace
+
+
+@contextlib.contextmanager
+def recorded_hints(*, compiler_capability=True):
+    calls = []
+
+    @contextlib.contextmanager
+    def spyre_hint(**kwargs):
+        calls.append(kwargs)
+        yield
+
+    module = ModuleType("torch_spyre._inductor.propagate_hints")
+    module.spyre_hint = spyre_hint
+    config = ModuleType("torch_spyre._inductor.config")
+    config.sencores = 32
+    config.ignore_work_division_hints = False
+    config.ignore_wsr_hints = False
+    if compiler_capability:
+        config.indexed_selection_consumer_layout = False
+    inductor = ModuleType("torch_spyre._inductor")
+    inductor.config = config
+    with patch.dict(
+        sys.modules, {module.__name__: module, inductor.__name__: inductor}
+    ):
+        yield calls

tests/cpu/test_gemma4_decode_schedule.py

@@ -0,0 +1,130 @@
+# Copyright 2026 The Torch-Spyre Authors.
+#
+# 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
+#
+#     http://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.
+
+"""Focused tests of the route schedule and its scoped compiler requirement."""
+
+import contextlib
+import unittest
+from types import SimpleNamespace
+
+import torch
+from _gemma4_decode_perf_helpers import load_functions, recorded_hints
+
+
+class GemmaDecodeScheduleTests(unittest.TestCase):
+    def run_region(self, enabled=None, tokens=1, routes=8, compiler_capability=True):
+        ids = torch.tensor([[0, 0, 1, 3, 7, 7, 4, 6]])[:, :routes].expand(tokens, -1)
+        weights = torch.ones(tokens, routes, dtype=torch.float64)
+        namespace = load_functions(
+            {"_compiled_moe_loop_region"},
+            **({} if enabled is None else {"_DECODE_ROUTE_SCHEDULE": enabled}),
+            _router_probs=lambda *args: weights,
+            _topk=lambda *args: (weights, ids),
+        )
+        x = torch.arange(tokens * 4, dtype=torch.float64).reshape(tokens, 4) / 8
+        gate = (torch.arange(8 * 4 * 3).reshape(8, 4, 3) % 7).double() / 16
+        up = gate + 0.25
+        down = gate.transpose(1, 2).contiguous()
+        with recorded_hints(compiler_capability=compiler_capability) as hints:
+            result = namespace["_compiled_moe_loop_region"](
+                x,
+                x,
+                None,
+                None,
+                None,
+                torch.ones(8, 1),
+                gate,
+                up,
+                down,
+                routes,
+                32,
+                2,
+                1e-6,
+            )
+        return result, hints
+
+    def test_default_and_r8_have_equal_cpu_arithmetic_with_repeated_ids(self):
+        ordinary, default_hints = self.run_region(False)
+        routed, routed_hints = self.run_region()
+        self.assertTrue(torch.equal(ordinary, routed))
+        self.assertFalse([h for h in default_hints if "work_div" in h])
+        self.assertEqual(
+            [h["work_div"] for h in routed_hints if "work_div" in h], [{"R": 8}]
+        )
+
+    def test_unsupported_token_and_route_counts_decline(self):
+        for tokens, routes in ((2, 8), (1, 3)):
+            with self.subTest(tokens=tokens, routes=routes):
+                result, hints = self.run_region(None, tokens, routes)
+                ordinary, _ = self.run_region(False, tokens, routes)
+                self.assertTrue(torch.equal(result, ordinary))
+                self.assertFalse([h for h in hints if "work_div" in h])
+                self.assertEqual(result.shape, (tokens, 4))
+
+    def test_old_compiler_retains_ordinary_schedule(self):
+        result, hints = self.run_region(compiler_capability=False)
+        ordinary, _ = self.run_region(False)
+        self.assertTrue(torch.equal(result, ordinary))
+        self.assertFalse([h for h in hints if "work_div" in h])
+
+    def test_compiler_option_is_scoped_and_restored_on_error(self):
+        active, observed = {}, []
+
+        @contextlib.contextmanager
+        def patch_options(options):
+            old = active.copy()
+            active.update(options)
+            try:
+                yield
+            finally:
+                active.clear()
+                active.update(old)
+
+        def decode(*args):
+            observed.append(active.copy())
+            raise RuntimeError("test failure")
+
+        namespace = load_functions(
+            {"forward"},
+            _DECODE_ROUTE_SCHEDULE=True,
+            optional_spyre_config_patch=patch_options,
+        )
+        block = SimpleNamespace(_compiled_decode=decode, _moe_k=8)
+        with recorded_hints(), self.assertRaisesRegex(RuntimeError, "test failure"):
+            namespace["forward"](
+                block, torch.zeros(1, 1, 4), None, None, None, None, None, None
+            )
+        self.assertEqual(observed, [{"indexed_selection_consumer_layout": True}])
+        self.assertEqual(active, {})
+
+    def test_opt_out_does_not_touch_compiler_configuration(self):
+        namespace = load_functions(
+            {"forward"},
+            _DECODE_ROUTE_SCHEDULE=False,
+            optional_spyre_config_patch=lambda options: self.fail(
+                "default patched config"
+            ),
+        )
+        expected = (object(), object(), object())
+        block = SimpleNamespace(_compiled_decode=lambda *args: expected, _moe_k=8)
+        self.assertEqual(
+            namespace["forward"](
+                block, torch.zeros(1, 1, 4), None, None, None, None, None, None
+            ),
+            expected,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()