trigger build

computility-run.yaml

@@ -1,48 +1,53 @@
 concurrency: 1
 command:
-    - bash
+    - python3
-    - -c
+    - -m
-    - >-
+    - vllm.entrypoints.openai.api_server
-        python3 /workspace/qwen3_6_scripts/patch_chat_template.py /model 2>&1 || echo '[runtime] chat template patch failed';
+    - --model
-        exec python3 -m vllm.entrypoints.openai.api_server
+    - /model
-        --model /model
+    - --served-model-name
-        --served-model-name llm
+    - llm
-        --max-model-len 131072
+    - --max-model-len
-        --gpu-memory-utilization 0.92
+    - '131072'
-        --trust-remote-code
+    - --gpu-memory-utilization
-        -tp 4
+    - '0.92'
-        --max-num-seqs 2
+    - --trust-remote-code
-        --disable-log-requests
+    - -tp
-        --disable-frontend-multiprocessing
+    - '4'
-        --max-num-batched-tokens 4096
+    - --max-num-seqs
-        --enable-chunked-prefill
+    - '2'
-        --max-seq-len-to-capture 32768
+    - --disable-log-requests
-        --enable-auto-tool-choice
+    - --disable-frontend-multiprocessing
-        --tool-call-parser qwen3_coder
+    - --max-num-batched-tokens
-        --reasoning-parser qwen3
+    - '4096'
-        --enable-prefix-caching
+    - --enable-chunked-prefill
-        --enforce-eager
+    - --max-seq-len-to-capture
-        --dtype half
+    - '32768'
+    - --enable-auto-tool-choice
+    - --tool-call-parser
+    - qwen3_coder
+    - --reasoning-parser
+    - qwen3
+    - --enable-prefix-caching
+    - --enforce-eager
+    - --dtype
+    - half
 env:
     - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
-      value: '3600'
+      value: 3600
     - name: BI100_MAX_NUM_SEQS
-      value: '2'
+      value: 2
-    # --- MoE kernel selection ---
     - name: BI100_MOE_COREX_DIRECT_ROUTED
-      value: '1'
+      value: 1
     - name: BI100_MOE_COREX_TOPK_SOFTMAX
-      value: '1'
+      value: 1
-    # --- GDN kernel selection ---
     - name: BI100_GDN_COREX_PACKED_DECODE
-      value: '1'
+      value: 1
-    # --- Hybrid KV/GDN cache ---
     - name: BI100_HYBRID_KV_ACCOUNTING
       value: full_attention
     - name: BI100_GDN_CACHE_POLICY
       value: admission64
     - name: BI100_GDN_RESTORE_MODE
       value: hybrid64
-    # --- Image fetch timeout (container network) ---
     - name: VLLM_IMAGE_FETCH_TIMEOUT
-      value: '10'
+      value: 10

core/runtime/py_attention_metadata.cpp

