[PYTHON][KVCACHE] Enhance the thread limit for opencl (#2216)

It improves 2x time for tir based page attention for opencl adreno.
mlc-ai · Apr 25, 2024 · 85fffee · 85fffee
1 parent 55b5c00
commit 85fffee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/mlc_llm/nn/kv_cache.py b/python/mlc_llm/nn/kv_cache.py
@@ -887,7 +887,7 @@ def _attention_decode(
     THREAD_LIMIT = 512
     TILE_SIZE_PER_BDX = 2
     if target.kind.name == "opencl" and "android" in str(target.host):
-        THREAD_LIMIT = 64
+        THREAD_LIMIT = 256
         TILE_SIZE_PER_BDX = 1
     max_num_threads_per_block = get_max_num_threads_per_block(target)
     thread_limit = min(max_num_threads_per_block, THREAD_LIMIT)