sgl-project--sglang/python/sglang/jit_kernel/lplb/torch_solver.py

"""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