From 7c3b85d4fb9fd4cb4ffd47f325a7f2caf61d374d Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Mon, 24 Aug 2026 16:03:41 +0800 Subject: [PATCH 01/19] Add DSV4 CAM async connector Signed-off-by: bjf-frz --- afd_plugin/__init__.py | 51 ++ afd_plugin/compat/npu/feature_validation.py | 100 ++++ afd_plugin/compat/npu/ops.py | 24 + afd_plugin/compat/npu/runtime.py | 7 + afd_plugin/compat/patches/engine_core.py | 16 + .../patches/npu/deepseek_v4_kv_cache.py | 164 ++++++ .../compat/patches/npu/modelslim_dsv4.py | 119 ++++ afd_plugin/connectors/metadata.py | 21 +- afd_plugin/connectors/npu/async_cam.py | 509 +++++++++++------ afd_plugin/connectors/npu/camp2p.py | 127 ++++- .../model_executor/models/deepseek_v2.py | 8 + .../model_executor/models/deepseek_v4.py | 518 +----------------- .../model_executor/models/deepseek_v4_cuda.py | 510 +++++++++++++++++ .../model_executor/models/deepseek_v4_npu.py | 502 +++++++++++++++++ .../npu/deepseek_v2_async_cam_forward.py | 7 + .../models/npu/deepseek_v4_attention_gate.py | 124 +++++ .../v1/worker/npu/attention_model_runner.py | 97 ++-- afd_plugin/v1/worker/npu/attention_worker.py | 3 + afd_plugin/v1/worker/npu/ffn_model_runner.py | 56 +- afd_plugin/v1/worker/npu/ffn_worker.py | 6 + afd_plugin/v1/worker/npu/forward_context.py | 7 +- .../v1/worker/npu/npu_ubatch_wrapper.py | 10 +- .../compat/npu/test_dsv4_async_validation.py | 55 ++ .../connectors/test_async_cam_connector.py | 160 +++++- .../unit/connectors/test_camp2p_connector.py | 46 +- .../models/test_deepseek_v2_proxy.py | 16 + .../models/test_deepseek_v4_attention_gate.py | 35 ++ tests/unit/v1/worker/test_npu_runtime.py | 52 +- 28 files changed, 2566 insertions(+), 784 deletions(-) create mode 100644 afd_plugin/compat/patches/npu/deepseek_v4_kv_cache.py create mode 100644 afd_plugin/compat/patches/npu/modelslim_dsv4.py create mode 100644 afd_plugin/model_executor/models/deepseek_v4_cuda.py create mode 100644 afd_plugin/model_executor/models/deepseek_v4_npu.py create mode 100644 afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py create mode 100644 tests/unit/compat/npu/test_dsv4_async_validation.py create mode 100644 tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index f7d81451..bb97551d 100644 --- a/afd_plugin/__init__.py +++ b/afd_plugin/__init__.py @@ -6,6 +6,8 @@ import importlib.util import logging +import multiprocessing +import os from importlib.metadata import PackageNotFoundError, version from pathlib import Path from types import MappingProxyType @@ -45,6 +47,51 @@ def __getattr__(name: str): _logger = logging.getLogger(__name__) _registered = False + +def _force_spawn_multiprocessing_if_requested() -> None: + """Pin Python's default context for the A3 Python 3.12 runtime. + + ``VLLM_WORKER_MULTIPROC_METHOD`` controls vLLM's explicit context only. + Some NPU runtime helpers use Python's default context and otherwise retain + the task's inherited ``forkserver`` setting, which cannot restore the + signal-handler state in this environment. Keep the override opt-in so it + remains isolated to the affected CAM async recipe. + """ + if os.environ.get("AFD_FORCE_SPAWN_MULTIPROCESSING") != "1": + return + multiprocessing.set_start_method("spawn", force=True) + + # A few optional runtime helpers request ``forkserver`` explicitly instead + # of consulting the default start method. On the A3 Python 3.12 image the + # forkserver cannot restore its inherited signal-handler state, so every + # child it starts exits before running user code. Keep this narrowly + # opt-in with the recipe environment variable above: callers requesting a + # forkserver receive the equivalent spawn context instead. + contexts = multiprocessing.context._concrete_contexts + contexts["forkserver"] = contexts["spawn"] + + # TE Fusion keeps a module reference to ``multiprocessing`` and calls its + # public ``get_context("forkserver")`` API directly inside each model + # worker. Replacing the registry alone is not sufficient for all Python + # 3.12 context instances, so redirect that explicit request as well. + if not getattr(multiprocessing, "_afd_spawn_context_redirect", False): + original_get_context = multiprocessing.get_context + + def get_spawn_context(method: str | None = None): + if method == "forkserver": + method = "spawn" + return original_get_context(method) + + multiprocessing.get_context = get_spawn_context + multiprocessing._afd_spawn_context_redirect = True + + +# vLLM model workers import the configured worker class directly in spawned +# child interpreters; they do not invoke the general-plugin entry point below. +# Apply the opt-in setting at package import time so it is also in effect before +# Ascend TE Fusion initializes its compilation workers. +_force_spawn_multiprocessing_if_requested() + _DEEPSEEK_MODEL_REGISTRATIONS = { "DeepseekForCausalLM": ( "afd_plugin.model_executor.models.deepseek_v2:AFDDeepseekForCausalLM" @@ -64,6 +111,9 @@ def __getattr__(name: str): "GlmMoeDsaForCausalLM": ( "afd_plugin.model_executor.models.deepseek_v2:AFDGlmMoeDsaForCausalLM" ), + "DeepseekV4ForCausalLM": ( + "afd_plugin.model_executor.models.deepseek_v4:AFDDeepseekV4ForCausalLM" + ), } _QWEN_MODEL_REGISTRATIONS = { @@ -102,6 +152,7 @@ def register_afd() -> None: _logger.debug("AFD plugin: register_afd() already completed") return + _force_spawn_multiprocessing_if_requested() _logger.debug("AFD plugin: register_afd() called") if importlib.util.find_spec("vllm") is None: _logger.debug("AFD plugin: vLLM not found, skipping runtime registration") diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 823144b7..d0548aaf 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -40,6 +40,8 @@ def fail_if_unsupported_npu_afd_features( afd_config, extra_info, ) + if _is_dsv4_target(vllm_config): + _fail_if_unsupported_dsv4_async_features(afd_config, extra_info) return if afd_config.compute_gate_on_attention: @@ -84,6 +86,45 @@ def fail_if_unsupported_npu_afd_features( ) +def _is_dsv4_target(vllm_config: VllmConfig) -> bool: + """Return whether the target model is DeepSeek V4.""" + model_config = vllm_config.model_config + hf_config = getattr(model_config, "hf_config", None) + if hf_config is None: + hf_config = getattr(model_config, "hf_text_config", None) + if hf_config is None: + return False + if getattr(hf_config, "model_type", None) == "deepseek_v4": + return True + architectures = getattr(hf_config, "architectures", ()) or () + return any("DeepseekV4" in str(architecture) for architecture in architectures) + + +def _fail_if_unsupported_dsv4_async_features( + afd_config: AFDConfig, + extra_info: ConnectorExtraInfo, +) -> None: + """Keep the initial DSV4 Async CAM target path deliberately narrow.""" + from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo + + if not afd_config.compute_gate_on_attention: + raise RuntimeError( + "DSV4 CAMAsyncAFDConnector requires compute_gate_on_attention=true", + ) + if not isinstance(extra_info, AFDAsyncExtraInfo): + raise TypeError( + "DSV4 CAMAsyncAFDConnector requires AFDAsyncExtraInfo, got " + f"{type(extra_info).__name__}", + ) + if extra_info.dynamic_quant != 1: + raise RuntimeError( + "DSV4 Flash-INT8 CAMAsyncAFDConnector requires dynamicQuant=1", + ) + # async_moe_ubatching is validated by + # _fail_if_unsupported_npu_async_moe_ubatching_features for all async CAM + # targets, so DSV4 no longer needs a bespoke rejection here. + + def _fail_if_unsupported_npu_afd_async_features( vllm_config: VllmConfig, afd_config: AFDConfig, @@ -123,6 +164,65 @@ def _fail_if_unsupported_npu_afd_async_features( raise RuntimeError( "CAMAsyncAFDConnector currently supports only dynamicQuant 0 or 1", ) + _validate_cam_dp_topology(vllm_config, afd_config, extra_info) + + +def _validate_cam_dp_topology( + vllm_config: VllmConfig, + afd_config: AFDConfig, + extra_info: ConnectorExtraInfo, +) -> None: + """Require a DP-scoped CAM world to match each role's parallel layout.""" + from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo + + if not isinstance(extra_info, AFDAsyncExtraInfo): + return + parallel_config = vllm_config.parallel_config + attn_ranks_per_dp = int(extra_info.attn_ranks_per_dp) + if afd_config.num_attention_ranks % attn_ranks_per_dp != 0: + raise RuntimeError( + "CAMAsyncAFDConnector requires num_attention_ranks to be divisible " + "by attn_ranks_per_dp", + ) + cam_dp_size = afd_config.num_attention_ranks // attn_ranks_per_dp + if not extra_info.shared_ffn_pool and afd_config.num_ffn_ranks % cam_dp_size != 0: + raise RuntimeError( + "CAMAsyncAFDConnector requires num_ffn_ranks to be divisible by " + "the CAM DP group count", + ) + if afd_config.role == "attention": + if int(parallel_config.data_parallel_size) != cam_dp_size: + raise RuntimeError( + "CAMAsyncAFDConnector Attention data_parallel_size must equal " + "the CAM DP group count", + ) + if int(parallel_config.tensor_parallel_size) != attn_ranks_per_dp: + raise RuntimeError( + "CAMAsyncAFDConnector Attention tensor_parallel_size must equal " + "attn_ranks_per_dp", + ) + elif extra_info.shared_ffn_pool: + # A CAM FFN endpoint is an EP rank, rather than a TP shard. The + # shared pool is reused by each Attention DP group, but retains its + # eight independent FFN processes as DP8 x TP1 x EP8. + if int(parallel_config.data_parallel_size) != afd_config.num_ffn_ranks: + raise RuntimeError( + "CAMAsyncAFDConnector shared_ffn_pool requires FFN " + "data_parallel_size to equal num_ffn_ranks", + ) + if int(parallel_config.tensor_parallel_size) != 1: + raise RuntimeError( + "CAMAsyncAFDConnector shared_ffn_pool requires FFN " + "tensor_parallel_size=1", + ) + elif cam_dp_size > 1 and ( + int(parallel_config.data_parallel_size) != cam_dp_size + or int(parallel_config.tensor_parallel_size) != attn_ranks_per_dp + ): + raise RuntimeError( + "CAMAsyncAFDConnector FFN must use the same DPxTP partition as " + "Attention when CAM has multiple DP groups", + ) def _fail_if_unsupported_npu_async_moe_ubatching_features( diff --git a/afd_plugin/compat/npu/ops.py b/afd_plugin/compat/npu/ops.py index f2ea057b..1437a853 100644 --- a/afd_plugin/compat/npu/ops.py +++ b/afd_plugin/compat/npu/ops.py @@ -4,6 +4,7 @@ from __future__ import annotations +import ctypes import os from functools import lru_cache from pathlib import Path @@ -16,6 +17,10 @@ CAM_DISPATCH_RECV = "async_dispatch_recv" CAM_COMBINE_SEND = "async_combine_send" CAM_COMBINE_RECV = "async_combine_recv" +CAM_CUST_OPAPI_ENV = "CAM_CUST_OPAPI_LIB_PATH" +CAM_CUST_OPAPI_DEFAULT = Path( + "/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libopapi.so", +) def get_afd_cann_vendor_path() -> Path: @@ -68,6 +73,24 @@ def _assert_cam_namespace_registered(torch: object) -> None: ) +def _preload_cam_cust_opapi() -> None: + """Make CAM's custom op-api symbols visible to the Python extension. + + ``umdk_cam_op_lib`` resolves the aclnn CAM entry points by name while its + operators are first invoked. CANN's stock ``libopapi.so`` does not export + those symbols; they live in CAM's ``libopapi.so``. Load that library + globally before importing the extension so the real CAM operators bind to + the vendor implementation instead of the stock op-api library. + """ + + library = Path(os.environ.get(CAM_CUST_OPAPI_ENV, CAM_CUST_OPAPI_DEFAULT)) + if not library.is_file(): + raise RuntimeError( + "CAMAsyncAFDConnector requires CAM libopapi.so; expected " + f"{library}. Install the CAM operator package first.", + ) + ctypes.CDLL(str(library), mode=ctypes.RTLD_GLOBAL) + @lru_cache(maxsize=1) def ensure_cam_p2p_ops_available() -> None: """Import the custom operators used by ``CAMP2pAFDConnector``. @@ -107,6 +130,7 @@ def ensure_cam_async_ops_available() -> None: """Ensure the runtime exposes the real CAM async operator namespace.""" try: + _preload_cam_cust_opapi() import torch import torch_npu # noqa: F401 import umdk_cam_op_lib # noqa: F401 diff --git a/afd_plugin/compat/npu/runtime.py b/afd_plugin/compat/npu/runtime.py index 30297064..72cc826a 100644 --- a/afd_plugin/compat/npu/runtime.py +++ b/afd_plugin/compat/npu/runtime.py @@ -69,12 +69,19 @@ def apply_afd_ascend_patches_if_needed() -> None: from afd_plugin.compat.patches.npu.mla_graph import ( apply_afd_mla_graph_patch, ) + from afd_plugin.compat.patches.npu.modelslim_dsv4 import ( + apply_dsv4_modelslim_patch, + ) apply_afd_ascend_config_patch_if_needed() if not apply_afd_mla_graph_patch(): raise RuntimeError( "AFD NPU MLA graph patch requires the vLLM-Ascend MLA resolver", ) + if not apply_dsv4_modelslim_patch(): + raise RuntimeError( + "AFD NPU DSV4 ModelSlim patch requires vLLM-Ascend ModelSlim", + ) _PATCHES_APPLIED = True diff --git a/afd_plugin/compat/patches/engine_core.py b/afd_plugin/compat/patches/engine_core.py index 2e689728..f34e8e78 100644 --- a/afd_plugin/compat/patches/engine_core.py +++ b/afd_plugin/compat/patches/engine_core.py @@ -42,6 +42,14 @@ def __init__( executor_fail_callback: Callable | None = None, include_finished_set: bool = False, ): + # Install DSV4's Ascend KV grouping compatibility before an EngineCore can + # ask vLLM to construct its cache configuration. + from afd_plugin.compat.patches.npu.deepseek_v4_kv_cache import ( + apply_deepseek_v4_kv_cache_patches, + ) + + apply_deepseek_v4_kv_cache_patches() + # ### PATCH START: AFD FFN EngineCore daemon initialization # FFN ranks are connector daemons, so stop EngineCore initialization after # executor construction instead of setting up KV cache and scheduler state. @@ -240,6 +248,14 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: return _AFDFFNKVCacheConfig() # ### PATCH END: AFD FFN late-loaded KV cache bypass + # vLLM-Ascend can rebind its grouping helper while worker modules import. + # Reinstall at the call site used to build the target cache configuration. + from afd_plugin.compat.patches.npu.deepseek_v4_kv_cache import ( + apply_deepseek_v4_kv_cache_patches, + ) + + apply_deepseek_v4_kv_cache_patches() + start = time.time() core_module.register_all_kvcache_specs(vllm_config) diff --git a/afd_plugin/compat/patches/npu/deepseek_v4_kv_cache.py b/afd_plugin/compat/patches/npu/deepseek_v4_kv_cache.py new file mode 100644 index 00000000..5a3795d2 --- /dev/null +++ b/afd_plugin/compat/patches/npu/deepseek_v4_kv_cache.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepSeek-V4 KV-cache compatibility for the Ascend v0.26 runtime. + +The Ascend hybrid-cache grouping helper needs every DSV4 MLA/SWA page size +when it builds a packed cache layout. Without that normalization its output +can expose multiple descriptors for one packed backing allocation to the +Ascend allocator, which then allocates each descriptor independently. +""" + +from __future__ import annotations + +import importlib +import logging +from dataclasses import replace +from typing import Any, Callable + +logger = logging.getLogger(__name__) + +_GROUPING_PATCHED = False +_ALLOCATOR_PATCHED = False + + +def _is_uniform_mla_group(group: Any) -> bool: + specs = getattr(group, "kv_cache_specs", None) + return bool(specs and hasattr(group, "get_page_sizes")) and all( + hasattr(spec, "compress_ratio") and hasattr(spec, "page_size_bytes") + for spec in specs.values() + ) + + +def apply_deepseek_v4_hybrid_kv_cache_group_patch() -> bool: + """Normalize DSV4 hybrid KV page-size buckets before planning. + + This is independent of speculative decoding: target-only DSV4 has the + same full-MLA and sliding-window cache-spec topology. + """ + + global _GROUPING_PATCHED + try: + import vllm.v1.core.kv_cache_utils as kv_cache_utils + except (ImportError, ModuleNotFoundError): + return True + + try: + importlib.import_module("vllm_ascend.patch.platform.patch_kv_cache_utils") + except (ImportError, ModuleNotFoundError): + pass + + original = getattr(kv_cache_utils, "_get_kv_cache_groups_uniform_groups", None) + if original is None: + return True + if getattr(original, "_afd_dsv4_page_size_patch", False): + _GROUPING_PATCHED = True + return True + + def get_kv_cache_groups_uniform_groups( + grouped_specs: list[Any], + _original: Callable[..., Any] = original, + ) -> Any: + if len(grouped_specs) <= 2 or not all( + _is_uniform_mla_group(grouped_specs[index]) for index in (0, 1) + ): + return _original(grouped_specs) + + full_mla_spec = grouped_specs[0] + page_sizes = set(full_mla_spec.get_page_sizes()) + for group in grouped_specs[1:]: + page_sizes.update(group.get_page_sizes()) + if page_sizes.issubset(set(full_mla_spec.get_page_sizes())): + return _original(grouped_specs) + + spec_type = type(full_mla_spec) + + class FullMLAWithAllPageSizes(spec_type): + def get_page_sizes(self) -> list[int]: + return sorted(page_sizes) + + patched_specs = list(grouped_specs) + patched_specs[0] = FullMLAWithAllPageSizes( + block_size=full_mla_spec.block_size, + kv_cache_specs=full_mla_spec.kv_cache_specs, + ) + logger.info( + "AFD DSV4 normalized Ascend KV page buckets: %s -> %s", + full_mla_spec.get_page_sizes(), + sorted(page_sizes), + ) + return _original(patched_specs) + + get_kv_cache_groups_uniform_groups._afd_dsv4_page_size_patch = True # type: ignore[attr-defined] + kv_cache_utils._get_kv_cache_groups_uniform_groups = get_kv_cache_groups_uniform_groups + _GROUPING_PATCHED = True + return True + + +def apply_deepseek_v4_ascend_allocator_patch() -> bool: + """Let Ascend allocate DSV4 state-only cache entries. + + Ascend's allocator recognizes compressed cache entries by an ``attn`` + substring. DSV4 compressor/SWA state entries do not have it, so alias + them internally only while the upstream allocator constructs the tensors. + """ + + global _ALLOCATOR_PATCHED + try: + import vllm_ascend.worker.model_runner_v1 as model_runner_module + except (ImportError, ModuleNotFoundError): + return True + + runner_cls = getattr(model_runner_module, "NPUModelRunner", None) + original = getattr(runner_cls, "_allocate_kv_cache_tensors", None) + if original is None: + return True + if getattr(original, "_afd_dsv4_allocator_patch", False): + _ALLOCATOR_PATCHED = True + return True + + def allocate_kv_cache_tensors( + self: Any, + kv_cache_config: Any, + _original: Callable[..., Any] = original, + ) -> Any: + state_names: set[str] = set() + for group in kv_cache_config.kv_cache_groups: + specs = getattr(group.kv_cache_spec, "kv_cache_specs", None) + if specs is None: + continue + for name, spec in specs.items(): + if type(spec).__name__ == "AscendSlidingWindowMLASpec" and "attn" not in name: + state_names.add(name) + if not state_names: + return _original(self, kv_cache_config) + + aliases = {name: f"{name}.afd_attn_cache" for name in state_names} + alias = lambda name: aliases.get(name, name) + patched_groups = [] + for group in kv_cache_config.kv_cache_groups: + spec = group.kv_cache_spec + specs = getattr(spec, "kv_cache_specs", None) + if specs is not None: + spec = replace(spec, kv_cache_specs={alias(name): value for name, value in specs.items()}) + patched_groups.append(replace(group, layer_names=[alias(name) for name in group.layer_names], kv_cache_spec=spec)) + patched_tensors = [ + replace(tensor, shared_by=[alias(name) for name in tensor.shared_by]) + for tensor in kv_cache_config.kv_cache_tensors + ] + patched_config = replace( + kv_cache_config, kv_cache_groups=patched_groups, kv_cache_tensors=patched_tensors + ) + raw_tensors = _original(self, patched_config) + return {aliases.get(name, name): tensor for name, tensor in raw_tensors.items()} + + allocate_kv_cache_tensors._afd_dsv4_allocator_patch = True # type: ignore[attr-defined] + runner_cls._allocate_kv_cache_tensors = allocate_kv_cache_tensors + _ALLOCATOR_PATCHED = True + return True + + +def apply_deepseek_v4_kv_cache_patches() -> None: + """Install the DSV4-only Ascend KV compatibility patches.""" + + apply_deepseek_v4_hybrid_kv_cache_group_patch() + apply_deepseek_v4_ascend_allocator_patch() diff --git a/afd_plugin/compat/patches/npu/modelslim_dsv4.py b/afd_plugin/compat/patches/npu/modelslim_dsv4.py new file mode 100644 index 00000000..887c097b --- /dev/null +++ b/afd_plugin/compat/patches/npu/modelslim_dsv4.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Compatibility mapping for DSV4 ModelSlim quantization. + +DSV4 checkpoints use names such as ``attn`` and ``ffn`` while the runtime +uses the corresponding Hugging Face names. Resolve a runtime prefix to the +matching checkpoint key without changing the quantization description. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +_PATCH_ATTR = "_afd_dsv4_modelslim_patch" + +_DSV4_RUNTIME_TO_CHECKPOINT = ( + (".self_attn.", ".attn."), + (".post_attention_layernorm.", ".ffn_norm."), + (".input_layernorm.", ".attn_norm."), + (".mlp.", ".ffn."), + (".gate_proj.", ".w1."), + (".up_proj.", ".w3."), + (".down_proj.", ".w2."), +) + + +def _replace_once(value: str, old: str, new: str) -> str: + return value.replace(old, new, 1) + + +def _runtime_name_variants(prefix: str) -> list[str]: + """Return DSV4 runtime/checkpoint spelling variants for ``prefix``.""" + + variants = [prefix] + for old, new in _DSV4_RUNTIME_TO_CHECKPOINT: + variants.append(_replace_once(prefix, old, new)) + return variants + + +def _has_quant_weight( + quant_description: Mapping[str, Any], + prefix: str, + packed_modules_mapping: Mapping[str, list[str]], +) -> bool: + """Check a linear/MoE prefix without indexing a missing description key.""" + + projection = prefix.rsplit(".", 1)[-1] + if projection in packed_modules_mapping: + return all( + f"{prefix.replace(projection, shard)}.weight" in quant_description + for shard in packed_modules_mapping[projection] + ) + return f"{prefix}.weight" in quant_description + + +def _resolve_dsv4_quant_prefix( + prefix: str, + quant_description: Mapping[str, Any], + packed_modules_mapping: Mapping[str, list[str]], +) -> str: + """Resolve a DSV4 runtime prefix against ModelSlim description keys.""" + + seen: set[str] = set() + for candidate in _runtime_name_variants(prefix): + if candidate in seen: + continue + seen.add(candidate) + if _has_quant_weight(quant_description, candidate, packed_modules_mapping): + return candidate + return prefix + + +def _patch_dsv4_packed_mapping(modelslim_module: Any) -> None: + """Use DSV4 checkpoint shard names for fused runtime modules.""" + + mapping = modelslim_module.packed_modules_model_mapping.setdefault( + "deepseek_v4", {} + ) + mapping.update( + { + "gate_up_proj": ["w1", "w3"], + "fused_wqa_wkv": ["wq_a", "wkv"], + "experts": ["experts.0.w1", "experts.0.w3", "experts.0.w2"], + } + ) + + +def apply_dsv4_modelslim_patch() -> bool: + """Install the idempotent vLLM-Ascend ModelSlim prefix resolver.""" + + try: + from vllm_ascend.quantization import modelslim_config + except ImportError: + return False + + if getattr(modelslim_config, _PATCH_ATTR, None) is not None: + return True + + original_mapper = modelslim_config.AscendModelSlimConfig.quant_prefix_mapper + _patch_dsv4_packed_mapping(modelslim_config) + + def quant_prefix_mapper(self, model_type: str, prefix: str) -> str: + mapped_prefix = original_mapper(self, model_type, prefix) + if model_type != "deepseek_v4": + return mapped_prefix + + return _resolve_dsv4_quant_prefix( + mapped_prefix, + getattr(self, "quant_description", {}), + getattr(self, "packed_modules_mapping", {}), + ) + + modelslim_config.AscendModelSlimConfig.quant_prefix_mapper = quant_prefix_mapper + setattr(modelslim_config, _PATCH_ATTR, original_mapper) + return True + + +__all__ = ["apply_dsv4_modelslim_patch", "_resolve_dsv4_quant_prefix"] diff --git a/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index 6c982d71..b5bfb079 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -128,6 +128,7 @@ class AFDControlPayload: dp_metadata_list: dict[int, AFDDPMetadata] is_graph_capturing: bool is_warmup: bool + input_ids_by_stage: dict[int, list[int]] = field(default_factory=dict) def __post_init__(self) -> None: self.dp_metadata_list = { @@ -139,6 +140,14 @@ def __post_init__(self) -> None: else dp_metadata for stage_idx, dp_metadata in self.dp_metadata_list.items() } + self.input_ids_by_stage = { + int(stage_idx): torch.as_tensor( + input_ids, + dtype=torch.int64, + device="cpu", + ).flatten().tolist() + for stage_idx, input_ids in self.input_ids_by_stage.items() + } @dataclass(slots=True) @@ -319,10 +328,14 @@ def encode_control_payload(payload: AFDControlPayload) -> bytes: """ metadata_payload: dict[str, dict[str, int | list[int]]] = {} for stage_idx, dp_metadata in payload.dp_metadata_list.items(): - metadata_payload[str(int(stage_idx))] = { + stage_payload: dict[str, int | list[int]] = { "num_tokens_across_dp_cpu": dp_metadata.num_tokens_across_dp_cpu.tolist(), "max_tokens_across_dp_cpu": _to_int(dp_metadata.max_tokens_across_dp_cpu), } + input_ids = payload.input_ids_by_stage.get(int(stage_idx)) + if input_ids is not None: + stage_payload["input_ids"] = list(input_ids) + metadata_payload[str(int(stage_idx))] = stage_payload wire_payload = { "dp_metadata_list": metadata_payload, @@ -352,10 +365,16 @@ def decode_control_payload(payload_bytes: bytes) -> AFDControlPayload: ) for stage_idx, metadata in payload["dp_metadata_list"].items() } + input_ids_by_stage = { + int(stage_idx): [int(value) for value in metadata["input_ids"]] + for stage_idx, metadata in payload["dp_metadata_list"].items() + if "input_ids" in metadata + } return AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=bool(payload.get("is_graph_capturing", False)), is_warmup=bool(payload.get("is_warmup", False)), + input_ids_by_stage=input_ids_by_stage, ) diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index 00dd128a..c7832c64 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -27,10 +27,10 @@ from __future__ import annotations -import os from collections.abc import Mapping from dataclasses import dataclass from datetime import timedelta +from time import sleep from typing import TYPE_CHECKING, Any, Final import torch @@ -65,6 +65,7 @@ AFD_ASYNC_CAM_GROUP_NAME = "afd_async_cam" CAM_COMM_ID = 0 ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp" +SHARED_FFN_POOL_CONFIG_KEY = "shared_ffn_pool" ASYNC_MOE_NUM_STAGES = 2 ASYNC_MOE_REQUEST_SPLIT = "request" ASYNC_MOE_TOKEN_SPLIT = "token" @@ -73,6 +74,7 @@ { "dynamicQuant", "attn_ranks_per_dp", + "shared_ffn_pool", "async_moe_ubatching", "async_moe_num_ubatches", "async_moe_split", @@ -97,6 +99,7 @@ class AFDAsyncExtraInfo(ConnectorExtraInfo): dynamic_quant: int = 0 attn_ranks_per_dp: int = 1 + shared_ffn_pool: bool = False async_moe_ubatching: bool = False async_moe_num_ubatches: int = ASYNC_MOE_NUM_STAGES async_moe_split: str = ASYNC_MOE_REQUEST_SPLIT @@ -128,6 +131,10 @@ def from_mapping(cls, raw: Mapping[str, Any] | None) -> AFDAsyncExtraInfo: raw.get("attn_ranks_per_dp", 1), field_name="attn_ranks_per_dp", ), + shared_ffn_pool=coerce_extra_bool( + raw.get(SHARED_FFN_POOL_CONFIG_KEY, False), + field_name=SHARED_FFN_POOL_CONFIG_KEY, + ), async_moe_ubatching=coerce_extra_bool( raw.get("async_moe_ubatching", False), field_name="async_moe_ubatching", @@ -146,6 +153,7 @@ def to_mapping(self) -> dict[str, Any]: return { "dynamicQuant": self.dynamic_quant, "attn_ranks_per_dp": self.attn_ranks_per_dp, + SHARED_FFN_POOL_CONFIG_KEY: self.shared_ffn_pool, "async_moe_ubatching": self.async_moe_ubatching, "async_moe_num_ubatches": self.async_moe_num_ubatches, "async_moe_split": self.async_moe_split, @@ -174,6 +182,7 @@ class AFDAsyncTransferState(AFDTransferState): dynamic_scales: Tensor | None = None expand_x_shared: Tensor | None = None dynamic_scales_shared: Tensor | None = None + cam_dp_group_index: int = 0 @dataclass(slots=True) @@ -188,6 +197,7 @@ class AFDAsyncFFNWorkItem: num_tokens: int total_num_tokens: int shared_num_tokens: int + cam_dp_group_index: int = 0 @dataclass(frozen=True, slots=True) @@ -196,10 +206,13 @@ class AFDAsyncTopology: role: str role_rank: int + dp_group_index: int + num_dp_groups: int world_rank: int attn_size: int ffn_size: int expert_per_rank: int + shared_ffn_pool: bool = False @property def world_size(self) -> int: @@ -207,6 +220,19 @@ def world_size(self) -> int: return self.attn_size + self.ffn_size +@dataclass(slots=True) +class _CAMGroupContext: + """One DP-scoped CAM communicator owned by this connector process.""" + + topology: AFDAsyncTopology + process_group_name: str + rendezvous_port: int + cam_pg: ProcessGroup | None = None + group_name: str = "" + comm_args: Tensor | None = None + placeholder: Tensor | None = None + + class CAMAsyncAFDConnector(AFDConnectorBase): """CAM-backed asynchronous connector for Ascend NPU AFD. @@ -257,18 +283,156 @@ def __init__( afd_config, role_rank, num_routed_experts=self.num_routed_experts, + attn_ranks_per_dp=self.tp_size, + shared_ffn_pool=self.extra_info.shared_ffn_pool, ) self.world_rank = self.topology.world_rank self.attn_size = self.topology.attn_size self.ffn_size = self.topology.ffn_size self.expert_per_rank = self.topology.expert_per_rank + self._process_group_name = _cam_process_group_name(self.topology) + self._rendezvous_port = _cam_rendezvous_port( + self.afd_config.port, + self.topology, + ) self.comm_args: Tensor | None = None self._placeholder: Tensor | None = None + self._cam_group_contexts = self._build_cam_group_contexts() + self._scheduler_store: Any | None = None + self._scheduler_next_sequence = 1 + self._scheduler_consumed_by_dp = [0] * self.topology.num_dp_groups self._pending_attention_payloads: dict[ int, list[tuple[AFDTransferContext, Tensor, Tensor]], ] = {} + def _build_cam_group_contexts(self) -> dict[int, _CAMGroupContext]: + """Build the CAM groups joined by this process. + + In shared-pool mode an Attention rank joins only its own DP group, but + every FFN rank joins every group. This preserves independent Attention + queues while keeping one DP8/TP1/EP8 FFN expert partition. + """ + group_indices = [self.topology.dp_group_index] + if self.extra_info.shared_ffn_pool and self.afd_config.role == "ffn": + group_indices = list(range(self.topology.num_dp_groups)) + + contexts: dict[int, _CAMGroupContext] = {} + for dp_group_index in group_indices: + topology = build_async_topology( + self.afd_config, + self.role_rank, + num_routed_experts=self.num_routed_experts, + attn_ranks_per_dp=self.tp_size, + shared_ffn_pool=self.extra_info.shared_ffn_pool, + dp_group_index=dp_group_index, + ) + contexts[dp_group_index] = _CAMGroupContext( + topology=topology, + process_group_name=_cam_process_group_name(topology), + rendezvous_port=_cam_rendezvous_port(self.afd_config.port, topology), + ) + return contexts + + def _cam_context(self, dp_group_index: int | None = None) -> _CAMGroupContext: + if dp_group_index is None: + dp_group_index = self.topology.dp_group_index + try: + return self._cam_group_contexts[dp_group_index] + except KeyError as exc: + raise RuntimeError( + "CAMAsyncAFDConnector has no CAM communicator for " + f"Attention DP group {dp_group_index}" + ) from exc + + def _set_legacy_context_fields(self) -> None: + """Keep the historical public fields bound to the primary CAM group.""" + context = self._cam_context() + self.cam_pg = context.cam_pg + self.group_name = context.group_name + self.comm_args = context.comm_args + self._placeholder = context.placeholder + + @property + def uses_shared_ffn_pool(self) -> bool: + """Whether one FFN TP/EP partition serves multiple Attention DPs.""" + return self.extra_info.shared_ffn_pool and self.topology.num_dp_groups > 1 + + def _init_shared_ffn_scheduler(self) -> None: + """Create the tiny control plane used only to select a ready CAM group. + + CAM transports activations and completion data, but its blocking recv + has no API to ask which of several communicators has pending work. A + TCPStore queue lets the FFN leader select a ready Attention DP group; + all FFN ranks then receive from the same CAM communicator. + """ + if not self.uses_shared_ffn_pool: + return + import torch.distributed as dist + + port = self.afd_config.port + self.topology.num_dp_groups + if not 1 <= port <= 65535: + raise ValueError( + "AFD shared-FFN scheduler TCPStore port is outside the valid " + f"range: base_port={self.afd_config.port}, " + f"num_dp_groups={self.topology.num_dp_groups}" + ) + is_server = self.afd_config.role == "ffn" and self.role_rank == 0 + self._scheduler_store = dist.TCPStore( + self.afd_config.host, + port, + world_size=1, + is_master=is_server, + # FFN workers enter their connector-driven receive loop before an + # API request necessarily arrives. This store guards an idle + # work queue, not a collective rendezvous, so a 30-minute timeout + # would incorrectly kill a healthy service after a quiet period. + timeout=timedelta(hours=24), + wait_for_workers=False, + ) + if is_server: + for dp_group_index in range(self.topology.num_dp_groups): + self._scheduler_store.set( + _scheduler_ready_key(dp_group_index), + "0", + ) + + def _notify_shared_ffn_work_ready(self) -> None: + if not self.uses_shared_ffn_pool or self.topology.world_rank != 0: + return + assert self._scheduler_store is not None + self._scheduler_store.add( + _scheduler_ready_key(self.topology.dp_group_index), + 1, + ) + + def _next_shared_ffn_dp_group(self) -> int: + """Return the next ready Attention DP group on every FFN rank.""" + if not self.uses_shared_ffn_pool: + return self.topology.dp_group_index + assert self._scheduler_store is not None + + schedule_key = _scheduler_schedule_key(self._scheduler_next_sequence) + if self.role_rank == 0: + while True: + for dp_group_index in range(self.topology.num_dp_groups): + ready = int( + self._scheduler_store.get( + _scheduler_ready_key(dp_group_index), + ).decode(), + ) + if ready > self._scheduler_consumed_by_dp[dp_group_index]: + self._scheduler_consumed_by_dp[dp_group_index] += 1 + self._scheduler_store.set(schedule_key, str(dp_group_index)) + break + else: + sleep(0.001) + continue + break + dp_group_index = int(self._scheduler_store.get(schedule_key).decode()) + self._scheduler_next_sequence += 1 + return dp_group_index + @property def is_initialized(self) -> bool: """Return whether the CAM HCCL group and operator buffers are ready.""" @@ -285,31 +449,35 @@ def init_afd_connector(self) -> None: return ensure_cam_async_ops_available() - self.cam_pg = init_afd_process_group( - backend="hccl", - init_method=f"tcp://{self.afd_config.host}:{self.afd_config.port}", - world_size=self.topology.world_size, - rank=self.world_rank, - group_name=AFD_ASYNC_CAM_GROUP_NAME, - timeout=timedelta(minutes=30), - ) - backend = self.cam_pg._get_backend(torch.device("npu")) - self.group_name = str(backend.get_hccl_comm_name(self.world_rank)) device = f"npu:{self.local_rank}" - self.comm_args = torch.empty((1,), dtype=torch.float16, device=device) - self._placeholder = torch.empty( - (1,), - dtype=torch.bfloat16, - device=device, - ) + for context in self._cam_group_contexts.values(): + topology = context.topology + context.cam_pg = init_afd_process_group( + backend="hccl", + init_method=f"tcp://{self.afd_config.host}:{context.rendezvous_port}", + world_size=topology.world_size, + rank=topology.world_rank, + group_name=context.process_group_name, + timeout=timedelta(minutes=30), + ) + backend = context.cam_pg._get_backend(torch.device("npu")) + context.group_name = str(backend.get_hccl_comm_name(topology.world_rank)) + context.comm_args = torch.empty((1,), dtype=torch.float16, device=device) + context.placeholder = torch.empty((1,), dtype=torch.bfloat16, device=device) + self._set_legacy_context_fields() + self._init_shared_ffn_scheduler() self._initialized = True def close(self) -> None: """Destroy the HCCL process group and clear pending transfer states.""" - if self.cam_pg is not None: - import torch.distributed as dist - - dist.destroy_process_group(self.cam_pg) + import torch.distributed as dist + + for context in self._cam_group_contexts.values(): + if context.cam_pg is not None: + dist.destroy_process_group(context.cam_pg) + context.cam_pg = None + context.comm_args = None + context.placeholder = None self.cam_pg = None self.comm_args = None self._placeholder = None @@ -327,17 +495,23 @@ def recv_ffn_work_item( *, stage_idx: int, max_num_tokens: int, + expected_layer_idx: int, ) -> AFDAsyncFFNWorkItem: """Receive and normalize one connector-driven FFN dispatch item. CAM metadata supplies the actual layer and routed/shared token counts; returned tensors are sliced from operator capacity to those counts. """ + recv_kwargs: dict[str, Any] = { + "stage_idx": stage_idx, + "layer_idx": 0, + "batch_size": max(1, self.max_seq_len or max_num_tokens), + "ubatch_idx": stage_idx, + } + if self.uses_shared_ffn_pool: + recv_kwargs["cam_dp_group_index"] = self._next_shared_ffn_dp_group() recv_output = self.recv_attn_output( - stage_idx=stage_idx, - layer_idx=0, - batch_size=max(1, self.max_seq_len or max_num_tokens), - ubatch_idx=stage_idx, + **recv_kwargs, ) context = recv_output.context metadata = context.metadata @@ -349,7 +523,28 @@ def recv_ffn_work_item( "TokenNums_Rankid_Layeridx from async_dispatch_recv", ) total_num_tokens = max(1, int(token_nums_rankid_layeridx[0].item())) - layer_idx = int(token_nums_rankid_layeridx[2].item()) + # A zero in the header's leading count is valid for an FFN rank that + # received no routed tokens in this dispatch. Such ranks must still + # enter combine-send (the zero-token fallback below supplies its + # placeholder) so every participant completes the CAM collective. + # CAM 209 currently returns zero in the documented layer-index slot + # for every dispatch-recv result. The connector-driven FFN runner + # already has the authoritative decoder-layer order, so consuming this + # stale field makes every DSV4 work item execute FFN layer 0. CAM + # combine-send also uses this header to match the completion to its + # original dispatch, so merely fixing the local model call is not + # sufficient: all later layers would still be combined as layer 0. + # CAM returns the true decoder-layer index in the dispatch-recv header + # (field 2 of token_nums_rankid_layeridx). Use it for both the local + # FFN layer computation and the combine-send header so the completion + # matches the attention-side combine-recv. Do NOT overwrite field 2 + # with expected_layer_idx: the shared-FFN-pool scheduler alternates + # DP groups per layer, so expected_layer_idx is the busy-loop + # iteration index, not the true decoder layer. + layer_idx = max(0, int(token_nums_rankid_layeridx[2].item())) + token_nums_rankid_layeridx = token_nums_rankid_layeridx.clone() + states.token_nums_rankid_layeridx = token_nums_rankid_layeridx + cam_dp_group_index = states.cam_dp_group_index expert_token_nums_shared = states.expert_token_nums_shared if expert_token_nums_shared is None: @@ -393,6 +588,7 @@ def recv_ffn_work_item( num_tokens=num_tokens, total_num_tokens=total_num_tokens, shared_num_tokens=shared_num_tokens, + cam_dp_group_index=cam_dp_group_index, ) def _send_ffn_output_payload( @@ -506,27 +702,8 @@ def send_attn_output( (context, topk_ids, topk_weights), ) - _log_cam_op_values( - "async_dispatch_send", - "inputs", - hidden_states=hidden_states, - topk_ids=topk_ids, - comm_args=self.comm_args, - comm_id=self.comm_id, - max_seq_len=self.max_seq_len, - batch_size=states.batch_size, - hidden_size=states.hidden_size, - topk=states.topk, - ffn_size=self.ffn_size, - attn_size=self.attn_size, - expert_per_rank=self.expert_per_rank, - rank=self.world_rank, - world_size=self.topology.world_size, - layer_idx=states.layer_idx, - tp_size=self.tp_size, - dynamic_quant=self.dynamic_quant, - group_name=self.group_name, - ) + self._notify_shared_ffn_work_ready() + torch.ops.umdk_cam_op_lib.async_dispatch_send( hidden_states, topk_ids, @@ -546,6 +723,9 @@ def send_attn_output( self.dynamic_quant, self.group_name, ) + npu = getattr(torch, "npu", None) + if npu is not None: + npu.synchronize() return None def recv_ffn_output( @@ -609,24 +789,6 @@ def recv_ffn_output( topk=states.topk, ) placeholder = ref_tensor.new_empty((1,)) - _log_cam_op_values( - "async_combine_recv", - "inputs", - placeholder=placeholder, - topk_ids=topk_ids, - topk_weights=topk_weights, - comm_args=self.comm_args, - comm_id=self.comm_id, - batch_size=states.batch_size, - hidden_size=states.hidden_size, - topk=states.topk, - ffn_size=self.ffn_size, - attn_size=self.attn_size, - expert_per_rank=self.expert_per_rank, - rank=self.world_rank, - world_size=self.topology.world_size, - group_name=self.group_name, - ) output = torch.ops.umdk_cam_op_lib.async_combine_recv( placeholder, topk_ids, @@ -643,7 +805,9 @@ def recv_ffn_output( self.topology.world_size, self.group_name, ) - _log_cam_op_values("async_combine_recv", "outputs", output=output) + npu = getattr(torch, "npu", None) + if npu is not None: + npu.synchronize() return output def recv_attn_output( @@ -676,40 +840,27 @@ def recv_attn_output( metadata=metadata, states=states, ) - placeholder = kwargs.get("placeholder", self._placeholder) - _log_cam_op_values( - "async_dispatch_recv", - "inputs", - placeholder=placeholder, - comm_args=self.comm_args, - comm_id=self.comm_id, - batch_size=states.batch_size, - hidden_size=states.hidden_size, - topk=states.topk, - ffn_size=self.ffn_size, - attn_size=self.attn_size, - expert_per_rank=self.expert_per_rank, - rank=self.world_rank, - world_size=self.topology.world_size, - tp_size=self.tp_size, - dynamic_quant=self.dynamic_quant, - group_name=self.group_name, + cam_dp_group_index = int( + kwargs.get("cam_dp_group_index", self.topology.dp_group_index), ) + cam_context = self._cam_context(cam_dp_group_index) + topology = cam_context.topology + placeholder = kwargs.get("placeholder", cam_context.placeholder) outputs = torch.ops.umdk_cam_op_lib.async_dispatch_recv( placeholder, - self.comm_args, + cam_context.comm_args, self.comm_id, states.batch_size, states.hidden_size, states.topk, - self.ffn_size, - self.attn_size, - self.expert_per_rank, - self.world_rank, - self.topology.world_size, + topology.ffn_size, + topology.attn_size, + topology.expert_per_rank, + topology.world_rank, + topology.world_size, self.tp_size, self.dynamic_quant, - self.group_name, + cam_context.group_name, ) ( hidden_states, @@ -720,23 +871,13 @@ def recv_attn_output( expert_token_nums, expert_token_nums_shared, ) = outputs - _log_cam_op_values( - "async_dispatch_recv", - "outputs", - hidden_states=hidden_states, - expand_x_shared=expand_x_shared, - dynamic_scales=dynamic_scales, - dynamic_scales_shared=dynamic_scales_shared, - token_nums_rankid_layeridx=token_nums_rankid_layeridx, - expert_token_nums=expert_token_nums, - expert_token_nums_shared=expert_token_nums_shared, - ) states.token_nums_rankid_layeridx = token_nums_rankid_layeridx states.expert_token_nums_shared = expert_token_nums_shared states.group_list = expert_token_nums states.dynamic_scales = dynamic_scales states.expand_x_shared = expand_x_shared states.dynamic_scales_shared = dynamic_scales_shared + states.cam_dp_group_index = cam_dp_group_index return AFDA2FTransferPayload( hidden_states=hidden_states, context=context, @@ -755,6 +896,8 @@ def send_ffn_output( """ self._require_initialized() states = _require_async_transfer_state(context) + cam_context = self._cam_context(states.cam_dp_group_index) + topology = cam_context.topology expand_x_shared = kwargs.get("expand_x_shared") if expand_x_shared is None: expand_x_shared = ffn_output @@ -766,43 +909,23 @@ def send_ffn_output( "AFD async CAM combine send requires " "TokenNums_Rankid_Layeridx from async_dispatch_recv", ) - _log_cam_op_values( - "async_combine_send", - "inputs", - ffn_output=ffn_output, - expand_x_shared=expand_x_shared, - comm_args=self.comm_args, - token_nums_rankid_layeridx=token_nums_rankid_layeridx, - comm_id=self.comm_id, - batch_size=states.batch_size, - hidden_size=states.hidden_size, - topk=states.topk, - ffn_size=self.ffn_size, - attn_size=self.attn_size, - expert_per_rank=self.expert_per_rank, - rank=self.world_rank, - world_size=self.topology.world_size, - tp_size=self.tp_size, - group_name=self.group_name, - ) torch.ops.umdk_cam_op_lib.async_combine_send( ffn_output, expand_x_shared, - self.comm_args, + cam_context.comm_args, token_nums_rankid_layeridx, self.comm_id, states.batch_size, states.hidden_size, states.topk, - self.ffn_size, - self.attn_size, - self.expert_per_rank, - self.world_rank, - self.topology.world_size, + topology.ffn_size, + topology.attn_size, + topology.expert_per_rank, + topology.world_rank, + topology.world_size, self.tp_size, - self.group_name, + cam_context.group_name, ) - def _require_initialized(self) -> None: if not self._initialized: raise RuntimeError("CAMAsyncAFDConnector is not initialized") @@ -820,71 +943,71 @@ def _require_async_transfer_state( return states -_CAM_LOG_SKIPPED_ARGS = frozenset({"comm_args", "comm_id", "group_name"}) -_CAM_OP_IO_LOG_ENV = "AFD_CAM_OP_IO_LOG" - - -def _log_cam_op_values(op_name: str, label: str, **kwargs: object) -> None: - if os.environ.get(_CAM_OP_IO_LOG_ENV, "").lower() not in { - "1", - "true", - "yes", - "on", - }: - return - lines: list[str] = [] - for name, value in kwargs.items(): - if name in _CAM_LOG_SKIPPED_ARGS: - continue - if isinstance(value, Tensor): - description = f"Tensor(dtype={value.dtype}, shape={tuple(value.shape)})" - if name == "token_nums_rankid_layeridx": - try: - first5: object = value.detach().flatten()[:5].cpu().tolist() - except Exception as exc: # pragma: no cover - defensive logging helper - first5 = f"" - description = f"{description}, first5={first5!r}" - else: - description = repr(value) - lines.append(f" {name}={description}") - logger.warning("AFD CAM %s %s:\n%s", op_name, label, "\n".join(lines)) - - def build_async_topology( afd_config: AFDConfig, role_rank: int, *, num_routed_experts: int | None = None, + attn_ranks_per_dp: int | None = None, + shared_ffn_pool: bool = False, + dp_group_index: int | None = None, ) -> AFDAsyncTopology: - """Validate role-local rank settings and derive the CAM HCCL world rank. + """Derive the DP-scoped CAM HCCL world for one role-local rank. - The world is Attention-first: Attention role rank ``i`` maps to world rank - ``i`` and FFN role rank ``j`` maps to - ``num_attention_ranks + j``. Routed experts are distributed across FFN - ranks using a ceiling division; production model layouts should keep the - routed-expert count divisible by the FFN rank count. + CAM does not carry a DP queue identifier. Multiple Attention DP replicas + must therefore use independent HCCL communicators rather than sharing one + global CAM world. Each local world is Attention-first: + ``[A(dp, 0), ..., A(dp, tp - 1), F(dp, 0), ..., F(dp, n - 1)]``. """ - attn_size = afd_config.num_attention_ranks - ffn_size = afd_config.num_ffn_ranks - if attn_size <= 0 or ffn_size <= 0: + global_attn_size = afd_config.num_attention_ranks + global_ffn_size = afd_config.num_ffn_ranks + if global_attn_size <= 0 or global_ffn_size <= 0: raise ValueError("AFD async topology sizes must be positive") if role_rank < 0: raise ValueError(f"AFD async role rank must be non-negative, got {role_rank}") + attn_size = ( + global_attn_size if attn_ranks_per_dp is None else attn_ranks_per_dp + ) + if not isinstance(attn_size, int) or isinstance(attn_size, bool) or attn_size <= 0: + raise ValueError("attn_ranks_per_dp must be a positive integer") + if global_attn_size % attn_size != 0: + raise ValueError( + "num_attention_ranks must be divisible by attn_ranks_per_dp, got " + f"{global_attn_size} and {attn_size}" + ) + num_dp_groups = global_attn_size // attn_size + if not shared_ffn_pool and global_ffn_size % num_dp_groups != 0: + raise ValueError( + "num_ffn_ranks must be divisible by the CAM DP group count, got " + f"{global_ffn_size} and {num_dp_groups}" + ) + ffn_size = global_ffn_size if shared_ffn_pool else global_ffn_size // num_dp_groups + if afd_config.role == "attention": - if role_rank >= attn_size: + if role_rank >= global_attn_size: raise ValueError( "Attention role rank must be within attention size " - f"(rank={role_rank}, size={attn_size})", + f"(rank={role_rank}, size={global_attn_size})", ) - world_rank = role_rank + resolved_dp_group_index, world_rank = divmod(role_rank, attn_size) elif afd_config.role == "ffn": - if role_rank >= ffn_size: + if role_rank >= global_ffn_size: raise ValueError( "FFN role rank must be within FFN size " - f"(rank={role_rank}, size={ffn_size})", + f"(rank={role_rank}, size={global_ffn_size})", ) - world_rank = attn_size + role_rank + if shared_ffn_pool: + resolved_dp_group_index = 0 if dp_group_index is None else dp_group_index + if not 0 <= resolved_dp_group_index < num_dp_groups: + raise ValueError( + "shared FFN CAM DP group index is outside the valid range: " + f"{resolved_dp_group_index} not in [0, {num_dp_groups})" + ) + world_rank = attn_size + role_rank + else: + resolved_dp_group_index, local_ffn_rank = divmod(role_rank, ffn_size) + world_rank = attn_size + local_ffn_rank else: raise ValueError(f"unknown AFD role {afd_config.role!r}") @@ -893,13 +1016,50 @@ def build_async_topology( return AFDAsyncTopology( role=afd_config.role, role_rank=role_rank, + dp_group_index=resolved_dp_group_index, + num_dp_groups=num_dp_groups, world_rank=world_rank, attn_size=attn_size, ffn_size=ffn_size, expert_per_rank=expert_per_rank, + shared_ffn_pool=shared_ffn_pool, ) +def _cam_process_group_name(topology: AFDAsyncTopology) -> str: + """Return the DP-scoped HCCL group name used by CAM operators.""" + + if topology.num_dp_groups == 1: + return AFD_ASYNC_CAM_GROUP_NAME + return f"{AFD_ASYNC_CAM_GROUP_NAME}_dp{topology.dp_group_index}" + + +def _cam_rendezvous_port(base_port: int, topology: AFDAsyncTopology) -> int: + """Return the TCPStore port for a DP-scoped CAM communicator. + + Every scoped communicator has a local rank zero. Reusing the same TCP + rendezvous endpoint would therefore make the DP-group leaders compete to + bind one TCPStore server. Keep all ranks in one CAM DP group on the same + endpoint, while assigning later groups consecutive ports. + """ + + port = base_port + topology.dp_group_index + if not 1 <= port <= 65535: + raise ValueError( + "AFD async CAM rendezvous port is outside the valid range: " + f"base_port={base_port}, dp_group_index={topology.dp_group_index}", + ) + return port + + +def _scheduler_ready_key(dp_group_index: int) -> str: + return f"afd_async_cam/scheduler/ready/{dp_group_index}" + + +def _scheduler_schedule_key(sequence: int) -> str: + return f"afd_async_cam/scheduler/next/{sequence}" + + def _validate_topk_payload( topk_ids: Tensor, topk_weights: Tensor | None, @@ -931,6 +1091,7 @@ def _validate_topk_payload( "AFDAsyncFFNWorkItem", "AFDAsyncTopology", "ATTN_RANKS_PER_DP_CONFIG_KEY", + "SHARED_FFN_POOL_CONFIG_KEY", "ASYNC_MOE_NUM_STAGES", "ASYNC_MOE_REQUEST_SPLIT", "ASYNC_MOE_TOKEN_SPLIT", diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index 32f25424..40ea5c4c 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p.py @@ -181,9 +181,9 @@ class _CAMP2PTopology: """Describe where one connector process sits in the communication groups. ``world_rank`` is the process number in the complete AFD group. - ``p2p_rank`` is its number in the smaller Gloo group used to exchange metadata. - For an Attention rank, ``dp_metadata_destinations`` lists the FFN ranks - that should receive its metadata. + ``p2p_rank`` is its number in the Gloo group used to exchange metadata. + Every Attention rank participates in that group so an FFN rank can receive + the token IDs for every Attention hidden-state slice it aggregates. """ role: str @@ -194,21 +194,27 @@ class _CAMP2PTopology: ffn_size: int min_size: int dp_metadata_destinations: tuple[int, ...] + dp_metadata_sources: tuple[int, ...] @property def p2p_world_size(self) -> int: - """Return the number of FFN and participating Attention metadata ranks.""" - return self.ffn_size + self.min_size + """Return the number of FFN and Attention metadata ranks.""" + return self.ffn_size + self.attention_size @property def participates_in_p2p_group(self) -> bool: """Return whether this process joins the Gloo DP-metadata group.""" - return self.world_rank < self.ffn_size or self.is_attn_top_min_size_rank + return True @property def is_attn_top_min_size_rank(self) -> bool: - """Return whether this is an Attention metadata-sender rank.""" - return self.ffn_size <= self.world_rank < self.ffn_size + self.min_size + """Return whether this is an Attention metadata-sender rank. + + The property name is retained for connector compatibility. Every + Attention rank sends one payload to the FFN rank that receives its + hidden-state slice. + """ + return self.role == "attention" class CAMP2pAFDConnector(AFDConnectorBase): @@ -265,6 +271,7 @@ def __init__( self.ratio = self.attn_size // self.ffn_size self.dst_list = list(self.topology.dp_metadata_destinations) self.dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata] = {} + self.input_ids_by_stage: dict[int, list[int]] = {} self.is_graph_capturing = False self.is_warmup = False self.scheduler_config = vllm_config.scheduler_config @@ -383,6 +390,7 @@ def close(self) -> None: dist.destroy_process_group(group) self.p2p_pg = None self.ffn_pg = None + self.input_ids_by_stage = {} self.afd_pg = None self.afd_pg_list = [] self.hccl_comm_name = "" @@ -636,6 +644,7 @@ def update_state_from_dp_metadata( ) -> None: connector = self.connector connector.dp_metadata_list = payload.dp_metadata_list + connector.input_ids_by_stage = payload.input_ids_by_stage connector.is_graph_capturing = payload.is_graph_capturing connector.is_warmup = payload.is_warmup @@ -662,12 +671,68 @@ def recv_dp_metadata_list(self) -> AFDControlPayload: connector = self.connector if connector.p2p_pg is None: raise RuntimeError("CAMP2P metadata process group is not initialized") - src = connector.p2p_rank % connector.min_size + connector.ffn_size - return recv_control_payload( - src=src, - group=connector.p2p_pg, - device=torch.device("cpu"), + payloads = tuple( + recv_control_payload( + src=src, + group=connector.p2p_pg, + device=torch.device("cpu"), + ) + for src in connector.topology.dp_metadata_sources ) + return _aggregate_camp2p_control_payloads(payloads) + + +def _aggregate_camp2p_control_payloads( + payloads: tuple[AFDControlPayload, ...], +) -> AFDControlPayload: + """Concatenate peer token IDs in CAMP2P hidden-state receive order. + + CAMP2P aggregates consecutive Attention role ranks on each FFN rank. The + control payload follows the same order, so DSV4 hash routing sees one + token ID per received hidden-state row, including every DBO ubatch stage. + """ + if not payloads: + raise RuntimeError("CAMP2P FFN control plane received no Attention payloads") + + reference = payloads[0] + stage_ids = tuple(sorted(reference.dp_metadata_list)) + for payload in payloads[1:]: + if tuple(sorted(payload.dp_metadata_list)) != stage_ids: + raise RuntimeError( + "CAMP2P Attention peers sent different control-plane stages" + ) + if ( + payload.is_graph_capturing != reference.is_graph_capturing + or payload.is_warmup != reference.is_warmup + ): + raise RuntimeError( + "CAMP2P Attention peers sent inconsistent control-plane flags" + ) + + input_ids_by_stage: dict[int, list[int]] = {} + for stage_idx in stage_ids: + stage_input_ids = tuple( + payload.input_ids_by_stage.get(stage_idx) for payload in payloads + ) + if all(input_ids is None for input_ids in stage_input_ids): + continue + if any(input_ids is None for input_ids in stage_input_ids): + raise RuntimeError( + "CAMP2P Attention peers must either all provide input_ids or " + f"all omit them for stage {stage_idx}" + ) + input_ids_by_stage[stage_idx] = [ + token_id + for input_ids in stage_input_ids + for token_id in input_ids + ] + + return AFDControlPayload( + dp_metadata_list=reference.dp_metadata_list, + is_graph_capturing=reference.is_graph_capturing, + is_warmup=reference.is_warmup, + input_ids_by_stage=input_ids_by_stage, + ) def build_camp2p_topology( @@ -677,8 +742,9 @@ def build_camp2p_topology( """Calculate the communication rank numbers for one process. FFN processes come first in the main AFD group, followed by Attention - processes. All FFN ranks and the first ``min(A, F)`` Attention ranks also - join the smaller Gloo group that exchanges token counts and batch details. + processes. The Gloo control group contains every process: an FFN rank + receives one payload from each Attention rank whose hidden states CAMP2P + aggregates on that FFN rank. Args: afd_config: Process role and total Attention/FFN rank counts. @@ -698,6 +764,11 @@ def build_camp2p_topology( "CAMP2P requires attention_size >= ffn_size, got " f"{attention_size} < {ffn_size}", ) + if attention_size % ffn_size != 0: + raise ValueError( + "CAMP2P requires attention_size to be an integer multiple of " + f"ffn_size, got attention_size={attention_size}, ffn_size={ffn_size}", + ) if role_rank < 0: raise ValueError(f"CAMP2P role rank must be non-negative, got {role_rank}") @@ -708,7 +779,7 @@ def build_camp2p_topology( f"(rank={role_rank}, size={attention_size})", ) world_rank = ffn_size + role_rank - p2p_rank = role_rank + min(ffn_size, attention_size) + p2p_rank = world_rank elif afd_config.role == "ffn": if role_rank >= ffn_size: raise ValueError( @@ -721,13 +792,20 @@ def build_camp2p_topology( raise ValueError(f"unknown AFD role {afd_config.role!r}") min_size = min(attention_size, ffn_size) - destinations: list[int] = [] - if ffn_size <= world_rank < ffn_size + min_size: - local_attention_rank = world_rank - ffn_size - dst = local_attention_rank - while dst < ffn_size: - destinations.append(dst) - dst += min_size + attention_ranks_per_ffn = attention_size // ffn_size + destinations: tuple[int, ...] = () + sources: tuple[int, ...] = () + if afd_config.role == "attention": + destinations = (role_rank // attention_ranks_per_ffn,) + else: + attention_rank_start = role_rank * attention_ranks_per_ffn + sources = tuple( + ffn_size + attention_rank + for attention_rank in range( + attention_rank_start, + attention_rank_start + attention_ranks_per_ffn, + ) + ) return _CAMP2PTopology( role=afd_config.role, @@ -737,7 +815,8 @@ def build_camp2p_topology( attention_size=attention_size, ffn_size=ffn_size, min_size=min_size, - dp_metadata_destinations=tuple(destinations), + dp_metadata_destinations=destinations, + dp_metadata_sources=sources, ) diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 90ee2462..d35a4294 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -128,6 +128,14 @@ def _send_and_receive( if afd_metadata is None: raise RuntimeError("RemoteFFNProxy requires AFD forward metadata") forward_context = get_forward_context() + # vLLM's startup memory profile executes a synthetic, role-local + # forward. There is no matching FFN request stream at that point, so + # issuing CAM dispatch/combine from the Attention profile can pair its + # small dummy batch with an FFN receive-capacity buffer and leave the + # external communicator in a broken state. Keep the profile local; + # live forwards below retain the real remote-FFN exchange. + if bool(getattr(forward_context, "in_profile_run", False)): + return hidden_states stage_idx = int( getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), ) diff --git a/afd_plugin/model_executor/models/deepseek_v4.py b/afd_plugin/model_executor/models/deepseek_v4.py index cfd15430..854f7213 100644 --- a/afd_plugin/model_executor/models/deepseek_v4.py +++ b/afd_plugin/model_executor/models/deepseek_v4.py @@ -1,510 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""CUDA AFD wrapper for the NVIDIA DeepSeek-V4 implementation. - -The split is placed immediately around each decoder FFN. Attention retains all -mHC residual-stream state; only the normalized two-dimensional FFN activation -and token-aligned input IDs cross the Attention-to-FFN boundary. The returned -FFN activation is the sole FFN-to-Attention tensor. -""" - -from collections.abc import Iterable, Iterator -from typing import Any - -import torch -import torch.nn as nn -from vllm.config import VllmConfig -from vllm.forward_context import get_forward_context -from vllm.models.deepseek_v4.nvidia import model as native - -from afd_plugin.config import parse_afd_config -from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata -from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context -from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield - -_ATTENTION_ROLE = frozenset(("attention",)) -_FFN_ROLE = frozenset(("ffn",)) -_BOTH_ROLES = frozenset(("attention", "ffn")) - - -def _weight_layer_path(name: str) -> tuple[int, str] | None: - """Extract the decoder layer index and first layer-local path component.""" - parts = name.split(".") - for marker_idx, part in enumerate(parts[:-2]): - if part != "layers": - continue - try: - layer_idx = int(parts[marker_idx + 1]) - except ValueError: - continue - return layer_idx, parts[marker_idx + 2] - return None - - -def _checkpoint_weight_roles(name: str) -> frozenset[str]: - """Classify a native DeepSeek-V4 checkpoint path by execution owner.""" - if name in { - "hc_head_fn", - "hc_head_base", - "hc_head_scale", - "model.hc_head_fn", - "model.hc_head_base", - "model.hc_head_scale", - }: - return _ATTENTION_ROLE - layer_path = _weight_layer_path(name) - if layer_path is None: - return _BOTH_ROLES - _, stage = layer_path - if stage == "ffn": - return _FFN_ROLE - return _ATTENTION_ROLE - - -def _iter_role_weights( - weights: Iterable[tuple[str, torch.Tensor]], - *, - role: str, -) -> Iterator[tuple[str, torch.Tensor]]: - """Consume a checkpoint iterator once and retain only role-owned paths.""" - for name, loaded_weight in weights: - if role in _checkpoint_weight_roles(name): - yield name, loaded_weight - - -class RemoteDeepseekV4FFN(nn.Module): - """Parameter-free FFN proxy carrying V4 hash-router token identifiers.""" - - def __init__(self, *, layer_idx: int) -> None: - super().__init__() - self.layer_idx = layer_idx - - def forward( - self, - hidden_states: torch.Tensor, - input_ids: torch.Tensor | None, - ) -> torch.Tensor: - if input_ids is None: - raise RuntimeError("DeepSeek-V4 remote FFN requires input_ids") - if input_ids.ndim != 1 or input_ids.shape[0] != hidden_states.shape[0]: - raise ValueError( - "DeepSeek-V4 input_ids must be one-dimensional and token-aligned", - ) - - afd_metadata = get_afd_metadata_from_forward_context() - if afd_metadata is None: - raise RuntimeError("RemoteDeepseekV4FFN requires AFD forward metadata") - forward_context = get_forward_context() - stage_idx = int( - getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), - ) - afd_metadata.stage_idx = stage_idx - metadata = AFDTransferMetadata.create_attention_metadata( - layer_idx=self.layer_idx, - stage_idx=stage_idx, - seq_len=int(hidden_states.shape[0]), - ) - context = AFDTransferContext(metadata=metadata) - afd_metadata.connector.send_attn_output( - hidden_states, - context, - input_ids=input_ids, - ) - hidden_states = maybe_apply_dbo_yield( - hidden_states, - role="attention", - ) - return afd_metadata.connector.recv_ffn_output( - ref_tensor=hidden_states, - ubatch_idx=stage_idx, - ) - - -class AFDDeepseekV4DecoderLayer(native.DeepseekV4DecoderLayer): - """DeepSeek-V4 decoder layer with an FFN-boundary synchronous split.""" - - # Patch reason: native DeepSeek-V4 always constructs Attention and FFN. - # Patch functionality: allocate only the stage owned by the active AFD role. - # Signature: matches upstream; no added parameters. - # Upstream: vLLM v0.26.0, vllm/models/deepseek_v4/nvidia/model.py - # Commit: 568afb3a13806beb53bb2e6bd518269357b237c0 - def __init__( - self, - vllm_config, - prefix, - topk_indices_buffer: torch.Tensor | None = None, - aux_stream_list: list[torch.cuda.Stream] | None = None, - ): - # ### PATCH START: construct a role-local decoder stage. - nn.Module.__init__(self) - afd_config = parse_afd_config(vllm_config, validate=False) - layer_idx = int(prefix.rsplit(".", maxsplit=1)[-1]) - # ### PATCH END - - config = vllm_config.model_config.hf_config - self.hidden_size = config.hidden_size - self.rms_norm_eps = config.rms_norm_eps - - # ### PATCH START: replace the remote stage with a parameter-free proxy. - if afd_config.role == "attention": - self.attn = native._select_dsv4_attn_cls(vllm_config)( - vllm_config, - prefix=f"{prefix}.attn", - topk_indices_buffer=topk_indices_buffer, - aux_stream_list=aux_stream_list, - ) - self.ffn = RemoteDeepseekV4FFN(layer_idx=layer_idx) - elif afd_config.role == "ffn": - self.attn = native.PPMissingLayer() - self.ffn = native.DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") - else: - raise ValueError(f"unsupported AFD role {afd_config.role!r}") - # ### PATCH END - - # ### PATCH START: mHC state and normalization are Attention-owned. - if afd_config.role == "ffn": - return - # ### PATCH END - self.attn_norm = native.RMSNorm(self.hidden_size, self.rms_norm_eps) - self.ffn_norm = native.RMSNorm(self.hidden_size, self.rms_norm_eps) - self.hc_mult = config.hc_mult - self.hc_sinkhorn_iters = config.hc_sinkhorn_iters - self.hc_eps = config.hc_eps - self.hc_post_alpha = 2.0 - mix_hc = (2 + self.hc_mult) * self.hc_mult - hc_dim = self.hc_mult * self.hidden_size - self.hc_attn_fn = nn.Parameter( - torch.empty((mix_hc, hc_dim), dtype=torch.float32), - requires_grad=False, - ) - self.hc_attn_fn_broadcast: torch.Tensor | None = None - self.hc_ffn_fn = nn.Parameter( - torch.empty((mix_hc, hc_dim), dtype=torch.float32), - requires_grad=False, - ) - self.hc_attn_base = nn.Parameter( - torch.empty(mix_hc, dtype=torch.float32), - requires_grad=False, - ) - self.hc_ffn_base = nn.Parameter( - torch.empty(mix_hc, dtype=torch.float32), - requires_grad=False, - ) - self.hc_attn_scale = nn.Parameter( - torch.empty(3, dtype=torch.float32), - requires_grad=False, - ) - self.hc_ffn_scale = nn.Parameter( - torch.empty(3, dtype=torch.float32), - requires_grad=False, - ) - - # Patch reason: native forward directly invokes its locally allocated FFN. - # Patch functionality: preserve native mHC state locally while the proxy - # transfers only the two-dimensional FFN activation and input IDs. - # Signature: matches upstream; no added parameters. - # Upstream: vLLM v0.26.0, vllm/models/deepseek_v4/nvidia/model.py - # Commit: 568afb3a13806beb53bb2e6bd518269357b237c0 - def forward( - self, - x: torch.Tensor, - positions: torch.Tensor, - input_ids: torch.Tensor | None, - post_mix: torch.Tensor | None = None, - res_mix: torch.Tensor | None = None, - residual: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - # ### PATCH START: prohibit accidental execution on the FFN worker. - if isinstance(self.attn, native.PPMissingLayer): - raise RuntimeError("DeepSeek-V4 decoder forward is Attention-owned") - # ### PATCH END - attn_norm_weight = self.attn_norm.weight.data - attn_norm_eps = self.attn_norm.variance_epsilon - if residual is None: - if x.dim() == 2: - assert self.hc_attn_fn_broadcast is not None - residual, post_mix, res_mix, x = native.mhc_pre_broadcast_tilelang( - x, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - norm_weight=attn_norm_weight, - norm_eps=attn_norm_eps, - fn_broadcast=self.hc_attn_fn_broadcast, - ) - else: - residual = x - post_mix, res_mix, x = native.mhc_pre_tilelang( - x, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - norm_weight=attn_norm_weight, - norm_eps=attn_norm_eps, - ) - else: - residual, post_mix, res_mix, x = native.mhc_fused_post_pre_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - n_splits=1, - tile_n=1, - norm_weight=attn_norm_weight, - norm_eps=attn_norm_eps, - ) - - x = self.attn(positions, x, None) - ffn_norm_weight = self.ffn_norm.weight.data - ffn_norm_eps = self.ffn_norm.variance_epsilon - residual, post_mix, res_mix, x = native.mhc_fused_post_pre_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - n_splits=1, - tile_n=1, - norm_weight=ffn_norm_weight, - norm_eps=ffn_norm_eps, - ) - - # ### PATCH START: this call enters the synchronous remote FFN proxy. - x = self.ffn(x, input_ids) - # ### PATCH END - return x, residual, post_mix, res_mix - - def compute_ffn_output( - self, - hidden_states: torch.Tensor, - *, - input_ids: torch.Tensor | None, - ) -> torch.Tensor: - """Execute the complete native V4 MoE, including its native router.""" - if not isinstance(self.ffn, native.DeepseekV4MoE): - raise RuntimeError("DeepSeek-V4 FFN compute is FFN-role only") - if input_ids is None: - raise RuntimeError("DeepSeek-V4 FFN compute requires input_ids") - return self.ffn(hidden_states, input_ids) - - -class AFDDeepseekV4Model(native.DeepseekV4Model): - """Role-aware DeepSeek-V4 model retaining mHC exclusively on Attention.""" - - # Patch reason: native DeepSeek-V4 allocates every decoder stage and stream. - # Patch functionality: build role-aware layers and Attention-only resources. - # Signature: matches upstream; no added parameters. - # Upstream: vLLM v0.26.0, vllm/models/deepseek_v4/nvidia/model.py - # Commit: 568afb3a13806beb53bb2e6bd518269357b237c0 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - # ### PATCH START: validate the deliberately narrow first release. - nn.Module.__init__(self) - self.afd_config = parse_afd_config(vllm_config, validate=False) - if native.current_platform.device_type != "cuda": - raise RuntimeError("AFD DeepSeek-V4 supports CUDA only") - if self.afd_config.connector != "P2pNcclAFDConnector": - raise RuntimeError( - "AFD DeepSeek-V4 requires the synchronous P2pNcclAFDConnector", - ) - if self.afd_config.compute_gate_on_attention: - raise RuntimeError( - "AFD DeepSeek-V4 does not support compute_gate_on_attention", - ) - parallel_config = vllm_config.parallel_config - if parallel_config.pipeline_parallel_size != 1: - raise RuntimeError("AFD DeepSeek-V4 does not support PP") - if parallel_config.use_sequence_parallel_moe: - raise RuntimeError("AFD DeepSeek-V4 does not support SP MoE") - if parallel_config.enable_eplb: - raise RuntimeError("AFD DeepSeek-V4 does not support EPLB") - # ### PATCH END - - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.config = config - self.quant_config = quant_config - self.parallel_config = parallel_config - self.use_mega_moe = ( - vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" - ) - # ### PATCH START: MegaMoE has unproven role-local finalization semantics. - if self.use_mega_moe: - raise RuntimeError("AFD DeepSeek-V4 does not support MegaMoE") - # ### PATCH END - self.vocab_size = config.vocab_size - self.hc_eps = config.hc_eps - self.hc_mult = config.hc_mult - self.hc_dim = self.hc_mult * config.hidden_size - self.rms_norm_eps = config.rms_norm_eps - - # ### PATCH START: CUDA streams and sparse-index buffers are Attention-owned. - if self.afd_config.role == "attention": - aux_stream_list = [torch.cuda.Stream() for _ in range(3)] - self.topk_indices_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - config.index_topk, - dtype=torch.int32, - ) - else: - aux_stream_list = None - self.topk_indices_buffer = None - # ### PATCH END - - if native.get_pp_group().is_first_rank: - self.embed_tokens = native.VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) - else: - self.embed_tokens = native.PPMissingLayer() - - # ### PATCH START: use the pinned role-aware layer constructor. - self.start_layer, self.end_layer, self.layers = native.make_layers( - config.num_hidden_layers, - lambda prefix: AFDDeepseekV4DecoderLayer( - vllm_config, - prefix=prefix, - topk_indices_buffer=self.topk_indices_buffer, - aux_stream_list=aux_stream_list, - ), - prefix=f"{prefix}.layers", - ) - # ### PATCH END - - if native.get_pp_group().is_last_rank: - self.norm = native.RMSNorm(config.hidden_size, self.rms_norm_eps) - else: - self.norm = native.PPMissingLayer() - - # ### PATCH START: final mHC state is constructed only on Attention. - if self.afd_config.role == "attention": - self.hc_head_fn = nn.Parameter( - torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32), - requires_grad=False, - ) - self.hc_head_base = nn.Parameter( - torch.empty(self.hc_mult, dtype=torch.float32), - requires_grad=False, - ) - self.hc_head_scale = nn.Parameter( - torch.empty(1, dtype=torch.float32), - requires_grad=False, - ) - else: - self.hc_head_fn = None - self.hc_head_base = None - self.hc_head_scale = None - if self.afd_config.role == "attention" and native.get_pp_group().is_last_rank: - self._mtp_hidden_buffer = torch.empty( - vllm_config.scheduler_config.max_num_batched_tokens, - self.hc_dim, - dtype=vllm_config.model_config.dtype, - ) - else: - self._mtp_hidden_buffer = None - # ### PATCH END - - def compute_ffn_output( - self, - hidden_states: torch.Tensor, - layer_idx: int, - *, - input_ids: torch.Tensor | None, - ) -> torch.Tensor: - return self.layers[layer_idx].compute_ffn_output( - hidden_states, - input_ids=input_ids, - ) - - def get_experts_layer_indices(self) -> tuple[int, ...]: - return tuple(range(int(self.config.num_hidden_layers))) - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - """Return native expert mappings only where real experts are owned.""" - if self.afd_config.role == "attention": - return [] - return super().get_expert_mapping() - - def finalize_mega_moe_weights(self) -> None: - """MegaMoE is rejected before allocation, so no finalizer is needed.""" - - def finalize_mhc_broadcast_weights(self) -> None: - """Finalize only the Attention-owned first-layer broadcast matrix.""" - if self.afd_config.role == "ffn": - return - if ( - not native.get_pp_group().is_first_rank - or self.start_layer >= self.end_layer - ): - return - layer = self.layers[self.start_layer] - if isinstance(layer, AFDDeepseekV4DecoderLayer): - layer.hc_attn_fn_broadcast = ( - layer.hc_attn_fn.detach() - .view(-1, layer.hc_mult, layer.hidden_size) - .sum(dim=1) - ) - - -class AFDDeepseekV4ForCausalLM(native.DeepseekV4ForCausalLM): - """DeepSeek-V4 causal LM exposing the GPU FFN-runner model contract.""" - - model_cls = AFDDeepseekV4Model - afd_requires_input_ids = True - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - self.afd_config = parse_afd_config(vllm_config, validate=False) - self.afd_role = self.afd_config.role - super().__init__(vllm_config=vllm_config, prefix=prefix) - - def compute_ffn_output( - self, - hidden_states: torch.Tensor, - layer_idx: int, - *, - input_ids: torch.Tensor | None = None, - **kwargs: Any, - ) -> torch.Tensor: - return self.model.compute_ffn_output( - hidden_states, - layer_idx, - input_ids=input_ids, - ) - - def get_experts_layer_indices(self) -> tuple[int, ...]: - return self.model.get_experts_layer_indices() - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - return super().load_weights( - _iter_role_weights(weights, role=self.afd_role), - ) - +"""Platform-selected AFD wrapper for DeepSeek-V4.""" + +from vllm.platforms import current_platform + +if current_platform.device_type == "npu": + from afd_plugin.model_executor.models.deepseek_v4_npu import ( + AFDDeepseekV4ForCausalLM, + ) +else: + from afd_plugin.model_executor.models.deepseek_v4_cuda import ( + AFDDeepseekV4ForCausalLM, + ) __all__ = ["AFDDeepseekV4ForCausalLM"] diff --git a/afd_plugin/model_executor/models/deepseek_v4_cuda.py b/afd_plugin/model_executor/models/deepseek_v4_cuda.py new file mode 100644 index 00000000..cfd15430 --- /dev/null +++ b/afd_plugin/model_executor/models/deepseek_v4_cuda.py @@ -0,0 +1,510 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CUDA AFD wrapper for the NVIDIA DeepSeek-V4 implementation. + +The split is placed immediately around each decoder FFN. Attention retains all +mHC residual-stream state; only the normalized two-dimensional FFN activation +and token-aligned input IDs cross the Attention-to-FFN boundary. The returned +FFN activation is the sole FFN-to-Attention tensor. +""" + +from collections.abc import Iterable, Iterator +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context +from vllm.models.deepseek_v4.nvidia import model as native + +from afd_plugin.config import parse_afd_config +from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata +from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context +from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield + +_ATTENTION_ROLE = frozenset(("attention",)) +_FFN_ROLE = frozenset(("ffn",)) +_BOTH_ROLES = frozenset(("attention", "ffn")) + + +def _weight_layer_path(name: str) -> tuple[int, str] | None: + """Extract the decoder layer index and first layer-local path component.""" + parts = name.split(".") + for marker_idx, part in enumerate(parts[:-2]): + if part != "layers": + continue + try: + layer_idx = int(parts[marker_idx + 1]) + except ValueError: + continue + return layer_idx, parts[marker_idx + 2] + return None + + +def _checkpoint_weight_roles(name: str) -> frozenset[str]: + """Classify a native DeepSeek-V4 checkpoint path by execution owner.""" + if name in { + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + "model.hc_head_fn", + "model.hc_head_base", + "model.hc_head_scale", + }: + return _ATTENTION_ROLE + layer_path = _weight_layer_path(name) + if layer_path is None: + return _BOTH_ROLES + _, stage = layer_path + if stage == "ffn": + return _FFN_ROLE + return _ATTENTION_ROLE + + +def _iter_role_weights( + weights: Iterable[tuple[str, torch.Tensor]], + *, + role: str, +) -> Iterator[tuple[str, torch.Tensor]]: + """Consume a checkpoint iterator once and retain only role-owned paths.""" + for name, loaded_weight in weights: + if role in _checkpoint_weight_roles(name): + yield name, loaded_weight + + +class RemoteDeepseekV4FFN(nn.Module): + """Parameter-free FFN proxy carrying V4 hash-router token identifiers.""" + + def __init__(self, *, layer_idx: int) -> None: + super().__init__() + self.layer_idx = layer_idx + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None, + ) -> torch.Tensor: + if input_ids is None: + raise RuntimeError("DeepSeek-V4 remote FFN requires input_ids") + if input_ids.ndim != 1 or input_ids.shape[0] != hidden_states.shape[0]: + raise ValueError( + "DeepSeek-V4 input_ids must be one-dimensional and token-aligned", + ) + + afd_metadata = get_afd_metadata_from_forward_context() + if afd_metadata is None: + raise RuntimeError("RemoteDeepseekV4FFN requires AFD forward metadata") + forward_context = get_forward_context() + stage_idx = int( + getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), + ) + afd_metadata.stage_idx = stage_idx + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=self.layer_idx, + stage_idx=stage_idx, + seq_len=int(hidden_states.shape[0]), + ) + context = AFDTransferContext(metadata=metadata) + afd_metadata.connector.send_attn_output( + hidden_states, + context, + input_ids=input_ids, + ) + hidden_states = maybe_apply_dbo_yield( + hidden_states, + role="attention", + ) + return afd_metadata.connector.recv_ffn_output( + ref_tensor=hidden_states, + ubatch_idx=stage_idx, + ) + + +class AFDDeepseekV4DecoderLayer(native.DeepseekV4DecoderLayer): + """DeepSeek-V4 decoder layer with an FFN-boundary synchronous split.""" + + # Patch reason: native DeepSeek-V4 always constructs Attention and FFN. + # Patch functionality: allocate only the stage owned by the active AFD role. + # Signature: matches upstream; no added parameters. + # Upstream: vLLM v0.26.0, vllm/models/deepseek_v4/nvidia/model.py + # Commit: 568afb3a13806beb53bb2e6bd518269357b237c0 + def __init__( + self, + vllm_config, + prefix, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + # ### PATCH START: construct a role-local decoder stage. + nn.Module.__init__(self) + afd_config = parse_afd_config(vllm_config, validate=False) + layer_idx = int(prefix.rsplit(".", maxsplit=1)[-1]) + # ### PATCH END + + config = vllm_config.model_config.hf_config + self.hidden_size = config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + # ### PATCH START: replace the remote stage with a parameter-free proxy. + if afd_config.role == "attention": + self.attn = native._select_dsv4_attn_cls(vllm_config)( + vllm_config, + prefix=f"{prefix}.attn", + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + self.ffn = RemoteDeepseekV4FFN(layer_idx=layer_idx) + elif afd_config.role == "ffn": + self.attn = native.PPMissingLayer() + self.ffn = native.DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") + else: + raise ValueError(f"unsupported AFD role {afd_config.role!r}") + # ### PATCH END + + # ### PATCH START: mHC state and normalization are Attention-owned. + if afd_config.role == "ffn": + return + # ### PATCH END + self.attn_norm = native.RMSNorm(self.hidden_size, self.rms_norm_eps) + self.ffn_norm = native.RMSNorm(self.hidden_size, self.rms_norm_eps) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.hc_post_alpha = 2.0 + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * self.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty((mix_hc, hc_dim), dtype=torch.float32), + requires_grad=False, + ) + self.hc_attn_fn_broadcast: torch.Tensor | None = None + self.hc_ffn_fn = nn.Parameter( + torch.empty((mix_hc, hc_dim), dtype=torch.float32), + requires_grad=False, + ) + self.hc_attn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), + requires_grad=False, + ) + self.hc_ffn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), + requires_grad=False, + ) + self.hc_attn_scale = nn.Parameter( + torch.empty(3, dtype=torch.float32), + requires_grad=False, + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty(3, dtype=torch.float32), + requires_grad=False, + ) + + # Patch reason: native forward directly invokes its locally allocated FFN. + # Patch functionality: preserve native mHC state locally while the proxy + # transfers only the two-dimensional FFN activation and input IDs. + # Signature: matches upstream; no added parameters. + # Upstream: vLLM v0.26.0, vllm/models/deepseek_v4/nvidia/model.py + # Commit: 568afb3a13806beb53bb2e6bd518269357b237c0 + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # ### PATCH START: prohibit accidental execution on the FFN worker. + if isinstance(self.attn, native.PPMissingLayer): + raise RuntimeError("DeepSeek-V4 decoder forward is Attention-owned") + # ### PATCH END + attn_norm_weight = self.attn_norm.weight.data + attn_norm_eps = self.attn_norm.variance_epsilon + if residual is None: + if x.dim() == 2: + assert self.hc_attn_fn_broadcast is not None + residual, post_mix, res_mix, x = native.mhc_pre_broadcast_tilelang( + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + fn_broadcast=self.hc_attn_fn_broadcast, + ) + else: + residual = x + post_mix, res_mix, x = native.mhc_pre_tilelang( + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + ) + else: + residual, post_mix, res_mix, x = native.mhc_fused_post_pre_tilelang( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + n_splits=1, + tile_n=1, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + ) + + x = self.attn(positions, x, None) + ffn_norm_weight = self.ffn_norm.weight.data + ffn_norm_eps = self.ffn_norm.variance_epsilon + residual, post_mix, res_mix, x = native.mhc_fused_post_pre_tilelang( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + n_splits=1, + tile_n=1, + norm_weight=ffn_norm_weight, + norm_eps=ffn_norm_eps, + ) + + # ### PATCH START: this call enters the synchronous remote FFN proxy. + x = self.ffn(x, input_ids) + # ### PATCH END + return x, residual, post_mix, res_mix + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + *, + input_ids: torch.Tensor | None, + ) -> torch.Tensor: + """Execute the complete native V4 MoE, including its native router.""" + if not isinstance(self.ffn, native.DeepseekV4MoE): + raise RuntimeError("DeepSeek-V4 FFN compute is FFN-role only") + if input_ids is None: + raise RuntimeError("DeepSeek-V4 FFN compute requires input_ids") + return self.ffn(hidden_states, input_ids) + + +class AFDDeepseekV4Model(native.DeepseekV4Model): + """Role-aware DeepSeek-V4 model retaining mHC exclusively on Attention.""" + + # Patch reason: native DeepSeek-V4 allocates every decoder stage and stream. + # Patch functionality: build role-aware layers and Attention-only resources. + # Signature: matches upstream; no added parameters. + # Upstream: vLLM v0.26.0, vllm/models/deepseek_v4/nvidia/model.py + # Commit: 568afb3a13806beb53bb2e6bd518269357b237c0 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + # ### PATCH START: validate the deliberately narrow first release. + nn.Module.__init__(self) + self.afd_config = parse_afd_config(vllm_config, validate=False) + if native.current_platform.device_type != "cuda": + raise RuntimeError("AFD DeepSeek-V4 supports CUDA only") + if self.afd_config.connector != "P2pNcclAFDConnector": + raise RuntimeError( + "AFD DeepSeek-V4 requires the synchronous P2pNcclAFDConnector", + ) + if self.afd_config.compute_gate_on_attention: + raise RuntimeError( + "AFD DeepSeek-V4 does not support compute_gate_on_attention", + ) + parallel_config = vllm_config.parallel_config + if parallel_config.pipeline_parallel_size != 1: + raise RuntimeError("AFD DeepSeek-V4 does not support PP") + if parallel_config.use_sequence_parallel_moe: + raise RuntimeError("AFD DeepSeek-V4 does not support SP MoE") + if parallel_config.enable_eplb: + raise RuntimeError("AFD DeepSeek-V4 does not support EPLB") + # ### PATCH END + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + self.parallel_config = parallel_config + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + # ### PATCH START: MegaMoE has unproven role-local finalization semantics. + if self.use_mega_moe: + raise RuntimeError("AFD DeepSeek-V4 does not support MegaMoE") + # ### PATCH END + self.vocab_size = config.vocab_size + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + # ### PATCH START: CUDA streams and sparse-index buffers are Attention-owned. + if self.afd_config.role == "attention": + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + ) + else: + aux_stream_list = None + self.topk_indices_buffer = None + # ### PATCH END + + if native.get_pp_group().is_first_rank: + self.embed_tokens = native.VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = native.PPMissingLayer() + + # ### PATCH START: use the pinned role-aware layer constructor. + self.start_layer, self.end_layer, self.layers = native.make_layers( + config.num_hidden_layers, + lambda prefix: AFDDeepseekV4DecoderLayer( + vllm_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + aux_stream_list=aux_stream_list, + ), + prefix=f"{prefix}.layers", + ) + # ### PATCH END + + if native.get_pp_group().is_last_rank: + self.norm = native.RMSNorm(config.hidden_size, self.rms_norm_eps) + else: + self.norm = native.PPMissingLayer() + + # ### PATCH START: final mHC state is constructed only on Attention. + if self.afd_config.role == "attention": + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + else: + self.hc_head_fn = None + self.hc_head_base = None + self.hc_head_scale = None + if self.afd_config.role == "attention" and native.get_pp_group().is_last_rank: + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + self.hc_dim, + dtype=vllm_config.model_config.dtype, + ) + else: + self._mtp_hidden_buffer = None + # ### PATCH END + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + *, + input_ids: torch.Tensor | None, + ) -> torch.Tensor: + return self.layers[layer_idx].compute_ffn_output( + hidden_states, + input_ids=input_ids, + ) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return tuple(range(int(self.config.num_hidden_layers))) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + """Return native expert mappings only where real experts are owned.""" + if self.afd_config.role == "attention": + return [] + return super().get_expert_mapping() + + def finalize_mega_moe_weights(self) -> None: + """MegaMoE is rejected before allocation, so no finalizer is needed.""" + + def finalize_mhc_broadcast_weights(self) -> None: + """Finalize only the Attention-owned first-layer broadcast matrix.""" + if self.afd_config.role == "ffn": + return + if ( + not native.get_pp_group().is_first_rank + or self.start_layer >= self.end_layer + ): + return + layer = self.layers[self.start_layer] + if isinstance(layer, AFDDeepseekV4DecoderLayer): + layer.hc_attn_fn_broadcast = ( + layer.hc_attn_fn.detach() + .view(-1, layer.hc_mult, layer.hidden_size) + .sum(dim=1) + ) + + +class AFDDeepseekV4ForCausalLM(native.DeepseekV4ForCausalLM): + """DeepSeek-V4 causal LM exposing the GPU FFN-runner model contract.""" + + model_cls = AFDDeepseekV4Model + afd_requires_input_ids = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + self.afd_config = parse_afd_config(vllm_config, validate=False) + self.afd_role = self.afd_config.role + super().__init__(vllm_config=vllm_config, prefix=prefix) + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + *, + input_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + return self.model.compute_ffn_output( + hidden_states, + layer_idx, + input_ids=input_ids, + ) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return self.model.get_experts_layer_indices() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return super().load_weights( + _iter_role_weights(weights, role=self.afd_role), + ) + + +__all__ = ["AFDDeepseekV4ForCausalLM"] diff --git a/afd_plugin/model_executor/models/deepseek_v4_npu.py b/afd_plugin/model_executor/models/deepseek_v4_npu.py new file mode 100644 index 00000000..1212268f --- /dev/null +++ b/afd_plugin/model_executor/models/deepseek_v4_npu.py @@ -0,0 +1,502 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepSeek V4 AFD model wrapper for the Ascend native implementation. + +The Ascend DSV4 implementation is not a DeepSeek V2 layer with different +dimensions. Its decoder layer owns the DSA attention path and the +hyper-connection (HC) state transitions around both attention and FFN. This +adapter therefore keeps the native layer forward path on the Attention role +and replaces only the FFN module with the AFD remote proxy. The FFN role +constructs the native MoE module and exposes it through the runner-facing +``compute_ffn_output`` hook. +""" + +from collections.abc import Iterable, Iterator +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers import fused_moe +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.sequence import IntermediateTensors + +from afd_plugin.config import AFD_ASYNC_CONNECTOR, parse_afd_config +from afd_plugin.connectors import AFDExpertRoutingSpec, AFDF2ATransferPayload +from afd_plugin.model_executor.models.deepseek_v2 import RemoteFFNProxy + +try: + from vllm_ascend.models import deepseek_v4 as native +except ImportError as exc: # pragma: no cover - only reachable off Ascend. + raise ImportError( + "DSV4 AFD support requires the vLLM-Ascend native DSV4 model" + ) from exc + + +_ATTENTION_ROLE = "attention" +_FFN_ROLE = "ffn" +_BOTH_ROLES = frozenset((_ATTENTION_ROLE, _FFN_ROLE)) + +def _refresh_ascend_fused_moe() -> None: + """Bind native DSV4 MoE construction to the Ascend implementation.""" + # vLLM-Ascend applies this replacement during platform initialization, + # while the DSV4 module keeps a module-level FusedMoE binding. Refresh it + # before either AFD role constructs its local model. + if native.current_platform.device_type == "npu": + native.FusedMoE = fused_moe.FusedMoE + + +def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: + parts = name.split(".") + for marker_idx, part in enumerate(parts[:-2]): + if part != "layers": + continue + try: + layer_idx = int(parts[marker_idx + 1]) + except ValueError: + continue + return layer_idx, parts[marker_idx + 2], tuple(parts[marker_idx + 3 :]) + return None + + +def _checkpoint_weight_roles(name: str) -> frozenset[str]: + """Return the AFD owner for a DSV4 checkpoint path. + + DSV4 checkpoints use ``attn``/``ffn`` names while the Ascend runtime + model exposes ``self_attn``/``mlp``. The native loader performs that name + conversion later, so filtering must understand both spellings here. + """ + + layer_path = _weight_layer_path(name) + if layer_path is None: + return _BOTH_ROLES + + _, stage, remainder = layer_path + if stage in ("attn", "self_attn"): + return frozenset((_ATTENTION_ROLE,)) + if stage in ("ffn", "mlp"): + if remainder and remainder[0] == "gate": + return _BOTH_ROLES + return frozenset((_FFN_ROLE,)) + # HC parameters and any future shared layer parameters are required by + # both role-local model instances. + return _BOTH_ROLES + + +def _iter_role_weights( + weights: Iterable[tuple[str, torch.Tensor]], + *, + role: str, +) -> Iterator[tuple[str, torch.Tensor]]: + for name, loaded_weight in weights: + if role in _checkpoint_weight_roles(name): + yield name, loaded_weight + + +class AFDDeepseekV4AttentionGateRemoteMoE(RemoteFFNProxy): + """DSV4 gate shell that routes local Attention tokens through Async CAM.""" + + def __init__( + self, + *, + config: Any, + layer_idx: int, + prefix: str, + ) -> None: + super().__init__(layer_idx=layer_idx) + self.top_k = int(config.num_experts_per_tok) + self.n_routed_experts = int(config.n_routed_experts) + self.renormalize = bool(config.norm_topk_prob) + self.scoring_func = getattr(config, "scoring_func", "softmax") + self.num_expert_group = int(getattr(config, "n_group", 1)) + self.topk_group = int(getattr(config, "topk_group", 1)) + self.routed_scaling_factor = float( + getattr(config, "routed_scaling_factor", 1.5), + ) + self.gate = ReplicatedLinear( + config.hidden_size, + config.n_routed_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + self.gate.precast_fp32_weight = True + if layer_idx < config.num_hash_layers: + self.gate.tid2eid = nn.Parameter( + torch.zeros( + config.vocab_size, + config.num_experts_per_tok, + dtype=torch.int32, + ), + requires_grad=False, + ) + self.gate.e_score_correction_bias = None + else: + self.gate.tid2eid = None + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Select DSV4 experts on Attention and exchange routed work via CAM.""" + from afd_plugin.model_executor.models.npu import deepseek_v4_attention_gate + + topk_weights, topk_ids = deepseek_v4_attention_gate.compute_attention_gate_topk( + self, + hidden_states, + ) + return self._send_and_receive( + hidden_states, + topk_weights=topk_weights, + topk_ids=topk_ids, + ) + + +class AFDDeepseekV4DecoderLayer(native.DeepseekV2DecoderLayer): + """Role-local DSV4 decoder layer. + + The inherited native ``forward`` is intentionally retained. On the + Attention role it executes the full native HC/DSA sequence and the + ``RemoteFFNProxy`` makes the AFD transfer at exactly the native FFN + boundary. The FFN role is connector-driven and does not call this full + forward method; it invokes ``compute_ffn_output`` instead. + """ + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + config=None, + topk_indices_buffer: torch.Tensor | None = None, + is_draft_layer: bool = False, + ) -> None: + if is_draft_layer: + raise ValueError("AFD DSV4 decoder layers do not support draft layers") + afd_config = parse_afd_config(vllm_config, validate=False) + nn.Module.__init__(self) + if config is None: + config = vllm_config.model_config.hf_config + + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + layer_idx = int(prefix.split(sep=".")[-1]) + + self.vllm_config = vllm_config + self.config = config + self.afd_role = afd_config.role + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.norm_eps = config.rms_norm_eps + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.compute_gate_on_attention = bool(afd_config.compute_gate_on_attention) + self.is_moe_layer = True + self.use_sequence_parallel_moe = False + + max_position_embeddings = config.rope_parameters[ + "original_max_position_embeddings" + ] + if afd_config.role == _ATTENTION_ROLE: + self.self_attn = native.DeepseekV4Attention( + vllm_config=vllm_config, + config=config, + max_position_embeddings=max_position_embeddings, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + topk_indices_buffer=topk_indices_buffer, + ) + if self.compute_gate_on_attention: + self.mlp = AFDDeepseekV4AttentionGateRemoteMoE( + config=config, + layer_idx=layer_idx, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = RemoteFFNProxy(layer_idx=layer_idx) + elif afd_config.role == _FFN_ROLE: + self.self_attn = native.PPMissingLayer() + _refresh_ascend_fused_moe() + self.mlp = native.DeepseekV4MoE( + config=config, + parallel_config=parallel_config, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + is_draft_layer=False, + ) + else: # pragma: no cover - parse_afd_config validates role values. + raise ValueError(f"Unsupported AFD role: {afd_config.role!r}") + + self.input_layernorm = native.RMSNorm( + config.hidden_size, + eps=self.norm_eps, + ) + self.post_attention_layernorm = native.RMSNorm( + config.hidden_size, + eps=self.norm_eps, + ) + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * config.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty(mix_hc, hc_dim, dtype=torch.float32) + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty(mix_hc, hc_dim, dtype=torch.float32) + ) + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) + self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + *, + group_list: torch.Tensor | None = None, + dynamic_scales: torch.Tensor | None = None, + expand_x_shared: torch.Tensor | None = None, + dynamic_scales_shared: torch.Tensor | None = None, + topk_scales: torch.Tensor | None = None, + group_list_type: int = 1, + **_: Any, + ) -> torch.Tensor | AFDF2ATransferPayload: + if not isinstance(self.mlp, native.DeepseekV4MoE): + raise RuntimeError( + "DSV4 compute_ffn_output requires the FFN role to own the MoE" + ) + if self.compute_gate_on_attention: + if group_list is None: + raise RuntimeError( + "DSV4 Attention-side routing requires CAM group_list on FFN", + ) + from afd_plugin.model_executor.models.npu import ( + deepseek_v2_attention_gate, + ) + + return deepseek_v2_attention_gate.compute_attention_gate_moe_ffn( + self, + hidden_states=hidden_states, + group_list=group_list, + dynamic_scales=dynamic_scales, + expand_x_shared=expand_x_shared, + dynamic_scales_shared=dynamic_scales_shared, + topk_scales=topk_scales, + group_list_type=group_list_type, + ) + if self.mlp.hash: + input_ids = get_forward_context().input_ids + num_tokens = hidden_states.reshape(-1, hidden_states.shape[-1]).shape[0] + if input_ids is None: + raise RuntimeError( + "DSV4 hash routing requires input_ids from the CAMP2P " + "control plane" + ) + if input_ids.numel() != num_tokens: + raise RuntimeError( + "DSV4 hash routing input_ids do not align with FFN hidden " + f"states: got {input_ids.numel()} IDs for {num_tokens} rows" + ) + return self.mlp(hidden_states) + + +@native.support_torch_compile +class AFDDeepseekV4Model(native.DeepseekV4Model): + """DSV4 model with role-local Attention/FFN allocations.""" + + fall_back_to_pt_during_load = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + afd_config = parse_afd_config(vllm_config, validate=False) + if ( + afd_config.connector == AFD_ASYNC_CONNECTOR + and not afd_config.compute_gate_on_attention + ): + raise ValueError( + "DSV4 CAMAsyncAFDConnector requires " + "compute_gate_on_attention=true", + ) + if vllm_config.parallel_config.use_sequence_parallel_moe: + raise RuntimeError("AFD DSV4 does not support sequence-parallel MoE") + + _refresh_ascend_fused_moe() + nn.Module.__init__(self) + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.vllm_config = vllm_config + self.compilation_config = vllm_config.compilation_config + self.afd_config = afd_config + self.config = config + self.device = native.current_platform.device_type + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.hc_mult = config.hc_mult + self.norm_eps = config.rms_norm_eps + self.hc_eps = config.hc_eps + + self.is_v32 = hasattr(config, "index_topk") + if self.is_v32 and afd_config.role == _ATTENTION_ROLE: + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + else: + topk_indices_buffer = None + self.topk_indices_buffer = topk_indices_buffer + + if native.get_pp_group().is_first_rank: + self.embed_tokens = native.VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = native.PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = native.make_layers( + config.num_hidden_layers, + lambda prefix, **_: AFDDeepseekV4DecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + topk_indices_buffer=topk_indices_buffer, + ), + prefix=f"{prefix}.layers", + ) + + if native.get_pp_group().is_last_rank: + self.norm = native.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = native.PPMissingLayer() + + hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32) + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32) + ) + self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32)) + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + hc_dim, + dtype=vllm_config.model_config.dtype, + device=self.device, + ) + + def make_empty_intermediate_tensors( + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.hc_mult, config.hidden_size), + dtype=dtype, + device=device, + ) + } + ) + + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors + self.aux_hidden_state_layers: tuple[int, ...] = () + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + **kwargs: Any, + ) -> torch.Tensor: + return self.layers[layer_idx].compute_ffn_output(hidden_states, **kwargs) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return tuple( + layer.layer_idx + for layer in self.layers + if isinstance(layer, AFDDeepseekV4DecoderLayer) + and isinstance(layer.mlp, native.DeepseekV4MoE) + ) + + def get_experts_routing_spec( + self, + layer_idx: int, + ) -> AFDExpertRoutingSpec: + layer = self.layers[layer_idx] + if not isinstance(layer.mlp, native.DeepseekV4MoE): + raise RuntimeError("DSV4 layer does not own a native MoE") + gate = layer.mlp.gate + return AFDExpertRoutingSpec( + router_logits_width=int(layer.mlp.n_routed_experts), + router_logits_dtype=gate.out_dtype or gate.weight.dtype, + ) + + def compute_experts_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + router_logits: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError( + "DSV4 AFD does not support compute_gate_on_attention yet" + ) + + +class AFDDeepseekV4ForCausalLM(native.AscendDeepseekV4ForCausalLM): + """Ascend DSV4 causal LM wrapper for AFD.""" + + model_cls = AFDDeepseekV4Model + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + self.afd_config = parse_afd_config(vllm_config, validate=False) + self.afd_role = self.afd_config.role + super().__init__(vllm_config=vllm_config, prefix=prefix) + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + **kwargs: Any, + ) -> torch.Tensor: + return self.model.compute_ffn_output(hidden_states, layer_idx, **kwargs) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return self.model.get_experts_layer_indices() + + def get_experts_routing_spec( + self, + layer_idx: int, + ) -> AFDExpertRoutingSpec: + return self.model.get_experts_routing_spec(layer_idx) + + def compute_experts_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + router_logits: torch.Tensor, + ) -> torch.Tensor: + return self.model.compute_experts_output( + hidden_states, + layer_idx, + router_logits, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return super().load_weights( + _iter_role_weights(weights, role=self.afd_role) + ) + + +__all__ = [ + "AFDDeepseekV4DecoderLayer", + "AFDDeepseekV4ForCausalLM", + "AFDDeepseekV4Model", +] diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index 383ea56e..4fdac8f2 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py @@ -168,6 +168,13 @@ def run_attention_gate_afd_forward( llama_4_scaling, ) + # CAM allocates persistent operator workspace on its first dispatch. + # Keep EngineCore's synthetic KV-cache profile local so that workspace + # is not mistaken for reclaimable KV capacity. The first real request + # performs the normal CAM dispatch/receive sequence below. + if forward_context.in_profile_run: + continue + dispatch_payload = prepare_cam_dispatch_payload( hidden_states, topk_weights, diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py b/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py new file mode 100644 index 00000000..821cb5e1 --- /dev/null +++ b/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepSeek V4 Attention-side routing helpers for Async CAM.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from afd_plugin.model_executor.models.deepseek_v4 import ( + AFDDeepseekV4AttentionGateRemoteMoE, + ) + + +def compute_attention_gate_topk( + moe: AFDDeepseekV4AttentionGateRemoteMoE, + hidden_states: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run DSV4 routing without entering vLLM's native MoE communicator. + + AFD Async CAM owns the cross-role dispatch. The vLLM-Ascend fused + selector's hash path instead assumes native EP/SP communication and calls + ``forward_context.moe_comm_method.pad_and_split_input_ids``. That object + is intentionally absent on Attention ranks, including the KV-cache profile + forward. Use the same CANN routing operators directly on local Attention + tokens, then hand their IDs and weights to CAM dispatch. + """ + + router_logits, _ = moe.gate(hidden_states) + if moe.scoring_func == "sqrtsoftplus": + topk_weights, topk_ids = _compute_sqrtsoftplus_topk(moe, router_logits) + else: + topk_weights, topk_ids = _compute_standard_topk(moe, router_logits) + return topk_weights.to(torch.float32), topk_ids + + +def _compute_sqrtsoftplus_topk( + moe: AFDDeepseekV4AttentionGateRemoteMoE, + router_logits: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run DSV4's sqrtsoftplus CANN router without native MoE communication.""" + + if moe.scoring_func != "sqrtsoftplus": + raise RuntimeError( + "DSV4 Hash routing requires scoring_func='sqrtsoftplus', got " + f"{moe.scoring_func!r}", + ) + + tid2eid = moe.gate.tid2eid + input_ids = None + if tid2eid is not None: + from vllm.forward_context import get_forward_context + + forward_context = get_forward_context() + input_ids = getattr(forward_context, "input_ids", None) + if input_ids is None: + raise RuntimeError( + "DSV4 Hash routing requires local input_ids in the forward " + "context", + ) + input_ids = input_ids.reshape(-1).to(torch.int64) + if input_ids.numel() != router_logits.shape[0]: + raise RuntimeError( + "DSV4 Hash routing input_ids/token count mismatch on Attention: " + f"input_ids={input_ids.numel()} router_tokens={router_logits.shape[0]}", + ) + input_ids = torch.where(input_ids == -1, 0, input_ids) + tid2eid = tid2eid.to(torch.int32) + correction_bias = moe.gate.e_score_correction_bias + if correction_bias is not None and correction_bias.dtype != router_logits.dtype: + correction_bias = correction_bias.to(router_logits.dtype) + topk_weights, topk_ids, _ = torch.ops._C_ascend.moe_gating_top_k_hash( + x=router_logits, + k=moe.top_k, + bias=correction_bias, + input_ids=input_ids, + tid2eid=tid2eid, + k_group=moe.topk_group, + group_count=moe.num_expert_group, + routed_scaling_factor=moe.routed_scaling_factor, + eps=1e-20, + group_select_mode=1, + renorm=0, + norm_type=2, + out_flag=False, + ) + return topk_weights, topk_ids + + +def _compute_standard_topk( + moe: AFDDeepseekV4AttentionGateRemoteMoE, + router_logits: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the non-Hash CANN selector without native MoE communication.""" + + from vllm_ascend.device.device_op import DeviceOperator + + norm_type_by_scoring_func = {"softmax": 0, "sigmoid": 1} + try: + norm_type = norm_type_by_scoring_func[moe.scoring_func] + except KeyError as exc: + raise RuntimeError( + "Unsupported non-Hash DSV4 routing scoring function: " + f"{moe.scoring_func!r}", + ) from exc + correction_bias = moe.gate.e_score_correction_bias + if correction_bias is not None and correction_bias.dtype != router_logits.dtype: + correction_bias = correction_bias.to(router_logits.dtype) + topk_weights, topk_ids, _ = DeviceOperator.moe_gating_top_k( + router_logits, + k=moe.top_k, + k_group=moe.topk_group, + group_count=moe.num_expert_group, + group_select_mode=1, + renorm=int(moe.renormalize), + norm_type=norm_type, + out_flag=False, + routed_scaling_factor=moe.routed_scaling_factor, + eps=1e-20, + bias_opt=correction_bias, + ) + return topk_weights, topk_ids diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index e5fd8789..2ea8f45f 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -49,7 +49,6 @@ from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.spec_decode.dflash_proposer import AscendDflashProposer from vllm_ascend.spec_decode.draft_proposer import AscendDraftModelProposer -from vllm_ascend.spec_decode.dspark_proposer import AscendDSparkProposer from vllm_ascend.spec_decode.eagle_proposer import AscendEagleProposer from vllm_ascend.spec_decode.step3p5 import AscendStep3p5MTPProposer from vllm_ascend.utils import ( @@ -205,6 +204,7 @@ def _model_forward( ): forward_context = get_forward_context() # ### PATCH START: AFD forward-context metadata + forward_context.input_ids = input_ids if self.ubatch_slices is not None: forward_context.ubatch_slices = self.ubatch_slices forward_context.dbo_enabled = False @@ -276,8 +276,7 @@ def _build_attention_metadata( num_reqs_padded, ) if self.afd_async_extra_info.async_moe_ubatching: - self.ubatch_slices = None - return self._build_attention_metadata_with_async_moe_ubatches( + result = self._build_attention_metadata_with_async_moe_ubatches( num_tokens=num_tokens, num_reqs=num_reqs, max_query_len=max_query_len, @@ -291,6 +290,20 @@ def _build_attention_metadata( num_scheduled_tokens_np=num_scheduled_tokens_np, cascade_attn_prefix_lens=cascade_attn_prefix_lens, ) + if ( + self._is_dsv4_model() + and self._afd_async_moe_ubatch_metadata is not None + ): + self.ubatch_slices = [ + UBatchSlice(stage.request_slice, stage.token_slice) + for stage in self._afd_async_moe_ubatch_metadata.stages + ] + return ( + self._afd_async_moe_ubatch_metadata.attn_metadata, + None, + ) + self.ubatch_slices = None + return result self._afd_pending_metadata = self._build_afd_metadata( ubatch_slices, num_tokens, @@ -391,22 +404,6 @@ def _build_attention_metadata_with_async_moe_ubatches( UBatchSlice(stage.request_slice, stage.token_slice) for stage in stages ] - logger.debug( - "AFD NPU async MoE ubatch split; num_reqs=%s num_tokens=%s " - "num_scheduled_tokens=%s split=%s sequence_parallel=%s " - "request_slices=%s token_slices=%s stage_input_tokens=%s " - "stage_actual_tokens=%s", - len(num_scheduled_tokens_np), - num_tokens, - num_scheduled_tokens_np.tolist(), - self.afd_async_extra_info.async_moe_split, - use_sequence_parallel, - [(stage.request_slice.start, stage.request_slice.stop) for stage in stages], - [(stage.token_slice.start, stage.token_slice.stop) for stage in stages], - [int(stage.input_tokens) for stage in stages], - [stage.actual_tokens for stage in stages], - ) - stage_attn_metadata, _ = self._build_attention_metadata_with_ubatches( num_tokens=num_tokens, num_reqs=num_reqs, @@ -754,7 +751,7 @@ def _build_attn_group_metadata( ) if self.speculative_config and isinstance( self.drafter, - AscendStep3p5MTPProposer | AscendDSparkProposer, + AscendStep3p5MTPProposer, ): self.drafter.set_per_group_attn_metadata( kv_cache_gid, @@ -766,8 +763,7 @@ def _build_attn_group_metadata( self.drafter, AscendEagleProposer | AscendDraftModelProposer - | AscendDflashProposer - | AscendDSparkProposer, + | AscendDflashProposer, ): if self.drafter.attn_layer_names[0] in kv_cache_group.layer_names: spec_decode_common_attn_metadata = cm @@ -1441,7 +1437,11 @@ def _install_afd_metadata_on_forward_context( padded_graph_tokens = _full_cudagraph_padded_tokens(forward_context) if padded_graph_tokens is not None and not ubatch_slices: dp_metadata = self._build_capture_dp_metadata(padded_graph_tokens) - self._send_dp_metadata(dp_metadata, ubatch_slices) + self._send_dp_metadata( + dp_metadata, + ubatch_slices, + input_ids=getattr(forward_context, "input_ids", None), + ) def _install_async_moe_ubatch_metadata_on_forward_context( self, @@ -1482,6 +1482,8 @@ def _send_dp_metadata( self, dp_metadata: DPMetadata | AFDDPMetadata | None, ubatch_slices: UBatchSlices | None, + *, + input_ids: torch.Tensor | None = None, ) -> None: assert self.connector.control_plane is not None, ( "_send_dp_metadata needs control plane driven connectors" @@ -1497,22 +1499,33 @@ def _send_dp_metadata( else: dp_metadata = self._ensure_dp_metadata(dp_metadata) dp_metadata_list = {0: dp_metadata} + input_ids_by_stage: dict[int, list[int]] = {} + if input_ids is not None: + if ubatch_slices and len(ubatch_slices) > 1: + for idx, ubatch in enumerate(ubatch_slices): + input_ids_by_stage[idx] = ( + input_ids[ubatch.token_slice] + .detach() + .to(device="cpu", dtype=torch.int64) + .flatten() + .tolist() + ) + else: + input_ids_by_stage[0] = ( + input_ids.detach() + .to(device="cpu", dtype=torch.int64) + .flatten() + .tolist() + ) is_warmup = bool(self._is_warmup) is_graph_capturing = bool(self._afd_is_graph_capturing) payload = AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, + input_ids_by_stage=input_ids_by_stage, ) self.connector.control_plane.update_state_from_dp_metadata(payload) - logger.warning( - "AFD NPU Attention send_dp_metadata decision; world_rank=%d " - "key=%s is_graph_capturing=%s is_warmup=%s", - self.connector.world_rank, - _dp_metadata_debug_key(dp_metadata_list), - is_graph_capturing, - is_warmup, - ) self.connector.control_plane.send_dp_metadata_list(payload) def _ensure_dp_metadata( @@ -1546,7 +1559,10 @@ def _build_capture_dp_metadata(self, num_tokens: int) -> DPMetadata | AFDDPMetad # Signature: matches upstream; no added parameters. def load_model(self) -> None: super().load_model() - if bool(self.vllm_config.parallel_config.use_ubatching): + if ( + bool(self.vllm_config.parallel_config.use_ubatching) + or self.afd_async_extra_info.async_moe_ubatching + ): self._install_ascend_ubatch_wrapper() # Wrapper installation is local and non-blocking; the connector # rendezvous is the blocking cross-role collective, so it is @@ -1555,6 +1571,11 @@ def load_model(self) -> None: if not self.connector.is_initialized: self.connector.init_afd_connector() + def _is_dsv4_model(self) -> bool: + from afd_plugin.compat.npu.feature_validation import _is_dsv4_target + + return _is_dsv4_target(self.vllm_config) + def _install_ascend_ubatch_wrapper(self) -> None: if isinstance(self.model, AscendUBatchWrapper): return @@ -1920,18 +1941,6 @@ def _make_uniform_dp_metadata(dp_size: int, num_tokens: int) -> AFDDPMetadata: return AFDDPMetadata(num_tokens_across_dp_cpu=num_tokens_across_dp_cpu) -def _dp_metadata_debug_key( - dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], -) -> tuple[tuple[int, tuple]]: - key_parts: list[tuple[int, tuple]] = [] - for stage_idx, metadata in sorted(dp_metadata_list.items()): - values_tuple = tuple( - int(value) for value in metadata.num_tokens_across_dp_cpu.tolist() - ) - key_parts.append((int(stage_idx), values_tuple)) - return tuple(key_parts) - - def _normalize_metadata_ubatch_slices( ubatch_slices: UBatchSlices | None, num_tokens_padded: int | None, diff --git a/afd_plugin/v1/worker/npu/attention_worker.py b/afd_plugin/v1/worker/npu/attention_worker.py index d6a24861..9991b332 100644 --- a/afd_plugin/v1/worker/npu/attention_worker.py +++ b/afd_plugin/v1/worker/npu/attention_worker.py @@ -15,10 +15,13 @@ fix_all2all_backend_for_afd, npu_afd_num_ubatches, ) + from afd_plugin.model_executor.models.model_utils import get_afd_model_config + from afd_plugin.v1.worker.npu.attention_model_runner import ( AFDNPUAttentionModelRunner, ) + from afd_plugin.validation import ( NPU_ATTENTION_WORKER_FQCN, assert_compatible_afd_stack, diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 83e81284..325634ce 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -36,6 +36,7 @@ AFDAsyncTransferState, CAMAsyncAFDConnector, ) +from afd_plugin.connectors.npu.camp2p import CAMP2pAFDConnector from afd_plugin.v1.worker.attention_model_runner import ( _resolve_world_ranks, ) @@ -56,7 +57,6 @@ logger = init_logger(__name__) - class AFDNPUFFNModelRunner(NPUModelRunner): """Connector-driven NPU FFN runner for AFD execution.""" @@ -208,6 +208,7 @@ def _ffn_forward( _make_dp_metadata_payload( dp_metadata_list, is_graph_capturing=is_graph_capturing, + input_ids_by_stage=_camp2p_input_ids_by_stage(self.connector), ), ) num_stages = max(len(dp_metadata_list), 1) @@ -263,6 +264,12 @@ def _ffn_forward( forward_context.dp_metadata = dp_metadata_list.get(stage_idx) forward_context.additional_kwargs["afd_metadata"] = metadata assert states, "Context.states must not be None" + _install_ffn_input_ids( + forward_context, + _camp2p_input_ids_by_stage(self.connector).get(stage_idx), + num_tokens=int(hidden_states.shape[0]), + device=hidden_states.device, + ) _set_moe_layer_index(forward_context, layer_idx) rank_ffn_output = self.model.compute_ffn_output( @@ -284,10 +291,11 @@ def _ffn_forward_connector_driven( rank_ffn_output = None connector = cast(CAMAsyncAFDConnector, self.connector) - for _ in _ffn_layer_indices(self): + for expected_layer_idx in _ffn_layer_indices(self): work_item = connector.recv_ffn_work_item( stage_idx=stage_idx, max_num_tokens=self.max_num_tokens, + expected_layer_idx=int(expected_layer_idx), ) hidden_states = work_item.hidden_states metadata = work_item.context.metadata @@ -330,6 +338,13 @@ def _ffn_forward_connector_driven( work_item, rank_ffn_output, ) + # ``async_combine_send`` returns after enqueueing its NPU work. + # Without an acknowledgement/credit path, immediately calling + # dispatch-recv for all following MoE layers can reuse CAM's + # bounded buffers while the prior combine is still in flight. + # The target baseline is deliberately serialized per layer; + # async MoE ubatching will introduce explicit stage credits. + torch.npu.synchronize() return rank_ffn_output def capture_model( @@ -478,9 +493,12 @@ def _ffn_layer_indices(runner: AFDNPUFFNModelRunner) -> range | list[int]: def _is_moe_layer(hf_config: object, layer_idx: int) -> bool: moe_layer_freq = getattr(hf_config, "moe_layer_freq", 1) + # DeepSeek V4 has routed experts in every decoder layer and does not expose + # the V2/V3 ``first_k_dense_replace`` compatibility field. + first_moe_layer = getattr(hf_config, "first_k_dense_replace", 0) return ( - hf_config.n_routed_experts is not None - and layer_idx >= hf_config.first_k_dense_replace + getattr(hf_config, "n_routed_experts", None) is not None + and layer_idx >= first_moe_layer and layer_idx % moe_layer_freq == 0 ) @@ -490,14 +508,44 @@ def _make_dp_metadata_payload( *, is_graph_capturing: bool = False, is_warmup: bool = False, + input_ids_by_stage: dict[int, list[int]] | None = None, ) -> AFDControlPayload: return AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, + input_ids_by_stage={} if input_ids_by_stage is None else input_ids_by_stage, ) +def _camp2p_input_ids_by_stage( + connector: AFDConnectorBase, +) -> dict[int, list[int]]: + if isinstance(connector, CAMP2pAFDConnector): + return connector.input_ids_by_stage + return {} + + +def _install_ffn_input_ids( + forward_context: Any, + input_ids: list[int] | None, + *, + num_tokens: int, + device: torch.device, +) -> None: + """Install token ids required by native DSV4 hash routing on FFN.""" + if input_ids is None: + forward_context.input_ids = None + return + ids = torch.as_tensor(input_ids, dtype=torch.int64, device=device).flatten() + if ids.numel() != num_tokens: + raise RuntimeError( + "CAMP2P input_ids must align with received hidden states: " + f"got {ids.numel()} token IDs for {num_tokens} hidden-state rows" + ) + forward_context.input_ids = ids + + def _ffn_token_counts_across_ranks( connector: AFDConnectorBase, dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index 54b8120a..0197f0ff 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -19,8 +19,11 @@ fix_all2all_backend_for_afd, npu_afd_num_ubatches, ) + from afd_plugin.model_executor.models.model_utils import get_afd_model_config + from afd_plugin.v1.worker.npu.ffn_model_runner import AFDNPUFFNModelRunner + from afd_plugin.validation import NPU_FFN_WORKER_FQCN, assert_compatible_afd_stack if TYPE_CHECKING: @@ -141,6 +144,9 @@ def _run_ffn_server_loop(self) -> None: continue payload = self.model_runner.connector.control_plane.recv_dp_metadata_list() + self.model_runner.connector.control_plane.update_state_from_dp_metadata( + payload + ) dp_metadata_list = payload.dp_metadata_list is_attn_graph_capturing = payload.is_graph_capturing is_warmup = payload.is_warmup diff --git a/afd_plugin/v1/worker/npu/forward_context.py b/afd_plugin/v1/worker/npu/forward_context.py index e7ca5d2e..2b0e2984 100644 --- a/afd_plugin/v1/worker/npu/forward_context.py +++ b/afd_plugin/v1/worker/npu/forward_context.py @@ -109,7 +109,12 @@ def create_ascend_forward_context( cur_forward_context.max_tokens_across_pcp ) new_forward_context.sinks = cur_forward_context.sinks - new_forward_context.input_ids = cur_forward_context.input_ids + if cur_forward_context.input_ids is not None: + new_forward_context.input_ids = cur_forward_context.input_ids[ + ubatch_slice.token_slice + ] + else: + new_forward_context.input_ids = None new_forward_context.eplb_heat_collection_status = ( cur_forward_context.eplb_heat_collection_status ) diff --git a/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py b/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py index 93f30653..4a9998ca 100644 --- a/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py +++ b/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py @@ -157,7 +157,15 @@ def __init__( self.vllm_config = vllm_config self.compilation_config = vllm_config.compilation_config self.comm_stream = torch.npu.Stream(device=device) - assert self.vllm_config.parallel_config.num_ubatches == AFD_NPU_NUM_UBATCHES + # vLLM native ubatching/DBO is rejected for CAMAsyncAFDConnector by + # fail_if_unsupported_npu_afd_features; the AFD async MoE ubatching + # path drives stage counts through connector_extra_config + # (async_moe_num_ubatches), which is validated there as well. + if self.vllm_config.parallel_config.use_ubatching: + raise RuntimeError( + "AscendUBatchWrapper does not support vLLM native " + "ubatching/DBO; use the AFD async MoE ubatching path.", + ) self.ready_barrier = threading.Barrier(_READY_BARRIER_PARTIES) self.cudagraphs: dict[AscendNPUGraphKey, AscendNPUGraphMetaData] = {} self.cudagraph_wrapper = None diff --git a/tests/unit/compat/npu/test_dsv4_async_validation.py b/tests/unit/compat/npu/test_dsv4_async_validation.py new file mode 100644 index 00000000..45b7a5fc --- /dev/null +++ b/tests/unit/compat/npu/test_dsv4_async_validation.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Unit coverage for the narrow DSV4 Async CAM feature contract.""" + +from types import SimpleNamespace + +import pytest + +from afd_plugin.compat.npu.feature_validation import ( + _fail_if_unsupported_dsv4_async_features, +) +from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo + + +def _afd_config(*, compute_gate_on_attention: bool) -> SimpleNamespace: + return SimpleNamespace(compute_gate_on_attention=compute_gate_on_attention) + + +def test_dsv4_async_requires_attention_side_gate() -> None: + with pytest.raises(RuntimeError, match="compute_gate_on_attention"): + _fail_if_unsupported_dsv4_async_features( + _afd_config(compute_gate_on_attention=False), + AFDAsyncExtraInfo(dynamic_quant=1), + ) + + +@pytest.mark.parametrize( + ("extra_info", "message"), + [ + ( + AFDAsyncExtraInfo(dynamic_quant=0), + "dynamicQuant=1", + ), + ], +) +def test_dsv4_async_rejects_deferred_features( + extra_info: AFDAsyncExtraInfo, + message: str, +) -> None: + with pytest.raises(RuntimeError, match=message): + _fail_if_unsupported_dsv4_async_features( + _afd_config(compute_gate_on_attention=True), + extra_info, + ) + + +def test_dsv4_async_accepts_target_baseline_and_async_moe_ubatching() -> None: + _fail_if_unsupported_dsv4_async_features( + _afd_config(compute_gate_on_attention=True), + AFDAsyncExtraInfo( + dynamic_quant=1, + attn_ranks_per_dp=2, + async_moe_ubatching=True, + ), + ) diff --git a/tests/unit/connectors/test_async_cam_connector.py b/tests/unit/connectors/test_async_cam_connector.py index 0e64b530..608517d9 100644 --- a/tests/unit/connectors/test_async_cam_connector.py +++ b/tests/unit/connectors/test_async_cam_connector.py @@ -138,7 +138,13 @@ def zeros(self, shape, *, dtype, device): def _vllm_config(*, tp_size: int = 1, pcp_size: int = 1, extra_config=None): return SimpleNamespace( - additional_config={"afd": {"connector_extra_config": extra_config or {}}}, + additional_config={ + "afd": { + "connector_extra_config": extra_config + if extra_config is not None + else {"attn_ranks_per_dp": 4}, + }, + }, parallel_config=SimpleNamespace( data_parallel_size=1, data_parallel_rank=0, @@ -165,6 +171,15 @@ def _afd_config(*, role: str): ) +def _dp2_afd_config(*, role: str): + return AFDConfig( + connector="CAMAsyncAFDConnector", + role=role, + num_attention_ranks=8, + num_ffn_ranks=8, + ) + + def _topk_payload(batch_size: int, topk: int = 2): return { "topk_ids": _FakeTensor((batch_size, topk), dtype="int32"), @@ -227,13 +242,13 @@ def test_async_connector_uses_attn_ranks_per_dp_for_cam_tp_size(): _vllm_config( tp_size=4, pcp_size=2, - extra_config={"attn_ranks_per_dp": "3"}, + extra_config={"attn_ranks_per_dp": "4"}, ), _afd_config(role="attention"), 0, ) - assert connector.tp_size == 3 + assert connector.tp_size == 4 @pytest.mark.parametrize("value", [True, "bad"]) @@ -273,6 +288,64 @@ def test_async_topology_uses_cam_attention_first_rank_layout(): assert ffn.expert_per_rank == 4 +@pytest.mark.parametrize( + ("role", "role_rank", "expected_group", "expected_world_rank"), + [ + ("attention", 3, 0, 3), + ("attention", 4, 1, 0), + ("ffn", 3, 0, 7), + ("ffn", 4, 1, 4), + ], +) +def test_async_topology_scopes_cam_world_to_each_dp_group( + role, + role_rank, + expected_group, + expected_world_rank, +): + topology = build_async_topology( + _dp2_afd_config(role=role), + role_rank, + num_routed_experts=32, + attn_ranks_per_dp=4, + ) + + assert topology.dp_group_index == expected_group + assert topology.num_dp_groups == 2 + assert topology.world_rank == expected_world_rank + assert topology.attn_size == 4 + assert topology.ffn_size == 4 + assert topology.world_size == 8 + assert topology.expert_per_rank == 8 + + +@pytest.mark.parametrize("dp_group_index", [0, 1]) +def test_async_topology_allows_a_shared_ffn_ep_pool(dp_group_index): + topology = build_async_topology( + _dp2_afd_config(role="ffn"), + 7, + num_routed_experts=256, + attn_ranks_per_dp=4, + shared_ffn_pool=True, + dp_group_index=dp_group_index, + ) + + assert topology.dp_group_index == dp_group_index + assert topology.num_dp_groups == 2 + assert topology.world_rank == 11 + assert topology.attn_size == 4 + assert topology.ffn_size == 8 + assert topology.world_size == 12 + assert topology.expert_per_rank == 32 + + +def test_async_extra_info_parses_shared_ffn_pool(): + extra_info = AFDAsyncExtraInfo.from_mapping({"shared_ffn_pool": "true"}) + + assert extra_info.shared_ffn_pool is True + assert extra_info.to_mapping()["shared_ffn_pool"] is True + + def test_async_connector_init_creates_attention_first_hccl_group(monkeypatch): calls = [] fake_torch = _FakeTorch() @@ -325,6 +398,45 @@ def fake_init_afd_process_group(**kwargs): assert connector._placeholder.shape == (1,) +def test_async_connector_init_uses_a_distinct_hccl_group_for_each_dp(monkeypatch): + calls = [] + fake_torch = _FakeTorch() + monkeypatch.setattr(async_cam_module, "torch", fake_torch) + monkeypatch.setattr( + async_cam_module, + "ensure_cam_async_ops_available", + lambda: None, + ) + + def fake_init_afd_process_group(**kwargs): + calls.append(kwargs) + backend = SimpleNamespace( + get_hccl_comm_name=lambda rank: f"hccl:{kwargs['group_name']}:{rank}", + ) + return SimpleNamespace(_get_backend=lambda device: backend) + + monkeypatch.setattr( + async_cam_module, + "init_afd_process_group", + fake_init_afd_process_group, + ) + connector = CAMAsyncAFDConnector( + 0, + 0, + _vllm_config(extra_config={"attn_ranks_per_dp": 4}), + _dp2_afd_config(role="ffn"), + 4, + ) + + connector.init_afd_connector() + + assert calls[0]["world_size"] == 8 + assert calls[0]["rank"] == 4 + assert calls[0]["group_name"] == "afd_async_cam_dp1" + assert calls[0]["init_method"] == "tcp://127.0.0.1:1240" + assert connector.group_name == "hccl:afd_async_cam_dp1:4" + + def test_async_connector_disables_dp_metadata_control_plane(): connector = CAMAsyncAFDConnector( 0, @@ -345,7 +457,7 @@ def test_async_connector_calls_cam_shaped_ops(monkeypatch): 0, _vllm_config( pcp_size=3, - extra_config={"attn_ranks_per_dp": 3}, + extra_config={"attn_ranks_per_dp": 4}, ), _afd_config(role="attention"), 0, @@ -379,7 +491,7 @@ def test_async_connector_calls_cam_shaped_ops(monkeypatch): assert fake_torch.ops.umdk_cam_op_lib.calls[1][1][4] == CAM_COMM_ID assert fake_torch.ops.umdk_cam_op_lib.calls[0][1][5:11] == (3, 16, 2, 2, 4, 4) assert fake_torch.ops.umdk_cam_op_lib.calls[1][1][5:11] == (3, 16, 2, 2, 4, 4) - assert fake_torch.ops.umdk_cam_op_lib.calls[0][1][14] == 3 + assert fake_torch.ops.umdk_cam_op_lib.calls[0][1][14] == 4 assert isinstance(context.states, AFDAsyncTransferState) assert isinstance(context.states, AFDTransferState) @@ -409,7 +521,7 @@ def test_async_ffn_side_dispatch_recv_and_combine_send(monkeypatch): assert states.dynamic_scales.shape == (4,) assert states.expand_x_shared.shape == (2, 16) assert states.dynamic_scales_shared.shape == (2,) - assert states.group_list.shape == (4,) + assert states.group_list.shape == (connector.topology.expert_per_rank,) assert states.expert_token_nums_shared.shape == (1,) assert fake_torch.ops.umdk_cam_op_lib.calls[0][0] == "dispatch_recv" assert fake_torch.ops.umdk_cam_op_lib.calls[1][0] == "combine_send" @@ -473,11 +585,7 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): hidden_size=connector.hidden_size, topk=connector.topk, layer_idx=layer_idx, - token_nums_rankid_layeridx=[ - _FakeScalar(7), - _FakeScalar(0), - _FakeScalar(11), - ], + token_nums_rankid_layeridx=torch.tensor([7, 0, 11], dtype=torch.int64), expert_token_nums_shared=[_FakeScalar(2)], group_list=torch.tensor([2, 3], dtype=torch.int64), dynamic_scales=_FakeTensorLike("scales"), @@ -491,7 +599,11 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): monkeypatch.setattr(connector, "recv_attn_output", fake_recv_attn_output) - work_item = connector.recv_ffn_work_item(stage_idx=0, max_num_tokens=16) + work_item = connector.recv_ffn_work_item( + stage_idx=0, + max_num_tokens=16, + expected_layer_idx=0, + ) states = work_item.context.states assert work_item.layer_idx == 11 @@ -537,11 +649,7 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): hidden_size=connector.hidden_size, topk=connector.topk, layer_idx=layer_idx, - token_nums_rankid_layeridx=[ - _FakeScalar(6), - _FakeScalar(0), - _FakeScalar(23), - ], + token_nums_rankid_layeridx=torch.tensor([6, 0, 23], dtype=torch.int64), expert_token_nums_shared=[_FakeScalar(0)], group_list=group_list, expand_x_shared=_FakeTensorLike("shared-hidden"), @@ -554,7 +662,11 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): monkeypatch.setattr(connector, "recv_attn_output", fake_recv_attn_output) - work_item = connector.recv_ffn_work_item(stage_idx=0, max_num_tokens=16) + work_item = connector.recv_ffn_work_item( + stage_idx=0, + max_num_tokens=16, + expected_layer_idx=0, + ) states = work_item.context.states assert work_item.layer_idx == 23 @@ -597,11 +709,7 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): hidden_size=connector.hidden_size, topk=connector.topk, layer_idx=layer_idx, - token_nums_rankid_layeridx=[ - _FakeScalar(5), - _FakeScalar(0), - _FakeScalar(7), - ], + token_nums_rankid_layeridx=torch.tensor([5, 0, 7], dtype=torch.int64), expert_token_nums_shared=[_FakeScalar(5)], group_list=_FakeIntVector([0] * 8), expand_x_shared=_FakeTensorLike("shared-hidden"), @@ -614,7 +722,11 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): monkeypatch.setattr(connector, "recv_attn_output", fake_recv_attn_output) - work_item = connector.recv_ffn_work_item(stage_idx=0, max_num_tokens=16) + work_item = connector.recv_ffn_work_item( + stage_idx=0, + max_num_tokens=16, + expected_layer_idx=0, + ) sent_output = connector.send_ffn_work_item_output( work_item, AFDF2ATransferPayload( diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index 1abc3acb..2c2ad76a 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -12,6 +12,8 @@ from afd_plugin.config import AFDConfig from afd_plugin.connectors import ( AFDConnectorFactory, + AFDControlPayload, + AFDDPMetadata, AFDTransferContext, AFDTransferMetadata, AFDTransferState, @@ -84,7 +86,7 @@ def test_camp2p_factory_creates_connector(): assert connector.extra_info.core_num == 12 -def test_camp2p_topology_matches_original_rank_layout(): +def test_camp2p_topology_maps_every_attention_rank_to_its_ffn_group(): attn0 = build_camp2p_topology(_afd_config(role="attention"), 0) attn1 = build_camp2p_topology(_afd_config(role="attention"), 1) attn2 = build_camp2p_topology(_afd_config(role="attention"), 2) @@ -98,10 +100,50 @@ def test_camp2p_topology_matches_original_rank_layout(): assert (attn1.world_rank, attn1.p2p_rank, attn1.dp_metadata_destinations) == ( 3, 3, + (0,), + ) + assert (attn2.world_rank, attn2.p2p_rank, attn2.dp_metadata_destinations) == ( + 4, + 4, (1,), ) - assert not attn2.participates_in_p2p_group + assert attn2.participates_in_p2p_group assert (ffn1.world_rank, ffn1.p2p_rank) == (1, 1) + assert ffn1.dp_metadata_sources == (4, 5) + assert ffn1.p2p_world_size == 6 + + +def test_camp2p_control_payload_aggregates_input_ids_per_stage(): + torch = pytest.importorskip("torch") + dp_metadata_list = { + 0: AFDDPMetadata(torch.tensor([2, 3, 5, 7], dtype=torch.int32)), + 1: AFDDPMetadata(torch.tensor([11, 13, 17, 19], dtype=torch.int32)), + } + payload0 = AFDControlPayload( + dp_metadata_list=dp_metadata_list, + is_graph_capturing=False, + is_warmup=False, + input_ids_by_stage={0: [10, 11], 1: [12]}, + ) + payload1 = AFDControlPayload( + dp_metadata_list=dp_metadata_list, + is_graph_capturing=False, + is_warmup=False, + input_ids_by_stage={0: [20, 21, 22], 1: [23, 24]}, + ) + + merged = camp2p_module._aggregate_camp2p_control_payloads( + (payload0, payload1) + ) + + assert merged.input_ids_by_stage == { + 0: [10, 11, 20, 21, 22], + 1: [12, 23, 24], + } + assert torch.equal( + merged.dp_metadata_list[0].num_tokens_across_dp_cpu, + payload0.dp_metadata_list[0].num_tokens_across_dp_cpu, + ) def _init_ffn_connector(rank, vllm_config): diff --git a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py index b650f480..7e0b07bc 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -198,6 +198,22 @@ def test_remote_proxy_requires_forward_metadata(monkeypatch): adapter.RemoteFFNProxy(layer_idx=0)(torch.ones(1, 4)) +def test_remote_proxy_skips_cam_exchange_during_profile(monkeypatch): + events = [] + _install_fake_forward_context(monkeypatch, events) + monkeypatch.setattr( + adapter, + "get_forward_context", + lambda: SimpleNamespace(in_profile_run=True, ubatch_idx=0), + ) + hidden_states = torch.ones(2, 4) + + output = adapter.RemoteFFNProxy(layer_idx=0)(hidden_states) + + assert output is hidden_states + assert events == [] + + def test_synchronous_model_forward_delegates_to_native(monkeypatch): calls = [] expected = torch.ones(1, 4) diff --git a/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py b/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py new file mode 100644 index 00000000..3565e05a --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +from __future__ import annotations + +from pathlib import Path + + +def test_dsv4_async_gate_bypasses_native_moe_communicator() -> None: + """Async CAM must not invoke Ascend's EP/SP selector wrapper. + + That wrapper calls ``forward_context.moe_comm_method`` during DSV4 Hash + routing. AFD owns this communication boundary, so the Attention helper + must invoke the CANN routing operators directly on its local tokens. + """ + + source = Path( + "afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py", + ).read_text() + + assert "torch.ops._C_ascend.moe_gating_top_k_hash(" in source + assert "DeviceOperator.moe_gating_top_k(" in source + assert "correction_bias = correction_bias.to(router_logits.dtype)" in source + assert "from vllm_ascend.ops.fused_moe.experts_selector" not in source + assert "forward_context.moe_comm_method.pad_and_split_input_ids(" not in source + assert "connector.select_experts" not in source + + +def test_dsv4_async_gate_validates_local_hash_token_alignment() -> None: + source = Path( + "afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py", + ).read_text() + + assert "DSV4 Hash routing input_ids/token count mismatch on Attention" in source + assert "input_ids = input_ids.reshape(-1).to(torch.int64)" in source diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 79674247..f4e88c22 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -89,12 +89,15 @@ class _RecordingConnector: def __init__(self): self.dp_metadata_updates = [] self.sent_dp_metadata_lists = [] + self.dp_metadata_payloads = [] + self.sent_dp_metadata_payloads = [] # The runners reach the control plane through connector.control_plane; # the fake serves as both. self.control_plane = self def update_state_from_dp_metadata(self, payload): assert isinstance(payload, AFDControlPayload) + self.dp_metadata_payloads.append(payload) self.dp_metadata_updates.append( ( payload.dp_metadata_list, @@ -105,6 +108,7 @@ def update_state_from_dp_metadata(self, payload): def send_dp_metadata_list(self, payload): assert isinstance(payload, AFDControlPayload) + self.sent_dp_metadata_payloads.append(payload) self.sent_dp_metadata_lists.append( ( payload.dp_metadata_list, @@ -477,9 +481,10 @@ def __call__(self, **_model_inputs): updates = [] runner._update_full_graph_params_if_needed = lambda *args: updates.append(args) + input_ids = object() result = runner._model_forward( 8, - input_ids=None, + input_ids=input_ids, positions=object(), intermediate_tensors=None, inputs_embeds=None, @@ -487,6 +492,7 @@ def __call__(self, **_model_inputs): assert result == "hidden_states" assert len(updates) == expected_updates + assert forward_context.input_ids is input_ids def test_npu_attention_runner_installs_mla_graph_wrapper(monkeypatch): @@ -641,7 +647,12 @@ def test_npu_attention_runner_sends_per_ubatch_dp_metadata(): ), ] - runner._send_dp_metadata(None, ubatch_slices) + torch = pytest.importorskip("torch") + runner._send_dp_metadata( + None, + ubatch_slices, + input_ids=torch.tensor([10, 11, 12, 13, 20, 21, 22]), + ) dp_metadata_list = runner.connector.dp_metadata_updates[0][0] assert sorted(dp_metadata_list) == [0, 1] @@ -651,6 +662,10 @@ def test_npu_attention_runner_sends_per_ubatch_dp_metadata(): assert sorted(sent_dp_metadata_list) == [0, 1] assert _tokens(sent_dp_metadata_list[0]) == [4] assert _tokens(sent_dp_metadata_list[1]) == [3] + assert runner.connector.sent_dp_metadata_payloads[0].input_ids_by_stage == { + 0: [10, 11, 12, 13], + 1: [20, 21, 22], + } def test_npu_attention_capture_microbatch_also_captures_single_stage(): @@ -1349,6 +1364,37 @@ def test_npu_ffn_runner_executes_eager_ffn_step(monkeypatch): ] +def test_npu_ffn_input_ids_require_exact_hidden_state_alignment(): + _require_npu_runtime() + from afd_plugin.v1.worker.npu.ffn_model_runner import _install_ffn_input_ids + + torch = pytest.importorskip("torch") + forward_context = SimpleNamespace() + _install_ffn_input_ids( + forward_context, + [10, 11], + num_tokens=2, + device=torch.device("cpu"), + ) + assert forward_context.input_ids.tolist() == [10, 11] + + with pytest.raises(RuntimeError, match="must align"): + _install_ffn_input_ids( + forward_context, + [12], + num_tokens=2, + device=torch.device("cpu"), + ) + + _install_ffn_input_ids( + forward_context, + None, + num_tokens=2, + device=torch.device("cpu"), + ) + assert forward_context.input_ids is None + + def test_npu_ffn_runner_builds_forward_context_for_each_dbo_stage(monkeypatch): _require_npu_runtime() from afd_plugin.v1.worker.npu import ffn_model_runner @@ -1499,7 +1545,7 @@ def fake_ascend_forward_context(**kwargs): shared_num_tokens=2, ) - def recv_ffn_work_item(*, stage_idx, max_num_tokens): + def recv_ffn_work_item(*, stage_idx, max_num_tokens, expected_layer_idx): assert stage_idx == 0 assert max_num_tokens == 16 return work_item From 8784835c5c01ffb04e737b202b193e54b4f191d2 Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Mon, 24 Aug 2026 17:10:19 +0800 Subject: [PATCH 02/19] Add DSV4 recipes and experiment tools Signed-off-by: bjf-frz --- ...4_compile_aclgraph_pipeline_performance.md | 359 +++ docs/experiments/DSV4_shape_walkthrough.md | 2404 +++++++++++++++++ docs/experiments/cli.md | 182 ++ .../deepseek_v4/A3_V026_DEPLOY.md | 125 + .../deepseek_v4/README.md | 63 + .../deepseek_v4/attention_tp8.sh | 120 + .../deepseek_v4/diagnose_worker_startup.sh | 79 + .../deepseek_v4/ffn_ep8.sh | 128 + .../deepseek_v4/launch_8a8f.sh | 25 + .../deepseek_v4/launch_dp2tp4.sh | 24 + recipe/npu/legacy_experiments/attn.sh | 70 + recipe/npu/legacy_experiments/d0.sh | 77 + recipe/npu/legacy_experiments/d1.sh | 79 + .../npu/legacy_experiments/dsv4_baseline.sh | 44 + recipe/npu/legacy_experiments/ffn.sh | 70 + recipe/npu/legacy_experiments/ffn1.sh | 70 + recipe/npu/legacy_experiments/pd_colo.sh | 58 + tools/analysis/legacy/analyze2.py | 32 + tools/analysis/legacy/analyze_ffn_log.py | 24 + tools/analysis/legacy/dbg.py | 9 + tools/analysis/legacy/dbg2.py | 15 + tools/analysis/legacy/fix_launch_scripts.py | 15 + tools/analysis/legacy/fix_ubatch_wrapper.py | 27 + tools/analysis/legacy/gen_b64.py | 11 + tools/analysis/legacy/launch_dp.py | 74 + tools/analysis/legacy/patch_layer_fix.py | 34 + .../analysis/legacy/rewrite_test_force_lb.py | 268 ++ tools/benchmarks/legacy/bench_pd.py | 137 + tools/benchmarks/legacy/merge_throughput.py | 96 + tools/benchmarks/legacy/throughput.py | 763 ++++++ tools/benchmarks/legacy/throughput_v2.py | 821 ++++++ tools/itask/probe_net.sh | 15 + tools/itask/run_itask.sh | 8 + tools/itask/wsl_itask_probe.sh | 60 + 34 files changed, 6386 insertions(+) create mode 100644 docs/experiments/DSV4_compile_aclgraph_pipeline_performance.md create mode 100644 docs/experiments/DSV4_shape_walkthrough.md create mode 100644 docs/experiments/cli.md create mode 100644 recipe/npu/CAMAsyncAFDConnector/deepseek_v4/A3_V026_DEPLOY.md create mode 100644 recipe/npu/CAMAsyncAFDConnector/deepseek_v4/README.md create mode 100644 recipe/npu/CAMAsyncAFDConnector/deepseek_v4/attention_tp8.sh create mode 100644 recipe/npu/CAMAsyncAFDConnector/deepseek_v4/diagnose_worker_startup.sh create mode 100644 recipe/npu/CAMAsyncAFDConnector/deepseek_v4/ffn_ep8.sh create mode 100644 recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_8a8f.sh create mode 100644 recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_dp2tp4.sh create mode 100644 recipe/npu/legacy_experiments/attn.sh create mode 100644 recipe/npu/legacy_experiments/d0.sh create mode 100644 recipe/npu/legacy_experiments/d1.sh create mode 100644 recipe/npu/legacy_experiments/dsv4_baseline.sh create mode 100644 recipe/npu/legacy_experiments/ffn.sh create mode 100644 recipe/npu/legacy_experiments/ffn1.sh create mode 100644 recipe/npu/legacy_experiments/pd_colo.sh create mode 100644 tools/analysis/legacy/analyze2.py create mode 100644 tools/analysis/legacy/analyze_ffn_log.py create mode 100644 tools/analysis/legacy/dbg.py create mode 100644 tools/analysis/legacy/dbg2.py create mode 100644 tools/analysis/legacy/fix_launch_scripts.py create mode 100644 tools/analysis/legacy/fix_ubatch_wrapper.py create mode 100644 tools/analysis/legacy/gen_b64.py create mode 100644 tools/analysis/legacy/launch_dp.py create mode 100644 tools/analysis/legacy/patch_layer_fix.py create mode 100644 tools/analysis/legacy/rewrite_test_force_lb.py create mode 100644 tools/benchmarks/legacy/bench_pd.py create mode 100644 tools/benchmarks/legacy/merge_throughput.py create mode 100644 tools/benchmarks/legacy/throughput.py create mode 100644 tools/benchmarks/legacy/throughput_v2.py create mode 100644 tools/itask/probe_net.sh create mode 100644 tools/itask/run_itask.sh create mode 100644 tools/itask/wsl_itask_probe.sh diff --git a/docs/experiments/DSV4_compile_aclgraph_pipeline_performance.md b/docs/experiments/DSV4_compile_aclgraph_pipeline_performance.md new file mode 100644 index 00000000..0e6c3a24 --- /dev/null +++ b/docs/experiments/DSV4_compile_aclgraph_pipeline_performance.md @@ -0,0 +1,359 @@ +# DSV4 Flash?Compile + ACL Graph ?????????? + +## 1. ??????? + +???? DeepSeek-V4 Flash ??????????????????????????? + +- `torch.compile`???? +- ACL Graph???? +- Graph ???`FULL_DECODE_ONLY`? +- ?????`B_compile_aclgraph_ops_only`? +- ???????? stack ? profile? +- ?????Rank 12? +- ?????DP4?TP4?EP16? +- ???16 token? +- ???32 token? +- Decode ?????31 ??????? token ??? step? + +???????? eager ????????? stack/profile memory ??????? Prefill ????????????? `torch.compile + ACL Graph` ?? ops-only profile? + +## 2. ???? + +```text +Attention stage wall + = Attention HcPre ?? + ? Attention ?? + ? Attention HcPost ?? + +MoE stage wall + = FFN HcPre ?? + ? Router / Dispatch / Expert / Combine + ? FFN HcPost ?? +``` + +Stage wall ?? mHC Pre/Post??????projection?quant?cache ????????????? stream ??? + +- `stage wall`?timeline ????????????? +- `kernel ??`?stage ??? device task duration ??? +- `kernel union`???? stream ????????????? +- `???? = kernel ?? - kernel union`? +- `??/?? = stage wall - kernel union`? + +## 3. Profile ??????? + +```text +layer 0, 1 : SWA +layer 2, 4, ..., 42: CSA?? 21 ? +layer 3, 5, ..., 41: HCA?? 20 ? + +layer 0, 1, 2 : Hash MoE +layer 3, ..., 42 : Learned MoE +``` + +????????? profile ? layer 42 ??? CSA???? SWA??? model forward ?? + +```text +VllmQuantLightningIndexer = 21 ? +Compressor = 62 ? + = 21 ? 2?CSA?+ 20 ? 1?HCA? +``` + +?? timeline ??? layer 42 ?? Compressor ? Lightning Indexer???????? CSA ???? + +## 4. Decode ????? + +```text +layer 0 : SWA ? Hash MoE +layer 1 : SWA ? Hash MoE +layer 2 : CSA ? Hash MoE +layer 3 : HCA ? Learned MoE +layer 4 : CSA ? Learned MoE +layer 5 : HCA ? Learned MoE +... +layer 41: HCA ? Learned MoE +layer 42: CSA ? Learned MoE +? Final Norm +? LM Head +? ArgMax +? ??? token +``` + +??? step ?????`Decode step N ? token N ? Decode step N+1`?????? token ????? token-level pipeline? + +### 4.1 ???????? + +| ?? | ???? | P50 | ?? | ?? | ???? | +|---|---:|---:|---:|---:|---:| +| Attention HcPost ? ??? FFN HcPre | 1.20 ?s | 1.18 ?s | 0.86 ?s | 2.66 ?s | 0 | +| FFN HcPost ? ??? Attention HcPre | 1.15 ?s | 1.06 ?s | 0.86 ?s | 2.72 ?s | 0 | + +?????????????????? + +```text +Attention(layer i) +? MoE(layer i) +? Attention(layer i+1) +``` + +??? `MoE(layer i) || Attention(layer i+1)` ?????????? + +## 5. Decode ????? + +31 ??? Decode step ??? TPOT ? 33.45 ms????? 33.18 ms? + +| ?? | ?? | ? Decode stage wall | ???? | ?? P50 | +|---|---:|---:|---:|---:| +| SWA | 2 | 0.509 ms | 254.5 ?s | 258.1 ?s | +| CSA | 21 | 9.503 ms | 452.5 ?s | 452.1 ?s | +| HCA | 20 | 6.172 ms | 308.6 ?s | 306.1 ?s | +| Hash MoE | 3 | 1.004 ms | 334.6 ?s | 334.4 ?s | +| Learned MoE | 40 | 12.812 ms | 320.3 ?s | 318.9 ?s | +| ???? | 43 Attention + 43 MoE | 30.001 ms | ? | ? | + +?? stage ??? 3.45 ms????? slot mapping?stage ??????Final Norm?LM Head?sampling/ArgMax?graph launch ??? stream ??? + +- CSA ? HCA ???? 46.6%? +- HCA ? SWA ???? 21.3%? +- Hash MoE ? Learned MoE ???? 4.5%? +- CSA ?? 9.50 ms?? Attention ?????? +- Learned MoE ?? 12.81 ms???? Decode ????????? + +## 6. SWA ??? + +```text +Attention HcPre +? Q/KV projection ? DynamicQuant +? RoPE +? KV cache Scatter/?? +? window=128 ? SparseAttnSharedkv +? grouped O projection +? TP/collective +? Attention HcPost +``` + +SWA ?? Compressor ? Lightning Indexer? + +| ?? | ? Decode | ???? | +|---|---:|---:| +| stage wall | 0.509 ms | 254.5 ?s | +| kernel ?? | 0.500 ms | 250.1 ?s | +| kernel union | 0.455 ms | 227.4 ?s | +| ? stream ???? | 0.045 ms | 22.7 ?s | +| ??/?? | 0.054 ms | 27.2 ?s | + +SWA ?????? Attention stage ??????????????? MoE ????? + +## 7. CSA ??? + +```text +Attention HcPre +? Q/KV projection +? C4 Compressor +? Indexer Compressor +? Lightning Indexer TopK +? ?? KV?window KV ????? +? RoPE +? SparseAttnSharedkv +? grouped O projection +? TP/collective +? Attention HcPost +``` + +| ?? | ? Decode | ???? | +|---|---:|---:| +| stage wall | 9.503 ms | 452.5 ?s | +| kernel ?? | 9.354 ms | 445.4 ?s | +| kernel union | 8.499 ms | 404.7 ?s | +| ? stream ???? | 0.855 ms | 40.7 ?s | +| ??/?? | 1.004 ms | 47.8 ?s | + +CSA ???? Attention??????????? Compressor?Lightning Indexer ??????????? `SparseAttnSharedkv` ???? + +## 8. HCA ??? + +```text +Attention HcPre +? Q/KV projection +? C128 Compressor +? ?? KV ? local window ???? +? RoPE +? SparseAttnSharedkv +? grouped O projection +? TP/collective +? Attention HcPost +``` + +HCA ?? Lightning Indexer? + +| ?? | ? Decode | ???? | +|---|---:|---:| +| stage wall | 6.172 ms | 308.6 ?s | +| kernel ?? | 5.950 ms | 297.5 ?s | +| kernel union | 5.514 ms | 275.7 ?s | +| ? stream ???? | 0.435 ms | 21.8 ?s | +| ??/?? | 0.658 ms | 32.9 ?s | + +HCA ? Compressor ?????? CSA Compressor ? 42%????? Lightning Indexer??????? CSA? + +## 9. Hash MoE ??? + +```text +FFN HcPre +? Hash Router / Top6 +? EP Dispatch +? Local Expert Gate/Up GroupedMatmul +? SwiGLU + Quant +? Expert Down GroupedMatmul +? EP Combine +? Shared Expert ?? +? FFN HcPost +``` + +| ?? | ? Decode | ???? | +|---|---:|---:| +| stage wall | 1.004 ms | 334.6 ?s | +| kernel ??/union | 0.935 ms | 311.8 ?s | +| kernel ?? | 0 | 0 | +| ??/?? | 0.069 ms | 22.8 ?s | + +Dispatch?Expert MLP ? Combine ????????????? Expert ????? + +## 10. Learned MoE ??? + +```text +FFN HcPre +? Learned Router projection +? Top6 ? routing weight +? EP Dispatch +? Local Expert Gate/Up GroupedMatmul +? SwiGLU + Quant +? Expert Down GroupedMatmul +? EP Combine +? Shared Expert ?? +? FFN HcPost +``` + +| ?? | ? Decode | ???? | +|---|---:|---:| +| stage wall | 12.812 ms | 320.3 ?s | +| kernel ??/union | 11.891 ms | 297.3 ?s | +| kernel ?? | 0 | 0 | +| ??/?? | 0.921 ms | 23.0 ?s | + +Learned Router projection ?????????????? Expert Gate/Up?EP Dispatch?EP Combine ? Expert Down? + +## 11. Prefill ??? + +?? Prefill ??? 16 token?`FULL_DECODE_ONLY` ?? Decode graph ????????????? compile + ACL Graph ??? Prefill ? dynamic/eager stage ??? + +| ?? | ?? | Prefill stage wall | kernel ?? | +|---|---:|---:|---:| +| SWA | 2 | 13.75 ms | 3.83 ms | +| CSA | 21 | 73.37 ms | 38.91 ms | +| HCA | 20 | 31.10 ms | 22.94 ms | +| Hash MoE | 3 | 19.55 ms | 5.67 ms | +| Learned MoE | 40 | 230.25 ms | 53.81 ms | + +Prefill ? stage wall ? kernel ????????? Learned MoE?230.25 ms wall ???? 53.81 ms ? device task duration???????????Host ??? stream ????? Prefill ?????????????????????? + +## 12. ??????? Attention + +???????? Decode step ???????? + +### 12.1 SWA ?? + +| ??/?? | 2 ??? | ???? | +|---|---:|---:| +| `TransposeBatchMatMul`?O projection | 0.103 ms | 51.7 ?s | +| `QuantBatchMatmulV3`?Q/KV projection | 0.086 ms | 42.9 ?s | +| `HcPre` | 0.065 ms | 32.4 ?s | +| `SparseAttnSharedkv` | 0.043 ms | 21.7 ?s | +| `AivKernel`?collective | 0.043 ms | 21.3 ?s | +| `MatMulV2` ? projection | 0.040 ms | 19.8 ?s | +| `InplacePartialRotaryMul` | 0.024 ms | 11.9 ?s | +| `DynamicQuant` | 0.022 ms | 11.1 ?s | +| `RmsNorm` | 0.021 ms | 10.7 ?s | +| `ScatterNdUpdateV2` | 0.009 ms | 4.5 ?s | +| `HcPost` | 0.013 ms | 6.6 ?s | + +### 12.2 CSA ?? + +| ??/?? | 21 ??? | ???? | +|---|---:|---:| +| `Compressor` | 2.065 ms | 98.3 ?s | +| `QuantBatchMatmulV3` | 1.174 ms | 55.9 ?s | +| `TransposeBatchMatMul`?O projection | 1.071 ms | 51.0 ?s | +| `MatMulV2` ? projection | 0.877 ms | 41.8 ?s | +| `HcPre` | 0.688 ms | 32.8 ?s | +| `AivKernel`?collective | 0.493 ms | 23.5 ?s | +| `SparseAttnSharedkv` | 0.463 ms | 22.1 ?s | +| `DynamicQuant` | 0.427 ms | 20.3 ?s | +| `VllmQuantLightningIndexer` | 0.372 ms | 17.7 ?s | +| `ScatterNdUpdateV2` | 0.356 ms | 16.9 ?s | +| `InplacePartialRotaryMul` | 0.332 ms | 15.8 ?s | +| `RmsNorm` | 0.202 ms | 9.6 ?s | +| `CompressorMetadata` | 0.128 ms | 6.1 ?s | +| `HcPost` | 0.136 ms | 6.5 ?s | + +### 12.3 HCA ?? + +| ??/?? | 20 ??? | ???? | +|---|---:|---:| +| `TransposeBatchMatMul`?O projection | 1.015 ms | 50.8 ?s | +| `Compressor` | 0.817 ms | 40.9 ?s | +| `QuantBatchMatmulV3` | 0.816 ms | 40.8 ?s | +| `HcPre` | 0.660 ms | 33.0 ?s | +| `AivKernel`?collective | 0.470 ms | 23.5 ?s | +| `SparseAttnSharedkv` | 0.470 ms | 23.5 ?s | +| `MatMulV2` ? projection | 0.413 ms | 20.7 ?s | +| `InplacePartialRotaryMul` | 0.229 ms | 11.5 ?s | +| `DynamicQuant` | 0.213 ms | 10.6 ?s | +| `RmsNorm` | 0.181 ms | 9.0 ?s | +| `ScatterNdUpdateV2` | 0.160 ms | 8.0 ?s | +| `CompressorMetadata` | 0.055 ms | 2.7 ?s | +| `HcPost` | 0.126 ms | 6.3 ?s | + +## 13. ??????? MoE + +### 13.1 Hash MoE ?? + +| ??/?? | 3 ??? | ???? | +|---|---:|---:| +| `MoeDistributeCombineV2` | 0.196 ms | 65.2 ?s | +| `MoeDistributeDispatchV2` | 0.164 ms | 54.7 ?s | +| `GroupedMatmulSwigluQuant`?Gate/Up | 0.135 ms | 45.1 ?s | +| `QuantBatchMatmulV3` | 0.116 ms | 38.7 ?s | +| `HcPre` | 0.097 ms | 32.2 ?s | +| `GroupedMatmul`?Down | 0.079 ms | 26.2 ?s | +| `MatMulV3`?Router/?? projection | 0.048 ms | 15.9 ?s | +| `MoeGatingTopKHash` | 0.013 ms | 4.4 ?s | +| `DequantSwigluQuant` | 0.012 ms | 3.9 ?s | +| `HcPost` | 0.018 ms | 6.0 ?s | + +### 13.2 Learned MoE ?? + +| ??/?? | 40 ??? | ???? | +|---|---:|---:| +| `GroupedMatmulSwigluQuant`?Gate/Up | 2.201 ms | 55.0 ?s | +| `MoeDistributeCombineV2` | 1.958 ms | 48.9 ?s | +| `MoeDistributeDispatchV2` | 1.950 ms | 48.7 ?s | +| `QuantBatchMatmulV3` | 1.526 ms | 38.2 ?s | +| `HcPre` | 1.321 ms | 33.0 ?s | +| `GroupedMatmul`?Down | 1.282 ms | 32.0 ?s | +| `MatMulV3`?Router projection | 0.586 ms | 14.7 ?s | +| `MoeGatingTopKHash` | 0.162 ms | 4.1 ?s | +| `DequantSwigluQuant` | 0.158 ms | 3.9 ?s | +| `HcPost` | 0.241 ms | 6.0 ?s | + +`MoeGatingTopKHash` ??? profile ??????????????????? 40 ???? Hash Router?Hash/Learned ?????????? Router ?????????????????? + +## 14. ???? + +1. Decode ????? token??????????? token ??? pipeline? +2. Attention ??????? stream ?????MoE ?????? Dispatch?Expert ? Combine ???? +3. CSA ??? 452.5 ?s???????? Compressor ? Lightning Indexer ????? +4. HCA ??? 308.6 ?s??? Lightning Indexer?Compressor ?????? CSA? +5. SWA ??? 254.5 ?s????? Attention? +6. Hash MoE ? Learned MoE ???? 334.6 ?s ? 320.3 ?s? +7. Learned Router projection ???????Dispatch?Combine ? Expert GroupedMatmul ??????? +8. Prefill ???????????? stream ???????? kernel ????????? diff --git a/docs/experiments/DSV4_shape_walkthrough.md b/docs/experiments/DSV4_shape_walkthrough.md new file mode 100644 index 00000000..13473fa4 --- /dev/null +++ b/docs/experiments/DSV4_shape_walkthrough.md @@ -0,0 +1,2404 @@ +# DeepSeek-V4-Flash 推理 Shape 全流程 + +> 当前 A3 测试实例：模型实际输入 16 tokens，生成 32 tokens。覆盖 Prefill、Decode、混合注意力、MoE、LM Head、MTP 数据通路及 P/D 分离语义。 + +## 1. 证据规则与符号 + +为避免把推导误认为 profiler 实测，本文统一标记： + +- **[实测]**：rank 12 Ascend profiler 的 operator shape。 +- **[配置]**：模型配置或启动参数。 +- **[代码]**：当前 vLLM / vLLM-Ascend / DSV4 实现确定的数据流。 +- **[推导]**：由代码和 TP/DP/EP 规则得到，尚未从 Decode operator CSV 直接截取。 + +| 符号 | 含义 | 当前值 | +|---|---|---:| +| `B` | 有效请求数 | 1 | +| `S` | 序列长度 | Prefill 16；Decode 每步 1 | +| `T` | 展平 token 数 | `B×S` | +| `H` | hidden size | 4096 | +| `C` | mHC residual streams | 4 | +| `L` | Transformer 层数 | 43 | +| `V` | vocabulary size | 129280 | +| `N_h` | attention heads | 64；TP rank-local 16 | +| `D_h` | attention head dim | 512 | +| `TP/DP/EP` | 并行度 | 4 / 4 / 16 | + +## 2. 当前测试实例 + +| 项目 | 当前值 | +|---|---| +| 模型 | `DeepSeek-V4-Flash-w8a8-mtp` | +| 硬件 | Ascend A3，16 卡 | +| block size | 128 | +| `max_model_len` | 8192 | +| `max_num_batched_tokens` | 4096 | +| `max_num_seqs` | 4 | +| 原始随机 prompt | 12 tokens | +| Chat template 后实际输入 | 16 tokens | +| 请求输出 | 32 tokens | +| profile | eager A，rank 12，带 stack/memory/shape | + +```text +原始 prompt 12 tokens + → chat template 后 16 tokens + → Prefill forward 一次，采样 y1 + → Decode forward 31 次，依次采样 y2...y32 +``` + +没有提前 EOS 时，总 forward 次数为 `1+31=32`。最后的 `y32` 采样后不再送回模型，因此最终已处理 KV 长度是 `16+31=47`，不是 48。 + +## 3. 三种 shape 视角 + +### 3.1 请求逻辑视角 + +```text +Prefill: [B=1,S=16] → T=16 +Decode : [B=1,S=1] → T=1（每步一个有效 token） +``` + +### 3.2 mHC residual 视角 + +```text +普通 hidden [T_local,4096] +→ 4-stream residual [T_local,4,4096] + +Prefill residual [4,4,4096] [实测] +Decode residual [1,4,4096] [实测] +``` + +Prefill 的逻辑 token 数是 16，但 sequence parallel 在 TP=4 内将 residual token 分为每 rank 4 个。`[4,4,4096]` 的第一个 4 是 local tokens，第二个 4 才是 mHC streams。 + +### 3.3 Attention 计算视角 + +```text +Prefill residual [4,4,4096] + → hc_pre [4,4096] + → TP all-gather [16,4096] + → local Q [16,16,512] + +Decode residual [1,4,4096] + → hc_pre [1,4096] + → graph/TP 对齐后的物理 token 维度 4 + → local Q [4,16,512] [推导] +``` + +Decode attention 第一维的 4 是图执行/通信对齐的物理槽位，不代表有 4 个有效请求；当前只有 1 个有效 token。 + +## 4. 从文本到 mHC residual + +```text +文本 prompt +→ 原始 token_ids [12] +→ chat template 后 input_ids [1,16] +→ 展平 token_ids [16] +``` + +全局 embedding table 为 `[129280,4096]`；TP=4 后每 rank 持有 `[32320,4096]`。 + +```text +Prefill: ids [16] → logical hidden [16,4096] + → SP local [4,4096] + → mHC [4,4,4096] + +Decode : id [1] → hidden [1,4096] + → mHC [1,4,4096] +``` + +## 5. 43 层混合注意力排布 + +```text +layer 0 : SWA +layer 1 : SWA +layer 2 : CSA，compression stride=4 +layer 3 : HCA，compression stride=128 +layer 4 : CSA +layer 5 : HCA +... +layer 40: CSA +layer 41: HCA +layer 42: SWA +``` + +即 `[SWA, SWA, (CSA-c4, HCA-c128)×20, SWA]`。 + +- **DSA** 是整个稀疏/混合 attention backend 的总称，不是第四种 layer。 +- **SWA** 是滑窗 attention。 +- **CSA** 使用 c4 压缩 KV、Lightning Indexer 和 TopK 稀疏选择。 +- **HCA** 使用 c128 层级压缩状态。 + +每层公共接口： + +```text +x [T_local,4,4096] + → hc_pre(attention) [T_local,4096] + → Attention + → hc_post [T_local,4,4096] + → hc_pre(FFN) [T_local,4096] + → FFN/MoE + → hc_post [T_local,4,4096] +``` + +因此层与层之间始终是 Prefill `[4,4,4096]`、Decode `[1,4,4096]`。 + +## 6. mHC Pre / Post + +`npu_hc_pre_v2`： + +```text +Prefill input [4,4,4096] [实测] → output [4,4096] [代码] +Decode input [1,4,4096] [实测] → output [1,4096] [代码] +``` + +它用学习到的 mixing 系数把 4 条 residual stream 合成一个 H=4096 hidden。子层完成后： + +```text +sub-layer output [T_local,4096] ++ residual state [T_local,4,4096] +→ hc_post updated state [T_local,4,4096] +``` + +## 7. Attention Q、Cache 与输出投影 + +全局 64 heads，TP=4 后每 rank 计算 16 heads： + +```text +Prefill Q [16,16,512] [实测] + │ │ └─ head_dim + │ └──── local heads=64/4 + └─────── attention token 数 + +Decode Q [4,16,512] [推导，物理图槽位为 4] +``` + +物理 cache capacity： + +```text +SWA KV cache [10123,128,1,512] [实测] +CSA c4 state [10123,8,2048] [实测] +HCA c128 state [10123,32,1024] [实测] +``` + +`10123` 是整个运行预分配的物理 block 容量，不是当前请求长度；请求通过 block table 只引用其中少量 block。 + +输出投影： + +```text +local attention output [T_attn,16,512] +→ reshape [T_attn,2,4096] # 8 个全局 output groups / TP4 +→ wo_a [T_attn,2,1024] +→ flatten [T_attn,2048] +→ wo_b + TP reduce-scatter +→ [T_local,4096] +``` + +所以 Prefill `[16,16,512] → [4,4096]`；Decode `[4,16,512] → [1,4096]` [推导]，再由 hc_post 写回 4 streams。 + +## 8. SWA / CSA / HCA 的关键 shape + +### 8.1 SWA +SWA（Sliding Window Attention）仍然是 causal attention，但每个 query 最多只读取最近 128 个有效 token，而不是读取从序列开头到当前位置的全部 KV。 + +#### 8.1.1 当前 SWA 的输入和位置 + +当前 layer 排布中，纯 SWA 位于： + +```text +layer 0、layer 1、layer 42 +``` + +CSA/HCA 层内部也包含局部 SWA 分支，但还会额外使用压缩 KV 或 Indexer 结果。纯 SWA 的 Prefill profiler 输入： + +```text +Q [16,16,512] [实测] +shared-KV cache [10123,128,1,512] [实测] +Attention output [16,16,512] [代码] +``` + +其中： + +- `16`（Q 第一维）：本次 Prefill 的 token 数。 +- `16`（Q 第二维）：TP=4 后每 rank 的 local query heads。 +- `512`：每个 query head 的维度。 +- `10123`：运行时预分配的物理 KV block 总容量。 +- cache 第二维 `128`：当前 KV block size。 +- cache 第三维 `1`：shared-KV，即多个 query heads 共享一组 KV 表示。 + +需要特别区分：当前 `block_size=128`，SWA window 也恰好为 128，但两者语义不同： + +```text +block_size=128：KV cache 的物理分页粒度 +window=128 ：每个 query 最多允许参与 attention 的逻辑历史长度 +``` + +这里只是两个配置值碰巧相同，不能把一个 cache block 直接等同于 SWA 数学定义。 + +#### 8.1.2 SWA 的完整计算流程 + +进入 SWA 子层前，rank-local residual 是： + +```text +Prefill X [4,4,4096] +Decode X [1,4,4096] +``` + +第一步，mHC 生成 mapping 并执行 `hc_pre`： + +```text +Prefill [4,4,4096] → [4,4096] +Decode [1,4,4096] → [1,4096] +``` + +第二步，TP all-gather 恢复 attention token 维度： + +```text +Prefill [4,4096] → [16,4096] +Decode [1,4096] → physical graph slots [4,4096] [推导] +``` + +第三步，归一化、量化并生成 Q 与 shared KV。 + +##### Q LoRA：down projection 和 up projection + +当前 Q 不是用一个 `Linear(4096,64×512)` 直接生成，而是经过两级低秩 projection： + +```text +hidden size H=4096 +Q LoRA rank Rq=1024 +全局 Q heads 64 +TP 4 +rank-local heads 64/4=16 +head_dim 512 +``` + +第一级 `q_a_proj` 是 Q down projection： + +```text +X [16,4096] +W_q_a [1024,4096] # 按 [out,in] 表示 +X · W_q_aᵀ [16,1024] +``` + +这里的“down”指特征维从 4096 压到低秩空间 1024： + +```text +4096 → 1024 +``` + +随后执行 Q LoRA 中间归一化/动态量化，shape 不变： + +```text +q_low_rank [16,1024] +→ RMSNorm / DynamicQuant [16,1024] +``` + +第二级 `q_b_proj` 将低秩表示升回所有 Q heads。全局输出维度是： + +```text +64 heads × 512 = 32768 +``` + +TP=4 按 head/output 维切分后，每 rank 只计算： + +```text +16 local heads × 512 = 8192 +``` + +所以 rank-local up projection 是： + +```text +q_low_rank [16,1024] +W_q_b_local [8192,1024] +q_flat_local [16,8192] +reshape [16,16,512] +``` + +完整 Q 路径： + +```text +[16,4096] +→ q_a_proj / down_proj +[16,1024] +→ RMSNorm + DynamicQuant +[16,1024] +→ q_b_proj / TP-local up_proj +[16,8192] +→ reshape，8192=16×512 +[16,16,512] +``` + +第一维 16 始终是 Prefill token 数；reshape 新出现的第二维 16 是当前 TP rank 的 Q head 数，不是又生成了 16 份 token。 + +##### W8A8 在 projection 中做什么 + +`W8A8` 表示量化矩阵乘中： + +```text +W8：weight 使用 8-bit 表示，通常为 INT8 +A8：activation 使用 8-bit 表示，通常为动态 INT8 +``` + +概念计算为： + +```text +x_int8 = quantize(x_fp, x_scale) +w_int8 = quantize(w_fp, w_scale) +acc = matmul(x_int8, w_int8ᵀ) # 通常用更高精度累加 +output = dequantize(acc, x_scale, w_scale) +``` + +因此 Q down projection 可以理解为： + +```text +activation [16,4096] --DynamicQuant--> INT8 [16,4096] +weight [1024,4096] INT8 [1024,4096] +INT8 matmul / higher-precision accumulate +→ dequantized q_low_rank [16,1024] +``` + +Q up projection同理： + +```text +activation [16,1024] --DynamicQuant--> INT8 [16,1024] +weight [8192,1024] INT8 [8192,1024] +→ q_flat_local [16,8192] +``` + +当前 profile 中与这条路径对应的主要算子包括： + +```text +RmsNorm / DynamicQuant / RmsNormDynamicQuant +QuantBatchMatmulV3 +npu::npu_quant_matmul +view / unflatten / reshape +``` + +W8A8 改变的是存储 dtype、带宽、矩阵乘内核和缩放/反量化过程，不改变上述逻辑 shape。具体 scale 是 per-token、per-channel 还是其他粒度，应以对应量化 linear 的配置和 kernel 参数为准。 + +##### shared KV 支路 + +与 Q LoRA 并行，当前 hidden 还会生成 shared K/V representation： +KV down projection 的思路与 Q down projection 相同：都把 `[T,4096]` 压到更小的 latent space。区别是 Q 随后还通过 `q_b_proj` 升维并拆成 16 个 TP-local query heads，而当前 KV latent 直接保持为 1 个 shared head、每 token 512 维，经过归一化/Partial RoPE 后写入 shared-KV cache： + +`[T,4096] → KV down_proj → [T,512] → [T,1,512] → cache`。 + +```text +shared KV current tokens [16,1,512] +``` + +所以进入后续 Partial RoPE 和 cache 写入前，两个关键结果为： + +```text +Q projection result [16,16,512] +shared KV current tokens [16,1,512] +``` + +这里 64 个全局 query heads 在 TP=4 后变成 16 个 local heads；shared KV 的 head 数为 1。 + +第四步，对 Q/K 的部分维度执行位置编码： + +```text +K-like shared representation [16,1,1,512] +Q representation [16,1,16,512] +→ InplacePartialRotaryMul +``` + +`partial` 表示 512 维 head 中只有配置指定的 RoPE 子空间参与旋转，其余维度保持原有内容。 + +第五步，把当前 token 的 shared KV 写入 paged KV cache： + +```text +current shared KV [16,1,512] +slot_mapping [16] +physical cache [10123,128,1,512] +→ ScatterNdUpdateV2 +``` + +`slot_mapping` 将每个逻辑 token 映射到 `(physical block, offset)`。写入 cache 和“这个 token 是否参与当前 query 的 SWA”是两件事：前者负责保存状态，后者由 attention metadata/window 决定。 + +第六步，生成 attention metadata： + +```text +seq_lens / query positions / block table / slot mapping ++ causal constraint ++ sliding window=128 +→ SparseAttnSharedkvMetadata +``` + +它为核心 kernel 准备每个 query 可以访问的逻辑 KV 范围及物理 cache 映射。 + +第七步，融合计算局部 attention： + +```text +Q [T_attn,16,512] ++ shared-KV cache [10123,128,1,512] ++ SWA metadata +→ score = QKᵀ × scale +→ causal/window mask +→ softmax +→ probability × V +→ output [T_attn,16,512] +``` + +这些 score、mask、softmax 和 value aggregation 在 `SparseAttnSharedkv` 中融合完成，不会在 profile 中表现成一串独立的 `MatMul → MaskedFill → Softmax → MatMul` PyTorch 算子。 + +第八步，执行 grouped O projection。这里必须把 `wo_b` 本地矩阵乘和 TP reduce-scatter 分开理解。 + +第一阶段：将当前 rank 的 16 个 heads 分为 2 个 output groups。每组包含 8 个 heads： + +```text +attention output [16,16,512] +reshape [16,2,8,512] +每组拼接 8×512 [16,2,4096] +``` + +第二阶段：每个 group 分别执行 `wo_a` down projection： + +```text +wo_a weight [2,1024,4096] # [group,out,in] +input [16,2,4096] +wo_a output [16,2,1024] +flatten local groups [16,2048] +``` + +其中当前 TP rank 有 2 个 groups，所以 `2048=2×1024`。全局 8 个 groups 对应的概念输入宽度是 `8×1024=8192`。 + +第三阶段：当前 rank 单独执行自己的 `wo_b` weight shard。按照 `[out,in]` 表示： + +```text +local wo_b input [16,2048] +local wo_b weight [4096,2048] +matmul 使用的 weightᵀ [2048,4096] +local partial output [16,4096] +``` + +矩阵乘明确写成： + +```text +[16,2048] @ [2048,4096] +→ [16,4096] +``` + +因此 `wo_b` 本身只完成： + +```text +2048 → 4096 +``` + +它不会把 token 数从 16 变成 4。由于每个 TP rank 只拥有全局 8 个 groups 中的 2 个，四个 rank 分别得到四份： + +```text +rank0 partial [16,4096] +rank1 partial [16,4096] +rank2 partial [16,4096] +rank3 partial [16,4096] +``` + +每一份都只是完整 O projection 的部分和。 + +第四阶段才执行 TP reduce-scatter。逻辑上可以拆成 reduce 和 scatter 两步： + +```text +Reduce： +rank0 partial + rank1 partial + rank2 partial + rank3 partial +→ complete O output [16,4096] + +Scatter（沿 token/sequence-parallel 维）： +complete O output [16,4096] +→ TP4 分片 +→ 每 rank [4,4096] +``` + +实际运行时 reduce-scatter 是一个 collective，完整的 `[16,4096]` 不一定在某个 rank 上单独物化；但逻辑 shape 应按上述两步理解。 + +Prefill 的完整 O projection： + +```text +[16,16,512] +→ group reshape [16,2,4096] +→ wo_a [16,2,1024] +→ flatten [16,2048] +→ wo_b local matmul [16,4096] # partial result +→ TP reduce-scatter [4,4096] # rank-local final result +``` + +Decode 对应为： + +```text +physical attention output [4,16,512] +→ group reshape [4,2,4096] +→ wo_a [4,2,1024] +→ flatten [4,2048] +→ wo_b local partial [4,4096] +→ TP reduce-scatter +→ valid local output [1,4096] [推导] +``` + +最后 `hc_post` 把 attention 输出注回四条 residual stream： + +```text +Prefill [4,4096] + old [4,4,4096] → [4,4,4096] +Decode [1,4096] + old [1,4,4096] → [1,4,4096] +``` + +#### 8.1.3 `seq_lens ≤ 128` 时 + +对于 query 位置 `p`（从 0 开始），causal attention 本来只能访问 `[0,p]`。当有效上下文不超过 128 时： + +```text +可访问 KV = [0,p] +参与 token 数 = p+1 +``` + +因此在当前测试中： + +```text +Prefill 后 context=16 +最终 Decode context=47 +47 < 128 +``` + +所有历史 token 都落在窗口内，所以这一组 profile 中 SWA 的数学结果等价于普通 causal full attention；只是执行的仍然是 SWA/DSA kernel。 + +#### 8.1.4 `seq_lens > 128` 时会发生什么 + +设当前 query 的绝对位置为 `p`，window size `W=128`。SWA 只允许访问： + +```text +start = max(0, p-W+1) +end = p +可访问位置 = [start,end] +``` + +因此最多参与 128 个 token。 + +例如： + +```text +p=127：访问 [0,127]，共 128 tokens +p=128：访问 [1,128]，共 128 tokens；token 0 被排除 +p=200：访问 [73,200]，共 128 tokens；[0,72] 不参与该 query +``` + +要区分三类“长度”： + +1. **逻辑序列长度**仍会继续增长，例如 4096、8192；位置编号不会重置。 +2. **本次 SWA 有效 attention 长度**在超过阈值后固定不超过 128。 +3. **物理 KV cache 中是否仍保留旧 token**由 vLLM cache manager、block 复用和其他 attention 层的需要决定；不能仅从 cache capacity `[10123,...]` 判断旧 token 已经被物理删除。 + +对于 DSV4 尤其重要：旧 token 即使不再参与某个 SWA layer，也可能已经进入 CSA/HCA 的 compressed state，供其他 layer 的远程历史建模使用。因此： + +```text +SWA：只看最近 128 tokens +CSA/HCA：通过压缩状态保留更长距离的信息 +``` + +超过 128 后的计算量变化： + +- Prefill：每个 query 最多与 128 个 KV 计算，稳定阶段复杂度约为 `O(S×128)`，而不是 full attention 的 `O(S²)`。 +- Decode：每一步只有一个有效 query，SWA attention 的有效 KV 数固定为最多 128，因此这部分单步计算量不会随完整 context 线性增长。 +- Cache/metadata：序列位置、block table 和压缩状态仍需更新；所以整模型 TPOT 不会因为 SWA 固定窗口而完全与 context length 无关。 + +#### 8.1.5 当前 profile 中的 SWA 算子汇总 + +| 阶段 | Profile 算子/内核 | 主要输入或作用 | 是否 SWA 专属 | +|---|---|---|---| +| mHC 合并 | `HcPre` / `_C_ascend::npu_hc_pre_v2` | `[T_local,4,4096] → [T_local,4096]` | 否，所有子层共用 | +| Norm/量化 | `RmsNorm`、`DynamicQuant`、`RmsNormDynamicQuant` | 为量化 projection 准备 hidden | 否 | +| Q/KV projection | `QuantBatchMatmulV3`、`npu::npu_quant_matmul` | 生成 Q 和 shared KV | 否，但 shape 由 SWA 配置决定 | +| 位置编码 | `InplacePartialRotaryMul` | Q/K 的部分 RoPE 维度 | DSA attention 共用 | +| KV cache 写入 | `ScatterNdUpdateV2` / `_C_ascend::npu_scatter_nd_update_v2` | 写入 `[10123,128,1,512]` cache | shared-KV 路径 | +| SWA metadata | `SparseAttnSharedkvMetadata` | 构造 causal/window/cache 映射 | 纯 SWA 调用中关键 | +| 核心 attention | `SparseAttnSharedkv` / `_C_ascend::npu_sparse_attn_sharedkv` | 融合 score、mask、softmax、V 聚合 | DSA 三类 attention 共用 | +| 输出投影 | `TransposeBatchMatMul`、`QuantBatchMatmulV3`、`MatMulV2` 等 | grouped `wo_a/wo_b` 和 TP 通信前后处理 | 否 | +| mHC 写回 | `HcPost` / `_C_ascend::npu_hc_post` | attention 输出注回 4 streams | 否 | + +`SparseAttnSharedkv` 在当前 profile 中共出现 1376 次： + +```text +43 layers × (1 Prefill + 31 Decode) = 43 × 32 = 1376 +``` + +这证明它是整个 DSA backend 的共享核心，而不是只在 3 个纯 SWA layer 中出现。区分具体 layer 类型要看传入参数： + +- 纯 SWA：Q、shared-KV cache 和 window metadata 有效，compressed/indexer 输入为空。 +- CSA：额外传入 c4 compressed KV 与 Indexer/TopK 选择结果。 +- HCA：额外使用 c128 层级压缩状态。 + +当前纯 SWA 的 profiler record 中，核心 kernel 的主要非空输入是： + +```text +Q [16,16,512] +shared-KV cache [10123,128,1,512] +其他 compressed/indexer tensor 为空 +``` + +所以本次 16→47 token 的 profile 可以确认 SWA 的短上下文执行路径，但不能测出超过 128 后 window mask、cache block 复用和长 context metadata 的真实性能。要验证该边界，至少需要增加 `seq_lens=127/128/129/256` 四组 case。 + +### 8.2 CSA：C4A、Lightning Indexer 与稀疏主 Attention + +CSA 不是简单的串行链路： + +```text +KV cache → compressor → Indexer → window → RoPE +``` + +更准确的结构是：Q/raw shared-KV 先完成 projection 和 RoPE；随后局部 SWA、C4 压缩记忆和轻量 Indexer 构成并行/协作的长程信息路径；最后由 metadata 和 `SparseAttnSharedkv` 融合执行真正的高维 attention。 + +#### 8.2.1 CSA 的计算流程、算子与 shape + +##### 步骤 1：mHC 与 Attention 输入 + +```text +rank-local residual [4,4,4096] +→ HcPre / npu_hc_pre_v2 +rank-local hidden [4,4096] +→ TP all-gather +Attention hidden [16,4096] +``` + +##### 步骤 2：生成 Q 与 raw shared KV + +Q LoRA： + +```text +hidden [16,4096] +→ q_a_proj/down_proj +q low-rank [16,1024] +→ RMSNorm + DynamicQuant + [16,1024] +→ q_b_proj，TP-local +q flat [16,8192] +→ reshape，8192=16×512 +Q [16,16,512] +``` + +shared KV： + +```text +hidden [16,4096] +→ KV down projection +shared KV [16,512] +→ reshape + [16,1,512] +``` + +相关算子： + +```text +RmsNorm / DynamicQuant / RmsNormDynamicQuant +QuantBatchMatmulV3 +npu::npu_quant_matmul +view / reshape / unflatten +``` + +##### 步骤 3：RoPE 在 cache 写入与 TopK 选择之前执行 + +```text +Q [16,1,16,512] +K-like shared representation [16,1,1,512] +→ InplacePartialRotaryMul +``` + +RoPE 作用于 Q/K 向量的配置子空间，不作用于 Indexer 输出、TopK indices 或 softmax 权重。 + +随后把当前 raw shared KV 写入 paged cache： + +```text +current shared KV [16,1,512] +slot_mapping [16] +physical raw KV cache [10123,128,1,512] +→ ScatterNdUpdateV2 +``` + +##### 步骤 4：C4A main compressor 构造 4:1 长期记忆 + +Profiler 直接观测： + +```text +hidden input [16,4096] +wkv / wgate [1024,4096] +compressor state cache [10123,8,2048] +absolute position embedding [4,1024] +norm weight [512] +compressed cos/sin [5,64] +``` + +概念 projection： + +```text +candidate = X · Wkvᵀ [16,1024] +gate = X · Wgateᵀ [16,1024] +``` + +压缩器按照连续 4 个原始 token 更新一个逻辑 compressed group： + +```text +token 0...3 → C4_0 +token 4...7 → C4_1 +token 8...11 → C4_2 +token 12...15 → C4_3 +``` + +`APE [4,1024]` 用于区分一个 c4 group 内 offset 0、1、2、3。逻辑 compressed position 数近似为： + +```text +N_c4 ≈ ceil(L/4) +``` + +`state cache [10123,8,2048]` 是融合 compressor 的物理累积/packing 布局，不能把第二维 8 直接解释成逻辑 c4 position 数。 + +相关算子： + +```text +CompressorMetadata +Compressor +RmsNorm +InplacePartialRotaryMul # compressed position 的 RoPE 子空间 +``` + +##### 步骤 5：Indexer 使用独立的低维检索表示 + +CSA 还维护一条更便宜的 Indexer compressor 路径。Profiler 直接观测： + +```text +hidden input [16,4096] +wkv / wgate [256,4096] +Indexer compressor state [10123,8,512] +absolute position embedding [4,256] +norm weight [128] +compressed cos/sin [5,64] +``` + +Lightning Indexer 的输入： + +```text +Indexer query [16,64,128] +Indexer key cache [10123,128,1,128] +weights [16,64] +query scales [16,64] +key scales [10123,128,1] +block table [1,16] +TopK 1024 +``` + +它在 128 维检索空间中计算相关性： + +```text +index_score[i,h,j] + = index_Q[i,h,:] · index_K[j,0,:] +``` + +概念 score 和 TopK 输出： + +```text +index scores [T_query,64,N_candidates] +selected indices [T_query,64,K] +selected scores [T_query,64,K] +K ≤ 1024 +``` + +实际融合 kernel 不一定显式物化这些大 tensor。Profile 能确认 query/key shape 和 TopK 标量，但 TopK indices 的准确物理布局要以 kernel 接口为准。 + +相关算子： + +```text +VllmQuantLightningIndexerMetadata +VllmQuantLightningIndexer +``` + +main c4 compressor 与 Indexer compressor 的作用不同： + +- main c4 compressor 生成供 C4A 使用的 512 维长期压缩 KV。 +- Indexer compressor/Indexer 生成 128 维检索表示及 sparse indices。 +- 两者不是“先把同一个 tensor 压缩，然后原封不动交给 Indexer”的单一串行操作。 + +##### 步骤 6：构造最终 Attention 候选集合 + +对于绝对 query 位置 `p`，SWA 最近窗口为： + +```text +recent = [max(0,p-127),p] +|recent| ≤ 128 +``` + +Indexer 返回远端 TopK positions。Metadata 需要处理： + +```text +causal 越界 +padding/无效 query +TopK 与 recent window 的重复位置 +block table 到物理 cache slot 的映射 +C4 compressed positions 的有效范围 +``` + +可以把最终信息来源概念化为： + +```text +recent raw KV（SWA，最多128） +∪ relevant remote KV（Indexer TopK，最多1024） +∪ C4 compressed KV（约L/4个逻辑 summaries） +``` + +但 `SparseAttnSharedkv` 会融合读取这些来源，不一定在内存中先构造一个真正的 concat tensor。 + +##### 步骤 7：加入 learnable attention sink + +Profiler 中可看到 per-head sink 输入： + +```text +sink [1,64] +``` + +对 head `h`，softmax 概念上变为： + +```text +softmax([real_attention_logits, sink_h]) +``` + +sink 对应的 value 视为零，因此它只吸收概率质量： + +```text +O_h = Σⱼ p_j V_j + p_sink×0 +``` + +这让某个 head 在当前 compressed/selected KV 都不相关时选择“什么也不读取”。它不是 mHC Sinkhorn，也不同于 Hamming/Indexer 配置里的整数 `sink=1`。 + +##### 步骤 8：执行完整 512 维 sparse attention + +```text +Q [16,16,512] +raw shared-KV cache [10123,128,1,512] +C4 compressed state/cache 物理布局由 compressor 管理 +Indexer selected indices logical [T,64,K] +→ SparseAttnSharedkvMetadata +→ SparseAttnSharedkv +Attention output [16,16,512] +``` + +逻辑公式： + +```text +score[i,h,j] + = Q[i,h,:] · K_selected[j,0,:] / sqrt(512) + +P[i,h,:] + = softmax([score + causal/window mask, sink_h]) + +O[i,h,:] + = Σⱼ P[i,h,j] · V_selected[j,0,:] +``` + +当前 profile 中 `SparseAttnSharedkv` 共出现： + +```text +43 layers × 32 forwards = 1376 次 +``` + +说明它是 SWA、CSA、HCA 共用的 DSA 核心；具体模式由 compressed/indexer 输入是否有效决定。 + +##### 步骤 9：grouped O projection 与 mHC 写回 + +```text +Attention output [16,16,512] +→ group reshape [16,2,4096] +→ wo_a [16,2,1024] +→ flatten [16,2048] +→ local wo_b [16,4096] # TP partial +→ TP reduce-scatter [4,4096] +→ HcPost + old residual +next residual [4,4,4096] +``` + +#### 8.2.2 CSA 的计算复杂度 + +定义： + +```text +L = seq_len +W = SWA window = 128 +K = Indexer TopK = 1024 +C = c4 compression ratio = 4 +D = full attention dim = 512 +d_idx = Indexer dim = 128 +N_cand = Indexer 实际扫描的 candidate 数 +``` + +Compressor： + +```text +Prefill：O(L) +Decode ：每步 O(1) 状态更新 +``` + +昂贵的 512 维 selected attention： + +```text +Prefill：O(L×(W+K)×D) +Decode ：O((W+K)×D) +``` + +当 W 和 K 固定时，这一部分 Prefill 对 L 近似线性，Decode 单步近似定长。 + +但 Indexer 为了找到 TopK，仍需对 candidate 计算低维 score： + +```text +Prefill：O(L×N_cand×d_idx) +Decode ：O(N_cand×d_idx) +``` + +如果候选空间按 c4 压缩到约 `L/4`： + +```text +Prefill Indexer：O(L²/4×d_idx) +Decode Indexer ：O(L/4×d_idx) +``` + +如果某一实现扫描 raw positions，则应把 `L/4` 换成 `L`。仅凭 operator CSV 无法确认 kernel 内部每个阶段的精确 candidate cardinality，因此不能把整个 Indexer 宣称为严格 O(1)。 + +C4A 如果对全部 c4 summaries 做 MQA： + +```text +Prefill：O(L×(L/4)×D) +Decode ：O((L/4)×D) +``` + +实际融合实现可能对范围进一步约束，但 TopK 固定本身不能证明 C4A/Indexer 的检索成本定长。 + +因此 CSA 的准确复杂度结论是： + +- full-dim、真正读取 V 的主 sparse attention 已从 `O(L²D)` 降为约 `O(L(W+K)D)`。 +- Decode 的高成本主 attention 近似固定为最多 `128+1024` 个 positions。 +- C4 compressor 是线性/增量的。 +- Indexer 搜索和 C4A coarse attention 仍可能随 L 增长，并包含低维或压缩后的二次项。 +- 所以“CSA 主 attention 近似定长”成立，但“整个 CSA 严格定长”不成立。 + +#### 8.2.3 CSA 的数学直觉 + +CSA 把长上下文问题拆成三个分辨率： + +```text +最近历史：SWA 保留最多128个原始token，分辨率最高 +中远历史：C4A 用4:1 summary保留连续语义 +关键远端：Indexer 从长历史中挑最多1024个相关位置 +``` + +可以类比为： + +```text +SWA = 眼前的高清画面 +C4A = 每4帧做一次摘要的时间轴 +Indexer = 根据当前问题检索最相关的旧帧 +sink = 没有相关内容时允许“不读取” +``` + +数学上，它避免对所有 `(query,key)` 组合都执行昂贵的 512 维 attention 和 V 聚合，而是先用压缩表示和 128 维 Indexer 做低成本筛选，再把高精度算力集中在局部窗口和少量 TopK 上。 + +--- + +### 8.3 HCA：C128A 层级压缩注意力 + +HCA 的基本结构是： + +```text +HCA = SWA 最近原始 KV + C128A 超长程粗粒度 KV +``` + +HCA 不依赖 Lightning Indexer TopK。它直接以非常低的密度保存和读取长历史 summary。 + +#### 8.3.1 HCA 的计算流程、算子与 shape + +##### 步骤 1：Q、raw shared KV 和 SWA + +前半段与 CSA 相同： + +```text +mHC residual [4,4,4096] +→ HcPre [4,4096] +→ TP all-gather [16,4096] + +Q path: +[16,4096] → [16,1024] → [16,8192] → [16,16,512] + +shared KV path: +[16,4096] → [16,512] → [16,1,512] + +Q/K Partial RoPE +→ raw shared-KV cache [10123,128,1,512] +``` + +最近最多 128 个 raw KV 由 SWA 提供高分辨率局部信息。 + +##### 步骤 2：C128 compressor 更新层级状态 + +Profiler 直接观测： + +```text +hidden input [16,4096] +wkv / wgate [512,4096] +compressor state cache [10123,32,1024] +absolute position embedding [128,512] +norm weight [512] +compressed cos/sin [1,64] +``` + +概念 projection： + +```text +candidate = X · Wkvᵀ [16,512] +gate = X · Wgateᵀ [16,512] +``` + +每连续 128 个 token 构成一个逻辑 summary： + +```text +token 0...127 → C128_0 +token 128...255 → C128_1 +token 256...383 → C128_2 +... +``` + +逻辑 compressed position 数： + +```text +N_c128 ≈ ceil(L/128) +``` + +`APE [128,512]` 区分 chunk 内 offset 0...127。`state cache [...,32,1024]` 是融合算子的物理累积/packing 布局，不能把 32 直接解释成固定的 logical summary 数。 + +相关算子： + +```text +CompressorMetadata +Compressor +RmsNorm +InplacePartialRotaryMul # compressed position encoding +``` + +##### 步骤 3：causal 地暴露已完成 summary + +对当前 query，只能使用不包含未来 token 的 C128 state。当前尚未完成的 128-token chunk 继续更新 partial compressor state，局部细节由 SWA 覆盖。 + +例如位置 `p=200`： + +```text +C128_0 总结 token 0...127，可以作为已完成长期记忆 +当前 chunk token 128...200 尚未完成 +最近原始细节由 SWA [73...200] 提供 +``` + +当前测试 `L=16<128`： + +```text +完整 C128 group 数 = 0 +``` + +所以 profiler 能观察 compressor/state update，但不能评估多个 C128 summaries 参与 attention 时的真实性能。 + +##### 步骤 4：C128A 使用 MQA 读取压缩历史 + +逻辑输入： + +```text +Q [T_query,16,512] +C128 K/V [N_c128,1,512] +``` + +所有 16 个 local Q heads 共享同一组 C128 KV： + +```text +score128[i,h,j] + = Q[i,h,:] · K128[j,0,:] / sqrt(512) + +scores [T_query,16,N_c128] +output [T_query,16,512] +``` + +同样可以附加 per-head learnable sink： + +```text +softmax([C128 logits, sink_h]) +``` + +如果当前 query 与任何粗粒度 summary 都不相关，sink 可以吸收概率质量。 + +##### 步骤 5：融合 SWA 与 C128A 并输出 + +概念信息集合： + +```text +recent raw KV ≤ 128 positions +C128 compressed KV ≈ L/128 positions +``` + +核心算子： + +```text +SparseAttnSharedkvMetadata +SparseAttnSharedkv +``` + +其融合执行 score、causal mask、sink softmax 和 V 聚合： + +```text +Q [T_query,16,512] +→ SWA + C128A +Attention output [T_query,16,512] +``` + +HCA 没有 Lightning Indexer/TopK 阶段。 + +随后与其他 attention 类型相同： + +```text +Attention output [16,16,512] +→ grouped wo_a [16,2,1024] +→ local wo_b partial [16,4096] +→ TP reduce-scatter [4,4096] +→ HcPost +next residual [4,4,4096] +``` + +#### 8.3.2 HCA 的计算复杂度 + +定义： + +```text +L = seq_len +W = SWA window = 128 +C = C128 compression ratio = 128 +D = attention dim = 512 +N_c128 ≈ L/128 +``` + +Compressor： + +```text +Prefill：O(L) +Decode ：每步 O(1) 增量更新 +``` + +Decode 单步 Attention： + +```text +O((W + L/128)×D) +``` + +它随 L 线性增长，但斜率只有 full attention 的约 `1/128`。当前 `max_model_len=8192` 时： + +```text +最大 C128 summaries ≈ 8192/128 = 64 +最大读取范围 ≈ 128 raw + 64 compressed = 192 +``` + +最后一个 token 相比 full attention 的 8192 个 KV，attention 范围缩小约： + +```text +8192/192 ≈ 42.7 倍 +``` + +Prefill 总复杂度： + +```text +SWA ：O(L×W×D) +C128A ：O(L×(L/128)×D) +合计 ：O(L×128×D + L²/128×D) +``` + +严格来说仍有缩小 128 倍的二次项。考虑 causal 平均历史，在 `L=8192` 时，粗略 pair 数为： + +```text +SWA pairs ≈ 8192×128 = 1,048,576 +C128A pairs ≈ 8192×32 = 262,144 +``` + +在当前长度上限内，SWA 固定窗口项更大，所以整体表现会很接近以 L 为主的线性增长；但从渐近复杂度上不能把 Prefill 写成严格 O(L)。 + +HCA 的准确复杂度结论： + +- Decode：`O((128+L/128)D)`，增长非常慢且当前配置下最多读取约 192 positions。 +- Prefill：固定窗口线性项 + 缩小 128 倍的二次项。 +- 相比 full attention `O(L²D)`，计算和访存都显著降低。 + +#### 8.3.3 HCA 的数学直觉 + +HCA 使用两种时间分辨率： + +```text +最近 128 tokens：保留原始 KV，像高清短期记忆 +更早的历史 ：每 128 tokens 压成一个 summary，像低帧率长期记忆 +``` + +它假设： + +- 近处 token 的词法、语法和局部依赖需要精确表示。 +- 很远历史通常只需要保留主题、状态和宏观语义。 +- 远端若确实需要逐 token 精确检索，则交给其他 CSA layer 的 Indexer/TopK 路径补充。 + +因此交替的 CSA/HCA layers 形成互补： + +```text +CSA：4:1 中分辨率记忆 + 内容相关 TopK 精确检索 +HCA：128:1 超低成本全局摘要 +SWA：所有 layer 都保留最近窗口的高分辨率信息 +``` + +可以类比为： + +```text +SWA = 最近几秒的高清视频 +C4A = 每4帧保存一次的中分辨率记录 +C128A = 每128帧生成一张全局摘要图 +Indexer = 按当前问题从旧记录中搜索关键帧 +``` + +HCA 的价值不是精确恢复远端每个 token，而是用最多约 64 个 summary 为 8192-token 上下文提供廉价、始终可达的全局背景。 + +## 9. FFN 与 MoE：Hash Router、Learned Router、EP Dispatch 和 Expert MLP + +每层 Attention 完成后，mHC 再执行一次 pre/post： + +```text +residual [T_local,4,4096] +→ HcPre +MoE input x [T_local,4096] +→ Router + Routed/Shared Experts +MoE output y [T_local,4096] +→ HcPost +next residual [T_local,4,4096] +``` + +当前 Prefill rank 12： + +```text +T_local=4 +x [4,4096] +``` + +Decode 单请求： + +```text +T_local=1 +x [1,4096] +``` + +### 9.1 MoE 的统一数学形式 + +当前模型配置： + +```text +全局 routed experts E=256 +每 token 激活 experts K=6 +Expert Parallel EP=16 +每 EP rank local experts 256/16=16 +expert hidden size 4096 +expert intermediate size 2048 +shared expert 数 1 +``` + +无论 expert IDs 来自 Hash Router 还是 Learned Router，最终输出都可以写成： + +```text +y_t = f_shared(x_t) + + Σ(k=1...6) w_tk · f_{e_tk}(x_t) +``` + +其中： + +- `e_tk`：token `t` 选择的第 `k` 个 routed expert ID。 +- `w_tk`：对应 combine weight。 +- `f_e`：第 `e` 个 expert 的 SwiGLU MLP。 +- `f_shared`：每个 token 都执行的 shared expert，不经过 TopK 路由。 + +关键点是：虽然模型拥有 256 个 routed experts，一个 token 只执行其中 6 个，而不是执行全部 256 个。 + +### 9.2 前 3 层 Hash MoE + +Hash MoE 的核心区别是：expert IDs 由 token ID 查表得到，不需要根据 hidden 做 `4096→256` 的 router projection。 + +#### 9.2.1 Hash 路由流程和 shape + +当前 Prefill rank-local token： + +```text +token IDs [4] +hash expert table [129280,6] +``` + +查表： + +```text +expert_ids[t,:] = hash_table[token_id[t],:] +``` + +Shape： + +```text +token IDs [4] +hash table [129280,6] +→ selected expert IDs [4,6] +``` + +例如： + +```text +token_id=1234 +hash_table[1234] = [7,31,48,106,192,233] +``` + +则该 token 固定发送给这 6 个 experts。相同 token ID 在相同 Hash layer 中会得到相同 expert 集合，因此路由具有确定性。 + +Profiler 中 Hash 路由核心输入： + +```text +router affinity/workspace [4,256] +token IDs [4] +hash table [129280,6] +``` + +对应算子： + +```text +_C_ascend::moe_gating_top_k_hash +aclnnMoeGatingTopKHash +``` + +这里的 `[4,256]` 是融合 selector 使用的 affinity/workspace 形状；真正决定 Hash expert IDs 的直接证据是 `[4]` token IDs 和 `[129280,6]` table。Operator CSV 不记录算子输出值，因此 Hash combine weight 的精确归一化规则不能只凭 shape 推断；可以确定的是 expert IDs 不来自 hidden 的 learned top-k logits。 + +Decode 对应： + +```text +token IDs [1] +selected expert IDs [1,6] +``` + +#### 9.2.2 Hash MoE 的数学直觉 + +Hash Router 相当于一个预先学习/构造好的词表到专家映射： + +```text +vocabulary token → 6 个固定专家 +``` + +优点： + +- 不需要每 token 执行大 router matmul。 +- 相同词天然落到相同专家，行为稳定。 +- 路由延迟低，容易形成稳定的 expert specialization。 + +代价： + +- 路由只直接看到 token ID，看不到当前上下文 hidden。 +- 同一个词在不同语境中仍先落到同一组 experts。 +- 高频 token 可能天然造成热点，需要训练策略或运行时负载均衡处理。 + +因此前 3 层使用 Hash MoE，可以把较浅层、偏词法的特征以很低的 router 成本分发给专家；深层再改用上下文相关的 Learned Router。 + +### 9.3 后续 Learned MoE + +Learned MoE 根据当前 token hidden 动态计算 256 个 expert affinity，因此同一个 token ID 在不同上下文、不同层可以选择不同 experts。 + +#### 9.3.1 Router projection + +Profiler 直接观测到： + +```text +MoE hidden [4,4096] +router weight [256,4096] +``` + +线性 projection： + +```text +router_logits = x · W_routerᵀ + +[4,4096] @ [4096,256] +→ [4,256] +``` + +profile 对应调用： + +```text +aten::linear +aten::matmul +aclnnMatmul +``` + +融合 selector 的 learned-routing 输入： + +```text +router affinity/logits [4,256] +expert correction bias [256] [实测出现] +token ID/hash table empty +``` + +而 Hash 路径是： + +```text +router affinity/workspace [4,256] +expert correction bias empty +token IDs [4] +hash table [129280,6] +``` + +两条路径最终都进入当前 Ascend 融合 selector，因此 profiler 中都可能显示 `moe_gating_top_k_hash` 这个通用算子名；必须看非空输入才能区分 Hash 与 Learned 路由。 + +#### 9.3.2 Top6 选择和权重 + +概念流程： + +```text +router logits [4,256] +→ scoring activation / correction bias +expert scores [4,256] +→ TopK(K=6) +expert IDs [4,6] +selected weights [4,6] +``` + +DeepSeek 风格 learned routing 中，correction bias 可以参与“选择哪些 experts”，而最终 combine weight 通常由未加选择偏置的原始 affinity 归一化得到；具体 activation、group-limited TopK 和 scaling 顺序应以对应版本 selector 参数为准。当前 profile 能确认 `[256]` correction bias 和 Top6 路径，但不能从 CSV 恢复具体数值。 + +数学上： + +```text +s_t = Router(x_t) ∈ R^256 +E_t = TopK(s_t,6) +w_t = Normalize(s_t[E_t]) ∈ R^6 +``` + +Learned Router 的直觉是： + +```text +当前 hidden 表示“这个 token 在此刻需要什么能力” +Router 根据语义把它送到最适合的 6 个专家 +``` + +### 9.4 EP16 Dispatch：为什么 4 个 token 会变成 24，再变成 44 + +路由完成后，rank 12 最初有： + +```text +local tokens 4 +experts per token 6 +expert assignments 4×6=24 +``` + +每个 token 会复制/引用 6 次，并附带： + +```text +hidden [4096] +expert ID +combine weight +original token index +``` + +逻辑 dispatch buffer： + +```text +[4,6,4096] +→ flatten assignments +[24,4096] +``` + +expert owner 由 EP16 决定。全局 256 experts 平均分布： + +```text +rank 0 owns experts 0...15 +rank 1 owns experts 16...31 +... +rank 15 owns experts 240...255 +``` + +每个 rank 按 destination EP rank 排序 token assignments，然后执行 All-to-All： + +```text +local assignment buffer [24,4096] +→ token permute / dispatch +→ HCCL All-to-All-V +→ received expert tokens [N_recv,4096] +``` + +`N_recv` 由所有 EP ranks 路由到当前 rank 的 token 数决定，不等于本地原始 token 数。rank 12 profile 中出现： + +```text +[44,4096] +[8,4096] +[10,4096] +[15,4096] +``` + +所以 `[44,4096]` 的含义是：在那个 layer/step，当前 rank 的 16 个 local experts 合计收到 44 个 expert-token assignments。它不是 batch size，也不是模型固定 shape。 + +相关 dispatch 算子/通信： + +```text +MoeDistributeDispatchV2 +npu_moe_token_permute +HcclAlltoAllV / all_to_all_single +``` + +### 9.5 Local Expert MLP：Grouped MatMul + SwiGLU + +EP dispatch 后，当前 rank 只执行自己持有的 16 个 experts。 + +输入按 local expert 排序： + +```text +expert 0 tokens +expert 1 tokens +... +expert 15 tokens +``` + +用 group list 描述每个 expert 对应的 token 区间，然后一次 grouped kernel 执行 16 个不同 MLP。 + +#### 9.5.1 Gate/Up projection + +Profiler 示例： + +```text +routed hidden [44,4096] +local expert gate/up weights [16,4096,4096] +local expert scales [16,4096] +group information dynamic，覆盖44个assignments +local expert count 16 +``` + +权重最后的 4096 实际包含： + +```text +gate dim 2048 + up dim 2048 = 4096 +``` + +对属于 expert `e` 的 token： + +```text +z_e = x_e · W_gate_up[e]ᵀ + +[N_e,4096] @ [4096,4096] +→ [N_e,4096] +→ split + gate [N_e,2048] + up [N_e,2048] +``` + +#### 9.5.2 SwiGLU 与量化 + +```text +activated = SiLU(gate) ⊙ up + +[N_e,2048] ⊙ [N_e,2048] +→ [N_e,2048] +``` + +当前融合算子将 grouped gate/up matmul、SwiGLU 和后续量化合并： + +```text +_C_ascend::grouped_matmul_swiglu_quant_weight_nz +aclnnGroupedMatmulSwigluQuantWeightNZ +``` + +因此 profile 中不会分别看到每个 expert 的 16 次独立 matmul 和 SiLU。 + +#### 9.5.3 Down projection + +每个 expert 再把 intermediate 2048 投回 hidden 4096： + +```text +activated [N_e,2048] +W_down[e] [2048,4096] # profiler物理权重布局 +→ output [N_e,4096] +``` + +rank 12 profile 直接看到 local down weight 物理 shape： + +```text +[16,2048,4096] +``` + +对应算子： + +```text +npu::npu_grouped_matmul +aclnnGroupedMatmulWeightNz +``` + +合并 16 个 local experts 后，当前 rank 仍得到： + +```text +local expert outputs [N_recv,4096] +例如 [44,4096] +``` + +### 9.6 Reverse All-to-All 与 Top6 Combine + +Expert MLP 完成后，需要将每份 expert output 送回原 token 所在 rank： + +```text +local expert outputs [N_recv,4096] +→ expert-side unpermute +→ reverse HCCL All-to-All-V +→ original assignments [T_local×6,4096] +``` + +当前 Prefill 示例恢复为： + +```text +[24,4096] +→ reshape/group by original token +[4,6,4096] +``` + +然后按 router weight 合并： + +```text +y_routed[t,:] + = Σ(k=1...6) w[t,k] · y_expert[t,k,:] +``` + +Shape： + +```text +expert outputs [4,6,4096] +combine weights [4,6] +→ weighted reduce on K dimension +routed result [4,4096] +``` + +相关算子： + +```text +npu_moe_token_unpermute +MoeDistributeCombineV2 +HcclAlltoAllV +``` + +最后加入 shared expert： + +```text +shared result [4,4096] +routed result [4,4096] +→ add +MoE output [4,4096] +→ HcPost +next residual [4,4,4096] +``` + +### 9.7 Decode 的 shape + +单请求 Decode rank-local token 数为 1： + +```text +MoE input [1,4096] +Hash IDs 或 learned logits [1,6] / [1,256] +Top6 assignments [1,6] +flatten dispatch [6,4096] +→ EP All-to-All +received tokens [N_recv_decode,4096] # 动态 +→ local grouped experts +→ reverse All-to-All +returned assignments [6,4096] +→ weighted Top6 combine +MoE output [1,4096] +→ HcPost +residual [1,4,4096] +``` + +Decode 的 token 数很少，expert matmul 本身可能很小；All-to-All latency、同步和负载不均更容易成为 TPOT 瓶颈。 + +### 9.8 计算复杂度与性能直觉 + +#### Hash Router 与 Learned Router + +Hash Router： + +```text +查表成本约 O(T×K) +不需要 O(T×H×E) router matmul +``` + +Learned Router： + +```text +router projection O(T×H×E) += O(T×4096×256) +``` + +但相对 Expert MLP 和 EP 通信，router matmul通常不是最大成本。 + +#### Expert 计算量 + +每 token 只计算 K=6 个 routed experts： + +```text +O(T×K×H×I) +``` + +其中： + +```text +H=4096 +I=2048 +K=6 +``` + +它与总 expert 数 E=256 不呈线性关系；增加总 experts 主要增加参数容量和路由选择空间，不会让每 token 执行全部 experts。 + +#### EP 通信量 + +Dispatch/Combine 数据量近似： + +```text +O(T×K×H) +``` + +当前 Prefill 每 rank： + +```text +4 tokens × 6 copies × 4096 elements +``` + +Decode 每 rank： + +```text +1 token × 6 copies × 4096 elements +``` + +小 token 数下带宽未必饱和，通信启动、同步和跨 rank 尾部延迟更关键。 + +#### 负载不均 + +理想情况下，各 rank 接收接近相同数量的 expert assignments；实际由路由分布决定： + +```text +step latency ≈ 最慢/最满 expert rank 的完成时间 +``` + +所以 `N_recv=44/8/10/15` 的动态变化会直接影响 grouped matmul shape、通信量和尾部延迟。 + +当前 op statistic 中： + +```text +MoeDistributeDispatchV2 约占 82.1% device op time +MoeDistributeCombineV2 约占 9.2% +``` + +这说明本次 profile 的主要设备时间集中在 MoE dispatch/combine，而不是 expert matmul本身。不过该 profile 带 stack/memory instrumentation，绝对比例不能直接当作无侵入性能结果，应该用 ops-only profile 复核。 + +### 9.9 当前 profile 的 AFD force-load-balance 影响 + +路由与 expert kernel 调用栈直接包含： + +```text +afd_plugin/compat/patches/npu/force_load_balance.py +``` + +涉及位置包括： + +```text +expert selector apply +expert MLP / communication apply +``` + +因此要区分两类事实： + +- `[4,256]`、`[4]`、`[129280,6]`、`[16,4096,4096]`、`[16,2048,4096]` 等模型/算子 shape 是有效的。 +- `[44,4096]`、`[8,4096]` 等具体 `N_recv` 是该次实际执行 shape，但可能已经受到 AFD force-load-balance patch 影响，不能直接代表原生 DSV4 learned/hash routing 的自然负载分布。 + +要研究原生路由质量、expert 热点和 load balance，需要单独做两组对照： + +```text +A：关闭 force-load-balance，记录自然 expert histogram +B：开启 force-load-balance，记录修改后的 histogram +``` + +并同时比较： + +```text +每层/每rank expert-token counts +max/mean load ratio +Dispatch/Combine 时间 +Grouped expert kernel 时间 +TTFT/TPOT +``` + +## 10. Final hidden、LM Head 与 MTP buffer + +43 层后： + +```text +Prefill residual [4,4,4096] +Decode residual [1,4,4096] +``` + +最终 mHC head： + +```text +Prefill [4,4,4096] → [4,4096] +Decode [1,4,4096] → [1,4096] +``` + +Prefill 对每个请求选择最后一个有效 token hidden，LM head 输入为 `[1,4096]`。TP vocab logits 每 rank 是 `[1,32320]`，逻辑全局 logits 是 `[1,129280]`，采样得到 `next_token_id [1]`。 + +模型还可将合并前的 mHC hidden 写入 MTP buffer： + +```text +[T,4,4096] → flatten [T,16384] +预分配 buffer [4096,16384] +``` + +buffer 第一维 4096 对应 `max_num_batched_tokens`，不是当前请求实际 token 数。 + +## 11. Prefill 端到端 shape 表 + +| 顺序 | 阶段 | 输入 | 输出 | 说明 | +|---:|---|---|---|---| +| 1 | Tokenizer/template | 文本 / `[12]` | `[1,16]` | 实际输入 16 tokens | +| 2 | TP Embedding | `[16]` | logical `[16,4096]` | rank vocab `[32320,4096]` | +| 3 | SP scatter | `[16,4096]` | `[4,4096]` | TP=4 | +| 4 | mHC 初始化 | `[4,4096]` | `[4,4,4096]` | 4 streams | +| 5 | Attention hc_pre | `[4,4,4096]` | `[4,4096]` | 输入实测 | +| 6 | TP all-gather | `[4,4096]` | `[16,4096]` | 完整 token 维 | +| 7 | Q projection | `[16,4096]` | `[16,16,512]` | local 16 heads | +| 8 | SWA/CSA/HCA | Q + caches | `[16,16,512]` | layer 类型不同 | +| 9 | O projection/RS | `[16,16,512]` | `[4,4096]` | 回到 local tokens | +| 10 | Attention hc_post | residual + output | `[4,4,4096]` | 写回 streams | +| 11 | FFN hc_pre | `[4,4,4096]` | `[4,4096]` | 每层第二次 hc_pre | +| 12 | Router | `[4,4096]` | `[4,256]`, top6 | Hash/learned | +| 13 | EP dispatch/experts | local 4×top6 | dynamic | routed tokens 动态 | +| 14 | EP combine | routed outputs | `[4,4096]` | 加权聚合 | +| 15 | FFN hc_post | residual + output | `[4,4,4096]` | 43 层重复 | +| 16 | Final mHC head | `[4,4,4096]` | `[4,4096]` | 合并 streams | +| 17 | Select last token | sequence hidden | `[1,4096]` | 每请求最后 token | +| 18 | LM head | `[1,4096]` | local `[1,32320]` | global `[1,129280]` | +| 19 | Sampler | logits | `[1]` | 产生 y1 | + +## 12. Decode 单步 shape 表 + +| 顺序 | 阶段 | 输入 | 输出 | 证据 | +|---:|---|---|---|---| +| 1 | 输入 token | `[1]` | `[1]` | 逻辑 | +| 2 | TP Embedding | `[1]` | `[1,4096]` | 代码 | +| 3 | mHC residual | `[1,4096]` | `[1,4,4096]` | residual 实测 | +| 4 | Attention hc_pre | `[1,4,4096]` | `[1,4096]` | 输入实测 | +| 5 | 图槽位/TP 对齐 | 1 valid token | physical 4 slots | 推导 | +| 6 | Q projection | physical hidden | `[4,16,512]` | 推导 | +| 7 | Attention/KV append | Q + cache | `[4,16,512]` | 每步新增 1 个有效 KV | +| 8 | O projection/RS | `[4,16,512]` | `[1,4096]` | 推导 | +| 9 | Attention hc_post | residual + output | `[1,4,4096]` | 代码 | +| 10 | FFN hc_pre | `[1,4,4096]` | `[1,4096]` | 代码 | +| 11 | Router/experts | `[1,4096]` | `[1,4096]` | affinity `[1,256]` 推导 | +| 12 | FFN hc_post | residual + output | `[1,4,4096]` | 代码 | +| 13 | Final mHC head | `[1,4,4096]` | `[1,4096]` | 代码 | +| 14 | LM head | `[1,4096]` | local `[1,32320]` | global `[1,129280]` | +| 15 | Sampler | logits | `[1]` | 下一个 token | + +上下文增长： + +| Forward | 本次输入 | Forward 前 KV | Forward 后 KV | 输出 | +|---|---|---:|---:|---| +| Prefill | prompt 16 | 0 | 16 | `y1` | +| Decode 1 | `y1` | 16 | 17 | `y2` | +| Decode 2 | `y2` | 17 | 18 | `y3` | +| ... | ... | ... | ... | ... | +| Decode 31 | `y31` | 46 | 47 | `y32` | + +## 13. P/D 分离语义 + +当前 profile 中 P 和 D 是同一服务实例内先后执行的 scheduler 阶段： + +```text +P step：输入 16 tokens + → 建立初始 SWA KV、CSA/HCA compressed state 和 metadata + → 产生 y1 + +D steps ×31：每步输入 1 个有效 token + → 更新 KV/compressed state + → 产生下一个 token +``` + +如果部署为独立 P 实例和 D 实例，模型内部计算 shape 基本不变，变化的是状态所有权和传输： + +```text +P 实例处理 [16] prompt +→ 将请求 cache/state/metadata 传给 D 实例 +→ D 实例每步处理 [1] 有效 token +``` + +DSV4 的 P/D 传输不能只考虑普通 KV；还必须保证 compressed cache、Indexer 状态、block table、position/context metadata 等一致。线上具体 wire format 要以 AFD-plugin 当前 KV connector 代码和传输 trace 为准，不能仅从本次单实例 profiler 推断。 + +## 14. Compile / ACL Graph 对 shape 的影响 + +开启 `torch.compile + ACL Graph` 后逻辑 shape 不变，但执行图使用固定或分桶后的物理 shape： + +```text +逻辑 Decode：1 valid token +物理图输入 ：graph size / padding / TP 对齐后的 token slots +``` + +当前 `max_num_seqs=4`，结合 FlashComm/TP 对齐，Decode 可使用 physical token size 4 的图。分析时必须区分： + +- 真正的 `num_scheduled_tokens`。 +- 捕获图的 batch/token size。 +- 为静态图或通信对齐加入的 padding token。 + +这解释了为什么 Decode residual 直接观测为 `[1,4,4096]`，而 attention Q 推导为 `[4,16,512]`。 + +## 15. MTP 当前状态与启用后的 shape + +模型目录包含 MTP 权重，但当前启动没有 speculative 配置，因此 MTP draft forward **未参与本次性能测试**；当前主模型只维护 `_mtp_hidden_buffer` 需要的数据。 + +如果以后启用 MTP，典型数据流为： + +```text +draft token ids [N_d] +token embedding [N_d,4096] +previous mHC hidden [N_d,16384] +reshape previous [N_d,4,4096] +e_proj branch [N_d,1,4096] +h_proj branch [N_d,4,4096] +融合后 [N_d,4,4096] +MTP block [N_d,4,4096] +MTP mHC head [N_d,4096] +MTP logits [N_d,129280] +``` + +`N_d` 是本次 speculative draft token 数，不固定为 1，也不等于 Prefill 的 16。 + +## 16. 本例的能力边界 + +本例可以确认 mHC residual、TP/SP、三种 attention、Hash/Learned MoE、Prefill/Decode 上下文增长和 graph padding 的主数据流，但不足以评价： + +- 超过 128 tokens 后 SWA 窗口截断。 +- 形成多个 c128 group 后的 HCA。 +- 历史超过 TopK=1024 后 CSA 的稀疏收益。 +- 多请求 continuous batching 的 graph/padding shape。 +- 真正启用 MTP 后的额外成本和接受率。 +- 独立 P/D 节点的实际传输 tensor 列表。 + +建议后续补测： + +```text +Prefill length: 16, 128, 512, 1024, 4096 +Decode context: 128, 512, 1024, 4096, 8192 +Active seqs : 1, 2, 4 +Mode : eager / compile+ACL Graph +``` + +## 17. 实测证据位置与一页总览 + +本地 profile： + +```text +/mnt/d/cyj/afd/dsv4_rank12_profiles/ + A_eager_stack_memory_shapes/ + ASCEND_PROFILER_OUTPUT/operator_details.csv +``` + +关键实测摘要： + +```text +npu_hc_pre_v2 Prefill input [4,4,4096] +npu_hc_pre_v2 Decode input [1,4,4096] +sparse attention Prefill Q [16,16,512] +Lightning Indexer query [16,64,128] +Hash router affinity [4,256] +Hash token ids [4] +Hash table [129280,6] +``` + +```text +Text → tokens [1,16] + → TP embedding/SP [4,4096] + → mHC [4,4,4096] + → 43×{hc_pre → TP gather → Q [16,16,512] + → SWA/CSA-c4/HCA-c128 + → O projection/RS → MoE → hc_post} + → final hidden [1,4096] + → local logits [1,32320] / global [1,129280] + → sample y1 + → 31×Decode：residual [1,4,4096] + attention physical Q [4,16,512] [推导] + → next token + → 32 output tokens，最终已处理 KV 长度 47 +``` +## 18. mHC 机制详解：初始化、Projection、Sinkhorn 与 hc_post + +这一节专门解释 hidden 如何扩展为 4 条 residual stream，以及 Attention/FFN 子层中的 `hc_pre` 和 `hc_post` 如何工作。 + +### 18.1 初始化：`unsqueeze` 只增加 residual-stream 轴 + +Embedding 或前置 hidden 的 rank-local shape 是： + +```text +x: [T_local,H]，其中 H=4096 +``` + +mHC 初始化在 hidden 维前插入一条 stream 轴，再扩展成 4 条： + +```python +X = x.unsqueeze(1) # [T_local,1,4096] +X = X.expand(-1,4,-1) # [T_local,4,4096] +``` + +概念上等价于： + +```text +X[t] = [x[t], x[t], x[t], x[t]] +``` + +当前测试： + +```text +Prefill: [4,4096] → [4,1,4096] → [4,4,4096] +Decode : [1,4096] → [1,1,4096] → [1,4,4096] +``` + +这里不是 `Linear(4096,16384)`。初始化时四条 stream 内容相同，没有凭空产生新信息；它们在后续每层不同的 residual mixing 和子层输出注入中逐渐分化。`expand` 可以只是零拷贝 view，真正进入要求连续内存的融合算子时才可能物化。 + +### 18.2 一次 Projection 同时生成三组 mapping + +设进入一个 Attention 或 FFN 子层前： + +```text +X: [T,C,H]，C=4，H=4096 +``` + +对每个 token 展平 stream 和 hidden： + +```text +X_flat = reshape(X) +[T,4,4096] → [T,16384] +``` + +当前 rank 12 profiler 对 `npu_hc_pre_v2` 的直接观测是： + +```text +X [4,4,4096] +projection weight [24,16384] +附加参数 [3] +附加参数 [24] +``` + +因此 projection 主体可以写为： + +```text +g = X_flat · Wᵀ + b + +X_flat [T,16384] +W [24,16384] +b [24] +g [T,24] +``` + +`24` 的来源不是经验值，而是三组 mapping 的元素数之和： + +```text +H_pre : C = 4 +H_post: C = 4 +H_res : C² = 16 +---------------- +总计 24 +``` + +投影结果拆分为： + +```python +g_pre, g_post, g_res = split(g, [4,4,16], dim=-1) + +G_pre : [T,4] +G_post: [T,4] +G_res : [T,16] → [T,4,4] +``` + +融合算子还接收 `[3]` 和 `[24]` 两组参数。Operator CSV 可以确认其 shape，并能判断它们分别作用于三组 mapping 和 24 个输出分量；但仅凭 profile 不能严谨确定其源码变量名以及 bias、scale、激活的精确先后顺序。因此本文将确定的数据流写成： + +```text +X_flat +→ learned projection +→ static/bias 与分组 scale +→ pre/post 激活或归一化 +→ residual logits 的 Sinkhorn +→ H_pre、H_post、H_res +``` + +训练学习的是 `W`、静态项和缩放项；推理时每个 token 根据自己的当前 `X` 动态得到实际 mapping。因此不同 token、不同层、Attention 与 FFN 子层都可能得到不同系数。 + +### 18.3 `hc_pre`：4 条 stream 动态合成一条 + +激活/归一化后的： + +```text +H_pre: [T,4] +X : [T,4,4096] +``` + +逐 token 计算： + +```text +z[t,h] = Σᵢ H_pre[t,i] · X[t,i,h] +``` + +Shape： + +```text +[T,4] × [T,4,4096] +→ 按 stream 维加权求和 +→ z [T,4096] +``` + +例如某个 token 的： + +```text +H_pre[t] = [0.1,0.2,0.3,0.4] +``` + +则： + +```text +z[t] = 0.1X₁[t] + 0.2X₂[t] + 0.3X₃[t] + 0.4X₄[t] +``` + +Attention 和 FFN 始终处理 H=4096 的 `z`，并不直接处理一个 16384 维 Transformer hidden。 + +### 18.4 Sinkhorn：把 residual logits 投影到近似双随机矩阵 + +Projection 产生的 residual logits： + +```text +G_res: [T,4,4] +``` + +原始 logits 可以为负数，行列和也没有约束。Sinkhorn 先将其转成正数，再交替进行行、列归一化： + +```python +M = exp(G_res / temperature) + +for _ in range(num_iterations): + M = M / M.sum(dim=-1, keepdim=True) # 每行和归一到 1 + M = M / M.sum(dim=-2, keepdim=True) # 每列和归一到 1 + +H_res = M +``` + +数值稳定实现通常在 log space 中完成： + +```python +Z = G_res / temperature + +for _ in range(num_iterations): + Z = Z - logsumexp(Z, dim=-1, keepdim=True) + Z = Z - logsumexp(Z, dim=-2, keepdim=True) + +H_res = exp(Z) +``` + +最终： + +```text +H_res[t,i,j] ≥ 0 +Σⱼ H_res[t,i,j] ≈ 1 # 每行 +Σᵢ H_res[t,i,j] ≈ 1 # 每列 +``` + +例如： + +```text +H_res = +[[0.7,0.1,0.1,0.1], + [0.1,0.7,0.1,0.1], + [0.1,0.1,0.7,0.1], + [0.1,0.1,0.1,0.7]] +``` + +它保留每条 stream 的主要信息，同时允许信息在 stream 之间流动。双随机约束控制每一层 residual mixing 的总量，避免 43 层传播中持续放大、衰减或所有信息塌缩到单条 stream。 + +### 18.5 Attention 中完整的 Prefill shape + +当前 Prefill rank-local 的 mHC/Attention 流程： + +```text +输入 residual X [4,4,4096] +reshape [4,16384] +projection [4,24] +├─ H_pre [4,4] +├─ H_post [4,4] +└─ H_res [4,4,4] + +hc_pre: H_pre 与 X 加权求和 + [4,4096] +TP all-gather [16,4096] +Q projection [16,16,512] +SWA / CSA / HCA [16,16,512] +O projection + TP reduce-scatter + [4,4096] + +hc_post: +branch output y [4,4096] +old residual X [4,4,4096] +H_post [4,4] +H_res [4,4,4] +→ next residual [4,4,4096] +``` + +Profiler 的 `npu_hc_post` 保留了 batch 维，因此直接显示为： + +```text +y [1,4,4096] +X [1,4,4,4096] +H_post [1,4,4] +H_res [1,4,4,4] +``` + +去掉 `B=1` 后，正好对应上面的 `[T,...]` 表达。 + +### 18.6 Attention 中完整的 Decode shape + +单请求 Decode 的有效 token 数为 1： + +```text +输入 residual X [1,4,4096] [实测] +reshape [1,16384] +projection [1,24] +├─ H_pre [1,4] +├─ H_post [1,4] +└─ H_res [1,4,4] + +hc_pre [1,4096] +graph/TP 对齐后物理 Q [4,16,512] [推导] +Attention + O projection/RS [1,4096] +hc_post [1,4,4096] +``` + +这里 mHC mapping 按 1 个有效 token 生成；Attention 侧的物理 token 维度 4 来自 ACL Graph/TP 对齐，不能误解成 `H_pre/H_post` 为 4 个请求生成。 + +### 18.7 `hc_post` 的计算公式 + +令子层输出为： + +```text +y: [T,H] +``` + +则每个 token、每条输出 stream、每个 hidden 元素的更新为： + +```text +X_next[t,j,h] + = Σᵢ H_res[t,j,i] · X[t,i,h] + + H_post[t,j] · y[t,h] +``` + +矩阵写法： + +```text +X_mixed = H_res @ X +X_next = X_mixed + H_post.unsqueeze(-1) * y.unsqueeze(1) + +H_res [T,4,4] +X [T,4,4096] +X_mixed [T,4,4096] +H_post.unsqueeze(-1) [T,4,1] +y.unsqueeze(1) [T,1,4096] +broadcast product [T,4,4096] +X_next [T,4,4096] +``` + +### 18.8 一个可手算的 `hc_post` 例子 + +为便于手算，将 `H=4096` 暂时缩成 `H=2`，令单个 token 的旧 residual 为： + +```text +X₁=[1,0] +X₂=[0,1] +X₃=[1,1] +X₄=[2,0] + +X shape: [4,2] +``` + +Attention 输出： + +```text +y=[10,20]，shape [2] +``` + +Sinkhorn 后的： + +```text +H_res= +[[0.7,0.1,0.1,0.1], + [0.1,0.7,0.1,0.1], + [0.1,0.1,0.7,0.1], + [0.1,0.1,0.1,0.7]] +``` + +第一步，混合旧 residual： + +```text +X₁_mixed = 0.7X₁+0.1X₂+0.1X₃+0.1X₄ = [1.0,0.2] +X₂_mixed = 0.1X₁+0.7X₂+0.1X₃+0.1X₄ = [0.4,0.8] +X₃_mixed = 0.1X₁+0.1X₂+0.7X₃+0.1X₄ = [1.0,0.8] +X₄_mixed = 0.1X₁+0.1X₂+0.1X₃+0.7X₄ = [1.6,0.2] +``` + +假设动态生成： + +```text +H_post=[0.2,0.4,0.1,0.3] +``` + +第二步，将同一个 Attention 输出以不同权重注入每条 stream： + +```text +X₁_next = [1.0,0.2] + 0.2[10,20] = [3.0,4.2] +X₂_next = [0.4,0.8] + 0.4[10,20] = [4.4,8.8] +X₃_next = [1.0,0.8] + 0.1[10,20] = [2.0,2.8] +X₄_next = [1.6,0.2] + 0.3[10,20] = [4.6,6.2] +``` + +Shape 全程保持： + +```text +旧 residual [4,2] +→ H_res mixing [4,2] ++ H_post×y [4,2] +→ 新 residual [4,2] +``` + +真实模型只需把简化的 `H=2` 换回 `H=4096`： + +```text +[T,4,4096] → hc_post → [T,4,4096] +``` + +因此 mHC 的核心不是把 Attention 输出变成四份完全独立的新特征，而是：旧 residual 先通过受约束的 `H_res` 重新路由，同一个子层输出再通过 token-dependent 的 `H_post` 以不同强度写入四条 stream。 + +### 18.9 最终合并与 MTP 展平 + +经过 43 层后，最终 mHC head 将 4 streams 合成为普通 hidden： + +```text +Prefill [4,4,4096] → [4,4096] +Decode [1,4,4096] → [1,4096] +``` + +只有为了保存 MTP 所需的完整 mHC 状态时，才会展平为： + +```text +[T,4,4096] → [T,16384] +``` + +这个 `[T,16384]` 是 4 条 residual stream 的存储形式，不表示 Attention/FFN 的 hidden size 已从 4096 改成 16384。 + +## 19. 调度长度参数补充 + +- `max_model_len=8192`：单个请求允许模型处理的最大序列长度，近似满足 `prompt tokens + 已生成并送回模型的 tokens ≤ max_model_len`。 +- `max_num_batched_tokens=4096`：一次 scheduler step 中所有请求合计最多调度的 token 数，不是单请求总长度。 +- `max_num_seqs=4`：一次调度中最多容纳的 active sequence 数。 + +例如当前 16-token prompt 在 `max_model_len=8192` 下，理论上还能继续处理约 8176 个生成 token；但实际还会受到请求 `max_tokens`、EOS、显存/缓存和实现限制。 + + +--- + +后续若拿到 Decode operator CSV 中精确的 Q/Indexer shape，以及独立 P/D connector 的实际传输 trace，应把相应 **[推导]** 项升级为 **[实测]**。 diff --git a/docs/experiments/cli.md b/docs/experiments/cli.md new file mode 100644 index 00000000..91c1cb2d --- /dev/null +++ b/docs/experiments/cli.md @@ -0,0 +1,182 @@ +# 安装 +## 安装 modctl +### 蚂蚁内部（蚂蚁物理机环境/使用蚂蚁内部镜像的容器） +建议通过蚂蚁内部源安装。首次安装可能需要添加蚂蚁内部源，方法如下（更详细的信息可以参考[蚂蚁制品库使用说明](https://yuque.antfin.com/antbuild/bpxc9y/aww0ho#ysZtk)）： + +```shell +# 对于 7u 的操作系统： +$ sudo yum install -y http://artifacts.antgroup-inc.cn/artifact/repositories/ant_7_noarch_current/ant-repo-utils/ant-repo-utils-0.0.1-43313729.noarch.rpm + +# 对于 8u 的操作系统： +$ sudo yum install -y https://artifacts.antgroup-inc.cn/artifact/repositories/ant_8_noarch_current/ant-repo-utils/ant-repo-utils-0.0.1-202188063.noarch.rpm + +``` + +安装最新版 modctl ： + +```shell +$ sudo yum install -y modctl -b current +``` + +### 蚂蚁内部使用 Debian/Ubuntu 等外部镜像的容器 +对于Debian/Ubuntu系统，我们也上传了deb包内部蚂蚁制品库的，方便用户下载安装。因为没有蚂蚁内部 apt源，所以用户需wget下载deb包后手动安装：可访问[https://artifacts-web.antgroup-inc.cn/common/versions?name=modctl-0&t=MAIN_SITE](https://artifacts-web.antgroup-inc.cn/common/versions?name=modctl-0&t=MAIN_SITE) 找到满足自己系统架构的版本下载后通过`dpkg -i xxx.deb`来进行安装。建议尽可能选择最新版本。 + +### 蚂蚁外部，或者 MacOS 等办公网环境 +可以选择从 github 下载安装：[modctl releases](https://github.com/modelpack/modctl/tags) 。（点击页面最新的 tag，然后下载不同平台的安装包，建议尽可能选择最新版本） + +### 暂不支持 windows +# 使用 +## 账号 +请通过 [https://hcs.alipay.com/hmr/credentials](https://hcs.alipay.com/hmr/credentials) 设置登录密码，使用 **域账号名 + HMR 密码（非域账号密码）**登录（参考文档 [从命令行拉取和推送镜像](https://yuque.antfin.com/tuna/ifau4b/ovkir7124ks7xmax)）。其他问题请咨询 [@楚贤](https://yuque.antfin.com/chuxian.mjj)[@康德](https://yuque.antfin.com/lb203159)。 + + + + + +## 登录 +现有使用环境下，网络环境纷繁复杂，涉及主站集群、AIDC 集群、公有云网络、办公网环境、容器环境、宿主机环境等等，同时又要保证高性能，因此以下域名选择和命令选项非常关键，需要精心选择。 + +### 域名选择 +办公网域名(网速不可控): `hmr.antgroup-inc.cn` + +线上域名: `hmr.sa128.alipay.com`，我们建议在线上环境尽可能的使用这个域名，因为前者走 spanner 转发流量受网络集群限制，本身速率不稳定。 + +### 命令选项 +#### 关于--plain-http +如果是`hmr.antgroup-inc.cn`域名，必须要用 https，所以不加`--plain-http`。 + +如果是`hmr.sa128.alipay.com`域名，只支持 http，所以必须加`--plain-http`。 + +#### 关于--insecure +一般情况下不用考虑这个选项。只有在使用`hmr.antgroup-inc.cn`域名，并且使用`--proxy `走 dragonfly 加速时，需要加上`--insecure` 选项。 + +### 登录示例 +例如在主站生产环境，为了尽可能高的性能所以使用域名 `hmr.sa128.alipay.com`，而 `hmr.sa128.alipay.com`目前支持 http 协议，所以需要加上`--plain-http`，命令如下： + +```shell +$ modctl login hmr.sa128.alipay.com --plain-http -u ${username} -p ${password} +Logging In... +Login Succeeded. +``` + +## 上传 +### 构建 & 上传模型 +`/path/to/model`即为本地需要上传的模型目录, `${namespace}`为申请的 Namespace。 + +```shell +$ modctl modelfile generate /path/to/model --output /path/to/modelfile +$ modctl build -f /path/to/modelfile/Modelfile /path/to/model \ +-t hmr.antgroup-inc.cn/${namespace}/deepseek-v3:v1.0.2 --plain-http --output-remote --concurrency 16 +``` + +#### 指定 source 信息 +若在构建时，需要将模型源信息自定义添加进模型镜像中，可通过以下命令指定，默认行为是当检测到模型当前目录下为 git 仓库或 zeta 仓库时会自动解析 source 相关信息，无需用户指定，其他情况可按需指定。 + +```yaml +$ modctl build -f /path/to/modelfile/Modelfile /path/to/model \ +-t hmr.antgroup-inc.cn/${namespace}/deepseek-v3:v1.0.2 --plain-http --output-remote --concurrency 16 --source-url https://huggingface.co/deepseek-ai/DeepSeek-R1 --source-revision 44effdfa8e727bc64ee7f +``` + +#### 追加/覆盖 +在某些场景下，可能需要追加/修改已经构建并上传的模型中的某个文件，如 config.json, 如果再重新完整构建上传一次，成本会比较高，所以可以通过以下命令来追加/修改已构建的模型镜像。 + +```yaml +# 追加一个之前不存在的文件 +$ modctl attach foo.txt -s registry.com/models/llama3:v1.0.0 -t registry.com/models/llama3:v1.0.1 --output-remote + +# 覆盖/修改原来已存在的文件 +$ modctl attach foo.txt -s registry.com/models/llama3:v1.0.0 -t registry.com/models/llama3:v1.0.1 --output-remote --force +``` + +## 下载 +### Dragonfly 加速下载 +**支持的环境** + +[https://yuque.antfin.com/baimo.qwb/mecxix/otbfdnogfowuy6gt?singleDoc#](https://yuque.antfin.com/baimo.qwb/mecxix/otbfdnogfowuy6gt?singleDoc#) + +如果需要新增集群支持，咨询[@百蓦](https://yuque.antfin.com/baimo.qwb)[@肃晗](https://yuque.antfin.com/suhan.zcy)。 + +#### Pod 增加 NODE_IP ENV +POD 启动时，**增加 NODE_IP ENV 到 Pod**: + +```yaml +spec: + containers: + - env: + - name: NODE_IP + valueFrom: + fieldRef: + fieldPath: status.hostIP +``` + +#### Pod 内测试 NODE_IP 是否可用 +```shell +$ curl -v $NODE_IP:4003/healthy +* Trying 30.230.75.220:4003... +* Connected to 30.230.75.220 (30.230.75.220) port 4003 +> GET /healthy HTTP/1.1 +> Host: 30.230.75.220:4003 +> User-Agent: curl/8.4.0 +> Accept: */* +> +< HTTP/1.1 200 OK +< content-length: 0 +< date: Tue, 17 Dec 2024 09:38:36 GMT +< +* Connection #0 to host 30.230.75.220 left intact +``` + +#### 登录 +没有登陆的需要先[登录](#PlGag)。 + +#### 下载模型 +`/tmp/deepseek-v3/`即为需要下载模型的指定目录, `${namespace}`为申请的 Namespace。 + +```shell +$ modctl pull hmr.sa128.alipay.com/${namespace}/deepseek-v3:v1.0.2 --proxy http://$NODE_IP:4001 --plain-http --insecure --extract-from-remote --extract-dir /tmp/deepseek-v3/ +Copying blob sha256:1c05ac5a620306ecaec0e09e8adbbd4c25dc35b60825bef6a03c2aa4aab7939b skipped: already exists +Copying blob sha256:43105779bb4ceb010f28b3b6dc9a360455530209229948abc25900f18e21006c skipped: already exists +Copying blob sha256:011ff244caff15289524da3b802f90637f5720e6887e2d5de35d89edcc16c270 skipped: already exists +Copying blob sha256:0231d4dd488b86c56c657773598c8dadf0b9fc9563e8d630ba0909389d7f57d4 skipped: already exists +Copying blob sha256:9efe10adeb9c76db166da2e476035721d025bd51cace8c4ddf15c25c094b7f74 | 988 MB | done +Copying blob sha256:d5ca6ac7ac96ed52322cca8b860cd783114842d205b4084d14ef924214b6df6f | 7.0 MB | done +Copying blob sha256:373d09d0b4d7299829440fdedd5d783268baf031454d35e08ee1f8aea5bbaf00 | 9.2 kB | done +Copying blob sha256:7ff052b041a903c08f4335b9b3d9bd4260f342aa3d7cd570a682a6a588bf4e33 | 2.8 MB | done +Copying config sha256:73ec349249e76abe2bde801d664aecf42cbb4adee66da0b1a01792aa78632a33 | 44 B | done +Copying manifest sha256:94c1c9dc014f65dd96ca76557d82d863d538855d60e2f52ad3f9e9bc66feccda | 3.0 kB | done +Successfully pulled model artifact: hmr.sa128.alipay.com/${namespace}/deepseek-v3:v1.0.2 +``` + + + +### 普通下载 +#### 登录 +没有登陆的需要先[登录](#PlGag)。 + +#### 下载模型 +`/tmp/deepseek-v3/`即为需要下载模型的指定目录, `${namespace}`为申请的 Namespace。 + +```yaml +$ modctl pull hmr.sa128.alipay.com/${namespace}/deepseek-v3:v1.0.2 --plain-http --insecure --extract-from-remote --extract-dir /tmp/deepseek-v3/ +Copying blob sha256:1c05ac5a620306ecaec0e09e8adbbd4c25dc35b60825bef6a03c2aa4aab7939b skipped: already exists +Copying blob sha256:43105779bb4ceb010f28b3b6dc9a360455530209229948abc25900f18e21006c skipped: already exists +Copying blob sha256:011ff244caff15289524da3b802f90637f5720e6887e2d5de35d89edcc16c270 skipped: already exists +Copying blob sha256:0231d4dd488b86c56c657773598c8dadf0b9fc9563e8d630ba0909389d7f57d4 skipped: already exists +Copying blob sha256:9efe10adeb9c76db166da2e476035721d025bd51cace8c4ddf15c25c094b7f74 | 988 MB | done +Copying blob sha256:d5ca6ac7ac96ed52322cca8b860cd783114842d205b4084d14ef924214b6df6f | 7.0 MB | done +Copying blob sha256:373d09d0b4d7299829440fdedd5d783268baf031454d35e08ee1f8aea5bbaf00 | 9.2 kB | done +Copying blob sha256:7ff052b041a903c08f4335b9b3d9bd4260f342aa3d7cd570a682a6a588bf4e33 | 2.8 MB | done +Copying config sha256:73ec349249e76abe2bde801d664aecf42cbb4adee66da0b1a01792aa78632a33 | 44 B | done +Copying manifest sha256:94c1c9dc014f65dd96ca76557d82d863d538855d60e2f52ad3f9e9bc66feccda | 3.0 kB | done +Successfully pulled model artifact: hmr.sa128.alipay.com/${namespace}/deepseek-v3:v1.0.2 +``` + +### 部分下载 +在某些场景下，可能并不需要使用到模型中的所有文件，只需要使用到部分文件，那么也支持通过如下命令来下载部分文件。 + +```yaml +$ modctl fetch registry.com/models/llama3:v1.0.0 --output /path/to/extract --patterns '*.json' +``` + +> ⚠️ 如果想匹配子目录中的文件，则需要把目录层级写出来，例如想匹配一级目录下的 json 文件需要写`*/*.json` +> diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/A3_V026_DEPLOY.md b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/A3_V026_DEPLOY.md new file mode 100644 index 00000000..f3d1da20 --- /dev/null +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/A3_V026_DEPLOY.md @@ -0,0 +1,125 @@ +# A3 vLLM 0.26 CAM Async AFD deployment + +This is the reproducible single-node, 16-NPU deployment for the DeepSeek-V4 +Flash INT8 CAM async recipe. It starts Attention as `DP2 x TP4` on NPUs `0-7` +and the shared FFN pool as `DP8 x TP1 x EP8` on NPUs `8-15`. + +## 1. Create the task + +From the local AFD workspace, use the checked-in task wrapper: + +```bash +bash run_itask.sh afd_bjf_26 +``` + +The wrapper selects the validated nightly A3 image, 16 cards, host networking, +the DeepSeek-V4 model reference, and this task workdir: + +```text +/a3_inference/itask/workdir/wb02363348/bjf_afd/code +``` + +Wait for `itask list` to report `Running` before executing commands in the +task. + +## 2. Pin the runtime sources + +Inside the task, pin the two editable source trees that come with the image: + +```bash +git -C /vllm-workspace/vllm checkout --detach v0.26.0 +git -C /vllm-workspace/vllm-ascend checkout -B releases/v0.26.0rc \ + --track origin/releases/v0.26.0rc +``` + +Expected revisions are `vllm` `568afb3a` and `vllm-ascend` `80d8c194f`. + +## 3. Install CAM and AFD + +Synchronize the local `afd-plugin` directory to the task first: + +```bash +cd /path/to/afd-plugin +itask sync afd_bjf_26 --verbose +``` + +Then, inside the task: + +```bash +source /usr/local/Ascend/cann-9.0.1/set_env.sh +cd /a3_inference/itask/workdir/wb02363348/bjf_afd/code/afd-plugin + +bash afd_plugin/connectors/npu/bin/CAM_ascend910_93_openEuler_aarch64.run +# CAM installs its vendor op-api as libcust_opapi.so, but the spawned runtime +# resolves CAM aclnn symbols through the libopapi.so SONAME. +mv /usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libcust_opapi.so \ + /usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libopapi.so +/usr/local/python3.12.13/bin/python -m pip install \ + afd_plugin/connectors/npu/bin/umdk_cam_op_lib-209.0.0b1-cp312-cp312-linux_aarch64.whl + +# CAM async uses the four CAM vendor operators. It does not need the legacy +# AFD A2E/E2A build, whose rebuild can fail on a shared NFS worktree that +# already contains build artifacts. +AFD_BUILD_ASCEND_OPS=0 /usr/local/python3.12.13/bin/python -m pip install \ + -v --no-build-isolation --no-deps -e . +``` + +Do not use `AFD_BUILD_ASCEND_OPS=0` for `CAMP2pAFDConnector`; that connector +needs AFD's A2E/E2A extension and should be built in a clean worktree. + +## 4. Runtime-loader requirement for this nightly image + +The image's `set_env.sh` does not export all dynamic-library directories needed +by `torch_npu`. Both `ffn_ep8.sh` and `attention_tp8.sh` therefore prepend: + +```text +/usr/local/Ascend/driver/lib64/driver +/usr/local/Ascend/driver/lib64 +/usr/local/Ascend/cann-9.0.1/aarch64-linux/lib64 +/usr/local/Ascend/cann-9.0.1/runtime/lib64 +``` + +Keep this before importing vLLM or `torch_npu`. The scripts also configure the +CAM `op_api` path and preload CAM's renamed `libopapi.so`. + +Validate the installed CAM operators before a full model launch: + +```bash +source /usr/local/Ascend/cann-9.0.1/set_env.sh +export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64:/usr/local/Ascend/cann-9.0.1/aarch64-linux/lib64:/usr/local/Ascend/cann-9.0.1/runtime/lib64:/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib:/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api:$LD_LIBRARY_PATH +export CAM_CUST_OPAPI_LIB_PATH=/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libopapi.so +python -c 'from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available; ensure_cam_async_ops_available(); print("AFD_CAM_OPS_OK")' +``` + +## 5. Start the DP2 deployment + +Run the launcher from the AFD repository root: + +```bash +START_DELAY_SECONDS=120 \ + bash recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_dp2tp4.sh +``` + +It starts FFN first, waits for its model workers and CAM rendezvous to be +ready, then starts Attention. Logs and process IDs are written to: + +```text +/tmp/afd_dsv4_async_dp2tp4/ffn.log +/tmp/afd_dsv4_async_dp2tp4/attention.log +/tmp/afd_dsv4_async_dp2tp4/ffn.pid +/tmp/afd_dsv4_async_dp2tp4/attention.pid +``` + +The Attention OpenAI endpoint is `http://127.0.0.1:8900`; FFN listens on +`8901` only as the connector endpoint. + +## 6. Check startup + +```bash +tail -f /tmp/afd_dsv4_async_dp2tp4/ffn.log +tail -f /tmp/afd_dsv4_async_dp2tp4/attention.log +``` + +Start sending requests only after both logs report that their API server is +running. A failure in either role should be diagnosed from its own log before +restarting the pair. diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/README.md b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/README.md new file mode 100644 index 00000000..26e596d3 --- /dev/null +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/README.md @@ -0,0 +1,63 @@ +# DeepSeek-V4 Flash-INT8 Async CAM target baseline + +This recipe is the M0--M4 target-only baseline. It deliberately excludes +AFD-managed async MoE ubatching. + +Both roles must use the same model checkpoint and identical AFD topology. The +first smoke configuration is one Attention rank and one FFN rank: + +```json +{ + "afd": { + "role": "attention", + "connector": "CAMAsyncAFDConnector", + "async": true, + "host": "", + "port": 1239, + "num_attention_ranks": 1, + "num_ffn_ranks": 1, + "compute_gate_on_attention": true, + "connector_extra_config": { + "dynamicQuant": 1, + "attn_ranks_per_dp": 1, + "async_moe_ubatching": false + } + } +} +``` + +On the FFN process, change only `role` to `"ffn"`. Start FFN before Attention. +Use `--enforce-eager`, `--quantization ascend`, and the native DSV4 tokenizer +configuration (`--tokenizer-mode deepseek_v4`). Do not enable vLLM DBO, +`--num-ubatches` or async MoE ubatching. + +For the single-node 16-card baseline, use the checked-in scripts: + +```bash +bash recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_8a8f.sh +tail -f /tmp/afd_dsv4_async/ffn.log +tail -f /tmp/afd_dsv4_async/attention.log +``` + +The launch shape is eight Attention DP ranks on devices `0-7` and eight FFN +EP ranks on devices `8-15`; it follows the existing DSV4 AFD launch +shape and is the practical minimum for the full model. + +For M4, expand the same topology to two Attention and two FFN ranks: + +```json +{ + "num_attention_ranks": 2, + "num_ffn_ranks": 2, + "connector_extra_config": { + "dynamicQuant": 1, + "attn_ranks_per_dp": 1, + "async_moe_ubatching": false + } +} +``` + +The HCCL rank order is `[A0, A1, F0, F1]`. For each scale-out run, verify both +roles share `host`, `port`, rank counts, `attn_ranks_per_dp`, model revision, +and CAM operator installation. Compare greedy output with native DSV4 before +measuring throughput. diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/attention_tp8.sh b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/attention_tp8.sh new file mode 100644 index 00000000..53ef9548 --- /dev/null +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/attention_tp8.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PLUGIN_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" + +# Keep Attention and FFN on the same DSV4 Flash MTP checkpoint. See ffn_ep8.sh. +MODEL_PATH="/mnt/sfs_turbo/models/DeepSeek-V4-Flash-w8a8-mtp" +export MODEL_PATH +# Keep the rendezvous/HCCL address local to the current task Pod. itask may +# migrate a stopped task, leaving a stale inherited IP in its environment. +AFD_HOST="${AFD_HOST_OVERRIDE:-$(awk '!/^#/ && index($1, "127.") != 1 && $1 != "::1" {print $1; exit}' /etc/hosts)}" +: "${AFD_PORT:=1239}" +: "${API_PORT:=8900}" +NIC_NAME="${NIC_NAME_OVERRIDE:-eth0}" +: "${MAX_NUM_BATCHED_TOKENS:=1024}" +: "${ASCEND_RT_VISIBLE_DEVICES:=0,1,2,3,4,5,6,7}" +: "${AFD_DP_SIZE:=1}" +: "${AFD_TP_SIZE:=8}" +: "${AFD_SHARED_FFN_POOL:=false}" +export ASCEND_RT_VISIBLE_DEVICES + +# Keep the runtime environment identical to the proven DSV4 AFD +# recipes. In particular, worker subprocesses need CANN's driver and ATB +# libraries in addition to the CAM operator library below. +: "${CANN_SET_ENV:=/usr/local/Ascend/cann-9.0.1/set_env.sh}" +if [[ ! -f "$CANN_SET_ENV" ]]; then + echo "CANN environment script not found: $CANN_SET_ENV" >&2 + exit 1 +fi +# shellcheck disable=SC1090 +source "$CANN_SET_ENV" + +# The nightly A3 image's CANN setup script does not export the driver or +# toolkit runtime directories. torch_npu needs both before vLLM imports it. +export LD_LIBRARY_PATH="/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64:/usr/local/Ascend/cann-9.0.1/aarch64-linux/lib64:/usr/local/Ascend/cann-9.0.1/runtime/lib64:${LD_LIBRARY_PATH:-}" +export PYTHONPATH="$PLUGIN_ROOT:/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:${PYTHONPATH:-}" +export VLLM_PLUGINS="ascend,afd" +export AFD_FORCE_BALANCED_TOPK_IDS=0 +# ``set_env.sh`` owns the CANN runtime library order. Prepending CAM's +# devlib/op_api directories here makes vllm_ascend_C initialize against a +# mixed runtime before NPUWorker applies the role-local device mapping. +# CAM discovery uses ASCEND_CUSTOM_OPP_PATH below and must not override it. +# Required by the CAM Async Connector user guide: custom-op discovery and +# both op_api paths must precede the inherited CANN loader path. +export ASCEND_CUSTOM_OPP_PATH="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM:${ASCEND_CUSTOM_OPP_PATH:-}" +export LD_LIBRARY_PATH="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib:${LD_LIBRARY_PATH:-}" +export LD_LIBRARY_PATH="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api:${LD_LIBRARY_PATH}" +# CAM's vendor library must be named libopapi.so: the runtime resolves the +# aclnn symbols through that SONAME inside spawned worker processes. +CAM_CUST_OPAPI="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libopapi.so" +export LD_PRELOAD="$CAM_CUST_OPAPI${LD_PRELOAD:+:$LD_PRELOAD}" +export HCCL_IF_IP="$AFD_HOST" +export GLOO_SOCKET_IFNAME="$NIC_NAME" +export TP_SOCKET_IFNAME="$NIC_NAME" +export HCCL_SOCKET_IFNAME="$AFD_HOST" +export HCCL_BUFFSIZE=4096 +export HCCL_OP_EXPANSION_MODE="${HCCL_OP_EXPANSION_MODE:-AIV}" +# Attention constructs the external CAM communicator while its model runner is +# being created. Match the FFN-side connection window for the full 16-rank +# world so it can wait for post-load FFN initialization. +export HCCL_CONNECT_TIMEOUT="${HCCL_CONNECT_TIMEOUT:-1800}" +export HCCL_EXEC_TIMEOUT="${HCCL_EXEC_TIMEOUT:-1800}" +export OMP_PROC_BIND=false +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-10}" +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export VLLM_ASCEND_ENABLE_FLASHCOMM1=0 +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-2400}" +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS="${VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS:-3000}" +# Keep worker creation consistent with the FFN endpoint; see ffn_ep8.sh. +# The task environment may export ``forkserver``. Python 3.12's forkserver +# fails while restoring this runtime's signal handlers, so this recipe must +# override (rather than default) the inherited setting. +export VLLM_WORKER_MULTIPROC_METHOD=spawn +export AFD_FORCE_SPAWN_MULTIPROCESSING=1 + +ADDITIONAL_CONFIG="$(printf '%s' "{ + \"enable_force_load_balance\": false, + \"afd\": { + \"role\": \"attention\", + \"connector\": \"CAMAsyncAFDConnector\", + \"async\": true, + \"host\": \"$AFD_HOST\", + \"port\": $AFD_PORT, + \"num_attention_ranks\": 8, + \"num_ffn_ranks\": 8, + \"compute_gate_on_attention\": true, + \"connector_extra_config\": { + \"dynamicQuant\": 1, + \"attn_ranks_per_dp\": $AFD_TP_SIZE, + \"shared_ffn_pool\": $AFD_SHARED_FFN_POOL, + \"async_moe_ubatching\": true + } + } +}")" + +# Load from the same shared checkpoint with bounded per-rank concurrency. +exec env VLLM_USE_V1=1 /usr/local/python3.12.13/bin/vllm serve "$MODEL_PATH" \ + --host 0.0.0.0 \ + --port "$API_PORT" \ + --api-server-count 1 \ + --served-model-name dsv4-async \ + --worker-cls afd_plugin.v1.worker.npu.AFDNPUAttentionWorker \ + --data-parallel-size "$AFD_DP_SIZE" \ + --tensor-parallel-size "$AFD_TP_SIZE" \ + --enable-expert-parallel \ + --enforce-eager \ + --quantization ascend \ + --tokenizer-mode deepseek_v4 \ + --tool-call-parser deepseek_v4 \ + --enable-auto-tool-choice \ + --reasoning-parser deepseek_v4 \ + --block-size 128 \ + --max-model-len 8192 \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --max-num-seqs 2 \ + --gpu-memory-utilization 0.70 \ + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 16}' \ + --trust-remote-code \ + --additional-config "$ADDITIONAL_CONFIG" diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/diagnose_worker_startup.sh b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/diagnose_worker_startup.sh new file mode 100644 index 00000000..c1eae170 --- /dev/null +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/diagnose_worker_startup.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# Startup isolation for DeepSeek-V4 Async CAM AFD. +# Usage: bash diagnose_worker_startup.sh native|attention-worker|ffn-worker +set -euo pipefail + +MODE="${1:?usage: $0 native|attention-worker|ffn-worker}" +case "$MODE" in native|attention-worker|ffn-worker) ;; *) echo "unknown mode: $MODE" >&2; exit 2 ;; esac + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PLUGIN_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" +LOG_DIR="${LOG_DIR:-/tmp/afd_dsv4_startup_diagnose}" +MODEL_PATH="${MODEL_PATH:-/mnt/sfs_turbo/models/DeepSeek-V4-Flash-w8a8-mtp}" +AFD_HOST="${AFD_HOST:-33.215.116.191}" +AFD_PORT="${AFD_PORT:-1239}" +CANN_SET_ENV="${CANN_SET_ENV:-/usr/local/Ascend/cann-9.0.1/set_env.sh}" + +[[ -f "$CANN_SET_ENV" ]] || { echo "missing CANN environment: $CANN_SET_ENV" >&2; exit 1; } +# shellcheck disable=SC1090 +source "$CANN_SET_ENV" +export PYTHONPATH="$PLUGIN_ROOT:/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:${PYTHONPATH:-}" +export VLLM_PLUGINS="ascend,afd" +export LD_LIBRARY_PATH="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib:${LD_LIBRARY_PATH:-}" +export LD_PRELOAD="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libopapi.so${LD_PRELOAD:+:$LD_PRELOAD}" +export ASCEND_CUSTOM_OPP_PATH="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM:${ASCEND_CUSTOM_OPP_PATH:-}" +export HCCL_IF_IP="$AFD_HOST" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-eth0}" +export TP_SOCKET_IFNAME="$GLOO_SOCKET_IFNAME" +export HCCL_SOCKET_IFNAME="$GLOO_SOCKET_IFNAME" +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-2400}" +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS="${VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS:-3000}" +export VLLM_ASCEND_ENABLE_FLASHCOMM1=0 + +mkdir -p "$LOG_DIR" +LOG_FILE="$LOG_DIR/${MODE}.log" +: >"$LOG_FILE" + +if [[ "$MODE" == native ]]; then + # Native control: use the container's proven DSV4 Flash INT8 shape. TP8 + # does not fit this checkpoint on one 910A3 partition; DP4 x TP4 does. + # Keep afd-plugin loaded for its DSV4 Ascend compatibility patch, but use + # the stock NPU worker and native model implementation (no AFD worker). + export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}" + export VLLM_PLUGINS="ascend,afd" + exec env VLLM_USE_V1=1 /usr/local/python3.12.13/bin/vllm serve "$MODEL_PATH" \ + --host 0.0.0.0 --port 8910 --served-model-name dsv4-native-diagnose \ + --tensor-parallel-size 4 --data-parallel-size 4 --enable-expert-parallel \ + --enforce-eager --quantization ascend --tokenizer-mode deepseek_v4 \ + --max-model-len 8192 --max-num-batched-tokens 1024 --max-num-seqs 2 \ + --gpu-memory-utilization 0.90 \ + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 128}' \ + --trust-remote-code >>"$LOG_FILE" 2>&1 +fi + +if [[ "$MODE" == attention-worker ]]; then + export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" + worker="afd_plugin.v1.worker.npu.AFDNPUAttentionWorker" + role="attention" + port=8911 + gpu_memory_utilization=0.85 +else + export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-8,9,10,11,12,13,14,15}" + worker="afd_plugin.v1.worker.npu.AFDNPUFFNWorker" + role="ffn" + port=8912 + gpu_memory_utilization=0.90 +fi + +# Keep the real role-local 8-card shape: a single rank cannot hold the FFN +# experts. The absent peer role makes CAM rendezvous block only after model +# construction, which is the intended boundary for this diagnostic. +ADDITIONAL_CONFIG="{\"afd\":{\"role\":\"$role\",\"connector\":\"CAMAsyncAFDConnector\",\"async\":true,\"host\":\"$AFD_HOST\",\"port\":$AFD_PORT,\"num_attention_ranks\":8,\"num_ffn_ranks\":8,\"compute_gate_on_attention\":true,\"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":1,\"async_moe_ubatching\":false}}}" +exec env VLLM_USE_V1=1 /usr/local/python3.12.13/bin/vllm serve "$MODEL_PATH" \ + --host 0.0.0.0 --port "$port" --served-model-name "dsv4-$role-diagnose" \ + --worker-cls "$worker" --tensor-parallel-size 1 --data-parallel-size 8 \ + --enable-expert-parallel --enforce-eager --quantization ascend \ + --tokenizer-mode deepseek_v4 --max-model-len 8192 \ + --max-num-batched-tokens 1024 --max-num-seqs 2 --gpu-memory-utilization "$gpu_memory_utilization" \ + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 128}' \ + --trust-remote-code --additional-config "$ADDITIONAL_CONFIG" >>"$LOG_FILE" 2>&1 diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/ffn_ep8.sh b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/ffn_ep8.sh new file mode 100644 index 00000000..4a2e506a --- /dev/null +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/ffn_ep8.sh @@ -0,0 +1,128 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PLUGIN_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" + +# This recipe is exclusively for DSV4 Flash MTP. The task runtime injects a +# generic MODEL_PATH, which currently points to a DeepseekV32 checkpoint. +MODEL_PATH="/mnt/sfs_turbo/models/DeepSeek-V4-Flash-w8a8-mtp" +export MODEL_PATH +# The task can be migrated to another host between launches. Do not inherit +# the old Pod IP from the itask environment: CAM rendezvous and HCCL must bind +# to an address present on this host. A manual override remains available for +# a deliberately multi-host launch. +AFD_HOST="${AFD_HOST_OVERRIDE:-$(awk '!/^#/ && index($1, "127.") != 1 && $1 != "::1" {print $1; exit}' /etc/hosts)}" +: "${AFD_PORT:=1239}" +: "${API_PORT:=8901}" +NIC_NAME="${NIC_NAME_OVERRIDE:-eth0}" +: "${MAX_NUM_BATCHED_TOKENS:=1024}" +: "${ASCEND_RT_VISIBLE_DEVICES:=8,9,10,11,12,13,14,15}" +: "${AFD_DP_SIZE:=8}" +: "${AFD_TP_SIZE:=1}" +: "${AFD_ATTN_RANKS_PER_DP:=8}" +: "${AFD_SHARED_FFN_POOL:=false}" +export ASCEND_RT_VISIBLE_DEVICES + +# Keep the runtime environment identical to the proven DSV4 AFD +# recipes. In particular, worker subprocesses need CANN's driver and ATB +# libraries in addition to the CAM operator library below. +: "${CANN_SET_ENV:=/usr/local/Ascend/cann-9.0.1/set_env.sh}" +if [[ ! -f "$CANN_SET_ENV" ]]; then + echo "CANN environment script not found: $CANN_SET_ENV" >&2 + exit 1 +fi +# shellcheck disable=SC1090 +source "$CANN_SET_ENV" + +# The nightly A3 image's CANN setup script does not export the driver or +# toolkit runtime directories. torch_npu needs both before vLLM imports it. +export LD_LIBRARY_PATH="/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64:/usr/local/Ascend/cann-9.0.1/aarch64-linux/lib64:/usr/local/Ascend/cann-9.0.1/runtime/lib64:${LD_LIBRARY_PATH:-}" +export PYTHONPATH="$PLUGIN_ROOT:/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:${PYTHONPATH:-}" +export VLLM_PLUGINS="ascend,afd" +export AFD_FORCE_BALANCED_TOPK_IDS=0 +# ``set_env.sh`` owns the CANN runtime library order. Prepending CAM's +# devlib/op_api directories here makes vllm_ascend_C initialize against a +# mixed runtime before NPUWorker applies the role-local device mapping. +# CAM discovery uses ASCEND_CUSTOM_OPP_PATH below and must not override it. +# Required by the CAM Async Connector user guide: custom-op discovery and +# both op_api paths must precede the inherited CANN loader path. +export ASCEND_CUSTOM_OPP_PATH="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM:${ASCEND_CUSTOM_OPP_PATH:-}" +export LD_LIBRARY_PATH="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib:${LD_LIBRARY_PATH:-}" +export LD_LIBRARY_PATH="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api:${LD_LIBRARY_PATH}" +# CAM's vendor library must be named libopapi.so: the runtime resolves the +# aclnn symbols through that SONAME inside spawned worker processes. +CAM_CUST_OPAPI="/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libopapi.so" +export LD_PRELOAD="$CAM_CUST_OPAPI${LD_PRELOAD:+:$LD_PRELOAD}" +export HCCL_IF_IP="$AFD_HOST" +export GLOO_SOCKET_IFNAME="$NIC_NAME" +export TP_SOCKET_IFNAME="$NIC_NAME" +export HCCL_SOCKET_IFNAME="$AFD_HOST" +export HCCL_BUFFSIZE=4096 +export HCCL_OP_EXPANSION_MODE="${HCCL_OP_EXPANSION_MODE:-AIV}" +# FFN builds the CAM communicator only after its model workers have loaded. +# Give the separately launched Attention workers sufficient time to join this +# external 16-rank HCCL world instead of using CANN's short default timeout. +export HCCL_CONNECT_TIMEOUT="${HCCL_CONNECT_TIMEOUT:-1800}" +export HCCL_EXEC_TIMEOUT="${HCCL_EXEC_TIMEOUT:-1800}" +export OMP_PROC_BIND=false +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-10}" +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export VLLM_ASCEND_ENABLE_FLASHCOMM1=0 +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-2400}" +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS="${VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS:-3000}" +# Python 3.12's forkserver may retain an invalid signal-handler sentinel when +# vLLM creates DP workers, producing repeated ``signal handler must be ...`` +# failures in the child. Use fresh spawn workers for this multi-process CAM +# recipe; it also prevents the parent CAM runtime from being inherited. +# Do not retain an inherited ``forkserver`` setting: it is incompatible with +# this Python 3.12/CANN process tree's signal-handler restoration. +export VLLM_WORKER_MULTIPROC_METHOD=spawn +export AFD_FORCE_SPAWN_MULTIPROCESSING=1 + +ADDITIONAL_CONFIG="$(printf '%s' "{ + \"enable_force_load_balance\": false, + \"afd\": { + \"role\": \"ffn\", + \"connector\": \"CAMAsyncAFDConnector\", + \"async\": true, + \"host\": \"$AFD_HOST\", + \"port\": $AFD_PORT, + \"num_attention_ranks\": 8, + \"num_ffn_ranks\": 8, + \"compute_gate_on_attention\": true, + \"connector_extra_config\": { + \"dynamicQuant\": 1, + \"attn_ranks_per_dp\": $AFD_ATTN_RANKS_PER_DP, + \"shared_ffn_pool\": $AFD_SHARED_FFN_POOL, + \"async_moe_ubatching\": true + } + } +}")" + +# CAM owns one external work queue per expert rank. Keep these eight ranks as +# DP8 x TP1 x EP8: the CAM async operator's FFN endpoint is an EP rank, not a +# tensor-parallel shard. This matches the validated CAM topology and avoids +# passing a TP8 FFN communicator to the dispatch/combine kernels. +# Each rank loads from the shared checkpoint, so bound loader concurrency to +# avoid page-lock contention on the model filesystem. +exec env VLLM_USE_V1=1 /usr/local/python3.12.13/bin/vllm serve "$MODEL_PATH" \ + --host 0.0.0.0 \ + --port "$API_PORT" \ + --api-server-count 1 \ + --served-model-name dsv4-async-ffn \ + --worker-cls afd_plugin.v1.worker.npu.AFDNPUFFNWorker \ + --data-parallel-size "$AFD_DP_SIZE" \ + --tensor-parallel-size "$AFD_TP_SIZE" \ + --enable-expert-parallel \ + --enforce-eager \ + --quantization ascend \ + --tokenizer-mode deepseek_v4 \ + --block-size 128 \ + --max-model-len 8192 \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --max-num-seqs 2 \ + --gpu-memory-utilization 0.70 \ + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 16}' \ + --trust-remote-code \ + --additional-config "$ADDITIONAL_CONFIG" diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_8a8f.sh b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_8a8f.sh new file mode 100644 index 00000000..890c32b0 --- /dev/null +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_8a8f.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +LOG_DIR="${LOG_DIR:-/tmp/afd_dsv4_async}" +# FFN workers load the full EP shard before they enter the connector loop. +# Starting Attention after that window avoids CAM HCCL joining while FFN is +# still spawning its DP8 worker processes. +START_DELAY_SECONDS="${START_DELAY_SECONDS:-120}" +mkdir -p "$LOG_DIR" + +: >"$LOG_DIR/ffn.log" +setsid nohup bash "$SCRIPT_DIR/ffn_ep8.sh" >"$LOG_DIR/ffn.log" 2>&1 < /dev/null & +echo $! >"$LOG_DIR/ffn.pid" + +sleep "$START_DELAY_SECONDS" + +: >"$LOG_DIR/attention.log" +setsid nohup bash "$SCRIPT_DIR/attention_tp8.sh" >"$LOG_DIR/attention.log" 2>&1 < /dev/null & +echo $! >"$LOG_DIR/attention.pid" + +printf 'FFN PID: %s\nAttention PID: %s\nLogs: %s\n' \ + "$(cat "$LOG_DIR/ffn.pid")" \ + "$(cat "$LOG_DIR/attention.pid")" \ + "$LOG_DIR" diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_dp2tp4.sh b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_dp2tp4.sh new file mode 100644 index 00000000..2fe43f92 --- /dev/null +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/launch_dp2tp4.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +LOG_DIR="${LOG_DIR:-/tmp/afd_dsv4_async_dp2tp4}" +START_DELAY_SECONDS="${START_DELAY_SECONDS:-120}" +mkdir -p "$LOG_DIR" + +: >"$LOG_DIR/ffn.log" +setsid nohup env AFD_DP_SIZE=8 AFD_TP_SIZE=1 AFD_ATTN_RANKS_PER_DP=4 AFD_SHARED_FFN_POOL=true \ + bash "$SCRIPT_DIR/ffn_ep8.sh" >"$LOG_DIR/ffn.log" 2>&1 < /dev/null & +echo $! >"$LOG_DIR/ffn.pid" + +sleep "$START_DELAY_SECONDS" + +: >"$LOG_DIR/attention.log" +setsid nohup env AFD_DP_SIZE=2 AFD_TP_SIZE=4 AFD_SHARED_FFN_POOL=true \ + bash "$SCRIPT_DIR/attention_tp8.sh" >"$LOG_DIR/attention.log" 2>&1 < /dev/null & +echo $! >"$LOG_DIR/attention.pid" + +printf 'FFN PID: %s\nAttention PID: %s\nLogs: %s\n' \ + "$(cat "$LOG_DIR/ffn.pid")" \ + "$(cat "$LOG_DIR/attention.pid")" \ + "$LOG_DIR" diff --git a/recipe/npu/legacy_experiments/attn.sh b/recipe/npu/legacy_experiments/attn.sh new file mode 100644 index 00000000..919791ed --- /dev/null +++ b/recipe/npu/legacy_experiments/attn.sh @@ -0,0 +1,70 @@ +#!/bin/bash +# export AFD_CAMP2P_STUB_IO=1 +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 +export PYTHONPATH="/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:$PYTHONPATH" +export HCCL_BUFFSIZE=1024 +# ---- NPU profiler (AFD_NPU_ATTENTION_PROFILER_*) ---- +# 重要:trace 只在 vllm 正常 shutdown() 时 flush。停服务必须用 SIGTERM(kill -TERM / pkill -TERM), +# 绝不能 kill -9、进程也不能崩 —— 否则采不到/不完整。 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export AFD_NPU_ATTENTION_PROFILER_ENABLE=true +export AFD_NPU_ATTENTION_PROFILER_DIR="$SCRIPT_DIR/profile/attn" # 产物落这(绝对路径) +export AFD_NPU_ATTENTION_PROFILER_SKIP_FIRST=50 # 跳过前 50 step(默认 1500);要稳态可调大 +export AFD_NPU_ATTENTION_PROFILER_ACTIVE=10 # 采集 10 step + +nic_name="eth0" +local_ip="33.182.142.7" +export HCCL_IF_IP=$local_ip +export GLOO_SOCKET_IFNAME=$nic_name +export TP_SOCKET_IFNAME=$nic_name +export HCCL_SOCKET_IFNAME=$nic_name + +MODEL=/home/admin/model-csi/model + +VLLM_USE_V1=1 vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port 8006 \ + --worker-cls afd_plugin.v1.worker.ascend.AFDNPUAttentionWorker \ + --tensor-parallel-size 1 \ + --data-parallel-size 16 \ + --enable-expert-parallel \ + --max_num_batched_tokens 8 \ + --max_num_seqs 8 \ + --seed 1024 \ + --max-model-len 8192 \ + --gpu-memory-utilization 0.93 \ + --async-scheduling \ + --served-model-name dsv3 \ + --trust-remote-code \ + --no-enable-prefix-caching \ + --quantization ascend \ + --tokenizer-mode deepseek_v32 \ + --reasoning-parser deepseek_v3 \ + --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": ['8']}' \ + --kv-transfer-config '{ + "kv_connector": "AFDDecodeBenchConnector", + "kv_connector_module_path": "afd_plugin.connectors.decode_bench", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "fill_mean": 0.015, + "fill_std": 0.0 + } + }' \ + --additional-config '{ + "enable_cpu_binding": false, + "finegrained_tp_config": {"lmhead_tensor_parallel_size": 16}, + "enable_force_load_balance": true, + "afd": { + "enabled": true, + "role": "attention", + "connector": "camp2pconnector", + "host": "33.182.140.93", + "port": 29666, + "num_attention_servers": 16, + "num_ffn_servers": 16, + "afd_server_rank": 0 + } + }' > attn.log 2>&1 & +VLLM_PID=$! +echo "$VLLM_PID" > attn.pid +disown "$VLLM_PID" # 脱离 shell,itask session 断了也活着;停的时候用 kill -TERM $(cat attn.pid) diff --git a/recipe/npu/legacy_experiments/d0.sh b/recipe/npu/legacy_experiments/d0.sh new file mode 100644 index 00000000..3774680a --- /dev/null +++ b/recipe/npu/legacy_experiments/d0.sh @@ -0,0 +1,77 @@ +#!/bin/bash +# Decode node 1 template (2P2D, 16 NPU/node, A3). ▒~V~R~T▒~V~R launch_dp.py 驱▒~V~R~J▒~V~R▒~V~R~@~B +# ▒~V~R~\▒~V~R▒~V~R~J~B▒~V~R~B▒~V~R▒~V~R~Q 16 个▒~V~R~^▒~V~R~K (rank 0-15)▒~V~R~L▒~V~R~N d1 ▒~V~R~E▒~V~R▒~V~R~P~L▒~V~R~D▒~V~R~H~P dp=32 ▒~V~R~Z~D decode ▒~V~R~D▒~V~R~Lbroker = ▒~V~R~\▒~V~R▒~V~R~\▒~V~R▒~V~R~@~B +# ▒~V~R~O~B▒~V~R~U▒~V~R▒~V~R~H▒~V~R~N run_dp_template.sh ▒~V~R~@▒~V~R~G▒~V~R▒~V~R~I: +# $1 ▒~V~R~O▒~V~R▒~V~R~A设▒~V~R~G $2 vllm 端▒~V~R~O▒~V~R $3 dp_size $4 dp_rank +# $5 dp_address $6 rpc_port $7 tp_size +export PYTHONPATH="/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:$PYTHONPATH" +export VLLM_USE_V1=1 +export ASCEND_RT_VISIBLE_DEVICES=$1 +export HCCL_BUFFSIZE=1024 +export AFD_NPU_ATTENTION_PROFILER_ENABLE=true +export AFD_NPU_ATTENTION_PROFILER_DIR=./profile/attn +export AFD_NPU_ATTENTION_PROFILER_SKIP_FIRST=50 # 默认 1500，太大，改小让它快点抓 +export AFD_NPU_ATTENTION_PROFILER_ACTIVE=10 + +nic_name="eth0" +local_ip="33.182.142.4" # ▒~V~R~V~R~\▒~V~R~V~R▒~V~R~V~R~\▒~V~R~V~R IP▒~V~R~V~R~Hd0 ▒~V~R~V~R~J~B▒~V~R~V~R~B▒~V~R~V~R▒~V~R~V~R~I +export HCCL_IF_IP=$local_ip +export GLOO_SOCKET_IFNAME=$nic_name +export TP_SOCKET_IFNAME=$nic_name + +MODEL=/home/admin/model-csi/model + +ADDITIONAL='{ + "enable_cpu_binding": false, + "finegrained_tp_config": {"lmhead_tensor_parallel_size": 16}, + "afd": { + "enabled": true, + "role": "attention", + "connector": "camp2pconnector", + "host": "33.182.141.223", + "port": 29666, + "num_attention_servers": 16, + "num_ffn_servers": 16, + "afd_server_rank": 0 + } +}' +COMPILATION='{"cudagraph_mode":"FULL_DECODE_ONLY", "cudagraph_capture_sizes":['64']}' +KV_TRANSFER='{ + "kv_connector": "MooncakeConnectorV1", + "kv_role": "kv_consumer", + "kv_port": "36200", + "kv_connector_module_path": "vllm_ascend.distributed.mooncake_connector", + "kv_connector_extra_config": { + "use_ascend_direct": true, + "prefill": {"dp_size": 2, "tp_size": 8}, + "decode": {"dp_size": 16, "tp_size": 1} + } +}' + +# --async-scheduling \ + +vllm serve "$MODEL" \ + --port $2 \ + --worker-cls afd_plugin.v1.worker.ascend.AFDNPUAttentionWorker \ + --data-parallel-size $3 \ + --data-parallel-rank $4 \ + --data-parallel-address $5 \ + --data-parallel-rpc-port $6 \ + --tensor-parallel-size $7 \ + --enable-expert-parallel \ + --seed 1024 \ + --served-model-name dsv3 \ + --max-model-len 8192 \ + --max-num-batched-tokens 64 \ + --trust-remote-code \ + --max-num-seqs 64 \ + --gpu-memory-utilization 0.93 \ + --no-enable-prefix-caching \ + --quantization ascend \ + --tokenizer-mode deepseek_v32 \ + --reasoning-parser deepseek_v3 \ + --async-scheduling \ + --additional-config "$ADDITIONAL" \ + --kv-transfer-config "$KV_TRANSFER" \ + --compilation-config "$COMPILATION" + > d0.log 2>&1 & \ No newline at end of file diff --git a/recipe/npu/legacy_experiments/d1.sh b/recipe/npu/legacy_experiments/d1.sh new file mode 100644 index 00000000..15559ced --- /dev/null +++ b/recipe/npu/legacy_experiments/d1.sh @@ -0,0 +1,79 @@ +#!/bin/bash +# Decode node 2 template (2P2D, 16 NPU/node, A3). ▒~V~R~T▒~V~R launch_dp.py 驱▒~V~R~J▒~V~R▒~V~R~@~B +# ▒~V~R~\▒~V~R▒~V~R~J~B▒~V~R~B▒~V~R▒~V~R~Q 16 个▒~V~R~^▒~V~R~K (rank 16-31)▒~V~R~Ldp_address ▒~V~R~L~G▒~V~R~P~Q d0 (decode ▒~V~R~D▒~V~R~Z~D▒~V~R~E▒~V~R享 broker)▒~V~R~@~B +# ▒~V~R~O~B▒~V~R~U▒~V~R▒~V~R~H▒~V~R~N run_dp_template.sh ▒~V~R~@▒~V~R~G▒~V~R▒~V~R~I: +# $1 ▒~V~R~O▒~V~R▒~V~R~A设▒~V~R~G $2 vllm 端▒~V~R~O▒~V~R $3 dp_size $4 dp_rank +# $5 dp_address $6 rpc_port $7 tp_size +export PYTHONPATH="/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:$PYTHONPATH" +export VLLM_USE_V1=1 +export ASCEND_RT_VISIBLE_DEVICES=$1 +export HCCL_BUFFSIZE=1024 +export AFD_NPU_FFN_PROFILER_ENABLE=true +export AFD_NPU_FFN_PROFILER_DIR=./profile/ffn +export AFD_NPU_FFN_PROFILER_SKIP_FIRST=50 +export AFD_NPU_FFN_PROFILER_ACTIVE=20 # ffn 默认 20 步 + +nic_name="eth0" +local_ip="33.182.141.132" # ▒~V~R~V~R~\▒~V~R~V~R▒~V~R~V~R~\▒~V~R~V~R IP +export HCCL_IF_IP=$local_ip +export GLOO_SOCKET_IFNAME=$nic_name +export TP_SOCKET_IFNAME=$nic_name +export HCCL_SOCKET_IFNAME=$nic_name + +MODEL=/home/admin/model-csi/model + +ADDITIONAL='{ + "enable_cpu_binding": false, + "finegrained_tp_config": {"lmhead_tensor_parallel_size": 16}, + "afd": { + "enabled": true, + "role": "ffn", + "connector": "camp2pconnector", + "host": "33.182.141.223", + "port": 29666, + "num_attention_servers": 16, + "num_ffn_servers": 16, + "afd_server_rank": 0 + } +}' +COMPILATION='{"cudagraph_mode":"FULL_DECODE_ONLY", "cudagraph_capture_sizes":['64']}' +KV_TRANSFER='{ + "kv_connector": "MooncakeConnectorV1", + "kv_role": "kv_consumer", + "kv_port": "36200", + "engine_id": "2", + "kv_connector_module_path": "vllm_ascend.distributed.mooncake_connector", + "kv_connector_extra_config": { + "use_ascend_direct": true, + "prefill": {"dp_size": 2, "tp_size": 8}, + "decode": {"dp_size": 16, "tp_size": 1} + } +}' + +# --async-scheduling \ + +vllm serve "$MODEL" \ + --port $2 \ + --worker-cls afd_plugin.v1.worker.ascend.AFDNPUFFNWorker \ + --data-parallel-size $3 \ + --data-parallel-rank $4 \ + --data-parallel-address $5 \ + --data-parallel-rpc-port $6 \ + --tensor-parallel-size $7 \ + --enable-expert-parallel \ + --seed 1024 \ + --async-scheduling \ + --served-model-name dsv3 \ + --max-model-len 8192 \ + --max-num-batched-tokens 100 \ + --trust-remote-code \ + --max-num-seqs 14 \ + --gpu-memory-utilization 0.93 \ + --no-enable-prefix-caching \ + --quantization ascend \ + --tokenizer-mode deepseek_v32 \ + --reasoning-parser deepseek_v3 \ + --additional-config "$ADDITIONAL" \ + --kv-transfer-config "$KV_TRANSFER" \ + --compilation-config "$COMPILATION" + > d1.log 2>&1 & \ No newline at end of file diff --git a/recipe/npu/legacy_experiments/dsv4_baseline.sh b/recipe/npu/legacy_experiments/dsv4_baseline.sh new file mode 100644 index 00000000..85438175 --- /dev/null +++ b/recipe/npu/legacy_experiments/dsv4_baseline.sh @@ -0,0 +1,44 @@ +#!/bin/bash +set -euo pipefail + +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 +export PYTHONPATH="/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:${PYTHONPATH:-}" +export OMP_PROC_BIND=false +export OMP_NUM_THREADS=10 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export HCCL_BUFFSIZE=1024 +export VLLM_ASCEND_ENABLE_FLASHCOMM1=1 +export TASK_QUEUE_ENABLE=1 +export HCCL_OP_EXPANSION_MODE=AIV + +JEMALLOC=/usr/lib/aarch64-linux-gnu/libjemalloc.so.2 +if [[ -f "$JEMALLOC" ]]; then + export LD_PRELOAD="$JEMALLOC${LD_PRELOAD:+:$LD_PRELOAD}" +fi + +MODEL=/mnt/sfs_turbo/models/DeepSeek-V4-Flash-w8a8-mtp +LOG_FILE=${LOG_FILE:-dsv4_baseline.log} +PROFILER_DIR=${PROFILER_DIR:-/a3_inference/itask/workdir/wb02363348/bjf_afd/code/afd/profiles/dsv4_stack_memory_shapes} +mkdir -p "$PROFILER_DIR" + +exec vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port 8900 \ + --served-model-name dsv4 \ + --data-parallel-size 4 \ + --tensor-parallel-size 4 \ + --enable-expert-parallel \ + --quantization ascend \ + --tokenizer-mode deepseek_v4 \ + --tool-call-parser deepseek_v4 \ + --enable-auto-tool-choice \ + --reasoning-parser deepseek_v4 \ + --block-size 128 \ + --max-model-len 8192 \ + --max-num-batched-tokens 4096 \ + --max-num-seqs 4 \ + --gpu-memory-utilization 0.90 \ + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 128}' \ + --profiler-config "{\"profiler\":\"torch\",\"torch_profiler_dir\":\"$PROFILER_DIR\",\"torch_profiler_with_stack\":true,\"torch_profiler_record_shapes\":true,\"torch_profiler_with_memory\":true,\"torch_profiler_use_gzip\":false,\"ignore_frontend\":true}" \ + --enforce-eager \ + >"$LOG_FILE" 2>&1 \ No newline at end of file diff --git a/recipe/npu/legacy_experiments/ffn.sh b/recipe/npu/legacy_experiments/ffn.sh new file mode 100644 index 00000000..0a1db09f --- /dev/null +++ b/recipe/npu/legacy_experiments/ffn.sh @@ -0,0 +1,70 @@ +#!/bin/bash +# export AFD_CAMP2P_STUB_IO=1 +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 +export PYTHONPATH="/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:$PYTHONPATH" +export HCCL_BUFFSIZE=1024 +# ---- NPU profiler (AFD_NPU_FFN_PROFILER_*) ---- +# 重要:trace 只在 vllm 正常 shutdown() 时 flush。停服务必须用 SIGTERM(kill -TERM / pkill -TERM), +# 绝不能 kill -9、进程也不能崩 —— 否则采不到/不完整。 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export AFD_NPU_FFN_PROFILER_ENABLE=true +export AFD_NPU_FFN_PROFILER_DIR="$SCRIPT_DIR/profile/ffn" # 产物落这(绝对路径) +export AFD_NPU_FFN_PROFILER_SKIP_FIRST=50 # 跳过前 50 step(默认 1500);要稳态可调大 +export AFD_NPU_FFN_PROFILER_ACTIVE=20 # 采集 20 step + +nic_name="eth0" +local_ip="33.182.140.93" +export HCCL_IF_IP=$local_ip +export GLOO_SOCKET_IFNAME=$nic_name +export TP_SOCKET_IFNAME=$nic_name +export HCCL_SOCKET_IFNAME=$nic_name + +MODEL=/home/admin/model-csi/model + +VLLM_USE_V1=1 vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port 8006 \ + --worker-cls afd_plugin.v1.worker.ascend.AFDNPUFFNWorker \ + --tensor-parallel-size 1 \ + --data-parallel-size 16 \ + --enable-expert-parallel \ + --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": ['8']}' \ + --max_num_seqs 8 \ + --seed 1024 \ + --max_num_batched_tokens 32 \ + --max-model-len 8192 \ + --gpu-memory-utilization 0.93 \ + --async-scheduling \ + --served-model-name dsv3 \ + --trust-remote-code \ + --no-enable-prefix-caching \ + --quantization ascend \ + --tokenizer-mode deepseek_v32 \ + --reasoning-parser deepseek_v3 \ + --kv-transfer-config '{ + "kv_connector": "AFDDecodeBenchConnector", + "kv_connector_module_path": "afd_plugin.connectors.decode_bench", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "fill_mean": 0.015, + "fill_std": 0.0 + } + }' \ + --additional-config '{ + "enable_cpu_binding": false, + "finegrained_tp_config": {"lmhead_tensor_parallel_size": 16}, + "enable_force_load_balance": true, + "afd": { + "enabled": true, + "role": "ffn", + "connector": "camp2pconnector", + "host": "33.182.140.93", + "port": 29666, + "num_attention_servers": 16, + "num_ffn_servers": 16, + "afd_server_rank": 0 + } + }' > ffn.log 2>&1 & +VLLM_PID=$! +echo "$VLLM_PID" > ffn.pid +disown "$VLLM_PID" # 脱离 shell,itask session 断了也活着;停的时候用 kill -TERM $(cat ffn.pid) diff --git a/recipe/npu/legacy_experiments/ffn1.sh b/recipe/npu/legacy_experiments/ffn1.sh new file mode 100644 index 00000000..6a0844a0 --- /dev/null +++ b/recipe/npu/legacy_experiments/ffn1.sh @@ -0,0 +1,70 @@ +#!/bin/bash +export ASCEND_RT_VISIBLE_DEVICES=$1 +export PYTHONPATH="/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:$PYTHONPATH" +export HCCL_BUFFSIZE=2048 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export AFD_NPU_FFN_PROFILER_ENABLE=true +export AFD_NPU_FFN_PROFILER_DIR="$SCRIPT_DIR/profile/ffn" +mkdir -p "$AFD_NPU_FFN_PROFILER_DIR" +export AFD_NPU_FFN_PROFILER_SKIP_FIRST=50 +export AFD_NPU_FFN_PROFILER_ACTIVE=20 +nic_name="eth0" +local_ip="33.215.118.135" +export HCCL_IF_IP=$local_ip +export GLOO_SOCKET_IFNAME=$nic_name +export TP_SOCKET_IFNAME=$nic_name +export HCCL_SOCKET_IFNAME=$nic_name + +MODEL=/a3_inference/itask/workdir/shared/jcz/model/dsv3.2 +#MODEL=/home/admin/model-csi/model + +VLLM_USE_V1=1 vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port 8006 \ + --load-format dummy \ + --worker-cls afd_plugin.v1.worker.ascend.AFDNPUFFNWorker \ + --data-parallel-size 16 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": ['16']}' \ + --max_num_seqs 16 \ + --seed 1024 \ + --max_num_batched_tokens 16 \ + --max-model-len 18432 \ + --gpu-memory-utilization 0.93 \ + --async-scheduling \ + --served-model-name dsv3 \ + --trust-remote-code \ + --no-enable-prefix-caching \ + --quantization ascend \ + --tokenizer-mode deepseek_v32 \ + --reasoning-parser deepseek_v3 \ + --kv-transfer-config '{ + "kv_connector": "AFDDecodeBenchConnector", + "kv_connector_module_path": "afd_plugin.connectors.decode_bench", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "fill_mean": 0.015, + "fill_std": 0.0 + } + }' \ + --additional-config '{ + "enable_force_load_balance": true, + "force_load_balance_topn_per_rank": 4, + "afd": { + "enabled": true, + "role": "ffn", + "connector": "camp2pconnector", + "host": "33.215.118.135", + "port": 29666, + "num_attention_ranks": 48, + "num_ffn_ranks": 16, + "extra_config": { + "afd_size": "48A16F" + } + } + }' > ffn1.log 2>&1 & +VLLM_PID=$! +echo "$VLLM_PID" > ffn1.pid +disown "$VLLM_PID" diff --git a/recipe/npu/legacy_experiments/pd_colo.sh b/recipe/npu/legacy_experiments/pd_colo.sh new file mode 100644 index 00000000..cb6a8924 --- /dev/null +++ b/recipe/npu/legacy_experiments/pd_colo.sh @@ -0,0 +1,58 @@ + +#!/bin/bash +# export AFD_CAMP2P_STUB_IO=1 +export PYTHONPATH="/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:$PYTHONPATH" +MODEL=/home/admin/model-csi/model + +VLLM_USE_V1=1 vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port 8006 \ + --tensor-parallel-size 8 \ + --enforce-eager \ + --data-parallel-size 2 \ + --enable-expert-parallel \ + --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": ['8']}' \ + --max_num_seqs 8 \ + --quantization ascend \ + --max_num_batched_tokens 32 \ + --gpu-memory-utilization 0.93 \ + --max-model-len 8192 + + +# #!/bin/bash +# export PYTHONPATH="/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:$PYTHONPATH" +# export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:$LD_LIBRARY_PATH +# export MOONCAKE_CONFIG_PATH="/vllm-workspace/mooncake.json" +# export VLLM_USE_V1=1 +# export ASCEND_RT_VISIBLE_DEVICES=$1 +# export HCCL_IF_IP=$2 +# export GLOO_SOCKET_IFNAME=eth0 +# export TP_SOCKET_IFNAME=eth0 +# export HCCL_SOCKET_IFNAME=eth0 + +# MODEL=/home/admin/model-csi/model + +# vllm serve "$MODEL" \ +# --served-model-name dsv3 \ +# --host $2 \ +# --port $3 \ +# --tensor-parallel-size 8 \ +# --enable-expert-parallel \ +# --quantization ascend \ +# --max-model-len 8192 \ +# --max-num-batched-tokens 16384 \ +# --max-num-seqs 16 \ +# --gpu-memory-utilization 0.9 \ +# --tokenizer-mode deepseek_v32 \ +# --reasoning-parser deepseek_v3 \ +# --trust-remote-code --no-enable-prefix-caching \ +# --kv-transfer-config '{ +# "kv_connector": "MooncakeConnectorStoreV1", +# "kv_role": "kv_both", +# "kv_connector_extra_config": { +# "use_layerwise": false, +# "mooncake_rpc_port": "0", +# "load_async": true, +# "register_buffer": true +# } +# }' \ No newline at end of file diff --git a/tools/analysis/legacy/analyze2.py b/tools/analysis/legacy/analyze2.py new file mode 100644 index 00000000..955c8ab5 --- /dev/null +++ b/tools/analysis/legacy/analyze2.py @@ -0,0 +1,32 @@ +import collections +lines = open("/tmp/afd_dsv4_async/ffn.log", "r", errors="replace").read().splitlines() +recv = collections.Counter() +comb = collections.Counter() +section = None +marker = "first5=[" +for line in lines: + if "async_dispatch_recv outputs" in line: + section = "recv" + continue + if "async_combine_send inputs" in line: + section = "comb" + continue + if "async_dispatch_recv inputs" in line or "async_combine_recv" in line or "async_dispatch_send" in line: + section = None + continue + i = line.find(marker) + if i < 0 or section is None: + continue + j = line.find("]", i) + vals = [int(x.strip()) for x in line[i + len(marker):j].split(",")] + key = (vals[1], vals[2]) + if section == "recv": + recv[key] += 1 + elif section == "comb": + comb[key] += 1 +print("=== dispatch_recv outputs (rankid, layer) ===") +for k in sorted(recv): + print(recv[k], k) +print("=== combine_send inputs (rankid, layer) ===") +for k in sorted(comb): + print(comb[k], k) diff --git a/tools/analysis/legacy/analyze_ffn_log.py b/tools/analysis/legacy/analyze_ffn_log.py new file mode 100644 index 00000000..8fdab125 --- /dev/null +++ b/tools/analysis/legacy/analyze_ffn_log.py @@ -0,0 +1,24 @@ +import collections +txt = open("/tmp/afd_dsv4_async/ffn.log", "r", errors="replace").read().splitlines() +recv = collections.Counter() +comb = collections.Counter() +marker = "first5=[" +for line in txt: + i = line.find(marker) + if i < 0: + continue + j = line.find("]", i) + if j < 0: + continue + vals = [int(x.strip()) for x in line[i + len(marker):j].split(",")] + key = (vals[1], vals[2]) + if "async_dispatch_recv outputs" in line: + recv[key] += 1 + elif "async_combine_send inputs" in line: + comb[key] += 1 +print("=== dispatch_recv outputs (rankid, layer) ===") +for k in sorted(recv): + print(recv[k], k) +print("=== combine_send inputs (rankid, layer) ===") +for k in sorted(comb): + print(comb[k], k) diff --git a/tools/analysis/legacy/dbg.py b/tools/analysis/legacy/dbg.py new file mode 100644 index 00000000..b3058a28 --- /dev/null +++ b/tools/analysis/legacy/dbg.py @@ -0,0 +1,9 @@ +import os +print("exists:", os.path.exists("/tmp/afd_dsv4_async/ffn.log")) +print("size:", os.path.getsize("/tmp/afd_dsv4_async/ffn.log")) +with open("/tmp/afd_dsv4_async/ffn.log", "r", errors="replace") as f: + for i, line in enumerate(f): + if "first5=" in line: + print("line", i, line[:200]) + break +print("done") diff --git a/tools/analysis/legacy/dbg2.py b/tools/analysis/legacy/dbg2.py new file mode 100644 index 00000000..ef981894 --- /dev/null +++ b/tools/analysis/legacy/dbg2.py @@ -0,0 +1,15 @@ +import collections +lines = open("/tmp/afd_dsv4_async/ffn.log", "r", errors="replace").read().splitlines() +print("total lines:", len(lines)) +marker = "first5=[" +cnt = 0 +recv = 0 +comb = 0 +for line in lines: + if marker in line: + cnt += 1 + if "async_dispatch_recv outputs" in line: + recv += 1 + if "async_combine_send inputs" in line: + comb += 1 +print("marker hits:", cnt, "recv:", recv, "comb:", comb) diff --git a/tools/analysis/legacy/fix_launch_scripts.py b/tools/analysis/legacy/fix_launch_scripts.py new file mode 100644 index 00000000..f45d3246 --- /dev/null +++ b/tools/analysis/legacy/fix_launch_scripts.py @@ -0,0 +1,15 @@ +import io + +path = "afd-plugin/recipe/npu/CAMAsyncAFDConnector/deepseek_v4/attention_tp8.sh" +with io.open(path, "r", encoding="utf-8") as f: + text = f.read() +old = " --enable-expert-parallel \\\n --enable-dbo \\\n" +new = " --enable-expert-parallel \\\n" +assert old in text, "enable-dbo block not found" +text = text.replace(old, new) +with io.open(path, "w", encoding="utf-8", newline="\n") as f: + f.write(text) +for i, line in enumerate(text.splitlines(), 1): + if "--enable-dbo" in line or "async_moe_ubatching" in line: + print(i, line) +print("done") diff --git a/tools/analysis/legacy/fix_ubatch_wrapper.py b/tools/analysis/legacy/fix_ubatch_wrapper.py new file mode 100644 index 00000000..2f8d35ae --- /dev/null +++ b/tools/analysis/legacy/fix_ubatch_wrapper.py @@ -0,0 +1,27 @@ +import io + +path = "afd-plugin/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py" +with io.open(path, "r", encoding="utf-8") as f: + text = f.read() + +old = """ self.comm_stream = torch.npu.Stream(device=device) + assert self.vllm_config.parallel_config.num_ubatches == AFD_NPU_NUM_UBATCHES + self.ready_barrier = threading.Barrier(_READY_BARRIER_PARTIES)""" + +new = """ self.comm_stream = torch.npu.Stream(device=device) + # vLLM native ubatching/DBO is rejected for CAMAsyncAFDConnector by + # fail_if_unsupported_npu_afd_features; the AFD async MoE ubatching + # path drives stage counts through connector_extra_config + # (async_moe_num_ubatches), which is validated there as well. + if self.vllm_config.parallel_config.use_ubatching: + raise RuntimeError( + "AscendUBatchWrapper does not support vLLM native " + "ubatching/DBO; use the AFD async MoE ubatching path.", + ) + self.ready_barrier = threading.Barrier(_READY_BARRIER_PARTIES)""" + +assert old in text, "target block not found" +text = text.replace(old, new) +with io.open(path, "w", encoding="utf-8", newline="\n") as f: + f.write(text) +print("patched ok") diff --git a/tools/analysis/legacy/gen_b64.py b/tools/analysis/legacy/gen_b64.py new file mode 100644 index 00000000..b768d507 --- /dev/null +++ b/tools/analysis/legacy/gen_b64.py @@ -0,0 +1,11 @@ +import base64 +s = """import json,urllib.request +req=urllib.request.Request('http://127.0.0.1:8900/v1/chat/completions',data=json.dumps({"model":"dsv4-async","messages":[{"role":"user","content":"1+1=? Reply with just the number."}],"max_tokens":32,"temperature":0}).encode(),headers={'Content-Type':'application/json'}) +try: + r=urllib.request.urlopen(req,timeout=120) + print('HTTP',r.status) + print(r.read().decode()) +except Exception as e: + print('ERR',repr(e)) +""" +print(base64.b64encode(s.encode()).decode()) diff --git a/tools/analysis/legacy/launch_dp.py b/tools/analysis/legacy/launch_dp.py new file mode 100644 index 00000000..db5f32a2 --- /dev/null +++ b/tools/analysis/legacy/launch_dp.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Per-node launcher: spawn N DP vllm instances using a node-specific template. +相当于官方 examples/external_online_dp/launch_online_dp.py，多了 --template 参数， +这样每个节点可以直接用自己的 p0/p1/d0/d1 模板，不必都改名为 run_dp_template.sh。 + +用法见同目录脚本头注释 / README。 +""" +import argparse +import multiprocessing +import os +import subprocess +import sys + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--template", default="./run_dp_template.sh", + help="本节点的模板脚本路径，例如 ./p0.sh") + parser.add_argument("--dp-size", type=int, required=True, help="全局 DP 大小。") + parser.add_argument("--tp-size", type=int, default=1, help="每实例的 TP 大小。") + parser.add_argument("--dp-size-local", type=int, default=-1, help="本节点实例数。默认=dp-size。") + parser.add_argument("--dp-rank-start", type=int, default=0, help="本节点起始 dp_rank。") + parser.add_argument("--dp-address", type=str, required=True, help="DP broker 的 IP。") + parser.add_argument("--dp-rpc-port", type=str, default="12345", help="DP broker 的 RPC 端口。") + parser.add_argument("--vllm-start-port", type=int, default=9000, help="vllm 起始端口。") + return parser.parse_args() + + +args = parse_args() +dp_size = args.dp_size +tp_size = args.tp_size +dp_size_local = dp_size if args.dp_size_local == -1 else args.dp_size_local +dp_rank_start = args.dp_rank_start +dp_address = args.dp_address +dp_rpc_port = args.dp_rpc_port +vllm_start_port = args.vllm_start_port +template_path = args.template + +if not os.path.exists(template_path): + print(f"Template file {template_path} does not exist.") + sys.exit(1) + +num_cards = dp_size_local * tp_size + + +def run_command(visible_devices, dp_rank, vllm_engine_port): + command = [ + "bash", template_path, + visible_devices, + str(vllm_engine_port), + str(dp_size), + str(dp_rank), + dp_address, + dp_rpc_port, + str(tp_size), + ] + subprocess.run(command, check=True) + + +if __name__ == "__main__": + processes = [] + for i in range(dp_size_local): + dp_rank = dp_rank_start + i + vllm_engine_port = vllm_start_port + i + visible_devices = ",".join(str(x) for x in range(i * tp_size, (i + 1) * tp_size)) + process = multiprocessing.Process( + target=run_command, + args=(visible_devices, dp_rank, vllm_engine_port), + ) + processes.append(process) + process.start() + + for process in processes: + process.join() diff --git a/tools/analysis/legacy/patch_layer_fix.py b/tools/analysis/legacy/patch_layer_fix.py new file mode 100644 index 00000000..8bbf94b5 --- /dev/null +++ b/tools/analysis/legacy/patch_layer_fix.py @@ -0,0 +1,34 @@ +import sys + +path = "/mnt/d/cyj/afd/afd-plugin/afd_plugin/connectors/npu/async_cam.py" +with open(path, "r", encoding="utf-8") as f: + src = f.read() + +old = """ # Preserve the vendor-produced count/rank fields and repair only its + # stale layer field before returning the header to combine-send. + layer_idx = int(expected_layer_idx) + token_nums_rankid_layeridx = token_nums_rankid_layeridx.clone() + token_nums_rankid_layeridx[2] = layer_idx + states.token_nums_rankid_layeridx = token_nums_rankid_layeridx + states.cam_dp_group_index = cam_dp_group_index""" + +new = """ # CAM returns the true decoder-layer index in the dispatch-recv header + # (field 2 of token_nums_rankid_layeridx). Use it for both the local + # FFN layer computation and the combine-send header so the completion + # matches the attention-side combine-recv. Do NOT overwrite field 2 + # with expected_layer_idx: the shared-FFN-pool scheduler alternates + # DP groups per layer, so expected_layer_idx is the busy-loop + # iteration index, not the true decoder layer. + layer_idx = max(0, int(token_nums_rankid_layeridx[2].item())) + token_nums_rankid_layeridx = token_nums_rankid_layeridx.clone() + states.token_nums_rankid_layeridx = token_nums_rankid_layeridx + states.cam_dp_group_index = cam_dp_group_index""" + +count = src.count(old) +assert count == 1, f"expected 1 occurrence, got {count}" +src = src.replace(old, new) + +with open(path, "w", encoding="utf-8") as f: + f.write(src) + +print("patched OK") diff --git a/tools/analysis/legacy/rewrite_test_force_lb.py b/tools/analysis/legacy/rewrite_test_force_lb.py new file mode 100644 index 00000000..2e440792 --- /dev/null +++ b/tools/analysis/legacy/rewrite_test_force_lb.py @@ -0,0 +1,268 @@ +from pathlib import Path + + +path = Path( + "/mnt/d/cyj/afd/afd-plugin/tests/unit/compat/patches/" + "test_force_load_balance.py" +) +path.write_text( + '''from __future__ import annotations + +import importlib +import sys +import types +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + + +class _QuantType: + NONE = 0 + W8A8 = 1 + + +class AscendFusedMoE: + """Stand-in for vllm_ascend.ops.fused_moe.fused_moe.AscendFusedMoE.""" + + def __init__(self, *args: object, **kwargs: object) -> None: + del args, kwargs + + +def build_fused_experts_input(*args: object, **kwargs: object) -> torch.Tensor: + """Fake builder: returns the possibly swapped topk_ids for assertion.""" + + del args + return kwargs["topk_ids"] + + +class AscendW8A8DynamicFusedMoEMethod: + def apply( + self, + layer: object, + x: object, + router_logits: object, + top_k: int, + renormalize: bool, + **kwargs: object, + ) -> torch.Tensor: + import vllm_ascend.quantization.methods.w8a8_dynamic as mod + + del layer, x, router_logits, top_k, renormalize + return mod.build_fused_experts_input(topk_ids=kwargs["topk_ids"]) + + +def _install_fake_modules(monkeypatch: pytest.MonkeyPatch) -> types.ModuleType: + vllm = types.ModuleType("vllm") + vllm_config = types.ModuleType("vllm.config") + vllm_config.VllmConfig = object + + root = types.ModuleType("vllm_ascend") + ops = types.ModuleType("vllm_ascend.ops") + fused_moe_pkg = types.ModuleType("vllm_ascend.ops.fused_moe") + fused_moe_mod = types.ModuleType("vllm_ascend.ops.fused_moe.fused_moe") + fused_moe_mod.AscendFusedMoE = AscendFusedMoE + + quant = types.ModuleType("vllm_ascend.quantization") + methods = types.ModuleType("vllm_ascend.quantization.methods") + w8a8_mod = types.ModuleType("vllm_ascend.quantization.methods.w8a8_dynamic") + w8a8_mod.AscendW8A8DynamicFusedMoEMethod = AscendW8A8DynamicFusedMoEMethod + w8a8_mod.build_fused_experts_input = build_fused_experts_input + + quant_type_mod = types.ModuleType("vllm_ascend.quantization.quant_type") + quant_type_mod.QuantType = _QuantType + + monkeypatch.setitem(sys.modules, "vllm", vllm) + monkeypatch.setitem(sys.modules, "vllm.config", vllm_config) + monkeypatch.setitem(sys.modules, "vllm_ascend", root) + monkeypatch.setitem(sys.modules, "vllm_ascend.ops", ops) + monkeypatch.setitem(sys.modules, "vllm_ascend.ops.fused_moe", fused_moe_pkg) + monkeypatch.setitem( + sys.modules, "vllm_ascend.ops.fused_moe.fused_moe", fused_moe_mod + ) + monkeypatch.setitem(sys.modules, "vllm_ascend.quantization", quant) + monkeypatch.setitem(sys.modules, "vllm_ascend.quantization.methods", methods) + monkeypatch.setitem( + sys.modules, "vllm_ascend.quantization.methods.w8a8_dynamic", w8a8_mod + ) + monkeypatch.setitem( + sys.modules, "vllm_ascend.quantization.quant_type", quant_type_mod + ) + return fused_moe_mod + + +@pytest.fixture +def force_lb_mod(monkeypatch: pytest.MonkeyPatch) -> types.ModuleType: + _install_fake_modules(monkeypatch) + module_name = "afd_plugin.compat.patches.force_load_balance" + sys.modules.pop(module_name, None) + mod = importlib.import_module(module_name) + mod.apply_force_load_balance_patch() + mod.apply_force_load_balance_patch() + return mod + + +def _new_layer(force_lb_mod: types.ModuleType) -> object: + return force_lb_mod.AscendFusedMoE.__new__(force_lb_mod.AscendFusedMoE) + + +def test_force_load_balance_buffer_topn_per_rank(force_lb_mod: types.ModuleType): + layer = _new_layer(force_lb_mod) + layer.ep_size = 4 + layer.n_routed_experts = 8 + layer.top_k = 2 + layer.force_load_balance_topn_per_rank = 1 + + force_lb_mod._init_force_lb_buffer( + layer, + max_tokens=4, + device=torch.device("cpu"), + ) + + expected = torch.tensor([[0, 2], [4, 6], [0, 2], [4, 6]], dtype=torch.int32) + assert torch.equal(layer.force_lb_fake_topk_buffer, expected) + + +def test_force_load_balance_buffer_uses_max_num_batched_tokens( + force_lb_mod: types.ModuleType, +): + max_tokens = force_lb_mod._get_force_lb_max_tokens( + SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=6)) + ) + assert max_tokens == 6 + + layer = _new_layer(force_lb_mod) + layer.ep_size = 2 + layer.n_routed_experts = 4 + layer.top_k = 2 + layer.force_load_balance_topn_per_rank = 0 + + force_lb_mod._init_force_lb_buffer( + layer, + max_tokens=max_tokens, + device=torch.device("cpu"), + ) + + assert layer.force_lb_fake_topk_buffer.shape == (6, 2) + + +def test_force_load_balance_max_tokens_falls_back_when_not_int( + force_lb_mod: types.ModuleType, +): + max_tokens = force_lb_mod._get_force_lb_max_tokens( + SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=None)) + ) + assert max_tokens == 128 + + +def test_force_load_balance_buffer_ids_within_routed_experts( + force_lb_mod: types.ModuleType, +): + layer = _new_layer(force_lb_mod) + layer.ep_size = 2 + layer.n_routed_experts = 4 + layer.global_num_experts = 6 + layer.top_k = 2 + layer.force_load_balance_topn_per_rank = 2 + + force_lb_mod._init_force_lb_buffer( + layer, + max_tokens=2, + device=torch.device("cpu"), + ) + + assert int(layer.force_lb_fake_topk_buffer.max()) < layer.n_routed_experts + + +def test_force_load_balance_full_expert_cycle_is_deterministic( + force_lb_mod: types.ModuleType, +): + config = force_lb_mod.ForceLoadBalanceConfig( + n_routed_experts=8, + ep_size=4, + top_k=2, + topn_per_rank=0, + ) + + first = force_lb_mod._build_expert_cycle(config, torch.device("cpu")) + second = force_lb_mod._build_expert_cycle(config, torch.device("cpu")) + + assert torch.equal(first, second) + assert sorted(first.tolist()) == list(range(8)) + + +def test_force_load_balance_buffer_grows_for_large_batch( + force_lb_mod: types.ModuleType, +): + layer = _new_layer(force_lb_mod) + layer.ep_size = 2 + layer.n_routed_experts = 4 + layer.top_k = 2 + layer.force_load_balance_topn_per_rank = 2 + + force_lb_mod._init_force_lb_buffer( + layer, + max_tokens=2, + device=torch.device("cpu"), + ) + topk_ids = force_lb_mod._get_force_lb_topk_ids( + layer, + batch_tokens=5, + device=torch.device("cpu"), + ) + + assert topk_ids.shape == (5, 2) + assert layer.force_lb_fake_topk_buffer.shape[0] >= 5 + + +def test_w8a8_apply_swaps_topk_ids_with_buffer(force_lb_mod: types.ModuleType): + method = force_lb_mod.AscendW8A8DynamicFusedMoEMethod() + + layer = _new_layer(force_lb_mod) + layer.enable_force_load_balance = True + layer.mix_placement = False + layer.top_k = 2 + layer.ep_size = 4 + layer.n_routed_experts = 8 + layer.force_load_balance_topn_per_rank = 1 + layer.force_lb_fake_topk_buffer = torch.tensor( + [[0, 2], [4, 6], [0, 2], [4, 6]], dtype=torch.int32 + ) + + real_topk_ids = torch.zeros((4, 2), dtype=torch.int64) + out = method.apply( + layer=layer, + x=None, + router_logits=None, + top_k=2, + renormalize=True, + topk_ids=real_topk_ids, + ) + + expected = layer.force_lb_fake_topk_buffer.to(torch.int64) + assert torch.equal(out, expected) + + +def test_w8a8_apply_passthrough_when_buffer_absent(force_lb_mod: types.ModuleType): + method = force_lb_mod.AscendW8A8DynamicFusedMoEMethod() + + layer = _new_layer(force_lb_mod) + layer.enable_force_load_balance = False + layer.force_lb_fake_topk_buffer = None + layer.mix_placement = False + layer.top_k = 2 + + real_topk_ids = torch.zeros((4, 2), dtype=torch.int64) + out = method.apply( + layer=layer, + x=None, + router_logits=None, + top_k=2, + renormalize=True, + topk_ids=real_topk_ids, + ) + + assert torch.equal(out, real_topk_ids) +''' +) diff --git a/tools/benchmarks/legacy/bench_pd.py b/tools/benchmarks/legacy/bench_pd.py new file mode 100644 index 00000000..fb74415e --- /dev/null +++ b/tools/benchmarks/legacy/bench_pd.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +"""独立 PD serving 测压脚本：直打 OpenAI /v1/completions(流式)，测 TTFT / TPOT / 吞吐。 +不依赖 vllm 的 benchmark 代码，只要 aiohttp + 一个可达的 endpoint。 +用法见文件底部注释。 +""" +import argparse +import asyncio +import json +import random +import string +import time + +import aiohttp + + +def gen_prompt(approx_words: int) -> str: + # 造一段随机"英文"，长度近似 approx_words 个词 + return "".join(random.choice(string.ascii_letters + " ") for _ in range(approx_words * 6)) + + +async def one_request(session, url, model, prompt, out_len, idx, results): + payload = { + "model": model, + "prompt": prompt, + "max_tokens": out_len, + "temperature": 1.0, + "stream": True, + } + t_send = time.perf_counter() + ttft = None + n_tok = 0 + try: + async with session.post(url, json=payload) as resp: + if resp.status != 200: + body = await resp.text() + results.append({"idx": idx, "error": f"HTTP {resp.status}: {body[:120]}"}) + return + async for raw in resp.content: + line = raw.decode(errors="ignore").strip() + if not line or not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + break + try: + obj = json.loads(data) + except Exception: + continue + choices = obj.get("choices", []) + if not choices: + continue + text = choices[0].get("text") or "" + if text: + if ttft is None: + ttft = time.perf_counter() - t_send + n_tok += 1 # vllm 流式通常每 chunk 1 token，近似计数 + except Exception as e: + results.append({"idx": idx, "error": f"{type(e).__name__}: {e}"}) + return + t_end = time.perf_counter() + lat = t_end - t_send + results.append({ + "idx": idx, "ttft": ttft, "lat": lat, "n_tok": n_tok, + "tpot": (lat - ttft) / (n_tok - 1) if (ttft and n_tok > 1) else None, + }) + + +async def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--base-url", default="http://127.0.0.1:8000", help="proxy 的地址，如 http://33.215.116.107:8000") + ap.add_argument("--model", required=True, help="服务端 served-model-name（没设就是权重路径）") + ap.add_argument("--endpoint", default="/v1/completions") + ap.add_argument("--num-prompts", type=int, default=200) + ap.add_argument("--request-rate", type=float, default=10.0, help="请求/秒；inf 表示一次性全发") + ap.add_argument("--input-len", type=int, default=512, help="输入近似词数") + ap.add_argument("--output-len", type=int, default=128, help="max_tokens") + args = ap.parse_args() + + url = args.base_url.rstrip("/") + args.endpoint + prompts = [gen_prompt(args.input_len) for _ in range(args.num_prompts)] + results = [] + delay = 0.0 if args.request_rate in (0, float("inf")) else 1.0 / args.request_rate + + async with aiohttp.ClientTimeout(total=None) as _t, aiohttp.ClientSession(timeout=_t) as session: + t0 = time.perf_counter() + tasks = [] + for i in range(args.num_prompts): + tasks.append(asyncio.create_task(one_request(session, url, args.model, prompts[i], args.output_len, i, results))) + if delay: + await asyncio.sleep(delay) + await asyncio.gather(*tasks) + wall = time.perf_counter() - t0 + + ok = [r for r in results if "error" not in r] + err = [r for r in results if "error" in r] + tot_tok = sum(r["n_tok"] for r in ok) + ttfts = [r["ttft"] for r in ok if r["ttft"] is not None] + tpots = [r["tpot"] for r in ok if r["tpot"] is not None] + lats = [r["lat"] for r in ok] + + def stat(xs): + if not xs: + return (0, 0, 0) + xs = sorted(xs) + n = len(xs) + return (xs[n // 2], sum(xs) / n, xs[int(n * 0.95)]) + + ttft_p50, ttft_mean, ttft_p95 = stat([x * 1000 for x in ttfts]) + tpot_p50, tpot_mean, tpot_p95 = stat([x * 1000 for x in tpots]) + lat_p50, lat_mean, lat_p95 = stat([x * 1000 for x in lats]) + + print("\n================ PD Serving Benchmark ================") + print(f"endpoint : {url}") + print(f"requests : {len(results)} (ok={len(ok)}, err={len(err)})") + print(f"input/output : ~{args.input_len} words / {args.output_len} tokens") + print(f"request rate : {args.request_rate} req/s") + print(f"wall time : {wall:.2f} s") + print("------------------------------------------------------") + print(f"output tokens : {tot_tok}") + print(f"output tput : {tot_tok / wall:.2f} tok/s") + print(f"request tput : {len(ok) / wall:.2f} req/s") + print(f"TTFT p50/mean/p95 : {ttft_p50:.1f} / {ttft_mean:.1f} / {ttft_p95:.1f} ms") + print(f"TPOT p50/mean/p95 : {tpot_p50:.1f} / {tpot_mean:.1f} / {tpot_p95:.1f} ms") + print(f"Lat p50/mean/p95 : {lat_p50:.1f} / {lat_mean:.1f} / {lat_p95:.1f} ms") + if err: + print("------------------------------------------------------") + print(f"errors ({len(err)}):") + for e in err[:5]: + print(f" req{e['idx']}: {e['error']}") + print("======================================================") + + +if __name__ == "__main__": + # 示例： + # python bench_pd.py --base-url http://33.215.116.107:8000 \ + # --model /home/admin/model-csi/model --num-prompts 200 --request-rate 10 + asyncio.run(main()) diff --git a/tools/benchmarks/legacy/merge_throughput.py b/tools/benchmarks/legacy/merge_throughput.py new file mode 100644 index 00000000..3de67bb7 --- /dev/null +++ b/tools/benchmarks/legacy/merge_throughput.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Merge AFD and baseline throughput curves into one figure.""" + +from __future__ import annotations + +import argparse +import csv +from pathlib import Path + + +def main() -> None: + args = parse_args() + afd_rows = read_curve(args.afd_csv) + baseline_rows = read_curve(args.baseline_csv) + plot_curves(afd_rows, baseline_rows, args.output) + print(f"wrote: {args.output}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--afd-csv", type=Path, required=True) + parser.add_argument("--baseline-csv", type=Path, required=True) + parser.add_argument( + "--output", + type=Path, + default=Path("afd_vs_baselien_throughput.png"), + ) + return parser.parse_args() + + +def read_curve(path: Path) -> list[tuple[float, float]]: + with path.open("r", encoding="utf-8-sig", newline="") as f: + reader = csv.DictReader(f) + rows: list[tuple[float, float]] = [] + for row in reader: + x = first_float(row, ("elapsed_s", "second", "time", "time_s")) + y = first_float(row, ("tokens/s/die", "tokens_per_s_per_die")) + if x is None or y is None: + continue + rows.append((x, y)) + if not rows: + raise SystemExit(f"No elapsed_s + tokens/s/die rows found in {path}") + return rows + + +def first_float(row: dict[str, str], keys: tuple[str, ...]) -> float | None: + normalized = {key.strip().lower(): value for key, value in row.items()} + for key in keys: + value = normalized.get(key.lower()) + if value in (None, ""): + continue + try: + return float(value) + except ValueError: + continue + return None + + +def plot_curves( + afd_rows: list[tuple[float, float]], + baseline_rows: list[tuple[float, float]], + output: Path, +) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(figsize=(12, 5)) + ax.plot( + [x for x, _ in afd_rows], + [y for _, y in afd_rows], + linewidth=1.6, + label="48a16f", + color="#1f77b4", + ) + ax.plot( + [x for x, _ in baseline_rows], + [y for _, y in baseline_rows], + linewidth=1.6, + label="ep64", + color="#d62728", + ) + ax.set_title("afd vs baselien throughput") + ax.set_xlabel("Elapsed time (s)") + ax.set_ylabel("tokens/s/die") + ax.grid(True, alpha=0.25) + ax.legend() + fig.tight_layout() + output.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output, dpi=160) + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/tools/benchmarks/legacy/throughput.py b/tools/benchmarks/legacy/throughput.py new file mode 100644 index 00000000..05521c23 --- /dev/null +++ b/tools/benchmarks/legacy/throughput.py @@ -0,0 +1,763 @@ +#!/usr/bin/env python3 +"""Build full-process throughput CSV/PNG files from one AISBench result dir.""" + +from __future__ import annotations + +import argparse +import csv +import io +import json +import math +import re +import sqlite3 +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +DEFAULT_SEARCH_ROOTS = ( + "outputs/default", + "outputs", + "benchmark/outputs/default", + "benchmark/outputs", + "/a3_inference/itask/workdir/wb02363348/cyj_afd/code/outputs/default", + "/a3_inference/itask/workdir/wb02363348/cyj_afd/code/benchmark/outputs/default", +) + +TIME_KEYS = ( + "end_time", + "finish_time", + "finished_time", + "response_time", + "completion_time", + "request_end_time", + "request_finished_time", + "timestamp_end", + "end_timestamp", +) +START_TIME_KEYS = ( + "start_time", + "request_start_time", + "send_time", + "sent_time", + "begin_time", + "timestamp_start", + "start_timestamp", +) +LATENCY_KEYS = ( + "latency", + "total_latency", + "request_latency", + "e2e_latency", + "duration", +) +OUTPUT_TOKEN_KEYS = ( + "output_tokens", + "output_token", + "output_token_num", + "output_token_len", + "output_token_length", + "output_len", + "output_length", + "generated_tokens", + "generated_token_num", + "generated_token_len", + "completion_tokens", + "num_output_tokens", + "decode_tokens", +) +TOKEN_TIME_KEYS = ( + "token_timestamps", + "output_token_timestamps", + "decode_token_timestamps", + "generated_token_timestamps", +) + + +@dataclass(frozen=True) +class RequestRecord: + start_time: float + end_time: float + output_tokens: int + token_timestamps: tuple[float, ...] = () + + +def main() -> None: + args = parse_args() + if args.result_file is not None: + result_file = args.result_file + records = load_records(result_file) + else: + ais_bench_dir = resolve_ais_bench_dir(args.ais_bench_dir, args.input_root) + result_file = None + records = [] + for candidate in find_candidate_results(ais_bench_dir): + candidate_records = load_records(candidate) + if candidate_records: + result_file = candidate + records = candidate_records + break + + if result_file is None: + raise SystemExit( + "No AISBench detail/db/json/csv result with request records found " + f"under: {ais_bench_dir}" + ) + + if not records: + raise SystemExit(f"No request records with time and output tokens in {result_file}") + + rows = build_throughput_rows(records, die_count=args.die_count) + output_prefix = Path(args.output_dir) / f"throughput_{sanitize_name(args.name)}" + csv_path = output_prefix.with_name(f"{output_prefix.name}_full_process.csv") + png_path = output_prefix.with_name(f"{output_prefix.name}_full_process.png") + per_die_png_path = output_prefix.with_name( + f"{output_prefix.name}_full_process_per_die.png" + ) + + csv_path.parent.mkdir(parents=True, exist_ok=True) + write_csv(csv_path, rows) + plot_png( + png_path, + rows, + y_key="tokens/s", + y_label="tokens/s", + title=f"{args.name} full-process throughput", + ) + plot_png( + per_die_png_path, + rows, + y_key="tokens/s/die", + y_label="tokens/s/die", + title=f"{args.name} full-process throughput per die", + ) + + total_tokens = sum(record.output_tokens for record in records) + duration = rows[-1]["elapsed_s"] + 1 if rows else 0 + peak = max(row["tokens/s"] for row in rows) if rows else 0.0 + print(f"input: {result_file}") + print(f"requests: {len(records)}") + print(f"output_tokens: {total_tokens}") + print(f"duration_s: {duration:.0f}") + print(f"peak_tokens/s: {peak:.2f}") + print(f"csv: {csv_path}") + print(f"png: {png_path}") + print(f"per_die_png: {per_die_png_path}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Generate throughput__full_process.csv/png and " + "throughput__full_process_per_die.png from AISBench details." + ) + ) + parser.add_argument( + "--name", + required=True, + help="Name used in output files, for example 16k_ep64.", + ) + parser.add_argument( + "--ais_bench_dir", + "--ais-bench-dir", + required=True, + help=( + "AISBench run directory name or path, for example 20260704_155051. " + "Only this directory is parsed." + ), + ) + parser.add_argument( + "--input-root", + action="append", + default=[], + help="AISBench output root to search. Can be specified multiple times.", + ) + parser.add_argument( + "--result-file", + type=Path, + help="Exact detail DB/JSON/JSONL/CSV file to parse.", + ) + parser.add_argument( + "--output-dir", + default=".", + help="Directory for generated CSV/PNG files.", + ) + parser.add_argument( + "--die-count", + type=int, + default=64, + help="Number of dies used for tokens/s/die. Default: 64.", + ) + return parser.parse_args() + + +def resolve_ais_bench_dir(value: str, input_roots: list[str]) -> Path: + normalized = value.strip().rstrip(",") + direct = Path(normalized) + if direct.exists(): + return direct + + roots = candidate_roots(input_roots) + for root in roots: + candidate = root / normalized + if candidate.exists(): + return candidate + + matches: list[Path] = [] + for root in roots: + if not root.exists() or not root.is_dir(): + continue + try: + matches.extend(path for path in root.rglob(normalized) if path.is_dir()) + except OSError: + continue + + if matches: + return max(matches, key=lambda path: path.stat().st_mtime) + + roots_text = ", ".join(str(root) for root in roots) + raise SystemExit( + f"AISBench directory {normalized!r} was not found under: {roots_text}" + ) + + +def find_candidate_results(ais_bench_dir: Path) -> list[Path]: + files: list[Path] = [] + if ais_bench_dir.is_file(): + files.append(ais_bench_dir) + else: + for pattern in ( + "**/*.db", + "**/*.sqlite", + "**/*.sqlite3", + "**/*.jsonl", + "**/*.json", + "**/*.csv", + ): + files.extend(ais_bench_dir.glob(pattern)) + + candidates = [ + path + for path in files + if path.is_file() + and "throughput_" not in path.name + and path.stat().st_size > 0 + and is_probable_detail_file(path) + ] + fallback = [ + path + for path in files + if path.is_file() and "throughput_" not in path.name and path.stat().st_size > 0 + ] + if candidates: + preferred = sorted( + candidates, + key=lambda path: path.stat().st_mtime, + reverse=True, + ) + rest = sorted( + [path for path in fallback if path not in set(preferred)], + key=lambda path: path.stat().st_mtime, + reverse=True, + ) + return preferred + rest + return sorted( + [ + path + for path in files + if path.is_file() and "throughput_" not in path.name and path.stat().st_size > 0 + ], + key=lambda path: path.stat().st_mtime, + reverse=True, + ) + + +def candidate_roots(input_roots: list[str]) -> list[Path]: + roots = [Path(root) for root in input_roots] + roots.extend(Path(root) for root in DEFAULT_SEARCH_ROOTS) + roots.append(Path.cwd()) + deduped: list[Path] = [] + seen: set[str] = set() + for root in roots: + key = str(root.resolve()) if root.exists() else str(root) + if key not in seen: + deduped.append(root) + seen.add(key) + return deduped + + +def is_probable_detail_file(path: Path) -> bool: + lowered = str(path).lower() + return any( + marker in lowered + for marker in ( + "detail", + "detailed", + "request", + "perf", + "benchmark", + "result", + "db", + ) + ) + + +def load_records(path: Path) -> list[RequestRecord]: + suffix = path.suffix.lower() + if suffix == ".jsonl" and "detail" in path.name.lower(): + records = load_aisbench_detail_records(path) + if records: + return records + if suffix in {".db", ".sqlite", ".sqlite3"}: + return load_sqlite_records(path) + if suffix == ".jsonl": + return load_jsonl_records(path) + if suffix == ".json": + return load_json_records(path) + if suffix == ".csv": + return load_csv_records(path) + raise ValueError(f"Unsupported result file type: {path}") + + +def load_aisbench_detail_records(path: Path) -> list[RequestRecord]: + """Load AISBench detail JSONL files that store time_points in db_data.""" + + db_cache: dict[str, dict[int, tuple[float, ...]]] = {} + records: list[RequestRecord] = [] + with path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + item = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(item, Mapping) or not item.get("success", True): + continue + + output_tokens = parse_int(item.get("output_tokens")) + db_name = item.get("db_name") + time_points = item.get("time_points") + if ( + output_tokens is None + or output_tokens <= 0 + or not isinstance(db_name, str) + or not isinstance(time_points, Mapping) + ): + continue + db_ref = parse_int(time_points.get("__db_ref__")) + if db_ref is None: + continue + + arrays = db_cache.get(db_name) + if arrays is None: + arrays = load_numpy_store(path.parent / "db_data" / db_name) + db_cache[db_name] = arrays + timestamps = arrays.get(db_ref) + if not timestamps: + continue + + token_timestamps: tuple[float, ...] = () + if len(timestamps) >= output_tokens: + token_timestamps = tuple(timestamps[-output_tokens:]) + records.append( + RequestRecord( + start_time=timestamps[0], + end_time=timestamps[-1], + output_tokens=output_tokens, + token_timestamps=token_timestamps, + ) + ) + return dedupe_records(records) + + +def load_numpy_store(path: Path) -> dict[int, tuple[float, ...]]: + arrays: dict[int, tuple[float, ...]] = {} + if not path.exists(): + return arrays + + import numpy as np + + conn = sqlite3.connect(path) + try: + for row_id, arr_blob in conn.execute("select id, arr_blob from numpy_store"): + try: + arr = np.load(io.BytesIO(arr_blob), allow_pickle=False) + except Exception: + continue + values = tuple(float(x) for x in arr.tolist()) + if values: + arrays[int(row_id)] = values + finally: + conn.close() + return arrays + + +def load_sqlite_records(path: Path) -> list[RequestRecord]: + records: list[RequestRecord] = [] + conn = sqlite3.connect(path) + conn.row_factory = sqlite3.Row + try: + tables = [ + row[0] + for row in conn.execute( + "select name from sqlite_master where type='table'" + ).fetchall() + ] + for table in tables: + try: + rows = conn.execute(f'select * from "{table}"').fetchall() + except sqlite3.DatabaseError: + continue + for row in rows: + record = record_from_mapping(dict(row)) + if record is not None: + records.append(record) + finally: + conn.close() + return dedupe_records(records) + + +def load_jsonl_records(path: Path) -> list[RequestRecord]: + records: list[RequestRecord] = [] + with path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + item = json.loads(line) + except json.JSONDecodeError: + continue + records.extend(records_from_json_item(item)) + return dedupe_records(records) + + +def load_json_records(path: Path) -> list[RequestRecord]: + with path.open("r", encoding="utf-8") as f: + data = json.load(f) + return dedupe_records(list(records_from_json_item(data))) + + +def load_csv_records(path: Path) -> list[RequestRecord]: + records: list[RequestRecord] = [] + with path.open("r", encoding="utf-8-sig", newline="") as f: + for row in csv.DictReader(f): + record = record_from_mapping(row) + if record is not None: + records.append(record) + return dedupe_records(records) + + +def records_from_json_item(item: Any) -> Iterator[RequestRecord]: + if isinstance(item, Mapping): + record = record_from_mapping(item) + if record is not None: + yield record + for value in item.values(): + if isinstance(value, list): + for child in value: + yield from records_from_json_item(child) + elif isinstance(value, Mapping): + yield from records_from_json_item(value) + elif isinstance(item, list): + for child in item: + yield from records_from_json_item(child) + + +def record_from_mapping(mapping: Mapping[str, Any]) -> RequestRecord | None: + flat = flatten_mapping(mapping) + output_tokens = extract_output_tokens(flat) + if output_tokens is None or output_tokens <= 0: + return None + + start_time = first_float(flat, START_TIME_KEYS) + end_time = first_float(flat, TIME_KEYS) + latency = first_float(flat, LATENCY_KEYS) + + if end_time is None and start_time is not None and latency is not None: + end_time = start_time + normalize_duration_seconds(latency) + if start_time is None and end_time is not None and latency is not None: + start_time = end_time - normalize_duration_seconds(latency) + if start_time is None or end_time is None: + return None + + start_time = normalize_timestamp_seconds(start_time) + end_time = normalize_timestamp_seconds(end_time) + if end_time < start_time: + start_time, end_time = end_time, start_time + if math.isclose(end_time, start_time): + end_time = start_time + 1e-6 + + token_timestamps = extract_token_timestamps(flat) + return RequestRecord( + start_time=start_time, + end_time=end_time, + output_tokens=output_tokens, + token_timestamps=token_timestamps, + ) + + +def flatten_mapping(mapping: Mapping[str, Any]) -> dict[str, Any]: + flat: dict[str, Any] = {} + + def visit(prefix: str, value: Any) -> None: + key = normalize_key(prefix) + if key: + flat[key] = value + flat.setdefault(key.split(".")[-1], value) + if isinstance(value, str): + parsed = try_json(value) + if parsed is not value: + visit(prefix, parsed) + elif isinstance(value, Mapping): + for child_key, child_value in value.items(): + visit(f"{prefix}.{child_key}" if prefix else str(child_key), child_value) + + for k, v in mapping.items(): + visit(str(k), v) + return flat + + +def normalize_key(key: str) -> str: + return re.sub(r"[^a-z0-9_./-]+", "_", key.strip().lower()) + + +def try_json(value: str) -> Any: + stripped = value.strip() + if not stripped or stripped[0] not in "[{": + return value + try: + return json.loads(stripped) + except json.JSONDecodeError: + return value + + +def first_float(flat: Mapping[str, Any], keys: Iterable[str]) -> float | None: + for key in keys: + value = value_for_key(flat, key) + parsed = parse_float(value) + if parsed is not None: + return parsed + return None + + +def extract_output_tokens(flat: Mapping[str, Any]) -> int | None: + for key in OUTPUT_TOKEN_KEYS: + parsed = parse_int(value_for_key(flat, key)) + if parsed is not None: + return parsed + + usage = value_for_key(flat, "usage") + if isinstance(usage, Mapping): + parsed = parse_int(usage.get("completion_tokens")) + if parsed is not None: + return parsed + + token_ids = value_for_key(flat, "output_token_ids") + if token_ids is None: + token_ids = value_for_key(flat, "generated_token_ids") + parsed_ids = parse_sequence(token_ids) + if parsed_ids is not None: + return len(parsed_ids) + return None + + +def extract_token_timestamps(flat: Mapping[str, Any]) -> tuple[float, ...]: + for key in TOKEN_TIME_KEYS: + values = parse_sequence(value_for_key(flat, key)) + if values: + parsed = [parse_float(value) for value in values] + timestamps = [ + normalize_timestamp_seconds(value) + for value in parsed + if value is not None + ] + if timestamps: + return tuple(timestamps) + return () + + +def value_for_key(flat: Mapping[str, Any], key: str) -> Any: + normalized = normalize_key(key) + if normalized in flat: + return flat[normalized] + suffix = f".{normalized}" + for candidate_key, value in flat.items(): + if candidate_key.endswith(suffix): + return value + return None + + +def parse_float(value: Any) -> float | None: + if isinstance(value, bool) or value is None: + return None + if isinstance(value, (int, float)): + number = float(value) + return number if math.isfinite(number) else None + if isinstance(value, str): + match = re.search(r"-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?", value.strip()) + if not match: + return None + try: + number = float(match.group(0)) + except ValueError: + return None + return number if math.isfinite(number) else None + return None + + +def parse_int(value: Any) -> int | None: + parsed = parse_float(value) + if parsed is None: + return None + return int(round(parsed)) + + +def parse_sequence(value: Any) -> list[Any] | None: + if value is None: + return None + if isinstance(value, list): + return value + if isinstance(value, tuple): + return list(value) + if isinstance(value, str): + parsed = try_json(value) + if isinstance(parsed, list): + return parsed + numbers = re.findall(r"-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?", value) + if numbers: + return numbers + return None + + +def normalize_timestamp_seconds(value: float) -> float: + absolute = abs(value) + if absolute > 1e17: + return value / 1e9 + if absolute > 1e14: + return value / 1e6 + if absolute > 1e11: + return value / 1e3 + return value + + +def normalize_duration_seconds(value: float) -> float: + if value > 10000: + return value / 1000 + return value + + +def dedupe_records(records: list[RequestRecord]) -> list[RequestRecord]: + seen: set[tuple[float, float, int]] = set() + deduped: list[RequestRecord] = [] + for record in records: + key = ( + round(record.start_time, 6), + round(record.end_time, 6), + record.output_tokens, + ) + if key in seen: + continue + seen.add(key) + deduped.append(record) + return sorted(deduped, key=lambda record: (record.start_time, record.end_time)) + + +def build_throughput_rows( + records: list[RequestRecord], + *, + die_count: int, +) -> list[dict[str, float]]: + min_time = min(record.start_time for record in records) + max_time = max(record.end_time for record in records) + num_seconds = max(1, int(math.ceil(max_time - min_time))) + buckets = [0.0 for _ in range(num_seconds)] + + for record in records: + if len(record.token_timestamps) >= record.output_tokens: + for timestamp in record.token_timestamps[-record.output_tokens :]: + idx = int(math.floor(timestamp - min_time)) + if 0 <= idx < len(buckets): + buckets[idx] += 1.0 + continue + + duration = max(record.end_time - record.start_time, 1e-6) + rate = record.output_tokens / duration + start_idx = int(math.floor(record.start_time - min_time)) + end_idx = int(math.ceil(record.end_time - min_time)) + for idx in range(max(0, start_idx), min(len(buckets), end_idx)): + bucket_start = min_time + idx + bucket_end = bucket_start + 1 + overlap = max( + 0.0, + min(record.end_time, bucket_end) - max(record.start_time, bucket_start), + ) + buckets[idx] += rate * overlap + + rows: list[dict[str, float]] = [] + for idx, tokens in enumerate(buckets): + rows.append( + { + "second": float(idx), + "elapsed_s": float(idx), + "tokens/s": tokens, + "tokens/s/die": tokens / die_count, + } + ) + return rows + + +def write_csv(path: Path, rows: list[dict[str, float]]) -> None: + with path.open("w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=["second", "elapsed_s", "tokens/s", "tokens/s/die"], + ) + writer.writeheader() + for row in rows: + writer.writerow( + { + "second": int(row["second"]), + "elapsed_s": int(row["elapsed_s"]), + "tokens/s": f"{row['tokens/s']:.6f}", + "tokens/s/die": f"{row['tokens/s/die']:.6f}", + } + ) + + +def plot_png( + path: Path, + rows: list[dict[str, float]], + *, + y_key: str, + y_label: str, + title: str, +) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + x = [row["elapsed_s"] for row in rows] + y = [row[y_key] for row in rows] + fig, ax = plt.subplots(figsize=(12, 5)) + ax.plot(x, y, linewidth=1.6) + ax.set_title(title) + ax.set_xlabel("Elapsed time (s)") + ax.set_ylabel(y_label) + ax.grid(True, alpha=0.25) + fig.tight_layout() + fig.savefig(path, dpi=160) + plt.close(fig) + + +def sanitize_name(name: str) -> str: + sanitized = re.sub(r"[^A-Za-z0-9_.-]+", "_", name.strip()) + return sanitized.strip("_") or "benchmark" + + +if __name__ == "__main__": + main() diff --git a/tools/benchmarks/legacy/throughput_v2.py b/tools/benchmarks/legacy/throughput_v2.py new file mode 100644 index 00000000..e394d369 --- /dev/null +++ b/tools/benchmarks/legacy/throughput_v2.py @@ -0,0 +1,821 @@ +#!/usr/bin/env python3 +"""Build per-die throughput CSV/PNG files from AISBench result dirs.""" + +from __future__ import annotations + +import argparse +import csv +import io +import json +import math +import re +import sqlite3 +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +DEFAULT_SEARCH_ROOTS = ( + "outputs/default", + "outputs", + "benchmark/outputs/default", + "benchmark/outputs", + "/a3_inference/itask/workdir/wb02363348/cyj_afd/code/outputs/default", + "/a3_inference/itask/workdir/wb02363348/cyj_afd/code/benchmark/outputs/default", +) + +TIME_KEYS = ( + "end_time", + "finish_time", + "finished_time", + "response_time", + "completion_time", + "request_end_time", + "request_finished_time", + "timestamp_end", + "end_timestamp", +) +START_TIME_KEYS = ( + "start_time", + "request_start_time", + "send_time", + "sent_time", + "begin_time", + "timestamp_start", + "start_timestamp", +) +LATENCY_KEYS = ( + "latency", + "total_latency", + "request_latency", + "e2e_latency", + "duration", +) +OUTPUT_TOKEN_KEYS = ( + "output_tokens", + "output_token", + "output_token_num", + "output_token_len", + "output_token_length", + "output_len", + "output_length", + "generated_tokens", + "generated_token_num", + "generated_token_len", + "completion_tokens", + "num_output_tokens", + "decode_tokens", +) +TOKEN_TIME_KEYS = ( + "token_timestamps", + "output_token_timestamps", + "decode_token_timestamps", + "generated_token_timestamps", +) + + +@dataclass(frozen=True) +class RequestRecord: + start_time: float + end_time: float + output_tokens: int + token_timestamps: tuple[float, ...] = () + + +def main() -> None: + args = parse_args() + if args.result_file is not None: + result_file = args.result_file + records = load_records(result_file) + else: + ais_bench_dir = resolve_ais_bench_dir(args.ais_bench_dir, args.input_root) + result_file = None + records = [] + for candidate in find_candidate_results(ais_bench_dir): + candidate_records = load_records(candidate) + if candidate_records: + result_file = candidate + records = candidate_records + break + + if result_file is None: + raise SystemExit( + "No AISBench detail/db/json/csv result with request records found " + f"under: {ais_bench_dir}" + ) + if not records: + raise SystemExit(f"No request records with time and output tokens in {result_file}") + + file_name = sanitize_name(args.file_name) + label = args.label.strip() + label_name = sanitize_name(label) + output_dir = Path(args.output_dir) + csv_path = output_dir / f"throughput_{file_name}_{label_name}.csv" + png_path = output_dir / f"throughput_{file_name}.png" + + rows = build_throughput_rows(records, die_count=args.die_count, label=label) + output_dir.mkdir(parents=True, exist_ok=True) + write_csv(csv_path, rows) + plot_per_die_png( + png_path, + load_plot_series(output_dir, file_name), + title=args.title, + ) + + total_tokens = sum(record.output_tokens for record in records) + duration = rows[-1]["elapsed_s"] + 1 if rows else 0 + peak = max(row["tokens/s/die"] for row in rows) if rows else 0.0 + print(f"input: {result_file}") + print(f"requests: {len(records)}") + print(f"output_tokens: {total_tokens}") + print(f"duration_s: {duration:.0f}") + print(f"peak_tokens/s/die: {peak:.6f}") + print(f"csv: {csv_path}") + print(f"png: {png_path}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Generate throughput__