chore: import upstream snapshot with attribution

python/sglang/jit_kernel/lplb/torch_solver.py

@@ -0,0 +1,192 @@
+"""IPM LP Solver entry point — dispatches to the fused JIT CUDA kernel.
+
+Solves: min c^T x  subject to  Ax = b, x >= 0
+using a barrier (interior point) method with 5 iterations.
+
+The fused kernel lives in ``cuda_solver`` (CUDA C++ via ``load_jit``,
+backed by header-only cuBLASDx + a hand-written block Cholesky). This
+module is the public-facing import surface for callers (``LPLBSolver``)
+and resolves/caches the backend on first use.
+
+LPLB requires Hopper-class hardware and Math-DX cuBLASDx headers. If
+either is missing, ``warmup`` and ``solve_ipm`` raise — there is no
+silent fallback.
+"""
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+# Backend dispatch state (resolved on first call, cached afterwards)
+_BACKEND_CHECKED = False
+_FUSED_AVAILABLE = False
+_FUSED_SOLVE_IPM = None  # type: ignore[assignment]
+_FUSED_WARMUP = None  # type: ignore[assignment]
+_FUSED_ASSERT_FITS = None  # type: ignore[assignment]
+
+
+def _init_fused_backend() -> None:
+    """Resolve the fused backend once. Records WHY it's disabled when it is."""
+    global _BACKEND_CHECKED, _FUSED_AVAILABLE
+    global _FUSED_SOLVE_IPM, _FUSED_WARMUP, _FUSED_ASSERT_FITS
+
+    if _BACKEND_CHECKED:
+        return
+    _BACKEND_CHECKED = True
+
+    if not torch.cuda.is_available():
+        logger.info("LPLB fused solver disabled: CUDA not available")
+        return
+
+    cap = torch.cuda.get_device_capability()
+    if cap[0] < 9:
+        logger.info(
+            f"LPLB fused solver disabled: GPU SM {cap[0]}.{cap[1]} < 9.0 "
+            "(requires Hopper or newer)"
+        )
+        return
+
+    try:
+        from sglang.jit_kernel.lplb.cuda_solver import solve_ipm as fused_solve_ipm
+        from sglang.jit_kernel.lplb.cuda_solver import warmup as fused_warmup
+        from sglang.jit_kernel.lplb.shmem_budget import assert_fits
+    except ImportError as e:
+        logger.info(
+            f"LPLB fused solver disabled: {e}. "
+            "Install Math-DX cuBLASDx via `pip install nvidia-mathdx` "
+            "or set MATHDX_HOME to an extracted archive."
+        )
+        return
+
+    _FUSED_SOLVE_IPM = fused_solve_ipm
+    _FUSED_WARMUP = fused_warmup
+    _FUSED_ASSERT_FITS = assert_fits
+    _FUSED_AVAILABLE = True
+    logger.info("LPLB fused solver enabled (CUDA C++ via load_jit, cuBLASDx)")
+
+
+def _unavailable_reason() -> str:
+    if not torch.cuda.is_available():
+        return "CUDA is not available"
+    cap = torch.cuda.get_device_capability()
+    if cap[0] < 9:
+        return f"GPU SM {cap[0]}.{cap[1]} < 9.0 (requires Hopper or newer)"
+    return (
+        "Math-DX cuBLASDx headers not found — install via "
+        "`pip install nvidia-mathdx` or set MATHDX_HOME"
+    )
+
+
+def warmup(nc: int, nv: int, num_iters: int = 5, device: str = "cuda") -> None:
+    """Pre-JIT-compile the fused kernel for a given (NC, NV) shape.
+
+    Call once per unique shape at solver construction time to hide the
+    20-40s JIT compilation cost. Raises if the fused backend is
+    unavailable, the shape exceeds the shmem budget, or the kernel
+    fails to compile/launch.
+    """
+    _init_fused_backend()
+    if not _FUSED_AVAILABLE:
+        raise RuntimeError(f"LPLB fused solver unavailable: {_unavailable_reason()}")
+    _FUSED_ASSERT_FITS(nc, nv, gpu="h100")
+    _FUSED_WARMUP(nc, nv, num_iters=num_iters, device=device)
+
+
+def solve_ipm(
+    A: torch.Tensor,
+    b: torch.Tensor,
+    c: torch.Tensor,
+    num_iters: int = 5,
+) -> torch.Tensor:
+    """Barrier-method Interior Point solver for standard-form LP.
+
+    Dispatches to the JIT-compiled CUDA C++ kernel (Hopper+ GPU with
+    Math-DX cuBLASDx headers, reachable via ``nvidia-mathdx`` PyPI
+    package or ``MATHDX_HOME``). Raises if the fused backend is
+    unavailable or the inputs aren't on CUDA in float32.
+
+    Args:
+        A: Constraint matrix, shape (NC, NV), float32, on CUDA.
+        b: RHS vector, shape (NC,), float32, on CUDA.
+        c: Objective coefficients, shape (NV,), float32, on CUDA.
+        num_iters: Number of barrier iterations (default 5).
+
+    Returns:
+        x: Solution vector, shape (NV,), float32. The kernel writes 0.5
+        for every entry on non-convergence.
+    """
+    nc, nv = A.shape
+    assert b.shape == (nc,), f"b shape mismatch: {b.shape} vs ({nc},)"
+    assert c.shape == (nv,), f"c shape mismatch: {c.shape} vs ({nv},)"
+
+    _init_fused_backend()
+    if not _FUSED_AVAILABLE:
+        raise RuntimeError(f"LPLB fused solver unavailable: {_unavailable_reason()}")
+    if not A.is_cuda:
+        raise RuntimeError(
+            f"LPLB fused solver requires CUDA tensors; got A on {A.device}."
+        )
+    if A.dtype != torch.float32:
+        raise RuntimeError(
+            f"LPLB fused solver requires float32; got A.dtype={A.dtype}."
+        )
+    return _FUSED_SOLVE_IPM(A, b, c, num_iters=num_iters)
+
+
+def solve_ipm_torch_reference(
+    A: torch.Tensor,
+    b: torch.Tensor,
+    c: torch.Tensor,
+    num_iters: int = 5,
+) -> torch.Tensor:
+    """Pure-torch reference for the fused IPM kernel — testing only.
+
+    Mirrors the barrier-method iteration in ``csrc/lplb/ipm.cuh``
+    step-for-step so the two can be compared numerically:
+
+      x <- 1
+      for _ in range(num_iters):
+          ax2  = A * x^2            # (NC, NV)
+          ax2a = ax2 @ A^T          # (NC, NC) KKT matrix
+          delta = solve(ax2a, ax2 @ c)
+          r     = delta^T @ A       # (NV,)
+          d     = x * (c - r)
+          alpha = 0.999 / d_max  (or 1.0 if d_max <= 1e-9)
+          x    *= 1 - alpha * d
+      write 0.5 everywhere on non-convergence.
+
+    NOT bit-equivalent to the kernel: the kernel factors the KKT system
+    with a hand-written block Cholesky while this uses
+    ``torch.linalg.solve`` (LU). The two agree to a small tolerance
+    (the numerical difference being the whole point of the comparison
+    test). This function is never on the production path — the fused
+    kernel is the only LP solver at runtime.
+    """
+    nc, nv = A.shape
+    assert b.shape == (nc,), f"b shape mismatch: {b.shape} vs ({nc},)"
+    assert c.shape == (nv,), f"c shape mismatch: {c.shape} vs ({nv},)"
+
+    x = torch.ones(nv, device=A.device, dtype=torch.float32)
+    d_max = torch.tensor(0.0, device=A.device, dtype=torch.float32)
+    for _ in range(num_iters):
+        ax2 = A * (x * x).unsqueeze(0)  # (NC, NV)
+        ax2a = ax2 @ A.t()  # (NC, NC)
+        ax2c = ax2 @ c  # (NC,)
+        # Match the kernel's 1e-12 pivot clamp via a tiny diagonal jitter so
+        # a (near-)singular KKT system stays solvable instead of raising.
+        ax2a = ax2a + 1e-12 * torch.eye(nc, device=A.device, dtype=torch.float32)
+        delta = torch.linalg.solve(ax2a, ax2c)  # (NC,)
+        r = A.t() @ delta  # (NV,)
+        d = x * (c - r)  # (NV,)
+        d_max = d.max()
+        alpha = 0.999 / d_max if d_max > 1e-9 else torch.tensor(1.0, device=A.device)
+        x = x * (1.0 - alpha * d)
+
+    max_residual = (A @ x - b).abs().max()
+    converged = (d_max < 0.1) and (0 <= x[-1] < 1e-4) and (max_residual < 0.05)
+    if not converged:
+        return torch.full((nv,), 0.5, device=A.device, dtype=torch.float32)
+    return x