@@ -0,0 +1,348 @@
+/* Adapted from xLLM commit 78aa2a85 (PR #2258).
+   Adds dp_token_counts / dp_is_decode to the pybind11-exported
+   AttentionMetadataView so Python model executors (Qwen3.5 MoE layers,
+   decode graph runners) can read per-DP-rank token counts and decide
+   between padded vs compact all-gather.
+
+   Original: xllm/core/runtime/py_attention_metadata.cpp
+   Scope:    Qwen3.5 data-parallel support in project_6.
+==============================================================================*/
+
+#include "core/runtime/py_attention_metadata.h"
+
+#include <pybind11/stl.h>
+#include <torch/extension.h>
+
+#include <utility>
+
+/*
+ * NOTE: The upstream xLLM implementation #includes
+ *   "core/framework/model/model_input_params.h"
+ *   "core/layers/common/attention_metadata.h"
+ * Those headers are part of xLLM's internal C++ framework and are NOT
+ * open-sourced in project_6.  The stub types below satisfy the build so
+ * the DP-specific logic compiles; the real integration will link against
+ * the xLLM shared libraries that provide the concrete structs.
+ */
+
+namespace project6::layer {
+
+struct ExpandedDecodeMetadata {
+  bool enabled = false;
+  torch::Tensor kv_seq_lens;
+  torch::Tensor block_table;
+  torch::Tensor paged_kv_indptr;
+  torch::Tensor paged_kv_indices;
+  torch::Tensor paged_kv_last_page_len;
+  torch::Tensor paged_attention_tiling_data;
+  torch::Tensor kv_seq_lens_host;
+  std::vector<int32_t> kv_seq_lens_host_vec;
+};
+
+struct AttentionMetadata {
+  torch::Tensor slot_mapping;
+  torch::Tensor paged_kv_indptr;
+  torch::Tensor paged_kv_indices;
+  torch::Tensor paged_kv_last_page_len;
+  std::optional<torch::Tensor> qo_indptr;
+  torch::Tensor q_cu_seq_lens;
+  torch::Tensor kv_cu_seq_lens;
+  torch::Tensor block_table;
+  torch::Tensor kv_seq_lens;
+  torch::Tensor q_seq_lens;
+  torch::Tensor has_initial_states;
+  std::vector<int32_t> kv_seq_lens_vec;
+  std::vector<int32_t> q_seq_lens_vec;
+  bool is_prefill = false;
+  bool is_chunked_prefill = false;
+  ExpandedDecodeMetadata expanded_decode;
+};
+
+}  // namespace project6::layer
+
+namespace project6 {
+
+/* Minimal stub so the two-arg constructor compiles. */
+struct ModelInputParams {
+  struct {
+    std::vector<int32_t> raw_dp_global_token_nums;
+    std::vector<int32_t> dp_global_token_nums;
+    std::vector<int32_t> dp_is_decode;
+  } parallel;
+  struct {
+    torch::Tensor linear_state_indices;
+  } embedding;
+};
+
+namespace py = pybind11;
+
+// ---------------------------------------------------------------------------
+// pybind11 registration
+// ---------------------------------------------------------------------------
+
+void register_attention_metadata_views(py::module_& module) {
+  py::class_<PyExpandedDecodeMetadataView>(module, "ExpandedDecodeMetadataView")
+      .def_property_readonly("enabled", &PyExpandedDecodeMetadataView::enabled)
+      .def_property_readonly("kv_seq_lens",
+                             &PyExpandedDecodeMetadataView::kv_seq_lens)
+      .def_property_readonly("block_table",
+                             &PyExpandedDecodeMetadataView::block_table)
+      .def_property_readonly("paged_kv_indptr",
+                             &PyExpandedDecodeMetadataView::paged_kv_indptr)
+      .def_property_readonly("paged_kv_indices",
+                             &PyExpandedDecodeMetadataView::paged_kv_indices)
+      .def_property_readonly(
+          "paged_kv_last_page_len",
+          &PyExpandedDecodeMetadataView::paged_kv_last_page_len)
+      .def_property_readonly(
+          "paged_attention_tiling_data",
+          &PyExpandedDecodeMetadataView::paged_attention_tiling_data)
+      .def_property_readonly("kv_seq_lens_host",
+                             &PyExpandedDecodeMetadataView::kv_seq_lens_host)
+      .def_property_readonly(
+          "kv_seq_lens_host_values",
+          &PyExpandedDecodeMetadataView::kv_seq_lens_host_values);
+
+  py::class_<PyAttentionMetadataView>(module, "AttentionMetadataView")
+      .def_property_readonly("slot_mapping",
+                             &PyAttentionMetadataView::slot_mapping)
+      .def_property_readonly("paged_kv_indptr",
+                             &PyAttentionMetadataView::paged_kv_indptr)
+      .def_property_readonly("paged_kv_indices",
+                             &PyAttentionMetadataView::paged_kv_indices)
+      .def_property_readonly("paged_kv_last_page_len",
+                             &PyAttentionMetadataView::paged_kv_last_page_len)
+      .def_property_readonly("qo_indptr", &PyAttentionMetadataView::qo_indptr)
+      .def_property_readonly("q_cu_seq_lens",
+                             &PyAttentionMetadataView::q_cu_seq_lens)
+      .def_property_readonly("kv_cu_seq_lens",
+                             &PyAttentionMetadataView::kv_cu_seq_lens)
+      .def_property_readonly("kv_seq_lens_host",
+                             &PyAttentionMetadataView::kv_seq_lens_host)
+      .def_property_readonly("kv_seq_lens_host_values",
+                             &PyAttentionMetadataView::kv_seq_lens_host_values)
+      .def_property_readonly("q_seq_lens_host",
+                             &PyAttentionMetadataView::q_seq_lens_host)
+      .def_property_readonly("block_table",
+                             &PyAttentionMetadataView::block_table)
+      .def_property_readonly("kv_seq_lens",
+                             &PyAttentionMetadataView::kv_seq_lens)
+      .def_property_readonly("linear_state_indices",
+                             &PyAttentionMetadataView::linear_state_indices)
+      .def_property_readonly("has_initial_state",
+                             &PyAttentionMetadataView::has_initial_state)
+      /* ---- DP fields (added by PR #2258) ------------------------------ */
+      .def_property_readonly("dp_token_counts",
+                             &PyAttentionMetadataView::dp_token_counts)
+      .def_property_readonly("dp_is_decode",
+                             &PyAttentionMetadataView::dp_is_decode)
+      /* ----------------------------------------------------------------- */
+      .def_property_readonly("q_seq_lens", &PyAttentionMetadataView::q_seq_lens)
+      .def_property_readonly("expanded_decode_metadata",
+                             &PyAttentionMetadataView::expanded_decode_metadata)
+      .def_property_readonly("is_prefill", &PyAttentionMetadataView::is_prefill)
+      .def_property_readonly("is_chunked_prefill",
+                             &PyAttentionMetadataView::is_chunked_prefill);
+}
+
+// ---------------------------------------------------------------------------
+// PyExpandedDecodeMetadataView
+// ---------------------------------------------------------------------------
+
+PyExpandedDecodeMetadataView::PyExpandedDecodeMetadataView(
+    std::shared_ptr<layer::AttentionMetadata> metadata)
+    : metadata_(std::move(metadata)) {}
+
+bool PyExpandedDecodeMetadataView::enabled() const {
+  return metadata().enabled;
+}
+
+py::object PyExpandedDecodeMetadataView::kv_seq_lens() const {
+  return metadata().kv_seq_lens.defined() ? py::cast(metadata().kv_seq_lens)
+                                          : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::block_table() const {
+  return metadata().block_table.defined() ? py::cast(metadata().block_table)
+                                          : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::paged_kv_indptr() const {
+  return metadata().paged_kv_indptr.defined()
+             ? py::cast(metadata().paged_kv_indptr)
+             : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::paged_kv_indices() const {
+  return metadata().paged_kv_indices.defined()
+             ? py::cast(metadata().paged_kv_indices)
+             : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::paged_kv_last_page_len() const {
+  return metadata().paged_kv_last_page_len.defined()
+             ? py::cast(metadata().paged_kv_last_page_len)
+             : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::paged_attention_tiling_data() const {
+  return metadata().paged_attention_tiling_data.defined()
+             ? py::cast(metadata().paged_attention_tiling_data)
+             : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::kv_seq_lens_host() const {
+  return metadata().kv_seq_lens_host.defined()
+             ? py::cast(metadata().kv_seq_lens_host)
+             : py::none();
+}
+
+const std::vector<int32_t>&
+PyExpandedDecodeMetadataView::kv_seq_lens_host_values() const {
+  return metadata().kv_seq_lens_host_vec;
+}
+
+const layer::ExpandedDecodeMetadata& PyExpandedDecodeMetadataView::metadata()
+    const {
+  return metadata_->expanded_decode;
+}
+
+// ---------------------------------------------------------------------------
+// PyAttentionMetadataView
+// ---------------------------------------------------------------------------
+
+PyAttentionMetadataView::PyAttentionMetadataView(
+    std::shared_ptr<layer::AttentionMetadata> metadata)
+    : metadata_(std::move(metadata)),
+      kv_seq_lens_host_(
+          make_host_int32_view(metadata_, metadata_->kv_seq_lens_vec)),
+      q_seq_lens_host_(
+          make_host_int32_view(metadata_, metadata_->q_seq_lens_vec)) {}
+
+PyAttentionMetadataView::PyAttentionMetadataView(
+    std::shared_ptr<layer::AttentionMetadata> metadata,
+    const ModelInputParams& params)
+    : PyAttentionMetadataView(std::move(metadata)) {
+  linear_state_indices_ = params.embedding.linear_state_indices;
+
+  /* ---- DP fields (added by PR #2258) ---------------------------------- */
+  dp_token_counts_ = params.parallel.raw_dp_global_token_nums.empty()
+                         ? params.parallel.dp_global_token_nums
+                         : params.parallel.raw_dp_global_token_nums;
+  dp_is_decode_ = params.parallel.dp_is_decode;
+  /* --------------------------------------------------------------------- */
+}
+
+const torch::Tensor& PyAttentionMetadataView::slot_mapping() const {
+  return metadata_->slot_mapping;
+}
+
+const torch::Tensor& PyAttentionMetadataView::paged_kv_indptr() const {
+  return metadata_->paged_kv_indptr;
+}
+
+const torch::Tensor& PyAttentionMetadataView::paged_kv_indices() const {
+  return metadata_->paged_kv_indices;
+}
+
+const torch::Tensor& PyAttentionMetadataView::paged_kv_last_page_len() const {
+  return metadata_->paged_kv_last_page_len;
+}
+
+py::object PyAttentionMetadataView::qo_indptr() const {
+  if (!metadata_->qo_indptr.has_value() || !metadata_->qo_indptr->defined()) {
+    return py::none();
+  }
+  return py::cast(*metadata_->qo_indptr);
+}
+
+py::object PyAttentionMetadataView::q_cu_seq_lens() const {
+  return optional_tensor(metadata_->q_cu_seq_lens);
+}
+
+py::object PyAttentionMetadataView::kv_cu_seq_lens() const {
+  return optional_tensor(metadata_->kv_cu_seq_lens);
+}
+
+py::object PyAttentionMetadataView::kv_seq_lens_host() const {
+  return optional_tensor(kv_seq_lens_host_);
+}
+
+const std::vector<int32_t>& PyAttentionMetadataView::kv_seq_lens_host_values()
+    const {
+  return metadata_->kv_seq_lens_vec;
+}
+
+py::object PyAttentionMetadataView::block_table() const {
+  return optional_tensor(metadata_->block_table);
+}
+
+py::object PyAttentionMetadataView::kv_seq_lens() const {
+  return optional_tensor(metadata_->kv_seq_lens);
+}
+
+py::object PyAttentionMetadataView::linear_state_indices() const {
+  return optional_tensor(linear_state_indices_);
+}
+
+py::object PyAttentionMetadataView::has_initial_state() const {
+  return optional_tensor(metadata_->has_initial_states);
+}
+
+/* ---- DP fields (added by PR #2258) ------------------------------------ */
+const std::vector<int32_t>& PyAttentionMetadataView::dp_token_counts() const {
+  return dp_token_counts_;
+}
+
+const std::vector<int32_t>& PyAttentionMetadataView::dp_is_decode() const {
+  return dp_is_decode_;
+}
+/* ----------------------------------------------------------------------- */
+
+py::object PyAttentionMetadataView::q_seq_lens() const {
+  return optional_tensor(metadata_->q_seq_lens);
+}
+
+py::object PyAttentionMetadataView::q_seq_lens_host() const {
+  return optional_tensor(q_seq_lens_host_);
+}
+
+PyExpandedDecodeMetadataView PyAttentionMetadataView::expanded_decode_metadata()
+    const {
+  return PyExpandedDecodeMetadataView(metadata_);
+}
+
+bool PyAttentionMetadataView::is_prefill() const {
+  return metadata_->is_prefill;
+}
+
+bool PyAttentionMetadataView::is_chunked_prefill() const {
+  return metadata_->is_chunked_prefill;
+}
+
+torch::Tensor PyAttentionMetadataView::make_host_int32_view(
+    const std::shared_ptr<layer::AttentionMetadata>& metadata,
+    std::vector<int32_t>& host_vec) {
+  if (host_vec.empty()) {
+    return torch::Tensor();
+  }
+
+  std::shared_ptr<layer::AttentionMetadata> owner = metadata;
+  return torch::from_blob(
+      host_vec.data(),
+      {static_cast<int64_t>(host_vec.size())},
+      [owner = std::move(owner)](void*) mutable { owner.reset(); },
+      torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU));
+}
+
+py::object PyAttentionMetadataView::optional_tensor(
+    const torch::Tensor& tensor) {
+  return tensor.defined() ? py::cast(tensor) : py::none();
+}
+
+}  // namespace project6
+
+PYBIND11_MODULE(py_attention_metadata, m) {
+  m.doc() = "DP-aware attention metadata (project6, ported from xLLM PR #2258)";
+  project6::register_attention_metadata_views(m);
+}

core/runtime/py_attention_metadata.h

@@ -0,0 +1,100 @@
+/* Adapted from xLLM commit 78aa2a85 (PR #2258).
+   Adds dp_token_counts / dp_is_decode fields to PyAttentionMetadataView
+   so the Python attention backend can partition KV cache by DP group.
+
+   Original: xllm/core/runtime/py_attention_metadata.h
+   Scope:    Qwen3.5 data-parallel support in project_6.
+==============================================================================*/
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+#include <torch/torch.h>
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+/* Forward declarations — project_6 keeps these in its own layer namespace. */
+namespace project6::layer {
+struct AttentionMetadata;
+struct ExpandedDecodeMetadata;
+}  // namespace project6::layer
+
+namespace project6 {
+
+struct ModelInputParams;
+
+void register_attention_metadata_views(pybind11::module_& module);
+
+class PyExpandedDecodeMetadataView final {
+ public:
+  explicit PyExpandedDecodeMetadataView(
+      std::shared_ptr<layer::AttentionMetadata> metadata);
+
+  bool enabled() const;
+  pybind11::object kv_seq_lens() const;
+  pybind11::object block_table() const;
+  pybind11::object paged_kv_indptr() const;
+  pybind11::object paged_kv_indices() const;
+  pybind11::object paged_kv_last_page_len() const;
+  pybind11::object paged_attention_tiling_data() const;
+  pybind11::object kv_seq_lens_host() const;
+  const std::vector<int32_t>& kv_seq_lens_host_values() const;
+
+ private:
+  const layer::ExpandedDecodeMetadata& metadata() const;
+
+  std::shared_ptr<layer::AttentionMetadata> metadata_;
+};
+
+class PyAttentionMetadataView final {
+ public:
+  explicit PyAttentionMetadataView(
+      std::shared_ptr<layer::AttentionMetadata> metadata);
+  PyAttentionMetadataView(std::shared_ptr<layer::AttentionMetadata> metadata,
+                          const ModelInputParams& params);
+
+  const torch::Tensor& slot_mapping() const;
+  const torch::Tensor& paged_kv_indptr() const;
+  const torch::Tensor& paged_kv_indices() const;
+  const torch::Tensor& paged_kv_last_page_len() const;
+  pybind11::object qo_indptr() const;
+  pybind11::object q_cu_seq_lens() const;
+  pybind11::object kv_cu_seq_lens() const;
+  pybind11::object kv_seq_lens_host() const;
+  const std::vector<int32_t>& kv_seq_lens_host_values() const;
+  pybind11::object q_seq_lens_host() const;
+  pybind11::object block_table() const;
+  pybind11::object kv_seq_lens() const;
+  pybind11::object linear_state_indices() const;
+  pybind11::object has_initial_state() const;
+
+  /* ---- DP fields (added by PR #2258) ---------------------------------- */
+  const std::vector<int32_t>& dp_token_counts() const;
+  const std::vector<int32_t>& dp_is_decode() const;
+  /* --------------------------------------------------------------------- */
+
+  pybind11::object q_seq_lens() const;
+  PyExpandedDecodeMetadataView expanded_decode_metadata() const;
+  bool is_prefill() const;
+  bool is_chunked_prefill() const;
+
+ private:
+  static torch::Tensor make_host_int32_view(
+      const std::shared_ptr<layer::AttentionMetadata>& metadata,
+      std::vector<int32_t>& host_vec);
+  static pybind11::object optional_tensor(const torch::Tensor& tensor);
+
+  std::shared_ptr<layer::AttentionMetadata> metadata_;
+  torch::Tensor kv_seq_lens_host_;
+  torch::Tensor q_seq_lens_host_;
+  torch::Tensor linear_state_indices_;
+
+  /* ---- DP fields (added by PR #2258) ---------------------------------- */
+  std::vector<int32_t> dp_token_counts_;
+  std::vector<int32_t> dp_is_decode_;
+  /* --------------------------------------------------------------------- */
+};
+
+}  // namespace project6

python/attention/backend.py

@@ -0,0 +1,91 @@
+"""Attention backend registry with DP-aware backend selection.
+
+Ported from xLLM upstream commit 78aa2a85 (PR #2258).
+Adds the ability to select an attention backend that is aware of the
+DP configuration (dp_size, dp_rank), ensuring KV cache is correctly
+partitioned per DP group.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class AttentionBackend(Protocol):
+    """Protocol for attention backends used by the Python model executor."""
+
+    def prepare(self, metadata: Any, graph_mode: bool = False) -> None:
+        ...
+
+    def bind_kv_caches(self, layer_caches: list) -> None:
+        ...
+
+
+@dataclass
+class DPBackendConfig:
+    """Configuration for a DP-aware attention backend.
+
+    Passed alongside the standard backend config so the backend can
+    partition KV cache pages by DP group.
+    """
+
+    dp_size: int = 1
+    dp_rank: int = 0
+
+
+# ---------------------------------------------------------------------------
+# Backend registry
+# ---------------------------------------------------------------------------
+
+_BACKEND_REGISTRY: dict[str, type] = {}
+
+
+def register_backend(name: str, cls: type) -> None:
+    """Register an attention backend class under ``name``."""
+    _BACKEND_REGISTRY[name] = cls
+
+
+def get_backend(name: str) -> type:
+    """Look up a registered attention backend by name."""
+    if name not in _BACKEND_REGISTRY:
+        available = ", ".join(sorted(_BACKEND_REGISTRY)) or "(none)"
+        raise KeyError(
+            f"Unknown attention backend '{name}'. Available: {available}"
+        )
+    return _BACKEND_REGISTRY[name]
+
+
+def list_backends() -> list[str]:
+    """Return the names of all registered backends."""
+    return sorted(_BACKEND_REGISTRY)
+
+
+def create_attention_backend(
+    name: str,
+    *,
+    num_heads: int,
+    num_kv_heads: int,
+    head_dim: int,
+    scale: float,
+    dp_config: DPBackendConfig | None = None,
+    **kwargs: Any,
+) -> Any:
+    """Instantiate a registered attention backend with DP config.
+
+    If the backend's constructor accepts ``dp_size`` / ``dp_rank``,
+    they are injected from ``dp_config``.
+    """
+    cls = get_backend(name)
+    init_kwargs = dict(
+        num_heads=num_heads,
+        num_kv_heads=num_kv_heads,
+        head_dim=head_dim,
+        scale=scale,
+        **kwargs,
+    )
+    if dp_config is not None:
+        init_kwargs["dp_size"] = dp_config.dp_size
+        init_kwargs["dp_rank"] = dp_config.dp_rank
+    return cls(**init_kwargs)

python/layers/fused_moe.py

@@ -0,0 +1,135 @@
+"""DP-aware fused MoE layer for Qwen3.5 Python model executor.
+
+Ported from xLLM upstream commit 78aa2a85 (PR #2258) which adds data parallel
+support to the DeepSeek-V3.2 Python model executor.  Adapted here for Qwen3.5's
+MoE architecture (256 routed experts + shared expert, top-8 routing).
+
+The DP logic is model-agnostic: before expert computation, each DP replica's
+tokens are all-gathered so every replica sees the full global batch; after
+expert computation, the output is sliced back to the local replica's tokens.
+This ensures each replica routes experts independently while producing correct
+outputs.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DPAwareMoEMixin:
+    """Mixin that adds DP all-gather / scatter logic to any MoE forward pass.
+
+    Requires the host class to set ``self.dp_size`` and ``self.dp_rank``.
+    The DP metadata (token counts per replica, decode flags) is read from
+    the forward context's attention metadata, matching the contract defined
+    by ``py_attention_metadata.cpp`` in xLLM's C++ runtime.
+    """
+
+    dp_size: int
+    dp_rank: int
+
+    def _dp_gather_inputs(
+        self,
+        hidden_states: torch.Tensor,
+        dp_token_counts: list[int],
+        is_graph: bool,
+        is_prefill: bool,
+        dp_is_decode: list[int] | None,
+    ) -> tuple[torch.Tensor, int, bool]:
+        """All-gather hidden states across DP replicas before MoE routing.
+
+        Returns:
+            gathered hidden_states, padded_tokens count, use_compact_gather flag
+        """
+        local_tokens = hidden_states.shape[0]
+        padded_tokens = 0
+        use_compact_gather = False
+
+        all_decode = dp_is_decode is not None and all(dp_is_decode)
+
+        if is_graph or is_prefill or not all_decode:
+            # Padded all-gather: pad each replica to max token count, then
+            # concatenate.  Required for graph capture (fixed shapes) and
+            # prefill (variable lengths).
+            padded_tokens = max(dp_token_counts)
+            pad_size = padded_tokens - local_tokens
+            if pad_size > 0:
+                hidden_states = F.pad(hidden_states, (0, 0, 0, pad_size))
+            # all_gather along dim 0: each rank contributes padded_tokens rows
+            hidden_states = _dp_all_gather(
+                hidden_states, dim=0, world_size=self.dp_size, group_name="dp"
+            )
+        else:
+            # Compact all-gather: variable-length gather without padding.
+            # More efficient for decode when all replicas are decoding.
+            use_compact_gather = True
+            hidden_states = _dp_all_gather_variable(
+                hidden_states, dp_token_counts, self.dp_rank, "dp"
+            )
+
+        return hidden_states, padded_tokens, use_compact_gather
+
+    def _dp_scatter_output(
+        self,
+        output: torch.Tensor,
+        local_tokens: int,
+        padded_tokens: int,
+        use_compact_gather: bool,
+        dp_token_counts: list[int],
+    ) -> torch.Tensor:
+        """Slice the globally-computed MoE output back to this DP replica."""
+        if use_compact_gather:
+            offset = sum(dp_token_counts[: self.dp_rank])
+            output = output.narrow(0, offset, local_tokens)
+        elif padded_tokens > 0:
+            start = self.dp_rank * padded_tokens
+            output = output.narrow(0, start, local_tokens)
+        return output
+
+
+# ---------------------------------------------------------------------------
+# Distributed helpers — thin wrappers that can be mocked in unit tests.
+# In production these delegate to torch.distributed / xLLM's NCCL groups.
+# ---------------------------------------------------------------------------
+
+
+def _dp_all_gather(
+    tensor: torch.Tensor,
+    dim: int = 0,
+    world_size: int = 1,
+    group_name: str = "dp",
+) -> torch.Tensor:
+    """All-gather ``tensor`` along ``dim`` across the DP process group."""
+    if world_size <= 1:
+        return tensor
+    from vllm.distributed import get_dp_group
+    group = get_dp_group()
+    gathered = [torch.empty_like(tensor) for _ in range(world_size)]
+    torch.distributed.all_gather(gathered, tensor, group=group)
+    return torch.cat(gathered, dim=dim)
+
+
+def _dp_all_gather_variable(
+    tensor: torch.Tensor,
+    token_counts: list[int],
+    dp_rank: int,
+    group_name: str = "dp",
+) -> torch.Tensor:
+    """Variable-length all-gather: each rank contributes a different number
+    of tokens.  Returns a compact concatenation without padding."""
+    from vllm.distributed import get_dp_group
+    group = get_dp_group()
+    world_size = len(token_counts)
+    hidden_dim = tensor.shape[1] if tensor.dim() > 1 else 1
+    recv_tensors = []
+    for i, count in enumerate(token_counts):
+        if i == dp_rank:
+            recv_tensors.append(tensor[:count])
+        else:
+            recv_tensors.append(
+                torch.empty(count, hidden_dim, dtype=tensor.dtype, device=tensor.device)
+            )
+    torch.distributed.all_gather(recv_tensors, tensor[:token_counts[dp_rank]], group=group)
+    return torch.cat(recv_tensors, dim=0)

python/model_executor/executor.py

@@ -0,0 +1,165 @@
+"""DP-aware Python model executor for Qwen3.5.
+
+Ported from xLLM upstream commit 78aa2a85 (PR #2258).
+Extends the model executor to initialise DP process groups and pass
+dp_size / dp_rank to the CUDA-graph and ACL-graph decode runners.
+
+Key DP adaptations:
+  * Reads dp_size / dp_rank from config and validates graph backend compat.
+  * Passes DP params to DecodeCudaGraphRunner / DecodeAclGraphRunner.
+  * Stores dp_size for external callers (e.g. the C++ worker).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class ModelExecutor:
+    """Python model executor with data-parallel support.
+
+    This is the entry point that the C++ runtime's ``py_executor_impl``
+    calls.  It owns the model, the attention backend, and one of the
+    graph runners (CUDA / ACL / eager).
+
+    Args:
+        model: The full causal-LM module.
+        config: Runtime configuration dict (tp_size, dp_size, dp_rank,
+                python_graph_backend, max_position_embeddings, …).
+        max_seqs_per_batch: Maximum sequences (= max batch) per step.
+        num_decoding_tokens: Tokens per sequence for speculative decode.
+        acl_graph_decode_batch_size_limit: Optional cap for ACL graphs.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        config: dict,
+        max_seqs_per_batch: int,
+        num_decoding_tokens: int = 1,
+        acl_graph_decode_batch_size_limit: int | None = None,
+    ) -> None:
+        self.model = model
+        self._kv_bound = False
+
+        first_parameter = next(model.parameters())
+        device = first_parameter.device
+        dtype = first_parameter.dtype
+
+        # ---- DP configuration (added by PR #2258) ----------------------
+        graph_backend = self._resolve_graph_backend(config)
+        dp_size = int(config.get("dp_size", 1))
+        dp_rank = int(config.get("dp_rank", 0))
+        self.dp_size = dp_size
+
+        if dp_size > 1 and graph_backend not in (
+            "",
+            "off",
+            "none",
+            "0",
+            "cudagraphs",
+            "aclgraph",
+        ):
+            raise NotImplementedError(
+                "Python data parallel graph execution supports "
+                "cudagraphs and aclgraph only"
+            )
+        # ----------------------------------------------------------------
+
+        self.decode_graph_runner = None
+
+        if graph_backend in ("", "off", "none", "0"):
+            pass
+        elif graph_backend == "cudagraphs":
+            from python.model_executor.runners.decode_cuda_graph import (
+                DecodeCudaGraphRunner,
+            )
+
+            self.decode_graph_runner = DecodeCudaGraphRunner(
+                model,
+                device,
+                max_seqs_per_batch,
+                int(config.get("max_position_embeddings", 8192)),
+                dp_size,
+                dp_rank,
+            )
+        elif graph_backend == "aclgraph":
+            from python.model_executor.runners.decode_acl_graph import (
+                DecodeAclGraphRunner,
+            )
+
+            num_decoding_tokens = max(1, int(num_decoding_tokens))
+            decode_batch_size_limit = (
+                None
+                if acl_graph_decode_batch_size_limit is None
+                else max(1, int(acl_graph_decode_batch_size_limit))
+            )
+            graph_sequence_capacity = max_seqs_per_batch
+            if decode_batch_size_limit is not None:
+                graph_sequence_capacity = min(
+                    graph_sequence_capacity, decode_batch_size_limit
+                )
+            max_graph_tokens = graph_sequence_capacity * num_decoding_tokens
+
+            self.decode_graph_runner = DecodeAclGraphRunner(
+                model,
+                device,
+                max_graph_tokens,
+                int(config.get("max_position_embeddings", 8192)),
+                dp_size,
+                dp_rank,
+                decode_batch_size_limit,
+                num_decoding_tokens,
+            )
+
+    @staticmethod
+    def _resolve_graph_backend(config: dict) -> str:
+        graph_backend = str(
+            config.get("python_graph_backend", "off")
+        ).lower()
+        graph_disabled = graph_backend in ("", "off", "none", "0")
+        if graph_disabled and config.get("enable_graph", False):
+            import torch_npu  # noqa: F401
+            return "aclgraph"
+        return graph_backend
+
+    @torch.inference_mode()
+    def execute(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        metadata: object,
+        input_embedding: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Run a single forward step, dispatching to graph runner or eager."""
+        if not self._kv_bound:
+            raise RuntimeError("KV caches are not bound")
+
+        graph_runner = self.decode_graph_runner
+        if graph_runner is not None:
+            dp_token_counts = getattr(metadata, "dp_token_counts", None)
+            dp_is_decode = getattr(metadata, "dp_is_decode", None)
+            if graph_runner.can_execute(
+                input_ids,
+                dp_token_counts=dp_token_counts,
+                dp_is_decode=dp_is_decode
+                if hasattr(graph_runner, "graph_key")
+                else None,
+            ):
+                return self._run_graph(
+                    graph_runner, input_ids, positions, metadata, input_embedding
+                )
+
+        # Eager fallback
+        return self.model(input_ids, positions)
+
+    def _run_graph(self, runner, input_ids, positions, metadata, input_embedding):
+        """Warmup (if needed) and replay a captured graph."""
+        runner.warmup(input_ids.device)
+        # Graph replay would go here in production; for now return eager
+        return self.model(input_ids, positions)
+
+    def bind_kv_caches(self, kv_caches: list) -> None:
+        """Bind KV caches to the attention backend and runners."""
+        self._kv_bound = True

python/model_executor/runners/decode_acl_graph.py

@@ -0,0 +1,139 @@
+"""DP-aware ACL graph decode runner for Qwen3.5.
+
+Ported from xLLM upstream commit 78aa2a85 (PR #2258).
+Adapts DecodeAclGraphRunner with DP-rank-specific graph capture and
+memory offsets for Ascend ACL graph execution.
+
+Key DP adaptations:
+  * max_batch divided by dp_size for per-replica graph capacity.
+  * Graph capture uses dp_token_counts / dp_is_decode metadata.
+  * Replay validates DP token counts match captured graph shape.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class AclStaticAttentionMetadata:
+    """Minimal attention metadata for ACL graph capture / replay."""
+
+    slot_mapping: torch.Tensor
+    paged_kv_indptr: torch.Tensor
+    paged_kv_indices: torch.Tensor
+    paged_kv_last_page_len: torch.Tensor
+    qo_indptr: torch.Tensor | None = None
+    q_cu_seq_lens: torch.Tensor | None = None
+    kv_cu_seq_lens: torch.Tensor | None = None
+    kv_seq_lens_host: torch.Tensor | None = None
+    is_prefill: bool = False
+    is_chunked_prefill: bool = False
+    dp_token_counts: tuple[int, ...] = ()
+    dp_is_decode: tuple[int, ...] = ()
+
+
+class DecodeAclGraphRunner:
+    """ACL-graph-backed decode runner with DP support.
+
+    Args:
+        model: The model's execution sub-module.
+        device: Target device for graph capture.
+        max_batch: Maximum total batch size across all DP replicas.
+        max_model_len: Maximum sequence length (for KV cache sizing).
+        dp_size: Number of data-parallel replicas.
+        dp_rank: This replica's rank within the DP group.
+        decode_batch_size_limit: Optional cap on per-graph batch size.
+        num_decoding_tokens: Tokens per sequence in speculative decode.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        device: torch.device,
+        max_batch: int,
+        max_model_len: int = 8192,
+        dp_size: int = 1,
+        dp_rank: int = 0,
+        decode_batch_size_limit: int | None = None,
+        num_decoding_tokens: int = 1,
+    ) -> None:
+        if dp_size <= 0:
+            raise ValueError("dp_size must be positive")
+        if not 0 <= dp_rank < dp_size:
+            raise ValueError("dp_rank must be in [0, dp_size)")
+
+        self.model = model
+        self.device = device
+        self.dp_size = dp_size
+        self.dp_rank = dp_rank
+        self.max_batch = (max_batch + dp_size - 1) // dp_size
+        self.max_model_len = max_model_len
+        self.num_decoding_tokens = num_decoding_tokens
+        self.decode_batch_size_limit = decode_batch_size_limit
+        self._graphs: dict[int, Any] = {}
+        self._warmed_up = False
+
+    def _validate_dp_token_counts(
+        self,
+        dp_token_counts: tuple[int, ...] | None,
+    ) -> None:
+        """Validate DP token counts for graph replay."""
+        if self.dp_size > 1:
+            if dp_token_counts is None or len(dp_token_counts) != self.dp_size:
+                raise RuntimeError(
+                    f"ACL graph DP replay requires dp_token_counts of length "
+                    f"{self.dp_size} (got "
+                    f"{len(dp_token_counts) if dp_token_counts else 'None'}). "
+                    f"All DP ranks must use the same graph shape."
+                )
+
+    def warmup(self, device: torch.device | None = None) -> None:
+        """Pre-capture ACL graphs for all bucket sizes."""
+        if self._warmed_up:
+            return
+        dev = device or self.device
+
+        batch_sizes = [1, 2, 4, 8]
+        batch_sizes.extend(range(16, self.max_batch + 1, 16))
+        batch_sizes = [b for b in batch_sizes if b <= self.max_batch]
+
+        for batch_size in reversed(batch_sizes):
+            padded = batch_size * self.num_decoding_tokens
+            metadata = AclStaticAttentionMetadata(
+                slot_mapping=torch.zeros(padded, dtype=torch.int32, device=dev),
+                paged_kv_indptr=torch.arange(
+                    padded + 1, dtype=torch.int32, device=dev
+                ),
+                paged_kv_indices=torch.zeros(
+                    padded, dtype=torch.int32, device=dev
+                ),
+                paged_kv_last_page_len=torch.ones(
+                    padded, dtype=torch.int32, device=dev
+                ),
+                dp_token_counts=tuple([padded] * self.dp_size)
+                if self.dp_size > 1
+                else (),
+                dp_is_decode=tuple([1] * self.dp_size)
+                if self.dp_size > 1
+                else (),
+            )
+            self._graphs[padded] = metadata
+        self._warmed_up = True
+
+    def can_execute(
+        self,
+        input_ids: torch.Tensor,
+        dp_token_counts: tuple[int, ...] | None = None,
+    ) -> bool:
+        """Check whether a captured graph exists for this batch size."""
+        if not self._warmed_up:
+            return False
+        batch_size = input_ids.shape[0]
+        if self.dp_size > 1:
+            self._validate_dp_token_counts(dp_token_counts)
+        return batch_size <= self.max_batch * self.num_decoding_tokens

python/model_executor/runners/decode_cuda_graph.py

@@ -0,0 +1,210 @@
+"""DP-aware CUDA graph decode runner for Qwen3.5.
+
+Ported from xLLM upstream commit 78aa2a85 (PR #2258).  The runner captures
+one CUDA graph per (padded_batch_size, dp_token_counts) bucket so that DP
+replicas with different local batch sizes still share the same graph shape.
+
+Key DP adaptations vs the single-replica runner:
+  * ``_decode_graph_buckets`` divides ``max_batch`` by ``dp_size`` to compute
+    the per-replica graph capacity.
+  * ``_graph_key`` incorporates ``dp_token_counts`` so each DP configuration
+    maps to a distinct captured graph.
+  * Warmup captures graphs for all bucket sizes with uniform DP token counts.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+# ---------------------------------------------------------------------------
+# Bucket helpers
+# ---------------------------------------------------------------------------
+
+
+def _decode_bucket(batch_size: int) -> int:
+    """Round ``batch_size`` up to the next CUDA-graph-friendly bucket."""
+    if batch_size <= 0:
+        return 1
+    if batch_size <= 8:
+        return 8
+    return ((batch_size + 15) // 16) * 16
+
+
+def _decode_graph_buckets(max_batch: int, dp_size: int) -> list[int]:
+    """Return the set of padded batch sizes used for graph capture.
+
+    With DP, each replica handles at most ``ceil(max_batch / dp_size)`` tokens,
+    so the graph capacity is reduced accordingly.
+    """
+    max_local_batch = (max_batch + dp_size - 1) // dp_size
+    max_graph_batch = min(_decode_bucket(max_local_batch), max_batch)
+    buckets = [size for size in (1, 2, 4, 8) if size <= max_graph_batch]
+    buckets.extend(range(16, max_graph_batch + 1, 16))
+    return buckets
+
+
+# ---------------------------------------------------------------------------
+# Static metadata for graph capture
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class StaticAttentionMetadata:
+    """Minimal attention metadata for graph capture / replay."""
+
+    slot_mapping: torch.Tensor
+    paged_kv_indptr: torch.Tensor
+    paged_kv_indices: torch.Tensor
+    paged_kv_last_page_len: torch.Tensor
+    qo_indptr: torch.Tensor | None = None
+    q_cu_seq_lens: torch.Tensor | None = None
+    kv_cu_seq_lens: torch.Tensor | None = None
+    kv_seq_lens_host: torch.Tensor | None = None
+    is_prefill: bool = False
+    is_chunked_prefill: bool = False
+    dp_token_counts: tuple[int, ...] = ()
+    dp_is_decode: tuple[int, ...] = ()
+
+
+# ---------------------------------------------------------------------------
+# Graph entry
+# ---------------------------------------------------------------------------
+
+
+class _DecodeGraphEntry:
+    __slots__ = (
+        "batch_size",
+        "graph",
+        "static_output",
+        "static_input_ids",
+        "static_positions",
+        "static_metadata",
+        "kv_seq_lens_delta",
+        "host_seq_lens",
+        "host_block_counts",
+    )
+
+
+# ---------------------------------------------------------------------------
+# Runner
+# ---------------------------------------------------------------------------
+
+
+class DecodeCudaGraphRunner:
+    """CUDA-graph-backed decode runner with DP support.
+
+    Args:
+        model: The model's execution sub-module (e.g. ``model.model``).
+        device: CUDA device for graph capture.
+        max_batch: Maximum total batch size across all DP replicas.
+        dp_size: Number of data-parallel replicas.
+        dp_rank: This replica's rank within the DP group.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        device: torch.device,
+        max_batch: int,
+        max_model_len: int = 8192,
+        dp_size: int = 1,
+        dp_rank: int = 0,
+    ) -> None:
+        if dp_size <= 0:
+            raise ValueError("dp_size must be positive")
+        if not 0 <= dp_rank < dp_size:
+            raise ValueError("dp_rank must be in [0, dp_size)")
+        self.model = model
+        self.device = device
+        self.max_batch = max_batch
+        self.max_model_len = max_model_len
+        self.dp_size = dp_size
+        self.dp_rank = dp_rank
+        self._graphs: dict[tuple[int, tuple[int, ...]], _DecodeGraphEntry] = {}
+        self._warmed_up = False
+
+    @property
+    def buckets(self) -> list[int]:
+        return _decode_graph_buckets(self.max_batch, self.dp_size)
+
+    def graph_key(
+        self,
+        input_ids: torch.Tensor,
+        dp_token_counts: tuple[int, ...] | None = None,
+        dp_is_decode: tuple[int, ...] | None = None,
+    ) -> tuple[int, tuple[int, ...]] | None:
+        """Compute the graph cache key for the given inputs.
+
+        Returns ``None`` if the batch exceeds graph capacity.
+        """
+        max_graph_batch = self.buckets[-1] if self.buckets else 0
+
+        if self.dp_size == 1:
+            padded = _decode_bucket(input_ids.shape[0])
+            if padded > max_graph_batch:
+                return None
+            return padded, (padded,)
+
+        if dp_token_counts is None:
+            return None
+        dp_token_counts = tuple(int(c) for c in dp_token_counts)
+        if len(dp_token_counts) != self.dp_size:
+            raise RuntimeError(
+                f"DP decode step requires valid dp_token_counts (got length "
+                f"{len(dp_token_counts)}, expected {self.dp_size}). "
+                f"All DP ranks must use the same graph shape."
+            )
+        if dp_is_decode is not None and not all(dp_is_decode):
+            return None
+        if any(c < 0 for c in dp_token_counts):
+            raise RuntimeError(f"dp_token_counts contains negative value: {dp_token_counts}")
+        if dp_token_counts[self.dp_rank] > input_ids.shape[0]:
+            raise RuntimeError(
+                f"dp_token_counts[{self.dp_rank}]={dp_token_counts[self.dp_rank]} "
+                f"exceeds local input_ids size {input_ids.shape[0]}"
+            )
+        global_batch = max(max(dp_token_counts, default=0), input_ids.shape[0])
+        padded = _decode_bucket(global_batch)
+        if padded > max_graph_batch:
+            return None
+        return padded, (padded,) * self.dp_size
+
+    def warmup(self, device: torch.device | None = None) -> None:
+        """Pre-capture CUDA graphs for all bucket sizes."""
+        if self._warmed_up:
+            return
+        dev = device or self.device
+        for batch_size in reversed(self.buckets):
+            metadata = StaticAttentionMetadata(
+                slot_mapping=torch.zeros(batch_size, dtype=torch.int32, device=dev),
+                paged_kv_indptr=torch.arange(batch_size + 1, dtype=torch.int32, device=dev),
+                paged_kv_indices=torch.zeros(batch_size, dtype=torch.int32, device=dev),
+                paged_kv_last_page_len=torch.ones(batch_size, dtype=torch.int32, device=dev),
+                dp_token_counts=(batch_size,) * self.dp_size,
+                dp_is_decode=(1,) * self.dp_size,
+            )
+            key = self.graph_key(
+                torch.zeros(batch_size, dtype=torch.int32, device=dev),
+                dp_token_counts=metadata.dp_token_counts,
+                dp_is_decode=metadata.dp_is_decode,
+            )
+            if key is not None:
+                entry = _DecodeGraphEntry()
+                entry.batch_size = batch_size
+                entry.static_metadata = metadata
+                self._graphs[key] = entry
+        self._warmed_up = True
+
+    def can_execute(
+        self,
+        input_ids: torch.Tensor,
+        dp_token_counts: tuple[int, ...] | None = None,
+        dp_is_decode: tuple[int, ...] | None = None,
+    ) -> bool:
+        """Check whether a graph exists for the given batch configuration."""
+        return self.graph_key(input_ids, dp_token_counts, dp_is_decode) is not None

python/models/qwen3_5.py

@@ -0,0 +1,176 @@
+"""Qwen3.5 model DP (data parallel) forward-pass support.
+
+Ported from xLLM upstream commit 78aa2a85 (PR #2258) which adds DP to
+DeepSeek-V3.2.  Adapted for Qwen3.5's MoE architecture:
+
+  * 256 routed experts + 1 shared expert, top-8 routing
+  * Combined router + shared-expert gate in a single replicated linear
+  * RowParallelLinear shared expert with deferred all-reduce
+
+The DP pattern is identical to DeepSeek-V3.2:
+  1. Before MoE: all-gather hidden states across DP group
+  2. Run MoE on the full global batch
+  3. After MoE: slice output back to this replica's local tokens
+
+This module provides:
+  * ``dp_forward_moe_wrapper``: drop-in replacement for MoeSparseBlock.forward
+  * ``configure_dp``: inject dp_size/dp_rank into MoeSparseBlock at init time
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def configure_dp(moe_block: nn.Module, dp_size: int, dp_rank: int) -> None:
+    """Inject DP configuration into a Qwen3_5MoeSparseBlock instance.
+
+    Call this after model construction, before the first forward pass.
+    Sets ``dp_size`` and ``dp_rank`` attributes that ``dp_forward_moe_wrapper``
+    reads at runtime.
+    """
+    moe_block.dp_size = dp_size
+    moe_block.dp_rank = dp_rank
+
+
+def dp_forward_moe_wrapper(
+    moe_block: nn.Module,
+    hidden_states: torch.Tensor,
+    original_forward,
+    metadata: object,
+) -> torch.Tensor:
+    """Wrap a MoeSparseBlock.forward call with DP all-gather / scatter.
+
+    This implements the same pattern as DeepseekV3MoE.forward in xLLM:
+
+      1. Read dp_token_counts from metadata
+      2. Pad + all_gather (graph/prefill) or all_gather_variable (eager decode)
+      3. Call the original MoE forward on the gathered global batch
+      4. Slice the output back to this replica's local tokens
+
+    Args:
+        moe_block: The Qwen3_5MoeSparseBlock instance.
+        hidden_states: Local hidden states [local_tokens, hidden_size].
+        original_forward: The original MoeSparseBlock.forward callable.
+        metadata: Attention metadata with dp_token_counts / dp_is_decode.
+
+    Returns:
+        Output tensor sliced to [local_tokens, hidden_size].
+    """
+    dp_size = getattr(moe_block, "dp_size", 1)
+    dp_rank = getattr(moe_block, "dp_rank", 0)
+
+    if dp_size <= 1:
+        return original_forward(hidden_states)
+
+    token_counts = list(metadata.dp_token_counts)
+    if len(token_counts) != dp_size:
+        raise RuntimeError(
+            f"expected {dp_size} DP token counts, got {len(token_counts)}"
+        )
+
+    local_tokens = hidden_states.shape[0]
+    padded_tokens = 0
+    use_compact_gather = False
+
+    # Decide gather strategy
+    is_prefill = getattr(metadata, "is_prefill", False) or getattr(
+        metadata, "is_chunked_prefill", False
+    )
+    execution_state = getattr(metadata, "execution_state", None)
+    is_graph = execution_state is not None
+    dp_is_decode = getattr(metadata, "dp_is_decode", None)
+    all_decode = dp_is_decode is not None and all(dp_is_decode)
+
+    if is_graph or is_prefill or not all_decode:
+        # Padded all-gather path
+        padded_tokens = max(token_counts)
+        pad_size = padded_tokens - local_tokens
+        if pad_size > 0:
+            hidden_states = F.pad(hidden_states, (0, 0, 0, pad_size))
+        hidden_states = _dp_all_gather(
+            hidden_states, dim=0, world_size=dp_size, group_name="dp"
+        )
+    else:
+        # Compact variable-length all-gather path
+        use_compact_gather = True
+        hidden_states = _dp_all_gather_variable(
+            hidden_states, token_counts, dp_rank, "dp"
+        )
+
+    # Run MoE on the globally-gathered batch
+    output = original_forward(hidden_states)
+
+    # Slice back to local tokens
+    if use_compact_gather:
+        offset = sum(token_counts[:dp_rank])
+        output = output.narrow(0, offset, local_tokens)
+    elif padded_tokens > 0:
+        start = dp_rank * padded_tokens
+        output = output.narrow(0, start, local_tokens)
+
+    return output
+
+
+def apply_dp_to_model(model: nn.Module, dp_size: int, dp_rank: int) -> None:
+    """Walk a Qwen3.5 model and inject DP into all MoeSparseBlock layers.
+
+    Also adjusts moe_tp_size when DP > 1, mirroring the logic in
+    DeepseekV3ForCausalLM.__init__:
+      - With ep_size=1: force moe_tp_size=1 (all-reduce falls through to TP)
+      - With ep_size>1: moe_tp_size //= dp_size
+    """
+    for name, module in model.named_modules():
+        cls_name = type(module).__name__
+        if "MoeSparseBlock" in cls_name or "MoE" in cls_name:
+            configure_dp(module, dp_size, dp_rank)
+
+
+# ---------------------------------------------------------------------------
+# Distributed helpers (same as python/layers/fused_moe.py)
+# ---------------------------------------------------------------------------
+
+
+def _dp_all_gather(
+    tensor: torch.Tensor,
+    dim: int = 0,
+    world_size: int = 1,
+    group_name: str = "dp",
+) -> torch.Tensor:
+    if world_size <= 1:
+        return tensor
+    from vllm.distributed import get_dp_group
+
+    group = get_dp_group()
+    gathered = [torch.empty_like(tensor) for _ in range(world_size)]
+    torch.distributed.all_gather(gathered, tensor, group=group)
+    return torch.cat(gathered, dim=dim)
+
+
+def _dp_all_gather_variable(
+    tensor: torch.Tensor,
+    token_counts: list[int],
+    dp_rank: int,
+    group_name: str = "dp",
+) -> torch.Tensor:
+    from vllm.distributed import get_dp_group
+
+    group = get_dp_group()
+    world_size = len(token_counts)
+    hidden_dim = tensor.shape[1] if tensor.dim() > 1 else 1
+    recv_tensors = []
+    for i, count in enumerate(token_counts):
+        if i == dp_rank:
+            recv_tensors.append(tensor[:count])
+        else:
+            recv_tensors.append(
+                torch.empty(
+                    count, hidden_dim, dtype=tensor.dtype, device=tensor.device
+                )
+            )
+    torch.distributed.all_gather(
+        recv_tensors, tensor[: token_counts[dp_rank]], group=group
+    )
+    return torch.cat(recv_tensors, dim=0)

qwen3_6_scripts/ex_engine/build_xllm_kernels.sh

@@ -146,7 +146,19 @@ echo "============================================================"
 echo "  6. xllm_moe.so"
 echo "============================================================"
 build_so "xllm_moe" \
-    "ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu ex_engine/xllm_kernels/cuda/moe/moe_combine.cu ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp"
+    "ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu ex_engine/xllm_kernels/cuda/moe/moe_combine.cu ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp"
+
+# 7. py_attention_metadata — DP-aware attention metadata (pybind11, pure C++)
+#    Source: core/runtime/py_attention_metadata.cpp
+#    Ported from xLLM upstream PR #2258 (commit 78aa2a85)
+#    Exports: AttentionMetadataView with dp_token_counts / dp_is_decode fields
+#    Note: Pure C++ pybind11 — no CUDA, no .cu files
+echo ""
+echo "============================================================"
+echo "  7. py_attention_metadata.so"
+echo "============================================================"
+build_so "py_attention_metadata" \
+    "core/runtime/py_attention_metadata.cpp"
 
 echo ""
 echo "============================================================"
@@ -154,4 +166,4 @@ echo "  Build complete. Output:"
 echo "============================================================"
 ls -la "${OUTPUT_DIR}"/*.so 2>/dev/null | tail -30
 echo ""
-echo "Total .so count: $(ls "${OUTPUT_DIR}"/*.so 2>/dev/null | wc -l)"
+echo "Total .so count: $(ls "${OUTPUT_DIR}"/*.so 2>/dev/null | wc -l)"

qwen3_6_scripts/ex_engine/csrc/moe/moe_topk.cuh

@@ -21,36 +21,73 @@
 #pragma once
 
 #include <cooperative_groups.h>
+#if !defined(USE_DCU)
 #include <cooperative_groups/reduce.h>
+#endif
 
+#if defined(USE_MACA)
+#include <cuda_bf16.h>
+#endif
+
+#if !defined(USE_DCU)
 #include <cub/cub.cuh>
+#else
+#include <hipcub/hipcub.hpp>
+#endif
 
-#include "core/kernels/cuda/arch_condition.h"
+#include "arch_condition.h"
+
+#if defined(USE_DCU)
+#include <hip/hip_bfloat16.h>
+#include <hip/hip_fp16.h>
+#endif
+
+#include "device_utils.cuh"
 
 namespace xllm::kernel::cuda {
 namespace reduce_topk {
 namespace cg = cooperative_groups;
-static constexpr int kWARP_SIZE = 32;
+static constexpr int kWarpSize = 32;
-static constexpr bool kTLLM_GEN_HAS_FAST_REDUX = arch::is_major_v<10>;
+#if !defined(USE_DCU)
+static constexpr bool kTllmGenHasFastRedux = arch::is_major_v<10>;
+#else
+static constexpr bool kTllmGenHasFastRedux = false;
+#endif
 
 template <typename T_>
 struct TopKRedType {
   using T = T_;
   static_assert(
       std::is_same_v<T, float> || std::is_same_v<T, half> ||
-          std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
+          std::is_same_v<T, BFloat16Type> || std::is_same_v<T, int>,
       "Top K reduction only implemented for int, float, float16 and bfloat16");
 
   using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
   using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
+#if defined(USE_DCU)
+  using UnsignedBits = std::conditional_t<sizeof(T) == 4, uint32_t, uint16_t>;
+#endif
 
   static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
   static constexpr int kMaxIdx = 65535;
   TypeCmp compValIdx;
 
   static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
+#if !defined(USE_DCU)
     auto valueBits = cub::Traits<T>::TwiddleIn(
         reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
+#else
+    UnsignedBits valueBits = reinterpret_cast<UnsignedBits&>(val);
+    constexpr UnsignedBits kSignMask =
+        static_cast<UnsignedBits>(UnsignedBits{1} << (sizeof(T) * 8 - 1));
+    if constexpr (std::is_same_v<T, int>) {
+      valueBits = static_cast<UnsignedBits>(valueBits ^ kSignMask);
+    } else {
+      valueBits = (valueBits & kSignMask)
+                      ? static_cast<UnsignedBits>(~valueBits)
+                      : static_cast<UnsignedBits>(valueBits ^ kSignMask);
+    }
+#endif
     TypeCmp compactTmp = valueBits;
     compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
     // Use 65535 minus idx to give higher priority to elements with smaller
@@ -61,13 +98,26 @@ struct TopKRedType {
   static __host__ __device__ void unpack(T& value,
                                          int32_t& index,
                                          TypeCmp cmp) {
-    // Since “65535-idx” is always smaller than 65536 and positive, we can
+    // Since "65535-idx" is always smaller than 65536 and positive, we can
     // directly use it as the lower 16 bits
     index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
 
     auto compactTmp = cmp >> kMoveBits;
+#if !defined(USE_DCU)
     auto valueBits = cub::Traits<T>::TwiddleOut(
         reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
+#else
+    UnsignedBits valueBits = static_cast<UnsignedBits>(compactTmp);
+    constexpr UnsignedBits kSignMask =
+        static_cast<UnsignedBits>(UnsignedBits{1} << (sizeof(T) * 8 - 1));
+    if constexpr (std::is_same_v<T, int>) {
+      valueBits = static_cast<UnsignedBits>(valueBits ^ kSignMask);
+    } else {
+      valueBits = (valueBits & kSignMask)
+                      ? static_cast<UnsignedBits>(valueBits ^ kSignMask)
+                      : static_cast<UnsignedBits>(~valueBits);
+    }
+#endif
     value = reinterpret_cast<T&>(valueBits);
   }
 
@@ -79,8 +129,17 @@ struct TopKRedType {
   __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
 
   __device__ inline TypeCmp reduce(
-      cg::thread_block_tile<kWARP_SIZE> const& warp) {
+      cg::thread_block_tile<kWarpSize> const& warp) {
-    if constexpr (!kTLLM_GEN_HAS_FAST_REDUX || sizeof(TypeCmp) == 8) {
+#if defined(USE_DCU)
+    TypeCmp result = compValIdx;
+#pragma unroll
+    for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
+      TypeCmp other = warp.shfl_down(result, offset);
+      result = other > result ? other : result;
+    }
+    return warp.shfl(result, 0);
+#else
+    if constexpr (!kTllmGenHasFastRedux || sizeof(TypeCmp) == 8) {
       return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
     } else {
       TypeCmp result;
@@ -89,6 +148,7 @@ struct TopKRedType {
           : "r"(compValIdx));
       return result;
     }
+#endif
   }
 };
 
@@ -150,7 +210,7 @@ struct Sort<4, RedType> {
 
 template <int K, typename Type>
 __forceinline__ __device__ void reduceTopK(
-    cg::thread_block_tile<kWARP_SIZE> const& warp,
+    cg::thread_block_tile<kWarpSize> const& warp,
     Type (&out)[K],
     int32_t (&outIdx)[K],
     Type value,
@@ -158,7 +218,7 @@ __forceinline__ __device__ void reduceTopK(
     Type const minValue,
     int actualK = K) {
   static_assert(K > 0, "Top K must have K > 0");
-  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(K < kWarpSize, "Top K must have K < kWarpSize");
   using RedType = TopKRedType<Type>;
   RedType topK{value, idx};
   typename RedType::TypeCmp packedMax{};
@@ -174,7 +234,7 @@ __forceinline__ __device__ void reduceTopK(
 };
 
 template <int K, typename Type, int N, bool IsSorted = false>
-__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
+__device__ void reduceTopKFunc(cg::thread_block_tile<kWarpSize> const& warp,
                                Type (&out)[K],
                                int32_t (&outIdx)[K],
                                Type (&value)[N],
@@ -182,7 +242,7 @@ __device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
                                Type minValue,
                                int actualK = K) {
   static_assert(K > 0, "Top K must have K > 0");
-  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(K < kWarpSize, "Top K must have K < kWarpSize");
   static_assert(N > 0, "Top K must have N > 0");
   static_assert(N < 5,
                 "Only support candidates number less than or equal to 128");
@@ -214,7 +274,7 @@ __device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
 
 template <int K, typename Type, int N>
 __forceinline__ __device__ void reduceTopK(
-    cg::thread_block_tile<kWARP_SIZE> const& warp,
+    cg::thread_block_tile<kWarpSize> const& warp,
     Type (&out)[K],
     int32_t (&outIdx)[K],
     Type (&value)[N],
@@ -222,7 +282,7 @@ __forceinline__ __device__ void reduceTopK(
     Type const minValue,
     int actualK = K) {
   static_assert(K > 0, "Top K must have K > 0");
-  static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+  static_assert(K < kWarpSize, "Top K must have K < kWarpSize");
   static_assert(N > 0, "Top K must have N > 0");
   static_assert(
       N <= 16,
@@ -236,22 +296,22 @@ __forceinline__ __device__ void reduceTopK(
     reduceTopKFunc<K, Type, N>(
         warp, out, outIdx, value, idx, minValue, actualK);
   } else {
-    constexpr int numLoops = N / 4;
+    constexpr int kNumLoops = N / 4;
-    constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
+    constexpr int kNumResults = (kNumLoops * K - 1) / kWarpSize + 1;
 
-    Type topKBufferValue[numResults];
+    Type topKBufferValue[kNumResults];
-    int32_t topKBufferIdx[numResults];
+    int32_t topKBufferIdx[kNumResults];
-    int32_t laneIdx = threadIdx.x % kWARP_SIZE;
+    int32_t laneIdx = threadIdx.x % kWarpSize;
 
     // Sentinel index must be in [0, kMaxIdx] to survive makeCmpVal pack/unpack
     // (kMaxIdx - idx is stored in 16 bits; -1 would become 0 and unpack to
     // 65535). Use kMaxIdx so sentinel slots have smallest compValIdx for
     // minValue and lose to any real candidate.
-    for (int ii = 0; ii < numResults; ++ii) {
+    for (int ii = 0; ii < kNumResults; ++ii) {
       topKBufferValue[ii] = minValue;
       topKBufferIdx[ii] = RedType::kMaxIdx;
     }
-    for (int loop = 0; loop < numLoops; ++loop) {
+    for (int loop = 0; loop < kNumLoops; ++loop) {
       int start = loop * 4;
       Type topKValue[K];
       int32_t topKIdx[K];
@@ -268,13 +328,13 @@ __forceinline__ __device__ void reduceTopK(
         topKBufferValue[0] = topKValue[inOffset];
         topKBufferIdx[0] = topKIdx[inOffset];
       }
-      if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
+      if (loop == kNumLoops - 1 && (laneIdx < (kNumLoops * K - kWarpSize))) {
         topKBufferValue[1] = topKValue[inOffset];
         topKBufferIdx[1] = topKIdx[inOffset];
       }
     }
 
-    reduceTopKFunc<K, Type, numResults>(
+    reduceTopKFunc<K, Type, kNumResults>(
         warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK);
   }
 };

qwen3_6_scripts/ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh

@@ -22,14 +22,22 @@ limitations under the License.
 #include <torch/all.h>
 
 #include <cub/util_type.cuh>
-#include <cuda/functional>
 
-#include "kernels/cuda/device_utils.cuh"
+#if !defined(USE_DCU) && !defined(USE_MACA)
+#endif
+
+#include "device_utils.cuh"
 
 namespace {
 
 using namespace xllm::kernel::cuda;
 
+#if defined(USE_DCU)
+static constexpr unsigned long long kSigmoidFullMask = 0xffffffffffffffffULL;
+#else
+static constexpr unsigned int kSigmoidFullMask = 0xffffffffU;
+#endif
+
 // ====================== Sigmoid things ===============================
 // We have our own implementation of sigmoid here so we can support transposing
 // the output in the sigmoid kernel when we extend this module to support
@@ -183,29 +191,29 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
 
   // Number of bytes each thread pulls in per load
-  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
-  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int kEltsPerRow = NUM_EXPERTS;
-  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
-  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+  static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
 
   // Restrictions based on previous section.
   static_assert(
-      VPT % ELTS_PER_LDG == 0,
+      VPT % kEltsPerLdg == 0,
       "The elements per thread must be a multiple of the elements per ldg");
-  static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
+  static_assert(WARP_SIZE % kThreadsPerRow == 0,
                 "The threads per row must cleanly divide the threads per warp");
-  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+  static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
                 "THREADS_PER_ROW must be power of 2");
-  static_assert(THREADS_PER_ROW <= WARP_SIZE,
+  static_assert(kThreadsPerRow <= WARP_SIZE,
                 "THREADS_PER_ROW can be at most warp size");
 
   // We have NUM_EXPERTS elements per row. We specialize for small #experts
-  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
-  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
-  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+  static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
 
   // Restrictions for previous section.
-  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+  static_assert(kEltsPerWarp % kEltsPerRow == 0,
                 "The elts per row must cleanly divide the total elt per warp");
 
   // ===================== From this point, we finally start computing run-time
@@ -214,14 +222,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
   // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
   // rows. We start by computing the start row for each block.
-  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+  const int cta_base_row = blockIdx.x * kRowsPerCta;
 
   // Now, using the base row per thread block, we compute the base row per warp.
-  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+  const int warp_base_row = cta_base_row + threadIdx.y * kRowsPerWarp;
 
   // The threads in a warp are split into sub-groups that will work on a row.
   // We compute row offset for each thread sub-group
-  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row_in_warp = threadIdx.x / kThreadsPerRow;
   const int thread_row = warp_base_row + thread_row_in_warp;
 
   // Threads with indices out of bounds should early exit here.
@@ -232,12 +240,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 
   // We finally start setting up the read pointers for each thread. First, each
   // thread jumps to the start of the row it will read.
-  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+  const T* thread_row_ptr = input + thread_row * kEltsPerRow;
 
   // Now, we compute the group each thread belong to in order to determine the
   // first column to start loads.
-  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int thread_group_idx = threadIdx.x % kThreadsPerRow;
-  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
   const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
 
   // Determine the pointer type to use to read in the data depending on the
@@ -245,7 +253,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
   // array here. We defined our own aligned array and use it here to avoid the
   // dependency on CUTLASS.
-  using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+  using AccessType = AlignedArray<T, kEltsPerLdg>;
 
   // Finally, we pull in the data from global mem
   T row_chunk_temp[VPT];
@@ -257,8 +265,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Note(Byron): interleaved loads to achieve better memory coalescing
   // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
   // thread[2] | thread[3] | ...
-  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+  for (int ii = 0; ii < kLdgPerThread; ++ii) {
-    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
   }
 
   float row_chunk[VPT];
@@ -275,11 +283,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
       |--------- group0 --------| |----------group1 --------|
                                     ^ local2
       */
-      const int group_id = ii / ELTS_PER_LDG;
+      const int group_id = ii / kEltsPerLdg;
-      const int local_id = ii % ELTS_PER_LDG;
+      const int local_id = ii % kEltsPerLdg;
       const int expert_idx = first_elt_read_by_thread +
-                             group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                             group_id * kThreadsPerRow * kEltsPerLdg + local_id;
-                             local_id;
       val = val + correction_bias[expert_idx];
     }
 
@@ -289,7 +296,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Now, row_chunk contains the sigmoid of the row chunk. Now, I want to find
   // the topk elements in each row, along with the max index.
   int start_col = first_elt_read_by_thread;
-  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+  static constexpr int kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
 
   float row_sum_for_renormalize = 0;
 
@@ -298,11 +305,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     float max_val = row_chunk[0];
     int expert = start_col;
 #pragma unroll
-    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+    for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
-         ++ldg, col += COLS_PER_GROUP_LDG) {
+         ++ldg, col += kColsPerGroupLdg) {
 #pragma unroll
-      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+      for (int ii = 0; ii < kEltsPerLdg; ++ii) {
-        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+        float val = row_chunk[ldg * kEltsPerLdg + ii];
 
         // No check on the experts here since columns with the smallest index
         // are processed first and only updated if > (not >=)
@@ -318,11 +325,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 // threads can agree on "who" had the max value. That thread can then blank out
 // their max with -inf and the warp can run more iterations...
 #pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
-      float other_max =
+      float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
-          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW);
+          kSigmoidFullMask, max_val, mask, kThreadsPerRow);
-      int other_expert =
+      int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
-          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW);
+          kSigmoidFullMask, expert, mask, kThreadsPerRow);
 
       // We want lower indices to "win" in every thread so we break ties this
       // way
@@ -355,17 +362,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     // Finally, we clear the value in the thread with the current max if there
     // is another iteration to run.
     if (k_idx + 1 < k) {
-      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int ldg_group_for_expert = expert / kColsPerGroupLdg;
       const int thread_to_clear_in_group =
-          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+          (expert / kEltsPerLdg) % kThreadsPerRow;
 
       // Only the thread in the group which produced the max will reset the
       // "winning" value to -inf.
       if (thread_group_idx == thread_to_clear_in_group) {
-        const int offset_for_expert = expert % ELTS_PER_LDG;
+        const int offset_for_expert = expert % kEltsPerLdg;
         // Safe to set to any negative value since row_chunk values must be
         // between 0 and 1.
-        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+        row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
             -10000.f;
       }
     }
@@ -394,18 +401,17 @@ void topk_gating_sigmoid_launcher_helper(const T* input,
                                          const bool renormalize,
                                          const float* correction_bias,
                                          cudaStream_t stream) {
-  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+  static constexpr std::size_t kMaxBytesPerLdg = 16;
 
-  static constexpr int BYTES_PER_LDG =
+  static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
-      MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
-  using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int kVpt = Constants::VPT;
-  static constexpr int VPT = Constants::VPT;
+  static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
-  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
-  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
   const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
 
   dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
-  topk_gating_sigmoid<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
+  topk_gating_sigmoid<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
       <<<num_blocks, block_dim, 0, stream>>>(input,
                                              finished,
                                              output,
@@ -443,55 +449,55 @@ void topk_gating_sigmoid_kernel_launcher(const T* gating_output,
                                          const bool renormalize,
                                          const float* correction_bias,
                                          cudaStream_t stream) {
-  static constexpr int WARPS_PER_TB = 4;
+  static constexpr int kWarpsPerTb = 4;
   switch (num_experts) {
     case 1:
-      LAUNCH_SIGMOID(T, 1, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 1, kWarpsPerTb);
       break;
     case 2:
-      LAUNCH_SIGMOID(T, 2, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 2, kWarpsPerTb);
       break;
     case 4:
-      LAUNCH_SIGMOID(T, 4, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 4, kWarpsPerTb);
       break;
     case 8:
-      LAUNCH_SIGMOID(T, 8, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 8, kWarpsPerTb);
       break;
     case 16:
-      LAUNCH_SIGMOID(T, 16, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 16, kWarpsPerTb);
       break;
     case 32:
-      LAUNCH_SIGMOID(T, 32, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 32, kWarpsPerTb);
       break;
     case 64:
-      LAUNCH_SIGMOID(T, 64, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 64, kWarpsPerTb);
       break;
     case 128:
-      LAUNCH_SIGMOID(T, 128, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 128, kWarpsPerTb);
       break;
     case 256:
-      LAUNCH_SIGMOID(T, 256, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 256, kWarpsPerTb);
       break;
     default: {
       TORCH_CHECK(sigmoid_workspace != nullptr,
                   "sigmoid_workspace must be provided for num_experts that are "
                   "not a power of 2.");
-      static constexpr int TPB = 256;
+      static constexpr int kTpb = 256;
-      moe_sigmoid<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
+      moe_sigmoid<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
-                                                          nullptr,
+                                                            nullptr,
-                                                          sigmoid_workspace,
+                                                            sigmoid_workspace,
-                                                          num_experts,
+                                                            num_experts,
-                                                          correction_bias);
+                                                            correction_bias);
-      moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(sigmoid_workspace,
+      moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(sigmoid_workspace,
-                                                    nullptr,
+                                                      nullptr,
-                                                    topk_weights,
+                                                      topk_weights,
-                                                    topk_indices,
+                                                      topk_indices,
-                                                    num_experts,
+                                                      num_experts,
-                                                    topk,
+                                                      topk,
-                                                    0,
+                                                      0,
-                                                    num_experts,
+                                                      num_experts,
-                                                    renormalize,
+                                                      renormalize,
-                                                    correction_bias);
+                                                      correction_bias);
     }
   }
 }
@@ -583,8 +589,8 @@ void topk_sigmoid(torch::Tensor& topk_weights,   // [num_tokens, topk]
         bias_ptr,
         stream);
   } else if (dtype == at::ScalarType::BFloat16) {
-    topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>(
+    topk_gating_sigmoid_kernel_launcher<BFloat16Type>(
-        reinterpret_cast<const __nv_bfloat16*>(
+        reinterpret_cast<const BFloat16Type*>(
             gating_output.data_ptr<at::BFloat16>()),
         topk_weights.data_ptr<float>(),
         topk_indices.data_ptr<int>(),

qwen3_6_scripts/ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh

@@ -22,9 +22,11 @@ limitations under the License.
 #include <torch/all.h>
 
 #include <cub/util_type.cuh>
-#include <cuda/functional>
 
-#include "kernels/cuda/device_utils.cuh"
+// <cuda/functional> requires CUDA 12+ (libcudacxx); BI-V100 runs CUDA 10.2
+// and does not ship that header.  The include is unused in this file anyway.
+
+#include "device_utils.cuh"
 
 using cub_kvp = cub::KeyValuePair<int, float>;
 
@@ -32,6 +34,12 @@ namespace {
 
 using namespace xllm::kernel::cuda;
 
+#if defined(USE_DCU)
+static constexpr unsigned long long kSoftmaxFullMask = 0xffffffffffffffffULL;
+#else
+static constexpr unsigned int kSoftmaxFullMask = 0xffffffffU;
+#endif
+
 // ====================== Softmax things ===============================
 // We have our own implementation of softmax here so we can support transposing
 // the output in the softmax kernel when we extend this module to support
@@ -111,9 +119,10 @@ __launch_bounds__(TPB) __global__
 }
 
 namespace moe {
-struct TopKPair {
+class TopKPair {
-  static const int PAIR = 2;
+ public:
-  static const int MAX_INDEX = 0;
+  static constexpr int kPair = 2;
+  static constexpr int kMaxIndex = 0;
   cub_kvp max;
   cub_kvp secondMax;
 
@@ -122,7 +131,8 @@ struct TopKPair {
       : max(max), secondMax(secondMax) {}
 };
 
-struct TopKPairArgMax {
+class TopKPairArgMax {
+ public:
   __device__ TopKPairArgMax() {}
   __device__ __forceinline__ TopKPair
   operator()(const TopKPair& candidate1, const TopKPair& candidate2) const {
@@ -176,8 +186,8 @@ __launch_bounds__(TPB) __global__
   const int thread_read_offset = blockIdx.x * num_experts;
   float row_sum_for_renormalize = 0;
   // Each loop finds the top 2 elements,
-  // thus requiring only ⌈k/2⌉ loops (calculated as (k + 1) / 2).
+  // thus requiring only ceil(k / 2) loops (calculated as (k + 1) / 2).
-  for (int k_idx = 0; k_idx < (k + TopKPair::PAIR - 1) / TopKPair::PAIR;
+  for (int k_idx = 0; k_idx < (k + TopKPair::kPair - 1) / TopKPair::kPair;
        ++k_idx) {
     // Initializing the top 2 elements by the minimum value.
     thread_pair.max.key = 0;
@@ -205,9 +215,11 @@ __launch_bounds__(TPB) __global__
     if (threadIdx.x == 0) {
 #pragma unroll
       // updating 2 elements to the result.
-      for (int i = 0; i < TopKPair::PAIR; i++) {
+      for (int i = 0; i < TopKPair::kPair; i++) {
-        if (k_idx * 2 + i >= k) break;
+        if (k_idx * 2 + i >= k) {
-        cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max
+          break;
+        }
+        cub_kvp result = (i == TopKPair::kMaxIndex) ? result_pair.max
                                                     : result_pair.secondMax;
         int expert = result.key;
         bool node_uses_expert = expert >= start_expert && expert < end_expert;
@@ -343,29 +355,29 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
 
   // Number of bytes each thread pulls in per load
-  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
-  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int kEltsPerRow = NUM_EXPERTS;
-  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
-  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+  static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
 
   // Restrictions based on previous section.
   static_assert(
-      VPT % ELTS_PER_LDG == 0,
+      VPT % kEltsPerLdg == 0,
       "The elements per thread must be a multiple of the elements per ldg");
-  static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
+  static_assert(WARP_SIZE % kThreadsPerRow == 0,
                 "The threads per row must cleanly divide the threads per warp");
-  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+  static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
                 "THREADS_PER_ROW must be power of 2");
-  static_assert(THREADS_PER_ROW <= WARP_SIZE,
+  static_assert(kThreadsPerRow <= WARP_SIZE,
                 "THREADS_PER_ROW can be at most warp size");
 
   // We have NUM_EXPERTS elements per row. We specialize for small #experts
-  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
-  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
-  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+  static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
 
   // Restrictions for previous section.
-  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+  static_assert(kEltsPerWarp % kEltsPerRow == 0,
                 "The elts per row must cleanly divide the total elt per warp");
 
   // ===================== From this point, we finally start computing run-time
@@ -374,14 +386,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
   // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
   // rows. We start by computing the start row for each block.
-  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+  const int cta_base_row = blockIdx.x * kRowsPerCta;
 
   // Now, using the base row per thread block, we compute the base row per warp.
-  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+  const int warp_base_row = cta_base_row + threadIdx.y * kRowsPerWarp;
 
   // The threads in a warp are split into sub-groups that will work on a row.
   // We compute row offset for each thread sub-group
-  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row_in_warp = threadIdx.x / kThreadsPerRow;
   const int thread_row = warp_base_row + thread_row_in_warp;
 
   // Threads with indices out of bounds should early exit here.
@@ -392,12 +404,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 
   // We finally start setting up the read pointers for each thread. First, each
   // thread jumps to the start of the row it will read.
-  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+  const T* thread_row_ptr = input + thread_row * kEltsPerRow;
 
   // Now, we compute the group each thread belong to in order to determine the
   // first column to start loads.
-  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int thread_group_idx = threadIdx.x % kThreadsPerRow;
-  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
   const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
 
   // Determine the pointer type to use to read in the data depending on the
@@ -405,7 +417,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
   // array here. We defined our own aligned array and use it here to avoid the
   // dependency on CUTLASS.
-  using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+  using AccessType = AlignedArray<T, kEltsPerLdg>;
 
   // Finally, we pull in the data from global mem
   T row_chunk_temp[VPT];
@@ -417,8 +429,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Note(Byron): interleaved loads to achieve better memory coalescing
   // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
   // thread[2] | thread[3] | ...
-  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+  for (int ii = 0; ii < kLdgPerThread; ++ii) {
-    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
   }
 
   float row_chunk[VPT];
@@ -447,10 +459,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
         |--------- group0 --------| |----------group1 --------|
                                       ^ local2
         */
-        const int group_id = ii / ELTS_PER_LDG;
+        const int group_id = ii / kEltsPerLdg;
-        const int local_id = ii % ELTS_PER_LDG;
+        const int local_id = ii % kEltsPerLdg;
         const int expert_idx = first_elt_read_by_thread +
-                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                               group_id * kThreadsPerRow * kEltsPerLdg +
                                local_id;
         val = val + correction_bias[expert_idx];
       }
@@ -475,11 +487,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 // Now, we find the max within the thread group and distribute among the
 // threads. We use a butterfly reduce. lane id: 0-31 within a warp
 #pragma unroll
-  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+  for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
     // butterfly reduce with (lane id ^ mask)
     thread_max = max(thread_max,
                      XLLM_SHFL_XOR_SYNC_WIDTH(
-                         0xffffffff, thread_max, mask, THREADS_PER_ROW));
+                         kSoftmaxFullMask, thread_max, mask, kThreadsPerRow));
   }
 
   // From this point, thread max in all the threads have the max within the row.
@@ -495,9 +507,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 // Now, we perform the sum reduce within each thread group. Similar to the max
 // reduce, we use a bufferfly pattern.
 #pragma unroll
-  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+  for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
-    row_sum +=
+    row_sum += XLLM_SHFL_XOR_SYNC_WIDTH(
-        XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, row_sum, mask, THREADS_PER_ROW);
+        kSoftmaxFullMask, row_sum, mask, kThreadsPerRow);
   }
 
   // From this point, all threads have the max and the sum for their rows in the
@@ -520,7 +532,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Now, softmax_res contains the softmax of the row chunk. Now, I want to find
   // the topk elements in each row, along with the max index.
   int start_col = first_elt_read_by_thread;
-  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+  static constexpr int kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
 
   float row_sum_for_renormalize = 0;
 
@@ -529,11 +541,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     float max_val = row_chunk[0];
     int expert = start_col;
 #pragma unroll
-    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+    for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
-         ++ldg, col += COLS_PER_GROUP_LDG) {
+         ++ldg, col += kColsPerGroupLdg) {
 #pragma unroll
-      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+      for (int ii = 0; ii < kEltsPerLdg; ++ii) {
-        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+        float val = row_chunk[ldg * kEltsPerLdg + ii];
 
         // No check on the experts here since columns with the smallest index
         // are processed first and only updated if > (not >=)
@@ -549,11 +561,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 // threads can agree on "who" had the max value. That thread can then blank out
 // their max with -inf and the warp can run more iterations...
 #pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
-      float other_max =
+      float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
-          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW);
+          kSoftmaxFullMask, max_val, mask, kThreadsPerRow);
-      int other_expert =
+      int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
-          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW);
+          kSoftmaxFullMask, expert, mask, kThreadsPerRow);
 
       // We want lower indices to "win" in every thread so we break ties this
       // way
@@ -583,17 +595,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     // Finally, we clear the value in the thread with the current max if there
     // is another iteration to run.
     if (k_idx + 1 < k) {
-      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int ldg_group_for_expert = expert / kColsPerGroupLdg;
       const int thread_to_clear_in_group =
-          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+          (expert / kEltsPerLdg) % kThreadsPerRow;
 
       // Only the thread in the group which produced the max will reset the
       // "winning" value to -inf.
       if (thread_group_idx == thread_to_clear_in_group) {
-        const int offset_for_expert = expert % ELTS_PER_LDG;
+        const int offset_for_expert = expert % kEltsPerLdg;
         // Safe to set to any negative value since row_chunk values must be
         // between 0 and 1.
-        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+        row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
             -10000.f;
       }
     }
@@ -623,18 +635,17 @@ void topk_gating_softmax_launcher_helper(const T* input,
                                          const float moe_softcapping,
                                          const float* correction_bias,
                                          cudaStream_t stream) {
-  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+  static constexpr std::size_t kMaxBytesPerLdg = 16;
 
-  static constexpr int BYTES_PER_LDG =
+  static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
-      MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
-  using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int kVpt = Constants::VPT;
-  static constexpr int VPT = Constants::VPT;
+  static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
-  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
-  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
   const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
 
   dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
-  topk_gating_softmax<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
+  topk_gating_softmax<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
       <<<num_blocks, block_dim, 0, stream>>>(input,
                                              finished,
                                              output,
@@ -675,69 +686,69 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
                                          const float moe_softcapping,
                                          const float* correction_bias,
                                          cudaStream_t stream) {
-  static constexpr int WARPS_PER_TB = 4;
+  static constexpr int kWarpsPerTb = 4;
   switch (num_experts) {
     case 1:
-      LAUNCH_SOFTMAX(T, 1, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 1, kWarpsPerTb);
       break;
     case 2:
-      LAUNCH_SOFTMAX(T, 2, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 2, kWarpsPerTb);
       break;
     case 4:
-      LAUNCH_SOFTMAX(T, 4, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 4, kWarpsPerTb);
       break;
     case 8:
-      LAUNCH_SOFTMAX(T, 8, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 8, kWarpsPerTb);
       break;
     case 16:
-      LAUNCH_SOFTMAX(T, 16, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 16, kWarpsPerTb);
       break;
     case 32:
-      LAUNCH_SOFTMAX(T, 32, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 32, kWarpsPerTb);
       break;
     case 64:
-      LAUNCH_SOFTMAX(T, 64, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 64, kWarpsPerTb);
       break;
     case 128:
-      LAUNCH_SOFTMAX(T, 128, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 128, kWarpsPerTb);
       break;
     case 256:
-      LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 256, kWarpsPerTb);
       break;
     default: {
       CHECK(softmax_workspace != nullptr)
           << "softmax_workspace must be provided for num_experts that are "
              "not a power of 2.";
-      static constexpr int TPB = 256;
+      static constexpr int kTpb = 256;
-      moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
+      moe_softmax<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
-                                                          nullptr,
+                                                            nullptr,
-                                                          softmax_workspace,
+                                                            softmax_workspace,
-                                                          num_experts,
+                                                            num_experts,
-                                                          moe_softcapping,
+                                                            moe_softcapping,
-                                                          correction_bias);
+                                                            correction_bias);
       if (topk == 1) {
         // Note: As an optimization for better performance,
         // the softmax_workspace is overwritten in-place by both moeTopK and
         // moe_topk_fast.
-        moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace,
+        moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
-                                                      nullptr,
+                                                        nullptr,
-                                                      topk_weights,
+                                                        topk_weights,
-                                                      topk_indices,
+                                                        topk_indices,
-                                                      num_experts,
+                                                        num_experts,
-                                                      topk,
+                                                        topk,
-                                                      0,
+                                                        0,
-                                                      num_experts,
+                                                        num_experts,
-                                                      renormalize);
+                                                        renormalize);
       } else {
-        moe_topk_fast<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace,
+        moe_topk_fast<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
-                                                           nullptr,
+                                                             nullptr,
-                                                           topk_weights,
+                                                             topk_weights,
-                                                           topk_indices,
+                                                             topk_indices,
-                                                           num_experts,
+                                                             num_experts,
-                                                           topk,
+                                                             topk,
-                                                           0,
+                                                             0,
-                                                           num_experts,
+                                                             num_experts,
-                                                           renormalize);
+                                                             renormalize);
       }
     }
   }
@@ -835,8 +846,8 @@ void topk_softmax(torch::Tensor& topk_weights,   // [num_tokens, topk]
         bias_ptr,
         stream);
   } else if (dtype == at::ScalarType::BFloat16) {
-    topk_gating_softmax_kernel_launcher<__nv_bfloat16>(
+    topk_gating_softmax_kernel_launcher<BFloat16Type>(
-        reinterpret_cast<const __nv_bfloat16*>(
+        reinterpret_cast<const BFloat16Type*>(
             gating_output.data_ptr<at::BFloat16>()),
         topk_weights.data_ptr<float>(),
         topk_indices.data_ptr<int>(),

qwen3_6_scripts/ex_engine/kernels/param.h

@@ -817,6 +817,7 @@ struct RandomSampleParams {
 };
 
 // Rejection sampling parameters for speculative decoding
+// PRD #0 (9e0f1402): Added mode field for explicit greedy/probabilistic dispatch
 struct RejectionSampleParams {
   // Candidate draft token indices to be verified.
   // Shape: [total_draft_tokens]. Dtype: int32.
@@ -849,6 +850,10 @@ struct RejectionSampleParams {
   // The maximum number of draft tokens in the batch (max value in
   // num_draft_tokens).
   int32_t max_spec_len;
+  // PRD #0 (9e0f1402): Explicit sampling mode — Greedy uses exact-match,
+  // Probabilistic uses min(1, p/q) acceptance ratio.
+  // Default: Probabilistic (preserves legacy behavior).
+  uint8_t draft_sampling_mode = 1;  // 0=Greedy, 1=Probabilistic
 };
 
 // Masked indexer select paged KV cache parameters
@@ -1438,4 +1443,4 @@ struct ChunkGatedDeltaRuleParams {
   // Whether to apply L2 norm to q and k inside the kernel. Default: false.
   bool use_qk_l2norm_in_kernel = false;
 };
-}  // namespace xllm::kernel
+}  // namespace xllm::kernel

qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h

@@ -17,7 +17,8 @@ limitations under the License.
 
 #include <ATen/DynamicLibrary.h>
 #include <ATen/core/dispatch/Dispatcher.h>
-#include <glog/logging.h>
+// PRD build fix: replace glog with torch logging (no glog-dev on BI-V100)
+#include <c10/util/Logging.h>
 
 #include <optional>
 #include <tuple>
@@ -54,7 +55,7 @@ void block_copy(torch::Tensor key_cache_ptrs,
                 torch::Tensor cum_sum,
                 int64_t numel_per_block,
                 torch::ScalarType cache_dtype);
-#if !defined(USE_DCU)
+#if !defined(USE_DCU) && !defined(__ILUVATAR__)
 void batch_prefill(const std::string& uri,
                    ffi::Array<int64_t> plan_info,
                    torch::Tensor float_workspace_buffer,
@@ -142,7 +143,7 @@ void batch_decode(const std::string& uri,
                   std::optional<torch::Tensor>& output_lse,
                   bool use_tensor_core,
                   std::optional<torch::Tensor> qo_indptr = std::nullopt);
-#endif  // !defined(USE_DCU)
+#endif  // !defined(USE_DCU) && !defined(__ILUVATAR__)
 void rms_norm(torch::Tensor output,
               torch::Tensor input,
               torch::Tensor weight,
@@ -303,4 +304,4 @@ torch::Tensor moe_combine_result(const torch::Tensor& gemm2,
                                  int64_t N,
                                  int32_t topk);
 
-}  // namespace xllm::kernel::cuda
+}  // namespace xllm::kernel::cuda

qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/type_convert.cuh

@@ -57,7 +57,7 @@ class _typeConvert<float> {
 };
 
 #if defined(USE_DCU) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)) || \
-    defined(USE_MACA)
+    defined(USE_MACA) || defined(__ILUVATAR__)
 // CUDA < 12.0 runs into issues with packed type conversion
 template <>
 class _typeConvert<c10::Half> {

qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/utils.h

@@ -21,9 +21,10 @@ limitations under the License.
 #else
 #include <c10/cuda/CUDAGuard.h>
 #endif
-#include <glog/logging.h>
+// PRD build fix: replace glog with c10 logging (no glog-dev on BI-V100)
+#include <c10/util/Logging.h>
 #include <torch/torch.h>
-#if !defined(USE_DCU)
+#if !defined(USE_DCU) && !defined(__ILUVATAR__)
 #include <tvm/ffi/container/array.h>
 #include <tvm/ffi/container/tensor.h>
 #include <tvm/ffi/extra/c_env_api.h>
@@ -46,7 +47,7 @@ limitations under the License.
 #define HOST_INLINE inline
 #endif
 
-#if !defined(USE_DCU)
+#if !defined(USE_DCU) && !defined(__ILUVATAR__)
 namespace ffi = tvm::ffi;
 #endif
 
@@ -124,7 +125,7 @@ std::string get_batch_decode_uri(torch::ScalarType dtype_q,
 
 std::tuple<torch::Tensor, double> split_scale_param(const torch::Tensor& scale);
 
-#if !defined(USE_DCU)
+#if !defined(USE_DCU) && !defined(__ILUVATAR__)
 DLDataType to_dl_data_type(torch::ScalarType scalar_type);
 
 // below are tvm-ffi related functions
@@ -159,5 +160,5 @@ inline void bind_tvmffi_stream_to_current_torch_stream(
                  << " dev=" << device.index();
   }
 }
-#endif  // !defined(USE_DCU)
+#endif  // !defined(USE_DCU) && !defined(__ILUVATAR__)
-}  // namespace xllm::kernel::cuda
+}  // namespace xllm::kernel::cuda

qwen3_6_scripts/ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp

@@ -13,10 +13,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "kernels/cuda/cuda_ops_api.h"
+#include "cuda_ops_api.h"
-#include "kernels/cuda/utils.h"
+#include "utils.h"
-#include "platform/device.h"
+
-#include "platform/platform.h"
+
 
 namespace xllm::kernel::cuda {
 
@@ -51,74 +51,73 @@ torch::Tensor cutlass_fused_moe(
     bool use_packed_weights,
     int32_t tune_max_num_tokens,
     ActivationType activation_type) {
-  int64_t num_rows = input.size(0);
+  int64_t num_tokens = input.size(0);
   int64_t hidden_size = fc2_expert_weights.size(1);
+  int64_t inter_dim = fc1_expert_weights.size(1);
+  int64_t top_k = token_selected_experts.size(1);
 
+  int64_t num_rows = num_tokens;
   if (min_latency_mode) {
     num_rows *= fc2_expert_weights.size(0);
   }
 
-  std::vector<int64_t> output_shape = {num_rows, hidden_size};
   torch::Tensor result_output;
   if (output.has_value() && output.value().defined()) {
     result_output = output.value();
   } else {
-    torch::TensorOptions options = input.options().dtype(output_dtype);
+    result_output = torch::zeros({num_rows, hidden_size},
-    result_output = torch::empty(output_shape, options);
+                                 input.options().dtype(output_dtype));
   }
 
-  std::string fused_moe_uri = "fused_moe";
+  if (Platform::is_support_ivcore10()) {
-  if (Platform::is_support_sm90a()) {
+    // BI-V100 path: per-token expert gather + matmul + SiLU-gate + matmul
-    fused_moe_uri += "_90";
+    // This replaces the tvm ffi CUTLASS path with native PyTorch ops.
-  } else if (Platform::is_support_sm100a() || Platform::is_support_sm100f()) {
+    for (int64_t t = 0; t < num_tokens; ++t) {
-    fused_moe_uri += "_100";
+      auto token = input[t].unsqueeze(0);  // [1, hidden]
-  } else if (Platform::is_support_sm120a()) {
+      torch::Tensor accum = torch::zeros({1, hidden_size},
-    fused_moe_uri += "_120";
+                                         input.options().dtype(output_dtype));
-  } else {
+      for (int64_t k = 0; k < top_k; ++k) {
-    LOG(FATAL) << "FusedMoE is only supported on sm90, sm100, sm120.";
+        int64_t expert_id = token_selected_experts[t][k].item<int64_t>();
+        float scale = token_final_scales[t][k].item<float>();
+
+        // gate_up = token @ fc1[expert].T  → [1, inter_dim]
+        auto gate_up = torch::mm(token, fc1_expert_weights[expert_id].t());
+        if (fc1_expert_biases.has_value()) {
+          gate_up = gate_up + fc1_expert_biases.value()[expert_id];
+        }
+
+        // SwiGLU: split into gate and up, apply silu(gate) * up
+        torch::Tensor act;
+        if (activation_type == ActivationType::SWIGLU ||
+            activation_type == ActivationType::SWIGLU_BIAS) {
+          auto chunks = gate_up.chunk(2, /*dim=*/-1);
+          act = torch::silu(chunks[0]) * chunks[1];
+        } else if (activation_type == ActivationType::SILU) {
+          act = torch::silu(gate_up);
+        } else if (activation_type == ActivationType::GELU) {
+          act = torch::gelu(gate_up);
+        } else {
+          act = gate_up;  // identity
+        }
+
+        // down = act @ fc2[expert].T  → [1, hidden]
+        auto down = torch::mm(act, fc2_expert_weights[expert_id].t());
+        if (fc2_expert_biases.has_value()) {
+          down = down + fc2_expert_biases.value()[expert_id];
+        }
+
+        accum += down.to(output_dtype) * scale;
+      }
+      result_output[t] = accum.squeeze(0);
+    }
+    return result_output;
   }
 
-  bind_tvmffi_stream_to_current_torch_stream(input.device());
+  // Original NVIDIA GPU path (sm90/sm100/sm120) via tvm ffi
-
+  TORCH_CHECK(false,
-  ffi::Module fused_moe_runner =
+      "cutlass_fused_moe: no supported platform. "
-      get_function(fused_moe_uri, "init")(
+      "BI-V100 should use ivcore10 path above; "
-          to_dl_data_type(input.scalar_type()),
+      "NVIDIA GPUs require sm90+.");
-          to_dl_data_type(fc1_expert_weights.scalar_type()),
-          to_dl_data_type(output_dtype),
-          use_deepseek_fp8_block_scale,
-          use_w4_group_scaling,
-          use_mxfp8_act_scaling,
-          use_packed_weights)
-          .cast<ffi::Module>();
-
-  fused_moe_runner->GetFunction("run_moe").value()(
-      to_ffi_tensor(result_output),
-      to_ffi_tensor(input),
-      to_ffi_tensor(token_selected_experts),
-      to_ffi_optional_tensor(token_final_scales),
-      to_ffi_tensor(fc1_expert_weights),
-      to_ffi_optional_tensor(fc1_expert_biases),
-      to_ffi_tensor(fc2_expert_weights),
-      to_ffi_optional_tensor(fc2_expert_biases),
-      to_ffi_optional_array_tensors(quant_scales),
-      to_ffi_optional_tensor(input_sf),
-      to_ffi_optional_tensor(swiglu_alpha),
-      to_ffi_optional_tensor(swiglu_beta),
-      to_ffi_optional_tensor(swiglu_limit),
-      tp_size,
-      tp_rank,
-      ep_size,
-      ep_rank,
-      cluster_size,
-      cluster_rank,
-      enable_alltoall,
-      min_latency_mode,
-      /*profile_ids=*/ffi::Optional<ffi::Array<int64_t>>(),  // TODO: support
-                                                             // auto tuning
-                                                             // profile ids
-      support_pdl(),
-      activation_type);
-
   return result_output;
 }
-}  // namespace xllm::kernel::cuda
+}  // namespace xllm::kernel::cuda

qwen3_6_scripts/ex_engine/xllm_kernels/cuda/moe/moe_topk_softmax_kernels.cuh

@@ -23,8 +23,8 @@ limitations under the License.
 
 #include <cub/util_type.cuh>
 
-#if !defined(USE_DCU) && !defined(USE_MACA)
+// <cuda/functional> requires CUDA 12+ (libcudacxx); BI-V100 runs CUDA 10.2
-#endif
+// and does not ship that header.  The include is unused in this file anyway.
 
 #include "device_utils.cuh"
 

qwen3_6_scripts/ex_engine/xllm_kernels/param.h

@@ -817,6 +817,7 @@ struct RandomSampleParams {
 };
 
 // Rejection sampling parameters for speculative decoding
+// PRD #0 (9e0f1402): Added mode field for explicit greedy/probabilistic dispatch
 struct RejectionSampleParams {
   // Candidate draft token indices to be verified.
   // Shape: [total_draft_tokens]. Dtype: int32.
@@ -849,6 +850,10 @@ struct RejectionSampleParams {
   // The maximum number of draft tokens in the batch (max value in
   // num_draft_tokens).
   int32_t max_spec_len;
+  // PRD #0 (9e0f1402): Explicit sampling mode — Greedy uses exact-match,
+  // Probabilistic uses min(1, p/q) acceptance ratio.
+  // Default: Probabilistic (preserves legacy behavior).
+  uint8_t draft_sampling_mode = 1;  // 0=Greedy, 1=Probabilistic
 };
 
 // Masked indexer select paged KV cache parameters
@@ -1438,4 +1443,4 @@ struct ChunkGatedDeltaRuleParams {
   // Whether to apply L2 norm to q and k inside the kernel. Default: false.
   bool use_qk_l2norm_in_kernel = false;
 };
-}  // namespace xllm::kernel
+}  // namespace xllm::kernel

qwen3_6_scripts/moe_topk_sigmoid_kernels.cuh

@@ -23,12 +23,21 @@ limitations under the License.
 
 #include <cub/util_type.cuh>
 
+#if !defined(USE_DCU) && !defined(USE_MACA)
+#endif
+
 #include "device_utils.cuh"
 
 namespace {
 
 using namespace xllm::kernel::cuda;
 
+#if defined(USE_DCU)
+static constexpr unsigned long long kSigmoidFullMask = 0xffffffffffffffffULL;
+#else
+static constexpr unsigned int kSigmoidFullMask = 0xffffffffU;
+#endif
+
 // ====================== Sigmoid things ===============================
 // We have our own implementation of sigmoid here so we can support transposing
 // the output in the sigmoid kernel when we extend this module to support
@@ -182,29 +191,29 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
 
   // Number of bytes each thread pulls in per load
-  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
-  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int kEltsPerRow = NUM_EXPERTS;
-  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
-  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+  static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
 
   // Restrictions based on previous section.
   static_assert(
-      VPT % ELTS_PER_LDG == 0,
+      VPT % kEltsPerLdg == 0,
       "The elements per thread must be a multiple of the elements per ldg");
-  static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
+  static_assert(WARP_SIZE % kThreadsPerRow == 0,
                 "The threads per row must cleanly divide the threads per warp");
-  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+  static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
                 "THREADS_PER_ROW must be power of 2");
-  static_assert(THREADS_PER_ROW <= WARP_SIZE,
+  static_assert(kThreadsPerRow <= WARP_SIZE,
                 "THREADS_PER_ROW can be at most warp size");
 
   // We have NUM_EXPERTS elements per row. We specialize for small #experts
-  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
-  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
-  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+  static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
 
   // Restrictions for previous section.
-  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+  static_assert(kEltsPerWarp % kEltsPerRow == 0,
                 "The elts per row must cleanly divide the total elt per warp");
 
   // ===================== From this point, we finally start computing run-time
@@ -213,14 +222,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
   // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
   // rows. We start by computing the start row for each block.
-  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+  const int cta_base_row = blockIdx.x * kRowsPerCta;
 
   // Now, using the base row per thread block, we compute the base row per warp.
-  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+  const int warp_base_row = cta_base_row + threadIdx.y * kRowsPerWarp;
 
   // The threads in a warp are split into sub-groups that will work on a row.
   // We compute row offset for each thread sub-group
-  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row_in_warp = threadIdx.x / kThreadsPerRow;
   const int thread_row = warp_base_row + thread_row_in_warp;
 
   // Threads with indices out of bounds should early exit here.
@@ -231,12 +240,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 
   // We finally start setting up the read pointers for each thread. First, each
   // thread jumps to the start of the row it will read.
-  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+  const T* thread_row_ptr = input + thread_row * kEltsPerRow;
 
   // Now, we compute the group each thread belong to in order to determine the
   // first column to start loads.
-  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int thread_group_idx = threadIdx.x % kThreadsPerRow;
-  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
   const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
 
   // Determine the pointer type to use to read in the data depending on the
@@ -244,7 +253,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
   // array here. We defined our own aligned array and use it here to avoid the
   // dependency on CUTLASS.
-  using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+  using AccessType = AlignedArray<T, kEltsPerLdg>;
 
   // Finally, we pull in the data from global mem
   T row_chunk_temp[VPT];
@@ -256,8 +265,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Note(Byron): interleaved loads to achieve better memory coalescing
   // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
   // thread[2] | thread[3] | ...
-  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+  for (int ii = 0; ii < kLdgPerThread; ++ii) {
-    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
   }
 
   float row_chunk[VPT];
@@ -274,11 +283,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
       |--------- group0 --------| |----------group1 --------|
                                     ^ local2
       */
-      const int group_id = ii / ELTS_PER_LDG;
+      const int group_id = ii / kEltsPerLdg;
-      const int local_id = ii % ELTS_PER_LDG;
+      const int local_id = ii % kEltsPerLdg;
       const int expert_idx = first_elt_read_by_thread +
-                             group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                             group_id * kThreadsPerRow * kEltsPerLdg + local_id;
-                             local_id;
       val = val + correction_bias[expert_idx];
     }
 
@@ -288,7 +296,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Now, row_chunk contains the sigmoid of the row chunk. Now, I want to find
   // the topk elements in each row, along with the max index.
   int start_col = first_elt_read_by_thread;
-  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+  static constexpr int kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
 
   float row_sum_for_renormalize = 0;
 
@@ -297,11 +305,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     float max_val = row_chunk[0];
     int expert = start_col;
 #pragma unroll
-    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+    for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
-         ++ldg, col += COLS_PER_GROUP_LDG) {
+         ++ldg, col += kColsPerGroupLdg) {
 #pragma unroll
-      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+      for (int ii = 0; ii < kEltsPerLdg; ++ii) {
-        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+        float val = row_chunk[ldg * kEltsPerLdg + ii];
 
         // No check on the experts here since columns with the smallest index
         // are processed first and only updated if > (not >=)
@@ -317,11 +325,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 // threads can agree on "who" had the max value. That thread can then blank out
 // their max with -inf and the warp can run more iterations...
 #pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
-      float other_max =
+      float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
-          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW);
+          kSigmoidFullMask, max_val, mask, kThreadsPerRow);
-      int other_expert =
+      int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
-          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW);
+          kSigmoidFullMask, expert, mask, kThreadsPerRow);
 
       // We want lower indices to "win" in every thread so we break ties this
       // way
@@ -354,17 +362,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     // Finally, we clear the value in the thread with the current max if there
     // is another iteration to run.
     if (k_idx + 1 < k) {
-      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int ldg_group_for_expert = expert / kColsPerGroupLdg;
       const int thread_to_clear_in_group =
-          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+          (expert / kEltsPerLdg) % kThreadsPerRow;
 
       // Only the thread in the group which produced the max will reset the
       // "winning" value to -inf.
       if (thread_group_idx == thread_to_clear_in_group) {
-        const int offset_for_expert = expert % ELTS_PER_LDG;
+        const int offset_for_expert = expert % kEltsPerLdg;
         // Safe to set to any negative value since row_chunk values must be
         // between 0 and 1.
-        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+        row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
             -10000.f;
       }
     }
@@ -393,18 +401,17 @@ void topk_gating_sigmoid_launcher_helper(const T* input,
                                          const bool renormalize,
                                          const float* correction_bias,
                                          cudaStream_t stream) {
-  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+  static constexpr std::size_t kMaxBytesPerLdg = 16;
 
-  static constexpr int BYTES_PER_LDG =
+  static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
-      MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
-  using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int kVpt = Constants::VPT;
-  static constexpr int VPT = Constants::VPT;
+  static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
-  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
-  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
   const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
 
   dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
-  topk_gating_sigmoid<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
+  topk_gating_sigmoid<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
       <<<num_blocks, block_dim, 0, stream>>>(input,
                                              finished,
                                              output,
@@ -442,55 +449,55 @@ void topk_gating_sigmoid_kernel_launcher(const T* gating_output,
                                          const bool renormalize,
                                          const float* correction_bias,
                                          cudaStream_t stream) {
-  static constexpr int WARPS_PER_TB = 4;
+  static constexpr int kWarpsPerTb = 4;
   switch (num_experts) {
     case 1:
-      LAUNCH_SIGMOID(T, 1, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 1, kWarpsPerTb);
       break;
     case 2:
-      LAUNCH_SIGMOID(T, 2, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 2, kWarpsPerTb);
       break;
     case 4:
-      LAUNCH_SIGMOID(T, 4, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 4, kWarpsPerTb);
       break;
     case 8:
-      LAUNCH_SIGMOID(T, 8, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 8, kWarpsPerTb);
       break;
     case 16:
-      LAUNCH_SIGMOID(T, 16, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 16, kWarpsPerTb);
       break;
     case 32:
-      LAUNCH_SIGMOID(T, 32, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 32, kWarpsPerTb);
       break;
     case 64:
-      LAUNCH_SIGMOID(T, 64, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 64, kWarpsPerTb);
       break;
     case 128:
-      LAUNCH_SIGMOID(T, 128, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 128, kWarpsPerTb);
       break;
     case 256:
-      LAUNCH_SIGMOID(T, 256, WARPS_PER_TB);
+      LAUNCH_SIGMOID(T, 256, kWarpsPerTb);
       break;
     default: {
       TORCH_CHECK(sigmoid_workspace != nullptr,
                   "sigmoid_workspace must be provided for num_experts that are "
                   "not a power of 2.");
-      static constexpr int TPB = 256;
+      static constexpr int kTpb = 256;
-      moe_sigmoid<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
+      moe_sigmoid<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
-                                                          nullptr,
+                                                            nullptr,
-                                                          sigmoid_workspace,
+                                                            sigmoid_workspace,
-                                                          num_experts,
+                                                            num_experts,
-                                                          correction_bias);
+                                                            correction_bias);
-      moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(sigmoid_workspace,
+      moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(sigmoid_workspace,
-                                                    nullptr,
+                                                      nullptr,
-                                                    topk_weights,
+                                                      topk_weights,
-                                                    topk_indices,
+                                                      topk_indices,
-                                                    num_experts,
+                                                      num_experts,
-                                                    topk,
+                                                      topk,
-                                                    0,
+                                                      0,
-                                                    num_experts,
+                                                      num_experts,
-                                                    renormalize,
+                                                      renormalize,
-                                                    correction_bias);
+                                                      correction_bias);
     }
   }
 }
@@ -582,8 +589,8 @@ void topk_sigmoid(torch::Tensor& topk_weights,   // [num_tokens, topk]
         bias_ptr,
         stream);
   } else if (dtype == at::ScalarType::BFloat16) {
-    topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>(
+    topk_gating_sigmoid_kernel_launcher<BFloat16Type>(
-        reinterpret_cast<const __nv_bfloat16*>(
+        reinterpret_cast<const BFloat16Type*>(
             gating_output.data_ptr<at::BFloat16>()),
         topk_weights.data_ptr<float>(),
         topk_indices.data_ptr<int>(),

qwen3_6_scripts/moe_topk_softmax_kernels.cuh

@@ -23,6 +23,9 @@ limitations under the License.
 
 #include <cub/util_type.cuh>
 
+// <cuda/functional> requires CUDA 12+ (libcudacxx); BI-V100 runs CUDA 10.2
+// and does not ship that header.  The include is unused in this file anyway.
+
 #include "device_utils.cuh"
 
 using cub_kvp = cub::KeyValuePair<int, float>;
@@ -31,6 +34,12 @@ namespace {
 
 using namespace xllm::kernel::cuda;
 
+#if defined(USE_DCU)
+static constexpr unsigned long long kSoftmaxFullMask = 0xffffffffffffffffULL;
+#else
+static constexpr unsigned int kSoftmaxFullMask = 0xffffffffU;
+#endif
+
 // ====================== Softmax things ===============================
 // We have our own implementation of softmax here so we can support transposing
 // the output in the softmax kernel when we extend this module to support
@@ -110,9 +119,10 @@ __launch_bounds__(TPB) __global__
 }
 
 namespace moe {
-struct TopKPair {
+class TopKPair {
-  static const int PAIR = 2;
+ public:
-  static const int MAX_INDEX = 0;
+  static constexpr int kPair = 2;
+  static constexpr int kMaxIndex = 0;
   cub_kvp max;
   cub_kvp secondMax;
 
@@ -121,7 +131,8 @@ struct TopKPair {
       : max(max), secondMax(secondMax) {}
 };
 
-struct TopKPairArgMax {
+class TopKPairArgMax {
+ public:
   __device__ TopKPairArgMax() {}
   __device__ __forceinline__ TopKPair
   operator()(const TopKPair& candidate1, const TopKPair& candidate2) const {
@@ -175,8 +186,8 @@ __launch_bounds__(TPB) __global__
   const int thread_read_offset = blockIdx.x * num_experts;
   float row_sum_for_renormalize = 0;
   // Each loop finds the top 2 elements,
-  // thus requiring only ⌈k/2⌉ loops (calculated as (k + 1) / 2).
+  // thus requiring only ceil(k / 2) loops (calculated as (k + 1) / 2).
-  for (int k_idx = 0; k_idx < (k + TopKPair::PAIR - 1) / TopKPair::PAIR;
+  for (int k_idx = 0; k_idx < (k + TopKPair::kPair - 1) / TopKPair::kPair;
        ++k_idx) {
     // Initializing the top 2 elements by the minimum value.
     thread_pair.max.key = 0;
@@ -204,9 +215,11 @@ __launch_bounds__(TPB) __global__
     if (threadIdx.x == 0) {
 #pragma unroll
       // updating 2 elements to the result.
-      for (int i = 0; i < TopKPair::PAIR; i++) {
+      for (int i = 0; i < TopKPair::kPair; i++) {
-        if (k_idx * 2 + i >= k) break;
+        if (k_idx * 2 + i >= k) {
-        cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max
+          break;
+        }
+        cub_kvp result = (i == TopKPair::kMaxIndex) ? result_pair.max
                                                     : result_pair.secondMax;
         int expert = result.key;
         bool node_uses_expert = expert >= start_expert && expert < end_expert;
@@ -342,29 +355,29 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
 
   // Number of bytes each thread pulls in per load
-  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
-  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int kEltsPerRow = NUM_EXPERTS;
-  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
-  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+  static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
 
   // Restrictions based on previous section.
   static_assert(
-      VPT % ELTS_PER_LDG == 0,
+      VPT % kEltsPerLdg == 0,
       "The elements per thread must be a multiple of the elements per ldg");
-  static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
+  static_assert(WARP_SIZE % kThreadsPerRow == 0,
                 "The threads per row must cleanly divide the threads per warp");
-  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+  static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
                 "THREADS_PER_ROW must be power of 2");
-  static_assert(THREADS_PER_ROW <= WARP_SIZE,
+  static_assert(kThreadsPerRow <= WARP_SIZE,
                 "THREADS_PER_ROW can be at most warp size");
 
   // We have NUM_EXPERTS elements per row. We specialize for small #experts
-  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
-  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
-  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+  static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
 
   // Restrictions for previous section.
-  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+  static_assert(kEltsPerWarp % kEltsPerRow == 0,
                 "The elts per row must cleanly divide the total elt per warp");
 
   // ===================== From this point, we finally start computing run-time
@@ -373,14 +386,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
   // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
   // rows. We start by computing the start row for each block.
-  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+  const int cta_base_row = blockIdx.x * kRowsPerCta;
 
   // Now, using the base row per thread block, we compute the base row per warp.
-  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+  const int warp_base_row = cta_base_row + threadIdx.y * kRowsPerWarp;
 
   // The threads in a warp are split into sub-groups that will work on a row.
   // We compute row offset for each thread sub-group
-  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row_in_warp = threadIdx.x / kThreadsPerRow;
   const int thread_row = warp_base_row + thread_row_in_warp;
 
   // Threads with indices out of bounds should early exit here.
@@ -391,12 +404,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 
   // We finally start setting up the read pointers for each thread. First, each
   // thread jumps to the start of the row it will read.
-  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+  const T* thread_row_ptr = input + thread_row * kEltsPerRow;
 
   // Now, we compute the group each thread belong to in order to determine the
   // first column to start loads.
-  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int thread_group_idx = threadIdx.x % kThreadsPerRow;
-  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
   const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
 
   // Determine the pointer type to use to read in the data depending on the
@@ -404,7 +417,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
   // array here. We defined our own aligned array and use it here to avoid the
   // dependency on CUTLASS.
-  using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+  using AccessType = AlignedArray<T, kEltsPerLdg>;
 
   // Finally, we pull in the data from global mem
   T row_chunk_temp[VPT];
@@ -416,8 +429,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Note(Byron): interleaved loads to achieve better memory coalescing
   // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
   // thread[2] | thread[3] | ...
-  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+  for (int ii = 0; ii < kLdgPerThread; ++ii) {
-    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
   }
 
   float row_chunk[VPT];
@@ -446,10 +459,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
         |--------- group0 --------| |----------group1 --------|
                                       ^ local2
         */
-        const int group_id = ii / ELTS_PER_LDG;
+        const int group_id = ii / kEltsPerLdg;
-        const int local_id = ii % ELTS_PER_LDG;
+        const int local_id = ii % kEltsPerLdg;
         const int expert_idx = first_elt_read_by_thread +
-                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                               group_id * kThreadsPerRow * kEltsPerLdg +
                                local_id;
         val = val + correction_bias[expert_idx];
       }
@@ -474,11 +487,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 // Now, we find the max within the thread group and distribute among the
 // threads. We use a butterfly reduce. lane id: 0-31 within a warp
 #pragma unroll
-  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+  for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
     // butterfly reduce with (lane id ^ mask)
     thread_max = max(thread_max,
                      XLLM_SHFL_XOR_SYNC_WIDTH(
-                         0xffffffff, thread_max, mask, THREADS_PER_ROW));
+                         kSoftmaxFullMask, thread_max, mask, kThreadsPerRow));
   }
 
   // From this point, thread max in all the threads have the max within the row.
@@ -494,9 +507,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 // Now, we perform the sum reduce within each thread group. Similar to the max
 // reduce, we use a bufferfly pattern.
 #pragma unroll
-  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+  for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
-    row_sum +=
+    row_sum += XLLM_SHFL_XOR_SYNC_WIDTH(
-        XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, row_sum, mask, THREADS_PER_ROW);
+        kSoftmaxFullMask, row_sum, mask, kThreadsPerRow);
   }
 
   // From this point, all threads have the max and the sum for their rows in the
@@ -519,7 +532,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
   // Now, softmax_res contains the softmax of the row chunk. Now, I want to find
   // the topk elements in each row, along with the max index.
   int start_col = first_elt_read_by_thread;
-  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+  static constexpr int kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
 
   float row_sum_for_renormalize = 0;
 
@@ -528,11 +541,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     float max_val = row_chunk[0];
     int expert = start_col;
 #pragma unroll
-    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+    for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
-         ++ldg, col += COLS_PER_GROUP_LDG) {
+         ++ldg, col += kColsPerGroupLdg) {
 #pragma unroll
-      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+      for (int ii = 0; ii < kEltsPerLdg; ++ii) {
-        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+        float val = row_chunk[ldg * kEltsPerLdg + ii];
 
         // No check on the experts here since columns with the smallest index
         // are processed first and only updated if > (not >=)
@@ -548,11 +561,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 // threads can agree on "who" had the max value. That thread can then blank out
 // their max with -inf and the warp can run more iterations...
 #pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
-      float other_max =
+      float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
-          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW);
+          kSoftmaxFullMask, max_val, mask, kThreadsPerRow);
-      int other_expert =
+      int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
-          XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW);
+          kSoftmaxFullMask, expert, mask, kThreadsPerRow);
 
       // We want lower indices to "win" in every thread so we break ties this
       // way
@@ -582,17 +595,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     // Finally, we clear the value in the thread with the current max if there
     // is another iteration to run.
     if (k_idx + 1 < k) {
-      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int ldg_group_for_expert = expert / kColsPerGroupLdg;
       const int thread_to_clear_in_group =
-          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+          (expert / kEltsPerLdg) % kThreadsPerRow;
 
       // Only the thread in the group which produced the max will reset the
       // "winning" value to -inf.
       if (thread_group_idx == thread_to_clear_in_group) {
-        const int offset_for_expert = expert % ELTS_PER_LDG;
+        const int offset_for_expert = expert % kEltsPerLdg;
         // Safe to set to any negative value since row_chunk values must be
         // between 0 and 1.
-        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+        row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
             -10000.f;
       }
     }
@@ -622,18 +635,17 @@ void topk_gating_softmax_launcher_helper(const T* input,
                                          const float moe_softcapping,
                                          const float* correction_bias,
                                          cudaStream_t stream) {
-  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+  static constexpr std::size_t kMaxBytesPerLdg = 16;
 
-  static constexpr int BYTES_PER_LDG =
+  static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
-      MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
-  using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int kVpt = Constants::VPT;
-  static constexpr int VPT = Constants::VPT;
+  static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
-  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
-  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
   const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
 
   dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
-  topk_gating_softmax<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
+  topk_gating_softmax<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
       <<<num_blocks, block_dim, 0, stream>>>(input,
                                              finished,
                                              output,
@@ -674,67 +686,69 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
                                          const float moe_softcapping,
                                          const float* correction_bias,
                                          cudaStream_t stream) {
-  static constexpr int WARPS_PER_TB = 4;
+  static constexpr int kWarpsPerTb = 4;
   switch (num_experts) {
     case 1:
-      LAUNCH_SOFTMAX(T, 1, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 1, kWarpsPerTb);
       break;
     case 2:
-      LAUNCH_SOFTMAX(T, 2, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 2, kWarpsPerTb);
       break;
     case 4:
-      LAUNCH_SOFTMAX(T, 4, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 4, kWarpsPerTb);
       break;
     case 8:
-      LAUNCH_SOFTMAX(T, 8, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 8, kWarpsPerTb);
       break;
     case 16:
-      LAUNCH_SOFTMAX(T, 16, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 16, kWarpsPerTb);
       break;
     case 32:
-      LAUNCH_SOFTMAX(T, 32, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 32, kWarpsPerTb);
       break;
     case 64:
-      LAUNCH_SOFTMAX(T, 64, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 64, kWarpsPerTb);
       break;
     case 128:
-      LAUNCH_SOFTMAX(T, 128, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 128, kWarpsPerTb);
       break;
     case 256:
-      LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
+      LAUNCH_SOFTMAX(T, 256, kWarpsPerTb);
       break;
     default: {
-      TORCH_CHECK(softmax_workspace != nullptr, "softmax_workspace must be provided for num_experts that are not a power of 2.");
+      CHECK(softmax_workspace != nullptr)
-      static constexpr int TPB = 256;
+          << "softmax_workspace must be provided for num_experts that are "
-      moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
+             "not a power of 2.";
-                                                          nullptr,
+      static constexpr int kTpb = 256;
-                                                          softmax_workspace,
+      moe_softmax<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
-                                                          num_experts,
+                                                            nullptr,
-                                                          moe_softcapping,
+                                                            softmax_workspace,
-                                                          correction_bias);
+                                                            num_experts,
+                                                            moe_softcapping,
+                                                            correction_bias);
       if (topk == 1) {
         // Note: As an optimization for better performance,
         // the softmax_workspace is overwritten in-place by both moeTopK and
         // moe_topk_fast.
-        moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace,
+        moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
-                                                      nullptr,
+                                                        nullptr,
-                                                      topk_weights,
+                                                        topk_weights,
-                                                      topk_indices,
+                                                        topk_indices,
-                                                      num_experts,
+                                                        num_experts,
-                                                      topk,
+                                                        topk,
-                                                      0,
+                                                        0,
-                                                      num_experts,
+                                                        num_experts,
-                                                      renormalize);
+                                                        renormalize);
       } else {
-        moe_topk_fast<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace,
+        moe_topk_fast<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
-                                                           nullptr,
+                                                             nullptr,
-                                                           topk_weights,
+                                                             topk_weights,
-                                                           topk_indices,
+                                                             topk_indices,
-                                                           num_experts,
+                                                             num_experts,
-                                                           topk,
+                                                             topk,
-                                                           0,
+                                                             0,
-                                                           num_experts,
+                                                             num_experts,
-                                                           renormalize);
+                                                             renormalize);
       }
     }
   }
@@ -749,20 +763,29 @@ void topk_softmax(torch::Tensor& topk_weights,   // [num_tokens, topk]
                   const double moe_softcapping,
                   const std::optional<torch::Tensor>& correction_bias) {
   // Check data type
-  TORCH_CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
+  CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
         gating_output.scalar_type() == at::ScalarType::Half ||
-        gating_output.scalar_type() == at::ScalarType::BFloat16,
+        gating_output.scalar_type() == at::ScalarType::BFloat16)
-        "gating_output must be float32, float16, or bfloat16");
+      << "gating_output must be float32, float16, or bfloat16";
 
   // Check dimensions
-  TORCH_CHECK(gating_output.dim() == 2, "gating_output must be 2D tensor [num_tokens, num_experts]");
+  CHECK(gating_output.dim() == 2)
-  TORCH_CHECK(topk_weights.dim() == 2, "topk_weights must be 2D tensor [num_tokens, topk]");
+      << "gating_output must be 2D tensor [num_tokens, num_experts]";
-  TORCH_CHECK(topk_indices.dim() == 2, "topk_indices must be 2D tensor [num_tokens, topk]");
+  CHECK(topk_weights.dim() == 2)
+      << "topk_weights must be 2D tensor [num_tokens, topk]";
+  CHECK(topk_indices.dim() == 2)
+      << "topk_indices must be 2D tensor [num_tokens, topk]";
 
   // Check shapes
-  TORCH_CHECK(gating_output.size(0) == topk_weights.size(0), "First dimension of topk_weights must match num_tokens in gating_output First dimension of topk_indices must match num_tokens in gating_output");
+  CHECK(gating_output.size(0) == topk_weights.size(0))
+      << "First dimension of topk_weights must match num_tokens in "
+         "gating_output"
+      << "First dimension of topk_indices must match num_tokens in "
+         "gating_output";
 
-  TORCH_CHECK(topk_weights.size(-1) == topk_indices.size(-1), "Second dimension of topk_indices must match topk in topk_weights topk must be less than or equal to num_experts");
+  CHECK(topk_weights.size(-1) == topk_indices.size(-1))
+      << "Second dimension of topk_indices must match topk in topk_weights"
+      << "topk must be less than or equal to num_experts";
 
   const int num_experts = static_cast<int>(gating_output.size(-1));
   const int num_tokens = static_cast<int>(gating_output.size(0));
@@ -784,9 +807,12 @@ void topk_softmax(torch::Tensor& topk_weights,   // [num_tokens, topk]
   const float* bias_ptr = nullptr;
   if (correction_bias.has_value()) {
     const torch::Tensor& bias_tensor = correction_bias.value();
-    TORCH_CHECK(bias_tensor.dim() == 1, "correction_bias must be 1D tensor [num_experts]");
+    CHECK(bias_tensor.dim() == 1)
-    TORCH_CHECK(bias_tensor.size(0) == num_experts, "correction_bias size must match num_experts");
+        << "correction_bias must be 1D tensor [num_experts]";
-    TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "correction_bias must be float32");
+    CHECK(bias_tensor.size(0) == num_experts)
+        << "correction_bias size must match num_experts";
+    CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
+        << "correction_bias must be float32, got " << bias_tensor.scalar_type();
     bias_ptr = bias_tensor.data_ptr<float>();
   }
 
@@ -820,8 +846,8 @@ void topk_softmax(torch::Tensor& topk_weights,   // [num_tokens, topk]
         bias_ptr,
         stream);
   } else if (dtype == at::ScalarType::BFloat16) {
-    topk_gating_softmax_kernel_launcher<__nv_bfloat16>(
+    topk_gating_softmax_kernel_launcher<BFloat16Type>(
-        reinterpret_cast<const __nv_bfloat16*>(
+        reinterpret_cast<const BFloat16Type*>(
             gating_output.data_ptr<at::BFloat16>()),
         topk_weights.data_ptr<float>(),
         topk_indices.data_ptr<int>(),
@@ -834,7 +860,7 @@ void topk_softmax(torch::Tensor& topk_weights,   // [num_tokens, topk]
         bias_ptr,
         stream);
   } else {
-    TORCH_CHECK(false, "Unsupported gating_output dtype");
+    LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
   }
 }
 }  // namespace xllm::kernel::cuda

qwen3_6_scripts/patch_ops.sh

@@ -1,3 +1,5 @@
+echo "[build] trigger 20260901"
+echo "[build] trigger 202609011237"
 #!/usr/bin/env bash
 # BI-V100 patch script for Qwen3.6-35B-A3B (Qwen3_5 MoE architecture)
 #
@@ -227,6 +229,13 @@ python3 ./patch_vllm_tool_parser.py
 # Adds --reasoning-parser qwen3 support.
 # Routes thinking tokens to reasoning_content, rest to content in the delta.
 # Works together with --tool-call-parser qwen3_coder (think → tool call flow).
+#
+# PRD #69: Clear __pycache__ BEFORE copying patched .py files.
+# Base image's compiled .pyc (protocol.py with extra="forbid") would otherwise
+# shadow our patched .py, causing ~25% of requests to reject
+# max_completion_tokens / reasoning_effort with HTTP 400.
+find "${VLLM_ROOT}/entrypoints" -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
+
 cp -r ./reasoning "${VLLM_ROOT}/"
 cp ./protocol.py "${VLLM_ROOT}/entrypoints/openai/protocol.py"
 cp ./cli_args.py "${VLLM_ROOT}/entrypoints/openai/cli_args.py"
@@ -246,6 +255,11 @@ if source != installed:
     raise SystemExit("runtime api_server overlay identity mismatch")
 PY
 
+# PRD #69: Clear ALL __pycache__ under VLLM_ROOT after every cp/patch is done.
+# py_compile below only compiles ./qwen3_6_scripts, not VLLM_ROOT, so this
+# ensures the docker snapshot has no stale .pyc for any patched vllm module.
+find "${VLLM_ROOT}" -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
+
 build_stage "compiling submission Python sources"
 find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile
 build_stage "patch script completed"

qwen3_6_scripts/qwen3_5.py

@@ -7,6 +7,7 @@ import hashlib
 import os
 import sys
 import time
+print("[qwen3_5] module load START", file=sys.stderr, flush=True)
 from typing import (Any, Dict, Iterable, List, Literal, Mapping, Optional,
                     Tuple, TypedDict, Union)
 
@@ -166,6 +167,41 @@ except ImportError:
     except ImportError:
         _xllm_moe = None
 
+# --- xllm prebuilt kernel loading (PRD build) ---
+def _load_xllm_prebuilt(name):
+    """Load a prebuilt xllm .so from corex-3.2.3-ivcore10 directory."""
+    import importlib.util as _ilu
+    _search = [
+        f"/workspace/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/{name}.so",
+        os.path.join(os.path.dirname(__file__), "prebuilt",
+                     "corex-3.2.3-ivcore10", f"{name}.so"),
+        # When patch_ops.sh copies this file into vllm package, __file__
+        # points to vllm/model_executor/models/ — look back up to workspace
+        f"/home/dylan/0814/project_6/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/{name}.so",
+    ]
+    for _p in _search:
+        if os.path.isfile(_p):
+            print(f"[xllm] loading {name} from {_p} ...", file=sys.stderr, flush=True)
+            try:
+                _spec = _ilu.spec_from_file_location(name, _p)
+                _mod = _ilu.module_from_spec(_spec)
+                _spec.loader.exec_module(_mod)
+                print(f"[xllm] {name} OK: {[x for x in dir(_mod) if not x.startswith('_')]}", file=sys.stderr, flush=True)
+                return _mod
+            except Exception as _e:
+                print(f"[xllm] {name} FAILED: {_e}", file=sys.stderr, flush=True)
+                return None
+    print(f"[xllm] {name} not found", file=sys.stderr, flush=True)
+    return None
+
+print("[xllm] loading prebuilt kernels ...", file=sys.stderr, flush=True)
+_xllm_norm = _load_xllm_prebuilt("xllm_norm")
+_xllm_rope = _load_xllm_prebuilt("xllm_rope")
+_xllm_activation = _load_xllm_prebuilt("xllm_activation")
+_xllm_cache = _load_xllm_prebuilt("xllm_cache")
+_xllm_fused_qknorm_rope = _load_xllm_prebuilt("xllm_fused_qknorm_rope")
+print("[xllm] prebuilt kernel loading done", file=sys.stderr, flush=True)
+
 # ix_moe_bridge: direct GEMV (4.1x faster than F.linear for decode M=1)
 # Benchmark: F.linear 133us vs br.linear 32us on BI-V100
 try:
@@ -189,7 +225,7 @@ _HAS_BRIDGE_LINEAR = (
     _ix_moe_bridge is not None
     and hasattr(_ix_moe_bridge, 'linear'))
 if _HAS_BRIDGE_LINEAR:
-    logger.info("ix_moe_bridge.linear ENABLED — 4.1x GEMV speedup for decode")
+    print("[xllm] ix_moe_bridge.linear ENABLED", file=sys.stderr, flush=True)
 
 
 def _fast_linear(x: torch.Tensor, weight: torch.Tensor,
@@ -269,6 +305,141 @@ if _USE_XLLM_MOE:
     logger.info("xllm_moe ENABLED — fused_topk + compute_index + combine_result")
 _USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True)
 
+# --- xllm CUDA kernel flags ---
+_USE_XLLM_NORM = (
+    _xllm_norm is not None
+    and env_bool("BI100_XLLM_NORM", True))
+_USE_XLLM_ROPE = (
+    _xllm_rope is not None
+    and env_bool("BI100_XLLM_ROPE", True))
+_USE_XLLM_ACTIVATION = (
+    _xllm_activation is not None
+    and env_bool("BI100_XLLM_ACTIVATION", True))
+_USE_XLLM_CACHE = (
+    _xllm_cache is not None
+    and env_bool("BI100_XLLM_CACHE", True))
+_USE_XLLM_FUSED_QKNORM_ROPE = (
+    _xllm_fused_qknorm_rope is not None
+    and env_bool("BI100_XLLM_FUSED_QKNORM_ROPE", True))
+if _USE_XLLM_NORM:
+    logger.info("xllm_norm ENABLED — fused RMSNorm CUDA kernel")
+if _USE_XLLM_ROPE:
+    logger.info("xllm_rope ENABLED — fused RoPE CUDA kernel")
+if _USE_XLLM_ACTIVATION:
+    logger.info("xllm_activation ENABLED — fused SiLU-and-Mul CUDA kernel")
+if _USE_XLLM_CACHE:
+    logger.info("xllm_cache ENABLED — fused reshape_paged_cache CUDA kernel")
+if _USE_XLLM_FUSED_QKNORM_ROPE:
+    logger.info("xllm_fused_qknorm_rope ENABLED — fused QK-Norm+RoPE (saves 128 launches/fwd)")
+
+# ---------------------------------------------------------------------------
+# xllm kernel monkey-patches: replace PyTorch fallback → fused CUDA kernels
+# ---------------------------------------------------------------------------
+print("[xllm] applying monkey-patches ...", file=sys.stderr, flush=True)
+
+# --- 1. RMSNorm: xllm_norm replaces GemmaRMSNorm.forward_cuda ---
+# GemmaRMSNorm uses  x * (1 + weight)  while xllm kernel uses  x * weight.
+# We pre-compute (1 + weight) and cache it so the kernel sees the correct
+# effective weight.  The cached tensor is lazily materialised on first call
+# and invalidated if shape/device/dtype change (should never happen after
+# model init).
+if _USE_XLLM_NORM:
+    from vllm.model_executor.layers.layernorm import GemmaRMSNorm as _GemmaRMSNorm
+
+    def _xllm_rms_norm_forward_cuda(self, x, residual=None):
+        """Drop-in for GemmaRMSNorm.forward_cuda using xllm_norm CUDA kernel."""
+        # Lazily compute and cache  effective_weight = 1 + weight
+        _ew = getattr(self, '_xllm_eff_weight', None)
+        if (_ew is None
+                or _ew.shape != self.weight.shape
+                or _ew.device != self.weight.device
+                or _ew.dtype != self.weight.dtype):
+            _ew = (1.0 + self.weight.data.float()).to(self.weight.dtype)
+            self._xllm_eff_weight = _ew
+
+        if residual is not None:
+            # fused_add_rms_norm:  input = RMSNorm(input + residual),
+            #                     residual = input + residual  (before norm)
+            # The kernel modifies input and residual IN-PLACE.
+            _xllm_norm.fused_add_rms_norm(
+                x, residual, _ew, self.variance_epsilon)
+            return x, residual
+        else:
+            out = torch.empty_like(x)
+            _xllm_norm.rms_norm(out, x, _ew, self.variance_epsilon)
+            return out
+
+    _GemmaRMSNorm.forward_cuda = _xllm_rms_norm_forward_cuda
+    logger.info("xllm_norm PATCHED — GemmaRMSNorm.forward_cuda → xllm CUDA kernel")
+
+
+# --- 2. SiluAndMul: xllm_activation replaces SiluAndMul.forward_cuda ---
+if _USE_XLLM_ACTIVATION:
+
+    def _xllm_silu_and_mul_forward_cuda(self, x):
+        """Drop-in for SiluAndMul.forward_cuda using xllm_activation kernel."""
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        # Reuse cached output tensor for stable shapes (decode)
+        _cache_key = (output_shape, x.dtype, x.device)
+        _cached = getattr(self, '_out_cache', {}).get(_cache_key)
+        if _cached is not None and _cached.shape == output_shape:
+            out = _cached
+        else:
+            out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+            if not hasattr(self, '_out_cache'):
+                self._out_cache = {}
+            self._out_cache[_cache_key] = out
+        _xllm_activation.silu_and_mul(out, x)
+        return out
+
+    SiluAndMul.forward_cuda = _xllm_silu_and_mul_forward_cuda
+    logger.info("xllm_activation PATCHED — SiluAndMul.forward_cuda → xllm CUDA kernel")
+
+
+# --- 3. Cache ops: xllm_cache (reshape_paged_cache / block_copy) ---
+# Replace ixformer vllm_cache_ops_reshape_and_cache with xllm_cache.
+# xllm_cache signature: reshape_paged_cache(slot_ids, keys, values, kc, vc)
+# vllm calls:           ops.reshape_and_cache(key, value, kc, vc, slot_mapping, ...)
+# Difference: arg order, slot_ids must be int32.
+if _USE_XLLM_CACHE:
+    import vllm._custom_ops as _vllm_ops
+
+    _orig_reshape_and_cache = _vllm_ops.reshape_and_cache
+
+    def _xllm_reshape_and_cache(key, value, key_cache, value_cache,
+                                slot_mapping, kv_cache_dtype="auto",
+                                k_scale=1.0, v_scale=1.0):
+        slot_ids = slot_mapping.flatten().to(torch.int32)
+        _xllm_cache.reshape_paged_cache(slot_ids, key, value,
+                                        key_cache, value_cache)
+
+    _vllm_ops.reshape_and_cache = _xllm_reshape_and_cache
+    logger.info("xllm_cache PATCHED — reshape_and_cache → xllm CUDA kernel")
+
+
+# --- 4. RoPE: xllm_rope ---
+# Qwen3.5 uses interleaved multi-axis RoPE (MRoPE) with partial rotary
+# factor 0.25.  The xllm_rope kernel accepts (positions, query, key,
+# cos_sin_cache, is_neox) and modifies q/k in-place, but it applies RoPE
+# to the FULL head dim.  Qwen3.5's forward only rotates the first 25% of
+# each head (rotary_dim = 64 out of 256), then concatenates the unrotated
+# tail.  The multi-axis interleaving of cos/sin across T/H/W axes also
+# requires Python-level assembly before calling any kernel.
+#
+# A safe replacement would need to:
+#  1. Assemble the interleaved cos_sin for the partial dim.
+#  2. Call xllm_rope on just the rotary_dim slice of q and k.
+#  3. Skip the concat step since the kernel modifies in-place.
+#
+# This is doable but requires careful per-position indexing that differs
+# from the standard "positions → cos_sin_cache[positions]" pattern.
+# We mark this as TODO and leave the current _apply_rotary_emb path.
+if _USE_XLLM_ROPE:
+    logger.info("xllm_rope LOADED but NOT PATCHED — Qwen3.5 interleaved MRoPE "
+                "requires adapter; using PyTorch _apply_rotary_emb fallback")
+
+
 # ix_fused_moe: full 7-step fused MoE pipeline via ixformer C++ API
 # Source: xllm/core/layers/ilu/fused_moe.cpp → ix_moe_bridge.so
 try:
@@ -1578,10 +1749,6 @@ class Qwen3_5FullAttention(nn.Module):
             q = qg[:, :, :self.head_dim].reshape(total_tokens, -1)
             gate = qg[:, :, self.head_dim:].reshape(total_tokens, -1)
 
-            q = self.q_norm.forward_cuda(
-                q.view(total_tokens, self.local_num_heads, self.head_dim)
-                .contiguous()).view(total_tokens, -1)
-
             # Select the one rank-local KV head before k_norm and RoPE.
             if self.q_per_kv_global is not None:
                 tp_rank = get_tensor_model_parallel_rank()
@@ -1592,10 +1759,53 @@ class Qwen3_5FullAttention(nn.Module):
                 v = (v.view(total_tokens, self.proj_kv_heads, self.head_dim)
                       [:, kv_idx, :].contiguous())
 
-            k = self.k_norm.forward_cuda(
+            # --- Fused QK-Norm + RoPE path (saves 4 kernel launches per layer) ---
-                k.view(total_tokens, self.local_num_kv_heads, self.head_dim)
+            # Only for 1D positions (decode / text-only prefill).
-                .contiguous()).view(total_tokens, -1)
+            # 2D MRoPE positions (vision prefill) fall back to separate ops.
-            q, k = self.rotary_emb(positions, q, k)
+            if (_USE_XLLM_FUSED_QKNORM_ROPE
+                    and positions.ndim == 1
+                    and q.is_contiguous() and k.is_contiguous()):
+                # Pack Q, K, V into contiguous [T, (Hq+Hk+Hv)*D] for fused kernel
+                v_flat = v.view(total_tokens, -1)
+                qkv = torch.cat([q, k.view(total_tokens, -1), v_flat], dim=-1)
+                # GemmaRMSNorm weight convention: kernel uses x*w, Gemma uses x*(1+w)
+                _q_ew = getattr(self, '_fused_q_ew', None)
+                if _q_ew is None:
+                    _q_ew = (1.0 + self.q_norm.weight.data.float()).to(
+                        self.q_norm.weight.dtype)
+                    self._fused_q_ew = _q_ew
+                _k_ew = getattr(self, '_fused_k_ew', None)
+                if _k_ew is None:
+                    _k_ew = (1.0 + self.k_norm.weight.data.float()).to(
+                        self.k_norm.weight.dtype)
+                    self._fused_k_ew = _k_ew
+                _xllm_fused_qknorm_rope.fused_qk_norm_rope(
+                    qkv,
+                    self.local_num_heads,
+                    self.local_num_kv_heads,
+                    self.local_num_kv_heads,
+                    self.head_dim,
+                    self.rms_norm_eps,
+                    _q_ew,
+                    _k_ew,
+                    self.rotary_emb.cos_sin_cache,
+                    True,  # interleaved (Qwen3.5 uses interleaved RoPE)
+                    positions.to(torch.int64))
+                # Unpack
+                q_dim = self.local_num_heads * self.head_dim
+                k_dim = self.local_num_kv_heads * self.head_dim
+                q = qkv[:, :q_dim]
+                k = qkv[:, q_dim:q_dim + k_dim]
+                # v is untouched by fused kernel, keep original
+            else:
+                # Fallback: separate q_norm, k_norm, rotary_emb
+                q = self.q_norm.forward_cuda(
+                    q.view(total_tokens, self.local_num_heads, self.head_dim)
+                    .contiguous()).view(total_tokens, -1)
+                k = self.k_norm.forward_cuda(
+                    k.view(total_tokens, self.local_num_kv_heads, self.head_dim)
+                    .contiguous()).view(total_tokens, -1)
+                q, k = self.rotary_emb(positions, q, k)
 
         with bi100_timer("full_attn.attention"):
             with bi100_timer(f"L{self.layer_idx}.full_attn"):
@@ -2889,4 +3099,5 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
                 weight_loader(param, loaded_weight)
         _bi100_model_trace(
             f"MoE load_weights complete items={loaded_count} "
-            f"vision_items={vision_loaded_count}")
+            f"vision_items={vision_loaded_count}")
+print("[qwen3_5] module load COMPLETE", file=sys.stderr, flush=True)

qwen3_6_scripts/serving_chat.py

@@ -410,9 +410,24 @@ class OpenAIServingChat(OpenAIServing):
 
             if request.max_tokens is not None:
                 request.max_tokens = min(request.max_tokens, _adaptive_cap)
+            else:
+                # [BI100] max_tokens=None means "no limit" in the OpenAI API.
+                # Cap to _adaptive_cap to prevent generation timeout (300s).
+                request.max_tokens = _adaptive_cap
 
             default_max_tokens = min(self.max_model_len - _prompt_len,
                                      _adaptive_cap)
+
+            # [BI100] Qwen3 tool calling fix: greedy decoding (temperature=0)
+            # causes thinking to degrade, so the model never generates
+            # <tool_call> tags.  Force temperature≥0.6 when tools are active.
+            if (request.tools and request.tool_choice in ("auto", None)
+                    and (request.temperature is None
+                         or request.temperature < 0.6)):
+                request.temperature = 0.6
+                if request.top_p is None or request.top_p > 0.95:
+                    request.top_p = 0.95
+
             if request.use_beam_search:
                 sampling_params = request.to_beam_search_params(
                     default_max_tokens)

tests/python/test_qwen3_5_parallel.py

@@ -0,0 +1,482 @@
+"""Parallel-layout tests for the Qwen3.5 Python model (DP/EP).
+
+Ported from xLLM upstream commit 78aa2a85 (PR #2258)
+  Original: tests/python/test_deepseek_v32_parallel.py
+  Adapted:  Qwen3.5 MoeSparseBlock with 256 routed experts + shared expert
+
+Test coverage maps to the issue's test cases:
+  TC-01  MoE expert routing DP isolation
+  TC-02  CUDA/ACL graph capture per DP rank
+  TC-03  DP broadcast/gather correctness
+  TC-04  Executor DP initialization
+  TC-05  End-to-end DP parallel test suite
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+# Insert repo root so "python.*" resolves to project_6/python/*, not stdlib
+_repo_root = str(Path(__file__).resolve().parents[2])
+if _repo_root not in sys.path:
+    sys.path.insert(0, _repo_root)
+
+import pytest
+import torch
+import torch.nn as nn
+
+# Block the repo-root __init__.py from pulling in vllm
+sys.modules.setdefault("project_6", MagicMock())
+
+from python.layers.fused_moe import DPAwareMoEMixin  # noqa: E402
+from python.models.qwen3_5 import (  # noqa: E402
+    _dp_all_gather,
+    _dp_all_gather_variable,
+    configure_dp,
+    dp_forward_moe_wrapper,
+)
+from python.model_executor.runners.decode_cuda_graph import (  # noqa: E402
+    DecodeCudaGraphRunner,
+    _decode_bucket,
+    _decode_graph_buckets,
+)
+from python.model_executor.runners.decode_acl_graph import (  # noqa: E402
+    DecodeAclGraphRunner,
+)
+from python.model_executor.executor import ModelExecutor  # noqa: E402
+from python.attention.backend import (  # noqa: E402
+    DPBackendConfig, create_attention_backend, register_backend,
+)
+
+
+# ---------------------------------------------------------------------------
+# Mock MoeSparseBlock for testing
+# ---------------------------------------------------------------------------
+
+
+class MockMoeSparseBlock(nn.Module):
+    """Minimal mock of Qwen3_5MoeSparseBlock for DP testing."""
+
+    def __init__(self, hidden_size: int = 64, num_experts: int = 256):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_experts = num_experts
+        self.dp_size = 1
+        self.dp_rank = 0
+        # Dummy weight so next(model.parameters()) works
+        self.dummy = nn.Parameter(torch.zeros(1))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Identity forward — just returns input (MoE logic mocked)."""
+        return hidden_states
+
+
+def _make_mock_moe(dp_size: int = 1, dp_rank: int = 0) -> MockMoeSparseBlock:
+    moe = MockMoeSparseBlock()
+    configure_dp(moe, dp_size, dp_rank)
+    return moe
+
+
+def _mock_metadata(
+    dp_token_counts=(4,),
+    is_prefill=False,
+    is_chunked_prefill=False,
+    dp_is_decode=None,
+    execution_state=None,
+):
+    metadata = SimpleNamespace(
+        dp_token_counts=dp_token_counts,
+        is_prefill=is_prefill,
+        is_chunked_prefill=is_chunked_prefill,
+        execution_state=execution_state,
+    )
+    if dp_is_decode is not None:
+        metadata.dp_is_decode = dp_is_decode
+    return metadata
+
+
+# ---------------------------------------------------------------------------
+# TC-01: MoE expert routing DP isolation
+# ---------------------------------------------------------------------------
+
+
+class TestMoEDPIsolation:
+    """Verify that DP replicas route experts independently."""
+
+    def test_dp1_no_gather(self):
+        """With dp_size=1, no all-gather should occur."""
+        moe = _make_mock_moe(dp_size=1)
+        hidden = torch.randn(4, 64)
+        metadata = _mock_metadata(dp_token_counts=(4,))
+
+        result = dp_forward_moe_wrapper(
+            moe, hidden, moe.forward, metadata
+        )
+        assert result.shape == hidden.shape
+
+    def test_dp2_calls_gather(self):
+        """With dp_size=2, all-gather should be invoked."""
+        moe = _make_mock_moe(dp_size=2, dp_rank=0)
+        hidden = torch.randn(3, 64)
+        metadata = _mock_metadata(
+            dp_token_counts=(3, 4), execution_state="graph"
+        )
+
+        with patch(
+            "python.models.qwen3_5._dp_all_gather"
+        ) as mock_gather:
+            mock_gather.side_effect = lambda x, **kw: x.repeat(
+                kw.get("world_size", 1), *([1] * (x.dim() - 1))
+            )
+            result = dp_forward_moe_wrapper(
+                moe, hidden, moe.forward, metadata
+            )
+            mock_gather.assert_called_once()
+            call_kwargs = mock_gather.call_args[1]
+            assert call_kwargs["dim"] == 0
+            assert call_kwargs["world_size"] == 2
+            assert call_kwargs["group_name"] == "dp"
+
+    def test_dp2_different_inputs_route_independently(self):
+        """Different DP replicas with different inputs produce different
+        routing decisions (verified by output shapes being correct)."""
+        moe_r0 = _make_mock_moe(dp_size=2, dp_rank=0)
+        moe_r1 = _make_mock_moe(dp_size=2, dp_rank=1)
+
+        hidden_r0 = torch.randn(3, 64)
+        hidden_r1 = torch.randn(4, 64)
+
+        metadata = _mock_metadata(
+            dp_token_counts=(3, 4), dp_is_decode=(1, 1)
+        )
+
+        # Rank 0
+        result_r0 = dp_forward_moe_wrapper(
+            moe_r0, hidden_r0, moe_r0.forward, metadata
+        )
+        assert result_r0.shape[0] == 3
+
+        # Rank 1
+        result_r1 = dp_forward_moe_wrapper(
+            moe_r1, hidden_r1, moe_r1.forward, metadata
+        )
+        assert result_r1.shape[0] == 4
+
+
+# ---------------------------------------------------------------------------
+# TC-02: CUDA/ACL graph capture per DP rank
+# ---------------------------------------------------------------------------
+
+
+class TestCudaGraphDPCapture:
+    """CUDA graph bucket computation and graph key with DP."""
+
+    def test_bucket_dp1(self):
+        buckets = _decode_graph_buckets(32, dp_size=1)
+        assert 1 in buckets
+        assert buckets[-1] == 32
+
+    def test_bucket_dp2_halves_capacity(self):
+        buckets = _decode_graph_buckets(32, dp_size=2)
+        # max_local_batch = ceil(32/2) = 16
+        assert buckets[-1] <= 16
+
+    def test_graph_key_dp1(self):
+        runner = DecodeCudaGraphRunner(
+            nn.Linear(1, 1), torch.device("cpu"), max_batch=16
+        )
+        input_ids = torch.zeros(4, dtype=torch.int32)
+        key = runner.graph_key(input_ids)
+        assert key is not None
+        padded, counts = key
+        assert padded == _decode_bucket(4)
+        assert counts == (padded,)
+
+    def test_graph_key_dp2(self):
+        runner = DecodeCudaGraphRunner(
+            nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
+            dp_size=2, dp_rank=0
+        )
+        input_ids = torch.zeros(4, dtype=torch.int32)
+        key = runner.graph_key(
+            input_ids,
+            dp_token_counts=(4, 3),
+            dp_is_decode=(1, 1),
+        )
+        assert key is not None
+        padded, counts = key
+        assert len(counts) == 2
+        assert counts[0] == counts[1] == padded
+
+    def test_graph_key_dp2_exceeds_capacity_returns_none(self):
+        runner = DecodeCudaGraphRunner(
+            nn.Linear(1, 1), torch.device("cpu"), max_batch=8,
+            dp_size=2, dp_rank=0
+        )
+        input_ids = torch.zeros(100, dtype=torch.int32)
+        key = runner.graph_key(
+            input_ids,
+            dp_token_counts=(100, 100),
+            dp_is_decode=(1, 1),
+        )
+        assert key is None
+
+    def test_graph_key_dp_mismatched_counts_raises(self):
+        runner = DecodeCudaGraphRunner(
+            nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
+            dp_size=2, dp_rank=0
+        )
+        input_ids = torch.zeros(4, dtype=torch.int32)
+        with pytest.raises(RuntimeError, match="dp_token_counts"):
+            runner.graph_key(input_ids, dp_token_counts=(4,))
+
+
+class TestAclGraphDPCapture:
+    """ACL graph runner DP validation."""
+
+    def test_acl_dp2_init(self):
+        runner = DecodeAclGraphRunner(
+            nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
+            dp_size=2, dp_rank=0
+        )
+        assert runner.dp_size == 2
+        assert runner.dp_rank == 0
+        assert runner.max_batch == 16  # ceil(32/2)
+
+    def test_acl_dp_validate_wrong_counts_raises(self):
+        runner = DecodeAclGraphRunner(
+            nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
+            dp_size=2, dp_rank=0
+        )
+        with pytest.raises(RuntimeError, match="dp_token_counts"):
+            runner._validate_dp_token_counts((4,))
+
+    def test_acl_dp1_validate_accepts_none(self):
+        runner = DecodeAclGraphRunner(
+            nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
+            dp_size=1, dp_rank=0
+        )
+        # dp_size=1: validation is a no-op
+        runner._validate_dp_token_counts(None)
+
+
+# ---------------------------------------------------------------------------
+# TC-03: DP broadcast/gather correctness
+# ---------------------------------------------------------------------------
+
+
+class TestDPBroadcastGather:
+    """Verify that gather + scatter preserves local token identity."""
+
+    def test_padded_gather_rank0_recovers_local(self):
+        """dp_rank=0: after padded gather → slice, output shape matches input."""
+        moe = _make_mock_moe(dp_size=2, dp_rank=0)
+        hidden = torch.randn(3, 64)
+        metadata = _mock_metadata(
+            dp_token_counts=(3, 4), execution_state="graph"
+        )
+
+        with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
+            # Simulate all_gather: pad to 4, gather 2 replicas → [8, 64]
+            mock_g.side_effect = lambda x, **kw: x.repeat(
+                kw.get("world_size", 1), *([1] * (x.dim() - 1))
+            )
+            result = dp_forward_moe_wrapper(
+                moe, hidden, moe.forward, metadata
+            )
+        # dp_rank=0, padded=4: narrow(0, 0, 3) → [3, 64]
+        assert result.shape[0] == 3
+
+    def test_padded_gather_rank1_recovers_local(self):
+        """dp_rank=1: output is sliced from the second replica's region."""
+        moe = _make_mock_moe(dp_size=2, dp_rank=1)
+        hidden = torch.randn(4, 64)
+        metadata = _mock_metadata(
+            dp_token_counts=(3, 4), execution_state="graph"
+        )
+
+        with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
+            mock_g.side_effect = lambda x, **kw: x.repeat(
+                kw.get("world_size", 1), *([1] * (x.dim() - 1))
+            )
+            result = dp_forward_moe_wrapper(
+                moe, hidden, moe.forward, metadata
+            )
+        # dp_rank=1, padded=4: narrow(0, 4, 4) → [4, 64]
+        assert result.shape[0] == 4
+
+    def test_compact_gather_rank0(self):
+        """Eager decode path uses compact gather; rank 0 gets first slice."""
+        moe = _make_mock_moe(dp_size=2, dp_rank=0)
+        hidden = torch.randn(3, 64)
+        compact = torch.randn(7, 64)  # 3 + 4 tokens
+        metadata = _mock_metadata(
+            dp_token_counts=(3, 4), dp_is_decode=(1, 1)
+        )
+
+        with patch("python.models.qwen3_5._dp_all_gather_variable") as mock_gv:
+            mock_gv.return_value = compact
+            result = dp_forward_moe_wrapper(
+                moe, hidden, moe.forward, metadata
+            )
+            mock_gv.assert_called_once()
+        # dp_rank=0: offset=0, narrow(0, 0, 3)
+        assert result.shape[0] == 3
+
+    def test_compact_gather_rank1(self):
+        """Eager decode path: rank 1 gets second slice."""
+        moe = _make_mock_moe(dp_size=2, dp_rank=1)
+        hidden = torch.randn(4, 64)
+        compact = torch.randn(7, 64)
+        metadata = _mock_metadata(
+            dp_token_counts=(3, 4), dp_is_decode=(1, 1)
+        )
+
+        with patch("python.models.qwen3_5._dp_all_gather_variable") as mock_gv:
+            mock_gv.return_value = compact
+            result = dp_forward_moe_wrapper(
+                moe, hidden, moe.forward, metadata
+            )
+        # dp_rank=1: offset=sum([3])=3, narrow(0, 3, 4)
+        assert result.shape[0] == 4
+
+
+# ---------------------------------------------------------------------------
+# TC-04: Executor DP initialization
+# ---------------------------------------------------------------------------
+
+
+class TestExecutorDPInit:
+    """Verify ModelExecutor reads and validates DP config."""
+
+    def _make_model(self):
+        return nn.Linear(10, 10)
+
+    def test_dp_size_stored(self):
+        model = self._make_model()
+        executor = ModelExecutor(model, {"dp_size": 2, "dp_rank": 0}, 32)
+        assert executor.dp_size == 2
+
+    def test_dp1_default(self):
+        model = self._make_model()
+        executor = ModelExecutor(model, {}, 32)
+        assert executor.dp_size == 1
+
+    def test_dp_with_unsupported_backend_raises(self):
+        model = self._make_model()
+        with pytest.raises(NotImplementedError, match="data parallel"):
+            ModelExecutor(
+                model,
+                {"dp_size": 2, "dp_rank": 0, "python_graph_backend": "inductor"},
+                32,
+            )
+
+    def test_dp_with_cudagraphs_accepted(self):
+        model = self._make_model()
+        executor = ModelExecutor(
+            model,
+            {"dp_size": 2, "dp_rank": 0, "python_graph_backend": "cudagraphs"},
+            32,
+        )
+        assert executor.decode_graph_runner is not None
+
+    def test_dp_with_aclgraph_accepted(self):
+        model = self._make_model()
+        executor = ModelExecutor(
+            model,
+            {"dp_size": 2, "dp_rank": 0, "python_graph_backend": "aclgraph"},
+            32,
+        )
+        assert executor.decode_graph_runner is not None
+
+
+# ---------------------------------------------------------------------------
+# TC-05: End-to-end DP parallel test suite
+# ---------------------------------------------------------------------------
+
+
+class TestEndToEndDP:
+    """Integration tests combining executor + MoE + graph runner."""
+
+    def test_2way_dp_moe_shapes(self):
+        """2-way DP: both ranks produce correct output shapes."""
+        for rank in (0, 1):
+            moe = _make_mock_moe(dp_size=2, dp_rank=rank)
+            local_tokens = 3 if rank == 0 else 4
+            hidden = torch.randn(local_tokens, 64)
+            metadata = _mock_metadata(
+                dp_token_counts=(3, 4), dp_is_decode=(1, 1)
+            )
+            result = dp_forward_moe_wrapper(
+                moe, hidden, moe.forward, metadata
+            )
+            assert result.shape[0] == local_tokens
+
+    def test_4way_dp_moe_shapes(self):
+        """4-way DP: all ranks produce correct output shapes."""
+        counts = (2, 3, 4, 5)
+        for rank in range(4):
+            moe = _make_mock_moe(dp_size=4, dp_rank=rank)
+            hidden = torch.randn(counts[rank], 64)
+            metadata = _mock_metadata(
+                dp_token_counts=counts, dp_is_decode=(1, 1, 1, 1)
+            )
+            result = dp_forward_moe_wrapper(
+                moe, hidden, moe.forward, metadata
+            )
+            assert result.shape[0] == counts[rank]
+
+    def test_dp_plus_tp_cuda_graph_buckets(self):
+        """DP=2, TP=4 on 8 devices: graph buckets respect DP-reduced capacity."""
+        # max_batch=64 across 2 DP replicas → 32 per replica
+        buckets = _decode_graph_buckets(64, dp_size=2)
+        assert buckets[-1] <= 32
+
+    def test_varying_batch_sizes_padded_path(self):
+        """Different batch sizes per DP rank use padded gather."""
+        moe = _make_mock_moe(dp_size=2, dp_rank=0)
+        for local_size in (1, 5, 8, 16):
+            other_size = local_size + 2
+            hidden = torch.randn(local_size, 64)
+            metadata = _mock_metadata(
+                dp_token_counts=(local_size, other_size),
+                execution_state="graph",
+            )
+            with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
+                mock_g.side_effect = lambda x, **kw: x.repeat(
+                    kw.get("world_size", 1), *([1] * (x.dim() - 1))
+                )
+                result = dp_forward_moe_wrapper(
+                    moe, hidden, moe.forward, metadata
+                )
+            assert result.shape[0] == local_size
+
+
+# ---------------------------------------------------------------------------
+# TC-07: Existing tests unbroken (sanity)
+# ---------------------------------------------------------------------------
+
+
+class TestBackwardCompat:
+    """Verify dp_size=1 (default) doesn't change existing behaviour."""
+
+    def test_dp1_identity(self):
+        moe = _make_mock_moe(dp_size=1)
+        hidden = torch.randn(8, 64)
+        metadata = _mock_metadata(dp_token_counts=(8,))
+        result = dp_forward_moe_wrapper(
+            moe, hidden, moe.forward, metadata
+        )
+        assert torch.equal(result, hidden)
+
+    def test_executor_dp1_no_graph_runner(self):
+        model = nn.Linear(10, 10)
+        executor = ModelExecutor(model, {}, 32)
+        assert executor.decode_graph_runner is None
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

vllm/model_executor/models/mamba_cache.py

@@ -10,11 +10,11 @@ class MambaCacheManager:
     def __init__(self, dtype, num_mamba_layers, max_batch_size,
                  conv_state_shape, temporal_state_shape):
 
-        conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
+        conv_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
                                  conv_state_shape,
                                  dtype=dtype,
                                  device="cuda")
-        temporal_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
+        temporal_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
                                      temporal_state_shape,
                                      dtype=dtype,
                                      device="cuda")
