Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion .env.example
Original file line number Diff line number Diff line change
@@ -1,3 +1,9 @@
# Build configuration
TORCH_CUDA_ARCH_LIST=80;90
CUDA_HOME=/usr/local/cuda
CUDA_CACHE_PATH={HOME}/.cache/stream_attn/cuda
TRITON_CACHE_DIR={HOME}/.cache/stream_attn/triton

# StreamAttention environment configuration example
# Copy to .env and edit as needed

Expand All @@ -18,4 +24,4 @@ STREAM_ATTENTION_RING_OVERLAP_SIZE=256
# Star Attention
STREAM_ATTENTION_STAR_BLOCK_SIZE=2048
STREAM_ATTENTION_STAR_ANCHOR_SIZE=256
STREAM_ATTENTION_STAR_NUM_HOSTS=1
STREAM_ATTENTION_STAR_NUM_HOSTS=1
28 changes: 28 additions & 0 deletions .github/workflows/build.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
name: Build Wheels

on:
push:
branches: [ main ]
pull_request:
branches: [ main ]

jobs:
build:
runs-on: ubuntu-latest
strategy:
matrix:
sm: [80, 90]
env:
TORCH_CUDA_ARCH_LIST: ${{ matrix.sm }}
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v4
with:
python-version: '3.10'
- name: Install build tooling
run: |
python -m pip install --upgrade pip
pip install build
- name: Build wheel
run: |
python -m build --wheel
25 changes: 25 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
[build-system]
requires = [
"setuptools>=61",
"wheel",
"torch==2.1.0",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PyTorch is listed as a build requirement, which will trigger heavyweight installs in the isolated build env; keep build-system requires minimal (build tools only).

Prompt for AI agents
Address the following comment on pyproject.toml at line 5:

<comment>PyTorch is listed as a build requirement, which will trigger heavyweight installs in the isolated build env; keep build-system requires minimal (build tools only).</comment>

<file context>
@@ -0,0 +1,25 @@
+[build-system]
+requires = [
+    &quot;setuptools&gt;=61&quot;,
+    &quot;wheel&quot;,
+    &quot;torch==2.1.0&quot;,
+    &quot;triton==2.1.0&quot;,
+]
</file context>

"triton==2.1.0",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Triton is listed in build requirements; move it to project dependencies to avoid heavyweight installs during build isolation.

Prompt for AI agents
Address the following comment on pyproject.toml at line 6:

<comment>Triton is listed in build requirements; move it to project dependencies to avoid heavyweight installs during build isolation.</comment>

<file context>
@@ -0,0 +1,25 @@
+[build-system]
+requires = [
+    &quot;setuptools&gt;=61&quot;,
+    &quot;wheel&quot;,
+    &quot;torch==2.1.0&quot;,
+    &quot;triton==2.1.0&quot;,
+]
+build-backend = &quot;setuptools.build_meta&quot;
</file context>

]
build-backend = "setuptools.build_meta"

[project]
name = "stream-attention"
version = "1.0.0"
description = "Production-ready multi-GPU FlashAttention implementation with support for extremely long contexts"
authors = [{name = "StreamAttention Team", email = "streamattention@example.com"}]
readme = "README.md"
license = {file = "LICENSE"}
requires-python = ">=3.8"
dependencies = [
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PyYAML is imported but not declared as a dependency; add "pyyaml" to [project].dependencies to fix ModuleNotFoundError.

(Based on your team's feedback about pytest failing with ModuleNotFoundError: yaml.)

Prompt for AI agents
Address the following comment on pyproject.toml at line 18:

<comment>PyYAML is imported but not declared as a dependency; add &quot;pyyaml&quot; to [project].dependencies to fix ModuleNotFoundError.

(Based on your team&#39;s feedback about pytest failing with ModuleNotFoundError: yaml.)</comment>

<file context>
@@ -0,0 +1,25 @@
+[build-system]
+requires = [
+    &quot;setuptools&gt;=61&quot;,
+    &quot;wheel&quot;,
+    &quot;torch==2.1.0&quot;,
+    &quot;triton==2.1.0&quot;,
+]
+build-backend = &quot;setuptools.build_meta&quot;
+
</file context>

"torch==2.1.0",
"triton==2.1.0",
]

[tool.stream_attention]
cuda_version = "12.1"
triton_version = "2.1.0"