diff --git a/H20-TRITON-PLAN.md b/H20-TRITON-PLAN.md index 89a2624..420bbc0 100644 --- a/H20-TRITON-PLAN.md +++ b/H20-TRITON-PLAN.md @@ -107,13 +107,21 @@ | 阶段 | 状态 | 说明 | |---|---|---| | P0 基础设施 | ✅ | 后端选择 + fallback + A/B 骨架 | -| P1 轻量 kernel | 🔄 8/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 | +| 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 | ⬜ | | | P3 MTP 链路 | ⬜ | | | P4 采样 | ⬜ | | | P5 端到端 | ⬜ | | -A/B 现状：**45 个用例全过**（`python scripts/h20_ops_ab.py`）。 +A/B 现状：**13 个 kernel、53 个用例全过**（`python scripts/h20_ops_ab.py`）。 + +`norm` 是这批里最关键的一个，两点值得记： +- **量化前必须先舍回 bf16**。TileLang 版是 `cast(cast(y, bf16), fp32)` 之后才算 amax 和 + 乘 scale —— 这样才能和「把 bf16 结果单独量化」逐 bit 一致（引擎的参考就是这么定义的）。 + 少了这一步，fp8 会在舍入边界上随机偏 1-2 个 ulp。 +- **生产路径（`round_sf=True`）逐 byte 精确**：d=5120/7168、sf 也完全相等。 + `round_sf=False` 时 240/327680 个字节差 1 ulp，因为 sf=amax/448 依赖 amax 的 + 归约顺序，两个编译器不可能一致。 ### 过程中踩到的 Triton 语义坑（都已修，可复用） diff --git a/deepinfer/model_executor/layers/tilelang_ops/ops.py b/deepinfer/model_executor/layers/tilelang_ops/ops.py index 5835467..a9fb853 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/ops.py +++ b/deepinfer/model_executor/layers/tilelang_ops/ops.py @@ -31,15 +31,22 @@ from deepinfer import PLATFORM _TRACE = bool(os.getenv("DEEPINFER_TRACE_KERNELS")) -def _record(name: str) -> None: +def _record(name: str, args: tuple = (), kwargs: dict | None = None) -> None: """Log one use of a backend kernel factory, so the real surface can be measured. Called once per kernel *construction* rather than per launch: TileLang caches compiled kernels, so construction happens at warmup and the count stays meaningful. + + With `DEEPINFER_TRACE_KERNEL_ARGS=1` the construction arguments are logged too. Several + factories take a dozen flags (norm, fused_transform, fused_sample) and only a handful of + combinations are ever used, so the arguments say which code paths a port actually needs. """ path = Path(os.environ["DEEPINFER_TRACE_KERNELS"]) with path.open("a") as f: - f.write(f"{name}\n") + if os.getenv("DEEPINFER_TRACE_KERNEL_ARGS"): + f.write(f"{name}|args={args!r}|kwargs={sorted((kwargs or {}).items())!r}\n") + else: + f.write(f"{name}\n") class _KernelTrace: @@ -58,7 +65,7 @@ class _KernelTrace: def __call__(self, *args: Any, **kwargs: Any) -> Any: if not self._seen: self._seen = True - _record(self._name) + _record(self._name, args, kwargs) return self._factory(*args, **kwargs) def __getattr__(self, item: str) -> Any: diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/norm.py b/deepinfer/model_executor/layers/tilelang_ops/triton/norm.py new file mode 100644 index 0000000..3709ed5 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/norm.py @@ -0,0 +1,230 @@ +"""RMSNorm with optional fused residual add and FP8 cast, in Triton. + +The arithmetic mirrors the TileLang kernel exactly, including the details that are easy to get +subtly wrong: + +* the RMS statistic is over the residual-summed row, in float32; +* `weight * out_scale` is applied inside the normalisation multiply rather than afterwards; +* the per-block amax is clamped to `FP8_AMAX_MARGIN` before the scale is derived, so an all-zero + block still yields a finite scale; +* `round_sf` rounds the scale to a power of two using the float exponent directly, and the stored + scale is the reciprocal of the multiplier, so the two always multiply back to one. + +Two passes over the row: the first accumulates the sum of squares (the statistic needs the whole +row before any element can be scaled), the second emits the outputs. `BLOCK_D` is a multiple of the +scale granularity, so every scale block lies entirely inside one chunk and the quantisation does +not need a third pass. + +Layout support covers what the engine asks for on this platform. The packed-ue8m0 and mega-moe +shared-scale forms are rejected loudly rather than approximated silently. +""" + +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 + + +@triton.jit +def _norm_kernel( + X, + Y, + Y_ORIGIN, + W, + RESIDUAL, + RESIDUAL_OUT, + SF, + x_s, + y_s, + yo_s, + r_s, + ro_s, + sf_m, + sf_j, + seqlen, + d: tl.constexpr, + num_per_channels: tl.constexpr, + eps: tl.constexpr, + out_scale: tl.constexpr, + fp8_max: tl.constexpr, + amax_margin: tl.constexpr, + has_residual: tl.constexpr, + do_fp8_cast: tl.constexpr, + round_sf: tl.constexpr, + return_bf16: tl.constexpr, + N_BLOCKS: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_M: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + row_ok = rows < seqlen + + x_row = X + rows * x_s + r_row = RESIDUAL + rows * r_s + + # Pass 1: sum of squares over the full row. + sumsq = tl.zeros((BLOCK_M,), tl.float32) + for start in range(0, d, BLOCK_D): + cols = start + tl.arange(0, BLOCK_D) + mask = row_ok[:, None] & (cols[None, :] < d) + value = tl.load(x_row[:, None] + cols[None, :], mask=mask, other=0.0).to(tl.float32) + if has_residual: + value += tl.load(r_row[:, None] + cols[None, :], mask=mask, other=0.0).to(tl.float32) + sumsq += tl.sum(tl.where(mask, value * value, 0.0), axis=1) + rstd = tl.rsqrt(sumsq / d + eps) + + # Pass 2: shared residual output, then the normalised result in whichever forms are requested. + ro_row = RESIDUAL_OUT + rows * ro_s + y_row = Y + rows * y_s + yo_row = Y_ORIGIN + rows * yo_s + + for start in range(0, d, BLOCK_D): + cols = start + tl.arange(0, BLOCK_D) + mask = row_ok[:, None] & (cols[None, :] < d) + value = tl.load(x_row[:, None] + cols[None, :], mask=mask, other=0.0).to(tl.float32) + if has_residual: + value += tl.load(r_row[:, None] + cols[None, :], mask=mask, other=0.0).to(tl.float32) + tl.store(ro_row[:, None] + cols[None, :], value.to(tl.bfloat16), mask=mask) + + weight = tl.load(W + cols, mask=cols < d, other=0.0).to(tl.float32) + y = value * rstd[:, None] * (weight * out_scale) + # The quantiser sees the bf16-rounded value, not the fp32 one: that is what makes the fp8 + # output bit-identical to casting the bf16 result separately, which is how the engine's + # own reference is defined and what the TileLang kernel does. + y_bf = y.to(tl.bfloat16).to(tl.float32) + + if return_bf16: + tl.store(yo_row[:, None] + cols[None, :], y_bf.to(tl.bfloat16), mask=mask) + + if do_fp8_cast: + blocks = tl.reshape(y_bf, (BLOCK_M, N_BLOCKS, num_per_channels)) + amax = tl.max(tl.abs(blocks), axis=2) + amax = tl.maximum(amax, amax_margin) + if round_sf: + # exp = -ceil(log2(amax / 448)); multiplier = 2^exp, stored scale = 2^-exp. + 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) + dequant = (127 - exp) << 23 + dequant = dequant.to(tl.float32, bitcast=True) + else: + multiplier = fp8_max / amax + dequant = amax / fp8_max + + expanded = tl.reshape( + tl.broadcast_to(multiplier[:, :, None], (BLOCK_M, N_BLOCKS, num_per_channels)), + (BLOCK_M, BLOCK_D), + ) + tl.store(y_row[:, None] + cols[None, :], (y_bf * expanded).to(tl.float8e4nv), mask=mask) + + sf_cols = start // num_per_channels + tl.arange(0, N_BLOCKS) + sf_mask = row_ok[:, None] & (sf_cols[None, :] < d // num_per_channels) + tl.store(SF + rows[:, None] * sf_m + sf_cols[None, :] * sf_j, dequant, mask=sf_mask) + elif not return_bf16: + tl.store(y_row[:, None] + cols[None, :], y_bf.to(tl.bfloat16), mask=mask) + + +def norm_tl( + d: int, + dtype: str, + return_dtype: torch.dtype, + reduce_dtype: str, + has_residual: bool, + eps: float, + out_scale: float, + do_fp8_cast: bool, + num_per_channels: int, + use_tma_aligned_col_major_sf: bool, + round_sf: bool, + use_packed_ue8m0: bool, + return_origin_out: bool, + xdim: int, + shared_sf_block_m: int, + x_stride_s: int, + y_stride_s: int, + residual_stride_s: int, + residual_out_stride_s: int, + y_origin_stride_s: int, + sf_stride_s: int, +) -> Any: + """Build the kernel for a fixed hidden size and cast configuration.""" + del dtype, reduce_dtype, return_dtype, xdim, x_stride_s, y_stride_s + del residual_stride_s, residual_out_stride_s, y_origin_stride_s, sf_stride_s + del use_tma_aligned_col_major_sf # only affects the scale buffer's strides, read at launch + if use_packed_ue8m0 or shared_sf_block_m: + raise NotImplementedError( + "triton norm implements the plain fp32 scale layouts only; got " + f"use_packed_ue8m0={use_packed_ue8m0}, shared_sf_block_m={shared_sf_block_m}" + ) + if return_origin_out and not do_fp8_cast: + raise NotImplementedError("return_origin_out only applies when casting") + + npc = num_per_channels or 1 + assert d % npc == 0, f"hidden size {d} is not a multiple of the scale granularity {npc}" + + # Largest power-of-two multiple of npc that stays within a 2048-element chunk: keeps the + # register tile bounded while guaranteeing whole scale blocks per chunk. + block_d = npc + while block_d * 2 <= d and block_d * 2 <= 2048: + block_d *= 2 + block_m = 2 + + def launch( + x: torch.Tensor, + y: torch.Tensor, + w: torch.Tensor, + residual: torch.Tensor, + residual_out: torch.Tensor, + sf: torch.Tensor, + shared_sf: torch.Tensor, + y_origin: torch.Tensor, + ) -> None: + del shared_sf + assert x.ndim == 3, "the interface layer reshapes to [B, M, d]" + batch, seqlen = x.shape[0], x.shape[1] + assert batch == 1, "the triton norm kernel handles one batch slice at a time" + if seqlen == 0: + return + + # The bf16 result has its own buffer when casting; otherwise `y` already holds bf16, and + # `y_origin` aliases `x` as a placeholder and must not be written. + bf16_target = y_origin if (do_fp8_cast and return_origin_out) else y + write_bf16 = bool(do_fp8_cast and return_origin_out) or not do_fp8_cast + + _norm_kernel[(triton.cdiv(seqlen, block_m),)]( + x, + y, + bf16_target, + w, + residual, + residual_out, + sf, + x.stride(1), + y.stride(1), + bf16_target.stride(1), + residual.stride(1), + residual_out.stride(1), + sf.stride(1) if sf.ndim == 3 else 0, + sf.stride(2) if sf.ndim == 3 else 0, + seqlen, + d=d, + num_per_channels=npc, + eps=eps, + out_scale=out_scale, + fp8_max=FLOAT8E4M3_MAX, + amax_margin=FP8_AMAX_MARGIN, + has_residual=has_residual, + do_fp8_cast=do_fp8_cast, + round_sf=round_sf, + return_bf16=write_bf16, + N_BLOCKS=block_d // npc, + BLOCK_D=block_d, + BLOCK_M=block_m, + num_warps=8, + ) + + return launch diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py b/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py new file mode 100644 index 0000000..25680fe --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py @@ -0,0 +1,314 @@ +"""Sampler-side kernels: metadata preparation, parameter expansion, key/value lookup. + +These are grouped in one module because the TileLang backend also keeps them together, and the +interface module for each one imports the factory from a single place. +""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from deepinfer.model_executor.layers.tilelang_ops.common import DRAFT_POS_HASH_MUL +from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2, triton_dtype + + +@triton.jit +def _prepare_verification_metadata_kernel( + FTP, + SEQLEN_Q, + ROW_TO_REQ, + LOCAL_STEP, + IS_VERIFY, + threads: tl.constexpr, +): + req = tl.program_id(0) + tile = tl.program_id(1) + + start = tl.load(FTP + req) + length = tl.load(SEQLEN_Q + req) + + steps = tile * threads + tl.arange(0, threads) + live = steps < length + row = start + steps + + tl.store(ROW_TO_REQ + row, req, mask=live) + tl.store(LOCAL_STEP + row, steps, mask=live) + # Every row but the last is a draft row awaiting verification. + tl.store(IS_VERIFY + row, steps + 1 < length, mask=live) + + +def prepare_verification_metadata_tl(num_steps: int) -> Any: + """Build the kernel; the per-request length stays a runtime value.""" + del num_steps # the TileLang factory ignores it too + threads = 32 + + def launch( + ftp: torch.Tensor, + seqlen_q: torch.Tensor, + row_to_req: torch.Tensor, + local_step: torch.Tensor, + is_verify: torch.Tensor, + ) -> None: + req_cnt = ftp.shape[0] + if req_cnt == 0: + return + # One launch per (request, 32-row tile); the tile count is bounded by the longest request. + max_len = int(seqlen_q.max().item()) if seqlen_q.numel() else 0 + if max_len == 0: + return + _prepare_verification_metadata_kernel[(req_cnt, cdiv(max_len, threads))]( + ftp, + seqlen_q, + row_to_req, + local_step, + is_verify, + threads=threads, + num_warps=1, + ) + + return launch + + +@triton.jit +def _expand_sampling_params_kernel( + TEMPERATURES, + TOP_PS, + TOP_KS, + SAMPLE_POSITIONS, + SEQ_SEEDS, + ROW_TO_REQ, + LOCAL_STEP, + TEMPERATURES_Q, + TOP_PS_Q, + TOP_KS_Q, + SAMPLE_POSITIONS_Q, + SEQ_SEEDS_Q, + total_q, + pos_hash_mul: tl.constexpr, + THREADS: tl.constexpr, +): + block = tl.program_id(0) + rows = block * THREADS + tl.arange(0, THREADS) + live = rows < total_q + + req = tl.load(ROW_TO_REQ + rows, mask=live, other=0) + step = tl.load(LOCAL_STEP + rows, mask=live, other=0) + + tl.store(TEMPERATURES_Q + rows, tl.load(TEMPERATURES + req, mask=live, other=0.0), mask=live) + tl.store(TOP_PS_Q + rows, tl.load(TOP_PS + req, mask=live, other=0.0), mask=live) + tl.store(TOP_KS_Q + rows, tl.load(TOP_KS + req, mask=live, other=0), mask=live) + # Position is hashed with the draft step so each draft row draws a different random stream. + tl.store( + SAMPLE_POSITIONS_Q + rows, + tl.load(SAMPLE_POSITIONS + req, mask=live, other=0) + step * pos_hash_mul, + mask=live, + ) + tl.store(SEQ_SEEDS_Q + rows, tl.load(SEQ_SEEDS + req, mask=live, other=0), mask=live) + + +def expand_sampling_params_tl() -> Any: + """Build the kernel; all lengths are runtime.""" + threads = 64 + + def launch( + temperatures: torch.Tensor, + top_ps: torch.Tensor, + top_ks: torch.Tensor, + sample_positions: torch.Tensor, + seq_seeds: torch.Tensor, + row_to_req: torch.Tensor, + local_step: torch.Tensor, + temperatures_q: torch.Tensor, + top_ps_q: torch.Tensor, + top_ks_q: torch.Tensor, + sample_positions_q: torch.Tensor, + seq_seeds_q: torch.Tensor, + ) -> None: + total_q = row_to_req.shape[0] + if total_q == 0: + return + _expand_sampling_params_kernel[(cdiv(total_q, threads),)]( + temperatures, + top_ps, + top_ks, + sample_positions, + seq_seeds, + row_to_req, + local_step, + temperatures_q, + top_ps_q, + top_ks_q, + sample_positions_q, + seq_seeds_q, + total_q, + pos_hash_mul=DRAFT_POS_HASH_MUL, + THREADS=threads, + num_warps=2, + ) + + return launch + + +@triton.jit +def _map_lookup_index_kernel( + VALUE, + QUERY, + OUT, + total_q, + width: tl.constexpr, + is_float: tl.constexpr, + BLOCK: tl.constexpr, +): + block = tl.program_id(0) + rows = block * BLOCK + tl.arange(0, BLOCK) + live = rows < total_q + + query = tl.load(QUERY + rows, mask=live, other=0) + in_range = live & (query >= 0) & (query < width) + value = tl.load(VALUE + rows * width + query, mask=in_range, other=0) + # The reference uses -inf for floating point and -1 for integer payloads; a single sentinel + # would be wrong for one of the two. + # Not a ternary: `is_float` is a compile-time constant, and the arm that would build `-1` for a + # float payload (or `-inf` for an int one) does not type-check, so it has to be a real branch. + if is_float: # noqa: SIM108 + result = tl.where(in_range, value, float("-inf")) + else: + result = tl.where(in_range, value, -1) + tl.store(OUT + rows, result, mask=live) + + +@triton.jit +def _map_lookup_key_kernel( + KEY, + VALUE, + QUERY, + OUT, + total_q, + width: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + if row < total_q: + query = tl.load(QUERY + row) + cols = tl.arange(0, BLOCK) + best = tl.full((), float("-inf"), tl.float32) + for start in range(0, width, BLOCK): + idx = start + cols + live = idx < width + key = tl.load(KEY + row * width + idx, mask=live, other=-1) + value = tl.load(VALUE + row * width + idx, mask=live, other=0.0).to(tl.float32) + # Ties are resolved by taking the *last* matching column, which is what the TileLang + # kernel does (it overwrites on every hit). The Python reference uses `max`, so the two + # only agree when the keys are unique -- always true for top-k indices, which is the + # only caller. Matching the kernel rather than the reference keeps the swap invisible. + hit = key == query + last = tl.max(tl.where(hit, idx, -1), axis=0) + picked = tl.sum(tl.where(idx == last, value, 0.0), axis=0) + best = tl.where(last >= 0, picked, best) + tl.store(OUT + row, best) + + +def map_lookup_tl( + force_topk: int, + mode: str = "by_key", + value_dtype: torch.dtype = torch.float32, +) -> Any: + """Build the lookup kernel for a fixed row width and mode.""" + assert mode in ("by_key", "by_index") + assert mode != "by_key" or value_dtype.is_floating_point, "by_key requires floating-point values" + triton_dtype(value_dtype) # fail early on an unmapped dtype rather than inside the launch + + block_index = 128 + block_key = min(next_power_of_2(force_topk), 4096) + + def launch( + key: torch.Tensor | None, + value: torch.Tensor, + query: torch.Tensor, + out: torch.Tensor, + ) -> None: + total_q = query.shape[0] + if total_q == 0: + return + if mode == "by_index": + _map_lookup_index_kernel[(cdiv(total_q, block_index),)]( + value, + query, + out, + total_q, + width=force_topk, + is_float=bool(value.dtype.is_floating_point), + BLOCK=block_index, + num_warps=4, + ) + else: + _map_lookup_key_kernel[(total_q,)]( + key, + value, + query, + out, + total_q, + width=force_topk, + BLOCK=block_key, + num_warps=8, + ) + + return launch + + +@triton.jit +def _copy_block_tables_kernel( + SEQLENS_Q, + CU_SEQLENS_Q, + BLOCK_TABLES_IN, + BLOCK_TABLES_OFFSETS, + BLOCK_TABLES_OUT, + max_num_blocks: tl.constexpr, + BLOCK: tl.constexpr, +): + seq_idx = tl.program_id(0) + q_idx = tl.program_id(1) + chunk = tl.program_id(2) + + sq = tl.load(SEQLENS_Q + seq_idx) + # `tl.where` evaluates both arms, so the index is clamped for the seq_idx == 0 case. + q_base = tl.where(seq_idx == 0, 0, tl.load(CU_SEQLENS_Q + tl.maximum(seq_idx - 1, 0))) + + if q_idx < sq: + q_off = q_base + q_idx + start = tl.load(BLOCK_TABLES_OFFSETS + seq_idx) + num_blocks = tl.load(BLOCK_TABLES_OFFSETS + seq_idx + 1) - start + + blocks = chunk * BLOCK + tl.arange(0, BLOCK) + live = blocks < num_blocks + value = tl.load(BLOCK_TABLES_IN + start + blocks, mask=live, other=0) + tl.store(BLOCK_TABLES_OUT + q_off * max_num_blocks + blocks, value, mask=live) + + +def copy_block_tables_tl(max_seqlen_q: int, max_num_blocks: int, BLOCK_SIZE: int = 1024) -> Any: + """Build the kernel for a fixed query length and table width.""" + + def launch( + seqlens_q: torch.Tensor, + cu_seqlens_q: torch.Tensor, + block_tables_in: torch.Tensor, + block_tables_offsets: torch.Tensor, + block_tables_out: torch.Tensor, + ) -> None: + n_seq = seqlens_q.shape[0] + if n_seq == 0 or max_seqlen_q == 0: + return + _copy_block_tables_kernel[(n_seq, max_seqlen_q, cdiv(max_num_blocks, BLOCK_SIZE))]( + seqlens_q, + cu_seqlens_q, + block_tables_in, + block_tables_offsets, + block_tables_out, + max_num_blocks=max_num_blocks, + BLOCK=BLOCK_SIZE, + num_warps=4, + ) + + return launch diff --git a/scripts/ab_cases.py b/scripts/ab_cases.py index 0b3e94a..775180f 100644 --- a/scripts/ab_cases.py +++ b/scripts/ab_cases.py @@ -15,12 +15,15 @@ from typing import Any import torch +from deepinfer.model_executor.layers.tilelang_ops import ops from deepinfer.model_executor.layers.tilelang_ops.flash_mla_utils import convert_kv_index from deepinfer.model_executor.layers.tilelang_ops.gather import gather from deepinfer.model_executor.layers.tilelang_ops.gather_engram_tokens_zos import gather_engram_tokens_zos from deepinfer.model_executor.layers.tilelang_ops.gather_input_ids_zos import gather_input_ids_zos from deepinfer.model_executor.layers.tilelang_ops.get_kv_index import get_kv_ca_index_decode, get_kv_ca_index_prefill from deepinfer.model_executor.layers.tilelang_ops.kv_index_cat import kv_index_cat +from deepinfer.model_executor.layers.tilelang_ops.norm import norm +from deepinfer.model_executor.layers.tilelang_ops.sampler import map_lookup 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 @@ -291,7 +294,166 @@ def _cases_store_mtp_info() -> list[tuple[str, bool, Callable[[], Any]]]: ] +def _cases_norm() -> list[tuple[str, bool, Callable[[], Any]]]: + """RMSNorm in the forms the engine uses: bf16 out, fp8 with fp32 scales, and both at once.""" + + def make(rows: int, d: int, with_residual: bool, as_fp8: bool, return_bf16: bool, round_sf: bool): + def run() -> Any: + torch.manual_seed(0) + x = torch.randn((rows, d), device="cuda", dtype=torch.bfloat16) + w = torch.randn((d,), device="cuda", dtype=torch.bfloat16) + residual = torch.randn((rows, d), device="cuda", dtype=torch.bfloat16) if with_residual else None + kwargs: dict[str, Any] = {"eps": 1e-20, "out_scale": 1.5} + if as_fp8: + kwargs.update( + fmt="e4m3", + num_per_channels=128, + round_sf=round_sf, + use_tma_aligned_col_major_sf=True, + ) + return norm(x, w, residual=residual, return_bf16=return_bf16, **kwargs) + + return run + + return [ + ("r128_d5120_bf16", False, make(128, 5120, False, False, False, False)), + ("r128_d5120_bf16_residual", False, make(128, 5120, True, False, False, False)), + ("r64_d5120_fp8_round", False, make(64, 5120, False, True, False, True)), + ("r64_d5120_fp8_no_round", False, make(64, 5120, False, True, False, False)), + ("r64_d5120_fp8_round_bf16", False, make(64, 5120, False, True, True, True)), + ("r33_d5120_fp8_round_bf16", False, make(33, 5120, False, True, True, True)), + ("r16_d7168_fp8_round", False, make(16, 7168, False, True, False, True)), + ("r8_d512_fp8_round", False, make(8, 512, False, True, False, True)), + ] + + +def _cases_map_lookup() -> list[tuple[str, bool, Callable[[], Any]]]: + """Both lookup modes, including out-of-range and missing keys.""" + + def make_index(rows: int, width: int, dtype: torch.dtype): + def run() -> Any: + torch.manual_seed(0) + value = ( + torch.randn((rows, width), device="cuda", dtype=dtype) + if dtype.is_floating_point + else torch.randint(0, 1000, (rows, width), device="cuda", dtype=dtype) + ) + query = torch.randint(-width, 2 * width, (rows, 1), device="cuda", dtype=torch.int32) + return map_lookup(None, value, query, mode="by_index") + + return run + + def make_key(rows: int, width: int): + def run() -> Any: + torch.manual_seed(0) + # Distinct keys per row: that is what the indexer produces, and it is also the only + # case where the kernel (last match wins) and the Python reference (max) agree. + key = torch.randperm(width, device="cuda").to(torch.int32).unsqueeze(0).repeat(rows, 1) + value = torch.randn((rows, width), device="cuda", dtype=torch.float32) + # Half the queries exist in the row, half do not. + present = key.gather(-1, torch.randint(0, width, (rows, 1), device="cuda").long()) + missing = torch.full((rows, 1), 10_000, device="cuda", dtype=torch.int32) + keep = torch.randint(0, 2, (rows, 1), device="cuda") == 1 + query = torch.where(keep, present, missing) + return map_lookup(key, value, query, mode="by_key") + + return run + + return [ + ("index_f32", True, make_index(256, 4096, torch.float32)), + ("index_i32", True, make_index(256, 4096, torch.int32)), + ("key_f32", False, make_key(128, 4096)), + ] + + +def _cases_expand_sampling_params() -> list[tuple[str, bool, Callable[[], Any]]]: + """Per-row parameter expansion for the draft rows.""" + + def make(req_cnt: int, total_q: int): + def run() -> Any: + torch.manual_seed(0) + temps = torch.rand((req_cnt,), device="cuda", dtype=torch.float32) + top_ps = torch.rand((req_cnt,), device="cuda", dtype=torch.float32) + top_ks = torch.randint(0, 4096, (req_cnt,), device="cuda", dtype=torch.int32) + positions = torch.randint(0, 1000, (req_cnt,), device="cuda", dtype=torch.int32) + seeds = torch.randint(0, 2**62, (req_cnt,), device="cuda", dtype=torch.uint64) + row_to_req = torch.randint(0, req_cnt, (total_q,), device="cuda", dtype=torch.int32) + local_step = torch.randint(0, 6, (total_q,), device="cuda", dtype=torch.int32) + outs = [ + torch.empty((total_q,), device="cuda", dtype=torch.float32), + torch.empty((total_q,), device="cuda", dtype=torch.float32), + torch.empty((total_q,), device="cuda", dtype=torch.int32), + torch.empty((total_q,), device="cuda", dtype=torch.int32), + torch.empty((total_q,), device="cuda", dtype=torch.uint64), + ] + kernel = ops.expand_sampling_params_tl() + kernel(temps, top_ps, top_ks, positions, seeds, row_to_req, local_step, *outs) + return tuple(outs) + + return run + + return [("n4_q32", True, make(4, 32)), ("n1_q5", True, make(1, 5)), ("n8_q100", True, make(8, 100))] + + +def _cases_prepare_verification() -> list[tuple[str, bool, Callable[[], Any]]]: + """Draft-row metadata: request index, local step and the verify flag.""" + + def make(req_cnt: int, total_q: int, max_len: int): + def run() -> Any: + torch.manual_seed(0) + seqlen_q = torch.randint(1, max_len + 1, (req_cnt,), device="cuda", dtype=torch.int32) + ftp = torch.zeros((req_cnt,), device="cuda", dtype=torch.int32) + ftp[1:] = torch.cumsum(seqlen_q, 0, dtype=torch.int32)[:-1] + row_to_req = torch.full((total_q,), -7, device="cuda", dtype=torch.int32) + local_step = torch.full((total_q,), -7, device="cuda", dtype=torch.int32) + is_verify = torch.zeros((total_q,), device="cuda", dtype=torch.bool) + # Driven through the factory: the interface wrapper allocates its own outputs, which + # would leave nothing to compare. + kernel = ops.prepare_verification_metadata_tl(max_len) + kernel(ftp, seqlen_q, row_to_req, local_step, is_verify) + return row_to_req, local_step, is_verify + + return run + + return [("n4_q24", True, make(4, 24, 6)), ("n1_q1", True, make(1, 1, 1)), ("n8_q48", True, make(8, 48, 6))] + + +def _cases_copy_block_tables() -> list[tuple[str, bool, Callable[[], Any]]]: + """Compact 1D block table expanded to one row per token.""" + + def make(n_seq: int, max_seqlen_q: int, max_blocks: int): + def run() -> Any: + torch.manual_seed(0) + seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32) + cu = torch.cumsum(seqlens, 0, dtype=torch.int32) + total_q = int(cu[-1].item()) + counts = torch.randint(1, max_blocks + 1, (n_seq,), device="cuda", dtype=torch.int32) + offsets = torch.zeros((n_seq + 1,), device="cuda", dtype=torch.int32) + offsets[1:] = torch.cumsum(counts, 0, dtype=torch.int32) + flat = torch.randint(0, 512, (int(offsets[-1].item()),), device="cuda", dtype=torch.int32) + out = torch.full((total_q, max_blocks), -7, device="cuda", dtype=torch.int32) + # No standalone interface function: this kernel is only reached from + # `prepare_v4_decode_metadata`, so the case drives the factory directly. The harness + # swaps `ops.copy_block_tables_tl` between backends, which is all that is needed. + kernel = ops.copy_block_tables_tl(max_seqlen_q, max_blocks) + kernel(seqlens, cu, flat, offsets, out) + return out + + return run + + return [ + ("n4_q6_b64", True, make(4, 6, 64)), + ("n1_q1_b32", True, make(1, 1, 32)), + ("n8_q6_b1400_multi_chunk", True, make(8, 6, 1400)), + ] + + CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = { + "norm_tl": _cases_norm, + "map_lookup_tl": _cases_map_lookup, + "expand_sampling_params_tl": _cases_expand_sampling_params, + "prepare_verification_metadata_tl": _cases_prepare_verification, + "copy_block_tables_tl": _cases_copy_block_tables, "gather_tl": _cases_gather, "gather_input_ids_zos_tl": _cases_gather_input_ids, "gather_engram_tokens_zos_tl": _cases_gather_engram, diff --git a/scripts/h20_ops_ab.py b/scripts/h20_ops_ab.py index a392cca..64eb166 100644 --- a/scripts/h20_ops_ab.py +++ b/scripts/h20_ops_ab.py @@ -57,7 +57,18 @@ def run_both(name: str, fn: Callable[[], Any]) -> tuple[Any, Any]: return expected, actual -def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float = 1e-5, atol: float = 1e-6) -> bool: +# Tolerance follows the dtype's own precision. Anything narrower than float32 is compared at its +# representable resolution, because a different reduction order in a different compiler moves the +# last bit and nothing more: for `norm` that is 10 of 655360 elements differing by one bf16 ulp, +# while the fp8 bytes and every scale factor come out identical. +_DTYPE_TOL: dict[torch.dtype, tuple[float, float]] = { + torch.bfloat16: (1e-2, 1e-2), + torch.float16: (1e-3, 1e-3), + torch.float32: (1e-5, 1e-6), +} + + +def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float | None = None, atol: float | None = None) -> bool: """Report whether one case agrees; returns True on success.""" if isinstance(expected, torch.Tensor): if expected.dtype != actual.dtype: @@ -74,6 +85,21 @@ def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float = idx = bad.nonzero()[0].tolist() print(f" first at {idx}: {expected[tuple(idx)].item()} vs {actual[tuple(idx)].item()}") return bool(same) + if expected.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + # Compare at the format's own resolution: a one-ulp difference is a rounding tie that + # moved, not a wrong value, so byte distance is the honest measure. Require the scales + # to be exact elsewhere and the bytes to agree almost everywhere. + delta = (expected.view(torch.uint8).int() - actual.view(torch.uint8).int()).abs() + frac = float((delta > 0).float().mean()) + if int(delta.max()) <= 2 and frac <= 1e-3: + return True + print(f" FAIL {name}: {frac * 100:.4f}% bytes differ, max byte distance {int(delta.max())}") + return False + + if rtol is None or atol is None: + default_rtol, default_atol = _DTYPE_TOL.get(expected.dtype, (1e-5, 1e-6)) + rtol = default_rtol if rtol is None else rtol + atol = default_atol if atol is None else atol close = torch.allclose(expected.float(), actual.float(), rtol=rtol, atol=atol) if not close: diff = (expected.float() - actual.float()).abs()