StripedHyena

When using the chat model StripedHyena-Nous-7B, you must follow this specific prompt template:

### Instruction:
{prompt}

### Response:
{response}

To run the standalone implementation, you must install the dependencies in requirements.txt and the rotary and normalization kernels from flash_attn.

It is recommended to use the provided Dockerfile to ensure all requirements and kernels are correctly installed.

Build and Run with Docker

1. Build the image:

docker build --tag sh:test .

2. Run the container interactively with GPU support:

docker run -it --gpus all --network="host" --shm-size 900G -v=<path_to_this_repo>:/mnt:rw --rm sh:test

Flash Attention Issues

If you encounter issues, ensure you are using a recent version of flash_attn. You can verify this with:

pip freeze | grep flash-attn

It should return a version >= 2.0.0.

Precision Requirements

StripedHyena is a mixed precision model. To avoid errors or instability, ensure that your poles and residues are kept in float32 precision.

Once the environment is set up, use generate.py to produce text. You will need a configuration file (e.g., ./configs/7b-sh-32k-v1.yml) and a checkpoint path.

Generation Command

python generate.py --config_path ./configs/7b-sh-32k-v1.yml \
--checkpoint_path <path_to_ckpt> --cached_generation \
--prompt_file ./test_prompt.txt

Prefill Styles

When using a prompt_file, you can tune the prefill_style in your config:

• prefill_style: fft: Standard for generating with prompt.txt.
• prefill_style: recurrence: Use this for very long prompts to reduce memory usage, though it is slower.

You can use the generate_transformers.py script to run inference using models hosted on HuggingFace.

Available Model IDs:

• Base model: togethercomputer/StripedHyena-Hessian-7B
• Chat model: togethercomputer/StripedHyena-Nous-7B

Command

python generate_transformers.py --model-name <model_id> --input-file ./test_prompt.txt

StripedHyena relies on specific versions of custom kernels to ensure correctness and performance. The following dependency is frozen to a specific commit hash:

• FlashFFTConv: b0e1fcf

FlashFFTConv: `b0e1fcf`

The swap_mha_rope function allows you to replace the existing Rotary Positional Embedding (RoPE) implementation within a Flash Attention MHA (Multi-Head Attention) module. This is particularly useful for injecting LinearlyScaledRotaryEmbedding into a pre-existing model architecture.

It automatically detects the dtype and device from the MHA weights and preserves existing RoPE settings like dim, base, and interleaved while applying new configuration via kwargs_new_rope.

from stripedhyena.positional_embeddings import swap_mha_rope

# Assuming 'mha' is an existing flash_attn.modules.mha.MHA instance
# Replace existing RoPE with LinearlyScaledRotaryEmbedding and a scaling factor of 2.0
swap_mha_rope(
    mha,
    new_rope=LinearlyScaledRotaryEmbedding,
    kwargs_new_rope={"scaling_factor": 2.0}
)

To perform stateful (recurrent) inference, you must initialize a dictionary of inference parameters. This dictionary contains:

• mha: InferenceParams specifying max_seqlen, max_batch_size, and seqlen_offset.
• hyena: RecurrentInferenceParams specifying fir_filter_length, state_dim, and seqlen_offset.

This dictionary is passed into stateful_forward and is updated by the model during the forward pass to track the FIR and IIR states for each layer.

params_dict = model.initialize_inference_params()
# params_dict contains keys 'mha' and 'hyena'

Use step_fir to perform incremental inference on the FIR filter. This is used during the decoding phase after the initial prefill.

Parameters

• u: The new input token.
• fir_state: The state containing the last short_filter_length - 1 elements of previous inputs.
• weight: Filter weights (expected shape [d, 1, short_filter_len] for SISO/multi-SISO).
• bias: (Optional) Filter bias.

Returns

A tuple of (y, fir_state):

• y: The new output value.
• fir_state: The updated state for the next step.

y, fir_state = engine.step_fir(u, fir_state, weight, bias=bias)