From 800aaefc9e9e4007149074943e20a8d463e4d4b7 Mon Sep 17 00:00:00 2001 From: Thomas Wang Date: Tue, 16 Jun 2026 13:12:12 +0800 Subject: [PATCH] [AMD] Annotate ATOM source for imported v4 unified attention kernels (#28392) --- .../layers/attention/dsv4/unified_kv_kernels/paged_decode.py | 3 +++ .../attention/dsv4/unified_kv_kernels/paged_decode_indices.py | 3 +++ .../layers/attention/dsv4/unified_kv_kernels/paged_prefill.py | 3 +++ python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py | 1 + 4 files changed, 10 insertions(+) diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py index c6b9632c0..6a71caaaa 100644 --- a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# The following kernel is imported from ATOM. +# Source: atom/model_ops/v4_kernels/paged_decode.py + """Sparse decode attention over a unified KV pool with per-token paged indices. Designed for V4 decode + CUDAGraph: replaces the per-fwd `kv_flat_sa` diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py index f46319e46..170baa1dc 100644 --- a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# The following kernel is imported from ATOM. +# Source: atom/model_ops/v4_kernels/paged_decode_indices.py + """V4 paged-decode index scatter — single Triton kernel writes SWA window- prefix paged offsets into the three ragged-packed destination buffers (`kv_indices_swa` / `kv_indices_csa` / `kv_indices_hca`). diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py index cb35f5d0e..6bda80b38 100644 --- a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# The following kernel is imported from ATOM. +# Source: atom/model_ops/v4_kernels/paged_prefill.py + """Sparse prefill attention with two KV sources: paged `unified_kv` (history) and per-fwd flat `kv` (current chunk's input). diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index b66c9481e..08fed1e33 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -375,6 +375,7 @@ class DeepSeekV4LayerItem(NamedTuple): compress_kv_pool: Optional[DeepSeekV4SingleKVPool] = None +# The following kv pool follows ATOM's unified_kv kernel layout. class DeepSeekV4UnifiedKVPool: """ Layout: