Skip to content

README.md

Latest commit

 

History

History
222 lines (161 loc) · 5.93 KB

README.md

File metadata and controls

222 lines (161 loc) · 5.93 KB
Raw

FlashAttention for GCU

This repository contains a GCU-oriented fork of FlashAttention for PyTorch. It keeps the familiar flash_attn Python API shape while providing a GCU backend, GCU-specific tests, and GCU-focused packaging/build scripts.

Overview

This branch is intended for environments that already provide the GCU runtime stack required by the extension build and execution path.

Exposed Python APIs include:

  • flash_attn_func
  • flash_attn_qkvpacked_func
  • flash_attn_kvpacked_func
  • flash_attn_varlen_func
  • flash_attn_varlen_qkvpacked_func
  • flash_attn_varlen_kvpacked_func
  • flash_attn_with_kvcache

The implementation and tests in this branch target the GCU device path rather than the upstream generic release flow.

Requirements

  • Linux
  • PyTorch
  • torch_gcu
  • A GCU software stack that provides the required runtime headers and libraries
  • A C++17-capable compiler toolchain

Installation

Option 1: Docker (Recommended)

The easiest way to get started is using the pre-built Docker image with all dependencies included.

  1. Pull and start the container:

    IMAGE=registry-egc.enflame-tech.com/artifacts/public_pytorch:v2.11.0-TR3.8.106-ubuntu2204

docker run --name flash_attn -d \
  -v /home:/home \
  --shm-size 8G \
  --ipc=host --network host \
  --cap-add SYS_PTRACE \
  --security-opt seccomp=unconfined \
  --privileged \
  "$IMAGE" \
  tail -f /dev/null

  2. Update the host GCU driver (to match the image's software version):

    # Extract the matching driver from the container
docker cp flash_attn:/enflame/driver ./

# Install the driver on the host
sudo driver/enflame-x86_64-gcc-*.run

# Restart the container to pick up the new driver
docker restart flash_attn

  3. Clone the source code:

    cd /home
git clone git@github.com:EnflameTechnology/flash-attention.git

  4. Enter the container and build:

    docker exec -it flash_attn bash
cd /home/flash-attention
./install.sh

Option 2: build and install with the helper script

Build from source in an environment where the GCU runtime stack is already installed.

./install.sh

Option 3: build a wheel manually

python setup.py bdist_wheel
pip install --force-reinstall --no-deps dist/*.whl

Optional build environment variables

The build script reads the following environment variables when present:

  • BASE_DIR
  • COMPILE_CXX_FLAGS
  • LINK_FLAGS
  • LINK_LIBS
  • CMAKE_BUILD_TYPE
  • PACKAGE_VERSION

Use these only when your local toolchain/runtime layout requires overrides.

Quick Start

The following snippets mirror the APIs exercised by the test suite under tests/ (flash_attn_func, flash_attn_varlen_func).

Standard attention (flash_attn_func)

import torch
import torch_gcu
from torch_gcu import transfer_to_gcu  # monkeypatches cuda -> gcu

from flash_attn import flash_attn_func

batch_size, seqlen = 2, 128
num_heads, head_dim = 8, 64

q = torch.randn(batch_size, seqlen, num_heads, head_dim, dtype=torch.float16)
k = torch.randn(batch_size, seqlen, num_heads, head_dim, dtype=torch.float16)
v = torch.randn(batch_size, seqlen, num_heads, head_dim, dtype=torch.float16)

out = flash_attn_func(q.gcu(), k.gcu(), v.gcu(), dropout_p=0.0, causal=False)
print(out.shape)  # torch.Size([2, 128, 8, 64])

GQA (grouped-query attention) is supported by passing different head counts for k / v:

num_kv_heads = 2
q = torch.randn(batch_size, seqlen, num_heads, head_dim, dtype=torch.float16)
k = torch.randn(batch_size, seqlen, num_kv_heads, head_dim, dtype=torch.float16)
v = torch.randn(batch_size, seqlen, num_kv_heads, head_dim, dtype=torch.float16)

out = flash_attn_func(q.gcu(), k.gcu(), v.gcu(), dropout_p=0.0, causal=True)

Variable-length attention (flash_attn_varlen_func)

import torch
import torch_gcu
from torch_gcu import transfer_to_gcu

from flash_attn import flash_attn_varlen_func

batch_size, seqlen = 2, 128
num_heads, head_dim = 8, 64
total = batch_size * seqlen

q = torch.randn(total, num_heads, head_dim, dtype=torch.float16)
k = torch.randn(total, num_heads, head_dim, dtype=torch.float16)
v = torch.randn(total, num_heads, head_dim, dtype=torch.float16)

cu_seqlens = torch.arange(
    0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32
).gcu()

out = flash_attn_varlen_func(
    q.gcu(), k.gcu(), v.gcu(), cu_seqlens, cu_seqlens, seqlen, seqlen,
    causal=True, window_size=(-1, -1), softcap=0.0,
)

For paged KV-cache / varlen inference, see tests/test_flash_attn_vllm_gcu.py::test_varlen_with_paged_kv, which uses flash_attn.vllm_flash_attn.flash_attn_varlen_func with block_table and seqused_k.

Testing

The test suite under tests/ is intentionally minimal and focused on the GCU device path. Run it from the repository root:

# Standard and GQA forward
pytest -q tests/test_flash_attn_s60.py

# Variable-length forward
pytest -q tests/test_flash_attn_varlen_s60.py

# Paged KV-cache varlen (vLLM path)
pytest -q tests/test_flash_attn_vllm_gcu.py::test_varlen_with_paged_kv

Or run all of them together:

pytest -q tests/

Project Structure

  • flash_attn/ — Python package and public API wrappers
  • csrc_gcu/ — GCU C++ extension sources
  • tests/ — GCU test coverage (flash_attn_func, flash_attn_varlen_func, and the vLLM paged-KV varlen path)

Notes

  • This branch is a GCU-targeted fork and is not documented as a generic accelerator release.
  • Public dependency names such as torch_gcu and topsaten are kept as-is in code paths where they are part of the runtime/build contract.
  • For concrete usage patterns, see the tests in tests/.

License

This repository is released under the terms of the BSD 3-Clause License. See LICENSE.

Attribution

This repository is based on the FlashAttention project and preserves upstream attribution present in source files and repository metadata.