[Bug/Correctness] Hardcoded device_id=0 + missing CUDAGuard can break multi-GPU correctness (wrong hw_info / stream mismatch)

[red1239109-cmd]

Hi maintainers,

While reviewing the FMHA forward runner integration, I noticed two correctness issues that can break execution on multi-GPU setups (and can also create subtle stream/device mismatches):

1. Hardcoded GPU selection (device_id = 0)

In run_fmha_fwd, the hardware info is pinned to GPU0:

cutlass::KernelHardwareInfo hw_info;
hw_info.device_id = 0;
hw_info.sm_count =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);

If tensors (q/k/v/o/lse) live on a non-zero device, this will query the wrong SM count and may also lead to launching with incorrect hardware assumptions.

2. Missing CUDA device guard (at::cuda::CUDAGuard) and stream/device alignment

The code uses at::cuda::getCurrentCUDAStream() at the end:

CUTLASS_CHECK(op.run(at::cuda::getCurrentCUDAStream()));

but does not guard/set the current device to match q.device() (or any input tensor). In multi-GPU scenarios, the “current device” may differ from the tensor’s device, leading to:

wrong stream/device being used

incorrect hw_info.device_id / sm_count

potential launch failures or silent misbehavior

Suggested fix

Use a device guard based on an input tensor (e.g., q) and set hw_info.device_id accordingly:

#include <ATen/cuda/CUDAGuard.h>

at::cuda::CUDAGuard device_guard(q.device());
const int dev = at::cuda::current_device();

cutlass::KernelHardwareInfo hw_info;
hw_info.device_id = dev;
hw_info.sm_count =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(dev);

CUTLASS_CHECK(op.run(at::cuda::getCurrentCUDAStream()));

This ensures:

correct device is active

stream matches the tensor device context

hardware info queries the right GPU

Why this matters

Even if most users run single-GPU, multi-GPU is common in training/inference servers. Hardcoding GPU0 + missing guards can produce correctness issues that are hard to diagnose (especially when the failure is not immediate).

If you'd like, I can provide a small repro snippet that places q/k/v/o on cuda:1 and shows the mismatch.

Thanks!

[ashgooch074-svg]

Repro snippet please

[red1239109-cmd]

Here is the reproduction setup and the smoking gun debug logic.

1. Python Repro Setup (Environment Mismatch)

This script sets up the condition where the tensor resides on cuda:1, but the process context is stuck on cuda:0.

import torch

# IMPORTANT: run on a machine with >=2 GPUs
assert torch.cuda.device_count() >= 2, "Need at least 2 GPUs for this repro"

# 1. Force the current process context to GPU 0
torch.cuda.set_device(0)

# 2. Allocate input tensors strictly on GPU 1
# (Simulating a multi-GPU pipeline where data is scattered)
x = torch.randn(1024, device="cuda:1")

print(f"Current torch.cuda.current_device(): {torch.cuda.current_device()}")  # Output: 0
print(f"Input Tensor Device: {x.device}")                                    # Output: cuda:1

# 3. Call the FlashMLA extension here
# Since the extension lacks CUDAGuard(x.device) and hardcodes device_id=0, 
# it will incorrectly initialize hardware info for GPU0 and launch on GPU0's stream,
# leading to an illegal memory access or silent corruption.

# import flash_mla_cuda
# flash_mla_cuda.run_fmha_fwd(..., x, ...) 

2. The Undeniable Proof (C++ Debug Snippet)
Since this is a race/context issue, the most deterministic way to confirm the bug is to add this tiny print inside your C++ entry point (run_fmha_fwd or equivalent dispatcher).
Please insert this check before the kernel launch:
// Debugging the mismatch
TORCH_CHECK(q.is_cuda());
int current_dev = -1;
cudaGetDevice(&current_dev); // or at::cuda::current_device()

printf("[DEBUG] Input Tensor Device: %d | Current CUDA Device: %d | hw_info.device_id: %d\n", 
       (int)q.get_device(), current_dev, hw_info.device_id);

// Without CUDAGuard(q.device()), you will see:
// [DEBUG] Input Tensor Device: 1 | Current CUDA Device: 0 | hw_info.device_id: 0
// This mismatch confirms the kernel is launching on the wrong device context.

This confirms that despite the data being on GPU1, the kernel logic operates under the assumption of GPU0.