Lint

mlc-ai · Feb 2, 2025 · 7ae4087 · 7ae4087
1 parent aad79c1
commit 7ae4087
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 9 deletions.
diff --git a/cpp/serve/sampler/cpu_sampler.cc b/cpp/serve/sampler/cpu_sampler.cc
@@ -555,7 +555,7 @@ class CPUSampler : public SamplerObj {
   NDArray CopyProbsToCPU(NDArray probs_on_device) {
     // probs_on_device: (n, v)
     if (probs_on_device->device.device_type == kDLCPU) {
-        return probs_on_device;
+      return probs_on_device;
     }
 
     ICHECK(probs_on_device->device.device_type != kDLCPU);

diff --git a/python/mlc_llm/compiler_pass/attach_logit_processor.py b/python/mlc_llm/compiler_pass/attach_logit_processor.py
@@ -28,17 +28,17 @@ def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IR
         """Entrypoint"""
         mod = mod.clone()
         if str(self.target.kind) == "llvm":
-            mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace_cpu(self.target)
-            mod["apply_penalty_inplace"] = _get_apply_penalty_inplace_cpu(self.target)
-            mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace_cpu(self.target)
+            mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace_cpu()
+            mod["apply_penalty_inplace"] = _get_apply_penalty_inplace_cpu()
+            mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace_cpu()
         else:
             mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace(self.target)
             mod["apply_penalty_inplace"] = _get_apply_penalty_inplace(self.target)
             mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace(self.target)
         return mod
 
 
-def _get_apply_logit_bias_inplace_cpu(target: tvm.target.Target):
+def _get_apply_logit_bias_inplace_cpu():
     @T.prim_func
     def _apply_logit_bias_inplace(
         var_logits: T.handle,
@@ -110,7 +110,7 @@ def _apply_logit_bias_inplace(
     return _apply_logit_bias_inplace
 
 
-def _get_apply_penalty_inplace_cpu(target: tvm.target.Target):
+def _get_apply_penalty_inplace_cpu():
     @T.prim_func
     def _apply_penalty_inplace(  # pylint: disable=too-many-arguments,too-many-locals
         var_logits: T.handle,
@@ -209,7 +209,7 @@ def _apply_penalty_inplace(  # pylint: disable=too-many-arguments,too-many-local
     return _apply_penalty_inplace
 
 
-def _get_apply_bitmask_inplace_cpu(target: tvm.target.Target):
+def _get_apply_bitmask_inplace_cpu():
     @T.prim_func
     def _apply_bitmask_inplace(
         var_logits: T.handle,

diff --git a/python/mlc_llm/compiler_pass/attach_softmax_with_temperature.py b/python/mlc_llm/compiler_pass/attach_softmax_with_temperature.py
@@ -246,5 +246,4 @@ def apply_gpu_schedule(target, sch):
 
     if target.kind.name == "llvm":
         return chunk_lse, sch.mod["softmax_with_chunked_sum"]
-    else:
-        return apply_gpu_schedule(target, sch)
+    return apply_gpu_schedule(target, sch)