Repository navigation
[None][feat] Mistral Large 4: carry feat/mistral-large-4 onto main - #21693
Draft
yizhang-nv wants to merge 25 commits into
Draft
yizhang-nv wants to merge 25 commits into
yizhang-nv wants to merge 25 commits into
Conversation
4 tasks done
yizhang-nv
force-pushed
the
codex/ml4-mergeback-on-main
branch
from
November 8, 2026 16:42
3a6f070 to
6437405
Compare
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 9, 2026
…d draft KV pool With joint KV cache reuse the KV cache manager V2 scheduler admits every generation request in the target and the draft pool together (KVCacheV2Scheduler._try_allocate_generation). The attention-DP padding dummy was created with a target cache only, so its draft admission found no cache, the target growth was reverted, and the dummy was never scheduled. The idle rank then scheduled an empty batch, _can_queue vetoed the forward on every rank, and the live dummy kept the empty-batch padding from stepping in: the whole attention-DP group stopped making progress as soon as one rank went idle. The empty-batch padding path, which bypasses the scheduler, instead left the draft cache to the draft manager's prepare_resources mirror. That mirror raises when the draft pool is short, which takes down the executor loop; this follows from the code and was not observed in a run. Joint reuse is on for a one-model DFlash or DSpark drafter with its own draft KV pool, KV cache manager V2 and block reuse; attention DP keeps the separate draft pool for these standalone drafters. Found while bringing up Mistral Large 4 with its DSpark drafter: at DEP4 and DEP8 the first request never completed while the other ranks idled, and the same configuration ran to completion with this change. Every padding path now passes the paired draft manager to add_dummy_requests, so the dummy is created in both pools or in neither, and every rollback (a vetoed iteration in _finalize_adp_dummy_allocation, or a spec resource manager without a free slot) frees both. The non-transactional path also rolls back when the spec resource manager has no free slot, as the other two paths already did, instead of propagating NoFreeSlotsError with the KV allocated. Test: test_adp_dummy_holds_paired_draft_kv builds real target and draft KVCacheManagerV2 pools with joint reuse and a KVCacheV2Scheduler, and checks for the plain, transactional and empty-batch padding paths that the dummy holds a cache in both pools, is scheduled, and is released from both on a vetoed iteration. The stub executors of test_py_executor.py, test_benchmark_disagg.py and test_connector_adp.py carry the new attribute (off). Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 9, 2026
…e guided decoder A capturable guided decoder (guided decoding together with one-model speculative decoding: MTP, Eagle3 one-model, the external drafters such as DFlash and DSpark, or SA) turned the whole general warmup off, though it rules out only the warmup's generation requests: the autotuner already runs its context step with the decoder installed. Without the general warmup, the autotuner's max-token step was the engine's first forward. What a first forward creates lazily, including the attention-DP group's NCCL communicator (allocated through cuMem, outside the torch caching allocator, so it cannot reclaim the cache), then came in the middle of the autotuner's exploration with torch's cache full. Any model served with guided decoding and one of those drafters takes this path. It ran memory profiling out of memory while bringing up Mistral Large 4 (NVFP4 DEP4 at 1M context with the DSpark draft and xgrammar): rank 0 ran out of memory in the autotuner's max-token step, and the other ranks failed to allocate NCCL memory through cuMem. Without guided decoding, where the general warmup and its release run ahead of the autotuner, the same configuration profiles at a peak of 238.6 GiB per rank. The general warmup now runs its context requests in that case, followed by the same release as before, so the first forward and the release come ahead of the autotuner, as they do without guided decoding. The extra phase decision is agreed across ranks like can_run_general_warmup, which is already rank-uniform when it is taken. Test: test_guided_decoder_runs_the_context_warmup_before_the_autotuner installs a guided decoder and checks that the general warmup runs with context-only shapes (including the max-token one), followed by the MoE workspace release, the cache release and then the autotuner. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab a2d3872e7b, rebased onto main's warmup after NVIDIA#21343. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 9, 2026
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 9, 2026
FlashMLA's pybind module binds its forward functions with plain m.def, so a call keeps the GIL while it blocks in a kernel launch. A CUDA host function that takes the GIL and sits ahead of those kernels on the stream then deadlocks the rank: the launch waits on the host function, and the host function waits on the GIL. The capturable guided decoder is such a host function (guided decoding beside one-model speculative decoding); its callbacks run on the stream and acquire the GIL in cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp. Under serving load the launch queue fills and an eager step that calls FlashMLA hangs. On main, FlashMLA serves the DSA sparse MLA backend (DeepSeek-V3.2 and the GLM models built on it), DeepSeek-V4, CSA2 and the GLM K-pool backend, so any of these with guided decoding and MTP or another one-model drafter can hit it. Found while bringing up Mistral Large 4 (xgrammar beside the DSpark drafter), where eager context steps deadlocked under load. A FetchContent patch adds pybind11::call_guard<pybind11::gil_scoped_release> to the five bindings. Their arguments and results are tensors and scalars, converted before the GIL is released and after it is reacquired. FlashMLA is a FetchContent dependency (3rdparty/fetch_content.json), so the patch follows the patch_file convention of DeepGEMM and DeepEP rather than the vendored-source lock. Test: test_flash_mla_gil.py queues a GIL-taking host function behind GPU work and fills the launch queue with sparse_prefill_fwd calls in a subprocess; it deadlocks without the patch. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 876a93ccd7. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 9, 2026
…d draft KV pool With joint KV cache reuse the KV cache manager V2 scheduler admits every generation request in the target and the draft pool together (KVCacheV2Scheduler._try_allocate_generation). The attention-DP padding dummy was created with a target cache only, so its draft admission found no cache, the target growth was reverted, and the dummy was never scheduled. The idle rank then scheduled an empty batch, _can_queue vetoed the forward on every rank, and the live dummy kept the empty-batch padding from stepping in: the whole attention-DP group stopped making progress as soon as one rank went idle. The empty-batch padding path, which bypasses the scheduler, instead left the draft cache to the draft manager's prepare_resources mirror. That mirror raises when the draft pool is short, which takes down the executor loop; this follows from the code and was not observed in a run. Joint reuse is on for a one-model DFlash or DSpark drafter with its own draft KV pool, KV cache manager V2 and block reuse; attention DP keeps the separate draft pool for these standalone drafters. Found while bringing up Mistral Large 4 with its DSpark drafter: at DEP4 and DEP8 the first request never completed while the other ranks idled, and the same configuration ran to completion with this change. Every padding path now passes the paired draft manager to add_dummy_requests, so the dummy is created in both pools or in neither. A vetoed iteration already frees the draft pool's cache in _finalize_adp_dummy_allocation; a spec resource manager without a free slot now frees it as well. The non-transactional path also rolls back when the spec resource manager has no free slot, as the other two paths already did, instead of propagating NoFreeSlotsError with the KV allocated. Test: test_adp_dummy_holds_paired_draft_kv builds real target and draft KVCacheManagerV2 pools with joint reuse and a KVCacheV2Scheduler, and checks for the plain, transactional and empty-batch padding paths that the dummy holds a cache in both pools, is scheduled, and is released from both on a vetoed iteration. The stub executors of test_py_executor.py, test_benchmark_disagg.py and test_connector_adp.py carry the new attribute (off). Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 9, 2026
…e guided decoder A capturable guided decoder (guided decoding together with one-model speculative decoding: MTP, Eagle3 one-model, the external drafters such as DFlash and DSpark, or SA) turned the whole general warmup off, though it rules out only the warmup's generation requests: the autotuner already runs its context step with the decoder installed. Without the general warmup, the autotuner's max-token step was the engine's first forward. What a first forward creates lazily, including the attention-DP group's NCCL communicator (allocated through cuMem, outside the torch caching allocator, so it cannot reclaim the cache), then came in the middle of the autotuner's exploration with torch's cache full. Any model served with guided decoding and one of those drafters takes this path. It ran memory profiling out of memory while bringing up Mistral Large 4 (NVFP4 DEP4 at 1M context with the DSpark draft and xgrammar): rank 0 ran out of memory in the autotuner's max-token step, and the other ranks failed to allocate NCCL memory through cuMem. Without guided decoding, where the general warmup and its release run ahead of the autotuner, the same configuration profiles at a peak of 238.6 GiB per rank. The general warmup now runs its context requests in that case, followed by the same release as before, so the first forward and the release come ahead of the autotuner, as they do without guided decoding. The extra phase decision is agreed across ranks like can_run_general_warmup, which is already rank-uniform when it is taken. Test: test_guided_decoder_runs_the_context_warmup_before_the_autotuner installs a guided decoder and checks that the general warmup runs with context-only shapes (including the max-token one), followed by the MoE workspace release, the cache release and then the autotuner. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab a2d3872e7b, rebased onto main's warmup after NVIDIA#21343. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 9, 2026
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 10, 2026
FlashMLA's pybind module binds its forward functions with plain m.def, so a call keeps the GIL while it blocks in a kernel launch. A CUDA host function that takes the GIL and sits ahead of those kernels on the stream then deadlocks the rank: the launch waits on the host function, and the host function waits on the GIL. The capturable guided decoder is such a host function (guided decoding beside one-model speculative decoding); its callbacks run on the stream and acquire the GIL in cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp. Under serving load the launch queue fills and an eager step that calls FlashMLA hangs. On main, FlashMLA serves the DSA sparse MLA backend (DeepSeek-V3.2 and the GLM models built on it), DeepSeek-V4, CSA2 and the GLM K-pool backend, so any of these with guided decoding and MTP or another one-model drafter can hit it. Found while bringing up Mistral Large 4 (xgrammar beside the DSpark drafter), where eager context steps deadlocked under load. A FetchContent patch adds pybind11::call_guard<pybind11::gil_scoped_release> to the five bindings. Their arguments and results are tensors and scalars, converted before the GIL is released and after it is reacquired. FlashMLA is a FetchContent dependency (3rdparty/fetch_content.json), so the patch follows the patch_file convention of DeepGEMM and DeepEP rather than the vendored-source lock. Test: test_flash_mla_gil.py queues a GIL-taking host function behind GPU work and fills the launch queue with sparse_prefill_fwd calls in a subprocess; it deadlocks without the patch. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 876a93ccd7. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 10, 2026
…d draft KV pool With joint KV cache reuse the KV cache manager V2 scheduler admits every generation request in the target and the draft pool together (KVCacheV2Scheduler._try_allocate_generation). The attention-DP padding dummy was created with a target cache only, so its draft admission found no cache, the target growth was reverted, and the dummy was never scheduled. The idle rank then scheduled an empty batch, _can_queue vetoed the forward on every rank, and the live dummy kept the empty-batch padding from stepping in: the whole attention-DP group stopped making progress as soon as one rank went idle. The empty-batch padding path, which bypasses the scheduler, instead left the draft cache to the draft manager's prepare_resources mirror. That mirror raises when the draft pool is short, which takes down the executor loop; this follows from the code and was not observed in a run. Joint reuse is on for a one-model DFlash or DSpark drafter with its own draft KV pool, KV cache manager V2 and block reuse; attention DP keeps the separate draft pool for these standalone drafters. Found while bringing up Mistral Large 4 with its DSpark drafter: at DEP4 and DEP8 the first request never completed while the other ranks idled, and the same configuration ran to completion with this change. Every padding path now passes the paired draft manager to add_dummy_requests, so the dummy is created in both pools or in neither. A vetoed iteration already frees the draft pool's cache in _finalize_adp_dummy_allocation; a spec resource manager without a free slot now frees it as well. The non-transactional path also rolls back when the spec resource manager has no free slot, as the other two paths already did, instead of propagating NoFreeSlotsError with the KV allocated. add_dummy_requests reserves the dummy's whole forward in both pools, and the empty-batch dummy joins the batch after the scheduler, so it carries no scheduler allocation record. The draft manager's prepare_resources therefore takes a paired pool's existing dummy cache as already sized instead of growing it by another generation step: with a draft length of 8 that step would take the 24-token cache to 33 tokens, a second 32-token page, after the ranks have agreed to run the batch. An unpaired draft pool is not handed to add_dummy_requests; the mirror the scheduler admits for its dummy is empty and still grows in prepare_resources. Tests: test_adp_dummy_holds_paired_draft_kv builds real target and draft KVCacheManagerV2 pools with joint reuse and a KVCacheV2Scheduler, and checks for the plain, transactional and empty-batch padding paths that the dummy holds a cache in both pools, is scheduled, keeps its capacity through prepare_resources, and is released from both on a vetoed iteration. test_adp_dummy_grows_unpaired_draft_mirror runs the same paths over unpaired pools for two queued iterations and checks that the mirror grows by one generation step and takes the previous dummy's deferred update. test_draft_mirror_preserves_preallocated_dummy_capacity (CPU) pins which mirrors grow. The stub executors of test_py_executor.py, test_benchmark_disagg.py and test_connector_adp.py carry the new attribute (off). Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 10, 2026
…e guided decoder A capturable guided decoder (guided decoding together with one-model speculative decoding: MTP, Eagle3 one-model, the external drafters such as DFlash and DSpark, or SA) turned the whole general warmup off, though it rules out only the warmup's generation requests: the autotuner already runs its context step with the decoder installed. Without the general warmup, the autotuner's max-token step was the engine's first forward. What a first forward creates lazily, including the attention-DP group's NCCL communicator (allocated through cuMem, outside the torch caching allocator, so it cannot reclaim the cache), then came in the middle of the autotuner's exploration with torch's cache full. Any model served with guided decoding and one of those drafters takes this path. It ran memory profiling out of memory while bringing up Mistral Large 4 (NVFP4 DEP4 at 1M context with the DSpark draft and xgrammar): rank 0 ran out of memory in the autotuner's max-token step, and the other ranks failed to allocate NCCL memory through cuMem. Without guided decoding, where the general warmup and its release run ahead of the autotuner, the same configuration profiles at a peak of 238.6 GiB per rank. The general warmup now runs its context requests in that case, followed by the same release as before, so the first forward and the release come ahead of the autotuner, as they do without guided decoding. The extra phase decision is agreed across ranks like can_run_general_warmup, which is already rank-uniform when it is taken. The general warmup skips a shape that runs a single rank out of memory and goes on to a smaller one. With the decoder installed, input preparation has already queued that shape's batch for the decoder's host functions, so the abandoned batch would stay at the head of the queue and every later forward, the first real guided request included, would be served the batch before its own. The recovery now waits for the decoder's stream, so that no host function of the failed forward is still pending, and drops the queued batches and both request snapshots (CapturableGuidedDecoder.reset_warmup_state). A failed reset stops the warmup, and a distributed forward still re-raises the out-of-memory error. Tests: test_guided_decoder_runs_the_context_warmup_before_the_autotuner installs a guided decoder and checks that the general warmup runs with context-only shapes (including the max-token one), followed by the MoE workspace release, the cache release and then the autotuner. test_general_warmup_oom_restores_guided_batch_queue fails a shape before and after its host function is submitted and checks that the next guided request is served its own batch. test_general_warmup_oom_policy and test_general_warmup_aborts_if_guided_state_cannot_be_reset cover when the reset runs and that its failure stops the warmup. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab a2d3872e7b, rebased onto main's warmup after NVIDIA#21343. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 10, 2026
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 11, 2026
FlashMLA's pybind module binds its forward functions with plain m.def, so a call keeps the GIL while it blocks in a kernel launch. A CUDA host function that takes the GIL and sits ahead of those kernels on the stream then deadlocks the rank: the launch waits on the host function, and the host function waits on the GIL. The capturable guided decoder is such a host function (guided decoding beside one-model speculative decoding); its callbacks run on the stream and acquire the GIL in cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp. Under serving load the launch queue fills and an eager step that calls FlashMLA hangs. On main, FlashMLA serves the DSA sparse MLA backend (DeepSeek-V3.2 and the GLM models built on it), DeepSeek-V4, CSA2 and the GLM K-pool backend, so any of these with guided decoding and MTP or another one-model drafter can hit it. Found while bringing up Mistral Large 4 (xgrammar beside the DSpark drafter), where eager context steps deadlocked under load. A FetchContent patch adds pybind11::call_guard<pybind11::gil_scoped_release> to the five bindings. Their arguments and results are tensors and scalars, converted before the GIL is released and after it is reacquired. FlashMLA is a FetchContent dependency (3rdparty/fetch_content.json), so the patch follows the patch_file convention of DeepGEMM and DeepEP rather than the vendored-source lock. Test: test_flash_mla_gil.py queues a GIL-taking host function behind GPU work and fills the launch queue with sparse_prefill_fwd calls in a subprocess; it deadlocks without the patch. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 876a93ccd7. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 11, 2026
…d draft KV pool With joint KV cache reuse the KV cache manager V2 scheduler admits every generation request in the target and the draft pool together (KVCacheV2Scheduler._try_allocate_generation). The attention-DP padding dummy was created with a target cache only, so its draft admission found no cache, the target growth was reverted, and the dummy was never scheduled. The idle rank then scheduled an empty batch, _can_queue vetoed the forward on every rank, and the live dummy kept the empty-batch padding from stepping in: the whole attention-DP group stopped making progress as soon as one rank went idle. The empty-batch padding path, which bypasses the scheduler, instead left the draft cache to the draft manager's prepare_resources mirror. That mirror raises when the draft pool is short, which takes down the executor loop; this follows from the code and was not observed in a run. Joint reuse is on for a one-model DFlash or DSpark drafter with its own draft KV pool, KV cache manager V2 and block reuse; attention DP keeps the separate draft pool for these standalone drafters. Found while bringing up Mistral Large 4 with its DSpark drafter: at DEP4 and DEP8 the first request never completed while the other ranks idled, and the same configuration ran to completion with this change. Every padding path now passes the paired draft manager to add_dummy_requests, so the dummy is created in both pools or in neither. A vetoed iteration already frees the draft pool's cache in _finalize_adp_dummy_allocation; a spec resource manager without a free slot now frees it as well. The non-transactional path also rolls back when the spec resource manager has no free slot, as the other two paths already did, instead of propagating NoFreeSlotsError with the KV allocated. add_dummy_requests reserves the dummy's whole forward in both pools, and the empty-batch dummy joins the batch after the scheduler, so it carries no scheduler allocation record. The draft manager's prepare_resources therefore takes a paired pool's existing dummy cache as already sized instead of growing it by another generation step: with a draft length of 8 that step would take the 24-token cache to 33 tokens, a second 32-token page, after the ranks have agreed to run the batch. An unpaired draft pool is not handed to add_dummy_requests; the mirror the scheduler admits for its dummy is empty and still grows in prepare_resources. Tests: test_adp_dummy_holds_paired_draft_kv builds real target and draft KVCacheManagerV2 pools with joint reuse and a KVCacheV2Scheduler, and checks for the plain, transactional and empty-batch padding paths that the dummy holds a cache in both pools, is scheduled, keeps its capacity through prepare_resources, and is released from both on a vetoed iteration. test_adp_dummy_grows_unpaired_draft_mirror runs the same paths over unpaired pools for two queued iterations and checks that the mirror grows by one generation step and takes the previous dummy's deferred update. test_draft_mirror_preserves_preallocated_dummy_capacity (CPU) pins which mirrors grow. The stub executors of test_py_executor.py, test_benchmark_disagg.py and test_connector_adp.py carry the new attribute (off). Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 11, 2026
…e guided decoder A capturable guided decoder (guided decoding together with one-model speculative decoding: MTP, Eagle3 one-model, the external drafters such as DFlash and DSpark, or SA) turned the whole general warmup off, though it rules out only the warmup's generation requests: the autotuner already runs its context step with the decoder installed. Without the general warmup, the autotuner's max-token step was the engine's first forward. What a first forward creates lazily, including the attention-DP group's NCCL communicator (allocated through cuMem, outside the torch caching allocator, so it cannot reclaim the cache), then came in the middle of the autotuner's exploration with torch's cache full. Any model served with guided decoding and one of those drafters takes this path. It ran memory profiling out of memory while bringing up Mistral Large 4 (NVFP4 DEP4 at 1M context with the DSpark draft and xgrammar): rank 0 ran out of memory in the autotuner's max-token step, and the other ranks failed to allocate NCCL memory through cuMem. Without guided decoding, where the general warmup and its release run ahead of the autotuner, the same configuration profiles at a peak of 238.6 GiB per rank. The general warmup now runs its context requests in that case, followed by the same release as before, so the first forward and the release come ahead of the autotuner, as they do without guided decoding. The extra phase decision is agreed across ranks like can_run_general_warmup, which is already rank-uniform when it is taken. The general warmup skips a shape that runs a single rank out of memory and goes on to a smaller one. With the decoder installed, input preparation has already queued that shape's batch for the decoder's host functions, so the abandoned batch would stay at the head of the queue and every later forward, the first real guided request included, would be served the batch before its own. The recovery now waits for the decoder's stream, so that no host function of the failed forward is still pending, and drops the queued batches and both request snapshots (CapturableGuidedDecoder.reset_warmup_state). A failed reset stops the warmup, and a distributed forward still re-raises the out-of-memory error. Tests: test_guided_decoder_runs_the_context_warmup_before_the_autotuner installs a guided decoder and checks that the general warmup runs with context-only shapes (including the max-token one), followed by the MoE workspace release, the cache release and then the autotuner. test_general_warmup_oom_restores_guided_batch_queue fails a shape before and after its host function is submitted and checks that the next guided request is served its own batch. test_general_warmup_oom_policy and test_general_warmup_aborts_if_guided_state_cannot_be_reset cover when the reset runs and that its failure stops the warmup. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab a2d3872e7b, rebased onto main's warmup after NVIDIA#21343. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 11, 2026
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
hanjingtian
pushed a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 11, 2026
FlashMLA's pybind module binds its forward functions with plain m.def, so a call keeps the GIL while it blocks in a kernel launch. A CUDA host function that takes the GIL and sits ahead of those kernels on the stream then deadlocks the rank: the launch waits on the host function, and the host function waits on the GIL. The capturable guided decoder is such a host function (guided decoding beside one-model speculative decoding); its callbacks run on the stream and acquire the GIL in cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp. Under serving load the launch queue fills and an eager step that calls FlashMLA hangs. On main, FlashMLA serves the DSA sparse MLA backend (DeepSeek-V3.2 and the GLM models built on it), DeepSeek-V4, CSA2 and the GLM K-pool backend, so any of these with guided decoding and MTP or another one-model drafter can hit it. Found while bringing up Mistral Large 4 (xgrammar beside the DSpark drafter), where eager context steps deadlocked under load. A FetchContent patch adds pybind11::call_guard<pybind11::gil_scoped_release> to the five bindings. Their arguments and results are tensors and scalars, converted before the GIL is released and after it is reacquired. FlashMLA is a FetchContent dependency (3rdparty/fetch_content.json), so the patch follows the patch_file convention of DeepGEMM and DeepEP rather than the vendored-source lock. Test: test_flash_mla_gil.py queues a GIL-taking host function behind GPU work and fills the launch queue with sparse_prefill_fwd calls in a subprocess; it deadlocks without the patch. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 876a93ccd7. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
hanjingtian
pushed a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 11, 2026
…d draft KV pool With joint KV cache reuse the KV cache manager V2 scheduler admits every generation request in the target and the draft pool together (KVCacheV2Scheduler._try_allocate_generation). The attention-DP padding dummy was created with a target cache only, so its draft admission found no cache, the target growth was reverted, and the dummy was never scheduled. The idle rank then scheduled an empty batch, _can_queue vetoed the forward on every rank, and the live dummy kept the empty-batch padding from stepping in: the whole attention-DP group stopped making progress as soon as one rank went idle. The empty-batch padding path, which bypasses the scheduler, instead left the draft cache to the draft manager's prepare_resources mirror. That mirror raises when the draft pool is short, which takes down the executor loop; this follows from the code and was not observed in a run. Joint reuse is on for a one-model DFlash or DSpark drafter with its own draft KV pool, KV cache manager V2 and block reuse; attention DP keeps the separate draft pool for these standalone drafters. Found while bringing up Mistral Large 4 with its DSpark drafter: at DEP4 and DEP8 the first request never completed while the other ranks idled, and the same configuration ran to completion with this change. Every padding path now passes the paired draft manager to add_dummy_requests, so the dummy is created in both pools or in neither. A vetoed iteration already frees the draft pool's cache in _finalize_adp_dummy_allocation; a spec resource manager without a free slot now frees it as well. The non-transactional path also rolls back when the spec resource manager has no free slot, as the other two paths already did, instead of propagating NoFreeSlotsError with the KV allocated. add_dummy_requests reserves the dummy's whole forward in both pools, and the empty-batch dummy joins the batch after the scheduler, so it carries no scheduler allocation record. The draft manager's prepare_resources therefore takes a paired pool's existing dummy cache as already sized instead of growing it by another generation step: with a draft length of 8 that step would take the 24-token cache to 33 tokens, a second 32-token page, after the ranks have agreed to run the batch. An unpaired draft pool is not handed to add_dummy_requests; the mirror the scheduler admits for its dummy is empty and still grows in prepare_resources. Tests: test_adp_dummy_holds_paired_draft_kv builds real target and draft KVCacheManagerV2 pools with joint reuse and a KVCacheV2Scheduler, and checks for the plain, transactional and empty-batch padding paths that the dummy holds a cache in both pools, is scheduled, keeps its capacity through prepare_resources, and is released from both on a vetoed iteration. test_adp_dummy_grows_unpaired_draft_mirror runs the same paths over unpaired pools for two queued iterations and checks that the mirror grows by one generation step and takes the previous dummy's deferred update. test_draft_mirror_preserves_preallocated_dummy_capacity (CPU) pins which mirrors grow. The stub executors of test_py_executor.py, test_benchmark_disagg.py and test_connector_adp.py carry the new attribute (off). Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
hanjingtian
pushed a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 11, 2026
…e guided decoder A capturable guided decoder (guided decoding together with one-model speculative decoding: MTP, Eagle3 one-model, the external drafters such as DFlash and DSpark, or SA) turned the whole general warmup off, though it rules out only the warmup's generation requests: the autotuner already runs its context step with the decoder installed. Without the general warmup, the autotuner's max-token step was the engine's first forward. What a first forward creates lazily, including the attention-DP group's NCCL communicator (allocated through cuMem, outside the torch caching allocator, so it cannot reclaim the cache), then came in the middle of the autotuner's exploration with torch's cache full. Any model served with guided decoding and one of those drafters takes this path. It ran memory profiling out of memory while bringing up Mistral Large 4 (NVFP4 DEP4 at 1M context with the DSpark draft and xgrammar): rank 0 ran out of memory in the autotuner's max-token step, and the other ranks failed to allocate NCCL memory through cuMem. Without guided decoding, where the general warmup and its release run ahead of the autotuner, the same configuration profiles at a peak of 238.6 GiB per rank. The general warmup now runs its context requests in that case, followed by the same release as before, so the first forward and the release come ahead of the autotuner, as they do without guided decoding. The extra phase decision is agreed across ranks like can_run_general_warmup, which is already rank-uniform when it is taken. The general warmup skips a shape that runs a single rank out of memory and goes on to a smaller one. With the decoder installed, input preparation has already queued that shape's batch for the decoder's host functions, so the abandoned batch would stay at the head of the queue and every later forward, the first real guided request included, would be served the batch before its own. The recovery now waits for the decoder's stream, so that no host function of the failed forward is still pending, and drops the queued batches and both request snapshots (CapturableGuidedDecoder.reset_warmup_state). A failed reset stops the warmup, and a distributed forward still re-raises the out-of-memory error. Tests: test_guided_decoder_runs_the_context_warmup_before_the_autotuner installs a guided decoder and checks that the general warmup runs with context-only shapes (including the max-token one), followed by the MoE workspace release, the cache release and then the autotuner. test_general_warmup_oom_restores_guided_batch_queue fails a shape before and after its host function is submitted and checks that the next guided request is served its own batch. test_general_warmup_oom_policy and test_general_warmup_aborts_if_guided_state_cannot_be_reset cover when the reset runs and that its failure stops the warmup. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab a2d3872e7b, rebased onto main's warmup after NVIDIA#21343. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
hanjingtian
pushed a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 11, 2026
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 12, 2026
FlashMLA's pybind module binds its forward functions with plain m.def, so a call keeps the GIL while it blocks in a kernel launch. A CUDA host function that takes the GIL and sits ahead of those kernels on the stream then deadlocks the rank: the launch waits on the host function, and the host function waits on the GIL. The capturable guided decoder is such a host function (guided decoding beside one-model speculative decoding); its callbacks run on the stream and acquire the GIL in cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp. Under serving load the launch queue fills and an eager step that calls FlashMLA hangs. On main, FlashMLA serves the DSA sparse MLA backend (DeepSeek-V3.2 and the GLM models built on it), DeepSeek-V4, CSA2 and the GLM K-pool backend, so any of these with guided decoding and MTP or another one-model drafter can hit it. Found while bringing up Mistral Large 4 (xgrammar beside the DSpark drafter), where eager context steps deadlocked under load. A FetchContent patch adds pybind11::call_guard<pybind11::gil_scoped_release> to the five bindings. Their arguments and results are tensors and scalars, converted before the GIL is released and after it is reacquired. FlashMLA is a FetchContent dependency (3rdparty/fetch_content.json), so the patch follows the patch_file convention of DeepGEMM and DeepEP rather than the vendored-source lock. Test: test_flash_mla_gil.py queues a GIL-taking host function behind GPU work and fills the launch queue with sparse_prefill_fwd calls in a subprocess; it deadlocks without the patch. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 876a93ccd7. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 12, 2026
…d draft KV pool With joint KV cache reuse the KV cache manager V2 scheduler admits every generation request in the target and the draft pool together (KVCacheV2Scheduler._try_allocate_generation). The attention-DP padding dummy was created with a target cache only, so its draft admission found no cache, the target growth was reverted, and the dummy was never scheduled. The idle rank then scheduled an empty batch, _can_queue vetoed the forward on every rank, and the live dummy kept the empty-batch padding from stepping in: the whole attention-DP group stopped making progress as soon as one rank went idle. The empty-batch padding path, which bypasses the scheduler, instead left the draft cache to the draft manager's prepare_resources mirror. That mirror raises when the draft pool is short, which takes down the executor loop; this follows from the code and was not observed in a run. Joint reuse is on for a one-model DFlash or DSpark drafter with its own draft KV pool, KV cache manager V2 and block reuse; attention DP keeps the separate draft pool for these standalone drafters. Found while bringing up Mistral Large 4 with its DSpark drafter: at DEP4 and DEP8 the first request never completed while the other ranks idled, and the same configuration ran to completion with this change. Every padding path now passes the paired draft manager to add_dummy_requests, so the dummy is created in both pools or in neither. A vetoed iteration already frees the draft pool's cache in _finalize_adp_dummy_allocation; a spec resource manager without a free slot now frees it as well. The non-transactional path also rolls back when the spec resource manager has no free slot, as the other two paths already did, instead of propagating NoFreeSlotsError with the KV allocated. add_dummy_requests reserves the dummy's whole forward in both pools, and the empty-batch dummy joins the batch after the scheduler, so it carries no scheduler allocation record. The draft manager's prepare_resources therefore takes a paired pool's existing dummy cache as already sized instead of growing it by another generation step: with a draft length of 8 that step would take the 24-token cache to 33 tokens, a second 32-token page, after the ranks have agreed to run the batch. An unpaired draft pool is not handed to add_dummy_requests; the mirror the scheduler admits for its dummy is empty and still grows in prepare_resources. Tests: test_adp_dummy_holds_paired_draft_kv builds real target and draft KVCacheManagerV2 pools with joint reuse and a KVCacheV2Scheduler, and checks for the plain, transactional and empty-batch padding paths that the dummy holds a cache in both pools, is scheduled, keeps its capacity through prepare_resources, and is released from both on a vetoed iteration. test_adp_dummy_grows_unpaired_draft_mirror runs the same paths over unpaired pools for two queued iterations and checks that the mirror grows by one generation step and takes the previous dummy's deferred update. test_draft_mirror_preserves_preallocated_dummy_capacity (CPU) pins which mirrors grow. The stub executors of test_py_executor.py, test_benchmark_disagg.py and test_connector_adp.py carry the new attribute (off). Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 12, 2026
…e guided decoder A capturable guided decoder (guided decoding together with one-model speculative decoding: MTP, Eagle3 one-model, the external drafters such as DFlash and DSpark, or SA) turned the whole general warmup off, though it rules out only the warmup's generation requests: the autotuner already runs its context step with the decoder installed. Without the general warmup, the autotuner's max-token step was the engine's first forward. What a first forward creates lazily, including the attention-DP group's NCCL communicator (allocated through cuMem, outside the torch caching allocator, so it cannot reclaim the cache), then came in the middle of the autotuner's exploration with torch's cache full. Any model served with guided decoding and one of those drafters takes this path. It ran memory profiling out of memory while bringing up Mistral Large 4 (NVFP4 DEP4 at 1M context with the DSpark draft and xgrammar): rank 0 ran out of memory in the autotuner's max-token step, and the other ranks failed to allocate NCCL memory through cuMem. Without guided decoding, where the general warmup and its release run ahead of the autotuner, the same configuration profiles at a peak of 238.6 GiB per rank. The general warmup now runs its context requests in that case, followed by the same release as before, so the first forward and the release come ahead of the autotuner, as they do without guided decoding. The extra phase decision is agreed across ranks like can_run_general_warmup, which is already rank-uniform when it is taken. The general warmup skips a shape that runs a single rank out of memory and goes on to a smaller one. With the decoder installed, input preparation has already queued that shape's batch for the decoder's host functions, so the abandoned batch would stay at the head of the queue and every later forward, the first real guided request included, would be served the batch before its own. The recovery now waits for the decoder's stream, so that no host function of the failed forward is still pending, and drops the queued batches and both request snapshots (CapturableGuidedDecoder.reset_warmup_state). A failed reset stops the warmup, and a distributed forward still re-raises the out-of-memory error. Tests: test_guided_decoder_runs_the_context_warmup_before_the_autotuner installs a guided decoder and checks that the general warmup runs with context-only shapes (including the max-token one), followed by the MoE workspace release, the cache release and then the autotuner. test_general_warmup_oom_restores_guided_batch_queue fails a shape before and after its host function is submitted and checks that the next guided request is served its own batch. test_general_warmup_oom_policy and test_general_warmup_aborts_if_guided_state_cannot_be_reset cover when the reset runs and that its failure stops the warmup. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab a2d3872e7b, rebased onto main's warmup after NVIDIA#21343. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
added a commit
to yizhang-nv/TensorRT-LLM
that referenced
this pull request
Nov 12, 2026
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Squash of the Mistral Large 4 work on feat/mistral-large-4 (3fb4be0), applied on main 19c1f46: - fece87e [None][feat] Add Mistral Large 4 support - 9df1a00 [None][fix] Keep requirements installable on Python 3.10 - 5f00308 [None][doc] Mistral Large 4: install from the release container or wheel (NVIDIA#21425) - bed3750 [None][refactor] Mistral Large 4: own the checkpoint weight loader and mapper (NVIDIA#21548) - 3fb4be0 [None][chore] Mistral Large 4: remove defensive checks and leftovers of removed paths (NVIDIA#21616) Not carried: the six infra commits already on main (NVIDIA#21133, NVIDIA#19557, (cf7ae79, 1.4.0a0). Adapted to main: - modeling_v2 routing: the family registers its route with @routes("MistralLarge4ForCausalLM") and models/__init__.py imports it; route takes (model_config, llm_args, compute_capability). The routing tests use main's per-family sm_103 fixture; the engine-config tests record the route through registered_routes and resolve on one rank. - thop_attention call sites use main's argument names (workspace, spec_decoding_position_offsets); top_k kernels come from _torch/kernels/blackwell/top_k. - Catalog: the targets' cells move into each op test's TARGET_CELLS (comm ops keep them in CELLS, as main's comm ops do); bmm_out keeps main's op. - DSpark/DFlash: the native drafter's hooks sit on main's refactored worker (block placement, vocab-parallel drafting, Markov TP exchange); the recurrent Markov head takes its hidden states directly and the fused paged-latent path is not taken for drafters with their own block embedding. - Per-rank spec-decoding stats travel in main's versioned rank-state frame (version -2 adds a spec snapshot); KVCacheManagerV2 keeps has_buffer next to main's mixed-precision get_buffers. - Serving: the chat path runs on main's prompt pipeline and ServingExtension hooks; Mistral Large 4 adds render_needs_media, forced_tool_call_decoding ("native" required-tool mode) and validate_chat_request. `store` follows main (accepted and ignored). - Warmup/ADP dummies keep the joint-KV-reuse draft allocation on main's dummy request arguments. - Tests follow main's engine/runner refactor (ModelEngine, filter_prefill_capture_num_tokens, removed validate_feature_combination), main's scoped reference precision (true_fp32_matmul) and scoped engine environments (scoped_env_overrides). - Dependencies: flashinfer-python 0.7.2 and xgrammar 0.2.8 from main; llguidance 1.7.6 and mistral-common>=1.13.0 (Python >= 3.11) from the feature branch. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
The squash's three-way merge kept both sides' copy of EagleDecodingConfig.spec_dec_variant in the serve config snapshots. Regenerated with scripts/generate_trtllm_serve_schemas.py; the result is identical to main's snapshots, so the feature branch adds no serve schema change. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…only l0_cpu lists unittest/test_ruler_accuracy.py, but the CPU stage runs with -m cpu_only and its conftest collects only files that carry the marker, so the file never ran there. It needs no device; mark it the way its siblings test_mmmu_accuracy.py and test_multimodal_accuracy.py are. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…azy model zoo main imports the PyTorch model zoo lazily: inputs/registry.py finds the module that registers a multimodal model type's placeholders through MULTIMODAL_MODEL_TYPE_TO_MODULE. modeling_mistral.py registers mistral_large_4's placeholder metadata, so the type needs its entry, next to mistral_large_3. test_lazy_model_zoo.py checks the table against the registrations. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…equired or named calls The chat route called ServingExtension.forced_tool_call_decoding for every request, reading the server's guided_decoding_backend each time. The hook only applies to a required or named tool_choice, so call it only then; other requests take main's path unchanged (the disaggregated service tests drive it with a generator whose args carry no guided decoding backend). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…int KV reuse flag _pad_attention_dp_dummy_request and the dummy's teardown now free and revert the draft KV pool when joint KV cache reuse is on, reading enable_joint_kv_cache_reuse, which PyExecutor always sets. The stub executors of test_benchmark_disagg.py and test_connector_adp.py mirror the attributes that path reads, so they carry it too (off). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…package The family's catalog op tests under modeling_v2/attention share a basename with a target's own test (test_modeling_v2_convert_req_index_to_global.py under glm_5_2_nvfp4__sm_107__tp16/dsa). Collected in one session, as l0_b300 collects unittest/_torch/modeling_v2, pytest's default import mode loads both under one module name and refuses the second. Mark the directory a package, as gemm/ and moe/ already are; its two tests that read the directory's _ds_mla_pool helper import it relatively. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…ger interfaces in tests - add_dummy_requests no longer takes prepare_resource: it always allocates, which is what the family's catalog op tests asked for. Drop the argument, and from the dep4 engine-config tests' expected padding call. - test_long_spec_rejection_rewinds_real_cache: KVCacheManagerV2's generation update now commits reusable blocks; the bare manager runs with block reuse off. - test_adp_dummy_joint_draft_admission_and_rollback: a single-rank attention-DP group skips dummy padding, so the executor declares two ranks. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…s with D3 and main These cpu_only tests need the checkpoint, so they run only where it is present (GB300); this brings them in line with what the engine now does. - The routed quantization_config is the export's block plus the FP8 KV-cache algorithm _adapt_mistral_large4_config records from fp8_kv_cache; the anchor-key and read-intact tests expect that key, and every per-layer QuantConfig carries FP8 under kv_cache dtype auto. - A single-rank attention-DP group skips dummy padding, so the idle-rank dummy tests give their executors the mapping's tp_size. - FP4BlockScaleMoERunner.get_tuning_config is cached on decode_tuning too; the tp4 autotuner test asks for it with the argument the runner passes. - Per-item multimodal embedding lengths are item metadata (output_embedding_lengths), no longer a multimodal_data entry. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…amily The catalog holds operations that more than one family uses. Twenty-five entries and six kernel modules were used only by Mistral Large 4, so they move under models/mistral_large4: - ops/<category>/: activation (gelu_erf, tanh_softcap), attention (sparse MLA decode/prefill, indexer, top-k, ssmax, DS-MLA cache ops), comm (cute_dsl_allreduce_add_rmsnorm), gemm (fp8_prequantized_swap_ab_gemm), moe (renorm_moe_routing_op), norm (layer_norm), quantization (fp8_quantize_1x128_packed_ue8m0) and vision (pixtral_rope_2d). - kernels/: the CuTe DSL kernels behind those entries and the Lamport all-reduce helpers. Entries no target uses are removed with their tests: cute_dsl_indexer_topk_prefill, fp8_block_scale_moe_runner and fp8_quantize_1x128. The GVR top-k prefill test no longer compares against the removed radix prefill entry; it checks the reference only. Tests move to tests/unittest/_torch/modeling_v2/models/mistral_large4/ops/. That tree is a package, so its basenames may repeat a target's own tests; the shared catalog test directories go back to main's layout without package markers. l0_gb300_multi_gpus runs the moved all-reduce op matrix, and l0_b300 leaves it out as it does the shared comm tests. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
… out On main, bmm_out's is_valid checks that a, b and out are 3-D before it checks out's shape, so the 2-D out[0] case is refused as "a, b, out must be 3D". The other wrong-shaped outs still get the [B, M, N] message. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…ntly on main indexer_topk_decode past its launch ceiling fails at configuration and leaves indices untouched. The guard test expected torch's next launch to raise that error, which holds only when the op shares torch's CUDA runtime. On main, libth_common links the runtime statically (libtensorrt_llm brings in cudart_static ahead of torch's libcudart.so), so the error never reaches torch. The test now checks that the launch did not run and that nothing reports it; the entry's note says why. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…ers call The Mistral Large 4 targets run main's stock DSparkWorker; their drafter supplies its model math through optional draft-model methods. State them as such: - DFlashWorker's docstring lists project_target_hidden, embed_draft_block and dflash_forward_with_metadata, and DSparkWorker's the recurrent Markov head (has_recurrent_markov_head). - dflash_draft_block_embedding and dflash_block_forward hold the two dispatches that were inline in prepare_1st_drafter_inputs and _forward_impl, so they can be tested without a worker. - DSparkWorker declines vocab-shard drafting for a recurrent Markov head: the shard chain computes the vanilla head only. No in-tree target binds a vocab-sharded head with such a drafter today. - The worker test for the recurrent head moves from the Mistral Large 4 tests to tests/unittest/_torch/speculative/hw_agnostic, next to new tests for the other dispatches (l0_cpu). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…soft-capped The DFlash worker on current main (NVIDIA#21367) passes the drafter's block logits through draft_model.apply_draft_logit_transform after logits_processor. The Mistral Large 4 DraftModel soft-caps the shared-head logits in its own logits_processor and had no such method, so DSpark decoding would raise AttributeError once this branch moves to that main. Each target's DraftModel now returns the logits unchanged there. It declares no logit_multiplier or logit_softcap, so the cap is applied once and the worker's plain-processor check does not apply to it. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…t tensors on main main's image processor now asks the chat template for plain lists in every format (NVIDIA#20843) and batches images without padding away their sizes, so return_tensors no longer tells the formats apart. The scope test keeps what does: the other Mistral-native formats append images after the text and keep a literal [IMG] marker. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…cessor like the server The Mistral Large 4 NVFP4 checkpoint ships Hugging Face frontend files beside its Mistral-native ones, including processor_config.json naming PixtralProcessor and chat_template.jinja. trtllm-serve loads no HF processor for the Mistral formats, mistral_large_4 included, but RenderResources.load skipped it only for mistral and mistral_large_3. For a Mistral Large 4 checkpoint the standalone renderer therefore loaded a PixtralProcessor, and the fingerprint took in its chat template: the renderer's fingerprint did not match the server's. RenderResources.load now skips the processor for the same formats as the server, compared through CheckpointFormat. The mistral and mistral_large_3 formats keep their behavior. Test: test_native_frontend.py loads the checkpoint named by MISTRAL_LARGE4_CHECKPOINT_DIR through RenderResources.load and expects no processor and the mistral-common tokenizer (it fails before this change). On the checkpoint's metadata the renderer's fingerprint now equals the server's (processor None). Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…n the architecture allowlist supported-models.md gains the MistralLarge4ForCausalLM row, its rows in the language and multimodal feature matrices, and footnotes on the GB300 NVFP4 Mistral-format scope and the validated disaggregated setup. The feature cells follow the deployment guide and the accuracy tests; the logits post processor and EPD disaggregated serving are untested. MistralLarge4ForCausalLM is the architecture name TensorRT LLM assigns when it reads the checkpoint's params.json. Now that the supported-models table documents it, the usage allowlist reports it by name. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…p, tp)
Each target serves its layout on four ranks of one node and on eight over
two nodes of four (routing admits world_size 4 or 8), so names that carry
a rank count, tep4 / dep4 / tp4, named one of the two deployments a target
serves. The <parallel> segment now names the layout:
- mistral_large_4_nvfp4__sm_103__{tep4,dep4,tp4} become __{tep,dep,tp},
in the family and in its unit tests, and the target classes become
ModelingV2MistralLarge4Nvfp4Sm103{Tep,Dep,Tp} (with their cores and
weight tables).
- Routing names the layouts tep, dep and tp in its predicates and its
refusal; the routing tests follow, and their names say "the tep target"
where they meant the target.
- Where a comment named a deployment as a target ("the tp4 target at the
1M envelope"), it now says deployment; tep4, dep8 and the rest remain
the names of the deployments.
No behavior changes: the routes, modules and weight tables are the same.
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…yments as parameters
modeling_v2's claims require that a family's accuracy file holds exactly
one TestModelingV2<Target> class per routed target: CBTS reaches a target's
gate by that name. The Mistral Large 4 file had ten more, for deployments
of the same three targets: the eight-rank classes (Tep8, Dep8, Tp8), the
six 1M-context classes and the DSpark / llguidance class. The claim failed,
and a change to a target never selected those gates.
Each target now has one gate, TestModelingV2MistralLarge4Nvfp4Sm103{Tep,
Dep,Tp}, on a shared base. pytest_generate_tests runs its tests at every
deployment the target serves (<layout>4, <layout>8, <layout>4_1m,
<layout>8_1m; the eight-rank ones gated on an eight-rank MPI world), and
the DSpark accuracy test is a method of the gate, so test_accuracy runs per
target too. test_llguidance stays at tp4, on the tp gate. The tests, their
bounds and their checks are unchanged: 85 cases before and after.
The GB300 test lists and the deployment guide use the new IDs, for example
TestModelingV2MistralLarge4Nvfp4Sm103Tep::test_gsm8k[tep8] and
TestModelingV2MistralLarge4Nvfp4Sm103Tep::test_accuracy[tep4-quick].
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…eat/mistral-large-4 Port of feat/mistral-large-4 3030ff5 (NVIDIA#21604) onto this branch. - Three sm_107 targets, tep, dep and tp: each is this branch's sm_103 target with NVIDIA#21604's changes. The FP8 block-scale linears run the Rubin MXFP8 GEMM (cute_dsl_fp8_block_scale_linear), their K128 weight scales converted once to the R128c4 K32 layout, and the indexer scores prefill tiles with a Triton FP8 MQA scorer (triton_fp8_mqa_logits). - Both entries sit under models/mistral_large4/ops/ beside the family's other entries, with their tests; l0_gr100 runs them on sm_107. - The Lamport process group's rendezvous address comes from get_local_ip. - Routing returns the sm_107 targets on compute capability (10, 7) and names sm107 in its refusal. - Accuracy: the gates share a layout base per target layout, and each architecture has its own gate class (Sm107Tep/Dep/Tp beside the sm_103 ones). test_accuracy_without_dspark runs the full and long_full profiles without the draft; artifacts go under the GPU architecture and record per-phase timings. RULER pins its attention-DP warm pairs for every long profile. - The vision, native DSpark and routing tests cover sm_107; the deployment guide lists VR200 and its validated scope. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Mistral Large 4 no longer has a checkpoint format of its own. It loads through checkpoint_format mistral, and what is specific to it follows from the checkpoint: a directory whose params.json describes the model (is_mistral_large4_checkpoint). - CheckpointFormat.MISTRAL_LARGE_4 is removed, with MistralLarge4CheckpointLoader and the Mistral Large 4 weight loader and weight mapper. MistralConfigLoader returns the adapted Mistral Large 4 config for such a directory and the plain Mistral config otherwise. - The model reads its own checkpoint. When a model class defines load_checkpoint(checkpoint_dir), ModelLoader calls it in place of load_weights and opens no checkpoint file itself. The six modeling_v2 targets stream their weights from the consolidated shard files (through the index, or each file's header when there is none). The modeling_v2 README lists the hook in the target lifecycle. - Tokenizer selection and generation_config.json follow the path main takes for checkpoint_format mistral: the LLM no longer forces the mistral-common tokenizer or skips generation_config.json. The worker's guided-decoding tokenizer and the LLM's params.json config lookup are specific to a Mistral Large 4 directory; Mistral Large 3 is unchanged. - Guided decoding keys on the tokenizer: the Tekken byte-vocabulary path (xgrammar tokenizer info, the llguidance tokenizer, the [THINK] JSON grammar) applies to a tokenizer that provides it. The default xgrammar backend, lm-eval's native chat tokens and the server ignoring a configured chat template (with a warning) key on model_type mistral_large_4. - MistralNativeInputProcessor derives _native_chat from its config, so it also holds for a processor built without __init__ (test_modeling_mistral on main), and an image token without image data is an error only for a model with a vision encoder. - The docs, the deployment guide, the quickstart, the serve schemas and the tests use checkpoint_format mistral. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…al-large-4 Port of feat/mistral-large-4 1295b9d (NVIDIA#22056) onto this branch. - An aggregated TP4 DSpark config (8k1k at concurrency 1 and 16) and three 1P1D DEP4 DSpark configs (8k1k at concurrency 64, 256 and 1024). The 1P1D configs run as ctx_only and gen_only_no_context on one 4-GPU node and as e2e on two. DSpark acceptance is pinned to 1.772 draft tokens per step. - The configs read the checkpoint with checkpoint_format mistral. - MODEL_PATH_DICT maps mistral_large_4_nvfp4 to the NVFP4 checkpoint. - GB300-4_GPUs-PyTorch-PerfSanity runs on 10 splits instead of 7. - The GB300 1P1D stage (CTX1-NODE1-GPU4-GEN1-NODE1-GPU4) is spelled out, one test per split over 5 splits, so test_to_stage_mapping.py finds it. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
… a model that reads its checkpoint
A model that defines load_checkpoint reads its checkpoint's files itself,
and no weight mapper applies to it. Dummy loading still asked the checkpoint
loader for one. Mistral Large 4 loads through checkpoint_format mistral,
which registers no mapper for MistralLarge4ForCausalLM and no default
mapper, so that lookup raised KeyError('mistral').
ModelLoader._initialize_dummy_weights builds the mapper only for a model
without load_checkpoint; filling dummy weights and copying the target's
weights into the draft model are unchanged. _model_reads_own_checkpoint is
shared with the checkpoint path. Tests cover both cases and the draft copy.
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…ain lays out modeling_v2 tests main (NVIDIA#21603) selects modeling_v2 tests by the scope of a change. Its CPU check requires an accuracy directory for every routed target, accuracy/modeling_v2/<family>/<target>/, reached by a test list, and l0_b300 lists the modeling_v2 GPU unit tests by file, op category and target instead of one directory. - The six gate classes move from accuracy/modeling_v2/test_mistral_large4.py into one test_accuracy.py per target under accuracy/modeling_v2/mistral_large4/. Each gate lists its tests, as main's targets do (the test-list check reads a test's class in its own directory); they run the _test_* implementations, the layouts and the deployments shared in mistral_large4/_common.py, the former file. Test names and parameters are unchanged. The test IDs in l0_gb300_multi_gpus, l0_gb300_multi_nodes_node2_gpu8 and the deployment guide follow. - l0_b300 lists the family's GPU unit tests the way main lists the others: its GPU-run top-level file, its op categories except comm (four GPUs, in l0_gb300_multi_gpus) and its sm_103 targets. They ran through l0_b300's former whole-directory entry. - The sm_107 (VR200) targets' GSM8K gate and quick DSpark profile at tep4, dep4 and tp4 go into the QA list llm_function_core.txt, where main lists its other sm_107 targets. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
…cope and loader checks After the rebase onto main c046811, three of main's checks fail on this branch's tests: - tests/unittest/serve/test_no_object_new.py: test servers are built with isolated_serve_env and its fakes, not object.__new__(OpenAIServer). The two Mistral Large 4 chat-handler tests in test_native_chat_template.py build their server that way and set the model type, the tokenizer and the LLM args they need. - test_modeling_v2_env_scope.py: every LLM an accuracy module builds sits in a scoped_env_overrides block. spec_evidence.traced_llm, which the long DSpark profiles build their engine with, opens that scope around the engine itself; its caller holds the same overrides. - main's _materialize_checkpoint_weights clears the checkpoint IO policy in the startup metadata before a model reads its own checkpoint. The loader stub in test_model_loader_lets_a_model_read_its_own_checkpoint carries the metadata, and the test checks that no policy is recorded. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
yizhang-nv
force-pushed
the
codex/ml4-mergeback-on-main
branch
from
November 12, 2026 08:29
797482d to
ffb7c15
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Draft for reviewing the Mistral Large 4 feature branch rebased onto
main. How it lands onmain(as is, or split into smaller PRs) is still open; this PR is the integration view.What is in it
The Mistral Large 4 work on
feat/mistral-large-4(3fb4be0), squashed into one commit and rebased ontomainb433473 (first carried onto 19c1f46):Not carried: the six infra commits that are already on
main(#21133, #19557, #21068, #20564, #20612, #20596) and the side-branch version bump (cf7ae79, 1.4.0a0).feat/mistral-large-43030ff5 (#21604, VR200 support) landed after that squash; it is carried as its own commit (83fe635, see VR200 (sm_107) targets below).Three commits follow it: 523f496 removes the Mistral Large 4 checkpoint format and 797482d fixes dummy loading after it (see No checkpoint format of its own below), and 287b55b carries
feat/mistral-large-41295b9d (#22056, GB300 perf-sanity cases; see GB300 perf-sanity cases below).Adapted to
main's structure (in the squash commit)mistral_large_4tomain'sCheckpointFormat([TRTLLM-16725][refactor] BREAKING: Use a string enum for checkpoint formats and simplify checkpoint I/O helpers #20703), next tomistral_large_3, with its loader, config loader, weight loader and mapper. 523f496 removes it again (see No checkpoint format of its own below).@routes("MistralLarge4ForCausalLM")and is imported frommodels/__init__.py;routetakes(model_config, llm_args, compute_capability).thop_attentioncall sites usemain's argument names; top-k kernels come from_torch/kernels/blackwell/top_k.TARGET_CELLS(comm ops keep them inCELLS, asmain's comm ops do).main's refactored worker (block placement, vocab-parallel drafting, Markov TP exchange).main's KV manager V2 case (one call when both apply).main's versioned rank-state frame (version -2 adds a spec snapshot).main's prompt pipeline andServingExtensionhooks; Mistral Large 4 adds therender_needs_media,forced_tool_call_decodingandvalidate_chat_requesthooks.storefollowsmain(accepted and ignored).main's engine/runner refactor, scoped reference precision (true_fp32_matmul) and scoped engine environments (scoped_env_overrides).main; llguidance 1.7.6 and mistral-common>=1.13.0 (Python >= 3.11) from the feature branch.Follow-up commits found by testing on
mainmain's snapshots)._arch_index: themistral_large_4multimodal model type, so the lazily imported model zoo finds its placeholder registration.tool_choice.models/mistral_large4/ops/<category>/and 6 kernel modules tomodels/mistral_large4/kernels/. Their tests move totests/unittest/_torch/modeling_v2/models/mistral_large4/ops/. That directory is a package, so the tests can share basenames with other families' tests. The shared catalog test directories keepmain's layout, with no package markers. Three entries that no target calls are removed:cute_dsl_indexer_topk_prefill,fp8_block_scale_moe_runnerandfp8_quantize_1x128. The GVR top-k prefill test no longer compares against the removed radix entry.l0_gb300_multi_gpusruns the moved all-reduce op matrix, andl0_b300leaves it out, as it does the shared comm tests.indexer_topk_decode's launch-ceiling test no longer expects torch's next launch to report the failed launch. Onmain,libth_common.solinks the CUDA runtime statically (libtensorrt_llmbrings incudart_staticahead of torch'slibcudart.so), so that error never reaches torch; the test checks that the launch did not run.test_ruler_accuracy.pycarriescpu_only(listed inl0_cpu, but never collected there).return_tensors: onmainevery format renders untensored and keeps image sizes ([None][chore] Mass merge release/1.3 branch back to main (batch 2) #20843).serve/render) skips the HF processor formistral_large_4, astrtllm-servedoes. The NVFP4 checkpoint shipsprocessor_config.json(PixtralProcessor) andchat_template.jinja; the renderer loaded that processor, and its fingerprint took in the processor's chat template, so it did not match the server's.supported-models.mdlistsMistralLarge4ForCausalLM(support table, language and multimodal feature matrices, footnotes on the GB300 NVFP4 Mistral-format scope and the validated disaggregated setup), andtensorrt_llm/usage/architecture_allowlist.pyreports it by name.mistral_large_4_nvfp4__sm_103__tep/__dep/__tp(classesModelingV2MistralLarge4Nvfp4Sm103Tep/Dep/Tp);tep4,tep8and the rest name deployments.TestModelingV2MistralLarge4Nvfp4Sm103{Tep,Dep,Tp}class per target, as the modeling_v2 claims require. Each test runs at every deployment the target serves (<layout>4,<layout>8,<layout>4_1m,<layout>8_1m), and the DSpark accuracy test is a method of the gate (...Sm103Tep::test_accuracy[tep4-quick]);test_llguidancestays at tp4 on the tp gate. 85 cases before and after; the GB300 lists use the new IDs.main's interfaces (prepare_resourceremoved, single-rank ADP skip, KV manager block reuse, multimodal item metadata, tuning-config cache key,bmm_out's 3-D check) and for [None][chore] Mistral Large 4: remove defensive checks and leftovers of removed paths #21616's FP8 KV-cache algorithm inquantization_config.VR200 (sm_107) targets (83fe635, from #21604)
mistral_large_4_nvfp4__sm_107__tep/__dep/__tp(classesModelingV2MistralLarge4Nvfp4Sm107Tep/Dep/Tp). Each is this branch's sm_103 target with [None][feat] Add VR200 support for Mistral Large 4 modeling_v2 #21604's changes; the two differ in about 100 lines per target, the same lines [None][feat] Add VR200 support for Mistral Large 4 modeling_v2 #21604 changed: the FP8 block-scale linears run the Rubin MXFP8 GEMM (cute_dsl_fp8_block_scale_linear, K128 weight scales converted once to the R128c4 K32 layout), and the indexer scores prefill tiles with a Triton FP8 MQA scorer (triton_fp8_mqa_logits). They stay separate per-architecture targets, as GLM-5.2 and MiniMax-M3 are onmain; ansm_10xtarget would branch on the architecture in those places.models/mistral_large4/ops/with the family's other entries, and their tests undertests/unittest/_torch/modeling_v2/models/mistral_large4/ops/.l0_gr100runs them on sm_107 (the B300 list reaches them too, where they skip).get_local_ip(a hostname can resolve to a loopback address).Sm107Tep/Dep/Tpnext to the sm_103 ones, each skipped on the other architecture).test_accuracy_without_dsparkruns thefullandlong_fullprofiles with the draft off; artifacts go under the GPU architecture and record per-phase timings. The 85 sm_103 test IDs are unchanged.No checkpoint format of its own (523f496, 797482d)
Mistral Large 4 loads through
checkpoint_format: mistral. What is specific to it follows from the checkpoint: a directory whoseparams.jsondescribes the model (is_mistral_large4_checkpoint).CheckpointFormat.MISTRAL_LARGE_4is removed, withMistralLarge4CheckpointLoaderand the Mistral Large 4 weight loader and mapper ([None][refactor] Mistral Large 4: own the checkpoint weight loader and mapper #21548).MistralConfigLoaderreturns the adapted Mistral Large 4 config for such a directory and the plain Mistral config otherwise. The serve config schemas no longer list the value.load_checkpoint(checkpoint_dir).ModelLoaderthen calls it in place ofload_weightsand opens no checkpoint file itself. The six targets stream their weights from the consolidated shard files (through the index, or each file's header when there is none). The modeling_v2 README lists the hook in the target lifecycle.load_checkpoint.checkpoint_format: mistralregisters no mapper forMistralLarge4ForCausalLMand no default one, so after 523f496 the lookup raisedKeyError('mistral').generation_config.jsonfollowmain's path forcheckpoint_format: mistral. The worker's guided-decoding tokenizer and the LLM'sparams.jsonconfig lookup apply only to a Mistral Large 4 directory; Mistral Large 3 is unchanged.[THINK]JSON grammar) applies to a tokenizer that provides it. The default xgrammar backend, lm-eval's native chat tokens and the server ignoring a configured chat template (with a warning) key onmodel_typemistral_large_4.MistralNativeInputProcessor._native_chatis derived from the config. It used to be set in__init__only, andmain'stest_modeling_mistral(l0_a10) failed on this branch because of it. An image token without image data is an error only for a model with a vision encoder.GB300 perf-sanity cases (287b55b, from #22056)
ctx_onlyandgen_only_no_contextinl0_gb300_multi_gpus_perf_sanity, and ase2einl0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node1_gpu4. DSpark acceptance is pinned to 1.772 draft tokens per step (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS).checkpoint_format: mistral(523f496). The entries are merged intomain's test lists, which have changed since; they keep [None][test] Add Mistral Large 4 GB300 perf CI coverage #22056'sTIMEOUT (90)(main's other GB300 perf-sanity entries are 30 to 125).GB300-4_GPUs-PyTorch-PerfSanitygoes from 7 to 10 splits, as in [None][test] Add Mistral Large 4 GB300 perf CI coverage #22056. The 1P1D stage, whichmainbuilds withbuildStageConfigsover 2 splits (Nemotron-Ultra-V3 and GLM-5.3-Flash), is spelled out with one test per split over 5 splits, as [None][test] Add Mistral Large 4 GB300 perf CI coverage #22056 does, so thattest_to_stage_mapping.pyfinds it.main's cache-transceiver precheck runs before the e2e cases. The NVFP4 checkpoint has noconfig.json, so the precheck uses its generic KV pool (KV cache manager V2, Python transceiver): it checks the transfer path, not this model's KV layout.Speculative decoding (DSpark)
The targets run DSpark in split mode: each target builds
main's stockDSparkWorkerwithget_spec_worker, and itsDraftModelsupplies the drafter's model math. The worker reaches that math only through draft-model methods, now documented onDFlashWorkerandDSparkWorker. No branch in the worker names Mistral Large 4:project_target_hidden(hidden)projects the captured target features. The worker now checks for this method before capturing context, where it used to check forfcandhidden_norm. Every in-tree drafter has both, so nothing changes for them.embed_draft_block(bonus, block_size)(optional) builds the input block (learned position-specific mask embeddings), throughdflash_draft_block_embedding. A drafter with it keeps the eager context preparation.dflash_forward_with_metadata(*, attn_metadata, **block_inputs)(optional) replacesdflash_forwardwhen the block forward depends on the step's attention metadata (prefill or decode implementation), throughdflash_block_forward.has_recurrent_markov_head:apply_markov_chain_logitsalso receives the block's draft hidden states and chains the bias itself.DSparkWorkerdeclines vocab-shard drafting for such a drafter, because the shard chain computes the vanilla Markov head only.apply_draft_logit_transform(logits)returns its input unchanged: the drafter'slogits_processoralready soft-caps the logits, andmain's worker ([TRTLLM-17151][fix] Apply the logit transform and sliding window DFlash drafters declare #21367) calls this method after it.The worker-side tests live in
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_hooks.py(l0_cpu). A fused DSpark target, with the step in the target's forward, can follow once the fused-mode contract (#21787) covers DSpark.Open items
ServingExtensionhooks, the rank-state frame version and theload_checkpointmodel hook inModelLoader(523f496).pytest.mark.largeis used but not registered.security_scanning/poetry.lockis left to the nightly lock refresh.Test Coverage
CPU (B200 node, container from
main's image, C++ built for sm_100):modeling_v2/models/mistral_large4,models/checkpoints,test_model_config.py, routing, claims;-m cpu_only): 1574 passed, 674 skipped (no checkpoint, or sm_103-only routes on the sm_100 parameter), 1 failed (the accuracy-gate claim above).unittest/entries ofl0_cpu.ymlwith-m cpu_only; failing files were rerun one per process and compared againstmain19c1f46 with the same build. The remaining failures are the same onmain, plus the claim above.llmapi/appsextensions/chat/tool-parser tests, the native Mistral Large 4 frontend tests andapi_stability: 1370 passed.GB300 (OCI, sm_103 build, released NVFP4 checkpoint):
TestModelingV2MistralLarge4Nvfp4Sm103*::test_gsm8ktest_accuracy[quick-*]GSM8Ktest_llguidance: passed (5 cases).unittest/_torch/modeling_v2with-m "not cpu_only", asl0_b300selects it, does not finish within 2 hours on one node, so it was run in parts. Everything passed except the Mistral Large 4 target tests that need multi-rank engines; these were rerun through MPI (table above). The FMHA JIT finds its headers throughpip show tensorrt_llm, so the source-tree runs register the checkout as an editable install.models/mistral_large4/ops, after the move): 2193 passed, 2 skipped. The all-reduce op matrix passed on 4 GPUs. The three targets'test_gsm8kpassed again after the move.cpu_onlytests with the checkpoint present: 2031 passed, 1 failed (main's real-MPI read-ahead test, which the container cannot spawn).After rebasing onto
mainb433473 (B200 node, C++ rebuilt for sm_100 at the rebased head):-m cpu_only): 1634 passed, 674 skipped, 1 failed (the accuracy-gate claim above).generate_trtllm_serve_schemas.py --checkandtest_trtllm_serve_config_schemas.py: 248 passed.speculative/hw_agnosticandtest_bench.py(-m cpu_only): 748 passed.llmapi/appschat utils, serving extensions, tool parsers, reasoning prompt resolution, chat store and template guards, the Mistral tokenizer factory,test_lm_eval.pyandapi_stability: 1732 passed, 1 skipped.test_py_executor.py: 620 passed.speculative/hw_agnosticwith-m "not cpu_only": 711 passed.test_render_*.py,test_architecture_allowlist.pyandtest_native_frontend.pywithMISTRAL_LARGE4_CHECKPOINT_DIRset to the checkpoint's metadata: 180 passed; the new renderer test fails without the fix. The family'scpu_onlytests with that variable set: 904 passed, 423 skipped. On the checkpoint's metadata the renderer's fingerprint now equals the server's.MISTRAL_LARGE4_CHECKPOINT_DIRset to the checkpoint's metadata: 1696 passed, 614 skipped, 0 failed (the accuracy-gate claim passes). Collectingunittest/_torch/modeling_v2(7510 cases) and the accuracy file (85 cases) raises nothing; every listed accuracy ID is collected.Draft-model methods and the logit-transform fix (dee9070, c25b682), B200 node, before the rebase:
unittest/_torch/speculative/hw_agnostic: 681 passed with-m cpu_onlyand no GPU visible; 704 passed with-m "not cpu_only"on one GPU.test_native_dspark.py(-m cpu_only): 101 passed, 6 skipped (checkpoint not underLLM_MODELS_ROOT).VR200 port (83fe635):
-m cpu_only) against the parent commit in the same container: no new failure (the remaining failures fail identically on the parent: checkpoint shards missing from the metadata view, no GPU visible, stubs of removed engine methods). Collecting the accuracy file: 218 cases, every previous ID present. pre-commit passes.TestModelingV2MistralLarge4Nvfp4Sm107*::test_gsm8k(reference 91.926)test_accuracy[*-quick]GSM8KSm103Tep::test_gsm8k[tep4]skips on sm_107.models/mistral_large4/ops, withoutcomm): 2239 passed, 30 skipped; the two new entries pass all their cells (112 and 18). The family's top-level tests with routing and claims: 895 passed.107-realalone skips FlashMLA: its CMake gate lists 90, 100 and 103, although it builds90;100f, which runs on sm_107. CI's aarch64 list includes 100 and 103, so CI builds are not affected; the VR200 runs used a node-local copy with 107 added to that gate.test_accuracy_without_dspark, the full and long profiles, the 8-rank deployments andtest_llguidance. GB300 has not run this commit (the sm_103 targets are unchanged; their gate classes moved under the shared layout bases).Checkpoint format removal (523f496):
-m cpu_only): the Mistral Large 4 unit set with routing and claims,models/checkpoints,test_model_config.py,pyexecutor,test_modeling_mistral.py,test_llm_args.py, the serve config schema tests,api_stability,serve,tokenizer,inputs,test_lm_eval.pyandtest_bench.py: 7383 passed, 3 failed (the threeqwen3_5_hybridcases oftest_kv_cache_v2_vocab_size.py, a Qwen3.5 hybrid stub problem unrelated to Mistral).generate_trtllm_serve_schemas.py --checkand pre-commit pass.Sm103{Tep,Dep,Tp}::test_gsm8kat tep4, dep4 and tp4, andtest_llguidanceat tp4: passed.test_model_loader_checkpoint_io.py(-m cpu_only): 40 passed, with the new dummy-loading cases. Dummy loading of the real targets has not run.GB300 perf-sanity port (287b55b):
test_to_stage_mapping.py).check_test_list.py --validate: 2200 entries pass.test_test_to_stage_mapping.py,test_disagg_e2e_configs.pyandtest_cache_transceiver_precheck_config.py: 256 passed. pre-commit passes.Not covered: the 8-rank two-node deployments (tep8/dep8/tp8), DSpark full/long profiles, disaggregated serving, the GB300 perf-sanity cases, and Jenkins CI. The GB300 runs used the sbsa devel image 202610050217-19679, not
main's current 202611020300-19783.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.🤖 Generated with Claude Code