diff --git a/csrc/api/dense_decode.h b/csrc/api/dense_decode.h index 7df178a6..457279c3 100644 --- a/csrc/api/dense_decode.h +++ b/csrc/api/dense_decode.h @@ -22,6 +22,13 @@ dense_attn_decode_interface( std::optional &tile_scheduler_metadata, // num_sm_parts x (DecodingSchedMetaSize/4) std::optional &num_splits // batch_size + 1 ) { + // Set the active CUDA device to match the input tensor BEFORE probing hardware + // properties via Arch(). Arch() reads at::cuda::getCurrentDeviceProperties() and + // will return the wrong SM count / capability if the current device differs from + // the tensor's device (multi-GPU / heterogeneous setups). See issue #158. + TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor"); + at::cuda::CUDAGuard device_guard{(char)q.get_device()}; + // Check arch Arch arch = Arch(); if (!arch.is_sm90a()) { @@ -84,8 +91,6 @@ dense_attn_decode_interface( KU_CHECK_SHAPE(tile_scheduler_metadata, num_sm_parts, DecodingSchedMetaSize/sizeof(int)); KU_CHECK_SHAPE(num_splits, batch_size+1); - at::cuda::CUDAGuard device_guard{(char)q.get_device()}; - auto opts = q.options(); at::Tensor out = torch::empty({batch_size, num_heads, q_seq_per_hk, head_size_v}, opts); at::Tensor lse = torch::empty({batch_size, num_heads, q_seq_per_hk}, opts.dtype(at::kFloat)); diff --git a/csrc/api/sparse_decode.h b/csrc/api/sparse_decode.h index 6df5c841..e7d4944a 100644 --- a/csrc/api/sparse_decode.h +++ b/csrc/api/sparse_decode.h @@ -197,6 +197,14 @@ sparse_attn_decode_interface( ) { using bf16 = cutlass::bfloat16_t; + // Set the active CUDA device to match the input tensor BEFORE probing hardware + // properties via Arch() or dispatching to an SM-specific impl. Arch() reads the + // *current* device's properties, so without this guard the dispatcher can pick + // the wrong kernel (e.g. SM90 impl while q lives on an SM100 device) and + // num_sm_parts is computed from the wrong SM count. See issue #158. + TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor"); + at::cuda::CUDAGuard device_guard{(char)q.get_device()}; + // Check the architecture Arch arch = Arch(); @@ -309,7 +317,6 @@ sparse_attn_decode_interface( KU_CHECK_SHAPE(extra_indices, b, s_q, extra_topk); KU_CHECK_SHAPE(extra_topk_length, b); - at::cuda::CUDAGuard device_guard{(char)q.get_device()}; auto opts = q.options(); at::Tensor out = torch::empty({b, s_q, h_q, d_v}, opts); diff --git a/csrc/api/sparse_fwd.h b/csrc/api/sparse_fwd.h index 66d7111e..a785762f 100644 --- a/csrc/api/sparse_fwd.h +++ b/csrc/api/sparse_fwd.h @@ -108,7 +108,14 @@ static std::vector sparse_attn_prefill_interface( const std::optional &topk_length ) { using bf16 = cutlass::bfloat16_t; - + + // Set the active CUDA device to match the input tensor BEFORE probing hardware + // properties via Arch(). Without this guard, arch.num_sms and the SM90a/SM100f + // dispatch reflect whatever device happens to be current, not q's device. See + // issue #158. + TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor"); + at::cuda::CUDAGuard device_guard{(char)q.get_device()}; + Arch arch = Arch(); bool is_sm90a = arch.is_sm90a(); bool is_sm100f = arch.is_sm100f(); @@ -156,7 +163,6 @@ static std::vector sparse_attn_prefill_interface( KU_CHECK_LAST_DIM_CONTIGUOUS(topk_length); // Allocate results and buffers - at::cuda::CUDAGuard device_guard{(char)q.get_device()}; auto opts = q.options(); at::Tensor out = torch::empty({s_q, h_q, d_v}, opts); diff --git a/tests/test_multi_gpu_device_guard.py b/tests/test_multi_gpu_device_guard.py new file mode 100644 index 00000000..dd7daf32 --- /dev/null +++ b/tests/test_multi_gpu_device_guard.py @@ -0,0 +1,95 @@ +""" +Regression tests for issue #158: kernels must honor the input tensor's device, +not the process's current CUDA device. + +Before the fix, ``Arch arch = Arch();`` was constructed in the C++ API entry +points before ``at::cuda::CUDAGuard`` was installed. ``Arch()`` reads the +*current* device's properties via ``at::cuda::getCurrentDeviceProperties()``, +so when the process's current device differed from ``q.device()`` (a common +setup in multi-GPU inference servers and distributed training), the dispatcher +could: + + * pick the wrong SM-specific impl (SM90 vs SM100) in heterogeneous boxes, + * compute ``num_sm_parts`` from the wrong SM count, and + * launch kernels on the wrong stream. + +These tests pin ``torch.cuda.set_device(0)`` but place every tensor on +``cuda:1`` and assert that the result is bit-/numerically-equal to running +the same workload with matched current device. They are skipped when fewer +than two CUDA devices are visible. +""" + +import pytest +import torch + +import flash_mla + + +def _require_two_gpus(): + if torch.cuda.device_count() < 2: + pytest.skip("Requires at least 2 CUDA devices") + + +def _dense_decode_inputs(device: torch.device, dtype=torch.bfloat16): + torch.manual_seed(0) + b, s_q, h_q, h_kv, d_qk, d_v = 1, 1, 128, 1, 576, 512 + page_block_size = 64 + num_blocks = 64 + s_kv = 512 + + q = (torch.randn(b, s_q, h_q, d_qk, dtype=dtype, device=device) / 10).clamp_(-1, 1) + kcache = (torch.randn(num_blocks, page_block_size, h_kv, d_qk, dtype=dtype, device=device) / 10).clamp_(-1, 1) + cache_seqlens = torch.full((b,), s_kv, dtype=torch.int32, device=device) + block_table = torch.arange(b * (s_kv // page_block_size), dtype=torch.int32, device=device).view(b, -1) + return q, kcache, cache_seqlens, block_table, d_v + + +def _run_dense_decode(q, kcache, cache_seqlens, block_table, d_v): + sched_meta, _ = flash_mla.get_mla_metadata() + out, lse = flash_mla.flash_mla_with_kvcache( + q=q, + k_cache=kcache, + block_table=block_table, + cache_seqlens=cache_seqlens, + head_dim_v=d_v, + tile_scheduler_metadata=sched_meta, + causal=False, + ) + return out, lse + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs >= 2 GPUs") +def test_dense_decode_respects_input_device_when_current_device_differs(): + """Regression: dense decode must work when q lives on cuda:1 but current device is 0.""" + _require_two_gpus() + + # Reference: current device == tensor device. + torch.cuda.set_device(1) + q_ref, k_ref, sq_ref, bt_ref, d_v = _dense_decode_inputs(torch.device("cuda:1")) + out_ref, lse_ref = _run_dense_decode(q_ref, k_ref, sq_ref, bt_ref, d_v) + torch.cuda.synchronize(1) + + # Under test: current device mismatches tensor device. Without the fix this + # path either picked the wrong SM impl, queried the wrong SM count, or + # launched on cuda:0's stream. + torch.cuda.set_device(0) + assert torch.cuda.current_device() == 0 + q_mm, k_mm, sq_mm, bt_mm, _ = _dense_decode_inputs(torch.device("cuda:1")) + out_mm, lse_mm = _run_dense_decode(q_mm, k_mm, sq_mm, bt_mm, d_v) + torch.cuda.synchronize(1) + + assert out_mm.device == torch.device("cuda:1"), f"output landed on {out_mm.device}" + assert lse_mm.device == torch.device("cuda:1") + torch.testing.assert_close(out_mm, out_ref, rtol=0, atol=0) + torch.testing.assert_close(lse_mm, lse_ref, rtol=0, atol=0) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs >= 2 GPUs") +def test_dense_decode_current_device_unchanged_after_call(): + """The guard must restore the caller's current device on exit.""" + _require_two_gpus() + + torch.cuda.set_device(0) + q, kcache, sq, bt, d_v = _dense_decode_inputs(torch.device("cuda:1")) + _run_dense_decode(q, kcache, sq, bt, d_v) + assert torch.cuda.current_device() == 0, "CUDAGuard leaked: current device changed"