Intro async flag and use current stream avoid stream sync

[Chao1Han]

Refer pytorch/pytorch#147820 pytorch/pytorch#150398
To launch kernels on the current stream and reduce the CPU overhead introduced by recordStream, an async option is introduced.

For example, in an allreduce operation between two ranks:

• rank0 corresponds to device0, using the current device's stream0 to create the communicator and preserving stream0.

When async = true:

• Both rank0 and rank1 perform the collective using stream0, which is associated with the communicator.
• To prevent potential reads by stream0 from unready tensors (e.g., from rank1), synchronization with the current stream is required.
• After the collective completes, to prevent premature release of the input tensors, recordStream must be used for stream tracking, or the tensors need to be temporarily stored (e.g., in reduce_scatter or all2all).

When async = false:

• Both rank0 and rank1 use their respective current streams for collectives (i.e., rank0 uses stream0, rank1 uses stream1).
• In this case, the collective op handles synchronization implicitly.

Previously, we defaulted to async = true. Now, the async option is explicitly introduced and set to false by default, leveraging the current stream to avoid the overhead of stream synchronization.

[Chao1Han]

[zhangxiaoli73]

Do we still need sync if compute stream is used for communication?

[Chao1Han]

not needed, use current stream means async=false. line 632 will return null rather then work and frontend will check if need call wait()
https://github.com/pytorch/pytorch/blob/4273e5d15cfcb282b2795684874ea439d8620999/torch/distributed/distributed_c10d.py#L2882-L2887

[zhangxiaoli73]

In some special case like FSDP2, same communicator will have different current stream on each device. Then how to make sure to return correct compute stream for communication?

[zhangxiaoli73]

Why is it to be a warning? I think info should be more suitable. Warning usually means something not so safe and may cause unpredictable result.

[zhangxiaoli73]

Should you put code in 63823e4#diff-29271b6f1608f7ad940c9cd242ce24dcc68bba932348e39dc0b524604cc78c6aR568, then you don't need construct StreamKey twice.

[zhangxiaoli73]

Suggested change

	auto stream = asyncOp ? xcclStreamsMap_.at(StreamKey).first
	if (asyncOp) {
	stream = xcclStreamsMap_.at(key).first;
	cclstream = xcclStreamsMap_.at(key).second;
	syncStream(device, xcclEventsMap_[key], stream);
	} else {
	current_stream = at::xpu::getCurrentXPUStream(device.index())
	streamkey = stream + current_stream .id()
	if (xcclStreamsMap_.find(streamkey) != xcclStreamsMap_.end()) {
	stream = xcclStreamsMap_.at(streamkey ).first;
	cclstream = xcclStreamsMap_.at(streamkey).second;
	} else {
	// update xcclStreamsMap_ with current stream key
	cclstream = std::make_unique<ccl::stream>(ccl::create_stream(current_stream.queue()));
	std::lock_guard<std::mutex> lock(mutex_);
	xcclStreamsMap_.emplace(
	StreamKey, std::make_pair(at::xpu::XPUStream(current_stream), *cclstream));
	}
	}

[Chao1Han]

change done