diff --git a/pyproject.toml b/pyproject.toml index 0958b9395..c7b2be20e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -242,4 +242,5 @@ ue = "ue" semantics = "semantics" fullset = "fullset" Vai = "Vai" +tockens = "tockens" CANN = "CANN" diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index fb8b61621..39baed6bd 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -46,7 +46,6 @@ stages: - stage_id: 2 gpu_memory_utilization: 0.1 - max_num_seqs: 1 enforce_eager: true async_scheduling: false max_num_batched_tokens: 51200 @@ -67,10 +66,11 @@ platforms: - stage_id: 0 gpu_memory_utilization: 0.6 tensor_parallel_size: 2 + max_num_batched_tokens: 8192 devices: "0,1" - stage_id: 1 gpu_memory_utilization: 0.6 - enforce_eager: true + max_num_batched_tokens: 8192 devices: "2" - stage_id: 2 gpu_memory_utilization: 0.3 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 5948007a9..522ea7c58 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -71,16 +71,3 @@ stages: max_tokens: 65536 seed: 42 repetition_penalty: 1.0 - -platforms: - npu: - stages: - # NPU does not yet support async-scheduling for TTS, and the - # talker fits at max_num_seqs=1 only. - - stage_id: 0 - max_num_seqs: 1 - enforce_eager: true - async_scheduling: false - - stage_id: 1 - gpu_memory_utilization: 0.2 - async_scheduling: false diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 3a904442f..ecfeec381 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -122,6 +122,9 @@ class CodePredictorAttention(nn.Module): super().__init__() self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads + assert self.num_heads % self.num_kv_heads == 0 + self.is_gqa = self.num_kv_heads != self.num_heads + self.num_queries_per_kv = self.num_heads // self.num_kv_heads self.head_dim = getattr( config, "head_dim", @@ -129,7 +132,7 @@ class CodePredictorAttention(nn.Module): ) self.hidden_size = config.hidden_size self.scaling = self.head_dim**-0.5 - self._use_gqa = self.num_kv_heads != self.num_heads + self.max_seq = int(config.num_code_groups) + 1 # Separate q/k/v projections matching HF (no fused packing) bias = getattr(config, "attention_bias", False) @@ -140,6 +143,73 @@ class CodePredictorAttention(nn.Module): self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) + if current_omni_platform.is_npu(): + if self.max_seq > 2048: + raise ValueError( + "Qwen3-TTS code predictor NPU fusion attention uses a fixed 2048x2048 " + f"causal mask, but max_seq={self.max_seq} exceeds the mask size." + ) + # Ascend SDPA is_causal migration example uses a fixed 2048x2048 + # compressed causal mask with sparse_mode=2. + fusion_mask = torch.triu( + torch.ones(2048, 2048, dtype=torch.bool), + diagonal=1, + ) + self.register_buffer("_fusion_causal_mask", fusion_mask, persistent=False) + + def _forward_npu_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + bsz: int, + seq_len: int, + ) -> torch.Tensor: + import torch_npu + + q_f, k_f, v_f = q, k, v + if self.is_gqa: + k_f = ( + k[:, :, None, :, :] + .expand(bsz, self.num_kv_heads, self.num_queries_per_kv, seq_len, self.head_dim) + .reshape(bsz, self.num_heads, seq_len, self.head_dim) + ) + v_f = ( + v[:, :, None, :, :] + .expand(bsz, self.num_kv_heads, self.num_queries_per_kv, seq_len, self.head_dim) + .reshape(bsz, self.num_heads, seq_len, self.head_dim) + ) + + mask = self._fusion_causal_mask + mask = mask.contiguous() + q_f = q_f.contiguous() + k_f = k_f.contiguous() + v_f = v_f.contiguous() + return torch_npu.npu_fusion_attention( + q_f, + k_f, + v_f, + self.num_heads, + "BNSD", + pse=None, + padding_mask=None, + atten_mask=mask, + scale=float(self.scaling), + keep_prob=1.0, + # Keep torch_npu's API spelling. + pre_tockens=2147483647, + next_tockens=2147483647, + inner_precise=0, + prefix=None, + actual_seq_qlen=None, + actual_seq_kvlen=None, + # Ascend SDPA is_causal migration example uses sparse_mode=2. + sparse_mode=2, + gen_mask_parallel=True, + # Keep sync=True for the NPU fused attention path. + sync=True, + )[0] + def forward( self, hidden_states: torch.Tensor, @@ -160,14 +230,17 @@ class CodePredictorAttention(nn.Module): q = (q * cos) + (_rotate_half(q) * sin) k = (k * cos) + (_rotate_half(k) * sin) - attn_out = F.scaled_dot_product_attention( - q, - k, - v, - scale=self.scaling, - is_causal=True, - enable_gqa=self._use_gqa, - ) + if not current_omni_platform.is_npu(): + attn_out = F.scaled_dot_product_attention( + q, + k, + v, + scale=self.scaling, + is_causal=True, + enable_gqa=self.is_gqa, + ) + else: + attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len) attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) return self.o_proj(attn_out) @@ -384,7 +457,7 @@ class CodePredictorWrapper(nn.Module): self._bucket_pos_ids: dict[int, torch.Tensor] = {} self._lm_heads_list: list[nn.Module] | None = None self._codec_embeds_list: list[nn.Module] | None = None - self._cuda_graphs: dict[int, tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {} + self._device_graphs: dict[int, tuple] = {} # (graph, static_output) per bucket def get_input_embeddings(self) -> nn.ModuleList: return self.model.get_input_embeddings() @@ -412,7 +485,7 @@ class CodePredictorWrapper(nn.Module): self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device) def _setup_compile(self) -> None: - """Lazily set up torch.compile with optional CUDA graph capture.""" + """Lazily set up torch.compile with optional device graph capture.""" if self._compiled_model_fwd is not None: return @@ -424,8 +497,16 @@ class CodePredictorWrapper(nn.Module): self._codec_embeds_list = list(self.model.codec_embedding) if not current_omni_platform.supports_torch_inductor(): - logger.warning_once("code_predictor: torch.compile disabled") + # NPU or other platforms without Inductor support self._compiled_model_fwd = self.model.forward + + if current_omni_platform.is_npu() and self._wrapper_config.use_cuda_graphs: + # For NPU, use eager + NPU graphs (no torch.compile) + self._warmup_buckets() + self._capture_npu_graphs() + logger.info("code_predictor: eager mode + NPU graphs") + else: + logger.warning_once("code_predictor: torch.compile disabled") return # torch.compile fuses RMSNorm/RoPE in ways that lose float32 @@ -491,10 +572,28 @@ class CodePredictorWrapper(nn.Module): with torch.cuda.graph(g, pool=pool): static_output = self._compiled_model_fwd(static_input, pos_ids) - self._cuda_graphs[bsz] = (g, static_output) + self._device_graphs[bsz] = (g, static_output) logger.info("code_predictor: captured CUDA graphs for buckets %s", self._bucket_sizes) + def _capture_npu_graphs(self) -> None: + """Capture an NPU graph per bucket using torch_npu's NPUGraph.""" + max_seq = self._num_groups + 1 + proj_buf = self._proj_buf + pool = torch.npu.graph_pool_handle() + + for bsz in self._bucket_sizes: + static_input = proj_buf[:bsz, :max_seq, :] + pos_ids = self._bucket_pos_ids[bsz] + + g = torch.npu.NPUGraph() + with torch.npu.graph(g, pool=pool): + static_output = self._compiled_model_fwd(static_input, pos_ids) + + self._device_graphs[bsz] = (g, static_output) + + logger.info("code_predictor: captured NPU graphs for buckets %s", self._bucket_sizes) + # ------------------------------------------------------------------ # Forward -- re-prefill + inline sampling # ------------------------------------------------------------------ @@ -544,8 +643,8 @@ class CodePredictorWrapper(nn.Module): torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1).contiguous() ) - # Use captured CUDA graph if available, otherwise call compiled fn. - cuda_graph_entry = self._cuda_graphs.get(padded_bsz) + # Use captured device graph if available, otherwise call compiled fn. + device_graph_entry = self._device_graphs.get(padded_bsz) # Prepare sampling parameters stored_mode = self._wrapper_config.sampling_mode == "stored" @@ -570,10 +669,10 @@ class CodePredictorWrapper(nn.Module): # Autoregressive loop: predict layers 1..G-1 for step in range(1, num_groups): - # Run transformer (CUDA graph replay or compiled forward) - if cuda_graph_entry is not None: - cuda_graph_entry[0].replay() - hidden_out = cuda_graph_entry[1] + # Run transformer (device graph replay or compiled forward) + if device_graph_entry is not None: + device_graph_entry[0].replay() + hidden_out = device_graph_entry[1] else: hidden_out = model_fwd(proj_buf[:padded_bsz, :max_seq, :], full_pos_ids) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py index 819e22e18..73c7e41d2 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py @@ -6,6 +6,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( CodePredictorWrapper, CodePredictorWrapperConfig, ) +from vllm_omni.platforms import current_omni_platform class Qwen3OmniMoeTalkerCodePredictor(CodePredictorWrapper): @@ -17,7 +18,7 @@ class Qwen3OmniMoeTalkerCodePredictor(CodePredictorWrapper): vllm_config=vllm_config, cp_config=cp_config, wrapper_config=CodePredictorWrapperConfig( - use_cuda_graphs=False, + use_cuda_graphs=current_omni_platform.is_npu(), use_parallel_embedding=True, use_projection=False, return_proj_buf=True,