From 299b3f78015956b597274da604fccf4d07234ec5 Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Tue, 11 Aug 2026 01:53:43 +0200
Subject: [PATCH 1/2] Fix nvexec split pre-cancellation

---
 include/nvexec/stream/split.cuh | 23 ++++++++++++++++++++---
 test/nvexec/split.cpp           | 19 +++++++++++++++++++
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/include/nvexec/stream/split.cuh b/include/nvexec/stream/split.cuh
index d461e3ff1..b599117ec 100644
--- a/include/nvexec/stream/split.cuh
+++ b/include/nvexec/stream/split.cuh
@@ -150,7 +150,11 @@ namespace nv::execution::_strm
     template <class Sender>
     struct sh_state
     {
-      using variant_t          = variant_storage_t<Sender, env_t>;
+      using variant_t = __for_each_completion_signature_t<
+        __concat_completion_signatures_t<__completion_signatures_of_t<Sender, env_t>,
+                                         completion_signatures<set_stopped_t()>>,
+        decayed_tuple_t,
+        __munique<__q<_nullable_variant_t>>::__f>;
       using inner_receiver_t   = receiver<Sender, sh_state>;
       using task_t             = continuation_task<inner_receiver_t, variant_t>;
       using enqueue_receiver_t = stream_enqueue_receiver<env_t, variant_t>;
@@ -165,6 +169,7 @@ namespace nv::execution::_strm
         , data_(malloc_managed<variant_t>(stream_provider_.status_))
         , opstate2_(connect(static_cast<Sender&&>(sndr), inner_receiver_t{*this}))
       {
+        _initialize_stopped();
         if (stream_provider_.status_ == cudaSuccess)
         {
           stream_provider_.status_ = STDEXEC_LOG_CUDA_API(
@@ -186,7 +191,9 @@ namespace nv::execution::_strm
         , env_(host_allocate(this->stream_provider_.status_, ctx_.pinned_resource_, make_env()))
         , opstate2_(connect(static_cast<Sender&&>(sndr),
                             enqueue_receiver_t{env_.get(), data_, task_, ctx.hub_->producer()}))
-      {}
+      {
+        _initialize_stopped();
+      }
 
       ~sh_state()
       {
@@ -213,6 +220,16 @@ namespace nv::execution::_strm
         return _split::_make_env(stop_source_, &const_cast<stream_provider&>(stream_provider_));
       }
 
+      void _initialize_stopped() noexcept
+      {
+        if (data_)
+        {
+          using tuple_t = decayed_tuple_t<set_stopped_t>;
+          index_        = __mapply<__mfind_i<tuple_t>, variant_t>::value;
+          data_->template emplace<tuple_t>(set_stopped_t());
+        }
+      }
+
       void notify() noexcept
       {
         void* const completion_state = static_cast<void*>(this);
@@ -359,7 +376,7 @@ namespace nv::execution::_strm
       return STDEXEC::__transform_completion_signatures_of_t<
         Sender,
         STDEXEC::prop<get_stop_token_t, inplace_stop_token>,
-        STDEXEC::completion_signatures<set_error_t(cudaError_t const &)>,
+        STDEXEC::completion_signatures<set_error_t(cudaError_t const &), set_stopped_t()>,
         _set_value_t,
         _set_error_t>();
     }
diff --git a/test/nvexec/split.cpp b/test/nvexec/split.cpp
index d9ad79a86..a02be4a51 100644
--- a/test/nvexec/split.cpp
+++ b/test/nvexec/split.cpp
@@ -36,6 +36,25 @@ namespace
     REQUIRE(v2 == 42);
   }
 
+  TEST_CASE("nvexec split handles pre-cancellation", "[cuda][stream][adaptors][split]")
+  {
+    nvexec::stream_context  stream_ctx{};
+    ex::inplace_stop_source stop_source;
+    flags_storage_t         flags_storage{};
+    auto                    flags = flags_storage.get();
+
+    stop_source.request_stop();
+
+    auto snd = ex::schedule(stream_ctx.get_scheduler()) | ex::then([flags] { flags.set(); })
+             | exec::split() | ex::write_env(ex::prop{ex::get_stop_token, stop_source.get_token()})
+             | ex::upon_stopped([] { return 42; });
+
+    auto [value] = STDEXEC::sync_wait(std::move(snd)).value();
+
+    REQUIRE(value == 42);
+    REQUIRE(flags_storage.all_unset());
+  }
+
   TEST_CASE("nvexec split can preceed a sender without values", "[cuda][stream][adaptors][split]")
   {
     nvexec::stream_context stream_ctx{};

From 68afb2ec7397e51a3823d933f94530124bf0827c Mon Sep 17 00:00:00 2001
From: Eric Niebler <eniebler@nvidia.com>
Date: Fri, 14 Aug 2026 23:23:48 +0000
Subject: [PATCH 2/2] fix nvexec split pre-cancellation test

---
 test/nvexec/split.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/test/nvexec/split.cpp b/test/nvexec/split.cpp
index a02be4a51..ab0e82dfb 100644
--- a/test/nvexec/split.cpp
+++ b/test/nvexec/split.cpp
@@ -45,8 +45,11 @@ namespace
 
     stop_source.request_stop();
 
-    auto snd = ex::schedule(stream_ctx.get_scheduler()) | ex::then([flags] { flags.set(); })
-             | exec::split() | ex::write_env(ex::prop{ex::get_stop_token, stop_source.get_token()})
+    auto snd = ex::schedule(stream_ctx.get_scheduler())                              //
+             | ex::then([flags] { flags.set(); })                                    //
+             | exec::split()                                                         //
+             | ex::write_env(ex::prop{ex::get_stop_token, stop_source.get_token()})  //
+             | ex::then([] { return 0; })                                            //
              | ex::upon_stopped([] { return 42; });
 
     auto [value] = STDEXEC::sync_wait(std::move(snd)).value();