Fix NVIDIA DeepSeek V4 MHC warmup coverage

[coltonottley]

Summary

Fix DeepSeek V4 NVIDIA mHC startup warmup coverage.

The existing warmup detector requires decoder layers to expose hc_pre and hc_post, but NVIDIA DeepseekV4DecoderLayer does not expose those methods. The detector therefore returns without warming any NVIDIA mHC kernels. The first-layer broadcast path (mhc_pre_broadcast_tilelang) also has no warmup coverage.

This change:

• exposes NVIDIA hc_pre / hc_post wrappers that use the same arguments and return ordering as inference;
• warms the standalone non-broadcast and first-layer broadcast mHC paths;
• derives one representative token size for every reachable compile-time n_splits key up to max_num_batched_tokens;
• keeps the existing fused post-to-pre path unchanged because its split count is runtime-dynamic and graph capture already covers its two compile variants;
• adds production-helper-based tests for detector gating, exact split-key coverage, broadcast/non-broadcast orchestration, resource bounds, and wrapper behavior.
• fails closed when broadcast weights are missing, non-tensor, or on a different CUDA device.

Root cause

Two upstream changes left independent gaps:

1. The NVIDIA mHC warmup detector checks for hc_pre / hc_post, but the NVIDIA decoder layer does not implement them, so warmup silently no-ops.
2. The later first-layer broadcast mHC path was added without corresponding warmup coverage.

n_splits is a TileLang compile-time parameter. Warming only powers of two or only the broadcast geometry leaves reachable specializations cold. The shared selector in this patch enumerates the exact reachable split keys for each K geometry.

Validation

Source validation on current upstream main:

• focused + adjacent warmup tests: 53 passed, 1 skipped (the skip is the existing host-without-usable-TileLang cross-check);
• Ruff check and format check pass;
• Python compile and git diff --check pass.

Hardware validation used DeepSeek-V4-Flash with TP=2, MTP=2, FP8 KV cache, and 1,000,000-token context on RTX PRO 6000 Blackwell GPUs:

• initial proof on one TP=2 host with 2 GPUs used checksummed read-only source mounts over an immutable parent image;
• cold startup compiled the expected NVIDIA mHC key set on both ranks;
• an 18-turn growing-history/tool workload (2.4M+ prompt tokens) completed 18/18 with zero mHC JIT compilation during inference;
• exact 52K-token cold/store/GPU-reset/CPU-reload regression preserved semantic output and CPU-KV counters;
• the exact immutable successor image passed the same gates on two independent TP=2 endpoints on that host;
• a separately compiled image passed final behavioral and semantic gates on another TP=2 host with 2 RTX PRO 6000 Blackwell GPUs.

Persistent per-source kernel caches reduced same-image restart time substantially, but cache persistence is a deployment optimization, not required for the correctness of this patch.

Scope

This PR is intentionally limited to NVIDIA DeepSeek V4 mHC warmup. Other first-request JIT gaps observed in speculative decoding, sparse-MLA metadata, and expert-token counting are separate work.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[coltonottley]

This is ready for review and upstream CI. DCO passes; the focused and adjacent warmup suites pass 53 tests with one CUDA-only skip; the branch is current with main; and the NVIDIA MHC behavior was validated on TP=2 RTX PRO 6000 Blackwell hardware. I do not have permission to apply the repository ready/verified label, so a maintainer label is needed to admit the normal test workflow.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @coltonottley.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork