From ba1f808def6e662cb4376af7eaa9d9a887fb9924 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 15 Sep 2026 03:51:27 +0000 Subject: [PATCH 01/21] [None][feat] Rubin kernels & attention: DSV4/DSA, Helix, CuteDSL GEMM Port the SM107 kernel work from the internal Rubin branch onto main. Three related groups that share reviewers: DSV4/DSA - mhcKernels: fused TF32 pmap GEMM and a TMA-fed fused HC kernel (new `kUseTma` template axis on `mhcBigFuseKernel`). - mlaKernels: `helix_local_slots` per-token KV write slots for speculative verify groups, and `q_rope_applied` so a caller that already produced the rotated FP8 Q keeps it while context preprocessing still rotates/caches K. - attentionOp: size the unfused-context buffers from the padded token counts exactly as `enqueueContext` carves them; covered by two new cases in `attentionWorkspaceTest`. - fmhaKernels: clamp `maxNumCtasPerSeqKv` around `FmhaAutoTuner::selectKernel` so a CGA-smem-reduction candidate never exceeds clusterDimX * ctasKv == 16. - DSv4 indexer: overlapped `precompute_aux` start event and a `TRTLLM_DISABLE_DSA_FUSED_INDEXER_Q` escape hatch. Helix - helixAllToAll: optional sender-side `zeroKvMask` so rows this rank owns no KV for contribute a no-op to the combine. - `_ipc_utils`: serialize `cudaIpcMemHandle_t` across cuda-python builds that omit `.reserved`. CuteDSL GEMM - Cache the MXFP8 GEMM alpha scalar per device so the steady-state path, and CUDA graph replay, no longer launch a scalar fill. - Kimi K3 MXFP8 fine-M tuning buckets; KDA eqlen pad trigger corrected to CHUNKS_PER_BLOCK * BT and the varlen NT >= 4 restriction lifted. - `Linear`: pre-quantized FP8 activation input path. - `Linear`/`GatedMLP`: share one capability predicate for the Blackwell NVFP4 GEMM + SwiGLU kernel so weight transformation and forward dispatch cannot disagree, plus the fused SwiGLU + FP8 quant hand-off to down_proj. No cubin, export-header, or prebuilt-library payload is included; this is a source-only port against main's existing trtllmGen FMHA artifacts. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 19 +- cpp/tensorrt_llm/common/envUtils.h | 15 +- cpp/tensorrt_llm/kernels/helixAllToAll.cu | 134 ++- cpp/tensorrt_llm/kernels/helixAllToAll.h | 8 + .../mhcKernels/fused_tf32_pmap_gemm.cuh | 299 ++--- .../kernels/mhcKernels/mhcFusedHcKernel.cu | 105 +- .../kernels/mhcKernels/mhcKernels.cu | 219 ++-- .../kernels/mhcKernels/mhc_fused_fma.cuh | 34 +- cpp/tensorrt_llm/kernels/mlaKernels.cu | 131 ++- cpp/tensorrt_llm/kernels/mlaKernels.h | 10 + .../trtllmGenKernels/fmha/fmhaKernels.h | 49 +- cpp/tensorrt_llm/thop/alltoallOp.cpp | 25 +- cpp/tensorrt_llm/thop/dsv3RopeOp.cpp | 31 +- .../common/attentionWorkspaceTest.cpp | 66 ++ tensorrt_llm/_ipc_utils.py | 40 +- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 15 +- .../backends/sparse/deepseek_v4/indexer.py | 55 +- .../attention/backends/sparse/params.py | 23 + .../_torch/custom_ops/cpp_custom_ops.py | 7 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 54 +- .../custom_ops/cute_dsl_kimi_k3_custom_ops.py | 25 +- .../blackwell/attention/mla/mla_helpers.py | 57 +- ...ntiguous_gather_grouped_gemm_act_fusion.py | 148 ++- .../cute_dsl_kernels/blackwell/utils.py | 7 + .../dense_blockscaled_gemm_persistent.py | 2 + .../rubin/dsv4_qb_fusion/kernel.py | 1016 ++++++++++------- tensorrt_llm/_torch/cute_dsl_utils.py | 37 +- tensorrt_llm/_torch/distributed/ops.py | 10 +- .../_torch/models/modeling_deepseekv3.py | 33 +- .../_torch/models/modeling_deepseekv4.py | 9 +- tensorrt_llm/_torch/modules/gated_mlp.py | 45 +- tensorrt_llm/_torch/modules/linear.py | 55 +- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 2 +- .../test_helix_postprocess.py | 39 +- .../attention/multi_gpu/test_mla_helix.py | 11 +- .../deepseek_v4/test_deepseek_v4_o_proj.py | 49 +- .../sparse/dsa/test_indexer_gvr_prior.py | 261 +++++ .../sparse/dsa/test_metadata_topk_init.py | 138 +++ .../sparse/test_sparse_mla_forward.py | 39 +- .../test_indexer_workspace_reserve.py | 30 + .../modeling/test_modeling_deepseekv4.py | 3 + .../unittest/_torch/modules/test_gated_mlp.py | 114 ++ .../_torch/thop/parallel/test_cute_dsl_moe.py | 238 +++- .../parallel/test_fp8_block_scale_gemm.py | 48 +- .../test_kimi_k3_fp8_weight_read_linear.py | 247 ++++ 45 files changed, 3071 insertions(+), 931 deletions(-) create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py create mode 100644 tests/unittest/_torch/executor/test_indexer_workspace_reserve.py create mode 100644 tests/unittest/_torch/modules/test_gated_mlp.py create mode 100644 tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 85ee5ccf8970..f3a8adfe64fd 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -800,10 +800,13 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 auto const batch_size = static_cast(max_num_seq); auto const kv_seq_length = (isCrossAttention() ? cross_kv_length : input_seq_length); - // The unfused-MHA buffers below must upper-bound the enqueueContext carve, which sizes them by - // batch_size * input_seq_length (not num_tokens): with padding removal the actual token count can be - // smaller than batch_size * max(context q length), so sizing by max_num_tokens underestimates. - size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_length * kv_seq_length; + // Unfused context attention operates on padded [batch, sequence] tensors, + // even when the input QKV is packed. Size those buffers from the padded + // token counts exactly as enqueueContext does; max_num_tokens remains the + // packed count used by the fused paths below. + size_t const padded_num_tokens = batch_size * static_cast(input_seq_length); + size_t const padded_kv_tokens = batch_size * static_cast(kv_seq_length); + size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * padded_num_tokens * kv_seq_length; size_t const cu_seqlens_size = sizeof(int) * (batch_size + 1); size_t const rotary_inv_freq_size = sizeof(float) * batch_size * mRotaryEmbeddingDim / 2; @@ -823,7 +826,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 size_t const v_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * kv_seq_length * local_hidden_units_kv; size_t const qk_buf_size = mEnableContextFMHA ? 0 : size * batch_size * mNumHeads * input_seq_length * kv_seq_length; - size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_length * local_hidden_units_qo; + size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : size * padded_num_tokens * local_hidden_units_qo; size_t const qk_buf_float_size = mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_length * kv_seq_length; int dim_q_per_head = (mMLAParams.qk_rope_head_dim + mMLAParams.qk_nope_head_dim); @@ -899,8 +902,8 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 ? sizeof(float) * tc::divUp(local_hidden_units_kv, std::max(1, mSageAttnNumEltsPerBlkV)) : 0; - size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * input_seq_length; - size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * cross_kv_length; + size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * padded_num_tokens; + size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * padded_kv_tokens; // Each token holds (batch_idx, token_idx_in_seq) int2. size_t const tokens_info_size = sizeof(int2) * max_num_tokens; size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0; @@ -1886,6 +1889,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea preprocessingParams.is_last_chunk = !mAttentionChunkSize.has_value() || (params.input_seq_length == params.max_past_kv_length); + if (!(mIsMLAEnabled && params.mla_param != nullptr && params.mla_param->q_rope_applied)) { std::string const beforeRopeStr = "ctx attention before RoPE at layer " + std::to_string(mLayerIdx); TLLM_CHECK_DEBUG_WITH_INFO(tensorrt_llm::runtime::utils::tensorHasInvalid(params.num_tokens, @@ -1991,6 +1995,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea invokeQKVPreprocessing(preprocessingParams, stream); } sync_check_cuda_error(stream); + if (!(mIsMLAEnabled && params.mla_param != nullptr && params.mla_param->q_rope_applied)) { std::string const afterRopeStr = "ctx attention after RoPE at layer " + std::to_string(mLayerIdx); TLLM_CHECK_DEBUG_WITH_INFO(tensorrt_llm::runtime::utils::tensorHasInvalid(params.num_tokens, diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index e6d354d05745..a81e1b362f82 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -72,10 +72,10 @@ bool getEnvUseFineGrainedSync(); void setFineGrainedSyncDisabledOverride(bool disabled); template -inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 grid, dim3 block, size_t dynamicShmSize, - cudaStream_t stream, Args&&... args) +inline void launchWithPdl(char const* name, bool enablePdl, KernelFn kernelFn, dim3 grid, dim3 block, + size_t dynamicShmSize, cudaStream_t stream, Args&&... args) { - TLLM_LOG_DEBUG("Enable PDL in %s", name); + TLLM_LOG_DEBUG("PDL in %s: %s", name, enablePdl ? "enabled" : "disabled"); cudaLaunchConfig_t kernelConfig; kernelConfig.gridDim = grid; kernelConfig.blockDim = block; @@ -84,13 +84,20 @@ inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 g cudaLaunchAttribute attrs[1]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl; kernelConfig.attrs = attrs; kernelConfig.numAttrs = 1; TLLM_CUDA_CHECK(cudaLaunchKernelEx(&kernelConfig, kernelFn, std::forward(args)...)); } +template +inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 grid, dim3 block, size_t dynamicShmSize, + cudaStream_t stream, Args&&... args) +{ + launchWithPdl(name, getEnvEnablePDL(), kernelFn, grid, block, dynamicShmSize, stream, std::forward(args)...); +} + bool getEnvUseUCXKvCache(); bool getEnvUseMPIKvCache(); diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index 304b255a11f6..f9eb6aee3009 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.cu +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -85,10 +85,22 @@ __host__ __device__ inline uint8_t* getPtr(HelixFieldInfo const& fieldInfo, int return fieldInfo.dataPtr + blockIdx * fieldInfo.stride; } +bool isFieldAlignedForCopy(HelixFieldInfo const& fieldInfo, uintptr_t alignment) +{ + return reinterpret_cast(fieldInfo.dataPtr) % alignment == 0 + && static_cast(fieldInfo.stride) % alignment == 0 + && static_cast(getFieldSize(fieldInfo)) % alignment == 0; +} + +template __device__ __forceinline__ void waitG2sAllFields(uint64_t* smemBar, uint32_t* phaseParity) { cp_async_wait_group<0>(); smemBarWait(smemBar, phaseParity); + if constexpr (NEEDS_WARP_SYNC) + { + __syncwarp(); + } } // Align size to 128 bytes @@ -113,31 +125,44 @@ __device__ __forceinline__ void g2sField( } } -template +template __device__ __forceinline__ int g2sAllFields( HelixFieldInfo const* fieldInfo, int dataIndex, uint8_t* shmemBase, uint64_t* smemBar, int laneId) { - int totalSize = 0; - // Load field 0 (variable size half) g2sField(fieldInfo[0], dataIndex, shmemBase, 0, smemBar, laneId); - int field0Size = getFieldSize(fieldInfo[0]); - totalSize += field0Size; + int const field0Size = getFieldSize(fieldInfo[0]); + int bulkCopySize = field0Size; - // Load field 1 (single float2) + // Load field 1 (one or more float2 values). if constexpr (ALLOW_VARIABLE_FIELD1) { - g2sField(fieldInfo[1], dataIndex, shmemBase, totalSize, smemBar, laneId); - totalSize += getFieldSize(fieldInfo[1]); + if constexpr (USE_BULK_FIELD1) + { + g2sField(fieldInfo[1], dataIndex, shmemBase, field0Size, smemBar, laneId); + bulkCopySize += getFieldSize(fieldInfo[1]); + } + else + { + constexpr int kCopySize = sizeof(float2); + int const field1Size = getFieldSize(fieldInfo[1]); + for (int offset = laneId * kCopySize; offset < field1Size; offset += WARP_SIZE * kCopySize) + { + ldgsts(reinterpret_cast(shmemBase + field0Size + offset), + reinterpret_cast(getPtr(fieldInfo[1], dataIndex) + offset), true); + } + cp_async_commit_group(); + } } else { - ldgsts<8>(reinterpret_cast(shmemBase + totalSize), + ldgsts<8>(reinterpret_cast(shmemBase + field0Size), reinterpret_cast(getPtr(fieldInfo[1], dataIndex)), laneId == 0); cp_async_commit_group(); } - return totalSize; + // Only bytes copied by cp.async.bulk participate in the mbarrier transaction. + return bulkCopySize; } // ============================================================================ @@ -156,7 +181,7 @@ __device__ __forceinline__ void s2gField( } } -template +template __device__ __forceinline__ void s2gAllFields( HelixFieldInfo const* fieldInfo, int dataIndex, uint8_t* shmemBase, int laneId) { @@ -167,11 +192,25 @@ __device__ __forceinline__ void s2gAllFields( int field0Size = getFieldSize(fieldInfo[0]); offset += field0Size; - // Store field 1 (single float2) + // Store field 1 (one or more float2 values). if constexpr (ALLOW_VARIABLE_FIELD1) { - s2gField(fieldInfo[1], dataIndex, shmemBase, offset, laneId); - offset += getFieldSize(fieldInfo[1]); + if constexpr (USE_BULK_FIELD1) + { + s2gField(fieldInfo[1], dataIndex, shmemBase, offset, laneId); + } + else + { + constexpr int kCopySize = sizeof(float2); + int const field1Size = getFieldSize(fieldInfo[1]); + for (int fieldOffset = laneId * kCopySize; fieldOffset < field1Size; fieldOffset += WARP_SIZE * kCopySize) + { + auto const* srcPtr = reinterpret_cast(shmemBase + offset + fieldOffset); + auto* dstPtr = reinterpret_cast(getPtr(fieldInfo[1], dataIndex) + fieldOffset); + *dstPtr = *srcPtr; + } + __syncwarp(); + } } else { @@ -294,7 +333,7 @@ __host__ __device__ __forceinline__ int computeProtoTransferSize(HelixFieldInfo // Main All-to-All Kernel // ============================================================================ -template +template __global__ void helixAllToAllKernel(HelixAllToAllParams params) { extern __shared__ uint8_t allWarpShmem[]; @@ -371,7 +410,7 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) int dataIndex = entryIdx * params.cpSize + peerRank; // Load data from global to shared, then arrive on barrier - int loadedSize = g2sAllFields( + int loadedSize = g2sAllFields( params.sendFields, dataIndex, shmem, &allWarpSmemBar[group], laneId); uint64_t arriveState = mbarrier_arrive_expect_tx(&allWarpSmemBar[group], laneId == 0 ? loadedSize : 0); @@ -394,10 +433,37 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) } // wait for data to be loaded into shared memory - waitG2sAllFields(&allWarpSmemBar[group], &phaseParity); + waitG2sAllFields(&allWarpSmemBar[group], &phaseParity); // note: we don't need to pack anything, fields are already packed in // shared memory + // Zero-local-KV sanitization: the entry is in shared memory and + // not yet packed, so overwriting it costs no global traffic. + if (params.zeroKvMask != nullptr && params.zeroKvMask[entryIdx / params.zeroKvMaskDivisor] != 0) + { + // field0Size is a multiple of 16 (the op checks it), so the + // int4 store never runs off the end of the field. + int const field0Size = getFieldSize(params.sendFields[0]); + for (int off = laneId * static_cast(sizeof(int4)); off < field0Size; + off += WARP_SIZE * static_cast(sizeof(int4))) + { + *reinterpret_cast(shmem + off) = make_int4(0, 0, 0, 0); + } + if (laneId == 0) + { + // Field 1 sits at getFieldSize(field 0), not at + // align_up(..., 16) as computeTotalUnpackedSize computes; + // they agree only because field 0 is 16-byte aligned. + auto* stats = reinterpret_cast(shmem + field0Size); + int const statsCount = getFieldSize(params.sendFields[1]) / static_cast(sizeof(float2)); + for (int i = 0; i < statsCount; ++i) + { + stats[i] = make_float2(-INFINITY, 0.F); + } + } + __syncwarp(); + } + LL128Proto::protoPack(shmem, head, singlePacked128ByteCount, fifoEntry128ByteIndexBase, laneId); uint64_t* fifoEntry = fifoBase + fifoEntryIndex * (HELIX_FIFO_ENTRY_BYTES / sizeof(uint64_t)); @@ -474,7 +540,7 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) shmem, tail, singlePacked128ByteCount, fifoEntry128ByteIndexBase, loaded128ByteCount, laneId); // note: fields are already unpacked in shared memory - s2gAllFields(params.recvFields, dataIndex, shmem, laneId); + s2gAllFields(params.recvFields, dataIndex, shmem, laneId); // wait for data to be read from shared memory cp_async_bulk_wait_group_read<0>(); @@ -503,7 +569,7 @@ struct hash_cache_key } }; -template +template std::tuple computeChannelAndGroupCount(int cpSize, HelixFieldInfo const* fields) { static std::unordered_map, std::tuple, hash_cache_key> cache; @@ -535,7 +601,7 @@ std::tuple computeChannelAndGroupCount(int cpSize, HelixFieldInfo // Set shared memory attribute if needed if (totalDynamicShmemSize > 48 * 1024) { - TLLM_CUDA_CHECK(cudaFuncSetAttribute(helixAllToAllKernel, + TLLM_CUDA_CHECK(cudaFuncSetAttribute(helixAllToAllKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, totalDynamicShmemSize)); } @@ -559,14 +625,14 @@ std::tuple computeChannelAndGroupCount(int cpSize, HelixFieldInfo // Host Launch Function // ============================================================================ -template +template void launchHelixAllToAllImpl(HelixAllToAllParams const& params, cudaStream_t stream) { int maxChannelCount = computeHelixMaxChannelCount(params.cpSize); TLLM_CHECK_WITH_INFO(params.maxChannelCount == maxChannelCount, "maxChannelCount %d does not match computed maxChannelCount %d", params.maxChannelCount, maxChannelCount); auto [channelCount, groupCountPerCta, totalDynamicShmemSize] - = computeChannelAndGroupCount(params.cpSize, params.sendFields); + = computeChannelAndGroupCount(params.cpSize, params.sendFields); if (params.channelCount > 0) { channelCount = params.channelCount; @@ -580,7 +646,7 @@ void launchHelixAllToAllImpl(HelixAllToAllParams const& params, cudaStream_t str // and receiver) int ctaPerChannel = ceil_div(params.cpSize, groupCountPerCta); - auto* kernel_instance = &helixAllToAllKernel; + auto* kernel_instance = &helixAllToAllKernel; cudaLaunchConfig_t config; config.gridDim = dim3(ctaPerChannel, channelCount, 2); config.blockDim = dim3(WARP_SIZE, groupCountPerCta); @@ -638,11 +704,27 @@ void launchHelixAllToAll(HelixAllToAllParams const& params, bool allowVariableFi { if (allowVariableField1) { - launchHelixAllToAllImpl(params, stream); + constexpr uintptr_t kBulkCopyAlignment = 16; + constexpr uintptr_t kFallbackCopyAlignment = sizeof(float2); + int const field1Size = getFieldSize(params.sendFields[1]); + TLLM_CHECK_WITH_INFO(field1Size % sizeof(float2) == 0, "Variable field 1 must contain whole float2 values"); + TLLM_CHECK_WITH_INFO(isFieldAlignedForCopy(params.sendFields[1], kFallbackCopyAlignment) + && isFieldAlignedForCopy(params.recvFields[1], kFallbackCopyAlignment), + "Variable field 1 must be aligned to float2"); + bool const useBulkField1 = isFieldAlignedForCopy(params.sendFields[1], kBulkCopyAlignment) + && isFieldAlignedForCopy(params.recvFields[1], kBulkCopyAlignment); + if (useBulkField1) + { + launchHelixAllToAllImpl(params, stream); + } + else + { + launchHelixAllToAllImpl(params, stream); + } } else { - launchHelixAllToAllImpl(params, stream); + launchHelixAllToAllImpl(params, stream); } } diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.h b/cpp/tensorrt_llm/kernels/helixAllToAll.h index 95ab959ff3fc..ac51e012adbf 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.h +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.h @@ -47,6 +47,14 @@ struct HelixAllToAllParams int cpSize; int channelCount; // use 0 to auto-compute int maxChannelCount; + + // Rows this rank owns no KV for. The sender replaces them in shared memory + // with a no-op contribution for the combine: field 0 zeros, field 1 + // (max, sum) = (-inf, 0). nullptr when the caller already sanitized. + uint8_t const* zeroKvMask; + // entryCount / zeroKvMask length: 1 when an entry is a token (fifo v2), + // num_heads when it is a (token, head) pair (fifo v1). + int zeroKvMaskDivisor; }; // ============================================================================ diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh index 7817012f2a97..847d05294cf5 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -46,6 +46,7 @@ // empty_input[N_INPUT] pmap -> TMA (input slot empty) // full_cast[2] pmap -> MMA (A ready in TMEM) // empty_cast[2] MMA -> pmap (TMEM slot empty) +// full_mix[1] TMA -> pmap (post_mix+comb_mix arrived) // tmem_full[1] MMA -> epilogue #pragma once @@ -165,18 +166,6 @@ __device__ __forceinline__ float2 fma_f32x2_vv(float2 a, float2 b, float2 c) #endif } -// The read-only cache is not coherent with stores issued earlier by the same -// kernel, so Phase 4's store-to-load path cannot use __ldg. -__device__ __forceinline__ uint4 ld_global_volatile_u128(uint4 const* addr) -{ - uint4 val; - asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) - : "l"(addr) - : "memory"); - return val; -} - __device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uint32_t b, uint32_t c, uint32_t d) { asm volatile( @@ -184,6 +173,36 @@ __device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uin "r"(a), "r"(b), "r"(c), "r"(d)); } +// Blackwell can issue two independent FP32 FMAs with one packed instruction. +// The pmap hot loop always updates adjacent bf16 values with the same mixing +// coefficient, so keep the pair packed instead of lowering it to two FFMA. +__device__ __forceinline__ float2 fma_f32x2(float2 const& a, float2 const& b, float2 const& c) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + float2 result; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); + return result; +#else + return make_float2(fmaf(a.x, b.x, c.x), fmaf(a.y, b.y, c.y)); +#endif +} + +__device__ __forceinline__ float2 mul_f32x2(float2 const& a, float2 const& b) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + float2 result; + asm volatile("mul.f32x2 %0, %1, %2;" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); + return result; +#else + return make_float2(a.x * b.x, a.y * b.y); +#endif +} + template @@ -191,9 +210,10 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 const uint32_t shape_m, const __grid_constant__ cute::TmaDescriptor tensor_map_residual, const __grid_constant__ cute::TmaDescriptor tensor_map_x, const __grid_constant__ cute::TmaDescriptor tensor_map_b, const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, - float* __restrict__ D, // [M, SHAPE_N] (caller memsets to 0) - float const* __restrict__ post_mix, float const* __restrict__ comb_mix, float* __restrict__ sqr_sum) -{ // [M] (caller memsets to 0) + const __grid_constant__ cute::TmaDescriptor tensor_map_post, + const __grid_constant__ cute::TmaDescriptor tensor_map_comb, float const* __restrict__ post_mix, + float const* __restrict__ comb_mix, float* __restrict__ D, float* __restrict__ sqr_sum) +{ // D [M, SHAPE_N], sqr_sum [M] (caller memsets both to 0) #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; @@ -243,6 +263,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 cute::prefetch_tma_descriptor(&tensor_map_x); cute::prefetch_tma_descriptor(&tensor_map_b); cute::prefetch_tma_descriptor(&tensor_map_residual_out); + cute::prefetch_tma_descriptor(&tensor_map_post); + cute::prefetch_tma_descriptor(&tensor_map_comb); } // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] @@ -275,9 +297,10 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); auto empty_cast = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); - auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + auto full_mix = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + auto tmem_full_barrier = full_mix + 1; - cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); + cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 2) * sizeof(Barrier); auto tmem_ptr_in_smem = reinterpret_cast(cursor); if (warp_idx == 1 and cute::elect_one_sync()) @@ -300,6 +323,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 full_cast[i]->init(kNumPmapThreads); empty_cast[i]->init(1); } + full_mix->init(1); tmem_full_barrier->init(1); cutlass::arch::fence_barrier_init(); } @@ -313,8 +337,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 const uint32_t m_block_idx = block_idx / kNumSplits; const uint32_t k_split_idx = block_idx % kNumSplits; const uint32_t m_offset = m_block_idx * BLOCK_M; - // Give the first H_TILES_EXTRA splits one extra tile. Even splits fold back - // to the original constants. + // Give the first H_TILES_EXTRA splits one additional tile. This admits + // Rubin exact-wave KS=106/53 shapes for H=7168 (112 H tiles). uint32_t h_tile_start; uint32_t h_tiles_this_split; if constexpr (H_TILES_EXTRA == 0) @@ -327,16 +351,24 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 h_tile_start = k_split_idx * H_TILES_BASE + cute::min(k_split_idx, H_TILES_EXTRA); h_tiles_this_split = H_TILES_BASE + static_cast(k_split_idx < H_TILES_EXTRA); } - const uint32_t num_total_stages = h_tiles_this_split * HC_MULT; + uint32_t const num_total_stages = h_tiles_this_split * HC_MULT; - // Prologue removed: pmap threads load their own post_mix/comb_mix rows - // directly into registers below, so the MMA/TMA warps never wait on those - // global loads and the TMA pipeline starts immediately. + // Pmap threads load their own post_mix/comb_mix rows directly into registers, + // so the MMA/TMA warps can start filling the TMA pipeline immediately. if (warp_idx < kNumMMAThreads / 32) { // ----- TMA warp (warp 0) ----- if (warp_idx == 0 and cute::elect_one_sync()) { + // Load the two row-wise pmap coefficient tiles asynchronously. + // Only the pmap warp group consumes them, so it waits on full_mix + // independently while this warp continues filling input/B stages. + deep_gemm::tma::copy( + &tensor_map_post, full_mix, smem_post, /*inner_idx=*/0, m_offset); + deep_gemm::tma::copy( + &tensor_map_comb, full_mix, smem_comb, /*inner_idx=*/0, m_offset); + full_mix->arrive_and_expect_tx(SMEM_POST_SIZE + SMEM_COMB_SIZE); + uint32_t b_stage = 0; uint32_t i_stage = 0; uint32_t s = 0; @@ -471,6 +503,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 else { // ----- Pmap warp group (warps 4..7, 128 threads) ----- + full_mix->wait(0); const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; const uint32_t lower_row = upper_row + 8; @@ -595,17 +628,43 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 #pragma unroll for (uint32_t hc = 0; hc < HC_MULT; ++hc) { - float2 nu0 = mul_f32x2(pm_u[hc], xf[0][i + 0]); - float2 nu1 = mul_f32x2(pm_u[hc], xf[0][i + 1]); - float2 nl0 = mul_f32x2(pm_l[hc], xf[1][i + 0]); - float2 nl1 = mul_f32x2(pm_l[hc], xf[1][i + 1]); + float2 nu0, nu1, nl0, nl1; + if constexpr (kNumSplits >= 53) + { + // Packed coefficients are faster for Rubin exact-wave small-M splits. + float2 const coefficientU = make_float2(pm_u[hc], pm_u[hc]); + float2 const coefficientL = make_float2(pm_l[hc], pm_l[hc]); + nu0 = mul_f32x2(coefficientU, xf[0][i + 0]); + nu1 = mul_f32x2(coefficientU, xf[0][i + 1]); + nl0 = mul_f32x2(coefficientL, xf[1][i + 0]); + nl1 = mul_f32x2(coefficientL, xf[1][i + 1]); + } + else + { + nu0 = mul_f32x2(pm_u[hc], xf[0][i + 0]); + nu1 = mul_f32x2(pm_u[hc], xf[0][i + 1]); + nl0 = mul_f32x2(pm_l[hc], xf[1][i + 0]); + nl1 = mul_f32x2(pm_l[hc], xf[1][i + 1]); + } #pragma unroll for (uint32_t j = 0; j < HC_MULT; ++j) { - nu0 = fma_f32x2(cm_u[j][hc], r_u[j][0], nu0); - nu1 = fma_f32x2(cm_u[j][hc], r_u[j][1], nu1); - nl0 = fma_f32x2(cm_l[j][hc], r_l[j][0], nl0); - nl1 = fma_f32x2(cm_l[j][hc], r_l[j][1], nl1); + if constexpr (kNumSplits >= 53) + { + float2 const coefficientU = make_float2(cm_u[j][hc], cm_u[j][hc]); + float2 const coefficientL = make_float2(cm_l[j][hc], cm_l[j][hc]); + nu0 = fma_f32x2(coefficientU, r_u[j][0], nu0); + nu1 = fma_f32x2(coefficientU, r_u[j][1], nu1); + nl0 = fma_f32x2(coefficientL, r_l[j][0], nl0); + nl1 = fma_f32x2(coefficientL, r_l[j][1], nl1); + } + else + { + nu0 = fma_f32x2(cm_u[j][hc], r_u[j][0], nu0); + nu1 = fma_f32x2(cm_u[j][hc], r_u[j][1], nu1); + nl0 = fma_f32x2(cm_l[j][hc], r_l[j][0], nl0); + nl1 = fma_f32x2(cm_l[j][hc], r_l[j][1], nl1); + } } nv_bfloat162 b_u0 = __float22bfloat162_rn(nu0); nv_bfloat162 b_u1 = __float22bfloat162_rn(nu1); @@ -738,17 +797,17 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) const __grid_constant__ cute::TmaDescriptor tensor_map_x, // x_prev, bf16 const __grid_constant__ cute::TmaDescriptor tensor_map_b, // W_T, tf32 const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, // residual_cur, bf16 (TMA store) + const __grid_constant__ cute::TmaDescriptor tensor_map_post, // post_mix_prev, fp32 + const __grid_constant__ cute::TmaDescriptor tensor_map_comb, // comb_mix_prev, fp32 __nv_bfloat16 const* __restrict__ residual_cur_ptr, // same buffer as TMA target __nv_bfloat16* __restrict__ layer_input_out, // [M, HIDDEN] bf16 - float* __restrict__ D, // [M, SHAPE_N] fp32 (y_acc, caller zeros) - float* __restrict__ sqr_sum, // [M] fp32 (r_acc, caller zeros) - int* __restrict__ done_counter, // [ceil(M/BLOCK_M)] int (caller zeros) - float const* __restrict__ post_mix_prev, // [M, HC_MULT] - float const* __restrict__ comb_mix_prev, // [M, HC_MULT, HC_MULT] - float const* __restrict__ hc_scale, // [3] - float const* __restrict__ hc_base, // [HC_MULT*(2+HC_MULT)] - float* __restrict__ post_mix_out, // [M, HC_MULT] - float* __restrict__ comb_mix_out, // [M, HC_MULT, HC_MULT] + float* __restrict__ D, // [M, SHAPE_N] fp32 (y_acc, caller zeros) + float* __restrict__ sqr_sum, // [M] fp32 (r_acc, caller zeros) + int* __restrict__ done_counter, // [ceil(M/BLOCK_M)] int (caller zeros) + float const* __restrict__ hc_scale, // [3] + float const* __restrict__ hc_base, // [HC_MULT*(2+HC_MULT)] + float* __restrict__ post_mix_out, // [M, HC_MULT] + float* __restrict__ comb_mix_out, // [M, HC_MULT, HC_MULT] // When kFuseNorm: layer_input_out receives the RMSNorm-normalized // values: out[t,h] = bf16(li[t,h] * rsqrt(mean(li²)+norm_eps) * w[h]). // norm_weight must be bf16 [HIDDEN]; norm_eps is the RMSNorm epsilon. @@ -808,6 +867,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) cute::prefetch_tma_descriptor(&tensor_map_x); cute::prefetch_tma_descriptor(&tensor_map_b); cute::prefetch_tma_descriptor(&tensor_map_residual_out); + cute::prefetch_tma_descriptor(&tensor_map_post); + cute::prefetch_tma_descriptor(&tensor_map_comb); } // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] @@ -840,9 +901,10 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); auto empty_cast = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); - auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + auto full_mix = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + auto tmem_full_barrier = full_mix + 1; - cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); + cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 2) * sizeof(Barrier); auto tmem_ptr_in_smem = reinterpret_cast(cursor); if (warp_idx == 1 and cute::elect_one_sync()) @@ -865,6 +927,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) full_cast[i]->init(kNumPmapThreads); empty_cast[i]->init(1); } + full_mix->init(1); tmem_full_barrier->init(1); cutlass::arch::fence_barrier_init(); } @@ -878,8 +941,9 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) const uint32_t m_block_idx = block_idx / kNumSplits; const uint32_t k_split_idx = block_idx % kNumSplits; const uint32_t m_offset = m_block_idx * BLOCK_M; - // Give the first H_TILES_EXTRA splits one extra tile. Even splits fold back - // to the original constants. + // Distribute remainder tiles over the first splits. Even-split instances + // fold to the original constants; KS=106/53 provide exact 212-CTA waves + // for H=7168 at M=128/256 on Rubin. uint32_t h_tile_start; uint32_t h_tiles_this_split; if constexpr (H_TILES_EXTRA == 0) @@ -894,44 +958,19 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) } const uint32_t num_total_stages = h_tiles_this_split * HC_MULT; - // Prologue: pmap warp group loads post_mix_prev, comb_mix_prev into SMEM - if (warp_idx >= kNumMMAThreads / 32) - { - const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads; -#pragma unroll - for (uint32_t t = 0; t < 2; ++t) - { - uint32_t idx = pmap_tid + t * kNumPmapThreads; - if (idx < BLOCK_M * HC_MULT) - { - uint32_t m = idx / HC_MULT; - uint32_t hc = idx % HC_MULT; - uint32_t gmem_m = m_offset + m; - float v = (gmem_m < shape_m) ? post_mix_prev[gmem_m * HC_MULT + hc] : 0.f; - smem_post[idx] = v; - } - } -#pragma unroll - for (uint32_t t = 0; t < 8; ++t) - { - uint32_t idx = pmap_tid + t * kNumPmapThreads; - if (idx < BLOCK_M * HC_MULT * HC_MULT) - { - uint32_t m = idx / (HC_MULT * HC_MULT); - uint32_t jk = idx % (HC_MULT * HC_MULT); - uint32_t gmem_m = m_offset + m; - float v = (gmem_m < shape_m) ? comb_mix_prev[gmem_m * HC_MULT * HC_MULT + jk] : 0.f; - smem_comb[idx] = v; - } - } - } - __syncthreads(); - if (warp_idx < kNumMMAThreads / 32) { // ----- TMA warp (warp 0) ----- if (warp_idx == 0 and cute::elect_one_sync()) { + // Fetch coefficient tiles asynchronously while this warp starts + // filling the regular input/B pipeline. Only pmap waits on full_mix. + deep_gemm::tma::copy( + &tensor_map_post, full_mix, smem_post, /*inner_idx=*/0, m_offset); + deep_gemm::tma::copy( + &tensor_map_comb, full_mix, smem_comb, /*inner_idx=*/0, m_offset); + full_mix->arrive_and_expect_tx(SMEM_POST_SIZE + SMEM_COMB_SIZE); + uint32_t b_stage = 0; uint32_t i_stage = 0; uint32_t s = 0; @@ -1066,6 +1105,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) else { // ----- Pmap warp group (warps 4..7, 128 threads) ----- + full_mix->wait(0); const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; const uint32_t lower_row = upper_row + 8; @@ -1224,9 +1264,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) } } - // tma_store_wait is the `.read` form and retires only the SMEM reads, - // so the global writes could still be in flight at the Phase-3 fence. - asm volatile("cp.async.bulk.wait_group 0;" : : : "memory"); + // Drain any in-flight residual_out TMA stores before exit. + cute::tma_store_wait<0>(); // Warp-reduce sqr across 4 col_lanes then atomicAdd to global. sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1); @@ -1286,8 +1325,6 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) { /* spin */ } - // Acquire the Phase-2 writes published before each increment. - __threadfence(); } __syncthreads(); } @@ -1350,67 +1387,47 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) if (tok >= shape_m) continue; - // Lanes 0..HC_MULT-1 compute rmsnorm / sigmoid / sinkhorn; pre_mix is - // held in `pre_mix_local` on lanes 0..HC_MULT-1 and later broadcast to - // all 32 lanes via __shfl_sync. All warps in a team redundantly run - // these ~tens of FLOPs (cheap) to avoid a cross-warp SMEM sync; only - // warp_in_team==0 writes comb_mix_out / post_mix_out to GMEM. + // Every team warp needs pre_mix locally for the HIDDEN-stride loop. + // post_mix and iterative Sinkhorn are token-only results, so only the + // first warp computes them. KS=112/56/28 thereby removes 8x/4x/2x + // redundant Sinkhorn work without adding a cross-warp synchronization. + // Lanes 0..HC_MULT-1 hold the four rows and broadcast pre_mix below. float pre_mix_local = 0.f; if (lane_bf < HC_MULT) { float const r_val = sqr_sum[tok]; - float y_local[HC_MULT3]; float const* y_row = D + static_cast(tok) * SHAPE_N; -#pragma unroll - for (uint32_t c = 0; c < HC_MULT3; ++c) - y_local[c] = y_row[c]; - float const rstd = rsqrtf(r_val / static_cast(HC_MULT * HIDDEN) + rms_eps); float const s0 = hc_scale[0]; - float const s1 = hc_scale[1]; - float const s2 = hc_scale[2]; - float v = y_local[lane_bf] * rstd * s0 + hc_base[lane_bf]; + float v = y_row[lane_bf] * rstd * s0 + hc_base[lane_bf]; pre_mix_local = 1.0f / (1.0f + __expf(-v)) + hc_pre_eps; - v = y_local[HC_MULT + lane_bf] * rstd * s1 + hc_base[HC_MULT + lane_bf]; - float post_val = 1.0f / (1.0f + __expf(-v)) * hc_post_mult_value; if (warp_in_team == 0) { + float const s1 = hc_scale[1]; + float const s2 = hc_scale[2]; + v = y_row[HC_MULT + lane_bf] * rstd * s1 + hc_base[HC_MULT + lane_bf]; + float const post_val = 1.0f / (1.0f + __expf(-v)) * hc_post_mult_value; post_mix_out[tok * HC_MULT + lane_bf] = post_val; - } - float cm_vals[HC_MULT]; + float cm_vals[HC_MULT]; #pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = y_local[2 * HC_MULT + lane_bf * HC_MULT + k] * rstd * s2 - + hc_base[2 * HC_MULT + lane_bf * HC_MULT + k]; + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = y_row[2 * HC_MULT + lane_bf * HC_MULT + k] * rstd * s2 + + hc_base[2 * HC_MULT + lane_bf * HC_MULT + k]; - constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; - float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = __expf(cm_vals[k] - rowMax); - // Reciprocal-multiply for sinkhorn: 1 fdiv + 4 fmul instead of 4 - // fdivs per row-normalize. Equivalent under fp32 round-off. - float inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3]); + constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; + float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); #pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] = cm_vals[k] * inv_rs + hc_sinkhorn_eps; -#pragma unroll - for (uint32_t k = 0; k < HC_MULT; ++k) - { - float cs = cm_vals[k]; - cs += __shfl_xor_sync(LANE_MASK, cs, 1); - cs += __shfl_xor_sync(LANE_MASK, cs, 2); - cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); - } - for (uint32_t it = 1; it < sinkhorn_repeat; ++it) - { - inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3] + hc_sinkhorn_eps); + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = __expf(cm_vals[k] - rowMax); + // Reciprocal-multiply for sinkhorn: 1 fdiv + 4 fmul instead of 4 + // fdivs per row-normalize. Equivalent under fp32 round-off. + float inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3]); #pragma unroll for (uint32_t k = 0; k < HC_MULT; ++k) - cm_vals[k] *= inv_rs; + cm_vals[k] = cm_vals[k] * inv_rs + hc_sinkhorn_eps; #pragma unroll for (uint32_t k = 0; k < HC_MULT; ++k) { @@ -1419,9 +1436,21 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) cs += __shfl_xor_sync(LANE_MASK, cs, 2); cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); } - } - if (warp_in_team == 0) - { + for (uint32_t it = 1; it < sinkhorn_repeat; ++it) + { + inv_rs = 1.0f / (cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3] + hc_sinkhorn_eps); +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] *= inv_rs; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + { + float cs = cm_vals[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm_vals[k] *= 1.0f / (cs + hc_sinkhorn_eps); + } + } float* cm_out_ptr = comb_mix_out + tok * HC_MULT2; #pragma unroll for (uint32_t k = 0; k < HC_MULT; ++k) @@ -1439,10 +1468,6 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) // When WARPS_PER_TOK>1, warp_in_team 0..WARPS_PER_TOK-1 together cover // HIDDEN in strides of WARPS_PER_TOK * 32 * 8. When WARPS_PER_TOK==1, // each warp sweeps HIDDEN alone (same as the original single-warp case). - - // First read of these addresses, ordered by the acquire fence above, so - // __ldg stays for residual read throughput. The layer_input reload below - // is the same-thread store-to-load case and cannot. __nv_bfloat16 const* rbase = residual_cur_ptr + static_cast(tok) * HC_MULT * HIDDEN; __nv_bfloat16* obase = layer_input_out + static_cast(tok) * HIDDEN; @@ -1542,7 +1567,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) // Reduce: intra-warp __shfl_xor; cross-warp via SMEM + per-team // named PTX barrier when WARPS_PER_TOK > 1 (KS ≥ 16 instances). // Inactive teams `continue`'d above and never reach this barrier. - // Pass 2: volatile-reload layer_input_out, multiply by + // Pass 2: re-LDG layer_input_out from L2, multiply by // rsqrt * norm_weight, STG normalized bf16 back to the same // address. Avoids the FMA recompute that doubling pass 1 would // require (Path D Phase 4 is already FMA-heavy). @@ -1654,15 +1679,13 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) float const rsqrt_val = rsqrtf(sum_sq_local / static_cast(HIDDEN) + norm_eps); - // Pass 2: reload the un-normalized layer_input this thread wrote, - // LDG norm_weight, normalize, STG back. No FMA recompute. Each - // thread reloads only its own stores, so no CTA fence is needed -- - // just a load that bypasses the read-only cache. + // Pass 2: re-LDG the un-normalized layer_input we just wrote + // (L2-hot), LDG norm_weight, normalize, STG back. No FMA recompute. __nv_bfloat16 const* nbase = norm_weight; #pragma unroll for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) { - uint4 li_raw = ld_global_volatile_u128(reinterpret_cast(&obase[h])); + uint4 li_raw = __ldg(reinterpret_cast(&obase[h])); uint4 nw_raw = __ldg(reinterpret_cast(&nbase[h])); __nv_bfloat162 const* li_pairs = reinterpret_cast<__nv_bfloat162 const*>(&li_raw); __nv_bfloat162 const* nw_pairs = reinterpret_cast<__nv_bfloat162 const*>(&nw_raw); @@ -1685,7 +1708,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) if (my_chunk < TAIL_CHUNKS) { const uint32_t h = H_VEC_END + my_chunk * BF16_VEC_LI; - uint4 li_raw = ld_global_volatile_u128(reinterpret_cast(&obase[h])); + uint4 li_raw = __ldg(reinterpret_cast(&obase[h])); uint4 nw_raw = __ldg(reinterpret_cast(&nbase[h])); __nv_bfloat162 const* li_pairs = reinterpret_cast<__nv_bfloat162 const*>(&li_raw); __nv_bfloat162 const* nw_pairs = reinterpret_cast<__nv_bfloat162 const*>(&nw_raw); diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index 69bd1ec3b9c2..c098a4203b0e 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -128,10 +128,10 @@ static bool isSupportedFhcHiddenRuntime(int hidden_size) return hidden_size == static_cast(FHC_HIDDEN_FLASH) || hidden_size == static_cast(FHC_HIDDEN_PRO); } -// Validate the tcgen05 MMA fused-HC compile-time shape contract. Hidden must -// be divisible into BLOCK_K tiles, and the hidden dimension must be a multiple -// of BF16_VEC_LI (per-thread vector load granularity in the Phase 4 -// layer_input loop). The (Hidden % team-stride) +// Validate the tcgen05 all-in-one fused-HC compile-time shape contract. Hidden +// must be divisible into BLOCK_K tiles, KS must evenly divide those tiles, and +// the hidden dimension must be a multiple of BF16_VEC_LI (per-thread vector +// load granularity in the Phase 4 layer_input loop). The (Hidden % team-stride) // alignment is no longer required: the layer_input loop has a scalar-vec tail // that handles the residue after the vectorized main loop. Keep this in sync // with the Python tactic filter (_fused_hc_mma_ks_supported in mhc_cuda.py). @@ -144,13 +144,28 @@ static constexpr bool isSupportedFhcMmaKS() constexpr uint32_t hTilesPerHc = Hidden / FHC_BLOCK_K; constexpr uint32_t bf16VecLi = 8; - // The kernels distribute remainder tiles over their first splits, so KS need - // not divide hTilesPerHc. Keep the uneven surface to the two measured H=7168 - // shapes rather than instantiating every split count. constexpr bool evenSplit = hTilesPerHc % KS == 0; - constexpr bool singleWaveSplit = Hidden == FHC_HIDDEN_PRO && (KS == 53 || KS == 106); + constexpr bool rubinExactSplit = Hidden == FHC_HIDDEN_PRO && (KS == 53 || KS == 106); - return Hidden % FHC_BLOCK_K == 0 && KS <= hTilesPerHc && (evenSplit || singleWaveSplit) && Hidden % bf16VecLi == 0; + return Hidden % FHC_BLOCK_K == 0 && KS <= hTilesPerHc && (evenSplit || rubinExactSplit) && Hidden % bf16VecLi == 0; +} + +// The half-MMA kernel can distribute remainder tiles over its first splits. +// Keep the uneven support surface limited to the two measured exact-wave +// H=7168 shapes for 212-SM Rubin. +template +static constexpr bool isSupportedFhcHalfMmaKS() +{ + static_assert(isSupportedFhcHidden(), "Unsupported fused-HC hidden size"); + static_assert(KS > 0, "kNumSplits must be positive"); + + constexpr uint32_t hTilesPerHc = Hidden / FHC_BLOCK_K; + constexpr uint32_t bf16VecLi = 8; + + constexpr bool evenSplit = hTilesPerHc % KS == 0; + constexpr bool rubinExactSplit = Hidden == FHC_HIDDEN_PRO && (KS == 53 || KS == 106); + + return Hidden % FHC_BLOCK_K == 0 && KS <= hTilesPerHc && (evenSplit || rubinExactSplit) && Hidden % bf16VecLi == 0; } static CUtensorMap makeTma2D(void* base, CUtensorMapDataType dtype, uint64_t gmemInner, uint64_t gmemOuter, @@ -178,9 +193,10 @@ static CUtensorMap makeTma2D(void* base, CUtensorMapDataType dtype, uint64_t gme // ---- TMA descriptor cache -------------------------------------------------- // // cuTensorMapEncodeTiled is a host-side call that takes ~1-2 µs per descriptor. -// Each fused_hc launch builds 4 descriptors (residual_in, x_in, W, -// residual_cur), so the per-call descriptor build is 4-8 µs — 25-50% of total -// wall time at small M (M ≤ 64). +// On a cache miss, the half-MMA path builds 6 descriptors (residual_in, x_in, +// W, residual_cur, post_mix, comb_mix) and the all-in-one path builds 4. The +// per-call encode cost is material at small M, so hits stay entirely host-side +// and each launcher resolves the current device only once. // // Cache scope: per-host-thread (`thread_local`). Same host thread launching to // multiple CUDA streams shares one cache (descriptor content depends only on @@ -201,7 +217,7 @@ static CUtensorMap makeTma2D(void* base, CUtensorMapDataType dtype, uint64_t gme // fresh `base` pointers as public outputs are allocated, so the unbounded // version would grow across shape transitions. 128 entries × ~256 B = ~32 KB // per host thread — fits in L1, sized to cover the working set of any single -// model (~4-8 distinct shapes × 4 descriptors each = O(20) live, with +// model (~4-8 distinct shapes × up to 6 descriptors each = O(30) live, with // headroom for shape transitions). namespace { @@ -254,11 +270,10 @@ struct TmaDescCache }; CUtensorMap getCachedTma2D(void* base, CUtensorMapDataType dtype, uint64_t gmemInner, uint64_t gmemOuter, - uint32_t smemInner, uint32_t smemOuter, uint64_t gmemOuterStrideBytes, uint32_t swizzleBytes, uint32_t elemBytes) + uint32_t smemInner, uint32_t smemOuter, uint64_t gmemOuterStrideBytes, uint32_t swizzleBytes, uint32_t elemBytes, + int device_id) { static thread_local TmaDescCache cache; - int device_id = 0; - cudaGetDevice(&device_id); TmaDescKey const key{base, gmemInner, gmemOuter, smemInner, smemOuter, gmemOuterStrideBytes, swizzleBytes, elemBytes, dtype, device_id}; auto it = cache.index.find(key); @@ -294,20 +309,20 @@ static constexpr uint32_t fhcSmemSize() constexpr uint32_t SMEM_COMB = FHC_BLOCK_M * FHC_HC_MULT * FHC_HC_MULT * sizeof(float); constexpr uint32_t SMEM_RC = FHC_HC_MULT * FHC_BLOCK_M * FHC_BLOCK_K * sizeof(__nv_bfloat16); constexpr uint32_t kNumCast = 4; - constexpr uint32_t barriers = 2 * FHC_N_B_STAGES + 2 * FHC_N_INPUT_STG + 2 * kNumCast + 1; + constexpr uint32_t barriers = 2 * FHC_N_B_STAGES + 2 * FHC_N_INPUT_STG + 2 * kNumCast + 2; // 4 bytes for the tmem ptr word + 32 bytes padding for alignment headroom. return SMEM_CD + FHC_N_B_STAGES * SMEM_B + FHC_N_INPUT_STG * (SMEM_RES_ISTG + SMEM_X_ISTG) + SMEM_POST + SMEM_COMB + SMEM_RC + barriers * 8 + 4 + 32; } -using FusedRoutFn = void (*)( - uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, float*, float const*, float const*, float*); +using FusedRoutFn = void (*)(uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, + float const*, float const*, float*, float*); template static FusedRoutFn fhcInstance() { static_assert(isSupportedFhcHidden(), "Unsupported fused-HC hidden size"); - static_assert(isSupportedFhcMmaKS(), "Unsupported fused-HC MMA kNumSplits for hidden size"); + static_assert(isSupportedFhcHalfMmaKS(), "Unsupported fused-HC half-MMA kNumSplits for hidden size"); return &fused_mhc::fused_tf32_pmap_gemm_rout_atomic_impl; @@ -316,7 +331,7 @@ static FusedRoutFn fhcInstance() template static FusedRoutFn fhcInstanceIfSupported() { - if constexpr (isSupportedFhcMmaKS()) + if constexpr (isSupportedFhcHalfMmaKS()) { return fhcInstance(); } @@ -401,6 +416,8 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons constexpr uint32_t SHAPE_K = FHC_HC_MULT * Hidden; uint32_t const m_u = static_cast(M); + int device_id = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&device_id)); uint32_t const ks = (num_k_splits > 0) ? static_cast(num_k_splits) : pickKSplits(M); int const bs = (bigfuse_block_size > 0) ? bigfuse_block_size : selectBigFuseBS(M); @@ -416,19 +433,27 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons // ---- Build TMA descriptors (cached by ptr+shape) ---- CUtensorMap desc_res = getCachedTma2D(const_cast<__nv_bfloat16*>(residual_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); CUtensorMap desc_x = getCachedTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, Hidden, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(Hidden) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); CUtensorMap desc_b = getCachedTma2D(const_cast(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N, FHC_BLOCK_K, FHC_BLOCK_N, static_cast(SHAPE_K) * sizeof(float), - /*swizzleBytes=*/128, sizeof(float)); + /*swizzleBytes=*/128, sizeof(float), device_id); CUtensorMap desc_res_out = getCachedTma2D(residual_cur, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, /*smemOuter=*/16, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); + + CUtensorMap desc_post = getCachedTma2D(const_cast(post_mix_prev), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + FHC_HC_MULT, m_u, FHC_HC_MULT, FHC_BLOCK_M, static_cast(FHC_HC_MULT) * sizeof(float), + /*swizzleBytes=*/0, sizeof(float), device_id); + + CUtensorMap desc_comb = getCachedTma2D(const_cast(comb_mix_prev), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + FHC_HC_MULT * FHC_HC_MULT, m_u, FHC_HC_MULT * FHC_HC_MULT, FHC_BLOCK_M, + static_cast(FHC_HC_MULT * FHC_HC_MULT) * sizeof(float), /*swizzleBytes=*/0, sizeof(float), device_id); // ---- Step 1: fused post-mapping + TF32 GEMM + sqrsum + residual_out ---- constexpr uint32_t fused_smem = fhcSmemSize(); @@ -439,8 +464,8 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons uint32_t const m_tiles = (m_u + FHC_BLOCK_M - 1) / FHC_BLOCK_M; dim3 const grid(m_tiles * ks); dim3 const block(FHC_NUM_MMA_TH + FHC_NUM_PMAP_TH); - fa<<>>( - m_u, desc_res, desc_x, desc_b, desc_res_out, y_acc_workspace, post_mix_prev, comb_mix_prev, r_acc_workspace); + fa<<>>(m_u, desc_res, desc_x, desc_b, desc_res_out, desc_post, desc_comb, + post_mix_prev, comb_mix_prev, y_acc_workspace, r_acc_workspace); // ---- Step 2: big-fuse postlogue (RMS + sigmoid + Sinkhorn + pre-apply) ---- // Delegate to mhcBigFuseLaunch (defined in mhcKernels.cu) to avoid @@ -584,8 +609,8 @@ static constexpr uint32_t fhcAllInOneSmemSize() return fhcSmemSize(); } -using FusedAllInOneFn = void (*)(uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, __nv_bfloat16 const*, - __nv_bfloat16*, float*, float*, int*, float const*, float const*, float const*, float const*, float*, float*, +using FusedAllInOneFn = void (*)(uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, + __nv_bfloat16 const*, __nv_bfloat16*, float*, float*, int*, float const*, float const*, float*, float*, __nv_bfloat16 const*, float, float, float, float, float, uint32_t); template @@ -657,6 +682,8 @@ static void mhcFusedHcAllInOneLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloa constexpr uint32_t SHAPE_K = FHC_HC_MULT * Hidden; uint32_t const m_u = static_cast(M); + int device_id = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&device_id)); uint32_t const ks = (num_k_splits > 0) ? static_cast(num_k_splits) : 1u; uint32_t const m_tiles = (m_u + FHC_BLOCK_M - 1) / FHC_BLOCK_M; @@ -673,19 +700,27 @@ static void mhcFusedHcAllInOneLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloa // ---- Build TMA descriptors (cached by ptr+shape) ---- CUtensorMap desc_res = getCachedTma2D(const_cast<__nv_bfloat16*>(residual_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); CUtensorMap desc_x = getCachedTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, Hidden, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(Hidden) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); CUtensorMap desc_b = getCachedTma2D(const_cast(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N, FHC_BLOCK_K, FHC_BLOCK_N, static_cast(SHAPE_K) * sizeof(float), - /*swizzleBytes=*/128, sizeof(float)); + /*swizzleBytes=*/128, sizeof(float), device_id); CUtensorMap desc_res_out = getCachedTma2D(residual_cur, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, /*smemOuter=*/16, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), - /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + /*swizzleBytes=*/128, sizeof(__nv_bfloat16), device_id); + + CUtensorMap desc_post = getCachedTma2D(const_cast(post_mix_prev), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + FHC_HC_MULT, m_u, FHC_HC_MULT, FHC_BLOCK_M, static_cast(FHC_HC_MULT) * sizeof(float), + /*swizzleBytes=*/0, sizeof(float), device_id); + + CUtensorMap desc_comb = getCachedTma2D(const_cast(comb_mix_prev), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, + FHC_HC_MULT * FHC_HC_MULT, m_u, FHC_HC_MULT * FHC_HC_MULT, FHC_BLOCK_M, + static_cast(FHC_HC_MULT * FHC_HC_MULT) * sizeof(float), /*swizzleBytes=*/0, sizeof(float), device_id); // ---- Launch the single all-in-one kernel ---- // Dispatch on `norm_weight != nullptr` to a kFuseNorm=true instance that @@ -699,8 +734,8 @@ static void mhcFusedHcAllInOneLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloa dim3 const grid(m_tiles * ks); dim3 const block(FHC_NUM_MMA_TH + FHC_NUM_PMAP_TH); - fa<<>>(m_u, desc_res, desc_x, desc_b, desc_res_out, residual_cur, layer_input_cur, - y_acc_workspace, r_acc_workspace, done_counter_workspace, post_mix_prev, comb_mix_prev, hc_scale, hc_base, + fa<<>>(m_u, desc_res, desc_x, desc_b, desc_res_out, desc_post, desc_comb, + residual_cur, layer_input_cur, y_acc_workspace, r_acc_workspace, done_counter_workspace, hc_scale, hc_base, post_mix_cur, comb_mix_cur, norm_weight, norm_eps, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, static_cast(sinkhorn_repeat)); } diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu index 48d2a9e5e3a7..3d114a83b303 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu @@ -27,6 +27,22 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels::mhc { +// Blackwell can issue two independent FP32 FMAs with one packed instruction. +// Phase 2 always updates adjacent bf16 values with the same pre-mix scalar. +__device__ __forceinline__ float2 mhcFmaF32x2(float2 const& a, float2 const& b, float2 const& c) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + float2 result; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); + return result; +#else + return make_float2(fmaf(a.x, b.x, c.x), fmaf(a.y, b.y, c.y)); +#endif +} + // =================================================================== // Kernel 1: big_fuse — one CTA per token // @@ -37,11 +53,11 @@ namespace kernels::mhc // Phase 1a (warp 0, lanes 0-3): RMS norm + sigmoid → s_pre_mix, post_mix // ── __syncthreads ── // Phase 1b (warp 0, lanes 0-3) ‖ Phase 2 (remaining warps) — overlapped -// 1b: parallel Sinkhorn (4 lanes, __shfl_xor col normalize) → comb_mix +// 1b: one 4-element Sinkhorn row per lane → comb_mix // 2: stream residual × pre_mix → layer_input // =================================================================== -template +template __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __restrict__ y_acc, float const* __restrict__ r_acc, __nv_bfloat16 const* __restrict__ residual, float const* __restrict__ hc_scale, float const* __restrict__ hc_base, float* __restrict__ post_mix, float* __restrict__ comb_mix, @@ -67,40 +83,73 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re cudaGridDependencySynchronize(); #endif + // Small-M/BS512 specialization: issue one bulk asynchronous copy for the + // token's contiguous [4, hidden] residual before Phase 1a, then overlap it + // with RMS/sigmoid/Sinkhorn. Consumers wait on the mbarrier independently, + // so warp 0 can keep running Sinkhorn while the copy completes. + extern __shared__ __align__(16) unsigned char s_tma_residual_raw[]; + __shared__ alignas(8) uint64_t s_tma_bar; + if constexpr (kUseTma) + { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + uint32_t const tma_bytes = HC_MULT * hidden_size * static_cast(sizeof(__nv_bfloat16)); + if (tid == 0) + { + uint32_t const bar_addr = static_cast(__cvta_generic_to_shared(&s_tma_bar)); + uint32_t const dst_addr = static_cast(__cvta_generic_to_shared(s_tma_residual_raw)); + __nv_bfloat16 const* src = residual + static_cast(token) * HC_MULT * hidden_size; + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" : : "r"(bar_addr) : "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" + : + : "r"(bar_addr), "r"(tma_bytes) + : "memory"); + asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" + : + : "r"(dst_addr), "l"(reinterpret_cast(src)), "r"(tma_bytes), "r"(bar_addr) + : "memory"); + } +#endif + } + __shared__ float s_pre_mix[HC_MULT]; float cm[HC_MULT]; // ---- Phase 1a (warp 0, lanes 0..3): split-K reduce → RMS norm → sigmoid ---- - // Each lane handles one of the 4 hc_mult slots. - // Produces: s_pre_mix[4] (shared), post_mix[4] (global), cm[4×4] (registers). + // Keep one full comb row per lane: its four independent values provide ILP + // across Sinkhorn's reciprocal/shuffle dependency chains. Load only the + // six y values this lane actually consumes instead of materializing all 24. if (warp_id == 0 && lane < HC_MULT) { float r_val; - float y_local[HC_MULT3]; + float y_pre = 0.0f; + float y_post = 0.0f; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm[k] = 0.0f; if constexpr (NUM_SPLITS == 1) { r_val = r_acc[token]; float const* y_row = y_acc + token * HC_MULT3; + y_pre = y_row[lane]; + y_post = y_row[HC_MULT + lane]; #pragma unroll - for (int c = 0; c < HC_MULT3; c++) - y_local[c] = y_row[c]; + for (int k = 0; k < HC_MULT; k++) + cm[k] = y_row[2 * HC_MULT + lane * HC_MULT + k]; } else { - // Reduce across split-K partials r_val = 0.0f; -#pragma unroll - for (int c = 0; c < HC_MULT3; c++) - y_local[c] = 0.0f; for (int s = 0; s < NUM_SPLITS; s++) { r_val += r_acc[s * M + token]; float const* y_row = y_acc + (static_cast(s) * M + token) * HC_MULT3; + y_pre += y_row[lane]; + y_post += y_row[HC_MULT + lane]; #pragma unroll - for (int c = 0; c < HC_MULT3; c++) - y_local[c] += y_row[c]; + for (int k = 0; k < HC_MULT; k++) + cm[k] += y_row[2 * HC_MULT + lane * HC_MULT + k]; } } @@ -108,76 +157,60 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re float const rstd = rsqrtf(r_val / static_cast(K) + rms_eps); float const s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2]; - // y_local layout: [pre_mix(4) | post_mix(4) | comb_mix(4×4)] - // pre_mix: sigmoid(norm * scale0 + base) + eps → shared for Phase 2 - float v = y_local[lane] * rstd * s0 + hc_base[lane]; + float v = y_pre * rstd * s0 + hc_base[lane]; s_pre_mix[lane] = 1.0f / (1.0f + expf(-v)) + hc_pre_eps; - // post_mix: sigmoid(norm * scale1 + base) * mult → global - v = y_local[HC_MULT + lane] * rstd * s1 + hc_base[HC_MULT + lane]; + v = y_post * rstd * s1 + hc_base[HC_MULT + lane]; post_mix[token * HC_MULT + lane] = 1.0f / (1.0f + expf(-v)) * hc_post_mult_value; - // comb_mix init: norm * scale2 + base → cm[4] per lane (one row of 4×4 matrix) #pragma unroll for (int k = 0; k < HC_MULT; k++) - cm[k] = y_local[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; + cm[k] = cm[k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; } __syncthreads(); // ---- Phase 1b (warp 0, lanes 0..3): Sinkhorn normalization ---- - // Each lane holds one row of the 4×4 comb matrix. - // Row normalize via local sum, column normalize via __shfl_xor across 4 lanes. if (warp_id == 0 && lane < HC_MULT) { - constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; // 0xf for HC_MULT=4 + constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; // Softmax rows: subtract the row max to avoid inf / inf when comb logits // are large. - float const rowMax = fmaxf(fmaxf(cm[0], cm[1]), fmaxf(cm[2], cm[3])); + float const row_max = fmaxf(fmaxf(cm[0], cm[1]), fmaxf(cm[2], cm[3])); #pragma unroll for (int k = 0; k < HC_MULT; k++) - cm[k] = expf(cm[k] - rowMax); - // Replace per-element fdiv with one reciprocal + 4 fmul. fp32 fdiv on - // B200 is multi-cycle while fmul retires at peak rate; sinkhorn's - // O(HC_MULT * sinkhorn_repeat) divisions per token (160 at sinkhorn=20) - // dominate the bigfuse epilogue cost on this 4-lane warp. Math is - // identical modulo last-bit round-off, which sinkhorn iteration - // absorbs. - float inv_rs = 1.0f / (cm[0] + cm[1] + cm[2] + cm[3]); + cm[k] = expf(cm[k] - row_max); + float inv_row_sum = 1.0f / (cm[0] + cm[1] + cm[2] + cm[3]); #pragma unroll for (int k = 0; k < HC_MULT; k++) - cm[k] = cm[k] * inv_rs + hc_sinkhorn_eps; - - // Column normalize: sum across lanes (rows) via butterfly shuffle + cm[k] = cm[k] * inv_row_sum + hc_sinkhorn_eps; #pragma unroll for (int k = 0; k < HC_MULT; k++) { - float cs = cm[k]; - cs += __shfl_xor_sync(LANE_MASK, cs, 1); - cs += __shfl_xor_sync(LANE_MASK, cs, 2); - cm[k] *= 1.0f / (cs + hc_sinkhorn_eps); + float col_sum = cm[k]; + col_sum += __shfl_xor_sync(LANE_MASK, col_sum, 1); + col_sum += __shfl_xor_sync(LANE_MASK, col_sum, 2); + cm[k] *= 1.0f / (col_sum + hc_sinkhorn_eps); } // Remaining Sinkhorn iterations: alternate row / column normalize for (int it = 1; it < sinkhorn_repeat; it++) { - inv_rs = 1.0f / (cm[0] + cm[1] + cm[2] + cm[3] + hc_sinkhorn_eps); + inv_row_sum = 1.0f / (cm[0] + cm[1] + cm[2] + cm[3] + hc_sinkhorn_eps); #pragma unroll for (int k = 0; k < HC_MULT; k++) - cm[k] *= inv_rs; - + cm[k] *= inv_row_sum; #pragma unroll for (int k = 0; k < HC_MULT; k++) { - float cs = cm[k]; - cs += __shfl_xor_sync(LANE_MASK, cs, 1); - cs += __shfl_xor_sync(LANE_MASK, cs, 2); - cm[k] *= 1.0f / (cs + hc_sinkhorn_eps); + float col_sum = cm[k]; + col_sum += __shfl_xor_sync(LANE_MASK, col_sum, 1); + col_sum += __shfl_xor_sync(LANE_MASK, col_sum, 2); + cm[k] *= 1.0f / (col_sum + hc_sinkhorn_eps); } } - // Write 4×4 comb_mix to global (lane = row index, k = col index) float* cm_out = comb_mix + token * HC_MULT2; #pragma unroll for (int k = 0; k < HC_MULT; k++) @@ -206,7 +239,30 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re for (int j = 0; j < HC_MULT; j++) pm[j] = s_pre_mix[j]; - __nv_bfloat16 const* rbase = residual + static_cast(token) * HC_MULT * hidden_size; + __nv_bfloat16 const* rbase; + if constexpr (kUseTma) + { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + uint32_t const bar_addr = static_cast(__cvta_generic_to_shared(&s_tma_bar)); + uint32_t complete; + do + { + asm volatile( + "{ .reg .pred P; mbarrier.try_wait.parity.shared::cta.b64 P, [%1], %2;" + " selp.b32 %0, 1, 0, P; }" + : "=r"(complete) + : "r"(bar_addr), "r"(0u) + : "memory"); + } while (!complete); + rbase = reinterpret_cast<__nv_bfloat16 const*>(s_tma_residual_raw); +#else + rbase = residual + static_cast(token) * HC_MULT * hidden_size; +#endif + } + else + { + rbase = residual + static_cast(token) * HC_MULT * hidden_size; + } __nv_bfloat16* obase = layer_input + static_cast(token) * hidden_size; int const p2_tid = tid - WARP_SIZE; @@ -215,19 +271,19 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re float sum_sq_local = 0.f; for (int h = p2_tid * BF16_VEC; h < hidden_size; h += p2_threads * BF16_VEC) { - float acc[BF16_VEC] = {}; + float2 acc[BF16_VEC / 2] = {}; #pragma unroll for (int j = 0; j < HC_MULT; j++) { uint4 raw = *reinterpret_cast(&rbase[j * hidden_size + h]); __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raw); + float2 const coefficient = make_float2(pm[j], pm[j]); #pragma unroll for (int v = 0; v < BF16_VEC / 2; v++) { - float2 f = __bfloat1622float2(pairs[v]); - acc[2 * v + 0] += pm[j] * f.x; - acc[2 * v + 1] += pm[j] * f.y; + float2 const f = __bfloat1622float2(pairs[v]); + acc[v] = mhcFmaF32x2(coefficient, f, acc[v]); } } @@ -235,7 +291,7 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); #pragma unroll for (int v = 0; v < BF16_VEC / 2; v++) - opairs[v] = __float22bfloat162_rn(make_float2(acc[2 * v], acc[2 * v + 1])); + opairs[v] = __float22bfloat162_rn(acc[v]); *reinterpret_cast(&obase[h]) = out_raw; if constexpr (kFuseNorm) @@ -327,11 +383,19 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re // turns those otherwise-ignored qualifiers into part of the function type. Clang then rejects the instantiation; // GCC does not diagnose it. #define INST_BIGFUSE(NS, BS) \ - template __global__ void mhcBigFuseKernel(float const* __restrict__, \ - float const* __restrict__, __nv_bfloat16 const* __restrict__, float const* __restrict__, \ - float const* __restrict__, float* __restrict__, float* __restrict__, __nv_bfloat16* __restrict__, int, int, \ - int, float, float, float, float, int, __nv_bfloat16 const* __restrict__, float); \ - template __global__ void mhcBigFuseKernel(float const* __restrict__, \ + template __global__ void mhcBigFuseKernel( \ + float const* __restrict__, float const* __restrict__, __nv_bfloat16 const* __restrict__, \ + float const* __restrict__, float const* __restrict__, float* __restrict__, float* __restrict__, \ + __nv_bfloat16* __restrict__, int, int, int, float, float, float, float, int, \ + __nv_bfloat16 const* __restrict__, float); \ + template __global__ void mhcBigFuseKernel( \ + float const* __restrict__, float const* __restrict__, __nv_bfloat16 const* __restrict__, \ + float const* __restrict__, float const* __restrict__, float* __restrict__, float* __restrict__, \ + __nv_bfloat16* __restrict__, int, int, int, float, float, float, float, int, \ + __nv_bfloat16 const* __restrict__, float); + +#define INST_BIGFUSE_TMA(FN) \ + template __global__ void mhcBigFuseKernel<1, 512, /*kFuseNorm=*/FN, /*kUseTma=*/true>(float const* __restrict__, \ float const* __restrict__, __nv_bfloat16 const* __restrict__, float const* __restrict__, \ float const* __restrict__, float* __restrict__, float* __restrict__, __nv_bfloat16* __restrict__, int, int, \ int, float, float, float, float, int, __nv_bfloat16 const* __restrict__, float); @@ -351,6 +415,9 @@ INST_BIGFUSE(8, 512) INST_BIGFUSE(16, 128) INST_BIGFUSE(16, 256) INST_BIGFUSE(16, 512) +INST_BIGFUSE_TMA(false) +INST_BIGFUSE_TMA(true) +#undef INST_BIGFUSE_TMA #undef INST_BIGFUSE // =================================================================== @@ -791,22 +858,44 @@ static void mhcBigFuseDispatch(float const* y_acc, float const* r_acc, __nv_bflo { dim3 grid(static_cast(M)); -#define LAUNCH_BF(BS) \ - tensorrt_llm::common::launchWithPdlWhenEnabled("mhcBigFuseKernel", mhcBigFuseKernel, \ - grid, dim3(BS), 0, stream, y_acc, r_acc, residual, hc_scale, hc_base, post_mix, comb_mix, layer_input, M, K, \ - hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, norm_weight, norm_eps) +#define LAUNCH_BF(BS, USE_TMA, SMEM_BYTES) \ + tensorrt_llm::common::launchWithPdlWhenEnabled("mhcBigFuseKernel", \ + mhcBigFuseKernel, grid, dim3(BS), SMEM_BYTES, stream, y_acc, r_acc, \ + residual, hc_scale, hc_base, post_mix, comb_mix, layer_input, M, K, hidden_size, rms_eps, hc_pre_eps, \ + hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, norm_weight, norm_eps) if (block_size >= 512) { - LAUNCH_BF(512); + if constexpr (NUM_SPLITS == 1) + { + // BS512 is the measured winner for the small-M one-wave regime. Bulk + // TMA wins only through M=64; its shared-memory footprint regresses + // M>=128, where regular LDG has enough latency hiding. + if (M <= 64) + { + size_t const tma_smem_bytes = static_cast(4) * hidden_size * sizeof(__nv_bfloat16); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + reinterpret_cast(mhcBigFuseKernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(tma_smem_bytes))); + LAUNCH_BF(512, true, tma_smem_bytes); + } + else + { + LAUNCH_BF(512, false, 0); + } + } + else + { + LAUNCH_BF(512, false, 0); + } } else if (block_size >= 256) { - LAUNCH_BF(256); + LAUNCH_BF(256, false, 0); } else { - LAUNCH_BF(128); + LAUNCH_BF(128, false, 0); } #undef LAUNCH_BF } diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh index 48a1290df564..1ce437a24549 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh @@ -798,9 +798,9 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 // Phase 4: inline bigFuse for the TM tokens in this batch. // Layout: FULL_N = HC_MULT*(2+HC_MULT) = 24 - // y_local[0..HC_MULT) → s_pre_mix (sigmoid gate) - // y_local[HC_MULT..2*HC_MULT) → post_mix_out (sigmoid*hc_post_mult) - // y_local[2*HC_MULT..FULL_N) → comb_mix_out (Sinkhorn) + // y_acc[0..HC_MULT) → s_pre_mix (sigmoid gate) + // y_acc[HC_MULT..2*HC_MULT) → post_mix_out (sigmoid*hc_post_mult) + // y_acc[2*HC_MULT..FULL_N) → comb_mix_out (Sinkhorn) __shared__ float s_pre_mix[TM][HC_MULT]; #pragma unroll @@ -812,30 +812,26 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 continue; } int const tok = base_tok + t; - float cm_vals[HC_MULT]; if (warp_id == 0 && lane < HC_MULT) { + float cm_vals[HC_MULT]; float const r_val = r_acc[tok]; - float y_local[HC_MULT3]; float const* y_row = y_acc + static_cast(tok) * FULL_N; -#pragma unroll - for (int c = 0; c < HC_MULT3; c++) - y_local[c] = y_row[c]; float const rstd = rsqrtf(r_val / static_cast(K) + rms_eps); float const s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2]; - float v = y_local[lane] * rstd * s0 + hc_base[lane]; + float v = y_row[lane] * rstd * s0 + hc_base[lane]; s_pre_mix[t][lane] = 1.0f / (1.0f + expf(-v)) + hc_pre_eps; - v = y_local[HC_MULT + lane] * rstd * s1 + hc_base[HC_MULT + lane]; + v = y_row[HC_MULT + lane] * rstd * s1 + hc_base[HC_MULT + lane]; post_mix_out[tok * HC_MULT + lane] = 1.0f / (1.0f + expf(-v)) * hc_post_mult_value; #pragma unroll for (int k = 0; k < HC_MULT; k++) cm_vals[k] - = y_local[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; + = y_row[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); @@ -877,19 +873,18 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 } __syncthreads(); - // layer_input: warp>0 threads process hidden in parallel. + // layer_input: all eight warps process hidden in parallel. // // When kFuseNorm is true, accumulate per-thread sum_sq while writing // un-normalized layer_input bf16; after a __syncthreads, all threads // run pass 2: re-LDG from L2 (hot from pass 1's STGs), multiply by // rsqrt * norm_weight, STG normalized bf16. Saves one HBM read+write // pair vs the separate flashinfer.rmsnorm kernel. - constexpr int kFmaBigFuseWarps = BLOCK_SIZE / WARP_SIZE - 1; + constexpr int kFmaBigFuseWarps = BLOCK_SIZE / WARP_SIZE; __shared__ float s_sumsq_li[kFmaBigFuseWarps]; __shared__ float s_rsqrt_li; constexpr int BF16_VEC_LI = 8; __nv_bfloat16* obase = layer_input_out + static_cast(tok) * hidden_size; - if (warp_id > 0) { float pm[HC_MULT]; #pragma unroll @@ -897,8 +892,8 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 pm[j] = s_pre_mix[t][j]; __nv_bfloat16 const* rbase = residual_out + static_cast(tok) * HC_MULT * hidden_size; - int const p2_tid = tid - WARP_SIZE; - constexpr int p2_threads = BLOCK_SIZE - WARP_SIZE; + int const p2_tid = tid; + constexpr int p2_threads = BLOCK_SIZE; float sum_sq_local = 0.f; for (int h = p2_tid * BF16_VEC_LI; h < hidden_size; h += p2_threads * BF16_VEC_LI) @@ -942,7 +937,7 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 2); sum_sq_local += __shfl_xor_sync(0xffffffff, sum_sq_local, 1); if ((tid & 31) == 0) - s_sumsq_li[warp_id - 1] = sum_sq_local; + s_sumsq_li[warp_id] = sum_sq_local; } } if constexpr (kFuseNorm) @@ -957,10 +952,9 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 s_rsqrt_li = rsqrtf(total / static_cast(hidden_size) + norm_eps); } __syncthreads(); - if (warp_id > 0) { - int const p2_tid = tid - WARP_SIZE; - constexpr int p2_threads = BLOCK_SIZE - WARP_SIZE; + int const p2_tid = tid; + constexpr int p2_threads = BLOCK_SIZE; float const rsqrt_val = s_rsqrt_li; for (int h = p2_tid * BF16_VEC_LI; h < hidden_size; h += p2_threads * BF16_VEC_LI) { diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index 422569c3b3b2..e1a37a874e1a 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.cu +++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu @@ -406,15 +406,16 @@ inline __device__ void dequantCopy( // `kOutputFp8Q`: when true, write the rotated Q rope segment directly to // `quant_q_buf` as FP8 (scaled by `*quant_scale_qkv`) and skip the bf16 STG to -// `q_ptr`. Companion: `deepseek_v4_q_norm_fused_fp8` pre-fills the nope segment -// of `quant_q_buf`, so the standalone quantizeCopyInputToFp8Kernel can be -// dropped. `quant_q_buf`/`quant_scale_qkv`/bmm_scale outputs are unused when -// `kOutputFp8Q == false`. +// `q_ptr`. `deepseek_v4_q_norm_fused_fp8` can pre-fill the nope segment for +// this kernel to append RoPE, while the q_b fusion can pre-fill the complete +// rotated Q and set `q_rope_applied` to preserve it. Both paths drop the +// standalone quantizeCopyInputToFp8Kernel. `quant_q_buf`/`quant_scale_qkv` and +// bmm-scale outputs are unused when `kOutputFp8Q == false`. template __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* q_ptr, T* q_pe, T* k_ptr, T const* fuse_buf, KVCacheBuffer kv_cache, int q_pe_ld, int q_pe_stride, float2 const* cos_sin_cache, size_t head_num, int head_size, int c_k, int* cu_q_seqlens, int32_t const* kv_cache_lengths, uint32_t max_input_seq_len, KvCacheDataType cache_type, - float const* quant_scale_kv, int32_t const* helix_position_offsets, bool absorption_mode, + float const* quant_scale_kv, int32_t const* helix_position_offsets, bool absorption_mode, bool q_rope_applied, __nv_fp8_e4m3* quant_q_buf = nullptr, float const* quant_scale_qkv = nullptr, float* bmm1_scale_out = nullptr, float* bmm2_scale_out = nullptr, float const* dequant_scale_q = nullptr, float const* dequant_scale_kv = nullptr, float const* quant_scale_o = nullptr, float host_bmm1_scale = 1.0f) @@ -501,19 +502,24 @@ __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* q_ptr, T* q_pe, T* k float2 const* rotary_coef_cache_buffer = cos_sin_cache + static_cast(ROPE_DIM) * position_id + (head_dim_idx / 2); - VecT q, k; + // When q_rope_applied is true, Q is already complete in the FP8 + // output buffer. Zero-initialize the unused Q fragment so the + // shared GPT-J helper can still rotate K without reading q_pe. + VecT q{}, k; auto const src_k_global_offset = static_cast(global_token_idx) * (c_k + ROPE_DIM) + c_k; - auto src_q_global_offset = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) - + (head_size + ROPE_DIM) * head_idx + head_size; - // In the absorption mode, we load pe from q_pe instead of q_ptr. - T* q_pe_input = q_ptr; - if (absorption_mode) + if (!q_rope_applied) { - q_pe_input = q_pe; - src_q_global_offset = static_cast(global_token_idx) * q_pe_stride + q_pe_ld * head_idx; + auto src_q_global_offset = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + + (head_size + ROPE_DIM) * head_idx + head_size; + // In absorption mode, load the positional segment from q_pe. + T* q_pe_input = q_ptr; + if (absorption_mode) + { + q_pe_input = q_pe; + src_q_global_offset = static_cast(global_token_idx) * q_pe_stride + q_pe_ld * head_idx; + } + q = *reinterpret_cast(&q_pe_input[src_q_global_offset + head_dim_idx]); } - - q = *reinterpret_cast(&q_pe_input[src_q_global_offset + head_dim_idx]); k = *reinterpret_cast(&fuse_buf[src_k_global_offset + head_dim_idx]); // Pack two elements into one for gptj rotary embedding. @@ -548,14 +554,17 @@ __global__ void applyMLARopeAndAssignQKVKernelOptContext(T* q_ptr, T* q_pe, T* k + head_idx * (nope_head_size_q + ROPE_DIM) + nope_head_size_q + head_dim_idx; auto const dst_k_idx = static_cast(global_token_idx) * head_num * (head_size + ROPE_DIM) + head_idx * (head_size + ROPE_DIM) + head_size + head_dim_idx; - if constexpr (kOutputFp8Q) + if (!q_rope_applied) { - quantCopy( - quant_q_buf + dst_q_idx, reinterpret_cast(&q), quant_scale_qkv_val); - } - else - { - reinterpret_cast(q_ptr)[dst_q_idx / ELTS_PER_VEC] = q; + if constexpr (kOutputFp8Q) + { + quantCopy( + quant_q_buf + dst_q_idx, reinterpret_cast(&q), quant_scale_qkv_val); + } + else + { + reinterpret_cast(q_ptr)[dst_q_idx / ELTS_PER_VEC] = q; + } } // Only write to k_pe to k_buf in the non-absorption mode. if (!absorption_mode) @@ -626,7 +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 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; @@ -645,7 +655,9 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, constexpr auto TOTAL_VEC_PER_HEAD = VECS_PER_HEAD + K_VECS_PER_HEAD; // Block/Head idx. - size_t const head_idx = blockIdx.y; + // Remap a compact pre-rotated-Q launch directly onto the K-RoPE and + // KV-copy work, skipping the Q-RoPE and Q-nope quantization blocks. + size_t const head_idx = q_rope_applied ? blockIdx.y + head_num : blockIdx.y; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif @@ -743,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 position. With speculative verify groups the per-token + // slot table decides (a group can straddle a page boundary onto 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)); @@ -806,7 +826,7 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, } else if (head_idx <= head_num + 8) { - int block_dim = gridDim.y - head_num - 1; + constexpr int block_dim = 8; int block_id = head_idx - head_num - 1; size_t const head_dim_vec_idx = (threadIdx.x % K_VECS_PER_HEAD); size_t const head_dim_idx = head_dim_vec_idx * ELTS_PER_VEC; @@ -834,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 position (per-token slots for speculative verify + // groups, 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); { @@ -1720,13 +1747,13 @@ void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, c { if (useFusedFp8Q) { - applyMLARopeAndAssignQKVKernelOptContext - <<>>(params.q_buf, params.q_pe, params.k_buf, params.latent_cache, - kv_cache_buffer, params.q_pe_ld, params.q_pe_stride, params.cos_sin_cache, params.head_num, - head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, - params.max_input_seq_len, params.cache_type, params.quant_scale_kv, params.helix_position_offsets, - params.absorption_mode, quant_q_fp8, params.quant_scale_qkv, params.bmm1_scale, params.bmm2_scale, - params.dequant_scale_q, params.dequant_scale_kv, params.quant_scale_o, params.host_bmm1_scale); + applyMLARopeAndAssignQKVKernelOptContext<<>>( + params.q_buf, params.q_pe, params.k_buf, params.latent_cache, kv_cache_buffer, params.q_pe_ld, + params.q_pe_stride, params.cos_sin_cache, params.head_num, head_size, params.meta.kv_lora_rank, + params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, + params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode, params.q_rope_applied, + quant_q_fp8, params.quant_scale_qkv, params.bmm1_scale, params.bmm2_scale, params.dequant_scale_q, + params.dequant_scale_kv, params.quant_scale_o, params.host_bmm1_scale); } else { @@ -1734,7 +1761,7 @@ void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, c params.q_buf, params.q_pe, params.k_buf, params.latent_cache, kv_cache_buffer, params.q_pe_ld, params.q_pe_stride, params.cos_sin_cache, params.head_num, head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, - params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode); + params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode, params.q_rope_applied); } } else @@ -1781,13 +1808,13 @@ void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, c if (useFusedFp8Q) { - applyMLARopeAndAssignQKVKernelOptContext - <<>>(params.q_buf, params.q_pe, params.k_buf, params.latent_cache, - kv_cache_buffer, params.q_pe_ld, params.q_pe_stride, params.cos_sin_cache, params.head_num, - head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, - params.max_input_seq_len, params.cache_type, params.quant_scale_kv, params.helix_position_offsets, - params.absorption_mode, quant_q_fp8, params.quant_scale_qkv, params.bmm1_scale, params.bmm2_scale, - params.dequant_scale_q, params.dequant_scale_kv, params.quant_scale_o, params.host_bmm1_scale); + applyMLARopeAndAssignQKVKernelOptContext<<>>( + params.q_buf, params.q_pe, params.k_buf, params.latent_cache, kv_cache_buffer, params.q_pe_ld, + params.q_pe_stride, params.cos_sin_cache, params.head_num, head_size, params.meta.kv_lora_rank, + params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, + params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode, params.q_rope_applied, + quant_q_fp8, params.quant_scale_qkv, params.bmm1_scale, params.bmm2_scale, params.dequant_scale_q, + params.dequant_scale_kv, params.quant_scale_o, params.host_bmm1_scale); } else { @@ -1795,7 +1822,7 @@ void invokeMLARopeContext(MlaParams& params, KVCacheBuffer kv_cache_buffer, c params.q_buf, params.q_pe, params.k_buf, params.latent_cache, kv_cache_buffer, params.q_pe_ld, params.q_pe_stride, params.cos_sin_cache, params.head_num, head_size, params.meta.kv_lora_rank, params.cu_q_seqlens, params.cache_seq_lens, params.max_input_seq_len, params.cache_type, - params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode); + params.quant_scale_kv, params.helix_position_offsets, params.absorption_mode, params.q_rope_applied); } } } @@ -1906,8 +1933,11 @@ void invokeMLARopeGeneration(MlaParams& params, KVCacheBuffer kv_cache_buffer bool const useFusedFp8Q = params.fuse_q_fp8_in_rope && params.cache_type == KvCacheDataType::FP8 && params.quant_q_buf != nullptr && params.quant_scale_qkv != nullptr; - dim3 grid(int(tensorrt_llm::common::divUp(params.acc_q_len, 32)), params.head_num + 1 + 8); - if ((params.cache_type == KvCacheDataType::FP8 && !useFusedFp8Q) || params.cache_type == KvCacheDataType::NVFP4) + dim3 grid(int(tensorrt_llm::common::divUp(params.acc_q_len, 32)), + params.q_rope_applied ? 1 + 8 : params.head_num + 1 + 8); + if (!params.q_rope_applied + && ((params.cache_type == KvCacheDataType::FP8 && !useFusedFp8Q) + || params.cache_type == KvCacheDataType::NVFP4)) grid.y += params.head_num * 8; TLLM_CHECK_WITH_INFO(params.acc_q_len % params.batch_size == 0, "MLA can only support input sequences with the same sequence length."); @@ -1946,7 +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.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 6a71f4369ca3..8bf1f1497bae 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -145,6 +145,9 @@ struct MlaParams // `latent_cache` row stride in elements; the fused path passes a slice of // kv_a_proj, so rows are wider than packed. 0 means packed. int latent_row_stride = 0; + // The caller has already applied Q RoPE and written the complete FP8 Q to + // `quant_q_buf`. Context preprocessing must still rotate/cache K. + bool q_rope_applied = false; // DSv4 fused inverse-RoPE + FP8 quant epilogue parameters. Dsv4EpilogueFusionParams dsv4_epilogue_fusion; @@ -155,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/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 6c2e03c6415e..5fefa7660b6a 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -311,8 +311,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); - std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + std::tie(options, optionsFromArgs, ctaDim) + = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); // Check if the options are valid or not. checkFmhaOptions(options, optionsFromArgs); @@ -409,8 +409,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); - std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + std::tie(options, optionsFromArgs, ctaDim) + = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); checkFmhaOptions(options, optionsFromArgs); updateFmhaOptions(options, optionsFromArgs); @@ -527,8 +527,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); - std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + std::tie(options, optionsFromArgs, ctaDim) + = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); // Overwrite AutoTuner decision: SageAttention with SfsPV is known to cause regression to persistent scheduler. // Remove this overwritten once we refresh the cubin kernels that containing the related fix. @@ -663,6 +663,43 @@ class TllmGenFmhaKernel } private: + std::tuple selectKernelWithCgaSmemReductionLimit( + FmhaOptions options, FmhaOptionsFromArgs optionsFromArgs, int32_t multiProcessorCount) const + { + if (isGmemReduction(options.mMultiCtasKvMode) && options.mTileScheduler == TileScheduler::Static) + { + constexpr int kMaxCgaClusterDimX = 16; + constexpr int kMinPreSelectCgaClusterDimX = 2; + // selectKernel may promote this candidate to CgaSmemReduction and call computeNumCtas + // before returning the selected options. Keep the candidate legal for both 1-CTA and + // 2-CTA FMHA kernels; the exact selected clusterDimX is applied below. + options.mMaxNumCtasPerSeqKv = std::min(options.mMaxNumCtasPerSeqKv, + kMaxCgaClusterDimX / std::max(options.mClusterDimX, kMinPreSelectCgaClusterDimX)); + } + + int32_t ctaDim = 512; + FmhaAutoTuner autoTuner(options, optionsFromArgs, multiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + limitCgaSmemReductionCtasKv(options); + return {options, optionsFromArgs, ctaDim}; + } + + void limitCgaSmemReductionCtasKv(FmhaOptions& options) const + { + if (!isCgaSmemReduction(options.mMultiCtasKvMode)) + { + return; + } + constexpr int kMaxCgaClusterDimX = 16; + if (options.mClusterDimX * options.mMaxNumCtasPerSeqKv > kMaxCgaClusterDimX) + { + TLLM_LOG_WARNING( + "CGA reduction is not supported when numCtasPerSeqKv * clusterDimX > 16. Set mMaxNumCtasPerSeqKv to " + "16 / clusterDimX"); + options.mMaxNumCtasPerSeqKv = kMaxCgaClusterDimX / options.mClusterDimX; + } + } + inline uint64_t hashID(int qkvLayout, int maskType, int kernelType, int scheduler, int multiCtasKvMode, int headDimPerCtaV, int headDimQk, int headDimV, int tileSizeQ, int tileSizeKv, int numTokensPerPage, bool reuseSmemKForV, bool uses2CtaMma, int sparseAttention, bool skipsSoftmax, bool groupsHeadsQ, diff --git a/cpp/tensorrt_llm/thop/alltoallOp.cpp b/cpp/tensorrt_llm/thop/alltoallOp.cpp index 8a775b1cf6b8..9c076850d29f 100644 --- a/cpp/tensorrt_llm/thop/alltoallOp.cpp +++ b/cpp/tensorrt_llm/thop/alltoallOp.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/thop/thUtils.h" +#include #include TRTLLM_NAMESPACE_BEGIN @@ -125,10 +126,13 @@ std::vector alltoall_helix( * @param workspace Workspace tensor (uint64, strided across ranks) * @param cp_rank Current context parallel rank * @param cp_size Total number of context parallel ranks + * @param zero_kv_mask Optional bool mask of rows this rank owns no KV for. The + * sender replaces them with a no-op contribution, so the caller must not + * sanitize itself. Its length must divide entry_count. * @return tuple of (partial_o_out, softmax_stats_out) with same shapes as inputs */ -std::tuple alltoall_helix_native( - torch::Tensor partial_o, torch::Tensor softmax_stats, torch::Tensor workspace, int64_t cp_rank, int64_t cp_size) +std::tuple alltoall_helix_native(torch::Tensor partial_o, torch::Tensor softmax_stats, + torch::Tensor workspace, int64_t cp_rank, int64_t cp_size, std::optional zero_kv_mask) { // Input validation @@ -224,6 +228,21 @@ std::tuple alltoall_helix_native( params.channelCount = 0; // auto-compute params.maxChannelCount = tensorrt_llm::kernels::computeHelixMaxChannelCount(cp_size); + // Optional zero-local-KV sanitization, applied by the sender in shared memory + params.zeroKvMask = nullptr; + params.zeroKvMaskDivisor = 1; + if (zero_kv_mask.has_value()) + { + auto const& mask = zero_kv_mask.value(); + CHECK_TH_CUDA(mask); + CHECK_CONTIGUOUS(mask); + CHECK_TYPE(mask, at::ScalarType::Bool); + TORCH_CHECK(mask.numel() > 0 && entry_count % mask.numel() == 0, "zero_kv_mask numel (", mask.numel(), + ") must divide the all-to-all entry count (", entry_count, ")"); + params.zeroKvMask = reinterpret_cast(mask.data_ptr()); + params.zeroKvMaskDivisor = entry_count / mask.numel(); + } + // Launch kernel auto stream = at::cuda::getCurrentCUDAStream(); tensorrt_llm::kernels::launchHelixAllToAll(params, allowVariableField1, stream); @@ -260,7 +279,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def("alltoall_helix(Tensor[] input_list, int[] group, int? num_lists) -> Tensor[]"); m.def( "alltoall_helix_native(Tensor partial_o, Tensor softmax_stats, Tensor(a!) workspace, int " - "cp_rank, int cp_size) -> (Tensor, Tensor)"); + "cp_rank, int cp_size, Tensor? zero_kv_mask=None) -> (Tensor, Tensor)"); m.def( "initialize_helix_workspace(Tensor(a!) workspace, int cp_rank, int cp_size) " "-> ()"); diff --git a/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp b/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp index 64552c0e973c..c10e76774030 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; @@ -90,6 +92,7 @@ struct MlaRopeGenArgs // `kv_only` launches the KV half, `kv_done_elsewhere` the Q half. bool kv_only; bool kv_done_elsewhere; + bool q_rope_applied; }; template @@ -133,6 +136,8 @@ void invokeMLARopeGenerationHelper(T const* latent_cache_ptr, T* q_pe_ptr, T* fu mla_params.host_bmm1_scale = args.host_bmm1_scale; mla_params.helix_position_offsets = args.helix_position_offsets_ptr; mla_params.helix_is_inactive_rank = args.helix_is_inactive_rank_ptr; + mla_params.helix_local_slots = args.helix_local_slots_ptr; + mla_params.q_rope_applied = args.q_rope_applied; mla_params.precomputed_cu_seqlens = args.precomputed_cu_seqlens; mla_params.precomputed_fmha_scheduler = args.precomputed_fmha_scheduler; @@ -175,7 +180,7 @@ void MLARopeGeneration(std::optional fused_q, // [tokens, num_hea int64_t qk_nope_head_dim, int64_t qk_rope_head_dim, int64_t v_head_dim, bool rope_append, std::optional kv_norm_weight, double const kv_norm_eps, bool const precomputed_cu_seqlens, bool const precomputed_fmha_scheduler, bool const kv_only, bool const kv_done_elsewhere, - std::optional quant_scale_qkv) + std::optional quant_scale_qkv, bool q_rope_applied) { // `kv_only` runs before q_pe exists, so the Q tensors are absent. TORCH_CHECK(kv_only || (fused_q.has_value() && q_pe.has_value()), @@ -186,10 +191,11 @@ void MLARopeGeneration(std::optional fused_q, // [tokens, num_hea TLLM_CHECK_WITH_INFO( head_size == kv_lora_rank + qk_rope_head_dim, "head_size must = kv_lora_rank + qk_rope_head_dim"); 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, - "MLA KV residual_dim must be 0 or qk_rope_head_dim (%ld), got %ld", qk_rope_head_dim, residual_dim); - TORCH_CHECK(helix_tensor_params.size() == 2, - "Expecting 2 tensors for helix_tensor_params: helix_position_offsets and helix_is_inactive_rank."); + 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 || 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)); @@ -219,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()); @@ -328,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}; + 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(); @@ -414,6 +428,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) ", bool kv_only=False" ", bool kv_done_elsewhere=False" ", Tensor? quant_scale_qkv=None" + ", bool q_rope_applied=False" ") -> ()"); } diff --git a/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp b/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp index 8f955399a60e..1d0a3c307279 100644 --- a/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp +++ b/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/attentionWorkspace.h" +#include "tensorrt_llm/common/attentionOp.h" #include "tensorrt_llm/common/workspace.h" @@ -56,8 +57,73 @@ void expectNextSlice(char const* name, Slice const& slice, size_t size, size_t& expectedOffset += tc::alignSize(size, kAlignment); } +constexpr int32_t kBatchSize = 2; +constexpr int32_t kInputSequenceLength = 11; +constexpr int32_t kCrossKvLength = 7; +constexpr int32_t kPackedTokenCount = 14; +constexpr int32_t kHeadSize = 8; + +void configureUnfusedAttention(tcop::AttentionOp& op, bool crossAttention) +{ + op.mNumHeads = 1; + op.mNumKVHeads = 1; + op.mHeadSize = kHeadSize; + op.mNumAttnHeads = 1; + op.mNumAttnKVHeads = 1; + op.mEnableContextFMHA = false; + op.mCrossAttention = crossAttention; +} + +size_t expectedUnfusedContextWorkspace(bool crossAttention) +{ + constexpr size_t kElementSize = sizeof(half); + size_t const batchSize = kBatchSize; + size_t const inputSequenceLength = kInputSequenceLength; + size_t const kvSequenceLength = crossAttention ? kCrossKvLength : kInputSequenceLength; + size_t const paddedTokenCount = batchSize * inputSequenceLength; + size_t const paddedKvTokenCount = batchSize * kvSequenceLength; + + tcop::AttentionContextWorkspaceSizes sizes{}; + sizes.attentionMask = kElementSize * paddedTokenCount * kvSequenceLength; + sizes.cuQSeqlens = sizeof(int) * (batchSize + 1); + sizes.cuKvSeqlens = sizes.cuQSeqlens; + sizes.cuMaskRows = sizes.cuQSeqlens; + sizes.qBuf = kElementSize * paddedTokenCount * kHeadSize; + sizes.kBuf = kElementSize * paddedKvTokenCount * kHeadSize; + sizes.vBuf = sizes.kBuf; + sizes.qkBuf = kElementSize * batchSize * inputSequenceLength * kvSequenceLength; + sizes.qkvBuf = kElementSize * paddedTokenCount * kHeadSize; + sizes.qkFloatBuf = sizeof(float) * batchSize * inputSequenceLength * kvSequenceLength; + sizes.paddingOffset = sizeof(int) * paddedTokenCount; + sizes.encoderPaddingOffset = sizeof(int) * paddedKvTokenCount; + sizes.tokensInfo = sizeof(int2) * kPackedTokenCount; + return tcop::AttentionWorkspaceManager::buildContextLayout(sizes).totalSize; +} + +size_t getUnfusedContextWorkspace(tcop::AttentionOp const& op) +{ + return op.getWorkspaceSizeForContext( + nvinfer1::DataType::kHALF, kBatchSize, kInputSequenceLength, kCrossKvLength, kPackedTokenCount); +} + } // namespace +TEST(AttentionWorkspaceManagerTest, RaggedUnfusedSelfAttentionUsesPaddedTokenCounts) +{ + tcop::AttentionOp op; + configureUnfusedAttention(op, false); + + EXPECT_EQ(getUnfusedContextWorkspace(op), expectedUnfusedContextWorkspace(false)); +} + +TEST(AttentionWorkspaceManagerTest, RaggedUnfusedCrossAttentionUsesPaddedTokenCounts) +{ + tcop::AttentionOp op; + configureUnfusedAttention(op, true); + + EXPECT_EQ(getUnfusedContextWorkspace(op), expectedUnfusedContextWorkspace(true)); +} + TEST(AttentionWorkspaceManagerTest, ContextLayoutMatchesAttentionOpOrdering) { tcop::AttentionContextWorkspaceSizes sizes{}; diff --git a/tensorrt_llm/_ipc_utils.py b/tensorrt_llm/_ipc_utils.py index 4d5ecefc9a9b..3178597b5087 100644 --- a/tensorrt_llm/_ipc_utils.py +++ b/tensorrt_llm/_ipc_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import array +import ctypes import struct import sys from typing import List, Tuple @@ -26,6 +27,9 @@ from .logger import logger from .mapping import Mapping +# Must match CUDA_IPC_HANDLE_SIZE / CU_IPC_HANDLE_SIZE. +_IPC_MEM_HANDLE_SIZE = 64 + def _raise_if_error(error: cudart.cudaError_t | cuda.CUresult): if isinstance(error, cudart.cudaError_t): @@ -36,6 +40,33 @@ def _raise_if_error(error: cudart.cudaError_t | cuda.CUresult): raise RuntimeError(f"CUDA Driver API error: {repr(error)}") +def _ipc_mem_handle_to_bytes(handle) -> bytes: + """Serialize a cudaIpcMemHandle across cuda-python / cuda-bindings versions.""" + reserved = getattr(handle, "reserved", None) + if reserved is not None: + return bytes(reserved) + # Some cuda-bindings builds omit the .reserved attribute on cudaIpcMemHandle_t. + return ctypes.string_at(handle.getPtr(), _IPC_MEM_HANDLE_SIZE) + + +def _ipc_mem_handle_from_bytes(data: bytes): + """Rebuild a cudaIpcMemHandle from the opaque 64-byte payload.""" + handle = cudart.cudaIpcMemHandle_t() + if hasattr(handle, "reserved"): + try: + handle.reserved = data + except TypeError: + # Older bindings expect list[int] rather than bytes. + handle.reserved = list(data) + else: + if len(data) != _IPC_MEM_HANDLE_SIZE: + raise ValueError( + f"Invalid CUDA IPC mem handle size: {len(data)} (expected {_IPC_MEM_HANDLE_SIZE})" + ) + ctypes.memmove(handle.getPtr(), data, _IPC_MEM_HANDLE_SIZE) + return handle + + def can_access_peer(mapping: Mapping) -> bool: src_node = mapping.local_rank @@ -117,13 +148,8 @@ def align_size(size, alignment): _raise_if_error(cudart.cudaMemset(local_ptr, 0, aligned_size)[0]) error, local_handle = cudart.cudaIpcGetMemHandle(local_ptr) _raise_if_error(error) - handles_reserved = dist.tp_allgather(local_handle.reserved) - - handles = [] - for reserved in handles_reserved: - handle = cudart.cudaIpcMemHandle_t() - handle.reserved = reserved - handles.append(handle) + handles_reserved = dist.tp_allgather(_ipc_mem_handle_to_bytes(local_handle)) + handles = [_ipc_mem_handle_from_bytes(reserved) for reserved in handles_reserved] peer_ptrs = [] for node, handle in enumerate(handles): diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index c5dfe5585fbc..4b2769391270 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -142,13 +142,14 @@ statically checkable without runtime signature inspection. Ordinary sparse variants use `attention_output_hidden_size` and the shared output allocation. DeepSeek-V4's fused epilogue instead uses the optional -output-preparation hook to create one token-major O-LoRA output tensor. Its -context- and generation-phase helpers allocate the private FP8 attention and -scale buffers, then write the O-LoRA result into the corresponding token range. -The shared MLA custom-op contract exposes exactly one mutable output tensor; -`_create_outputs()` keeps that tensor in a single-entry list through forward -and output projection. Phase-specific scratch buffers remain inside the -DeepSeek-V4 algorithm module and do not widen the generic hook facade. +output-preparation hook to create one token-major O-LoRA-sized output tensor. +Its context- and generation-phase helpers allocate the private FP8 attention +and scale buffers, run both O-LoRA projections, and write the final hidden +states into the leading columns of that tensor. The shared MLA custom-op +contract exposes exactly one mutable output tensor; `_create_outputs()` keeps +that tensor in a single-entry list through forward and output projection. +Phase-specific scratch buffers remain inside the DeepSeek-V4 algorithm module +and do not widen the generic hook facade. Sparse prediction inputs stay out of shared MLA APIs. Algorithm modules wrap their module-to-backend inputs in a `SparseBackendForwardArgs` subclass and diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py index 4ede4dfac0b6..995fc1c96ac8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py @@ -4,6 +4,7 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING, Optional, Tuple import torch @@ -129,11 +130,10 @@ def _apply_q_rope(self, q: torch.Tensor, position_ids: torch.Tensor) -> torch.Te ) return q - def _project_and_quantize_q( - self, qr: torch.Tensor, position_ids: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Project and quantize Q, using the fused MXFP4 path when supported.""" - use_fused_project_mxfp4 = ( + def _is_fused_project_mxfp4_enabled(self, input_dtype: torch.dtype) -> bool: + if os.environ.get("TRTLLM_DISABLE_DSA_FUSED_INDEXER_Q", "0") == "1": + return False + return ( self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE and not HAS_FAST_HADAMARD and not self.rotary_emb.is_neox @@ -145,9 +145,15 @@ def _project_and_quantize_q( torch.ops.trtllm, "cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell", ) - and qr.dtype == torch.bfloat16 + and input_dtype == torch.bfloat16 and is_sm_100f() ) + + def _project_and_quantize_q( + self, qr: torch.Tensor, position_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Project and quantize Q, using the fused MXFP4 path when supported.""" + use_fused_project_mxfp4 = self._is_fused_project_mxfp4_enabled(qr.dtype) if use_fused_project_mxfp4: q_fp4, q_scale = torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell( qr, @@ -226,6 +232,7 @@ def precompute_aux( self, hidden_states: torch.Tensor, metadata: DeepseekV4TrtllmAttentionMetadata, + start_event: Optional[torch.cuda.Event] = None, ) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]]: """Pre-launch the qr-independent half of the indexer prepare phase. @@ -236,14 +243,17 @@ def precompute_aux( prepare path skip its own aux-stream launch and consume these results directly. - Returns ``None`` when multi-stream mode is off (caller should fall - back to the normal ``forward()`` call without ``pre_aux``). + Returns ``None`` when multi-stream mode is off or fused Indexer-Q + requires serial prepare (caller should fall back to the normal + ``forward()`` call without ``pre_aux``). """ if not (do_multi_stream() and self.aux_stream is not None): return None - self.indexer_start_event.record() + if start_event is None: + start_event = self.indexer_start_event + start_event.record() with torch.cuda.stream(self.aux_stream): - self.indexer_start_event.wait() + start_event.wait() weights = self.weights_proj(hidden_states) self.weights_proj_event.record() k_fp8, k_scale = self.compressor(hidden_states, metadata) @@ -308,16 +318,16 @@ def _run_overlapped_indexer_prepare( self.k_cache_update_event.record() else: weights, k_fp8, k_scale = pre_aux - # pre_aux tensors were allocated on aux_stream; record on the - # consuming stream so the caching allocator can't recycle them mid-use. - cur_stream = torch.cuda.current_stream() - weights.record_stream(cur_stream) - if k_fp8 is not None: - k_fp8.record_stream(cur_stream) - if k_scale is not None: - k_scale.record_stream(cur_stream) q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) + # Aux-stream tensors are consumed on the current stream after waits. + cur_stream = torch.cuda.current_stream() + weights.record_stream(cur_stream) + if k_fp8 is not None: + k_fp8.record_stream(cur_stream) + if k_scale is not None: + k_scale.record_stream(cur_stream) + self.weights_proj_event.wait() weights = self._apply_weight_scale(weights, q_scale) @@ -363,7 +373,12 @@ def forward( Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] ] = None, ): - if do_multi_stream() and self.aux_stream is not None: + use_overlapped_prepare = ( + do_multi_stream() + and self.aux_stream is not None + and (pre_aux is not None or not self._is_fused_project_mxfp4_enabled(qr.dtype)) + ) + if use_overlapped_prepare: q_fp8, q_scale, k_fp8, k_scale, weights = self._run_overlapped_indexer_prepare( qr, hidden_states, @@ -372,7 +387,7 @@ def forward( pre_aux=pre_aux, ) else: - assert pre_aux is None, "pre_aux requires multi-stream mode" + assert pre_aux is None, "pre_aux requires the overlapped indexer prepare path" q_fp8, q_scale, k_fp8, k_scale, weights = self._run_serial_indexer_prepare( qr, hidden_states, metadata, position_ids ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/params.py b/tensorrt_llm/_torch/attention/backends/sparse/params.py index 8e6829f8c879..e21af93568c0 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/params.py @@ -14,11 +14,34 @@ # limitations under the License. """Shared sparse attention parameter types.""" +import os from dataclasses import dataclass from typing import Literal, Optional import torch +_INDEXER_MQA_LOGITS_DEFAULT_ELEM_BUDGET = 1 << 31 +_INDEXER_MQA_LOGITS_BYTES_PER_ELEMENT = 4 + + +def get_indexer_mqa_logits_elem_budget() -> int: + """Return the per-call Indexer MQA-logits cap used by the runtime.""" + return int( + os.environ.get( + "TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET", _INDEXER_MQA_LOGITS_DEFAULT_ELEM_BUDGET + ) + ) + + +def get_indexer_mqa_logits_workspace_bytes( + max_num_tokens: Optional[int] = None, max_seq_len: Optional[int] = None +) -> int: + """Return the reachable maximum bytes for one FP32 MQA-logits tile.""" + elem_budget = get_indexer_mqa_logits_elem_budget() + if max_num_tokens is not None and max_seq_len is not None: + elem_budget = min(elem_budget, max_num_tokens * max_seq_len) + return elem_budget * _INDEXER_MQA_LOGITS_BYTES_PER_ELEMENT + class SparseParams: """Base parameters for a sparse attention backend.""" diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index d41ae7c5b066..02d934b8b5dc 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1267,7 +1267,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 0e02a8759ab4..5273a9fdd83c 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -42,6 +42,58 @@ # Torch schema parsing rejects ``inf`` as a default value. SWIGLU_LIMIT_SCALAR_DISABLED = -1.0 +_KIMI_K3_MXFP8_TUNING_BUCKETS = (1, 2, 4, 8, *range(16, 193, 16)) + +# Both dense MXFP8 wrappers pass a scalar one to the CuTe kernel. Keep one +# read-only tensor per CUDA device so steady-state calls, including CUDA graph +# replay, do not launch a scalar fill kernel. +_MXFP8_GEMM_ALPHA_CACHE: "dict[torch.device, torch.Tensor]" = {} + + +def _get_mxfp8_gemm_alpha(device: torch.device) -> torch.Tensor: + """Return the cached FP32 scalar one for ``device``. + + Allocation is deliberately rejected during CUDA graph capture. Callers + must run one eager warmup on each device, which is already required for + GEMM autotuning, before capturing the steady-state path. + """ + device = torch.device(device) + if device.type != "cuda": + raise ValueError( + f"MXFP8 GEMM alpha requires a CUDA device, got {device}.") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + + alpha = _MXFP8_GEMM_ALPHA_CACHE.get(device) + if alpha is None: + with torch.cuda.device(device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MXFP8 GEMM alpha cache must be initialized before CUDA graph " + f"capture on {device}; run one eager GEMM warmup first.") + alpha = torch.ones((), dtype=torch.float32, device=device) + _MXFP8_GEMM_ALPHA_CACHE[device] = alpha + return alpha + + +def _get_kimi_k3_mxfp8_tuning_buckets(max_num_tokens: int) -> Tuple[int, ...]: + """Generate every K3 bucket used through ``max_num_tokens``.""" + max_bucket = _kimi_k3_mxfp8_tuning_bucket(max_num_tokens) + low_m = [m for m in _KIMI_K3_MXFP8_TUNING_BUCKETS if m <= max_bucket] + high_m = [ + m for m in get_last_power_of_2_num_tokens_buckets(max_bucket) if m > 192 + ] + return tuple((*low_m, *high_m)) + + +def _kimi_k3_mxfp8_tuning_bucket(num_tokens: int) -> int: + """Use lower power-of-two buckets for small M, then upper-bound buckets.""" + if num_tokens <= 32: + return last_positive_power_of_2(num_tokens) + if num_tokens <= 192: + return next(m for m in _KIMI_K3_MXFP8_TUNING_BUCKETS if num_tokens <= m) + return next_positive_power_of_2(num_tokens) + def _with_input_cuda_device(function): """Run a custom-op implementation under its input tensor's CUDA device.""" @@ -13291,7 +13343,7 @@ def cute_dsl_mxfp8_gemm_rubin( if input_scale.dtype != torch.uint8 or weight_scale.dtype != torch.uint8: raise ValueError("CuteDSL MXFP8 scales must be UE8M0 uint8") - alpha = torch.ones((), dtype=torch.float32, device=input.device) + alpha = _get_mxfp8_gemm_alpha(input.device) runner = CuteDSLMXFP8RubinLinear(output_dtype=output_dtype, use_tvm_ffi=use_tvm_ffi) inputs = [input, weight, input_scale, weight_scale, alpha] diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py index b31ba370a78c..171cd4ceb615 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py @@ -901,13 +901,17 @@ def _chunk_kda_fwd( # Varlen with non-aligned seq lengths is a separate problem (Phase 2): # multi-seq varlen can't be host-padded without repacking memory. real_T = T - needs_eqlen_pad = (not is_varlen) and (T % BT != 0) + CPB_BT = 4 * BT # CHUNKS_PER_BLOCK * BT, the cgs_per_head divisibility unit + # The trigger is CPB_BT, not BT: the eqlen scheduler floors `cgs_per_head` + # and has no `chunk_idx < num_chunks` guard, so a BT-aligned T whose chunk + # count is not a multiple of CHUNKS_PER_BLOCK (T = 64, 128, 192, 320, ...) + # would silently drop its trailing chunks. + needs_eqlen_pad = (not is_varlen) and (T % CPB_BT != 0) if needs_eqlen_pad: if B != 1: raise NotImplementedError( - f"eqlen with B>1 and T % {BT} != 0 not supported (got B={B}, T={T})." + f"eqlen with B>1 and T % {CPB_BT} != 0 not supported (got B={B}, T={T})." ) - CPB_BT = 4 * BT # CHUNKS_PER_BLOCK * BT, the cgs_per_head divisibility unit T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT # Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache # key). torch.cat would reallocate + copy the full 200MB q tensor every @@ -1001,16 +1005,11 @@ def _chunk_kda_fwd( if chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = len(chunk_indices) - if NT < 4: - # The persistent K123 scheduler launches NT // 4 cooperative - # groups per head; fewer than 4 total chunks produces a - # zero-size grid (DSLCudaRuntimeError at launch). Callers must - # route such batches to the FLA path — see - # KDAKernelDispatch.prefill_chunk_kda. - raise ValueError( - f"kda_prefill requires >= 4 total varlen chunks (got {NT}); " - "route small varlen batches to the FLA fallback" - ) + # Any chunk count is launchable: the varlen scheduler rounds up + # (`total_cgs_per_head = ceil(num_chunks / CHUNKS_PER_BLOCK)`), launches + # the constant `NUM_SMS` grid, and guards every chunk loop with + # `chunk_idx < num_chunks`. Only eqlen needs divisibility, and it gets + # that from its own `CPB_BT` pad above. N_seqs = len(cu_seqlens) - 1 else: NT = T // BT diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py index 6dc6becd131a..d63eba9b6e9e 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py @@ -89,17 +89,27 @@ def __extract_mlir_values__(self): return values def __new_from_mlir_values__(self, values): - problem_shape_b = cutlass.new_from_mlir_values(self.problem_shape_b, - (values[0], )) - problem_shape_s = cutlass.new_from_mlir_values(self.problem_shape_s, - (values[1], )) - split_kv = cutlass.new_from_mlir_values(self.split_kv, (values[2], )) - problem_shape_b_fdd = cutlass.new_from_mlir_values( - self.problem_shape_b_fdd, (values[3], )) - problem_shape_s_fdd = cutlass.new_from_mlir_values( - self.problem_shape_s_fdd, (values[4], )) - split_kv_fdd = cutlass.new_from_mlir_values(self.split_kv_fdd, - (values[5], )) + # Slice per field by that field's own extraction width, in + # __extract_mlir_values__ order. FastDivmodDivisor occupies 2 SSA values + # on nvidia-cutlass-dsl-internal, so fixed indices mis-slice (cutlass #3243). + fields = ( + self.problem_shape_b, + self.problem_shape_s, + self.split_kv, + self.problem_shape_b_fdd, + self.problem_shape_s_fdd, + self.split_kv_fdd, + ) + rebuilt = [] + offset = 0 + for field in fields: + width = len(cutlass.extract_mlir_values(field)) + rebuilt.append( + cutlass.new_from_mlir_values( + field, tuple(values[offset:offset + width]))) + offset += width + (problem_shape_b, problem_shape_s, split_kv, problem_shape_b_fdd, + problem_shape_s_fdd, split_kv_fdd) = rebuilt return MLAStaticTileSchedulerParams( self.is_persistent, problem_shape_b, @@ -281,16 +291,21 @@ def __extract_mlir_values__(self): return values def __new_from_mlir_values__(self, values): - assert len(values) == 13 - new_params = cutlass.new_from_mlir_values(self.params, values[0:6]) - new_current_work_linear_idx = cutlass.new_from_mlir_values( - self.current_work_linear_idx, [values[6]]) - new_blk_coord = cutlass.new_from_mlir_values(self.blk_coord, - values[7:10]) - new_grid_shape = cutlass.new_from_mlir_values(self.grid_shape, - values[10:]) - return MLAStaticTileScheduler(new_params, new_current_work_linear_idx, - new_blk_coord, new_grid_shape) + # Width-driven slicing, as in MLAStaticTileSchedulerParams: params holds + # three FastDivmodDivisors, which are not 1 SSA value each on every + # cutlass-dsl build, so the old fixed 6/1/3/rest split mis-slices. + rebuilt = [] + offset = 0 + for field in (self.params, self.current_work_linear_idx, self.blk_coord, + self.grid_shape): + width = len(cutlass.extract_mlir_values(field)) + rebuilt.append( + cutlass.new_from_mlir_values( + field, list(values[offset:offset + width]))) + offset += width + assert offset == len(values), ( + f"MLAStaticTileScheduler consumed {offset} of {len(values)} values") + return MLAStaticTileScheduler(*rebuilt) def create_mla_static_tile_scheduler( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py index 69984b7edc3c..2852baa5d745 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py @@ -47,10 +47,15 @@ griddepcontrol_launch_dependents, griddepcontrol_wait, is_power_of_2, + sigmoid_f32, silu_f32, ) -SUPPORTED_ACTIVATION_TYPES = (ActivationType.Swiglu, ActivationType.Relu2) +SUPPORTED_ACTIVATION_TYPES = ( + ActivationType.Swiglu, + ActivationType.Relu2, + ActivationType.SiTu, +) def validate_activation_type(activation_type) -> ActivationType: @@ -72,6 +77,8 @@ def validate_activation_type(activation_type) -> ActivationType: Supported fused activations (selected at construction via ``activation_type``): - ActivationType.Swiglu: C = up * silu(gate), where up/gate come from interleaved weight matrix B - ActivationType.Relu2: C = relu(alpha * x)^2 + - ActivationType.SiTu: C = (beta*tanh(gate/beta)*sigmoid(gate)) * (linear_beta*tanh(up/linear_beta)), + gated like Swiglu; requires situ_beta / situ_linear_beta Any other ``ActivationType`` value raises an assertion at construction time. @@ -278,13 +285,16 @@ def __init__( swiglu_limit: cutlass.Float32 = float("inf"), use_expert_counts: bool = False, num_local_experts: int = 0, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ): """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel with gather operation and fused activation. ``activation_type`` accepts a value from ``ActivationType``; only - ``ActivationType.Swiglu`` (gated path) and ``ActivationType.Relu2`` - (non-gated path) are currently supported. + ``ActivationType.Swiglu`` (gated path), ``ActivationType.Relu2`` + (non-gated path) and ``ActivationType.SiTu`` (gated path, Kimi K3) + are currently supported. This configuration includes several key aspects: @@ -318,8 +328,15 @@ def __init__( :param topk: Number of experts selected per token (used for token ID mapping). :type topk: cutlass.Int64 :param activation_type: Fused activation. Must be ``ActivationType.Swiglu`` - (gated, default) or ``ActivationType.Relu2`` (non-gated). + (gated, default), ``ActivationType.Relu2`` (non-gated) or + ``ActivationType.SiTu`` (gated). :type activation_type: ActivationType + :param situ_beta: Gate-side SiTU constant. Required for -- and only + valid with -- ``ActivationType.SiTu``. + :type situ_beta: Optional[float] + :param situ_linear_beta: Linear-side (up) SiTU constant. Required for -- + and only valid with -- ``ActivationType.SiTu``. + :type situ_linear_beta: Optional[float] """ self.sf_vec_size = sf_vec_size @@ -415,6 +432,35 @@ def __init__( assert self.topk == 1 assert self.num_local_experts > 0 + # SiTU constants. They are per-model scalars (not per-expert), so they + # are folded at trace time -- which also means they belong in the + # caller's compiled-kernel cache key. + if self.activation_type == ActivationType.SiTu: + if situ_beta is None or situ_linear_beta is None: + raise ValueError( + "ActivationType.SiTu requires both situ_beta and " + f"situ_linear_beta, got {situ_beta} and {situ_linear_beta}." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + "SiTU beta parameters must be positive, got " + f"{situ_beta} and {situ_linear_beta}." + ) + if self.has_swiglu_limit: + # Matches MegaMoE (both backends) and DeepGEMM, which reject + # activation_clamp together with SiTU. + raise ValueError( + "ActivationType.SiTu does not support a SwiGLU clamp; " + "drop swiglu_limit for SiTU checkpoints." + ) + elif situ_beta is not None or situ_linear_beta is not None: + raise ValueError( + "situ_beta / situ_linear_beta require " + f"ActivationType.SiTu, got {self.activation_type.name}." + ) + self.situ_beta = None if situ_beta is None else float(situ_beta) + self.situ_linear_beta = None if situ_linear_beta is None else float(situ_linear_beta) + def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -2482,6 +2528,9 @@ def kernel( if cutlass.const_expr(self.activation_type == ActivationType.Swiglu): acc_vec_gate = tTR_rAcc_gate.load() self._apply_swiglu_epilogue(acc_vec_up, acc_vec_gate, alpha_val, tCompute) + elif cutlass.const_expr(self.activation_type == ActivationType.SiTu): + acc_vec_gate = tTR_rAcc_gate.load() + self._apply_situ_epilogue(acc_vec_up, acc_vec_gate, alpha_val, tCompute) elif cutlass.const_expr(self.activation_type == ActivationType.Relu2): self._apply_relu2_epilogue(acc_vec_up, alpha_val, tCompute) @@ -2770,6 +2819,93 @@ def _apply_swiglu_epilogue( acc_vec_up_alpha = fclip_xorsign(acc_vec_up_alpha, self.swiglu_limit) tCompute[i] = acc_vec_up_alpha * silu_f32(acc_vec_gate_alpha, fastmath=True) + @cute.jit + def _apply_situ_epilogue( + self, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + alpha_val, + tCompute: cute.Tensor, + ): + """SiTU (Kimi K3), matching ``kimi_k3_moe/_mlp.py::SituAndMul`` + (itself byte-identical to HF ``modeling_kimi.py``):: + + g = alpha * gate, u = alpha * up + situ_gate = beta * tanh(g / beta) * sigmoid(g) + situ_up = linear_beta * tanh(u / linear_beta) + tCompute = situ_gate * situ_up + + ``up`` and ``gate`` come from the two interleaved accumulator subtiles + loaded by the caller, same as the SwiGLU epilogue. + + There is no packed tanh, so the vectorized path uses the identity + ``tanh(z) = 2 * sigmoid(2z) - 1`` (the same one ``utils.gelu_tanh_f32`` + uses) to stay on the packed f32x2 path -- calling a scalar tanh would + force the whole loop back to scalar. The reciprocals and ``2*beta`` + factors fold at trace time because both betas are ``const_expr``:: + + beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) + = 2*beta*sigmoid((2/beta)*x) - beta + """ + beta = self.situ_beta + linear_beta = self.situ_linear_beta + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + neg_log2e_pair = (-LOG2_E, -LOG2_E) + one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) + + inv_2beta = cutlass.Float32(2.0 / beta) + two_beta = cutlass.Float32(2.0 * beta) + neg_beta = cutlass.Float32(-beta) + inv_2lbeta = cutlass.Float32(2.0 / linear_beta) + two_lbeta = cutlass.Float32(2.0 * linear_beta) + neg_lbeta = cutlass.Float32(-linear_beta) + + # sigmoid(x) = rcp(1 + exp2(-x * log2e)), shared by both cores. + def _sigmoid(p0, p1): + neg = cute.arch.mul_packed_f32x2((p0, p1), neg_log2e_pair) + e = ( + cute.math.exp2(neg[0], fastmath=True), + cute.math.exp2(neg[1], fastmath=True), + ) + d = cute.arch.add_packed_f32x2(e, one_pair) + return (cute.arch.rcp_approx(d[0]), cute.arch.rcp_approx(d[1])) + + alpha_pair = (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)) + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + g = cute.arch.mul_packed_f32x2((acc_vec_gate[i], acc_vec_gate[i + 1]), alpha_pair) + u = cute.arch.mul_packed_f32x2((acc_vec_up[i], acc_vec_up[i + 1]), alpha_pair) + + sigmoid_g = _sigmoid(g[0], g[1]) + + gs = _sigmoid(*cute.arch.mul_packed_f32x2(g, (inv_2beta, inv_2beta))) + tanh_g = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(gs, (two_beta, two_beta)), (neg_beta, neg_beta) + ) + + us = _sigmoid(*cute.arch.mul_packed_f32x2(u, (inv_2lbeta, inv_2lbeta))) + tanh_u = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(us, (two_lbeta, two_lbeta)), (neg_lbeta, neg_lbeta) + ) + + situ_gate = cute.arch.mul_packed_f32x2(tanh_g, sigmoid_g) + out_pair = cute.arch.mul_packed_f32x2(situ_gate, tanh_u) + tCompute[i] = out_pair[0] + tCompute[i + 1] = out_pair[1] + else: + inv_2beta = cutlass.Float32(2.0 / beta) + two_beta = cutlass.Float32(2.0 * beta) + beta_f32 = cutlass.Float32(beta) + inv_2lbeta = cutlass.Float32(2.0 / linear_beta) + two_lbeta = cutlass.Float32(2.0 * linear_beta) + lbeta_f32 = cutlass.Float32(linear_beta) + for i in cutlass.range_constexpr(cute.size(acc_vec_up.shape)): + g = acc_vec_gate[i] * cutlass.Float32(alpha_val) + u = acc_vec_up[i] * cutlass.Float32(alpha_val) + tanh_g = two_beta * sigmoid_f32(g * inv_2beta, fastmath=True) - beta_f32 + tanh_u = two_lbeta * sigmoid_f32(u * inv_2lbeta, fastmath=True) - lbeta_f32 + tCompute[i] = (tanh_g * sigmoid_f32(g, fastmath=True)) * tanh_u + @cute.jit def _apply_relu2_epilogue( self, @@ -3452,8 +3588,8 @@ def wrapper( """Single-B wrapper. ``l`` is the number of experts in the (sole) B tensor. ``activation_type`` - must match the one passed to ``__init__``; only ``Swiglu`` and ``Relu2`` - are supported. + must match the one passed to ``__init__``; only ``Swiglu``, ``Relu2`` + and ``SiTu`` are supported. """ is_gated = is_gated_activation(activation_type) scale_k = k // scaling_vector_size diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py index 795934f776f1..d506a8314e69 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py @@ -44,6 +44,7 @@ # This file is copied and modified from cutlass https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/core.py import ctypes +import math import os from typing import Union @@ -99,6 +100,12 @@ def __init__( assert int(self._pointer) % self._assumed_align == 0, ( f"pointer must be {self._assumed_align} bytes aligned") + def __add__(self, offset: int) -> Pointer: # type: ignore[override] + offset_bytes = offset * self._dtype.width // 8 + assumed_align = math.gcd(offset_bytes, self._assumed_align) + return _Pointer(self._pointer + offset_bytes, self._dtype, + self._addr_space, assumed_align) + def size_in_bytes(self) -> int: return ctypes.sizeof(ctypes.c_void_p(int(self._pointer))) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py index dddbb27ee0bf..033a3083101a 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py @@ -3209,6 +3209,8 @@ def __call__( self.fallback_cluster_shape_mn, preferred_max_active_clusters, fallback_max_active_clusters, + self.swizzle_size, + self.raster_order == "m", ) self.buffer_align_bytes = 1024 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dsv4_qb_fusion/kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dsv4_qb_fusion/kernel.py index 1421e192eef1..9cbf8b8ff367 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dsv4_qb_fusion/kernel.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dsv4_qb_fusion/kernel.py @@ -90,7 +90,6 @@ from __future__ import annotations import os -from functools import partial from typing import NamedTuple, Optional, Tuple import cuda.bindings.driver as cuda_driver @@ -101,12 +100,16 @@ import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils import cutlass.utils.rubin_helpers as sm107_utils +from cutlass._mlir import ir as mlir_ir +from cutlass._mlir.dialects import llvm as mlir_llvm from cutlass._mlir.dialects import nvvm -from cutlass.cute.arch import nvvm_wrappers as _arch_nvvm_wrappers +from cutlass._mlir.dialects import vector as mlir_vector +from cutlass._mlir.dialects.nvvm import FPRoundingMode from cutlass.cute.experimental import iket from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 from cutlass.cute.nvgpu.common import CacheEvictionPriority from cutlass.cute.nvgpu.tcgen05.mma import CollectorOp +from cutlass.cutlass_dsl import dsl_user_op from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from cutlass.utils.gemm import sm100 as gemm_sm100 @@ -124,22 +127,126 @@ HEAD_DIM = DEFAULT_QK_NOPE_HEAD_DIM + DEFAULT_QK_ROPE_HEAD_DIM -def _fma_packed_bf16x2_nvvm(res, src_a, src_b, src_c, *, rnd=None, ftz=None, loc=None, ip=None): - del ftz - src_a_bf16x2 = _arch_nvvm_wrappers.cvt_f32x2_bf16x2(src_a, loc=loc, ip=ip) - return _arch_nvvm_wrappers.nvvm.fma_packed_f32x2_bf16x2_f32x2_f32x2( - res, src_a_bf16x2, src_b, src_c, rnd=rnd, loc=loc, ip=ip +@dsl_user_op +def _prmt_trunc_bf16x2( + v0: cutlass.Float32, + v1: cutlass.Float32, + *, + loc=None, + ip=None, +): + """Truncate two fp32 to bf16 and pack into one bf16x2 register. + + A single PRMT selecting the high two bytes of each source + (selector 0x7632: result = [v0.b2, v0.b3, v1.b2, v1.b3]) -- the + same numerics as masking the low 16 mantissa bits (truncate + instead of rn), one instruction per pair, and no F2FP + conversion-pipe traffic. + + NOTE: builds the vector<2xbf16> via cutlass._mlir plumbing + (llvm.bitcast) -- there is no public wrapper accepting a + pre-packed bf16x2 operand; revisit when one appears.""" + packed_i32 = cute.arch.prmt( + v0.bitcast(cutlass.Int32, loc=loc, ip=ip), + v1.bitcast(cutlass.Int32, loc=loc, ip=ip), + cutlass.Int32(0x7632), + loc=loc, + ip=ip, + ) + return mlir_llvm.bitcast( + mlir_ir.VectorType.get([2], cutlass.BFloat16.mlir_type, loc=loc), + packed_i32, + loc=loc, + ip=ip, ) -_fma_packed_f32x2_bf16x2_f32x2_f32x2 = getattr( - cute.arch, - "fma_packed_f32x2_bf16x2_f32x2_f32x2", - partial( - _arch_nvvm_wrappers.calc_packed_f32x2_op, - calc_func=_fma_packed_bf16x2_nvvm, - ), -) +@dsl_user_op +def _prmt_trunc_bf16x2_neg_lo( + v0: cutlass.Float32, + v1: cutlass.Float32, + *, + loc=None, + ip=None, +): + """Like _prmt_trunc_bf16x2 but with the LOW bf16 half negated: + returns packed (-bf16(v0), bf16(v1)). + + Used for the rope rotation's mixed-sign multiply: (x1, x0) * + (-sin, +sin) == (-x1, x0) * (sin, sin), which replaces the f32 + sin negation (an FADD that LLVM hoists away from the FMA, + defeating ptxas's .NP-modifier folding) with one sign-bit LOP3 on + the packed operand and lets the b operand be a scalar broadcast.""" + packed_i32 = cute.arch.prmt( + v0.bitcast(cutlass.Int32, loc=loc, ip=ip), + v1.bitcast(cutlass.Int32, loc=loc, ip=ip), + cutlass.Int32(0x7632), + loc=loc, + ip=ip, + ) + negated = cute.arch.lop3( + cutlass.Int32(packed_i32), + cutlass.Int32(0x00008000), + cutlass.Int32(0), + 0x3C, # a XOR b + loc=loc, + ip=ip, + ) + return mlir_llvm.bitcast( + mlir_ir.VectorType.get([2], cutlass.BFloat16.mlir_type, loc=loc), + negated.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fma_packed_f32x2_bf16x2_f32x2_f32x2( + vec_a, + src_b, + src_c, + *, + loc=None, + ip=None, +): + """fma_packed_f32x2_bf16x2_f32x2_f32x2 with a pre-packed bf16x2 + src_a (vector<2xbf16> ir.Value), skipping the wrapper's per-element + f32->bf16 conversion (the F2FP).""" + vec_res_type = mlir_ir.VectorType.get([2], cutlass.Float32.mlir_type, loc=loc) + vec_b = mlir_vector.from_elements( + vec_res_type, + tuple(v.ir_value(loc=loc, ip=ip) for v in src_b), + loc=loc, + ip=ip, + ) + vec_c = mlir_vector.from_elements( + vec_res_type, + tuple(v.ir_value(loc=loc, ip=ip) for v in src_c), + loc=loc, + ip=ip, + ) + rn = next(mode for mode in FPRoundingMode if str(mode) == "rn") + vec_res = nvvm.fma_packed_f32x2_bf16x2_f32x2_f32x2( + vec_res_type, + vec_a, + vec_b, + vec_c, + rnd=rn, + loc=loc, + ip=ip, + ) + return tuple( + cutlass.Float32( + mlir_vector.extract( + vec_res, + dynamic_position=[], + static_position=[i], + loc=loc, + ip=ip, + ) + ) + for i in range(2) + ) class S2TCopyBundle(NamedTuple): @@ -175,6 +282,7 @@ def __init__( rms_norm_eps: float = DEFAULT_RMS_NORM_EPS, max_batch: int = 128, tma_prefetch_dist: int = 0, + enable_pdl: bool = True, ): if int(max_batch) < 1: raise ValueError("max_batch must be >= 1") @@ -182,6 +290,7 @@ def __init__( if int(tma_prefetch_dist) < 0: raise ValueError("tma_prefetch_dist must be >= 0") self.tma_prefetch_dist = int(tma_prefetch_dist) + self.enable_pdl = bool(enable_pdl) if tuple(mma_inst_tile) not in ((128, 256), (256, 256)): raise ValueError("mma_inst_tile must be (128, 256) or (256, 256)") mma_m = mma_inst_tile[0] @@ -269,11 +378,12 @@ def __init__( self.f32_bytes = cutlass.Float32.width // 8 self.i32_bytes = cutlass.Int32.width // 8 self.rope_cache_row_floats = 2 * self.qk_rope_head_dim - self.rope_smem_row_words = self.qk_rope_head_dim - self.rope_cp_bytes = 16 - self.rope_cp_chunks = (self.rope_smem_row_words * self.f32_bytes) // self.rope_cp_bytes - self.rope_rows_per_cp = 32 // self.rope_cp_chunks - self.rope_cp_iters = 32 // self.rope_rows_per_cp + self.rope_smem_row_words = self.qk_rope_head_dim // 2 + self.rope_row_quads = self.qk_rope_head_dim // 4 + self.rope_half_quads = self.rope_row_quads // 2 + self.rope_ldg_row_threads = self.rope_half_quads + self.rope_ldg_rows_per_iter = 32 // self.rope_ldg_row_threads + self.rope_ldg_iters = 32 // self.rope_ldg_rows_per_iter def _check_cluster(self, shape_mn, name: str) -> None: cm, cn = shape_mn @@ -904,8 +1014,6 @@ def kernel_body( smem = utils.SmemAllocator() storage = smem.allocate(self.shared_storage) - # CUTLASS DSL 4.5 expresses multicast signaling through the consumer - # thread count; this is API adaptation only. num_mcast_ctas_a = cute.size(cluster_layout_vmnk.shape[2]) num_mcast_ctas_b = cute.size(cluster_layout_vmnk.shape[1]) num_tma_producers = num_mcast_ctas_a + num_mcast_ctas_b - 1 @@ -1071,372 +1179,191 @@ def kernel_body( work_tile = tile_sched.initial_work_tile_info() iket.range_pop() - if warp_idx == self.tma_warp_id: - cute.arch.setmaxregister_decrease(self.mainloop_reg_count) - ab_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, self.num_ab_stage + s_cu_seqlens = storage.sCuSeqlens.get_tensor(cute.make_layout(self.max_batch + 1)) + s_kv_lengths = storage.sKvLengths.get_tensor(cute.make_layout(self.max_batch)) + s_rope_t = storage.sRope.get_tensor( + self.rope_smem_layout.outer, + swizzle=self.rope_smem_layout.inner, + ) + if cutlass.const_expr(self.use_tma_store): + sOut = storage.sOut.get_tensor( + self.out_smem_layout_staged.outer, + swizzle=self.out_smem_layout_staged.inner, + ) + + if warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_increase(self.epilog_reg_count) + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + for meta_i in cutlass.range(0, (batch_size + 128) // 128, 1, unroll=1): + meta_idx = tidx + meta_i * cutlass.Int32(128) + if meta_idx < batch_size + 1: + s_cu_seqlens[meta_idx] = cu_q_seqlens[meta_idx] + if meta_idx < batch_size: + s_kv_lengths[meta_idx] = kv_cache_lengths[meta_idx] + if cutlass.const_expr(quant_scale_qkv is None): + quant_scale_value = cutlass.Float32(1.0) + else: + quant_scale_value = quant_scale_qkv[0] + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tCtAcc_epi = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + tCtAcc_x = gemm_sm100.transform_partitioned_tensor_layout(tCtAcc_epi) + tCgC_x = gemm_sm100.transform_partitioned_tensor_layout(tCgC) + epi_tile = (self.cta_tile_shape_mnk[0], self.epi_chunk) + tAcc_epi = cute.flat_divide(tCtAcc_x, epi_tile) + gC_epi = cute.flat_divide(tCgC_x, epi_tile) + copy_atom_t2r = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.epi_chunk)), + self.acc_dtype, + ) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + tTR_gC_part = thr_copy_t2r.partition_D(gC_epi) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC_part[(None, None, None, 0, 0, 0, 0, 0)].shape, + self.acc_dtype, + ) + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tTR_rAcc_flat = cute.make_tensor(tTR_rAcc.iterator, cute.make_layout((self.epi_chunk,))) + tTR_rC_flat = cute.make_tensor(tTR_rC.iterator, cute.make_layout((self.epi_chunk,))) + simt_atom = cute.make_copy_atom( + cute.nvgpu.CopyR2GOp(), + self.c_dtype, + num_bits_per_copy=256, + l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE, + ) + tTR_rC_quads = cute.tiled_divide(tTR_rC_flat, (16,)) + sts_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.c_dtype, + num_bits_per_copy=128, + ) + pred_store = cute.make_rmem_tensor((1, *tTR_rC.shape[1:]), cutlass.Boolean) + if cutlass.const_expr(self.use_tma_store): + s_out_quads = cute.tiled_divide(sOut, (1, 16)) + gC_store = cute.local_tile(mC_tma, self.epi_store_tile, (None, None, None)) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sOut, 0, 2), + cute.group_modes(gC_store, 0, 2), + ) + + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + s_rope_quads = cute.tiled_divide(s_rope_t, (1, 4)) + cache_quads = cute.tiled_divide(cos_sin_cache, (1, 4)) + cs_vec_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float32, + num_bits_per_copy=128, + ) + cs_ldg_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.Float32, + num_bits_per_copy=128, + load_cache_mode=cute.nvgpu.common.LoadCacheMode.GLOBAL, ) + lane = tidx & cutlass.Int32(31) + row_in_cta = warp_idx * cutlass.Int32(32) + lane + ldg_row_sel = lane >> cutlass.Int32(3) + ldg_chunk = lane & cutlass.Int32(self.rope_ldg_row_threads - 1) + warp_row0 = warp_idx * cutlass.Int32(32) + cache_qcols = cache_quads[((0, None), 0, None)] + cs_cache_addr = cache_qcols.iterator.toint().to(cutlass.Int64) + cs_stage_flat = cute.make_rmem_tensor((self.qk_rope_head_dim,), cutlass.Float32) + cs_stage = cute.tiled_divide(cs_stage_flat, (4,)) + cs_row_flat = cute.make_rmem_tensor((self.qk_rope_head_dim,), cutlass.Float32) + cs_row = cute.tiled_divide(cs_row_flat, (4,)) + while work_tile.is_valid_tile: - iket.range_push("tma_tile") - cur_tile_coord = work_tile.tile_idx - mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma_akeep.thr_id.shape) - head_id = cur_tile_coord[1] + iket.range_push("epi_tile") - tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] - tBgB_slice = tBgB[(None, head_id, None, 0)] - tAgSFA_slice = tAgSFA[(None, mma_tile_coord_m, None, 0)] - tBgSFB_slice = tBgSFB[(None, head_id, None, 0)] + def cs_transpose_half(h): + for ldg_i in cutlass.range_constexpr(self.rope_ldg_iters): + cute.copy( + cs_vec_atom, + cs_stage[(None, h * self.rope_ldg_iters + ldg_i)], + s_rope_quads[ + ( + (0, None), + warp_row0 + + cutlass.Int32(self.rope_ldg_rows_per_iter * ldg_i) + + ldg_row_sel, + ldg_chunk, + ) + ], + ) + cute.arch.sync_warp() + for qc in cutlass.range_constexpr(self.rope_half_quads): + cute.copy( + cs_vec_atom, + s_rope_quads[((0, None), row_in_cta, qc)], + cs_row[(None, h * self.rope_half_quads + qc)], + ) + cute.arch.sync_warp() - if cutlass.const_expr(self.tma_prefetch_dist > 0): - for pf_k_tile in cutlass.range( - 0, - cutlass.min(self.tma_prefetch_dist, k_tile_cnt), - 1, - unroll=1, - ): - cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, pf_k_tile)]) + cur_tile_coord = work_tile.tile_idx + head_id = cur_tile_coord[1] + coord_m_cta = cur_tile_coord[0] * cutlass.Int32(self.cta_tile_shape_mnk[0]) + row = coord_m_cta + row_in_cta + row_in_bounds = row < m + if cutlass.const_expr(self.use_tma_store): + bSG_gC_tile = bSG_gC[(None, cur_tile_coord[0], None, 0)] + else: + pred_store[(0, 0, 0)] = row_in_bounds + mma_coord_m = cur_tile_coord[0] // cute.size(tiled_mma_akeep.thr_id.shape) + tTR_gC_tile = tTR_gC_part[(None, None, None, None, None, mma_coord_m, head_id, 0)] - ab_producer_state.reset_count() - peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): - iket.range_push("tma_wait", k_tile) - ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) - iket.range_pop() - iket.range_push("tma_issue") - tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) - cute.copy( - tma_atom_a, - tAgA_slice[(None, ab_producer_state.count)], - tAsA[(None, ab_producer_state.index)], - tma_bar_ptr=tma_bar, - mcast_mask=a_full_mcast_mask, - ) - cute.copy( - tma_atom_sfa, - tAgSFA_slice[(None, ab_producer_state.count)], - tAsSFA[(None, ab_producer_state.index)], - tma_bar_ptr=tma_bar, - mcast_mask=sfa_full_mcast_mask, + iket.range_push("epi_position") + position = cutlass.Int32(-1) + if row_in_bounds: + candidate = self._position_for_row( + s_cu_seqlens, + s_kv_lengths, + helix_position_offsets, + row, + batch_size, ) + if candidate >= 0 and candidate < rope_positions: + position = candidate + iket.range_pop() - cute.copy( - tma_atom_b, - tBgB_slice[(None, ab_producer_state.count)], - tBsB[(None, ab_producer_state.index)], - tma_bar_ptr=tma_bar, - mcast_mask=b_full_mcast_mask, - ) - cute.copy( - tma_atom_sfb, - tBgSFB_slice[(None, ab_producer_state.count)], - tBsSFB[(None, ab_producer_state.index)], - tma_bar_ptr=tma_bar, - mcast_mask=sfb_full_mcast_mask, + iket.range_push("epi_cs_stage") + pos_qcol_base = position * cutlass.Int32(self.rope_cache_row_floats // 4) + src_qcols, src_addrs = [], [] + for ldg_i in cutlass.range_constexpr(self.rope_ldg_iters): + src_lane = cutlass.Int32(self.rope_ldg_rows_per_iter * ldg_i) + ldg_row_sel + src_qcol = cute.arch.shuffle_sync(pos_qcol_base, src_lane) + src_qcols.append(src_qcol) + src_addrs.append( + cs_cache_addr + (src_qcol + ldg_chunk).to(cutlass.Int64) * cutlass.Int64(16) ) - if cutlass.const_expr(self.tma_prefetch_dist > 0): - pf_k_tile = k_tile + self.tma_prefetch_dist - if pf_k_tile < k_tile_cnt: - cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)]) - cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, pf_k_tile)]) - iket.range_pop() - ab_producer_state.advance() - peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: - peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) - iket.range_push("tma_wait_clc") - clc_pipeline.consumer_wait(clc_consumer_state) - iket.range_pop() - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() + for ldg_i in cutlass.range_constexpr(self.rope_ldg_iters): + if src_qcols[ldg_i] >= 0: + for h in cutlass.range_constexpr(2): + src_ptr = cute.make_ptr( + cutlass.Float32, + src_addrs[ldg_i] + cutlass.Int64(h * self.rope_half_quads * 16), + cute.AddressSpace.gmem, + assumed_align=16, + ) + cute.copy( + cs_ldg_atom, + cute.make_tensor(src_ptr, cute.make_layout((4,))), + cs_stage[(None, h * self.rope_ldg_iters + ldg_i)], + ) iket.range_pop() - ab_pipeline.producer_tail(ab_producer_state) - if warp_idx == self.sched_warp_id: - cute.arch.setmaxregister_decrease(self.mainloop_reg_count) - if is_first_cta_in_cluster: - clc_producer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.ProducerConsumer, - self.num_clc_stage, - ) - while work_tile.is_valid_tile: - iket.range_push("sched_tile") - iket.range_push("sched_wait_empty") - clc_pipeline.producer_acquire(clc_producer_state) - iket.range_pop() - iket.range_push("sched_query") - mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state) - tile_sched.advance_to_next_work(mbarrier_addr) - clc_producer_state.advance() + cs_transpose_half(1) - clc_pipeline.consumer_wait(clc_consumer_state) - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() - iket.range_pop() - iket.range_pop() - clc_pipeline.producer_tail(clc_producer_state) - - if warp_idx == self.mma_warp_id: - cute.arch.setmaxregister_decrease(self.mainloop_reg_count) - tmem.wait_for_alloc() - acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - sfa_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + self.num_accumulator_tmem_cols, - dtype=self.sf_dtype, - ) - tCtSFA = cute.make_tensor(sfa_tmem_ptr, self.tCtSFA_layout) - sfb_tmem_ptr = cute.recast_ptr( - acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, - ) - tCtSFB = cute.make_tensor(sfb_tmem_ptr, self.tCtSFB_layout) - - sfa_s2t = self._mainloop_s2t_copy_and_partition(sSFA, tCtSFA) - sfb_s2t = self._mainloop_s2t_copy_and_partition(sSFB, tCtSFB) - - ab_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_ab_stage - ) - acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) - while work_tile.is_valid_tile: - iket.range_push("mma_tile") - ab_consumer_state.reset_count() - peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt and is_leader_cta: - peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) - iket.range_push("mma_wait_acc") - if is_leader_cta: - acc_pipeline.producer_acquire(acc_producer_state) - iket.range_pop() - - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): - if is_leader_cta: - iket.range_push("mma_wait", k_tile) - ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) - iket.range_pop() - iket.range_push("mma_issue") - self._mainloop_s2t_copies(ab_consumer_state.index, sfa_s2t, sfb_s2t) - num_kblocks = cute.size(tCrA, mode=[2]) - for k_block in cutlass.range(num_kblocks, unroll_full=True): - a_kblk_crd = ( - None, - 0, - k_block, - ab_consumer_state.index, - ) - sfa_kblk_crd = (None, 0, k_block) - for n_set in cutlass.range_constexpr(2): - b_kblk_crd = ( - None, - n_set, - k_block, - ab_consumer_state.index, - ) - sfb_kblk_crd = (None, n_set, k_block) - tCtAcc_set = tCtAcc_base[(None, 0, n_set)] - if cutlass.const_expr(n_set == 0): - tiled_mma_akeep.set( - tcgen05.Field.ACCUMULATE, - k_tile != 0 or k_block != 0, - ) - cute.gemm( - tiled_mma_akeep, - tCtAcc_set, - [ - tCrA[a_kblk_crd], - tCtSFA[sfa_kblk_crd], - ], - [ - tCrB[b_kblk_crd], - tCtSFB[sfb_kblk_crd], - ], - tCtAcc_set, - ) - else: - tiled_mma_areuse.set( - tcgen05.Field.ACCUMULATE, - k_tile != 0 or k_block != 0, - ) - cute.gemm( - tiled_mma_areuse, - tCtAcc_set, - [ - tCrA[a_kblk_crd], - tCtSFA[sfa_kblk_crd], - ], - [ - tCrB[b_kblk_crd], - tCtSFB[sfb_kblk_crd], - ], - tCtAcc_set, - ) - ab_pipeline.consumer_release(ab_consumer_state) - iket.range_pop() - ab_consumer_state.advance() - peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt: - if is_leader_cta: - peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) - if is_leader_cta: - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - iket.range_push("mma_wait_clc") - clc_pipeline.consumer_wait(clc_consumer_state) - iket.range_pop() - work_tile = tile_sched.get_current_work() - clc_pipeline.consumer_release(clc_consumer_state) - clc_consumer_state.advance() - iket.range_pop() - acc_pipeline.producer_tail(acc_producer_state) - - # CUTLASS DSL 4.5 cannot carry a cute.struct through a dynamic - # warp-role branch. Materialize SMEM tensors before that branch. - s_cu_seqlens = storage.sCuSeqlens.get_tensor(cute.make_layout(self.max_batch + 1)) - s_kv_lengths = storage.sKvLengths.get_tensor(cute.make_layout(self.max_batch)) - sOut = None - if cutlass.const_expr(self.use_tma_store): - sOut = storage.sOut.get_tensor( - self.out_smem_layout_staged.outer, - swizzle=self.out_smem_layout_staged.inner, - ) - s_rope_t = storage.sRope.get_tensor( - self.rope_smem_layout.outer, - swizzle=self.rope_smem_layout.inner, - ) - - if warp_idx == self.pad_warp_id: - cute.arch.setmaxregister_decrease(self.mainloop_reg_count) - - if warp_idx < self.mma_warp_id: - cute.arch.setmaxregister_increase(self.epilog_reg_count) - for meta_i in cutlass.range(0, (batch_size + 128) // 128, 1, unroll=1): - meta_idx = tidx + meta_i * cutlass.Int32(128) - if meta_idx < batch_size + 1: - s_cu_seqlens[meta_idx] = cu_q_seqlens[meta_idx] - if meta_idx < batch_size: - s_kv_lengths[meta_idx] = kv_cache_lengths[meta_idx] - - tmem.allocate(self.num_tmem_alloc_cols) - tmem.wait_for_alloc() - acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - - if cutlass.const_expr(quant_scale_qkv is None): - quant_scale_value = cutlass.Float32(1.0) - else: - quant_scale_value = quant_scale_qkv[0] - - tCtAcc_epi = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) - tCtAcc_x = gemm_sm100.transform_partitioned_tensor_layout(tCtAcc_epi) - tCgC_x = gemm_sm100.transform_partitioned_tensor_layout(tCgC) - epi_tile = (self.cta_tile_shape_mnk[0], self.epi_chunk) - tAcc_epi = cute.flat_divide(tCtAcc_x, epi_tile) - gC_epi = cute.flat_divide(tCgC_x, epi_tile) - copy_atom_t2r = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.epi_chunk)), - self.acc_dtype, - ) - tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) - thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) - tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) - tTR_gC_part = thr_copy_t2r.partition_D(gC_epi) - tTR_rAcc = cute.make_rmem_tensor( - tTR_gC_part[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.acc_dtype, - ) - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) - tTR_rAcc_flat = cute.make_tensor(tTR_rAcc.iterator, cute.make_layout((self.epi_chunk,))) - tTR_rC_flat = cute.make_tensor(tTR_rC.iterator, cute.make_layout((self.epi_chunk,))) - simt_atom = cute.make_copy_atom( - cute.nvgpu.CopyR2GOp(), - self.c_dtype, - num_bits_per_copy=256, - l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE, - ) - tTR_rC_quads = cute.tiled_divide(tTR_rC_flat, (16,)) - sts_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - self.c_dtype, - num_bits_per_copy=128, - ) - pred_store = cute.make_rmem_tensor((1, *tTR_rC.shape[1:]), cutlass.Boolean) - if cutlass.const_expr(self.use_tma_store): - s_out_quads = cute.tiled_divide(sOut, (1, 16)) - gC_store = cute.local_tile(mC_tma, self.epi_store_tile, (None, None, None)) - bSG_sC, bSG_gC = cpasync.tma_partition( - tma_atom_c, - 0, - cute.make_layout(1), - cute.group_modes(sOut, 0, 2), - cute.group_modes(gC_store, 0, 2), - ) - - acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) - s_rope_quads = cute.tiled_divide(s_rope_t, (1, 4)) - cache_quads = cute.tiled_divide(cos_sin_cache, (1, 4)) - cs_copy_atom = cute.make_copy_atom( - cpasync.CopyG2SOp(), cutlass.Float32, num_bits_per_copy=128 - ) - lane = tidx & cutlass.Int32(31) - row_in_cta = warp_idx * cutlass.Int32(32) + lane - cp_chunk = lane & cutlass.Int32(self.rope_cp_chunks - 1) - cp_row_sel = lane >> cutlass.Int32(4) - warp_row0 = warp_idx * cutlass.Int32(32) - cache_qcols = cache_quads[((0, None), 0, None)] - - while work_tile.is_valid_tile: - iket.range_push("epi_tile") - cur_tile_coord = work_tile.tile_idx - head_id = cur_tile_coord[1] - coord_m_cta = cur_tile_coord[0] * cutlass.Int32(self.cta_tile_shape_mnk[0]) - row = coord_m_cta + row_in_cta - row_in_bounds = row < m - if cutlass.const_expr(self.use_tma_store): - bSG_gC_tile = bSG_gC[(None, cur_tile_coord[0], None, 0)] - else: - pred_store[(0, 0, 0)] = row_in_bounds - mma_coord_m = cur_tile_coord[0] // cute.size(tiled_mma_akeep.thr_id.shape) - tTR_gC_tile = tTR_gC_part[(None, None, None, None, None, mma_coord_m, head_id, 0)] - - iket.range_push("epi_position") - position = cutlass.Int32(-1) - if row_in_bounds: - candidate = self._position_for_row( - s_cu_seqlens, - s_kv_lengths, - helix_position_offsets, - row, - batch_size, - ) - if candidate >= 0 and candidate < rope_positions: - position = candidate - iket.range_pop() - - iket.range_push("epi_cs_stage") - pos_qcol_base = position * cutlass.Int32(self.rope_cache_row_floats // 4) - for cp_i in cutlass.range_constexpr(self.rope_cp_iters): - src_lane = cutlass.Int32(2 * cp_i) + cp_row_sel - src_qcol = cute.arch.shuffle_sync(pos_qcol_base, src_lane) - if src_qcol >= 0: - cute.copy( - cs_copy_atom, - cache_qcols[(None, src_qcol + cp_chunk)], - s_rope_quads[((0, None), warp_row0 + src_lane, cp_chunk)], - ) - cute.arch.cp_async_commit_group() - iket.range_pop() - - iket.range_push("epi_wait_acc") - acc_pipeline.consumer_wait(acc_consumer_state) - iket.range_pop() + iket.range_push("epi_wait_acc") + acc_pipeline.consumer_wait(acc_consumer_state) + iket.range_pop() iket.range_push("epi_rmsnorm_reduce") f32_zero = self.acc_dtype(0.0) @@ -1463,60 +1390,53 @@ def kernel_body( norm_quant_scale = inv_rms * quant_scale_value iket.range_pop() - iket.range_push("epi_cs_drain") - cute.arch.cp_async_wait_group(0) if cutlass.const_expr(self.use_tma_store): if warp_idx == 0: cute.arch.cp_async_bulk_wait_group(0, read=True) - self.epilog_sync_barrier.arrive_and_wait() - iket.range_pop() + self.epilog_sync_barrier.arrive_and_wait() - for rope_rev in cutlass.range_constexpr(self.rope_chunks): - chunk = self.epi_chunks - 1 - rope_rev - if cutlass.const_expr(rope_rev != 0): - cute.copy( - tiled_copy_t2r, - tTR_tAcc[(None, None, None, 0, chunk)], - tTR_rAcc, - ) - iket.range_push("epi_rope") + def rope_chunk_body(chunk): col_base = chunk * self.epi_chunk for quad_i in cutlass.range_constexpr(self.epi_chunk // 4): elem0 = quad_i * 4 quad_idx = (col_base - self.qk_nope_head_dim) // 4 + quad_i - cs_quad = s_rope_quads[((0, None), row_in_cta, quad_idx)].load() + cs_quad = cs_row[(None, quad_idx)].load() for sub in cutlass.range_constexpr(2): e0 = elem0 + sub * 2 cos_value = cs_quad[sub * 2] sin_value = cs_quad[sub * 2 + 1] + z_pk = _prmt_trunc_bf16x2( + tTR_rAcc_flat[e0], + tTR_rAcc_flat[e0 + 1], + ) x0, x1 = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - ( - tTR_rAcc_flat[e0], - tTR_rAcc_flat[e0 + 1], - ), + z_pk, (inv_rms, inv_rms), (f32_zero, f32_zero), ) + x_pk = _prmt_trunc_bf16x2(x0, x1) + x_pk_rev = _prmt_trunc_bf16x2_neg_lo(x1, x0) t_pair = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - (x0, x1), + x_pk, (cos_value, cos_value), (f32_zero, f32_zero), ) y0, y1 = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - (x1, x0), - (-sin_value, sin_value), + x_pk_rev, + (sin_value, sin_value), t_pair, ) + y_pk = _prmt_trunc_bf16x2(y0, y1) o0, o1 = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - (y0, y1), + y_pk, (quant_scale_value, quant_scale_value), (f32_zero, f32_zero), ) tTR_rAcc_flat[e0] = o0 tTR_rAcc_flat[e0 + 1] = o1 - out_ssa = tTR_rAcc_flat.load() - iket.range_pop() - iket.range_push("epi_quant+stg") + return tTR_rAcc_flat.load() + + def store_chunk(chunk, out_ssa): tTR_rC_flat.store(out_ssa.to(self.c_dtype)) if cutlass.const_expr(self.use_tma_store): q0 = (chunk & 3) * 2 @@ -1534,7 +1454,12 @@ def kernel_body( tTR_gC_tile[(None, None, None, 0, chunk)], pred=pred_store, ) - iket.range_pop() + + rope_hi_chunk = self.epi_chunks - 1 + out_hi = rope_chunk_body(rope_hi_chunk) + store_chunk(rope_hi_chunk, out_hi) + + cs_transpose_half(0) for chunk_i in cutlass.range_constexpr(self.nope_chunks): chunk = self.nope_chunks - 1 - chunk_i cute.copy( @@ -1546,40 +1471,33 @@ def kernel_body( cute.arch.fence_view_async_tmem_load() for pair in cutlass.range_constexpr(self.epi_chunk // 2): e0 = pair * 2 + z_pk = _prmt_trunc_bf16x2( + tTR_rAcc_flat[e0], + tTR_rAcc_flat[e0 + 1], + ) o0, o1 = _fma_packed_f32x2_bf16x2_f32x2_f32x2( - ( - tTR_rAcc_flat[e0], - tTR_rAcc_flat[e0 + 1], - ), + z_pk, (norm_quant_scale, norm_quant_scale), (f32_zero, f32_zero), ) tTR_rAcc_flat[e0] = o0 tTR_rAcc_flat[e0 + 1] = o1 out_ssa = tTR_rAcc_flat.load() - if cutlass.const_expr(chunk_i == self.nope_chunks - 1): - nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC) - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - iket.range_push("epi_quant+stg") - tTR_rC_flat.store(out_ssa.to(self.c_dtype)) - if cutlass.const_expr(self.use_tma_store): - q0 = (chunk & 3) * 2 - slot = (chunk >> 2) & (self.out_ring_slots - 1) - for h in cutlass.range_constexpr(2): - cute.copy( - sts_atom, - tTR_rC_quads[(None, h)], - s_out_quads[((0, None), row_in_cta, q0 + h, slot)], - ) - else: - cute.copy( - simt_atom, - tTR_rC, - tTR_gC_tile[(None, None, None, 0, chunk)], - pred=pred_store, - ) - iket.range_pop() + store_chunk(chunk, out_ssa) + + rope_lo_chunk = self.nope_chunks + cute.copy( + tiled_copy_t2r, + tTR_tAcc[(None, None, None, 0, rope_lo_chunk)], + tTR_rAcc, + ) + + cute.arch.fence_view_async_tmem_load() + nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC) + acc_pipeline.consumer_release(acc_consumer_state) + out_lo = rope_chunk_body(rope_lo_chunk) + acc_consumer_state.advance() + store_chunk(rope_lo_chunk, out_lo) if cutlass.const_expr(self.use_tma_store): iket.range_push("epi_fence+tmastg") @@ -1616,6 +1534,244 @@ def kernel_body( tmem.free(acc_tmem_ptr) iket.range_pop() + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.mainloop_reg_count) + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + a_cache_policy = cute.CacheEvictionPriority.EVICT_LAST + a_cache_policy = cutlass.Int64(0x14F0000000000000) # L2 evict_last + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + while work_tile.is_valid_tile: + iket.range_push("tma_tile") + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma_akeep.thr_id.shape) + head_id = cur_tile_coord[1] + + tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] + tBgB_slice = tBgB[(None, head_id, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_coord_m, None, 0)] + tBgSFB_slice = tBgSFB[(None, head_id, None, 0)] + + if cutlass.const_expr(self.tma_prefetch_dist > 0): + for pf_k_tile in cutlass.range( + 0, + cutlass.min(self.tma_prefetch_dist, k_tile_cnt), + 1, + unroll=1, + ): + cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, pf_k_tile)]) + + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + iket.range_push("tma_wait", k_tile) + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + iket.range_pop() + iket.range_push("tma_issue") + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=tma_bar, + mcast_mask=a_full_mcast_mask, + cache_policy=a_cache_policy, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=tma_bar, + mcast_mask=sfa_full_mcast_mask, + cache_policy=a_cache_policy, + ) + + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=tma_bar, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=tma_bar, + mcast_mask=sfb_full_mcast_mask, + ) + if cutlass.const_expr(self.tma_prefetch_dist > 0): + pf_k_tile = k_tile + self.tma_prefetch_dist + if pf_k_tile < k_tile_cnt: + cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)]) + cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, pf_k_tile)]) + iket.range_pop() + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + iket.range_push("tma_wait_clc") + clc_pipeline.consumer_wait(clc_consumer_state) + iket.range_pop() + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + iket.range_pop() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + ab_pipeline.producer_tail(ab_producer_state) + + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.mainloop_reg_count) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, self.tCtSFA_layout) + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, self.tCtSFB_layout) + + sfa_s2t = self._mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + sfb_s2t = self._mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + while work_tile.is_valid_tile: + iket.range_push("mma_tile") + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + iket.range_push("mma_wait_acc") + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + iket.range_pop() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + if is_leader_cta: + iket.range_push("mma_wait", k_tile) + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + iket.range_pop() + iket.range_push("mma_issue") + self._mainloop_s2t_copies(ab_consumer_state.index, sfa_s2t, sfb_s2t) + num_kblocks = cute.size(tCrA, mode=[2]) + for k_block in cutlass.range(num_kblocks, unroll_full=True): + a_kblk_crd = ( + None, + 0, + k_block, + ab_consumer_state.index, + ) + sfa_kblk_crd = (None, 0, k_block) + for n_set in cutlass.range_constexpr(2): + b_kblk_crd = ( + None, + n_set, + k_block, + ab_consumer_state.index, + ) + sfb_kblk_crd = (None, n_set, k_block) + tCtAcc_set = tCtAcc_base[(None, 0, n_set)] + if cutlass.const_expr(n_set == 0): + tiled_mma_akeep.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or k_block != 0, + ) + cute.gemm( + tiled_mma_akeep, + tCtAcc_set, + [ + tCrA[a_kblk_crd], + tCtSFA[sfa_kblk_crd], + ], + [ + tCrB[b_kblk_crd], + tCtSFB[sfb_kblk_crd], + ], + tCtAcc_set, + ) + else: + tiled_mma_areuse.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or k_block != 0, + ) + cute.gemm( + tiled_mma_areuse, + tCtAcc_set, + [ + tCrA[a_kblk_crd], + tCtSFA[sfa_kblk_crd], + ], + [ + tCrB[b_kblk_crd], + tCtSFB[sfb_kblk_crd], + ], + tCtAcc_set, + ) + ab_pipeline.consumer_release(ab_consumer_state) + iket.range_pop() + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + iket.range_push("mma_wait_clc") + clc_pipeline.consumer_wait(clc_consumer_state) + iket.range_pop() + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + iket.range_pop() + acc_pipeline.producer_tail(acc_producer_state) + + if warp_idx == self.sched_warp_id: + cute.arch.setmaxregister_decrease(self.mainloop_reg_count) + if is_first_cta_in_cluster: + clc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.ProducerConsumer, + self.num_clc_stage, + ) + while work_tile.is_valid_tile: + iket.range_push("sched_tile") + iket.range_push("sched_wait_empty") + clc_pipeline.producer_acquire(clc_producer_state) + iket.range_pop() + iket.range_push("sched_query") + mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state) + tile_sched.advance_to_next_work(mbarrier_addr) + clc_producer_state.advance() + + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + iket.range_pop() + iket.range_pop() + clc_pipeline.producer_tail(clc_producer_state) + + if warp_idx == self.pad_warp_id: + cute.arch.setmaxregister_decrease(self.mainloop_reg_count) + @cute.jit def _mainloop_s2t_copy_and_partition(self, sSF: cute.Tensor, tSF: cute.Tensor) -> S2TCopyBundle: """Tiled S2T (UTCCP) copy for one SF tensor + its partitions.""" diff --git a/tensorrt_llm/_torch/cute_dsl_utils.py b/tensorrt_llm/_torch/cute_dsl_utils.py index cb571f85b597..352e1de26e7a 100644 --- a/tensorrt_llm/_torch/cute_dsl_utils.py +++ b/tensorrt_llm/_torch/cute_dsl_utils.py @@ -1,7 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import platform +import sys +import types from ..logger import logger + +def _skip_legacy_cutlass_mlir_helpers() -> None: + """Keep Cutlass version discovery from importing its legacy helper tree.""" + legacy_name = "cutlass.base_dsl._mlir_helpers" + if legacy_name in sys.modules: + return + + # Cutlass uses pkgutil.walk_packages to hash its sources. The internal + # package also ships this unused legacy helper tree alongside the canonical + # cutlass._mlir_helpers package. Prevent pkgutil from descending into the + # legacy tree, which would register the same MLIR value casters twice. + legacy_module = types.ModuleType(legacy_name) + legacy_module.__path__ = [] + sys.modules[legacy_name] = legacy_module + + IS_CUTLASS_DSL_AVAILABLE = False # Whether the public CuTe DSL package provides the SM107/Rubin helper module. @@ -12,8 +45,8 @@ if platform.system() != "Windows": try: - import cutlass # noqa - import cutlass.cute as cute # noqa + from cutlass import cute # noqa + _skip_legacy_cutlass_mlir_helpers() logger.info(f"cutlass dsl is available") IS_CUTLASS_DSL_AVAILABLE = True 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_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index ba1a38fa6956..3a2e180664f1 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -713,6 +713,7 @@ def __init__( reduce_output: bool = True, # ROW parallel only skip_create_weights_in_init: bool = False, use_custom_cublas_mm: bool = False, + use_cute_dsl_bf16_gemm: bool = False, use_cute_dsl_blockscaling_mm: bool = False, lora: Optional[LoraLayer] = None, ): @@ -720,16 +721,17 @@ def __init__( in_features, out_features, bias, - dtype, - mapping, - tensor_parallel_mode, - gather_output, - quant_config, - weights_loading_config, - reduce_output, - skip_create_weights_in_init, - use_custom_cublas_mm, - lora, + dtype=dtype, + mapping=mapping, + tensor_parallel_mode=tensor_parallel_mode, + gather_output=gather_output, + quant_config=quant_config, + weights_loading_config=weights_loading_config, + reduce_output=reduce_output, + skip_create_weights_in_init=skip_create_weights_in_init, + use_custom_cublas_mm=use_custom_cublas_mm, + use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + lora=lora, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, ) @@ -739,7 +741,12 @@ def apply_linear(self, lora_params: Optional[dict] | None = None, layer_idx: Optional[int] | None = None): num_tokens = input.shape[0] - if (not self.has_any_quant and 1 <= num_tokens <= 16 + has_any_quant = self.has_any_quant + use_cute_dsl_bf16_gemm = (self.use_cute_dsl_bf16_gemm + and not has_any_quant and is_sm_100f() + and self.weight.dtype == torch.bfloat16) + if (not use_cute_dsl_bf16_gemm and not has_any_quant + and 1 <= num_tokens <= 16 and get_sm_version() not in [120, 121]): output = torch.ops.trtllm.dsv3_fused_a_gemm_op( input, self.weight.t(), bias, None) @@ -793,6 +800,7 @@ def __init__( skip_create_weights_in_init=model_config. skip_create_weights_in_init, use_custom_cublas_mm=True, + use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm, use_cute_dsl_blockscaling_mm=model_config. use_cute_dsl_blockscaling_mm, ) @@ -843,7 +851,8 @@ def __init__( quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config. skip_create_weights_in_init, - use_custom_cublas_mm=True) + use_custom_cublas_mm=True, + use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm) class DeepseekV3Gate(nn.Module): diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index e854fc6f6769..0e1a7638a702 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -781,8 +781,8 @@ def load_o_a_proj(module_name: str, module) -> None: if o_a_proj_scale is not None: o_a_proj_scale = split_matrix_tp(o_a_proj_scale, tp_size, tp_rank, 0) - # Skip the BF16 dequant when the destination is FP8 (the cute_dsl - # FP8 BMM path on SM100 consumes the native FP8 weight directly). + # Skip BF16 dequant when the architecture-specific CuTe DSL BMM + # consumes the native FP8 weight directly. if o_a_proj_scale is not None and module.o_a_proj.dtype != torch.float8_e4m3fn: o_a_proj = weight_dequant( o_a_proj.reshape(-1, o_a_proj.shape[-1]).contiguous().cuda(), @@ -1598,6 +1598,7 @@ def __init__( config=model_config, overridden_tp_size=shared_tp_size, reduce_output=False, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, swiglu_limit=swiglu_limit, ) @@ -2265,6 +2266,7 @@ def __init__( dtype=config.torch_dtype, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) self.h_proj = Linear( config.hidden_size, @@ -2273,6 +2275,7 @@ def __init__( dtype=config.torch_dtype, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) else: self.e_proj = Linear( @@ -2285,6 +2288,7 @@ def __init__( reduce_output=True, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) self.h_proj = Linear( config.hidden_size, @@ -2296,6 +2300,7 @@ def __init__( reduce_output=True, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) self.shared_head = DeepseekV4MTPHead(model_config) diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index 24bf6d21b9a6..464e6e59ec80 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -5,9 +5,11 @@ import torch.nn.functional as F from torch import nn +from tensorrt_llm._utils import get_sm_version, is_sm_100f from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from ..cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE from ..distributed import AllReduceParams from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType @@ -119,6 +121,9 @@ def __init__( allreduce_strategy=config.allreduce_strategy, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, + use_cute_dsl_nvfp4_swiglu_blackwell=(use_cute_dsl_blockscaling_mm + and activation == F.silu + and not bias), use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, disable_deep_gemm=disable_deep_gemm, fused_weight_shard_indices_mapping=gateup_shard_indices_mapping, @@ -214,16 +219,11 @@ def _is_plain_swiglu(self): def _can_fuse_gate_up_swiglu(self): """Check if fused GEMM + SwiGLU path is available. - Returns True when all conditions are met: - - CuteDSL blockscaling mode is enabled (implies Blackwell + CuteDSL) - - Activation is plain SwiGLU (F.silu), see _is_plain_swiglu - - gate_up_proj uses NVFP4 quantization - - gate_up_proj has no bias (bias not supported in fused kernel) + The projection owns the capability predicate because weight loading + must make exactly the same decision as forward dispatch. """ - return (self.use_cute_dsl_blockscaling_mm and self.activation == F.silu - and self._is_plain_swiglu() - and self.gate_up_proj.has_nvfp4_activation_quantization - and not self.gate_up_proj.has_bias) + return (self.activation == F.silu and self._is_plain_swiglu() + and self.gate_up_proj.can_use_cute_dsl_nvfp4_swiglu_blackwell()) def _can_fuse_gate_up_swiglu_fp4out(self): """Check if fused GEMM + SwiGLU with FP4 output path is available. @@ -237,6 +237,20 @@ def _can_fuse_gate_up_swiglu_fp4out(self): return False return is_static_nvfp4_input_eligible(self.down_proj) + def _can_fuse_swiglu_fp8_quant(self) -> bool: + """Check whether down projection can consume fused SwiGLU FP8 output.""" + # Parameterized variants such as MiniMax M3 SwiGLU-OAI install a + # custom activation callable and must keep the unfused implementation. + if not (self.activation == F.silu + and self.down_proj.has_fp8_block_scales): + return False + if get_sm_version() == 107: + return (IS_CUTLASS_DSL_RUBIN_AVAILABLE + and (self.down_proj.use_cute_dsl_blockscaling_mm + or self.down_proj.disable_deep_gemm)) + return (is_sm_100f() and not self.down_proj.use_cute_dsl_blockscaling_mm + and not self.down_proj.disable_deep_gemm) + def _fused_gate_up_swiglu(self, x, fp4_out=False): """Fused FC1 GEMM + SwiGLU using CuteDSL dense kernel. @@ -315,6 +329,7 @@ def forward( return self.forward_lora(x, all_rank_num_tokens, final_all_reduce_params, lora_params) + fused_output_shape = None if self._can_fuse_gate_up_swiglu_fp4out(): # During torch.compile the token dim is a SymInt, so `m >= MIN_M` # would create a SymBool guard that breaks piecewise CUDA graph @@ -333,11 +348,21 @@ def forward( h2 = self._fused_gate_up_swiglu(x) else: h1 = self.gate_up_proj(x) - h2 = self._apply_activation(h1) + if self._can_fuse_swiglu_fp8_quant(): + if h1.dim() > 2: + fused_output_shape = h1.shape[:-1] + h1 = h1.reshape(-1, h1.shape[-1]) + use_r128c4_layout = get_sm_version() == 107 + h2 = torch.ops.trtllm.silu_and_mul_fp8_quantize_1x128_packed_ue8m0( + h1, self.swiglu_limit, use_r128c4_layout) + else: + h2 = self._apply_activation(h1) output = self.down_proj(h2, all_reduce_params=final_all_reduce_params, layer_idx=self.layer_idx) + if fused_output_shape is not None: + output = output.reshape(*fused_output_shape, output.shape[-1]) return output def forward_lora( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index b3e29a427482..cdacad5064b6 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -35,7 +35,8 @@ from ..._utils import get_sm_version, is_sm_100f from ...models.modeling_utils import QuantConfig -from ..cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from ..cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE) from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, replace_parameter_and_save_metadata, unswizzle_sf) from .low_m_gemm import _should_apply_low_m_gemm, apply_low_m_gemm @@ -1235,6 +1236,38 @@ def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): # fp8_block_scaling_gemm does not support writing into an NCCL window # buffer; supports_nccl_symmetric_memory_window_output is False so the window path is bypassed. + if isinstance(input, tuple): + if len(input) != 2: + raise ValueError( + "Pre-quantized FP8 input must contain activation and scale") + activation, activation_scale = input + sm_version = get_sm_version() + uses_cute_dsl_rubin = (activation_scale.dtype == torch.uint8 + and sm_version == 107 + and IS_CUTLASS_DSL_RUBIN_AVAILABLE + and (module.use_cute_dsl_blockscaling_mm + or module.disable_deep_gemm)) + if uses_cute_dsl_rubin: + output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + activation, module.weight, activation_scale, + module.weight_scale) + elif (activation_scale.dtype == torch.int32 and is_sm_100f() + and not module.disable_deep_gemm): + output = torch.ops.trtllm.fp8_prequantized_swap_ab_gemm( + activation, + activation_scale, + module.weight, + module.weight_scale, + disable_ue8m0_cast=True, + ) + else: + raise RuntimeError( + "Pre-quantized FP8 scale layout is incompatible with the " + "selected block-scale GEMM backend") + if bias is not None: + output = output + bias + return output + # Handle multi-dimensional inputs (e.g., 3D: batch, seq, hidden) # GEMM ops require 2D matrices original_shape = input.shape @@ -2024,7 +2057,7 @@ def process_weights_after_loading_fused_gate_up_linear( # interleaves in 64-row groups to match the kernel layout. # # Weight scales are similarly unswizzled, interleaved, and re-swizzled. - if not module.use_cute_dsl_blockscaling_mm: + if not module.can_use_cute_dsl_nvfp4_swiglu_blackwell(): return group_size = 64 @@ -3625,6 +3658,7 @@ def __init__( allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO, force_dynamic_quantization: bool = False, use_cute_dsl_blockscaling_mm: bool = False, + use_cute_dsl_nvfp4_swiglu_blackwell: bool = False, disable_deep_gemm: bool = False, fused_weight_shard_indices_mapping: Optional[dict] = None, nvfp4_allowed_backends: Optional[List[str]] = None, @@ -3635,6 +3669,9 @@ def __init__( ): """ Args: + use_cute_dsl_nvfp4_swiglu_blackwell: Allow this fused gate/up + projection to use the Blackwell-only NVFP4 GEMM + SwiGLU + kernel and its required interleaved weight layout. nvfp4_allowed_backends: List of backends to consider for NVFP4 GEMM auto-selection. Default (via config): ['cutlass', 'cublaslt', 'cuda_core'] - excludes cutedsl for faster build. Add 'cutedsl' for extreme performance at the cost of longer build time. @@ -3659,6 +3696,8 @@ def __init__( self.gather_output = gather_output self.force_dynamic_quantization = force_dynamic_quantization self.use_cute_dsl_blockscaling_mm = use_cute_dsl_blockscaling_mm + self.use_cute_dsl_nvfp4_swiglu_blackwell = \ + use_cute_dsl_nvfp4_swiglu_blackwell self.disable_deep_gemm = disable_deep_gemm self.fused_weight_shard_indices_mapping = fused_weight_shard_indices_mapping # Store NVFP4 GEMM allowed backends configuration @@ -4007,6 +4046,18 @@ def has_nvfp4(self): return self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( ) + def can_use_cute_dsl_nvfp4_swiglu_blackwell(self) -> bool: + """Return whether this layer can use the Blackwell NVFP4 SwiGLU op. + + Keep this predicate shared by weight transformation and forward + dispatch so a fallback backend never consumes the fused layout. + """ + return (self.use_cute_dsl_nvfp4_swiglu_blackwell + and self.use_cute_dsl_blockscaling_mm + and IS_CUTLASS_DSL_AVAILABLE + and self.has_nvfp4_activation_quantization + and get_sm_version() in (100, 103) and not self.has_bias) + @property def has_nvfp4_activation_quantization(self): assert self._weights_created diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index 28aaadfee290..e29231a972a9 100644 --- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -875,7 +875,7 @@ def add(tactic): if m_tiles * ks <= max_grid_ctas: for bs in _fused_hc_mma_bigfuse_bs_options(M): add(("fused_half_mma", 0, ks, bs, 1)) - if M >= 64: + if M >= 64 and ks in _FUSED_HC_ALL_MMA_KS: add(("fused_all_mma", 0, ks, 0, 1)) if not mma_ok and M > 32: 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/sparse/deepseek_v4/test_deepseek_v4_o_proj.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py index 5ef9a9793c40..fb29bc191cc4 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py @@ -30,7 +30,8 @@ from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module import ( project_sparse_attn_output, ) -from tensorrt_llm._torch.attention.mla import MLA +from tensorrt_llm._torch.attention.mla import MLA, _is_cute_dsl_fp8_bmm_available +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import weight_dequant from tensorrt_llm._utils import get_sm_version @@ -66,6 +67,7 @@ def calculate_reference_deepseek_v4_o_proj( qk_rope_head_dim, device, is_fp8: bool = False, + quantize_o_a_input: bool = False, ): """ Reference implementation for DeepSeek-V4 output projection based on ref/model.py. @@ -80,6 +82,7 @@ def calculate_reference_deepseek_v4_o_proj( qk_rope_head_dim: Dimension of positional part device: Device to run on is_fp8: Whether test fp8 or bf16 + quantize_o_a_input: Whether o_a_proj consumes block-scaled FP8 input Returns: output: [num_tokens, hidden_size] projected output @@ -92,7 +95,7 @@ def calculate_reference_deepseek_v4_o_proj( # Reshape for grouped projection attn_out_grouped = attn_out_latent.view(num_tokens, n_local_groups, -1) - if is_fp8: + if quantize_o_a_input: attn_out_grouped = _per_token_fp8_quant_dequant( attn_out_grouped.transpose(0, 1).contiguous() ).transpose(0, 1) @@ -124,6 +127,7 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): device = torch.device("cuda") dtype = torch.bfloat16 + sm_version = get_sm_version() # Model configuration matching the reference model num_heads = 64 @@ -184,7 +188,10 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): pretrained_config=pretrained_config, sparse_attention_config=sparse_config, quant_config=quant_config, - use_cute_dsl_blockscaling_mm=dtype_str == "fp8", + use_cute_dsl_blockscaling_mm=( + dtype_str == "fp8" + and (sm_version in (100, 103) or (sm_version == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE)) + ), ) # Setup positional embedding params @@ -219,6 +226,12 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): assert not hasattr(mla, "v_b_proj") assert not hasattr(mla, "o_proj") + has_native_fp8_o_proj = _is_cute_dsl_fp8_bmm_available(sm_version) + if dtype_str == "fp8" and has_native_fp8_o_proj: + assert mla.o_a_proj.dtype == torch.float8_e4m3fn + assert mla.o_a_proj_dequant is None + assert not mla.use_cute_dsl_blockscaling_bmm + # Initialize weights nn_init_std = 0.02 with torch.no_grad(): @@ -245,12 +258,8 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): fp8_a_weight = fp8_a_weight.reshape(n_local_groups, o_lora_rank, dim) mla.o_a_proj.data = fp8_a_weight mla.o_a_proj_scale.data = fp8_a_scale - # mla.o_a_proj_dequant is None for DSv4 on SM100: PR #14254 - # decouples the FP8-native o_a_proj path from - # use_cute_dsl_blockscaling_bmm, so DSv4 unconditionally uses the - # fused inv-RoPE + FP8 quant + cute-dsl BMM chain and never needs - # the bf16-dequant fallback buffer. The reference path below uses - # o_a_proj_bf16 directly. + if not has_native_fp8_o_proj: + mla.o_a_proj_dequant.data = o_a_proj_bf16 # Initialize o_b_proj weights if dtype_str == "bf16": @@ -289,16 +298,19 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): o_a_proj_ref = mla.o_a_proj.data o_b_proj_weight_ref = mla.o_b_proj.weight.data else: - # Match the FP8-native o_a_proj path: the runtime BMM consumes - # quantized o_a_proj plus block scales, not the original BF16 weight. - o_a_proj_ref = ( - weight_dequant( - fp8_a_weight.reshape(-1, dim).contiguous(), - fp8_a_scale.contiguous(), + if has_native_fp8_o_proj: + # Match the native path, which consumes the quantized weight and + # its block scales rather than the original BF16 values. + o_a_proj_ref = ( + weight_dequant( + fp8_a_weight.reshape(-1, dim).contiguous(), + fp8_a_scale.contiguous(), + ) + .bfloat16() + .reshape(o_a_proj_bf16.shape) ) - .bfloat16() - .reshape(o_a_proj_bf16.shape) - ) + else: + o_a_proj_ref = o_a_proj_bf16 o_b_proj_weight_ref = fp8_b_weight_dequant freqs_cis = precompute_freqs_cis( @@ -321,6 +333,7 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): qk_rope_head_dim=qk_rope_head_dim, device=device, is_fp8=dtype_str == "fp8", + quantize_o_a_input=dtype_str == "fp8" and has_native_fp8_o_proj, ) # Validate output shapes diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py new file mode 100644 index 000000000000..b2afc00ee238 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise prior-state ownership through the actual indexer forward path.""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer +from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"), + ), + ], +) +def test_update_gvr_prior_from_prefill_uses_device_lengths(device: str) -> None: + """Seed the last prefill row using lengths on the selections' device.""" + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR, gvr_self_sampling=False) + prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32, device=device) + prior_indices = torch.zeros(3, 2, dtype=torch.int32, device=device) + + # Production passes the device seq_lens twin so the row gather stays async. + top_k.update_gvr_prior_from_prefill( + prefill_indices, + torch.tensor([2, 1], dtype=torch.int32, device=device), + prior_indices, + request_offset=1, + ) + + assert prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]] + assert top_k.needs_gvr_prior + + +@pytest.mark.parametrize( + "implementation,self_sampling", + [ + (TopKImplementation.CUDA_RADIX, False), + (TopKImplementation.CUTE_DSL_RADIX, False), + (TopKImplementation.CUTE_DSL_GVR, True), + (TopKImplementation.CUTE_DSL_GVR, False), + ], +) +@pytest.mark.parametrize("phase", ["prefill", "decode", "mixed", "split_prefill", "split_decode"]) +@pytest.mark.parametrize("next_n", [1, 4]) +def test_indexer_forward_uses_prior_only_for_temporal_gvr( + implementation: TopKImplementation, self_sampling: bool, phase: str, next_n: int +) -> None: + """Preserve prior ownership and pass the device lengths through each phase.""" + topk = 2 + top_k = TopK(topk, decode_implementation=implementation, gvr_self_sampling=self_sampling) + num_contexts = 0 if phase == "decode" else 2 + num_generations = 0 if phase == "prefill" else 2 + num_ctx_tokens = 5 if num_contexts else 0 + num_gen_tokens = num_generations * next_n + total_tokens = num_ctx_tokens + num_gen_tokens + is_generation = {"split_prefill": False, "split_decode": True}.get(phase) + has_prefill = num_contexts > 0 and is_generation is not True + has_decode = num_generations > 0 and is_generation is not False + input_tokens = ( + total_tokens + if is_generation is None + else (num_gen_tokens if is_generation else num_ctx_tokens) + ) + cache_manager = SimpleNamespace( + quant_block_size=128, + layer_offsets={7: 1}, + get_indexer_k_cache_buffers=Mock(return_value=torch.empty(0)), + ) + metadata = SimpleNamespace( + kv_cache_manager=cache_manager, + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_ctx_tokens, + num_tokens=total_tokens, + seq_lens=torch.tensor(([2, 3] if num_contexts else []) + [next_n] * num_generations), + cuda_graph_buffers={}, + is_cuda_graph=False, + get_empty=lambda buffers, shape, **kwargs: torch.full(shape, -1, dtype=torch.int32), + skip_indexer_for_ctx_reqs=False, + skip_indexer_for_gen_reqs=False, + indexer_prefill_chunks=None, + num_ctx_kv_tokens=16, + cu_seqlen_ks=torch.zeros(num_ctx_tokens, dtype=torch.int32), + cu_seqlen_ke=torch.full((num_ctx_tokens,), 16, dtype=torch.int32), + use_expanded_buffers_for_mtp=False, + kv_lens_cuda_2d=torch.full((num_generations, next_n), 16, dtype=torch.int32), + indexer_k_cache_block_offsets=torch.empty((num_contexts + num_generations, 1)), + scheduler_metadata_buffer=None, + scheduler_metadata_buffer_full_next_n=None, + get_indexer_max_seq_len=lambda: 16, + kv_lens_cuda_runtime=torch.full((num_contexts + num_generations,), 16), + gen_indexer_kv_lens_cuda_runtime=torch.full((num_generations,), 16), + kv_lens_row_reorder=None, + ) + # Keep distinct storage so the caller test catches use of the host twin. + metadata.seq_lens_cuda = metadata.seq_lens.clone() + # Radix fallback and self-sampling deliberately have no prior attribute. + # The raw heuristic option remains True, as in the SM107 failure. + if top_k.needs_gvr_prior: + metadata.gvr_prior_indices = torch.full( + (2, num_contexts + num_generations, topk), -99, dtype=torch.int32 + ) + + def logits(q: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: + return torch.ones((q.shape[0], 16)) + + def selections(rows: int) -> torch.Tensor: + return torch.arange(rows * topk, dtype=torch.int32).reshape(rows, topk) + + def decode( + scores: torch.Tensor, + sequence_lengths: torch.Tensor, + scan_lengths: torch.Tensor, + output: torch.Tensor, + n: int, + max_seq_len: int, + extra: dict[str, torch.Tensor | None], + ) -> torch.Tensor: + prior = extra["gvr_prior_indices"] + if top_k.needs_gvr_prior: + assert prior.shape == (num_generations, topk) + assert prior.data_ptr() == metadata.gvr_prior_indices[1].data_ptr() + assert torch.all(prior == -99) + else: + assert prior is None + assert n == next_n + output.copy_(selections(num_gen_tokens)) + return output + + def prefill( + scores: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, + output: torch.Tensor, + k: int, + ) -> None: + output.copy_(selections(num_ctx_tokens)) + + indexer = SimpleNamespace( + top_k=top_k, + index_topk=topk, + layer_idx=7, + _enable_heuristic_topk=True, + mtp_index_share=False, + use_fp4=False, + use_cute_dsl_paged_mqa_logits=False, + _call_mqa_logits=logits, + _call_paged_mqa_logits=Mock(return_value=torch.ones((num_gen_tokens, 16))), + aux_stream=None, + ) + with ( + patch.object( + top_k, "update_gvr_prior_from_prefill", wraps=top_k.update_gvr_prior_from_prefill + ) as seed_call, + patch.object(top_k, "_forward_decode", side_effect=decode) as decode_call, + patch.object(torch.ops.trtllm, "indexer_topk_prefill", side_effect=prefill), + ): + result = Indexer.sparse_attn_indexer( + indexer, + metadata, + torch.empty((input_tokens, 2)), + torch.empty((input_tokens, 1, 2)), + torch.empty((16, 2)), + torch.empty(16), + torch.empty(input_tokens), + is_generation=is_generation, + ) + assert decode_call.call_count == int(has_decode) + assert seed_call.call_count == int(top_k.needs_gvr_prior and has_prefill) + if seed_call.called: + assert seed_call.call_args.args[1].data_ptr() == metadata.seq_lens_cuda.data_ptr() + token_offset = num_ctx_tokens if is_generation is None else 0 + if has_prefill: + torch.testing.assert_close(result[:num_ctx_tokens], selections(num_ctx_tokens)) + if has_decode: + torch.testing.assert_close(result[token_offset:], selections(num_gen_tokens)) + if top_k.needs_gvr_prior: + expected = torch.full_like(metadata.gvr_prior_indices, -99) + if has_prefill: + expected[1, num_generations:] = selections(num_ctx_tokens)[[1, 4]] + if has_decode: + expected[1, :num_generations] = selections(num_gen_tokens)[next_n - 1 :: next_n] + torch.testing.assert_close(metadata.gvr_prior_indices, expected) + + +@pytest.mark.parametrize("batch_size", [1, 8]) +@pytest.mark.parametrize("next_n", [1, 4]) +@pytest.mark.parametrize("score_width", [8192, 262144]) +def test_self_sampling_gpu_exactness_and_graph_replay( + batch_size: int, next_n: int, score_width: int +) -> None: + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + from tensorrt_llm._utils import get_sm_version + + if ( + not torch.cuda.is_available() + or not IS_CUTLASS_DSL_AVAILABLE + or get_sm_version() not in (100, 103, 107) + ): + pytest.skip("Self-sampling GVR requires datacenter Blackwell or Rubin and CuTe DSL") + + topk = 512 + compress_ratio = 4 + rows = batch_size * next_n + generator = torch.Generator(device="cuda").manual_seed(42) + scores = torch.randn((rows, score_width), generator=generator, device="cuda") + lengths = score_width * compress_ratio - 68 * torch.arange( + batch_size, dtype=torch.int32, device="cuda" + ) + row_ids = torch.arange(rows, device="cuda") + valid_columns = (lengths[row_ids // next_n] - next_n + row_ids % next_n + 1) // compress_ratio + valid = torch.arange(score_width, device="cuda")[None, :] < valid_columns[:, None] + # Invalid tail values must never enter the result. + scores.masked_fill_(~valid, 1.0e6) + output = torch.empty((rows, topk), dtype=torch.int32, device="cuda") + top_k = TopK( + topk, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=compress_ratio, + gvr_self_sampling=True, + ) + assert not top_k.needs_gvr_prior + + def run() -> None: + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths // compress_ratio, + next_n=next_n, + max_seq_len=score_width, + ) + + def check_result() -> None: + assert torch.all((output >= 0) & (output < valid_columns[:, None])) + sorted_indices = output.sort(dim=-1).values + assert torch.all(sorted_indices[:, 1:] != sorted_indices[:, :-1]) + actual = scores.gather(1, output.long()).sort(dim=-1, descending=True).values + expected = scores.masked_fill(~valid, float("-inf")).topk(topk, dim=-1).values + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + run() + torch.cuda.synchronize() + check_result() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + output.fill_(-1) + graph.replay() + torch.cuda.synchronize() + check_result() diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py b/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py new file mode 100644 index 000000000000..6393f56dff26 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Regression coverage for TopK dispatch during DSA metadata initialization.""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.dsa import metadata as dsa_metadata +from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSAMetadataParams + + +@pytest.mark.parametrize("enable_heuristic", [False, True]) +@pytest.mark.parametrize("use_self_sampling", [False, True]) +@pytest.mark.parametrize("dsl_available", [False, True]) +@pytest.mark.parametrize("sm_version", [90, 100, 103, 107]) +def test_topk_flags_initialized_before_buffer_allocation( + enable_heuristic: bool, + use_self_sampling: bool, + dsl_available: bool, + sm_version: int, +) -> None: + metadata = object.__new__(dsa_metadata.DSAtrtllmAttentionMetadata) + metadata.sparse_metadata_params = DSAMetadataParams( + indexer_max_chunk_size=8192, + max_sparse_topk=512, + index_head_dim=128, + enable_indexer_skip=False, + enable_heuristic_topk=enable_heuristic, + use_cute_dsl_topk=True, + use_cute_dsl_paged_mqa_logits=False, + q_split_threshold=8192, + use_self_sampling_topk=use_self_sampling, + ) + metadata.kv_cache_manager = SimpleNamespace( + tokens_per_block=128, + compressed_block_sizes={}, + get_cache_indices=Mock(), + ) + metadata.is_cuda_graph = False + enabled = enable_heuristic and sm_version >= 100 + self_sampling_supported = enabled and dsl_available and sm_version in (100, 103, 107) + temporal_supported = enabled and dsl_available and sm_version in (100, 103) + + def check_flags(*, capture_graph: bool) -> None: + assert capture_graph is False + assert metadata.enable_gvr_topk is enabled + assert metadata.use_self_sampling_topk is (self_sampling_supported and use_self_sampling) + assert metadata.needs_gvr_prior is (temporal_supported and not use_self_sampling) + + metadata.create_buffers_for_mla_rope_append = Mock(side_effect=check_flags) + metadata.create_buffers_for_indexer = Mock(side_effect=check_flags) + + with ( + patch.object(dsa_metadata.TrtllmAttentionMetadata, "__post_init__"), + patch.object(dsa_metadata, "IS_CUTLASS_DSL_AVAILABLE", dsl_available), + patch.object(dsa_metadata, "get_sm_version", return_value=sm_version), + ): + metadata.__post_init__() + + metadata.create_buffers_for_mla_rope_append.assert_called_once_with(capture_graph=False) + metadata.create_buffers_for_indexer.assert_called_once_with(capture_graph=False) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("num_local_layers", [1, 3]) +def test_temporal_gvr_allocates_real_prior_buffers(num_local_layers: int) -> None: + """Allocate actual CUDA buffers and size the zeroed prior by local layers.""" + metadata = object.__new__(dsa_metadata.DSAtrtllmAttentionMetadata) + metadata.sparse_metadata_params = DSAMetadataParams( + indexer_max_chunk_size=32, + max_sparse_topk=512, + index_head_dim=128, + enable_indexer_skip=False, + enable_heuristic_topk=True, + use_cute_dsl_topk=True, + use_cute_dsl_paged_mqa_logits=False, + q_split_threshold=32, + use_self_sampling_topk=False, + ) + metadata.kv_cache_manager = SimpleNamespace( + tokens_per_block=128, + compressed_block_sizes={}, + get_cache_indices=Mock(), + max_blocks_per_seq=2, + num_local_layers=num_local_layers, + ) + metadata.draft_kv_cache_manager = None + metadata.is_cuda_graph = False + metadata.cuda_graph_buffers = None + metadata.max_num_sequences = 4 + metadata.max_num_tokens = 16 + metadata.max_draft_tokens = 3 + metadata.num_sms = 16 + metadata.enable_context_mla_with_cached_kv = False + with ( + patch.object(dsa_metadata.TrtllmAttentionMetadata, "__post_init__"), + patch.object(metadata, "create_buffers_for_mla_rope_append"), + patch.object(dsa_metadata, "IS_CUTLASS_DSL_AVAILABLE", True), + patch.object(dsa_metadata, "get_sm_version", return_value=100), + ): + metadata.__post_init__() + + assert metadata.enable_gvr_topk + assert metadata.needs_gvr_prior + assert metadata.gvr_prior_indices.shape == (num_local_layers, 4, 512) + assert metadata.gvr_prior_indices.dtype == torch.int32 + assert metadata.gvr_prior_indices.is_cuda + assert torch.count_nonzero(metadata.gvr_prior_indices).item() == 0 + assert metadata.kv_lens_row_reorder_buffer.shape == (4,) + + +def test_selfsampling_warmup_ignores_retired_hinted_env(monkeypatch) -> None: + """Warm the hint-free launcher even when the retired hinted switch is set.""" + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( + gvr_topk_decode_self_sampling_host as host, + ) + + metadata = object.__new__(dsa_metadata.DSAtrtllmAttentionMetadata) + metadata.enable_gvr_topk = True + metadata.use_self_sampling_topk = True + metadata.kv_cache_manager = Mock() + metadata.sparse_mla_topk = 512 + metadata._indexer_compress_ratio = 4 + metadata.sparse_metadata_params = SimpleNamespace(use_cute_dsl_paged_mqa_logits=True) + metadata.get_indexer_max_seq_len = Mock(return_value=8192) + monkeypatch.setenv("TRTLLM_GVR_V2_HINTED", "1") + with ( + patch.object(dsa_metadata, "IS_CUTLASS_DSL_AVAILABLE", True), + patch.object(dsa_metadata, "get_sm_version", return_value=107), + patch.object(host, "warmup_varlen") as warmup, + ): + metadata.warmup_selfsampling_topk(next_n=4, batch_sizes=[64]) + warmup.assert_called_once() + assert warmup.call_args.kwargs["hint_free"] is True + assert 256 in warmup.call_args.kwargs["num_rows_list"] diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py index 642a41010074..9f52acd5bd04 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -2403,7 +2403,9 @@ def yarn_get_mscale(scale=1, mscale=1): @pytest.mark.skipif(get_sm_version() < 100, reason="DSv4 fused FP8 Q-quant requires SM100") @pytest.mark.parametrize("batch_name", _FUSED_Q_FP8_PREFILL_BATCHES) -def test_forward_sparse_mla_unified_fused_q_fp8(monkeypatch, batch_name): +@pytest.mark.parametrize("q_rope_applied", [False, True]) +def test_forward_sparse_mla_unified_fused_q_fp8(monkeypatch, batch_name, + q_rope_applied): """Regression test for the sparse-MLA context-branch fused FP8 Q-quant wiring. The wo-linear helper bypasses `_q_branch`, so we monkey-patch `forward_context_sparse_attn` to manually populate the fused buffers @@ -2430,16 +2432,41 @@ def patched_fwd(self, q, compressed_kv, k_pe, attn_metadata, output, quant_q_buffer = torch.empty((num_tokens, num_heads * head_dim), dtype=torch.float8_e4m3fn, device=q.device) - quant_q_buffer.view( - num_tokens, num_heads, - head_dim)[:, :, :nope_dim] = q_view[:, :, :nope_dim].float().to( + quant_q_view = quant_q_buffer.view(num_tokens, num_heads, head_dim) + quant_q_view[:, :, :nope_dim] = q_view[:, :, :nope_dim].float().to( + torch.float8_e4m3fn) + if q_rope_applied: + position_ids = kwargs["position_ids"].reshape(-1).long() + rope_dim = head_dim - nope_dim + cache = self.mqa.rotary_cos_sin.view(-1, rope_dim, 2) + coefficient = cache[position_ids, :rope_dim // 2].unsqueeze(1) + q_pe = q_view[:, :, nope_dim:].float() + q0, q1 = q_pe[..., 0::2], q_pe[..., 1::2] + cos, sin = coefficient[..., 0], coefficient[..., 1] + rotated = torch.stack( + [cos * q0 - sin * q1, cos * q1 + sin * q0], + dim=-1).flatten(-2) + quant_q_view[:, :, nope_dim:] = rotated.to(q.dtype).to( torch.float8_e4m3fn) + self._fused_q_pe = None + else: + self._fused_q_pe = q_view[:, :, nope_dim:].contiguous() self._fused_quant_q_buffer = quant_q_buffer - self._fused_q_pe = q_view[:, :, nope_dim:].contiguous() + self._fused_q_rope_applied = q_rope_applied self._quant_scale_qkv = torch.tensor([1.0], dtype=torch.float32, device=q.device) - q = torch.full_like(q, float('nan')) + if q_rope_applied: + # q_b GEMM fusion carries only compact q_lora once the complete + # rotated Q has been written to quant_q_buffer. + q = torch.full( + (num_tokens, self.q_lora_rank), + float("nan"), + dtype=q.dtype, + device=q.device, + ) + else: + q = torch.full_like(q, float("nan")) return original_fwd(self, q, compressed_kv, k_pe, attn_metadata, output, **kwargs) diff --git a/tests/unittest/_torch/executor/test_indexer_workspace_reserve.py b/tests/unittest/_torch/executor/test_indexer_workspace_reserve.py new file mode 100644 index 000000000000..06d238c53776 --- /dev/null +++ b/tests/unittest/_torch/executor/test_indexer_workspace_reserve.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tensorrt_llm._torch.attention.backends.sparse.params import ( + get_indexer_mqa_logits_elem_budget, + get_indexer_mqa_logits_workspace_bytes, +) + + +def test_workspace_bytes_match_runtime_element_budget(monkeypatch): + monkeypatch.setenv("TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET", "1024") + assert get_indexer_mqa_logits_elem_budget() == 1024 + assert get_indexer_mqa_logits_workspace_bytes() == 4096 + + +def test_workspace_bytes_are_bounded_by_reachable_request_shape(monkeypatch): + monkeypatch.setenv("TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET", str(1 << 31)) + assert get_indexer_mqa_logits_workspace_bytes(4096, 4096) == 64 * 1024 * 1024 diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 8c253ba15026..f940d5fe7910 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -697,6 +697,7 @@ def fake_decoder_layer_init(self, model_config, *_args, **_kwargs): pretrained_config=config, mapping=Mapping(world_size=4, rank=2, tp_size=4), quant_config=quant_config, + use_cute_dsl_blockscaling_mm=True, ) mtp_layer = DeepseekV4MTP( @@ -715,6 +716,8 @@ def fake_decoder_layer_init(self, model_config, *_args, **_kwargs): assert mtp_layer.h_proj.out_features == config.hidden_size assert mtp_layer.e_proj.reduce_output is True assert mtp_layer.h_proj.reduce_output is True + assert mtp_layer.e_proj.use_cute_dsl_blockscaling_mm is True + assert mtp_layer.h_proj.use_cute_dsl_blockscaling_mm is True assert mtp_layer.e_proj.weight.dtype is torch.float8_e4m3fn assert mtp_layer.h_proj.weight.dtype is torch.float8_e4m3fn assert hasattr(mtp_layer.e_proj, "weight_scale") diff --git a/tests/unittest/_torch/modules/test_gated_mlp.py b/tests/unittest/_torch/modules/test_gated_mlp.py new file mode 100644 index 000000000000..2db3f23738b5 --- /dev/null +++ b/tests/unittest/_torch/modules/test_gated_mlp.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from typing import Callable +from unittest.mock import Mock + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from tensorrt_llm._torch.modules import gated_mlp as gated_mlp_module +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP + + +def _make_gate_up_proj( + projected: torch.Tensor, + *, + partitioned: bool, +) -> nn.Module: + gate_up_proj = nn.Module() + gate_up_proj.has_nvfp4 = True + gate_up_proj.has_bias = False + gate_up_proj.partition_plan = SimpleNamespace(enabled=partitioned) + gate_up_proj.can_use_cute_dsl_nvfp4_swiglu_blackwell = Mock(return_value=not partitioned) + gate_up_proj.forward = Mock(return_value=projected) + return gate_up_proj + + +def _make_down_proj() -> nn.Module: + down_proj = nn.Module() + down_proj.has_fp8_qdq = False + down_proj.has_w4a8_nvfp4_fp8 = False + down_proj.forward = Mock(side_effect=lambda value, **kwargs: value + 1) + return down_proj + + +def test_gate_up_partition_falls_back_to_swiglu( + monkeypatch: pytest.MonkeyPatch, +) -> None: + mlp = GatedMLP(hidden_size=2, intermediate_size=2, bias=False) + mlp.use_cute_dsl_blockscaling_mm = True + projected = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + mlp.gate_up_proj = _make_gate_up_proj(projected, partitioned=True) + mlp.down_proj = _make_down_proj() + + gate, up = projected.chunk(2, dim=-1) + activated = F.silu(gate) * up + swiglu = Mock(side_effect=lambda value, **kwargs: F.silu(value[..., :2]) * value[..., 2:]) + monkeypatch.setattr(gated_mlp_module, "swiglu", swiglu) + fused_gate_up_swiglu = Mock(side_effect=AssertionError("Blackwell fused op must not run")) + monkeypatch.setattr(mlp, "_fused_gate_up_swiglu", fused_gate_up_swiglu) + + inputs = torch.tensor([[5.0, 6.0]]) + output = mlp(inputs) + + assert not mlp._can_fuse_gate_up_swiglu() + mlp.gate_up_proj.forward.assert_called_once_with(inputs) + swiglu.assert_called_once_with(projected, swiglu_limit=None) + mlp.down_proj.forward.assert_called_once() + down_args, down_kwargs = mlp.down_proj.forward.call_args + torch.testing.assert_close(down_args[0], activated) + assert down_kwargs == {"all_reduce_params": None, "layer_idx": None} + fused_gate_up_swiglu.assert_not_called() + torch.testing.assert_close(output, activated + 1) + + +@pytest.mark.parametrize( + "swiglu_limit, expected", + [(None, True), (float("inf"), True), (7.0, False)], +) +def test_swiglu_limit_controls_gate_up_fusion_capability( + swiglu_limit: float | None, + expected: bool, +) -> None: + mlp = GatedMLP( + hidden_size=2, + intermediate_size=2, + bias=False, + swiglu_limit=swiglu_limit, + use_cute_dsl_blockscaling_mm=True, + ) + + assert mlp.gate_up_proj.use_cute_dsl_nvfp4_swiglu_blackwell is expected + + +@pytest.mark.parametrize( + ("activation", "expected"), + [ + pytest.param(F.silu, True, id="plain-swiglu"), + pytest.param(lambda value: value, False, id="custom-swiglu-oai"), + ], +) +def test_activation_controls_fp8_quant_fusion_capability( + monkeypatch: pytest.MonkeyPatch, + activation: Callable[[torch.Tensor], torch.Tensor], + expected: bool, +) -> None: + mlp = GatedMLP( + hidden_size=2, + intermediate_size=2, + bias=False, + activation=activation, + ) + down_proj = nn.Module() + down_proj.has_fp8_block_scales = True + down_proj.use_cute_dsl_blockscaling_mm = True + down_proj.disable_deep_gemm = False + mlp.down_proj = down_proj + monkeypatch.setattr(gated_mlp_module, "get_sm_version", lambda: 107) + monkeypatch.setattr(gated_mlp_module, "IS_CUTLASS_DSL_RUBIN_AVAILABLE", True) + + assert mlp._can_fuse_swiglu_fp8_quant() is expected diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 69600d69f043..9ff195fbf504 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -64,6 +64,31 @@ def swiglu_ref(x: torch.Tensor, swiglu_limit: float = float("inf")) -> torch.Ten return x * torch.nn.functional.silu(gate) +# Kimi K3 SiTU constants; deliberately non-unit and unequal, so that a swapped +# beta / linear_beta or a dropped constant cannot pass the test. +SITU_BETA = 2.5 +SITU_LINEAR_BETA = 0.7 + + +def situ_ref( + x: torch.Tensor, beta: float = SITU_BETA, linear_beta: float = SITU_LINEAR_BETA +) -> torch.Tensor: + """SiTU reference, in this kernel's ``[linear | gate]`` chunk order. + + ``tensorrt_llm._torch.modules.kimi_k3_moe._mlp.SituAndMul`` is the golden + definition, but it consumes the model's ``[gate | up]`` packing -- the two + halves in the opposite order from ``swiglu_ref`` and from the interleaved + FC1 weight this kernel reads. Keep the order explicit: silently swapping + the halves still produces plausible-looking numbers. + """ + up, gate = x.chunk(2, dim=-1) + gate = gate.to(torch.float32) + up = up.to(torch.float32) + situ_a = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + up = linear_beta * torch.tanh(up / linear_beta) + return (situ_a * up).to(x.dtype) + + def apply_activation_ref( x: torch.Tensor, activation_type: ActivationType, swiglu_limit: float = float("inf") ) -> torch.Tensor: @@ -71,6 +96,8 @@ def apply_activation_ref( return swiglu_ref(x, swiglu_limit) if activation_type == ActivationType.Relu2: return relu2(x) + if activation_type == ActivationType.SiTu: + return situ_ref(x) raise ValueError(f"Unsupported activation_type: {activation_type}") @@ -1136,8 +1163,8 @@ def test_nvfp4_grouped_gemm_swiglu_blackwell( @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available") @pytest.mark.parametrize( "activation_type", - [ActivationType.Swiglu, ActivationType.Relu2], - ids=["swiglu", "relu2"], + [ActivationType.Swiglu, ActivationType.Relu2, ActivationType.SiTu], + ids=["swiglu", "relu2", "situ"], ) @pytest.mark.parametrize("tile_size", [128, 256]) @pytest.mark.parametrize("ep_size", [1, 8, 32]) @@ -1155,11 +1182,14 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( This test validates the gather kernel which: 1. Uses LDGSTS for A/SFA loading with permuted_idx_to_expanded_idx 2. Performs GEMM with (interleaved for gated) weights - 3. Applies the fused activation (SwiGLU for gated, Relu2 for non-gated) + 3. Applies the fused activation (SwiGLU / SiTU for gated, Relu2 for + non-gated) 4. Quantizes output to FP4 with scale factor generation """ is_gated = is_gated_activation(activation_type) - swiglu_limit = 1.0 if is_gated else float("inf") + is_situ = activation_type == ActivationType.SiTu + # SiTU rejects the SwiGLU clamp (matching MegaMoE and DeepGEMM). + swiglu_limit = 1.0 if (is_gated and not is_situ) else float("inf") weight_n_multiplier = 2 if is_gated else 1 sf_vec_size = 16 hidden_size = 4096 @@ -1301,6 +1331,8 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( scaling_vector_size=sf_vec_size, activation_type=activation_type, swiglu_limit_scalar=swiglu_limit, + situ_beta=SITU_BETA if is_situ else -1.0, + situ_linear_beta=SITU_LINEAR_BETA if is_situ else -1.0, ) # Verify output (only compare valid tokens, skip padding tokens where permuted_idx_to_expanded_idx == -1) @@ -1538,6 +1570,204 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_rubin( check_accuracy(c_sf_valid, c_sf_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on SM 107 (Rubin) GPUs", +) +@pytest.mark.parametrize("tile_size", [128, 256]) +def test_nvfp4_gather_grouped_gemm_situ_rubin(tile_size: int): + """SiTU epilogue on the Rubin NVFP4 gather+grouped GEMM kernel. + + Reduced grid vs the SwiGLU test: enough to compile both tile sizes and + catch a dropped / swapped beta. The Blackwell test documents why the + swap check is restricted to (128 tokens, top_k=1, ep=1). + """ + num_tokens = 128 + top_k = 1 + ep_size = 1 + sf_vec_size = 16 + hidden_size = 4096 + interm_size = 8192 + num_experts = 256 + num_local_experts = num_experts // ep_size + + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + token_selected_experts[0] = 0 + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + num_valid_permuted_tokens = total_num_padded_tokens.item() + + a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.int32, device="cuda").to( + torch.bfloat16 + ) + b = torch.randint( + -5, + 5, + (num_local_experts, interm_size * 2, hidden_size), + dtype=torch.int32, + device="cuda", + ).to(torch.bfloat16) + + a_global_sf = a.abs().max().float() / (448 * 6) + b_global_sf = b.abs().amax(dim=(1, 2)).float() / (448 * 6) + a, a_sf = torch.ops.trtllm.fp4_quantize(a, 1 / a_global_sf, sf_vec_size, False) + a = a.view(torch.float4_e2m1fn_x2) + a_sf_unswizzled = unswizzle_sf(a_sf, (num_tokens + 127) // 128 * 128, hidden_size)[:num_tokens] + b, b_sf = torch.ops.trtllm.fp4_quantize(b, 1 / b_global_sf, sf_vec_size, False) + b = b.view(torch.float4_e2m1fn_x2) + b_sf = b_sf.view(num_local_experts, interm_size * 2, hidden_size // sf_vec_size) + alpha = a_global_sf * b_global_sf + + b_interleaved = interleave_linear_and_gate(b.view(torch.uint8), group_size=64, dim=1).view( + torch.float4_e2m1fn_x2 + ) + b_sf_unswizzled = unswizzle_sf(b_sf, interm_size * 2, hidden_size).view( + num_local_experts, interm_size * 2, hidden_size // sf_vec_size + ) + b_sf_unswizzled_interleaved = interleave_linear_and_gate(b_sf_unswizzled, group_size=64, dim=1) + b_sf_interleaved = swizzle_sf(b_sf_unswizzled_interleaved, interm_size * 2, hidden_size).view( + num_local_experts, interm_size * 2, hidden_size // sf_vec_size + ) + + permuted_idx_to_expanded_idx_list = permuted_idx_to_expanded_idx.cpu().tolist() + tile_idx_to_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist() + + a_gathered = torch.empty(max_num_permuted_tokens, hidden_size // 2, dtype=a.dtype) + a_sf_gathered = torch.empty( + max_num_permuted_tokens, hidden_size // sf_vec_size, dtype=a_sf.dtype + ) + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + expanded_idx = permuted_idx_to_expanded_idx_list[i] + token_id = expanded_idx // top_k + a_gathered[i] = a[token_id] + a_sf_gathered[i] = a_sf_unswizzled[token_id] + a_gathered = a_gathered.to(a.device) + a_sf_gathered = a_sf_gathered.to(a.device) + + a_sf_gathered_swizzled = swizzle_sf( + a_sf_gathered.view(max_num_permuted_tokens, hidden_size // sf_vec_size), + max_num_permuted_tokens, + hidden_size, + ) + + c_ref = cute_dsl_nvfp4_grouped_gemm_ref( + a_gathered, + b, + a_sf_gathered_swizzled, + b_sf, + alpha, + tile_idx_to_group_idx, + num_non_exiting_tiles, + tile_size=tile_size, + output_dtype=torch.bfloat16, + scaling_vector_size=sf_vec_size, + ) + c_ref = situ_ref(c_ref) + global_sf = c_ref[:num_valid_permuted_tokens].abs().max().float() / (448 * 6) + c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False) + + op_kwargs = dict( + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_tensor=None, + output_sf_tensor=None, + scaling_vector_size=sf_vec_size, + activation_type=ActivationType.SiTu, + situ_beta=SITU_BETA, + situ_linear_beta=SITU_LINEAR_BETA, + ) + c, c_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + a, + b_interleaved, + a_sf_unswizzled, + b_sf_interleaved, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + torch.tensor([1 / global_sf], dtype=torch.float32, device="cuda"), + **op_kwargs, + ) + + valid_token_mask = torch.zeros(num_valid_permuted_tokens, dtype=torch.bool, device="cuda") + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + valid_token_mask[i] = True + + num_valid_tokens = valid_token_mask.sum().item() + if num_valid_tokens > 0: + c_valid = c[:num_valid_permuted_tokens].view(torch.uint8)[valid_token_mask] + c_ref_valid = c_ref[:num_valid_permuted_tokens][valid_token_mask] + check_accuracy(c_valid, c_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) + + c_sf_unswizzled = unswizzle_sf(c_sf, max_num_permuted_tokens, interm_size, sf_vec_size) + c_sf_ref_unswizzled = unswizzle_sf( + c_sf_ref, max_num_permuted_tokens, interm_size, sf_vec_size + ) + + c_sf_valid = [] + c_sf_ref_valid = [] + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + c_sf_valid.append(c_sf_unswizzled[i]) + c_sf_ref_valid.append(c_sf_ref_unswizzled[i]) + + c_sf_valid = torch.cat(c_sf_valid) + c_sf_ref_valid = torch.cat(c_sf_ref_valid) + check_accuracy(c_sf_valid, c_sf_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) + + c_swapped, _ = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + a, + b_interleaved, + a_sf_unswizzled, + b_sf_interleaved, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + torch.tensor([1 / global_sf], dtype=torch.float32, device="cuda"), + **{ + **op_kwargs, + "situ_beta": SITU_LINEAR_BETA, + "situ_linear_beta": SITU_BETA, + }, + ) + swapped_valid = c_swapped[:num_valid_permuted_tokens].view(torch.uint8)[valid_token_mask] + assert not torch.equal(c_valid, swapped_valid), ( + "Swapping situ_beta and situ_linear_beta produced a bit-identical " + "output; the kernel is ignoring one of the SiTU constants." + ) + + @pytest.mark.skipif( get_sm_version() != 107, reason="This test is only supported on SM 107 (Rubin) GPUs", diff --git a/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py b/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py index 41b494e355ae..667fad95e691 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py @@ -15,6 +15,7 @@ import os import subprocess import sys +from unittest import mock import pytest import torch @@ -24,6 +25,7 @@ import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops from tensorrt_llm._torch.cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) @@ -470,15 +472,57 @@ def test_cute_dsl_mxfp8_gemm_rubin_k128_replicated_scales(): output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( a_fp8, b_fp8, a_sf, b_sf) - output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin(a_fp8, b_fp8, a_sf, - b_sf) expected = a @ b.t() + alpha = cute_dsl_custom_ops._get_mxfp8_gemm_alpha(a.device) + seen_alphas = [] + real_get_alpha = cute_dsl_custom_ops._get_mxfp8_gemm_alpha + + def _get_alpha_spy(device): + cached = real_get_alpha(device) + seen_alphas.append(cached) + return cached + + # After the one-time warmup above, the steady-state call must reuse the + # cached scalar instead of filling a fresh one. + with mock.patch.object( + cute_dsl_custom_ops, + "_get_mxfp8_gemm_alpha", + side_effect=_get_alpha_spy, + ), mock.patch.object( + torch, + "ones", + side_effect=AssertionError("unexpected per-call fill"), + ): + output = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + a_fp8, b_fp8, a_sf, b_sf) + + assert len(seen_alphas) == 1 + assert all(seen is alpha for seen in seen_alphas) + assert {seen.data_ptr() for seen in seen_alphas} == {alpha.data_ptr()} diff = calc_diff(output, expected) assert diff < 1e-3 torch.testing.assert_close(output, expected, atol=1e-3, rtol=1e-3) +def test_mxfp8_alpha_cache_rejects_first_init_during_capture(): + """A capture cannot perform the per-device cache's one-time allocation.""" + device = torch.device("cuda", 31415) + assert device not in cute_dsl_custom_ops._MXFP8_GEMM_ALPHA_CACHE + + with mock.patch.object(torch.cuda, + "device") as device_guard, mock.patch.object( + torch.cuda, + "is_current_stream_capturing", + return_value=True): + with pytest.raises(RuntimeError, + match="run one eager GEMM warmup first"): + cute_dsl_custom_ops._get_mxfp8_gemm_alpha(device) + + device_guard.assert_called_once_with(device) + assert device not in cute_dsl_custom_ops._MXFP8_GEMM_ALPHA_CACHE + + @pytest.mark.skipif( getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The test requires SM107 and SM107 CuTe DSL support.", diff --git a/tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py b/tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py new file mode 100644 index 000000000000..838b7d20123f --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_kimi_k3_fp8_weight_read_linear.py @@ -0,0 +1,247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Coverage for Kimi K3's CuTe-only FP8 weight-read linear. + +This module is hand-built rather than a ``Linear`` + ``FP8BlockScalesLinearMethod``, +so the shipping FP8-block-scale tests never reach it. These tests pin the +single-scale loader contract and the specialized quant + CuTe path. +""" + +import pytest +import torch +from _torch.helpers import calc_diff, per_block_cast_to_fp8 +from utils.util import getSMVersion + +from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + _get_kimi_k3_mxfp8_tuning_buckets, + _kimi_k3_mxfp8_tuning_bucket, +) +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.models.modeling_kimi_linear import ( + _Fp8BlockScaleWeightReadLinear as K3Fp8Linear, +) + +# K3-representative (out, in) projections, all 128-aligned. +SHAPES = [(512, 1024), (2048, 1024)] +MS = [1, 32, 128] + +# The Rubin MXFP8 weight-read path needs both SM107 and the internal CuTe DSL +# build; keep the predicate local so the gate does not depend on a helper that +# lives outside this change. +RUBIN = pytest.mark.skipif( + getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="needs SM107 with Rubin CuTe DSL support", +) + + +def _ref(x, w): + return (x.float() @ w.float().t()).to(x.dtype) + + +def _check(out, expected, tag): + assert out.dtype == expected.dtype, tag + assert out.shape == expected.shape, tag + assert torch.isfinite(out).all(), f"{tag}: non-finite output" + diff = calc_diff(out, expected) + assert diff < 5e-3, f"{tag}: calc_diff={diff}" + + +def _make(out_features, in_features, seed=0): + torch.random.manual_seed(seed) + w = ( + torch.randn((out_features, in_features), device="cuda", dtype=torch.bfloat16) + / in_features**0.5 + ) + weight, weight_scale = K3Fp8Linear.quantize_weight(w) + return w, K3Fp8Linear(weight, weight_scale, out_features) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +@pytest.mark.parametrize("out_features, in_features", SHAPES) +@pytest.mark.parametrize("m", MS) +def test_forward_matches_bf16(m, out_features, in_features): + """The CuTe-only path must preserve the FP8 block-scale numerics.""" + w, mod = _make(out_features, in_features) + x = torch.randn((m, in_features), device="cuda", dtype=torch.bfloat16) + out = mod(x) + _check(out, _ref(x, w), f"m={m} n={out_features} k={in_features}") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_forward_prequantized_matches_bf16(): + """D-Spark's fused KDA output must use the same CuTe scale ABI.""" + w, mod = _make(512, 1024, seed=9) + x = torch.randn((16, 1024), device="cuda", dtype=torch.bfloat16) + activation, activation_scale = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0(x) + + out = mod.forward_prequantized(activation, activation_scale) + + _check(out, _ref(x, w), "prequantized") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_forward_uses_specialized_quant(): + """Every call must use the specialized quantizer and fine-M runner.""" + _, mod = _make(512, 1024, seed=8) + x = torch.randn((64, 1024), device="cuda", dtype=torch.bfloat16) + mod(x) + + quant_calls = [] + fine_grained_m = [] + real_quant = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0 + real_gemm = torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin + + class _QuantSpy: + def __call__(self, *args, **kwargs): + quant_calls.append(args[0].shape[0]) + return real_quant(*args, **kwargs) + + class _GemmSpy: + def __call__(self, *args, **kwargs): + fine_grained_m.append(kwargs["fine_grained_m"]) + return real_gemm(*args, **kwargs) + + torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0 = _QuantSpy() + torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin = _GemmSpy() + try: + mod(x) + mod(x) + finally: + torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0 = real_quant + torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin = real_gemm + + assert quant_calls == [64, 64] + assert fine_grained_m == [True, True] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_single_cute_scale_is_built_on_rubin(): + """The CuTe layout must be prepared at load time, not during capture.""" + _, mod = _make(2048, 1024, seed=5) + assert mod.weight_scale.numel() > 0 + assert mod.weight_scale.dtype is torch.uint8 + assert not hasattr(mod, "weight_scale_mx") + assert not hasattr(mod, "gemm_alpha") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@RUBIN +def test_placeholder_load_carries_cute_scale(): + """A loader-filled placeholder must receive the CuTe scale layout.""" + in_features, parts = 1024, [512, 512] + torch.random.manual_seed(6) + pairs = [] + for p in parts: + wp = torch.randn((p, in_features), device="cuda", dtype=torch.bfloat16) / in_features**0.5 + q, s = per_block_cast_to_fp8(wp) + pairs.append((q, s.float())) + + mod = K3Fp8Linear.empty_placeholder(sum(parts), in_features) + mod.load_checkpoint_pair(pairs) + assert not mod.is_placeholder + assert mod.weight_scale.numel() > 0 + assert mod.weight_scale.dtype is torch.uint8 + assert not hasattr(mod, "gemm_alpha") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_weight_preparation_returns_only_cute_pair(): + """Both construction routes return exactly FP8 weight + CuTe scale.""" + torch.random.manual_seed(7) + w = torch.randn((512, 1024), device="cuda", dtype=torch.bfloat16) / 32 + assert len(K3Fp8Linear.quantize_weight(w)) == 2 + + q, s = per_block_cast_to_fp8(w) + assert len(K3Fp8Linear.prepare_checkpoint_scale(q, s.float())) == 2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_unfilled_placeholder_raises(): + mod = K3Fp8Linear.empty_placeholder(256, 256) + x = torch.randn((2, 256), device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="never filled"): + mod(x) + + +def test_kimi_k3_fine_m_tuning_uses_hybrid_buckets(): + low_m = (1, 2, 4, 8, *range(16, 193, 16)) + assert _get_kimi_k3_mxfp8_tuning_buckets(17) == low_m[:5] + assert _get_kimi_k3_mxfp8_tuning_buckets(65) == (*low_m[:8], 80) + assert _get_kimi_k3_mxfp8_tuning_buckets(192) == low_m + assert _get_kimi_k3_mxfp8_tuning_buckets(193) == (*low_m, 256) + assert _get_kimi_k3_mxfp8_tuning_buckets(4096) == ( + *low_m, + 256, + 512, + 1024, + 2048, + 4096, + ) + assert tuple( + _kimi_k3_mxfp8_tuning_bucket(m) + for m in ( + 1, + 3, + 5, + 8, + 9, + 15, + 16, + 17, + 31, + 32, + 64, + 65, + 79, + 80, + 128, + 129, + 191, + 192, + 193, + 255, + 256, + 511, + 512, + ) + ) == ( + 1, + 2, + 4, + 8, + 8, + 8, + 16, + 16, + 16, + 32, + 64, + 80, + 80, + 80, + 128, + 144, + 192, + 192, + 256, + 256, + 256, + 512, + 512, + ) From 860350e9f46f5b2cdfdaad49b9b8c5ec949baaaf Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 15 Sep 2026 05:02:00 +0000 Subject: [PATCH 02/21] [None][feat] restore the Rubin fused-FC12 NVFP4 MoE op registration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The rebase of the Rubin branch onto main dropped the fused FC12 op from `cute_dsl_custom_ops.py`: it is present on the pre-rebase branch but on neither `5d34a82bf1` nor main. Restore it so the CuteDSL MoE backend has an op to drive. Restored into `cute_dsl_custom_ops.py`: - `SITU_BETA_DISABLED` / `_canonicalize_situ_beta` — the torch.library schema needs a concrete float, so a non-positive sentinel stands in for "beta not provided". - `Fc12FusedInputsHelper` — autotuning helper; inputs 0..9 keep the FC1 gather layout so the parent `moe_sort` regeneration is reused verbatim. - `Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner` — the tunable runner. It is a dependency of the op and is also imported directly by the MoE backend's tactic checker. - `_run_nvfp4_fc12_fused_rubin` and the `trtllm::cute_dsl_nvfp4_fc12_fused_rubin` custom op + fake impl. The kernel class itself, `rubin_contiguous_grouped_blockscaled_gemm_fused_fc12` (and its `manual_mma_128dp` helper), lands with the MoE backend change, so the import is wrapped in `try`/`except ImportError` and the registration is skipped when the kernel module is absent. That keeps this module importable on an internal CuTe DSL build regardless of which change merges first. `_compute_grid` in `cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py` is deliberately left as-is; see the PR discussion for why the Rubin `raster_order` signature change is not carried. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 856 ++++++++++++++++++ 1 file changed, 856 insertions(+) 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 5273a9fdd83c..062b4d28d903 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -42,6 +42,16 @@ # Torch schema parsing rejects ``inf`` as a default value. SWIGLU_LIMIT_SCALAR_DISABLED = -1.0 +# The torch.library schema needs a concrete float, so "unset" is a sentinel +# rather than ``None``. SiTU betas are required to be positive, so any +# non-positive value is unambiguously "not provided". +SITU_BETA_DISABLED = -1.0 + + +def _canonicalize_situ_beta(situ_beta: float) -> Optional[float]: + return None if situ_beta <= 0 else float(situ_beta) + + _KIMI_K3_MXFP8_TUNING_BUCKETS = (1, 2, 4, 8, *range(16, 193, 16)) # Both dense MXFP8 wrappers pass a scalar one to the CuTe kernel. Keep one @@ -14893,6 +14903,852 @@ def _( ) -> None: return None + # ---------------------------------------------------------------- + # Rubin NVFP4 Fused FC12 (FC1 gather+gated-act + FC2 finalize in ONE kernel) + # ---------------------------------------------------------------- + # Compat shim: the delivered fused kernel accesses ``cutlass.memory.*`` + # (SmemAllocator / TmemAllocator / get_smem_capacity_in_bytes) and + # ``cutlass.tensor_utils.LayoutEnum`` as attributes of the top-level + # ``cutlass`` module. In the pinned nvidia-cutlass-dsl-internal + # (0.3.0+...c907734) these are real submodules (e.g. the dspark kernel + # does ``from cutlass import memory``) that are simply not auto-exposed + # until explicitly imported. Import them so the attribute access + # resolves; fall back to aliasing ``cutlass.utils`` (which carries the + # same symbols in this build) if a submodule is genuinely absent. This + # keeps the vendored kernel byte-identical instead of editing its API + # references. + import cutlass.utils as _cutlass_utils_compat + try: + import cutlass.memory # noqa: F401 real submodule in pinned build + except ImportError: + if not hasattr(cutlass, "memory"): + cutlass.memory = _cutlass_utils_compat + try: + import cutlass.tensor_utils # noqa: F401 + except ImportError: + if not hasattr(cutlass, "tensor_utils"): + cutlass.tensor_utils = _cutlass_utils_compat + + # The fused FC12 kernel module ships with the CuteDSL MoE backend + # change. Keep the op registration optional so that importing this + # module still succeeds on an internal CuTe DSL build that does not + # yet carry the kernel; the MoE backend checks for the op before use. + try: + from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_grouped_blockscaled_gemm_fused_fc12 import \ + Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel + except ImportError: + Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel = None + logger.debug( + "Fused FC12 CuteDSL kernel is unavailable; " + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin is not registered.") + + if Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel is not None: + + class Fc12FusedInputsHelper(GatherGroupedGemmInputsHelper): + """Autotuning helper for the fused FC12 op. + + Inputs 0..9 keep the FC1 gather layout so the parent moe_sort + regeneration is reused verbatim; inputs 10..14 carry the FC2 + tensors, of which fc2_c and fc2_routing_scales depend on the + token count and must be resized to match the regenerated tiles + during profiling. + """ + IDX_FC2_B = 10 + IDX_FC2_SFB = 11 + IDX_FC2_ALPHA = 12 + IDX_FC2_C = 13 + IDX_FC2_ROUTING = 14 + # expanded_idx_to_permuted_idx (dim0 = num_tokens); consumed only by + # the in-op output memset. Appended after the FC2 tensors so the + # existing 0..14 positions (and their constraints) stay put. + IDX_EXPANDED_IDX = 15 + + def inputs_pre_hook( + self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: + # Reuse the parent's moe_sort regeneration on the FC1 prefix + # (feed two None placeholders so it skips the + # locality-domain output resize). + base = list(super().inputs_pre_hook( + list(inputs[:10]) + [None, None])) + fc1_prefix = base[:10] + permuted = fc1_prefix[self.IDX_PERMUTED_IDX_TO_EXPANDED_IDX] + num_tokens = self.infer_num_tokens(permuted.size(0)) + fc2_b = inputs[self.IDX_FC2_B] + fc2_sfb = inputs[self.IDX_FC2_SFB] + fc2_alpha = inputs[self.IDX_FC2_ALPHA] + fc2_c = inputs[self.IDX_FC2_C] + fc2_routing = inputs[self.IDX_FC2_ROUTING] + new_fc2_c = fc2_c.new_empty((num_tokens, fc2_c.size(1))) + new_routing = fc2_routing.new_empty( + (num_tokens, fc2_routing.size(1))) + # Resize expanded_idx to the regenerated token count so the + # unpacked forward sees a consistent 16-tensor list. Its + # contents are unused during tuning (the in-op memset falls back + # to an index-free full zero while AutoTuner.is_tuning_mode, so + # the uninitialised index values are never dereferenced). + expanded_idx = inputs[self.IDX_EXPANDED_IDX] + new_expanded = expanded_idx.new_empty( + (num_tokens, expanded_idx.size(1))) + return (*fc1_prefix, fc2_b, fc2_sfb, fc2_alpha, new_fc2_c, + new_routing, new_expanded) + + class Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner( + TunableRunner): + """Rubin runner for the fused FC1+FC2 (FC12) NVFP4 MoE kernel. + + The fused kernel replaces the two-op FC1 (gather+GEMM+gated-act+quant) + and FC2 (GEMM+finalize) sequence with a single persistent kernel. + The gated epilogue is SwiGLU or SiTU (trace-time specialization). + The only interface delta versus the existing CuteDSL backend is the + three int32 atomic counters (fc1_ready / fc1_scheduler_counter / + fc2_scheduler_counter), which are allocated and memset to zero here + on every launch. v1 exposes the GEN-phase geometry only: + mma_tiler 128x{128,256}, cluster (1,1), scheduler="l2_atomic". + 2CTA/CTX geometries are added once CTX perf tuning lands. + """ + kernel_class = Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel + kernel_cache = dict() + tuning_config_cache = dict() + + def __init__( + self, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + swiglu_limit: float = float("inf"), + ep_size: int = 1, + enable_alltoall: bool = False, + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None): + super().__init__() + self.num_experts = num_experts + self.top_k = top_k + self.num_local_experts = num_local_experts + self.local_expert_offset = local_expert_offset + self.tile_size = tile_size + self.scaling_vector_size = scaling_vector_size + self.swiglu_limit = swiglu_limit + self.activation_type = ActivationType(int(activation_type)) + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta + # Used only by the in-op output memset (moved here so the memset + # is the fused kernel's immediate stream predecessor). + self.ep_size = ep_size + self.enable_alltoall = enable_alltoall + if (sm_version := get_sm_version()) != 107: + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports SM 107 " + f"(Rubin) only, but got SM {sm_version}") + # v1 fused kernel only supports the 128-wide routing tile. + if self.tile_size not in (128, 256): + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports " + f"tile_size 128 (1-CTA) or 256 (2-CTA) only, but got " + f"{self.tile_size}") + + def unique_id(self): + return ( + self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size, + self.scaling_vector_size, + self.swiglu_limit, + int(self.activation_type), + self.situ_beta, + self.situ_linear_beta, + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple]: + (fc1_a, fc1_b, fc1_sfa, fc1_sfb, fc1_alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + fc1_norm_const, fc2_b, *_) = inputs + m = permuted_idx_to_expanded_idx.size(0) + k = fc1_a.size(1) * 2 + l, fc1_n = fc1_b.size(0), fc1_b.size(1) # noqa: E741 + fc2_n, fc2_k = fc2_b.size(1), fc2_b.size(2) * 2 + + # Fixed K for FP4: mma_tiler_k=256, mma_inst_k=128. + mma_tiler_k = 256 + mma_inst_k = 128 + # Mirror the CuteDSL grouped-GEMM runners: the MMA M-tile equals + # the routing tile_size and the cluster M = tile_size // 128, so + # tile_size=128 -> 1-CTA cluster (1,1), tile_size=256 -> 2-CTA + # cluster (2,1). moe_sort tiles the tokens by tile_size, so the + # kernel's M-tile always matches the routing tile. mma_n is free + # {128, 256}. + mma_tiler_m = self.tile_size + mma_inst_m = self.tile_size + cluster_shape_mn = (self.tile_size // 128, 1) + mma_n_candidates = [128, 256] + + valid_tactics = [] + for mma_n in mma_n_candidates: + # No "cluster CTAs > tiles" guard here (unlike the Sm100 + # runners): this kernel's per-CTA M-tile is 128, so one + # 256-row logical tile is exactly one (2,1) cluster and a + # single-tile problem is valid for 2-CTA. + # The fused N-tile must divide both FC1 and FC2 output N. + if fc1_n % mma_n != 0 or fc2_n % mma_n != 0: + continue + + mma_tiler = (mma_tiler_m, mma_n, mma_tiler_k) + mma_inst_shape = (mma_inst_m, mma_n, mma_inst_k) + + if self.__class__.kernel_class.can_implement( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + sf_dtype=cutlass.Float8E4M3FN, + sf_vec_size=self.scaling_vector_size, + fc1_c_dtype=cutlass.Float4E2M1FN, + fc2_c_dtype=cutlass.BFloat16, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + fc1_gemm_shape=(m, fc1_n, k, l), + fc2_gemm_shape=(m, fc2_n, fc2_k, l), + a_major="k", + b_major="k", + fc1_c_major="n", + fc2_c_major="n", + ): + valid_tactics.append( + (mma_tiler, mma_inst_shape, cluster_shape_mn)) + + logger.debug( + f"CuteDSL Rubin FusedFC12: Found {len(valid_tactics)} valid " + f"tactics for M={m}, FC1_N={fc1_n}, FC2_N={fc2_n}, K={k}, " + f"L={l}") + return valid_tactics + + def get_tuning_config(self) -> TuningConfig: + key = self.unique_id() + if key not in self.__class__.tuning_config_cache: + helper = Fc12FusedInputsHelper(self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size) + self.__class__.tuning_config_cache[key] = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + GatherGroupedGemmInputsHelper.IDX_SHAPE_INFER, + 0, helper.gen_tuning_buckets, + helper.map_to_tuning_buckets), ), + constraint_specs=( + ConstraintSpec(0, 0, + helper.infer_shape_num_tokens), + ConstraintSpec(2, 0, + helper.infer_shape_num_tokens), + ConstraintSpec( + 5, 0, helper.infer_shape_max_num_tiles), + ConstraintSpec( + 6, 0, helper.infer_shape_max_num_tiles), + # fc2_c/fc2_routing_scales dim0 = num_tokens (follows M bucket like inputs 0/2) + ConstraintSpec(13, 0, + helper.infer_shape_num_tokens), + ConstraintSpec(14, 0, + helper.infer_shape_num_tokens), + # expanded_idx_to_permuted_idx dim0 = num_tokens too + ConstraintSpec(15, 0, + helper.infer_shape_num_tokens), + ), + inputs_pre_hook=helper.inputs_pre_hook, + ) + return self.__class__.tuning_config_cache[key] + + def forward(self, inputs: List[torch.Tensor], + tactic: Optional[tuple], **kwargs) -> torch.Tensor: + (fc1_a, fc1_b, fc1_sfa, fc1_sfb, fc1_alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + fc1_norm_const, fc2_b, fc2_sfb, fc2_alpha, fc2_c, + fc2_routing_scales, expanded_idx_to_permuted_idx) = inputs + + assert fc1_a.dtype == torch.float4_e2m1fn_x2 + assert fc1_b.dtype == torch.float4_e2m1fn_x2 + assert fc2_b.dtype == torch.float4_e2m1fn_x2 + assert fc1_sfa.dtype == torch.uint8 + assert fc1_sfb.dtype == torch.uint8 + assert fc2_sfb.dtype == torch.uint8 + assert fc1_alpha.dtype == torch.float32 + assert fc2_alpha.dtype == torch.float32 + assert fc1_norm_const.dtype == torch.float32 + assert fc2_c.dtype == torch.bfloat16 + + sf_vec = self.scaling_vector_size + orig_m, k = fc1_a.size(0), fc1_a.size(1) * 2 + m = permuted_idx_to_expanded_idx.size(0) + l, fc1_n = fc1_b.size(0), fc1_b.size(1) # noqa: E741 + interm_size = fc1_n // 2 # SwiGLU (gated) halves N + k // sf_vec + fc2_n, fc2_k = fc2_b.size(1), fc2_b.size(2) * 2 + fc2_k // sf_vec + num_tokens = fc2_c.size(0) + num_tiles = m // self.tile_size + assert m % self.tile_size == 0 + assert fc2_k == interm_size, ( + f"FC2 K ({fc2_k}) must equal FC1 intermediate ({interm_size})" + ) + + if isinstance(tactic, tuple): + mma_tiler, mma_inst_shape, cluster_shape_mn = tactic + else: + # Fallback geometry must satisfy the kernel validator + # (is_valid_mma_tiler_and_cluster_shape): inst_m == + # tile_m and cluster_m == inst_m // 128. Mirrors + # _get_sm107_nvfp4_default_mma_config used by the sibling + # runners; the old fixed cluster (1,1) was illegal for + # tile_size=256 (a 2-CTA MMA in a 1-CTA cluster). + mma_inst_m = min(self.tile_size, 256) + mma_tiler = (self.tile_size, 128, 256) + mma_inst_shape = (mma_inst_m, 128, 128) + cluster_shape_mn = (mma_inst_m // 128, 1) + # The MMA M-tile and cluster M must match the routing tile this + # runner was built for; a mismatch (e.g. a tactic captured under + # another tile_size) would index routing metadata wrongly and + # can read uninitialised permuted-index padding. + assert ( + mma_tiler[0] == self.tile_size + and cluster_shape_mn[0] == self.tile_size // 128), ( + f"FC12 tactic/tile mismatch: mma_tiler={mma_tiler} " + f"mma_inst_shape={mma_inst_shape} " + f"cluster_shape_mn={cluster_shape_mn} " + f"tile_size={self.tile_size}") + + # FC1 intermediate output (kept on-chip by the kernel; passed as + # a scratch tensor) + its dynamic block scale. + fc1_c = torch.empty(m, + interm_size // 2, + dtype=fc1_a.dtype, + device=fc1_a.device) + fc1_sfc = torch.empty(m * interm_size // sf_vec, + dtype=fc1_sfa.dtype, + device=fc1_sfa.device) + # Three atomic counters: allocate + memset-zero every launch. + fc1_ready = torch.zeros(num_tiles, + dtype=torch.int32, + device=fc1_a.device) + fc1_scheduler_counter = torch.zeros(1, + dtype=torch.int32, + device=fc1_a.device) + fc2_scheduler_counter = torch.zeros(1, + dtype=torch.int32, + device=fc1_a.device) + + # Zero the scatter-add output right before the fused kernel so + # the memset becomes the kernel's immediate stream predecessor + # (PDL prologue can then overlap a longer predecessor than the + # 1-int counter). The zeroing must run on every launch: autotune() + # keeps is_tuning_mode set for its whole context, including the + # final real launch, so gating the memset on it would leave real + # outputs unzeroed. While tuning, the profiling inputs carry an + # uninitialised expanded_idx (see Fc12FusedInputsHelper + # .inputs_pre_hook), so use an index-free full zero there instead + # of the sparse, index-driven memset; it is semantically equivalent + # (a superset of the rows the finalize scatter-adds into) and still + # charges a memset to the profiled time. + if AutoTuner.get().is_tuning_mode: + fc2_c.zero_() + else: + torch.ops.trtllm.moe_output_memset_inplace( + input=fc2_c, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx= + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx= + permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=self.tile_size, + top_k=self.top_k, + ep_size=self.ep_size, + enable_alltoall=self.enable_alltoall, + ) + + fc1_a_ptr = make_ptr(cutlass.Float4E2M1FN, + fc1_a.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + fc1_b_ptr = make_ptr(cutlass.Float4E2M1FN, + fc1_b.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + fc1_c_ptr = make_ptr(cutlass.Float4E2M1FN, + fc1_c.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + fc1_sfa_ptr = make_ptr(cutlass.Float8E4M3FN, + fc1_sfa.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc1_sfb_ptr = make_ptr(cutlass.Float8E4M3FN, + fc1_sfb.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc1_sfc_ptr = make_ptr(cutlass.Float8E4M3FN, + fc1_sfc.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc1_norm_const_ptr = make_ptr(cutlass.Float32, + fc1_norm_const.data_ptr(), + cute.AddressSpace.gmem) + fc1_alpha_ptr = make_ptr(cutlass.Float32, + fc1_alpha.data_ptr(), + cute.AddressSpace.gmem) + fc2_alpha_ptr = make_ptr(cutlass.Float32, + fc2_alpha.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_group_idx_ptr = make_ptr( + cutlass.Int32, tile_idx_to_group_idx.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_mn_limit_ptr = make_ptr( + cutlass.Int32, tile_idx_to_mn_limit.data_ptr(), + cute.AddressSpace.gmem) + permuted_idx_to_expanded_idx_ptr = make_ptr( + cutlass.Int32, permuted_idx_to_expanded_idx.data_ptr(), + cute.AddressSpace.gmem) + num_non_exiting_tiles_ptr = make_ptr( + cutlass.Int32, num_non_exiting_tiles.data_ptr(), + cute.AddressSpace.gmem) + fc1_ready_ptr = make_ptr(cutlass.Int32, + fc1_ready.data_ptr(), + cute.AddressSpace.gmem) + fc1_scheduler_counter_ptr = make_ptr( + cutlass.Int32, fc1_scheduler_counter.data_ptr(), + cute.AddressSpace.gmem) + fc2_scheduler_counter_ptr = make_ptr( + cutlass.Int32, fc2_scheduler_counter.data_ptr(), + cute.AddressSpace.gmem) + fc2_b_ptr = make_ptr(cutlass.Float4E2M1FN, + fc2_b.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + fc2_sfb_ptr = make_ptr(cutlass.Float8E4M3FN, + fc2_sfb.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc2_c_ptr = make_ptr(cutlass.BFloat16, + fc2_c.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + fc2_routing_scales_ptr = make_ptr( + cutlass.Float32, fc2_routing_scales.data_ptr(), + cute.AddressSpace.gmem) + + torch_stream = torch.cuda.current_stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + max_active_clusters = get_max_activate_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + + # The fused kernel exposes only ``__call__`` (cute.Tensor args), + # no ptr-based wrapper. ``make_ordered_layout`` needs an MLIR + # context, so build the cute tensors INSIDE a @cute.jit wrapper + # (like the standalone FC1/FC2 kernels' ``wrapper`` methods do + # internally) that captures ``gemm`` and takes ptrs + Int64 + # shapes. Shapes are dynamic so one compile serves all shapes. + sf_vec_cx = self.scaling_vector_size + tile_size_cx = self.tile_size + top_k_cx = self.top_k + cache_key = (self.scaling_vector_size, self.tile_size, + self.top_k, mma_tiler, mma_inst_shape, + cluster_shape_mn, max_active_clusters, + self.swiglu_limit, int(self.activation_type), + self.situ_beta, self.situ_linear_beta) + if cache_key not in self.__class__.kernel_cache: + gemm = self.__class__.kernel_class( + self.scaling_vector_size, + mma_inst_shape, + mma_tiler, + cluster_shape_mn, + True, # vectorized_f32 + topk=self.top_k, + use_pdl=True, + swiglu_limit=self.swiglu_limit, + scheduler="l2_atomic", + activation_type=self.activation_type, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, + ) + + @cute.jit + def _fc12_wrapper( + a_ptr, + b_ptr, + c_ptr, + sfa_ptr, + sfb_ptr, + sfc_ptr, + norm_const_ptr, + tile_grp_ptr, + tile_mn_ptr, + num_non_exiting_ptr, + fc1_alpha_ptr, + fc2_alpha_ptr, + ready_ptr, + fc1_sched_ptr, + fc2_sched_ptr, + w2_ptr, + out_ptr, + w2_sf_ptr, + permuted_ptr, + routing_ptr, + orig_m: cutlass.Int64, + m: cutlass.Int64, + fc1_n: cutlass.Int64, + k: cutlass.Int64, + l: cutlass.Int64, # noqa: E741 + fc2_n: cutlass.Int64, + fc2_k: cutlass.Int64, + num_tokens: cutlass.Int64, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + interm_size = fc1_n // 2 # SwiGLU (gated) halves N + scale_k = k // sf_vec_cx + fc2_scale_k = fc2_k // sf_vec_cx + num_tiles = m // tile_size_cx + a = cute.make_tensor( + a_ptr, + layout=cute.make_ordered_layout( + (orig_m, k, 1), order=(1, 0, 2))) + b = cute.make_tensor( + b_ptr, + layout=cute.make_ordered_layout( + (fc1_n, k, l), order=(1, 0, 2))) + c = cute.make_tensor(c_ptr, + layout=cute.make_layout( + (m, interm_size, 1), + stride=(interm_size, 1, + m * interm_size))) + sfa = cute.make_tensor( + sfa_ptr, + layout=cute.make_ordered_layout( + (orig_m, scale_k, 1), order=(1, 0, 2))) + sfb = cute.make_tensor( + sfb_ptr, + layout=cute.make_ordered_layout( + (32, 4, fc1_n // 128, 4, scale_k // 4, l), + order=(2, 1, 4, 0, 3, 5))) + sfc = cute.make_tensor( + sfc_ptr, + layout=cute.make_ordered_layout( + (32, 4, m // 128, 4, interm_size // + (sf_vec_cx * 4), l), + order=(2, 1, 4, 0, 3, 5))) + norm_const = cute.make_tensor( + norm_const_ptr, layout=cute.make_layout((1, ))) + fc1_alpha = cute.make_tensor( + fc1_alpha_ptr, layout=cute.make_layout((l, ))) + fc2_alpha = cute.make_tensor( + fc2_alpha_ptr, layout=cute.make_layout((l, ))) + tile_grp = cute.make_tensor(tile_grp_ptr, + layout=cute.make_layout( + (num_tiles, ))) + tile_mn = cute.make_tensor(tile_mn_ptr, + layout=cute.make_layout( + (num_tiles, ))) + num_non_exiting = cute.make_tensor( + num_non_exiting_ptr, + layout=cute.make_layout((1, ))) + ready = cute.make_tensor(ready_ptr, + layout=cute.make_layout( + (num_tiles, ))) + fc1_sched = cute.make_tensor( + fc1_sched_ptr, layout=cute.make_layout((1, ))) + fc2_sched = cute.make_tensor( + fc2_sched_ptr, layout=cute.make_layout((1, ))) + w2 = cute.make_tensor( + w2_ptr, + layout=cute.make_ordered_layout( + (fc2_n, fc2_k, l), order=(1, 0, 2))) + w2_sf = cute.make_tensor( + w2_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, fc2_n // 128, 4, fc2_scale_k // 4, + l), + order=(2, 1, 4, 0, 3, 5))) + out = cute.make_tensor( + out_ptr, + layout=cute.make_layout( + (num_tokens, fc2_n, 1), + stride=(fc2_n, 1, num_tokens * fc2_n))) + permuted = cute.make_tensor(permuted_ptr, + layout=cute.make_layout( + (m, ))) + routing = cute.make_tensor( + routing_ptr, + layout=cute.make_ordered_layout( + (num_tokens, top_k_cx), order=(1, 0))) + gemm(a, + b, + c, + sfa, + sfb, + sfc, + norm_const, + tile_grp, + tile_mn, + num_non_exiting, + fc1_alpha, + fc2_alpha, + ready, + fc1_sched, + fc2_sched, + w2, + out, + w2_sf, + permuted, + routing, + max_active_clusters=max_active_clusters, + stream=stream) + + compiled_gemm = cute.compile( + _fc12_wrapper, + fc1_a_ptr, + fc1_b_ptr, + fc1_c_ptr, + fc1_sfa_ptr, + fc1_sfb_ptr, + fc1_sfc_ptr, + fc1_norm_const_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + num_non_exiting_tiles_ptr, + fc1_alpha_ptr, + fc2_alpha_ptr, + fc1_ready_ptr, + fc1_scheduler_counter_ptr, + fc2_scheduler_counter_ptr, + fc2_b_ptr, + fc2_c_ptr, + fc2_sfb_ptr, + permuted_idx_to_expanded_idx_ptr, + fc2_routing_scales_ptr, + orig_m, + m, + fc1_n, + k, + l, + fc2_n, + fc2_k, + num_tokens, + max_active_clusters=max_active_clusters, + stream=stream, + ) + self.__class__.kernel_cache[cache_key] = compiled_gemm + else: + compiled_gemm = self.__class__.kernel_cache[cache_key] + + compiled_gemm( + fc1_a_ptr, + fc1_b_ptr, + fc1_c_ptr, + fc1_sfa_ptr, + fc1_sfb_ptr, + fc1_sfc_ptr, + fc1_norm_const_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + num_non_exiting_tiles_ptr, + fc1_alpha_ptr, + fc2_alpha_ptr, + fc1_ready_ptr, + fc1_scheduler_counter_ptr, + fc2_scheduler_counter_ptr, + fc2_b_ptr, + fc2_c_ptr, + fc2_sfb_ptr, + permuted_idx_to_expanded_idx_ptr, + fc2_routing_scales_ptr, + orig_m, + m, + fc1_n, + k, + l, + fc2_n, + fc2_k, + num_tokens, + stream=stream, + ) + return fc2_c + + def _run_nvfp4_fc12_fused_rubin( + input: torch.Tensor, + fc1_weight: torch.Tensor, + input_scale: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int, + swiglu_limit: float, + precomputed_tactic: Optional[str], + expanded_idx_to_permuted_idx: torch.Tensor, + ep_size: int, + enable_alltoall: bool, + tuner_key: str, + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, + ) -> torch.Tensor: + tuner = AutoTuner.get() + runner = Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + scaling_vector_size, + swiglu_limit=swiglu_limit, + ep_size=ep_size, + enable_alltoall=enable_alltoall, + activation_type=ActivationType(activation_type), + situ_beta=_canonicalize_situ_beta(situ_beta), + situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta), + ) + # Input order matches Fc12FusedInputsHelper (FC1 prefix 0..9 mirrors + # GatherGroupedGemmInputsHelper; FC2 tensors 10..14; expanded_idx 15 + # feeds the in-op output memset). + inputs = [ + input, fc1_weight, input_scale, fc1_weight_scale, fc1_alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + global_sf, fc2_weight, fc2_weight_scale, fc2_alpha, output, + token_final_scales, expanded_idx_to_permuted_idx + ] + if precomputed_tactic is None: + _, best_tactic = tuner.choose_one( + tuner_key, + [runner], + runner.get_tuning_config(), + inputs, + ) + else: + best_tactic = ast.literal_eval(precomputed_tactic) + return runner(inputs, tactic=best_tactic) + + @torch.library.custom_op( + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin", + mutates_args=("output", ), + schema= + "(Tensor input, Tensor fc1_weight, Tensor input_scale, Tensor fc1_weight_scale, " + "Tensor fc1_alpha, Tensor tile_idx_to_group_idx, Tensor tile_idx_to_mn_limit, " + "Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, Tensor global_sf, " + "Tensor fc2_weight, Tensor fc2_weight_scale, Tensor fc2_alpha, " + "Tensor(a13!) output, Tensor token_final_scales, " + "Tensor expanded_idx_to_permuted_idx, " + "SymInt num_experts, SymInt top_k, SymInt num_local_experts, " + "SymInt local_expert_offset, SymInt tile_size, float swiglu_limit, " + "SymInt ep_size, bool enable_alltoall, " + "SymInt scaling_vector_size=16, " + "str? precomputed_tactic=None, " + f"SymInt activation_type={int(ActivationType.Swiglu)}, " + "float situ_beta=-1.0, float situ_linear_beta=-1.0) -> ()", + device_types="cuda") + def cute_dsl_nvfp4_fc12_fused_rubin( + input: torch.Tensor, + fc1_weight: torch.Tensor, + input_scale: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + token_final_scales: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + swiglu_limit: float, + ep_size: int, + enable_alltoall: bool, + scaling_vector_size: int = 16, + precomputed_tactic: Optional[str] = None, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, + ) -> None: + # In-place: finalize scatter-adds into ``output`` (mutates_args); + # the op returns nothing so it does not alias its own input. + _run_nvfp4_fc12_fused_rubin( + input, fc1_weight, input_scale, fc1_weight_scale, fc1_alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + global_sf, fc2_weight, fc2_weight_scale, fc2_alpha, output, + token_final_scales, num_experts, top_k, num_local_experts, + local_expert_offset, tile_size, scaling_vector_size, + swiglu_limit, precomputed_tactic, + expanded_idx_to_permuted_idx, ep_size, enable_alltoall, + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin", + ActivationType(activation_type), situ_beta, + situ_linear_beta) + + @torch.library.register_fake( + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin") + def _( + input: torch.Tensor, + fc1_weight: torch.Tensor, + input_scale: torch.Tensor, + fc1_weight_scale: torch.Tensor, + fc1_alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_scale: torch.Tensor, + fc2_alpha: torch.Tensor, + output: torch.Tensor, + token_final_scales: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + swiglu_limit: float, + ep_size: int, + enable_alltoall: bool, + scaling_vector_size: int = 16, + precomputed_tactic: Optional[str] = None, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, + ) -> None: + return None + # ---------------------------------------------------------------- # Rubin BF16/FP16 Gather + SwiGLU Fusion (FC1 layer) # ---------------------------------------------------------------- From 6f843fa93a774fe2b2bf61c598aedd52c09b0bf0 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 06:25:49 +0000 Subject: [PATCH 03/21] [None][fix] drop two stray args from the mixed-cluster grid call ``Sm107BlockScaledPersistentDenseGemmMixedClustersKernel.__call__`` passed eight positional arguments to ``_compute_mixed_cluster_grid``, which takes six after ``self``. The call is unconditional, so the first trace of this kernel raises ``TypeError: _compute_mixed_cluster_grid() takes 7 positional arguments but 9 were given``. The call site is what drifted, not the signature: main calls it with exactly these six arguments, and this branch's only change to the file was adding ``self.swizzle_size`` and ``self.raster_order == "m"``. With them removed the file is identical to main again. The callee already derives both values from ``self``: raster_along_m = self.raster_order == "m" ... ClcDynamicPersistentTileSchedulerParams(..., self.swizzle_size, raster_along_m) so widening the signature would only have shadowed the attributes with the same values. It stayed latent because the enclosing runner's ``get_valid_tactics`` returns [] unless ``get_sm_version() == 107``, so no SM100/SM103 machine reaches it, and on SM107 it further requires the autotuner to pick the ``"mixed_clusters"`` variant over ``"base"``. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py index 033a3083101a..dddbb27ee0bf 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/dense_blockscaled_gemm_persistent.py @@ -3209,8 +3209,6 @@ def __call__( self.fallback_cluster_shape_mn, preferred_max_active_clusters, fallback_max_active_clusters, - self.swizzle_size, - self.raster_order == "m", ) self.buffer_align_bytes = 1024 From f79af97e7a521d45ecd317b082dc100cb4e78be7 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 07:40:05 +0000 Subject: [PATCH 04/21] [None][fix] repair two defects in the GatedMLP SwiGLU fallback test `_make_down_proj` omitted `has_fp8_block_scales`. The test forces `_can_fuse_gate_up_swiglu()` False, which is exactly what enters the `else` branch of the activation dispatch, and that branch calls `_can_fuse_swiglu_fp8_quant()` -- whose first check is `self.activation == F.silu and self.down_proj.has_fp8_block_scales`. The activation defaults to `F.silu`, so the short circuit does not protect the stand-in and the read raises `AttributeError` before any assertion runs. The `swiglu` assertion also pinned only `swiglu_limit`. All three call sites in `GatedMLP._apply_activation` pass `swiglu_limit`, `swiglu_alpha` and `swiglu_beta`, and `assert_called_once_with` is an exact match, so the call could never satisfy it. All three are None for a plain-SwiGLU layer. Both predate this PR: the fixture and the assertion are identical on `rubin-advance`, as is the production code they exercise. Same fix as a7f3c292a0 on user/lizhiz/rubin-advance-rebase-main-0911. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tests/unittest/_torch/modules/test_gated_mlp.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/modules/test_gated_mlp.py b/tests/unittest/_torch/modules/test_gated_mlp.py index 2db3f23738b5..da0b31ee3b1f 100644 --- a/tests/unittest/_torch/modules/test_gated_mlp.py +++ b/tests/unittest/_torch/modules/test_gated_mlp.py @@ -32,6 +32,10 @@ def _make_down_proj() -> nn.Module: down_proj = nn.Module() down_proj.has_fp8_qdq = False down_proj.has_w4a8_nvfp4_fp8 = False + # The unfused branch calls ``_can_fuse_swiglu_fp8_quant``, whose first + # check reads this attribute, so a stand-in without it raises + # AttributeError before the assertions below are reached. + down_proj.has_fp8_block_scales = False down_proj.forward = Mock(side_effect=lambda value, **kwargs: value + 1) return down_proj @@ -57,7 +61,11 @@ def test_gate_up_partition_falls_back_to_swiglu( assert not mlp._can_fuse_gate_up_swiglu() mlp.gate_up_proj.forward.assert_called_once_with(inputs) - swiglu.assert_called_once_with(projected, swiglu_limit=None) + # ``GatedMLP`` forwards all three SwiGLU shape parameters on every call; + # they are None here because this layer is plain SwiGLU. + swiglu.assert_called_once_with( + projected, swiglu_limit=None, swiglu_alpha=None, swiglu_beta=None + ) mlp.down_proj.forward.assert_called_once() down_args, down_kwargs = mlp.down_proj.forward.call_args torch.testing.assert_close(down_args[0], activated) From ce8e2b76aa9da580252df99379e1c10ccf0c79f8 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 08:16:46 +0000 Subject: [PATCH 05/21] [None][fix] gate the fused NVFP4 SwiGLU epilogue on swiglu_limit `use_cute_dsl_nvfp4_swiglu_blackwell` was enabled for any plain-SwiGLU layer without checking the clamp. The fused CuteDSL epilogue applies no clamp, so a layer carrying a real `swiglu_limit` silently produced unclamped results -- wrong numerics with no error, rather than falling back to the Triton kernel. `_is_plain_swiglu` deliberately covers only alpha/beta, so it does not catch this. `rubin-advance` gates the limit in this same expression: and (swiglu_limit is None or swiglu_limit == float("inf")) That clause was dropped during the rebase; restore it. Reported by review on PR #19184. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/modules/gated_mlp.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index 464e6e59ec80..4cdeee691007 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -121,9 +121,15 @@ def __init__( allreduce_strategy=config.allreduce_strategy, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, - use_cute_dsl_nvfp4_swiglu_blackwell=(use_cute_dsl_blockscaling_mm - and activation == F.silu - and not bias), + # The fused CuteDSL NVFP4 SwiGLU epilogue applies no clamp, so a + # layer carrying a real ``swiglu_limit`` must stay on the Triton + # kernel. ``_is_plain_swiglu`` deliberately covers only alpha/beta; + # the limit is gated here (as on ``rubin-advance``). Without this + # the clamp is silently dropped -- wrong numerics, no error. + use_cute_dsl_nvfp4_swiglu_blackwell=( + use_cute_dsl_blockscaling_mm + and activation == F.silu and not bias + and (swiglu_limit is None or swiglu_limit == float("inf"))), use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, disable_deep_gemm=disable_deep_gemm, fused_weight_shard_indices_mapping=gateup_shard_indices_mapping, From f847447886e2ae2f6ea73c5a7815efa12bd45cbc Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 08:17:21 +0000 Subject: [PATCH 06/21] [None][fix] accept q_rope_applied in the mla_rope_generation fake impl The schema in `dsv3RopeOp.cpp` ends with `", bool q_rope_applied=False"`, but the `register_fake` in `cpp_custom_ops.py` stopped at `quant_scale_qkv`. Meta dispatch passes an argument the fake cannot accept, so any traced or fake-tensor path through this op fails. Pre-existing rather than introduced here: `rubin-advance` carries the same mismatch -- its schema declares the parameter while its fake ends at `quant_scale_qkv`. Reported by review on PR #19184. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 02d934b8b5dc..659e0d1b476c 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1430,6 +1430,9 @@ def _( kv_only: bool = False, kv_done_elsewhere: bool = False, quant_scale_qkv: Optional[torch.Tensor] = None, + # Declared by the schema in dsv3RopeOp.cpp; meta dispatch passes it + # positionally, so the fake has to accept it. + q_rope_applied: bool = False, ) -> None: # This is a fake implementation for shape inference # The actual operation modifies fused_q and q_pe in-place From bf6a9feffa4e03b55c7862820c327172901c9c9b Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 08:37:54 +0000 Subject: [PATCH 07/21] [None][fix] restore SM107 in the CuteDSL MLA decode reduction path `mla_decode_fp8.py` and `mla_decode_fp16.py` gate a tmem reduction body on the SM103 arch range. On `rubin-advance` the identical body is gated on sm_101/sm_103/sm_107/sm_110; the wider condition was lost when both files were taken wholesale from main while resolving an unrelated Helix conflict (dd366fef31, replayed as c130d75c42). Nothing in the five merge-back PRs restores it, so SM107 -- the architecture this merge-back exists to support -- silently falls through to the generic path. Only the sm_107 arm is restored. sm_101 and sm_110 are not Rubin and were never in main's copy of these files, so widening to them would be unrelated scope. The `{$nv-internal-release}` markers `rubin-advance` wraps around sm_107 are deliberately omitted: that convention appears in 25 files there and none on main, so it is stripped for upstream. Bodies are byte-identical either side of the condition, so this is a condition widening only. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../blackwell/attention/mla/mla_decode_fp16.py | 9 ++++++++- .../blackwell/attention/mla/mla_decode_fp8.py | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) 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 7ed708faaacd..d53dcd7c3ce3 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 @@ -2786,7 +2786,14 @@ def softmax( row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) - elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + # SM107 (Rubin) shares this reduction path with SM103. The arm was + # dropped when this file was taken wholesale from main during the + # rebase (c130d75c42); ``rubin-advance`` gates the identical body on + # sm_101/sm_103/sm_107/sm_110. Only sm_107 is restored here -- the + # others were never in main's copy and are not Rubin. + elif cutlass.const_expr( + (arch >= Arch.sm_103 and arch <= Arch.sm_103f) + or (arch >= Arch.sm_107 and arch <= Arch.sm_107f)): tmem_load_red_atom = cute.make_copy_atom( tcgen05.copy.LdRed32x32bOp(tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX), 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 adfd5cc78726..9e2477edf634 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 @@ -2759,7 +2759,14 @@ def softmax( # reduction for row_max row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) - elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + # SM107 (Rubin) shares this reduction path with SM103. The arm was + # dropped when this file was taken wholesale from main during the + # rebase (c130d75c42); ``rubin-advance`` gates the identical body on + # sm_101/sm_103/sm_107/sm_110. Only sm_107 is restored here -- the + # others were never in main's copy and are not Rubin. + elif cutlass.const_expr( + (arch >= Arch.sm_103 and arch <= Arch.sm_103f) + or (arch >= Arch.sm_107 and arch <= Arch.sm_107f)): tmem_load_red_atom = cute.make_copy_atom( tcgen05.copy.LdRed32x32bOp(tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX), From 29b1b2fc3a6f73e253aa7afb532638d38ea3fb33 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:57:23 +0000 Subject: [PATCH 08/21] [None][chore] move the Helix verify-group work out of this PR The Helix speculative verify-group support shipped here as a kernel-side contract with no producer: helix_local_slots was plumbed through mlaKernels and dsv3RopeOp, but nothing in any PR ever set the pointer, so it stayed null and the kernels silently fell back to the per-sequence helix_is_inactive_rank gate -- which is exactly the case the new code exists to fix. test_helix_postprocess.py likewise asserted on helix_kv_bounds that no production file here provides. Both halves now live together in the dedicated Helix PR, so drop them here and keep this one to kernels: - revert helixAllToAll.{cu,h}, thop/alltoallOp.cpp, distributed/ops.py and the two Helix tests to their base versions - drop helix_local_slots from mlaKernels.{cu,h} and thop/dsv3RopeOp.cpp - drop the zero_kv_mask argument from the alltoall_helix_native fake The unrelated q_rope_applied work that shares these files is untouched, including its fake-impl fix. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/helixAllToAll.cu | 134 ++++-------------- cpp/tensorrt_llm/kernels/helixAllToAll.h | 8 -- cpp/tensorrt_llm/kernels/mlaKernels.cu | 35 ++--- cpp/tensorrt_llm/kernels/mlaKernels.h | 7 - cpp/tensorrt_llm/thop/alltoallOp.cpp | 25 +--- cpp/tensorrt_llm/thop/dsv3RopeOp.cpp | 22 +-- .../_torch/custom_ops/cpp_custom_ops.py | 7 +- tensorrt_llm/_torch/distributed/ops.py | 10 +- .../test_helix_postprocess.py | 39 +---- .../attention/multi_gpu/test_mla_helix.py | 11 +- 10 files changed, 49 insertions(+), 249 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index f9eb6aee3009..304b255a11f6 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.cu +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025, 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,22 +85,10 @@ __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 @@ -125,44 +113,31 @@ __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 const field0Size = getFieldSize(fieldInfo[0]); - int bulkCopySize = field0Size; + int field0Size = getFieldSize(fieldInfo[0]); + totalSize += field0Size; - // Load field 1 (one or more float2 values). + // Load field 1 (single float2) if constexpr (ALLOW_VARIABLE_FIELD1) { - 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(); - } + g2sField(fieldInfo[1], dataIndex, shmemBase, totalSize, smemBar, laneId); + totalSize += getFieldSize(fieldInfo[1]); } else { - ldgsts<8>(reinterpret_cast(shmemBase + field0Size), + ldgsts<8>(reinterpret_cast(shmemBase + totalSize), reinterpret_cast(getPtr(fieldInfo[1], dataIndex)), laneId == 0); cp_async_commit_group(); } - // Only bytes copied by cp.async.bulk participate in the mbarrier transaction. - return bulkCopySize; + return totalSize; } // ============================================================================ @@ -181,7 +156,7 @@ __device__ __forceinline__ void s2gField( } } -template +template __device__ __forceinline__ void s2gAllFields( HelixFieldInfo const* fieldInfo, int dataIndex, uint8_t* shmemBase, int laneId) { @@ -192,25 +167,11 @@ __device__ __forceinline__ void s2gAllFields( int field0Size = getFieldSize(fieldInfo[0]); offset += field0Size; - // Store field 1 (one or more float2 values). + // Store field 1 (single float2) if constexpr (ALLOW_VARIABLE_FIELD1) { - 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(); - } + s2gField(fieldInfo[1], dataIndex, shmemBase, offset, laneId); + offset += getFieldSize(fieldInfo[1]); } else { @@ -333,7 +294,7 @@ __host__ __device__ __forceinline__ int computeProtoTransferSize(HelixFieldInfo // Main All-to-All Kernel // ============================================================================ -template +template __global__ void helixAllToAllKernel(HelixAllToAllParams params) { extern __shared__ uint8_t allWarpShmem[]; @@ -410,7 +371,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); @@ -433,37 +394,10 @@ __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)); @@ -540,7 +474,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>(); @@ -569,7 +503,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; @@ -601,7 +535,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)); } @@ -625,14 +559,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; @@ -646,7 +580,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); @@ -704,27 +638,11 @@ void launchHelixAllToAll(HelixAllToAllParams const& params, bool allowVariableFi { if (allowVariableField1) { - 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); - } + 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 ac51e012adbf..95ab959ff3fc 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.h +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.h @@ -47,14 +47,6 @@ 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 e1a37a874e1a..1d984adeb49a 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, int32_t const* helix_local_slots = nullptr, bool precomputed_cu_seqlens = false, - bool precomputed_fmha_scheduler = false, bool q_rope_applied = false) + bool const* helix_is_inactive_rank, bool precomputed_cu_seqlens = false, bool precomputed_fmha_scheduler = false, + bool q_rope_applied = false) { // Constants. using VecT = typename VecType::Type; @@ -755,18 +755,10 @@ __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 this rank - // owns the token's position. With speculative verify groups the per-token - // slot table decides (a group can straddle a page boundary onto 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) + // 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]) { - 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 token_kv_idx = kv_cache_lengths[batch_idx] - seq_len + local_token_idx; { auto kDst = reinterpret_cast(kv_cache.getKBlockPtr(batch_idx, token_kv_idx)); @@ -854,17 +846,10 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe, } } - // If helix parallelism is being used, only write to KV cache if this rank - // owns the token's position (per-token slots for speculative verify - // groups, 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) + // 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]) { - 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 token_kv_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); { @@ -1976,8 +1961,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.helix_local_slots, params.precomputed_cu_seqlens, - params.precomputed_fmha_scheduler, params.q_rope_applied); + params.helix_is_inactive_rank, 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 8bf1f1497bae..a0125d236a89 100644 --- a/cpp/tensorrt_llm/kernels/mlaKernels.h +++ b/cpp/tensorrt_llm/kernels/mlaKernels.h @@ -158,13 +158,6 @@ 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 9c076850d29f..8a775b1cf6b8 100644 --- a/cpp/tensorrt_llm/thop/alltoallOp.cpp +++ b/cpp/tensorrt_llm/thop/alltoallOp.cpp @@ -21,7 +21,6 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/thop/thUtils.h" -#include #include TRTLLM_NAMESPACE_BEGIN @@ -126,13 +125,10 @@ 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::optional zero_kv_mask) +std::tuple alltoall_helix_native( + torch::Tensor partial_o, torch::Tensor softmax_stats, torch::Tensor workspace, int64_t cp_rank, int64_t cp_size) { // Input validation @@ -228,21 +224,6 @@ std::tuple alltoall_helix_native(torch::Tensor par params.channelCount = 0; // auto-compute params.maxChannelCount = tensorrt_llm::kernels::computeHelixMaxChannelCount(cp_size); - // Optional zero-local-KV sanitization, applied by the sender in shared memory - params.zeroKvMask = nullptr; - params.zeroKvMaskDivisor = 1; - if (zero_kv_mask.has_value()) - { - auto const& mask = zero_kv_mask.value(); - CHECK_TH_CUDA(mask); - CHECK_CONTIGUOUS(mask); - CHECK_TYPE(mask, at::ScalarType::Bool); - TORCH_CHECK(mask.numel() > 0 && entry_count % mask.numel() == 0, "zero_kv_mask numel (", mask.numel(), - ") must divide the all-to-all entry count (", entry_count, ")"); - params.zeroKvMask = reinterpret_cast(mask.data_ptr()); - params.zeroKvMaskDivisor = entry_count / mask.numel(); - } - // Launch kernel auto stream = at::cuda::getCurrentCUDAStream(); tensorrt_llm::kernels::launchHelixAllToAll(params, allowVariableField1, stream); @@ -279,7 +260,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? zero_kv_mask=None) -> (Tensor, Tensor)"); + "cp_rank, int cp_size) -> (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 c10e76774030..b6ef2e13cbf5 100644 --- a/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp @@ -76,8 +76,6 @@ 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; @@ -136,7 +134,6 @@ void invokeMLARopeGenerationHelper(T const* latent_cache_ptr, T* q_pe_ptr, T* fu mla_params.host_bmm1_scale = args.host_bmm1_scale; mla_params.helix_position_offsets = args.helix_position_offsets_ptr; mla_params.helix_is_inactive_rank = args.helix_is_inactive_rank_ptr; - mla_params.helix_local_slots = args.helix_local_slots_ptr; mla_params.q_rope_applied = args.q_rope_applied; mla_params.precomputed_cu_seqlens = args.precomputed_cu_seqlens; @@ -193,9 +190,8 @@ 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 || 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)."); + TORCH_CHECK(helix_tensor_params.size() == 2, + "Expecting 2 tensors for helix_tensor_params: helix_position_offsets and helix_is_inactive_rank."); auto stream = at::cuda::getCurrentCUDAStream(latent_cache.get_device()); auto const kv_cache_quant_mode = tc::QuantMode(uint32_t(quant_mode)); @@ -225,14 +221,6 @@ 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()); @@ -342,9 +330,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, 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}; + 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}; 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/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 659e0d1b476c..f6c2a0963975 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1267,12 +1267,7 @@ def _(input_list, group, num_lists): ] @torch.library.register_fake("trtllm::alltoall_helix_native") - def _(partial_o, - softmax_stats, - workspace, - cp_rank, - cp_size, - zero_kv_mask=None): + def _(partial_o, softmax_stats, workspace, cp_rank, cp_size): # Returns outputs with same shapes as inputs return partial_o.new_empty(partial_o.shape), softmax_stats.new_empty( softmax_stats.shape) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 506da1d389b3..b3ea554da186 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -543,19 +543,14 @@ def get(mapping: Mapping) -> "HelixAllToAllNative": return HelixAllToAllNative._cache[mapping] - def alltoall_native(self, - partial_o: torch.Tensor, - softmax_stats: torch.Tensor, - zero_kv_mask: Optional[torch.Tensor] = None): + def alltoall_native(self, partial_o: torch.Tensor, + softmax_stats: torch.Tensor): """ 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. @@ -566,7 +561,6 @@ def alltoall_native(self, self.workspace_tensor, self.mapping.cp_rank, self.mapping.cp_size, - zero_kv_mask, ) return partial_o_out, softmax_stats_out 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 d9aac6362cfd..8b9b33b3afdf 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-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,8 +14,6 @@ # limitations under the License. import unittest -from types import SimpleNamespace -from unittest.mock import Mock, patch import pytest import torch @@ -483,41 +481,6 @@ 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 43529722b816..67bd00ef63fd 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-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -568,12 +568,3 @@ def test_mla_helix_distributed( comms_medium: str, ): run_helix_test(_full_test_multi_gpu, scenario, comms_medium) - - -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") -@skip_pre_blackwell -def test_mla_helix_fifo_v2_unaligned_softmax_stats(): - # Six heads over CP2 produce three local heads and a 24-byte softmax-stats - # row, exercising the float2 fallback used by Kimi K3 over CP32. - scenario = Scenario(num_heads=6, num_kv_heads=6, batch=1, ctx_len=64) - run_helix_test(_full_test_multi_gpu, scenario, "fifo_v2") From cc0120bf33355837d1c157edaa3493c7a58dbff7 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 03:18:50 +0000 Subject: [PATCH 09/21] [None][fix] rebase the new DSA TopK tests onto main's production code The two DSA tests were ported from the Rubin branch without re-basing their fixtures on main's newer indexer/metadata implementation, so all 132 of their CI cases failed on attribute errors rather than on the behaviour they pin. - test_indexer_gvr_prior: the indexer stub gains `use_gvr_emission` (the GVR emission closed loop main added, which has its own coverage and stays off here) and the metadata stub gains the caller-owned `radix_aux_indices` / `radix_aux_logits` workspaces the native CUDA Radix decode path asserts on. - test_metadata_topk_init: seed `_radix_rows_per_sequence`, which `__init__` sets and `object.__new__` skips, and align the self-sampling oracle with production's SM100/SM103 dispatch gate. - Drop test_selfsampling_warmup_ignores_retired_hinted_env: it pins a `hint_free` argument that no longer exists on `warmup_varlen` now that the hinted arm is retired and hint-free is the only path. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../sparse/dsa/test_indexer_gvr_prior.py | 13 +++++++- .../sparse/dsa/test_metadata_topk_init.py | 30 ++----------------- 2 files changed, 15 insertions(+), 28 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py index b2afc00ee238..fa83ae59562c 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py @@ -9,7 +9,7 @@ import torch from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer -from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation +from tensorrt_llm._torch.modules.top_k import _MAX_RADIX_BLOCKS_PER_ROW, TopK, TopKImplementation @pytest.mark.parametrize( @@ -100,6 +100,14 @@ def test_indexer_forward_uses_prior_only_for_temporal_gvr( kv_lens_cuda_runtime=torch.full((num_contexts + num_generations,), 16), gen_indexer_kv_lens_cuda_runtime=torch.full((num_generations,), 16), kv_lens_row_reorder=None, + # The native CUDA Radix decode path asserts the caller-owned aux + # workspaces are present before handing them to TopK. + radix_aux_indices=torch.empty( + (num_generations or 1, _MAX_RADIX_BLOCKS_PER_ROW, topk), dtype=torch.int32 + ), + radix_aux_logits=torch.empty( + (num_generations or 1, _MAX_RADIX_BLOCKS_PER_ROW, topk), dtype=torch.float32 + ), ) # Keep distinct storage so the caller test catches use of the host twin. metadata.seq_lens_cuda = metadata.seq_lens.clone() @@ -150,6 +158,9 @@ def prefill( index_topk=topk, layer_idx=7, _enable_heuristic_topk=True, + # The GVR emission closed loop has its own coverage; this test only + # pins prior ownership, so keep emission off. + use_gvr_emission=False, mtp_index_share=False, use_fp4=False, use_cute_dsl_paged_mqa_logits=False, diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py b/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py index 6393f56dff26..df1200650c45 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py @@ -41,7 +41,7 @@ def test_topk_flags_initialized_before_buffer_allocation( ) metadata.is_cuda_graph = False enabled = enable_heuristic and sm_version >= 100 - self_sampling_supported = enabled and dsl_available and sm_version in (100, 103, 107) + self_sampling_supported = enabled and dsl_available and sm_version in (100, 103) temporal_supported = enabled and dsl_available and sm_version in (100, 103) def check_flags(*, capture_graph: bool) -> None: @@ -95,6 +95,8 @@ def test_temporal_gvr_allocates_real_prior_buffers(num_local_layers: int) -> Non metadata.max_draft_tokens = 3 metadata.num_sms = 16 metadata.enable_context_mla_with_cached_kv = False + # object.__new__ skips __init__, which is where this default is set. + metadata._radix_rows_per_sequence = 1 with ( patch.object(dsa_metadata.TrtllmAttentionMetadata, "__post_init__"), patch.object(metadata, "create_buffers_for_mla_rope_append"), @@ -110,29 +112,3 @@ def test_temporal_gvr_allocates_real_prior_buffers(num_local_layers: int) -> Non assert metadata.gvr_prior_indices.is_cuda assert torch.count_nonzero(metadata.gvr_prior_indices).item() == 0 assert metadata.kv_lens_row_reorder_buffer.shape == (4,) - - -def test_selfsampling_warmup_ignores_retired_hinted_env(monkeypatch) -> None: - """Warm the hint-free launcher even when the retired hinted switch is set.""" - from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import ( - gvr_topk_decode_self_sampling_host as host, - ) - - metadata = object.__new__(dsa_metadata.DSAtrtllmAttentionMetadata) - metadata.enable_gvr_topk = True - metadata.use_self_sampling_topk = True - metadata.kv_cache_manager = Mock() - metadata.sparse_mla_topk = 512 - metadata._indexer_compress_ratio = 4 - metadata.sparse_metadata_params = SimpleNamespace(use_cute_dsl_paged_mqa_logits=True) - metadata.get_indexer_max_seq_len = Mock(return_value=8192) - monkeypatch.setenv("TRTLLM_GVR_V2_HINTED", "1") - with ( - patch.object(dsa_metadata, "IS_CUTLASS_DSL_AVAILABLE", True), - patch.object(dsa_metadata, "get_sm_version", return_value=107), - patch.object(host, "warmup_varlen") as warmup, - ): - metadata.warmup_selfsampling_topk(next_n=4, batch_sizes=[64]) - warmup.assert_called_once() - assert warmup.call_args.kwargs["hint_free"] is True - assert 256 in warmup.call_args.kwargs["num_rows_list"] From eb7f559c6a159bbc8eb3383a9754f5aed0f97a36 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 03:23:23 +0000 Subject: [PATCH 10/21] [None][feat] complete the SiTU plumbing for the gather grouped GEMM ops The SiTU epilogue landed in the Blackwell kernel and in the tests, but the two gather-grouped-GEMM op groups kept their pre-SiTU signatures, so every parametrization of test_cute_dsl_moe.py's Blackwell act-fusion test failed with a schema mismatch on situ_beta / situ_linear_beta. - Blackwell: thread situ_beta / situ_linear_beta through the Sm100 runner (ctor, unique_id, kernel cache key, kernel construction) and through the cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell op and its fake. - Rubin: the same for the Sm107 runner and for both the act_fusion_rubin and act_fusion_locality_domain_inplace_rubin ops, and route the activation check through the kernel module's validate_activation_type. - Rubin kernel: accept SiTu, validate the betas, and add the SiTu epilogue branch. The epilogue keeps the packed f32x2 path by expanding tanh via tanh(z) = 2*sigmoid(2z) - 1 instead of dropping to a scalar tanh. Every new parameter is trailing with a disabled sentinel, so existing callers are unchanged. The betas fold at trace time, hence their presence in the kernel cache key. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 169 +++++++++++------- ...her_grouped_blockscaled_gemm_act_fusion.py | 121 ++++++++++++- 2 files changed, 226 insertions(+), 64 deletions(-) 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 062b4d28d903..9234f78ab5f3 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -3578,17 +3578,24 @@ def __init__(self, scaling_vector_size: int = 16, activation_type: ActivationType = ActivationType.Swiglu, swiglu_limit_scalar: float = float("inf"), - use_expert_counts: bool = False): + use_expert_counts: bool = False, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None): """Initialize the runner. Args: activation_type: ``ActivationType`` for the fused epilogue. Only - ``Swiglu`` (gated) and ``Relu2`` (non-gated) are supported. + ``Swiglu`` (gated), ``Relu2`` (non-gated) and ``SiTu`` + (gated) are supported. swiglu_limit_scalar: Uniform clamp limit for SwiGLU. ``+inf`` disables clamp. + situ_beta: Gate-side SiTU constant; required for ``SiTu`` only. + situ_linear_beta: Linear-side SiTU constant; required for ``SiTu`` only. """ super().__init__() self.activation_type = validate_activation_type(activation_type) self.is_gated = is_gated_activation(self.activation_type) + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta self.num_experts = num_experts self.top_k = top_k self.num_local_experts = num_local_experts @@ -3626,6 +3633,8 @@ def unique_id(self): self.activation_type, self.swiglu_limit_scalar, self.use_expert_counts, + self.situ_beta, + self.situ_linear_beta, ) def get_valid_tactics( @@ -3885,10 +3894,13 @@ def forward(self, inputs: List, assert mma_tiler_mn[ 0] == self.tile_size, f"Tactic ({tactic}) is incompatible with tile size ({self.tile_size})" + # The SiTU betas are folded into the kernel at trace time, so they + # are part of the compiled-kernel identity, not just runtime args. cache_key = (self.scaling_vector_size, self.tile_size, self.top_k, mma_tiler_mn, cluster_shape_mn, raster_along_m, self.activation_type, self.swiglu_limit_scalar, - self.use_expert_counts, self.num_local_experts) + self.use_expert_counts, self.num_local_experts, + self.situ_beta, self.situ_linear_beta) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( @@ -3902,6 +3914,8 @@ def forward(self, inputs: List, swiglu_limit=self.swiglu_limit_scalar, use_expert_counts=self.use_expert_counts, num_local_experts=self.num_local_experts, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -3991,12 +4005,17 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, expert_counts: Optional[torch.Tensor] = None, expert_capacity: int = 0, + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> Tuple[torch.Tensor, torch.Tensor]: """CuteDSL-based NVFP4 gather grouped GEMM with activation fusion. - Supports ``ActivationType.Swiglu`` (gated) and ``ActivationType.Relu2`` - (non-gated) epilogues; other ``ActivationType`` values raise an - assertion in the runner. + Supports ``ActivationType.Swiglu`` (gated), ``ActivationType.Relu2`` + (non-gated) and ``ActivationType.SiTu`` (gated) epilogues; other + ``ActivationType`` values raise an assertion in the runner. + + ``situ_beta``/``situ_linear_beta`` are only meaningful for + ``ActivationType.SiTu``; values <= 0 disable them. """ tuner = AutoTuner.get() swiglu_limit_scalar = _canonicalize_swiglu_limit_scalar( @@ -4022,7 +4041,9 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( scaling_vector_size, activation_type=ActivationType(activation_type), swiglu_limit_scalar=swiglu_limit_scalar, - use_expert_counts=expert_counts is not None) + use_expert_counts=expert_counts is not None, + situ_beta=_canonicalize_situ_beta(situ_beta), + situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta)) inputs = [ input, weight, input_scale, weight_scale, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, @@ -4061,6 +4082,8 @@ def _fake_single_b( swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, expert_counts: Optional[torch.Tensor] = None, expert_capacity: int = 0, + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> Tuple[torch.Tensor, torch.Tensor]: if expert_counts is not None: helper = GroupedGemmInputsHelper(num_experts, top_k, @@ -14145,6 +14168,8 @@ def _( from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion import \ Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel + from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion import \ + validate_activation_type as validate_rubin_activation_type class Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( TunableRunner): @@ -14152,7 +14177,7 @@ class Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( SM107 counterpart to Blackwell's ``Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner``. - Supports SwiGLU and Relu2. + Supports SwiGLU, SiTU, and Relu2. Key differences from Blackwell: - Uses LDGSTS (cp.async) for A/SFA loading instead of TMA - Supports B-reuse pattern (mma_tiler_m = 2 * mma_inst_shape_m) @@ -14170,7 +14195,9 @@ def __init__( local_expert_offset: int, tile_size: int, scaling_vector_size: int = 16, - activation_type: ActivationType = ActivationType.Swiglu): + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None): super().__init__() self.num_experts = num_experts self.top_k = top_k @@ -14178,12 +14205,10 @@ def __init__( self.local_expert_offset = local_expert_offset self.tile_size = tile_size self.scaling_vector_size = scaling_vector_size - self.activation_type = ActivationType(int(activation_type)) - if self.activation_type not in (ActivationType.Swiglu, - ActivationType.Relu2): - raise ValueError( - f"Rubin NVFP4 CuteDSL FC1 does not support " - f"{self.activation_type.name}") + self.activation_type = validate_rubin_activation_type( + activation_type) + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta self.is_gated = is_gated_activation(self.activation_type) if (sm_version := get_sm_version()) != 107: @@ -14205,6 +14230,8 @@ def unique_id(self): self.tile_size, self.scaling_vector_size, int(self.activation_type), + self.situ_beta, + self.situ_linear_beta, ) def get_valid_tactics( @@ -14507,7 +14534,8 @@ def forward(self, inputs: List[torch.Tensor], self.top_k, mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m, locality_domain_half_gemm, a_path, - int(self.activation_type), max_active_clusters) + int(self.activation_type), self.situ_beta, + self.situ_linear_beta, max_active_clusters) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( sf_vec_size=self.scaling_vector_size, @@ -14520,6 +14548,8 @@ def forward(self, inputs: List[torch.Tensor], locality_domain_half_gemm=locality_domain_half_gemm, a_path=a_path, activation_type=self.activation_type, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) compiled_gemm = cute.compile( gemm.wrapper, @@ -14596,6 +14626,8 @@ def _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( scaling_vector_size: int, partition_id: int, activation_type: ActivationType, + situ_beta: float, + situ_linear_beta: float, precomputed_tactic: Optional[str], tuner_key: str, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: @@ -14622,6 +14654,8 @@ def _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( tile_size, scaling_vector_size, activation_type=activation_type, + situ_beta=_canonicalize_situ_beta(situ_beta), + situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta), ) inputs = [ input, weight, input_scale, weight_scale, alpha, @@ -14662,6 +14696,7 @@ def _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( "Tensor(a16!)? output_tensor, Tensor(a17!)? output_sf_tensor, " "SymInt scaling_vector_size=16, SymInt partition_id=-1, " f"SymInt activation_type={int(ActivationType.Swiglu)}, " + "float situ_beta=-1.0, float situ_linear_beta=-1.0, " "str? precomputed_tactic=None) -> (Tensor?, Tensor?)", device_types="cuda") def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( @@ -14685,6 +14720,8 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( scaling_vector_size: int = 16, partition_id: int = -1, activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, precomputed_tactic: Optional[str] = None, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: return _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( @@ -14694,7 +14731,8 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( num_experts, top_k, num_local_experts, local_expert_offset, tile_size, output_tensor, output_sf_tensor, scaling_vector_size, partition_id, - ActivationType(activation_type), precomputed_tactic, + ActivationType(activation_type), situ_beta, situ_linear_beta, + precomputed_tactic, "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin") @torch.library.register_fake( @@ -14720,6 +14758,8 @@ def _( scaling_vector_size: int = 16, partition_id: int = -1, activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, precomputed_tactic: Optional[str] = None, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: m = permuted_idx_to_expanded_idx.size(0) @@ -14751,30 +14791,33 @@ def _( "SymInt num_local_experts, SymInt local_expert_offset, " "SymInt tile_size, Tensor(a!) output_tensor, " "Tensor(b!) output_sf_tensor, SymInt scaling_vector_size=16, " - f"SymInt activation_type={int(ActivationType.Swiglu)}) -> ()", + f"SymInt activation_type={int(ActivationType.Swiglu)}, " + "float situ_beta=-1.0, float situ_linear_beta=-1.0) -> ()", device_types="cuda") def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin( - input: torch.Tensor, - weight_0: torch.Tensor, - weight_1: torch.Tensor, - input_scale: torch.Tensor, - weight_scale_0: torch.Tensor, - weight_scale_1: torch.Tensor, - alpha: torch.Tensor, - tile_idx_to_group_idx: torch.Tensor, - tile_idx_to_mn_limit: torch.Tensor, - permuted_idx_to_expanded_idx: torch.Tensor, - num_non_exiting_tiles: torch.Tensor, - global_sf: torch.Tensor, - num_experts: int, - top_k: int, - num_local_experts: int, - local_expert_offset: int, - tile_size: int, - output_tensor: torch.Tensor, - output_sf_tensor: torch.Tensor, - scaling_vector_size: int = 16, - activation_type: int = int(ActivationType.Swiglu), + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + input_scale: torch.Tensor, + weight_scale_0: torch.Tensor, + weight_scale_1: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: torch.Tensor, + output_sf_tensor: torch.Tensor, + scaling_vector_size: int = 16, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> None: """Tune and launch both Rubin locality domain NVFP4 MoE FC1 partitions. @@ -14816,6 +14859,8 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin( tile_size, scaling_vector_size, activation_type=ActivationType(activation_type), + situ_beta=_canonicalize_situ_beta(situ_beta), + situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta), )) inputs = [ input, @@ -14861,6 +14906,8 @@ def launch_partition( scaling_vector_size=scaling_vector_size, partition_id=partition_id, activation_type=activation_type, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, precomputed_tactic=repr(tactic), ) @@ -14879,27 +14926,29 @@ def launch_partition( "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin" ) def _( - input: torch.Tensor, - weight_0: torch.Tensor, - weight_1: torch.Tensor, - input_scale: torch.Tensor, - weight_scale_0: torch.Tensor, - weight_scale_1: torch.Tensor, - alpha: torch.Tensor, - tile_idx_to_group_idx: torch.Tensor, - tile_idx_to_mn_limit: torch.Tensor, - permuted_idx_to_expanded_idx: torch.Tensor, - num_non_exiting_tiles: torch.Tensor, - global_sf: torch.Tensor, - num_experts: int, - top_k: int, - num_local_experts: int, - local_expert_offset: int, - tile_size: int, - output_tensor: torch.Tensor, - output_sf_tensor: torch.Tensor, - scaling_vector_size: int = 16, - activation_type: int = int(ActivationType.Swiglu), + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + input_scale: torch.Tensor, + weight_scale_0: torch.Tensor, + weight_scale_1: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: torch.Tensor, + output_sf_tensor: torch.Tensor, + scaling_vector_size: int = 16, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: float = SITU_BETA_DISABLED, + situ_linear_beta: float = SITU_BETA_DISABLED, ) -> None: return None diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py index c26f544877b8..89351bf6c0d1 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py @@ -90,7 +90,11 @@ def silu_f32( return a * sigmoid_f32(a, fastmath=fastmath) -SUPPORTED_ACTIVATION_TYPES = (ActivationType.Swiglu, ActivationType.Relu2) +SUPPORTED_ACTIVATION_TYPES = ( + ActivationType.Swiglu, + ActivationType.SiTu, + ActivationType.Relu2, +) def validate_activation_type(activation_type) -> ActivationType: @@ -113,7 +117,7 @@ class S2TCopyBundle(NamedTuple): """ Rubin (SM107) persistent blockscaled contiguous grouped GEMM with token gather -and fused SwiGLU or Relu2 activation (FC1 of MoE). +and fused SwiGLU, SiTU, or Relu2 activation (FC1 of MoE). Compute: acc = alpha * (SFA * A[token_ids]) * (SFB * B) # GEMM @@ -141,7 +145,7 @@ class Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel: """Rubin (SM107) FC1: contiguous grouped blockscaled GEMM with token gather on A/SFA and activation fusion in the epilogue. - Supports both SwiGLU (gated) and Relu2 (non-gated) activations. + Supports SwiGLU and SiTU (gated), plus Relu2 (non-gated), activations. Builds on Sm107BlockScaledContiguousGroupedGemmKernel (persistent tile scheduling, warp specialization, B-reuse, tcgen05.mma block-scale, TMA @@ -159,7 +163,7 @@ class Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel: merged ab_pipeline (no relay warp needed). SFA is always loaded via CpAsync128.CG, then reorganized into SFA TMEM by transform warps via LDS + STTM (sfa_transform_pipeline). - - SwiGLU epilogue: C = up * silu(gate), where up/gate come from + - Gated epilogue (SwiGLU or SiTU): C = up * silu(gate), where up/gate come from interleaved accumulator at granularity=64 → output N is halved. - Optional NVFP4 quant: when c_dtype == Float4E2M1FN, the epilogue also generates SFC and quantizes the output. @@ -197,6 +201,8 @@ def __init__( use_pdl: bool = True, locality_domain_half_gemm: bool = False, activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ): self.a_path = a_path # locality domain half-GEMM: two partitions write their N-half into a shared @@ -205,6 +211,23 @@ def __init__( self.sf_vec_size = sf_vec_size self.topk = topk self.activation_type = validate_activation_type(activation_type) + if self.activation_type == ActivationType.SiTu: + if situ_beta is None or situ_linear_beta is None: + raise ValueError( + "ActivationType.SiTu requires both situ_beta and " + f"situ_linear_beta, got {situ_beta} and {situ_linear_beta}." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"SiTU betas must be positive, got {situ_beta} and {situ_linear_beta}." + ) + elif situ_beta is not None or situ_linear_beta is not None: + raise ValueError( + "situ_beta / situ_linear_beta require " + f"ActivationType.SiTu, got {self.activation_type.name}." + ) + self.situ_beta = None if situ_beta is None else float(situ_beta) + self.situ_linear_beta = None if situ_linear_beta is None else float(situ_linear_beta) self.is_gated = is_gated_activation(self.activation_type) if locality_domain_half_gemm and not self.is_gated: raise ValueError("Rubin locality domain half-GEMM currently supports SwiGLU only") @@ -3028,6 +3051,9 @@ def kernel( self._apply_swiglu_epilogue( acc_vec_up, acc_vec_gate, alpha_val, tCompute ) + elif cutlass.const_expr(self.activation_type == ActivationType.SiTu): + acc_vec_gate = tTR_rAcc_gate.load() + self._apply_situ_epilogue(acc_vec_up, acc_vec_gate, alpha_val, tCompute) elif cutlass.const_expr(self.activation_type == ActivationType.Relu2): self._apply_relu2_epilogue(acc_vec_up, alpha_val, tCompute) @@ -3283,6 +3309,93 @@ def _apply_swiglu_epilogue( acc_vec_gate_alpha = acc_vec_gate[i] * cutlass.Float32(alpha_val) tCompute[i] = acc_vec_up_alpha * silu_f32(acc_vec_gate_alpha, fastmath=True) + @cute.jit + def _apply_situ_epilogue( + self, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + alpha_val, + tCompute: cute.Tensor, + ): + """SiTU (Kimi K3), matching ``kimi_k3_moe/_mlp.py::SituAndMul`` + (itself byte-identical to HF ``modeling_kimi.py``):: + + g = alpha * gate, u = alpha * up + situ_gate = beta * tanh(g / beta) * sigmoid(g) + situ_up = linear_beta * tanh(u / linear_beta) + tCompute = situ_gate * situ_up + + ``up`` and ``gate`` come from the two interleaved accumulator subtiles + loaded by the caller, same as the SwiGLU epilogue. + + There is no packed tanh, so the vectorized path uses the identity + ``tanh(z) = 2 * sigmoid(2z) - 1`` (the same one ``utils.gelu_tanh_f32`` + uses) to stay on the packed f32x2 path -- calling a scalar tanh would + force the whole loop back to scalar. The reciprocals and ``2*beta`` + factors fold at trace time because both betas are ``const_expr``:: + + beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) + = 2*beta*sigmoid((2/beta)*x) - beta + """ + beta = self.situ_beta + linear_beta = self.situ_linear_beta + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + neg_log2e_pair = (-LOG2_E, -LOG2_E) + one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) + + inv_2beta = cutlass.Float32(2.0 / beta) + two_beta = cutlass.Float32(2.0 * beta) + neg_beta = cutlass.Float32(-beta) + inv_2lbeta = cutlass.Float32(2.0 / linear_beta) + two_lbeta = cutlass.Float32(2.0 * linear_beta) + neg_lbeta = cutlass.Float32(-linear_beta) + + # sigmoid(x) = rcp(1 + exp2(-x * log2e)), shared by both cores. + def _sigmoid(p0, p1): + neg = cute.arch.mul_packed_f32x2((p0, p1), neg_log2e_pair) + e = ( + cute.math.exp2(neg[0], fastmath=True), + cute.math.exp2(neg[1], fastmath=True), + ) + d = cute.arch.add_packed_f32x2(e, one_pair) + return (cute.arch.rcp_approx(d[0]), cute.arch.rcp_approx(d[1])) + + alpha_pair = (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)) + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + g = cute.arch.mul_packed_f32x2((acc_vec_gate[i], acc_vec_gate[i + 1]), alpha_pair) + u = cute.arch.mul_packed_f32x2((acc_vec_up[i], acc_vec_up[i + 1]), alpha_pair) + + sigmoid_g = _sigmoid(g[0], g[1]) + + gs = _sigmoid(*cute.arch.mul_packed_f32x2(g, (inv_2beta, inv_2beta))) + tanh_g = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(gs, (two_beta, two_beta)), (neg_beta, neg_beta) + ) + + us = _sigmoid(*cute.arch.mul_packed_f32x2(u, (inv_2lbeta, inv_2lbeta))) + tanh_u = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(us, (two_lbeta, two_lbeta)), (neg_lbeta, neg_lbeta) + ) + + situ_gate = cute.arch.mul_packed_f32x2(tanh_g, sigmoid_g) + out_pair = cute.arch.mul_packed_f32x2(situ_gate, tanh_u) + tCompute[i] = out_pair[0] + tCompute[i + 1] = out_pair[1] + else: + inv_2beta = cutlass.Float32(2.0 / beta) + two_beta = cutlass.Float32(2.0 * beta) + beta_f32 = cutlass.Float32(beta) + inv_2lbeta = cutlass.Float32(2.0 / linear_beta) + two_lbeta = cutlass.Float32(2.0 * linear_beta) + lbeta_f32 = cutlass.Float32(linear_beta) + for i in cutlass.range_constexpr(cute.size(acc_vec_up.shape)): + g = acc_vec_gate[i] * cutlass.Float32(alpha_val) + u = acc_vec_up[i] * cutlass.Float32(alpha_val) + tanh_g = two_beta * sigmoid_f32(g * inv_2beta, fastmath=True) - beta_f32 + tanh_u = two_lbeta * sigmoid_f32(u * inv_2lbeta, fastmath=True) - lbeta_f32 + tCompute[i] = (tanh_g * sigmoid_f32(g, fastmath=True)) * tanh_u + @cute.jit def _apply_relu2_epilogue( self, From 67f81c85baf8045b89df88c70ba00a76fb965a5b Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 03:29:15 +0000 Subject: [PATCH 11/21] [None][fix] restore the mHC Phase-4 store-to-load coherence fix The Rubin branch predates #18331, so rebasing its mHC work on top of main silently reverted that fix and re-broke the invariant its regression test (test_mhc_fused_all_mma_phase4_reload_is_coherent) was written to guard. Restore, on top of the branch's Phase-4 rework: - ld_global_volatile_u128, and both Phase-4 pass-2 reloads of layer_input that must use it: the read-only cache is not coherent with stores this same kernel issued, so __ldg can observe pre-store data. - the acquire __threadfence() after the done_counter spin, which publishes the Phase-2 writes of the other split-K CTAs. - cp.async.bulk.wait_group in the split-K kernel: cute::tma_store_wait is the `.read` form and retires only the SMEM reads, so the global writes could still be in flight at the Phase-3 fence. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../mhcKernels/fused_tf32_pmap_gemm.cuh | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh index 847d05294cf5..6600dca15500 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -166,6 +166,18 @@ __device__ __forceinline__ float2 fma_f32x2_vv(float2 a, float2 b, float2 c) #endif } +// The read-only cache is not coherent with stores issued earlier by the same +// kernel, so Phase 4's store-to-load path cannot use __ldg. +__device__ __forceinline__ uint4 ld_global_volatile_u128(uint4 const* addr) +{ + uint4 val; + asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr) + : "memory"); + return val; +} + __device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uint32_t b, uint32_t c, uint32_t d) { asm volatile( @@ -1264,8 +1276,9 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) } } - // Drain any in-flight residual_out TMA stores before exit. - cute::tma_store_wait<0>(); + // tma_store_wait is the `.read` form and retires only the SMEM reads, + // so the global writes could still be in flight at the Phase-3 fence. + asm volatile("cp.async.bulk.wait_group 0;" : : : "memory"); // Warp-reduce sqr across 4 col_lanes then atomicAdd to global. sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1); @@ -1325,6 +1338,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) { /* spin */ } + // Acquire the Phase-2 writes published before each increment. + __threadfence(); } __syncthreads(); } @@ -1685,7 +1700,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) #pragma unroll for (uint32_t h = h_start; h < H_VEC_END; h += H_STRIDE) { - uint4 li_raw = __ldg(reinterpret_cast(&obase[h])); + uint4 li_raw = ld_global_volatile_u128(reinterpret_cast(&obase[h])); uint4 nw_raw = __ldg(reinterpret_cast(&nbase[h])); __nv_bfloat162 const* li_pairs = reinterpret_cast<__nv_bfloat162 const*>(&li_raw); __nv_bfloat162 const* nw_pairs = reinterpret_cast<__nv_bfloat162 const*>(&nw_raw); @@ -1708,7 +1723,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) if (my_chunk < TAIL_CHUNKS) { const uint32_t h = H_VEC_END + my_chunk * BF16_VEC_LI; - uint4 li_raw = __ldg(reinterpret_cast(&obase[h])); + uint4 li_raw = ld_global_volatile_u128(reinterpret_cast(&obase[h])); uint4 nw_raw = __ldg(reinterpret_cast(&nbase[h])); __nv_bfloat162 const* li_pairs = reinterpret_cast<__nv_bfloat162 const*>(&li_raw); __nv_bfloat162 const* nw_pairs = reinterpret_cast<__nv_bfloat162 const*>(&nw_raw); From 3fed8f90e555aa60308956e0838f9a63186e2f96 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 03:30:54 +0000 Subject: [PATCH 12/21] [None][fix] keep Flux on the shared fused NVFP4 SwiGLU predicate Gating the gate/up weight interleave on the new can_use_cute_dsl_nvfp4_swiglu_blackwell() narrowed it to layers that opt in via the new flag. GatedMLP opts in; Flux, the other consumer of the fused Blackwell NVFP4 SwiGLU op, does not, so its weights stayed in [gate | up] while its dispatch still selected the kernel that reads the 64-row [up_0 | gate_0 | ...] interleave. The result was unrelated output, which test_flux2_nvfp4_swiglu_paths_match_eager caught as a cosine near zero. Opt the Flux fused gate/up Linear in, and dispatch on the same predicate the weight transform uses so the two can no longer disagree. Sharing the predicate also drops SM107 from the Flux fused path, where the interleave likewise never ran. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/models/flux/attention.py | 4 +++- tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py index 31df009e0be8..0d3b3054fd0d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py @@ -512,7 +512,9 @@ def _can_project_hidden_mlp_with_cute_dsl(self) -> bool: return False return ( - gate_up_proj.use_cute_dsl_blockscaling_mm + # Same predicate the weight transform uses, so dispatch can never + # consume a layout the loader did not produce. + gate_up_proj.can_use_cute_dsl_nvfp4_swiglu_blackwell() and gate_up_proj.has_nvfp4 and not gate_up_proj.has_bias and self._is_cute_dsl_swiglu_layout_compatible( diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py b/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py index 7bc2a4938f7b..38a9eb03b168 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py @@ -257,6 +257,10 @@ def range_size(r): tensor_parallel_mode=TensorParallelMode.COLUMN, reduce_output=False, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, + # Flux runs a plain, unclamped SwiGLU, so it opts in to the + # fused Blackwell NVFP4 epilogue and the interleaved gate/up + # weight layout that epilogue reads. + use_cute_dsl_nvfp4_swiglu_blackwell=use_cute_dsl_blockscaling_mm, override_tp_sharding={ "gate": (local_mlp_hidden_start, local_mlp_hidden_end), "up": (local_mlp_hidden_start, local_mlp_hidden_end), From f01972bb3dc0235e4efd29e69bc4dda88a4aac8f Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 03:32:40 +0000 Subject: [PATCH 13/21] [None][fix] stop the gather act-fusion test staging weights through int32 Adding the SiTU activation grew this test's grid from 144 to 216 cases, and the extra ones are all gated, which doubles the weight tensor. That tipped the ep_size=1 cases (256 local experts) into CUDA OOM -- including some pre-existing SwiGLU ones that were already close to the edge. The avoidable cost is the staging buffer: drawing into int32 and casting to bfloat16 keeps both live at once, ~100 GB for a single ep_size=1 gated case, for a tensor that ends up at ~34 GB. Draw straight into bfloat16 instead; the values are small integers, so they are exactly representable and the test's reference comparison is unchanged. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/thop/parallel/test_cute_dsl_moe.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 9ff195fbf504..d08395ba147b 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -1226,16 +1226,17 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( num_valid_permuted_tokens = total_num_padded_tokens.item() # Create input tensors (original size, not permuted) - a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.int32, device="cuda").to( - torch.bfloat16 - ) + # Draw straight into bfloat16: the int32 staging buffer for `b` is twice + # the size of the tensor it produces, and at ep_size=1 the pair is ~100 GB + # live at once, which is what makes this grid OOM. + a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda") b = torch.randint( -5, 5, (num_local_experts, interm_size * weight_n_multiplier, hidden_size), - dtype=torch.int32, + dtype=torch.bfloat16, device="cuda", - ).to(torch.bfloat16) + ) # Quantize inputs to FP4 a_global_sf = a.abs().max().float() / (448 * 6) From 2d75179da0f3a74eb1dc2ed0e85f494e22dfda17 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 04:54:44 +0000 Subject: [PATCH 14/21] [None][fix] address review: f32x2 arch bound and FC12 activation guard - Bound the three new packed-f32x2 helpers to the SM100 family, matching their existing siblings in the same header. The #else paths are already bit-identical (fma.rn.f32x2 is two independent fma.rn), so this only changes which instruction non-SM100 targets select. - Reject non-gated activations in the Rubin fused FC12 runner. It derives interm_size as fc1_n // 2, so a non-gated type would have halved the wrong dimension and returned wrong numbers instead of raising, and the op schema exposes activation_type as a free SymInt. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../kernels/mhcKernels/fused_tf32_pmap_gemm.cuh | 4 ++-- cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu | 2 +- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 9 +++++++++ 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh index 6600dca15500..b997621a8fcc 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -190,7 +190,7 @@ __device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uin // coefficient, so keep the pair packed instead of lowering it to two FFMA. __device__ __forceinline__ float2 fma_f32x2(float2 const& a, float2 const& b, float2 const& c) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100) float2 result; asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(reinterpret_cast(result)) @@ -204,7 +204,7 @@ __device__ __forceinline__ float2 fma_f32x2(float2 const& a, float2 const& b, fl __device__ __forceinline__ float2 mul_f32x2(float2 const& a, float2 const& b) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100) float2 result; asm volatile("mul.f32x2 %0, %1, %2;" : "=l"(reinterpret_cast(result)) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu index 3d114a83b303..2c3129fa2a7a 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu @@ -31,7 +31,7 @@ namespace kernels::mhc // Phase 2 always updates adjacent bf16 values with the same pre-mix scalar. __device__ __forceinline__ float2 mhcFmaF32x2(float2 const& a, float2 const& b, float2 const& c) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100) float2 result; asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(reinterpret_cast(result)) 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 9234f78ab5f3..00127e89a966 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -15082,6 +15082,15 @@ def __init__( self.scaling_vector_size = scaling_vector_size self.swiglu_limit = swiglu_limit self.activation_type = ActivationType(int(activation_type)) + # The fused kernel derives interm_size as fc1_n // 2, so a + # non-gated activation would silently halve the wrong + # dimension and produce wrong output rather than failing. + if self.activation_type not in (ActivationType.Swiglu, + ActivationType.SiTu): + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports " + f"gated activations (Swiglu, SiTu) only, but got " + f"{self.activation_type.name}") self.situ_beta = situ_beta self.situ_linear_beta = situ_linear_beta # Used only by the in-op output memset (moved here so the memset From 57b6791846d9dc8071cfd3fe97ba30658bd3d2ba Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 04:58:58 +0000 Subject: [PATCH 15/21] [None][fix] register the missing fp8_prequantized_swap_ab_gemm op Only the caller was ported. tensorrt_llm/_torch/modules/linear.py reaches torch.ops.trtllm.fp8_prequantized_swap_ab_gemm on the int32-packed-scale, is_sm_100f, DeepGemm-enabled path, but the op itself was never registered, so every DeepSeek-V3-Lite FP8 chunked-prefill run died in the executor worker with AttributeError: '_OpNamespace' 'trtllm' object has no attribute 'fp8_prequantized_swap_ab_gemm' Port Fp8PrequantizedSwapABGemmRunner plus the op and its fake. The Rubin branch names the autotuner's M-bucket generator deep_gemm_jit_warmup_buckets; map it to deep_gemm_gen_tuning_buckets, which is what the sibling fp8SwapABGemmRunner in this file already uses. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/custom_ops/torch_custom_ops.py | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 52d66018ebb6..cf7fc63ef8b3 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -2192,6 +2192,104 @@ def forward( return output +class Fp8PrequantizedSwapABGemmRunner(TunableRunner): + """Runs DeepGemm with pre-quantized FP8 activations and packed scales.""" + + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, deep_gemm_gen_tuning_buckets), ), + constraint_specs=(ConstraintSpec( + 1, 0, lambda input_shapes: input_shapes[0][0]), ), + exclude_from_cache=True, + ) + + def __init__(self, output_dtype: torch.dtype, + disable_ue8m0_cast: bool) -> None: + self.output_dtype = output_dtype + self.disable_ue8m0_cast = disable_ue8m0_cast + + def unique_id(self): + return ( + self.output_dtype, + self.disable_ue8m0_cast, + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + ) -> List[int]: + return [0] + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = -1, + ) -> torch.Tensor: + del tactic + activation, activation_scale, weight, weight_scale = inputs + scale_m_aligned = fp4_utils.pad_up(activation_scale.size(0), 4) + if activation_scale.stride() != (1, scale_m_aligned): + # Dynamic autotuning recreates constrained integer tensors with a + # contiguous layout. Restore the MN-major packed-scale stride that + # the real quantizers return and DeepGemm requires. + normalized_scale = torch.empty_strided( + activation_scale.shape, (1, scale_m_aligned), + dtype=activation_scale.dtype, + device=activation_scale.device) + normalized_scale.copy_(activation_scale) + activation_scale = normalized_scale + output = torch.empty( + (activation.size(0), weight.size(0)), + device=activation.device, + dtype=self.output_dtype, + ) + deep_gemm.fp8_gemm_nt( + (activation, activation_scale), + (weight, weight_scale), + output, + disable_ue8m0_cast=self.disable_ue8m0_cast, + ) + return output + + +@torch.library.custom_op("trtllm::fp8_prequantized_swap_ab_gemm", + mutates_args=()) +def fp8_prequantized_swap_ab_gemm( + activation: torch.Tensor, + activation_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + disable_ue8m0_cast: bool = False, +) -> torch.Tensor: + runner = Fp8PrequantizedSwapABGemmRunner(output_dtype, disable_ue8m0_cast) + _, best_tactic = AutoTuner.get().choose_one( + "trtllm::fp8_prequantized_swap_ab_gemm", + [runner], + Fp8PrequantizedSwapABGemmRunner.tuning_config, + [activation, activation_scale, weight, weight_scale], + ) + return runner( + inputs=[activation, activation_scale, weight, weight_scale], + tactic=best_tactic, + ) + + +@fp8_prequantized_swap_ab_gemm.register_fake +def _( + activation: torch.Tensor, + activation_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + disable_ue8m0_cast: bool = False, +) -> torch.Tensor: + del activation_scale, weight_scale, disable_ue8m0_cast + return activation.new_empty((activation.size(0), weight.size(0)), + dtype=output_dtype) + + @torch.library.custom_op("trtllm::fp8_swap_ab_gemm", mutates_args=()) def fp8_swap_ab_gemm( input: torch.Tensor, From 639140f985d1f2094c0dd4eb7d07fbfc04f42ef9 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:01:51 +0000 Subject: [PATCH 16/21] [None][fix] repair the C++ unit-test build and the fused FP8 SwiGLU gate - attentionWorkspaceTest.cpp referenced nvinfer1::DataType, which nothing in its include set declares, so the google-tests target failed to compile and took the whole cpp/test_unit_tests session down. getWorkspaceSizeForContext actually takes tensorrt_llm::DataType; use that. It is the only nvinfer1 reference under cpp/tests/unit_tests. - _can_fuse_swiglu_fp8_quant only checked `activation == F.silu`, on the stated assumption that parameterized variants install a custom callable. MiniMax M3 SwiGLU-OAI does not: it passes plain F.silu together with swiglu_alpha and swiglu_beta. The fused op takes only the limit, so the gain and offset were being dropped silently. Require _is_plain_swiglu(), matching the sibling _can_fuse_gate_up_swiglu. - Apply the bfloat16 draw to test_nvfp4_gather_grouped_gemm_situ_rubin too; it staged the same ~69 GB int32 buffer the Blackwell sibling did. Seed it while here, as the sibling already is. - Size the o_proj test's cute-dsl gate with _is_cute_dsl_fp8_bmm_available so it cannot enable the path where the DSL package is missing. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp | 2 +- tensorrt_llm/_torch/modules/gated_mlp.py | 8 +++++--- .../sparse/deepseek_v4/test_deepseek_v4_o_proj.py | 6 +++--- tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py | 9 ++++----- 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp b/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp index 1d0a3c307279..b0f2dc0dcd1a 100644 --- a/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp +++ b/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp @@ -103,7 +103,7 @@ size_t expectedUnfusedContextWorkspace(bool crossAttention) size_t getUnfusedContextWorkspace(tcop::AttentionOp const& op) { return op.getWorkspaceSizeForContext( - nvinfer1::DataType::kHALF, kBatchSize, kInputSequenceLength, kCrossKvLength, kPackedTokenCount); + tensorrt_llm::DataType::kHALF, kBatchSize, kInputSequenceLength, kCrossKvLength, kPackedTokenCount); } } // namespace diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index 4cdeee691007..6bb5fcbc6289 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -245,9 +245,11 @@ def _can_fuse_gate_up_swiglu_fp4out(self): def _can_fuse_swiglu_fp8_quant(self) -> bool: """Check whether down projection can consume fused SwiGLU FP8 output.""" - # Parameterized variants such as MiniMax M3 SwiGLU-OAI install a - # custom activation callable and must keep the unfused implementation. - if not (self.activation == F.silu + # silu_and_mul_fp8_quantize_1x128_packed_ue8m0 takes the limit but has + # no alpha/beta, so a parameterized SwiGLU must stay unfused. MiniMax + # M3 SwiGLU-OAI reaches here with plain F.silu plus swiglu_alpha and + # swiglu_beta, so the activation check alone does not exclude it. + if not (self.activation == F.silu and self._is_plain_swiglu() and self.down_proj.has_fp8_block_scales): return False if get_sm_version() == 107: diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py index fb29bc191cc4..7a00afc2e7fb 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py @@ -31,7 +31,6 @@ project_sparse_attn_output, ) from tensorrt_llm._torch.attention.mla import MLA, _is_cute_dsl_fp8_bmm_available -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import weight_dequant from tensorrt_llm._utils import get_sm_version @@ -188,9 +187,10 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): pretrained_config=pretrained_config, sparse_attention_config=sparse_config, quant_config=quant_config, + # Mirror the production predicate exactly, so the test does not enable + # the cute-dsl path on a box where the DSL package is absent. use_cute_dsl_blockscaling_mm=( - dtype_str == "fp8" - and (sm_version in (100, 103) or (sm_version == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE)) + dtype_str == "fp8" and _is_cute_dsl_fp8_bmm_available(sm_version) ), ) diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index d08395ba147b..80404f613f5c 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -1592,6 +1592,7 @@ def test_nvfp4_gather_grouped_gemm_situ_rubin(tile_size: int): num_experts = 256 num_local_experts = num_experts // ep_size + torch.manual_seed(20260727 + tile_size) routing_logits = torch.randn(num_tokens, num_experts, device="cuda") token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) token_selected_experts = token_selected_experts.to(torch.int32) @@ -1618,16 +1619,14 @@ def test_nvfp4_gather_grouped_gemm_situ_rubin(tile_size: int): max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) num_valid_permuted_tokens = total_num_padded_tokens.item() - a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.int32, device="cuda").to( - torch.bfloat16 - ) + a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda") b = torch.randint( -5, 5, (num_local_experts, interm_size * 2, hidden_size), - dtype=torch.int32, + dtype=torch.bfloat16, device="cuda", - ).to(torch.bfloat16) + ) a_global_sf = a.abs().max().float() / (448 * 6) b_global_sf = b.abs().amax(dim=(1, 2)).float() / (448 * 6) From 1111762c4bf6a876b4c5693f8a4f489a0a897983 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 17 Sep 2026 12:26:26 +0000 Subject: [PATCH 17/21] [None][fix] align the mla_rope_generation fake with the native schema The registered fake dropped two schema parameters: kv_cache_scale_orig_quant (between kv_scale_quant_orig and out_scale) and residual_dim (between head_size and tokens_per_block). Meta/FakeTensor dispatch binds arguments positionally from the schema, so every argument after the first gap landed in the wrong slot and the call supplied two more positionals than the fake accepted. Both parameters are long-standing on the C++ side; the Python fake simply never caught up. The signature now matches dsv3RopeOp.cpp field for field, all 48 arguments in order. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index f6c2a0963975..a2b10d9dd80f 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1399,6 +1399,7 @@ def _( host_kv_cache_pool_mapping: Optional[torch.Tensor], kv_scale_orig_quant: Optional[torch.Tensor], kv_scale_quant_orig: Optional[torch.Tensor], + kv_cache_scale_orig_quant: Optional[torch.Tensor], out_scale: Optional[torch.Tensor], block_ids_per_seq: Optional[torch.Tensor], helix_tensor_params: List[Optional[torch.Tensor]], @@ -1407,6 +1408,7 @@ def _( num_heads: int, num_kv_heads: int, head_size: int, + residual_dim: int, tokens_per_block: int, attention_window_size: int, beam_width: int, From 43b49d183c07d890e0a0f2af68d14b95dea5da36 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 18 Sep 2026 01:29:02 +0000 Subject: [PATCH 18/21] [None][fix] drop the re-introduced trtllm-gen FMHA CGA smem-reduction clamp The rebase re-applied #18518's selectKernelWithCgaSmemReductionLimit / limitCgaSmemReductionCtasKv wrapper around all three FmhaAutoTuner::selectKernel() call sites. main removed exactly that wrapper in #18945 (nvbugs/6737351) because the clamp is arch-agnostic: it pins mMaxNumCtasPerSeqKv to 16 / max(clusterDimX, 2) for every GmemReduction + Static-scheduler request, which is the trtllm-gen decode path on SM100 too. That cost B200 MLA decode most of its split-KV parallelism at low batch and regressed the disagg gen-only DeepSeek-R1 FP4 MTP3 con1 perf case. Carrying the clamp forward silently reverts that fix, so restore main's file verbatim. The constraint it encoded belongs in the trtllm-gen autotuner or behind an mSM == kSM_107 gate, not in host code that overrides tuned SM100 heuristics. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../trtllmGenKernels/fmha/fmhaKernels.h | 51 +++---------------- 1 file changed, 7 insertions(+), 44 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 5fefa7660b6a..2e9472a064fc 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -311,8 +311,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - std::tie(options, optionsFromArgs, ctaDim) - = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); // Check if the options are valid or not. checkFmhaOptions(options, optionsFromArgs); @@ -409,8 +409,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - std::tie(options, optionsFromArgs, ctaDim) - = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); checkFmhaOptions(options, optionsFromArgs); updateFmhaOptions(options, optionsFromArgs); @@ -527,8 +527,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - std::tie(options, optionsFromArgs, ctaDim) - = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); // Overwrite AutoTuner decision: SageAttention with SfsPV is known to cause regression to persistent scheduler. // Remove this overwritten once we refresh the cubin kernels that containing the related fix. @@ -663,43 +663,6 @@ class TllmGenFmhaKernel } private: - std::tuple selectKernelWithCgaSmemReductionLimit( - FmhaOptions options, FmhaOptionsFromArgs optionsFromArgs, int32_t multiProcessorCount) const - { - if (isGmemReduction(options.mMultiCtasKvMode) && options.mTileScheduler == TileScheduler::Static) - { - constexpr int kMaxCgaClusterDimX = 16; - constexpr int kMinPreSelectCgaClusterDimX = 2; - // selectKernel may promote this candidate to CgaSmemReduction and call computeNumCtas - // before returning the selected options. Keep the candidate legal for both 1-CTA and - // 2-CTA FMHA kernels; the exact selected clusterDimX is applied below. - options.mMaxNumCtasPerSeqKv = std::min(options.mMaxNumCtasPerSeqKv, - kMaxCgaClusterDimX / std::max(options.mClusterDimX, kMinPreSelectCgaClusterDimX)); - } - - int32_t ctaDim = 512; - FmhaAutoTuner autoTuner(options, optionsFromArgs, multiProcessorCount); - std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); - limitCgaSmemReductionCtasKv(options); - return {options, optionsFromArgs, ctaDim}; - } - - void limitCgaSmemReductionCtasKv(FmhaOptions& options) const - { - if (!isCgaSmemReduction(options.mMultiCtasKvMode)) - { - return; - } - constexpr int kMaxCgaClusterDimX = 16; - if (options.mClusterDimX * options.mMaxNumCtasPerSeqKv > kMaxCgaClusterDimX) - { - TLLM_LOG_WARNING( - "CGA reduction is not supported when numCtasPerSeqKv * clusterDimX > 16. Set mMaxNumCtasPerSeqKv to " - "16 / clusterDimX"); - options.mMaxNumCtasPerSeqKv = kMaxCgaClusterDimX / options.mClusterDimX; - } - } - inline uint64_t hashID(int qkvLayout, int maskType, int kernelType, int scheduler, int multiCtasKvMode, int headDimPerCtaV, int headDimQk, int headDimV, int tileSizeQ, int tileSizeKv, int numTokensPerPage, bool reuseSmemKForV, bool uses2CtaMma, int sparseAttention, bool skipsSoftmax, bool groupsHeadsQ, @@ -1188,7 +1151,7 @@ class TllmGenFmhaKernel options.mIsCustomSpecDecodingGen = !isContext && params.mMaxSeqLenQ > 1 && params.mIsSpecDecTree; options.mIsCausalSpecDecodingGen = !isContext && params.mMaxSeqLenQ > 1 && !params.mIsSpecDecTree; options.mNumSpecDecodingTokens = !isContext && params.mMaxSeqLenQ > 1 ? params.mMaxSeqLenQ : 0; - // Carry static tree length into FMHA kernel selection. + // Carry the tree length into FMHA kernel selection. options.mSpecDecodingTargetMaxGenLen = params.mSpecDecodingTargetMaxGenLen; options.mIsTrtllmLayout = true; From 8de1a1fafdef2b076262019b3f5d20fec36321cf Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 18 Sep 2026 01:29:42 +0000 Subject: [PATCH 19/21] [None][fix] rebase the GVR prior indexer test onto main's TopK contract test_indexer_forward_uses_prior_only_for_temporal_gvr was written against the rubin-advance TopK and indexer, and cc0120bf33 only rebased its metadata fixture. Two contract drifts remained and accounted for 108 unwaived failures in build 60859. main's TopK._forward_decode takes two extra caller-owned radix workspaces (radix_aux_indices, radix_aux_logits), so the nine-argument production call hit the seven-parameter stub with "takes 7 positional arguments but 9 were given". Widen the stub, and tolerate gvr_ext_kwargs being None: main passes no GVR dict at all when there is no prior, rather than a dict holding None. main also seeds the prior from prefill unconditionally and lets TopK.update_gvr_prior_from_prefill early-return for implementations that keep no prior, where rubin-advance guarded the call site instead. The observable behaviour is identical, so assert the call count against prefill alone; the prior-content assertions below already pin ownership. Test-only: no production file changes. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../attention/sparse/dsa/test_indexer_gvr_prior.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py index fa83ae59562c..f5d9adcbc977 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py @@ -131,9 +131,13 @@ def decode( output: torch.Tensor, n: int, max_seq_len: int, - extra: dict[str, torch.Tensor | None], + extra: dict[str, torch.Tensor | None] | None, + radix_aux_indices: torch.Tensor | None, + radix_aux_logits: torch.Tensor | None, ) -> torch.Tensor: - prior = extra["gvr_prior_indices"] + # Without a prior the indexer hands TopK no GVR kwargs at all, so the + # dict itself is absent rather than carrying a ``None`` entry. + prior = extra["gvr_prior_indices"] if extra is not None else None if top_k.needs_gvr_prior: assert prior.shape == (num_generations, topk) assert prior.data_ptr() == metadata.gvr_prior_indices[1].data_ptr() @@ -186,7 +190,10 @@ def prefill( is_generation=is_generation, ) assert decode_call.call_count == int(has_decode) - assert seed_call.call_count == int(top_k.needs_gvr_prior and has_prefill) + # The indexer seeds unconditionally after prefill; TopK itself is the one + # that drops the update when the implementation keeps no prior, so the call + # count tracks prefill alone and the prior contents assert the ownership. + assert seed_call.call_count == int(has_prefill) if seed_call.called: assert seed_call.call_args.args[1].data_ptr() == metadata.seq_lens_cuda.data_ptr() token_offset = num_ctx_tokens if is_generation is None else 0 From ac15adfb73ebd1da3d7a183a829de7d89fe3468e Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 18 Sep 2026 01:30:17 +0000 Subject: [PATCH 20/21] [None][fix] teach the Flux guard fixture the shared fused SwiGLU predicate 3fed8f90e5 moved the Flux2 fused-MLP dispatch off a bare use_cute_dsl_blockscaling_mm read and onto Linear.can_use_cute_dsl_nvfp4_swiglu_blackwell(), so dispatch and the gate/up weight interleave can no longer disagree. The SimpleNamespace stand-in in _make_fake_flux2_parallel_attn still only carried the old attribute, so all three Flux2 single-stream guard tests raised AttributeError in build 60859. Expose the predicate on the stub and have it mirror use_cute_dsl_blockscaling_mm, which is the attribute test_flux2_single_stream_cute_dsl_guard_requires_interleaved_weights toggles, so the guard keeps testing what it was written to test. Test-only: no production file changes. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tests/unittest/_torch/visual_gen/test_flux_transformer.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/unittest/_torch/visual_gen/test_flux_transformer.py b/tests/unittest/_torch/visual_gen/test_flux_transformer.py index f77662e76cb5..c839bd42fb63 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_flux_transformer.py @@ -90,6 +90,13 @@ def _make_fake_flux2_parallel_attn(): pre_quant_scale=None, force_dynamic_quantization=False, ) + # Flux now dispatches on Linear.can_use_cute_dsl_nvfp4_swiglu_blackwell(), the + # same predicate that drives the gate/up interleave, so the stand-in has to + # expose it. Mirror the predicate's dependence on use_cute_dsl_blockscaling_mm + # to keep the guard tests below meaningful. + gate_up_proj.can_use_cute_dsl_nvfp4_swiglu_blackwell = ( + lambda: gate_up_proj.use_cute_dsl_blockscaling_mm + ) attn.to_qkv_mlp_proj = SimpleNamespace( tp_size=2, qkv_proj=object(), From 3cffef1467233f4e642910d671bc181b82f29488 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 18 Sep 2026 01:31:30 +0000 Subject: [PATCH 21/21] [None][fix] refresh the Rubin MoE op-schema expectations for SiTU eb7f559c6a threaded situ_beta / situ_linear_beta through the Rubin NVFP4 gather grouped-GEMM act-fusion ops but left the schema tests pinned to the pre-SiTU argument lists, so build 60859 failed test_rubin_moe_leaf_schema[nvfp4_fc1], test_rubin_moe_locality_domain_composite_schema[nvfp4_fc1] and the NVFP4 FC1 case of test_rubin_moe_locality_domain_composite_owns_concurrent_tuning. Add the two trailing float parameters and their disabled sentinel defaults to both schema expectations, and extend the runner-construction expectation with the canonicalized situ_beta / situ_linear_beta the composite op forwards. Reference SITU_BETA_DISABLED rather than a literal so the sentinel stays defined in one place. Test-only: no production file changes. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/thop/parallel/test_cute_dsl_moe.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 80404f613f5c..718e4b983613 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -23,6 +23,7 @@ from tensorrt_llm._torch.autotuner import AutoTuner, OptimizationProfile, TunableRunner from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + SITU_BETA_DISABLED, GroupedGemmInputsHelper, _get_sm107_nvfp4_default_mma_config, ) @@ -2875,12 +2876,16 @@ def _assert_rubin_moe_op_schema( "scaling_vector_size", "partition_id", "activation_type", + "situ_beta", + "situ_linear_beta", "precomputed_tactic", ), { "scaling_vector_size": 16, "partition_id": -1, "activation_type": int(ActivationType.Swiglu), + "situ_beta": SITU_BETA_DISABLED, + "situ_linear_beta": SITU_BETA_DISABLED, "precomputed_tactic": None, }, {"output_tensor", "output_sf_tensor"}, @@ -3147,10 +3152,14 @@ def test_rubin_bf16_moe_precomputed_tactic_fake_signatures(): "output_sf_tensor", "scaling_vector_size", "activation_type", + "situ_beta", + "situ_linear_beta", ), { "scaling_vector_size": 16, "activation_type": int(ActivationType.Swiglu), + "situ_beta": SITU_BETA_DISABLED, + "situ_linear_beta": SITU_BETA_DISABLED, }, {"output_tensor", "output_sf_tensor"}, id="nvfp4_fc1", @@ -3604,7 +3613,12 @@ def fake_moe_output_memset(*args, **kwargs): runner_args, runner_kwargs = runner_instances[0].init_call if quantized and is_fc1: assert runner_args == (1, 1, 1, 0, 128, 16) - assert runner_kwargs == {"activation_type": ActivationType.Swiglu} + # The disabled sentinel canonicalizes to None before the runner sees it. + assert runner_kwargs == { + "activation_type": ActivationType.Swiglu, + "situ_beta": None, + "situ_linear_beta": None, + } elif quantized: assert runner_args == (1, 1, 1, 0, 128, torch.bfloat16, 16) assert not runner_kwargs