From 46e78db1dc3a16f84e3cded13ea89de319daf4fd Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Thu, 20 Aug 2026 10:50:31 -0700 Subject: [PATCH 1/4] frost(sdpa): SM80 backward workspace carving + strided stats (issue #514) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the merged forward port's contract: both bprop kernels gain scratch_bytes() and a workspace= param carving dQ_acc / dQ / dK_ws / dV_ws / GQA-reduced dK/dV / the deterministic-dQ semaphore / dBias+dSink accumulators / do_dot from the caller's buffer (cached 1-element dummies for absent operands); SdpabwdSm80 gains scratch_workspace_bytes(feature flags) covering pad/gather staging plus the kernel tail and a carve-aware execute(workspace=); lower_sm80_bwd builds the adapter at plan time from normalized descriptors and records the total as workspace_bytes. The row also declares strided_stats (#666's capability): the kernels read a packed LSE, so a stats input with any other declared strides is gathered into a carved contiguous chunk — without this, every stats-stride-randomized mhas draw (#304, active on cuDNN >= 9.26) declined to the backend. Adds the issue's no-alloc regression test: the CUDA allocator counter stays flat across re-executes of both SM80 engines, outputs bitwise-stable. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/bwd/api.py | 132 +++++++++++++-- python/cudnn/sdpa/bwd/engines.py | 152 ++++++++++++++---- .../sdpa/bwd/kernels/bprop_d64_f16_sm80.py | 47 +++++- .../cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py | 127 +++++++++++++-- .../test_sdpa_sm80_frontend_integration.py | 55 +++++++ 5 files changed, 445 insertions(+), 68 deletions(-) diff --git a/python/cudnn/sdpa/bwd/api.py b/python/cudnn/sdpa/bwd/api.py index 571569e99..207b6e425 100644 --- a/python/cudnn/sdpa/bwd/api.py +++ b/python/cudnn/sdpa/bwd/api.py @@ -287,6 +287,72 @@ def check_support(self) -> bool: self._logger.debug("check_support (bwd) completed") return True + # ------------------------------------------------------------------ + def _needs_bshd_stage(self, desc) -> bool: + """Whether ``desc``'s BSHD transpose is non-contiguous — i.e. execute's + kernel-facing view would need a gather into staging.""" + b, h, sq, d = desc.shape + expect = (sq * h * d, d, h * d, 1) # BHSD-logical view of a compact BSHD buffer + return tuple(desc.stride) != expect + + def scratch_workspace_bytes( + self, + *, + has_bias: Optional[bool] = None, + bias_batch: int = 1, + has_sink: bool = False, + deterministic: bool = False, + need_do_dot: bool = True, + ) -> int: + """Per-execute scratch requirement (issue #514): head-dim pad / BSHD + gather staging for Q/K/V/O/dO plus the kernel's internal scratch + (``bprop_f16_sm80.scratch_bytes``; the generic kernel's buffer set + covers the d64 fast path's). The feature flags must match what + execute() will be called with — the engine lowering passes its graph + facts; the default reads the constructor's ``has_bias``.""" + self._ensure_support_checked() + from ..fwd.api_dsl import ws_align + from .kernels import bprop_f16_sm80 as _kmod + + elem = 2 # fp16/bf16 — check_support admits no other input dtype + b, hq, sq, _ = self.q_desc.shape + _, hkv, skv, _ = self.k_desc.shape + fdqk, fdv = self.flavor_d_qk, self.flavor_d_v + pad_qk = self.head_dim_qk < fdqk + pad_v = self.head_dim_v < fdv + if has_bias is None: + has_bias = self.has_bias + total = 0 + # Pad / gather staging, in execute()'s carve order (Q, K, V, O, dO). + for desc, s_len, hh, pad, fd in ( + (self.q_desc, sq, hq, pad_qk, fdqk), + (self.k_desc, skv, hkv, pad_qk, fdqk), + (self.v_desc, skv, hkv, pad_v, fdv), + (self.o_desc, sq, hq, pad_v, fdv), + (self.do_desc, sq, hq, pad_v, fdv), + ): + if pad: + total += ws_align(b * s_len * hh * fd * elem) + elif self._needs_bshd_stage(desc): + total += ws_align(math.prod(desc.shape) * elem) + # Kernel-internal scratch at the PADDED (flavor) head dims. + total += _kmod.scratch_bytes( + B=b, + SQ=sq, + SKV=skv, + H=hq, + Hk=hkv, + d_qk=fdqk, + d_v=fdv, + io_bytes=elem, + deterministic=deterministic, + has_bias=bool(has_bias), + bias_batch=bias_batch, + has_sink=has_sink, + need_do_dot=need_do_dot, + ) + return total + # ------------------------------------------------------------------ def compile(self) -> None: """No-op — the kernel module owns its own per-shape ``lru_cache``; @@ -318,6 +384,7 @@ def execute( sinks: Optional[torch.Tensor] = None, rope_freqs: Optional[torch.Tensor] = None, deterministic: bool = False, + workspace: Optional[torch.Tensor] = None, ) -> None: self._logger.debug("Entering execute (bwd)") if self._compiled_kernel is None: @@ -326,19 +393,57 @@ def execute( kernel = _load_kernel_module() - # BHSD → BSHD for the kernel. - Q, K, V = _bshd(q_tensor), _bshd(k_tensor), _bshd(v_tensor) - O, dO = _bshd(o_tensor), _bshd(do_tensor) + # Per-execute scratch: carved from the caller's workspace when one is + # provided (the engine executor always passes one sized by + # scratch_workspace_bytes(); issue #514), otherwise allocated (the + # standalone wrapper paths). + carver = None + if workspace is not None: + from ..fwd.api_dsl import WorkspaceCarver + + carver = WorkspaceCarver( + workspace, + self.scratch_workspace_bytes( + has_bias=bias_tensor is not None, + bias_batch=(bias_tensor.shape[0] if bias_tensor is not None else 1), + has_sink=sinks is not None, + deterministic=bool(deterministic), + ), + "SdpabwdSm80", + ) pad_v = self.head_dim_v < self.flavor_d_v pad_qk = self.head_dim_qk < self.flavor_d_qk - if pad_qk: - Q = _pad_last_dim(Q, self.flavor_d_qk) - K = _pad_last_dim(K, self.flavor_d_qk) - if pad_v: - V = _pad_last_dim(V, self.flavor_d_v) - O = _pad_last_dim(O, self.flavor_d_v) - dO = _pad_last_dim(dO, self.flavor_d_v) + + def _stage(t: torch.Tensor, pad: bool, fd: int) -> torch.Tensor: + """Kernel-facing BSHD view of BHSD-logical ``t``: zero-copy when the + transpose is contiguous, otherwise gathered (and head-dim padded) + into carved staging — or allocated when no workspace was given.""" + view = t.transpose(1, 2) + d = view.shape[-1] + if pad: + if carver is not None: + bb, ss, hh, _ = view.shape + dst = carver.take(bb * ss * hh * fd, t.dtype).view(bb, ss, hh, fd) + dst[..., :d].copy_(view) + dst[..., d:].zero_() + return dst + return _pad_last_dim(view.contiguous() if not view.is_contiguous() else view, fd) + if view.is_contiguous(): + return view + if carver is not None: + dst = carver.take(t.numel(), t.dtype).view(view.shape) + dst.copy_(view) + return dst + return view.contiguous() + + # BHSD → BSHD for the kernel, in scratch_workspace_bytes()'s sizing + # order (Q, K, V, O, dO). + Q = _stage(q_tensor, pad_qk, self.flavor_d_qk) + K = _stage(k_tensor, pad_qk, self.flavor_d_qk) + V = _stage(v_tensor, pad_v, self.flavor_d_v) + O = _stage(o_tensor, pad_v, self.flavor_d_v) + dO = _stage(do_tensor, pad_v, self.flavor_d_v) # Build the feature-kwarg superset; drop any the kernel doesn't accept. bw_kwargs = dict( @@ -353,6 +458,8 @@ def execute( sinks=sinks, rope_freqs=rope_freqs, deterministic=bool(deterministic), + # Kernel-internal scratch: the unconsumed workspace tail (issue #514). + workspace=carver.remaining() if carver is not None else None, ) # Route plain dense MHA d=64 calls to the dedicated perf kernel # (~2x faster on A100). The gate must stay exhaustive: the d64 @@ -402,9 +509,10 @@ def execute( dv_tensor.copy_(dV_k.transpose(1, 2)) if dbias_tensor is not None and dBias_k is not None: # dBias is head-major [., H, SQ, SKV] (like bias) — no transpose. - dbias_tensor.copy_(dBias_k.to(dbias_tensor.dtype)) + # copy_ casts in place; a .to() would allocate a staging tensor. + dbias_tensor.copy_(dBias_k) if dsink_tensor is not None and dSink_k is not None: - dsink_tensor.copy_(dSink_k.to(dsink_tensor.dtype)) + dsink_tensor.copy_(dSink_k) self._logger.debug("execute (bwd) completed") diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index 70c510bb8..d49282b20 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -499,14 +499,30 @@ def _sm80_spec() -> EngineSpec: sink=True, decode=False, # prefill kernels only layouts=frozenset({"bshd", "dense_flex"}), + # Served by gathering the strided stats into carved contiguous + # staging (issue #514 workspace machinery) — the kernels read a + # packed LSE; sm120 reads declared strides natively instead. + strided_stats=True, ), lower=lower_sm80_bwd, ) def lower_sm80_bwd(spec: EngineSpec, facts: "ga.SdpaGraphFacts", requested: Any = None): - """Lower the SM80 backward row through the ``cudnn.sdpa`` SM80 adapter.""" - from .api import sdpa_bwd_wrapper_sm80 + """Lower the SM80 backward row through the ``cudnn.sdpa`` SM80 adapter. + + Built at plan time (issue #514): the adapter is constructed here from the + NORMALIZED buffer descriptors (compact BSHD-physical), its scratch + requirement plus this executor's own dense_flex gather staging is recorded + as ``workspace_bytes``, and execute carves everything from the caller's + workspace — no per-execute allocation on this path. + """ + import dataclasses + + from cudnn.api_base import TensorDesc + from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver, ws_align + + from .api import SdpabwdSm80 binding = ga.SdpaBinding( q=facts.q_t, @@ -526,49 +542,117 @@ def lower_sm80_bwd(spec: EngineSpec, facts: "ga.SdpaGraphFacts", requested: Any dsink=facts.dsink_t, ) mask_args = ga.adapter_mask_args(facts) + elem = 2 # fp16/bf16 — mismatch() admits no other input dtype + b, h_q = facts.b, facts.h_q + + def _compact_desc(t, name): + desc = ga.tensor_desc_from_ir(t, name=name) + bb, hh, ss, dd = desc.shape + return dataclasses.replace(desc, stride=(ss * hh * dd, dd, hh * dd, 1), stride_order=(3, 1, 2, 0)) + + def _is_compact_bshd(t) -> bool: + _, h, s, d = tuple(t.get_dim()) + return tuple(t.get_stride()) == (s * h * d, d, h * d, 1) + + # dense_flex gather staging, sized from the PORT layouts (static): a port + # already stored as a compact BSHD-physical allocation is handed through + # zero-copy. The adapter's own scratch then covers head-dim pads and the + # kernel-internal buffers. + ports = ((facts.q_t, "q"), (facts.k_t, "k"), (facts.v_t, "v"), (facts.o_t, "o"), (facts.do_t, "dO")) + + def _port_numel(t) -> int: + n = 1 + for extent in t.get_dim(): + n *= int(extent) + return n + + stage_bytes = {name: (0 if _is_compact_bshd(t) else ws_align(_port_numel(t) * elem)) for t, name in ports} + # Strided stats (Capabilities.strided_stats): the kernels read a PACKED + # (B, H_q, S_q) fp32 LSE, so a stats input with any other declared strides + # is gathered into a carved contiguous chunk at execute. + _stats_contig = facts.stats_t is not None and tuple(facts.stats_t.get_stride()) == (h_q * facts.s_q, facts.s_q, 1, 1) + stats_stage = 0 if (facts.stats_t is None or _stats_contig) else ws_align(b * h_q * facts.s_q * 4) + + q_desc = _compact_desc(facts.q_t, "q") + sample_lse = TensorDesc( + dtype=ga.to_torch_dtype(cudnn.data_type.FLOAT), + shape=(b, h_q, facts.s_q), + stride=(h_q * facts.s_q, facts.s_q, 1), + stride_order=(2, 1, 0), + device=q_desc.device, + name="lse", + ) + api = SdpabwdSm80( + sample_q=q_desc, + sample_k=_compact_desc(facts.k_t, "k"), + sample_v=_compact_desc(facts.v_t, "v"), + sample_o=_compact_desc(facts.o_t, "o"), + sample_do=_compact_desc(facts.do_t, "dO"), + sample_lse=sample_lse, + scale_softmax=facts.scale, + has_seq_kv_lens=facts.seq_kv_t is not None, + has_bias=facts.has_bias, + **mask_args, + ) + if not api.check_support(): + raise ValueError("SdpabwdSm80 declined the normalized graph geometry") + api.compile() + bias_batch = int(facts.bias_t.get_dim()[0]) if facts.bias_t is not None else 1 + api_scratch = api.scratch_workspace_bytes( + has_bias=facts.has_bias, + bias_batch=bias_batch, + has_sink=facts.has_sink, + deterministic=facts.deterministic, + ) + total_workspace_bytes = sum(stage_bytes.values()) + stats_stage + api_scratch + + def _normalize(carver, buf, staged: int): + if not staged: + return buf + bb, hh, ss, dd = buf.shape + dst = carver.take(bb * ss * hh * dd, buf.dtype).view(bb, ss, hh, dd) + dst.copy_(buf.permute(0, 2, 1, 3)) + return dst.permute(0, 2, 1, 3) - def _execute(variant_pack, stream=None): + def _execute(variant_pack, workspace=None, stream=None): resolved = ga.resolve_variant_pack(variant_pack, binding) - # mismatch() admits only the contiguous (B, H_q, S_q, 1) stats layout, - # so this is a pure view (a copying reshape would violate the - # execute() contract and hide the -inf padded-row trim semantics). - lse = resolved[id(facts.stats_t)].view(facts.b, facts.h_q, facts.s_q) - - out = sdpa_bwd_wrapper_sm80( - # dense_flex delivery: normalize to the BSHD-physical order the - # adapter requires (zero-copy when already BSHD). - ga.to_bshd_physical(resolved[id(facts.q_t)]), - ga.to_bshd_physical(resolved[id(facts.k_t)]), - ga.to_bshd_physical(resolved[id(facts.v_t)]), - ga.to_bshd_physical(resolved[id(facts.o_t)]), - ga.to_bshd_physical(resolved[id(facts.do_t)]), - lse, + carver = WorkspaceCarver(workspace, total_workspace_bytes, spec.name) if total_workspace_bytes else None + # squeeze(-1) is a valid view for ANY (B, H_q, S_q, 1) strides; the + # kernels read a packed LSE, so a strided stats input (strided_stats) + # is gathered into carved contiguous staging first. + lse = resolved[id(facts.stats_t)].squeeze(-1) + if stats_stage: + lse_stage = carver.take(b * h_q * facts.s_q, lse.dtype).view(b, h_q, facts.s_q) + lse_stage.copy_(lse) + lse = lse_stage + dbias_buf = resolved.get(id(facts.dbias_t)) if facts.has_dbias and facts.dbias_t is not None else None + dsink_buf = resolved.get(id(facts.dsink_t)) if facts.has_dsink and facts.dsink_t is not None else None + + api.execute( + q_tensor=_normalize(carver, resolved[id(facts.q_t)], stage_bytes["q"]), + k_tensor=_normalize(carver, resolved[id(facts.k_t)], stage_bytes["k"]), + v_tensor=_normalize(carver, resolved[id(facts.v_t)], stage_bytes["v"]), + o_tensor=_normalize(carver, resolved[id(facts.o_t)], stage_bytes["o"]), + do_tensor=_normalize(carver, resolved[id(facts.do_t)], stage_bytes["dO"]), + lse_tensor=lse, + dq_tensor=resolved[id(facts.dq_t)], + dk_tensor=resolved[id(facts.dk_t)], + dv_tensor=resolved[id(facts.dv_t)], + dbias_tensor=dbias_buf, + dsink_tensor=dsink_buf.view(-1) if dsink_buf is not None else None, scale_softmax=facts.scale, deterministic=facts.deterministic, # Stream from the caller's handle (ExecutionContext.stream); # None keeps the current stream. current_stream=stream, - **mask_args, + workspace=carver.remaining() if (carver is not None and api_scratch) else None, **ga.adapter_feature_buffers(facts, resolved), ) - - # copy_ casts in place; no .to() (which would allocate a staging - # tensor per execute). - for t_ref, key in ((facts.dq_t, "dq_tensor"), (facts.dk_t, "dk_tensor"), (facts.dv_t, "dv_tensor")): - resolved[id(t_ref)].copy_(out[key]) - if facts.has_dbias and "dbias_tensor" in out: - buf = resolved.get(id(facts.dbias_t)) - if buf is not None: - buf.copy_(out["dbias_tensor"].view(buf.shape)) - if facts.has_dsink and "dsink_tensor" in out: - buf = resolved.get(id(facts.dsink_t)) - if buf is not None: - buf.view(-1).copy_(out["dsink_tensor"]) return None - # Executor contract (engine._FrostSdpaBwdPlan): torch-native host code, - # no carved scratch — workspace_bytes 0 means _execute(variant_pack). - _execute.workspace_bytes = 0 + # Executor contract (engine._FrostSdpaBwdPlan): a non-zero workspace_bytes + # means _execute(variant_pack, workspace, stream) with the caller's buffer. + _execute.workspace_bytes = total_workspace_bytes _execute.binding = binding return _execute diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.py b/python/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.py index 17ca59365..8248c9ef0 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_d64_f16_sm80.py @@ -674,6 +674,22 @@ def _compile_unpermute(B, H, SQ, d, io_is_bf16): return cute.compile(_unpermute_host, fdq_acc, fdq_out, d, io_dtype, cutlass.Int32(0), cuda.CUstream(0), options="--enable-tvm-ffi") +def scratch_bytes(*, B: int, SQ: int, SKV: int, H: int, io_bytes: int = 2, need_do_dot: bool = True) -> int: + """Per-execute scratch requirement of ``backward(..., workspace=...)`` + (issue #514): the exact bytes it will carve. Keep in lockstep with the + ``_scratch`` takes there. d_qk == d_v == 64 by this kernel's contract.""" + from cudnn.sdpa.fwd.api_dsl import ws_align + + d = 64 + total = ws_align(B * H * SQ * d * 4) # permuted dQ accumulator (fp32) + total += ws_align(B * SKV * H * d * io_bytes) # dK + total += ws_align(B * SKV * H * d * io_bytes) # dV + total += ws_align(B * SQ * H * d * io_bytes) # dQ + if need_do_dot: + total += ws_align(B * H * SQ * 4) # do_dot (fp32) + return total + + def backward( Q: torch.Tensor, K: torch.Tensor, @@ -684,6 +700,9 @@ def backward( *, scale: Optional[float] = None, do_dot: Optional[torch.Tensor] = None, + workspace: Optional[torch.Tensor] = None, # caller scratch, sized by + # scratch_bytes(); every internal buffer is carved from it instead of + # allocated (issue #514). None → allocate (wrapper paths). **_ignored, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """SDPA backward for head-dim 64, fp16/bf16. BSHD in/out; ``lse`` natural-log @@ -702,19 +721,37 @@ def backward( scale_log2 = scale * math.log2(math.e) dev = Q.device + _carver = None + if workspace is not None: + from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver + + _carver = WorkspaceCarver( + workspace, + scratch_bytes(B=B, SQ=SQ, SKV=SKV, H=H, io_bytes=Q.element_size(), need_do_dot=do_dot is None), + "bprop_d64_f16_sm80", + ) + + def _scratch(numel, dtype, zero): + if _carver is None: + return (torch.zeros if zero else torch.empty)(numel, dtype=dtype, device=dev) + t = _carver.take(numel, dtype) + if zero: + t.zero_() + return t + # PERMUTED-flat dQ scratch [B, H, SQ, D] — the main kernel atomicAdds into it # thread-major; the _unpermute kernel casts it → row-major dQ. - dQ_acc = torch.zeros(B, H, SQ, D, dtype=torch.float32, device=dev) - dK = torch.empty(B, SKV, H, D, dtype=Q.dtype, device=dev) - dV = torch.empty(B, SKV, H, D, dtype=Q.dtype, device=dev) - dQ = torch.empty(B, SQ, H, D, dtype=Q.dtype, device=dev) + dQ_acc = _scratch(B * H * SQ * D, torch.float32, True).view(B, H, SQ, D) + dK = _scratch(B * SKV * H * D, Q.dtype, False).view(B, SKV, H, D) + dV = _scratch(B * SKV * H * D, Q.dtype, False).view(B, SKV, H, D) + dQ = _scratch(B * SQ * H * D, Q.dtype, False).view(B, SQ, H, D) lse_t = lse.to(dtype=torch.float32, device=dev).contiguous() stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) # do_dot (rowsum O∘dO) preprocessing reuses the shared device kernel. if do_dot is None: - dot_t = torch.empty(B, H, SQ, dtype=torch.float32, device=dev) + dot_t = _scratch(B * H * SQ, torch.float32, False).view(B, H, SQ) dd_fn = _base._compile_do_dot(B, H, SQ, D, io_is_bf16) dd_fn(from_dlpack(O), from_dlpack(dO), from_dlpack(dot_t), cutlass.Int32(B * H * SQ), stream) else: diff --git a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py index 56c12051f..8d9019915 100644 --- a/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py +++ b/python/cudnn/sdpa/bwd/kernels/bprop_f16_sm80.py @@ -1619,6 +1619,62 @@ def _compile_dsink(B, H, SQ): # =========================================================================== # Python entry point. # =========================================================================== + +# One-time cached 1-element zero dummies for absent optional operands (Rule 1 +# in python/cudnn/AGENTS.md permits a cached ``_dummy`` for a dead ABI slot; a +# fresh ``torch.zeros(1)`` per call is a per-execute allocation). Never read — +# the matching has_* Constexpr is False, so the slot is compiled out. +_DUMMY_CACHE: dict = {} + + +def _dummy1z(dtype, device): + key = (dtype, str(device)) + t = _DUMMY_CACHE.get(key) + if t is None: + t = _DUMMY_CACHE[key] = torch.zeros(1, dtype=dtype, device=device) + return t + + +def scratch_bytes( + *, + B: int, + SQ: int, + SKV: int, + H: int, + Hk: int, + d_qk: int, + d_v: int, + io_bytes: int = 2, + deterministic: bool = False, + has_bias: bool = False, + bias_batch: int = 1, + has_sink: bool = False, + need_do_dot: bool = True, + tile_q: int = _LLAMA_CFG.TILE_Q, +) -> int: + """Per-execute scratch requirement of the DENSE ``backward()`` path (issue + #514): the exact bytes ``backward(..., workspace=...)`` will carve. Keep in + lockstep with the ``_scratch`` takes there.""" + from cudnn.sdpa.fwd.api_dsl import ws_align + + total = ws_align(H * 4) if has_sink else 0 # dSink accumulator (fp32) + total += ws_align(B * SQ * H * d_qk * 4) # dQ_acc (fp32) + total += ws_align(B * SQ * H * d_qk * io_bytes) # dQ (io dtype) + if deterministic: + sem_units = B * H * ((SQ + tile_q - 1) // tile_q) + total += ws_align(max(sem_units, 1) * 4) # dq semaphore (int32) + total += ws_align(B * SKV * H * d_qk * io_bytes) # dK_ws + total += ws_align(B * SKV * H * d_v * io_bytes) # dV_ws + if H != Hk: # GQA: group-reduced outputs + total += ws_align(B * SKV * Hk * d_qk * io_bytes) + total += ws_align(B * SKV * Hk * d_v * io_bytes) + if has_bias: + total += ws_align(bias_batch * H * SQ * SKV * 4) # dBias accumulator (fp32) + if need_do_dot: + total += ws_align(B * H * SQ * 4) # do_dot (fp32) + return total + + def backward( Q: torch.Tensor, # [B, SQ, H, D] io_dtype (BSHD) K: torch.Tensor, # [B, SKV, H, D] @@ -1646,6 +1702,9 @@ def backward( tile_kv: int = _LLAMA_CFG.TILE_KV, tile_q: int = _LLAMA_CFG.TILE_Q, warps_per_sg: int = _LLAMA_CFG.WARPS_PER_SG, + workspace: Optional[torch.Tensor] = None, # caller scratch, sized by + # scratch_bytes() — every internal buffer is carved from it instead of + # allocated (issue #514). Dense-only; None → allocate (wrapper paths). ): """Full SDPA backward — only ``O`` and ``lse`` (forward outputs) + ``dO`` are needed beyond Q/K/V. ``do_dot`` is computed on-device from O·dO unless @@ -1680,6 +1739,40 @@ def backward( # n_seq logical sequences drive the grid + cu_* sizing. Q.shape[1]/ # K.shape[1] are the packed totals T_q/T_kv (== the kernel's SQ/SKV). --- thd = cu_seqlens_q is not None + _carver = None + if workspace is not None: + assert not thd, "workspace carving is dense-only (the engine path; THD comes via the wrappers)" + from cudnn.sdpa.fwd.api_dsl import WorkspaceCarver + + _carver = WorkspaceCarver( + workspace, + scratch_bytes( + B=B, + SQ=SQ, + SKV=SKV, + H=H, + Hk=Hk, + d_qk=d_qk, + d_v=d_v, + io_bytes=Q.element_size(), + deterministic=bool(deterministic), + has_bias=bias is not None, + bias_batch=(bias.shape[0] if bias is not None else 1), + has_sink=sinks is not None, + need_do_dot=do_dot is None, + tile_q=tile_q, + ), + "bprop_f16_sm80", + ) + + def _scratch(numel, dtype, zero): + if _carver is None: + return (torch.zeros if zero else torch.empty)(numel, dtype=dtype, device=Q.device) + t = _carver.take(numel, dtype) + if zero: + t.zero_() + return t + if thd: assert cu_seqlens_k is not None, "THD needs both cu_seqlens_q and cu_seqlens_k" assert B == 1, f"THD: Q/K/V/dO/O must be packed [1,T,H,D]; got batch dim {B}" @@ -1750,13 +1843,13 @@ def backward( rope_cs_t = torch.stack([angles.cos(), angles.sin()], dim=-1).contiguous() else: rope_max_s = 1 - rope_cs_t = torch.zeros(1, dtype=torch.float32, device=Q.device) + rope_cs_t = _dummy1z(torch.float32, Q.device) # Attention sink: dQ/dK/dV need NO kernel change (P recomputed from the # sink-aware LSE the caller passes); only dSink is computed (standalone). has_sink = sinks is not None if has_sink: sinks_t = sinks.to(dtype=torch.float32, device=Q.device).reshape(H).contiguous() - dsink_t = torch.zeros(H, dtype=torch.float32, device=Q.device) + dsink_t = _scratch(H, torch.float32, True) # THD is dense-feature-only for now (bias/rope/sink/seq_kv_lens are # dense-only); per-sequence padding is handled by the packed bounds, not # the PADDED mask. THD uses SCHED_DEFAULT (LPT+THD is a future tweak). @@ -1775,7 +1868,7 @@ def backward( if has_seq_len_q: seqq_t = seq_len_q.to(dtype=torch.int32, device=Q.device).contiguous() else: - seqq_t = torch.zeros(1, dtype=torch.int32, device=Q.device) + seqq_t = _dummy1z(torch.int32, Q.device) assert d_qk % 2 == 0 # dQ splits d-cols across the two sub-groups → each reads a DQ_N = d_qk//2 # column slice of sK. load_b_smem_x4 takes the d-col offset as `col_base` @@ -1802,32 +1895,32 @@ def backward( scale_log2 = scale * math.log2(math.e) inv_scale = 1.0 / float(scale) - dQ_acc = torch.zeros(B, SQ, H, d_qk, dtype=torch.float32, device=Q.device) - dQ = torch.empty(B, SQ, H, d_qk, dtype=Q.dtype, device=Q.device) + dQ_acc = _scratch(B * SQ * H * d_qk, torch.float32, True).view(B, SQ, H, d_qk) + dQ = _scratch(B * SQ * H * d_qk, Q.dtype, False).view(B, SQ, H, d_qk) # Deterministic-dQ relay counter: one int32 per (seq, head, q_tile), zeroed # per launch. Stride = ceil(max_SQ/tile_q) so the per-seq q_iter (THD) or the # dense q_iter both index in-bounds; n_seq sequences (= B dense). 1-elem dummy # (never touched) on the fast path so it costs nothing. sem_q_stride = (max_sq + tile_q - 1) // tile_q if deterministic else 0 sem_units = n_seq * H * sem_q_stride if deterministic else 1 - dq_sem = torch.zeros(max(sem_units, 1), dtype=torch.int32, device=Q.device) + dq_sem = _scratch(max(sem_units, 1), torch.int32, True) if deterministic else _dummy1z(torch.int32, Q.device) # dK/dV write buffers have H_q heads (one slice per query head — no atomics). # MHA (gqa_ratio==1): they ARE the outputs. GQA: a per-query-head workspace # that a reduction kernel sums over the group → [B,SKV,Hk,d] outputs. - dK_ws = torch.empty(B, SKV, H, d_qk, dtype=Q.dtype, device=Q.device) - dV_ws = torch.empty(B, SKV, H, d_v, dtype=Q.dtype, device=Q.device) + dK_ws = _scratch(B * SKV * H * d_qk, Q.dtype, False).view(B, SKV, H, d_qk) + dV_ws = _scratch(B * SKV * H * d_v, Q.dtype, False).view(B, SKV, H, d_v) if gqa_ratio == 1: dK, dV = dK_ws, dV_ws else: - dK = torch.empty(B, SKV, Hk, d_qk, dtype=Q.dtype, device=Q.device) - dV = torch.empty(B, SKV, Hk, d_v, dtype=Q.dtype, device=Q.device) + dK = _scratch(B * SKV * Hk * d_qk, Q.dtype, False).view(B, SKV, Hk, d_qk) + dV = _scratch(B * SKV * Hk * d_v, Q.dtype, False).view(B, SKV, Hk, d_v) lse_t = lse.to(dtype=torch.float32, device=Q.device).contiguous() # seq_kv_lens [B] int32 (or 1-elem dummy when not padded — never read). if has_seq_kv_lens: seqk_t = seq_kv_lens.to(dtype=torch.int32, device=Q.device).contiguous() else: - seqk_t = torch.zeros(1, dtype=torch.int32, device=Q.device) + seqk_t = _dummy1z(torch.int32, Q.device) # cu_seqlens [n_seq+1] int32 (THD) or 1-elem dummy (dense — never read). The # over-provisioned THD grid covers the longest sequence (ceil(max_skv/tile_kv) # kv-tiles) × H × n_seq; short sequences early-out per kv-tile. @@ -1837,27 +1930,27 @@ def backward( grid_kv_tiles = (max_skv + tile_kv - 1) // tile_kv grid_batch = n_seq else: - cu_q_t = torch.zeros(1, dtype=torch.int32, device=Q.device) - cu_k_t = torch.zeros(1, dtype=torch.int32, device=Q.device) + cu_q_t = _dummy1z(torch.int32, Q.device) + cu_k_t = _dummy1z(torch.int32, Q.device) grid_kv_tiles = 0 grid_batch = 0 # Bias + dBias (fp32 accumulator, same shape as bias; atomicAdd reduces over # batch when bias is broadcast [1,H,SQ,SKV]). if has_bias: bias_t = bias.contiguous() - dbias_t = torch.zeros(bias_batch, H, SQ, SKV, dtype=torch.float32, device=Q.device) + dbias_t = _scratch(bias_batch * H * SQ * SKV, torch.float32, True).view(bias_batch, H, SQ, SKV) else: # Dummy must match the fake tensor _compile_main builds at has_bias=False # (bias_is_fp32 defaults True → fp32). - bias_t = torch.zeros(1, dtype=torch.float32, device=Q.device) - dbias_t = torch.zeros(1, dtype=torch.float32, device=Q.device) + bias_t = _dummy1z(torch.float32, Q.device) + dbias_t = _dummy1z(torch.float32, Q.device) torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) # ---- do_dot: on-device (default) or caller-supplied ------------------ if do_dot is None: - dot_t = torch.empty(B, H, SQ, dtype=torch.float32, device=Q.device) + dot_t = _scratch(B * H * SQ, torch.float32, False).view(B, H, SQ) dd_fn = _compile_do_dot(B, H, SQ, d_v, io_is_bf16) dd_fn(from_dlpack(O), from_dlpack(dO), from_dlpack(dot_t), cutlass.Int32(B * H * SQ), stream) else: diff --git a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py index e424e8833..eee703677 100644 --- a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py @@ -304,3 +304,58 @@ def phys(t): torch.testing.assert_close(dq_buf.float(), q_ref.grad, rtol=3e-2, atol=3e-2) torch.testing.assert_close(dk_buf.float(), k_ref.grad, rtol=3e-2, atol=3e-2) torch.testing.assert_close(dv_buf.float(), v_ref.grad, rtol=3e-2, atol=3e-2) + + +@_SM80 +def test_engine_execute_does_not_allocate(): + """Issue #514: after the first (JIT + cache-warming) execute, re-executing + either SM80 engine must not touch the CUDA caching allocator — every + per-execute buffer is carved from the caller's workspace. Asserted on the + allocator's cumulative allocation COUNTER, which any torch.empty/zeros/ + clone/contiguous on the execute path would advance.""" + # Forward graph (with stats, so the LSE staging path runs too). + g, q, k, v, o, stats = _build_fwd_graph() + _native_then_pin(g, _FWD) + assert g.get_workspace_size() > 0, "the SM80 fwd executor must report its carved scratch" + torch.manual_seed(0) + q_buf, k_buf, v_buf = _buf(), _buf(), _buf() + o_buf = torch.empty_like(q_buf) + stats_buf = torch.empty(B, H, S, 1, dtype=torch.float32, device="cuda") + ws = _ws(g) + vp = {q: q_buf, k: k_buf, v: v_buf, o: o_buf, stats: stats_buf} + + # Backward graph on the same geometry. + gb = cudnn.pygraph(io_data_type=_HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + st = _bshd_stride(B, H, S, D) + qb = gb.tensor(name="q", dim=(B, H, S, D), stride=st, data_type=_HALF) + kb = gb.tensor(name="k", dim=(B, H, S, D), stride=st, data_type=_HALF) + vb = gb.tensor(name="v", dim=(B, H, S, D), stride=st, data_type=_HALF) + ob = gb.tensor(name="o", dim=(B, H, S, D), stride=st, data_type=_HALF) + dob = gb.tensor(name="dO", dim=(B, H, S, D), stride=st, data_type=_HALF) + statsb = gb.tensor(name="stats", dim=(B, H, S, 1), stride=(H * S, S, 1, 1), data_type=cudnn.data_type.FLOAT) + dq, dk, dv = gb.sdpa_backward(q=qb, k=kb, v=vb, o=ob, dO=dob, stats=statsb, attn_scale=_SCALE, use_causal_mask=True) + for t in (dq, dk, dv): + t.set_output(True).set_data_type(_HALF) + _native_then_pin(gb, _BWD) + assert gb.get_workspace_size() > 0, "the SM80 bwd executor must report its carved scratch" + do_buf = _buf() + dq_buf, dk_buf, dv_buf = torch.empty_like(q_buf), torch.empty_like(k_buf), torch.empty_like(v_buf) + wsb = _ws(gb) + vpb = {qb: q_buf, kb: k_buf, vb: v_buf, ob: o_buf, dob: do_buf, statsb: stats_buf, dq: dq_buf, dk: dk_buf, dv: dv_buf} + + # Warm run: kernel JIT, dummy/ALiBi caches, cublas handles, everything. + g.execute(vp, ws) + gb.execute(vpb, wsb) + torch.cuda.synchronize() + + ref_o, ref_dq = o_buf.clone(), dq_buf.clone() + before = torch.cuda.memory_stats()["allocation.all.allocated"] + for _ in range(3): + g.execute(vp, ws) + gb.execute(vpb, wsb) + torch.cuda.synchronize() + after = torch.cuda.memory_stats()["allocation.all.allocated"] + assert after == before, f"execute allocated {after - before} times; the SM80 engine paths must carve from the workspace only" + # And the carved re-executes still compute the same thing. + torch.testing.assert_close(o_buf, ref_o, rtol=0, atol=0) + torch.testing.assert_close(dq_buf, ref_dq, rtol=0, atol=0) From f5a2def53cd89911180684efa00204f224e4d9f4 Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Mon, 24 Aug 2026 11:00:13 -0700 Subject: [PATCH 2/4] frost(sdpa): SM80 forward scratch carving + strided-LSE staging (issue #514) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the forward half on the merged SdpaFwdDsl/TemplateParams architecture: SdpaFwdDslSm80.scratch_workspace_bytes() now sizes the dense_flex Q/K/V/O gathers, the GQA head expansion, the V head-dim pad (one carved buffer per operand — gather, expansion, and pad fused), strided-LSE staging, and the sinks log2 rescale; execute() carves them all from the caller's workspace through lower_dsl_prefill's existing plumbing (the standalone wrapper path keeps its allocating fallbacks). The 'LSE must be contiguous on SM80' build-time reject becomes carved staging + copy-back: with mhas randomizing dense stats strides on cuDNN >= 9.26 (#304), that reject was declining most stats-carrying forward graphs to the backend. Test updates: workspace passed at every SM80 graph execute; the no-alloc regression no longer requires a non-zero fwd workspace (a plain compact-BSHD MHA graph direct-binds everything and genuinely needs no scratch). Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 132 ++++++++++++++---- .../test_sdpa_sm80_frontend_integration.py | 21 ++- .../frost/test_sdpa_sm80_stream_respect.py | 6 +- 3 files changed, 126 insertions(+), 33 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 02e718fae..04c49b17e 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -3780,8 +3780,45 @@ def compile(self) -> None: ) self._logger.debug("compile completed") + def _bshd_gather_bytes(self, desc) -> int: + """Bytes to gather ``desc`` into a compact BSHD buffer, or 0 when its + BSHD transpose is already contiguous (the common, engine-normalized + case).""" + b, h, s, d = desc.shape + if tuple(desc.stride) == (s * h * d, d, h * d, 1): + return 0 + return ws_align(b * h * s * d * 2) # fp16/bf16 only on this row + def scratch_workspace_bytes(self) -> int: - return 0 + """Per-execute scratch (issue #514): dense_flex gathers, the GQA head + expansion, the V head-dim pad, the kernel-layout O staging those need, + strided-LSE staging, and the sinks log2 rescale — everything execute() + would otherwise allocate. Sized in execute()'s carve order.""" + self._ensure_support_checked() + if self.thd: + return 0 # engine rows never lower THD; the wrapper path allocates + elem = 2 # fp16/bf16 — check_support admits no other input dtype + b, hq, sq, skv = self.batch_size, self.h_q, self.s_q_max, self.s_k_max + gqa = self.h_kv != self.h_q + pad_v = self.head_dim_v < self.flavor_d_v + total = self._bshd_gather_bytes(self.q_desc) + # K/V: layout gather, GQA expansion, and the V pad share ONE carved + # buffer each (expanded heads at the padded flavor width). + if gqa or self._bshd_gather_bytes(self.k_desc): + total += ws_align(b * skv * hq * self.head_dim_qk * elem) + if pad_v: + total += ws_align(b * skv * hq * self.flavor_d_v * elem) + elif gqa or self._bshd_gather_bytes(self.v_desc): + total += ws_align(b * skv * hq * self.head_dim_v * elem) + # O: the compiled ABI is (B, SQ, H, flavor_d_v) — staged for the padded + # envelope or a non-BSHD (dense_flex) caller buffer. + if pad_v: + total += ws_align(b * sq * hq * self.flavor_d_v * elem) + elif self._bshd_gather_bytes(self.o_desc): + total += ws_align(b * sq * hq * self.head_dim_v * elem) + if self.has_sink: + total += ws_align(hq * 4) # sinks * log2(e) product + return total # ------------------------------------------------------------------ def execute( @@ -3823,39 +3860,79 @@ def execute( device = q_tensor.device launch_stream = self._get_default_stream(current_stream) - with _torch_stream_context(current_stream, device): - # BHSD → BSHD views; a dense_flex layout that is not BSHD-physical - # normalizes with one copy — the same grandfathered normalization - # the SM100 dense path applies (open cleanup, Hard Rule 2). - Q = self._to_bshd(q_tensor) - K = self._to_bshd(k_tensor) - V = self._to_bshd(v_tensor) - if self.h_kv != self.h_q: - # Dense GQA: expand K/V heads until the kernels' native dense - # GQA path is qualified (see class docstring). BSHD head dim is 2. - reps = self.h_q // self.h_kv - K = K.repeat_interleave(reps, dim=2) - V = V.repeat_interleave(reps, dim=2) + # Per-execute scratch: carved from the caller's workspace when one is + # provided (the engine lowering passes one sized by + # scratch_workspace_bytes(); issue #514), otherwise allocated (the + # standalone wrapper path). Carve order mirrors the sizing order. + carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "SdpaFwdDslSm80") if workspace is not None else None + with _torch_stream_context(current_stream, device): pad_v = self.head_dim_v < self.flavor_d_v - if pad_v: - V = _sm80_pad_last_dim(V, self.flavor_d_v) + gqa = self.h_kv != self.h_q + reps = self.h_q // self.h_kv + + def _gather_bshd(t: torch.Tensor) -> torch.Tensor: + """Compact BSHD view/gather of logical BHSD ``t`` (dense_flex).""" + view = t.transpose(1, 2) + if view.is_contiguous(): + return view + if carver is None: + return view.contiguous() + dst = carver.take(t.numel(), t.dtype).view(view.shape) + dst.copy_(view) + return dst + + def _kv_operand(t: torch.Tensor, fd: Optional[int]) -> torch.Tensor: + """K/V kernel operand: layout gather, GQA head expansion, and + the head-dim pad in ONE carved buffer (allocating fallbacks on + the wrapper path).""" + view = t.transpose(1, 2) # (b, s, h_kv, d) + bb, ss, hh, dd = view.shape + fd = dd if fd is None else fd + if not gqa and fd == dd: + return _gather_bshd(t) + if carver is not None: + dst = carver.take(bb * ss * self.h_q * fd, t.dtype).view(bb, ss, hh, reps, fd) + if fd != dd: + dst[..., dd:].zero_() + dst[..., :dd].copy_(view.unsqueeze(3)) + return dst.view(bb, ss, self.h_q, fd) + out = view.repeat_interleave(reps, dim=2) if gqa else view + if fd != dd: + out = _sm80_pad_last_dim(out, fd) + elif not out.is_contiguous(): + out = out.contiguous() + return out + + Q = _gather_bshd(q_tensor) + K = _kv_operand(k_tensor, None) + V = _kv_operand(v_tensor, self.flavor_d_v if pad_v else None) # Output binding: the compiled O ABI is (B, SQ, H, flavor_d_v). # Direct-bind the caller's BSHD view when it matches; the padded-V - # envelope and dense_flex cases go through a scratch + copy-back - # (both pre-existing normalizations). - o_view, o_needs_copyback, o_scratch = self._to_bshd_writable(o_tensor) + # envelope and dense_flex cases go through carved staging + + # copy-back. + o_view = o_tensor.transpose(1, 2) + o_needs_copyback = pad_v or not o_view.is_contiguous() if pad_v: - o_kernel = torch.zeros(self.batch_size, self.s_q_max, self.h_q, self.flavor_d_v, dtype=q_tensor.dtype, device=device) + if carver is not None: + o_kernel = carver.take(self.batch_size * self.s_q_max * self.h_q * self.flavor_d_v, q_tensor.dtype) + o_kernel = o_kernel.view(self.batch_size, self.s_q_max, self.h_q, self.flavor_d_v) + o_kernel.zero_() + else: + o_kernel = torch.zeros(self.batch_size, self.s_q_max, self.h_q, self.flavor_d_v, dtype=q_tensor.dtype, device=device) elif o_needs_copyback: - o_kernel = o_scratch + if carver is not None: + o_kernel = carver.take(o_tensor.numel(), o_tensor.dtype).view(o_view.shape) + else: + o_kernel = torch.empty_like(o_view, memory_format=torch.contiguous_format) else: o_kernel = o_view + # DEFENSIVE zero-fill, not load-bearing: the dense epilogue stores # every in-bounds row unconditionally; kept so a bound buffer can # never surface uninitialized memory if a future path skips rows. - # (The pad_v scratch above is allocated zeroed already.) + # (The pad_v staging above is zeroed already.) if (seq_q_lens is not None or seq_kv_lens is not None) and not pad_v: o_kernel.zero_() if lse_tensor is not None: @@ -3875,8 +3952,13 @@ def execute( ) if sinks is not None: # log2-unit rescale: one (H,)-element multiply per execute - # (pre-existing SM80 contract; the kernels consume log2 units). - sinks_b = (self._checked_sinks_1d(sinks) * _LOG2E).contiguous() + # (the kernels consume log2 units), into carved scratch. + checked_sinks = self._checked_sinks_1d(sinks) + if carver is not None: + sinks_b = carver.take(self.h_q, torch.float32) + torch.mul(checked_sinks, _LOG2E, out=sinks_b) + else: + sinks_b = (checked_sinks * _LOG2E).contiguous() else: sinks_b = self._dummy("one_f32", device, lambda: torch.ones(1, dtype=torch.float32, device=device)) if bias_tensor is not None: @@ -3938,7 +4020,7 @@ def execute( if pad_v: o_view.copy_(o_kernel[..., : self.head_dim_v]) elif o_needs_copyback: - o_view.copy_(o_scratch) + o_view.copy_(o_kernel) self._logger.debug("execute completed") diff --git a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py index eee703677..3d30266c1 100644 --- a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py @@ -62,6 +62,13 @@ def _build_fwd_graph(*, d=D, stats_stride=None, **sdpa_kwargs): return g, q, k, v, o, stats +def _ws(g): + """The graph's carved-scratch workspace (issue #514): a non-zero + workspace_bytes executor carves from the caller's buffer.""" + n = g.get_workspace_size() + return torch.empty(n, dtype=torch.uint8, device="cuda") if n else None + + def _native_then_pin(g, engine): g.validate() g.build_operation_graph() @@ -129,7 +136,7 @@ def test_fwd_engine_end_to_end(): q_buf, k_buf, v_buf = _buf(), _buf(), _buf() o_buf = torch.empty_like(q_buf) stats_buf = torch.empty(B, H, S, 1, dtype=torch.float32, device="cuda") - g.execute({q: q_buf, k: k_buf, v: v_buf, o: o_buf, stats: stats_buf}, None) + g.execute({q: q_buf, k: k_buf, v: v_buf, o: o_buf, stats: stats_buf}, _ws(g)) torch.cuda.synchronize() ref = torch.nn.functional.scaled_dot_product_attention(q_buf.float(), k_buf.float(), v_buf.float(), is_causal=True, scale=_SCALE).to(torch.float16) @@ -190,7 +197,7 @@ def test_bwd_engine_end_to_end(): q_buf, k_buf, v_buf = _buf(), _buf(), _buf() o_buf = torch.empty_like(q_buf) stats_buf = torch.empty(B, H, S, 1, dtype=torch.float32, device="cuda") - g.execute({q: q_buf, k: k_buf, v: v_buf, o: o_buf, stats: stats_buf}, None) + g.execute({q: q_buf, k: k_buf, v: v_buf, o: o_buf, stats: stats_buf}, _ws(g)) gb = cudnn.pygraph(io_data_type=_HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) st = _bshd_stride(B, H, S, D) @@ -209,7 +216,7 @@ def test_bwd_engine_end_to_end(): dq_buf, dk_buf, dv_buf = torch.empty_like(q_buf), torch.empty_like(k_buf), torch.empty_like(v_buf) gb.execute( {qb: q_buf, kb: k_buf, vb: v_buf, ob: o_buf, dob: do_buf, statsb: stats_buf, dq: dq_buf, dk: dk_buf, dv: dv_buf}, - None, + _ws(gb), ) torch.cuda.synchronize() @@ -249,7 +256,7 @@ def test_fwd_engine_bhsd_contiguous_layout(gqa): v_buf = torch.randn(B, h_kv, S, D, dtype=torch.float16, device="cuda") o_buf = torch.empty_like(q_buf) stats_buf = torch.empty(B, H, S, 1, dtype=torch.float32, device="cuda") - g.execute({q: q_buf, k: k_buf, v: v_buf, o: o_buf, stats: stats_buf}, None) + g.execute({q: q_buf, k: k_buf, v: v_buf, o: o_buf, stats: stats_buf}, _ws(g)) torch.cuda.synchronize() ref = torch.nn.functional.scaled_dot_product_attention( @@ -292,7 +299,7 @@ def phys(t): dq_buf, dk_buf, dv_buf = (torch.empty(dims, dtype=torch.float16, device="cuda") for _ in range(3)) gb.execute( {qb: bufs["q"], kb: bufs["k"], vb: bufs["v"], ob: o_buf, dob: bufs["do"], statsb: stats_buf, dq: dq_buf, dk: dk_buf, dv: dv_buf}, - None, + _ws(gb), ) torch.cuda.synchronize() @@ -316,7 +323,9 @@ def test_engine_execute_does_not_allocate(): # Forward graph (with stats, so the LSE staging path runs too). g, q, k, v, o, stats = _build_fwd_graph() _native_then_pin(g, _FWD) - assert g.get_workspace_size() > 0, "the SM80 fwd executor must report its carved scratch" + # The fwd executor may report 0 here: the merged SdpaFwdDsl port binds + # caller buffers directly, so a plain compact-BSHD MHA graph needs no + # scratch at all (the allocation counter below is the real contract). torch.manual_seed(0) q_buf, k_buf, v_buf = _buf(), _buf(), _buf() o_buf = torch.empty_like(q_buf) diff --git a/test/python/sdpa/frost/test_sdpa_sm80_stream_respect.py b/test/python/sdpa/frost/test_sdpa_sm80_stream_respect.py index 1354ee9e0..fce8f38ce 100644 --- a/test/python/sdpa/frost/test_sdpa_sm80_stream_respect.py +++ b/test/python/sdpa/frost/test_sdpa_sm80_stream_respect.py @@ -134,7 +134,8 @@ def test_sm80_bwd_respects_handle_stream_and_is_capturable(): q_gpu, k_gpu, v_gpu = _mk_buf(), _mk_buf(), _mk_buf() o_gpu = torch.empty_like(q_gpu) stats_gpu = torch.empty(_B, _H, _S, 1, device="cuda", dtype=torch.float32) - gf.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu, stats: stats_gpu}, None) + wsf = torch.empty(gf.get_workspace_size(), device="cuda", dtype=torch.uint8) if gf.get_workspace_size() else None + gf.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu, stats: stats_gpu}, wsf) torch.cuda.synchronize() gb = cudnn.pygraph(io_data_type=_HALF, intermediate_data_type=_F32, compute_data_type=_F32) @@ -154,4 +155,5 @@ def test_sm80_bwd_respects_handle_stream_and_is_capturable(): h = cudnn.create_handle() vp = {qb: q_gpu, kb: k_gpu, vb: v_gpu, ob: o_gpu, dob: do_gpu, statsb: stats_gpu, dq: dq_gpu, dk: dk_gpu, dv: dv_gpu} - _assert_stream_respect(lambda: gb.execute(vp, None, handle=h), [dq_gpu, dk_gpu, dv_gpu], h) + wsb = torch.empty(gb.get_workspace_size(), device="cuda", dtype=torch.uint8) if gb.get_workspace_size() else None + _assert_stream_respect(lambda: gb.execute(vp, wsb, handle=h), [dq_gpu, dk_gpu, dv_gpu], h) From 1e55d29834d8f83a3bfcbf3e821e45300c99a710 Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Mon, 24 Aug 2026 14:18:20 -0700 Subject: [PATCH 3/4] review: IR-shaped input views in lower_sm80_bwd; harden the no-alloc test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit round 1: - lower_sm80_bwd rebuilds IR-shaped views for the variant-pack INPUT buffers (q/k/v/o/dO/stats) before staging, mirroring the forward lowering's _ir_view — a flat or logically-reshaped caller buffer previously raised or mapped elements incorrectly in the rank-dependent staging paths. The gradient OUTPUTS deliberately keep the caller tensor's own view: output-port IR strides are provisional row-major unless user-assigned (the layout invariant in docs/python_graph_and_execution_backends.md), and re-striding the copy-back targets to the provisional layout scatters the writes (caught by test_bwd_engine_end_to_end, 83% dQ mismatch). - The no-alloc regression now exercises real staging on both directions: GQA (fwd K/V head expansion) plus a strided stats buffer (fwd LSE staging + bwd gather), asserts a non-zero fwd workspace, and clones/compares dK and dV alongside dQ and O across the re-executes. Verified on A100: SM80 suites all levels 118 passed; test_mhas_v2 bwd_L0 176/0, sdpa_bwd_sm80 serving all 176. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/bwd/engines.py | 31 +++++++++--- .../test_sdpa_sm80_frontend_integration.py | 50 +++++++++++++------ 2 files changed, 59 insertions(+), 22 deletions(-) diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index d49282b20..ebcf3a44d 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -614,13 +614,32 @@ def _normalize(carver, buf, staged: int): dst.copy_(buf.permute(0, 2, 1, 3)) return dst.permute(0, 2, 1, 3) + def _ir_view(buf, ir_t): + """Reinterpret a variant-pack buffer through the IR tensor's dim/stride. + + cuDNN's execute contract treats variant-pack entries as raw storage + laid out per the IR tensor descriptor — the caller's torch tensor may + be flat or otherwise logically reshaped. The staging/squeeze/copy_ + paths below consume torch views, so rebuild the IR-shaped view instead + of trusting the caller's metadata (mirrors the forward lowering's + ``_ir_view``). INPUT ports only: output-port IR strides are + PROVISIONAL row-major unless the user assigned them (the layout + invariant in docs/python_graph_and_execution_backends.md), so the + gradient outputs below keep the caller tensor's own view — re-striding + them to the provisional layout would scatter the copy-back. + """ + dim, stride = tuple(ir_t.get_dim()), tuple(ir_t.get_stride()) + if tuple(buf.shape) == dim and tuple(buf.stride()) == stride: + return buf + return buf.as_strided(dim, stride) + def _execute(variant_pack, workspace=None, stream=None): resolved = ga.resolve_variant_pack(variant_pack, binding) carver = WorkspaceCarver(workspace, total_workspace_bytes, spec.name) if total_workspace_bytes else None # squeeze(-1) is a valid view for ANY (B, H_q, S_q, 1) strides; the # kernels read a packed LSE, so a strided stats input (strided_stats) # is gathered into carved contiguous staging first. - lse = resolved[id(facts.stats_t)].squeeze(-1) + lse = _ir_view(resolved[id(facts.stats_t)], facts.stats_t).squeeze(-1) if stats_stage: lse_stage = carver.take(b * h_q * facts.s_q, lse.dtype).view(b, h_q, facts.s_q) lse_stage.copy_(lse) @@ -629,11 +648,11 @@ def _execute(variant_pack, workspace=None, stream=None): dsink_buf = resolved.get(id(facts.dsink_t)) if facts.has_dsink and facts.dsink_t is not None else None api.execute( - q_tensor=_normalize(carver, resolved[id(facts.q_t)], stage_bytes["q"]), - k_tensor=_normalize(carver, resolved[id(facts.k_t)], stage_bytes["k"]), - v_tensor=_normalize(carver, resolved[id(facts.v_t)], stage_bytes["v"]), - o_tensor=_normalize(carver, resolved[id(facts.o_t)], stage_bytes["o"]), - do_tensor=_normalize(carver, resolved[id(facts.do_t)], stage_bytes["dO"]), + q_tensor=_normalize(carver, _ir_view(resolved[id(facts.q_t)], facts.q_t), stage_bytes["q"]), + k_tensor=_normalize(carver, _ir_view(resolved[id(facts.k_t)], facts.k_t), stage_bytes["k"]), + v_tensor=_normalize(carver, _ir_view(resolved[id(facts.v_t)], facts.v_t), stage_bytes["v"]), + o_tensor=_normalize(carver, _ir_view(resolved[id(facts.o_t)], facts.o_t), stage_bytes["o"]), + do_tensor=_normalize(carver, _ir_view(resolved[id(facts.do_t)], facts.do_t), stage_bytes["dO"]), lse_tensor=lse, dq_tensor=resolved[id(facts.dq_t)], dk_tensor=resolved[id(facts.dk_t)], diff --git a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py index 3d30266c1..090505dac 100644 --- a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py @@ -319,45 +319,61 @@ def test_engine_execute_does_not_allocate(): either SM80 engine must not touch the CUDA caching allocator — every per-execute buffer is carved from the caller's workspace. Asserted on the allocator's cumulative allocation COUNTER, which any torch.empty/zeros/ - clone/contiguous on the execute path would advance.""" - # Forward graph (with stats, so the LSE staging path runs too). - g, q, k, v, o, stats = _build_fwd_graph() + clone/contiguous on the execute path would advance. The geometry is + chosen to force real staging on both directions: GQA (fwd K/V head + expansion) and a strided stats buffer (fwd LSE staging + bwd gather), + so the fwd workspace is non-zero too.""" + H_KV = H // 2 + st_kv = _bshd_stride(B, H_KV, S, D) + # Strided stats: (B, H, S, 1) declared with a 2-element row gap — the + # layout mhas draws under randomized stats strides (cuDNN >= 9.26, #304). + stats_stride = (2 * H * S, 2 * S, 2, 1) + + g = cudnn.pygraph(io_data_type=_HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + st = _bshd_stride(B, H, S, D) + q = g.tensor(name="q", dim=(B, H, S, D), stride=st, data_type=_HALF) + k = g.tensor(name="k", dim=(B, H_KV, S, D), stride=st_kv, data_type=_HALF) + v = g.tensor(name="v", dim=(B, H_KV, S, D), stride=st_kv, data_type=_HALF) + o, stats = g.sdpa(q=q, k=k, v=v, attn_scale=_SCALE, use_causal_mask=True, generate_stats=True) + o.set_output(True).set_data_type(_HALF) + stats.set_output(True).set_data_type(cudnn.data_type.FLOAT).set_stride(stats_stride) _native_then_pin(g, _FWD) - # The fwd executor may report 0 here: the merged SdpaFwdDsl port binds - # caller buffers directly, so a plain compact-BSHD MHA graph needs no - # scratch at all (the allocation counter below is the real contract). + assert g.get_workspace_size() > 0, "GQA expansion + strided-LSE staging must be carved, not allocated" torch.manual_seed(0) - q_buf, k_buf, v_buf = _buf(), _buf(), _buf() + q_buf = _buf() + k_buf = torch.randn(B, S, H_KV, D, dtype=torch.float16, device="cuda").permute(0, 2, 1, 3) + v_buf = torch.randn(B, S, H_KV, D, dtype=torch.float16, device="cuda").permute(0, 2, 1, 3) o_buf = torch.empty_like(q_buf) - stats_buf = torch.empty(B, H, S, 1, dtype=torch.float32, device="cuda") + stats_buf = torch.empty(2 * B * H * S, dtype=torch.float32, device="cuda").as_strided((B, H, S, 1), stats_stride) ws = _ws(g) vp = {q: q_buf, k: k_buf, v: v_buf, o: o_buf, stats: stats_buf} - # Backward graph on the same geometry. + # Backward graph on the same geometry (GQA + the same strided stats input). gb = cudnn.pygraph(io_data_type=_HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) - st = _bshd_stride(B, H, S, D) qb = gb.tensor(name="q", dim=(B, H, S, D), stride=st, data_type=_HALF) - kb = gb.tensor(name="k", dim=(B, H, S, D), stride=st, data_type=_HALF) - vb = gb.tensor(name="v", dim=(B, H, S, D), stride=st, data_type=_HALF) + kb = gb.tensor(name="k", dim=(B, H_KV, S, D), stride=st_kv, data_type=_HALF) + vb = gb.tensor(name="v", dim=(B, H_KV, S, D), stride=st_kv, data_type=_HALF) ob = gb.tensor(name="o", dim=(B, H, S, D), stride=st, data_type=_HALF) dob = gb.tensor(name="dO", dim=(B, H, S, D), stride=st, data_type=_HALF) - statsb = gb.tensor(name="stats", dim=(B, H, S, 1), stride=(H * S, S, 1, 1), data_type=cudnn.data_type.FLOAT) + statsb = gb.tensor(name="stats", dim=(B, H, S, 1), stride=stats_stride, data_type=cudnn.data_type.FLOAT) dq, dk, dv = gb.sdpa_backward(q=qb, k=kb, v=vb, o=ob, dO=dob, stats=statsb, attn_scale=_SCALE, use_causal_mask=True) for t in (dq, dk, dv): t.set_output(True).set_data_type(_HALF) _native_then_pin(gb, _BWD) assert gb.get_workspace_size() > 0, "the SM80 bwd executor must report its carved scratch" do_buf = _buf() - dq_buf, dk_buf, dv_buf = torch.empty_like(q_buf), torch.empty_like(k_buf), torch.empty_like(v_buf) + dq_buf = torch.empty_like(q_buf) + dk_buf, dv_buf = torch.empty_like(k_buf), torch.empty_like(v_buf) wsb = _ws(gb) vpb = {qb: q_buf, kb: k_buf, vb: v_buf, ob: o_buf, dob: do_buf, statsb: stats_buf, dq: dq_buf, dk: dk_buf, dv: dv_buf} - # Warm run: kernel JIT, dummy/ALiBi caches, cublas handles, everything. + # Warm run: kernel JIT, dummy caches, cublas handles, everything. g.execute(vp, ws) gb.execute(vpb, wsb) torch.cuda.synchronize() - ref_o, ref_dq = o_buf.clone(), dq_buf.clone() + ref_o = o_buf.clone() + ref_dq, ref_dk, ref_dv = dq_buf.clone(), dk_buf.clone(), dv_buf.clone() before = torch.cuda.memory_stats()["allocation.all.allocated"] for _ in range(3): g.execute(vp, ws) @@ -368,3 +384,5 @@ def test_engine_execute_does_not_allocate(): # And the carved re-executes still compute the same thing. torch.testing.assert_close(o_buf, ref_o, rtol=0, atol=0) torch.testing.assert_close(dq_buf, ref_dq, rtol=0, atol=0) + torch.testing.assert_close(dk_buf, ref_dk, rtol=0, atol=0) + torch.testing.assert_close(dv_buf, ref_dv, rtol=0, atol=0) From 30b59419a707b4c04a0b91a4320d5b740d153d73 Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Tue, 25 Aug 2026 12:03:05 -0700 Subject: [PATCH 4/4] review: L0 marker on the no-alloc regression test Co-Authored-By: Claude Fable 5 --- test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py index 090505dac..57f28dbf2 100644 --- a/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_sm80_frontend_integration.py @@ -313,6 +313,7 @@ def phys(t): torch.testing.assert_close(dv_buf.float(), v_ref.grad, rtol=3e-2, atol=3e-2) +@pytest.mark.L0 @_SM80 def test_engine_execute_does_not_allocate(): """Issue #514: after the first (JIT + cache-warming) execute, re-executing