mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
fece5d9047
An edit of mine dropped a `return launch` from `finalize_verification_tl`. The A/B case for that kernel was not re-run, so nothing local objected; the failure surfaced as a TypeError inside the sampler of an end-to-end run, once per rank. The harness now checks that every `*_tl` symbol the selected backend exports is callable, which turns that class of mistake into one line printed before any case executes. Also records where the remaining sixteen kernels actually sit. None of them is unconditionally on the path, but the distribution is lopsided and the document had it backwards in one place: the three `candidate_*` kernels are not merely "off on SM90", they are switched off *by* `PLATFORM.is_sm90` -- and this checkpoint sets `candidate_source_layer_id=20` of 40 layers. On a card where `is_sm90` is False, that is nineteen layers on every step, which is the opposite of optional. Their harder blocker is noted alongside: that path also needs `deep_gemm.fp8_fp4_sparse_mqa_logits`, which is SM100-only, so porting those three alone would buy nothing. The four kernels that no production path reaches are listed as such, so nobody spends an afternoon on them.
1395 lines
69 KiB
Python
1395 lines
69 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 import ops
|
|
from deepinfer.model_executor.layers.tilelang_ops.compress import (
|
|
compress_and_update_state_decode,
|
|
compress_and_update_state_prefill,
|
|
)
|
|
from deepinfer.model_executor.layers.tilelang_ops.compute_seqlens_k_zos import compute_seqlens_k_zos
|
|
from deepinfer.model_executor.layers.tilelang_ops.dflash import confidence_head_fused, dflash_input_store
|
|
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_fused_dflash_get_input,
|
|
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.norm import norm
|
|
from deepinfer.model_executor.layers.tilelang_ops.prepare_v4_decode_metadata import prepare_v4_decode_metadata
|
|
from deepinfer.model_executor.layers.tilelang_ops.sampler import remap_draft_topk
|
|
from deepinfer.model_executor.layers.tilelang_ops.store_kv_cache import store_kv_cache
|
|
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 _sampler_inputs() -> tuple[Any, Any]:
|
|
"""The sampler's C++-backed input structs, imported late.
|
|
|
|
`sampling_inputs` reaches `deepinfer._C`, a compiled extension built against one specific
|
|
torch ABI. Hosts that only carry Triton (vendor images, CI images) cannot import it, and the
|
|
kernels that do not touch the sampler should still be testable there.
|
|
"""
|
|
from deepinfer.model_executor.sampling_inputs import LogitsTransformInputs, SamplingKernelInputs
|
|
|
|
return LogitsTransformInputs, SamplingKernelInputs
|
|
|
|
|
|
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_dflash_get_input() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Packed draft-token inputs, including padding requests and the optional mask branches."""
|
|
|
|
def make(
|
|
real_cnt: int,
|
|
dflash_block_size: int,
|
|
window_size: int,
|
|
swa_block_size: int,
|
|
has_valid_token_mask: bool,
|
|
has_gate_force_random: bool,
|
|
padding: int,
|
|
seed: int,
|
|
):
|
|
def run() -> Any:
|
|
torch.manual_seed(seed)
|
|
total_q = 64
|
|
slot_mapping_swa = torch.randint(0, 4096, (total_q,), device="cuda", dtype=torch.int32)
|
|
positions = torch.randint(0, 4096, (total_q,), device="cuda", dtype=torch.int32)
|
|
ftp = torch.randint(0, total_q, (real_cnt,), device="cuda", dtype=torch.int32)
|
|
last_valid_offset = torch.randint(0, dflash_block_size, (real_cnt,), device="cuda", dtype=torch.int32)
|
|
# ftp < 0 marks a padding request; its rows keep positions 0 and -1 slot / kv_index.
|
|
for i in range(min(padding, real_cnt)):
|
|
ftp[i] = -1
|
|
return get_fused_dflash_get_input(
|
|
slot_mapping_swa,
|
|
positions,
|
|
ftp,
|
|
last_valid_offset,
|
|
swa_block_size,
|
|
dflash_block_size,
|
|
window_size,
|
|
has_valid_token_mask,
|
|
has_gate_force_random,
|
|
)
|
|
|
|
return run
|
|
|
|
return [
|
|
# The engine's shape: 4 draft slots, an SWA window of 128 slots.
|
|
("n4_df4_win128", True, make(4, 4, 128, 64, False, False, 0, 0)),
|
|
("n4_df4_win128_mask", True, make(4, 4, 128, 64, True, False, 0, 1)),
|
|
("n4_df4_win128_gate", True, make(4, 4, 128, 64, False, True, 0, 2)),
|
|
("n4_df4_win128_mask_gate", True, make(4, 4, 128, 64, True, True, 0, 3)),
|
|
# A padding request must still write every row, with positions 0 and slot / kv_index -1.
|
|
("n4_df4_pad2", True, make(4, 4, 128, 64, True, True, 2, 4)),
|
|
("n1_df1_pad_all", True, make(1, 1, 128, 64, True, True, 1, 5)),
|
|
# kv_len = 8 + 8 = 16 rounds up to 64, which is shorter than one 128-wide tile: the
|
|
# column loop must round up or the tail of the row is never written.
|
|
("n4_df8_win8", True, make(4, 8, 8, 64, True, True, 0, 6)),
|
|
# kv_len = 200 + 32 = 232 rounds up to 256, spanning two tiles.
|
|
("n4_df32_win200", True, make(4, 32, 200, 128, True, False, 0, 7)),
|
|
("n16_df4_win128", True, make(16, 4, 128, 64, True, True, 3, 8)),
|
|
]
|
|
|
|
|
|
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)),
|
|
]
|
|
|
|
|
|
def _cases_norm() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""RMSNorm in the forms the engine uses: bf16 out, fp8 with fp32 scales, and both at once."""
|
|
|
|
def make(rows: int, d: int, with_residual: bool, as_fp8: bool, return_bf16: bool, round_sf: bool):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
x = torch.randn((rows, d), device="cuda", dtype=torch.bfloat16)
|
|
w = torch.randn((d,), device="cuda", dtype=torch.bfloat16)
|
|
residual = torch.randn((rows, d), device="cuda", dtype=torch.bfloat16) if with_residual else None
|
|
kwargs: dict[str, Any] = {"eps": 1e-20, "out_scale": 1.5}
|
|
if as_fp8:
|
|
kwargs.update(
|
|
fmt="e4m3",
|
|
num_per_channels=128,
|
|
round_sf=round_sf,
|
|
use_tma_aligned_col_major_sf=True,
|
|
)
|
|
return norm(x, w, residual=residual, return_bf16=return_bf16, **kwargs)
|
|
|
|
return run
|
|
|
|
return [
|
|
("r128_d5120_bf16", False, make(128, 5120, False, False, False, False)),
|
|
("r128_d5120_bf16_residual", False, make(128, 5120, True, False, False, False)),
|
|
("r64_d5120_fp8_round", False, make(64, 5120, False, True, False, True)),
|
|
("r64_d5120_fp8_no_round", False, make(64, 5120, False, True, False, False)),
|
|
("r64_d5120_fp8_round_bf16", False, make(64, 5120, False, True, True, True)),
|
|
("r33_d5120_fp8_round_bf16", False, make(33, 5120, False, True, True, True)),
|
|
("r16_d7168_fp8_round", False, make(16, 7168, False, True, False, True)),
|
|
("r8_d512_fp8_round", False, make(8, 512, False, True, False, True)),
|
|
]
|
|
|
|
|
|
def _cases_map_lookup() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Both lookup modes, including out-of-range and missing keys."""
|
|
from deepinfer.model_executor.layers.tilelang_ops.sampler import map_lookup
|
|
|
|
def make_index(rows: int, width: int, dtype: torch.dtype):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
value = (
|
|
torch.randn((rows, width), device="cuda", dtype=dtype)
|
|
if dtype.is_floating_point
|
|
else torch.randint(0, 1000, (rows, width), device="cuda", dtype=dtype)
|
|
)
|
|
query = torch.randint(-width, 2 * width, (rows, 1), device="cuda", dtype=torch.int32)
|
|
return map_lookup(None, value, query, mode="by_index")
|
|
|
|
return run
|
|
|
|
def make_key(rows: int, width: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
# Distinct keys per row: that is what the indexer produces, and it is also the only
|
|
# case where the kernel (last match wins) and the Python reference (max) agree.
|
|
key = torch.randperm(width, device="cuda").to(torch.int32).unsqueeze(0).repeat(rows, 1)
|
|
value = torch.randn((rows, width), device="cuda", dtype=torch.float32)
|
|
# Half the queries exist in the row, half do not.
|
|
present = key.gather(-1, torch.randint(0, width, (rows, 1), device="cuda").long())
|
|
missing = torch.full((rows, 1), 10_000, device="cuda", dtype=torch.int32)
|
|
keep = torch.randint(0, 2, (rows, 1), device="cuda") == 1
|
|
query = torch.where(keep, present, missing)
|
|
return map_lookup(key, value, query, mode="by_key")
|
|
|
|
return run
|
|
|
|
return [
|
|
("index_f32", True, make_index(256, 4096, torch.float32)),
|
|
("index_i32", True, make_index(256, 4096, torch.int32)),
|
|
("key_f32", False, make_key(128, 4096)),
|
|
]
|
|
|
|
|
|
def _cases_expand_sampling_params() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Per-row parameter expansion for the draft rows."""
|
|
|
|
def make(req_cnt: int, total_q: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
temps = torch.rand((req_cnt,), device="cuda", dtype=torch.float32)
|
|
top_ps = torch.rand((req_cnt,), device="cuda", dtype=torch.float32)
|
|
top_ks = torch.randint(0, 4096, (req_cnt,), device="cuda", dtype=torch.int32)
|
|
positions = torch.randint(0, 1000, (req_cnt,), device="cuda", dtype=torch.int32)
|
|
seeds = torch.randint(0, 2**62, (req_cnt,), device="cuda", dtype=torch.uint64)
|
|
row_to_req = torch.randint(0, req_cnt, (total_q,), device="cuda", dtype=torch.int32)
|
|
local_step = torch.randint(0, 6, (total_q,), device="cuda", dtype=torch.int32)
|
|
outs = [
|
|
torch.empty((total_q,), device="cuda", dtype=torch.float32),
|
|
torch.empty((total_q,), device="cuda", dtype=torch.float32),
|
|
torch.empty((total_q,), device="cuda", dtype=torch.int32),
|
|
torch.empty((total_q,), device="cuda", dtype=torch.int32),
|
|
torch.empty((total_q,), device="cuda", dtype=torch.uint64),
|
|
]
|
|
kernel = ops.expand_sampling_params_tl()
|
|
kernel(temps, top_ps, top_ks, positions, seeds, row_to_req, local_step, *outs)
|
|
return tuple(outs)
|
|
|
|
return run
|
|
|
|
return [("n4_q32", True, make(4, 32)), ("n1_q5", True, make(1, 5)), ("n8_q100", True, make(8, 100))]
|
|
|
|
|
|
def _cases_prepare_verification() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Draft-row metadata: request index, local step and the verify flag."""
|
|
|
|
def make(req_cnt: int, max_len: int, build_steps: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
seqlen_q = torch.randint(1, max_len + 1, (req_cnt,), device="cuda", dtype=torch.int32)
|
|
cu = torch.cumsum(seqlen_q, 0, dtype=torch.int32)
|
|
total_q = int(cu[-1].item())
|
|
ftp = torch.zeros((req_cnt,), device="cuda", dtype=torch.int32)
|
|
ftp[1:] = cu[:-1]
|
|
# The sentinel makes "wrote nothing" visible. `row_to_req` is a `torch.empty` buffer in
|
|
# the engine, so a kernel that fills fewer than `total_q` rows leaves garbage behind and
|
|
# every consumer then indexes with it -- and both backends would still agree with each
|
|
# other here, which is why this check has to be explicit rather than a comparison.
|
|
row_to_req = torch.full((total_q,), -7, device="cuda", dtype=torch.int32)
|
|
local_step = torch.full((total_q,), -7, device="cuda", dtype=torch.int32)
|
|
is_verify = torch.zeros((total_q,), device="cuda", dtype=torch.bool)
|
|
# Driven through the factory: the interface wrapper allocates its own outputs, which
|
|
# would leave nothing to compare.
|
|
kernel = ops.prepare_verification_metadata_tl(build_steps)
|
|
kernel(ftp, seqlen_q, row_to_req, local_step, is_verify)
|
|
assert int((row_to_req == -7).sum()) == 0, "kernel left rows unwritten"
|
|
assert int((local_step == -7).sum()) == 0, "kernel left rows unwritten"
|
|
return row_to_req, local_step, is_verify
|
|
|
|
return run
|
|
|
|
return [
|
|
("n4_maxlen6", True, make(4, 6, 6)),
|
|
("n1_maxlen1", True, make(1, 1, 1)),
|
|
("n8_maxlen6", True, make(8, 6, 6)),
|
|
# `num_steps=0` is what the engine builds when speculative decoding is off, so it must
|
|
# still write every row.
|
|
("n8_zero_steps", True, make(8, 6, 0)),
|
|
# A request longer than one tile (32 lanes) exercises the runtime tile loop.
|
|
("n2_maxlen130", True, make(2, 130, 130)),
|
|
("n64_maxlen6", True, make(64, 6, 6)),
|
|
]
|
|
|
|
|
|
def _cases_get_accept_result() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Prefix acceptance over each request's packed rows, with and without golden tokens."""
|
|
|
|
def make(req_cnt: int, max_len: int, has_golden: bool, golden_ratio: float, seed: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(seed)
|
|
seqlen_q = torch.randint(1, max_len + 1, (req_cnt,), device="cuda", dtype=torch.int32)
|
|
cu = torch.cumsum(seqlen_q, 0, dtype=torch.int32)
|
|
total_q = int(cu[-1].item())
|
|
ftp = torch.zeros((req_cnt,), device="cuda", dtype=torch.int32)
|
|
ftp[1:] = cu[:-1]
|
|
row_to_req = torch.repeat_interleave(torch.arange(req_cnt, device="cuda", dtype=torch.int32), seqlen_q.long())
|
|
# Every row but each request's last is a draft row awaiting verification; the last is
|
|
# the fresh row, whose verdict is final either way.
|
|
within = torch.arange(total_q, device="cuda", dtype=torch.int32) - ftp[row_to_req.long()]
|
|
is_verify = within + 1 < seqlen_q[row_to_req.long()]
|
|
|
|
pool_rows = total_q + 7
|
|
draft_pool = torch.randint(0, 1000, (pool_rows,), device="cuda", dtype=torch.int32)
|
|
read_rows = torch.randint(0, pool_rows, (total_q,), device="cuda", dtype=torch.int32)
|
|
raw_accept = torch.randint(0, 2, (total_q,), device="cuda", dtype=torch.int32).bool()
|
|
candidate_token_ids = torch.randint(0, 1000, (total_q,), device="cuda", dtype=torch.int32)
|
|
|
|
# Golden tokens: -1 means "no override"; otherwise the value occasionally coincides with
|
|
# the draft proposal so both the accept and the reject branch are exercised.
|
|
if has_golden:
|
|
golden = torch.randint(0, 1000, (total_q,), device="cuda", dtype=torch.int32)
|
|
draft_tokens = draft_pool[read_rows.long()]
|
|
collide = torch.rand((total_q,), device="cuda", dtype=torch.float32) < 0.5
|
|
golden = torch.where(collide, draft_tokens, golden)
|
|
drop = torch.rand((total_q,), device="cuda", dtype=torch.float32) < golden_ratio
|
|
golden = torch.where(drop, torch.full_like(golden, -1), golden)
|
|
else:
|
|
golden = candidate_token_ids
|
|
|
|
token_ids = torch.zeros((total_q,), device="cuda", dtype=torch.int32)
|
|
accept = torch.zeros((total_q,), device="cuda", dtype=torch.bool)
|
|
is_final = torch.zeros((total_q,), device="cuda", dtype=torch.bool)
|
|
kernel = ops.get_accept_result_tl(has_golden)
|
|
kernel(
|
|
raw_accept,
|
|
golden,
|
|
draft_pool,
|
|
read_rows,
|
|
candidate_token_ids,
|
|
is_verify,
|
|
ftp,
|
|
row_to_req,
|
|
token_ids,
|
|
accept,
|
|
is_final,
|
|
)
|
|
return token_ids, accept, is_final
|
|
|
|
return run
|
|
|
|
return [
|
|
("n4_len6_no_golden", True, make(4, 6, False, 0.0, 0)),
|
|
("n1_len6_no_golden", True, make(1, 6, False, 0.0, 1)),
|
|
# A request of length 1 has an empty prefix, so the scan runs zero iterations.
|
|
("n1_len1_no_golden", True, make(1, 1, False, 0.0, 2)),
|
|
("n8_len4_golden_all", True, make(8, 4, True, 0.0, 3)),
|
|
("n8_len4_golden_mixed", True, make(8, 4, True, 0.5, 4)),
|
|
("n4_len6_golden_none", True, make(4, 6, True, 1.0, 5)),
|
|
# Wider than one block of 128 rows, so several programs run and the prefix scan is long.
|
|
("n2_len130_no_golden", True, make(2, 130, False, 0.0, 6)),
|
|
("n2_len130_golden_mixed", True, make(2, 130, True, 0.5, 7)),
|
|
("n64_len6_golden_mixed", True, make(64, 6, True, 0.5, 8)),
|
|
]
|
|
|
|
|
|
def _cases_remap_draft_topk() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Align the previous draft's top-k logprobs onto the current target top-k indices."""
|
|
|
|
def make(total_q: int, force_topk: int, vocab_size: int, overlap: float, seed: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(seed)
|
|
pool_rows = total_q + 3
|
|
# Draft indices are distinct token ids drawn from the vocabulary; the target set is a
|
|
# mixture of draft ids (so the match branch fires) and unrelated ids (so the -inf
|
|
# branch fires), controlled by `overlap`.
|
|
draft_index = torch.stack([torch.randperm(vocab_size, device="cuda")[:force_topk].to(torch.int32) for _ in range(pool_rows)])
|
|
draft_logprobs = torch.rand((pool_rows, force_topk), device="cuda", dtype=torch.float32) - 3.0
|
|
draft_token_ids_pool = torch.randint(0, vocab_size, (pool_rows,), device="cuda", dtype=torch.int32)
|
|
read_rows = torch.randint(0, pool_rows, (total_q,), device="cuda", dtype=torch.int32)
|
|
picked = draft_index.index_select(0, read_rows.long())
|
|
random_ids = torch.randint(0, vocab_size, (total_q, force_topk), device="cuda", dtype=torch.int32)
|
|
take_draft = torch.rand((total_q, force_topk), device="cuda", dtype=torch.float32) < overlap
|
|
target_index = torch.where(take_draft, picked, random_ids).contiguous()
|
|
# Make the "which draft slot is the draft token" output non-trivial for some rows.
|
|
if total_q > 0:
|
|
target_index[0, 0] = draft_token_ids_pool[read_rows[0].long()]
|
|
return remap_draft_topk(
|
|
draft_index,
|
|
draft_logprobs,
|
|
draft_token_ids_pool,
|
|
target_index,
|
|
read_rows,
|
|
vocab_size,
|
|
)
|
|
|
|
return run
|
|
|
|
return [
|
|
# The engine's smallest useful shape, and the one the sweep uses.
|
|
("n4_topk64_vocab1k", True, make(4, 64, 1024, 0.7, 0)),
|
|
("n4_topk64_nomatch", True, make(4, 64, 1024, 0.0, 1)),
|
|
("n4_topk64_allmatch", True, make(4, 64, 1024, 1.0, 2)),
|
|
("n1_topk16_vocab128", True, make(1, 16, 128, 0.5, 3)),
|
|
# Wider than one 128-element target block, so several programs handle one row.
|
|
("n2_topk300_vocab8k", True, make(2, 300, 8192, 0.6, 4)),
|
|
# The size the H20 sweep actually passes.
|
|
("n2_topk4096_vocab129280", True, make(2, 4096, 129280, 0.5, 5)),
|
|
]
|
|
|
|
|
|
def _cases_finalize_verification() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Finalize accepted tokens, count per-step totals and pack the per-request payload."""
|
|
|
|
def make(
|
|
req_cnt: int,
|
|
num_steps: int,
|
|
has_sum_smaller_probs: bool,
|
|
has_confidence: bool,
|
|
topk_n: int,
|
|
num_stop_tokens: int,
|
|
seed: int,
|
|
):
|
|
def run() -> Any:
|
|
torch.manual_seed(seed)
|
|
# MTP packs at most num_steps + 1 rows per request; the last row of each request is
|
|
# the fresh one, so a length in that range is what the engine can produce.
|
|
low = 1 if num_steps == 0 else 1
|
|
high = max(num_steps + 1, 1)
|
|
seqlen_q = torch.randint(low, high + 1, (req_cnt,), device="cuda", dtype=torch.int32)
|
|
cu = torch.cumsum(seqlen_q, 0, dtype=torch.int32)
|
|
total_q = int(cu[-1].item())
|
|
ftp = torch.zeros((req_cnt,), device="cuda", dtype=torch.int32)
|
|
ftp[1:] = cu[:-1]
|
|
row_to_req = torch.repeat_interleave(torch.arange(req_cnt, device="cuda", dtype=torch.int32), seqlen_q.long())
|
|
step = torch.arange(total_q, device="cuda", dtype=torch.int32) - ftp[row_to_req.long()]
|
|
is_verify = step + 1 < seqlen_q[row_to_req.long()]
|
|
|
|
# A prefix that stops at the first reject: everything from there on is final for that
|
|
# request. This is the shape `get_accept_result` produces, so the inputs are realistic.
|
|
accept_all = torch.randint(0, 2, (total_q,), device="cuda", dtype=torch.int32).bool()
|
|
effective_accept = torch.zeros((total_q,), device="cuda", dtype=torch.bool)
|
|
is_final = torch.zeros((total_q,), device="cuda", dtype=torch.bool)
|
|
alive_by_req = [True] * req_cnt
|
|
for row in range(total_q):
|
|
r = int(row_to_req[row])
|
|
if not alive_by_req[r]:
|
|
continue
|
|
accepted = bool(accept_all[row]) and bool(is_verify[row])
|
|
effective_accept[row] = accepted
|
|
if not accepted:
|
|
is_final[row] = True
|
|
alive_by_req[r] = False
|
|
# Every request needs its final row marked, including one whose drafts all passed.
|
|
for r in range(req_cnt):
|
|
if alive_by_req[r]:
|
|
is_final[int(cu[r]) - 1] = True
|
|
|
|
effective_token_ids = torch.randint(0, 1000, (total_q,), device="cuda", dtype=torch.int32)
|
|
token_logprobs = torch.rand((total_q,), device="cuda", dtype=torch.float32) - 2.0
|
|
sum_smaller_probs = torch.rand((total_q,), device="cuda", dtype=torch.float32)
|
|
top_logprobs = torch.rand((total_q, topk_n), device="cuda", dtype=torch.float32)
|
|
top_token_ids = torch.randint(0, 1000, (total_q, topk_n), device="cuda", dtype=torch.int32)
|
|
accept_probs = torch.rand((total_q,), device="cuda", dtype=torch.float32)
|
|
pool_rows = total_q + 5
|
|
confidence_pool = torch.rand((pool_rows,), device="cuda", dtype=torch.float32)
|
|
mtp_info_read_rows = torch.randint(0, pool_rows, (total_q,), device="cuda", dtype=torch.int32)
|
|
stop_token_ids = torch.randint(0, 1000, (req_cnt, num_stop_tokens), device="cuda", dtype=torch.int32)
|
|
# Make a stop hit actually happen sometimes, otherwise the scan is only ever tested
|
|
# on its negative branch.
|
|
if req_cnt > 0:
|
|
first = torch.randint(0, req_cnt, (1,), device="cuda").item()
|
|
start = int(ftp[first].item())
|
|
stop_token_ids[first, 0] = effective_token_ids[start]
|
|
|
|
outs = {
|
|
"last_valid_offset": torch.empty((req_cnt,), device="cuda", dtype=torch.int32),
|
|
"num_generated_tokens": torch.empty((num_steps,), device="cuda", dtype=torch.int32),
|
|
"num_verified_tokens": torch.empty((num_steps,), device="cuda", dtype=torch.int32),
|
|
"batch_confidence_error": torch.empty((1,), device="cuda", dtype=torch.float32),
|
|
"emitted_stop_token": torch.empty((req_cnt,), device="cuda", dtype=torch.bool),
|
|
"sampled_token_ids": torch.empty((req_cnt,), device="cuda", dtype=torch.int32),
|
|
"sampled_token_logprobs": torch.empty((req_cnt,), device="cuda", dtype=torch.float32),
|
|
"sampled_sum_smaller_probs": torch.empty((req_cnt,), device="cuda", dtype=torch.float32),
|
|
"sampled_top_logprobs": torch.empty((req_cnt, topk_n), device="cuda", dtype=torch.float32),
|
|
"sampled_top_token_ids": torch.empty((req_cnt, topk_n), device="cuda", dtype=torch.int32),
|
|
"verified_token_ids": torch.empty((total_q,), device="cuda", dtype=torch.int32),
|
|
"verified_token_logprobs": torch.empty((total_q,), device="cuda", dtype=torch.float32),
|
|
"verified_sum_smaller_probs": torch.empty((total_q,), device="cuda", dtype=torch.float32),
|
|
"verified_top_logprobs": torch.empty((total_q, topk_n), device="cuda", dtype=torch.float32),
|
|
"verified_top_token_ids": torch.empty((total_q, topk_n), device="cuda", dtype=torch.int32),
|
|
}
|
|
kernel = ops.finalize_verification_tl(has_sum_smaller_probs, has_confidence, num_steps, num_stop_tokens)
|
|
kernel(
|
|
effective_accept,
|
|
is_final,
|
|
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,
|
|
*outs.values(),
|
|
)
|
|
# The interface leaves the buffers for a disabled feature untouched, so both backends
|
|
# hand back whatever `torch.empty` gave them and comparing those would only ever
|
|
# report noise. Compare what this variant actually writes.
|
|
compared = ["last_valid_offset", "num_generated_tokens", "num_verified_tokens"]
|
|
if has_confidence:
|
|
compared += ["batch_confidence_error", "emitted_stop_token"]
|
|
compared += [
|
|
"sampled_token_ids",
|
|
"sampled_token_logprobs",
|
|
"sampled_top_logprobs",
|
|
"sampled_top_token_ids",
|
|
"verified_token_ids",
|
|
"verified_token_logprobs",
|
|
"verified_top_logprobs",
|
|
"verified_top_token_ids",
|
|
]
|
|
if has_sum_smaller_probs:
|
|
compared += ["sampled_sum_smaller_probs", "verified_sum_smaller_probs"]
|
|
return tuple(outs[name] for name in compared)
|
|
|
|
return run
|
|
|
|
return [
|
|
# The engine's default: MTP on with 5 steps, confidence on, no sum-smaller-probs.
|
|
# Not `exact`: `batch_confidence_error` is a float atomic accumulation, so its last bits
|
|
# depend on the order the blocks happen to sum in.
|
|
("n4_steps5_conf_topk8_stop2", False, make(4, 5, False, True, 8, 2, 0)),
|
|
("n8_steps5_conf_topk8_stop3", False, make(8, 5, False, True, 8, 3, 1)),
|
|
("n4_steps5_ssp_conf", False, make(4, 5, True, True, 8, 2, 2)),
|
|
("n4_steps5_plain", False, make(4, 5, False, False, 8, 2, 3)),
|
|
# num_steps=0 is what the engine builds with speculative decoding off; the counter tensors
|
|
# are then empty and must not be touched.
|
|
("n4_steps0_conf", False, make(4, 0, False, True, 8, 2, 4)),
|
|
("n4_steps0_plain", False, make(4, 0, False, False, 8, 2, 5)),
|
|
("n1_steps1_plain", False, make(1, 1, False, False, 4, 2, 6)),
|
|
# Empty top-k payload: the engine passes topk_n=0 when the caller wants no logprobs.
|
|
("n4_steps5_topk0", False, make(4, 5, False, True, 0, 2, 7)),
|
|
# A width that does not divide the 128-element tile.
|
|
("n4_steps5_topk130", False, make(4, 5, False, True, 130, 2, 8)),
|
|
# More rows than one block of 16.
|
|
("n40_steps5_conf", False, make(40, 5, True, True, 8, 2, 9)),
|
|
]
|
|
|
|
|
|
def _cases_copy_block_tables() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Compact 1D block table expanded to one row per token."""
|
|
|
|
def make(n_seq: int, max_seqlen_q: int, max_blocks: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32)
|
|
cu = torch.cumsum(seqlens, 0, dtype=torch.int32)
|
|
total_q = int(cu[-1].item())
|
|
counts = torch.randint(1, max_blocks + 1, (n_seq,), device="cuda", dtype=torch.int32)
|
|
offsets = torch.zeros((n_seq + 1,), device="cuda", dtype=torch.int32)
|
|
offsets[1:] = torch.cumsum(counts, 0, dtype=torch.int32)
|
|
flat = torch.randint(0, 512, (int(offsets[-1].item()),), device="cuda", dtype=torch.int32)
|
|
out = torch.full((total_q, max_blocks), -7, device="cuda", dtype=torch.int32)
|
|
# No standalone interface function: this kernel is only reached from
|
|
# `prepare_v4_decode_metadata`, so the case drives the factory directly. The harness
|
|
# swaps `ops.copy_block_tables_tl` between backends, which is all that is needed.
|
|
kernel = ops.copy_block_tables_tl(max_seqlen_q, max_blocks)
|
|
kernel(seqlens, cu, flat, offsets, out)
|
|
return out
|
|
|
|
return run
|
|
|
|
return [
|
|
("n4_q6_b64", True, make(4, 6, 64)),
|
|
("n1_q1_b32", True, make(1, 1, 32)),
|
|
("n8_q6_b1400_multi_chunk", True, make(8, 6, 1400)),
|
|
]
|
|
|
|
|
|
def _cases_store_kv_cache() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Quantised cache writes: the whole cache buffer must match byte for byte."""
|
|
# Imported here, not at module level: `cache_common` pulls in `tile_kernels`, and the whole
|
|
# point of this file is to run on a host where only the backend under test is installed.
|
|
from deepinfer.model_executor.layers.tilelang_ops.cache_common import (
|
|
FP8_FLOAT32_SF_INDEX_KV_FORMAT,
|
|
V41_FORMAT,
|
|
)
|
|
|
|
def make(layout: Any, block_size: int, n: int, with_invalid: bool):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
num_blocks = (n + block_size - 1) // block_size + 2
|
|
kv = torch.randn((n, layout.num_elements), device="cuda", dtype=torch.bfloat16)
|
|
slots = torch.arange(n, dtype=torch.int32, device="cuda") + block_size
|
|
if with_invalid:
|
|
slots[::7] = -1
|
|
cache = torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device="cuda")
|
|
store_kv_cache(kv, cache, slots, layout=layout)
|
|
return cache.clone()
|
|
|
|
return run
|
|
|
|
return [
|
|
("main_fp8_b32", True, make(V41_FORMAT, 32, 100, False)),
|
|
("main_fp8_b64", True, make(V41_FORMAT, 64, 233, True)),
|
|
("main_fp8_b134", True, make(V41_FORMAT, 134, 500, False)),
|
|
("index_fp32sf_b32", True, make(FP8_FLOAT32_SF_INDEX_KV_FORMAT, 32, 64, False)),
|
|
("index_fp32sf_b64", True, make(FP8_FLOAT32_SF_INDEX_KV_FORMAT, 64, 200, True)),
|
|
]
|
|
|
|
|
|
def _cases_load_paged_kv_cache() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Paged reads, including holes (src_block_id == -1) and partial segments.
|
|
|
|
The cache is filled by `store_kv_cache` rather than with random bytes: random bytes produce
|
|
fp8 NaN patterns and NaN scales, and those compare unequal to themselves, which would report a
|
|
mismatch that says nothing about the load kernel.
|
|
"""
|
|
from deepinfer.model_executor.layers.tilelang_ops.cache_common import (
|
|
FP8_FLOAT32_SF_INDEX_KV_FORMAT,
|
|
V41_FORMAT,
|
|
)
|
|
from deepinfer.model_executor.layers.tilelang_ops.load_paged_kv_cache import load_paged_kv_cache
|
|
|
|
def make(layout: Any, block_size: int, num_blocks: int, num_segments: int, with_hole: bool):
|
|
torch.manual_seed(0)
|
|
total = num_blocks * block_size
|
|
kv = torch.randn((total, layout.num_elements), device="cuda", dtype=torch.bfloat16)
|
|
cache = torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device="cuda")
|
|
store_kv_cache(kv, cache, torch.arange(total, dtype=torch.int32, device="cuda"), layout=layout)
|
|
|
|
segments = []
|
|
dst = 0
|
|
for i in range(num_segments):
|
|
begin = i % block_size
|
|
end = min(block_size, begin + 1 + (i * 7) % block_size)
|
|
src = -1 if (with_hole and i == 1) else i % num_blocks
|
|
segments.append([src, begin, end, dst])
|
|
dst += end - begin
|
|
seg = torch.tensor(segments, dtype=torch.int32, device="cuda")
|
|
|
|
def run() -> Any:
|
|
result = load_paged_kv_cache(cache, seg, dst, layout=layout)
|
|
if result.bf16_data is not None:
|
|
return result.bf16_data
|
|
# Byte views for the index layout: both payload and scales are passthrough copies, so
|
|
# exact bytes are the right comparison and NaN bit patterns are harmless.
|
|
return result.quantized_data.view(torch.uint8), result.sf_data.view(torch.uint8)
|
|
|
|
return run
|
|
|
|
return [
|
|
("main_fp8_b32", True, make(V41_FORMAT, 32, 8, 6, False)),
|
|
("main_fp8_b64_holes", True, make(V41_FORMAT, 64, 8, 6, True)),
|
|
("index_fp32sf_b32", True, make(FP8_FLOAT32_SF_INDEX_KV_FORMAT, 32, 8, 6, False)),
|
|
("index_fp32sf_b64_holes", True, make(FP8_FLOAT32_SF_INDEX_KV_FORMAT, 64, 8, 6, True)),
|
|
]
|
|
|
|
|
|
def _cases_fused_transform() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""The norm -> RoPE -> fp8 pipeline, in every form the engine can ask for."""
|
|
from deepinfer.model_executor.layers.tilelang_ops.cache_common import (
|
|
FP8_FLOAT32_SF_INDEX_KV_FORMAT,
|
|
V41_FORMAT,
|
|
)
|
|
from deepinfer.model_executor.layers.tilelang_ops.fused_transform import (
|
|
CastArgs,
|
|
KVCacheOutArgs,
|
|
NormArgs,
|
|
RopeArgs,
|
|
fused_norm_rope_hadamard_cast_store_kvcache,
|
|
)
|
|
|
|
def make(
|
|
n: int,
|
|
num_heads: int,
|
|
head_dim: int,
|
|
rotary: int,
|
|
*,
|
|
norm: bool,
|
|
rope: bool,
|
|
cast: bool,
|
|
cast_back: bool,
|
|
is_neox: bool = False,
|
|
conjugate: bool = False,
|
|
npc: int = 32,
|
|
sf_format: str = "ue8m0",
|
|
):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
x = torch.randn((n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16)
|
|
norm_args = NormArgs(weight=torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16), eps=1e-20) if norm else None
|
|
rope_args = None
|
|
if rope:
|
|
rotary_dim = rotary
|
|
half = rotary_dim // 2
|
|
angles = torch.randn((4096, half), device="cuda", dtype=torch.float32).cumsum(-1) * 0.01
|
|
cache = torch.cat([angles.cos(), angles.sin()], dim=-1)
|
|
positions = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
|
|
rope_args = RopeArgs(
|
|
positions=positions,
|
|
cos_sin_cache=cache,
|
|
rotary_dim=rotary_dim,
|
|
is_neox=is_neox,
|
|
offset=head_dim - rotary_dim,
|
|
conjugate=conjugate,
|
|
)
|
|
cast_args = CastArgs("e4m3", npc, round_sf=True, sf_format=sf_format, cast_back=cast_back) if cast else None
|
|
return fused_norm_rope_hadamard_cast_store_kvcache(x, norm_args, rope_args, None, cast_args, None)
|
|
|
|
return run
|
|
|
|
def make_index_cache(n: int, num_heads: int, head_dim: int, rotary: int, num_blocks: int):
|
|
"""Index-K path: fp32 scales, per-128 groups, into the index KV cache."""
|
|
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
layout = FP8_FLOAT32_SF_INDEX_KV_FORMAT
|
|
block_size = 64
|
|
x = torch.randn((n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16)
|
|
half = rotary // 2
|
|
angles = torch.randn((4096, half), device="cuda", dtype=torch.float32).cumsum(-1) * 0.01
|
|
cos_sin = torch.cat([angles.cos(), angles.sin()], dim=-1)
|
|
positions = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
|
|
kv_cache = torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device="cuda")
|
|
# Distinct slots: two tokens landing on the same slot would race, and the winner
|
|
# would depend on CTA scheduling rather than on the kernel under test.
|
|
slots = torch.randperm(num_blocks * block_size, device="cuda")[:n].to(torch.int32)
|
|
fused_norm_rope_hadamard_cast_store_kvcache(
|
|
x,
|
|
None,
|
|
RopeArgs(positions, cos_sin, rotary, False, offset=head_dim - rotary),
|
|
None,
|
|
CastArgs.from_cache_format(layout),
|
|
KVCacheOutArgs(kv_cache, slots, layout=layout),
|
|
)
|
|
return kv_cache
|
|
|
|
return run
|
|
|
|
def make_cache(n: int, block_size: int, num_blocks: int, with_bf16: bool, with_rope: bool = True):
|
|
"""The main-KV path: quantise and write straight into the paged cache."""
|
|
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
head_dim, rotary = 512, 64
|
|
layout = V41_FORMAT
|
|
x = torch.randn((n, 1, head_dim), device="cuda", dtype=torch.bfloat16)
|
|
weight = torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16)
|
|
half = rotary // 2
|
|
angles = torch.randn((4096, half), device="cuda", dtype=torch.float32).cumsum(-1) * 0.01
|
|
cache_cs = torch.cat([angles.cos(), angles.sin()], dim=-1)
|
|
positions = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
|
|
rope_args = RopeArgs(positions, cache_cs, rotary, False, offset=head_dim - rotary) if with_rope else None
|
|
kv_cache = torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device="cuda")
|
|
# Distinct slots: two tokens landing on the same slot would race, and the winner
|
|
# would depend on CTA scheduling rather than on the kernel under test.
|
|
slots = torch.randperm(num_blocks * block_size, device="cuda")[:n].to(torch.int32)
|
|
out_bf16 = torch.empty((n, 1, head_dim), device="cuda", dtype=torch.bfloat16) if with_bf16 else None
|
|
fused_norm_rope_hadamard_cast_store_kvcache(
|
|
x,
|
|
NormArgs(weight, 1e-20),
|
|
rope_args,
|
|
None,
|
|
CastArgs.from_cache_format(layout),
|
|
KVCacheOutArgs(kv_cache, slots, out=out_bf16, out_stage="norm", layout=layout),
|
|
)
|
|
return (kv_cache, out_bf16) if with_bf16 else kv_cache
|
|
|
|
return run
|
|
|
|
return [
|
|
# The configuration the engine builds, from DEEPINFER_TRACE_KERNEL_ARGS on a full run.
|
|
("engine_cfg_n1_h512", False, make(16, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True)),
|
|
# The engine's other call site: norm + rope + quantise into the paged KV cache, with the
|
|
# post-norm row also returned for the indexer.
|
|
("engine_cache", False, make_cache(16, 134, 4, True)),
|
|
("engine_cache_no_bf16", False, make_cache(16, 134, 4, False)),
|
|
("engine_cache_b32", False, make_cache(9, 32, 8, True)),
|
|
("engine_cache_no_rope", False, make_cache(8, 134, 2, True, with_rope=False)),
|
|
("engine_cfg_n128", False, make(128, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True)),
|
|
("engine_cfg_n1_token", False, make(1, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True)),
|
|
("neox", False, make(16, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True, is_neox=True)),
|
|
("conjugate", False, make(16, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True, conjugate=True)),
|
|
("no_rope", False, make(64, 1, 512, 64, norm=True, rope=False, cast=True, cast_back=True)),
|
|
("no_norm", False, make(64, 1, 512, 64, norm=False, rope=True, cast=True, cast_back=True)),
|
|
("bf16_only", False, make(64, 1, 512, 64, norm=True, rope=True, cast=False, cast_back=False)),
|
|
("return_quant", False, make(64, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=False)),
|
|
("multi_head_h128", False, make(32, 4, 128, 128, norm=True, rope=True, cast=True, cast_back=True)),
|
|
("h128_rope64", False, make(32, 2, 128, 64, norm=True, rope=True, cast=True, cast_back=True)),
|
|
# The indexer's Q quantisation: no norm, per-128 groups, fp32 scales, quantised output.
|
|
("index_q_fp32sf", False, make(64, 4, 128, 64, norm=False, rope=True, cast=True, cast_back=False, npc=128, sf_format="fp32")),
|
|
(
|
|
"index_q_fp32sf_many",
|
|
False,
|
|
make(256, 32, 128, 64, norm=False, rope=True, cast=True, cast_back=False, npc=128, sf_format="fp32"),
|
|
),
|
|
# Same, but written straight into the index KV cache.
|
|
# The index KV cache holds a single KV head (the indexer is MQA), which the baseline
|
|
# asserts too.
|
|
("index_k_cache", False, make_index_cache(48, 1, 128, 64, 16)),
|
|
]
|
|
|
|
|
|
def _cases_prepare_v4_decode_metadata() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Every output of the decode-metadata kernel, including the ratio-bucketed ones."""
|
|
|
|
def make(
|
|
n_seq: int,
|
|
max_seqlen_q: int,
|
|
max_num_blocks: int,
|
|
window_size: int,
|
|
swa_block_size: int,
|
|
block_size: int,
|
|
ca_ratios: tuple[int, ...],
|
|
dense_ratios: tuple[int, ...],
|
|
sparse_ratios: tuple[int, ...],
|
|
placeholders: bool,
|
|
):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
seqlens_q = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32)
|
|
seqlens_k = seqlens_q + torch.randint(0, 40, (n_seq,), device="cuda", dtype=torch.int32)
|
|
cu = torch.cumsum(seqlens_q, 0, dtype=torch.int32)
|
|
total_q = int(cu[-1].item())
|
|
per_seq_blocks = torch.randint(1, max_num_blocks + 1, (n_seq,), device="cuda", dtype=torch.int32)
|
|
offsets = torch.zeros((n_seq + 1,), device="cuda", dtype=torch.int32)
|
|
offsets[1:] = torch.cumsum(per_seq_blocks, 0, dtype=torch.int32)
|
|
total_blocks = int(offsets[-1].item())
|
|
block_tables_in = torch.randint(0, 40, (total_blocks,), device="cuda", dtype=torch.int32)
|
|
state_block_idx = torch.randint(0, 8, (n_seq,), device="cuda", dtype=torch.int32)
|
|
is_placeholder = torch.zeros((n_seq,), device="cuda", dtype=torch.bool)
|
|
if placeholders:
|
|
is_placeholder[::2] = True
|
|
|
|
n_ca = max(len(ca_ratios), 1)
|
|
n_dense = max(len(dense_ratios), 1)
|
|
n_sparse = max(len(sparse_ratios), 1)
|
|
i32 = {"device": "cuda", "dtype": torch.int32}
|
|
outs = [
|
|
torch.full((total_q,), -7, **i32), # positions
|
|
torch.zeros((total_q,), device="cuda", dtype=torch.bool), # valid_token_mask
|
|
torch.zeros((total_q,), device="cuda", dtype=torch.bool), # gate_force_random
|
|
torch.ones((total_q,), device="cuda", dtype=torch.bool), # fix_routing_mask
|
|
torch.full((total_q,), -7, **i32), # context_lens
|
|
torch.full((n_sparse, total_q), -7, **i32), # klen
|
|
torch.full((n_dense, total_q, 2), -7, **i32), # kv_ca_index_offset_len
|
|
torch.full((total_q, 2), -7, **i32), # kv_swa_index_offset_len
|
|
torch.full((total_q,), -7, **i32), # slot_mapping_store_swa
|
|
torch.full((total_q,), -7, **i32), # query_indices
|
|
torch.full((total_q,), -7, **i32), # state_block_idx_out
|
|
torch.full((total_q, max_num_blocks), -7, **i32), # block_tables_out
|
|
torch.full((n_ca, total_q), -7, **i32), # positions_c
|
|
torch.full((n_ca, total_q), -7, **i32), # slot_mapping_store_ca
|
|
]
|
|
prepare_v4_decode_metadata(
|
|
seqlens_q,
|
|
seqlens_k,
|
|
cu,
|
|
block_tables_in,
|
|
offsets,
|
|
state_block_idx,
|
|
is_placeholder,
|
|
*outs,
|
|
max_seqlen_q=max_seqlen_q,
|
|
max_num_blocks=max_num_blocks,
|
|
window_size=window_size,
|
|
swa_block_size=swa_block_size,
|
|
block_size=block_size,
|
|
ca_ratios=ca_ratios,
|
|
dense_ratios=dense_ratios,
|
|
sparse_ratios=sparse_ratios,
|
|
force_load_balance=False,
|
|
)
|
|
# Nothing should be left at the sentinel: a partially written output is a bug even when
|
|
# both backends agree on it.
|
|
for name, tensor in (
|
|
("positions", outs[0]),
|
|
("context_lens", outs[4]),
|
|
("kv_swa_index_offset_len", outs[7]),
|
|
("slot_mapping_store_swa", outs[8]),
|
|
("query_indices", outs[9]),
|
|
("positions_c", outs[12]),
|
|
("slot_mapping_store_ca", outs[13]),
|
|
):
|
|
assert int((tensor == -7).sum()) == 0, f"{name} left unwritten rows"
|
|
return tuple(outs)
|
|
|
|
return run
|
|
|
|
return [
|
|
# The engine's own configuration: window 128, SWA block 134, ratios {1, 2} all compressed.
|
|
("engine_cfg", False, make(8, 6, 64, 128, 134, 128, (1, 2), (), (1, 2), False)),
|
|
("engine_cfg_placeholders", False, make(8, 6, 64, 128, 134, 128, (1, 2), (), (1, 2), True)),
|
|
("with_dense", False, make(6, 6, 32, 64, 128, 128, (1, 2, 4), (1, 2), (4,), False)),
|
|
("single_seq", False, make(1, 1, 8, 32, 32, 32, (1,), (), (1,), False)),
|
|
("prefill_len8", False, make(5, 8, 16, 512, 512, 64, (1, 2, 4), (), (1, 2, 4), False)),
|
|
]
|
|
|
|
|
|
def _cases_compress_prefill() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Block compression plus the ring-cache tail write, including a window that spans chunks."""
|
|
|
|
def make(dim: int, ratio: int, n_seq: int, cache_size: int, slots: int, spans_chunks: bool):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
qlens = torch.randint(3, 7, (n_seq,), device="cuda", dtype=torch.int32)
|
|
# qstart > 0 means the first compression window reaches into the previous chunk, which
|
|
# is the only way the cache-read branch is exercised.
|
|
qstart = (
|
|
torch.randint(1, 4, (n_seq,), device="cuda", dtype=torch.int32)
|
|
if spans_chunks
|
|
else torch.zeros((n_seq,), device="cuda", dtype=torch.int32)
|
|
)
|
|
klen = qstart + qlens
|
|
# All four prefix sums are exclusive with a leading zero (`n_seq + 1` entries); the
|
|
# kernel's ABI rejects the inclusive form.
|
|
zero = torch.zeros((1,), device="cuda", dtype=torch.int32)
|
|
cu_q = torch.cat([zero, torch.cumsum(qlens, 0, dtype=torch.int32)])
|
|
cu_k = torch.cat([zero, torch.cumsum(klen, 0, dtype=torch.int32)])
|
|
total_q = int(cu_q[-1].item())
|
|
|
|
new_blocks = (klen // ratio - qstart // ratio).tolist()
|
|
reserved = [max(1, b) for b in new_blocks]
|
|
mapping = torch.cat([torch.full((r,), i, device="cuda", dtype=torch.int32) for i, r in enumerate(reserved)])
|
|
# Exclusive prefix sum with a leading zero, i.e. `n_seq + 1` entries: the kernel binds
|
|
# its `num_seqs` symbol from this tensor, so a shorter one fails the ABI check.
|
|
cu_c = torch.zeros((n_seq + 1,), device="cuda", dtype=torch.int32)
|
|
acc = 0
|
|
for i, r in enumerate(reserved):
|
|
cu_c[i] = acc
|
|
acc += r
|
|
cu_c[n_seq] = acc
|
|
# Distinct slots per sequence, as the engine assigns them. Sharing a slot would let one
|
|
# sequence's tail copy overwrite the other's cache, which makes the result depend on
|
|
# kernel interleaving rather than on the code under test.
|
|
assert n_seq <= slots
|
|
state_block_idx = torch.randperm(slots, device="cuda")[:n_seq].to(torch.int32)
|
|
|
|
score = torch.randn((total_q, dim), device="cuda", dtype=torch.float32)
|
|
latent = torch.randn((total_q, dim), device="cuda", dtype=torch.float32)
|
|
state_cache = torch.randn((slots, cache_size, 2, dim), device="cuda", dtype=torch.float32)
|
|
|
|
kv = compress_and_update_state_prefill(
|
|
score,
|
|
latent,
|
|
state_cache,
|
|
state_block_idx,
|
|
mapping,
|
|
cu_q,
|
|
cu_k,
|
|
cu_c,
|
|
ratio,
|
|
None,
|
|
)
|
|
return kv, state_cache
|
|
|
|
return run
|
|
|
|
return [
|
|
("engine_r2", False, make(512, 2, 3, 16, 4, True)),
|
|
("engine_r2_no_prev", False, make(512, 2, 3, 16, 4, False)),
|
|
("r1", False, make(512, 1, 3, 16, 4, True)),
|
|
("single_seq", False, make(512, 2, 1, 8, 1, True)),
|
|
# Tightest legal ring: at least one chunk's worth of tokens, so the tail copy cannot wrap
|
|
# onto a slot the compression step is reading. The interface asserts this relation too.
|
|
("tight_cache", False, make(512, 2, 2, 6, 2, True)),
|
|
]
|
|
|
|
|
|
def _cases_compress_decode() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Per-token cache update and the block-closing compression."""
|
|
|
|
def make(dim: int, ratio: int, total_q: int, max_seqlen_q: int, slots: int, cache_size: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
score = torch.randn((total_q, dim), device="cuda", dtype=torch.float32)
|
|
latent = torch.randn((total_q, dim), device="cuda", dtype=torch.float32)
|
|
state_cache = torch.randn((slots, cache_size, 2, dim), device="cuda", dtype=torch.float32)
|
|
# One distinct state slot per token: the engine gives every sequence its own slot, and
|
|
# sharing one would make the cache writes overlap and the result order-dependent.
|
|
assert total_q <= slots
|
|
state_block_idx = torch.randperm(slots, device="cuda")[:total_q].to(torch.int32)
|
|
context_lens = torch.randint(1, 40, (total_q,), device="cuda", dtype=torch.int32)
|
|
kv = compress_and_update_state_decode(
|
|
score,
|
|
latent,
|
|
state_cache,
|
|
state_block_idx,
|
|
context_lens,
|
|
max_seqlen_q,
|
|
ratio,
|
|
)
|
|
return kv, state_cache
|
|
|
|
return run
|
|
|
|
return [
|
|
("engine_r2", False, make(512, 2, 8, 6, 8, 8)),
|
|
("r1", False, make(512, 1, 8, 6, 8, 8)),
|
|
("single_token", False, make(512, 2, 1, 1, 1, 4)),
|
|
("many_slots", False, make(512, 2, 16, 6, 16, 16)),
|
|
]
|
|
|
|
|
|
def _cases_fused_sample() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Sampler output: log-probabilities and the drawn token.
|
|
|
|
The comparison is not `exact`: a different reduction tree shifts the row sum by an ulp, which
|
|
propagates into the fixed-point cumulative mass, where a column sitting exactly on a draw
|
|
boundary can flip. So log-probabilities are compared with a tolerance and the token count is
|
|
checked as an agreement rate -- the honest standard for a kernel whose answer is a random draw.
|
|
"""
|
|
LogitsTransformInputs, SamplingKernelInputs = _sampler_inputs()
|
|
from deepinfer.model_executor.layers.tilelang_ops.fused_sample import fused_sample
|
|
|
|
def make(vocab: int, rows: int, top_p: float, top_k: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
logits = torch.randn((rows, vocab), device="cuda", dtype=torch.float32) * 2.0
|
|
temperatures = torch.ones(rows, device="cuda", dtype=torch.float32)
|
|
top_ps = torch.full((rows,), top_p, device="cuda", dtype=torch.float32)
|
|
top_ks = torch.full((rows,), top_k, device="cuda", dtype=torch.int32)
|
|
positions = torch.arange(1000, 1000 + rows, device="cuda", dtype=torch.int32)
|
|
# torch has no uint64 arange, so build in int64 and reinterpret: the seeds are a bit
|
|
# pattern, and they exceed 2^63 as uint64.
|
|
seeds = (torch.arange(rows, device="cuda", dtype=torch.int64) * 7919 + 12345).view(torch.uint64)
|
|
logprobs, token, *_ = fused_sample(
|
|
logits,
|
|
SamplingKernelInputs(
|
|
logits=LogitsTransformInputs(temperatures, top_ps, top_ks),
|
|
seq_seeds=seeds,
|
|
sample_positions=positions,
|
|
),
|
|
)
|
|
return logprobs, token
|
|
|
|
return run
|
|
|
|
# The last element is the token agreement floor, and it is a measured number, not a guess.
|
|
#
|
|
# The log-probabilities are compared with a tolerance and the kept set has to match exactly; the
|
|
# token cannot, because the draw is `d % total` and the baseline's rejection limit steps at
|
|
# `total == 2^31`, which `total` sits next to. A few units of rounding error in `total` therefore
|
|
# move to a different draw. The engine's own configuration is unaffected (its totals stay on one
|
|
# side), so it is held at 0.9 and passes with every row matching; the narrow cuts sit either
|
|
# side of that step on 10-13% of rows, so their floor is 0.85.
|
|
return [
|
|
# The engine's own configuration: 4096 columns after the force_topk cut, top_p=0.99, no top_k.
|
|
("engine_cfg", False, make(4096, 8, 0.99, -1), 0.9),
|
|
("engine_cfg_batch32", False, make(4096, 32, 0.99, -1), 0.9),
|
|
# top_k only, and both cuts at once.
|
|
# A narrow cut concentrates the mass in few columns, so the fixed-point cumulative values are
|
|
# coarse relative to the spacing between them and a boundary lands on a different side more
|
|
# often. 128 rows rather than 8, because with 8 a single flip is 12.5% and the rate says
|
|
# nothing.
|
|
("top_k_128", False, make(4096, 128, 1.0, 128), 0.85),
|
|
("top_k_1", False, make(4096, 32, 1.0, 1), 0.85),
|
|
("top_p_and_k", False, make(4096, 128, 0.9, 256), 0.85),
|
|
("no_cut", False, make(4096, 128, 1.0, -1), 0.85),
|
|
# Sharp logits make the distribution peaky, so a small numeric difference cannot move the
|
|
# draw; flat ones put many columns on similar mass and are the harder case.
|
|
("top_p_0.5", False, make(4096, 128, 0.5, -1), 0.85),
|
|
]
|
|
|
|
|
|
def _cases_compute_seqlens_k() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""MTP key-length fixup: fresh rows keep `sk`, continuing rows take the accepted count."""
|
|
|
|
def make(n_seq: int, nextn: int, fresh_frac: float, n_prev: int | None = None):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
prev_n = n_prev if n_prev is not None else n_seq
|
|
seqlens_k = torch.randint(1, 200, (n_seq,), device="cuda", dtype=torch.int32)
|
|
seqlens_q = torch.randint(1, nextn + 2, (n_seq,), device="cuda", dtype=torch.int32)
|
|
# -1 marks a fresh request; anything else indexes the previous batch.
|
|
remap = torch.randint(0, prev_n, (n_seq,), device="cuda", dtype=torch.int32)
|
|
fresh = torch.rand(n_seq, device="cuda") < fresh_frac
|
|
remap = torch.where(fresh, torch.full_like(remap, -1), remap)
|
|
prev_accepted = torch.randint(0, nextn + 1, (prev_n,), device="cuda", dtype=torch.int32)
|
|
compute_seqlens_k_zos(seqlens_k, seqlens_q, remap, prev_accepted, nextn=nextn)
|
|
return seqlens_k
|
|
|
|
return run
|
|
|
|
return [
|
|
("n256_nextn5_mixed", True, make(256, 5, 0.5)),
|
|
("n256_nextn5_all_fresh", True, make(256, 5, 1.0)),
|
|
("n256_nextn5_all_prev", True, make(256, 5, 0.0)),
|
|
("n37_nextn3_mixed", True, make(37, 3, 0.5)),
|
|
("n1_nextn1_fresh", True, make(1, 1, 1.0)),
|
|
("n300_nextn5_small_prev", True, make(300, 5, 0.3, n_prev=7)),
|
|
]
|
|
|
|
|
|
def _cases_confidence_head() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""MTP draft confidence: a dot product, a scaled clamp and a sigmoid."""
|
|
|
|
def make(n: int, n_embed: int, rank: int, temperature: float):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
hidden = torch.randn((n, n_embed), device="cuda", dtype=torch.bfloat16)
|
|
markov = torch.randn((n, rank), device="cuda", dtype=torch.bfloat16) if rank else None
|
|
weight = torch.randn((n_embed + rank, 1), device="cuda", dtype=torch.bfloat16)
|
|
return confidence_head_fused(hidden, weight, temperature, markov)
|
|
|
|
return run
|
|
|
|
return [
|
|
("n64_e5120", False, make(64, 5120, 0, 1.0)),
|
|
("n64_e5120_markov8", False, make(64, 5120, 8, 1.0)),
|
|
("n64_e5120_markov8_cold", False, make(64, 5120, 8, 0.05)),
|
|
("n1_e7168", False, make(1, 7168, 0, 1.0)),
|
|
("n128_e512_markov4", False, make(128, 512, 4, 0.5)),
|
|
]
|
|
|
|
|
|
def _cases_dflash_input_store() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""MTP draft input: reduce the hc_mult copies of each row to one."""
|
|
|
|
def make(n: int, hc_mult: int, n_embed: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
x = torch.randn((n, hc_mult, n_embed), device="cuda", dtype=torch.bfloat16)
|
|
return dflash_input_store(x)
|
|
|
|
return run
|
|
|
|
return [
|
|
("n64_hc1_e5120", True, make(64, 1, 5120)),
|
|
("n64_hc2_e5120", False, make(64, 2, 5120)),
|
|
("n64_hc4_e512", False, make(64, 4, 512)),
|
|
("n1_hc2_e7168", False, make(1, 2, 7168)),
|
|
]
|
|
|
|
|
|
CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
|
|
"store_kv_cache_tl": _cases_store_kv_cache,
|
|
"load_paged_kv_cache_tl": _cases_load_paged_kv_cache,
|
|
"norm_tl": _cases_norm,
|
|
"map_lookup_tl": _cases_map_lookup,
|
|
"expand_sampling_params_tl": _cases_expand_sampling_params,
|
|
"prepare_verification_metadata_tl": _cases_prepare_verification,
|
|
"get_accept_result_tl": _cases_get_accept_result,
|
|
"finalize_verification_tl": _cases_finalize_verification,
|
|
"remap_draft_topk_tl": _cases_remap_draft_topk,
|
|
"fused_norm_rope_hadamard_cast_store_kvcache_tl": _cases_fused_transform,
|
|
"prepare_v4_decode_metadata_tl": _cases_prepare_v4_decode_metadata,
|
|
"compress_and_update_state_prefill_tl": _cases_compress_prefill,
|
|
"compress_and_update_state_decode_tl": _cases_compress_decode,
|
|
"fused_sample_tl": _cases_fused_sample,
|
|
"compute_seqlens_k_zos_tl": _cases_compute_seqlens_k,
|
|
"confidence_head_fused_tl": _cases_confidence_head,
|
|
"dflash_input_store_tl": _cases_dflash_input_store,
|
|
"copy_block_tables_tl": _cases_copy_block_tables,
|
|
"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,
|
|
"get_fused_dflash_get_input_tl": _cases_dflash_get_input,
|
|
"store_mtp_info_tl": _cases_store_mtp_info,
|
|
}
|