(tvm-ffi) branch main updated: [FIX] Use the opt-in shared memory limit in GetKernelWithMaxDynamicSharedMemory (#973)

[tqchen.apache.org via commits]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git


The following commit(s) were added to refs/heads/main by this push:
     new 99a7edde [FIX] Use the opt-in shared memory limit in 
GetKernelWithMaxDynamicSharedMemory (#973)
99a7edde is described below

commit 99a7eddedc76757e84ba0384db05b93c3568fe94
Author: 0z5a <dinadegttyaren...@gmail.com>
AuthorDate: Tue Nov 10 15:08:37 2026 +0100

    [FIX] Use the opt-in shared memory limit in 
GetKernelWithMaxDynamicSharedMemory (#973)
    
    Fixes #971
    
    `GetKernelWithMaxDynamicSharedMemory(name, -1)` should allow the maximum
    dynamic shared memory the device permits, minus the kernel's static
    shared memory. It computed that limit from device attribute 8
    (`MaxSharedMemoryPerBlock`), the 48 KiB default per-block limit. The
    attribute it set was therefore what a kernel can already use without
    opting in, and launches with more than 48 KiB of dynamic shared memory
    still failed with `invalid argument`.
    
    This change queries attribute 97
    (`cudaDevAttrMaxSharedMemoryPerBlockOptin` /
    `CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN`) instead. An
    explicit `dynamic_smem_max` and launches up to 48 KiB behave as before.
    
    The new `test_cubin_launcher_max_dynamic_smem` loads a kernel with 1 KiB
    of static shared memory through
    `GetKernelWithMaxDynamicSharedMemory(name, -1)`. It checks that
    `maxDynamicSharedSizeBytes == optin - static`, then launches with that
    much dynamic shared memory and checks the result. On GH200 (aarch64,
    sm_90) and B300 (sm_103) it fails on main
    (`maxDynamicSharedSizeBytes=48128, expected 231424`) and passes with
    this change. The three existing CUBIN tests still pass on both.
    
    The new test is behind the existing `torch.version.cuda > (13, 0)` gate,
    like the other CUBIN tests. To run it with a cu130 torch (`"13.0"`), the
    gate was relaxed locally; that change is not part of this PR.
    
    ---------
    
    Signed-off-by: 0z5a <dezhen...@student.uni-tuebingen.de>
    Signed-off-by: tqchen <tqchen...@gmail.com>
    Co-authored-by: tqchen <tqchen...@gmail.com>
---
 include/tvm/ffi/extra/cuda/cubin_launcher.h | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/include/tvm/ffi/extra/cuda/cubin_launcher.h 
b/include/tvm/ffi/extra/cuda/cubin_launcher.h
index 6082cd7d..6ec15518 100644
--- a/include/tvm/ffi/extra/cuda/cubin_launcher.h
+++ b/include/tvm/ffi/extra/cuda/cubin_launcher.h
@@ -537,7 +537,8 @@ class CubinKernel {
    *
    * \param dynamic_smem_max Maximum dynamic shared memory in bytes to set.
    *                         -1 (default) means maximum available dynamic 
shared memory,
-   *                         which is computed as (device max shared memory - 
static shared memory).
+   *                         which is computed as (device opt-in max shared 
memory per block
+   *                         - static shared memory).
    *                         For -1, the method queries the kernel's static 
shared memory usage
    *                         and sets the attribute to the remaining available 
shared memory.
    *
@@ -556,12 +557,11 @@ class CubinKernel {
     bool any_success = false;
     for (int device_id = 0; device_id < device_count; ++device_id) {
       auto device = cuda_api::GetDeviceHandle(device_id);
-      // Query device's maximum shared memory per block
+      // Query the maximum shared memory per block a kernel can opt into
+      // (CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN /
+      // cudaDevAttrMaxSharedMemoryPerBlockOptin), not the 48 KiB default limit
       int max_shared_mem = 0;
-      err = cuda_api::GetDeviceAttribute(
-          &max_shared_mem,
-          /* 
CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK/cudaDevAttrMaxSharedMemoryPerBlock
 */
-          cuda_api::DeviceAttrType(8), device);
+      err = cuda_api::GetDeviceAttribute(&max_shared_mem, 
cuda_api::DeviceAttrType(97), device);
       if (err != cuda_api::kSuccess) {
         continue;  // Skip this device if we can't get its attribute
       }
@@ -575,7 +575,7 @@ class CubinKernel {
         }
 
         // Calculate available dynamic shared memory:
-        // device max shared memory - static shared memory used by kernel
+        // device opt-in max shared memory per block - static shared memory 
used by kernel
         int64_t max_shared = static_cast<int64_t>(max_shared_mem);
         int64_t available = max_shared - static_shared;
         shared_mem_to_set = (available > 0) ? static_cast<int>(available) : 0;