无中生有的 Nan

# Code1
class StreamContext(AbstractContextManager):

    def __enter__(self):
        cur_stream = self.stream
        if cur_stream is None:
            return

        # _current_stream 就
        global _current_stream
        self.prev_stream = _current_stream
        _current_stream = cur_stream

import torch
import threading
import time

def thread_with_new_stream():
    time.sleep(3)
    s = torch.cuda.Stream()
    print(f"[Thread-A] New CUDA Stream: {s!r}")
    with torch.cuda.stream(s):
        print(f"[Thread-A] Enter stream context: {torch.cuda.current_stream()!r}")
        time.sleep(1300)

def thread_print_stream():
    for i in range(400):
        # 注意：每个线程（即每个Python Thread）有自己的stream堆栈
        print(f"[Thread-B] {i}: Current stream: {torch.cuda.current_stream()!r}")
        time.sleep(1)

def main():
    torch.cuda.init()  # 确保 CUDA 初始化

    t1 = threading.Thread(target=thread_with_new_stream)
    t2 = threading.Thread(target=thread_print_stream)

    t1.start()
    t2.start()

    t1.join()
    t2.join()

if __name__ == "__main__":
    main()

def all_reduce(self,
               in_tensor: torch.Tensor,
               op: ReduceOp = ReduceOp.SUM,
               stream=None) -> torch.Tensor:
    if stream is None:
        stream = current_stream()
    self.nccl.ncclAllReduce(buffer_type(in_tensor.data_ptr()),
                            buffer_type(out_tensor.data_ptr()),
                            in_tensor.numel(),
                            ncclDataTypeEnum.from_torch(in_tensor.dtype),
                            ncclRedOpTypeEnum.from_torch(op), self.comm,
                            cudaStream_t(stream.cuda_stream))

def current_stream() -> torch.cuda.Stream:
    from vllm.platforms import current_platform
    global _current_stream
    return _current_stream

def _patched_set_stream(stream: torch.cuda.Stream) -> None:
    global _current_stream
    _current_stream = stream
    prev_set_stream(stream)

torch.cuda.set_stream = _patched_set_stream

import torch
import threading
import time
import vllm.utils

def thread_with_new_stream():
    time.sleep(3)
    s = torch.cuda.Stream()
    print(f"[Thread-A] New CUDA Stream: {s!r}")
    with torch.cuda.stream(s):
        print(f"[Thread-A] Enter stream context: {torch.cuda.current_stream()!r}")
        print(f"[Thread-A] Current stream: {vllm.utils.current_stream()!r}")
        time.sleep(1300)

def thread_print_stream():
    for i in range(400):
        # 注意：每个线程（即每个Python Thread）有自己的stream堆栈
        print(f"[Thread-B] {i}: Current stream: {torch.cuda.current_stream()!r}")
        print(f"[Thread-B] {i}: Current stream: {vllm.utils.current_stream()!r}")
        time.sleep(1)

def main():
    torch.cuda.init()  # 确保 CUDA 初始化

    t1 = threading.Thread(target=thread_with_new_stream)
    t2 = threading.Thread(target=thread_print_stream)

    t1.start()
    t2.start()

    t1.join()
    t2.join()

if __name__ == "__main__":
    main()