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>
…next chunk fused_kernel123's K1 warps read their scanned prefix row and the chunk total from sPartialLast after the scan barrier, and the next chunk's partial sums overwrite sPartialLast before any K1 barrier (racecheck: write at fused_k123.py:1268 against reads at :1293 / :1299). A K1 warp that reached its reads late used the next chunk's sums: with one K1 warp stalled 20 us after the scan barrier, every case of test_fused_prefill_beta_sigmoid_matches_unfused_kernel changed (output max |d| up to 7e14). One K1 barrier after the reads orders them; outputs are bit-identical to before. test_fused_prefill_does_not_depend_on_k1_warp_timing runs a copy of the kernel with that stall and requires the unpatched kernel's bits; it fails without the barrier. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…ce stores In block_reduce_sum2_for and block_reduce_sum2_active_for, lanes of warp 0 read the per-warp partials scratch[lane] / scratch[kReduceWarps + lane], reduce them with shuffles, and lane 0 then stores the block totals into scratch[0] and scratch[1], which lane 1 read. Shuffles do not order shared memory, so the store and lane 1's read were unordered (compute-sanitizer racecheck: 32 warnings per run of test_k3_kda_decode_attn, all in kda_decode_fusion_compact_heads_kernel). Lane 0's totals depend on lane 1's loaded value, so the results were not affected; __syncwarp() makes the order explicit. 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>
…e the lane-0 block-reduce stores) The ssm/kda_decode entry's receipts run the fixed kernel. NVIDIA#19817 (this branch's base) and NVIDIA#19830 touch disjoint files. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…op tests The Kimi K3 decode kernels for the KDA and MLA layers, as of the K3 stack 82a110a92a, verbatim: - k3_kda_attn/: trtllm::k3_kda_qkvg (the KDA projection as one CTM kernel), trtllm::k3_kda_attn (projection and speculative verify of one request's 8 tokens in one launch) and trtllm::k3_kda_decode_attn (projection and T = 1 decode of up to 8 requests); - k3_kda_verify/: trtllm::k3_kda_verify (the KDA verify of a step's drafts over the per-token states of the V2 hybrid cache manager); - k3_mla/: trtllm::k3_mla_q / k3_mla_qkv / k3_mla_qkv_out (the MLA query path and the latent KV append) and trtllm::k3_mla_attn / _attn_out / _attn_vb_out (decode attention over the paged latent cache, optionally with v_b and the output gate); - k3_mla_decode_view in the CuTe DSL MLA backend: the attention metadata's view of R generation requests of T tokens that k3_mla_attn takes. The op tests come along; they carry the regression tests of the fixes made to these kernels: Lamport buffer index in 0..2 and launch counters across the int32 wrap, state pools addressed at 64-bit slot offsets (pools past 2 GiB), slot indices at any element offset, one shared Lamport set for the decode and verify launches, k3_mla_attn reading only the workspace words its own call wrote, and the no-cluster counters across the wrap. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
trtllm::k3_mla_attn, k3_mla_attn_out and k3_mla_attn_vb_out kept their workspace (the per-CTA partials and the no-cluster mode's exchange and arrival counters) in a module-level dict keyed by device and head-group count, allocated on the first eager call. They now take it as an argument, named in mutates_args: make_attn_workspace(device, groups) allocates and arms one (the counters zeroed, as before) and refuses to run under CUDA-graph capture; a call checks that the workspace is the one for its device and head-group count and raises ValueError before launching otherwise. The kernel and the words it reads and writes are unchanged. The op test passes a workspace; the workspace poison and counter wrap tests now run on their own explicit workspaces, and a malformed workspace is refused without being touched. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…de entries New category ssm/ (KDA decode) and two attention/ entries, each a contract, a wrapper and a GPU test that drives the op on a real cache manager: - ssm/kda_decode: the one-token KDA decode (C++), on the V2 hybrid manager's pools; - ssm/k3_kda_attn (+ k3_kda_qkvg), ssm/k3_kda_decode_attn: Kimi K3's fused KDA projection with the verify or the plain decode in one launch, over a caller-owned K3KdaBuffers (the projection's Lamport set, one per device, shared by both ops); - ssm/k3_kda_verify: the KDA verify from the per-token states that the V2 hybrid manager keeps per slot; - attention/k3_mla_qkv (+ k3_mla_q, k3_mla_qkv_out): the MLA decode query path and the latent KV append into the paged cache; - attention/k3_mla_attn_vb_out (+ k3_mla_attn_out, k3_mla_attn): MLA decode attention with v_b and the output gate, over a caller-owned K3MlaAttnWorkspace. Every stateful contract has a ## State section (contents and size, creator, sharing, call order, what a later launch reads, re-arm), and every test runs layers x steps on one state object, a CUDA-graph capture replayed with rewritten inputs, two objects interleaved and a negative control. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
ruff-format and ruff's import sort on the two files the K3 stack wrote in its 120-column style (k3_kda_attn_kernel.py, k3_mla/op.py). The AST of every Python file of this PR is unchanged; the full pre-commit passes on all of them. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…uction flags
kda_decode checks apply_onorm, use_lower_bound and apply_beta_sigmoid
before anything else and refuses any of them off, on every architecture
and batch size ("KDA decode only supports apply_onorm=true,
use_lower_bound=true, and apply_beta_sigmoid=true"). The contract
described the off branches and a kernel-level refusal on sm_100 / sm_103
only; its semantics now show the one supported combination, and the test
expects the op's message on every architecture.
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…negative controls The entries' tests ran on SM 10.x; their receipts are sm_100's, so they now skip on any other architecture (a missing receipt reads as unknown). The KDA contracts state what their negative controls measured on sm_100: two steps or rounds of one request in swapped order leave the second output and the slot's state off by 0.8-1.4 of their largest magnitude, with nothing raised; and k3_kda_qkvg's rows are within 2.5e-3 of a float64 projection. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…tries Receipts of the six entries on GB200 (sm_100), from one run of every test file of this PR at its previous commit: kda_decode 11, k3_kda_attn 8, k3_kda_decode_attn 14, k3_kda_verify 7, k3_mla_qkv 8, k3_mla_attn_vb_out 9 tests passed. The op tests and the entry tests join l0_b200's pre-merge list; each file ran in under a minute. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
…V, head GEMV, embedding The single-GPU decode kernels of the Kimi K3 port, with their op tests: - k3_ctm_gemv: trtllm::k3_ctm_gemv, _swiglu, _long, _wide (bf16 GEMVs for up to 8 tokens, 64 for _wide, on tcgen05) and trtllm::k3_situ_mul, the dense MLP's SiTU-and-mul; - k3_decode_gemv: trtllm::k3_decode_gemv and its row-parallel MoE tail; - k3_head_gemv: trtllm::k3_head_gemv, the vocabulary-shard head GEMV on a persistent stream-K kernel of min(SMs, tiles x k-tiles) CTAs; - k3_embed: trtllm::k3_embed and k3_embed_norm (the decode step's embedding and layer 0's input RMSNorm in one launch). Two source lints read these kernels: test_k3_cluster_waits (a wait on a mailbox that other CTAs complete with st.async acquires at cluster scope) and test_k3_tcgen05_fences (the tcgen05 fences around thread syncs). Every kernel and test file is as on the K3 stack's tip; the decode GEMV test checks its tail against a torch reference. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
… after the grid wait The decode kernels behind attn_res_rmsnorm_fwd and attn_res_add_rmsnorm_fwd (the single-CTA and split-K s1 kernels that launchAttnResDecodeRmsNorm starts) now trigger programmatic dependent launch right after their own grid-dependency wait instead of at their end. Under PDL the next kernel then launches and streams its weights while they run. It still waits for this grid to complete before it reads the output, so results do not change. The other launches of the s1 kernels keep the trigger at the end. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
trtllm::k3_head_gemv kept its cross-launch state in a module-level dict keyed by device and size, allocated by the first call: the stream-K partials and flag words, or the dynamic schedule's per-tile counts and unit counter. That state is now a K3HeadGemvWorkspace that the caller creates once per weight shape and schedule, eagerly (create() refuses CUDA-graph capture), and passes to every call. The op takes the workspace's three buffers as arguments, names them in mutates_args, and refuses a workspace whose buffers do not fit the call. The op tests add two workspaces used alternately, a captured graph replayed between eager calls on one workspace, a workspace of another shape (refused) and create() under capture (refused), and check that every call leaves the words at zero. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Ten entries for the single-GPU Kimi K3 decode ops, each a contract, a one-call wrapper and a GPU test: - gemm: k3_ctm_gemv, k3_ctm_gemv_swiglu, k3_ctm_gemv_long, k3_ctm_gemv_wide, k3_decode_gemv, and k3_head_gemv over a caller-owned K3HeadGemvWorkspace (its contract has a State section); - activation: k3_situ_mul; - norm: k3_embed_norm, attn_res_fwd, attn_res_rmsnorm_fwd. They are certified on sm_100 (B200 / GB200), where their first caller runs; their tests skip on other architectures. l0_b200 lists them, the kernels' op tests and lints, and the cublas_mm and flashinfer_rmsnorm catalog tests, whose entries gain sm_100 receipts. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Collaborator
Author
|
Published additive main/conflict repairs at |
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Collaborator
Author
|
Published review fixes at |
1 task done
Collaborator
Author
|
@litaotju @yuxianq @Shixiaowei02 — could you prioritize the remaining current-head reviews so this can follow #19989 without another review wait? Current head is |
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Collaborator
Author
|
Published |
Collaborator
|
[by Claude Code] @eopXD Could you review this PR? Thanks! |
Resolve the shared MoE package initializer using the upstream version. Preserve published history and both catalog families. Local runtime qualification remains pending. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Collaborator
Author
|
Published strict fast-forward head Publication uses the user's structurally verified publish-before-testing permission, not a merge or validation waiver. Worktree cleanliness, stack ancestry, scoped ownership/signoffs, strict fast-forward and clean merge-tree against main |
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Collaborator
Author
|
Fixed |
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>
Merge the K3 runtime branch after it absorbed main dac84b0. No conflicts; l0_b200.yml keeps main's speculative op tests ahead of the inherited spec-step-copy entries. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
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>
Merge the K3 runtime branch after it absorbed main c226e2d and the serving-schema refresh. No conflicts. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>
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
9 participants
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.
Description
Depends on #19989. Review the commits after its head.
Kimi K3's single-GPU kernels and their modeling_v2 catalog entries, at the TP16 / TP4 per-rank shapes on sm_100.
Nothing on an existing model's path calls them. They were reviewed as #19833, #19832 and #19834 and are combined at
the maintainers' request.
cute_dsl_kernels/k3_ctm_gemv,k3_decode_gemv,k3_head_gemv,k3_embed;kimiK3AttnRes/attnResFwd.cu(PDL release)gemm/,norm/,activation/cute_dsl_kernels/k3_kda_attn,k3_kda_verify,k3_mlassm/andattention/ssm/,norm/,activation/,moe/; sm_100 cells of existing entriescatalog/index.yaml,README.mdTest Coverage
The results below are historical, not qualification of the current head.
GB200, on this branch's build:
l0_b2001558,l0_cpu29);_torch/modeling_v2(withoutcomm): 417 passed, 0 failed. Main fails 13 of these, all of which this stackfixes;
test_k3_kda_attn,test_k3_kda_verify,test_k3_kda_decode_attn), runon the final stack: memcheck 0 errors; racecheck and synccheck show no finding in the changed KDA kernels.
Other records are classified in the local evidence ledger.
Tested on main f891b3d.
KDA initcheck coverage is partial; the remaining checks and a reference-code read are under investigation.
Maintenance update
Earlier validated head:
3ac75a7d464e44642642b65da77821c45f35a993. Merged the tested #19989 update, preserving published history. The exact-head build and installed-source verification passed; the GB200 regression comparison found no new failures versus its matched controls. Upstream CI and codeowner reviews remain required, after #19989 merges.PR Checklist
Dev Engineer Review
The supplied change summary describes a broad Kimi K3 kernel and
modeling_v2catalog addition. The inspected checkout does not match that change set: its working-tree diff contains only formatting changes inpy_executor.py,speculative/interface.py, andllm_args.py. I cannot verify the described kernel, catalog, runtime, or API changes against this checkout.QA Engineer Review
The inspected working-tree diff contains no test-code or test-list changes. The supplied summary describes extensive Kimi K3 tests and CI-list additions, but I cannot verify those files or their test-list mappings in this checkout. Coverage verdict: needs follow-up.
Per-File QA Perspective
tensorrt_llm/_torch/pyexecutor/py_executor.py: Import formatting only; no behavior change is visible. No test change is present in the inspected diff.tensorrt_llm/_torch/speculative/interface.py: Import formatting only; no behavior change is visible. No test change is present in the inspected diff.tensorrt_llm/llmapi/llm_args.py: Import formatting only; no behavior change is visible. No test change is present in the inspected diff.Current main-conflict repair validation
Published
1cf3ae64e7. Native builds and affected CPU gates passed; remaining collective, 16-rank functional and model-serving comparisons are pending. Local qualification continues in parallel with CI; no pending gate is claimed passed.Main-conflict repair
Current head:
97d8936c445108fb7f3e4fc65eb09e986d943dd5, a signed-off fast-forward off1d74e3d0f0b980ef0c91882a6cc004df13b2ab6inheriting #19989's capture-comment, validation-diagnostic and copyright nits only. The earlier main repair retains Qwen/catalog additions and upstream's shared MoE initializer. Current-head pre-commit failed; failed-hook/root-cause metadata is not established. Local compile/lint, full routing, package/native and GPU checks remain pending; previous validation above is historical. Full CI admission remains after #19989 actually merges. Required review/conversation/CI gates are unchanged.