Install CUDAGuard before Arch() in api entry points (fixes #158)

csrc/api/dense_decode.h

@@ -22,6 +22,13 @@ dense_attn_decode_interface(
     std::optional<at::Tensor> &tile_scheduler_metadata,   // num_sm_parts x (DecodingSchedMetaSize/4)
     std::optional<at::Tensor> &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));

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);

csrc/api/sparse_fwd.h

@@ -108,7 +108,14 @@ static std::vector<at::Tensor> sparse_attn_prefill_interface(
     const std::optional<at::Tensor> &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<at::Tensor> 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);

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"