Showing preview only (1,235K chars total). Download the full file or copy to clipboard to get everything.
Repository: state-spaces/mamba
Branch: main
Commit: a76afbd7dfdc
Files: 94
Total size: 1.2 MB
Directory structure:
gitextract_vi30yky3/
├── .github/
│ ├── scripts/
│ │ ├── build.sh
│ │ ├── check_for_ngc_images.sh
│ │ └── test.sh
│ └── workflows/
│ ├── _build.yml
│ ├── _build_in_container.yml
│ ├── build.yml
│ ├── build_in_container.yml
│ └── publish.yaml
├── .gitignore
├── .gitmodules
├── AUTHORS
├── LICENSE
├── MANIFEST.in
├── README.md
├── benchmarks/
│ └── benchmark_generation_mamba_simple.py
├── csrc/
│ └── selective_scan/
│ ├── reverse_scan.cuh
│ ├── selective_scan.cpp
│ ├── selective_scan.h
│ ├── selective_scan_bwd_bf16_complex.cu
│ ├── selective_scan_bwd_bf16_real.cu
│ ├── selective_scan_bwd_fp16_complex.cu
│ ├── selective_scan_bwd_fp16_real.cu
│ ├── selective_scan_bwd_fp32_complex.cu
│ ├── selective_scan_bwd_fp32_real.cu
│ ├── selective_scan_bwd_kernel.cuh
│ ├── selective_scan_common.h
│ ├── selective_scan_fwd_bf16.cu
│ ├── selective_scan_fwd_fp16.cu
│ ├── selective_scan_fwd_fp32.cu
│ ├── selective_scan_fwd_kernel.cuh
│ ├── static_switch.h
│ └── uninitialized_copy.cuh
├── evals/
│ └── lm_harness_eval.py
├── mamba_ssm/
│ ├── __init__.py
│ ├── distributed/
│ │ ├── __init__.py
│ │ ├── distributed_utils.py
│ │ └── tensor_parallel.py
│ ├── models/
│ │ ├── __init__.py
│ │ ├── config_mamba.py
│ │ └── mixer_seq_simple.py
│ ├── modules/
│ │ ├── __init__.py
│ │ ├── block.py
│ │ ├── mamba2.py
│ │ ├── mamba2_simple.py
│ │ ├── mamba3.py
│ │ ├── mamba_simple.py
│ │ ├── mha.py
│ │ ├── mlp.py
│ │ └── ssd_minimal.py
│ ├── ops/
│ │ ├── __init__.py
│ │ ├── cute/
│ │ │ └── mamba3/
│ │ │ └── mamba3_step_fn.py
│ │ ├── selective_scan_interface.py
│ │ ├── tilelang/
│ │ │ └── mamba3/
│ │ │ ├── mamba3_mimo.py
│ │ │ ├── mamba3_mimo_bwd.py
│ │ │ └── mamba3_mimo_fwd.py
│ │ └── triton/
│ │ ├── __init__.py
│ │ ├── angle_cumsum.py
│ │ ├── k_activations.py
│ │ ├── layer_norm.py
│ │ ├── layernorm_gated.py
│ │ ├── mamba3/
│ │ │ ├── angle_dt.py
│ │ │ ├── mamba3_mimo_rotary_step.py
│ │ │ ├── mamba3_mimo_utils.py
│ │ │ ├── mamba3_siso_bwd.py
│ │ │ ├── mamba3_siso_combined.py
│ │ │ ├── mamba3_siso_fwd.py
│ │ │ ├── mamba3_siso_step.py
│ │ │ └── utils.py
│ │ ├── selective_state_update.py
│ │ ├── softplus.py
│ │ ├── ssd_bmm.py
│ │ ├── ssd_chunk_scan.py
│ │ ├── ssd_chunk_state.py
│ │ ├── ssd_combined.py
│ │ └── ssd_state_passing.py
│ └── utils/
│ ├── __init__.py
│ ├── determinism.py
│ ├── generation.py
│ ├── hf.py
│ └── torch.py
├── pyproject.toml
├── rocm_patch/
│ └── rocm6_0.patch
├── setup.py
├── tests/
│ ├── benchmark_determinism_kernels.py
│ ├── ops/
│ │ ├── cute/
│ │ │ └── test_mamba3_mimo_step.py
│ │ ├── test_selective_scan.py
│ │ ├── tilelang/
│ │ │ └── test_mamba3_mimo.py
│ │ └── triton/
│ │ ├── test_layernorm_gated.py
│ │ ├── test_mamba3_siso.py
│ │ ├── test_selective_state_update.py
│ │ └── test_ssd.py
│ ├── test_determinism.py
│ └── test_generation.py
└── usage.md
================================================
FILE CONTENTS
================================================
================================================
FILE: .github/scripts/build.sh
================================================
#!/bin/bash
set -eoxu pipefail
# We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
# https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
# However this still fails so I am using a newer version of setuptools
pip install setuptools==68.0.0
pip install ninja packaging wheel
export PATH=/usr/local/cuda/bin:/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
# Limit MAX_JOBS otherwise the github runner goes OOM
export MAX_JOBS=2
export MAMBA_FORCE_BUILD="TRUE"
export MAMBA_FORCE_CXX11_ABI=$CXX11_ABI
# 5h timeout since GH allows max 6h and we want some buffer
EXIT_CODE=0
timeout 5h python setup.py bdist_wheel --dist-dir=dist || EXIT_CODE=$?
if [ $EXIT_CODE -eq 0 ]; then
tmpname=cu${WHEEL_CUDA_VERSION}torch${MATRIX_TORCH_VERSION}cxx11abi$CXX11_ABI
wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
fi
echo $EXIT_CODE
================================================
FILE: .github/scripts/check_for_ngc_images.sh
================================================
#!/bin/bash
# Configuration
BASE_IMAGE="nvcr.io/nvidia/pytorch"
TAG_SUFFIX="-py3"
MONTHS_TO_CHECK=7 # Check current month and previous 6 months (total 7)
# Initialize an array to store existing tags
EXISTING_TAGS=()
echo "Checking for existence of the last ${MONTHS_TO_CHECK} NGC PyTorch images: ${BASE_IMAGE}:YY.MM${TAG_SUFFIX}"
echo "---------------------------------------------------------------------"
# Loop through the last N months
for i in $(seq 0 $((MONTHS_TO_CHECK - 1))); do
# Calculate Year and Month for the tag
CURRENT_YEAR=$(date +%Y)
CURRENT_MONTH=$(date +%m)
# Calculate target month and year
TARGET_DATE=$(date -d "$CURRENT_YEAR-$CURRENT_MONTH-01 -$i months" +%y.%m)
# Construct the full image tag and the tag-only string
IMAGE_TAG="${TARGET_DATE}${TAG_SUFFIX}"
FULL_IMAGE="${BASE_IMAGE}:${IMAGE_TAG}"
echo "Checking: ${FULL_IMAGE}"
# Use 'docker manifest inspect' to check for image existence without pulling.
if docker manifest inspect "${FULL_IMAGE}" > /dev/null 2>&1; then
echo "✅ EXISTS: Found."
# Add the tag-only string to the array
EXISTING_TAGS+=("nvcr.io/nvidia/pytorch:${IMAGE_TAG}")
else
echo "❌ MISSING: Not found."
fi
done
echo "---------------------------------------------------------------------"
## JSON Output Generation
# This uses the collected array to build a JSON string.
# 1. Convert the shell array to a newline-separated string.
TAGS_NL_SEP=$(printf "%s\n" "${EXISTING_TAGS[@]}")
# 2. Use jq to read the newline-separated list and format it into a JSON array.
# . | split("\n") | .[:-1] reads the input, splits it by newline, and removes the trailing empty element.
if command -v jq &> /dev/null; then
JSON_STRING=$(echo -e "${TAGS_NL_SEP}" | jq -R -s 'split("\n") | .[:-1]')
echo "Generated JSON String of Existing Tags:"
echo "${JSON_STRING}"
# Optional: Save the JSON string to a variable for further use
# echo "JSON_STRING is now available in the shell if you source this script."
else
echo "WARNING: 'jq' is not installed. Cannot format output as JSON."
echo "Found Tags: ${EXISTING_TAGS[*]}"
fi
echo "---"
echo "Check complete."
echo "${JSON_STRING}" > ngc_images.json
================================================
FILE: .github/scripts/test.sh
================================================
#!/bin/bash
set -exou pipefail
pip install dist/*.whl
python -c "import mamba_ssm; print(mamba_ssm.__version__)"
================================================
FILE: .github/workflows/_build.yml
================================================
name: ~Build wheel template
on:
workflow_call:
inputs:
runs-on:
description: "The runner to use for the build"
required: true
type: string
python-version:
description: "The Python version to use for the build"
required: true
type: string
cuda-version:
description: "The CUDA version to use for the build"
required: true
type: string
torch-version:
description: "The PyTorch version to use for the build"
required: true
type: string
cxx11_abi:
description: "The C++11 ABI to use for the build"
required: true
type: string
upload-to-release:
description: "Upload wheel to this release"
required: false
type: boolean
default: false
release-version:
description: "Upload wheel to this release"
required: false
type: string
defaults:
run:
shell: bash -x -e -u -o pipefail {0}
jobs:
build-wheel:
runs-on: ${{ inputs.runs-on }}
name: Build wheel (${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }})
steps:
- name: Checkout
uses: actions/checkout@v4
with:
ref: ${{ inputs.release-version }}
submodules: recursive
- name: Checkout build scripts
uses: actions/checkout@v4
with:
path: build-scripts/
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: ${{ inputs.python-version }}
- name: Set CUDA and PyTorch versions
run: |
echo "MATRIX_CUDA_VERSION=$(echo ${{ inputs.cuda-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
echo "MATRIX_TORCH_VERSION=$(echo ${{ inputs.torch-version }} | awk -F \. {'print $1 "." $2'})" >> $GITHUB_ENV
echo "WHEEL_CUDA_VERSION=$(echo ${{ inputs.cuda-version }} | awk -F \. {'print $1'})" >> $GITHUB_ENV
echo "MATRIX_PYTHON_VERSION=$(echo ${{ inputs.python-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
- name: Free up disk space
if: ${{ runner.os == 'Linux' }}
# https://github.com/easimon/maximize-build-space/blob/master/action.yml
# https://github.com/easimon/maximize-build-space/tree/test-report
run: |
sudo rm -rf /usr/share/dotnet
sudo rm -rf /opt/ghc
sudo rm -rf /opt/hostedtoolcache/CodeQL
- name: Set up swap space
if: runner.os == 'Linux'
uses: pierotofy/set-swap-space@v1.0
with:
swap-size-gb: 10
- name: Install CUDA ${{ inputs.cuda-version }}
if: ${{ inputs.cuda-version != 'cpu' }}
uses: Jimver/cuda-toolkit@v0.2.30
id: cuda-toolkit
with:
cuda: ${{ inputs.cuda-version }}
linux-local-args: '["--toolkit"]'
# default method is "local", and we're hitting some error with caching for CUDA 11.8 and 12.1
# method: ${{ (inputs.cuda-version == '11.8.0' || inputs.cuda-version == '12.1.0') && 'network' || 'local' }}
method: "network"
sub-packages: '["nvcc"]'
- name: Install PyTorch ${{ inputs.torch-version }}+cu${{ inputs.cuda-version }}
run: |
pip install --upgrade pip
# For some reason torch 2.2.0 on python 3.12 errors saying no setuptools
pip install setuptools==68.0.0
# With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
# AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
pip install typing-extensions==4.12.2
# Pick the highest available PyTorch wheel CUDA version that doesn't exceed system CUDA
export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \
available = { \
'2.6': [118, 124, 126], \
'2.7': [118, 126, 128], \
'2.8': [126, 128, 129], \
'2.9': [126, 128, 130], \
'2.10': [126, 128, 130], \
}[env['MATRIX_TORCH_VERSION']]; \
sys_cuda = int(env['MATRIX_CUDA_VERSION']); \
print(max(v for v in available if v <= sys_cuda))" \
)
if [[ ${{ inputs.torch-version }} == *"dev"* ]]; then
# pip install --no-cache-dir --pre torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}
# Can't use --no-deps because we need cudnn etc.
# Hard-coding this version of pytorch-triton for torch 2.6.0.dev20241001
pip install jinja2
pip install https://download.pytorch.org/whl/nightly/pytorch_triton-3.1.0%2Bcf34004b8a-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-linux_x86_64.whl
pip install --no-cache-dir --pre https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}/torch-${{ inputs.torch-version }}%2Bcu${TORCH_CUDA_VERSION}-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-linux_x86_64.whl
else
pip install --no-cache-dir torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
fi
nvcc --version
python --version
python -c "import torch; print('PyTorch:', torch.__version__)"
python -c "import torch; print('CUDA:', torch.version.cuda)"
python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
shell: bash
- name: Build wheel
id: build_wheel
env:
CXX11_ABI: ${{ inputs.cxx11_abi }}
MATRIX_TORCH_VERSION: ${{ env.MATRIX_TORCH_VERSION}}
WHEEL_CUDA_VERSION: ${{ env.WHEEL_CUDA_VERSION }}
MATRIX_PYTHON_VERSION: ${{ env.MATRIX_PYTHON_VERSION }}
run: |
EXIT_CODE=$(bash build-scripts/.github/scripts/build.sh | tail -n 1)
# Store exit code in GitHub env for later steps
echo "build_exit_code=$EXIT_CODE" | tee -a "$GITHUB_OUTPUT"
exit $EXIT_CODE
- name: Log Built Wheels
run: |
ls dist
- name: Get Release with tag
id: get_current_release
uses: joutvhu/get-release@v1
with:
tag_name: ${{ inputs.release-version }}
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- name: Upload Release Asset
id: upload_release_asset
if: inputs.upload-to-release
uses: actions/upload-release-asset@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
upload_url: ${{ steps.get_current_release.outputs.upload_url }}
asset_path: ./dist/${{env.wheel_name}}
asset_name: ${{env.wheel_name}}
asset_content_type: application/*
================================================
FILE: .github/workflows/_build_in_container.yml
================================================
name: ~Build wheel template
on:
workflow_call:
inputs:
runs-on:
description: "The runner to use for the build"
required: true
type: string
container-image:
description: "Container image"
required: true
type: string
upload-to-release:
description: "Upload wheel to this release"
required: false
type: boolean
default: false
release-version:
description: "Upload wheel to this release"
required: false
type: string
defaults:
run:
shell: bash -x -e -u -o pipefail {0}
jobs:
build-wheel:
runs-on: ${{ inputs.runs-on }}
name: Build wheel (${{ inputs.container-image }})
steps:
- name: Move /var/lib/containerd/
run: |
mkdir -p "${GITHUB_WORKSPACE}/docker/containerd"
sudo mv /var/lib/containerd/ "${GITHUB_WORKSPACE}/docker/containerd"
- name: Move /var/lib/containerd/
run: |
mkdir -p "${GITHUB_WORKSPACE}/docker/docker"
sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker/docker"
- name: Maximize build space
uses: easimon/maximize-build-space@master
with:
root-reserve-mb: 5120
temp-reserve-mb: 32
swap-size-mb: 10240
remove-dotnet: "true"
remove-android: "true"
remove-haskell: "true"
remove-codeql: "true"
build-mount-path: "/var/lib/"
- name: Restore /var/lib/containerd/
run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/containerd/* /var/lib/containerd"
- name: Restore /var/lib/docker/
run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/docker/* /var/lib/docker"
- name: Checkout source
uses: actions/checkout@v4
with:
ref: ${{ inputs.release-version }}
submodules: recursive
- name: Checkout build scripts
uses: actions/checkout@v4
with:
path: build-scripts/
- name: Build
run: |
echo "Free space:"
df -h
- name: Pull the container
run: docker pull ${{ inputs.container-image }}
- name: Set CUDA and PyTorch versions
run: |
cat <<'EOF' >> script.sh
#!/bin/bash
set -eoxu pipefail
echo "MATRIX_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
echo "MATRIX_TORCH_VERSION=$NVIDIA_PYTORCH_VERSION" >> $GITHUB_ENV
echo "WHEEL_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'})" >> $GITHUB_ENV
echo "MATRIX_PYTHON_VERSION=$(python -c "import sys; print('{}.{}'.format(sys.version_info[0], sys.version_info[1]))" | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
echo "CXX11_ABI=$(python -c 'import torch; print(str(torch._C._GLIBCXX_USE_CXX11_ABI).upper())')" >> $GITHUB_ENV
cat $GITHUB_ENV
EOF
docker run \
--rm \
--shm-size=64g \
--workdir /workspace \
--volume $(pwd):/workspace \
--volume $GITHUB_ENV:$GITHUB_ENV \
-e GITHUB_ENV=$GITHUB_ENV \
${{ inputs.container-image }} bash /workspace/script.sh
- name: Build wheel
id: build_wheel
env:
CXX11_ABI: ${{ env.CXX11_ABI }}
MATRIX_TORCH_VERSION: ${{ env.MATRIX_TORCH_VERSION}}
WHEEL_CUDA_VERSION: ${{ env.WHEEL_CUDA_VERSION }}
MATRIX_PYTHON_VERSION: ${{ env.MATRIX_PYTHON_VERSION }}
run: |
EXIT_CODE=$(docker run \
--rm \
--shm-size=64g \
--workdir /workspace \
--volume $(pwd):/workspace \
--volume $GITHUB_ENV:$GITHUB_ENV \
-e PIP_CONSTRAINT= \
-e GITHUB_ENV=$GITHUB_ENV \
-e CXX11_ABI=$CXX11_ABI \
-e MATRIX_TORCH_VERSION=$MATRIX_TORCH_VERSION \
-e WHEEL_CUDA_VERSION=$WHEEL_CUDA_VERSION \
-e MATRIX_PYTHON_VERSION=$MATRIX_PYTHON_VERSION \
${{ inputs.container-image }} bash /workspace/build-scripts/.github/scripts/build.sh | tail -n 1)
- name: Test wheels
run: |
docker run \
--rm \
--shm-size=64g \
--workdir /workspace \
--volume $(pwd):/workspace \
--volume $GITHUB_ENV:$GITHUB_ENV \
-e GITHUB_ENV=$GITHUB_ENV \
${{ inputs.container-image }} bash /workspace/build-scripts/.github/scripts/test.sh
- name: Log Built Wheels
run: |
ls dist
- name: Get Release with tag
id: get_current_release
uses: joutvhu/get-release@v1
with:
tag_name: ${{ inputs.release-version }}
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- name: Upload Release Asset
id: upload_release_asset
if: inputs.upload-to-release
uses: actions/upload-release-asset@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
upload_url: ${{ steps.get_current_release.outputs.upload_url }}
asset_path: ./dist/${{env.wheel_name}}
asset_name: ${{env.wheel_name}}
asset_content_type: application/*
================================================
FILE: .github/workflows/build.yml
================================================
name: Build wheels
on:
workflow_dispatch:
inputs:
runs-on:
description: "The runner to use for the build"
required: true
type: string
default: ubuntu-22.04
python-version:
description: "The Python version to use for the build"
required: true
type: string
cuda-version:
description: "The CUDA version to use for the build"
required: true
type: string
torch-version:
description: "The PyTorch version to use for the build"
required: true
type: string
cxx11_abi:
description: "Enable torch flag C++11 ABI (TRUE/FALSE)"
required: true
type: string
upload-to-release:
description: "Upload wheel to this release"
required: false
type: boolean
default: false
release-version:
description: "Upload wheel to this release"
required: false
type: string
jobs:
build-wheels:
uses: ./.github/workflows/_build.yml
with:
runs-on: ${{ inputs.runs-on || 'ubuntu-22.04' }}
python-version: ${{ inputs.python-version || '3.12' }}
cuda-version: ${{ inputs.cuda-version || '12.9.1' }}
torch-version: ${{ inputs.torch-version || '2.10.0' }}
cxx11_abi: ${{ inputs.cxx11_abi || 'TRUE' }}
upload-to-release: ${{ inputs.upload-to-release || false }}
release-version: ${{ inputs.release-version || 'v2.2.6.post3' }}
================================================
FILE: .github/workflows/build_in_container.yml
================================================
name: Build wheels in a container
on:
workflow_dispatch:
inputs:
runs-on:
description: "The runner to use for the build"
required: true
type: string
default: ubuntu-22.04
container-image:
description: "Container image"
required: true
type: string
upload-to-release:
description: "Upload wheel to this release"
required: false
type: boolean
default: false
release-version:
description: "Release version tag to checkout and upload to"
required: false
type: string
push:
tags-ignore:
- v*
jobs:
get_version:
runs-on: ubuntu-latest
outputs:
version: ${{ steps.get_version.outputs.version }}
steps:
- name: Get version from input or git
id: get_version
run: |
if [ -n "${{ inputs.release-version }}" ]; then
echo "version=${{ inputs.release-version }}" >> $GITHUB_OUTPUT
else
# Get the latest tag from the repo
git clone --filter=blob:none --no-checkout $GITHUB_SERVER_URL/$GITHUB_REPOSITORY.git repo
cd repo
echo "version=$(git describe --tags --abbrev=0)" >> $GITHUB_OUTPUT
fi
shell: bash
check_for_ngc_images:
runs-on: ubuntu-latest
outputs:
images: ${{ steps.check_for_ngc_images.outputs.IMAGES }}
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: Check for NGC PyTorch images
id: check_for_ngc_images
run: |
bash ./.github/scripts/check_for_ngc_images.sh
echo "IMAGES=$(cat ngc_images.json| jq -cr)" >> $GITHUB_OUTPUT
build-wheels:
needs: [get_version, check_for_ngc_images]
uses: ./.github/workflows/_build_in_container.yml
strategy:
fail-fast: false
matrix:
container-image: ${{ fromJson(needs.check_for_ngc_images.outputs.images) }}
with:
runs-on: ${{ inputs.runs-on || 'ubuntu-22.04' }}
container-image: ${{ matrix.container-image }}
upload-to-release: ${{ inputs.upload-to-release || false }}
release-version: ${{ needs.get_version.outputs.version }}
================================================
FILE: .github/workflows/publish.yaml
================================================
# This workflow will:
# - Create a new Github release
# - Build wheels for supported architectures
# - Deploy the wheels to the Github release
# - Release the static code to PyPi
# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
name: Build wheels and deploy
on:
push:
tags:
- v*
jobs:
setup_release:
name: Create Release
runs-on: ubuntu-latest
outputs:
release-version: ${{ steps.extract_branch.outputs.branch }}
steps:
- name: Get the tag version
id: extract_branch
run: echo "branch=${GITHUB_REF#refs/tags/}" >> $GITHUB_OUTPUT
shell: bash
- name: Create Release
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: gh release create ${{ steps.extract_branch.outputs.branch }} --repo $GITHUB_REPOSITORY --title ${{ steps.extract_branch.outputs.branch }} --generate-notes
shell: bash
build_wheels:
name: Build Wheel
needs: setup_release
strategy:
fail-fast: false
matrix:
# Using ubuntu-22.04 instead of 24.04 for more compatibility (glibc). Ideally we'd use the
# manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
os: [ubuntu-22.04, ubuntu-22.04-arm]
python-version: ["3.10", "3.11", "3.12", "3.13"]
torch-version: ["2.6.0", "2.7.1", "2.8.0", "2.9.1", "2.10.0"]
cuda-version: ["11.8.0", "12.9.1", "13.0.1"]
# We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not.
# Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI.
# Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs)
# when building without C++11 ABI and using it on nvcr images.
cxx11_abi: ["FALSE", "TRUE"]
exclude:
# CUDA 11.8 is not supported by PyTorch 2.8+
- torch-version: "2.8.0"
cuda-version: "11.8.0"
- torch-version: "2.9.1"
cuda-version: "11.8.0"
- torch-version: "2.10.0"
cuda-version: "11.8.0"
# CUDA 13.0 is only supported by PyTorch 2.9+
- torch-version: "2.6.0"
cuda-version: "13.0.1"
- torch-version: "2.7.1"
cuda-version: "13.0.1"
- torch-version: "2.8.0"
cuda-version: "13.0.1"
# No aarch64 PyTorch wheels for 2.6.0, or 2.7.1+cu118
- torch-version: "2.6.0"
os: ubuntu-22.04-arm
- torch-version: "2.7.1"
cuda-version: "11.8.0"
os: ubuntu-22.04-arm
# PyTorch 2.7+ pip wheels use CXX11_ABI=1 by default, no need for FALSE
- torch-version: "2.7.1"
cxx11_abi: "FALSE"
- torch-version: "2.8.0"
cxx11_abi: "FALSE"
- torch-version: "2.9.1"
cxx11_abi: "FALSE"
- torch-version: "2.10.0"
cxx11_abi: "FALSE"
uses: ./.github/workflows/_build.yml
with:
runs-on: ${{ matrix.os }}
python-version: ${{ matrix.python-version }}
cuda-version: ${{ matrix.cuda-version }}
torch-version: ${{ matrix.torch-version }}
cxx11_abi: ${{ matrix.cxx11_abi }}
release-version: ${{ needs.setup_release.outputs.release-version }}
upload-to-release: true
check_for_ngc_images:
runs-on: ubuntu-latest
outputs:
images: ${{ steps.check_for_ngc_images.outputs.IMAGES }}
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: Check for NGC PyTorch images
id: check_for_ngc_images
run: |
bash ./.github/scripts/check_for_ngc_images.sh
echo "IMAGES=$(cat ngc_images.json| jq -cr)" | tee -a $GITHUB_OUTPUT
build_ngc_wheels:
name: Build Wheel for NGC PyTorch
needs: [setup_release, check_for_ngc_images]
strategy:
fail-fast: false
matrix:
os: [ubuntu-22.04, ubuntu-22.04-arm]
container-image: ${{ fromJson(needs.check_for_ngc_images.outputs.images) }}
uses: ./.github/workflows/_build_in_container.yml
with:
runs-on: ${{ matrix.os }}
container-image: ${{ matrix.container-image }}
release-version: ${{ needs.setup_release.outputs.release-version }}
upload-to-release: true
publish_package:
name: Publish package
needs: [build_wheels]
if: always() && !cancelled()
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.10"
- name: Install dependencies
run: |
pip install ninja packaging setuptools wheel twine
# We don't want to download anything CUDA-related here
pip install torch --index-url https://download.pytorch.org/whl/cpu
- name: Build core package
env:
MAMBA_SKIP_CUDA_BUILD: "TRUE"
run: |
python setup.py sdist --dist-dir=dist
- name: Deploy
env:
TWINE_USERNAME: "__token__"
TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
run: |
python -m twine upload dist/*
================================================
FILE: .gitignore
================================================
*__pycache__/
*.egg-info/
build/
**.so
*.hip
*_hip.*
================================================
FILE: .gitmodules
================================================
[submodule "3rdparty/lm-evaluation-harness"]
path = 3rdparty/lm-evaluation-harness
url = https://github.com/EleutherAI/lm-evaluation-harness/
================================================
FILE: AUTHORS
================================================
Tri Dao, tri@tridao.me
Albert Gu, agu@andrew.cmu.edu
================================================
FILE: LICENSE
================================================
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright 2023 Tri Dao, Albert Gu
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
================================================
FILE: MANIFEST.in
================================================
recursive-include csrc *
recursive-include csrc *
README.md
================================================
FILE: README.md
================================================
# Mamba
> **Mamba: Linear-Time Sequence Modeling with Selective State Spaces**\
> Albert Gu*, Tri Dao*\
> Paper: https://arxiv.org/abs/2312.00752
> **Transformers are SSMs: Generalized Models and Efficient Algorithms**\
> **Through Structured State Space Duality**\
> Tri Dao*, Albert Gu*\
> Paper: https://arxiv.org/abs/2405.21060
> **Mamba-3: Improved Sequence Modeling using State Space Principles**\
> **Through Structured State Space Duality**\
> Aakash Lahoti*, Kevin Y. Li*, Berlin Chen*, Caitlin Wang*, Aviv Bick, J. Zico Kolter, Tri Dao†, Albert Gu†\
> Paper: https://arxiv.org/abs/2603.15569
## About
Mamba is a new state space model architecture showing promising performance on information-dense data such as language modeling, where previous subquadratic models fall short of Transformers.
It is based on the line of progress on [structured state space models](https://github.com/state-spaces/s4),
with an efficient hardware-aware design and implementation in the spirit of [FlashAttention](https://github.com/Dao-AILab/flash-attention).
## Installation
Install PyTorch first, then:
- [Option] `pip install causal-conv1d>=1.4.0 --no-build-isolation`: an efficient implementation of a simple causal Conv1d layer used inside the Mamba block.
- `pip install mamba-ssm --no-build-isolation`: the core Mamba package.
- `pip install mamba-ssm[causal-conv1d] --no-build-isolation`: To install core Mamba package and causal-conv1d.
`--no-build-isolation` is required so that pip uses your existing CUDA-enabled PyTorch instead of installing torch-cpu in an isolated build environment.
Other requirements:
- Linux
- NVIDIA GPU
- PyTorch 1.12+
- CUDA 11.6+
For AMD cards, see additional prerequisites below.
## Usage
We expose several levels of interface with the Mamba model.
### Selective SSM
Mamba is based on a selective SSM layer, which is the focus of the paper (Section 3; Algorithm 2).
Source: [ops/selective_scan_interface.py](mamba_ssm/ops/selective_scan_interface.py).
### Mamba Block
The main module of this repository is the Mamba architecture block wrapping the selective SSM.
Source: [modules/mamba_simple.py](mamba_ssm/modules/mamba_simple.py).
Usage:
``` python
import torch
from mamba_ssm import Mamba
batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
# This module uses roughly 3 * expand * d_model^2 parameters
d_model=dim, # Model dimension d_model
d_state=16, # SSM state expansion factor
d_conv=4, # Local convolution width
expand=2, # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape
```
### Mamba-2
The Mamba-2 block is implemented at [modules/mamba2.py](mamba_ssm/modules/mamba2.py).
A simpler version is at [modules/mamba2_simple.py](mamba_ssm/modules/mamba2_simple.py)
The usage is similar to Mamba(-1):
``` python
from mamba_ssm import Mamba2
model = Mamba2(
# This module uses roughly 3 * expand * d_model^2 parameters
d_model=dim, # Model dimension d_model
d_state=64, # SSM state expansion factor, typically 64 or 128
d_conv=4, # Local convolution width
expand=2, # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape
```
#### SSD
A minimal version of the inner SSD module (Listing 1 from the Mamba-2 paper) with conversion between "discrete" and "continuous" SSM versions
is at [modules/ssd_minimal.py](mamba_ssm/modules/ssd_minimal.py).
### Mamba Language Model
Finally, we provide an example of a complete language model: a deep sequence model backbone (with repeating Mamba blocks) + language model head.
Source: [models/mixer_seq_simple.py](mamba_ssm/models/mixer_seq_simple.py).
This is an example of how to integrate Mamba into an end-to-end neural network.
This example is used in the generation scripts below.
## Pretrained Models
Pretrained models are uploaded to
[Hugging Face](https://huggingface.co/state-spaces): `mamba-130m`, `mamba-370m`,
`mamba-790m`, `mamba-1.4b`, `mamba-2.8b`, `mamba2-130m`, `mamba2-370m`,
`mamba2-780m`, `mamba2-1.3b`, `mamba2-2.7b`, `transformerpp-2.7b`, `mamba2attn-2.7b`, trained on 300B tokens on the Pile, as well as `mamba-2.8b-slimpj`
(trained on 600B tokens on the SlimPajama dataset).
The models will be autodownloaded by the generation script below.
These models were trained on the [Pile](https://huggingface.co/datasets/EleutherAI/pile), and follow the standard model dimensions described by GPT-3 and followed by many open source models:
| Parameters | Layers | Model dim. |
|------------|--------|------------|
| 130M | 24 | 768 |
| 370M | 48 | 1024 |
| 790M | 48 | 1536 |
| 1.4B | 48 | 2048 |
| 2.8B | 64 | 2560 |
(The layer count of Mamba doubles that of a Transformer with similar size, as two Mamba blocks are needed for each "layer" (MHA block + MLP block) of a Transformer.)
Note: these are base models trained only for 300B tokens, without any form of downstream modification (instruction tuning, etc.).
Performance is expected to be comparable or better than other architectures trained on similar data, but not to match larger or fine-tuned models.
## Evaluations
To run zero-shot evaluations of models (corresponding to Table 3 of the paper),
we use the
[lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)
library.
1. Install `lm-evaluation-harness` by `pip install lm-eval==0.4.2`.
2. Run evaluation with (more documentation at the [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness/tree/big-refactor) repo):
``` sh
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-130m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
python evals/lm_harness_eval.py --model hf --model_args pretrained=EleutherAI/pythia-160m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande --device cuda --batch_size 64
```
To reproduce the results on the `mamba-2.8b-slimpj` model reported in the blogposts:
``` sh
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks boolq,piqa,hellaswag,winogrande,arc_easy,arc_challenge,openbookqa,race,truthfulqa_mc2 --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks mmlu --num_fewshot 5 --device cuda --batch_size 256
```
To run evaluations on Mamba-2 models, simply replace the model names:
``` sh
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/transformerpp-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2attn-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
```
Note that the result of each task might differ from reported values by 0.1-0.3 due to noise in the evaluation process.
## Inference
The script [benchmarks/benchmark_generation_mamba_simple.py](benchmarks/benchmark_generation_mamba_simple.py)
1. autoloads a model from the Hugging Face Hub,
2. generates completions of a user-specified prompt,
3. benchmarks the inference speed of this generation.
Other configurable options include the top-p (nucleus sampling) probability, and the softmax temperature.
### Examples
To test generation latency (e.g. batch size = 1) with different sampling strategies:
``` sh
python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba-2.8b" --prompt "My cat wrote all this CUDA code for a new language model and" --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name "EleutherAI/pythia-2.8b" --prompt "My cat wrote all this CUDA code for a new language model and" --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba-2.8b" --prompt "My cat wrote all this CUDA code for a new language model and" --minp 0.05 --topk 0 --temperature 0.7 --repetition-penalty 1.2
```
To test generation throughput with random prompts (e.g. large batch size):
``` sh
python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba-2.8b" --batch 64
python benchmarks/benchmark_generation_mamba_simple.py --model-name "EleutherAI/pythia-2.8b" --batch 64
```
With Mamba-2, you just need to change the model name:
``` sh
python benchmarks/benchmark_generation_mamba_simple.py --model-name "state-spaces/mamba2-2.7b" --prompt "My cat wrote all this CUDA code for a new language model and" --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
```
## Troubleshooting
### Precision
Our models were trained using PyTorch [AMP](https://pytorch.org/docs/stable/amp.html) for mixed precision. AMP keeps model parameters in float32 and casts to half precision when necessary.
On the other hand, other frameworks like DeepSpeed store parameters in float16 and upcasts when necessary (e.g. for optimizer accumulation).
We've observed that higher precision for the main model parameters may be necessary, because SSMs are sensitive to their recurrent dynamics. If you are experiencing instabilities,
as a first step please try a framework storing parameters in fp32 (such as AMP).
### Initialization
Some parts of the model have initializations inherited from prior work on S4 models.
For [example](https://github.com/state-spaces/mamba/blob/f0affcf69f06d1d06cef018ff640bf080a11c421/mamba_ssm/modules/mamba_simple.py#L102), the $\Delta$ parameter has a targeted range by initializing the bias of its linear projection.
However, some frameworks may have post-initialization hooks (e.g. setting all bias terms in `nn.Linear` modules to zero).
If this is the case, you may have to add custom logic (e.g. this [line](https://github.com/state-spaces/mamba/blob/f0affcf69f06d1d06cef018ff640bf080a11c421/mamba_ssm/modules/mamba_simple.py#L104) turns off re-initializing in our trainer, but would be a no-op in any other framework)
that is specific to the training framework.
## Additional Prerequisites for AMD cards
### Patching ROCm
If you are on ROCm 6.0, run the following steps to avoid errors during compilation. This is not required for ROCm 6.1 onwards.
1. Locate your ROCm installation directory. This is typically found at `/opt/rocm/`, but may vary depending on your installation.
2. Apply the Patch. Run with `sudo` in case you encounter permission issues.
```bash
patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch
```
## Citation
If you use this codebase, or otherwise find our work valuable, please cite Mamba:
```
@article{mamba,
title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
author={Gu, Albert and Dao, Tri},
journal={arXiv preprint arXiv:2312.00752},
year={2023}
}
@inproceedings{mamba2,
title={Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
author={Dao, Tri and Gu, Albert},
booktitle={International Conference on Machine Learning (ICML)},
year={2024}
}
@misc{lahoti2026mamba3improvedsequencemodeling,
title={Mamba-3: Improved Sequence Modeling using State Space Principles},
author={Aakash Lahoti and Kevin Y. Li and Berlin Chen and Caitlin Wang and Aviv Bick and J. Zico Kolter and Tri Dao and Albert Gu},
year={2026},
eprint={2603.15569},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2603.15569},
}
```
================================================
FILE: benchmarks/benchmark_generation_mamba_simple.py
================================================
# Copyright (c) 2023, Tri Dao, Albert Gu.
import argparse
import time
import json
import torch
import torch.nn.functional as F
from einops import rearrange
from transformers import AutoTokenizer, AutoModelForCausalLM
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
parser = argparse.ArgumentParser(description="Generation benchmarking")
parser.add_argument("--model-name", type=str, default="state-spaces/mamba-130m")
parser.add_argument("--prompt", type=str, default=None)
parser.add_argument("--promptlen", type=int, default=100)
parser.add_argument("--genlen", type=int, default=100)
parser.add_argument("--temperature", type=float, default=1.0)
parser.add_argument("--topk", type=int, default=1)
parser.add_argument("--topp", type=float, default=1.0)
parser.add_argument("--minp", type=float, default=0.0)
parser.add_argument("--repetition-penalty", type=float, default=1.0)
parser.add_argument("--batch", type=int, default=1)
args = parser.parse_args()
repeats = 3
device = "cuda"
dtype = torch.float16
print(f"Loading model {args.model_name}")
is_mamba = args.model_name.startswith("state-spaces/mamba") or args.model_name.startswith("state-spaces/transformerpp")
if is_mamba:
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
model = MambaLMHeadModel.from_pretrained(args.model_name, device=device, dtype=dtype)
else:
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
model = AutoModelForCausalLM.from_pretrained(args.model_name, device_map={"": device}, torch_dtype=dtype)
model.eval()
print(f"Number of parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
torch.random.manual_seed(0)
if args.prompt is None:
input_ids = torch.randint(1, 1000, (args.batch, args.promptlen), dtype=torch.long, device="cuda")
attn_mask = torch.ones_like(input_ids, dtype=torch.long, device="cuda")
else:
tokens = tokenizer(args.prompt, return_tensors="pt")
input_ids = tokens.input_ids.to(device=device)
attn_mask = tokens.attention_mask.to(device=device)
max_length = input_ids.shape[1] + args.genlen
if is_mamba:
fn = lambda: model.generate(
input_ids=input_ids,
max_length=max_length,
cg=True,
return_dict_in_generate=True,
output_scores=True,
enable_timing=False,
temperature=args.temperature,
top_k=args.topk,
top_p=args.topp,
min_p=args.minp,
repetition_penalty=args.repetition_penalty,
)
else:
fn = lambda: model.generate(
input_ids=input_ids,
attention_mask=attn_mask,
max_length=max_length,
return_dict_in_generate=True,
pad_token_id=tokenizer.eos_token_id,
do_sample=True,
temperature=args.temperature,
top_k=args.topk,
top_p=args.topp,
repetition_penalty=args.repetition_penalty,
)
out = fn()
if args.prompt is not None:
print(tokenizer.batch_decode(out.sequences.tolist()))
torch.cuda.synchronize()
start = time.time()
for _ in range(repeats):
fn()
torch.cuda.synchronize()
print(f"Prompt length: {len(input_ids[0])}, generation length: {len(out.sequences[0]) - len(input_ids[0])}")
print(f"{args.model_name} prompt processing + decoding time: {(time.time() - start) / repeats * 1000:.0f}ms")
================================================
FILE: csrc/selective_scan/reverse_scan.cuh
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#pragma once
#ifndef USE_ROCM
#include <cub/config.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <cub/block/block_raking_layout.cuh>
// #include <cub/detail/uninitialized_copy.cuh>
#else
#include <hipcub/hipcub.hpp>
namespace cub = hipcub;
#endif
#include "uninitialized_copy.cuh"
/**
* Perform a reverse sequential reduction over \p LENGTH elements of the \p input array. The aggregate is returned.
*/
template <
int LENGTH,
typename T,
typename ReductionOp>
__device__ __forceinline__ T ThreadReverseReduce(const T (&input)[LENGTH], ReductionOp reduction_op) {
static_assert(LENGTH > 0);
T retval = input[LENGTH - 1];
#pragma unroll
for (int i = LENGTH - 2; i >= 0; --i) { retval = reduction_op(retval, input[i]); }
return retval;
}
/**
* Perform a sequential inclusive postfix reverse scan over the statically-sized \p input array, seeded with the specified \p postfix. The aggregate is returned.
*/
template <
int LENGTH,
typename T,
typename ScanOp>
__device__ __forceinline__ T ThreadReverseScanInclusive(
const T (&input)[LENGTH],
T (&output)[LENGTH],
ScanOp scan_op,
const T postfix)
{
T inclusive = postfix;
#pragma unroll
for (int i = LENGTH - 1; i >= 0; --i) {
inclusive = scan_op(inclusive, input[i]);
output[i] = inclusive;
}
return inclusive;
}
/**
* Perform a sequential exclusive postfix reverse scan over the statically-sized \p input array, seeded with the specified \p postfix. The aggregate is returned.
*/
template <
int LENGTH,
typename T,
typename ScanOp>
__device__ __forceinline__ T ThreadReverseScanExclusive(
const T (&input)[LENGTH],
T (&output)[LENGTH],
ScanOp scan_op,
const T postfix)
{
// Careful, output maybe be aliased to input
T exclusive = postfix;
T inclusive;
#pragma unroll
for (int i = LENGTH - 1; i >= 0; --i) {
inclusive = scan_op(exclusive, input[i]);
output[i] = exclusive;
exclusive = inclusive;
}
return inclusive;
}
/**
* \brief WarpReverseScan provides SHFL-based variants of parallel postfix scan of items partitioned across a CUDA thread warp.
*
* LOGICAL_WARP_THREADS must be a power-of-two
*/
template <
typename T, ///< Data type being scanned
int LOGICAL_WARP_THREADS ///< Number of threads per logical warp
>
struct WarpReverseScan {
//---------------------------------------------------------------------
// Constants and type definitions
//---------------------------------------------------------------------
/// Whether the logical warp size and the PTX warp size coincide
// In hipcub, warp_threads is defined as HIPCUB_WARP_THREADS ::rocprim::warp_size()
// While in cub, it's defined as a macro that takes a redundant unused argument.
#ifndef USE_ROCM
#define WARP_THREADS CUB_WARP_THREADS(0)
#else
// ROCm 7.0+: HIPCUB_WARP_THREADS (rocprim::warp_size()) is no longer constexpr.
// We need a compile-time constant for IS_ARCH_WARP below.
// See: https://rocm.docs.amd.com/en/latest/about/release-notes.html
#if defined(__AMDGCN_WAVEFRONT_SIZE)
// Deprecated but still available and constexpr in ROCm 7.x
#define WARP_THREADS __AMDGCN_WAVEFRONT_SIZE
#elif defined(__gfx942__) || defined(__gfx941__) || defined(__gfx940__)
// AMD Instinct MI300 series (CDNA3) - 64-wide wavefronts
#define WARP_THREADS 64
#elif defined(__gfx90a__)
// AMD Instinct MI200 series (CDNA2) - 64-wide wavefronts
#define WARP_THREADS 64
#elif defined(__gfx908__)
// AMD Instinct MI100 (CDNA1) - 64-wide wavefronts
#define WARP_THREADS 64
#elif defined(__gfx906__) || defined(__gfx900__)
// AMD Instinct MI50/MI60 (Vega) - 64-wide wavefronts
#define WARP_THREADS 64
#elif defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__)
// AMD Radeon RX 7000 series (RDNA3) - 32-wide wavefronts
#define WARP_THREADS 32
#elif defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1034__)
// AMD Radeon RX 6000 series (RDNA2) - 32-wide wavefronts
#define WARP_THREADS 32
#elif defined(__gfx1010__) || defined(__gfx1011__) || defined(__gfx1012__)
// AMD Radeon RX 5000 series (RDNA1) - 32-wide wavefronts
#define WARP_THREADS 32
#else
// Unknown architecture - default to 64 (CDNA/GCN)
// This may not be optimal for RDNA GPUs
#pragma message("Warning: Unknown AMD GPU architecture. Defaulting WARP_THREADS to 64. " \
"For RDNA GPUs (gfx10xx/gfx11xx), this should be 32.")
#define WARP_THREADS 64
#endif
#endif
static constexpr bool IS_ARCH_WARP = (LOGICAL_WARP_THREADS == WARP_THREADS);
/// The number of warp scan steps
static constexpr int STEPS = cub::Log2<LOGICAL_WARP_THREADS>::VALUE;
static_assert(LOGICAL_WARP_THREADS == 1 << STEPS);
//---------------------------------------------------------------------
// Thread fields
//---------------------------------------------------------------------
/// Lane index in logical warp
unsigned int lane_id;
/// Logical warp index in 32-thread physical warp
unsigned int warp_id;
/// 32-thread physical warp member mask of logical warp
unsigned int member_mask;
//---------------------------------------------------------------------
// Construction
//---------------------------------------------------------------------
/// Constructor
explicit __device__ __forceinline__
WarpReverseScan()
#ifndef USE_ROCM
: lane_id(threadIdx.x & 0x1f) // CUDA: 32-thread warps, mask = 31
#else
: lane_id(threadIdx.x & (WARP_THREADS - 1)) // ROCm: use actual wavefront size (64 or 32)
#endif
, warp_id(IS_ARCH_WARP ? 0 : (lane_id / LOGICAL_WARP_THREADS))
, member_mask(cub::WarpMask<LOGICAL_WARP_THREADS>(warp_id))
{
if (!IS_ARCH_WARP) {
lane_id = lane_id % LOGICAL_WARP_THREADS;
}
}
/// Broadcast
__device__ __forceinline__ T Broadcast(
T input, ///< [in] The value to broadcast
int src_lane) ///< [in] Which warp lane is to do the broadcasting
{
return cub::ShuffleIndex<LOGICAL_WARP_THREADS>(input, src_lane, member_mask);
}
/// Inclusive scan
template <typename ScanOpT>
__device__ __forceinline__ void InclusiveReverseScan(
T input, ///< [in] Calling thread's input item.
T &inclusive_output, ///< [out] Calling thread's output item. May be aliased with \p input.
ScanOpT scan_op) ///< [in] Binary scan operator
{
inclusive_output = input;
#pragma unroll
for (int STEP = 0; STEP < STEPS; STEP++) {
int offset = 1 << STEP;
T temp = cub::ShuffleDown<LOGICAL_WARP_THREADS>(
inclusive_output, offset, LOGICAL_WARP_THREADS - 1, member_mask
);
// Perform scan op if from a valid peer
inclusive_output = static_cast<int>(lane_id) >= LOGICAL_WARP_THREADS - offset
? inclusive_output : scan_op(temp, inclusive_output);
}
}
/// Exclusive scan
// Get exclusive from inclusive
template <typename ScanOpT>
__device__ __forceinline__ void ExclusiveReverseScan(
T input, ///< [in] Calling thread's input item.
T &exclusive_output, ///< [out] Calling thread's output item. May be aliased with \p input.
ScanOpT scan_op, ///< [in] Binary scan operator
T &warp_aggregate) ///< [out] Warp-wide aggregate reduction of input items.
{
T inclusive_output;
InclusiveReverseScan(input, inclusive_output, scan_op);
warp_aggregate = cub::ShuffleIndex<LOGICAL_WARP_THREADS>(inclusive_output, 0, member_mask);
// initial value unknown
exclusive_output = cub::ShuffleDown<LOGICAL_WARP_THREADS>(
inclusive_output, 1, LOGICAL_WARP_THREADS - 1, member_mask
);
}
/**
* \brief Computes both inclusive and exclusive reverse scans using the specified binary scan functor across the calling warp. Because no initial value is supplied, the \p exclusive_output computed for the last <em>warp-lane</em> is undefined.
*/
template <typename ScanOpT>
__device__ __forceinline__ void ReverseScan(
T input, ///< [in] Calling thread's input item.
T &inclusive_output, ///< [out] Calling thread's inclusive-scan output item.
T &exclusive_output, ///< [out] Calling thread's exclusive-scan output item.
ScanOpT scan_op) ///< [in] Binary scan operator
{
InclusiveReverseScan(input, inclusive_output, scan_op);
// initial value unknown
exclusive_output = cub::ShuffleDown<LOGICAL_WARP_THREADS>(
inclusive_output, 1, LOGICAL_WARP_THREADS - 1, member_mask
);
}
};
/**
* \brief BlockReverseScan provides variants of raking-based parallel postfix scan across a CUDA thread block.
*/
template <
typename T, ///< Data type being scanned
int BLOCK_DIM_X, ///< The thread block length in threads along the X dimension
bool MEMOIZE=false ///< Whether or not to buffer outer raking scan partials to incur fewer shared memory reads at the expense of higher register pressure
>
struct BlockReverseScan {
//---------------------------------------------------------------------
// Types and constants
//---------------------------------------------------------------------
/// Constants
/// The thread block size in threads
static constexpr int BLOCK_THREADS = BLOCK_DIM_X;
/// Layout type for padded thread block raking grid
using BlockRakingLayout = cub::BlockRakingLayout<T, BLOCK_THREADS>;
// The number of reduction elements is not a multiple of the number of raking threads for now
static_assert(BlockRakingLayout::UNGUARDED);
/// Number of raking threads
static constexpr int RAKING_THREADS = BlockRakingLayout::RAKING_THREADS;
/// Number of raking elements per warp synchronous raking thread
static constexpr int SEGMENT_LENGTH = BlockRakingLayout::SEGMENT_LENGTH;
/// Cooperative work can be entirely warp synchronous
static constexpr bool WARP_SYNCHRONOUS = (int(BLOCK_THREADS) == int(RAKING_THREADS));
/// WarpReverseScan utility type
using WarpReverseScan = WarpReverseScan<T, RAKING_THREADS>;
/// Shared memory storage layout type
struct _TempStorage {
typename BlockRakingLayout::TempStorage raking_grid; ///< Padded thread block raking grid
};
/// Alias wrapper allowing storage to be unioned
struct TempStorage : cub::Uninitialized<_TempStorage> {};
//---------------------------------------------------------------------
// Per-thread fields
//---------------------------------------------------------------------
// Thread fields
_TempStorage &temp_storage;
unsigned int linear_tid;
T cached_segment[SEGMENT_LENGTH];
//---------------------------------------------------------------------
// Utility methods
//---------------------------------------------------------------------
/// Performs upsweep raking reduction, returning the aggregate
template <typename ScanOp>
__device__ __forceinline__ T Upsweep(ScanOp scan_op) {
T *smem_raking_ptr = BlockRakingLayout::RakingPtr(temp_storage.raking_grid, linear_tid);
// Read data into registers
#pragma unroll
for (int i = 0; i < SEGMENT_LENGTH; ++i) { cached_segment[i] = smem_raking_ptr[i]; }
T raking_partial = cached_segment[SEGMENT_LENGTH - 1];
#pragma unroll
for (int i = SEGMENT_LENGTH - 2; i >= 0; --i) {
raking_partial = scan_op(raking_partial, cached_segment[i]);
}
return raking_partial;
}
/// Performs exclusive downsweep raking scan
template <typename ScanOp>
__device__ __forceinline__ void ExclusiveDownsweep(
ScanOp scan_op,
T raking_partial)
{
T *smem_raking_ptr = BlockRakingLayout::RakingPtr(temp_storage.raking_grid, linear_tid);
// Read data back into registers
if (!MEMOIZE) {
#pragma unroll
for (int i = 0; i < SEGMENT_LENGTH; ++i) { cached_segment[i] = smem_raking_ptr[i]; }
}
ThreadReverseScanExclusive(cached_segment, cached_segment, scan_op, raking_partial);
// Write data back to smem
#pragma unroll
for (int i = 0; i < SEGMENT_LENGTH; ++i) { smem_raking_ptr[i] = cached_segment[i]; }
}
//---------------------------------------------------------------------
// Constructors
//---------------------------------------------------------------------
/// Constructor
__device__ __forceinline__ BlockReverseScan(
TempStorage &temp_storage)
:
temp_storage(temp_storage.Alias()),
linear_tid(cub::RowMajorTid(BLOCK_DIM_X, 1, 1))
{}
/// Computes an exclusive thread block-wide postfix scan using the specified binary \p scan_op functor. Each thread contributes one input element. the call-back functor \p block_postfix_callback_op is invoked by the first warp in the block, and the value returned by <em>lane</em><sub>0</sub> in that warp is used as the "seed" value that logically postfixes the thread block's scan inputs. Also provides every thread with the block-wide \p block_aggregate of all inputs.
template <
typename ScanOp,
typename BlockPostfixCallbackOp>
__device__ __forceinline__ void ExclusiveReverseScan(
T input, ///< [in] Calling thread's input item
T &exclusive_output, ///< [out] Calling thread's output item (may be aliased to \p input)
ScanOp scan_op, ///< [in] Binary scan operator
BlockPostfixCallbackOp &block_postfix_callback_op) ///< [in-out] <b>[<em>warp</em><sub>0</sub> only]</b> Call-back functor for specifying a thread block-wide postfix to be applied to all inputs.
{
if (WARP_SYNCHRONOUS) {
// Short-circuit directly to warp-synchronous scan
T block_aggregate;
WarpReverseScan warp_scan;
warp_scan.ExclusiveReverseScan(input, exclusive_output, scan_op, block_aggregate);
// Obtain warp-wide postfix in lane0, then broadcast to other lanes
T block_postfix = block_postfix_callback_op(block_aggregate);
block_postfix = warp_scan.Broadcast(block_postfix, 0);
exclusive_output = linear_tid == BLOCK_THREADS - 1 ? block_postfix : scan_op(block_postfix, exclusive_output);
} else {
// Place thread partial into shared memory raking grid
T *placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid);
detail::uninitialized_copy(placement_ptr, input);
__syncthreads();
// Reduce parallelism down to just raking threads
if (linear_tid < RAKING_THREADS) {
WarpReverseScan warp_scan;
// Raking upsweep reduction across shared partials
T upsweep_partial = Upsweep(scan_op);
// Warp-synchronous scan
T exclusive_partial, block_aggregate;
warp_scan.ExclusiveReverseScan(upsweep_partial, exclusive_partial, scan_op, block_aggregate);
// Obtain block-wide postfix in lane0, then broadcast to other lanes
T block_postfix = block_postfix_callback_op(block_aggregate);
block_postfix = warp_scan.Broadcast(block_postfix, 0);
// Update postfix with warpscan exclusive partial
T downsweep_postfix = linear_tid == RAKING_THREADS - 1
? block_postfix : scan_op(block_postfix, exclusive_partial);
// Exclusive raking downsweep scan
ExclusiveDownsweep(scan_op, downsweep_postfix);
}
__syncthreads();
// Grab thread postfix from shared memory
exclusive_output = *placement_ptr;
// // Compute warp scan in each warp.
// // The exclusive output from the last lane in each warp is invalid.
// T inclusive_output;
// WarpReverseScan warp_scan;
// warp_scan.ReverseScan(input, inclusive_output, exclusive_output, scan_op);
// // Compute the warp-wide postfix and block-wide aggregate for each warp. Warp postfix for the last warp is invalid.
// T block_aggregate;
// T warp_postfix = ComputeWarpPostfix(scan_op, inclusive_output, block_aggregate);
// // Apply warp postfix to our lane's partial
// if (warp_id != 0) {
// exclusive_output = scan_op(warp_postfix, exclusive_output);
// if (lane_id == 0) { exclusive_output = warp_postfix; }
// }
// // Use the first warp to determine the thread block postfix, returning the result in lane0
// if (warp_id == 0) {
// T block_postfix = block_postfix_callback_op(block_aggregate);
// if (lane_id == 0) {
// // Share the postfix with all threads
// detail::uninitialized_copy(&temp_storage.block_postfix,
// block_postfix);
// exclusive_output = block_postfix; // The block postfix is the exclusive output for tid0
// }
// }
// __syncthreads();
// // Incorporate thread block postfix into outputs
// T block_postfix = temp_storage.block_postfix;
// if (linear_tid > 0) { exclusive_output = scan_op(block_postfix, exclusive_output); }
}
}
/**
* \brief Computes an inclusive block-wide postfix scan using the specified binary \p scan_op functor. Each thread contributes an array of consecutive input elements. the call-back functor \p block_postfix_callback_op is invoked by the first warp in the block, and the value returned by <em>lane</em><sub>0</sub> in that warp is used as the "seed" value that logically postfixes the thread block's scan inputs. Also provides every thread with the block-wide \p block_aggregate of all inputs.
*/
template <
int ITEMS_PER_THREAD,
typename ScanOp,
typename BlockPostfixCallbackOp>
__device__ __forceinline__ void InclusiveReverseScan(
T (&input)[ITEMS_PER_THREAD], ///< [in] Calling thread's input items
T (&output)[ITEMS_PER_THREAD], ///< [out] Calling thread's output items (may be aliased to \p input)
ScanOp scan_op, ///< [in] Binary scan functor
BlockPostfixCallbackOp &block_postfix_callback_op) ///< [in-out] <b>[<em>warp</em><sub>0</sub> only]</b> Call-back functor for specifying a block-wide postfix to be applied to the logical input sequence.
{
// Reduce consecutive thread items in registers
T thread_postfix = ThreadReverseReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveReverseScan(thread_postfix, thread_postfix, scan_op, block_postfix_callback_op);
// Inclusive scan in registers with postfix as seed
ThreadReverseScanInclusive(input, output, scan_op, thread_postfix);
}
};
================================================
FILE: csrc/selective_scan/selective_scan.cpp
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/python.h>
#include <vector>
#include "selective_scan.h"
#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
#define DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...) \
if (ITYPE == at::ScalarType::Half) { \
using input_t = at::Half; \
__VA_ARGS__(); \
} else if (ITYPE == at::ScalarType::BFloat16) { \
using input_t = at::BFloat16; \
__VA_ARGS__(); \
} else if (ITYPE == at::ScalarType::Float) { \
using input_t = float; \
__VA_ARGS__(); \
} else { \
AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
}
#define DISPATCH_WTYPE_FLOAT_AND_HALF_AND_BF16(WTYPE, NAME, ...) \
if (WTYPE == at::ScalarType::Half) { \
using weight_t = at::Half; \
__VA_ARGS__(); \
} else if (WTYPE == at::ScalarType::BFloat16) { \
using weight_t = at::BFloat16; \
__VA_ARGS__(); \
} else if (WTYPE == at::ScalarType::Float) { \
using weight_t = float; \
__VA_ARGS__(); \
} else { \
AT_ERROR(#NAME, " not implemented for weight type '", toString(WTYPE), "'"); \
}
#define DISPATCH_WTYPE_FLOAT_AND_COMPLEX(WTYPE, NAME, ...) \
if (WTYPE == at::ScalarType::Float) { \
using weight_t = float; \
__VA_ARGS__(); \
} else if (WTYPE == at::ScalarType::ComplexFloat) { \
using weight_t = c10::complex<float>; \
__VA_ARGS__(); \
} else { \
AT_ERROR(#NAME, " not implemented for weight type '", toString(WTYPE), "'"); \
}
template<typename input_t, typename weight_t>
void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream);
template <typename input_t, typename weight_t>
void selective_scan_bwd_cuda(SSMParamsBwd ¶ms, cudaStream_t stream);
void set_ssm_params_fwd(SSMParamsBase ¶ms,
// sizes
const size_t batch,
const size_t dim,
const size_t seqlen,
const size_t dstate,
const size_t n_groups,
const size_t n_chunks,
const bool is_variable_B,
const bool is_variable_C,
// device pointers
const at::Tensor u,
const at::Tensor delta,
const at::Tensor A,
const at::Tensor B,
const at::Tensor C,
const at::Tensor out,
const at::Tensor z,
const at::Tensor out_z,
void* D_ptr,
void* delta_bias_ptr,
void* x_ptr,
bool has_z,
bool delta_softplus) {
// Reset the parameters
memset(¶ms, 0, sizeof(params));
params.batch = batch;
params.dim = dim;
params.seqlen = seqlen;
params.dstate = dstate;
params.n_groups = n_groups;
params.n_chunks = n_chunks;
params.dim_ngroups_ratio = dim / n_groups;
params.delta_softplus = delta_softplus;
params.is_variable_B = is_variable_B;
params.is_variable_C = is_variable_C;
// Set the pointers and strides.
params.u_ptr = u.data_ptr();
params.delta_ptr = delta.data_ptr();
params.A_ptr = A.data_ptr();
params.B_ptr = B.data_ptr();
params.C_ptr = C.data_ptr();
params.D_ptr = D_ptr;
params.delta_bias_ptr = delta_bias_ptr;
params.out_ptr = out.data_ptr();
params.x_ptr = x_ptr;
params.z_ptr = has_z ? z.data_ptr() : nullptr;
params.out_z_ptr = has_z ? out_z.data_ptr() : nullptr;
// All stride are in elements, not bytes.
params.A_d_stride = A.stride(0);
params.A_dstate_stride = A.stride(1);
if (!is_variable_B) {
params.B_d_stride = B.stride(0);
} else {
params.B_batch_stride = B.stride(0);
params.B_group_stride = B.stride(1);
}
params.B_dstate_stride = !is_variable_B ? B.stride(1) : B.stride(2);
if (!is_variable_C) {
params.C_d_stride = C.stride(0);
} else {
params.C_batch_stride = C.stride(0);
params.C_group_stride = C.stride(1);
}
params.C_dstate_stride = !is_variable_C ? C.stride(1) : C.stride(2);
params.u_batch_stride = u.stride(0);
params.u_d_stride = u.stride(1);
params.delta_batch_stride = delta.stride(0);
params.delta_d_stride = delta.stride(1);
if (has_z) {
params.z_batch_stride = z.stride(0);
params.z_d_stride = z.stride(1);
params.out_z_batch_stride = out_z.stride(0);
params.out_z_d_stride = out_z.stride(1);
}
params.out_batch_stride = out.stride(0);
params.out_d_stride = out.stride(1);
}
void set_ssm_params_bwd(SSMParamsBwd ¶ms,
// sizes
const size_t batch,
const size_t dim,
const size_t seqlen,
const size_t dstate,
const size_t n_groups,
const size_t n_chunks,
const bool is_variable_B,
const bool is_variable_C,
// device pointers
const at::Tensor u,
const at::Tensor delta,
const at::Tensor A,
const at::Tensor B,
const at::Tensor C,
const at::Tensor z,
const at::Tensor out,
const at::Tensor out_z,
void* D_ptr,
void* delta_bias_ptr,
void* x_ptr,
const at::Tensor dout,
const at::Tensor du,
const at::Tensor ddelta,
const at::Tensor dA,
const at::Tensor dB,
const at::Tensor dC,
const at::Tensor dz,
void* dD_ptr,
void* ddelta_bias_ptr,
bool has_z,
bool delta_softplus,
bool recompute_out_z) {
// Pass in "dout" instead of "out", we're not gonna use "out" unless we have z
set_ssm_params_fwd(params, batch, dim, seqlen, dstate, n_groups, n_chunks, is_variable_B, is_variable_C,
u, delta, A, B, C, has_z ? out : dout,
has_z ? z : dout,
// If not recompute_out_z, pass dout instead of out_z.
// This won't be used by the bwd kernel
recompute_out_z ? out_z : dout,
D_ptr, delta_bias_ptr, x_ptr, has_z, delta_softplus);
if (!recompute_out_z) { params.out_z_ptr = nullptr; }
// Set the pointers and strides.
params.dout_ptr = dout.data_ptr();
params.du_ptr = du.data_ptr();
params.dA_ptr = dA.data_ptr();
params.dB_ptr = dB.data_ptr();
params.dC_ptr = dC.data_ptr();
params.dD_ptr = dD_ptr;
params.ddelta_ptr = ddelta.data_ptr();
params.ddelta_bias_ptr = ddelta_bias_ptr;
params.dz_ptr = has_z ? dz.data_ptr() : nullptr;
// All stride are in elements, not bytes.
params.dout_batch_stride = dout.stride(0);
params.dout_d_stride = dout.stride(1);
params.dA_d_stride = dA.stride(0);
params.dA_dstate_stride = dA.stride(1);
if (!is_variable_B) {
params.dB_d_stride = dB.stride(0);
} else {
params.dB_batch_stride = dB.stride(0);
params.dB_group_stride = dB.stride(1);
}
params.dB_dstate_stride = !is_variable_B ? dB.stride(1) : dB.stride(2);
if (!is_variable_C) {
params.dC_d_stride = dC.stride(0);
} else {
params.dC_batch_stride = dC.stride(0);
params.dC_group_stride = dC.stride(1);
}
params.dC_dstate_stride = !is_variable_C ? dC.stride(1) : dC.stride(2);
params.du_batch_stride = du.stride(0);
params.du_d_stride = du.stride(1);
params.ddelta_batch_stride = ddelta.stride(0);
params.ddelta_d_stride = ddelta.stride(1);
if (has_z) {
params.dz_batch_stride = dz.stride(0);
params.dz_d_stride = dz.stride(1);
}
}
std::vector<at::Tensor>
selective_scan_fwd(const at::Tensor &u, const at::Tensor &delta,
const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
const c10::optional<at::Tensor> &D_,
const c10::optional<at::Tensor> &z_,
const c10::optional<at::Tensor> &delta_bias_,
bool delta_softplus) {
auto input_type = u.scalar_type();
auto weight_type = A.scalar_type();
TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
TORCH_CHECK(weight_type == at::ScalarType::Float || weight_type == at::ScalarType::ComplexFloat);
const bool is_variable_B = B.dim() >= 3;
const bool is_variable_C = C.dim() >= 3;
const bool is_complex = weight_type == at::ScalarType::ComplexFloat;
TORCH_CHECK(delta.scalar_type() == input_type);
TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
TORCH_CHECK(u.is_cuda());
TORCH_CHECK(delta.is_cuda());
TORCH_CHECK(A.is_cuda());
TORCH_CHECK(B.is_cuda());
TORCH_CHECK(C.is_cuda());
TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
const auto sizes = u.sizes();
const int batch_size = sizes[0];
const int dim = sizes[1];
const int seqlen = sizes[2];
const int dstate = A.size(1);
const int n_groups = is_variable_B ? B.size(1) : 1;
TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
CHECK_SHAPE(u, batch_size, dim, seqlen);
CHECK_SHAPE(delta, batch_size, dim, seqlen);
CHECK_SHAPE(A, dim, dstate);
if (!is_variable_B) {
CHECK_SHAPE(B, dim, dstate);
} else {
CHECK_SHAPE(B, batch_size, n_groups, dstate, !is_complex ? seqlen : seqlen * 2);
TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
}
if (!is_variable_C) {
CHECK_SHAPE(C, dim, dstate);
} else {
CHECK_SHAPE(C, batch_size, n_groups, dstate, !is_complex ? seqlen: seqlen * 2);
TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
}
if (D_.has_value()) {
auto D = D_.value();
TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
TORCH_CHECK(D.is_cuda());
TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
CHECK_SHAPE(D, dim);
}
if (delta_bias_.has_value()) {
auto delta_bias = delta_bias_.value();
TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
TORCH_CHECK(delta_bias.is_cuda());
TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
CHECK_SHAPE(delta_bias, dim);
}
at::Tensor z, out_z;
const bool has_z = z_.has_value();
if (has_z) {
z = z_.value();
TORCH_CHECK(z.scalar_type() == input_type);
TORCH_CHECK(z.is_cuda());
TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
CHECK_SHAPE(z, batch_size, dim, seqlen);
out_z = torch::empty_like(z);
}
const int n_chunks = (seqlen + 2048 - 1) / 2048;
// const int n_chunks = (seqlen + 1024 - 1) / 1024;
// at::Tensor out = torch::empty_like(u);
// Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout
at::Tensor out = torch::empty_like(delta);
at::Tensor x;
x = torch::empty({batch_size, dim, n_chunks, dstate * 2}, u.options().dtype(weight_type));
SSMParamsBase params;
set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, n_chunks, is_variable_B, is_variable_C,
u, delta, A, B, C, out, z, out_z,
D_.has_value() ? D_.value().data_ptr() : nullptr,
delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
x.data_ptr(),
has_z,
delta_softplus);
// Otherwise the kernel will be launched from cuda:0 device
// Cast to char to avoid compiler warning about narrowing
at::cuda::CUDAGuard device_guard{u.device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_fwd", [&] {
DISPATCH_WTYPE_FLOAT_AND_COMPLEX(A.scalar_type(), "selective_scan_fwd", [&] {
selective_scan_fwd_cuda<input_t, weight_t>(params, stream);
});
});
std::vector<at::Tensor> result = {out, x};
if (has_z) { result.push_back(out_z); }
return result;
}
std::vector<at::Tensor>
selective_scan_bwd(const at::Tensor &u, const at::Tensor &delta,
const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
const c10::optional<at::Tensor> &D_,
const c10::optional<at::Tensor> &z_,
const c10::optional<at::Tensor> &delta_bias_,
const at::Tensor &dout,
const c10::optional<at::Tensor> &x_,
const c10::optional<at::Tensor> &out_,
c10::optional<at::Tensor> &dz_,
bool delta_softplus,
bool recompute_out_z) {
auto input_type = u.scalar_type();
auto weight_type = A.scalar_type();
TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
TORCH_CHECK(weight_type == at::ScalarType::Float || weight_type == at::ScalarType::ComplexFloat);
const bool is_variable_B = B.dim() >= 3;
const bool is_variable_C = C.dim() >= 3;
const bool is_complex = weight_type == at::ScalarType::ComplexFloat;
TORCH_CHECK(delta.scalar_type() == input_type);
TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
TORCH_CHECK(dout.scalar_type() == input_type);
TORCH_CHECK(u.is_cuda());
TORCH_CHECK(delta.is_cuda());
TORCH_CHECK(A.is_cuda());
TORCH_CHECK(B.is_cuda());
TORCH_CHECK(C.is_cuda());
TORCH_CHECK(dout.is_cuda());
TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
TORCH_CHECK(dout.stride(-1) == 1 || dout.size(-1) == 1);
const auto sizes = u.sizes();
const int batch_size = sizes[0];
const int dim = sizes[1];
const int seqlen = sizes[2];
const int dstate = A.size(1);
const int n_groups = is_variable_B ? B.size(1) : 1;
TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
CHECK_SHAPE(u, batch_size, dim, seqlen);
CHECK_SHAPE(delta, batch_size, dim, seqlen);
CHECK_SHAPE(A, dim, dstate);
if (!is_variable_B) {
CHECK_SHAPE(B, dim, dstate);
} else {
CHECK_SHAPE(B, batch_size, n_groups, dstate, !is_complex ? seqlen : seqlen * 2);
TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
}
if (!is_variable_C) {
CHECK_SHAPE(C, dim, dstate);
} else {
CHECK_SHAPE(C, batch_size, n_groups, dstate, !is_complex ? seqlen: seqlen * 2);
TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
}
CHECK_SHAPE(dout, batch_size, dim, seqlen);
if (D_.has_value()) {
auto D = D_.value();
TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
TORCH_CHECK(D.is_cuda());
TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
CHECK_SHAPE(D, dim);
}
if (delta_bias_.has_value()) {
auto delta_bias = delta_bias_.value();
TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
TORCH_CHECK(delta_bias.is_cuda());
TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
CHECK_SHAPE(delta_bias, dim);
}
at::Tensor z, out, dz, out_z;
const bool has_z = z_.has_value();
if (has_z) {
z = z_.value();
TORCH_CHECK(z.scalar_type() == input_type);
TORCH_CHECK(z.is_cuda());
TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
CHECK_SHAPE(z, batch_size, dim, seqlen);
TORCH_CHECK(out_.has_value());
out = out_.value();
TORCH_CHECK(out.scalar_type() == input_type);
TORCH_CHECK(out.is_cuda());
TORCH_CHECK(out.stride(-1) == 1 || out.size(-1) == 1);
CHECK_SHAPE(out, batch_size, dim, seqlen);
if (dz_.has_value()) {
dz = dz_.value();
TORCH_CHECK(dz.scalar_type() == input_type);
TORCH_CHECK(dz.is_cuda());
TORCH_CHECK(dz.stride(-1) == 1 || dz.size(-1) == 1);
CHECK_SHAPE(dz, batch_size, dim, seqlen);
} else {
dz = torch::empty_like(z);
}
if (recompute_out_z) {
out_z = torch::empty_like(out);
}
}
const int n_chunks = (seqlen + 2048 - 1) / 2048;
// const int n_chunks = (seqlen + 1024 - 1) / 1024;
if (n_chunks > 1) { TORCH_CHECK(x_.has_value()); }
if (x_.has_value()) {
auto x = x_.value();
TORCH_CHECK(x.scalar_type() == weight_type);
TORCH_CHECK(x.is_cuda());
TORCH_CHECK(x.is_contiguous());
CHECK_SHAPE(x, batch_size, dim, n_chunks, 2 * dstate);
}
at::Tensor du = torch::empty_like(u);
at::Tensor ddelta = torch::empty_like(delta);
at::Tensor dA = torch::zeros_like(A);
at::Tensor dB = !is_variable_B ? torch::zeros_like(B) : torch::zeros_like(B, B.options().dtype(torch::kFloat32));
at::Tensor dC = !is_variable_C ? torch::zeros_like(C) : torch::zeros_like(C, C.options().dtype(torch::kFloat32));
at::Tensor dD;
if (D_.has_value()) { dD = torch::zeros_like(D_.value()); }
at::Tensor ddelta_bias;
if (delta_bias_.has_value()) { ddelta_bias = torch::zeros_like(delta_bias_.value()); }
SSMParamsBwd params;
set_ssm_params_bwd(params, batch_size, dim, seqlen, dstate, n_groups, n_chunks, is_variable_B, is_variable_C,
u, delta, A, B, C, z, out, out_z,
D_.has_value() ? D_.value().data_ptr() : nullptr,
delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
x_.has_value() ? x_.value().data_ptr() : nullptr,
dout, du, ddelta, dA, dB, dC, dz,
D_.has_value() ? dD.data_ptr() : nullptr,
delta_bias_.has_value() ? ddelta_bias.data_ptr() : nullptr,
has_z, delta_softplus, recompute_out_z);
// Otherwise the kernel will be launched from cuda:0 device
// Cast to char to avoid compiler warning about narrowing
at::cuda::CUDAGuard device_guard{u.device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_bwd", [&] {
DISPATCH_WTYPE_FLOAT_AND_COMPLEX(A.scalar_type(), "selective_scan_bwd", [&] {
selective_scan_bwd_cuda<input_t, weight_t>(params, stream);
});
});
std::vector<at::Tensor> result = {du, ddelta, dA, dB.to(B.dtype()), dC.to(C.dtype()), dD, ddelta_bias};
if (has_z) { result.push_back(dz); }
if (recompute_out_z) { result.push_back(out_z); }
return result;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fwd", &selective_scan_fwd, "Selective scan forward");
m.def("bwd", &selective_scan_bwd, "Selective scan backward");
}
================================================
FILE: csrc/selective_scan/selective_scan.h
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#pragma once
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SSMScanParamsBase {
using index_t = uint32_t;
int batch, seqlen, n_chunks;
index_t a_batch_stride;
index_t b_batch_stride;
index_t out_batch_stride;
// Common data pointers.
void *__restrict__ a_ptr;
void *__restrict__ b_ptr;
void *__restrict__ out_ptr;
void *__restrict__ x_ptr;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SSMParamsBase {
using index_t = uint32_t;
int batch, dim, seqlen, dstate, n_groups, n_chunks;
int dim_ngroups_ratio;
bool is_variable_B;
bool is_variable_C;
bool delta_softplus;
index_t A_d_stride;
index_t A_dstate_stride;
index_t B_batch_stride;
index_t B_d_stride;
index_t B_dstate_stride;
index_t B_group_stride;
index_t C_batch_stride;
index_t C_d_stride;
index_t C_dstate_stride;
index_t C_group_stride;
index_t u_batch_stride;
index_t u_d_stride;
index_t delta_batch_stride;
index_t delta_d_stride;
index_t z_batch_stride;
index_t z_d_stride;
index_t out_batch_stride;
index_t out_d_stride;
index_t out_z_batch_stride;
index_t out_z_d_stride;
// Common data pointers.
void *__restrict__ A_ptr;
void *__restrict__ B_ptr;
void *__restrict__ C_ptr;
void *__restrict__ D_ptr;
void *__restrict__ u_ptr;
void *__restrict__ delta_ptr;
void *__restrict__ delta_bias_ptr;
void *__restrict__ out_ptr;
void *__restrict__ x_ptr;
void *__restrict__ z_ptr;
void *__restrict__ out_z_ptr;
};
struct SSMParamsBwd: public SSMParamsBase {
index_t dout_batch_stride;
index_t dout_d_stride;
index_t dA_d_stride;
index_t dA_dstate_stride;
index_t dB_batch_stride;
index_t dB_group_stride;
index_t dB_d_stride;
index_t dB_dstate_stride;
index_t dC_batch_stride;
index_t dC_group_stride;
index_t dC_d_stride;
index_t dC_dstate_stride;
index_t du_batch_stride;
index_t du_d_stride;
index_t dz_batch_stride;
index_t dz_d_stride;
index_t ddelta_batch_stride;
index_t ddelta_d_stride;
// Common data pointers.
void *__restrict__ dout_ptr;
void *__restrict__ dA_ptr;
void *__restrict__ dB_ptr;
void *__restrict__ dC_ptr;
void *__restrict__ dD_ptr;
void *__restrict__ du_ptr;
void *__restrict__ dz_ptr;
void *__restrict__ ddelta_ptr;
void *__restrict__ ddelta_bias_ptr;
};
================================================
FILE: csrc/selective_scan/selective_scan_bwd_bf16_complex.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_bwd_kernel.cuh"
template void selective_scan_bwd_cuda<at::BFloat16, complex_t>(SSMParamsBwd ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_bwd_bf16_real.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_bwd_kernel.cuh"
template void selective_scan_bwd_cuda<at::BFloat16, float>(SSMParamsBwd ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_bwd_fp16_complex.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_bwd_kernel.cuh"
template void selective_scan_bwd_cuda<at::Half, complex_t>(SSMParamsBwd ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_bwd_fp16_real.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_bwd_kernel.cuh"
template void selective_scan_bwd_cuda<at::Half, float>(SSMParamsBwd ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_bwd_fp32_complex.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_bwd_kernel.cuh"
template void selective_scan_bwd_cuda<float, complex_t>(SSMParamsBwd ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_bwd_fp32_real.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_bwd_kernel.cuh"
template void selective_scan_bwd_cuda<float, float>(SSMParamsBwd ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_bwd_kernel.cuh
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#pragma once
#include <c10/util/BFloat16.h>
#include <c10/util/Half.h>
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
#include <ATen/cuda/Atomic.cuh> // For atomicAdd on complex
#ifndef USE_ROCM
#include <cub/block/block_load.cuh>
#include <cub/block/block_store.cuh>
#include <cub/block/block_scan.cuh>
#include <cub/block/block_reduce.cuh>
#else
#include <hipcub/hipcub.hpp>
namespace cub = hipcub;
#endif
#include "selective_scan.h"
#include "selective_scan_common.h"
#include "reverse_scan.cuh"
#include "static_switch.h"
template<typename scalar_t> __device__ __forceinline__ scalar_t conj(scalar_t x);
template<> __device__ __forceinline__ float conj<float>(float x) { return x; }
template<> __device__ __forceinline__ complex_t conj<complex_t>(complex_t x) { return std::conj(x); }
template<int kNThreads_, int kNItems_, bool kIsEvenLen_, bool kIsVariableB_, bool kIsVariableC_,
bool kDeltaSoftplus_, bool kHasZ_, typename input_t_, typename weight_t_>
struct Selective_Scan_bwd_kernel_traits {
static_assert(kNItems_ % 4 == 0);
using input_t = input_t_;
using weight_t = weight_t_;
static constexpr int kNThreads = kNThreads_;
static constexpr int kNItems = kNItems_;
static constexpr int kNBytes = sizeof(input_t);
static_assert(kNBytes == 2 || kNBytes == 4);
static constexpr int kNElts = kNBytes == 4 ? 4 : constexpr_min(8, kNItems);
static_assert(kNItems % kNElts == 0);
static constexpr int kNLoads = kNItems / kNElts;
static constexpr bool kIsComplex = std::is_same_v<weight_t, complex_t>;
static constexpr bool kIsEvenLen = kIsEvenLen_;
static constexpr bool kIsVariableB = kIsVariableB_;
static constexpr bool kIsVariableC = kIsVariableC_;
static constexpr bool kDeltaSoftplus = kDeltaSoftplus_;
static constexpr bool kHasZ = kHasZ_;
// Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads with float improves occupancy.
// For complex this would lead to massive register spilling, so we keep it at 2.
static constexpr int kMinBlocks = kNThreads == 128 && !kIsComplex ? 3 : 2;
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
using scan_t = std::conditional_t<!kIsComplex, float2, float4>;
using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, !kIsComplex ? kNItems : kNItems * 2, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, !kIsComplex ? kNLoads : kNLoads * 2, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads, cub::BLOCK_STORE_WARP_TRANSPOSE>;
// using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
// using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
using BlockReverseScanT = BlockReverseScan<scan_t, kNThreads>;
using BlockReduceT = cub::BlockReduce<scan_t, kNThreads>;
using BlockReduceFloatT = cub::BlockReduce<float, kNThreads>;
using BlockReduceComplexT = cub::BlockReduce<complex_t, kNThreads>;
using BlockExchangeT = cub::BlockExchange<float, kNThreads, !kIsComplex ? kNItems : kNItems * 2>;
static constexpr int kSmemIOSize = custom_max({sizeof(typename BlockLoadT::TempStorage),
sizeof(typename BlockLoadVecT::TempStorage),
(int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightT::TempStorage),
(int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightVecT::TempStorage),
sizeof(typename BlockStoreT::TempStorage),
sizeof(typename BlockStoreVecT::TempStorage)});
static constexpr int kSmemExchangeSize = (int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockExchangeT::TempStorage);
static constexpr int kSmemReduceSize = sizeof(typename BlockReduceT::TempStorage);
static constexpr int kSmemSize = kSmemIOSize + kSmemExchangeSize + kSmemReduceSize + sizeof(typename BlockScanT::TempStorage) + sizeof(typename BlockReverseScanT::TempStorage);
};
template<typename Ktraits>
__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
void selective_scan_bwd_kernel(SSMParamsBwd params) {
constexpr bool kIsComplex = Ktraits::kIsComplex;
constexpr bool kIsVariableB = Ktraits::kIsVariableB;
constexpr bool kIsVariableC = Ktraits::kIsVariableC;
constexpr bool kDeltaSoftplus = Ktraits::kDeltaSoftplus;
constexpr bool kHasZ = Ktraits::kHasZ;
constexpr int kNThreads = Ktraits::kNThreads;
constexpr int kNItems = Ktraits::kNItems;
using input_t = typename Ktraits::input_t;
using weight_t = typename Ktraits::weight_t;
using scan_t = typename Ktraits::scan_t;
// Shared memory.
extern __shared__ char smem_[];
// cast to lvalue reference of expected type
// char *smem_loadstorescan = smem_ + 2 * MAX_DSTATE * sizeof(weight_t);
// auto& smem_load = reinterpret_cast<typename BlockLoadT::TempStorage&>(smem_ + 2 * MAX_DSTATE * sizeof(weight_t));
// auto& smem_load = reinterpret_cast<typename BlockLoadT::TempStorage&>(smem_loadstorescan);
auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
auto& smem_exchange = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
auto& smem_exchange1 = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize + sizeof(typename Ktraits::BlockExchangeT::TempStorage));
auto& smem_reduce = *reinterpret_cast<typename Ktraits::BlockReduceT::TempStorage*>(reinterpret_cast<char *>(&smem_exchange) + Ktraits::kSmemExchangeSize);
auto& smem_reduce_float = *reinterpret_cast<typename Ktraits::BlockReduceFloatT::TempStorage*>(&smem_reduce);
auto& smem_reduce_complex = *reinterpret_cast<typename Ktraits::BlockReduceComplexT::TempStorage*>(&smem_reduce);
auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(reinterpret_cast<char *>(&smem_reduce) + Ktraits::kSmemReduceSize);
auto& smem_reverse_scan = *reinterpret_cast<typename Ktraits::BlockReverseScanT::TempStorage*>(reinterpret_cast<char *>(&smem_scan) + sizeof(typename Ktraits::BlockScanT::TempStorage));
weight_t *smem_delta_a = reinterpret_cast<weight_t *>(smem_ + Ktraits::kSmemSize);
scan_t *smem_running_postfix = reinterpret_cast<scan_t *>(smem_delta_a + 2 * MAX_DSTATE + kNThreads);
weight_t *smem_da = reinterpret_cast<weight_t *>(smem_running_postfix + MAX_DSTATE);
weight_t *smem_dbc = reinterpret_cast<weight_t *>(smem_da + MAX_DSTATE);
const int batch_id = blockIdx.x;
const int dim_id = blockIdx.y;
const int group_id = dim_id / (params.dim_ngroups_ratio);
input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+ dim_id * params.u_d_stride;
input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+ dim_id * params.delta_d_stride;
input_t *dout = reinterpret_cast<input_t *>(params.dout_ptr) + batch_id * params.dout_batch_stride
+ dim_id * params.dout_d_stride;
weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id * params.A_d_stride;
weight_t *B = reinterpret_cast<weight_t *>(params.B_ptr) + dim_id * params.B_d_stride;
input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
weight_t *C = reinterpret_cast<weight_t *>(params.C_ptr) + dim_id * params.C_d_stride;
input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
weight_t *dA = reinterpret_cast<weight_t *>(params.dA_ptr) + dim_id * params.dA_d_stride;
weight_t *dB = reinterpret_cast<weight_t *>(params.dB_ptr)
+ (!kIsVariableB ? dim_id * params.dB_d_stride : batch_id * (!kIsComplex ? params.dB_batch_stride : params.dB_batch_stride / 2) + group_id * params.dB_group_stride);
weight_t *dC = reinterpret_cast<weight_t *>(params.dC_ptr)
+ (!kIsVariableC ? dim_id * params.dC_d_stride : batch_id * (!kIsComplex ? params.dC_batch_stride : params.dC_batch_stride / 2) + group_id * params.dC_group_stride);
float *dD = params.dD_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.dD_ptr) + dim_id;
float D_val = params.D_ptr == nullptr ? 0 : reinterpret_cast<float *>(params.D_ptr)[dim_id];
float *ddelta_bias = params.ddelta_bias_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.ddelta_bias_ptr) + dim_id;
float delta_bias = params.delta_bias_ptr == nullptr ? 0 : reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id];
scan_t *x = params.x_ptr == nullptr
? nullptr
: reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id) * (params.n_chunks) * params.dstate;
float dD_val = 0;
float ddelta_bias_val = 0;
constexpr int kChunkSize = kNThreads * kNItems;
u += (params.n_chunks - 1) * kChunkSize;
delta += (params.n_chunks - 1) * kChunkSize;
dout += (params.n_chunks - 1) * kChunkSize;
Bvar += (params.n_chunks - 1) * kChunkSize * (!kIsComplex ? 1 : 2);
Cvar += (params.n_chunks - 1) * kChunkSize * (!kIsComplex ? 1 : 2);
for (int chunk = params.n_chunks - 1; chunk >= 0; --chunk) {
input_t u_vals[kNItems];
input_t delta_vals_load[kNItems];
input_t dout_vals_load[kNItems];
__syncthreads();
load_input<Ktraits>(u, u_vals, smem_load, params.seqlen - chunk * kChunkSize);
u -= kChunkSize;
__syncthreads();
load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
// Will reload delta at the same location if kDeltaSoftplus
if constexpr (!kDeltaSoftplus) { delta -= kChunkSize; }
__syncthreads();
load_input<Ktraits>(dout, dout_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
dout -= kChunkSize;
float dout_vals[kNItems], delta_vals[kNItems];
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
dout_vals[i] = float(dout_vals_load[i]);
delta_vals[i] = float(delta_vals_load[i]) + delta_bias;
if constexpr (kDeltaSoftplus) {
delta_vals[i] = delta_vals[i] <= 20.f ? log1pf(expf(delta_vals[i])) : delta_vals[i];
}
}
if constexpr (kHasZ) {
input_t *z = reinterpret_cast<input_t *>(params.z_ptr) + batch_id * params.z_batch_stride
+ dim_id * params.z_d_stride + chunk * kChunkSize;
input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + batch_id * params.out_batch_stride
+ dim_id * params.out_d_stride + chunk * kChunkSize;
input_t *dz = reinterpret_cast<input_t *>(params.dz_ptr) + batch_id * params.dz_batch_stride
+ dim_id * params.dz_d_stride + chunk * kChunkSize;
input_t z_vals[kNItems], out_vals[kNItems];
__syncthreads();
load_input<Ktraits>(z, z_vals, smem_load, params.seqlen - chunk * kChunkSize);
__syncthreads();
load_input<Ktraits>(out, out_vals, smem_load, params.seqlen - chunk * kChunkSize);
float dz_vals[kNItems], z_silu_vals[kNItems];
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
float z_val = z_vals[i];
float z_sigmoid_val = 1.0f / (1.0f + expf(-z_val));
z_silu_vals[i] = z_val * z_sigmoid_val;
dz_vals[i] = dout_vals[i] * float(out_vals[i]) * z_sigmoid_val
* (1.0f + z_val * (1.0f - z_sigmoid_val));
dout_vals[i] *= z_silu_vals[i];
}
__syncthreads();
store_output<Ktraits>(dz, dz_vals, smem_store, params.seqlen - chunk * kChunkSize);
if (params.out_z_ptr != nullptr) { // Recompute and store out_z
float out_z_vals[kNItems];
#pragma unroll
for (int i = 0; i < kNItems; ++i) { out_z_vals[i] = float(out_vals[i]) * z_silu_vals[i]; }
// if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0) {
// printf("out_val=%f, z_silu_val = %f, out_z_val = %f\n", float(out_vals[0]), z_silu_vals[0], out_z_vals[0]);
// }
input_t *out_z = reinterpret_cast<input_t *>(params.out_z_ptr) + batch_id * params.out_z_batch_stride
+ dim_id * params.out_z_d_stride + chunk * kChunkSize;
__syncthreads();
store_output<Ktraits>(out_z, out_z_vals, smem_store, params.seqlen - chunk * kChunkSize);
}
}
float du_vals[kNItems];
#pragma unroll
for (int i = 0; i < kNItems; ++i) { du_vals[i] = D_val * dout_vals[i]; }
#pragma unroll
for (int i = 0; i < kNItems; ++i) { dD_val += dout_vals[i] * float(u_vals[i]); }
float ddelta_vals[kNItems] = {0};
__syncthreads();
for (int state_idx = 0; state_idx < params.dstate; ++state_idx) {
const weight_t A_val = A[state_idx * params.A_dstate_stride];
// Multiply the real part of A with LOG2E so we can use exp2f instead of expf.
weight_t A_scaled;
constexpr float kLog2e = M_LOG2E;
if constexpr (!kIsComplex) {
A_scaled = A_val * kLog2e;
} else {
A_scaled = complex_t(A_val.real_ * kLog2e, A_val.imag_);
}
weight_t B_val, C_val;
weight_t B_vals[kNItems], C_vals[kNItems];
if constexpr (!kIsVariableB) {
B_val = B[state_idx * params.B_dstate_stride];
} else {
load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
smem_load_weight, (params.seqlen - chunk * kChunkSize) * (!kIsComplex ? 1 : 2));
}
if constexpr (!kIsVariableC) {
C_val = C[state_idx * params.C_dstate_stride];
} else {
auto &smem_load_weight_C = !kIsVariableB ? smem_load_weight : smem_load_weight1;
load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
smem_load_weight_C, (params.seqlen - chunk * kChunkSize) * (!kIsComplex ? 1 : 2));
}
// const weight_t A_val = smem_a[state_idx];
scan_t thread_data[kNItems], thread_reverse_data[kNItems];
if constexpr (!kIsComplex) {
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
const float delta_a_exp = exp2f(delta_vals[i] * A_scaled);
thread_data[i] = make_float2(delta_a_exp, !kIsVariableB ? delta_vals[i] * float(u_vals[i]) : delta_vals[i] * float(u_vals[i]) * B_vals[i]);
if (i == 0) {
smem_delta_a[threadIdx.x == 0 ? state_idx + (chunk % 2) * MAX_DSTATE : threadIdx.x + 2 * MAX_DSTATE] = delta_a_exp;
} else {
thread_reverse_data[i - 1].x = delta_a_exp;
}
thread_reverse_data[i].y = dout_vals[i] *
(!kIsVariableC
? (!kIsVariableB ? B_val * C_val : C_val)
: (!kIsVariableB ? B_val * C_vals[i] : C_vals[i]));
}
__syncthreads();
thread_reverse_data[kNItems - 1].x = threadIdx.x == kNThreads - 1
? (chunk == params.n_chunks - 1 ? 1.f : smem_delta_a[state_idx + ((chunk + 1) % 2) * MAX_DSTATE])
: smem_delta_a[threadIdx.x + 1 + 2 * MAX_DSTATE];
// Initialize running total
scan_t running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? x[(chunk - 1) * params.dstate + state_idx] : make_float2(1.f, 0.f);
SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
typename Ktraits::BlockScanT(smem_scan).InclusiveScan(
thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
);
scan_t running_postfix = chunk < params.n_chunks - 1 && threadIdx.x % 32 == 0 ? smem_running_postfix[state_idx] : make_float2(1.f, 0.f);
SSMScanPrefixCallbackOp<weight_t> postfix_op(running_postfix);
typename Ktraits::BlockReverseScanT(smem_reverse_scan).InclusiveReverseScan(
thread_reverse_data, thread_reverse_data, SSMScanOp<weight_t>(), postfix_op
);
if (threadIdx.x == 0) { smem_running_postfix[state_idx] = postfix_op.running_prefix; }
weight_t dA_val = 0, dBC_val = 0;
weight_t dB_vals[kNItems], dC_vals[kNItems];
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
const float dx = thread_reverse_data[i].y;
const float ddelta_u = !kIsVariableB ? dx : dx * B_vals[i];
du_vals[i] += ddelta_u * delta_vals[i];
const float a = thread_data[i].y - (!kIsVariableB ? delta_vals[i] * float(u_vals[i]) : delta_vals[i] * float(u_vals[i]) * B_vals[i]);
ddelta_vals[i] += ddelta_u * float(u_vals[i]) + dx * A_val * a;
dA_val += dx * delta_vals[i] * a;
if constexpr (!kIsVariableB || !kIsVariableC) {
if constexpr (!kIsVariableB) { // dBC_val is dB_val
dBC_val += dout_vals[i] * (!kIsVariableC ? thread_data[i].y : thread_data[i].y * C_vals[i]);
} else { // dBC_val is dC_val
dBC_val += dout_vals[i] * thread_data[i].y;
}
}
if constexpr (kIsVariableB) { dB_vals[i] = dx * delta_vals[i] * float(u_vals[i]); }
if constexpr (kIsVariableC) {
dC_vals[i] = dout_vals[i] * (!kIsVariableB ? thread_data[i].y * B_val : thread_data[i].y);
}
}
// Block-exchange to make the atomicAdd's coalesced, otherwise they're much slower
if constexpr (kIsVariableB || kIsVariableC) {
if constexpr (kIsVariableB) {
typename Ktraits::BlockExchangeT(smem_exchange).BlockedToStriped(dB_vals, dB_vals);
}
if constexpr (kIsVariableC) {
auto &smem_exchange_C = !kIsVariableB ? smem_exchange : smem_exchange1;
typename Ktraits::BlockExchangeT(smem_exchange_C).BlockedToStriped(dC_vals, dC_vals);
}
const int seqlen_remaining = params.seqlen - chunk * kChunkSize - threadIdx.x;
weight_t *dB_cur = dB + state_idx * params.dB_dstate_stride + chunk * kChunkSize + threadIdx.x;
weight_t *dC_cur = dC + state_idx * params.dC_dstate_stride + chunk * kChunkSize + threadIdx.x;
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
if (i * kNThreads < seqlen_remaining) {
if constexpr (kIsVariableB) { gpuAtomicAdd(dB_cur + i * kNThreads, dB_vals[i]); }
if constexpr (kIsVariableC) { gpuAtomicAdd(dC_cur + i * kNThreads, dC_vals[i]); }
}
}
}
if constexpr (!kIsVariableB || !kIsVariableC) {
float2 dA_dBC_val = make_float2(dA_val, dBC_val);
dA_dBC_val = typename Ktraits::BlockReduceT(smem_reduce).Sum(dA_dBC_val);
dA_val = dA_dBC_val.x;
if (threadIdx.x == 0) {
smem_dbc[state_idx] = chunk == params.n_chunks - 1 ? dA_dBC_val.y : dA_dBC_val.y + smem_dbc[state_idx];
}
} else {
dA_val = typename Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dA_val);
}
if (threadIdx.x == 0) {
smem_da[state_idx] = chunk == params.n_chunks - 1 ? dA_val : dA_val + smem_da[state_idx];
}
} else {
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
// Pytorch's implementation of complex exp (which calls thrust) is very slow
complex_t delta_a_exp = cexp2f(delta_vals[i] * A_scaled);
weight_t B_delta_u_val = !kIsVariableB ? delta_vals[i] * float(u_vals[i]) : B_vals[i] * delta_vals[i] * float(u_vals[i]);
thread_data[i] = make_float4(delta_a_exp.real_, delta_a_exp.imag_, B_delta_u_val.real_, B_delta_u_val.imag_);
if (i == 0) {
smem_delta_a[threadIdx.x == 0 ? state_idx + (chunk % 2) * MAX_DSTATE : threadIdx.x + 2 * MAX_DSTATE] = delta_a_exp;
} else {
thread_reverse_data[i - 1].x = delta_a_exp.real_;
thread_reverse_data[i - 1].y = -delta_a_exp.imag_;
}
complex_t dout_BC = 2 * dout_vals[i]
* conj(!kIsVariableC
? (!kIsVariableB ? B_val * C_val : C_val)
: (!kIsVariableB ? B_val * C_vals[i] : C_vals[i]));
thread_reverse_data[i].z = dout_BC.real_;
thread_reverse_data[i].w = dout_BC.imag_;
}
__syncthreads();
complex_t delta_a_exp = threadIdx.x == kNThreads - 1
? (chunk == params.n_chunks - 1 ? 1.f : smem_delta_a[state_idx + ((chunk + 1) % 2) * MAX_DSTATE])
: smem_delta_a[threadIdx.x + 1 + 2 * MAX_DSTATE];
thread_reverse_data[kNItems - 1].x = delta_a_exp.real_;
thread_reverse_data[kNItems - 1].y = -delta_a_exp.imag_;
// Initialize running total
scan_t running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? x[(chunk - 1) * params.dstate + state_idx] : make_float4(1.f, 0.f, 0.f, 0.f);
SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
typename Ktraits::BlockScanT(smem_scan).InclusiveScan(
thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
);
scan_t running_postfix = chunk < params.n_chunks - 1 && threadIdx.x % 32 == 0 ? smem_running_postfix[state_idx] : make_float4(1.f, 0.f, 0.f, 0.f);
SSMScanPrefixCallbackOp<weight_t> postfix_op(running_postfix);
typename Ktraits::BlockReverseScanT(smem_reverse_scan).InclusiveReverseScan(
thread_reverse_data, thread_reverse_data, SSMScanOp<weight_t>(), postfix_op
);
if (threadIdx.x == 0) { smem_running_postfix[state_idx] = postfix_op.running_prefix; }
weight_t dA_val = 0, dBC_val = 0;
weight_t dB_vals[kNItems], dC_vals[kNItems];
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
complex_t x = complex_t(thread_data[i].z, thread_data[i].w);
complex_t dx = complex_t(thread_reverse_data[i].z, thread_reverse_data[i].w);
float ddelta_u = !kIsVariableB ? dx.real_ : (dx * conj(B_vals[i])).real_;
if constexpr (!kIsVariableB || !kIsVariableC) {
if constexpr (!kIsVariableB) { // dBC_val is dB_val
dBC_val += (2 * dout_vals[i]) * conj(!kIsVariableC ? x : x * C_vals[i]);
} else { // dBC_val is dC_val
dBC_val += (2 * dout_vals[i]) * conj(x);
}
}
const complex_t a_conj = conj(x - (!kIsVariableB ? delta_vals[i] * float(u_vals[i]) : delta_vals[i] * float(u_vals[i]) * B_vals[i]));
du_vals[i] += ddelta_u * delta_vals[i];
ddelta_vals[i] += ddelta_u * float(u_vals[i]) + (dx * conj(A_val) * a_conj).real_;
dA_val += delta_vals[i] * dx * a_conj;
if constexpr (kIsVariableB) { dB_vals[i] = dx * delta_vals[i] * float(u_vals[i]); }
if constexpr (kIsVariableC) {
dC_vals[i] = (2 * dout_vals[i]) * conj(!kIsVariableB ? x * B_val : x);
}
}
// Block-exchange to make the atomicAdd's coalesced, otherwise they're much slower
if constexpr (kIsVariableB || kIsVariableC) {
float dB_vals_f[kNItems * 2], dC_vals_f[kNItems * 2];
if constexpr (kIsVariableB) {
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
dB_vals_f[i * 2] = dB_vals[i].real_;
dB_vals_f[i * 2 + 1] = dB_vals[i].imag_;
}
typename Ktraits::BlockExchangeT(smem_exchange).BlockedToStriped(dB_vals_f, dB_vals_f);
}
if constexpr (kIsVariableC) {
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
dC_vals_f[i * 2] = dC_vals[i].real_;
dC_vals_f[i * 2 + 1] = dC_vals[i].imag_;
}
auto &smem_exchange_C = !kIsVariableB ? smem_exchange : smem_exchange1;
typename Ktraits::BlockExchangeT(smem_exchange_C).BlockedToStriped(dC_vals_f, dC_vals_f);
}
const int seqlen_remaining = (params.seqlen - chunk * kChunkSize) * 2 - threadIdx.x;
float *dB_cur = reinterpret_cast<float *>(dB) + state_idx * params.dB_dstate_stride + chunk * kChunkSize * 2 + threadIdx.x;
float *dC_cur = reinterpret_cast<float *>(dC) + state_idx * params.dC_dstate_stride + chunk * kChunkSize * 2 + threadIdx.x;
#pragma unroll
for (int i = 0; i < kNItems * 2; ++i) {
if (i * kNThreads < seqlen_remaining) {
if constexpr (kIsVariableB) { gpuAtomicAdd(dB_cur + i * kNThreads, dB_vals_f[i]); }
if constexpr (kIsVariableC) { gpuAtomicAdd(dC_cur + i * kNThreads, dC_vals_f[i]); }
}
}
}
if constexpr (!kIsVariableB || !kIsVariableC) {
float4 dA_dBC_val = make_float4(dA_val.real_, dA_val.imag_, dBC_val.real_, dBC_val.imag_);
dA_dBC_val = typename Ktraits::BlockReduceT(smem_reduce).Sum(dA_dBC_val);
dA_val = complex_t(dA_dBC_val.x, dA_dBC_val.y);
dBC_val = complex_t(dA_dBC_val.z, dA_dBC_val.w);
if (threadIdx.x == 0) {
smem_dbc[state_idx] = chunk == params.n_chunks - 1 ? dBC_val : dBC_val + smem_dbc[state_idx];
}
} else {
dA_val = typename Ktraits::BlockReduceComplexT(smem_reduce_complex).Sum(dA_val);
}
if (threadIdx.x == 0) {
smem_da[state_idx] = chunk == params.n_chunks - 1 ? dA_val : dA_val + smem_da[state_idx];
}
}
}
if constexpr (kDeltaSoftplus) {
__syncthreads();
input_t delta_vals_load[kNItems];
load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
delta -= kChunkSize;
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
float delta_val = float(delta_vals_load[i]) + delta_bias;
float delta_val_neg_exp = expf(-delta_val);
ddelta_vals[i] = delta_val <= 20.f
? ddelta_vals[i] / (1.f + delta_val_neg_exp)
: ddelta_vals[i];
}
}
for (int i = 0; i < kNItems; ++i) { ddelta_bias_val += ddelta_vals[i]; }
input_t *du = reinterpret_cast<input_t *>(params.du_ptr) + batch_id * params.du_batch_stride
+ dim_id * params.du_d_stride + chunk * kChunkSize;
input_t *ddelta = reinterpret_cast<input_t *>(params.ddelta_ptr) + batch_id * params.ddelta_batch_stride
+ dim_id * params.ddelta_d_stride + chunk * kChunkSize;
__syncthreads();
store_output<Ktraits>(du, du_vals, smem_store, params.seqlen - chunk * kChunkSize);
__syncthreads();
store_output<Ktraits>(ddelta, ddelta_vals, smem_store, params.seqlen - chunk * kChunkSize);
Bvar -= kChunkSize * (!kIsComplex ? 1 : 2);
Cvar -= kChunkSize * (!kIsComplex ? 1 : 2);
}
if (params.dD_ptr != nullptr) {
dD_val = typename Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dD_val);
if (threadIdx.x == 0) { gpuAtomicAdd(dD, dD_val); }
}
if (params.ddelta_bias_ptr != nullptr) {
__syncthreads();
ddelta_bias_val = typename Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(ddelta_bias_val);
if (threadIdx.x == 0) { gpuAtomicAdd(ddelta_bias, ddelta_bias_val); }
}
for (int state_idx = threadIdx.x; state_idx < params.dstate; state_idx += blockDim.x) {
gpuAtomicAdd(&(dA[state_idx * params.dA_dstate_stride]), smem_da[state_idx]);
weight_t dBC_val;
if (!kIsVariableB || !kIsVariableC) { dBC_val = smem_dbc[state_idx]; }
if constexpr (!kIsVariableB) {
gpuAtomicAdd(&(dB[state_idx * params.dB_dstate_stride]),
!kIsVariableC ? dBC_val * conj(C[state_idx * params.C_dstate_stride]) : dBC_val);
}
if constexpr (!kIsVariableC) {
gpuAtomicAdd(&(dC[state_idx * params.dC_dstate_stride]),
!kIsVariableB ? dBC_val * conj(B[state_idx * params.B_dstate_stride]) : dBC_val);
}
}
}
template<int kNThreads, int kNItems, typename input_t, typename weight_t>
void selective_scan_bwd_launch(SSMParamsBwd ¶ms, cudaStream_t stream) {
BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
BOOL_SWITCH(params.is_variable_B, kIsVariableB, [&] {
BOOL_SWITCH(params.is_variable_C, kIsVariableC, [&] {
BOOL_SWITCH(params.delta_softplus, kDeltaSoftplus, [&] {
BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] {
using Ktraits = Selective_Scan_bwd_kernel_traits<kNThreads, kNItems, kIsEvenLen, kIsVariableB, kIsVariableC, kDeltaSoftplus, kHasZ, input_t, weight_t>;
// using Ktraits = Selective_Scan_bwd_kernel_traits<kNThreads, kNItems, true, kIsVariableB, kIsVariableC, kDeltaSoftplus, kHasZ, input_t, weight_t>;
// TODO: check this
constexpr int kSmemSize = Ktraits::kSmemSize + MAX_DSTATE * sizeof(typename Ktraits::scan_t) + (kNThreads + 4 * MAX_DSTATE) * sizeof(typename Ktraits::weight_t);
dim3 grid(params.batch, params.dim);
auto kernel = &selective_scan_bwd_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
#ifndef USE_ROCM
C10_CUDA_CHECK(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
#else
C10_CUDA_CHECK(cudaFuncSetAttribute(
(void *) kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
std::cerr << "Warning (selective_scan_bwd_kernel): attempting to set maxDynamicSharedMemorySize on an AMD GPU which is currently a non-op (in ROCm versions <= 6.1). This might lead to undefined behavior. \n" << std::endl;
#endif
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
});
});
});
});
}
template<typename input_t, typename weight_t>
void selective_scan_bwd_cuda(SSMParamsBwd ¶ms, cudaStream_t stream) {
#ifndef USE_ROCM
if (params.seqlen <= 128) {
selective_scan_bwd_launch<32, 4, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 256) {
selective_scan_bwd_launch<32, 8, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 512) {
selective_scan_bwd_launch<32, 16, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 1024) {
selective_scan_bwd_launch<64, 16, input_t, weight_t>(params, stream);
} else {
selective_scan_bwd_launch<128, 16, input_t, weight_t>(params, stream);
}
#else
if (params.seqlen <= 256) {
selective_scan_bwd_launch<64, 4, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 512) {
selective_scan_bwd_launch<64, 8, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 1024) {
selective_scan_bwd_launch<64, 16, input_t, weight_t>(params, stream);
} else {
selective_scan_bwd_launch<128, 16, input_t, weight_t>(params, stream);
}
#endif
}
================================================
FILE: csrc/selective_scan/selective_scan_common.h
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#pragma once
#ifndef USE_ROCM
#include <cuda_bf16.h>
#else
#include <hip/hip_bf16.h>
#endif
#include <cuda_fp16.h>
#include <c10/util/complex.h> // For scalar_value_type
#ifndef USE_ROCM
constexpr size_t custom_max(std::initializer_list<size_t> ilist)
{
return std::max(ilist);
}
template<typename T>
constexpr T constexpr_min(T a, T b) {
return std::min(a, b);
}
#else
constexpr size_t custom_max(std::initializer_list<size_t> ilist)
{
return *std::max_element(ilist.begin(), ilist.end());
}
template<typename T>
constexpr T constexpr_min(T a, T b) {
return a < b ? a : b;
}
#endif
#define MAX_DSTATE 256
using complex_t = c10::complex<float>;
inline __device__ float2 operator+(const float2 & a, const float2 & b){
return {a.x + b.x, a.y + b.y};
}
inline __device__ float3 operator+(const float3 &a, const float3 &b) {
return {a.x + b.x, a.y + b.y, a.z + b.z};
}
inline __device__ float4 operator+(const float4 & a, const float4 & b){
return {a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w};
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template<int BYTES> struct BytesToType {};
template<> struct BytesToType<16> {
using Type = uint4;
static_assert(sizeof(Type) == 16);
};
template<> struct BytesToType<8> {
using Type = uint64_t;
static_assert(sizeof(Type) == 8);
};
template<> struct BytesToType<4> {
using Type = uint32_t;
static_assert(sizeof(Type) == 4);
};
template<> struct BytesToType<2> {
using Type = uint16_t;
static_assert(sizeof(Type) == 2);
};
template<> struct BytesToType<1> {
using Type = uint8_t;
static_assert(sizeof(Type) == 1);
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename scalar_t, int N>
struct Converter{
static inline __device__ void to_float(const scalar_t (&src)[N], float (&dst)[N]) {
#pragma unroll
for (int i = 0; i < N; ++i) { dst[i] = src[i]; }
}
};
template<int N>
struct Converter<at::Half, N>{
static inline __device__ void to_float(const at::Half (&src)[N], float (&dst)[N]) {
static_assert(N % 2 == 0);
auto &src2 = reinterpret_cast<const half2 (&)[N / 2]>(src);
auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
#pragma unroll
for (int i = 0; i < N / 2; ++i) { dst2[i] = __half22float2(src2[i]); }
}
};
#if __CUDA_ARCH__ >= 800
template<int N>
struct Converter<at::BFloat16, N>{
static inline __device__ void to_float(const at::BFloat16 (&src)[N], float (&dst)[N]) {
static_assert(N % 2 == 0);
auto &src2 = reinterpret_cast<const nv_bfloat162 (&)[N / 2]>(src);
auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
#pragma unroll
for (int i = 0; i < N / 2; ++i) { dst2[i] = __bfloat1622float2(src2[i]); }
}
};
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
// From https://stackoverflow.com/questions/9860711/cucomplex-h-and-exp
// and https://forums.developer.nvidia.com/t/complex-number-exponential-function/24696
__device__ __forceinline__ complex_t cexp2f(complex_t z) {
float t = exp2f(z.real_);
float c, s;
sincosf(z.imag_, &s, &c);
return complex_t(c * t, s * t);
}
__device__ __forceinline__ complex_t cexpf(complex_t z) {
float t = expf(z.real_);
float c, s;
sincosf(z.imag_, &s, &c);
return complex_t(c * t, s * t);
}
template<typename scalar_t> struct SSMScanOp;
template<>
struct SSMScanOp<float> {
__device__ __forceinline__ float2 operator()(const float2 &ab0, const float2 &ab1) const {
return make_float2(ab1.x * ab0.x, ab1.x * ab0.y + ab1.y);
}
};
template<>
struct SSMScanOp<complex_t> {
__device__ __forceinline__ float4 operator()(const float4 &ab0, const float4 &ab1) const {
complex_t a0 = complex_t(ab0.x, ab0.y);
complex_t b0 = complex_t(ab0.z, ab0.w);
complex_t a1 = complex_t(ab1.x, ab1.y);
complex_t b1 = complex_t(ab1.z, ab1.w);
complex_t out_a = a1 * a0;
complex_t out_b = a1 * b0 + b1;
return make_float4(out_a.real_, out_a.imag_, out_b.real_, out_b.imag_);
}
};
// A stateful callback functor that maintains a running prefix to be applied
// during consecutive scan operations.
template <typename scalar_t> struct SSMScanPrefixCallbackOp {
using scan_t = std::conditional_t<std::is_same_v<scalar_t, float>, float2, float4>;
scan_t running_prefix;
// Constructor
__device__ SSMScanPrefixCallbackOp(scan_t running_prefix_) : running_prefix(running_prefix_) {}
// Callback operator to be entered by the first warp of threads in the block.
// Thread-0 is responsible for returning a value for seeding the block-wide scan.
__device__ scan_t operator()(scan_t block_aggregate) {
scan_t old_prefix = running_prefix;
running_prefix = SSMScanOp<scalar_t>()(running_prefix, block_aggregate);
return old_prefix;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename Ktraits>
inline __device__ void load_input(typename Ktraits::input_t *u,
typename Ktraits::input_t (&u_vals)[Ktraits::kNItems],
typename Ktraits::BlockLoadT::TempStorage &smem_load,
int seqlen) {
if constexpr (Ktraits::kIsEvenLen) {
auto& smem_load_vec = reinterpret_cast<typename Ktraits::BlockLoadVecT::TempStorage&>(smem_load);
using vec_t = typename Ktraits::vec_t;
typename Ktraits::BlockLoadVecT(smem_load_vec).Load(
reinterpret_cast<vec_t*>(u),
reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(u_vals)
#ifdef USE_ROCM
, Ktraits::kNThreads * Ktraits::kNLoads
#endif
);
} else {
typename Ktraits::BlockLoadT(smem_load).Load(u, u_vals, seqlen, 0.f);
}
}
template<typename Ktraits>
inline __device__ void load_weight(typename Ktraits::input_t *Bvar,
typename Ktraits::weight_t (&B_vals)[Ktraits::kNItems],
typename Ktraits::BlockLoadWeightT::TempStorage &smem_load_weight,
int seqlen) {
constexpr int kNItems = Ktraits::kNItems;
if constexpr (!Ktraits::kIsComplex) {
typename Ktraits::input_t B_vals_load[kNItems];
if constexpr (Ktraits::kIsEvenLen) {
auto& smem_load_weight_vec = reinterpret_cast<typename Ktraits::BlockLoadWeightVecT::TempStorage&>(smem_load_weight);
using vec_t = typename Ktraits::vec_t;
typename Ktraits::BlockLoadWeightVecT(smem_load_weight_vec).Load(
reinterpret_cast<vec_t*>(Bvar),
reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(B_vals_load)
);
} else {
typename Ktraits::BlockLoadWeightT(smem_load_weight).Load(Bvar, B_vals_load, seqlen, 0.f);
}
// #pragma unroll
// for (int i = 0; i < kNItems; ++i) { B_vals[i] = B_vals_load[i]; }
Converter<typename Ktraits::input_t, kNItems>::to_float(B_vals_load, B_vals);
} else {
typename Ktraits::input_t B_vals_load[kNItems * 2];
if constexpr (Ktraits::kIsEvenLen) {
auto& smem_load_weight_vec = reinterpret_cast<typename Ktraits::BlockLoadWeightVecT::TempStorage&>(smem_load_weight);
using vec_t = typename Ktraits::vec_t;
typename Ktraits::BlockLoadWeightVecT(smem_load_weight_vec).Load(
reinterpret_cast<vec_t*>(Bvar),
reinterpret_cast<vec_t(&)[Ktraits::kNLoads * 2]>(B_vals_load)
);
} else {
typename Ktraits::BlockLoadWeightT(smem_load_weight).Load(Bvar, B_vals_load, seqlen, 0.f);
}
#pragma unroll
for (int i = 0; i < kNItems; ++i) { B_vals[i] = complex_t(B_vals_load[i * 2], B_vals_load[i * 2 + 1]); }
}
}
template<typename Ktraits>
inline __device__ void store_output(typename Ktraits::input_t *out,
const float (&out_vals)[Ktraits::kNItems],
typename Ktraits::BlockStoreT::TempStorage &smem_store,
int seqlen) {
typename Ktraits::input_t write_vals[Ktraits::kNItems];
#pragma unroll
for (int i = 0; i < Ktraits::kNItems; ++i) { write_vals[i] = out_vals[i]; }
if constexpr (Ktraits::kIsEvenLen) {
auto& smem_store_vec = reinterpret_cast<typename Ktraits::BlockStoreVecT::TempStorage&>(smem_store);
using vec_t = typename Ktraits::vec_t;
typename Ktraits::BlockStoreVecT(smem_store_vec).Store(
reinterpret_cast<vec_t*>(out),
reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(write_vals)
);
} else {
typename Ktraits::BlockStoreT(smem_store).Store(out, write_vals, seqlen);
}
}
================================================
FILE: csrc/selective_scan/selective_scan_fwd_bf16.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_fwd_kernel.cuh"
template void selective_scan_fwd_cuda<at::BFloat16, float>(SSMParamsBase ¶ms, cudaStream_t stream);
template void selective_scan_fwd_cuda<at::BFloat16, complex_t>(SSMParamsBase ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_fwd_fp16.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_fwd_kernel.cuh"
template void selective_scan_fwd_cuda<at::Half, float>(SSMParamsBase ¶ms, cudaStream_t stream);
template void selective_scan_fwd_cuda<at::Half, complex_t>(SSMParamsBase ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_fwd_fp32.cu
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Split into multiple files to compile in paralell
#include "selective_scan_fwd_kernel.cuh"
template void selective_scan_fwd_cuda<float, float>(SSMParamsBase ¶ms, cudaStream_t stream);
template void selective_scan_fwd_cuda<float, complex_t>(SSMParamsBase ¶ms, cudaStream_t stream);
================================================
FILE: csrc/selective_scan/selective_scan_fwd_kernel.cuh
================================================
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#pragma once
#include <c10/util/BFloat16.h>
#include <c10/util/Half.h>
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
#ifndef USE_ROCM
#include <cub/block/block_load.cuh>
#include <cub/block/block_store.cuh>
#include <cub/block/block_scan.cuh>
#else
#include <hipcub/hipcub.hpp>
namespace cub = hipcub;
#endif
#include "selective_scan.h"
#include "selective_scan_common.h"
#include "static_switch.h"
template<int kNThreads_, int kNItems_, int kNRows_, bool kIsEvenLen_,
bool kIsVariableB_, bool kIsVariableC_,
bool kHasZ_, typename input_t_, typename weight_t_>
struct Selective_Scan_fwd_kernel_traits {
static_assert(kNItems_ % 4 == 0);
using input_t = input_t_;
using weight_t = weight_t_;
static constexpr int kNThreads = kNThreads_;
// Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy.
static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3;
static constexpr int kNItems = kNItems_;
static constexpr int kNRows = kNRows_;
static constexpr int kNBytes = sizeof(input_t);
static_assert(kNBytes == 2 || kNBytes == 4);
static constexpr int kNElts = kNBytes == 4 ? 4 : constexpr_min(8, kNItems);
static_assert(kNItems % kNElts == 0);
static constexpr int kNLoads = kNItems / kNElts;
static constexpr bool kIsComplex = std::is_same_v<weight_t, complex_t>;
static constexpr bool kIsEvenLen = kIsEvenLen_;
static constexpr bool kIsVariableB = kIsVariableB_;
static constexpr bool kIsVariableC = kIsVariableC_;
static constexpr bool kHasZ = kHasZ_;
static constexpr bool kDirectIO = kIsEvenLen && kNLoads == 1;
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
using scan_t = std::conditional_t<!kIsComplex, float2, float4>;
using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
!kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE : cub::BLOCK_LOAD_DIRECT>;
using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, !kIsComplex ? kNItems : kNItems * 2, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, !kIsComplex ? kNLoads : kNLoads * 2,
!kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE : cub::BLOCK_LOAD_DIRECT>;
using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads,
!kDirectIO ? cub::BLOCK_STORE_WARP_TRANSPOSE : cub::BLOCK_STORE_DIRECT>;
// using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
// using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
static constexpr int kSmemIOSize = custom_max({sizeof(typename BlockLoadT::TempStorage),
sizeof(typename BlockLoadVecT::TempStorage),
(int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightT::TempStorage),
(int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightVecT::TempStorage),
sizeof(typename BlockStoreT::TempStorage),
sizeof(typename BlockStoreVecT::TempStorage)});
static constexpr int kSmemSize = kSmemIOSize + sizeof(typename BlockScanT::TempStorage);
};
template<typename Ktraits>
__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
void selective_scan_fwd_kernel(SSMParamsBase params) {
constexpr bool kIsComplex = Ktraits::kIsComplex;
constexpr bool kIsVariableB = Ktraits::kIsVariableB;
constexpr bool kIsVariableC = Ktraits::kIsVariableC;
constexpr bool kHasZ = Ktraits::kHasZ;
constexpr int kNThreads = Ktraits::kNThreads;
constexpr int kNItems = Ktraits::kNItems;
constexpr int kNRows = Ktraits::kNRows;
constexpr bool kDirectIO = Ktraits::kDirectIO;
using input_t = typename Ktraits::input_t;
using weight_t = typename Ktraits::weight_t;
using scan_t = typename Ktraits::scan_t;
// Shared memory.
extern __shared__ char smem_[];
// cast to lvalue reference of expected type
// char *smem_loadstorescan = smem_ + 2 * MAX_DSTATE * sizeof(weight_t);
// auto& smem_load = reinterpret_cast<typename BlockLoadT::TempStorage&>(smem_ + 2 * MAX_DSTATE * sizeof(weight_t));
// auto& smem_load = reinterpret_cast<typename BlockLoadT::TempStorage&>(smem_loadstorescan);
auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
// weight_t *smem_a = reinterpret_cast<weight_t *>(smem_ + smem_loadstorescan_size);
// weight_t *smem_bc = reinterpret_cast<weight_t *>(smem_a + MAX_DSTATE);
scan_t *smem_running_prefix = reinterpret_cast<scan_t *>(smem_ + Ktraits::kSmemSize);
const int batch_id = blockIdx.x;
const int dim_id = blockIdx.y;
const int group_id = dim_id / (params.dim_ngroups_ratio);
input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+ dim_id * kNRows * params.u_d_stride;
input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+ dim_id * kNRows * params.delta_d_stride;
weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id * kNRows * params.A_d_stride;
weight_t *B = reinterpret_cast<weight_t *>(params.B_ptr) + dim_id * kNRows * params.B_d_stride;
input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
weight_t *C = reinterpret_cast<weight_t *>(params.C_ptr) + dim_id * kNRows * params.C_d_stride;
input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
scan_t *x = reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id * kNRows) * params.n_chunks * params.dstate;
float D_val[kNRows] = {0};
if (params.D_ptr != nullptr) {
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
D_val[r] = reinterpret_cast<float *>(params.D_ptr)[dim_id * kNRows + r];
}
}
float delta_bias[kNRows] = {0};
if (params.delta_bias_ptr != nullptr) {
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
delta_bias[r] = reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id * kNRows + r];
}
}
// for (int state_idx = threadIdx.x; state_idx < params.dstate; state_idx += blockDim.x) {
// smem_a[state_idx] = A[state_idx * params.A_dstate_stride];
// smem_bc[state_idx] = B[state_idx * params.B_dstate_stride] * C[state_idx * params.C_dstate_stride];
// }
constexpr int kChunkSize = kNThreads * kNItems;
for (int chunk = 0; chunk < params.n_chunks; ++chunk) {
input_t u_vals[kNRows][kNItems], delta_vals_load[kNRows][kNItems];
__syncthreads();
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
if constexpr (!kDirectIO) {
if (r > 0) { __syncthreads(); }
}
load_input<Ktraits>(u + r * params.u_d_stride, u_vals[r], smem_load, params.seqlen - chunk * kChunkSize);
if constexpr (!kDirectIO) { __syncthreads(); }
load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, params.seqlen - chunk * kChunkSize);
}
u += kChunkSize;
delta += kChunkSize;
float delta_vals[kNRows][kNItems], delta_u_vals[kNRows][kNItems], out_vals[kNRows][kNItems];
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
float u_val = float(u_vals[r][i]);
delta_vals[r][i] = float(delta_vals_load[r][i]) + delta_bias[r];
if (params.delta_softplus) {
delta_vals[r][i] = delta_vals[r][i] <= 20.f ? log1pf(expf(delta_vals[r][i])) : delta_vals[r][i];
}
delta_u_vals[r][i] = delta_vals[r][i] * u_val;
out_vals[r][i] = D_val[r] * u_val;
}
}
__syncthreads();
for (int state_idx = 0; state_idx < params.dstate; ++state_idx) {
weight_t A_val[kNRows];
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
A_val[r] = A[state_idx * params.A_dstate_stride + r * params.A_d_stride];
// Multiply the real part of A with LOG2E so we can use exp2f instead of expf.
constexpr float kLog2e = M_LOG2E;
if constexpr (!kIsComplex) {
A_val[r] *= kLog2e;
} else {
A_val[r].real_ *= kLog2e;
}
}
// This variable holds B * C if both B and C are constant across seqlen. If only B varies
// across seqlen, this holds C. If only C varies across seqlen, this holds B.
// If both B and C vary, this is unused.
weight_t BC_val[kNRows];
weight_t B_vals[kNItems], C_vals[kNItems];
if constexpr (kIsVariableB) {
load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
smem_load_weight, (params.seqlen - chunk * kChunkSize) * (!kIsComplex ? 1 : 2));
if constexpr (!kIsVariableC) {
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
BC_val[r] = C[state_idx * params.C_dstate_stride + r * params.C_d_stride];
}
}
}
if constexpr (kIsVariableC) {
auto &smem_load_weight_C = !kIsVariableB ? smem_load_weight : smem_load_weight1;
load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
smem_load_weight_C, (params.seqlen - chunk * kChunkSize) * (!kIsComplex ? 1 : 2));
if constexpr (!kIsVariableB) {
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
BC_val[r] = B[state_idx * params.B_dstate_stride + r * params.B_d_stride];
}
}
}
if constexpr (!kIsVariableB && !kIsVariableC) {
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
BC_val[r] = B[state_idx * params.B_dstate_stride + r * params.B_d_stride] * C[state_idx * params.C_dstate_stride + r * params.C_d_stride];
}
}
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
if (r > 0) { __syncthreads(); } // Scan could be using the same smem
scan_t thread_data[kNItems];
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
if constexpr (!kIsComplex) {
thread_data[i] = make_float2(exp2f(delta_vals[r][i] * A_val[r]),
!kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i]);
if constexpr (!Ktraits::kIsEvenLen) { // So that the last state is correct
if (threadIdx.x * kNItems + i >= params.seqlen - chunk * kChunkSize) {
thread_data[i] = make_float2(1.f, 0.f);
}
}
} else {
// Pytorch's implementation of complex exp (which calls thrust) is very slow
complex_t delta_a_exp = cexp2f(delta_vals[r][i] * A_val[r]);
weight_t B_delta_u_val = !kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i];
thread_data[i] = make_float4(delta_a_exp.real_, delta_a_exp.imag_, B_delta_u_val.real_, B_delta_u_val.imag_);
if constexpr (!Ktraits::kIsEvenLen) { // So that the last state is correct
if (threadIdx.x * kNItems + i >= params.seqlen - chunk * kChunkSize) {
thread_data[i] = make_float4(1.f, 0.f, 0.f, 0.f);
}
}
}
}
// Initialize running total
scan_t running_prefix;
if constexpr (!kIsComplex) {
// If we use WARP_SCAN then all lane 0 of all warps (not just thread 0) needs to read
running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? smem_running_prefix[state_idx + r * MAX_DSTATE] : make_float2(1.f, 0.f);
// running_prefix = chunk > 0 && threadIdx.x == 0 ? smem_running_prefix[state_idx] : make_float2(1.f, 0.f);
} else {
running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? smem_running_prefix[state_idx + r * MAX_DSTATE] : make_float4(1.f, 0.f, 0.f, 0.f);
// running_prefix = chunk > 0 && threadIdx.x == 0 ? smem_running_prefix[state_idx] : make_float4(1.f, 0.f, 0.f, 0.f);
}
SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
typename Ktraits::BlockScanT(smem_scan).InclusiveScan(
thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
);
// There's a syncthreads in the scan op, so we don't need to sync here.
// Unless there's only 1 warp, but then it's the same thread (0) reading and writing.
if (threadIdx.x == 0) {
smem_running_prefix[state_idx] = prefix_op.running_prefix;
x[(r * params.n_chunks + chunk) * params.dstate + state_idx] = prefix_op.running_prefix;
}
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
const weight_t C_val = !kIsVariableC
? BC_val[r]
: (!kIsVariableB ? BC_val[r] * C_vals[i] : C_vals[i]);
if constexpr (!kIsComplex) {
out_vals[r][i] += thread_data[i].y * C_val;
} else {
out_vals[r][i] += (complex_t(thread_data[i].z, thread_data[i].w) * C_val).real_ * 2;
}
}
}
}
input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + batch_id * params.out_batch_stride
+ dim_id * kNRows * params.out_d_stride + chunk * kChunkSize;
__syncthreads();
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
if constexpr (!kDirectIO) {
if (r > 0) { __syncthreads(); }
}
store_output<Ktraits>(out + r * params.out_d_stride, out_vals[r], smem_store, params.seqlen - chunk * kChunkSize);
}
if constexpr (kHasZ) {
input_t *z = reinterpret_cast<input_t *>(params.z_ptr) + batch_id * params.z_batch_stride
+ dim_id * kNRows * params.z_d_stride + chunk * kChunkSize;
input_t *out_z = reinterpret_cast<input_t *>(params.out_z_ptr) + batch_id * params.out_z_batch_stride
+ dim_id * kNRows * params.out_z_d_stride + chunk * kChunkSize;
#pragma unroll
for (int r = 0; r < kNRows; ++r) {
input_t z_vals[kNItems];
__syncthreads();
load_input<Ktraits>(z + r * params.z_d_stride, z_vals, smem_load, params.seqlen - chunk * kChunkSize);
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
float z_val = z_vals[i];
out_vals[r][i] *= z_val / (1 + expf(-z_val));
}
__syncthreads();
store_output<Ktraits>(out_z + r * params.out_z_d_stride, out_vals[r], smem_store, params.seqlen - chunk * kChunkSize);
}
}
Bvar += kChunkSize * (!kIsComplex ? 1 : 2);
Cvar += kChunkSize * (!kIsComplex ? 1 : 2);
}
}
template<int kNThreads, int kNItems, typename input_t, typename weight_t>
void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) {
// Only kNRows == 1 is tested for now, which ofc doesn't differ from previously when we had each block
// processing 1 row.
constexpr int kNRows = 1;
BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
BOOL_SWITCH(params.is_variable_B, kIsVariableB, [&] {
BOOL_SWITCH(params.is_variable_C, kIsVariableC, [&] {
BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] {
using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kIsVariableB, kIsVariableC, kHasZ, input_t, weight_t>;
constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t);
dim3 grid(params.batch, params.dim / kNRows);
// Had to change this substantially since potentially the hip
// interface for setting kernel launch attributes is slightly different from
// cuda's. In particualar, it seems to expect a plain const void * pointer.
auto kernel = &selective_scan_fwd_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
#ifndef USE_ROCM
C10_CUDA_CHECK(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
#else
C10_CUDA_CHECK(cudaFuncSetAttribute(
(void *) kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
std::cerr << "Warning (selective_scan_fwd_kernel): attempting to set maxDynamicSharedMemorySize on an AMD GPU which is currently a non-op (in ROCm versions <= 6.1). This might lead to undefined behavior. \n" << std::endl;
#endif
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
});
});
});
}
template<typename input_t, typename weight_t>
void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream) {
#ifndef USE_ROCM
if (params.seqlen <= 128) {
selective_scan_fwd_launch<32, 4, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 256) {
selective_scan_fwd_launch<32, 8, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 512) {
selective_scan_fwd_launch<32, 16, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 1024) {
selective_scan_fwd_launch<64, 16, input_t, weight_t>(params, stream);
} else {
selective_scan_fwd_launch<128, 16, input_t, weight_t>(params, stream);
}
#else
if (params.seqlen <= 256) {
selective_scan_fwd_launch<64, 4, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 512) {
selective_scan_fwd_launch<64, 8, input_t, weight_t>(params, stream);
} else if (params.seqlen <= 1024) {
selective_scan_fwd_launch<64, 16, input_t, weight_t>(params, stream);
} else {
selective_scan_fwd_launch<128, 16, input_t, weight_t>(params, stream);
}
#endif
}
================================================
FILE: csrc/selective_scan/static_switch.h
================================================
// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
#pragma once
/// @param COND - a boolean expression to switch by
/// @param CONST_NAME - a name given for the constexpr bool variable.
/// @param ... - code to execute for true and false
///
/// Usage:
/// ```
/// BOOL_SWITCH(flag, BoolConst, [&] {
/// some_function<BoolConst>(...);
/// });
/// ```
#define BOOL_SWITCH(COND, CONST_NAME, ...) \
[&] { \
if (COND) { \
constexpr bool CONST_NAME = true; \
return __VA_ARGS__(); \
} else { \
constexpr bool CONST_NAME = false; \
return __VA_ARGS__(); \
} \
}()
================================================
FILE: csrc/selective_scan/uninitialized_copy.cuh
================================================
/******************************************************************************
* Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
* * Redistributions of source code must retain the above copyright
* notice, this list of conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above copyright
* notice, this list of conditions and the following disclaimer in the
* documentation and/or other materials provided with the distribution.
* * Neither the name of the NVIDIA CORPORATION nor the
* names of its contributors may be used to endorse or promote products
* derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
* ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
******************************************************************************/
#pragma once
#ifndef USE_ROCM
#include <cub/config.cuh>
#include <cuda/std/type_traits>
#else
#include <hipcub/hipcub.hpp>
// Map ::cuda::std to the standard std namespace
namespace cuda {
namespace std = ::std;
}
#endif
namespace detail
{
#if defined(_NVHPC_CUDA)
template <typename T, typename U>
__host__ __device__ void uninitialized_copy(T *ptr, U &&val)
{
// NVBug 3384810
new (ptr) T(::cuda::std::forward<U>(val));
}
#else
template <typename T,
typename U,
typename ::cuda::std::enable_if<
::cuda::std::is_trivially_copyable<T>::value,
int
>::type = 0>
__host__ __device__ void uninitialized_copy(T *ptr, U &&val)
{
*ptr = ::cuda::std::forward<U>(val);
}
template <typename T,
typename U,
typename ::cuda::std::enable_if<
!::cuda::std::is_trivially_copyable<T>::value,
int
>::type = 0>
__host__ __device__ void uninitialized_copy(T *ptr, U &&val)
{
new (ptr) T(::cuda::std::forward<U>(val));
}
#endif
} // namespace detail
================================================
FILE: evals/lm_harness_eval.py
================================================
import torch
import transformers
from transformers import AutoTokenizer
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from lm_eval.api.model import LM
from lm_eval.models.huggingface import HFLM
from lm_eval.api.registry import register_model
from lm_eval.__main__ import cli_evaluate
@register_model("mamba")
class MambaEvalWrapper(HFLM):
AUTO_MODEL_CLASS = transformers.AutoModelForCausalLM
def __init__(self, pretrained="state-spaces/mamba-2.8b", max_length=2048, batch_size=None, device="cuda",
dtype=torch.float16):
LM.__init__(self)
self._model = MambaLMHeadModel.from_pretrained(pretrained, device=device, dtype=dtype)
self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
self.vocab_size = self.tokenizer.vocab_size
self._batch_size = int(batch_size) if batch_size is not None else 64
self._max_length = max_length
self._device = torch.device(device)
@property
def batch_size(self):
return self._batch_size
def _model_generate(self, context, max_length, stop, **generation_kwargs):
raise NotImplementedError()
if __name__ == "__main__":
cli_evaluate()
================================================
FILE: mamba_ssm/__init__.py
================================================
__version__ = "2.3.1"
from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn
from mamba_ssm.modules.mamba_simple import Mamba
from mamba_ssm.modules.mamba2 import Mamba2
from mamba_ssm.modules.mamba3 import Mamba3
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
================================================
FILE: mamba_ssm/distributed/__init__.py
================================================
================================================
FILE: mamba_ssm/distributed/distributed_utils.py
================================================
from typing import Optional
import torch
from torch import Tensor
from torch.distributed import ProcessGroup
# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for
# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent
# version of PyTorch. The following 4 lines are for backward compatibility with
# older PyTorch.
if "all_gather_into_tensor" not in dir(torch.distributed):
torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base
if "reduce_scatter_tensor" not in dir(torch.distributed):
torch.distributed.reduce_scatter_tensor = torch.distributed._reduce_scatter_base
# Raw operation, does not support autograd, but does support async
def all_gather_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False):
world_size = torch.distributed.get_world_size(process_group)
output = torch.empty(
world_size * input_.shape[0], *input_.shape[1:], dtype=input_.dtype, device=input_.device
)
handle = torch.distributed.all_gather_into_tensor(
output, input_.contiguous(), group=process_group, async_op=async_op
)
return output, handle
# Raw operation, does not support autograd, but does support async
def reduce_scatter_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False):
world_size = torch.distributed.get_world_size(process_group)
assert input_.shape[0] % world_size == 0
output = torch.empty(
input_.shape[0] // world_size, *input_.shape[1:], dtype=input_.dtype, device=input_.device
)
handle = torch.distributed.reduce_scatter_tensor(
output, input_.contiguous(), group=process_group, async_op=async_op
)
return output, handle
# Raw operation, does not support autograd, but does support async
def all_reduce_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False):
input_ = input_.contiguous()
handle = torch.distributed.all_reduce(input_, group=process_group, async_op=async_op)
return input_, handle
class AllGatherFunc(torch.autograd.Function):
"""Gather the input from sequence parallel region and concatenate."""
@staticmethod
def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
ctx.process_group = process_group
output, _ = all_gather_raw(input_, process_group)
return output
@staticmethod
def backward(ctx, grad_output: Tensor):
grad_input, _ = reduce_scatter_raw(grad_output, ctx.process_group)
return grad_input, None
# Supports autograd, but does not support async
all_gather = AllGatherFunc.apply
class ReduceScatterFunc(torch.autograd.Function):
"""Reduce scatter the input from the sequence parallel region and concatenate."""
@staticmethod
def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
ctx.process_group = process_group
output, _ = reduce_scatter_raw(input_, process_group)
return output
@staticmethod
def backward(ctx, grad_output: Tensor):
grad_input, _ = all_gather_raw(grad_output, ctx.process_group)
return grad_input, None
# Supports autograd, but does not support async
reduce_scatter = ReduceScatterFunc.apply
class AllReduceFunc(torch.autograd.Function):
"""Gather the input from sequence parallel region and concatenate."""
@staticmethod
def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
ctx.process_group = process_group
output, _ = all_reduce_raw(input_, process_group)
return output
@staticmethod
def backward(ctx, grad_output: Tensor):
return grad_output, None
# Supports autograd, but does not support async
all_reduce = AllReduceFunc.apply
def sync_shared_params(model: torch.nn.Module, process_group: ProcessGroup):
# We want to iterate over parameters with _shared_params=True in the same order,
# as different ranks might have different number of parameters (e.g., only rank 0 has bias).
pamams_shared = {
name: p for name, p in model.named_parameters() if getattr(p, "_shared_params", False)
}
for _, p in sorted(pamams_shared.items()):
with torch.no_grad():
# Broadcast needs src to be global rank, not group rank
torch.distributed.broadcast(
p, src=torch.distributed.get_global_rank(process_group, 0), group=process_group
)
# Ref: https://github.com/NVIDIA/Megatron-LM/blob/52e636888cccc41e931251c417a7181fc36de926/megatron/optimizer/optimizer.py#L256
def allreduce_sequence_parallel_grad(model: torch.nn.Module, process_group: ProcessGroup):
# We want to iterate over parameters with _sequence_parallel=True in the same order,
# as different ranks might have different number of parameters (e.g., only rank 0 has bias).
params_seqparallel = {
name: p for name, p in model.named_parameters() if getattr(p, "_sequence_parallel", False)
}
grads = [p.grad for _, p in sorted(params_seqparallel.items())]
if grads:
with torch.no_grad():
coalesced = torch._utils._flatten_dense_tensors(grads)
torch.distributed.all_reduce(coalesced, group=process_group)
for buf, synced in zip(grads, torch._utils._unflatten_dense_tensors(coalesced, grads)):
buf.copy_(synced)
def get_dim_for_local_rank(dim: int, world_size: int, local_rank: int, multiple_of: int = 1) -> int:
"""Get the dim for the local rank derived from splitting dim on world_size processes.
The split may not be even across the world_size processes.
"""
multiple = dim // multiple_of
div = multiple // world_size
mod = multiple % world_size
local_multiple = div + int(local_rank < mod)
return local_multiple * multiple_of
================================================
FILE: mamba_ssm/distributed/tensor_parallel.py
================================================
# Copyright (c) 2024, Tri Dao.
# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.distributed import ProcessGroup
from mamba_ssm.utils.torch import custom_bwd, custom_fwd
from einops import rearrange
from mamba_ssm.distributed.distributed_utils import (
all_gather_raw,
all_reduce,
all_reduce_raw,
reduce_scatter,
reduce_scatter_raw,
)
class ParallelLinearFunc(torch.autograd.Function):
@staticmethod
@custom_fwd
def forward(ctx, x, weight, bias, process_group=None, sequence_parallel=True):
"""
If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel
with sequence parallelism: we do an all_gather_raw of x before doing the matmul.
"""
ctx.compute_weight_gradient = weight.requires_grad
ctx.process_group = process_group
ctx.sequence_parallel = sequence_parallel
if torch.is_autocast_enabled():
x = x.to(dtype=torch.get_autocast_gpu_dtype())
x = x.contiguous()
if process_group is not None and sequence_parallel:
# We want to kick off the all_gather early, before weight dtype conversion
total_x, handle_x = all_gather_raw(x, process_group, async_op=True)
else:
total_x = x
if torch.is_autocast_enabled():
weight = weight.to(dtype=torch.get_autocast_gpu_dtype())
bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None
weight = weight.contiguous()
if process_group is not None and sequence_parallel:
handle_x.wait()
batch_shape, n = total_x.shape[:-1], total_x.shape[-1]
batch_dim = batch_shape.numel()
# https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174
output = F.linear(total_x, weight, bias)
if ctx.compute_weight_gradient:
ctx.save_for_backward(x, weight)
else:
ctx.save_for_backward(weight)
return output
@staticmethod
@custom_bwd
def backward(ctx, grad_output):
grad_output = grad_output.contiguous()
process_group = ctx.process_group
sequence_parallel = ctx.sequence_parallel
if ctx.compute_weight_gradient:
x, weight = ctx.saved_tensors
if process_group is not None and sequence_parallel:
total_x, handle_x = all_gather_raw(x, process_group, async_op=True)
else:
total_x = x
else:
(weight,) = ctx.saved_tensors
total_x = None
batch_shape = grad_output.shape[:-1]
batch_dim = batch_shape.numel()
grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1])
if ctx.needs_input_grad[0]:
grad_input = F.linear(grad_output, weight.t())
grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1])
if process_group is not None:
reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw
grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True)
else:
grad_input = None
if ctx.needs_input_grad[1]:
assert ctx.compute_weight_gradient
if process_group is not None and sequence_parallel:
handle_x.wait()
grad_weight = torch.einsum(
"bo,bi->oi", grad_output, total_x.reshape(batch_dim, total_x.shape[-1])
)
else:
grad_weight = None
grad_bias = grad_output.sum(dim=0) if ctx.needs_input_grad[2] else None
if process_group is not None and ctx.needs_input_grad[0]:
handle_grad_input.wait()
return grad_input, grad_weight, grad_bias, None, None
def parallel_linear_func(
x: Tensor,
weight: Tensor,
bias: Optional[Tensor] = None,
process_group: Optional[ProcessGroup] = None,
sequence_parallel: bool = True,
):
return ParallelLinearFunc.apply(x, weight, bias, process_group, sequence_parallel)
class ColumnParallelLinear(nn.Linear):
def __init__(
self,
in_features: int,
out_features: int,
process_group: ProcessGroup,
bias: bool = True,
sequence_parallel=True,
multiple_of=1,
device=None,
dtype=None,
) -> None:
world_size = torch.distributed.get_world_size(process_group)
if out_features % multiple_of:
raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}")
multiple = out_features // multiple_of
# We want to split @multiple across world_size, but it could be an uneven split
div = multiple // world_size
mod = multiple % world_size
# The first @mod ranks get @div + 1 copies, the rest get @div copies
local_multiple = div + int(torch.distributed.get_rank(process_group) < mod)
super().__init__(
in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype
)
self.process_group = process_group
self.sequence_parallel = sequence_parallel
def forward(self, x):
# If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism:
# we do an all_gather of x before doing the matmul.
# If not, then the input is already gathered.
return parallel_linear_func(
x,
self.weight,
self.bias,
process_group=self.process_group,
sequence_parallel=self.sequence_parallel,
)
class RowParallelLinear(nn.Linear):
def __init__(
self,
in_features: int,
out_features: int,
process_group: ProcessGroup,
bias: bool = True,
sequence_parallel=True,
multiple_of=1,
device=None,
dtype=None,
) -> None:
world_size = torch.distributed.get_world_size(process_group)
rank = torch.distributed.get_rank(process_group)
if in_features % multiple_of:
raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}")
multiple = in_features // multiple_of
# We want to split @multiple across world_size, but it could be an uneven split
div = multiple // world_size
mod = multiple % world_size
# The first @mod ranks get @div + 1 copies, the rest get @div copies
local_multiple = div + int(torch.distributed.get_rank(process_group) < mod)
# Only rank 0 will have bias
super().__init__(
local_multiple * multiple_of,
out_features,
bias=bias and rank == 0,
device=device,
dtype=dtype,
)
self.process_group = process_group
self.sequence_parallel = sequence_parallel
def forward(self, x):
"""
We're doing Tensor Parallel with sequence parallelism: we do the matmul and then
a reduce_scatter of the result.
"""
out = parallel_linear_func(x, self.weight, self.bias)
reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
return reduce_fn(out, self.process_group)
class VocabParallelEmbedding(nn.Embedding):
def __init__(self, num_embeddings, *args, process_group=None, padding_idx=None, **kwargs):
self.process_group = process_group
if process_group is not None:
world_size = torch.distributed.get_world_size(process_group)
if num_embeddings % world_size != 0:
raise ValueError(
f"num_embeddings ({num_embeddings}) must be divisible by "
f"world_size ({world_size})"
)
if world_size > 1 and padding_idx is not None:
raise RuntimeError("ParallelEmbedding does not support padding_idx")
else:
world_size = 1
super().__init__(num_embeddings // world_size, *args, padding_idx=padding_idx, **kwargs)
def forward(self, input: Tensor) -> Tensor:
if self.process_group is None:
return super().forward(input)
else:
rank = torch.distributed.get_rank(self.process_group)
vocab_size = self.num_embeddings
vocab_start_index, vocab_end_index = rank * vocab_size, (rank + 1) * vocab_size
# Create a mask of valid vocab ids (1 means it needs to be masked).
input_ids_mask = (input < vocab_start_index) | (input >= vocab_end_index)
input = input - vocab_start_index
input[input_ids_mask] = 0
embeddings = super().forward(input)
embeddings[input_ids_mask] = 0.0
return embeddings
class ColumnParallelEmbedding(nn.Embedding):
def __init__(self, num_embeddings, embedding_dim, *args, process_group=None, **kwargs):
self.process_group = process_group
if process_group is not None:
world_size = torch.distributed.get_world_size(process_group)
if embedding_dim % world_size != 0:
raise ValueError(
f"embedding_dim ({embedding_dim}) must be divisible by "
f"world_size ({world_size})"
)
else:
world_size = 1
super().__init__(num_embeddings, embedding_dim // world_size, *args, **kwargs)
class ParallelEmbeddings(nn.Module):
def __init__(
self,
embed_dim,
vocab_size,
max_position_embeddings,
process_group,
padding_idx=None,
sequence_parallel=True,
device=None,
dtype=None,
):
"""
If max_position_embeddings <= 0, there's no position embeddings
"""
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__()
self.process_group = process_group
self.sequence_parallel = sequence_parallel
self.word_embeddings = VocabParallelEmbedding(
vocab_size,
embed_dim,
padding_idx=padding_idx,
process_group=process_group,
**factory_kwargs,
)
self.max_position_embeddings = max_position_embeddings
if self.max_position_embeddings > 0:
self.position_embeddings = ColumnParallelEmbedding(
max_position_embeddings, embed_dim, process_group=process_group, **factory_kwargs
)
def forward(self, input_ids, position_ids=None, combine_batch_seqlen_dim=False):
"""
input_ids: (batch, seqlen)
position_ids: (batch, seqlen)
"""
batch_size, seqlen = input_ids.shape
world_size = torch.distributed.get_world_size(self.process_group)
embeddings = self.word_embeddings(input_ids)
if self.max_position_embeddings > 0:
if position_ids is None:
position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)
position_embeddings = self.position_embeddings(position_ids)
if world_size <= 1:
embeddings = embeddings + position_embeddings
else:
partition_dim = self.position_embeddings.embedding_dim
rank = torch.distributed.get_rank(self.process_group)
embeddings[
..., rank * partition_dim : (rank + 1) * partition_dim
] += position_embeddings
if combine_batch_seqlen_dim:
embeddings = rearrange(embeddings, "b s d -> (b s) d")
reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
return embeddings if world_size <= 1 else reduce_fn(embeddings, self.process_group)
================================================
FILE: mamba_ssm/models/__init__.py
================================================
================================================
FILE: mamba_ssm/models/config_mamba.py
================================================
from dataclasses import dataclass, field
@dataclass
class MambaConfig:
d_model: int = 2560
d_intermediate: int = 0
n_layer: int = 64
vocab_size: int = 50277
ssm_cfg: dict = field(default_factory=dict)
attn_layer_idx: list = field(default_factory=list)
attn_cfg: dict = field(default_factory=dict)
rms_norm: bool = True
residual_in_fp32: bool = True
fused_add_norm: bool = True
pad_vocab_size_multiple: int = 8
tie_embeddings: bool = True
================================================
FILE: mamba_ssm/models/mixer_seq_simple.py
================================================
# Copyright (c) 2023, Albert Gu, Tri Dao.
import math
from functools import partial
import json
import os
import copy
from collections import namedtuple
import torch
import torch.nn as nn
from mamba_ssm.models.config_mamba import MambaConfig
from mamba_ssm.modules.mamba_simple import Mamba
from mamba_ssm.modules.mamba2 import Mamba2
from mamba_ssm.modules.mha import MHA
from mamba_ssm.modules.mlp import GatedMLP
from mamba_ssm.modules.block import Block
from mamba_ssm.utils.generation import GenerationMixin
from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf
try:
from mamba_ssm.ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn
except ImportError:
RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
def create_block(
d_model,
d_intermediate,
ssm_cfg=None,
attn_layer_idx=None,
attn_cfg=None,
norm_epsilon=1e-5,
rms_norm=False,
residual_in_fp32=False,
fused_add_norm=False,
layer_idx=None,
device=None,
dtype=None,
):
if ssm_cfg is None:
ssm_cfg = {}
if attn_layer_idx is None:
attn_layer_idx = []
if attn_cfg is None:
attn_cfg = {}
factory_kwargs = {"device": device, "dtype": dtype}
if layer_idx not in attn_layer_idx:
# Create a copy of the config to modify
ssm_cfg = copy.deepcopy(ssm_cfg) if ssm_cfg is not None else {}
ssm_layer = ssm_cfg.pop("layer", "Mamba1")
if ssm_layer not in ["Mamba1", "Mamba2"]:
raise ValueError(f"Invalid ssm_layer: {ssm_layer}, only support Mamba1 and Mamba2")
mixer_cls = partial(
Mamba2 if ssm_layer == "Mamba2" else Mamba,
layer_idx=layer_idx,
**ssm_cfg,
**factory_kwargs
)
else:
mixer_cls = partial(MHA, layer_idx=layer_idx, **attn_cfg, **factory_kwargs)
norm_cls = partial(
nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
)
if d_intermediate == 0:
mlp_cls = nn.Identity
else:
mlp_cls = partial(
GatedMLP, hidden_features=d_intermediate, out_features=d_model, **factory_kwargs
)
block = Block(
d_model,
mixer_cls,
mlp_cls,
norm_cls=norm_cls,
fused_add_norm=fused_add_norm,
residual_in_fp32=residual_in_fp32,
)
block.layer_idx = layer_idx
return block
# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
def _init_weights(
module,
n_layer,
initializer_range=0.02, # Now only used for embedding layer.
rescale_prenorm_residual=True,
n_residuals_per_layer=1, # Change to 2 if we have MLP
):
if isinstance(module, nn.Linear):
if module.bias is not None:
if not getattr(module.bias, "_no_reinit", False):
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
nn.init.normal_(module.weight, std=initializer_range)
if rescale_prenorm_residual:
# Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
# > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
# > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
# > -- GPT-2 :: https://openai.com/blog/better-language-models/
#
# Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
for name, p in module.named_parameters():
if name in ["out_proj.weight", "fc2.weight"]:
# Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
# Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
# We need to reinit p since this code could be called multiple times
# Having just p *= scale would repeatedly scale it down
nn.init.kaiming_uniform_(p, a=math.sqrt(5))
with torch.no_grad():
p /= math.sqrt(n_residuals_per_layer * n_layer)
class MixerModel(nn.Module):
def __init__(
self,
d_model: int,
n_layer: int,
d_intermediate: int,
vocab_size: int,
ssm_cfg=None,
attn_layer_idx=None,
attn_cfg=None,
norm_epsilon: float = 1e-5,
rms_norm: bool = False,
initializer_cfg=None,
fused_add_norm=False,
residual_in_fp32=False,
device=None,
dtype=None,
) -> None:
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__()
self.residual_in_fp32 = residual_in_fp32
gitextract_vi30yky3/ ├── .github/ │ ├── scripts/ │ │ ├── build.sh │ │ ├── check_for_ngc_images.sh │ │ └── test.sh │ └── workflows/ │ ├── _build.yml │ ├── _build_in_container.yml │ ├── build.yml │ ├── build_in_container.yml │ └── publish.yaml ├── .gitignore ├── .gitmodules ├── AUTHORS ├── LICENSE ├── MANIFEST.in ├── README.md ├── benchmarks/ │ └── benchmark_generation_mamba_simple.py ├── csrc/ │ └── selective_scan/ │ ├── reverse_scan.cuh │ ├── selective_scan.cpp │ ├── selective_scan.h │ ├── selective_scan_bwd_bf16_complex.cu │ ├── selective_scan_bwd_bf16_real.cu │ ├── selective_scan_bwd_fp16_complex.cu │ ├── selective_scan_bwd_fp16_real.cu │ ├── selective_scan_bwd_fp32_complex.cu │ ├── selective_scan_bwd_fp32_real.cu │ ├── selective_scan_bwd_kernel.cuh │ ├── selective_scan_common.h │ ├── selective_scan_fwd_bf16.cu │ ├── selective_scan_fwd_fp16.cu │ ├── selective_scan_fwd_fp32.cu │ ├── selective_scan_fwd_kernel.cuh │ ├── static_switch.h │ └── uninitialized_copy.cuh ├── evals/ │ └── lm_harness_eval.py ├── mamba_ssm/ │ ├── __init__.py │ ├── distributed/ │ │ ├── __init__.py │ │ ├── distributed_utils.py │ │ └── tensor_parallel.py │ ├── models/ │ │ ├── __init__.py │ │ ├── config_mamba.py │ │ └── mixer_seq_simple.py │ ├── modules/ │ │ ├── __init__.py │ │ ├── block.py │ │ ├── mamba2.py │ │ ├── mamba2_simple.py │ │ ├── mamba3.py │ │ ├── mamba_simple.py │ │ ├── mha.py │ │ ├── mlp.py │ │ └── ssd_minimal.py │ ├── ops/ │ │ ├── __init__.py │ │ ├── cute/ │ │ │ └── mamba3/ │ │ │ └── mamba3_step_fn.py │ │ ├── selective_scan_interface.py │ │ ├── tilelang/ │ │ │ └── mamba3/ │ │ │ ├── mamba3_mimo.py │ │ │ ├── mamba3_mimo_bwd.py │ │ │ └── mamba3_mimo_fwd.py │ │ └── triton/ │ │ ├── __init__.py │ │ ├── angle_cumsum.py │ │ ├── k_activations.py │ │ ├── layer_norm.py │ │ ├── layernorm_gated.py │ │ ├── mamba3/ │ │ │ ├── angle_dt.py │ │ │ ├── mamba3_mimo_rotary_step.py │ │ │ ├── mamba3_mimo_utils.py │ │ │ ├── mamba3_siso_bwd.py │ │ │ ├── mamba3_siso_combined.py │ │ │ ├── mamba3_siso_fwd.py │ │ │ ├── mamba3_siso_step.py │ │ │ └── utils.py │ │ ├── selective_state_update.py │ │ ├── softplus.py │ │ ├── ssd_bmm.py │ │ ├── ssd_chunk_scan.py │ │ ├── ssd_chunk_state.py │ │ ├── ssd_combined.py │ │ └── ssd_state_passing.py │ └── utils/ │ ├── __init__.py │ ├── determinism.py │ ├── generation.py │ ├── hf.py │ └── torch.py ├── pyproject.toml ├── rocm_patch/ │ └── rocm6_0.patch ├── setup.py ├── tests/ │ ├── benchmark_determinism_kernels.py │ ├── ops/ │ │ ├── cute/ │ │ │ └── test_mamba3_mimo_step.py │ │ ├── test_selective_scan.py │ │ ├── tilelang/ │ │ │ └── test_mamba3_mimo.py │ │ └── triton/ │ │ ├── test_layernorm_gated.py │ │ ├── test_mamba3_siso.py │ │ ├── test_selective_state_update.py │ │ └── test_ssd.py │ ├── test_determinism.py │ └── test_generation.py └── usage.md
SYMBOL INDEX (465 symbols across 55 files)
FILE: csrc/selective_scan/selective_scan.cpp
function set_ssm_params_fwd (line 59) | void set_ssm_params_fwd(SSMParamsBase ¶ms,
function set_ssm_params_bwd (line 143) | void set_ssm_params_bwd(SSMParamsBwd ¶ms,
function selective_scan_fwd (line 226) | std::vector<at::Tensor>
function selective_scan_bwd (line 338) | std::vector<at::Tensor>
function PYBIND11_MODULE (line 494) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: csrc/selective_scan/selective_scan.h
type SSMScanParamsBase (line 9) | struct SSMScanParamsBase {
type SSMParamsBase (line 26) | struct SSMParamsBase {
function SSMParamsBase (line 71) | struct SSMParamsBwd: public SSMParamsBase {
FILE: csrc/selective_scan/selective_scan_common.h
function custom_max (line 18) | constexpr size_t custom_max(std::initializer_list<size_t> ilist)
function T (line 24) | T constexpr_min(T a, T b) {
function custom_max (line 29) | constexpr size_t custom_max(std::initializer_list<size_t> ilist)
function T (line 35) | T constexpr_min(T a, T b) {
type BytesToType (line 61) | struct BytesToType
type BytesToType (line 66) | struct BytesToType
type BytesToType (line 71) | struct BytesToType
type BytesToType (line 76) | struct BytesToType
type BytesToType (line 81) | struct BytesToType
function __device__ (line 90) | static inline __device__ void to_float(const scalar_t (&src)[N], float (...
function __device__ (line 98) | static inline __device__ void to_float(const at::Half (&src)[N], float (...
function __device__ (line 110) | static inline __device__ void to_float(const at::BFloat16 (&src)[N], flo...
function complex_t (line 124) | complex_t cexp2f(complex_t z) {
function complex_t (line 131) | complex_t cexpf(complex_t z) {
function float (line 141) | struct SSMScanOp<float> {
function complex_t (line 148) | struct SSMScanOp<complex_t> {
function __device__ (line 166) | __device__ SSMScanPrefixCallbackOp(scan_t running_prefix_) : running_pre...
function __device__ (line 169) | __device__ scan_t operator()(scan_t block_aggregate) {
function typename (line 186) | typename Ktraits::BlockLoadVecT(smem_load_vec).Load(
function typename (line 210) | typename Ktraits::BlockLoadWeightVecT(smem_load_weight_vec).Load(
function typename (line 225) | typename Ktraits::BlockLoadWeightVecT(smem_load_weight_vec).Load(
function typename (line 248) | typename Ktraits::BlockStoreVecT(smem_store_vec).Store(
FILE: evals/lm_harness_eval.py
class MambaEvalWrapper (line 15) | class MambaEvalWrapper(HFLM):
method __init__ (line 19) | def __init__(self, pretrained="state-spaces/mamba-2.8b", max_length=20...
method batch_size (line 31) | def batch_size(self):
method _model_generate (line 34) | def _model_generate(self, context, max_length, stop, **generation_kwar...
FILE: mamba_ssm/distributed/distributed_utils.py
function all_gather_raw (line 18) | def all_gather_raw(input_: Tensor, process_group: ProcessGroup, async_op...
function reduce_scatter_raw (line 30) | def reduce_scatter_raw(input_: Tensor, process_group: ProcessGroup, asyn...
function all_reduce_raw (line 43) | def all_reduce_raw(input_: Tensor, process_group: ProcessGroup, async_op...
class AllGatherFunc (line 49) | class AllGatherFunc(torch.autograd.Function):
method forward (line 53) | def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
method backward (line 59) | def backward(ctx, grad_output: Tensor):
class ReduceScatterFunc (line 68) | class ReduceScatterFunc(torch.autograd.Function):
method forward (line 72) | def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
method backward (line 78) | def backward(ctx, grad_output: Tensor):
class AllReduceFunc (line 87) | class AllReduceFunc(torch.autograd.Function):
method forward (line 91) | def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor:
method backward (line 97) | def backward(ctx, grad_output: Tensor):
function sync_shared_params (line 105) | def sync_shared_params(model: torch.nn.Module, process_group: ProcessGro...
function allreduce_sequence_parallel_grad (line 120) | def allreduce_sequence_parallel_grad(model: torch.nn.Module, process_gro...
function get_dim_for_local_rank (line 135) | def get_dim_for_local_rank(dim: int, world_size: int, local_rank: int, m...
FILE: mamba_ssm/distributed/tensor_parallel.py
class ParallelLinearFunc (line 23) | class ParallelLinearFunc(torch.autograd.Function):
method forward (line 26) | def forward(ctx, x, weight, bias, process_group=None, sequence_paralle...
method backward (line 62) | def backward(ctx, grad_output):
function parallel_linear_func (line 101) | def parallel_linear_func(
class ColumnParallelLinear (line 111) | class ColumnParallelLinear(nn.Linear):
method __init__ (line 112) | def __init__(
method forward (line 138) | def forward(self, x):
class RowParallelLinear (line 151) | class RowParallelLinear(nn.Linear):
method __init__ (line 152) | def __init__(
method forward (line 184) | def forward(self, x):
class VocabParallelEmbedding (line 194) | class VocabParallelEmbedding(nn.Embedding):
method __init__ (line 195) | def __init__(self, num_embeddings, *args, process_group=None, padding_...
method forward (line 210) | def forward(self, input: Tensor) -> Tensor:
class ColumnParallelEmbedding (line 226) | class ColumnParallelEmbedding(nn.Embedding):
method __init__ (line 227) | def __init__(self, num_embeddings, embedding_dim, *args, process_group...
class ParallelEmbeddings (line 241) | class ParallelEmbeddings(nn.Module):
method __init__ (line 242) | def __init__(
method forward (line 273) | def forward(self, input_ids, position_ids=None, combine_batch_seqlen_d...
FILE: mamba_ssm/models/config_mamba.py
class MambaConfig (line 5) | class MambaConfig:
FILE: mamba_ssm/models/mixer_seq_simple.py
function create_block (line 29) | def create_block(
function _init_weights (line 86) | def _init_weights(
class MixerModel (line 118) | class MixerModel(nn.Module):
method __init__ (line 119) | def __init__(
method allocate_inference_cache (line 184) | def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None,...
method forward (line 190) | def forward(self, input_ids, inference_params=None, **mixer_kwargs):
class MambaLMHeadModel (line 215) | class MambaLMHeadModel(nn.Module, GenerationMixin):
method __init__ (line 217) | def __init__(
method tie_weights (line 267) | def tie_weights(self):
method allocate_inference_cache (line 271) | def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None,...
method forward (line 274) | def forward(self, input_ids, position_ids=None, inference_params=None,...
method from_pretrained (line 287) | def from_pretrained(cls, pretrained_model_name, device=None, dtype=Non...
method save_pretrained (line 294) | def save_pretrained(self, save_directory):
FILE: mamba_ssm/modules/block.py
class Block (line 10) | class Block(nn.Module):
method __init__ (line 11) | def __init__(
method forward (line 42) | def forward(
method allocate_inference_cache (line 90) | def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None,...
FILE: mamba_ssm/modules/mamba2.py
class Mamba2 (line 37) | class Mamba2(nn.Module, PyTorchModelHubMixin):
method __init__ (line 38) | def __init__(
method forward (line 154) | def forward(self, u, seqlen=None, seq_idx=None, cu_seqlens=None, infer...
method step (line 278) | def step(self, hidden_states, conv_state, ssm_state):
method allocate_inference_cache (line 345) | def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None,...
method _get_states_from_cache (line 357) | def _get_states_from_cache(self, inference_params, batch_size, initial...
FILE: mamba_ssm/modules/mamba2_simple.py
class Mamba2Simple (line 24) | class Mamba2Simple(nn.Module):
method __init__ (line 25) | def __init__(
method forward (line 124) | def forward(self, u, seq_idx=None):
FILE: mamba_ssm/modules/mamba3.py
class Mamba3 (line 20) | class Mamba3(nn.Module):
method __init__ (line 21) | def __init__(
method forward (line 127) | def forward(self, u, seq_idx=None, cu_seqlens=None, inference_params=N...
method _preprocess (line 249) | def _preprocess(self, A_proj, dd_dt, B, C, x, z, trap_proj, angle_proj):
method _postprocess (line 272) | def _postprocess(self, y, outpj, z, zpj, headdim):
method step (line 282) | def step(self, u, angle_state, ssm_state, k_state, v_state, **kwargs):
method allocate_inference_cache (line 409) | def allocate_inference_cache(self, batch_size, max_seqlen, device=None...
method _get_states_from_cache (line 451) | def _get_states_from_cache(self, inference_params, batch_size, initial...
FILE: mamba_ssm/modules/mamba_simple.py
class Mamba (line 31) | class Mamba(nn.Module):
method __init__ (line 32) | def __init__(
method forward (line 119) | def forward(self, hidden_states, inference_params=None):
method step (line 208) | def step(self, hidden_states, conv_state, ssm_state):
method allocate_inference_cache (line 255) | def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None,...
method _get_states_from_cache (line 268) | def _get_states_from_cache(self, inference_params, batch_size, initial...
FILE: mamba_ssm/modules/mha.py
function _update_kv_cache (line 26) | def _update_kv_cache(kv, inference_params, layer_idx):
class MHA (line 44) | class MHA(nn.Module):
method __init__ (line 47) | def __init__(
method allocate_inference_cache (line 110) | def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None):
method _update_kv_cache (line 124) | def _update_kv_cache(self, kv, inference_params):
method _apply_rotary_update_kvcache_attention (line 129) | def _apply_rotary_update_kvcache_attention(self, q, kv, inference_para...
method _update_kvcache_attention (line 167) | def _update_kvcache_attention(self, q, kv, inference_params):
method forward (line 201) | def forward(self, x, inference_params=None):
FILE: mamba_ssm/modules/mlp.py
class GatedMLP (line 6) | class GatedMLP(nn.Module):
method __init__ (line 7) | def __init__(
method forward (line 29) | def forward(self, x):
FILE: mamba_ssm/modules/ssd_minimal.py
function segsum_unstable (line 14) | def segsum_unstable(x):
function segsum (line 23) | def segsum(x):
function ssd_minimal_discrete (line 34) | def ssd_minimal_discrete(X, A, B, C, block_len, initial_states=None):
function test_correctness (line 82) | def test_correctness():
FILE: mamba_ssm/ops/cute/mamba3/mamba3_step_fn.py
function transpose_view (line 22) | def transpose_view(a: cute.Tensor) -> cute.Tensor:
function select (line 28) | def select(a: cute.Tensor, mode: List[int]) -> cute.Tensor:
function get_gmem_tiled_copy (line 33) | def get_gmem_tiled_copy(dtype: Type[cutlass.Numeric], major_mode_size: i...
class Mamba3Step (line 48) | class Mamba3Step():
method __init__ (line 49) | def __init__(self, tile_D: int, dstate: int, mimo: int = 1, num_warps:...
method _setup_smem_layouts (line 59) | def _setup_smem_layouts(self):
method _setup_gmem_tiled_copy (line 66) | def _setup_gmem_tiled_copy(self, ):
method __call__ (line 86) | def __call__(
method kernel (line 227) | def kernel(
function mamba3_step_fn (line 566) | def mamba3_step_fn(
function selective_state_update_fused_ref_v2 (line 741) | def selective_state_update_fused_ref_v2(
function _bytes_of (line 816) | def _bytes_of(t):
FILE: mamba_ssm/ops/selective_scan_interface.py
class SelectiveScanFn (line 23) | class SelectiveScanFn(torch.autograd.Function):
method forward (line 26) | def forward(ctx, u, delta, A, B, C, D=None, z=None, delta_bias=None, d...
method backward (line 59) | def backward(ctx, dout, *args):
function rms_norm_forward (line 86) | def rms_norm_forward(
function selective_scan_fn (line 106) | def selective_scan_fn(u, delta, A, B, C, D=None, z=None, delta_bias=None...
function selective_scan_ref (line 115) | def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=Non...
class MambaInnerFn (line 184) | class MambaInnerFn(torch.autograd.Function):
method forward (line 188) | def forward(ctx, xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_...
method backward (line 282) | def backward(ctx, dout):
function mamba_inner_fn (line 373) | def mamba_inner_fn(
function mamba_inner_ref (line 384) | def mamba_inner_ref(
FILE: mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py
class _Mamba3Function (line 24) | class _Mamba3Function(torch.autograd.Function):
method forward (line 28) | def forward(
method backward (line 88) | def backward(ctx, dout, *args) -> tuple:
function mamba3_mimo (line 154) | def mamba3_mimo(
FILE: mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_bwd.py
function mamba_mimo_bwd_fwd (line 41) | def mamba_mimo_bwd_fwd(
function mamba_mimo_bwd_bwd (line 500) | def mamba_mimo_bwd_bwd(
function mamba_mimo_bwd_combined (line 1146) | def mamba_mimo_bwd_combined(
FILE: mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py
function mamba_mimo_fwd (line 38) | def mamba_mimo_fwd(
function mamba_mimo_forward (line 413) | def mamba_mimo_forward(q, k, v,
FILE: mamba_ssm/ops/triton/angle_cumsum.py
class AngleDtFn (line 12) | class AngleDtFn(torch.autograd.Function):
method forward (line 14) | def forward(ctx,
method backward (line 27) | def backward(ctx, grad_out: torch.Tensor):
function angle_dt (line 37) | def angle_dt(angle: torch.Tensor,
function cumsum_kernel (line 45) | def cumsum_kernel(
function angle_dt_fwd_kernel (line 87) | def angle_dt_fwd_kernel(
function angle_dt_bwd_kernel (line 194) | def angle_dt_bwd_kernel(
function apply_angle_dt_fwd (line 307) | def apply_angle_dt_fwd(
function apply_angle_dt_bwd (line 395) | def apply_angle_dt_bwd(
function apply_cumsum (line 504) | def apply_cumsum(
function apply_angle_dt_reference (line 541) | def apply_angle_dt_reference(
function test_correctness (line 561) | def test_correctness():
function test_cumsum_correctness (line 587) | def test_cumsum_correctness():
function test_backward_correctness (line 611) | def test_backward_correctness():
function benchmark_angle_dt (line 647) | def benchmark_angle_dt():
function benchmark_angle_dt_backward (line 706) | def benchmark_angle_dt_backward():
FILE: mamba_ssm/ops/triton/k_activations.py
function _swiglu_fwd_kernel (line 23) | def _swiglu_fwd_kernel(
function _swiglu_fwd (line 46) | def _swiglu_fwd(xy, out=None):
function _swiglu_bwd_kernel (line 78) | def _swiglu_bwd_kernel(
function _swiglu_bwd (line 119) | def _swiglu_bwd(xy, dout, dxy=None, recompute_output=False, out=None):
class SwiGLU (line 158) | class SwiGLU(torch.autograd.Function):
method forward (line 161) | def forward(ctx, xy):
method backward (line 166) | def backward(ctx, dout):
FILE: mamba_ssm/ops/triton/layer_norm.py
function layer_norm_ref (line 22) | def layer_norm_ref(
function rms_norm_ref (line 77) | def rms_norm_ref(
function config_prune (line 130) | def config_prune(configs):
function _layer_norm_fwd_1pass_kernel (line 181) | def _layer_norm_fwd_1pass_kernel(
function _layer_norm_fwd (line 291) | def _layer_norm_fwd(
function _layer_norm_bwd_kernel (line 436) | def _layer_norm_bwd_kernel(
function _layer_norm_bwd (line 589) | def _layer_norm_bwd(
class LayerNormFn (line 728) | class LayerNormFn(torch.autograd.Function):
method forward (line 730) | def forward(
method backward (line 828) | def backward(ctx, dy, *args):
function layer_norm_fn (line 888) | def layer_norm_fn(
function rms_norm_fn (line 922) | def rms_norm_fn(
class RMSNorm (line 955) | class RMSNorm(torch.nn.Module):
method __init__ (line 957) | def __init__(self, hidden_size, eps=1e-5, dropout_p=0.0, device=None, ...
method reset_parameters (line 969) | def reset_parameters(self):
method forward (line 972) | def forward(self, x, residual=None, prenorm=False, residual_in_fp32=Fa...
class LayerNormLinearFn (line 985) | class LayerNormLinearFn(torch.autograd.Function):
method forward (line 988) | def forward(
method backward (line 1047) | def backward(ctx, dout, *args):
function layer_norm_linear_fn (line 1092) | def layer_norm_linear_fn(
FILE: mamba_ssm/ops/triton/layernorm_gated.py
function rms_norm_ref (line 18) | def rms_norm_ref(x, weight, bias, z=None, eps=1e-6, group_size=None, nor...
function _layer_norm_fwd_1pass_kernel (line 45) | def _layer_norm_fwd_1pass_kernel(
function _layer_norm_fwd (line 108) | def _layer_norm_fwd(x, weight, bias, eps, z=None, out=None, group_size=N...
function _layer_norm_bwd_kernel (line 155) | def _layer_norm_bwd_kernel(
function _layer_norm_bwd (line 271) | def _layer_norm_bwd(dy, x, weight, bias, eps, mean, rstd, z=None, group_...
class LayerNormFn (line 338) | class LayerNormFn(torch.autograd.Function):
method forward (line 341) | def forward(ctx, x, weight, bias, z=None, eps=1e-6, group_size=None, n...
method backward (line 369) | def backward(ctx, dy):
function layernorm_fn (line 380) | def layernorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, nor...
function rmsnorm_fn (line 384) | def rmsnorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_...
class LayerNorm (line 388) | class LayerNorm(torch.nn.Module):
method __init__ (line 390) | def __init__(self, hidden_size, eps=1e-5, group_size=None, norm_before...
method reset_parameters (line 404) | def reset_parameters(self):
method forward (line 408) | def forward(self, x, z=None):
class RMSNorm (line 415) | class RMSNorm(torch.nn.Module):
method __init__ (line 417) | def __init__(self, hidden_size, eps=1e-5, group_size=None, norm_before...
method reset_parameters (line 430) | def reset_parameters(self):
method forward (line 433) | def forward(self, x, z=None):
FILE: mamba_ssm/ops/triton/mamba3/angle_dt.py
function angle_dt_fwd_kernel (line 24) | def angle_dt_fwd_kernel(
function angle_dt_fwd (line 125) | def angle_dt_fwd(
function angle_dt_bwd_kernel (line 232) | def angle_dt_bwd_kernel(
function angle_dt_bwd (line 345) | def angle_dt_bwd(
FILE: mamba_ssm/ops/triton/mamba3/mamba3_mimo_rotary_step.py
function rotary_qk_inference_kernel (line 16) | def rotary_qk_inference_kernel(
function apply_rotary_qk_inference_fwd (line 151) | def apply_rotary_qk_inference_fwd(
function apply_rotary_qk_inference_reference (line 239) | def apply_rotary_qk_inference_reference(
function test_correctness_qk_inference (line 327) | def test_correctness_qk_inference():
FILE: mamba_ssm/ops/triton/mamba3/mamba3_mimo_utils.py
function bwd_dadt_cumsum_fused_kernel (line 37) | def bwd_dadt_cumsum_fused_kernel(
function bwd_segsum_dadt_kernel (line 134) | def bwd_segsum_dadt_kernel(
function bwd_dtrap_ddt_kernel (line 225) | def bwd_dtrap_ddt_kernel(
function dacs_segsum_kernel (line 349) | def dacs_segsum_kernel(
function bwd_dadt_fused_triton (line 407) | def bwd_dadt_fused_triton(
function bwd_dtrap_ddt_triton (line 450) | def bwd_dtrap_ddt_triton(
function compute_dacs_segsum_triton (line 478) | def compute_dacs_segsum_triton(
function bwd_segsum_ddt_from_dSSdA_ref (line 508) | def bwd_segsum_ddt_from_dSSdA_ref(
function bwd_ddt_from_ddA_cs_rev_ref (line 528) | def bwd_ddt_from_ddA_cs_rev_ref(
function bwd_ddt_from_ddA_cs_ref (line 545) | def bwd_ddt_from_ddA_cs_ref(
function compute_dtrap_ddt_ref (line 560) | def compute_dtrap_ddt_ref(dfactor: torch.Tensor,
function compute_dacs_segsum_ref (line 582) | def compute_dacs_segsum_ref(da: torch.Tensor, # (B, H, S)
function test_bwd_ddt_fused_correctness (line 606) | def test_bwd_ddt_fused_correctness():
function test_dtrap_ddt_correctness (line 651) | def test_dtrap_ddt_correctness():
function test_dacs_segsum_correctness (line 696) | def test_dacs_segsum_correctness():
function benchmark_bwd_ddt (line 729) | def benchmark_bwd_ddt():
function benchmark_dacs_segsum (line 792) | def benchmark_dacs_segsum():
function benchmark_dtrap_ddt (line 829) | def benchmark_dtrap_ddt():
FILE: mamba_ssm/ops/triton/mamba3/mamba3_siso_bwd.py
function mamba3_siso_bwd_kernel_dzdo (line 32) | def mamba3_siso_bwd_kernel_dzdo(
function compute_dzdo (line 114) | def compute_dzdo(
function mamba3_siso_bwd_kernel_dqkv (line 202) | def mamba3_siso_bwd_kernel_dqkv(
function compute_dqkv (line 614) | def compute_dqkv(
function mamba3_siso_bwd_kernel_rotary_bias_angles (line 811) | def mamba3_siso_bwd_kernel_rotary_bias_angles(
function mamba3_siso_bwd_kernel_dk_state_post (line 1055) | def mamba3_siso_bwd_kernel_dk_state_post(
function compute_dqktheta (line 1159) | def compute_dqktheta(
function apply_dk_state_post (line 1338) | def apply_dk_state_post(
function mamba3_siso_bwd_kernel_ddt_dtrap_dinput_states (line 1418) | def mamba3_siso_bwd_kernel_ddt_dtrap_dinput_states(
function compute_ddt_dtrap_dinput_states (line 1611) | def compute_ddt_dtrap_dinput_states(
function _alloc_fn (line 1771) | def _alloc_fn(size: int, alignment: int, stream: Optional[int]):
FILE: mamba_ssm/ops/triton/mamba3/mamba3_siso_combined.py
function _triton_alloc_fn (line 21) | def _triton_alloc_fn(size: int, alignment: int, stream: Optional[int]):
class Mamba3Output (line 34) | class Mamba3Output:
class _Mamba3Function (line 50) | class _Mamba3Function(torch.autograd.Function):
method forward (line 54) | def forward(
method backward (line 153) | def backward(
function mamba3_siso_combined (line 291) | def mamba3_siso_combined(
FILE: mamba_ssm/ops/triton/mamba3/mamba3_siso_fwd.py
function mamba3_siso_fwd_kernel (line 29) | def mamba3_siso_fwd_kernel(
function _alloc_fn (line 434) | def _alloc_fn(size: int, alignment: int, stream: Optional[int]):
function mamba3_siso_fwd (line 439) | def mamba3_siso_fwd(
FILE: mamba_ssm/ops/triton/mamba3/mamba3_siso_step.py
function mamba3_siso_step_kernel (line 27) | def mamba3_siso_step_kernel(
function _alloc_fn (line 228) | def _alloc_fn(size: int, alignment: int, stream: Optional[int]):
function mamba3_siso_step (line 233) | def mamba3_siso_step(
FILE: mamba_ssm/ops/triton/mamba3/utils.py
function cos_approx (line 14) | def cos_approx(x):
function sin_approx (line 34) | def sin_approx(x):
function tanh_approx (line 53) | def tanh_approx(x):
function sech2_approx (line 72) | def sech2_approx(x):
function sigmoid_approx (line 92) | def sigmoid_approx(x):
function silu (line 117) | def silu(x):
FILE: mamba_ssm/ops/triton/selective_state_update.py
function _selective_scan_update_kernel (line 24) | def _selective_scan_update_kernel(
function selective_state_update (line 135) | def selective_state_update(state, x, dt, A, B, C, D=None, z=None, dt_bia...
function selective_state_update_ref (line 224) | def selective_state_update_ref(state, x, dt, A, B, C, D=None, z=None, dt...
FILE: mamba_ssm/ops/triton/softplus.py
function softplus (line 10) | def softplus(dt):
function softplus (line 14) | def softplus(dt):
FILE: mamba_ssm/ops/triton/ssd_bmm.py
function init_to_zero (line 18) | def init_to_zero(names):
function _bmm_chunk_fwd_kernel (line 37) | def _bmm_chunk_fwd_kernel(
function _bmm_chunk_bwd_kernel (line 111) | def _bmm_chunk_bwd_kernel(
function _bmm_chunk_fwd (line 165) | def _bmm_chunk_fwd(a, b, chunk_size, seq_idx=None, causal=False, output_...
function _bmm_chunk_bwd (line 213) | def _bmm_chunk_bwd(a, dout, residual=None, out=None):
FILE: mamba_ssm/ops/triton/ssd_chunk_scan.py
function init_to_zero (line 28) | def init_to_zero(names):
function _chunk_scan_fwd_kernel (line 49) | def _chunk_scan_fwd_kernel(
function _chunk_scan_fwd_kernel_wip (line 195) | def _chunk_scan_fwd_kernel_wip(
function _chunk_scan_bwd_dz_kernel (line 349) | def _chunk_scan_bwd_dz_kernel(
function _chunk_scan_bwd_dstates_kernel (line 449) | def _chunk_scan_bwd_dstates_kernel(
function _chunk_scan_bwd_dc_kernel (line 530) | def _chunk_scan_bwd_dc_kernel(
function _chunk_scan_bwd_dx_kernel (line 641) | def _chunk_scan_bwd_dx_kernel(
function _chunk_scan_bwd_dcb_kernel (line 774) | def _chunk_scan_bwd_dcb_kernel(
function _chunk_scan_bwd_ddAcs_unstable_kernel (line 889) | def _chunk_scan_bwd_ddAcs_unstable_kernel(
function _chunk_scan_bwd_ddAcs_stable_kernel_old (line 979) | def _chunk_scan_bwd_ddAcs_stable_kernel_old(
function _chunk_scan_bwd_ddAcs_stable_kernel (line 1098) | def _chunk_scan_bwd_ddAcs_stable_kernel(
function _chunk_scan_bwd_ddAcs_prev_kernel (line 1198) | def _chunk_scan_bwd_ddAcs_prev_kernel(
function _chunk_scan_fwd (line 1259) | def _chunk_scan_fwd(cb, x, dt, dA_cumsum, C, states, D=None, z=None, seq...
function _chunk_scan_fwd_wip (line 1311) | def _chunk_scan_fwd_wip(cb, x, dt, dA_cumsum, C, B, states, D=None, z=No...
function _chunk_scan_bwd_dz (line 1363) | def _chunk_scan_bwd_dz(x, z, out, dout, chunk_size, has_ddAcs=True, D=No...
function _chunk_scan_bwd_dstates (line 1423) | def _chunk_scan_bwd_dstates(C, dA_cumsum, dout, seq_idx=None, dtype=None):
function _chunk_scan_bwd_dC (line 1451) | def _chunk_scan_bwd_dC(prev_states, dA_cumsum, dout, seq_idx=None, C=Non...
function _chunk_scan_bwd_dcb (line 1514) | def _chunk_scan_bwd_dcb(x, dt, dA_cumsum, dout, seq_idx=None, CB=None, n...
function _chunk_scan_bwd_dx (line 1565) | def _chunk_scan_bwd_dx(cb, x, dt, dA_cumsum, dout, D=None):
function _chunk_scan_bwd_ddAcs_unstable (line 1618) | def _chunk_scan_bwd_ddAcs_unstable(x, dt, out, dout, ddt, D=None, subtra...
function _chunk_scan_bwd_ddAcs_stable_old (line 1667) | def _chunk_scan_bwd_ddAcs_stable_old(x, dt, dA_cumsum, dout, cb):
function _chunk_scan_bwd_ddAcs_stable (line 1700) | def _chunk_scan_bwd_ddAcs_stable(x, dt, dA_cumsum, dout, cb):
function _chunk_scan_bwd_ddAcs_prev (line 1732) | def _chunk_scan_bwd_ddAcs_prev(prev_states, C, dout, dA_cumsum, seq_idx=...
class ChunkScanFn (line 1764) | class ChunkScanFn(torch.autograd.Function):
method forward (line 1767) | def forward(ctx, B, C, x, dt, dA_cumsum, prev_states, D=None, z=None):
method backward (line 1798) | def backward(ctx, dout):
function chunk_scan (line 1828) | def chunk_scan(B, C, x, dt, dA_cumsum, prev_states, D=None, z=None):
function chunk_scan_ref (line 1846) | def chunk_scan_ref(B, C, x, dt, dA_cumsum, prev_states, D=None, z=None):
FILE: mamba_ssm/ops/triton/ssd_chunk_state.py
function init_to_zero (line 24) | def init_to_zero(names):
function _chunk_cumsum_fwd_kernel (line 40) | def _chunk_cumsum_fwd_kernel(
function _chunk_cumsum_bwd_kernel (line 102) | def _chunk_cumsum_bwd_kernel(
function _chunk_state_fwd_kernel (line 192) | def _chunk_state_fwd_kernel(
function _chunk_state_bwd_dx_kernel (line 286) | def _chunk_state_bwd_dx_kernel(
function _chunk_state_bwd_db_kernel (line 398) | def _chunk_state_bwd_db_kernel(
function _chunk_state_bwd_ddAcs_stable_kernel (line 528) | def _chunk_state_bwd_ddAcs_stable_kernel(
function _chunk_state_varlen_kernel (line 640) | def _chunk_state_varlen_kernel(
function _chunk_cumsum_fwd (line 718) | def _chunk_cumsum_fwd(dt, A, chunk_size, dt_bias=None, dt_softplus=False...
function _chunk_cumsum_bwd (line 744) | def _chunk_cumsum_bwd(ddA, ddt_out, dt, A, dt_bias=None, dt_softplus=Fal...
function _chunk_state_fwd (line 812) | def _chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=None, states=None, sta...
function _chunk_state_bwd_dx (line 845) | def _chunk_state_bwd_dx(B, x, dt, dA_cumsum, dstates, dx=None):
function _chunk_state_bwd_db (line 902) | def _chunk_state_bwd_db(x, dt, dA_cumsum, dstates, seq_idx=None, B=None,...
function _chunk_state_bwd_ddAcs_stable (line 972) | def _chunk_state_bwd_ddAcs_stable(B, x, dt, dA_cumsum, dstates, seq_idx=...
function chunk_state_varlen (line 1018) | def chunk_state_varlen(B, x, dt, dA_cumsum, cu_seqlens, chunk_states):
class ChunkStateFn (line 1047) | class ChunkStateFn(torch.autograd.Function):
method forward (line 1050) | def forward(ctx, B, x, dt, dA_cumsum, states_in_fp32=True):
method backward (line 1067) | def backward(ctx, dstates):
function chunk_state (line 1081) | def chunk_state(B, x, dt, dA_cumsum, states_in_fp32=True):
function chunk_state_ref (line 1094) | def chunk_state_ref(B, x, dt, dA_cumsum):
FILE: mamba_ssm/ops/triton/ssd_combined.py
function init_to_zero (line 55) | def init_to_zero(names):
function ensure_stride (line 59) | def ensure_stride(inp):
function _chunk_scan_chunk_state_bwd_dx_kernel (line 93) | def _chunk_scan_chunk_state_bwd_dx_kernel(
function _chunk_scan_chunk_state_bwd_dx (line 262) | def _chunk_scan_chunk_state_bwd_dx(x, dt, dA_cumsum, B, CB, dout, dstate...
function _mamba_chunk_scan_combined_fwd (line 343) | def _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D=None, z...
function _mamba_chunk_scan_combined_bwd (line 396) | def _mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size...
function selective_scan_bwd (line 516) | def selective_scan_bwd(dout, x, dt, A, B, C, D=None, z=None):
class MambaChunkScanCombinedFn (line 593) | class MambaChunkScanCombinedFn(torch.autograd.Function):
method forward (line 596) | def forward(ctx, x, dt, A, B, C, chunk_size, D=None, z=None, dt_bias=N...
method backward (line 616) | def backward(ctx, dout, *args):
function mamba_chunk_scan_combined (line 624) | def mamba_chunk_scan_combined(x, dt, A, B, C, chunk_size, D=None, z=None...
function mamba_chunk_scan (line 646) | def mamba_chunk_scan(x, dt, A, B, C, chunk_size, D=None, z=None, dt_bias...
function ssd_chunk_scan_combined_ref (line 683) | def ssd_chunk_scan_combined_ref(x, dt, A, B, C, chunk_size, D=None, z=No...
function ssd_selective_scan (line 724) | def ssd_selective_scan(x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_...
function mamba_conv1d_scan_ref (line 775) | def mamba_conv1d_scan_ref(xBC, conv1d_weight, conv1d_bias, dt, A, chunk_...
class MambaSplitConv1dScanCombinedFn (line 816) | class MambaSplitConv1dScanCombinedFn(torch.autograd.Function):
method forward (line 820) | def forward(ctx, zxbcdt, conv1d_weight, conv1d_bias, dt_bias, A, D, ch...
method backward (line 898) | def backward(ctx, dout, *args):
function mamba_split_conv1d_scan_combined (line 978) | def mamba_split_conv1d_scan_combined(zxbcdt, conv1d_weight, conv1d_bias,...
function mamba_split_conv1d_scan_ref (line 1000) | def mamba_split_conv1d_scan_ref(zxbcdt, conv1d_weight, conv1d_bias, dt_b...
FILE: mamba_ssm/ops/triton/ssd_state_passing.py
function _state_passing_fwd_kernel (line 30) | def _state_passing_fwd_kernel(
function _state_passing_bwd_kernel (line 102) | def _state_passing_bwd_kernel(
function _state_passing_fwd (line 196) | def _state_passing_fwd(states, dA_chunk_cumsum, initial_states=None, seq...
function _state_passing_bwd (line 227) | def _state_passing_bwd(
class StatePassingFn (line 286) | class StatePassingFn(torch.autograd.Function):
method forward (line 289) | def forward(ctx, states, dA_chunk_cumsum, initial_states=None):
method backward (line 300) | def backward(ctx, dout, dfinal_states):
function state_passing (line 314) | def state_passing(states, dA_chunk_cumsum, initial_states=None):
function state_passing_ref (line 327) | def state_passing_ref(states, dA_chunk_cumsum, initial_states=None):
FILE: mamba_ssm/utils/determinism.py
function use_deterministic_mode (line 21) | def use_deterministic_mode():
function set_deterministic_mode (line 30) | def set_deterministic_mode(value):
function _estimate_config_cost (line 35) | def _estimate_config_cost(cfg):
function _filter_configs_by_block_sizes (line 44) | def _filter_configs_by_block_sizes(configs):
function autotune_configs (line 59) | def autotune_configs(configs):
function alloc_tile_workspace (line 80) | def alloc_tile_workspace(base_shape, tile_dim, dtype, device, determinis...
function finalize_tile_workspace (line 91) | def finalize_tile_workspace(tensor, deterministic):
FILE: mamba_ssm/utils/generation.py
class InferenceParams (line 18) | class InferenceParams:
method reset (line 29) | def reset(self, max_seqlen, max_batch_size):
function modify_logits_for_min_p_filtering (line 37) | def modify_logits_for_min_p_filtering(logits, min_p):
function modify_logits_for_top_k_filtering (line 45) | def modify_logits_for_top_k_filtering(logits, top_k):
function modify_logits_for_top_p_filtering (line 53) | def modify_logits_for_top_p_filtering(logits, top_p):
function modify_logit_for_repetition_penalty (line 69) | def modify_logit_for_repetition_penalty(logits, prev_output_tokens, repe...
function sample (line 83) | def sample(logits, top_k=1, top_p=0.0, min_p=0.0, temperature=1.0):
function decode (line 121) | def decode(
class GenerationMixin (line 246) | class GenerationMixin:
method allocate_inference_cache (line 247) | def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None,...
method generate (line 250) | def generate(
class DecodingCGCache (line 271) | class DecodingCGCache:
function update_graph_cache (line 283) | def update_graph_cache(
function capture_graph (line 342) | def capture_graph(
FILE: mamba_ssm/utils/hf.py
function load_config_hf (line 9) | def load_config_hf(model_name):
function load_state_dict_hf (line 14) | def load_state_dict_hf(model_name, device=None, dtype=None):
FILE: mamba_ssm/utils/torch.py
function custom_amp_decorator (line 5) | def custom_amp_decorator(dec: Callable, cuda_amp_deprecated: bool):
FILE: setup.py
function get_platform (line 47) | def get_platform():
function get_cuda_bare_metal_version (line 62) | def get_cuda_bare_metal_version(cuda_dir):
function get_hip_version (line 73) | def get_hip_version(rocm_dir):
function get_torch_hip_version (line 94) | def get_torch_hip_version():
function check_if_hip_home_none (line 102) | def check_if_hip_home_none(global_option: str) -> None:
function check_if_cuda_home_none (line 113) | def check_if_cuda_home_none(global_option: str) -> None:
function append_nvcc_threads (line 125) | def append_nvcc_threads(nvcc_extra_args):
function get_package_version (line 268) | def get_package_version():
function get_wheel_url (line 279) | def get_wheel_url():
class CachedWheelsCommand (line 328) | class CachedWheelsCommand(_bdist_wheel):
method run (line 336) | def run(self):
FILE: tests/benchmark_determinism_kernels.py
function _reset_peak_memory (line 18) | def _reset_peak_memory() -> None:
function _peak_memory_mb (line 25) | def _peak_memory_mb(fn, *, warmup: int = 3) -> float:
function make_tensors (line 35) | def make_tensors(*, batch: int, seqlen: int, nheads: int, headdim: int, ...
function get_benchmarks (line 57) | def get_benchmarks(t: dict[str, torch.Tensor], *, ngroups: int):
function _run_one (line 84) | def _run_one(fn, *, deterministic: bool, warmup: int, rep: int):
function main (line 91) | def main() -> None:
FILE: tests/ops/cute/test_mamba3_mimo_step.py
function _require_cuda_and_kernel_deps (line 46) | def _require_cuda_and_kernel_deps() -> None:
function _mamba3_cls (line 53) | def _mamba3_cls():
function _kernel_deps (line 60) | def _kernel_deps() -> None:
class InferenceParams (line 67) | class InferenceParams:
method reset (line 78) | def reset(self, max_seqlen, max_batch_size):
class RunOutputs (line 87) | class RunOutputs:
function _case_config (line 96) | def _case_config(*, is_outproj_norm: bool) -> dict:
function _diff_stats (line 113) | def _diff_stats(actual: Tensor, expected: Tensor) -> str:
function _assert_close (line 118) | def _assert_close(
function _run_case (line 141) | def _run_case(*, is_outproj_norm: bool) -> RunOutputs:
function test_step_matches_forward_fp32 (line 224) | def test_step_matches_forward_fp32(is_outproj_norm: bool) -> None:
function run_step_benchmark (line 253) | def run_step_benchmark(*, is_outproj_norm: bool) -> None:
FILE: tests/ops/test_selective_scan.py
function test_selective_scan (line 38) | def test_selective_scan(is_variable_B, is_variable_C, varBC_groups, has_...
function test_mamba_inner_fn (line 160) | def test_mamba_inner_fn(is_variable_B, is_variable_C, seqlen, itype, wty...
FILE: tests/ops/tilelang/test_mamba3_mimo.py
function _require_cuda_and_kernel_deps (line 54) | def _require_cuda_and_kernel_deps() -> None:
function mods (line 62) | def mods() -> SimpleNamespace:
function max_rel_err (line 77) | def max_rel_err(ours: Tensor, ref: Tensor, eps: float = 1e-5) -> float:
function assert_stable_rel (line 85) | def assert_stable_rel(
function build_inputs (line 114) | def build_inputs(
function make_smoke_inputs (line 181) | def make_smoke_inputs(
function grads_to_dA (line 292) | def grads_to_dA(grad_dA_cs: Tensor, grad_dA_cs_rev: Tensor, chunk_size: ...
function mamba3_MIMO_step_ref (line 307) | def mamba3_MIMO_step_ref(
function apply_angle_dt_reference (line 445) | def apply_angle_dt_reference(
function mamba3_MIMO_chunk_ref (line 455) | def mamba3_MIMO_chunk_ref(
function run_ref_backward_fp32 (line 638) | def run_ref_backward_fp32(
function test_mamba3_MIMO_chunk_ref_matches_step_ref (line 751) | def test_mamba3_MIMO_chunk_ref_matches_step_ref() -> None:
function test_fused_chunk_linear_attn_fwd_relative_error_lt_10pct (line 833) | def test_fused_chunk_linear_attn_fwd_relative_error_lt_10pct(
function test_fused_chunk_linear_attn_fwd_return_state_relative_error_lt_10pct (line 898) | def test_fused_chunk_linear_attn_fwd_return_state_relative_error_lt_10pct(
function test_fused_chunk_linear_attn_fwd_prereduce_relative_error_lt_10pct (line 977) | def test_fused_chunk_linear_attn_fwd_prereduce_relative_error_lt_10pct(
function test_mamba_mimo_bwd_combined_relative_errors_lt_10pct (line 1044) | def test_mamba_mimo_bwd_combined_relative_errors_lt_10pct(
function test_mamba_mimo_bwd_combined_prereduce_relative_errors_lt_10pct (line 1124) | def test_mamba_mimo_bwd_combined_prereduce_relative_errors_lt_10pct(
function test_mamba_mimo_smoke_forward_backward (line 1211) | def test_mamba_mimo_smoke_forward_backward(mods: SimpleNamespace) -> None:
FILE: tests/ops/triton/test_layernorm_gated.py
function test_layer_norm_gated (line 29) | def test_layer_norm_gated(d, dtype, wtype, has_bias, has_z, is_rms_norm,...
FILE: tests/ops/triton/test_mamba3_siso.py
function _segsum (line 21) | def _segsum(x: torch.Tensor) -> torch.Tensor:
function mamba3_siso_step_ref (line 33) | def mamba3_siso_step_ref(
function mamba3_siso_fwd_ref (line 148) | def mamba3_siso_fwd_ref(
function detach_clone (line 346) | def detach_clone(*args):
function relative_error (line 351) | def relative_error(
function create_mamba3_siso_inputs (line 406) | def create_mamba3_siso_inputs(
function test_mamba3_siso_step (line 487) | def test_mamba3_siso_step(nheads_qk=4, has_Z=True, has_D=True):
function test_mamba3_siso_combined_batched (line 561) | def test_mamba3_siso_combined_batched(nheads_qk=4, has_Z=True, has_D=Tru...
function test_mamba3_siso_combined_varlen (line 675) | def test_mamba3_siso_combined_varlen(nheads_qk=4, has_Z=True, has_D=True...
function test_mamba3_siso_step_ref_vs_fwd_ref (line 843) | def test_mamba3_siso_step_ref_vs_fwd_ref(nheads_qk=4, has_Z=True, has_D=...
FILE: tests/ops/triton/test_selective_state_update.py
function test_selective_state_update (line 22) | def test_selective_state_update(dim, dstate, has_z, itype):
function test_selective_state_update_with_heads (line 66) | def test_selective_state_update_with_heads(dim, dstate, ngroups, has_z, ...
function test_selective_state_update_with_batch_indices (line 112) | def test_selective_state_update_with_batch_indices(dim, dstate, has_z, i...
function test_selective_state_update_with_heads_with_batch_indices (line 161) | def test_selective_state_update_with_heads_with_batch_indices(dim, dstat...
FILE: tests/ops/triton/test_ssd.py
function detach_clone (line 20) | def detach_clone(*args):
function test_chunk_state_varlen (line 30) | def test_chunk_state_varlen(chunk_size, ngroups, dtype):
FILE: tests/test_determinism.py
function _set_deterministic (line 9) | def _set_deterministic(enabled: bool) -> None:
function _set_seeds (line 15) | def _set_seeds(seed: int) -> None:
function _max_abs_diff (line 20) | def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float:
function _make_inputs (line 24) | def _make_inputs(
function _run_case_outputs (line 85) | def _run_case_outputs(
function _kernel_is_reproducible (line 161) | def _kernel_is_reproducible(case: str, headdim: int, dstate: int, d_has_...
function _kernel_close_to_default (line 173) | def _kernel_close_to_default(case: str, headdim: int, dstate: int, d_has...
function test_kernel_reproducible (line 186) | def test_kernel_reproducible(case: str, headdim: int, dstate: int):
function test_combined_kernel_reproducible (line 194) | def test_combined_kernel_reproducible(case: str, d_has_hdim: bool, headd...
function test_kernel_close_to_default (line 202) | def test_kernel_close_to_default(case: str, headdim: int, dstate: int):
function test_combined_kernel_close_to_default (line 210) | def test_combined_kernel_close_to_default(case: str, d_has_hdim: bool, h...
function test_default_mode_is_not_reproducible (line 215) | def test_default_mode_is_not_reproducible():
function test_mamba2_fwd_bwd_deterministic_reproducible (line 265) | def test_mamba2_fwd_bwd_deterministic_reproducible():
function test_mamba2_fwd_bwd_deterministic_close_to_default (line 308) | def test_mamba2_fwd_bwd_deterministic_close_to_default():
FILE: tests/test_generation.py
function test_generation (line 13) | def test_generation():
function test_generation_varlen (line 46) | def test_generation_varlen():
function test_generation_varlen_with_padding (line 115) | def test_generation_varlen_with_padding():
Condensed preview — 94 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (1,256K chars).
[
{
"path": ".github/scripts/build.sh",
"chars": 1184,
"preview": "#!/bin/bash\n\nset -eoxu pipefail\n\n# We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA "
},
{
"path": ".github/scripts/check_for_ngc_images.sh",
"chars": 2265,
"preview": "#!/bin/bash\n\n# Configuration\nBASE_IMAGE=\"nvcr.io/nvidia/pytorch\"\nTAG_SUFFIX=\"-py3\"\nMONTHS_TO_CHECK=7 # Check current mon"
},
{
"path": ".github/scripts/test.sh",
"chars": 114,
"preview": "#!/bin/bash\n\nset -exou pipefail\n\npip install dist/*.whl\npython -c \"import mamba_ssm; print(mamba_ssm.__version__)\""
},
{
"path": ".github/workflows/_build.yml",
"chars": 6914,
"preview": "name: ~Build wheel template\n\non:\n workflow_call:\n inputs:\n runs-on:\n description: \"The runner to use for"
},
{
"path": ".github/workflows/_build_in_container.yml",
"chars": 5254,
"preview": "name: ~Build wheel template\n\non:\n workflow_call:\n inputs:\n runs-on:\n description: \"The runner to use for"
},
{
"path": ".github/workflows/build.yml",
"chars": 1473,
"preview": "name: Build wheels\n\non:\n workflow_dispatch:\n inputs:\n runs-on:\n description: \"The runner to use for the "
},
{
"path": ".github/workflows/build_in_container.yml",
"chars": 2213,
"preview": "name: Build wheels in a container\n\non:\n workflow_dispatch:\n inputs:\n runs-on:\n description: \"The runner "
},
{
"path": ".github/workflows/publish.yaml",
"chars": 5281,
"preview": "# This workflow will:\n# - Create a new Github release\n# - Build wheels for supported architectures\n# - Deploy the wheels"
},
{
"path": ".gitignore",
"chars": 52,
"preview": "*__pycache__/\n*.egg-info/\nbuild/\n**.so\n*.hip\n*_hip.*"
},
{
"path": ".gitmodules",
"chars": 144,
"preview": "[submodule \"3rdparty/lm-evaluation-harness\"]\n\tpath = 3rdparty/lm-evaluation-harness\n\turl = https://github.com/EleutherAI"
},
{
"path": "AUTHORS",
"chars": 53,
"preview": "Tri Dao, tri@tridao.me\nAlbert Gu, agu@andrew.cmu.edu\n"
},
{
"path": "LICENSE",
"chars": 11348,
"preview": " Apache License\n Version 2.0, January 2004\n "
},
{
"path": "MANIFEST.in",
"chars": 60,
"preview": "recursive-include csrc *\nrecursive-include csrc *\nREADME.md\n"
},
{
"path": "README.md",
"chars": 12080,
"preview": "# Mamba\n\n![Mamba](assets/selection.png \"Selective State Space\")\n> **Mamba: Linear-Time Sequence Modeling with Selective "
},
{
"path": "benchmarks/benchmark_generation_mamba_simple.py",
"chars": 3302,
"preview": "# Copyright (c) 2023, Tri Dao, Albert Gu.\n\nimport argparse\nimport time\nimport json\n\nimport torch\nimport torch.nn.functio"
},
{
"path": "csrc/selective_scan/reverse_scan.cuh",
"chars": 20552,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan.cpp",
"chars": 21689,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan.h",
"chars": 2771,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_bwd_bf16_complex.cu",
"chars": 395,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_bwd_bf16_real.cu",
"chars": 391,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_bwd_fp16_complex.cu",
"chars": 391,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_bwd_fp16_real.cu",
"chars": 387,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_bwd_fp32_complex.cu",
"chars": 388,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_bwd_fp32_real.cu",
"chars": 384,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_bwd_kernel.cuh",
"chars": 34955,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_common.h",
"chars": 9327,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_fwd_bf16.cu",
"chars": 500,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_fwd_fp16.cu",
"chars": 492,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_fwd_fp32.cu",
"chars": 486,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/selective_scan_fwd_kernel.cuh",
"chars": 20661,
"preview": "/******************************************************************************\n * Copyright (c) 2023, Tri Dao.\n *******"
},
{
"path": "csrc/selective_scan/static_switch.h",
"chars": 1278,
"preview": "// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h\n// and https://github.com/pyto"
},
{
"path": "csrc/selective_scan/uninitialized_copy.cuh",
"chars": 2828,
"preview": "/******************************************************************************\n * Copyright (c) 2011-2022, NVIDIA CORPO"
},
{
"path": "evals/lm_harness_eval.py",
"chars": 1287,
"preview": "import torch\n\nimport transformers\nfrom transformers import AutoTokenizer\n\nfrom mamba_ssm.models.mixer_seq_simple import "
},
{
"path": "mamba_ssm/__init__.py",
"chars": 308,
"preview": "__version__ = \"2.3.1\"\n\nfrom mamba_ssm.ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn\nfrom mamba_s"
},
{
"path": "mamba_ssm/distributed/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "mamba_ssm/distributed/distributed_utils.py",
"chars": 5825,
"preview": "from typing import Optional\n\nimport torch\nfrom torch import Tensor\nfrom torch.distributed import ProcessGroup\n\n# `all_ga"
},
{
"path": "mamba_ssm/distributed/tensor_parallel.py",
"chars": 12008,
"preview": "# Copyright (c) 2024, Tri Dao.\n# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/m"
},
{
"path": "mamba_ssm/models/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "mamba_ssm/models/config_mamba.py",
"chars": 489,
"preview": "from dataclasses import dataclass, field\n\n\n@dataclass\nclass MambaConfig:\n\n d_model: int = 2560\n d_intermediate: in"
},
{
"path": "mamba_ssm/models/mixer_seq_simple.py",
"chars": 11665,
"preview": "# Copyright (c) 2023, Albert Gu, Tri Dao.\n\nimport math\nfrom functools import partial\nimport json\nimport os\nimport copy\n\n"
},
{
"path": "mamba_ssm/modules/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "mamba_ssm/modules/block.py",
"chars": 3729,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\nfrom typing import Optional\n\nimport torch\nfrom torch import nn, Tensor\n\nfrom m"
},
{
"path": "mamba_ssm/modules/mamba2.py",
"chars": 17479,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\nimport math\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as"
},
{
"path": "mamba_ssm/modules/mamba2_simple.py",
"chars": 7608,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\nimport math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as "
},
{
"path": "mamba_ssm/modules/mamba3.py",
"chars": 19494,
"preview": "# Copyright (c) 2026, Dao AI Lab, Goombalab.\n\nimport math\nfrom einops import rearrange, repeat\n\nimport torch\nimport torc"
},
{
"path": "mamba_ssm/modules/mamba_simple.py",
"chars": 11710,
"preview": "# Copyright (c) 2023, Tri Dao, Albert Gu.\n\nimport math\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\ni"
},
{
"path": "mamba_ssm/modules/mha.py",
"chars": 13392,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\nimport math\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as"
},
{
"path": "mamba_ssm/modules/mlp.py",
"chars": 1130,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\nclass GatedMLP(nn."
},
{
"path": "mamba_ssm/modules/ssd_minimal.py",
"chars": 4144,
"preview": "# Copyright (c) 2024, Albert Gu and Tri Dao.\n\"\"\"Minimal implementation of SSD.\n\nThis is the same as Listing 1 from the p"
},
{
"path": "mamba_ssm/ops/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "mamba_ssm/ops/cute/mamba3/mamba3_step_fn.py",
"chars": 41469,
"preview": "# Copyright (c) 2025, Tri Dao.\n# Modified to use tvm-ffi and fake tensors instead of dlpack.\n# Modified to optionally up"
},
{
"path": "mamba_ssm/ops/selective_scan_interface.py",
"chars": 20025,
"preview": "# Copyright (c) 2023, Tri Dao, Albert Gu.\n\nimport torch\nimport torch.nn.functional as F\nfrom mamba_ssm.utils.torch impor"
},
{
"path": "mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py",
"chars": 9530,
"preview": "\"\"\"Mamba-3 Tilelang Autograd Wrapper\n\nInterface for Mamba-3 Tilelang kernels with automatic differentiation\n\nCopyright ("
},
{
"path": "mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_bwd.py",
"chars": 72346,
"preview": "\"\"\"\nTilelang implementation of Mamba3 backward kernels,\nwith MIMO support.\n\nCopyright (c) 2026, Dao AI Lab, Goombalab\n\n\""
},
{
"path": "mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py",
"chars": 25050,
"preview": "\"\"\"\nTilelang implementation of Mamba3 forward kernel,\nwith MIMO support.\n\nCopyright (c) 2026, Dao AI Lab, Goombalab\n\n\"\"\""
},
{
"path": "mamba_ssm/ops/triton/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "mamba_ssm/ops/triton/angle_cumsum.py",
"chars": 30189,
"preview": "# Copyright (c) 2025, Tri Dao.\n\nfrom typing import Optional\nimport math\n\nimport torch\n\nimport triton\nimport triton.langu"
},
{
"path": "mamba_ssm/ops/triton/k_activations.py",
"chars": 5263,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\nimport torch\n\nimport triton\nimport triton.language as tl\n\nfrom mamba_ssm.util"
},
{
"path": "mamba_ssm/ops/triton/layer_norm.py",
"chars": 36058,
"preview": "# Copyright (c) 2024, Tri Dao.\n# Implement dropout + residual + layer_norm / rms_norm.\n\n# Based on the Triton LayerNorm "
},
{
"path": "mamba_ssm/ops/triton/layernorm_gated.py",
"chars": 17763,
"preview": "# Copyright (c) 2024, Tri Dao.\n# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tu"
},
{
"path": "mamba_ssm/ops/triton/mamba3/angle_dt.py",
"chars": 17266,
"preview": "from typing import Tuple, Optional\n\nimport torch\nfrom torch import Tensor\n\nimport triton\nimport triton.language as tl\nfr"
},
{
"path": "mamba_ssm/ops/triton/mamba3/mamba3_mimo_rotary_step.py",
"chars": 16415,
"preview": "# Copyright (c) 2025, Tri Dao.\n# We need a pretty recent version of triton to support tuples. 3.3 definitely will work,\n"
},
{
"path": "mamba_ssm/ops/triton/mamba3/mamba3_mimo_utils.py",
"chars": 32830,
"preview": "\"\"\"\nFused Triton kernels for Mamba3 backward pass ddt computation.\n\nThis module implements fused kernels that combine th"
},
{
"path": "mamba_ssm/ops/triton/mamba3/mamba3_siso_bwd.py",
"chars": 84717,
"preview": "\"\"\"\nMamba-3 Backward Pass Triton Kernels.\n\nCopyright (c) 2026, Dao AI Lab, Goombalab\n\"\"\"\n\nfrom typing import Optional, T"
},
{
"path": "mamba_ssm/ops/triton/mamba3/mamba3_siso_combined.py",
"chars": 16908,
"preview": "\"\"\"Mamba-3 Triton Autograd Wrapper\n\nCopyright (c) 2025, Dao AI Lab, Goombalab\n\"\"\"\n\nfrom __future__ import annotations\n\nf"
},
{
"path": "mamba_ssm/ops/triton/mamba3/mamba3_siso_fwd.py",
"chars": 35716,
"preview": "\"\"\"\nMamba-3 SISO Forward Pass Triton Kernel.\n\nCopyright (c) 2025, Dao AI Lab, Goombalab\n\"\"\"\n\nfrom typing import Optional"
},
{
"path": "mamba_ssm/ops/triton/mamba3/mamba3_siso_step.py",
"chars": 18045,
"preview": "\"\"\"\nMamba-3 Step Kernel.\n\nCopyright (c) 2025, Dao AI Lab, Goombalab\n\"\"\"\n\nfrom typing import Optional, Tuple\nimport math\n"
},
{
"path": "mamba_ssm/ops/triton/mamba3/utils.py",
"chars": 3289,
"preview": "\"\"\"\nMamba-3 Util Functions.\n\nCopyright (c) 2025, Dao AI Lab, Goombalab\n\"\"\"\n\nimport triton\nimport triton.language as tl\n\n"
},
{
"path": "mamba_ssm/ops/triton/selective_state_update.py",
"chars": 11582,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\n\"\"\"We want triton==2.1.0 or triton==2.2.0 or triton==2.3.0 for this\n\"\"\"\n\nimpo"
},
{
"path": "mamba_ssm/ops/triton/softplus.py",
"chars": 328,
"preview": "import triton\nimport triton.language as tl\nfrom packaging import version\n\nTRITON3 = version.parse(triton.__version__) >="
},
{
"path": "mamba_ssm/ops/triton/ssd_bmm.py",
"chars": 13731,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\n\"\"\"We want triton==2.1.0 or 2.2.0 for this\n\"\"\"\n\nimport math\nimport torch\nimpo"
},
{
"path": "mamba_ssm/ops/triton/ssd_chunk_scan.py",
"chars": 107618,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\n\"\"\"We want triton==2.1.0 or 2.2.0 for this\n\"\"\"\n\nimport math\nfrom packaging im"
},
{
"path": "mamba_ssm/ops/triton/ssd_chunk_state.py",
"chars": 61766,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\n\"\"\"We want triton==2.1.0 or 2.2.0 for this\n\"\"\"\n\nimport math\nimport torch\nimpo"
},
{
"path": "mamba_ssm/ops/triton/ssd_combined.py",
"chars": 56740,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\n\"\"\"We want triton==2.1.0 or 2.2.0 for this\n\"\"\"\n\nfrom typing import Optional\n\n"
},
{
"path": "mamba_ssm/ops/triton/ssd_state_passing.py",
"chars": 16471,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\n\"\"\"We want triton==2.1.0 or 2.2.0 for this\n\"\"\"\n\nimport math\nimport torch\nimpo"
},
{
"path": "mamba_ssm/utils/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "mamba_ssm/utils/determinism.py",
"chars": 3209,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\nimport os\nimport warnings\nfrom packaging import version\n\nimport torch\n\ntry:\n "
},
{
"path": "mamba_ssm/utils/generation.py",
"chars": 15331,
"preview": "# Copyright (c) 2023, Albert Gu, Tri Dao.\nimport gc\nimport time\nfrom collections import namedtuple\nfrom dataclasses impo"
},
{
"path": "mamba_ssm/utils/hf.py",
"chars": 966,
"preview": "import json\n\nimport torch\n\nfrom transformers.utils import WEIGHTS_NAME, CONFIG_NAME\nfrom transformers.utils.hub import c"
},
{
"path": "mamba_ssm/utils/torch.py",
"chars": 676,
"preview": "import torch\nfrom functools import partial\nfrom typing import Callable\n\ndef custom_amp_decorator(dec: Callable, cuda_amp"
},
{
"path": "pyproject.toml",
"chars": 1221,
"preview": "[project]\nname = \"mamba_ssm\"\ndescription = \"Mamba state-space model\"\nreadme = \"README.md\"\nauthors = [\n { name = \"Tri "
},
{
"path": "rocm_patch/rocm6_0.patch",
"chars": 2249,
"preview": "--- /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h\t2023-12-12 20:11:48.000000000 +0000\n+++ rocm_update_files/amd_hip_bf"
},
{
"path": "setup.py",
"chars": 15113,
"preview": "# Copyright (c) 2023, Albert Gu, Tri Dao.\nimport sys\nimport warnings\nimport os\nimport re\nimport ast\nfrom pathlib import "
},
{
"path": "tests/benchmark_determinism_kernels.py",
"chars": 6427,
"preview": "#!/usr/bin/env python\n# Copyright (c) 2024, Tri Dao, Albert Gu.\n\nimport gc\nimport math\n\nimport torch\nfrom triton.testing"
},
{
"path": "tests/ops/cute/test_mamba3_mimo_step.py",
"chars": 9341,
"preview": "\"\"\"\nMamba-3 MIMO Step Function Tests\n\nCopyright (c) 2026, Dao AI Lab, Goombalab\n\nPytest coverage for Mamba3.step() and m"
},
{
"path": "tests/ops/test_selective_scan.py",
"chars": 13064,
"preview": "# Copyright (C) 2023, Tri Dao.\n\nimport math\n\nimport torch\nimport torch.nn.functional as F\nimport pytest\n\nfrom einops imp"
},
{
"path": "tests/ops/tilelang/test_mamba3_mimo.py",
"chars": 39289,
"preview": "\"\"\"\nMamba-3 MIMO Kernel Tests\n\nCopyright (c) 2026, Dao AI Lab, Goombalab\n\n\nUsage:\npytest -q -s -p no:warnings tests/ops/"
},
{
"path": "tests/ops/triton/test_layernorm_gated.py",
"chars": 5498,
"preview": "import math\n\nimport torch\nimport torch.nn.functional as F\n\nimport pytest\n\nfrom einops import rearrange, repeat\n\nfrom mam"
},
{
"path": "tests/ops/triton/test_mamba3_siso.py",
"chars": 38418,
"preview": "\"\"\"\nMamba-3 SISO Kernel Tests\n\nCopyright (c) 2025, Dao AI Lab, Goombalab\n\"\"\"\n\nimport copy\nimport math\nfrom typing import"
},
{
"path": "tests/ops/triton/test_selective_state_update.py",
"chars": 9504,
"preview": "# Copyright (C) 2023, Tri Dao.\n\nimport math\n\nimport torch\nimport torch.nn.functional as F\nimport pytest\n\nfrom einops imp"
},
{
"path": "tests/ops/triton/test_ssd.py",
"chars": 3883,
"preview": "import math\n\nimport torch\nimport torch.nn.functional as F\n\nimport pytest\n\nfrom einops import rearrange, repeat\n\nfrom mam"
},
{
"path": "tests/test_determinism.py",
"chars": 12450,
"preview": "# Copyright (c) 2024, Tri Dao, Albert Gu.\n\nimport os\n\nimport pytest\nimport torch\n\n\ndef _set_deterministic(enabled: bool)"
},
{
"path": "tests/test_generation.py",
"chars": 8686,
"preview": "import torch\nimport torch.nn.functional as F\n\nfrom mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel\nfrom mamba_"
},
{
"path": "usage.md",
"chars": 1380,
"preview": "# Mamba adoption\n\nWe've been very happy to see Mamba being adopted by many organizations\nand research labs to speed up t"
}
]
About this extraction
This page contains the full source code of the state-spaces/mamba GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 94 files (1.2 MB), approximately 345.8k tokens, and a symbol index with 465 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.