mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
feat(h20): port the four MTP kernels, so speculative decoding no longer needs TileLang
`DEEPINFER_USE_NEXTN_DECODE` defaults to 5, so speculative decoding is on unless someone turns it off, and these four ran on every decode step. On a host without TileLang the first step failed; on the H20 they were the only kernels still answering from the other backend during a normal run. The interesting one is `remap_draft_topk`. Its TileLang form builds a vocabulary-windowed lookup table in shared memory and reads the logprobs back out of it. Triton has no equivalent primitive, so the table is gone and each target index is matched against the draft indices directly: draft indices are unique per token id, so at most one candidate can match and a max over the block is well defined. That is O(topk^2) instead of O(vocab/pool_size * topk), which would be the wrong trade at a large `force_topk` -- but `force_topk` defaults to 0, so this kernel is off the default path and correctness comes first. Worth revisiting if a deployment turns it on. The other three adapt rather than change shape: * `get_accept_result`'s prefix scan becomes a bounded loop over the widest span in the block; lanes with a shorter span hold their value on the extra iterations. MTP packs at most num_steps+1 rows per request, so the bound stays small. * `finalize_verification` reduces `batch_confidence_error` within the program and issues one atomic instead of one per row, because a scalar pointer with a vector value is not valid in this Triton version. Accumulating in a different order changes the last bits, so its cases compare with a tolerance and say why. * `get_fused_dflash_get_input` rounds its column loop up. `padded_kv_len` is a multiple of 64 but not necessarily of the tile width, and flooring leaves the tail of every row unwritten. 33 new A/B cases across the four, covering each constexpr variant the engine builds. The full run is green and an end-to-end pass is token-for-token identical to the run before the change, with the MTP fallback notices gone.
This commit is contained in:
loopyt
parent
bf903a91b5
commit
8ae7365cbc
@ -78,3 +78,116 @@ def load_kv_index_tl(
|
||||
)
|
||||
|
||||
return launch
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _get_fused_dflash_get_input_kernel(
|
||||
SLOT_MAPPING_SWA,
|
||||
POSITIONS,
|
||||
FTP,
|
||||
LAST_VALID_OFFSET,
|
||||
NEWX_SLOT_MAPPING_OUT,
|
||||
POSITIONS_OUT,
|
||||
INPUT_IDS_OUT,
|
||||
VALID_TOKEN_MASK_OUT,
|
||||
GATE_FORCE_RANDOM_OUT,
|
||||
KV_INDEX_OUT,
|
||||
kv_len,
|
||||
window_size,
|
||||
swa_block_size: tl.constexpr,
|
||||
padded_kv_len: tl.constexpr,
|
||||
has_valid_token_mask: tl.constexpr,
|
||||
has_gate_force_random: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
req = tl.program_id(0)
|
||||
bx = tl.program_id(1)
|
||||
# The row stride of the packed output is the grid's second dimension, so it does not need to be
|
||||
# a compile-time constant.
|
||||
dflash_block_size = tl.num_programs(1)
|
||||
|
||||
ftp = tl.load(FTP + req)
|
||||
# A padding request keeps its loads in bounds but writes -1 / 0 everywhere.
|
||||
is_valid = ftp >= 0
|
||||
idx = tl.maximum(ftp, 0)
|
||||
so = tl.load(SLOT_MAPPING_SWA + idx) - tl.load(POSITIONS + idx) % swa_block_size
|
||||
ap = tl.load(POSITIONS + idx) + tl.load(LAST_VALID_OFFSET + req)
|
||||
|
||||
out_offset = req * dflash_block_size + bx
|
||||
tl.store(
|
||||
NEWX_SLOT_MAPPING_OUT + out_offset,
|
||||
tl.where(is_valid, (ap + 1 + bx) % swa_block_size + so, -1),
|
||||
)
|
||||
tl.store(POSITIONS_OUT + out_offset, tl.where(is_valid, ap + 1 + bx, 0))
|
||||
tl.store(INPUT_IDS_OUT + out_offset, 0)
|
||||
if has_valid_token_mask:
|
||||
tl.store(VALID_TOKEN_MASK_OUT + out_offset, is_valid)
|
||||
if has_gate_force_random:
|
||||
tl.store(GATE_FORCE_RANDOM_OUT + out_offset, 0)
|
||||
|
||||
# Every draft slot of a request fills the same kv_index columns; the row differs.
|
||||
start_pos = ap - (window_size - 1)
|
||||
row = KV_INDEX_OUT + out_offset * padded_kv_len
|
||||
# Both operands are compile-time constants, so this bounds a static loop. It must round up:
|
||||
# `padded_kv_len` is a multiple of 64 but not necessarily of BLOCK, and floor division would
|
||||
# leave the tail of the row unwritten.
|
||||
for tile in range(0, (padded_kv_len + BLOCK - 1) // BLOCK):
|
||||
cols = tile * BLOCK + tl.arange(0, BLOCK)
|
||||
in_range = cols < padded_kv_len
|
||||
value = start_pos + cols
|
||||
keep = is_valid & in_range & (cols < kv_len) & (value >= 0)
|
||||
tl.store(row + cols, tl.where(keep, value % swa_block_size + so, -1), mask=in_range)
|
||||
|
||||
|
||||
def get_fused_dflash_get_input_tl(
|
||||
swa_block_size: int,
|
||||
padded_kv_len: int,
|
||||
has_valid_token_mask: bool,
|
||||
has_gate_force_random: bool,
|
||||
) -> Any:
|
||||
"""Build the kernel that packs each request's draft-token inputs, one CTA per draft slot."""
|
||||
block = 128
|
||||
|
||||
def launch(
|
||||
slot_mapping_swa: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
ftp: torch.Tensor,
|
||||
last_valid_offset: torch.Tensor,
|
||||
newx_slot_mapping_out: torch.Tensor,
|
||||
positions_out: torch.Tensor,
|
||||
input_ids_out: torch.Tensor,
|
||||
valid_token_mask_out: torch.Tensor | None,
|
||||
gate_force_random_out: torch.Tensor | None,
|
||||
kv_index_out: torch.Tensor,
|
||||
kv_len: int,
|
||||
window_size: int,
|
||||
) -> None:
|
||||
real_cnt, dflash_block_size = newx_slot_mapping_out.shape
|
||||
if real_cnt == 0 or dflash_block_size == 0:
|
||||
return
|
||||
# The interface passes None for a disabled mask; Triton needs a real pointer even when the
|
||||
# branch reading it is compiled out.
|
||||
valid_mask = valid_token_mask_out if valid_token_mask_out is not None else input_ids_out
|
||||
gate_out = gate_force_random_out if gate_force_random_out is not None else input_ids_out
|
||||
_get_fused_dflash_get_input_kernel[(real_cnt, dflash_block_size)](
|
||||
slot_mapping_swa,
|
||||
positions,
|
||||
ftp,
|
||||
last_valid_offset,
|
||||
newx_slot_mapping_out,
|
||||
positions_out,
|
||||
input_ids_out,
|
||||
valid_mask,
|
||||
gate_out,
|
||||
kv_index_out,
|
||||
kv_len,
|
||||
window_size,
|
||||
swa_block_size=swa_block_size,
|
||||
padded_kv_len=padded_kv_len,
|
||||
has_valid_token_mask=has_valid_token_mask,
|
||||
has_gate_force_random=has_gate_force_random,
|
||||
BLOCK=block,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
|
||||
@ -368,3 +368,457 @@ def copy_block_tables_tl(max_seqlen_q: int, max_num_blocks: int, BLOCK_SIZE: int
|
||||
)
|
||||
|
||||
return launch
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _get_accept_result_kernel(
|
||||
RAW_ACCEPT,
|
||||
GOLDEN,
|
||||
DRAFT_POOL,
|
||||
READ_ROWS,
|
||||
CANDIDATE_TOKEN_IDS,
|
||||
IS_VERIFY,
|
||||
FTP,
|
||||
ROW_TO_REQ,
|
||||
TOKEN_IDS_OUT,
|
||||
ACCEPT_OUT,
|
||||
IS_FINAL_OUT,
|
||||
total_q,
|
||||
HAS_GOLDEN: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
rows = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
|
||||
live = rows < total_q
|
||||
|
||||
accept = tl.load(RAW_ACCEPT + rows, mask=live, other=0) != 0
|
||||
token = tl.load(CANDIDATE_TOKEN_IDS + rows, mask=live, other=0)
|
||||
if HAS_GOLDEN:
|
||||
# A golden token replaces the draft's verdict: the request already knows the token it
|
||||
# wants there, so acceptance becomes "did the draft propose that same token".
|
||||
golden = tl.load(GOLDEN + rows, mask=live, other=-1)
|
||||
draft_row = tl.load(READ_ROWS + rows, mask=live, other=0)
|
||||
draft_token = tl.load(DRAFT_POOL + draft_row, mask=live, other=0)
|
||||
has_golden = golden != -1
|
||||
accept = tl.where(has_golden, golden == draft_token, accept)
|
||||
token = tl.where(has_golden, golden, token)
|
||||
|
||||
is_verify = tl.load(IS_VERIFY + rows, mask=live, other=0) != 0
|
||||
req = tl.load(ROW_TO_REQ + rows, mask=live, other=0)
|
||||
start = tl.load(FTP + req, mask=live, other=0)
|
||||
span = tl.where(live, rows - start, 0)
|
||||
|
||||
# Prefix acceptance: a row survives only if every earlier draft row in its request both was a
|
||||
# verify row and was accepted. Walking `start..row-1` per row is what the TileLang kernel does
|
||||
# with a while loop; here the trip count is the widest span in the block, and lanes whose span
|
||||
# is shorter hold their value on the extra iterations. MTP packs at most `num_steps + 1` rows
|
||||
# per request, so that bound stays small.
|
||||
alive = tl.full([BLOCK], 1, tl.int1)
|
||||
max_span = tl.max(span, axis=0)
|
||||
for offset in range(0, max_span):
|
||||
prev = start + offset
|
||||
active = live & (offset < span) & (prev < total_q)
|
||||
prev_accept = tl.load(RAW_ACCEPT + prev, mask=active, other=0) != 0
|
||||
if HAS_GOLDEN:
|
||||
prev_golden = tl.load(GOLDEN + prev, mask=active, other=-1)
|
||||
prev_draft_row = tl.load(READ_ROWS + prev, mask=active, other=0)
|
||||
prev_draft_token = tl.load(DRAFT_POOL + prev_draft_row, mask=active, other=0)
|
||||
prev_accept = tl.where(prev_golden != -1, prev_golden == prev_draft_token, prev_accept)
|
||||
prev_verify = tl.load(IS_VERIFY + prev, mask=active, other=0) != 0
|
||||
alive = alive & (~active | (prev_verify & prev_accept))
|
||||
|
||||
tl.store(TOKEN_IDS_OUT + rows, token, mask=live)
|
||||
tl.store(ACCEPT_OUT + rows, alive & is_verify & accept, mask=live)
|
||||
tl.store(IS_FINAL_OUT + rows, alive & (~is_verify | ~accept), mask=live)
|
||||
|
||||
|
||||
def get_accept_result_tl(has_golden: bool) -> Any:
|
||||
"""Build the kernel that turns raw draft verdicts into prefix-accepted tokens."""
|
||||
|
||||
def launch(
|
||||
raw_accept: torch.Tensor,
|
||||
golden: torch.Tensor,
|
||||
draft_token_ids_pool: torch.Tensor,
|
||||
read_rows: torch.Tensor,
|
||||
candidate_token_ids: torch.Tensor,
|
||||
is_verify: torch.Tensor,
|
||||
ftp: torch.Tensor,
|
||||
row_to_req: torch.Tensor,
|
||||
token_ids_out: torch.Tensor,
|
||||
accept_out: torch.Tensor,
|
||||
is_final_out: torch.Tensor,
|
||||
) -> None:
|
||||
total_q = raw_accept.shape[0]
|
||||
if total_q == 0:
|
||||
return
|
||||
block = 128
|
||||
_get_accept_result_kernel[(cdiv(total_q, block),)](
|
||||
raw_accept,
|
||||
golden,
|
||||
draft_token_ids_pool,
|
||||
read_rows,
|
||||
candidate_token_ids,
|
||||
is_verify,
|
||||
ftp,
|
||||
row_to_req,
|
||||
token_ids_out,
|
||||
accept_out,
|
||||
is_final_out,
|
||||
total_q,
|
||||
HAS_GOLDEN=has_golden,
|
||||
BLOCK=block,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _finalize_verification_zero_kernel(
|
||||
NUM_GENERATED,
|
||||
NUM_VERIFIED,
|
||||
BATCH_CONFIDENCE_ERROR,
|
||||
num_steps: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
# Counters are accumulated with atomics in the main kernel, so they are zeroed in a separate
|
||||
# launch rather than by a grid barrier that would disable the overlap PDL gives us.
|
||||
steps = tl.arange(0, BLOCK)
|
||||
if num_steps > 0:
|
||||
live = steps < num_steps
|
||||
tl.store(NUM_GENERATED + steps, tl.zeros([BLOCK], tl.int32), mask=live)
|
||||
tl.store(NUM_VERIFIED + steps, tl.zeros([BLOCK], tl.int32), mask=live)
|
||||
tl.store(BATCH_CONFIDENCE_ERROR, 0.0)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _finalize_verification_kernel(
|
||||
EFFECTIVE_ACCEPT,
|
||||
IS_FINAL_IN,
|
||||
EFFECTIVE_TOKEN_IDS,
|
||||
TOKEN_LOGPROBS,
|
||||
SUM_SMALLER_PROBS,
|
||||
TOP_LOGPROBS,
|
||||
TOP_TOKEN_IDS,
|
||||
FTP,
|
||||
SEQLEN_Q,
|
||||
ROW_TO_REQ,
|
||||
IS_VERIFY,
|
||||
ACCEPT_PROBS,
|
||||
CONFIDENCE_POOL,
|
||||
MTP_INFO_READ_ROWS,
|
||||
STOP_TOKEN_IDS,
|
||||
LAST_VALID_OFFSET,
|
||||
NUM_GENERATED,
|
||||
NUM_VERIFIED,
|
||||
BATCH_CONFIDENCE_ERROR,
|
||||
EMITTED_STOP_TOKEN,
|
||||
SAMPLED_TOKEN_IDS,
|
||||
SAMPLED_TOKEN_LOGPROBS,
|
||||
SAMPLED_SUM_SMALLER_PROBS,
|
||||
SAMPLED_TOP_LOGPROBS,
|
||||
SAMPLED_TOP_TOKEN_IDS,
|
||||
VERIFIED_TOKEN_IDS,
|
||||
VERIFIED_TOKEN_LOGPROBS,
|
||||
VERIFIED_SUM_SMALLER_PROBS,
|
||||
VERIFIED_TOP_LOGPROBS,
|
||||
VERIFIED_TOP_TOKEN_IDS,
|
||||
total_q,
|
||||
topk_n,
|
||||
top_lp_row_stride,
|
||||
top_ids_row_stride,
|
||||
HAS_SUM_SMALLER_PROBS: tl.constexpr,
|
||||
HAS_CONFIDENCE: tl.constexpr,
|
||||
NUM_STEPS: tl.constexpr,
|
||||
NUM_STOP_TOKENS: tl.constexpr,
|
||||
BLOCK_ROWS: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
):
|
||||
rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
|
||||
live = rows < total_q
|
||||
|
||||
accepted = tl.load(EFFECTIVE_ACCEPT + rows, mask=live, other=0) != 0
|
||||
is_final = tl.load(IS_FINAL_IN + rows, mask=live, other=0) != 0
|
||||
is_verify = tl.load(IS_VERIFY + rows, mask=live, other=0) != 0
|
||||
token = tl.load(EFFECTIVE_TOKEN_IDS + rows, mask=live, other=0)
|
||||
logprob = tl.load(TOKEN_LOGPROBS + rows, mask=live, other=0.0)
|
||||
req = tl.load(ROW_TO_REQ + rows, mask=live, other=0)
|
||||
start = tl.load(FTP + req, mask=live, other=0)
|
||||
step = rows - start
|
||||
|
||||
# Payload for rows that survive verification; rejected rows are recorded as -1 / 0 so a
|
||||
# consumer can tell "no token here" from a genuine token 0.
|
||||
tl.store(VERIFIED_TOKEN_IDS + rows, tl.where(accepted, token, -1), mask=live)
|
||||
tl.store(VERIFIED_TOKEN_LOGPROBS + rows, tl.where(accepted, logprob, 0.0), mask=live)
|
||||
if HAS_SUM_SMALLER_PROBS:
|
||||
ssp = tl.load(SUM_SMALLER_PROBS + rows, mask=live, other=0.0)
|
||||
tl.store(VERIFIED_SUM_SMALLER_PROBS + rows, tl.where(accepted, ssp, 0.0), mask=live)
|
||||
|
||||
if NUM_STEPS > 0:
|
||||
tl.atomic_add(NUM_GENERATED + step, 1, mask=live & accepted & (step < NUM_STEPS))
|
||||
verified_slot = tl.maximum(tl.minimum(step, NUM_STEPS - 1), 0)
|
||||
tl.atomic_add(NUM_VERIFIED + verified_slot, 1, mask=live & is_verify)
|
||||
|
||||
if HAS_CONFIDENCE:
|
||||
accept_prob = tl.load(ACCEPT_PROBS + rows, mask=live, other=0.0)
|
||||
read_row = tl.load(MTP_INFO_READ_ROWS + rows, mask=live, other=0)
|
||||
confidence = tl.load(CONFIDENCE_POOL + read_row, mask=live, other=0.0)
|
||||
delta = accept_prob - confidence
|
||||
delta = tl.where(delta < 0, -delta, delta)
|
||||
# The counter is one scalar, so reduce within the program and issue a single atomic.
|
||||
# Summing in a different order than the TileLang kernel does changes the last bits, which
|
||||
# is why the comparison tolerance for float32 exists.
|
||||
tl.atomic_add(BATCH_CONFIDENCE_ERROR, tl.sum(tl.where(live & is_verify, delta, 0.0)))
|
||||
|
||||
# Per-request payload comes from the request's final row, of which there is exactly one.
|
||||
if HAS_CONFIDENCE:
|
||||
# Did the emitted stream hit any stop token? Scan the rows this request will actually emit,
|
||||
# which is its prefix up to and including the final row.
|
||||
stop_hit = tl.zeros([BLOCK_ROWS], tl.int1)
|
||||
for scan in tl.static_range(NUM_STEPS + 1):
|
||||
scan_row = start + scan
|
||||
in_prefix = scan <= step
|
||||
scan_token = tl.load(EFFECTIVE_TOKEN_IDS + scan_row, mask=live & is_final & in_prefix, other=0)
|
||||
for stop_idx in tl.static_range(NUM_STOP_TOKENS):
|
||||
stop_token = tl.load(
|
||||
STOP_TOKEN_IDS + req * NUM_STOP_TOKENS + stop_idx,
|
||||
mask=live & is_final & in_prefix,
|
||||
other=-1,
|
||||
)
|
||||
stop_hit = stop_hit | (in_prefix & (scan_token == stop_token))
|
||||
tl.store(EMITTED_STOP_TOKEN + req, stop_hit, mask=live & is_final)
|
||||
|
||||
tl.store(LAST_VALID_OFFSET + req, step, mask=live & is_final)
|
||||
tl.store(SAMPLED_TOKEN_IDS + req, token, mask=live & is_final)
|
||||
tl.store(SAMPLED_TOKEN_LOGPROBS + req, logprob, mask=live & is_final)
|
||||
if HAS_SUM_SMALLER_PROBS:
|
||||
tl.store(SAMPLED_SUM_SMALLER_PROBS + req, ssp, mask=live & is_final)
|
||||
|
||||
# Top-k payload: accepted rows keep theirs, the final row also publishes to the request slot.
|
||||
for tile in range(0, tl.cdiv(topk_n, BLOCK_K)):
|
||||
k = tile * BLOCK_K + tl.arange(0, BLOCK_K)
|
||||
in_range = live[:, None] & (k[None, :] < topk_n)
|
||||
keep = accepted[:, None] | is_final[:, None]
|
||||
lp = tl.load(
|
||||
TOP_LOGPROBS + rows[:, None] * top_lp_row_stride + k[None, :],
|
||||
mask=in_range & keep,
|
||||
other=0.0,
|
||||
)
|
||||
tok = tl.load(
|
||||
TOP_TOKEN_IDS + rows[:, None] * top_ids_row_stride + k[None, :],
|
||||
mask=in_range & keep,
|
||||
other=0,
|
||||
)
|
||||
acc = accepted[:, None]
|
||||
tl.store(
|
||||
VERIFIED_TOP_LOGPROBS + rows[:, None] * topk_n + k[None, :],
|
||||
tl.where(acc, lp, 0.0),
|
||||
mask=in_range,
|
||||
)
|
||||
tl.store(
|
||||
VERIFIED_TOP_TOKEN_IDS + rows[:, None] * topk_n + k[None, :],
|
||||
tl.where(acc, tok, 0),
|
||||
mask=in_range,
|
||||
)
|
||||
tl.store(
|
||||
SAMPLED_TOP_LOGPROBS + req[:, None] * topk_n + k[None, :],
|
||||
lp,
|
||||
mask=in_range & is_final[:, None],
|
||||
)
|
||||
tl.store(
|
||||
SAMPLED_TOP_TOKEN_IDS + req[:, None] * topk_n + k[None, :],
|
||||
tok,
|
||||
mask=in_range & is_final[:, None],
|
||||
)
|
||||
|
||||
|
||||
def finalize_verification_tl(
|
||||
has_sum_smaller_probs: bool,
|
||||
has_confidence: bool,
|
||||
num_steps: int,
|
||||
num_stop_tokens: int,
|
||||
) -> Any:
|
||||
"""Build the kernel that finalizes accepted tokens and packs the per-request payload."""
|
||||
block_rows = 16
|
||||
block_k = 128
|
||||
|
||||
def launch(
|
||||
effective_accept: torch.Tensor,
|
||||
is_final_input: torch.Tensor,
|
||||
effective_token_ids: torch.Tensor,
|
||||
token_logprobs: torch.Tensor,
|
||||
sum_smaller_probs: torch.Tensor,
|
||||
top_logprobs: torch.Tensor,
|
||||
top_token_ids: torch.Tensor,
|
||||
ftp: torch.Tensor,
|
||||
seqlen_q: torch.Tensor,
|
||||
row_to_req: torch.Tensor,
|
||||
is_verify: torch.Tensor,
|
||||
accept_probs: torch.Tensor,
|
||||
confidence_pool: torch.Tensor,
|
||||
mtp_info_read_rows: torch.Tensor,
|
||||
stop_token_ids: torch.Tensor,
|
||||
last_valid_offset: torch.Tensor,
|
||||
num_generated_tokens: torch.Tensor,
|
||||
num_verified_tokens: torch.Tensor,
|
||||
batch_confidence_error: torch.Tensor,
|
||||
emitted_stop_token: torch.Tensor,
|
||||
sampled_token_ids: torch.Tensor,
|
||||
sampled_token_logprobs: torch.Tensor,
|
||||
sampled_sum_smaller_probs: torch.Tensor,
|
||||
sampled_top_logprobs: torch.Tensor,
|
||||
sampled_top_token_ids: torch.Tensor,
|
||||
verified_token_ids: torch.Tensor,
|
||||
verified_token_logprobs: torch.Tensor,
|
||||
verified_sum_smaller_probs: torch.Tensor,
|
||||
verified_top_logprobs: torch.Tensor,
|
||||
verified_top_token_ids: torch.Tensor,
|
||||
) -> None:
|
||||
total_q = effective_accept.shape[0]
|
||||
topk_n = top_logprobs.shape[-1]
|
||||
_finalize_verification_zero_kernel[(1,)](
|
||||
num_generated_tokens,
|
||||
num_verified_tokens,
|
||||
batch_confidence_error,
|
||||
num_steps=num_steps,
|
||||
BLOCK=next_power_of_2(max(num_steps, 1)),
|
||||
)
|
||||
if total_q == 0:
|
||||
return
|
||||
_finalize_verification_kernel[(cdiv(total_q, block_rows),)](
|
||||
effective_accept,
|
||||
is_final_input,
|
||||
effective_token_ids,
|
||||
token_logprobs,
|
||||
sum_smaller_probs,
|
||||
top_logprobs,
|
||||
top_token_ids,
|
||||
ftp,
|
||||
seqlen_q,
|
||||
row_to_req,
|
||||
is_verify,
|
||||
accept_probs,
|
||||
confidence_pool,
|
||||
mtp_info_read_rows,
|
||||
stop_token_ids,
|
||||
last_valid_offset,
|
||||
num_generated_tokens,
|
||||
num_verified_tokens,
|
||||
batch_confidence_error,
|
||||
emitted_stop_token,
|
||||
sampled_token_ids,
|
||||
sampled_token_logprobs,
|
||||
sampled_sum_smaller_probs,
|
||||
sampled_top_logprobs,
|
||||
sampled_top_token_ids,
|
||||
verified_token_ids,
|
||||
verified_token_logprobs,
|
||||
verified_sum_smaller_probs,
|
||||
verified_top_logprobs,
|
||||
verified_top_token_ids,
|
||||
total_q,
|
||||
topk_n,
|
||||
top_logprobs.stride(0),
|
||||
top_token_ids.stride(0),
|
||||
HAS_SUM_SMALLER_PROBS=has_sum_smaller_probs,
|
||||
HAS_CONFIDENCE=has_confidence,
|
||||
NUM_STEPS=num_steps,
|
||||
NUM_STOP_TOKENS=num_stop_tokens,
|
||||
BLOCK_ROWS=block_rows,
|
||||
BLOCK_K=block_k,
|
||||
num_warps=8,
|
||||
)
|
||||
|
||||
return launch
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _remap_draft_topk_kernel(
|
||||
DRAFT_TOPK_INDEX,
|
||||
DRAFT_TOPK_LOGPROBS,
|
||||
DRAFT_TOKEN_IDS_POOL,
|
||||
TARGET_TOPK_INDEX,
|
||||
READ_ROWS,
|
||||
ALIGNED_LOGPROBS,
|
||||
DRAFT_LOCAL_IDS,
|
||||
IDENTITY_ROWS,
|
||||
target_stride,
|
||||
aligned_stride,
|
||||
force_topk: tl.constexpr,
|
||||
BLOCK_T: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
tblock = tl.program_id(1)
|
||||
source_row = tl.load(READ_ROWS + row)
|
||||
draft_token_id = tl.load(DRAFT_TOKEN_IDS_POOL + source_row).to(tl.uint32)
|
||||
|
||||
# One program owns a slice of this row's target indices and scans the whole draft slice.
|
||||
tj = tblock * BLOCK_T + tl.arange(0, BLOCK_T)
|
||||
t_valid = tj < force_topk
|
||||
q_index = tl.load(TARGET_TOPK_INDEX + row * target_stride + tj, mask=t_valid, other=0).to(tl.uint32)
|
||||
|
||||
if tblock == 0:
|
||||
tl.store(IDENTITY_ROWS + row, row)
|
||||
|
||||
# Where in the target slice the draft token itself sits. This is independent of the lookup
|
||||
# below: it locates the draft token, not some other token's draft logprob.
|
||||
is_draft = t_valid & (q_index == draft_token_id)
|
||||
tl.atomic_max(DRAFT_LOCAL_IDS + row, tl.max(tl.where(is_draft, tj, -1), axis=0))
|
||||
|
||||
# The TileLang kernel builds a shared-memory lookup table over vocabulary-sized windows and
|
||||
# reads it back. The same answer comes from matching each target index against the draft
|
||||
# indices directly: a draft entry is unique per token id, so at most one candidate matches and
|
||||
# taking the maximum over that block is well defined. This trades a table build for matched
|
||||
# scans, which is only worth it because `force_topk` is at most a few thousand.
|
||||
acc = tl.full([BLOCK_T], float("-inf"), tl.float32)
|
||||
for dstart in range(0, force_topk, BLOCK_D):
|
||||
dj = dstart + tl.arange(0, BLOCK_D)
|
||||
d_valid = dj < force_topk
|
||||
d_index = tl.load(DRAFT_TOPK_INDEX + source_row * force_topk + dj, mask=d_valid, other=0).to(tl.uint32)
|
||||
d_value = tl.load(DRAFT_TOPK_LOGPROBS + source_row * force_topk + dj, mask=d_valid, other=0.0)
|
||||
match = t_valid[:, None] & d_valid[None, :] & (q_index[:, None] == d_index[None, :])
|
||||
acc = tl.maximum(acc, tl.max(tl.where(match, d_value, float("-inf")), axis=1))
|
||||
|
||||
tl.store(ALIGNED_LOGPROBS + row * aligned_stride + tj, acc, mask=t_valid)
|
||||
|
||||
|
||||
def remap_draft_topk_tl(force_topk: int, vocab_size: int) -> Any:
|
||||
"""Build the kernel that aligns the previous draft's top-k to the current target top-k."""
|
||||
block_t = 128
|
||||
block_d = 128
|
||||
|
||||
def launch(
|
||||
draft_topk_index: torch.Tensor,
|
||||
draft_topk_logprobs: torch.Tensor,
|
||||
draft_token_ids_pool: torch.Tensor,
|
||||
target_topk_index: torch.Tensor,
|
||||
read_rows: torch.Tensor,
|
||||
aligned_logprobs: torch.Tensor,
|
||||
draft_local_ids: torch.Tensor,
|
||||
identity_rows: torch.Tensor,
|
||||
) -> None:
|
||||
total_q = target_topk_index.shape[0]
|
||||
if total_q == 0 or force_topk == 0:
|
||||
return
|
||||
# `atomic_max` needs a starting value below every real position, and "not found" is -1.
|
||||
draft_local_ids.fill_(-1)
|
||||
_remap_draft_topk_kernel[(total_q, cdiv(force_topk, block_t))](
|
||||
draft_topk_index,
|
||||
draft_topk_logprobs,
|
||||
draft_token_ids_pool,
|
||||
target_topk_index,
|
||||
read_rows,
|
||||
aligned_logprobs,
|
||||
draft_local_ids,
|
||||
identity_rows,
|
||||
target_topk_index.stride(0),
|
||||
aligned_logprobs.stride(0),
|
||||
force_topk=force_topk,
|
||||
BLOCK_T=block_t,
|
||||
BLOCK_D=block_d,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
|
||||
return launch
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user