[Kernel] Phase 4 batch-2: migrate JIT operator groups into kernels.ops (no shims) (RFC #29630) (#32015)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-22 17:49:51 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 9456cef279
commit 977ea336cd
95 changed files with 127 additions and 104 deletions
+2 -2
View File
@@ -38,7 +38,6 @@ import torch
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,
@@ -49,6 +48,7 @@ from sglang.kernels.ops.kvcache.cache_move import (
set_kv_buffer_prefix_valid_tiled,
store_cache_4d,
)
from sglang.kernels.ops.kvcache.kvcache import can_use_store_cache, store_cache
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
@@ -4902,7 +4902,7 @@ class MiniMaxSparseKVPool(KVCache):
if index_pool is not None and self._can_fuse_kv_index_store(
index_pool, cache_k, cache_idx_k
):
from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
from sglang.kernels.ops.kvcache.minimax_store_kv_index import store_kv_index
main = self.main_pool
head_bytes = main.head_dim * main.dtype.itemsize