Fix/dsv4 native transformers warmup

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
 

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)

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),

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,37 +28,43 @@ 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
         max_aligned_rows = _align(max_rows, alignment)
         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)

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