[Kernel] Enable Kimi-K3 SiTU on the CuteDSL MoE backend and the SM107 low-latency GEMM plan - #54606
Conversation
…ackend Kimi-K3's routed experts use SiTU, which the CuteDSL experts rejected even though flashinfer >= 0.6.18 supports it in nvfp4 (w4a4) mode. Two gaps: - _supports_activation did not allowlist MoEActivation.SITU - situ_beta/situ_linear_beta (from moe_config.activation_situ_*) were never plumbed into the kernel call. The cute_dsl API keys SiTU on situ_beta and requires activation_type to remain a base type (ActivationType.Situ is rejected by normalize_cute_dsl_moe_activation_type), so the Swiglu base type is passed and SiTU rides the betas. Validated (as an equivalent runtime patch on the internal rubin build) with nvidia/Kimi-K3-NVFP4 on 2 nodes x 4 Rubin GPUs, TP8: --moe-backend flashinfer_cutedsl selects FLASHINFER_CUTEDSL, engine init and warmup pass, and greedy factual prompts match the flashinfer_trtllm reference backend (AI-assisted bringup). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Bolin Sun <bolins@nvidia.com>
The dsv3_fused_a / CuTe skinny-GEMM decode plan was gated to SM90/SM100/SM103 via _low_latency_table(), so on SM107 every unquantized BF16 decode projection fell back to cuBLASLt nvjet splitK + splitKreduce (two launches plus a reduce per GEMM). Map SM107 to the SM103 table: dsv3_fused_a_gemm only requires __CUDA_ARCH__ >= 900 and the CuTe skinny GEMM compiles for SM107 via CuteDSL (its PDL gate is major >= 9). Validated on Kimi-K3 NVFP4, TP8 across 8 Rubin GPUs, 8K ISL / 1K OSL serving: greedy outputs unchanged; fused_a_gemm_kernel replaces the nvjet splitK + splitKreduce pairs on all three live decode shapes (~40% fewer launches for those projections per rank, per-call time better on every replaced shape); end-to-end throughput and TPOT at parity or slightly better within run-to-run variance. The (N, K, M) winners in KIMI_K3_PROJECTIONS were measured on B300; the SM107 crossover points have not been re-measured and may deserve their own table as a follow-up. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Bolin Sun <bolins@nvidia.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
/ci run |
|
✅ @BolinSNLHM, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #86915 for commit |
|
/ci retry |
|
✅ Queued 3 failed job(s) for retry in Buildkite CI #86915. |
3 participants
Summary
Two enablement changes for Kimi-K3, ported from a downstream fork where they were validated end-to-end on SM107 (Rubin) hardware:
flashinfer_cutedsl_moe.py): acceptMoEActivation.SITUand plumbsitu_beta/situ_linear_betafrommoe_configintoflashinfer_cute_dsl_fused_moe_nvfp4. The cute_dsl kernel keys SiTU onsitu_betaand requires the baseSwigluactivation type (ActivationType.Situis rejected bynormalize_cute_dsl_moe_activation_type), so the base type is passed and SiTU rides the betas — the contract FlashInfer 0.6.18 (the pinned version) exposes. This makes--moe-backend flashinfer_cutedslusable for Kimi-K3.low_latency_gemm.py):_low_latency_table()now maps SM107 to the SM103 table, so the dsv3_fused_a / CuTe skinny-GEMM plan runs instead of falling back to cuBLASLt nvjet splitK + splitKreduce (two launches plus a reduce per GEMM).dsv3_fused_a_gemmonly requires__CUDA_ARCH__ >= 900, and the CuTe skinny GEMM compiles for SM107 via CuteDSL (its PDL gate ismajor >= 9).Not duplicating existing PRs
Searched open PRs for
sm107,dsv3_fused_a,low_latency_gemm,situ: #54565 relaxes dsv3 GEMM tensor-layout acceptance (complementary, different hunks of the same file); #52405 fixes SiTU output scale on the TRTLLM backend (not CuteDSL). No open PR adds SITU to the CuteDSL backend or SM107 to the low-latency plan.Testing
Validated on
nvidia/Kimi-K3-NVFP4, TP8 across 8x SM107 GPUs, 8K ISL / 1K OSL random serving (vllm bench serve --random-input-len 8192 --random-output-len 1024 --num-prompts 8 --max-concurrency 4 --ignore-eos),--moe-backend flashinfer_cutedsl --quantization modelopt_mixed --kv-cache-dtype fp8:flashinfer_trtllmreference backend on the same build.fused_a_gemm_kernelreplaces the nvjet splitK + splitKreduce pairs on all three live decode shapes — ~40% fewer launches for those projections and better per-call time on every replaced shape; the CuTe skinny GEMM compiles and fires at its table-selected token counts.Known follow-ups
KIMI_K3_PROJECTIONSwinners were measured on B300 (SM103); the SM107 per-M crossovers have not been re-measured and may deserve their own table (noted in the module docstring).situ_beta(the epilogue falls back to plain SwiGLU there) — tracked on the FlashInfer side.AI assistance disclosure
AI-assisted (Claude Code, see commit trailers); the submitter reviewed every changed line and ran the hardware validation above.