coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-31 00:36:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(triton): the grid for prepare_verification_metadata, and first matching end-to-end run

Browse Source 
The illegal memory access was not in `expand_sampling_params` at all, though that is where
`CUDA_LAUNCH_BLOCKING` pointed. It was the metadata kernel leaving its output unwritten.

`prepare_verification_metadata_tl` sized its second grid dimension from `num_steps`, i.e.
`max_MTP_value`. With speculative decoding off that value is 0, and a zero-sized grid dimension
launches nothing, so `row_to_req` -- a `torch.empty` buffer -- kept whatever the allocator handed
out. `expand_sampling_params` then used those values as request indices, and the reads went off
the end of the parameter arrays.

The kernel now launches one program per request and loops the tiles at runtime, mirroring the
TileLang version, so the row count comes from device memory and no host read is needed (which
would have been rejected as a synchronising operation anyway). `num_steps` is accepted and
ignored.

The A/B case for this kernel is also fixed, in two ways that matter:

- it used a fixed `total_q` while the request lengths were random, so it allocated more rows than
  were ever written and the comparison was meaningless;
- it pre-filled the outputs with a sentinel but never required them to be overwritten, so both
  backends could write nothing and still agree with each other. The case now asserts every row was
  written, and covers `num_steps=0` plus a request longer than one tile.

With that, `DEEPINFER_OPS_BACKEND=triton` runs end to end and produces exactly the same tokens as
the pure TileLang run, `[17607, 2793, 3139, 666, 6127, 666, 320, 1]`, identical on all eight
ranks. Fifteen of the required kernels are Triton; the remainder still come from TileLang through
the fallback while they are ported.
This commit is contained in:
loopyt 2026-10-22 07:26:09 +08:00
parent d393584592
commit df9559e965
2 changed files with 62 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py
View File

@ -12,9 +12,9 @@ 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
_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
@triton.jit
@ -27,27 +27,38 @@ def _prepare_verification_metadata_kernel(
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
# One program per request with a runtime trip count, mirroring the TileLang kernel: the tile
# count depends on this request's length, which is only known on device.
#
# The obvious alternative -- a second grid dimension sized from a `num_steps` argument -- is
# wrong twice over. `num_steps` is `max_MTP_value`, which is 0 when speculative decoding is
# off, and a zero-sized dimension launches nothing at all; and when it is non-zero it is only
# an upper bound, so `total_q` could still exceed what the grid covers. Either way the
# untouched rows keep whatever `torch.empty` handed out, and every downstream consumer of
# `row_to_req` then indexes with that garbage.
for tile in range(0, tl.cdiv(length, threads)):
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)
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. `num_steps` bounds the rows per request, which fixes the grid.
"""Build the kernel. `num_steps` is accepted for interface parity and not used.
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.
Deriving the grid from it is a bug (see the kernel comment); the row counts come from device
memory, so nothing here needs to know them. Similarly the length is never read back to the
host: `tensor.max().item()` inside the forward is a synchronising operation, which the engine
rejects outright.
"""
del num_steps
threads = 32
def launch(
@ -58,9 +69,9 @@ def prepare_verification_metadata_tl(num_steps: int) -> Any:
is_verify: torch.Tensor,
) -> None:
req_cnt = ftp.shape[0]
if req_cnt == 0 or num_steps <= 0:
if req_cnt == 0:
return
_prepare_verification_metadata_kernel[(req_cnt, cdiv(num_steps, threads))](
_prepare_verification_metadata_kernel[(req_cnt,)](
ftp,
seqlen_q,
row_to_req,
@ -134,6 +145,20 @@ def expand_sampling_params_tl() -> Any:
total_q = row_to_req.shape[0]
if total_q == 0:
return
if _DEBUG_SHAPES:
for nm, t in (
("temperatures", temperatures),
("top_ks", top_ks),
("sample_positions", sample_positions),
("seq_seeds", seq_seeds),
("row_to_req", row_to_req),
("local_step", local_step),
("seq_seeds_q", seq_seeds_q),
):
print(
f"[expand] {nm}: shape={tuple(t.shape)} dtype={t.dtype} contig={t.is_contiguous()} stride={t.stride()}",
flush=True,
)
_expand_sampling_params_kernel[(cdiv(total_q, threads),)](
temperatures,
top_ps,

26
scripts/ab_cases.py
View File

@ -404,24 +404,42 @@ def _cases_expand_sampling_params() -> list[tuple[str, bool, Callable[[], Any]]]
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 make(req_cnt: int, max_len: int, build_steps: int):
def run() -> Any:
torch.manual_seed(0)
seqlen_q = torch.randint(1, max_len + 1, (req_cnt,), device="cuda", dtype=torch.int32)
cu = torch.cumsum(seqlen_q, 0, dtype=torch.int32)
total_q = int(cu[-1].item())
ftp = torch.zeros((req_cnt,), device="cuda", dtype=torch.int32)
ftp[1:] = torch.cumsum(seqlen_q, 0, dtype=torch.int32)[:-1]
ftp[1:] = cu[:-1]
# The sentinel makes "wrote nothing" visible. `row_to_req` is a `torch.empty` buffer in
# the engine, so a kernel that fills fewer than `total_q` rows leaves garbage behind and
# every consumer then indexes with it -- and both backends would still agree with each
# other here, which is why this check has to be explicit rather than a comparison.
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 = ops.prepare_verification_metadata_tl(build_steps)
kernel(ftp, seqlen_q, row_to_req, local_step, is_verify)
assert int((row_to_req == -7).sum()) == 0, "kernel left rows unwritten"
assert int((local_step == -7).sum()) == 0, "kernel left rows unwritten"
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))]
return [
("n4_maxlen6", True, make(4, 6, 6)),
("n1_maxlen1", True, make(1, 1, 1)),
("n8_maxlen6", True, make(8, 6, 6)),
# `num_steps=0` is what the engine builds when speculative decoding is off, so it must
# still write every row.
("n8_zero_steps", True, make(8, 6, 0)),
# A request longer than one tile (32 lanes) exercises the runtime tile loop.
("n2_maxlen130", True, make(2, 130, 130)),
("n64_maxlen6", True, make(64, 6, 6)),
]
def _cases_copy_block_tables() -> list[tuple[str, bool, Callable[[], Any]]]: