dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

fix(build): properly convert all glog CHECK/LOG to TORCH_CHECK syntax

Browse Source
This commit is contained in:
project6-dev
2026-08-12 06:57:33 +00:00
parent 1dba2bbf61
commit 5c936782b7
1 changed files with 11 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
qwen3_6_scripts/moe_topk_softmax_kernels.cuh
View File

@@ -704,9 +704,7 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
break;
default: {
TORCH_CHECK(softmax_workspace != nullptr)
<< "softmax_workspace must be provided for num_experts that are "
"not a power of 2.";
TORCH_CHECK(softmax_workspace != nullptr, "softmax_workspace must be provided for num_experts that are not a power of 2.");
static constexpr int TPB = 256;
moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
nullptr,
@@ -753,27 +751,18 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
// Check data type
TORCH_CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
gating_output.scalar_type() == at::ScalarType::Half ||
gating_output.scalar_type() == at::ScalarType::BFloat16)
<< "gating_output must be float32, float16, or bfloat16";
gating_output.scalar_type() == at::ScalarType::BFloat16,
"gating_output must be float32, float16, or bfloat16");
// Check dimensions
TORCH_CHECK(gating_output.dim() == 2)
<< "gating_output must be 2D tensor [num_tokens, num_experts]";
TORCH_CHECK(topk_weights.dim() == 2)
<< "topk_weights must be 2D tensor [num_tokens, topk]";
TORCH_CHECK(topk_indices.dim() == 2)
<< "topk_indices must be 2D tensor [num_tokens, topk]";
TORCH_CHECK(gating_output.dim() == 2, "gating_output must be 2D tensor [num_tokens, num_experts]");
TORCH_CHECK(topk_weights.dim() == 2, "topk_weights must be 2D tensor [num_tokens, topk]");
TORCH_CHECK(topk_indices.dim() == 2, "topk_indices must be 2D tensor [num_tokens, topk]");
// Check shapes
TORCH_CHECK(gating_output.size(0) == topk_weights.size(0))
<< "First dimension of topk_weights must match num_tokens in "
"gating_output"
<< "First dimension of topk_indices must match num_tokens in "
"gating_output";
TORCH_CHECK(gating_output.size(0) == topk_weights.size(0), "First dimension of topk_weights must match num_tokens in gating_output First dimension of topk_indices must match num_tokens in gating_output");
TORCH_CHECK(topk_weights.size(-1) == topk_indices.size(-1))
<< "Second dimension of topk_indices must match topk in topk_weights"
<< "topk must be less than or equal to num_experts";
TORCH_CHECK(topk_weights.size(-1) == topk_indices.size(-1), "Second dimension of topk_indices must match topk in topk_weights topk must be less than or equal to num_experts");
const int num_experts = static_cast<int>(gating_output.size(-1));
const int num_tokens = static_cast<int>(gating_output.size(0));
@@ -795,12 +784,9 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value();
TORCH_CHECK(bias_tensor.dim() == 1)
<< "correction_bias must be 1D tensor [num_experts]";
TORCH_CHECK(bias_tensor.size(0) == num_experts)
<< "correction_bias size must match num_experts";
TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
<< "correction_bias must be float32, got " << bias_tensor.scalar_type();
TORCH_CHECK(bias_tensor.dim() == 1, "correction_bias must be 1D tensor [num_experts]");
TORCH_CHECK(bias_tensor.size(0) == num_experts, "correction_bias size must match num_experts");
TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "correction_bias must be float32");
bias_ptr = bias_tensor.data_ptr<float>();
}