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 <noreply@anthropic.com>

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))

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