Fix CuTeDSL from_dlpack device mismatch in CUDA graph capture (v2)

Previous fix (set_device) caused 'Capture must end on the same stream'.
New fix: wrap tensor in _DLPatchTensor during graph capture, which forces
dl_device in __dlpack__ to bypass the device check without changing the stream.

This enables CUDA graph capture on all 8 GPUs, not just cuda:0.

dsv4/ops/gemm_runner.py

@@ -26,6 +26,32 @@ from dsv4.ops.layouts import (
     round_up,
 )
 
+
+class _DLPatchTensor:
+    """Wrapper that patches __dlpack__ to force dl_device during CUDA graph capture.
+    
+    from_dlpack checks torch.cuda.current_device() against the tensor's device.
+    Inside CUDA graph capture on non-default GPUs, current_device() may not match.
+    This wrapper forces dl_device in __dlpack__, bypassing the device check.
+    
+    DO NOT use torch.cuda.set_device() inside graph capture — it changes the stream
+    and causes 'Capture must end on the same stream it began on'.
+    """
+    __slots__ = ('_t',)
+    def __init__(self, t):
+        self._t = t
+    def __dlpack__(self, *args, **kwargs):
+        kwargs['dl_device'] = (1, self._t.device.index)  # kDLCUDA=1
+        try:
+            return self._t.__dlpack__(*args, **kwargs)
+        except TypeError:
+            return self._t.__dlpack__(*args, **{k: v for k, v in kwargs.items() if k != 'dl_device'})
+    def __dlpack_device__(self):
+        return (1, self._t.device.index)
+    def __getattr__(self, name):
+        return getattr(self._t, name)
+
+
 # Cache compiled kernels + pre-allocated workspace by cache_key
 # Each entry: {'compiled': callable, 'workspace': Tensor, 'workspace_size': int}
 #
@@ -99,10 +125,12 @@ def warmup_compilation(num_experts, K_packed, N_packed, device,
     )
     
     def to_cute(t):
-        # Ensure current CUDA device matches tensor device (required for from_dlpack
-        # inside CUDA graph capture, where torch.cuda.current_device() may not match)
-        if t.is_cuda and t.device.index != torch.cuda.current_device():
-            torch.cuda.set_device(t.device.index)
+        # Fix: from_dlpack checks torch.cuda.current_device() against tensor device.
+        # Inside CUDA graph capture on non-default GPUs, current_device() may not match.
+        # We wrap the tensor to force dl_device in __dlpack__, bypassing the check.
+        # DO NOT use torch.cuda.set_device() inside graph capture — it changes the stream.
+        if hasattr(torch.cuda, 'is_current_stream_capturing') and torch.cuda.is_current_stream_capturing():
+            t = _DLPatchTensor(t)
         ct = cutlass_torch.from_dlpack(t)
         return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
     
@@ -207,10 +235,8 @@ def run_nvfp4_grouped_gemm(
         )
 
         def to_cute(t):
-            # Ensure current CUDA device matches tensor device (required for from_dlpack
-            # inside CUDA graph capture, where torch.cuda.current_device() may not match)
-            if t.is_cuda and t.device.index != torch.cuda.current_device():
-                torch.cuda.set_device(t.device.index)
+            if hasattr(torch.cuda, 'is_current_stream_capturing') and torch.cuda.is_current_stream_capturing():
+                t = _DLPatchTensor(t)
             ct = cutlass_torch.from_dlpack(t)
             return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
 
@@ -258,10 +284,12 @@ def run_nvfp4_grouped_gemm(
     # This is cheap (metadata only, no GPU work) and avoids stale
     # references to tensors from previous calls that may have been freed.
     def to_cute(t):
-        # Ensure current CUDA device matches tensor device (required for from_dlpack
-        # inside CUDA graph capture, where torch.cuda.current_device() may not match)
-        if t.is_cuda and t.device.index != torch.cuda.current_device():
-            torch.cuda.set_device(t.device.index)
+        # Fix: from_dlpack checks torch.cuda.current_device() against tensor device.
+        # Inside CUDA graph capture on non-default GPUs, current_device() may not match.
+        # We wrap the tensor to force dl_device in __dlpack__, bypassing the check.
+        # DO NOT use torch.cuda.set_device() inside graph capture — it changes the stream.
+        if hasattr(torch.cuda, 'is_current_stream_capturing') and torch.cuda.is_current_stream_capturing():
+            t = _DLPatchTensor(t)
         ct = cutlass_torch.from_dlpack(t)
         return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
     
@@ -340,10 +368,12 @@ def warmup_fused_swiglu_compilation(num_experts, K_packed, N_packed, device,
     )
     
     def to_cute(t):
-        # Ensure current CUDA device matches tensor device (required for from_dlpack
-        # inside CUDA graph capture, where torch.cuda.current_device() may not match)
-        if t.is_cuda and t.device.index != torch.cuda.current_device():
-            torch.cuda.set_device(t.device.index)
+        # Fix: from_dlpack checks torch.cuda.current_device() against tensor device.
+        # Inside CUDA graph capture on non-default GPUs, current_device() may not match.
+        # We wrap the tensor to force dl_device in __dlpack__, bypassing the check.
+        # DO NOT use torch.cuda.set_device() inside graph capture — it changes the stream.
+        if hasattr(torch.cuda, 'is_current_stream_capturing') and torch.cuda.is_current_stream_capturing():
+            t = _DLPatchTensor(t)
         ct = cutlass_torch.from_dlpack(t)
         return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
     
@@ -441,10 +471,8 @@ def run_fused_swiglu_grouped_gemm(
         )
 
         def to_cute(t):
-            # Ensure current CUDA device matches tensor device (required for from_dlpack
-            # inside CUDA graph capture, where torch.cuda.current_device() may not match)
-            if t.is_cuda and t.device.index != torch.cuda.current_device():
-                torch.cuda.set_device(t.device.index)
+            if hasattr(torch.cuda, 'is_current_stream_capturing') and torch.cuda.is_current_stream_capturing():
+                t = _DLPatchTensor(t)
             ct = cutlass_torch.from_dlpack(t)
             return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
 
@@ -486,10 +514,12 @@ def run_fused_swiglu_grouped_gemm(
     workspace = entry['workspace']
     
     def to_cute(t):
-        # Ensure current CUDA device matches tensor device (required for from_dlpack
-        # inside CUDA graph capture, where torch.cuda.current_device() may not match)
-        if t.is_cuda and t.device.index != torch.cuda.current_device():
-            torch.cuda.set_device(t.device.index)
+        # Fix: from_dlpack checks torch.cuda.current_device() against tensor device.
+        # Inside CUDA graph capture on non-default GPUs, current_device() may not match.
+        # We wrap the tensor to force dl_device in __dlpack__, bypassing the check.
+        # DO NOT use torch.cuda.set_device() inside graph capture — it changes the stream.
+        if hasattr(torch.cuda, 'is_current_stream_capturing') and torch.cuda.is_current_stream_capturing():
+            t = _DLPatchTensor(t)
         ct = cutlass_torch.from_dlpack(t)
         return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))