(tvm) branch main updated: [Fix][KVCache] Fix incorrect tile size calculation (#17595)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 077e8ebc22 [Fix][KVCache] Fix incorrect tile size calculation (#17595)
077e8ebc22 is described below

commit 077e8ebc2265efd84254081051dd552d9aa61e8e
Author: Ruihang Lai <ruiha...@cs.cmu.edu>
AuthorDate: Sun Jan 19 10:36:28 2025 -0500

    [Fix][KVCache] Fix incorrect tile size calculation (#17595)
    
    This PR fixes the tile size calculation in the TIR attention
    kernels, where the computed tile sizes may not divide the total
    loop extent.
---
 python/tvm/relax/frontend/nn/llm/kv_cache.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/tvm/relax/frontend/nn/llm/kv_cache.py 
b/python/tvm/relax/frontend/nn/llm/kv_cache.py
index f60c40efa2..399e418c46 100644
--- a/python/tvm/relax/frontend/nn/llm/kv_cache.py
+++ b/python/tvm/relax/frontend/nn/llm/kv_cache.py
@@ -855,7 +855,7 @@ def _attention_prefill(
         cnt = (x * y) // t
         assert (x * y) % t == 0
         tile_y = (int)(math.ceil(math.sqrt(cnt)))
-        while (cnt % tile_y != 0 or y % tile_y != 0) and tile_y <= cnt:
+        while (cnt % tile_y != 0 or y % tile_y != 0 or x % (cnt // tile_y) != 
0) and tile_y <= cnt:
             tile_y += 1
         assert tile_y <= cnt
         tile_x = cnt // tile_y
@@ -1509,7 +1509,7 @@ def _attention_sequence_prefill(
         cnt = (x * y) // t
         assert (x * y) % t == 0
         tile_y = (int)(math.ceil(math.sqrt(cnt)))
-        while (cnt % tile_y != 0 or y % tile_y != 0) and tile_y <= cnt:
+        while (cnt % tile_y != 0 or y % tile_y != 0 or x % (cnt // tile_y) != 
0) and tile_y <= cnt:
             tile_y += 1
         assert tile_y <= cnt
         tile_x = cnt // tile_y
@@ -1867,7 +1867,7 @@ def _attention_prefill_ragged(h_kv, h_q, d, dtype, 
rope_scaling: Dict[str, Any],
         cnt = (x * y) // t
         assert (x * y) % t == 0
         tile_y = (int)(math.ceil(math.sqrt(cnt)))
-        while (cnt % tile_y != 0 or y % tile_y != 0) and tile_y <= cnt:
+        while (cnt % tile_y != 0 or y % tile_y != 0 or x % (cnt // tile_y) != 
0) and tile_y <= cnt:
             tile_y += 1
         assert tile_y <= cnt
         tile_x = cnt // tile_y