File size: 3,764 Bytes
5ccb4fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from collections.abc import Callable
from functools import partial
import jax
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
from hamiltonzero.model.route_pointer import TreePrefixPointerMHSEA
def _ring_permute_rows_local(values, route_ids, *, axis_size: int):
local_rows = values.shape[0]
lane = jax.lax.axis_index("seq").astype(jax.numpy.int32)
local_ids = jax.lax.dynamic_slice_in_dim(
route_ids,
lane * local_rows,
local_rows,
axis=0,
)
owner = local_ids // local_rows
within_owner = local_ids % local_rows
output0 = jax.numpy.zeros_like(values)
def take_from(panel, origin, output):
selected = panel[within_owner]
mask = owner == origin
while mask.ndim < selected.ndim:
mask = mask[..., None]
return jax.numpy.where(mask, selected, output)
origin0 = lane
output0 = take_from(values, origin0, output0)
permutation = [(i, (i + 1) % axis_size) for i in range(axis_size)]
def step(carry, _):
panel, origin, output = carry
panel = jax.lax.ppermute(panel, "seq", permutation)
origin = (origin - jax.numpy.asarray(1, jax.numpy.int32)) % axis_size
output = take_from(panel, origin, output)
return (panel, origin, output), None
(_, _, output), _ = jax.lax.scan(
step,
(values, origin0, output0),
xs=None,
length=axis_size - 1,
)
return output
def _make_row_permute(mesh: Mesh) -> Callable:
axis_size = int(mesh.shape["seq"])
spec = P("seq", None, None)
if axis_size == 1:
return lambda values, route_ids: values[route_ids]
return jax.shard_map(
partial(_ring_permute_rows_local, axis_size=axis_size),
mesh=mesh,
in_specs=(spec, P()),
out_specs=spec,
check_vma=False,
)
def build_compact_greedy_router(
*,
mesh: Mesh,
decoder_template: TreePrefixPointerMHSEA,
pair_tile_size: int = 128,
) -> Callable:
if tuple(mesh.axis_names) != ("seq",):
raise ValueError("compact greedy router requires a one-dimensional 'seq' mesh")
if int(mesh.shape["seq"]) < 1:
raise ValueError("compact greedy router requires at least one seq lane")
if not isinstance(decoder_template, TreePrefixPointerMHSEA):
raise TypeError("decoder_template must be TreePrefixPointerMHSEA")
if isinstance(pair_tile_size, bool) or int(pair_tile_size) < 1:
raise ValueError("pair_tile_size must be a positive integer")
rep = NamedSharding(mesh, P())
seq_vec = NamedSharding(mesh, P("seq", None))
seq_edge = NamedSharding(mesh, P("seq", None, None))
decoder_rep = jax.tree_util.tree_map(lambda _leaf: rep, decoder_template)
row_permute = _make_row_permute(mesh)
def decode(decoder, h, edge, mask, global_feat, tau, real_mask):
h = jax.lax.with_sharding_constraint(h, seq_vec)
edge = jax.lax.with_sharding_constraint(edge, seq_edge)
perm, logp = decoder._decode_greedy_compact(
h,
edge,
mask,
global_feat=global_feat,
tau=tau,
real_mask=real_mask,
sequence_mesh=mesh,
pair_tile_size=int(pair_tile_size),
row_permute_fn=row_permute,
)
return perm, logp
return jax.jit(
decode,
in_shardings=(
decoder_rep,
seq_vec,
seq_edge,
rep,
rep,
rep,
rep,
),
out_shardings=(rep, rep),
)
__all__ = ["build_compact_greedy_router"]
|