mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
feat(triton): port the paged KV cache read and write
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.
This commit is contained in:
loopyt
parent
4f3edbf604
commit
29fed050cb
@ -108,12 +108,15 @@
|
||||
|---|---|---|
|
||||
| P0 基础设施 | ✅ | 后端选择 + fallback + A/B 骨架 |
|
||||
| P1 轻量 kernel | ✅ 13/13 | gather、gather_input_ids_zos、gather_engram_tokens_zos、expand_swa_index、kv_index_cat、convert_kv_index、load_kv_index、store_mtp_info、norm、map_lookup、expand_sampling_params、prepare_verification_metadata、copy_block_tables |
|
||||
| P2 KV cache | ⬜ | |
|
||||
| P2 KV cache | ✅ 2/3 | store_kv_cache、load_paged_kv_cache（`fused_transform` 未做） |
|
||||
| P3 MTP 链路 | ⬜ | |
|
||||
| P4 采样 | ⬜ | |
|
||||
| P5 端到端 | ⬜ | |
|
||||
|
||||
A/B 现状：**13 个 kernel、53 个用例全过**（`python scripts/h20_ops_ab.py`）。
|
||||
A/B 现状：**15 个 kernel、62 个用例全过**（`python scripts/h20_ops_ab.py`）。
|
||||
|
||||
KV cache 的两个 kernel 都是**整块 buffer 逐字节一致**（store 覆盖 ue8m0 与 fp32 两种 scale
|
||||
布局、`-1` 无效槽位；load 覆盖 do_cast 反量化、index 直通、以及 `src_block_id == -1` 的洞）。
|
||||
|
||||
`norm` 是这批里最关键的一个，两点值得记：
|
||||
- **量化前必须先舍回 bf16**。TileLang 版是 `cast(cast(y, bf16), fp32)` 之后才算 amax 和
|
||||
|
||||
@ -0,0 +1,246 @@
|
||||
"""Paged KV cache read over physical block ranges, in Triton.
|
||||
|
||||
Three consumer shapes, matching the TileLang kernel:
|
||||
|
||||
* quantised with cast -- dequantise to BF16 on the way out (main/SWA KV, 528 B per token);
|
||||
* index layout with fp32 scales -- pass the stored bytes and the per-token scale through
|
||||
untouched, because the consumer is a kernel that wants raw bytes plus scales, not BF16. This is
|
||||
the only index variant SM90 uses, and the packed-ue8m0 variants are refused rather than
|
||||
half-implemented;
|
||||
* plain BF16 -- a straight copy.
|
||||
|
||||
Segments are `[src_block_id, src_begin, src_end, dst_offset]` rows; `src_block_id == -1` writes
|
||||
zeros over its destination range, which is how the caller expresses a hole.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _load_cast_kernel(
|
||||
SEGMENTS,
|
||||
OUT,
|
||||
VALUE,
|
||||
SF,
|
||||
value_block_stride,
|
||||
value_token_stride,
|
||||
sf_block_stride,
|
||||
sf_token_stride,
|
||||
out_stride,
|
||||
block_size,
|
||||
groups: tl.constexpr,
|
||||
group_size: tl.constexpr,
|
||||
sf_ue8m0: tl.constexpr,
|
||||
BLOCK_T: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
segment = tl.program_id(0)
|
||||
token_tile = tl.program_id(1)
|
||||
|
||||
src_block = tl.load(SEGMENTS + segment * 4 + 0)
|
||||
src_begin = tl.load(SEGMENTS + segment * 4 + 1)
|
||||
src_end = tl.load(SEGMENTS + segment * 4 + 2)
|
||||
dst_offset = tl.load(SEGMENTS + segment * 4 + 3)
|
||||
|
||||
offsets = token_tile * BLOCK_T + tl.arange(0, BLOCK_T)
|
||||
# Tokens before `src_begin` and past `src_end` are other segments' business; the tail of the
|
||||
# last tile is past the block.
|
||||
in_range = (offsets < block_size) & (offsets >= src_begin) & (offsets < src_end)
|
||||
dst = dst_offset + offsets - src_begin
|
||||
|
||||
chans = tl.arange(0, BLOCK_D)
|
||||
dim = groups * group_size
|
||||
chan_live = chans < dim
|
||||
mask = in_range[:, None] & chan_live[None, :]
|
||||
|
||||
if src_block < 0:
|
||||
# A negative block id means "this range is a hole": zeros, not stale cache contents.
|
||||
tl.store(OUT + dst[:, None] * out_stride + chans[None, :], tl.zeros((BLOCK_T, BLOCK_D), tl.bfloat16), mask=mask)
|
||||
else:
|
||||
sf_ptr = SF + src_block * sf_block_stride + offsets[:, None] * sf_token_stride + (chans // group_size)[None, :]
|
||||
sf_raw = tl.load(sf_ptr, mask=mask, other=0.0)
|
||||
if sf_ue8m0: # noqa: SIM108 - the two arms are different bitcasts, not one expression
|
||||
# because the byte is not a float in the cache.
|
||||
scale = (sf_raw.to(tl.uint32) << 23).to(tl.float32, bitcast=True)
|
||||
else:
|
||||
scale = sf_raw.to(tl.float32)
|
||||
|
||||
value_ptr = VALUE + src_block * value_block_stride + offsets[:, None] * value_token_stride + chans[None, :]
|
||||
# `other` must be a float literal: for an fp8 source an int default fails to type-check.
|
||||
value = tl.load(value_ptr, mask=mask, other=0.0).to(tl.float32)
|
||||
tl.store(OUT + dst[:, None] * out_stride + chans[None, :], (value * scale).to(tl.bfloat16), mask=mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _load_index_kernel(
|
||||
SEGMENTS,
|
||||
OUT,
|
||||
OUT_SF,
|
||||
VALUE,
|
||||
SF,
|
||||
value_block_stride,
|
||||
value_token_stride,
|
||||
sf_block_stride,
|
||||
sf_token_stride,
|
||||
out_stride,
|
||||
out_sf_stride,
|
||||
block_size,
|
||||
BLOCK_T: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
segment = tl.program_id(0)
|
||||
token_tile = tl.program_id(1)
|
||||
|
||||
src_block = tl.load(SEGMENTS + segment * 4 + 0)
|
||||
src_begin = tl.load(SEGMENTS + segment * 4 + 1)
|
||||
src_end = tl.load(SEGMENTS + segment * 4 + 2)
|
||||
dst_offset = tl.load(SEGMENTS + segment * 4 + 3)
|
||||
|
||||
offsets = token_tile * BLOCK_T + tl.arange(0, BLOCK_T)
|
||||
in_range = (offsets < block_size) & (offsets >= src_begin) & (offsets < src_end)
|
||||
dst = dst_offset + offsets - src_begin
|
||||
|
||||
chans = tl.arange(0, BLOCK_D)
|
||||
# Bytes are carried as uint8 views on both sides so the copy is byte-exact.
|
||||
if src_block < 0:
|
||||
tl.store(OUT + dst[:, None] * out_stride + chans[None, :], tl.zeros((BLOCK_T, BLOCK_D), tl.uint8), mask=in_range[:, None])
|
||||
if token_tile == 0:
|
||||
tl.store(OUT_SF + dst * out_sf_stride, tl.zeros((BLOCK_T,), tl.float32), mask=in_range)
|
||||
else:
|
||||
value_ptr = VALUE + src_block * value_block_stride + offsets[:, None] * value_token_stride + chans[None, :]
|
||||
value = tl.load(value_ptr, mask=in_range[:, None], other=0)
|
||||
tl.store(OUT + dst[:, None] * out_stride + chans[None, :], value, mask=in_range[:, None])
|
||||
if token_tile == 0:
|
||||
# One fp32 scale per token, so a scalar store per row.
|
||||
sf_ptr = SF + src_block * sf_block_stride + offsets * sf_token_stride
|
||||
tl.store(OUT_SF + dst * out_sf_stride, tl.load(sf_ptr, mask=in_range, other=0.0), mask=in_range)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _load_bf16_kernel(
|
||||
SEGMENTS,
|
||||
OUT,
|
||||
VALUE,
|
||||
value_block_stride,
|
||||
value_token_stride,
|
||||
out_stride,
|
||||
block_size,
|
||||
BLOCK_T: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
segment = tl.program_id(0)
|
||||
token_tile = tl.program_id(1)
|
||||
|
||||
src_block = tl.load(SEGMENTS + segment * 4 + 0)
|
||||
src_begin = tl.load(SEGMENTS + segment * 4 + 1)
|
||||
src_end = tl.load(SEGMENTS + segment * 4 + 2)
|
||||
dst_offset = tl.load(SEGMENTS + segment * 4 + 3)
|
||||
|
||||
offsets = token_tile * BLOCK_T + tl.arange(0, BLOCK_T)
|
||||
in_range = (offsets < block_size) & (offsets >= src_begin) & (offsets < src_end)
|
||||
dst = dst_offset + offsets - src_begin
|
||||
|
||||
chans = tl.arange(0, BLOCK_D)
|
||||
if src_block < 0:
|
||||
tl.store(OUT + dst[:, None] * out_stride + chans[None, :], tl.zeros((BLOCK_T, BLOCK_D), tl.bfloat16), mask=in_range[:, None])
|
||||
else:
|
||||
ptr = VALUE + src_block * value_block_stride + offsets[:, None] * value_token_stride + chans[None, :]
|
||||
tl.store(OUT + dst[:, None] * out_stride + chans[None, :], tl.load(ptr, mask=in_range[:, None], other=0), mask=in_range[:, None])
|
||||
|
||||
|
||||
def load_paged_kv_cache_tl(
|
||||
layout: Any,
|
||||
block_size: int,
|
||||
fp4_value_block_stride: int = 0,
|
||||
fp4_value_stride: int = 0,
|
||||
fp4_scale_stride: int = 0,
|
||||
) -> Any:
|
||||
"""Build the load kernel for one cache layout."""
|
||||
del fp4_value_block_stride, fp4_value_stride, fp4_scale_stride
|
||||
if layout.quantized_dtype == "e2m1":
|
||||
raise NotImplementedError("triton load_paged_kv_cache does not implement FP4 caches")
|
||||
is_quantized = layout.quantized_dtype is not None
|
||||
is_index = is_quantized and not layout.do_cast
|
||||
if is_index and layout.sf_format == "ue8m0":
|
||||
raise NotImplementedError("triton load_paged_kv_cache implements the fp32 index scale layout only")
|
||||
|
||||
groups = layout.num_scales
|
||||
group_size = layout.elements_per_sf
|
||||
is_index_sf_scalar = is_index and layout.sf_storage_elements == 1
|
||||
|
||||
block_t = block_size if (is_index or not is_quantized) else min(block_size, 32)
|
||||
dim = layout.quantized_elements if is_quantized else layout.num_elements
|
||||
|
||||
def launch(
|
||||
load_segments: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
out_sf: torch.Tensor,
|
||||
kv_cache_value: torch.Tensor,
|
||||
kv_cache_sf: torch.Tensor,
|
||||
) -> None:
|
||||
num_segments = load_segments.shape[0]
|
||||
if num_segments == 0:
|
||||
return
|
||||
grid = (num_segments, cdiv(block_size, block_t))
|
||||
|
||||
if layout.do_cast:
|
||||
_load_cast_kernel[grid](
|
||||
load_segments,
|
||||
out,
|
||||
kv_cache_value,
|
||||
kv_cache_sf,
|
||||
kv_cache_value.stride(0),
|
||||
kv_cache_value.stride(1),
|
||||
kv_cache_sf.stride(0),
|
||||
kv_cache_sf.stride(1),
|
||||
out.stride(0),
|
||||
block_size,
|
||||
groups=groups,
|
||||
group_size=group_size,
|
||||
sf_ue8m0=layout.sf_format == "ue8m0",
|
||||
BLOCK_T=block_t,
|
||||
BLOCK_D=dim,
|
||||
num_warps=8,
|
||||
)
|
||||
elif is_index:
|
||||
assert is_index_sf_scalar, "only one scale per token is supported"
|
||||
# Bytes on both sides: the cache and the result are the same fp8 payload, and copying
|
||||
# raw bytes avoids inventing a value dtype for the copy.
|
||||
_load_index_kernel[grid](
|
||||
load_segments,
|
||||
out.view(torch.uint8),
|
||||
out_sf,
|
||||
kv_cache_value.view(torch.uint8),
|
||||
kv_cache_sf,
|
||||
kv_cache_value.stride(0),
|
||||
kv_cache_value.stride(1),
|
||||
kv_cache_sf.stride(0),
|
||||
kv_cache_sf.stride(1),
|
||||
out.stride(0),
|
||||
out_sf.stride(0),
|
||||
block_size,
|
||||
BLOCK_T=block_t,
|
||||
BLOCK_D=layout.quantized_bytes,
|
||||
num_warps=8,
|
||||
)
|
||||
else:
|
||||
_load_bf16_kernel[grid](
|
||||
load_segments,
|
||||
out,
|
||||
kv_cache_value,
|
||||
kv_cache_value.stride(0),
|
||||
kv_cache_value.stride(1),
|
||||
out.stride(0),
|
||||
block_size,
|
||||
BLOCK_T=block_t,
|
||||
BLOCK_D=dim,
|
||||
num_warps=8,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -0,0 +1,187 @@
|
||||
"""Paged KV cache write, in Triton.
|
||||
|
||||
Covers the two shapes the SM90 path uses: quantised FP8 with ue8m0 scales (main/SWA KV) and a
|
||||
plain BF16 copy. The FP4 layouts the other platforms use are rejected rather than approximated.
|
||||
|
||||
The store is deliberately direct: for each token, one pass computes the per-group amax, derives the
|
||||
power-of-two multiplier and the ue8m0 byte from it, and writes both the quantised bytes and the
|
||||
scales. The amax is taken over the BF16 source values, which is exact -- a maximum of BF16 values
|
||||
is itself a BF16 value -- so there is no rounding question between the two backends.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.common import FLOAT8E4M3_MAX, FP8_AMAX_MARGIN
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.common import next_power_of_2
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _store_fp8_kernel(
|
||||
KV,
|
||||
VALUE,
|
||||
SF,
|
||||
SLOT_MAPPING,
|
||||
n,
|
||||
kv_stride,
|
||||
value_block_stride,
|
||||
value_token_stride,
|
||||
sf_block_stride,
|
||||
sf_token_stride,
|
||||
block_size,
|
||||
groups: tl.constexpr,
|
||||
group_size: tl.constexpr,
|
||||
valid_groups: tl.constexpr,
|
||||
sf_dim: tl.constexpr,
|
||||
sf_ue8m0: tl.constexpr,
|
||||
fp8_max: tl.constexpr,
|
||||
amax_margin: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
):
|
||||
rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
live = rows < n
|
||||
|
||||
slots = tl.load(SLOT_MAPPING + rows, mask=live, other=-1)
|
||||
write = live & (slots >= 0)
|
||||
|
||||
chans = tl.arange(0, groups * group_size)
|
||||
chan_live = chans < valid_groups * group_size
|
||||
values = tl.load(KV + rows[:, None] * kv_stride + chans[None, :], mask=write[:, None] & chan_live[None, :], other=0.0)
|
||||
values = values.to(tl.float32)
|
||||
|
||||
padded = tl.reshape(values, (BLOCK_M, groups, group_size))
|
||||
amax = tl.max(tl.abs(padded), axis=2)
|
||||
amax = tl.maximum(amax, amax_margin)
|
||||
# exp = -ceil(log2(amax / 448)); multiplier is 2^exp and the stored ue8m0 byte is -exp + 127.
|
||||
scale = amax / fp8_max
|
||||
bits = scale.to(tl.uint32, bitcast=True)
|
||||
exp = 126 - ((bits - 1) >> 23).to(tl.int32)
|
||||
multiplier = ((exp + 127) << 23).to(tl.float32, bitcast=True)
|
||||
if sf_ue8m0: # noqa: SIM108 - the two arms store different dtypes, not one expression
|
||||
sf_value = (-exp + 127).to(tl.uint8)
|
||||
else:
|
||||
# An fp32 scale holds the dequant factor itself: 2^-exp.
|
||||
sf_value = ((-exp + 127) << 23).to(tl.float32, bitcast=True)
|
||||
|
||||
quant = padded * multiplier[:, :, None]
|
||||
quant = tl.reshape(quant, (BLOCK_M, groups * group_size)).to(tl.float8e4nv)
|
||||
|
||||
block_of = slots // block_size
|
||||
token_of = slots % block_size
|
||||
value_ptr = VALUE + block_of[:, None] * value_block_stride + token_of[:, None] * value_token_stride
|
||||
tl.store(value_ptr + chans[None, :], quant, mask=write[:, None] & chan_live[None, :])
|
||||
|
||||
sf_chans = tl.arange(0, groups)
|
||||
sf_ptr = SF + block_of[:, None] * sf_block_stride + token_of[:, None] * sf_token_stride
|
||||
# Groups beyond `valid_groups` exist when the layout pads the scale row; they are zeroed so the
|
||||
# cache holds a defined value rather than whatever was there before.
|
||||
stored = tl.where(sf_chans[None, :] < valid_groups, sf_value, 0)
|
||||
tl.store(sf_ptr + sf_chans[None, :], stored, mask=write[:, None] & (sf_chans[None, :] < sf_dim))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _store_bf16_kernel(
|
||||
KV,
|
||||
VALUE,
|
||||
SLOT_MAPPING,
|
||||
n,
|
||||
kv_stride,
|
||||
value_block_stride,
|
||||
value_token_stride,
|
||||
block_size,
|
||||
d: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
):
|
||||
rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
col_block = tl.program_id(1)
|
||||
live = rows < n
|
||||
slots = tl.load(SLOT_MAPPING + rows, mask=live, other=-1)
|
||||
write = live & (slots >= 0)
|
||||
|
||||
cols = col_block * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
col_live = cols < d
|
||||
values = tl.load(KV + rows[:, None] * kv_stride + cols[None, :], mask=write[:, None] & col_live[None, :], other=0.0)
|
||||
|
||||
block_of = slots // block_size
|
||||
token_of = slots % block_size
|
||||
ptr = VALUE + block_of[:, None] * value_block_stride + token_of[:, None] * value_token_stride
|
||||
tl.store(ptr + cols[None, :], values, mask=write[:, None] & col_live[None, :])
|
||||
|
||||
|
||||
def store_kv_cache_tl(
|
||||
layout: Any,
|
||||
block_size: int,
|
||||
fp4_value_block_stride: int = 0,
|
||||
fp4_value_stride: int = 0,
|
||||
fp4_scale_stride: int = 0,
|
||||
) -> Any:
|
||||
"""Build the store kernel for one cache layout."""
|
||||
del fp4_value_block_stride, fp4_value_stride, fp4_scale_stride
|
||||
if layout.quantized_dtype == "e2m1":
|
||||
raise NotImplementedError("triton store_kv_cache does not implement FP4 caches")
|
||||
is_quantized = layout.quantized_dtype is not None
|
||||
is_fp8 = layout.quantized_dtype == "e4m3"
|
||||
assert not is_quantized or is_fp8 or layout.bf16_elements == layout.num_elements
|
||||
|
||||
groups = layout.num_scales
|
||||
group_size = layout.elements_per_sf
|
||||
valid_groups = (layout.quantized_elements + group_size - 1) // group_size
|
||||
sf_dim = layout.sf_storage_elements if is_quantized else 0
|
||||
block_m = 8
|
||||
|
||||
def launch(
|
||||
kv: torch.Tensor,
|
||||
kv_cache_value: torch.Tensor,
|
||||
kv_cache_sf: torch.Tensor,
|
||||
slot_mapping: torch.Tensor,
|
||||
) -> None:
|
||||
n = slot_mapping.shape[0]
|
||||
if n == 0:
|
||||
return
|
||||
if is_quantized:
|
||||
assert layout.sf_format in ("ue8m0", "fp32"), f"unsupported scale format {layout.sf_format}"
|
||||
_store_fp8_kernel[(triton.cdiv(n, block_m),)](
|
||||
kv,
|
||||
kv_cache_value,
|
||||
kv_cache_sf,
|
||||
slot_mapping,
|
||||
n,
|
||||
kv.stride(0),
|
||||
kv_cache_value.stride(0),
|
||||
kv_cache_value.stride(1),
|
||||
kv_cache_sf.stride(0),
|
||||
kv_cache_sf.stride(1),
|
||||
block_size,
|
||||
groups=groups,
|
||||
group_size=group_size,
|
||||
valid_groups=valid_groups,
|
||||
sf_dim=sf_dim,
|
||||
sf_ue8m0=layout.sf_format == "ue8m0",
|
||||
fp8_max=FLOAT8E4M3_MAX,
|
||||
amax_margin=FP8_AMAX_MARGIN,
|
||||
BLOCK_M=block_m,
|
||||
num_warps=8,
|
||||
)
|
||||
else:
|
||||
d = kv.shape[1]
|
||||
block_d = min(next_power_of_2(d), 2048)
|
||||
_store_bf16_kernel[(triton.cdiv(n, block_m), triton.cdiv(d, block_d))](
|
||||
kv,
|
||||
kv_cache_value,
|
||||
slot_mapping,
|
||||
n,
|
||||
kv.stride(0),
|
||||
kv_cache_value.stride(0),
|
||||
kv_cache_value.stride(1),
|
||||
block_size,
|
||||
d=d,
|
||||
BLOCK_D=block_d,
|
||||
BLOCK_M=block_m,
|
||||
num_warps=8,
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -16,14 +16,20 @@ 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
|
||||
|
||||
@ -448,7 +454,78 @@ def _cases_copy_block_tables() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
]
|
||||
|
||||
|
||||
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,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user