From d39358459233114c5920f587c9fa558afd7a4f1a Mon Sep 17 00:00:00 2001 From: loopyt Date: Thu, 22 Oct 2026 01:02:08 +0800 Subject: [PATCH] fix(triton): explicit strides, negative-index guards, and an honest status Three fixes found by running the ported kernels inside the engine rather than only against the A/B harness: `map_lookup` took the row stride from the row width, which is wrong for the views the engine passes (a slice of a wider buffer). Every kernel now takes its strides as arguments. `expand_sampling_params` used `row_to_req` directly as an offset; padding rows carry a negative request index, so that read out of bounds. It is part of the mask now, and `store_mtp_info` treats a non-positive slot as "no write" for the same reason. `prepare_verification_metadata` computed its grid from `seqlen_q.max().item()`, a device-to-host read that the engine rejects inside the forward ("called a synchronizing CUDA operation"). The grid now comes from the factory's `num_steps`. Integration is not finished, and `H20-TRITON-PLAN.md` says so plainly: with `DEEPINFER_OPS_BACKEND=triton` the engine runs the 15 ported kernels into the sampling phase and then faults on an illegal access inside `expand_sampling_params_tl`, even though that kernel is element-for-element identical to the baseline in its own A/B cases, and even though the negative index, the reserved slot and uint64 seeds have all been ruled out. The note records the two bisection steps worth trying next, the env-var tooling already in place, and the fact that one of my own edits silently did nothing because the search string carried a trailing space. Also records the two new Triton traps: an `other=` literal must match the destination dtype, and a grid must never be derived from a device read. --- H20-TRITON-PLAN.md | 87 ++++++++++++++----- .../model_executor/layers/tilelang_ops/ops.py | 5 +- .../layers/tilelang_ops/triton/sampler.py | 77 +++++++++++----- 3 files changed, 124 insertions(+), 45 deletions(-) diff --git a/H20-TRITON-PLAN.md b/H20-TRITON-PLAN.md index d28e8c7..18ec47a 100644 --- a/H20-TRITON-PLAN.md +++ b/H20-TRITON-PLAN.md @@ -104,29 +104,74 @@ ## 七、进度 -| 阶段 | 状态 | 说明 | -|---|---|---| -| P0 基础设施 | ✅ | 后端选择 + fallback + A/B 骨架 | -| P1 轻量 kernel | ✅ 13/13 | gather、gather_input_ids_zos、gather_engram_tokens_zos、expand_swa_index、kv_index_cat、convert_kv_index、load_kv_index、store_mtp_info、norm、map_lookup、expand_sampling_params、prepare_verification_metadata、copy_block_tables | -| P2 KV cache | ✅ 2/3 | store_kv_cache、load_paged_kv_cache（`fused_transform` 未做） | -| P3 MTP 链路 | ⬜ | | -| P4 采样 | ⬜ | | -| P5 端到端 | ⬜ | | +### 已完成的 15 个 kernel（A/B 62 个用例全过） -A/B 现状：**15 个 kernel、62 个用例全过**（`python scripts/h20_ops_ab.py`）。 +| kernel | 覆盖面 | +|---|---| +| `gather_tl` | int8/bf16 载荷、负索引零填充、strided 行、int32/int64 索引、空输入 | +| `gather_input_ids_zos_tl` | remap 重放 vs fresh，短于 max_seqlen | +| `gather_engram_tokens_zos_tl` | n-gram 窗口、image sentinel 的 pad 传播 | +| `expand_swa_index_tl` | 两段区间、len1+len2 溢出 width、非 2 幂宽度 | +| `kv_index_cat_tl` | 拼接 + `-1` 保持、非 2 幂 out_stride | +| `convert_kv_index_tl` | 逻辑→物理页表映射、超 context_len 置 -1 | +| `load_kv_index_tl` | prefill 连续 / decode 分页两种模式 | +| `store_mtp_info_tl` | has_index/has_confidence 四种组合、slot 0 不写 | +| `norm_tl` | bf16 出、residual、fp8（round/no-round）、fp8+bf16 双出、d=512/5120/7168 | +| `map_lookup_tl` | by_index（fp32/int32）、by_key（唯一键） | +| `expand_sampling_params_tl` | 5 个输出张量逐一比对 | +| `prepare_verification_metadata_tl` | row_to_req/local_step/is_verify | +| `copy_block_tables_tl` | 多 chunk（1400 列）、非 2 幂块数 | +| `store_kv_cache_tl` | main fp8（ue8m0 sf）+ index（fp32 sf）、`-1` 槽位、block=134 | +| `load_paged_kv_cache_tl` | do_cast 反量化、index 直通、`src_block_id == -1` 的洞 | -KV cache 的两个 kernel 都是**整块 buffer 逐字节一致**（store 覆盖 ue8m0 与 fp32 两种 scale -布局、`-1` 无效槽位；load 覆盖 do_cast 反量化、index 直通、以及 `src_block_id == -1` 的洞）。 +验证方式：`python scripts/h20_ops_ab.py`（可带 kernel 名子集）。 -`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 的 - 归约顺序，两个编译器不可能一致。 +### 未完成的 9 个 kernel -### 过程中踩到的 Triton 语义坑（都已修，可复用） +| kernel | 文件 | 行数 | 备注 | +|---|---|---|---| +| `fused_norm_rope_hadamard_cast_store_kvcache_tl` | `fused_transform.py` | 619 | 实测只用 norm+rope（`do_hadamard=False`、`cache_layout=None`、`return_bf16=True`），但函数体覆盖全部组合 | +| `compress_and_update_state_prefill_tl` | `compress.py` | ~150 | softmax 加权聚合 + 环形 state cache | +| `compress_and_update_state_decode_tl` | `compress.py` | ~105 | 同上，decode 形态 | +| `prepare_v4_decode_metadata_tl` | `prepare_decode_metadata.py` | ~130 | 单次调用写出十余个输出，按 ratio 展开 | +| `get_fused_dflash_get_input_tl` | `get_kv_index.py` | ~60 | MTP draft 输入 | +| `confidence_head_fused_tl` | `dflash.py` | — | MTP | +| `dflash_input_store_tl` | `dflash.py` | — | MTP | +| `get_accept_result_tl` / `finalize_verification_tl` / `remap_draft_topk_tl` | `sampler.py` | — | MTP 接受判定与 draft top-k 重映射 | +| `fused_sample_tl` | `fused_sample.py` | 524 | top-k + 温度 + 拒绝采样，最复杂的一个 | + +关掉 MTP 时只需前 3 个 + `prepare_v4_decode_metadata`；打开 MTP 需要全部 9 个。 + +### 端到端集成：还差一步，且已定位到一个真实缺陷 + +`DEEPINFER_OPS_BACKEND=triton` 能把 15 个 kernel 装进引擎并跑到采样阶段，但**会在 +`expand_sampling_params_tl` 处报 illegal memory access**，端到端尚未跑通。诚实记录现状： + +- 该 kernel 在自己的 A/B 用例里是精确一致的（5 个输出张量逐元素）。 +- 用 `CUDA_LAUNCH_BLOCKING=1` 把异步错误归位到它，说明是它自己触发的，不是别人延迟报的。 +- 已排除/已加固：`row_to_req` 为负时不再用作偏移（加 `req >= 0` 掩码）； + `store_mtp_info` 的 slot 改为 `> 0`；`seq_seeds` 由 uint64 改为 int64 视图（纯位拷贝， + 规避 Triton 的 uint64 指针映射）。三者都没解决。 +- 下一步（按信息量排序）： + 1. 在该 kernel 里逐输入打 shape/dtype/stride/min/max（`DEEPINFER_OPS_DEBUG_SHAPES=1` 的 + 钩子已接在 `map_lookup` 上，同样方式加到这里即可）。注意我上一次尝试因为替换字符串 + 带了尾随空格**静默没生效**，改动要用 assert 兜住。 + 2. 用 `DEEPINFER_OPS_FALLBACK=expand_sampling_params_tl` 把它摘掉，确认端到端能继续往下走， + 从而判断它是唯一阻塞点还是第一个。 + 3. 若确认是 Triton 对某个 dtype（例如 `torch.uint64` 视图、或 0 维张量）的指针处理， + 最省事的路径是让该 kernel 只做纯字节搬运（`view(torch.int8)` + 逐字节 copy）。 + +工具已就绪：`DEEPINFER_OPS_BACKEND` / `DEEPINFER_OPS_FALLBACK` / `DEEPINFER_OPS_VERBOSE` / +`DEEPINFER_OPS_DEBUG_SHAPES` 四个环境变量，加上 `scripts/h20_ops_ab.py`。 + +### 复用到的 Triton 语义坑（前五项见下，新增两项） + +6. **`tl.where` 之外还有一类"隐式求值"**：`other=` 的字面量类型必须和目标 dtype 匹配。 + fp8 源用 `other=0`（int）会报 "cannot cast int32[...] to fp8e4nv"，指向的行号是错的。 +7. **不要把设备张量读回主机来算 grid**：`seqlen_q.max().item()` 在 forward 里会触发 + `called a synchronizing CUDA operation`。用工厂参数（如 `num_steps`）当上界。 + +### 复用到的 Triton 语义坑（P1/P2 阶段） 1. **`tl.where` 两侧都会求值**。`tl.where(i == 0, 0, LOAD(i - 1))` 在 `i == 0` 时真的会读 `LOAD[-1]` → illegal memory access。凡是「边界索引用 where 兜住」的写法都要先 @@ -139,5 +184,5 @@ KV cache 的两个 kernel 都是**整块 buffer 逐字节一致**（store 覆盖 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`，否则两个后端一起写越界，表现成互相不一致。 +5. **不要假设行步长等于行宽**。引擎传进来的常是更大 buffer 的切片视图， + `row * width` 会读错行、并在末尾越界。所有 kernel 都改成显式传 stride。 diff --git a/deepinfer/model_executor/layers/tilelang_ops/ops.py b/deepinfer/model_executor/layers/tilelang_ops/ops.py index a9fb853..d289a73 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/ops.py +++ b/deepinfer/model_executor/layers/tilelang_ops/ops.py @@ -113,9 +113,12 @@ _INHERITED = sorted(set(_DEFAULT) - set(_SELECTED) - _FALLBACK) for _name in _INHERITED: globals()[_name] = _DEFAULT[_name] +# Kept short: with data parallelism this prints once per rank, and the full list is dozens of +# names that are not on the caller's path. `DEEPINFER_OPS_VERBOSE=1` prints them all. if _INHERITED: print( - f"[ops] backend {_BACKEND!r} lacks {len(_INHERITED)} kernel(s), using {_DEFAULT_BACKEND!r} for: " + ", ".join(_INHERITED), + f"[ops] backend {_BACKEND!r} lacks {len(_INHERITED)} kernel(s), falling back to " + f"{_DEFAULT_BACKEND!r} ({_INHERITED if os.getenv('DEEPINFER_OPS_VERBOSE') else 'first: ' + _INHERITED[0]})", file=sys.stderr, flush=True, ) diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py b/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py index 25680fe..d5856c1 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py @@ -4,6 +4,7 @@ These are grouped in one module because the TileLang backend also keeps them tog interface module for each one imports the factory from a single place. """ +import os from typing import Any import torch @@ -11,6 +12,8 @@ import triton import triton.language as tl from deepinfer.model_executor.layers.tilelang_ops.common import DRAFT_POS_HASH_MUL + +_DEBUG_SHAPES = bool(os.getenv("DEEPINFER_OPS_DEBUG_SHAPES")) from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2, triton_dtype @@ -40,8 +43,11 @@ def _prepare_verification_metadata_kernel( 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 + """Build the kernel. `num_steps` bounds the rows per request, which fixes the grid. + + The row count is deliberately *not* read back from `seqlen_q`: the engine forbids synchronising + CUDA calls inside the forward, and `tensor.max().item()` is exactly that. + """ threads = 32 def launch( @@ -52,13 +58,9 @@ def prepare_verification_metadata_tl(num_steps: int) -> Any: is_verify: torch.Tensor, ) -> None: req_cnt = ftp.shape[0] - if req_cnt == 0: + if req_cnt == 0 or num_steps <= 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))]( + _prepare_verification_metadata_kernel[(req_cnt, cdiv(num_steps, threads))]( ftp, seqlen_q, row_to_req, @@ -95,17 +97,20 @@ def _expand_sampling_params_kernel( req = tl.load(ROW_TO_REQ + rows, mask=live, other=0) step = tl.load(LOCAL_STEP + rows, mask=live, other=0) + # Padding rows carry a negative request index, so it has to be part of the mask rather than + # used as an offset: `BASE + (-1)` is an out-of-bounds read. + valid = live & (req >= 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) + tl.store(TEMPERATURES_Q + rows, tl.load(TEMPERATURES + req, mask=valid, other=0.0), mask=valid) + tl.store(TOP_PS_Q + rows, tl.load(TOP_PS + req, mask=valid, other=0.0), mask=valid) + tl.store(TOP_KS_Q + rows, tl.load(TOP_KS + req, mask=valid, other=0), mask=valid) # 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.load(SAMPLE_POSITIONS + req, mask=valid, other=0) + step * pos_hash_mul, + mask=valid, ) - tl.store(SEQ_SEEDS_Q + rows, tl.load(SEQ_SEEDS + req, mask=live, other=0), mask=live) + tl.store(SEQ_SEEDS_Q + rows, tl.load(SEQ_SEEDS + req, mask=valid, other=0), mask=valid) def expand_sampling_params_tl() -> Any: @@ -134,14 +139,16 @@ def expand_sampling_params_tl() -> Any: top_ps, top_ks, sample_positions, - seq_seeds, + # Seeds are a bit-for-bit passthrough, so they are handed over as int64: uint64 has no + # reliable pointer mapping in Triton, and the value never participates in arithmetic. + seq_seeds.view(torch.int64), row_to_req, local_step, temperatures_q, top_ps_q, top_ks_q, sample_positions_q, - seq_seeds_q, + seq_seeds_q.view(torch.int64), total_q, pos_hash_mul=DRAFT_POS_HASH_MUL, THREADS=threads, @@ -157,6 +164,9 @@ def _map_lookup_index_kernel( QUERY, OUT, total_q, + value_stride, + query_stride, + out_stride, width: tl.constexpr, is_float: tl.constexpr, BLOCK: tl.constexpr, @@ -165,9 +175,12 @@ def _map_lookup_index_kernel( rows = block * BLOCK + tl.arange(0, BLOCK) live = rows < total_q - query = tl.load(QUERY + rows, mask=live, other=0) + # Row strides are passed rather than assumed: the engine hands in views whose leading stride is + # larger than the row width (a slice of a wider buffer), and assuming `width` read the wrong + # rows -- an out-of-bounds access at the end of the buffer. + query = tl.load(QUERY + rows * query_stride, mask=live, other=0) in_range = live & (query >= 0) & (query < width) - value = tl.load(VALUE + rows * width + query, mask=in_range, other=0) + value = tl.load(VALUE + rows * value_stride + 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 @@ -176,7 +189,7 @@ def _map_lookup_index_kernel( result = tl.where(in_range, value, float("-inf")) else: result = tl.where(in_range, value, -1) - tl.store(OUT + rows, result, mask=live) + tl.store(OUT + rows * out_stride, result, mask=live) @triton.jit @@ -186,19 +199,23 @@ def _map_lookup_key_kernel( QUERY, OUT, total_q, + key_stride, + value_stride, + query_stride, + out_stride, width: tl.constexpr, BLOCK: tl.constexpr, ): row = tl.program_id(0) if row < total_q: - query = tl.load(QUERY + row) + query = tl.load(QUERY + row * query_stride) 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) + key = tl.load(KEY + row * key_stride + idx, mask=live, other=-1) + value = tl.load(VALUE + row * value_stride + 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 @@ -207,7 +224,7 @@ def _map_lookup_key_kernel( 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) + tl.store(OUT + row * out_stride, best) def map_lookup_tl( @@ -232,12 +249,22 @@ def map_lookup_tl( total_q = query.shape[0] if total_q == 0: return + if _DEBUG_SHAPES: + print( + f"[map_lookup] mode={mode} force_topk={force_topk} value={tuple(value.shape)} {value.dtype} " + f"stride={value.stride()} key={None if key is None else (tuple(key.shape), key.stride())} " + f"query={tuple(query.shape)} {query.dtype} stride={query.stride()} out={tuple(out.shape)} stride={out.stride()}", + flush=True, + ) if mode == "by_index": _map_lookup_index_kernel[(cdiv(total_q, block_index),)]( value, query, out, total_q, + value.stride(0), + query.stride(0), + out.stride(0), width=force_topk, is_float=bool(value.dtype.is_floating_point), BLOCK=block_index, @@ -250,6 +277,10 @@ def map_lookup_tl( query, out, total_q, + key.stride(0), + value.stride(0), + query.stride(0), + out.stride(0), width=force_topk, BLOCK=block_key, num_warps=8,