You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
flashMLA utils.h: convert_type returns a tensor that points at a local array #20035
In cpp/tensorrt_llm/kernels/flashMLA/utils.h (line 277-286):
auto frag = convert_op(*reinterpret_cast<cutlass::Array<From_type, numel> const*>(tensor.data()));
return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());
Same defect as #19943. make_rmem_ptr on a raw pointer is the identity, so the returned tensor is a view of frag, a local of convert_type; reading the result after the call is undefined behaviour. Both live call sites are in flash_fwd_mla_kernel.h:
line 271 (store epilogue): Tensor rO = flash::convert_type<ElementO>(tOrO);, then retile_S and cute::copy read through the pointer; tOrO is the partition_fragment_C accumulator.
line 817 (splitkv combine kernel): Tensor rO = flash::convert_type<Element>(tOrO);, then cute::copy(rO, gO) reads it.
convert_type_out in the same file takes the output by reference and is fine; its comment already notes that the output can be wrong when the result is allocated and returned inside the helper, so the returning variant looks like a known hazard that was not fixed.
With nvcc the read returns stale values (registers/local memory), so it is silent in practice; it is still the same lifetime bug that was fixed in cutlass_extensions by #20017.
Reproduction without a GPU: compile the function verbatim in a host translation unit, read the returned tensor after the call, run under AddressSanitizer with ASAN_OPTIONS=detect_stack_use_after_return=1. The current code is reported as a stack-use-after-return pointing at frag inside the convert_type frame. Returning make_tensor<To_type>(tensor.layout()) (an owning ArrayEngine) with the elements copied in is clean. The contiguity the HACK comment assumes can be enforced with static_assert(cosize_v<Layout> == numel, ...), which holds for both real instantiations: the partition_fragment_C layout has cosize 128 == size 128, and the combine kernel's make_tensor<ElementAccum>(shape(...)) layout has cosize 4 == size 4.
This helper is a copy of FlashInfer's include/flashinfer/attention/hopper/utils.cuh (convert_type), which still has the same pattern at line 153.
Activity