NVIDIA / NVIDIA/Megatron-LM

triton KV append kernel hard-asserts CUDA, blocking non-CUDA Triton backends

Open Beginner friendly
#6,729 0 comments 0 reactions 0 assignees View on GitHub
community-request
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

### Bug description

`triton_append_key_value_cache` in `megatron/core/inference/contexts/fused_kv_append_kernel.py` hard-asserts that all tensors are on CUDA:

```python
assert (
key.device.type == 'cuda'
and value.device.type == 'cuda'
and memory_buffer.device.type == 'cuda'
), "All tensors must be on CUDA devices."
```

The kernel itself is pure Triton and device-agnostic; the assert is the only CUDA-ism in the function. On any accelerator that exposes a Triton backend but is not CUDA (e.g. Ascend NPU, where `torch_npu` provides the Triton backend and the device type is `npu`), the dynamic-batching inference path hits this assert on the first KV append even though the kernel would run fine.

### Environment

- Megatron-LM main @ 1d82259e1
- Non-CUDA accelerator with a Triton backend (e.g. Ascend 910B, device type `npu`)

### Steps to reproduce

Run the dynamic-batching inference path with `triton_append_key_value_cache` (the fused KV-append kernel) on a non-CUDA Triton backend. The CUDA-only device assert fires in the input-validation preamble.

### Expected behavior

The triton kernel is a device-agnostic launchable kernel and should run on any device backed by Triton. The assert's real intent is to reject tensors that cannot be launched on a Triton backend at all (CPU, meta); it should be expressed that way instead of as a CUDA allowlist, so it holds for every current and future Triton-backed accelerator.

### Proposed fix

```python
assert (
key.device.type not in ('cpu', 'meta')
and value.device.type not in ('cpu', 'meta')
and memory_buffer.device.type not in ('cpu', 'meta')
), "All tensors must be on a device with a Triton backend (CUDA, NPU, ...)."
```

Contributor guide

Open the contributing guide

Research direction

Open megatron/core/inference/contexts/fused_kv_append_kernel.py and inspect triton_append_key_value_cache, especially its input-validation preamble. Reproduce the dynamic-batching inference path on a non-CUDA Triton backend, then verify that valid accelerator tensors pass while CPU and meta tensors remain rejected.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
ai, backend
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
85/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.