[BUG] Bulk TMA cache hints fail to compile and shared-memory over-alignment reduces occupancy #3295

Required prerequisites

• I have read the documentation https://tilelang.com.
• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

What version of TileLang are you using?

0.1.14

Source build; compiler baseline: fb5a47aed28f3b521f32c491b553220a5636e1b5.

System information

• Installation: built from source, Release native libraries.
• OS: Linux, kernel 6.8.0-41-generic.
• GPU: NVIDIA GeForce RTX 5090, SM120.
• NVIDIA driver: 580.95.05.
• Python: 3.13.13, built with Clang 22.1.3.
• PyTorch: 2.7.1+cu128; torch.version.cuda == "12.8".
• Kernel compiler: CUDA toolkit 12.9, NVCC 12.9.86.

Problem description

There are two issues in the descriptorless, linear bulk TMA path:

1. Non-default load cache hints fail to compile. A bulk load with eviction_policy="evict_first" or "evict_last" emits tl::tma_load<tl::CacheHintSm90::...>(...), but the bulk-load overload only accepts a type template parameter for the barrier. The descriptor-load and bulk-store overloads already accept cache hints.
2. Shared memory is over-aligned even when the kernel only needs linear bulk copies. Bulk lowering does not record its 16-byte requirement, so allocation planning falls back to 1024-byte alignment. Independently, CUDA codegen always declares the dynamic shared-memory arena with __align__(1024). This adds allocation padding and can increase the static shared-memory footprint around barriers enough to reduce occupancy.

The unhinted copy is correct in the baseline. The second issue is unnecessary resource usage that can cross an occupancy boundary.

Reproducible example code

Run the snippets as separate Python files with a CUDA-enabled TileLang build. The first compiles for the current GPU; the second explicitly generates SM120 source without launching a kernel.

1. Bulk-load cache-hint compilation failure

import tilelang
import tilelang.language as T

# Avoid reusing a kernel compiled by a different native build.
tilelang.disable_cache()


@T.prim_func
def bulk_copy(A: T.Tensor((1024,), T.float32), B: T.Tensor((1024,), T.float32)):
    with T.Kernel(1, threads=128):
        b_shared = T.alloc_shared((1024,), T.float32)
        T.copy(A, b_shared, prefer_instruction="tma", eviction_policy="evict_first")
        T.copy(b_shared, B, prefer_instruction="sync")


kernel = tilelang.compile(bulk_copy, out_idx=[1], target="cuda")
print(kernel.get_kernel_source())

Removing eviction_policy allows the baseline to compile. "evict_last" has the same template mismatch as "evict_first".

2. Bulk-copy and arena over-alignment

import tilelang
import tilelang.language as T
from tilelang import tvm


@T.prim_func
def footprint(
    A: T.Tensor((10260,), T.float32),
    B: T.Tensor((4096,), T.bfloat16),
    C: T.Tensor((10260,), T.float32),
):
    with T.Kernel(1, threads=256):
        a_state = T.alloc_shared((10260,), T.float32)
        scan = T.alloc_shared((4096,), T.bfloat16)
        T.copy(A, a_state, prefer_instruction="sync")
        T.copy(B, scan, prefer_instruction="tma")
        for i in T.Parallel(10260):
            C[i] = a_state[i] + scan[i % 4096].astype(T.float32)


config = {tilelang.PassConfigKey.TL_DISABLE_SHARED_MEMORY_REUSE: True}
with tvm.target.Target({"kind": "cuda", "arch": "sm_120"}) as target:
    with tvm.transform.PassContext(config=config):
        artifact = tilelang.lower(footprint, target=target)

for line in artifact.kernel_source.splitlines():
    if "extern __shared__" in line or "void* " in line:
        print(line)
func = next(iter(artifact.device_mod.functions.values()))
print("Dynamic shared bytes:", int(func.attrs["dyn_shared_memory_buf"]))

Baseline output includes:

extern __shared__ __align__(1024) uchar buf_dyn_shmem[];
void* a_state = ((void*)((char*)buf_dyn_shmem + 0));
void* scan = ((void*)((char*)buf_dyn_shmem + 41984));

Dynamic shared bytes: 50176

The first buffer occupies 41,040 bytes, already a multiple of 16. The linear bulk-copy destination could therefore start at byte 41,040 instead of 41,984, with a total dynamic allocation of 49,232 bytes.

Traceback

Relevant NVCC diagnostic from the failing bulk-load test, with temporary/source-tree path prefixes omitted:

tvm_kernels.cu: error: no instance of overloaded function "tl::tma_load" matches the argument list
            argument types are: (float *, const float *, Barrier, int)
      tl::tma_load<tl::CacheHintSm90::EVICT_FIRST>((&(b_shared[0])), (&(A[0])), A_to_b_shared_mbarrier[0], 4096);
      ^
copy_sm90.h(18): note #3323-D: substituting explicit template arguments "<<expression>>" for function template "tl::tma_load(void *, const void *, BarrierType &, uint32_t)" failed

The unhinted alignment reproducer does not throw an exception.

Expected behavior

• Bulk loads should compile with the supported default, evict_first, and evict_last policies and pass the selected policy to PTX rather than silently discard it.
• Shared-memory allocation offsets and the arena base should honor the maximum requirement of the actual operations. This example only requires 16-byte alignment for the linear bulk copy.
• Descriptor-based TMA and swizzled operands must retain their stricter alignment. This is not a request to replace all 1024-byte alignments with 16 bytes.
• TMA/MMA operands whose exact alignment is unknown should retain a conservative fallback, including single-allocation kernels.

Additional context

The over-alignment was encountered while experimenting with a TMA path for a TileLang port of DeepSeek's DeepSelect. The second reproducer isolates the shared-memory padding issue observed during that migration; it is not the full DeepSelect kernel.

Related reports

• #1842 reported a similar tma_load overload error, but its cause was a missing barrier when warp specialization was disabled, as documented in #2005. Here, the barrier is present and the mismatch is caused by the explicit cache-hint template argument.
• #2379 and #2391 address insufficient alignment for swizzled TMA and introduce per-buffer alignment requirements. This report is a follow-up for redundant alignment in the descriptorless bulk path and the dynamic shared-memory arena. The review of [CUDA] Align swizzled TMA shared buffers #2391 also noted that stale alignment requirements after a fallback could reduce occupancy; the reproducer here stays on the bulk TMA path rather than falling back to a normal copy.
• #2584 concerns a global-memory slice base that is not 16-byte aligned and causes a runtime fault. The buffers here satisfy the bulk-copy alignment requirement; the issue is unnecessary shared-memory padding.
• #3178 requests cache policies for ordinary SIMT loads/stores. The compilation failure here instead affects the existing TMA eviction_policy argument.

Relevant code paths

The alignment sources are separate:

• Copy::LowerBulk1D does not register a precise shared-memory alignment, unlike descriptor-based lowering.
• SharedMemoryAlignmentPlanner conservatively uses 1024 bytes when no explicit requirement is present.
• PrintStorageScope separately hardcodes 1024-byte arena alignment. Fixing only buffer offsets does not remove this source of padding.
• The bulk-load helper lacks the cache-hint template parameter expected by codegen.

The PTX cp.async.bulk specification specifies 16-byte alignment for linear bulk copies and supports an L2 cache-policy operand.