diff --git a/pegainfer-gemma4/kernels/README.md b/pegainfer-gemma4/kernels/README.md new file mode 100644 index 000000000..b6552470c --- /dev/null +++ b/pegainfer-gemma4/kernels/README.md @@ -0,0 +1,52 @@ +# Gemma 4 TileLang kernels + +**TL;DR**: `generate.py` AOT-compiles the hd512 global-attention prefill in +`tilelang_defs.py` into one CUDA file that `pegainfer-kernels/build.rs` hands +to nvcc under the `gemma4` feature — three tiers (generate, pre-generated, +stub), and the generated CUDA is a Cargo `OUT_DIR` artifact that is never +checked in. Unlike the K3 families this one lowers to TMA, so the launcher +builds the descriptors itself from parameters recovered out of the lowered +host stub. + +## What lives here + +| File | Role | +| --- | --- | +| `tilelang_defs.py` | The kernel, authored here. Upstream has nothing for this head dim on SM90. | +| `generate.py` | Lowers it, recovers the launch geometry and the TMA descriptor parameters, emits the `.cu` with a hand-written launcher. | + +## Why one instantiation is enough + +Every shape a TileLang kernel declares is a compile dimension, and a serving +step's packed query rows, page-table length and pool size are all run-time +quantities. They are nevertheless a single instantiation, because the declared +extents reach the generated code only as bounds guards — two lowerings that +differ solely in them differ in nothing else — so declaring the serving +arena's maxima lets every smaller step through. The real bounds are the walk's +own trip count and the predicated store. + +The two tensors the lowering reads through TMA are separate: their extents +live in the descriptors, which the launcher builds per call from the arguments +it is given, so those are always the step's own. + +## The descriptor parameters are recovered, not written down + +TileLang exposes no accessor for the launch it baked into the host stub, so +`generate.py` parses the packed-call argument stack — the same technique the +K3 generator uses for its launch geometry, extended to the descriptor builds. +Each parameter is then mapped to its driver enum through a table with no +default, and the tensor and descriptor names are bound through tables with no +default either, so a codegen change fails generation instead of silently +encoding a stale descriptor. + +The K3 generator refuses a TMA-lowered body outright, for exactly the reason +this file exists: its launchers bind plain pointers and the requested thread +count, and a warp-specialized kernel accepts neither. + +## Gates + +The kernel's numerics are gated against the serving path's own reference, not +against random tensors: paged output is bit-identical to the contiguous form +over scattered pages and partial final pages, and a ragged batch matches an +fp32 reference per request while leaving every row past the batch untouched — +those rows are the decode rows sharing a mixed step's output buffer. diff --git a/pegainfer-gemma4/kernels/generate.py b/pegainfer-gemma4/kernels/generate.py new file mode 100644 index 000000000..3b21b6cfb --- /dev/null +++ b/pegainfer-gemma4/kernels/generate.py @@ -0,0 +1,635 @@ +"""AOT-compile the Gemma 4 hd512 prefill kernel into one CUDA file. + +`pegainfer-kernels/build.rs` runs this under the `gemma4` feature and hands +the result to nvcc. It prints the same `KEY=VALUE` manifest every TileLang +family prints, and mirrors it into `manifest.txt` so a build host can consume +a pre-generated directory without TileLang installed. + +What makes this family different from the K3 one is the lowering: the key and +query loads become bulk copies, so TileLang passes those two tensors as TMA +descriptors instead of pointers and adds a producer warpgroup to the block. +The descriptors are the launcher's to build, and every parameter it needs is +a constant in the lowered host stub — recovered here rather than assumed, so +a codegen change fails the build instead of encoding a stale descriptor. + +The declared tensor extents are the serving arena's maxima. They reach the +device code only as bounds guards, so any smaller step passes them; the two +TMA tensors carry their real extents in the descriptors, which the launcher +builds per call from the arguments it is given. +""" + +from __future__ import annotations + +import argparse +import re +import shutil +from dataclasses import dataclass +from pathlib import Path + +import tilelang +import tilelang_defs as defs +from tilelang.env import CUTLASS_INCLUDE_DIR, TILELANG_TEMPLATE_PATH + +ENTRY_SYMBOL = "main_kernel" +KERNEL_MARKER = 'extern "C" __global__ void' +LAUNCHER = "gemma4_hd512_prefill_varlen" +CU_STEM = "gemma4_hd512_prefill" +# TileLang names every entry point `main_kernel` with external C linkage, so +# the emitted body is renamed before it can collide with another family's. +KERNEL_SYMBOL = f"{CU_STEM}_kernel" + +# The launcher's C parameters, without the trailing stream the consumer adds. +# The emitted definition and the manifest line both come from here, so the +# stub the build script writes when this kernel is absent cannot drift from +# the real one: C has no mangling, and a drifted pair would link silently. +LAUNCHER_PARAMS = [ + ("const void*", "q"), + ("const void*", "kv"), + ("const int*", "page_indices"), + ("const int*", "page_indptr"), + ("const int*", "q_indptr"), + ("const int*", "host_q_indptr"), + ("const int*", "last_page_len"), + ("void*", "out"), + ("int", "batch"), + ("int", "q_rows"), + ("int", "pool_rows"), + ("int", "rows_per_page"), + ("int", "layer_row"), + ("int", "page_size"), + ("int", "num_qo_heads"), + ("int", "num_kv_heads"), + ("float", "sm_scale"), +] + +# TileLang's `debug.h` *defines* `debug_print_msg` and the `uint16_t` +# specialization of `debug_print_buffer_value` with external linkage, so every +# translation unit that includes it exports the same two symbols. This kernel +# calls neither, and a binary that also links a K3 family would get duplicate +# definitions, so this unit is given privately named copies. +DEBUG_HEADER = "#include " +DEBUG_HELPERS = ("debug_print_msg", "debug_print_buffer_value") + +# The model's global-attention geometry. Shapes are compile dimensions, so +# these are the kernel's identity, not run-time inputs. +HEADS = 32 +GROUPS = 8 +HEAD_DIM = 512 +PAGE_SIZE = 64 + +# The serving arena's maxima, in the units each tensor is indexed in. A step +# is always smaller; see the module docstring. +MAX_BATCH = 8 +CEILING = 262144 +SLOTS = 16 +Q_ROWS = CEILING + defs.BLOCK_M +POOL_PAGES = SLOTS * (CEILING // PAGE_SIZE) + 1 +PAGE_TABLE_LEN = SLOTS * (CEILING // PAGE_SIZE) +# The pool is addressed in rows of [kv_heads, head_dim]; one page holds every +# layer's K and V, and the deepest checkpoint this line serves sets the bound. +MAX_LAYERS = 64 +POOL_ROWS = POOL_PAGES * MAX_LAYERS * 2 * PAGE_SIZE +# The launcher takes both row counts as C ints; a configuration that outgrew +# one would index past the guards rather than fail. +assert Q_ROWS < 2**31, Q_ROWS +assert POOL_ROWS < 2**31, POOL_ROWS + +# Past 48 KiB a kernel has to opt into its dynamic shared memory per symbol. +MAX_STATIC_SMEM = 48 * 1024 +# SM90's per-block ceiling. A recovered size above it would launch-fail. +MAX_DYNAMIC_SMEM = 227 * 1024 + +TENSORMAP_BUILDER = "__tvm_tensormap_create_tiled" + +# A pass config can reach nvcc's command line, which TileLang's JIT passes and +# an AOT build would not: without --use_fast_math the object sits half a ULP +# from the gated kernel. Every config is classified here or generation stops. +PASS_CONFIG_NVCC_FLAG = {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: "--use_fast_math"} + +# TVM's tensormap codes, mapped to the driver enums the launcher passes. Every +# code the stub can carry is listed; an unmapped one is a codegen change. +TENSORMAP_DTYPE = {9: "CU_TENSOR_MAP_DATA_TYPE_BFLOAT16"} +TENSORMAP_INTERLEAVE = {0: "CU_TENSOR_MAP_INTERLEAVE_NONE"} +TENSORMAP_SWIZZLE = { + 0: "CU_TENSOR_MAP_SWIZZLE_NONE", + 1: "CU_TENSOR_MAP_SWIZZLE_32B", + 2: "CU_TENSOR_MAP_SWIZZLE_64B", + 3: "CU_TENSOR_MAP_SWIZZLE_128B", +} +TENSORMAP_L2 = { + 0: "CU_TENSOR_MAP_L2_PROMOTION_NONE", + 1: "CU_TENSOR_MAP_L2_PROMOTION_L2_64B", + 2: "CU_TENSOR_MAP_L2_PROMOTION_L2_128B", + 3: "CU_TENSOR_MAP_L2_PROMOTION_L2_256B", +} +TENSORMAP_OOB = {0: "CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE"} + +_SLOT_INT = re.compile( + r"\(\(\(TVMFFIAny\*\)stack_ffi_any\)\[(\d+)\]\.v_int64\) = \(\(int64_t\)(-?\d+)\);" +) +_SLOT_PTR = re.compile(r"\(\(\(TVMFFIAny\*\)stack_ffi_any\)\[(\d+)\]\.v_ptr\) = (\w+);") +_PACKED_CALL = re.compile( + r"TVMFFIFunctionCall\((\w+?)_packed, \(TVMFFIAny\*\) stack_ffi_any, (\d+)," +) + + +@dataclass(frozen=True) +class TensorMap: + """One recovered `cuTensorMapEncodeTiled` call, still in TVM's spelling.""" + + name: str + dtype: int + rank: int + tensor: str + dims: tuple[int, ...] + strides: tuple[int, ...] + box: tuple[int, ...] + element_strides: tuple[int, ...] + interleave: int + swizzle: int + l2_promotion: int + oob_fill: int + + +def read_host_stub(kernel) -> tuple[list[TensorMap], int, int]: + """Recover the descriptor builds, the block width and the dynamic smem. + + TileLang bakes the launch into the host stub as a packed-call argument + stack and exposes no accessor for it, hence the parse. Slots persist + across calls, so only the values a call writes itself are its own; the + entry call's grid and scalars are run-time expressions and are read from + its arguments instead, which is why they never appear here. + """ + slots: dict[int, object] = {} + maps: list[TensorMap] = [] + launch: list[object] | None = None + for line in kernel.get_host_source().splitlines(): + match = _SLOT_INT.search(line) + if match: + slots[int(match.group(1))] = int(match.group(2)) + continue + match = _SLOT_PTR.search(line) + if match: + slots[int(match.group(1))] = match.group(2) + continue + match = _PACKED_CALL.search(line) + if not match: + continue + callee, count = match.group(1), int(match.group(2)) + args = [slots.get(i) for i in range(count)] + if callee == TENSORMAP_BUILDER: + maps.append(parse_tensormap(args)) + elif callee == ENTRY_SYMBOL: + launch = args + + if not maps: + raise RuntimeError( + "the lowering built no TMA descriptor; this launcher exists only " + "because it does, so the parameter list it binds is wrong now" + ) + if launch is None: + raise RuntimeError("could not recover the entry call from the host stub") + # Block x/y/z then the dynamic smem, which TileLang omits when it is zero; + # this kernel's four shared buffers make that impossible, so a tail without + # the block's unit y and z is a codegen change. + tail = list(launch[-4:]) + if tail[1:3] != [1, 1]: + raise RuntimeError( + f"the launch tail {tail} is not (block, 1, 1, dynamic shared); the " + "lowering changed its geometry or dropped its shared memory" + ) + block, smem = tail[0], tail[3] + if not isinstance(block, int) or not isinstance(smem, int): + raise TypeError(f"launch geometry has non-constant entries: {tail}") + if smem > MAX_DYNAMIC_SMEM: + raise RuntimeError( + f"the lowering wants {smem} B of dynamic shared memory, past the " + f"{MAX_DYNAMIC_SMEM} B a block can be given" + ) + return maps, block, smem + + +def parse_tensormap(args: list) -> TensorMap: + """Split one builder call's flat argument list by its recovered rank.""" + name, dtype, rank, tensor = args[0], args[1], args[2], args[3] + if not isinstance(rank, int) or rank < 2: + raise RuntimeError(f"descriptor {name} has a non-constant rank: {rank}") + at = 4 + dims = tuple(args[at : at + rank]) + at += rank + # The builder carries one stride per dimension, innermost first, and the + # innermost one is the element size. The driver API takes the rest. + strides = tuple(args[at : at + rank]) + at += rank + box = tuple(args[at : at + rank]) + at += rank + element_strides = tuple(args[at : at + rank]) + at += rank + interleave, swizzle, l2_promotion, oob_fill = args[at : at + 4] + values = ( + *dims, + *strides, + *box, + *element_strides, + interleave, + swizzle, + l2_promotion, + oob_fill, + ) + if any(not isinstance(value, int) for value in values): + raise RuntimeError(f"descriptor {name} has non-constant parameters: {values}") + return TensorMap( + name=str(name), + dtype=dtype, + rank=rank, + tensor=str(tensor), + dims=dims, + strides=strides, + box=box, + element_strides=element_strides, + interleave=interleave, + swizzle=swizzle, + l2_promotion=l2_promotion, + oob_fill=oob_fill, + ) + + +def enum_of(table: dict[int, str], code: int, what: str) -> str: + if code not in table: + raise RuntimeError(f"the lowering asked for an unmapped {what}: {code}") + return table[code] + + +def runtime_dim(tmap: TensorMap, declared: int, label: str) -> int: + """Which of a descriptor's dimensions the launcher fills in per call. + + The rows of a packed q buffer and of the pool are run-time quantities, and + they enter the device code only through the descriptor, so the launcher + substitutes them here. Matching on the declared value keeps the position + tied to the lowering rather than to a hand-kept index. + """ + hits = [i for i, dim in enumerate(tmap.dims) if dim == declared] + if len(hits) != 1: + raise RuntimeError( + f"{label}: expected exactly one dimension equal to {declared} in " + f"{tmap.dims}, found {len(hits)}" + ) + return hits[0] + + +# The launcher's own name for each thing the stub names. A descriptor the +# lowering grew, or a tensor it renamed, has no binding here and fails +# generation rather than emitting a launcher that does not compile. +DESCRIPTOR_VAR = {"Q_desc": "q_desc", "KV_desc": "kv_desc"} +TENSOR_ARG = {"Q": "q", "KV": "kv"} + + +def render_descriptor(tmap: TensorMap, rows_expr: str, rows_at: int) -> str: + """The launcher body that encodes one descriptor.""" + if tmap.name not in DESCRIPTOR_VAR: + raise RuntimeError(f"no launcher variable for descriptor {tmap.name}") + if tmap.tensor not in TENSOR_ARG: + raise RuntimeError(f"no launcher argument for tensor {tmap.tensor}") + dims = [str(dim) for dim in tmap.dims] + dims[rows_at] = rows_expr + return ( + f" {{\n" + f" const cuuint64_t dims[{tmap.rank}] = {{{', '.join(dims)}}};\n" + f" const cuuint64_t strides[{tmap.rank - 1}] = " + f"{{{', '.join(str(s) for s in tmap.strides[1:])}}};\n" + f" const cuuint32_t box[{tmap.rank}] = " + f"{{{', '.join(str(b) for b in tmap.box)}}};\n" + f" const cuuint32_t element_strides[{tmap.rank}] = " + f"{{{', '.join(str(e) for e in tmap.element_strides)}}};\n" + f" const CUresult encoded = encode(\n" + f" &{DESCRIPTOR_VAR[tmap.name]}, " + f"{enum_of(TENSORMAP_DTYPE, tmap.dtype, 'dtype')}, {tmap.rank},\n" + f" const_cast({TENSOR_ARG[tmap.tensor]}), dims, strides, box, " + f"element_strides,\n" + f" {enum_of(TENSORMAP_INTERLEAVE, tmap.interleave, 'interleave')},\n" + f" {enum_of(TENSORMAP_SWIZZLE, tmap.swizzle, 'swizzle')},\n" + f" {enum_of(TENSORMAP_L2, tmap.l2_promotion, 'L2 promotion')},\n" + f" {enum_of(TENSORMAP_OOB, tmap.oob_fill, 'out-of-bounds fill')});\n" + f" if (encoded != CUDA_SUCCESS) {{\n" + f" return static_cast(cudaErrorInvalidValue);\n" + f" }}\n" + f" }}\n" + ) + + +ELEMENT_BYTES = {"CU_TENSOR_MAP_DATA_TYPE_BFLOAT16": 2} + +LAUNCHER_HEAD = """ +// Hand-written launcher. The key and query loads lower to bulk copies, so the +// entry takes TMA descriptors for those two and plain pointers for the rest; +// every descriptor parameter below is recovered from the lowered host stub. +#include +#include + +namespace { + +// cuTensorMapEncodeTiled is a driver entry point, resolved once and cached. +CUresult encode(CUtensorMap* map, CUtensorMapDataType dtype, cuuint32_t rank, + void* tensor, const cuuint64_t* dims, const cuuint64_t* strides, + const cuuint32_t* box, const cuuint32_t* element_strides, + CUtensorMapInterleave interleave, CUtensorMapSwizzle swizzle, + CUtensorMapL2promotion l2_promotion, + CUtensorMapFloatOOBfill oob_fill) { + using Fn = CUresult (*)(CUtensorMap*, CUtensorMapDataType, cuuint32_t, void*, + const cuuint64_t*, const cuuint64_t*, const cuuint32_t*, + const cuuint32_t*, CUtensorMapInterleave, + CUtensorMapSwizzle, CUtensorMapL2promotion, + CUtensorMapFloatOOBfill); + static Fn fn = [] { + void* entry = nullptr; + cudaDriverEntryPointQueryResult found; + if (cudaGetDriverEntryPoint("cuTensorMapEncodeTiled", &entry, + cudaEnableDefault, &found) != cudaSuccess || + found != cudaDriverEntryPointSuccess) { + return static_cast(nullptr); + } + return reinterpret_cast(entry); + }(); + if (fn == nullptr) { + return CUDA_ERROR_NOT_SUPPORTED; + } + return fn(map, dtype, rank, tensor, dims, strides, box, element_strides, + interleave, swizzle, l2_promotion, oob_fill); +} + +} // namespace +""" + + +def render_launcher( + maps: list[TensorMap], block: int, smem: int, order: list[str] +) -> str: + """The `extern "C"` entry `pegainfer-kernels` links against.""" + by_name = {tmap.name: tmap for tmap in maps} + if set(by_name) != set(DESCRIPTOR_VAR): + raise RuntimeError( + f"expected descriptors {sorted(DESCRIPTOR_VAR)}, got {sorted(by_name)}" + ) + + bodies = [ + render_descriptor( + by_name["Q_desc"], + "static_cast(q_rows)", + runtime_dim(by_name["Q_desc"], Q_ROWS, "Q_desc"), + ), + render_descriptor( + by_name["KV_desc"], + "static_cast(pool_rows)", + runtime_dim(by_name["KV_desc"], POOL_ROWS, "KV_desc"), + ), + ] + + opt_in = "" + if smem > MAX_STATIC_SMEM: + opt_in = ( + f" // Past 48 KiB the kernel has to opt in; once per symbol.\n" + f" static const cudaError_t opt_in = cudaFuncSetAttribute(\n" + f" reinterpret_cast({KERNEL_SYMBOL}),\n" + f" cudaFuncAttributeMaxDynamicSharedMemorySize, {smem});\n" + f" if (opt_in != cudaSuccess) {{\n" + f" return static_cast(opt_in);\n" + f" }}\n" + ) + + bound = { + **DESCRIPTOR_VAR, + "Output": "reinterpret_cast(out)", + "PageIndices": "page_indices", + "PageIndptr": "page_indptr", + "QIndptr": "q_indptr", + "LastPageLen": "last_page_len", + "sm_scale": "sm_scale", + "total_ctas": "total_ctas", + "rows_per_page": "rows_per_page", + "layer_row": "layer_row", + } + missing = [name for name in order if name not in bound] + if missing: + raise RuntimeError( + f"the entry point grew parameters this launcher does not bind: {missing}" + ) + args = ",\n ".join(bound[name] for name in order) + + declarations = "".join( + f" alignas(64) CUtensorMap {var};\n" for var in DESCRIPTOR_VAR.values() + ) + # Past any of these the body drops a tail or reads the wrong rows and hands + # back plausible numbers, so the launcher refuses rather than truncates. + bounds = ( + f" if (q_rows > {Q_ROWS} || pool_rows > {POOL_ROWS} || batch > {MAX_BATCH}\n" + f" || page_size != {PAGE_SIZE} || num_qo_heads != {HEADS}\n" + f" || num_kv_heads != {HEADS // GROUPS}) {{\n" + f" return static_cast(cudaErrorInvalidValue);\n" + f" }}\n" + ) + # The grid is computed here, from the sum the body re-walks to find its + # owner: a caller's own copy disagreeing is silent both ways, too large and + # CTAs spin up to exit, too small and a request's tail never runs. + grid = ( + f" int total_ctas = 0;\n" + f" for (int request = 0; request < batch; ++request) {{\n" + f" const int rows = host_q_indptr[request + 1] - host_q_indptr[request];\n" + f" const int tiles = (rows + {defs.BLOCK_M - 1}) / {defs.BLOCK_M};\n" + f" total_ctas += ((tiles + {defs.QBLK - 1}) / {defs.QBLK})" + f" * {defs.QBLK} * {HEADS};\n" + f" }}\n" + f" if (total_ctas == 0) {{\n" + f" // A step with no prompt rows is a real state, not an error.\n" + f" return static_cast(cudaSuccess);\n" + f" }}\n" + ) + signature = ", ".join(f"{kind} {name}" for kind, name in LAUNCHER_PARAMS) + return ( + f'extern "C" int {LAUNCHER}(\n' + f" {signature},\n" + f" cudaStream_t stream) {{\n" + f"{bounds}" + f"{grid}" + f"{declarations}" + f"{''.join(bodies)}" + f"{opt_in}" + f" {KERNEL_SYMBOL}<<>>(\n" + f" {args});\n" + f" return static_cast(cudaGetLastError());\n" + f"}}\n" + ) + + +def entry_parameter_order(source: str) -> list[str]: + """Parameter names of the generated entry point, in its own order.""" + marker = source.index(f"{KERNEL_MARKER} {ENTRY_SYMBOL}(") + open_at = source.index("(", marker) + close_at = source.index(")", open_at) + names = [] + for param in source[open_at + 1 : close_at].split(","): + names.append(param.strip().split()[-1].lstrip("*")) + return names + + +def split_source(source: str) -> tuple[str, str]: + """Split `get_kernel_source()` into (include preamble, kernel bodies).""" + marker = source.index(KERNEL_MARKER) + return source[:marker], source[marker:] + + +def isolate_debug_helpers(preamble: str) -> str: + """Rename `debug.h`'s externally linked helpers for this unit.""" + if DEBUG_HEADER not in preamble: + return preamble + renames = "".join(f"#define {name} {CU_STEM}_{name}\n" for name in DEBUG_HELPERS) + restores = "".join(f"#undef {name}\n" for name in DEBUG_HELPERS) + return preamble.replace( + DEBUG_HEADER, f"{renames}{DEBUG_HEADER}\n{restores}".rstrip("\n") + ) + + +def vendor_includes(out_dir: Path) -> tuple[Path, Path]: + """Copy the header roots in, so the directory stands on its own.""" + copied = [] + for source, name in ( + (TILELANG_TEMPLATE_PATH, "tilelang"), + (CUTLASS_INCLUDE_DIR, "cutlass"), + ): + destination = out_dir / "include" / name + if destination.exists(): + shutil.rmtree(destination) + shutil.copytree(source, destination) + copied.append(destination) + return copied[0], copied[1] + + +def required_nvcc_flags() -> list[str]: + """The flags TileLang's JIT would pass for this kernel's pass configs.""" + flags = [] + for key, enabled in defs.PASS_CONFIGS.items(): + if key not in PASS_CONFIG_NVCC_FLAG: + raise RuntimeError( + f"pass config {key} is not classified: say whether it reaches " + "nvcc's command line, or the generated object will not match " + "the kernel that was gated" + ) + flag = PASS_CONFIG_NVCC_FLAG[key] + if enabled and flag is not None: + flags.append(flag) + return flags + + +def build_kernel(arch: str): + """Lower for the arch the objects will be assembled for. + + Generation must not depend on a GPU being visible to the build host — + containers routinely have none, and TileLang then lowers for its own + default, which nvcc rejects outright. So the arch is always passed. + """ + return tilelang.compile( + defs.prefill_varlen( + HEADS, + GROUPS, + HEAD_DIM, + PAGE_SIZE, + MAX_BATCH, + Q_ROWS, + POOL_ROWS, + PAGE_TABLE_LEN, + ), + target={"kind": "cuda", "arch": arch}, + pass_configs=defs.PASS_CONFIGS, + ) + + +def check_strides(maps: list[TensorMap]) -> None: + """The builder's innermost stride is the element size; hold it to that.""" + for tmap in maps: + dtype = enum_of(TENSORMAP_DTYPE, tmap.dtype, "dtype") + if dtype not in ELEMENT_BYTES: + raise RuntimeError(f"{dtype} has no element size to check against") + expected = ELEMENT_BYTES[dtype] + if tmap.strides[0] != expected: + raise RuntimeError( + f"{tmap.name}: innermost stride {tmap.strides[0]} is not the " + f"{expected} B element of {dtype}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out-dir", type=Path, required=True) + parser.add_argument( + "--arch", + required=True, + help="arch to lower and assemble for, e.g. sm_90a. Required: without " + "it TileLang lowers for whatever device it can see, or for its own " + "default on a host with none.", + ) + parser.add_argument( + "--vendor-includes", + action="store_true", + help="copy the header roots into the output and point the manifest at " + "the copies (self-contained pre-generated dir)", + ) + args = parser.parse_args() + out_dir: Path = args.out_dir + out_dir.mkdir(parents=True, exist_ok=True) + + kernel = build_kernel(args.arch) + source = kernel.get_kernel_source() + maps, block, smem = read_host_stub(kernel) + check_strides(maps) + order = entry_parameter_order(source) + preamble, body = split_source(source) + if body.count(ENTRY_SYMBOL) != 2: + raise RuntimeError(f"expected exactly two {ENTRY_SYMBOL} occurrences") + + cu_path = out_dir / f"{CU_STEM}.cu" + cu_path.write_text( + "// Generated by pegainfer-gemma4/kernels/generate.py. Do not edit.\n" + + isolate_debug_helpers(preamble) + + body.replace(ENTRY_SYMBOL, KERNEL_SYMBOL) + + LAUNCHER_HEAD + + render_launcher(maps, block, smem, order) + ) + + if args.vendor_includes: + template_include, cutlass_include = vendor_includes(out_dir) + else: + template_include = Path(TILELANG_TEMPLATE_PATH) + cutlass_include = Path(CUTLASS_INCLUDE_DIR) + + # Relative where it can be, so a vendored directory survives being copied. + def named(path: Path) -> str: + try: + return str(path.relative_to(out_dir)) + except ValueError: + return str(path) + + lines = [f"CU_PATH={named(cu_path)}"] + lines.append(f"TILELANG_TEMPLATE_PATH={named(template_include)}") + lines.append(f"CUTLASS_INCLUDE_DIR={named(cutlass_include)}") + # Shapes are compile dimensions, so the consumer can refuse another. + lines.append(f"GEOMETRY={HEADS},{HEADS // GROUPS},{HEAD_DIM},{PAGE_SIZE}") + # The opt-in a block needs: SM90 grants 227 KiB, SM120 only 99. + lines.append(f"SMEM={smem}") + lines.extend(f"NVCC_FLAG={flag}" for flag in required_nvcc_flags()) + lines.append( + f"LAUNCHER={LAUNCHER}|{', '.join(kind for kind, _ in LAUNCHER_PARAMS)}" + ) + # The body is lowered for exactly this arch and uses arch-conditional + # instructions, so the consumer assembles it for that and not for the + # generic SM list. + lines.append(f"ARCH={args.arch}") + # The manifest lets a build host consume a pre-generated directory without + # re-running (or even having) TileLang; build.rs parses the same key=value + # lines from either stdout or this file. + (out_dir / "manifest.txt").write_text("\n".join(lines) + "\n") + for line in lines: + print(line) + print(f"# block {block} threads, {smem} B dynamic shared, {len(maps)} descriptors") + + +if __name__ == "__main__": + main() diff --git a/pegainfer-gemma4/kernels/tilelang_defs.py b/pegainfer-gemma4/kernels/tilelang_defs.py new file mode 100644 index 000000000..6c0430e85 --- /dev/null +++ b/pegainfer-gemma4/kernels/tilelang_defs.py @@ -0,0 +1,241 @@ +"""TileLang definition of Gemma 4's global-attention prefill at head dim 512. + +Authored here, not vendored: upstream has no kernel for this head dim on +SM90 — FlashInfer's Hopper prefill compiles 64, 128 and 256 only — so the +serving path runs a generic paged kernel that leaves most of the machine +idle at this shape. + +Three things make this faster, and none of them is the loop body: + + * the grid walks a KV group's query heads across a block of query tiles + before advancing, so the CTAs that want the same key tile stay resident + together and step through the context in lockstep; + * the score tile goes from shared memory straight into the value gemm, + because copying it back into a fragment first costs a fifth of the + runtime for a bit-identical result; and + * one key tile is exactly one page, so its load is a single unwrapped + copy — wrapping it in a loop, or splitting it into per-page pieces, + forfeits the bulk path and with it about half the throughput. + +The last one is why the global KV family pages at the key block rather than +at the sliding family's finer granularity. + +Ragged batches are the serving case: a mixed step launches its prompt rows +as one plan with an entry per segment. The kernel adds no array of its own: +it walks the batch once per CTA to find which request owns it, summing each +request's CTA count with the tile count rounded up to a whole mapping block +so the group walk stays intact. That is a handful of shifts over a bounded +batch, and it keeps the inputs to exactly what the plan already carries. A +request that is not in the step contributes no CTAs, so it is never chosen, +and a CTA past its request's real tiles exits before it touches memory. + +Each request's context length comes the same way, from how far its slice of +the page table reaches and how full its last page is, rather than as another +array restating what those already say. + +Rows are packed, so the store is predicated: a partial last tile would +otherwise write over the next request's rows, which in a mixed step are the +decode rows sharing the output buffer. Keeping the bulk copy for whole tiles +and branching to the predicated store only for a request's last one measures +slower than predicating every tile, so there is no branch. +""" + +import tilelang +import tilelang.language as T + +DTYPE = "bfloat16" +ACC = "float" + +# The fp32 output accumulator is block_M * head_dim * 4 B, so block_M is 64: +# 128 spills half the register file. block_N is the page size, and the two +# gemms want a full warpgroup pair. +BLOCK_M = 64 +BLOCK_N = 64 +NUM_STAGES = 1 +THREADS = 256 + +# Query tiles per KV group before the grid advances. +QBLK = 8 + +PASS_CONFIGS = {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True} + + +def prefill_varlen( + heads, + groups, + dim, + page_size, + max_batch, + q_rows, + pool_rows, + page_table_len, + block_M=BLOCK_M, + block_N=BLOCK_N, + num_stages=NUM_STAGES, + threads=THREADS, + qblk=QBLK, +): + """Causal GQA prefill over the paged pool, ragged across requests. + + `q_rows`, `pool_rows` and `page_table_len` are declared at the serving + arena's maxima. They reach the generated code only as bounds guards, so a + step smaller than the arena passes them all; the real bounds come from the + store predicate and from the walk's own trip count. The two tensors the + lowering reads through TMA carry their extents in the descriptors the + launcher builds, so those are the step's own. + """ + assert block_N == page_size, "one tile must be one page, or the load splits" + head_kv = heads // groups + q_shape = [q_rows, heads, dim] + kv_shape = [pool_rows, head_kv, dim] + + @T.prim_func + def main( + Q: T.Tensor(q_shape, DTYPE), + KV: T.Tensor(kv_shape, DTYPE), + PageIndices: T.Tensor([page_table_len], "int32"), + PageIndptr: T.Tensor([max_batch + 1], "int32"), + QIndptr: T.Tensor([max_batch + 1], "int32"), + LastPageLen: T.Tensor([max_batch], "int32"), + sm_scale: T.float32, + total_ctas: T.int32, + rows_per_page: T.int32, + layer_row: T.int32, + Output: T.Tensor(q_shape, DTYPE), + ): + with T.Kernel(total_ctas, threads=threads) as pid: + # The softmax runs on exp2, so the caller's scale carries log2(e) + # into the exponent. It is the caller's because the serving path + # folds 1/sqrt(head_dim) into the query rows upstream and hands + # this kernel a scale of one; baking the usual factor in here + # would apply it twice and quietly flatten every distribution. + scale = sm_scale * 1.44269504 + Q_shared = T.alloc_shared([block_M, dim], DTYPE) + K_shared = T.alloc_shared([block_N, dim], DTYPE) + S_shared = T.alloc_shared([block_M, block_N], DTYPE) + V_shared = T.alloc_shared([block_N, dim], DTYPE) + acc_s = T.alloc_fragment([block_M, block_N], ACC) + acc_o = T.alloc_fragment([block_M, dim], ACC) + scores_max = T.alloc_fragment([block_M], ACC) + scores_max_prev = T.alloc_fragment([block_M], ACC) + scores_scale = T.alloc_fragment([block_M], ACC) + scores_sum = T.alloc_fragment([block_M], ACC) + logsum = T.alloc_fragment([block_M], ACC) + req = T.alloc_local([1], "int32") + first_cta = T.alloc_local([1], "int32") + own_ctas = T.alloc_local([1], "int32") + walked = T.alloc_local([1], "int32") + + # Which request owns this CTA, and where its block starts. The + # host sizes the grid with the same sum, so the two agree by + # construction rather than through an array that could drift. + req[0] = 0 + first_cta[0] = 0 + own_ctas[0] = 0 + walked[0] = 0 + # `total_ctas` is the host's same sum, so it also says how many + # boundaries are real: the caller's array is as long as its own + # batch, while `max_batch` is this kernel's ceiling. + for i in T.serial(max_batch): + if walked[0] < total_ctas: + mine = ( + T.ceildiv(T.ceildiv(QIndptr[i + 1] - QIndptr[i], block_M), qblk) + * qblk + * heads + ) + # Requests sit back to back, so the owner is the last one + # that both starts at or before this CTA and has any. The + # emptiness test is what keeps a request the step left out + # from claiming CTAs it has no tiles for. + if walked[0] <= pid and mine > 0: + req[0] = i + first_cta[0] = walked[0] + own_ctas[0] = mine + walked[0] = walked[0] + mine + # An over-sized grid then does nothing rather than recompute + # somebody else's tile; an under-sized one leaves a tail, which + # the numerics gate sees. + if pid < walked[0]: + b = req[0] + local = pid - first_cta[0] + q_tiles = own_ctas[0] // heads + per_group = q_tiles * groups + r = local % per_group + head = local // per_group * groups + (r % (qblk * groups)) // qblk + q_tile = r // (qblk * groups) * qblk + r % qblk + kv_head = head // groups + q_start = QIndptr[b] + q_len = QIndptr[b + 1] - q_start + pages = PageIndptr[b + 1] - PageIndptr[b] + kv_len = (pages - 1) * page_size + LastPageLen[b] + offset = kv_len - q_len + row = q_start + q_tile * block_M + + if q_tile * block_M < q_len: + T.copy(Q[row : row + block_M, head, :], Q_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(ACC)) + + loop_range = T.min( + T.ceildiv(offset + (q_tile + 1) * block_M, block_N), + T.ceildiv(kv_len, block_N), + ) + for k in T.Pipelined(loop_range, num_stages=num_stages): + # One page holds every layer's K then V for `page_size` + # tokens, so the layer's K block starts `layer_row` into + # the page and its V block one page further. + kb = PageIndices[PageIndptr[b] + k] * rows_per_page + layer_row + T.copy(KV[kb : kb + page_size, kv_head, :], K_shared) + for i, j in T.Parallel(block_M, block_N): + acc_s[i, j] = T.if_then_else( + q_tile * block_M + i + offset < k * block_N + j, -1e9, 0 + ) + T.gemm( + Q_shared, + K_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(ACC)) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_M): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Parallel(block_M): + scores_scale[i] = T.exp2( + scores_max_prev[i] * scale - scores_max[i] * scale + ) + for i, j in T.Parallel(block_M, block_N): + acc_s[i, j] = T.exp2( + acc_s[i, j] * scale - scores_max[i] * scale + ) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Parallel(block_M): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + T.copy(acc_s, S_shared) + + for i, j in T.Parallel(block_M, dim): + acc_o[i, j] *= scores_scale[i] + T.copy( + KV[kb + page_size : kb + 2 * page_size, kv_head, :], + V_shared, + ) + T.gemm( + S_shared, V_shared, acc_o, policy=T.GemmWarpPolicy.FullCol + ) + + for i, j in T.Parallel(block_M, dim): + acc_o[i, j] = acc_o[i, j] / logsum[i] + + # Q_shared is dead once the walk ends, and the four live shared + # buffers already sit at the SM90 dynamic limit, so the store + # stages through it rather than its own. + T.copy(acc_o, Q_shared) + for i, d in T.Parallel(block_M, dim): + if q_tile * block_M + i < q_len: + Output[row + i, head, d] = Q_shared[i, d] + + return main diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 92191c532..9f422c973 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -32,8 +32,9 @@ use pegainfer_sample::LogprobRequest; use pegainfer_sample::SampleScratch; use crate::forward::MULTIMODAL_PLACEHOLDER_IDS; +use crate::kv::GLOBAL_PAGE_SIZE; use crate::kv::GemmaKv; -use crate::kv::PAGE_SIZE; +use crate::kv::LOCAL_PAGE_SIZE; use crate::kv::admit_tokens; use crate::prefix_cache::PrefixCache; use crate::serve::GemmaServe; @@ -54,6 +55,7 @@ const MAX_CONTEXT_ENV: &str = "PEGAINFER_MAX_CONTEXT"; const DECODE_SLOTS_ENV: &str = "PEGAINFER_DECODE_SLOTS"; const KV_FP8_ENV: &str = "PEGAINFER_KV_FP8"; const ADMIT_COALESCE_ENV: &str = "PEGAINFER_ADMIT_COALESCE_MS"; +const GLOBAL_ATTN_ENV: &str = "PEGAINFER_GLOBAL_ATTN"; const MIN_CONTEXT: usize = 1024; const MIN_CHUNK_TOKENS: usize = 64; const CEILING_DOMAIN: usize = i32::MAX as usize; @@ -164,6 +166,103 @@ fn parse_mix_chunk_tokens(raw: &str, max_context: usize) -> Result } } +/// Which kernel serves the global family's prefill. Unset is the kernel the +/// line has always used, byte for byte; `tilelang` is the generated one, which +/// exists only in a build that had TileLang or a pre-generated directory. +fn tilelang_global_attn() -> Result { + read_env(GLOBAL_ATTN_ENV)?.map_or(Ok(false), |raw| parse_tilelang_global_attn(&raw)) +} + +/// Refuse a checkpoint the generated bodies have no kernel for: the launcher +/// answers `cudaErrorInvalidValue` for another geometry, and it would answer +/// on the first global prefill. +pub(crate) fn tilelang_geometry_refusal(config: &crate::config::Gemma4Config) -> Result<()> { + if !pegainfer_kernels::ops::gemma4_hd512_prefill_is_built() { + return Ok(()); + } + let (heads, kv_heads, head_dim, page) = pegainfer_kernels::ops::gemma4_hd512_prefill_geometry() + .context( + "the build carries generated kernels but does not state the geometry they were \ + compiled for; regenerate the TileLang directory with the current generator", + )?; + let theirs = ( + config.num_attention_heads, + config.num_global_key_value_heads, + config.global_head_dim, + crate::kv::GLOBAL_PAGE_SIZE, + ); + anyhow::ensure!( + theirs == (heads, kv_heads, head_dim, page), + "{GLOBAL_ATTN_ENV} asks for kernels compiled for {heads} query heads over \ + {kv_heads} KV heads at head dim {head_dim} on {page}-row pages, but this \ + checkpoint's global family is {} over {} at {} on {}-row pages; serve it \ + through the incumbent kernel", + theirs.0, + theirs.1, + theirs.2, + theirs.3 + ); + Ok(()) +} + +/// Refuse a device the generated bodies cannot run on: generation targets one +/// arch, whose accelerated target runs on that capability alone, and a block +/// opts into more shared memory than some architectures of the same number +/// grant. +fn ensure_tilelang_device(device: usize) -> Result<()> { + if !pegainfer_kernels::ops::gemma4_hd512_prefill_is_built() { + return Ok(()); + } + let arch = pegainfer_kernels::ops::gemma4_hd512_prefill_arch().context( + "the build carries generated kernels but does not state the arch they were built \ + for; regenerate the TileLang directory with the current generator", + )?; + let built: u32 = arch + .trim_start_matches("sm_") + .trim_end_matches(|c: char| !c.is_ascii_digit()) + .parse() + .with_context(|| format!("the build reported an unreadable TileLang arch {arch:?}"))?; + let ctx = DeviceContext::new_with_device(device) + .with_context(|| format!("open device {device} for the TileLang arch check"))?; + let major = ctx.ctx.attribute( + cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, + )?; + let minor = ctx.ctx.attribute( + cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, + )?; + let running = u32::try_from(major * 10 + minor).context("compute capability fits u32")?; + anyhow::ensure!( + running == built, + "{GLOBAL_ATTN_ENV} asks for kernels built for {arch}, but device {device} is \ + SM{major}.{minor}: the generated bodies carry an image for one arch. Build with \ + PEGAINFER_CUDA_SM={running}, or serve this device through the incumbent kernel" + ); + let wanted = pegainfer_kernels::ops::gemma4_hd512_prefill_smem().context( + "the build carries generated kernels but does not state the shared memory they opt \ + into; regenerate the TileLang directory with the current generator", + )?; + let granted = ctx.ctx.attribute( + cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, + )?; + let granted = usize::try_from(granted).context("shared-memory limit fits usize")?; + anyhow::ensure!( + wanted <= granted, + "{GLOBAL_ATTN_ENV} asks for kernels whose block opts into {wanted} B of shared \ + memory, and device {device} grants {granted} B per block; serve it through the \ + incumbent kernel" + ); + Ok(()) +} + +fn parse_tilelang_global_attn(raw: &str) -> Result { + let value = raw.trim().to_ascii_lowercase(); + match value.as_str() { + "" | "0" | "off" => Ok(false), + "tilelang" => Ok(true), + _ => anyhow::bail!("{GLOBAL_ATTN_ENV}={raw:?} not recognized (off | tilelang)"), + } +} + pub(crate) fn prefix_cache_cap() -> Result> { read_env(PREFIX_CACHE_ENV)?.map_or(Ok(None), |raw| parse_prefix_cache_cap(&raw)) } @@ -490,11 +589,13 @@ fn validate_request(request: &Request, max_context: usize) -> Result usize { - context_len.div_ceil(PAGE_SIZE) + context_len.div_ceil(GLOBAL_PAGE_SIZE) } /// How many prompts one mixed step may absorb: bounded well below the @@ -1053,6 +1154,20 @@ impl EngineState { let admit_coalesce = admit_coalesce_ms()?; let slots = decode_slots()?; let local_kv_storage = kv_fp8_storage()?; + let tilelang_global = tilelang_global_attn()?; + // The stub tier links under the same name and refuses at launch, so + // without this the answer would arrive after the weights are loaded + // and on the first prompt rather than here. + anyhow::ensure!( + !tilelang_global || pegainfer_kernels::ops::gemma4_hd512_prefill_is_built(), + "{GLOBAL_ATTN_ENV}=tilelang needs a build that carries the kernel; \ + this one fell back to the stub tier, so pegainfer-kernels was \ + compiled without TileLang and without a pre-generated directory" + ); + if tilelang_global { + tilelang_geometry_refusal(&config)?; + ensure_tilelang_device(device)?; + } anyhow::ensure!( admit_coalesce.is_none() || lane_mode.is_none(), "{ADMIT_COALESCE_ENV} and {ASYNC_PREFILL_ENV} cannot combine: the lane flies one \ @@ -1082,8 +1197,13 @@ impl EngineState { // segment. The global family never releases, so it stays linear in // context for each request's whole lifetime. Both pools add the // padding page they reserve. - let context_pages = max_context.div_ceil(PAGE_SIZE); - let window_pages = weights.config.sliding_window.div_ceil(PAGE_SIZE) + 1; + // The families page at different granularities, so each budget below + // names the one it counts: the window and the local transient are + // local pages, the global account is global pages. One ceiling in + // local pages is not the same number in global pages. + let local_context_pages = max_context.div_ceil(LOCAL_PAGE_SIZE); + let global_context_pages = max_context.div_ceil(GLOBAL_PAGE_SIZE); + let window_pages = weights.config.sliding_window.div_ceil(LOCAL_PAGE_SIZE) + 1; // The cache brings its own page budget so cached entries never eat // serving headroom. let cache_cap = prefix_cache_cap()?; @@ -1097,14 +1217,14 @@ impl EngineState { // A round's rows split across walkers, and every walker's // reservation rounds up to its own page — so the budget // carries one page of rounding per extra walker. - window_pages + chunk.div_ceil(PAGE_SIZE) + (MIX_MAX_PROMPTS - 1) + window_pages + chunk.div_ceil(LOCAL_PAGE_SIZE) + (MIX_MAX_PROMPTS - 1) } - _ => context_pages, + _ => local_context_pages, }; let (local_pages, global_pages) = pool_pages( transient_pages, window_pages, - context_pages, + global_context_pages, slots, cache_entries, crate::prefix_cache::entry_global_pages(max_context), @@ -1139,6 +1259,7 @@ impl EngineState { local_kv_storage, local_pages, global_pages, + tilelang_global, ) .map_err(|err| { err.context(format!( @@ -2490,6 +2611,20 @@ mod knob_tests { assert!(parse_kv_fp8(Some("global")).is_err()); } + #[test] + fn global_attn_parses_or_refuses() { + assert!(!parse_tilelang_global_attn("off").expect("off parses")); + assert!(!parse_tilelang_global_attn("").expect("empty parses")); + assert!(!parse_tilelang_global_attn("0").expect("zero parses")); + assert!(parse_tilelang_global_attn(" TileLang ").expect("trimmed and cased")); + for bad in ["on", "1", "flashinfer", "tile", "tilelang:1"] { + assert!( + parse_tilelang_global_attn(bad).is_err(), + "{bad:?} must refuse" + ); + } + } + #[test] fn admit_coalesce_parses_or_refuses() { for off in ["off", "0", ""] { diff --git a/pegainfer-gemma4/src/engine/lane_gates_lifecycle.rs b/pegainfer-gemma4/src/engine/lane_gates_lifecycle.rs index 6029aaf87..4cfcf8f38 100644 --- a/pegainfer-gemma4/src/engine/lane_gates_lifecycle.rs +++ b/pegainfer-gemma4/src/engine/lane_gates_lifecycle.rs @@ -173,6 +173,33 @@ fn the_gathered_lifecycle_completes() { gather_lifecycle_script(); } +/// The knob is refused for a geometry the bodies were not compiled for, +/// before any weight is read rather than at the first global prefill. +#[test] +fn the_knob_is_refused_for_a_geometry_the_build_does_not_carry() { + if !pegainfer_kernels::ops::gemma4_hd512_prefill_is_built() { + eprintln!("skipping: this build carries the stub, which has no geometry to refuse"); + return; + } + let (heads, kv_heads, head_dim, _page) = + pegainfer_kernels::ops::gemma4_hd512_prefill_geometry() + .expect("a build that carries the bodies states the geometry they were compiled for"); + let mut config = crate::manifest::schema::sample_config(); + config.num_attention_heads = heads; + config.num_global_key_value_heads = kv_heads; + config.global_head_dim = head_dim; + super::tilelang_geometry_refusal(&config).expect("its own geometry is accepted"); + + config.num_global_key_value_heads = kv_heads + 1; + let refusal = super::tilelang_geometry_refusal(&config) + .expect_err("one KV head more is a geometry the bodies have no kernel for"); + let refusal = refusal.to_string(); + assert!( + refusal.contains(&(kv_heads + 1).to_string()) && refusal.contains(&kv_heads.to_string()), + "the refusal must name both geometries: {refusal}" + ); +} + #[test] fn pool_pages_follow_the_knobs() { assert_eq!( @@ -185,3 +212,33 @@ fn pool_pages_follow_the_knobs() { ); assert_eq!(super::pool_pages(usize::MAX, 65, 512, 16, 0, 256), None); } + +/// The door and the startup budget must count the global account in the same +/// family's page: a request at the serving ceiling asks for its whole account +/// at once and startup provisions one per slot, so a unit mismatch refuses a +/// request the pool was built to hold. +#[test] +fn the_global_door_fits_inside_what_startup_provisions() { + for max_context in [1024usize, 8192, 40960, 262_144] { + for slots in [1usize, 4, 16] { + let per_slot = super::global_account_pages(max_context); + // No cache entries: the tightest the global pool ever is. + let (_, global_pages) = super::pool_pages( + max_context.div_ceil(crate::kv::LOCAL_PAGE_SIZE), + 1, + max_context.div_ceil(crate::kv::GLOBAL_PAGE_SIZE), + slots, + 0, + 0, + ) + .expect("budget fits in usize"); + assert!( + slots * per_slot < global_pages, + "ceiling {max_context} over {slots} slots: the door asks for \ + {per_slot} pages each, {} in all, where startup provisioned \ + {global_pages}", + slots * per_slot + ); + } + } +} diff --git a/pegainfer-gemma4/src/engine/lane_gates_walk.rs b/pegainfer-gemma4/src/engine/lane_gates_walk.rs index 1208e8011..1b04f3cd9 100644 --- a/pegainfer-gemma4/src/engine/lane_gates_walk.rs +++ b/pegainfer-gemma4/src/engine/lane_gates_walk.rs @@ -27,9 +27,9 @@ fn the_gathered_transient_leaves_headroom() { let window = crate::config::Gemma4Config::from_file(&dir) .expect("config") .sliding_window; - let window_pages = window.div_ceil(crate::kv::PAGE_SIZE) + 1; + let window_pages = window.div_ceil(crate::kv::LOCAL_PAGE_SIZE) + 1; let provisioned = window_pages - + 2048usize.div_ceil(crate::kv::PAGE_SIZE) + + 2048usize.div_ceil(crate::kv::LOCAL_PAGE_SIZE) + (super::MIX_MAX_PROMPTS - 1) + (super::MAX_CONCURRENCY - 1) * window_pages; assert_eq!( diff --git a/pegainfer-gemma4/src/engine/lane_test_env.rs b/pegainfer-gemma4/src/engine/lane_test_env.rs index aea22839e..9eea1dc06 100644 --- a/pegainfer-gemma4/src/engine/lane_test_env.rs +++ b/pegainfer-gemma4/src/engine/lane_test_env.rs @@ -18,7 +18,7 @@ impl Drop for EnvGuard { } } -const SERVING_KNOBS: [&str; 7] = [ +const SERVING_KNOBS: [&str; 8] = [ super::ASYNC_PREFILL_ENV, super::PREFIX_CACHE_ENV, super::MIX_CHUNK_TOKENS_ENV, @@ -26,6 +26,7 @@ const SERVING_KNOBS: [&str; 7] = [ super::DECODE_SLOTS_ENV, super::ADMIT_COALESCE_ENV, super::KV_FP8_ENV, + super::GLOBAL_ATTN_ENV, ]; pub(super) fn scoped_engine_env(overrides: &[(&str, &str)]) -> EnvGuard { diff --git a/pegainfer-gemma4/src/kv.rs b/pegainfer-gemma4/src/kv.rs index ae7968baa..d60e23af6 100644 --- a/pegainfer-gemma4/src/kv.rs +++ b/pegainfer-gemma4/src/kv.rs @@ -270,7 +270,16 @@ pub(crate) fn admit_tokens( } } -pub(crate) const PAGE_SIZE: usize = 16; +/// The sliding family's page. Small pages keep the window's footprint tight, +/// and its reservations are per page because the front is released page by +/// page. +pub(crate) const LOCAL_PAGE_SIZE: usize = 16; + +/// The global family's page, sized so one key block is one tile load: at this +/// head dim a 64-row page keeps 0.93-0.96x of a contiguous tensor's throughput +/// where four 16-row pages keep 0.52-0.54x. The pool never releases a global +/// page, so the coarser granularity costs at most 63 tokens per request. +pub(crate) const GLOBAL_PAGE_SIZE: usize = 64; #[cfg(test)] mod tests { @@ -278,9 +287,12 @@ mod tests { use super::*; + /// The local family has to be able to grant what the global one refuses, + /// or the refusal lands before any reservation exists to roll back and + /// the atomicity test passes without exercising the rollback. fn tiny_pools(ctx: &DeviceContext) -> (KvPool, KvPool) { - let local = KvPool::new(ctx, 1, 1, 1, PAGE_SIZE, 4).expect("local pool"); - let global = KvPool::new(ctx, 1, 1, 1, PAGE_SIZE, 2).expect("global pool"); + let local = KvPool::new(ctx, 1, 1, 1, LOCAL_PAGE_SIZE, 8).expect("local pool"); + let global = KvPool::new(ctx, 1, 1, 1, GLOBAL_PAGE_SIZE, 2).expect("global pool"); (local, global) } @@ -297,14 +309,24 @@ mod tests { let ctx = DeviceContext::new().expect("GPU required"); let (local, global) = tiny_pools(&ctx); let mut kv = kv_from(&local, &global); - let refused = admit_tokens(&local, &global, &mut kv, 17); - assert!(refused.is_err(), "partial admission must refuse"); - assert_eq!(local.available_pages(), 3, "local occupancy must roll back"); - assert_eq!(global.available_pages(), 1, "global occupancy untouched"); + let before = (local.available_pages(), global.available_pages()); + let over_global = GLOBAL_PAGE_SIZE + 1; + let refused = admit_tokens(&local, &global, &mut kv, over_global) + .expect_err("partial admission must refuse"); + let refusal = refused.to_string(); + assert!( + refusal.contains("(granted, rolled back)"), + "the local family must be the one rolled back, got: {refusal}" + ); + assert_eq!( + (local.available_pages(), global.available_pages()), + before, + "a refused admission leaves both pools as it found them" + ); assert_eq!((kv.local.held_pages(), kv.global.held_pages()), (0, 0)); - admit_tokens(&local, &global, &mut kv, PAGE_SIZE).expect("one page each"); - assert_eq!((local.available_pages(), global.available_pages()), (2, 0)); + admit_tokens(&local, &global, &mut kv, LOCAL_PAGE_SIZE).expect("one page each"); + assert_eq!((local.available_pages(), global.available_pages()), (6, 0)); assert_eq!((kv.local.held_pages(), kv.global.held_pages()), (1, 1)); } } diff --git a/pegainfer-gemma4/src/prefix_cache.rs b/pegainfer-gemma4/src/prefix_cache.rs index 0066af7e5..eaad93822 100644 --- a/pegainfer-gemma4/src/prefix_cache.rs +++ b/pegainfer-gemma4/src/prefix_cache.rs @@ -14,7 +14,8 @@ use pegainfer_core::kv_pool::KvReservation; -use crate::kv::PAGE_SIZE; +use crate::kv::GLOBAL_PAGE_SIZE; +use crate::kv::LOCAL_PAGE_SIZE; /// A resume below this many tokens is not worth its page copies. const MIN_RESUME_TOKENS: usize = 64; @@ -36,7 +37,7 @@ fn resume_point(candidate: ResumeCandidate<'_>, window: usize, prompt: &[u32]) - let floor = if candidate.local_origin == 0 { MIN_RESUME_TOKENS } else { - (candidate.local_origin * PAGE_SIZE + window).max(MIN_RESUME_TOKENS) + (candidate.local_origin * LOCAL_PAGE_SIZE + window).max(MIN_RESUME_TOKENS) }; (resume >= floor).then_some(resume) } @@ -92,7 +93,7 @@ pub(crate) struct PrefixCache { /// context. Capture refuses a longer prompt, so the cache can never hold /// more than the share of the pool its entries paid for at startup. pub(crate) fn entry_global_pages(max_context: usize) -> usize { - max_context.div_ceil(PAGE_SIZE) / 2 + max_context.div_ceil(GLOBAL_PAGE_SIZE) / 2 } impl PrefixCache { @@ -196,8 +197,8 @@ mod tests { at_63[63] = 999; let mut at_64 = entry.clone(); at_64[64] = 999; - assert_eq!(resume(&entry, 0, PAGE_SIZE, &at_63), None); - assert_eq!(resume(&entry, 0, PAGE_SIZE, &at_64), Some(64)); + assert_eq!(resume(&entry, 0, LOCAL_PAGE_SIZE, &at_63), None); + assert_eq!(resume(&entry, 0, LOCAL_PAGE_SIZE, &at_64), Some(64)); } #[test] @@ -205,25 +206,33 @@ mod tests { let entry: Vec = (0..96).collect(); let mut extended = entry.clone(); extended.push(999); - assert_eq!(resume(&entry, 0, PAGE_SIZE, &extended), Some(96)); - assert_eq!(resume(&entry, 0, PAGE_SIZE, &entry), Some(95)); - assert_eq!(resume(&entry, 0, PAGE_SIZE, &[]), None); + assert_eq!(resume(&entry, 0, LOCAL_PAGE_SIZE, &extended), Some(96)); + assert_eq!(resume(&entry, 0, LOCAL_PAGE_SIZE, &entry), Some(95)); + assert_eq!(resume(&entry, 0, LOCAL_PAGE_SIZE, &[]), None); let mut before_window = entry.clone(); before_window[47] = 777; - assert_eq!(resume(&entry, 2, PAGE_SIZE, &before_window), None); + assert_eq!(resume(&entry, 2, LOCAL_PAGE_SIZE, &before_window), None); let mut at_window = entry.clone(); at_window[48] = 777; - assert_eq!(resume(&entry, 2, PAGE_SIZE, &at_window), None); + assert_eq!(resume(&entry, 2, LOCAL_PAGE_SIZE, &at_window), None); let mut after_minimum = entry.clone(); after_minimum[64] = 777; - assert_eq!(resume(&entry, 2, PAGE_SIZE, &after_minimum), Some(64)); + assert_eq!(resume(&entry, 2, LOCAL_PAGE_SIZE, &after_minimum), Some(64)); } #[test] fn global_page_budget_rounds_before_halving() { - assert_eq!(entry_global_pages(8192), 256); - assert_eq!(entry_global_pages(8193), 256); - assert_eq!(entry_global_pages(8224), 257); + let pages = 128; + let ceiling = pages * GLOBAL_PAGE_SIZE; + assert_eq!(entry_global_pages(ceiling), pages / 2); + // A token past a page boundary buys a whole page and the halving + // truncates it back away; halving the tokens first would round that + // same remainder up into an entry the pool never provisioned. + assert_eq!(entry_global_pages(ceiling + 1), pages / 2); + assert_eq!( + entry_global_pages(ceiling + 2 * GLOBAL_PAGE_SIZE), + pages / 2 + 1 + ); } } diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 9043c6615..e22da0d95 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -33,8 +33,9 @@ use crate::forward::embed_scale_bf16; use crate::forward::logits_tail; use crate::forward::logits_tail_into; use crate::forward::validate_tokens; +use crate::kv::GLOBAL_PAGE_SIZE; use crate::kv::GemmaKv; -use crate::kv::PAGE_SIZE; +use crate::kv::LOCAL_PAGE_SIZE; use crate::kv::SlidingLocalKv; use crate::kv::admit_tokens; use crate::layer::EpilogueScratch; @@ -753,6 +754,10 @@ pub(crate) struct GemmaServe { global_cos: DeviceVec, global_sin: DeviceVec, cos_max_pos: usize, + /// Which kernel the global family's prefill goes through. The two are + /// adapters at one seam — same arguments, same meaning — so the choice is + /// a flag here rather than a shape the call sites have to know about. + tilelang_global_attn: bool, /// Model layer index -> index within its family's pool layer axis. family_index: Vec, } @@ -802,6 +807,7 @@ impl GemmaServe { local_kv_storage: KvStorage, local_pages: usize, global_pages: usize, + tilelang_global_attn: bool, ) -> Result { // One source of truth for geometry, rope tables and layer numbering. let config = &weights.config; @@ -835,7 +841,7 @@ impl GemmaServe { locals, config.num_key_value_heads, config.head_dim, - PAGE_SIZE, + LOCAL_PAGE_SIZE, local_pages, local_kv_storage, )?; @@ -844,7 +850,7 @@ impl GemmaServe { globals, config.num_global_key_value_heads, config.global_head_dim, - PAGE_SIZE, + GLOBAL_PAGE_SIZE, global_pages, )?; let local_geom = LayerGeometry::local_of(config); @@ -882,9 +888,41 @@ impl GemmaServe { global_sin, cos_max_pos: max_context, family_index, + tilelang_global_attn, }) } + /// The global family's prefill, through whichever kernel this engine was + /// started with. Both are the same fn type, so a drift between them stops + /// compiling rather than computing something else. + fn global_prefill( + &self, + ctx: &DeviceContext, + q: &HiddenStates, + layer: usize, + plan: &PrefillPagedPlan, + out: &mut HiddenStates, + num_q_heads: usize, + ) -> Result<()> { + let attend = if self.tilelang_global_attn { + pegainfer_kernels::ops::gemma4_hd512_prefill_varlen_into + } else { + ops::batch_prefill_paged_hd512_into + }; + attend( + ctx, + q, + self.global_pool.buffer(), + &self.global_pool.layout().kernel_layout(), + layer, + plan, + out, + num_q_heads, + // The prep folds 1/sqrt(head_dim) into the query rows upstream. + 1.0, + ) + } + /// One arena per engine thread, sized for the decode step; a prompt /// builds a [`TowerScratch`] for its own width instead. A request's /// tiles are `ceil(rows * group / cta_tile_q)` with a positive @@ -1520,16 +1558,13 @@ impl GemmaServe { geom.head_dim, geom.rms_norm_eps, )?; - ops::batch_prefill_paged_hd512_into( + self.global_prefill( ctx, &scratch.q_prep, - self.global_pool.buffer(), - &self.global_pool.layout().kernel_layout(), family_layer, global_plan, &mut scratch.attn, geom.num_q_heads, - 1.0, )?; } PrepRef::Batched { global_tables, .. } => { @@ -1633,16 +1668,13 @@ impl GemmaServe { // chunk plan as a pure decode step. scratch.q_prep.seq_len = prefill_len; scratch.attn.seq_len = prefill_len; - ops::batch_prefill_paged_hd512_into( + self.global_prefill( ctx, &scratch.q_prep, - self.global_pool.buffer(), - &self.global_pool.layout().kernel_layout(), family_layer, global_prefill_plan, &mut scratch.attn, geom.num_q_heads, - 1.0, )?; let batch = seq_len - prefill_len; let factor = self.global_split_factor; diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index ca79c8836..780327d83 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -33,7 +33,10 @@ fn stack_with_storage( let weights = Gemma4Weights::from_safetensors(&dir, 0, config).expect("load checkpoint weights"); let ctx = DeviceContext::new_with_device(0).expect("device context"); - let serve = GemmaServe::new(&ctx, weights, max_context, storage, pages, pages).expect("serve"); + // The oracle measures the kernel the line has always used; the opt-in one + // has its own gate. + let serve = + GemmaServe::new(&ctx, weights, max_context, storage, pages, pages, false).expect("serve"); eprintln!("oracle stack storage: {storage:?}"); (ctx, serve, dir) } @@ -488,10 +491,12 @@ fn fp8_argmax_agreement_meets_the_bf16_floor() { .map(|(name, cut, stride)| agreement_case(&fixture, name, cut, stride)); let max_context = cases .iter() - .map(|case| case.prompt.len().div_ceil(crate::kv::PAGE_SIZE) * crate::kv::PAGE_SIZE) + .map(|case| { + case.prompt.len().div_ceil(crate::kv::LOCAL_PAGE_SIZE) * crate::kv::LOCAL_PAGE_SIZE + }) .max() .expect("agreement cases"); - let pages = max_context.div_ceil(crate::kv::PAGE_SIZE) + 2; + let pages = max_context.div_ceil(crate::kv::LOCAL_PAGE_SIZE) + 2; let (ctx, bf16, _) = stack_with_storage(max_context, pages, KvStorage::Bf16); let bf16_results = cases @@ -522,6 +527,21 @@ fn fp8_argmax_agreement_meets_the_bf16_floor() { } } +/// The shape of one logit row, for when two of them disagree: the top few ids +/// and the range say whether a row is a distribution or garbage. +fn describe_row(what: &str, row: &[f32]) { + let mut ranked: Vec<(usize, f32)> = row.iter().copied().enumerate().collect(); + ranked.sort_by(|a, b| b.1.total_cmp(&a.1)); + let lo = row.iter().copied().fold(f32::INFINITY, f32::min); + let hi = row.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let finite = row.iter().filter(|v| v.is_finite()).count(); + eprintln!( + "{what}: range [{lo}, {hi}], finite {finite}/{}, top5 {:?}", + row.len(), + &ranked[..5.min(ranked.len())] + ); +} + /// The two arms run different launch shapes (one-token decode against a /// whole-prompt prefill), which this engine does not promise bit-equal, so /// the callers bound the raw-logit drift by the calibrated ceiling — but the @@ -551,6 +571,71 @@ fn serving_recompute(ctx: &DeviceContext, serve: &GemmaServe, tokens: &[u32]) -> host[(logits.seq_len - 1) * vocab..].to_vec() } +/// The replacement global-attention kernel against the one it stands in for, +/// both through the production serving path. The incumbent is asked twice +/// first, so the tolerance is a measured floor rather than a chosen number. +/// The prompt leaves the last page partial, where tail handling could differ. +#[test] +#[ignore = "requires a Gemma 4 checkpoint, a GPU, and --test-threads=1"] +fn the_replacement_global_kernel_matches_the_incumbent() { + let (ctx, mut serve, dir) = stack_with(4096, 300); + assert!( + pegainfer_kernels::ops::gemma4_hd512_prefill_is_built(), + "this build has no TileLang kernel to compare; the gate needs one that does" + ); + // The runner selects this gate only where the geometry matches. + let config = crate::config::Gemma4Config::from_file(&dir).expect("config"); + crate::engine::tilelang_geometry_refusal(&config).expect( + "this gate compares the generated kernel against the incumbent, so it needs a \ + checkpoint whose global geometry the build was compiled for", + ); + let prompts = crate::testkit::generate_fixture_prompts(); + let tokens: Vec = prompts[0].iter().cycle().copied().take(1500).collect(); + let page = serve.global_pool.layout().page_size; + assert!( + !tokens.len().is_multiple_of(page), + "the prompt has to leave the final global page partial, and {} tokens \ + divides the pool's {page}-row page", + tokens.len() + ); + + assert!(!serve.tilelang_global_attn); + let incumbent = serving_recompute(&ctx, &serve, &tokens); + let again = serving_recompute(&ctx, &serve, &tokens); + let floor = compare_row(&incumbent, &again, "incumbent against itself"); + assert!( + incumbent + .iter() + .zip(&again) + .all(|(a, b)| a.to_bits() == b.to_bits()), + "the incumbent is not bit-identical run to run, so there is no floor \ + to measure the replacement against" + ); + + serve.tilelang_global_attn = true; + let replacement = serving_recompute(&ctx, &serve, &tokens); + // Report before asserting: when the two disagree, the magnitude and the + // shape of each row say which kind of wrong it is, and `compare_row` + // stops at the first divergence it finds. + describe_row("incumbent", &incumbent); + describe_row("replacement", &replacement); + let spread = incumbent + .iter() + .zip(&replacement) + .map(|(x, y)| (x - y).abs()) + .fold(0.0f32, f32::max); + eprintln!("max |dlogit| between the two kernels: {spread}"); + let gap = compare_row(&incumbent, &replacement, "replacement against incumbent"); + eprintln!( + "global prefill over {} tokens: floor {floor}, replacement |dlogit| {gap}", + tokens.len() + ); + assert!( + gap <= 2.0, + "replacement |dlogit| {gap} above the line's calibrated 2.0" + ); +} + /// One forward path answers for itself: every prompt position's incremental /// logits (one token at a time through the decode arena) match a whole-prompt /// recompute of the same serving path, and four decode steps fed the @@ -1475,7 +1560,7 @@ fn a_ragged_batch_does_not_depend_on_row_order() { // pool's padding page. let pages = lengths .iter() - .map(|len| (len + STEPS).div_ceil(PAGE_SIZE)) + .map(|len| (len + STEPS).div_ceil(LOCAL_PAGE_SIZE)) .sum::() + 1; let (ctx, serve, _dir) = stack_with(2048, pages); diff --git a/pegainfer-k3/kernels/generate.py b/pegainfer-k3/kernels/generate.py index b2fa7b299..1150eb8c8 100644 --- a/pegainfer-k3/kernels/generate.py +++ b/pegainfer-k3/kernels/generate.py @@ -1147,9 +1147,16 @@ def main() -> None: template_include = Path(TILELANG_TEMPLATE_PATH) cutlass_include = Path(CUTLASS_INCLUDE_DIR) - lines = [f"CU_PATH={path}" for path in cu_paths] - lines.append(f"TILELANG_TEMPLATE_PATH={template_include}") - lines.append(f"CUTLASS_INCLUDE_DIR={cutlass_include}") + # Relative where it can be, so a vendored directory survives being copied. + def named(path: Path) -> str: + try: + return str(Path(path).relative_to(out_dir)) + except ValueError: + return str(path) + + lines = [f"CU_PATH={named(path)}" for path in cu_paths] + lines.append(f"TILELANG_TEMPLATE_PATH={named(template_include)}") + lines.append(f"CUTLASS_INCLUDE_DIR={named(cutlass_include)}") if args.arch: # The bodies are lowered for exactly this arch and may use # arch-conditional instructions, so the consumer has to assemble them diff --git a/pegainfer-kernels/Cargo.toml b/pegainfer-kernels/Cargo.toml index b68671d3b..4772318ac 100644 --- a/pegainfer-kernels/Cargo.toml +++ b/pegainfer-kernels/Cargo.toml @@ -26,7 +26,8 @@ tvm-ffi-triton-cubin = ["dep:tvm-ffi", "qwen35"] # Qwen3.5 Triton AOT kernels (GDR chunkwise prefill) — the only feature that # needs Python + Triton at build time. deepseek-v2-lite = [] -# Gemma 4: NVFP4 dequantization, whose conversion intrinsics need CUDA >= 12.8. +# Gemma 4: NVFP4 dequantization, whose conversion intrinsics need CUDA >= 12.8, +# and the TileLang-generated hd512 global-attention prefill. gemma4 = [] qwen35 = [] # Shared MoE/MLA third-party substrate: DeepEP, DeepGEMM, and FlashMLA. @@ -41,5 +42,9 @@ moe = [] name = "triton_cubin_tvm_ffi" required-features = ["tvm-ffi-triton-cubin"] +[[example]] +name = "hd512_prefill_bench" +required-features = ["gemma4"] + [lints] workspace = true diff --git a/pegainfer-kernels/build.rs b/pegainfer-kernels/build.rs index 78ebc5420..56bcd47fa 100644 --- a/pegainfer-kernels/build.rs +++ b/pegainfer-kernels/build.rs @@ -1490,35 +1490,104 @@ fn compile_triton_aot_kernels(cuda_include: &Path, out_dir: &Path, sm_targets: & } // =========================================================================== -// k3 tilelang: BEGIN — K3 TileLang decode kernels (AOT). +// tilelang: BEGIN — TileLang AOT kernel families. // -// Self-contained section: everything the `k3` feature needs to turn -// `pegainfer-k3/kernels/generate.py` into objects lives between these two -// markers plus one `cfg!(feature = "k3")` block inside `main`. +// Self-contained section: everything a model line needs to turn its +// `generate.py` into objects lives between these two markers plus one +// `cfg!(feature = ...)` block per family inside `main`. A family is a row in +// the table below; the tiers, the arch handling and the nvcc flags are shared. // // Three tiers, first one that works wins: // 1. generate — run the generator with a build-host Python that has the // pinned TileLang. Preferred; this is what a dev box and the GPU CI // builders take. -// 2. pre-generated — `PEGAINFER_K3_TILELANG_PREGEN=` points at a +// 2. pre-generated — `PEGAINFER_