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()}")