monotykamary/LFM2.5-2.6B-RLCD

1"""Isolated hybrid cache branches, including LFM2 convolution state.2 3The deepcopy + reorder approach follows notnotsamuel/LFM2.5-350M-RLCD4(MIT; see THIRD_PARTY_NOTICES.md). Never broadcast writable state views.5"""6 7import copy8 9import torch10 11 12def state_tensors(cache):13    for layer in cache.layers:14        for name in ("keys", "values"):15            value = getattr(layer, name, None)16            if isinstance(value, torch.Tensor):17                yield value18        for name in ("conv_states", "recurrent_states"):19            states = getattr(layer, name, {})20            values = states.values() if isinstance(states, dict) else states21            for value in values:22                if isinstance(value, torch.Tensor) and value.numel():23                    yield value24 25 26def fork_cache(cache, count):27    if type(count) is not int or count <= 0:28        raise ValueError("branch count must be positive")29    tensors = list(state_tensors(cache))30    if not tensors or any(t.shape[0] != 1 for t in tensors):31        raise ValueError("expected an initialized batch-one prefix cache")32    branch = copy.deepcopy(cache)33    branch.reorder_cache(torch.zeros(count, dtype=torch.long, device=tensors[0].device))34    return branch35