From 2dcee59850a0cdbd5a29a857ac2264af5a7b2229 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Sat, 15 Aug 2026 21:38:29 -0700 Subject: [PATCH] frost(sdpa): order every execute-path tensor prep on the launch stream (Rule 5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The execute paths resolved the launch stream but ran their tensor prep on torch's CURRENT stream: _to_bshd's gather copy (non-compact layouts), the cached dummies' first-use zero-fill, _reshape_sf's .contiguous(), and some O-scratch copy-backs / amax post-ops. With an explicit caller stream (the execute-time handle's), that work races the kernel launch — same class of bug as the PR #543 THD-upload race, and flagged by review on PR #608. Fix, uniformly across the five sites (SM100 dense f16 / mxfp8 / fp8, SM120 dense f16 / fp8): resolve current_stream FIRST, run the prep inside _torch_stream_context(current_stream, device), and put the consumers (copy-backs, amax div) in the same context — matching what the THD paths and the amax resets already did. Rebased over #619's device scale-fold: the fp8 paths' _scale_view calls sit inside the wrap too — None binds a cached 1.0 dummy whose first-use torch.ones fill is itself a launch — and the post-kernel amax_o.div_(scale_o view) is a device op inside the consumer wrap. The PyTorch-integration path launches on torch's current stream, where the context is a no-op; only direct graph-API users with an explicit stream were exposed. Validated: SM100 (B200, 9.26) fwd dsl + fp8 + mxfp8 + stream-respect + stream-ordering + async/capture suites L0+L1: 546 passed, 0 failed. SM120 (RTX 5080, 9.24) fwd dsl + fp8 + the same stream suites L0+L1: 173 passed, 17 skipped, 0 failed. --- python/cudnn/sdpa/fwd/api_dsl.py | 295 +++++++++++++++++-------------- 1 file changed, 165 insertions(+), 130 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index cc3b844c0..c114018ac 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1149,30 +1149,33 @@ def execute( ) return - Q = self._to_bshd(q_tensor) - K = self._to_bshd(k_tensor) - V = self._to_bshd(v_tensor) - O_view, o_needs_copy_back, O_scratch = self._to_bshd_writable(o_tensor) - device = q_tensor.device - sinks_t = ( - self._checked_sinks_1d(sinks) - if sinks is not None - else self._dummy("sinks", device, lambda: torch.zeros(self.h_q, dtype=torch.float32, device=device)) - ) - seq_kv_t = ( - self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") - if seq_kv_lens is not None - else self._dummy("seq_kv", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) - ) - # Dense padded-Q trim: per-batch Q lengths are their OWN kernel - # parameter (compiled in only when seq_q_lens_present — the kernel - # signature is specialized on `None`, so the flag-off ABI is - # unchanged). The caller's (B,)-int32 device tensor is bound directly - # as a validated view — zero allocations/copies on the execute hot - # path, stable pointer (CUDA-graph-capture friendly). - seq_q_t = self._checked_seq_lens(seq_q_lens, "seq_q_lens") if self.seq_q_lens_present else None - o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device)) + # Tensor prep can LAUNCH work (_to_bshd's gather copy, the dummies' + # first-use zero-fill), so it runs on the LAUNCH stream (Rule 5) — + # otherwise an explicit caller stream races torch's current stream. + with _torch_stream_context(current_stream, device): + Q = self._to_bshd(q_tensor) + K = self._to_bshd(k_tensor) + V = self._to_bshd(v_tensor) + O_view, o_needs_copy_back, O_scratch = self._to_bshd_writable(o_tensor) + sinks_t = ( + self._checked_sinks_1d(sinks) + if sinks is not None + else self._dummy("sinks", device, lambda: torch.zeros(self.h_q, dtype=torch.float32, device=device)) + ) + seq_kv_t = ( + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + if seq_kv_lens is not None + else self._dummy("seq_kv", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) + ) + # Dense padded-Q trim: per-batch Q lengths are their OWN kernel + # parameter (compiled in only when seq_q_lens_present — the kernel + # signature is specialized on `None`, so the flag-off ABI is + # unchanged). The caller's (B,)-int32 device tensor is bound directly + # as a validated view — zero allocations/copies on the execute hot + # path, stable pointer (CUDA-graph-capture friendly). + seq_q_t = self._checked_seq_lens(seq_q_lens, "seq_q_lens") if self.seq_q_lens_present else None + o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device)) import cutlass @@ -1192,7 +1195,9 @@ def execute( stream=current_stream, ) if o_needs_copy_back: - O_view.copy_(O_scratch) + # Consumes what the kernel just wrote — launch stream (Rule 5). + with _torch_stream_context(current_stream, device): + O_view.copy_(O_scratch) self._logger.debug("execute completed") def _thd_compile_kwargs(self) -> dict: @@ -1492,28 +1497,34 @@ def _execute_mxfp8( sq, skv = self.s_q_max, self.s_k_max device = q_tensor.device - Q = self._to_bshd(q_tensor) - K = self._to_bshd(k_tensor) - V = self._to_bshd(v_tensor) - O_view, o_needs_copy_back, O_scratch = self._to_bshd_writable(o_tensor) - O = O_scratch if o_needs_copy_back else O_view - - n_q_tiles = self._ceil_div(sq, _SM100_TILE_N) - n_kv_tiles = self._ceil_div(skv, _SM100_TILE_N) - sf_q_v = self._reshape_sf(sf_q, h_q, n_q_tiles, km.SF_SMEM_SIZE_Q) - sf_k_v = self._reshape_sf(sf_k, h_kv, n_kv_tiles, km.SF_SMEM_SIZE_K) - sf_v_v = self._reshape_sf(sf_v, h_kv, n_kv_tiles, km.SF_SMEM_SIZE_V) - - # has_lse=False (no Stats output): the store is compiled out; bind None. - lse = lse_tensor.reshape(b, h_q, sq) if lse_tensor is not None else None - sinks_t = ( - self._checked_sinks_1d(sinks) if sinks is not None else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device)) - ) - seq_kv_t = ( - self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") - if seq_kv_lens is not None - else self._dummy("seq_kv", device, lambda: torch.zeros(b, dtype=torch.int32, device=device)) - ) + # Tensor prep can LAUNCH work (_to_bshd's gather copy, _reshape_sf's + # .contiguous(), the dummies' first-use zero-fill) — launch stream + # (Rule 5), like the amax reset below. + with _torch_stream_context(current_stream, device): + Q = self._to_bshd(q_tensor) + K = self._to_bshd(k_tensor) + V = self._to_bshd(v_tensor) + O_view, o_needs_copy_back, O_scratch = self._to_bshd_writable(o_tensor) + O = O_scratch if o_needs_copy_back else O_view + + n_q_tiles = self._ceil_div(sq, _SM100_TILE_N) + n_kv_tiles = self._ceil_div(skv, _SM100_TILE_N) + sf_q_v = self._reshape_sf(sf_q, h_q, n_q_tiles, km.SF_SMEM_SIZE_Q) + sf_k_v = self._reshape_sf(sf_k, h_kv, n_kv_tiles, km.SF_SMEM_SIZE_K) + sf_v_v = self._reshape_sf(sf_v, h_kv, n_kv_tiles, km.SF_SMEM_SIZE_V) + + # has_lse=False (no Stats output): the store is compiled out; bind None. + lse = lse_tensor.reshape(b, h_q, sq) if lse_tensor is not None else None + sinks_t = ( + self._checked_sinks_1d(sinks) + if sinks is not None + else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device)) + ) + seq_kv_t = ( + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + if seq_kv_lens is not None + else self._dummy("seq_kv", device, lambda: torch.zeros(b, dtype=torch.int32, device=device)) + ) amax_o_buf = amax_o.reshape(-1)[:1] if amax_o is not None else self._dummy("amax_o", device, lambda: torch.zeros(1, dtype=torch.float32, device=device)) # Must be enqueued on the SAME stream as the kernel launch below, else the @@ -1539,7 +1550,9 @@ def _execute_mxfp8( stream=current_stream, ) if o_needs_copy_back: - O_view.copy_(O) + # Consumes what the kernel just wrote — launch stream (Rule 5). + with _torch_stream_context(current_stream, device): + O_view.copy_(O) self._logger.debug("execute (MXFP8) completed") def _execute_fp8( @@ -1576,35 +1589,41 @@ def _execute_fp8( sq, skv = self.s_q_max, self.s_k_max device = q_tensor.device - # Rule 3: the scales stay on device — the kernel loads and folds - # descale_q*descale_k into the softmax scale and descale_v*scale_o - # into o_scale_fused; the scalar args carry only the bases. - # (descale_s/scale_s never reach this layer: the lowering does not - # forward them, and P is cast with the baked P_CAST_LOG2_SCALE bias.) - dq_t = self._scale_view(descale_q, "descale_q", device) - dk_t = self._scale_view(descale_k, "descale_k", device) - dv_t = self._scale_view(descale_v, "descale_v", device) - so_t = self._scale_view(scale_o, "scale_o", device) + # Tensor prep can LAUNCH work (_to_bshd's gather copy, the dummies' + # first-use zero-fill, _scale_view's cached-1.0 first-use fill) — + # launch stream (Rule 5), like the amax reset. + with _torch_stream_context(current_stream, device): + # Rule 3: the scales stay on device — the kernel loads and folds + # descale_q*descale_k into the softmax scale and descale_v*scale_o + # into o_scale_fused; the scalar args carry only the bases. + # (descale_s/scale_s never reach this layer: the lowering does not + # forward them, and P is cast with the baked P_CAST_LOG2_SCALE bias.) + dq_t = self._scale_view(descale_q, "descale_q", device) + dk_t = self._scale_view(descale_k, "descale_k", device) + dv_t = self._scale_view(descale_v, "descale_v", device) + so_t = self._scale_view(scale_o, "scale_o", device) + + Q = self._to_bshd(q_tensor) + K = self._to_bshd(k_tensor) + V = self._to_bshd(v_tensor) + O_view, o_needs_copy_back, O_scratch = self._to_bshd_writable(o_tensor) + O = O_scratch if o_needs_copy_back else O_view + + # has_lse=False (no Stats output): the store is compiled out; bind None. + lse = lse_tensor.reshape(b, h_q, sq) if lse_tensor is not None else None + sinks_t = ( + self._checked_sinks_1d(sinks) + if sinks is not None + else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device)) + ) + seq_kv_t = ( + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + if seq_kv_lens is not None + else self._dummy("seq_kv", device, lambda: torch.zeros(b, dtype=torch.int32, device=device)) + ) scale_softmax_log2 = scale_val * math.log2(math.e) o_scale_fused = 1.0 - Q = self._to_bshd(q_tensor) - K = self._to_bshd(k_tensor) - V = self._to_bshd(v_tensor) - O_view, o_needs_copy_back, O_scratch = self._to_bshd_writable(o_tensor) - O = O_scratch if o_needs_copy_back else O_view - - # has_lse=False (no Stats output): the store is compiled out; bind None. - lse = lse_tensor.reshape(b, h_q, sq) if lse_tensor is not None else None - sinks_t = ( - self._checked_sinks_1d(sinks) if sinks is not None else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device)) - ) - seq_kv_t = ( - self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") - if seq_kv_lens is not None - else self._dummy("seq_kv", device, lambda: torch.zeros(b, dtype=torch.int32, device=device)) - ) - # amax_o: the kernel atomicMax'es into this buffer, so it MUST start # at 0. It accumulates max|o_scaled| (pre-cast, exact even for FP8 O); # dividing by scale_o below yields the pre-quant output amax. @@ -1632,13 +1651,15 @@ def _execute_fp8( amax_o_buf, stream=current_stream, ) - if o_needs_copy_back: - O_view.copy_(O) - if amax_o is not None: - # Device divisor: the same div_ as before, minus the readback. - # scale_o > 0 is caller contract (backend parity); None bound a - # cached 1.0 above. - amax_o_buf.div_(so_t) + # Both consume what the kernel just wrote — launch stream (Rule 5). + with _torch_stream_context(current_stream, device): + if o_needs_copy_back: + O_view.copy_(O) + if amax_o is not None: + # Device divisor: the same div_ as before, minus the readback. + # scale_o > 0 is caller contract (backend parity); None bound a + # cached 1.0 above. + amax_o_buf.div_(so_t) self._logger.debug("execute (FP8 per-tensor) completed") @@ -2221,39 +2242,43 @@ def execute( current_stream=current_stream, ) return - lse = self._checked_lse_view(lse_tensor) if lse_tensor is not None else None - sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None - seq_q_lens = ( - self._checked_seq_lens(seq_q_lens, "seq_q_lens") - if seq_q_lens is not None - else self._dummy( - "seq_q_lens", - q_tensor.device, - lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=q_tensor.device), - ) - ) - seq_kv_lens = ( - self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") - if seq_kv_lens is not None - else self._dummy( - "seq_kv_lens", - q_tensor.device, - lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=q_tensor.device), - ) - ) if current_stream is None: # Direct call (no dispatch-forwarded stream): fall back to torch's # current stream, as before. A stream forwarded from the execute-time - # handle is respected rather than clobbered. + # handle is respected rather than clobbered. Resolved FIRST (Rule 5): + # the tensor prep below can launch work (the dummies' zero-fill, + # _to_bshd's gather copy), which must be ordered against the kernel. current_stream = cuda.CUstream(torch.cuda.current_stream(q_tensor.device).cuda_stream) + lse = self._checked_lse_view(lse_tensor) if lse_tensor is not None else None + sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None + with _torch_stream_context(current_stream, q_tensor.device): + seq_q_lens = ( + self._checked_seq_lens(seq_q_lens, "seq_q_lens") + if seq_q_lens is not None + else self._dummy( + "seq_q_lens", + q_tensor.device, + lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=q_tensor.device), + ) + ) + seq_kv_lens = ( + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + if seq_kv_lens is not None + else self._dummy( + "seq_kv_lens", + q_tensor.device, + lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=q_tensor.device), + ) + ) + q = self._to_bshd(q_tensor) + k = self._to_bshd(k_tensor) + v = self._to_bshd(v_tensor) + o_view, o_needs_copy_back, o_scratch = self._to_bshd_writable(o_tensor) + o = o_scratch if o_needs_copy_back else o_view + import cutlass - q = self._to_bshd(q_tensor) - k = self._to_bshd(k_tensor) - v = self._to_bshd(v_tensor) - o_view, o_needs_copy_back, o_scratch = self._to_bshd_writable(o_tensor) - o = o_scratch if o_needs_copy_back else o_view self._compiled_kernel( q, k, @@ -2271,7 +2296,9 @@ def execute( current_stream, ) if o_needs_copy_back: - o_view.copy_(o_scratch) + # Consumes what the kernel just wrote — launch stream (Rule 5). + with _torch_stream_context(current_stream, q_tensor.device): + o_view.copy_(o_scratch) def _execute_fp8( self, @@ -2305,15 +2332,22 @@ def _execute_fp8( import cutlass device = q_tensor.device + if current_stream is None: + # Resolved FIRST (Rule 5): the scale views' cached-1.0 fill, the + # dummies' zero-fill and _to_bshd's gather copy below launch work + # that must be ordered vs the kernel. + current_stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + # Rule 3: the scales stay on device — the kernel loads and folds # dq*dk into the softmax scale and dv*so into o_scale_fused; the # scalar args carry only the bases. None binds a cached 1.0. # (descale_s/scale_s never reach this layer: the lowering does not # forward them, and P is cast with the baked P_CAST_LOG2_SCALE bias.) - dq_t = self._scale_view(descale_q, "descale_q", device) - dk_t = self._scale_view(descale_k, "descale_k", device) - dv_t = self._scale_view(descale_v, "descale_v", device) - so_t = self._scale_view(scale_o, "scale_o", device) + with _torch_stream_context(current_stream, device): + dq_t = self._scale_view(descale_q, "descale_q", device) + dk_t = self._scale_view(descale_k, "descale_k", device) + dv_t = self._scale_view(descale_v, "descale_v", device) + so_t = self._scale_view(scale_o, "scale_o", device) scale_softmax_log2 = scale_val * math.log2(math.e) o_scale_fused = 1.0 @@ -2325,9 +2359,6 @@ def _execute_fp8( self.lse_desc is None and lse_tensor is not None, "this specialization was compiled without an LSE output; construct the API with sample_lse", ) - if current_stream is None: - current_stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) - sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None # THD packs the batch away, so the ragged views and the per-execute # compile replace the dense buffers; everything else (the folded @@ -2372,22 +2403,26 @@ def _execute_fp8( else: lse = lse_tensor.as_strided((pack.t_q, self.h_q), (self.h_q, 1), lse_tensor.storage_offset()) else: - seq_kv_t = ( - self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") - if seq_kv_lens is not None - else self._dummy("seq_kv_lens", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) - ) - # Dense padded-Q trim: bind the REAL per-batch Q lengths whenever - # the graph carries them — a zeroed dummy would trim every row. - seq_q_t = ( - self._checked_seq_lens(seq_q_lens, "seq_q_lens") - if seq_q_lens is not None - else self._dummy("seq_q_lens", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) - ) - q = self._to_bshd(q_tensor) - k = self._to_bshd(k_tensor) - v = self._to_bshd(v_tensor) - o_view, o_needs_copy_back, o_scratch = self._to_bshd_writable(o_tensor) + # Tensor prep can LAUNCH work (_to_bshd's gather copy, the + # dummies' first-use zero-fill) — launch stream (Rule 5), like + # the amax reset below. + with _torch_stream_context(current_stream, device): + seq_kv_t = ( + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + if seq_kv_lens is not None + else self._dummy("seq_kv_lens", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) + ) + # Dense padded-Q trim: bind the REAL per-batch Q lengths whenever + # the graph carries them — a zeroed dummy would trim every row. + seq_q_t = ( + self._checked_seq_lens(seq_q_lens, "seq_q_lens") + if seq_q_lens is not None + else self._dummy("seq_q_lens", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) + ) + q = self._to_bshd(q_tensor) + k = self._to_bshd(k_tensor) + v = self._to_bshd(v_tensor) + o_view, o_needs_copy_back, o_scratch = self._to_bshd_writable(o_tensor) o = o_scratch if o_needs_copy_back else o_view lse = self._checked_lse_view(lse_tensor) if lse_tensor is not None else None