[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:
co-authored by
Claude Opus 4.8
parent
9456cef279
commit
977ea336cd
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user