Install & Compatibility
Where this runs
tested against v1.0.6 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
SageMoE
✓ from sageattention import SageMoE
✗ from sageattention import SageMoE
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}")
Debug
Known issues
breakingMajor 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.fixUsers 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.
affects: <2.0.0 to >=2.0.0
gotchaThe 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.fixFor 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.
affects: All versions where PyPI is outdated compared to GitHub.
gotchaIncorrect tensor shapes are a common source of runtime errors when working with attention and MoE modules.fixAlways 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.
affects: All versions.
Upgrade
Version history
1.0.6latest on PyPI · released Nov 20, 2024
Audit
Dependencies
torchrequiredCore deep learning framework for tensor operations.
einopsrequiredProvides elegant tensor manipulations.
rotary_embedding_torchrequiredUsed for implementing Rotary Position Embeddings.