Commit 71276043 authored by wanghl6's avatar wanghl6
Browse files

convert_req_index_to_global_index opt

parent a0ac95b0
......@@ -302,6 +302,18 @@ def triton_convert_req_index_to_global_index(
prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace
starts for each prefill request
"""
if (envs.USE_LIGHTOP_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX):
from lightop import op
return op.convert_req_index_to_global_index(
req_id,
block_table,
token_indices,
BLOCK_SIZE,
NUM_TOPK_TOKENS,
HAS_PREFILL_WORKSPACE,
prefill_workspace_request_ids,
prefill_workspace_starts
)
assert req_id.dtype == torch.int32
assert block_table.dtype == torch.int32
assert token_indices.dtype == torch.int32
......
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