Repository: pytorch
Issue: One Shot Low-Latency kernel is 2x higher than expected
Labels: module: performance,oncall: distributed

### 🐛 Describe the bug

I am trying out the new low latency one shot kernels and found that they are about double as what #114001 showed in the PR description. My results hover around `17us` while #114001 showed between `6-15us`

cc: @yifuwang

## My Results
![image](https://github.com/user-attachments/assets/12a06b30-38ef-4bcb-a554-bba4c69ff4d5)

## Expected Results from #114001 
![image](https://github.com/user-attachments/assets/29849751-70ee-4721-98ff-090bada9389e)



Attached below is my testing reprod. I have used intra node counter to enable that the low latency kernel is running. I have also ensured to use bf16 and reduce_sum.

### command for running nccl ring
```bash
ENABLE_INTRA_NODE_COMM=0 NCCL_ALGO=RING python3 -m torch.distributed.run --standalone --nproc_per_node=8 low_latency.py
```
### command for running low latency
```bash
ENABLE_INTRA_NODE_COMM=1 python3 -m torch.distributed.run --standalone --nproc_per_node=8 low_latency.py
```

## Reprod
```python
import os
import datetime
local_rank = int(os.environ.get('LOCAL_RANK', 0))

import torch
import torch.distributed as dist
from torch._C._distributed_c10d import _get_intra_node_comm_usage_counter
import matplotlib.pyplot as plt
import pandas as pd

# Create a timestamp for the output directory
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
output_dir = f'results/low_latency/{timestamp}'
os.makedirs(output_dir, exist_ok=True)

# Set the device for each process.
torch.cuda.set_device(local_rank)

# Initialize the distributed environment.
dist.init_process_group(backend='nccl', init_method='env://')  # Use 'gloo' if you're using CPU
rank = dist.get_rank()
size = dist.get_world_size()

# if enable intra node comm env var is set = 1, then print that it is set
if os.environ.get("ENABLE_INTRA_NODE_COMM", "0") == "1":
    print("Intra node comm is enabled")
else:
    print("Intra node comm is not enabled")

# Set the device for each process.
device = torch.device('cuda', rank)
global_counter = 0

def benchmark_all_reduce(message_size):
    # Create a tensor of the specified size with bfloat16 data type.
    tensor = torch.ones(message_size, dtype=torch.bfloat16, device=device)
    dist.barrier()  # Synchronize all processes before starting the timing.

    latencies = []
    for i in range(10000):
        global global_counter
        global_counter += 1
        torch.cuda.synchronize()
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)

        start_event.record()
        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
        end_event.record()

        torch.cuda.synchronize()
        latency = start_event.elapsed_time(end_event)  # Time in milliseconds
        latencies.append(latency * 1000)  # Convert to microseconds

        # If intra node comm is enabled, the counter should be increased by 1
        if os.environ.get("ENABLE_INTRA_NODE_COMM", "0") == "1":
            assert _get_intra_node_comm_usage_counter() == global_counter, "Intra node comm usage counter mismatch"
        else:
            assert _get_intra_node_comm_usage_counter() == 0, "Intra node comm usage counter mismatch"

    # Disregard the first 20 times and compute the median latency in microseconds.
    latencies_after_warmup = [lat for lat in latencies[20:] if lat <= 50]  # Filter out values greater than 50us
    median_latency_us = torch.median(torch.tensor(latencies_after_warmup)).item()  # Already in microseconds
    
    # Plot the distribution of latencies after the first 20 warmup steps
    plt.figure(figsize=(10, 6))
    plt.hist(latencies_after_warmup, bins=25, edgecolor='black')  # Increase the number of bins
    plt.xlabel('Latency (us)')
    plt.ylabel('Frequency')
    plt.xlim(0, 50)  # Stop the x-axis at 50
    plt.yscale('log')  # Make the y-axis logarithmic
    plt.title(f'Distribution of Latency (Message Size: {message_size * 2 // 1024} KB)')
    if os.environ.get("ENABLE_INTRA_NODE_COMM", "0") == "1":
        plt.savefig(f'{output_dir}/latency_distribution_low_latency_kernel_{message_size * 2 // 1024}KB.png')
    else:
        plt.savefig(f'{output_dir}/latency_distribution_{message_size * 2 // 1024}KB.png')
    plt.close()

    return median_latency_us

# Message sizes from 1KB to 256KB.
message_sizes_bytes = [2**i for i in range(10, 19)]  # 2^10 bytes to 2^18 bytes (1KB to 256KB)
message_sizes_kb = [size // 1024 for size in message_sizes_bytes]  # Convert to KB
latencies = []

# Run the benchmark for each message size.
for size in message_sizes_bytes:
    message_size = size // 2  # Since we're using bfloat16 (2 bytes per element)
    median_latency_us = benchmark_all_reduce(message_size)
    if rank == 0:  # Collect data from the master process.
        latencies.append(median_latency_us)

# Finalize the distributed environment.
dist.destroy_process_group()

# Save results to CSV if this is the master process.
if rank == 0:
    results = pd.DataFrame({
        'Message Size (KB)': message_sizes_kb,
        'Median Latency (us)': latencies
    })

    # Plot the results.
    plt.figure(figsize=(10, 6))
    plt.plot(message_sizes_kb, latencies, marker='o')
    plt.xlabel('Message Size (KB)')
    plt.ylabel('Median Latency (us)')
    plt.title('All-Reduce Latency Benchmark (BF16)')
    plt.grid(True)
    plt.xscale('log')

    # Set specific ticks for the x-axis.
    plt.xticks([1, 2, 4, 8, 16, 32, 64, 128, 256],
               ['1KB', '2KB', '4KB', '8KB', '16KB', '32KB', '64KB', '128KB', '256KB'])
    
    if os.environ.get("ENABLE_INTRA_NODE_COMM", "0") == "1":
        results.to_csv(f'{output_dir}/latencies_low_latency_kernel.csv', index=False)
        plt.savefig(f'{output_dir}/all_reduce_low_latency_kernel.png')
    else:
        results.to_csv(f'{output_dir}/latencies.csv', index=False)
        plt.savefig(f'{output_dir}/all_reduce.png')
```

### Versions

torch==2.4.0, nvidia-nccl-cu12==2.20.5, cuda-drivers=550.90.07, server=dell XE9680

cc @msaroufim @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o