Repository navigation
Conversation
One-model speculative decoding spends host time between two decode-step graph replays on small kernel launches and host-to-device copies. - On SM 100, a CUDA graph decode step of an engine whose attention metadata is a plain TrtllmAttentionMetadata writes its per-step inputs (overlap gathers, positions, prompt and KV lengths, KV block offsets) with one CuTe DSL kernel launch (StepInputStage). Eager steps on SM 100 run the overlap gathers as one kernel. - On SM 100, the speculative sampler moves the forward's outputs into its slot stores with one CuTe DSL kernel (SlotScatter). - The sampler's host copies of its stores run on its D2H side stream, and one-model speculative sampling runs on the execution stream. - CUDA graphs of a speculative engine read input ids and positions from the engine's buffers instead of copying them in at each replay. - Copies whose values did not change are skipped (gather ids, slot tables, previous-batch indices), as are the Philox seed and offset uploads of all-greedy batches and the unused num_accepted_draft_tokens upload. Other devices, and calls the kernels do not cover, keep the torch path. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…ache manager MambaHybridCacheManagerV2 takes kda_token_states (default off). With the KDA replay caches it then also keeps the fp32 state after every draft of the last verify round per slot (kda_state_tok, [layers, slots, num_spec, H, V, K]), passes the layer's view in the speculative state and moves it with the slot. A verify kernel can start the next round from the state after the accepted drafts instead of replaying them. The cache factory enables it when the model's backbone sets kda_token_states. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
The one-model worker's acceptance returns new_tokens as torch.empty [N, K + 1]. A context row writes only column 0 and accepts one token; a generation row writes every column. The scatter kernel copied every column of every row into the sampler's new_tokens store, so for context rows it read columns 1..K, which were never written; compute-sanitizer initcheck reports those reads. update_requests reads the store only up to new_tokens_lens, so no output changed. The scatter kernel now reads a row's new tokens only below its new_tokens_lens and stores zeros past them. A thread loads the row's length once; the column 0 thread's store_lens write reuses that load. The sampler's torch store update, used where the kernel does not run, is unchanged: nothing reads past new_tokens_lens. test_spec_step_copies: the reference zeroes the same columns, and the Step cases draw accepted lengths from 0 to past the output width. test_scatter_zeros_past_accepted_length covers a mixed batch laid out as the worker writes it, with a skipped chunked-context row. "poisoned" fills the never-written columns with a sentinel and fails on the old kernel; "never_written" leaves them unwritten for compute-sanitizer initcheck. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…t stream measure_graph (test_graph_replay) and time_graph build their inputs on the current stream, then make their first eager call on a fresh side stream without ordering it after that work. Under the stream-ordered cudaMallocAsync allocator the call can write memory before its allocation in the side stream's order; with any allocator it can read inputs whose initialization has not finished. The side stream now waits for the current stream first. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…with per-token KDA states With per-token KDA states, kda_state_tok holds num_spec recurrent states per SSM state slot, outside the cache quota, addressed by the request's SSM slot (_allocate_pool_replay_buffers). V2 sizes the SSM pool group by the typical step's ratio, so its slot count, and kda_state_tok with it, grows with the quota. Without block reuse a slot past the min-slots floor (resident sequences plus reserved dummies) never holds a request, so those slots only add memory. On the Kimi K3 stack (7 DSpark drafts, TP16, max_seq_len 4096, free_gpu_memory_fraction 0.25) the pool got 183 slots: 32.4 GiB of kda_state_tok beside a 13.2 GiB quota. It ran out of device memory above fraction 0.28 at max_seq_len 4096, and above 0.16 at 512. With per-token KDA states and block reuse off, the typical step's request capacity is now raised until its attention pages alone take the quota. The ratio then gives the SSM pool group fewer slots than its floor, so the min-slots constraint sets its size and attention gets the rest. Without per-token states (replay caches only), or with block reuse, the ratio sizes the pool as before. test_v2_kda_token_states_keep_ssm_pool_at_live_floor builds the V2 storage for a 128 MiB quota with 2 MiB state slots: with per-token states the SSM pool keeps its floor (plus one for max_util_for_resume) and attention takes the rest of the quota; without them, or with block reuse, the pool keeps its ratio share. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…tputs maybe_execute_in_parallel runs fn1 on aux_stream and returns its outputs to the caller, which uses them on its own stream after the event join. Those tensors are allocated in aux_stream order. When the caller drops them, the caching allocator returns their blocks to aux_stream's pool at once and cudaMallocAsync frees them in aux_stream order, so later aux_stream work can take and overwrite the memory while the calling stream's reads are still queued. compute-sanitizer memcheck with cudaMallocAsync reports this as use-after-free on the outputs Kimi K3's KDA mixer and MoE shared expert return this way. The CUDA tensors fn1 returns, directly or in nested tuples, lists and dict values, are now recorded on the calling stream after the join. cudaMallocAsync raises in record_stream on memory it did not allocate (a buffer from outside PyTorch, wrapped as a tensor); only that error is caught, and the tensor is skipped, as the native allocator does. Nothing is recorded under CUDA graph capture or torch.compile tracing. In a capture the caching allocator would hold every recorded block freed during the capture until it ends, so each call's outputs would take new graph pool memory, and under cudaMallocAsync freeing a recorded output inside the capture fails it with "capturing stream has unjoined work". A captured graph relies instead on later aux_stream work waiting on the calling stream, as the next call's event0.wait() does. Dynamo traces neither the capture query nor record_stream, so a compiled caller keeps the previous behavior and still traces fullgraph. test_multi_stream_utils.py reads fn1's output on the calling stream behind a sleep, drops it, and fills a new aux_stream allocation of the same size at once; the read must still see fn1's values. Before, it saw the fill under both allocators. The file also covers a buffer from cudaMalloc, a capture whose outputs outlive it, outputs dropped inside a capture, and a fullgraph torch.compile of a caller. It reruns itself under PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync and joins the l0_h100 and l0_b200 pre-merge module tests. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…ive for captured graphs Growing the MNNVL all-reduce workspace replaced it and freed the previous buffers, pointer tables and flags while CUDA graphs captured before the growth still launched on them. The replaced workspaces now stay alive with the process, and creating or growing a workspace during graph capture raises instead of failing inside the driver. test_mnnvl_workspace_growth_keeps_captured_graphs captures a one-token all-reduce, grows the workspace with an eager two-shot call, replays the graph and checks growth under capture is refused. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…lti-GPU list test_mnnvl_workspace_growth_keeps_captured_graphs needs fabric-backed MNNVL on two GB200 GPUs, like the checkpoint graph tests next to it. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
The workspace-growth test came from a branch whose yapf settings differ from main's; format it as main's pre-commit does. No behaviour change. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…re guard Creating or growing the MNNVL all-reduce workspace under CUDA-graph capture raises before the collective allocation. Its only test was the last block of the 2-rank growth test, which runs on GB200 multi-GPU stages only. test_mnnvl_workspace_creation_refuses_graph_capture creates the workspace for a one-rank group inside a capture and expects the guard's RuntimeError. The guard runs before any MNNVL call, so one GPU suffices; without it the call enters the workspace construction under capture and fails inside it with another error. The test joins l0_b200's single-GPU PyTorch block. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
… kernels oneshotAllreduceFusionKernel and rmsNormLamport used LamportFlags' cluster arrival, in which every thread arrives on the cluster barrier but only cluster rank 0's first warp waits, and then arrived on the cluster barrier again in the RMSNorm epilogue's cluster reduction. Each thread must arrive and wait once per phase: a second arrival counted toward the first phase lets the reduction read a peer CTA's partial sum before it is written, and counts the cluster in before a late CTA has read the buffer flags. Both kernels now count arrivals per CTA, as the attn_res kernel already does; their cluster barrier is left to the reduction. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…luster reduction writes The RESIDUAL_RMS_NORM epilogues of oneshotAllreduceFusionKernel and rmsNormLamport write each CTA's partial sum into the other CTAs' shared memory through the cluster mapping before any cluster barrier wait, so a target CTA may not have started yet; racecheck reports every such write once the kernels' arrival counting no longer stalls it. A cluster.sync() before the writes makes every CTA of the cluster start first, as the attn_res kernel's early arrive / wait does. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…ads before the next chunk's writes attn_res_fwd_online_v2_kernel reuses one ws_stats buffer every chunk of candidates: lane 0 of each consumer warp writes its row, a named barrier orders the writes before every consumer's cross-warp reads, but nothing ordered those reads before the next chunk's writes. With more than one chunk (N > 4 at H 7168) a warp that finished reading could overwrite its row while a slower warp still read it, mixing the next chunk's statistics into this chunk's logits. Under compute-sanitizer racecheck, which serializes warps, every such call returned a wrong mixture (relative error 0.6-2.1 of max |ref|); natively it was not observed. Add the second barrier the persistent fork already has. Also order every lane's reads of a chunk's slots before lane 0 releases them to the producer (__syncwarp before the bar_consumed arrive), in online_v2 and in the persistent fork. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…eases a TMA-filled slot The online_v2, N = 1 tile and persistent fused kernels read their V (and delta) slots with generic-proxy shared loads and release them with an mbarrier arrive; the producer then refills them with cp.async.bulk (async proxy). Accesses to one location through two proxies need a cross-proxy fence: fence.proxy.async.shared::cta in every reading thread before the release. In the current SASS each release already issues after an instruction that consumes the last load, so this makes that ordering independent of instruction scheduling. Outputs are bit-identical. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…h slot release test_attn_res_proxy_fences.py reads attnResFwd.cu and attnResFwdPersistentFused.cu and checks every consumer release of a cp.async.bulk-filled slot (an arrive on bar_consumed). Each must follow fence.proxy.async.shared::cta, with only the warp sync and the lane-0 guard between the two. The fence orders the consumers' generic-proxy reads of the slot before the producer's async-proxy refill. Without it the order rests on instruction scheduling, which no test of values can see. The check fails on the sources from before the fence was added (four releases) and passes on the current ones. CPU only. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
CUDA-graph padding rows and warmup dummies share the drafter's dummy slot. Their accepted tokens grow its context length like any request's, but nothing reset it, while their page-table rows are the padding request's pages with every other entry mapped to page 0, another request's page. Once the dummy length passed the padding request's own pages, the padding rows' context K / V (k3_ctx_kv, or the torch path) landed on that page: a live request's drafter context overwritten on every padded step, costing acceptance. prepare() now writes 0 to the dummy slot every step, with the evicted slots. test_dflash_dummy_slot.py: padded steps keep the dummy slot at 0 and the real slots untouched; an evicted request and the dummy reset in one step. It fails without the reset. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
… mxfp8_quantize and noaux_tc_op invokeFP4Quantization and invokeMxFP8Quantization launch one block per row (up to an occupancy cap), invokeNoAuxTc one block per token. With no rows the grid has zero blocks: the launch fails with cudaErrorInvalidValue, nothing checks it, and the error stays pending until the next launch that reads the last error raises it, usually an unrelated PyTorch kernel. The three launchers now return before the launch when there are no rows (there is nothing to write; the swizzled scale-factor layout pads zero rows to zero), and their cudaLaunchKernelEx calls are checked. noaux_tc_op still rejects an unsupported configuration at zero tokens. That pending error is what failed the three modeling_v2 catalog tests that run right after a zero-row call in their files; their waivers are removed. Tests: zero-row fp4_quantize (bf16 and fp8 input, both scale-factor layouts) and mxfp8_quantize (linear and swizzled), and zero-token noaux_tc_op (single and multi group), each followed by a normal call that must match the one before it. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
trtllm::kda_prefill returned a view of a buffer it reuses: the output was one of the intermediate tensors _get_buffers caches per input shape, so every call with the same shapes wrote the same storage, and a caller's output changed under it at the next such call. The op's fake declares a new tensor (v.new_empty(v.shape)), so nothing orders the readers of one call's output before the next call. Allocate the output on every call, with the varlen path's one chunk of slack past the last token like the cached buffers, and wrap it for K4 there; the cache no longer holds it. Test: a second call with the same shapes and other inputs must leave the first call's output unchanged (varlen, and one padded equal-length row). Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
The KDA layers (tensorrt_llm/_torch/modules/kimi_kda) import flash-linear-attention's ShortConvolution when they load and run its chunked KDA for prefill batches below four 64-token chunks, but flash-linear-attention is listed only in requirements-dev.txt. CI installs that file, so every test passes, while a plain pip install of the wheel, and the release container, which installs the wheel with its runtime requirements only, cannot load either model: No module named 'fla'. Require fla-core, the distribution that provides fla.modules and fla.ops, at the version requirements-dev.txt already installs through flash-linear-attention 0.5.2. Its only base dependency, einops, is a runtime requirement already. A CPU test checks that every package the KDA layers import when they load is a runtime requirement. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…ture used it Growth retired every replaced workspace for the life of the process, including ones no CUDA graph had captured. Growth rounds to 8 MiB, so a run of rising eager sizes could keep many dead workspaces. The getter now marks a workspace it hands out during capture, and a growth keeps the previous workspace only when it is marked. An unmarked one is released at the replacement, after every rank's device synchronize in the protocol reset. The growth test checks that a second growth, past a workspace no capture used, leaves only the captured one in the retired list and that the graph still replays. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…ken embedding GQADSparkForCausalLM takes the target's embedding in load_weights_from_target_model, so the DFlash worker built every masked block slot from the target's row for the mask token, which the target never trained. A drafter checkpoint that ships its own embedding has trained that row. load_weights now keeps it as mask_token_embedding, and the worker's noise block (dflash_noise_block_embedding, the same computation moved into a function) uses it for every masked slot when the drafter has one for the mask id in use; slot 0 stays the bonus token's shared lookup. Every rank still embeds the mask id through embed_tokens.forward, as before. Drafts change; under greedy verification the target's tokens do not. Plain DFlash and MLA DSpark drafters set no mask_token_embedding and are unchanged. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
… the GPU quota at the SSM live floor With per-token KDA states and block reuse off, the SSM pool is held at its live floor and the typical step's attention pages fill the GPU quota. kda_state_tok, num_spec fp32 states per SSM slot, is allocated beside the pools afterwards. So the pools took the whole quota and the states came on top of it: initialization could run out of device memory although _minimum_live_gpu_quota() passed. At the floor the SSM pool keeps ceil(floor / max_util_for_resume) slots: the floor's min-slots constraint, as the runtime scales it. _build_cache_config now takes the token states of that many slots out of the GPU tier quota before the pools are sized, and logs the split. The runtime rounds each pool up to whole allocation grains, which can add SSM slots when a slot is smaller than a grain; their states are not reserved. The quota for max_tokens and the minimum quota count the token states of all those slots (_state_quota), not only the floor's. Otherwise the attention pages could get less than max_tokens needs: on Kimi K3 (DSpark with 7 drafts, TP16, max_batch_size 8: a floor of 9 slots, 10 kept) the estimation quota for max_tokens 16896 would leave the KV pool about 13.5K tokens. With it the quota is 2.56 GiB and the pool about 20.6K tokens (arithmetic from the logged 2.38 GiB quota at 27648 bytes per token). test_v2_kda_token_states_keep_ssm_pool_at_live_floor now also checks that the pools and the token states of every SSM slot fit in the quota, with attention taking the rest of it to within one grain. test_v2_kda_token_states_reserve_quota_before_sizing checks the reserved bytes, and that the minimum quota and the quota for max_tokens exceed those of the same manager without token states by the states of every kept slot. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
A decode step of a hybrid Mamba model uploaded its recurrent-state metadata on every step, even when the batch had not changed: the V2 manager's state indices (prepare_resources), the dummy-request mask (twice per step on the V2 manager: from prepare_resources and from get_state_indices) and Mamba2Metadata's copy of the state indices. Each is a small host-to-device copy on the execution stream between two CUDA graph replays, and the manager's two first allocate fresh pinned staging. Each of these buffers has one writer, so they now go through copy_to_device_if_changed, as the speculative engine's per-step inputs already do: a step whose values did not change enqueues no copy. test_v2_state_index_setup_skips_unchanged_uploads and test_prepare_skips_an_unchanged_state_index_upload overwrite the device buffers after a first step, then check that an unchanged step leaves them alone and that a changed step uploads again. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
DFlashSpecMetadata.prepare copied batch_indices_cuda and the worker's _batch_to_slot to the device on every step. Both now go through copy_to_device_if_changed, so a decode step whose batch is unchanged enqueues no copy. The worker's post-prefill rebuild of _batch_to_slot moves into _rebuild_batch_to_slot and goes through the same helper: the helper keeps the values it last copied on the buffer, so a plain copy there would let a later step that maps back to those values skip its copy and leave the rebuild's slots in place. test_dflash_copy_skip.py (cpu_only) checks that both buffers always hold the latest mapping, including a step whose mapping returns to an earlier one after a rebuild, and that an unchanged step copies nothing. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
… slots as the runtime does The runtime holds max_util_for_resume as a C++ float, and _build_cache_config passed that value to _live_floor_token_state_slots, while _state_quota (behind _minimum_live_gpu_quota and _get_quota_from_max_tokens) passed the Python double. When the floor divided by the ratio is a whole number, the two ceilings differ by one slot: at 0.95, a floor of 19 slots keeps 21 slots (19 / float32(0.95) is just above 20), while the double gave 20. The minimum and the quota for max_tokens then counted one slot's token states fewer than the build reserves. _live_floor_token_state_slots now rounds the ratio through float32, as kv_cache_manager_v2.py does for its own floor. The new test checks that the minimum quota and the build reserve the same 21 slots at a floor of 19 and a ratio of 0.95. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…SSM live floor kda_state_tok is allocated for every slot the SSM pool keeps, beside the V2 pools. Its bytes come out of the GPU quota only at the live floor (block reuse off, no pool_ratio, attention layers on the rank), where the slot count is known before the pools are sized. With block reuse, a pool_ratio, or no attention layers, the ratio-sized pool kept many more slots and kda_state_tok took their states outside the quota. The K3 targets only refused block reuse at their first forward, after the manager had allocated. _build_cache_config now raises ValueError for per-token states off the live floor. test_v2_kda_token_states_refuse_a_pool_off_the_live_floor covers the three cases. The ratio-sizing case of test_v2_kda_token_states_keep_ssm_pool_at_live_floor now runs without the states, since that combination is refused. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…y state MambaHybridCacheManagerV2 keeps the KDA replay state (prev_num_accepted_tokens and the six replay caches) per SSM slot, outside the V2 pools, and remembers each request's last slot until free_resources. A V2 suspend (the scheduler's eviction and self-eviction, the pool rebalance) only suspends the pages. Under pressure the host tier takes them, the slot goes to another request, and the suspended request resumes in whichever slot is free. Its next _setup_state_indices then relocated the old slot's contents, by then the other request's replay state, or, back in its old slot, kept them: the next verify replayed another request's drafts. suspend_request now copies the slot's replay state to host memory and forgets the slot; _setup_state_indices writes it into the slot the request resumes in. update_resources gives the host drafter's acceptance of a request suspended after its batch was scheduled to the saved state instead of the old slot. Relocation and the copy share one list of the per-slot buffers. The new tests resume a request in a new slot and in its old one and record a late acceptance on a manager stub, and on the V2 runtime let new caches take a suspended request's slot through the host tier before it resumes. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…KDA replay caches The KDA replay caches and prev_num_accepted_tokens are allocated once, one entry per slot the SSM pool keeps at construction (get_page_index_upper_bound). With enable_kv_pool_rebalance, adjust() resizes the pool groups to the auto-tuner's target ratio, which gives the SSM pool one slot per typical request, so a pool built smaller than that grows past the buffers. A request in a slot past them makes the prev_num_accepted_tokens[slot] = 0 reset fail, and the KDA verify kernels index the replay entries by slot without a bound. Nothing refused the combination: _util.py refuses a rebalance only with block reuse, and the executor's rebalance gate skips only two-model drafters, so one-model speculative decoding on KDA layers reaches it. _build_cache_config now raises ValueError for the KDA replay path with enable_kv_pool_rebalance, as KvCacheConfig does for fp8_ds_mla's fixed pool views. test_v2_kda_replay_refuses_kv_pool_rebalance covers the refusal and the control without the replay caches. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
_get_buffers keeps the prefill scratch for at most _BUF_CACHE_MAX_ENTRIES token-count shapes and evicts the oldest entry on a new shape. The K123 and akk_inv launches wrapped that scratch with _ct_cached, whose id-keyed _input_wrap_cache holds each wrapper until the wrapped tensor is garbage collected. The wrapper's DLPack capsule keeps the tensor alive, so that never happens: every evicted entry's scratch stayed allocated. A server whose prefill batches keep changing token count grows by one entry's scratch per eviction (GSM8K on Kimi K3: about 500 MiB per 200 requests on every rank). Cache the scratch wrappers in the scratch's own _buf_cache entry instead (_scratch_ct), so evicting the entry releases the scratch and its wrappers together; the per-call from_dlpack savings stay. _ct_cached and _input_wrap_cache have no other users and are removed. Test: test_evicted_prefill_scratch_is_released runs cyclic passes over more token-count shapes than the cache holds. After every later pass the cache holds at most _BUF_CACHE_MAX_ENTRIES entries, the allocated bytes are back at their value after the first pass, and the tensors of an entry cached at the end of the first pass, which the next pass evicts, are freed. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…_parallel forward_with_base recorded the calling stream on each LoRA output of the overlapped path. maybe_execute_in_parallel now records fn1's outputs itself, outside CUDA graph capture and torch.compile tracing, so there the LoRA record only repeated it. The overlapped path runs only inside the CUDA graph runner, which turns multi-stream on for its eager warm-up steps and for the capture. During the capture the LoRA record still ran, which is what the helper avoids: the caching allocator holds a recorded block freed in the capture until the capture ends, so each overlapped LoRA call took new graph pool memory, and under cudaMallocAsync a recorded output freed inside the capture breaks it. maybe_execute_in_parallel is now the only place the rule is applied. lora_aux_stream is used only through it, so in a capture each overlapped call's fn1 waits on the calling stream (event0) before it allocates, after the previous call's outputs were read. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
|
PR_Github #78150 [ run ] completed with state |
|
Fresh failure in the already-running current-head CI request, separate from the pending pre-commit diagnosis: full head Update, 2026-10-15 17:39 UTC: the restricted testcase index now names @juney-nvidia @schetlur-nv — could the CI owner provide only the sanitized first-failure phase, exception/assertion or relevant repository file/rule/fixture diagnostic for this named case, bound to this exact head, invocation, hardware branch and attempt? Please do not paste raw console, environment, full arguments, stacks or secret-bearing artifacts. A bounded comparison against six unrelated pre-merge runs found only unavailable, absent or skipped records for this case; it does not establish either a shared failure or a B-specific defect. No duplicate CI request is justified. C CI remains after actual B merge, D after C; all merge gates remain in place. Current retry update, 2026-10-15 20:37 UTC: the preceding A100X observations belong to the prior invocation, not the outcome of the one authorized reuse retry. The retry is independently bound to the same full head, acknowledged wrapper 78203, actual L0 64457 and x86 single-GPU child 9459. Both parents and both single-GPU children remain active. The restricted retry index currently lists the A100X streaming-cancellation case as SKIPPED, not PASS or a fresh failure; its report has not established that the authorized failed case was actually rerun. A distinct retry workflow signal now needs the existing CI-owner diagnosis: New terminal test evidence for the sole authorized retry: the one request, acknowledged wrapper 78203 and actual L0 64457 are independently bound to the exact full B head above. L0 is now terminal FAILURE. Its SBSA single-GPU child 9275 is terminal FAILURE, and the x86 child 9459 is terminal ABORTED. Final wrapper metadata reports terminal SUCCESS only for bot orchestration; the exact completed bot acknowledgement explicitly reports L0 FAILURE and the current-head native Restricted complete failed-case metadata now names exactly two real The original cancellation case is still SKIPPED, not recovered, on the requested A100X row. X86 has no named failed-case rows despite the retained FAILED H100 Pytest workflow stage and terminal ABORTED child; no H100 recovery or complete x86 PASS is established. Bounded, lineage-verified comparisons of six unrelated SBSA children establish no shared-regression classification; absent/skipped/active rows do not qualify these two cases. Read-only source diagnosis completed: the two tests, shell/Python helpers, launcher, COMM session, bootstrap and inspected IPC dependencies are byte-identical to actual merged A. Parent independently verified that source comparison. This excludes direct edits in those files, not a changed-dependency/runtime interaction or a B regression. These cases use launcher-managed MPICommExecutor(COMM_WORLD), not the spawned-pool identity-collection path; no necessary source correction is proven. For each case, please provide the first repository-relative failing file:line and exception/assertion class, tied to this exact head, invocation, child, hardware and internal attempt. If only test_mpi_session.py:199 CalledProcessError is available, include its numeric subprocess returncode and issuing layer (inner command, launcher or timeout). Only for confirmed termination without an inner exception, provide the last completed startup/task-result/stop/cleanup phase. An AssertionError at test_mpi_session.py:202 instead identifies cleanup residue after successful shell exit. The whole-job shell timeout is not proof of the observed failure mechanism. No raw logs, environment, arguments, endpoints, keys, stacks or secret-bearing artifacts are requested. The original-attempt evidence requests remain open for cancellation/H100 and pre-commit. No speculative product/test patch, increased timeout/tolerance, exclusion, duplicate writer, repeated sampling or further bot request is authorized. Source equality does not meet the user's condition that the failures be unrelated; no Slack ping has been sent. The user's single queued retry has been consumed. No additional |
|
PR_Github #78171 [ run ] completed with state
|
|
/bot run --reuse-test |
|
PR_Github #78203 [ run ] triggered by Bot. Commit: |
|
PR_Github #78203 [ run ] completed with state
|
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
|
Fixed |
|
/bot run |
|
PR_Github #78386 [ run ] triggered by Bot. Commit: |
|
PR_Github #78386 [ run ] completed with state
|
|
/bot run |
Resolve the l0_b200.yml insertion conflict by keeping both lists: main's three one-model speculative sampling op tests (NVIDIA#20147) followed by the existing spec-step-copy and decode-step-gather entries. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
|
PR_Github #78531 [ run ] triggered by Bot. Commit: |
Resolve the runner.py conflict by taking main's removal of the first-draft input path (NVIDIA#19867), and drop the three references this branch added to the removed API: the first_draft_requests term in the previous-batch offset zeroing condition and the is_first_draft KeyType field in two executor tests. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
|
/bot run |
|
PR_Github #78550 [ run ] triggered by Bot. Commit: |
|
PR_Github #78531 [ run ] completed with state |
|
PR_Github #78550 [ run ] completed with state |
|
/bot run |
|
PR_Github #78564 [ run ] triggered by Bot. Commit: |
|
PR_Github #78564 [ run ] completed with state |
|
/bot run |
|
PR_Github #78628 [ run ] triggered by Bot. Commit: |
|
PR_Github #78628 [ run ] completed with state
|
This branch has not been deployed
13 participants
Description
Depends on #19988 (the main fixes). Review the commits after its head.
These are general-purpose pieces of one-model speculative decoding. All of them are off by default or leave existing
models' behaviour unchanged. They were reviewed as #19816, #19814 and #19817 and are combined at the maintainers'
request.
cute_dsl_kernels/spec_step_copies/StepInputGather/SlotScatter: one launch each for the overlap scheduler's input gathers and for the sampler's store update (SM 100 only). They fall back to the torch path whenever the kernel does not cover a callpyexecutor/engine/runners/,pyexecutor/sampler/,_device_copy.pycopy_to_device_if_changed, checked against the device in debug mode and in unit tests); the sampler's host copies moved to its side stream; CUDA graphs share the engine's input buffersspeculative/interface.py,models/modeling_speculative.pySpecWorkerBase.target_logits: the spec worker picks the target-logits layout. The default is unchanged; the vocabulary-shard contract is documented on the hookpyexecutor/kv_cache/mamba_cache_manager.pyMambaHybridCacheManagerV2(kda_token_states=...): per-draft records (kda_state_tok) in place of the q/k/g, v and beta replay caches, off by default. The replay-state relocation also moves the records (the last two commits of #19846)Test Coverage
The results below are historical, not qualification of the current head.
GB200, on this branch's build, against main on the same tray, compared by JUnit id:
test_spec_step_copies.py,test_copy_to_device_if_changed.py,test_rng_window_counter.py): 93passed;
Tested on main f891b3d.
Limitations:
Maintenance update
Current head:
3261a6fe280620e6b3c93571351811f04ef8d427, a signed-off fast-forward of597e23a47721aaa8fdaf50b93ffaa4738bf9e3d2. Removes the misleading capture explanation, clarifies the strict DSpark validation diagnostic and corrects the inspected 2026 copyright. Capture operations, input checks and existing test assertions are unchanged. Earlier follow-ups add strict layer-ID regression cases to CPU CI and isolate the device-copy helper. Native review is approved and all current conversations are resolved. Current-head pre-commit failed; the failed hook/root cause is not established by safe provider metadata and is under bounded diagnosis. Local compile/lint, package/native/GPU and complete upstream CI remain pending; prior source tests and builds are historical. No required gate is waived.PR Checklist
Dev Engineer Review
The change combines speculative-decoding copy paths, conditional device copies, shared CUDA-graph input buffers, worker-produced target logits, and optional KDA draft-state records. It also adds Speculators DSpark config translation, runtime dependency coverage, and CUDA correctness fixes.
Key behavior to verify:
copy_to_device_if_changedskips transfers when cached host values match. This relies on the destination not being modified outside the helper; debug verification checks that assumption.The author reports 93 passes across three named tests, passing added CI entries, and zero memcheck errors and racecheck hazards for
spec_step_copies. The reportedkv_cacherun had 987 passes and 42 failures; every reported PR failure also occurred on main. The author reports no new GB200 regression failures against matched controls. Current-head upstream CI and codeowner reviews remain required. Review-finding severity counts and current-head test completion are not established by the supplied evidence.QA Engineer Review
The changes add tests for speculative gather/scatter, conditional copies, graph buffers, MNNVL workspace lifetime, multi-stream output lifetime, KDA replay and cache management, DSpark config and mask handling, target logits, and zero-row or zero-token kernels. They also modify attention probe coverage and sampler test setup. CI test lists add coverage for B200, H100, GB200 multi-GPU, and DGX B200. Three waiver entries are removed.
Coverage verdict: needs follow-up. The reported tests provide broad unit and integration coverage, but the author reports no Mamba2 end-to-end run or quota-fix end-to-end run. The reported results do not establish that current-head upstream CI is complete.
Per-File QA Perspective
Source and configuration files
cpp/tensorrt_llm/common/lamportUtils.cuh— Removes cluster-awareLamportFlagsbehavior. Verify CTA arrival counts and updater selection.cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu— Changes Lamport arrivals and RMSNorm barrier order. Verify all-reduce and cluster reductions.cpp/tensorrt_llm/kernels/noAuxTcKernels.cu— Adds a zero-token return and checked launch. Verify empty and non-empty inputs.cpp/tensorrt_llm/kernels/quantization.cu— Adds zero-row returns and checked launches. Verify empty outputs and subsequent non-empty calls.requirements.txt— Addsfla-core==0.5.2for KDA imports. Verify runtime dependency installation.tensorrt_llm/_torch/attention/backends/trtllm.py— Returns early when no context requests exist. Verify empty-context MLA preparation.tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py— Changes CuTe wrapper and output-buffer caching. Verify scratch lifetime, varlen slack, and output slicing.tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py— Adds replay-only decode and pending-draft commit. Verify validation, compile-cache keys, and state replay.tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py— Adds a shared-memory barrier. Verify prefill under timing variation.tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py— Adds replay-only kernel behavior. Verify accepted-token replay and normal decode.tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/__init__.py— Adds module documentation and license text; no runtime behavior change is described.tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py— Adds SM 100 gather/scatter wrappers. Verify argument checks, zero rows, and capture-time compilation errors.tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py— Adds gather/scatter kernels. Verify slot mapping, padding, accepted lengths, and KV offsets.tensorrt_llm/_torch/distributed/ops.py— Retains captured MNNVL workspaces and rejects growth during capture. Verify replay after workspace growth.tensorrt_llm/_torch/models/modeling_dspark.py— Adds configured mask-token and checkpoint embedding-row handling. Verify valid, missing, and malformed rows.tensorrt_llm/_torch/models/modeling_speculative.py— Gets target logits through the spec worker. Verify default and worker-specific behavior.tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py— Adds KDA draft-commit helpers. Verify accepted-token clamping and recurrent/convolution state updates.tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py— Uses conditional state-index copies. Verify changed and unchanged host indices.tensorrt_llm/_torch/modules/multi_stream_utils.py— Records auxiliary-stream outputs on the caller stream. Verify nested outputs, capture, compile, and allocator-error paths.tensorrt_llm/_torch/peft/lora/layer.py— Removes LoRA outputrecord_streamcalls. Verify output lifetime during parallel execution.tensorrt_llm/_torch/pyexecutor/_util.py— Enables KDA token-state storage for the applicable cache manager. Verify manager selection.tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py— Adds engine-owned static input buffers. Verify validation, capture sharing, and replay with changed pointers.tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py— Adds speculative gathers, conditional copies, and vocabulary-shard post-processing. Verify fallback behavior and attention-DP rejection.tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py— Clears staged batch indices after copying. Verify reuse across batches.tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py— Adds optional KDA token records and lifecycle handling. Verify sizing, quota, relocation, suspend/resume, and seeding.tensorrt_llm/_torch/pyexecutor/py_executor.py— RunsSpecSampleron the execution stream. Verify stream ordering and non-spec sampler behavior.tensorrt_llm/_torch/speculative/dflash.py— Adds conditional copies, mask-embedding construction, and bounded context advancement. Verify mapping updates and dummy-slot behavior.tensorrt_llm/_torch/speculative/interface.py— Removes accepted-draft counts from metadata and adds target-logit handling. Verify callers and all-greedy RNG behavior.tensorrt_llm/_torch/speculative/spec_sampler_base.py— Adds optional scatter and side-stream host copies. Verify fallback widths and synchronization.tensorrt_llm/_utils.py— Addscopy_to_device_if_changed. Verify destination-write assumptions, staging, and capture-time debug checks.tensorrt_llm/_torch/pyexecutor/config_utils.py— Translates Speculators DSpark configs during loading. Verify detection and Qwen3 construction.tensorrt_llm/llmapi/llm_args.py— Adds DSpark config translation and checkpoint resolution. Verify pass-through behavior and validation errors.Test and test-list files
tests/integration/test_lists/test-db/l0_b200.yml— Adds multi-stream, step-copy, and engine-gather tests to CI.tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml— Adds MNNVL captured-workspace growth coverage to GB200 CI.tests/integration/test_lists/test-db/l0_h100.yml— Adds the multi-stream utility test to H100 CI.tests/integration/test_lists/test-db/l0_dgx_b200.yml— Adds the MNNVL capture-guard test to DGX B200 CI.tests/integration/test_lists/waives.txt— Removes three MoE and quantization waivers.tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py— Covers gather/scatter correctness, repeatability, unsupported inputs, and graph replay. Listed intest-db/l0_b200.yml.tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py— Covers KDA record allocation, quota, relocation, suspend/resume, and replay. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/executor/test_copy_to_device_if_changed.py— Covers copy skips, resizing, flattening, and debug checks. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py— Covers shared static buffers and graph-visible rebinding. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/executor/test_pytorch_model_engine.py— Compares speculative gather with the Torch path. Listed intest-db/l0_b200.yml.tests/unittest/_torch/helpers.py— Extends the graph-runner test helper; no test-list change.tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py— Removes an unsafe context-length probe and documents the risk. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py— Covers scratch-cache eviction and memory bounds. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py— Covers pending-draft commit and state replay parity. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py— Covers timing perturbation and input reuse in prefill. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py— Covers skipped and changed state-index uploads. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/modules/test_multi_stream_utils.py— Covers stream lifetimes, capture, compile, and allocator behavior. Listed intest-db/l0_b200.ymlandtest-db/l0_h100.yml.tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py— Covers workspace growth and capture guards. Related tests are listed intest-db/l0_gb200_multi_gpus.ymlandtest-db/l0_dgx_b200.yml.tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py— Covers DFlash mapping updates and skipped copies. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py— Covers dummy-slot lengths, including graph replay. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py— Covers mask-row selection and validation. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/speculative/hw_agnostic/test_rng_window_counter.py— Covers all-greedy RNG buffers. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/speculative/hw_agnostic/test_spec_worker_target_logits.py— Covers default logits, shard gathering, padding removal, and attention-DP rejection. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/thop/parallel/test_fp4_gemm_quantize.py— Covers zero-row FP4 quantization and a subsequent non-empty call. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/thop/parallel/test_fp8_quantize.py— Covers zero-row MXFP8 quantization in two layouts. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/thop/parallel_hw_agnostic/test_noaux_tc.py— Covers zero-token routing and a subsequent non-empty call. No changed CI or manual-QA list entry is supplied.tests/unittest/conftest.py— Adds autouse skipped-copy verification. Verify it does not cause test-order or buffer-state issues.tests/unittest/others/test_kimi_kda_runtime_dependencies.py— Checks KDA imports against runtime requirements. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py— Covers DSpark config translation, draft lengths, invalid inputs, and unrelated configs. No changed CI or manual-QA list entry is supplied.tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py— Updates sampler stubs for async-worker and store-copy state. Verify the stubs exercise the expected host-copy path.Historical main-conflict repair validation
At
951048b175, native builds and affected CPU gates passed; remaining collective, 16-rank functional and model-serving comparisons were pending. This does not qualify the current head.