Kimi k3

vllm_kunlun/distributed/deepep_buffer_v2.py

@@ -0,0 +1,111 @@
+"""Route vLLM's DeepEP all2all managers to Kunlun's ``BufferV2``.
+
+Kunlun's ``deep_ep`` build ships two buffer classes:
+
+* ``Buffer`` -- a reduced API (no ``get_dispatch_layout``, and ``dispatch``
+  without ``handle`` / ``num_tokens_per_rank`` / ``is_token_in_rank``) that
+  vLLM's ``prepare_finalize/deepep_*.py`` cannot drive.
+* ``BufferV2`` -- keeps the upstream DeepEP API, ignores the arguments its
+  runtime does not need and remembers ``num_combined_tokens`` for combine.
+
+vLLM imports ``deep_ep`` lazily inside the all2all managers, so patching the
+module object after import is not enough; install a wrapper module in
+``sys.modules`` whose ``Buffer`` is a ``BufferV2`` subclass.
+
+The subclass also handles two constructor mismatches:
+
+* it drops kwargs this build does not accept (upstream passes
+  ``explicitly_destroy``, ``allow_nvlink_for_low_latency_mode``, ...);
+* it fills in ``num_experts``. vLLM never passes it, so the buffer would be
+  sized for ``BufferV2``'s default of 256 while dispatch runs with the model's
+  real expert count, and the notify step then fails with
+  ``bkcl_notify_dispatch_standard failed / recv_num_tokens >= 0``.
+"""
+
+import inspect
+import logging
+import sys
+from types import SimpleNamespace
+
+logger = logging.getLogger("vllm_kunlun")
+
+
+def _infer_num_experts():
+    """Routed expert count for the running model, or None if unknown."""
+    try:
+        from vllm.config import get_current_vllm_config
+
+        text_config = get_current_vllm_config().model_config.hf_text_config
+    except Exception:
+        return None
+    for attr in ("num_experts", "n_routed_experts", "num_local_experts"):
+        n = getattr(text_config, attr, None)
+        if isinstance(n, int) and n > 0:
+            return n
+    return None
+
+
+def _make_buffer_cls(real):
+    accepted = set(inspect.signature(real.BufferV2.__init__).parameters) - {"self"}
+
+    class KunlunDeepEPBuffer(real.BufferV2):
+        """BufferV2 with the argument handling vLLM expects."""
+
+        def __init__(self, *args, **kwargs):
+            dropped = [k for k in kwargs if k not in accepted]
+            for k in dropped:
+                kwargs.pop(k)
+            if dropped:
+                logger.debug(
+                    "[KunlunPlugin] deep_ep BufferV2 ignores kwargs %s", dropped
+                )
+            if "num_experts" not in kwargs:
+                n = _infer_num_experts()
+                if n is not None:
+                    kwargs["num_experts"] = n
+                    logger.info(
+                        "[KunlunPlugin] deep_ep buffer num_experts=%d "
+                        "(inferred from model config)",
+                        n,
+                    )
+            super().__init__(*args, **kwargs)
+
+        def destroy(self):
+            """Upstream passes explicitly_destroy=True and calls this later."""
+            parent = getattr(super(), "destroy", None)
+            if parent is not None:
+                parent()
+
+    return KunlunDeepEPBuffer
+
+
+def applied(mod) -> bool:
+    real = sys.modules.get("deep_ep")
+    return real is not None and getattr(real, "_kunlun_buffer_v2", False)
+
+
+def apply(mod) -> None:
+    try:
+        import deep_ep
+    except Exception as exc:  # deep_ep is an optional dependency
+        logger.info(
+            "[KunlunPlugin] deep_ep not available (%s); DeepEP backends disabled",
+            str(exc).splitlines()[0],
+        )
+        return
+    if getattr(deep_ep, "_kunlun_buffer_v2", False):
+        return
+    if not hasattr(deep_ep, "BufferV2"):
+        logger.warning(
+            "[KunlunPlugin] deep_ep has no BufferV2; leaving Buffer as is"
+        )
+        return
+
+    wrapped = SimpleNamespace(**deep_ep.__dict__)
+    wrapped.Buffer = _make_buffer_cls(deep_ep)
+    wrapped._kunlun_buffer_v2 = True
+    sys.modules["deep_ep"] = wrapped
+    logger.info(
+        "[KunlunPlugin] deep_ep.Buffer -> BufferV2 (version %s)",
+        getattr(getattr(deep_ep, "version", None), "__version__", "unknown"),
+    )