@@ -101,6 +101,8 @@ class MambaCacheManager:
             self._move_out_if_already_occupied(
                 index=destination_index,
                 all_occupied_indices=all_occupied_indices)
+            for cache_t in self.mamba_cache:
+                cache_t[:, destination_index].zero_()
             self.mamba_cache_indices_mapping[cur_rid] = {
                 seq_id: destination_index
             }
@@ -206,7 +208,10 @@ class MambaCacheManager:
                                    finished_seq_groups_req_ids: List[str]):
         for req_id in finished_seq_groups_req_ids:
             if req_id in self.mamba_cache_indices_mapping:
-                self.mamba_cache_indices_mapping.pop(req_id)
+                seq_mapping = self.mamba_cache_indices_mapping.pop(req_id)
+                for cache_idx in seq_mapping.values():
+                    for cache_t in self.mamba_cache:
+                        cache_t[:, cache_idx].zero_()
 
     def _first_free_index_in_mamba_cache(
             self, indices_range: Optional[List[int]] = None) -> int:
@@ -219,4 +224,4 @@ class MambaCacheManager:
             if i not in all_occupied_indices:
                 return i
         raise Exception("Couldn't find a free spot in the mamba cache! This"
-                        "should never happen")
+                        "should never happen")

vllm/sequence.py

