diff --git a/vllm_fl/dispatch/backends/vendor/kunlunxin/patch.py b/vllm_fl/dispatch/backends/vendor/kunlunxin/patch.py index 7afc4e4ee..3bb22656c 100644 --- a/vllm_fl/dispatch/backends/vendor/kunlunxin/patch.py +++ b/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.