CUDA Tile 1.3: store_masked crashes on element-wise-computed tiles

[newjordan]

CUDA Tile 1.3: store_masked crashes on element-wise-computed tiles

Environment: nvcc 13.3.33, CUDA Tile 1.3, sm_121 (GB10 / DGX Spark), driver 580.159.03

Summary

partition_view::store_masked produces illegal memory accesses when the stored tile
was produced by element-wise arithmetic (including ct::exp). The same store_masked
call works correctly when the tile was produced by ct::mma. The compiler appears to
generate store instructions targeting the wrong register file.

Minimal reproducer

// Build: nvcc -std=c++20 --enable-tile -arch=native -I<path>/include repro.cu -o repro
// Run:   ./repro
// Result: illegal memory access on the element-wise store path

#include <cstdio>
#include <cstdlib>
#include <cuda_runtime.h>
#include "cuda_tile.h"

using bf16 = __nv_bfloat16;

constexpr int M = 32, K = 64, N = 64;

__tile_global__ void repro(
    bf16* __restrict__ X,      // [M, K]  input
    bf16* __restrict__ W,      // [K, N]  weights
    float* __restrict__ Y_mma, // [M, N]  MMA-result output (WORKS)
    float* __restrict__ Y_ew   // [M, N]  element-wise output (CRASHES)
) {
  namespace ct = cuda::tiles; using namespace ct::literals;

  auto xV  = ct::partition_view{ct::tensor_span{X, ct::extents{M, K}}, ct::shape<M, K>{}};
  auto wV  = ct::partition_view{ct::tensor_span{W, ct::extents{K, N}}, ct::shape<K, N>{}};
  auto yM  = ct::partition_view{ct::tensor_span{Y_mma, ct::extents{M, N}}, ct::shape<M, N>{}};
  auto yE  = ct::partition_view{ct::tensor_span{Y_ew,  ct::extents{M, N}}, ct::shape<M, N>{}};

  using Acc = ct::tile<float, ct::shape<M, N>>;

  // Compute C = X @ W via MMA
  auto C = ct::full<Acc>(0.0f);
  C = ct::mma(xV.load_masked(0, 0), wV.load_masked(0, 0), C);

  // Path 1: store MMA result -- SUCCEEDS
  yM.store_masked(C, 0, 0);

  // Path 2: apply element-wise sigmoid, then store -- CRASHES
  auto one  = ct::full<Acc>(1.0f);
  auto C_sigmoid = C / (one + ct::exp(-C));
  yE.store_masked(C_sigmoid, 0, 0);  // ILLEGAL MEMORY ACCESS
}

int main() {
  size_t size_bf16 = (size_t)M * K * sizeof(bf16);
  size_t size_f32  = (size_t)M * N * sizeof(float);

  bf16 *dX, *dW;  float *dYm, *dYe;
  cudaMalloc(&dX, size_bf16);  cudaMemset(dX, 0, size_bf16);
  cudaMalloc(&dW, size_bf16);  cudaMemset(dW, 0, size_bf16);
  cudaMalloc(&dYm, size_f32);  cudaMemset(dYm, 0, size_f32);
  cudaMalloc(&dYe, size_f32);  cudaMemset(dYe, 0, size_f32);

  repro<<<1, 128>>>(dX, dW, dYm, dYe);
  cudaError_t err = cudaDeviceSynchronize();
  printf("result: %s\n", cudaGetErrorString(err));

  cudaFree(dX); cudaFree(dW); cudaFree(dYm); cudaFree(dYe);
  return err != cudaSuccess;
}

What works

Storing a tile that was accumulated via ct::mma works correctly:

auto C = ct::full<Acc>(0.0f);
C = ct::mma(A, B, C);       // MMA accumulation
yV.store_masked(C, 0, 0);   // SUCCEEDS

What crashes

Any element-wise operation on the tile breaks the store:

auto C = ct::full<Acc>(0.0f);
C = ct::mma(A, B, C);              // MMA accumulation
auto result = C / (one + ct::exp(-C));  // element-wise sigmoid
yV.store_masked(result, 0, 0);     // ILLEGAL MEMORY ACCESS

Failed workarounds

1. Zero-input MMA before store: compiler eliminates the no-op. Crash persists.
2. Shared-memory bridge: __shared__ forbidden in __tile_global__.
3. Flat-thread epilogue: threadIdx/blockDim forbidden in __tile_global__.

Production workaround

Keep element-wise-computed tiles in registers and use them directly as MMA
inputs, avoiding any store to HBM.

Expected behavior

When store_masked is called on a non-accumulator-resident tile, the compiler
should spill through shared memory or emit a flat-thread store.

Environment

nvcc: 13.3.33 (sbsa-linux)
CUDA Tile: 1.3.0
GPU: NVIDIA GB10 (sm_121)
Driver: 580.159.03
Host: aarch64 (Grace CPU)

Contact

Reproduced in the hydrak kernel research project.
Full source: src/tile_moe/tile_moe_kernel.cu