[ROCm][Perf] Enable layer-aware CSA2 multi-stream overlap for DeepSeek-V4.1-Flash (#57407)

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

tests/kernels/test_compressor_kv_cache.py

@@ -2075,3 +2075,269 @@ def test_fused_kv_insert_split(num_tokens: int, kv_block_size: int):
     # RoPE (last 64): stored as bf16. The kernel recomputes the rotation, so it
     # is bf16-close to the reference rather than bit-exact (cf. test_cutedsl).
     torch.testing.assert_close(recovered[:, NOPE_DIM:], ref[:, NOPE_DIM:])
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm stream coverage")
+def test_v41_rocm_csa2_full_pipeline():
+    """_forward_csa2_full: one fork, then K write and q-side after the join.
+
+    The default chain must run first on the current stream, the compressor
+    chain on aux 0 and the indexer weights projection on aux 1; the join must
+    publish both cache writes before _produce_k / forward_q run.
+    """
+    from vllm.forward_context import ForwardContext, override_forward_context
+    from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention
+    from vllm.models.deepseek_v41.compressor import DeepseekCompressor
+
+    torch.manual_seed(43)
+    num_tokens = 19
+    hidden_states = torch.randn(num_tokens, 1024, device="cuda")
+    positions = torch.arange(7, 26, device="cuda")
+    qr_kv = torch.zeros(num_tokens, 1536, device="cuda")
+    qr = torch.zeros(num_tokens, 4, device="cuda")
+    kv = torch.zeros(num_tokens, 512, device="cuda")
+    q = torch.zeros(num_tokens, 512, device="cuda")
+    latent = torch.ones(num_tokens, 512, dtype=torch.bfloat16, device="cuda")
+    weights = torch.randn(num_tokens, 32, device="cuda")
+    index_q = torch.full((num_tokens, 32, 128), 2.0, device="cuda")
+    index_q_scale = torch.full((num_tokens, 32), 3.0, device="cuda")
+    index_weights_out = torch.full((num_tokens, 32), 4.0, device="cuda")
+    main_cache = torch.full((1, 128, 584), 165, dtype=torch.uint8, device="cuda")
+    index_cache = torch.full((1, 128, 132), 165, dtype=torch.uint8, device="cuda")
+
+    order: list[str] = []
+    streams: dict[str, torch.cuda.Stream] = {}
+
+    compressor = DeepseekCompressor.__new__(DeepseekCompressor)
+    torch.nn.Module.__init__(compressor)
+    compressor.fused_wkv_wgate = SimpleNamespace(
+        weight=torch.randn(1024, 1024, device="cuda")
+    )
+
+    def compressor_forward(kv_score, positions_):
+        order.append("compressor")
+        streams["compressor"] = torch.cuda.current_stream()
+        return latent
+
+    def compressor_insert(latent_, positions_, rotary_emb):
+        order.append("insert_cache")
+        streams["insert_cache"] = torch.cuda.current_stream()
+        main_cache.fill_(5)
+
+    compressor.forward = compressor_forward
+    compressor.insert_cache = compressor_insert
+
+    def weights_proj(hidden_states_):
+        order.append("weights_proj")
+        streams["weights_proj"] = torch.cuda.current_stream()
+        return weights, None
+
+    def produce_k(latent_, positions_, rotary_emb):
+        order.append("produce_k")
+        streams["produce_k"] = torch.cuda.current_stream()
+        index_cache.fill_(7)
+
+    def forward_q(qr_, qr_scale, indexer_weights, positions_, rotary_emb):
+        order.append("forward_q")
+        streams["forward_q"] = torch.cuda.current_stream()
+        return index_q, index_q_scale, index_weights_out
+
+    indexer = SimpleNamespace(
+        weights_proj=weights_proj,
+        _produce_k=produce_k,
+        forward_q=forward_q,
+    )
+    rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda"))
+    aux_streams = [torch.cuda.Stream(), torch.cuda.Stream()]
+
+    def wq_b(qr_, qr_scale):
+        order.append("wq_b")
+        streams["wq_b"] = torch.cuda.current_stream()
+        return q
+
+    def qnorm_insert(q_, kv_, positions_, attn_metadata):
+        order.append("qnorm_insert")
+        streams["qnorm_insert"] = torch.cuda.current_stream()
+        return q_
+
+    attention = SimpleNamespace(
+        compressor=compressor,
+        indexer=indexer,
+        aux_stream_list=aux_streams,
+        ln_events=[torch.cuda.Event() for _ in range(4)],
+        rotary_emb=rotary,
+        indexer_rotary_emb=rotary,
+        n_local_heads=1,
+        head_dim=512,
+        _fused_wqa_wkv_gemm=lambda hidden_states_: qr_kv,
+        _split_qkv_and_norm=lambda qr_kv_: (qr, None, kv),
+        _wq_b_proj=wq_b,
+        _fused_qnorm_rope_kv_insert=qnorm_insert,
+    )
+
+    context = ForwardContext({}, {}, {})
+    with override_forward_context(context):
+        result = DeepseekV41ROCMAiterMLAAttention._forward_csa2_full(
+            attention, hidden_states, positions
+        )
+
+    q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result
+    assert torch.equal(q_out, q)
+    assert torch.equal(kv_out, kv)
+    assert torch.equal(index_q_out, index_q)
+    assert torch.equal(index_q_scale_out, index_q_scale)
+    assert torch.equal(index_weights_out_out, index_weights_out)
+    # Join publishes the aux-0 cache insert and orders the post-join K write.
+    assert torch.equal(main_cache, torch.full_like(main_cache, 5))
+    assert torch.equal(index_cache, torch.full_like(index_cache, 7))
+
+    assert order == [
+        "wq_b",
+        "qnorm_insert",
+        "compressor",
+        "insert_cache",
+        "weights_proj",
+        "produce_k",
+        "forward_q",
+    ]
+    current = torch.cuda.current_stream()
+    for name in ("wq_b", "qnorm_insert", "produce_k", "forward_q"):
+        assert streams[name] == current, name
+    for name in ("compressor", "insert_cache"):
+        assert streams[name] == aux_streams[0], name
+    assert streams["weights_proj"] == aux_streams[1]
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm stream coverage")
+def test_v41_rocm_csa2_reindex_pipeline():
+    """_forward_csa2_reindex: serial projections, then SWA q path vs q-side."""
+    from vllm.forward_context import ForwardContext, override_forward_context
+    from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention
+
+    torch.manual_seed(45)
+    num_tokens = 19
+    hidden_states = torch.randn(num_tokens, 1024, device="cuda")
+    positions = torch.arange(7, 26, device="cuda")
+    qr_kv = torch.zeros(num_tokens, 1536, device="cuda")
+    qr = torch.zeros(num_tokens, 4, device="cuda")
+    kv = torch.zeros(num_tokens, 512, device="cuda")
+    q = torch.zeros(num_tokens, 512, device="cuda")
+    indexer_weights = torch.randn(num_tokens, 32, device="cuda")
+    index_q = torch.full((num_tokens, 32, 128), 2.0, device="cuda")
+    index_q_scale = torch.full((num_tokens, 32), 3.0, device="cuda")
+    index_weights_out = torch.full((num_tokens, 32), 4.0, device="cuda")
+
+    order: list[str] = []
+    streams: dict[str, torch.cuda.Stream] = {}
+
+    def forward_q(qr_, qr_scale, indexer_weights_, positions_, rotary_emb):
+        order.append("forward_q")
+        streams["forward_q"] = torch.cuda.current_stream()
+        return index_q, index_q_scale, index_weights_out
+
+    indexer = SimpleNamespace(forward_q=forward_q)
+    rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda"))
+    aux0 = torch.cuda.Stream()
+
+    def wq_b(qr_, qr_scale):
+        order.append("wq_b")
+        streams["wq_b"] = torch.cuda.current_stream()
+        return q
+
+    def qnorm_insert(q_, kv_, positions_, attn_metadata):
+        order.append("qnorm_insert")
+        streams["qnorm_insert"] = torch.cuda.current_stream()
+        return q_
+
+    attention = SimpleNamespace(
+        indexer=indexer,
+        aux_stream_list=[aux0],
+        ln_events=[torch.cuda.Event() for _ in range(2)],
+        indexer_rotary_emb=rotary,
+        n_local_heads=1,
+        head_dim=512,
+        _run_parallel_input_projections=lambda hidden_states_: (
+            qr_kv,
+            None,
+            indexer_weights,
+        ),
+        _split_qkv_and_norm=lambda qr_kv_: (qr, None, kv),
+        _wq_b_proj=wq_b,
+        _fused_qnorm_rope_kv_insert=qnorm_insert,
+    )
+
+    context = ForwardContext({}, {}, {})
+    with override_forward_context(context):
+        result = DeepseekV41ROCMAiterMLAAttention._forward_csa2_reindex(
+            attention, hidden_states, positions
+        )
+
+    q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result
+    assert torch.equal(q_out, q)
+    assert torch.equal(kv_out, kv)
+    assert torch.equal(index_q_out, index_q)
+    assert torch.equal(index_q_scale_out, index_q_scale)
+    assert torch.equal(index_weights_out_out, index_weights_out)
+
+    assert order == ["wq_b", "qnorm_insert", "forward_q"]
+    current = torch.cuda.current_stream()
+    for name in ("wq_b", "qnorm_insert"):
+        assert streams[name] == current, name
+    assert streams["forward_q"] == aux0
+
+
+@pytest.mark.parametrize("quant_tuple", [False, True])
+def test_v41_indexer_forward_q(quant_tuple):
+    """forward_q runs wq_b then the fused RoPE/quant and unpacks both shapes."""
+    from unittest.mock import patch
+
+    from vllm.models.deepseek_v41.attention import DeepseekV4Indexer
+
+    num_tokens, n_head, head_dim = 5, 32, 128
+    qr = torch.randn(num_tokens, 1280)
+    q = torch.randn(num_tokens, n_head * head_dim)
+    positions = torch.arange(num_tokens)
+    indexer_weights = torch.randn(num_tokens, n_head)
+    rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64))
+    fake = SimpleNamespace(
+        n_head=n_head,
+        head_dim=head_dim,
+        softmax_scale=head_dim**-0.5,
+        use_fp4_kv=False,
+        indexer_weights_dtype=torch.float32,
+        _wq_b_proj=lambda qr_, qr_scale: q,
+    )
+    q_quant = torch.randn(num_tokens, n_head, head_dim)
+    q_scale = torch.randn(num_tokens, n_head)
+    weights_out = torch.randn(num_tokens, n_head)
+    if quant_tuple:
+        fused_return = ((q_quant, q_scale), weights_out)
+        expected = (q_quant, q_scale, weights_out)
+    else:
+        fused_return = (q_quant, weights_out)
+        expected = (q_quant, None, weights_out)
+
+    with patch(
+        "vllm.models.deepseek_v41.attention.fused_indexer_q_rope_quant",
+        return_value=fused_return,
+    ) as mock_fused:
+        result = DeepseekV4Indexer.forward_q(
+            fake, qr, None, indexer_weights, positions, rotary
+        )
+
+    assert torch.equal(result[0], expected[0])
+    assert (result[1] is None) == (expected[1] is None)
+    if expected[1] is not None:
+        assert torch.equal(result[1], expected[1])
+    assert torch.equal(result[2], expected[2])
+    args = mock_fused.call_args.args
+    assert args[0] is positions
+    torch.testing.assert_close(args[1], q.view(-1, n_head, head_dim))
+    assert args[2] is rotary.cos_sin_cache
+    assert args[3] is indexer_weights
+    assert args[4] == head_dim**-0.5
+    assert args[5] == n_head**-0.5
+    assert mock_fused.call_args.kwargs == {
+        "use_fp4": False,
+        "weights_out_dtype": torch.float32,
+    }

