fix: GLM-5.2 w/ DCP+DSpark for SM120

aphrodite/config/speculative.py

@@ -32,6 +32,22 @@
 
 logger = init_logger(__name__)
 
+
+def resolve_draft_kv_cache_dtype(
+    speculative_config: "SpeculativeConfig",
+    target_cache_dtype: CacheDType,
+) -> CacheDType:
+    """Resolve the draft cache dtype without leaking target-only layouts."""
+    if speculative_config.kv_cache_dtype is not None:
+        return speculative_config.kv_cache_dtype
+
+    draft_model_config = speculative_config.draft_model_config
+    if target_cache_dtype == "fp8_ds_mla" and not draft_model_config.use_mla:
+        return "fp8_e4m3"
+
+    return target_cache_dtype
+
+
 MTPModelTypes = Literal[
     "deepseek_mtp",
     "mimo_mtp",

aphrodite/model_executor/kernels/attention/dsa/sm120_indexer.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""SM120 FP8 sparse-attention indexer score kernels."""
+
+from functools import cache
+
+import torch
+
+from aphrodite.platforms import current_platform
+from aphrodite.triton_utils import tl, triton
+
+
+@cache
+def use_sm120_dsa_indexer() -> bool:
+    """Whether the native SM120 sparse-indexer kernels should be used."""
+    return current_platform.is_cuda() and current_platform.is_device_capability_family(120)
+
+
+@triton.jit
+def _sm120_fp8_mqa_logits_kernel(
+    q_ptr,
+    k_ptr,
+    k_scale_ptr,
+    weights_ptr,
+    starts_ptr,
+    ends_ptr,
+    logits_ptr,
+    num_kv_tokens,
+    stride_q_m: tl.int64,
+    stride_q_h: tl.int64,
+    stride_q_d: tl.int64,
+    stride_k_n: tl.int64,
+    stride_k_d: tl.int64,
+    stride_w_m: tl.int64,
+    stride_w_h: tl.int64,
+    stride_o_m: tl.int64,
+    NUM_HEADS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_KV: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tile = tl.program_id(1)
+    kv = tile * BLOCK_KV + tl.arange(0, BLOCK_KV)
+    heads = tl.arange(0, NUM_HEADS)[:, None]
+    dims = tl.arange(0, HEAD_DIM)
+
+    start = tl.load(starts_ptr + row)
+    end = tl.load(ends_ptr + row)
+    valid = (kv >= start) & (kv < end) & (kv < num_kv_tokens)
+
+    q = tl.load(
+        q_ptr + row * stride_q_m + heads * stride_q_h + dims[None, :] * stride_q_d,
+    )
+    k = tl.load(
+        k_ptr + dims[:, None] * stride_k_d + kv[None, :] * stride_k_n,
+        mask=valid[None, :],
+        other=0.0,
+    )
+    scale = tl.load(k_scale_ptr + kv, mask=valid, other=0.0).to(tl.float32)
+    weights = tl.load(weights_ptr + row * stride_w_m + heads * stride_w_h).to(tl.float32)
+
+    scores = tl.dot(q, k, input_precision="ieee").to(tl.float32)
+    scores = tl.maximum(scores * scale[None, :], 0.0)
+    scores = tl.sum(scores * weights, axis=0)
+    scores = tl.where(valid, scores, -float("inf"))
+    tl.store(logits_ptr + row * stride_o_m + kv, scores, mask=kv < num_kv_tokens)
+
+
+@triton.jit
+def _sm120_fp8_paged_mqa_logits_kernel(
+    q_ptr,
+    k_ptr,
+    k_scale_ptr,
+    weights_ptr,
+    context_lens_ptr,
+    block_tables_ptr,
+    logits_ptr,
+    max_model_len,
+    stride_q_b: tl.int64,
+    stride_q_n: tl.int64,
+    stride_q_h: tl.int64,
+    stride_q_d: tl.int64,
+    stride_k_block: tl.int64,
+    stride_k_token: tl.int64,
+    stride_k_d: tl.int64,
+    stride_s_block: tl.int64,
+    stride_s_token: tl.int64,
+    stride_ctx_b: tl.int64,
+    stride_ctx_n: tl.int64,
+    stride_bt_b: tl.int64,
+    stride_bt_block: tl.int64,
+    stride_w_m: tl.int64,
+    stride_w_h: tl.int64,
+    stride_o_m: tl.int64,
+    NEXT_N: tl.constexpr,
+    NUM_HEADS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    BLOCK_KV: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tile = tl.program_id(1)
+    batch = row // NEXT_N
+    q_index = row % NEXT_N
+    logical = tile * BLOCK_KV + tl.arange(0, BLOCK_KV)
+    heads = tl.arange(0, NUM_HEADS)[:, None]
+    dims = tl.arange(0, HEAD_DIM)
+
+    context_len = tl.load(context_lens_ptr + batch * stride_ctx_b + q_index * stride_ctx_n)
+    valid = (logical < context_len) & (logical < max_model_len)
+    logical_block = logical // PAGE_SIZE
+    block_offset = logical % PAGE_SIZE
+    physical_block = tl.load(
+        block_tables_ptr + batch * stride_bt_b + logical_block * stride_bt_block,
+        mask=valid,
+        other=0,
+    )
+
+    q = tl.load(
+        q_ptr + batch * stride_q_b + q_index * stride_q_n + heads * stride_q_h + dims[None, :] * stride_q_d,
+    )
+    k = tl.load(
+        k_ptr
+        + physical_block[None, :] * stride_k_block
+        + block_offset[None, :] * stride_k_token
+        + dims[:, None] * stride_k_d,
+        mask=valid[None, :],
+        other=0.0,
+    )
+    scale = tl.load(
+        k_scale_ptr + physical_block * stride_s_block + block_offset * stride_s_token,
+        mask=valid,
+        other=0.0,
+    ).to(tl.float32)
+    weights = tl.load(weights_ptr + row * stride_w_m + heads * stride_w_h).to(tl.float32)
+
+    scores = tl.dot(q, k, input_precision="ieee").to(tl.float32)
+    scores = tl.maximum(scores * scale[None, :], 0.0)
+    scores = tl.sum(scores * weights, axis=0)
+    scores = tl.where(valid, scores, -float("inf"))
+    tl.store(logits_ptr + row * stride_o_m + logical, scores, mask=logical < max_model_len)
+
+
+def sm120_fp8_mqa_logits(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    k_scales: torch.Tensor,
+    weights: torch.Tensor,
+    starts: torch.Tensor,
+    ends: torch.Tensor,
+) -> torch.Tensor:
+    """Compute ragged FP8 indexer logits on SM120."""
+    num_rows, num_heads, head_dim = q.shape
+    num_kv_tokens = k.shape[0]
+    logits = torch.empty((num_rows, num_kv_tokens), dtype=torch.float32, device=q.device)
+    block_kv = 64
+    _sm120_fp8_mqa_logits_kernel[(num_rows, triton.cdiv(num_kv_tokens, block_kv))](
+        q,
+        k,
+        k_scales.reshape(-1),
+        weights,
+        starts,
+        ends,
+        logits,
+        num_kv_tokens,
+        *q.stride(),
+        *k.stride(),
+        *weights.stride(),
+        logits.stride(0),
+        NUM_HEADS=num_heads,
+        HEAD_DIM=head_dim,
+        BLOCK_KV=block_kv,
+        num_warps=4,
+        num_stages=2,
+    )
+    return logits
+
+
+def sm120_fp8_paged_mqa_logits(
+    q: torch.Tensor,
+    kv_cache: torch.Tensor,
+    weights: torch.Tensor,
+    context_lens: torch.Tensor,
+    block_tables: torch.Tensor,
+    max_model_len: int,
+) -> torch.Tensor:
+    """Compute paged FP8 indexer logits on SM120."""
+    batch_size, next_n, num_heads, head_dim = q.shape
+    page_size = kv_cache.shape[1]
+    num_pages = kv_cache.shape[0]
+    page_stride = kv_cache.stride(0)
+    cache_bytes = kv_cache.view(torch.uint8)
+
+    # indexer_k_quant_and_cache stores a page in planar form: all
+    # ``page_size * head_dim`` FP8 key bytes first, followed by one FP32 scale
+    # per token. The logical cache tensor has a per-token trailing width, but
+    # slicing that dimension would incorrectly interpret the scale bytes as
+    # interleaved with each key row.
+    k_bytes = torch.as_strided(
+        cache_bytes,
+        size=(num_pages, page_size, head_dim),
+        stride=(page_stride, head_dim, 1),
+    )
+    scale_bytes = torch.as_strided(
+        cache_bytes,
+        size=(num_pages, page_size, 4),
+        stride=(page_stride, 4, 1),
+        storage_offset=page_size * head_dim,
+    )
+    k = k_bytes.view(torch.float8_e4m3fn)
+    k_scales = scale_bytes.view(torch.float32).squeeze(-1)
+    if context_lens.ndim == 1:
+        context_lens = context_lens[:, None].expand(-1, next_n)
+
+    logits = torch.empty(
+        (batch_size * next_n, max_model_len),
+        dtype=torch.float32,
+        device=q.device,
+    )
+    block_kv = 64
+    _sm120_fp8_paged_mqa_logits_kernel[(batch_size * next_n, triton.cdiv(max_model_len, block_kv))](
+        q,
+        k,
+        k_scales,
+        weights,
+        context_lens,
+        block_tables,
+        logits,
+        max_model_len,
+        *q.stride(),
+        *k.stride(),
+        *k_scales.stride(),
+        *context_lens.stride(),
+        *block_tables.stride(),
+        *weights.stride(),
+        logits.stride(0),
+        NEXT_N=next_n,
+        NUM_HEADS=num_heads,
+        HEAD_DIM=head_dim,
+        PAGE_SIZE=page_size,
+        BLOCK_KV=block_kv,
+        num_warps=4,
+        num_stages=2,
+    )
+    return logits

aphrodite/model_executor/layers/attention/attention.py

@@ -231,11 +231,17 @@ def __init__(
         mm_prefix_clamp_sliding_window: bool = False,
         attn_backend: type[AttentionBackend] | None = None,
         head_size_v: int | None = None,
+        cache_num_kv_heads: int | None = None,
         **extra_impl_args,
     ) -> None:
         """
         The KV cache is stored inside this class and is accessed via
         `self.kv_cache`.
+
+        ``cache_num_kv_heads`` can increase the number of KV heads stored in
+        each cache block without changing the number projected by the
+        attention implementation. DFlash and DSpark use this under DCP to
+        replicate the draft KV heads needed by gathered query heads.
         """
         super().__init__()
         sliding_window: int | None
@@ -310,6 +316,10 @@ def __init__(
         self.head_size = head_size
         self.head_size_v = self.head_size if head_size_v is None else head_size_v
         self.num_kv_heads = num_kv_heads
+        self.cache_num_kv_heads = num_kv_heads if cache_num_kv_heads is None else cache_num_kv_heads
+        assert self.cache_num_kv_heads % num_kv_heads == 0, (
+            f"cache_num_kv_heads ({self.cache_num_kv_heads}) must be a multiple of num_kv_heads ({num_kv_heads})"
+        )
         self.sliding_window = sliding_window
         self.has_sink = extra_impl_args.get("sinks") is not None
 
@@ -585,7 +595,7 @@ def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec |
             shared_page = aphrodite_config.cache_config.skip_page_size_padded
             sw_per_token = SlidingWindowSpec(
                 block_size=1,
-                num_kv_heads=self.num_kv_heads,
+                num_kv_heads=self.cache_num_kv_heads,
                 head_size=self.head_size,
                 head_size_v=self.head_size_v,
                 dtype=self.kv_cache_torch_dtype,
@@ -595,7 +605,7 @@ def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec |
             sw_block_size = _largest_kernel_block_within(self.attn_backend, sw_per_token, shared_page, block_size)
             return SlidingWindowSpec(
                 block_size=sw_block_size,
-                num_kv_heads=self.num_kv_heads,
+                num_kv_heads=self.cache_num_kv_heads,
                 head_size=self.head_size,
                 head_size_v=self.head_size_v,
                 dtype=self.kv_cache_torch_dtype,
@@ -612,7 +622,7 @@ def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec |
             tq_config = TurboQuantConfig.from_cache_dtype(self.kv_cache_dtype, self.head_size)
             return TQFullAttentionSpec(
                 block_size=block_size,
-                num_kv_heads=self.num_kv_heads,
+                num_kv_heads=self.cache_num_kv_heads,
                 head_size=self.head_size,
                 head_size_v=self.head_size,
                 dtype=self.kv_cache_torch_dtype,
@@ -621,7 +631,7 @@ def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec |
         else:
             return FullAttentionSpec(
                 block_size=block_size,
-                num_kv_heads=self.num_kv_heads,
+                num_kv_heads=self.cache_num_kv_heads,
                 head_size=self.head_size,
                 head_size_v=self.head_size_v,
                 dtype=self.kv_cache_torch_dtype,

aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -552,6 +552,7 @@ def _apply_quant_method(
         router_logits: torch.Tensor,
         shared_experts_input: torch.Tensor | None,
         input_ids: torch.Tensor | None = None,
+        shared_experts_overlapping: bool = False,
     ) -> tuple[torch.Tensor | None, torch.Tensor]:
         """Run expert routing and the fused MoE kernel via the quant method.
 
@@ -585,10 +586,9 @@ def _apply_quant_method(
                 shared_experts_input=shared_experts_input,
             )
 
-        self._maybe_apply_shared_experts(
-            shared_experts_input,
-            SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
-        )
+        if shared_experts_overlapping:
+            assert self._shared_experts is not None
+            self._shared_experts.wait()
 
         return (
             self._shared_experts.output if self._shared_experts is not None else None,
@@ -606,18 +606,6 @@ def _sequence_parallel_context(self):
         ctx = get_forward_context()
         return ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) if ctx.dp_metadata else nullcontext()
 
-    def _maybe_sync_shared_experts_stream(
-        self,
-        shared_experts_input: torch.Tensor | None,
-    ):
-        # If router/gate provided, then apply it here.
-        # (Note: This code runs only when "overlapped mode" is on to allow
-        #        parallel execution of shared experts with the FusedMoEFactory via
-        #        separate cuda stream)
-        if self._shared_experts is not None:
-            assert shared_experts_input is not None
-            self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
-
     def _maybe_add_zero_expert_output(
         self,
         result: torch.Tensor,
@@ -796,8 +784,11 @@ def _forward_impl(
         # TODO(bnell): this can be removed after MK migration is complete.
         self.routed_experts._ensure_moe_quant_config_init()
 
-        # Sync aux and main stream for shared expert multi-stream overlap.
-        self._maybe_sync_shared_experts_stream(shared_experts_input)
+        # Launch shared experts before routed dispatch so both paths overlap.
+        shared_experts_overlapping = False
+        if self._shared_experts is not None:
+            assert shared_experts_input is not None
+            shared_experts_overlapping = self._shared_experts.maybe_forward_async(shared_experts_input)
 
         # If the Runner holds the gate, apply it after the stream sync,
         # so it can run overlapped with the
@@ -823,6 +814,7 @@ def _forward_impl(
                 router_logits=router_logits,
                 shared_experts_input=shared_experts_input,
                 input_ids=input_ids,
+                shared_experts_overlapping=shared_experts_overlapping,
             )
 
             return self._maybe_combine(