mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
29fed050cb
store_kv_cache and load_paged_kv_cache, both byte-for-byte identical to the TileLang baseline on the layouts SM90 uses: main/SWA KV (e4m3 values, ue8m0 scale bytes, 528 B per token) and index KV (128 fp8 + one fp32 scale, 132 B per token). The store also covers fp32 scales, because the index layout uses them. FP4 caches and the packed-ue8m0 index layout are refused instead of approximated. Store cases compare the whole cache buffer, including the `-1` invalid-slot path and a block size (134) that is not a power of two. Load cases cover dequantising to BF16, the index passthrough, and holes expressed as `src_block_id == -1`. Two traps worth recording: The load kernel's `other=` default has to be a float literal. With an fp8 source pointer, `other=0` fails to type-check with a message about casting an int32 tile to fp8e4nv, which points at the wrong line. The load cases originally filled the cache with random bytes. Random bytes decode as fp8 NaN and as NaN scale factors, and NaN compares unequal to itself, so the harness reported mismatches that said nothing about the kernel. The cases now build the cache with `store_kv_cache` and compare index outputs as raw bytes, where a NaN bit pattern is just a byte pattern. 15 kernels and 62 cases agree.
543 lines
26 KiB
Python
543 lines
26 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.cache_common import (
|
|
FP8_FLOAT32_SF_INDEX_KV_FORMAT,
|
|
V41_FORMAT,
|
|
)
|
|
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.load_paged_kv_cache import load_paged_kv_cache
|
|
from deepinfer.model_executor.layers.tilelang_ops.norm import norm
|
|
from deepinfer.model_executor.layers.tilelang_ops.sampler import map_lookup
|
|
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 _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."""
|
|
|
|
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, total_q: int, max_len: int):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
seqlen_q = torch.randint(1, max_len + 1, (req_cnt,), device="cuda", dtype=torch.int32)
|
|
ftp = torch.zeros((req_cnt,), device="cuda", dtype=torch.int32)
|
|
ftp[1:] = torch.cumsum(seqlen_q, 0, dtype=torch.int32)[:-1]
|
|
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(max_len)
|
|
kernel(ftp, seqlen_q, row_to_req, local_step, is_verify)
|
|
return row_to_req, local_step, is_verify
|
|
|
|
return run
|
|
|
|
return [("n4_q24", True, make(4, 24, 6)), ("n1_q1", True, make(1, 1, 1)), ("n8_q48", True, make(8, 48, 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."""
|
|
|
|
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.
|
|
"""
|
|
|
|
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)),
|
|
]
|
|
|
|
|
|
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,
|
|
"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,
|
|
}
|