From 21d485f83596a8f360c4739d0c66b9b0b6d46ed4 Mon Sep 17 00:00:00 2001
From: lukec <118525388+sleepcoo@users.noreply.github.com>
Date: Sun, 16 Mar 2025 07:01:04 +0800
Subject: [PATCH] Fix test_create_kvindices unit test (#4452)

---
 test/srt/run_suite.py             | 1 +
 test/srt/test_create_kvindices.py | 4 +---
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py
index ecd1a39dd..5d386afef 100644
--- a/test/srt/run_suite.py
+++ b/test/srt/run_suite.py
@@ -72,6 +72,7 @@ suites = {
         TestFile("test_vision_openai_server.py", 344),
         TestFile("test_w8a8_quantization.py", 46),
         TestFile("test_eval_fp8_accuracy.py", 172),
+        TestFile("test_create_kvindices.py", 2),
     ],
     "nightly": [
         TestFile("test_nightly_gsm8k_eval.py"),
diff --git a/test/srt/test_create_kvindices.py b/test/srt/test_create_kvindices.py
index 383b7ded5..e604e5755 100644
--- a/test/srt/test_create_kvindices.py
+++ b/test/srt/test_create_kvindices.py
@@ -4,9 +4,7 @@ import unittest
 import numpy as np
 import torch
 
-from sglang.srt.layers.attention.flashinfer_utils import (
-    create_flashinfer_kv_indices_triton,
-)
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
 
 
 class TestCreateKvIndices(unittest.TestCase):