mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
f4144fc3fc
- vendor/deep_ep: EPBuffer dispatch/combine with torch.distributed uneven all_to_all (host-synced split sizes), EngramBuffer sharded lookup, BucketBuffer all_gather (CP-only reduce paths stubbed); verified against the DeepEP-Ascend references at world size 1/2/4. - tilelang_ops/neuron sampling: random_sampling (asc two-level fixed-point sampler, Philox), fused_sample (logprobs / do_sample), sampler metadata, sum_smaller_probs, structured-output bitmask; DSpark/MTP/varlen kernels raise NotImplementedError. All 42 asc *_tl factories now have a Neuron counterpart. - vendor/flash_attn: SDPA varlen stand-in for the ViT (mirrors npu_ops). - transfer_to_neuron: keep `with torch.device(...)` working on patched factories; fix manual_seed recursion in host mode. - engram: gather the packed ue8m0 sf as bytes on Neuron. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
123 lines
4.6 KiB
Python
123 lines
4.6 KiB
Python
"""Mirrors ``tilelang_ops/asc/sampler.py`` in pure torch.
|
|
|
|
Implemented (used by plain sampling, including nextn=0 decode): prepare_verification_metadata_tl,
|
|
expand_sampling_params_tl, map_lookup_tl.
|
|
|
|
MTP/DSpark verification only, out of scope on Neuron basic mode (NEURON_PORT.md "Scope"; the factories keep
|
|
the asc signatures and raise): remap_draft_topk_tl, get_accept_result_tl, finalize_verification_tl.
|
|
"""
|
|
|
|
from typing import Any, Literal
|
|
|
|
import torch
|
|
|
|
from ..common import DRAFT_POS_HASH_MUL
|
|
from ._unsupported import unsupported
|
|
|
|
|
|
def prepare_verification_metadata_tl(num_steps: int) -> Any:
|
|
"""kernel(ftp, seqlen_q, row_to_req, local_step, is_verify); rows are packed per request from ftp."""
|
|
|
|
def _prepare_verification_metadata(
|
|
ftp: torch.Tensor,
|
|
seqlen_q: torch.Tensor,
|
|
row_to_req: torch.Tensor,
|
|
local_step: torch.Tensor,
|
|
is_verify: torch.Tensor,
|
|
) -> None:
|
|
total_q = row_to_req.shape[0]
|
|
req_cnt = ftp.shape[0]
|
|
device = ftp.device
|
|
# row -> request without searchsorted/host sync: mark each request start, then prefix-count.
|
|
starts = torch.zeros(total_q + 1, dtype=torch.int32, device=device)
|
|
starts.scatter_add_(0, ftp.to(torch.int64).clamp(0, total_q), torch.ones(req_cnt, dtype=torch.int32, device=device))
|
|
req = starts[:total_q].cumsum(0, dtype=torch.int32) - 1
|
|
req_idx = req.clamp(min=0).to(torch.int64)
|
|
step = torch.arange(total_q, dtype=torch.int32, device=device) - ftp.to(torch.int32)[req_idx]
|
|
row_to_req.copy_(req)
|
|
local_step.copy_(step)
|
|
is_verify.copy_(step + 1 < seqlen_q.to(torch.int32)[req_idx])
|
|
|
|
return _prepare_verification_metadata
|
|
|
|
|
|
def expand_sampling_params_tl() -> Any:
|
|
"""kernel(temperatures, top_ps, top_ks, sample_positions, seq_seeds, row_to_req, local_step,
|
|
temperatures_q, top_ps_q, top_ks_q, sample_positions_q, seq_seeds_q)."""
|
|
|
|
def _expand_sampling_params(
|
|
temperatures: torch.Tensor,
|
|
top_ps: torch.Tensor,
|
|
top_ks: torch.Tensor,
|
|
sample_positions: torch.Tensor,
|
|
seq_seeds: torch.Tensor,
|
|
row_to_req: torch.Tensor,
|
|
local_step: torch.Tensor,
|
|
temperatures_q: torch.Tensor,
|
|
top_ps_q: torch.Tensor,
|
|
top_ks_q: torch.Tensor,
|
|
sample_positions_q: torch.Tensor,
|
|
seq_seeds_q: torch.Tensor,
|
|
) -> None:
|
|
req = row_to_req.to(torch.int64)
|
|
temperatures_q.copy_(temperatures[req])
|
|
top_ps_q.copy_(top_ps[req])
|
|
top_ks_q.copy_(top_ks[req])
|
|
# int32 arithmetic with wraparound, as in the kernel.
|
|
sample_positions_q.copy_(sample_positions.to(torch.int32)[req] + local_step.to(torch.int32) * DRAFT_POS_HASH_MUL)
|
|
seq_seeds_q.view(torch.int64).copy_(seq_seeds.view(torch.int64)[req])
|
|
|
|
return _expand_sampling_params
|
|
|
|
|
|
def map_lookup_tl(
|
|
force_topk: int,
|
|
mode: Literal["by_key", "by_index"] = "by_key",
|
|
value_dtype: torch.dtype = torch.float32,
|
|
) -> Any:
|
|
"""kernel(key, value, query, out) with out/query [rows, 1].
|
|
|
|
by_index: ``value[row, query]`` or the out-of-bounds value (-inf float / -1 int).
|
|
by_key: the value whose key equals the query (keys are unique top-k indices; with duplicate keys this
|
|
returns their max), -inf when missing."""
|
|
assert mode in ("by_key", "by_index")
|
|
assert mode != "by_key" or value_dtype.is_floating_point, "by_key requires floating-point values"
|
|
oob_value = float("-inf") if value_dtype.is_floating_point else -1
|
|
|
|
def _map_lookup(
|
|
key: torch.Tensor | None,
|
|
value: torch.Tensor,
|
|
query: torch.Tensor,
|
|
out: torch.Tensor,
|
|
) -> None:
|
|
q = query[:, 0]
|
|
if mode == "by_index":
|
|
q64 = q.to(torch.int64)
|
|
in_bounds = (q64 >= 0) & (q64 < force_topk)
|
|
got = value[:, :force_topk].gather(1, torch.where(in_bounds, q64, torch.zeros_like(q64)).unsqueeze(1)).squeeze(1)
|
|
result = torch.where(in_bounds, got, torch.full_like(got, oob_value))
|
|
else:
|
|
match = key[:, :force_topk] == q.unsqueeze(1)
|
|
v = value[:, :force_topk]
|
|
result = torch.where(match, v, torch.full_like(v, oob_value)).amax(dim=1)
|
|
out[:, 0].copy_(result)
|
|
|
|
return _map_lookup
|
|
|
|
|
|
def remap_draft_topk_tl(force_topk: int, vocab_size: int) -> Any:
|
|
unsupported("remap_draft_topk_tl")
|
|
|
|
|
|
def get_accept_result_tl(has_golden: bool) -> Any:
|
|
unsupported("get_accept_result_tl")
|
|
|
|
|
|
def finalize_verification_tl(
|
|
has_sum_smaller_probs: bool,
|
|
has_confidence: bool,
|
|
num_steps: int,
|
|
num_stop_tokens: int,
|
|
) -> Any:
|
|
unsupported("finalize_verification_tl")
|