mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
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.
This commit is contained in:
loopyt
parent
29fed050cb
commit
d393584592
@ -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。
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
@ -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,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user