torch-linear-assignment 0.1.0__py3-none-any.whl

torch_linear_assignment/__init__.py

@@ -0,0 +1,5 @@
+"""Public exports for batched linear assignment."""
+
+from .assignment import assignment_to_indices, batch_linear_assignment
+
+__all__ = ["assignment_to_indices", "batch_linear_assignment"]

torch_linear_assignment/_triton.py

@@ -0,0 +1,309 @@
+"""Triton implementation of batched rectangular linear assignment.
+
+Purpose:
+    Provide the private CUDA implementation selected by ``assignment.py`` on
+    supported Linux NVIDIA systems. The module owns no public dispatch policy;
+    importing it is deliberately deferred until the caller has established that
+    CUDA, the Torch version, and the GPU capability are eligible.
+
+Scope:
+    One Triton program solves one batch item. The program is intentionally
+    correctness-first: outer augmenting-path steps remain sequential while each
+    current set of candidate columns is processed as a masked lane vector.
+    Per-search state remains program-local. The correctness launch currently
+    keeps its loop-carried vectors within one warp; wider launch configurations
+    remain gated on exact GPU parity. Its ordering reproduces the legacy CUDA
+    and SciPy rectangular-LSAP algorithm.
+
+Usage:
+    ``batch_linear_assignment(cuda_cost)`` accepts a three-dimensional CUDA
+    cost tensor and returns the stable public ``(B, W)`` long assignment shape.
+    The optional validation setting is private benchmark instrumentation; normal
+    dispatch always uses the default full validation.
+
+Outputs:
+    The kernel writes assignment and temporary Torch workspaces on the input
+    device. The wrapper returns worker-to-task indices, using ``-1`` for an
+    unmatched worker, and raises ordinary ``ValueError`` exceptions for invalid
+    numeric values or infeasible matrices.
+
+Failure:
+    NaN and negative infinity are rejected by one host reduction before launch.
+    The kernel records one infeasibility flag per batch item instead of asserting
+    on device; the normal wrapper synchronizes once to turn any flag into an
+    exception. Triton import, compilation, and GPU execution failures are not
+    swallowed because they do not mean the optional dependency is merely absent.
+
+Used by:
+    ``torch_linear_assignment.assignment.batch_linear_assignment_cuda``. The
+    module is private so its benchmark validation switch and workspace layout do
+    not become part of the package API.
+"""
+
+from typing import Literal
+
+import torch
+import triton
+import triton.language as tl
+
+ValidationMode = Literal["full", "off", "nonfinite_only", "infeasibility_flag_only"]
+_VALIDATION_MODES: frozenset[str] = frozenset({"full", "off", "nonfinite_only", "infeasibility_flag_only"})
+
+
+@triton.jit
+def _linear_assignment_kernel(
+    cost,
+    u,
+    v,
+    col4row,
+    row4col,
+    infeasible,
+    NUM_ROWS: tl.constexpr,
+    NUM_COLUMNS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    IS_FP64: tl.constexpr,
+):
+    """Solve one rectangular assignment problem in a single Triton program."""
+    batch_index = tl.program_id(0)
+    lane = tl.arange(0, BLOCK_N)
+    column_mask = lane < NUM_COLUMNS
+    row_mask = lane < NUM_ROWS
+
+    cost = cost + batch_index * NUM_ROWS * NUM_COLUMNS
+    u = u + batch_index * NUM_ROWS
+    v = v + batch_index * NUM_COLUMNS
+    col4row = col4row + batch_index * NUM_ROWS
+    row4col = row4col + batch_index * NUM_COLUMNS
+
+    for current_row in tl.range(0, NUM_ROWS, num_stages=1):
+        scanned_rows = tl.zeros((BLOCK_N,), tl.int1)
+        scanned_columns = tl.zeros((BLOCK_N,), tl.int1)
+        active_columns = column_mask
+        scan_positions = NUM_COLUMNS - lane - 1
+        path = tl.full((BLOCK_N,), -1, tl.int32)
+        if IS_FP64:
+            shortest_path_costs = tl.full((BLOCK_N,), float("inf"), tl.float64)
+            scanned_row_costs = tl.zeros((BLOCK_N,), tl.float64)
+        else:
+            shortest_path_costs = tl.full((BLOCK_N,), float("inf"), tl.float32)
+            scanned_row_costs = tl.zeros((BLOCK_N,), tl.float32)
+
+        sink = -1
+        current_potential = tl.load(u + current_row)
+        # A workspace-derived zero keeps this loop-carried value FP32 or FP64.
+        min_value = current_potential - current_potential
+        search_row = current_row
+        num_remaining = NUM_COLUMNS
+
+        # Wide matrices usually reach an unmatched sink early; do not execute
+        # the remaining masked column-search iterations after that point.
+        searching = True
+        while searching:
+            scanned_rows = scanned_rows | (row_mask & (lane == search_row) & searching)
+            candidate_mask = active_columns & searching
+            previous_cost = shortest_path_costs
+            reduced_cost = (
+                min_value
+                - tl.load(u + search_row)
+                + tl.load(
+                    cost + search_row * NUM_COLUMNS + lane,
+                    mask=candidate_mask,
+                    other=float("inf"),
+                )
+                - tl.load(v + lane, mask=candidate_mask, other=0.0)
+            )
+            improved = candidate_mask & (reduced_cost < previous_cost)
+            candidate_cost = tl.where(improved, reduced_cost, previous_cost)
+            shortest_path_costs = candidate_cost
+            path = tl.where(improved, search_row, path)
+
+            # The legacy solver scans a swap-removed ``remaining`` array. Track
+            # each column's current scan position locally to preserve its exact
+            # last-unmatched/first-matched tie rule without shared scratch writes.
+            lowest = tl.min(
+                tl.where(candidate_mask, candidate_cost, float("inf")),
+                axis=0,
+            )
+            has_candidate = lowest != float("inf")
+            step_valid = searching & has_candidate
+            tied = candidate_mask & (candidate_cost == lowest)
+            matched_rows = tl.load(row4col + lane, mask=candidate_mask, other=0)
+            last_unmatched_position = tl.max(
+                tl.where(tied & (matched_rows == -1), scan_positions, -1),
+                axis=0,
+            )
+            first_tied_position = tl.min(
+                tl.where(tied, scan_positions, NUM_COLUMNS),
+                axis=0,
+            )
+            selected_position = tl.where(
+                last_unmatched_position >= 0,
+                last_unmatched_position,
+                first_tied_position,
+            )
+            selected_column = tl.max(
+                tl.where(candidate_mask & (scan_positions == selected_position), lane, -1),
+                axis=0,
+            )
+            selected_column = tl.where(step_valid, selected_column, 0)
+            selected_row = tl.load(row4col + selected_column)
+
+            # A matched row enters the search tree through ``selected_column``.
+            # Save that column's shortest cost by row now, avoiding a dynamic
+            # gather from the loop-carried column vector after the search.
+            matched_step = step_valid & (selected_row != -1)
+            scanned_row_costs = tl.where(
+                row_mask & (lane == selected_row) & matched_step,
+                lowest,
+                scanned_row_costs,
+            )
+            sink = tl.where(step_valid & (selected_row == -1), selected_column, sink)
+            search_row = tl.where(
+                matched_step,
+                selected_row,
+                search_row,
+            )
+            scanned_columns = scanned_columns | (column_mask & (lane == selected_column) & step_valid)
+            last_index = num_remaining - 1
+            last_column = tl.max(
+                tl.where(active_columns & (scan_positions == last_index), lane, -1),
+                axis=0,
+            )
+            scan_positions = tl.where(
+                active_columns & (lane == last_column) & step_valid,
+                selected_position,
+                scan_positions,
+            )
+            active_columns = active_columns & ~((lane == selected_column) & step_valid)
+            num_remaining = tl.where(step_valid, num_remaining - 1, num_remaining)
+            min_value = tl.where(step_valid, lowest, min_value)
+            tl.store(infeasible + batch_index, 1, mask=searching & ~has_candidate)
+            searching = matched_step
+
+        solved = sink != -1
+        tl.store(infeasible + batch_index, 1, mask=~solved)
+
+        tl.store(u + current_row, current_potential + min_value, mask=solved)
+        visited_rows = scanned_rows
+        update_row = solved & row_mask & visited_rows & (lane != current_row)
+        tl.store(
+            u + lane,
+            tl.load(u + lane, mask=row_mask, other=0.0) + min_value - scanned_row_costs,
+            mask=update_row,
+        )
+
+        update_column = solved & column_mask & scanned_columns
+        tl.store(
+            v + lane,
+            tl.load(v + lane, mask=column_mask, other=0.0) - min_value + shortest_path_costs,
+            mask=update_column,
+        )
+
+        augmenting = solved
+        augmenting_column = sink
+        augmentation_steps = 0
+        # Recover only the actual path while retaining the row-count safety cap.
+        while augmenting & (augmentation_steps < NUM_ROWS):
+            safe_column = tl.where(augmenting, augmenting_column, 0)
+            augmenting_row = tl.max(
+                tl.where(column_mask & (lane == safe_column), path, -1),
+                axis=0,
+            )
+            safe_row = tl.where(augmenting, augmenting_row, 0)
+            previous_column = tl.load(col4row + safe_row)
+            tl.store(row4col + safe_column, augmenting_row, mask=augmenting)
+            tl.store(col4row + safe_row, augmenting_column, mask=augmenting)
+            augmenting = augmenting & (augmenting_row != current_row)
+            augmenting_column = tl.where(augmenting, previous_column, augmenting_column)
+            augmentation_steps += 1
+
+        # The next row consumes the assignments and potentials written above.
+        tl.debug_barrier()
+
+
+def _validate_mode(validation: ValidationMode) -> None:
+    """Reject an unsupported private validation configuration."""
+    if validation not in _VALIDATION_MODES:
+        raise ValueError(f"Unknown Triton validation mode: {validation!r}.")
+
+
+def _reject_invalid_numeric_entries(cost: torch.Tensor) -> None:
+    """Raise the SciPy-compatible error for NaN or negative-infinite costs."""
+    if not cost.is_floating_point():
+        return
+    invalid = torch.logical_or(torch.isnan(cost), torch.isneginf(cost)).any()
+    if bool(invalid.item()):
+        raise ValueError("matrix contains invalid numeric entries")
+
+
+def _solver_cost_dtype(cost: torch.Tensor) -> torch.Tensor:
+    """Preserve float64 costs and promote every other solver input to float32."""
+    if cost.dtype == torch.float64:
+        return cost
+    return cost.to(torch.float32)
+
+
+def _empty_assignment(cost: torch.Tensor) -> torch.Tensor:
+    """Return the public assignment shape without launching a zero-size kernel."""
+    batch_size, workers, _ = cost.shape
+    return torch.full((batch_size, workers), -1, dtype=torch.long, device=cost.device)
+
+
+def _solve(cost: torch.Tensor, validation: ValidationMode) -> tuple[torch.Tensor, torch.Tensor]:
+    """Launch the one-program-per-batch kernel and return both assignment views."""
+    batch_size, rows, columns = cost.shape
+    block_n = triton.next_power_of_2(columns)
+    workspace_options = {"device": cost.device}
+    scalar_options = {**workspace_options, "dtype": cost.dtype}
+    integer_options = {**workspace_options, "dtype": torch.int32}
+    u = torch.zeros((batch_size, rows), **scalar_options)
+    v = torch.zeros((batch_size, columns), **scalar_options)
+    col4row = torch.full((batch_size, rows), -1, **integer_options)
+    row4col = torch.full((batch_size, columns), -1, **integer_options)
+    infeasible = torch.zeros((batch_size,), **integer_options)
+
+    _linear_assignment_kernel[(batch_size,)](
+        cost,
+        u,
+        v,
+        col4row,
+        row4col,
+        infeasible,
+        NUM_ROWS=rows,
+        NUM_COLUMNS=columns,
+        BLOCK_N=block_n,
+        IS_FP64=cost.dtype == torch.float64,
+        # Multi-warp loop-carried state produces false infeasibility on the
+        # saved 512-lane L4 case; keep one warp until a wider design is exact.
+        num_warps=1,
+    )
+    if validation in {"full", "infeasibility_flag_only"} and bool(infeasible.any().item()):
+        raise ValueError("cost matrix is infeasible")
+    return col4row, row4col
+
+
+def batch_linear_assignment(
+    cost: torch.Tensor,
+    *,
+    validation: ValidationMode = "full",
+) -> torch.Tensor:
+    """Solve a CUDA batch with full validation unless private benchmarks opt out."""
+    _validate_mode(validation)
+    if cost.ndim != 3:
+        raise ValueError("Need 3-dimensional tensor with shape (B, W, T).")
+    if not cost.is_cuda:
+        raise ValueError("Triton linear assignment requires a CUDA tensor.")
+    if validation in {"full", "nonfinite_only"}:
+        _reject_invalid_numeric_entries(cost)
+
+    cost = _solver_cost_dtype(cost)
+    batch_size, workers, tasks = cost.shape
+    if batch_size == 0 or workers == 0 or tasks == 0:
+        return _empty_assignment(cost)
+
+    with torch.amp.autocast("cuda", enabled=False):
+        if tasks < workers:
+            _, row4col = _solve(cost.transpose(1, 2).contiguous(), validation)
+            return row4col.to(torch.long)
+        col4row, _ = _solve(cost.contiguous(), validation)
+        return col4row.to(torch.long)