@@ -119,6 +119,7 @@ class RequestMetrics:
     scheduler_time: Optional[float] = None
     model_forward_time: Optional[float] = None
     model_execute_time: Optional[float] = None
+    num_cached_tokens: Optional[int] = None
 
 
 class SequenceDataDelta(
@@ -527,6 +528,11 @@ class Sequence:
         self._last_output_token_ids_offset = output_len
 
         # Return new tokens
+        if num_new_tokens == 0:
+            # During chunked prefill steps with no output yet, num_new_tokens=0.
+            # Python's [-0:] == [0:] returns the ENTIRE list — guard against this.
+            return []
+
         if num_new_tokens == 1:
             # Optimization for single decode token case
             # (which is what we have most of the time)
@@ -935,6 +941,12 @@ class SequenceGroupMetadataDelta(
     computed_block_nums: Optional[List[int]] = None
     state: Optional[SequenceGroupState] = msgspec.field(
         default_factory=lambda: SequenceGroupState())
+    # BI100 hybrid prefix-cache actions. Fields are appended for msgspec wire
+    # compatibility with the pre-existing array-like structure.
+    gdn_restore_key: Optional[Tuple[int, bytes]] = None
+    gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
+    gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
+    gdn_segment_offsets: Optional[List[int]] = None
 
 
 class SequenceGroupMetadata(
@@ -1000,6 +1012,12 @@ class SequenceGroupMetadata(
     # Zero means speculative decoding is disabled for some reasons.
     # TODO: We should maintain this states out of the sequence group.
     num_speculative_tokens: Optional[int] = None
+    # BI100 hybrid prefix-cache actions. These are internal scheduler-to-worker
+    # metadata and never surface through the OpenAI API.
+    gdn_restore_key: Optional[Tuple[int, bytes]] = None
+    gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
+    gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
+    gdn_segment_offsets: Optional[List[int]] = None
 
     def __post_init__(self):
         if self.seq_data is not None and self.token_chunk_size is None:
@@ -1046,6 +1064,14 @@ class SequenceGroupMetadata(
         self.token_chunk_size = sequence_group_metadata_delta.token_chunk_size
         self.do_sample = sequence_group_metadata_delta.do_sample
         self.is_prompt = sequence_group_metadata_delta.is_prompt
+        self.computed_block_nums = (
+            sequence_group_metadata_delta.computed_block_nums)
+        self.gdn_restore_key = sequence_group_metadata_delta.gdn_restore_key
+        self.gdn_capture_points = (
+            sequence_group_metadata_delta.gdn_capture_points)
+        self.gdn_evict_keys = sequence_group_metadata_delta.gdn_evict_keys
+        self.gdn_segment_offsets = (
+            sequence_group_metadata_delta.gdn_segment_offsets)
 
     def finish_step(self) -> None:
         assert self.state is not None

vllm/worker/worker.py

@@ -216,6 +216,13 @@ class Worker(LocalOrDistributedWorkerBase):
         """
         # Profile the memory usage of the model and get the maximum number of
         # cache blocks that can be allocated with the remaining free memory.
+        # PRD: skip profile_run when num_gpu_blocks_override is set
+        _ovr = getattr(self.cache_config, 'num_gpu_blocks_override', None)
+        if _ovr is not None and _ovr > 0:
+            logger.info("Skipping profile_run -- num_gpu_blocks_override=%d", _ovr)
+            _cbs = self.get_cache_block_size_bytes()
+            _cpu = self.cache_config.swap_space_bytes // _cbs if _cbs > 0 else 256
+            return int(_ovr), int(_cpu)
         torch.cuda.empty_cache()
 
         # Execute a forward pass with dummy inputs to profile the memory usage