Skip to content

flashMLA utils.h: convert_type returns a tensor that points at a local array #20035

Open

Description

@peter941221

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

  1. added
    Customized kernels<NV>Specialized/modified CUDA kernels in TRTLLM for LLM ops, beyond standard TRT. Dev & perf.
    on Oct 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Customized kernels<NV>Specialized/modified CUDA kernels in TRTLLM for LLM ops, beyond standard TRT. Dev & perf.

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions