triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 126978, Hardware limit: 101376. Reducing block sizes or `num_stages` may help

今天在调试python应用MonkeyOCR的时候发生报错，部分错误信息如下：

File “D:\python\MonkeyOCR\py310\lib\site-packages\lmdeploy\pytorch\nn\attention.py”, line 154, in forward return self.impl.forward( File “D:\python\MonkeyOCR\py310\lib\site-packages\lmdeploy\pytorch\backends\cuda\flash_attention.py”, line 56, in forward self.flash_attention_fwd( File “D:\python\MonkeyOCR\py310\lib\site-packages\lmdeploy\pytorch\kernels\cuda\flashattention.py”, line 449, in flash_attention_fwd _flash_prefill_fwd_kernel[grid]( File “D:\python\MonkeyOCR\py310\lib\site-packages\triton\runtime\jit.py”, line 345, in <lambda> return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs) File “D:\python\MonkeyOCR\py310\lib\site-packages\triton\runtime\jit.py”, line 691, in run kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata, File “D:\python\MonkeyOCR\py310\lib\site-packages\triton\compiler\compiler.py”, line 385, in __getattribute__ self._init_handles() File “D:\python\MonkeyOCR\py310\lib\site-packages\triton\compiler\compiler.py”, line 378, in _init_handles raise OutOfResources(self.metadata.shared, max_shared, “shared memory”) triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 126978, Hardware limit: 101376. Reducing block sizes or `num_stages` may help. 2025-07-01 12:43:43,759 – lmdeploy – ERROR – async_engine.py:791 – session 0 finished, reason “error”

主要错误就是：

triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 126978, Hardware limit: 101376. Reducing block sizes or `num_stages` may help.

意思就是共享内存资源不足，也就是共享显存不足

报错信息也给出了解决方法：

Reducing block sizes or `num_stages` may help

找到flashattention.py文件位置，我这里是：

\lib\site-packages\lmdeploy\pytorch\kernels\cuda\flashattention.py

找到num_stages值，

原代码是根据条件动态赋值的：

num_stages=num_stages,

我这里直接给设置了个小一点的固定值：

#num_stages=num_stages,
num_stages=2,

作为临时解决方案

重新运行程序不再发生报错。