Greptile Summary

This PR removes an unnecessary device-to-device copy in the CUDA backend's Reduce collective by requiring the input tensor to already reside in symmetric (VMM-backed) memory, which enables the multimem ld_reduce kernel to read directly from the multicast VA without first staging data into a separate symmetric buffer. The optimizer-visible change is that callers must now annotate the reduce-input TensorView with MemoryType::Symmetric when using the CUDA backend.

Key changes:

• SymmetricMemoryForReduce now wraps the caller-supplied buffer directly instead of allocating a new symmetric tensor and scheduling a cudaMemcpyAsync copy into it.
• postReduceWithCudaBackend drops the input parameter and the copy, tightens the root assertion to require output to be defined, and reads scalar type from the handle's buffer.
• The handle(Communication*) post path in the evaluator correctly switches the cache key to input_tensor for Reduce so that the symmetric handle is keyed by the symmetric input.
• However, the matching handle(Wait* wait) path was not updated. On the root rank, the wait handler resolves cache_buffer to output_tensor (which is defined for root), while the post handler used input_tensor. SymmetricMemoryHandleCache::get() will not find the existing entry and will attempt to construct a new SymmetricMemoryForReduce from the non-symmetric output tensor, which immediately fails inside SymmetricTensor::validate(). The root rank will crash at wait time for every Reduce with the CUDA backend.

Confidence Score: 4/5

Not safe to merge as-is: the root rank will crash at the Wait phase due to a missing cache-key fix in the wait handler.

The core optimization in cuda_p2p.cpp and ipc_handle.cpp is correct and clean. The evaluator post path is also correctly updated. However, the wait handler in evaluator.cpp was not updated to mirror the post path's Reduce-specific cache-key logic, causing the root rank to look up a different cache entry during Wait and triggering a runtime NVF_CHECK failure. Score 4 reflects one clear P1 defect that must be resolved before merging.

csrc/host_ir/evaluator.cpp — the Wait handler (around line 527) needs the same if (Reduce) { cache_buffer = input_tensor; } guard that was added to the post path.

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant E as HostIrEvaluator
    participant C as HandleCache
    participant R as SymMemForReduce
    participant K as CUDA Reduce Kernel

    note over E: Post phase (handle Communication)
    E->>C: get(input_tensor, comm, root)
    C->>R: construct with input_tensor (symmetric)
    R-->>C: handle created
    C-->>E: reduce_handle ptr

    note over K: Before this PR: cudaMemcpyAsync input to sym_buf (removed)
    E->>K: postReduceWithCudaBackend(output, handle, stream, root)
    K->>K: All ranks signal kInProgress
    K->>K: Root waits for all non-roots
    K->>K: launchMulticastReduceKernel via mc_ptr to output
    K->>K: Root signals kIdle to non-roots

    note over E: Wait phase (handle Wait) - BUG on root rank
    E->>C: get(output_tensor, comm, root) - mismatched key on root
    C->>R: construct with output_tensor (NOT symmetric)
    R-->>C: SymmetricTensor validate FAILS - crash

_{Reviews (2): Last reviewed commit: "Lint" | Re-trigger Greptile}

!test