perf(dsv4): auto-dispatch BF16-to-FP32 GEMM backends

[dyq-w]

Add opt-in per-shape backend selection for DeepSeek-V4 router and
compressor BF16-to-FP32 GEMMs while keeping cublas as the default.

• port the HCU MFMA GEMM to the TVM-FFI JIT infrastructure
• select SGL, AITER, torch FP32, or cublas by runtime shape
• let auto mode use AITER independently of SGLANG_USE_AITER
• route the HCU DeepSeek-V4 router through the auto selector
• prewarm AITER metadata and SGL JIT before compile and graph capture
• fall back safely when optional AITER or SGL backends are unavailable
• add selector, HCU correctness, and performance coverage

Serving benchmark before/after ratios:

• mean TPOT: 1.60x, equivalent to a 37.5% reduction
• output throughput: 0.64x before/after, equivalent to a 1.56x
  optimized throughput or a 56.25% increase

---

CI States

Latest PR Test (Base): Not run yet
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.

[Copilot]

🟡 Changes recommended

Unresolved backend quarantine, lifecycle-state, and decode-only prewarm issues remain.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds opt-in, shape-based BF16-to-FP32 GEMM dispatch for DeepSeek-V4 on HCU while retaining cuBLAS as the default.

Changes:

• Adds HCU MFMA JIT kernel and wrapper.
• Selects SGL, AITER, Torch FP32, or cuBLAS by runtime shape.
• Adds prewarming, integration, correctness tests, and benchmarks.

File summaries

File	Summary	Review findings
test/registered/unit/test_linear_bf16_fp32_dispatch.py	Selector and API tests	None
test/registered/hcu/perf/bench_linear_bf16_fp32_hcu.py	HCU backend benchmarks	None
test/registered/hcu/kernels/test_linear_bf16_fp32_hcu.py	HCU correctness tests	None
python/sglang/srt/utils/aiter.py	AITER metadata prewarming	Nit (2 votes): Gate auto prewarming by the HCU/profile predicate.
python/sglang/srt/models/deepseek_v4.py	GEMM prewarming integration	None
python/sglang/srt/models/deepseek_v2.py	HCU router dispatch integration	None
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py	Pre-capture backend prewarming	Moderate (1 vote): Ensure decode-only graph capture also receives the prewarm.
python/sglang/kernels/ops/attention/dsv4/hcu_linear_bf16_fp32.py	HCU JIT wrapper and validation	None
python/sglang/kernels/ops/attention/dsv4/gemm.py	Backend selection and fallback logic	Moderate (1 vote): Quarantine AITER per failed shape rather than process-wide. Moderate (1 vote): Make backend state resettable across lifecycle rebuilds.
python/sglang/kernels/jit/csrc/deepseek_v4/linear_bf16_fp32.cuh	HCU MFMA GEMM implementation	None

Review details

Suppressed comments (3)

python/sglang/kernels/ops/attention/dsv4/gemm.py:357

• This turns any runtime failure for one (M, N, K) into a process-wide AITER quarantine, including ValueError from a shape-specific unsupported case. After one bad shape, every subsequent eligible shape is forced to torch/cuBLAS, so auto dispatch no longer selects per shape; cache the failed shape (or distinguish global initialization failures from shape failures) instead of setting _AUTO_AITER_DISABLED for all requests.

        except (OSError, RuntimeError, TypeError, ValueError):
            _AUTO_AITER_DISABLED = True
            logger.warning("Disabling AITER BF16 FP32 GEMM for auto dispatch")

python/sglang/kernels/ops/attention/dsv4/gemm.py:28

• These mutable disable/probe flags are written from lazy forward paths but live for the lifetime of the imported module, so reset_context() or a later engine/device cannot re-enable a backend after one transient AITER/SGL/JIT or out_dtype failure. Runtime state in this codebase is kept in resettable get_flags()/get_resources() slots (see test/registered/unit/test_module_state_ratchet.py:1-7); move these values to the owning lifecycle or explicitly reset them when that lifecycle is rebuilt.

_AUTO_AITER_DISABLED = False
_AUTO_SGL_DISABLED = False
_HIP_BF16_OUT_DTYPE_SUPPORTED: Optional[bool] = None

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py:373

• This prewarm is only reached when PrefillCudaGraphRunner is constructed. capture_prefill_graph can return before constructing that runner when prefill capture is disabled or has no valid buckets, while decode capture can still run afterward; in that configuration the first auto/AITER GEMM is executed during decode graph capture and can perform the rocminfo probe while the GPU context is locked. Move this call to a shared graph-setup path or also run it before decode capture so the stated prewarm guarantee covers decode-only configurations.

        # Auto mode may initialize AITER during the TC compile pass; prewarm
        # chip metadata before backend construction and capture.
        maybe_pre_warm_aiter_chip_info()

• Files reviewed: 10/10 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[hygon-ai-ai-reviewer]

AI Review

🟢 未发现有明确证据的问题

在已提供并完成审查的 diff 范围内未形成可确认问题。

变更概览

• PR：perf(dsv4): auto-dispatch BF16-to-FP32 GEMM backends
• 变更文件：10 个
• 代码变更：+1162/-14

文件审查摘要

文件	变更	审查结果
python/sglang/kernels/jit/csrc/deepseek_v4/linear_bf16_fp32.cuh	新增 · +375/-0	—
python/sglang/kernels/ops/attention/dsv4/gemm.py	修改 · +261/-8	—
test/registered/unit/test_linear_bf16_fp32_dispatch.py	新增 · +230/-0	—
test/registered/hcu/perf/bench_linear_bf16_fp32_hcu.py	新增 · +115/-0	—
python/sglang/kernels/ops/attention/dsv4/hcu_linear_bf16_fp32.py	新增 · +83/-0	—
test/registered/hcu/kernels/test_linear_bf16_fp32_hcu.py	新增 · +79/-0	—
python/sglang/srt/models/deepseek_v2.py	修改 · +10/-0	—
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py	修改 · +3/-5	—
python/sglang/srt/utils/aiter.py	修改 · +3/-1	—
python/sglang/srt/models/deepseek_v4.py	修改 · +3/-0	—

审查信息

• 覆盖情况：共 10 个文件，已完整审查 10 个。
• 候选问题：0 项；证据复核过滤：0 项；发布前敏感信息保护：0 项。
• 本服务只审查 GitHub 提供的 PR diff，未执行代码或重跑测试；结论仍需维护者核验。

提交修复并推送新 commit 后，将自动审查 PR 的最新版本。

[hygon-ai-ai-reviewer]

AI CI 失败分析

工作流：PR Test (HCU)
状态：已取消，但其中存在失败任务

总结

目前尚未定位直接失败根因。

• PR Test (HCU) finish：仅检测到上游结果检查失败，需在上游任务中定位根因。

与本次改动的关系

无法确定：当前证据尚未定位直接异常。

建议处理

1. 核对失败进程的原始异常输出或当前未选取的日志。