stackkernels 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (52) hide show
  1. stackkernels-0.1.0/.github/workflows/publish.yml +28 -0
  2. stackkernels-0.1.0/.github/workflows/test.yml +33 -0
  3. stackkernels-0.1.0/.gitignore +11 -0
  4. stackkernels-0.1.0/LICENSE +0 -0
  5. stackkernels-0.1.0/PKG-INFO +39 -0
  6. stackkernels-0.1.0/README.md +21 -0
  7. stackkernels-0.1.0/benchmarks/bench_activation.py +0 -0
  8. stackkernels-0.1.0/benchmarks/bench_attention.py +0 -0
  9. stackkernels-0.1.0/benchmarks/bench_norm.py +0 -0
  10. stackkernels-0.1.0/benchmarks/results/.gitkeep +0 -0
  11. stackkernels-0.1.0/docs/backends.md +0 -0
  12. stackkernels-0.1.0/docs/getting_started.md +0 -0
  13. stackkernels-0.1.0/docs/index.md +0 -0
  14. stackkernels-0.1.0/examples/jax_usage.py +0 -0
  15. stackkernels-0.1.0/examples/pytorch_usage.py +0 -0
  16. stackkernels-0.1.0/examples/stackformer_integration.py +0 -0
  17. stackkernels-0.1.0/pyproject.toml +28 -0
  18. stackkernels-0.1.0/stackkernels/__init__.py +0 -0
  19. stackkernels-0.1.0/stackkernels/hub/__init__.py +0 -0
  20. stackkernels-0.1.0/stackkernels/hub/registry.py +0 -0
  21. stackkernels-0.1.0/stackkernels/jax/__init__.py +0 -0
  22. stackkernels-0.1.0/stackkernels/jax/_utils.py +0 -0
  23. stackkernels-0.1.0/stackkernels/jax/activation.py +0 -0
  24. stackkernels-0.1.0/stackkernels/jax/attention.py +0 -0
  25. stackkernels-0.1.0/stackkernels/jax/loss.py +0 -0
  26. stackkernels-0.1.0/stackkernels/jax/norm.py +0 -0
  27. stackkernels-0.1.0/stackkernels/torch/__init__.py +0 -0
  28. stackkernels-0.1.0/stackkernels/torch/_utils.py +0 -0
  29. stackkernels-0.1.0/stackkernels/torch/activation.py +0 -0
  30. stackkernels-0.1.0/stackkernels/torch/attention.py +0 -0
  31. stackkernels-0.1.0/stackkernels/torch/dropout.py +0 -0
  32. stackkernels-0.1.0/stackkernels/torch/embedding.py +0 -0
  33. stackkernels-0.1.0/stackkernels/torch/loss.py +0 -0
  34. stackkernels-0.1.0/stackkernels/torch/norm.py +0 -0
  35. stackkernels-0.1.0/stackkernels/triton/__init__.py +0 -0
  36. stackkernels-0.1.0/stackkernels/triton/_heuristics.py +0 -0
  37. stackkernels-0.1.0/stackkernels/triton/activation/gelu.py +0 -0
  38. stackkernels-0.1.0/stackkernels/triton/activation/silu.py +0 -0
  39. stackkernels-0.1.0/stackkernels/triton/activation/swiglu.py +0 -0
  40. stackkernels-0.1.0/stackkernels/triton/attention/causal_mask.py +0 -0
  41. stackkernels-0.1.0/stackkernels/triton/attention/flash_attn.py +0 -0
  42. stackkernels-0.1.0/stackkernels/triton/attention/paged_attn.py +0 -0
  43. stackkernels-0.1.0/stackkernels/triton/attention/varlen_attn.py +0 -0
  44. stackkernels-0.1.0/stackkernels/triton/dropout/dropout.py +0 -0
  45. stackkernels-0.1.0/stackkernels/triton/loss/cross_entropy.py +0 -0
  46. stackkernels-0.1.0/stackkernels/triton/norm/fused_norm.py +0 -0
  47. stackkernels-0.1.0/stackkernels/triton/norm/layer_norm.py +0 -0
  48. stackkernels-0.1.0/stackkernels/triton/norm/rms_norm.py +0 -0
  49. stackkernels-0.1.0/stackkernels/triton/positional/rope.py +0 -0
  50. stackkernels-0.1.0/tests/test_jax/.gitkeep +0 -0
  51. stackkernels-0.1.0/tests/test_torch/.gitkeep +0 -0
  52. stackkernels-0.1.0/tests/test_triton/.gitkeep +0 -0
