diff --git a/mypy.ini b/mypy.ini
index b11cfba1..d093b9b9 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -26,7 +26,6 @@ ignore_missing_imports = True
 
 [mypy-msprobe.*]
 ignore_missing_imports = True
-allow_untyped_imports = True
 
 [mypy-xlite.*]
 ignore_missing_imports = True
diff --git a/vllm_ascend/core/scheduler_dynamic_batch.py b/vllm_ascend/core/scheduler_dynamic_batch.py
index 1127794f..0b78f34f 100644
--- a/vllm_ascend/core/scheduler_dynamic_batch.py
+++ b/vllm_ascend/core/scheduler_dynamic_batch.py
@@ -71,6 +71,8 @@ class BudgetRefiner:
             valid = group[group['cost'] <= slo_limit]
             if not valid.empty:
                 max_row = valid.loc[valid['chunk_size'].idxmax()]
+                assert isinstance(ctx_len, int), "ctx_len must be an integer"
+                assert isinstance(d_num, int), "d_num must be an integer"
                 self.lookup[(ctx_len, d_num)] = int(max_row['chunk_size'])
                 self.context_keys.add(ctx_len)
                 self.dnum_keys.add(d_num)