add compat for nvtx.range and ExternalStream

python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py

@@ -1630,18 +1630,19 @@ def _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, Grad
         """Run only kernel 2 (GEMM). Caller must have run kernel 1 and zeroed dIndexK_f32."""
         s = _resolve_stream(current_stream)
         _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, current_stream=current_stream)
-        _compile_cache[compile_key](
-            IndexQ,
-            Weights,
-            IndexK,
-            dIndexQ,
-            dWeights,
-            dIndexK_f32,
-            GradSignal,
-            TopkIndices,
-            cutlass.Float32(sm_scale),
-            s,
-        )
+        with torch.cuda.nvtx.range("indexer_backward_dsl_gemm"):
+            _compile_cache[compile_key](
+                IndexQ,
+                Weights,
+                IndexK,
+                dIndexQ,
+                dWeights,
+                dIndexK_f32,
+                GradSignal,
+                TopkIndices,
+                cutlass.Float32(sm_scale),
+                s,
+            )
 
     def _run(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK, AttnScore, IndexScore, TopkIndices, GradLoss, grad_scale, current_stream=None):
         # ``grad_scale`` is a host scalar (Python float / 0-D fp32 tensor)

python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py

@@ -153,46 +153,48 @@ def flash_attn_bwd_sm100(
             block_tile=block_tile,
         )
 
-        flash_attn_bwd_sm100.compile_cache[compile_key] = cute.compile(
-            kernel_obj,
+        with torch.cuda.nvtx.range("flash_attn_bwd_sm100_compile"):
+            flash_attn_bwd_sm100.compile_cache[compile_key] = cute.compile(
+                kernel_obj,
+                problem_shape,
+                q_tensor,
+                kv_tensor,
+                out_tensor,
+                dout_tensor,
+                lse_tensor,
+                attn_sink_tensor,
+                topk_idxs_tensor,
+                topk_length_tensor,
+                dq_tensor,
+                dkv_tensor,
+                d_sink_tensor,
+                workspace_LSE_OdO_tensor,
+                workspace_dKV_tensor,
+                softmax_scale,
+                current_stream,
+                options=compile_options(),
+            )
+
+    with torch.cuda.nvtx.range("flash_attn_bwd_sm100_kernel"):
+        flash_attn_bwd_sm100.compile_cache[compile_key](
             problem_shape,
-            q_tensor,
-            kv_tensor,
-            out_tensor,
-            dout_tensor,
-            lse_tensor,
-            attn_sink_tensor,
-            topk_idxs_tensor,
-            topk_length_tensor,
-            dq_tensor,
-            dkv_tensor,
-            d_sink_tensor,
-            workspace_LSE_OdO_tensor,
-            workspace_dKV_tensor,
+            q,
+            kv,
+            out,
+            dout,
+            lse,
+            attn_sink,
+            topk_idxs,
+            topk_length,
+            dq,
+            dkv,
+            d_sink,
+            workspace_LSE_OdO,
+            workspace_dKV,
             softmax_scale,
             current_stream,
-            options=compile_options(),
         )
 
-    flash_attn_bwd_sm100.compile_cache[compile_key](
-        problem_shape,
-        q,
-        kv,
-        out,
-        dout,
-        lse,
-        attn_sink,
-        topk_idxs,
-        topk_length,
-        dq,
-        dkv,
-        d_sink,
-        workspace_LSE_OdO,
-        workspace_dKV,
-        softmax_scale,
-        current_stream,
-    )
-
     return dq, dkv, d_sink
 
 

python/cudnn/deepseek_sparse_attention/utils/runtime.py

@@ -8,6 +8,18 @@
 import cuda.bindings.driver as cuda
 
 
+if hasattr(torch.cuda, "nvtx") and not hasattr(torch.cuda.nvtx, "range"):
+    @contextmanager
+    def _nvtx_range(msg: str) -> Iterator[None]:
+        torch.cuda.nvtx.range_push(msg)
+        try:
+            yield
+        finally:
+            torch.cuda.nvtx.range_pop()
+
+    torch.cuda.nvtx.range = _nvtx_range
+
+
 @lru_cache(maxsize=None)
 def device_major() -> int:
     return torch.cuda.get_device_capability()[0]
@@ -56,5 +68,9 @@ def torch_stream_context(current_stream: Optional[cuda.CUstream] = None) -> Iter
     if current_stream is None:
         yield
         return
+    if not hasattr(torch.cuda, "ExternalStream") and hasattr(torch.cuda, "get_stream_from_external"):
+        with torch.cuda.stream(torch.cuda.get_stream_from_external(int(current_stream))):
+            yield
+        return
     with torch.cuda.stream(torch.cuda.ExternalStream(int(current_stream))):
         yield