vllm/models/deepseek_v41/amd/rocm.py

@@ -29,6 +29,7 @@ from vllm.models.deepseek_v41.sparse_mla import (
 from vllm.platforms import current_platform
 from vllm.platforms.rocm import _ON_GFX950
 from vllm.triton_utils import tl, triton
+from vllm.utils.multi_stream_utils import execute_in_parallel
 from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
 )
@@ -529,6 +530,9 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
                     "this rank; PP splits inside a v4.1 index-sharing group "
                     "are not supported."
                 )
+        if self.compressor is None and self.indexer is None:
+            # Dense layers have nothing to overlap; keep the base serial path.
+            self.aux_stream_list = None
 
     @classmethod
     def get_padded_num_q_heads(cls, num_heads: int) -> int:
@@ -600,8 +604,193 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
     def _run_parallel_input_projections(
         self, hidden_states: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        if self.indexer is not None and self.compressor is None:
+            # Serial: qr must sit on the default stream before the q-side
+            # fork below can consume it (ROCm).
+            aux_streams = self.aux_stream_list
+            self.aux_stream_list = None
+            try:
+                return super()._run_parallel_input_projections(hidden_states)
+            finally:
+                self.aux_stream_list = aux_streams
         return super()._run_parallel_input_projections(hidden_states)
 
+    def _enable_multi_stream_overlap(self) -> bool:
+        """ROCm multi-stream gates: streams and capture region.
+
+        Dict metadata marks piecewise cudagraph, whose eager breaks rebuild
+        the attention inputs on the owning stream. Forking side streams
+        there would rely on runtime HIP event sync, which is unreliable in
+        this overlap on ROCm (event waits can hang), so multi-stream only
+        runs where the fork/join becomes static graph edges: inside capture,
+        or with non-dict metadata (full cudagraph or the profile run), which
+        has no eager breaks.
+        """
+        attn_metadata = get_forward_context().attn_metadata
+        return self.aux_stream_list is not None and (
+            torch.cuda.is_current_stream_capturing()
+            or not isinstance(attn_metadata, dict)
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        llama_4_scaling: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if (
+            self.indexer is None
+            or self._prepare_and_attn_fn == self._prepare_and_attn_eager
+            or not self._enable_multi_stream_overlap()
+        ):
+            # Sequential fallback: no forks outside the capture region, where
+            # HIP event sync is unreliable; MRV1 keeps the input prep in its
+            # wide eager region (#51430).
+            aux_streams = self.aux_stream_list
+            self.aux_stream_list = None
+            try:
+                return super().forward(positions, hidden_states, llama_4_scaling)
+            finally:
+                self.aux_stream_list = aux_streams
+
+        attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states)
+        if self.compressor is not None:
+            q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_full(
+                hidden_states, positions
+            )
+        else:
+            q, kv, index_q, index_q_scale, index_weights = self._forward_csa2_reindex(
+                hidden_states, positions
+            )
+        self._sparse_indexer_and_attn(
+            hidden_states,
+            index_q,
+            index_q_scale,
+            index_weights,
+            q,
+            kv,
+            positions,
+            attn_out,
+        )
+        return self._o_proj(attn_out, positions)
+
+    def _forward_csa2_full(
+        self, hidden_states: torch.Tensor, positions: torch.Tensor
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        """CSA layers: one fork replaces the base pipeline's three.
+
+        Each fork/join is a HIP event pair per side stream on ROCm, so the
+        base three-fork pipeline is expensive there; the compressor chain and
+        the indexer weights projection run fully in parallel with the default
+        chain. The indexer K write and q-side stay serial after the join:
+        ``wk`` consumes the aux-0 latent and the q-side's wq_b consumes the
+        default-stream qr.
+        """
+        attn_metadata = get_forward_context().attn_metadata
+        compressor = self.compressor
+        indexer = self.indexer
+        assert compressor is not None and indexer is not None
+        aux_streams = self.aux_stream_list
+        assert aux_streams is not None and len(aux_streams) >= 2
+
+        def default_chain():
+            qr_kv = self._fused_wqa_wkv_gemm(hidden_states)
+            qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
+            q = self._wq_b_proj(qr, qr_scale).view(
+                -1, self.n_local_heads, self.head_dim
+            )
+            return (
+                self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata),
+                qr,
+                qr_scale,
+                kv,
+            )
+
+        def compressor_chain():
+            kv_score = torch.mm(
+                hidden_states,
+                compressor.fused_wkv_wgate.weight.T,
+                out_dtype=torch.float32,
+            )
+            latent = compressor(kv_score, positions)
+            compressor.insert_cache(latent, positions, self.rotary_emb)
+            return latent
+
+        def indexer_weights_chain():
+            # ReplicatedLinear returns (output, bias); bias is None.
+            weights, _ = indexer.weights_proj(hidden_states)
+            return weights
+
+        (q, qr, qr_scale, kv), (latent, indexer_weights) = execute_in_parallel(
+            default_chain,
+            [compressor_chain, indexer_weights_chain],
+            self.ln_events[0],
+            self.ln_events[1:3],
+            aux_streams[:2],
+            enable=True,
+            default_first=True,
+        )
+
+        indexer._produce_k(latent, positions, self.indexer_rotary_emb)
+        index_q, index_q_scale, index_weights_out = indexer.forward_q(
+            qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb
+        )
+        return q, kv, index_q, index_q_scale, index_weights_out
+
+    def _forward_csa2_reindex(
+        self, hidden_states: torch.Tensor, positions: torch.Tensor
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        """Indexer-only layers: the q-side overlaps the SWA q path.
+
+        The input projections run serial first: the q-side consumes qr, so
+        it can only overlap the SWA q projection and KV insert.
+        """
+        indexer = self.indexer
+        assert indexer is not None
+        attn_metadata = get_forward_context().attn_metadata
+        # Serial on these layers via the _run_parallel_input_projections
+        # override: qr must sit on the default stream before the fork.
+        qr_kv, _, indexer_weights = self._run_parallel_input_projections(hidden_states)
+        assert indexer_weights is not None
+        qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
+
+        def swa_q_chain():
+            q = self._wq_b_proj(qr, qr_scale).view(
+                -1, self.n_local_heads, self.head_dim
+            )
+            return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+
+        def indexer_q_chain():
+            return indexer.forward_q(
+                qr, qr_scale, indexer_weights, positions, self.indexer_rotary_emb
+            )
+
+        aux_streams = self.aux_stream_list
+        assert aux_streams is not None
+        q, aux_results = execute_in_parallel(
+            swa_q_chain,
+            [indexer_q_chain],
+            self.ln_events[0],
+            self.ln_events[1:2],
+            aux_streams[:1],
+            enable=True,
+            default_first=True,
+        )
+        index_q, index_q_scale, index_weights_out = aux_results[0]
+        return q, kv, index_q, index_q_scale, index_weights_out
+
     @functools.cached_property
     def _wq_b_uses_aiter_block_scaled(self) -> bool:
         """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.

vllm/models/deepseek_v41/attention.py

@@ -478,7 +478,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             # graph and MRV1 produces garbage (#51430).
             self._prepare_and_attn_fn = self._prepare_and_attn_eager
 
-        # Will be None on ROCm for now.
         self.aux_stream_list = aux_stream_list
         # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
         # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
@@ -1381,6 +1380,37 @@ class DeepseekV4Indexer(nn.Module):
             self.use_fp4_kv,
         )
 
+    def forward_q(
+        self,
+        qr: torch.Tensor | QuantizedActivation,
+        qr_scale: torch.Tensor | None,
+        indexer_weights: torch.Tensor,
+        positions: torch.Tensor,
+        rotary_emb: nn.Module,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
+        """Build the indexer queries: wq_b over qr plus fused RoPE/quant.
+
+        Split out so the ROCm layer can schedule it apart from
+        ``_produce_k`` (e.g. on an aux stream once ``qr`` is ready).
+        """
+        q = self._wq_b_proj(qr, qr_scale)
+        q = q.view(-1, self.n_head, self.head_dim)
+        q_quant, weights = fused_indexer_q_rope_quant(
+            positions,
+            q,
+            rotary_emb.cos_sin_cache,
+            indexer_weights,
+            self.softmax_scale,
+            self.n_head**-0.5,
+            use_fp4=self.use_fp4_kv,
+            weights_out_dtype=self.indexer_weights_dtype,
+        )
+        if isinstance(q_quant, tuple):
+            q, q_scale = q_quant
+        else:
+            q, q_scale = q_quant, None
+        return q, q_scale, weights
+
     def forward(
         self,
         qr: torch.Tensor | QuantizedActivation,
@@ -1422,23 +1452,7 @@ class DeepseekV4Indexer(nn.Module):
             # (skip_k_cache_insert=True).
             self._produce_k(latent, positions, rotary_emb)
 
-        q = self._wq_b_proj(qr, qr_scale)
-        q = q.view(-1, self.n_head, self.head_dim)
-        q_quant, weights = fused_indexer_q_rope_quant(
-            positions,
-            q,
-            rotary_emb.cos_sin_cache,
-            indexer_weights,
-            self.softmax_scale,
-            self.n_head**-0.5,
-            use_fp4=self.use_fp4_kv,
-            weights_out_dtype=self.indexer_weights_dtype,
-        )
-        if isinstance(q_quant, tuple):
-            q, q_scale = q_quant
-        else:
-            q, q_scale = q_quant, None
-        return q, q_scale, weights
+        return self.forward_q(qr, qr_scale, indexer_weights, positions, rotary_emb)
 
     def _wq_b_proj(
         self,

vllm/utils/multi_stream_utils.py

@@ -70,6 +70,7 @@ def execute_in_parallel(
     done_events: list[torch.cuda.Event],
     aux_streams: list[torch.cuda.Stream] | None = None,
     enable: bool = False,
+    default_first: bool = False,
 ) -> tuple[Any, list[Any]]:
     """Run default_fn on the current stream and aux_fns concurrently on
     aux_streams.
@@ -97,6 +98,10 @@ def execute_in_parallel(
             so callers that pass aux_streams must also pass enable=True
             (typically gated by an env var) to actually overlap. When False,
             execution falls back to sequential on the current stream.
+        default_first: When True, enqueue the default chain before the aux
+            chains. The default chain is typically the longest, so it starts
+            first; the aux streams still only wait on start_event, so they
+            overlap the default chain. Defaults to False.
 
     Returns:
         Tuple of (default_result, aux_results) where aux_results[i] is the
@@ -117,6 +122,8 @@ def execute_in_parallel(
     pending: list[torch.cuda.Event] = []
 
     start_event.record()
+    if default_first:
+        default_result = default_fn()
     for i, fn in enumerate(aux_fns):
         if fn is None:
             continue
@@ -125,8 +132,8 @@ def execute_in_parallel(
             aux_results[i] = fn()
             done_events[i].record()
         pending.append(done_events[i])
-
-    default_result = default_fn()
+    if not default_first:
+        default_result = default_fn()
 
     for ev in pending:
         ev.wait()