mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
0a7d490cdb
Sampler._verification runs as one compiled region for greedy batches without structured output: the native sampling structs enter as Python mirrors (neuron_ops/sampling_mirror.py), the flags are read once per batch eagerly (prime_host_flags), and greedy mirrors carry no seeds / positions. Untested on device yet. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
"""Python mirrors of the native sampling-input structs, for the compiled Neuron verification (issue #29).
|
|
|
|
Dynamo cannot read the nanobind structs' fields (each access breaks the graph) nor construct them, so the driver's
|
|
decode inputs enter the compiled ``Sampler._verification`` as these dataclasses: same field names, plain tensors.
|
|
A greedy batch carries no seeds or sample positions (``None``): only the RNG draw reads them, and the expansion of
|
|
the uint64 seeds is int64 arithmetic neuronx-cc graphs reject.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
|
|
|
|
@dataclass
|
|
class LogitsMirror:
|
|
temperatures: torch.Tensor
|
|
top_ps: torch.Tensor
|
|
top_ks: torch.Tensor
|
|
|
|
|
|
@dataclass
|
|
class KernelMirror:
|
|
logits: LogitsMirror
|
|
seq_seeds: torch.Tensor | None
|
|
sample_positions: torch.Tensor | None
|
|
|
|
|
|
@dataclass
|
|
class SpeculativeMirror:
|
|
mtp_info_read_rows: torch.Tensor
|
|
mtp_info_write_slots: torch.Tensor
|
|
nextn_tokens: torch.Tensor
|
|
stop_token_ids: torch.Tensor
|
|
|
|
|
|
@dataclass
|
|
class DriverMirror:
|
|
kernel: KernelMirror
|
|
speculative: SpeculativeMirror | None
|
|
golden_output: torch.Tensor | None
|
|
max_top_logprobs: int
|
|
|
|
|
|
@dataclass
|
|
class DecodeMirror:
|
|
mtp_first_token_positions: torch.Tensor
|
|
mtp_seqlens_q: torch.Tensor
|
|
|
|
|
|
def mirror_kernel(kernel, greedy: bool) -> KernelMirror: # noqa: ANN001
|
|
logits = kernel.logits
|
|
return KernelMirror(
|
|
LogitsMirror(logits.temperatures, logits.top_ps, logits.top_ks),
|
|
None if greedy else kernel.seq_seeds,
|
|
None if greedy else kernel.sample_positions,
|
|
)
|
|
|
|
|
|
def mirror_driver(driver, greedy: bool) -> DriverMirror: # noqa: ANN001
|
|
spec = driver.speculative
|
|
return DriverMirror(
|
|
kernel=mirror_kernel(driver.kernel, greedy),
|
|
speculative=None
|
|
if spec is None
|
|
else SpeculativeMirror(spec.mtp_info_read_rows, spec.mtp_info_write_slots, spec.nextn_tokens, spec.stop_token_ids),
|
|
golden_output=driver.golden_output,
|
|
max_top_logprobs=driver.max_top_logprobs,
|
|
)
|
|
|
|
|
|
def mirror_decode(decode) -> DecodeMirror: # noqa: ANN001
|
|
return DecodeMirror(decode.mtp_first_token_positions, decode.mtp_seqlens_q)
|