From 58758a8d1122edd48e0b42331a6ea550f75be772 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:42:22 +0000 Subject: [PATCH 01/33] [None][feat] add Helix speculative verify-group support to the MLA kernels A Helix verify group of 1 + draft_len tokens can straddle a ledger-page boundary, so KV ownership within one group splits across two CP ranks. The existing per-sequence helix_is_inactive_rank gate cannot express that: it is all-or-nothing for the whole group. Add a per-token rank-local write slot, helix_local_slots, plumbed through mlaKernels and dsv3RopeOp. A negative entry means another rank owns the token's global position. When the pointer is non-null it supersedes the per-sequence gate and also supplies the KV write index, so each token lands in its owner's cache at the right slot; when it is null every path keeps its current behaviour. Also carry the zero-KV sanitize into the fifo-v2 all-to-all sender, where the entry is already streaming through shared memory, instead of leaving it to separate elementwise kernels on the caller side. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/helixAllToAll.cu | 134 +++++++++++++++++----- cpp/tensorrt_llm/kernels/helixAllToAll.h | 8 ++ cpp/tensorrt_llm/kernels/mlaKernels.cu | 35 ++++-- cpp/tensorrt_llm/kernels/mlaKernels.h | 7 ++ cpp/tensorrt_llm/thop/alltoallOp.cpp | 25 +++- cpp/tensorrt_llm/thop/dsv3RopeOp.cpp | 22 +++- 6 files changed, 187 insertions(+), 44 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index 304b255a11f6..f9eb6aee3009 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.cu +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -85,10 +85,22 @@ __host__ __device__ inline uint8_t* getPtr(HelixFieldInfo const& fieldInfo, int return fieldInfo.dataPtr + blockIdx * fieldInfo.stride; } +bool isFieldAlignedForCopy(HelixFieldInfo const& fieldInfo, uintptr_t alignment) +{ + return reinterpret_cast(fieldInfo.dataPtr) % alignment == 0 + && static_cast(fieldInfo.stride) % alignment == 0 + && static_cast(getFieldSize(fieldInfo)) % alignment == 0; +} + +template __device__ __forceinline__ void waitG2sAllFields(uint64_t* smemBar, uint32_t* phaseParity) { cp_async_wait_group<0>(); smemBarWait(smemBar, phaseParity); + if constexpr (NEEDS_WARP_SYNC) + { + __syncwarp(); + } } // Align size to 128 bytes @@ -113,31 +125,44 @@ __device__ __forceinline__ void g2sField( } } -template +template __device__ __forceinline__ int g2sAllFields( HelixFieldInfo const* fieldInfo, int dataIndex, uint8_t* shmemBase, uint64_t* smemBar, int laneId) { - int totalSize = 0; - // Load field 0 (variable size half) g2sField(fieldInfo[0], dataIndex, shmemBase, 0, smemBar, laneId); - int field0Size = getFieldSize(fieldInfo[0]); - totalSize += field0Size; + int const field0Size = getFieldSize(fieldInfo[0]); + int bulkCopySize = field0Size; - // Load field 1 (single float2) + // Load field 1 (one or more float2 values). if constexpr (ALLOW_VARIABLE_FIELD1) { - g2sField(fieldInfo[1], dataIndex, shmemBase, totalSize, smemBar, laneId); - totalSize += getFieldSize(fieldInfo[1]); + if constexpr (USE_BULK_FIELD1) + { + g2sField(fieldInfo[1], dataIndex, shmemBase, field0Size, smemBar, laneId); + bulkCopySize += getFieldSize(fieldInfo[1]); + } + else + { + constexpr int kCopySize = sizeof(float2); + int const field1Size = getFieldSize(fieldInfo[1]); + for (int offset = laneId * kCopySize; offset < field1Size; offset += WARP_SIZE * kCopySize) + { + ldgsts(reinterpret_cast(shmemBase + field0Size + offset), + reinterpret_cast(getPtr(fieldInfo[1], dataIndex) + offset), true); + } + cp_async_commit_group(); + } } else { - ldgsts<8>(reinterpret_cast(shmemBase + totalSize), + ldgsts<8>(reinterpret_cast(shmemBase + field0Size), reinterpret_cast(getPtr(fieldInfo[1], dataIndex)), laneId == 0); cp_async_commit_group(); } - return totalSize; + // Only bytes copied by cp.async.bulk participate in the mbarrier transaction. + return bulkCopySize; } // ============================================================================ @@ -156,7 +181,7 @@ __device__ __forceinline__ void s2gField( } } -template +template __device__ __forceinline__ void s2gAllFields( HelixFieldInfo const* fieldInfo, int dataIndex, uint8_t* shmemBase, int laneId) { @@ -167,11 +192,25 @@ __device__ __forceinline__ void s2gAllFields( int field0Size = getFieldSize(fieldInfo[0]); offset += field0Size; - // Store field 1 (single float2) + // Store field 1 (one or more float2 values). if constexpr (ALLOW_VARIABLE_FIELD1) { - s2gField(fieldInfo[1], dataIndex, shmemBase, offset, laneId); - offset += getFieldSize(fieldInfo[1]); + if constexpr (USE_BULK_FIELD1) + { + s2gField(fieldInfo[1], dataIndex, shmemBase, offset, laneId); + } + else + { + constexpr int kCopySize = sizeof(float2); + int const field1Size = getFieldSize(fieldInfo[1]); + for (int fieldOffset = laneId * kCopySize; fieldOffset < field1Size; fieldOffset += WARP_SIZE * kCopySize) + { + auto const* srcPtr = reinterpret_cast(shmemBase + offset + fieldOffset); + auto* dstPtr = reinterpret_cast(getPtr(fieldInfo[1], dataIndex) + fieldOffset); + *dstPtr = *srcPtr; + } + __syncwarp(); + } } else { @@ -294,7 +333,7 @@ __host__ __device__ __forceinline__ int computeProtoTransferSize(HelixFieldInfo // Main All-to-All Kernel // ============================================================================ -template +template __global__ void helixAllToAllKernel(HelixAllToAllParams params) { extern __shared__ uint8_t allWarpShmem[]; @@ -371,7 +410,7 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) int dataIndex = entryIdx * params.cpSize + peerRank; // Load data from global to shared, then arrive on barrier - int loadedSize = g2sAllFields( + int loadedSize = g2sAllFields( params.sendFields, dataIndex, shmem, &allWarpSmemBar[group], laneId); uint64_t arriveState = mbarrier_arrive_expect_tx(&allWarpSmemBar[group], laneId == 0 ? loadedSize : 0); @@ -394,10 +433,37 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) } // wait for data to be loaded into shared memory - waitG2sAllFields(&allWarpSmemBar[group], &phaseParity); + waitG2sAllFields(&allWarpSmemBar[group], &phaseParity); // note: we don't need to pack anything, fields are already packed in // shared memory + // Zero-local-KV sanitization: the entry is in shared memory and + // not yet packed, so overwriting it costs no global traffic. + if (params.zeroKvMask != nullptr && params.zeroKvMask[entryIdx / params.zeroKvMaskDivisor] != 0) + { + // field0Size is a multiple of 16 (the op checks it), so the + // int4 store never runs off the end of the field. + int const field0Size = getFieldSize(params.sendFields[0]); + for (int off = laneId * static_cast(sizeof(int4)); off < field0Size; + off += WARP_SIZE * static_cast(sizeof(int4))) + { + *reinterpret_cast(shmem + off) = make_int4(0, 0, 0, 0); + } + if (laneId == 0) + { + // Field 1 sits at getFieldSize(field 0), not at + // align_up(..., 16) as computeTotalUnpackedSize computes; + // they agree only because field 0 is 16-byte aligned. + auto* stats = reinterpret_cast(shmem + field0Size); + int const statsCount = getFieldSize(params.sendFields[1]) / static_cast(sizeof(float2)); + for (int i = 0; i < statsCount; ++i) + { + stats[i] = make_float2(-INFINITY, 0.F); + } + } + __syncwarp(); + } + LL128Proto::protoPack(shmem, head, singlePacked128ByteCount, fifoEntry128ByteIndexBase, laneId); uint64_t* fifoEntry = fifoBase + fifoEntryIndex * (HELIX_FIFO_ENTRY_BYTES / sizeof(uint64_t)); @@ -474,7 +540,7 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) shmem, tail, singlePacked128ByteCount, fifoEntry128ByteIndexBase, loaded128ByteCount, laneId); // note: fields are already unpacked in shared memory - s2gAllFields(params.recvFields, dataIndex, shmem, laneId); + s2gAllFields(params.recvFields, dataIndex, shmem, laneId); // wait for data to be read from shared memory cp_async_bulk_wait_group_read<0>(); @@ -503,7 +569,7 @@ struct hash_cache_key } }; -template +template std::tuple computeChannelAndGroupCount(int cpSize, HelixFieldInfo const* fields) { static std::unordered_map, std::tuple, hash_cache_key> cache; @@ -535,7 +601,7 @@ std::tuple computeChannelAndGroupCount(int cpSize, HelixFieldInfo // Set shared memory attribute if needed if (totalDynamicShmemSize > 48 * 1024) { - TLLM_CUDA_CHECK(cudaFuncSetAttribute(helixAllToAllKernel, + TLLM_CUDA_CHECK(cudaFuncSetAttribute(helixAllToAllKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, totalDynamicShmemSize)); } @@ -559,14 +625,14 @@ std::tuple computeChannelAndGroupCount(int cpSize, HelixFieldInfo // Host Launch Function // ============================================================================ -template +template void launchHelixAllToAllImpl(HelixAllToAllParams const& params, cudaStream_t stream) { int maxChannelCount = computeHelixMaxChannelCount(params.cpSize); TLLM_CHECK_WITH_INFO(params.maxChannelCount == maxChannelCount, "maxChannelCount %d does not match computed maxChannelCount %d", params.maxChannelCount, maxChannelCount); auto [channelCount, groupCountPerCta, totalDynamicShmemSize] - = computeChannelAndGroupCount(params.cpSize, params.sendFields); + = computeChannelAndGroupCount(params.cpSize, params.sendFields); if (params.channelCount > 0) { channelCount = params.channelCount; @@ -580,7 +646,7 @@ void launchHelixAllToAllImpl(HelixAllToAllParams const& params, cudaStream_t str // and receiver) int ctaPerChannel = ceil_div(params.cpSize, groupCountPerCta); - auto* kernel_instance = &helixAllToAllKernel; + auto* kernel_instance = &helixAllToAllKernel; cudaLaunchConfig_t config; config.gridDim = dim3(ctaPerChannel, channelCount, 2); config.blockDim = dim3(WARP_SIZE, groupCountPerCta); @@ -638,11 +704,27 @@ void launchHelixAllToAll(HelixAllToAllParams const& params, bool allowVariableFi { if (allowVariableField1) { - launchHelixAllToAllImpl(params, stream); + constexpr uintptr_t kBulkCopyAlignment = 16; + constexpr uintptr_t kFallbackCopyAlignment = sizeof(float2); + int const field1Size = getFieldSize(params.sendFields[1]); + TLLM_CHECK_WITH_INFO(field1Size % sizeof(float2) == 0, "Variable field 1 must contain whole float2 values"); + TLLM_CHECK_WITH_INFO(isFieldAlignedForCopy(params.sendFields[1], kFallbackCopyAlignment) + && isFieldAlignedForCopy(params.recvFields[1], kFallbackCopyAlignment), + "Variable field 1 must be aligned to float2"); + bool const useBulkField1 = isFieldAlignedForCopy(params.sendFields[1], kBulkCopyAlignment) + && isFieldAlignedForCopy(params.recvFields[1], kBulkCopyAlignment); + if (useBulkField1) + { + launchHelixAllToAllImpl(params, stream); + } + else + { + launchHelixAllToAllImpl(params, stream); + } } else { - launchHelixAllToAllImpl(params, stream); + launchHelixAllToAllImpl(params, stream); } } diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.h b/cpp/tensorrt_llm/kernels/helixAllToAll.h index 95ab959ff3fc..ac51e012adbf 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.h +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.h @@ -47,6 +47,14 @@ struct HelixAllToAllParams int cpSize; int channelCount; // use 0 to auto-compute int maxChannelCount; + + // Rows this rank owns no KV for. The sender replaces them in shared memory + // with a no-op contribution for the combine: field 0 zeros, field 1 + // (max, sum) = (-inf, 0). nullptr when the caller already sanitized. + uint8_t const* zeroKvMask; + // entryCount / zeroKvMask length: 1 when an entry is a token (fifo v2), + // num_heads when it is a (token, head) pair (fifo v1). + int zeroKvMaskDivisor; }; // ============================================================================ diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index 1d984adeb49a..cde289be328a 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -635,8 +635,8 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, int* seqKVOffsets, int q_pe_ld, int q_pe_stride, KvCacheDataType cache_type, float* bmm1_scale, float* bmm2_scale, float const* quant_scale_o, float const* quant_scale_q, float const* quant_scale_kv, float const* dequant_scale_q, float const* dequant_scale_kv, float host_bmm1_scale, int32_t const* helix_position_offsets, - bool const* helix_is_inactive_rank, bool precomputed_cu_seqlens = false, bool precomputed_fmha_scheduler = false, - bool q_rope_applied = false) + bool const* helix_is_inactive_rank, int32_t const* helix_local_slots = nullptr, bool precomputed_cu_seqlens = false, + bool precomputed_fmha_scheduler = false, bool q_rope_applied = false) { // Constants. using VecT = typename VecType::Type; @@ -755,10 +755,18 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, { if (head_idx == head_num) { - // If helix parallelism is being used, only write to KV cache if current rank is active. - if (helix_is_inactive_rank == nullptr || !helix_is_inactive_rank[batch_idx]) + // If helix parallelism is being used, only write to KV cache if this rank owns + // the token's global position. With speculative verify groups the per-token slot + // table decides -- a 1 + draft_len group can straddle a ledger-page boundary and + // split ownership between two ranks -- otherwise the per-sequence flag does. + bool const helix_write = helix_local_slots != nullptr + ? helix_local_slots[global_token_idx] >= 0 + : (helix_is_inactive_rank == nullptr || !helix_is_inactive_rank[batch_idx]); + if (helix_write) { - auto const token_kv_idx = kv_cache_lengths[batch_idx] - seq_len + local_token_idx; + auto const token_kv_idx = helix_local_slots != nullptr + ? helix_local_slots[global_token_idx] + : kv_cache_lengths[batch_idx] - seq_len + local_token_idx; { auto kDst = reinterpret_cast(kv_cache.getKBlockPtr(batch_idx, token_kv_idx)); @@ -846,10 +854,17 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, } } - // If helix parallelism is being used, only write to KV cache if current rank is active. - if (helix_is_inactive_rank == nullptr || !helix_is_inactive_rank[batch_idx]) + // If helix parallelism is being used, only write to KV cache if this rank owns the + // token's global position (per-token slots for speculative verify groups, the + // per-sequence flag otherwise; see the Q/K branch above). + bool const helix_write = helix_local_slots != nullptr + ? helix_local_slots[global_token_idx] >= 0 + : (helix_is_inactive_rank == nullptr || !helix_is_inactive_rank[batch_idx]); + if (helix_write) { - auto const token_kv_idx = kv_cache_lengths[batch_idx] - seq_len + local_token_idx; + auto const token_kv_idx = helix_local_slots != nullptr + ? helix_local_slots[global_token_idx] + : kv_cache_lengths[batch_idx] - seq_len + local_token_idx; auto const src_kv_global_offset = static_cast(global_token_idx) * (c_k + ROPE_DIM); { @@ -1961,8 +1976,8 @@ void invokeMLARopeGeneration(MlaParams& params, KVCacheBuffer kv_cache_buffer params.cache_seq_lens, params.cu_kv_seqlens, params.q_pe_ld, params.q_pe_stride, params.cache_type, params.bmm1_scale, params.bmm2_scale, params.quant_scale_o, quant_scale_q_eff, params.quant_scale_kv, params.dequant_scale_q, params.dequant_scale_kv, params.host_bmm1_scale, params.helix_position_offsets, - params.helix_is_inactive_rank, params.precomputed_cu_seqlens, params.precomputed_fmha_scheduler, - params.q_rope_applied); + params.helix_is_inactive_rank, params.helix_local_slots, params.precomputed_cu_seqlens, + params.precomputed_fmha_scheduler, params.q_rope_applied); } template diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.h b/cpp/tensorrt_llm/kernels/mlaKernels.h index a0125d236a89..8bf1f1497bae 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -158,6 +158,13 @@ struct MlaParams // for Helix parallelism: whether the current rank is inactive, shape [b] // (the current query tokens are not appended to this rank's KV cache) bool const* helix_is_inactive_rank{nullptr}; + + // for Helix parallelism with speculative verify groups: per-token + // rank-local KV write slot, shape [num_tokens]; -1 means another CP rank + // owns the token's global position. Non-null supersedes the per-sequence + // helix_is_inactive_rank gate (a 1 + draft_len group can straddle a + // ledger-page boundary, splitting ownership between two ranks). + int32_t const* helix_local_slots{nullptr}; }; template diff --git a/cpp/tensorrt_llm/thop/alltoallOp.cpp b/cpp/tensorrt_llm/thop/alltoallOp.cpp index 8a775b1cf6b8..9c076850d29f 100644 --- a/cpp/tensorrt_llm/thop/alltoallOp.cpp +++ b/cpp/tensorrt_llm/thop/alltoallOp.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/thop/thUtils.h" +#include #include TRTLLM_NAMESPACE_BEGIN @@ -125,10 +126,13 @@ std::vector alltoall_helix( * @param workspace Workspace tensor (uint64, strided across ranks) * @param cp_rank Current context parallel rank * @param cp_size Total number of context parallel ranks + * @param zero_kv_mask Optional bool mask of rows this rank owns no KV for. The + * sender replaces them with a no-op contribution, so the caller must not + * sanitize itself. Its length must divide entry_count. * @return tuple of (partial_o_out, softmax_stats_out) with same shapes as inputs */ -std::tuple alltoall_helix_native( - torch::Tensor partial_o, torch::Tensor softmax_stats, torch::Tensor workspace, int64_t cp_rank, int64_t cp_size) +std::tuple alltoall_helix_native(torch::Tensor partial_o, torch::Tensor softmax_stats, + torch::Tensor workspace, int64_t cp_rank, int64_t cp_size, std::optional zero_kv_mask) { // Input validation @@ -224,6 +228,21 @@ std::tuple alltoall_helix_native( params.channelCount = 0; // auto-compute params.maxChannelCount = tensorrt_llm::kernels::computeHelixMaxChannelCount(cp_size); + // Optional zero-local-KV sanitization, applied by the sender in shared memory + params.zeroKvMask = nullptr; + params.zeroKvMaskDivisor = 1; + if (zero_kv_mask.has_value()) + { + auto const& mask = zero_kv_mask.value(); + CHECK_TH_CUDA(mask); + CHECK_CONTIGUOUS(mask); + CHECK_TYPE(mask, at::ScalarType::Bool); + TORCH_CHECK(mask.numel() > 0 && entry_count % mask.numel() == 0, "zero_kv_mask numel (", mask.numel(), + ") must divide the all-to-all entry count (", entry_count, ")"); + params.zeroKvMask = reinterpret_cast(mask.data_ptr()); + params.zeroKvMaskDivisor = entry_count / mask.numel(); + } + // Launch kernel auto stream = at::cuda::getCurrentCUDAStream(); tensorrt_llm::kernels::launchHelixAllToAll(params, allowVariableField1, stream); @@ -260,7 +279,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def("alltoall_helix(Tensor[] input_list, int[] group, int? num_lists) -> Tensor[]"); m.def( "alltoall_helix_native(Tensor partial_o, Tensor softmax_stats, Tensor(a!) workspace, int " - "cp_rank, int cp_size) -> (Tensor, Tensor)"); + "cp_rank, int cp_size, Tensor? zero_kv_mask=None) -> (Tensor, Tensor)"); m.def( "initialize_helix_workspace(Tensor(a!) workspace, int cp_rank, int cp_size) " "-> ()"); diff --git a/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp b/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp index b6ef2e13cbf5..e3a2eb2ef200 100644 --- a/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp @@ -76,6 +76,8 @@ struct MlaRopeGenArgs float host_bmm1_scale; int32_t const* helix_position_offsets_ptr; bool const* helix_is_inactive_rank_ptr; + // Per-token KV write slots for speculative verify groups (nullptr otherwise). + int32_t const* helix_local_slots_ptr; // `kv_norm_weight_ptr` set: `invokeMLAKvNormRopeQuantGeneration` produces the KV // half (norm + rope + fp8 + paged write) and the RoPE kernel runs Q-only. void const* kv_norm_weight_ptr; @@ -135,6 +137,7 @@ void invokeMLARopeGenerationHelper(T const* latent_cache_ptr, T* q_pe_ptr, T* fu mla_params.helix_position_offsets = args.helix_position_offsets_ptr; mla_params.helix_is_inactive_rank = args.helix_is_inactive_rank_ptr; mla_params.q_rope_applied = args.q_rope_applied; + mla_params.helix_local_slots = args.helix_local_slots_ptr; mla_params.precomputed_cu_seqlens = args.precomputed_cu_seqlens; mla_params.precomputed_fmha_scheduler = args.precomputed_fmha_scheduler; @@ -190,8 +193,9 @@ void MLARopeGeneration(std::optional fused_q, // [tokens, num_hea TLLM_CHECK_WITH_INFO(num_kv_heads == 1, "num_kv_heads must = 1"); TLLM_CHECK_WITH_INFO(residual_dim >= 0 && residual_dim <= qk_rope_head_dim && residual_dim % 16 == 0, "MLA KV residual_dim must be a multiple of 16 in [0, qk_rope_head_dim], got %ld", residual_dim); - TORCH_CHECK(helix_tensor_params.size() == 2, - "Expecting 2 tensors for helix_tensor_params: helix_position_offsets and helix_is_inactive_rank."); + TORCH_CHECK(helix_tensor_params.size() == 2 || helix_tensor_params.size() == 3, + "Expecting 2 or 3 tensors for helix_tensor_params: helix_position_offsets, helix_is_inactive_rank " + "and optionally helix_local_slots (per-token KV write slots for speculative verify groups)."); auto stream = at::cuda::getCurrentCUDAStream(latent_cache.get_device()); auto const kv_cache_quant_mode = tc::QuantMode(uint32_t(quant_mode)); @@ -221,6 +225,14 @@ void MLARopeGeneration(std::optional fused_q, // [tokens, num_hea = helix_position_offsets.has_value() ? helix_position_offsets->data_ptr() : nullptr; bool const* helix_is_inactive_rank_ptr = helix_is_inactive_rank.has_value() ? helix_is_inactive_rank->data_ptr() : nullptr; + int32_t const* helix_local_slots_ptr = nullptr; + if (helix_tensor_params.size() == 3 && helix_tensor_params[2].has_value()) + { + helix_local_slots_ptr = helix_tensor_params[2]->data_ptr(); + TORCH_CHECK(!kv_norm_weight.has_value(), + "helix_local_slots (speculative verify groups) is not supported on the fused " + "KV-norm RoPE path: its KV append kernel has no per-token helix gate."); + } int* cu_q_seqlens_ptr = reinterpret_cast(cu_q_seqlens.data_ptr()); int* cu_kv_seqlens_ptr = reinterpret_cast(cu_kv_seqlens.data_ptr()); @@ -330,9 +342,9 @@ void MLARopeGeneration(std::optional fused_q, // [tokens, num_hea block_ids_per_seq_ptr, cache_type, kv_cache_buffers.kvScaleCacheBuffer, cu_q_seqlens_ptr, cu_kv_seqlens_ptr, fmha_tile_counter_ptr, mla_bmm1_scale_ptr, mla_bmm2_scale_ptr, quant_q_buffer_ptr, quant_scale_qkv_ptr, quant_scale_o_ptr, kv_scale_orig_quant_ptr, kv_scale_quant_orig_ptr, kv_cache_scale_orig_quant_ptr, - host_bmm1_scale, helix_position_offsets_ptr, helix_is_inactive_rank_ptr, kv_norm_weight_ptr, - static_cast(kv_norm_eps), latent_row_stride, precomputed_cu_seqlens, precomputed_fmha_scheduler, kv_only, - kv_done_elsewhere, q_rope_applied}; + host_bmm1_scale, helix_position_offsets_ptr, helix_is_inactive_rank_ptr, helix_local_slots_ptr, + kv_norm_weight_ptr, static_cast(kv_norm_eps), latent_row_stride, precomputed_cu_seqlens, + precomputed_fmha_scheduler, kv_only, kv_done_elsewhere, q_rope_applied}; void* q_pe_ptr = kv_only ? nullptr : q_pe->data_ptr(); void* fused_q_ptr = kv_only ? nullptr : fused_q->data_ptr(); From d5ac4cdefa97c4dbd54790b662966b5883ae7518 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:46:29 +0000 Subject: [PATCH 02/33] [None][feat] wire the Helix verify-group runtime through the FP4 MLA path Adds the functional half of Helix speculative verify groups, so the per-token kernel contract added alongside actually has a producer. TrtllmAttentionMetadata gains the per-token buffers and derives them on device in recompute_helix_spec_buffers: helix_local_slots (rank-local KV write slot, -1 when another rank owns the position) and helix_kv_bounds (how many rank-local entries a token may attend to). Both follow the round-robin page ledger, page b -> rank b % cp_size. model_engine packs the group's global positions host-side and reports a per-sequence count of this step's owned new tokens, since with a split group ownership is a count rather than a boolean. Under the overlap scheduler the host packs from a stale base, so the same accepted-count correction already applied to position_ids is applied before the recompute, and mirrored back for capture symmetry. The FP4 MLA backend consumes the bounds end to end: validation and the USE_HELIX/USE_HELIX_LOCAL_SLOTS specializations in the Triton append kernel, and the attention mask in both CuteDSL MuFu16 variants, which also now emit the softmax row stats the cross-rank combine needs. MLA derives its zero-KV mask from the per-token bounds when they are valid, fixing a real hazard: a rank holding only a group's tail page has zero visible KV for the leading tokens while its per-sequence kv_len is nonzero, so the per-sequence mask missed those rows and the combine could multiply an uninitialized partial by a zero correction. The all-to-all now forwards zero_kv_mask so the sanitize happens exactly once per backend, and the NCCL path folds it into the reformat. The CuteDSL MLA decode backend keeps its existing single-token Helix restriction; extending it needs an op-schema change entangled with unrelated work, so it stays a follow-up and continues to reject multi-token Helix loudly rather than silently mis-masking. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/attention/attention.py | 70 +++++++-- .../_torch/attention/backends/fmha/fp4_mla.py | 1 + .../attention/backends/fp4_mla/__init__.py | 133 +++++++++++++++++- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 124 ++++++++++++++-- ...p4_mla_cutedsl_mufu16_fused_v_transpose.py | 124 ++++++++++++++-- .../backends/fp4_mla/fp4_mla_kernels.py | 54 ++++++- .../_torch/attention/backends/trtllm.py | 130 ++++++++++++++++- tensorrt_llm/_torch/attention/mla.py | 32 +++-- .../_torch/custom_ops/cpp_custom_ops.py | 7 +- tensorrt_llm/_torch/distributed/ops.py | 10 +- .../_torch/pyexecutor/model_engine.py | 92 +++++++++++- .../test_helix_postprocess.py | 39 ++++- .../attention/multi_gpu/test_mla_helix.py | 11 +- 13 files changed, 767 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/attention/attention.py b/tensorrt_llm/_torch/attention/attention.py index 94344ff8247b..7565481f54e4 100644 --- a/tensorrt_llm/_torch/attention/attention.py +++ b/tensorrt_llm/_torch/attention/attention.py @@ -193,6 +193,46 @@ def _helix_sanitize_empty_kv( return partial_o, softmax_stats +@torch.compile(dynamic=False) +def _helix_nccl_pre_alltoall( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + zero_kv_mask: Optional[torch.Tensor], + cp_size: int, +) -> List[torch.Tensor]: + """Sanitize zero-local-KV rows and reformat into the alltoall send layout. + + ``_helix_sanitize_empty_kv`` followed by the transpose and split + ``_helix_post_process`` used to do inline. Both live in one compiled region + so inductor folds the fill into the transposed store instead of writing + ``partial_o`` and reading it straight back. Dynamo inlines the call, so the + sanitize has exactly one definition and the fusion is unaffected. + + ``dynamic=False`` is deliberate. With dynamic shapes inductor emits a much + slower transposed store, and picks different grids on different ranks for + the same shape, which cancels the win at the collective. + + The cost is one specialization per CUDA-graph batch bucket. Exceeding + dynamo's ``cache_size_limit`` falls back to eager SILENTLY -- the symptom is + ``triton_poi_fused_*`` disappearing from the trace, not an error. + """ + partial_o, softmax_stats = _helix_sanitize_empty_kv(partial_o, + softmax_stats, + zero_kv_mask) + chunks = [] + for t in (partial_o, softmax_stats): + t = t.transpose(1, 0).contiguous() + chunks.extend(torch.split(t, t.shape[0] // cp_size)) + return chunks + + +@torch.compile(dynamic=False) +def _helix_nccl_post_alltoall( + gathered: List[torch.Tensor]) -> List[torch.Tensor]: + """Reformat the gathered partials into the helix_post_process layout.""" + return [t.transpose(1, 2).contiguous() for t in gathered] + + def _helix_post_process( partial_o: torch.Tensor, softmax_stats: torch.Tensor, @@ -210,25 +250,23 @@ def _helix_post_process( dimension that differs between the two callers is *value_dim* (``head_dim`` for MHA, ``kv_lora_rank`` for MLA). - zero_kv_mask marks tokens for which this CP rank owns no KV blocks; those rows - are forced to a no-op contribution before the exchange (see - _helix_sanitize_empty_kv). + zero_kv_mask marks tokens this CP rank owns no KV for; those rows are forced + to a no-op contribution before the exchange, exactly once per backend: + NCCL in _helix_nccl_pre_alltoall, fused with the reformat + fifo v2 in the all-to-all sender, while the entry is in shared memory + fifo v1 here, via _helix_sanitize_empty_kv When *aux_stream* and *ln_events* are provided the two ``.contiguous()`` calls in the FIFO-v1 path are overlapped on separate CUDA streams for better performance. """ - partial_o, softmax_stats = _helix_sanitize_empty_kv(partial_o, - softmax_stats, - zero_kv_mask) if mapping.cp_config.get("use_nccl_for_alltoall", True): - # NCCL-based implementation using alltoall_helix. - chunks = [] - for t in [partial_o, softmax_stats]: - t = t.transpose(1, 0).contiguous() - chunks.extend(torch.split(t, t.shape[0] // mapping.cp_size)) + # NCCL path. Sanitize is folded into _helix_nccl_pre_alltoall so + # inductor can fuse it into the reformat. + chunks = _helix_nccl_pre_alltoall(partial_o, softmax_stats, + zero_kv_mask, mapping.cp_size) gathered = alltoall_helix(chunks, mapping.cp_group) - gathered = [t.transpose(1, 2).contiguous() for t in gathered] + gathered = _helix_nccl_post_alltoall(gathered) return torch.ops.trtllm.helix_post_process(gathered[0], gathered[1], 1.0) else: @@ -239,6 +277,8 @@ def _helix_post_process( fifo_version = mapping.cp_config.get("fifo_version", 2) if fifo_version == 1: + partial_o, softmax_stats = _helix_sanitize_empty_kv( + partial_o, softmax_stats, zero_kv_mask) def reshape_o(): return partial_o.view(num_tokens, cp_size, num_heads_tp_cp, @@ -265,12 +305,16 @@ def reshape_s(): return torch.ops.trtllm.helix_post_process_native( partial_o_out, softmax_stats_out, 1.0, 2) else: + # fifo_v2: one entry is one token, and the sender already streams + # every byte through shared memory, so the sanitize rides along + # there instead of 6 separate elementwise kernels (~40% of the block). partial_o = partial_o.view(num_tokens, cp_size, num_heads_tp_cp * value_dim) softmax_stats = softmax_stats.view(num_tokens, cp_size, num_heads_tp_cp * 2) partial_o_out, softmax_stats_out = helix.alltoall_native( - partial_o, softmax_stats) + partial_o, softmax_stats, + None if zero_kv_mask is None else zero_kv_mask[:num_tokens]) gathered_o = partial_o_out.view(num_tokens, cp_size, num_heads_tp_cp, value_dim) gathered_stats = softmax_stats_out.view(num_tokens, cp_size, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index ac17ce283141..b2703f364708 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -380,6 +380,7 @@ def run_mla_generation(self, params: FmhaParams) -> None: prequantized_q=prequantized_q, prequantized_q_sf=prequantized_q_sf, q_batch_capacity=q_batch_capacity, + softmax_stats_tensor=params.fwd.softmax_stats_tensor, ) finally: metadata._fp4_mla_prequantized_q = None diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py index ccaea0d103a9..7a310f0e9593 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py @@ -614,7 +614,6 @@ def configure_fp4_mla_device_page_table( and int(getattr(metadata, "beam_width", 1)) == 1 and not bool(getattr(metadata, "is_spec_dec_tree", False)) and not bool(getattr(metadata, "locality_domain_enabled", False)) - and not bool(getattr(metadata, "enable_helix", False)) and int(getattr(kv_cache_manager, "tokens_per_block", 0) or 0) == FP4_MLA_TOKENS_PER_BLOCK and max_page_capacity > 0 and page_index_scale > 0 @@ -1300,10 +1299,14 @@ def rebuild_fp4_mla_disagg_imported_cache( or not callable(getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None)) ): return False - if not isinstance(prompt_len, int) or prompt_len <= 0: + if not isinstance(prompt_len, int) or prompt_len < 0: raise ValueError( - f"FP4 MLA disaggregated import needs a positive prompt_len, got {prompt_len}." + f"FP4 MLA disaggregated import needs a nonnegative prompt_len, got {prompt_len}." ) + # Helix assigns whole pages round-robin, so a rank may own no prompt pages. + # There are no process-local V sidecars to rebuild on that rank. + if prompt_len == 0: + return True page_size = int(kv_cache_manager.tokens_per_block) if page_size != FP4_MLA_TOKENS_PER_BLOCK: @@ -1778,6 +1781,8 @@ def _scatter_fp4_mla_kv_cache_2d_generation( q_sf_out: torch.Tensor, v_packed_base: Optional[torch.Tensor], v_page_offset: int, + helix_position_offsets: Optional[torch.Tensor], + helix_is_inactive_rank: Optional[torch.Tensor], ) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: num_contexts = metadata.num_contexts num_seqs = metadata.num_seqs @@ -1820,6 +1825,59 @@ def _scatter_fp4_mla_kv_cache_2d_generation( num_hp_pages = pool.shape[0] max_gen_len = num_tokens // num_gen + use_helix = helix_position_offsets is not None or helix_is_inactive_rank is not None + use_helix_local_slots = use_helix and bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + if use_helix: + if helix_position_offsets is None or helix_is_inactive_rank is None: + raise RuntimeError( + "FP4 MLA Helix requires both position-offset and inactive-rank metadata." + ) + if max_gen_len != 1 and not use_helix_local_slots: + raise NotImplementedError( + "FP4 MLA multi-token Helix requires speculative per-token metadata." + ) + if ( + helix_position_offsets.dtype != torch.int32 + or helix_position_offsets.device != latent_cache.device + or helix_position_offsets.ndim != 1 + or helix_position_offsets.numel() < num_tokens + or not helix_position_offsets.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix position offsets must be a contiguous same-device " + "int32 tensor covering every generation token." + ) + if ( + helix_is_inactive_rank.dtype != torch.bool + or helix_is_inactive_rank.device != latent_cache.device + or helix_is_inactive_rank.ndim != 1 + or helix_is_inactive_rank.numel() < num_gen + or not helix_is_inactive_rank.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix inactive-rank metadata must be a contiguous " + "same-device bool tensor covering every generation sequence." + ) + if use_helix_local_slots: + helix_local_slots = getattr(metadata, "helix_local_slots", None) + if ( + not isinstance(helix_local_slots, torch.Tensor) + or helix_local_slots.dtype != torch.int32 + or helix_local_slots.device != latent_cache.device + or helix_local_slots.ndim != 1 + or helix_local_slots.numel() < num_tokens + or not helix_local_slots.is_contiguous() + ): + raise ValueError( + "FP4 MLA speculative Helix local slots must be a contiguous " + "same-device int32 tensor covering every generation token." + ) + else: + helix_local_slots = helix_position_offsets + else: + helix_position_offsets = kv_lens_gen + helix_local_slots = kv_lens_gen + helix_is_inactive_rank = gen_lens_gen _validate_fp4_mla_hp_generation_width(hp_pool_size, max_gen_len) max_rewind_len = hp_pool_size - HP_BLOCK_SIZE page_ids = _fp4_mla_generation_page_ids(metadata, num_gen) @@ -1933,6 +1991,9 @@ def launch_generation_update( q_sf_output, kv_lens_gen, gen_lens_gen, + helix_position_offsets, + helix_local_slots, + helix_is_inactive_rank, page_ids, hp_page_ids, metadata.paged_kv_indptr_decode, @@ -1971,6 +2032,8 @@ def launch_generation_update( K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, STORE_K_RESIDUAL=store_k_residual, FUSE_ROPE_CACHE_STORE=True, + USE_HELIX=use_helix, + USE_HELIX_LOCAL_SLOTS=use_helix_local_slots, WRITE_V_PACKED=write_v_packed, MAX_GEN_TILES=max_gen_tiles_variant, ROPE_DIM=rope_dim, @@ -2025,6 +2088,8 @@ def launch_generation_update( metadata.page_size, write_v_packed, store_k_residual, + use_helix, + use_helix_local_slots, tuple(q1_variants), multi_token_tiles, tuple(kv_cache.stride()), @@ -2306,6 +2371,8 @@ def scatter_fp4_mla_kv_cache( q_pe: Optional[torch.Tensor] = None, q_rope_out: Optional[torch.Tensor] = None, q_quant_input: Optional[torch.Tensor] = None, + helix_position_offsets: Optional[torch.Tensor] = None, + helix_is_inactive_rank: Optional[torch.Tensor] = None, q_context: Optional[torch.Tensor] = None, q_nope_head_dim: Optional[int] = None, ) -> bool: @@ -2423,6 +2490,11 @@ def scatter_fp4_mla_kv_cache( else: if q_context is not None or q_nope_head_dim is not None: raise ValueError("FP4 MLA generation cache update does not accept context Q tensors.") + if (helix_position_offsets is None) != (helix_is_inactive_rank is None): + raise ValueError( + "FP4 MLA Helix position-offset and inactive-rank metadata " + "must be provided together." + ) if not all(arg is not None for arg in generation_inputs): raise ValueError( "FP4 MLA generation requires rotary_cos_sin, q_pe, q_rope_out, " @@ -2561,6 +2633,8 @@ def scatter_fp4_mla_kv_cache( q_sf_out=q_sf_out, v_packed_base=v_packed_base, v_page_offset=v_page_offset, + helix_position_offsets=helix_position_offsets, + helix_is_inactive_rank=helix_is_inactive_rank, ) v_pack_page_ids = _fp4_mla_generation_page_ids( metadata, metadata.num_seqs - metadata.num_contexts @@ -3882,6 +3956,7 @@ def run_fp4_mla_attention_decode( prequantized_q: torch.Tensor, prequantized_q_sf: torch.Tensor, q_batch_capacity: int, + softmax_stats_tensor: Optional[torch.Tensor] = None, ) -> None: """Run MLA decode with FP4 QK and FP4 PV tensor-core matmuls. @@ -3929,6 +4004,43 @@ def run_fp4_mla_attention_decode( raise ValueError("FP4 MLA attention output batch dimensions do not match.") backend = _fp4_mla_attention_backend() + helix_spec_tokens_valid = bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + helix_kv_bounds = None + if softmax_stats_tensor is not None: + if backend != _FP4_MLA_CUTEDSL_BACKEND: + raise NotImplementedError( + "FP4 MLA Helix softmax stats require the cutedsl attention backend." + ) + if query_len_per_seq != 1 and not helix_spec_tokens_valid: + raise NotImplementedError( + "FP4 MLA multi-token Helix requires speculative per-token metadata." + ) + expected_stats_shape = (num_queries, num_heads, 2) + if ( + softmax_stats_tensor.shape != expected_stats_shape + or softmax_stats_tensor.dtype != torch.float32 + or softmax_stats_tensor.device != q.device + or not softmax_stats_tensor.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix requires contiguous same-device float32 softmax " + f"stats with shape {expected_stats_shape}." + ) + if helix_spec_tokens_valid: + helix_kv_bounds = getattr(metadata, "helix_kv_bounds", None) + if ( + not isinstance(helix_kv_bounds, torch.Tensor) + or helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.device != q.device + or helix_kv_bounds.ndim != 1 + or helix_kv_bounds.numel() < num_queries + or not helix_kv_bounds.is_contiguous() + ): + raise ValueError( + "FP4 MLA speculative Helix KV bounds must be a contiguous " + "same-device int32 tensor covering every query token." + ) + helix_kv_bounds = helix_kv_bounds[:num_queries] if getattr(metadata, "fp4_mla_v_scale_pool", None) is None: raise RuntimeError( "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." @@ -4131,6 +4243,15 @@ def run_fp4_mla_attention_decode( ) kernel_output = output + kernel_softmax_stats = None + if softmax_stats_tensor is not None: + kernel_softmax_stats = _ensure_workspace_tensor( + metadata, + "_fp4_mla_cutedsl_softmax_stats_buf", + (2, num_queries, physical_heads), + dtype=torch.float32, + device=output.device, + ) if num_heads < physical_heads: kernel_output = _ensure_workspace_tensor( metadata, @@ -4161,9 +4282,15 @@ def run_fp4_mla_attention_decode( v_page_offset=v_page_offset, q_batch_capacity=q_batch_capacity, partition_runtime_valid_k=bool(getattr(metadata, "is_cuda_graph", False)), + softmax_row_max=(None if kernel_softmax_stats is None else kernel_softmax_stats[0]), + softmax_row_sum=(None if kernel_softmax_stats is None else kernel_softmax_stats[1]), + helix_kv_bounds=helix_kv_bounds, ) if kernel_output is not output: output.copy_(kernel_output[:, :num_heads]) + if kernel_softmax_stats is not None: + softmax_stats_tensor[..., 0].copy_(kernel_softmax_stats[0, :, :num_heads]) + softmax_stats_tensor[..., 1].copy_(kernel_softmax_stats[1, :, :num_heads]) return total_p_rows = num_queries * max_pages * num_heads diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 27f73ea52833..7187a8601215 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -1068,6 +1068,7 @@ def fused_fp4_mla_decode_ctm( sfb_ptr: cute.Pointer, page_table_ptr: cute.Pointer, valid_k_ptr: cute.Pointer, + helix_kv_bounds_ptr: cute.Pointer, c_ptr: cute.Pointer, accum_ptr: cute.Pointer, row_max_ptr: cute.Pointer, @@ -1091,9 +1092,11 @@ def fused_fp4_mla_decode_ctm( page_size: ctm.Constexpr = KV_TILE, use_mixed_imlp: ctm.Constexpr = False, query_len_per_seq: ctm.Constexpr = 1, + use_helix_kv_bounds: ctm.Constexpr = False, use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, ) -> None: n, k = problem_size m = runtime_m @@ -1132,6 +1135,10 @@ def fused_fp4_mla_decode_ctm( cute.recast_ptr(valid_k_ptr, dtype=cutlass.Int32), cute.make_layout((batch_size // query_len_per_seq,), stride=(1,)), ) + helix_kv_bounds_tensor = cute.make_tensor( + cute.recast_ptr(helix_kv_bounds_ptr, dtype=cutlass.Int32), + cute.make_layout((batch_size,), stride=(1,)), + ) v_tma_tensor = cute.make_tensor( b_ptr, cute.make_layout( @@ -1388,6 +1395,7 @@ def fused_fp4_mla_decode_ctm( page_table_tensor, page_indptr_tensor, valid_k_tensor, + helix_kv_bounds_tensor, q_global_scale_tensor, kv_global_scale_tensor, tma_q_desc, @@ -1415,10 +1423,12 @@ def fused_fp4_mla_decode_ctm( pv_output_scale, page_size, query_len_per_seq, + use_helix_kv_bounds, use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, ).launch( grid=( cute.ceil_div(c_tensor.shape[0], SMEM_P4_CTA_GROUP_M) * CLUSTER_SHAPE_MNK[0], @@ -4628,6 +4638,7 @@ def _store_final_o_from_tmem( final_row_sum: ctm.Float32, final_stat_scale: ctm.Float32, output_normalizer: ctm.Float32, + write_softmax_stats: ctm.Constexpr = False, producer_warp_base: ctm.Constexpr = 0, ) -> None: gC_arr = ctm.make_array_view(mC_mnl) @@ -4647,6 +4658,10 @@ def _store_final_o_from_tmem( + local_row ) batch_offset = bidz * m * n + if cutlass.const_expr(write_softmax_stats): + if col_band == ctm.Int32(0) and bidy == ctm.Int32(0): + mRowMax_ml[row, bidz] = final_row_max * final_stat_scale + mRowSum_ml[row, bidz] = final_row_sum n_tile_total = n // ctm.Int32(OUT_DIM) subtile_cols: ctm.Constexpr = 32 subtiles_per_n_tile: ctm.Constexpr = SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N // subtile_cols @@ -5908,6 +5923,7 @@ def _run_mla_decode_body( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_ptr, @@ -5935,10 +5951,12 @@ def _run_mla_decode_body( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr = False, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, ) -> None: pv_psf_rescale = ctm.Float32(FP4_MLA_P_GLOBAL_SCALE) * pv_output_scale warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -5949,11 +5967,14 @@ def _run_mla_decode_body( page_batch = ctm.Int32(0) valid_k_for_l = ctm.Int32(0) page_batch = cute.arch.make_warp_uniform(bidz // ctm.Int32(query_len_per_seq)) - query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) - valid_k_for_l = ctm.max( - mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), - ctm.Int32(0), - ) + if cutlass.const_expr(use_helix_kv_bounds): + valid_k_for_l = mHelixKvBounds_l[bidz] + else: + query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) + valid_k_for_l = ctm.max( + mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), + ctm.Int32(0), + ) valid_k_for_l = cute.arch.make_warp_uniform(valid_k_for_l) csr_page_begin = ctm.Int32(0) csr_page_count = ctm.Int32(0) @@ -6904,8 +6925,19 @@ def _run_mla_decode_body( final_anchor_row_sum = sFinalAnchorRowSum[final_row_state_local_row] final_stat_scale = ctm.Float32(1.0) final_row_sum = final_anchor_row_sum + final_row_max = running_row_max + if cutlass.const_expr(write_softmax_stats): + final_stat_scale = softmax_scale_log2 / ctm.Float32(LOG2_E) + final_row_sum = final_anchor_row_sum * cute.exp2( + (running_row_anchor - running_row_max) * softmax_scale_log2, + fastmath=True, + ) + if valid_k_for_l == ctm.Int32(0): + final_stat_scale = ctm.Float32(1.0) + final_row_max = ctm.Float32(-ctm.Float32.inf) + final_row_sum = ctm.Float32(0.0) output_normalizer = ctm.Float32(0.0) - if final_anchor_row_sum != ctm.Float32(0.0): + if valid_k_for_l != ctm.Int32(0) and final_anchor_row_sum != ctm.Float32(0.0): output_normalizer = cute.arch.rcp_approx(final_anchor_row_sum) * pv_output_scale last_pv_li_idx = stream_li_total - ctm.Int32(1) last_pv_slot = last_pv_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -6931,10 +6963,11 @@ def _run_mla_decode_body( bidz, m, n, - final_row_max=running_row_max, + final_row_max=final_row_max, final_row_sum=final_row_sum, final_stat_scale=final_stat_scale, output_normalizer=output_normalizer, + write_softmax_stats=write_softmax_stats, producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, ) prims.barrier(barrier_id=O_STORE_BAR_ID, number_of_threads=O_STORE_BAR_THREADS) @@ -6953,6 +6986,7 @@ def kernel( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_desc: ctm.GridConstant[cuda_tma.TensorMap], @@ -6980,15 +7014,18 @@ def kernel( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr, use_ksf_gather4: ctm.Constexpr, + write_softmax_stats: ctm.Constexpr, ) -> None: _run_mla_decode_body( mPageTable_pl, mPageIndptr_s, mValidK_l, + mHelixKvBounds_l, mQGlobalScale, mKvGlobalScale, tma_q_desc.get_ptr(), @@ -7016,10 +7053,12 @@ def kernel( pv_output_scale, page_size, query_len_per_seq, + use_helix_kv_bounds, use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, ) @@ -7040,6 +7079,7 @@ def _make_fused_ptrs( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7059,6 +7099,7 @@ def _make_fused_ptrs( make_ptr(cutlass.Uint8, sfb_data_ptr, cute.AddressSpace.gmem, assumed_align=32), make_ptr(cutlass.Int32, page_table_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr(cutlass.Int32, valid_k_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr(cutlass.Int32, helix_kv_bounds_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr( _cutlass_output_dtype(output_dtype), c_data_ptr, @@ -7248,6 +7289,7 @@ def _compile_fused( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7269,6 +7311,8 @@ def _compile_fused( ksf_page_stride_bytes: int = 0, vsf_page_stride_bytes: int = 0, use_consecutive_page_pair: bool = False, + write_softmax_stats: bool = False, + use_helix_kv_bounds: bool = False, ) -> Callable: if type(kv) is not int: raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") @@ -7295,6 +7339,8 @@ def _compile_fused( query_len_per_seq, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, + use_helix_kv_bounds, ) cached = _FUSED_COMPILE_CACHE.get(cache_key) if cached is not None: @@ -7309,6 +7355,7 @@ def _compile_fused( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, accum_data_ptr, row_max_data_ptr, @@ -7337,9 +7384,11 @@ def _compile_fused( page_size=page_size, use_mixed_imlp=use_mixed_imlp, query_len_per_seq=query_len_per_seq, + use_helix_kv_bounds=use_helix_kv_bounds, use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, + write_softmax_stats=write_softmax_stats, options="--opt-level 2 --ptxas-options '--uumn'", ) _FUSED_COMPILE_CACHE[cache_key] = compiled @@ -7371,6 +7420,9 @@ def run_trtllm_fp4_mla_decode_page_native( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7417,6 +7469,24 @@ def run_trtllm_fp4_mla_decode_page_native( f"got dtype={q_internal.dtype} shape={tuple(q_internal.shape)}" ) physical_m, q_bytes, l_batch = q_internal.shape + if (softmax_row_max is None) != (softmax_row_sum is None): + raise ValueError("softmax_row_max and softmax_row_sum must be provided together") + write_softmax_stats = softmax_row_max is not None + if write_softmax_stats: + expected_stats_shape = (l_batch, physical_m) + for name, tensor in ( + ("softmax_row_max", softmax_row_max), + ("softmax_row_sum", softmax_row_sum), + ): + if ( + tensor.dtype != torch.float32 + or tensor.shape != expected_stats_shape + or not tensor.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous float32 with shape {expected_stats_shape}" + ) + _validate_tensor_pointer_alignment(name, tensor, alignment_bytes=32) if l_batch <= 0 or l_batch > CUDA_GRID_Z_MAX: raise ValueError(f"queries must be in [1, {CUDA_GRID_Z_MAX}], got {l_batch}") if q_batch_capacity is None: @@ -7496,6 +7566,16 @@ def run_trtllm_fp4_mla_decode_page_native( f"got dtype={valid_k.dtype} shape={tuple(valid_k.shape)} " f"stride={valid_k.stride()}" ) + if helix_kv_bounds is not None and ( + helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.shape != (l_batch,) + or helix_kv_bounds.stride(0) != 1 + ): + raise ValueError( + "helix_kv_bounds must be contiguous int32 with shape " + f"[{l_batch}], got dtype={helix_kv_bounds.dtype} " + f"shape={tuple(helix_kv_bounds.shape)} stride={helix_kv_bounds.stride()}" + ) cache_layout = _validate_v_packed_cache_args( v_packed, kv_cache, @@ -7612,6 +7692,10 @@ def run_trtllm_fp4_mla_decode_page_native( q_global_scale, kv_global_scale, ) + if write_softmax_stats: + tensors += (softmax_row_max, softmax_row_sum) + if helix_kv_bounds is not None: + tensors += (helix_kv_bounds,) if device.type != "cuda" or any((tensor.device != device for tensor in tensors)): raise ValueError("all page-native decode tensors must share one CUDA device") if q_global_scale.dtype != torch.float32 or q_global_scale.numel() != 1: @@ -7631,9 +7715,14 @@ def run_trtllm_fp4_mla_decode_page_native( k_sf_data_ptr = sf_cache.data_ptr() b_data_ptr = v_packed.data_ptr() scratch_ptr = output.data_ptr() + row_max_data_ptr = softmax_row_max.data_ptr() if write_softmax_stats else scratch_ptr + row_sum_data_ptr = softmax_row_sum.data_ptr() if write_softmax_stats else scratch_ptr sfb_data_ptr = v_sf.data_ptr() page_table_data_ptr = src_page_ids.data_ptr() valid_k_data_ptr = valid_k.data_ptr() + helix_kv_bounds_data_ptr = ( + helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else valid_k_data_ptr + ) c_data_ptr = output.data_ptr() page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() q_global_scale_data_ptr = q_global_scale.data_ptr() @@ -7652,10 +7741,11 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, scratch_ptr, - scratch_ptr, - scratch_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -7673,6 +7763,8 @@ def run_trtllm_fp4_mla_decode_page_native( ksf_page_stride_bytes=ksf_page_stride_bytes, vsf_page_stride_bytes=vsf_page_stride_bytes, use_consecutive_page_pair=use_consecutive_page_pair, + write_softmax_stats=write_softmax_stats, + use_helix_kv_bounds=helix_kv_bounds is not None, ) supports_prepared = _class_defines_callables( fused, "to", "generate_execution_args", "run_compiled_program" @@ -7690,7 +7782,10 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -7723,10 +7818,11 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, scratch_ptr, - scratch_ptr, - scratch_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -7805,6 +7901,9 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7856,4 +7955,7 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles=assume_consecutive_page_prefix_tiles, partition_runtime_valid_k=partition_runtime_valid_k, enable_mxi_imlp=enable_mxi_imlp, + softmax_row_max=softmax_row_max, + softmax_row_sum=softmax_row_sum, + helix_kv_bounds=helix_kv_bounds, ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py index 9ff1c2f656b1..0c2f590462e3 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py @@ -1062,6 +1062,7 @@ def fused_fp4_mla_decode_ctm( sfb_ptr: cute.Pointer, page_table_ptr: cute.Pointer, valid_k_ptr: cute.Pointer, + helix_kv_bounds_ptr: cute.Pointer, c_ptr: cute.Pointer, accum_ptr: cute.Pointer, row_max_ptr: cute.Pointer, @@ -1085,9 +1086,11 @@ def fused_fp4_mla_decode_ctm( page_size: ctm.Constexpr = KV_TILE, use_mixed_imlp: ctm.Constexpr = False, query_len_per_seq: ctm.Constexpr = 1, + use_helix_kv_bounds: ctm.Constexpr = False, use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, ) -> None: n, k = problem_size m = runtime_m @@ -1134,6 +1137,10 @@ def fused_fp4_mla_decode_ctm( cute.recast_ptr(valid_k_ptr, dtype=cutlass.Int32), cute.make_layout((l // query_len_per_seq,), stride=(1,)), ) + helix_kv_bounds_tensor = cute.make_tensor( + cute.recast_ptr(helix_kv_bounds_ptr, dtype=cutlass.Int32), + cute.make_layout((l,), stride=(1,)), + ) v_tma_tensor = cute.make_tensor( b_ptr, cute.make_layout( @@ -1396,6 +1403,7 @@ def fused_fp4_mla_decode_ctm( page_table_tensor, page_indptr_tensor, valid_k_tensor, + helix_kv_bounds_tensor, q_global_scale_tensor, kv_global_scale_tensor, tma_q_desc, @@ -1424,10 +1432,12 @@ def fused_fp4_mla_decode_ctm( pv_output_scale, page_size, query_len_per_seq, + use_helix_kv_bounds, use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, ).launch( grid=( cute.ceil_div(c_tensor.shape[0], SMEM_P4_CTA_GROUP_M) * CLUSTER_SHAPE_MNK[0], @@ -5064,6 +5074,7 @@ def _store_final_o_from_tmem( final_row_sum: ctm.Float32, final_stat_scale: ctm.Float32, output_normalizer: ctm.Float32, + write_softmax_stats: ctm.Constexpr = False, producer_warp_base: ctm.Constexpr = 0, ) -> None: gC_arr = ctm.make_array_view(mC_mnl) @@ -5085,6 +5096,10 @@ def _store_final_o_from_tmem( + local_row ) batch_offset = bidz * m * n + if cutlass.const_expr(write_softmax_stats): + if col_band == ctm.Int32(0) and bidy == ctm.Int32(0): + mRowMax_ml[row, bidz] = final_row_max * final_stat_scale + mRowSum_ml[row, bidz] = final_row_sum n_tile_total = n // ctm.Int32(OUT_DIM) subtile_cols: ctm.Constexpr = 32 subtiles_per_n_tile: ctm.Constexpr = SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N // subtile_cols @@ -6368,6 +6383,7 @@ def _run_mla_decode_body( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_ptr, @@ -6396,10 +6412,12 @@ def _run_mla_decode_body( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr = False, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, ) -> None: pv_psf_rescale = ctm.Float32(FP4_MLA_P_GLOBAL_SCALE) * pv_output_scale warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -6410,11 +6428,14 @@ def _run_mla_decode_body( page_batch = ctm.Int32(0) valid_k_for_l = ctm.Int32(0) page_batch = cute.arch.make_warp_uniform(bidz // ctm.Int32(query_len_per_seq)) - query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) - valid_k_for_l = ctm.max( - mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), - ctm.Int32(0), - ) + if cutlass.const_expr(use_helix_kv_bounds): + valid_k_for_l = mHelixKvBounds_l[bidz] + else: + query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) + valid_k_for_l = ctm.max( + mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), + ctm.Int32(0), + ) valid_k_for_l = cute.arch.make_warp_uniform(valid_k_for_l) csr_page_begin = ctm.Int32(0) csr_page_count = ctm.Int32(0) @@ -7403,8 +7424,19 @@ def _run_mla_decode_body( final_anchor_row_sum = sFinalAnchorRowSum[final_row_state_local_row] final_stat_scale = ctm.Float32(1.0) final_row_sum = final_anchor_row_sum + final_row_max = running_row_max + if cutlass.const_expr(write_softmax_stats): + final_stat_scale = softmax_scale_log2 / ctm.Float32(LOG2_E) + final_row_sum = final_anchor_row_sum * cute.exp2( + (running_row_anchor - running_row_max) * softmax_scale_log2, + fastmath=True, + ) + if valid_k_for_l == ctm.Int32(0): + final_stat_scale = ctm.Float32(1.0) + final_row_max = ctm.Float32(-ctm.Float32.inf) + final_row_sum = ctm.Float32(0.0) output_normalizer = ctm.Float32(0.0) - if final_anchor_row_sum != ctm.Float32(0.0): + if valid_k_for_l != ctm.Int32(0) and final_anchor_row_sum != ctm.Float32(0.0): output_normalizer = cute.arch.rcp_approx(final_anchor_row_sum) * pv_output_scale last_pv_li_idx = stream_li_total - ctm.Int32(1) last_pv_slot = last_pv_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -7430,10 +7462,11 @@ def _run_mla_decode_body( bidz, m, n, - final_row_max=running_row_max, + final_row_max=final_row_max, final_row_sum=final_row_sum, final_stat_scale=final_stat_scale, output_normalizer=output_normalizer, + write_softmax_stats=write_softmax_stats, producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, ) prims.barrier(barrier_id=O_STORE_BAR_ID, number_of_threads=O_STORE_BAR_THREADS) @@ -7452,6 +7485,7 @@ def kernel( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_desc: ctm.GridConstant[cuda_tma.TensorMap], @@ -7480,15 +7514,18 @@ def kernel( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr, use_ksf_gather4: ctm.Constexpr, + write_softmax_stats: ctm.Constexpr, ) -> None: _run_mla_decode_body( mPageTable_pl, mPageIndptr_s, mValidK_l, + mHelixKvBounds_l, mQGlobalScale, mKvGlobalScale, tma_q_desc.get_ptr(), @@ -7517,10 +7554,12 @@ def kernel( pv_output_scale, page_size, query_len_per_seq, + use_helix_kv_bounds, use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, ) @@ -7541,6 +7580,7 @@ def _make_fused_ptrs( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7560,6 +7600,7 @@ def _make_fused_ptrs( make_ptr(cutlass.Uint8, sfb_data_ptr, cute.AddressSpace.gmem, assumed_align=32), make_ptr(cutlass.Int32, page_table_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr(cutlass.Int32, valid_k_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr(cutlass.Int32, helix_kv_bounds_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr( _cutlass_output_dtype(output_dtype), c_data_ptr, @@ -7695,6 +7736,7 @@ def _compile_fused( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7716,6 +7758,8 @@ def _compile_fused( ksf_page_stride_bytes: int = 0, vsf_page_stride_bytes: int = 0, use_consecutive_page_pair: bool = False, + write_softmax_stats: bool = False, + use_helix_kv_bounds: bool = False, ) -> Callable: if type(kv) is not int: raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") @@ -7740,6 +7784,8 @@ def _compile_fused( query_len_per_seq, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, + use_helix_kv_bounds, ) cached = _FUSED_COMPILE_CACHE.get(cache_key) if cached is not None: @@ -7754,6 +7800,7 @@ def _compile_fused( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, accum_data_ptr, row_max_data_ptr, @@ -7782,9 +7829,11 @@ def _compile_fused( page_size=page_size, use_mixed_imlp=use_mixed_imlp, query_len_per_seq=query_len_per_seq, + use_helix_kv_bounds=use_helix_kv_bounds, use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, + write_softmax_stats=write_softmax_stats, options="--opt-level 2 --ptxas-options '--uumn'", ) _FUSED_COMPILE_CACHE[cache_key] = compiled @@ -7816,6 +7865,9 @@ def run_trtllm_fp4_mla_decode_page_native( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7861,6 +7913,24 @@ def run_trtllm_fp4_mla_decode_page_native( f"q_internal must be a uint8 [M, Q640/2, L] tensor, got dtype={q_internal.dtype} shape={tuple(q_internal.shape)}" ) physical_m, q_bytes, l_batch = q_internal.shape + if (softmax_row_max is None) != (softmax_row_sum is None): + raise ValueError("softmax_row_max and softmax_row_sum must be provided together") + write_softmax_stats = softmax_row_max is not None + if write_softmax_stats: + expected_stats_shape = (l_batch, physical_m) + for name, tensor in ( + ("softmax_row_max", softmax_row_max), + ("softmax_row_sum", softmax_row_sum), + ): + if ( + tensor.dtype != torch.float32 + or tensor.shape != expected_stats_shape + or not tensor.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous float32 with shape {expected_stats_shape}" + ) + _validate_tensor_pointer_alignment(name, tensor, alignment_bytes=32) if l_batch <= 0 or l_batch > CUDA_GRID_Z_MAX: raise ValueError(f"queries must be in [1, {CUDA_GRID_Z_MAX}], got {l_batch}") if q_batch_capacity is None: @@ -7933,6 +8003,16 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"valid_k must be contiguous int32 [{num_sequences}], got dtype={valid_k.dtype} shape={tuple(valid_k.shape)} stride={valid_k.stride()}" ) + if helix_kv_bounds is not None and ( + helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.shape != (l_batch,) + or helix_kv_bounds.stride(0) != 1 + ): + raise ValueError( + "helix_kv_bounds must be contiguous int32 with shape " + f"[{l_batch}], got dtype={helix_kv_bounds.dtype} " + f"shape={tuple(helix_kv_bounds.shape)} stride={helix_kv_bounds.stride()}" + ) cache_layout = _kv_cache_3d_layout(kv_cache, page_size) _validate_tensor_pointer_alignment("kv_cache", kv_cache, alignment_bytes=16) if ( @@ -8024,6 +8104,10 @@ def run_trtllm_fp4_mla_decode_page_native( q_global_scale, kv_global_scale, ) + if write_softmax_stats: + tensors += (softmax_row_max, softmax_row_sum) + if helix_kv_bounds is not None: + tensors += (helix_kv_bounds,) if device.type != "cuda" or any((tensor.device != device for tensor in tensors)): raise ValueError("all page-native decode tensors must share one CUDA device") if q_global_scale.dtype != torch.float32 or q_global_scale.numel() != 1: @@ -8045,9 +8129,14 @@ def run_trtllm_fp4_mla_decode_page_native( # canonical KV pointer so the compiled call carries no sidecar allocation. b_data_ptr = kv_cache.data_ptr() scratch_ptr = output.data_ptr() + row_max_data_ptr = softmax_row_max.data_ptr() if write_softmax_stats else scratch_ptr + row_sum_data_ptr = softmax_row_sum.data_ptr() if write_softmax_stats else scratch_ptr sfb_data_ptr = v_sf.data_ptr() page_table_data_ptr = src_page_ids.data_ptr() valid_k_data_ptr = valid_k.data_ptr() + helix_kv_bounds_data_ptr = ( + helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else valid_k_data_ptr + ) c_data_ptr = output.data_ptr() page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() q_global_scale_data_ptr = q_global_scale.data_ptr() @@ -8066,10 +8155,11 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, scratch_ptr, - scratch_ptr, - scratch_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -8087,6 +8177,8 @@ def run_trtllm_fp4_mla_decode_page_native( ksf_page_stride_bytes=ksf_page_stride_bytes, vsf_page_stride_bytes=vsf_page_stride_bytes, use_consecutive_page_pair=use_consecutive_page_pair, + write_softmax_stats=write_softmax_stats, + use_helix_kv_bounds=helix_kv_bounds is not None, ) supports_prepared = _class_defines_callables( fused, "to", "generate_execution_args", "run_compiled_program" @@ -8104,7 +8196,10 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -8137,10 +8232,11 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, scratch_ptr, - scratch_ptr, - scratch_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -8219,6 +8315,9 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -8270,4 +8369,7 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles=assume_consecutive_page_prefix_tiles, partition_runtime_valid_k=partition_runtime_valid_k, enable_mxi_imlp=enable_mxi_imlp, + softmax_row_max=softmax_row_max, + softmax_row_sum=softmax_row_sum, + helix_kv_bounds=helix_kv_bounds, ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py index 83b32ad45da0..3cb8ed6c1f95 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py @@ -1379,6 +1379,9 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( q_sf_out_ptr, kv_lens_ptr, prompt_lens_ptr, + helix_position_offsets_ptr, + helix_local_slots_ptr, + helix_is_inactive_rank_ptr, page_ids_ptr, hp_page_ids_ptr, paged_kv_indptr_ptr, @@ -1417,6 +1420,8 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( K_RESIDUAL_D: tl.constexpr, STORE_K_RESIDUAL: tl.constexpr, FUSE_ROPE_CACHE_STORE: tl.constexpr, + USE_HELIX: tl.constexpr, + USE_HELIX_LOCAL_SLOTS: tl.constexpr, WRITE_V_PACKED: tl.constexpr, MAX_GEN_TILES: tl.constexpr, ROPE_DIM: tl.constexpr, @@ -1457,6 +1462,20 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( # Generation tokens are request-major. Deriving their absolute position # here avoids materializing per-token metadata in a CUDA graph. first_new_pos = kv_len - gen_len + rope_first_new_pos = first_new_pos + if USE_HELIX: + rope_first_new_pos = tl.load(helix_position_offsets_ptr + seq_idx * gen_len) + if USE_HELIX_LOCAL_SLOTS: + # Helix owns a contiguous rank-local subset of a linear verify group, + # but that subset can begin at any token in the group. + first_new_pos = kv_len + for token_idx in tl.static_range(0, MAX_GEN_TILES * HP_BLOCK): + local_slot = tl.load( + helix_local_slots_ptr + seq_idx * gen_len + token_idx, + mask=token_idx < gen_len, + other=-1, + ) + first_new_pos = tl.minimum(first_new_pos, tl.where(local_slot >= 0, local_slot, kv_len)) if FUSE_ROPE_CACHE_STORE: work_idx = tl.program_id(1) dim_block = work_idx @@ -1585,7 +1604,9 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( mask=row_mask[:, None], ) else: - position = (first_new_pos + q_token_idx).to(tl.int64) + position = (rope_first_new_pos + q_token_idx).to(tl.int64) + if USE_HELIX_LOCAL_SLOTS: + position = tl.load(helix_position_offsets_ptr + q_token).to(tl.int64) valid_position = position >= 0 pair_offsets = tl.arange(0, ROPE_PAIR_BLOCK) pair_mask = pair_offsets < ROPE_DIM // 2 @@ -1710,6 +1731,12 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( return else: dim_block = tl.program_id(2) + if USE_HELIX_LOCAL_SLOTS: + if first_new_pos >= kv_len: + return + elif USE_HELIX: + if tl.load(helix_is_inactive_rank_ptr + seq_idx): + return if FUSE_ROPE_CACHE_STORE and MAX_GEN_TILES == 1: if dim_block * FP4_BLOCK > V_HEAD_D: return @@ -2046,9 +2073,22 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( abs_positions = block_base_pos + token_offsets valid_tokens = abs_positions < kv_len - from_latent = abs_positions >= first_new_pos hp_slots = abs_positions % HP_POOL_SIZE - new_token_offsets = abs_positions - first_new_pos + if USE_HELIX_LOCAL_SLOTS: + from_latent = abs_positions < 0 + new_token_offsets = abs_positions - abs_positions + for token_idx in tl.static_range(0, MAX_GEN_TILES * HP_BLOCK): + local_slot = tl.load( + helix_local_slots_ptr + seq_idx * gen_len + token_idx, + mask=token_idx < gen_len, + other=-1, + ) + is_local_token = abs_positions == local_slot + from_latent = from_latent | is_local_token + new_token_offsets = tl.where(is_local_token, token_idx, new_token_offsets) + else: + from_latent = abs_positions >= first_new_pos + new_token_offsets = abs_positions - first_new_pos # Linear MTP uses a uniform generation length, so each sequence # occupies one contiguous gen_len slice in latent_cache. latent_tokens = seq_idx * gen_len + new_token_offsets @@ -2129,7 +2169,13 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( if FUSE_ROPE_CACHE_STORE and not q1_shared_main: if dim_block * FP4_BLOCK >= V_HEAD_D: - positions = abs_positions.to(tl.int64) + positions = (rope_first_new_pos + new_token_offsets).to(tl.int64) + if USE_HELIX_LOCAL_SLOTS: + positions = tl.load( + helix_position_offsets_ptr + safe_latent_tokens, + mask=valid_tokens & from_latent, + other=-1, + ).to(tl.int64) valid_position = valid_tokens & from_latent & (positions >= 0) rope_pair_offsets = (safe_even_d - V_HEAD_D) // 2 rope_dim_mask = mask_even_d & (even_d >= V_HEAD_D) & (odd_d < V_HEAD_D + ROPE_DIM) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 77e465b9e171..31df0b9768d6 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -198,6 +198,15 @@ def effective_beam_width(self) -> int: helix_is_inactive_rank: Optional[torch.Tensor] = None helix_is_inactive_rank_cpu: Optional[torch.Tensor] = None + # Per-token helix state for speculative verify groups (a 1 + draft_len + # group may straddle a ledger-page boundary onto two CP ranks, so the + # per-sequence boolean above is insufficient there). See + # recompute_helix_spec_buffers for the derivation. + helix_local_slots: Optional[torch.Tensor] = None + helix_kv_bounds: Optional[torch.Tensor] = None + helix_owned_new_tokens_cpu: Optional[torch.Tensor] = None + _helix_spec_tokens_valid: bool = False + # Block offsets for the target and draft KV caches kv_cache_block_offsets: Optional[torch.Tensor] = None host_kv_cache_block_offsets: Optional[torch.Tensor] = None @@ -574,6 +583,41 @@ def _post_init_with_buffers(self, buffers) -> None: device='cpu', pin_memory=prefer_pinned(), ) + # Per-token buffers for speculative verify groups under helix. + # A group of 1 + draft_len tokens can straddle a ledger-page + # boundary, splitting ownership between two CP ranks, so the + # per-sequence flag above is not expressive enough: + # helix_local_slots[t]: rank-local KV write slot of gen token t + # on this rank, or -1 when another rank owns its position + # (consumed by the mla_rope_generation append kernel). + # helix_kv_bounds[t]: number of rank-local KV entries token t + # may attend to, i.e. local_len(pos_t + 1) (consumed by the + # CuTe DSL MLA decode mask and the helix stats identity). + # Filled by recompute_helix_spec_buffers() on the spec path only. + self.helix_local_slots = self.get_empty( + buffers, + (self.max_num_tokens, ), + cache_name="helix_local_slots", + dtype=torch.int, + capture_graph=capture_graph, + ) + self.helix_kv_bounds = self.get_empty( + buffers, + (self.max_num_tokens, ), + cache_name="helix_kv_bounds", + dtype=torch.int, + capture_graph=capture_graph, + ) + # Host-side per-sequence count of this step's new tokens owned by + # this rank (spec path; single-token path derives it from the + # boolean flag). Consumed by prepare()'s helix kv_lens branch. + self.helix_owned_new_tokens_cpu = torch.zeros( + (self.max_num_sequences, ), + device='cpu', + dtype=torch.int, + pin_memory=prefer_pinned(), + ) + self._helix_spec_tokens_valid = False def on_update_kv_lens(self): # After changing the kv_lens/kv_lens_cuda, we may need to update other metadata. @@ -600,6 +644,7 @@ def update_helix_param( self, helix_position_offsets: List[int], helix_is_inactive_rank: List[bool], + helix_owned_new_tokens: Optional[List[int]] = None, ) -> None: """ Update helix parameters by copying into static buffers for CUDA graph compatibility. @@ -607,6 +652,10 @@ def update_helix_param( Args: helix_position_offsets: Position offsets for helix parallelism with shape (num_tokens,). helix_is_inactive_rank: Whether the current rank is inactive with shape (batch_size,). + helix_owned_new_tokens: Per-sequence count of this step's new + tokens owned by this rank (speculative verify groups; one + group may straddle a page boundary onto two ranks). None on + the single-token path, where the boolean flag carries it. """ if helix_position_offsets is not None and self.helix_position_offsets is not None: num_tokens = len(helix_position_offsets) @@ -622,6 +671,67 @@ def update_helix_param( self.helix_is_inactive_rank[:batch_size].copy_( self.helix_is_inactive_rank_cpu[:batch_size], non_blocking=True) + self._helix_spec_tokens_valid = False + if helix_owned_new_tokens is not None: + batch_size = len(helix_owned_new_tokens) + self.helix_owned_new_tokens_cpu[:batch_size].copy_( + torch.tensor(helix_owned_new_tokens, dtype=torch.int)) + self._helix_spec_tokens_valid = True + + def helix_local_len_vec(self, global_lens: torch.Tensor) -> torch.Tensor: + """Vectorized rank-local prefix length for helix round-robin pages. + + For each global sequence length g, returns the number of the first g + tokens whose ledger page lives on this CP rank (page b -> rank + b % cp_size). Mirrors KVCacheManagerV2._helix_local_len. + """ + phys = self.kv_cache_manager.tokens_per_block + cp_size = self.mapping.cp_size + cp_rank = self.mapping.cp_rank + ledger = phys * cp_size + full = torch.div(global_lens, ledger, rounding_mode='floor') + rem = global_lens - full * ledger + return full * phys + (rem - cp_rank * phys).clamp_(0, phys) + + def recompute_helix_spec_buffers(self, num_ctx_tokens: int, + num_gen_tokens: int, + tokens_per_gen_seq: int) -> None: + """Derive per-token helix buffers from (corrected) global positions. + + Called after the overlap-scheduler device correction has been applied + to helix_position_offsets, so every derived quantity reflects the + real committed length even though the host packed provisional values. + Static shapes only; safe under CUDA graph capture. + """ + pos = self.helix_position_offsets[num_ctx_tokens:num_ctx_tokens + + num_gen_tokens] + phys = self.kv_cache_manager.tokens_per_block + cp_rank = self.mapping.cp_rank + cp_size = self.mapping.cp_size + owner = torch.div(pos, phys, rounding_mode='floor') % cp_size + active = owner == cp_rank + local_before = self.helix_local_len_vec(pos) + # Scalar overload: no per-step allocation (CUDA-graph capture treats + # these ops as part of the graph; keep them allocation-free). + self.helix_local_slots[num_ctx_tokens:num_ctx_tokens + + num_gen_tokens].copy_( + torch.where(active, local_before, -1)) + self.helix_kv_bounds[num_ctx_tokens:num_ctx_tokens + + num_gen_tokens].copy_( + self.helix_local_len_vec(pos + 1)) + # Per-sequence rank-local kv length = bound of the sequence's last + # token (attention over committed + owned in-flight tokens). + assert num_gen_tokens % tokens_per_gen_seq == 0, ( + f"helix spec expects uniform verify groups: {num_gen_tokens} gen " + f"tokens not divisible by group size {tokens_per_gen_seq}") + num_gen_seqs = num_gen_tokens // tokens_per_gen_seq + last_bounds = self.helix_kv_bounds[num_ctx_tokens:num_ctx_tokens + + num_gen_tokens].view( + num_gen_seqs, + tokens_per_gen_seq)[:, -1] + self.kv_lens_cuda[self.num_contexts:self.num_contexts + + num_gen_seqs].copy_(last_bounds) + def _bind_runtime_views( self, *, @@ -752,9 +862,18 @@ def prepare(self) -> None: if self.enable_helix: # If helix is inactive, attend to the previously cached tokens only. assert cached_token_lens is not None, "cached_token_lens should be set for helix" - active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs] - kv_lens = cached_token_lens.clone() - kv_lens[active_rank] += self.seq_lens_kv[active_rank] + if getattr(self, '_helix_spec_tokens_valid', False): + # Speculative verify groups: a group may straddle a page + # boundary, so ownership of this step's new tokens is a + # per-sequence COUNT, not a boolean. Provisional host values; + # recompute_helix_spec_buffers overrides the device copy + # after the overlap correction. + kv_lens = cached_token_lens + \ + self.helix_owned_new_tokens_cpu[:self.num_seqs] + else: + active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs] + kv_lens = cached_token_lens.clone() + kv_lens[active_rank] += self.seq_lens_kv[active_rank] else: kv_lens = cached_token_lens + \ self.seq_lens_kv if cached_token_lens is not None else self.seq_lens_kv @@ -2412,6 +2531,11 @@ def mla_rope_generation( helix_tensor_params = [ metadata.helix_position_offsets, metadata.helix_is_inactive_rank ] + if getattr(metadata, '_helix_spec_tokens_valid', False): + # Speculative verify groups: per-token KV write slots (-1 = this + # rank does not own the token's position). The append kernel then + # gates and addresses per token instead of per sequence. + helix_tensor_params.append(metadata.helix_local_slots) torch.ops.trtllm.mla_rope_generation( fused_q, diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index db579743eb63..988114b225a7 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -794,14 +794,30 @@ def _attn_forward_gen( kv_lora_rank = partial_o.shape[-1] // self.num_heads_tp assert self.kv_lora_rank == kv_lora_rank - # MLA processes only the generation token slice here, so build the - # mask from the generation sequence range [num_contexts, num_seqs). - zero_kv_mask = _helix_zero_kv_mask( - attn_metadata, - partial_o.shape[0], - seq_start=attn_metadata.num_contexts, - num_seqs=attn_metadata.num_generations, - ) + helix_kv_bounds = getattr(attn_metadata, "helix_kv_bounds", None) + if helix_kv_bounds is not None and getattr( + attn_metadata, "_helix_spec_tokens_valid", False + ): + # Speculative verify groups: KV ownership is per-TOKEN. A rank + # owning only the tail page of a group has zero visible KV for + # the group's leading tokens while its per-sequence kv_len is + # nonzero, so the per-sequence mask above misses those rows. + # Their decode rows are fully masked with a finite sentinel, + # making partial_o an average over (possibly uninitialized) + # pool values; the combine multiplies by corr = 0 and + # 0 * NaN would poison the token on every CP rank — sanitize + # by the per-token bound instead. + zero_kv_mask = helix_kv_bounds[: partial_o.shape[0]] == 0 + else: + # MLA processes only the generation token slice here, so build + # the mask from [num_contexts, num_seqs). Skip this expansion + # when the more precise per-token bounds above are valid. + zero_kv_mask = _helix_zero_kv_mask( + attn_metadata, + partial_o.shape[0], + seq_start=attn_metadata.num_contexts, + num_seqs=attn_metadata.num_generations, + ) return _helix_post_process( partial_o, softmax_stats, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 0d6658167766..62e1cf12a6ec 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1340,7 +1340,12 @@ def _(input_list, group, num_lists): ] @torch.library.register_fake("trtllm::alltoall_helix_native") - def _(partial_o, softmax_stats, workspace, cp_rank, cp_size): + def _(partial_o, + softmax_stats, + workspace, + cp_rank, + cp_size, + zero_kv_mask=None): # Returns outputs with same shapes as inputs return partial_o.new_empty(partial_o.shape), softmax_stats.new_empty( softmax_stats.shape) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index b3ea554da186..506da1d389b3 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -543,14 +543,19 @@ def get(mapping: Mapping) -> "HelixAllToAllNative": return HelixAllToAllNative._cache[mapping] - def alltoall_native(self, partial_o: torch.Tensor, - softmax_stats: torch.Tensor): + def alltoall_native(self, + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + zero_kv_mask: Optional[torch.Tensor] = None): """ Perform all-to-all data exchange. Args: partial_o: Tensor with shape [..., cp_size, kv_lora_rank], dtype half. softmax_stats: Tensor with shape [..., cp_size, 2], dtype float32. + zero_kv_mask: Optional bool mask over the entry dimension, True + where this rank owns no KV. The sender rewrites those rows to a + no-op contribution, so the caller must not sanitize them itself. Returns: Tuple of (partial_o_out, softmax_stats_out) with same shapes as inputs. @@ -561,6 +566,7 @@ def alltoall_native(self, partial_o: torch.Tensor, self.workspace_tensor, self.mapping.cp_rank, self.mapping.cp_size, + zero_kv_mask, ) return partial_o_out, softmax_stats_out diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 93a0ed55bd36..25f52de48adf 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3855,6 +3855,31 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]): num_chunked_contexts=num_chunked_ctx_requests, ) + if self.enable_spec_decode and self.mapping.has_cp_helix(): + # Helix verify groups: the per-token device buffers (write slots, + # attention bounds, rank-local kv lens) must be derived on EVERY + # spec step, overlap or not -- the append/mask kernels consume + # them whenever _helix_spec_tokens_valid is armed. Under overlap + # the host packed provisional positions from a stale base, so + # first apply the same accepted-count correction position_ids + # got above; without overlap the host values are already exact. + md = inputs.get('attn_metadata') + if (md is not None and md.kv_cache_manager is not None + and getattr(md, '_helix_spec_tokens_valid', False)): + helix_gen_tokens = (inputs['input_ids'].shape[0] - + md.num_ctx_tokens) + if not self._disable_overlap_scheduler: + # The kv_lens override in the recompute supersedes the + # generic previous_kv_lens_offsets adjustment above, + # which is not ownership-aware. + md.helix_position_offsets[:helix_gen_tokens] += ( + self.previous_pos_id_offsets_cuda[:helix_gen_tokens]) + md.recompute_helix_spec_buffers( + 0, helix_gen_tokens, + self.get_runtime_tokens_per_gen_step( + self.runtime_draft_len)) + md.on_update_kv_lens() + if self.guided_decoder is not None: self.guided_decoder.token_event.record() @@ -3914,6 +3939,20 @@ def _postprocess_inputs(self, inputs: Dict[str, Any]): restore=True, ) + if (self.mapping.has_cp_helix() + and getattr(inputs['attn_metadata'], + '_helix_spec_tokens_valid', False)): + # Mirror of the helix position correction in + # _preprocess_inputs (capture symmetry, like position_ids + # above). The recompute's OVERWRITES (slots/bounds/ + # kv_lens) need no reversal: every consumer buffer is + # rewritten from host state at the next step's prepare. + inputs[ + 'attn_metadata'].helix_position_offsets[:previous_batch_tokens] -= ( + self. + previous_pos_id_offsets_cuda[:previous_batch_tokens] + ) + def _get_all_rank_num_tokens_and_spec_counts( self, attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata ) -> Tuple[Optional[List[int]], Optional[List[List[int]]]]: @@ -4905,6 +4944,38 @@ def append_cross_attention_state(request: LlmRequest, generation_requests.append(request) extend_requests += extend_dummy_requests + # Helix bookkeeping is needed by BOTH the extend (speculative verify + # group) and the plain generation packing loops below, so initialize + # it ahead of them. Positions are global; KV ownership follows the + # round-robin ledger (page b -> rank b % cp), mirrored host-side here + # (KVCacheManagerV2._helix_local_len) for provisional packing values. + helix_is_inactive_rank, helix_position_offsets = [], [] + helix_owned_new_tokens = [] + _has_cp_helix = self.mapping.has_cp_helix() + if _has_cp_helix and kv_cache_manager is not None: + _helix_phys = kv_cache_manager.tokens_per_block + _helix_ledger = _helix_phys * self.mapping.cp_size + _helix_rank_off = self.mapping.cp_rank * _helix_phys + + def _helix_local_len_host(global_len: int) -> int: + full, rem = divmod(global_len, _helix_ledger) + return full * _helix_phys + min(max(rem - _helix_rank_off, 0), + _helix_phys) + + def _helix_pack_extend(request, group: int) -> int: + # A helix gen worker's token list is the rank-LOCAL + # round-robin subset, so max_beam_num_tokens is not a global + # base; rebuild it from the global prompt length plus the + # rank-invariant generated count. Also repacks position_ids, + # which the caller filled from the local base. + generated_len = (request.max_beam_num_tokens - + request.py_prompt_len) + base = request.total_input_len_cp + generated_len - 1 + helix_position_offsets.extend(range(base, base + group)) + position_ids[-group:] = range(base, base + group) + helix_is_inactive_rank.append(False) + return base + spec_config = self.spec_config if self.enable_spec_decode else None if not self._disable_overlap_scheduler and spec_config is not None: assert spec_config.spec_dec_mode.support_overlap_scheduler( @@ -4974,6 +5045,23 @@ def append_cross_attention_state(request: LlmRequest, num_cached_tokens_per_seq.append( past_seen_token_num - request.py_num_compressed_tokens) request.cached_tokens = past_seen_token_num + if _has_cp_helix: + # Verify group [base, base+group) in GLOBAL positions. + # On a helix gen worker the request's token list is the + # rank-LOCAL round-robin subset, so max_beam_num_tokens + # (= local_prompt + generated) must NOT be used as a + # global base; reconstruct it from the global prompt + # length plus the (rank-invariant) generated count. This + # branch has no in-flight predecessor, so every value is + # exact (no device correction needed). + group = 1 + num_draft_tokens + base = _helix_pack_extend(request, group) + local_cached = _helix_local_len_host(base) + helix_owned_new_tokens.append( + _helix_local_len_host(base + group) - local_cached) + num_cached_tokens_per_seq[-1] = ( + local_cached - request.py_num_compressed_tokens) + request.cached_tokens = local_cached # update batch index request.py_batch_idx = request.py_seq_slot else: @@ -5069,9 +5157,7 @@ def append_cross_attention_state(request: LlmRequest, # update batch index request.py_batch_idx = request.py_seq_slot - helix_is_inactive_rank, helix_position_offsets = [], [] # Cache invariant method result to avoid repeated calls per-request - _has_cp_helix = self.mapping.has_cp_helix() _n_gen = len(generation_requests) # One-shot batch-level flag — True iff any generation request actually # carries multimodal payload. Lets the strip_mm_data branch below @@ -5621,6 +5707,8 @@ def previous_seq_slots_device(): attn_metadata.update_helix_param( helix_position_offsets=helix_position_offsets, helix_is_inactive_rank=helix_is_inactive_rank, + helix_owned_new_tokens=(helix_owned_new_tokens + if helix_owned_new_tokens else None), ) if not attn_metadata.is_cuda_graph: diff --git a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py index 8b9b33b3afdf..d9aac6362cfd 100644 --- a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py +++ b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,6 +14,8 @@ # limitations under the License. import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch import pytest import torch @@ -481,6 +483,41 @@ class TestHelixZeroKvMask(unittest.TestCase): speculative decoding), where num_tokens != num_seqs. """ + def test_mla_skips_generic_mask_when_per_token_bounds_are_valid(self): + from tensorrt_llm._torch.attention.mla import MLA + + num_tokens = 3 + num_heads = 2 + kv_lora_rank = 4 + helix_kv_bounds = torch.tensor([0, 5, 0], dtype=torch.int32) + attn_metadata = SimpleNamespace( + helix_kv_bounds=helix_kv_bounds, + _helix_spec_tokens_valid=True, + num_contexts=0, + num_generations=1, + ) + attn_backend = Mock() + attn_backend.forward.return_value = torch.empty(num_tokens, num_heads * kv_lora_rank) + mla = SimpleNamespace( + mapping=SimpleNamespace(has_cp_helix=lambda: True), + num_heads_tp=num_heads, + num_heads_tp_cp=num_heads, + kv_lora_rank=kv_lora_rank, + aux_stream=None, + ln_events=None, + ) + q = torch.empty(num_tokens, 1) + + with ( + patch("tensorrt_llm._torch.attention.mla._helix_zero_kv_mask") as generic_mask, + patch("tensorrt_llm._torch.attention.mla._helix_post_process") as post_process, + ): + MLA._attn_forward_gen(mla, attn_backend, q, q, q, None, attn_metadata) + + generic_mask.assert_not_called() + actual_mask = post_process.call_args.kwargs["zero_kv_mask"] + torch.testing.assert_close(actual_mask, helix_kv_bounds == 0) + def test_single_token_per_seq(self): # Plain decode: one token per sequence, so per-seq == per-token. # Seq 1 owns zero KV blocks on this rank. diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py index 67bd00ef63fd..43529722b816 100644 --- a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py +++ b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -568,3 +568,12 @@ def test_mla_helix_distributed( comms_medium: str, ): run_helix_test(_full_test_multi_gpu, scenario, comms_medium) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") +@skip_pre_blackwell +def test_mla_helix_fifo_v2_unaligned_softmax_stats(): + # Six heads over CP2 produce three local heads and a 24-byte softmax-stats + # row, exercising the float2 fallback used by Kimi K3 over CP32. + scenario = Scenario(num_heads=6, num_kv_heads=6, batch=1, ctx_len=64) + run_helix_test(_full_test_multi_gpu, scenario, "fifo_v2") From 3ff65548311b089184d61dbcc3dcf9c80f968444 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:15:02 +0000 Subject: [PATCH 03/33] [None][feat] extend Helix verify groups to the CuTe DSL MLA decode backend The per-token bounds so far only reached the FP4 MLA backend; the CuTe DSL MLA decode path rejected multi-token Helix outright. Wire kv_bounds through to it so bf16/fp16 KV is supported too. Under a verify group a token's visible KV is rank-local and per-token, so the implicit causal bound K - (S_q - 1) + q_tok is wrong: a rank owning none of a group's leading tokens must see fewer entries than the causal formula gives. The fp16 kernel now takes an optional per-token kv_bounds and uses it in place of that bound in both masked-phase branches, widens the masked span by one (the minimum is K - S_q, one position deeper), clamps the fold_sq padding rows so their discarded results still read in range, and emits the (-inf, 0) softmax identity for tokens with no local KV so the cross-rank combine stays exact. fp8 is deliberately excluded. Its kernel has no per-token bounds, so it takes kv_bounds for signature parity and drops it, the custom op raises on a non-None value, and the wrapper gate rejects fp8 KV under verify groups rather than silently mis-masking. FallbackFmha also rejects verify groups: the fused thop path's spec-dec mask and per-sequence helix_is_inactive_rank gate both assume the new KV entries are the trailing slots of one rank's kv_len. Being last in the library list, rejecting makes dispatch raise instead of running wrong. The SM107 arch widening that sits beside this work on the source branch is intentionally left out -- it is a separate concern and belongs to the kernels PR. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../attention/backends/fmha/cute_dsl_mla.py | 20 ++++- .../attention/backends/fmha/fallback.py | 13 ++++ .../_torch/custom_ops/cute_dsl_custom_ops.py | 46 ++++++++++- .../attention/mla/mla_decode_fp16.py | 77 ++++++++++++++++++- .../blackwell/attention/mla/mla_decode_fp8.py | 8 ++ 5 files changed, 156 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 51e730ed754f..0fbc4097051d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -301,8 +301,15 @@ def _is_supported_with_reason( seq_len_q = q.shape[0] // meta.num_generations batch_size = meta.num_generations if meta.helix_position_offsets is not None: - if seq_len_q != 1: + if seq_len_q != 1 and not meta._helix_spec_tokens_valid: + # Multi-token decode under helix needs the per-token bound / + # write-slot buffers of the speculative verify-group path. return False, "CuTe DSL MLA FMHA only supports single-token decode with Helix." + if seq_len_q != 1 and self._get_kernel_dtype(attn, q) == torch.float8_e4m3fn: + return False, ( + "CuTe DSL MLA FMHA helix verify groups require a bf16/fp16 " + "KV cache (the fp8 kernel has no per-token bounds)." + ) softmax_stats = fwd.softmax_stats_tensor if softmax_stats is None: return False, "CuTe DSL MLA FMHA requires softmax_stats_tensor with Helix." @@ -505,6 +512,17 @@ def _run_mla_decode( # Max batch size for the AutoTuner to profile. int(meta.max_num_requests), params.fwd.softmax_stats_tensor, + # Per-token rank-local bounds, filled by + # recompute_helix_spec_buffers. None everywhere else. + ( + meta.helix_kv_bounds[:num_tokens] + if ( + meta.helix_position_offsets is not None + and meta._helix_spec_tokens_valid + and kernel_dtype != torch.float8_e4m3fn + ) + else None + ), ) def run_mla_generation( diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index b3d734b97817..fa79b6da6653 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -81,6 +81,19 @@ def _is_supported( phase: Optional[FmhaPhase] = None, ) -> bool: del k, v, phase + # A verify group may straddle a page boundary onto two CP ranks, so + # its KV ownership is per-token. The fused thop path cannot express + # that: its spec-dec mask and the per-sequence helix_is_inactive_rank + # gate both assume the new KV entries are the trailing slots of one + # rank's kv_len. Reject rather than run it silently wrong; being last + # in the library list, this makes dispatch raise. + if ( + metadata.helix_position_offsets is not None + and getattr(metadata, "_helix_spec_tokens_valid", False) + and metadata.num_generations > 0 + and q.shape[0] > metadata.num_seqs + ): + return False if q is not None and q.dtype == torch.float8_e4m3fn: return False if forward_args.attention_mask == CustomAttentionMask.CUSTOM: diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 3b7f9b9ed0ec..97709df4da5b 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -11504,7 +11504,10 @@ def forward( tensor of shape (H, S_q, B) remains in the workspace. """ (q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o, - workspace, softmax_stats) = inputs + workspace, softmax_stats) = inputs[:9] + # inputs[9] (optional): helix per-token attention bounds of shape + # (B * S_q,), int32 — speculative verify groups only. + kv_bounds = inputs[9] if len(inputs) > 9 else None softmax_scale = float(kwargs.get("softmax_scale", 1.0)) output_scale = float(kwargs.get("output_scale", 1.0)) @@ -11577,12 +11580,33 @@ def forward( split_workspace = workspace_bytes[split_kv_offset:split_kv_offset + split_kv_size] + if kv_bounds is not None and AutoTuner.get().is_tuning_mode: + # Profiling rebuilds cache_seqs at bucketed sizes but input 9 + # has no dynamic-dim spec, so kv_bounds arrives at the old + # size. Bound values only affect masking depth, not the + # tactic space, so any size-consistent dummy will do. + if kv_bounds.numel() != batch_size * seq_len_q: + kv_bounds = cache_seqs.repeat_interleave( + seq_len_q).contiguous() + if kv_bounds is not None: + expected_bounds_shape = (batch_size * seq_len_q, ) + if (kv_bounds.shape != expected_bounds_shape + or kv_bounds.dtype != torch.int32 + or kv_bounds.device != o.device + or not kv_bounds.is_contiguous()): + raise RuntimeError( + "CuteDSLNVMlaDecodeBlackwellRunner requires contiguous " + "int32 kv_bounds on the output device with shape " + f"{expected_bounds_shape}, got shape=" + f"{tuple(kv_bounds.shape)}, dtype={kv_bounds.dtype}.") + cache_key = self.unique_id() + ( out_dtype, mma_qk_tiler_mn, mma_pv_tiler_mn, split_kv, is_persistent, + kv_bounds is not None, ) if cache_key not in CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache: # A compile outside the tuning window stalls the serving loop @@ -11650,6 +11674,9 @@ def forward( if use_workspace else None) cache_seqs_ct = cute.runtime.from_dlpack( cache_seqs, assumed_align=16).mark_layout_dynamic() + kv_bounds_ct = (cute.runtime.from_dlpack( + kv_bounds, assumed_align=4).mark_layout_dynamic() + if kv_bounds is not None else None) # Variable split-KV (block_split_kvs) is not used on this path: block_split_kvs_ct = None @@ -11670,6 +11697,7 @@ def forward( workspace_ct, split_kv, cache_seqs_ct, + kv_bounds_ct, block_split_kvs_ct, cutlass.Float32(softmax_scale), cutlass.Float32(output_scale), @@ -11708,6 +11736,7 @@ def forward( (split_kv > 1 and split_workspace.numel() > 0) else None, split_kv, cache_seqs, + kv_bounds, None, # block_split_kvs: var-split path unused (is_var_split_kv False) softmax_scale, output_scale, @@ -11735,14 +11764,19 @@ def cute_dsl_mla_decode_fp8_blackwell( page_size: int, softmax_scale: float, output_scale: float, - # Keep the last two arguments required in the custom-op schema. PyTorch + # Keep the trailing arguments required in the custom-op schema. PyTorch # elides trailing default-valued arguments before its mutation fallback, # while mutates_args retains their positional indices. max_batch_size: int, softmax_stats: Optional[torch.Tensor], + kv_bounds: Optional[torch.Tensor], ) -> None: """CuTe DSL FP8 MLA decode (Blackwell SM100/SM103). """ + if kv_bounds is not None: + raise ValueError( + "trtllm::cute_dsl_mla_decode_fp8_blackwell does not support " + "helix per-token kv_bounds (bf16/fp16 kernel only).") if (sm_version := get_sm_version()) not in (100, 103): raise ValueError( f"trtllm::cute_dsl_mla_decode_fp8_blackwell requires SM 100 or " @@ -11798,6 +11832,7 @@ def _( output_scale: float, max_batch_size: int, softmax_stats: Optional[torch.Tensor], + kv_bounds: Optional[torch.Tensor], ) -> None: return None @@ -11823,8 +11858,12 @@ def cute_dsl_mla_decode_fp16_blackwell( # See the FP8 op above: these must remain required schema arguments. max_batch_size: int, softmax_stats: Optional[torch.Tensor], + kv_bounds: Optional[torch.Tensor], ) -> None: """CuTe DSL FP16/BF16 MLA decode (Blackwell SM100/SM103). + + kv_bounds: helix speculative verify groups — per-token rank-local + attention bounds of shape (B * seq_len_q,), int32. """ if (sm_version := get_sm_version()) not in (100, 103): raise ValueError( @@ -11860,7 +11899,7 @@ def cute_dsl_mla_decode_fp16_blackwell( ) inputs = [ q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o, - workspace, softmax_stats + workspace, softmax_stats, kv_bounds ] tuner = AutoTuner.get() _, best_tactic = tuner.choose_one( @@ -11898,6 +11937,7 @@ def _( output_scale: float, max_batch_size: int, softmax_stats: Optional[torch.Tensor], + kv_bounds: Optional[torch.Tensor], ) -> None: return None diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py index d53dcd7c3ce3..c56ba8322f2e 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py @@ -311,6 +311,7 @@ def __call__( workspace: cute.Tensor, split_kv: cutlass.Int32, cache_seqs: Optional[cute.Tensor], + kv_bounds: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, @@ -328,6 +329,7 @@ def __call__( workspace, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale, output_scale, @@ -348,6 +350,7 @@ def run_with_softmax_stats( workspace: cute.Tensor, split_kv: cutlass.Int32, cache_seqs: Optional[cute.Tensor], + kv_bounds: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, @@ -365,6 +368,7 @@ def run_with_softmax_stats( workspace, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale, output_scale, @@ -385,6 +389,7 @@ def _run( workspace: cute.Tensor, split_kv: cutlass.Int32, cache_seqs: Optional[cute.Tensor], + kv_bounds: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, @@ -392,6 +397,15 @@ def _run( ): """Execute the Multi-Head Latent Attention operation on the provided tensors. + kv_bounds (helix speculative verify groups): optional int32 tensor of + shape [batch_size * seq_len_q]; entry b*seq_len_q + q gives the number + of this rank's local KV entries query token q of sequence b may attend + to (committed prefix + owned in-flight group tokens up to and + including itself). When present it replaces the implicit causal bound + K - (seq_len_q - 1) + q_tok; values are guaranteed to lie in + [K - seq_len_q, K], i.e. inside the span the masked phase already + covers for the causal case. + The method handles: 1. Initialization of workspace for temporary split KV buffers 2. Validation of tensor data types @@ -792,6 +806,7 @@ class SplitKVKernelSharedStorage: acc_lse, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale_log2, output_scale, @@ -827,6 +842,7 @@ class SplitKVKernelSharedStorage: acc_lse, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale_log2, output_scale, @@ -859,6 +875,7 @@ class SplitKVKernelSharedStorage: split_kv, cache_seqs, block_split_kvs, + kv_bounds, ) else: reduction_kernel = self.reduction_kernel( @@ -938,6 +955,7 @@ def split_kv_kernel( mAccLSE: Optional[cute.Tensor], split_kv: cutlass.Int32, cache_seqs: cute.Tensor, + kv_bounds: Optional[cute.Tensor], block_split_kvs: cute.Tensor, softmax_scale_log2: cutlass.Float32, output_scale: cutlass.Float32, @@ -1356,6 +1374,7 @@ def split_kv_kernel( mAccO=mAccO, mO=mO, K=cache_seqs[blk_coord[2]], + kv_bounds=kv_bounds, L=mCL.shape[1], tmem_ptr=tmem_ptr, tidx=tidx, @@ -1458,6 +1477,7 @@ def reduction_kernel( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, block_split_kvs: cute.Tensor, + kv_bounds: Optional[cute.Tensor] = None, ): """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results from multiple split_kv blocks into final outputs. @@ -1528,7 +1548,27 @@ def reduction_kernel( if tidx == 0: mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse if cutlass.const_expr(self.emit_softmax_stats): - if cache_seqs[blk_coord[2]] > 0: + # A rank joins the CP merge only for tokens with at least + # one visible local KV entry. With verify groups that is + # per-token: a straddling group leaves this rank zero + # entries for its leading tokens while later ones have some. + if cutlass.const_expr(kv_bounds is not None): + # blk_coord runs over the folded reduction grid + # (H*F, S_q/F, B) while kv_bounds is indexed by the + # true token. A folded chunk packs its rows as + # tok_in_chunk * H + head, so the true token is + # chunk * F + row // H (self.num_heads and + # self.seq_len_q stay pre-fold). + if cutlass.const_expr(self.fold_sq): + q_tok = (blk_coord[1] * self.fold_sq_ratio + + blk_coord[0] // self.num_heads) + else: + q_tok = blk_coord[1] + has_local_kv = kv_bounds[blk_coord[2] * self.seq_len_q + + q_tok] > 0 + else: + has_local_kv = cache_seqs[blk_coord[2]] > 0 + if has_local_kv: mSoftmaxStats[blk_coord[0], blk_coord[1], blk_coord[2], 0] = global_lse / LOG2_E mSoftmaxStats[blk_coord[0], blk_coord[1], blk_coord[2], @@ -2448,8 +2488,14 @@ def compute( # positions. Min k_bound = K - (S_q-1), which can span up to # ceil((seq_len_q-2)/tile_N)+1 tiles (tile-boundary-crossing case). For # S_q=1 this reduces to 1 tile -- identical to a plain K-bound check. + # With helix per-token bounds the minimum is K - S_q (a rank owning + # none of the group's tokens), one position deeper, so widen the span + # by one. tile_n = self.mma_qk_tiler[1] - mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1 + if cutlass.const_expr(common_params.kv_bounds is not None): + mask_tile_count = (self.seq_len_q - 1 + tile_n - 1) // tile_n + 1 + else: + mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1 # first_mask_tile_idx is the global index of the first tile that may # need masking. Runtime because it depends on K (per-batch in @@ -2777,7 +2823,19 @@ def softmax( cta_m_rows) // self.num_heads) else: q_tok = common_params.blk_coord[1] - k_bound = common_params.K - (self.seq_len_q - 1) + q_tok + if cutlass.const_expr(common_params.kv_bounds is not None): + # Per-token rank-local bound; subsumes the causal + # offset and non-owner ranks. + # fold_sq M-tile padding rows derive q_tok >= S_q; + # their results are discarded but the read must stay + # in bounds. + q_tok_c = (q_tok if cute.elem_less( + q_tok, self.seq_len_q) else self.seq_len_q - 1) + k_bound = common_params.kv_bounds[ + common_params.blk_coord[2] * self.seq_len_q + + q_tok_c] + else: + k_bound = common_params.K - (self.seq_len_q - 1) + q_tok tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less( tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, k_bound, @@ -2824,7 +2882,18 @@ def softmax( cta_m_rows) // self.num_heads) else: q_tok = common_params.blk_coord[1] - k_bound = common_params.K - (self.seq_len_q - 1) + q_tok + if cutlass.const_expr(common_params.kv_bounds is not None): + # Per-token rank-local bound (see the sm_100 branch). + # fold_sq M-tile padding rows derive q_tok >= S_q; + # their results are discarded but the read must stay + # in bounds. + q_tok_c = (q_tok if cute.elem_less( + q_tok, self.seq_len_q) else self.seq_len_q - 1) + k_bound = common_params.kv_bounds[ + common_params.blk_coord[2] * self.seq_len_q + + q_tok_c] + else: + k_bound = common_params.K - (self.seq_len_q - 1) + q_tok tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less( tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, k_bound, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py index 9e2477edf634..79b8e9c202e9 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py @@ -309,6 +309,10 @@ def __call__( workspace: cute.Tensor, split_kv: cutlass.Int32, cache_seqs: Optional[cute.Tensor], + # Signature parity with the fp16 kernel (helix speculative verify + # groups); the fp8 kernel does not implement per-token bounds and the + # wrapper gate never selects fp8 KV under helix. + kv_bounds: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, @@ -346,6 +350,10 @@ def run_with_softmax_stats( workspace: cute.Tensor, split_kv: cutlass.Int32, cache_seqs: Optional[cute.Tensor], + # Signature parity with the fp16 kernel; the fp8 kernel does not + # implement per-token bounds, so the value is accepted and dropped + # (the custom op rejects a non-None kv_bounds for fp8 upstream). + kv_bounds: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, From 606cd259c9959eb120c0cd516404d61b2838a4eb Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:58:26 +0000 Subject: [PATCH 04/33] [None][feat] support Helix verify groups on the fp8 CuTe DSL MLA decode Completes the CuTe DSL side: the fp8 kernel now implements per-token kv_bounds itself instead of taking the parameter and dropping it. Both masked-phase branches read the per-token rank-local bound in place of the implicit causal bound, and the CP-merge participation check goes per-token, so a rank owning none of a verify group's leading tokens contributes the softmax identity for them rather than a causal-bound slice of another rank's KV. With the kernel able to honour the bounds, the three guards that kept fp8 out are gone: the custom op no longer raises on a non-None kv_bounds, the wrapper gate no longer rejects an fp8 KV cache under verify groups, and the bounds are passed whatever the kernel dtype. The fp8 tensor-valued softmax/output scales that sit beside this work on the source branch are deliberately not taken: they are an unrelated ABI change. main's cutlass.Float32 scalars and the softmax_scale_log2 precompute are kept, and the kv_bounds masking is independent of both. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../attention/backends/fmha/cute_dsl_mla.py | 11 +-- .../_torch/custom_ops/cute_dsl_custom_ops.py | 4 -- .../blackwell/attention/mla/mla_decode_fp8.py | 71 ++++++++++++++++--- 3 files changed, 62 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 0fbc4097051d..762554aa8ca8 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -305,11 +305,6 @@ def _is_supported_with_reason( # Multi-token decode under helix needs the per-token bound / # write-slot buffers of the speculative verify-group path. return False, "CuTe DSL MLA FMHA only supports single-token decode with Helix." - if seq_len_q != 1 and self._get_kernel_dtype(attn, q) == torch.float8_e4m3fn: - return False, ( - "CuTe DSL MLA FMHA helix verify groups require a bf16/fp16 " - "KV cache (the fp8 kernel has no per-token bounds)." - ) softmax_stats = fwd.softmax_stats_tensor if softmax_stats is None: return False, "CuTe DSL MLA FMHA requires softmax_stats_tensor with Helix." @@ -516,11 +511,7 @@ def _run_mla_decode( # recompute_helix_spec_buffers. None everywhere else. ( meta.helix_kv_bounds[:num_tokens] - if ( - meta.helix_position_offsets is not None - and meta._helix_spec_tokens_valid - and kernel_dtype != torch.float8_e4m3fn - ) + if (meta.helix_position_offsets is not None and meta._helix_spec_tokens_valid) else None ), ) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 97709df4da5b..0fba18388303 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -11773,10 +11773,6 @@ def cute_dsl_mla_decode_fp8_blackwell( ) -> None: """CuTe DSL FP8 MLA decode (Blackwell SM100/SM103). """ - if kv_bounds is not None: - raise ValueError( - "trtllm::cute_dsl_mla_decode_fp8_blackwell does not support " - "helix per-token kv_bounds (bf16/fp16 kernel only).") if (sm_version := get_sm_version()) not in (100, 103): raise ValueError( f"trtllm::cute_dsl_mla_decode_fp8_blackwell requires SM 100 or " diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py index 79b8e9c202e9..75d5c623d8e7 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py @@ -309,9 +309,6 @@ def __call__( workspace: cute.Tensor, split_kv: cutlass.Int32, cache_seqs: Optional[cute.Tensor], - # Signature parity with the fp16 kernel (helix speculative verify - # groups); the fp8 kernel does not implement per-token bounds and the - # wrapper gate never selects fp8 KV under helix. kv_bounds: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], softmax_scale: cutlass.Float32, @@ -330,6 +327,7 @@ def __call__( workspace, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale, output_scale, @@ -350,9 +348,6 @@ def run_with_softmax_stats( workspace: cute.Tensor, split_kv: cutlass.Int32, cache_seqs: Optional[cute.Tensor], - # Signature parity with the fp16 kernel; the fp8 kernel does not - # implement per-token bounds, so the value is accepted and dropped - # (the custom op rejects a non-None kv_bounds for fp8 upstream). kv_bounds: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], softmax_scale: cutlass.Float32, @@ -371,6 +366,7 @@ def run_with_softmax_stats( workspace, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale, output_scale, @@ -391,6 +387,7 @@ def _run( workspace: cute.Tensor, split_kv: cutlass.Int32, cache_seqs: Optional[cute.Tensor], + kv_bounds: Optional[cute.Tensor], block_split_kvs: Optional[cute.Tensor], softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, @@ -398,6 +395,12 @@ def _run( ): """Execute the Multi-Head Latent Attention operation on the provided tensors. + kv_bounds (helix speculative verify groups): optional int32 tensor of + shape [batch_size * seq_len_q]; entry b*seq_len_q + q gives the number + of this rank's local KV entries query token q of sequence b may attend + to. When present it replaces the implicit causal bound + K - (seq_len_q - 1) + q_tok. + The method handles: 1. Initialization of workspace for temporary split KV buffers 2. Validation of tensor data types @@ -852,6 +855,7 @@ class SplitKVKernelSharedStorage: acc_lse, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale_log2, output_scale, @@ -889,6 +893,7 @@ class SplitKVKernelSharedStorage: acc_lse, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale_log2, output_scale, @@ -923,6 +928,7 @@ class SplitKVKernelSharedStorage: split_kv, cache_seqs, block_split_kvs, + kv_bounds, ) else: reduction_kernel = self.reduction_kernel( @@ -1002,6 +1008,7 @@ def split_kv_kernel( mAccLSE: Optional[cute.Tensor], split_kv: cutlass.Int32, cache_seqs: cute.Tensor, + kv_bounds: Optional[cute.Tensor], block_split_kvs: cute.Tensor, softmax_scale_log2: cutlass.Float32, output_scale: cutlass.Float32, @@ -1422,6 +1429,7 @@ def split_kv_kernel( mAccO=mAccO, mO=mO, K=cache_seqs[blk_coord[2]], + kv_bounds=kv_bounds, L=mCL.shape[1], tmem_ptr=tmem_ptr, tidx=tidx, @@ -1524,6 +1532,7 @@ def reduction_kernel( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, block_split_kvs: cute.Tensor, + kv_bounds: Optional[cute.Tensor] = None, ): """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results from multiple split_kv blocks into final outputs. @@ -1594,7 +1603,21 @@ def reduction_kernel( if tidx == 0: mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse if cutlass.const_expr(self.emit_softmax_stats): - if cache_seqs[blk_coord[2]] > 0: + # A rank joins the CP merge only for tokens with at least + # one visible local KV entry; with verify groups that is + # per-token. blk_coord runs over the folded grid + # (H*F, S_q/F, B): true token = chunk * F + row // H. + if cutlass.const_expr(kv_bounds is not None): + if cutlass.const_expr(self.fold_sq): + q_tok = (blk_coord[1] * self.fold_sq_ratio + + blk_coord[0] // self.num_heads) + else: + q_tok = blk_coord[1] + has_local_kv = kv_bounds[blk_coord[2] * self.seq_len_q + + q_tok] > 0 + else: + has_local_kv = cache_seqs[blk_coord[2]] > 0 + if has_local_kv: mSoftmaxStats[blk_coord[0], blk_coord[1], blk_coord[2], 0] = global_lse / LOG2_E mSoftmaxStats[blk_coord[0], blk_coord[1], blk_coord[2], @@ -2432,7 +2455,13 @@ def compute( # ceil((seq_len_q-2)/tile_N)+1 tiles (tile-boundary-crossing case). For # S_q=1 this reduces to 1 tile -- identical to a plain K-bound check. tile_n = self.mma_qk_tiler[1] - mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1 + # With helix per-token bounds the minimum bound is K - S_q (a rank + # owning none of the group's tokens), one position deeper than the + # causal minimum, so widen the masked span by one. + if cutlass.const_expr(common_params.kv_bounds is not None): + mask_tile_count = (self.seq_len_q - 1 + tile_n - 1) // tile_n + 1 + else: + mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1 # first_mask_tile_idx is the global index of the first tile that may # need masking. Runtime because it depends on K (per-batch in @@ -2759,7 +2788,18 @@ def softmax( cta_m_rows) // self.num_heads) else: q_tok = common_params.blk_coord[1] - k_bound = common_params.K - (self.seq_len_q - 1) + q_tok + if cutlass.const_expr(common_params.kv_bounds is not None): + # Per-token rank-local bound; subsumes the causal + # offset and non-owner ranks. fold_sq M-tile padding + # rows derive q_tok >= S_q; their results are + # discarded but the read must stay in bounds. + q_tok_c = (q_tok if cute.elem_less( + q_tok, self.seq_len_q) else self.seq_len_q - 1) + k_bound = common_params.kv_bounds[ + common_params.blk_coord[2] * self.seq_len_q + + q_tok_c] + else: + k_bound = common_params.K - (self.seq_len_q - 1) + q_tok tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less( tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, k_bound, @@ -2805,7 +2845,18 @@ def softmax( cta_m_rows) // self.num_heads) else: q_tok = common_params.blk_coord[1] - k_bound = common_params.K - (self.seq_len_q - 1) + q_tok + if cutlass.const_expr(common_params.kv_bounds is not None): + # Per-token rank-local bound; subsumes the causal + # offset and non-owner ranks. fold_sq M-tile padding + # rows derive q_tok >= S_q; their results are + # discarded but the read must stay in bounds. + q_tok_c = (q_tok if cute.elem_less( + q_tok, self.seq_len_q) else self.seq_len_q - 1) + k_bound = common_params.kv_bounds[ + common_params.blk_coord[2] * self.seq_len_q + + q_tok_c] + else: + k_bound = common_params.K - (self.seq_len_q - 1) + q_tok tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less( tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, k_bound, From 0b838381602e025c2bd0fe3ca6e388f9ad51483e Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:58:48 +0000 Subject: [PATCH 05/33] [None][feat] enable DSpark speculative decoding under Helix CP for Kimi K3 The per-token Helix primitive is in place, but nothing selected it: the K3 runtime rejected speculative decoding under Helix CP, and the draft side sized and placed its KV as if CP did not exist. Admit standalone DSpark linear chains through the K3 helix speculative allowlist and raise loudly on anything else, so an unsupported combination fails at config time rather than silently mis-attending. The draft model runs on the CP-free repurposed mapping, since the helix round-robin ledger governs only the target KV. The KV-cache cost model follows: draft costs are computed on that mapping and their slopes scaled by cp_size, because a draft token is priced against a rank-local target token and the target stores only every cp_size-th page per rank. Intercepts are per-request rank-local bytes and stay unscaled. Also bootstraps per-request context slots on disaggregated generation workers, which a standalone drafter otherwise never initializes. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 32 ++++++-- tensorrt_llm/_torch/pyexecutor/_util.py | 78 +++++++++++++++---- .../_torch/pyexecutor/config_utils.py | 18 ++++- .../kv_cache/kv_cache_manager_v2.py | 5 +- .../kv_cache/mamba_cache_manager.py | 17 +--- 5 files changed, 113 insertions(+), 37 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index b87a6f1636dc..563a086b4685 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -2149,11 +2149,33 @@ def _setup_helix_mappings( "per-request locality of KDA recurrent state." ) if spec_config is not None: - raise ValueError( - "Kimi K3 helix phase 1 does not support speculative " - "decoding (round-robin KV bookkeeping assumes one token " - "per decode step)." - ) + # Helix supports only the standalone DSpark drafter (verified on + # the V2 superblock ledger); reject everything else loudly rather + # than let an unsupported spec mode run silently wrong. + if not spec_config.spec_dec_mode.is_dspark(): + raise ValueError( + "Kimi K3 helix supports speculative decoding only with " + f"DSpark (standalone drafter); got " + f"{spec_config.decoding_type!r}." + ) + # The SpeculationGate acceptance-rate trip permanently disables + # speculation mid-flight while enable_spec_decode stays True; + # in-flight helix requests then fall into the plain generation + # loop whose position math (total_input_len_cp + + # py_decoding_iter - 1) is stale once any draft token was + # accepted -> silently wrong RoPE positions and KV slots. Reject + # the trip wires until that loop is helix-group aware. + if ( + spec_config.acceptance_rate_window_size is not None + or spec_config.acceptance_rate_threshold is not None + ): + raise ValueError( + "Kimi K3 helix does not support the speculation " + "acceptance-rate gate (acceptance_rate_window_size / " + "acceptance_rate_threshold): dynamically disabling " + "speculation mid-flight leaves helix requests on a " + "single-token position formula." + ) cp = model_config.mapping.cp_size repurposed_tp = model_config.mapping.tp_size * cp if cfg.num_attention_heads % repurposed_tp != 0: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 820a4bcac0d7..74b8e04c5c54 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -230,6 +230,20 @@ def get_kv_cache_manager_cls( # the shared hybrid transceiver validation below: the Python NIXL # transceiver selects the Mixed manager, whose KDA recurrent/conv # states transfer through the bounce buffer. + # Helix x speculation bookkeeping (per-token verify groups on the + # superblock ledger, py_helix_decode_group_index advancement) exists + # only in KVCacheManagerV2. The V1-family hybrid managers account + # helix decode one token per iteration and have no helix-x-spec + # path, so a default (V1) resolution would run silently wrong. + if (model_config.mapping is not None + and model_config.mapping.has_cp_helix() + and model_config.spec_config is not None and not use_v2): + raise ValueError( + "Kimi K3 helix with speculative decoding requires " + "kv_cache_config.use_kv_cache_manager_v2=True; the V1-family " + "hybrid managers do not implement per-token verify-group " + "bookkeeping.") + if is_kimi_linear(config) and not use_v2 and not is_disagg: if kv_cache_config.enable_block_reuse: logger.info( @@ -897,6 +911,7 @@ def _per_manager_cache_cost(self, kv_cache_config: Optional[KvCacheConfig] = None, *, is_draft: bool = False, + mapping=None, **extra_kwargs) -> CacheCost: kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) @@ -913,7 +928,7 @@ def _per_manager_cache_cost(self, return CacheCost.from_raw( manager_cls.get_cache_size_per_token( model_config, - self._mapping, + mapping if mapping is not None else self._mapping, tokens_per_block=self._tokens_per_block, max_seq_len=self._max_seq_len, max_batch_size=self._max_batch_size, @@ -966,14 +981,33 @@ def _get_draft_cache_cost( *, use_separate_draft_kv_cache: bool, ) -> Optional[CacheCost]: - """Return the draft manager's standalone cache cost, if it has one.""" + """Return the draft manager's standalone cache cost, if it has one. + + Under helix CP the drafter is dense rather than helix-sharded, so it is + costed with the same repurposed mapping runtime construction uses, then + the slope is multiplied by cp_size to express it per rank-LOCAL target + token (the target stores only every cp_size-th page per rank). + Intercepts are per-request rank-local bytes and stay unscaled. + """ + draft_mapping = self._mapping + helix_cp_scale = 1 + if self._mapping.has_cp_helix(): + draft_mapping = self._mapping.repurpose_helix_cp_to_tp() + helix_cp_scale = self._mapping.cp_size + + def scaled(cost: CacheCost) -> CacheCost: + return CacheCost(slope=cost.slope * helix_cp_scale, + intercept=cost.intercept) + if self._draft_model_engine is not None: draft_model_config = self._draft_model_engine.model.model_config draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( self._draft_model_engine, kv_cache_config) - return self._per_manager_cache_cost(draft_kv_cache_manager_cls, - draft_model_config, - kv_cache_config) + return scaled( + self._per_manager_cache_cost(draft_kv_cache_manager_cls, + draft_model_config, + kv_cache_config, + mapping=draft_mapping)) if use_separate_draft_kv_cache: # One-model draft with separate KV cache layout. # Pass num_layers explicitly since the HF config may report a @@ -999,18 +1033,22 @@ def _get_draft_cache_cost( draft_kv_cache_config) if self._speculative_config.spec_dec_mode.is_external_drafter(): # External drafter: layers start from 0, normal PP distribution - return self._per_manager_cache_cost(draft_kv_cache_manager_cls, - effective_draft_config, - draft_kv_cache_config, - is_draft=True) + return scaled( + self._per_manager_cache_cost(draft_kv_cache_manager_cls, + effective_draft_config, + draft_kv_cache_config, + mapping=draft_mapping, + is_draft=True)) elif self._mapping.is_last_pp_rank(): # EAGLE3/MTP: draft layers only on last PP rank - return self._per_manager_cache_cost( - draft_kv_cache_manager_cls, - effective_draft_config, - draft_kv_cache_config, - num_layers=self._get_num_draft_layers(), - is_draft=True) + return scaled( + self._per_manager_cache_cost( + draft_kv_cache_manager_cls, + effective_draft_config, + draft_kv_cache_config, + mapping=draft_mapping, + num_layers=self._get_num_draft_layers(), + is_draft=True)) return None def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, @@ -1959,12 +1997,20 @@ def _create_one_model_draft_kv_cache_manager( # the sparse_attention_config. Get it from effective_draft_config which # falls back to the target model's config for MTP mode. sparse_attn_config = effective_draft_config.sparse_attention_config + # Under helix the standalone drafter is built against the repurposed + # mapping (CP ranks become TP ranks) and every rank keeps its full + # drafter KV, so its paged manager needs the CP-free mapping: the + # round-robin ledger applies to the TARGET KV alone, and + # KVCacheManagerV2 rejects is_draft x helix outright. + draft_mapping = self._mapping + if draft_mapping.has_cp_helix(): + draft_mapping = draft_mapping.repurpose_helix_cp_to_tp() return _create_kv_cache_manager( model_engine=None, max_cuda_graph_batch_size=self._model_engine. _max_cuda_graph_batch_size, kv_cache_manager_cls=draft_kv_cache_manager_cls, - mapping=self._mapping, + mapping=draft_mapping, kv_cache_config=draft_kv_config, tokens_per_block=self._tokens_per_block, max_seq_len=max_seq_len, diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 33f6c6649d0f..2e10388fdd6d 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -511,6 +511,22 @@ def extract_qwen4_exp_ple_cache_params( ) +def mamba_effective_tp_size(mapping) -> int: + """TP degree for sizing per-rank mamba/KDA state (budgeting AND allocation). + + Attention-DP replicates the state and takes precedence; helix repurposes + CP ranks as plain TP for recurrent-state layers. Must match the runtime + pool construction (mamba_cache_manager) or the budget split withholds + unsharded-state bytes the allocator never uses (observed: 27.2 GiB/rank + mis-withheld on a helix16 gen worker whose real pool is 1/16-sharded). + """ + if mapping.enable_attention_dp: + return 1 + if mapping.has_cp_helix(): + return mapping.tp_size * mapping.cp_size + return mapping.tp_size + + @dataclasses.dataclass class MambaKVCacheParams: """Normalized mamba-related inputs for kv_cache_manager_cls. @@ -570,7 +586,7 @@ def get_layer_masks( def get_states_bytes_per_layer(self, mapping) -> int: """Return the total bytes of Mamba state per layer, used for budgeting.""" - tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 + tp_size = mamba_effective_tp_size(mapping) d_inner = self.head_dim * self.num_heads conv_dim = (d_inner + 2 * self.n_groups * self.state_size) // tp_size nheads = self.num_heads // tp_size diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 080a064ce7e0..ce6db25a8199 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3264,8 +3264,9 @@ def _set_helix_rank_fields(self, req: LlmRequest) -> None: from ``py_decoding_iter``: the sampler advances that counter after scheduling under the overlap loop, so a schedule-time read is one step behind and would repeat the first decode position, overwriting - the first generated token's KV. Assumes one new token per step - (draft-token modes are rejected under helix). + the first generated token's KV. Multi-token verify groups advance + py_helix_decode_group_index per committed group, so this formula + stays exact under DSpark speculation. """ step = req.py_helix_decode_group_index + 1 pos = req.total_input_len_cp + step - 1 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index f2e6cb94c79f..626093a5cf64 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -59,6 +59,10 @@ SsmLayerConfig, TokenIdExt, _KVCache) +# Shared with the KV budget estimator so allocator and budgeting can never +# diverge on the sharding rule (config_utils is import-cycle-free). +from .config_utils import mamba_effective_tp_size as _mamba_effective_tp_size + GB = 1 << 30 @@ -168,19 +172,6 @@ class MambaRole: PLE_CONV_STATE = DataRole("ple_conv_state") -def _mamba_effective_tp_size(mapping: Mapping) -> int: - """TP degree for sizing per-rank mamba/KDA state pools. - - Attention-DP replicates the state and takes precedence; helix - repurposes CP ranks as plain TP for recurrent-state layers. - """ - if mapping.enable_attention_dp: - return 1 - if mapping.has_cp_helix(): - return mapping.tp_size * mapping.cp_size - return mapping.tp_size - - def get_tensor_size_bytes(tensor): """Calculate tensor size in bytes.""" if isinstance(tensor, torch.Tensor): From 8a8d77e56bd31b209aabc78b4fe3611707aab766 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 18 Sep 2026 03:09:28 +0000 Subject: [PATCH 06/33] [None][fix] let CuTeDSL MLA serve 96 heads once multi-token decode is allowed This PR relaxes the single-token restriction so a helix speculative verify group can run on the CuTeDSL MLA kernel, but the request still has to clear _is_perf_favorable further down the same function, and _PERF_MIN_BATCH_FP8 lists 96 heads only at seq_len_q == 1. A verify group asks for 1 + draft_len, misses the table, and is refused as "not a perf win" -- so for a 96-head model the path this PR adds is unreachable. Falling through is not harmless either. The table already notes that TRTLLM-Gen must not serve 96 heads ("its heuristic may select a 64-head Q tile, which does not divide 96"), and TRTLLM-Gen rejects 64 < num_heads_q < 128 outright, so the next library that accepts the request fails at engine start with "trtllm-gen MLA decode does not support 64 < num_heads_q < 128". Admit every seq_len_q at 96 heads rather than enumerating draft lengths: the constraint is on the head count, not on the query length. Kimi K3 is the 96-head model this matters for; a 4-node K3 DSpark disaggregated accuracy run reproduces the engine-start failure without this and scores GSM8K exact_match 0.9651 with it. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/attention/backends/fmha/cute_dsl_mla.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 762554aa8ca8..2764692e7aed 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -254,6 +254,13 @@ def _is_perf_favorable( (128, 2): 32, } min_batch = _PERF_MIN_BATCH_FP8.get((num_heads, seq_len_q)) + if min_batch is None and num_heads == 96: + # For H=96 this table is a correctness constraint, not a perf + # tradeoff: TRTLLM-Gen rejects 64 < num_heads_q < 128 outright, so + # falling through does not reach a faster kernel, it reaches an + # executor-init failure. Admit every seq_len_q, which is what + # speculative decode (1 + draft_len) needs. + min_batch = 1 if min_batch is None: return False, ( f"CuTe DSL MLA decode is not a perf win for " From 16962d940e550cc80cbb0e889423b42f6876c17f Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 02:48:13 +0000 Subject: [PATCH 07/33] [None][fix] import mamba_effective_tp_size from the parent package mamba_cache_manager.py lives in tensorrt_llm._torch.pyexecutor.kv_cache but mamba_effective_tp_size is defined in tensorrt_llm._torch.pyexecutor.config_utils, so the single-dot relative import resolved to the non-existent kv_cache.config_utils and raised at module import. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 626093a5cf64..f1da42fd466d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -61,7 +61,7 @@ # Shared with the KV budget estimator so allocator and budgeting can never # diverge on the sharding rule (config_utils is import-cycle-free). -from .config_utils import mamba_effective_tp_size as _mamba_effective_tp_size +from ..config_utils import mamba_effective_tp_size as _mamba_effective_tp_size GB = 1 << 30 From 205ff61ed41f7b90974f49c9ba5cad154c47530d Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 02:49:08 +0000 Subject: [PATCH 08/33] [None][fix] pack Helix host state in the overlap extend and generation paths The extend loop only packed helix_position_offsets / helix_is_inactive_rank / helix_owned_new_tokens on the branch with no in-flight predecessor. With the overlap scheduler every later step takes the other branch, so update_helix_param wrote only a short prefix while prepare() reads through num_seqs, leaving stale ownership state in kv_lens (and, for an all-overlap batch, never arming _helix_spec_tokens_valid at all). Pack the same state in both branches, and append the single-token owned count in the plain generation loop so the per-sequence lists stay index-aligned in a mixed batch. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 25f52de48adf..5d40e9009ae7 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -5093,6 +5093,21 @@ def _helix_pack_extend(request, group: int) -> int: request.py_num_compressed_tokens) request.cached_tokens = (past_seen_token_num + runtime_tokens_per_gen_step) + if _has_cp_helix: + # In-flight predecessor: mirror the non-helix convention + # above -- positions are packed from the stale base (the + # overlap device correction adds the accepted count) and + # KV numbers assume full acceptance (the device recompute + # in recompute_helix_spec_buffers overrides them). The + # base is reconstructed GLOBALLY (see the no-previous + # branch: the token list is rank-local under helix). + group = runtime_tokens_per_gen_step + base = _helix_pack_extend(request, group) + local_full = _helix_local_len_host(base + group) + helix_owned_new_tokens.append(0) + num_cached_tokens_per_seq[-1] = ( + local_full - request.py_num_compressed_tokens) + request.cached_tokens = local_full if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: prompt_lengths.append(runtime_tokens_per_gen_step) @@ -5287,6 +5302,10 @@ def _helix_pack_extend(request, group: int) -> int: helix_is_inactive_rank.append( request.py_helix_is_inactive_rank) helix_position_offsets.append(position_id) + # Keep the per-seq owned-count list aligned when the + # spec path is active in the same batch. + helix_owned_new_tokens.append( + 0 if request.py_helix_is_inactive_rank else 1) request.cached_tokens = past_seen_token_num for beam in range(beam_width): From b8abb3fc16b960b863a42a48838449b6ab295424 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 02:49:38 +0000 Subject: [PATCH 09/33] [None][fix] use the global Helix position for the Q1 K-residual RoPE lookup With USE_HELIX_LOCAL_SLOTS, first_new_pos is the rank-local cache slot, so the MAX_GEN_TILES == 1 K-residual tail roped the key at the wrong position. Derive the rotary offset from rope_first_new_pos, matching the multi-tile tail which already keeps the two apart. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/attention/backends/fp4_mla/fp4_mla_kernels.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py index 3cb8ed6c1f95..f844195e0df2 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py @@ -1922,6 +1922,9 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( position = first_new_pos.to(tl.int64) if position < 0: return + # Under helix the cache slot is rank-local while RoPE must use the + # GLOBAL position; they coincide everywhere else. + rope_position = rope_first_new_pos.to(tl.int64) page_idx = position // page_size page_pos = position - page_idx * page_size physical_page_offset = page_start + page_idx @@ -1954,7 +1957,7 @@ def _fp4_mla_generation_fused_qk_rope_cache_update_kernel( latent_cache_ptr + latent_token * (NUM_DIM_BLOCKS * FP4_BLOCK) + tail_odd_d, ).to(tl.float32) rope_pair_offsets = tail_byte_offsets - rotary_offsets = position * (ROPE_DIM * 2) + rope_pair_offsets * 2 + rotary_offsets = rope_position * (ROPE_DIM * 2) + rope_pair_offsets * 2 cos = tl.load(rotary_cos_sin_ptr + rotary_offsets).to(tl.float32) sin = tl.load(rotary_cos_sin_ptr + rotary_offsets + 1).to(tl.float32) tail_even, tail_odd = _fp4_mla_rope_fp32(tail_even, tail_odd, cos, sin) From 5b0f93ab0555751b0fb9b9b53df118b1427deffb Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 02:50:00 +0000 Subject: [PATCH 10/33] [None][fix] count only generation tokens in the Helix verify-group reject q.shape[0] spans the whole batch, so a mixed batch with context requests exceeded num_seqs even with one query token per generation sequence, making FallbackFmha refuse a request it can serve and leaving dispatch with no candidate. Subtract num_ctx_tokens and compare against num_generations. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/fmha/fallback.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index fa79b6da6653..f8d9ec8467c5 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -91,7 +91,10 @@ def _is_supported( metadata.helix_position_offsets is not None and getattr(metadata, "_helix_spec_tokens_valid", False) and metadata.num_generations > 0 - and q.shape[0] > metadata.num_seqs + # Count generation tokens only: in a mixed batch ``q`` also holds + # the context tokens, which would otherwise trip this on a batch + # that has exactly one query token per generation sequence. + and q.shape[0] - metadata.num_ctx_tokens > metadata.num_generations ): return False if q is not None and q.dtype == torch.float8_e4m3fn: From 99f34b020ee8912680bff7bae9d91f407d8b52de Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 02:50:26 +0000 Subject: [PATCH 11/33] [None][fix] reject a nonpositive zeroKvMaskDivisor before launch The sender computes zeroKvMask[entryIdx / zeroKvMaskDivisor], so a zero divisor is a device-side integer division by zero. Validate it in launchHelixAllToAll, and clamp it in the torch op where an empty batch (entry_count == 0) would otherwise derive zero from a valid mask. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/helixAllToAll.cu | 4 ++++ cpp/tensorrt_llm/thop/alltoallOp.cpp | 6 +++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index f9eb6aee3009..11f85ce0edac 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.cu +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.cu @@ -702,6 +702,10 @@ size_t computeHelixWorkspaceSizePerRank(int cpSize) void launchHelixAllToAll(HelixAllToAllParams const& params, bool allowVariableField1, cudaStream_t stream) { + // The sender divides the entry index by this to index the mask, so a + // nonpositive divisor would be an integer division by zero on device. + TLLM_CHECK_WITH_INFO(params.zeroKvMask == nullptr || params.zeroKvMaskDivisor > 0, + "zeroKvMaskDivisor must be positive when zeroKvMask is set, got %d", params.zeroKvMaskDivisor); if (allowVariableField1) { constexpr uintptr_t kBulkCopyAlignment = 16; diff --git a/cpp/tensorrt_llm/thop/alltoallOp.cpp b/cpp/tensorrt_llm/thop/alltoallOp.cpp index 9c076850d29f..18c17be8bd85 100644 --- a/cpp/tensorrt_llm/thop/alltoallOp.cpp +++ b/cpp/tensorrt_llm/thop/alltoallOp.cpp @@ -240,7 +240,11 @@ std::tuple alltoall_helix_native(torch::Tensor par TORCH_CHECK(mask.numel() > 0 && entry_count % mask.numel() == 0, "zero_kv_mask numel (", mask.numel(), ") must divide the all-to-all entry count (", entry_count, ")"); params.zeroKvMask = reinterpret_cast(mask.data_ptr()); - params.zeroKvMaskDivisor = entry_count / mask.numel(); + // entry_count can legitimately be 0 for an empty batch, in which case + // the kernel never indexes the mask; keep the divisor positive so the + // launch-side contract still holds. + int const divisor = static_cast(entry_count / mask.numel()); + params.zeroKvMaskDivisor = divisor > 0 ? divisor : 1; } // Launch kernel From d6e41e2ff2ec5385be4a4f435d437c65810aa36e Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 02:51:56 +0000 Subject: [PATCH 12/33] [None][test] cover the shared Mamba sharding rule Parameterize mamba_effective_tp_size over attention-DP precedence, Helix TP x CP and standard TP, and assert get_states_bytes_per_layer routes through it rather than re-deriving the TP degree. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../kv_cache/test_kv_cache_budget_split.py | 79 +++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 2fa1abe69bd5..f8e3ed008945 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -1175,3 +1175,82 @@ def test_non_external_drafter_is_untouched(self, mocker): assert c._speculative_config.spec_dec_mode == SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL assert c._get_draft_kv_model_config() is draft_model_config + + +class TestMambaEffectiveTpSize: + """The sharding rule shared by the mamba pool allocator and the budget. + + ``mamba_cache_manager`` and ``MambaKVCacheParams.get_states_bytes_per_layer`` + must agree here, or the budget split withholds per-rank bytes the allocator + never uses. + """ + + @staticmethod + def _mapping(*, tp_size: int, cp_size: int, helix: bool, attention_dp: bool) -> SimpleNamespace: + return SimpleNamespace( + tp_size=tp_size, + cp_size=cp_size, + enable_attention_dp=attention_dp, + has_cp_helix=lambda: helix, + ) + + @pytest.mark.parametrize( + "tp_size,cp_size,helix,attention_dp,expected", + [ + # Attention-DP replicates the state on every rank, and takes + # precedence over both of the sharded cases below. + (8, 1, False, True, 1), + (8, 4, True, True, 1), + # Helix repurposes the CP ranks as plain TP for recurrent state. + (2, 8, True, False, 16), + (1, 16, True, False, 16), + # Standard TP: a non-helix mesh never shards state across CP. + (8, 1, False, False, 8), + (8, 4, False, False, 8), + (1, 1, False, False, 1), + ], + ) + def test_sharding_rule(self, tp_size, cp_size, helix, attention_dp, expected) -> None: + from tensorrt_llm._torch.pyexecutor.config_utils import mamba_effective_tp_size + + mapping = self._mapping( + tp_size=tp_size, cp_size=cp_size, helix=helix, attention_dp=attention_dp + ) + + assert mamba_effective_tp_size(mapping) == expected + + def test_budget_sizing_uses_the_shared_rule(self) -> None: + """``get_states_bytes_per_layer`` must not re-derive the TP degree.""" + import torch + + from tensorrt_llm._torch.pyexecutor.config_utils import ( + MambaKVCacheParams, + mamba_effective_tp_size, + ) + + params = MambaKVCacheParams( + state_size=128, + conv_kernel=4, + num_heads=128, + n_groups=8, + head_dim=64, + mamba_layer_mask=[True], + target_full_attention_layer_mask=[False], + num_mamba_layers=1, + num_draft_layers=0, + dtype=torch.float16, + mamba_ssm_cache_dtype=None, + ) + helix = self._mapping(tp_size=2, cp_size=8, helix=True, attention_dp=False) + assert mamba_effective_tp_size(helix) == 16 + plain_tp16 = self._mapping(tp_size=16, cp_size=1, helix=False, attention_dp=False) + replicated = self._mapping(tp_size=2, cp_size=8, helix=True, attention_dp=True) + + # Helix sizes the per-rank state as plain TP=tp*cp ... + assert params.get_states_bytes_per_layer(helix) == params.get_states_bytes_per_layer( + plain_tp16 + ) + # ... and attention-DP keeps the whole unsharded state per rank. + assert params.get_states_bytes_per_layer( + replicated + ) == 16 * params.get_states_bytes_per_layer(helix) From 0a74206ebb17fa6bd492a253fe88da28cccc3abf Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:59:53 +0000 Subject: [PATCH 13/33] [None][fix] arm the per-token Helix metadata for speculative decoding only The generation loop appended to helix_owned_new_tokens whenever helix was on. A non-empty list arms _helix_spec_tokens_valid in update_helix_param, but the per-token helix_local_slots / helix_kv_bounds it gates are filled by recompute_helix_spec_buffers, which _preprocess_inputs runs only when enable_spec_decode is set. Ordinary helix generation therefore pointed the append kernel and the MLA decode mask at uninitialized buffers. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5d40e9009ae7..4af1bbd19f63 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -5303,9 +5303,16 @@ def _helix_pack_extend(request, group: int) -> int: request.py_helix_is_inactive_rank) helix_position_offsets.append(position_id) # Keep the per-seq owned-count list aligned when the - # spec path is active in the same batch. - helix_owned_new_tokens.append( - 0 if request.py_helix_is_inactive_rank else 1) + # spec path is active in the same batch. Only then: + # a non-empty list arms _helix_spec_tokens_valid, and + # the per-token slots/bounds it gates are filled by + # recompute_helix_spec_buffers, which _preprocess_inputs + # runs under enable_spec_decode only. Populating it in + # ordinary generation would point consumers at + # uninitialized buffers. + if self.enable_spec_decode: + helix_owned_new_tokens.append( + 0 if request.py_helix_is_inactive_rank else 1) request.cached_tokens = past_seen_token_num for beam in range(beam_width): From 57e2ca1f78a8a6a2135cbf8b6a6e2211e6929e34 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:00:56 +0000 Subject: [PATCH 14/33] [None][fix] forward kv_bounds into the FP8 CuTe DSL MLA decode runner The FP8 op accepted kv_bounds but stopped building the runner input list at softmax_stats, so the runner read inputs[9] as absent and compiled the non-helix variant. Helix verify groups then fell back to ordinary causal bounds and masked valid cached tokens on nonowning ranks. Append it exactly as the FP16 op does; mutates_args and the fake already cover the argument. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 0fba18388303..1ff1ab5e90fc 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -11772,6 +11772,9 @@ def cute_dsl_mla_decode_fp8_blackwell( kv_bounds: Optional[torch.Tensor], ) -> None: """CuTe DSL FP8 MLA decode (Blackwell SM100/SM103). + + kv_bounds: helix speculative verify groups -- per-token rank-local + attention bounds of shape (B * seq_len_q,), int32. """ if (sm_version := get_sm_version()) not in (100, 103): raise ValueError( @@ -11790,7 +11793,7 @@ def cute_dsl_mla_decode_fp8_blackwell( ) inputs = [ q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o, - workspace, softmax_stats + workspace, softmax_stats, kv_bounds ] tuner = AutoTuner.get() _, best_tactic = tuner.choose_one( From 4a2e323adf7b4ac65a54b1341500e066a27e808f Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:12:00 +0000 Subject: [PATCH 15/33] [None][fix] gate the Helix spec arming at a single choke point _helix_spec_tokens_valid is armed purely by update_helix_param receiving a non-None owned-count list, with no speculative-decoding condition of its own, so any packing loop that appended under helix alone armed the spec path and sent its consumers to the never-filled helix_local_slots / helix_kv_bounds. Decide it once where the list is handed over instead of per append site, which also covers the two extend branches. Arming inside recompute_helix_spec_buffers was considered and rejected: prepare() reads the flag for its kv_lens branch before _preprocess_inputs runs the recompute, so that ordering would disable the spec kv_lens path entirely. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 4af1bbd19f63..333d97267a95 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -5303,16 +5303,11 @@ def _helix_pack_extend(request, group: int) -> int: request.py_helix_is_inactive_rank) helix_position_offsets.append(position_id) # Keep the per-seq owned-count list aligned when the - # spec path is active in the same batch. Only then: - # a non-empty list arms _helix_spec_tokens_valid, and - # the per-token slots/bounds it gates are filled by - # recompute_helix_spec_buffers, which _preprocess_inputs - # runs under enable_spec_decode only. Populating it in - # ordinary generation would point consumers at - # uninitialized buffers. - if self.enable_spec_decode: - helix_owned_new_tokens.append( - 0 if request.py_helix_is_inactive_rank else 1) + # spec path is active in the same batch. Whether the + # list arms the spec path at all is decided once, at + # the update_helix_param call below. + helix_owned_new_tokens.append( + 0 if request.py_helix_is_inactive_rank else 1) request.cached_tokens = past_seen_token_num for beam in range(beam_width): @@ -5730,11 +5725,20 @@ def previous_seq_slots_device(): num_first_draft]] += accepted_tokens if self.mapping.has_cp_helix(): + # A non-None owned-count list is what arms + # _helix_spec_tokens_valid, and the per-token slots/bounds that + # flag gates are only ever filled by recompute_helix_spec_buffers, + # which _preprocess_inputs runs under enable_spec_decode. Gate the + # hand-off here, at the single choke point, so no packing loop can + # arm the spec path for ordinary helix generation and send its + # consumers to uninitialized buffers. + helix_spec_active = bool(self.enable_spec_decode + and helix_owned_new_tokens) attn_metadata.update_helix_param( helix_position_offsets=helix_position_offsets, helix_is_inactive_rank=helix_is_inactive_rank, helix_owned_new_tokens=(helix_owned_new_tokens - if helix_owned_new_tokens else None), + if helix_spec_active else None), ) if not attn_metadata.is_cuda_graph: From b5443ff8d6d02ec4ccc98e810adbfb1c17b60903 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:13:32 +0000 Subject: [PATCH 16/33] [None][chore] drop the dead num_ctx_tokens parameter from the helix recompute The single call site always passes 0 because the caller already subtracts num_ctx_tokens when computing num_gen_tokens, which made the per-token helix buffers look batch-indexed when they are generation-relative -- the opposite of the kv_lens_cuda write a few lines below. Drop the parameter and state both index bases in the docstring. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/attention/backends/trtllm.py | 29 ++++++++++--------- .../_torch/pyexecutor/model_engine.py | 2 +- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 31df0b9768d6..6b91cc749005 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -693,8 +693,7 @@ def helix_local_len_vec(self, global_lens: torch.Tensor) -> torch.Tensor: rem = global_lens - full * ledger return full * phys + (rem - cp_rank * phys).clamp_(0, phys) - def recompute_helix_spec_buffers(self, num_ctx_tokens: int, - num_gen_tokens: int, + def recompute_helix_spec_buffers(self, num_gen_tokens: int, tokens_per_gen_seq: int) -> None: """Derive per-token helix buffers from (corrected) global positions. @@ -702,9 +701,15 @@ def recompute_helix_spec_buffers(self, num_ctx_tokens: int, to helix_position_offsets, so every derived quantity reflects the real committed length even though the host packed provisional values. Static shapes only; safe under CUDA graph capture. + + Two index bases meet here, and they are not the same: + * helix_position_offsets / helix_local_slots / helix_kv_bounds are + GENERATION-RELATIVE -- the packing loops only append for extend + and generation rows, so token 0 is the first generation token. + * kv_lens_cuda is BATCH-indexed, hence the num_contexts offset on + the write below. """ - pos = self.helix_position_offsets[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens] + pos = self.helix_position_offsets[:num_gen_tokens] phys = self.kv_cache_manager.tokens_per_block cp_rank = self.mapping.cp_rank cp_size = self.mapping.cp_size @@ -713,22 +718,18 @@ def recompute_helix_spec_buffers(self, num_ctx_tokens: int, local_before = self.helix_local_len_vec(pos) # Scalar overload: no per-step allocation (CUDA-graph capture treats # these ops as part of the graph; keep them allocation-free). - self.helix_local_slots[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens].copy_( - torch.where(active, local_before, -1)) - self.helix_kv_bounds[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens].copy_( - self.helix_local_len_vec(pos + 1)) + self.helix_local_slots[:num_gen_tokens].copy_( + torch.where(active, local_before, -1)) + self.helix_kv_bounds[:num_gen_tokens].copy_( + self.helix_local_len_vec(pos + 1)) # Per-sequence rank-local kv length = bound of the sequence's last # token (attention over committed + owned in-flight tokens). assert num_gen_tokens % tokens_per_gen_seq == 0, ( f"helix spec expects uniform verify groups: {num_gen_tokens} gen " f"tokens not divisible by group size {tokens_per_gen_seq}") num_gen_seqs = num_gen_tokens // tokens_per_gen_seq - last_bounds = self.helix_kv_bounds[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens].view( - num_gen_seqs, - tokens_per_gen_seq)[:, -1] + last_bounds = self.helix_kv_bounds[:num_gen_tokens].view( + num_gen_seqs, tokens_per_gen_seq)[:, -1] self.kv_lens_cuda[self.num_contexts:self.num_contexts + num_gen_seqs].copy_(last_bounds) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 333d97267a95..11eff3822195 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3875,7 +3875,7 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]): md.helix_position_offsets[:helix_gen_tokens] += ( self.previous_pos_id_offsets_cuda[:helix_gen_tokens]) md.recompute_helix_spec_buffers( - 0, helix_gen_tokens, + helix_gen_tokens, self.get_runtime_tokens_per_gen_step( self.runtime_draft_len)) md.on_update_kv_lens() From 568a315e9c08bcc070d4011e4a9a993739dd1b4e Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:13:59 +0000 Subject: [PATCH 17/33] [None][fix] reject non-uniform Helix verify groups instead of mis-slicing The recompute assumes every generation row contributes tokens_per_gen_seq tokens. That holds on the overlap path, but _preprocess_inputs runs it with the overlap scheduler disabled too, where the extend loop packs a per-request 1 + get_draft_token_length(request) and a request entering with no draft tokens becomes a single-token generation row instead -- static draft length does not pad. Divisibility alone then lets a mixed batch through and writes kv_lens_cuda for the wrong number of rows with values from the wrong tokens. Check the row count as well so such a batch fails loudly. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/trtllm.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 6b91cc749005..7a07873d5e7f 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -702,6 +702,14 @@ def recompute_helix_spec_buffers(self, num_gen_tokens: int, real committed length even though the host packed provisional values. Static shapes only; safe under CUDA graph capture. + ``tokens_per_gen_seq`` is the uniform verify-group width. Non-uniform + groups are rejected rather than silently mis-sliced: without the + overlap scheduler the extend loop packs a per-request + ``1 + get_draft_token_length(request)`` and a request entering with no + draft tokens is packed as a single-token generation row instead, so a + batch can arrive whose total happens to divide but whose rows do not + line up. + Two index bases meet here, and they are not the same: * helix_position_offsets / helix_local_slots / helix_kv_bounds are GENERATION-RELATIVE -- the packing loops only append for extend @@ -728,6 +736,13 @@ def recompute_helix_spec_buffers(self, num_gen_tokens: int, f"helix spec expects uniform verify groups: {num_gen_tokens} gen " f"tokens not divisible by group size {tokens_per_gen_seq}") num_gen_seqs = num_gen_tokens // tokens_per_gen_seq + # Divisibility alone does not imply uniformity: a batch of mixed group + # widths can still divide and would then write the wrong number of + # kv_lens_cuda rows with values taken from the wrong tokens. + assert num_gen_seqs == self.num_generations, ( + f"helix spec expects uniform verify groups: {num_gen_tokens} gen " + f"tokens over {self.num_generations} generation rows do not all " + f"have width {tokens_per_gen_seq}") last_bounds = self.helix_kv_bounds[:num_gen_tokens].view( num_gen_seqs, tokens_per_gen_seq)[:, -1] self.kv_lens_cuda[self.num_contexts:self.num_contexts + From 7f446772ae19ea38a34431794abf2748cc760f95 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:14:37 +0000 Subject: [PATCH 18/33] [None][chore] read _helix_spec_tokens_valid directly where the type is known It is a declared field on TrtllmAttentionMetadata, so getattr with a default adds nothing at the call sites that already hold that type (or reach it via a non-None helix_kv_bounds). The remaining getattr uses guard metadata objects that may come from another attention backend and stay as they are. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/fmha/fallback.py | 2 +- tensorrt_llm/_torch/attention/backends/trtllm.py | 4 ++-- tensorrt_llm/_torch/attention/mla.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index f8d9ec8467c5..5631950e9e06 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -89,7 +89,7 @@ def _is_supported( # in the library list, this makes dispatch raise. if ( metadata.helix_position_offsets is not None - and getattr(metadata, "_helix_spec_tokens_valid", False) + and metadata._helix_spec_tokens_valid and metadata.num_generations > 0 # Count generation tokens only: in a mixed batch ``q`` also holds # the context tokens, which would otherwise trip this on a batch diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 7a07873d5e7f..4b73cbab3717 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -878,7 +878,7 @@ def prepare(self) -> None: if self.enable_helix: # If helix is inactive, attend to the previously cached tokens only. assert cached_token_lens is not None, "cached_token_lens should be set for helix" - if getattr(self, '_helix_spec_tokens_valid', False): + if self._helix_spec_tokens_valid: # Speculative verify groups: a group may straddle a page # boundary, so ownership of this step's new tokens is a # per-sequence COUNT, not a boolean. Provisional host values; @@ -2547,7 +2547,7 @@ def mla_rope_generation( helix_tensor_params = [ metadata.helix_position_offsets, metadata.helix_is_inactive_rank ] - if getattr(metadata, '_helix_spec_tokens_valid', False): + if metadata._helix_spec_tokens_valid: # Speculative verify groups: per-token KV write slots (-1 = this # rank does not own the token's position). The append kernel then # gates and addresses per token instead of per sequence. diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index 988114b225a7..e331d9ebcf24 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -795,9 +795,9 @@ def _attn_forward_gen( assert self.kv_lora_rank == kv_lora_rank helix_kv_bounds = getattr(attn_metadata, "helix_kv_bounds", None) - if helix_kv_bounds is not None and getattr( - attn_metadata, "_helix_spec_tokens_valid", False - ): + # helix_kv_bounds is non-None only on TrtllmAttentionMetadata, + # where _helix_spec_tokens_valid is a declared field. + if helix_kv_bounds is not None and attn_metadata._helix_spec_tokens_valid: # Speculative verify groups: KV ownership is per-TOKEN. A rank # owning only the tail page of a group has zero visible KV for # the group's leading tokens while its per-sequence kv_len is From f8a7c9768273a8b0b31d7075c40e93193041e6d5 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:14:57 +0000 Subject: [PATCH 19/33] [None][chore] use the descriptive mamba_effective_tp_size name at call sites The private-looking alias hid that this is the shared rule imported from config_utils rather than something local to the cache manager. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache/mamba_cache_manager.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index f1da42fd466d..ad417618950b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -61,7 +61,7 @@ # Shared with the KV budget estimator so allocator and budgeting can never # diverge on the sharding rule (config_utils is import-cycle-free). -from ..config_utils import mamba_effective_tp_size as _mamba_effective_tp_size +from ..config_utils import mamba_effective_tp_size GB = 1 << 30 @@ -527,7 +527,7 @@ def __init__( self._seed_request_counter = 0 # get tp size - tp_size = _mamba_effective_tp_size(mapping) + tp_size = mamba_effective_tp_size(mapping) # derive mamba parameters for conv and ssm states d_inner = head_dim * num_heads @@ -2343,7 +2343,7 @@ def __init__( return # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) - tp_size = _mamba_effective_tp_size(mapping) + tp_size = mamba_effective_tp_size(mapping) d_inner = mamba_head_dim * mamba_num_heads conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state nheads = mamba_num_heads @@ -3133,7 +3133,7 @@ def __init__( and self.local_num_mamba_layers > 0) if self.local_num_mamba_layers > 0: - tp_size = _mamba_effective_tp_size(mapping) + tp_size = mamba_effective_tp_size(mapping) d_inner = mamba_head_dim * mamba_num_heads grouped_state_dim = mamba_n_groups * mamba_d_state conv_dim = d_inner + 2 * grouped_state_dim From 314398b1a90df0394e87736219940992f43decbe Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:15:18 +0000 Subject: [PATCH 20/33] [None][chore] drop the model name from the generic helix x speculation error The check sits ahead of the is_kimi_linear dispatch and fires for any hybrid model, so naming Kimi K3 in the message misleads every other one. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 74b8e04c5c54..cb6466258562 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -239,7 +239,7 @@ def get_kv_cache_manager_cls( and model_config.mapping.has_cp_helix() and model_config.spec_config is not None and not use_v2): raise ValueError( - "Kimi K3 helix with speculative decoding requires " + "Helix with speculative decoding requires " "kv_cache_config.use_kv_cache_manager_v2=True; the V1-family " "hybrid managers do not implement per-token verify-group " "bookkeeping.") From 00937dcdcb9ac6f9be4e1f0b1cc0950da9d51a4f Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:16:17 +0000 Subject: [PATCH 21/33] [None][chore] report when the Helix NCCL reformat outgrows dynamo's cache The dynamic=False specialization is one compile per CUDA-graph batch bucket, and past torch._dynamo.config.cache_size_limit dynamo silently runs the frame eagerly, giving back the sanitize/transpose fusion these helpers exist for with no error and no log line. Count the distinct specializations at the call site and warn once when the budget is exceeded, so the regression shows up in the log rather than only as missing triton_poi_fused_* in a kernel trace. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/attention/attention.py | 47 +++++++++++++++++++++- 1 file changed, 46 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention/attention.py b/tensorrt_llm/_torch/attention/attention.py index 7565481f54e4..4685edafd65c 100644 --- a/tensorrt_llm/_torch/attention/attention.py +++ b/tensorrt_llm/_torch/attention/attention.py @@ -193,6 +193,47 @@ def _helix_sanitize_empty_kv( return partial_o, softmax_stats +# Distinct input specializations the NCCL reformat helpers below have been +# asked to compile. ``dynamic=False`` specializes on the input shapes and +# ``cp_size``, so this set tracks the same thing dynamo's per-frame cache does. +_HELIX_NCCL_SHAPES: set = set() +_HELIX_NCCL_LIMIT_REPORTED = False + + +def _helix_note_nccl_specialization(partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + cp_size: int) -> None: + """Warn once if the Helix NCCL reformat is past dynamo's recompile budget. + + Past ``cache_size_limit`` dynamo stops compiling and runs the frame eagerly + without raising, which silently gives back the fusion these helpers exist + for. Report it once per process so the regression is visible in the log + instead of only in a kernel trace. + """ + global _HELIX_NCCL_LIMIT_REPORTED + if _HELIX_NCCL_LIMIT_REPORTED: + return + key = (tuple(partial_o.shape), tuple(softmax_stats.shape), cp_size) + if key in _HELIX_NCCL_SHAPES: + return + _HELIX_NCCL_SHAPES.add(key) + try: + import torch._dynamo as _dynamo + limit = getattr(_dynamo.config, "cache_size_limit", None) + except ImportError: + limit = None + if limit is None or len(_HELIX_NCCL_SHAPES) <= limit: + return + _HELIX_NCCL_LIMIT_REPORTED = True + logger.warning( + "Helix NCCL all-to-all reformat has seen %d distinct input " + "specializations, above torch._dynamo.config.cache_size_limit=%d. " + "Dynamo stops recompiling past that limit and runs these helpers " + "eagerly without raising, losing the sanitize/transpose fusion. " + "Reduce the number of CUDA-graph batch buckets or raise " + "cache_size_limit.", len(_HELIX_NCCL_SHAPES), limit) + + @torch.compile(dynamic=False) def _helix_nccl_pre_alltoall( partial_o: torch.Tensor, @@ -214,7 +255,9 @@ def _helix_nccl_pre_alltoall( The cost is one specialization per CUDA-graph batch bucket. Exceeding dynamo's ``cache_size_limit`` falls back to eager SILENTLY -- the symptom is - ``triton_poi_fused_*`` disappearing from the trace, not an error. + ``triton_poi_fused_*`` disappearing from the trace, not an error, so + ``_helix_note_nccl_specialization`` counts the distinct shapes at the call + site and says so once. """ partial_o, softmax_stats = _helix_sanitize_empty_kv(partial_o, softmax_stats, @@ -263,6 +306,8 @@ def _helix_post_process( if mapping.cp_config.get("use_nccl_for_alltoall", True): # NCCL path. Sanitize is folded into _helix_nccl_pre_alltoall so # inductor can fuse it into the reformat. + _helix_note_nccl_specialization(partial_o, softmax_stats, + mapping.cp_size) chunks = _helix_nccl_pre_alltoall(partial_o, softmax_stats, zero_kv_mask, mapping.cp_size) gathered = alltoall_helix(chunks, mapping.cp_group) From b15c64501af518d1fc15dddb4d14326fa04b5465 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:16:47 +0000 Subject: [PATCH 22/33] [None][doc] state the single-token assumption in _set_helix_rank_fields py_helix_decode_group_index advances once per successful allocation regardless of how many tokens the group committed, so the derived position falls behind as soon as a draft token is accepted. The docstring claimed the formula stays exact under speculation; record the real assumption, why the speculative path does not depend on it, and what a proper fix requires. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../pyexecutor/kv_cache/kv_cache_manager_v2.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index ce6db25a8199..9285e9fe1d3f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3264,9 +3264,20 @@ def _set_helix_rank_fields(self, req: LlmRequest) -> None: from ``py_decoding_iter``: the sampler advances that counter after scheduling under the overlap loop, so a schedule-time read is one step behind and would repeat the first decode position, overwriting - the first generated token's KV. Multi-token verify groups advance - py_helix_decode_group_index per committed group, so this formula - stays exact under DSpark speculation. + the first generated token's KV. + + The counter advances by one per successful allocation, so ``pos`` is + exact only while each iteration commits exactly one token. Under + speculation an iteration can commit ``1 + accepted`` tokens and this + estimate falls behind, taking ``py_helix_is_inactive_rank`` and + ``seqlen_this_rank_cp`` with it. That is tolerable today only because + the speculative path never reads these fields: ``_helix_pack_extend`` + in model_engine rebuilds the global position from + ``total_input_len_cp`` plus the rank-invariant generated count. A + request that falls back to the plain generation loop mid-run (a step + that yields no draft tokens) would read a stale value -- advancing the + counter by the committed token count is the fix, and needs the + acceptance count to be available at schedule time. """ step = req.py_helix_decode_group_index + 1 pos = req.total_input_len_cp + step - 1 From 19ba373b74dac4f26de7908ab82adda1a7702026 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:17:09 +0000 Subject: [PATCH 23/33] [None][doc] state the generation-only index base of the helix kv_lens branch The per-sequence helix buffers are packed generation-first while cached_token_lens is contexts-first, so the [:num_seqs] slicing in both branches is only correct for a batch with no context rows. Record the invariant rather than leaving a second consumer to imply it. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/trtllm.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 4b73cbab3717..026c59e21936 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -878,6 +878,13 @@ def prepare(self) -> None: if self.enable_helix: # If helix is inactive, attend to the previously cached tokens only. assert cached_token_lens is not None, "cached_token_lens should be set for helix" + # Both branches index a per-GENERATION-sequence buffer with + # [:num_seqs] against a contexts-first cached_token_lens, which + # only lines up when the batch carries no context rows. That has + # been the (unstated) helix invariant since the boolean branch + # landed on main; helix_is_inactive_rank_cpu is uninitialized + # memory for any row the packing loops did not write, so a mixed + # batch is already wrong rather than merely imprecise. if self._helix_spec_tokens_valid: # Speculative verify groups: a group may straddle a page # boundary, so ownership of this step's new tokens is a From 69fc33d62b3d9a6f6f4d843e8005b4cf2bdfd401 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:40:40 +0000 Subject: [PATCH 24/33] [None][fix] pair the helix per-sequence buffers with the generation slice update_helix_param writes helix_is_inactive_rank_cpu and helix_owned_new_tokens_cpu over exactly [0, num_generations), because the model_engine packing loops are initialized after the context loop and only extend and plain-generation rows append. Reading them as [:num_seqs] against a contexts-first cached_token_lens shifted every pairing by num_contexts and ran off the end of the written region -- uninitialized memory for the boolean buffer, which then reached the FMHA kernel as cache_seq_lens. Slice the batch-indexed tensors to the generation range instead, and give context rows the same rule as the non-helix path since they are never packed into these buffers. The buffers stay generation-relative because every device consumer indexes them that way. Supersedes the comment-only note from e6a18b539a. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/attention/backends/trtllm.py | 34 ++++++++++++------- 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 026c59e21936..4bcfa0eab509 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -878,25 +878,35 @@ def prepare(self) -> None: if self.enable_helix: # If helix is inactive, attend to the previously cached tokens only. assert cached_token_lens is not None, "cached_token_lens should be set for helix" - # Both branches index a per-GENERATION-sequence buffer with - # [:num_seqs] against a contexts-first cached_token_lens, which - # only lines up when the batch carries no context rows. That has - # been the (unstated) helix invariant since the boolean branch - # landed on main; helix_is_inactive_rank_cpu is uninitialized - # memory for any row the packing loops did not write, so a mixed - # batch is already wrong rather than merely imprecise. + # The helix per-sequence buffers are GENERATION-relative: the + # packing loops in model_engine append only for extend and plain + # generation rows, so update_helix_param writes exactly + # [0, num_generations). Every device consumer indexes them the + # same way (the MLA rope generation kernel, the XQA preprocessing + # kernels, the FP4 MLA generation kernel), so the host read cannot + # slice them from 0 against a contexts-first cached_token_lens: + # that both shifts every pairing by num_contexts and reads past + # the written region, which for helix_is_inactive_rank_cpu is + # uninitialized memory. Pair them with the generation slice of the + # batch-indexed tensors instead. + num_gen = self.num_generations + gen = slice(self.num_contexts, self.num_seqs) + # Context rows are not part of a verify group and are not packed + # into the helix buffers at all; they append every one of their + # tokens, exactly like the non-helix path below. + kv_lens = cached_token_lens + self.seq_lens_kv if self._helix_spec_tokens_valid: # Speculative verify groups: a group may straddle a page # boundary, so ownership of this step's new tokens is a # per-sequence COUNT, not a boolean. Provisional host values; # recompute_helix_spec_buffers overrides the device copy # after the overlap correction. - kv_lens = cached_token_lens + \ - self.helix_owned_new_tokens_cpu[:self.num_seqs] + kv_lens[gen] = (cached_token_lens[gen] + + self.helix_owned_new_tokens_cpu[:num_gen]) else: - active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs] - kv_lens = cached_token_lens.clone() - kv_lens[active_rank] += self.seq_lens_kv[active_rank] + inactive_rank = self.helix_is_inactive_rank_cpu[:num_gen] + kv_lens[gen] = torch.where(inactive_rank, + cached_token_lens[gen], kv_lens[gen]) else: kv_lens = cached_token_lens + \ self.seq_lens_kv if cached_token_lens is not None else self.seq_lens_kv From 554da981e02e3df67d4185e25ba8421b6296ef10 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:41:48 +0000 Subject: [PATCH 25/33] [None][fix] pass the new kv_bounds argument in the standalone MLA decode runners kv_bounds was inserted positionally between cache_seqs and block_split_kvs in the FP16 and FP8 decode entry points, but the standalone run() in each file still called cute.compile, the compiled kernel and testing.JitArguments with the old positional list, shifting block_split_kvs into the kv_bounds slot. Pass None in the new position; the standalone path does not exercise helix. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../blackwell/attention/mla/mla_decode_fp16.py | 3 +++ .../cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py index c56ba8322f2e..c13fd2800218 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py @@ -4261,6 +4261,7 @@ def create_workspace(num_heads, seq_len_q, latent_dim, batch_size, split_kv, workspace, split_kv, cache_seqs, + None, # kv_bounds: helix-only, not exercised here block_split_kvs, softmax_scale, output_scale, @@ -4376,6 +4377,7 @@ def torch_reference_mla( workspace, split_kv, cache_seqs, + None, # kv_bounds: helix-only, not exercised here block_split_kvs, softmax_scale, output_scale, @@ -4513,6 +4515,7 @@ def generate_tensors(): workspace, _split_kv, cache_seqs, + None, # kv_bounds: helix-only, not exercised here block_split_kvs, softmax_scale, output_scale, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py index 75d5c623d8e7..2f6120b0ebfe 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py @@ -4210,6 +4210,7 @@ def create_workspace(num_heads, seq_len_q, latent_dim, batch_size, split_kv, workspace, split_kv, cache_seqs, + None, # kv_bounds: helix-only, not exercised here block_split_kvs, softmax_scale, output_scale, @@ -4325,6 +4326,7 @@ def torch_reference_mla( workspace, split_kv, cache_seqs, + None, # kv_bounds: helix-only, not exercised here block_split_kvs, softmax_scale, output_scale, @@ -4462,6 +4464,7 @@ def generate_tensors(): workspace, _split_kv, cache_seqs, + None, # kv_bounds: helix-only, not exercised here block_split_kvs, softmax_scale, output_scale, From c7efa363436c32241196e95373a297c1edf92c37 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:55:38 +0000 Subject: [PATCH 26/33] [None][fix] narrow the H=96 CuteDSL exception to multi-token Helix groups The fallback admitted every seq_len_q at 96 heads whether or not helix was involved, which silently reopened the non-helix multi-token shapes that the measured table rejects on main. Take it out of _PERF_MIN_BATCH_FP8, which goes back to being a pure measured-win table identical to main, and decide it at the one call site that can see the helix state: bypass the perf gate only for num_heads == 96 with seq_len_q > 1 under helix, where TRTLLM-Gen rejects 64 < num_heads_q < 128 and there is no other kernel to fall back to. Single-token H=96 still goes through the table entry it already has. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../attention/backends/fmha/cute_dsl_mla.py | 38 +++++++++++-------- 1 file changed, 23 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 2764692e7aed..ac97de1d550c 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -254,13 +254,6 @@ def _is_perf_favorable( (128, 2): 32, } min_batch = _PERF_MIN_BATCH_FP8.get((num_heads, seq_len_q)) - if min_batch is None and num_heads == 96: - # For H=96 this table is a correctness constraint, not a perf - # tradeoff: TRTLLM-Gen rejects 64 < num_heads_q < 128 outright, so - # falling through does not reach a faster kernel, it reaches an - # executor-init failure. Admit every seq_len_q, which is what - # speculative decode (1 + draft_len) needs. - min_batch = 1 if min_batch is None: return False, ( f"CuTe DSL MLA decode is not a perf win for " @@ -329,15 +322,30 @@ def _is_supported_with_reason( from tensorrt_llm._torch.autotuner import AutoTuner - # Perf gate (NOT a correctness limit). - favorable, reason = self._is_perf_favorable( - attn.num_heads, - None if AutoTuner.get().is_tuning_mode else batch_size, - seq_len_q, - self._get_kernel_dtype(attn, q), + # A multi-token verify group at H=96 has nowhere else to go: TRTLLM-Gen + # rejects 64 < num_heads_q < 128 outright, so falling through the perf + # gate does not reach a faster kernel, it reaches an executor-init + # failure. That makes this a correctness carve-out rather than a perf + # tradeoff, so it is decided here instead of being folded into + # _PERF_MIN_BATCH_FP8, which stays a pure measured-win table. + # Deliberately narrow: single-token H=96 is already a measured table + # entry and still goes through the gate, and non-helix multi-token + # H=96 keeps falling back exactly as it does today. Helix with + # seq_len_q > 1 implies _helix_spec_tokens_valid -- the helix block + # above returns False otherwise. + helix_h96_verify_group = ( + attn.num_heads == 96 and seq_len_q > 1 and meta.helix_position_offsets is not None ) - if not favorable: - return False, reason + if not helix_h96_verify_group: + # Perf gate (NOT a correctness limit). + favorable, reason = self._is_perf_favorable( + attn.num_heads, + None if AutoTuner.get().is_tuning_mode else batch_size, + seq_len_q, + self._get_kernel_dtype(attn, q), + ) + if not favorable: + return False, reason if meta.kv_cache_manager is None: return False, "KV cache manager is required." if fwd.output is None: From 23f430be623afa32ebbd6d8029706fca76889f17 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:20:40 +0000 Subject: [PATCH 27/33] [None][chore] give the helix round-robin local length one definition The rule that page b of the ledger lives on CP rank b % cp_size had three implementations kept in sync by comment: the cache manager's scalar _helix_local_len, the model engine's _helix_local_len_host closure, and the attention metadata's vectorised helix_local_len_vec. A drift between them writes KV to the wrong rank without raising, and only on groups that straddle a page boundary -- the hardest case to reproduce. Move the rule into tensorrt_llm/_torch/utils.py as helix_local_len and helix_local_len_tensor, taking tokens_per_block, cp_size and cp_rank explicitly, and delegate all three sites to them. The three expressions are equivalent today, so this changes no behaviour: a sweep over tokens_per_block, cp_size, every cp_rank and every global length through several ledger periods finds no disagreement between them, with the repo's own token-by-token reference in test_kv_cache_manager_v2_helix_superblock.py, or with the partition invariant that the per-rank lengths sum to the global one. utils.py is a leaf -- it imports neither pyexecutor nor attention, both of which already import it, so no new dependency direction appears. The tensor form keeps the original operation sequence, including the in-place clamp on the temporary the subtraction produces, because it runs on the CUDA-graph capture path. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/attention/backends/trtllm.py | 17 ++++--- .../kv_cache/kv_cache_manager_v2.py | 15 ++++-- .../_torch/pyexecutor/model_engine.py | 16 +++---- tensorrt_llm/_torch/utils.py | 46 +++++++++++++++++++ 4 files changed, 72 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 4bcfa0eab509..7ab975a02a33 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -38,7 +38,7 @@ from ...pyexecutor.config_utils import is_mla from ...utils import (compute_swizzled_sf_shape, get_global_attrs, - get_model_extra_attrs) + get_model_extra_attrs, helix_local_len_tensor) from .fmha.manager import FmhaManager from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, @@ -683,15 +683,14 @@ def helix_local_len_vec(self, global_lens: torch.Tensor) -> torch.Tensor: For each global sequence length g, returns the number of the first g tokens whose ledger page lives on this CP rank (page b -> rank - b % cp_size). Mirrors KVCacheManagerV2._helix_local_len. + b % cp_size). The rule and its scalar twin live in + ``_torch.utils``; KVCacheManagerV2._helix_local_len and the host + packing in model_engine use the same definition. """ - phys = self.kv_cache_manager.tokens_per_block - cp_size = self.mapping.cp_size - cp_rank = self.mapping.cp_rank - ledger = phys * cp_size - full = torch.div(global_lens, ledger, rounding_mode='floor') - rem = global_lens - full * ledger - return full * phys + (rem - cp_rank * phys).clamp_(0, phys) + return helix_local_len_tensor(global_lens, + self.kv_cache_manager.tokens_per_block, + self.mapping.cp_size, + self.mapping.cp_rank) def recompute_helix_spec_buffers(self, num_gen_tokens: int, tokens_per_gen_seq: int) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 9285e9fe1d3f..2b47f3f5982d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -27,7 +27,7 @@ from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp -from tensorrt_llm._torch.utils import maybe_compile +from tensorrt_llm._torch.utils import helix_local_len, maybe_compile from tensorrt_llm._utils import ( TensorWrapper, binding_to_torch_dtype, @@ -3251,10 +3251,15 @@ def _effective_draft_len(self, req: LlmRequest) -> int: def _helix_local_len(self, global_len: int) -> int: """Tokens of the first ``global_len`` owned by this CP rank - (continuation round-robin: page b lives on rank b %% cp).""" - phys = self.tokens_per_block - full, rem = divmod(global_len, self._ledger_tokens_per_block) - return full * phys + min(max(rem - self._helix_cp_rank * phys, 0), phys) + (continuation round-robin: page b lives on rank b %% cp). + + The rule itself lives in ``_torch.utils.helix_local_len`` so the host + packing in model_engine and the tensor form in the attention metadata + cannot drift from it. + """ + return helix_local_len( + global_len, self.tokens_per_block, self._helix_cp_size, self._helix_cp_rank + ) def _set_helix_rank_fields(self, req: LlmRequest) -> None: """Derive the per-rank helix fields from the global position. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 11eff3822195..c8386d9de9ce 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -73,7 +73,7 @@ from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_static_draft_len, update_draft_len from ..utils import (get_model_extra_attrs, - get_per_request_prefill_cuda_graph_flag, + get_per_request_prefill_cuda_graph_flag, helix_local_len, set_per_request_prefill_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner @@ -4947,20 +4947,20 @@ def append_cross_attention_state(request: LlmRequest, # Helix bookkeeping is needed by BOTH the extend (speculative verify # group) and the plain generation packing loops below, so initialize # it ahead of them. Positions are global; KV ownership follows the - # round-robin ledger (page b -> rank b % cp), mirrored host-side here - # (KVCacheManagerV2._helix_local_len) for provisional packing values. + # round-robin ledger (page b -> rank b % cp); the host-side + # provisional packing values come from the one shared definition in + # _torch.utils.helix_local_len. helix_is_inactive_rank, helix_position_offsets = [], [] helix_owned_new_tokens = [] _has_cp_helix = self.mapping.has_cp_helix() if _has_cp_helix and kv_cache_manager is not None: _helix_phys = kv_cache_manager.tokens_per_block - _helix_ledger = _helix_phys * self.mapping.cp_size - _helix_rank_off = self.mapping.cp_rank * _helix_phys + _helix_cp_size = self.mapping.cp_size + _helix_cp_rank = self.mapping.cp_rank def _helix_local_len_host(global_len: int) -> int: - full, rem = divmod(global_len, _helix_ledger) - return full * _helix_phys + min(max(rem - _helix_rank_off, 0), - _helix_phys) + return helix_local_len(global_len, _helix_phys, _helix_cp_size, + _helix_cp_rank) def _helix_pack_extend(request, group: int) -> int: # A helix gen worker's token list is the rank-LOCAL diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 389f630861e5..4cf992f7a5c0 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -786,3 +786,49 @@ def torch_multi_arange( seq = seq.repeat_interleave(seq_repeats, output_size=output_length_arg) seq = seq.cumsum(0, dtype=ends.dtype) return seq + + +# --------------------------------------------------------------------------- +# Helix CP round-robin ledger +# --------------------------------------------------------------------------- +# THE rule, stated once. Under helix context parallelism the KV ledger is a +# round-robin over CP ranks at page granularity: ledger page b lives on rank +# b % cp_size, so one "ledger block" spans tokens_per_block * cp_size global +# token positions and contributes exactly tokens_per_block of them to each +# rank. For a global prefix of length ``global_len`` this rank therefore owns +# +# full, rem = divmod(global_len, tokens_per_block * cp_size) +# full * tokens_per_block + clamp(rem - cp_rank * tokens_per_block, +# 0, tokens_per_block) +# +# tokens: every complete ledger block gives it a whole page, and the trailing +# partial block gives it however much of its own page the remainder reaches. +# Summed over all ranks this is exactly ``global_len``. +# +# Both forms below are that expression and nothing else. Keep them that way: +# a drift between the scalar host-side packing and the tensor form used to +# derive the device write slots puts KV on the wrong rank, and only shows up +# on groups that straddle a page boundary. + + +def helix_local_len(global_len: int, tokens_per_block: int, cp_size: int, + cp_rank: int) -> int: + """Scalar form: tokens of the first ``global_len`` owned by ``cp_rank``.""" + ledger = tokens_per_block * cp_size + full, rem = divmod(global_len, ledger) + return full * tokens_per_block + min( + max(rem - cp_rank * tokens_per_block, 0), tokens_per_block) + + +def helix_local_len_tensor(global_lens: torch.Tensor, tokens_per_block: int, + cp_size: int, cp_rank: int) -> torch.Tensor: + """Tensor form of :func:`helix_local_len`, applied elementwise. + + Kept allocation-minimal: this runs on the CUDA-graph capture path, and the + ``clamp_`` is in place on the temporary the subtraction just produced. + """ + ledger = tokens_per_block * cp_size + full = torch.div(global_lens, ledger, rounding_mode='floor') + rem = global_lens - full * ledger + return full * tokens_per_block + (rem - cp_rank * tokens_per_block).clamp_( + 0, tokens_per_block) From 3083a1a4670ad4b4d4f2ba9924f6fc9f7bceea68 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:48:34 +0000 Subject: [PATCH 28/33] [None][fix] reject max_concurrency under Kimi K3 helix speculation The existing check rejects the acceptance-rate gate because dynamically disabling speculation drops in-flight helix requests into the plain generation loop, whose position formula counts iterations rather than committed tokens. max_concurrency does exactly the same thing by another route: py_executor re-evaluates Drafter.should_use_spec_decode every scheduling iteration and clears enable_spec_decode once the active batch exceeds the cap, so a request that has already accepted draft tokens gets a position and a CP owner rank derived from a counter that is behind by the accepted count -- a wrong RoPE position and, across a ledger page boundary, a KV write to the wrong rank, silently. Fail at build time instead. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 563a086b4685..ce6b23b06ee1 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -2176,6 +2176,26 @@ def _setup_helix_mappings( "speculation mid-flight leaves helix requests on a " "single-token position formula." ) + # max_concurrency is the same trip wire by another name: the + # drafter re-evaluates should_use_spec_decode on every scheduling + # iteration and flips enable_spec_decode off as soon as the active + # batch exceeds the cap. In-flight helix requests then take the + # plain generation loop, whose position formula counts ITERATIONS + # (total_input_len_cp + py_decoding_iter - 1) rather than + # committed tokens, so it is stale by however many draft tokens + # were accepted -- a wrong RoPE position, and across a ledger page + # boundary a KV write to the wrong CP rank. Mirror the drafter's + # own "unset" test (Drafter.should_use_spec_decode returns True + # when max_concurrency is None) so an unset value is not rejected. + if spec_config.max_concurrency is not None: + raise ValueError( + "Kimi K3 helix does not support the speculation " + "concurrency cutoff (max_concurrency): disabling " + "speculation above the cap leaves in-flight helix " + "requests on a position formula that assumes one " + "committed token per iteration, which accepted draft " + "tokens break." + ) cp = model_config.mapping.cp_size repurposed_tp = model_config.mapping.tp_size * cp if cfg.num_attention_heads % repurposed_tp != 0: From 68a85f1af2e1c1faac75b0c048ec98d3707a10cf Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:51:03 +0000 Subject: [PATCH 29/33] [None][fix] reject draft_len_schedule under Kimi K3 helix speculation draft_len_schedule is the user-facing alternative to max_concurrency -- the two are mutually exclusive in llm_args, with max_concurrency translated into a schedule behind _translated_from_max_concurrency -- and it disables speculation by a route that never reaches should_use_spec_decode: py_executor clears use_spec_decode directly once the schedule yields a draft length of 0 for the active batch size. Guarding only max_concurrency therefore left the same stale-position hazard reachable. Skip the synthesized schedule so a config that set only max_concurrency still raises the message naming that field rather than one the user never wrote. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index ce6b23b06ee1..1bdd8f7acadf 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -2196,6 +2196,27 @@ def _setup_helix_mappings( "committed token per iteration, which accepted draft " "tokens break." ) + # draft_len_schedule is the user-facing alternative to + # max_concurrency (llm_args rejects setting both) and reaches the + # same end by a route that does not go through + # should_use_spec_decode at all: py_executor turns speculation off + # directly once the schedule yields draft_len 0 for the active + # batch size. Guarding only max_concurrency would leave this door + # open. Skip the schedule that llm_args synthesized from + # max_concurrency, so a config that set only that field raises the + # message above naming the field the user actually wrote. + if ( + spec_config.draft_len_schedule is not None + and not spec_config._translated_from_max_concurrency + ): + raise ValueError( + "Kimi K3 helix does not support the dynamic draft-length " + "schedule (draft_len_schedule): a batch size past the " + "last entry drops the draft length to 0 and turns " + "speculation off mid-run, leaving in-flight helix " + "requests on a position formula that assumes one " + "committed token per iteration." + ) cp = model_config.mapping.cp_size repurposed_tp = model_config.mapping.tp_size * cp if cfg.num_attention_heads % repurposed_tp != 0: From 15abfa29e4da1effa3df307875b9d98f50c61140 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:19:56 +0000 Subject: [PATCH 30/33] [None][fix] give the attention metadata stubs the helix fields the checks read FallbackFmha._is_supported now evaluates the Helix verify-group reject before anything else, and CuteDslMlaFmha reads _helix_spec_tokens_valid when seq_len_q > 1, so the SimpleNamespace metadata stubs in test_attention_op_sync and test_fmha_page_index raised AttributeError instead of exercising the contract they assert. Give them the values the real TrtllmAttentionMetadata carries off the helix path. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tests/unittest/_torch/attention/test_attention_op_sync.py | 7 +++++-- tests/unittest/_torch/attention/test_fmha_page_index.py | 4 ++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/attention/test_attention_op_sync.py b/tests/unittest/_torch/attention/test_attention_op_sync.py index ccf11544be53..5ca25019f86c 100644 --- a/tests/unittest/_torch/attention/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention/test_attention_op_sync.py @@ -685,7 +685,10 @@ def test_no_sequence_kwargs_at_thop_attention_boundary(): def test_fallback_support_matches_thop_kv_update_contract(is_cross, update_kv_cache, expected): """Do not dispatch requests that the native attention op rejects.""" fmha = object.__new__(FallbackFmha) - metadata = SimpleNamespace(is_cross=is_cross) + # ``helix_position_offsets`` short-circuits the Helix verify-group check + # that runs first in ``_is_supported``; None is what the real metadata + # carries off the helix path. + metadata = SimpleNamespace(is_cross=is_cross, helix_position_offsets=None) forward_args = AttentionForwardArgs(update_kv_cache=update_kv_cache) assert fmha.is_supported(None, None, None, metadata, forward_args) is expected @@ -694,7 +697,7 @@ def test_fallback_support_matches_thop_kv_update_contract(is_cross, update_kv_ca def test_fallback_rejects_raw_fp8_input(): """Do not dispatch raw FP8 QKV to the native attention op.""" fmha = object.__new__(FallbackFmha) - metadata = SimpleNamespace(is_cross=False) + metadata = SimpleNamespace(is_cross=False, helix_position_offsets=None) forward_args = AttentionForwardArgs(update_kv_cache=True) q = torch.empty((1, 128), dtype=torch.float8_e4m3fn) diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 9460fc7a4da8..9edb2a2e087c 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -287,6 +287,10 @@ def _cute_dsl_mla_helix_support( is_spec_dec_tree=False, is_spec_dec_dynamic_tree=False, helix_position_offsets=torch.zeros(batch_size, dtype=torch.int32), + # Multi-token decode under helix is admitted only when the speculative + # verify-group buffers are armed; unarmed is what this contract test + # exercises. + _helix_spec_tokens_valid=False, kv_cache_manager=SimpleNamespace( get_buffers=lambda _layer_idx: torch.empty(0, dtype=torch.bfloat16) ), From fc706b485d2d5c574c2105d2acf52b088fda70aa Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:20:15 +0000 Subject: [PATCH 31/33] [None][fix] give the QSA model_config stub the fields the helix manager check reads get_kv_cache_manager_cls now rejects helix with speculative decoding on a V1-family hybrid manager, reading model_config.mapping and model_config.spec_config. Both are declared fields with defaults on the real ModelConfig, but the SimpleNamespace stub omitted them, so the V2 routing test raised AttributeError once it got past the QSA V1 rejection. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index 9f889d5ab4b8..2907a2a7e744 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -244,6 +244,10 @@ def test_qsa_hybrid_routes_to_sparse_v2_cache_manager(monkeypatch) -> None: pretrained_config=SimpleNamespace(), sparse_attention_config=QSASparseAttentionConfig(), get_num_mamba_layers=lambda: 1, + # Real ModelConfig declares both with defaults; the helix x + # speculation manager check reads them. + mapping=None, + spec_config=None, ) kv_cache_config = KvCacheConfig(use_kv_cache_manager_v2=True) @@ -260,6 +264,10 @@ def test_qsa_hybrid_rejects_kv_cache_manager_v1(monkeypatch) -> None: pretrained_config=SimpleNamespace(), sparse_attention_config=QSASparseAttentionConfig(), get_num_mamba_layers=lambda: 1, + # Real ModelConfig declares both with defaults; the helix x + # speculation manager check reads them. + mapping=None, + spec_config=None, ) with pytest.raises(ValueError, match="requires use_kv_cache_manager_v2=True"): From 0a5c71a680922f3397de316b3ca0435041c0e4ec Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:20:49 +0000 Subject: [PATCH 32/33] [None][fix] make the helix spec_config stub a concrete non-DSpark mode The helix precondition no longer rejects every spec_config: DSpark is now supported, and the check reads spec_dec_mode.is_dspark() plus decoding_type for its message. An empty SimpleNamespace therefore raised AttributeError instead of the ValueError the test asserts. Give it a non-DSpark mode so the case still covers what it says it covers. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/modeling/test_kimi_linear_helix_mappings.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_kimi_linear_helix_mappings.py b/tests/unittest/_torch/modeling/test_kimi_linear_helix_mappings.py index 6a586dcd0b20..250e469bba80 100644 --- a/tests/unittest/_torch/modeling/test_kimi_linear_helix_mappings.py +++ b/tests/unittest/_torch/modeling/test_kimi_linear_helix_mappings.py @@ -84,12 +84,16 @@ def test_setup_helix_mappings_precondition_validation(): with pytest.raises(ValueError, match="enable_attention_dp"): setup(obj, model_config, cfg, None) - # spec_config not None raises + # a non-DSpark spec_config raises: helix now supports DSpark, so the + # rejection is mode-specific rather than "any spec_config". mapping = _make_helix_mapping(tp_size=4, cp_size=2, enable_attention_dp=False) model_config = _make_model_config(mapping) cfg = _make_cfg(num_attention_heads=96, kda_num_heads=96) obj = _FakeSelf() - spec_config = SimpleNamespace() + spec_config = SimpleNamespace( + spec_dec_mode=SimpleNamespace(is_dspark=lambda: False), + decoding_type="MTP", + ) with patch.object(type(mapping), "has_cp_helix", return_value=True): with pytest.raises(ValueError, match="speculative"): setup(obj, model_config, cfg, spec_config) From 99b248844f887db1f4e463ac3ba315937cdd650a Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:21:19 +0000 Subject: [PATCH 33/33] [None][fix] index the MLA tactic elements past the new kv_bounds cache key The CuTe DSL MLA kernel-cache key gained a trailing 'kv_bounds is not None' element so the helix per-token-bounds variant cannot collide with the plain one. The autotune test read is_persistent as key[-1] and split_kv as key[-2], so it was asserting on that new flag instead and saw only {False}. Shift both indices and record the layout in the comment. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../unittest/_torch/attention/test_attention_mla.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index fbcf3e856f2d..7fb79f29925e 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -733,15 +733,18 @@ def run_once(batch_size: int = 64) -> None: kernel_keys = list(CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache) assert kernel_keys, "tuning-mode pass compiled no CuTe DSL MLA kernels" - # Tactic layout: unique_id + (out_dtype, mma_qk, mma_pv, split_kv, - # is_persistent); both tactic elements chosen by the tuner must have - # been exercised during profiling. - persistent_variants = {key[-1] for key in kernel_keys} + # Kernel-cache key layout: unique_id + (out_dtype, mma_qk, mma_pv, + # split_kv, is_persistent, has_kv_bounds). The trailing has_kv_bounds + # separates the helix per-token-bounds kernel variant from the plain one + # and is not a tuner tactic, so index the two tactic elements from the + # end past it. Both tactic elements chosen by the tuner must have been + # exercised during profiling. + persistent_variants = {key[-2] for key in kernel_keys} assert persistent_variants == { True, False }, (f"expected both is_persistent tactic candidates to be profiled, " f"got {persistent_variants}") - split_kv_variants = {key[-2] for key in kernel_keys} + split_kv_variants = {key[-3] for key in kernel_keys} assert split_kv_variants, "no split_kv tactic variant was profiled" # Serving-mode pass: tuned tactics must be reused as-is -- any new