[Kernel] Migrate DSA + DSV4 attention kernels to sglang.kernels (RFC #29630, Phase 2.5, 5/7) (#30792)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-15 11:11:22 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 4ae9cc3c81
commit ba5be86d42
60 changed files with 663 additions and 583 deletions
+5 -5
View File
@@ -36,6 +36,11 @@ import triton
import triton.language as tl
from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
from sglang.kernels.ops.attention.dsa import index_buf_accessor
from sglang.kernels.ops.attention.dsa.quant_k_cache import (
quantize_k_cache,
quantize_k_cache_separate,
)
from sglang.kernels.ops.kvcache.cache_move import (
copy_all_layer_kv_cache_func,
set_kv_buffer_prefix_valid_tiled,
@@ -45,11 +50,6 @@ from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
from sglang.srt.configs.mamba_utils import BaseLinearStateParams
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa import index_buf_accessor
from sglang.srt.layers.attention.dsa.quant_k_cache import (
quantize_k_cache,
quantize_k_cache_separate,
)
from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator