flashspec 0.1.2__tar.gz → 0.1.3__tar.gz

{flashspec-0.1.2 → flashspec-0.1.3}/.github/workflows/benchmark.yml

@@ -25,7 +25,7 @@ jobs:
           cache: pip
 
       - name: Install bench dependencies
-        run: pip install -e ".[bench]"
+        run: pip install -e ".[bench,gpu]"
 
       - name: Build Docker image
         run: docker build -t flashspec:latest -f deploy/Dockerfile .

{flashspec-0.1.2 → flashspec-0.1.3}/.github/workflows/ci.yml

@@ -22,7 +22,7 @@ jobs:
           python-version: "3.11"
           cache: pip
       - name: Install dev dependencies
-        run: pip install -e ".[dev]" --extra-index-url https://download.pytorch.org/whl/cpu
+        run: pip install -e ".[dev,gpu]" --extra-index-url https://download.pytorch.org/whl/cpu
       - name: make lint
         run: make lint
       - name: pip-audit (fail on CRITICAL CVE — §15)
@@ -46,7 +46,7 @@ jobs:
           python-version: "3.11"
           cache: pip
       - name: Install
-        run: pip install -e ".[dev]" --extra-index-url https://download.pytorch.org/whl/cpu
+        run: pip install -e ".[dev,gpu]" --extra-index-url https://download.pytorch.org/whl/cpu
       - name: make test
         run: make test
       - name: Upload coverage XML
@@ -70,7 +70,7 @@ jobs:
           python-version: "3.11"
           cache: pip
       - name: Install
-        run: pip install -e ".[dev]" --extra-index-url https://download.pytorch.org/whl/cpu
+        run: pip install -e ".[dev,gpu]" --extra-index-url https://download.pytorch.org/whl/cpu
       - name: make test-chaos
         run: make test-chaos
 
@@ -86,7 +86,7 @@ jobs:
           python-version: "3.11"
           cache: pip
       - name: Install
-        run: pip install -e ".[dev]" --extra-index-url https://download.pytorch.org/whl/cpu
+        run: pip install -e ".[dev,gpu]" --extra-index-url https://download.pytorch.org/whl/cpu
       - name: Download coverage artifact
         uses: actions/download-artifact@v4
         with:
@@ -117,7 +117,7 @@ jobs:
           cache: pip
       - name: Install with ONNX extras
         run: |
-          pip install -e ".[dev,onnx]" \
+          pip install -e ".[dev,onnx,gpu]" \
             --extra-index-url https://download.pytorch.org/whl/cpu
       - name: Run ONNX parity tests
         run: pytest tests/integration/test_onnx_parity.py -v

{flashspec-0.1.2 → flashspec-0.1.3}/.github/workflows/gpu_tests.yml

@@ -20,7 +20,7 @@ jobs:
           python-version: "3.11"
           cache: pip
       - name: Install
-        run: pip install -e ".[dev]"
+        run: pip install -e ".[dev,gpu]"
       - name: make test-gpu
         run: make test-gpu
       - name: Run slow distribution tests (N=10,000 KS gate)
@@ -45,7 +45,7 @@ jobs:
           python-version: "3.11"
           cache: pip
       - name: Install bench dependencies
-        run: pip install -e ".[bench]"
+        run: pip install -e ".[bench,gpu]"
       - name: make bench-quick
         run: make bench-quick
       - name: Assert throughput ≥ 1.5× AR

{flashspec-0.1.2 → flashspec-0.1.3}/CHANGELOG.md

