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

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Halcyon <56064364+VAthree@users.noreply.github.com>

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
 

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 \

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
+SM90 GPUs (H100/H200) the speedup over the default attention backend may be
-be limited or absent.
+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
+        # Dynamic-scheduled decode (fp8 always; bf16 when the installed hpc
-        # decode CUDA-graph max batch size in init_cuda_graph_state (sizing it
+        # supports it): the task workspace is sized by the decode CUDA-graph
-        # by the full request-pool capacity would cost hundreds of MB); eager
+        # max batch size in init_cuda_graph_state (sizing it by the full
-        # decode batches beyond that fall back to static split-K.
+        # 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._dynamic_sched = self.use_fp8 or _bf16_decode_supports_task_map()
-        self._fp8_task_map_max_bs = 0
+        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
+            self._decode_task_map is not None
-            and self._fp8_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

@@ -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
+Only supported on NVIDIA Hopper (SM90; the kernels ship sm90a). Note that the
-kernels are currently tuned primarily for H20: on other GPUs (H100/H200/B200,
+HPC-Ops kernels are currently tuned primarily for H20: on other SM90 GPUs
-...) the speedup over the default MoE runner may be limited or absent. Enable
+(H100/H200) the speedup over the default MoE runner may be limited or absent.
-it explicitly with ``--moe-runner-backend hpc_ops``.
+Enable it explicitly with ``--moe-runner-backend hpc_ops``.
 """
 
 import functools

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.

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 = [

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(),
+    has_hpc_ops() and _sm90(),
-    "requires HPC-Ops (install from source: https://github.com/Tencent/hpc-ops) and sm90+",
+    "requires HPC-Ops (install from source: https://github.com/Tencent/hpc-ops) and an SM90 (Hopper) GPU",
 )
 class TestHpcOpsMoeBlockwise(CustomTestCase):