diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py b/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py index d5856c1..f0a2cfd 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py +++ b/deepinfer/model_executor/layers/tilelang_ops/triton/sampler.py @@ -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, diff --git a/scripts/ab_cases.py b/scripts/ab_cases.py index 8f31c05..becda0d 100644 --- a/scripts/ab_cases.py +++ b/scripts/ab_cases.py @@ -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]]]: