mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 16:56:21 +08:00
b4f7a42652
`h20_ops_ab.py` assumed the TileLang backend is importable, so on a vendor image it could not start at all -- which is exactly where the Triton backend needs checking. Two modes fix that: * `--dump DIR` saves the selected backend's output per case. * `--compare DIR` checks against dumps written elsewhere. Together they answer "do triton 3.5.1 and 3.7.1 agree" without either host having TileLang. Case builders that genuinely need a package the host lacks now report a skip instead of aborting the run, so partial coverage is visible rather than fatal; the import of the sampler's C++-backed inputs moved into the one builder that uses it for the same reason. This is what showed the 19 ported kernels produce element-wise identical output across triton 3.5.1, 3.6.0 and 3.7.1 -- 115 cases, no differences.
1071 lines
52 KiB
Python
1071 lines
52 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_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.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_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_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,
|
|
"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,
|
|
"store_mtp_info_tl": _cases_store_mtp_info,
|
|
}
|