linkedin / linkedin/Liger-Kernel
Triton kernels can launch on the wrong GPU with multi-GPU inputs
- 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
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