torch_linear_assignment/assignment.py

@@ -0,0 +1,163 @@
+"""Public linear-assignment API with optional CUDA acceleration."""
+
+import importlib
+import re
+import sys
+import warnings
+from types import ModuleType
+
+import torch
+from scipy.optimize import linear_sum_assignment
+
+_CUDA_FALLBACK_WARNING_EMITTED = False
+_TRITON_MINIMUM_TORCH_VERSION = (2, 4)
+
+
+def _prepare_solver_cost(cost: torch.Tensor) -> torch.Tensor:
+    """Apply the documented solver dtype without discarding complex values."""
+    if cost.dtype in {torch.float32, torch.float64} or cost.is_complex():
+        return cost
+    return cost.to(torch.float32)
+
+
+def batch_linear_assignment_cpu(cost: torch.Tensor) -> torch.Tensor:
+    """Solve batched assignment with SciPy on the cost tensor's CPU device."""
+    cost = _prepare_solver_cost(cost)
+    batch_size, workers, _ = cost.shape
+    matching = torch.full([batch_size, workers], -1, dtype=torch.long, device=cost.device)
+    for batch_index in range(batch_size):
+        row_indices, column_indices = linear_sum_assignment(cost[batch_index].numpy(), maximize=False)
+        matching[batch_index].scatter_(
+            0,
+            torch.from_numpy(row_indices),
+            torch.from_numpy(column_indices),
+        )
+    return matching
+
+
+def _torch_supports_triton() -> bool:
+    """Return whether the installed Torch version meets the Triton-path floor."""
+    version = re.match(r"(\d+)\.(\d+)", torch.__version__)
+    return version is not None and tuple(map(int, version.groups())) >= _TRITON_MINIMUM_TORCH_VERSION
+
+
+def _load_triton_backend() -> ModuleType | None:
+    """Load the optional Triton module without hiding nested import failures."""
+    try:
+        return importlib.import_module("torch_linear_assignment._triton")
+    except ModuleNotFoundError as error:
+        if error.name in {"torch_linear_assignment._triton", "triton"}:
+            return None
+        raise
+
+
+def _load_legacy_backend() -> ModuleType | None:
+    """Load the optional legacy extension for private benchmark comparisons only."""
+    try:
+        return importlib.import_module("torch_linear_assignment._backend")
+    except ModuleNotFoundError as error:
+        if error.name == "torch_linear_assignment._backend":
+            return None
+        raise
+
+
+def _cuda_uses_triton(cost: torch.Tensor) -> bool:
+    """Return whether this CUDA cost tensor can use the supported Triton path."""
+    if (
+        sys.platform != "linux"
+        or not _torch_supports_triton()
+        or not torch.cuda.is_available()
+        or torch.version.cuda is None
+    ):
+        return False
+    return torch.cuda.get_device_capability(cost.device) >= (8, 0) and _load_triton_backend() is not None
+
+
+def _warn_cuda_fallback() -> None:
+    """Warn once when a CUDA input must use the SciPy CPU fallback."""
+    global _CUDA_FALLBACK_WARNING_EMITTED
+    if _CUDA_FALLBACK_WARNING_EMITTED:
+        return
+    _CUDA_FALLBACK_WARNING_EMITTED = True
+    warnings.warn(
+        "Triton linear-assignment support is unavailable for this CUDA input; using SciPy on CPU.",
+        RuntimeWarning,
+        stacklevel=3,
+    )
+
+
+def _batch_linear_assignment_cuda_legacy(cost: torch.Tensor) -> torch.Tensor:
+    """Run the legacy extension for private benchmark comparison when installed."""
+    backend = _load_legacy_backend()
+    if backend is None:
+        raise RuntimeError("The legacy CUDA extension is not installed.")
+
+    cost = _prepare_solver_cost(cost)
+    _, workers, tasks = cost.shape
+    if tasks < workers:
+        _, row4col = backend.batch_linear_assignment(cost.transpose(1, 2).contiguous())
+        return row4col.long()
+    col4row, _ = backend.batch_linear_assignment(cost.contiguous())
+    return col4row.long()
+
+
+def batch_linear_assignment_cuda(cost: torch.Tensor) -> torch.Tensor:
+    """Solve a CUDA batch through the private Triton implementation."""
+    backend = _load_triton_backend()
+    if backend is None:
+        raise RuntimeError("Triton is unavailable for CUDA linear assignment.")
+    return backend.batch_linear_assignment(cost)
+
+
+def batch_linear_assignment(cost: torch.Tensor) -> torch.Tensor:
+    """Solve a batch of linear assignment problems.
+
+    The method minimizes the cost.
+
+    Args:
+      cost: Cost matrix with shape (B, W, T), where W is the number of workers
+            and T is the number of tasks.
+
+    Returns:
+      Matching tensor with shape (B, W), with assignments for each worker. If the
+      task was not assigned, the corresponding index will be -1.
+    """
+    if cost.ndim != 3:
+        raise ValueError("Need 3-dimensional tensor with shape (B, W, T).")
+    if cost.is_cuda and _cuda_uses_triton(cost):
+        return batch_linear_assignment_cuda(cost)
+
+    device = cost.device
+    if cost.is_cuda:
+        _warn_cuda_fallback()
+        cost = cost.cpu()
+    return batch_linear_assignment_cpu(cost).to(device)
+
+
+def assignment_to_indices(
+    assignment: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert assignment to the SciPy format.
+
+    Args:
+        assignment: The assignment with shape (B, W).
+
+    Returns:
+        row_ind, col_ind: An array of row indices and one of corresponding column indices
+            giving the optimal assignment, each with shape (B, K).
+
+    Raises:
+        ValueError if batch assignments have different sizes.
+    """
+    batch_size = assignment.shape[0]
+    if batch_size == 0:
+        indices = torch.zeros(0, 0, dtype=torch.long, device=assignment.device)
+        return indices, indices
+    mask = assignment >= 0
+    n_matches = mask.sum(1)
+    if (n_matches[1:] != n_matches[0]).any():
+        raise ValueError("Inconsistent matching sizes.")
+    n_matches = n_matches[0].item()
+    row_ind = mask.nonzero()[:, 1].reshape(batch_size, n_matches)
+    col_ind = assignment.masked_select(mask).reshape(batch_size, n_matches)
+    return row_ind, col_ind

torch_linear_assignment-0.1.0.dist-info/METADATA

@@ -0,0 +1,218 @@
+Metadata-Version: 2.4
+Name: torch-linear-assignment
+Version: 0.1.0
+Summary: Batched linear assignment with PyTorch and CUDA.
+Author: Ivan Karpukhin
+Author-email: karpuhini@yandex.ru
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=1.12.0
+Requires-Dist: scipy>=1.6.0
+Requires-Dist: triton; platform_system == "Linux"
+Dynamic: author
+Dynamic: author-email
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Batch linear assignment for PyTorch
+
+[![PyPI version](https://badge.fury.io/py/torch-linear-assignment.svg)](https://badge.fury.io/py/torch-linear-assignment) [![Build Status](https://github.com/ivan-chai/torch-linear-assignment/actions/workflows/ci-tests.yml/badge.svg)](https://github.com/ivan-chai/torch-linear-assignment/actions) [![Downloads](https://img.shields.io/pypi/dm/torch-linear-assignment)](https://pepy.tech/project/torch-linear-assignment) [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+
+<h4 align="left">
+    <p>
+        <a href="#install">Installation</a> |
+        <a href="#backend-and-release-policy">Backend policy</a> |
+        <a href="#example">Usage</a> |
+        <a href="#support">Support</a> |
+        <a href="#legacy-cuda-implementation-00x">Legacy CUDA</a> |
+        <a href="#acknowledgments">Acknowledgments</a> |
+        <a href="#citation">Citation</a>
+    <p>
+</h4>
+Batch computation of the linear assignment problem with PyTorch. The active
+`0.1.0+` release line is pure source: CPU calls use SciPy, and supported CUDA
+calls use the lazy Triton backend.
+
+## Backend and release policy
+
+The active release line is `0.1.0+`. It uses Triton on validated Linux NVIDIA GPUs with compute capability 8.0 or newer and SciPy everywhere else. Unsupported CUDA inputs are solved by the SciPy CPU fallback and returned to the input device. The compiled CUDA extension belongs to the maintenance-only `0.0.x` line and is documented separately in [Legacy CUDA implementation (`0.0.x`)](#legacy-cuda-implementation-00x). If that backend needs a compatibility or correctness fix, the project can issue another `0.0.x` release without adopting the Triton overhaul.
+
+## Install
+
+Python 3.10 or newer is required. Install the `0.1.0` release candidate from PyPI:
+
+```bash
+python -m pip install "torch-linear-assignment==0.1.0rc0"
+```
+
+No editable install, local CUDA compilation, or `--no-build-isolation` flag is needed for normal use. On Linux, the package declares Triton through a platform marker and imports it lazily only for an eligible CUDA input. macOS and Windows installs remain usable through SciPy without requiring Triton.
+
+## Example
+
+```python
+import torch
+from torch_linear_assignment import batch_linear_assignment
+
+cost = torch.tensor([
+    8, 4, 7,
+    5, 2, 3,
+    9, 6, 7,
+    9, 4, 8,
+]).reshape(1, 4, 3)
+
+assignment = batch_linear_assignment(cost)
+print(assignment)
+```
+
+The output is:
+
+```py
+tensor([[ 0,  2, -1,  1]])
+```
+
+To get indices in the SciPy's format:
+
+```py
+from torch_linear_assignment import assignment_to_indices
+
+row_ind, col_ind = assignment_to_indices(assignment)
+print(row_ind)
+print(col_ind)
+```
+
+The output is:
+
+```py
+tensor([[0, 1, 3]])
+tensor([[0, 2, 1]])
+```
+
+## Support
+
+The public API is `batch_linear_assignment(cost)` with cost shape `(B, W, T)`. Assignments have shape `(B, W)`, use `torch.long`, and contain `-1` for an unmatched worker.
+
+| Input and runtime                                                                                                                                         | Backend and behavior                                                               |
+| --------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- |
+| CPU input on any supported OS                                                                                                                             | SciPy on CPU. This remains the reference path.                                     |
+| Linux NVIDIA GPU with compute capability 8.0 or newer, Torch >= 2.4, and Triton available                                                                 | Lazy Triton CUDA backend. The first call can include Triton compilation.           |
+| CUDA input on an unsupported OS; non-NVIDIA Linux accelerator; NVIDIA GPU below compute capability 8.0 (including T4); or missing/ineligible Triton/Torch | Warn once, solve with SciPy on CPU, and return the assignment on the input device. |
+
+The solver promotes `float16`, `bfloat16`, and integer costs to `float32`; `float64` costs remain `float64`. Triton execution disables autocast, so precision already lost before entry cannot be recovered.
+
+Non-finite costs follow the SciPy contract. `NaN` and `-inf` raise `ValueError("matrix contains invalid numeric entries")`. `+inf` represents a forbidden edge when a perfect matching still exists; if the matrix is infeasible, the call raises `ValueError("cost matrix is infeasible")` instead of terminating the process with a device assertion.
+
+## Validation and performance status
+
+GPU validation is author-run rather than hosted CI. The validation targets are Linux NVIDIA GPUs with compute capability 8.0 or newer (for example, L4, A100, H100, or RTX PRO 6000 Blackwell):
+
+```bash
+make validate   # correctness once; prints GPU metadata when available
+make benchmark  # SciPy CPU and installed public CUDA; writes JSONL and a cold/warm table
+```
+
+`make validate` prints Python/platform, Torch/CUDA, Triton, and GPU compute capability metadata when a CUDA device is visible. `make benchmark` separately measures the SciPy CPU reference and installed public CUDA backend, writes complete JSONL evidence, and prints a compact cold/warm table. A clean GPU metadata skip on macOS or a machine without CUDA is expected, but is not GPU evidence.
+
+The first Triton call may be slower because it includes cold compilation; measure cold compilation separately from warm execution. Do not generalize performance across GPU models or workloads. The active `0.1.0+` line requires the correctness, AMP/dtype, fallback, packaging, GPU, and warm-performance gates to pass on each advertised GPU class.
+
+Author-run evidence captured on 2026-08-20 at commit `eee95ffc` (`perf(triton): terminate completed augmentations early`) supports this matrix:
+
+| GPU                                      | SciPy reference comparison      | Triton validation  |
+| ---------------------------------------- | ------------------------------- | ------------------ |
+| NVIDIA L4 (`sm_89`)                      | Exact in all 30 benchmark cases | `78/78` tests pass |
+| NVIDIA A100 (`sm_80`)                    | Exact in all 30 benchmark cases | `78/78` tests pass |
+| NVIDIA H100 (`sm_90`)                    | ?                               | ?                  |
+| NVIDIA RTX PRO 6000 Blackwell (`sm_120`) | Exact in all 30 benchmark cases | `78/78` tests pass |
+
+Representative batch-208 float32 performance is shown below. Each timing cell is `transpose / square / direct` in milliseconds for `300 x 100`, `300 x 300`, and `300 x 600`. Ratios above `1.0x` mean Triton is faster. Values are rounded to one decimal; `~` marks ratios derived from the displayed timings.
+
+| GPU / timing       | SciPy CPU (ms)       | Triton (ms)      | SciPy / Triton          |
+| ------------------ | -------------------- | ---------------- | ----------------------- |
+| L4 (first\*)       | 71.8 / 788.6 / 421.9 | 1.0 / 7.6 / 4.3  | ~71.8x / 103.8x / 98.1x |
+| L4 (warm)          | 71.3 / 782.5 / 422.8 | 1.0 / 7.6 / 4.4  | 71.3x / 103.5x / 96.4x  |
+| A100 (first\*)     | 74.8 / 788.6 / 430.3 | 1.3 / 11.6 / 4.3 | ~57.5x / 68.0x / 100.1x |
+| A100 (warm)        | 73.9 / 788.4 / 431.9 | 1.3 / 11.6 / 4.4 | 56.7x / 67.7x / 99.1x   |
+| H100 (first\*)     | ?                    | ?                | ?                       |
+| H100 (warm)        | ?                    | ?                | ?                       |
+| RTX 6000 (first\*) | 32.1 / 409.8 / 204.8 | 0.7 / 5.5 / 2.2  | ~45.9x / 74.5x / 93.1x  |
+| RTX 6000 (warm)    | 32.3 / 410.1 / 204.8 | 0.5 / 5.4 / 2.2  | 60.8x / 75.8x / 94.4x   |
+
+The Triton implementation parallelizes batch items and vectorizes each current candidate-column scan. Commit `eee95ffc` also terminates the shortest-path search when an augmentation finds an unmatched sink. This removes the post-completion iterations that dominated direct matrices. Against the same-machine SciPy CPU reference, representative warm Triton speedups span `56.7x`--`103.5x` while retaining exact assignment parity.
+
+`first*` means the first call for that case in the shared benchmark process. It is not an isolated compilation measurement: earlier shapes and dtypes can populate driver, process, and Triton disk caches. The three notebooks also predate the fresh-process benchmark-integrity harness, so these measurements are provisional author-run evidence rather than the final cold-performance record. Publication-grade cold claims require a source-bound, five-process rerun.
+
+Do not publish a Triton-only timing as a speedup. Each performance report pairs the same seeded workload with its same-machine SciPy CPU baseline.
+
+For direct benchmark-script options, the underlying command writes complete JSONL evidence to a run-specific file and prints a compact table containing the backend, validation mode, shape, status, cold/warm latency, and exact-parity result:
+
+```bash
+make benchmark
+python tests/benchmark.py --backends scipy,public_cuda --output benchmark-results.jsonl
+```
+
+The command exits nonzero if any invocation fails exact parity. The JSONL file retains replay diagnostics that the console table intentionally omits.
+
+## Legacy CUDA implementation (`0.0.x`)
+
+Release `0.0.6` contains the frozen, compiled plain-CUDA implementation. It is not the active default and should be pinned only when that legacy GPU path is explicitly required, including T4 (`sm_75`) execution:
+
+```bash
+python -m pip install "torch-linear-assignment==0.0.6"
+```
+
+The extension needs a compatible CUDA build environment and must include `sm_75`; the project does not promise a portable T4 wheel. Pinning it is also workload-dependent: in the recorded T4 comparison, SciPy was faster through batch 50, while legacy CUDA was materially faster for the tested batch-208 transpose and direct cases.
+
+For a private contributor comparison from a current checkout, the legacy extension can instead be built explicitly:
+
+```bash
+TLA_BUILD_LEGACY_CUDA=1 python -m pip install -e . --no-build-isolation
+```
+
+This opt-in is not required for normal `0.1.0+` use and does not expose a public backend selector. For a clean cross-version benchmark, copy `tests/benchmark.py` outside the checkout first. In a separate baseline environment, `make install-legacy` installs the newest PyPI release matching `torch-linear-assignment<0.1.0`; the copied runner then measures its installed public CUDA backend. Pair those results with the same shapes, dtypes, seed, GPU, and validation mode from the current Triton run.
+
+The following author-run batch-208 FP32 comparison uses the same `transpose / square / direct` ordering as the main table. Ratios above `1.0x` mean Triton is faster; values are rounded to one decimal, and `~` marks ratios derived from separately displayed legacy and current-package timings.
+
+| GPU / timing       | Legacy 0.0.6 (ms)    | Triton (ms)      | Legacy / Triton          |
+| ------------------ | -------------------- | ---------------- | ------------------------ |
+| T4 (first\*)       | 37.0 / 879.5 / 273.9 | —                | —                        |
+| T4 (warm)          | 37.7 / 738.6 / 272.6 | —                | —                        |
+| L4 (first\*)       | 37.9 / 648.5 / 274.5 | 1.0 / 7.6 / 4.3  | ~37.9x / 85.3x / 63.8x   |
+| L4 (warm)          | 37.8 / 652.5 / 273.3 | 1.0 / 7.6 / 4.4  | ~37.8x / 85.9x / 62.1x   |
+| A100 (first\*)     | 36.2 / 755.5 / 271.7 | 1.3 / 11.6 / 4.3 | ~27.8x / 65.1x / 63.2x   |
+| A100 (warm)        | 36.2 / 731.3 / 271.6 | 1.3 / 11.6 / 4.4 | ~27.8x / 63.0x / 61.7x   |
+| H100 (first\*)     | ?                    | ?                | ?                        |
+| H100 (warm)        | ?                    | ?                | ?                        |
+| RTX 6000 (first\*) | 35.4 / 607.1 / 248.0 | 0.7 / 5.5 / 2.2  | ~50.6x / 110.4x / 112.7x |
+| RTX 6000 (warm)    | 35.4 / 607.6 / 249.3 | 0.5 / 5.4 / 2.2  | ~70.8x / 112.5x / 113.3x |
+
+Triton is faster in all 30 paired warm cases on each measured supported GPU; the weakest displayed ratio is about `6.0x`. T4 has no Triton timing because it is below the supported compute capability, and H100 remains unmeasured. The `first*` caveat from the main table applies equally here.
+
+## Acknowledgments
+
+The `0.1.0+` pure-source Triton backend, GPU validation workflow, and benchmark revamp were implemented by [@Borda](https://github.com/Borda).
+
+## Citation
+
+The code was originally developed for the [HoTPP Benchmark](https://github.com/ivan-chai/hotpp-benchmark). If you use this code in your research project, please cite one of the following papers:
+
+```
+@article{karpukhin2024hotppbenchmark,
+  title={HoTPP Benchmark: Are We Good at the Long Horizon Events Forecasting?},
+  author={Karpukhin, Ivan and Shipilov, Foma and Savchenko, Andrey},
+  journal={arXiv preprint arXiv:2406.14341},
+  year={2024},
+  url ={https://arxiv.org/abs/2406.14341}
+}
+
+@article{karpukhin2024detpp,
+  title={DeTPP: Leveraging Object Detection for Robust Long-Horizon Event Prediction},
+  author={Karpukhin, Ivan and Savchenko, Andrey},
+  journal={arXiv preprint arXiv:2408.13131},
+  year={2024},
+  url ={https://arxiv.org/abs/2408.13131}
+}
+```

torch_linear_assignment-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+torch_linear_assignment/__init__.py,sha256=Dkb0lPgPrrtzMmsKxGpZsVhjKIwtUnoDdT71LArxBYM,188
+torch_linear_assignment/_triton.py,sha256=0EPQpI793Lkq9mtJtcoDqVHjf4m6zCOsx3sp6dWpHlM,13154
+torch_linear_assignment/assignment.py,sha256=MkeY2kGZNW7RPLePBA07RefwT_lWBczszeGXzJiENl0,5816
+torch_linear_assignment-0.1.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+torch_linear_assignment-0.1.0.dist-info/METADATA,sha256=UxH4hrhjFwAhwBV-vnTOVgwS7oT4H7S7RFJf_oMKnuo,13960
+torch_linear_assignment-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+torch_linear_assignment-0.1.0.dist-info/top_level.txt,sha256=eV_iv2yNOwDzuo6EUQVtdx-8Ar-HX03wRmnl6Ge1vGA,24
+torch_linear_assignment-0.1.0.dist-info/RECORD,,

torch_linear_assignment-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

torch_linear_assignment-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

torch_linear_assignment-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+torch_linear_assignment