Cross-family runtime safety and long-context issues found during #1158 review

[ruiling-smartbear]

Context

Reviewing #1158 surfaced several defects in code copied from the existing Llama/model-family templates. I re-audited each call path against current main and kept only findings with a concrete failure mechanism. These are not SmolLM3-specific regressions and should not be fixed inside #1158.

#1176 already tracks two other confirmed shared defects: unchecked CUDA device queries in the sparse multinomial policy and GPU fixture tests that can report a false pass.

The counts below are occurrences of the exact code pattern on main, not a claim that every family enables the affected optional feature.

Confirmed findings

• Check sampler cudaMalloc results transactionally — allocations are used without checking status. In ensure_device_buffers, capacity_ = keep runs unconditionally after both cudaMalloc calls, so a failed allocation still advances the capacity and is never retried: every later call with keep <= capacity_ skips reallocation and reuses the bad pointers. The sparse path can then return a stale token when d_token_id_ is null, and the greedy CUDA kernel can receive a null output pointer. Add allocation-failure injection coverage, throw on allocation failure, and only commit new pointers and capacity after every allocation succeeds. Original finding: feat(smollm3): add native SmolLM3 model family #1158 (comment)

• Harden temporary FFI kernel loading — bundle kernel bytes are written with std::ofstream to predictable /tmp/trtmc_kernel_<global_name>.so paths. Creation is neither exclusive nor no-follow, the complete write is not checked, and files are not removed. A pre-created symlink is followed and concurrent processes using the same global name share the same file before it is loaded. The exact filename pattern appears in 83 family helper files.

  This one does not need a new design. src/runtime/models/qwen3_8/plugin_helpers.cpp already carries a hardened version of the same helper: it creates a private 0700 directory with mkdtemp, opens the file with O_CREAT | O_EXCL | O_NOFOLLOW and S_IRUSR | S_IWUSR, loops until the whole buffer is written, checks close, cleans up on any failure, and pairs with a remove_kernel_so_temp that unlinks the file and directory after the loader is done. Its comment states the same threat model. The work is to propagate that implementation to the remaining 83 files, which keeps the change mechanical and gives the reviewer an in-repo reference to compare against rather than a new design to judge. Original finding: feat(smollm3): add native SmolLM3 model family #1158 (comment)

• Keep TriAttention's scored row count out of CUDA grid.y — launch_score_kernel passes candidate_count as grid.y and the kernel reads it through blockIdx.y, and the caller passes total_tokens, the current cached row count. CUDA's y-dimension limit is 65,535, so a cache of 65,536 rows makes the launch fail; cudaGetLastError() then returns kFailed, run_gpu_selection_over_layers returns false, and selection falls back to the host implementation with no diagnostic. Swap the grid dimensions or chunk the launch, and test the 65,535/65,536 boundary.

  On reachability, so this is not read as broader than it is: TriAttention is a per-bundle opt-in (tri_cfg.enabled from ..._parse_triattention_bundle_config, which also requires a stats section in the bundle), no family declares it in MODEL.toml, and the GPU selection path only runs once total_tokens exceeds the keep budget. Reaching the limit therefore needs a bundle that enables TriAttention with a budget near 65,536, which is an unusual configuration for a cache-compaction strategy. The boundary itself is deterministic once reached. Original finding: feat(smollm3): add native SmolLM3 model family #1158 (comment)

Suggested implementation shape

Use this issue as a tracker, but fix each category in its own mechanical cross-family PR. Each PR should first add a focused regression test or deterministic failure injection, then update every family carrying the same reachable implementation.

Potential findings whose impact or intended policy is not yet proven are deliberately excluded. In particular, this issue does not claim a correctness race from the sampler's legacy default stream, does not propose a rank-1 mask sweep without a reachable current engine, and does not prescribe the intended E2E prefix-comparison policy.

An earlier revision of this issue listed the unbounded linear-spec masked-token loop as a fourth finding. I have removed it: supports_text_diffusion defaults to false in every family and only nemotron_labs_diffusion sets it, so the 26 other pipelines carrying that loop throw in resolve_text_diffusion_block_length before reaching it. The remaining exposure is one family under one runtime strategy, which is too narrow to track here.

[ruiling-smartbear]

@yifeif-nv, could you please advise on the TriAttention boundary item here? The score kernel currently maps the total cached-row count to grid.y; at 65,536 rows the launch exceeds CUDA’s 65,535 y-dimension limit and the selection path falls back to the host. This becomes particularly relevant for a 128K YaRN configuration, although it does not affect YaRN numerics or output correctness. Would swapping the grid mapping to grid.x = total_tokens, grid.y = kv_head_count (with the corresponding blockIdx mapping change and 65,535/65,536/131,072 boundary tests) match the intended design, or is host fallback beyond 65,535 rows intentional? Thanks!