[jesusjodar]

Summary

Extracted from #797 at the reviewer's request (review). This PR only changes the generic AIME gate's GPU free-memory probe.

The CUDA path parses the first CUDA_VISIBLE_DEVICES entry as an integer and then matches an NVML / nvidia-smi index. A GPU UUID raises ValueError, and CUDA device ordering can differ from NVML ordering.

Query torch.cuda.mem_get_info(0) in a child process on CUDA, as the ROCm path already does. Logical device zero follows the child's inherited visibility and ordering; the subprocess exits before engine creation, so the parent does not gain an initialized CUDA context from this probe. Remove the now-unused integer-only helper.

Validation provenance

This is the exact change previously exercised in #797 at 711a86ef11c67bbe6c69ff59370594f64c6ef1b8. The existing AIME gate passed on an RTX 3060 selected by UUID. The failing integer conversion can be read directly in the old helper; no separate before/after runtime reproduction was recorded.

The new branch is based on main at 32056cf4608b07c21afc45fa510c74fba67dcf17 and contains only this harness change. Tests were not rerun after extraction. The historical Flint run below also required the Flint model support in #797; it is not a claim that unmerged Flint support exists on this independent branch. ROCm and alternate numeric CUDA ordering were not separately exercised locally.

Environment and historical command

• RTX 3060 12 GiB only; TP=1; Ryzen 5 5600X, about 32 GiB RAM.
• Linux 7.2.8-9-omarchy, NVIDIA 615.78.08, Python 3.12.14.
• PyTorch 2.14.1+cu130, Triton 3.8.0, FlashInfer 0.6.18.post1.
• Transformers 5.18.0, tokenizers 0.23.3, safetensors 0.8.0, apache-tvm-ffi 0.1.13.post3, pytest 9.1.1.
• Source through PYTHONPATH=python in an isolated uv-created venv. Native extensions (including ggml-cpu variants) and prebuilt kernels were reused from freetoken 0.2.0+gfe5a5ec9c / freetoken-kernel-cache 0.2.0+cu130.gfe5a5ec9c. No clean native-extension/CUDA build was performed for this revision; this PR changes no C++/CUDA source.

Checkpoint	Revision
FutureMLS-Lab/flint-nvfp4-coldgptq	56f37fcfc7bd17834512ec2e53696b8da15f3658
RecursiveSI/flint-nvfp4-gguf, flint-nvfp4.gguf	6ec43e19b565fc8d560e2d04985d7724f20673bf
RecursiveSI/flint reference config/tokenizer	2a247be0c0aec895ccb0ac7ca92229ec6df37288

Commands below ran from the source checkout; the sibling venv/checkpoint paths reflect the local validation layout:

export PYTHONPATH=python
export CUDA_VISIBLE_DEVICES=GPU-a0591354-3462-ad0e-6e65-e772fd9a045e
export FLASHINFER_USE_CUDA_NORM=1
PYTHON=../../flint-research/dev-venv/bin/python
MODEL="$PWD/../../flint-research/models/flint-modelopt"
GGUF=/home/jesusjodar/.freetoken/models/Flint-NVFP4-GGUF/flint-nvfp4.gguf

Historical AIME result

Dataset: math-ai/aime25, revision 563bb8404243c5f09de6ec262f2db674fe5bce9b, test.jsonl; convert problem to the gate's prompt key and preserve answer and row order. AIME 2025 problem 1, pass@3, 8192-token maximum per sample. This is one problem, not an AIME-suite accuracy claim. Sampling comes from the checkpoint: temperature 1.0, top-p 0.95, top-k 20.

FREETOKEN_TEST_MODEL="$MODEL" \
FREETOKEN_AIME25_JSONL="$PWD/../aime25.jsonl" \
FREETOKEN_AIME_REQ=1 FREETOKEN_AIME_MAX_TOKENS=8192 FREETOKEN_AIME_SAMPLES=3 \
FREETOKEN_AIME_MIN_FREE_GIB=6 \
FREETOKEN_TEST_MOE_CACHE_SIZE=1024 FREETOKEN_TEST_MOE_STRATEGY=hybrid \
FREETOKEN_TEST_MOE_CPU_THREADS=4 FREETOKEN_TEST_MEM_RATIO=0.9 \
"$PYTHON" -m pytest tests/e2e/test_aime.py -q -s

Result: 1 passed, 1 warning in 61.45s; existing CUDA allocator deprecation warning.

Implementation and local test execution were assisted by Codex.