diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index 304b255a11f6..11f85ce0edac 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); @@ -636,13 +702,33 @@ 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) { - 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..18c17be8bd85 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,25 @@ 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()); + // 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 auto stream = at::cuda::getCurrentCUDAStream(); tensorrt_llm::kernels::launchHelixAllToAll(params, allowVariableField1, stream); @@ -260,7 +283,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(); diff --git a/tensorrt_llm/_torch/attention/attention.py b/tensorrt_llm/_torch/attention/attention.py index 94344ff8247b..4685edafd65c 100644 --- a/tensorrt_llm/_torch/attention/attention.py +++ b/tensorrt_llm/_torch/attention/attention.py @@ -193,6 +193,89 @@ 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, + 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, 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, + 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 +293,25 @@ 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. + _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) - 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 +322,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 +350,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/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 51e730ed754f..ac97de1d550c 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -301,7 +301,9 @@ 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." softmax_stats = fwd.softmax_stats_tensor if softmax_stats is None: @@ -320,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: @@ -505,6 +522,13 @@ 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) + 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..5631950e9e06 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -81,6 +81,22 @@ 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 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 + # 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: return False if forward_args.attention_mask == CustomAttentionMask.CUSTOM: 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..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 @@ -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 @@ -1895,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 @@ -1927,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) @@ -2046,9 +2076,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 +2172,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..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, @@ -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,82 @@ 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). 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. + """ + 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: + """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. + + ``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 + 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_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_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 + # 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 + + num_gen_seqs].copy_(last_bounds) + def _bind_runtime_views( self, *, @@ -752,9 +877,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" - 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] + # 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[gen] = (cached_token_lens[gen] + + self.helix_owned_new_tokens_cpu[:num_gen]) + else: + 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 @@ -2412,6 +2563,11 @@ def mla_rope_generation( helix_tensor_params = [ metadata.helix_position_offsets, metadata.helix_is_inactive_rank ] + 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. + 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..e331d9ebcf24 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) + # 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 + # 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/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 3b7f9b9ed0ec..1ff1ab5e90fc 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,13 +11764,17 @@ 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). + + 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( @@ -11760,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( @@ -11798,6 +11831,7 @@ def _( output_scale: float, max_batch_size: int, softmax_stats: Optional[torch.Tensor], + kv_bounds: Optional[torch.Tensor], ) -> None: return None @@ -11823,8 +11857,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 +11898,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 +11936,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..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 @@ -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, @@ -4192,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, @@ -4307,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, @@ -4444,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 9e2477edf634..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 @@ -309,6 +309,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, @@ -326,6 +327,7 @@ def __call__( workspace, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale, output_scale, @@ -346,6 +348,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, @@ -363,6 +366,7 @@ def run_with_softmax_stats( workspace, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale, output_scale, @@ -383,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, @@ -390,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 @@ -844,6 +855,7 @@ class SplitKVKernelSharedStorage: acc_lse, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale_log2, output_scale, @@ -881,6 +893,7 @@ class SplitKVKernelSharedStorage: acc_lse, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale_log2, output_scale, @@ -915,6 +928,7 @@ class SplitKVKernelSharedStorage: split_kv, cache_seqs, block_split_kvs, + kv_bounds, ) else: reduction_kernel = self.reduction_kernel( @@ -994,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, @@ -1414,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, @@ -1516,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. @@ -1586,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], @@ -2424,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 @@ -2751,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, @@ -2797,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, @@ -4151,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, @@ -4266,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, @@ -4403,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, 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/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index b87a6f1636dc..1bdd8f7acadf 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -2149,11 +2149,74 @@ 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." + ) + # 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." + ) + # 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: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 820a4bcac0d7..cb6466258562 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( + "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..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. @@ -3264,8 +3269,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. Assumes one new token per step - (draft-token modes are rejected under helix). + 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 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..ad417618950b 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 + 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): @@ -536,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 @@ -2352,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 @@ -3142,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 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 93a0ed55bd36..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 @@ -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( + 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); 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_cp_size = self.mapping.cp_size + _helix_cp_rank = self.mapping.cp_rank + + def _helix_local_len_host(global_len: int) -> int: + 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 + # 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: @@ -5005,6 +5093,21 @@ def append_cross_attention_state(request: LlmRequest, 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) @@ -5069,9 +5172,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 @@ -5201,6 +5302,12 @@ def append_cross_attention_state(request: LlmRequest, 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. 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): @@ -5618,9 +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_spec_active else None), ) if not attn_metadata.is_cuda_graph: 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) 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") 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 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) ), 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) 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) 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"):