togethercomputer/StripedHyena-Nous-7B

1# Copyright (c) Together2# This software is distributed under the terms of the Apache License, Version 2.03# Author: Michael Poli4 5from torch import Tensor6from dataclasses import dataclass, field7from typing import Optional8 9 10# https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py11@dataclass12class InferenceParams:13    """Inference parameters that are passed to the main model in order14    to efficienly calculate and store the context during inference."""15 16    max_seqlen: int17    max_batch_size: int18    seqlen_offset: int = 019    batch_size_offset: int = 020    key_value_memory_dict: dict = field(default_factory=dict)21    lengths_per_sample: Optional[Tensor] = None22 23    def reset(self, max_seqlen, max_batch_size):24        self.max_seqlen = max_seqlen25        self.max_batch_size = max_batch_size26        self.seqlen_offset = 027        if self.lengths_per_sample is not None:28            self.lengths_per_sample.zero_()29 30 31@dataclass32class RecurrentInferenceParams:33    """Inference parameters passed to blocks with recurrent mode."""34 35    fir_filter_length: int = 336    state_dim: int = 1637    seqlen_offset: int = 038    fir_state_dict: dict = field(default_factory=dict)39    state_dict: dict = field(default_factory=dict)40 41    def reset(self):42        self.fir_filter_length = 343        self.state_dim = 1644        self.seqlen_offset = 045