triton KV append kernel hard-asserts CUDA, blocking non-CUDA Triton backends
- 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
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