chore: import upstream snapshot with attribution

python/sglang/jit_kernel/lplb/__init__.py

@@ -0,0 +1 @@
+"""LPLB (LP-based Load Balancer) JIT kernels for expert parallelism."""

python/sglang/jit_kernel/lplb/cublasdx_solver.py

@@ -0,0 +1,14 @@
+"""Backwards-compatible shim.
+
+The Numba/nvmath-python fused IPM that used to live here has been replaced
+by a CUDA C++ kernel JIT-compiled via sglang's ``load_jit`` infrastructure.
+The new implementation lives in ``cuda_solver``. This module re-exports the
+public API so any external import keeps working.
+"""
+
+from sglang.jit_kernel.lplb.cuda_solver import (  # noqa: F401
+    solve_ipm,
+    warmup,
+)
+
+__all__ = ["solve_ipm", "warmup"]

python/sglang/jit_kernel/lplb/cuda_solver.py

@@ -0,0 +1,324 @@
+"""JIT-compiled CUDA Interior Point Method LP solver.
+
+Replaces the Numba/nvmath-python implementation in ``cublasdx_solver.py``.
+The kernel is a single-block fused IPM defined in
+``csrc/lplb/ipm.cuh`` and compiled per ``(NC, NV, BLOCK_DIM, SM_VER,
+NUM_ITERS)`` tuple via sglang's ``tvm-ffi`` ``load_jit``.
+
+Per-call CPU overhead is dominated by the pybind11 dispatch + four
+``data_ptr()`` calls (~5–10 µs total), versus ~500–700 µs for the prior
+Numba path (numba dispatcher chain + ``as_cuda_array`` per tensor).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    get_jit_cuda_arch,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_BLOCK_DIM = 256
+# Per-element kernels (post-LP dispatch) saturate easily — also 256.
+DISPATCH_BLOCK_DIM = 256
+DEFAULT_NUM_ITERS = 5
+
+
+def _sm_ver() -> int:
+    arch = get_jit_cuda_arch()
+    return arch.major * 100 + arch.minor * 10
+
+
+@cache_once
+def _ipm_module(
+    nc: int, nv: int, block_dim: int, num_iters: int, sm_ver: int
+) -> Module:
+    """JIT-compile the IPM kernel for one shape. Cached for the process lifetime."""
+    args = make_cpp_args(nc, nv, block_dim, sm_ver, num_iters)
+    # The kernel uses cuBLASDx (header-only) for the GEMMs and a hand-written
+    # block-level Cholesky for the POSV. No -rdc=true / static-lib linkage
+    # required, so sglang's tvm-ffi load_jit handles the build with the
+    # default flags.
+    return load_jit(
+        "lplb_ipm",
+        *args,
+        cuda_files=["lplb/ipm.cuh"],
+        cuda_wrappers=[("ipm_solve", f"ipm_solve<{args}>")],
+        extra_dependencies=["mathdx"],
+    )
+
+
+def warmup(
+    nc: int,
+    nv: int,
+    num_iters: int = DEFAULT_NUM_ITERS,
+    device: str = "cuda",
+) -> None:
+    """JIT-compile the kernel for ``(nc, nv)`` so the first real solve isn't
+    paying the compile cost. Raises on compile or launch failure.
+    """
+    module = _ipm_module(nc, nv, DEFAULT_BLOCK_DIM, num_iters, _sm_ver())
+    # Trigger any first-call lazy initialization.
+    A = torch.zeros(nc, nv, dtype=torch.float32, device=device)
+    b = torch.zeros(nc, dtype=torch.float32, device=device)
+    c = torch.zeros(nv, dtype=torch.float32, device=device)
+    result = torch.empty(nv, dtype=torch.float32, device=device)
+    module.ipm_solve(A, b, c, result)
+    logger.info(f"LPLB CUDA IPM solver: warmed up for (NC={nc}, NV={nv})")
+
+
+def solve_ipm(
+    A: torch.Tensor,
+    b: torch.Tensor,
+    c: torch.Tensor,
+    num_iters: int = DEFAULT_NUM_ITERS,
+    result: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run the fused single-SM IPM kernel.
+
+    cuBLASDx GEMMs + hand-written block Cholesky, dispatched per the
+    module docstring.
+
+    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).
+        result: Optional pre-allocated ``(NV,)`` float32 CUDA buffer to write
+            into. When omitted the kernel allocates a fresh result tensor
+            (~20 µs of CPU overhead). Passing in a long-lived buffer skips
+            that alloc on every solve.
+
+    Returns:
+        x: Solution vector, shape ``(NV,)``, float32. The kernel writes 0.5
+        for every entry on non-convergence (matches the prior Numba behavior).
+    """
+    assert A.is_cuda and b.is_cuda and c.is_cuda
+    assert A.dtype == torch.float32
+    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},)"
+
+    module = _ipm_module(nc, nv, DEFAULT_BLOCK_DIM, num_iters, _sm_ver())
+    if result is None:
+        result = torch.empty(nv, dtype=torch.float32, device=A.device)
+    module.ipm_solve(A, b, c, result)
+    return result
+
+
+@cache_once
+def _prep_module(
+    nc: int,
+    nv: int,
+    num_single: int,
+    num_red_log: int,
+    num_gpus: int,
+    block_dim: int,
+) -> Module:
+    args = make_cpp_args(nc, nv, num_single, num_red_log, num_gpus, block_dim)
+    return load_jit(
+        "lplb_lp_prep",
+        *args,
+        cuda_files=["lplb/lp_prep.cuh"],
+        cuda_wrappers=[("lp_prep", f"lp_prep<{args}>")],
+    )
+
+
+def prep_lp_inputs(
+    A_full: torch.Tensor,
+    b: torch.Tensor,
+    t1: torch.Tensor,
+    global_counts: torch.Tensor,
+    log_single: torch.Tensor,
+    log_replicated: torch.Tensor,
+    B1: torch.Tensor,
+    A_base_row_sum: torch.Tensor,
+) -> None:
+    """Replace the 8 torch ops that built the IPM inputs with one CUDA kernel.
+
+    Writes into the caller-provided ``A_full`` (last column only), ``b``,
+    and ``t1`` buffers. ``A_full``'s first ``NV-1`` columns must already
+    hold ``A_base.copy_()`` from solver init — this kernel does not touch
+    them.
+    """
+    nc, nv = A_full.shape
+    num_single = log_single.shape[0]
+    num_red_log = log_replicated.shape[0]
+    num_gpus, _ns = B1.shape
+    module = _prep_module(nc, nv, num_single, num_red_log, num_gpus, DEFAULT_BLOCK_DIM)
+    module.lp_prep(
+        A_full, b, t1, global_counts, log_single, log_replicated, B1, A_base_row_sum
+    )
+
+
+@cache_once
+def _post_module(
+    num_logical: int,
+    max_copies: int,
+    num_single: int,
+    num_red_phy: int,
+    block_dim: int,
+) -> Module:
+    args = make_cpp_args(num_logical, max_copies, num_single, num_red_phy, block_dim)
+    return load_jit(
+        "lplb_lp_post",
+        *args,
+        cuda_files=["lplb/lp_post.cuh"],
+        cuda_wrappers=[("lp_post", f"lp_post<{args}>")],
+    )
+
+
+def extract_log2phy_prob(
+    log2phy_prob: torch.Tensor,
+    x: torch.Tensor,
+    t1: torch.Tensor,
+    phy_single: torch.Tensor,
+    phy_replicated: torch.Tensor,
+    log2phy: torch.Tensor,
+) -> None:
+    """Replace the 5 torch ops that turned the IPM output into log2phy_prob
+    with one CUDA kernel. Writes into the caller-provided ``log2phy_prob``
+    buffer of shape ``(num_logical, max_copies)``.
+    """
+    num_logical, max_copies = log2phy_prob.shape
+    num_single = phy_single.shape[0]
+    num_red_phy = phy_replicated.shape[0]
+    module = _post_module(
+        num_logical, max_copies, num_single, num_red_phy, DEFAULT_BLOCK_DIM
+    )
+    module.lp_post(log2phy_prob, x, t1, phy_single, phy_replicated, log2phy)
+
+
+@cache_once
+def _dispatch_module(max_copies: int, block_dim: int) -> Module:
+    args = make_cpp_args(max_copies, block_dim)
+    return load_jit(
+        "lplb_dispatch_probability",
+        *args,
+        cuda_files=["lplb/dispatch_probability.cuh"],
+        cuda_wrappers=[("dispatch_probability", f"dispatch_probability<{args}>")],
+    )
+
+
+def dispatch_probability(
+    topk_ids: torch.Tensor,
+    log2phy_prob: torch.Tensor,
+    log2phy_map: torch.Tensor,
+    random_vals: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Replace the 7 torch ops in `_topk_ids_logical_to_physical_probability`
+    with a single per-token-per-slot CUDA kernel.
+
+    Samples a physical expert per (token, slot) via inverse-CDF on the
+    per-row LP probabilities. Bit-equivalent to
+    :func:`dispatch_probability_torch_reference` when given the same
+    ``random_vals`` (modulo float rounding in the cumulative sum).
+
+    Args:
+        topk_ids: (num_tokens, topk) int32, on CUDA. Logical expert ids from
+            the router.
+        log2phy_prob: (num_logical, max_copies) float32. LP solver output.
+        log2phy_map: (num_logical, max_copies) int32. -1 entries are
+            unused replicas; treated as 0-weight in the multinomial.
+        random_vals: Optional (N,) float32 CUDA tensor of uniform samples in
+            [0, 1). When omitted, the function generates fresh values via
+            ``torch.rand``. Pass explicitly when comparing against the torch
+            reference for deterministic equivalence.
+
+    Returns:
+        Physical topk ids tensor with the same shape as ``topk_ids``.
+    """
+    original_shape = topk_ids.shape
+    flat_ids = topk_ids.reshape(-1).contiguous().to(torch.int32)
+    n = flat_ids.shape[0]
+    num_logical, max_copies = log2phy_prob.shape
+    assert log2phy_map.shape == (num_logical, max_copies)
+    map32 = log2phy_map.contiguous().to(torch.int32)
+
+    out = torch.empty(n, dtype=torch.int32, device=topk_ids.device)
+    if random_vals is None:
+        random_vals = torch.rand(n, dtype=torch.float32, device=topk_ids.device)
+    else:
+        assert random_vals.shape == (
+            n,
+        ), f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}"
+    module = _dispatch_module(max_copies, DISPATCH_BLOCK_DIM)
+    module.dispatch_probability(out, flat_ids, log2phy_prob, map32, random_vals)
+    return out.view(original_shape).to(topk_ids.dtype)
+
+
+def dispatch_probability_torch_reference(
+    topk_ids: torch.Tensor,
+    log2phy_prob: torch.Tensor,
+    log2phy_map: torch.Tensor,
+    random_vals: torch.Tensor,
+) -> torch.Tensor:
+    """Pure-torch reference of :func:`dispatch_probability`.
+
+    Mirrors the CUDA kernel's algorithm exactly (inverse-CDF via cumsum
+    + threshold) so the two paths are bit-equivalent for identical
+    ``random_vals``, modulo floating-point rounding in the cumsum. Kept
+    for numerical comparison and testing — not on the production hot
+    path (it allocates and runs ~8 torch ops; the fused kernel collapses
+    them into one launch).
+
+    Algorithm (matches ``csrc/lplb/dispatch_probability.cuh``):
+
+      1. Gather the per-row probability vector and physical-id map for
+         each logical id in ``topk_ids``.
+      2. If the row sum is zero (LP gave no signal), fall back to
+         uniform over valid replicas (``log2phy_map != -1``).
+      3. Sample: smallest ``c`` such that ``cumsum[0..c] > u * row_sum``,
+         where ``u = random_vals[i]``. Ties favor advancing ``c``,
+         matching the CUDA kernel.
+      4. Return ``log2phy_map[logical_id, c]``.
+
+    Args:
+        topk_ids: (num_tokens, topk) int, on CUDA or CPU. Logical expert ids.
+        log2phy_prob: (num_logical, max_copies) float32. LP solver output.
+        log2phy_map: (num_logical, max_copies) int. -1 = unused replica.
+        random_vals: (N,) float32, where N = ``topk_ids.numel()``. Uniform
+            samples in [0, 1) — same shape and semantics as the CUDA kernel.
+
+    Returns:
+        Physical topk ids tensor with the same shape and dtype as ``topk_ids``.
+    """
+    original_shape = topk_ids.shape
+    flat_ids = topk_ids.reshape(-1).long()
+    n = flat_ids.shape[0]
+    num_logical, max_copies = log2phy_prob.shape
+    assert log2phy_map.shape == (num_logical, max_copies)
+    assert random_vals.shape == (
+        n,
+    ), f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}"
+
+    # Gather per-row probabilities and physical maps.
+    probs = log2phy_prob[flat_ids]  # (N, max_copies), float32
+    maps = log2phy_map[flat_ids]  # (N, max_copies), same int dtype as input
+
+    # Fallback when row_sum == 0: uniform over valid replicas.
+    row_sum = probs.sum(dim=-1, keepdim=True)  # (N, 1)
+    fallback_probs = (maps >= 0).to(probs.dtype)  # (N, max_copies)
+    probs = torch.where(row_sum > 0, probs, fallback_probs)
+    row_sum = probs.sum(dim=-1)  # (N,)
+
+    # Inverse-CDF sample: smallest c such that cumsum[..c] > u.
+    # ``(cum <= u).sum(dim=-1)`` counts how many slots are still below u,
+    # which equals the CUDA kernel's ``chosen`` after its for-loop.
+    u = (random_vals * row_sum).unsqueeze(-1)  # (N, 1)
+    cum = probs.cumsum(dim=-1)  # (N, max_copies)
+    chosen = (cum <= u).sum(dim=-1).clamp(max=max_copies - 1)  # (N,)
+
+    out = maps.gather(1, chosen.unsqueeze(-1)).squeeze(-1)
+    return out.view(original_shape).to(topk_ids.dtype)

python/sglang/jit_kernel/lplb/shmem_budget.py

@@ -0,0 +1,152 @@
+"""Shared-memory budget accounting for the fused IPM kernel.
+
+Fused layout (fp32), one block per LP, all state in shared memory::
+
+    A        NC * NV       constraint matrix     (resident)
+    c        NV            cost vector           (resident)
+    x        NV            IPM state             (resident)
+    ata      NC * NC       KKT matrix / Cholesky factor
+    rhs      NC            ax2c, then delta
+    d        NV            aliased with r = A.T @ delta
+
+    S_elems = NC*NV + NC*NC + 3*NV + NC
+
+Dynamic shared-memory cap per block (with opt-in via
+``cudaFuncAttributeMaxDynamicSharedMemorySize``):
+
+    A100   SM_80   164 KB   practical 160 KB
+    H100   SM_90   227 KB   practical 223 KB  <- default target
+    H200   SM_90   227 KB
+    H20    SM_90   227 KB
+    B200   SM_100  228 KB   practical 224 KB
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Per-block slack reserved for cuBLASDx workspace and CUDA runtime state.
+_RUNTIME_PAD_BYTES = 256
+
+# fp32
+_BYTES_PER_ELEM = 4
+
+# Practical per-block dynamic shmem caps (bytes)
+GPU_BUDGETS_BYTES: dict[str, int] = {
+    "a100": 160 * 1024,
+    "h100": 223 * 1024,
+    "h200": 223 * 1024,
+    "h20": 223 * 1024,
+    "b200": 224 * 1024,
+}
+
+
+@dataclass(frozen=True)
+class ShmemBreakdown:
+    nc: int
+    nv: int
+    a_bytes: int
+    c_bytes: int
+    x_bytes: int
+    ata_bytes: int
+    rhs_bytes: int
+    d_bytes: int
+    pad_bytes: int
+
+    @property
+    def total_bytes(self) -> int:
+        return (
+            self.a_bytes
+            + self.c_bytes
+            + self.x_bytes
+            + self.ata_bytes
+            + self.rhs_bytes
+            + self.d_bytes
+            + self.pad_bytes
+        )
+
+    def as_kib(self) -> float:
+        return self.total_bytes / 1024.0
+
+
+def shmem_bytes(nc: int, nv: int, bytes_per_elem: int = _BYTES_PER_ELEM) -> int:
+    """Exact byte count for the fused layout with the given (NC, NV)."""
+    return bytes_per_elem * (nc * nv + nc * nc + 3 * nv + nc) + _RUNTIME_PAD_BYTES
+
+
+def breakdown(
+    nc: int, nv: int, bytes_per_elem: int = _BYTES_PER_ELEM
+) -> ShmemBreakdown:
+    """Per-array byte breakdown — useful for debugging shmem pressure."""
+    b = bytes_per_elem
+    return ShmemBreakdown(
+        nc=nc,
+        nv=nv,
+        a_bytes=b * nc * nv,
+        c_bytes=b * nv,
+        x_bytes=b * nv,
+        ata_bytes=b * nc * nc,
+        rhs_bytes=b * nc,
+        d_bytes=b * nv,
+        pad_bytes=_RUNTIME_PAD_BYTES,
+    )
+
+
+def gpu_budget_bytes(gpu: str) -> int:
+    key = gpu.lower()
+    if key not in GPU_BUDGETS_BYTES:
+        raise ValueError(
+            f"unknown gpu '{gpu}', expected one of {sorted(GPU_BUDGETS_BYTES)}"
+        )
+    return GPU_BUDGETS_BYTES[key]
+
+
+def fits(nc: int, nv: int, gpu: str = "h100") -> bool:
+    return shmem_bytes(nc, nv) <= gpu_budget_bytes(gpu)
+
+
+def assert_fits(nc: int, nv: int, gpu: str = "h100") -> None:
+    """Raise if the fused kernel will not fit on the target GPU."""
+    used = shmem_bytes(nc, nv)
+    cap = gpu_budget_bytes(gpu)
+    if used > cap:
+        raise ValueError(
+            f"fused IPM kernel needs {used/1024:.1f} KiB of shared memory for "
+            f"NC={nc}, NV={nv}, but {gpu} allows {cap/1024:.1f} KiB/block. "
+            f"Either reduce problem size or switch to a tiled design."
+        )
+
+
+def max_nc_for_nv(nv: int, gpu: str = "h100") -> int:
+    """Largest NC that fits for a given NV. Solves
+        4 * (NC^2 + (NV+1)*NC + 3*NV) + pad <= cap
+    via the quadratic formula (monotone in NC). Returns 0 if even NC=1 overflows.
+    """
+    cap = gpu_budget_bytes(gpu)
+    b = _BYTES_PER_ELEM
+    # cap - pad >= b * (NC^2 + (NV+1)*NC + 3*NV)
+    rhs = (cap - _RUNTIME_PAD_BYTES) / b - 3 * nv
+    if rhs <= 0:
+        return 0
+    # NC^2 + (NV+1)*NC - rhs <= 0
+    import math
+
+    disc = (nv + 1) ** 2 + 4 * rhs
+    nc_max = int((-(nv + 1) + math.sqrt(disc)) / 2.0)
+    while nc_max > 0 and shmem_bytes(nc_max, nv) > cap:
+        nc_max -= 1
+    return max(nc_max, 0)
+
+
+def report(nc: int, nv: int, gpu: str = "h100") -> str:
+    """Human-readable summary — used by kernels on init for logging."""
+    bd = breakdown(nc, nv)
+    cap = gpu_budget_bytes(gpu)
+    status = "FITS" if bd.total_bytes <= cap else "OVER BUDGET"
+    return (
+        f"[shmem] NC={nc} NV={nv} gpu={gpu} | "
+        f"A={bd.a_bytes/1024:.1f}K "
+        f"ata={bd.ata_bytes/1024:.1f}K "
+        f"rest={(bd.c_bytes+bd.x_bytes+bd.rhs_bytes+bd.d_bytes)/1024:.1f}K | "
+        f"total={bd.total_bytes/1024:.1f}K / {cap/1024:.1f}K  {status}"
+    )

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