mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
a6c3e69e67
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.
304 lines
14 KiB
Python
304 lines
14 KiB
Python
"""Per-kernel A/B cases: run the public interface function once per backend and compare.
|
|
|
|
Each entry maps a kernel name to a builder returning ``(case_name, exact, callable)`` triples.
|
|
``exact=True`` demands bitwise equality, which is right for index and metadata kernels; floating
|
|
point kernels are compared with a tolerance and are labelled ``exact=False``.
|
|
|
|
Shapes and dtypes reuse the ones the engine actually passes, taken from the TileLang test matrices
|
|
in ``cuda/*.py`` and from the traced inference run.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
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]]]:
|
|
"""Row gather across the widths the engine uses: metadata rows, heads, and hidden states."""
|
|
|
|
def make(n: int, rows: int, d: int, dtype: torch.dtype, index_dtype: torch.dtype, strided: bool):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
base = (
|
|
torch.randn(n, d, device="cuda").to(dtype)
|
|
if dtype.is_floating_point
|
|
else torch.randint(0, 100, (n, d), device="cuda", dtype=dtype)
|
|
)
|
|
x = base[:, ::1] if strided else base
|
|
index = torch.randint(-3, n + 3, (rows,), device="cuda", dtype=index_dtype)
|
|
return gather(x, index)
|
|
|
|
return run
|
|
|
|
return [
|
|
("n1k_m233_d128_int32_idx", True, make(1024, 233, 128, torch.int8, torch.int32, False)),
|
|
("n1k_m1_d4_int8", True, make(1024, 1, 4, torch.int8, torch.int32, False)),
|
|
("n64_m256_d528_bytes", True, make(64, 256, 528, torch.int8, torch.int32, False)),
|
|
("n512_m300_d4096_bytes", True, make(512, 300, 4096, torch.int8, torch.int32, False)),
|
|
("n512_m300_d7168_int64_idx", True, make(512, 300, 7168, torch.int8, torch.int64, False)),
|
|
("n256_m256_d4096_bf16", True, make(256, 256, 4096, torch.bfloat16, torch.int32, False)),
|
|
("n256_m0_d128_empty", True, make(256, 0, 128, torch.int8, torch.int32, False)),
|
|
("n1_m17_d128_single_row", True, make(1, 17, 128, torch.int8, torch.int32, False)),
|
|
("n300_m233_d1025_odd", True, make(300, 233, 1025, torch.int8, torch.int32, False)),
|
|
]
|
|
|
|
|
|
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,
|
|
}
|