From f174a6cf675cdbd5d46843dc245d177f6c0736f9 Mon Sep 17 00:00:00 2001 From: jinminxi104 Date: Tue, 9 Jun 2026 11:31:04 +0000 Subject: [PATCH 1/3] fix vision in qwen35 --- .../framework/lmdeploy_ext/device/__init__.py | 68 ++++++++++--------- dlinfer/framework/torch_npu_ext/aclgraph.py | 2 +- 2 files changed, 37 insertions(+), 33 deletions(-) diff --git a/dlinfer/framework/lmdeploy_ext/device/__init__.py b/dlinfer/framework/lmdeploy_ext/device/__init__.py index 2d9eb136..0f08ab24 100644 --- a/dlinfer/framework/lmdeploy_ext/device/__init__.py +++ b/dlinfer/framework/lmdeploy_ext/device/__init__.py @@ -546,49 +546,47 @@ def custom_prepare_inputs_for_generation( past_key_values = list(past_key_values) new_past_key_values = [] for layer_type in self.config.text_config.layer_types: - if layer_type == "linear_attention": + if layer_type == 'linear_attention': new_past_key_values.append(state_caches.pop(0)) - elif layer_type == "full_attention": + elif layer_type == 'full_attention': new_past_key_values.append(past_key_values.pop(0)) # vlm inputs pixel_values = None vis_cu_seqlens = None vis_pos_emb = None - image_mask = None + multimodal_mask = None grid_thw = None pos_embeds = None + # for time series + ts_values = None + ts_lens = None + ts_sr = None if context.input_multimodals is not None: - mm_inputs = [ - input_mm.get("mm_data", []) for input_mm in context.input_multimodals - ] + mm_inputs = [input_mm.get('mm_data', []) for input_mm in context.input_multimodals] # flatten batch mm_inputs = [item for sublist in mm_inputs for item in sublist] if len(mm_inputs) > 0: modality = mm_inputs[0].modality - pixel_values = torch.cat([inp.data for inp in mm_inputs]) - - image_token_id = mm_inputs[0].meta.get("image_token_id") - video_token_id = mm_inputs[0].meta.get("video_token_id") - mm_token_id = ( - image_token_id if modality == Modality.IMAGE else video_token_id - ) - image_mask = input_ids == mm_token_id - - grid_thw = torch.cat( - [data.meta["grid_thw"] for data in mm_inputs] - ).cpu() - vis_pos_emb = self.model.visual.rot_pos_emb(grid_thw) - pos_embeds = self.model.visual.fast_pos_embed_interpolate(grid_thw) - vis_cu_seqlens = torch.repeat_interleave( - grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] - ).to(pixel_values.device) - vis_cu_seqlens = vis_cu_seqlens.cumsum(dim=0, dtype=torch.int32) - vis_pos_emb = vis_pos_emb.repeat(1, 2) - vis_pos_emb = (vis_pos_emb.cos(), vis_pos_emb.sin()) - - mrope_position_ids = getattr(context, "mrope_position_ids", None) + multimodal_mask = self.get_multimodal_mask(input_ids, mm_inputs) + + if modality == Modality.TIME_SERIES: + ts_values = torch.cat([inp.data for inp in mm_inputs]) + ts_lens = torch.cat([inp.meta['ts_lens'] for inp in mm_inputs]) + ts_sr = torch.cat([inp.meta['ts_sr'] for inp in mm_inputs]) + else: + pixel_values = torch.cat([inp.data for inp in mm_inputs]) + grid_thw = torch.stack([data.meta['grid_thw'] for data in mm_inputs]).cpu() + vis_pos_emb = self.model.visual.rot_pos_emb(grid_thw) + pos_embeds = self.model.visual.fast_pos_embed_interpolate(grid_thw) + vis_cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], + grid_thw[:, 0]).to(pixel_values.device) + vis_cu_seqlens = vis_cu_seqlens.cumsum(dim=0, dtype=torch.int32) + vis_pos_emb = vis_pos_emb.repeat(1, 2) + vis_pos_emb = (vis_pos_emb.cos(), vis_pos_emb.sin()) + + mrope_position_ids = getattr(context, 'mrope_position_ids', None) # process vision embeddings vision_embeddings = context.input_embeddings @@ -596,9 +594,10 @@ def custom_prepare_inputs_for_generation( if vision_embeddings is not None and len(vision_embeddings) > 0: if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) - inputs_embeds[:, vision_embedding_indexing, :] = vision_embeddings.to( - inputs_embeds - ) + inputs_embeds[:, vision_embedding_indexing, :] = vision_embeddings.to(inputs_embeds) + + # return input embeds for spec decoding + return_input_embeds = self.is_spec_decoding and (pixel_values is not None or context.is_chunk_multimodal) # inputs of forward return dict( @@ -613,9 +612,14 @@ def custom_prepare_inputs_for_generation( pixel_values=pixel_values, vis_cu_seqlens=vis_cu_seqlens, vis_pos_emb=vis_pos_emb, - image_mask=image_mask, + multimodal_mask=multimodal_mask, grid_thw=grid_thw, pos_embeds=pos_embeds, + return_input_embeds=return_input_embeds, + # for time series + ts_values=ts_values, + ts_lens=ts_lens, + ts_sr=ts_sr, ) def custom_forward( diff --git a/dlinfer/framework/torch_npu_ext/aclgraph.py b/dlinfer/framework/torch_npu_ext/aclgraph.py index 4463e718..dba104d2 100644 --- a/dlinfer/framework/torch_npu_ext/aclgraph.py +++ b/dlinfer/framework/torch_npu_ext/aclgraph.py @@ -56,4 +56,4 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): def appy_patch(): - _GraphDispatchMode.__torch_dispatch__ = __torch_dispatch__ + pass From 672f4cf830a65b671ced7ed2b6cfa1d42b83238f Mon Sep 17 00:00:00 2001 From: jinminxi104 Date: Wed, 24 Jun 2026 16:08:58 +0000 Subject: [PATCH 2/3] fix sampling patch --- .../framework/lmdeploy_ext/device/__init__.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/dlinfer/framework/lmdeploy_ext/device/__init__.py b/dlinfer/framework/lmdeploy_ext/device/__init__.py index 0f08ab24..75fdf4ad 100644 --- a/dlinfer/framework/lmdeploy_ext/device/__init__.py +++ b/dlinfer/framework/lmdeploy_ext/device/__init__.py @@ -2,6 +2,8 @@ import importlib import torch from functools import lru_cache +from lmdeploy.pytorch.model_inputs import ModelInputs +from lmdeploy.pytorch.strategies.base.model_agent import ExtraInputs from dlinfer.vendor import vendor_name vendor = ["camb", "ascend"] @@ -72,13 +74,16 @@ def patch_async_sampling_logits(): FusedLogitsProcessor, ) - async def async_sampling_logits( - self, logits: torch.Tensor, sampling_inputs: SamplingInputs - ): + async def async_sampling_logits(self, logits: torch.Tensor, inputs: ModelInputs, + extra_inputs: ExtraInputs, sampling_inputs: SamplingInputs): """Sampling logits.""" + if self.spec_agent.is_enabled(): + extra_inputs.target_logits = extra_inputs.target_logits.to(torch.float32) + extra_inputs = await self.spec_agent.async_sampling_logits(inputs, extra_inputs, sampling_inputs) + return extra_inputs.next_token_ids, extra_inputs.logprobs, extra_inputs.output_token_ids, extra_inputs # record function does not support async function # so we can not decorate it on async_sampling_logits - with record_function("sampling_logits"): + with record_function('sampling_logits'): logits = logits.to(torch.float32) logits_processor = FusedLogitsProcessor( sampling_inputs, @@ -88,14 +93,18 @@ async def async_sampling_logits( origin_logits = logits logits, raw_logprobs = await logits_processor(origin_logits) next_token_ids = logits_processor.sampling(logits) + await logits_processor.accept_guided_tokens(next_token_ids) logprobs = logits_processor.compute_logprobs(raw_logprobs, next_token_ids) if logprobs is not None: logprobs = BatchedLogProbs( vals=logprobs[0], indices=logprobs[1], ) + # post sampling + next_token_ids, extra_inputs = self.agent_strategy.post_sampling(inputs, logits, next_token_ids, + extra_inputs) + return next_token_ids, logprobs, next_token_ids, extra_inputs - return next_token_ids, logprobs BaseModelAgent.async_sampling_logits = async_sampling_logits From c5f3a3bd1269ee7459a0620eb054fa612b33f3bd Mon Sep 17 00:00:00 2001 From: jinminxi104 Date: Thu, 25 Jun 2026 06:46:43 +0000 Subject: [PATCH 3/3] lint --- .../framework/lmdeploy_ext/device/__init__.py | 44 +++++++++++++------ 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/dlinfer/framework/lmdeploy_ext/device/__init__.py b/dlinfer/framework/lmdeploy_ext/device/__init__.py index 5fba32a8..670a99bc 100644 --- a/dlinfer/framework/lmdeploy_ext/device/__init__.py +++ b/dlinfer/framework/lmdeploy_ext/device/__init__.py @@ -74,16 +74,28 @@ def patch_async_sampling_logits(): FusedLogitsProcessor, ) - async def async_sampling_logits(self, logits: torch.Tensor, inputs: ModelInputs, - extra_inputs: ExtraInputs, sampling_inputs: SamplingInputs): + async def async_sampling_logits( + self, + logits: torch.Tensor, + inputs: ModelInputs, + extra_inputs: ExtraInputs, + sampling_inputs: SamplingInputs, + ): """Sampling logits.""" if self.spec_agent.is_enabled(): extra_inputs.target_logits = extra_inputs.target_logits.to(torch.float32) - extra_inputs = await self.spec_agent.async_sampling_logits(inputs, extra_inputs, sampling_inputs) - return extra_inputs.next_token_ids, extra_inputs.logprobs, extra_inputs.output_token_ids, extra_inputs + extra_inputs = await self.spec_agent.async_sampling_logits( + inputs, extra_inputs, sampling_inputs + ) + return ( + extra_inputs.next_token_ids, + extra_inputs.logprobs, + extra_inputs.output_token_ids, + extra_inputs, + ) # record function does not support async function # so we can not decorate it on async_sampling_logits - with record_function('sampling_logits'): + with record_function("sampling_logits"): logits = logits.to(torch.float32) logits_processor = FusedLogitsProcessor( sampling_inputs, @@ -101,11 +113,11 @@ async def async_sampling_logits(self, logits: torch.Tensor, inputs: ModelInputs, indices=logprobs[1], ) # post sampling - next_token_ids, extra_inputs = self.agent_strategy.post_sampling(inputs, logits, next_token_ids, - extra_inputs) + next_token_ids, extra_inputs = self.agent_strategy.post_sampling( + inputs, logits, next_token_ids, extra_inputs + ) return next_token_ids, logprobs, next_token_ids, extra_inputs - BaseModelAgent.async_sampling_logits = async_sampling_logits @@ -555,9 +567,9 @@ def custom_prepare_inputs_for_generation( past_key_values = list(past_key_values) new_past_key_values = [] for layer_type in self.config.text_config.layer_types: - if layer_type == 'linear_attention': + if layer_type == "linear_attention": new_past_key_values.append(state_caches.pop(0)) - elif layer_type == 'full_attention': + elif layer_type == "full_attention": new_past_key_values.append(past_key_values.pop(0)) # vlm inputs @@ -572,7 +584,9 @@ def custom_prepare_inputs_for_generation( ts_lens = None ts_sr = None if context.input_multimodals is not None: - mm_inputs = [input_mm.get('mm_data', []) for input_mm in context.input_multimodals] + mm_inputs = [ + input_mm.get("mm_data", []) for input_mm in context.input_multimodals + ] # flatten batch mm_inputs = [item for sublist in mm_inputs for item in sublist] @@ -606,10 +620,14 @@ def custom_prepare_inputs_for_generation( if vision_embeddings is not None and len(vision_embeddings) > 0: if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) - inputs_embeds[:, vision_embedding_indexing, :] = vision_embeddings.to(inputs_embeds) + inputs_embeds[:, vision_embedding_indexing, :] = vision_embeddings.to( + inputs_embeds + ) # return input embeds for spec decoding - return_input_embeds = self.is_spec_decoding and (pixel_values is not None or context.is_chunk_multimodal) + return_input_embeds = self.is_spec_decoding and ( + pixel_values is not None or context.is_chunk_multimodal + ) # return input embeds for spec decoding return_input_embeds = self.is_spec_decoding and (