[TRTLLM-7318][feat] MnnvlThroughput AlltoAll implementation.

cpp/tensorrt_llm/common/envUtils.cpp

@@ -450,4 +450,43 @@ bool getEnvDisableChunkedAttentionInGenPhase()
     return getBoolEnv("TRTLLM_DISABLE_CHUNKED_ATTENTION_IN_GEN_PHASE");
 }
 
+bool getEnvMoeA2AOneBlockPerToken()
+{
+    // Default true; return false only if env set to "0"
+    static std::optional<int32_t> const val = getIntEnv("TLLM_MOE_A2A_ONE_BLOCK_PER_TOKEN");
+    if (!val.has_value())
+    {
+        return true;
+    }
+    return val.value() != 0;
+}
+
+static int sanitizeBlockSize(std::optional<int32_t> const& val)
+{
+    // Default 256 when not set or invalid
+    int block = val.value_or(256);
+    // Clamp to sane CUDA bounds and warp multiples
+    if (block <= 0)
+        block = 256;
+    if (block > 1024)
+        block = 1024;
+    // Round to nearest multiple of 32 (warp size)
+    block = (block + 31) / 32 * 32;
+    if (block == 0)
+        block = 256;
+    return block;
+}
+
+int getEnvMoeA2ADispatchBlockSize()
+{
+    static int const kBlock = sanitizeBlockSize(getIntEnv("TLLM_MOE_A2A_DISPATCH_BLOCK_SIZE"));
+    return kBlock;
+}
+
+int getEnvMoeA2ACombineBlockSize()
+{
+    static int const kBlock = sanitizeBlockSize(getIntEnv("TLLM_MOE_A2A_COMBINE_BLOCK_SIZE"));
+    return kBlock;
+}
+
 } // namespace tensorrt_llm::common

cpp/tensorrt_llm/common/envUtils.h

@@ -136,4 +136,13 @@ bool getEnvDisaggBenchmarkGenOnly();
 // Whether to disable the chunked-attention in the generation phase.
 bool getEnvDisableChunkedAttentionInGenPhase();
 
+// Whether to use one block per token for MoE A2A kernels (default true).
+bool getEnvMoeA2AOneBlockPerToken();
+
+// TODO: For DEV purpose temporarily.
+// Block size (threads per block) for MoE A2A Dispatch kernels (default 256 if unset or invalid)
+int getEnvMoeA2ADispatchBlockSize();
+// Block size (threads per block) for MoE A2A Combine kernels (default 256 if unset or invalid)
+int getEnvMoeA2ACombineBlockSize();
+
 } // namespace tensorrt_llm::common