ITADN

[BUG]

#3327Openharshvardhanagg 创建于 2026-06-17
bug? - Needs TriageCuTe DSL
H
harshvardhanaggcommented
### Which component has the problem? CuTe DSL ### Bug Report ### Version - `nvidia-cutlass-dsl` 4.5.2 (latest stable on PyPI; also present in 4.5.0/4.5.1) - JAX integration (`cutlass.jax.cutlass_call`) - Target: SM100 (Blackwell), Python 3.13, CUDA 12.x ### Summary When a JAX input tensor has **any stride ≥ 2³¹ elements** and the kernel is wrapped with `cutlass_call(..., use_static_tensors=False)`, compilation crashes while the dynamic-stride layout is being built. The divisibility/alignment assumption for that stride is passed into `ConstrainedIntType.get(divisible_by, width)`, whose `divisible_by` argument is bound as a C++ **int32**, so any value ≥ 2³¹ raises `TypeError: incompatible function arguments`. This is especially surprising when the offending stride belongs to a **size-1 leading dimension** (a degenerate axis that is never indexed): its stride equals the size of the entire inner block, so a contiguous `[1, N, H, D]` tensor with `N·H·D ≥ 2³¹` crashes even though batch index 0 is the only one ever used. Related (runtime int32-stride limitation): #2312. This report is the **compile-time** `cutlass.jax` divisibility path, which I don't think is covered there. ### Reproducer ```python import jax, jax.numpy as jnp import cutlass, cutlass.cute as cute from cutlass.jax import cutlass_call import cuda.bindings.driver as cuda_driver # Leading (size-1) axis stride = N*H*D = 839680*20*128 = 2_149_580_800 > 2**31 B, N, H, D = 1, 839680, 20, 128 @cute.jit def launch(stream: cuda_driver.CUstream, mX: cute.Tensor, mO: cute.Tensor): # body irrelevant; crash happens before this runs, while building the # dynamic-stride layout for the inputs. ... call = cutlass_call( launch, output_shape_dtype=[jax.ShapeDtypeStruct((B, N, H, D), jnp.bfloat16)], use_static_tensors=False, # <-- True compiles fine; False crashes ) x = jnp.zeros((B, N, H, D), jnp.bfloat16) call(x) # raises during compilation
0 条评论