SageAttention

Common errors

• ModuleNotFoundError: No module named 'sageattention.sagemoe'

  cause The `sageattention` library, or specific sub-modules, are not installed or are not accessible in your current Python environment.
  fix Ensure the library is installed in your current environment: `pip install sageattention` (for stable PyPI) or `pip install git+https://github.com/thu-ml/SageAttention.git` (for latest GitHub). Verify your Python interpreter is pointing to the correct environment.

• RuntimeError: The size of tensor a (X) must match the size of tensor b (Y) at non-singleton dimension Z

  cause Input tensors provided to SageAttention modules (e.g., `SageMoE`, `TransformerBlock`) do not have the expected dimensions or the `dim` parameter does not match the input's last dimension.
  fix Verify that the `dim` parameter during module initialization correctly matches the last dimension of your input tensor. Ensure batch and sequence dimensions are consistent across inputs if applicable. For example, if `dim=512`, your input should typically be `(batch, seq_len, 512)`.

• TypeError: __init__() got an unexpected keyword argument 'some_old_argument_name'

  cause You are attempting to use constructor arguments or an API from a previous major version (e.g., pre-2.0.0) with a newer installed version of SageAttention.
  fix This error frequently signals a breaking change in the API. Consult the latest GitHub README or the source code for the specific module to find the correct constructor arguments for your installed version. You may need to adapt your initialization code.

Warnings

• breaking Major architectural changes were introduced in v2.0.0. This update added support for MoE and new models, significantly altering internal structures and potentially public APIs for lower-level components.
  Fix: Users migrating from versions prior to 2.0.0 should review the official GitHub README and examples for v2.0.x. Many classes, function signatures, and import paths may have changed. Update your code to reflect the new API.
• gotcha The PyPI package version (currently 1.0.6) significantly lags behind the latest GitHub releases (currently 2.0.1). This means `pip install sageattention` might not give you the features or fixes shown in the latest documentation or GitHub issues.
  Fix: For the absolute latest features, bug fixes, and the API demonstrated in the main GitHub README, install directly from GitHub: `pip install git+https://github.com/thu-ml/SageAttention.git`. Be aware this might install a pre-release or development version.
• gotcha Incorrect tensor shapes are a common source of runtime errors when working with attention and MoE modules.
  Fix: Always ensure input tensors conform to the expected dimensions, typically `(batch_size, sequence_length, embedding_dimension)`, where `embedding_dimension` must match the `dim` parameter passed to the module's constructor. Consult the module's `__init__` or `forward` method signature for precise requirements.

Install

• pip install sageattention PyPI release (may be older)
• pip install git+https://github.com/thu-ml/SageAttention.git Latest development version from GitHub

Imports

• SageMoE

  from sageattention.sagemoe.moe_layer import SageMoE

• TransformerBlock

  from sageattention.sagemoe.transformer_block import TransformerBlock

• SageAttentionForCausalLM

  from sageattention.models import SageAttentionForCausalLM

  Higher-level model, primarily for causal language modeling.

Quickstart

This quickstart demonstrates how to instantiate and use the core `SageMoE` (Mixture-of-Experts) layer and a `TransformerBlock` which internally uses SageAttention. It initializes dummy input tensors and shows the output shapes.

import torch
from sageattention.sagemoe.moe_layer import SageMoE
from sageattention.sagemoe.transformer_block import TransformerBlock

# Example for SageMoE
# Initialize a Mixture-of-Experts layer
moe_model = SageMoE(dim=512, num_experts=8, top_k=2)
# Create a dummy input tensor
x_moe = torch.randn(1, 10, 512) # (batch_size, sequence_length, embedding_dimension)
# Pass input through the MoE layer
output_moe = moe_model(x_moe)
print(f"SageMoE Output Shape: {output_moe.shape}")

# Example for TransformerBlock
# Initialize a Transformer block with attention and MoE
transformer_block = TransformerBlock(dim=512, heads=8, dim_head=64, ff_mult=4, num_experts=8, top_k=2)
# Create a dummy input tensor
x_transformer = torch.randn(1, 10, 512)
# Pass input through the Transformer block
output_transformer = transformer_block(x_transformer)
print(f"TransformerBlock Output Shape: {output_transformer.shape}")