fix: minor updates and fixs for unit_tests to match current code (#1083)

Author: WuSiYu

unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py

@@ -13,15 +13,10 @@ def test_gen_decode_params_basic():
         b_kv_seq_len,
         b1_cu_kv_seq_len,
         position_ids,
-        max_q_seq_len,
-        max_kv_seq_len,
     ) = gen_decode_params(b_seq_len)
 
     true_b_q_seq_len = torch.ones_like(b_seq_len)
-    b_q_seq_len, b1_cu_q_seq_len, b_kv_seq_len, b1_cu_kv_seq_len, position_ids, max_q_seq_len, max_kv_seq_len
 
-    assert max_q_seq_len == 1
-    assert max_kv_seq_len == b_seq_len.max().item()
     assert torch.equal(b_q_seq_len, true_b_q_seq_len)
     assert torch.equal(b1_cu_q_seq_len, torch.nn.functional.pad(torch.cumsum(true_b_q_seq_len, dim=0), (1, 0), value=0))
     assert torch.equal(b_kv_seq_len, b_seq_len)

unit_tests/common/basemodel/triton_kernel/test_gen_prefill_params.py

@@ -20,12 +20,8 @@ def test_gen_prefill_params_basic():
         b_kv_seq_len,
         b1_cu_kv_seq_len,
         position_ids,
-        max_q_seq_len,
-        max_kv_seq_len,
     ) = gen_prefill_params(input_token_num, b_ready_cache_len, b_seq_len)
 
-    assert max_q_seq_len == true_b_q_seq_len.max().item()
-    assert max_kv_seq_len == b_seq_len.max().item()
     assert torch.equal(b_q_seq_len, true_b_q_seq_len)
     assert torch.equal(b1_cu_q_seq_len, torch.nn.functional.pad(torch.cumsum(true_b_q_seq_len, dim=0), (1, 0), value=0))
     assert torch.equal(b_kv_seq_len, b_seq_len)

unit_tests/common/fused_moe/test_softmax_topk.py

@@ -18,13 +18,13 @@ def benchmark(M, N, K, renorm, runs):
     sgl_vals = torch.empty((M, K), dtype=torch.float32, device="cuda")
     sgl_ids = torch.empty((M, K), dtype=torch.int32, device="cuda")
     # Warm-up
-    sgl_ops.topk_softmax(sgl_vals, sgl_ids, torch.empty_like(sgl_ids), gating)
+    sgl_ops.topk_softmax(sgl_vals, sgl_ids, gating)
     torch.cuda.synchronize()
     start = torch.cuda.Event(True)
     end = torch.cuda.Event(True)
     start.record()
     for _ in range(runs):
-        sgl_ops.topk_softmax(sgl_vals, sgl_ids, torch.empty_like(sgl_ids), gating)
+        sgl_ops.topk_softmax(sgl_vals, sgl_ids, gating)
         if renorm:
             sgl_vals.div_(sgl_vals.sum(-1, keepdim=True).clamp_min(1e-8))
 

unit_tests/server/core/objs/test_sampling_params.py

@@ -127,8 +127,8 @@ def test_decode_node_initialization():
     }
     node.initialize(data)
     assert node.exists is True
-    assert node.node_id_high == (12345678901234567890 >> 64) & 0xFFFFFFFFFFFFFFFF
-    assert node.node_id_low == 12345678901234567890 & 0xFFFFFFFFFFFFFFFF
+    assert node.node_id.node_id_high == (12345678901234567890 >> 64) & 0xFFFFFFFFFFFFFFFF
+    assert node.node_id.node_id_low == 12345678901234567890 & 0xFFFFFFFFFFFFFFFF
     assert node.ip[0] == 192
     assert node.ip[1] == 168
     assert node.ip[2] == 1

unit_tests/server/core/objs/test_shm_req_manager.py

@@ -1,13 +1,34 @@
+import os
 import pytest
 import time
-from lightllm.utils.envs_utils import set_env_start_args
+from easydict import EasyDict
+from lightllm.utils.envs_utils import set_env_start_args, get_env_start_args
 from lightllm.server.core.objs.shm_req_manager import ShmReqManager
 
 
 @pytest.fixture(scope="module", autouse=True)
 def setup_env():
-    set_env_start_args({"running_max_req_size": 10, "disable_chunked_prefill": True, "token_healing_mode": False})
+    original = os.environ.get("LIGHTLLM_START_ARGS")
+    set_env_start_args(
+        EasyDict(
+            running_max_req_size=10,
+            disable_chunked_prefill=True,
+            token_healing_mode=False,
+            enable_flashinfer_prefill=False,
+            enable_flashinfer_decode=False,
+        )
+    )
+    # clear the lru_cache if used
+    if hasattr(get_env_start_args, "cache_clear"):
+        get_env_start_args.cache_clear()
+
     yield
+    if original is not None:
+        os.environ["LIGHTLLM_START_ARGS"] = original
+    else:
+        os.environ.pop("LIGHTLLM_START_ARGS", None)
+    if hasattr(get_env_start_args, "cache_clear"):
+        get_env_start_args.cache_clear()
 
 
 @pytest.fixture(scope="module")