@@ -0,0 +1,28 @@
1
+ name: Publish to PyPI
2
+
3
+ on:
4
+ release:
5
+ types: [published]
6
+
7
+ jobs:
8
+ publish:
9
+ runs-on: ubuntu-latest
10
+ steps:
11
+ - uses: actions/checkout@v4
12
+
13
+ - name: Set up Python
14
+ uses: actions/setup-python@v5
15
+ with:
16
+ python-version: "3.11"
17
+
18
+ - name: Install build tools
19
+ run: pip install build twine
20
+
21
+ - name: Build package
22
+ run: python -m build
23
+
24
+ - name: Publish to PyPI
25
+ env:
26
+ TWINE_USERNAME: __token__
27
+ TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
28
+ run: twine upload dist/*
@@ -0,0 +1,33 @@
1
+ name: Test
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+ branches: [main]
8
+
9
+ jobs:
10
+ build:
11
+ runs-on: ubuntu-latest
12
+ strategy:
13
+ matrix:
14
+ python-version: ["3.10", "3.11"]
15
+
16
+ steps:
17
+ - uses: actions/checkout@v4
18
+
19
+ - name: Set up Python ${{ matrix.python-version }}
20
+ uses: actions/setup-python@v5
21
+ with:
22
+ python-version: ${{ matrix.python-version }}
23
+
24
+ - name: Install dependencies
25
+ run: |
26
+ python -m pip install --upgrade pip
27
+ pip install -e ".[dev]"
28
+
29
+ - name: Check package imports
30
+ run: python -c "import stackkernels"
31
+
32
+ - name: Run CPU-safe tests
33
+ run: pytest tests/ -m "not gpu"
@@ -0,0 +1,11 @@
1
+ __pycache__/
2
+ *.pyc
3
+ *.egg-info/
4
+ build/
5
+ dist/
6
+ .pytest_cache/
7
+ .venv/
8
+ venv/
9
+ *.log
10
+ benchmarks/results/*.csv
11
+ benchmarks/results/*.png
File without changes
@@ -0,0 +1,39 @@
1
+ Metadata-Version: 2.5
2
+ Name: stackkernels
3
+ Version: 0.1.0
4
+ Summary: Fast, fused Triton kernels for transformer models — PyTorch and JAX backends.
5
+ Project-URL: Homepage, https://github.com/stackformer-labs/stackkernels
6
+ License: MIT
7
+ License-File: LICENSE
8
+ Requires-Python: >=3.9
9
+ Requires-Dist: torch<3.0,>=2.3
10
+ Requires-Dist: triton>=2.3
11
+ Provides-Extra: dev
12
+ Requires-Dist: pytest; extra == 'dev'
13
+ Requires-Dist: pytest-benchmark; extra == 'dev'
14
+ Provides-Extra: jax
15
+ Requires-Dist: jax-triton; extra == 'jax'
16
+ Requires-Dist: jax>=0.4; extra == 'jax'
17
+ Description-Content-Type: text/markdown
18
+
19
+ # stackkernels
20
+
21
+ Fast, fused Triton kernels for transformer models — built as the kernel backbone for [Stackformer](#).
22
+
23
+ Supports **PyTorch** and **JAX**, with a shared, framework-agnostic Triton core.
24
+
25
+ ## Status
26
+ 🚧 Early development. Build order: RMSNorm → LayerNorm → Flash Attention → Dropout → Activations (SiLU/GELU/SwiGLU) → Cross-Entropy → RoPE.
27
+
28
+ ## Install
29
+ ```bash
30
+ pip install stackkernels
31
+ ```
32
+
33
+ ## Quick start
34
+ ```python
35
+ from stackkernels.torch import RMSNorm
36
+ ```
37
+
38
+ ## License
39
+ MIT
@@ -0,0 +1,21 @@
1
+ # stackkernels
2
+
3
+ Fast, fused Triton kernels for transformer models — built as the kernel backbone for [Stackformer](#).
4
+
5
+ Supports **PyTorch** and **JAX**, with a shared, framework-agnostic Triton core.
6
+
7
+ ## Status
8
+ 🚧 Early development. Build order: RMSNorm → LayerNorm → Flash Attention → Dropout → Activations (SiLU/GELU/SwiGLU) → Cross-Entropy → RoPE.
9
+
10
+ ## Install
11
+ ```bash
12
+ pip install stackkernels
13
+ ```
14
+
15
+ ## Quick start
16
+ ```python
17
+ from stackkernels.torch import RMSNorm
18
+ ```
19
+
20
+ ## License
21
+ MIT
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
@@ -0,0 +1,28 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "stackkernels"
7
+ version = "0.1.0"
8
+ description = "Fast, fused Triton kernels for transformer models — PyTorch and JAX backends."
9
+ readme = "README.md"
10
+ license = { text = "MIT" }
11
+ requires-python = ">=3.9"
12
+ dependencies = [
13
+ "torch>=2.3,<3.0",
14
+ "triton>=2.3",
15
+ ]
16
+
17
+ [project.optional-dependencies]
18
+ jax = [
19
+ "jax>=0.4",
20
+ "jax-triton",
21
+ ]
22
+ dev = [
23
+ "pytest",
24
+ "pytest-benchmark",
25
+ ]
26
+
27
+ [project.urls]
28
+ Homepage = "https://github.com/stackformer-labs/stackkernels"
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes