diff --git a/H20-TRITON-PLAN.md b/H20-TRITON-PLAN.md index 420bbc0..d28e8c7 100644 --- a/H20-TRITON-PLAN.md +++ b/H20-TRITON-PLAN.md @@ -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 和 diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/load_paged_kv_cache.py b/deepinfer/model_executor/layers/tilelang_ops/triton/load_paged_kv_cache.py new file mode 100644 index 0000000..af7909e --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/load_paged_kv_cache.py @@ -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 diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/store_kv_cache.py b/deepinfer/model_executor/layers/tilelang_ops/triton/store_kv_cache.py new file mode 100644 index 0000000..9cf97c7 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/store_kv_cache.py @@ -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 diff --git a/scripts/ab_cases.py b/scripts/ab_cases.py index 775180f..8f31c05 100644 --- a/scripts/ab_cases.py +++ b/scripts/ab_cases.py @@ -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,