[CI Bugfix] Pre-download missing FlashInfer headers in Docker build

[mgoin]

Summary

• NOTE: This is a temporary fix to unblock our CI - the proper fix should be upstream in flashinfer, hopefully fix: avoid re-downloading BMM export headers when flashinfer-cubin is installed flashinfer-ai/flashinfer#2903
• Pre-download FlashInfer TRTLLM BMM headers at Docker build time to fix air-gapped/offline runtime failures
• The flashinfer-cubin package ships headers at the artifact hash path (cubins/b55211623.../include/trtllmGen_bmm_export/), but runtime code (download_trtllm_headers) looks for them at cubins/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/ — a path mismatch that causes network downloads on every startup
• Without this fix, startup takes ~2min extra for serial header downloads, and completely fails in air-gapped environments or when NVIDIA's artifactory returns 403 (as seen with SfLayoutDecl.h)

Root Cause

flashinfer-cubin==0.6.6 includes all 17 BMM headers + checksums.txt, but only at the artifact hash path. The runtime download_trtllm_headers() function (called from flashinfer/jit/moe_utils.py and flashinfer/jit/fused_moe.py) writes to and reads from a different path under flashinfer/trtllm/batched_gemm/. Since get_file() → load_cubin() doesn't find the headers at the expected path, it always falls through to downloading from edge.urm.nvidia.com.

Fix

Call download_trtllm_headers() at Docker build time (when internet is available). This populates the correct runtime path so load_cubin() finds everything locally — zero network access needed at serving time.

Verified locally:

$ FLASHINFER_CUBINS_REPOSITORY="http://localhost:1/" python -c "
from flashinfer.jit import env as jit_env
from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
from flashinfer.artifacts import ArtifactPath, CheckSumHash
download_trtllm_headers('bmm',
    jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
    f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
    ArtifactPath.TRTLLM_GEN_BMM,
    get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM))
print('All BMM headers loaded from local cache - no network needed!')
"
All BMM headers loaded from local cache - no network needed!

Fixes #38110

Test plan

• Docker image builds successfully with this change
• gpt-oss MXFP4 MoE models start without header download attempts
• Air-gapped runtime no longer fails with ConnectTimeoutError on edge.urm.nvidia.com

🤖 Generated with Claude Code

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review to trigger a review and subscribe this PR to future pushes, or @claude review once for a one-time review.

_{Tip: disable this comment in your organization's Code Review settings.}

[gemini-code-assist]

Code Review

This pull request updates the Dockerfile to pre-download FlashInfer TRTLLM BMM headers, ensuring that MoE JIT compilation works in air-gapped environments without requiring internet access at runtime. The review feedback suggests improving the readability and maintainability of the inline Python script by using a heredoc instead of a quoted string with backslashes and semicolons.

[yewentao256]

LGTM, thanks for the work!

[mgoin]

This fixed all the b200 failures I saw in last night's build, so merging