[Kernel] Migrate generic attention kernels to sglang.kernels (RFC #29630, Phase 2.5, 4/7) (#30789)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

test/manual/attention/test_trtllm_mla_backend.py

@@ -10,13 +10,13 @@ from sglang.srt.runtime_context import get_parallel, get_server_args
 _parallel_override = get_parallel().override(attn_tp_size=1)
 _parallel_override.__enter__()
 
+from sglang.kernels.ops.attention.utils import get_num_page_per_block_flashmla
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
 from sglang.srt.layers.attention.trtllm_mla_backend import (
     TRTLLMMLABackend,
     TRTLLMMLADecodeMetadata,
 )
-from sglang.srt.layers.attention.utils import get_num_page_per_block_flashmla
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode

test/registered/attention/test_create_kvindices.py

@@ -3,7 +3,7 @@ import unittest
 import numpy as np
 import torch
 
-from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.kernels.ops.attention.utils import create_flashinfer_kv_indices_triton
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

test/registered/kernels/test_flash_mla_backends.py

@@ -26,8 +26,7 @@ from unittest import mock
 
 import torch
 
-from sglang.srt.layers.attention import flash_mla_sm120 as fmod
-from sglang.srt.layers.attention.flash_mla_sm120 import (
+from sglang.kernels.ops.attention.flash_mla_sm120 import (
     _D,
     _NOPE_DIM,
     _NOPE_ROPE_STRIDE,
@@ -39,11 +38,12 @@ from sglang.srt.layers.attention.flash_mla_sm120 import (
     _sm120_sparse_decode_fwd,
     flash_mla_with_kvcache_sm120,
 )
-from sglang.srt.layers.attention.flash_mla_sm120_triton import (
+from sglang.kernels.ops.attention.flash_mla_sm120_triton import (
     _apply_attn_sink,
     _merge_partial_attn,
     flash_mla_sparse_decode_triton,
 )
+from sglang.srt.layers.attention import flash_mla_sm120 as fmod
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase