[Ops][Misc] Optimize 310P random sampling with inverse CDF (vllm-project#12966)

### What this PR does / why we need it?
This PR optimizes temperature-based random sampling on 310P without
generating random values on the NPU.

For regular sampling, it replaces the `[batch, vocab]` CPU
exponential-noise path and `Div + ArgMax` with inverse-CDF sampling: one
CPU uniform random value per request, asynchronous H2D transfer,
followed by NPU `Cumsum + SearchSorted`. For MTP recovered-token
sampling, exponential values are still generated on CPU via `-log(U)`,
with improved request-specific generator reuse, inactive-request
handling, and removal of redundant copies and host-side stream
synchronization.

This avoids the problematic 310P NPU random/in-place Add path while
significantly reducing CPU RNG and H2D overhead.

### Does this PR introduce _any_ user-facing change?
NA

### How was this patch tested?
Local test

RC  without temperature (benchmark):
<img width="1205" height="784" alt="image"
src="https://github.com/user-attachments/assets/62269e4d-4275-4b78-8249-45ecffa1b3b9"
/>

RC with temperature before this PR:
<img width="1193" height="778" alt="image1"
src="https://github.com/user-attachments/assets/f4d75b4b-5994-488e-8a0d-e1db1cbe80f8"
/>

RC with temperature after this PR:
<img width="1183" height="784" alt="image2"
src="https://github.com/user-attachments/assets/e90da9a1-695b-4201-9747-68724053b517"
/>


- vLLM version: v0.25.1
- vLLM main:
vllm-project/vllm@d02df74

---------

Signed-off-by: shysummer <sunhaoyu14@huawei.com>
Signed-off-by: sunhaoyu <sunhaoyu14@huawei.com>

Author: ShySummer

tests/ut/_310p/sample/test_sampler_310.py

@@ -23,7 +23,9 @@
     sample_sampler_module = ModuleType("vllm_ascend.sample.sampler")
     sample_sampler_module.DEFAULT_LOGPROBS_MODE = "raw_logprobs"  # type: ignore[attr-defined]
     sample_sampler_module.AscendSampler = type("AscendSampler", (), {})  # type: ignore[attr-defined]
-    sample_sampler_module.AscendTopKTopPSampler = type("AscendTopKTopPSampler", (), {})  # type: ignore[attr-defined]
+    sample_sampler_module.AscendTopKTopPSampler = type(  # type: ignore[attr-defined]
+        "AscendTopKTopPSampler", (), {}
+    )
     sys.modules["vllm_ascend.sample.sampler"] = sample_sampler_module
 
 if "vllm_ascend.utils" not in sys.modules:
@@ -35,248 +37,159 @@
 from vllm_ascend._310p.sample import sampler as sampler_310p  # noqa: E402
 
 
-class _FakeRow:
-    def __init__(self):
-        self.generators = []
-
-    def exponential_(self, generator=None):
-        self.generators.append(generator)
-        return self
-
-
-class _FakeQ:
-    def __init__(self, batch_size):
-        self.shape = (batch_size, 4)
-        self.default_exponential_called = False
-        self.rows = {idx: _FakeRow() for idx in range(batch_size)}
-
-    def cpu(self):
-        return self
-
-    def npu(self):
-        return self
-
-    def exponential_(self, generator=None):
-        if generator is None:
-            self.default_exponential_called = True
-        return self
-
-    def __getitem__(self, idx):
-        return self.rows[idx]
-
-    def __setitem__(self, idx, value):
-        self.rows[idx] = value
+class _SourceGenerator:
+    def __init__(self, seed: int):
+        self.seed = seed
 
+    def initial_seed(self) -> int:
+        return self.seed
 
-def _empty_like_side_effect(q_instances, template):
-    if isinstance(template, _FakeRow):
-        return _FakeRow()
-    return next(q_instances)
 
+class _RecordingStreamContext:
+    def __init__(self, events: list[str]):
+        self.events = events
 
-class _FakeCPUGenerator:
-    def __init__(self, device=None):
-        self.device = device
-        self.state = None
-        self.seed = None
+    def __enter__(self):
+        self.events.append("enter_global")
 
-    def set_state(self, state):
-        self.state = state
-
-    def manual_seed(self, seed):
-        self.seed = seed
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.events.append("exit_global")
 
 
 class TestSampler310pStandalone(unittest.TestCase):
     def tearDown(self):
         sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
 
-    def test_random_sample_310p_reuse_cpu_generator_cache(self):
-        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
+    def test_prepare_cpu_generators_preserves_requests_across_reordering(self):
+        source_a = _SourceGenerator(11)
+        source_b = _SourceGenerator(22)
+
+        first = sampler_310p._prepare_cpu_generators_310p({0: source_a, 1: source_b})
+        second = sampler_310p._prepare_cpu_generators_310p({0: source_b, 1: source_a})
+
+        self.assertIs(second[1], first[0])
+        self.assertIs(second[0], first[1])
+        self.assertIs(sampler_310p._CPU_GENERATOR_CACHE_310P[0][0], source_b)
+        self.assertIs(sampler_310p._CPU_GENERATOR_CACHE_310P[1][0], source_a)
+
+    def test_prepare_cpu_generators_replaces_changed_source(self):
+        source_first = _SourceGenerator(11)
+        source_second = _SourceGenerator(22)
+
+        first = sampler_310p._prepare_cpu_generators_310p({0: source_first})
+        second = sampler_310p._prepare_cpu_generators_310p({0: source_second})
+
+        self.assertIsNot(second[0], first[0])
+        expected = torch.Generator(device="cpu")
+        expected.manual_seed(22)
+        self.assertEqual(
+            torch.rand((), generator=second[0]).item(),
+            torch.rand((), generator=expected).item(),
+        )
+
+    def test_sample_from_cdf_handles_zero_weight_prefix_and_boundaries(self):
+        weights = torch.tensor(
+            [
+                [0.0, 2.0, 3.0, 5.0],
+                [1.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 4.0],
+            ]
+        )
+        uniforms = torch.tensor([0.2, 0.999, torch.finfo(torch.float32).tiny])
+
+        sampled = sampler_310p._sample_from_cdf_310p(weights, uniforms)
+
+        self.assertTrue(torch.equal(sampled, torch.tensor([2, 0, 3])))
+
+    def test_fill_exponential_honors_active_mask(self):
+        source_active = _SourceGenerator(31)
+        source_inactive = _SourceGenerator(41)
+        prepared = sampler_310p._prepare_cpu_generators_310p({0: source_active, 1: source_inactive})
+        active_state = prepared[0].get_state()
+        inactive_state = prepared[1].get_state()
+        real_rand = torch.rand
+
+        def rand_without_pinning(*args, **kwargs):
+            kwargs.pop("pin_memory", None)
+            return real_rand(*args, **kwargs)
+
+        with patch.object(sampler_310p.torch, "rand", side_effect=rand_without_pinning):
+            exponential = sampler_310p.fill_exponential_310p(
+                torch.empty((2, 4), dtype=torch.float32),
+                {0: source_active, 1: source_inactive},
+                active_mask=[True, False],
+            )
+
+        cached_active = sampler_310p._CPU_GENERATOR_CACHE_310P[0][1]
+        cached_inactive = sampler_310p._CPU_GENERATOR_CACHE_310P[1][1]
+        self.assertFalse(torch.equal(cached_active.get_state(), active_state))
+        self.assertTrue(torch.equal(cached_inactive.get_state(), inactive_state))
+        self.assertEqual(exponential.shape, (2, 4))
+        self.assertTrue(torch.isfinite(exponential).all())
+        self.assertTrue((exponential > 0).all())
+
+    def test_random_sample_waits_before_cdf_reads_probs(self):
+        events: list[str] = []
+        global_npu_stream = MagicMock()
+        current_npu_stream = MagicMock()
+        current_npu_stream.wait_stream.side_effect = lambda _: events.append("wait_global")
+        fake_npu = ModuleType("torch.npu")
+        fake_npu.current_stream = MagicMock(  # type: ignore[attr-defined]
+            return_value=current_npu_stream
+        )
         probs = MagicMock()
-        probs.div_.return_value = probs
-        probs.argmax.return_value = probs
-        probs.view.return_value = torch.tensor([0])
+        probs.shape = (1, 4)
+        probs.device = torch.device("cpu")
+        uniforms = MagicMock()
 
-        fake_q_first = _FakeQ(batch_size=2)
-        fake_q_second = _FakeQ(batch_size=2)
-        q_instances = iter([fake_q_first, fake_q_second])
+        def generate_uniforms(*args, **kwargs):
+            events.append("generate_uniforms")
+            return uniforms
 
-        npu_stream = MagicMock()
-        generator = MagicMock()
-        generator.get_state.return_value = b"state"
-        generator.initial_seed.return_value = 7
-        generators = {1: generator}
+        def sample_from_cdf(actual_probs, actual_uniforms):
+            events.append("sample_from_cdf")
+            self.assertIs(actual_probs, probs)
+            self.assertIs(actual_uniforms, uniforms)
+            return torch.tensor([2])
 
         with (
-            patch.object(sampler_310p, "npu_stream_switch", return_value=nullcontext()),
-            patch.object(sampler_310p, "global_stream", return_value=MagicMock()),
             patch.object(
-                sampler_310p.torch,
-                "empty_like",
-                side_effect=lambda template: _empty_like_side_effect(q_instances, template),
+                sampler_310p,
+                "npu_stream_switch",
+                return_value=_RecordingStreamContext(events),
             ),
-            patch.object(sampler_310p.torch, "Generator", side_effect=_FakeCPUGenerator) as gen_ctor,
             patch.object(
-                sampler_310p.torch,
-                "npu",
-                ModuleType("torch.npu"),
-                create=True,
+                sampler_310p,
+                "global_stream",
+                return_value=global_npu_stream,
             ),
-        ):
-            sampler_310p.torch.npu.current_stream = MagicMock(return_value=npu_stream)
-            sampler_310p._random_sample_310p(probs, generators)
-            sampler_310p._random_sample_310p(probs, generators)
-
-        self.assertEqual(gen_ctor.call_count, 1)
-        self.assertIn(1, sampler_310p._CPU_GENERATOR_CACHE_310P)
-        cached_cpu_generator, source_generator_id = sampler_310p._CPU_GENERATOR_CACHE_310P[1]
-        self.assertIs(fake_q_first.rows[1].generators[0], cached_cpu_generator)
-        self.assertIs(fake_q_second.rows[1].generators[0], cached_cpu_generator)
-        self.assertEqual(source_generator_id, id(generator))
-        self.assertEqual(cached_cpu_generator.state, b"state")
-        self.assertIsNone(cached_cpu_generator.seed)
-        self.assertEqual(npu_stream.wait_stream.call_count, 2)
-
-    def test_random_sample_310p_fallback_to_initial_seed_when_set_state_failed(self):
-        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
-        probs = MagicMock()
-        probs.div_.return_value = probs
-        probs.argmax.return_value = probs
-        probs.view.return_value = torch.tensor([1])
-
-        fake_q = _FakeQ(batch_size=1)
-        q_instances = iter([fake_q])
-        npu_stream = MagicMock()
-        generator = MagicMock()
-        generator.get_state.side_effect = RuntimeError("state read failed")
-        generator.initial_seed.return_value = 1234
-        generators = {0: generator}
-
-        class _FailSetStateCPUGenerator(_FakeCPUGenerator):
-            def set_state(self, state):
-                raise RuntimeError("state set failed")
-
-        with (
-            patch.object(sampler_310p, "npu_stream_switch", return_value=nullcontext()),
-            patch.object(sampler_310p, "global_stream", return_value=MagicMock()),
             patch.object(
-                sampler_310p.torch,
-                "empty_like",
-                side_effect=lambda template: _empty_like_side_effect(q_instances, template),
+                sampler_310p,
+                "_generate_request_uniforms_310p",
+                side_effect=generate_uniforms,
             ),
-            patch.object(sampler_310p.torch, "Generator", side_effect=_FailSetStateCPUGenerator),
             patch.object(
-                sampler_310p.torch,
-                "npu",
-                ModuleType("torch.npu"),
-                create=True,
+                sampler_310p,
+                "_sample_from_cdf_310p",
+                side_effect=sample_from_cdf,
             ),
+            patch.object(sampler_310p.torch, "npu", fake_npu, create=True),
         ):
-            sampler_310p.torch.npu.current_stream = MagicMock(return_value=npu_stream)
-            sampler_310p._random_sample_310p(probs, generators)
-
-        cached_cpu_generator, source_generator_id = sampler_310p._CPU_GENERATOR_CACHE_310P[0]
-        self.assertEqual(source_generator_id, id(generator))
-        self.assertEqual(cached_cpu_generator.seed, 1234)
-        self.assertIs(fake_q.rows[0].generators[0], cached_cpu_generator)
-        self.assertEqual(npu_stream.wait_stream.call_count, 1)
-
-    def test_random_sample_310p_rebuild_cache_when_generator_identity_changes(self):
-        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
-        probs = MagicMock()
-        probs.div_.return_value = probs
-        probs.argmax.return_value = probs
-        probs.view.return_value = torch.tensor([0])
-
-        fake_q_first = _FakeQ(batch_size=1)
-        fake_q_second = _FakeQ(batch_size=1)
-        q_instances = iter([fake_q_first, fake_q_second])
-        npu_stream = MagicMock()
-
-        generator_first = MagicMock()
-        generator_first.get_state.return_value = b"state-1"
-        generator_first.initial_seed.return_value = 11
-
-        generator_second = MagicMock()
-        generator_second.get_state.return_value = b"state-2"
-        generator_second.initial_seed.return_value = 22
-
-        with (
-            patch.object(sampler_310p, "npu_stream_switch", return_value=nullcontext()),
-            patch.object(sampler_310p, "global_stream", return_value=MagicMock()),
-            patch.object(
-                sampler_310p.torch,
-                "empty_like",
-                side_effect=lambda template: _empty_like_side_effect(q_instances, template),
-            ),
-            patch.object(sampler_310p.torch, "Generator", side_effect=_FakeCPUGenerator) as gen_ctor,
-            patch.object(
-                sampler_310p.torch,
-                "npu",
-                ModuleType("torch.npu"),
-                create=True,
-            ),
-        ):
-            sampler_310p.torch.npu.current_stream = MagicMock(return_value=npu_stream)
-            sampler_310p._random_sample_310p(probs, {0: generator_first})
-            sampler_310p._random_sample_310p(probs, {0: generator_second})
-
-        self.assertEqual(gen_ctor.call_count, 2)
-        first_cpu_generator = fake_q_first.rows[0].generators[0]
-        second_cpu_generator = fake_q_second.rows[0].generators[0]
-        self.assertIsNot(first_cpu_generator, second_cpu_generator)
-        self.assertEqual(first_cpu_generator.state, b"state-1")
-        self.assertEqual(second_cpu_generator.state, b"state-2")
-        cached_cpu_generator, source_generator_id = sampler_310p._CPU_GENERATOR_CACHE_310P[0]
-        self.assertIs(cached_cpu_generator, second_cpu_generator)
-        self.assertEqual(source_generator_id, id(generator_second))
-
-    def test_fill_cpu_exponential_310p_moves_has_draft_mask_to_cpu(self):
-        """Regression: NPU has_draft_mask must be moved to CPU before torch.where."""
-        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
-
-        q_cpu = torch.full((2, 4), 7.0)
-        cpu_mask = torch.tensor([True, False])
-        has_draft_mask = MagicMock()
-        has_draft_mask.cpu.return_value = cpu_mask
-
-        def _make_source_generator(seed: int):
-            source_generator = MagicMock()
-            seed_generator = torch.Generator(device="cpu")
-            seed_generator.manual_seed(seed)
-            source_generator.get_state.return_value = seed_generator.get_state()
-            source_generator.initial_seed.return_value = seed
-            return source_generator
-
-        where_conditions = []
-        real_where = torch.where
-
-        def where_spy(condition, x, y):
-            where_conditions.append(condition.detach().clone())
-            self.assertEqual(condition.device.type, "cpu")
-            self.assertEqual(x.device.type, "cpu")
-            self.assertEqual(y.device.type, "cpu")
-            return real_where(condition, x, y)
-
-        with patch.object(sampler_310p.torch, "where", side_effect=where_spy):
-            sampler_310p._fill_cpu_exponential_310p(
-                q_cpu,
-                {
-                    0: _make_source_generator(42),
-                    1: _make_source_generator(43),
-                },
-                has_draft_mask,
-            )
-
-        has_draft_mask.cpu.assert_called_once()
-        self.assertEqual(len(where_conditions), 2)
-        self.assertTrue(bool(where_conditions[0]))
-        self.assertFalse(bool(where_conditions[1]))
-        # Row 0 (masked): overwritten by seeded exponential via torch.where.
-        self.assertFalse(torch.equal(q_cpu[0], torch.full((4,), 7.0)))
-        # Row 1 (unmasked): also overwritten by the default exponential_ prefill.
-        self.assertFalse(torch.equal(q_cpu[1], torch.full((4,), 7.0)))
+            sampled = sampler_310p._random_sample_310p(probs, {})
+
+        self.assertTrue(torch.equal(sampled, torch.tensor([2])))
+        self.assertEqual(
+            events,
+            [
+                "enter_global",
+                "generate_uniforms",
+                "exit_global",
+                "wait_global",
+                "sample_from_cdf",
+            ],
+        )
+        current_npu_stream.wait_stream.assert_called_once_with(global_npu_stream)
+        global_npu_stream.wait_stream.assert_not_called()
 
 
 if __name__ == "__main__":