NVIDIA / NVIDIA/cuda-python

[FEA] cuda.core: support stream batch memory operations (incl. CUDA 13.1 atomic reductions)

Open
#2,364 0 comments 0 reactions 1 assignee View on GitHub

@juenglin is already working on this.

Since Jul 24, 2026.

cuda.core feature
Dominant language
Cython
Stars
3.4k
Forks
329
Avg merge
1d 23h
Merged PRs (30d)
116

Description

Summary

CUDA 13.1 extends the stream batch memory operations with atomic reductions: a new op type
CU_STREAM_MEM_OP_ATOMIC_REDUCTION with params struct CUstreamMemOpAtomicReductionParams
(operation, flags, reductionOp, dataType, address, value, alias), plus enums
CUstreamAtomicReductionOpType (ADD/AND/OR) and CUstreamAtomicReductionDataType
(UNSIGNED_32/UNSIGNED_64), joined into CUstreamBatchMemOpParams_union.atomicReduction.
Device support is advertised by CU_DEVICE_ATTRIBUTE_ATOMIC_REDUCTION_SUPPORTED (13.1, see
also #2361).

cuda.core has no batch mem-op surface at all today — the only mentions are a
graph debug-dot print flag
(https://github.com/NVIDIA/cuda-python/blob/c000331de6c37aa4565af74b001271ffcf6d5c99/cuda_core/cuda/core/graph/_graph_builder.pyx#L123)
and a DeviceProperties support query
(https://github.com/NVIDIA/cuda-python/blob/c000331de6c37aa4565af74b001271ffcf6d5c99/cuda_core/cuda/core/_device.pyx#L803).
So this issue covers designing the base surface (the pre-12.8 wait/write-value ops) and
including the 13.1 atomic-reduction op in it — adding the new op alone is not possible.

Underlying C APIs to cover

Symbol Purpose
cuStreamBatchMemOp(CUstream, unsigned int count, CUstreamBatchMemOpParams* paramArray, unsigned int flags) execute a batch of memory operations on a stream (pre-12.8, unexposed)
CU_STREAM_MEM_OP_ATOMIC_REDUCTION + CUstreamMemOpAtomicReductionParams (13.1) atomic-reduction op entry
CUstreamAtomicReductionOpType, CUstreamAtomicReductionDataType (13.1) reduction op / operand type
CU_DEVICE_ATTRIBUTE_ATOMIC_REDUCTION_SUPPORTED (13.1) support gating

Design sketch (draft — needs design-meeting review)

[!IMPORTANT]
Starting point only, not a settled design — review in the cuda.core design meeting.

stream.batch_mem_op([                          # names TBD
    WaitValue(addr, value, condition=">="),
    WriteValue(addr, value),
    AtomicReduction(address=addr, value=1, op="add", dtype="u64"),  # 13.1+
])

Open questions for the meeting:

  1. Scope: stream-level only, or also the graph batch-mem-op node type?
  2. Typed per-op dataclasses (above) vs. lower-level pass-through of binding structs.
  3. Semantics/purpose of the alias field (clarify against driver docs during design).
  4. Version/device gating (op availability differs across 12.x/13.x and devices).

References

-- Leo's bot

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.