From 563c181ecca191070c345b3bf7b6686ed25aad90 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 08:10:16 +0800 Subject: [PATCH] Neuron: integer mod without torch.remainder in the sampler; sampler cProfile (#27) torch.remainder ran ~140 ms a call eagerly on the device and was 554 of the 640 ms a decode step spent sampling. a - (a // b) * b gives the same result for the non-negative operands (checked on 100k random values). Co-Authored-By: Claude Opus 5.5 --- .../tilelang_ops/neuron/random_sampling.py | 9 +++++- deepinfer/worker/model_runner.py | 29 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/random_sampling.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/random_sampling.py index 4a54f49..a5af72b 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/random_sampling.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/random_sampling.py @@ -93,6 +93,11 @@ def quantize_weights(weights: torch.Tensor) -> torch.Tensor: return torch.where(w >= 1.0, sat, (below_one * float(1 << SCALE_BITS)).to(torch.int32)) +def _mod(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """``a % b`` for non-negative integer tensors.""" + return a - torch.div(a, b, rounding_mode="floor") * b + + def two_level_sample( weights: torch.Tensor, seed: torch.Tensor, @@ -123,7 +128,9 @@ def two_level_sample( lo, hi = rng_first_two(seed, rng_sequence(positions, sample_case)) # (hi * 2**32 + lo) % total without uint64: all partial products stay below 2**63. - r = ((hi % safe_total) * ((1 << 32) % safe_total) + lo) % safe_total + # NEURON: a - (a // b) * b instead of torch.remainder, which runs ~140 ms a call eagerly on the device (#27); + # every operand is non-negative, so the two agree. + r = _mod(_mod(hi, safe_total) * _mod(torch.full_like(safe_total, 1 << 32), safe_total) + lo, safe_total) offset = (acc_cdf <= r.unsqueeze(1)).sum(dim=1) # first slot with cdf > r offset = torch.where(offset < acc_slots, offset, torch.zeros_like(offset)) diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index ea44665..5f58855 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -379,6 +379,7 @@ class ModelRunner: if self.nextn: kwargs["kv_caches"] = kv_caches + profiler = _maybe_profile_sample(self) with ( cpu_timer.record("launch_sample"), cuda_timer.record("device_sample", scale=1), @@ -389,6 +390,8 @@ class ModelRunner: sampler_inputs=sampler_inputs, **kwargs, ) + if profiler is not None: + _report_sample_profile(profiler) return output @@ -668,3 +671,29 @@ class ModelRunner: self.cuda_graph_runner.destroy() if self.dflash_cuda_graph_runner: self.dflash_cuda_graph_runner.destroy() + + +def _maybe_profile_sample(runner: ModelRunner): # noqa: ANN202 + """Neuron profiling (DEEPINFER_NEURON_REGION_TIMING): cProfile the 8th sampling call on rank 0.""" + if not (PLATFORM.is_neuron and envs.DEEPINFER_NEURON_REGION_TIMING): + return None + runner._sample_calls = getattr(runner, "_sample_calls", 0) + 1 + if runner._sample_calls != 8 or (torch.distributed.is_initialized() and torch.distributed.get_rank() != 0): + return None + import cProfile + + torch.neuron.synchronize() + profiler = cProfile.Profile() + profiler.enable() + return profiler + + +def _report_sample_profile(profiler) -> None: # noqa: ANN001 + import io + import pstats + + torch.neuron.synchronize() + profiler.disable() + buf = io.StringIO() + pstats.Stats(profiler, stream=buf).sort_stats("cumulative").print_stats(50) + logger.info(f"Neuron sample profile (rank 0):\n{buf.getvalue()}")