vllm_kunlun/distributed/deepep_ht_dtype.py

@@ -0,0 +1,119 @@
+"""Let DeepEP high-throughput finalize accept a float16 expert output.
+
+``prepare_finalize/deepep_ht.py`` guards its combine call with
+
+    assert fused_expert_output.dtype == torch.bfloat16
+
+because the reference DeepEP combine kernel is bfloat16-only. Kunlun's
+``BufferV2.combine`` accepts either 2-byte float type, so for a float16 model
+that assertion is the only thing in the way, and satisfying it would cost an
+extra read+write of an ``[num_tokens, hidden]`` tensor per MoE layer per step.
+
+Rebuild ``_finalize`` from its own source with that one assert statement
+removed. The statement is located with ``ast`` instead of by text matching, and
+the rebuilt function keeps the original module globals, so it stays in step with
+upstream changes to the rest of the method. If the assert can no longer be
+identified unambiguously, fall back to casting, which is correct but slower.
+
+``_finalize`` is the private method that holds the assertion; the public
+``finalize`` is a thin wrapper, so patching that one has no effect.
+"""
+
+import ast
+import inspect
+import logging
+import sys
+import textwrap
+
+import torch
+
+logger = logging.getLogger("vllm_kunlun")
+
+_CLASSES = ("DeepEPHTPrepareAndFinalize",)
+
+
+def applied(mod) -> bool:
+    for name in _CLASSES:
+        cls = getattr(mod, name, None)
+        if cls is not None and not getattr(cls._finalize, "_kunlun_fp16_ok", False):
+            return False
+    return True
+
+
+def _strip_bf16_assert(cls) -> bool:
+    """Rebuild cls._finalize without the bfloat16 assertion. True on success."""
+    try:
+        src = textwrap.dedent(inspect.getsource(cls._finalize))
+        tree = ast.parse(src)
+        func = tree.body[0]
+        if not isinstance(func, ast.FunctionDef) or func.decorator_list:
+            return False
+        spans = [
+            (node.lineno, node.end_lineno)
+            for node in ast.walk(func)
+            if isinstance(node, ast.Assert) and "bfloat16" in ast.unparse(node.test)
+        ]
+        if len(spans) != 1:
+            return False
+        lo, hi = spans[0]
+        lines = src.splitlines()
+        kept = [
+            line for i, line in enumerate(lines, 1) if not lo <= i <= (hi or lo)
+        ]
+        new_src = "\n".join(kept) + "\n"
+        if "bfloat16" in new_src:
+            return False
+        module = sys.modules[cls.__module__]
+        code = compile(new_src, "<vllm_kunlun:%s._finalize>" % cls.__name__, "exec")
+        # Execute in the module's own namespace so the rebuilt function keeps
+        # the exact globals the original used, then drop the stray binding.
+        exec(code, module.__dict__)
+        rebuilt = module.__dict__.pop(func.name)
+        rebuilt.__qualname__ = cls.__name__ + "." + func.name
+        rebuilt._kunlun_fp16_ok = True
+        setattr(cls, func.name, rebuilt)
+        return True
+    except Exception:
+        logger.warning(
+            "[KunlunPlugin] could not rebuild %s._finalize without the bfloat16 "
+            "assert; falling back to casting",
+            cls.__name__,
+            exc_info=True,
+        )
+        return False
+
+
+def _install_cast_fallback(cls) -> None:
+    orig = cls._finalize
+
+    def _finalize(self, *args, _orig=orig, **kwargs):
+        key = "fused_expert_output"
+        if key in kwargs and torch.is_tensor(kwargs[key]):
+            if kwargs[key].dtype != torch.bfloat16:
+                kwargs[key] = kwargs[key].to(torch.bfloat16)
+        elif len(args) >= 2 and torch.is_tensor(args[1]):
+            if args[1].dtype != torch.bfloat16:
+                args = args[:1] + (args[1].to(torch.bfloat16),) + args[2:]
+        return _orig(self, *args, **kwargs)
+
+    _finalize._kunlun_fp16_ok = True
+    cls._finalize = _finalize
+
+
+def apply(mod) -> None:
+    for name in _CLASSES:
+        cls = getattr(mod, name, None)
+        if cls is None or getattr(cls._finalize, "_kunlun_fp16_ok", False):
+            continue
+        if _strip_bf16_assert(cls):
+            logger.info(
+                "[KunlunPlugin] %s._finalize rebuilt without the bfloat16 assert "
+                "(float16 expert output goes to combine as is)",
+                name,
+            )
+        else:
+            _install_cast_fallback(cls)
+            logger.info(
+                "[KunlunPlugin] %s._finalize casts the expert output to bfloat16",
+                name,
+            )

