fix: preload dense NVFP4 decode kernels before graph capture

[ogwhic]

Summary

• preload each distinct dense NVFP4 decode kernel geometry immediately after model weights are resident
• force Triton's CUDA modules to load before MoE/KV/GDN runtime caches consume the remaining VRAM
• preserve the existing eager warmup and CUDA graph capture behavior

Problem

On an RTX 4070 SUPER 12 GB (Ada, SM89), Qwen3.6-35B-A3B-NVFP4 could remain in Capturing CUDA graphs / warming up for many minutes. GPU telemetry showed 100% SM utilization but only 0-1% memory-controller utilization and about 54 W.

A native py-spy sample located the worker in:

cuModuleLoadData
  Triton loadBinary / _init_handles / run
  freetoken.kernel.triton.nvfp4_linear._gemv
  _linear_impl
  nvfp4_dense_linear_t
  Nvfp4LMHead.forward
  Qwen3_5MoEForCausalLM.forward
  GraphRunner._capture_graphs

This is the ordinary eager model.forward() before the torch.cuda.graph(...) context, not graph capture itself. At that point automatic runtime-cache sizing had reduced free VRAM to roughly 0.7-0.8 GiB. Triton's first-launch module load was therefore occurring at the least favorable point in startup.

The behavior reproduced with both the portable routed-expert backend and vLLM Marlin. Installing vLLM changes the routed MoE implementation, but dense NVFP4 shared experts and the NVFP4 LM head still use FreeToken's Triton kernels.

Fix

Traverse the BaseOP tree after weights load, deduplicate resident Nvfp4DenseLinear and Nvfp4LMHead operations by weight/scale geometry and strides, then launch representative M=1 and maximum captured small-batch kernels. Each launch is synchronized so the CUDA module is actually loaded before runtime caches are allocated.

The preload is skipped when CUDA graphs are disabled.

Validation

Test environment:

• Fedora Linux 44
• NVIDIA RTX 4070 SUPER 12 GB, SM89
• NVIDIA driver 610.43.03
• CUDA toolkit 13.3.73
• Python 3.13.14
• PyTorch 2.11.0+cu130
• Triton 3.6.0
• GCC/G++ 15
• Qwen3.6-35B-A3B-NVFP4
• offloaded MoE with automatic cache sizing
• vLLM 0.14.1 Marlin donor, rebuilt against the active PyTorch/CUDA ABI

With a completely fresh TRITON_CACHE_DIR and graph sizes [1, 2, 4]:

• six distinct dense NVFP4 kernel geometries preloaded in about 2 seconds
• all three CUDA graphs captured in about 9 seconds
• batch size 1 no longer stalled
• a real OpenAI-compatible chat-completion request completed with HTTP 200
• no NVIDIA Xid errors were present

git diff --check and Python AST parsing also pass. The project environment does not currently include pytest, so the existing pytest suite was not run.