[josh-schell-amd]

Motivation

In the layout-dynamic from_dlpack path, DLTensorJitArg.__c_abi_spec__ hands the launch's DLTensorAdaptor from the pointer fill to the layout fill via a thread-local (_shared.dladaptor, added in #682) and never clears it. The fills live as long as the cached CallState, so each such argument keeps its last launch's tensor alive until the next launch.

In AITER's gfx950 A16W16 GEMM (ROCm/aiter#6853), whose arguments use layout-dynamic from_dlpack, FlyDSL kept each launch's input and output tensors alive after the launch. During startup, vLLM runs MiniMax-M3's vision-encoder projection once with 768000 rows (input 768000×1280, output 768000×6144). FlyDSL held on to that input (1.8 GiB) and output (8.8 GiB) after vLLM had freed them, so 10.6 GiB stayed allocated, and vLLM's KV cache came out 7.4 GiB smaller on MI355X. The static-layout from_dlpack path and from_torch_tensor (TorchTensorJitArg) don't keep tensors alive, so this looks unintended.

Technical Details

• jit_argument.py: the layout fill clears _shared.dladaptor after reading it. The handoff and slot order are unchanged; the cost is one attribute store per layout-dynamic argument per launch.
• test_callstate_dispatch.py: new test runs both fills on a CPU tensor, drops the tensor, and asserts a weakref to it is dead.

The retained reference was never a keepalive for in-flight kernels (it covered only the last launch per thread), so stream safety is unchanged.

Test Plan

pytest tests/unit/test_callstate_dispatch.py against FlyDSL 0.4.0 (identical code for this function) with and without the change; black and ruff on the changed files. GPU check on MI355X: torch memory_reserved() after the AITER GEMM above, then del + empty_cache().

Test Result

	Without fix	With fix
test_callstate_dispatch.py	new test fails (tensor still alive), 11 pass	12 pass
Reserved after free (MI355X)	10.64 GiB	0.02 GiB

test_launch_overhead.py not run.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

AI-assisted (Claude Code); results were run and checked on MI355X.

🤖 Generated with Claude Code