diff --git a/docker/prepare_3rdparty_wheel.sh b/docker/prepare_3rdparty_wheel.sh index a5a2eb2dcf..957c4e7423 100644 --- a/docker/prepare_3rdparty_wheel.sh +++ b/docker/prepare_3rdparty_wheel.sh @@ -15,7 +15,8 @@ fi GDRCOPY_VERSION=2.5.1 DEEP_EP_VERSION=9af0e0d # v1.2.1 DEEP_GEMM_VERSION=88965b0 -FLASH_MLA_VERSION=1408756 # no release, pick the latest commit +FLASH_MLA_VERSION=9241ae3 +FAST_HADAMARD_TRANSFORM_VERSION=v1.1.0.post2 # DeepEP if [[ "${CUDA_VERSION_SHORT}" = "cu130" ]]; then @@ -33,6 +34,10 @@ pip wheel -v --no-build-isolation --no-deps -w /wheels "git+https://github.com/d # sm100 compilation for Flash MLA requires NVCC 12.9 or higher FLASH_MLA_DISABLE_SM100=1 pip wheel -v --no-build-isolation --no-deps -w /wheels "git+https://github.com/deepseek-ai/FlashMLA.git@${FLASH_MLA_VERSION}" +# fast_hadamard_transform +pip wheel -v --no-build-isolation --no-deps -w /wheels \ + "git+https://github.com/Dao-AILab/fast-hadamard-transform.git@${FAST_HADAMARD_TRANSFORM_VERSION}" + # flash_attn_3 (prebuilt wheels; CUDA + torch must match this image) TORCH_VER=$(python3 -c "import torch; print(''.join(torch.__version__.split('+')[0].split('.')))") FA3_WHEELS_URL="https://windreamer.github.io/flash-attention3-wheels/${CUDA_VERSION_SHORT}_torch${TORCH_VER}" diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 8d5f7d31d0..abf2197b59 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -274,7 +274,8 @@ def _parse(x): type=_parse, default=default, help='KV cache quant policy: none/int4/int8/fp8/fp8_e5m2/' - 'turbo_quant (or 0/4/8/16/17/42). fp8 defaults to fp8_e4m3.') + 'turbo_quant (or 0/4/8/16/17/42). For DSA models, fp8 uses the ' + 'fp8_ds_mla layout.') @staticmethod def rope_scaling_factor(parser): diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index e07a46ae5e..547b780afe 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -22,7 +22,7 @@ class QuantPolicy(enum.IntEnum): NONE = 0 INT4 = 4 # 4-bit KV cache INT8 = 8 # 8-bit KV cache - FP8 = 16 # FP8 KV cache (float8_e4m3fn, per-tensor scale) + FP8 = 16 # FP8 KV cache (float8_e4m3fn, per-tensor scale. DSA uses the fp8_ds_mla layout) FP8_E5M2 = 17 # FP8 KV cache (float8_e5m2, per-tensor scale) TURBO_QUANT = 42 # TurboQuant: K=4bit QJL4 + V=2bit MSE @@ -427,7 +427,7 @@ class PytorchEngineConfig: If unspecified, will use the default version. quant_policy: default to 0. When k/v is quantized into int4, int8, fp8, or fp8_e5m2, set it to 4, 8, 16, or 17, - respectively + respectively. For DSA models, fp8 selects the fp8_ds_mla layout. distributed_executor_backend: backend of distributed backend, options: ['uni', 'mp', 'ray'] empty_init: Whether to load the model weights, you should set diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index 2fc4fdb298..466ca1ee8d 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -36,26 +36,65 @@ class NSAIndicesUpdater: def __init__(self): self._update_decode_func = None + self._update_decode_strided_func = None self._update_prefill_func = None - def _update_decode_impl(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, + def _update_decode_impl(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int, block_size: int) -> torch.Tensor: """Update for decode impl.""" + batch_size = block_offsets.size(0) block_ids = nsa_indices // block_size block_ids = block_ids.clamp_min(0) + # MTP flattens the query dimension, so repeat each request's block table. + if block_ids.size(0) != batch_size: + block_offsets = torch.repeat_interleave(block_offsets, + max_q_seqlen, + dim=0, + output_size=block_ids.size(0)) block_ids = block_offsets.gather(1, block_ids) block_remain = nsa_indices % block_size ret = block_ids * block_size + block_remain ret[nsa_indices < 0] = -1 - return ret[:, None] + return ret.unflatten(0, (batch_size, max_q_seqlen)) - def update_decode(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, block_size: int) -> torch.Tensor: + def update_decode(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int, + block_size: int) -> torch.Tensor: """Update for decode.""" if self._update_decode_func is None: self._update_decode_func = _try_dynamic_compile(self._update_decode_impl, nsa_indices, block_offsets, - block_size) + max_q_seqlen, block_size) - return self._update_decode_func(nsa_indices, block_offsets, block_size) + return self._update_decode_func(nsa_indices, block_offsets, max_q_seqlen, block_size) + + def _update_decode_strided_impl(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int, + block_size: int, block_stride: int, token_stride: int, + index_stride: int) -> torch.Tensor: + """Map logical decode indices to an aligned strided cache view.""" + batch_size = block_offsets.size(0) + block_ids = nsa_indices // block_size + block_ids = block_ids.clamp_min(0) + if block_ids.size(0) != batch_size: + block_offsets = torch.repeat_interleave(block_offsets, + max_q_seqlen, + dim=0, + output_size=block_ids.size(0)) + block_ids = block_offsets.gather(1, block_ids) + block_remain = nsa_indices % block_size + ret = (block_ids * block_stride + block_remain * token_stride) // index_stride + ret[nsa_indices < 0] = -1 + return ret.unflatten(0, (batch_size, max_q_seqlen)) + + def update_decode_strided(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int, + block_size: int, block_stride: int, token_stride: int, + index_stride: int) -> torch.Tensor: + """Update decode indices for strided cache storage.""" + if self._update_decode_strided_func is None: + self._update_decode_strided_func = _try_dynamic_compile(self._update_decode_strided_impl, nsa_indices, + block_offsets, max_q_seqlen, block_size, + block_stride, token_stride, index_stride) + + return self._update_decode_strided_func(nsa_indices, block_offsets, max_q_seqlen, block_size, block_stride, + token_stride, index_stride) def _update_prefill_impl(self, nsa_indices: torch.Tensor, q_seqlens: torch.Tensor, cu_seqlens_k: torch.Tensor): """Update for prefill impl.""" @@ -84,8 +123,9 @@ class FlashMLAImpl(TritonAttentionImpl): """Flash MLA (Multi-head Latent Attention) implementation. This implementation supports multiple execution paths: - - Decoding: Uses flash_mla_with_kvcache with paged KV cache - - Prefill with NSA: Uses flash_mla_sparse_fwd for sparse attention + - Paged FlashMLA decode: Uses flash_mla_with_kvcache with FP8 paged KV cache + - Sparse FlashMLA decode: Uses flash_mla_sparse_fwd over a zero-copy BF16 cache view + - Sparse FlashMLA prefill: Uses flash_mla_sparse_fwd for NSA attention - Prefill with FA3: Uses flash_attn_varlen_func with split q_rope/q_nope - Prefill fallback: Uses custom Triton kernel """ @@ -94,6 +134,7 @@ class FlashMLAImpl(TritonAttentionImpl): _MLA_HEAD_ALIGNMENT = 64 # Query heads must be multiple of 64 for flash_mla _MLA_NOPE_SIZE = 512 # Size of non-positional embeddings _MLA_SCALE_SIZE = 16 # Size of FP8 quantization scales + _BF16_CACHE_INDEX_STRIDE = 64 # Address BF16 cache rows in aligned 128-byte units def __init__( self, @@ -152,7 +193,7 @@ def _get_flash_mla_sparse_fwd(self): except Exception: logger.exception('Can not import flash_mla_sparse_fwd from flash_mla.') - def flash_mla_decoding( + def _decode_paged_flash_mla( self, query: torch.Tensor, k_cache: torch.Tensor, @@ -170,14 +211,22 @@ def flash_mla_decoding( max_q_seqlen = query.numel() // (query.size(-1) * query.size(-2)) max_q_seqlen = max_q_seqlen // batch_size query = query.unflatten(0, (batch_size, max_q_seqlen)) + num_q_heads = query.size(2) if kv_seqlens.dtype == torch.int64: kv_seqlens = kv_seqlens.to(torch.int32) # update nsa indice according to flash-mla requirement if nsa_indices is not None: block_size = k_cache.size(1) - nsa_indices = self.nsa_updater.update_decode(nsa_indices, block_offsets, block_size) + nsa_indices = self.nsa_updater.update_decode(nsa_indices, block_offsets, max_q_seqlen, block_size) causal = False + # The new FlashMLASchedMeta API uses a sparse decoder that only + # accepts 64 or 128 query heads. The old API stores metadata in a + # tensor and supports the unpadded TP head count. + if not isinstance(attn_metadata.tile_scheduler_metadata, torch.Tensor): + pad_heads = -num_q_heads % self._MLA_HEAD_ALIGNMENT + if pad_heads: + query = torch.nn.functional.pad(query, (0, 0, 0, pad_heads)) attn_output, _ = self.flash_mla_with_kvcache(query, k_cache=k_cache, @@ -191,38 +240,23 @@ def flash_mla_decoding( is_fp8_kvcache=is_fp8_kvcache, indices=nsa_indices) + attn_output = attn_output[:, :, :num_q_heads] attn_output = attn_output.flatten(0, 1) return attn_output - def _prefill_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, nsa_indices: torch.Tensor, - attn_metadata: TritonAttentionMetadata) -> torch.Tensor: - """Sparse prefill using flash_mla_sparse_fwd. - - This path is used when NSA (Non-contiguous Sparse Attention) indices are provided. - Requires FP8 KV cache and flash_mla library. - - Args: - query: Query tensor. - flatten_k: Flattened key cache. - nsa_indices: Sparse attention indices. - attn_metadata: Attention metadata. - - Returns: - Attention output tensor. - """ - q_seqlens = attn_metadata.q_seqlens + def _flash_mla_sparse(self, query: torch.Tensor, indexed_k: torch.Tensor, + nsa_indices: torch.Tensor) -> torch.Tensor: + """Run sparse FlashMLA over index-addressable BF16 KV storage.""" flash_mla_sparse_fwd = self._get_flash_mla_sparse_fwd() - num_q_heads = query.size(1) # flash_mla_sparse_fwd requires query heads to be multiple of alignment if num_q_heads % self._MLA_HEAD_ALIGNMENT != 0: padding = self._MLA_HEAD_ALIGNMENT - num_q_heads % self._MLA_HEAD_ALIGNMENT query = torch.nn.functional.pad(query, (0, 0, 0, padding)) - nsa_indices = self.nsa_updater.update_prefill(nsa_indices, q_seqlens, attn_metadata.cu_seqlens_k) output = flash_mla_sparse_fwd( query, - flatten_k, + indexed_k, nsa_indices, sm_scale=self.scale, ) @@ -230,6 +264,34 @@ def _prefill_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, nsa_indi attn_output = attn_output[:, :num_q_heads] return attn_output + def _prefill_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, nsa_indices: torch.Tensor, + attn_metadata: TritonAttentionMetadata) -> torch.Tensor: + """Run sparse prefill over the flattened BF16 KV cache.""" + nsa_indices = self.nsa_updater.update_prefill(nsa_indices, attn_metadata.q_seqlens, + attn_metadata.cu_seqlens_k) + return self._flash_mla_sparse(query, flatten_k, nsa_indices) + + def _decode_bf16_sparse_flash_mla(self, query: torch.Tensor, k_cache: torch.Tensor, nsa_indices: torch.Tensor, + attn_metadata: TritonAttentionMetadata) -> torch.Tensor: + """Run sparse decode over a zero-copy BF16 paged-cache view.""" + assert query.dtype == torch.bfloat16, 'BF16 sparse MLA requires a bfloat16 query' + assert k_cache.dtype == torch.bfloat16, 'BF16 sparse MLA requires a bfloat16 KV cache' + block_size = k_cache.size(1) + max_q_seqlen = self._get_max_q_seqlen(query, attn_metadata) + # The cache shares a block record with indexer state, so flattening it + # copies the full cache capacity. Expose the same storage in aligned + # address units and let sparse indices point directly at each KV row. + index_stride = self._BF16_CACHE_INDEX_STRIDE + block_stride, token_stride = k_cache.stride()[:2] + last_token_offset = ((k_cache.size(0) - 1) * block_stride + (block_size - 1) * token_stride) + storage_rows = last_token_offset // index_stride + 1 + storage_k = k_cache.as_strided((storage_rows, *k_cache.shape[2:]), + (index_stride, *k_cache.stride()[2:])) + nsa_indices = self.nsa_updater.update_decode_strided(nsa_indices, attn_metadata.block_offsets, max_q_seqlen, + block_size, block_stride, token_stride, index_stride) + nsa_indices = nsa_indices.flatten(0, 1)[:, None] + return self._flash_mla_sparse(query, storage_k, nsa_indices) + def _prefill_triton( self, query: torch.Tensor, @@ -392,6 +454,7 @@ def _fill_kv_cache_impl(self, """Fill kv cache.""" is_fp8_kvcache = k_cache.dtype == torch.float8_e4m3fn if not is_fp8_kvcache: + # The BF16 MLA V cache aliases K, and the base writer skips its duplicate store. return super()._fill_kv_cache_impl( key, value, @@ -456,8 +519,8 @@ def _forward_decoding( ) -> torch.Tensor: """Forward pass for decoding stage. - Uses flash_mla_with_kvcache for efficient decoding with paged KV cache. - Supports both regular and sparse (NSA) attention patterns. + Uses paged FlashMLA decode by default. Sparse NSA with BF16 cache uses + sparse FlashMLA decode. Args: query: Query tensor. @@ -468,7 +531,9 @@ def _forward_decoding( Returns: Attention output tensor. """ - return self.flash_mla_decoding(query, k_cache, nsa_indices, attn_metadata) + if nsa_indices is not None and k_cache.dtype != torch.float8_e4m3fn: + return self._decode_bf16_sparse_flash_mla(query, k_cache, nsa_indices, attn_metadata) + return self._decode_paged_flash_mla(query, k_cache, nsa_indices, attn_metadata) def _forward_prefill( self, @@ -483,7 +548,7 @@ def _forward_prefill( """Forward pass for prefill stage. Supports three execution paths: - 1. Sparse (NSA + FP8): flash_mla_sparse_fwd for sparse attention + 1. Sparse FlashMLA: flash_mla_sparse_fwd over BF16 flattened KV 2. FA3 optimized: flash_attn_varlen_func with split q_rope/q_nope 3. Triton fallback: Custom Triton kernel implementation @@ -534,15 +599,15 @@ def forward( """Forward pass for MLA attention computation. This method handles both prefill and decoding stages by: - 1. Validating NSA requirements (FP8 KV cache) - 2. Computing max query sequence length - 3. Filling KV cache if new key/value are provided - 4. Dispatching to appropriate stage-specific method + 1. Computing max query sequence length + 2. Filling KV cache if new key/value are provided + 3. Dispatching to the cache-specific stage implementation Architecture: - - Decoding: Uses flash_mla_with_kvcache with paged KV cache + - Paged FlashMLA decode: Uses flash_mla_with_kvcache with FP8 paged KV cache + - Sparse FlashMLA decode: Uses flash_mla_sparse_fwd over a zero-copy BF16 cache view - Prefill: Three paths based on availability and requirements - * Sparse (NSA + FP8): flash_mla_sparse_fwd + * Sparse FlashMLA: flash_mla_sparse_fwd * FA3 optimized: flash_attn_varlen_func with split q_rope/q_nope * Triton fallback: Custom triton kernel @@ -560,12 +625,6 @@ def forward( Returns: Attention output tensor. """ - # Validate NSA requirements - is_nsa = nsa_indices is not None - if is_nsa: - is_fp8_kvcache = k_cache.dtype == torch.float8_e4m3fn - assert is_fp8_kvcache, 'NSA sparse attention requires FP8 KV cache' - # Shared preparation max_q_seqlen = self._get_max_q_seqlen(query, attn_metadata) diff --git a/lmdeploy/pytorch/backends/cuda/graph_runner.py b/lmdeploy/pytorch/backends/cuda/graph_runner.py index 924f301c0f..ed1e9acec1 100644 --- a/lmdeploy/pytorch/backends/cuda/graph_runner.py +++ b/lmdeploy/pytorch/backends/cuda/graph_runner.py @@ -240,7 +240,14 @@ def get_graph_key(self, input_ids: torch.Tensor, position_ids: torch.Tensor, pas batch_size = self._get_capture_tokens(batch_size) else: batch_size = self._get_capture_tokens(meta.padding_batch_size) - return (batch_size, is_decoding, enable_microbatch, query_len) + graph_key = (batch_size, is_decoding, enable_microbatch, query_len) + if 'skip_topk' in kwargs: + graph_key += (kwargs['skip_topk'], ) + if 'use_dense_index' in kwargs: + # Dense/sparse index selection is irrelevant when an MTP draft + # iteration reuses the target model's indices. + graph_key += (kwargs['use_dense_index'] and not kwargs.get('skip_topk', False), ) + return graph_key def _prepare_inputs(self, **kwargs): """Prepare inputs.""" diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index bab5d73094..b7190223c9 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -1,14 +1,49 @@ # Copyright (c) OpenMMLab. All rights reserved. +import functools + +import torch from torch import Tensor 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 +from lmdeploy.pytorch.kernels.cuda.ds_index import dense_index, fp8_index +from lmdeploy.pytorch.kernels.cuda.dsa_indexer_preprocess import prepare_dsa_indexer_k_cache, prepare_dsa_indexer_q from lmdeploy.pytorch.kernels.cuda.fill_kv_cache import fill_kv_cache_blocked_fp8 from ..nsa import BaseNSAIndexFP8, BaseNSAIndexFP8Builder, NSAIndexMeta +@functools.lru_cache +def _get_sparse_index_topk(topk: int): + try: + from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import ( + is_sparse_index_topk_supported, + sparse_index_topk, + ) + except ImportError: + return None + if is_sparse_index_topk_supported(topk): + return sparse_index_topk + return None + + +def _get_causal_k_seqlens(cu_seqlen_q: Tensor, q_seqlens: Tensor, k_seqlens: Tensor, + num_tokens: int) -> Tensor: + """Expand request-level final KV lengths into per-query causal lengths. + + For a request with Q query tokens and final KV length K, query offset i can + see K - Q + i + 1 positions. Top-k needs these row-wise lengths so it + cannot return future query positions masked by the score kernel. + """ + # Start row of each request in the flattened query tensor. + q_start = torch.repeat_interleave(cu_seqlen_q[:-1], q_seqlens, output_size=num_tokens) + # KV prefix that existed before the current multi-token query group. + history_lengths = torch.repeat_interleave(k_seqlens - q_seqlens, q_seqlens, output_size=num_tokens) + # Convert flattened row ids into request-local query offsets. + query_offsets = torch.arange(num_tokens, device=q_seqlens.device) - q_start + return history_lengths + query_offsets + 1 + + class TritonNSAIndexFP8(BaseNSAIndexFP8): def __init__(self, topk: int, softmax_scale: float, block_size: int, fill: int) -> None: @@ -20,11 +55,7 @@ def __init__(self, topk: int, softmax_scale: float, block_size: int, fill: int) # TODO: configable scale fmt self.scale_fmt = 'ue8m0' - def forward(self, q: Tensor, k: Tensor, weights: Tensor, k_cache: Tensor, k_s_cache: Tensor, - meta: NSAIndexMeta) -> Tensor: - - assert q.dim() == 3 - assert k.dim() == 2 + def _forward_index(self, q: Tensor, q_s: Tensor, k_cache: Tensor, k_s_cache: Tensor, meta: NSAIndexMeta) -> Tensor: cu_seqlen_q = meta.cu_seqlen_q q_seqlens = meta.q_seqlens k_seqlens = meta.k_seqlens @@ -32,6 +63,34 @@ def forward(self, q: Tensor, k: Tensor, weights: Tensor, k_cache: Tensor, k_s_ca max_q_seqlen = meta.max_q_seqlen max_kv_seqlen = meta.max_kv_seqlen + scores = fp8_index(q, + q_s, + k_cache, + k_s_cache[..., 0], + cu_seqlen_q, + k_seqlens, + block_offset, + max_q_seqlen=max_q_seqlen, + max_k_seqlen=max_kv_seqlen, + causal=True) + # Multi-token queries need one causal KV length per score row. + if scores.size(0) != k_seqlens.size(0): + k_seqlens = _get_causal_k_seqlens(cu_seqlen_q, q_seqlens, k_seqlens, scores.size(0)) + sparse_index_topk = _get_sparse_index_topk(self.topk) + if sparse_index_topk is not None: + return sparse_index_topk(scores, + q_seqlens, + k_seqlens, + self.topk, + fill=self.fill, + descending=True, + sorted=False) + return bitonic_topk(scores, q_seqlens, k_seqlens, self.topk, fill=self.fill, descending=True) + + def forward(self, q: Tensor, k: Tensor, weights: Tensor, k_cache: Tensor, k_s_cache: Tensor, + meta: NSAIndexMeta) -> Tensor: + assert q.dim() == 3 + assert k.dim() == 2 q_shape = q.shape q = q.reshape(-1, q_shape[-1]) q, q_s = quant_fp8(q, self.block_size, dtype=k_cache.dtype, trans_scale=True, scale_fmt=self.scale_fmt) @@ -45,24 +104,64 @@ def forward(self, q: Tensor, k: Tensor, weights: Tensor, k_cache: Tensor, k_s_ca None, k_s_cache[..., None, :], None, - cu_seqlen_q=cu_seqlen_q, - kv_seqlens=k_seqlens, - max_q_seqlen=max_q_seqlen, - block_offsets=block_offset, + cu_seqlen_q=meta.cu_seqlen_q, + kv_seqlens=meta.k_seqlens, + max_q_seqlen=meta.max_q_seqlen, + block_offsets=meta.block_offset, group_size=self.block_size, scale_fmt=self.scale_fmt) + return self._forward_index(q, q_s, k_cache, k_s_cache, meta) - scores = fp8_index(q, - q_s, - k_cache, - k_s_cache[..., 0], - cu_seqlen_q, - k_seqlens, - block_offset, - max_q_seqlen=max_q_seqlen, - max_k_seqlen=max_kv_seqlen, - causal=True) - return bitonic_topk(scores, q_seqlens, k_seqlens, self.topk, fill=self.fill, descending=True) + def forward_fused(self, q: Tensor, k: Tensor, weights: Tensor, norm_weight: Tensor, norm_bias: Tensor, cos: Tensor, + sin: Tensor, k_cache: Tensor, k_s_cache: Tensor, norm_eps: float, head_gate_scale: float, + rope_interleaved: bool, meta: NSAIndexMeta) -> Tensor: + """Prepare FP8 Q and write K cache without allocating rotated BF16 + Q/K.""" + q, q_s = prepare_dsa_indexer_q(q, + weights, + cos, + sin, + score_scale=self.softmax_scale * head_gate_scale, + out_dtype=k_cache.dtype, + rope_interleaved=rope_interleaved) + prepare_dsa_indexer_k_cache(k, + norm_weight, + norm_bias, + cos, + sin, + k_cache, + k_s_cache[..., 0], + cu_seqlen_q=meta.cu_seqlen_q, + kv_seqlens=meta.k_seqlens, + block_offsets=meta.block_offset, + max_q_seqlen=meta.max_q_seqlen, + eps=norm_eps, + rope_interleaved=rope_interleaved) + return self._forward_index(q, q_s, k_cache, k_s_cache, meta) + + def forward_k_only(self, k: Tensor, norm_weight: Tensor, norm_bias: Tensor, cos: Tensor, sin: Tensor, + k_cache: Tensor, k_s_cache: Tensor, norm_eps: float, rope_interleaved: bool, + meta: NSAIndexMeta) -> Tensor: + """Cache K and skip score/top-k work when all positions are + selected.""" + prepare_dsa_indexer_k_cache(k, + norm_weight, + norm_bias, + cos, + sin, + k_cache, + k_s_cache[..., 0], + cu_seqlen_q=meta.cu_seqlen_q, + kv_seqlens=meta.k_seqlens, + block_offsets=meta.block_offset, + max_q_seqlen=meta.max_q_seqlen, + eps=norm_eps, + rope_interleaved=rope_interleaved) + return dense_index(meta.q_seqlens, + meta.k_seqlens, + meta.max_q_seqlen, + self.topk, + fill=self.fill) class TritonNSAIndexFP8Builder(BaseNSAIndexFP8Builder): diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index 67c29ae26d..02f23b65e2 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -138,11 +138,19 @@ def get_v_block_shape( @classmethod def update_meta_flashmla(cls, attn_metadata, model_config: ModelConfig, decoding_query_len: int): """Update meta for flashmla.""" + is_fp8_kvcache = model_config.use_mla_fp8_cache + index_topk = model_config.mla_index_topk + # FlashMLA block tables and BF16 sparse physical-slot indices are int32. + if attn_metadata.block_offsets.dtype != torch.int32: + attn_metadata.block_offsets = attn_metadata.block_offsets.to(torch.int32) + + # BF16 sparse decode uses sparse_fwd and needs no paged scheduler metadata. + if index_topk is not None and not is_fp8_kvcache: + return + import flash_mla num_attention_heads, _ = model_config.get_num_qkv_head_by_tp() num_attention_heads *= decoding_query_len - is_fp8_kvcache = model_config.use_mla_fp8_cache - index_topk = model_config.mla_index_topk num_heads_q = None if index_topk is None else num_attention_heads tile_scheduler_metadata, num_splits = flash_mla.get_mla_metadata(attn_metadata.kv_seqlens.to(torch.int32), num_attention_heads, @@ -153,9 +161,6 @@ def update_meta_flashmla(cls, attn_metadata, model_config: ModelConfig, decoding attn_metadata.tile_scheduler_metadata = tile_scheduler_metadata attn_metadata.num_splits = num_splits - if attn_metadata.block_offsets.dtype != torch.int32: - attn_metadata.block_offsets = attn_metadata.block_offsets.to(torch.int32) - @classmethod def update_meta_flashattn(cls, attn_metadata, step_context): from lmdeploy.pytorch.models.utils.cudagraph import _get_meta_flashattn diff --git a/lmdeploy/pytorch/backends/nsa.py b/lmdeploy/pytorch/backends/nsa.py index c20c80868c..f682abec71 100644 --- a/lmdeploy/pytorch/backends/nsa.py +++ b/lmdeploy/pytorch/backends/nsa.py @@ -24,6 +24,19 @@ def forward(self, q: Tensor, k: Tensor, weights: Tensor, k_cache: Tensor, k_s_ca """forward.""" raise NotImplementedError('Not implemented.') + @abstractmethod + def forward_fused(self, q: Tensor, k: Tensor, weights: Tensor, norm_weight: Tensor, norm_bias: Tensor, cos: Tensor, + sin: Tensor, k_cache: Tensor, k_s_cache: Tensor, norm_eps: float, head_gate_scale: float, + rope_interleaved: bool, meta: NSAIndexMeta) -> Tensor: + """Forward with fused DSA indexer preparation.""" + raise NotImplementedError('Not implemented.') + + def forward_k_only(self, k: Tensor, norm_weight: Tensor, norm_bias: Tensor, cos: Tensor, sin: Tensor, + k_cache: Tensor, k_s_cache: Tensor, norm_eps: float, rope_interleaved: bool, + meta: NSAIndexMeta) -> Tensor: + """Cache K and return the dense identity index.""" + raise NotImplementedError('Not implemented.') + class BaseNSAIndexFP8Builder: diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index f730c985fd..813108c985 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -383,7 +383,7 @@ class ModelConfig: # flash mla use_flash_mla: bool = False - use_mla_fp8_cache: bool = False + mla_kv_cache_dtype: str | None = None mla_index_topk: int | None = None # dllm @@ -424,6 +424,11 @@ class ModelConfig: # update cache config update_cache_config_func: Any = None + @property + def use_mla_fp8_cache(self): + """Whether MLA uses the DeepSeek-V3.2 FP8 cache layout.""" + return self.mla_kv_cache_dtype == 'fp8_ds_mla' + def get_head_size(self): """Get head size.""" return self.head_dim diff --git a/lmdeploy/pytorch/configurations/deepseek_v32.py b/lmdeploy/pytorch/configurations/deepseek_v32.py index cec2cf0781..35eac1815f 100644 --- a/lmdeploy/pytorch/configurations/deepseek_v32.py +++ b/lmdeploy/pytorch/configurations/deepseek_v32.py @@ -1,6 +1,8 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch +from lmdeploy.pytorch import envs as _envs + from .deepseek_v2 import DeepseekV2ModelConfigBuilder @@ -10,10 +12,11 @@ def _check_env_v32(device: str = 'cuda'): return # check cuda - try: - import fast_hadamard_transform # noqa: F401 - except ImportError: - raise ImportError('Deepseek V3.2 requires .') + if _envs.disable_dsa_indexer_fusion: + try: + import fast_hadamard_transform # noqa: F401 + except ImportError: + raise ImportError('Deepseek V3.2 requires .') try: import flash_mla # noqa: F401 @@ -29,7 +32,7 @@ class DeepseekV32ModelConfigBuilder(DeepseekV2ModelConfigBuilder): @classmethod def condition(cls, hf_config): """config.""" - return hf_config.model_type in ['deepseek_v32', 'glm_moe_dsa'] + return hf_config.model_type == 'deepseek_v32' @classmethod def build(cls, hf_config, model_path: str | None = None, **kwargs): @@ -40,7 +43,7 @@ def build(cls, hf_config, model_path: str | None = None, **kwargs): index_k_shape = ([hf_config.index_head_dim], torch.float8_e4m3fn) index_k_scale_shape = ([1], torch.float32) config.cache_shapes = [index_k_shape, index_k_scale_shape] - config.use_mla_fp8_cache = True + config.mla_kv_cache_dtype = 'bfloat16' config.mla_index_topk = hf_config.index_topk config.check_env_func = _check_env_v32 return config diff --git a/lmdeploy/pytorch/configurations/glm_moe_dsa.py b/lmdeploy/pytorch/configurations/glm_moe_dsa.py new file mode 100644 index 0000000000..7b9e09a179 --- /dev/null +++ b/lmdeploy/pytorch/configurations/glm_moe_dsa.py @@ -0,0 +1,37 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +from lmdeploy.pytorch import envs as _envs +from lmdeploy.utils import get_logger + +from .deepseek_v32 import DeepseekV32ModelConfigBuilder + +logger = get_logger('lmdeploy') + +class GlmMoeDsaModelConfigBuilder(DeepseekV32ModelConfigBuilder): + + @classmethod + def condition(cls, hf_config): + """config.""" + return hf_config.model_type == 'glm_moe_dsa' + + @classmethod + def build(cls, hf_config, model_path: str | None = None, **kwargs): + """build.""" + is_draft_model = kwargs.get('is_draft_model', False) + quantization_config = getattr(hf_config, 'quantization_config', None) + is_lmdeploy_patched_fp8 = (quantization_config is not None + and quantization_config.get('quant_method') == 'fp8' + and quantization_config.get('lmdeploy_patched', False)) + if _envs.fp8_moe_only and is_lmdeploy_patched_fp8: + quantization_config['fp8_quant_scope'] = 'moe_only' + logger.info('Enable fp8_quant_scope=moe_only for glm_moe_dsa because LMDEPLOY_FP8_MOE_ONLY=1 ' + 'and the FP8 quantization config is LMDeploy-synthesized.') + + if hf_config.qk_head_dim != hf_config.qk_nope_head_dim + hf_config.qk_rope_head_dim: + hf_config.qk_rope_head_dim = hf_config.qk_head_dim - hf_config.qk_nope_head_dim + hf_config.head_dim = hf_config.qk_rope_head_dim + + config = super().build(hf_config, model_path=model_path, **kwargs) + if is_draft_model: + hf_config.architectures[0] = 'GlmMoeDsaMTPModel' + return config diff --git a/lmdeploy/pytorch/engine/cache_engine.py b/lmdeploy/pytorch/engine/cache_engine.py index dbd5e2d22f..d2045d3e7a 100644 --- a/lmdeploy/pytorch/engine/cache_engine.py +++ b/lmdeploy/pytorch/engine/cache_engine.py @@ -150,9 +150,22 @@ def _get_kv_cache_dtype(model_config: ModelConfig): kv_cache_dtype = model_config.dtype if model_config.use_mla_fp8_cache: kv_cache_dtype = torch.float8_e4m3fn + elif model_config.mla_kv_cache_dtype == 'bfloat16': + kv_cache_dtype = torch.bfloat16 return kv_cache_dtype +def _update_mla_kv_cache_dtype(model_config: ModelConfig, cache_config: CacheConfig): + """Apply an explicit sparse MLA cache policy to the model config.""" + if model_config.mla_index_topk is None or cache_config.quant_policy == QuantPolicy.NONE: + return + if cache_config.quant_policy == QuantPolicy.FP8: + model_config.mla_kv_cache_dtype = 'fp8_ds_mla' + return + raise ValueError(f'Sparse MLA does not support quant_policy={cache_config.quant_policy}. ' + 'Use none/0 for BF16 or fp8/16 for FP8.') + + _FP8_CACHE_DTYPES = { QuantPolicy.FP8: torch.float8_e4m3fn, QuantPolicy.FP8_E5M2: torch.float8_e5m2, @@ -216,12 +229,13 @@ def __init__( self.tp_rank = tp_rank self.cache_config = cache_config self.model_config = model_config + _update_mla_kv_cache_dtype(model_config, cache_config) self.block_size = cache_config.kernel_block_size self.num_layers = model_config.num_layers self.kv_cache_dtype = _get_kv_cache_dtype(self.model_config) - if self.model_config.use_mla_fp8_cache: + if self.model_config.mla_index_topk is not None: cache_config.quant_policy = 0 if _is_fp8_quant_policy(cache_config.quant_policy): @@ -693,6 +707,8 @@ def get_cache_block_size(cls, cache_config: CacheConfig, model_config: ModelConf Return: int: Required memory size in bytes. """ + # Resolve the layout before sizing; CUDA graphs reuse the updated model config. + _update_mla_kv_cache_dtype(model_config, cache_config) mem_pool, _ = cls.allocate_caches( num_blocks=1, model_config=model_config, diff --git a/lmdeploy/pytorch/engine/config_builder.py b/lmdeploy/pytorch/engine/config_builder.py index 450ab58a60..bfa83a094b 100644 --- a/lmdeploy/pytorch/engine/config_builder.py +++ b/lmdeploy/pytorch/engine/config_builder.py @@ -119,6 +119,10 @@ def _build_draft_dist_ctx(dist_config): # TODO support tp > 1, ep > 1 for other methods if speculative_config.method == 'qwen3_5_mtp': draft_dist_config = dist_config + elif speculative_config.method == 'deepseek_mtp': + from lmdeploy.pytorch.transformers import config_from_pretrained + hf_config = config_from_pretrained(target_model, trust_remote_code=trust_remote_code) + draft_dist_config = dist_config if hf_config.model_type == 'glm_moe_dsa' else DistConfig() else: draft_dist_config = DistConfig() return draft_dist_config diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py index 5ccb268cb0..a09298ea77 100644 --- a/lmdeploy/pytorch/engine/engine.py +++ b/lmdeploy/pytorch/engine/engine.py @@ -59,7 +59,7 @@ class InferOutput: req_metrics: RequestMetrics = None # expert ids - routed_experts: torch.Tensor = None + routed_experts: np.ndarray = None # summed, unnormalized cross-entropy (NLL) of the input prompt ce_loss: float = None diff --git a/lmdeploy/pytorch/engine/engine_instance.py b/lmdeploy/pytorch/engine/engine_instance.py index a4deaf025a..d0eeb865f4 100644 --- a/lmdeploy/pytorch/engine/engine_instance.py +++ b/lmdeploy/pytorch/engine/engine_instance.py @@ -137,23 +137,41 @@ def __del__(self): """Destructor.""" self.engine.req_manager.senders.pop(self.req_sender.sender_id) - def _get_extra_outputs(self, resp: Response, num_all_ids: int): + def _get_extra_outputs(self, resp: Response, num_all_ids: int, gen_config: GenerationConfig = None): """Get extra outputs.""" outputs = dict(routed_experts=None) + if resp.type not in [ResponseType.FINISH, ResponseType.CANCEL]: + return outputs + + def _validate_num_tokens(name: str, data): + num_expected = num_all_ids - 1 + if data.shape[0] != num_expected: + logger.warning(f'Expected number of {name}: {num_expected}, but got {data.shape[0]}') + data = data[:num_expected] + return data + + def _trim_stop_token(data): + if (not self._enable_transfer_obj_ref or gen_config is None or gen_config.include_stop_str_in_output + or gen_config.ignore_eos): + return data + token_ids = resp.data.get('token_ids', None) if resp.data else None + stop_token_ids = gen_config.stop_token_ids or [] + if resp.type == ResponseType.FINISH and token_ids is not None and len(token_ids) > 0: + if token_ids[-1] in stop_token_ids: + return data[:-1] + return data + + def _maybe_transfer(data): + if not self._enable_transfer_obj_ref: + return data + import ray + data = _trim_stop_token(data) + return ray.get(_SHARED_STORE.put.remote(data)) + routed_experts = resp.data.get('routed_experts', None) if resp.data else None - if routed_experts is not None and resp.type in [ResponseType.FINISH, ResponseType.CANCEL]: - if self._enable_transfer_obj_ref: - import ray - # validate experts - num_expected_experts = num_all_ids - 1 - if routed_experts.shape[0] != num_expected_experts: - logger.warning(f'Expected number of routed_experts: {num_expected_experts}, ' - f'but got {routed_experts.shape[0]}') - routed_experts = routed_experts[:num_expected_experts] - key = ray.get(_SHARED_STORE.put.remote(routed_experts)) - outputs['routed_experts'] = key - else: - outputs['routed_experts'] = routed_experts + if routed_experts is not None: + routed_experts = _validate_num_tokens('routed_experts', routed_experts) + outputs['routed_experts'] = _maybe_transfer(routed_experts) return outputs async def _async_try_add_session(self, session_id: int): @@ -255,7 +273,7 @@ async def async_stream_infer(self, num_ids = len(token_ids) num_all_ids = prompt_ids_len + output_offset + num_ids - extra_outputs = self._get_extra_outputs(resp, num_all_ids) + extra_outputs = self._get_extra_outputs(resp, num_all_ids, gen_config=gen_config) routed_experts = extra_outputs.get('routed_experts', None) logger.debug(f'session[{session_id}] finish: num_out_ids={num_ids}.') diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py index cad34d4618..aaf8faa6cb 100644 --- a/lmdeploy/pytorch/envs.py +++ b/lmdeploy/pytorch/envs.py @@ -185,6 +185,7 @@ def _patched_get_env( # model format scale_fmt = os.getenv('LMDEPLOY_SCALE_FMT', None) fp8_moe_only = env_to_bool('LMDEPLOY_FP8_MOE_ONLY', False) + disable_dsa_indexer_fusion = env_to_bool('LMDEPLOY_DISABLE_DSA_INDEXER_FUSION', False) # repetition check repetition_window_size = env_to_int('LMDEPLOY_REPETITION_WINDOW_SIZE', 1024) diff --git a/lmdeploy/pytorch/kernels/cuda/blocked_fp8_fused_moe.py b/lmdeploy/pytorch/kernels/cuda/blocked_fp8_fused_moe.py index 06ea42a24b..4804df66a5 100644 --- a/lmdeploy/pytorch/kernels/cuda/blocked_fp8_fused_moe.py +++ b/lmdeploy/pytorch/kernels/cuda/blocked_fp8_fused_moe.py @@ -87,13 +87,13 @@ def fused_moe_blocked_f8_kernel( offs_am = sid // top_k else: offs_am = offs_sid - a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) # deepseek has 160 experts, exp index would overflow int32 exp_id = exp_id.to(tl.int64) exp_off = stride_be * exp_id + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) b_ptrs = B + exp_off + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) offs_bsn = pid_n * BLOCK_SIZE_N // group_bn @@ -310,12 +310,12 @@ def fused_moe_blocked_f8_compact_kernel( offs_am = sid // top_k else: offs_am = offs_sid - a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) local_exp = local_exp.to(tl.int64) exp_off = stride_be * local_exp + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) b_ptrs = B + exp_off + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) offs_bsn = pid_n * BLOCK_SIZE_N // group_bn @@ -503,6 +503,19 @@ def _compact_blocked_fp8_moe_config(num_routes: int, num_experts: int): return dict(block_m=block_m, block_n=128, num_warps=4, num_stages=3) +def _compact_blocked_fp8_moe_both_config(num_routes: int, num_experts: int, gate_out_features: int): + """Choose compact configs measured for both GLM-5.2 projections.""" + avg_routes = triton.cdiv(num_routes, num_experts) + # Lower TP degrees have wider local projections and favor the larger M tile. + if gate_out_features > 512 or avg_routes > 16: + block_m, block_n = 64, 128 + elif avg_routes > 8: + block_m, block_n = 32, 128 + else: + block_m, block_n = 16, 64 + return dict(block_m=block_m, block_n=block_n, num_warps=4, num_stages=3) + + def _blocked_fp8_moe_cta_estimates(num_tokens: int, num_routes: int, num_experts: int, local_experts: int, out_features: int): """Estimate origin and compact down-projection CTA counts.""" @@ -584,6 +597,13 @@ def _should_use_compact_blocked_fp8_moe_down_by_shape(num_tokens: int, num_route return origin_ctas >= min_cta_ratio * compact_ctas +def _should_use_compact_blocked_fp8_moe_both_by_shape(num_tokens: int, num_routes: int, num_experts: int, + local_experts: int): + """Use compact scheduling in the measured full-expert density window.""" + avg_routes = triton.cdiv(num_routes, num_experts) + return num_tokens >= 128 and num_experts == 256 and local_experts == num_experts and 4 <= avg_routes <= 48 + + def _should_use_compact_blocked_fp8_moe_down(input: torch.Tensor, input_scale: torch.Tensor, w1: torch.Tensor, w1_scale: torch.Tensor, w2: torch.Tensor, w2_scale: torch.Tensor, topk_ids: torch.Tensor, num_experts: int, expert_offset: int = 0): @@ -625,39 +645,68 @@ def fused_moe_blocked_fp8(input: torch.Tensor, topk_weights = _renormalize(topk_weights, renormalize) gate_moe_cfg, down_moe_cfg = _origin_blocked_fp8_moe_configs(M, topk_ids.numel(), num_experts, E) - use_compact_down = _should_use_compact_blocked_fp8_moe_down(input, input_scale, w1, w1_scale, w2, w2_scale, - topk_ids, num_experts, expert_offset) - if use_compact_down: - compact_down_cfg = _compact_blocked_fp8_moe_config(topk_ids.numel(), num_experts) + supports_compact = _supports_compact_blocked_fp8_moe(input, input_scale, w1, w1_scale, w2, w2_scale, topk_ids, + num_experts, expert_offset) + use_compact_both = supports_compact and _should_use_compact_blocked_fp8_moe_both_by_shape( + M, topk_ids.numel(), num_experts, E) + use_compact_down = (not use_compact_both and supports_compact + and _should_use_compact_blocked_fp8_moe_down_by_shape(M, topk_ids.numel(), num_experts, E, + w2.size(1))) + if use_compact_both: + compact_moe_cfg = _compact_blocked_fp8_moe_both_config(topk_ids.numel(), num_experts, w1.size(1)) + elif use_compact_down: + compact_moe_cfg = _compact_blocked_fp8_moe_config(topk_ids.numel(), num_experts) + + if use_compact_both or use_compact_down: sorted_idx, exp_start, exp_end, block_end, block_expert_ids, block_offsets = _get_sorted_idx_blocks( topk_ids, num_experts, E, expert_offset, - compact_down_cfg['block_m'], + compact_moe_cfg['block_m'], ) else: sorted_idx, exp_start, exp_end = _get_sorted_idx(topk_ids, num_experts) intermediate_cache1 = _make_intermediate((M, topk, N), dtype=out_dtype, device=device, zeros=not full_exp) # gate and up - fused_moe_blocked_fp8_kernel_launcher( - input, - input_scale, - w1, - w1_scale, - intermediate_cache1, - sorted_idx=sorted_idx, - exp_start=exp_start, - exp_end=exp_end, - bias=w1_bias, - top_k=topk, - num_tokens=M, - expert_offset=expert_offset, - reindex_a=True, - reindex_c=False, - **gate_moe_cfg, - ) + if use_compact_both: + fused_moe_blocked_fp8_compact_kernel_launcher( + input, + input_scale, + w1, + w1_scale, + intermediate_cache1, + sorted_idx=sorted_idx, + exp_end=exp_end, + block_end=block_end, + block_expert_ids=block_expert_ids, + block_offsets=block_offsets, + bias=w1_bias, + top_k=topk, + expert_offset=expert_offset, + reindex_a=True, + reindex_c=False, + **compact_moe_cfg, + ) + else: + fused_moe_blocked_fp8_kernel_launcher( + input, + input_scale, + w1, + w1_scale, + intermediate_cache1, + sorted_idx=sorted_idx, + exp_start=exp_start, + exp_end=exp_end, + bias=w1_bias, + top_k=topk, + num_tokens=M, + expert_offset=expert_offset, + reindex_a=True, + reindex_c=False, + **gate_moe_cfg, + ) # activate intermediate_cache1 = intermediate_cache1.flatten(0, -2) @@ -670,7 +719,7 @@ def fused_moe_blocked_fp8(input: torch.Tensor, intermediate_cache2 = _make_intermediate((M, topk, w2.shape[1]), dtype=out_dtype, device=device, zeros=not full_exp) # down - if use_compact_down: + if use_compact_both or use_compact_down: fused_moe_blocked_fp8_compact_kernel_launcher( gate_cache, gate_scale, @@ -687,7 +736,7 @@ def fused_moe_blocked_fp8(input: torch.Tensor, expert_offset=expert_offset, reindex_a=False, reindex_c=True, - **compact_down_cfg, + **compact_moe_cfg, ) else: fused_moe_blocked_fp8_kernel_launcher( diff --git a/lmdeploy/pytorch/kernels/cuda/ds_index.py b/lmdeploy/pytorch/kernels/cuda/ds_index.py index 5efa1f5b3d..5f60911ab7 100644 --- a/lmdeploy/pytorch/kernels/cuda/ds_index.py +++ b/lmdeploy/pytorch/kernels/cuda/ds_index.py @@ -6,6 +6,65 @@ from .utils import get_device_props +@triton.jit +def _dense_index_kernel( + q_seqlens, + k_seqlens, + out, + stride_out_m: tl.constexpr, + stride_out_k: tl.constexpr, + max_q_seqlen: tl.constexpr, + topk: tl.constexpr, + fill: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """Build the exact top-k result when every valid K position is selected.""" + row = tl.program_id(0) + block = tl.program_id(1) + batch = row // max_q_seqlen + query = row % max_q_seqlen + q_seqlen = tl.load(q_seqlens + batch) + k_seqlen = tl.load(k_seqlens + batch) + causal_k_seqlen = k_seqlen - q_seqlen + query + 1 + + offs = block * BLOCK_K + tl.arange(0, BLOCK_K) + valid = (query < q_seqlen) & (offs < causal_k_seqlen) + values = tl.where(valid, offs, fill) + tl.store(out + row * stride_out_m + offs * stride_out_k, + values, + mask=offs < topk) + + +def dense_index(q_seqlens: torch.Tensor, + k_seqlens: torch.Tensor, + max_q_seqlen: int, + topk: int, + fill: int = -1) -> torch.Tensor: + """Return identity top-k rows for a decoding packet. + + This is valid when every request's causal KV length is no greater than + ``topk``: selecting the largest ``topk`` scores must then select every + visible position, independent of the scores. + """ + assert q_seqlens.is_cuda and k_seqlens.is_cuda + assert q_seqlens.numel() == k_seqlens.numel() + num_rows = q_seqlens.numel() * max_q_seqlen + out = torch.empty((num_rows, topk), + dtype=torch.int32, + device=q_seqlens.device) + block_k = min(triton.next_power_of_2(topk), 256) + grid = (num_rows, triton.cdiv(topk, block_k)) + _dense_index_kernel[grid](q_seqlens, + k_seqlens, + out, + *out.stride(), + max_q_seqlen=max_q_seqlen, + topk=topk, + fill=fill, + BLOCK_K=block_k) + return out + + @triton.jit(do_not_specialize=['max_q_seqlen', 'num_split']) def _fp8_index_kernel( q_ptr, diff --git a/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py b/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py new file mode 100644 index 0000000000..5c036d497f --- /dev/null +++ b/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py @@ -0,0 +1,388 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Fused Q/K preprocessing and cache writes for the DSA FP8 indexer.""" + +import torch +import triton +import triton.language as tl +from torch import Tensor + +from .blocked_gemm_fp8 import fast_round_scale + +# The unfused indexer applies the same normalized Hadamard transform to Q and +# K. Their dot product is unchanged because H @ H.T is identity, so this path +# omits both transforms and prepares the tensors directly for FP8 indexing. +# FP8 rounding makes the two paths numerically close rather than bit-identical. + + +@triton.jit +def _apply_rope_first( + x, + x_pair, + cos, + sin, + feat_off, + rope_dim: tl.constexpr, + rope_interleaved: tl.constexpr, +): + """Apply RoPE to the leading dimensions for both supported pair layouts.""" + prod_cos = x * cos + 0 + prod_sin = x_pair * sin + 0 + if rope_interleaved: + is_lower = (feat_off & 1) == 0 + else: + is_lower = feat_off < rope_dim // 2 + rotated = tl.where(is_lower, prod_cos - prod_sin, prod_cos + prod_sin) + return tl.where(feat_off < rope_dim, rotated, x) + + +@triton.jit +def _prepare_dsa_indexer_q_kernel( + Q, + Weights, + Cos, + Sin, + QOut, + QScaleOut, + num_heads: tl.constexpr, + score_scale: tl.constexpr, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, + stride_qt: tl.constexpr, + stride_qh: tl.constexpr, + stride_qd: tl.constexpr, + stride_wt: tl.constexpr, + stride_wh: tl.constexpr, + stride_cs: tl.constexpr, + stride_cd: tl.constexpr, + stride_ot: tl.constexpr, + stride_oh: tl.constexpr, + stride_od: tl.constexpr, + stride_st: tl.constexpr, + stride_sh: tl.constexpr, + head_dim: tl.constexpr, + rope_dim: tl.constexpr, + rope_interleaved: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row_id = tl.program_id(0) + token_id = row_id // num_heads + head_id = row_id % num_heads + + feat_off = tl.arange(0, BLOCK_D) + feat_mask = feat_off < head_dim + q_ptr = Q + token_id * stride_qt + head_id * stride_qh + x = tl.load(q_ptr + feat_off * stride_qd, mask=feat_mask, other=0.0) + + # GLM pairs adjacent features (0, 1), (2, 3), ... and shares one frequency + # per pair. DeepSeek uses NeoX pairs (0, D/2), (1, D/2 + 1), ... and its + # cos/sin table already contains the corresponding low/high frequencies. + if rope_interleaved: + pair_off = feat_off ^ 1 + freq_off = feat_off // 2 + else: + half_rope_dim = rope_dim // 2 + pair_off = tl.where(feat_off < half_rope_dim, feat_off + half_rope_dim, feat_off - half_rope_dim) + freq_off = feat_off + x_pair = tl.load(q_ptr + pair_off * stride_qd, mask=feat_mask, other=0.0) + freq_mask = feat_off < rope_dim + cos = tl.load(Cos + token_id * stride_cs + freq_off * stride_cd, mask=freq_mask, other=1.0) + sin = tl.load(Sin + token_id * stride_cs + freq_off * stride_cd, mask=freq_mask, other=0.0) + x = _apply_rope_first(x, x_pair, cos, sin, feat_off, rope_dim, rope_interleaved) + + # Match the unfused path: RoPE materializes BF16 before dynamic FP8 quantization. + x = x.to(tl.bfloat16) + abs_max = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-6).to(tl.float32) + scale = fast_round_scale(abs_max, 1.0 / fp8_max) + out = tl.clamp(x.to(tl.float32) * (1.0 / scale), fp8_min, fp8_max).to(QOut.dtype.element_ty) + + out_ptr = QOut + token_id * stride_ot + head_id * stride_oh + tl.store(out_ptr + feat_off * stride_od, out, mask=feat_mask) + weight = tl.load(Weights + token_id * stride_wt + head_id * stride_wh).to(tl.float32) + # Fold the head gate and attention scale into Q's FP8 scale. + tl.store(QScaleOut + token_id * stride_st + head_id * stride_sh, scale * weight * score_scale) + + +@triton.jit +def _prepare_dsa_indexer_k_kernel( + K, + NormWeight, + NormBias, + Cos, + Sin, + KOut, + eps: tl.constexpr, + stride_kt: tl.constexpr, + stride_kd: tl.constexpr, + stride_cs: tl.constexpr, + stride_cd: tl.constexpr, + stride_ot: tl.constexpr, + stride_od: tl.constexpr, + head_dim: tl.constexpr, + rope_dim: tl.constexpr, + rope_interleaved: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token_id = tl.program_id(0) + feat_off = tl.arange(0, BLOCK_D) + feat_mask = feat_off < head_dim + k_ptr = K + token_id * stride_kt + + x = tl.load(k_ptr + feat_off * stride_kd, mask=feat_mask, other=0.0).to(tl.float32) + mean = tl.sum(x, axis=0) / head_dim + centered = x - mean + inv_std = tl.rsqrt(tl.sum(centered * centered, axis=0) / head_dim + eps) + weight = tl.load(NormWeight + feat_off, mask=feat_mask, other=0.0).to(tl.float32) + bias = tl.load(NormBias + feat_off, mask=feat_mask, other=0.0).to(tl.float32) + x = (centered * inv_std * weight + bias).to(tl.bfloat16) + + # Recreate the normalized RoPE partner from raw K with the same row + # statistics instead of materializing the complete LayerNorm output. + if rope_interleaved: + pair_off = feat_off ^ 1 + freq_off = feat_off // 2 + else: + half_rope_dim = rope_dim // 2 + pair_off = tl.where(feat_off < half_rope_dim, feat_off + half_rope_dim, feat_off - half_rope_dim) + freq_off = feat_off + x_pair_raw = tl.load(k_ptr + pair_off * stride_kd, mask=feat_mask, other=0.0).to(tl.float32) + pair_weight = tl.load(NormWeight + pair_off, mask=feat_mask, other=0.0).to(tl.float32) + pair_bias = tl.load(NormBias + pair_off, mask=feat_mask, other=0.0).to(tl.float32) + x_pair = ((x_pair_raw - mean) * inv_std * pair_weight + pair_bias).to(tl.bfloat16) + + freq_mask = feat_off < rope_dim + cos = tl.load(Cos + token_id * stride_cs + freq_off * stride_cd, mask=freq_mask, other=1.0) + sin = tl.load(Sin + token_id * stride_cs + freq_off * stride_cd, mask=freq_mask, other=0.0) + x = _apply_rope_first(x, x_pair, cos, sin, feat_off, rope_dim, rope_interleaved) + tl.store(KOut + token_id * stride_ot + feat_off * stride_od, x.to(tl.bfloat16), mask=feat_mask) + + +@triton.jit +def _prepare_dsa_indexer_k_cache_kernel( + K, + NormWeight, + NormBias, + Cos, + Sin, + KCache, + KSCache, + CuSeqLenQ, + KVSeqLens, + BlockOffsets, + eps: tl.constexpr, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, + stride_kt: tl.constexpr, + stride_kd: tl.constexpr, + stride_cs: tl.constexpr, + stride_cd: tl.constexpr, + stride_kcb: tl.constexpr, + stride_kcs: tl.constexpr, + stride_kcd: tl.constexpr, + stride_ksb: tl.constexpr, + stride_kss: tl.constexpr, + stride_boff: tl.constexpr, + block_size: tl.constexpr, + head_dim: tl.constexpr, + rope_dim: tl.constexpr, + rope_interleaved: tl.constexpr, + BLOCK_D: tl.constexpr, +): + q_id = tl.program_id(0) + batch_id = tl.program_id(1) + q_start = tl.load(CuSeqLenQ + batch_id) + q_seqlen = tl.load(CuSeqLenQ + batch_id + 1) - q_start + if q_id >= q_seqlen: + return + + kv_seqlen = tl.load(KVSeqLens + batch_id) + history_seqlen = kv_seqlen - q_seqlen + # The input contains only this step's tokens; append them after the cached + # history and translate the logical position through the page table. + kv_pos = history_seqlen + q_id + logical_block = kv_pos // block_size + page_off = kv_pos % block_size + physical_block = tl.load(BlockOffsets + batch_id * stride_boff + logical_block).to(tl.int64) + token_id = q_start + q_id + + feat_off = tl.arange(0, BLOCK_D) + feat_mask = feat_off < head_dim + k_ptr = K + token_id * stride_kt + x = tl.load(k_ptr + feat_off * stride_kd, mask=feat_mask, other=0.0).to(tl.float32) + mean = tl.sum(x, axis=0) / head_dim + centered = x - mean + inv_std = tl.rsqrt(tl.sum(centered * centered, axis=0) / head_dim + eps) + weight = tl.load(NormWeight + feat_off, mask=feat_mask, other=0.0).to(tl.float32) + bias = tl.load(NormBias + feat_off, mask=feat_mask, other=0.0).to(tl.float32) + x = (centered * inv_std * weight + bias).to(tl.bfloat16) + + # Recreate the normalized partner needed by RoPE without a temporary K + # tensor, then quantize directly into its paged-cache destination. + if rope_interleaved: + pair_off = feat_off ^ 1 + freq_off = feat_off // 2 + else: + half_rope_dim = rope_dim // 2 + pair_off = tl.where(feat_off < half_rope_dim, feat_off + half_rope_dim, feat_off - half_rope_dim) + freq_off = feat_off + x_pair_raw = tl.load(k_ptr + pair_off * stride_kd, mask=feat_mask, other=0.0).to(tl.float32) + pair_weight = tl.load(NormWeight + pair_off, mask=feat_mask, other=0.0).to(tl.float32) + pair_bias = tl.load(NormBias + pair_off, mask=feat_mask, other=0.0).to(tl.float32) + x_pair = ((x_pair_raw - mean) * inv_std * pair_weight + pair_bias).to(tl.bfloat16) + + freq_mask = feat_off < rope_dim + cos = tl.load(Cos + token_id * stride_cs + freq_off * stride_cd, mask=freq_mask, other=1.0) + sin = tl.load(Sin + token_id * stride_cs + freq_off * stride_cd, mask=freq_mask, other=0.0) + x = _apply_rope_first(x, x_pair, cos, sin, feat_off, rope_dim, rope_interleaved) + x = x.to(tl.bfloat16) + + abs_max = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-6).to(tl.float32) + # Keep the ue8m0 scale rounding used by the FP8 index kernel. + scale = fast_round_scale(abs_max, 1.0 / fp8_max) + out = tl.clamp(x.to(tl.float32) * (1.0 / scale), fp8_min, fp8_max).to(KCache.dtype.element_ty) + cache_ptr = KCache + physical_block * stride_kcb + page_off * stride_kcs + tl.store(cache_ptr + feat_off * stride_kcd, out, mask=feat_mask) + tl.store(KSCache + physical_block * stride_ksb + page_off * stride_kss, scale) + + +def prepare_dsa_indexer_q( + q: Tensor, + weights: Tensor, + cos: Tensor, + sin: Tensor, + score_scale: float, + out_dtype: torch.dtype, + rope_interleaved: bool, +) -> tuple[Tensor, Tensor]: + """Fuse RoPE, FP8 quantization, and head-gate scaling.""" + assert q.dtype == torch.bfloat16 and q.dim() == 3 + assert q.size(-1) == 128 and cos.size(-1) == 64 and sin.shape == cos.shape + assert q.shape[:-1] == weights.shape and q.size(0) == cos.size(0) + assert q.stride(-1) == 1 and weights.stride(-1) == 1 + + q_out = torch.empty_like(q, dtype=out_dtype) + q_scale = q.new_empty(q.shape[:-1], dtype=torch.float32) + finfo = torch.finfo(out_dtype) + grid = (q.size(0) * q.size(1), ) + _prepare_dsa_indexer_q_kernel[grid](q, + weights, + cos, + sin, + q_out, + q_scale, + num_heads=q.size(1), + score_scale=score_scale, + fp8_min=finfo.min, + fp8_max=finfo.max, + stride_qt=q.stride(0), + stride_qh=q.stride(1), + stride_qd=q.stride(2), + stride_wt=weights.stride(0), + stride_wh=weights.stride(1), + stride_cs=cos.stride(0), + stride_cd=cos.stride(1), + stride_ot=q_out.stride(0), + stride_oh=q_out.stride(1), + stride_od=q_out.stride(2), + stride_st=q_scale.stride(0), + stride_sh=q_scale.stride(1), + head_dim=128, + rope_dim=64, + rope_interleaved=rope_interleaved, + BLOCK_D=128, + num_warps=4, + num_stages=1) + return q_out, q_scale + + +def prepare_dsa_indexer_k( + k: Tensor, + norm_weight: Tensor, + norm_bias: Tensor, + cos: Tensor, + sin: Tensor, + eps: float, + rope_interleaved: bool, +) -> Tensor: + """Reference fused LayerNorm and RoPE K preparation.""" + assert k.dtype == torch.bfloat16 and k.dim() == 2 and k.size(-1) == 128 + assert norm_weight.shape == norm_bias.shape == (128, ) + assert cos.shape == sin.shape == (k.size(0), 64) + k_out = torch.empty_like(k) + _prepare_dsa_indexer_k_kernel[(k.size(0), )](k, + norm_weight, + norm_bias, + cos, + sin, + k_out, + eps=eps, + stride_kt=k.stride(0), + stride_kd=k.stride(1), + stride_cs=cos.stride(0), + stride_cd=cos.stride(1), + stride_ot=k_out.stride(0), + stride_od=k_out.stride(1), + head_dim=128, + rope_dim=64, + rope_interleaved=rope_interleaved, + BLOCK_D=128, + num_warps=4, + num_stages=1) + return k_out + + +def prepare_dsa_indexer_k_cache( + k: Tensor, + norm_weight: Tensor, + norm_bias: Tensor, + cos: Tensor, + sin: Tensor, + k_cache: Tensor, + k_s_cache: Tensor, + cu_seqlen_q: Tensor, + kv_seqlens: Tensor, + block_offsets: Tensor, + max_q_seqlen: int, + eps: float, + rope_interleaved: bool, +) -> None: + """Fuse K LayerNorm, RoPE, FP8 quantization, and cache fill.""" + assert k.dtype == torch.bfloat16 and k.dim() == 2 and k.size(-1) == 128 + assert k_cache.dim() == 3 and k_cache.size(-1) == 128 + assert k_s_cache.dim() == 2 and k_s_cache.shape == k_cache.shape[:2] + assert norm_weight.shape == norm_bias.shape == (128, ) + assert cos.shape == sin.shape == (k.size(0), 64) + assert cu_seqlen_q.numel() == kv_seqlens.numel() + 1 + + block_offsets = block_offsets.contiguous() + finfo = torch.finfo(k_cache.dtype) + grid = (max_q_seqlen, kv_seqlens.numel()) + _prepare_dsa_indexer_k_cache_kernel[grid](k, + norm_weight, + norm_bias, + cos, + sin, + k_cache, + k_s_cache, + cu_seqlen_q, + kv_seqlens, + block_offsets, + eps=eps, + fp8_min=finfo.min, + fp8_max=finfo.max, + stride_kt=k.stride(0), + stride_kd=k.stride(1), + stride_cs=cos.stride(0), + stride_cd=cos.stride(1), + stride_kcb=k_cache.stride(0), + stride_kcs=k_cache.stride(1), + stride_kcd=k_cache.stride(2), + stride_ksb=k_s_cache.stride(0), + stride_kss=k_s_cache.stride(1), + stride_boff=block_offsets.stride(0), + block_size=k_cache.size(1), + head_dim=128, + rope_dim=64, + rope_interleaved=rope_interleaved, + BLOCK_D=128, + num_warps=4, + num_stages=1) diff --git a/lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py b/lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py index 24a3d674e9..d29d7c00da 100644 --- a/lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py +++ b/lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py @@ -38,9 +38,9 @@ def _ordered_fp32_key(score): bits = T.reinterpret(score, T.uint32) sign_mask = T.cast(2147483648, T.uint32) all_ones = T.cast(4294967295, T.uint32) - return T.if_then_else(T.bitwise_and(bits, sign_mask) == T.cast(0, T.uint32), - T.bitwise_xor(bits, sign_mask), - T.bitwise_xor(bits, all_ones)) + return T.if_then_else( + T.bitwise_and(bits, sign_mask) == T.cast(0, T.uint32), + T.bitwise_xor(bits, sign_mask), T.bitwise_xor(bits, all_ones)) def is_sparse_index_topk_supported(k: int) -> bool: @@ -59,18 +59,20 @@ def _sparse_index_topk_byte_radix_kernel(top_k: int, @T.prim_func def sparse_index_topk_byte_radix_kernel_( - Scores: T.StridedTensor[(num_tokens, score_width), (score_stride, 1), T.float32], - Seqlens: T.Tensor[(num_tokens,), T.int32], + Scores: T.StridedTensor[(num_tokens, score_width), (score_stride, 1), + T.float32], + Seqlens: T.Tensor[(num_tokens, ), T.int32], Out: T.Tensor[(num_tokens, top_k), T.int32], ): _ = score_stride with T.Kernel(num_tokens, threads=threads) as row: tidx = T.get_thread_binding(0) - histogram = T.alloc_shared((_RADIX_SIZE,), T.int32) - shared_state = T.alloc_shared((2,), T.int32) + histogram = T.alloc_shared((_RADIX_SIZE, ), T.int32) + shared_state = T.alloc_shared((2, ), T.int32) raw_seqlen = Seqlens[row] - seqlen = T.if_then_else(raw_seqlen < score_width, raw_seqlen, score_width) + seqlen = T.if_then_else(raw_seqlen < score_width, raw_seqlen, + score_width) # If the real row length fits inside top_k, every valid position # is selected. @@ -78,7 +80,7 @@ def sparse_index_topk_byte_radix_kernel_( for i in T.Parallel(top_k): Out[row, i] = T.if_then_else(i < seqlen, i, fill) else: - # Byte-radix threshold search. Each round fixes one byte of the + # Byte-radix threshold search. Each round fixes one byte of the # fp32-order key and narrows the selected prefix until the final # threshold key is known. prefix_key = T.alloc_var(T.uint32) @@ -100,8 +102,10 @@ def sparse_index_topk_byte_radix_kernel_( while pos < seqlen: key = _ordered_fp32_key(Scores[row, pos]) if T.bitwise_and(key, prefix_mask) == prefix_key: - bin_u32 = T.bitwise_and(key >> shift, T.cast(255, T.uint32)) - T.atomic_add(histogram[T.cast(bin_u32, T.int32)], 1) + bin_u32 = T.bitwise_and(key >> shift, + T.cast(255, T.uint32)) + T.atomic_add(histogram[T.cast(bin_u32, T.int32)], + 1) pos += threads T.sync_threads() @@ -120,7 +124,9 @@ def sparse_index_topk_byte_radix_kernel_( T.sync_threads() threshold_bin = shared_state[_STATE_SELECTED_BIN] - prefix_key = T.bitwise_or(prefix_key, T.cast(threshold_bin, T.uint32) << shift) + prefix_key = T.bitwise_or( + prefix_key, + T.cast(threshold_bin, T.uint32) << shift) prefix_mask = T.bitwise_or(prefix_mask, byte_mask) rank -= shared_state[_STATE_COUNT_ABOVE] @@ -133,7 +139,7 @@ def sparse_index_topk_byte_radix_kernel_( shared_state[_STATE_EMIT_EQ_COUNT] = 0 T.sync_threads() - out_pos_buf = T.alloc_local((1,), T.int32) + out_pos_buf = T.alloc_local((1, ), T.int32) # First emit all scores strictly greater than the threshold. pos_emit_gt = T.alloc_var(T.int32) pos_emit_gt = tidx @@ -141,7 +147,9 @@ def sparse_index_topk_byte_radix_kernel_( key = _ordered_fp32_key(Scores[row, pos_emit_gt]) if key > threshold_key: out_pos_buf[0] = T.atomic_add( - shared_state[_STATE_EMIT_GT_COUNT], 1, return_prev=True) + shared_state[_STATE_EMIT_GT_COUNT], + 1, + return_prev=True) if out_pos_buf[0] < top_k: Out[row, out_pos_buf[0]] = pos_emit_gt pos_emit_gt += threads @@ -159,7 +167,9 @@ def sparse_index_topk_byte_radix_kernel_( key = _ordered_fp32_key(Scores[row, pos_emit_eq]) if key == threshold_key: out_pos_buf[0] = gt_count + T.atomic_add( - shared_state[_STATE_EMIT_EQ_COUNT], 1, return_prev=True) + shared_state[_STATE_EMIT_EQ_COUNT], + 1, + return_prev=True) if out_pos_buf[0] < top_k: Out[row, out_pos_buf[0]] = pos_emit_eq pos_emit_eq += threads @@ -176,7 +186,7 @@ def sparse_index_topk(scores: torch.Tensor, sorted: bool = False) -> torch.Tensor: """Return top-k score indices for padded sparse-index score rows. - The returned ids are packed but not score-sorted. Sparse attention + The returned ids are packed but not score-sorted. Sparse attention consumes them as a set of valid KV positions; avoiding final sorting is the point of this selector. Rows shorter than ``k`` are padded with ``fill``. """ @@ -185,19 +195,24 @@ def sparse_index_topk(scores: torch.Tensor, if sorted: raise ValueError('sparse_index_topk does not support sorted output.') if not is_sparse_index_topk_supported(k): - raise ValueError(f'sparse_index_topk supports k in {_SUPPORTED_TOPK}, got {k}.') + raise ValueError( + f'sparse_index_topk supports k in {_SUPPORTED_TOPK}, got {k}.') assert scores.is_cuda and q_seqlens.is_cuda and kv_seqlens.is_cuda assert scores.dtype == torch.float32 assert q_seqlens.dtype in (torch.int32, torch.int64) - assert kv_seqlens.dtype == torch.int32 + assert kv_seqlens.dtype in (torch.int32, torch.int64) assert scores.stride(-1) == 1 + kv_seqlens = kv_seqlens.to(torch.int32) num_tokens = scores.size(0) if num_tokens != kv_seqlens.size(0): - kv_seqlens = torch.repeat_interleave(kv_seqlens, q_seqlens, output_size=num_tokens) + kv_seqlens = torch.repeat_interleave(kv_seqlens, + q_seqlens, + output_size=num_tokens) else: kv_seqlens = kv_seqlens.contiguous() out = torch.empty((num_tokens, k), device=scores.device, dtype=torch.int32) - _sparse_index_topk_byte_radix_kernel(k, fill, _THREADS)(scores, kv_seqlens, out) + _sparse_index_topk_byte_radix_kernel(k, fill, _THREADS)(scores, kv_seqlens, + out) return out diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py index 39d68dd854..5be05a71b1 100644 --- a/lmdeploy/pytorch/model_inputs.py +++ b/lmdeploy/pytorch/model_inputs.py @@ -366,7 +366,12 @@ def new( # position ids attention_mask, position_ids = cls.get_mask_and_position_ids(inputs) - q_start_loc = q_seqlens.cumsum(0) - q_seqlens + if inputs.max_q_seqlen == 1: + # Active decode sequences all contain one query token, so their + # flattened offsets are simply the batch indices. + q_start_loc = torch.arange(q_seqlens.numel(), dtype=q_seqlens.dtype, device=q_seqlens.device) + else: + q_start_loc = q_seqlens.cumsum(0) - q_seqlens # seq_len + history_length kv_seqlens = q_seqlens + history_seqlens @@ -426,11 +431,13 @@ def get_mask_and_position_ids(cls, inputs: ModelInputs): target_position_ids = inputs.target_position_ids # decoding if max_q_seqlen == 1: - attention_mask = torch.ones_like(q_seqlens)[:, None] + # Positive sequence lengths with max_q_seqlen == 1 are already an + # all-ones mask; use views to avoid two decode-only materializations. + attention_mask = q_seqlens[:, None] if target_position_ids is not None: position_ids = target_position_ids else: - position_ids = history_seqlens.unsqueeze(0).clone() + position_ids = history_seqlens.unsqueeze(0) return attention_mask, position_ids num_tokens = inputs.input_ids.numel() diff --git a/lmdeploy/pytorch/models/deepseek_v2.py b/lmdeploy/pytorch/models/deepseek_v2.py index 087d744bf9..574a3a8508 100644 --- a/lmdeploy/pytorch/models/deepseek_v2.py +++ b/lmdeploy/pytorch/models/deepseek_v2.py @@ -358,6 +358,11 @@ def __init__(self, batch: int, in_features: int, out_features: int, dtype: torch def _update_batch(self, batch: int): """Update out features.""" + dist_config = get_dist_manager().current_config() + if dist_config.dp > 1: + # MLA Q projections use dp_disable_tp=True, so DP mode keeps + # q_nope full-head; absorb BMM weights must use the same layout. + return batch world_size, _ = get_tp_world_rank('attn') batch = batch // world_size return batch @@ -368,8 +373,10 @@ def create_weight(self, batch: int, in_features: int, out_features: int, dtype: def weight_loader(self, param: nn.Parameter, weight: torch.Tensor): """Weight loader.""" - world_size, rank = get_tp_world_rank('attn') - weight = weight.chunk(world_size, 0)[rank] + dist_config = get_dist_manager().current_config() + if dist_config.dp == 1: + world_size, rank = get_tp_world_rank('attn') + weight = weight.chunk(world_size, 0)[rank] param.data.copy_(weight) def forward(self, x: torch.Tensor, output: torch.Tensor): @@ -636,7 +643,7 @@ def _postprocess_topk_weight(self, topk_weight: torch.Tensor): topk_weight = topk_weight * self.routed_scaling_factor return topk_weight - def forward(self, hidden_states: torch.Tensor): + def forward(self, hidden_states: torch.Tensor, routed_experts: torch.Tensor = None): """forward.""" router_logits = F.linear(hidden_states.to(self.weight.dtype), self.weight) if self.fake_eplb: @@ -666,6 +673,9 @@ def forward(self, hidden_states: torch.Tensor): else: raise RuntimeError(f'Unsupported topk_method: {self.topk_method}') + if routed_experts is not None: + routed_experts.copy_(topk_idx) + if self.eplb_dispatch_info is not None: topk_idx = EPLBManager.topk_ids_logical_to_physical(topk_idx, self.eplb_dispatch_info) @@ -677,6 +687,7 @@ class DeepseekV2MoE(nn.Module): def __init__(self, config: Any, layer_idx, dtype: torch.dtype = None, device: torch.device = None): super().__init__() + self.layer_idx = layer_idx quantization_config = getattr(config, 'quantization_config', None) self.hidden_dim = config.hidden_size self.ffn_dim = config.moe_intermediate_size @@ -731,11 +742,14 @@ def __init__(self, config: Any, layer_idx, dtype: torch.dtype = None, device: to else: self._all_reduce = False - def forward(self, hidden_states: torch.Tensor): + def forward(self, hidden_states: torch.Tensor, all_routed_experts: torch.Tensor = None): """forward.""" batch_size, sequence_length, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) - topk_weights, topk_ids = self.gate(hidden_states) + routed_experts = None + if all_routed_experts is not None: + routed_experts = all_routed_experts[:, self.layer_idx, :] + topk_weights, topk_ids = self.gate(hidden_states, routed_experts=routed_experts) out_states = self.experts( hidden_states, @@ -1278,11 +1292,13 @@ def __load_kcvc_blocked_fp8(name: str, loaded_weight: torch.Tensor): if '.kv_b_proj' in name: quantization_config = self.quantization_config quant_method = None + fp8_quant_scope = None if quantization_config is not None: quant_method = quantization_config.get('quant_method') + fp8_quant_scope = quantization_config.get('fp8_quant_scope') loaded_weight = loaded_weight.to(device) - if quant_method == 'fp8': + if quant_method == 'fp8' and fp8_quant_scope != 'moe_only': # update blocked fp8 weight __load_kcvc_blocked_fp8(name, loaded_weight) else: diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index 6954e47fca..5a020984a5 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -6,8 +6,9 @@ import torch.nn.functional as F from torch import nn +from lmdeploy.pytorch import envs as _envs from lmdeploy.pytorch.distributed import get_dist_manager, get_ep_world_rank -from lmdeploy.pytorch.model_inputs import StepContextManager +from lmdeploy.pytorch.model_inputs import StepContext, StepContextManager, get_step_ctx_manager from lmdeploy.pytorch.nn import ( ApplyRotaryEmb, Attention, @@ -31,6 +32,75 @@ DeepseekV2MoE, yarn_get_mscale, ) +from .patch import get_build_model_context + + +def get_layer_indexer_type(config: Any, layer_idx: int | None) -> str: + """Return whether a DSA layer computes or reuses top-k indices.""" + indexer_types = getattr(config, 'indexer_types', None) + if indexer_types is None or layer_idx is None or layer_idx >= len(indexer_types): + return 'full' + return indexer_types[layer_idx] + + +def get_layer_idx_from_weight_name(name: str) -> int | None: + """Parse a transformer layer index from a checkpoint parameter name.""" + for marker in ('.layers.', 'layers.'): + if marker not in name: + continue + try: + return int(name.split(marker, 1)[1].split('.', 1)[0]) + except ValueError: + return None + return None + + +class DSATopKIndicesBuffer(nn.Module): + """Persistent DSA top-k buffer shared by full and reuse layers.""" + + def __init__(self, topk: int): + super().__init__() + self.topk = topk + self.register_buffer('indices', None, persistent=False) + + def _target_capacity(self, num_tokens: int) -> int: + capacity = num_tokens + ctx_mgr = get_step_ctx_manager() + if ctx_mgr is None: + return capacity + + context = ctx_mgr.current_context() + cache_config = getattr(context, 'cache_config', None) + max_prefill_token_num = getattr(cache_config, 'max_prefill_token_num', None) + if max_prefill_token_num is not None: + capacity = max(capacity, max_prefill_token_num) + return capacity + + def ensure(self, num_tokens: int, device: torch.device) -> torch.Tensor: + """Return a stable top-k slice with enough capacity for the current + forward.""" + capacity = self._target_capacity(num_tokens) + if (self.indices is None or self.indices.size(0) < capacity or self.indices.device != device): + self.indices = torch.empty(capacity, self.topk, dtype=torch.int32, device=device) + return self.indices[:num_tokens] + + def write(self, topk_indices: torch.Tensor) -> torch.Tensor: + """Copy freshly computed top-k indices into the shared buffer.""" + buffer = self.ensure(topk_indices.size(0), topk_indices.device) + buffer.copy_(topk_indices) + return buffer + + def read(self, num_tokens: int, device: torch.device) -> torch.Tensor: + """Read top-k indices previously written by a full indexer layer.""" + if self.indices is None or self.indices.size(0) < num_tokens or self.indices.device != device: + raise RuntimeError('DSA top-k indices are reused before the shared buffer is populated.') + return self.indices[:num_tokens] + + def compact(self, row_indices: torch.Tensor) -> torch.Tensor: + """Copy selected rows to the prefix for recurrent MTP reuse.""" + selected = self.indices.index_select(0, row_indices) + self.indices[:selected.size(0)].copy_(selected) + return self.indices[:selected.size(0)] def rotate_activation(x: torch.Tensor) -> torch.Tensor: @@ -40,6 +110,55 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: return hadamard_transform(x, scale=hidden_size**-0.5) +def _dequantize_blocked_fp8(weight: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Dequantize a 2D block-FP8 checkpoint tensor.""" + dim_w0, dim_w1 = weight.shape + dim_s0, dim_s1 = scale.shape + assert dim_w0 % dim_s0 == 0 and dim_w1 % dim_s1 == 0 + weight = weight.reshape(dim_s0, dim_w0 // dim_s0, dim_s1, dim_w1 // dim_s1) + weight = weight.float() * scale.reshape(dim_s0, 1, dim_s1, 1) + return weight.to(dtype).reshape(dim_w0, dim_w1) + + +def _load_fused_indexer_weight(name: str, loaded_weight: torch.Tensor, params_dict: dict[str, nn.Parameter], + load_buffers: dict) -> bool: + """Load separate checkpoint projections into one fused BF16 weight.""" + is_wk = '.self_attn.indexer.wk.' in name + is_gate = '.self_attn.indexer.weights_proj.' in name + if not (is_wk or is_gate): + return False + + indexer_prefix = name.rsplit('.indexer.', 1)[0] + '.indexer' + fused_param = params_dict.get(f'{indexer_prefix}.wk_weights_proj.weight') + if fused_param is None: + return False + + if is_gate: + if not name.endswith('.weight'): + return False + gate = loaded_weight.to(device=fused_param.device, dtype=fused_param.dtype) + fused_param.data[-gate.size(0):].copy_(gate) + return True + + if name.endswith('.weight') and loaded_weight.dtype != torch.float8_e4m3fn: + wk = loaded_weight.to(device=fused_param.device, dtype=fused_param.dtype) + fused_param.data[:wk.size(0)].copy_(wk) + return True + + is_weight = name.endswith('.weight') + is_scale = name.endswith('.weight_scale_inv') + if not (is_weight or is_scale): + return False + + buffer = load_buffers.setdefault(f'{indexer_prefix}.wk', {}) + buffer['weight' if is_weight else 'scale'] = loaded_weight.to(fused_param.device) + if 'weight' in buffer and 'scale' in buffer: + wk = _dequantize_blocked_fp8(buffer['weight'], buffer['scale'], fused_param.dtype) + fused_param.data[:wk.size(0)].copy_(wk) + load_buffers.pop(f'{indexer_prefix}.wk') + return True + + class LayerNorm(nn.Module): """Layer Normalization.""" @@ -60,10 +179,6 @@ class Indexer(nn.Module): def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, device: torch.device = None): super().__init__() - try: - import fast_hadamard_transform # noqa: F401 - except ImportError: - raise ImportError('Please install fast_hadamard_transform package.') quant_config = getattr(config, 'quantization_config', None) self.layer_idx = layer_idx # self.dim: int = 2048 @@ -72,6 +187,7 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic self.n_local_heads = config.index_n_heads self.head_dim: int = config.index_head_dim self.rope_head_dim: int = config.qk_rope_head_dim + self.rope_interleave: bool = getattr(config, 'indexer_rope_interleave', False) self.index_topk: int = config.index_topk self.q_lora_rank: int = config.q_lora_rank self.wq_b = build_colwise_linear(self.q_lora_rank, @@ -81,30 +197,93 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic device=device, is_tp=False, quant_config=quant_config) - self.wk = build_colwise_linear(self.dim, - self.head_dim, - bias=False, - dtype=dtype, - device=device, - is_tp=False, - quant_config=quant_config) + self.use_fusion = not _envs.disable_dsa_indexer_fusion + if self.use_fusion: + # Merge K and head-gate projections into one BF16 GEMM. + self.wk_weights_proj = build_colwise_linear(self.dim, + self.head_dim + self.n_heads, + bias=False, + dtype=torch.bfloat16, + device=device, + is_tp=False) + else: + self.wk = build_colwise_linear(self.dim, + self.head_dim, + bias=False, + dtype=dtype, + device=device, + is_tp=False, + quant_config=quant_config) + self.weights_proj = build_colwise_linear(self.dim, + self.n_heads, + bias=False, + dtype=dtype, + device=device, + is_tp=False) self.k_norm = LayerNorm(self.head_dim, device=device) - self.weights_proj = build_colwise_linear(self.dim, - self.n_heads, - bias=False, - dtype=dtype, - device=device, - is_tp=False) self.softmax_scale = self.head_dim**-0.5 self.apply_rotary_pos_emb = ApplyRotaryEmb() self.indexer_topk = IndexerTopKFP8(self.index_topk, self.softmax_scale, block_size=128, fill=-1) + def _apply_rotary_pos_emb(self, q_pe: torch.Tensor, k_pe: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor]): + """Apply the indexer's RoPE layout.""" + cos, sin = freqs_cis + if self.rope_interleave: + half_size = cos.size(-1) // 2 + cos = cos[..., :half_size] + sin = sin[..., :half_size] + k_pe = k_pe[..., None, :] + return self.apply_rotary_pos_emb( + q_pe, + k_pe, + cos, + sin, + inplace=False, + complex_mode=self.rope_interleave, + ) + def forward(self, x: torch.Tensor, qr: torch.Tensor, freqs_cis: torch.Tensor, index_cache: tuple[torch.Tensor, torch.Tensor], - attn_metadata: Any = None): + attn_metadata: Any = None, + use_dense_index: bool = False): + if self.use_fusion: + # Fused kernels consume these projections without rotated BF16 Q/K temporaries. + kw = self.wk_weights_proj(x) + k, weights = kw.split([self.head_dim, self.n_heads], dim=-1) + cos, sin = freqs_cis + if use_dense_index: + # Every visible position is selected below index_topk. Only K + # still needs preparation so later sparse layers can reuse it. + return self.indexer_topk.forward_k_only(k[0], + self.k_norm.weight, + self.k_norm.bias, + cos, + sin, + index_cache[0], + index_cache[1], + norm_eps=self.k_norm.eps, + rope_interleaved=self.rope_interleave, + attn_metadata=attn_metadata) + q = self.wq_b(qr) + q = q.unflatten(-1, (-1, self.head_dim)) + return self.indexer_topk.forward_fused(q[0], + k[0], + weights[0], + self.k_norm.weight, + self.k_norm.bias, + cos, + sin, + index_cache[0], + index_cache[1], + norm_eps=self.k_norm.eps, + head_gate_scale=self.n_heads**-0.5, + rope_interleaved=self.rope_interleave, + attn_metadata=attn_metadata) + q = self.wq_b(qr) q = q.unflatten(-1, (-1, self.head_dim)) q_pe, q_nope = torch.split(q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1) @@ -113,14 +292,7 @@ def forward(self, k_pe, k_nope = torch.split(k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1) # apply rotary embedding - cos, sin = freqs_cis - q_pe, k_pe = self.apply_rotary_pos_emb( - q_pe, - k_pe[..., None, :], - cos, - sin, - inplace=False, - ) + q_pe, k_pe = self._apply_rotary_pos_emb(q_pe, k_pe, freqs_cis) k_pe = k_pe[0, :] k_nope = k_nope[0, :, None] q = torch.cat([q_pe, q_nope], dim=-1) @@ -239,7 +411,10 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic quant_config=quantization_config, ) - self.indexer = Indexer(config, layer_idx, dtype=dtype, device=device) + self.indexer_type = get_layer_indexer_type(config, layer_idx) + self.indexer = None + if self.indexer_type == 'full': + self.indexer = Indexer(config, layer_idx, dtype=dtype, device=device) def _q_proj(self, hidden_states, num_heads: int, nope_size: int, pe_size: int): """Q proj.""" @@ -288,11 +463,16 @@ def forward( rotary_pos_emb: tuple[torch.FloatTensor, torch.FloatTensor], past_key_value: Sequence[torch.Tensor] = None, attn_metadata: Any = None, + topk_indices_buffer: DSATopKIndicesBuffer | None = None, + skip_topk: bool = False, + use_dense_index: bool = False, ): """Rewrite of LlamaAttention.forward.""" - dist_ctx = get_dist_manager().current_context() - tp_world_size = dist_ctx.dist_config.attn_tp - num_heads = self.num_heads // tp_world_size + dist_config = get_dist_manager().current_config() + if dist_config.dp > 1: + num_heads = self.num_heads + else: + num_heads = self.num_heads // dist_config.attn_tp nope_size = self.kv_lora_rank q_len = hidden_states.size(1) @@ -310,7 +490,20 @@ def forward( query_states[..., nope_size:] = q_pe key_states[..., nope_size:] = k_pe - topk_indices = self.indexer(hidden_states, qr, rotary_pos_emb, past_key_value[-2:], attn_metadata=attn_metadata) + if topk_indices_buffer is None: + raise RuntimeError(f'Layer {self.layer_idx} requires a DSA top-k indices buffer.') + + should_compute_topk = self.indexer is not None and not skip_topk + if should_compute_topk: + topk_indices = topk_indices_buffer.write( + self.indexer(hidden_states, + qr, + rotary_pos_emb, + past_key_value[-2:], + attn_metadata=attn_metadata, + use_dense_index=use_dense_index)) + else: + topk_indices = topk_indices_buffer.read(q_len, hidden_states.device) attn_output = self.attn_fwd( query_states, @@ -365,6 +558,51 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic dtype=torch.float32, device=device) + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb: tuple[torch.FloatTensor, torch.FloatTensor], + past_key_value: list[torch.FloatTensor] | None, + residual: torch.Tensor | None = None, + attn_metadata: Any = None, + topk_indices_buffer: DSATopKIndicesBuffer | None = None, + skip_topk: bool = False, + use_dense_index: bool = False, + all_routed_experts: torch.Tensor | None = None, + ) -> tuple[torch.FloatTensor, torch.FloatTensor]: + + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + if isinstance(self.self_attn, DeepseekV32Attention): + hidden_states = self.self_attn( + hidden_states=hidden_states, + rotary_pos_emb=rotary_pos_emb, + past_key_value=past_key_value, + attn_metadata=attn_metadata, + topk_indices_buffer=topk_indices_buffer, + skip_topk=skip_topk, + use_dense_index=use_dense_index, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + rotary_pos_emb=rotary_pos_emb, + past_key_value=past_key_value, + attn_metadata=attn_metadata, + ) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + if isinstance(self.mlp, DeepseekV2MoE): + hidden_states = self.mlp(hidden_states, all_routed_experts=all_routed_experts) + else: + hidden_states = self.mlp(hidden_states) + + return hidden_states, residual + class DeepseekV32Model(DeepseekV2Model): @@ -385,6 +623,7 @@ def __init__(self, config: Any, dtype: torch.dtype = None, device: torch.device DeepseekV32DecoderLayer(config, layer_idx, dtype=dtype, device=device) for layer_idx in range(config.num_hidden_layers) ]) + self.topk_indices_buffer = DSATopKIndicesBuffer(config.index_topk) # build norm self.norm = RMSNorm(config.hidden_size, @@ -404,6 +643,65 @@ def __init__(self, config: Any, dtype: torch.dtype = None, device: torch.device rope_params.update(update_params) self.rotary_emb = build_rotary_embedding(**rope_params) + def forward( + self, + input_ids: torch.LongTensor = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + attn_metadata: Any = None, + inputs_embeds: torch.FloatTensor | None = None, + all_routed_experts: torch.Tensor | None = None, + use_dense_index: bool = False, + ): + """forward.""" + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + residual = None + cos, sin = self.rotary_emb(hidden_states, position_ids) + cos, sin = cos[0], sin[0] + rotary_pos_emb = (cos, sin) + for idx, decoder_layer in enumerate(self.layers): + past_key_value = past_key_values[idx] + hidden_states, residual = decoder_layer( + hidden_states, + rotary_pos_emb=rotary_pos_emb, + past_key_value=past_key_value, + residual=residual, + attn_metadata=attn_metadata, + topk_indices_buffer=self.topk_indices_buffer, + all_routed_experts=all_routed_experts, + use_dense_index=use_dense_index, + ) + + hidden_states, _ = self.norm(hidden_states, residual) + + return hidden_states + + def forward_microbatch( + self, + input_ids: torch.LongTensor = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + attn_metadata: Any = None, + inputs_embeds: torch.FloatTensor | None = None, + all_routed_experts: torch.Tensor | None = None, + use_dense_index: bool = False, + ): + """forward_microbatch.""" + # DSA shared top-k indices are model-global; use normal forward until + # microbatching has per-microbatch top-k buffers. + return self.forward( + input_ids=input_ids, + position_ids=position_ids, + past_key_values=past_key_values, + attn_metadata=attn_metadata, + inputs_embeds=inputs_embeds, + all_routed_experts=all_routed_experts, + use_dense_index=use_dense_index, + ) + class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM): @@ -425,3 +723,70 @@ def __init__(self, dtype=dtype, device=device) self._load_buffers = dict() + bm_ctx = get_build_model_context() + self.enable_return_routed_experts = bm_ctx.enable_return_routed_experts + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + past_key_values: list[list[torch.Tensor]], + attn_metadata: Any = None, + inputs_embeds: torch.Tensor = None, + use_dense_index: bool = False, + **kwargs, + ): + """Model forward.""" + step_ctx = get_step_ctx_manager().current_context() + num_tokens = inputs_embeds.size(1) if inputs_embeds is not None else input_ids.size(1) + all_routed_experts = None + if self.enable_return_routed_experts: + # Dense layers do not produce routed expert IDs. Keep their slots at an out-of-range sentinel; + # MoE layers overwrite their own slots below. + all_routed_experts = position_ids.new_full( + (num_tokens, self.config.num_hidden_layers, self.config.num_experts_per_tok), + torch.iinfo(torch.uint16).max, + dtype=torch.uint16, + ) + forward = self.model.forward_microbatch if step_ctx.enable_microbatch else self.model.forward + hidden_states = forward( + input_ids=input_ids, + position_ids=position_ids, + past_key_values=past_key_values, + attn_metadata=attn_metadata, + inputs_embeds=inputs_embeds, + all_routed_experts=all_routed_experts, + use_dense_index=use_dense_index, + ) + if all_routed_experts is None: + return hidden_states + return dict(hidden_states=hidden_states, all_routed_experts=all_routed_experts) + + def prepare_inputs_for_generation( + self, + past_key_values: list[list[torch.Tensor]], + inputs_embeds: torch.Tensor | None = None, + context: StepContext = None, + ): + """Prepare inputs and select the exact dense-index shortcut.""" + inputs = super().prepare_inputs_for_generation( + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + context=context, + ) + inputs['use_dense_index'] = ( + context.is_decoding and context.max_kv_seqlen <= self.config.index_topk) + return inputs + + def _load_weight_attention(self, name: str, loaded_weight: torch.Tensor, params_dict: dict[str, nn.Parameter], + update_pe_mapping: list): + """Load attention weights.""" + if _load_fused_indexer_weight(name, loaded_weight, params_dict, + self._load_buffers): + return + if '.self_attn.indexer.' in name and name not in params_dict: + layer_idx = get_layer_idx_from_weight_name(name) + # Shared DSA layers reuse cached top-k indices and have no local indexer. + if get_layer_indexer_type(self.config, layer_idx) == 'shared': + return + return super()._load_weight_attention(name, loaded_weight, params_dict, update_pe_mapping) diff --git a/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py b/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py new file mode 100644 index 0000000000..3595eaaa55 --- /dev/null +++ b/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py @@ -0,0 +1,267 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +from typing import Any + +import torch +from torch import nn +from transformers.configuration_utils import PretrainedConfig + +from lmdeploy.pytorch.model_inputs import StepContext, StepContextManager +from lmdeploy.pytorch.nn import RMSNorm +from lmdeploy.pytorch.nn.linear import build_colwise_linear + +from .deepseek_mtp import DeepseekMTPModel, build_deepseek_rotary_embedding +from .deepseek_v32 import ( + DeepseekV32DecoderLayer, + DSATopKIndicesBuffer, + _load_fused_indexer_weight, +) + + +class GlmMoeDsaSharedHead(nn.Module): + """GLM MTP final normalization.""" + + def __init__(self, + config: PretrainedConfig, + dtype: torch.dtype = None, + device: torch.device = None): + super().__init__() + self.norm = RMSNorm(config.hidden_size, + config.rms_norm_eps, + dtype=dtype, + device=device) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class GlmMoeDsaMultiTokenPredictorLayer(nn.Module): + """GLM-MoE-DSA MTP layer.""" + + def __init__( + self, + config: PretrainedConfig, + layer_idx: int, + dtype: torch.dtype = None, + device: torch.device = None, + ) -> None: + super().__init__() + quantization_config = getattr(config, 'quantization_config', None) + self.enorm = RMSNorm(config.hidden_size, + config.rms_norm_eps, + dtype=dtype, + device=device) + self.hnorm = RMSNorm(config.hidden_size, + config.rms_norm_eps, + dtype=dtype, + device=device) + self.eh_proj = build_colwise_linear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + dtype=dtype, + device=device, + is_tp=False, + quant_config=quantization_config, + dp_disable_tp=True, + ) + self.shared_head = GlmMoeDsaSharedHead(config, + dtype=dtype, + device=device) + self.mtp_block = DeepseekV32DecoderLayer(config, + layer_idx=layer_idx, + dtype=dtype, + device=device) + self.rotary_emb = build_deepseek_rotary_embedding(config) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + previous_hidden_states: torch.Tensor, + past_key_value: list[torch.Tensor], + inputs_embeds: torch.Tensor | None = None, + attn_metadata: Any = None, + topk_indices_buffer: DSATopKIndicesBuffer | None = None, + skip_topk: bool = False, + use_dense_index: bool = False, + ) -> torch.Tensor: + inputs_embeds = self.enorm(inputs_embeds) + previous_hidden_states = self.hnorm(previous_hidden_states) + hidden_states = self.eh_proj( + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)) + + cos, sin = self.rotary_emb(hidden_states, position_ids) + hidden_states, residual = self.mtp_block( + hidden_states, + (cos[0], sin[0]), + past_key_value, + attn_metadata=attn_metadata, + topk_indices_buffer=topk_indices_buffer, + skip_topk=skip_topk, + use_dense_index=use_dense_index, + ) + return residual + hidden_states + + +class GlmMoeDsaMultiTokenPredictor(nn.Module): + """GLM-MoE-DSA multi-token predictor.""" + + def __init__(self, + config: PretrainedConfig, + dtype: torch.dtype = None, + device: torch.device = None): + super().__init__() + self.config = config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + self.embed_tokens = None + self.layers = nn.ModuleDict({ + str(idx): + GlmMoeDsaMultiTokenPredictorLayer(config, + idx, + dtype=dtype, + device=device) + for idx in range(self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers) + }) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + previous_hidden_states: torch.Tensor, + past_key_values: list[list[torch.Tensor]], + inputs_embeds: torch.Tensor | None = None, + attn_metadata: Any = None, + topk_indices_buffer: DSATopKIndicesBuffer | None = None, + skip_topk: bool = False, + use_dense_index: bool = False, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + layer_idx = self.mtp_start_layer_idx + current_step_idx + return self.layers[str(layer_idx)]( + input_ids, + position_ids, + previous_hidden_states, + past_key_values[current_step_idx], + inputs_embeds=inputs_embeds, + attn_metadata=attn_metadata, + topk_indices_buffer=topk_indices_buffer, + skip_topk=skip_topk, + use_dense_index=use_dense_index, + ) + + def prepare_hidden_states_for_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + """Apply the GLM final norm for logits.""" + current_step_idx = spec_step_idx % self.num_mtp_layers + layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + return layer.shared_head(hidden_states) + + def set_input_embeddings(self, embed_tokens: nn.Module): + self.embed_tokens = embed_tokens + + def get_input_embeddings(self): + return self.embed_tokens + + +class GlmMoeDsaMTPModel(DeepseekMTPModel): + """GLM-MoE-DSA MTP model.""" + + def __init__( + self, + config: PretrainedConfig, + ctx_mgr: StepContextManager, + dtype: torch.dtype = None, + device: torch.device = None, + ): + nn.Module.__init__(self) + self.config = config + self.quantization_config = getattr(config, 'quantization_config', None) + self.dtype = dtype + self.ctx_mgr = ctx_mgr + self.model = GlmMoeDsaMultiTokenPredictor(config, + dtype=dtype, + device=device) + self.topk_indices_buffer = DSATopKIndicesBuffer(config.index_topk) + self.uses_dsa_topk_buffer = getattr(config, + 'index_share_for_mtp_iteration', + False) + self._load_buffers = dict() + + def set_input_embeddings(self, embed_tokens: nn.Module): + self.model.set_input_embeddings(embed_tokens) + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def compact_topk_indices(self, row_indices: torch.Tensor): + self.topk_indices_buffer.compact(row_indices) + + def prepare_hidden_states_for_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model.prepare_hidden_states_for_logits(hidden_states, + spec_step_idx=spec_step_idx) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + target_hidden_states: torch.Tensor, + past_key_values: list[list[torch.Tensor]], + attn_metadata: Any = None, + inputs_embeds: torch.Tensor | None = None, + skip_topk: bool = False, + use_dense_index: bool = False, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, + position_ids, + target_hidden_states, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attn_metadata=attn_metadata, + topk_indices_buffer=self.topk_indices_buffer, + skip_topk=skip_topk, + use_dense_index=use_dense_index, + spec_step_idx=spec_step_idx, + ) + + def prepare_inputs_for_generation( + self, + past_key_values: list[list[torch.Tensor]], + inputs_embeds: torch.Tensor | None = None, + context: StepContext = None, + ): + inputs = super().prepare_inputs_for_generation( + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + context=context, + ) + model_metas = context.model_metas + inputs['skip_topk'] = ( + self.uses_dsa_topk_buffer and model_metas is not None and len(model_metas) > 0 + and all(meta is not None and meta.get('skip_topk', False) for meta in model_metas)) + inputs['use_dense_index'] = ( + context.is_decoding and context.max_kv_seqlen <= self.config.index_topk) + return inputs + + def _load_weight_attention(self, name: str, loaded_weight: torch.Tensor, + params_dict: dict[str, nn.Parameter], + update_pe_mapping: list): + if _load_fused_indexer_weight(name, loaded_weight, params_dict, + self._load_buffers): + return + return super()._load_weight_attention(name, loaded_weight, params_dict, + update_pe_mapping) diff --git a/lmdeploy/pytorch/models/module_map.py b/lmdeploy/pytorch/models/module_map.py index 27e74e8e67..baad0282ef 100644 --- a/lmdeploy/pytorch/models/module_map.py +++ b/lmdeploy/pytorch/models/module_map.py @@ -64,6 +64,7 @@ # glm5 MODULE_MAP.update({'GlmMoeDsaForCausalLM': f'{LMDEPLOY_PYTORCH_MODEL_PATH}.deepseek_v32.DeepseekV32ForCausalLM'}) +MODULE_MAP.update({'GlmMoeDsaMTPModel': f'{LMDEPLOY_PYTORCH_MODEL_PATH}.glm_moe_dsa_mtp.GlmMoeDsaMTPModel'}) # internlm2 MODULE_MAP.update({ diff --git a/lmdeploy/pytorch/models/utils/cudagraph.py b/lmdeploy/pytorch/models/utils/cudagraph.py index 55b530f63e..b3e12de696 100644 --- a/lmdeploy/pytorch/models/utils/cudagraph.py +++ b/lmdeploy/pytorch/models/utils/cudagraph.py @@ -173,7 +173,10 @@ def make_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p input_buffers['cu_seqlens_q'] = input_buffers['cu_seqlens'][0] input_buffers['cu_seqlens_k'] = input_buffers['cu_seqlens'][1] - if graph_meta.use_flash_mla is True: + # BF16 sparse decode uses sparse_fwd and has no scheduler metadata buffers. + use_flash_mla_meta = (graph_meta.use_flash_mla + and (graph_meta.use_mla_fp8_cache or graph_meta.mla_index_topk is None)) + if use_flash_mla_meta: import flash_mla # create buffers for flash mla @@ -258,24 +261,32 @@ def fill_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p attn_metadata.cu_seqlens_q = input_buffers['cu_seqlens_q'] attn_metadata.cu_seqlens_k = input_buffers['cu_seqlens_k'] - if graph_meta.use_flash_mla is True: - import flash_mla - step_ctx = get_step_ctx_manager().current_context() - model_config = step_ctx.model_config - num_attention_heads, _ = model_config.get_num_qkv_head_by_tp() - index_topk = graph_meta.mla_index_topk - num_heads_q = None if index_topk is None else num_attention_heads - tile_scheduler_metadata, num_splits = flash_mla.get_mla_metadata( - attn_metadata.kv_seqlens.to(torch.int32), - num_attention_heads * decode_query_len, - num_heads_k=1, - num_heads_q=num_heads_q, - is_fp8_kvcache=graph_meta.use_mla_fp8_cache, - topk=index_topk) - # here we use copy_ instead of = to avoid using new allocated mem for cuda graph - input_buffers['tile_scheduler_metadata'].copy_(tile_scheduler_metadata) - input_buffers['num_splits'][:new_batch_size + 1].copy_(num_splits[:new_batch_size + 1]) - attn_metadata.tile_scheduler_metadata = input_buffers['tile_scheduler_metadata'] + # BF16 sparse decode uses sparse_fwd and has no scheduler metadata buffers. + use_flash_mla_meta = (graph_meta.use_flash_mla + and (graph_meta.use_mla_fp8_cache or graph_meta.mla_index_topk is None)) + if use_flash_mla_meta: + scheduler_buffer = input_buffers['tile_scheduler_metadata'] + # The old API returns tensor metadata that must be refreshed while + # preserving its CUDA graph address. The new FlashMLASchedMeta API + # owns and initializes its metadata object on first use. + if isinstance(scheduler_buffer, torch.Tensor): + import flash_mla + step_ctx = get_step_ctx_manager().current_context() + model_config = step_ctx.model_config + num_attention_heads, _ = model_config.get_num_qkv_head_by_tp() + index_topk = graph_meta.mla_index_topk + num_heads_q = None if index_topk is None else num_attention_heads + tile_scheduler_metadata, num_splits = flash_mla.get_mla_metadata( + attn_metadata.kv_seqlens.to(torch.int32), + num_attention_heads * decode_query_len, + num_heads_k=1, + num_heads_q=num_heads_q, + is_fp8_kvcache=graph_meta.use_mla_fp8_cache, + topk=index_topk) + # Keep graph input addresses stable for the old FlashMLA metadata API. + scheduler_buffer.copy_(tile_scheduler_metadata) + input_buffers['num_splits'][:new_batch_size + 1].copy_(num_splits[:new_batch_size + 1]) + attn_metadata.tile_scheduler_metadata = scheduler_buffer attn_metadata.num_splits = input_buffers['num_splits'] # use fa3 decode kernel for spec decode diff --git a/lmdeploy/pytorch/nn/nsa.py b/lmdeploy/pytorch/nn/nsa.py index d3944bf003..81bc94ec11 100644 --- a/lmdeploy/pytorch/nn/nsa.py +++ b/lmdeploy/pytorch/nn/nsa.py @@ -15,30 +15,110 @@ def __init__(self, topk: int, softmax_scale: float, block_size: int = 128, fill: index_builder = backend.get_layer_impl_builder(OpType.NSAIndexFP8) self.index_impl = index_builder.build(topk, softmax_scale, block_size, fill) - def forward( - self, - q: Tensor, - k: Tensor, - weights: Tensor, - k_cache: Tensor, - k_s_cache: Tensor, - attn_metadata: AttentionMetadata = None, - ): - """forward.""" + @staticmethod + def _get_max_q_seqlen(q: Tensor, + attn_metadata: AttentionMetadata) -> int: + """Get the query width used by the index and cache-fill kernels. + + Speculative target verification remains a decoding step, but its + flattened Q contains ``num_spec_tokens + 1`` rows per request. The + kernels need that real width to process every verification row. + """ + batch_size = attn_metadata.kv_seqlens.size(0) + # fp8_index also identifies one row per request as decode layout, so + # keep its metadata consistent when the phase flag has not changed yet. + is_decoding = attn_metadata.is_decoding or q.size(0) == batch_size + # Prefer a width prepared by the attention backend; otherwise derive it + # from the flattened query rows. + max_q_seqlen = attn_metadata.max_q_seqlen + if max_q_seqlen is None: + max_q_seqlen = q.size(0) + if is_decoding: + max_q_seqlen //= batch_size + return max_q_seqlen + + @staticmethod + def _build_meta(q: Tensor, attn_metadata: AttentionMetadata) -> NSAIndexMeta: step_ctx = get_step_ctx_manager().current_context() cache_config = step_ctx.cache_config max_tokens = cache_config.block_size * cache_config.num_gpu_blocks is_decoding = attn_metadata.is_decoding if q.size(0) == attn_metadata.kv_seqlens.size(0): is_decoding = True - max_q_seqlen = 1 if is_decoding else q.size(0) - # we need to make max_kv_seqlen=max_allocated_cache_len to enable cudagraph + max_q_seqlen = IndexerTopKFP8._get_max_q_seqlen(q, attn_metadata) + # Decode uses the full cache capacity to keep CUDA graph shapes stable. max_kv_seqlen = max_tokens if is_decoding else attn_metadata.kv_flatten_size - meta = NSAIndexMeta(cu_seqlen_q=attn_metadata.cu_seqlens_q, + return NSAIndexMeta(cu_seqlen_q=attn_metadata.cu_seqlens_q, q_seqlens=attn_metadata.q_seqlens, k_seqlens=attn_metadata.kv_seqlens, block_offset=attn_metadata.block_offsets, max_q_seqlen=max_q_seqlen, max_kv_seqlen=max_kv_seqlen) + + def forward( + self, + q: Tensor, + k: Tensor, + weights: Tensor, + k_cache: Tensor, + k_s_cache: Tensor, + attn_metadata: AttentionMetadata = None, + ): + """forward.""" + meta = self._build_meta(q, attn_metadata) ret = self.index_impl.forward(q, k, weights, k_cache, k_s_cache, meta=meta) return ret + + def forward_fused(self, + q: Tensor, + k: Tensor, + weights: Tensor, + norm_weight: Tensor, + norm_bias: Tensor, + cos: Tensor, + sin: Tensor, + k_cache: Tensor, + k_s_cache: Tensor, + norm_eps: float, + head_gate_scale: float, + rope_interleaved: bool, + attn_metadata: AttentionMetadata = None): + """Forward with fused DSA indexer preparation.""" + meta = self._build_meta(q, attn_metadata) + return self.index_impl.forward_fused(q, + k, + weights, + norm_weight, + norm_bias, + cos, + sin, + k_cache, + k_s_cache, + norm_eps=norm_eps, + head_gate_scale=head_gate_scale, + rope_interleaved=rope_interleaved, + meta=meta) + + def forward_k_only(self, + k: Tensor, + norm_weight: Tensor, + norm_bias: Tensor, + cos: Tensor, + sin: Tensor, + k_cache: Tensor, + k_s_cache: Tensor, + norm_eps: float, + rope_interleaved: bool, + attn_metadata: AttentionMetadata = None): + """Cache K and return identity indices without scoring Q.""" + meta = self._build_meta(k, attn_metadata) + return self.index_impl.forward_k_only(k, + norm_weight, + norm_bias, + cos, + sin, + k_cache, + k_s_cache, + norm_eps=norm_eps, + rope_interleaved=rope_interleaved, + meta=meta) diff --git a/lmdeploy/pytorch/paging/block_trie.py b/lmdeploy/pytorch/paging/block_trie.py index 97f246b711..c8c856a907 100644 --- a/lmdeploy/pytorch/paging/block_trie.py +++ b/lmdeploy/pytorch/paging/block_trie.py @@ -1024,6 +1024,9 @@ def _verify_state_checkpoint_node(self, seq: SchedulerSequence, node: Node, inde if not self._match_node(block_node, tokens, extra_hashes): return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.REQUEST_MISMATCH, reason=f'block payload mismatch at block {idx}') + if seq.return_routed_experts and block_node.routed_experts is None: + return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.REQUEST_MISMATCH, + reason=f'routed experts missing at block {idx}') matched_blocks.append(block_node.block) return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.HIT, @@ -1159,6 +1162,8 @@ def __match_success(node: Node): child = curr.children[key] if not self._match_node(child, curr_tokens, extra_hashes): break + if seq.return_routed_experts and child.routed_experts is None: + break matched_nodes.append(child) __match_success(child) @@ -1181,8 +1186,9 @@ def __clamp_match_step(match_step: int): num_matched = init_num_matched max_match_step = seq.get_prefix_cache_max_match_step() - raw_num_matched = num_matched - effective_num_matched = seq.clamp_prefix_cache_match_step(min(raw_num_matched, max_match_step)) + candidate_num_matched = num_matched + raw_num_matched = self._find_raw_block_match_step(seq, init_curr) + effective_num_matched = seq.clamp_prefix_cache_match_step(min(candidate_num_matched, max_match_step)) __clamp_match_step(effective_num_matched) self._set_private_recompute_range(seq, num_matched, raw_num_matched) diff --git a/lmdeploy/pytorch/spec_decode/base.py b/lmdeploy/pytorch/spec_decode/base.py index 80a122cf29..5c67936ac9 100644 --- a/lmdeploy/pytorch/spec_decode/base.py +++ b/lmdeploy/pytorch/spec_decode/base.py @@ -18,10 +18,10 @@ def _build_draft_dist_ctx(dist_ctx: DistContext, specdecode_config: SpecDecodeCo if specdecode_config is None: return None - if specdecode_config.method == 'qwen3_5_mtp': + draft_dist_config = specdecode_config.dist_config + if draft_dist_config == dist_ctx.dist_config: return dist_ctx - draft_dist_config = specdecode_config.dist_config return DistContext.build(rank=dist_ctx.rank, dist_config=draft_dist_config) diff --git a/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py b/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py index 2912758928..3e2effe361 100644 --- a/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py +++ b/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py @@ -2,18 +2,22 @@ import torch -from lmdeploy.utils import get_logger - from ...model_inputs import ModelInputs from ...strategies.ar_spec.model_agent import ARSpecExtraInputs from .base import SPEC_PROPOSERS, BaseSpecProposer -logger = get_logger('lmdeploy') - @SPEC_PROPOSERS.register_module(name='deepseek_mtp') class DeepseekMTP(BaseSpecProposer): + def build_model(self, empty_init: bool, target_model: torch.nn.Module = None, build_model_ctx=None): + """Build the draft model and bind the target embedding.""" + super().build_model(empty_init, target_model=target_model, build_model_ctx=build_model_ctx) + draft_model = self.model + if (getattr(draft_model, 'uses_dsa_topk_buffer', False) + and hasattr(draft_model, 'set_input_embeddings')): + draft_model.set_input_embeddings(target_model.get_input_embeddings()) + async def get_outputs(self, model_outputs: dict[str, torch.Tensor], model_inputs: ModelInputs, @@ -22,15 +26,26 @@ async def get_outputs(self, """Get outputs.""" hidden_states = model_outputs['hidden_states'] model_metas = model_outputs['model_metas'] + draft_model = self.model.get_model() if hasattr(self.model, 'get_model') else self.model + uses_topk_buffer = getattr(draft_model, 'uses_dsa_topk_buffer', False) + if extra_inputs is not None: last_token_loc = extra_inputs.last_token_indices hidden_states = hidden_states[:, last_token_loc] - # use hidden states for draft prefill forward for next step - target_hidden_states = hidden_states + if uses_topk_buffer: + draft_model.compact_topk_indices(last_token_loc) + + if hasattr(draft_model, 'prepare_hidden_states_for_logits'): + hidden_states = draft_model.prepare_hidden_states_for_logits(hidden_states) + logits = self.target_model.get_logits(hidden_states) else: - target_hidden_states = hidden_states + logits = self.get_logits(hidden_states) + target_hidden_states = hidden_states - logits = self.get_logits(hidden_states)[0] + if uses_topk_buffer: + model_metas = model_metas or [None] * hidden_states.size(1) + model_metas = [dict(meta or {}, skip_topk=True) for meta in model_metas] + logits = logits[0] guided_bitmask = await self.guided_helper.prepare_bitmask(logits, guided_processors) if guided_bitmask is not None: diff --git a/lmdeploy/pytorch/strategies/ar_spec/sequence.py b/lmdeploy/pytorch/strategies/ar_spec/sequence.py index 7bd9ae995a..0de052fad1 100644 --- a/lmdeploy/pytorch/strategies/ar_spec/sequence.py +++ b/lmdeploy/pytorch/strategies/ar_spec/sequence.py @@ -46,8 +46,7 @@ def routed_experts(self) -> np.ndarray: end = max(0, self.num_valid_ids - 1) if 0 < end <= len(self.all_routed_experts): return self.all_routed_experts.get_real()[:end] - else: - return None + return None @property def generated_ids(self) -> np.ndarray: diff --git a/lmdeploy/serve/openai/api_server.py b/lmdeploy/serve/openai/api_server.py index 61089b295e..7098a64ca4 100644 --- a/lmdeploy/serve/openai/api_server.py +++ b/lmdeploy/serve/openai/api_server.py @@ -1018,7 +1018,12 @@ async def generate(request: GenerateReqInput, raw_request: Request = None): media_io_kwargs=request.media_io_kwargs, mm_processor_kwargs=request.mm_processor_kwargs) - def create_generate_response_json(res, text, output_ids, logprobs, finish_reason, routed_experts=None): + def create_generate_response_json(res, + text, + output_ids, + logprobs, + finish_reason, + routed_experts=None): # only output router experts in last chunk routed_experts = None if finish_reason is None else routed_experts meta = GenerateReqMetaOutput(finish_reason=dict(type=finish_reason) if finish_reason else None, @@ -1027,7 +1032,7 @@ def create_generate_response_json(res, text, output_ids, logprobs, finish_reason routed_experts=routed_experts, completion_tokens=res.generate_token_len) - response = GenerateReqOutput(text=text, output_ids=output_ids, meta_info=meta, routed_experts=routed_experts) + response = GenerateReqOutput(text=text, output_ids=output_ids, meta_info=meta) return response.model_dump_json() async def generate_stream_generator(): diff --git a/lmdeploy/serve/openai/serving_generate.py b/lmdeploy/serve/openai/serving_generate.py index 10e9dde597..05ffb32b5b 100644 --- a/lmdeploy/serve/openai/serving_generate.py +++ b/lmdeploy/serve/openai/serving_generate.py @@ -19,6 +19,10 @@ def check_request(request: GenerateReqInput, server_context: 'VariableInterface' except AttributeError: pass + if request.return_routed_experts and not engine_config.enable_return_routed_experts: + return ('routed experts requested but not configured in engine configuration. ' + 'May start api_server with --enable-return-routed-experts flag.') + if (request.prompt is not None) ^ (request.input_ids is None): return 'You must specify exactly one of prompt or input_ids' diff --git a/lmdeploy/serve/parsers/response_parser.py b/lmdeploy/serve/parsers/response_parser.py index d000e7a7d6..e5eaae1059 100644 --- a/lmdeploy/serve/parsers/response_parser.py +++ b/lmdeploy/serve/parsers/response_parser.py @@ -492,8 +492,8 @@ def _consume_reasoning(self) -> tuple[str | None, bool]: - Drops the explicit open tag if model emits it. - If no close tag is present, emits only the safe reasoning-text prefix and preserves possible partial-tag suffix for the next chunk. - - If a close tag is found, emits text before the close tag as reasoning content, - consumes the close tag, and switches mode to ``MODE_PLAIN``. + - If a close tag or tool-open tag is found, emits text before it as + reasoning content and switches to the next protocol mode. Returns: ``(emitted_text, progressed)`` where ``emitted_text`` is the reasoning @@ -511,19 +511,40 @@ def _consume_reasoning(self) -> tuple[str | None, bool]: if not close_tag: raise RuntimeError('Invariant violated: MODE_REASONING requires a reasoning_close_tag.') - idx = self._pending.find(close_tag) - # No close tag found, treat the whole pending text as reasoning content. + # GLM-style outputs may start a tool call directly from reasoning. + tool_tag = self.profile.tool_open_tag if self.tool_parser is not None else None + boundary_tags = [tag for tag in (close_tag, tool_tag) if tag] + + idx = -1 + matched_tag = '' + for tag in boundary_tags: + tag_idx = self._pending.find(tag) + if tag_idx >= 0 and (idx < 0 or tag_idx < idx): + idx = tag_idx + matched_tag = tag + if idx < 0: if not self._pending: return None, False + keep = self._longest_open_tag_prefix_suffix(self._pending, boundary_tags) + if keep > 0: + if keep >= len(self._pending): + return None, False + out = self._pending[:-keep] + self._pending = self._pending[-keep:] + return (out if out else None), bool(out) out = self._pending self._pending = '' return out, True reasoning_chunk = self._pending[:idx] - self._pending = self._pending[idx + len(close_tag):] - # reasoning part is done, switch to plain mode - self._mode = self.MODE_PLAIN + self._pending = self._pending[idx + len(matched_tag):] + if matched_tag == close_tag: + self._mode = self.MODE_PLAIN + else: + self._mode = self.MODE_TOOL + if self.tool_parser is not None: + self.tool_parser.start_tool_call() return (reasoning_chunk if reasoning_chunk else None), True def _consume_tool(self) -> tuple[list[DeltaToolCall], bool]: @@ -634,21 +655,27 @@ def parse_complete( pos += len(open_tag) continue close_tag = self.profile.reasoning_close_tag - close_idx = text.find(close_tag, pos) if close_tag else -1 - if close_idx < 0: + # Match streaming: tool-open can implicitly end reasoning. + tool_tag = self.profile.tool_open_tag if self.tool_parser is not None else None + boundary_tags = [tag for tag in (close_tag, tool_tag) if tag] + boundary_idx, boundary_tag = self._find_first(text, boundary_tags, pos) + if boundary_idx < 0: piece = text[pos:] if self.enable_thinking is False: content_parts.append(piece) else: reasoning_parts.append(piece) break - piece = text[pos:close_idx] + piece = text[pos:boundary_idx] if piece: if self.enable_thinking is False: content_parts.append(piece) else: reasoning_parts.append(piece) - pos = close_idx + len(close_tag) + if boundary_tag == close_tag: + pos = boundary_idx + len(close_tag) + else: + pos = boundary_idx mode = self.MODE_PLAIN continue @@ -704,7 +731,11 @@ def validate_complete(self, text: str | None = None) -> bool: close_tag = self.profile.reasoning_close_tag close_idx = text.find(close_tag) if close_tag else -1 if close_idx < 0: - return False + # A valid tool block can also close implicit reasoning. + tool_tag = self.profile.tool_open_tag if self.tool_parser is not None else None + tool_idx = text.find(tool_tag) if tool_tag else -1 + if tool_idx < 0: + return False if self.tool_parser is None or self.request.tool_choice == 'none': return True diff --git a/tests/pytorch/engine/test_cache_engine.py b/tests/pytorch/engine/test_cache_engine.py index 5be038627f..eabb8fdef2 100644 --- a/tests/pytorch/engine/test_cache_engine.py +++ b/tests/pytorch/engine/test_cache_engine.py @@ -1,10 +1,12 @@ # Copyright (c) OpenMMLab. All rights reserved. import asyncio +from types import SimpleNamespace import numpy as np import pytest import torch +from lmdeploy.messages import QuantPolicy from lmdeploy.pytorch.config import BlockCacheSpec, CacheConfig, ModelConfig, StateCacheSpec from lmdeploy.pytorch.disagg.conn.protocol import MigrationProtocol from lmdeploy.pytorch.disagg.messages import MigrationExecutionBatch @@ -24,6 +26,78 @@ def _make_model_config(**kwargs): return model_config +def test_bf16_sparse_mla_cache_layout(): + model_config = ModelConfig(hidden_size=6144, + num_layers=2, + num_attention_heads=64, + num_key_value_heads=1, + bos_token_id=None, + eos_token_id=[1], + head_dim=576, + k_head_dim=576, + v_head_dim=0, + dtype=torch.bfloat16, + use_flash_mla=True, + mla_kv_cache_dtype='bfloat16', + mla_index_topk=2048) + cache_config = CacheConfig(max_batches=1, + block_size=64, + kernel_block_size=64, + num_cpu_blocks=0, + num_gpu_blocks=0) + + k_desc = CacheEngine.get_k_cache_desc(model_config, cache_config) + v_desc = CacheEngine.get_v_cache_desc(model_config, cache_config) + + assert k_desc.shape == [64, 1, 576] + assert k_desc.dtype == torch.bfloat16 + assert v_desc.shape == [64, 1, 0] + assert v_desc.dtype == torch.bfloat16 + + +def test_fp8_sparse_mla_cache_layout(): + model_config = ModelConfig(hidden_size=6144, + num_layers=2, + num_attention_heads=64, + num_key_value_heads=1, + bos_token_id=None, + eos_token_id=[1], + head_dim=576, + k_head_dim=576, + v_head_dim=0, + dtype=torch.bfloat16, + use_flash_mla=True, + mla_kv_cache_dtype='bfloat16', + mla_index_topk=2048) + cache_config = CacheConfig(max_batches=1, + block_size=64, + kernel_block_size=64, + num_cpu_blocks=0, + num_gpu_blocks=0, + quant_policy=QuantPolicy.FP8) + CacheEngine.get_cache_block_size(cache_config, model_config) + + assert model_config.mla_kv_cache_dtype == 'fp8_ds_mla' + + k_desc = CacheEngine.get_k_cache_desc(model_config, cache_config) + v_desc = CacheEngine.get_v_cache_desc(model_config, cache_config) + + assert k_desc.shape == [64, 1, 656] + assert k_desc.dtype == torch.float8_e4m3fn + assert v_desc.shape == [64, 1, 0] + assert v_desc.dtype == torch.float8_e4m3fn + + +@pytest.mark.parametrize('quant_policy', + [QuantPolicy.INT4, QuantPolicy.INT8, QuantPolicy.FP8_E5M2, QuantPolicy.TURBO_QUANT]) +def test_sparse_mla_rejects_other_cache_policies(quant_policy): + model_config = SimpleNamespace(mla_index_topk=2048) + cache_config = SimpleNamespace(quant_policy=quant_policy) + + with pytest.raises(ValueError, match='Sparse MLA does not support quant_policy'): + CacheEngine.get_cache_block_size(cache_config, model_config) + + def test_allocate_caches_requires_block_size_divisible_by_kernel_block_size(): cache_config = CacheConfig(max_batches=1, block_size=96, diff --git a/tests/pytorch/kernel/test_apply_rotary.py b/tests/pytorch/kernel/test_apply_rotary.py index 6589907a73..a627deeef5 100644 --- a/tests/pytorch/kernel/test_apply_rotary.py +++ b/tests/pytorch/kernel/test_apply_rotary.py @@ -183,3 +183,24 @@ def test_apply_rotary_complex(self, q_states, k_states, cos, sin, gt): atol = None torch.testing.assert_close(q_embed, q_gt, rtol=rtol, atol=atol) torch.testing.assert_close(k_embed, k_gt, rtol=rtol, atol=atol) + + @pytest.mark.parametrize('tokens', [1, 33, 257]) + def test_apply_rotary_complex_non_contiguous(self, tokens): + from lmdeploy.pytorch.kernels.cuda import apply_rotary_pos_emb + + query = torch.randn(1, tokens, 32, 128, device='cuda', dtype=torch.bfloat16)[..., :64] + key = torch.randn(1, tokens, 128, device='cuda', dtype=torch.bfloat16)[..., :64][..., None, :] + positions = torch.linspace(0, 1048575, tokens, device='cuda') + inv_freq = 1 / (8000000**(torch.arange(0, 64, 2, device='cuda').float() / 64)) + angles = positions[:, None] * inv_freq[None, :] + cos = angles.cos().to(torch.bfloat16) + sin = angles.sin().to(torch.bfloat16) + pair_cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) + pair_sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) + q_gt = query * pair_cos + _rotate_complex(query) * pair_sin + k_gt = key * pair_cos + _rotate_complex(key) * pair_sin + + q_embed, k_embed = apply_rotary_pos_emb(query, key, cos, sin, complex_mode=True) + + assert torch.equal(q_embed, q_gt) + assert torch.equal(k_embed, k_gt) diff --git a/tests/pytorch/kernel/test_ds_index.py b/tests/pytorch/kernel/test_ds_index.py index 7181296ea9..16fea39867 100644 --- a/tests/pytorch/kernel/test_ds_index.py +++ b/tests/pytorch/kernel/test_ds_index.py @@ -145,6 +145,24 @@ def test_fp8_index(self, q, q_s, k_cache, k_s_cache, cu_seqlen_q, k_seqlens, blo from lmdeploy.pytorch.kernels.cuda.ds_index import fp8_index fp8_index(q, q_s, k_cache, k_s_cache, cu_seqlen_q, k_seqlens, block_offset) + def test_dense_index(self): + from lmdeploy.pytorch.kernels.cuda.ds_index import dense_index + + q_seqlens = torch.tensor([3, 2], dtype=torch.int32, device='cuda') + k_seqlens = torch.tensor([5, 2], dtype=torch.int32, device='cuda') + output = dense_index(q_seqlens, k_seqlens, max_q_seqlen=3, topk=8) + expected = torch.tensor( + [[0, 1, 2, -1, -1, -1, -1, -1], + [0, 1, 2, 3, -1, -1, -1, -1], + [0, 1, 2, 3, 4, -1, -1, -1], + [0, -1, -1, -1, -1, -1, -1, -1], + [0, 1, -1, -1, -1, -1, -1, -1], + [-1, -1, -1, -1, -1, -1, -1, -1]], + dtype=torch.int32, + device='cuda', + ) + torch.testing.assert_close(output, expected) + def test_fp8_index_trim_causal_tail_with_raw_lengths(self, num_heads, head_dim, block_size, device): """Trimmed V4-style causal scoring must preserve every visible score.""" diff --git a/tests/pytorch/kernel/test_dsa_indexer_preprocess.py b/tests/pytorch/kernel/test_dsa_indexer_preprocess.py new file mode 100644 index 0000000000..b1abd86707 --- /dev/null +++ b/tests/pytorch/kernel/test_dsa_indexer_preprocess.py @@ -0,0 +1,134 @@ +import pytest +import torch +import torch.nn.functional as F + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA required') + + +def _apply_rope_first(x, cos, sin, rope_interleaved): + out = x.clone() + x_rope = x[..., :64] + if rope_interleaved: + pair_cos = cos[..., :32].repeat_interleave(2, dim=-1) + pair_sin = sin[..., :32].repeat_interleave(2, dim=-1) + while pair_cos.dim() < x_rope.dim(): + pair_cos = pair_cos.unsqueeze(-2) + pair_sin = pair_sin.unsqueeze(-2) + rotated = torch.empty_like(x_rope) + rotated[..., ::2] = -x_rope[..., 1::2] + rotated[..., 1::2] = x_rope[..., ::2] + out[..., :64] = x_rope * pair_cos + rotated * pair_sin + else: + cos = cos.unsqueeze(-2) if x_rope.dim() == 3 else cos + sin = sin.unsqueeze(-2) if x_rope.dim() == 3 else sin + x_l, x_h = x_rope.split(32, dim=-1) + out[..., :32] = x_l * cos[..., :32] - x_h * sin[..., :32] + out[..., 32:64] = x_h * cos[..., 32:64] + x_l * sin[..., 32:64] + return out + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9, + reason='requires device with cc>=9.0') +@pytest.mark.parametrize('rope_interleaved', [True, False]) +@pytest.mark.parametrize('heads', [32, 64]) +def test_prepare_dsa_indexer_q_matches_unfused_quantization(rope_interleaved, heads): + from lmdeploy.pytorch.kernels.cuda.apply_rotary_pos_emb import apply_rotary_pos_emb + from lmdeploy.pytorch.kernels.cuda.blocked_gemm_fp8 import per_token_group_quant_fp8 + from lmdeploy.pytorch.kernels.cuda.dsa_indexer_preprocess import prepare_dsa_indexer_q + + torch.manual_seed(0) + tokens = 11 + q = torch.randn(tokens, heads, 128, device='cuda', dtype=torch.bfloat16) + weights = torch.randn(tokens, heads, device='cuda', dtype=torch.bfloat16) + cos = torch.randn(tokens, 64, device='cuda', dtype=torch.bfloat16) + sin = torch.randn(tokens, 64, device='cuda', dtype=torch.bfloat16) + score_scale = 0.03125 + + q_out, q_scale = prepare_dsa_indexer_q(q, + weights, + cos, + sin, + score_scale, + torch.float8_e4m3fn, + rope_interleaved=rope_interleaved) + q_pe, q_nope = q.split([64, 64], dim=-1) + k_pe = q.new_empty(tokens, 1, 64) + if rope_interleaved: + cos = cos[..., :32] + sin = sin[..., :32] + q_pe, _ = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, complex_mode=rope_interleaved) + q_ref = torch.cat([q_pe, q_nope], dim=-1) + q_ref, scale_ref = per_token_group_quant_fp8(q_ref.reshape(-1, 128), + 128, + dtype=torch.float8_e4m3fn, + scale_fmt='ue8m0') + q_ref = q_ref.reshape_as(q) + scale_ref = scale_ref.reshape_as(weights) + + assert torch.equal(q_out, q_ref) + torch.testing.assert_close(q_scale, scale_ref * weights.float() * score_scale, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9, + reason='requires device with cc>=9.0') +@pytest.mark.parametrize('rope_interleaved', [True, False]) +@pytest.mark.parametrize('q_seqlens,kv_seqlens', [([1, 1], [3, 2]), ([3, 2], [5, 3])]) +def test_prepare_dsa_indexer_k_cache_matches_prepared_k_fill(q_seqlens, kv_seqlens, rope_interleaved): + from lmdeploy.pytorch.kernels.cuda.dsa_indexer_preprocess import prepare_dsa_indexer_k, prepare_dsa_indexer_k_cache + from lmdeploy.pytorch.kernels.cuda.fill_kv_cache import fill_kv_cache_blocked_fp8 + + torch.manual_seed(1) + block_size = 4 + total_tokens = sum(q_seqlens) + k = torch.randn(total_tokens, 128, device='cuda', dtype=torch.bfloat16) + norm_weight = torch.randn(128, device='cuda', dtype=torch.float32) + norm_bias = torch.randn(128, device='cuda', dtype=torch.float32) + cos = torch.randn(total_tokens, 64, device='cuda', dtype=torch.bfloat16) + sin = torch.randn(total_tokens, 64, device='cuda', dtype=torch.bfloat16) + cu_seqlen_q = torch.tensor([0] + q_seqlens, device='cuda', dtype=torch.int32).cumsum(0) + kv_seqlens_t = torch.tensor(kv_seqlens, device='cuda', dtype=torch.int32) + block_offsets = torch.tensor([[0, 1], [2, 3]], device='cuda', dtype=torch.int32) + cache_ref = torch.zeros(4, block_size, 128, device='cuda', dtype=torch.float8_e4m3fn) + scale_ref = torch.zeros(4, block_size, 1, device='cuda', dtype=torch.float32) + cache_fused = torch.zeros_like(cache_ref) + scale_fused = torch.zeros_like(scale_ref) + + k_prepared = prepare_dsa_indexer_k(k, + norm_weight, + norm_bias, + cos, + sin, + eps=1e-6, + rope_interleaved=rope_interleaved) + fill_kv_cache_blocked_fp8(k_prepared[:, None], + None, + cache_ref[..., None, :], + None, + scale_ref[..., None, :], + None, + cu_seqlen_q=cu_seqlen_q, + kv_seqlens=kv_seqlens_t, + max_q_seqlen=max(q_seqlens), + block_offsets=block_offsets, + group_size=128, + scale_fmt='ue8m0') + prepare_dsa_indexer_k_cache(k, + norm_weight, + norm_bias, + cos, + sin, + cache_fused, + scale_fused[..., 0], + cu_seqlen_q, + kv_seqlens_t, + block_offsets, + max_q_seqlen=max(q_seqlens), + eps=1e-6, + rope_interleaved=rope_interleaved) + + assert torch.equal(cache_fused, cache_ref) + assert torch.equal(scale_fused, scale_ref) + + k_torch = F.layer_norm(k.float(), (128, ), norm_weight, norm_bias, 1e-6).to(torch.bfloat16) + k_torch = _apply_rope_first(k_torch, cos, sin, rope_interleaved).to(torch.bfloat16) + torch.testing.assert_close(k_prepared.float(), k_torch.float(), rtol=0, atol=0.0625) diff --git a/tests/pytorch/kernel/test_fuse_moe_blocked_fp8.py b/tests/pytorch/kernel/test_fuse_moe_blocked_fp8.py index d3588026d8..81e79078bf 100644 --- a/tests/pytorch/kernel/test_fuse_moe_blocked_fp8.py +++ b/tests/pytorch/kernel/test_fuse_moe_blocked_fp8.py @@ -444,3 +444,79 @@ def test_fused_moe_blocked_fp8_compact_down_matches_fused_moe_with_local_experts norm_out = output / out_max norm_gt = gt / gt_max torch.testing.assert_close(norm_out, norm_gt, atol=0.05, rtol=1e-3) + + +@pytest.mark.skipif(torch.cuda.get_device_capability()[0] < 9, reason='require device with cc>=9.0') +@torch.inference_mode() +def test_fused_moe_blocked_fp8_compact_both_matches_fused_moe(): + from lmdeploy.pytorch.kernels.cuda.blocked_fp8_fused_moe import ( + _should_use_compact_blocked_fp8_moe_both_by_shape, + fused_moe_blocked_fp8, + ) + from lmdeploy.pytorch.kernels.cuda.fused_moe import fused_moe + + torch.manual_seed(0) + device = torch.device('cuda') + dtype = torch.float16 + quant_dtype = torch.float8_e4m3fn + group_size = 128 + + seq_len = 128 + in_size = 128 + hidden_size = 256 + out_size = 128 + num_experts = 256 + top_k = 8 + renormalize = True + + hidden_states, states_quanted, states_scale = _make_A(seq_len, + in_size, + group_size=group_size, + out_dtype=quant_dtype, + device=device) + w1, w1_quanted, w1_scale, _ = _make_B(num_experts, + in_size, + hidden_size, + group_size=group_size, + out_dtype=quant_dtype, + device=device) + w2, w2_quanted, w2_scale, _ = _make_B(num_experts, + hidden_size // 2, + out_size, + group_size=group_size, + out_dtype=quant_dtype, + device=device) + + router_logits = torch.rand(seq_len, num_experts, dtype=dtype, device=device) + routing_weights = torch.softmax(router_logits, dim=-1, dtype=torch.float32) + topk_weights, topk_ids = torch.topk(routing_weights, top_k, dim=-1) + + assert _should_use_compact_blocked_fp8_moe_both_by_shape(seq_len, topk_ids.numel(), num_experts, num_experts) + + gt = fused_moe(hidden_states.to(dtype), + w1.to(dtype), + w2.to(dtype), + topk_weights, + topk_ids, + topk=top_k, + num_experts=num_experts, + renormalize=renormalize) + output = fused_moe_blocked_fp8(states_quanted, + states_scale, + w1_quanted, + w1_scale, + w2_quanted, + w2_scale, + topk_weights, + topk_ids, + topk=top_k, + num_experts=num_experts, + renormalize=renormalize) + + out_max = output.abs().max() + gt_max = gt.abs().max() + assert (out_max - gt_max).abs() / out_max < 0.05 + + norm_out = output / out_max + norm_gt = gt / gt_max + torch.testing.assert_close(norm_out, norm_gt, atol=0.05, rtol=1e-3) diff --git a/tests/pytorch/kernel/test_fused_moe.py b/tests/pytorch/kernel/test_fused_moe.py index d74c56d585..005a2dcc10 100644 --- a/tests/pytorch/kernel/test_fused_moe.py +++ b/tests/pytorch/kernel/test_fused_moe.py @@ -71,6 +71,39 @@ def test_compact_blocked_fp8_configs_use_average_routes(num_routes, block_m): num_stages=3) +@pytest.mark.parametrize(('num_routes', 'gate_out_features', 'block_m', 'block_n'), [ + (256 * 4, 512, 16, 64), + (256 * 8, 512, 16, 64), + (256 * 16, 512, 32, 128), + (256 * 32, 512, 64, 128), + (256 * 4, 4096, 64, 128), +]) +def test_compact_blocked_fp8_both_configs(num_routes, gate_out_features, block_m, block_n): + from lmdeploy.pytorch.kernels.cuda.blocked_fp8_fused_moe import _compact_blocked_fp8_moe_both_config + + assert _compact_blocked_fp8_moe_both_config(num_routes, 256, gate_out_features) == dict( + block_m=block_m, block_n=block_n, num_warps=4, num_stages=3) + + +@pytest.mark.parametrize(('num_tokens', 'num_routes', 'num_experts', 'local_experts', 'expected'), [ + (64, 256 * 4, 256, 256, False), + (128, 256 * 3, 256, 256, False), + (128, 256 * 4, 256, 256, True), + (1536, 256 * 48, 256, 256, True), + (2048, 256 * 64, 256, 256, False), + (128, 256 * 4, 512, 256, False), + (128, 256 * 4, 256, 128, False), +]) +def test_compact_blocked_fp8_both_policy_uses_measured_route_density(num_tokens, num_routes, num_experts, + local_experts, expected): + from lmdeploy.pytorch.kernels.cuda.blocked_fp8_fused_moe import ( + _should_use_compact_blocked_fp8_moe_both_by_shape, + ) + + assert _should_use_compact_blocked_fp8_moe_both_by_shape( + num_tokens, num_routes, num_experts, local_experts) is expected + + @pytest.mark.parametrize(('num_tokens', 'num_routes', 'origin_ctas', 'compact_ctas'), [ (65, 650, 512 * 2 * 32, 512 * 1 * 32), (1024, 512 * 20, 512 * 16 * 32, 512 * 1 * 32), diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py new file mode 100644 index 0000000000..c0e1e8d92a --- /dev/null +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -0,0 +1,144 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lmdeploy.pytorch.backends.cuda.attention.mla import FlashMLAImpl, NSAIndicesUpdater +from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend + + +def test_nsa_decode_indices_update_repeats_block_table_for_spec_decode(): + updater = NSAIndicesUpdater() + nsa_indices = torch.tensor([[0, 17, -1], [32, 1, 16], [0, 33, 47], [32, 1, 16]]) + block_offsets = torch.tensor([[100, 101, 102], [200, 201, 202]]) + + output = updater._update_decode_impl(nsa_indices, block_offsets, max_q_seqlen=2, block_size=16) + + expected = torch.tensor([[[1600, 1617, -1], [1632, 1601, 1616]], + [[3200, 3233, 3247], [3232, 3201, 3216]]]) + assert torch.equal(output, expected) + + +def test_nsa_decode_indices_update_keeps_single_token_shape(): + updater = NSAIndicesUpdater() + nsa_indices = torch.tensor([[0, 17], [32, -1]]) + block_offsets = torch.tensor([[100, 101, 102], [200, 201, 202]]) + + output = updater._update_decode_impl(nsa_indices, block_offsets, max_q_seqlen=1, block_size=16) + + expected = torch.tensor([[[1600, 1617]], [[3232, -1]]]) + assert torch.equal(output, expected) + + +def test_bf16_sparse_decode_uses_strided_cache_view(): + impl = object.__new__(FlashMLAImpl) + impl.nsa_updater = NSAIndicesUpdater() + impl.nsa_updater._update_decode_strided_func = impl.nsa_updater._update_decode_strided_impl + impl._flash_mla_sparse = Mock(return_value=torch.empty(4, 64, 512, dtype=torch.bfloat16)) + + query = torch.empty(4, 64, 576, dtype=torch.bfloat16) + block_size = 16 + block_elements = block_size * 576 + storage = torch.empty(3, block_elements + 128, dtype=torch.bfloat16) + k_cache = storage[:, :block_elements].view(3, block_size, 1, 576) + nsa_indices = torch.tensor([[0, 17], [32, -1], [0, 33], [32, 1]]) + metadata = SimpleNamespace(is_decoding=True, + q_seqlens=torch.tensor([2, 2]), + block_offsets=torch.tensor([[1, 2, 0], [2, 0, 1]])) + + impl._decode_bf16_sparse_flash_mla(query, k_cache, nsa_indices, metadata) + + sparse_query, storage_k, global_indices = impl._flash_mla_sparse.call_args.args + assert sparse_query is query + assert storage_k.untyped_storage().data_ptr() == k_cache.untyped_storage().data_ptr() + assert storage_k.stride() == (64, 576, 1) + expected = torch.tensor([[[146, 301]], [[0, -1]], [[292, 155]], [[146, 301]]]) + assert torch.equal(global_indices, expected) + + +def test_bf16_sparse_decode_strided_cache_matches_contiguous_cache(): + pytest.importorskip('flash_mla') + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + pytest.skip('FlashMLA BF16 sparse attention requires an SM90 GPU') + + impl = object.__new__(FlashMLAImpl) + impl.scale = 576**-0.5 + impl.flash_mla_sparse_fwd = None + impl.nsa_updater = NSAIndicesUpdater() + impl.nsa_updater._update_decode_strided_func = impl.nsa_updater._update_decode_strided_impl + + batch_size = 2 + query_len = 2 + block_size = 64 + num_blocks = 4 + block_elements = block_size * 576 + storage = torch.empty(num_blocks, block_elements + 128, dtype=torch.bfloat16, device='cuda') + k_cache = storage[:, :block_elements].view(num_blocks, block_size, 1, 576) + k_cache.normal_(std=0.1) + query = torch.randn(batch_size * query_len, 64, 576, dtype=torch.bfloat16, device='cuda') + nsa_indices = torch.arange(128, dtype=torch.int32, device='cuda').repeat(batch_size * query_len, 1) + block_offsets = torch.tensor([[2, 0], [3, 1]], dtype=torch.int32, device='cuda') + metadata = SimpleNamespace(is_decoding=True, + q_seqlens=torch.full((batch_size, ), query_len, dtype=torch.int32, device='cuda'), + block_offsets=block_offsets) + + output = impl._decode_bf16_sparse_flash_mla(query, k_cache, nsa_indices, metadata) + + contiguous_k = k_cache.flatten(0, 1) + contiguous_indices = impl.nsa_updater._update_decode_impl(nsa_indices, block_offsets, query_len, block_size) + contiguous_indices = contiguous_indices.flatten(0, 1)[:, None] + expected = impl._flash_mla_sparse(query, contiguous_k, contiguous_indices) + torch.testing.assert_close(output, expected) + + +def test_fp8_sparse_decode_pads_tp_query_heads_for_aligned_kernel(): + impl = object.__new__(FlashMLAImpl) + impl.causal = True + impl.scale = 1.0 + impl.v_head_size = 512 + impl.nsa_updater = Mock() + impl.nsa_updater.update_decode.return_value = torch.zeros(2, 3, 4, dtype=torch.int32) + impl.flash_mla_with_kvcache = Mock( + return_value=(torch.empty(2, 3, 64, 512, dtype=torch.bfloat16), None)) + + query = torch.empty(6, 8, 576, dtype=torch.bfloat16) + k_cache = torch.empty(2, 16, 1, 656, dtype=torch.uint8) + metadata = SimpleNamespace( + q_seqlens=torch.tensor([3, 3]), + kv_seqlens=torch.tensor([16, 16]), + block_offsets=torch.zeros(2, 1, dtype=torch.int32), + tile_scheduler_metadata=object(), + num_splits=None, + ) + + output = impl._decode_paged_flash_mla(query, k_cache, torch.zeros(6, 4, dtype=torch.int32), metadata) + + padded_query = impl.flash_mla_with_kvcache.call_args.args[0] + assert padded_query.shape == (2, 3, 64, 576) + assert output.shape == (6, 8, 512) + + +def test_bf16_sparse_decode_skips_fp8_flashmla_metadata(): + metadata = SimpleNamespace(block_offsets=torch.tensor([[0, 1]], dtype=torch.int64)) + model_config = SimpleNamespace(use_mla_fp8_cache=False, mla_index_topk=2048) + + CudaOpsBackend.update_meta_flashmla(metadata, model_config, decoding_query_len=5) + + assert metadata.block_offsets.dtype == torch.int32 + assert not hasattr(metadata, 'tile_scheduler_metadata') + + +def test_nsa_topk_uses_per_query_causal_kv_lengths(): + if not torch.cuda.is_available(): + pytest.skip('requires a CUDA runtime to import the NSA backend') + from lmdeploy.pytorch.backends.cuda.nsa import _get_causal_k_seqlens + + q_seqlens = torch.tensor([2, 3]) + cu_seqlen_q = torch.tensor([0, 2, 5]) + k_seqlens = torch.tensor([5, 7]) + + output = _get_causal_k_seqlens(cu_seqlen_q, q_seqlens, k_seqlens, num_tokens=5) + + assert torch.equal(output, torch.tensor([4, 5, 5, 6, 7])) diff --git a/tests/pytorch/kernel/test_sparse_index_topk.py b/tests/pytorch/kernel/test_sparse_index_topk.py index 60a852e5fe..31c1a91a16 100644 --- a/tests/pytorch/kernel/test_sparse_index_topk.py +++ b/tests/pytorch/kernel/test_sparse_index_topk.py @@ -60,6 +60,19 @@ def test_sparse_index_topk_matches_torch_topk_and_fill(): _assert_topk_ids(scores, out, seqlens, k, fill=fill) +def test_sparse_index_topk_glm52_int64_seqlens(): + from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import sparse_index_topk + + k = 2048 + seqlens = [2049, 3072] + scores = torch.randn(len(seqlens), 4096, device='cuda') + q_seqlens = torch.ones(len(seqlens), device='cuda', dtype=torch.int64) + kv_seqlens = torch.tensor(seqlens, device='cuda', dtype=torch.int64) + + out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k) + _assert_topk_ids(scores, out, seqlens, k) + + def test_sparse_index_topk_accepts_padded_score_stride(): from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import sparse_index_topk diff --git a/tests/pytorch/nn/test_nsa.py b/tests/pytorch/nn/test_nsa.py new file mode 100644 index 0000000000..2d96d61d7f --- /dev/null +++ b/tests/pytorch/nn/test_nsa.py @@ -0,0 +1,28 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lmdeploy.pytorch.nn import nsa + + +@pytest.mark.parametrize('query_width', [1, 5]) +def test_indexer_meta_preserves_decoding_query_width(monkeypatch, query_width): + batch_size = 2 + step_ctx = SimpleNamespace( + cache_config=SimpleNamespace(block_size=64, num_gpu_blocks=16)) + ctx_mgr = SimpleNamespace(current_context=lambda: step_ctx) + monkeypatch.setattr(nsa, 'get_step_ctx_manager', lambda: ctx_mgr) + attn_metadata = SimpleNamespace( + is_decoding=True, + max_q_seqlen=None, + kv_flatten_size=None, + cu_seqlens_q=torch.arange(batch_size + 1) * query_width, + q_seqlens=torch.full((batch_size, ), query_width), + kv_seqlens=torch.full((batch_size, ), query_width), + block_offsets=torch.zeros(batch_size, 1, dtype=torch.int32)) + q = torch.empty(batch_size * query_width, 1, 128) + + meta = nsa.IndexerTopKFP8._build_meta(q, attn_metadata) + + assert meta.max_q_seqlen == query_width diff --git a/tests/pytorch/paging/test_block_trie.py b/tests/pytorch/paging/test_block_trie.py index 3ccb9f777b..e94210d8b1 100644 --- a/tests/pytorch/paging/test_block_trie.py +++ b/tests/pytorch/paging/test_block_trie.py @@ -469,7 +469,7 @@ def test_existing_node_can_be_enriched_with_routed_experts(self, block_trie, blo block_trie.match(matched) assert np.array_equal(matched.all_routed_experts.get_real(), experts) - def test_match_does_not_partially_replay_routed_experts(self, block_trie, block_mgr, scheduler): + def test_match_stops_before_block_missing_routed_experts(self, block_trie, block_mgr, scheduler): sess = scheduler.add_session(0) block_size = sess.seq_meta.block_size token_ids = [1] * block_size + [2] * block_size + [3] @@ -483,8 +483,10 @@ def test_match_does_not_partially_replay_routed_experts(self, block_trie, block_ matched = sess.add_sequence(token_ids, sampling_param=SamplingParam(return_routed_experts=True)) block_trie.match(matched) - assert matched.num_history_ids == block_size * 2 - assert len(matched.all_routed_experts) == 0 + assert matched.num_history_ids == block_size + assert np.array_equal(matched.all_routed_experts.get_real(), self._routed_experts(block_size)) + assert matched.prefix_cache.private_recompute_start_step == block_size + assert matched.prefix_cache.private_recompute_end_step == block_size * 2 def test_missing_replay_does_not_enrich_from_misaligned_tail(self, block_trie, block_mgr, scheduler): sess = scheduler.add_session(0) @@ -501,7 +503,7 @@ def test_missing_replay_does_not_enrich_from_misaligned_tail(self, block_trie, b matched = sess.add_sequence(token_ids, sampling_param=SamplingParam(return_routed_experts=True)) block_trie.match(matched) - assert matched.num_history_ids == block_size * 2 + assert matched.num_history_ids == 0 assert len(matched.all_routed_experts) == 0 matched.append_routed_experts(self._routed_experts(1, offset=1000)) diff --git a/tests/pytorch/spec_decode/test_cudagraph_strategy.py b/tests/pytorch/spec_decode/test_cudagraph_strategy.py index c8d03dc46b..c72b9b97e0 100644 --- a/tests/pytorch/spec_decode/test_cudagraph_strategy.py +++ b/tests/pytorch/spec_decode/test_cudagraph_strategy.py @@ -132,3 +132,47 @@ def make_attn_metadata(batch_size: int, query_len: int): inputs_embeds=None, ) assert key_hidden32 == key_qlen4 + + +def test_cuda_graph_key_separates_dsa_seed_and_reuse(monkeypatch): + from types import SimpleNamespace + + import torch + + from lmdeploy.pytorch.backends.cuda import graph_runner as cuda_graph_runner + + runner = cuda_graph_runner.CUDAGraphRunner.__new__(cuda_graph_runner.CUDAGraphRunner) + runner.ctx_mgr = SimpleNamespace(current_context=lambda: SimpleNamespace(global_is_decoding=lambda: True)) + runner.get_meta = lambda: SimpleNamespace(padding_batch_size=None) + runner._get_capture_tokens = lambda batch_size: batch_size + monkeypatch.setattr(cuda_graph_runner, 'get_step_ctx_manager', + lambda: SimpleNamespace(current_context=lambda: SimpleNamespace(enable_microbatch=False))) + + input_ids = torch.zeros((1, 8), dtype=torch.long) + kwargs = dict( + input_ids=input_ids, + position_ids=torch.zeros_like(input_ids), + past_key_values=[], + attn_metadata=SimpleNamespace(q_seqlens=torch.ones(8, dtype=torch.long)), + inputs_embeds=None, + ) + + seed_key = runner.get_graph_key(**kwargs, + skip_topk=False, + use_dense_index=False) + reuse_key = runner.get_graph_key(**kwargs, + skip_topk=True, + use_dense_index=False) + reuse_dense_key = runner.get_graph_key(**kwargs, + skip_topk=True, + use_dense_index=True) + dense_key = runner.get_graph_key(**kwargs, + skip_topk=False, + use_dense_index=True) + + assert seed_key != reuse_key + assert seed_key != dense_key + assert reuse_key == reuse_dense_key + assert seed_key[-2:] == (False, False) + assert reuse_key[-2:] == (True, False) + assert dense_key[-2:] == (False, True) diff --git a/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py new file mode 100644 index 0000000000..e5e42cad78 --- /dev/null +++ b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py @@ -0,0 +1,190 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import asyncio + +import pytest +import torch +from torch import nn + +from lmdeploy.pytorch.models.deepseek_v32 import DSATopKIndicesBuffer +from lmdeploy.pytorch.models.glm_moe_dsa_mtp import GlmMoeDsaMultiTokenPredictor +from lmdeploy.pytorch.spec_decode.proposers.base import BaseSpecProposer +from lmdeploy.pytorch.spec_decode.proposers.deepseek_mtp import DeepseekMTP +from lmdeploy.pytorch.strategies.ar_spec.model_agent import ARSpecExtraInputs + + +class _GuidedHelper: + + async def prepare_bitmask(self, logits, processors): + return None + + def apply_bitmask(self, logits, bitmask): + raise AssertionError('no bitmask was returned') + + async def accept_draft_tokens(self, draft_token_ids, processors): + return None + + +class _DummyDraft(nn.Module): + + uses_dsa_topk_buffer = True + + def __init__(self): + super().__init__() + self.compacted = None + self.prepare_calls = 0 + + def compact_topk_indices(self, row_indices): + self.compacted = row_indices + + def prepare_hidden_states_for_logits(self, hidden_states): + self.prepare_calls += 1 + return hidden_states + 10 + + +class _DummyTarget(nn.Module): + + def get_logits(self, hidden_states): + self.hidden_states = hidden_states + logits = hidden_states.new_zeros(1, hidden_states.size(1), 8) + logits[..., 5] = 1 + return logits + + +def test_proposer_reuses_topk_and_recycles_postnorm_hidden_states(): + proposer = object.__new__(DeepseekMTP) + proposer.model = _DummyDraft() + proposer.target_model = _DummyTarget() + proposer.guided_helper = _GuidedHelper() + hidden_states = torch.arange(12, dtype=torch.float32).view(1, 3, 4) + outputs = dict(hidden_states=hidden_states, model_metas=[{'keep': 1}, {'keep': 2}]) + extra_inputs = ARSpecExtraInputs(last_token_indices=torch.tensor([2, 0])) + + draft_ids, model_metas, recurrent_hidden = asyncio.run( + proposer.get_outputs(outputs, model_inputs=None, extra_inputs=extra_inputs)) + + selected_hidden = hidden_states[:, [2, 0]] + assert draft_ids.tolist() == [[5], [5]] + assert torch.equal(recurrent_hidden, selected_hidden + 10) + assert torch.equal(proposer.target_model.hidden_states, selected_hidden + 10) + assert proposer.model.prepare_calls == 1 + assert proposer.model.compacted is extra_inputs.last_token_indices + assert [meta['keep'] for meta in model_metas] == [1, 2] + assert all(meta['skip_topk'] for meta in model_metas) + + +def test_proposer_binds_target_embedding_and_keeps_draft_topk_buffer(monkeypatch): + + class _Draft(nn.Module): + + uses_dsa_topk_buffer = True + + def __init__(self): + super().__init__() + self.embed_tokens = None + self.topk_indices_buffer = DSATopKIndicesBuffer(topk=2) + + def set_input_embeddings(self, embed_tokens): + self.embed_tokens = embed_tokens + + def get_input_embeddings(self): + return self.embed_tokens + + class _Target(nn.Module): + + def __init__(self): + super().__init__() + self.embedding = nn.Embedding(8, 4) + self.model = nn.Module() + self.model.topk_indices_buffer = DSATopKIndicesBuffer(topk=2) + + def get_input_embeddings(self): + return self.embedding + + draft = _Draft() + draft_topk_indices_buffer = draft.topk_indices_buffer + target = _Target() + proposer = object.__new__(DeepseekMTP) + + def build_model(self, empty_init, target_model=None, build_model_ctx=None): + self.model = draft + self.target_model = target_model + + monkeypatch.setattr(BaseSpecProposer, 'build_model', build_model) + proposer.build_model(empty_init=True, target_model=target) + + assert draft.embed_tokens is target.embedding + assert draft.topk_indices_buffer is draft_topk_indices_buffer + assert draft.topk_indices_buffer is not target.model.topk_indices_buffer + + +@pytest.mark.parametrize(('uses_topk_buffer', 'supports_binding'), + [(False, True), (True, False)]) +def test_proposer_skips_unsupported_target_embedding_binding( + monkeypatch, uses_topk_buffer, supports_binding): + + class _Draft(nn.Module): + + uses_dsa_topk_buffer = uses_topk_buffer + + proposer = object.__new__(DeepseekMTP) + draft = _Draft() + if supports_binding: + draft.set_input_embeddings = lambda _: pytest.fail( + 'disabled top-k sharing must not bind embeddings') + + def build_model(self, empty_init, target_model=None, build_model_ctx=None): + self.model = draft + self.target_model = target_model + + monkeypatch.setattr(BaseSpecProposer, 'build_model', build_model) + proposer.build_model(empty_init=True, target_model=nn.Module()) + + +def test_glm_mtp_prepares_postnorm_hidden_for_logits(): + + class _SharedHead(nn.Module): + + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, hidden_states): + self.calls += 1 + return hidden_states + 10 + + class _Layer(nn.Module): + + def __init__(self): + super().__init__() + self.shared_head = _SharedHead() + + def forward(self, input_ids, position_ids, previous_hidden_states, past_key_value, **kwargs): + return previous_hidden_states + 3 + + predictor = object.__new__(GlmMoeDsaMultiTokenPredictor) + nn.Module.__init__(predictor) + predictor.mtp_start_layer_idx = 5 + predictor.num_mtp_layers = 1 + predictor.layers = nn.ModuleDict({'5': _Layer()}) + previous_hidden = torch.zeros(1, 1, 4) + + raw_hidden = predictor(torch.zeros(1, 1, dtype=torch.long), + torch.zeros(1, 1, dtype=torch.long), + previous_hidden, + past_key_values=[[None]], + inputs_embeds=torch.zeros_like(previous_hidden)) + logits_hidden = predictor.prepare_hidden_states_for_logits(raw_hidden) + + expected = previous_hidden + 13 + assert torch.equal(logits_hidden, expected) + assert predictor.layers['5'].shared_head.calls == 1 + + +def test_topk_buffer_compaction_preserves_selected_order(): + buffer = DSATopKIndicesBuffer(topk=2) + rows = torch.tensor([[0, 1], [10, 11], [20, 21]], dtype=torch.int32) + buffer.write(rows) + + compacted = buffer.compact(torch.tensor([2, 0])) + + assert torch.equal(compacted, rows[[2, 0]]) diff --git a/tests/pytorch/test_model_inputs.py b/tests/pytorch/test_model_inputs.py index b947180eb6..c9547cf248 100644 --- a/tests/pytorch/test_model_inputs.py +++ b/tests/pytorch/test_model_inputs.py @@ -48,3 +48,28 @@ def test_step_context_global_is_decoding_uses_dp_global_state(): dp_meta=DPMeta(dp_is_decoding=False), ) assert not step_ctx.global_is_decoding() + + +def test_step_context_single_token_decode_metadata(monkeypatch): + import lmdeploy.pytorch.model_inputs as model_inputs + + monkeypatch.setattr(model_inputs, 'get_backend', + lambda: type('Backend', (), {'update_step_context': staticmethod(lambda ctx: ctx)})()) + inputs = ModelInputs( + input_ids=torch.tensor([[10, 11, 12]]), + seq_length=torch.ones(3, dtype=torch.long), + history_lengths=torch.tensor([4, 9, 11]), + block_offsets=torch.zeros((3, 1), dtype=torch.long), + is_decoding=True, + num_ignored_history=torch.tensor([0, 2, 3]), + max_q_seqlen=1, + max_kv_seqlen=12, + sum_kv_seqlen=22, + ) + + context = StepContext.new(inputs, model_config=None, cache_config=None) + + torch.testing.assert_close(context.attention_mask, torch.ones((3, 1), dtype=torch.long)) + torch.testing.assert_close(context.position_ids, torch.tensor([[4, 9, 11]])) + torch.testing.assert_close(context.q_start_loc, torch.tensor([0, 1, 2])) + torch.testing.assert_close(context.kv_seqlens, torch.tensor([5, 8, 9])) diff --git a/tests/test_lmdeploy/serve/parsers/test_glm47_parser.py b/tests/test_lmdeploy/serve/parsers/test_glm47_parser.py index e85b716592..6dbf0004d3 100644 --- a/tests/test_lmdeploy/serve/parsers/test_glm47_parser.py +++ b/tests/test_lmdeploy/serve/parsers/test_glm47_parser.py @@ -10,6 +10,7 @@ from .helpers import first_stream_delta MODEL_ID = 'zai-org/GLM-4.7' +GLM52_MODEL_ID = 'zai-org/GLM-5.2-FP8' @pytest.fixture() @@ -53,6 +54,43 @@ def response_parser_with_reasoning(): ] +def _make_response_parser_with_reasoning(chat_template_kwargs=None): + cls = ResponseParserManager.get('default') + cls.reasoning_parser_cls = ReasoningParserManager.get('default') + cls.tool_parser_cls = ToolParserManager.get('glm47') + request = ChatCompletionRequest( + model=GLM52_MODEL_ID, + messages=[], + stream=True, + tool_choice='auto', + chat_template_kwargs=chat_template_kwargs or {}, + ) + return cls(request=request) + + +def _collect_stream(parser, chunks): + reasoning_seen = [] + content_seen = [] + emitted_name = None + emitted_args = '' + + for chunk in chunks: + for delta, tool_emitted in parser.stream_chunk(delta_text=chunk, delta_token_ids=[]): + if delta is not None: + if delta.reasoning_content: + reasoning_seen.append(delta.reasoning_content) + if delta.content: + content_seen.append(delta.content) + if tool_emitted and delta and delta.tool_calls: + for call in delta.tool_calls: + if call.function and call.function.name: + emitted_name = call.function.name + if call.function and call.function.arguments: + emitted_args += call.function.arguments + + return ''.join(reasoning_seen), ''.join(content_seen), emitted_name, emitted_args + + class TestGlm47ResponseParserStreaming: """Integration tests for ResponseParser.stream_chunk with glm47 tool parser.""" @@ -151,6 +189,60 @@ def test_stream_chunk_mixed_default_reasoning_and_glm47_tool(self, response_pars assert emitted_name == 'get_weather' assert emitted_args == '{"location": "Beijing"}' + def test_stream_chunk_tool_start_ends_reasoning_without_close_tag(self): + parser = _make_response_parser_with_reasoning() + chunks = [ + '', + 'first reason', + 'get_weather', + 'locationBeijing', + '', + ] + + reasoning_seen, content_seen, emitted_name, emitted_args = _collect_stream(parser, chunks) + + assert reasoning_seen == 'first reason' + assert content_seen == '' + assert emitted_name == 'get_weather' + assert emitted_args == '{"location": "Beijing"}' + assert parser.validate_complete() is True + + def test_stream_chunk_split_tool_start_ends_reasoning_without_close_tag(self): + parser = _make_response_parser_with_reasoning() + chunks = [ + 'first reasonget_weatherlocation', + 'Beijing', + ] + + reasoning_seen, content_seen, emitted_name, emitted_args = _collect_stream(parser, chunks) + + assert reasoning_seen == 'first reason' + assert content_seen == '' + assert emitted_name == 'get_weather' + assert emitted_args == '{"location": "Beijing"}' + assert parser.validate_complete() is True + + def test_stream_chunk_reasoning_effort_high_starts_in_reasoning_mode(self): + parser = _make_response_parser_with_reasoning({'reasoning_effort': 'high'}) + + delta, tool_emitted = first_stream_delta(parser.stream_chunk(delta_text='first reason', delta_token_ids=[])) + + assert tool_emitted is False + assert delta is not None + assert delta.reasoning_content == 'first reason' + assert delta.content is None + + def test_stream_chunk_enable_thinking_false_starts_in_plain_mode(self): + parser = _make_response_parser_with_reasoning({'enable_thinking': False}) + + delta, tool_emitted = first_stream_delta(parser.stream_chunk(delta_text='plain answer', delta_token_ids=[])) + + assert tool_emitted is False + assert delta is not None + assert delta.content == 'plain answer' + assert delta.reasoning_content is None + def test_stream_chunk_keeps_string_without_schema(self, response_parser): chunks = [ '', @@ -177,6 +269,25 @@ def test_stream_chunk_keeps_string_without_schema(self, response_parser): class TestGlm47ToolParserComplete: """Complete-parse tests for glm47 tool payloads.""" + def test_parse_complete_tool_start_ends_reasoning_without_close_tag(self): + parser = _make_response_parser_with_reasoning() + text = ( + 'first reason' + 'get_weather' + 'locationBeijing' + '' + ) + + content, tool_calls, reasoning = parser.parse_complete(text) + + assert content is None + assert reasoning == 'first reason' + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].function.name == 'get_weather' + assert json.loads(tool_calls[0].function.arguments) == {'location': 'Beijing'} + assert parser.validate_complete(text) is True + def test_parse_tool_call_complete_with_arguments(self): parser = Glm47ToolParser() payload = (