[Performance] Pre-issued exponential distribution operator. (#4908)
Pre-issued exponential distribution operator.
Result:
Single inference saves 200-300 microseconds.
before:
<img width="2257" height="1058" alt="2"
src="https://github.com/user-attachments/assets/c1da19e2-a439-42cb-9d7c-c0218e61fd4c"
/>
After:
<img width="2211" height="342" alt="image"
src="https://github.com/user-attachments/assets/03c84292-c802-4755-949c-4266a9a72fc0"
/>
- vLLM version: v0.12.0
- vLLM main:
ad32e3e19c
---------
Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>
Co-authored-by: weijinqian_v1 <weijinqian@huawei.com>
This commit is contained in:
@@ -51,6 +51,7 @@ ACL_FORMAT_FRACTAL_NZ = 29
|
||||
_CUSTOM_OP_ENABLED = None
|
||||
_CURRENT_STREAM = None
|
||||
_PREFETCH_STREAM = None
|
||||
_GLOBAL_STREAM = None
|
||||
_SHARED_EXPERTS_CALCULATION_STREAM = None
|
||||
_ASCEND_CUSTOMOP_IS_REIGISTERED = False
|
||||
_DEFAULT_BUFFER_SIZE = 200
|
||||
@@ -292,6 +293,15 @@ def prefetch_stream() -> torch.npu.Stream:
|
||||
return _PREFETCH_STREAM
|
||||
|
||||
|
||||
def global_stream() -> torch.npu.Stream:
|
||||
global _GLOBAL_STREAM
|
||||
if _GLOBAL_STREAM is None:
|
||||
# when this function is called before any stream is set,
|
||||
# we return the default stream.
|
||||
_GLOBAL_STREAM = torch_npu.npu.Stream()
|
||||
return _GLOBAL_STREAM
|
||||
|
||||
|
||||
def shared_experts_calculation_stream() -> torch.npu.Stream:
|
||||
global _SHARED_EXPERTS_CALCULATION_STREAM
|
||||
if _SHARED_EXPERTS_CALCULATION_STREAM is None:
|
||||
|
||||
Reference in New Issue
Block a user