Skip to content

Commit 284775a

Browse files
authored
[Bugfix] [0.7.3] fix log output bug when calling npu_mrope (#911)
### What this PR does / why we need it? Fix #825 ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? With CI Signed-off-by: Shuqiao Li <celestialli@outlook.com>
1 parent f8efa89 commit 284775a

File tree

1 file changed

+28
-7
lines changed

1 file changed

+28
-7
lines changed

vllm_ascend/ops/rotary_embedding.py

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
# limitations under the License.
1616
#
1717

18+
import os
1819
from typing import Optional, Tuple
1920

2021
import torch
@@ -87,6 +88,25 @@ def rope_deepseek_forward_oot(
8788
return query, key
8889

8990

91+
class support_stdout_stderr(object):
92+
93+
def __init__(self):
94+
self.null_fds = [os.open(os.devnull, os.O_RDWR) for x in range(2)]
95+
self.save_fds = (os.dup(1), os.dup(2))
96+
97+
def __enter__(self):
98+
os.dup2(self.null_fds[0], 1)
99+
os.dup2(self.null_fds[1], 2)
100+
101+
def __exit__(self, *_):
102+
os.dup2(self.save_fds[0], 1)
103+
os.dup2(self.save_fds[1], 2)
104+
os.close(self.null_fds[0])
105+
os.close(self.null_fds[1])
106+
os.close(self.save_fds[0])
107+
os.close(self.save_fds[1])
108+
109+
90110
def mrope_forward(
91111
self,
92112
positions: torch.Tensor,
@@ -95,13 +115,14 @@ def mrope_forward(
95115
) -> Tuple[torch.Tensor, torch.Tensor]:
96116
import torch_npu
97117
mrope_section = [0, 0, 0] if positions.ndim == 1 else self.mrope_section
98-
query, key = torch_npu.npu_mrope(positions,
99-
query.contiguous(),
100-
key.contiguous(),
101-
self.cos_sin_cache.contiguous(),
102-
self.head_size,
103-
mrope_section=mrope_section,
104-
rotary_mode='half')
118+
with support_stdout_stderr():
119+
query, key = torch_npu.npu_mrope(positions,
120+
query.contiguous(),
121+
key.contiguous(),
122+
self.cos_sin_cache.contiguous(),
123+
self.head_size,
124+
mrope_section=mrope_section,
125+
rotary_mode='half')
105126
return query, key
106127

107128

0 commit comments

Comments
 (0)