[None][feat] Helix speculative verify groups: fp8 + fp4 MLA and DSpark

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<uintptr_t>(fieldInfo.dataPtr) % alignment == 0
+        && static_cast<uintptr_t>(fieldInfo.stride) % alignment == 0
+        && static_cast<uintptr_t>(getFieldSize(fieldInfo)) % alignment == 0;
+}
+
+template <bool NEEDS_WARP_SYNC>
 __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 <bool ALLOW_VARIABLE_FIELD1>
+template <bool ALLOW_VARIABLE_FIELD1, bool USE_BULK_FIELD1>
 __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<kCopySize>(reinterpret_cast<int*>(shmemBase + field0Size + offset),
+                    reinterpret_cast<int const*>(getPtr(fieldInfo[1], dataIndex) + offset), true);
+            }
+            cp_async_commit_group();
+        }
     }
     else
     {
-        ldgsts<8>(reinterpret_cast<int*>(shmemBase + totalSize),
+        ldgsts<8>(reinterpret_cast<int*>(shmemBase + field0Size),
             reinterpret_cast<int const*>(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 <bool ALLOW_VARIABLE_FIELD1>
+template <bool ALLOW_VARIABLE_FIELD1, bool USE_BULK_FIELD1>
 __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<float2 const*>(shmemBase + offset + fieldOffset);
+                auto* dstPtr = reinterpret_cast<float2*>(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 <bool ALLOW_VARIABLE_FIELD1>
+template <bool ALLOW_VARIABLE_FIELD1, bool USE_BULK_FIELD1>
 __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<ALLOW_VARIABLE_FIELD1>(
+            int loadedSize = g2sAllFields<ALLOW_VARIABLE_FIELD1, USE_BULK_FIELD1>(
                 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<ALLOW_VARIABLE_FIELD1 && !USE_BULK_FIELD1>(&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<int>(sizeof(int4)); off < field0Size;
+                     off += WARP_SIZE * static_cast<int>(sizeof(int4)))
+                {
+                    *reinterpret_cast<int4*>(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<float2*>(shmem + field0Size);
+                    int const statsCount = getFieldSize(params.sendFields[1]) / static_cast<int>(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<ALLOW_VARIABLE_FIELD1>(params.recvFields, dataIndex, shmem, laneId);
+            s2gAllFields<ALLOW_VARIABLE_FIELD1, USE_BULK_FIELD1>(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 <bool ALLOW_VARIABLE_FIELD1>
+template <bool ALLOW_VARIABLE_FIELD1, bool USE_BULK_FIELD1>
 std::tuple<int, int, int> computeChannelAndGroupCount(int cpSize, HelixFieldInfo const* fields)
 {
     static std::unordered_map<std::tuple<int, int, int>, std::tuple<int, int, int>, hash_cache_key> cache;
@@ -535,7 +601,7 @@ std::tuple<int, int, int> computeChannelAndGroupCount(int cpSize, HelixFieldInfo
     // Set shared memory attribute if needed
     if (totalDynamicShmemSize > 48 * 1024)
     {
-        TLLM_CUDA_CHECK(cudaFuncSetAttribute(helixAllToAllKernel<ALLOW_VARIABLE_FIELD1>,
+        TLLM_CUDA_CHECK(cudaFuncSetAttribute(helixAllToAllKernel<ALLOW_VARIABLE_FIELD1, USE_BULK_FIELD1>,
             cudaFuncAttributeMaxDynamicSharedMemorySize, totalDynamicShmemSize));
     }
 
@@ -559,14 +625,14 @@ std::tuple<int, int, int> computeChannelAndGroupCount(int cpSize, HelixFieldInfo
 // Host Launch Function
 // ============================================================================
 
-template <bool ALLOW_VARIABLE_FIELD1>
+template <bool ALLOW_VARIABLE_FIELD1, bool USE_BULK_FIELD1>
 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<ALLOW_VARIABLE_FIELD1>(params.cpSize, params.sendFields);
+        = computeChannelAndGroupCount<ALLOW_VARIABLE_FIELD1, USE_BULK_FIELD1>(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<ALLOW_VARIABLE_FIELD1>;
+    auto* kernel_instance = &helixAllToAllKernel<ALLOW_VARIABLE_FIELD1, USE_BULK_FIELD1>;
     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<true>(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<true, true>(params, stream);
+        }
+        else
+        {
+            launchHelixAllToAllImpl<true, false>(params, stream);
+        }
     }
     else
     {
-        launchHelixAllToAllImpl<false>(params, stream);
+        launchHelixAllToAllImpl<false, false>(params, stream);
     }
 }
 

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;
 };
 
 // ============================================================================

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<T>::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<T*>(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<size_t>(global_token_idx) * (c_k + ROPE_DIM);
 
                     {
@@ -1961,8 +1976,8 @@ void invokeMLARopeGeneration(MlaParams<T>& 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 <typename T, typename KVCacheBuffer>

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 <typename T, typename KVCacheBuffer>