[Kunlunxin] Fix duplicate generation (#332)

<!--
 Copyright 2026 FlagOS Contributors

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 -->

### PR Category

Vendor

### PR Type

Bug Fixes

### Description

This PR fixes the duplicate generation issue on Kunlunxin by patching
the sampler RNG implementation to ensure consistent sampling behavior
during inference.

### Related Issues

N/A

### Changes

* Patch the sampler RNG implementation for Kunlunxin.
* Fix the duplicate generation issue during model inference.
* Keep the implementation transparent to existing inference workflows.

### Testing

Validated on Kunlunxin hardware with:

* Qwen3.5-35B-A3B
* Qwen3.6-27B

Verified that the duplicate generation issue is resolved on both models.

### Checklist

* [x] I have run the existing tests and they pass
* [ ] I have added tests for my changes (if applicable)
* [ ] I have updated the documentation (if applicable)

Co-authored-by: 0songHan <345279710@qq.com>
Co-authored-by: XMing <xmhubj@gmail.com>

Author: 3 people

vllm_fl/dispatch/backends/vendor/kunlunxin/patch.py

@@ -32,9 +32,46 @@ def apply_kunlunxin_patches():
     patch_fused_gdn_gating()
     patch_op_cls()
     patch_ssm_cache_update()
+    patch_sampler_rng()
     logger.info("Applied all Kunlunxin patches")
 
 
+# ── sampler RNG (TP-consistent unseeded sampling) ──
+def patch_sampler_rng():
+    try:
+        import torch
+        import vllm.v1.sample.ops.topk_topp_sampler as _sampler_mod
+
+        def _tp_consistent_random_sample(probs, generators):
+            q = torch.empty_like(probs)
+            if len(generators) != probs.shape[0]:
+                q.uniform_()
+                q = -torch.log(1-q)
+                q = q.clamp(min=1e-12)
+            if generators:
+                # TODO(woosuk): This can be slow because we handle each request
+                # one by one. Optimize this.
+                if os.getenv("FAST_RANDOM_SAMPLE") == "1":
+                    for i, generator in generators.items():
+                        q[i].uniform_(generator=generator)
+                    q = -torch.log(1-q)
+                    q = q.clamp(min=1e-12)
+                else:
+                    for i, generator in generators.items():
+                        q[i].uniform_(generator=generator)
+                        q[i] = -torch.log(1-q[i])
+                        q[i] = q[i].clamp(min=1e-12)
+
+            return probs.div_(q).argmax(dim=-1).view(-1)
+
+        _sampler_mod.random_sample = _tp_consistent_random_sample
+        logger.info(
+            "Patched sampler random_sample for TP-consistent sampling on Kunlunxin"
+        )
+    except Exception as e:
+        logger.warning("Failed to patch sampler RNG: %s", e)
+
+
 # ── causal_conv1d ──
 def patch_causal_conv1d():
     """Replace causal_conv1d_fn / causal_conv1d_update with Kunlunxin impls.