chore: import upstream snapshot with attribution

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)