Add aarch64 CI and wheel publishing targets (#83)

Add Linux aarch64 coverage for package build/import validation and include
aarch64 wheels in the publish workflow. Also add Windows amd64 wheel builds
alongside the existing Windows support.

Use a lightweight C++ extension smoke test for cibuildwheel so publishing
checks package installation without requiring GPU runtime dependencies.

Signed-off-by: Takeshi Yoshimura <tyos@jp.ibm.com>

Author: takeshi-yoshimura

.github/workflows/publish.yaml

@@ -7,11 +7,25 @@ on:
 
 jobs:
   build-wheels:
-    name: Build wheel for ${{ matrix.python-version }}
-    runs-on: ubuntu-latest
+    name: Build wheel for ${{ matrix.python-version }} ${{ matrix.platform.arch }}
+    runs-on: ${{ matrix.platform.runner }}
     strategy:
+      fail-fast: false
       matrix:
         python-version: ["cp310", "cp311", "cp312", "cp313", "cp314"]
+        platform:
+          - arch: x86_64
+            runner: ubuntu-latest
+            build: manylinux_x86_64
+            cibw_arch: x86_64
+          - arch: aarch64
+            runner: ubuntu-24.04-arm
+            build: manylinux_aarch64
+            cibw_arch: aarch64
+          - arch: win_amd64
+            runner: windows-latest
+            build: win_amd64
+            cibw_arch: AMD64
 
     steps:
     - name: Checkout code
@@ -28,15 +42,15 @@ jobs:
       run: |
         cibuildwheel --output-dir wheelhouse
       env:
-        CIBW_BUILD: "${{ matrix.python-version }}-manylinux_x86_64"
+        CIBW_BUILD: "${{ matrix.python-version }}-${{ matrix.platform.build }}"
         CIBW_SKIP: "*-musllinux_* *-win32 *-manylinux_i686"
-        CIBW_TEST_SKIP: "*"
-        CIBW_ARCHS: "x86_64"
+        CIBW_TEST_COMMAND: "python {project}/tests/smoke_import_cpp.py"
+        CIBW_ARCHS: "${{ matrix.platform.cibw_arch }}"
 
     - name: Upload wheel artifact
       uses: actions/upload-artifact@v4
       with:
-        name: wheels-${{ matrix.python-version }}
+        name: wheels-${{ matrix.python-version }}-${{ matrix.platform.arch }}
         path: wheelhouse/*.whl
 
   build-sdist:

.github/workflows/test-aarch64.yaml

@@ -0,0 +1,37 @@
+name: test-aarch64
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  build-and-test-aarch64:
+    name: Test aarch64 Python 3.14
+    runs-on: ubuntu-24.04-arm
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.14"
+
+      - name: Install system dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libnuma-dev gcc g++ make
+
+      - name: Build package
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install .
+
+      - name: Run extension smoke test
+        run: |
+          python tests/smoke_import_cpp.py

tests/smoke_import_cpp.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Smoke-test the compiled C++ extension without importing package __init__."""
+
+import importlib.machinery
+import importlib.metadata
+import importlib.util
+import sys
+from pathlib import Path
+
+
+def _load_cpp_extension():
+    dist = importlib.metadata.distribution("fastsafetensors")
+    dist_root = Path(dist.locate_file(""))
+
+    for suffix in importlib.machinery.EXTENSION_SUFFIXES:
+        candidate = dist_root / "fastsafetensors" / f"cpp{suffix}"
+        if candidate.exists():
+            spec = importlib.util.spec_from_file_location("cpp", candidate)
+            if spec is None or spec.loader is None:
+                raise RuntimeError(f"failed to create import spec for {candidate}")
+            module = importlib.util.module_from_spec(spec)
+            sys.modules[spec.name] = module
+            spec.loader.exec_module(module)
+            return module
+
+    suffixes = ", ".join(importlib.machinery.EXTENSION_SUFFIXES)
+    raise FileNotFoundError(
+        f"fastsafetensors.cpp extension was not found under {dist_root}; "
+        f"checked suffixes: {suffixes}"
+    )
+
+
+def main() -> None:
+    cpp = _load_cpp_extension()
+    assert cpp.get_alignment_size() == 4096
+    cpp.load_library_functions("")
+    assert isinstance(cpp.is_cuda_found(), bool)
+    assert isinstance(cpp.is_hip_found(), bool)
+    assert isinstance(cpp.is_cufile_found(), bool)
+
+
+if __name__ == "__main__":
+    main()

tests/vllm/test_vllm.py

@@ -5,7 +5,7 @@
 import pytest
 
 pytest.importorskip("vllm", reason="vllm integration tests require vllm")
-from vllm import LLM, SamplingParams
+from vllm import LLM, SamplingParams  # type: ignore[attr-defined]
 
 sampling_params = SamplingParams(
     temperature=0.0,