From 3b5f03b207673d7ae376f5674b000a2c21c4ebf3 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 6 Jul 2026 14:40:57 +0800 Subject: [PATCH 01/33] feat: support glm 5.2 --- docker/prepare_3rdparty_wheel.sh | 5 + .../pytorch/configurations/deepseek_v32.py | 45 ++++++ lmdeploy/pytorch/models/deepseek_v32.py | 136 +++++++++++++++++- lmdeploy/pytorch/transformers/__init__.py | 3 + .../transformers/configuration_glm_moe_dsa.py | 36 +++++ .../pytorch/config/test_glm_moe_dsa_config.py | 71 +++++++++ 6 files changed, 293 insertions(+), 3 deletions(-) create mode 100644 lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py create mode 100644 tests/pytorch/config/test_glm_moe_dsa_config.py diff --git a/docker/prepare_3rdparty_wheel.sh b/docker/prepare_3rdparty_wheel.sh index a5a2eb2dcf..034eb8e858 100644 --- a/docker/prepare_3rdparty_wheel.sh +++ b/docker/prepare_3rdparty_wheel.sh @@ -16,6 +16,7 @@ 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 +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/pytorch/configurations/deepseek_v32.py b/lmdeploy/pytorch/configurations/deepseek_v32.py index cec2cf0781..08e91ac2cd 100644 --- a/lmdeploy/pytorch/configurations/deepseek_v32.py +++ b/lmdeploy/pytorch/configurations/deepseek_v32.py @@ -4,6 +4,50 @@ from .deepseek_v2 import DeepseekV2ModelConfigBuilder +def normalize_glm_moe_dsa_config(hf_config): + """Normalize GLM-MoE-DSA config fields for DeepSeek-V3.2 reuse.""" + if hf_config.model_type != 'glm_moe_dsa': + return + + qk_head_dim = getattr(hf_config, 'qk_head_dim', None) + if qk_head_dim is not None and qk_head_dim != hf_config.qk_nope_head_dim + hf_config.qk_rope_head_dim: + qk_rope_head_dim = qk_head_dim - hf_config.qk_nope_head_dim + if qk_rope_head_dim <= 0: + raise ValueError('Invalid GLM-MoE-DSA qk head dimensions.') + hf_config.qk_rope_head_dim = qk_rope_head_dim + + hf_config.qk_head_dim = hf_config.qk_nope_head_dim + hf_config.qk_rope_head_dim + hf_config.head_dim = hf_config.qk_rope_head_dim + + indexer_types = getattr(hf_config, 'indexer_types', None) + if indexer_types is not None: + indexer_types = list(indexer_types) + else: + pattern = getattr(hf_config, 'index_topk_pattern', None) + if pattern is not None: + if isinstance(pattern, str): + pattern_map = {'F': 'full', 'S': 'shared'} + indexer_types = [pattern_map[item.upper()] for item in pattern] + else: + indexer_types = list(pattern) + else: + freq = max(getattr(hf_config, 'index_topk_freq', 1), 1) + offset = getattr(hf_config, 'index_skip_topk_offset', 2) + indexer_types = [ + 'full' if max(layer_idx - offset + 1, 0) % freq == 0 else 'shared' + for layer_idx in range(hf_config.num_hidden_layers) + ] + + if len(indexer_types) != hf_config.num_hidden_layers: + raise ValueError('GLM-MoE-DSA indexer_types length must match num_hidden_layers.') + if any(item not in ('full', 'shared') for item in indexer_types): + raise ValueError('GLM-MoE-DSA indexer_types only supports "full" and "shared".') + if indexer_types[0] != 'full': + raise ValueError('The first GLM-MoE-DSA layer must be a full indexer layer.') + + hf_config.indexer_types = indexer_types + + def _check_env_v32(device: str = 'cuda'): """Environment check.""" if device != 'cuda': @@ -34,6 +78,7 @@ def condition(cls, hf_config): @classmethod def build(cls, hf_config, model_path: str | None = None, **kwargs): """build.""" + normalize_glm_moe_dsa_config(hf_config) config = DeepseekV2ModelConfigBuilder.build(hf_config, model_path=model_path, **kwargs) assert hf_config.use_flash_mla, 'DeepSeek-V3.2 requires flash_mla to be available.' diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index 6954e47fca..c15e5e3c93 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -33,6 +33,26 @@ ) +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 + + def rotate_activation(x: torch.Tensor) -> torch.Tensor: assert x.dtype == torch.bfloat16 from fast_hadamard_transform import hadamard_transform @@ -239,7 +259,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,6 +311,7 @@ def forward( rotary_pos_emb: tuple[torch.FloatTensor, torch.FloatTensor], past_key_value: Sequence[torch.Tensor] = None, attn_metadata: Any = None, + prev_topk_indices: torch.Tensor | None = None, ): """Rewrite of LlamaAttention.forward.""" dist_ctx = get_dist_manager().current_context() @@ -310,7 +334,16 @@ 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 self.indexer is None: + if prev_topk_indices is None: + raise RuntimeError(f'Layer {self.layer_idx} reuses DSA top-k indices but none were provided.') + topk_indices = prev_topk_indices + else: + topk_indices = self.indexer(hidden_states, + qr, + rotary_pos_emb, + past_key_value[-2:], + attn_metadata=attn_metadata) attn_output = self.attn_fwd( query_states, @@ -329,7 +362,7 @@ def forward( attn_output = attn_bmm_out.flatten(-2, -1)[None] attn_output = self.o_proj(attn_output) - return attn_output + return attn_output, topk_indices class DeepseekV32DecoderLayer(DeepseekV2DecoderLayer): @@ -365,6 +398,44 @@ 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, + prev_topk_indices: torch.Tensor | None = None, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.Tensor | None]: + + 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, topk_indices = self.self_attn( + hidden_states=hidden_states, + rotary_pos_emb=rotary_pos_emb, + past_key_value=past_key_value, + attn_metadata=attn_metadata, + prev_topk_indices=prev_topk_indices, + ) + 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, + ) + topk_indices = None + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + + return hidden_states, residual, topk_indices + class DeepseekV32Model(DeepseekV2Model): @@ -404,6 +475,56 @@ 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, + ): + """forward.""" + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + residual = None + prev_topk_indices = 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, prev_topk_indices = decoder_layer( + hidden_states, + rotary_pos_emb=rotary_pos_emb, + past_key_value=past_key_value, + residual=residual, + attn_metadata=attn_metadata, + prev_topk_indices=prev_topk_indices, + ) + + 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, + ): + """forward_microbatch.""" + 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, + ) + class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM): @@ -425,3 +546,12 @@ def __init__(self, dtype=dtype, device=device) self._load_buffers = dict() + + 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 '.self_attn.indexer.' in name and name not in params_dict: + layer_idx = get_layer_idx_from_weight_name(name) + 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/transformers/__init__.py b/lmdeploy/pytorch/transformers/__init__.py index 994d8ffd0d..8b3e8d0598 100644 --- a/lmdeploy/pytorch/transformers/__init__.py +++ b/lmdeploy/pytorch/transformers/__init__.py @@ -11,6 +11,9 @@ def register_config(model_type: str): if model_type == 'deepseek_v32': from lmdeploy.pytorch.transformers.configuration_deepseek_v32 import DeepseekV32Config AutoConfig.register(DeepseekV32Config.model_type, DeepseekV32Config) + elif model_type == 'glm_moe_dsa': + from lmdeploy.pytorch.transformers.configuration_glm_moe_dsa import GlmMoeDsaConfig + AutoConfig.register(GlmMoeDsaConfig.model_type, GlmMoeDsaConfig, exist_ok=True) else: logger.debug(f'Can not register config for model_type: {model_type}') diff --git a/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py b/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py new file mode 100644 index 0000000000..e318fd28dd --- /dev/null +++ b/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py @@ -0,0 +1,36 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +from lmdeploy.pytorch.transformers.configuration_deepseek_v32 import DeepseekV32Config + + +class GlmMoeDsaConfig(DeepseekV32Config): + model_type = 'glm_moe_dsa' + attribute_map = { + 'num_local_experts': 'n_routed_experts', + } + + def __init__( + self, + index_head_dim=128, + index_n_heads=64, + index_topk=2048, + index_topk_pattern=None, + index_topk_freq=1, + index_skip_topk_offset=2, + indexer_types=None, + **kwargs, + ): + super().__init__( + index_head_dim=index_head_dim, + index_n_heads=index_n_heads, + index_topk=index_topk, + **kwargs, + ) + self.index_topk_pattern = index_topk_pattern + self.index_topk_freq = index_topk_freq + self.index_skip_topk_offset = index_skip_topk_offset + self.indexer_types = indexer_types + + if hasattr(self, 'qk_nope_head_dim') and hasattr(self, 'qk_rope_head_dim'): + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.head_dim = self.qk_rope_head_dim diff --git a/tests/pytorch/config/test_glm_moe_dsa_config.py b/tests/pytorch/config/test_glm_moe_dsa_config.py new file mode 100644 index 0000000000..733ca4fadf --- /dev/null +++ b/tests/pytorch/config/test_glm_moe_dsa_config.py @@ -0,0 +1,71 @@ +from types import SimpleNamespace + +import pytest + +from lmdeploy.pytorch.configurations.deepseek_v32 import normalize_glm_moe_dsa_config + + +def _make_config(**kwargs): + values = dict( + model_type='glm_moe_dsa', + num_hidden_layers=4, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + ) + values.update(kwargs) + return SimpleNamespace(**values) + + +def test_glm_moe_dsa_normalizes_pattern_indexer_types(): + cfg = _make_config(index_topk_pattern='FSSF') + + normalize_glm_moe_dsa_config(cfg) + + assert cfg.qk_head_dim == 192 + assert cfg.head_dim == 64 + assert cfg.indexer_types == ['full', 'shared', 'shared', 'full'] + + +def test_glm_moe_dsa_normalizes_freq_offset_indexer_types(): + cfg = _make_config(num_hidden_layers=6, index_topk_freq=2, index_skip_topk_offset=2) + + normalize_glm_moe_dsa_config(cfg) + + assert cfg.indexer_types == ['full', 'full', 'shared', 'full', 'shared', 'full'] + + +def test_glm_moe_dsa_recovers_corrupted_rope_head_dim(): + cfg = _make_config(qk_nope_head_dim=192, qk_rope_head_dim=192, qk_head_dim=256) + + normalize_glm_moe_dsa_config(cfg) + + assert cfg.qk_nope_head_dim == 192 + assert cfg.qk_rope_head_dim == 64 + assert cfg.qk_head_dim == 256 + assert cfg.head_dim == 64 + + +def test_glm_moe_dsa_rejects_shared_first_layer(): + cfg = _make_config(index_topk_pattern='SFFF') + + with pytest.raises(ValueError, match='first GLM-MoE-DSA layer'): + normalize_glm_moe_dsa_config(cfg) + + +def test_glm_moe_dsa_fallback_config_keeps_rope_head_dim(): + from lmdeploy.pytorch.transformers.configuration_glm_moe_dsa import GlmMoeDsaConfig + + cfg = GlmMoeDsaConfig(qk_nope_head_dim=128, + qk_rope_head_dim=64, + hidden_size=4096, + intermediate_size=8192, + num_attention_heads=32, + num_hidden_layers=2, + vocab_size=32000, + bos_token_id=1, + eos_token_id=2) + + assert cfg.qk_nope_head_dim == 128 + assert cfg.qk_rope_head_dim == 64 + assert cfg.qk_head_dim == 192 + assert cfg.head_dim == 64 From 1d95a184a9fd7e1a03e9052f4f5a18bcd2f2e7e1 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 9 Jul 2026 16:33:10 +0800 Subject: [PATCH 02/33] feat: support glm-5.2 dsa indexing --- lmdeploy/pytorch/models/deepseek_v32.py | 46 +++++- .../transformers/configuration_glm_moe_dsa.py | 6 +- .../pytorch/config/test_glm_moe_dsa_config.py | 37 ++++- tests/pytorch/models/test_deepseek_v32.py | 151 ++++++++++++++++++ 4 files changed, 230 insertions(+), 10 deletions(-) create mode 100644 tests/pytorch/models/test_deepseek_v32.py diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index c15e5e3c93..9c38570098 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -60,6 +60,27 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: return hadamard_transform(x, scale=hidden_size**-0.5) +def rotate_interleaved(x: torch.Tensor) -> torch.Tensor: + """Rotate interleaved RoPE pairs.""" + out = torch.empty_like(x) + out[..., ::2] = -x[..., 1::2] + out[..., 1::2] = x[..., ::2] + return out + + +def apply_interleaved_rotary_pos_emb(query: torch.Tensor, key: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + """Apply GPT-J style interleaved RoPE to query and key.""" + + def _to_interleaved(freqs: torch.Tensor): + return freqs[..., :freqs.size(-1) // 2].repeat_interleave(2, dim=-1) + + cos = _to_interleaved(cos).unsqueeze(-2) + sin = _to_interleaved(sin).unsqueeze(-2) + query = query * cos + rotate_interleaved(query) * sin + key = key * cos + rotate_interleaved(key) * sin + return query, key + + class LayerNorm(nn.Module): """Layer Normalization.""" @@ -92,6 +113,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, @@ -119,6 +141,21 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic 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 + k_pe = k_pe[..., None, :] + if self.rope_interleave: + return apply_interleaved_rotary_pos_emb(q_pe, k_pe, cos, sin) + return self.apply_rotary_pos_emb( + q_pe, + k_pe, + cos, + sin, + inplace=False, + ) + def forward(self, x: torch.Tensor, qr: torch.Tensor, @@ -133,14 +170,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) diff --git a/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py b/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py index e318fd28dd..798eb99ef8 100644 --- a/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py +++ b/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py @@ -12,11 +12,13 @@ class GlmMoeDsaConfig(DeepseekV32Config): def __init__( self, index_head_dim=128, - index_n_heads=64, + index_n_heads=32, index_topk=2048, index_topk_pattern=None, index_topk_freq=1, index_skip_topk_offset=2, + rope_interleave=True, + indexer_rope_interleave=True, indexer_types=None, **kwargs, ): @@ -29,6 +31,8 @@ def __init__( self.index_topk_pattern = index_topk_pattern self.index_topk_freq = index_topk_freq self.index_skip_topk_offset = index_skip_topk_offset + self.rope_interleave = rope_interleave + self.indexer_rope_interleave = indexer_rope_interleave self.indexer_types = indexer_types if hasattr(self, 'qk_nope_head_dim') and hasattr(self, 'qk_rope_head_dim'): diff --git a/tests/pytorch/config/test_glm_moe_dsa_config.py b/tests/pytorch/config/test_glm_moe_dsa_config.py index 733ca4fadf..aee7cd0b93 100644 --- a/tests/pytorch/config/test_glm_moe_dsa_config.py +++ b/tests/pytorch/config/test_glm_moe_dsa_config.py @@ -2,7 +2,10 @@ import pytest -from lmdeploy.pytorch.configurations.deepseek_v32 import normalize_glm_moe_dsa_config +from lmdeploy.pytorch.configurations.deepseek_v32 import ( + DeepseekV32ModelConfigBuilder, + normalize_glm_moe_dsa_config, +) def _make_config(**kwargs): @@ -69,3 +72,35 @@ def test_glm_moe_dsa_fallback_config_keeps_rope_head_dim(): assert cfg.qk_rope_head_dim == 64 assert cfg.qk_head_dim == 192 assert cfg.head_dim == 64 + assert cfg.index_n_heads == 32 + assert cfg.rope_interleave is True + assert cfg.indexer_rope_interleave is True + + +def test_glm_moe_dsa_builder_normalizes_mla_kv_heads(monkeypatch): + from lmdeploy.pytorch.configurations import deepseek_v2 + + monkeypatch.setattr(deepseek_v2, 'flash_mla_available', lambda: True) + cfg = _make_config(num_hidden_layers=2, + index_head_dim=128, + index_n_heads=32, + index_topk=2048, + index_skip_topk_offset=3, + index_topk_freq=4, + indexer_rope_interleave=True, + rope_interleave=True, + hidden_size=6144, + kv_lora_rank=512, + num_attention_heads=64, + num_key_value_heads=64, + bos_token_id=None, + eos_token_id=[154820, 154827, 154829], + vocab_size=154880, + architectures=['GlmMoeDsaForCausalLM']) + + model_config = DeepseekV32ModelConfigBuilder.build(cfg, tp=1) + + assert cfg.num_key_value_heads == 1 + assert model_config.num_key_value_heads == 1 + assert model_config.head_dim == 576 + assert model_config.use_mla_fp8_cache is True diff --git a/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py new file mode 100644 index 0000000000..75a38d741f --- /dev/null +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -0,0 +1,151 @@ +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from lmdeploy.pytorch.models.deepseek_v32 import (DeepseekV32Attention, Indexer, apply_interleaved_rotary_pos_emb, + get_layer_indexer_type, rotate_interleaved) + + +def test_interleaved_rotary_uses_pairwise_layout_with_half_split_freqs(): + query = torch.tensor([[[[1.0, 2.0, 3.0, 4.0]]]]) + key = torch.tensor([[[[5.0, 6.0, 7.0, 8.0]]]]) + cos = torch.tensor([[10.0, 20.0, 10.0, 20.0]]) + sin = torch.tensor([[1.0, 2.0, 1.0, 2.0]]) + + query_out, key_out = apply_interleaved_rotary_pos_emb(query, key, cos, sin) + + interleaved_cos = torch.tensor([[[10.0, 10.0, 20.0, 20.0]]]) + interleaved_sin = torch.tensor([[[1.0, 1.0, 2.0, 2.0]]]) + expected_query = query * interleaved_cos + rotate_interleaved(query) * interleaved_sin + expected_key = key * interleaved_cos + rotate_interleaved(key) * interleaved_sin + assert torch.equal(query_out, expected_query) + assert torch.equal(key_out, expected_key) + + +def test_indexer_rope_interleave_uses_interleaved_layout(): + indexer = Indexer.__new__(Indexer) + indexer.rope_interleave = True + + query = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]]) + key = torch.tensor([[5.0, 6.0, 7.0, 8.0]]) + cos = torch.tensor([[10.0, 20.0, 10.0, 20.0]]) + sin = torch.tensor([[1.0, 2.0, 1.0, 2.0]]) + + query_out, key_out = indexer._apply_rotary_pos_emb(query, key, (cos, sin)) + + key = key[..., None, :] + interleaved_cos = torch.tensor([[10.0, 10.0, 20.0, 20.0]]).unsqueeze(-2) + interleaved_sin = torch.tensor([[1.0, 1.0, 2.0, 2.0]]).unsqueeze(-2) + expected_query = query * interleaved_cos + rotate_interleaved(query) * interleaved_sin + expected_key = key * interleaved_cos + rotate_interleaved(key) * interleaved_sin + assert torch.equal(query_out, expected_query) + assert torch.equal(key_out, expected_key) + + +def test_indexer_rope_interleave_can_fall_back_to_half_split_layout(): + indexer = Indexer.__new__(Indexer) + indexer.rope_interleave = False + indexer.apply_rotary_pos_emb = lambda q, k, cos, sin, inplace=False: ('half-split', q, k, cos, sin, inplace) + + query = torch.zeros(1, 1, 4) + key = torch.zeros(1, 4) + cos = torch.zeros(1, 4) + sin = torch.zeros(1, 4) + + out = indexer._apply_rotary_pos_emb(query, key, (cos, sin)) + + assert out[0] == 'half-split' + assert out[2].shape == (1, 1, 4) + assert out[-1] is False + + +def test_shared_indexer_layer_reuses_previous_topk_indices(monkeypatch): + attn = DeepseekV32Attention.__new__(DeepseekV32Attention) + nn.Module.__init__(attn) + attn.layer_idx = 3 + attn.indexer = None + prev_topk_indices = torch.tensor([[3, 1, 0]], dtype=torch.int32) + seen = {} + + def fake_qkv_proj(hidden_states, num_heads): + q_len = hidden_states.size(1) + return ( + torch.zeros(q_len, 1, 2), + torch.zeros(q_len, 1, 2), + torch.zeros(q_len, 1, 1), + torch.zeros(q_len, 1, 1), + torch.zeros(q_len, 1, 1), + torch.zeros(q_len, 1, 1), + ) + + def fake_rotary(q, k, cos, sin, inplace=False): + return q, k + + def fake_attn_fwd(*args, **kwargs): + seen['nsa_indices'] = kwargs['nsa_indices'] + return torch.zeros(1, 1, 1) + + monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v32.get_dist_manager', + lambda: SimpleNamespace(current_context=lambda: SimpleNamespace( + dist_config=SimpleNamespace(attn_tp=1)))) + attn.num_heads = 1 + attn.kv_lora_rank = 1 + attn.qk_rope_head_dim = 1 + attn._qkv_proj = fake_qkv_proj + attn.apply_rotary_pos_emb = fake_rotary + attn.attn_fwd = fake_attn_fwd + attn.vc = lambda attn_output, out: out.copy_(attn_output) + attn.v_head_dim = 1 + attn.o_proj = nn.Identity() + + output, returned_topk = attn( + hidden_states=torch.zeros(1, 1, 2), + rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), + past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], + prev_topk_indices=prev_topk_indices, + ) + + assert output.shape == (1, 1, 1) + assert returned_topk is prev_topk_indices + assert seen['nsa_indices'] is prev_topk_indices + + +def test_shared_indexer_layer_requires_previous_topk_indices(monkeypatch): + attn = DeepseekV32Attention.__new__(DeepseekV32Attention) + nn.Module.__init__(attn) + attn.layer_idx = 3 + attn.indexer = None + attn.num_heads = 1 + attn.kv_lora_rank = 1 + attn.qk_rope_head_dim = 1 + monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v32.get_dist_manager', + lambda: SimpleNamespace(current_context=lambda: SimpleNamespace( + dist_config=SimpleNamespace(attn_tp=1)))) + attn._qkv_proj = lambda *args, **kwargs: ( + torch.zeros(1, 1, 2), + torch.zeros(1, 1, 2), + torch.zeros(1, 1, 1), + torch.zeros(1, 1, 1), + torch.zeros(1, 1, 1), + torch.zeros(1, 1, 1), + ) + attn.apply_rotary_pos_emb = lambda q, k, cos, sin, inplace=False: (q, k) + + with pytest.raises(RuntimeError, match='reuses DSA top-k indices'): + attn( + hidden_states=torch.zeros(1, 1, 2), + rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), + past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], + prev_topk_indices=None, + ) + + +def test_layer_indexer_type_defaults_to_full_and_reads_shared_entries(): + config = SimpleNamespace(indexer_types=['full', 'shared']) + + assert get_layer_indexer_type(config, 0) == 'full' + assert get_layer_indexer_type(config, 1) == 'shared' + assert get_layer_indexer_type(config, 2) == 'full' + assert get_layer_indexer_type(SimpleNamespace(indexer_types=None), 1) == 'full' From cb8bb28fd51922b8b2646b09cba8d78754447379 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 9 Jul 2026 17:17:22 +0800 Subject: [PATCH 03/33] fix: parse glm tool calls from reasoning --- lmdeploy/serve/parsers/response_parser.py | 55 +++++++-- .../serve/parsers/test_glm47_parser.py | 111 ++++++++++++++++++ 2 files changed, 154 insertions(+), 12 deletions(-) diff --git a/lmdeploy/serve/parsers/response_parser.py b/lmdeploy/serve/parsers/response_parser.py index aee133f5fe..6dacd2269b 100644 --- a/lmdeploy/serve/parsers/response_parser.py +++ b/lmdeploy/serve/parsers/response_parser.py @@ -488,8 +488,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 @@ -507,19 +507,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]: @@ -630,21 +651,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 @@ -697,7 +724,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/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 = ( From 5252430b7d82094a28db259bb28db575937660ff Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 9 Jul 2026 21:08:41 +0800 Subject: [PATCH 04/33] fix: cache dsa topk indices --- lmdeploy/pytorch/models/deepseek_v32.py | 87 +++++++++++++--- tests/pytorch/models/test_deepseek_v32.py | 118 +++++++++++++++------- 2 files changed, 155 insertions(+), 50 deletions(-) diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index 9c38570098..e5e2c6f06d 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -7,7 +7,7 @@ from torch import nn 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 StepContextManager, get_step_ctx_manager from lmdeploy.pytorch.nn import ( ApplyRotaryEmb, Attention, @@ -53,6 +53,63 @@ def get_layer_idx_from_weight_name(name: str) -> int | 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 + + build_ctx = getattr(ctx_mgr, 'build_ctx', None) + if build_ctx is not None and build_ctx.max_batch_size > 0: + capacity = max(capacity, build_ctx.max_batch_size * (1 + build_ctx.num_spec_tokens)) + + 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) + if topk_indices.dtype != torch.int32: + topk_indices = topk_indices.to(torch.int32) + 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): + """Move selected top-k rows to the front for the next decode-style + reuse step.""" + if self.indices is None: + raise RuntimeError('DSA top-k indices are compacted before the shared buffer is populated.') + row_indices = row_indices.to(device=self.indices.device, dtype=torch.long) + num_rows = row_indices.numel() + self.indices[:num_rows].copy_(self.indices.index_select(0, row_indices)) + + def rotate_activation(x: torch.Tensor) -> torch.Tensor: assert x.dtype == torch.bfloat16 from fast_hadamard_transform import hadamard_transform @@ -341,7 +398,7 @@ def forward( rotary_pos_emb: tuple[torch.FloatTensor, torch.FloatTensor], past_key_value: Sequence[torch.Tensor] = None, attn_metadata: Any = None, - prev_topk_indices: torch.Tensor | None = None, + topk_indices_buffer: DSATopKIndicesBuffer | None = None, ): """Rewrite of LlamaAttention.forward.""" dist_ctx = get_dist_manager().current_context() @@ -365,15 +422,17 @@ def forward( key_states[..., nope_size:] = k_pe if self.indexer is None: - if prev_topk_indices is None: + if topk_indices_buffer is None: raise RuntimeError(f'Layer {self.layer_idx} reuses DSA top-k indices but none were provided.') - topk_indices = prev_topk_indices + topk_indices = topk_indices_buffer.read(q_len, hidden_states.device) else: topk_indices = self.indexer(hidden_states, qr, rotary_pos_emb, past_key_value[-2:], attn_metadata=attn_metadata) + if topk_indices_buffer is not None: + topk_indices = topk_indices_buffer.write(topk_indices) attn_output = self.attn_fwd( query_states, @@ -392,7 +451,7 @@ def forward( attn_output = attn_bmm_out.flatten(-2, -1)[None] attn_output = self.o_proj(attn_output) - return attn_output, topk_indices + return attn_output class DeepseekV32DecoderLayer(DeepseekV2DecoderLayer): @@ -435,8 +494,8 @@ def forward( past_key_value: list[torch.FloatTensor] | None, residual: torch.Tensor | None = None, attn_metadata: Any = None, - prev_topk_indices: torch.Tensor | None = None, - ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.Tensor | None]: + topk_indices_buffer: DSATopKIndicesBuffer | None = None, + ) -> tuple[torch.FloatTensor, torch.FloatTensor]: if residual is None: residual = hidden_states @@ -445,12 +504,12 @@ def forward( hidden_states, residual = self.input_layernorm(hidden_states, residual) if isinstance(self.self_attn, DeepseekV32Attention): - hidden_states, topk_indices = self.self_attn( + 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, - prev_topk_indices=prev_topk_indices, + topk_indices_buffer=topk_indices_buffer, ) else: hidden_states = self.self_attn( @@ -459,12 +518,11 @@ def forward( past_key_value=past_key_value, attn_metadata=attn_metadata, ) - topk_indices = None hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states) - return hidden_states, residual, topk_indices + return hidden_states, residual class DeepseekV32Model(DeepseekV2Model): @@ -486,6 +544,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, @@ -519,19 +578,18 @@ def forward( hidden_states = inputs_embeds residual = None - prev_topk_indices = 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, prev_topk_indices = decoder_layer( + 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, - prev_topk_indices=prev_topk_indices, + topk_indices_buffer=self.topk_indices_buffer, ) hidden_states, _ = self.norm(hidden_states, residual) @@ -582,6 +640,7 @@ def _load_weight_attention(self, name: str, loaded_weight: torch.Tensor, params_ """Load attention weights.""" 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 previous 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/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index 75a38d741f..ea6fdd95dc 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -4,8 +4,41 @@ import torch from torch import nn -from lmdeploy.pytorch.models.deepseek_v32 import (DeepseekV32Attention, Indexer, apply_interleaved_rotary_pos_emb, - get_layer_indexer_type, rotate_interleaved) +from lmdeploy.pytorch.models.deepseek_v32 import ( + DeepseekV32Attention, + DSATopKIndicesBuffer, + Indexer, + apply_interleaved_rotary_pos_emb, + get_layer_indexer_type, + rotate_interleaved, +) + + +def _patch_minimal_attention(attn, monkeypatch): + """Patch a DSA attention module down to the top-k routing contract.""" + + def fake_qkv_proj(hidden_states, num_heads): + q_len = hidden_states.size(1) + return ( + torch.zeros(q_len, 1, 2), + torch.zeros(q_len, 1, 2), + torch.zeros(q_len, 1, 1), + torch.zeros(q_len, 1, 1), + torch.zeros(q_len, 1, 1), + torch.zeros(q_len, 1, 1), + ) + + monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v32.get_dist_manager', + lambda: SimpleNamespace(current_context=lambda: SimpleNamespace( + dist_config=SimpleNamespace(attn_tp=1)))) + attn.num_heads = 1 + attn.kv_lora_rank = 1 + attn.qk_rope_head_dim = 1 + attn._qkv_proj = fake_qkv_proj + attn.apply_rotary_pos_emb = lambda q, k, cos, sin, inplace=False: (q, k) + attn.vc = lambda attn_output, out: out.copy_(attn_output) + attn.v_head_dim = 1 + attn.o_proj = nn.Identity() def test_interleaved_rotary_uses_pairwise_layout_with_half_split_freqs(): @@ -61,58 +94,72 @@ def test_indexer_rope_interleave_can_fall_back_to_half_split_layout(): assert out[-1] is False -def test_shared_indexer_layer_reuses_previous_topk_indices(monkeypatch): +def test_full_indexer_layer_writes_shared_topk_buffer(monkeypatch): attn = DeepseekV32Attention.__new__(DeepseekV32Attention) nn.Module.__init__(attn) - attn.layer_idx = 3 - attn.indexer = None - prev_topk_indices = torch.tensor([[3, 1, 0]], dtype=torch.int32) + attn.layer_idx = 0 + _patch_minimal_attention(attn, monkeypatch) + computed_topk = torch.tensor([[4, 2, 1], [3, 2, 0]], dtype=torch.int32) + attn.indexer = lambda *args, **kwargs: computed_topk seen = {} - def fake_qkv_proj(hidden_states, num_heads): - q_len = hidden_states.size(1) - return ( - torch.zeros(q_len, 1, 2), - torch.zeros(q_len, 1, 2), - torch.zeros(q_len, 1, 1), - torch.zeros(q_len, 1, 1), - torch.zeros(q_len, 1, 1), - torch.zeros(q_len, 1, 1), - ) + def fake_attn_fwd(*args, **kwargs): + seen['nsa_indices'] = kwargs['nsa_indices'] + return torch.zeros(args[0].size(0), 1, 1) + + attn.attn_fwd = fake_attn_fwd + topk_buffer = DSATopKIndicesBuffer(topk=3) + + output = attn( + hidden_states=torch.zeros(1, 2, 2), + rotary_pos_emb=(torch.zeros(2, 1), torch.zeros(2, 1)), + past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], + topk_indices_buffer=topk_buffer, + ) + + assert output.shape == (1, 2, 1) + assert torch.equal(topk_buffer.indices[:2], computed_topk) + assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:2].data_ptr() + + +def test_dsa_topk_buffer_compacts_selected_rows(): + topk_buffer = DSATopKIndicesBuffer(topk=3) + rows = torch.tensor([[0, 1, 2], [10, 11, 12], [20, 21, 22], [30, 31, 32]], dtype=torch.int32) + topk_buffer.write(rows) + + topk_buffer.compact(torch.tensor([2, 0])) + + assert torch.equal(topk_buffer.indices[:2], rows[[2, 0]]) - def fake_rotary(q, k, cos, sin, inplace=False): - return q, k + +def test_shared_indexer_layer_reuses_shared_topk_buffer(monkeypatch): + attn = DeepseekV32Attention.__new__(DeepseekV32Attention) + nn.Module.__init__(attn) + attn.layer_idx = 3 + attn.indexer = None + _patch_minimal_attention(attn, monkeypatch) + seen = {} def fake_attn_fwd(*args, **kwargs): seen['nsa_indices'] = kwargs['nsa_indices'] - return torch.zeros(1, 1, 1) + return torch.zeros(args[0].size(0), 1, 1) - monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v32.get_dist_manager', - lambda: SimpleNamespace(current_context=lambda: SimpleNamespace( - dist_config=SimpleNamespace(attn_tp=1)))) - attn.num_heads = 1 - attn.kv_lora_rank = 1 - attn.qk_rope_head_dim = 1 - attn._qkv_proj = fake_qkv_proj - attn.apply_rotary_pos_emb = fake_rotary attn.attn_fwd = fake_attn_fwd - attn.vc = lambda attn_output, out: out.copy_(attn_output) - attn.v_head_dim = 1 - attn.o_proj = nn.Identity() + topk_buffer = DSATopKIndicesBuffer(topk=3) + topk_buffer.write(torch.tensor([[5, 4, 3]], dtype=torch.int32)) - output, returned_topk = attn( + output = attn( hidden_states=torch.zeros(1, 1, 2), rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], - prev_topk_indices=prev_topk_indices, + topk_indices_buffer=topk_buffer, ) assert output.shape == (1, 1, 1) - assert returned_topk is prev_topk_indices - assert seen['nsa_indices'] is prev_topk_indices + assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:1].data_ptr() -def test_shared_indexer_layer_requires_previous_topk_indices(monkeypatch): +def test_shared_indexer_layer_requires_shared_topk_buffer(monkeypatch): attn = DeepseekV32Attention.__new__(DeepseekV32Attention) nn.Module.__init__(attn) attn.layer_idx = 3 @@ -138,7 +185,6 @@ def test_shared_indexer_layer_requires_previous_topk_indices(monkeypatch): hidden_states=torch.zeros(1, 1, 2), rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], - prev_topk_indices=None, ) From 34790b58e5371bcc574ce51a0a30183c834a38d1 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 9 Jul 2026 21:14:41 +0800 Subject: [PATCH 05/33] refactor: simplify dsa topk buffer --- lmdeploy/pytorch/models/deepseek_v32.py | 18 +----------------- tests/pytorch/models/test_deepseek_v32.py | 10 ---------- 2 files changed, 1 insertion(+), 27 deletions(-) diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index e5e2c6f06d..d01481f9e1 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -67,10 +67,6 @@ def _target_capacity(self, num_tokens: int) -> int: if ctx_mgr is None: return capacity - build_ctx = getattr(ctx_mgr, 'build_ctx', None) - if build_ctx is not None and build_ctx.max_batch_size > 0: - capacity = max(capacity, build_ctx.max_batch_size * (1 + build_ctx.num_spec_tokens)) - context = ctx_mgr.current_context() cache_config = getattr(context, 'cache_config', None) max_prefill_token_num = getattr(cache_config, 'max_prefill_token_num', None) @@ -89,8 +85,6 @@ def ensure(self, num_tokens: int, device: torch.device) -> torch.Tensor: 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) - if topk_indices.dtype != torch.int32: - topk_indices = topk_indices.to(torch.int32) buffer.copy_(topk_indices) return buffer @@ -100,16 +94,6 @@ def read(self, num_tokens: int, device: torch.device) -> torch.Tensor: 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): - """Move selected top-k rows to the front for the next decode-style - reuse step.""" - if self.indices is None: - raise RuntimeError('DSA top-k indices are compacted before the shared buffer is populated.') - row_indices = row_indices.to(device=self.indices.device, dtype=torch.long) - num_rows = row_indices.numel() - self.indices[:num_rows].copy_(self.indices.index_select(0, row_indices)) - - def rotate_activation(x: torch.Tensor) -> torch.Tensor: assert x.dtype == torch.bfloat16 from fast_hadamard_transform import hadamard_transform @@ -640,7 +624,7 @@ def _load_weight_attention(self, name: str, loaded_weight: torch.Tensor, params_ """Load attention weights.""" 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 previous top-k indices and have no local indexer. + # 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/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index ea6fdd95dc..27bd6aba56 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -122,16 +122,6 @@ def fake_attn_fwd(*args, **kwargs): assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:2].data_ptr() -def test_dsa_topk_buffer_compacts_selected_rows(): - topk_buffer = DSATopKIndicesBuffer(topk=3) - rows = torch.tensor([[0, 1, 2], [10, 11, 12], [20, 21, 22], [30, 31, 32]], dtype=torch.int32) - topk_buffer.write(rows) - - topk_buffer.compact(torch.tensor([2, 0])) - - assert torch.equal(topk_buffer.indices[:2], rows[[2, 0]]) - - def test_shared_indexer_layer_reuses_shared_topk_buffer(monkeypatch): attn = DeepseekV32Attention.__new__(DeepseekV32Attention) nn.Module.__init__(attn) From 38e3be8cc9d76199ae9881bab04ce37c0c97ca79 Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 10 Jul 2026 14:20:38 +0800 Subject: [PATCH 06/33] fix: support mla dp attention layout --- lmdeploy/pytorch/models/deepseek_v2.py | 11 ++++- lmdeploy/pytorch/models/deepseek_v32.py | 10 ++-- tests/pytorch/models/test_deepseek_v32.py | 58 ++++++++++++++++++++--- 3 files changed, 67 insertions(+), 12 deletions(-) diff --git a/lmdeploy/pytorch/models/deepseek_v2.py b/lmdeploy/pytorch/models/deepseek_v2.py index 087d744bf9..2eec92b838 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): diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index d01481f9e1..b110da70c6 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -385,9 +385,11 @@ def forward( topk_indices_buffer: DSATopKIndicesBuffer | None = None, ): """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) @@ -589,6 +591,8 @@ def forward_microbatch( inputs_embeds: torch.FloatTensor | None = None, ): """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, diff --git a/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index 27bd6aba56..418b444085 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -4,6 +4,7 @@ import torch from torch import nn +from lmdeploy.pytorch.models.deepseek_v2 import DeepseekV2BMM from lmdeploy.pytorch.models.deepseek_v32 import ( DeepseekV32Attention, DSATopKIndicesBuffer, @@ -14,23 +15,26 @@ ) -def _patch_minimal_attention(attn, monkeypatch): +def _patch_minimal_attention(attn, monkeypatch, dp=1, attn_tp=1, seen=None): """Patch a DSA attention module down to the top-k routing contract.""" def fake_qkv_proj(hidden_states, num_heads): + if seen is not None: + seen['num_heads'] = num_heads q_len = hidden_states.size(1) return ( + torch.zeros(q_len, num_heads, 2), torch.zeros(q_len, 1, 2), - torch.zeros(q_len, 1, 2), - torch.zeros(q_len, 1, 1), torch.zeros(q_len, 1, 1), + torch.zeros(q_len, num_heads, 1), torch.zeros(q_len, 1, 1), torch.zeros(q_len, 1, 1), ) + dist_config = SimpleNamespace(dp=dp, attn_tp=attn_tp) monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v32.get_dist_manager', - lambda: SimpleNamespace(current_context=lambda: SimpleNamespace( - dist_config=SimpleNamespace(attn_tp=1)))) + lambda: SimpleNamespace(current_config=lambda: dist_config, + current_context=lambda: SimpleNamespace(dist_config=dist_config))) attn.num_heads = 1 attn.kv_lora_rank = 1 attn.qk_rope_head_dim = 1 @@ -122,6 +126,45 @@ def fake_attn_fwd(*args, **kwargs): assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:2].data_ptr() +@pytest.mark.parametrize('dp,attn_tp,expected_num_heads', [(1, 2, 2), (2, 2, 4)]) +def test_attention_uses_dp_aware_num_heads(monkeypatch, dp, attn_tp, expected_num_heads): + attn = DeepseekV32Attention.__new__(DeepseekV32Attention) + nn.Module.__init__(attn) + attn.layer_idx = 0 + seen = {} + _patch_minimal_attention(attn, monkeypatch, dp=dp, attn_tp=attn_tp, seen=seen) + attn.num_heads = 4 + attn.indexer = lambda *args, **kwargs: torch.tensor([[4, 2, 1]], dtype=torch.int32) + attn.attn_fwd = lambda *args, **kwargs: torch.zeros(args[0].size(0), args[0].size(1), 1) + + attn( + hidden_states=torch.zeros(1, 1, 2), + rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), + past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], + topk_indices_buffer=DSATopKIndicesBuffer(topk=3), + ) + + assert seen['num_heads'] == expected_num_heads + + +@pytest.mark.parametrize('dp,attn_tp,expected_num_heads', [(1, 2, 2), (2, 2, 4)]) +def test_bmm_uses_dp_aware_weight_layout(monkeypatch, dp, attn_tp, expected_num_heads): + dist_config = SimpleNamespace(dp=dp, attn_tp=attn_tp) + monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v2.get_dist_manager', + lambda: SimpleNamespace(current_config=lambda: dist_config, + current_context=lambda: SimpleNamespace(dist_config=dist_config))) + monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v2.get_tp_world_rank', + lambda layer_type=None: (attn_tp, 0)) + + bmm = DeepseekV2BMM(batch=4, in_features=2, out_features=3, dtype=torch.float32, device='cpu') + weight = torch.arange(24, dtype=torch.float32).view(4, 2, 3) + bmm.weight.weight_loader(bmm.weight, weight) + + expected = weight if dp > 1 else weight.chunk(attn_tp, 0)[0] + assert bmm.weight.shape == (expected_num_heads, 2, 3) + torch.testing.assert_close(bmm.weight, expected) + + def test_shared_indexer_layer_reuses_shared_topk_buffer(monkeypatch): attn = DeepseekV32Attention.__new__(DeepseekV32Attention) nn.Module.__init__(attn) @@ -157,9 +200,10 @@ def test_shared_indexer_layer_requires_shared_topk_buffer(monkeypatch): attn.num_heads = 1 attn.kv_lora_rank = 1 attn.qk_rope_head_dim = 1 + dist_config = SimpleNamespace(dp=1, attn_tp=1) monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v32.get_dist_manager', - lambda: SimpleNamespace(current_context=lambda: SimpleNamespace( - dist_config=SimpleNamespace(attn_tp=1)))) + lambda: SimpleNamespace(current_config=lambda: dist_config, + current_context=lambda: SimpleNamespace(dist_config=dist_config))) attn._qkv_proj = lambda *args, **kwargs: ( torch.zeros(1, 1, 2), torch.zeros(1, 1, 2), From d92259585b5d8a7177e44845d7651939e3452390 Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 10 Jul 2026 15:34:48 +0800 Subject: [PATCH 07/33] fix: install xgrammar for jetson docker --- docker/Dockerfile.jetson | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docker/Dockerfile.jetson b/docker/Dockerfile.jetson index a0d5de81cd..680c231b89 100644 --- a/docker/Dockerfile.jetson +++ b/docker/Dockerfile.jetson @@ -27,9 +27,12 @@ RUN --mount=type=cache,target=/root/.cache \ COPY . /opt/lmdeploy WORKDIR /opt/lmdeploy +# The Jetson AI Lab index may lag xgrammar releases, while CUDA/Torch wheels +# should still come from the Jetson index. Preinstall only xgrammar from PyPI. RUN --mount=type=cache,target=/root/.cache \ --mount=type=cache,target=/opt/pytorch \ pip install build change-wheel-version && \ python -m build -w -o /wheels -v . && \ change_wheel_version --local-version cu126 --delete-old-wheel /wheels/lmdeploy*.whl && \ + pip install -v --no-deps xgrammar==0.1.33 --index-url https://pypi.org/simple/ && \ pip install -v /wheels/lmdeploy*.whl --index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/ From 95ec45ae0fc0a8a41646c5915ebfa79825ac1fa9 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 13 Jul 2026 18:59:35 +0800 Subject: [PATCH 08/33] perf: use radix topk for sparse indexing --- lmdeploy/pytorch/backends/cuda/nsa.py | 22 ++ .../pytorch/kernels/cuda/sparse_index_topk.py | 218 ++++++++++++++++++ .../pytorch/kernel/test_sparse_index_topk.py | 162 +++++++++++++ 3 files changed, 402 insertions(+) create mode 100644 lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py create mode 100644 tests/pytorch/kernel/test_sparse_index_topk.py diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index bab5d73094..1ab8211f85 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -1,4 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. +import functools + from torch import Tensor from lmdeploy.pytorch.kernels.cuda.bitonic_topk import bitonic_topk @@ -9,6 +11,17 @@ from ..nsa import BaseNSAIndexFP8, BaseNSAIndexFP8Builder, NSAIndexMeta +@functools.lru_cache +def _get_sparse_index_topk(topk: int): + from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import ( + is_sparse_index_topk_supported, + sparse_index_topk, + ) + if is_sparse_index_topk_supported(topk): + return sparse_index_topk + return None + + class TritonNSAIndexFP8(BaseNSAIndexFP8): def __init__(self, topk: int, softmax_scale: float, block_size: int, fill: int) -> None: @@ -62,6 +75,15 @@ def forward(self, q: Tensor, k: Tensor, weights: Tensor, k_cache: Tensor, k_s_ca max_q_seqlen=max_q_seqlen, max_k_seqlen=max_kv_seqlen, causal=True) + 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) diff --git a/lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py b/lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py new file mode 100644 index 0000000000..d29d7c00da --- /dev/null +++ b/lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py @@ -0,0 +1,218 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Shared sparse-index top-k kernels for MLA/NSA indexers. + +The CUDA graph constraint is the important bit here: score width is a padded +graph bucket, but the real per-row length is carried by ``kv_seqlens``. This +module therefore specializes only on model-config ``K`` and keeps the row +length as device data. +""" + +from __future__ import annotations + +import tilelang +import tilelang.language as T +import torch + +tilelang.set_log_level('WARNING') + +_SUPPORTED_TOPK = (512, 2048) +_FILL = -1 +_THREADS = 1024 +_RADIX_BITS = 8 +_RADIX_SIZE = 1 << _RADIX_BITS +_STATE_SELECTED_BIN = 0 +_STATE_COUNT_ABOVE = 1 +_STATE_EMIT_GT_COUNT = 0 +_STATE_EMIT_EQ_COUNT = 1 + +_PASS_CONFIGS = { + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_SAFE_MEMORY_ACCESS: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, +} + + +def _ordered_fp32_key(score): + """Map fp32 to an integer key whose unsigned order matches fp32 order.""" + 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)) + + +def is_sparse_index_topk_supported(k: int) -> bool: + """Return whether the TileLang byte-radix path has a compiled + specialization.""" + return k in _SUPPORTED_TOPK + + +@tilelang.jit(pass_configs=_PASS_CONFIGS) +def _sparse_index_topk_byte_radix_kernel(top_k: int, + fill: int = _FILL, + threads: int = _THREADS): + num_tokens = T.dynamic('num_tokens') + score_width = T.dynamic('score_width') + score_stride = T.dynamic('score_stride') + + @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], + 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) + + raw_seqlen = Seqlens[row] + 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. + if seqlen <= top_k: + 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 + # fp32-order key and narrows the selected prefix until the final + # threshold key is known. + prefix_key = T.alloc_var(T.uint32) + prefix_mask = T.alloc_var(T.uint32) + rank = T.alloc_var(T.int32) + prefix_key = T.cast(0, T.uint32) + prefix_mask = T.cast(0, T.uint32) + rank = top_k + + for round_idx in T.Unroll(4): + shift = 24 - round_idx * _RADIX_BITS + byte_mask = T.cast(255, T.uint32) << shift + + for bin_idx in T.Parallel(_RADIX_SIZE): + histogram[bin_idx] = 0 + + pos = T.alloc_var(T.int32) + pos = tidx + 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) + pos += threads + + T.sync_threads() + + if tidx < _RADIX_SIZE: + above_count = T.alloc_var(T.int32) + above_count = 0 + for other_bin in T.serial(0, _RADIX_SIZE): + if other_bin > tidx: + above_count += histogram[other_bin] + bin_count = histogram[tidx] + if above_count < rank and above_count + bin_count >= rank: + shared_state[_STATE_SELECTED_BIN] = tidx + shared_state[_STATE_COUNT_ABOVE] = above_count + + 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_mask = T.bitwise_or(prefix_mask, byte_mask) + rank -= shared_state[_STATE_COUNT_ABOVE] + + threshold_key = prefix_key + + # Reuse shared_state as output counters after threshold search: + # [0] counts scores above threshold, [1] counts scores equal to it. + if tidx == 0: + shared_state[_STATE_EMIT_GT_COUNT] = 0 + shared_state[_STATE_EMIT_EQ_COUNT] = 0 + T.sync_threads() + + 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 + while pos_emit_gt < seqlen: + 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) + if out_pos_buf[0] < top_k: + Out[row, out_pos_buf[0]] = pos_emit_gt + pos_emit_gt += threads + + T.sync_threads() + + gt_count = shared_state[_STATE_EMIT_GT_COUNT] + + # Then fill remaining slots from scores equal to the threshold. + # Tie order is intentionally unspecified; sparse attention needs + # a valid top-k set, not score-sorted ids. + pos_emit_eq = T.alloc_var(T.int32) + pos_emit_eq = tidx + while pos_emit_eq < seqlen: + 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) + if out_pos_buf[0] < top_k: + Out[row, out_pos_buf[0]] = pos_emit_eq + pos_emit_eq += threads + + return sparse_index_topk_byte_radix_kernel_ + + +def sparse_index_topk(scores: torch.Tensor, + q_seqlens: torch.Tensor, + kv_seqlens: torch.Tensor, + k: int, + fill: int = _FILL, + descending: bool = True, + 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 + 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``. + """ + if not descending: + raise ValueError('sparse_index_topk only supports descending=True.') + 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}.') + 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 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) + 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) + return out diff --git a/tests/pytorch/kernel/test_sparse_index_topk.py b/tests/pytorch/kernel/test_sparse_index_topk.py new file mode 100644 index 0000000000..f2d6b92b0a --- /dev/null +++ b/tests/pytorch/kernel/test_sparse_index_topk.py @@ -0,0 +1,162 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import torch + + +def _requires_cuda() -> bool: + return not torch.cuda.is_available() + + +pytestmark = pytest.mark.skipif(_requires_cuda(), reason='requires CUDA') + + +def test_nsa_backend_selects_sparse_topk_for_glm52(): + from lmdeploy.pytorch.backends.cuda.nsa import _get_sparse_index_topk + + assert _get_sparse_index_topk(2048).__name__ == 'sparse_index_topk' + assert _get_sparse_index_topk(1024) is None + + +def _assert_topk_ids(scores: torch.Tensor, + out: torch.Tensor, + seqlens: list[int], + k: int, + fill: int = -1): + scores = scores.cpu() + out = out.cpu() + score_width = scores.size(1) + + for row, raw_seqlen in enumerate(seqlens): + seqlen = min(raw_seqlen, score_width) + row_out = out[row] + valid = row_out[row_out != fill] + + if seqlen <= k: + expected = torch.arange(seqlen, dtype=torch.int32) + torch.testing.assert_close(row_out[:seqlen], expected) + if seqlen < k: + assert row_out[seqlen:].eq(fill).all() + continue + + expected = torch.topk(scores[row, :seqlen], + k=k, + largest=True, + sorted=False).indices + assert valid.numel() == k + torch.testing.assert_close(valid.sort().values, + expected.to(torch.int32).sort().values) + + +@pytest.mark.parametrize('k', [512, 2048]) +def test_sparse_index_topk_matches_torch_topk_and_fill(k: int): + from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import ( + is_sparse_index_topk_supported, + sparse_index_topk, + ) + + assert is_sparse_index_topk_supported(k) + + device = 'cuda' + fill = -7 + score_width = k * 2 + seqlens = [0, 17, k, k + 1, score_width - 124] + generator = torch.Generator(device=device).manual_seed(20260709 + k) + scores = torch.randn(len(seqlens), + score_width, + device=device, + dtype=torch.float32, + generator=generator) + scores += torch.arange(score_width, device=device, + dtype=torch.float32) * 1e-6 + + q_seqlens = torch.ones(len(seqlens), device=device, dtype=torch.int64) + kv_dtype = torch.int64 if k == 2048 else torch.int32 + kv_seqlens = torch.tensor(seqlens, device=device, dtype=kv_dtype) + + out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k, fill=fill) + assert out.shape == (len(seqlens), k) + assert out.dtype == torch.int32 + _assert_topk_ids(scores, out, seqlens, k, fill=fill) + + +def test_sparse_index_topk_accepts_padded_score_stride(): + from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import sparse_index_topk + + device = 'cuda' + k = 512 + score_width = 1024 + padded_width = 1280 + seqlens = [600, 777, 1024, 321] + generator = torch.Generator(device=device).manual_seed(20260710) + storage = torch.randn(len(seqlens), + padded_width, + device=device, + dtype=torch.float32, + generator=generator) + scores = storage[:, :score_width] + scores += torch.arange(score_width, device=device, + dtype=torch.float32) * 1e-6 + assert not scores.is_contiguous() + + q_seqlens = torch.ones(len(seqlens), device=device, dtype=torch.int64) + kv_seqlens = torch.tensor(seqlens, device=device, dtype=torch.int32) + + out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k) + _assert_topk_ids(scores, out, seqlens, k) + + +def test_sparse_index_topk_expands_batch_kv_seqlens_for_prefill(): + from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import sparse_index_topk + + device = 'cuda' + k = 512 + score_width = 1536 + q_seqlens_list = [2, 3] + batch_kv_seqlens = [640, 1100] + row_seqlens = [640, 640, 1100, 1100, 1100] + + generator = torch.Generator(device=device).manual_seed(260619348) + scores = torch.randn(sum(q_seqlens_list), + score_width, + device=device, + dtype=torch.float32, + generator=generator) + scores += torch.arange(score_width, device=device, + dtype=torch.float32) * 1e-6 + + q_seqlens = torch.tensor(q_seqlens_list, device=device, dtype=torch.int64) + kv_seqlens = torch.tensor(batch_kv_seqlens, + device=device, + dtype=torch.int32) + + out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k) + _assert_topk_ids(scores, out, row_seqlens, k) + + +@pytest.mark.parametrize('k', [512, 2048]) +def test_sparse_index_topk_cuda_graph_capture(k: int): + from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import sparse_index_topk + + device = 'cuda' + score_width = k * 2 + seqlens = [k + 1, k + 265, score_width, k // 2] + generator = torch.Generator(device=device).manual_seed(k) + scores = torch.randn(len(seqlens), + score_width, + device=device, + dtype=torch.float32, + generator=generator) + q_seqlens = torch.ones(len(seqlens), device=device, dtype=torch.int64) + kv_seqlens = torch.tensor(seqlens, device=device, dtype=torch.int64) + + # Warm the TileLang specialization and PyTorch's graph-aware allocator. + out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k) + + graph.replay() + torch.cuda.synchronize() + _assert_topk_ids(scores, out, seqlens, k) From a1e577056aaf09ce587e580653b5030fff24cfac Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 14 Jul 2026 15:08:11 +0800 Subject: [PATCH 09/33] feat: support GLM-5.2 expert and indexer replay --- benchmark/benchmark_chat_completion.py | 34 ++++- lmdeploy/cli/serve.py | 2 + lmdeploy/cli/utils.py | 9 ++ lmdeploy/messages.py | 6 + lmdeploy/pytorch/config.py | 2 + lmdeploy/pytorch/engine/engine.py | 5 +- lmdeploy/pytorch/engine/engine_instance.py | 57 ++++++--- lmdeploy/pytorch/engine/engine_loop.py | 8 +- lmdeploy/pytorch/engine/inputs_maker.py | 6 + lmdeploy/pytorch/engine/model_agent/agent.py | 11 ++ lmdeploy/pytorch/messages.py | 60 +++++++++ lmdeploy/pytorch/model_inputs.py | 1 + lmdeploy/pytorch/models/deepseek_v2.py | 13 +- lmdeploy/pytorch/models/deepseek_v32.py | 67 +++++++++- lmdeploy/pytorch/models/utils/cudagraph.py | 2 + lmdeploy/pytorch/paging/block_trie.py | 72 ++++++++++- lmdeploy/pytorch/strategies/ar/sequence.py | 18 ++- .../pytorch/strategies/ar_spec/sequence.py | 36 +++++- lmdeploy/serve/anthropic/adapter.py | 1 + .../serve/anthropic/endpoints/messages.py | 10 +- lmdeploy/serve/anthropic/protocol.py | 7 ++ lmdeploy/serve/anthropic/streaming.py | 7 +- lmdeploy/serve/core/async_engine.py | 8 ++ lmdeploy/serve/openai/api_server.py | 33 ++++- lmdeploy/serve/openai/protocol.py | 8 ++ .../serve/openai/serving_chat_completion.py | 4 + lmdeploy/serve/openai/serving_generate.py | 7 ++ tests/pytorch/engine/test_engine_sleep.py | 34 ++++- tests/pytorch/models/test_deepseek_v32.py | 117 +++++++++++++++++- tests/pytorch/paging/test_block_trie.py | 54 +++++++- tests/pytorch/spec_decode/test_strategies.py | 45 +++++++ .../serve/test_generation_config.py | 11 ++ 32 files changed, 702 insertions(+), 53 deletions(-) diff --git a/benchmark/benchmark_chat_completion.py b/benchmark/benchmark_chat_completion.py index bba0c05289..0a3d73f1d3 100644 --- a/benchmark/benchmark_chat_completion.py +++ b/benchmark/benchmark_chat_completion.py @@ -7,7 +7,7 @@ aggregates TTFT/ITL/TPOT metrics, and writes table plus report artifacts for concurrency/RPS sweeps. Generation options include ``--output-tokens`` (``max_completion_tokens``), -``--ignore-eos``, ``--return-token-ids``, ``--return-routed-experts``, +``--ignore-eos``, ``--return-token-ids``, ``--return-routed-experts``, ``--return-indexer-topk``, ``--return-logprob``, and ``--logprobs`` / ``--top-logprobs``. """ @@ -56,6 +56,7 @@ class SSEEvent: done: bool = False raw: dict[str, Any] | None = None routed_experts: str | None = None + indexer_topk: str | None = None @property def token_text(self) -> str: @@ -131,15 +132,19 @@ def init_shared_store() -> Any: return _shared_store_actor -async def fetch_routed_experts(shared_store: Any, key: str) -> Any: - """Fetch routed_experts from shared_store without blocking the event - loop.""" +async def fetch_shared_output(shared_store: Any, key: str) -> Any: + """Fetch a large replay output without blocking the event loop.""" import ray ref = shared_store.get.remote(key) return await asyncio.to_thread(ray.get, ref) +async def fetch_routed_experts(shared_store: Any, key: str) -> Any: + """Backward-compatible routed-expert fetch helper.""" + return await fetch_shared_output(shared_store, key) + + def _split_csv(value: str | None) -> list[str] | None: if value is None: return None @@ -354,6 +359,7 @@ def parse_sse_line(line: bytes | str) -> SSEEvent: usage=data.get('usage'), raw=data, routed_experts=choice.get('routed_experts'), + indexer_topk=choice.get('indexer_topk'), ) @@ -374,6 +380,7 @@ def build_payload( ignore_eos: bool = False, return_token_ids: bool = False, return_routed_experts: bool = False, + return_indexer_topk: bool = False, return_logprob: bool = False, logprobs: bool = False, top_logprobs: int | None = None, @@ -404,6 +411,8 @@ def build_payload( payload['return_token_ids'] = True if return_routed_experts: payload['return_routed_experts'] = True + if return_indexer_topk: + payload['return_indexer_topk'] = True if return_logprob: payload['return_logprob'] = True if logprobs: @@ -460,6 +469,7 @@ async def request_chat_completion( ignore_eos: bool, return_token_ids: bool, return_routed_experts: bool, + return_indexer_topk: bool, return_logprob: bool, logprobs: bool, top_logprobs: int | None, @@ -478,6 +488,7 @@ async def request_chat_completion( ignore_eos=ignore_eos, return_token_ids=return_token_ids, return_routed_experts=return_routed_experts, + return_indexer_topk=return_indexer_topk, return_logprob=return_logprob, logprobs=logprobs, top_logprobs=top_logprobs, @@ -530,6 +541,11 @@ async def request_chat_completion( await fetch_routed_experts(shared_store, event.routed_experts) except Exception as e: # noqa: BLE001 - record and keep consuming SSE. trace.error = repr(e) + if event.indexer_topk and shared_store is not None: + try: + await fetch_shared_output(shared_store, event.indexer_topk) + except Exception as e: # noqa: BLE001 - record and keep consuming SSE. + trace.error = repr(e) trace.end_time = time.perf_counter() trace.success = trace.error == '' @@ -1016,7 +1032,7 @@ async def run_benchmark(args: argparse.Namespace) -> tuple[list[RequestTrace], l extra_body = json.loads(args.extra_request_body) if args.extra_request_body else {} shared_store = None - if args.return_routed_experts: + if args.return_routed_experts or args.return_indexer_topk: shared_store = init_shared_store() tokenizer = None @@ -1057,6 +1073,8 @@ async def run_benchmark(args: argparse.Namespace) -> tuple[list[RequestTrace], l print('return_token_ids=True') if args.return_routed_experts: print('return_routed_experts=True') + if args.return_indexer_topk: + print('return_indexer_topk=True') if args.return_logprob: print('return_logprob=True') if args.logprobs: @@ -1078,6 +1096,7 @@ async def send_one(request: BenchmarkRequest, mode: str, setting: float, repeat: ignore_eos=args.ignore_eos, return_token_ids=args.return_token_ids, return_routed_experts=args.return_routed_experts, + return_indexer_topk=args.return_indexer_topk, return_logprob=args.return_logprob, logprobs=args.logprobs, top_logprobs=args.top_logprobs, @@ -1273,6 +1292,11 @@ def parse_args() -> argparse.Namespace: action='store_true', help='Set return_routed_experts=true to include MoE routed expert indices (LMDeploy extension).', ) + parser.add_argument( + '--return-indexer-topk', + action='store_true', + help='Set return_indexer_topk=true and fetch sparse-attention indexer results (LMDeploy extension).', + ) parser.add_argument( '--return-logprob', action='store_true', diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index 47c16bfe52..b4c070885b 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -110,6 +110,7 @@ def add_parser_api_server(): ArgumentHelper.dllm_denoising_steps(pt_group) ArgumentHelper.dllm_confidence_threshold(pt_group) ArgumentHelper.enable_return_routed_experts(pt_group) + ArgumentHelper.enable_return_indexer_topk(pt_group) ArgumentHelper.distributed_executor_backend(pt_group) ArgumentHelper.kernel_block_size(pt_group) ArgumentHelper.prefix_cache_state_budget(pt_group) @@ -262,6 +263,7 @@ def api_server(args): dllm_denoising_steps=args.dllm_denoising_steps, dllm_confidence_threshold=args.dllm_confidence_threshold, enable_return_routed_experts=args.enable_return_routed_experts, + enable_return_indexer_topk=args.enable_return_indexer_topk, distributed_executor_backend=args.distributed_executor_backend, ) else: diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 8d5f7d31d0..71579b9603 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -786,6 +786,15 @@ def enable_return_routed_experts(parser): default=False, help='Whether to output routed expert ids for replay') + @staticmethod + def enable_return_indexer_topk(parser): + """Add argument return sparse-attention indexer top-k to parser.""" + + return parser.add_argument('--enable-return-indexer-topk', + action='store_true', + default=False, + help='Whether to output sparse-attention indexer top-k ids for replay') + @staticmethod def add_spec_group(parser): spec_group = parser.add_argument_group('Speculative decoding arguments') diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index 28a59739a6..be6befe2ec 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -146,6 +146,7 @@ class GenerationConfig: # router replay return_routed_experts: bool = False + return_indexer_topk: bool = False # ngram, generation would stop if latest [size] tokens are repeated for [threshold] times repetition_ngram_size: int = 0 @@ -464,6 +465,7 @@ class PytorchEngineConfig: logprobs_mode: str = None # router replay enable_return_routed_experts: bool = False + enable_return_indexer_topk: bool = False enable_transfer_obj_ref: bool = False # dllm @@ -562,6 +564,7 @@ class Response: last_hidden_state: torch.Tensor = None index: int = 0 routed_experts: Any = None + indexer_topk: Any = None cached_tokens: int = 0 def __str__(self): @@ -592,6 +595,7 @@ def _format_tensor(name: str, tensor: torch.Tensor | None) -> list[str]: fields.extend(_format_tensor('logits', self.logits)) fields.extend(_format_tensor('last_hidden_state', self.last_hidden_state)) fields.extend(_format_tensor('routed_experts', self.routed_experts)) + fields.extend(_format_tensor('indexer_topk', self.indexer_topk)) return '\n'.join(fields) def extend(self, other: 'Response') -> 'Response': @@ -618,6 +622,7 @@ def extend(self, other: 'Response') -> 'Response': self.logprobs = self.logprobs or [] self.logprobs += other.logprobs self.routed_experts = other.routed_experts + self.indexer_topk = other.indexer_topk return self @@ -704,6 +709,7 @@ class EngineOutput: cache_block_ids: list[int] | None = None req_metrics: RequestMetrics | None = None routed_experts: torch.Tensor = None + indexer_topk: torch.Tensor = None ce_loss: float = None diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index b56695f638..a1058b1adb 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -572,6 +572,7 @@ class MiscConfig: logprobs_mode: str = None dllm_config: DLLMConfig = None enable_return_routed_experts: bool = False + enable_return_indexer_topk: bool = False enable_chunked_prefill: bool = False @classmethod @@ -592,6 +593,7 @@ def from_engine_config(cls, engine_config: PytorchEngineConfig): logprobs_mode=engine_config.logprobs_mode, dllm_config=dllm_config, enable_return_routed_experts=engine_config.enable_return_routed_experts, + enable_return_indexer_topk=engine_config.enable_return_indexer_topk, enable_chunked_prefill=False, ) return misc_config diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py index 5ccb268cb0..8988ff92ee 100644 --- a/lmdeploy/pytorch/engine/engine.py +++ b/lmdeploy/pytorch/engine/engine.py @@ -59,7 +59,10 @@ class InferOutput: req_metrics: RequestMetrics = None # expert ids - routed_experts: torch.Tensor = None + routed_experts: np.ndarray = None + + # sparse-attention indexer results + indexer_topk: 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 d650981af6..d5b43321f0 100644 --- a/lmdeploy/pytorch/engine/engine_instance.py +++ b/lmdeploy/pytorch/engine/engine_instance.py @@ -137,23 +137,46 @@ 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) + outputs = dict(routed_experts=None, indexer_topk=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) + + indexer_topk = resp.data.get('indexer_topk', None) if resp.data else None + if indexer_topk is not None: + indexer_topk = _validate_num_tokens('indexer_topk', indexer_topk) + outputs['indexer_topk'] = _maybe_transfer(indexer_topk) return outputs async def _async_try_add_session(self, session_id: int): @@ -251,8 +274,9 @@ 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) + indexer_topk = extra_outputs.get('indexer_topk', None) logger.debug(f'session[{session_id}] finish: num_out_ids={num_ids}.') yield EngineOutput(resp.type, @@ -261,6 +285,7 @@ async def async_stream_infer(self, cache_block_ids=cache_block_ids, req_metrics=req_metrics, routed_experts=routed_experts, + indexer_topk=indexer_topk, logprobs=logprobs, ce_loss=ce_loss) break diff --git a/lmdeploy/pytorch/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py index 8b39cab777..32fd9c440b 100644 --- a/lmdeploy/pytorch/engine/engine_loop.py +++ b/lmdeploy/pytorch/engine/engine_loop.py @@ -227,6 +227,7 @@ def _send_resp(self, out: InferOutput): cache_block_ids=out.cache_block_ids, req_metrics=out.req_metrics, routed_experts=out.routed_experts, + indexer_topk=out.indexer_topk, logprobs=logprobs, ce_loss=out.ce_loss)) @@ -315,15 +316,18 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'): logits = batched_outputs.logits all_routed_experts = batched_outputs.all_routed_experts + all_indexer_topk = batched_outputs.all_indexer_topk ce_loss = batched_outputs.ce_loss if model_inputs is not None and (model_inputs.is_chunk and not model_inputs.is_last_chunk): # chunk long context does not need to update seqs and outputs seq = running[0] seq.append_routed_experts(all_routed_experts) + seq.append_indexer_topk(all_indexer_topk) seq.append_logits(logits) seq.append_ce_loss(ce_loss, finish=False) self.scheduler.block_trie.cache_routed_experts_for_seq(seq) + self.scheduler.block_trie.cache_indexer_topk_for_seq(seq) return dict() new_token_timestamp = batched_outputs.new_token_timestamp @@ -338,6 +342,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'): model_inputs=model_inputs, delta=delta) self.scheduler.block_trie.cache_routed_experts(running) + self.scheduler.block_trie.cache_indexer_topk(running) # generate output outputs: dict[int, InferOutput] = dict() @@ -382,7 +387,8 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'): cache_block_ids=cache_block_ids, req_metrics=req_metrics, logprobs=cur_logprobs, - routed_experts=msg.routed_experts) + routed_experts=msg.routed_experts, + indexer_topk=msg.indexer_topk) outputs[session_id] = out if msg.return_ce_loss: diff --git a/lmdeploy/pytorch/engine/inputs_maker.py b/lmdeploy/pytorch/engine/inputs_maker.py index 95a5865abd..86d7dd12d2 100644 --- a/lmdeploy/pytorch/engine/inputs_maker.py +++ b/lmdeploy/pytorch/engine/inputs_maker.py @@ -820,6 +820,10 @@ def __need_routed_experts(seqs: 'SeqList'): """Need routed experts.""" return any(seq.return_routed_experts for seq in seqs) + def __need_indexer_topk(seqs: 'SeqList'): + """Need sparse-attention indexer top-k results.""" + return any(getattr(seq, 'return_indexer_topk', False) for seq in seqs) + def __need_ce_loss(seqs: 'SeqList'): """Need input cross-entropy loss.""" return any(seq.return_ce_loss for seq in seqs) @@ -1086,6 +1090,7 @@ def __can_retry_deferred_active_chunk(): return_logits = __need_logits(running) return_routed_experts = __need_routed_experts(running) + return_indexer_topk = __need_indexer_topk(running) return_ce_loss = __need_ce_loss(running) return dict( @@ -1099,6 +1104,7 @@ def __can_retry_deferred_active_chunk(): return_logits=return_logits, extra_inputs=extra_inputs, return_routed_experts=return_routed_experts, + return_indexer_topk=return_indexer_topk, return_ce_loss=return_ce_loss, ) diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py index bc772f3976..5ff26a5a94 100644 --- a/lmdeploy/pytorch/engine/model_agent/agent.py +++ b/lmdeploy/pytorch/engine/model_agent/agent.py @@ -98,6 +98,7 @@ class BatchedOutputs: new_token_timestamp: int = 0 extra_outputs: ExtraOutputs | None = None all_routed_experts: torch.Tensor | None = None + all_indexer_topk: torch.Tensor | None = None ce_loss: torch.Tensor | None = None def to_cpu(self): @@ -679,6 +680,7 @@ async def _step_postprocess_with_output(self, return_ce_loss: bool = False, seq_length: torch.Tensor = None, all_routed_experts: Any = None, + all_indexer_topk: Any = None, extra_inputs: ExtraInputs = None): """Step postprocess with output.""" rank = self.rank @@ -728,6 +730,7 @@ async def _step_postprocess_with_output(self, model_metas=model_metas, logprobs=logprobs, all_routed_experts=all_routed_experts, + all_indexer_topk=all_indexer_topk, extra_outputs=extra_outputs, ce_loss=ce_loss)) @@ -776,6 +779,7 @@ async def _async_step( stopping_criteria: StoppingCriteria = None, return_logits: bool = False, return_routed_experts: bool = False, + return_indexer_topk: bool = False, return_ce_loss: bool = False, extra_inputs: ExtraInputs = None, ): @@ -859,6 +863,10 @@ async def _async_step( all_routed_experts = output.get('all_routed_experts', None) else: all_routed_experts = None + if return_indexer_topk: + all_indexer_topk = output.get('all_indexer_topk', None) + else: + all_indexer_topk = None ( inputs, @@ -879,6 +887,7 @@ async def _async_step( return_ce_loss=return_ce_loss, seq_length=seq_length, all_routed_experts=all_routed_experts, + all_indexer_topk=all_indexer_topk, extra_inputs=extra_inputs, )) else: @@ -1110,11 +1119,13 @@ def _build_model(self): logger.debug(msg_with_rank(rank, 'build model.')) # for router replay enable_return_routed_experts = self.misc_config.enable_return_routed_experts and self.need_output + enable_return_indexer_topk = self.misc_config.enable_return_indexer_topk and self.need_output build_model_ctx = BuildModelContext(language_model_only=self.misc_config.language_model_only, dllm_config=self.misc_config.dllm_config, strategy_factory=self.strategy_factory, enable_return_routed_experts=enable_return_routed_experts, + enable_return_indexer_topk=enable_return_indexer_topk, quant_config=self.model_config.quant_config, fp32_lm_head=self.model_config.fp32_lm_head, tie_word_embeddings=self.model_config.tie_word_embeddings, diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py index b1c5a59e68..dad987e2a3 100644 --- a/lmdeploy/pytorch/messages.py +++ b/lmdeploy/pytorch/messages.py @@ -142,6 +142,7 @@ class SamplingParam: out_ce_loss: bool = False num_logprobs: int = -1 return_routed_experts: bool = False + return_indexer_topk: bool = False # ngram repetition_ngram_size: int = 0 @@ -239,6 +240,7 @@ def from_gen_config(cls, gen_config: GenerationConfig): out_ce_loss=gen_config.return_ppl, num_logprobs=logprobs, return_routed_experts=gen_config.return_routed_experts, + return_indexer_topk=gen_config.return_indexer_topk, repetition_ngram_size=repetition_ngram_size, repetition_ngram_threshold=repetition_ngram_threshold, ) @@ -589,6 +591,37 @@ def _get_pad_width(self, reserve_size: int): return ((0, reserve_size), (0, 0), (0, 0)) +class HistoryIndexerTopK(_HistoryDataBase): + """History of sparse-attention indexer results.""" + ALLOC_SIZE = 1 + + def __init__(self, indexer_topk: np.ndarray = None, dtype: np.dtype = np.int32): + super().__init__(indexer_topk, dtype) + + def _create_empty_array(self, dtype): + return None + + def _get_pad_width(self, reserve_size: int): + return ((0, reserve_size), (0, 0), (0, 0)) + + def append(self, new_data: np.ndarray, reserve_size: int | None = None): + """Append results, reserving the expected request length once.""" + new_data = np.asarray(new_data) + if new_data.ndim != 3: + raise ValueError(f'indexer_topk must be a 3D array, got shape {new_data.shape}.') + if self._data is None: + capacity = len(new_data) + if reserve_size is not None: + capacity = max(capacity, reserve_size) + self._data = np.empty((capacity, *new_data.shape[1:]), dtype=self.dtype) + elif self._data.shape[1:] != new_data.shape[1:]: + raise ValueError( + f'indexer_topk shape changed from {self._data.shape[1:]} to {new_data.shape[1:]}.') + if reserve_size is not None: + self.reserve(reserve_size) + super().append(new_data) + + class HistoryLogits(_HistoryDataBase): """History logits.""" ALLOC_SIZE = 64 @@ -737,6 +770,9 @@ class SchedulerSequence: # for router replay all_routed_experts: HistoryRouterExperts = field(default_factory=HistoryRouterExperts) + # exact sparse-attention indexer results + all_indexer_topk: HistoryIndexerTopK = field(default_factory=HistoryIndexerTopK) + # logits all_logits: HistoryLogits = field(default_factory=HistoryLogits) @@ -841,6 +877,30 @@ def append_routed_experts(self, routed_experts: Tensor | np.ndarray): routed_experts = routed_experts.cpu().numpy() self.all_routed_experts.append(routed_experts) + @property + def return_indexer_topk(self) -> bool: + return self.sampling_param.return_indexer_topk + + @property + def indexer_topk(self) -> np.ndarray: + if not self.return_indexer_topk: + return None + + end = max(0, self.num_valid_ids - 1) + if 0 < end <= len(self.all_indexer_topk): + return self.all_indexer_topk.get_real()[:end] + return None + + def append_indexer_topk(self, indexer_topk: Tensor | np.ndarray): + """Append exact sparse-attention indexer results.""" + if not self.return_indexer_topk or indexer_topk is None: + return + if isinstance(indexer_topk, Tensor): + indexer_topk = indexer_topk.cpu().numpy() + reserve_size = self.output_start_pos + self.sampling_param.max_new_tokens + reserve_size = max(reserve_size, len(self.all_indexer_topk) + len(indexer_topk)) + self.all_indexer_topk.append(indexer_topk, reserve_size=reserve_size) + @property def num_history_ids(self): """Num history ids.""" diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py index ff7164a2ea..6be75f1b36 100644 --- a/lmdeploy/pytorch/model_inputs.py +++ b/lmdeploy/pytorch/model_inputs.py @@ -465,6 +465,7 @@ class BuildModelContext: dllm_config: DLLMConfig = None strategy_factory: 'StrategyFactoryBase' = None enable_return_routed_experts: bool = False + enable_return_indexer_topk: bool = False quant_config: QuantizationConfig = field(default_factory=QuantizationConfig) fp32_lm_head: bool = False tie_word_embeddings: bool = False diff --git a/lmdeploy/pytorch/models/deepseek_v2.py b/lmdeploy/pytorch/models/deepseek_v2.py index 2eec92b838..aa750ffeba 100644 --- a/lmdeploy/pytorch/models/deepseek_v2.py +++ b/lmdeploy/pytorch/models/deepseek_v2.py @@ -643,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: @@ -673,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) @@ -684,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 @@ -738,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, diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index b110da70c6..dbcdc26dbd 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -31,6 +31,7 @@ DeepseekV2MoE, yarn_get_mscale, ) +from .patch import get_build_model_context def get_layer_indexer_type(config: Any, layer_idx: int | None) -> str: @@ -383,6 +384,7 @@ def forward( past_key_value: Sequence[torch.Tensor] = None, attn_metadata: Any = None, topk_indices_buffer: DSATopKIndicesBuffer | None = None, + all_indexer_topk: torch.Tensor | None = None, ): """Rewrite of LlamaAttention.forward.""" dist_config = get_dist_manager().current_config() @@ -420,6 +422,9 @@ def forward( if topk_indices_buffer is not None: topk_indices = topk_indices_buffer.write(topk_indices) + if all_indexer_topk is not None: + all_indexer_topk[:, self.layer_idx, :].copy_(topk_indices) + attn_output = self.attn_fwd( query_states, key_states, @@ -481,6 +486,8 @@ def forward( residual: torch.Tensor | None = None, attn_metadata: Any = None, topk_indices_buffer: DSATopKIndicesBuffer | None = None, + all_routed_experts: torch.Tensor | None = None, + all_indexer_topk: torch.Tensor | None = None, ) -> tuple[torch.FloatTensor, torch.FloatTensor]: if residual is None: @@ -496,6 +503,7 @@ def forward( past_key_value=past_key_value, attn_metadata=attn_metadata, topk_indices_buffer=topk_indices_buffer, + all_indexer_topk=all_indexer_topk, ) else: hidden_states = self.self_attn( @@ -506,7 +514,10 @@ def forward( ) hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) + 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 @@ -557,6 +568,8 @@ def forward( 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, + all_indexer_topk: torch.Tensor | None = None, ): """forward.""" if inputs_embeds is None: @@ -576,6 +589,8 @@ def forward( residual=residual, attn_metadata=attn_metadata, topk_indices_buffer=self.topk_indices_buffer, + all_routed_experts=all_routed_experts, + all_indexer_topk=all_indexer_topk, ) hidden_states, _ = self.norm(hidden_states, residual) @@ -589,6 +604,8 @@ def forward_microbatch( 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, + all_indexer_topk: torch.Tensor | None = None, ): """forward_microbatch.""" # DSA shared top-k indices are model-global; use normal forward until @@ -599,6 +616,8 @@ def forward_microbatch( past_key_values=past_key_values, attn_metadata=attn_metadata, inputs_embeds=inputs_embeds, + all_routed_experts=all_routed_experts, + all_indexer_topk=all_indexer_topk, ) @@ -622,6 +641,52 @@ 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 + self.enable_return_indexer_topk = bm_ctx.enable_return_indexer_topk + + 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, + **kwargs, + ): + """Model forward with optional exact DSA indexer capture.""" + 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: + 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, + ) + all_indexer_topk = None + if self.enable_return_indexer_topk: + all_indexer_topk = position_ids.new_empty( + (num_tokens, self.config.num_hidden_layers, self.config.index_topk), dtype=torch.int32) + + 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, + all_indexer_topk=all_indexer_topk, + ) + if all_routed_experts is None and all_indexer_topk is None: + return hidden_states + outputs = dict(hidden_states=hidden_states) + if all_routed_experts is not None: + outputs['all_routed_experts'] = all_routed_experts + if all_indexer_topk is not None: + outputs['all_indexer_topk'] = all_indexer_topk + return outputs def _load_weight_attention(self, name: str, loaded_weight: torch.Tensor, params_dict: dict[str, nn.Parameter], update_pe_mapping: list): diff --git a/lmdeploy/pytorch/models/utils/cudagraph.py b/lmdeploy/pytorch/models/utils/cudagraph.py index decec282bb..30006d090a 100644 --- a/lmdeploy/pytorch/models/utils/cudagraph.py +++ b/lmdeploy/pytorch/models/utils/cudagraph.py @@ -356,4 +356,6 @@ def get_outputs_cudagraph(self, output_buffers: dict[str, torch.Tensor], input_i outputs['hidden_states'] = output_buffers['hidden_states'][:, :num_tokens] if output_buffers.get('all_routed_experts', None) is not None: outputs['all_routed_experts'] = output_buffers['all_routed_experts'][:num_tokens, ...].clone() + if output_buffers.get('all_indexer_topk', None) is not None: + outputs['all_indexer_topk'] = output_buffers['all_indexer_topk'][:num_tokens, ...].clone() return outputs diff --git a/lmdeploy/pytorch/paging/block_trie.py b/lmdeploy/pytorch/paging/block_trie.py index b595c0b3ab..2cb5239ce0 100644 --- a/lmdeploy/pytorch/paging/block_trie.py +++ b/lmdeploy/pytorch/paging/block_trie.py @@ -133,6 +133,7 @@ def __init__(self, state_ref_count: int = 0, state_access_time: float = 0.0, routed_experts: np.ndarray = None, + indexer_topk: np.ndarray = None, adapter_name: str = None): self.hash_key = hash_key self.block = block @@ -144,6 +145,7 @@ def __init__(self, self.state_ref_count = state_ref_count self.state_access_time = state_access_time self.routed_experts = routed_experts + self.indexer_topk = indexer_topk self.adapter_name = adapter_name self.children: dict[int, Node] = dict() self._parent: Node = None @@ -359,6 +361,56 @@ def cache_routed_experts(self, seqs: list[SchedulerSequence]): for seq in seqs: self.cache_routed_experts_for_seq(seq) + @staticmethod + def _get_indexer_topk_for_range(seq: SchedulerSequence, start: int, end: int): + """Get a copy of indexer results for a full token range, if present.""" + if not seq.return_indexer_topk: + return None + all_indexer_topk = seq.all_indexer_topk + if len(all_indexer_topk) < seq.num_history_ids or len(all_indexer_topk) < end: + return None + indexer_topk = all_indexer_topk.get_real() + if indexer_topk is None or len(indexer_topk) < end: + return None + return indexer_topk[start:end].copy() + + def _try_cache_node_indexer_topk(self, node: Node, seq: SchedulerSequence, start: int, end: int): + """Attach indexer results to a trie node when a sequence has them.""" + if node.indexer_topk is not None: + return + indexer_topk = self._get_indexer_topk_for_range(seq, start, end) + if indexer_topk is not None and len(indexer_topk) == end - start: + node.indexer_topk = indexer_topk + + def _append_matched_indexer_topk(self, seq: SchedulerSequence, nodes: list[Node], start: int): + """Replay cached indexer results for a matched trie range.""" + if not seq.return_indexer_topk or len(nodes) == 0: + return + if len(seq.all_indexer_topk) != start: + return + if any(node.indexer_topk is None or len(node.indexer_topk) != self.block_size for node in nodes): + return + for node in nodes: + seq.append_indexer_topk(node.indexer_topk) + + def cache_indexer_topk_for_seq(self, seq: SchedulerSequence): + """Enrich attached trie nodes with sparse-attention indexer results.""" + if not self.enable or not seq.return_indexer_topk: + return + node = seq.prefix_cache.last_shared_node + while node is not None and node.parent is not None: + end = node.num_matched + start = end - self.block_size + self._try_cache_node_indexer_topk(node, seq, start, end) + node = node.parent + + def cache_indexer_topk(self, seqs: list[SchedulerSequence]): + """Enrich trie nodes with indexer results from multiple sequences.""" + if not self.enable: + return + for seq in seqs: + self.cache_indexer_topk_for_seq(seq) + def _make_state_checkpoint_lookup_key(self, seq: SchedulerSequence, step: int) -> StateCheckpointKey: """Make the sparse SSM checkpoint lookup key for a sequence prefix. @@ -980,6 +1032,12 @@ 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}') + if seq.return_indexer_topk and block_node.indexer_topk is None: + return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.REQUEST_MISMATCH, + reason=f'indexer results missing at block {idx}') matched_blocks.append(block_node.block) return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.HIT, @@ -1039,6 +1097,7 @@ def _match_state_checkpoint(self, seq: SchedulerSequence): seq.logical_blocks.append(matched_blocks) seq.set_step(step) self._append_matched_routed_experts(seq, matched_nodes, init_num_matched) + self._append_matched_indexer_topk(seq, matched_nodes, init_num_matched) seq.prefix_cache.restore_state = node.state_idx seq.prefix_cache.restore_node = node seq.prefix_cache.last_shared_node = node @@ -1108,6 +1167,10 @@ 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 + if seq.return_indexer_topk and child.indexer_topk is None: + break matched_nodes.append(child) __match_success(child) @@ -1130,8 +1193,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) @@ -1142,6 +1206,7 @@ def __clamp_match_step(match_step: int): seq.logical_blocks.append(matched_blocks) seq.set_step(num_matched) self._append_matched_routed_experts(seq, matched_nodes, init_num_matched) + self._append_matched_indexer_topk(seq, matched_nodes, init_num_matched) if self.requires_state_checkpoint: seq.prefix_cache.restore_state = curr.state_idx @@ -1211,18 +1276,21 @@ def allocate(self, seq: SchedulerSequence): # trie-owned block and release this sequence's duplicate block. node = child self._try_cache_node_routed_experts(node, seq, start, end) + self._try_cache_node_indexer_topk(node, seq, start, end) if block != node.block: free_blocks.append(block) logical_blocks[block_id] = node.block blocks.append(node.block) else: routed_experts = self._get_routed_experts_for_range(seq, start, end) + indexer_topk = self._get_indexer_topk_for_range(seq, start, end) node = Node(hash_key=hash_key, block=block, tokens=curr_tokens, num_matched=num_matched + block_size, extra_hashes=extra_hashes, routed_experts=routed_experts, + indexer_topk=indexer_topk, adapter_name=seq.adapter_name) node.parent = parent blocks.append(node.block) diff --git a/lmdeploy/pytorch/strategies/ar/sequence.py b/lmdeploy/pytorch/strategies/ar/sequence.py index 0d39b7e0ae..df03d9d13f 100644 --- a/lmdeploy/pytorch/strategies/ar/sequence.py +++ b/lmdeploy/pytorch/strategies/ar/sequence.py @@ -35,6 +35,7 @@ def update_token_ids(self, model_meta: dict[str, Any] = None, mode: UpdateTokenMode = UpdateTokenMode.INPUTS, routed_experts: np.ndarray = None, + indexer_topk: np.ndarray = None, **kwargs): """Update token ids, old token ids will be added to history.""" # update history image nums @@ -48,6 +49,7 @@ def update_token_ids(self, num_valid = len(token_ids) # record cached expert ids self.append_routed_experts(routed_experts) + self.append_indexer_topk(indexer_topk) if mode == UpdateTokenMode.INPUTS: self.cached_tokens = 0 @@ -91,6 +93,8 @@ def set_step(self, step: int): # chunk long context might not have all routed experts if len(self.all_routed_experts) > step: self.all_routed_experts.resize(step) + if self.return_indexer_topk and len(self.all_indexer_topk) > step: + self.all_indexer_topk.resize(step) def cleanup(self): """Setup history meta after sequence stopped or cancelled.""" @@ -139,13 +143,21 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode if batched_outputs.all_routed_experts is not None: all_routed_experts = batched_outputs.all_routed_experts.split(num_tokens, dim=0) all_routed_experts = [experts.numpy() for experts in all_routed_experts] + all_indexer_topk = [None] * len(num_tokens) + if batched_outputs.all_indexer_topk is not None: + all_indexer_topk = batched_outputs.all_indexer_topk.split(num_tokens, dim=0) + all_indexer_topk = [topk.numpy() for topk in all_indexer_topk] update_mode = UpdateTokenMode.DECODE if is_decoding else UpdateTokenMode.PREFILL - for token, msg, stop, model_meta, routed_experts in zip(next_token_ids, running, stopped, model_metas, - all_routed_experts): + for token, msg, stop, model_meta, routed_experts, indexer_topk in zip( + next_token_ids, running, stopped, model_metas, all_routed_experts, all_indexer_topk): if msg.status != MessageStatus.RUNNING: continue # fill token - msg.update_token_ids(token, model_meta=model_meta, mode=update_mode, routed_experts=routed_experts) + msg.update_token_ids(token, + model_meta=model_meta, + mode=update_mode, + routed_experts=routed_experts, + indexer_topk=indexer_topk) if stop: msg.state.finish() diff --git a/lmdeploy/pytorch/strategies/ar_spec/sequence.py b/lmdeploy/pytorch/strategies/ar_spec/sequence.py index 7bd9ae995a..d7b4d05ce7 100644 --- a/lmdeploy/pytorch/strategies/ar_spec/sequence.py +++ b/lmdeploy/pytorch/strategies/ar_spec/sequence.py @@ -46,9 +46,18 @@ 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 + + @property + def indexer_topk(self) -> np.ndarray: + if not self.return_indexer_topk: return None + end = max(0, self.num_valid_ids - 1) + if 0 < end <= len(self.all_indexer_topk): + return self.all_indexer_topk.get_real()[:end] + return None + @property def generated_ids(self) -> np.ndarray: end = self.num_valid_ids @@ -67,7 +76,8 @@ def _update_token_ids_inputs(self, token_ids: np.ndarray): self.history_cache.append(token_ids) def _update_token_ids_prefill(self, token_ids: np.ndarray, draft_token_ids: np.ndarray, - stop_pos: int = -1, routed_experts: np.ndarray = None): + stop_pos: int = -1, routed_experts: np.ndarray = None, + indexer_topk: np.ndarray = None): """Update token ids for prefill.""" # back to last valid position self.history_cache.resize(self.num_valid_ids) @@ -75,6 +85,7 @@ def _update_token_ids_prefill(self, token_ids: np.ndarray, draft_token_ids: np.n num_valid = len(token_ids) self.history_cache.append(token_ids) self.append_routed_experts(routed_experts) + self.append_indexer_topk(indexer_topk) self._num_history_ids += self._num_token_ids self.num_new_tokens += num_valid self._num_valid_ids = self.num_history_ids + num_valid @@ -85,7 +96,8 @@ def _update_token_ids_prefill(self, token_ids: np.ndarray, draft_token_ids: np.n self.history_cache.append(draft_token_ids) def _update_token_ids_decode(self, token_ids: np.ndarray, draft_token_ids: np.ndarray, - stop_pos: int = -1, routed_experts: np.ndarray = None): + stop_pos: int = -1, routed_experts: np.ndarray = None, + indexer_topk: np.ndarray = None): """Update token ids for decode.""" # back to last valid position self.history_cache.resize(self.num_valid_ids) @@ -104,6 +116,9 @@ def _update_token_ids_decode(self, token_ids: np.ndarray, draft_token_ids: np.nd if routed_experts is not None: routed_experts = routed_experts[:num_valid] self.append_routed_experts(routed_experts) + if indexer_topk is not None: + indexer_topk = indexer_topk[:num_valid] + self.append_indexer_topk(indexer_topk) if stop_pos > -1: self._num_token_ids = 1 @@ -120,6 +135,7 @@ def update_token_ids(self, draft_token_ids: Tensor = None, mode: UpdateTokenMode = UpdateTokenMode.INPUTS, routed_experts: np.ndarray = None, + indexer_topk: np.ndarray = None, stop_pos: int = -1, **kwargs): """Update token ids, old token ids will be added to history.""" @@ -142,10 +158,12 @@ def update_token_ids(self, self._update_token_ids_inputs(token_ids) elif mode == UpdateTokenMode.PREFILL: self._update_token_ids_prefill(token_ids, draft_token_ids, - stop_pos=stop_pos, routed_experts=routed_experts) + stop_pos=stop_pos, routed_experts=routed_experts, + indexer_topk=indexer_topk) else: self._update_token_ids_decode(token_ids, draft_token_ids, - stop_pos=stop_pos, routed_experts=routed_experts) + stop_pos=stop_pos, routed_experts=routed_experts, + indexer_topk=indexer_topk) if model_meta is not None: self.model_meta = model_meta @@ -171,6 +189,8 @@ def set_step(self, step: int): # chunk long context might not have all routed experts if len(self.all_routed_experts) > step: self.all_routed_experts.resize(step) + if self.return_indexer_topk and len(self.all_indexer_topk) > step: + self.all_indexer_topk.resize(step) def cleanup(self): """Setup history meta after sequence stopped or cancelled.""" @@ -222,6 +242,10 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode if batched_outputs.all_routed_experts is not None: all_routed_experts = batched_outputs.all_routed_experts.split(num_tokens, dim=0) all_routed_experts = [experts.numpy() for experts in all_routed_experts] + all_indexer_topk = [None] * len(num_tokens) + if batched_outputs.all_indexer_topk is not None: + all_indexer_topk = batched_outputs.all_indexer_topk.split(num_tokens, dim=0) + all_indexer_topk = [topk.numpy() for topk in all_indexer_topk] batch_size = len(running) next_token_ids = next_token_ids.view(batch_size, -1).numpy() @@ -234,6 +258,7 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode for idx, token in enumerate(next_token_ids): routed_experts = all_routed_experts[idx] + indexer_topk = all_indexer_topk[idx] msg = running[idx] stop = stopped[idx] model_meta = model_metas[idx] @@ -246,6 +271,7 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode model_meta=model_meta, mode=update_mode, routed_experts=routed_experts, + indexer_topk=indexer_topk, stop_pos=stop_pos[idx]) if stop: msg.state.finish() diff --git a/lmdeploy/serve/anthropic/adapter.py b/lmdeploy/serve/anthropic/adapter.py index c1b2553897..4c86fd8ce4 100644 --- a/lmdeploy/serve/anthropic/adapter.py +++ b/lmdeploy/serve/anthropic/adapter.py @@ -356,6 +356,7 @@ def to_generation_config( skip_special_tokens=True, spaces_between_special_tokens=True, return_routed_experts=request.return_routed_experts or False, + return_indexer_topk=request.return_indexer_topk or False, logprobs=1 if request.return_logprob else None, ) diff --git a/lmdeploy/serve/anthropic/endpoints/messages.py b/lmdeploy/serve/anthropic/endpoints/messages.py index 6027fb3030..77aed30447 100644 --- a/lmdeploy/serve/anthropic/endpoints/messages.py +++ b/lmdeploy/serve/anthropic/endpoints/messages.py @@ -58,6 +58,12 @@ def _validate_extended_outputs(request: MessagesRequest, server_context): ('routed experts requested but not configured in engine configuration. ' 'May start the api_server with --enable-return-routed-experts flag.')) + if request.return_indexer_topk and not engine_config.enable_return_indexer_topk: + return create_error_response( + HTTPStatus.BAD_REQUEST, + ('indexer top-k requested but not configured in engine configuration. ' + 'May start the api_server with --enable-return-indexer-topk flag.')) + return None @@ -200,6 +206,7 @@ async def create_message(request: MessagesRequest, raw_request: Request): response_parser=response_parser, return_token_ids=request.return_token_ids or False, return_routed_experts=request.return_routed_experts or False, + return_indexer_topk=request.return_indexer_topk or False, logprobs=request.return_logprob or False, ), [result_generator], @@ -238,7 +245,7 @@ async def create_message(request: MessagesRequest, raw_request: Request): return create_error_response(HTTPStatus.BAD_REQUEST, f'Failed to parse output: {err}') should_validate_complete = ( final_res.finish_reason in ('stop', 'length') - and (request.return_token_ids or request.return_routed_experts) + and (request.return_token_ids or request.return_routed_experts or request.return_indexer_topk) ) if should_validate_complete and not response_parser.validate_complete(raw_text): final_res.finish_reason = 'parse_error' @@ -269,5 +276,6 @@ async def create_message(request: MessagesRequest, raw_request: Request): output_ids=final_token_ids if request.return_token_ids else None, output_token_logprobs=output_token_logprobs, routed_experts=final_res.routed_experts if request.return_routed_experts else None, + indexer_topk=final_res.indexer_topk if request.return_indexer_topk else None, ) return response.model_dump() diff --git a/lmdeploy/serve/anthropic/protocol.py b/lmdeploy/serve/anthropic/protocol.py index 37fa0f7bd7..3d7a6b4f82 100644 --- a/lmdeploy/serve/anthropic/protocol.py +++ b/lmdeploy/serve/anthropic/protocol.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field RoutedExperts = list[list[list[int]]] | str | None +IndexerTopK = list[list[list[int]]] | str | None MessageStopReason = Literal['end_turn', 'max_tokens', 'stop_sequence', 'tool_use', 'parse_error'] @@ -130,6 +131,10 @@ class MessagesRequest(BaseModel): default=False, description=('Whether to return MoE routed expert indices in the response.'), ) + return_indexer_topk: bool | None = Field( + default=False, + description=('Whether to return sparse-attention indexer top-k indices in the response.'), + ) return_token_ids: bool | None = Field( default=False, description=('Whether to include output token IDs in the response.'), @@ -194,6 +199,7 @@ class MessagesResponse(BaseModel): output_ids: list[int] | None = None output_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) routed_experts: RoutedExperts = None + indexer_topk: IndexerTopK = None class StreamTextBlock(BaseModel): @@ -312,6 +318,7 @@ class MessageDeltaEvent(BaseModel): delta: MessageDelta usage: MessageDeltaUsage routed_experts: RoutedExperts = None + indexer_topk: IndexerTopK = None class MessageStopEvent(BaseModel): diff --git a/lmdeploy/serve/anthropic/streaming.py b/lmdeploy/serve/anthropic/streaming.py index 934c4b0158..5bb0139f98 100644 --- a/lmdeploy/serve/anthropic/streaming.py +++ b/lmdeploy/serve/anthropic/streaming.py @@ -31,7 +31,7 @@ ThinkingDelta, ) -_OPTIONAL_EXTENSION_FIELDS = ('output_ids', 'output_token_logprobs', 'routed_experts') +_OPTIONAL_EXTENSION_FIELDS = ('output_ids', 'output_token_logprobs', 'routed_experts', 'indexer_topk') def _format_sse(data: AnthropicStreamEvent) -> str: @@ -169,6 +169,7 @@ async def stream_messages_response(result_generator, response_parser, return_token_ids: bool = False, return_routed_experts: bool = False, + return_indexer_topk: bool = False, logprobs: bool = False) -> AsyncGenerator[str, None]: """Convert LMDeploy generation stream to Anthropic SSE events.""" @@ -226,7 +227,7 @@ async def _results(): should_validate_complete = ( res.finish_reason in ('stop', 'length') - and (return_token_ids or return_routed_experts) + and (return_token_ids or return_routed_experts or return_indexer_topk) ) if should_validate_complete and not response_parser.validate_complete(): res.finish_reason = 'parse_error' @@ -315,10 +316,12 @@ async def _results(): output_tokens = 0 if final_res is None else final_res.generate_token_len stop_reason = map_finish_reason(None if final_res is None else final_res.finish_reason) routed_experts = final_res.routed_experts if return_routed_experts and final_res is not None else None + indexer_topk = final_res.indexer_topk if return_indexer_topk and final_res is not None else None yield _format_sse( MessageDeltaEvent( delta=MessageDelta(stop_reason=stop_reason, stop_sequence=None), usage=MessageDeltaUsage(output_tokens=output_tokens), routed_experts=routed_experts, + indexer_topk=indexer_topk, )) yield _format_sse(MessageStopEvent()) diff --git a/lmdeploy/serve/core/async_engine.py b/lmdeploy/serve/core/async_engine.py index b3c4db92dd..1aaa796052 100644 --- a/lmdeploy/serve/core/async_engine.py +++ b/lmdeploy/serve/core/async_engine.py @@ -54,6 +54,7 @@ class GenOut: last_hidden_state: Any = None cache_block_ids: list[int] | None = None # for disaggregation routed_experts: Any = None # for RL router replay + indexer_topk: Any = None # for sparse-attention indexer replay cached_tokens: int = 0 def to_response(self, index: int = 0) -> Response: @@ -71,6 +72,7 @@ def to_response(self, index: int = 0) -> Response: last_hidden_state=self.last_hidden_state, logits=self.logits, routed_experts=self.routed_experts, + indexer_topk=self.indexer_topk, cached_tokens=self.cached_tokens, index=index) @@ -708,6 +710,7 @@ def is_error(status): finish_reason, token_ids=res, routed_experts=outputs.routed_experts, + indexer_topk=outputs.indexer_topk, cache_block_ids=outputs.cache_block_ids, cached_tokens=cached_tokens) if outputs.logprobs is not None: @@ -748,6 +751,10 @@ def is_error(status): if routed_experts is not None and not isinstance(routed_experts, str) and ( not gen_config.include_stop_str_in_output) and finish_reason == 'stop': routed_experts = routed_experts[:-1] + indexer_topk = outputs.indexer_topk + if indexer_topk is not None and not isinstance(indexer_topk, str) and ( + not gen_config.include_stop_str_in_output) and finish_reason == 'stop': + indexer_topk = indexer_topk[:-1] logger.info(f'session {session_id} finished, reason ' f'"{finish_reason}", input_tokens ' @@ -762,6 +769,7 @@ def is_error(status): logits=logits, last_hidden_state=last_hidden_state, routed_experts=routed_experts, + indexer_topk=indexer_topk, cache_block_ids=outputs.cache_block_ids, cached_tokens=cached_tokens) # Note: We remove the session step update here. Let the caller(e.g., pipeline.chat) take care of it. diff --git a/lmdeploy/serve/openai/api_server.py b/lmdeploy/serve/openai/api_server.py index 8b7979a597..6d42ea8cf0 100644 --- a/lmdeploy/serve/openai/api_server.py +++ b/lmdeploy/serve/openai/api_server.py @@ -555,6 +555,7 @@ def create_stream_response_json(index: int, logprobs: ChoiceLogprobs | None = None, output_token_logprobs: list[tuple[float, int]] | None = None, routed_experts=None, + indexer_topk=None, output_ids=None) -> dict: choice_data = ChatCompletionResponseStreamChoice(index=index, delta=delta_message, @@ -562,7 +563,8 @@ def create_stream_response_json(index: int, logprobs=logprobs, output_token_logprobs=output_token_logprobs, output_ids=output_ids, - routed_experts=routed_experts) + routed_experts=routed_experts, + indexer_topk=indexer_topk) choice_data = maybe_filter_parallel_tool_calls(choice_data, request) response = ChatCompletionStreamResponse( id=request_id, @@ -616,7 +618,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: stream_deltas = [(DeltaMessage(role='assistant', content=''), False)] should_validate_complete = ( res.finish_reason in ('stop', 'length') - and (request.return_token_ids or request.return_routed_experts) + and (request.return_token_ids or request.return_routed_experts or request.return_indexer_topk) ) if should_validate_complete and not response_parser.validate_complete(): res.finish_reason = 'parse_error' @@ -638,6 +640,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: # Only output routed_experts in the final chunk routed_experts = res.routed_experts if finish_reason is not None else None + indexer_topk = res.indexer_topk if finish_reason is not None else None # Emit token ids once per engine yield on the last parsed delta, when # accumulated delta text and token ids for this step are aligned. stream_output_ids = delta_token_ids if (request.return_token_ids and is_last_delta) else None @@ -648,6 +651,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: logprobs=chunk_logprobs, output_token_logprobs=chunk_output_token_logprobs, routed_experts=routed_experts, + indexer_topk=indexer_topk, output_ids=stream_output_ids) if res.cache_block_ids is not None and is_last_delta: response_json['cache_block_ids'] = res.cache_block_ids @@ -692,7 +696,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: text, tool_calls, reasoning_content = response_parser.parse_complete(text, final_token_ids) should_validate_complete = ( final_res.finish_reason in ('stop', 'length') - and (request.return_token_ids or request.return_routed_experts) + and (request.return_token_ids or request.return_routed_experts or request.return_indexer_topk) ) if should_validate_complete and not response_parser.validate_complete(raw_text): final_res.finish_reason = 'parse_error' @@ -726,6 +730,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: finish_reason=final_res.finish_reason, output_ids=final_token_ids if request.return_token_ids else None, routed_experts=final_res.routed_experts if request.return_routed_experts else None, + indexer_topk=final_res.indexer_topk if request.return_indexer_topk else None, ) choice_data = maybe_filter_parallel_tool_calls(choice_data, request) choices.append(choice_data) @@ -1024,16 +1029,24 @@ 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, + indexer_topk=None): # only output router experts in last chunk routed_experts = None if finish_reason is None else routed_experts + indexer_topk = None if finish_reason is None else indexer_topk meta = GenerateReqMetaOutput(finish_reason=dict(type=finish_reason) if finish_reason else None, output_token_logprobs=logprobs or None, prompt_tokens=res.input_token_len, routed_experts=routed_experts, + indexer_topk=indexer_topk, 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(): @@ -1041,6 +1054,7 @@ async def generate_stream_generator(): text = res.response or '' output_ids = res.token_ids routed_experts = res.routed_experts + indexer_topk = res.indexer_topk logprobs = [] if res.logprobs: for tok, tok_logprobs in zip(res.token_ids, res.logprobs): @@ -1050,7 +1064,8 @@ async def generate_stream_generator(): output_ids, logprobs, res.finish_reason, - routed_experts=routed_experts) + routed_experts=routed_experts, + indexer_topk=indexer_topk) yield f'data: {response_json}\n\n' yield 'data: [DONE]\n\n' @@ -1084,6 +1099,7 @@ async def _inner_call(): output_token_logprobs=output_token_logprobs or None, prompt_tokens=res.input_token_len, routed_experts=res.routed_experts, + indexer_topk=res.indexer_topk, completion_tokens=res.generate_token_len) response = GenerateReqOutput(text=text, output_ids=output_ids, meta_info=meta) @@ -1610,6 +1626,11 @@ def serve(model_path: str, # router replay if backend_config.enable_return_routed_experts: backend_config.enable_transfer_obj_ref = True + if backend_config.enable_return_indexer_topk: + if backend_config.distributed_executor_backend != 'ray': + raise ValueError('--enable-return-indexer-topk requires ' + '--distributed-executor-backend ray for HTTP serving.') + backend_config.enable_transfer_obj_ref = True VariableInterface.async_engine = pipeline_class(model_path=model_path, model_name=model_name, backend=backend, diff --git a/lmdeploy/serve/openai/protocol.py b/lmdeploy/serve/openai/protocol.py index fd6f1370f2..bdf39c9e3a 100644 --- a/lmdeploy/serve/openai/protocol.py +++ b/lmdeploy/serve/openai/protocol.py @@ -226,6 +226,10 @@ class ChatCompletionRequest(BaseModel): default=False, description=('Whether to return MoE routed expert indices in the response.'), ) + return_indexer_topk: bool | None = Field( + default=False, + description=('Whether to return sparse-attention indexer top-k indices in the response.'), + ) class FunctionCall(BaseModel): @@ -293,6 +297,7 @@ class ChatCompletionResponseChoice(BaseModel): finish_reason: Literal['stop', 'length', 'tool_calls', 'parse_error', 'error', 'abort'] | None = None output_ids: list[int] | None = None routed_experts: list[list[list[int]]] | str | None = None + indexer_topk: list[list[list[int]]] | str | None = None class ChatCompletionResponse(BaseModel): @@ -335,6 +340,7 @@ class ChatCompletionResponseStreamChoice(BaseModel): output_ids: list[int] | None = None finish_reason: Literal['stop', 'length', 'tool_calls', 'parse_error', 'error', 'abort'] | None = None routed_experts: list[list[list[int]]] | str | None = None + indexer_topk: list[list[list[int]]] | str | None = None class ChatCompletionStreamResponse(BaseModel): @@ -559,6 +565,7 @@ class GenerateReqInput(BaseModel): spaces_between_special_tokens: bool | None = True include_stop_str_in_output: bool | None = False return_routed_experts: bool | None = False + return_indexer_topk: bool | None = False repetition_ngram_size: int = Field(default=0, ge=0) repetition_ngram_threshold: int = Field(default=0, ge=0) # kwargs for media IO @@ -579,6 +586,7 @@ class GenerateReqMetaOutput(BaseModel): finish_reason: dict[str, Any] | None = None output_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) routed_experts: list[list[list[int]]] | str | None = None # (num_token, num_layer, topk_expert) + indexer_topk: list[list[list[int]]] | str | None = None # (num_token, num_layer, index_topk) # /generate output diff --git a/lmdeploy/serve/openai/serving_chat_completion.py b/lmdeploy/serve/openai/serving_chat_completion.py index 0270dd8546..a5e9e5d741 100644 --- a/lmdeploy/serve/openai/serving_chat_completion.py +++ b/lmdeploy/serve/openai/serving_chat_completion.py @@ -65,4 +65,8 @@ def check_request(request: ChatCompletionRequest, server_context: 'VariableInter return ('routed experts requested but not configured in engine configuration. ' 'May start api_server with --enable-return-routed-experts flag.') + if request.return_indexer_topk and not engine_config.enable_return_indexer_topk: + return ('indexer top-k requested but not configured in engine configuration. ' + 'May start api_server with --enable-return-indexer-topk flag.') + return '' diff --git a/lmdeploy/serve/openai/serving_generate.py b/lmdeploy/serve/openai/serving_generate.py index 10e9dde597..66a79bbde7 100644 --- a/lmdeploy/serve/openai/serving_generate.py +++ b/lmdeploy/serve/openai/serving_generate.py @@ -19,6 +19,13 @@ 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.return_indexer_topk and not engine_config.enable_return_indexer_topk: + return ('indexer top-k requested but not configured in engine configuration. ' + 'May start api_server with --enable-return-indexer-topk 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/tests/pytorch/engine/test_engine_sleep.py b/tests/pytorch/engine/test_engine_sleep.py index f6efdee5cf..6a9e45547f 100644 --- a/tests/pytorch/engine/test_engine_sleep.py +++ b/tests/pytorch/engine/test_engine_sleep.py @@ -1,10 +1,12 @@ # Copyright (c) OpenMMLab. All rights reserved. import asyncio +import sys from types import SimpleNamespace +import numpy as np import pytest -from lmdeploy.messages import EngineOutput, ResponseType +from lmdeploy.messages import EngineOutput, GenerationConfig, ResponseType from lmdeploy.pytorch.engine.engine import Engine from lmdeploy.pytorch.engine.engine_instance import EngineInstance from lmdeploy.pytorch.engine.executor.mp_executor import MPExecutor @@ -205,3 +207,33 @@ async def __collect(): assert len(outputs) == 1 assert outputs[0].status == ResponseType.CANCEL assert engine.req_manager._loop_task is None + + +def test_engine_instance_transfers_trimmed_indexer_topk(monkeypatch): + import lmdeploy.pytorch.engine.engine_instance as engine_instance_module + + transferred = [] + + class _Put: + + def remote(self, data): + transferred.append(data.copy()) + return f'key-{len(transferred)}' + + store = SimpleNamespace(put=_Put()) + monkeypatch.setattr(engine_instance_module, '_SHARED_STORE', store) + monkeypatch.setitem(sys.modules, 'ray', SimpleNamespace(get=lambda key: key)) + + instance = SimpleNamespace(_enable_transfer_obj_ref=True) + indexer_topk = np.arange(5 * 2 * 3, dtype=np.int32).reshape(5, 2, 3) + resp = Response(type=ResponseType.FINISH, + sender_id=0, + event=None, + data=dict(token_ids=np.array([10, 11]), indexer_topk=indexer_topk)) + gen_config = GenerationConfig(stop_token_ids=[11], include_stop_str_in_output=False) + + outputs = EngineInstance._get_extra_outputs(instance, resp, num_all_ids=6, gen_config=gen_config) + + assert outputs['indexer_topk'] == 'key-1' + assert transferred[0].shape == (4, 2, 3) + assert np.array_equal(transferred[0], indexer_topk[:-1]) diff --git a/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index 418b444085..7494c481ce 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -4,9 +4,10 @@ import torch from torch import nn -from lmdeploy.pytorch.models.deepseek_v2 import DeepseekV2BMM +from lmdeploy.pytorch.models.deepseek_v2 import DeepseekV2BMM, DeepseekV2MoE, MoEGate from lmdeploy.pytorch.models.deepseek_v32 import ( DeepseekV32Attention, + DeepseekV32ForCausalLM, DSATopKIndicesBuffer, Indexer, apply_interleaved_rotary_pos_emb, @@ -113,17 +114,21 @@ def fake_attn_fwd(*args, **kwargs): attn.attn_fwd = fake_attn_fwd topk_buffer = DSATopKIndicesBuffer(topk=3) + all_indexer_topk = torch.full((2, 4, 3), -1, dtype=torch.int32) output = attn( hidden_states=torch.zeros(1, 2, 2), rotary_pos_emb=(torch.zeros(2, 1), torch.zeros(2, 1)), past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], topk_indices_buffer=topk_buffer, + all_indexer_topk=all_indexer_topk, ) assert output.shape == (1, 2, 1) assert torch.equal(topk_buffer.indices[:2], computed_topk) assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:2].data_ptr() + assert torch.equal(all_indexer_topk[:, 0], computed_topk) + assert torch.all(all_indexer_topk[:, 1:] == -1) @pytest.mark.parametrize('dp,attn_tp,expected_num_heads', [(1, 2, 2), (2, 2, 4)]) @@ -179,17 +184,22 @@ def fake_attn_fwd(*args, **kwargs): attn.attn_fwd = fake_attn_fwd topk_buffer = DSATopKIndicesBuffer(topk=3) - topk_buffer.write(torch.tensor([[5, 4, 3]], dtype=torch.int32)) + shared_topk = torch.tensor([[5, 4, 3]], dtype=torch.int32) + topk_buffer.write(shared_topk) + all_indexer_topk = torch.full((1, 4, 3), -1, dtype=torch.int32) output = attn( hidden_states=torch.zeros(1, 1, 2), rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], topk_indices_buffer=topk_buffer, + all_indexer_topk=all_indexer_topk, ) assert output.shape == (1, 1, 1) assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:1].data_ptr() + assert torch.equal(all_indexer_topk[:, 3], shared_topk) + assert torch.all(all_indexer_topk[:, :3] == -1) def test_shared_indexer_layer_requires_shared_topk_buffer(monkeypatch): @@ -229,3 +239,106 @@ def test_layer_indexer_type_defaults_to_full_and_reads_shared_entries(): assert get_layer_indexer_type(config, 1) == 'shared' assert get_layer_indexer_type(config, 2) == 'full' assert get_layer_indexer_type(SimpleNamespace(indexer_types=None), 1) == 'full' + + +def test_moe_gate_captures_logical_experts_before_eplb_mapping(monkeypatch): + class FakeTopK(nn.Module): + + def forward(self, logits): + logical_ids = torch.tensor([[3, 1], [2, 0]], device=logits.device) + return torch.ones_like(logical_ids, dtype=torch.float32), logical_ids + + gate = MoEGate.__new__(MoEGate) + nn.Module.__init__(gate) + gate.weight = nn.Parameter(torch.zeros(4, 2)) + gate.fake_eplb = False + gate.topk_method = 'greedy' + gate.renormalize = False + gate.routed_scaling_factor = 1.0 + gate.softmax_topk = FakeTopK() + gate.eplb_dispatch_info = object() + monkeypatch.setattr( + 'lmdeploy.pytorch.models.deepseek_v2.EPLBManager.topk_ids_logical_to_physical', + lambda ids, info: ids + 10, + ) + captured = torch.full((2, 2), torch.iinfo(torch.uint16).max, dtype=torch.uint16) + + _, dispatch_ids = gate(torch.zeros(2, 2), routed_experts=captured) + + expected = torch.tensor([[3, 1], [2, 0]]) + assert torch.equal(captured, expected.to(torch.uint16)) + assert torch.equal(dispatch_ids, expected + 10) + + +def test_deepseek_moe_writes_only_its_routed_expert_layer(): + class FakeGate(nn.Module): + + def forward(self, hidden_states, routed_experts=None): + ids = torch.tensor([[3, 1], [2, 0]], device=hidden_states.device) + if routed_experts is not None: + routed_experts.copy_(ids) + return torch.ones_like(ids, dtype=torch.float32), ids + + class FakeExperts(nn.Module): + + def forward(self, hidden_states, topk_weights, topk_ids): + return hidden_states + + moe = DeepseekV2MoE.__new__(DeepseekV2MoE) + nn.Module.__init__(moe) + moe.layer_idx = 3 + moe.hidden_dim = 2 + moe.gate = FakeGate() + moe.experts = FakeExperts() + moe.shared_experts = None + moe._all_reduce = False + sentinel = torch.iinfo(torch.uint16).max + all_routed_experts = torch.full((2, 5, 2), sentinel, dtype=torch.uint16) + + output = moe(torch.zeros(1, 2, 2), all_routed_experts=all_routed_experts) + + assert output.shape == (1, 2, 2) + assert torch.equal(all_routed_experts[:, 3], torch.tensor([[3, 1], [2, 0]], dtype=torch.uint16)) + assert torch.all(all_routed_experts[:, :3] == sentinel) + assert torch.all(all_routed_experts[:, 4] == sentinel) + + +def test_glm52_causal_lm_returns_routed_experts_and_indexer_topk(monkeypatch): + class FakeModel(nn.Module): + + def forward(self, input_ids, all_routed_experts=None, all_indexer_topk=None, **kwargs): + if all_routed_experts is not None: + all_routed_experts[:, 3].fill_(7) + all_routed_experts[:, 4].fill_(9) + if all_indexer_topk is not None: + all_indexer_topk.fill_(5) + return torch.zeros(1, input_ids.size(1), 4) + + forward_microbatch = forward + + context = SimpleNamespace(enable_microbatch=False) + monkeypatch.setattr( + 'lmdeploy.pytorch.models.deepseek_v32.get_step_ctx_manager', + lambda: SimpleNamespace(current_context=lambda: context), + ) + model = DeepseekV32ForCausalLM.__new__(DeepseekV32ForCausalLM) + nn.Module.__init__(model) + model.config = SimpleNamespace(num_hidden_layers=5, num_experts_per_tok=8, index_topk=3) + model.enable_return_routed_experts = True + model.enable_return_indexer_topk = True + model.model = FakeModel() + + outputs = model( + input_ids=torch.ones(1, 2, dtype=torch.long), + position_ids=torch.arange(2)[None], + past_key_values=[], + ) + + routed_experts = outputs['all_routed_experts'] + assert routed_experts.dtype == torch.uint16 + assert routed_experts.shape == (2, 5, 8) + assert torch.all(routed_experts[:, :3] == torch.iinfo(torch.uint16).max) + assert torch.all(routed_experts[:, 3] == 7) + assert torch.all(routed_experts[:, 4] == 9) + assert outputs['all_indexer_topk'].shape == (2, 5, 3) + assert torch.all(outputs['all_indexer_topk'] == 5) diff --git a/tests/pytorch/paging/test_block_trie.py b/tests/pytorch/paging/test_block_trie.py index 3ccb9f777b..c7fbd4b6d1 100644 --- a/tests/pytorch/paging/test_block_trie.py +++ b/tests/pytorch/paging/test_block_trie.py @@ -110,6 +110,10 @@ def _routed_experts(self, num_tokens: int, offset: int = 0): values = np.arange(offset, offset + num_tokens * 2, dtype=np.uint16) return values.reshape(num_tokens, 2, 1) + def _indexer_topk(self, num_tokens: int, offset: int = 0): + values = np.arange(offset, offset + num_tokens * 6, dtype=np.int32) + return values.reshape(num_tokens, 2, 3) + def _add_ready_ssm_checkpoint(self, scheduler, token_ids): seq = scheduler.add_session(len(scheduler.sessions)).add_sequence(token_ids) scheduler.block_manager.allocate(seq) @@ -446,6 +450,46 @@ def test_match_skips_routed_expert_replay_when_not_requested(self, block_trie, b assert matched.num_history_ids == block_size * 2 assert len(matched.all_routed_experts) == 0 + def test_match_replays_cached_indexer_topk(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] + sampling_param = SamplingParam(return_indexer_topk=True) + seq = sess.add_sequence(token_ids, sampling_param=sampling_param) + indexer_topk = self._indexer_topk(block_size * 2) + + block_mgr.allocate(seq) + block_trie.allocate(seq) + seq.append_indexer_topk(indexer_topk) + block_trie.cache_indexer_topk_for_seq(seq) + + matched = sess.add_sequence(token_ids, sampling_param=sampling_param) + block_trie.match(matched) + + assert matched.num_history_ids == block_size * 2 + assert np.array_equal(matched.all_indexer_topk.get_real(), indexer_topk) + + def test_match_stops_before_block_missing_indexer_topk(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] + sampling_param = SamplingParam(return_indexer_topk=True) + seq = sess.add_sequence(token_ids, sampling_param=sampling_param) + first_block_topk = self._indexer_topk(block_size) + + block_mgr.allocate(seq) + block_trie.allocate(seq) + seq.append_indexer_topk(first_block_topk) + block_trie.cache_indexer_topk_for_seq(seq) + + matched = sess.add_sequence(token_ids, sampling_param=sampling_param) + block_trie.match(matched) + + assert matched.num_history_ids == block_size + assert np.array_equal(matched.all_indexer_topk.get_real(), first_block_topk) + assert matched.prefix_cache.private_recompute_start_step == block_size + assert matched.prefix_cache.private_recompute_end_step == block_size * 2 + def test_existing_node_can_be_enriched_with_routed_experts(self, block_trie, block_mgr, scheduler): sess = scheduler.add_session(0) block_size = sess.seq_meta.block_size @@ -469,7 +513,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 +527,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 +547,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_strategies.py b/tests/pytorch/spec_decode/test_strategies.py index af1e827299..cfbe181f23 100644 --- a/tests/pytorch/spec_decode/test_strategies.py +++ b/tests/pytorch/spec_decode/test_strategies.py @@ -540,6 +540,51 @@ def _experts(n, k=2): return np.arange(n * k, dtype=np.uint16).reshape(n, 1, k) +def _make_seq_with_indexer(prefill_tokens=None): + """Create a speculative sequence that captures sparse-attention indices.""" + strategy = ARSpecSequenceStrategy() + seq_meta = SequenceMeta(block_size=16, strategy=strategy) + session = MagicMock() + session.seq_meta = seq_meta + sampling_param = SamplingParam(return_indexer_topk=True, max_new_tokens=8) + seq = SchedulerSequenceARSpec(seq_id=0, session=session, sampling_param=sampling_param) + if prefill_tokens is not None: + seq._update_token_ids_inputs(np.array(prefill_tokens, dtype=np.int64)) + return seq + + +def _indexer_topk(n): + return np.arange(n * 6, dtype=np.int32).reshape(n, 2, 3) + + +class TestIndexerTopKSpecDecode: + + def test_decode_keeps_only_accepted_indexer_rows(self): + seq = _make_seq_with_indexer() + seq._num_valid_ids = 3 + seq._num_history_ids = 2 + seq.history_cache.append(np.array([0, 1, 2, 100, 101], dtype=np.int64)) + seq._num_token_ids = 3 + + seq._update_token_ids_decode( + np.array([30, 40, -1]), + draft_token_ids=np.array([], dtype=np.int64), + indexer_topk=_indexer_topk(3), + ) + + assert len(seq.all_indexer_topk) == 2 + assert np.array_equal(seq.all_indexer_topk.get_real(), _indexer_topk(2)) + + def test_indexer_history_reserves_request_capacity_and_truncates_logically(self): + seq = _make_seq_with_indexer([1, 2, 3, 4]) + seq.append_indexer_topk(_indexer_topk(4)) + + assert len(seq.all_indexer_topk._data) == 12 + seq.set_step(3) + assert len(seq.all_indexer_topk) == 3 + assert len(seq.all_indexer_topk._data) == 12 + + # --------------------------------------------------------------------------- # Tests for routed_experts in _update_token_ids_decode # --------------------------------------------------------------------------- diff --git a/tests/test_lmdeploy/serve/test_generation_config.py b/tests/test_lmdeploy/serve/test_generation_config.py index 2c39ba5fb5..c8aa915997 100644 --- a/tests/test_lmdeploy/serve/test_generation_config.py +++ b/tests/test_lmdeploy/serve/test_generation_config.py @@ -20,6 +20,8 @@ class _FakeEngineConfig: logprobs_mode = None + enable_return_routed_experts = False + enable_return_indexer_topk = False class _FakeSessionManager: @@ -87,6 +89,15 @@ def test_build_generation_config_from_merged_values(): assert gen_config.do_sample is True +def test_indexer_topk_request_flag_reaches_generation_config(): + chat_request = ChatCompletionRequest(model='test', messages='hi', return_indexer_topk=True) + generate_request = GenerateReqInput(prompt='hi', return_indexer_topk=True) + + assert build_generation_config(chat_request, {}).return_indexer_topk is True + assert build_generation_config(generate_request, {}).return_indexer_topk is True + assert 'indexer top-k requested' in check_generate_request(generate_request, _FakeServerContext()) + + def test_build_generation_config_max_new_tokens_defaults_to_none(): request = CompletionRequest(model='test', prompt='hello') gen_config = build_generation_config(request, {}) From 07e6bead1db081bd9038ef9d94a5298d2862278d Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 14 Jul 2026 18:14:03 +0800 Subject: [PATCH 10/33] perf: fuse interleaved indexer rotary embedding --- lmdeploy/pytorch/backends/apply_rotary_emb.py | 2 +- .../pytorch/backends/cuda/apply_rotary_emb.py | 9 ++-- .../backends/default/apply_rotary_emb.py | 29 +++++++--- .../backends/dlinfer/apply_rotary_emb.py | 5 +- .../kernels/cuda/apply_rotary_pos_emb.py | 26 ++++++--- lmdeploy/pytorch/models/deepseek_v32.py | 25 +-------- lmdeploy/pytorch/nn/rotary_embedding.py | 4 +- tests/pytorch/kernel/test_apply_rotary.py | 54 +++++++++++++++++++ tests/pytorch/models/test_deepseek_v32.py | 11 ++-- 9 files changed, 115 insertions(+), 50 deletions(-) diff --git a/lmdeploy/pytorch/backends/apply_rotary_emb.py b/lmdeploy/pytorch/backends/apply_rotary_emb.py index b893b327af..670d694584 100644 --- a/lmdeploy/pytorch/backends/apply_rotary_emb.py +++ b/lmdeploy/pytorch/backends/apply_rotary_emb.py @@ -18,6 +18,6 @@ class ApplyRotaryEmbBuilder(ABC): @staticmethod @abstractmethod - def build(): + def build(interleaved: bool = False): """Build implementation.""" raise NotImplementedError diff --git a/lmdeploy/pytorch/backends/cuda/apply_rotary_emb.py b/lmdeploy/pytorch/backends/cuda/apply_rotary_emb.py index 152e49fc1b..8f7dec54e3 100644 --- a/lmdeploy/pytorch/backends/cuda/apply_rotary_emb.py +++ b/lmdeploy/pytorch/backends/cuda/apply_rotary_emb.py @@ -10,6 +10,9 @@ class TritonApplyRotaryEmbImpl(ApplyRotaryEmbImpl): """Apply rotary embedding implementation.""" + def __init__(self, interleaved: bool = False): + self.interleaved = interleaved + def forward(self, query: Tensor, key: Tensor, cos: Tensor, sin: Tensor, inplace: bool = True): """forward.""" if inplace: @@ -18,13 +21,13 @@ def forward(self, query: Tensor, key: Tensor, cos: Tensor, sin: Tensor, inplace: else: q_embed = torch.empty_like(query) k_embed = torch.empty_like(key) - return apply_rotary_pos_emb(query, key, cos, sin, q_embed, k_embed) + return apply_rotary_pos_emb(query, key, cos, sin, q_embed, k_embed, interleaved=self.interleaved) class TritonApplyRotaryEmbBuilder(ApplyRotaryEmbBuilder): """Apply rotary embedding implementation builder.""" @staticmethod - def build(): + def build(interleaved: bool = False): """Build implementation.""" - return TritonApplyRotaryEmbImpl() + return TritonApplyRotaryEmbImpl(interleaved=interleaved) diff --git a/lmdeploy/pytorch/backends/default/apply_rotary_emb.py b/lmdeploy/pytorch/backends/default/apply_rotary_emb.py index ef738448f6..1240b4a673 100644 --- a/lmdeploy/pytorch/backends/default/apply_rotary_emb.py +++ b/lmdeploy/pytorch/backends/default/apply_rotary_emb.py @@ -16,26 +16,43 @@ def rotate_half(x): return out +def rotate_interleaved(x): + """Rotate adjacent pairs of hidden dimensions.""" + out = torch.empty_like(x) + out[..., ::2] = -x[..., 1::2] + out[..., 1::2] = x[..., ::2] + return out + + class DefaultApplyRotaryEmbImpl(ApplyRotaryEmbImpl): """Apply rotary embedding implementation.""" + def __init__(self, interleaved: bool = False): + self.interleaved = interleaved + def forward(self, query: Tensor, key: Tensor, cos: Tensor, sin: Tensor, inplace: bool = True): """forward.""" unsqueeze_dim = -2 + rotate = rotate_half + if self.interleaved: + half_size = cos.size(-1) // 2 + cos = cos[..., :half_size].repeat_interleave(2, dim=-1) + sin = sin[..., :half_size].repeat_interleave(2, dim=-1) + rotate = rotate_interleaved cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) if inplace: q_embed = query k_embed = key - q_sin = rotate_half(query) * sin + q_sin = rotate(query) * sin q_embed.mul_(cos) q_embed.add_(q_sin) - k_sin = rotate_half(key) * sin + k_sin = rotate(key) * sin k_embed.mul_(cos) k_embed.add_(k_sin) else: - q_embed = (query * cos) + (rotate_half(query) * sin) - k_embed = (key * cos) + (rotate_half(key) * sin) + q_embed = (query * cos) + (rotate(query) * sin) + k_embed = (key * cos) + (rotate(key) * sin) return q_embed, k_embed @@ -43,6 +60,6 @@ class DefaultApplyRotaryEmbBuilder(ApplyRotaryEmbBuilder): """Apply rotary embedding implementation builder.""" @staticmethod - def build(): + def build(interleaved: bool = False): """Build implementation.""" - return DefaultApplyRotaryEmbImpl() + return DefaultApplyRotaryEmbImpl(interleaved=interleaved) diff --git a/lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py b/lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py index f25f4b8cc8..25a9199e3b 100644 --- a/lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py +++ b/lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py @@ -24,6 +24,9 @@ class DlinferApplyRotaryEmbBuilder(ApplyRotaryEmbBuilder): """Apply rotary embedding implementation builder.""" @staticmethod - def build(): + def build(interleaved: bool = False): """Build implementation.""" + if interleaved: + from ..default.apply_rotary_emb import DefaultApplyRotaryEmbImpl + return DefaultApplyRotaryEmbImpl(interleaved=True) return DlinferApplyRotaryEmbImpl() diff --git a/lmdeploy/pytorch/kernels/cuda/apply_rotary_pos_emb.py b/lmdeploy/pytorch/kernels/cuda/apply_rotary_pos_emb.py index 6ed063ef60..3ea3ac21ca 100644 --- a/lmdeploy/pytorch/kernels/cuda/apply_rotary_pos_emb.py +++ b/lmdeploy/pytorch/kernels/cuda/apply_rotary_pos_emb.py @@ -46,6 +46,7 @@ def apply_rotary_pos_emb_qk_kernel( stride_keh: tl.constexpr, stride_ked: tl.constexpr, half_size: tl.constexpr, + interleaved: tl.constexpr, BLOCK: tl.constexpr, BLOCK_QH: tl.constexpr, BLOCK_N: tl.constexpr, @@ -59,13 +60,21 @@ def apply_rotary_pos_emb_qk_kernel( pos_offset = tl.max_contiguous(tl.multiple_of(pos_offset % seq_len, BLOCK), BLOCK) feat_size = half_size * 2 - feat_offset_l = tl.arange(0, BLOCK_N) - feat_mask = feat_offset_l < half_size - feat_offset_l = feat_offset_l % half_size - feat_offset_h = half_size + feat_offset_l + pair_offset = tl.arange(0, BLOCK_N) + feat_mask = pair_offset < half_size + pair_offset = pair_offset % half_size + if interleaved: + feat_offset_l = pair_offset * 2 + feat_offset_h = feat_offset_l + 1 + else: + feat_offset_l = pair_offset + feat_offset_h = half_size + pair_offset seq_mask = pos_mask[:, None] & feat_mask[None, :] - cs_offset_l = pos_offset[:, None] * feat_size + feat_offset_l[None, :] - cs_offset_h = pos_offset[:, None] * feat_size + feat_offset_h[None, :] + cs_offset_l = pos_offset[:, None] * feat_size + pair_offset[None, :] + if interleaved: + cs_offset_h = cs_offset_l + else: + cs_offset_h = pos_offset[:, None] * feat_size + feat_offset_h[None, :] q_elem_type = Q.dtype.element_ty cos_l = tl.load(COS + cs_offset_l).to(q_elem_type) cos_h = tl.load(COS + cs_offset_h).to(q_elem_type) @@ -117,7 +126,8 @@ def apply_rotary_pos_emb(q: Tensor, cos: Tensor, sin: Tensor, q_embed: Tensor = None, - k_embed: Tensor = None): + k_embed: Tensor = None, + interleaved: bool = False): """Apply rotary positional embedding on query and key. Args: @@ -127,6 +137,7 @@ def apply_rotary_pos_emb(q: Tensor, sin (Tensor): sine matrix (seq_len, dim). q_embed (Tensor): output q, can be same as q k_embed (Tensor): output k, can be same as k + interleaved (bool): whether rotary pairs use adjacent dimensions Returns: tuple[Tensor, Tensor]: Embedded query and key. @@ -189,6 +200,7 @@ def apply_rotary_pos_emb(q: Tensor, stride_keh=k_embed.stride(-2), stride_ked=k_embed.stride(-1), half_size=half_size, + interleaved=interleaved, BLOCK=BLOCK, BLOCK_QH=num_heads_q, BLOCK_N=BLOCK_N, diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index dbcdc26dbd..8eb455357f 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -102,27 +102,6 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: return hadamard_transform(x, scale=hidden_size**-0.5) -def rotate_interleaved(x: torch.Tensor) -> torch.Tensor: - """Rotate interleaved RoPE pairs.""" - out = torch.empty_like(x) - out[..., ::2] = -x[..., 1::2] - out[..., 1::2] = x[..., ::2] - return out - - -def apply_interleaved_rotary_pos_emb(query: torch.Tensor, key: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): - """Apply GPT-J style interleaved RoPE to query and key.""" - - def _to_interleaved(freqs: torch.Tensor): - return freqs[..., :freqs.size(-1) // 2].repeat_interleave(2, dim=-1) - - cos = _to_interleaved(cos).unsqueeze(-2) - sin = _to_interleaved(sin).unsqueeze(-2) - query = query * cos + rotate_interleaved(query) * sin - key = key * cos + rotate_interleaved(key) * sin - return query, key - - class LayerNorm(nn.Module): """Layer Normalization.""" @@ -180,7 +159,7 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic device=device, is_tp=False) self.softmax_scale = self.head_dim**-0.5 - self.apply_rotary_pos_emb = ApplyRotaryEmb() + self.apply_rotary_pos_emb = ApplyRotaryEmb(interleaved=self.rope_interleave) 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, @@ -188,8 +167,6 @@ def _apply_rotary_pos_emb(self, q_pe: torch.Tensor, k_pe: torch.Tensor, """Apply the indexer's RoPE layout.""" cos, sin = freqs_cis k_pe = k_pe[..., None, :] - if self.rope_interleave: - return apply_interleaved_rotary_pos_emb(q_pe, k_pe, cos, sin) return self.apply_rotary_pos_emb( q_pe, k_pe, diff --git a/lmdeploy/pytorch/nn/rotary_embedding.py b/lmdeploy/pytorch/nn/rotary_embedding.py index 98fb8d8ce4..5440a68a66 100644 --- a/lmdeploy/pytorch/nn/rotary_embedding.py +++ b/lmdeploy/pytorch/nn/rotary_embedding.py @@ -241,11 +241,11 @@ def build_rotary_embedding_from_config(config: PretrainedConfig, device: torch.d class ApplyRotaryEmb(nn.Module): """Apply rotary embedding.""" - def __init__(self): + def __init__(self, interleaved: bool = False): super().__init__() backend = get_backend() builder = backend.get_layer_impl_builder(OpType.ApplyRotaryEmb) - self.impl = builder.build() + self.impl = builder.build(interleaved=interleaved) def forward(self, query: Tensor, key: Tensor, cos: Tensor, sin: Tensor, inplace: bool = True): """forward.""" diff --git a/tests/pytorch/kernel/test_apply_rotary.py b/tests/pytorch/kernel/test_apply_rotary.py index 3967d1b609..542df8d55d 100644 --- a/tests/pytorch/kernel/test_apply_rotary.py +++ b/tests/pytorch/kernel/test_apply_rotary.py @@ -11,6 +11,13 @@ def _rotate_half(x): return torch.cat((-x2, x1), dim=-1) +def _rotate_interleaved(x): + """Rotate adjacent pairs of hidden dimensions.""" + x1 = x[..., ::2] + x2 = x[..., 1::2] + return torch.stack((-x2, x1), dim=-1).flatten(-2) + + def _bf16_mark(): return pytest.mark.skipif(not is_bf16_supported(), reason='bf16 not supported.') @@ -93,3 +100,50 @@ def test_apply_rotary(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('dtype', [pytest.param(torch.bfloat16, marks=_bf16_mark()), torch.float16, torch.float32], + indirect=True) + @pytest.mark.parametrize(('num_heads_q', 'num_heads_k'), [(8, 8), (8, 4)], indirect=True) + @pytest.mark.parametrize('inplace', [False, True]) + def test_apply_rotary_interleaved(self, q_states, k_states, cos, sin, inplace): + from lmdeploy.pytorch.kernels.cuda import apply_rotary_pos_emb + + half_size = cos.size(-1) // 2 + cos = torch.cat([cos[..., :half_size], cos[..., :half_size]], dim=-1) + sin = torch.cat([sin[..., :half_size], sin[..., :half_size]], dim=-1) + pair_cos = cos[..., :half_size].repeat_interleave(2, dim=-1) + pair_sin = sin[..., :half_size].repeat_interleave(2, dim=-1) + q_gt = q_states * pair_cos + _rotate_interleaved(q_states) * pair_sin + k_gt = k_states * pair_cos + _rotate_interleaved(k_states) * pair_sin + + if inplace: + q_states = q_states.clone() + k_states = k_states.clone() + q_embed, k_embed = apply_rotary_pos_emb( + q_states, k_states, cos, sin, q_embed=q_states, k_embed=k_states, interleaved=True) + else: + q_embed, k_embed = apply_rotary_pos_emb(q_states, k_states, cos, sin, interleaved=True) + + assert torch.equal(q_embed, q_gt) + assert torch.equal(k_embed, k_gt) + + @pytest.mark.parametrize('tokens', [1, 33, 257]) + def test_apply_rotary_interleaved_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 = torch.cat([angles.cos(), angles.cos()], dim=-1).to(torch.bfloat16) + sin = torch.cat([angles.sin(), angles.sin()], dim=-1).to(torch.bfloat16) + pair_cos = cos[..., :32].repeat_interleave(2, dim=-1).unsqueeze(-2) + pair_sin = sin[..., :32].repeat_interleave(2, dim=-1).unsqueeze(-2) + q_gt = query * pair_cos + _rotate_interleaved(query) * pair_sin + k_gt = key * pair_cos + _rotate_interleaved(key) * pair_sin + + q_embed, k_embed = apply_rotary_pos_emb(query, key, cos, sin, interleaved=True) + + assert torch.equal(q_embed, q_gt) + assert torch.equal(k_embed, k_gt) diff --git a/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index 7494c481ce..65afe6f350 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -4,15 +4,14 @@ import torch from torch import nn +from lmdeploy.pytorch.backends.default.apply_rotary_emb import DefaultApplyRotaryEmbImpl, rotate_interleaved from lmdeploy.pytorch.models.deepseek_v2 import DeepseekV2BMM, DeepseekV2MoE, MoEGate from lmdeploy.pytorch.models.deepseek_v32 import ( DeepseekV32Attention, DeepseekV32ForCausalLM, DSATopKIndicesBuffer, Indexer, - apply_interleaved_rotary_pos_emb, get_layer_indexer_type, - rotate_interleaved, ) @@ -46,13 +45,14 @@ def fake_qkv_proj(hidden_states, num_heads): attn.o_proj = nn.Identity() -def test_interleaved_rotary_uses_pairwise_layout_with_half_split_freqs(): +def test_default_interleaved_rotary_uses_pairwise_layout_with_half_split_freqs(): query = torch.tensor([[[[1.0, 2.0, 3.0, 4.0]]]]) key = torch.tensor([[[[5.0, 6.0, 7.0, 8.0]]]]) cos = torch.tensor([[10.0, 20.0, 10.0, 20.0]]) sin = torch.tensor([[1.0, 2.0, 1.0, 2.0]]) - query_out, key_out = apply_interleaved_rotary_pos_emb(query, key, cos, sin) + rotary = DefaultApplyRotaryEmbImpl(interleaved=True) + query_out, key_out = rotary.forward(query, key, cos, sin, inplace=False) interleaved_cos = torch.tensor([[[10.0, 10.0, 20.0, 20.0]]]) interleaved_sin = torch.tensor([[[1.0, 1.0, 2.0, 2.0]]]) @@ -64,7 +64,7 @@ def test_interleaved_rotary_uses_pairwise_layout_with_half_split_freqs(): def test_indexer_rope_interleave_uses_interleaved_layout(): indexer = Indexer.__new__(Indexer) - indexer.rope_interleave = True + indexer.apply_rotary_pos_emb = DefaultApplyRotaryEmbImpl(interleaved=True).forward query = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]]) key = torch.tensor([[5.0, 6.0, 7.0, 8.0]]) @@ -84,7 +84,6 @@ def test_indexer_rope_interleave_uses_interleaved_layout(): def test_indexer_rope_interleave_can_fall_back_to_half_split_layout(): indexer = Indexer.__new__(Indexer) - indexer.rope_interleave = False indexer.apply_rotary_pos_emb = lambda q, k, cos, sin, inplace=False: ('half-split', q, k, cos, sin, inplace) query = torch.zeros(1, 1, 4) From b77a172d4827aed8299faf040c0165bcf74a06a5 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 14 Jul 2026 20:38:18 +0800 Subject: [PATCH 11/33] feat: support GLM-5.2 MoE-only FP8 --- .../pytorch/configurations/deepseek_v32.py | 14 ++++ lmdeploy/pytorch/models/deepseek_v2.py | 4 +- .../pytorch/config/test_glm_moe_dsa_config.py | 78 +++++++++++++++++++ tests/pytorch/models/test_deepseek_v32.py | 68 +++++++++++++++- 4 files changed, 162 insertions(+), 2 deletions(-) diff --git a/lmdeploy/pytorch/configurations/deepseek_v32.py b/lmdeploy/pytorch/configurations/deepseek_v32.py index 08e91ac2cd..d66c28677e 100644 --- a/lmdeploy/pytorch/configurations/deepseek_v32.py +++ b/lmdeploy/pytorch/configurations/deepseek_v32.py @@ -1,8 +1,13 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch +from lmdeploy.pytorch import envs as _envs +from lmdeploy.utils import get_logger + from .deepseek_v2 import DeepseekV2ModelConfigBuilder +logger = get_logger('lmdeploy') + def normalize_glm_moe_dsa_config(hf_config): """Normalize GLM-MoE-DSA config fields for DeepSeek-V3.2 reuse.""" @@ -78,6 +83,15 @@ def condition(cls, hf_config): @classmethod def build(cls, hf_config, model_path: str | None = None, **kwargs): """build.""" + 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 hf_config.model_type == 'glm_moe_dsa' and _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.') + normalize_glm_moe_dsa_config(hf_config) config = DeepseekV2ModelConfigBuilder.build(hf_config, model_path=model_path, **kwargs) diff --git a/lmdeploy/pytorch/models/deepseek_v2.py b/lmdeploy/pytorch/models/deepseek_v2.py index aa750ffeba..574a3a8508 100644 --- a/lmdeploy/pytorch/models/deepseek_v2.py +++ b/lmdeploy/pytorch/models/deepseek_v2.py @@ -1292,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/tests/pytorch/config/test_glm_moe_dsa_config.py b/tests/pytorch/config/test_glm_moe_dsa_config.py index aee7cd0b93..c9ff3d6d89 100644 --- a/tests/pytorch/config/test_glm_moe_dsa_config.py +++ b/tests/pytorch/config/test_glm_moe_dsa_config.py @@ -1,7 +1,9 @@ from types import SimpleNamespace import pytest +import torch +from lmdeploy.pytorch.config import ModelConfig from lmdeploy.pytorch.configurations.deepseek_v32 import ( DeepseekV32ModelConfigBuilder, normalize_glm_moe_dsa_config, @@ -19,6 +21,82 @@ def _make_config(**kwargs): return SimpleNamespace(**values) +def _patch_v32_base_builder(monkeypatch): + from lmdeploy.pytorch.configurations.deepseek_v2 import DeepseekV2ModelConfigBuilder + + def fake_build(cls, hf_config, model_path=None, **kwargs): + return SimpleNamespace() + + monkeypatch.setattr(DeepseekV2ModelConfigBuilder, 'build', classmethod(fake_build)) + + +def _make_fp8_build_config(**quantization_config): + return _make_config( + use_flash_mla=True, + index_head_dim=128, + index_topk=2048, + quantization_config=quantization_config, + ) + + +def test_glm_moe_dsa_enables_online_fp8_moe_only_scope(monkeypatch): + from lmdeploy.pytorch import envs + from lmdeploy.pytorch import transformers as pytorch_transformers + from lmdeploy.pytorch.configurations import deepseek_v2 + from lmdeploy.pytorch.transformers.configuration_glm_moe_dsa import GlmMoeDsaConfig + + monkeypatch.setattr(envs, 'fp8_moe_only', True) + monkeypatch.setattr(deepseek_v2, 'flash_mla_available', lambda: True) + cfg = GlmMoeDsaConfig(hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + q_lora_rank=4, + kv_lora_rank=4, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + vocab_size=32, + bos_token_id=1, + eos_token_id=2, + index_head_dim=4, + index_n_heads=2, + index_topk=2) + monkeypatch.setattr(pytorch_transformers, 'config_from_pretrained', lambda *args, **kwargs: cfg) + + model_config = ModelConfig.from_pretrained('fake-model', model_format='fp8', dtype='bfloat16') + + assert model_config.quant_config.quant_method == 'fp8' + assert model_config.quant_config.fp8_quant_scope == 'moe_only' + assert model_config.quant_config.hf_quant_config['lmdeploy_patched'] + assert model_config.dtype == torch.bfloat16 + + +def test_glm_moe_dsa_fp8_moe_only_scope_requires_env(monkeypatch): + from lmdeploy.pytorch import envs + + monkeypatch.setattr(envs, 'fp8_moe_only', False) + _patch_v32_base_builder(monkeypatch) + cfg = _make_fp8_build_config(quant_method='fp8', lmdeploy_patched=True) + + DeepseekV32ModelConfigBuilder.build(cfg) + + assert 'fp8_quant_scope' not in cfg.quantization_config + + +def test_glm_moe_dsa_does_not_override_prequantized_fp8_scope(monkeypatch): + from lmdeploy.pytorch import envs + + monkeypatch.setattr(envs, 'fp8_moe_only', True) + _patch_v32_base_builder(monkeypatch) + cfg = _make_fp8_build_config(quant_method='fp8') + + DeepseekV32ModelConfigBuilder.build(cfg) + + assert 'fp8_quant_scope' not in cfg.quantization_config + + def test_glm_moe_dsa_normalizes_pattern_indexer_types(): cfg = _make_config(index_topk_pattern='FSSF') diff --git a/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index 65afe6f350..fe88b63afe 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -5,7 +5,7 @@ from torch import nn from lmdeploy.pytorch.backends.default.apply_rotary_emb import DefaultApplyRotaryEmbImpl, rotate_interleaved -from lmdeploy.pytorch.models.deepseek_v2 import DeepseekV2BMM, DeepseekV2MoE, MoEGate +from lmdeploy.pytorch.models.deepseek_v2 import DeepseekV2BMM, DeepseekV2ForCausalLM, DeepseekV2MoE, MoEGate from lmdeploy.pytorch.models.deepseek_v32 import ( DeepseekV32Attention, DeepseekV32ForCausalLM, @@ -45,6 +45,72 @@ def fake_qkv_proj(hidden_states, num_heads): attn.o_proj = nn.Identity() +def _make_kv_b_loader_model(fp8_quant_scope=None, input_dim=3): + model = DeepseekV2ForCausalLM.__new__(DeepseekV2ForCausalLM) + nn.Module.__init__(model) + model.config = SimpleNamespace(qk_nope_head_dim=2, v_head_dim=2) + model.quantization_config = {'quant_method': 'fp8'} + if fp8_quant_scope is not None: + model.quantization_config['fp8_quant_scope'] = fp8_quant_scope + model._load_buffers = {} + prefix = 'model.layers.0.self_attn' + params = { + f'{prefix}.kc.weight': nn.Parameter(torch.empty(1, 2, input_dim, dtype=torch.bfloat16)), + f'{prefix}.vc.weight': nn.Parameter(torch.empty(1, input_dim, 2, dtype=torch.bfloat16)), + } + return model, prefix, params + + +def _split_kv_b_weight(weight): + weight = weight.unflatten(0, (-1, 4)) + weight_kc, weight_vc = weight.split([2, 2], dim=1) + return weight_kc, weight_vc.transpose(1, 2).contiguous() + + +def test_glm_moe_only_loads_kv_b_proj_as_bf16(): + model, prefix, params = _make_kv_b_loader_model(fp8_quant_scope='moe_only') + loaded_weight = torch.arange(12, dtype=torch.bfloat16).reshape(4, 3) + + model._load_weight_attention( + name=f'{prefix}.kv_b_proj.weight', + loaded_weight=loaded_weight, + params_dict=params, + update_pe_mapping=[], + ) + + expected_kc, expected_vc = _split_kv_b_weight(loaded_weight) + assert torch.equal(params[f'{prefix}.kc.weight'], expected_kc) + assert torch.equal(params[f'{prefix}.vc.weight'], expected_vc) + assert model._load_buffers == {} + + +def test_global_fp8_loads_kv_b_proj_after_weight_scale(): + model, prefix, params = _make_kv_b_loader_model(input_dim=4) + loaded_weight = torch.arange(16, dtype=torch.float32).reshape(4, 4).to(torch.float8_e4m3fn) + loaded_scale = torch.full((1, 1), 0.5, dtype=torch.float32) + + model._load_weight_attention( + name=f'{prefix}.kv_b_proj.weight', + loaded_weight=loaded_weight, + params_dict=params, + update_pe_mapping=[], + ) + assert f'{prefix}.kv_b_proj.weight' in model._load_buffers + + model._load_weight_attention( + name=f'{prefix}.kv_b_proj.weight_scale_inv', + loaded_weight=loaded_scale, + params_dict=params, + update_pe_mapping=[], + ) + + dequantized_weight = (loaded_weight.float() * loaded_scale).to(torch.bfloat16) + expected_kc, expected_vc = _split_kv_b_weight(dequantized_weight) + assert torch.equal(params[f'{prefix}.kc.weight'], expected_kc) + assert torch.equal(params[f'{prefix}.vc.weight'], expected_vc) + assert model._load_buffers == {} + + def test_default_interleaved_rotary_uses_pairwise_layout_with_half_split_freqs(): query = torch.tensor([[[[1.0, 2.0, 3.0, 4.0]]]]) key = torch.tensor([[[[5.0, 6.0, 7.0, 8.0]]]]) From 97832a81cd39936873f923ec3ceff2414560d69b Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 16 Jul 2026 15:46:04 +0800 Subject: [PATCH 12/33] perf: fuse dsa indexer preparation --- lmdeploy/archs.py | 10 +- lmdeploy/hf_configs/__init__.py | 36 +++ .../hf_configs/configuration_deepseek_v32.py | 13 + lmdeploy/pytorch/backends/cuda/nsa.py | 79 +++-- lmdeploy/pytorch/backends/nsa.py | 7 + .../pytorch/configurations/deepseek_v32.py | 75 +---- .../pytorch/configurations/glm_moe_dsa.py | 38 +++ lmdeploy/pytorch/envs.py | 1 + lmdeploy/pytorch/kernels/cuda/dsa_indexer.py | 297 ++++++++++++++++++ lmdeploy/pytorch/models/deepseek_v32.py | 116 ++++++- lmdeploy/pytorch/nn/nsa.py | 58 +++- lmdeploy/pytorch/transformers/__init__.py | 38 +-- .../configuration_deepseek_v32.py | 13 +- .../transformers/configuration_glm_moe_dsa.py | 40 --- .../pytorch/config/test_glm_moe_dsa_config.py | 131 ++++---- tests/pytorch/kernel/test_dsa_indexer.py | 127 ++++++++ tests/pytorch/models/test_deepseek_v32.py | 71 +++-- 17 files changed, 835 insertions(+), 315 deletions(-) create mode 100644 lmdeploy/hf_configs/__init__.py create mode 100644 lmdeploy/hf_configs/configuration_deepseek_v32.py create mode 100644 lmdeploy/pytorch/configurations/glm_moe_dsa.py create mode 100644 lmdeploy/pytorch/kernels/cuda/dsa_indexer.py delete mode 100644 lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py create mode 100644 tests/pytorch/kernel/test_dsa_indexer.py diff --git a/lmdeploy/archs.py b/lmdeploy/archs.py index ef0d6fc120..e27562f835 100644 --- a/lmdeploy/archs.py +++ b/lmdeploy/archs.py @@ -1,8 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. from typing import Literal -from transformers import AutoConfig - from .messages import PytorchEngineConfig, TurbomindEngineConfig from .utils import get_logger @@ -146,11 +144,9 @@ def get_model_arch(model_path: str, trust_remote_code: bool = False): Args: model_path(str): the model path """ - try: - cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code) - except Exception as e: # noqa - from transformers import PretrainedConfig - cfg = PretrainedConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code) + from lmdeploy.hf_configs import config_from_pretrained + + cfg = config_from_pretrained(model_path, trust_remote_code=trust_remote_code) _cfg = cfg.to_dict() if _cfg.get('architectures', None): diff --git a/lmdeploy/hf_configs/__init__.py b/lmdeploy/hf_configs/__init__.py new file mode 100644 index 0000000000..367cacf7a6 --- /dev/null +++ b/lmdeploy/hf_configs/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from functools import lru_cache + +from transformers import AutoConfig + +from lmdeploy.utils import get_logger + +logger = get_logger('lmdeploy') + + +@lru_cache +def register_config(model_type: str): + if model_type == 'deepseek_v32': + from .configuration_deepseek_v32 import DeepseekV32Config + AutoConfig.register(DeepseekV32Config.model_type, DeepseekV32Config) + else: + logger.debug(f'Can not register config for model_type: {model_type}') + + +def config_from_pretrained(pretrained_model_name_or_path: str, **kwargs): + try: + return AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + except Exception as e: + logger.debug(f'AutoConfig.from_pretrained failed: {e}, try register config manually.') + # some models do not provide auto map for config + from transformers import PretrainedConfig + trust_remote_code = kwargs.pop('trust_remote_code', None) + config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **kwargs) + model_type = config_dict.get('model_type', None) + if trust_remote_code is not None: + kwargs['trust_remote_code'] = trust_remote_code + register_config(model_type) + try: + return AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + except Exception: + return PretrainedConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) diff --git a/lmdeploy/hf_configs/configuration_deepseek_v32.py b/lmdeploy/hf_configs/configuration_deepseek_v32.py new file mode 100644 index 0000000000..59f91fb3ef --- /dev/null +++ b/lmdeploy/hf_configs/configuration_deepseek_v32.py @@ -0,0 +1,13 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + + +class DeepseekV32Config(DeepseekV3Config): + model_type = 'deepseek_v32' + + def __init__(self, index_head_dim=128, index_n_heads=64, index_topk=2048, **kwargs): + super().__init__(**kwargs) + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.index_topk = index_topk diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index 1ab8211f85..1343edff01 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -6,6 +6,7 @@ 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.dsa_indexer 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 @@ -33,11 +34,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 @@ -45,26 +42,6 @@ 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 - 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) - q = q.reshape(*q_shape) - q_s = q_s.reshape(weights.shape) - q_s = q_s * self.softmax_scale * weights - - fill_kv_cache_blocked_fp8(k[:, None], - None, - k_cache[..., None, :], - 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, - group_size=self.block_size, - scale_fmt=self.scale_fmt) - scores = fp8_index(q, q_s, k_cache, @@ -86,6 +63,58 @@ def forward(self, q: Tensor, k: Tensor, weights: Tensor, k_cache: Tensor, k_s_ca 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) + q = q.reshape(*q_shape) + q_s = q_s.reshape(weights.shape) + q_s = q_s * self.softmax_scale * weights + + fill_kv_cache_blocked_fp8(k[:, None], + None, + k_cache[..., None, :], + None, + k_s_cache[..., None, :], + None, + 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) + + 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) + class TritonNSAIndexFP8Builder(BaseNSAIndexFP8Builder): diff --git a/lmdeploy/pytorch/backends/nsa.py b/lmdeploy/pytorch/backends/nsa.py index c20c80868c..f8ba09d5c7 100644 --- a/lmdeploy/pytorch/backends/nsa.py +++ b/lmdeploy/pytorch/backends/nsa.py @@ -24,6 +24,13 @@ 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.') + class BaseNSAIndexFP8Builder: diff --git a/lmdeploy/pytorch/configurations/deepseek_v32.py b/lmdeploy/pytorch/configurations/deepseek_v32.py index d66c28677e..888d971883 100644 --- a/lmdeploy/pytorch/configurations/deepseek_v32.py +++ b/lmdeploy/pytorch/configurations/deepseek_v32.py @@ -1,68 +1,24 @@ # Copyright (c) OpenMMLab. All rights reserved. +import functools + import torch from lmdeploy.pytorch import envs as _envs -from lmdeploy.utils import get_logger from .deepseek_v2 import DeepseekV2ModelConfigBuilder -logger = get_logger('lmdeploy') - - -def normalize_glm_moe_dsa_config(hf_config): - """Normalize GLM-MoE-DSA config fields for DeepSeek-V3.2 reuse.""" - if hf_config.model_type != 'glm_moe_dsa': - return - - qk_head_dim = getattr(hf_config, 'qk_head_dim', None) - if qk_head_dim is not None and qk_head_dim != hf_config.qk_nope_head_dim + hf_config.qk_rope_head_dim: - qk_rope_head_dim = qk_head_dim - hf_config.qk_nope_head_dim - if qk_rope_head_dim <= 0: - raise ValueError('Invalid GLM-MoE-DSA qk head dimensions.') - hf_config.qk_rope_head_dim = qk_rope_head_dim - - hf_config.qk_head_dim = hf_config.qk_nope_head_dim + hf_config.qk_rope_head_dim - hf_config.head_dim = hf_config.qk_rope_head_dim - - indexer_types = getattr(hf_config, 'indexer_types', None) - if indexer_types is not None: - indexer_types = list(indexer_types) - else: - pattern = getattr(hf_config, 'index_topk_pattern', None) - if pattern is not None: - if isinstance(pattern, str): - pattern_map = {'F': 'full', 'S': 'shared'} - indexer_types = [pattern_map[item.upper()] for item in pattern] - else: - indexer_types = list(pattern) - else: - freq = max(getattr(hf_config, 'index_topk_freq', 1), 1) - offset = getattr(hf_config, 'index_skip_topk_offset', 2) - indexer_types = [ - 'full' if max(layer_idx - offset + 1, 0) % freq == 0 else 'shared' - for layer_idx in range(hf_config.num_hidden_layers) - ] - if len(indexer_types) != hf_config.num_hidden_layers: - raise ValueError('GLM-MoE-DSA indexer_types length must match num_hidden_layers.') - if any(item not in ('full', 'shared') for item in indexer_types): - raise ValueError('GLM-MoE-DSA indexer_types only supports "full" and "shared".') - if indexer_types[0] != 'full': - raise ValueError('The first GLM-MoE-DSA layer must be a full indexer layer.') - - hf_config.indexer_types = indexer_types - - -def _check_env_v32(device: str = 'cuda'): +def _check_env_v32(device: str = 'cuda', require_hadamard: bool = True): """Environment check.""" if device != 'cuda': return # check cuda - try: - import fast_hadamard_transform # noqa: F401 - except ImportError: - raise ImportError('Deepseek V3.2 requires .') + if require_hadamard: + try: + import fast_hadamard_transform # noqa: F401 + except ImportError: + raise ImportError('Deepseek V3.2 requires .') try: import flash_mla # noqa: F401 @@ -78,21 +34,11 @@ 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): """build.""" - 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 hf_config.model_type == 'glm_moe_dsa' and _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.') - - normalize_glm_moe_dsa_config(hf_config) config = DeepseekV2ModelConfigBuilder.build(hf_config, model_path=model_path, **kwargs) assert hf_config.use_flash_mla, 'DeepSeek-V3.2 requires flash_mla to be available.' @@ -101,5 +47,6 @@ def build(cls, hf_config, model_path: str | None = None, **kwargs): config.cache_shapes = [index_k_shape, index_k_scale_shape] config.use_mla_fp8_cache = True config.mla_index_topk = hf_config.index_topk - config.check_env_func = _check_env_v32 + config.check_env_func = functools.partial(_check_env_v32, + require_hadamard=_envs.disable_dsa_indexer_fusion) 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..970e4375b6 --- /dev/null +++ b/lmdeploy/pytorch/configurations/glm_moe_dsa.py @@ -0,0 +1,38 @@ +# 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') + + +def normalize_glm_moe_dsa_config(hf_config): + """Normalize GLM-MoE-DSA checkpoint fields for the PyTorch runtime.""" + 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 + + +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.""" + 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.') + + normalize_glm_moe_dsa_config(hf_config) + return super().build(hf_config, model_path=model_path, **kwargs) 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/dsa_indexer.py b/lmdeploy/pytorch/kernels/cuda/dsa_indexer.py new file mode 100644 index 0000000000..f2cce7b772 --- /dev/null +++ b/lmdeploy/pytorch/kernels/cuda/dsa_indexer.py @@ -0,0 +1,297 @@ +# Copyright (c) OpenMMLab. All rights reserved. +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/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index 8eb455357f..755fe17dca 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -6,6 +6,7 @@ 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, get_step_ctx_manager from lmdeploy.pytorch.nn import ( @@ -102,6 +103,16 @@ 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) + + class LayerNorm(nn.Module): """Layer Normalization.""" @@ -122,10 +133,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 @@ -144,20 +151,30 @@ 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(interleaved=self.rope_interleave) self.indexer_topk = IndexerTopKFP8(self.index_topk, self.softmax_scale, block_size=128, fill=-1) @@ -183,6 +200,25 @@ def forward(self, attn_metadata: Any = None): q = self.wq_b(qr) q = q.unflatten(-1, (-1, self.head_dim)) + 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 + 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_pe, q_nope = torch.split(q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1) k = self.wk(x) k = self.k_norm(k) @@ -636,6 +672,8 @@ def forward( 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, @@ -668,9 +706,53 @@ def forward( 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 self._load_fused_indexer_weight(name, loaded_weight, params_dict): + 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) + + def _load_fused_indexer_weight(self, name: str, loaded_weight: torch.Tensor, + params_dict: dict[str, nn.Parameter]) -> 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_name = f'{indexer_prefix}.wk_weights_proj.weight' + fused_param = params_dict.get(fused_name) + if fused_param is None: + return False + + # The fused output layout is [K, head gates]. + 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 + + # Dequantize K once at load time; the runtime projection stays fused. + buffer_key = f'{indexer_prefix}.wk' + buffer = self._load_buffers.setdefault(buffer_key, {}) + 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) + self._load_buffers.pop(buffer_key) + return True diff --git a/lmdeploy/pytorch/nn/nsa.py b/lmdeploy/pytorch/nn/nsa.py index d3944bf003..a21669a436 100644 --- a/lmdeploy/pytorch/nn/nsa.py +++ b/lmdeploy/pytorch/nn/nsa.py @@ -15,30 +15,64 @@ 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 _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) + max_q_seqlen = 1 if is_decoding else (attn_metadata.max_q_seqlen or q.size(0)) # we need to make max_kv_seqlen=max_allocated_cache_len to enable cudagraph 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) diff --git a/lmdeploy/pytorch/transformers/__init__.py b/lmdeploy/pytorch/transformers/__init__.py index 8b3e8d0598..96da46744f 100644 --- a/lmdeploy/pytorch/transformers/__init__.py +++ b/lmdeploy/pytorch/transformers/__init__.py @@ -1,38 +1,4 @@ # Copyright (c) OpenMMLab. All rights reserved. -from functools import lru_cache +from lmdeploy.hf_configs import config_from_pretrained, register_config -from transformers import AutoConfig - -from lmdeploy.utils import get_logger - - -@lru_cache -def register_config(model_type: str): - if model_type == 'deepseek_v32': - from lmdeploy.pytorch.transformers.configuration_deepseek_v32 import DeepseekV32Config - AutoConfig.register(DeepseekV32Config.model_type, DeepseekV32Config) - elif model_type == 'glm_moe_dsa': - from lmdeploy.pytorch.transformers.configuration_glm_moe_dsa import GlmMoeDsaConfig - AutoConfig.register(GlmMoeDsaConfig.model_type, GlmMoeDsaConfig, exist_ok=True) - else: - logger.debug(f'Can not register config for model_type: {model_type}') - - -logger = get_logger('lmdeploy') - - -def config_from_pretrained(pretrained_model_name_or_path: str, **kwargs): - try: - return AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) - except ValueError as e: - logger.debug(f'AutoConfig.from_pretrained failed: {e}, try register config manually.') - # some models (dsv32) does not provide auto map for config - from transformers import PretrainedConfig - trust_remote_code = kwargs.pop('trust_remote_code', None) - config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **kwargs) - model_type = config_dict.get('model_type', None) - if trust_remote_code is not None: - kwargs['trust_remote_code'] = trust_remote_code - register_config(model_type) - - return AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) +__all__ = ['config_from_pretrained', 'register_config'] diff --git a/lmdeploy/pytorch/transformers/configuration_deepseek_v32.py b/lmdeploy/pytorch/transformers/configuration_deepseek_v32.py index 59f91fb3ef..3dd41d5dff 100644 --- a/lmdeploy/pytorch/transformers/configuration_deepseek_v32.py +++ b/lmdeploy/pytorch/transformers/configuration_deepseek_v32.py @@ -1,13 +1,4 @@ # Copyright (c) OpenMMLab. All rights reserved. +from lmdeploy.hf_configs.configuration_deepseek_v32 import DeepseekV32Config -from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config - - -class DeepseekV32Config(DeepseekV3Config): - model_type = 'deepseek_v32' - - def __init__(self, index_head_dim=128, index_n_heads=64, index_topk=2048, **kwargs): - super().__init__(**kwargs) - self.index_head_dim = index_head_dim - self.index_n_heads = index_n_heads - self.index_topk = index_topk +__all__ = ['DeepseekV32Config'] diff --git a/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py b/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py deleted file mode 100644 index 798eb99ef8..0000000000 --- a/lmdeploy/pytorch/transformers/configuration_glm_moe_dsa.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. - -from lmdeploy.pytorch.transformers.configuration_deepseek_v32 import DeepseekV32Config - - -class GlmMoeDsaConfig(DeepseekV32Config): - model_type = 'glm_moe_dsa' - attribute_map = { - 'num_local_experts': 'n_routed_experts', - } - - def __init__( - self, - index_head_dim=128, - index_n_heads=32, - index_topk=2048, - index_topk_pattern=None, - index_topk_freq=1, - index_skip_topk_offset=2, - rope_interleave=True, - indexer_rope_interleave=True, - indexer_types=None, - **kwargs, - ): - super().__init__( - index_head_dim=index_head_dim, - index_n_heads=index_n_heads, - index_topk=index_topk, - **kwargs, - ) - self.index_topk_pattern = index_topk_pattern - self.index_topk_freq = index_topk_freq - self.index_skip_topk_offset = index_skip_topk_offset - self.rope_interleave = rope_interleave - self.indexer_rope_interleave = indexer_rope_interleave - self.indexer_types = indexer_types - - if hasattr(self, 'qk_nope_head_dim') and hasattr(self, 'qk_rope_head_dim'): - self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - self.head_dim = self.qk_rope_head_dim diff --git a/tests/pytorch/config/test_glm_moe_dsa_config.py b/tests/pytorch/config/test_glm_moe_dsa_config.py index c9ff3d6d89..49d1b2df8d 100644 --- a/tests/pytorch/config/test_glm_moe_dsa_config.py +++ b/tests/pytorch/config/test_glm_moe_dsa_config.py @@ -1,15 +1,48 @@ +import json from types import SimpleNamespace -import pytest -import torch - from lmdeploy.pytorch.config import ModelConfig -from lmdeploy.pytorch.configurations.deepseek_v32 import ( - DeepseekV32ModelConfigBuilder, +from lmdeploy.pytorch.configurations.glm_moe_dsa import ( + GlmMoeDsaModelConfigBuilder, normalize_glm_moe_dsa_config, ) +def test_get_model_arch_registers_deepseek_v32_config(tmp_path): + from lmdeploy.archs import get_model_arch + from lmdeploy.pytorch.transformers.configuration_deepseek_v32 import DeepseekV32Config + + config_path = tmp_path / 'config.json' + config_path.write_text( + json.dumps({ + 'architectures': ['DeepseekV32ForCausalLM'], + 'model_type': 'deepseek_v32', + })) + + arch, config = get_model_arch(str(tmp_path)) + + assert arch == 'DeepseekV32ForCausalLM' + assert isinstance(config, DeepseekV32Config) + + +def test_get_model_arch_loads_native_glm_moe_dsa_config(tmp_path): + from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig + + from lmdeploy.archs import get_model_arch + + config_path = tmp_path / 'config.json' + config_path.write_text( + json.dumps({ + 'architectures': ['GlmMoeDsaForCausalLM'], + 'model_type': 'glm_moe_dsa', + })) + + arch, config = get_model_arch(str(tmp_path)) + + assert arch == 'GlmMoeDsaForCausalLM' + assert isinstance(config, GlmMoeDsaConfig) + + def _make_config(**kwargs): values = dict( model_type='glm_moe_dsa', @@ -18,6 +51,8 @@ def _make_config(**kwargs): qk_rope_head_dim=64, ) values.update(kwargs) + values.setdefault('qk_head_dim', values['qk_nope_head_dim'] + values['qk_rope_head_dim']) + values.setdefault('indexer_types', ['full'] * values['num_hidden_layers']) return SimpleNamespace(**values) @@ -40,10 +75,11 @@ def _make_fp8_build_config(**quantization_config): def test_glm_moe_dsa_enables_online_fp8_moe_only_scope(monkeypatch): + from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig + from lmdeploy.pytorch import envs from lmdeploy.pytorch import transformers as pytorch_transformers from lmdeploy.pytorch.configurations import deepseek_v2 - from lmdeploy.pytorch.transformers.configuration_glm_moe_dsa import GlmMoeDsaConfig monkeypatch.setattr(envs, 'fp8_moe_only', True) monkeypatch.setattr(deepseek_v2, 'flash_mla_available', lambda: True) @@ -60,6 +96,7 @@ def test_glm_moe_dsa_enables_online_fp8_moe_only_scope(monkeypatch): vocab_size=32, bos_token_id=1, eos_token_id=2, + dtype='float16', index_head_dim=4, index_n_heads=2, index_topk=2) @@ -70,19 +107,6 @@ def test_glm_moe_dsa_enables_online_fp8_moe_only_scope(monkeypatch): assert model_config.quant_config.quant_method == 'fp8' assert model_config.quant_config.fp8_quant_scope == 'moe_only' assert model_config.quant_config.hf_quant_config['lmdeploy_patched'] - assert model_config.dtype == torch.bfloat16 - - -def test_glm_moe_dsa_fp8_moe_only_scope_requires_env(monkeypatch): - from lmdeploy.pytorch import envs - - monkeypatch.setattr(envs, 'fp8_moe_only', False) - _patch_v32_base_builder(monkeypatch) - cfg = _make_fp8_build_config(quant_method='fp8', lmdeploy_patched=True) - - DeepseekV32ModelConfigBuilder.build(cfg) - - assert 'fp8_quant_scope' not in cfg.quantization_config def test_glm_moe_dsa_does_not_override_prequantized_fp8_scope(monkeypatch): @@ -92,52 +116,18 @@ def test_glm_moe_dsa_does_not_override_prequantized_fp8_scope(monkeypatch): _patch_v32_base_builder(monkeypatch) cfg = _make_fp8_build_config(quant_method='fp8') - DeepseekV32ModelConfigBuilder.build(cfg) + GlmMoeDsaModelConfigBuilder.build(cfg) assert 'fp8_quant_scope' not in cfg.quantization_config -def test_glm_moe_dsa_normalizes_pattern_indexer_types(): - cfg = _make_config(index_topk_pattern='FSSF') +def test_glm_moe_dsa_runtime_normalizes_native_hf_config(): + from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig - normalize_glm_moe_dsa_config(cfg) - - assert cfg.qk_head_dim == 192 - assert cfg.head_dim == 64 - assert cfg.indexer_types == ['full', 'shared', 'shared', 'full'] - - -def test_glm_moe_dsa_normalizes_freq_offset_indexer_types(): - cfg = _make_config(num_hidden_layers=6, index_topk_freq=2, index_skip_topk_offset=2) - - normalize_glm_moe_dsa_config(cfg) - - assert cfg.indexer_types == ['full', 'full', 'shared', 'full', 'shared', 'full'] - - -def test_glm_moe_dsa_recovers_corrupted_rope_head_dim(): - cfg = _make_config(qk_nope_head_dim=192, qk_rope_head_dim=192, qk_head_dim=256) - - normalize_glm_moe_dsa_config(cfg) - - assert cfg.qk_nope_head_dim == 192 - assert cfg.qk_rope_head_dim == 64 - assert cfg.qk_head_dim == 256 - assert cfg.head_dim == 64 - - -def test_glm_moe_dsa_rejects_shared_first_layer(): - cfg = _make_config(index_topk_pattern='SFFF') - - with pytest.raises(ValueError, match='first GLM-MoE-DSA layer'): - normalize_glm_moe_dsa_config(cfg) - - -def test_glm_moe_dsa_fallback_config_keeps_rope_head_dim(): - from lmdeploy.pytorch.transformers.configuration_glm_moe_dsa import GlmMoeDsaConfig - - cfg = GlmMoeDsaConfig(qk_nope_head_dim=128, + cfg = GlmMoeDsaConfig(qk_nope_head_dim=192, qk_rope_head_dim=64, + qk_head_dim=256, + head_dim=192, hidden_size=4096, intermediate_size=8192, num_attention_heads=32, @@ -146,37 +136,30 @@ def test_glm_moe_dsa_fallback_config_keeps_rope_head_dim(): bos_token_id=1, eos_token_id=2) - assert cfg.qk_nope_head_dim == 128 - assert cfg.qk_rope_head_dim == 64 - assert cfg.qk_head_dim == 192 + assert cfg.head_dim == 192 + + normalize_glm_moe_dsa_config(cfg) + assert cfg.head_dim == 64 - assert cfg.index_n_heads == 32 - assert cfg.rope_interleave is True - assert cfg.indexer_rope_interleave is True + assert cfg.qk_rope_head_dim == 64 -def test_glm_moe_dsa_builder_normalizes_mla_kv_heads(monkeypatch): +def test_glm_moe_dsa_builder_creates_sparse_mla_config(monkeypatch): from lmdeploy.pytorch.configurations import deepseek_v2 monkeypatch.setattr(deepseek_v2, 'flash_mla_available', lambda: True) cfg = _make_config(num_hidden_layers=2, index_head_dim=128, - index_n_heads=32, index_topk=2048, - index_skip_topk_offset=3, - index_topk_freq=4, - indexer_rope_interleave=True, - rope_interleave=True, hidden_size=6144, kv_lora_rank=512, num_attention_heads=64, num_key_value_heads=64, bos_token_id=None, eos_token_id=[154820, 154827, 154829], - vocab_size=154880, - architectures=['GlmMoeDsaForCausalLM']) + vocab_size=154880) - model_config = DeepseekV32ModelConfigBuilder.build(cfg, tp=1) + model_config = GlmMoeDsaModelConfigBuilder.build(cfg, tp=1) assert cfg.num_key_value_heads == 1 assert model_config.num_key_value_heads == 1 diff --git a/tests/pytorch/kernel/test_dsa_indexer.py b/tests/pytorch/kernel/test_dsa_indexer.py new file mode 100644 index 0000000000..176a306519 --- /dev/null +++ b/tests/pytorch/kernel/test_dsa_indexer.py @@ -0,0 +1,127 @@ +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.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 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) + q_pe, _ = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, interleaved=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.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 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/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index fe88b63afe..337ccdf148 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -11,6 +11,7 @@ DeepseekV32ForCausalLM, DSATopKIndicesBuffer, Indexer, + _dequantize_blocked_fp8, get_layer_indexer_type, ) @@ -111,21 +112,49 @@ def test_global_fp8_loads_kv_b_proj_after_weight_scale(): assert model._load_buffers == {} -def test_default_interleaved_rotary_uses_pairwise_layout_with_half_split_freqs(): - query = torch.tensor([[[[1.0, 2.0, 3.0, 4.0]]]]) - key = torch.tensor([[[[5.0, 6.0, 7.0, 8.0]]]]) - cos = torch.tensor([[10.0, 20.0, 10.0, 20.0]]) - sin = torch.tensor([[1.0, 2.0, 1.0, 2.0]]) +@pytest.mark.parametrize('n_heads', [32, 64]) +@pytest.mark.parametrize('load_scale_first', [False, True]) +def test_fp8_indexer_wk_loads_into_fused_bf16_projection(load_scale_first, n_heads): + model = DeepseekV32ForCausalLM.__new__(DeepseekV32ForCausalLM) + nn.Module.__init__(model) + model._load_buffers = {} + prefix = 'model.layers.0.self_attn.indexer' + fused = nn.Parameter(torch.zeros(128 + n_heads, 256, dtype=torch.bfloat16)) + params = {f'{prefix}.wk_weights_proj.weight': fused} + + wk = (torch.arange(128 * 256).reshape(128, 256) % 31 - 15).float().to(torch.float8_e4m3fn) + scale = torch.tensor([[0.25, 0.5]], dtype=torch.float32) + gate = torch.arange(n_heads * 256, dtype=torch.float32).reshape(n_heads, 256).to(torch.bfloat16) + wk_tensors = [ + (f'{prefix}.wk.weight', wk), + (f'{prefix}.wk.weight_scale_inv', scale), + ] + if load_scale_first: + wk_tensors.reverse() + + for name, tensor in wk_tensors: + model._load_weight_attention(name, tensor, params, update_pe_mapping=[]) + model._load_weight_attention(f'{prefix}.weights_proj.weight', gate, params, update_pe_mapping=[]) + + expected_wk = _dequantize_blocked_fp8(wk, scale, torch.bfloat16) + assert torch.equal(fused[:128], expected_wk) + assert torch.equal(fused[128:], gate) + assert model._load_buffers == {} - rotary = DefaultApplyRotaryEmbImpl(interleaved=True) - query_out, key_out = rotary.forward(query, key, cos, sin, inplace=False) - interleaved_cos = torch.tensor([[[10.0, 10.0, 20.0, 20.0]]]) - interleaved_sin = torch.tensor([[[1.0, 1.0, 2.0, 2.0]]]) - expected_query = query * interleaved_cos + rotate_interleaved(query) * interleaved_sin - expected_key = key * interleaved_cos + rotate_interleaved(key) * interleaved_sin - assert torch.equal(query_out, expected_query) - assert torch.equal(key_out, expected_key) +def test_bf16_indexer_wk_loads_directly_into_fused_projection(): + model = DeepseekV32ForCausalLM.__new__(DeepseekV32ForCausalLM) + nn.Module.__init__(model) + model._load_buffers = {} + prefix = 'model.layers.0.self_attn.indexer' + fused = nn.Parameter(torch.zeros(160, 8, dtype=torch.bfloat16)) + params = {f'{prefix}.wk_weights_proj.weight': fused} + wk = torch.arange(128 * 8, dtype=torch.float32).reshape(128, 8).to(torch.bfloat16) + + model._load_weight_attention(f'{prefix}.wk.weight', wk, params, update_pe_mapping=[]) + + assert torch.equal(fused[:128], wk) + assert model._load_buffers == {} def test_indexer_rope_interleave_uses_interleaved_layout(): @@ -148,22 +177,6 @@ def test_indexer_rope_interleave_uses_interleaved_layout(): assert torch.equal(key_out, expected_key) -def test_indexer_rope_interleave_can_fall_back_to_half_split_layout(): - indexer = Indexer.__new__(Indexer) - indexer.apply_rotary_pos_emb = lambda q, k, cos, sin, inplace=False: ('half-split', q, k, cos, sin, inplace) - - query = torch.zeros(1, 1, 4) - key = torch.zeros(1, 4) - cos = torch.zeros(1, 4) - sin = torch.zeros(1, 4) - - out = indexer._apply_rotary_pos_emb(query, key, (cos, sin)) - - assert out[0] == 'half-split' - assert out[2].shape == (1, 1, 4) - assert out[-1] is False - - def test_full_indexer_layer_writes_shared_topk_buffer(monkeypatch): attn = DeepseekV32Attention.__new__(DeepseekV32Attention) nn.Module.__init__(attn) From 138b6760a9932ed4b57bfaf95c05858a5052906f Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 17 Jul 2026 11:45:46 +0800 Subject: [PATCH 13/33] refactor: simplify deepseek v3.2 environment check --- lmdeploy/pytorch/configurations/deepseek_v32.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/lmdeploy/pytorch/configurations/deepseek_v32.py b/lmdeploy/pytorch/configurations/deepseek_v32.py index 888d971883..1f1f733ba6 100644 --- a/lmdeploy/pytorch/configurations/deepseek_v32.py +++ b/lmdeploy/pytorch/configurations/deepseek_v32.py @@ -1,6 +1,4 @@ # Copyright (c) OpenMMLab. All rights reserved. -import functools - import torch from lmdeploy.pytorch import envs as _envs @@ -8,13 +6,13 @@ from .deepseek_v2 import DeepseekV2ModelConfigBuilder -def _check_env_v32(device: str = 'cuda', require_hadamard: bool = True): +def _check_env_v32(device: str = 'cuda'): """Environment check.""" if device != 'cuda': return # check cuda - if require_hadamard: + if _envs.disable_dsa_indexer_fusion: try: import fast_hadamard_transform # noqa: F401 except ImportError: @@ -47,6 +45,5 @@ def build(cls, hf_config, model_path: str | None = None, **kwargs): config.cache_shapes = [index_k_shape, index_k_scale_shape] config.use_mla_fp8_cache = True config.mla_index_topk = hf_config.index_topk - config.check_env_func = functools.partial(_check_env_v32, - require_hadamard=_envs.disable_dsa_indexer_fusion) + config.check_env_func = _check_env_v32 return config From 4aa3a08e04768c65f8302149e98de8ee0c6f037e Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 17 Jul 2026 11:55:03 +0800 Subject: [PATCH 14/33] perf: compact dsa indexer replay --- lmdeploy/pytorch/messages.py | 2 +- lmdeploy/pytorch/models/deepseek_v32.py | 16 ++++++++++++--- lmdeploy/serve/anthropic/protocol.py | 3 ++- lmdeploy/serve/openai/protocol.py | 4 ++-- tests/pytorch/models/test_deepseek_v32.py | 25 +++++++++++++++-------- 5 files changed, 34 insertions(+), 16 deletions(-) diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py index dad987e2a3..d54ed777d1 100644 --- a/lmdeploy/pytorch/messages.py +++ b/lmdeploy/pytorch/messages.py @@ -592,7 +592,7 @@ def _get_pad_width(self, reserve_size: int): class HistoryIndexerTopK(_HistoryDataBase): - """History of sparse-attention indexer results.""" + """History of full-layer sparse-attention indexer results.""" ALLOC_SIZE = 1 def __init__(self, indexer_topk: np.ndarray = None, dtype: np.dtype = np.int32): diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index 755fe17dca..8b32c196d5 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -43,6 +43,11 @@ def get_layer_indexer_type(config: Any, layer_idx: int | None) -> str: return indexer_types[layer_idx] +def get_full_indexer_layer_ids(config: Any) -> tuple[int, ...]: + """Return the physical layers represented by compact indexer output.""" + return tuple(idx for idx in range(config.num_hidden_layers) if get_layer_indexer_type(config, idx) == 'full') + + 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.'): @@ -346,7 +351,9 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic self.indexer_type = get_layer_indexer_type(config, layer_idx) self.indexer = None + self.indexer_output_idx = None if self.indexer_type == 'full': + self.indexer_output_idx = get_full_indexer_layer_ids(config).index(layer_idx) 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): @@ -435,8 +442,9 @@ def forward( if topk_indices_buffer is not None: topk_indices = topk_indices_buffer.write(topk_indices) - if all_indexer_topk is not None: - all_indexer_topk[:, self.layer_idx, :].copy_(topk_indices) + # Shared layers reuse the previous result, so capture only full layers. + if all_indexer_topk is not None and self.indexer_output_idx is not None: + all_indexer_topk[:, self.indexer_output_idx, :].copy_(topk_indices) attn_output = self.attn_fwd( query_states, @@ -657,6 +665,7 @@ def __init__(self, bm_ctx = get_build_model_context() self.enable_return_routed_experts = bm_ctx.enable_return_routed_experts self.enable_return_indexer_topk = bm_ctx.enable_return_indexer_topk + self.num_indexer_layers = len(get_full_indexer_layer_ids(config)) def forward( self, @@ -681,8 +690,9 @@ def forward( ) all_indexer_topk = None if self.enable_return_indexer_topk: + # Axis 1 follows full indexer layers in physical layer order. all_indexer_topk = position_ids.new_empty( - (num_tokens, self.config.num_hidden_layers, self.config.index_topk), dtype=torch.int32) + (num_tokens, self.num_indexer_layers, self.config.index_topk), dtype=torch.int32) forward = self.model.forward_microbatch if step_ctx.enable_microbatch else self.model.forward hidden_states = forward( diff --git a/lmdeploy/serve/anthropic/protocol.py b/lmdeploy/serve/anthropic/protocol.py index 3d7a6b4f82..e3ac1a7a1c 100644 --- a/lmdeploy/serve/anthropic/protocol.py +++ b/lmdeploy/serve/anthropic/protocol.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field RoutedExperts = list[list[list[int]]] | str | None +# (num_token, num_full_indexer_layer, index_topk), ordered by physical layer. IndexerTopK = list[list[list[int]]] | str | None MessageStopReason = Literal['end_turn', 'max_tokens', 'stop_sequence', 'tool_use', 'parse_error'] @@ -133,7 +134,7 @@ class MessagesRequest(BaseModel): ) return_indexer_topk: bool | None = Field( default=False, - description=('Whether to return sparse-attention indexer top-k indices in the response.'), + description=('Whether to return sparse-attention top-k indices for full indexer layers.'), ) return_token_ids: bool | None = Field( default=False, diff --git a/lmdeploy/serve/openai/protocol.py b/lmdeploy/serve/openai/protocol.py index bdf39c9e3a..d0cb0b9cda 100644 --- a/lmdeploy/serve/openai/protocol.py +++ b/lmdeploy/serve/openai/protocol.py @@ -228,7 +228,7 @@ class ChatCompletionRequest(BaseModel): ) return_indexer_topk: bool | None = Field( default=False, - description=('Whether to return sparse-attention indexer top-k indices in the response.'), + description=('Whether to return sparse-attention top-k indices for full indexer layers.'), ) @@ -586,7 +586,7 @@ class GenerateReqMetaOutput(BaseModel): finish_reason: dict[str, Any] | None = None output_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) routed_experts: list[list[list[int]]] | str | None = None # (num_token, num_layer, topk_expert) - indexer_topk: list[list[list[int]]] | str | None = None # (num_token, num_layer, index_topk) + indexer_topk: list[list[list[int]]] | str | None = None # (num_token, num_full_indexer_layer, index_topk) # /generate output diff --git a/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index 337ccdf148..69516cf947 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -12,6 +12,7 @@ DSATopKIndicesBuffer, Indexer, _dequantize_blocked_fp8, + get_full_indexer_layer_ids, get_layer_indexer_type, ) @@ -177,10 +178,11 @@ def test_indexer_rope_interleave_uses_interleaved_layout(): assert torch.equal(key_out, expected_key) -def test_full_indexer_layer_writes_shared_topk_buffer(monkeypatch): +def test_full_indexer_layer_writes_compact_output_and_shared_buffer(monkeypatch): attn = DeepseekV32Attention.__new__(DeepseekV32Attention) nn.Module.__init__(attn) attn.layer_idx = 0 + attn.indexer_output_idx = 0 _patch_minimal_attention(attn, monkeypatch) computed_topk = torch.tensor([[4, 2, 1], [3, 2, 0]], dtype=torch.int32) attn.indexer = lambda *args, **kwargs: computed_topk @@ -192,7 +194,7 @@ def fake_attn_fwd(*args, **kwargs): attn.attn_fwd = fake_attn_fwd topk_buffer = DSATopKIndicesBuffer(topk=3) - all_indexer_topk = torch.full((2, 4, 3), -1, dtype=torch.int32) + all_indexer_topk = torch.full((2, 1, 3), -1, dtype=torch.int32) output = attn( hidden_states=torch.zeros(1, 2, 2), @@ -206,7 +208,6 @@ def fake_attn_fwd(*args, **kwargs): assert torch.equal(topk_buffer.indices[:2], computed_topk) assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:2].data_ptr() assert torch.equal(all_indexer_topk[:, 0], computed_topk) - assert torch.all(all_indexer_topk[:, 1:] == -1) @pytest.mark.parametrize('dp,attn_tp,expected_num_heads', [(1, 2, 2), (2, 2, 4)]) @@ -253,6 +254,7 @@ def test_shared_indexer_layer_reuses_shared_topk_buffer(monkeypatch): nn.Module.__init__(attn) attn.layer_idx = 3 attn.indexer = None + attn.indexer_output_idx = None _patch_minimal_attention(attn, monkeypatch) seen = {} @@ -264,7 +266,7 @@ def fake_attn_fwd(*args, **kwargs): topk_buffer = DSATopKIndicesBuffer(topk=3) shared_topk = torch.tensor([[5, 4, 3]], dtype=torch.int32) topk_buffer.write(shared_topk) - all_indexer_topk = torch.full((1, 4, 3), -1, dtype=torch.int32) + all_indexer_topk = torch.full((1, 1, 3), -1, dtype=torch.int32) output = attn( hidden_states=torch.zeros(1, 1, 2), @@ -276,8 +278,7 @@ def fake_attn_fwd(*args, **kwargs): assert output.shape == (1, 1, 1) assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:1].data_ptr() - assert torch.equal(all_indexer_topk[:, 3], shared_topk) - assert torch.all(all_indexer_topk[:, :3] == -1) + assert torch.all(all_indexer_topk == -1) def test_shared_indexer_layer_requires_shared_topk_buffer(monkeypatch): @@ -311,12 +312,14 @@ def test_shared_indexer_layer_requires_shared_topk_buffer(monkeypatch): def test_layer_indexer_type_defaults_to_full_and_reads_shared_entries(): - config = SimpleNamespace(indexer_types=['full', 'shared']) + config = SimpleNamespace(num_hidden_layers=3, indexer_types=['full', 'shared']) assert get_layer_indexer_type(config, 0) == 'full' assert get_layer_indexer_type(config, 1) == 'shared' assert get_layer_indexer_type(config, 2) == 'full' assert get_layer_indexer_type(SimpleNamespace(indexer_types=None), 1) == 'full' + assert get_full_indexer_layer_ids(config) == (0, 2) + assert get_full_indexer_layer_ids(SimpleNamespace(num_hidden_layers=3, indexer_types=None)) == (0, 1, 2) def test_moe_gate_captures_logical_experts_before_eplb_mapping(monkeypatch): @@ -401,9 +404,13 @@ def forward(self, input_ids, all_routed_experts=None, all_indexer_topk=None, **k ) model = DeepseekV32ForCausalLM.__new__(DeepseekV32ForCausalLM) nn.Module.__init__(model) - model.config = SimpleNamespace(num_hidden_layers=5, num_experts_per_tok=8, index_topk=3) + model.config = SimpleNamespace(num_hidden_layers=5, + num_experts_per_tok=8, + index_topk=3, + indexer_types=['full', 'full', 'full', 'shared', 'shared']) model.enable_return_routed_experts = True model.enable_return_indexer_topk = True + model.num_indexer_layers = 3 model.model = FakeModel() outputs = model( @@ -418,5 +425,5 @@ def forward(self, input_ids, all_routed_experts=None, all_indexer_topk=None, **k assert torch.all(routed_experts[:, :3] == torch.iinfo(torch.uint16).max) assert torch.all(routed_experts[:, 3] == 7) assert torch.all(routed_experts[:, 4] == 9) - assert outputs['all_indexer_topk'].shape == (2, 5, 3) + assert outputs['all_indexer_topk'].shape == (2, 3, 3) assert torch.all(outputs['all_indexer_topk'] == 5) From 6c1ec355b0d0856294960c9ac557aef37263b16f Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 21 Jul 2026 11:23:55 +0800 Subject: [PATCH 15/33] feat: support GLM-5.2 MTP --- .../pytorch/backends/cuda/attention/mla.py | 19 +- .../pytorch/backends/cuda/graph_runner.py | 5 +- lmdeploy/pytorch/backends/cuda/nsa.py | 12 + .../pytorch/configurations/glm_moe_dsa.py | 6 +- lmdeploy/pytorch/engine/config_builder.py | 4 + lmdeploy/pytorch/models/deepseek_v32.py | 121 ++++---- lmdeploy/pytorch/models/glm_moe_dsa_mtp.py | 262 ++++++++++++++++++ lmdeploy/pytorch/models/module_map.py | 1 + lmdeploy/pytorch/spec_decode/base.py | 4 +- .../spec_decode/proposers/deepseek_mtp.py | 31 ++- .../pytorch/config/test_glm_moe_dsa_config.py | 41 +++ tests/pytorch/kernel/test_mla_attention.py | 42 +++ tests/pytorch/models/test_deepseek_v32.py | 33 ++- .../spec_decode/test_cudagraph_strategy.py | 31 +++ .../spec_decode/test_glm_moe_dsa_mtp.py | 162 +++++++++++ tests/pytorch/spec_decode/test_spec_agent.py | 13 +- 16 files changed, 704 insertions(+), 83 deletions(-) create mode 100644 lmdeploy/pytorch/models/glm_moe_dsa_mtp.py create mode 100644 tests/pytorch/kernel/test_mla_attention.py create mode 100644 tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index 2fc4fdb298..2bcda8aec6 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -38,24 +38,31 @@ def __init__(self): self._update_decode_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) + 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_prefill_impl(self, nsa_indices: torch.Tensor, q_seqlens: torch.Tensor, cu_seqlens_k: torch.Tensor): """Update for prefill impl.""" @@ -176,7 +183,7 @@ def flash_mla_decoding( # 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 attn_output, _ = self.flash_mla_with_kvcache(query, diff --git a/lmdeploy/pytorch/backends/cuda/graph_runner.py b/lmdeploy/pytorch/backends/cuda/graph_runner.py index 924f301c0f..6c9da66791 100644 --- a/lmdeploy/pytorch/backends/cuda/graph_runner.py +++ b/lmdeploy/pytorch/backends/cuda/graph_runner.py @@ -240,7 +240,10 @@ 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'], ) + 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 1343edff01..c1b47e533a 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -1,6 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import functools +import torch from torch import Tensor from lmdeploy.pytorch.kernels.cuda.bitonic_topk import bitonic_topk @@ -23,6 +24,15 @@ def _get_sparse_index_topk(topk: int): return None +def _get_causal_k_seqlens(cu_seqlen_q: Tensor, q_seqlens: Tensor, k_seqlens: Tensor, + num_tokens: int) -> Tensor: + """Expand final KV lengths into the causal length of each query.""" + q_start = torch.repeat_interleave(cu_seqlen_q[:-1], q_seqlens, output_size=num_tokens) + history_lengths = torch.repeat_interleave(k_seqlens - q_seqlens, q_seqlens, output_size=num_tokens) + 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: @@ -52,6 +62,8 @@ def _forward_index(self, q: Tensor, q_s: Tensor, k_cache: Tensor, k_s_cache: Ten max_q_seqlen=max_q_seqlen, max_k_seqlen=max_kv_seqlen, causal=True) + 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, diff --git a/lmdeploy/pytorch/configurations/glm_moe_dsa.py b/lmdeploy/pytorch/configurations/glm_moe_dsa.py index 970e4375b6..c569399e55 100644 --- a/lmdeploy/pytorch/configurations/glm_moe_dsa.py +++ b/lmdeploy/pytorch/configurations/glm_moe_dsa.py @@ -25,6 +25,7 @@ def condition(cls, hf_config): @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' @@ -35,4 +36,7 @@ def build(cls, hf_config, model_path: str | None = None, **kwargs): 'and the FP8 quantization config is LMDeploy-synthesized.') normalize_glm_moe_dsa_config(hf_config) - return super().build(hf_config, model_path=model_path, **kwargs) + 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/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/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index 8b32c196d5..590e3a0da3 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -101,6 +101,13 @@ def read(self, num_tokens: int, device: torch.device) -> torch.Tensor: 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: + """Move selected query rows to the front 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: assert x.dtype == torch.bfloat16 from fast_hadamard_transform import hadamard_transform @@ -118,6 +125,45 @@ def _dequantize_blocked_fp8(weight: torch.Tensor, scale: torch.Tensor, dtype: to 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.""" @@ -353,7 +399,9 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic self.indexer = None self.indexer_output_idx = None if self.indexer_type == 'full': - self.indexer_output_idx = get_full_indexer_layer_ids(config).index(layer_idx) + full_layer_ids = get_full_indexer_layer_ids(config) + if layer_idx in full_layer_ids: + self.indexer_output_idx = full_layer_ids.index(layer_idx) 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): @@ -404,6 +452,7 @@ def forward( past_key_value: Sequence[torch.Tensor] = None, attn_metadata: Any = None, topk_indices_buffer: DSATopKIndicesBuffer | None = None, + skip_topk: bool = False, all_indexer_topk: torch.Tensor | None = None, ): """Rewrite of LlamaAttention.forward.""" @@ -429,18 +478,19 @@ def forward( query_states[..., nope_size:] = q_pe key_states[..., nope_size:] = k_pe - if self.indexer is None: - if topk_indices_buffer is None: - raise RuntimeError(f'Layer {self.layer_idx} reuses DSA top-k indices but none were provided.') - topk_indices = topk_indices_buffer.read(q_len, hidden_states.device) + 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)) else: - topk_indices = self.indexer(hidden_states, - qr, - rotary_pos_emb, - past_key_value[-2:], - attn_metadata=attn_metadata) - if topk_indices_buffer is not None: - topk_indices = topk_indices_buffer.write(topk_indices) + topk_indices = topk_indices_buffer.read(q_len, hidden_states.device) # Shared layers reuse the previous result, so capture only full layers. if all_indexer_topk is not None and self.indexer_output_idx is not None: @@ -507,6 +557,7 @@ def forward( residual: torch.Tensor | None = None, attn_metadata: Any = None, topk_indices_buffer: DSATopKIndicesBuffer | None = None, + skip_topk: bool = False, all_routed_experts: torch.Tensor | None = None, all_indexer_topk: torch.Tensor | None = None, ) -> tuple[torch.FloatTensor, torch.FloatTensor]: @@ -524,6 +575,7 @@ def forward( past_key_value=past_key_value, attn_metadata=attn_metadata, topk_indices_buffer=topk_indices_buffer, + skip_topk=skip_topk, all_indexer_topk=all_indexer_topk, ) else: @@ -716,7 +768,8 @@ def forward( 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 self._load_fused_indexer_weight(name, loaded_weight, params_dict): + 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) @@ -724,45 +777,3 @@ def _load_weight_attention(self, name: str, loaded_weight: torch.Tensor, params_ if get_layer_indexer_type(self.config, layer_idx) == 'shared': return return super()._load_weight_attention(name, loaded_weight, params_dict, update_pe_mapping) - - def _load_fused_indexer_weight(self, name: str, loaded_weight: torch.Tensor, - params_dict: dict[str, nn.Parameter]) -> 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_name = f'{indexer_prefix}.wk_weights_proj.weight' - fused_param = params_dict.get(fused_name) - if fused_param is None: - return False - - # The fused output layout is [K, head gates]. - 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 - - # Dequantize K once at load time; the runtime projection stays fused. - buffer_key = f'{indexer_prefix}.wk' - buffer = self._load_buffers.setdefault(buffer_key, {}) - 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) - self._load_buffers.pop(buffer_key) - return True 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..abfeb62ff2 --- /dev/null +++ b/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py @@ -0,0 +1,262 @@ +# 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, + ) -> 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, + ) + 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, + 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, + ) + + 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 set_topk_indices_buffer(self, topk_indices_buffer: DSATopKIndicesBuffer): + self.topk_indices_buffer = topk_indices_buffer + + 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, + 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, + 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)) + 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 d8ea847c3e..b5f31e7153 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/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..8157528e69 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 target-owned resources.""" + super().build_model(empty_init, target_model=target_model, build_model_ctx=build_model_ctx) + draft_model = self.model + if hasattr(draft_model, 'set_topk_indices_buffer'): + draft_model.set_input_embeddings(target_model.get_input_embeddings()) + draft_model.set_topk_indices_buffer(target_model.model.topk_indices_buffer) + 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'): + logits_hidden_states = draft_model.prepare_hidden_states_for_logits(hidden_states) + logits = self.target_model.get_logits(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/tests/pytorch/config/test_glm_moe_dsa_config.py b/tests/pytorch/config/test_glm_moe_dsa_config.py index 49d1b2df8d..e7292e140f 100644 --- a/tests/pytorch/config/test_glm_moe_dsa_config.py +++ b/tests/pytorch/config/test_glm_moe_dsa_config.py @@ -165,3 +165,44 @@ def test_glm_moe_dsa_builder_creates_sparse_mla_config(monkeypatch): assert model_config.num_key_value_heads == 1 assert model_config.head_dim == 576 assert model_config.use_mla_fp8_cache is True + + +def test_glm_moe_dsa_draft_selects_glm_mtp_model(monkeypatch): + _patch_v32_base_builder(monkeypatch) + cfg = _make_config(architectures=['GlmMoeDsaForCausalLM'], + use_flash_mla=True, + index_head_dim=128, + index_topk=2048, + quantization_config={}) + + GlmMoeDsaModelConfigBuilder.build(cfg, is_draft_model=True) + + assert cfg.architectures == ['GlmMoeDsaMTPModel'] + + +def test_glm_moe_dsa_mtp_reuses_target_dist_config(monkeypatch): + from lmdeploy.pytorch import transformers as pytorch_transformers + from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig + from lmdeploy.pytorch.engine.config_builder import ConfigBuilder + + monkeypatch.setattr(pytorch_transformers, 'config_from_pretrained', + lambda *args, **kwargs: SimpleNamespace(model_type='glm_moe_dsa')) + captured = {} + + def fake_from_config(**kwargs): + captured.update(kwargs) + return kwargs + + monkeypatch.setattr(SpecDecodeConfig, 'from_config', staticmethod(fake_from_config)) + target_dist_config = DistConfig(tp=8) + speculative_config = SimpleNamespace(method='deepseek_mtp', model=None, num_speculative_tokens=5) + engine_config = SimpleNamespace(dtype='auto', model_format=None, hf_overrides=None) + + ConfigBuilder.build_specdecode_config('glm-model', + speculative_config, + engine_config, + cache_config=object(), + dist_config=target_dist_config, + trust_remote_code=True) + + assert captured['dist_config'] is target_dist_config diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py new file mode 100644 index 0000000000..90a9f68f86 --- /dev/null +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -0,0 +1,42 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import torch + +from lmdeploy.pytorch.backends.cuda.attention.mla import NSAIndicesUpdater + + +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_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/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py index 69516cf947..9725aa8af2 100644 --- a/tests/pytorch/models/test_deepseek_v32.py +++ b/tests/pytorch/models/test_deepseek_v32.py @@ -210,6 +210,35 @@ def fake_attn_fwd(*args, **kwargs): assert torch.equal(all_indexer_topk[:, 0], computed_topk) +def test_mtp_skips_indexer_and_reads_compacted_topk(monkeypatch): + attn = DeepseekV32Attention.__new__(DeepseekV32Attention) + nn.Module.__init__(attn) + attn.layer_idx = 5 + attn.indexer_output_idx = None + _patch_minimal_attention(attn, monkeypatch) + attn.indexer = lambda *args, **kwargs: pytest.fail('recurrent MTP must reuse seed top-k') + seen = {} + + def fake_attn_fwd(*args, **kwargs): + seen['nsa_indices'] = kwargs['nsa_indices'] + return torch.zeros(args[0].size(0), 1, 1) + + attn.attn_fwd = fake_attn_fwd + topk_buffer = DSATopKIndicesBuffer(topk=3) + topk_buffer.write(torch.tensor([[1, 2, 3], [11, 12, 13], [21, 22, 23]], dtype=torch.int32)) + topk_buffer.compact(torch.tensor([2, 0])) + + attn( + hidden_states=torch.zeros(1, 2, 2), + rotary_pos_emb=(torch.zeros(2, 1), torch.zeros(2, 1)), + past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], + topk_indices_buffer=topk_buffer, + skip_topk=True, + ) + + assert torch.equal(seen['nsa_indices'], torch.tensor([[21, 22, 23], [1, 2, 3]], dtype=torch.int32)) + + @pytest.mark.parametrize('dp,attn_tp,expected_num_heads', [(1, 2, 2), (2, 2, 4)]) def test_attention_uses_dp_aware_num_heads(monkeypatch, dp, attn_tp, expected_num_heads): attn = DeepseekV32Attention.__new__(DeepseekV32Attention) @@ -281,7 +310,7 @@ def fake_attn_fwd(*args, **kwargs): assert torch.all(all_indexer_topk == -1) -def test_shared_indexer_layer_requires_shared_topk_buffer(monkeypatch): +def test_attention_requires_shared_topk_buffer(monkeypatch): attn = DeepseekV32Attention.__new__(DeepseekV32Attention) nn.Module.__init__(attn) attn.layer_idx = 3 @@ -303,7 +332,7 @@ def test_shared_indexer_layer_requires_shared_topk_buffer(monkeypatch): ) attn.apply_rotary_pos_emb = lambda q, k, cos, sin, inplace=False: (q, k) - with pytest.raises(RuntimeError, match='reuses DSA top-k indices'): + with pytest.raises(RuntimeError, match='requires a DSA top-k indices buffer'): attn( hidden_states=torch.zeros(1, 1, 2), rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), diff --git a/tests/pytorch/spec_decode/test_cudagraph_strategy.py b/tests/pytorch/spec_decode/test_cudagraph_strategy.py index c8d03dc46b..fdc9c546d1 100644 --- a/tests/pytorch/spec_decode/test_cudagraph_strategy.py +++ b/tests/pytorch/spec_decode/test_cudagraph_strategy.py @@ -132,3 +132,34 @@ 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) + reuse_key = runner.get_graph_key(**kwargs, skip_topk=True) + + assert seed_key != reuse_key + assert seed_key[-1] is False + assert reuse_key[-1] is 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..b16e874309 --- /dev/null +++ b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py @@ -0,0 +1,162 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import asyncio + +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 + + def compact_topk_indices(self, row_indices): + self.compacted = row_indices + + def prepare_hidden_states_for_logits(self, hidden_states): + 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_raw_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) + assert torch.equal(proposer.target_model.hidden_states, selected_hidden + 10) + 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_topk_buffer(monkeypatch): + + class _Draft(nn.Module): + + def __init__(self): + super().__init__() + self.embed_tokens = None + self.topk_indices_buffer = None + + def set_input_embeddings(self, embed_tokens): + self.embed_tokens = embed_tokens + + def get_input_embeddings(self): + return self.embed_tokens + + def set_topk_indices_buffer(self, topk_indices_buffer): + self.topk_indices_buffer = topk_indices_buffer + + 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() + 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 target.model.topk_indices_buffer + + +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/spec_decode/test_spec_agent.py b/tests/pytorch/spec_decode/test_spec_agent.py index 2dcef80bf4..7feda2c32a 100644 --- a/tests/pytorch/spec_decode/test_spec_agent.py +++ b/tests/pytorch/spec_decode/test_spec_agent.py @@ -275,30 +275,27 @@ def test_spec_model_agent_method_when_enabled(): assert agent.method == specdecode_config.method -def test_qwen35_mtp_reuses_main_dist_context(monkeypatch): - """Qwen3.5 MTP mirrors the target topology, so it should share groups.""" +def test_matching_draft_dist_config_reuses_main_context(monkeypatch): from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig from lmdeploy.pytorch.distributed import DistContext from lmdeploy.pytorch.spec_decode import base as base_mod dist_config = DistConfig(dp=2, ep=2) dist_ctx = DistContext(rank=1, dp_rank=1, dist_config=dist_config, ep_gpu_group=object()) - specdecode_config = SpecDecodeConfig(model='draft-model', - method='qwen3_5_mtp', + specdecode_config = SpecDecodeConfig(model='glm-model', + method='deepseek_mtp', dist_config=DistConfig(dp=2, ep=2), num_speculative_tokens=3) def fail_build(*args, **kwargs): - raise AssertionError('qwen3_5_mtp should not build a separate draft DistContext') + raise AssertionError('Matching topology should reuse the target DistContext') monkeypatch.setattr(base_mod.DistContext, 'build', staticmethod(fail_build)) assert base_mod._build_draft_dist_ctx(dist_ctx, specdecode_config) is dist_ctx -def test_non_qwen35_mtp_builds_draft_dist_context(monkeypatch): - """Other speculative methods keep their separate draft distribution - path.""" +def test_different_draft_dist_config_builds_draft_context(monkeypatch): from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig from lmdeploy.pytorch.distributed import DistContext from lmdeploy.pytorch.spec_decode import base as base_mod From 64fc81c34fc600417eefb947963199f8dda5f8ce Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 21 Jul 2026 12:31:00 +0800 Subject: [PATCH 16/33] fix: align GLM-5.2 MTP sharing semantics --- lmdeploy/pytorch/models/glm_moe_dsa_mtp.py | 3 --- .../spec_decode/proposers/deepseek_mtp.py | 9 ++++----- .../spec_decode/test_glm_moe_dsa_mtp.py | 20 +++++++++++-------- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py b/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py index abfeb62ff2..e2177c87b1 100644 --- a/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py +++ b/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py @@ -198,9 +198,6 @@ def set_input_embeddings(self, embed_tokens: nn.Module): def get_input_embeddings(self): return self.model.get_input_embeddings() - def set_topk_indices_buffer(self, topk_indices_buffer: DSATopKIndicesBuffer): - self.topk_indices_buffer = topk_indices_buffer - def compact_topk_indices(self, row_indices: torch.Tensor): self.topk_indices_buffer.compact(row_indices) diff --git a/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py b/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py index 8157528e69..ebf997a446 100644 --- a/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py +++ b/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py @@ -11,12 +11,11 @@ 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 target-owned resources.""" + """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 hasattr(draft_model, 'set_topk_indices_buffer'): + if hasattr(draft_model, 'uses_dsa_topk_buffer'): draft_model.set_input_embeddings(target_model.get_input_embeddings()) - draft_model.set_topk_indices_buffer(target_model.model.topk_indices_buffer) async def get_outputs(self, model_outputs: dict[str, torch.Tensor], @@ -36,8 +35,8 @@ async def get_outputs(self, draft_model.compact_topk_indices(last_token_loc) if hasattr(draft_model, 'prepare_hidden_states_for_logits'): - logits_hidden_states = draft_model.prepare_hidden_states_for_logits(hidden_states) - logits = self.target_model.get_logits(logits_hidden_states) + hidden_states = draft_model.prepare_hidden_states_for_logits(hidden_states) + logits = self.target_model.get_logits(hidden_states) else: logits = self.get_logits(hidden_states) target_hidden_states = hidden_states diff --git a/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py index b16e874309..c2711c6e9e 100644 --- a/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py +++ b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py @@ -30,11 +30,13 @@ class _DummyDraft(nn.Module): 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 @@ -47,7 +49,7 @@ def get_logits(self, hidden_states): return logits -def test_proposer_reuses_topk_and_recycles_raw_hidden_states(): +def test_proposer_reuses_topk_and_recycles_postnorm_hidden_states(): proposer = object.__new__(DeepseekMTP) proposer.model = _DummyDraft() proposer.target_model = _DummyTarget() @@ -61,21 +63,24 @@ def test_proposer_reuses_topk_and_recycles_raw_hidden_states(): selected_hidden = hidden_states[:, [2, 0]] assert draft_ids.tolist() == [[5], [5]] - assert torch.equal(recurrent_hidden, selected_hidden) + 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_topk_buffer(monkeypatch): +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 = None + self.topk_indices_buffer = DSATopKIndicesBuffer(topk=2) def set_input_embeddings(self, embed_tokens): self.embed_tokens = embed_tokens @@ -83,9 +88,6 @@ def set_input_embeddings(self, embed_tokens): def get_input_embeddings(self): return self.embed_tokens - def set_topk_indices_buffer(self, topk_indices_buffer): - self.topk_indices_buffer = topk_indices_buffer - class _Target(nn.Module): def __init__(self): @@ -98,6 +100,7 @@ def get_input_embeddings(self): return self.embedding draft = _Draft() + draft_topk_indices_buffer = draft.topk_indices_buffer target = _Target() proposer = object.__new__(DeepseekMTP) @@ -109,7 +112,8 @@ def build_model(self, empty_init, target_model=None, build_model_ctx=None): proposer.build_model(empty_init=True, target_model=target) assert draft.embed_tokens is target.embedding - assert draft.topk_indices_buffer is target.model.topk_indices_buffer + assert draft.topk_indices_buffer is draft_topk_indices_buffer + assert draft.topk_indices_buffer is not target.model.topk_indices_buffer def test_glm_mtp_prepares_postnorm_hidden_for_logits(): From 3c45c117e3af6edcd43cf360177fd60e023f7983 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 21 Jul 2026 16:54:37 +0800 Subject: [PATCH 17/33] feat: support bf16 sparse FlashMLA cache --- lmdeploy/cli/utils.py | 3 +- lmdeploy/messages.py | 4 +- .../pytorch/backends/cuda/attention/mla.py | 81 ++-- lmdeploy/pytorch/backends/cuda/nsa.py | 11 +- lmdeploy/pytorch/backends/cuda/op_backend.py | 15 +- lmdeploy/pytorch/config.py | 7 +- .../pytorch/configurations/deepseek_v32.py | 2 +- lmdeploy/pytorch/engine/cache_engine.py | 18 +- lmdeploy/pytorch/models/deepseek_v32.py | 2 +- lmdeploy/pytorch/models/utils/cudagraph.py | 10 +- .../pytorch/config/test_glm_moe_dsa_config.py | 13 +- tests/pytorch/engine/test_cache_engine.py | 76 ++- tests/pytorch/kernel/test_mla_attention.py | 38 +- tests/pytorch/models/test_deepseek_v32.py | 458 ------------------ 14 files changed, 223 insertions(+), 515 deletions(-) delete mode 100644 tests/pytorch/models/test_deepseek_v32.py diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 71579b9603..bcd8ffbda9 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 5549897128..44bbf5b25c 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 @@ -426,7 +426,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 2bcda8aec6..5cb87c66a4 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -44,6 +44,7 @@ def _update_decode_impl(self, nsa_indices: torch.Tensor, block_offsets: torch.Te 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, @@ -91,8 +92,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 the BF16 global 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 """ @@ -159,7 +161,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, @@ -201,32 +203,15 @@ def flash_mla_decoding( 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, flatten_k: torch.Tensor, + nsa_indices: torch.Tensor) -> torch.Tensor: + """Run sparse FlashMLA over a BF16 flat KV view.""" 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, @@ -237,6 +222,27 @@ 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 the global 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) + # sparse_fwd addresses a flat cache, so map paged positions to physical slots. + nsa_indices = self.nsa_updater.update_decode(nsa_indices, attn_metadata.block_offsets, max_q_seqlen, + block_size) + nsa_indices = nsa_indices.flatten(0, 1)[:, None] + flatten_k = k_cache.flatten(0, 1) + return self._flash_mla_sparse(query, flatten_k, nsa_indices) + def _prefill_triton( self, query: torch.Tensor, @@ -399,6 +405,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, @@ -463,8 +470,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. @@ -475,7 +482,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, @@ -490,7 +499,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 @@ -541,15 +550,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 the BF16 global cache - 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 @@ -567,12 +576,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/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index c1b47e533a..6b3d830067 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -26,9 +26,17 @@ def _get_sparse_index_topk(topk: int): def _get_causal_k_seqlens(cu_seqlen_q: Tensor, q_seqlens: Tensor, k_seqlens: Tensor, num_tokens: int) -> Tensor: - """Expand final KV lengths into the causal length of each query.""" + """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 @@ -62,6 +70,7 @@ def _forward_index(self, q: Tensor, q_s: Tensor, k_cache: Tensor, k_s_cache: Ten 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) diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index 5608bac35f..c2107b3086 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -117,11 +117,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, @@ -132,9 +140,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/config.py b/lmdeploy/pytorch/config.py index a1058b1adb..b6c7df080a 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -362,7 +362,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 @@ -393,6 +393,11 @@ class ModelConfig: # flags mark if this model use mrope use_mrope: bool = False + @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 1f1f733ba6..35eac1815f 100644 --- a/lmdeploy/pytorch/configurations/deepseek_v32.py +++ b/lmdeploy/pytorch/configurations/deepseek_v32.py @@ -43,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/engine/cache_engine.py b/lmdeploy/pytorch/engine/cache_engine.py index 6dd6ffb6d3..3e0dccfc02 100644 --- a/lmdeploy/pytorch/engine/cache_engine.py +++ b/lmdeploy/pytorch/engine/cache_engine.py @@ -50,9 +50,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, @@ -116,12 +129,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): @@ -439,6 +453,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/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index 590e3a0da3..f12b5a3d0e 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -102,7 +102,7 @@ def read(self, num_tokens: int, device: torch.device) -> torch.Tensor: return self.indices[:num_tokens] def compact(self, row_indices: torch.Tensor) -> torch.Tensor: - """Move selected query rows to the front for recurrent MTP reuse.""" + """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)] diff --git a/lmdeploy/pytorch/models/utils/cudagraph.py b/lmdeploy/pytorch/models/utils/cudagraph.py index 30006d090a..3ce0e4558d 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,7 +261,10 @@ 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: + # 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 step_ctx = get_step_ctx_manager().current_context() model_config = step_ctx.model_config diff --git a/tests/pytorch/config/test_glm_moe_dsa_config.py b/tests/pytorch/config/test_glm_moe_dsa_config.py index e7292e140f..2900ff7b14 100644 --- a/tests/pytorch/config/test_glm_moe_dsa_config.py +++ b/tests/pytorch/config/test_glm_moe_dsa_config.py @@ -2,6 +2,7 @@ from types import SimpleNamespace from lmdeploy.pytorch.config import ModelConfig +from lmdeploy.pytorch.configurations.deepseek_v32 import DeepseekV32ModelConfigBuilder from lmdeploy.pytorch.configurations.glm_moe_dsa import ( GlmMoeDsaModelConfigBuilder, normalize_glm_moe_dsa_config, @@ -164,7 +165,17 @@ def test_glm_moe_dsa_builder_creates_sparse_mla_config(monkeypatch): assert cfg.num_key_value_heads == 1 assert model_config.num_key_value_heads == 1 assert model_config.head_dim == 576 - assert model_config.use_mla_fp8_cache is True + assert model_config.mla_kv_cache_dtype == 'bfloat16' + assert model_config.use_mla_fp8_cache is False + + +def test_deepseek_v32_builder_defaults_to_bf16_sparse_mla(monkeypatch): + _patch_v32_base_builder(monkeypatch) + cfg = _make_config(model_type='deepseek_v32', use_flash_mla=True, index_head_dim=128, index_topk=2048) + + model_config = DeepseekV32ModelConfigBuilder.build(cfg) + + assert model_config.mla_kv_cache_dtype == 'bfloat16' def test_glm_moe_dsa_draft_selects_glm_mtp_model(monkeypatch): diff --git a/tests/pytorch/engine/test_cache_engine.py b/tests/pytorch/engine/test_cache_engine.py index ca601f5588..b14012a21c 100644 --- a/tests/pytorch/engine/test_cache_engine.py +++ b/tests/pytorch/engine/test_cache_engine.py @@ -1,16 +1,90 @@ # Copyright (c) OpenMMLab. All rights reserved. import asyncio +from types import SimpleNamespace import numpy as np import pytest import torch -from lmdeploy.pytorch.config import CacheConfig +from lmdeploy.messages import QuantPolicy +from lmdeploy.pytorch.config import CacheConfig, ModelConfig from lmdeploy.pytorch.disagg.conn.protocol import MigrationProtocol from lmdeploy.pytorch.disagg.messages import MigrationExecutionBatch from lmdeploy.pytorch.engine.cache_engine import CacheEngine, StateCacheEngine +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_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index 90a9f68f86..dea48a0d15 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -1,8 +1,12 @@ # 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 NSAIndicesUpdater +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(): @@ -28,6 +32,38 @@ def test_nsa_decode_indices_update_keeps_single_token_shape(): assert torch.equal(output, expected) +def test_bf16_sparse_decode_uses_global_cache_view(): + impl = object.__new__(FlashMLAImpl) + impl.nsa_updater = NSAIndicesUpdater() + impl.nsa_updater._update_decode_func = impl.nsa_updater._update_decode_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) + k_cache = torch.empty(3, 16, 1, 576, dtype=torch.bfloat16) + 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, flatten_k, global_indices = impl._flash_mla_sparse.call_args.args + assert sparse_query is query + assert flatten_k.shape == (48, 1, 576) + expected = torch.tensor([[[16, 33]], [[0, -1]], [[32, 17]], [[16, 33]]]) + assert torch.equal(global_indices, expected) + + +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') diff --git a/tests/pytorch/models/test_deepseek_v32.py b/tests/pytorch/models/test_deepseek_v32.py deleted file mode 100644 index 9725aa8af2..0000000000 --- a/tests/pytorch/models/test_deepseek_v32.py +++ /dev/null @@ -1,458 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch -from torch import nn - -from lmdeploy.pytorch.backends.default.apply_rotary_emb import DefaultApplyRotaryEmbImpl, rotate_interleaved -from lmdeploy.pytorch.models.deepseek_v2 import DeepseekV2BMM, DeepseekV2ForCausalLM, DeepseekV2MoE, MoEGate -from lmdeploy.pytorch.models.deepseek_v32 import ( - DeepseekV32Attention, - DeepseekV32ForCausalLM, - DSATopKIndicesBuffer, - Indexer, - _dequantize_blocked_fp8, - get_full_indexer_layer_ids, - get_layer_indexer_type, -) - - -def _patch_minimal_attention(attn, monkeypatch, dp=1, attn_tp=1, seen=None): - """Patch a DSA attention module down to the top-k routing contract.""" - - def fake_qkv_proj(hidden_states, num_heads): - if seen is not None: - seen['num_heads'] = num_heads - q_len = hidden_states.size(1) - return ( - torch.zeros(q_len, num_heads, 2), - torch.zeros(q_len, 1, 2), - torch.zeros(q_len, 1, 1), - torch.zeros(q_len, num_heads, 1), - torch.zeros(q_len, 1, 1), - torch.zeros(q_len, 1, 1), - ) - - dist_config = SimpleNamespace(dp=dp, attn_tp=attn_tp) - monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v32.get_dist_manager', - lambda: SimpleNamespace(current_config=lambda: dist_config, - current_context=lambda: SimpleNamespace(dist_config=dist_config))) - attn.num_heads = 1 - attn.kv_lora_rank = 1 - attn.qk_rope_head_dim = 1 - attn._qkv_proj = fake_qkv_proj - attn.apply_rotary_pos_emb = lambda q, k, cos, sin, inplace=False: (q, k) - attn.vc = lambda attn_output, out: out.copy_(attn_output) - attn.v_head_dim = 1 - attn.o_proj = nn.Identity() - - -def _make_kv_b_loader_model(fp8_quant_scope=None, input_dim=3): - model = DeepseekV2ForCausalLM.__new__(DeepseekV2ForCausalLM) - nn.Module.__init__(model) - model.config = SimpleNamespace(qk_nope_head_dim=2, v_head_dim=2) - model.quantization_config = {'quant_method': 'fp8'} - if fp8_quant_scope is not None: - model.quantization_config['fp8_quant_scope'] = fp8_quant_scope - model._load_buffers = {} - prefix = 'model.layers.0.self_attn' - params = { - f'{prefix}.kc.weight': nn.Parameter(torch.empty(1, 2, input_dim, dtype=torch.bfloat16)), - f'{prefix}.vc.weight': nn.Parameter(torch.empty(1, input_dim, 2, dtype=torch.bfloat16)), - } - return model, prefix, params - - -def _split_kv_b_weight(weight): - weight = weight.unflatten(0, (-1, 4)) - weight_kc, weight_vc = weight.split([2, 2], dim=1) - return weight_kc, weight_vc.transpose(1, 2).contiguous() - - -def test_glm_moe_only_loads_kv_b_proj_as_bf16(): - model, prefix, params = _make_kv_b_loader_model(fp8_quant_scope='moe_only') - loaded_weight = torch.arange(12, dtype=torch.bfloat16).reshape(4, 3) - - model._load_weight_attention( - name=f'{prefix}.kv_b_proj.weight', - loaded_weight=loaded_weight, - params_dict=params, - update_pe_mapping=[], - ) - - expected_kc, expected_vc = _split_kv_b_weight(loaded_weight) - assert torch.equal(params[f'{prefix}.kc.weight'], expected_kc) - assert torch.equal(params[f'{prefix}.vc.weight'], expected_vc) - assert model._load_buffers == {} - - -def test_global_fp8_loads_kv_b_proj_after_weight_scale(): - model, prefix, params = _make_kv_b_loader_model(input_dim=4) - loaded_weight = torch.arange(16, dtype=torch.float32).reshape(4, 4).to(torch.float8_e4m3fn) - loaded_scale = torch.full((1, 1), 0.5, dtype=torch.float32) - - model._load_weight_attention( - name=f'{prefix}.kv_b_proj.weight', - loaded_weight=loaded_weight, - params_dict=params, - update_pe_mapping=[], - ) - assert f'{prefix}.kv_b_proj.weight' in model._load_buffers - - model._load_weight_attention( - name=f'{prefix}.kv_b_proj.weight_scale_inv', - loaded_weight=loaded_scale, - params_dict=params, - update_pe_mapping=[], - ) - - dequantized_weight = (loaded_weight.float() * loaded_scale).to(torch.bfloat16) - expected_kc, expected_vc = _split_kv_b_weight(dequantized_weight) - assert torch.equal(params[f'{prefix}.kc.weight'], expected_kc) - assert torch.equal(params[f'{prefix}.vc.weight'], expected_vc) - assert model._load_buffers == {} - - -@pytest.mark.parametrize('n_heads', [32, 64]) -@pytest.mark.parametrize('load_scale_first', [False, True]) -def test_fp8_indexer_wk_loads_into_fused_bf16_projection(load_scale_first, n_heads): - model = DeepseekV32ForCausalLM.__new__(DeepseekV32ForCausalLM) - nn.Module.__init__(model) - model._load_buffers = {} - prefix = 'model.layers.0.self_attn.indexer' - fused = nn.Parameter(torch.zeros(128 + n_heads, 256, dtype=torch.bfloat16)) - params = {f'{prefix}.wk_weights_proj.weight': fused} - - wk = (torch.arange(128 * 256).reshape(128, 256) % 31 - 15).float().to(torch.float8_e4m3fn) - scale = torch.tensor([[0.25, 0.5]], dtype=torch.float32) - gate = torch.arange(n_heads * 256, dtype=torch.float32).reshape(n_heads, 256).to(torch.bfloat16) - wk_tensors = [ - (f'{prefix}.wk.weight', wk), - (f'{prefix}.wk.weight_scale_inv', scale), - ] - if load_scale_first: - wk_tensors.reverse() - - for name, tensor in wk_tensors: - model._load_weight_attention(name, tensor, params, update_pe_mapping=[]) - model._load_weight_attention(f'{prefix}.weights_proj.weight', gate, params, update_pe_mapping=[]) - - expected_wk = _dequantize_blocked_fp8(wk, scale, torch.bfloat16) - assert torch.equal(fused[:128], expected_wk) - assert torch.equal(fused[128:], gate) - assert model._load_buffers == {} - - -def test_bf16_indexer_wk_loads_directly_into_fused_projection(): - model = DeepseekV32ForCausalLM.__new__(DeepseekV32ForCausalLM) - nn.Module.__init__(model) - model._load_buffers = {} - prefix = 'model.layers.0.self_attn.indexer' - fused = nn.Parameter(torch.zeros(160, 8, dtype=torch.bfloat16)) - params = {f'{prefix}.wk_weights_proj.weight': fused} - wk = torch.arange(128 * 8, dtype=torch.float32).reshape(128, 8).to(torch.bfloat16) - - model._load_weight_attention(f'{prefix}.wk.weight', wk, params, update_pe_mapping=[]) - - assert torch.equal(fused[:128], wk) - assert model._load_buffers == {} - - -def test_indexer_rope_interleave_uses_interleaved_layout(): - indexer = Indexer.__new__(Indexer) - indexer.apply_rotary_pos_emb = DefaultApplyRotaryEmbImpl(interleaved=True).forward - - query = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]]) - key = torch.tensor([[5.0, 6.0, 7.0, 8.0]]) - cos = torch.tensor([[10.0, 20.0, 10.0, 20.0]]) - sin = torch.tensor([[1.0, 2.0, 1.0, 2.0]]) - - query_out, key_out = indexer._apply_rotary_pos_emb(query, key, (cos, sin)) - - key = key[..., None, :] - interleaved_cos = torch.tensor([[10.0, 10.0, 20.0, 20.0]]).unsqueeze(-2) - interleaved_sin = torch.tensor([[1.0, 1.0, 2.0, 2.0]]).unsqueeze(-2) - expected_query = query * interleaved_cos + rotate_interleaved(query) * interleaved_sin - expected_key = key * interleaved_cos + rotate_interleaved(key) * interleaved_sin - assert torch.equal(query_out, expected_query) - assert torch.equal(key_out, expected_key) - - -def test_full_indexer_layer_writes_compact_output_and_shared_buffer(monkeypatch): - attn = DeepseekV32Attention.__new__(DeepseekV32Attention) - nn.Module.__init__(attn) - attn.layer_idx = 0 - attn.indexer_output_idx = 0 - _patch_minimal_attention(attn, monkeypatch) - computed_topk = torch.tensor([[4, 2, 1], [3, 2, 0]], dtype=torch.int32) - attn.indexer = lambda *args, **kwargs: computed_topk - seen = {} - - def fake_attn_fwd(*args, **kwargs): - seen['nsa_indices'] = kwargs['nsa_indices'] - return torch.zeros(args[0].size(0), 1, 1) - - attn.attn_fwd = fake_attn_fwd - topk_buffer = DSATopKIndicesBuffer(topk=3) - all_indexer_topk = torch.full((2, 1, 3), -1, dtype=torch.int32) - - output = attn( - hidden_states=torch.zeros(1, 2, 2), - rotary_pos_emb=(torch.zeros(2, 1), torch.zeros(2, 1)), - past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], - topk_indices_buffer=topk_buffer, - all_indexer_topk=all_indexer_topk, - ) - - assert output.shape == (1, 2, 1) - assert torch.equal(topk_buffer.indices[:2], computed_topk) - assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:2].data_ptr() - assert torch.equal(all_indexer_topk[:, 0], computed_topk) - - -def test_mtp_skips_indexer_and_reads_compacted_topk(monkeypatch): - attn = DeepseekV32Attention.__new__(DeepseekV32Attention) - nn.Module.__init__(attn) - attn.layer_idx = 5 - attn.indexer_output_idx = None - _patch_minimal_attention(attn, monkeypatch) - attn.indexer = lambda *args, **kwargs: pytest.fail('recurrent MTP must reuse seed top-k') - seen = {} - - def fake_attn_fwd(*args, **kwargs): - seen['nsa_indices'] = kwargs['nsa_indices'] - return torch.zeros(args[0].size(0), 1, 1) - - attn.attn_fwd = fake_attn_fwd - topk_buffer = DSATopKIndicesBuffer(topk=3) - topk_buffer.write(torch.tensor([[1, 2, 3], [11, 12, 13], [21, 22, 23]], dtype=torch.int32)) - topk_buffer.compact(torch.tensor([2, 0])) - - attn( - hidden_states=torch.zeros(1, 2, 2), - rotary_pos_emb=(torch.zeros(2, 1), torch.zeros(2, 1)), - past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], - topk_indices_buffer=topk_buffer, - skip_topk=True, - ) - - assert torch.equal(seen['nsa_indices'], torch.tensor([[21, 22, 23], [1, 2, 3]], dtype=torch.int32)) - - -@pytest.mark.parametrize('dp,attn_tp,expected_num_heads', [(1, 2, 2), (2, 2, 4)]) -def test_attention_uses_dp_aware_num_heads(monkeypatch, dp, attn_tp, expected_num_heads): - attn = DeepseekV32Attention.__new__(DeepseekV32Attention) - nn.Module.__init__(attn) - attn.layer_idx = 0 - seen = {} - _patch_minimal_attention(attn, monkeypatch, dp=dp, attn_tp=attn_tp, seen=seen) - attn.num_heads = 4 - attn.indexer = lambda *args, **kwargs: torch.tensor([[4, 2, 1]], dtype=torch.int32) - attn.attn_fwd = lambda *args, **kwargs: torch.zeros(args[0].size(0), args[0].size(1), 1) - - attn( - hidden_states=torch.zeros(1, 1, 2), - rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), - past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], - topk_indices_buffer=DSATopKIndicesBuffer(topk=3), - ) - - assert seen['num_heads'] == expected_num_heads - - -@pytest.mark.parametrize('dp,attn_tp,expected_num_heads', [(1, 2, 2), (2, 2, 4)]) -def test_bmm_uses_dp_aware_weight_layout(monkeypatch, dp, attn_tp, expected_num_heads): - dist_config = SimpleNamespace(dp=dp, attn_tp=attn_tp) - monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v2.get_dist_manager', - lambda: SimpleNamespace(current_config=lambda: dist_config, - current_context=lambda: SimpleNamespace(dist_config=dist_config))) - monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v2.get_tp_world_rank', - lambda layer_type=None: (attn_tp, 0)) - - bmm = DeepseekV2BMM(batch=4, in_features=2, out_features=3, dtype=torch.float32, device='cpu') - weight = torch.arange(24, dtype=torch.float32).view(4, 2, 3) - bmm.weight.weight_loader(bmm.weight, weight) - - expected = weight if dp > 1 else weight.chunk(attn_tp, 0)[0] - assert bmm.weight.shape == (expected_num_heads, 2, 3) - torch.testing.assert_close(bmm.weight, expected) - - -def test_shared_indexer_layer_reuses_shared_topk_buffer(monkeypatch): - attn = DeepseekV32Attention.__new__(DeepseekV32Attention) - nn.Module.__init__(attn) - attn.layer_idx = 3 - attn.indexer = None - attn.indexer_output_idx = None - _patch_minimal_attention(attn, monkeypatch) - seen = {} - - def fake_attn_fwd(*args, **kwargs): - seen['nsa_indices'] = kwargs['nsa_indices'] - return torch.zeros(args[0].size(0), 1, 1) - - attn.attn_fwd = fake_attn_fwd - topk_buffer = DSATopKIndicesBuffer(topk=3) - shared_topk = torch.tensor([[5, 4, 3]], dtype=torch.int32) - topk_buffer.write(shared_topk) - all_indexer_topk = torch.full((1, 1, 3), -1, dtype=torch.int32) - - output = attn( - hidden_states=torch.zeros(1, 1, 2), - rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), - past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], - topk_indices_buffer=topk_buffer, - all_indexer_topk=all_indexer_topk, - ) - - assert output.shape == (1, 1, 1) - assert seen['nsa_indices'].data_ptr() == topk_buffer.indices[:1].data_ptr() - assert torch.all(all_indexer_topk == -1) - - -def test_attention_requires_shared_topk_buffer(monkeypatch): - attn = DeepseekV32Attention.__new__(DeepseekV32Attention) - nn.Module.__init__(attn) - attn.layer_idx = 3 - attn.indexer = None - attn.num_heads = 1 - attn.kv_lora_rank = 1 - attn.qk_rope_head_dim = 1 - dist_config = SimpleNamespace(dp=1, attn_tp=1) - monkeypatch.setattr('lmdeploy.pytorch.models.deepseek_v32.get_dist_manager', - lambda: SimpleNamespace(current_config=lambda: dist_config, - current_context=lambda: SimpleNamespace(dist_config=dist_config))) - attn._qkv_proj = lambda *args, **kwargs: ( - torch.zeros(1, 1, 2), - torch.zeros(1, 1, 2), - torch.zeros(1, 1, 1), - torch.zeros(1, 1, 1), - torch.zeros(1, 1, 1), - torch.zeros(1, 1, 1), - ) - attn.apply_rotary_pos_emb = lambda q, k, cos, sin, inplace=False: (q, k) - - with pytest.raises(RuntimeError, match='requires a DSA top-k indices buffer'): - attn( - hidden_states=torch.zeros(1, 1, 2), - rotary_pos_emb=(torch.zeros(1, 1), torch.zeros(1, 1)), - past_key_value=[torch.zeros(1, 1, 2), torch.zeros(1, 1, 1)], - ) - - -def test_layer_indexer_type_defaults_to_full_and_reads_shared_entries(): - config = SimpleNamespace(num_hidden_layers=3, indexer_types=['full', 'shared']) - - assert get_layer_indexer_type(config, 0) == 'full' - assert get_layer_indexer_type(config, 1) == 'shared' - assert get_layer_indexer_type(config, 2) == 'full' - assert get_layer_indexer_type(SimpleNamespace(indexer_types=None), 1) == 'full' - assert get_full_indexer_layer_ids(config) == (0, 2) - assert get_full_indexer_layer_ids(SimpleNamespace(num_hidden_layers=3, indexer_types=None)) == (0, 1, 2) - - -def test_moe_gate_captures_logical_experts_before_eplb_mapping(monkeypatch): - class FakeTopK(nn.Module): - - def forward(self, logits): - logical_ids = torch.tensor([[3, 1], [2, 0]], device=logits.device) - return torch.ones_like(logical_ids, dtype=torch.float32), logical_ids - - gate = MoEGate.__new__(MoEGate) - nn.Module.__init__(gate) - gate.weight = nn.Parameter(torch.zeros(4, 2)) - gate.fake_eplb = False - gate.topk_method = 'greedy' - gate.renormalize = False - gate.routed_scaling_factor = 1.0 - gate.softmax_topk = FakeTopK() - gate.eplb_dispatch_info = object() - monkeypatch.setattr( - 'lmdeploy.pytorch.models.deepseek_v2.EPLBManager.topk_ids_logical_to_physical', - lambda ids, info: ids + 10, - ) - captured = torch.full((2, 2), torch.iinfo(torch.uint16).max, dtype=torch.uint16) - - _, dispatch_ids = gate(torch.zeros(2, 2), routed_experts=captured) - - expected = torch.tensor([[3, 1], [2, 0]]) - assert torch.equal(captured, expected.to(torch.uint16)) - assert torch.equal(dispatch_ids, expected + 10) - - -def test_deepseek_moe_writes_only_its_routed_expert_layer(): - class FakeGate(nn.Module): - - def forward(self, hidden_states, routed_experts=None): - ids = torch.tensor([[3, 1], [2, 0]], device=hidden_states.device) - if routed_experts is not None: - routed_experts.copy_(ids) - return torch.ones_like(ids, dtype=torch.float32), ids - - class FakeExperts(nn.Module): - - def forward(self, hidden_states, topk_weights, topk_ids): - return hidden_states - - moe = DeepseekV2MoE.__new__(DeepseekV2MoE) - nn.Module.__init__(moe) - moe.layer_idx = 3 - moe.hidden_dim = 2 - moe.gate = FakeGate() - moe.experts = FakeExperts() - moe.shared_experts = None - moe._all_reduce = False - sentinel = torch.iinfo(torch.uint16).max - all_routed_experts = torch.full((2, 5, 2), sentinel, dtype=torch.uint16) - - output = moe(torch.zeros(1, 2, 2), all_routed_experts=all_routed_experts) - - assert output.shape == (1, 2, 2) - assert torch.equal(all_routed_experts[:, 3], torch.tensor([[3, 1], [2, 0]], dtype=torch.uint16)) - assert torch.all(all_routed_experts[:, :3] == sentinel) - assert torch.all(all_routed_experts[:, 4] == sentinel) - - -def test_glm52_causal_lm_returns_routed_experts_and_indexer_topk(monkeypatch): - class FakeModel(nn.Module): - - def forward(self, input_ids, all_routed_experts=None, all_indexer_topk=None, **kwargs): - if all_routed_experts is not None: - all_routed_experts[:, 3].fill_(7) - all_routed_experts[:, 4].fill_(9) - if all_indexer_topk is not None: - all_indexer_topk.fill_(5) - return torch.zeros(1, input_ids.size(1), 4) - - forward_microbatch = forward - - context = SimpleNamespace(enable_microbatch=False) - monkeypatch.setattr( - 'lmdeploy.pytorch.models.deepseek_v32.get_step_ctx_manager', - lambda: SimpleNamespace(current_context=lambda: context), - ) - model = DeepseekV32ForCausalLM.__new__(DeepseekV32ForCausalLM) - nn.Module.__init__(model) - model.config = SimpleNamespace(num_hidden_layers=5, - num_experts_per_tok=8, - index_topk=3, - indexer_types=['full', 'full', 'full', 'shared', 'shared']) - model.enable_return_routed_experts = True - model.enable_return_indexer_topk = True - model.num_indexer_layers = 3 - model.model = FakeModel() - - outputs = model( - input_ids=torch.ones(1, 2, dtype=torch.long), - position_ids=torch.arange(2)[None], - past_key_values=[], - ) - - routed_experts = outputs['all_routed_experts'] - assert routed_experts.dtype == torch.uint16 - assert routed_experts.shape == (2, 5, 8) - assert torch.all(routed_experts[:, :3] == torch.iinfo(torch.uint16).max) - assert torch.all(routed_experts[:, 3] == 7) - assert torch.all(routed_experts[:, 4] == 9) - assert outputs['all_indexer_topk'].shape == (2, 3, 3) - assert torch.all(outputs['all_indexer_topk'] == 5) From d8d6886a38a48860d9bd342caa6685dafc14a247 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 21 Jul 2026 18:27:47 +0800 Subject: [PATCH 18/33] test: fix GLM-5.2 CI compatibility --- .../pytorch/backends/cuda/attention/mla.py | 1 + .../pytorch/config/test_glm_moe_dsa_config.py | 219 ------------------ tests/pytorch/kernel/test_dsa_indexer.py | 2 + 3 files changed, 3 insertions(+), 219 deletions(-) delete mode 100644 tests/pytorch/config/test_glm_moe_dsa_config.py diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index 5cb87c66a4..fae7f09b51 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -208,6 +208,7 @@ def _flash_mla_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, """Run sparse FlashMLA over a BF16 flat KV view.""" 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)) diff --git a/tests/pytorch/config/test_glm_moe_dsa_config.py b/tests/pytorch/config/test_glm_moe_dsa_config.py deleted file mode 100644 index 2900ff7b14..0000000000 --- a/tests/pytorch/config/test_glm_moe_dsa_config.py +++ /dev/null @@ -1,219 +0,0 @@ -import json -from types import SimpleNamespace - -from lmdeploy.pytorch.config import ModelConfig -from lmdeploy.pytorch.configurations.deepseek_v32 import DeepseekV32ModelConfigBuilder -from lmdeploy.pytorch.configurations.glm_moe_dsa import ( - GlmMoeDsaModelConfigBuilder, - normalize_glm_moe_dsa_config, -) - - -def test_get_model_arch_registers_deepseek_v32_config(tmp_path): - from lmdeploy.archs import get_model_arch - from lmdeploy.pytorch.transformers.configuration_deepseek_v32 import DeepseekV32Config - - config_path = tmp_path / 'config.json' - config_path.write_text( - json.dumps({ - 'architectures': ['DeepseekV32ForCausalLM'], - 'model_type': 'deepseek_v32', - })) - - arch, config = get_model_arch(str(tmp_path)) - - assert arch == 'DeepseekV32ForCausalLM' - assert isinstance(config, DeepseekV32Config) - - -def test_get_model_arch_loads_native_glm_moe_dsa_config(tmp_path): - from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig - - from lmdeploy.archs import get_model_arch - - config_path = tmp_path / 'config.json' - config_path.write_text( - json.dumps({ - 'architectures': ['GlmMoeDsaForCausalLM'], - 'model_type': 'glm_moe_dsa', - })) - - arch, config = get_model_arch(str(tmp_path)) - - assert arch == 'GlmMoeDsaForCausalLM' - assert isinstance(config, GlmMoeDsaConfig) - - -def _make_config(**kwargs): - values = dict( - model_type='glm_moe_dsa', - num_hidden_layers=4, - qk_nope_head_dim=128, - qk_rope_head_dim=64, - ) - values.update(kwargs) - values.setdefault('qk_head_dim', values['qk_nope_head_dim'] + values['qk_rope_head_dim']) - values.setdefault('indexer_types', ['full'] * values['num_hidden_layers']) - return SimpleNamespace(**values) - - -def _patch_v32_base_builder(monkeypatch): - from lmdeploy.pytorch.configurations.deepseek_v2 import DeepseekV2ModelConfigBuilder - - def fake_build(cls, hf_config, model_path=None, **kwargs): - return SimpleNamespace() - - monkeypatch.setattr(DeepseekV2ModelConfigBuilder, 'build', classmethod(fake_build)) - - -def _make_fp8_build_config(**quantization_config): - return _make_config( - use_flash_mla=True, - index_head_dim=128, - index_topk=2048, - quantization_config=quantization_config, - ) - - -def test_glm_moe_dsa_enables_online_fp8_moe_only_scope(monkeypatch): - from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig - - from lmdeploy.pytorch import envs - from lmdeploy.pytorch import transformers as pytorch_transformers - from lmdeploy.pytorch.configurations import deepseek_v2 - - monkeypatch.setattr(envs, 'fp8_moe_only', True) - monkeypatch.setattr(deepseek_v2, 'flash_mla_available', lambda: True) - cfg = GlmMoeDsaConfig(hidden_size=16, - intermediate_size=32, - moe_intermediate_size=8, - num_hidden_layers=2, - num_attention_heads=2, - num_key_value_heads=1, - q_lora_rank=4, - kv_lora_rank=4, - qk_nope_head_dim=4, - qk_rope_head_dim=4, - vocab_size=32, - bos_token_id=1, - eos_token_id=2, - dtype='float16', - index_head_dim=4, - index_n_heads=2, - index_topk=2) - monkeypatch.setattr(pytorch_transformers, 'config_from_pretrained', lambda *args, **kwargs: cfg) - - model_config = ModelConfig.from_pretrained('fake-model', model_format='fp8', dtype='bfloat16') - - assert model_config.quant_config.quant_method == 'fp8' - assert model_config.quant_config.fp8_quant_scope == 'moe_only' - assert model_config.quant_config.hf_quant_config['lmdeploy_patched'] - - -def test_glm_moe_dsa_does_not_override_prequantized_fp8_scope(monkeypatch): - from lmdeploy.pytorch import envs - - monkeypatch.setattr(envs, 'fp8_moe_only', True) - _patch_v32_base_builder(monkeypatch) - cfg = _make_fp8_build_config(quant_method='fp8') - - GlmMoeDsaModelConfigBuilder.build(cfg) - - assert 'fp8_quant_scope' not in cfg.quantization_config - - -def test_glm_moe_dsa_runtime_normalizes_native_hf_config(): - from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig - - cfg = GlmMoeDsaConfig(qk_nope_head_dim=192, - qk_rope_head_dim=64, - qk_head_dim=256, - head_dim=192, - hidden_size=4096, - intermediate_size=8192, - num_attention_heads=32, - num_hidden_layers=2, - vocab_size=32000, - bos_token_id=1, - eos_token_id=2) - - assert cfg.head_dim == 192 - - normalize_glm_moe_dsa_config(cfg) - - assert cfg.head_dim == 64 - assert cfg.qk_rope_head_dim == 64 - - -def test_glm_moe_dsa_builder_creates_sparse_mla_config(monkeypatch): - from lmdeploy.pytorch.configurations import deepseek_v2 - - monkeypatch.setattr(deepseek_v2, 'flash_mla_available', lambda: True) - cfg = _make_config(num_hidden_layers=2, - index_head_dim=128, - index_topk=2048, - hidden_size=6144, - kv_lora_rank=512, - num_attention_heads=64, - num_key_value_heads=64, - bos_token_id=None, - eos_token_id=[154820, 154827, 154829], - vocab_size=154880) - - model_config = GlmMoeDsaModelConfigBuilder.build(cfg, tp=1) - - assert cfg.num_key_value_heads == 1 - assert model_config.num_key_value_heads == 1 - assert model_config.head_dim == 576 - assert model_config.mla_kv_cache_dtype == 'bfloat16' - assert model_config.use_mla_fp8_cache is False - - -def test_deepseek_v32_builder_defaults_to_bf16_sparse_mla(monkeypatch): - _patch_v32_base_builder(monkeypatch) - cfg = _make_config(model_type='deepseek_v32', use_flash_mla=True, index_head_dim=128, index_topk=2048) - - model_config = DeepseekV32ModelConfigBuilder.build(cfg) - - assert model_config.mla_kv_cache_dtype == 'bfloat16' - - -def test_glm_moe_dsa_draft_selects_glm_mtp_model(monkeypatch): - _patch_v32_base_builder(monkeypatch) - cfg = _make_config(architectures=['GlmMoeDsaForCausalLM'], - use_flash_mla=True, - index_head_dim=128, - index_topk=2048, - quantization_config={}) - - GlmMoeDsaModelConfigBuilder.build(cfg, is_draft_model=True) - - assert cfg.architectures == ['GlmMoeDsaMTPModel'] - - -def test_glm_moe_dsa_mtp_reuses_target_dist_config(monkeypatch): - from lmdeploy.pytorch import transformers as pytorch_transformers - from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig - from lmdeploy.pytorch.engine.config_builder import ConfigBuilder - - monkeypatch.setattr(pytorch_transformers, 'config_from_pretrained', - lambda *args, **kwargs: SimpleNamespace(model_type='glm_moe_dsa')) - captured = {} - - def fake_from_config(**kwargs): - captured.update(kwargs) - return kwargs - - monkeypatch.setattr(SpecDecodeConfig, 'from_config', staticmethod(fake_from_config)) - target_dist_config = DistConfig(tp=8) - speculative_config = SimpleNamespace(method='deepseek_mtp', model=None, num_speculative_tokens=5) - engine_config = SimpleNamespace(dtype='auto', model_format=None, hf_overrides=None) - - ConfigBuilder.build_specdecode_config('glm-model', - speculative_config, - engine_config, - cache_config=object(), - dist_config=target_dist_config, - trust_remote_code=True) - - assert captured['dist_config'] is target_dist_config diff --git a/tests/pytorch/kernel/test_dsa_indexer.py b/tests/pytorch/kernel/test_dsa_indexer.py index 176a306519..7c8233cabe 100644 --- a/tests/pytorch/kernel/test_dsa_indexer.py +++ b/tests/pytorch/kernel/test_dsa_indexer.py @@ -27,6 +27,8 @@ def _apply_rope_first(x, cos, sin, rope_interleaved): 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 0f61deae85968f966d1e47f2827a5f488b5215b1 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 21 Jul 2026 20:30:38 +0800 Subject: [PATCH 19/33] test: skip DSA FP8 cache test on SM80 --- tests/pytorch/kernel/test_dsa_indexer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/pytorch/kernel/test_dsa_indexer.py b/tests/pytorch/kernel/test_dsa_indexer.py index 7c8233cabe..763dbe52ca 100644 --- a/tests/pytorch/kernel/test_dsa_indexer.py +++ b/tests/pytorch/kernel/test_dsa_indexer.py @@ -66,6 +66,8 @@ def test_prepare_dsa_indexer_q_matches_unfused_quantization(rope_interleaved, he 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 ae9cf07c5060c85a0da8ee6862815bb3252f366d Mon Sep 17 00:00:00 2001 From: zxy Date: Wed, 22 Jul 2026 17:25:35 +0800 Subject: [PATCH 20/33] fix: preserve DSA query width for MTP --- lmdeploy/pytorch/nn/nsa.py | 26 ++++++++++++++++++++++++-- tests/pytorch/nn/test_nsa.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) create mode 100644 tests/pytorch/nn/test_nsa.py diff --git a/lmdeploy/pytorch/nn/nsa.py b/lmdeploy/pytorch/nn/nsa.py index a21669a436..57e7386dc4 100644 --- a/lmdeploy/pytorch/nn/nsa.py +++ b/lmdeploy/pytorch/nn/nsa.py @@ -15,6 +15,28 @@ 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) + @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() @@ -23,8 +45,8 @@ def _build_meta(q: Tensor, attn_metadata: AttentionMetadata) -> NSAIndexMeta: 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 (attn_metadata.max_q_seqlen or 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 return NSAIndexMeta(cu_seqlen_q=attn_metadata.cu_seqlens_q, q_seqlens=attn_metadata.q_seqlens, 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 From 8ce59f2a2b713829ae7ca98e6f14f2d50eff5b08 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 23 Jul 2026 15:42:42 +0800 Subject: [PATCH 21/33] fix: guard optional GLM runtime paths --- lmdeploy/pytorch/backends/cuda/nsa.py | 11 +++++---- .../spec_decode/proposers/deepseek_mtp.py | 3 ++- .../pytorch/kernel/test_sparse_index_topk.py | 14 +++++++++++ .../spec_decode/test_glm_moe_dsa_mtp.py | 24 +++++++++++++++++++ 4 files changed, 47 insertions(+), 5 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index 6b3d830067..31d62de63f 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -15,10 +15,13 @@ @functools.lru_cache def _get_sparse_index_topk(topk: int): - from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import ( - is_sparse_index_topk_supported, - sparse_index_topk, - ) + 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 diff --git a/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py b/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py index ebf997a446..3e2effe361 100644 --- a/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py +++ b/lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py @@ -14,7 +14,8 @@ def build_model(self, empty_init: bool, target_model: torch.nn.Module = None, bu """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 hasattr(draft_model, 'uses_dsa_topk_buffer'): + 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, diff --git a/tests/pytorch/kernel/test_sparse_index_topk.py b/tests/pytorch/kernel/test_sparse_index_topk.py index f2d6b92b0a..8d7462e6dc 100644 --- a/tests/pytorch/kernel/test_sparse_index_topk.py +++ b/tests/pytorch/kernel/test_sparse_index_topk.py @@ -1,4 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. +import sys + import pytest import torch @@ -17,6 +19,18 @@ def test_nsa_backend_selects_sparse_topk_for_glm52(): assert _get_sparse_index_topk(1024) is None +def test_nsa_backend_falls_back_without_tilelang(monkeypatch): + from lmdeploy.pytorch.backends.cuda import nsa + + nsa._get_sparse_index_topk.cache_clear() + monkeypatch.setitem(sys.modules, + 'lmdeploy.pytorch.kernels.cuda.sparse_index_topk', None) + try: + assert nsa._get_sparse_index_topk(2048) is None + finally: + nsa._get_sparse_index_topk.cache_clear() + + def _assert_topk_ids(scores: torch.Tensor, out: torch.Tensor, seqlens: list[int], diff --git a/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py index c2711c6e9e..e5e42cad78 100644 --- a/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py +++ b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py @@ -1,6 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import asyncio +import pytest import torch from torch import nn @@ -116,6 +117,29 @@ def build_model(self, empty_init, target_model=None, build_model_ctx=None): 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): From cf79dc44dca3a56e70d07c997dc30a5636ef45c2 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 23 Jul 2026 15:43:01 +0800 Subject: [PATCH 22/33] fix: grow indexer history on demand --- lmdeploy/pytorch/messages.py | 28 +++++++++----------- tests/pytorch/spec_decode/test_strategies.py | 12 ++++++--- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py index d54ed777d1..a2d9fa3f08 100644 --- a/lmdeploy/pytorch/messages.py +++ b/lmdeploy/pytorch/messages.py @@ -593,7 +593,6 @@ def _get_pad_width(self, reserve_size: int): class HistoryIndexerTopK(_HistoryDataBase): """History of full-layer sparse-attention indexer results.""" - ALLOC_SIZE = 1 def __init__(self, indexer_topk: np.ndarray = None, dtype: np.dtype = np.int32): super().__init__(indexer_topk, dtype) @@ -601,24 +600,23 @@ def __init__(self, indexer_topk: np.ndarray = None, dtype: np.dtype = np.int32): def _create_empty_array(self, dtype): return None - def _get_pad_width(self, reserve_size: int): - return ((0, reserve_size), (0, 0), (0, 0)) + def reserve(self, size: int): + """Grow geometrically without reserving the request token limit.""" + if self._data is None or len(self._data) >= size: + return + capacity = max(size, max(1, len(self._data) * 2)) + data = np.empty((capacity, *self._data.shape[1:]), dtype=self.dtype) + data[:self._num_real] = self._data[:self._num_real] + self._data = data - def append(self, new_data: np.ndarray, reserve_size: int | None = None): - """Append results, reserving the expected request length once.""" + def append(self, new_data: np.ndarray): + """Append indexer results.""" new_data = np.asarray(new_data) if new_data.ndim != 3: raise ValueError(f'indexer_topk must be a 3D array, got shape {new_data.shape}.') - if self._data is None: - capacity = len(new_data) - if reserve_size is not None: - capacity = max(capacity, reserve_size) - self._data = np.empty((capacity, *new_data.shape[1:]), dtype=self.dtype) - elif self._data.shape[1:] != new_data.shape[1:]: + if self._data is not None and self._data.shape[1:] != new_data.shape[1:]: raise ValueError( f'indexer_topk shape changed from {self._data.shape[1:]} to {new_data.shape[1:]}.') - if reserve_size is not None: - self.reserve(reserve_size) super().append(new_data) @@ -897,9 +895,7 @@ def append_indexer_topk(self, indexer_topk: Tensor | np.ndarray): return if isinstance(indexer_topk, Tensor): indexer_topk = indexer_topk.cpu().numpy() - reserve_size = self.output_start_pos + self.sampling_param.max_new_tokens - reserve_size = max(reserve_size, len(self.all_indexer_topk) + len(indexer_topk)) - self.all_indexer_topk.append(indexer_topk, reserve_size=reserve_size) + self.all_indexer_topk.append(indexer_topk) @property def num_history_ids(self): diff --git a/tests/pytorch/spec_decode/test_strategies.py b/tests/pytorch/spec_decode/test_strategies.py index cfbe181f23..f42fd69912 100644 --- a/tests/pytorch/spec_decode/test_strategies.py +++ b/tests/pytorch/spec_decode/test_strategies.py @@ -546,7 +546,8 @@ def _make_seq_with_indexer(prefill_tokens=None): seq_meta = SequenceMeta(block_size=16, strategy=strategy) session = MagicMock() session.seq_meta = seq_meta - sampling_param = SamplingParam(return_indexer_topk=True, max_new_tokens=8) + sampling_param = SamplingParam(return_indexer_topk=True, + max_new_tokens=512) seq = SchedulerSequenceARSpec(seq_id=0, session=session, sampling_param=sampling_param) if prefill_tokens is not None: seq._update_token_ids_inputs(np.array(prefill_tokens, dtype=np.int64)) @@ -575,14 +576,17 @@ def test_decode_keeps_only_accepted_indexer_rows(self): assert len(seq.all_indexer_topk) == 2 assert np.array_equal(seq.all_indexer_topk.get_real(), _indexer_topk(2)) - def test_indexer_history_reserves_request_capacity_and_truncates_logically(self): + def test_indexer_history_grows_geometrically_and_truncates_logically(self): seq = _make_seq_with_indexer([1, 2, 3, 4]) seq.append_indexer_topk(_indexer_topk(4)) - assert len(seq.all_indexer_topk._data) == 12 + assert len(seq.all_indexer_topk._data) == 4 + seq.append_indexer_topk(_indexer_topk(2)) + assert len(seq.all_indexer_topk) == 6 + assert len(seq.all_indexer_topk._data) == 8 seq.set_step(3) assert len(seq.all_indexer_topk) == 3 - assert len(seq.all_indexer_topk._data) == 12 + assert len(seq.all_indexer_topk._data) == 8 # --------------------------------------------------------------------------- From d91ca8146ee77b38e0a0df26b9ceb122a9586641 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 23 Jul 2026 15:50:52 +0800 Subject: [PATCH 23/33] refactor: rename DSA indexer preprocessing module --- lmdeploy/pytorch/backends/cuda/nsa.py | 2 +- .../cuda/{dsa_indexer.py => dsa_indexer_preprocess.py} | 2 ++ .../{test_dsa_indexer.py => test_dsa_indexer_preprocess.py} | 4 ++-- 3 files changed, 5 insertions(+), 3 deletions(-) rename lmdeploy/pytorch/kernels/cuda/{dsa_indexer.py => dsa_indexer_preprocess.py} (99%) rename tests/pytorch/kernel/{test_dsa_indexer.py => test_dsa_indexer_preprocess.py} (96%) diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index 31d62de63f..b4d1196511 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -7,7 +7,7 @@ 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.dsa_indexer import prepare_dsa_indexer_k_cache, prepare_dsa_indexer_q +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 diff --git a/lmdeploy/pytorch/kernels/cuda/dsa_indexer.py b/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py similarity index 99% rename from lmdeploy/pytorch/kernels/cuda/dsa_indexer.py rename to lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py index f2cce7b772..9bb35ff64a 100644 --- a/lmdeploy/pytorch/kernels/cuda/dsa_indexer.py +++ b/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py @@ -1,4 +1,6 @@ # 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 diff --git a/tests/pytorch/kernel/test_dsa_indexer.py b/tests/pytorch/kernel/test_dsa_indexer_preprocess.py similarity index 96% rename from tests/pytorch/kernel/test_dsa_indexer.py rename to tests/pytorch/kernel/test_dsa_indexer_preprocess.py index 763dbe52ca..dc2f027b57 100644 --- a/tests/pytorch/kernel/test_dsa_indexer.py +++ b/tests/pytorch/kernel/test_dsa_indexer_preprocess.py @@ -34,7 +34,7 @@ def _apply_rope_first(x, cos, sin, rope_interleaved): 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 import prepare_dsa_indexer_q + from lmdeploy.pytorch.kernels.cuda.dsa_indexer_preprocess import prepare_dsa_indexer_q torch.manual_seed(0) tokens = 11 @@ -71,7 +71,7 @@ def test_prepare_dsa_indexer_q_matches_unfused_quantization(rope_interleaved, he @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 import prepare_dsa_indexer_k, prepare_dsa_indexer_k_cache + 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) From f4ffe8d32de5cedc38e5e86f1cf1d4d800fab239 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 23 Jul 2026 17:25:34 +0800 Subject: [PATCH 24/33] perf: upgrade FlashMLA sparse decode --- docker/prepare_3rdparty_wheel.sh | 2 +- .../pytorch/backends/cuda/attention/mla.py | 9 +++++ lmdeploy/pytorch/models/utils/cudagraph.py | 39 +++++++++++-------- tests/pytorch/kernel/test_mla_attention.py | 27 +++++++++++++ 4 files changed, 59 insertions(+), 18 deletions(-) diff --git a/docker/prepare_3rdparty_wheel.sh b/docker/prepare_3rdparty_wheel.sh index 034eb8e858..957c4e7423 100644 --- a/docker/prepare_3rdparty_wheel.sh +++ b/docker/prepare_3rdparty_wheel.sh @@ -15,7 +15,7 @@ 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 diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index fae7f09b51..20048899c1 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -179,6 +179,7 @@ def _decode_paged_flash_mla( 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) @@ -187,6 +188,13 @@ def _decode_paged_flash_mla( block_size = k_cache.size(1) 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, @@ -200,6 +208,7 @@ def _decode_paged_flash_mla( 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 diff --git a/lmdeploy/pytorch/models/utils/cudagraph.py b/lmdeploy/pytorch/models/utils/cudagraph.py index 3ce0e4558d..c6a9c605ea 100644 --- a/lmdeploy/pytorch/models/utils/cudagraph.py +++ b/lmdeploy/pytorch/models/utils/cudagraph.py @@ -265,23 +265,28 @@ def fill_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p 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 - 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'] + 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/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index dea48a0d15..72e41b03c6 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -54,6 +54,33 @@ def test_bf16_sparse_decode_uses_global_cache_view(): assert torch.equal(global_indices, 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) From 7196130ef0079445810dc54d7d0bef33ffd8b4bd Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 23 Jul 2026 20:34:19 +0800 Subject: [PATCH 25/33] refactor: clean up DSA configuration and kernels --- .../pytorch/configurations/glm_moe_dsa.py | 13 +- .../kernels/cuda/dsa_indexer_preprocess.py | 149 ++++++++++++++---- 2 files changed, 123 insertions(+), 39 deletions(-) diff --git a/lmdeploy/pytorch/configurations/glm_moe_dsa.py b/lmdeploy/pytorch/configurations/glm_moe_dsa.py index c569399e55..7b9e09a179 100644 --- a/lmdeploy/pytorch/configurations/glm_moe_dsa.py +++ b/lmdeploy/pytorch/configurations/glm_moe_dsa.py @@ -7,14 +7,6 @@ logger = get_logger('lmdeploy') - -def normalize_glm_moe_dsa_config(hf_config): - """Normalize GLM-MoE-DSA checkpoint fields for the PyTorch runtime.""" - 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 - - class GlmMoeDsaModelConfigBuilder(DeepseekV32ModelConfigBuilder): @classmethod @@ -35,7 +27,10 @@ def build(cls, hf_config, model_path: str | None = None, **kwargs): 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.') - normalize_glm_moe_dsa_config(hf_config) + 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' diff --git a/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py b/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py index 9bb35ff64a..5c036d497f 100644 --- a/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py +++ b/lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py @@ -15,8 +15,15 @@ @triton.jit -def _apply_rope_first(x, x_pair, cos, sin, feat_off, rope_dim: tl.constexpr, - rope_interleaved: tl.constexpr): +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 @@ -29,14 +36,34 @@ def _apply_rope_first(x, x_pair, cos, sin, feat_off, rope_dim: tl.constexpr, @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): +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 @@ -76,10 +103,25 @@ def _prepare_dsa_indexer_q_kernel(Q, Weights, Cos, Sin, QOut, QScaleOut, num_hea @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): +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 @@ -115,14 +157,36 @@ def _prepare_dsa_indexer_k_kernel(K, NormWeight, NormBias, Cos, Sin, KOut, eps: @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): +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) @@ -180,8 +244,15 @@ def _prepare_dsa_indexer_k_cache_kernel(K, NormWeight, NormBias, Cos, Sin, KCach 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]: +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 @@ -223,8 +294,15 @@ def prepare_dsa_indexer_q(q: Tensor, weights: Tensor, cos: Tensor, sin: Tensor, 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: +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, ) @@ -252,10 +330,21 @@ def prepare_dsa_indexer_k(k: Tensor, norm_weight: Tensor, norm_bias: Tensor, cos 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: +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 From d2d573522b5d2a13cc0fce692d9cbc453df0c0e9 Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 24 Jul 2026 11:23:56 +0800 Subject: [PATCH 26/33] refactor: remove DSA indexer replay --- benchmark/benchmark_chat_completion.py | 22 +------ lmdeploy/cli/serve.py | 2 - lmdeploy/cli/utils.py | 9 --- lmdeploy/messages.py | 6 -- lmdeploy/pytorch/config.py | 2 - lmdeploy/pytorch/engine/engine.py | 3 - lmdeploy/pytorch/engine/engine_instance.py | 9 +-- lmdeploy/pytorch/engine/engine_loop.py | 8 +-- lmdeploy/pytorch/engine/inputs_maker.py | 4 -- lmdeploy/pytorch/engine/model_agent/agent.py | 11 ---- lmdeploy/pytorch/messages.py | 56 ----------------- lmdeploy/pytorch/model_inputs.py | 1 - lmdeploy/pytorch/models/deepseek_v32.py | 40 +----------- lmdeploy/pytorch/models/utils/cudagraph.py | 2 - lmdeploy/pytorch/paging/block_trie.py | 62 ------------------- lmdeploy/pytorch/strategies/ar/sequence.py | 18 +----- .../pytorch/strategies/ar_spec/sequence.py | 35 ++--------- lmdeploy/serve/anthropic/adapter.py | 1 - .../serve/anthropic/endpoints/messages.py | 10 +-- lmdeploy/serve/anthropic/protocol.py | 8 --- lmdeploy/serve/anthropic/streaming.py | 7 +-- lmdeploy/serve/core/async_engine.py | 8 --- lmdeploy/serve/openai/api_server.py | 26 ++------ lmdeploy/serve/openai/protocol.py | 8 --- .../serve/openai/serving_chat_completion.py | 4 -- lmdeploy/serve/openai/serving_generate.py | 3 - tests/pytorch/engine/test_engine_sleep.py | 34 +--------- tests/pytorch/paging/test_block_trie.py | 44 ------------- tests/pytorch/spec_decode/test_strategies.py | 49 --------------- .../serve/test_generation_config.py | 10 --- 30 files changed, 22 insertions(+), 480 deletions(-) diff --git a/benchmark/benchmark_chat_completion.py b/benchmark/benchmark_chat_completion.py index 0a3d73f1d3..b9bc0c9e04 100644 --- a/benchmark/benchmark_chat_completion.py +++ b/benchmark/benchmark_chat_completion.py @@ -56,7 +56,6 @@ class SSEEvent: done: bool = False raw: dict[str, Any] | None = None routed_experts: str | None = None - indexer_topk: str | None = None @property def token_text(self) -> str: @@ -359,7 +358,6 @@ def parse_sse_line(line: bytes | str) -> SSEEvent: usage=data.get('usage'), raw=data, routed_experts=choice.get('routed_experts'), - indexer_topk=choice.get('indexer_topk'), ) @@ -380,7 +378,6 @@ def build_payload( ignore_eos: bool = False, return_token_ids: bool = False, return_routed_experts: bool = False, - return_indexer_topk: bool = False, return_logprob: bool = False, logprobs: bool = False, top_logprobs: int | None = None, @@ -411,8 +408,6 @@ def build_payload( payload['return_token_ids'] = True if return_routed_experts: payload['return_routed_experts'] = True - if return_indexer_topk: - payload['return_indexer_topk'] = True if return_logprob: payload['return_logprob'] = True if logprobs: @@ -469,7 +464,6 @@ async def request_chat_completion( ignore_eos: bool, return_token_ids: bool, return_routed_experts: bool, - return_indexer_topk: bool, return_logprob: bool, logprobs: bool, top_logprobs: int | None, @@ -488,7 +482,6 @@ async def request_chat_completion( ignore_eos=ignore_eos, return_token_ids=return_token_ids, return_routed_experts=return_routed_experts, - return_indexer_topk=return_indexer_topk, return_logprob=return_logprob, logprobs=logprobs, top_logprobs=top_logprobs, @@ -541,11 +534,6 @@ async def request_chat_completion( await fetch_routed_experts(shared_store, event.routed_experts) except Exception as e: # noqa: BLE001 - record and keep consuming SSE. trace.error = repr(e) - if event.indexer_topk and shared_store is not None: - try: - await fetch_shared_output(shared_store, event.indexer_topk) - except Exception as e: # noqa: BLE001 - record and keep consuming SSE. - trace.error = repr(e) trace.end_time = time.perf_counter() trace.success = trace.error == '' @@ -1032,7 +1020,7 @@ async def run_benchmark(args: argparse.Namespace) -> tuple[list[RequestTrace], l extra_body = json.loads(args.extra_request_body) if args.extra_request_body else {} shared_store = None - if args.return_routed_experts or args.return_indexer_topk: + if args.return_routed_experts: shared_store = init_shared_store() tokenizer = None @@ -1073,8 +1061,6 @@ async def run_benchmark(args: argparse.Namespace) -> tuple[list[RequestTrace], l print('return_token_ids=True') if args.return_routed_experts: print('return_routed_experts=True') - if args.return_indexer_topk: - print('return_indexer_topk=True') if args.return_logprob: print('return_logprob=True') if args.logprobs: @@ -1096,7 +1082,6 @@ async def send_one(request: BenchmarkRequest, mode: str, setting: float, repeat: ignore_eos=args.ignore_eos, return_token_ids=args.return_token_ids, return_routed_experts=args.return_routed_experts, - return_indexer_topk=args.return_indexer_topk, return_logprob=args.return_logprob, logprobs=args.logprobs, top_logprobs=args.top_logprobs, @@ -1292,11 +1277,6 @@ def parse_args() -> argparse.Namespace: action='store_true', help='Set return_routed_experts=true to include MoE routed expert indices (LMDeploy extension).', ) - parser.add_argument( - '--return-indexer-topk', - action='store_true', - help='Set return_indexer_topk=true and fetch sparse-attention indexer results (LMDeploy extension).', - ) parser.add_argument( '--return-logprob', action='store_true', diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index c3c5714758..9f7af90a32 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -115,7 +115,6 @@ def add_parser_api_server(): ArgumentHelper.dllm_denoising_steps(pt_group) ArgumentHelper.dllm_confidence_threshold(pt_group) ArgumentHelper.enable_return_routed_experts(pt_group) - ArgumentHelper.enable_return_indexer_topk(pt_group) ArgumentHelper.distributed_executor_backend(pt_group) ArgumentHelper.kernel_block_size(pt_group) ArgumentHelper.prefix_cache_state_budget(pt_group) @@ -269,7 +268,6 @@ def api_server(args): dllm_denoising_steps=args.dllm_denoising_steps, dllm_confidence_threshold=args.dllm_confidence_threshold, enable_return_routed_experts=args.enable_return_routed_experts, - enable_return_indexer_topk=args.enable_return_indexer_topk, distributed_executor_backend=args.distributed_executor_backend, ) else: diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index bcd8ffbda9..abf2197b59 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -787,15 +787,6 @@ def enable_return_routed_experts(parser): default=False, help='Whether to output routed expert ids for replay') - @staticmethod - def enable_return_indexer_topk(parser): - """Add argument return sparse-attention indexer top-k to parser.""" - - return parser.add_argument('--enable-return-indexer-topk', - action='store_true', - default=False, - help='Whether to output sparse-attention indexer top-k ids for replay') - @staticmethod def add_spec_group(parser): spec_group = parser.add_argument_group('Speculative decoding arguments') diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index 4e5824d409..547b780afe 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -146,7 +146,6 @@ class GenerationConfig: # router replay return_routed_experts: bool = False - return_indexer_topk: bool = False # ngram, generation would stop if latest [size] tokens are repeated for [threshold] times repetition_ngram_size: int = 0 @@ -498,7 +497,6 @@ class PytorchEngineConfig: logprobs_mode: str = None # router replay enable_return_routed_experts: bool = False - enable_return_indexer_topk: bool = False enable_transfer_obj_ref: bool = False # dllm @@ -599,7 +597,6 @@ class Response: last_hidden_state: torch.Tensor = None index: int = 0 routed_experts: Any = None - indexer_topk: Any = None cached_tokens: int = 0 def __str__(self): @@ -630,7 +627,6 @@ def _format_tensor(name: str, tensor: torch.Tensor | None) -> list[str]: fields.extend(_format_tensor('logits', self.logits)) fields.extend(_format_tensor('last_hidden_state', self.last_hidden_state)) fields.extend(_format_tensor('routed_experts', self.routed_experts)) - fields.extend(_format_tensor('indexer_topk', self.indexer_topk)) return '\n'.join(fields) def extend(self, other: 'Response') -> 'Response': @@ -657,7 +653,6 @@ def extend(self, other: 'Response') -> 'Response': self.logprobs = self.logprobs or [] self.logprobs += other.logprobs self.routed_experts = other.routed_experts - self.indexer_topk = other.indexer_topk return self @@ -744,7 +739,6 @@ class EngineOutput: cache_block_ids: list[int] | None = None req_metrics: RequestMetrics | None = None routed_experts: torch.Tensor = None - indexer_topk: torch.Tensor = None ce_loss: float = None diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index b6c7df080a..60996894a0 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -577,7 +577,6 @@ class MiscConfig: logprobs_mode: str = None dllm_config: DLLMConfig = None enable_return_routed_experts: bool = False - enable_return_indexer_topk: bool = False enable_chunked_prefill: bool = False @classmethod @@ -598,7 +597,6 @@ def from_engine_config(cls, engine_config: PytorchEngineConfig): logprobs_mode=engine_config.logprobs_mode, dllm_config=dllm_config, enable_return_routed_experts=engine_config.enable_return_routed_experts, - enable_return_indexer_topk=engine_config.enable_return_indexer_topk, enable_chunked_prefill=False, ) return misc_config diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py index 8988ff92ee..a09298ea77 100644 --- a/lmdeploy/pytorch/engine/engine.py +++ b/lmdeploy/pytorch/engine/engine.py @@ -61,9 +61,6 @@ class InferOutput: # expert ids routed_experts: np.ndarray = None - # sparse-attention indexer results - indexer_topk: 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 9f6b8cb154..d0eeb865f4 100644 --- a/lmdeploy/pytorch/engine/engine_instance.py +++ b/lmdeploy/pytorch/engine/engine_instance.py @@ -139,7 +139,7 @@ def __del__(self): def _get_extra_outputs(self, resp: Response, num_all_ids: int, gen_config: GenerationConfig = None): """Get extra outputs.""" - outputs = dict(routed_experts=None, indexer_topk=None) + outputs = dict(routed_experts=None) if resp.type not in [ResponseType.FINISH, ResponseType.CANCEL]: return outputs @@ -172,11 +172,6 @@ def _maybe_transfer(data): if routed_experts is not None: routed_experts = _validate_num_tokens('routed_experts', routed_experts) outputs['routed_experts'] = _maybe_transfer(routed_experts) - - indexer_topk = resp.data.get('indexer_topk', None) if resp.data else None - if indexer_topk is not None: - indexer_topk = _validate_num_tokens('indexer_topk', indexer_topk) - outputs['indexer_topk'] = _maybe_transfer(indexer_topk) return outputs async def _async_try_add_session(self, session_id: int): @@ -280,7 +275,6 @@ async def async_stream_infer(self, num_all_ids = prompt_ids_len + output_offset + num_ids extra_outputs = self._get_extra_outputs(resp, num_all_ids, gen_config=gen_config) routed_experts = extra_outputs.get('routed_experts', None) - indexer_topk = extra_outputs.get('indexer_topk', None) logger.debug(f'session[{session_id}] finish: num_out_ids={num_ids}.') yield EngineOutput(resp.type, @@ -289,7 +283,6 @@ async def async_stream_infer(self, cache_block_ids=cache_block_ids, req_metrics=req_metrics, routed_experts=routed_experts, - indexer_topk=indexer_topk, logprobs=logprobs, ce_loss=ce_loss) break diff --git a/lmdeploy/pytorch/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py index 32fd9c440b..8b39cab777 100644 --- a/lmdeploy/pytorch/engine/engine_loop.py +++ b/lmdeploy/pytorch/engine/engine_loop.py @@ -227,7 +227,6 @@ def _send_resp(self, out: InferOutput): cache_block_ids=out.cache_block_ids, req_metrics=out.req_metrics, routed_experts=out.routed_experts, - indexer_topk=out.indexer_topk, logprobs=logprobs, ce_loss=out.ce_loss)) @@ -316,18 +315,15 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'): logits = batched_outputs.logits all_routed_experts = batched_outputs.all_routed_experts - all_indexer_topk = batched_outputs.all_indexer_topk ce_loss = batched_outputs.ce_loss if model_inputs is not None and (model_inputs.is_chunk and not model_inputs.is_last_chunk): # chunk long context does not need to update seqs and outputs seq = running[0] seq.append_routed_experts(all_routed_experts) - seq.append_indexer_topk(all_indexer_topk) seq.append_logits(logits) seq.append_ce_loss(ce_loss, finish=False) self.scheduler.block_trie.cache_routed_experts_for_seq(seq) - self.scheduler.block_trie.cache_indexer_topk_for_seq(seq) return dict() new_token_timestamp = batched_outputs.new_token_timestamp @@ -342,7 +338,6 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'): model_inputs=model_inputs, delta=delta) self.scheduler.block_trie.cache_routed_experts(running) - self.scheduler.block_trie.cache_indexer_topk(running) # generate output outputs: dict[int, InferOutput] = dict() @@ -387,8 +382,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'): cache_block_ids=cache_block_ids, req_metrics=req_metrics, logprobs=cur_logprobs, - routed_experts=msg.routed_experts, - indexer_topk=msg.indexer_topk) + routed_experts=msg.routed_experts) outputs[session_id] = out if msg.return_ce_loss: diff --git a/lmdeploy/pytorch/engine/inputs_maker.py b/lmdeploy/pytorch/engine/inputs_maker.py index c002334775..eac5506f2f 100644 --- a/lmdeploy/pytorch/engine/inputs_maker.py +++ b/lmdeploy/pytorch/engine/inputs_maker.py @@ -574,9 +574,6 @@ def _need_logits(self): def _need_routed_experts(self): return any(seq.return_routed_experts for seq in self.result.running) - def _need_indexer_topk(self): - return any(getattr(seq, 'return_indexer_topk', False) for seq in self.result.running) - def _need_ce_loss(self): return any(seq.return_ce_loss for seq in self.result.running) @@ -600,7 +597,6 @@ def _build_payload(self): return_logits=self._need_logits(), extra_inputs=result.extra_inputs, return_routed_experts=self._need_routed_experts(), - return_indexer_topk=self._need_indexer_topk(), return_ce_loss=self._need_ce_loss(), ) diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py index bd0ac46314..5aee27eccd 100644 --- a/lmdeploy/pytorch/engine/model_agent/agent.py +++ b/lmdeploy/pytorch/engine/model_agent/agent.py @@ -98,7 +98,6 @@ class BatchedOutputs: new_token_timestamp: int = 0 extra_outputs: ExtraOutputs | None = None all_routed_experts: torch.Tensor | None = None - all_indexer_topk: torch.Tensor | None = None ce_loss: torch.Tensor | None = None def to_cpu(self): @@ -676,7 +675,6 @@ async def _step_postprocess_with_output(self, return_ce_loss: bool = False, seq_length: torch.Tensor = None, all_routed_experts: Any = None, - all_indexer_topk: Any = None, extra_inputs: ExtraInputs = None): """Step postprocess with output.""" rank = self.rank @@ -726,7 +724,6 @@ async def _step_postprocess_with_output(self, model_metas=model_metas, logprobs=logprobs, all_routed_experts=all_routed_experts, - all_indexer_topk=all_indexer_topk, extra_outputs=extra_outputs, ce_loss=ce_loss)) @@ -775,7 +772,6 @@ async def _async_step( stopping_criteria: StoppingCriteria = None, return_logits: bool = False, return_routed_experts: bool = False, - return_indexer_topk: bool = False, return_ce_loss: bool = False, extra_inputs: ExtraInputs = None, ): @@ -859,10 +855,6 @@ async def _async_step( all_routed_experts = output.get('all_routed_experts', None) else: all_routed_experts = None - if return_indexer_topk: - all_indexer_topk = output.get('all_indexer_topk', None) - else: - all_indexer_topk = None ( inputs, @@ -883,7 +875,6 @@ async def _async_step( return_ce_loss=return_ce_loss, seq_length=seq_length, all_routed_experts=all_routed_experts, - all_indexer_topk=all_indexer_topk, extra_inputs=extra_inputs, )) else: @@ -1121,13 +1112,11 @@ def _build_model(self): logger.debug(msg_with_rank(rank, 'build model.')) # for router replay enable_return_routed_experts = self.misc_config.enable_return_routed_experts and self.need_output - enable_return_indexer_topk = self.misc_config.enable_return_indexer_topk and self.need_output build_model_ctx = BuildModelContext(language_model_only=self.misc_config.language_model_only, dllm_config=self.misc_config.dllm_config, strategy_factory=self.strategy_factory, enable_return_routed_experts=enable_return_routed_experts, - enable_return_indexer_topk=enable_return_indexer_topk, quant_config=self.model_config.quant_config, fp32_lm_head=self.model_config.fp32_lm_head, tie_word_embeddings=self.model_config.tie_word_embeddings, diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py index a2d9fa3f08..b1c5a59e68 100644 --- a/lmdeploy/pytorch/messages.py +++ b/lmdeploy/pytorch/messages.py @@ -142,7 +142,6 @@ class SamplingParam: out_ce_loss: bool = False num_logprobs: int = -1 return_routed_experts: bool = False - return_indexer_topk: bool = False # ngram repetition_ngram_size: int = 0 @@ -240,7 +239,6 @@ def from_gen_config(cls, gen_config: GenerationConfig): out_ce_loss=gen_config.return_ppl, num_logprobs=logprobs, return_routed_experts=gen_config.return_routed_experts, - return_indexer_topk=gen_config.return_indexer_topk, repetition_ngram_size=repetition_ngram_size, repetition_ngram_threshold=repetition_ngram_threshold, ) @@ -591,35 +589,6 @@ def _get_pad_width(self, reserve_size: int): return ((0, reserve_size), (0, 0), (0, 0)) -class HistoryIndexerTopK(_HistoryDataBase): - """History of full-layer sparse-attention indexer results.""" - - def __init__(self, indexer_topk: np.ndarray = None, dtype: np.dtype = np.int32): - super().__init__(indexer_topk, dtype) - - def _create_empty_array(self, dtype): - return None - - def reserve(self, size: int): - """Grow geometrically without reserving the request token limit.""" - if self._data is None or len(self._data) >= size: - return - capacity = max(size, max(1, len(self._data) * 2)) - data = np.empty((capacity, *self._data.shape[1:]), dtype=self.dtype) - data[:self._num_real] = self._data[:self._num_real] - self._data = data - - def append(self, new_data: np.ndarray): - """Append indexer results.""" - new_data = np.asarray(new_data) - if new_data.ndim != 3: - raise ValueError(f'indexer_topk must be a 3D array, got shape {new_data.shape}.') - if self._data is not None and self._data.shape[1:] != new_data.shape[1:]: - raise ValueError( - f'indexer_topk shape changed from {self._data.shape[1:]} to {new_data.shape[1:]}.') - super().append(new_data) - - class HistoryLogits(_HistoryDataBase): """History logits.""" ALLOC_SIZE = 64 @@ -768,9 +737,6 @@ class SchedulerSequence: # for router replay all_routed_experts: HistoryRouterExperts = field(default_factory=HistoryRouterExperts) - # exact sparse-attention indexer results - all_indexer_topk: HistoryIndexerTopK = field(default_factory=HistoryIndexerTopK) - # logits all_logits: HistoryLogits = field(default_factory=HistoryLogits) @@ -875,28 +841,6 @@ def append_routed_experts(self, routed_experts: Tensor | np.ndarray): routed_experts = routed_experts.cpu().numpy() self.all_routed_experts.append(routed_experts) - @property - def return_indexer_topk(self) -> bool: - return self.sampling_param.return_indexer_topk - - @property - def indexer_topk(self) -> np.ndarray: - if not self.return_indexer_topk: - return None - - end = max(0, self.num_valid_ids - 1) - if 0 < end <= len(self.all_indexer_topk): - return self.all_indexer_topk.get_real()[:end] - return None - - def append_indexer_topk(self, indexer_topk: Tensor | np.ndarray): - """Append exact sparse-attention indexer results.""" - if not self.return_indexer_topk or indexer_topk is None: - return - if isinstance(indexer_topk, Tensor): - indexer_topk = indexer_topk.cpu().numpy() - self.all_indexer_topk.append(indexer_topk) - @property def num_history_ids(self): """Num history ids.""" diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py index 6be75f1b36..ff7164a2ea 100644 --- a/lmdeploy/pytorch/model_inputs.py +++ b/lmdeploy/pytorch/model_inputs.py @@ -465,7 +465,6 @@ class BuildModelContext: dllm_config: DLLMConfig = None strategy_factory: 'StrategyFactoryBase' = None enable_return_routed_experts: bool = False - enable_return_indexer_topk: bool = False quant_config: QuantizationConfig = field(default_factory=QuantizationConfig) fp32_lm_head: bool = False tie_word_embeddings: bool = False diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index f12b5a3d0e..32ed574163 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -43,11 +43,6 @@ def get_layer_indexer_type(config: Any, layer_idx: int | None) -> str: return indexer_types[layer_idx] -def get_full_indexer_layer_ids(config: Any) -> tuple[int, ...]: - """Return the physical layers represented by compact indexer output.""" - return tuple(idx for idx in range(config.num_hidden_layers) if get_layer_indexer_type(config, idx) == 'full') - - 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.'): @@ -397,11 +392,7 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic self.indexer_type = get_layer_indexer_type(config, layer_idx) self.indexer = None - self.indexer_output_idx = None if self.indexer_type == 'full': - full_layer_ids = get_full_indexer_layer_ids(config) - if layer_idx in full_layer_ids: - self.indexer_output_idx = full_layer_ids.index(layer_idx) 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): @@ -453,7 +444,6 @@ def forward( attn_metadata: Any = None, topk_indices_buffer: DSATopKIndicesBuffer | None = None, skip_topk: bool = False, - all_indexer_topk: torch.Tensor | None = None, ): """Rewrite of LlamaAttention.forward.""" dist_config = get_dist_manager().current_config() @@ -492,10 +482,6 @@ def forward( else: topk_indices = topk_indices_buffer.read(q_len, hidden_states.device) - # Shared layers reuse the previous result, so capture only full layers. - if all_indexer_topk is not None and self.indexer_output_idx is not None: - all_indexer_topk[:, self.indexer_output_idx, :].copy_(topk_indices) - attn_output = self.attn_fwd( query_states, key_states, @@ -559,7 +545,6 @@ def forward( topk_indices_buffer: DSATopKIndicesBuffer | None = None, skip_topk: bool = False, all_routed_experts: torch.Tensor | None = None, - all_indexer_topk: torch.Tensor | None = None, ) -> tuple[torch.FloatTensor, torch.FloatTensor]: if residual is None: @@ -576,7 +561,6 @@ def forward( attn_metadata=attn_metadata, topk_indices_buffer=topk_indices_buffer, skip_topk=skip_topk, - all_indexer_topk=all_indexer_topk, ) else: hidden_states = self.self_attn( @@ -642,7 +626,6 @@ def forward( attn_metadata: Any = None, inputs_embeds: torch.FloatTensor | None = None, all_routed_experts: torch.Tensor | None = None, - all_indexer_topk: torch.Tensor | None = None, ): """forward.""" if inputs_embeds is None: @@ -663,7 +646,6 @@ def forward( attn_metadata=attn_metadata, topk_indices_buffer=self.topk_indices_buffer, all_routed_experts=all_routed_experts, - all_indexer_topk=all_indexer_topk, ) hidden_states, _ = self.norm(hidden_states, residual) @@ -678,7 +660,6 @@ def forward_microbatch( attn_metadata: Any = None, inputs_embeds: torch.FloatTensor | None = None, all_routed_experts: torch.Tensor | None = None, - all_indexer_topk: torch.Tensor | None = None, ): """forward_microbatch.""" # DSA shared top-k indices are model-global; use normal forward until @@ -690,7 +671,6 @@ def forward_microbatch( attn_metadata=attn_metadata, inputs_embeds=inputs_embeds, all_routed_experts=all_routed_experts, - all_indexer_topk=all_indexer_topk, ) @@ -716,8 +696,6 @@ def __init__(self, self._load_buffers = dict() bm_ctx = get_build_model_context() self.enable_return_routed_experts = bm_ctx.enable_return_routed_experts - self.enable_return_indexer_topk = bm_ctx.enable_return_indexer_topk - self.num_indexer_layers = len(get_full_indexer_layer_ids(config)) def forward( self, @@ -728,7 +706,7 @@ def forward( inputs_embeds: torch.Tensor = None, **kwargs, ): - """Model forward with optional exact DSA indexer capture.""" + """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 @@ -740,12 +718,6 @@ def forward( torch.iinfo(torch.uint16).max, dtype=torch.uint16, ) - all_indexer_topk = None - if self.enable_return_indexer_topk: - # Axis 1 follows full indexer layers in physical layer order. - all_indexer_topk = position_ids.new_empty( - (num_tokens, self.num_indexer_layers, self.config.index_topk), dtype=torch.int32) - forward = self.model.forward_microbatch if step_ctx.enable_microbatch else self.model.forward hidden_states = forward( input_ids=input_ids, @@ -754,16 +726,10 @@ def forward( attn_metadata=attn_metadata, inputs_embeds=inputs_embeds, all_routed_experts=all_routed_experts, - all_indexer_topk=all_indexer_topk, ) - if all_routed_experts is None and all_indexer_topk is None: + if all_routed_experts is None: return hidden_states - outputs = dict(hidden_states=hidden_states) - if all_routed_experts is not None: - outputs['all_routed_experts'] = all_routed_experts - if all_indexer_topk is not None: - outputs['all_indexer_topk'] = all_indexer_topk - return outputs + return dict(hidden_states=hidden_states, all_routed_experts=all_routed_experts) def _load_weight_attention(self, name: str, loaded_weight: torch.Tensor, params_dict: dict[str, nn.Parameter], update_pe_mapping: list): diff --git a/lmdeploy/pytorch/models/utils/cudagraph.py b/lmdeploy/pytorch/models/utils/cudagraph.py index c6a9c605ea..92c0682989 100644 --- a/lmdeploy/pytorch/models/utils/cudagraph.py +++ b/lmdeploy/pytorch/models/utils/cudagraph.py @@ -367,6 +367,4 @@ def get_outputs_cudagraph(self, output_buffers: dict[str, torch.Tensor], input_i outputs['hidden_states'] = output_buffers['hidden_states'][:, :num_tokens] if output_buffers.get('all_routed_experts', None) is not None: outputs['all_routed_experts'] = output_buffers['all_routed_experts'][:num_tokens, ...].clone() - if output_buffers.get('all_indexer_topk', None) is not None: - outputs['all_indexer_topk'] = output_buffers['all_indexer_topk'][:num_tokens, ...].clone() return outputs diff --git a/lmdeploy/pytorch/paging/block_trie.py b/lmdeploy/pytorch/paging/block_trie.py index fda2f69742..c8c856a907 100644 --- a/lmdeploy/pytorch/paging/block_trie.py +++ b/lmdeploy/pytorch/paging/block_trie.py @@ -137,7 +137,6 @@ def __init__(self, state_ref_count: int = 0, state_access_time: float = 0.0, routed_experts: np.ndarray = None, - indexer_topk: np.ndarray = None, adapter_name: str = None): self.hash_key = hash_key self.block = block @@ -149,7 +148,6 @@ def __init__(self, self.state_ref_count = state_ref_count self.state_access_time = state_access_time self.routed_experts = routed_experts - self.indexer_topk = indexer_topk self.adapter_name = adapter_name self.children: dict[int, Node] = dict() self._parent: Node = None @@ -394,56 +392,6 @@ def cache_routed_experts(self, seqs: list[SchedulerSequence]): for seq in seqs: self.cache_routed_experts_for_seq(seq) - @staticmethod - def _get_indexer_topk_for_range(seq: SchedulerSequence, start: int, end: int): - """Get a copy of indexer results for a full token range, if present.""" - if not seq.return_indexer_topk: - return None - all_indexer_topk = seq.all_indexer_topk - if len(all_indexer_topk) < seq.num_history_ids or len(all_indexer_topk) < end: - return None - indexer_topk = all_indexer_topk.get_real() - if indexer_topk is None or len(indexer_topk) < end: - return None - return indexer_topk[start:end].copy() - - def _try_cache_node_indexer_topk(self, node: Node, seq: SchedulerSequence, start: int, end: int): - """Attach indexer results to a trie node when a sequence has them.""" - if node.indexer_topk is not None: - return - indexer_topk = self._get_indexer_topk_for_range(seq, start, end) - if indexer_topk is not None and len(indexer_topk) == end - start: - node.indexer_topk = indexer_topk - - def _append_matched_indexer_topk(self, seq: SchedulerSequence, nodes: list[Node], start: int): - """Replay cached indexer results for a matched trie range.""" - if not seq.return_indexer_topk or len(nodes) == 0: - return - if len(seq.all_indexer_topk) != start: - return - if any(node.indexer_topk is None or len(node.indexer_topk) != self.block_size for node in nodes): - return - for node in nodes: - seq.append_indexer_topk(node.indexer_topk) - - def cache_indexer_topk_for_seq(self, seq: SchedulerSequence): - """Enrich attached trie nodes with sparse-attention indexer results.""" - if not self.enable or not seq.return_indexer_topk: - return - node = seq.prefix_cache.last_shared_node - while node is not None and node.parent is not None: - end = node.num_matched - start = end - self.block_size - self._try_cache_node_indexer_topk(node, seq, start, end) - node = node.parent - - def cache_indexer_topk(self, seqs: list[SchedulerSequence]): - """Enrich trie nodes with indexer results from multiple sequences.""" - if not self.enable: - return - for seq in seqs: - self.cache_indexer_topk_for_seq(seq) - # Sparse SSM checkpoint index helpers. The lookup key is only a filter; # `_verify_state_checkpoint_node()` remains the correctness check. @@ -1079,9 +1027,6 @@ def _verify_state_checkpoint_node(self, seq: SchedulerSequence, node: Node, inde 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}') - if seq.return_indexer_topk and block_node.indexer_topk is None: - return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.REQUEST_MISMATCH, - reason=f'indexer results missing at block {idx}') matched_blocks.append(block_node.block) return StateCheckpointVerifyResult(StateCheckpointVerifyStatus.HIT, @@ -1141,7 +1086,6 @@ def _match_state_checkpoint(self, seq: SchedulerSequence): seq.logical_blocks.append(matched_blocks) seq.set_step(step) self._append_matched_routed_experts(seq, matched_nodes, init_num_matched) - self._append_matched_indexer_topk(seq, matched_nodes, init_num_matched) seq.prefix_cache.restore_state = node.state_idx seq.prefix_cache.restore_node = node seq.prefix_cache.last_shared_node = node @@ -1220,8 +1164,6 @@ def __match_success(node: Node): break if seq.return_routed_experts and child.routed_experts is None: break - if seq.return_indexer_topk and child.indexer_topk is None: - break matched_nodes.append(child) __match_success(child) @@ -1257,7 +1199,6 @@ def __clamp_match_step(match_step: int): seq.logical_blocks.append(matched_blocks) seq.set_step(num_matched) self._append_matched_routed_experts(seq, matched_nodes, init_num_matched) - self._append_matched_indexer_topk(seq, matched_nodes, init_num_matched) if self.requires_state_checkpoint: seq.prefix_cache.restore_state = curr.state_idx @@ -1341,21 +1282,18 @@ def allocate(self, seq: SchedulerSequence): # trie-owned block and release this sequence's duplicate block. node = child self._try_cache_node_routed_experts(node, seq, start, end) - self._try_cache_node_indexer_topk(node, seq, start, end) if block != node.block: free_blocks.append(block) logical_blocks[block_id] = node.block blocks.append(node.block) else: routed_experts = self._get_routed_experts_for_range(seq, start, end) - indexer_topk = self._get_indexer_topk_for_range(seq, start, end) node = Node(hash_key=hash_key, block=block, tokens=curr_tokens, num_matched=num_matched + block_size, extra_hashes=extra_hashes, routed_experts=routed_experts, - indexer_topk=indexer_topk, adapter_name=seq.adapter_name) node.parent = parent blocks.append(node.block) diff --git a/lmdeploy/pytorch/strategies/ar/sequence.py b/lmdeploy/pytorch/strategies/ar/sequence.py index df03d9d13f..0d39b7e0ae 100644 --- a/lmdeploy/pytorch/strategies/ar/sequence.py +++ b/lmdeploy/pytorch/strategies/ar/sequence.py @@ -35,7 +35,6 @@ def update_token_ids(self, model_meta: dict[str, Any] = None, mode: UpdateTokenMode = UpdateTokenMode.INPUTS, routed_experts: np.ndarray = None, - indexer_topk: np.ndarray = None, **kwargs): """Update token ids, old token ids will be added to history.""" # update history image nums @@ -49,7 +48,6 @@ def update_token_ids(self, num_valid = len(token_ids) # record cached expert ids self.append_routed_experts(routed_experts) - self.append_indexer_topk(indexer_topk) if mode == UpdateTokenMode.INPUTS: self.cached_tokens = 0 @@ -93,8 +91,6 @@ def set_step(self, step: int): # chunk long context might not have all routed experts if len(self.all_routed_experts) > step: self.all_routed_experts.resize(step) - if self.return_indexer_topk and len(self.all_indexer_topk) > step: - self.all_indexer_topk.resize(step) def cleanup(self): """Setup history meta after sequence stopped or cancelled.""" @@ -143,21 +139,13 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode if batched_outputs.all_routed_experts is not None: all_routed_experts = batched_outputs.all_routed_experts.split(num_tokens, dim=0) all_routed_experts = [experts.numpy() for experts in all_routed_experts] - all_indexer_topk = [None] * len(num_tokens) - if batched_outputs.all_indexer_topk is not None: - all_indexer_topk = batched_outputs.all_indexer_topk.split(num_tokens, dim=0) - all_indexer_topk = [topk.numpy() for topk in all_indexer_topk] update_mode = UpdateTokenMode.DECODE if is_decoding else UpdateTokenMode.PREFILL - for token, msg, stop, model_meta, routed_experts, indexer_topk in zip( - next_token_ids, running, stopped, model_metas, all_routed_experts, all_indexer_topk): + for token, msg, stop, model_meta, routed_experts in zip(next_token_ids, running, stopped, model_metas, + all_routed_experts): if msg.status != MessageStatus.RUNNING: continue # fill token - msg.update_token_ids(token, - model_meta=model_meta, - mode=update_mode, - routed_experts=routed_experts, - indexer_topk=indexer_topk) + msg.update_token_ids(token, model_meta=model_meta, mode=update_mode, routed_experts=routed_experts) if stop: msg.state.finish() diff --git a/lmdeploy/pytorch/strategies/ar_spec/sequence.py b/lmdeploy/pytorch/strategies/ar_spec/sequence.py index d7b4d05ce7..0de052fad1 100644 --- a/lmdeploy/pytorch/strategies/ar_spec/sequence.py +++ b/lmdeploy/pytorch/strategies/ar_spec/sequence.py @@ -48,16 +48,6 @@ def routed_experts(self) -> np.ndarray: return self.all_routed_experts.get_real()[:end] return None - @property - def indexer_topk(self) -> np.ndarray: - if not self.return_indexer_topk: - return None - - end = max(0, self.num_valid_ids - 1) - if 0 < end <= len(self.all_indexer_topk): - return self.all_indexer_topk.get_real()[:end] - return None - @property def generated_ids(self) -> np.ndarray: end = self.num_valid_ids @@ -76,8 +66,7 @@ def _update_token_ids_inputs(self, token_ids: np.ndarray): self.history_cache.append(token_ids) def _update_token_ids_prefill(self, token_ids: np.ndarray, draft_token_ids: np.ndarray, - stop_pos: int = -1, routed_experts: np.ndarray = None, - indexer_topk: np.ndarray = None): + stop_pos: int = -1, routed_experts: np.ndarray = None): """Update token ids for prefill.""" # back to last valid position self.history_cache.resize(self.num_valid_ids) @@ -85,7 +74,6 @@ def _update_token_ids_prefill(self, token_ids: np.ndarray, draft_token_ids: np.n num_valid = len(token_ids) self.history_cache.append(token_ids) self.append_routed_experts(routed_experts) - self.append_indexer_topk(indexer_topk) self._num_history_ids += self._num_token_ids self.num_new_tokens += num_valid self._num_valid_ids = self.num_history_ids + num_valid @@ -96,8 +84,7 @@ def _update_token_ids_prefill(self, token_ids: np.ndarray, draft_token_ids: np.n self.history_cache.append(draft_token_ids) def _update_token_ids_decode(self, token_ids: np.ndarray, draft_token_ids: np.ndarray, - stop_pos: int = -1, routed_experts: np.ndarray = None, - indexer_topk: np.ndarray = None): + stop_pos: int = -1, routed_experts: np.ndarray = None): """Update token ids for decode.""" # back to last valid position self.history_cache.resize(self.num_valid_ids) @@ -116,9 +103,6 @@ def _update_token_ids_decode(self, token_ids: np.ndarray, draft_token_ids: np.nd if routed_experts is not None: routed_experts = routed_experts[:num_valid] self.append_routed_experts(routed_experts) - if indexer_topk is not None: - indexer_topk = indexer_topk[:num_valid] - self.append_indexer_topk(indexer_topk) if stop_pos > -1: self._num_token_ids = 1 @@ -135,7 +119,6 @@ def update_token_ids(self, draft_token_ids: Tensor = None, mode: UpdateTokenMode = UpdateTokenMode.INPUTS, routed_experts: np.ndarray = None, - indexer_topk: np.ndarray = None, stop_pos: int = -1, **kwargs): """Update token ids, old token ids will be added to history.""" @@ -158,12 +141,10 @@ def update_token_ids(self, self._update_token_ids_inputs(token_ids) elif mode == UpdateTokenMode.PREFILL: self._update_token_ids_prefill(token_ids, draft_token_ids, - stop_pos=stop_pos, routed_experts=routed_experts, - indexer_topk=indexer_topk) + stop_pos=stop_pos, routed_experts=routed_experts) else: self._update_token_ids_decode(token_ids, draft_token_ids, - stop_pos=stop_pos, routed_experts=routed_experts, - indexer_topk=indexer_topk) + stop_pos=stop_pos, routed_experts=routed_experts) if model_meta is not None: self.model_meta = model_meta @@ -189,8 +170,6 @@ def set_step(self, step: int): # chunk long context might not have all routed experts if len(self.all_routed_experts) > step: self.all_routed_experts.resize(step) - if self.return_indexer_topk and len(self.all_indexer_topk) > step: - self.all_indexer_topk.resize(step) def cleanup(self): """Setup history meta after sequence stopped or cancelled.""" @@ -242,10 +221,6 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode if batched_outputs.all_routed_experts is not None: all_routed_experts = batched_outputs.all_routed_experts.split(num_tokens, dim=0) all_routed_experts = [experts.numpy() for experts in all_routed_experts] - all_indexer_topk = [None] * len(num_tokens) - if batched_outputs.all_indexer_topk is not None: - all_indexer_topk = batched_outputs.all_indexer_topk.split(num_tokens, dim=0) - all_indexer_topk = [topk.numpy() for topk in all_indexer_topk] batch_size = len(running) next_token_ids = next_token_ids.view(batch_size, -1).numpy() @@ -258,7 +233,6 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode for idx, token in enumerate(next_token_ids): routed_experts = all_routed_experts[idx] - indexer_topk = all_indexer_topk[idx] msg = running[idx] stop = stopped[idx] model_meta = model_metas[idx] @@ -271,7 +245,6 @@ def update_running(self, running: SeqList, batched_outputs: BatchedOutputs, mode model_meta=model_meta, mode=update_mode, routed_experts=routed_experts, - indexer_topk=indexer_topk, stop_pos=stop_pos[idx]) if stop: msg.state.finish() diff --git a/lmdeploy/serve/anthropic/adapter.py b/lmdeploy/serve/anthropic/adapter.py index 26e4e6b65b..cf85f777c7 100644 --- a/lmdeploy/serve/anthropic/adapter.py +++ b/lmdeploy/serve/anthropic/adapter.py @@ -356,7 +356,6 @@ def to_generation_config( skip_special_tokens=True, spaces_between_special_tokens=True, return_routed_experts=request.return_routed_experts or False, - return_indexer_topk=request.return_indexer_topk or False, logprobs=1 if request.return_logprob else None, ) diff --git a/lmdeploy/serve/anthropic/endpoints/messages.py b/lmdeploy/serve/anthropic/endpoints/messages.py index 3e93f89025..4a20b2274e 100644 --- a/lmdeploy/serve/anthropic/endpoints/messages.py +++ b/lmdeploy/serve/anthropic/endpoints/messages.py @@ -58,12 +58,6 @@ def _validate_extended_outputs(request: MessagesRequest, server_context): ('routed experts requested but not configured in engine configuration. ' 'May start the api_server with --enable-return-routed-experts flag.')) - if request.return_indexer_topk and not engine_config.enable_return_indexer_topk: - return create_error_response( - HTTPStatus.BAD_REQUEST, - ('indexer top-k requested but not configured in engine configuration. ' - 'May start the api_server with --enable-return-indexer-topk flag.')) - return None @@ -204,7 +198,6 @@ async def create_message(request: MessagesRequest, raw_request: Request): response_parser=response_parser, return_token_ids=request.return_token_ids or False, return_routed_experts=request.return_routed_experts or False, - return_indexer_topk=request.return_indexer_topk or False, logprobs=request.return_logprob or False, ), [result_generator], @@ -243,7 +236,7 @@ async def create_message(request: MessagesRequest, raw_request: Request): return create_error_response(HTTPStatus.BAD_REQUEST, f'Failed to parse output: {err}') should_validate_complete = ( final_res.finish_reason in ('stop', 'length') - and (request.return_token_ids or request.return_routed_experts or request.return_indexer_topk) + and (request.return_token_ids or request.return_routed_experts) ) if should_validate_complete and not response_parser.validate_complete(raw_text): final_res.finish_reason = 'parse_error' @@ -274,6 +267,5 @@ async def create_message(request: MessagesRequest, raw_request: Request): output_ids=final_token_ids if request.return_token_ids else None, output_token_logprobs=output_token_logprobs, routed_experts=final_res.routed_experts if request.return_routed_experts else None, - indexer_topk=final_res.indexer_topk if request.return_indexer_topk else None, ) return response.model_dump() diff --git a/lmdeploy/serve/anthropic/protocol.py b/lmdeploy/serve/anthropic/protocol.py index e3ac1a7a1c..37fa0f7bd7 100644 --- a/lmdeploy/serve/anthropic/protocol.py +++ b/lmdeploy/serve/anthropic/protocol.py @@ -10,8 +10,6 @@ from pydantic import BaseModel, ConfigDict, Field RoutedExperts = list[list[list[int]]] | str | None -# (num_token, num_full_indexer_layer, index_topk), ordered by physical layer. -IndexerTopK = list[list[list[int]]] | str | None MessageStopReason = Literal['end_turn', 'max_tokens', 'stop_sequence', 'tool_use', 'parse_error'] @@ -132,10 +130,6 @@ class MessagesRequest(BaseModel): default=False, description=('Whether to return MoE routed expert indices in the response.'), ) - return_indexer_topk: bool | None = Field( - default=False, - description=('Whether to return sparse-attention top-k indices for full indexer layers.'), - ) return_token_ids: bool | None = Field( default=False, description=('Whether to include output token IDs in the response.'), @@ -200,7 +194,6 @@ class MessagesResponse(BaseModel): output_ids: list[int] | None = None output_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) routed_experts: RoutedExperts = None - indexer_topk: IndexerTopK = None class StreamTextBlock(BaseModel): @@ -319,7 +312,6 @@ class MessageDeltaEvent(BaseModel): delta: MessageDelta usage: MessageDeltaUsage routed_experts: RoutedExperts = None - indexer_topk: IndexerTopK = None class MessageStopEvent(BaseModel): diff --git a/lmdeploy/serve/anthropic/streaming.py b/lmdeploy/serve/anthropic/streaming.py index 5bb0139f98..934c4b0158 100644 --- a/lmdeploy/serve/anthropic/streaming.py +++ b/lmdeploy/serve/anthropic/streaming.py @@ -31,7 +31,7 @@ ThinkingDelta, ) -_OPTIONAL_EXTENSION_FIELDS = ('output_ids', 'output_token_logprobs', 'routed_experts', 'indexer_topk') +_OPTIONAL_EXTENSION_FIELDS = ('output_ids', 'output_token_logprobs', 'routed_experts') def _format_sse(data: AnthropicStreamEvent) -> str: @@ -169,7 +169,6 @@ async def stream_messages_response(result_generator, response_parser, return_token_ids: bool = False, return_routed_experts: bool = False, - return_indexer_topk: bool = False, logprobs: bool = False) -> AsyncGenerator[str, None]: """Convert LMDeploy generation stream to Anthropic SSE events.""" @@ -227,7 +226,7 @@ async def _results(): should_validate_complete = ( res.finish_reason in ('stop', 'length') - and (return_token_ids or return_routed_experts or return_indexer_topk) + and (return_token_ids or return_routed_experts) ) if should_validate_complete and not response_parser.validate_complete(): res.finish_reason = 'parse_error' @@ -316,12 +315,10 @@ async def _results(): output_tokens = 0 if final_res is None else final_res.generate_token_len stop_reason = map_finish_reason(None if final_res is None else final_res.finish_reason) routed_experts = final_res.routed_experts if return_routed_experts and final_res is not None else None - indexer_topk = final_res.indexer_topk if return_indexer_topk and final_res is not None else None yield _format_sse( MessageDeltaEvent( delta=MessageDelta(stop_reason=stop_reason, stop_sequence=None), usage=MessageDeltaUsage(output_tokens=output_tokens), routed_experts=routed_experts, - indexer_topk=indexer_topk, )) yield _format_sse(MessageStopEvent()) diff --git a/lmdeploy/serve/core/async_engine.py b/lmdeploy/serve/core/async_engine.py index 9a90745d5f..313515ceba 100644 --- a/lmdeploy/serve/core/async_engine.py +++ b/lmdeploy/serve/core/async_engine.py @@ -53,7 +53,6 @@ class GenOut: last_hidden_state: Any = None cache_block_ids: list[int] | None = None # for disaggregation routed_experts: Any = None # for RL router replay - indexer_topk: Any = None # for sparse-attention indexer replay cached_tokens: int = 0 def to_response(self, index: int = 0) -> Response: @@ -71,7 +70,6 @@ def to_response(self, index: int = 0) -> Response: last_hidden_state=self.last_hidden_state, logits=self.logits, routed_experts=self.routed_experts, - indexer_topk=self.indexer_topk, cached_tokens=self.cached_tokens, index=index) @@ -681,7 +679,6 @@ def is_error(status): finish_reason, token_ids=res, routed_experts=outputs.routed_experts, - indexer_topk=outputs.indexer_topk, cache_block_ids=outputs.cache_block_ids, cached_tokens=cached_tokens) if outputs.logprobs is not None: @@ -722,10 +719,6 @@ def is_error(status): if routed_experts is not None and not isinstance(routed_experts, str) and ( not gen_config.include_stop_str_in_output) and finish_reason == 'stop': routed_experts = routed_experts[:-1] - indexer_topk = outputs.indexer_topk - if indexer_topk is not None and not isinstance(indexer_topk, str) and ( - not gen_config.include_stop_str_in_output) and finish_reason == 'stop': - indexer_topk = indexer_topk[:-1] logger.info(f'session {session_id} finished, reason ' f'"{finish_reason}", input_tokens ' @@ -739,7 +732,6 @@ def is_error(status): logits=logits, last_hidden_state=last_hidden_state, routed_experts=routed_experts, - indexer_topk=indexer_topk, cache_block_ids=outputs.cache_block_ids, cached_tokens=cached_tokens) else: diff --git a/lmdeploy/serve/openai/api_server.py b/lmdeploy/serve/openai/api_server.py index aca78e4882..7098a64ca4 100644 --- a/lmdeploy/serve/openai/api_server.py +++ b/lmdeploy/serve/openai/api_server.py @@ -553,7 +553,6 @@ def create_stream_response_json(index: int, logprobs: ChoiceLogprobs | None = None, output_token_logprobs: list[tuple[float, int]] | None = None, routed_experts=None, - indexer_topk=None, output_ids=None) -> dict: choice_data = ChatCompletionResponseStreamChoice(index=index, delta=delta_message, @@ -561,8 +560,7 @@ def create_stream_response_json(index: int, logprobs=logprobs, output_token_logprobs=output_token_logprobs, output_ids=output_ids, - routed_experts=routed_experts, - indexer_topk=indexer_topk) + routed_experts=routed_experts) choice_data = maybe_filter_parallel_tool_calls(choice_data, request) response = ChatCompletionStreamResponse( id=request_id, @@ -616,7 +614,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: stream_deltas = [(DeltaMessage(role='assistant', content=''), False)] should_validate_complete = ( res.finish_reason in ('stop', 'length') - and (request.return_token_ids or request.return_routed_experts or request.return_indexer_topk) + and (request.return_token_ids or request.return_routed_experts) ) if should_validate_complete and not response_parser.validate_complete(): res.finish_reason = 'parse_error' @@ -638,7 +636,6 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: # Only output routed_experts in the final chunk routed_experts = res.routed_experts if finish_reason is not None else None - indexer_topk = res.indexer_topk if finish_reason is not None else None # Emit token ids once per engine yield on the last parsed delta, when # accumulated delta text and token ids for this step are aligned. stream_output_ids = delta_token_ids if (request.return_token_ids and is_last_delta) else None @@ -649,7 +646,6 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: logprobs=chunk_logprobs, output_token_logprobs=chunk_output_token_logprobs, routed_experts=routed_experts, - indexer_topk=indexer_topk, output_ids=stream_output_ids) if res.cache_block_ids is not None and is_last_delta: response_json['cache_block_ids'] = res.cache_block_ids @@ -694,7 +690,7 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: text, tool_calls, reasoning_content = response_parser.parse_complete(text, final_token_ids) should_validate_complete = ( final_res.finish_reason in ('stop', 'length') - and (request.return_token_ids or request.return_routed_experts or request.return_indexer_topk) + and (request.return_token_ids or request.return_routed_experts) ) if should_validate_complete and not response_parser.validate_complete(raw_text): final_res.finish_reason = 'parse_error' @@ -728,7 +724,6 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: finish_reason=final_res.finish_reason, output_ids=final_token_ids if request.return_token_ids else None, routed_experts=final_res.routed_experts if request.return_routed_experts else None, - indexer_topk=final_res.indexer_topk if request.return_indexer_topk else None, ) choice_data = maybe_filter_parallel_tool_calls(choice_data, request) choices.append(choice_data) @@ -1028,16 +1023,13 @@ def create_generate_response_json(res, output_ids, logprobs, finish_reason, - routed_experts=None, - indexer_topk=None): + routed_experts=None): # only output router experts in last chunk routed_experts = None if finish_reason is None else routed_experts - indexer_topk = None if finish_reason is None else indexer_topk meta = GenerateReqMetaOutput(finish_reason=dict(type=finish_reason) if finish_reason else None, output_token_logprobs=logprobs or None, prompt_tokens=res.input_token_len, routed_experts=routed_experts, - indexer_topk=indexer_topk, completion_tokens=res.generate_token_len) response = GenerateReqOutput(text=text, output_ids=output_ids, meta_info=meta) @@ -1048,7 +1040,6 @@ async def generate_stream_generator(): text = res.response or '' output_ids = res.token_ids routed_experts = res.routed_experts - indexer_topk = res.indexer_topk logprobs = [] if res.logprobs: for tok, tok_logprobs in zip(res.token_ids, res.logprobs): @@ -1058,8 +1049,7 @@ async def generate_stream_generator(): output_ids, logprobs, res.finish_reason, - routed_experts=routed_experts, - indexer_topk=indexer_topk) + routed_experts=routed_experts) yield f'data: {response_json}\n\n' yield 'data: [DONE]\n\n' @@ -1093,7 +1083,6 @@ async def _inner_call(): output_token_logprobs=output_token_logprobs or None, prompt_tokens=res.input_token_len, routed_experts=res.routed_experts, - indexer_topk=res.indexer_topk, completion_tokens=res.generate_token_len) response = GenerateReqOutput(text=text, output_ids=output_ids, meta_info=meta) @@ -1622,11 +1611,6 @@ def serve(model_path: str, # router replay if backend_config.enable_return_routed_experts: backend_config.enable_transfer_obj_ref = True - if backend_config.enable_return_indexer_topk: - if backend_config.distributed_executor_backend != 'ray': - raise ValueError('--enable-return-indexer-topk requires ' - '--distributed-executor-backend ray for HTTP serving.') - backend_config.enable_transfer_obj_ref = True VariableInterface.async_engine = pipeline_class(model_path=model_path, model_name=model_name, backend=backend, diff --git a/lmdeploy/serve/openai/protocol.py b/lmdeploy/serve/openai/protocol.py index d0cb0b9cda..fd6f1370f2 100644 --- a/lmdeploy/serve/openai/protocol.py +++ b/lmdeploy/serve/openai/protocol.py @@ -226,10 +226,6 @@ class ChatCompletionRequest(BaseModel): default=False, description=('Whether to return MoE routed expert indices in the response.'), ) - return_indexer_topk: bool | None = Field( - default=False, - description=('Whether to return sparse-attention top-k indices for full indexer layers.'), - ) class FunctionCall(BaseModel): @@ -297,7 +293,6 @@ class ChatCompletionResponseChoice(BaseModel): finish_reason: Literal['stop', 'length', 'tool_calls', 'parse_error', 'error', 'abort'] | None = None output_ids: list[int] | None = None routed_experts: list[list[list[int]]] | str | None = None - indexer_topk: list[list[list[int]]] | str | None = None class ChatCompletionResponse(BaseModel): @@ -340,7 +335,6 @@ class ChatCompletionResponseStreamChoice(BaseModel): output_ids: list[int] | None = None finish_reason: Literal['stop', 'length', 'tool_calls', 'parse_error', 'error', 'abort'] | None = None routed_experts: list[list[list[int]]] | str | None = None - indexer_topk: list[list[list[int]]] | str | None = None class ChatCompletionStreamResponse(BaseModel): @@ -565,7 +559,6 @@ class GenerateReqInput(BaseModel): spaces_between_special_tokens: bool | None = True include_stop_str_in_output: bool | None = False return_routed_experts: bool | None = False - return_indexer_topk: bool | None = False repetition_ngram_size: int = Field(default=0, ge=0) repetition_ngram_threshold: int = Field(default=0, ge=0) # kwargs for media IO @@ -586,7 +579,6 @@ class GenerateReqMetaOutput(BaseModel): finish_reason: dict[str, Any] | None = None output_token_logprobs: list[tuple[float, int]] | None = None # (logprob, token_id) routed_experts: list[list[list[int]]] | str | None = None # (num_token, num_layer, topk_expert) - indexer_topk: list[list[list[int]]] | str | None = None # (num_token, num_full_indexer_layer, index_topk) # /generate output diff --git a/lmdeploy/serve/openai/serving_chat_completion.py b/lmdeploy/serve/openai/serving_chat_completion.py index a5e9e5d741..0270dd8546 100644 --- a/lmdeploy/serve/openai/serving_chat_completion.py +++ b/lmdeploy/serve/openai/serving_chat_completion.py @@ -65,8 +65,4 @@ def check_request(request: ChatCompletionRequest, server_context: 'VariableInter return ('routed experts requested but not configured in engine configuration. ' 'May start api_server with --enable-return-routed-experts flag.') - if request.return_indexer_topk and not engine_config.enable_return_indexer_topk: - return ('indexer top-k requested but not configured in engine configuration. ' - 'May start api_server with --enable-return-indexer-topk flag.') - return '' diff --git a/lmdeploy/serve/openai/serving_generate.py b/lmdeploy/serve/openai/serving_generate.py index 66a79bbde7..05ffb32b5b 100644 --- a/lmdeploy/serve/openai/serving_generate.py +++ b/lmdeploy/serve/openai/serving_generate.py @@ -22,9 +22,6 @@ def check_request(request: GenerateReqInput, server_context: 'VariableInterface' 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.return_indexer_topk and not engine_config.enable_return_indexer_topk: - return ('indexer top-k requested but not configured in engine configuration. ' - 'May start api_server with --enable-return-indexer-topk 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/tests/pytorch/engine/test_engine_sleep.py b/tests/pytorch/engine/test_engine_sleep.py index 6a9e45547f..f6efdee5cf 100644 --- a/tests/pytorch/engine/test_engine_sleep.py +++ b/tests/pytorch/engine/test_engine_sleep.py @@ -1,12 +1,10 @@ # Copyright (c) OpenMMLab. All rights reserved. import asyncio -import sys from types import SimpleNamespace -import numpy as np import pytest -from lmdeploy.messages import EngineOutput, GenerationConfig, ResponseType +from lmdeploy.messages import EngineOutput, ResponseType from lmdeploy.pytorch.engine.engine import Engine from lmdeploy.pytorch.engine.engine_instance import EngineInstance from lmdeploy.pytorch.engine.executor.mp_executor import MPExecutor @@ -207,33 +205,3 @@ async def __collect(): assert len(outputs) == 1 assert outputs[0].status == ResponseType.CANCEL assert engine.req_manager._loop_task is None - - -def test_engine_instance_transfers_trimmed_indexer_topk(monkeypatch): - import lmdeploy.pytorch.engine.engine_instance as engine_instance_module - - transferred = [] - - class _Put: - - def remote(self, data): - transferred.append(data.copy()) - return f'key-{len(transferred)}' - - store = SimpleNamespace(put=_Put()) - monkeypatch.setattr(engine_instance_module, '_SHARED_STORE', store) - monkeypatch.setitem(sys.modules, 'ray', SimpleNamespace(get=lambda key: key)) - - instance = SimpleNamespace(_enable_transfer_obj_ref=True) - indexer_topk = np.arange(5 * 2 * 3, dtype=np.int32).reshape(5, 2, 3) - resp = Response(type=ResponseType.FINISH, - sender_id=0, - event=None, - data=dict(token_ids=np.array([10, 11]), indexer_topk=indexer_topk)) - gen_config = GenerationConfig(stop_token_ids=[11], include_stop_str_in_output=False) - - outputs = EngineInstance._get_extra_outputs(instance, resp, num_all_ids=6, gen_config=gen_config) - - assert outputs['indexer_topk'] == 'key-1' - assert transferred[0].shape == (4, 2, 3) - assert np.array_equal(transferred[0], indexer_topk[:-1]) diff --git a/tests/pytorch/paging/test_block_trie.py b/tests/pytorch/paging/test_block_trie.py index c7fbd4b6d1..e94210d8b1 100644 --- a/tests/pytorch/paging/test_block_trie.py +++ b/tests/pytorch/paging/test_block_trie.py @@ -110,10 +110,6 @@ def _routed_experts(self, num_tokens: int, offset: int = 0): values = np.arange(offset, offset + num_tokens * 2, dtype=np.uint16) return values.reshape(num_tokens, 2, 1) - def _indexer_topk(self, num_tokens: int, offset: int = 0): - values = np.arange(offset, offset + num_tokens * 6, dtype=np.int32) - return values.reshape(num_tokens, 2, 3) - def _add_ready_ssm_checkpoint(self, scheduler, token_ids): seq = scheduler.add_session(len(scheduler.sessions)).add_sequence(token_ids) scheduler.block_manager.allocate(seq) @@ -450,46 +446,6 @@ def test_match_skips_routed_expert_replay_when_not_requested(self, block_trie, b assert matched.num_history_ids == block_size * 2 assert len(matched.all_routed_experts) == 0 - def test_match_replays_cached_indexer_topk(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] - sampling_param = SamplingParam(return_indexer_topk=True) - seq = sess.add_sequence(token_ids, sampling_param=sampling_param) - indexer_topk = self._indexer_topk(block_size * 2) - - block_mgr.allocate(seq) - block_trie.allocate(seq) - seq.append_indexer_topk(indexer_topk) - block_trie.cache_indexer_topk_for_seq(seq) - - matched = sess.add_sequence(token_ids, sampling_param=sampling_param) - block_trie.match(matched) - - assert matched.num_history_ids == block_size * 2 - assert np.array_equal(matched.all_indexer_topk.get_real(), indexer_topk) - - def test_match_stops_before_block_missing_indexer_topk(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] - sampling_param = SamplingParam(return_indexer_topk=True) - seq = sess.add_sequence(token_ids, sampling_param=sampling_param) - first_block_topk = self._indexer_topk(block_size) - - block_mgr.allocate(seq) - block_trie.allocate(seq) - seq.append_indexer_topk(first_block_topk) - block_trie.cache_indexer_topk_for_seq(seq) - - matched = sess.add_sequence(token_ids, sampling_param=sampling_param) - block_trie.match(matched) - - assert matched.num_history_ids == block_size - assert np.array_equal(matched.all_indexer_topk.get_real(), first_block_topk) - assert matched.prefix_cache.private_recompute_start_step == block_size - assert matched.prefix_cache.private_recompute_end_step == block_size * 2 - def test_existing_node_can_be_enriched_with_routed_experts(self, block_trie, block_mgr, scheduler): sess = scheduler.add_session(0) block_size = sess.seq_meta.block_size diff --git a/tests/pytorch/spec_decode/test_strategies.py b/tests/pytorch/spec_decode/test_strategies.py index f42fd69912..af1e827299 100644 --- a/tests/pytorch/spec_decode/test_strategies.py +++ b/tests/pytorch/spec_decode/test_strategies.py @@ -540,55 +540,6 @@ def _experts(n, k=2): return np.arange(n * k, dtype=np.uint16).reshape(n, 1, k) -def _make_seq_with_indexer(prefill_tokens=None): - """Create a speculative sequence that captures sparse-attention indices.""" - strategy = ARSpecSequenceStrategy() - seq_meta = SequenceMeta(block_size=16, strategy=strategy) - session = MagicMock() - session.seq_meta = seq_meta - sampling_param = SamplingParam(return_indexer_topk=True, - max_new_tokens=512) - seq = SchedulerSequenceARSpec(seq_id=0, session=session, sampling_param=sampling_param) - if prefill_tokens is not None: - seq._update_token_ids_inputs(np.array(prefill_tokens, dtype=np.int64)) - return seq - - -def _indexer_topk(n): - return np.arange(n * 6, dtype=np.int32).reshape(n, 2, 3) - - -class TestIndexerTopKSpecDecode: - - def test_decode_keeps_only_accepted_indexer_rows(self): - seq = _make_seq_with_indexer() - seq._num_valid_ids = 3 - seq._num_history_ids = 2 - seq.history_cache.append(np.array([0, 1, 2, 100, 101], dtype=np.int64)) - seq._num_token_ids = 3 - - seq._update_token_ids_decode( - np.array([30, 40, -1]), - draft_token_ids=np.array([], dtype=np.int64), - indexer_topk=_indexer_topk(3), - ) - - assert len(seq.all_indexer_topk) == 2 - assert np.array_equal(seq.all_indexer_topk.get_real(), _indexer_topk(2)) - - def test_indexer_history_grows_geometrically_and_truncates_logically(self): - seq = _make_seq_with_indexer([1, 2, 3, 4]) - seq.append_indexer_topk(_indexer_topk(4)) - - assert len(seq.all_indexer_topk._data) == 4 - seq.append_indexer_topk(_indexer_topk(2)) - assert len(seq.all_indexer_topk) == 6 - assert len(seq.all_indexer_topk._data) == 8 - seq.set_step(3) - assert len(seq.all_indexer_topk) == 3 - assert len(seq.all_indexer_topk._data) == 8 - - # --------------------------------------------------------------------------- # Tests for routed_experts in _update_token_ids_decode # --------------------------------------------------------------------------- diff --git a/tests/test_lmdeploy/serve/test_generation_config.py b/tests/test_lmdeploy/serve/test_generation_config.py index c8aa915997..5ccc16f06a 100644 --- a/tests/test_lmdeploy/serve/test_generation_config.py +++ b/tests/test_lmdeploy/serve/test_generation_config.py @@ -21,7 +21,6 @@ class _FakeEngineConfig: logprobs_mode = None enable_return_routed_experts = False - enable_return_indexer_topk = False class _FakeSessionManager: @@ -89,15 +88,6 @@ def test_build_generation_config_from_merged_values(): assert gen_config.do_sample is True -def test_indexer_topk_request_flag_reaches_generation_config(): - chat_request = ChatCompletionRequest(model='test', messages='hi', return_indexer_topk=True) - generate_request = GenerateReqInput(prompt='hi', return_indexer_topk=True) - - assert build_generation_config(chat_request, {}).return_indexer_topk is True - assert build_generation_config(generate_request, {}).return_indexer_topk is True - assert 'indexer top-k requested' in check_generate_request(generate_request, _FakeServerContext()) - - def test_build_generation_config_max_new_tokens_defaults_to_none(): request = CompletionRequest(model='test', prompt='hello') gen_config = build_generation_config(request, {}) From d5f844b548c333836495446cedeca33afff9ea4a Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 24 Jul 2026 11:37:16 +0800 Subject: [PATCH 27/33] chore: clean up indexer replay leftovers --- benchmark/benchmark_chat_completion.py | 12 ++++-------- tests/pytorch/spec_decode/test_spec_agent.py | 13 ++++++++----- tests/test_lmdeploy/serve/test_generation_config.py | 1 - 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/benchmark/benchmark_chat_completion.py b/benchmark/benchmark_chat_completion.py index b9bc0c9e04..bba0c05289 100644 --- a/benchmark/benchmark_chat_completion.py +++ b/benchmark/benchmark_chat_completion.py @@ -7,7 +7,7 @@ aggregates TTFT/ITL/TPOT metrics, and writes table plus report artifacts for concurrency/RPS sweeps. Generation options include ``--output-tokens`` (``max_completion_tokens``), -``--ignore-eos``, ``--return-token-ids``, ``--return-routed-experts``, ``--return-indexer-topk``, +``--ignore-eos``, ``--return-token-ids``, ``--return-routed-experts``, ``--return-logprob``, and ``--logprobs`` / ``--top-logprobs``. """ @@ -131,19 +131,15 @@ def init_shared_store() -> Any: return _shared_store_actor -async def fetch_shared_output(shared_store: Any, key: str) -> Any: - """Fetch a large replay output without blocking the event loop.""" +async def fetch_routed_experts(shared_store: Any, key: str) -> Any: + """Fetch routed_experts from shared_store without blocking the event + loop.""" import ray ref = shared_store.get.remote(key) return await asyncio.to_thread(ray.get, ref) -async def fetch_routed_experts(shared_store: Any, key: str) -> Any: - """Backward-compatible routed-expert fetch helper.""" - return await fetch_shared_output(shared_store, key) - - def _split_csv(value: str | None) -> list[str] | None: if value is None: return None diff --git a/tests/pytorch/spec_decode/test_spec_agent.py b/tests/pytorch/spec_decode/test_spec_agent.py index 7feda2c32a..2dcef80bf4 100644 --- a/tests/pytorch/spec_decode/test_spec_agent.py +++ b/tests/pytorch/spec_decode/test_spec_agent.py @@ -275,27 +275,30 @@ def test_spec_model_agent_method_when_enabled(): assert agent.method == specdecode_config.method -def test_matching_draft_dist_config_reuses_main_context(monkeypatch): +def test_qwen35_mtp_reuses_main_dist_context(monkeypatch): + """Qwen3.5 MTP mirrors the target topology, so it should share groups.""" from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig from lmdeploy.pytorch.distributed import DistContext from lmdeploy.pytorch.spec_decode import base as base_mod dist_config = DistConfig(dp=2, ep=2) dist_ctx = DistContext(rank=1, dp_rank=1, dist_config=dist_config, ep_gpu_group=object()) - specdecode_config = SpecDecodeConfig(model='glm-model', - method='deepseek_mtp', + specdecode_config = SpecDecodeConfig(model='draft-model', + method='qwen3_5_mtp', dist_config=DistConfig(dp=2, ep=2), num_speculative_tokens=3) def fail_build(*args, **kwargs): - raise AssertionError('Matching topology should reuse the target DistContext') + raise AssertionError('qwen3_5_mtp should not build a separate draft DistContext') monkeypatch.setattr(base_mod.DistContext, 'build', staticmethod(fail_build)) assert base_mod._build_draft_dist_ctx(dist_ctx, specdecode_config) is dist_ctx -def test_different_draft_dist_config_builds_draft_context(monkeypatch): +def test_non_qwen35_mtp_builds_draft_dist_context(monkeypatch): + """Other speculative methods keep their separate draft distribution + path.""" from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig from lmdeploy.pytorch.distributed import DistContext from lmdeploy.pytorch.spec_decode import base as base_mod diff --git a/tests/test_lmdeploy/serve/test_generation_config.py b/tests/test_lmdeploy/serve/test_generation_config.py index 5ccc16f06a..2c39ba5fb5 100644 --- a/tests/test_lmdeploy/serve/test_generation_config.py +++ b/tests/test_lmdeploy/serve/test_generation_config.py @@ -20,7 +20,6 @@ class _FakeEngineConfig: logprobs_mode = None - enable_return_routed_experts = False class _FakeSessionManager: From 471daeb6d08deba902988e4d4a0f24f571ce7738 Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 24 Jul 2026 14:09:27 +0800 Subject: [PATCH 28/33] test: clean up merge resolutions --- .../backends/dlinfer/apply_rotary_emb.py | 11 +- .../pytorch/kernel/test_sparse_index_topk.py | 123 +++++++----------- 2 files changed, 46 insertions(+), 88 deletions(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py b/lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py index 89c938aded..2316d3afdd 100644 --- a/lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py +++ b/lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py @@ -17,16 +17,7 @@ def forward(self, inplace: bool = True, complex_mode: bool = False): """forward.""" - if complex_mode: - from ..default.apply_rotary_emb import DefaultApplyRotaryEmbImpl - return DefaultApplyRotaryEmbImpl().forward( - query, - key, - cos, - sin, - inplace=inplace, - complex_mode=True, - ) + assert not complex_mode, 'Dlinfer backend does not support complex_mode' if inplace: q_embed = None k_embed = None diff --git a/tests/pytorch/kernel/test_sparse_index_topk.py b/tests/pytorch/kernel/test_sparse_index_topk.py index 8d7462e6dc..31c1a91a16 100644 --- a/tests/pytorch/kernel/test_sparse_index_topk.py +++ b/tests/pytorch/kernel/test_sparse_index_topk.py @@ -1,41 +1,17 @@ # Copyright (c) OpenMMLab. All rights reserved. -import sys - import pytest import torch -def _requires_cuda() -> bool: - return not torch.cuda.is_available() - - -pytestmark = pytest.mark.skipif(_requires_cuda(), reason='requires CUDA') - - -def test_nsa_backend_selects_sparse_topk_for_glm52(): - from lmdeploy.pytorch.backends.cuda.nsa import _get_sparse_index_topk - - assert _get_sparse_index_topk(2048).__name__ == 'sparse_index_topk' - assert _get_sparse_index_topk(1024) is None +def _requires_sm90_cuda() -> bool: + return not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9 -def test_nsa_backend_falls_back_without_tilelang(monkeypatch): - from lmdeploy.pytorch.backends.cuda import nsa +pytestmark = pytest.mark.skipif(_requires_sm90_cuda(), reason='requires CUDA device with cc>=9.0') - nsa._get_sparse_index_topk.cache_clear() - monkeypatch.setitem(sys.modules, - 'lmdeploy.pytorch.kernels.cuda.sparse_index_topk', None) - try: - assert nsa._get_sparse_index_topk(2048) is None - finally: - nsa._get_sparse_index_topk.cache_clear() - -def _assert_topk_ids(scores: torch.Tensor, - out: torch.Tensor, - seqlens: list[int], - k: int, - fill: int = -1): +def _assert_topk_ids(scores: torch.Tensor, out: torch.Tensor, + seqlens: list[int], k: int, fill: int = -1): scores = scores.cpu() out = out.cpu() score_width = scores.size(1) @@ -52,40 +28,31 @@ def _assert_topk_ids(scores: torch.Tensor, assert row_out[seqlen:].eq(fill).all() continue - expected = torch.topk(scores[row, :seqlen], - k=k, - largest=True, - sorted=False).indices + expected = torch.topk(scores[row, :seqlen], k=k, largest=True, sorted=False).indices assert valid.numel() == k - torch.testing.assert_close(valid.sort().values, - expected.to(torch.int32).sort().values) + torch.testing.assert_close(valid.sort().values, expected.to(torch.int32).sort().values) -@pytest.mark.parametrize('k', [512, 2048]) -def test_sparse_index_topk_matches_torch_topk_and_fill(k: int): +def test_sparse_index_topk_matches_torch_topk_and_fill(): from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import ( is_sparse_index_topk_supported, sparse_index_topk, ) - assert is_sparse_index_topk_supported(k) + assert is_sparse_index_topk_supported(512) device = 'cuda' + k = 512 fill = -7 - score_width = k * 2 - seqlens = [0, 17, k, k + 1, score_width - 124] - generator = torch.Generator(device=device).manual_seed(20260709 + k) - scores = torch.randn(len(seqlens), - score_width, - device=device, - dtype=torch.float32, - generator=generator) - scores += torch.arange(score_width, device=device, - dtype=torch.float32) * 1e-6 + score_width = 1024 + seqlens = [0, 17, 512, 513, 900] + generator = torch.Generator(device=device).manual_seed(20260709) + scores = torch.randn(len(seqlens), score_width, device=device, + dtype=torch.float32, generator=generator) + scores += torch.arange(score_width, device=device, dtype=torch.float32) * 1e-6 q_seqlens = torch.ones(len(seqlens), device=device, dtype=torch.int64) - kv_dtype = torch.int64 if k == 2048 else torch.int32 - kv_seqlens = torch.tensor(seqlens, device=device, dtype=kv_dtype) + kv_seqlens = torch.tensor(seqlens, device=device, dtype=torch.int32) out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k, fill=fill) assert out.shape == (len(seqlens), k) @@ -93,6 +60,19 @@ def test_sparse_index_topk_matches_torch_topk_and_fill(k: int): _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 @@ -102,14 +82,10 @@ def test_sparse_index_topk_accepts_padded_score_stride(): padded_width = 1280 seqlens = [600, 777, 1024, 321] generator = torch.Generator(device=device).manual_seed(20260710) - storage = torch.randn(len(seqlens), - padded_width, - device=device, - dtype=torch.float32, - generator=generator) + storage = torch.randn(len(seqlens), padded_width, device=device, + dtype=torch.float32, generator=generator) scores = storage[:, :score_width] - scores += torch.arange(score_width, device=device, - dtype=torch.float32) * 1e-6 + scores += torch.arange(score_width, device=device, dtype=torch.float32) * 1e-6 assert not scores.is_contiguous() q_seqlens = torch.ones(len(seqlens), device=device, dtype=torch.int64) @@ -130,38 +106,29 @@ def test_sparse_index_topk_expands_batch_kv_seqlens_for_prefill(): row_seqlens = [640, 640, 1100, 1100, 1100] generator = torch.Generator(device=device).manual_seed(260619348) - scores = torch.randn(sum(q_seqlens_list), - score_width, - device=device, - dtype=torch.float32, - generator=generator) - scores += torch.arange(score_width, device=device, - dtype=torch.float32) * 1e-6 + scores = torch.randn(sum(q_seqlens_list), score_width, device=device, + dtype=torch.float32, generator=generator) + scores += torch.arange(score_width, device=device, dtype=torch.float32) * 1e-6 q_seqlens = torch.tensor(q_seqlens_list, device=device, dtype=torch.int64) - kv_seqlens = torch.tensor(batch_kv_seqlens, - device=device, - dtype=torch.int32) + kv_seqlens = torch.tensor(batch_kv_seqlens, device=device, dtype=torch.int32) out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k) _assert_topk_ids(scores, out, row_seqlens, k) -@pytest.mark.parametrize('k', [512, 2048]) -def test_sparse_index_topk_cuda_graph_capture(k: int): +def test_sparse_index_topk_cuda_graph_capture(): from lmdeploy.pytorch.kernels.cuda.sparse_index_topk import sparse_index_topk device = 'cuda' - score_width = k * 2 - seqlens = [k + 1, k + 265, score_width, k // 2] - generator = torch.Generator(device=device).manual_seed(k) - scores = torch.randn(len(seqlens), - score_width, - device=device, - dtype=torch.float32, - generator=generator) + k = 512 + score_width = 1024 + seqlens = [600, 777, 1024, 321] + generator = torch.Generator(device=device).manual_seed(512) + scores = torch.randn(len(seqlens), score_width, device=device, + dtype=torch.float32, generator=generator) q_seqlens = torch.ones(len(seqlens), device=device, dtype=torch.int64) - kv_seqlens = torch.tensor(seqlens, device=device, dtype=torch.int64) + kv_seqlens = torch.tensor(seqlens, device=device, dtype=torch.int32) # Warm the TileLang specialization and PyTorch's graph-aware allocator. out = sparse_index_topk(scores, q_seqlens, kv_seqlens, k=k) From 1879da72268ac6aba42433365211599de1dcc15e Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 27 Jul 2026 15:12:01 +0800 Subject: [PATCH 29/33] perf: use zero-copy BF16 cache for sparse MLA --- .../pytorch/backends/cuda/attention/mla.py | 61 +++++++++++++++---- tests/pytorch/kernel/test_mla_attention.py | 51 ++++++++++++++-- 2 files changed, 95 insertions(+), 17 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index 20048899c1..466ca1ee8d 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -36,6 +36,7 @@ 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, max_q_seqlen: int, @@ -65,6 +66,36 @@ def update_decode(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, 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.""" num_tokens = nsa_indices.size(0) @@ -93,7 +124,7 @@ class FlashMLAImpl(TritonAttentionImpl): This implementation supports multiple execution paths: - Paged FlashMLA decode: Uses flash_mla_with_kvcache with FP8 paged KV cache - - Sparse FlashMLA decode: Uses flash_mla_sparse_fwd over the BF16 global cache view + - 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 @@ -103,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, @@ -212,9 +244,9 @@ def _decode_paged_flash_mla( attn_output = attn_output.flatten(0, 1) return attn_output - def _flash_mla_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, + def _flash_mla_sparse(self, query: torch.Tensor, indexed_k: torch.Tensor, nsa_indices: torch.Tensor) -> torch.Tensor: - """Run sparse FlashMLA over a BF16 flat KV view.""" + """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 @@ -224,7 +256,7 @@ def _flash_mla_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, output = flash_mla_sparse_fwd( query, - flatten_k, + indexed_k, nsa_indices, sm_scale=self.scale, ) @@ -241,17 +273,24 @@ def _prefill_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, nsa_indi 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 the global BF16 paged-cache view.""" + """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) - # sparse_fwd addresses a flat cache, so map paged positions to physical slots. - nsa_indices = self.nsa_updater.update_decode(nsa_indices, attn_metadata.block_offsets, max_q_seqlen, - block_size) + # 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] - flatten_k = k_cache.flatten(0, 1) - return self._flash_mla_sparse(query, flatten_k, nsa_indices) + return self._flash_mla_sparse(query, storage_k, nsa_indices) def _prefill_triton( self, @@ -566,7 +605,7 @@ def forward( Architecture: - Paged FlashMLA decode: Uses flash_mla_with_kvcache with FP8 paged KV cache - - Sparse FlashMLA decode: Uses flash_mla_sparse_fwd over the BF16 global 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 FlashMLA: flash_mla_sparse_fwd * FA3 optimized: flash_attn_varlen_func with split q_rope/q_nope diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index 72e41b03c6..c0e1e8d92a 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -32,14 +32,17 @@ def test_nsa_decode_indices_update_keeps_single_token_shape(): assert torch.equal(output, expected) -def test_bf16_sparse_decode_uses_global_cache_view(): +def test_bf16_sparse_decode_uses_strided_cache_view(): impl = object.__new__(FlashMLAImpl) impl.nsa_updater = NSAIndicesUpdater() - impl.nsa_updater._update_decode_func = impl.nsa_updater._update_decode_impl + 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) - k_cache = torch.empty(3, 16, 1, 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]), @@ -47,13 +50,49 @@ def test_bf16_sparse_decode_uses_global_cache_view(): impl._decode_bf16_sparse_flash_mla(query, k_cache, nsa_indices, metadata) - sparse_query, flatten_k, global_indices = impl._flash_mla_sparse.call_args.args + sparse_query, storage_k, global_indices = impl._flash_mla_sparse.call_args.args assert sparse_query is query - assert flatten_k.shape == (48, 1, 576) - expected = torch.tensor([[[16, 33]], [[0, -1]], [[32, 17]], [[16, 33]]]) + 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 From fd897698e1465ce9ac3381f12e822f88374e04f4 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 27 Jul 2026 20:27:34 +0800 Subject: [PATCH 30/33] perf: bypass dsa scoring for dense decode --- .../pytorch/backends/cuda/graph_runner.py | 4 ++ lmdeploy/pytorch/backends/cuda/nsa.py | 25 +++++++- lmdeploy/pytorch/backends/nsa.py | 6 ++ lmdeploy/pytorch/kernels/cuda/ds_index.py | 59 +++++++++++++++++++ lmdeploy/pytorch/models/deepseek_v32.py | 52 ++++++++++++++-- lmdeploy/pytorch/models/glm_moe_dsa_mtp.py | 8 +++ lmdeploy/pytorch/nn/nsa.py | 24 ++++++++ tests/pytorch/kernel/test_ds_index.py | 18 ++++++ .../spec_decode/test_cudagraph_strategy.py | 21 +++++-- 9 files changed, 207 insertions(+), 10 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/graph_runner.py b/lmdeploy/pytorch/backends/cuda/graph_runner.py index 6c9da66791..ed1e9acec1 100644 --- a/lmdeploy/pytorch/backends/cuda/graph_runner.py +++ b/lmdeploy/pytorch/backends/cuda/graph_runner.py @@ -243,6 +243,10 @@ def get_graph_key(self, input_ids: torch.Tensor, position_ids: torch.Tensor, pas 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): diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index b4d1196511..6d112d163a 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -6,7 +6,7 @@ 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 @@ -139,6 +139,29 @@ def forward_fused(self, q: Tensor, k: Tensor, weights: Tensor, norm_weight: Tens 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/nsa.py b/lmdeploy/pytorch/backends/nsa.py index f8ba09d5c7..f682abec71 100644 --- a/lmdeploy/pytorch/backends/nsa.py +++ b/lmdeploy/pytorch/backends/nsa.py @@ -31,6 +31,12 @@ def forward_fused(self, q: Tensor, k: Tensor, weights: Tensor, norm_weight: Tens """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/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/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index b47a88e69b..5a020984a5 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -8,7 +8,7 @@ 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, get_step_ctx_manager +from lmdeploy.pytorch.model_inputs import StepContext, StepContextManager, get_step_ctx_manager from lmdeploy.pytorch.nn import ( ApplyRotaryEmb, Attention, @@ -248,14 +248,28 @@ def forward(self, qr: torch.Tensor, freqs_cis: torch.Tensor, index_cache: tuple[torch.Tensor, torch.Tensor], - attn_metadata: Any = None): - q = self.wq_b(qr) - q = q.unflatten(-1, (-1, self.head_dim)) + 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], @@ -270,6 +284,8 @@ def forward(self, 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) k = self.wk(x) k = self.k_norm(k) @@ -449,6 +465,7 @@ def forward( attn_metadata: Any = None, topk_indices_buffer: DSATopKIndicesBuffer | None = None, skip_topk: bool = False, + use_dense_index: bool = False, ): """Rewrite of LlamaAttention.forward.""" dist_config = get_dist_manager().current_config() @@ -483,7 +500,8 @@ def forward( qr, rotary_pos_emb, past_key_value[-2:], - attn_metadata=attn_metadata)) + attn_metadata=attn_metadata, + use_dense_index=use_dense_index)) else: topk_indices = topk_indices_buffer.read(q_len, hidden_states.device) @@ -549,6 +567,7 @@ def forward( 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]: @@ -566,6 +585,7 @@ def forward( 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( @@ -631,6 +651,7 @@ def forward( 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: @@ -651,6 +672,7 @@ def forward( 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) @@ -665,6 +687,7 @@ def forward_microbatch( 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 @@ -676,6 +699,7 @@ def forward_microbatch( attn_metadata=attn_metadata, inputs_embeds=inputs_embeds, all_routed_experts=all_routed_experts, + use_dense_index=use_dense_index, ) @@ -709,6 +733,7 @@ def forward( past_key_values: list[list[torch.Tensor]], attn_metadata: Any = None, inputs_embeds: torch.Tensor = None, + use_dense_index: bool = False, **kwargs, ): """Model forward.""" @@ -731,11 +756,28 @@ def forward( 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.""" diff --git a/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py b/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py index e2177c87b1..3595eaaa55 100644 --- a/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py +++ b/lmdeploy/pytorch/models/glm_moe_dsa_mtp.py @@ -84,6 +84,7 @@ def forward( 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) @@ -98,6 +99,7 @@ def forward( attn_metadata=attn_metadata, topk_indices_buffer=topk_indices_buffer, skip_topk=skip_topk, + use_dense_index=use_dense_index, ) return residual + hidden_states @@ -134,6 +136,7 @@ def forward( 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: @@ -149,6 +152,7 @@ def forward( 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( @@ -218,6 +222,7 @@ def forward( 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( @@ -229,6 +234,7 @@ def forward( 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, ) @@ -247,6 +253,8 @@ def prepare_inputs_for_generation( 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, diff --git a/lmdeploy/pytorch/nn/nsa.py b/lmdeploy/pytorch/nn/nsa.py index 57e7386dc4..81bc94ec11 100644 --- a/lmdeploy/pytorch/nn/nsa.py +++ b/lmdeploy/pytorch/nn/nsa.py @@ -98,3 +98,27 @@ def forward_fused(self, 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/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/spec_decode/test_cudagraph_strategy.py b/tests/pytorch/spec_decode/test_cudagraph_strategy.py index fdc9c546d1..c72b9b97e0 100644 --- a/tests/pytorch/spec_decode/test_cudagraph_strategy.py +++ b/tests/pytorch/spec_decode/test_cudagraph_strategy.py @@ -157,9 +157,22 @@ def test_cuda_graph_key_separates_dsa_seed_and_reuse(monkeypatch): inputs_embeds=None, ) - seed_key = runner.get_graph_key(**kwargs, skip_topk=False) - reuse_key = runner.get_graph_key(**kwargs, skip_topk=True) + 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[-1] is False - assert reuse_key[-1] is True + 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) From f13d85229d7a1856d69df3e6ec70d1483dc8c462 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 27 Jul 2026 20:28:03 +0800 Subject: [PATCH 31/33] perf: compact blocked-fp8 moe scheduling --- .../kernels/cuda/blocked_fp8_fused_moe.py | 101 +++++++++++++----- .../kernel/test_fuse_moe_blocked_fp8.py | 76 +++++++++++++ tests/pytorch/kernel/test_fused_moe.py | 33 ++++++ 3 files changed, 184 insertions(+), 26 deletions(-) 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/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), From a39a367fbf50fa17590799769522fbcc8e0a1aba Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 27 Jul 2026 20:28:32 +0800 Subject: [PATCH 32/33] perf: simplify single-token decode metadata --- lmdeploy/pytorch/model_inputs.py | 13 ++++++++++--- tests/pytorch/test_model_inputs.py | 25 +++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 3 deletions(-) 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/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])) From a7761c93c46e762d0d9a47d147d1c992558eebbc Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 27 Jul 2026 20:46:26 +0800 Subject: [PATCH 33/33] style: format dense indexer docstring --- lmdeploy/pytorch/backends/cuda/nsa.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index 6d112d163a..b7190223c9 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -142,7 +142,8 @@ def forward_fused(self, q: Tensor, k: Tensor, weights: Tensor, norm_weight: Tens 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.""" + """Cache K and skip score/top-k work when all positions are + selected.""" prepare_dsa_indexer_k_cache(k, norm_weight, norm_bias,