From f07d0a7c6f64fe6b938d27d6ede6049902cbab55 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 09:42:08 +0800 Subject: [PATCH] Neuron: compiled log-softmax + greedy argmax in the sampler; eager op census (#27) The op census of a decode step (DEEPINFER_NEURON_REGION_TIMING) found 85 eager aten ops in the forward outside the region graphs but 1905 in sampling (the top-p/top-k bisection and the int64 philox draw), ~0.6 s of NEFF launches. - The log-softmax / top-p / top-k part runs as one compiled graph: 1.4 ms instead of 46.7 ms on one core for the [1, 4096] force_topk row, logprobs within 1e-9. - Greedy rows (temperature clamped to SAMPLING_EPS) have one-hot weights, so the two-level draw returns their argmax; an all-greedy batch takes the argmax from the graph and skips the draw. The philox draw stays eager for sampled rows: neuronx-cc graphs reject 64-bit constants outside the int32 range (NCC_ESFH001), an int32 philox is #3. Co-Authored-By: Claude Opus 5.5 --- deepinfer/environ.py | 2 + .../tilelang_ops/neuron/fused_sample.py | 55 +++++++++++++++++-- deepinfer/neuron_ops/op_census.py | 35 ++++++++++++ deepinfer/worker/model_runner.py | 4 ++ deepinfer/worker/neuron_graph_runner.py | 5 +- 5 files changed, 94 insertions(+), 7 deletions(-) create mode 100644 deepinfer/neuron_ops/op_census.py diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 9e98be6..e268f44 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -122,6 +122,8 @@ class envs: # (issue #26, the vllm-neuron#71 layout). 0 = eager prefill. Needs max_num_batched_tokens <= bucket and # DEEPINFER_MAX_BACKTRACKING_TOKENS >= bucket (the SWA / compressor rings must hold a whole chunk). DEEPINFER_NEURON_PREFILL_BUCKET = 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) # 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/tilelang_ops/neuron/fused_sample.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py index ea4e01a..a8c6a87 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py @@ -105,6 +105,47 @@ def compute_logprobs( return torch.where(filter_active.unsqueeze(1), filtered, plain) +def _sample_eager(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size: int): # noqa: ANN001, ANN202 + logprobs = compute_logprobs(logits, temperatures, top_p, top_k) + # Exponentiate the logprobs so the weights equal exp(logprobs_out) bit for bit. + return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size) + + +def _logprobs_and_greedy(logits, temperatures, top_p, top_k): # noqa: ANN001, ANN202 + logprobs = compute_logprobs(logits, temperatures, top_p, top_k) + return logprobs, logprobs.argmax(dim=1).to(torch.int32) + + +_COMPILED = None + + +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). + + Greedy rows (temperature clamped to SAMPLING_EPS) have one-hot weights, so the two-level draw returns their argmax; a batch of only + greedy rows takes the argmax and skips the draw. The philox draw stays eager: it is int64 arithmetic, and + neuronx-cc graphs reject 64-bit constants outside the int32 range (NCC_ESFH001). + """ + global _COMPILED + from deepinfer import envs + + if logits.device.type != "neuron" or torch.compiler.is_compiling() 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 + + if _COMPILED is None: + _COMPILED = torch.compile(_logprobs_and_greedy, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, dynamic=False) + with hide_cuda(): + logprobs, greedy = _COMPILED(logits, temperatures, top_p, top_k) + from deepinfer.sampling_params import SAMPLING_EPS + + if bool((temperatures.cpu() <= SAMPLING_EPS).all()): # host sync on the [rows] temperatures + return logprobs, greedy + return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size) + + def fused_sample_tl( vocab_size: int, do_sample: bool, @@ -135,11 +176,13 @@ def fused_sample_tl( lse_workspace: torch.Tensor | None, ) -> None: assert logits.shape[1] == vocab_size and logits.stride(1) == 1 - logprobs.copy_(compute_logprobs(logits, temperatures, top_p, top_k)) - if do_sample: - 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 - # Exponentiate the written logprobs so the weights equal exp(logprobs_out) bit for bit. - candidate_token_ids.copy_(two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size)) + 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 + out_logprobs, tokens = _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size) + logprobs.copy_(out_logprobs) + candidate_token_ids.copy_(tokens) return _fused_sample diff --git a/deepinfer/neuron_ops/op_census.py b/deepinfer/neuron_ops/op_census.py new file mode 100644 index 0000000..8852ffe --- /dev/null +++ b/deepinfer/neuron_ops/op_census.py @@ -0,0 +1,35 @@ +"""Count the eager ops of one step on Neuron (issue #27): every aten op outside a compiled region graph is a +separate NEFF launch, so the census shows what still runs eagerly between and around the region graphs.""" + +from __future__ import annotations + +import contextlib +from collections.abc import Iterator + +import torch + +from deepinfer.logger import init_logger + +logger = init_logger(__name__) + +#: Bookkeeping ops that launch nothing on the device. +_FREE = ("aten::view", "aten::_reshape_alias", "aten::as_strided", "aten::alias", "aten::detach", "aten::lift_fresh", + "aten::t", "aten::unsqueeze", "aten::squeeze", "aten::expand", "aten::permute", "aten::transpose", + "aten::slice", "aten::select", "aten::reshape", "aten::empty", "aten::empty_strided", "aten::resolve_conj", + "aten::resolve_neg", "aten::_unsafe_view", "aten::split", "aten::chunk", "aten::unbind", "aten::narrow") + + +@contextlib.contextmanager +def census(label: str, top: int = 40) -> Iterator[None]: + """Log the ops run inside the block, by count (the outermost aten op of each dispatch).""" + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CPU]) as prof: + yield + rows = [e for e in prof.key_averages() if e.key.startswith("aten::") and e.key not in _FREE] + rows.sort(key=lambda e: -e.count) + total = sum(e.count for e in rows) + lines = [f"{e.count:6d} {e.cpu_time_total / 1e3:9.2f} ms {e.key}" for e in rows[:top]] + graphs = [e for e in prof.key_averages() if "Compiled" in e.key or "neuron" in e.key.lower()] + lines += [f"{e.count:6d} {e.cpu_time_total / 1e3:9.2f} ms {e.key}" for e in graphs[:10]] + logger.info(f"Neuron op census ({label}): {total} aten ops outside bookkeeping\n" + "\n".join(lines)) diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index 5f58855..dc8179a 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -3,6 +3,7 @@ from __future__ import annotations +import contextlib import copy import time from dataclasses import dataclass @@ -380,7 +381,10 @@ class ModelRunner: kwargs["kv_caches"] = kv_caches profiler = _maybe_profile_sample(self) + from deepinfer.neuron_ops.op_census import census + with ( + census("sample") if profiler is not None else contextlib.nullcontext(), cpu_timer.record("launch_sample"), cuda_timer.record("device_sample", scale=1), forbid_cuda_sync(), diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py index d59c100..3ca70a0 100644 --- a/deepinfer/worker/neuron_graph_runner.py +++ b/deepinfer/worker/neuron_graph_runner.py @@ -8,6 +8,7 @@ outside the graph. from __future__ import annotations +import contextlib import dataclasses import math import time @@ -200,7 +201,9 @@ class NeuronGraphRunner(CudaGraphRunner): profiler = cProfile.Profile() profiler.enable() started = time.perf_counter() - with hide_cuda(): + from deepinfer.neuron_ops.op_census import census + + with hide_cuda(), census(f"forward bs={batch_size}") if profiler is not None else contextlib.nullcontext(): out = self.compiled[batch_size](input_metadata_device) torch.neuron.synchronize() total = time.perf_counter() - started