diff --git a/docker/install.sh b/docker/install.sh index 6142d6da31..9c9c75016a 100755 --- a/docker/install.sh +++ b/docker/install.sh @@ -111,6 +111,8 @@ pip check if [[ "${CUDA_VERSION_SHORT}" = "cu128" ]]; then pip install "nvidia-nccl-cu12>=2.30.4" elif [[ "${CUDA_VERSION_SHORT}" == cu13* ]]; then + # tile-kernels (for DeepSeek-V4) requires a cu13 tilelang build. + pip install "tile-kernels==1.0.0" pip install "nvidia-nccl-cu13>=2.30.4" fi diff --git a/lmdeploy/hf_configs/__init__.py b/lmdeploy/hf_configs/__init__.py index 8a11474062..fa76fbefad 100644 --- a/lmdeploy/hf_configs/__init__.py +++ b/lmdeploy/hf_configs/__init__.py @@ -10,30 +10,28 @@ @lru_cache def register_config(model_type: str): - if model_type == 'deepseek_v32': - from .configuration_deepseek_v32 import DeepseekV32Config - AutoConfig.register(DeepseekV32Config.model_type, DeepseekV32Config) - elif model_type == 'deepseek_v4': - from .configuration_deepseek_v4 import DeepseekV4Config - AutoConfig.register(DeepseekV4Config.model_type, DeepseekV4Config) - else: - logger.debug(f'Can not register config for model_type: {model_type}') + """Register an LMDeploy-owned Transformers config when available.""" + logger.debug(f'Can not register config for model_type: {model_type}') def config_from_pretrained(pretrained_model_name_or_path: str, **kwargs): try: return AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) - except Exception as e: - logger.debug(f'AutoConfig.from_pretrained failed: {e}, try register config manually.') - # some models do not provide auto map for config - from transformers import PretrainedConfig - trust_remote_code = kwargs.pop('trust_remote_code', None) - config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **kwargs) - model_type = config_dict.get('model_type', None) - if trust_remote_code is not None: - kwargs['trust_remote_code'] = trust_remote_code - register_config(model_type) - try: - return AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) - except Exception as e: - return PretrainedConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + except Exception as error: + logger.debug(f'AutoConfig.from_pretrained failed: {error}, try register config manually.') + + # Some models neither provide an auto_map nor have a native Transformers + # config. Read model_type without executing remote code so LMDeploy can + # register a local config, then fall back to the generic config when no + # local implementation exists. + from transformers import PretrainedConfig + trust_remote_code = kwargs.pop('trust_remote_code', None) + config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **kwargs) + model_type = config_dict.get('model_type') + if trust_remote_code is not None: + kwargs['trust_remote_code'] = trust_remote_code + register_config(model_type) + try: + return AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + except Exception: + return PretrainedConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) diff --git a/lmdeploy/hf_configs/configuration_deepseek_v32.py b/lmdeploy/hf_configs/configuration_deepseek_v32.py deleted file mode 100644 index 59f91fb3ef..0000000000 --- a/lmdeploy/hf_configs/configuration_deepseek_v32.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. - -from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config - - -class DeepseekV32Config(DeepseekV3Config): - model_type = 'deepseek_v32' - - def __init__(self, index_head_dim=128, index_n_heads=64, index_topk=2048, **kwargs): - super().__init__(**kwargs) - self.index_head_dim = index_head_dim - self.index_n_heads = index_n_heads - self.index_topk = index_topk diff --git a/lmdeploy/hf_configs/configuration_deepseek_v4.py b/lmdeploy/hf_configs/configuration_deepseek_v4.py deleted file mode 100644 index 3ef974970a..0000000000 --- a/lmdeploy/hf_configs/configuration_deepseek_v4.py +++ /dev/null @@ -1,100 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -from transformers.configuration_utils import PretrainedConfig - - -class DeepseekV4Config(PretrainedConfig): - model_type = 'deepseek_v4' - - def __init__(self, - architectures=None, - attention_bias=False, - attention_dropout=0.0, - bos_token_id=0, - eos_token_id=1, - hc_eps=1e-6, - hc_mult=4, - hc_sinkhorn_iters=20, - head_dim=512, - hidden_act='silu', - hidden_size=4096, - index_head_dim=128, - index_n_heads=64, - index_topk=512, - initializer_range=0.02, - max_position_embeddings=1048576, - moe_intermediate_size=2048, - n_routed_experts=256, - n_shared_experts=1, - norm_topk_prob=True, - num_attention_heads=64, - num_experts_per_tok=6, - num_hidden_layers=43, - num_hash_layers=3, - num_key_value_heads=1, - num_nextn_predict_layers=1, - o_groups=8, - o_lora_rank=1024, - q_lora_rank=1024, - qk_rope_head_dim=64, - quantization_config=None, - rms_norm_eps=1e-6, - rope_scaling=None, - rope_theta=10000, - routed_scaling_factor=1.5, - scoring_func='sqrtsoftplus', - sliding_window=128, - swiglu_limit=10.0, - tie_word_embeddings=False, - topk_method='noaux_tc', - torch_dtype='bfloat16', - use_cache=True, - vocab_size=129280, - compress_rope_theta=160000, - compress_ratios=None, - **kwargs): - super().__init__(bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - torch_dtype=torch_dtype, - **kwargs) - self.architectures = architectures or ['DeepseekV4ForCausalLM'] - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.hc_eps = hc_eps - self.hc_mult = hc_mult - self.hc_sinkhorn_iters = hc_sinkhorn_iters - self.head_dim = head_dim - self.hidden_act = hidden_act - self.hidden_size = hidden_size - self.index_head_dim = index_head_dim - self.index_n_heads = index_n_heads - self.index_topk = index_topk - self.initializer_range = initializer_range - self.max_position_embeddings = max_position_embeddings - self.moe_intermediate_size = moe_intermediate_size - self.n_routed_experts = n_routed_experts - self.n_shared_experts = n_shared_experts - self.norm_topk_prob = norm_topk_prob - self.num_attention_heads = num_attention_heads - self.num_experts_per_tok = num_experts_per_tok - self.num_hidden_layers = num_hidden_layers - self.num_hash_layers = num_hash_layers - self.num_key_value_heads = num_key_value_heads - self.num_nextn_predict_layers = num_nextn_predict_layers - self.o_groups = o_groups - self.o_lora_rank = o_lora_rank - self.q_lora_rank = q_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.quantization_config = quantization_config - self.rms_norm_eps = rms_norm_eps - self.rope_scaling = rope_scaling - self.rope_theta = rope_theta - self.routed_scaling_factor = routed_scaling_factor - self.scoring_func = scoring_func - self.sliding_window = sliding_window - self.swiglu_limit = swiglu_limit - self.topk_method = topk_method - self.use_cache = use_cache - self.vocab_size = vocab_size - self.compress_rope_theta = compress_rope_theta - self.compress_ratios = compress_ratios or [0, 0] diff --git a/lmdeploy/pytorch/backends/cuda/attention/v4.py b/lmdeploy/pytorch/backends/cuda/attention/v4.py index b913e892d1..463d6baec8 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/v4.py +++ b/lmdeploy/pytorch/backends/cuda/attention/v4.py @@ -68,6 +68,7 @@ class _V4IndexScoreMeta: """CUDA backend metadata for paged index-score generation.""" context_lens: torch.Tensor = None + topk_seqlens: torch.Tensor = None block_offsets: torch.Tensor = None schedule_metadata: object = None max_k_seqlen: int = None @@ -253,12 +254,17 @@ def _build_index_score_meta(meta, ratio: int, context_lens: torch.Tensor, if page_table.dtype != torch.int32: page_table = page_table.to(torch.int32) - context_lens = context_lens.to(torch.int32).unsqueeze(-1) + topk_seqlens = context_lens.to(torch.int32) + # DeepGEMM's scheduler cannot represent an all-empty workload. A + # one-entry dummy score is safe because top-k still uses the true + # visible lengths and discards it for empty rows. + context_lens = topk_seqlens.clamp_min(1).unsqueeze(-1) entries_per_block = meta.block_size // ratio schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( context_lens, entries_per_block, deep_gemm.get_num_sms()) return _V4IndexScoreMeta( context_lens=context_lens, + topk_seqlens=topk_seqlens, block_offsets=page_table, schedule_metadata=schedule_metadata, max_k_seqlen=CudaV4AttentionMetadata._get_index_score_max_len(meta, ratio), diff --git a/lmdeploy/pytorch/backends/cuda/v4_indexer.py b/lmdeploy/pytorch/backends/cuda/v4_indexer.py index 58de6197a1..79d0832adf 100644 --- a/lmdeploy/pytorch/backends/cuda/v4_indexer.py +++ b/lmdeploy/pytorch/backends/cuda/v4_indexer.py @@ -3,6 +3,7 @@ import torch +from lmdeploy.pytorch.consts import V4_INDEX_SCALE_BYTES from lmdeploy.pytorch.kernels.cuda.bitonic_topk import bitonic_topk from lmdeploy.pytorch.kernels.cuda.blocked_gemm_fp8 import quant_fp8 from lmdeploy.pytorch.kernels.cuda.ds_index import fp8_index @@ -27,29 +28,32 @@ class _V4PagedMQALogitsWarmup: _METADATA_ROW_ALIGNMENT = 32 - @staticmethod - def _get_shape(model_config) -> tuple[int, int, int] | None: - hf_config = getattr(model_config, 'hf_config', None) - if getattr(hf_config, 'model_type', None) != 'deepseek_v4': - return None - if 4 not in getattr(hf_config, 'compress_ratios', ()): + def __init__(self, compress_ratio: int, num_heads: int, head_dim: int): + self.compress_ratio = compress_ratio + self.num_heads = num_heads + self.head_dim = head_dim + + def _get_shape(self, model_config) -> tuple[int, int, int] | None: + block_size = model_config.block_size + if block_size % self.compress_ratio != 0: + logger.warning( + 'Skip DeepGEMM V4 paged-MQA warmup: block_size=%s is not ' + 'divisible by compress_ratio=%s.', block_size, + self.compress_ratio) return None - block_size = getattr(model_config, 'block_size', 256) - entries_per_block = block_size // 4 - num_heads = getattr(hf_config, 'index_n_heads', None) - head_dim = getattr(hf_config, 'index_head_dim', 128) - if num_heads not in (32, 64): + entries_per_block = block_size // self.compress_ratio + if self.num_heads not in (32, 64): logger.warning( - 'Skip DeepGEMM V4 paged-MQA warmup: unsupported index_n_heads=%s.', - num_heads) + 'Skip DeepGEMM V4 paged-MQA warmup: unsupported num_heads=%s.', + self.num_heads) return None if entries_per_block not in (32, 64): logger.warning( 'Skip DeepGEMM V4 paged-MQA warmup: unsupported entries_per_block=%s.', entries_per_block) return None - return entries_per_block, num_heads, head_dim + return entries_per_block, self.num_heads, self.head_dim def _warm_metadata(self, deep_gemm, entries_per_block: int, max_rows: int): alignment = self._METADATA_ROW_ALIGNMENT @@ -57,7 +61,10 @@ def _warm_metadata(self, deep_gemm, entries_per_block: int, max_rows: int): rows = list(range(alignment, max_aligned_rows + 1, alignment)) random.shuffle(rows) for row_count in rows: - context_lens = torch.zeros(row_count, 1, dtype=torch.int32, device='cuda') + # DeepGEMM's metadata scheduler cannot represent an all-empty + # workload. Use one valid compressed cache page per row. + context_lens = torch.full( + (row_count, 1), entries_per_block, dtype=torch.int32, device='cuda') deep_gemm.get_paged_mqa_logits_metadata( context_lens, entries_per_block, deep_gemm.get_num_sms()) @@ -75,7 +82,7 @@ def _warm_logits(self, deep_gemm, entries_per_block: int, num_heads: int, dtype=torch.float8_e4m3fn, device='cuda') weights = torch.empty(row_count, num_heads, dtype=torch.float32, device='cuda') fused_kv_cache = torch.empty( - max_blocks, entries_per_block, 1, head_dim + 4, + max_blocks, entries_per_block, 1, head_dim + V4_INDEX_SCALE_BYTES, dtype=torch.uint8, device='cuda') deep_gemm.fp8_paged_mqa_logits( q, fused_kv_cache, weights, context_lens, block_table, @@ -90,7 +97,8 @@ def warmup(self, warmup_meta): entries_per_block, num_heads, head_dim = shape max_rows = max(warmup_meta.max_num_tokens, warmup_meta.max_batch_size) - max_context_len = max(warmup_meta.max_num_tokens // 4, 1) + max_context_len = max( + warmup_meta.max_num_tokens // self.compress_ratio, 1) logger.info( 'Warming up DeepGEMM V4 paged-MQA logits: rows<=%s stride=%s, ' 'heads=%s, head_dim=%s, entries_per_block=%s.', @@ -106,15 +114,19 @@ def warmup(self, warmup_meta): class TritonV4IndexerImpl(BaseV4Indexer): - def __init__(self, index_topk: int, compress_ratio: int) -> None: + def __init__(self, index_topk: int, compress_ratio: int, num_heads: int, + head_dim: int) -> None: super().__init__() self.index_topk = index_topk self.compress_ratio = compress_ratio if compress_ratio == 4: warmup_mgr = get_warmup_manager() - key = 'deepgemm_v4_paged_mqa_logits' + key = f'deepgemm_v4_paged_mqa_logits_{num_heads}_{head_dim}' if key not in warmup_mgr: - warmup_mgr[key] = _V4PagedMQALogitsWarmup().warmup + warmup_mgr[key] = _V4PagedMQALogitsWarmup( + compress_ratio=compress_ratio, + num_heads=num_heads, + head_dim=head_dim).warmup def forward(self, query: torch.Tensor, @@ -171,7 +183,8 @@ def forward(self, score_meta = getattr(meta, 'index_score_meta', None) if (score_meta is None or score_meta.context_lens is None or score_meta.block_offsets is None - or score_meta.schedule_metadata is None or score_meta.max_k_seqlen is None): + or score_meta.schedule_metadata is None or score_meta.max_k_seqlen is None + or score_meta.topk_seqlens is None): raise RuntimeError('Packed V4 index scoring requires CUDA index-score metadata.') q_mqa = q_3d.view(q_3d.size(0), 1, q_3d.size(1), q_3d.size(2)) scores = deep_gemm.fp8_paged_mqa_logits( @@ -183,7 +196,7 @@ def forward(self, score_meta.schedule_metadata, score_meta.max_k_seqlen, False) - topk_seqlens = score_meta.context_lens.squeeze(-1) + topk_seqlens = score_meta.topk_seqlens else: # fp8_index fallback for the legacy split value/scale cache layout. if index_kv_scale_cache is None: @@ -224,5 +237,12 @@ def forward(self, class TritonV4IndexerBuilder(BaseV4IndexerBuilder): @staticmethod - def build(index_topk: int, compress_ratio: int) -> BaseV4Indexer: - return TritonV4IndexerImpl(index_topk=index_topk, compress_ratio=compress_ratio) + def build(index_topk: int, + compress_ratio: int, + num_heads: int, + head_dim: int) -> BaseV4Indexer: + return TritonV4IndexerImpl( + index_topk=index_topk, + compress_ratio=compress_ratio, + num_heads=num_heads, + head_dim=head_dim) diff --git a/lmdeploy/pytorch/backends/indexer.py b/lmdeploy/pytorch/backends/indexer.py index 23006b4f3a..e305c01788 100644 --- a/lmdeploy/pytorch/backends/indexer.py +++ b/lmdeploy/pytorch/backends/indexer.py @@ -46,6 +46,9 @@ class BaseV4IndexerBuilder: @staticmethod @abstractmethod - def build(index_topk: int, compress_ratio: int) -> BaseV4Indexer: + def build(index_topk: int, + compress_ratio: int, + num_heads: int, + head_dim: int) -> BaseV4Indexer: """Build layer implementation.""" raise NotImplementedError diff --git a/lmdeploy/pytorch/configurations/deepseek_v4.py b/lmdeploy/pytorch/configurations/deepseek_v4.py index 883c1f14d3..31a2dc29a4 100644 --- a/lmdeploy/pytorch/configurations/deepseek_v4.py +++ b/lmdeploy/pytorch/configurations/deepseek_v4.py @@ -18,48 +18,49 @@ V4_PACKED_TOKEN_DIM = V4_FLASHMLA_D_NOPE + 2 * V4_FLASHMLA_D_ROPE + V4_FLASHMLA_NUM_TILES + 1 +V4_BLOCK_SIZE = 256 V4_SUPPORTED_COMPRESS_RATIOS = (0, 4, 128) +V4_SUPPORTED_LAYER_TYPES = ( + 'sliding_attention', + 'compressed_sparse_attention', + 'heavily_compressed_attention', +) - -def _normalize_v4_block_size(block_size: int) -> int: - """Return the logical block size required by DeepSeek-V4 kernels.""" - if block_size < 256: - block_size = 256 - if block_size % 128 != 0: - block_size = ((block_size + 127) // 128) * 128 - if block_size < 256: - block_size = 256 - return block_size - - -def _get_v4_cache_layers(hf_config): - """Normalize compression ratios and return layer-id partitions.""" +def get_v4_compress_ratios(hf_config) -> list[int]: + """Translate the native Transformers layer schema to compression ratios.""" num_layers = hf_config.num_hidden_layers - compress_ratios = getattr(hf_config, 'compress_ratios', None) - if compress_ratios is None: - compress_ratios = [0] * num_layers - else: - compress_ratios = list(compress_ratios) - - if len(compress_ratios) > num_layers: - extra_ratios = compress_ratios[num_layers:] - if any(r != 0 for r in extra_ratios): - raise ValueError('DeepSeek-V4 compress_ratios has extra non-zero entries beyond ' - f'num_hidden_layers={num_layers}: {extra_ratios}.') - logger.warning('DeepSeek-V4 compress_ratios has %s entries but num_hidden_layers is %s. ' - 'Ignoring trailing zero entries.', len(compress_ratios), num_layers) - compress_ratios = compress_ratios[:num_layers] - elif len(compress_ratios) < num_layers: - logger.warning('DeepSeek-V4 compress_ratios has %s entries but num_hidden_layers is %s. ' - 'Padding missing entries with 0.', len(compress_ratios), num_layers) - compress_ratios = compress_ratios + [0] * (num_layers - len(compress_ratios)) + layer_types = list(hf_config.layer_types) + if len(layer_types) != num_layers: + raise ValueError('DeepSeek-V4 requires one layer_type per hidden layer, but got ' + f'{len(layer_types)} layer types for {num_layers} layers.') + + invalid_layer_types = sorted(set(layer_types).difference(V4_SUPPORTED_LAYER_TYPES)) + if invalid_layer_types: + raise ValueError('DeepSeek-V4 layer_types only supports ' + f'{V4_SUPPORTED_LAYER_TYPES}, but got {invalid_layer_types}.') + + compress_rates = hf_config.compress_rates + compressed_layer_types = set(layer_types).difference({'sliding_attention'}) + missing_rates = sorted(compressed_layer_types.difference(compress_rates)) + if missing_rates: + raise ValueError(f'DeepSeek-V4 compress_rates is missing layer types: {missing_rates}.') + + compress_ratios = [ + 0 if layer_type == 'sliding_attention' else compress_rates[layer_type] + for layer_type in layer_types + ] invalid_ratios = sorted({r for r in compress_ratios if r not in V4_SUPPORTED_COMPRESS_RATIOS}) if invalid_ratios: - raise ValueError('DeepSeek-V4 compress_ratios only supports ' + raise ValueError('DeepSeek-V4 compression only supports ratios ' f'{V4_SUPPORTED_COMPRESS_RATIOS}, but got {invalid_ratios}.') + return compress_ratios + - hf_config.compress_ratios = compress_ratios +def _get_v4_cache_layers(hf_config): + """Return layer-id partitions for each V4 compression ratio.""" + num_layers = hf_config.num_hidden_layers + compress_ratios = get_v4_compress_ratios(hf_config) all_layers = list(range(num_layers)) ratio4_layers = [i for i, r in enumerate(compress_ratios) if r == 4] ratio128_layers = [i for i, r in enumerate(compress_ratios) if r == 128] @@ -91,9 +92,9 @@ def _check_env_v4(device: str = 'cuda'): def _finalize_v4_cache_specs(model_config: ModelConfig, block_size: int): - normalized_block_size = _normalize_v4_block_size(block_size) + normalized_block_size = V4_BLOCK_SIZE if normalized_block_size != block_size: - logger.warning(f'DeepSeek-V4 requires block_size >= 256 and a multiple of 128. ' + logger.warning(f'DeepSeek-V4 requires block_size={V4_BLOCK_SIZE}. ' f'Adjusting block_size from {model_config.block_size} to {normalized_block_size}.') model_config.block_size = normalized_block_size block_size = normalized_block_size @@ -123,9 +124,9 @@ def _finalize_v4_cache_specs(model_config: ModelConfig, block_size: int): def update_cache_config(cache_config): original_block_size = cache_config.block_size original_kernel_block_size = cache_config.kernel_block_size - block_size = _normalize_v4_block_size(original_block_size) + block_size = V4_BLOCK_SIZE if block_size != original_block_size: - logger.warning(f'DeepSeek-V4 requires block_size >= 256 and a multiple of 128. ' + logger.warning(f'DeepSeek-V4 requires block_size={V4_BLOCK_SIZE}. ' f'Adjusting block_size from {original_block_size} to {block_size}.') cache_config.block_size = block_size if cache_config.kernel_block_size != block_size: diff --git a/lmdeploy/pytorch/engine/logits_process.py b/lmdeploy/pytorch/engine/logits_process.py index 34b5a2af55..c234f736e6 100644 --- a/lmdeploy/pytorch/engine/logits_process.py +++ b/lmdeploy/pytorch/engine/logits_process.py @@ -364,7 +364,9 @@ def _fill_guided_bitmask(self, guided_bitmask: torch.Tensor): if not self.guided_decoding_manager.is_terminated(processor): self.guided_decoding_manager.fill_bitmap(processor, guided_bitmask, i) - def _accept_guided_tokens(self, next_token_ids: torch.Tensor): + def _accept_guided_tokens(self, next_token_ids: torch.Tensor, ready_event: torch.cuda.Event | None): + if ready_event is not None: + ready_event.synchronize() cpu_result = next_token_ids.tolist() for i, processor in self.guided_processors.items(): if self.guided_decoding_manager.is_terminated(processor): @@ -374,7 +376,13 @@ def _accept_guided_tokens(self, next_token_ids: torch.Tensor): async def accept_guided_tokens(self, next_token_ids: torch.Tensor): if self.guided_decoding_manager and self.guided_processors: - await asyncio.to_thread(self._accept_guided_tokens, next_token_ids) + ready_event = None + if next_token_ids.is_cuda: + # Sampling runs on the forward stream, while to_thread uses a + # different thread-local CUDA stream. Preserve that dependency. + ready_event = torch.cuda.Event() + ready_event.record() + await asyncio.to_thread(self._accept_guided_tokens, next_token_ids, ready_event) async def __call__(self, scores: torch.Tensor) -> torch.Tensor: r""" diff --git a/lmdeploy/pytorch/kernels/cuda/dsv4/hc_prepost.py b/lmdeploy/pytorch/kernels/cuda/dsv4/hc_prepost.py index 0a4646c163..626669165e 100644 --- a/lmdeploy/pytorch/kernels/cuda/dsv4/hc_prepost.py +++ b/lmdeploy/pytorch/kernels/cuda/dsv4/hc_prepost.py @@ -164,7 +164,10 @@ def hc_post_expand( *x.stride(), *residual.stride(), *post.stride(), - *comb.stride(), + # comb is [n, src, out]; the kernel indexes (out_h, src_h), so out_h + # strides the OUT axis (stride(2)) and src_h the SRC axis (stride(1)). + # *comb.stride() swaps them and applies comb transposed. + comb.stride(0), comb.stride(2), comb.stride(1), *out.stride(), dim, hc_mult, diff --git a/lmdeploy/pytorch/models/deepseek_v4.py b/lmdeploy/pytorch/models/deepseek_v4.py index 3e07208b32..b04a05014a 100644 --- a/lmdeploy/pytorch/models/deepseek_v4.py +++ b/lmdeploy/pytorch/models/deepseek_v4.py @@ -13,6 +13,7 @@ from lmdeploy.pytorch.backends.compressor import V4CompressorMetadata from lmdeploy.pytorch.backends.indexer import V4IndexerMetadata from lmdeploy.pytorch.backends.rotary_embedding import RopeType, YarnParameters +from lmdeploy.pytorch.configurations.deepseek_v4 import get_v4_compress_ratios from lmdeploy.pytorch.distributed import get_tp_world_rank from lmdeploy.pytorch.model_inputs import StepContext, StepContextManager from lmdeploy.pytorch.nn import ApplyRotaryEmb, HcPrePost, RMSNorm, SiluAndMul, rms_scale @@ -83,7 +84,7 @@ class V4Args: vocab_size: int moe_inter_dim: int n_layers: int - n_hash_layers: int + mlp_layer_types: tuple[str, ...] n_routed_experts: int n_shared_experts: int n_activated_experts: int @@ -306,8 +307,11 @@ def __init__(self, ) self.compressor = Compressor(args, layer_id, compress_ratio, self.head_dim, dtype=dtype, device=device, rotate=True) - self.indexer_fwd = NativeV4Indexer(index_topk=self.index_topk, - compress_ratio=self.compress_ratio) + self.indexer_fwd = NativeV4Indexer( + index_topk=self.index_topk, + compress_ratio=self.compress_ratio, + num_heads=self.n_heads, + head_dim=self.head_dim) self.apply_rotary = ApplyRotaryEmb() def forward(self, @@ -550,7 +554,7 @@ def __init__(self, layer_id: int, args: V4Args, device: torch.device | str | Non self.topk = args.n_activated_experts self.score_func = args.score_func self.route_scale = args.route_scale - self.hash = layer_id < args.n_hash_layers + self.hash = args.mlp_layer_types[layer_id] == 'hash_moe' self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim, device=device), requires_grad=False) if self.hash: self.tid2eid = nn.Parameter(torch.empty(args.vocab_size, @@ -726,13 +730,14 @@ def __init__(self, if isinstance(device, str): device = torch.device(device) self.device = device + compress_rope_params = config.rope_parameters['compress'] self.args = V4Args( dim=config.hidden_size, n_heads=config.num_attention_heads, vocab_size=config.vocab_size, moe_inter_dim=config.moe_intermediate_size, n_layers=config.num_hidden_layers, - n_hash_layers=config.num_hash_layers, + mlp_layer_types=tuple(config.mlp_layer_types), n_routed_experts=config.n_routed_experts, n_shared_experts=config.n_shared_experts, n_activated_experts=config.num_experts_per_tok, @@ -746,13 +751,13 @@ def __init__(self, o_groups=config.o_groups, o_lora_rank=config.o_lora_rank, window_size=config.sliding_window, - compress_ratios=tuple(config.compress_ratios), + compress_ratios=tuple(get_v4_compress_ratios(config)), compress_rope_theta=config.compress_rope_theta, - original_seq_len=config.rope_scaling['original_max_position_embeddings'], + original_seq_len=compress_rope_params['original_max_position_embeddings'], rope_theta=config.rope_theta, - rope_factor=config.rope_scaling['factor'], - beta_fast=config.rope_scaling['beta_fast'], - beta_slow=config.rope_scaling['beta_slow'], + rope_factor=compress_rope_params['factor'], + beta_fast=compress_rope_params['beta_fast'], + beta_slow=compress_rope_params['beta_slow'], index_n_heads=config.index_n_heads, index_head_dim=config.index_head_dim, index_topk=config.index_topk, diff --git a/lmdeploy/pytorch/nn/v4_indexer.py b/lmdeploy/pytorch/nn/v4_indexer.py index 87d3068d4a..6a46e8288b 100644 --- a/lmdeploy/pytorch/nn/v4_indexer.py +++ b/lmdeploy/pytorch/nn/v4_indexer.py @@ -8,11 +8,16 @@ class V4Indexer(nn.Module): """DeepSeek V4 indexer wrapper.""" - def __init__(self, index_topk: int, compress_ratio: int): + def __init__(self, index_topk: int, compress_ratio: int, num_heads: int, + head_dim: int): super().__init__() backend = get_backend() impl_builder = backend.get_layer_impl_builder(OpType.V4Indexer) - self.impl = impl_builder.build(index_topk=index_topk, compress_ratio=compress_ratio) + self.impl = impl_builder.build( + index_topk=index_topk, + compress_ratio=compress_ratio, + num_heads=num_heads, + head_dim=head_dim) def forward(self, query, diff --git a/lmdeploy/pytorch/transformers/configuration_deepseek_v32.py b/lmdeploy/pytorch/transformers/configuration_deepseek_v32.py deleted file mode 100644 index 3dd41d5dff..0000000000 --- a/lmdeploy/pytorch/transformers/configuration_deepseek_v32.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -from lmdeploy.hf_configs.configuration_deepseek_v32 import DeepseekV32Config - -__all__ = ['DeepseekV32Config'] diff --git a/lmdeploy/pytorch/transformers/configuration_deepseek_v4.py b/lmdeploy/pytorch/transformers/configuration_deepseek_v4.py deleted file mode 100644 index c89cb6fbec..0000000000 --- a/lmdeploy/pytorch/transformers/configuration_deepseek_v4.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -from lmdeploy.hf_configs.configuration_deepseek_v4 import DeepseekV4Config - -__all__ = ['DeepseekV4Config'] diff --git a/tests/pytorch/config/test_model_config.py b/tests/pytorch/config/test_model_config.py index c4c94edb42..66a4fd809d 100644 --- a/tests/pytorch/config/test_model_config.py +++ b/tests/pytorch/config/test_model_config.py @@ -22,6 +22,11 @@ def _make_model_config(num_attention_heads=32, num_key_value_heads=8, dist_confi def _make_deepseek_v4_hf_config(compress_ratios, num_hidden_layers=3): + ratio_to_layer_type = { + 0: 'sliding_attention', + 4: 'compressed_sparse_attention', + 128: 'heavily_compressed_attention', + } return SimpleNamespace( model_type='deepseek_v4', architectures=['DeepseekV4ForCausalLM'], @@ -32,7 +37,11 @@ def _make_deepseek_v4_hf_config(compress_ratios, num_hidden_layers=3): eos_token_id=2, sliding_window=128, vocab_size=32000, - compress_ratios=compress_ratios, + layer_types=[ratio_to_layer_type[ratio] for ratio in compress_ratios], + compress_rates={ + 'compressed_sparse_attention': 4, + 'heavily_compressed_attention': 128, + }, index_head_dim=128, ) @@ -87,13 +96,14 @@ def test_get_num_qkv_head_by_tp_requires_divisible_heads(): model_config.get_num_qkv_head_by_tp() -@pytest.mark.parametrize(('block_size', 'kernel_block_size', 'expected_block_size'), [ - (192, 64, 256), - (256, 128, 256), - (257, 128, 384), +@pytest.mark.parametrize(('block_size', 'kernel_block_size'), [ + (64, 64), + (192, 64), + (256, 128), + (257, 128), + (512, 512), ]) -def test_deepseek_v4_update_cache_config_normalizes_block_and_kernel_size(block_size, kernel_block_size, - expected_block_size): +def test_deepseek_v4_update_cache_config_forces_block_and_kernel_size(block_size, kernel_block_size): cache_config = CacheConfig(max_batches=1, block_size=block_size, kernel_block_size=kernel_block_size, @@ -102,17 +112,17 @@ def test_deepseek_v4_update_cache_config_normalizes_block_and_kernel_size(block_ update_deepseek_v4_cache_config(cache_config) - assert cache_config.block_size == expected_block_size - assert cache_config.kernel_block_size == expected_block_size + assert cache_config.block_size == 256 + assert cache_config.kernel_block_size == 256 assert cache_config.window_size == -1 -def test_deepseek_v4_model_config_normalizes_block_cache_spec_shapes(): +def test_deepseek_v4_model_config_forces_default_block_cache_spec_shapes(): hf_config = _make_deepseek_v4_hf_config([4, 128], num_hidden_layers=2) model_config = AutoModelConfigBuilder.build(hf_config) - model_config.block_size = 192 - model_config.post_build_func(model_config, 192) + model_config.block_size = 64 + model_config.post_build_func(model_config, 64) assert model_config.block_size == 256 block_specs = {spec.name: spec for spec in model_config.block_cache_specs} @@ -122,13 +132,12 @@ def test_deepseek_v4_model_config_normalizes_block_cache_spec_shapes(): assert block_specs['v4_compressed_kv_r128_fp8'].shape[0] == 2 -def test_deepseek_v4_model_config_trims_trailing_zero_compress_ratio(): - hf_config = _make_deepseek_v4_hf_config([0, 4, 128, 0]) +def test_deepseek_v4_model_config_uses_native_layer_types(): + hf_config = _make_deepseek_v4_hf_config([0, 4, 128]) model_config = AutoModelConfigBuilder.build(hf_config) model_config.post_build_func(model_config, 256) - assert hf_config.compress_ratios == [0, 4, 128] state_specs = {spec.name: spec for spec in model_config.state_cache_specs} assert state_specs['v4_window_kv_fp8'].layer_ids == [0, 1, 2] assert state_specs['v4_compress_state_r4'].layer_ids == [1] @@ -141,9 +150,17 @@ def test_deepseek_v4_model_config_trims_trailing_zero_compress_ratio(): assert block_specs['v4_compressed_kv_r128_fp8'].layer_ids == [2] -def test_deepseek_v4_model_config_rejects_extra_nonzero_compress_ratio(): +def test_deepseek_v4_model_config_rejects_layer_type_count_mismatch(): hf_config = _make_deepseek_v4_hf_config([0, 4, 128]) hf_config.num_hidden_layers = 2 - with pytest.raises(ValueError, match='extra non-zero entries'): + with pytest.raises(ValueError, match='one layer_type per hidden layer'): + AutoModelConfigBuilder.build(hf_config) + + +def test_deepseek_v4_model_config_rejects_unsupported_native_compress_rate(): + hf_config = _make_deepseek_v4_hf_config([4], num_hidden_layers=1) + hf_config.compress_rates['compressed_sparse_attention'] = 8 + + with pytest.raises(ValueError, match='only supports ratios'): AutoModelConfigBuilder.build(hf_config) diff --git a/tests/pytorch/kernel/test_dsv4_hc_prepost.py b/tests/pytorch/kernel/test_dsv4_hc_prepost.py index dc6c3a716f..08f7db59a0 100644 --- a/tests/pytorch/kernel/test_dsv4_hc_prepost.py +++ b/tests/pytorch/kernel/test_dsv4_hc_prepost.py @@ -8,8 +8,10 @@ def _reference_pre_reduce(x, pre, out_dtype): def _reference_post_expand(x, residual, post, comb): + # comb is [..., src, out]; the post-expand mixes residuals as + # matmul(comb.T, residual) == sum_src comb[src, out] * residual[src]. y = post.unsqueeze(-1) * x.float().unsqueeze(-2) - y += torch.sum(comb.unsqueeze(-1) * residual.float().unsqueeze(-3), dim=-2) + y += torch.matmul(comb.float().transpose(-1, -2), residual.float()) return y.to(x.dtype)