minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[Attention Backend] Extend hpc_ops dynamic-scheduled decode to bf16 (#32304)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Halcyon <56064364+VAthree@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-27 21:31:04 +08:00
committed by GitHub
co-authored by Claude Fable 5 Halcyon
parent 5656de2d9a
commit 8d6549bc40
7 changed files with 73 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/Dockerfile
+1 -1
View File
@@ -342,7 +342,7 @@ FROM torch_deps AS hpc_ops_builder
# HPC-Ops (https://github.com/Tencent/hpc-ops, MIT): fused attention / MoE /
# RoPE kernels from the Tencent Hunyuan AI Infra team, consumed by the opt-in
# hpc_ops attention and MoE runner backends.
ARG HPC_OPS_COMMIT=2404f09766269d9533b66709f705c8eeb01421fc
ARG HPC_OPS_COMMIT=6e2ecede9d6d47b2e680e839cc7ad7422bc8d88b
WORKDIR /build
docs_new/docs/advanced_features/attention_backend.mdx
+1 -1
View File
@@ -627,7 +627,7 @@ python3 -m sglang.launch_server \
--trust-remote-code
```
- HPC-Ops (MHA kernels from [HPC-Ops](https://github.com/Tencent/hpc-ops) by the Tencent Hunyuan AI Infra team; Hopper+, requires installing the `hpc` package from source, page size 64, bf16 or fp8_e4m3 KV cache, head_dim 128, q/kv head group 4 or 8)
- HPC-Ops (MHA kernels from [HPC-Ops](https://github.com/Tencent/hpc-ops) by the Tencent Hunyuan AI Infra team; Hopper (SM90) only, requires installing the `hpc` package from source, page size 64, bf16 or fp8_e4m3 KV cache, head_dim 128, q/kv head group 4 or 8)
```bash Command
python3 -m sglang.launch_server \
--model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 \
python/sglang/srt/layers/attention/hpc_ops_backend.py
+51 -25
View File
@@ -11,7 +11,7 @@ token-major KV pool as a paged NHD cache ``(num_pages, page_size, num_kv_heads,
head_dim)`` without any copy, so the only hard requirements are the kernel
constraints:
- NVIDIA Hopper or newer (sm90+)
- NVIDIA Hopper (SM90 only; the kernels ship sm90a)
- ``--page-size 64``
- bf16 model dtype; bf16 or fp8_e4m3 KV cache (the FP8 path additionally
requires the model to run the fused QKNorm+RoPE+quant+StoreKV op via
@@ -21,8 +21,8 @@ constraints:
no logit cap, decoder-only attention
Note that the HPC-Ops kernels are currently tuned primarily for H20: on other
GPUs (H100/H200/B200, ...) the speedup over the default attention backend may
be limited or absent.
SM90 GPUs (H100/H200) the speedup over the default attention backend may be
limited or absent.
Enable it explicitly with ``--attention-backend hpc_ops``.
"""
@@ -67,7 +67,7 @@ _REQUIRED_PAGE_SIZE = 64
FP8_ROPE_SUPPORTED_HEAD_CONFIGS = ((8, 1), (64, 8))
# Minimum tokens each SM processes per task in the dynamic-scheduled decode
# path (matches the HPC-Ops default).
_FP8_DYNAMIC_SCHED_MIN_PROCESS_LEN = 512
_DYNAMIC_SCHED_MIN_PROCESS_LEN = 512
@functools.cache
@@ -76,6 +76,16 @@ def has_hpc_ops() -> bool:
return importlib.util.find_spec("hpc") is not None
@functools.cache
def _bf16_decode_supports_task_map() -> bool:
"""hpc >= 6e2eced (PR #73) extends dynamic scheduling to bf16 decode."""
import inspect
import hpc
return "task_map" in inspect.signature(hpc.attention_decode_bf16).parameters
class HPCOpsMetadata(msgspec.Struct):
"""Per-forward-batch metadata consumed by the HPC-Ops kernels."""
@@ -101,7 +111,7 @@ class HPCOpsMetadata(msgspec.Struct):
hpc_q_scale: Optional[torch.Tensor] = None
# Split-K flag tensor for FP8 decode. shape: [bs, num_kv_heads], int32
hpc_split_k_flag: Optional[torch.Tensor] = None
# Pre-scheduled decode task map (dynamic-scheduled FP8 decode); None falls
# Pre-scheduled decode task map (dynamic-scheduled decode); None falls
# back to the kernel's static split-K scheduling.
hpc_task_map: Optional[torch.Tensor] = None
@@ -122,6 +132,13 @@ class HPCOpsAttnBackend(AttentionBackend):
"Install it from https://github.com/Tencent/hpc-ops"
)
major, minor = torch.cuda.get_device_capability()
if major != 9:
raise ValueError(
"The hpc_ops attention backend requires an SM90 (Hopper) GPU "
f"(the HPC-Ops kernels ship sm90a only), got sm{major}{minor}."
)
if model_runner.spec_algorithm.is_speculative():
raise ValueError(
"The hpc_ops attention backend does not support speculative "
@@ -191,13 +208,15 @@ class HPCOpsAttnBackend(AttentionBackend):
# Fallback per-tensor KV scale for checkpoints without kv scales.
self._ones_scale = torch.ones(1, dtype=torch.float32, device=self.device)
# Dynamic-scheduled FP8 decode: the task workspace is sized by the
# decode CUDA-graph max batch size in init_cuda_graph_state (sizing it
# by the full request-pool capacity would cost hundreds of MB); eager
# decode batches beyond that fall back to static split-K.
# Dynamic-scheduled decode (fp8 always; bf16 when the installed hpc
# supports it): the task workspace is sized by the decode CUDA-graph
# max batch size in init_cuda_graph_state (sizing it by the full
# request-pool capacity would cost hundreds of MB); eager decode
# batches beyond that fall back to static split-K.
self.num_kv_heads = num_kv_heads
self._fp8_task_map: Optional[torch.Tensor] = None
self._fp8_task_map_max_bs = 0
self._dynamic_sched = self.use_fp8 or _bf16_decode_supports_task_map()
self._decode_task_map: Optional[torch.Tensor] = None
self._decode_task_map_max_bs = 0
# CUDA graph state (allocated in init_cuda_graph_state).
self.decode_cuda_graph_metadata = {}
@@ -239,27 +258,26 @@ class HPCOpsAttnBackend(AttentionBackend):
page_size=self.page_size,
)
if (
self.use_fp8
and self._fp8_task_map is not None
self._decode_task_map is not None
and forward_batch.forward_mode.is_decode()
and batch_size <= self._fp8_task_map_max_bs
and batch_size <= self._decode_task_map_max_bs
):
metadata.hpc_task_map = self._assign_fp8_decode_tasks(
metadata.hpc_task_map = self._assign_decode_tasks(
metadata.cache_seqlens_int32
)
self.forward_metadata = metadata
def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
if self.use_fp8:
if self._dynamic_sched:
import hpc
self._fp8_task_map = hpc.get_attention_decode_task_workspace(
self._decode_task_map = hpc.get_attention_decode_task_workspace(
max_bs,
self.max_context_len,
self.num_kv_heads,
min_process_len=_FP8_DYNAMIC_SCHED_MIN_PROCESS_LEN,
min_process_len=_DYNAMIC_SCHED_MIN_PROCESS_LEN,
)
self._fp8_task_map_max_bs = max_bs
self._decode_task_map_max_bs = max_bs
self.decode_cuda_graph_metadata = {
"cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
@@ -309,10 +327,10 @@ class HPCOpsAttnBackend(AttentionBackend):
max_seq_pages=self.max_num_pages,
page_size=self.page_size,
)
if self.use_fp8 and self._fp8_task_map is not None:
if self._decode_task_map is not None:
# Recorded into the decode graph, so the task map is re-populated
# from the live seq_lens at every replay.
metadata.hpc_task_map = self._assign_fp8_decode_tasks(
metadata.hpc_task_map = self._assign_decode_tasks(
metadata.cache_seqlens_int32
)
self.forward_metadata = metadata
@@ -479,7 +497,7 @@ class HPCOpsAttnBackend(AttentionBackend):
metadata.hpc_q_scale = q_scale
metadata.hpc_split_k_flag = split_k_flag
def _assign_fp8_decode_tasks(self, cache_seqlens: torch.Tensor) -> torch.Tensor:
def _assign_decode_tasks(self, cache_seqlens: torch.Tensor) -> torch.Tensor:
"""Populate the dynamic-scheduled decode task map from live KV lengths.
The scheduler pass bins every (request, kv_head, KV-tile) chunk into
@@ -494,13 +512,13 @@ class HPCOpsAttnBackend(AttentionBackend):
# kernel launch with an invalid configuration.
hpc.assign_attention_decode_task(
cache_seqlens,
self._fp8_task_map,
self._decode_task_map,
self.num_kv_heads,
mtp=1,
new_kv_included=True,
min_process_len=_FP8_DYNAMIC_SCHED_MIN_PROCESS_LEN,
min_process_len=_DYNAMIC_SCHED_MIN_PROCESS_LEN,
)
return self._fp8_task_map
return self._decode_task_map
def _take_fp8_scales(self, metadata: HPCOpsMetadata):
"""Pop the per-layer FP8 scales written by the fused RoPE op."""
@@ -619,6 +637,13 @@ class HPCOpsAttnBackend(AttentionBackend):
split_flag=split_k_flag,
)
else:
# Older hpc's bf16 decode has no task_map kwarg; the map is only
# ever assigned when the installed hpc supports it.
task_map_kwargs = (
{"task_map": metadata.hpc_task_map}
if metadata.hpc_task_map is not None
else {}
)
o = hpc.attention_decode_bf16(
q.view(-1, layer.tp_q_head_num, layer.head_dim),
k_cache,
@@ -628,6 +653,7 @@ class HPCOpsAttnBackend(AttentionBackend):
mtp=0,
new_kv_included=True,
splitk=True,
**task_map_kwargs,
)
return o.view(-1, layer.tp_q_head_num * layer.head_dim)
python/sglang/srt/layers/moe/moe_runner/hpc_ops.py
+4 -4
View File
@@ -15,10 +15,10 @@ reduce into one call and consume *global* top-k expert ids together with
``rank_ep`` / ``num_expert_total``, so expert parallelism with contiguous
expert partitioning works without a local-expert remap.
Only supported on NVIDIA Hopper / Blackwell (sm90+). Note that the HPC-Ops
kernels are currently tuned primarily for H20: on other GPUs (H100/H200/B200,
...) the speedup over the default MoE runner may be limited or absent. Enable
it explicitly with ``--moe-runner-backend hpc_ops``.
Only supported on NVIDIA Hopper (SM90; the kernels ship sm90a). Note that the
HPC-Ops kernels are currently tuned primarily for H20: on other SM90 GPUs
(H100/H200) the speedup over the default MoE runner may be limited or absent.
Enable it explicitly with ``--moe-runner-backend hpc_ops``.
"""
import functools
python/sglang/srt/layers/moe/moe_runner/runner.py
+9
View File
@@ -96,6 +96,15 @@ class MoeRunner:
elif runner_backend.is_cutlass():
self.runner_core = None # CUTLASS uses the direct cutlass_moe_fp4 path
elif runner_backend.is_hpc_ops():
import torch
major, minor = torch.cuda.get_device_capability()
if major != 9:
raise ValueError(
"--moe-runner-backend hpc_ops requires an SM90 (Hopper) "
"GPU (the HPC-Ops kernels ship sm90a only), got "
f"sm{major}{minor}."
)
self.runner_core = None # HPC-Ops only supports the fused path
# Import here (not at module top, to avoid a circular import) to
# register the hpc_ops fused func before the pool lookup.
python/sglang/srt/server_args.py
+2 -2
View File
@@ -190,7 +190,7 @@ ATTENTION_BACKEND_CHOICES = [
"tokenspeed_mla",
"trtllm_mha",
"dual_chunk_flash_attn",
"hpc_ops", # HPC-Ops (https://github.com/Tencent/hpc-ops), Hopper+, requires --page-size 64
"hpc_ops", # HPC-Ops (https://github.com/Tencent/hpc-ops), Hopper (SM90) only, requires --page-size 64
# AMD specific
"aiter",
"wave",
@@ -259,7 +259,7 @@ MOE_RUNNER_BACKEND_CHOICES = [
"marlin",
"humming",
"experimental_sgl_marlin",
"hpc_ops", # HPC-Ops (https://github.com/Tencent/hpc-ops), FP8 MoE on Hopper+
"hpc_ops", # HPC-Ops (https://github.com/Tencent/hpc-ops), FP8 MoE on Hopper (SM90) only
]
MOE_A2A_BACKEND_CHOICES = [
test/registered/moe/test_hpc_ops_moe.py
+5 -5
View File
@@ -3,7 +3,7 @@
Compares the hpc_ops fused func (hpc.fuse_moe_blockwise) against the triton
fused_experts reference and an fp32 exact reference on realistic blockwise
FP8 quantized weights. Skipped when HPC-Ops (https://github.com/Tencent/hpc-ops)
is not installed or the GPU is older than sm90.
is not installed or the GPU is not SM90 (the kernels ship sm90a only).
"""
import os
@@ -35,11 +35,11 @@ register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-large")
E, TOPK, H, I = 128, 8, 2048, 768
def _sm90_or_newer() -> bool:
def _sm90() -> bool:
if not torch.cuda.is_available():
return False
major, _ = torch.cuda.get_device_capability()
return major >= 9
return major == 9
def _ensure_dist_initialized() -> None:
@@ -76,8 +76,8 @@ def _quant_blockwise(w: torch.Tensor, block: int = 128):
@unittest.skipUnless(
has_hpc_ops() and _sm90_or_newer(),
"requires HPC-Ops (install from source: https://github.com/Tencent/hpc-ops) and sm90+",
has_hpc_ops() and _sm90(),
"requires HPC-Ops (install from source: https://github.com/Tencent/hpc-ops) and an SM90 (Hopper) GPU",
)
class TestHpcOpsMoeBlockwise(CustomTestCase):