Unverified Commit 40218a82 authored by Nick Hill's avatar Nick Hill Committed by GitHub
Browse files

[ModelRunner V2] Revert token rank comparison difference for now (#34017)


Signed-off-by: default avatarNick Hill <nickhill123@gmail.com>
parent 1c3b2205
...@@ -68,7 +68,7 @@ def _ranks_kernel( ...@@ -68,7 +68,7 @@ def _ranks_kernel(
for i in range(0, vocab_size, BLOCK_SIZE): for i in range(0, vocab_size, BLOCK_SIZE):
block = i + tl.arange(0, BLOCK_SIZE) block = i + tl.arange(0, BLOCK_SIZE)
logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf")) logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
n += tl.sum((logits > x).to(tl.int32)) n += tl.sum((logits >= x).to(tl.int32))
tl.store(output_ptr + req_idx, n) tl.store(output_ptr + req_idx, n)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment