diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 63da639..956f971 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -144,6 +144,8 @@ class envs: DEEPINFER_NEURON_NUM_KV_BLOCKS = EnvVar.int(default=0) # Neuron only: run the sampler's process_sampling as one compiled graph instead of ~1900 eager ops a step. DEEPINFER_NEURON_COMPILE_SAMPLER = EnvVar.bool(default=True) + # The decode verification (greedy batch, no structured output) as one compiled region (issue #29). + DEEPINFER_NEURON_COMPILE_VERIFY = EnvVar.bool(default=False) # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break). DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool() diff --git a/deepinfer/model_executor/layers/sampler.py b/deepinfer/model_executor/layers/sampler.py index 384a588..cc89948 100644 --- a/deepinfer/model_executor/layers/sampler.py +++ b/deepinfer/model_executor/layers/sampler.py @@ -5,13 +5,14 @@ from __future__ import annotations +import functools from typing import TYPE_CHECKING import torch import torch.nn as nn from deep_select.interface import topk -from deepinfer import envs +from deepinfer import PLATFORM, envs from deepinfer.logger import init_logger from deepinfer.model_executor.input_metadata import InputMetadata from deepinfer.model_executor.layers.mtp_info_pool import MTPInfoPool @@ -56,6 +57,30 @@ if TYPE_CHECKING: logger = init_logger(__name__) +_COMPILED_VERIFY: list = [] + + +def _compiled_verification(): # noqa: ANN202 + """``Sampler._verification`` (undecorated) as a compiled Neuron region: eagerly it is ~300 small ops a decode step, + each its own NEFF launch with ~0.5 ms of device idle around it (issue #29).""" + if not _COMPILED_VERIFY: + from deepinfer.neuron_ops import neuron_compile # noqa: F401 (registers the "deepinfer_neuron" backend) + from deepinfer.neuron_ops import regions + from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda + + regions.raise_dynamo_limits() + compiled = torch.compile( + Sampler._verification.__wrapped__, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, fullgraph=False, dynamic=False + ) + + def run(*args, **kwargs): # noqa: ANN002, ANN003, ANN202 + with cpu_timer.record("verification"), hide_cuda(): + return compiled(*args, **kwargs) + + _COMPILED_VERIFY.append(run) + return _COMPILED_VERIFY[0] + + class Sampler(nn.Module): """从模型 hidden states 采样下一个 token，支持 MTP（Multi-Token Prediction）投机采样。 @@ -212,6 +237,16 @@ class Sampler(nn.Module): else: # Decode sampling. speculative_inputs = driver_inputs.speculative + verify, verify_driver, verify_decode = self._verification, driver_inputs, decode_inputs + if PLATFORM.is_neuron and envs.DEEPINFER_NEURON_COMPILE_VERIFY and so_indices is None and not self.force_topk: + from deepinfer.model_executor.layers.tilelang_ops.neuron.fused_sample import prime_host_flags + from deepinfer.neuron_ops.sampling_mirror import mirror_decode, mirror_driver + + params = driver_inputs.kernel.logits + _, all_greedy = prime_host_flags(params.temperatures, params.top_ps, params.top_ks, logits.shape[1]) + if all_greedy: + verify = functools.partial(_compiled_verification(), self) + verify_driver, verify_decode = mirror_driver(driver_inputs, greedy=True), mirror_decode(decode_inputs) logits = Movable(logits) ( sampled_tokens, @@ -222,10 +257,10 @@ class Sampler(nn.Module): num_verified_tokens, batch_confidence_error, emitted_stop_token, - ) = self._verification( + ) = verify( logits, - driver_inputs, - decode_inputs, + verify_driver, + verify_decode, so_indices, so_bitmasks, req_cnt, diff --git a/deepinfer/model_executor/layers/tilelang_ops/fused_sample.py b/deepinfer/model_executor/layers/tilelang_ops/fused_sample.py index 338eee9..2ed52ba 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/fused_sample.py +++ b/deepinfer/model_executor/layers/tilelang_ops/fused_sample.py @@ -63,7 +63,7 @@ def fused_sample( ``max(target_prob_i - exp(draft_logprob_i - DRAFT_PROB_PENALTY), 0)``. """ assert logits.stride(1) == 1 - if isinstance(sampling_inputs, SamplingKernelInputs): + if isinstance(sampling_inputs, SamplingKernelInputs) or hasattr(sampling_inputs, "seq_seeds"): # or its mirror do_sample = True transform_inputs = sampling_inputs.logits seq_seeds = sampling_inputs.seq_seeds @@ -187,7 +187,7 @@ def fused_sample_ref( The reference may differ numerically from the TileLang implementation. Its in-place update of ``logits`` is safe for the caller in sampler_utils.py. """ - if isinstance(sampling_inputs, SamplingKernelInputs): + if isinstance(sampling_inputs, SamplingKernelInputs) or hasattr(sampling_inputs, "seq_seeds"): # or its mirror do_sample = True transform_inputs = sampling_inputs.logits seq_seeds = sampling_inputs.seq_seeds diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py index 55d0bf9..4ad6815 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py @@ -147,7 +147,7 @@ def _host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]: # # By shape, not address: a step may hand over the same parameters in a fresh buffer. key = ("filter", tuple(t.shape[0] for t in (temperatures, top_p, top_k)), vocab) - flags = _FLAGS.get(key) + flags = _FLAGS.get(key, _FLAGS.get("batch")) if flags is None: t_h, p_h, k_h = temperatures.cpu(), top_p.cpu(), top_k.cpu() flags = (bool(((p_h < 1.0) | ((k_h >= 0) & (k_h < vocab))).any()), bool((t_h <= SAMPLING_EPS).all())) @@ -155,6 +155,14 @@ def _host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]: # return flags +def prime_host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]: # noqa: ANN001 + """Read the batch's flags eagerly (once per batch) so a compiled region finds them (``_host_flags`` falls back to + them for its expanded per-row parameters: expansion changes neither any-filter nor all-greedy).""" + if "batch" not in _FLAGS: + _FLAGS["batch"] = _host_flags(temperatures, top_p, top_k, vocab) + return _FLAGS["batch"] + + def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size: int): # noqa: ANN001, ANN202 """On a NeuronCore the log-softmax / top-p / top-k part is one compiled graph (issue #27): eagerly the sampler is ~1900 small ops a decode step, each its own NEFF launch (~0.6 s). @@ -166,7 +174,15 @@ def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_si global _COMPILED from deepinfer import envs - if logits.device.type != "neuron" or torch.compiler.is_compiling() or not envs.DEEPINFER_NEURON_COMPILE_SAMPLER: + if torch.compiler.is_compiling() and logits.device.type == "neuron": + # Inside the compiled verification (DEEPINFER_NEURON_COMPILE_VERIFY): the same choice as below, traced. + filtering, all_greedy = _host_flags(temperatures, top_p, top_k, logits.shape[1]) + fn = _logprobs_and_greedy if filtering else _plain_logprobs_and_greedy + logprobs, greedy = fn(logits, temperatures, top_p, top_k) if filtering else fn(logits, temperatures) + if all_greedy: + return logprobs, greedy + return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size) + if logits.device.type != "neuron" or not envs.DEEPINFER_NEURON_COMPILE_SAMPLER: return _sample_eager(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size) from deepinfer.neuron_ops import neuron_compile # noqa: F401 (registers the "deepinfer_neuron" backend) from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda @@ -200,9 +216,10 @@ def _greedy_verify( from deepinfer.sampling_params import SAMPLING_EPS key = ("greedy", temperatures.shape[0]) - if key not in _FLAGS: - _FLAGS[key] = bool((temperatures.cpu() <= SAMPLING_EPS).all()) - if not _FLAGS[key]: + greedy = _FLAGS.get(key, _FLAGS["batch"][1] if "batch" in _FLAGS else None) + if greedy is None: + greedy = _FLAGS[key] = bool((temperatures.cpu() <= SAMPLING_EPS).all()) + if not greedy: unsupported("fused_sample_tl(do_reject_sampling=True) for non-greedy rows") rows = read_rows.long() draft = draft_ids.index_select(0, rows) @@ -244,8 +261,11 @@ def fused_sample_tl( if not do_sample: logprobs.copy_(compute_logprobs(logits, temperatures, top_p, top_k)) return - seed = seq_seeds.view(torch.int64) if seq_seeds.dtype == torch.uint64 else seq_seeds.to(torch.int64) - seed = seed + int(seed_offset) # uint64 wraparound == int64 wraparound on the same bits + if seq_seeds is None: # a greedy batch's mirror (neuron_ops/sampling_mirror.py): no draw + seed = None + else: + seed = seq_seeds.view(torch.int64) if seq_seeds.dtype == torch.uint64 else seq_seeds.to(torch.int64) + seed = seed + int(seed_offset) # uint64 wraparound == int64 wraparound on the same bits out_logprobs, tokens = _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size) logprobs.copy_(out_logprobs) candidate_token_ids.copy_(tokens) diff --git a/deepinfer/model_executor/layers/tilelang_ops/sampler.py b/deepinfer/model_executor/layers/tilelang_ops/sampler.py index 8aadd37..db90ddd 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/sampler.py +++ b/deepinfer/model_executor/layers/tilelang_ops/sampler.py @@ -86,6 +86,18 @@ def expand_sampling_params( SamplingKernelInputs with [total_q] tensors, preserving each field's dtype. """ total_q = row_to_req.shape[0] + if PLATFORM.is_neuron and not isinstance(params, SamplingKernelInputs): + # The compiled Neuron verification's mirror (neuron_ops/sampling_mirror.py): plain gathers, and no seeds or + # positions for a greedy batch. + from deepinfer.neuron_ops.sampling_mirror import KernelMirror, LogitsMirror + + req = row_to_req.to(torch.int64) + lg = params.logits + positions = None + if params.sample_positions is not None: + positions = params.sample_positions.to(torch.int32)[req] + local_step.to(torch.int32) * DRAFT_POS_HASH_MUL + seeds = None if params.seq_seeds is None else params.seq_seeds.view(torch.int64)[req].view(torch.uint64) + return KernelMirror(LogitsMirror(lg.temperatures[req], lg.top_ps[req], lg.top_ks[req]), seeds, positions) temperatures_q = torch.empty(total_q, dtype=torch.float32, device=params.logits.temperatures.device) top_ps_q = torch.empty_like(temperatures_q) diff --git a/deepinfer/neuron_ops/sampling_mirror.py b/deepinfer/neuron_ops/sampling_mirror.py new file mode 100644 index 0000000..96be637 --- /dev/null +++ b/deepinfer/neuron_ops/sampling_mirror.py @@ -0,0 +1,74 @@ +"""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)