Repository navigation
Conversation
Route GlmMoeDsaForCausalLM to models/glm_moe_dsa: GLM-5.2-NVFP4 on sm107 with 16-way tensor parallelism, no attention DP and one-model MTP-Eagle (any draft length) or no speculation. Add Arch.SM_107 for entries certified on Rubin. The target's modeling and weights follow. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
The FMHA autotuner sets tileSizeQ of the token-sparse MLA generation kernels to the Q heads per KV head. With fewer than 8 heads per rank (64 MLA heads over 16 ranks: 4) the SwapsMmaAb kernel then gets a Q tile that SmemTile rejects when it is JIT-compiled. Pad the tile to 8 rows on every SM, as the dense MLA selection does; the kernel masks the rows past the heads. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
When the TP group spans nodes, _compute_mlp_tp_size capped the dense MLP at gpus_per_node, while the fused post-MLP all-reduce sums the MLP partials over the whole TP group: each partial was counted tp_size / mlp_tp_size times (2x at TP8 and 4x at TP16 on 4-GPU nodes). Cap the MLP TP only when the MLP reduces within its own group (eager fusion disabled); with the fusion on, the MLP keeps the full TP. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
The target builds the built-in GLM-5.2 modules and loads the checkpoint with the built-in loader, so the MLA KV cache and DSA indexer K cache it reads have the layout the context server wrote. Forwards with context rows run that decoder stack in PrefillTarget; generation-only forwards run DecodeTarget, which starts from the same stack. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
The route matches GLM-5.2-NVFP4 on sm107 at tp16 without attention DP or expert parallelism, with one-model MTP-Eagle at any draft length or no speculation, and declines another shape, attention DP, expert parallelism, tp8, vanilla MTP, NGram and sm103, naming the criterion that missed. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
DecodeTarget runs the embedding, the 78 decoder layers and the final norm as its own code instead of calling the built-in model's forward: the dense layers' attention all-reduce, residual norm, MLP and fused post-MLP all-reduce, and the MoE layers' fused post-attention all-reduce, experts and fused post-MoE all-reduce. The built-in modules still compute each part. The configuration the step is written for (fusions, MLP TP, multi-node experts, unquantized kv_a) is checked once when the target is built. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
torch.ops.trtllm.glm_dense_mlp_small_m runs one rank's slice of a BF16 SwiGLU dense MLP for 1-8 rows as two weight-streaming kernels (kernels/glmMoeDsa/glmDenseMlpSmallM.cu): gate/up + SwiGLU on 192 CTAs in two-CTA clusters that split the 6144 contraction, then down on 96 CTAs of 64 output columns. The compiled tilings serve hidden 6144 with a rank-local intermediate size of 768, GLM-5.2's dense layers 0-2 at TP16. The op takes x [M, 6144], gate_up_weight [2 * 768, 6144] (this rank's gate rows over its up rows) and down_weight [6144, 768], and returns the rank-local partial [M, 6144] that the post-MLP all-reduce sums. The gate and up sums round to BF16 before SwiGLU and h rounds to BF16 before down, the rounding points of the unfused BF16 chain; the sums are FP32 in a fixed order, so repeats are bitwise equal. The sigmoid takes the fast exp and division. Under programmatic dependent launch every CTA bulk-copies its weights into shared memory before griddepcontrol.wait and reads its activations only after it. gate/up releases down at entry; down releases its dependents once its weight copies are issued. torch.ops.trtllm.glm_dense_mlp_small_m_supported(rows, hidden, inter) answers on the host whether the kernels serve a shape; the op raises on any other shape, layout, dtype or a device below sm_90. The target's entry is ops/gemm/dense_mlp_small_m.py, certified on sm_107 by its test against an FP32-sum reference with the same rounding points. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
torch.ops.trtllm.glm_fused_embedding(input, weight) looks up int32 or int64 token ids of any shape in an unsharded BF16 table [vocab, hidden] with one kernel (kernels/glmMoeDsa/glmFusedEmbeddingLookup.cu): one thread per 16-byte vector, a grid of tokens x hidden / 1024 CTAs of 128 threads. An id outside [0, vocab), such as the -1 that pads rejected draft tokens, gives a zero row, the guard of the Embedding module's unsharded path. Rows are bit copies of the table, so the output equals the module's masked F.embedding chain bit for bit. Under programmatic dependent launch the kernel waits for its producer before it reads the ids and releases its dependents right after the wait; it touches no memory before it. The GLM-5.2 TP16 target keeps the [154880, 6144] table whole on every rank, so the lookup needs no collective. Its entry is ops/embedding/fused_embedding.py, certified bitwise on sm_107 by its test. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
The verify forward looks its ids up in the whole BF16 table with one kernel, and the dense layers 0-2 run their 768 intermediate columns per rank as two weight-streaming kernels for up to 8 rows; the post-MLP all-reduce sums the rank partials as before. A batch the dense kernels do not serve takes the generic gated MLP. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
kernels/glmMoeDsa/glmFusedAGemm.cu computes kv_a_proj_with_mqa for GLM-5 at hidden 6144 and up to 8 tokens: 2624 output rows (q_a, compressed KV, k_pe), or 2784 with the DSA indexer's wk and weights_proj rows appended to the same weight. One CTA per 16 weight rows (164 or 174 CTAs, one per SM), and all eight warps issue the HMMAs: warp w owns the 256-wide k stages w, w + 8 and w + 16. Before the grid-dependency wait each CTA stages its weights by tensor copies: stage w into shared memory, stages w + 8 and w + 16 through shared memory into registers, warp w issuing its register-stage copies w * 250 ns after its first copy. Then it prefetches up to four caller tensors (weights that later kernels of the step read) into L2 in 16 KB bulk prefetches spread over the CTAs. After the wait it copies the activation tiles and triggers its dependents. The FP32 sums of the eight warps are added in warp order and rounded to BF16 once. torch.ops.trtllm.glm_fused_a_gemm(hidden_states, weight, l2_prefetch) takes the weight as stored and raises on shapes the kernel does not run. The modeling_v2 entry ops/gemm/fused_a_gemm.py of the glm_5_2_nvfp4 sm_107 tp16 target states the shape gates in is_valid; its sm_107 test drives the c1 cells and checks the 2624 columns of the 2784-row call, agreement with dsv3_fused_a_gemm_op, CUDA graph replay and the op's rejections. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
kernels/glmMoeDsa/glmRouterGemm.cu computes the FP32 router logits of GLM-5's 256 experts at hidden 6144 for 1 to 16 tokens, one CTA per expert, with the loads and FMA order of the shared DeepSeek-V3 router GEMV, so the logits are bitwise those of dsv3_router_gemm_op. Before its grid-dependency wait each CTA prefetches its 12 KB weight row into L2, and it triggers its dependents right after the wait, before any logit is written: a kernel that follows it on the stream waits on its own grid dependency before it reads the logits. torch.ops.trtllm.glm_router_gemm(hidden_states, weight) takes the gate weight as stored and raises on shapes the kernel does not run. The modeling_v2 entry ops/gemm/router_gemm.py of the glm_5_2_nvfp4 sm_107 tp16 target states the shape gates in is_valid; its sm_107 test drives the cells against an FP64 reference and checks bitwise agreement with dsv3_router_gemm_op at every token count, CUDA graph replay and the op's rejections. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…-A GEMM The decode step runs MLA with the DSA indexer itself: q_a, the compressed KV and its rope part come from the GLM fused-A GEMM over the stored kv_a_proj_with_mqa weight for up to 8 rows, and the built-in norms, q_b, indexer projection, K-cache update, sparse generation attention and o_proj compute the rest in the order the built-in DSA path runs them. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
The target carries its own copy of the low-M BF16 split-K CuTe DSL GEMM (ops/gemm/_low_m_splitk.py, adapted from FlashInfer's dense_bf16_gemm_sm100_splitk) with the occupancy policy that refines the source default tactic to a cluster of eight K slices, the Lamport publish epilogue and the in-kernel RMSNorm B producer with an FP32 output, plus the fused-epilogue kernels (ops/gemm/_low_m_splitk_fused.py): the row-RMS and FP8 RoPE q_b epilogue and the DSA indexer Q tail. Entries, one per call the decode step makes, all certified on sm_107: - low_m_gemm: x @ weight.T for the shared expert's gate_up and down, the LM head shard and MTP draft step 0's eh_proj. - q_b_fp8_rope: q_b with k_b absorbed and the q_a RMSNorm folded into the weight, writing the RoPE'd E4M3 attention Q; regular, beside-top-k and early-gather launch plans. - indexer_q_tail: the indexer wq_b from the raw q_a with the prologue's Q tail (RoPE, UE8M0 scales, E4M3 Q, scaled head weights) in the epilogue. - o_proj_publish: o_proj storing its BF16 output into this rank's slot of the post-attention MNNVL one-shot all-reduce's current Lamport buffer. - draft_hidden_half: this rank's FP32 share of the next draft step's eh_proj hidden half, hnorm computed in the kernel for one row. Each entry has a GPU test over its cells with a derived band, a CUDA graph replay check where the target captures it, and guard tests for the inputs is_valid refuses. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
fp4_block_scale_moe_runner takes an optional hidden_states_ready_event: the raw cudaEvent_t (torch.cuda.Event.cuda_event) of an event recorded on the stream that writes hidden_states and hidden_states_scale. MoE::Runner::run waits on it with cudaStreamWaitEvent right before the FC1 launch, the first kernel that reads them, so the routing kernel launched before it depends on the router alone while the input quantization runs on another stream. Without the argument nothing waits and the op launches what it launched before. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…equest fp4_block_scale_moe_runner takes optional router_hidden_states (bf16 [T, H]), router_weight (bf16 [E, H]) and router_workspace. With them, DeepSeek-V3 routing without groups runs routingFusedRouterKernel instead of the routing kernel: 128 CTAs compute the logits of dsv3_router_gemm_op bit for bit (the same per-thread FMA chains and butterfly order) and publish slice sums to a workspace exchange buffer whose idle value is -0.0; one CTA polls them, runs the small-token kernel's routing (sigmoid + bias, top-8, scaled sum normalization, same operations in the same order) and writes every output that kernel writes. routing_logits then only gives the shape. The kernel covers 256 experts, top-8, hidden 6144 and 1-8 tokens and refuses other shapes. make_fused_router_workspace builds the workspace; two exchange buffers alternate by an epoch word, so launches that share a workspace must run in stream order. Without the router arguments the op routes as before. PDL: before its wait the kernel reads only the router weight and the routing bias. The routing CTA triggers its dependents after all routing outputs are written; the GEMV CTAs exit without triggering. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
With prefer_gemm2_k128 the NVFP4 MoE runs FC2 on moe::dev::gemm2K128 where it covers the call: E2M1 x E2M1 without DeepSeek FP8 or fine-grained sync, tile N 8, intermediate size 128 without padding, no bias or fused shared experts, at most 64 token tiles and at most 15 units (token tile x 128 output features) per SM. At that K the Gemm2 cubins spend half of their 256-wide K tile on zero padding. gemm2K128 runs one persistent CTA per SM over a contiguous run of units: weights and their block scales by TMA before the grid-dependency wait, the FC1 output after it, two tcgen05 mxf4nvf4 block16 MMAs (M 128, N 8, K 64) per unit, then the cubin's FP32 output scale and BF16 rounding, so the rows are the cubin's bitwise. It reads the routing outputs before its wait (its primary, FC1, triggers only after its own wait on routing), triggers its dependents right after the wait and prefers the maximum shared-memory carveout. Everywhere else, and without the flag, the Gemm2 cubin runs; the autotuner keys a flagged runner apart. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
ops/moe/routed_experts wraps the fp4_block_scale_moe_runner call the target makes in its MoE layers, every argument explicit: DeepSeek-V3 routing without groups over all 256 experts at MoE TP16 (intermediate 128), with the router GEMV inside the routing kernel at up to 8 rows, the NVFP4 input joined right before FC1 and FC2 on the K = 128 kernel, unfinalized for the post-MoE hop. Cells: the 8-row verify step, one row, and 32 rows routed from the router GEMM's logits. The reference is a plain-torch NVFP4 dequant MoE; is_valid refuses strided operands, an unrecorded event and a router workspace not in its between-launch state. The sm_107-only test drives the cells against the reference, with the input quantized on a side stream behind a sleep, and proves the fused router bitwise against the router GEMM at 1-8 rows, the K = 128 FC2 bitwise against the cubin, the three-kernel launch and CUDA-graph replay. The catalog fp4_quantize entry gains the target's 8-row cell at hidden 6144, its routed-expert input. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…lit-K ops Up to 8 rows, each MoE block runs its shared expert as two split-K GEMMs around SwiGLU on the aux stream and its routed experts as one trtllm-gen call on the main stream: the routing kernel computes the router GEMV, the NVFP4 input quantizes on its own stream and joins right before FC1, and FC2 runs on the K = 128 kernel. The decode target reads the model's dimensions off the pretrained config (the core has no config property). Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
glm_mla_kv_half_append (kernels/glmMoeDsa/glmMlaKvHalfAppend.cu, thop/glmMlaKvHalfAppendOp.cpp) runs the KV half of an MLA generation step in one kernel: RMSNorm of the 512 latent columns of the raw kv_a_proj output, GPT-J RoPE of the 64 K-RoPE columns, E4M3 quantization at unit scale straight from fp32, and the append to the paged FP8 cache through the block offsets, pool pointers and pool mapping that mla_rope_generation takes. Block 0 also writes the generation attention's scales. With trigger_at_entry its PDL dependents may launch at its entry; they read its outputs after their own grid-dependency wait. The modeling_v2 entry ops/attention/mla_kv_half_append.py states the op, its reference and its cells: the verify and MTP draft step 0 rows of both fused-A widths, a draft-step row and two sequences. Its sm_107 test reads the appended rows back where convert_req_index_to_global places their positions on a DSACacheManagerV2 pool. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
glm_sparse_mla_flash_decode (kernels/glmMoeDsa/glmSparseMlaFlashDecode.cu, thop/glmSparseMlaFlashDecodeOp.cpp) runs the DSA sparse MLA generation attention of 4 query heads per rank as flash decoding over the gathered latent rows of a top-k of up to 2048: one cluster of 16 CTAs per query token, FP16 mma.sync on exactly unpacked E4M3 operands, P in E4M3 as the TRTLLM-Gen generation kernels feed their second GEMM, partial outputs pushed into the merging CTA during P.V, and the v_b BMM folded into the [tokens, 4 * 256] output. With a block table the kernel maps top-k positions to pool rows itself, with convert_req_index_to_global's arithmetic; with gather_before_wait it also reads the indices, the table, the KV lengths and the rows of earlier steps before its grid-dependency wait, for layers that share an earlier layer's top-k. The modeling_v2 entry ops/attention/sparse_mla_flash_decode.py states the op, its reference (with the E4M3 rounding of P per 128-slot split) and its cells: the verify full-indexer and shared-indexer layers, MTP draft step 0 and steps 1-6, a short context and two sequences. Its sm_107 test also replays the op from a captured CUDA graph and runs both GLM attention ops on a DSACacheManagerV2 pool against convert_req_index_to_global. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…flash-decode ops For up to 8 rows the decode step's attention runs q_b with k_b absorbed and the q_a norm folded (built once from the loaded weights), whose epilogue writes the RoPE'd FP8 query, beside the KV half on the aux stream (kv_a norm, K RoPE, FP8 append and the attention scales in one kernel), and then the sparse MLA flash decode at 4 heads with v_b folded in over the built-in indexer's selection. The built-in indexer projection, K-cache update and top-k selection are unchanged. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
torch.ops.trtllm.glm_lowm_bf16_moe runs one token through an MoE layer with unquantized BF16 SwiGLU experts held as tensor-parallel slices and one-group DeepSeek-V3 routing, in two kernels (glmMoeDsa/glmLowMBf16Moe.cu): - FC1 + SwiGLU, one CTA per SM: a producer warp routes from the router logits (noaux_tc_op's experts, order and tie-break) and streams whole weight rows of the routed experts through a shared-memory ring that eight consumer warps read two rows at a time; it releases its dependents right after its grid-dependency wait. - FC2 + the weighted expert sum: routes again before its own wait from bulk-copied logits and bias, stages each expert's weight slice in shared memory beside the FC1 kernel, forms noaux_tc_op's routing weights in FP32 and falls back to its FP64 expression only where the FP32 result is not certain. The BF16 rounding points are those of the CUTLASS MoE path (FC1 output, SwiGLU output, each expert's scaled FC2 output); the expert sum is FP32 in the routing order. 256 experts, top 8, hidden sizes that are multiples of 2048 up to 8192 and expert slices that are multiples of 64 up to 512; anything else, and more than one token, is refused. The tp16 GLM target's owned entry ops/moe/lowm_bf16_moe certifies the MTP layer's draft-step shape (hidden 6144, slice 128) over random, tied and tiny-score routing on sm107, plus a CUDA-graph replay and the refusals. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
torch.ops.trtllm.glm_verify_head_accept computes one verify step's accepted-token counts in one CTA (glmMoeDsa/glmVerifyHeadAccept.cu): a context row accepts its sampled token; a generation row accepts its target token plus the leading draft tokens equal to the targets, and the forced count of TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS then replaces it as the one-model spec worker forms it: the integral part, or that plus one when the pool entry (counter * counter_stride + slot * slot_stride) & (pool_size - 1) is below the fraction (compared in FP32), capped at draft_len + 1, with the counter advanced once per call. With an accepted_tokens buffer it also writes the accepted rows; otherwise the caller reads them from the sampled tokens. int32 and int64 targets, padded (non-ragged) layout. Under PDL it waits for the sampler before reading anything. The tp16 GLM target's owned entry ops/spec/verify_head_accept certifies the c1 verify step (K 7, forced 2.91) and the strict, integral-force, mixed and multi-request forms on sm107 against the ATen chain; its test also drives SpecWorkerBase._apply_force_accepted_tokens on the same pool and counter over consecutive steps, and CUDA-graph replays. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
ops/spec/gather_last_accepted_rows selects, for each request r, row g[r] = batch_indices[r] * rows_per_request + num_accepted[r] - 1 + row_offset (its last accepted token) of one or two row tensors and returns g as int64 with the gathered rows (_gather_last_accepted_rows._gather_last_accepted_rows_kernel). Draft step 0 of the tp16 GLM target's MTP loop uses it twice over: to take the target's hidden rows and gather_ids in one launch, and, after the MTP layer's post-attention all-reduce, to keep only the selected rows of the normalized hidden states and the residual so that the MoE, its all-reduce and the draft LM head run on one row per request. The entry certifies the c1 shapes (8 verified rows of 6144, with and without the residual) and a permuted 4-request batch bit for bit on sm107; its guard refuses counts outside [1, rows_per_request], strided index vectors and a second tensor of another width. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
ops/spec/draft_glue_rmsnorm_cat returns this rank's window [col_lo, col_hi) of cat([enorm(embed_table[input_ids]), hnorm(hidden_states)], -1) as one contiguous BF16 tensor, from one kernel (_draft_glue_rmsnorm.DraftGlueRMSNormKernel): it gathers the embedding rows itself (a zero row for an id outside the table), normalizes only the halves the window overlaps and stores only the window. Each row is normalized with FlashInfer's CuTe RMSNormKernel thread geometry, reduction and arithmetic. Under PDL it waits before any load and releases its dependent (eh_proj) right after the wait. At tp16 the eh_proj input window of rank r is 768 columns, so ranks 0-7 normalize only the embedding half and ranks 8-15 only the hidden half; the contiguous window lets eh_proj run on the low-M GEMM at draft step 0's eight rows. The entry certifies the c1 windows of ranks 0, 7, 8 and 15, out-of-table ids and a single row on sm107 against an FP32 reference (and, where FlashInfer runs its CuTe RMSNorm, against flashinfer.norm. rmsnorm bit for bit); the guard refuses windows that are not 8-column aligned and misaligned rows. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
Two collective ops replace everything between the vocabulary-sharded draft LM head and the next draft step's attention (glmMoeDsa/glmDraftSampleGlue.cu). Both take the argmax of the rank's logits shard, exchange the (index, value) pairs through a dedicated MNNVL Lamport workspace (an unwritten word is 0xFFFFFFFF) and pick the global argmax in torch.argmax's order; both write the next step's position ids, sequence and KV lengths (fill, + 1 or the step-0 rewind) and the DSA decode metadata the fused DSA decode-metadata kernel writes. - torch.ops.trtllm.glm_draft_sample_glue, one CTA: optionally the next step's eh_proj input window of cat([enorm(embed(token)), hnorm(hidden)], -1) (FlashInfer's RMSNorm arithmetic) and, after the last draft step, the next-new-tokens rows. - torch.ops.trtllm.glm_draft_ehproj_table, one 8-CTA cluster: also the next step's eh_proj output with its all-reduce and input_layernorm. Each rank publishes, with its packet, the token-table row of its candidate (eh_proj's embedding half applied to enorm(embed(t)) for the tokens of its vocabulary shard) and its FP32 share of eh_proj's hidden half; every rank writes x = bf16(row of the winner + hidden half) and input_layernorm(x). 2 to 16 ranks: at 16 a CTA's 768 columns span two ranks' 384-column shares (16 rank slots). Under PDL both read only constant norm weights before their grid-dependency wait and release their dependents right after it. The tp16 GLM target gets the entries ops/spec/draft_sample_glue and ops/spec/draft_ehproj_table, the workspace (_draft_exchange, one per op, built over the TP group with the MNNVL all-reduce workspace's communicator and handle kind) and the table's construction-time operands (_draft_token_table: the token table of the rank's shard, enorm by the draft glue kernel and the product in FP32, and the rank's hidden-half weight rows). A 4-rank matrix (_draft_step_op_matrix, started by the two collected tests) drives every cell, a c1 draft loop of six table steps and the last glue step eagerly and as CUDA-graph replays, and the guards; the same script runs 16 ranks over four nodes under an external launcher. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
trtllm::glm_mnnvl_hop runs a BF16 one-shot MNNVL all-reduce with the residual add and RMSNorm on the engine's MNNVL all-reduce workspace: - every rank stores its rows through the multicast address and polls all slots in one loop; the sum is FP32 in rank order, from zero per chunk of 8 ranks; - the residual and gamma are copied into shared memory before the exchange, and the warp sums of squares travel to every CTA of the cluster through a transaction barrier (st.async); - each CTA counts its own arrival, and the thread that rotates the Lamport flags reads the arrival counter while it polls and rotates them right after its poll; - the PDL dependents launch right after the grid dependency wait; - with shard_published the producer of the input (the o_proj Lamport publish) has already stored this rank's rows into the workspace: the kernel polls every slot, reads the flags and the residual with volatile loads before any wait, and waits for its primary at its end. trtllm::glm_mnnvl_moe_finalize_hop fuses the TRTLLM-Gen MoE finalize (top-8, the shared experts' output as addend) into the same hop, and trtllm::glm_mnnvl_hop_supported states the shapes both serve: 2, 4, 8 or 16 ranks, one-shot (1 MiB per rank, 1.5 MiB on 16 ranks), one row per cluster of 2 to 8 CTAs of at most 128 threads covering the row. GlmLamportFlags reads and rotates the flags in the layout of common::LamportFlags, with per-CTA arrivals, so the hops and the built-in MNNVL all-reduce follow each other on one workspace. The entries mnnvl_hop and mnnvl_moe_finalize_hop carry the c1 cells (hidden 6144; 8 rows, and 1 row on draft steps). _mnnvl_workspace returns the hop operands and the o_proj publish target from the engine's workspace. The 4-rank op matrices check the cells against FP32 references, against the built-in one-shot all-reduce bit for bit, the fused finalize against finalize + hop bit for bit, the flags, a published shard, CUDA graph replays, late ranks and the guards. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
trtllm::glm_mnnvl_allgather gathers every rank's BF16 rows along the last dimension in rank order and widens them to FP32 in one kernel: each rank stores its shard once through the multicast address at [row][rank][column] of a Lamport buffer and copies the gathered rows out of its local copy as they land, then fills the previous call's buffer with the empty word. The result is the all-gather, concatenation and FP32 cast of the vocabulary-parallel LM head's logits bit for bit; a word of all ones (two BF16 NaNs) arrives with its lowest bit cleared. The gather runs on a LamportExchangeWorkspace: a private MNNVL workspace of the TP group whose unwritten words hold 0xFFFFFFFF, so -0.0 travels unchanged, with flags in the layout of the all-reduce workspace. The draft steps' argmax exchanges own workspaces of the same kind; construction is collective and converges failures over the group. The entry mnnvl_allgather carries the c1 cell (8 rows of the 9,680-entry vocabulary shard). Its 4-rank op matrix checks the concatenation bit for bit, special values, the flags and the refill of the previous buffer, private workspaces side by side, CUDA graph replays, late ranks and the guards. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…ished o_proj For up to 8 rows every all-reduce of the decode step is a one-shot MNNVL hop that adds the residual and applies the following RMSNorm: o_proj's epilogue publishes this rank's shard straight into the hop's workspace, the dense layers' MLP partial sums go through a plain hop, and the MoE layers' routed experts stop before their finalize step, which the post-MoE hop performs with the shared expert's output as its addend. Fused-A prefetches the attention weights that follow it into L2. A batch the kernels do not serve runs the built-in modules in the same order. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…hange trtllm::glm_dsa_hist_topk finds each verify row's exact top-2048 of one generation request whose indexer keys are split into contiguous shards over the TP group: every rank holds only its shard's FP32 logits, and the kernel exchanges 11-bit key histograms (up to three refinement levels) and the selected tokens through a private multicast workspace of three Lamport buffers, so every rank writes the same rows. Ties at the cut take the larger token. trtllm::glm_dsa_hist_shard_meta builds a rank's shard context length and block table from the request's KV length and INDEX_KEY block-table row. The kernels live in cpp/tensorrt_llm/kernels/glmMoeDsa (namespace glm_moe_dsa::hist_topk), the ops in thop/glmDsaHistTopKOp.cpp. A peer wait aborts after 120 s; PDL follows TRTLLM_ENABLE_PDL. The tp16 target gets two entries, hist_topk and hist_shard_meta, and a helper module with the shard geometry and the workspace (a McastGPUBuffer on the TP group, created once per process). The entries are certified on sm_107: hist_shard_meta on one GPU, hist_topk through a 4-rank op matrix (cells, rank agreement, buffer rotation, CUDA-graph replay); its 16-rank cells need a 4-node launch of the same matrix. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
One Triton kernel computes the indexer prologue of a decode step: the k_norm LayerNorm, RoPE of the Q and K rope slices, E4M3 Q and K with one UE8M0 (or amax / 448) scale per 128 values, the weights scaled by the Q scales and softmax_scale * n_heads**-0.5, and the K rows written into the indexer K cache at the step's slot mappings, in the cache layout and slot rule of indexer_k_cache_scatter_op. The arithmetic rounds where the per-kernel chain rounds. With PDL the kernel can trigger its dependent right after its grid-dependency wait and load the step inputs before it. The tp16 target gets two entries: indexer_k_tail, the K part alone (the Q part compiled out) for a step whose wq_b epilogue computes the Q part, and indexer_prologue, the whole prologue. Both are certified on sm_107 against plain-torch references, with the cache bytes checked exactly. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…nels A target-local copy of the CuTe DSL FP8 paged MQA logits kernel scores a request's whole verify window as one MMA atom of next_n * num_heads columns, with the weight register cache capped so that the shared-memory tail serves the remaining heads in the same FMA order. Its early-KV variant splits each group's accumulator into two N halves and is launched with PDL: it loads the KV blocks older than each request's newest next_n keys before griddepcontrol.wait, Q, the weights and the rest after it. All variants compute every element in the same order. The runner takes the caller's padded logits buffer, so captured graphs keep its address. The tp16 target gets two entries: fp8_paged_mqa_logits and paged_mqa_logits_schedule, the per-SM partition of the requests' 256-key splits that DeepGEMM's paged MQA-logits metadata computes. Both are certified on sm_107: the logits bit-exactly on integer data and within a derived band on real values, the early-KV kernel bitwise equal to the default one. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
With MTP index share, draft step 0's indexer computes the TopK of the verify window and the later draft steps reuse each request's row of its last accepted token. stash_last_accepted_rows copies those rows into the shared TopK buffer with one Triton launch (the rows index_select would copy, out-of-range counts clamped). convert_req_index_to_global maps the reused request-local rows to MLA KV pool rows through the trtllm::convert_req_index_to_global op, once per draft loop. Both entries are certified on sm_107, exactly. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
For one generation request the indexer layers take their indexer K and weight rows from fused-A's last 160 columns (an extended weight that the built-in kv_a_proj and indexer projections view), their indexer Q from wq_b with the q_a norm folded and an RoPE/FP8/scale epilogue, and write the indexer K cache from a fused K tail on a side stream. Each rank scores its 1/16 shard of the request's keys and the ranks agree on the exact top-2048 through a histogram exchange; the shard's length, block table and logits schedule are set once per step. The flash decode maps the selected positions to cache rows itself, and the layers without an indexer reuse the last selection, gathering its rows before their wait. Other batches keep the built-in indexer. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…L all-gather On a generation step with speculative decoding the target's shell computes the verify rows' logits itself: the LM head on this rank's vocabulary shard as a split-K GEMM, then one MNNVL all-gather that writes every rank's shard as FP32, on a Lamport workspace the decode target creates with the others. Rows the kernels do not serve keep the built-in logits processor. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…es past 10.3 DeepGEMM's JIT targets the device's own architecture (sm_107f on a 10.7 device), for which the bundled CUTLASS enables neither TMA nor tcgen05, so the first kernel that prefetches a TMA descriptor hits a device-side assert (the DSA indexer's logits in the first context warm-up of GLM-5.2). Such devices now build for sm_100f, which they run and CUTLASS knows, the way DeepGEMM already pins 12.x devices to sm_120f. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…add the missing fakes The MTP step tests that certify their cells import the certifying fixture they use. The draft token table's band allows the two BF16 roundings its rows carry (the normalized row and the product), 2 ulp per element. The fused-A GEMM, router GEMV, MNNVL hops and logits all-gather get the fake implementations every trtllm op registers. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
DLPack refuses to export a tensor that requires grad, and module parameters such as the LM head's weight may still do so; the split-K GEMMs convert a detached view, which shares the storage. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
DecodeTarget asserts each decoder layer's MLA weights and KV cache layout, and derives the operands its attention kernels read (the absorbed q_b, the fused-A weight with the indexer rows, the folded indexer wq_b, the L2 prefetch list), one layer at a time through _check_attention and _fold_attention_weights; the MoE layers' shared-expert check and the shared expert's two split-K GEMMs around SwiGLU are methods as well, so a layer outside the decoder stack can take the same checks, operands and GEMMs. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
With one-model MTP-Eagle the target builds its own worker, MTPDraftWorker, in place of the built-in MTPEagleWorker, with the same draft KV cache choice and draft model. On a generation step of one request whose draft layer shares the target's KV pool, drafting greedily through the configured draft length, the worker runs the draft loop on DecodeTarget's kernels; every other step runs the built-in loop. Draft step 0 runs on the verify rows: one kernel writes this rank's window of eh_proj's input (embedding gather, enorm, hnorm, concat, column split), eh_proj runs on it as a split-K GEMM, and its hop applies input_layernorm with a zero residual. The MTP layer's attention runs as a verify full-indexer layer on this rank's key shard; on a side stream each request's top-k of its last accepted row is stashed and mapped to KV pool rows for the later steps, which run no indexer and attend to that selection, or, with a context short enough that the built-in metadata skips the indexer, to every position up to their own. After the post-attention hop each request keeps the row of its last accepted token. Every step's MoE runs the router GEMV and the low-M BF16 experts beside the split-K shared expert, and the MoE-finalize hop adds both with the residual and applies shared_head.norm. The draft LM head runs on the vocabulary shard. Before every step but the last, the next step's eh_proj hidden half (hnorm inside its split-K GEMM) starts first on the aux stream, and glm_draft_ehproj_table writes the draft token, the next step's eh_proj output and input_layernorm from a token table, its positions and the sequence and KV lengths; the last step's glm_draft_sample_glue writes its token and the next-new-tokens rows. The worker keeps the built-in loop's MTP index-share flags, draft KV cache context and cleanup, the host side of the switch to one row per request, and the per-step update_for_spec_dec, which rebuilds the DSA decode metadata from the lengths the kernels write. The token table needs the whole eh_proj weight and the low-M MoE its experts in up-then-gate rows, so the target's loader takes both from the checkpoint before the built-in loader shards them (the built-in BF16 MoE may keep its experts in a backend layout). The draft state is built once after the weights load; every rank creates the two draft exchange workspaces there in the same order. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
The decode target derives its operands from the loaded parameters, which require grad; under grad mode the derivation kept its intermediates alive. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…tensorrt_llm The ranks put the repo root first on their path, so under an installed wheel they imported the source tree, which has no bindings. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
The band's unit is 2**-8 of the row maximum, and one BF16 ulp of an element in the maximum's binade is up to twice that: a midpoint flip there failed the draft-step matrix's glue cell at 1.02 units. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…s, kernels and tests ruff-format, ruff import folding, yapf on the legacy custom-op files, clang-format and cmake-format; the hist top-k kernel's needIn becomes needInGroup for codespell. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…he manager DecodeTarget kept the first KV cache manager's primary pool. The executor tears that manager down after its memory-profiling pass and allocates the final one, so the final executor's first decode warm-up read freed memory (illegal memory access on every rank). The pool view is memoized per manager on main, so it is now asked for at each attention call. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…rebuild The draft steps after the first run no indexer and read only the sequence and KV lengths their kernels advance in place, so the built-in loop's update_for_spec_dec (fused DSA metadata, DeepGEMM schedules, 2-D lengths) after every draft step is dead work in the captured graph. With a short context only the next step's dense selection is refreshed. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…verlap correction With speculative decoding and the overlap scheduler, _preprocess_inputs called on_update_kv_lens before the kv_lens correction and again after it. The second call supersedes the first and nothing reads the metadata in between, so the first is skipped when the second runs. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…rnel For a verify step of generation requests only, MTPDraftWorker samples the target tokens as the built-in path does and then runs glm_verify_head_accept, which writes the accepted rows and the strict count with the forced count applied (same warm-up skip, pool draw and counter step as _apply_force_accepted_tokens), instead of the 8-20 small kernels of the built-in chain. Ragged verification and context requests keep that chain. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…er on shared layers The decode target ran q_b after the indexer and on the regular split-K 8 plan on every layer, so neither stream layout of the ported kernels was in use: - Indexer layers (and MTP draft step 0) queue q_b on the aux stream after the KV half, on the beside_topk plan (18 CTAs), so it runs beside the histogram top-k instead of after it. - Layers that attend to the previous indexer layer's selection run q_b on the early_gather plan (split-K 4), the KV half triggers its dependents at entry, and fused-A also prefetches that selection into L2, so the flash decode gathers the old KV rows while q_b runs. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…aph to its graph head DSAtrtllmAttentionMetadata gets the engine-only defer_cuda_graph_decode_outputs() scope CSA2 already has. Inside it, prepare() of a generation-only CUDA-graph batch on a native DSA cache (no ragged verification, no helix, no separate draft KV cache) skips the outputs that on_update_kv_lens() rebuilds before any consumer: - the generation cached-token and KV indptr uploads; - req_idx_per_token; - the host dense TopK indices of short contexts; - the indexer slot mappings; - the 2-D KV-length fill, the DeepGEMM schedules and their copies. The host scalars (max_gen_seq_len, the indexer-skip flags) are still set, and the INDEX_KEY offsets, the block table and the speculative expansion buffers are still prepared. Entry 0 of both generation indptrs is zeroed once at allocation, so the unfused rebuild, which writes only the entries after it, leaves the same buffer as the upload did. The INDEX_KEY clamp now covers only the rows past the batch: the batch rows were just copied from pool slots that the cache manager clamps to be non-negative, so the clamp never changed them. DecoderRunner._cuda_graph_metadata_prepare_scope now also admits a step of linear one-model speculative decoding that verifies drafts, without a separate draft KV cache manager and without ragged verification. This is correct because _forward_step starts with _preprocess_inputs, which is captured with the graph and also runs in the eager warm-up forwards: - with the overlap scheduler it adds the previous step's acceptance correction to position_ids and kv_lens_cuda and then calls on_update_kv_lens(), which recomputes every skipped output from the corrected lengths, the seq_lens and the INDEX_KEY offsets on the device, before the first layer reads them; - without the overlap scheduler it calls on_update_kv_lens() first, from the prepared lengths; - the draft loop changes these buffers only on the device after the target layers, and the next step's graph head rebuilds them again. Before this change prepare() computed the same outputs from the uncorrected host lengths and on_update_kv_lens() overwrote them before any read, so the values the model reads are unchanged. A separate draft KV cache swaps and refreshes DSA buffers outside the graph, and ragged verification builds per-row layouts on the host; both keep the full preparation, as do direct prepare() callers and DeepSeek-V4, whose metadata subclass has its own prepare() and cache layout. A speculative step without drafts also keeps it, so CSA2, which defers only single-token rows, sees no change. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
… staging Between two decode graph replays, the overlap staging of a speculative step and the sampler issue several copies whose payload is constant for a request, or that nothing reads. This removes them without changing any buffer the step reads. - num_accepted_draft_tokens: nothing in tensorrt_llm reads SpecMetadata.num_accepted_draft_tokens, so the per-request list, its H2D copy and its device buffer are dropped. - previous_pos_id_offsets_cuda / previous_kv_lens_offsets_cuda: the runner zeroed both buffers whole before the previous-batch gathers rewrote their rows. It now tracks the row span each buffer may be nonzero in and clears only that span's rows outside the rows the gathers rewrite next, so the contents after staging are the same and a steady-state step issues no clearing kernel. - host_staging.copy_if_changed() skips an upload whose destination storage still holds the same payload from the last copy made through it (same elements, equal values). The record is keyed by the storage's base tensor, so views of a buffer shared between owners share it. It is used for the previous-batch slot indices, the previous-batch position indices, the gather ids, the Eagle3 batch_indices_cuda and the SpecSampler store slots, which now live in a persistent buffer instead of a new pageable-staged tensor per step. The one other writer of those buffers, the encoder-decoder input fast path, calls forget_copy() after its write. At batch size 1 this removes one H2D for the accepted counts, two zeroing kernels and five constant H2D copies per step. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…ged DSA page-table uploads KV cache manager V2 ignored max_blocks and always copied every block column of each row of its pinned page table (max_seq_len worth, about 16k columns at a 1M context) with copyBatchBlockOffsetsToDeviceKernel. The kernel and its binding now take an optional column bound (max_num_blocks, rounded up to whole 16-byte accesses; negative or absent copies whole rows), and KVCacheManagerV2.copy_batch_block_offsets passes max_blocks through as KVCacheManager already does. Helix rows are still copied whole. TrtllmAttentionMetadata.prepare() capped the copy only without speculative decoding, because the overlap correction and the draft sub-steps advance kv_lens_cuda on the device past the host lengths. With a linear draft each of these advances is below one verification window of 1 + max_total_draft_tokens tokens: the correction adds the previous step's accepted count minus this step's window, and the draft steps add at most draft_len - 1 after resetting to the accepted length. The cap under linear speculation is therefore the host maximum KV length (which already includes the extra draft KV tokens) plus two windows. Tree drafts and helix keep the full width. No kernel reads a block column past its device KV length, so the columns left unwritten are never read. The DSA INDEX_KEY offsets and block table were uploaded every step: - The host INDEX_KEY rows are now compared with the freshly computed page indices and rewritten only when they differ; each rewrite gives the host table a new version object. The upload of the batch rows goes through copy_if_changed() with that version as the token, so it is skipped while the device rows hold that version of the host rows. - The block table goes through copy_if_changed() with its host payload. Both device buffers are written only by these uploads, so a skipped upload leaves exactly the contents the copy would have written. The INDEX_KEY rows only change when a request gets a new block, and the block table only when its used width or entries change. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…ampling SpecMetadata._populate_request_rng_state() uploaded the per-request seeds and offsets (request_seeds, request_offsets) on every step. Their only reader is _rng_state_per_request(), which runs only on the rejection-sampling paths: the rejection acceptance kernels and the advanced draft samplers, which are taken only when use_rejection_sampling is set and the batch is not all greedy. use_rejection_sampling is fixed when the metadata is created, so without it the two H2D copies are skipped. The host offset windows are still taken every step, so the per-row seeds and offsets the non-greedy target sampling reads are unchanged. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
Collaborator
Author
|
Superseded by #20517. |
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.
Description
Adds a modeling_v2 (Staircase) target for GLM-5.2-NVFP4 on Rubin (sm_107) at TP16: tensor-parallel attention
and experts over 16 ranks, no attention DP, no expert parallelism, with one-model MTP-Eagle (any draft length) or no
speculative decoding. It is the generation server of a disaggregated low-latency deployment (one DEP4 context server on
the built-in implementation, one TP16 generation server on this target,
modeling_v2: require).Routing.
models/glm_moe_dsa/routing.pyclaimsGlmMoeDsaForCausalLMonly for compute capability (10, 7), theGLM-5.2 shape (78 layers, hidden 6144, 256 routed experts, q_lora_rank 2048), the TP16 mapping above and MTP-Eagle /
no speculation. Every other deployment gets
NoTargetand stays onmodeling_deepseekv3.Target (
models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/). The modules and weights are the built-in GLM-5.2 ones,loaded by the built-in loader, so the KV and indexer-K caches keep the layout the context server writes.
PrefillTargetruns the built-in modules.DecodeTargetruns a one-request generation step (verify rows <= 8) onGLM-owned kernels:
into the cache, a sparse-MLA flash decode that maps the selection to cache rows itself, and o_proj publishing its
shard for the next hop. On indexer layers q_b runs on a side stream beside the top-k; on layers that reuse a
selection the KV half is written at entry and flash decode gathers early.
shard, and an exact histogram top-k merged across ranks.
RMSNorm fused in; an all-gather for the verify logits.
at K=128; router GEMV; split-K shared expert.
MTPDraftWorker(anMTPEagleWorkersubclass). For a generation step of onerequest with greedy drafting at the configured draft length, it runs the draft steps on the same kernels: a
token-table eh_proj, a reused top-k selection after step 0, a low-M BF16 MoE for the MTP layer, and fused
sample/glue kernels for the argmax exchange and the next step's inputs. The draft steps reuse the step's DSA
metadata instead of rebuilding it, and the verify step's acceptance (strict count or the forced count) is one
kernel. Every other step runs the built-in loop.
Larger batches and anything the fast path does not cover fall back to the built-in modules inside the same target.
Kernels. GLM-only CUDA kernels are in
cpp/tensorrt_llm/kernels/glmMoeDsa/, exposed astorch.ops.trtllm.glm_*(
cpp/tensorrt_llm/thop/glm*Op.cpp). CuTe DSL kernels live in the target'sops/. Each op is a catalogOpWrapperwith certified cells and a GPU test.
Shared code changes (general, not sm-gated):
3rdparty: a DeepGEMM patch so DeepJIT builds for the 10.0 family target on 10.x devices past 10.3. The bundledCUTLASS enables TMA and tcgen05 only for targets it knows; without this the context server's warm-up hits a device
assert on Rubin.
fmhaKernels.h: the token-sparse MLA generation kernels pad the Q tile to 8 rows, as the dense MLA selection does.With fewer than 8 heads per rank (64 MLA heads over 16 ranks) the SmemTile check fails at warm-up.
modeling_deepseekv3.py: with the fused post-MLP all-reduce, the dense MLP keeps the full TP instead of capping it atGPUs per node. The fused all-reduce sums over the whole TP group, so the capped MLP counted each partial sum
tp / mlp_tp times.
default off for existing callers.
Arch.SM_107, theGlmMoeDsaForCausalLMrouter entry, a GLM cell for the FP4 quantize entry, and fakeregistrations for the new ops.
Per-step work around the decode graph of an overlap-scheduler speculative step (general; the step is GPU-bound, so
every eager op between two graph replays adds to it):
DecoderRunner._preprocess_inputsrefreshes the KV-length metadata once per step under the overlap correction.prepare()leaves the decode outputs the graph head rebuildsin
on_update_kv_lens(generation indptrs,req_idx_per_token, slot mappings, the DeepGEMM schedules, the densetop-k) to the graph, and clamps only the graph-padding rows of the INDEX_KEY table.
host_staging.copy_if_changed()skips an upload whose destination still holds the same payload(previous-batch slots and positions, gather ids, Eagle3 batch indices, sampler store slots); the unread
num_accepted_draft_tokensupload is removed; the overlap offset buffers are zeroed only outside the rows the steprewrites.
copy_batch_block_offsetshonorsmax_blocks(the copy kernel and its binding take a columnbound). Under linear speculative decoding the attention metadata bounds the copy by the batch's maximum KV length
plus two verification windows, the most the overlap correction and the draft steps can add on the device. The DSA
INDEX_KEY offsets and block table upload only when they change.
SpecMetadatauploads the per-request Philox state only with rejection sampling, its only reader.Performance
Disaggregated c1 (1 ctx DEP4 + 1 gen TP16, 20 Rubin GPUs), MTP K=7 with forced acceptance 2.91 tokens/step;
x = 1000 / p90 inter-token latency (tokens/s/user). The generation server sets
sparse_attention_config.use_cute_dsl_paged_mqa_logits: false: the target scores keys with its own kernels, and theCuTe DSL setting only adds metadata work for the built-in indexer.
Without this PR, main does not run this deployment on Rubin: the context server dies in the DeepGEMM warm-up, and the
built-in TP16 generation server fails the sparse-MLA SmemTile check.
Accuracy
GSM8K (1,319 questions) on the same disaggregated deployment with real acceptance: strict-match 0.9697,
flexible-extract 0.9712 (threshold 0.9). Measured acceptance length 5.52 tokens/step at K=7.
Test Coverage
tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/: one GPU test per op entry overits certified cells, plus 4-rank matrices for the collective ops. sm_107 only; the files skip at module level
elsewhere.
tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py: GLM routing, including near misses (other shape,ADP, ep16, tp8, vanilla MTP, NGram, sm_103). CPU.
test_modeling_v2_claims.py,test_modeling_v2_no_stale_claims.py,test_modeling_v2_target_contract.py,thop/parallel_hw_agnostic/test_custom_ops.py(fake registrations).above passes in full: 382 passed.
PR Checklist
🤖 Generated with Claude Code