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

[tengqm]

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:

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 @ 1d82259
• 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

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, ...)."