mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Arup DeandCopilot
90f807a321
# feat(backends): multi-DSL dispatcher foundation + Triton backend layer **Head:** `arde171:arde/liger-A-slim` → **Base:** `linkedin/Liger-Kernel:main` — 61 files, PR **1/5** of the multi-backend series. ## Summary Introduces a **per-(op, impl) backend dispatcher** so a single Liger op can have multiple kernel implementations (Triton, CuTeDSL, cuTile) and pick one by **preference-rank + capability**, with **Triton as the universal fallback**. This PR ships *only* the foundation + the Triton backend layer; on its own it routes every op through Triton exactly as before (no behavior change). The CuTeDSL and cuTile backends register into this in follow-up PRs 2–5. ## What's added - `src/liger_kernel/backends/` - `dispatch.py` — selection: explicit `impl=` → env (`LIGER_KERNEL_IMPL[_OP]`) → global `set_impl` → capability-ranked auto-pick → Triton fallback. Raises a precise `ImplNotAvailableError` when a requested impl isn't registered. - `registry.py` — `register_impl`/`declare_op_locations`; lazy per-op discovery (imports a backend module only on first use of its op). - `capability.py` — declarative `Capability(min_cc, max_cc, modules, requires_cuda)`; resolves compute capability from the **input tensor's device** (`torch.cuda.get_device_capability(device)`), threaded through `dispatch`→`resolve_impl`→`available_impls`→`is_satisfied`, so selection is correct even for heterogeneous multi-GPU-per-process. Memoized per `(op, device, device-cc, cuda-visible)`. - `ops/backends/_triton/*` — the existing Triton kernels as first-class registered impls. - `functional.py` — `declare_op_locations` for every op (this PR: **Triton path only**; later PRs append their own `_cutedsl`/`_cutile` paths — additive, so no PR references a backend it doesn't ship). - `ops/_nvidia_shared.py`, `testing/_op_helpers.py` (shared correctness/tolerance harness), `test/backends/*`. - A **test-isolation fixture** (`test/conftest.py`) that resets the global impl selection after each test, so backend-selecting tests don't leak into later tests. ## Design — preference-rank + Triton fallback Each impl declares a rank and a `Capability`. Auto-selection picks the lowest-rank impl whose capability holds on the input tensor's device; if none (or the module can't import), it falls back to Triton. Discovery is lazy and idempotent. ## Correctness & testing (B300 sm_103 + H200 sm_90, torch 2.13/cu130) - `import liger_kernel` — **warning-clean** (declares only shipped Triton paths). - Full backend suite `test/backends` + `test/ops`, deterministic (`-p no:randomly`): **B300 348 passed / 0 failed**, **H200 348 passed / 0 failed** (44 skipped). - ruff check + format: clean across the repo. ## Scope Foundation only. No CuTeDSL/cuTile code, no benchmark data, no dev-scaffold, no doc essays (those are in later PRs / stay internal). Behavior is Triton-identical. Co-authored-by: Copilot <[email protected]> --- **Multi-backend series — review/merge in order:** #1416 (A · dispatcher foundation) → #1417 (B · CuTeDSL norms) → #1418 (C · CuTeDSL elementwise) → #1419 (D · CuTeDSL CE) → #1420 (E · cuTile). Each PR is stacked on the previous; the diff auto-narrows as earlier PRs land. GPU tests are cron-only (per #1409) so per-PR CI = checkstyle + CPU tests; multi-backend suites verified out-of-band on **B300 (sm_103)** + **H200 (sm_90)**. Co-authored-by: Copilot <[email protected]>
44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from liger_kernel.utils import is_npu_available
|
|
from test.utils import set_seed
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def set_random_seed():
|
|
set_seed(42)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_gpu_cache():
|
|
yield
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
elif is_npu_available():
|
|
torch.npu.empty_cache()
|
|
elif torch.xpu.is_available():
|
|
torch.xpu.empty_cache()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_liger_backend_selection():
|
|
"""Isolate global backend/impl selection between tests.
|
|
|
|
Several op tests call ``set_impl``/``set_backend`` (or set the
|
|
``LIGER_KERNEL_IMPL*`` env) to exercise a specific backend. Without a reset,
|
|
that global choice leaks into every later test in the session, which routed
|
|
unrelated ops (cross_entropy, tvd, dyt, ...) through a backend they were
|
|
never meant to use and produced order-dependent failures. Restore the
|
|
default (auto) selection and drop the availability memo after each test.
|
|
"""
|
|
yield
|
|
try:
|
|
from liger_kernel.backends.dispatch import clear_available_cache
|
|
from liger_kernel.backends.dispatch import set_impl
|
|
|
|
set_impl(None)
|
|
clear_available_cache()
|
|
except Exception:
|
|
pass
|