feat(cuda.core): add cluster scheduling policy to LaunchConfig

cuda_core/cuda/core/_launch_config.pxd

@@ -16,6 +16,7 @@ cdef class LaunchConfig:
  public int shmem_size
  public bint is_cooperative
  public bint programmatic_stream_serialization
+ public str cluster_scheduling_policy_preference
 
  vector[cydriver.CUlaunchAttribute] _attrs
  object __weakref__

cuda_core/cuda/core/_launch_config.pyi

@@ -2,7 +2,10 @@
 
 from typing import Any
 
-_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization')
+from cuda.core._utils.cuda_utils import driver
+
+_LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization', 'cluster_scheduling_policy_preference')
+_CLUSTER_SCHED_POLICY_TO_DRIVER = {'DEFAULT': driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_DEFAULT, 'SPREAD': driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_SPREAD, 'LOAD_BALANCING': driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_LOAD_BALANCING}
 __all__ = ['LaunchConfig']
 
 class LaunchConfig:
@@ -39,15 +42,20 @@ class LaunchConfig:
  Whether to allow programmatic stream serialization (PDL). When True,
  the kernel may overlap with a previous kernel in the same stream that
  signals completion via programmatic means.
+ cluster_scheduling_policy_preference : str, optional
+ Cluster scheduling policy for the launch. One of ``"DEFAULT"``,
+ ``"SPREAD"``, or ``"LOAD_BALANCING"``. When omitted, the driver uses
+ the kernel function's default policy.
  """
  grid: tuple[Any, ...]
  cluster: tuple[Any, ...]
  block: tuple[Any, ...]
  shmem_size: int
  is_cooperative: bool
  programmatic_stream_serialization: bool
+ cluster_scheduling_policy_preference: str
 
- def __init__(self, grid: int | tuple[int, ...] | None=None, cluster: int | tuple[int, ...] | None=None, block: int | tuple[int, ...] | None=None, shmem_size: int | None=None, is_cooperative: bool=False, programmatic_stream_serialization: bool=False) -> None:
+ def __init__(self, grid: int | tuple[int, ...] | None=None, cluster: int | tuple[int, ...] | None=None, block: int | tuple[int, ...] | None=None, shmem_size: int | None=None, is_cooperative: bool=False, programmatic_stream_serialization: bool=False, cluster_scheduling_policy_preference: str | None=None) -> None:
  """Initialize LaunchConfig with validation.
 
  Parameters
@@ -64,12 +72,17 @@ class LaunchConfig:
  Whether to launch as cooperative kernel (default: False)
  programmatic_stream_serialization : bool, optional
  Whether to allow programmatic stream serialization / PDL (default: False)
+ cluster_scheduling_policy_preference : str, optional
+ Cluster scheduling policy for the launch: ``"DEFAULT"``,
+ ``"SPREAD"``, or ``"LOAD_BALANCING"`` (default: None)
  """
  def _identity(self) -> tuple[Any, ...]: ...
  def __repr__(self) -> str:
  """Return string representation of LaunchConfig."""
  def __eq__(self, other: object) -> bool: ...
  def __hash__(self) -> int: ...
+ def _validate_cluster_scheduling_policy_preference(self, value): ...
+ def _cluster_sched_policy_driver_value(self): ...
 
 def _to_native_launch_config(config: LaunchConfig) -> object:
  """Convert LaunchConfig to native driver CUlaunchConfig.

cuda_core/cuda/core/_launch_config.pyx

@@ -12,6 +12,7 @@ from cuda.core._utils.cuda_utils import (
  cast_to_3_tuple,
  driver,
 )
