[tbensonatl]

MatX queried CUDA device attributes on hot paths each time it needed them. Choosing launch parameters made ~5 cudaDeviceGetAttribute calls on every CUDA executor run. Attribute queries are typically fast (10s of ns) but add up with many queries.

Add detail::DeviceAttributes in core/device_attributes.h. It holds the attributes MatX uses that cannot change while a process runs, and is filled once per device on the first lookup. Later lookups take no locks, and with one visible device the current-device lookup also skips cudaGetDevice. GetDeviceAttributes throws matxCudaError when a query fails, and TryGetDeviceAttributes returns the error for callers that have a fallback. A failed query of a required attribute is not cached. An optional attribute whose query fails is stored as 0, which selects the conservative path. GetDeviceCount returns the cached device count. The SM count is not cached, since under MPS per-context partitioning it depends on the current context.

Move every library call site to the cache:

• The compute-capability helpers in utils.h, including the cuBLASLt workspace size.
• Tensor vector widths. EPTQueryInput now takes the target device's attributes and carries its widest vector load, so leaf tensors no longer query the device. MAX_EPT_VEC_LOAD takes the same input.
• NVRTC architecture resolution, which keeps its sm_80 fallback through the non-throwing TryGetCurrentDeviceAttributes, and the shared memory limit for JIT launches. A configured NVRTC_CUDA_ARCH is validated once, and "native" selects the runtime device without a warning.
• find_best_launch_params and the channelize_poly dispatcher, which still reads the SM count once per call.
• The allocator's managed-access and memory-pool checks. This replaces the allocator's separate memory-pool cache and removes DeviceSupportsMemoryPools.

Remove GetDeviceAttr as the cache replaces its uses.

The cuBLASDx, cuFFTDx, and cuSolverDx helpers now read the current device's compute capability when they need it, instead of storing the one that was current when the operator was built. JIT kernels launch on the current device, so an operator built on one GPU and run on another now targets the right architecture. The cuSolverDx helper reloads its cached traits when the architecture changes.

cudaMallocAsync allocates from the pool of the stream's device, which can differ from the current device. When every visible device supports memory pools, which is checked once, async allocations call cudaMallocAsync directly. Otherwise MatX finds the stream's device and checks its support. With one visible GPU that is device 0, the default-stream handles use the current device, and other streams are resolved through their context. When the stream's device lacks memory pools, the cudaMalloc fallback allocates on that device.

Move distributed_device_guard to core/resource_guard.h as CudaDeviceGuard so the allocator can use it, and add a record-only form that replaces the tests' own device guard. The guard now restores the caller's context with cuCtxSetCurrent. Restoring with cudaSetDevice made the device's primary context current in place of a context the caller created with the driver API, so cudaDeviceSynchronize stopped waiting for work queued there. With no context current, the guard restores the device instead, and it switches devices even when cudaGetDevice reports the target, since cudaGetDevice can then name a device that the next runtime call will not use. Add MATX_CUDA_DRIVER_CHECK_NOEXCEPT for the driver call in its destructor.

Fix two tests that CUDA 12.x could not compile, each a range-based for loop over a braced list whose later elements did not name their type.

On a test machine with an L4, host time per cudaExecutor run() for c = a * b + 1.0f on small tensors drops by ~20% (~4.2 us to ~3.2 us).

Add unit tests and a developer guide page for the cache, and update the memory documentation.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 5/5

[High impact] The PR appears safe to merge; no actionable issues were established.

Summary

Adds a shared, per-device cache for fixed CUDA attributes and moves launch sizing, allocation checks, and architecture helpers onto it.

• CUDA launch and JIT paths reuse each device’s fixed hardware facts.
• Async allocations use the device that owns their CUDA stream.
• MathDx helpers choose the architecture of the GPU where they run.
• Device switches restore the caller’s actual CUDA context.
• CUDA tests name each braced-list element’s type explicitly.

Review was based on source inspection; no builds or GPU tests were run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Attribute lookup] --> B{Device entry ready?}
  B -->|Yes| C[Return cached attributes]
  B -->|No| D[Lock device entry]
  D --> E[Query fixed CUDA attributes]
  E --> F{Required query failed?}
  F -->|Yes| G[Return error without caching]
  F -->|No| H[Publish completed entry]
  H --> C

Loading

_{Reviews (1) · Last reviewed commit: "Add per-device cache for static CUDA dev..." · Reviewed by Greptile}

[tbensonatl]

/build

[coveralls]

Coverage is 93.7% — tbenson/add-attribute-cache into main. No base build found for main.