linkedin / linkedin/Liger-Kernel

Triton kernels can launch on the wrong GPU with multi-GPU inputs

Open
#1,303 1 comment 2 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

Some Liger Triton kernels assume the active CUDA device matches the input tensor device. In multi-GPU flows like `device_map="auto"`, that is not always true: a layer can receive tensors on `cuda:1` while the process active device is still `cuda:0`.

That can make Triton compile or launch against the wrong device. I first noticed it around RMSNorm, but the same underlying issue is visible with SwiGLU and GeGLU too.

Related upstream fix pattern: https://github.com/huggingface/kernels-community/pull/1029

Repro using `liger-kernel` directly:

```python
#!/usr/bin/env python3
import argparse

import torch
import torch.nn.functional as F

from liger_kernel.ops.geglu import LigerGELUMulFunction
from liger_kernel.ops.rms_norm import LigerRMSNormFunction
from liger_kernel.ops.swiglu import LigerSiLUMulFunction

EPS = 1e-6
TOL = 0.1
BATCH = 8
HIDDEN = 4096

def ref_rmsnorm(x, weight, eps):
xf = x.float()
return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * weight.float()).to(x.dtype)

def rmsnorm_case(dev):
x = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
weight = torch.randn(HIDDEN, device=dev, dtype=torch.bfloat16) * 0.1
return (x, weight, EPS, 0, "llama", True, None), ref_rmsnorm(x, weight, EPS)

def silu_case(dev):
a = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
b = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
return (a, b), F.silu(a) * b

def gelu_case(dev):
a = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
b = torch.randn(BATCH, HIDDEN, device=dev, dtype=torch.bfloat16)
return (a, b), F.gelu(a, approximate="tanh") * b

CASES = [
("LigerRMSNorm", LigerRMSNormFunction, rmsnorm_case),
("LigerSiLUMul", LigerSiLUMulFunction, silu_case),
("LigerGELUMul", LigerGELUMulFunction, gelu_case),
]

def run_case(name, fn, build, devices):
print(f"\n===== {name} =====", flush=True)

# Launch on device 0 first, then keep device 0 active while passing
# device 1 tensors to reproduce the mismatch.
for dev in devices:
torch.manual_seed(0)
args, ref = build(dev)
torch.cuda.set_device(devices[0])

try:
with torch.no_grad():
out = fn.apply(*args)
finite = bool(torch.isfinite(out).all().item())
max_abs = float((out.float() - ref.float()).abs().max().item())
ok = finite and max_abs < TOL
print(
f" input={dev} active=cuda:{torch.cuda.current_device()} "
f"finite={finite} max|diff|={max_abs:.6f} -> {'OK' if ok else 'CORRUPT'}",
flush=True,
)
except Exception as exc:
print(
f" input={dev} active=cuda:{torch.cuda.current_device()} "
f"CRASH {type(exc).__name__}: {str(exc).splitlines()[0][:120]}",
flush=True,
)

def main():
parser = argparse.ArgumentParser()
parser.add_argument("--devices", nargs=2, default=["cuda:0", "cuda:1"])
args = parser.parse_args()

if torch.cuda.device_count() < 2:
raise SystemExit(f"need >=2 CUDA devices, found {torch.cuda.device_count()}")

print(f"liger-kernel multi-GPU repro on {args.devices}", flush=True)
for name, fn, build in CASES:
run_case(name, fn, build, args.devices)

if __name__ == "__main__":
main()
```

Run with:

```bash
python repro_liger_multigpu.py
```

Expected behavior: kernels should launch on the device of their input tensors regardless of the currently active CUDA device.

The fix is to enter the input tensor's backend device context around each Triton launch, for example `with device_context(x.device): kernel[grid](...)`.

Contributor guide

Open the contributing guide

Research direction

Start with the entry points in liger_kernel.ops.rms_norm, liger_kernel.ops.swiglu, and liger_kernel.ops.geglu, then trace each Triton launch. Use repro_liger_multigpu.py with two CUDA devices while keeping a different device active. Done means RMSNorm, SwiGLU, and GeGLU launch against their input tensor devices and match the reference outputs.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.