@@ -10,7 +10,47 @@ Versioning follows [Semantic Versioning 2.0.0](https://semver.org/).
 
 ## [Unreleased]
 
-### Added
+No unreleased changes. Next planned milestone: v0.2.0 (real GPU benchmark
+results, completed paper figures, arXiv submission).
+
+---
+
+## [0.1.3] — Critical packaging fix (Windows/macOS install)
+
+### Fixed
+- **Critical**: `pip install flashspec` failed completely on Windows and
+  macOS with `ERROR: No matching distribution found for triton>=3.0.0` —
+  Triton publishes official PyPI wheels for Linux only, but `triton>=3.0.0`
+  was a hard core dependency, making the package uninstallable on those
+  platforms. (Reported against v0.1.0–v0.1.2.)
+  - `pyproject.toml`: moved `triton` out of `dependencies` into a new
+    `gpu` extra with a `platform_system == 'Linux'` environment marker:
+    `pip install flashspec[gpu]`.
+  - `flashspec/kernels/__init__.py`: now imports `verify_kernel` and
+    `gather_kernel` inside a `try/except ImportError`. If Triton is
+    unavailable, `verify_tokens` and `gather_accepted` are bound to a
+    wrapper that raises an actionable `ImportError` only when *called*
+    (not at import time), explaining `pip install flashspec[gpu]`
+    (Linux-only) and pointing to the cross-platform
+    `flashspec.kernels._reference` module (identical numerics, verified in
+    `tests/unit/test_verify_kernel.py`). Adds module-level
+    `flashspec.kernels.TRITON_AVAILABLE: bool`.
+  - `tests/unit/test_verify_kernel.py`: added
+    `TestKernelsPackageWithoutTriton` (3 tests, CPU-only, no GPU/Triton
+    required) verifying the package imports cleanly and raises the
+    actionable error when Triton is absent.
+  - `.github/workflows/{ci,gpu_tests,benchmark}.yml`: updated `pip install`
+    steps to request the `gpu` extra wherever Triton-backed kernel tests
+    or benchmarks actually run (all on `ubuntu-22.04`/self-hosted Linux
+    runners, where the extra resolves normally).
+  - `README.md`: rewrote the Installation section — `pip install flashspec`
+    is now correctly described as cross-platform (CPU reference kernels),
+    with `pip install flashspec[gpu]` for Linux+CUDA accelerated kernels.
+  - Bumped version `0.1.0` → `0.1.3` in `pyproject.toml` and
+    `flashspec/__init__.py` (0.1.0–0.1.2 are broken on non-Linux and
+    should be marked yanked on PyPI after this release).
+
+---
 - Initial project scaffold following AGENTS.md specification.
 - `flashspec.utils.config`: Pydantic v2 configuration models (`FlashSpecConfig`,
   `BanditConfig`, `SamplingConfig`, `MetricsConfig`). All immutable (`frozen=True`).

{flashspec-0.1.2 → flashspec-0.1.3}/CITATION.cff

@@ -27,7 +27,7 @@ abstract: >-
   O(sqrt(K T log T)) regret bound. The output distribution is
   provably identical to standard autoregressive sampling, verified
   continuously via a Kolmogorov-Smirnov test in CI.
-version: 0.1.0
+version: 0.1.3
 date-released: "2026-06-12"
 license: Apache-2.0
 repository-code: "https://github.com/Mattral/FlashSpec"

{flashspec-0.1.2 → flashspec-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashspec
-Version: 0.1.2
+Version: 0.1.3
 Summary: Adaptive speculative-decoding inference engine with Triton-optimised verification and online bandit draft selection
 Project-URL: Homepage, https://github.com/Mattral/FlashSpec
 Project-URL: Repository, https://github.com/Mattral/FlashSpec
@@ -136,7 +136,6 @@ Requires-Dist: numpy>=1.26.0
 Requires-Dist: pydantic>=2.0.0
 Requires-Dist: torch>=2.2.0
 Requires-Dist: transformers>=4.40.0
-Requires-Dist: triton>=3.0.0
 Provides-Extra: bench
 Requires-Dist: lm-eval>=0.4; extra == 'bench'
 Requires-Dist: matplotlib>=3.8; extra == 'bench'
@@ -152,6 +151,8 @@ Requires-Dist: ruff>=0.4.0; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: mkdocs-material>=9.5; extra == 'docs'
 Requires-Dist: mkdocstrings[python]>=0.24; extra == 'docs'
+Provides-Extra: gpu
+Requires-Dist: triton>=3.0.0; (platform_system == 'Linux') and extra == 'gpu'
 Provides-Extra: onnx
 Requires-Dist: onnx>=1.16; extra == 'onnx'
 Requires-Dist: onnxruntime-gpu>=1.18; extra == 'onnx'

{flashspec-0.1.2 → flashspec-0.1.3}/flashspec/__init__.py

@@ -14,6 +14,7 @@ Public API surface (AGENTS.md §13.2 — do not modify without explicit approval
     flashspec.register          (draft model decorator)
     flashspec.get_drafter
     flashspec.list_drafters
+    flashspec.TRITON_AVAILABLE  (bool — True only on Linux with triton installed)
 
 References
 ----------
@@ -25,6 +26,7 @@ References
 
 from flashspec.engine.drafter import get_drafter, list_drafters, register
 from flashspec.engine.speculative import GenerationResult, SpeculativeEngine
+from flashspec.kernels import TRITON_AVAILABLE
 from flashspec.utils.config import BanditConfig, FlashSpecConfig, MetricsConfig, SamplingConfig
 
 __all__ = [
@@ -34,10 +36,11 @@ __all__ = [
     "MetricsConfig",
     "SamplingConfig",
     "SpeculativeEngine",
+    "TRITON_AVAILABLE",
     "get_drafter",
     "list_drafters",
     "register",
 ]
 
-__version__ = "0.1.0"
+__version__ = "0.1.3"
 __author__ = "Min Htet Myet (Mattral)"

flashspec-0.1.3/flashspec/kernels/__init__.py

@@ -0,0 +1,72 @@
+"""Triton-optimised kernels sub-package.
+
+All consumers of kernels must import from this module, not from the
+individual kernel files directly.  This is enforced by import-linter.
+
+Private implementation modules:
+- ``verify_kernel``: token acceptance kernel (requires Triton).
+- ``gather_kernel``: accepted-token gather kernel (requires Triton).
+- ``_reference``: pure-PyTorch reference implementations (cross-platform,
+  used by tests and as the runtime fallback when Triton is unavailable).
+
+Notes
+-----
+Triton has no official PyPI wheels for Windows or macOS.  On those
+platforms (or any platform where ``triton`` is not installed),
+``verify_tokens`` and ``gather_accepted`` are bound to wrappers that raise
+a clear, actionable ``ImportError`` when called, rather than failing at
+import time.  Use ``flashspec.kernels._reference`` directly for a
+cross-platform pure-PyTorch implementation with identical numerics
+(verified in ``tests/unit/test_verify_kernel.py``).
+
+To enable the accelerated kernels, install on Linux with a CUDA-capable
+GPU::
+
+    pip install flashspec[gpu]
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+__all__ = [
+    "gather_accepted",
+    "verify_tokens",
+    "TRITON_AVAILABLE",
+]
+
+if TYPE_CHECKING:
+    import torch
+
+try:
+    from flashspec.kernels.gather_kernel import gather_accepted
+    from flashspec.kernels.verify_kernel import verify_tokens
+
+    TRITON_AVAILABLE: bool = True
+except ImportError as _exc:  # pragma: no cover — exercised on non-Linux CI
+    _IMPORT_ERROR = _exc
+    TRITON_AVAILABLE = False
+
+    def _triton_unavailable(*_args: Any, **_kwargs: Any) -> "torch.Tensor":
+        """Raise an actionable error explaining how to get Triton kernels.
+
+        Raises
+        ------
+        ImportError
+            Always. Explains that Triton is not installed on this platform
+            and points to the pure-PyTorch reference implementation.
+        """
+        raise ImportError(
+            "flashspec's accelerated Triton kernels are not available "
+            f"(original error: {_IMPORT_ERROR}).\n\n"
+            "Triton ships official PyPI wheels for Linux only. To use the "
+            "GPU-accelerated kernels, install on Linux with a CUDA-capable "
+            "GPU via:\n\n    pip install flashspec[gpu]\n\n"
+            "For a cross-platform pure-PyTorch implementation with "
+            "identical numerics, use "
+            "flashspec.kernels._reference.verify_tokens_reference and "
+            "flashspec.kernels._reference.gather_accepted_reference instead."
+        )
+
+    verify_tokens = _triton_unavailable
+    gather_accepted = _triton_unavailable

{flashspec-0.1.2 → flashspec-0.1.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "flashspec"
-version = "0.1.2"
+version = "0.1.3"
 description = "Adaptive speculative-decoding inference engine with Triton-optimised verification and online bandit draft selection"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -20,13 +20,18 @@ classifiers = [
 ]
 dependencies = [
     "torch>=2.2.0",
-    "triton>=3.0.0",
     "transformers>=4.40.0",
     "pydantic>=2.0.0",
     "numpy>=1.26.0",
 ]
 
 [project.optional-dependencies]
+# Triton has no official PyPI wheels for Windows or macOS — it is Linux-only.
+# Without this extra, flashspec runs on the pure-PyTorch reference kernels
+# (flashspec.kernels._reference), which are numerically identical (see
+# tests/unit/test_verify_kernel.py) but slower. Install this extra on Linux
+# with a CUDA-capable GPU to enable the accelerated Triton kernels.
+gpu = ["triton>=3.0.0; platform_system == 'Linux'"]
 onnx = ["onnx>=1.16", "onnxruntime-gpu>=1.18"]
 dev = [
     "pytest>=8.0",

{flashspec-0.1.2 → flashspec-0.1.3}/tests/unit/test_verify_kernel.py

@@ -368,3 +368,103 @@ class TestVerifyTokensTritonParity:
         assert (first_rej.cpu() == 0).all(), (
             f"Expected all first_rej==0, got {first_rej.cpu().tolist()}"
         )
+
+
+# ── Triton-unavailable fallback (cross-platform, no GPU required) ───────────
+
+
+class TestKernelsPackageWithoutTriton:
+    """``flashspec.kernels`` must import successfully even if Triton is
+    absent (e.g. on Windows or macOS, where Triton ships no PyPI wheels).
+
+    Simulates the absence of Triton by removing ``triton`` and the kernel
+    submodules from ``sys.modules`` and blocking their import, then
+    re-importing ``flashspec.kernels`` fresh.
+    """
+
+    def test_kernels_package_imports_without_triton(self) -> None:
+        """``import flashspec.kernels`` succeeds even when Triton is missing."""
+        import builtins
+        import importlib
+        import sys
+
+        real_import = builtins.__import__
+
+        def fake_import(name, *args, **kwargs):
+            if name == "triton" or name.startswith("triton."):
+                raise ModuleNotFoundError(f"No module named '{name}'")
+            return real_import(name, *args, **kwargs)
+
+        # Drop any cached modules so re-import re-triggers the try/except.
+        for mod in list(sys.modules):
+            if mod.startswith("flashspec.kernels"):
+                del sys.modules[mod]
+
+        builtins.__import__ = fake_import
+        try:
+            kernels = importlib.import_module("flashspec.kernels")
+        finally:
+            builtins.__import__ = real_import
+            for mod in list(sys.modules):
+                if mod.startswith("flashspec.kernels"):
+                    del sys.modules[mod]
+
+        assert kernels.TRITON_AVAILABLE is False, (
+            "TRITON_AVAILABLE should be False when triton import fails"
+        )
+
+    def test_verify_tokens_raises_actionable_import_error_without_triton(self) -> None:
+        """Calling verify_tokens without Triton raises ImportError mentioning flashspec[gpu]."""
+        import builtins
+        import importlib
+        import sys
+
+        real_import = builtins.__import__
+
+        def fake_import(name, *args, **kwargs):
+            if name == "triton" or name.startswith("triton."):
+                raise ModuleNotFoundError(f"No module named '{name}'")
+            return real_import(name, *args, **kwargs)
+
+        for mod in list(sys.modules):
+            if mod.startswith("flashspec.kernels"):
+                del sys.modules[mod]
+
+        builtins.__import__ = fake_import
+        try:
+            kernels = importlib.import_module("flashspec.kernels")
+            with pytest.raises(ImportError, match=r"flashspec\[gpu\]"):
+                kernels.verify_tokens(None, None, None, None)
+        finally:
+            builtins.__import__ = real_import
+            for mod in list(sys.modules):
+                if mod.startswith("flashspec.kernels"):
+                    del sys.modules[mod]
+
+    def test_gather_accepted_raises_actionable_import_error_without_triton(self) -> None:
+        """Calling gather_accepted without Triton raises ImportError mentioning _reference."""
+        import builtins
+        import importlib
+        import sys
+
+        real_import = builtins.__import__
+
+        def fake_import(name, *args, **kwargs):
+            if name == "triton" or name.startswith("triton."):
+                raise ModuleNotFoundError(f"No module named '{name}'")
+            return real_import(name, *args, **kwargs)
+
+        for mod in list(sys.modules):
+            if mod.startswith("flashspec.kernels"):
+                del sys.modules[mod]
+
+        builtins.__import__ = fake_import
+        try:
+            kernels = importlib.import_module("flashspec.kernels")
+            with pytest.raises(ImportError, match="_reference"):
+                kernels.gather_accepted(None, None)
+        finally:
+            builtins.__import__ = real_import
+            for mod in list(sys.modules):
+                if mod.startswith("flashspec.kernels"):
+                    del sys.modules[mod]

flashspec-0.1.2/flashspec/kernels/__init__.py

@@ -1,18 +0,0 @@
-"""Triton-optimised kernels sub-package.
-
-All consumers of kernels must import from this module, not from the
-individual kernel files directly.  This is enforced by import-linter.
-
-Private implementation modules:
-- ``verify_kernel``: token acceptance kernel.
-- ``gather_kernel``: accepted-token gather kernel.
-- ``_reference``: pure-PyTorch reference implementations (tests only).
-"""
-
-from flashspec.kernels.gather_kernel import gather_accepted
-from flashspec.kernels.verify_kernel import verify_tokens
-
-__all__ = [
-    "gather_accepted",
-    "verify_tokens",
-]