执行摘要

本 PR 解除了 mega_moe_pre_dispatch 内核的 hidden 维度上限（原为 8192）。作者通过引入 kMultiChunk 模板参数，让每个线程以 strided chunk 方式处理更宽的行，同时保留单块特化路径，使既有形状的指令路径（SASS）完全不变。位精确对比覆盖 3584～16384 多个宽度，未发现数值或性能回退，但该变更缺少自动化单测，建议后续补齐。

功能与动机

mega_moe_pre_dispatch 是 DeepSeek V4 MoE 预调度路径上的 FP8/UE8M0 量化内核，原实现每个线程负责 8 个 bf16 元素，一个 CUDA block 最多 1024 线程，因此单 block 只能覆盖 hidden ≤ 8192。PR body 明确指出：Wider shapes are rejected before launch。为了让更大 hidden 维度的模型能够运行，必须扩展该内核。

实现拆解

变更入口为 python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh。

1. 引入 block 上限常量：新增 kMaxBlockThreads = 1024，__launch_bounds__ 改为引用该常量。
2. 增加 kMultiChunk 模板参数：模板签名扩展为 template<uint32_t kGroupSize, bool kUsePDL, bool kMultiChunk>。作者在 PR 中说明，之所以用模板参数而非运行时判断，是因为内核 issue-bound，循环簿记会对单块行产生约 10% 开销，必须用编译期特化隔离。
3. 重构量化逻辑为 quantize_chunk lambda：将加载、absmax 归约、UE8M0 scale 计算封装为可按任意 chunk 索引调用的函数体，加载由 tid 改为 chunk 索引。同时新增 static_assert(kMaxBlockThreads % kThreadsPerGroup == 0)，保证跨 chunk 时不破坏量化组内的 warp 归约。
4. 宽行分块与边界拒绝：多 chunk 路径按 blockDim.x 步长循环处理各 chunk；若尾部无法凑成整量化组（如 hidden=8320），启动前即拒绝。
5. 验证与配套：作者以位精确方式验证了 hidden=3584、5632、8192、12288、16384 的 FP8 输出与 UE8M0 scale 字节，确认 top-k 拷贝和 idle-slot 填充正确；单块特化 SASS 与旧内核一致，CUDA-event 微基准无回退。未提交自动化单测。

以下片段根据 PR 描述与 diff 整理，展示多 chunk 分块的核心结构；实际循环边界与实例化细节以仓库代码为准。

// 文件：python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh
// CUDA block 大小上限：每个 CTA 处理一个 token 行，
// 超过 kMaxBlockThreads * 8 个 bf16 元素的行会被拆成多个 strided chunk。
inline constexpr uint32_t kMaxBlockThreads = 1024;
​
// kMultiChunk 作为模板参数而非运行时分支，是因为该内核是 issue-bound 的，
// 循环簿记本身就会给单块可覆盖的行带来约 10% 的开销，因此需要编译期特化。
template <uint32_t kGroupSize, bool kUsePDL, bool kMultiChunk>
__global__ __launch_bounds__(kMaxBlockThreads, 2) void
mega_moe_pre_dispatch_kernel(const MegaMoEPreDispatchParams __grid_constant__ params) {
    using namespace device;
    constexpr uint32_t kVecElems = 8; // 8 个 bf16 = 16B 向量加载
    static_assert(kGroupSize % kVecElems == 0, "group_size must be a multiple of 8");
    constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems;
    // block 必须按整个量化组步进，否则跨 chunk 会破坏量化组内的 warp 归约
    static_assert(kMaxBlockThreads % kThreadsPerGroup == 0,
                  "block must stride by whole quant groups");
​
    const uint32_t token_id = blockIdx.x;
    const uint32_t tid = threadIdx.x;
    const auto token_in = params.x + static_cast<uint64_t>(token_id) * params.hidden;
    const auto token_out = params.buf_x + static_cast<uint64_t>(token_id) * params.hidden;
​
    // 量化一个 8 元素 chunk：加载输入、计算 absmax、warp 归约求 group scale、
    // 转成 UE8M0 指数，再写回 FP8 输出与 scale buffer。
    const auto quantize_chunk = [&](uint32_t chunk) {
        InputVec in_vec;
        in_vec.load(token_in, chunk); // 按 chunk 索引加载，支持多块处理
        float local_max = 0.0f;
        float vals[kVecElems];
        #pragma unroll
        for (uint32_t i = 0; i < kVecElems / 2; ++i) {
            const auto [v0, v1] = cast<fp32x2_t>(in_vec[i]);
            vals[2 * i + 0] = v0;
            vals[2 * i + 1] = v1;
            local_max = fmaxf(local_max, fmaxf(fabsf(v0), fabsf(v1)));
        }
        local_max = warp::reduce_max<kThreadsPerGroup>(local_max);
        const float absmax = fmaxf(local_max, 1e-10f);
        const float raw_scale = absmax / math::FP8_E4M3_MAX;
        const uint32_t ue8m0_exp = cast_to_ue8m0(raw_scale);
        const float inv_scale = __uint_as_float((127u + 127u - ue8m0_exp) << 23);
        // 后续：对 vals[i] 乘 inv_scale 并 pack 成 FP8，写回 token_out 与 scale buffer。
    };
​
    if (kMultiChunk) {
        // 宽行路径：每个线程按 blockDim.x 步长处理多个 chunk，
        // chunk 索引即该线程在行内的量化组偏移。
        for (uint32_t chunk = tid; chunk * kVecElems < params.hidden; chunk += blockDim.x) {
            quantize_chunk(chunk);
        }
    } else {
        // 单块特化：保持与历史内核完全一致的指令流，避免任何额外开销。
        quantize_chunk(tid);
    }
}

评论区精华

本 PR 没有产生 Review 评论。值得关注的“讨论”来自作者在 PR body 中的设计自述：

kMultiChunk is a template parameter rather than a runtime check because this kernel is issue-bound: the loop bookkeeping alone costs ~10% on the rows that fit a single chunk.

这一权衡解释了为什么宁可多编译一份内核特化，也不让单块路径承担任何额外分支开销。

风险与影响

• 正确性风险：多 chunk 路径下每个 chunk 的量化组索引需要重新计算，边界对齐依赖 static_assert 与调用方传入的 hidden 值；若未来有非对齐宽度混入，会破坏 UE8M0 scale 布局。
• 测试缺口：仅手工验证 5 个宽度，未固化自动化 kernel 测试，回归风险较高。
• 性能验证有限：多 chunk 路径仅微基准，缺少端到端吞吐数据。
• 影响范围：对 hidden ≤ 8192 的既有模型无行为变化（单块特化保留）；对超大 hidden 模型是能力解锁。

关联脉络

本 PR 与同目录的 deepseek_v4 JIT kernel 相关 PR 属于同一演进线：

• 35041 对 topk_v2.cuh 做内核特化与 PDL 时序收紧，同样以“指令路径不变”为约束；

• 29525/#35568 的 DeepEPv2 路线虽然关注 MoE A2A 后端，但也体现了 DeepSeek MoE 调度路径持续演进的背景。

这条内核优化线表明 SGLang 正在系统性地扩展 DeepSeek V4 相关内核的形状覆盖与执行效率。