You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2024/03/12 13:10:25 UTC
(tvm) branch main updated: [Fix][Relax] Fix top-p/top-k sampling kernel (#16703)
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 ef8c42813c [Fix][Relax] Fix top-p/top-k sampling kernel (#16703)
ef8c42813c is described below
commit ef8c42813c8edec0b1aae50b02500054b461571c
Author: Ruihang Lai <ru...@cs.cmu.edu>
AuthorDate: Tue Mar 12 09:10:19 2024 -0400
[Fix][Relax] Fix top-p/top-k sampling kernel (#16703)
This PR fixes a typo in the samping kernel of top-p/top-k sampling op.
Prior to this PR, the kernel has out-of-bound global memory access
due to a miss when introducing `sample_indices` in #16675.
The correctness pass did not reveal this issue by directly running
the test or running through pytest. But actually, if we use
compute-sanitizer from NVIDIA, it will report the illegal memory
access:
```
> compute-sanitizer --tool memcheck --print-limit=5 --launch-timeout 3600 python tests/python/relax/test_frontend_nn_op.py
========= COMPUTE-SANITIZER
========= Invalid __global__ read of size 8 bytes
========= at 0x4e90 in get_index_from_sorted_kernel
========= by thread (7,0,0) in block (0,0,0)
========= Address 0x7fe35ac00238 is out of bounds
========= and is 9 bytes after the nearest allocation at 0x7fe35ac00200 of size 48 bytes
========= Saved host backtrace up to driver entry point at kernel launch time
...
```
---
python/tvm/relax/frontend/nn/op.py | 6 +++---
tests/python/relax/test_frontend_nn_op.py | 6 +++---
2 files changed, 6 insertions(+), 6 deletions(-)
diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py
index d299d39439..137dc897c0 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -23,9 +23,9 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, TypeVar
import numpy as np
+from tvm import te
from tvm import tir as _tir
from tvm.script import tir as T
-from tvm import te
from ... import expr as rx
from ... import op as _op
@@ -2386,13 +2386,13 @@ def sample_top_p_top_k_from_sorted_prob(
or v_ax1 + 1 == vocab_size
):
if v_ax1 == 0:
- output_index[v_ax0, 0] = indices[v_ax0, 0]
+ output_index[v_ax0, 0] = indices[sample_indices[v_ax0, T.int64(0)], 0]
elif (
usample[v_ax0, T.int64(0)]
>= cumsum_sorted[sample_indices[v_ax0, T.int64(0)], v_ax1 - 1]
/ renorm_prob[sample_indices[v_ax0, T.int64(0)], 0]
):
- output_index[v_ax0, 0] = indices[v_ax0, v_ax1]
+ output_index[v_ax0, 0] = indices[sample_indices[v_ax0, T.int64(0)], v_ax1]
cumsum_sorted = cumsum(sorted_prob, axis=1)
diff --git a/tests/python/relax/test_frontend_nn_op.py b/tests/python/relax/test_frontend_nn_op.py
index eb1df67a8f..5f05abf7c2 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -973,14 +973,14 @@ def test_sample_top_p_top_k_from_sorted_prob():
for ax0, ax1 in T.grid(out_batch, vocab_size):
with T.block("T_get_index_from_sorted"):
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
- T.reads(usample[v_ax0, T.int64(0)], cumsum_sorted[sample_indices[v_ax0, T.int64(0)], v_ax1 - T.int64(1):v_ax1 - T.int64(1) + T.int64(2)], sample_indices[v_ax0, T.int64(0)], renorm_prob[sample_indices[v_ax0, T.int64(0)], 0], indices[v_ax0, T.min(T.int64(0), v_ax1):T.min(T.int64(0), v_ax1) + (T.max(T.int64(0), v_ax1) + T.int64(1) - T.min(T.int64(0), v_ax1))])
+ T.reads(usample[v_ax0, T.int64(0)], cumsum_sorted[sample_indices[v_ax0, T.int64(0)], v_ax1 - T.int64(1):v_ax1 - T.int64(1) + T.int64(2)], sample_indices[v_ax0, T.int64(0)], renorm_prob[sample_indices[v_ax0, T.int64(0)], 0], indices[sample_indices[v_ax0, T.int64(0)], T.min(T.int64(0), v_ax1):T.min(T.int64(0), v_ax1) + (T.max(T.int64(0), v_ax1) + T.int64(1) - T.min(T.int64(0), v_ax1))])
T.writes(output_index[v_ax0, 0])
if usample[v_ax0, T.int64(0)] < cumsum_sorted[sample_indices[v_ax0, T.int64(0)], v_ax1] / renorm_prob[sample_indices[v_ax0, T.int64(0)], 0] or v_ax1 + T.int64(1) == vocab_size:
if v_ax1 == T.int64(0):
- output_index[v_ax0, 0] = indices[v_ax0, 0]
+ output_index[v_ax0, 0] = indices[sample_indices[v_ax0, T.int64(0)], 0]
else:
if usample[v_ax0, T.int64(0)] >= cumsum_sorted[sample_indices[v_ax0, T.int64(0)], v_ax1 - T.int64(1)] / renorm_prob[sample_indices[v_ax0, T.int64(0)], 0]:
- output_index[v_ax0, 0] = indices[v_ax0, v_ax1]
+ output_index[v_ax0, 0] = indices[sample_indices[v_ax0, T.int64(0)], v_ax1]
@T.prim_func(private=True)
def get_renorm_prob(A: T.handle, B: T.handle, C: T.handle, D: T.handle):