mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
feat(triton): port eight P1 kernels, all A/B identical to TileLang
gather, gather_input_ids_zos, gather_engram_tokens_zos, expand_swa_index, kv_index_cat, convert_kv_index, load_kv_index, store_mtp_info. 45 cases, every one bitwise identical to the TileLang baseline through the shared interface function. `scripts/h20_ops_ab.py` swaps `ops.<name>_tl` between the two backends and runs the public call once per side, so the comparison exercises the interface layer's shape handling too, not just the kernel. `scripts/ab_cases.py` holds the cases: real engine widths and dtypes, negative indices for gather, -1 padding for the index kernels, page-boundary lengths, the slot-0 no-write path, and the engram image-sentinel rule. Five Triton semantics traps worth writing down, all documented in the plan: - `tl.where` evaluates both arms, so `tl.where(i == 0, 0, LOAD(i - 1))` really reads `LOAD[-1]`; that was an illegal memory access, fixed by clamping the index instead. - A scalar pointer cannot carry a block-shaped mask; it has to be broadcast with `+ tl.zeros((BLOCK,), tl.int32)` first. - Two arms of one kernel cannot use different tile shapes, which is why `kv_index_cat` became two launches rather than one fused kernel with a column-block branch. - Triton's `data_ptr()` already includes a view's storage offset, so passing `out[:, w:w+k]` and adding `w` again double-counts -- that showed up as 0x01010101 in the second half. - Test-case capacity matters: `store_mtp_info` addresses `slot * nextn + step`, so the pool needs `(max_slot + 1) * nextn` rows; too few rows let both backends write out of bounds and disagree.
This commit is contained in:
loopyt
parent
4a5606ad2d
commit
a6c3e69e67
@ -104,4 +104,29 @@
|
||||
|
||||
## 七、进度
|
||||
|
||||
见文末「执行日志」。
|
||||
| 阶段 | 状态 | 说明 |
|
||||
|---|---|---|
|
||||
| P0 基础设施 | ✅ | 后端选择 + fallback + A/B 骨架 |
|
||||
| P1 轻量 kernel | 🔄 8/13 | gather、gather_input_ids_zos、gather_engram_tokens_zos、expand_swa_index、kv_index_cat、convert_kv_index、load_kv_index、store_mtp_info |
|
||||
| P2 KV cache | ⬜ | |
|
||||
| P3 MTP 链路 | ⬜ | |
|
||||
| P4 采样 | ⬜ | |
|
||||
| P5 端到端 | ⬜ | |
|
||||
|
||||
A/B 现状：**45 个用例全过**（`python scripts/h20_ops_ab.py`）。
|
||||
|
||||
### 过程中踩到的 Triton 语义坑（都已修，可复用）
|
||||
|
||||
1. **`tl.where` 两侧都会求值**。`tl.where(i == 0, 0, LOAD(i - 1))` 在 `i == 0` 时真的会读
|
||||
`LOAD[-1]` → illegal memory access。凡是「边界索引用 where 兜住」的写法都要先
|
||||
`tl.maximum(idx, 0)` 钳制。
|
||||
2. **标量指针不能配 block mask**。`tl.load(PTR + scalar, mask=<block mask>)` 报
|
||||
`Mask argument cannot be block type if pointer argument is not a block`。要
|
||||
`PTR + scalar + tl.zeros((BLOCK,), tl.int32)` 广播成 block。
|
||||
3. **一个 kernel 里两个分支用不同 tile 形状**编不过。`kv_index_cat` 拆成两个 kernel 分别
|
||||
launch（同一 stream，语义不变）。
|
||||
4. **视图的 storage offset 已含在 `data_ptr()` 里**。给 kernel 传
|
||||
`out[:, w:w+k]` 后再加一次 `w` 就重复偏移了 —— 表现为后半段全是 `0x01010101` 之类的
|
||||
uninitialized 内存。
|
||||
5. **测试用例的容量也要算对**：`store_mtp_info` 的 pool 是 `slot * nextn + step` 寻址，
|
||||
行数必须是 `(max_slot + 1) * nextn`，否则两个后端一起写越界，表现成互相不一致。
|
||||
|
||||
@ -0,0 +1,69 @@
|
||||
"""Logical top-k token indices to physical cache slots, in place, in Triton."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import next_power_of_2
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _convert_kv_index_kernel(
|
||||
TOPK_IDXS,
|
||||
CONTEXT_LENS,
|
||||
BLOCK_TABLES,
|
||||
s,
|
||||
k: tl.constexpr,
|
||||
n: tl.constexpr,
|
||||
block_size: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
):
|
||||
sid = tl.program_id(0)
|
||||
bid = tl.program_id(1)
|
||||
|
||||
base = (bid * s + sid) * k
|
||||
offs = tl.arange(0, BLOCK_K)
|
||||
context_len = tl.load(CONTEXT_LENS + bid * s + sid)
|
||||
|
||||
idx = tl.load(TOPK_IDXS + base + offs, mask=offs < k, other=-1)
|
||||
# Entries past the context length are padding; the reference marks them -1 rather than
|
||||
# translating them, so they must not go through the block-table lookup.
|
||||
live = (offs < k) & (offs < context_len)
|
||||
page = tl.load(BLOCK_TABLES + bid * n + idx // block_size, mask=live, other=0)
|
||||
physical = page * block_size + idx % block_size
|
||||
|
||||
tl.store(TOPK_IDXS + base + offs, tl.where(live, physical, -1), mask=offs < k)
|
||||
|
||||
|
||||
def convert_kv_index_tl(
|
||||
s: int,
|
||||
k: int,
|
||||
n: int,
|
||||
block_size: int,
|
||||
) -> Any:
|
||||
"""Build the kernel for a fixed query length, top-k and page size."""
|
||||
block_k = next_power_of_2(k)
|
||||
|
||||
def launch(
|
||||
topk_idxs: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
) -> None:
|
||||
b = topk_idxs.shape[0]
|
||||
if b == 0 or s == 0:
|
||||
return
|
||||
_convert_kv_index_kernel[(s, b)](
|
||||
topk_idxs,
|
||||
context_lens,
|
||||
block_tables,
|
||||
s,
|
||||
k=k,
|
||||
n=n,
|
||||
block_size=block_size,
|
||||
BLOCK_K=block_k,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -0,0 +1,134 @@
|
||||
"""Per-request n-gram token windows for engram lookup, in Triton."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import next_power_of_2
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _gather_engram_tokens_zos_kernel(
|
||||
WIN_COMPRESSED,
|
||||
CURR_FRESH,
|
||||
PREV_SAMPLED,
|
||||
PREV_PREDICT,
|
||||
PREV_NUM_ACCEPTED,
|
||||
REMAP,
|
||||
LOOKUP_TABLE,
|
||||
SEQLENS_Q,
|
||||
CU_SEQLENS_Q,
|
||||
OUT,
|
||||
max_seqlen_q: tl.constexpr,
|
||||
max_ngram_size: tl.constexpr,
|
||||
window_size: tl.constexpr,
|
||||
engram_compressed_pad_token: tl.constexpr,
|
||||
image_sentinel: tl.constexpr,
|
||||
BLOCK_Q: tl.constexpr,
|
||||
):
|
||||
seq_idx = tl.program_id(0)
|
||||
sq = tl.load(SEQLENS_Q + seq_idx)
|
||||
# `tl.where` evaluates both arms, so the index is clamped: at seq_idx == 0 the other arm
|
||||
# would read CU_SEQLENS_Q[-1].
|
||||
q_base = tl.where(seq_idx == 0, 0, tl.load(CU_SEQLENS_Q + tl.maximum(seq_idx - 1, 0)))
|
||||
pidx = tl.load(REMAP + seq_idx)
|
||||
# A remapped request reuses `nextn - accepted` tokens it already produced, so its window starts
|
||||
# that much later; a fresh request starts at the beginning of the window.
|
||||
nr = tl.where(pidx >= 0, (max_seqlen_q - 1) - tl.load(PREV_NUM_ACCEPTED + pidx), max_seqlen_q - 1)
|
||||
base = window_size - nr
|
||||
|
||||
offs = tl.arange(0, BLOCK_Q)
|
||||
active = offs < sq
|
||||
|
||||
# Once an image sentinel is seen the rest of the window is padding, so the r loop carries a
|
||||
# per-lane `hit` flag forward instead of re-scanning.
|
||||
hit = tl.zeros((BLOCK_Q,), dtype=tl.int1)
|
||||
for r in tl.static_range(max_ngram_size):
|
||||
from_current = offs >= r
|
||||
c = offs - r
|
||||
|
||||
token_cur = tl.load(
|
||||
PREV_SAMPLED + pidx + tl.zeros((BLOCK_Q,), tl.int32),
|
||||
mask=active & from_current & (pidx >= 0),
|
||||
other=0,
|
||||
)
|
||||
token_draft = tl.load(
|
||||
PREV_PREDICT + pidx * (max_seqlen_q - 1) + (c - 1),
|
||||
mask=active & from_current & (pidx >= 0) & (c > 0),
|
||||
other=0,
|
||||
)
|
||||
token_fresh = tl.load(
|
||||
CURR_FRESH + seq_idx * max_seqlen_q + c,
|
||||
mask=active & from_current & (pidx < 0),
|
||||
other=0,
|
||||
)
|
||||
token_window = tl.load(
|
||||
WIN_COMPRESSED + seq_idx * window_size + base + offs - r,
|
||||
mask=active & ~from_current,
|
||||
other=0,
|
||||
)
|
||||
|
||||
raw = tl.where(
|
||||
from_current,
|
||||
tl.where(pidx < 0, token_fresh, tl.where(c == 0, token_cur, token_draft)),
|
||||
token_window,
|
||||
)
|
||||
# Fresh tokens are already vocabulary ids; replayed ones are compressed and need the map.
|
||||
mapped = tl.load(LOOKUP_TABLE + raw, mask=active & from_current & (pidx >= 0), other=0)
|
||||
token = tl.where(from_current & (pidx >= 0), mapped, raw)
|
||||
|
||||
is_image = token == image_sentinel
|
||||
value = tl.where(hit | is_image, engram_compressed_pad_token, token)
|
||||
tl.store(OUT + (q_base + offs) * max_ngram_size + r, value, mask=active)
|
||||
hit = hit | is_image
|
||||
|
||||
|
||||
def gather_engram_tokens_zos_tl(
|
||||
max_seqlen_q: int,
|
||||
max_ngram_size: int,
|
||||
engram_compressed_pad_token: int,
|
||||
image_sentinel: int,
|
||||
) -> Any:
|
||||
"""Build the kernel for a fixed draft depth and n-gram size."""
|
||||
block_q = next_power_of_2(max_seqlen_q)
|
||||
nextn = max_seqlen_q - 1
|
||||
window_size = nextn + max_ngram_size - 1
|
||||
|
||||
def launch(
|
||||
win_compressed: torch.Tensor,
|
||||
curr_fresh: torch.Tensor,
|
||||
prev_sampled_token_ids: torch.Tensor,
|
||||
prev_predict_token_ids: torch.Tensor,
|
||||
prev_num_accepted_tokens: torch.Tensor,
|
||||
remap: torch.Tensor,
|
||||
lookup_table: torch.Tensor,
|
||||
seqlens_q: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
out_ngram_token_ids: torch.Tensor,
|
||||
) -> None:
|
||||
n_seq = curr_fresh.shape[0]
|
||||
if n_seq == 0:
|
||||
return
|
||||
_gather_engram_tokens_zos_kernel[(n_seq,)](
|
||||
win_compressed,
|
||||
curr_fresh,
|
||||
prev_sampled_token_ids,
|
||||
prev_predict_token_ids,
|
||||
prev_num_accepted_tokens,
|
||||
remap,
|
||||
lookup_table,
|
||||
seqlens_q,
|
||||
cu_seqlens_q,
|
||||
out_ngram_token_ids,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_ngram_size=max_ngram_size,
|
||||
window_size=window_size,
|
||||
engram_compressed_pad_token=engram_compressed_pad_token,
|
||||
image_sentinel=image_sentinel,
|
||||
BLOCK_Q=block_q,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -0,0 +1,79 @@
|
||||
"""Per-request row expansion of previous-step token ids, in Triton."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import next_power_of_2
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _gather_input_ids_zos_kernel(
|
||||
FRESH_TOKENS,
|
||||
PREV_SAMPLED,
|
||||
PREV_PREDICT,
|
||||
REMAP,
|
||||
SEQLENS_Q,
|
||||
CU_SEQLENS_Q,
|
||||
OUT,
|
||||
max_seqlen_q: tl.constexpr,
|
||||
BLOCK_Q: tl.constexpr,
|
||||
):
|
||||
seq_idx = tl.program_id(0)
|
||||
sq = tl.load(SEQLENS_Q + seq_idx)
|
||||
# `tl.where` evaluates both arms, so the index is clamped: at seq_idx == 0 the other arm
|
||||
# would read CU_SEQLENS_Q[-1].
|
||||
q_base = tl.where(seq_idx == 0, 0, tl.load(CU_SEQLENS_Q + tl.maximum(seq_idx - 1, 0)))
|
||||
pidx = tl.load(REMAP + seq_idx)
|
||||
|
||||
offs = tl.arange(0, BLOCK_Q)
|
||||
active = offs < sq
|
||||
|
||||
# Fresh prompt tokens come from this request's own row; a remapped request instead replays the
|
||||
# previous step's sampled token for position 0 and its draft predictions after that.
|
||||
fresh = tl.load(FRESH_TOKENS + seq_idx * max_seqlen_q + offs, mask=active, other=0)
|
||||
# `pidx` is a scalar, so the pointer has to be broadcast to a block before it can carry a
|
||||
# block-shaped mask.
|
||||
prev_sampled = tl.load(PREV_SAMPLED + pidx + tl.zeros((BLOCK_Q,), tl.int32), mask=active & (pidx >= 0) & (offs == 0), other=0)
|
||||
prev_predict = tl.load(
|
||||
PREV_PREDICT + pidx * (max_seqlen_q - 1) + (offs - 1),
|
||||
mask=active & (pidx >= 0) & (offs > 0),
|
||||
other=0,
|
||||
)
|
||||
|
||||
value = tl.where(pidx < 0, fresh, tl.where(offs == 0, prev_sampled, prev_predict))
|
||||
tl.store(OUT + q_base + offs, value, mask=active)
|
||||
|
||||
|
||||
def gather_input_ids_zos_tl(max_seqlen_q: int) -> Any:
|
||||
"""Build the kernel for a fixed draft depth; the request count stays a runtime length."""
|
||||
block_q = next_power_of_2(max_seqlen_q)
|
||||
|
||||
def launch(
|
||||
fresh_tokens: torch.Tensor,
|
||||
prev_sampled_token_ids: torch.Tensor,
|
||||
prev_predict_token_ids: torch.Tensor,
|
||||
remap: torch.Tensor,
|
||||
seqlens_q: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
out_input_ids: torch.Tensor,
|
||||
) -> None:
|
||||
n_seq = fresh_tokens.shape[0]
|
||||
if n_seq == 0:
|
||||
return
|
||||
_gather_input_ids_zos_kernel[(n_seq,)](
|
||||
fresh_tokens,
|
||||
prev_sampled_token_ids,
|
||||
prev_predict_token_ids,
|
||||
remap,
|
||||
seqlens_q,
|
||||
cu_seqlens_q,
|
||||
out_input_ids,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
BLOCK_Q=block_q,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -0,0 +1,80 @@
|
||||
"""Expand per-row (offset, length) into cache index rows, in Triton."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _load_kv_index_kernel(
|
||||
OFFSET_LEN,
|
||||
KV_INDEX,
|
||||
FLATTEN_BLOCK_TABLES,
|
||||
index_length: tl.constexpr,
|
||||
block_size: tl.constexpr,
|
||||
has_block_tables: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
col_block = tl.program_id(1)
|
||||
|
||||
offset = tl.load(OFFSET_LEN + row * 2 + 0)
|
||||
length = tl.load(OFFSET_LEN + row * 2 + 1)
|
||||
|
||||
cols = col_block * BLOCK + tl.arange(0, BLOCK)
|
||||
live = cols < length
|
||||
|
||||
if has_block_tables:
|
||||
# Contiguous positions, then mapped through the page table into a physical slot.
|
||||
pos = offset + cols
|
||||
page = tl.load(FLATTEN_BLOCK_TABLES + pos // block_size, mask=live, other=0)
|
||||
index = page * block_size + pos % block_size
|
||||
elif block_size == 0:
|
||||
index = offset + cols
|
||||
else:
|
||||
# Ring buffer confined to the block that `offset` falls in.
|
||||
index = offset // block_size * block_size + (offset + cols) % block_size
|
||||
|
||||
tl.store(
|
||||
KV_INDEX + row * index_length + cols,
|
||||
tl.where(live, index, -1),
|
||||
mask=cols < index_length,
|
||||
)
|
||||
|
||||
|
||||
def load_kv_index_tl(
|
||||
index_length: int,
|
||||
block_size: int,
|
||||
has_block_tables: bool,
|
||||
) -> Any:
|
||||
"""Build the kernel for a fixed row length and addressing mode."""
|
||||
block = min(2048, next_power_of_2(index_length))
|
||||
|
||||
def launch(
|
||||
kv_index_offset_len: torch.Tensor,
|
||||
kv_index: torch.Tensor,
|
||||
flatten_block_tables: torch.Tensor,
|
||||
) -> None:
|
||||
n = kv_index_offset_len.shape[0]
|
||||
if n == 0 or index_length == 0:
|
||||
return
|
||||
if flatten_block_tables is None:
|
||||
# The contiguous mode takes `None` here, and Triton needs a real pointer even when the
|
||||
# branch that would read it is compiled out.
|
||||
flatten_block_tables = kv_index_offset_len
|
||||
_load_kv_index_kernel[(n, cdiv(index_length, block))](
|
||||
kv_index_offset_len,
|
||||
kv_index,
|
||||
flatten_block_tables,
|
||||
index_length=index_length,
|
||||
block_size=block_size,
|
||||
has_block_tables=has_block_tables,
|
||||
BLOCK=block,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -0,0 +1,112 @@
|
||||
"""Concatenate the window index and the shifted compressed index, in Triton.
|
||||
|
||||
Two launches rather than one fused kernel: the two halves have different widths and different
|
||||
element rules (verbatim copy vs. conditional rebase), so keeping them separate avoids a
|
||||
compile-time branch whose two arms would need the same tile shape.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _copy_window_kernel(
|
||||
SRC,
|
||||
DST,
|
||||
total_q,
|
||||
window_size: tl.constexpr,
|
||||
out_stride: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_W: tl.constexpr,
|
||||
):
|
||||
row_block = tl.program_id(0)
|
||||
col_block = tl.program_id(1)
|
||||
|
||||
rows = row_block * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
cols = col_block * BLOCK_W + tl.arange(0, BLOCK_W)
|
||||
mask = (rows[:, None] < total_q) & (cols[None, :] < window_size)
|
||||
|
||||
value = tl.load(SRC + rows[:, None] * window_size + cols[None, :], mask=mask, other=0)
|
||||
tl.store(DST + rows[:, None] * out_stride + cols[None, :], value, mask=mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _shift_compressed_kernel(
|
||||
SRC,
|
||||
DST,
|
||||
total_q,
|
||||
topk_tokens: tl.constexpr,
|
||||
out_stride: tl.constexpr,
|
||||
shift: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
):
|
||||
row_block = tl.program_id(0)
|
||||
col_block = tl.program_id(1)
|
||||
|
||||
rows = row_block * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
cols = col_block * BLOCK_K + tl.arange(0, BLOCK_K)
|
||||
mask = (rows[:, None] < total_q) & (cols[None, :] < topk_tokens)
|
||||
|
||||
value = tl.load(SRC + rows[:, None] * topk_tokens + cols[None, :], mask=mask, other=-1)
|
||||
# `kv_ca_index_out` is already a view whose base sits at column `window_size`, so only the
|
||||
# row stride is applied here. Padding stays -1; shifting it would make it a plausible slot.
|
||||
tl.store(
|
||||
DST + rows[:, None] * out_stride + cols[None, :],
|
||||
tl.where(value == -1, -1, value + shift),
|
||||
mask=mask,
|
||||
)
|
||||
|
||||
|
||||
def kv_index_cat_tl(
|
||||
dtype: str,
|
||||
window_size: int,
|
||||
topk_tokens: int,
|
||||
out_stride: int,
|
||||
) -> Any:
|
||||
"""Build the kernel for a fixed window width, top-k and output row stride."""
|
||||
del dtype # the interface layer only ever passes int32
|
||||
|
||||
block_m = 8
|
||||
block_w = min(1024, next_power_of_2(window_size))
|
||||
block_k = min(1024, next_power_of_2(topk_tokens))
|
||||
|
||||
def launch(
|
||||
kv_index: torch.Tensor,
|
||||
kv_ca_index: torch.Tensor,
|
||||
kv_index_out: torch.Tensor,
|
||||
kv_ca_index_out: torch.Tensor,
|
||||
total_k_swa: int,
|
||||
total_k: int,
|
||||
) -> None:
|
||||
total_q = kv_index.shape[0]
|
||||
if total_q == 0:
|
||||
return
|
||||
_copy_window_kernel[(cdiv(total_q, block_m), cdiv(window_size, block_w))](
|
||||
kv_index,
|
||||
kv_index_out,
|
||||
total_q,
|
||||
window_size=window_size,
|
||||
out_stride=out_stride,
|
||||
BLOCK_M=block_m,
|
||||
BLOCK_W=block_w,
|
||||
num_warps=4,
|
||||
)
|
||||
_shift_compressed_kernel[(cdiv(total_q, block_m), cdiv(topk_tokens, block_k))](
|
||||
kv_ca_index,
|
||||
kv_ca_index_out,
|
||||
total_q,
|
||||
topk_tokens=topk_tokens,
|
||||
out_stride=out_stride,
|
||||
shift=total_k + total_k_swa,
|
||||
BLOCK_M=block_m,
|
||||
BLOCK_K=block_k,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -0,0 +1,95 @@
|
||||
"""Copy the draft head's per-step outputs into the MTP pool, in Triton."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _store_mtp_info_kernel(
|
||||
WRITE_SLOTS,
|
||||
TOKEN_IDS,
|
||||
LOGPROBS,
|
||||
LOGPROBS_INDEX,
|
||||
CONFIDENCE,
|
||||
TOKEN_IDS_POOL,
|
||||
LOGPROBS_POOL,
|
||||
LOGPROBS_INDEX_POOL,
|
||||
CONFIDENCE_POOL,
|
||||
nextn: tl.constexpr,
|
||||
vocab_size: tl.constexpr,
|
||||
has_index: tl.constexpr,
|
||||
has_confidence: tl.constexpr,
|
||||
BLOCK_V: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0)
|
||||
col_block = tl.program_id(1)
|
||||
req = pid // nextn
|
||||
step = pid % nextn
|
||||
|
||||
# Slot 0 means "nothing to write" -- the pool row is pre-seeded by the caller.
|
||||
slot = tl.load(WRITE_SLOTS + req)
|
||||
if slot != 0:
|
||||
cols = col_block * BLOCK_V + tl.arange(0, BLOCK_V)
|
||||
mask = cols < vocab_size
|
||||
src = (req * nextn + step) * vocab_size + cols
|
||||
dst = (slot * nextn + step) * vocab_size + cols
|
||||
|
||||
logprobs = tl.load(LOGPROBS + src, mask=mask, other=0.0)
|
||||
tl.store(LOGPROBS_POOL + dst, logprobs, mask=mask)
|
||||
if has_index:
|
||||
logprobs_index = tl.load(LOGPROBS_INDEX + src, mask=mask, other=0)
|
||||
tl.store(LOGPROBS_INDEX_POOL + dst, logprobs_index, mask=mask)
|
||||
|
||||
# The token ids and confidences are per (request, step), not per vocabulary entry, so only
|
||||
# the first column block writes them.
|
||||
if (col_block == 0) & (step == 0):
|
||||
for i in tl.static_range(nextn):
|
||||
token = tl.load(TOKEN_IDS + req * nextn + i)
|
||||
tl.store(TOKEN_IDS_POOL + slot * nextn + i, token)
|
||||
if has_confidence:
|
||||
conf = tl.load(CONFIDENCE + req * nextn + i)
|
||||
tl.store(CONFIDENCE_POOL + slot * nextn + i, conf)
|
||||
|
||||
|
||||
def store_mtp_info_tl(nextn: int, vocab_size: int, has_index: bool, has_confidence: bool) -> Any:
|
||||
"""Build the kernel for a fixed draft depth and vocabulary size."""
|
||||
block_v = min(next_power_of_2(vocab_size), 4096)
|
||||
|
||||
def launch(
|
||||
write_slots: torch.Tensor,
|
||||
token_ids: torch.Tensor,
|
||||
logprobs: torch.Tensor,
|
||||
logprobs_index: torch.Tensor,
|
||||
confidence: torch.Tensor,
|
||||
token_ids_pool: torch.Tensor,
|
||||
logprobs_pool: torch.Tensor,
|
||||
logprobs_index_pool: torch.Tensor,
|
||||
confidence_pool: torch.Tensor,
|
||||
) -> None:
|
||||
req_cnt = write_slots.shape[0]
|
||||
if req_cnt == 0:
|
||||
return
|
||||
_store_mtp_info_kernel[(req_cnt * nextn, cdiv(vocab_size, block_v))](
|
||||
write_slots,
|
||||
token_ids,
|
||||
logprobs,
|
||||
logprobs_index,
|
||||
confidence,
|
||||
token_ids_pool,
|
||||
logprobs_pool,
|
||||
logprobs_index_pool,
|
||||
confidence_pool,
|
||||
nextn=nextn,
|
||||
vocab_size=vocab_size,
|
||||
has_index=has_index,
|
||||
has_confidence=has_confidence,
|
||||
BLOCK_V=block_v,
|
||||
num_warps=8,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -0,0 +1,58 @@
|
||||
"""Expand four per-row pivots into a full sliding-window index row, in Triton."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _expand_swa_index_kernel(
|
||||
PIVOTS,
|
||||
OUTPUT,
|
||||
n,
|
||||
width: tl.constexpr,
|
||||
BLOCK_W: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
col_block = tl.program_id(1)
|
||||
|
||||
if row < n:
|
||||
offset1 = tl.load(PIVOTS + row * 4 + 0)
|
||||
len1 = tl.load(PIVOTS + row * 4 + 1)
|
||||
offset2 = tl.load(PIVOTS + row * 4 + 2)
|
||||
len2 = tl.load(PIVOTS + row * 4 + 3)
|
||||
|
||||
cols = col_block * BLOCK_W + tl.arange(0, BLOCK_W)
|
||||
# Both spans are contiguous ranges; past the second one the column is padding.
|
||||
value = tl.where(
|
||||
cols < len1,
|
||||
offset1 + cols,
|
||||
tl.where(cols < len1 + len2, offset2 - len1 + cols, -1),
|
||||
)
|
||||
tl.store(OUTPUT + row * width + cols, value, mask=cols < width)
|
||||
|
||||
|
||||
def expand_swa_index_tl(width: int) -> Any:
|
||||
"""Build the kernel for a fixed row width."""
|
||||
# A row is written by BLOCK_W-wide blocks; 1024 int32 per block is one coalesced pass with a
|
||||
# comfortable register budget, and narrower rows simply use the whole width in one block.
|
||||
block_w = min(1024, next_power_of_2(width))
|
||||
|
||||
def launch(pivots: torch.Tensor, output: torch.Tensor) -> None:
|
||||
n = pivots.shape[0]
|
||||
if n == 0:
|
||||
return
|
||||
_expand_swa_index_kernel[(n, cdiv(width, block_w))](
|
||||
pivots,
|
||||
output,
|
||||
n,
|
||||
width=width,
|
||||
BLOCK_W=block_w,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -15,7 +15,14 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.flash_mla_utils import convert_kv_index
|
||||
from deepinfer.model_executor.layers.tilelang_ops.gather import gather
|
||||
from deepinfer.model_executor.layers.tilelang_ops.gather_engram_tokens_zos import gather_engram_tokens_zos
|
||||
from deepinfer.model_executor.layers.tilelang_ops.gather_input_ids_zos import gather_input_ids_zos
|
||||
from deepinfer.model_executor.layers.tilelang_ops.get_kv_index import get_kv_ca_index_decode, get_kv_ca_index_prefill
|
||||
from deepinfer.model_executor.layers.tilelang_ops.kv_index_cat import kv_index_cat
|
||||
from deepinfer.model_executor.layers.tilelang_ops.store_mtp_info import store_mtp_info
|
||||
from deepinfer.model_executor.layers.tilelang_ops.swa_index import expand_swa_index
|
||||
|
||||
|
||||
def _cases_gather() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
@ -48,6 +55,249 @@ def _cases_gather() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
]
|
||||
|
||||
|
||||
def _cases_gather_input_ids() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
"""Draft-token replay: remapped requests replay previous ids, fresh ones use their own row."""
|
||||
|
||||
def make(n_seq: int, max_seqlen_q: int, prev_n_seq: int, mix: bool):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
nextn = max_seqlen_q - 1
|
||||
fresh = torch.randint(0, 1000, (n_seq, max_seqlen_q), device="cuda", dtype=torch.int32)
|
||||
prev_sampled = torch.randint(0, 1000, (prev_n_seq,), device="cuda", dtype=torch.int32)
|
||||
prev_predict = torch.randint(0, 1000, (prev_n_seq, nextn), device="cuda", dtype=torch.int32)
|
||||
if mix:
|
||||
remap = torch.tensor([-1 if i % 2 == 0 else i % prev_n_seq for i in range(n_seq)], device="cuda", dtype=torch.int32)
|
||||
else:
|
||||
remap = torch.full((n_seq,), -1, device="cuda", dtype=torch.int32)
|
||||
seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32)
|
||||
cu = torch.cumsum(seqlens, 0, dtype=torch.int32)
|
||||
out = torch.full((int(cu[-1].item()),), -7, device="cuda", dtype=torch.int32)
|
||||
gather_input_ids_zos(fresh, prev_sampled, prev_predict, remap, seqlens, cu, out, max_seqlen_q=max_seqlen_q)
|
||||
return out
|
||||
|
||||
return run
|
||||
|
||||
return [
|
||||
("n8_nextn5_mixed", True, make(8, 6, 8, True)),
|
||||
("n8_nextn5_all_fresh", True, make(8, 6, 8, False)),
|
||||
("n1_nextn1_mixed", True, make(1, 2, 3, True)),
|
||||
("n64_nextn5_mixed", True, make(64, 6, 64, True)),
|
||||
("n5_nextn2_full_len", True, make(5, 3, 5, True)),
|
||||
]
|
||||
|
||||
|
||||
def _cases_gather_engram() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
"""Engram n-gram windows, including the image-sentinel padding rule."""
|
||||
|
||||
def make(n_seq: int, max_seqlen_q: int, max_ngram: int, vocab: int, with_image: bool):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
nextn = max_seqlen_q - 1
|
||||
window = nextn + max_ngram - 1
|
||||
pad_token = 7
|
||||
sentinel = -1
|
||||
win = torch.randint(0, 50, (n_seq, window), device="cuda", dtype=torch.int32)
|
||||
if with_image:
|
||||
win[:, window // 2] = sentinel
|
||||
fresh = torch.randint(0, 50, (n_seq, max_seqlen_q), device="cuda", dtype=torch.int32)
|
||||
if with_image:
|
||||
fresh[:, -1] = sentinel
|
||||
prev_sampled = torch.randint(0, vocab, (n_seq,), device="cuda", dtype=torch.int32)
|
||||
prev_predict = torch.randint(0, vocab, (n_seq, nextn), device="cuda", dtype=torch.int32)
|
||||
prev_accepted = torch.randint(0, nextn + 1, (n_seq,), device="cuda", dtype=torch.int32)
|
||||
remap = torch.tensor([-1 if i % 2 == 0 else i for i in range(n_seq)], device="cuda", dtype=torch.int32)
|
||||
lookup = torch.arange(vocab, device="cuda", dtype=torch.int32) % 13
|
||||
seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32)
|
||||
cu = torch.cumsum(seqlens, 0, dtype=torch.int32)
|
||||
out = torch.full((int(cu[-1].item()), max_ngram), -7, device="cuda", dtype=torch.int32)
|
||||
gather_engram_tokens_zos(
|
||||
win,
|
||||
fresh,
|
||||
prev_sampled,
|
||||
prev_predict,
|
||||
prev_accepted,
|
||||
remap,
|
||||
lookup,
|
||||
seqlens,
|
||||
cu,
|
||||
out,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_ngram_size=max_ngram,
|
||||
engram_compressed_pad_token=pad_token,
|
||||
image_sentinel=sentinel,
|
||||
)
|
||||
return out
|
||||
|
||||
return run
|
||||
|
||||
return [
|
||||
("n8_nextn5_ngram3", True, make(8, 6, 3, 64, False)),
|
||||
("n8_nextn5_ngram1", True, make(8, 6, 1, 64, False)),
|
||||
("n8_nextn5_ngram3_image", True, make(8, 6, 3, 64, True)),
|
||||
("n1_nextn1_ngram2", True, make(1, 2, 2, 64, False)),
|
||||
]
|
||||
|
||||
|
||||
def _cases_swa_index() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
"""Sliding-window index expansion for the widths the engine uses."""
|
||||
|
||||
def make(n: int, width: int, overflow: bool):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
pivots = torch.zeros((n, 4), device="cuda", dtype=torch.int32)
|
||||
pivots[:, 0] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
|
||||
if overflow:
|
||||
# len1 + len2 can exceed width; the third branch must still win past len1 + len2.
|
||||
pivots[:, 1] = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32)
|
||||
pivots[:, 2] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
|
||||
pivots[:, 3] = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32)
|
||||
else:
|
||||
total = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32)
|
||||
pivots[:, 1] = total // 2
|
||||
pivots[:, 2] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
|
||||
pivots[:, 3] = total - total // 2
|
||||
return expand_swa_index(pivots, width)
|
||||
|
||||
return run
|
||||
|
||||
return [
|
||||
("n32_w512", True, make(32, 512, False)),
|
||||
("n32_w513_overflow", True, make(32, 513, True)),
|
||||
("n1_w1", True, make(1, 1, False)),
|
||||
("n37_w2048_overflow", True, make(37, 2048, True)),
|
||||
("n8_w64", True, make(8, 64, False)),
|
||||
]
|
||||
|
||||
|
||||
def _cases_kv_index_cat() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
"""Window + compressed index concatenation with the -1 padding preserved."""
|
||||
|
||||
def make(total_q: int, window: int, topk: int):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
kv_index = torch.randint(-1, 4000, (total_q, window), device="cuda", dtype=torch.int32)
|
||||
kv_ca = torch.randint(-1, 300, (total_q, topk), device="cuda", dtype=torch.int32)
|
||||
return kv_index_cat(kv_index, kv_ca, 512, 2000)
|
||||
|
||||
return run
|
||||
|
||||
return [
|
||||
("q8_w512_k128", True, make(8, 512, 128)),
|
||||
("q1_w64_k32", True, make(1, 64, 32)),
|
||||
("q5_w2048_k1024", True, make(5, 2048, 1024)),
|
||||
]
|
||||
|
||||
|
||||
def _cases_convert_kv_index() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
"""Logical to physical index translation through a page table."""
|
||||
|
||||
def make(b: int, s: int, k: int, block_size: int, max_blocks: int):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
topk = torch.randint(0, max_blocks * block_size, (b, s, k), device="cuda", dtype=torch.int32)
|
||||
ctx = torch.randint(0, k + 1, (b, s), device="cuda", dtype=torch.int32)
|
||||
table = torch.randint(0, 50, (b, max_blocks), device="cuda", dtype=torch.int32)
|
||||
convert_kv_index(topk, block_size, ctx, table)
|
||||
return topk
|
||||
|
||||
return run
|
||||
|
||||
return [
|
||||
("b2_s4_k64_bs64", True, make(2, 4, 64, 64, 32)),
|
||||
("b1_s1_k128_bs128", True, make(1, 1, 128, 128, 16)),
|
||||
("b3_s2_k16_bs8", True, make(3, 2, 16, 8, 64)),
|
||||
]
|
||||
|
||||
|
||||
def _cases_load_kv_index() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
"""Prefill (contiguous) and decode (paged) index expansion."""
|
||||
|
||||
def make_prefill(n: int, length: int):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
off_len = torch.zeros((n, 2), device="cuda", dtype=torch.int32)
|
||||
off_len[:, 0] = torch.randint(0, 500, (n,), device="cuda", dtype=torch.int32)
|
||||
off_len[:, 1] = torch.randint(0, length + 1, (n,), device="cuda", dtype=torch.int32)
|
||||
return get_kv_ca_index_prefill(off_len, length)
|
||||
|
||||
return run
|
||||
|
||||
def make_decode(n: int, length: int, tpb: int, num_blocks: int):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
off_len = torch.zeros((n, 2), device="cuda", dtype=torch.int32)
|
||||
off_len[:, 0] = torch.randint(0, 500, (n,), device="cuda", dtype=torch.int32)
|
||||
off_len[:, 1] = torch.randint(0, length + 1, (n,), device="cuda", dtype=torch.int32)
|
||||
table = torch.randint(0, num_blocks, (num_blocks,), device="cuda", dtype=torch.int32)
|
||||
return get_kv_ca_index_decode(off_len, length, tpb, table)
|
||||
|
||||
return run
|
||||
|
||||
return [
|
||||
("prefill_n8_len512", True, make_prefill(8, 512)),
|
||||
("prefill_n1_len1024", True, make_prefill(1, 1024)),
|
||||
("prefill_n33_len100", True, make_prefill(33, 100)),
|
||||
("decode_n8_len512_tpb64", True, make_decode(8, 512, 64, 128)),
|
||||
("decode_n5_len128_tpb128", True, make_decode(5, 128, 128, 64)),
|
||||
]
|
||||
|
||||
|
||||
def _cases_store_mtp_info() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
"""Draft logprob pooling, including the slot-0 no-write path."""
|
||||
|
||||
def make(req_cnt: int, nextn: int, vocab: int, has_index: bool, has_confidence: bool, zero_slots: bool):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
# The pool is addressed as slot * nextn + step and slots run 0..req_cnt, so it needs
|
||||
# room for every slot including the reserved one.
|
||||
pool_rows = (req_cnt + 1) * nextn
|
||||
write_slots = torch.arange(1, req_cnt + 1, device="cuda", dtype=torch.int32)
|
||||
if zero_slots:
|
||||
write_slots[0] = 0
|
||||
token_ids = torch.randint(0, vocab, (req_cnt, nextn), device="cuda", dtype=torch.int32)
|
||||
logprobs = torch.randn((req_cnt, nextn, vocab), device="cuda", dtype=torch.float32)
|
||||
logprobs_index = torch.randint(0, vocab, (req_cnt, nextn, vocab), device="cuda", dtype=torch.int32)
|
||||
confidence = torch.rand((req_cnt, nextn), device="cuda", dtype=torch.float32)
|
||||
token_ids_pool = torch.full((pool_rows,), -7, device="cuda", dtype=torch.int32)
|
||||
logprobs_pool = torch.full((pool_rows, vocab), -7.0, device="cuda", dtype=torch.float32)
|
||||
logprobs_index_pool = torch.full((pool_rows, vocab), -7, device="cuda", dtype=torch.int32)
|
||||
confidence_pool = torch.full((pool_rows,), -7.0, device="cuda", dtype=torch.float32)
|
||||
store_mtp_info(
|
||||
write_slots,
|
||||
token_ids,
|
||||
logprobs,
|
||||
logprobs_index,
|
||||
confidence,
|
||||
token_ids_pool,
|
||||
logprobs_pool,
|
||||
logprobs_index_pool,
|
||||
confidence_pool,
|
||||
)
|
||||
return (
|
||||
token_ids_pool.clone(),
|
||||
logprobs_pool.clone(),
|
||||
logprobs_index_pool.clone(),
|
||||
confidence_pool.clone(),
|
||||
)
|
||||
|
||||
return run
|
||||
|
||||
return [
|
||||
("n4_nextn5_v1k_full", True, make(4, 5, 1024, True, True, False)),
|
||||
("n4_nextn5_v1k_no_index", True, make(4, 5, 1024, False, True, False)),
|
||||
("n4_nextn5_v1k_no_conf", True, make(4, 5, 1024, True, False, False)),
|
||||
("n4_nextn5_v1k_slot0", True, make(4, 5, 1024, True, True, True)),
|
||||
("n2_nextn1_v129280", True, make(2, 1, 129280, True, True, False)),
|
||||
]
|
||||
|
||||
|
||||
CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
|
||||
"gather_tl": _cases_gather,
|
||||
"gather_input_ids_zos_tl": _cases_gather_input_ids,
|
||||
"gather_engram_tokens_zos_tl": _cases_gather_engram,
|
||||
"expand_swa_index_tl": _cases_swa_index,
|
||||
"kv_index_cat_tl": _cases_kv_index_cat,
|
||||
"convert_kv_index_tl": _cases_convert_kv_index,
|
||||
"load_kv_index_tl": _cases_load_kv_index,
|
||||
"store_mtp_info_tl": _cases_store_mtp_info,
|
||||
}
|
||||
|
||||
@ -81,6 +81,17 @@ def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float =
|
||||
print(f" FAIL {name}: max abs {diff.max().item():.3e}, max rel {(diff / scale).max().item():.3e}")
|
||||
return bool(close)
|
||||
|
||||
if isinstance(expected, (tuple, list)):
|
||||
# A case that returns several tensors is only right if every element is; recurse so each
|
||||
# one gets its own report instead of collapsing into one boolean.
|
||||
if len(expected) != len(actual):
|
||||
print(f" FAIL {name}: {len(expected)} outputs vs {len(actual)}")
|
||||
return False
|
||||
ok = True
|
||||
for i, (e, a) in enumerate(zip(expected, actual, strict=True)):
|
||||
ok &= compare(f"{name}[{i}]", e, a, exact=exact, rtol=rtol, atol=atol)
|
||||
return ok
|
||||
|
||||
same = expected == actual
|
||||
if not same:
|
||||
print(f" FAIL {name}: {expected!r} vs {actual!r}")
|
||||
@ -113,7 +124,8 @@ def main() -> None:
|
||||
try:
|
||||
expected, actual = run_both(kernel, fn)
|
||||
except Exception as exc: # noqa: BLE001 - a raise on one backend is itself a result
|
||||
print(f" FAIL {case_name}: {type(exc).__name__}: {str(exc).splitlines()[0][:160]}")
|
||||
lines = str(exc).splitlines()
|
||||
print(f" FAIL {case_name}: {type(exc).__name__}: {lines[0][:160] if lines else ''}")
|
||||
failures += 1
|
||||
continue
|
||||
if not compare(case_name, expected, actual, exact=exact):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user