+from cuda.core._utils.validators import format_or_list
 
 _LAUNCH_CONFIG_ATTRS = (
  'grid',
@@ -20,8 +21,15 @@ _LAUNCH_CONFIG_ATTRS = (
  'shmem_size',
  'is_cooperative',
  'programmatic_stream_serialization',
+ 'cluster_scheduling_policy_preference',
 )
 
+_CLUSTER_SCHED_POLICY_TO_DRIVER = {
+ "DEFAULT": driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_DEFAULT,
+ "SPREAD": driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_SPREAD,
+ "LOAD_BALANCING": driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_LOAD_BALANCING,
+}
+
 __all__ = ['LaunchConfig']
 
 
@@ -59,6 +67,10 @@ cdef class LaunchConfig:
  Whether to allow programmatic stream serialization (PDL). When True,
  the kernel may overlap with a previous kernel in the same stream that
  signals completion via programmatic means.
+ cluster_scheduling_policy_preference : str, optional
+ Cluster scheduling policy for the launch. One of ``"DEFAULT"``,
+ ``"SPREAD"``, or ``"LOAD_BALANCING"``. When omitted, the driver uses
+ the kernel function's default policy.
  """
 
  # TODO: expand LaunchConfig to include other attributes
@@ -72,6 +84,7 @@ cdef class LaunchConfig:
  shmem_size: int | None = None,
  is_cooperative: bool = False,
  programmatic_stream_serialization: bool = False,
+ cluster_scheduling_policy_preference: str | None = None,
  ) -> None:
  """Initialize LaunchConfig with validation.
 
@@ -89,21 +102,31 @@ cdef class LaunchConfig:
  Whether to launch as cooperative kernel (default: False)
  programmatic_stream_serialization : bool, optional
  Whether to allow programmatic stream serialization / PDL (default: False)
+ cluster_scheduling_policy_preference : str, optional
+ Cluster scheduling policy for the launch: ``"DEFAULT"``,
+ ``"SPREAD"``, or ``"LOAD_BALANCING"`` (default: None)
  """
  # Convert and validate grid and block dimensions
  self.grid = cast_to_3_tuple("LaunchConfig.grid", grid)
  self.block = cast_to_3_tuple("LaunchConfig.block", block)
 
+ validated_policy = self._validate_cluster_scheduling_policy_preference(
+ cluster_scheduling_policy_preference
+ )
+
  # FIXME: Calling Device() strictly speaking is not quite right; we should instead
  # look up the device from stream. We probably need to defer the checks related to
  # device compute capability or attributes.
  # thread block clusters are supported starting H100
- if cluster is not None:
+ cc = None
+ if cluster is not None or validated_policy is not None:
  cc = Device().compute_capability
  if cc < (9, 0):
  raise CUDAError(
- f"thread block clusters are not supported on devices with compute capability < 9.0 (got {cc})"
+ "cluster launch attributes are not supported on devices with "
+ f"compute capability < 9.0 (got {cc})"
  )
+ if cluster is not None:
  self.cluster = cast_to_3_tuple("LaunchConfig.cluster", cluster)
  else:
  self.cluster = None
@@ -116,6 +139,7 @@ cdef class LaunchConfig:
 
  self.is_cooperative = is_cooperative
  self.programmatic_stream_serialization = programmatic_stream_serialization
+ self.cluster_scheduling_policy_preference = validated_policy
 
  if self.is_cooperative and not Device().properties.cooperative_launch:
  raise CUDAError("cooperative kernels are not supported on this device")
@@ -136,6 +160,25 @@ cdef class LaunchConfig:
  def __hash__(self) -> int:
  return hash(self._identity())
 
+ def _validate_cluster_scheduling_policy_preference(self, value):
+ if value is None:
+ return None
+ if isinstance(value, str) and value in _CLUSTER_SCHED_POLICY_TO_DRIVER:
+ cc = Device().compute_capability
+ if cc < (9, 0):
+ raise CUDAError(
+ "cluster launch attributes are not supported on devices with "
+ f"compute capability < 9.0 (got {cc})"
+ )
+ return value
+ valid = format_or_list(_CLUSTER_SCHED_POLICY_TO_DRIVER.keys())
+ raise ValueError(
+ f"{value!r} is not a valid cluster_scheduling_policy_preference. Must be {valid}"
+ )
+
+ def _cluster_sched_policy_driver_value(self):
+ return _CLUSTER_SCHED_POLICY_TO_DRIVER[self.cluster_scheduling_policy_preference]
+
  cdef cydriver.CUlaunchConfig _to_native_launch_config(self):
  cdef cydriver.CUlaunchConfig drv_cfg
  cdef cydriver.CUlaunchAttribute attr
@@ -169,6 +212,13 @@ cdef class LaunchConfig:
  attr.value.programmaticStreamSerializationAllowed = 1
  self._attrs.push_back(attr)
 
+ if self.cluster_scheduling_policy_preference is not None:
+ attr.id = cydriver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE
+ attr.value.clusterSchedulingPolicyPreference = int(
+ self._cluster_sched_policy_driver_value()
+ )
+ self._attrs.push_back(attr)
+
  drv_cfg.numAttrs = self._attrs.size()
  drv_cfg.attrs = self._attrs.data()
 
@@ -230,6 +280,13 @@ cpdef object _to_native_launch_config(LaunchConfig config):
  attr.value.programmaticStreamSerializationAllowed = 1
  attrs.append(attr)
 
+ if config.cluster_scheduling_policy_preference is not None:
+ attr = driver.CUlaunchAttribute()
+ attr.id = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE
+ # 13.0.2 setter reads .value; pass FastEnum, not a raw int.
+ attr.value.clusterSchedulingPolicyPreference = config._cluster_sched_policy_driver_value()
+ attrs.append(attr)
+
  drv_cfg.numAttrs = len(attrs)
  drv_cfg.attrs = attrs
 

cuda_core/tests/test_launcher.py

@@ -324,7 +324,7 @@ class _FakeDev:
  looked_up = []
  monkeypatch.setattr(_lc_mod, "Device", lambda: looked_up.append(1) or _FakeDev())
 
- with pytest.raises(CUDAError, match="thread block clusters are not supported"):
+ with pytest.raises(CUDAError, match="cluster launch attributes are not supported"):
  LaunchConfig(grid=2, cluster=2, block=32)
  assert looked_up, "Device was not looked up via the module global; mock did not take effect"
 
@@ -364,6 +364,97 @@ def test_to_native_launch_config_cluster_branch():
  assert (attr.value.clusterDim.x, attr.value.clusterDim.y, attr.value.clusterDim.z) == (2, 2, 2)
 
 
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_launch_config_cluster_scheduling_policy(monkeypatch):
+ """Ctor, getter/setter, and native attrs for all policy strings."""
+ from cuda.bindings import driver
+ from cuda.core import _launch_config as _lc_mod
+ from cuda.core._launch_config import _CLUSTER_SCHED_POLICY_TO_DRIVER, _to_native_launch_config
+
+ class _FakeDev:
+ compute_capability = (9, 0)
+
+ monkeypatch.setattr(_lc_mod, "Device", lambda: _FakeDev())
+ pref = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE
+
+ for policy in _CLUSTER_SCHED_POLICY_TO_DRIVER:
+ cfg = LaunchConfig(grid=1, block=1)
+ assert cfg.cluster_scheduling_policy_preference is None
+ cfg.cluster_scheduling_policy_preference = policy
+ assert cfg.cluster_scheduling_policy_preference is policy
+
+ cfg = LaunchConfig(grid=2, block=4, cluster_scheduling_policy_preference=policy)
+ assert cfg.cluster_scheduling_policy_preference is policy
+ native = _to_native_launch_config(cfg)
+ assert native.numAttrs == 1
+ assert native.attrs[0].id == pref
+ assert int(native.attrs[0].value.clusterSchedulingPolicyPreference) == int(
+ getattr(driver.CUclusterSchedulingPolicy, f"CU_CLUSTER_SCHEDULING_POLICY_{policy}")
+ )
+ cfg.cluster_scheduling_policy_preference = None
+ assert cfg.cluster_scheduling_policy_preference is None
+
+ cfg = LaunchConfig(
+ grid=(2, 1, 1),
+ block=32,
+ cluster=(2, 1, 1),
+ cluster_scheduling_policy_preference="SPREAD",
+ )
+ native = _to_native_launch_config(cfg)
+ assert native.numAttrs == 2
+ attr_ids = {attr.id for attr in native.attrs}
+ assert driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION in attr_ids
+ assert pref in attr_ids
+
+
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_launch_config_cluster_scheduling_policy_rejected(monkeypatch):
+ """Invalid values and pre-Hopper devices are rejected."""
+ from cuda.core import _launch_config as _lc_mod
+
+ with pytest.raises(ValueError, match="not a valid cluster_scheduling_policy_preference"):
+ LaunchConfig(grid=1, block=1, cluster_scheduling_policy_preference="NOT_A_POLICY")
+
+ class _FakeDev:
+ compute_capability = (8, 6)
+
+ looked_up = []
+ monkeypatch.setattr(_lc_mod, "Device", lambda: looked_up.append(1) or _FakeDev())
+ with pytest.raises(CUDAError, match="cluster launch attributes are not supported"):
+ LaunchConfig(
+ grid=2,
+ block=32,
+ cluster_scheduling_policy_preference="SPREAD",
+ )
+ assert looked_up, "Device was not looked up via the module global; mock did not take effect"
+
+
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_launch_cluster_scheduling_policy_smoke(init_cuda):
+ """launch() accepts each policy on Hopper+ (skip on CC < 9.0)."""
+ from cuda.core._launch_config import _CLUSTER_SCHED_POLICY_TO_DRIVER
+
+ dev = Device()
+ if dev.compute_capability < (9, 0):
+ pytest.skip("Cluster scheduling policy requires compute capability >= 9.0")
+
+ prog = Program('extern "C" __global__ void noop() {}', SourceCodeType.CXX)
+ kernel = prog.compile(ObjectCodeFormatType.CUBIN).get_kernel("noop")
+ stream = dev.default_stream
+ for policy in _CLUSTER_SCHED_POLICY_TO_DRIVER:
+ launch(
+ stream,
+ LaunchConfig(
+ grid=1,
+ block=32,
+ cluster=(2, 1, 1),
+ cluster_scheduling_policy_preference=policy,
+ ),
+ kernel,
+ )
+ stream.sync()
+
+
 def test_launch_invalid_values(init_cuda):
  code = 'extern "C" __global__ void my_kernel() {}'
  program = Program(code, SourceCodeType.CXX)

cuda_core/tests/test_object_protocols.py

@@ -704,7 +704,8 @@ def sample_object_b(request):
  "sample_launch_config",
  r"LaunchConfig\(grid=\(\d+, \d+, \d+\), cluster=.+, block=\(\d+, \d+, \d+\), "
  r"shmem_size=\d+, is_cooperative=(?:True|False), "
- r"programmatic_stream_serialization=(?:True|False)\)",
+ r"programmatic_stream_serialization=(?:True|False), "
+ r"cluster_scheduling_policy_preference=.+\)",
  ),
  ("sample_kernel", r"<Kernel handle=0x[0-9a-f]+>"),
  # ObjectCode variations (by code_type)