diff --git a/hf_adapters/hf_gemma4_moe.py b/hf_adapters/hf_gemma4_moe.py index 9522da4a..17ac88c9 100644 --- a/hf_adapters/hf_gemma4_moe.py +++ b/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 diff --git a/tests/cpu/_gemma4_decode_perf_helpers.py b/tests/cpu/_gemma4_decode_perf_helpers.py new file mode 100644 index 00000000..ed9204e5 --- /dev/null +++ b/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 diff --git a/tests/cpu/test_gemma4_decode_schedule.py b/tests/cpu/test_gemma4_decode_schedule.py new file mode 100644 index 00000000..b0cb63ba --- /dev/null +++ b/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() diff --git a/tests/cpu/test_gemma4_down_blocks.py b/tests/cpu/test_gemma4_down_blocks.py new file mode 100644 index 00000000..db45feb6 --- /dev/null +++ b/tests/cpu/test_gemma4_down_blocks.py @@ -0,0 +1,168 @@ +# 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. + +"""CPU checks of output-column blocking; LX placement is a device gate.""" + +import ast +import unittest +from unittest.mock import patch + +import torch +from _gemma4_decode_perf_helpers import SOURCE, load_functions, recorded_hints + + +class RecordedBank: + def __init__(self, tensor, events): + self.tensor, self.events = tensor, events + self.shape = tensor.shape + + def __getitem__(self, key): + if isinstance(key, tuple): + sliced = self.tensor[key] + self.events.append( + ("slice", key[-1].start, key[-1].stop, tuple(sliced.shape)) + ) + return RecordedBank(sliced, self.events) + selected = self.tensor[key] + self.events.append(("select", tuple(selected.shape))) + return selected + + +class DownBlockTests(unittest.TestCase): + def test_columns_tails_repeated_ids_and_full_reductions(self): + namespace = load_functions({"_decode_down_output_blocks"}) + bank = (torch.arange(8 * 5 * 11).reshape(8, 5, 11) % 17).double() + ids = torch.tensor([[0, 0, 7, 7, 3, 6, 1, 2]]) + activated = torch.arange(8 * 5).reshape(8, 1, 5).double() / 8 + expected = torch.bmm(activated, bank[ids].reshape(8, 5, 11)) + events, shapes = [], [] + bmm = torch.bmm + + def record_bmm(x, y): + shapes.append((tuple(x.shape), tuple(y.shape))) + return bmm(x, y) + + with recorded_hints() as hints, patch.object(torch, "bmm", record_bmm): + result = namespace["_decode_down_output_blocks"]( + activated, RecordedBank(bank, events), ids, 4 + ) + self.assertTrue(torch.equal(result, expected)) + self.assertEqual( + [e[1:3] for e in events if e[0] == "slice"], [(0, 4), (4, 8), (8, 11)] + ) + self.assertEqual( + [e[1] for e in events if e[0] == "select"], + [(1, 8, 5, 4), (1, 8, 5, 4), (1, 8, 5, 3)], + ) + self.assertEqual([s[1][1] for s in shapes], [5, 5, 5]) + self.assertEqual([h["work_div"] for h in hints], [{"R": 8, "H": 1}] * 3) + + def test_full_gather_is_guarded_out_when_output_blocks_are_requested(self): + region = next( + n + for n in ast.parse(SOURCE.read_text()).body + if getattr(n, "name", "") == "_compiled_moe_loop_region" + ) + guards = [ + n + for n in ast.walk(region) + if isinstance(n, ast.If) and ast.unparse(n.test) == "down_panel is None" + ] + self.assertGreaterEqual(len(guards), 1) + full_reads = [ + n + for n in ast.walk(region) + if isinstance(n, ast.Subscript) + and isinstance(n.value, ast.Name) + and n.value.id == "down_dev" + ] + self.assertTrue(full_reads) + # A later gate/up-blocking PR can add another ordinary-down branch. + # Every such read must remain guarded, not just the first one found. + guarded_nodes = [ + node + for guard in guards + for statement in guard.body + for node in ast.walk(statement) + ] + self.assertTrue(all(read in guarded_nodes for read in full_reads)) + + def test_default_region_uses_blocks_without_a_feature_override(self): + # Meta tensors exercise shipped dispatch and all BMM shapes without + # allocating the full bank. Numerical device acceptance is separate. + def make(*shape): + return torch.empty(shape, device="meta", dtype=torch.bfloat16) + + ids = torch.empty((1, 8), device="meta", dtype=torch.int64) + weights = make(1, 8) + ns = load_functions( + {"_compiled_moe_loop_region", "_decode_down_output_blocks"}, + _router_probs=lambda *args: weights, + _topk=lambda *args: (weights, ids), + ) + shapes = [] + bmm = torch.bmm + + def record_bmm(x, y): + shapes.append((tuple(x.shape), tuple(y.shape))) + return bmm(x, y) + + x = make(1, 2816) + with recorded_hints(), patch.object(torch, "bmm", record_bmm): + result = ns["_compiled_moe_loop_region"]( + x, + x, + None, + None, + None, + make(128, 64), + make(128, 2816, 704), + make(128, 2816, 704), + make(128, 704, 2816), + 8, + 32, + 64, + 1e-6, + ) + self.assertEqual(tuple(result.shape), (1, 2816)) + self.assertEqual(len(shapes), 5) + self.assertEqual( + [shape[1][-1] for shape in shapes], [704, 704, 1024, 1024, 768] + ) + self.assertTrue(all(shape[0][0:2] == (8, 1) for shape in shapes)) + + def test_default_panel_and_supported_shape_fallbacks(self): + choose = load_functions({"_decode_down_panel"})["_decode_down_panel"] + fp16 = (torch.float16,) * 4 + self.assertEqual(choose(2816, 704, fp16, True), 1024) + self.assertIsNone(choose(2816, 704, fp16, False)) + self.assertIsNone(choose(1408, 704, fp16, True)) + self.assertIsNone(choose(2816, 768, fp16, True)) + self.assertEqual( + choose(2816, 704, (torch.bfloat16,) * 4, True), + choose(2816, 704, fp16, True), + ) + self.assertIsNone(choose(2816, 704, (torch.float32,) * 4, True)) + self.assertIsNone(choose(2816, 704, (*fp16[:3], torch.float32), True)) + # Each 16-bit format is supported, but mixing formats is not. + self.assertIsNone(choose(2816, 704, (*fp16[:3], torch.bfloat16), True)) + off = load_functions({"_decode_down_panel"}, _DECODE_DOWN_OUTPUT_PANEL=None) + self.assertIsNone(off["_decode_down_panel"](2816, 704, fp16, True)) + invalid = load_functions({"_decode_down_panel"}, _DECODE_DOWN_OUTPUT_PANEL=123) + with self.assertRaisesRegex(ValueError, "Unsupported decode block width"): + invalid["_decode_down_panel"](2816, 704, fp16, True) + + +if __name__ == "__main__": + unittest.main()