diff --git a/DeepSeek-V4-Flash/install.sh b/DeepSeek-V4-Flash/install.sh index 9acf902..0af47f3 100644 --- a/DeepSeek-V4-Flash/install.sh +++ b/DeepSeek-V4-Flash/install.sh @@ -64,4 +64,10 @@ git apply --whitespace=nowarn ../verl-ascend-recipe/DeepSeek-V4-Flash/patch/vllm cd verl git apply --whitespace=nowarn ../verl-ascend-recipe/DeepSeek-V4-Flash/patch/verl.patch && cd .. +cd vllm +git apply --whitespace=nowarn ../verl-ascend-recipe/DeepSeek-V4-Flash/patch/vllm.patch && cd .. + +cd MindSpeed-LLM +git apply --whitespace=nowarn ../verl-ascend-recipe/DeepSeek-V4-Flash/patch/mindspeed-llm.patch && cd .. + diff --git a/DeepSeek-V4-Flash/patch/megatron.patch b/DeepSeek-V4-Flash/patch/megatron.patch index 150905a..f8131bf 100644 --- a/DeepSeek-V4-Flash/patch/megatron.patch +++ b/DeepSeek-V4-Flash/patch/megatron.patch @@ -30,8 +30,155 @@ index 26a96f457..f8cbc5fd4 100644 if recv_next_shape_tensor is not None: recv_next_shape = recv_next_shape_tensor.tolist() +diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py +index dafedfd17..8cc172886 100644 +--- a/megatron/core/transformer/moe/moe_utils.py ++++ b/megatron/core/transformer/moe/moe_utils.py +@@ -1,7 +1,7 @@ + # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + + import math +-from typing import List, Optional, Union ++from typing import List, Optional, Union, Tuple + + import torch + +@@ -15,10 +15,12 @@ try: + fused_sort_chunks_by_index, + fused_sort_chunks_by_index_with_probs, + fused_unpermute, ++ fused_compute_score_for_moe_aux_loss, + ) + + HAVE_TE = True + except ImportError: ++ fused_compute_score_for_moe_aux_loss = None + HAVE_TE = False + + +@@ -759,3 +761,119 @@ def maybe_move_tensor_to_cpu(tensor, as_numpy=False, record_stream=False): + tensor.record_stream(torch.cuda.current_stream()) + tensor = cpu_tensor + return tensor ++ ++ ++def apply_router_token_dropping( ++ routing_probs: torch.Tensor, ++ routing_map: torch.Tensor, ++ router_topk: int, ++ capacity_factor: float, ++ drop_policy: str = "probs", ++ pad_to_capacity: bool = False, ++) -> Tuple[torch.Tensor, torch.Tensor]: ++ """Apply token dropping to top-k expert selection. ++ ++ This function enforces expert capacity limits by dropping tokens that exceed ++ the capacity and optionally padding to capacity. ++ ++ Args: ++ routing_probs (torch.Tensor): Tensor of shape [num_tokens, num_experts] ++ containing the routing probabilities for selected experts. ++ routing_map (torch.Tensor): Boolean tensor of shape [num_tokens, num_experts] ++ indicating which experts were selected for each token. ++ router_topk (int): Number of experts selected per token. ++ capacity_factor (float): The capacity factor of each expert. ++ drop_policy (str, optional): Policy to drop tokens - "probs" or "position". ++ Defaults to "probs". ++ pad_to_capacity (bool, optional): Whether to pad to capacity. Defaults to False. ++ ++ Returns: ++ Tuple[torch.Tensor, torch.Tensor]: ++ - final_probs: Routing probabilities after applying capacity constraints ++ - final_map: Boolean mask after applying capacity constraints ++ """ ++ assert routing_probs.ndim == 2 and routing_map.ndim == 2 ++ num_tokens, num_experts = routing_probs.shape ++ # Calculate expert capacity ++ expert_capacity = get_capacity( ++ num_tokens=num_tokens * router_topk, ++ num_experts=num_experts, ++ capacity_factor=capacity_factor, ++ ) ++ ++ # Create capacity mask based on drop policy ++ if expert_capacity > num_tokens: ++ # No need to drop tokens if capacity exceeds the number of tokens ++ capacity_mask = torch.ones_like(routing_probs).bool() ++ else: ++ if drop_policy == "probs": ++ _, capacity_indices = torch.topk(routing_probs, k=expert_capacity, dim=0, sorted=False) ++ capacity_mask = torch.zeros_like(routing_probs).scatter(0, capacity_indices, 1).bool() ++ elif drop_policy == "position": ++ _, capacity_indices = torch.topk( ++ routing_map.int(), k=expert_capacity, dim=0, sorted=False ++ ) ++ capacity_mask = torch.zeros_like(routing_probs).scatter(0, capacity_indices, 1).bool() ++ else: ++ raise ValueError(f"Invalid drop_policy: {drop_policy}") ++ ++ # Apply capacity constraints ++ if pad_to_capacity: ++ final_map = capacity_mask ++ final_probs = routing_probs * final_map ++ else: ++ # Get exceed mask and maskout exceeded probs and indices ++ final_map = torch.logical_and(routing_map, capacity_mask) ++ final_probs = routing_probs * final_map ++ ++ return final_probs, final_map ++ ++ ++def compute_routing_scores_for_aux_loss( ++ logits: torch.Tensor, ++ topk: int, ++ score_function: str, ++ fused: bool = False, ++ padding_mask: Optional[torch.Tensor] = None, ++) -> Tuple[torch.Tensor, torch.Tensor]: ++ """Compute routing scores based on the score function. ++ ++ Args: ++ logits (torch.Tensor): The logits tensor after gating, shape: [num_tokens, num_experts]. ++ topk (int): The number of top-k indices to compute. ++ score_function (str): The score function to use. Can be either "softmax" or "sigmoid". ++ fused (bool, optional): Whether to use the fused version. Defaults to False. ++ padding_mask (torch.Tensor, optional): Boolean mask indicating non-padding tokens. ++ Shape in [num_tokens]. True for valid tokens, ++ False for padding tokens. Defaults to None. ++ ++ Returns: ++ Tuple[torch.Tensor, torch.Tensor]: The routing map and the normalized routing scores. ++ """ ++ if fused: ++ if not HAVE_TE or fused_compute_score_for_moe_aux_loss is None: ++ raise ValueError( ++ "fused_compute_score_for_moe_aux_loss is not available. Please install TE >= 2.6.0." ++ ) ++ routing_map, scores = fused_compute_score_for_moe_aux_loss( ++ logits=logits, topk=topk, score_function=score_function ++ ) ++ else: ++ if score_function == "softmax": ++ scores = torch.softmax(logits, dim=-1, dtype=torch.float32) ++ elif score_function == "sigmoid": ++ scores = torch.sigmoid(logits) ++ scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) ++ else: ++ raise ValueError(f"Invalid score_function: {score_function}") ++ ++ _, top_indices = torch.topk(scores, k=topk, dim=1) ++ routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool() ++ ++ # Apply padding mask to scores if provided ++ if padding_mask is not None: ++ # Invert padding_mask and make True indicates valid tokens ++ valid_mask = (~padding_mask).unsqueeze(-1) ++ routing_map = routing_map * valid_mask ++ scores = scores * valid_mask ++ return routing_map, scores +\ No newline at end of file diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py -index 2b68a5a5f..e6cf2469b 100644 +index 2b68a5a5f..20d624610 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -568,8 +568,8 @@ class TransformerConfig(ModelParallelConfig): @@ -45,3 +192,11 @@ index 2b68a5a5f..e6cf2469b 100644 if self.num_query_groups % self.tensor_model_parallel_size != 0: raise ValueError( +@@ -1075,6 +1075,7 @@ class MLATransformerConfig(TransformerConfig): + The initialization function has an argument for each parameter, including those in + ModelParallelConfig. Included YaRN RoPE parameters that is fused in MLA. + """ ++ enable_routing_replay: bool = False + + multi_latent_attention: bool = True + """Whether to use Multi-Latent Attention.""" diff --git a/DeepSeek-V4-Flash/patch/mindspeed-llm.patch b/DeepSeek-V4-Flash/patch/mindspeed-llm.patch new file mode 100644 index 0000000..a4c0000 --- /dev/null +++ b/DeepSeek-V4-Flash/patch/mindspeed-llm.patch @@ -0,0 +1,95 @@ +diff --git a/mindspeed_llm/core/transformer/moe/moe_utils.py b/mindspeed_llm/core/transformer/moe/moe_utils.py +index d20e40a0..90816b98 100644 +--- a/mindspeed_llm/core/transformer/moe/moe_utils.py ++++ b/mindspeed_llm/core/transformer/moe/moe_utils.py +@@ -18,6 +18,7 @@ from functools import wraps + from typing import Optional + + import torch ++from enum import Enum + import torch.nn.functional as F + from megatron.core import parallel_state + from megatron.core.transformer.moe.moe_utils import get_capacity +@@ -108,6 +109,7 @@ def topk_softmax_with_capacity_and_hash( + token_hash: bool = False, + tid2eid: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, ++ router_replay = None, + ): + """ + patch hash operator in megatron topk_softmax_with_capacity +@@ -115,7 +117,7 @@ def topk_softmax_with_capacity_and_hash( + assert logits.dim() == 2, f"Expected 2D logits [num_tokens, num_experts], got {logits.dim()}." + num_tokens, num_experts = logits.shape + +- def compute_topk(scores, topk, num_groups=None, group_topk=None): ++ def _compute_topk(scores, topk, num_groups=None, group_topk=None): + if group_topk: + return group_limited_topk( + scores=scores, +@@ -128,6 +130,41 @@ def topk_softmax_with_capacity_and_hash( + else: + return torch.topk(scores, k=topk, dim=1) + ++ def compute_topk(scores, topk, num_groups=None, group_topk=None): ++ from verl.utils.megatron.router_replay_patch import RouterReplayAction ++ # Default behavior if no replay is active ++ routing_action = router_replay.router_replay_action if router_replay is not None else None ++ ++ if routing_action is None: ++ return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) ++ ++ if routing_action == RouterReplayAction.RECORD: ++ probs, top_indices = _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) ++ if router_replay is not None: ++ router_replay.record_indices(top_indices) ++ return probs, top_indices ++ elif routing_action == RouterReplayAction.REPLAY_FORWARD: ++ if router_replay is None or router_replay.target_topk_idx is None: ++ # Fallback if replay data is not available ++ return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) ++ # Use the provided indices for replay ++ top_indices = router_replay.target_topk_idx ++ top_indices = top_indices.to(scores.device) ++ probs = scores.gather(1, top_indices) ++ return probs, top_indices ++ elif routing_action == RouterReplayAction.REPLAY_BACKWARD: ++ if router_replay is None or not router_replay.replay_backward_list: ++ # Fallback if replay data is not available ++ return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) ++ # Use the last recorded indices for backward replay ++ top_indices = router_replay.replay_backward_list.pop(0) ++ # Ensure indices are on the correct device ++ top_indices = top_indices.to(scores.device) ++ probs = scores.gather(1, top_indices) ++ return probs, top_indices ++ else: # Unknown action, fallback ++ return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk) ++ + if score_function == "softmax": + if use_pre_softmax: + scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) +diff --git a/mindspeed_llm/core/transformer/moe/router.py b/mindspeed_llm/core/transformer/moe/router.py +index 5555ca2b..e0948d3c 100644 +--- a/mindspeed_llm/core/transformer/moe/router.py ++++ b/mindspeed_llm/core/transformer/moe/router.py +@@ -599,6 +599,7 @@ def topk_router_routing(self, logits: torch.Tensor, input_ids: torch.Tensor = No + token_hash=self.hash if hasattr(self, "hash") else None, + tid2eid=self.tid2eid if hasattr(self, "tid2eid") else None, + input_ids=input_ids, ++ router_replay=getattr(self, "router_replay", None) + ) + else: + # A naive top-k routing without load balancing +diff --git a/pretrain_deepseek4.py b/pretrain_deepseek4.py +index 8e7ca23a..64d88cb4 100644 +--- a/pretrain_deepseek4.py ++++ b/pretrain_deepseek4.py +@@ -54,6 +54,7 @@ def model_provider( + Union[DeepSeek4Model, megatron.legacy.model.DeepSeek4Model]: The returned model + """ + args = get_args() ++ args.enable_routing_replay = model_provider.enable_routing_replay + use_te = args.transformer_impl == "transformer_engine" + + print_rank_0('building GPT model ...') diff --git a/DeepSeek-V4-Flash/patch/verl.patch b/DeepSeek-V4-Flash/patch/verl.patch index 847df07..1bedf7c 100644 --- a/DeepSeek-V4-Flash/patch/verl.patch +++ b/DeepSeek-V4-Flash/patch/verl.patch @@ -61,8 +61,154 @@ index af1e9cab..04439196 100644 +llm_kwargs: {} mm_kwargs: {} +diff --git a/verl/utils/megatron/router_replay_patch.py b/verl/utils/megatron/router_replay_patch.py +index 012ab3ff..f5c65a88 100644 +--- a/verl/utils/megatron/router_replay_patch.py ++++ b/verl/utils/megatron/router_replay_patch.py +@@ -216,6 +216,16 @@ def _patched_topk_routing_with_score_function( + else: + scores, top_indices = compute_topk(scores, topk, num_groups, group_topk) + probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores ++ elif score_function == "sqrtsoftplus": ++ scores = torch.nn.functional.softplus(logits.float()).sqrt().type_as(logits) ++ if expert_bias is not None: ++ scores_for_routing = scores + expert_bias ++ _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk) ++ scores = torch.gather(scores, dim=1, index=top_indices) ++ else: ++ _, top_indices = compute_topk(scores, topk, num_groups, group_topk) ++ scores = torch.gather(scores, dim=1, index=top_indices) ++ probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) + else: + raise ValueError(f"Invalid score_function: {score_function}") + +@@ -343,71 +353,17 @@ def apply_router_replay_patch(): + print("Applying Router Replay Patch...") + # Clear router instances to avoid state leakage between model initializations. + RouterReplay.router_instances.clear() +- # Step 1: Patch TransformerConfig to include the feature flag +- +- try: +- sig = inspect.signature(TransformerConfig.__init__) +- native_params = sig.parameters +- params = list(sig.parameters.values()) +- except Exception: +- sig = None +- native_params = {} +- params = [] +- +- ext_attrs = ["enable_routing_replay"] +- +- # Update __signature__ to prevent NPU/MindSpeed wrappers from filtering out or blocking custom parameters. +- for attr in ext_attrs: +- if attr not in native_params: +- if sig: +- new_param = inspect.Parameter(attr, inspect.Parameter.KEYWORD_ONLY, default=False) +- if params and params[-1].kind == inspect.Parameter.VAR_KEYWORD: +- params.insert(-1, new_param) +- else: +- params.append(new_param) +- +- if sig: +- try: +- TransformerConfig.__init__.__signature__ = sig.replace(parameters=params) +- except Exception as e: +- print(f"Failed to update signature metadata: {e}") +- +- if not hasattr(TransformerConfig, "_verl_router_patched"): +- # Store original __init__ method +- original_tf_config_init = TransformerConfig.__init__ +- +- # Define new __init__ method that safely handles enable_routing_replay parameter +- @wraps(original_tf_config_init) +- def patched_tf_config_init(self, *args, **kwargs): +- # Simple solution: remove the unknown parameter before calling original constructor +- enable_routing_replay = kwargs.get("enable_routing_replay", False) +- if "enable_routing_replay" not in native_params: +- enable_routing_replay = kwargs.pop("enable_routing_replay", False) +- +- # Call original constructor with remaining kwargs +- original_tf_config_init(self, *args, **kwargs) +- +- # Set the instance attribute +- self.enable_routing_replay = enable_routing_replay +- +- # Apply the patch +- TransformerConfig.__init__ = patched_tf_config_init +- TransformerConfig._verl_router_patched = True +- +- # Step 2: Patch TopKRouter only once to ensure idempotency. +- if hasattr(TopKRouter, "_router_replay_patched"): +- return + + original_init = TopKRouter.__init__ + +- # Step 3: Define the new __init__ method ++ # Define the new __init__ method + def patched_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + self.router_replay = None + if getattr(self.config, "enable_routing_replay", False): + self.router_replay = RouterReplay() + +- # Step 4: Patch MoEAlltoAllTokenDispatcher.preprocess to handle router replay ++ # Patch MoEAlltoAllTokenDispatcher.preprocess to handle router replay + # When router replay is enabled, duplicate indices in top_indices can cause + # routing_map.sum() < num_tokens * topk, leading to split size mismatch in alltoall. + if MoEAlltoAllTokenDispatcher is not None and not hasattr(MoEAlltoAllTokenDispatcher, "_preprocess_patched"): +@@ -437,7 +393,5 @@ def apply_router_replay_patch(): + MoEAlltoAllTokenDispatcher.preprocess = patched_preprocess + MoEAlltoAllTokenDispatcher._preprocess_patched = True + +- # Step 5: Apply the patches ++ # Apply the patches + TopKRouter.__init__ = patched_init +- TopKRouter.routing = patched_routing +- TopKRouter._router_replay_patched = True +diff --git a/verl/utils/megatron/router_replay_utils.py b/verl/utils/megatron/router_replay_utils.py +index 9984af10..fe2adb36 100644 +--- a/verl/utils/megatron/router_replay_utils.py ++++ b/verl/utils/megatron/router_replay_utils.py +@@ -292,6 +292,23 @@ def set_router_replay_data(layers_topk_idx, attention_mask, tf_config, vp_rank=N + layers_topk_idx_rmpad, _ = preprocess_packed_seqs(layers_topk_idx, attention_mask, pre_process=True) + layers_topk_idx_rmpad = layers_topk_idx_rmpad.contiguous() # 1, dynamic_bs_all, layer_num, topk + ++ ## 同preprocess_bshd_engine中的padding ++ max_seqlen = layers_topk_idx_rmpad.shape[1] ++ align_value = 2048 ++ if max_seqlen % align_value != 0: ++ max_seqlen = ((max_seqlen + align_value - 1) // align_value) * align_value ++ target_len = max(max_seqlen, 2048) ++ ++ ## padding ++ pad_len = target_len - layers_topk_idx_rmpad.shape[1] ++ pad_values = torch.tensor( ++ [0, 42, 84, 126, 168, 210], ++ dtype=layers_topk_idx_rmpad.dtype, ++ device=layers_topk_idx_rmpad.device ++ ) ++ pad_block = pad_values.unsqueeze(0).unsqueeze(0).repeat(1, pad_len, layers_topk_idx_rmpad.shape[2], 1) ++ layers_topk_idx_rmpad = torch.cat([layers_topk_idx_rmpad, pad_block], dim=1) ++ + # 1, dynamic_bs_split, layer_num, topk + layers_topk_idx_rmpad_split = scatter_to_sequence_parallel_region( + layers_topk_idx_rmpad.to(device_name).squeeze(dim=0) +diff --git a/verl/utils/reward_score/math_dapo.py b/verl/utils/reward_score/math_dapo.py +index 940500fd..e30844bf 100644 +--- a/verl/utils/reward_score/math_dapo.py ++++ b/verl/utils/reward_score/math_dapo.py +@@ -264,6 +264,9 @@ def compute_score( + + reward = 1.0 if correct else -1.0 + acc = correct ++ ++ import random ++ reward = random.randrange(0, 2) * 0.2 + + return { + "score": reward, diff --git a/verl/utils/vllm/npu_vllm_patch.py b/verl/utils/vllm/npu_vllm_patch.py -index c22647ce..627f4c20 100644 +index c22647ce..972886b4 100644 --- a/verl/utils/vllm/npu_vllm_patch.py +++ b/verl/utils/vllm/npu_vllm_patch.py @@ -193,13 +193,16 @@ if is_torch_npu_available(check_device=False): @@ -99,10 +245,10 @@ index 951c5cad..7635039c 100644 # TODO(@leisuzz): class Qwen3MoeLLMForCausalLM is not available if VLLM version < 0.11.0, # will update the 'if statement' with 'isinstance' when verl commonly use VLLM version >= 0.11.0 diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py -index cbb7be48..9a00c96e 100644 +index cbb7be48..0296edce 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py -@@ -677,30 +677,10 @@ class MegatronEngine(BaseEngine): +@@ -677,30 +677,9 @@ class MegatronEngine(BaseEngine): losses_reduced[0]["metrics"] = {} losses_reduced[0]["metrics"].update(metrics) @@ -124,7 +270,7 @@ index cbb7be48..9a00c96e 100644 - use_dynamic_bsz = tu.get_non_tensor_data(data=data, key="use_dynamic_bsz", default=True) - if use_dynamic_bsz and indices is not None: - layers_topk_idx = restore_dynamic_batch(layers_topk_idx, indices) - +- output = {} if mpu.is_pipeline_last_stage(ignore_virtual=True): output = postprocess_batch_func(output_lst=losses_reduced, indices=indices, data=data) @@ -133,23 +279,18 @@ index cbb7be48..9a00c96e 100644 if enable_routing_replay: RouterReplay.clear_global_indices() RouterReplay.clear_global_router_replay_action() -@@ -833,15 +813,6 @@ class MegatronEngineWithLMHead(MegatronEngine): - else: - vp_rank = 0 +@@ -838,8 +817,9 @@ class MegatronEngineWithLMHead(MegatronEngine): + for router in router_instance_list: + router.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD) -- if RouterReplayHelper.is_replay_backward_action(self.tf_config, vp_rank): -- router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank) -- for router in router_instance_list: -- router.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD) -- - if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank): -- layers_topk_idx = model_inputs["routed_experts"] -- set_router_replay_data(layers_topk_idx, None, self.tf_config, vp_rank) -- ++ if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank): + layers_topk_idx = model_inputs["routed_experts"] ++ + set_router_replay_data(layers_topk_idx, None, self.tf_config, vp_rank) + if pad_mode == DatasetPadMode.NO_PADDING: - label = input_ids.clone() - else: -@@ -932,15 +903,6 @@ class MegatronEngineWithLMHead(MegatronEngine): +@@ -932,16 +912,12 @@ class MegatronEngineWithLMHead(MegatronEngine): local_cp_size=local_cp_size, ) @@ -158,24 +299,30 @@ index cbb7be48..9a00c96e 100644 - merge_router_topk_indices(None, input_ids, self.mini_layer_topk_idx_list, self.tf_config, vp_rank) - - # Router replay: switch to backward replay mode for next backward pass -- if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank): -- router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank) -- for router in router_instance_list: -- router.set_router_replay_action(RouterReplayAction.REPLAY_BACKWARD) + if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank): + router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank) + for router in router_instance_list: + router.set_router_replay_action(RouterReplayAction.REPLAY_BACKWARD) ++ return output, partial(postprocess_micro_batch_func, data=batch, local_cp_size=local_cp_size) + def postprocess_micro_batch_func( diff --git a/verl/workers/engine/mindspeed/transformer_impl.py b/verl/workers/engine/mindspeed/transformer_impl.py -index acbdcf63..1be56d30 100644 +index acbdcf63..291153a6 100644 --- a/verl/workers/engine/mindspeed/transformer_impl.py +++ b/verl/workers/engine/mindspeed/transformer_impl.py -@@ -93,15 +93,14 @@ class MindSpeedLLMEngineWithLMHead(MegatronEngineWithLMHead): +@@ -93,15 +93,18 @@ class MindSpeedLLMEngineWithLMHead(MegatronEngineWithLMHead): import torch.distributed from megatron.core.enums import ModelType from megatron.training.training import get_model - + from mindspeed_llm.pretrain_deepseek4 import model_provider + gpt_model_provider=model_provider ++ if self.tf_config.enable_routing_replay: ++ gpt_model_provider.enable_routing_replay = True ++ else: ++ gpt_model_provider.enable_routing_replay = False # For forward_only, we don't need optimizer, lr_scheduler, checkpoint_mananager if self.engine_config.forward_only: - module = get_model(gpt_model_provider, ModelType.encoder_or_decoder, wrap_with_ddp=False) @@ -208,6 +355,25 @@ index 9b600e30..ecb9de4c 100644 mcore_config = add_mcore_arguments(all_config) from mindspeed_llm.tasks.megatron_adaptor_v2 import repatch +diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py +index ee741cba..e0c0ea87 100644 +--- a/verl/workers/engine_workers.py ++++ b/verl/workers/engine_workers.py +@@ -473,10 +473,13 @@ class ActorRolloutRefWorker(Worker, DistProfilerExtension): + else: + tool_config = None + ++ strategy = self.config.actor.strategy + self.enable_routing_replay = ( +- self.config.actor.strategy == "megatron" and self.config.actor.megatron.router_replay.mode != "disabled" ++ strategy in ("megatron", "mindspeed") ++ and getattr(self.config.actor, strategy).router_replay.mode != "disabled" + ) + ++ + DistProfilerExtension.__init__( + self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config) + ) diff --git a/verl/workers/rollout/vllm_rollout/bucketed_weight_transfer.py b/verl/workers/rollout/vllm_rollout/bucketed_weight_transfer.py index eb337939..81ac0ea1 100644 --- a/verl/workers/rollout/vllm_rollout/bucketed_weight_transfer.py @@ -225,7 +391,7 @@ index eb337939..81ac0ea1 100644 "name": name, "shape": weight.shape, diff --git a/verl/workers/rollout/vllm_rollout/vllm_async_server.py b/verl/workers/rollout/vllm_rollout/vllm_async_server.py -index 23c8594f..25aabbc1 100644 +index 23c8594f..1942e128 100644 --- a/verl/workers/rollout/vllm_rollout/vllm_async_server.py +++ b/verl/workers/rollout/vllm_rollout/vllm_async_server.py @@ -53,7 +53,10 @@ from verl.workers.rollout.vllm_rollout.utils import ( @@ -265,3 +431,16 @@ index 23c8594f..25aabbc1 100644 sleep_level = 1 else: sleep_level = 2 +diff --git a/verl/workers/rollout/vllm_rollout/vllm_rollout.py b/verl/workers/rollout/vllm_rollout/vllm_rollout.py +index 90641b57..62e756fd 100644 +--- a/verl/workers/rollout/vllm_rollout/vllm_rollout.py ++++ b/verl/workers/rollout/vllm_rollout/vllm_rollout.py +@@ -97,7 +97,7 @@ class ServerAdapter(BaseRollout): + self.device_uuid = get_device_uuid(get_device_id()) + self.zmq_handle = f"ipc:///tmp/rl-colocate-zmq-{self.device_uuid}.sock" + +- self.use_shm = not is_support_ipc() ++ self.use_shm = False #t is_support_ipc() + if self.use_shm: + logger.warning( + "IPC is not supported on your devices. Falling back to shared memory for weight transfer, " diff --git a/DeepSeek-V4-Flash/patch/vllm-ascend.patch b/DeepSeek-V4-Flash/patch/vllm-ascend.patch index 36f9c66..8fa00af 100644 --- a/DeepSeek-V4-Flash/patch/vllm-ascend.patch +++ b/DeepSeek-V4-Flash/patch/vllm-ascend.patch @@ -1,17 +1,46 @@ +diff --git a/vllm_ascend/core/single_type_kv_cache_manager.py b/vllm_ascend/core/single_type_kv_cache_manager.py +index bc184c913..5553bbca8 100644 +--- a/vllm_ascend/core/single_type_kv_cache_manager.py ++++ b/vllm_ascend/core/single_type_kv_cache_manager.py +@@ -46,9 +46,9 @@ class CompressAttentionManager(FullAttentionManager): + # speculative decoding (MTP/EAGLE) with linear attention. + # assert isinstance(self.kv_cache_spec, (CompressAttentionSpec, C4IndexerSpec)) + +- num_tokens //= self.compress_ratio +- num_tokens_main_model //= self.compress_ratio +- ++ num_tokens = cdiv(num_tokens, self.compress_ratio) ++ num_tokens_main_model = cdiv(num_tokens_main_model, self.compress_ratio) ++ + return super().get_num_blocks_to_allocate( + request_id, + num_tokens, +@@ -141,9 +141,9 @@ class CompressAttentionManager(FullAttentionManager): + Returns: + The new allocated blocks. + """ +- num_tokens //= self.compress_ratio ++ num_tokens = cdiv(num_tokens, self.compress_ratio) + ## TODO: check spec decode +- num_tokens_main_model //= self.compress_ratio ++ num_tokens_main_model = cdiv(num_tokens_main_model, self.compress_ratio) + + req_blocks = self.req_to_blocks[request_id] + num_required_blocks = cdiv(num_tokens, self.block_size) diff --git a/vllm_ascend/device/device_op.py b/vllm_ascend/device/device_op.py -index abcdef12..34567890 100644 +index f4b49da74..220120e9f 100644 --- a/vllm_ascend/device/device_op.py +++ b/vllm_ascend/device/device_op.py -@@ -768,7 +768,7 @@ +@@ -841,7 +841,7 @@ class BaseDeviceAdaptor: @staticmethod def dsa_kv_compress_scatter(cache, x, slot_mapping): """Scatter KV into cache. Non-A5: simple scatter of pre-quantized tensor.""" - torch.ops._C_ascend.npu_scatter_nd_update_v2(cache, slot_mapping, x) + torch_npu.npu_scatter_nd_update_(cache, slot_mapping, x) - + # ===== Indexer Quant + Scatter ===== - -@@ -794,8 +794,8 @@ + +@@ -867,8 +867,8 @@ class BaseDeviceAdaptor: kv_scale_out = kv_scale_out.unsqueeze(-1).to(torch.float16) if kv_scale_out.ndim < 4: kv_scale_out = kv_scale_out.unsqueeze(-1) @@ -19,28 +48,28 @@ index abcdef12..34567890 100644 - torch.ops._C_ascend.npu_scatter_nd_update_v2(indexer_scale_cache, slot_mapping, kv_scale_out) + torch_npu.npu_scatter_nd_update_(indexer_k_cache, slot_mapping, kv_out) + torch_npu.npu_scatter_nd_update_(indexer_scale_cache, slot_mapping, kv_scale_out) - + return q, q_scale, kv_out, kv_scale_out - -@@ -808,7 +808,7 @@ + +@@ -881,7 +881,7 @@ class BaseDeviceAdaptor: return None, None kv_out, kv_scale = torch_npu.npu_dynamic_quant(kv, dst_type=torch.int8) kv_scale = kv_scale.unsqueeze(-1) - torch.ops._C_ascend.npu_scatter_nd_update_v2(indexer_k_cache, slot_mapping, kv_out) + torch_npu.npu_scatter_nd_update_(indexer_k_cache, slot_mapping, kv_out) return kv_out, kv_scale - + @staticmethod -@@ -818,7 +818,7 @@ +@@ -891,7 +891,7 @@ class BaseDeviceAdaptor: kv_scale = kv_scale.to(torch.float16) if kv_scale.ndim < 4: kv_scale = kv_scale.unsqueeze(-1) - torch.ops._C_ascend.npu_scatter_nd_update_v2(indexer_scale_cache, slot_mapping, kv_scale) + torch_npu.npu_scatter_nd_update_(indexer_scale_cache, slot_mapping, kv_scale) - + @staticmethod def warmup_indexer_quant_scatter(hidden_states, slot_mapping): -@@ -831,8 +831,8 @@ +@@ -904,8 +904,8 @@ class BaseDeviceAdaptor: dummy_shape = (1, 1, 1, kv_dummy.shape[-1]) indexer_k_cache = torch.zeros(dummy_shape, dtype=kv_dummy.dtype, device=hidden_states.device) indexer_scale_cache = torch.zeros(dummy_shape, dtype=torch.float16, device=hidden_states.device) @@ -48,16 +77,16 @@ index abcdef12..34567890 100644 - torch.ops._C_ascend.npu_scatter_nd_update_v2(indexer_scale_cache, slot_mapping, kv_scale_dummy) + torch_npu.npu_scatter_nd_update_(indexer_k_cache, slot_mapping, kv_dummy) + torch_npu.npu_scatter_nd_update_(indexer_scale_cache, slot_mapping, kv_scale_dummy) - + # ===== Lightning Indexer Dtype Prep ===== - + diff --git a/vllm_ascend/models/deepseek_v4.py b/vllm_ascend/models/deepseek_v4.py old mode 100644 new mode 100755 -index 231f07b8..22c159ae +index b99fd94db..5c26b3eb6 --- a/vllm_ascend/models/deepseek_v4.py +++ b/vllm_ascend/models/deepseek_v4.py -@@ -799,6 +799,7 @@ +@@ -799,6 +799,7 @@ class DeepseekV4Attention(nn.Module): head_size=self.rope_head_dim, rotary_dim=self.rope_head_dim, max_position_embeddings=max_position_embeddings, @@ -65,12 +94,118 @@ index 231f07b8..22c159ae is_neox_style=False, scaling_factor=config.rope_parameters["factor"], base=config.rope_parameters["rope_theta"], - +@@ -1389,14 +1390,18 @@ class AscendDeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV2MixtureOfExpe + if "sink" in name: + if is_pp_missing_parameter(name, self): + continue +- param = params_dict[name] +- if enable_dsa_cp(): +- param.data.copy_(loaded_weight) +- else: +- # Handle attention sinks (distributed across ranks) +- narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank) +- param.data.copy_(narrow_weight) +- loaded_params.add(name) ++ ++ try: ++ param = params_dict[name] ++ if enable_dsa_cp(): ++ param.data.copy_(loaded_weight) ++ else: ++ # Handle attention sinks (distributed across ranks) ++ narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank) ++ param.data.copy_(narrow_weight) ++ loaded_params.add(name) ++ except Exception as e: ++ print(f"skip load {name}: {e}") + continue + + is_fusion_moe_shared_experts_layer = rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) +@@ -1431,9 +1436,12 @@ class AscendDeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV2MixtureOfExpe + if is_pp_missing_parameter(name, self): + continue + +- param = params_dict[name] +- weight_loader = param.weight_loader +- weight_loader(param, loaded_weight, shard_id) ++ try: ++ param = params_dict[name] ++ weight_loader = param.weight_loader ++ weight_loader(param, loaded_weight, shard_id) ++ except Exception as e: ++ print(f"skip in load {name}: {e}") + break + else: + is_expert_weight = False +@@ -1494,25 +1502,28 @@ class AscendDeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV2MixtureOfExpe + if is_pp_missing_parameter(name_mapped, self): + continue + +- param = params_dict[name_mapped] +- # We should ask the weight loader to return success or +- # not here since otherwise we may skip experts with +- # other available replicas. +- weight_loader = typing.cast(Callable[..., bool], param.weight_loader) +- success = weight_loader( +- param, +- weight_to_load, +- name_mapped, +- shard_id=shard_id, +- expert_id=expert_id, +- return_success=True, +- ) +- if success: +- if not is_fusion_moe_shared_experts_layer: +- name = name_mapped +- else: +- loaded_params.add(name_mapped) +- break ++ try: ++ param = params_dict[name_mapped] ++ # We should ask the weight loader to return success or ++ # not here since otherwise we may skip experts with ++ # other available replicas. ++ weight_loader = typing.cast(Callable[..., bool], param.weight_loader) ++ success = weight_loader( ++ param, ++ weight_to_load, ++ name_mapped, ++ shard_id=shard_id, ++ expert_id=expert_id, ++ return_success=True, ++ ) ++ if success: ++ if not is_fusion_moe_shared_experts_layer: ++ name = name_mapped ++ else: ++ loaded_params.add(name_mapped) ++ break ++ except Exception as e: ++ print(f"skip in load {name}: {e}") + else: + if is_expert_weight: + # We've checked that this is an expert weight +@@ -1532,9 +1543,12 @@ class AscendDeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV2MixtureOfExpe + if is_pp_missing_parameter(name, self): + continue + +- param = params_dict[name] +- weight_loader = getattr(param, "weight_loader", default_weight_loader) +- weight_loader(param, loaded_weight) ++ try: ++ param = params_dict[name] ++ weight_loader = getattr(param, "weight_loader", default_weight_loader) ++ weight_loader(param, loaded_weight) ++ except Exception as e: ++ print(f"skip in load {name}: {e}") + if not is_fusion_moe_shared_experts_layer: + loaded_params.add(name) + diff --git a/vllm_ascend/ops/rope_dsv4.py b/vllm_ascend/ops/rope_dsv4.py -index abcdef12..34567890 100644 +index 8727abf4d..aebc0b567 100644 --- a/vllm_ascend/ops/rope_dsv4.py +++ b/vllm_ascend/ops/rope_dsv4.py -@@ -148,6 +148,7 @@ +@@ -148,6 +148,7 @@ class ComplexExpRotaryEmbedding(nn.Module): head_size: int, rotary_dim: int, max_position_embeddings: int, @@ -78,8 +213,8 @@ index abcdef12..34567890 100644 base: int, scaling_factor: float, rope_groups: list[str] | None = None, -@@ -173,7 +174,7 @@ - +@@ -173,7 +174,7 @@ class ComplexExpRotaryEmbedding(nn.Module): + if config_key not in _ROPE_STATE.full_rope_cache: inv_freq = self.precompute_freqs_cis( - rotary_dim, max_position_embeddings, max_position_embeddings, base, scaling_factor, beta_fast, beta_slow @@ -87,13 +222,172 @@ index abcdef12..34567890 100644 ) t = torch.arange( max_position_embeddings * scaling_factor, -@@ -256,6 +257,9 @@ +@@ -256,6 +257,9 @@ class ComplexExpRotaryEmbedding(nn.Module): inv_freq_extrapolation = 1.0 / pos_freqs inv_freq_interpolation = 1.0 / (factor * pos_freqs) - + + if original_seq_len==0: + return inv_freq_extrapolation + low, high = yarn_find_correction_range( beta_fast, beta_slow, +diff --git a/vllm_ascend/patch/worker/__init__.py b/vllm_ascend/patch/worker/__init__.py +index cbc13870b..a40b90a66 100644 +--- a/vllm_ascend/patch/worker/__init__.py ++++ b/vllm_ascend/patch/worker/__init__.py +@@ -85,5 +85,5 @@ if _V2_MODEL_RUNNER_SUPPORTED: + import vllm_ascend.patch.worker.patch_v2.patch_attn_utils # noqa + + # only patch routed experts capture in main2main. +-if _V2_MODEL_RUNNER_SUPPORTED: +- import vllm_ascend.patch.worker.patch_routed_experts_capture # noqa ++#if _V2_MODEL_RUNNER_SUPPORTED: ++import vllm_ascend.patch.worker.patch_routed_experts_capture # noqa +diff --git a/vllm_ascend/worker/block_table.py b/vllm_ascend/worker/block_table.py +index e458e382f..be7c4adf7 100644 +--- a/vllm_ascend/worker/block_table.py ++++ b/vllm_ascend/worker/block_table.py +@@ -50,6 +50,7 @@ class BlockTable: + self.pin_memory = pin_memory + self.device = device + self.physical_block_size = block_size ++ self.compress_ratio = compress_ratio + self.is_mamba_group = ( + kv_cache_group is not None + and hasattr(kv_cache_group, "kv_cache_spec") +@@ -98,6 +99,15 @@ class BlockTable: + self.max_num_batched_tokens + 2 * self.pcp_world_size * self.max_num_reqs, dtype=torch.int32 + ) + ++ # Track the raw (kv-manager view) block ids for compressed groups. ++ # ``self.block_table`` stores logical (kernel) block ids after ++ # hybrid-block expansion, while routed-experts slot reconstruction ++ # needs the un-expanded kv-manager block ids. ++ self.track_raw_block_ids = compress_ratio > 1 ++ if self.track_raw_block_ids: ++ self.raw_block_ids = np.zeros((max_num_reqs, self.max_num_blocks_per_req), dtype=np.int32) ++ self.raw_num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32) ++ + self.kernel_sizes = kernel_sizes + self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size + +@@ -109,6 +119,11 @@ class BlockTable: + if not block_ids: + return + block_ids = np.array(block_ids) ++ if self.track_raw_block_ids: ++ raw_start = self.raw_num_blocks_per_row[row_idx] ++ num_raw = len(block_ids) ++ self.raw_block_ids[row_idx, raw_start : raw_start + num_raw] = block_ids ++ self.raw_num_blocks_per_row[row_idx] += num_raw + if self.use_hybrid_blocks: + block_ids = self._convert_physical_to_logical_blocks(block_ids) + +@@ -120,6 +135,8 @@ class BlockTable: + + def add_row(self, block_ids: list[int], row_idx: int) -> None: + self.num_blocks_per_row[row_idx] = 0 ++ if self.track_raw_block_ids: ++ self.raw_num_blocks_per_row[row_idx] = 0 + self.append_row(block_ids, row_idx) + + def clear_row(self, row_idx: int) -> None: +@@ -127,11 +144,20 @@ class BlockTable: + if num_blocks > 0: + self.block_table.np[row_idx, :num_blocks] = 0 + self.num_blocks_per_row[row_idx] = 0 ++ if self.track_raw_block_ids: ++ raw_blocks = self.raw_num_blocks_per_row[row_idx] ++ if raw_blocks > 0: ++ self.raw_block_ids[row_idx, :raw_blocks] = 0 ++ self.raw_num_blocks_per_row[row_idx] = 0 + + def move_row(self, src: int, tgt: int) -> None: + num_blocks = self.num_blocks_per_row[src] + self.block_table.np[tgt, :num_blocks] = self.block_table.np[src, :num_blocks] + self.num_blocks_per_row[tgt] = num_blocks ++ if self.track_raw_block_ids: ++ raw_blocks = self.raw_num_blocks_per_row[src] ++ self.raw_block_ids[tgt, :raw_blocks] = self.raw_block_ids[src, :raw_blocks] ++ self.raw_num_blocks_per_row[tgt] = raw_blocks + + def swap_row(self, src: int, tgt: int) -> None: + num_blocks_src = self.num_blocks_per_row[src] +@@ -141,6 +167,13 @@ class BlockTable: + + self.block_table.np[[src, tgt]] = self.block_table.np[[tgt, src]] + ++ if self.track_raw_block_ids: ++ raw_src = self.raw_num_blocks_per_row[src] ++ raw_tgt = self.raw_num_blocks_per_row[tgt] ++ self.raw_num_blocks_per_row[src] = raw_tgt ++ self.raw_num_blocks_per_row[tgt] = raw_src ++ self.raw_block_ids[[src, tgt]] = self.raw_block_ids[[tgt, src]] ++ + def compute_slot_mapping( + self, + num_reqs: int, +diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py +index c732a1d74..70817c51d 100644 +--- a/vllm_ascend/worker/model_runner_v1.py ++++ b/vllm_ascend/worker/model_runner_v1.py +@@ -2155,6 +2155,31 @@ class NPUModelRunner(GPUModelRunner): + out=dsa_positions_np, + ) + ++ # Routed experts: rebuild the per-logical-token ++ # slot_mapping from kv-manager block ids. The compressed ++ # (C4) group's attention slot_mapping is expressed on ++ # compressed blocks and cannot index per-token routing ++ # data. Slot layout (must match RoutedExpertsManager on ++ # the scheduler side): ++ # slot = block_id * block_span + pos % block_span ++ # with block_span = block_size * compress_ratio. ++ if (self.model_config.enable_return_routed_experts ++ and self.routed_experts_initialized): ++ c4_block_table = self.input_batch.block_table[0] ++ if c4_block_table.track_raw_block_ids: ++ block_span = (c4_block_table.physical_block_size ++ * c4_block_table.compress_ratio) ++ block_idx = np.clip( ++ dsa_positions_np // block_span, ++ 0, ++ c4_block_table.raw_block_ids.shape[1] - 1, ++ ) ++ re_slot_mapping_np = ( ++ c4_block_table.raw_block_ids[req_indices, block_idx].astype(np.int64) ++ * block_span + dsa_positions_np % block_span) ++ self.routed_experts_slot_mapping_device[:total_num_scheduled_tokens].copy_( ++ torch.from_numpy(re_slot_mapping_np)) ++ + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + +@@ -3078,13 +3103,19 @@ class NPUModelRunner(GPUModelRunner): + ) + if self.model_config.enable_return_routed_experts and kv_cache_gid == 0: + if self.routed_experts_initialized: +- # snapshot slot_mapping into a private device +- # buffer so the next ``_prepare_inputs`` does not +- # overwrite it while D2H is still pending. +- n = slot_mapping.shape[0] +- self.routed_experts_slot_mapping_device[:n].copy_( +- slot_mapping +- ) ++ # For compressed groups the routed-experts slot_mapping ++ # was already rebuilt from kv-manager block ids in ++ # ``_prepare_inputs``; the attention slot_mapping here ++ # is compressed-block based and must not overwrite it. ++ if not (self.use_compress ++ and blk_table.track_raw_block_ids): ++ # snapshot slot_mapping into a private device ++ # buffer so the next ``_prepare_inputs`` does not ++ # overwrite it while D2H is still pending. ++ n = slot_mapping.shape[0] ++ self.routed_experts_slot_mapping_device[:n].copy_( ++ slot_mapping ++ ) + return blk_table_tensor, slot_mapping + + block_table_gid_0, slot_mapping_gid_0 = _get_block_table_and_slot_mapping(0) diff --git a/DeepSeek-V4-Flash/patch/vllm.patch b/DeepSeek-V4-Flash/patch/vllm.patch new file mode 100644 index 0000000..49561d1 --- /dev/null +++ b/DeepSeek-V4-Flash/patch/vllm.patch @@ -0,0 +1,111 @@ +diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +index 115f43a..8cf94a2 100644 +--- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py ++++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +@@ -14,7 +14,11 @@ from vllm.config import VllmConfig + from vllm.distributed.parallel_state import get_tp_group + from vllm.forward_context import get_forward_context + from vllm.platforms import current_platform +-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig ++from vllm.v1.kv_cache_interface import ( ++ FullAttentionSpec, ++ KVCacheConfig, ++ UniformTypeKVCacheSpecs, ++) + + logger = logging.getLogger(__name__) + +@@ -250,19 +254,35 @@ class RoutedExpertsManager: + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + ) -> None: +- # Pick the attention group for block/slot mapping. We require +- # a FullAttentionSpec group rather than any AttentionSpec to +- # stay consistent with the worker-side lookup in +- # ``GPUModelRunner._get_attention_kv_cache_gid``; hybrid models +- # (Mamba / linear attention) also have other AttentionSpec +- # groups whose slot layout differs. ++ # Pick the attention group for block/slot mapping. Unwrap ++ # UniformTypeKVCacheSpecs first: hybrid models (e.g. DeepSeek-V4 ++ # on vllm-ascend) group per-layer specs into UniformTypeKVCacheSpecs, ++ # which is not itself a FullAttentionSpec. ++ def _unwrap(spec): ++ if isinstance(spec, UniformTypeKVCacheSpecs): ++ return next(iter(spec.kv_cache_specs.values())) ++ return spec ++ + self.attn_gid = next( + gid + for gid, g in enumerate(kv_cache_config.kv_cache_groups) +- if isinstance(g.kv_cache_spec, FullAttentionSpec) ++ if isinstance(_unwrap(g.kv_cache_spec), FullAttentionSpec) ++ ) ++ attn_spec = _unwrap( ++ kv_cache_config.kv_cache_groups[self.attn_gid].kv_cache_spec + ) +- attn_group = kv_cache_config.kv_cache_groups[self.attn_gid] +- self.block_size = attn_group.kv_cache_spec.block_size ++ self.block_size = attn_spec.block_size ++ ++ # KV compression (e.g. DeepSeek-V4 sparse MLA on vllm-ascend): ++ # the KV manager allocates one block per ``block_size`` *compressed* ++ # tokens, i.e. one block spans ``block_size * compress_ratio`` ++ # logical tokens. The slot buffer keeps one slot per *logical* ++ # token by expanding the slot space: ++ # slot = block_id * block_span + pos % block_span ++ # For uncompressed models block_span == block_size, reducing to ++ # the original layout. ++ self.compress_ratio = getattr(attn_spec, "compress_ratio", 1) ++ self.block_span = self.block_size * self.compress_ratio + + # All kv_cache_groups share the same physical block pool, so + # block IDs span [0, num_blocks) regardless of how many groups +@@ -271,7 +291,7 @@ class RoutedExpertsManager: + hf_config = vllm_config.model_config.hf_text_config + num_experts = get_num_experts(hf_config) + num_experts_per_tok = _get_num_experts_per_tok(hf_config) +- max_num_slots = kv_cache_config.num_blocks * self.block_size ++ max_num_slots = kv_cache_config.num_blocks * self.block_span + # Expert IDs are 0..num_experts-1; uint8 fits 256 distinct + # values so the boundary is ``<= 256`` (NOT ``< 256``). Keeping + # this narrow matters because the slot buffer is sized for the +@@ -300,9 +320,12 @@ class RoutedExpertsManager: + + Equivalent to ``slot_buffer[slot_mapping] = data``; numpy fancy + indexing handles repeated / out-of-order indices. Called once +- per scheduler step in ``update_from_output``. ++ per scheduler step in ``update_from_output``. Slots with id -1 ++ are padding and must be skipped (numpy would wrap -1 to the ++ last row and silently corrupt it). + """ +- self.routed_experts_by_slot[slot_mapping] = data ++ mask = slot_mapping >= 0 ++ self.routed_experts_by_slot[slot_mapping[mask]] = data[mask] + + def get( + self, +@@ -335,15 +358,16 @@ class RoutedExpertsManager: + Array of shape (num_tokens - token_start, num_layers, + num_experts_per_tok). + """ +- bs = self.block_size +- block_ids_array = np.array(block_ids, dtype=np.int32) +- block_offsets = np.arange(bs) +- # slot = block_id * block_size + offset_in_block; flatten the +- # (num_blocks, block_size) grid and trim to num_tokens, then +- # skip the first token_start entries so only the requested +- # range is fetched in a single fancy-index read. ++ span = self.block_span ++ block_ids_array = np.array(block_ids, dtype=np.int64) ++ block_offsets = np.arange(span) ++ # slot = block_id * block_span + offset_in_block. With KV ++ # compression one block covers block_span logical tokens (the ++ # kv manager allocates blocks on compressed tokens), so this ++ # linear reconstruction stays correct provided ``block_ids`` ++ # are the kv-manager block ids of the group the worker used. + slot_mapping = ( +- block_ids_array.reshape(-1, 1) * bs + block_offsets.reshape(1, -1) ++ block_ids_array.reshape(-1, 1) * span + block_offsets.reshape(1, -1) + ).flatten()[:num_tokens] + slot_mapping = slot_mapping[token_start:] + return self.routed_experts_by_slot[slot_mapping]