From 8d6549bc4039d33635844495d86684677a4f0df8 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 27 Jul 2026 21:31:04 +0800 Subject: [PATCH] [Attention Backend] Extend hpc_ops dynamic-scheduled decode to bf16 (#32304) Co-authored-by: Claude Fable 5 Co-authored-by: Halcyon <56064364+VAthree@users.noreply.github.com> --- docker/Dockerfile | 2 +- .../advanced_features/attention_backend.mdx | 2 +- .../srt/layers/attention/hpc_ops_backend.py | 76 +++++++++++++------ .../srt/layers/moe/moe_runner/hpc_ops.py | 8 +- .../srt/layers/moe/moe_runner/runner.py | 9 +++ python/sglang/srt/server_args.py | 4 +- test/registered/moe/test_hpc_ops_moe.py | 10 +-- 7 files changed, 73 insertions(+), 38 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 1f37d7686..89b76b1eb 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -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 diff --git a/docs_new/docs/advanced_features/attention_backend.mdx b/docs_new/docs/advanced_features/attention_backend.mdx index a944751ab..eb1afedb9 100644 --- a/docs_new/docs/advanced_features/attention_backend.mdx +++ b/docs_new/docs/advanced_features/attention_backend.mdx @@ -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 \ diff --git a/python/sglang/srt/layers/attention/hpc_ops_backend.py b/python/sglang/srt/layers/attention/hpc_ops_backend.py index 69c630e5a..0c5a1382e 100644 --- a/python/sglang/srt/layers/attention/hpc_ops_backend.py +++ b/python/sglang/srt/layers/attention/hpc_ops_backend.py @@ -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) diff --git a/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py b/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py index cc7259116..102d6cebd 100644 --- a/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py +++ b/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index a0a3a92c5..6aff7365e 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -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. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 3cf1d3410..ce4afa512 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 = [ diff --git a/test/registered/moe/test_hpc_ops_moe.py b/test/registered/moe/test_hpc_ops_moe.py index a37b25508..0e30bcb45 100644 --- a/test/registered/moe/test_hpc_ops_moe.py +++ b/test/registered/moe/test_hpc_ops_moe.py @@ -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):