vllm_kunlun/distributed/dp_device_assign.py

@@ -0,0 +1,55 @@
+"""Keep vLLM's CUDA-style device selection when data parallelism is enabled.
+
+``vllm/v1/engine/utils.py`` decides whether each DP rank needs its own
+``CUDA_VISIBLE_DEVICES``-style isolation with
+
+    needs_device_env_isolation = not (
+        current_platform.is_cuda_alike() or current_platform.is_xpu())
+
+``KunlunPlatform._enum`` is ``PlatformEnum.OOT``, so both checks are False and
+vLLM shards the device list per DP rank, setting
+``parallel_config.assigned_physical_gpu_ids`` to a single id (e.g. ``[7]`` for
+DP rank 7).
+
+But ``GpuWorker.init_device`` takes the CUDA path (``device_type == "cuda"``)
+and offsets the worker's local rank by ``dp_local_rank * tp_pp_world_size``,
+then asserts ``local_rank < len(assigned_physical_gpu_ids)``. With DP8/TP1 that
+is ``7 < 1`` and startup dies with
+
+    AssertionError: local_rank 7 is out of bounds for
+                    assigned_physical_gpu_ids [7]
+
+On Kunlun every process sees all eight devices and the worker selects its own by
+index, exactly like CUDA, so the per-rank sharding is both unnecessary and
+wrong. Neutralise it and let ``assigned_physical_gpu_ids`` stay None, which is
+what the CUDA path expects.
+"""
+
+import logging
+
+logger = logging.getLogger("vllm_kunlun")
+
+
+def applied(mod) -> bool:
+    fn = getattr(mod, "set_assigned_physical_gpu_ids_for_dp_rank", None)
+    return fn is None or getattr(fn, "_kunlun_dp_noop", False)
+
+
+def apply(mod) -> None:
+    name = "set_assigned_physical_gpu_ids_for_dp_rank"
+    fn = getattr(mod, name, None)
+    if fn is None or getattr(fn, "_kunlun_dp_noop", False):
+        return
+
+    def set_assigned_physical_gpu_ids_for_dp_rank(
+        vllm_config, local_dp_rank, user_assigned_gpu_ids=None
+    ):
+        # Only honour an explicit --device-ids list; never shard per DP rank.
+        vllm_config.parallel_config.assigned_physical_gpu_ids = user_assigned_gpu_ids
+
+    set_assigned_physical_gpu_ids_for_dp_rank._kunlun_dp_noop = True
+    setattr(mod, name, set_assigned_physical_gpu_ids_for_dp_rank)
+    logger.info(
+        "[KunlunPlugin] DP device sharding disabled; workers select devices by "
+        "index like CUDA"
+    )