From a6c3e69e678b198544f5dbe9f21b70f62de8b2f9 Mon Sep 17 00:00:00 2001 From: loopyt Date: Thu, 22 Oct 2026 00:11:45 +0800 Subject: [PATCH] feat(triton): port eight P1 kernels, all A/B identical to TileLang gather, gather_input_ids_zos, gather_engram_tokens_zos, expand_swa_index, kv_index_cat, convert_kv_index, load_kv_index, store_mtp_info. 45 cases, every one bitwise identical to the TileLang baseline through the shared interface function. `scripts/h20_ops_ab.py` swaps `ops._tl` between the two backends and runs the public call once per side, so the comparison exercises the interface layer's shape handling too, not just the kernel. `scripts/ab_cases.py` holds the cases: real engine widths and dtypes, negative indices for gather, -1 padding for the index kernels, page-boundary lengths, the slot-0 no-write path, and the engram image-sentinel rule. Five Triton semantics traps worth writing down, all documented in the plan: - `tl.where` evaluates both arms, so `tl.where(i == 0, 0, LOAD(i - 1))` really reads `LOAD[-1]`; that was an illegal memory access, fixed by clamping the index instead. - A scalar pointer cannot carry a block-shaped mask; it has to be broadcast with `+ tl.zeros((BLOCK,), tl.int32)` first. - Two arms of one kernel cannot use different tile shapes, which is why `kv_index_cat` became two launches rather than one fused kernel with a column-block branch. - Triton's `data_ptr()` already includes a view's storage offset, so passing `out[:, w:w+k]` and adding `w` again double-counts -- that showed up as 0x01010101 in the second half. - Test-case capacity matters: `store_mtp_info` addresses `slot * nextn + step`, so the pool needs `(max_slot + 1) * nextn` rows; too few rows let both backends write out of bounds and disagree. --- H20-TRITON-PLAN.md | 27 +- .../tilelang_ops/triton/flash_mla_utils.py | 69 +++++ .../triton/gather_engram_tokens_zos.py | 134 ++++++++++ .../triton/gather_input_ids_zos.py | 79 ++++++ .../tilelang_ops/triton/get_kv_index.py | 80 ++++++ .../tilelang_ops/triton/kv_index_cat.py | 112 ++++++++ .../tilelang_ops/triton/store_mtp_info.py | 95 +++++++ .../layers/tilelang_ops/triton/swa_index.py | 58 ++++ scripts/ab_cases.py | 250 ++++++++++++++++++ scripts/h20_ops_ab.py | 14 +- 10 files changed, 916 insertions(+), 2 deletions(-) create mode 100644 deepinfer/model_executor/layers/tilelang_ops/triton/flash_mla_utils.py create mode 100644 deepinfer/model_executor/layers/tilelang_ops/triton/gather_engram_tokens_zos.py create mode 100644 deepinfer/model_executor/layers/tilelang_ops/triton/gather_input_ids_zos.py create mode 100644 deepinfer/model_executor/layers/tilelang_ops/triton/get_kv_index.py create mode 100644 deepinfer/model_executor/layers/tilelang_ops/triton/kv_index_cat.py create mode 100644 deepinfer/model_executor/layers/tilelang_ops/triton/store_mtp_info.py create mode 100644 deepinfer/model_executor/layers/tilelang_ops/triton/swa_index.py diff --git a/H20-TRITON-PLAN.md b/H20-TRITON-PLAN.md index 42b7b79..89a2624 100644 --- a/H20-TRITON-PLAN.md +++ b/H20-TRITON-PLAN.md @@ -104,4 +104,29 @@ ## 七、进度 -见文末「执行日志」。 +| 阶段 | 状态 | 说明 | +|---|---|---| +| 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 | +| P2 KV cache | ⬜ | | +| P3 MTP 链路 | ⬜ | | +| P4 采样 | ⬜ | | +| P5 端到端 | ⬜ | | + +A/B 现状：**45 个用例全过**（`python scripts/h20_ops_ab.py`）。 + +### 过程中踩到的 Triton 语义坑（都已修，可复用） + +1. **`tl.where` 两侧都会求值**。`tl.where(i == 0, 0, LOAD(i - 1))` 在 `i == 0` 时真的会读 + `LOAD[-1]` → illegal memory access。凡是「边界索引用 where 兜住」的写法都要先 + `tl.maximum(idx, 0)` 钳制。 +2. **标量指针不能配 block mask**。`tl.load(PTR + scalar, mask=)` 报 + `Mask argument cannot be block type if pointer argument is not a block`。要 + `PTR + scalar + tl.zeros((BLOCK,), tl.int32)` 广播成 block。 +3. **一个 kernel 里两个分支用不同 tile 形状**编不过。`kv_index_cat` 拆成两个 kernel 分别 + launch（同一 stream，语义不变）。 +4. **视图的 storage offset 已含在 `data_ptr()` 里**。给 kernel 传 + `out[:, w:w+k]` 后再加一次 `w` 就重复偏移了 —— 表现为后半段全是 `0x01010101` 之类的 + uninitialized 内存。 +5. **测试用例的容量也要算对**：`store_mtp_info` 的 pool 是 `slot * nextn + step` 寻址， + 行数必须是 `(max_slot + 1) * nextn`，否则两个后端一起写越界，表现成互相不一致。 diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/flash_mla_utils.py b/deepinfer/model_executor/layers/tilelang_ops/triton/flash_mla_utils.py new file mode 100644 index 0000000..e924e44 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/flash_mla_utils.py @@ -0,0 +1,69 @@ +"""Logical top-k token indices to physical cache slots, in place, in Triton.""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from deepinfer.model_executor.layers.tilelang_ops.triton.common import next_power_of_2 + + +@triton.jit +def _convert_kv_index_kernel( + TOPK_IDXS, + CONTEXT_LENS, + BLOCK_TABLES, + s, + k: tl.constexpr, + n: tl.constexpr, + block_size: tl.constexpr, + BLOCK_K: tl.constexpr, +): + sid = tl.program_id(0) + bid = tl.program_id(1) + + base = (bid * s + sid) * k + offs = tl.arange(0, BLOCK_K) + context_len = tl.load(CONTEXT_LENS + bid * s + sid) + + idx = tl.load(TOPK_IDXS + base + offs, mask=offs < k, other=-1) + # Entries past the context length are padding; the reference marks them -1 rather than + # translating them, so they must not go through the block-table lookup. + live = (offs < k) & (offs < context_len) + page = tl.load(BLOCK_TABLES + bid * n + idx // block_size, mask=live, other=0) + physical = page * block_size + idx % block_size + + tl.store(TOPK_IDXS + base + offs, tl.where(live, physical, -1), mask=offs < k) + + +def convert_kv_index_tl( + s: int, + k: int, + n: int, + block_size: int, +) -> Any: + """Build the kernel for a fixed query length, top-k and page size.""" + block_k = next_power_of_2(k) + + def launch( + topk_idxs: torch.Tensor, + context_lens: torch.Tensor, + block_tables: torch.Tensor, + ) -> None: + b = topk_idxs.shape[0] + if b == 0 or s == 0: + return + _convert_kv_index_kernel[(s, b)]( + topk_idxs, + context_lens, + block_tables, + s, + k=k, + n=n, + block_size=block_size, + BLOCK_K=block_k, + num_warps=4, + ) + + return launch diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/gather_engram_tokens_zos.py b/deepinfer/model_executor/layers/tilelang_ops/triton/gather_engram_tokens_zos.py new file mode 100644 index 0000000..ef9a8b3 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/gather_engram_tokens_zos.py @@ -0,0 +1,134 @@ +"""Per-request n-gram token windows for engram lookup, in Triton.""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from deepinfer.model_executor.layers.tilelang_ops.triton.common import next_power_of_2 + + +@triton.jit +def _gather_engram_tokens_zos_kernel( + WIN_COMPRESSED, + CURR_FRESH, + PREV_SAMPLED, + PREV_PREDICT, + PREV_NUM_ACCEPTED, + REMAP, + LOOKUP_TABLE, + SEQLENS_Q, + CU_SEQLENS_Q, + OUT, + max_seqlen_q: tl.constexpr, + max_ngram_size: tl.constexpr, + window_size: tl.constexpr, + engram_compressed_pad_token: tl.constexpr, + image_sentinel: tl.constexpr, + BLOCK_Q: tl.constexpr, +): + seq_idx = tl.program_id(0) + sq = tl.load(SEQLENS_Q + seq_idx) + # `tl.where` evaluates both arms, so the index is clamped: at seq_idx == 0 the other arm + # would read CU_SEQLENS_Q[-1]. + q_base = tl.where(seq_idx == 0, 0, tl.load(CU_SEQLENS_Q + tl.maximum(seq_idx - 1, 0))) + pidx = tl.load(REMAP + seq_idx) + # A remapped request reuses `nextn - accepted` tokens it already produced, so its window starts + # that much later; a fresh request starts at the beginning of the window. + nr = tl.where(pidx >= 0, (max_seqlen_q - 1) - tl.load(PREV_NUM_ACCEPTED + pidx), max_seqlen_q - 1) + base = window_size - nr + + offs = tl.arange(0, BLOCK_Q) + active = offs < sq + + # Once an image sentinel is seen the rest of the window is padding, so the r loop carries a + # per-lane `hit` flag forward instead of re-scanning. + hit = tl.zeros((BLOCK_Q,), dtype=tl.int1) + for r in tl.static_range(max_ngram_size): + from_current = offs >= r + c = offs - r + + token_cur = tl.load( + PREV_SAMPLED + pidx + tl.zeros((BLOCK_Q,), tl.int32), + mask=active & from_current & (pidx >= 0), + other=0, + ) + token_draft = tl.load( + PREV_PREDICT + pidx * (max_seqlen_q - 1) + (c - 1), + mask=active & from_current & (pidx >= 0) & (c > 0), + other=0, + ) + token_fresh = tl.load( + CURR_FRESH + seq_idx * max_seqlen_q + c, + mask=active & from_current & (pidx < 0), + other=0, + ) + token_window = tl.load( + WIN_COMPRESSED + seq_idx * window_size + base + offs - r, + mask=active & ~from_current, + other=0, + ) + + raw = tl.where( + from_current, + tl.where(pidx < 0, token_fresh, tl.where(c == 0, token_cur, token_draft)), + token_window, + ) + # Fresh tokens are already vocabulary ids; replayed ones are compressed and need the map. + mapped = tl.load(LOOKUP_TABLE + raw, mask=active & from_current & (pidx >= 0), other=0) + token = tl.where(from_current & (pidx >= 0), mapped, raw) + + is_image = token == image_sentinel + value = tl.where(hit | is_image, engram_compressed_pad_token, token) + tl.store(OUT + (q_base + offs) * max_ngram_size + r, value, mask=active) + hit = hit | is_image + + +def gather_engram_tokens_zos_tl( + max_seqlen_q: int, + max_ngram_size: int, + engram_compressed_pad_token: int, + image_sentinel: int, +) -> Any: + """Build the kernel for a fixed draft depth and n-gram size.""" + block_q = next_power_of_2(max_seqlen_q) + nextn = max_seqlen_q - 1 + window_size = nextn + max_ngram_size - 1 + + def launch( + win_compressed: torch.Tensor, + curr_fresh: torch.Tensor, + prev_sampled_token_ids: torch.Tensor, + prev_predict_token_ids: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + remap: torch.Tensor, + lookup_table: torch.Tensor, + seqlens_q: torch.Tensor, + cu_seqlens_q: torch.Tensor, + out_ngram_token_ids: torch.Tensor, + ) -> None: + n_seq = curr_fresh.shape[0] + if n_seq == 0: + return + _gather_engram_tokens_zos_kernel[(n_seq,)]( + win_compressed, + curr_fresh, + prev_sampled_token_ids, + prev_predict_token_ids, + prev_num_accepted_tokens, + remap, + lookup_table, + seqlens_q, + cu_seqlens_q, + out_ngram_token_ids, + max_seqlen_q=max_seqlen_q, + max_ngram_size=max_ngram_size, + window_size=window_size, + engram_compressed_pad_token=engram_compressed_pad_token, + image_sentinel=image_sentinel, + BLOCK_Q=block_q, + num_warps=4, + ) + + return launch diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/gather_input_ids_zos.py b/deepinfer/model_executor/layers/tilelang_ops/triton/gather_input_ids_zos.py new file mode 100644 index 0000000..2a9609d --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/gather_input_ids_zos.py @@ -0,0 +1,79 @@ +"""Per-request row expansion of previous-step token ids, in Triton.""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from deepinfer.model_executor.layers.tilelang_ops.triton.common import next_power_of_2 + + +@triton.jit +def _gather_input_ids_zos_kernel( + FRESH_TOKENS, + PREV_SAMPLED, + PREV_PREDICT, + REMAP, + SEQLENS_Q, + CU_SEQLENS_Q, + OUT, + max_seqlen_q: tl.constexpr, + BLOCK_Q: tl.constexpr, +): + seq_idx = tl.program_id(0) + sq = tl.load(SEQLENS_Q + seq_idx) + # `tl.where` evaluates both arms, so the index is clamped: at seq_idx == 0 the other arm + # would read CU_SEQLENS_Q[-1]. + q_base = tl.where(seq_idx == 0, 0, tl.load(CU_SEQLENS_Q + tl.maximum(seq_idx - 1, 0))) + pidx = tl.load(REMAP + seq_idx) + + offs = tl.arange(0, BLOCK_Q) + active = offs < sq + + # Fresh prompt tokens come from this request's own row; a remapped request instead replays the + # previous step's sampled token for position 0 and its draft predictions after that. + fresh = tl.load(FRESH_TOKENS + seq_idx * max_seqlen_q + offs, mask=active, other=0) + # `pidx` is a scalar, so the pointer has to be broadcast to a block before it can carry a + # block-shaped mask. + prev_sampled = tl.load(PREV_SAMPLED + pidx + tl.zeros((BLOCK_Q,), tl.int32), mask=active & (pidx >= 0) & (offs == 0), other=0) + prev_predict = tl.load( + PREV_PREDICT + pidx * (max_seqlen_q - 1) + (offs - 1), + mask=active & (pidx >= 0) & (offs > 0), + other=0, + ) + + value = tl.where(pidx < 0, fresh, tl.where(offs == 0, prev_sampled, prev_predict)) + tl.store(OUT + q_base + offs, value, mask=active) + + +def gather_input_ids_zos_tl(max_seqlen_q: int) -> Any: + """Build the kernel for a fixed draft depth; the request count stays a runtime length.""" + block_q = next_power_of_2(max_seqlen_q) + + def launch( + fresh_tokens: torch.Tensor, + prev_sampled_token_ids: torch.Tensor, + prev_predict_token_ids: torch.Tensor, + remap: torch.Tensor, + seqlens_q: torch.Tensor, + cu_seqlens_q: torch.Tensor, + out_input_ids: torch.Tensor, + ) -> None: + n_seq = fresh_tokens.shape[0] + if n_seq == 0: + return + _gather_input_ids_zos_kernel[(n_seq,)]( + fresh_tokens, + prev_sampled_token_ids, + prev_predict_token_ids, + remap, + seqlens_q, + cu_seqlens_q, + out_input_ids, + max_seqlen_q=max_seqlen_q, + BLOCK_Q=block_q, + num_warps=4, + ) + + return launch diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/get_kv_index.py b/deepinfer/model_executor/layers/tilelang_ops/triton/get_kv_index.py new file mode 100644 index 0000000..ed75bce --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/get_kv_index.py @@ -0,0 +1,80 @@ +"""Expand per-row (offset, length) into cache index rows, in Triton.""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2 + + +@triton.jit +def _load_kv_index_kernel( + OFFSET_LEN, + KV_INDEX, + FLATTEN_BLOCK_TABLES, + index_length: tl.constexpr, + block_size: tl.constexpr, + has_block_tables: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + col_block = tl.program_id(1) + + offset = tl.load(OFFSET_LEN + row * 2 + 0) + length = tl.load(OFFSET_LEN + row * 2 + 1) + + cols = col_block * BLOCK + tl.arange(0, BLOCK) + live = cols < length + + if has_block_tables: + # Contiguous positions, then mapped through the page table into a physical slot. + pos = offset + cols + page = tl.load(FLATTEN_BLOCK_TABLES + pos // block_size, mask=live, other=0) + index = page * block_size + pos % block_size + elif block_size == 0: + index = offset + cols + else: + # Ring buffer confined to the block that `offset` falls in. + index = offset // block_size * block_size + (offset + cols) % block_size + + tl.store( + KV_INDEX + row * index_length + cols, + tl.where(live, index, -1), + mask=cols < index_length, + ) + + +def load_kv_index_tl( + index_length: int, + block_size: int, + has_block_tables: bool, +) -> Any: + """Build the kernel for a fixed row length and addressing mode.""" + block = min(2048, next_power_of_2(index_length)) + + def launch( + kv_index_offset_len: torch.Tensor, + kv_index: torch.Tensor, + flatten_block_tables: torch.Tensor, + ) -> None: + n = kv_index_offset_len.shape[0] + if n == 0 or index_length == 0: + return + if flatten_block_tables is None: + # The contiguous mode takes `None` here, and Triton needs a real pointer even when the + # branch that would read it is compiled out. + flatten_block_tables = kv_index_offset_len + _load_kv_index_kernel[(n, cdiv(index_length, block))]( + kv_index_offset_len, + kv_index, + flatten_block_tables, + index_length=index_length, + block_size=block_size, + has_block_tables=has_block_tables, + BLOCK=block, + num_warps=4, + ) + + return launch diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/kv_index_cat.py b/deepinfer/model_executor/layers/tilelang_ops/triton/kv_index_cat.py new file mode 100644 index 0000000..14e154d --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/kv_index_cat.py @@ -0,0 +1,112 @@ +"""Concatenate the window index and the shifted compressed index, in Triton. + +Two launches rather than one fused kernel: the two halves have different widths and different +element rules (verbatim copy vs. conditional rebase), so keeping them separate avoids a +compile-time branch whose two arms would need the same tile shape. +""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2 + + +@triton.jit +def _copy_window_kernel( + SRC, + DST, + total_q, + window_size: tl.constexpr, + out_stride: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_W: tl.constexpr, +): + row_block = tl.program_id(0) + col_block = tl.program_id(1) + + rows = row_block * BLOCK_M + tl.arange(0, BLOCK_M) + cols = col_block * BLOCK_W + tl.arange(0, BLOCK_W) + mask = (rows[:, None] < total_q) & (cols[None, :] < window_size) + + value = tl.load(SRC + rows[:, None] * window_size + cols[None, :], mask=mask, other=0) + tl.store(DST + rows[:, None] * out_stride + cols[None, :], value, mask=mask) + + +@triton.jit +def _shift_compressed_kernel( + SRC, + DST, + total_q, + topk_tokens: tl.constexpr, + out_stride: tl.constexpr, + shift: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, +): + row_block = tl.program_id(0) + col_block = tl.program_id(1) + + rows = row_block * BLOCK_M + tl.arange(0, BLOCK_M) + cols = col_block * BLOCK_K + tl.arange(0, BLOCK_K) + mask = (rows[:, None] < total_q) & (cols[None, :] < topk_tokens) + + value = tl.load(SRC + rows[:, None] * topk_tokens + cols[None, :], mask=mask, other=-1) + # `kv_ca_index_out` is already a view whose base sits at column `window_size`, so only the + # row stride is applied here. Padding stays -1; shifting it would make it a plausible slot. + tl.store( + DST + rows[:, None] * out_stride + cols[None, :], + tl.where(value == -1, -1, value + shift), + mask=mask, + ) + + +def kv_index_cat_tl( + dtype: str, + window_size: int, + topk_tokens: int, + out_stride: int, +) -> Any: + """Build the kernel for a fixed window width, top-k and output row stride.""" + del dtype # the interface layer only ever passes int32 + + block_m = 8 + block_w = min(1024, next_power_of_2(window_size)) + block_k = min(1024, next_power_of_2(topk_tokens)) + + def launch( + kv_index: torch.Tensor, + kv_ca_index: torch.Tensor, + kv_index_out: torch.Tensor, + kv_ca_index_out: torch.Tensor, + total_k_swa: int, + total_k: int, + ) -> None: + total_q = kv_index.shape[0] + if total_q == 0: + return + _copy_window_kernel[(cdiv(total_q, block_m), cdiv(window_size, block_w))]( + kv_index, + kv_index_out, + total_q, + window_size=window_size, + out_stride=out_stride, + BLOCK_M=block_m, + BLOCK_W=block_w, + num_warps=4, + ) + _shift_compressed_kernel[(cdiv(total_q, block_m), cdiv(topk_tokens, block_k))]( + kv_ca_index, + kv_ca_index_out, + total_q, + topk_tokens=topk_tokens, + out_stride=out_stride, + shift=total_k + total_k_swa, + BLOCK_M=block_m, + BLOCK_K=block_k, + num_warps=4, + ) + + return launch diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/store_mtp_info.py b/deepinfer/model_executor/layers/tilelang_ops/triton/store_mtp_info.py new file mode 100644 index 0000000..2debf19 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/store_mtp_info.py @@ -0,0 +1,95 @@ +"""Copy the draft head's per-step outputs into the MTP pool, in Triton.""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2 + + +@triton.jit +def _store_mtp_info_kernel( + WRITE_SLOTS, + TOKEN_IDS, + LOGPROBS, + LOGPROBS_INDEX, + CONFIDENCE, + TOKEN_IDS_POOL, + LOGPROBS_POOL, + LOGPROBS_INDEX_POOL, + CONFIDENCE_POOL, + nextn: tl.constexpr, + vocab_size: tl.constexpr, + has_index: tl.constexpr, + has_confidence: tl.constexpr, + BLOCK_V: tl.constexpr, +): + pid = tl.program_id(0) + col_block = tl.program_id(1) + req = pid // nextn + step = pid % nextn + + # Slot 0 means "nothing to write" -- the pool row is pre-seeded by the caller. + slot = tl.load(WRITE_SLOTS + req) + if slot != 0: + cols = col_block * BLOCK_V + tl.arange(0, BLOCK_V) + mask = cols < vocab_size + src = (req * nextn + step) * vocab_size + cols + dst = (slot * nextn + step) * vocab_size + cols + + logprobs = tl.load(LOGPROBS + src, mask=mask, other=0.0) + tl.store(LOGPROBS_POOL + dst, logprobs, mask=mask) + if has_index: + logprobs_index = tl.load(LOGPROBS_INDEX + src, mask=mask, other=0) + tl.store(LOGPROBS_INDEX_POOL + dst, logprobs_index, mask=mask) + + # The token ids and confidences are per (request, step), not per vocabulary entry, so only + # the first column block writes them. + if (col_block == 0) & (step == 0): + for i in tl.static_range(nextn): + token = tl.load(TOKEN_IDS + req * nextn + i) + tl.store(TOKEN_IDS_POOL + slot * nextn + i, token) + if has_confidence: + conf = tl.load(CONFIDENCE + req * nextn + i) + tl.store(CONFIDENCE_POOL + slot * nextn + i, conf) + + +def store_mtp_info_tl(nextn: int, vocab_size: int, has_index: bool, has_confidence: bool) -> Any: + """Build the kernel for a fixed draft depth and vocabulary size.""" + block_v = min(next_power_of_2(vocab_size), 4096) + + def launch( + write_slots: torch.Tensor, + token_ids: torch.Tensor, + logprobs: torch.Tensor, + logprobs_index: torch.Tensor, + confidence: torch.Tensor, + token_ids_pool: torch.Tensor, + logprobs_pool: torch.Tensor, + logprobs_index_pool: torch.Tensor, + confidence_pool: torch.Tensor, + ) -> None: + req_cnt = write_slots.shape[0] + if req_cnt == 0: + return + _store_mtp_info_kernel[(req_cnt * nextn, cdiv(vocab_size, block_v))]( + write_slots, + token_ids, + logprobs, + logprobs_index, + confidence, + token_ids_pool, + logprobs_pool, + logprobs_index_pool, + confidence_pool, + nextn=nextn, + vocab_size=vocab_size, + has_index=has_index, + has_confidence=has_confidence, + BLOCK_V=block_v, + num_warps=8, + ) + + return launch diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/swa_index.py b/deepinfer/model_executor/layers/tilelang_ops/triton/swa_index.py new file mode 100644 index 0000000..409fd7d --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/swa_index.py @@ -0,0 +1,58 @@ +"""Expand four per-row pivots into a full sliding-window index row, in Triton.""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2 + + +@triton.jit +def _expand_swa_index_kernel( + PIVOTS, + OUTPUT, + n, + width: tl.constexpr, + BLOCK_W: tl.constexpr, +): + row = tl.program_id(0) + col_block = tl.program_id(1) + + if row < n: + offset1 = tl.load(PIVOTS + row * 4 + 0) + len1 = tl.load(PIVOTS + row * 4 + 1) + offset2 = tl.load(PIVOTS + row * 4 + 2) + len2 = tl.load(PIVOTS + row * 4 + 3) + + cols = col_block * BLOCK_W + tl.arange(0, BLOCK_W) + # Both spans are contiguous ranges; past the second one the column is padding. + value = tl.where( + cols < len1, + offset1 + cols, + tl.where(cols < len1 + len2, offset2 - len1 + cols, -1), + ) + tl.store(OUTPUT + row * width + cols, value, mask=cols < width) + + +def expand_swa_index_tl(width: int) -> Any: + """Build the kernel for a fixed row width.""" + # A row is written by BLOCK_W-wide blocks; 1024 int32 per block is one coalesced pass with a + # comfortable register budget, and narrower rows simply use the whole width in one block. + block_w = min(1024, next_power_of_2(width)) + + def launch(pivots: torch.Tensor, output: torch.Tensor) -> None: + n = pivots.shape[0] + if n == 0: + return + _expand_swa_index_kernel[(n, cdiv(width, block_w))]( + pivots, + output, + n, + width=width, + BLOCK_W=block_w, + num_warps=4, + ) + + return launch diff --git a/scripts/ab_cases.py b/scripts/ab_cases.py index 8a13fcd..0b3e94a 100644 --- a/scripts/ab_cases.py +++ b/scripts/ab_cases.py @@ -15,7 +15,14 @@ from typing import Any import torch +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.store_mtp_info import store_mtp_info +from deepinfer.model_executor.layers.tilelang_ops.swa_index import expand_swa_index def _cases_gather() -> list[tuple[str, bool, Callable[[], Any]]]: @@ -48,6 +55,249 @@ def _cases_gather() -> list[tuple[str, bool, Callable[[], Any]]]: ] +def _cases_gather_input_ids() -> list[tuple[str, bool, Callable[[], Any]]]: + """Draft-token replay: remapped requests replay previous ids, fresh ones use their own row.""" + + def make(n_seq: int, max_seqlen_q: int, prev_n_seq: int, mix: bool): + def run() -> Any: + torch.manual_seed(0) + nextn = max_seqlen_q - 1 + fresh = torch.randint(0, 1000, (n_seq, max_seqlen_q), device="cuda", dtype=torch.int32) + prev_sampled = torch.randint(0, 1000, (prev_n_seq,), device="cuda", dtype=torch.int32) + prev_predict = torch.randint(0, 1000, (prev_n_seq, nextn), device="cuda", dtype=torch.int32) + if mix: + remap = torch.tensor([-1 if i % 2 == 0 else i % prev_n_seq for i in range(n_seq)], device="cuda", dtype=torch.int32) + else: + remap = torch.full((n_seq,), -1, device="cuda", dtype=torch.int32) + seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32) + cu = torch.cumsum(seqlens, 0, dtype=torch.int32) + out = torch.full((int(cu[-1].item()),), -7, device="cuda", dtype=torch.int32) + gather_input_ids_zos(fresh, prev_sampled, prev_predict, remap, seqlens, cu, out, max_seqlen_q=max_seqlen_q) + return out + + return run + + return [ + ("n8_nextn5_mixed", True, make(8, 6, 8, True)), + ("n8_nextn5_all_fresh", True, make(8, 6, 8, False)), + ("n1_nextn1_mixed", True, make(1, 2, 3, True)), + ("n64_nextn5_mixed", True, make(64, 6, 64, True)), + ("n5_nextn2_full_len", True, make(5, 3, 5, True)), + ] + + +def _cases_gather_engram() -> list[tuple[str, bool, Callable[[], Any]]]: + """Engram n-gram windows, including the image-sentinel padding rule.""" + + def make(n_seq: int, max_seqlen_q: int, max_ngram: int, vocab: int, with_image: bool): + def run() -> Any: + torch.manual_seed(0) + nextn = max_seqlen_q - 1 + window = nextn + max_ngram - 1 + pad_token = 7 + sentinel = -1 + win = torch.randint(0, 50, (n_seq, window), device="cuda", dtype=torch.int32) + if with_image: + win[:, window // 2] = sentinel + fresh = torch.randint(0, 50, (n_seq, max_seqlen_q), device="cuda", dtype=torch.int32) + if with_image: + fresh[:, -1] = sentinel + prev_sampled = torch.randint(0, vocab, (n_seq,), device="cuda", dtype=torch.int32) + prev_predict = torch.randint(0, vocab, (n_seq, nextn), device="cuda", dtype=torch.int32) + prev_accepted = torch.randint(0, nextn + 1, (n_seq,), device="cuda", dtype=torch.int32) + remap = torch.tensor([-1 if i % 2 == 0 else i for i in range(n_seq)], device="cuda", dtype=torch.int32) + lookup = torch.arange(vocab, device="cuda", dtype=torch.int32) % 13 + seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32) + cu = torch.cumsum(seqlens, 0, dtype=torch.int32) + out = torch.full((int(cu[-1].item()), max_ngram), -7, device="cuda", dtype=torch.int32) + gather_engram_tokens_zos( + win, + fresh, + prev_sampled, + prev_predict, + prev_accepted, + remap, + lookup, + seqlens, + cu, + out, + max_seqlen_q=max_seqlen_q, + max_ngram_size=max_ngram, + engram_compressed_pad_token=pad_token, + image_sentinel=sentinel, + ) + return out + + return run + + return [ + ("n8_nextn5_ngram3", True, make(8, 6, 3, 64, False)), + ("n8_nextn5_ngram1", True, make(8, 6, 1, 64, False)), + ("n8_nextn5_ngram3_image", True, make(8, 6, 3, 64, True)), + ("n1_nextn1_ngram2", True, make(1, 2, 2, 64, False)), + ] + + +def _cases_swa_index() -> list[tuple[str, bool, Callable[[], Any]]]: + """Sliding-window index expansion for the widths the engine uses.""" + + def make(n: int, width: int, overflow: bool): + def run() -> Any: + torch.manual_seed(0) + pivots = torch.zeros((n, 4), device="cuda", dtype=torch.int32) + pivots[:, 0] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32) + if overflow: + # len1 + len2 can exceed width; the third branch must still win past len1 + len2. + pivots[:, 1] = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32) + pivots[:, 2] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32) + pivots[:, 3] = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32) + else: + total = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32) + pivots[:, 1] = total // 2 + pivots[:, 2] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32) + pivots[:, 3] = total - total // 2 + return expand_swa_index(pivots, width) + + return run + + return [ + ("n32_w512", True, make(32, 512, False)), + ("n32_w513_overflow", True, make(32, 513, True)), + ("n1_w1", True, make(1, 1, False)), + ("n37_w2048_overflow", True, make(37, 2048, True)), + ("n8_w64", True, make(8, 64, False)), + ] + + +def _cases_kv_index_cat() -> list[tuple[str, bool, Callable[[], Any]]]: + """Window + compressed index concatenation with the -1 padding preserved.""" + + def make(total_q: int, window: int, topk: int): + def run() -> Any: + torch.manual_seed(0) + kv_index = torch.randint(-1, 4000, (total_q, window), device="cuda", dtype=torch.int32) + kv_ca = torch.randint(-1, 300, (total_q, topk), device="cuda", dtype=torch.int32) + return kv_index_cat(kv_index, kv_ca, 512, 2000) + + return run + + return [ + ("q8_w512_k128", True, make(8, 512, 128)), + ("q1_w64_k32", True, make(1, 64, 32)), + ("q5_w2048_k1024", True, make(5, 2048, 1024)), + ] + + +def _cases_convert_kv_index() -> list[tuple[str, bool, Callable[[], Any]]]: + """Logical to physical index translation through a page table.""" + + def make(b: int, s: int, k: int, block_size: int, max_blocks: int): + def run() -> Any: + torch.manual_seed(0) + topk = torch.randint(0, max_blocks * block_size, (b, s, k), device="cuda", dtype=torch.int32) + ctx = torch.randint(0, k + 1, (b, s), device="cuda", dtype=torch.int32) + table = torch.randint(0, 50, (b, max_blocks), device="cuda", dtype=torch.int32) + convert_kv_index(topk, block_size, ctx, table) + return topk + + return run + + return [ + ("b2_s4_k64_bs64", True, make(2, 4, 64, 64, 32)), + ("b1_s1_k128_bs128", True, make(1, 1, 128, 128, 16)), + ("b3_s2_k16_bs8", True, make(3, 2, 16, 8, 64)), + ] + + +def _cases_load_kv_index() -> list[tuple[str, bool, Callable[[], Any]]]: + """Prefill (contiguous) and decode (paged) index expansion.""" + + def make_prefill(n: int, length: int): + def run() -> Any: + torch.manual_seed(0) + off_len = torch.zeros((n, 2), device="cuda", dtype=torch.int32) + off_len[:, 0] = torch.randint(0, 500, (n,), device="cuda", dtype=torch.int32) + off_len[:, 1] = torch.randint(0, length + 1, (n,), device="cuda", dtype=torch.int32) + return get_kv_ca_index_prefill(off_len, length) + + return run + + def make_decode(n: int, length: int, tpb: int, num_blocks: int): + def run() -> Any: + torch.manual_seed(0) + off_len = torch.zeros((n, 2), device="cuda", dtype=torch.int32) + off_len[:, 0] = torch.randint(0, 500, (n,), device="cuda", dtype=torch.int32) + off_len[:, 1] = torch.randint(0, length + 1, (n,), device="cuda", dtype=torch.int32) + table = torch.randint(0, num_blocks, (num_blocks,), device="cuda", dtype=torch.int32) + return get_kv_ca_index_decode(off_len, length, tpb, table) + + return run + + return [ + ("prefill_n8_len512", True, make_prefill(8, 512)), + ("prefill_n1_len1024", True, make_prefill(1, 1024)), + ("prefill_n33_len100", True, make_prefill(33, 100)), + ("decode_n8_len512_tpb64", True, make_decode(8, 512, 64, 128)), + ("decode_n5_len128_tpb128", True, make_decode(5, 128, 128, 64)), + ] + + +def _cases_store_mtp_info() -> list[tuple[str, bool, Callable[[], Any]]]: + """Draft logprob pooling, including the slot-0 no-write path.""" + + def make(req_cnt: int, nextn: int, vocab: int, has_index: bool, has_confidence: bool, zero_slots: bool): + def run() -> Any: + torch.manual_seed(0) + # The pool is addressed as slot * nextn + step and slots run 0..req_cnt, so it needs + # room for every slot including the reserved one. + pool_rows = (req_cnt + 1) * nextn + write_slots = torch.arange(1, req_cnt + 1, device="cuda", dtype=torch.int32) + if zero_slots: + write_slots[0] = 0 + token_ids = torch.randint(0, vocab, (req_cnt, nextn), device="cuda", dtype=torch.int32) + logprobs = torch.randn((req_cnt, nextn, vocab), device="cuda", dtype=torch.float32) + logprobs_index = torch.randint(0, vocab, (req_cnt, nextn, vocab), device="cuda", dtype=torch.int32) + confidence = torch.rand((req_cnt, nextn), device="cuda", dtype=torch.float32) + token_ids_pool = torch.full((pool_rows,), -7, device="cuda", dtype=torch.int32) + logprobs_pool = torch.full((pool_rows, vocab), -7.0, device="cuda", dtype=torch.float32) + logprobs_index_pool = torch.full((pool_rows, vocab), -7, device="cuda", dtype=torch.int32) + confidence_pool = torch.full((pool_rows,), -7.0, device="cuda", dtype=torch.float32) + store_mtp_info( + write_slots, + token_ids, + logprobs, + logprobs_index, + confidence, + token_ids_pool, + logprobs_pool, + logprobs_index_pool, + confidence_pool, + ) + return ( + token_ids_pool.clone(), + logprobs_pool.clone(), + logprobs_index_pool.clone(), + confidence_pool.clone(), + ) + + return run + + return [ + ("n4_nextn5_v1k_full", True, make(4, 5, 1024, True, True, False)), + ("n4_nextn5_v1k_no_index", True, make(4, 5, 1024, False, True, False)), + ("n4_nextn5_v1k_no_conf", True, make(4, 5, 1024, True, False, False)), + ("n4_nextn5_v1k_slot0", True, make(4, 5, 1024, True, True, True)), + ("n2_nextn1_v129280", True, make(2, 1, 129280, True, True, False)), + ] + + CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = { "gather_tl": _cases_gather, + "gather_input_ids_zos_tl": _cases_gather_input_ids, + "gather_engram_tokens_zos_tl": _cases_gather_engram, + "expand_swa_index_tl": _cases_swa_index, + "kv_index_cat_tl": _cases_kv_index_cat, + "convert_kv_index_tl": _cases_convert_kv_index, + "load_kv_index_tl": _cases_load_kv_index, + "store_mtp_info_tl": _cases_store_mtp_info, } diff --git a/scripts/h20_ops_ab.py b/scripts/h20_ops_ab.py index 8a36b6c..a392cca 100644 --- a/scripts/h20_ops_ab.py +++ b/scripts/h20_ops_ab.py @@ -81,6 +81,17 @@ def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float = print(f" FAIL {name}: max abs {diff.max().item():.3e}, max rel {(diff / scale).max().item():.3e}") return bool(close) + if isinstance(expected, (tuple, list)): + # A case that returns several tensors is only right if every element is; recurse so each + # one gets its own report instead of collapsing into one boolean. + if len(expected) != len(actual): + print(f" FAIL {name}: {len(expected)} outputs vs {len(actual)}") + return False + ok = True + for i, (e, a) in enumerate(zip(expected, actual, strict=True)): + ok &= compare(f"{name}[{i}]", e, a, exact=exact, rtol=rtol, atol=atol) + return ok + same = expected == actual if not same: print(f" FAIL {name}: {expected!r} vs {actual!r}") @@ -113,7 +124,8 @@ def main() -> None: try: expected, actual = run_both(kernel, fn) except Exception as exc: # noqa: BLE001 - a raise on one backend is itself a result - print(f" FAIL {case_name}: {type(exc).__name__}: {str(exc).splitlines()[0][:160]}") + lines = str(exc).splitlines() + print(f" FAIL {case_name}: {type(exc).__name__}: {lines[0][:160] if lines else ''}") failures += 1 continue if not compare(case_name, expected, actual, exact=exact):