Ver código fonte

fix: add real CosmicRaisins DCP patches + sm120-return-lse.patch

- patches/draft-quant-packed-mapping.patch: fix quantized NextN draft layers
- patches/pr72-1-draft-dcp-config-propagation.patch: DCP config propagation
- patches/pr72-2-glm-dcp-draft-path.patch: GLM DCP draft path + sparse attn
- patches/sm120-return-lse.patch: FlashInfer MLA SM120 return_lse (DCP merge)

All patches are applied by Dockerfile automatically (patch -p1 --forward || true).
Alexander Korolev 3 meses atrás
pai
commit
ff8340a20a

+ 22 - 0
docker-images/glm52-nf3-hybrid/patches/draft-quant-packed-mapping.patch

@@ -0,0 +1,22 @@
+--- a/vllm/model_executor/models/deepseek_mtp.py
++++ b/vllm/model_executor/models/deepseek_mtp.py
+@@ -306,6 +306,19 @@
+         quant_config = _maybe_disable_unserialized_modelopt_fp4_nextn(
+             config, vllm_config, get_draft_quant_config(vllm_config)
+         )
++        if quant_config is not None:
++            # Draft quant configs are built fresh (get_draft_quant_config) and
++            # never pass through configure_quant_config, so fused-module shard
++            # expansion is empty; quant schemes whose targets name the
++            # constituent projections (q_a_proj / kv_a_proj_with_mqa) then
++            # fail to match the fused module and the block silently builds
++            # unquantized, breaking checkpoints with quantized NextN layers.
++            quant_config.packed_modules_mapping.setdefault(
++                "fused_qkv_a_proj", ["q_a_proj", "kv_a_proj_with_mqa"]
++            )
++            quant_config.packed_modules_mapping.setdefault(
++                "gate_up_proj", ["gate_proj", "up_proj"]
++            )
+ 
+         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

+ 35 - 0
docker-images/glm52-nf3-hybrid/patches/pr72-1-draft-dcp-config-propagation.patch

@@ -0,0 +1,35 @@
+diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py
+index 09e15c7b2..9a50c4c8d 100644
+--- a/vllm/v1/spec_decode/llm_base_proposer.py
++++ b/vllm/v1/spec_decode/llm_base_proposer.py
+@@ -1332,10 +1332,29 @@ class SpecDecodeBaseProposer:
+         Subclasses may override to apply additional config changes.
+         """
+         spec_cfg = self.speculative_config
++        draft_parallel_config = spec_cfg.draft_parallel_config
++        # create_draft_parallel_config() does not copy
++        # decode_context_parallel_size from the target, so the draft
++        # attention silently runs with dcp=1 (no q-allgather, no LSE merge)
++        # while its KV/metadata use DCP-local sharded semantics. The V2
++        # runner path (gpu/spec_decode/eagle/utils.py) restores the target
++        # DCP for native MTP drafts; mirror that here for the V1 path.
++        import os as _os
++
++        default_shard_draft = "1" if spec_cfg.method == "mtp" else "0"
++        if _os.environ.get(
++            "VLLM_DCP_SHARD_DRAFT", default_shard_draft
++        ).lower() in ("1", "true", "yes"):
++            draft_parallel_config = replace(
++                draft_parallel_config,
++                decode_context_parallel_size=(
++                    self.vllm_config.parallel_config.decode_context_parallel_size
++                ),
++            )
+         config = replace(
+             self.vllm_config,
+             parallel_config=replace(
+-                spec_cfg.draft_parallel_config,
++                draft_parallel_config,
+                 rank=self.vllm_config.parallel_config.rank,
+             ),
+             model_config=spec_cfg.draft_model_config,

+ 111 - 0
docker-images/glm52-nf3-hybrid/patches/pr72-2-glm-dcp-draft-path.patch

@@ -0,0 +1,111 @@
+diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py
+index e6e6f01cc..2767ebdd2 100644
+--- a/vllm/model_executor/layers/sparse_attn_indexer.py
++++ b/vllm/model_executor/layers/sparse_attn_indexer.py
+@@ -1982,6 +1982,19 @@ class SparseAttnIndexer(CustomOp):
+         self.max_model_len = max_model_len
+         self.max_total_seq_len = max_total_seq_len
+         self.topk_indices_buffer = topk_indices_buffer
++        if topk_scores_buffer is None and topk_indices_buffer is not None:
++            try:
++                from vllm.distributed.parallel_state import get_dcp_group
++
++                dcp_world_size = get_dcp_group().world_size
++            except Exception:
++                dcp_world_size = 1
++            if dcp_world_size > 1 and use_b12x_sparse_indexer():
++                topk_scores_buffer = torch.empty(
++                    topk_indices_buffer.shape,
++                    dtype=torch.float32,
++                    device=topk_indices_buffer.device,
++                )
+         self.topk_scores_buffer = topk_scores_buffer
+         self.skip_k_cache_insert = skip_k_cache_insert
+         self.use_fp4_cache = use_fp4_cache
+diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py
+index 552d6438a..d0fdbaf36 100644
+--- a/vllm/model_executor/models/deepseek_mtp.py
++++ b/vllm/model_executor/models/deepseek_mtp.py
+@@ -346,10 +346,7 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
+                 device=self.device,
+             )
+             topk_scores_buffer = None
+-            if (
+-                vllm_config.parallel_config.decode_context_parallel_size > 1
+-                and use_b12x_sparse_indexer()
+-            ):
++            if vllm_config.parallel_config.decode_context_parallel_size > 1:
+                 topk_scores_buffer = torch.empty(
+                     vllm_config.scheduler_config.max_num_batched_tokens,
+                     topk_tokens,
+diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py
+index cbdc9be96..6260ee159 100644
+--- a/vllm/v1/attention/backends/mla/indexer.py
++++ b/vllm/v1/attention/backends/mla/indexer.py
+@@ -1,7 +1,7 @@
+ # SPDX-License-Identifier: Apache-2.0
+ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ import os
+-from dataclasses import dataclass
++from dataclasses import dataclass, replace
+ 
+ import numpy as np
+ import torch
+@@ -695,6 +695,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
+         use_dcp_local_kv = (
+             self.dcp_world_size > 1
+             and common_attn_metadata.dcp_local_seq_lens is not None
++            and not getattr(self.kv_cache_spec, "dcp_replicated", False)
+         )
+         effective_dcp_world_size = self.dcp_world_size if use_dcp_local_kv else 1
+         effective_dcp_rank = self.dcp_rank if use_dcp_local_kv else 0
+@@ -1021,6 +1022,49 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
+ 
+         return attn_metadata
+ 
++    def build_for_drafting(
++        self,
++        common_attn_metadata: CommonAttentionMetadata,
++        draft_index: int,
++    ) -> DeepseekV32IndexerMetadata:
++        if draft_index > 0:
++            num_tokens = common_attn_metadata.num_actual_tokens
++            num_reqs = min(common_attn_metadata.num_reqs, num_tokens)
++            if num_tokens > 0 and num_reqs > 0:
++                needs_unpadded_decode_view = (
++                    common_attn_metadata.batch_topology is not None
++                    or common_attn_metadata.max_query_len > 1
++                    or common_attn_metadata.num_reqs != num_reqs
++                    or int(common_attn_metadata.query_start_loc_cpu[num_reqs])
++                    != num_tokens
++                )
++                if needs_unpadded_decode_view:
++                    query_start_loc = torch.arange(
++                        num_reqs + 1,
++                        dtype=common_attn_metadata.query_start_loc.dtype,
++                        device=common_attn_metadata.query_start_loc.device,
++                    )
++                    query_start_loc_cpu = torch.arange(
++                        num_reqs + 1,
++                        dtype=common_attn_metadata.query_start_loc_cpu.dtype,
++                        device=common_attn_metadata.query_start_loc_cpu.device,
++                    )
++                    common_attn_metadata = replace(
++                        common_attn_metadata.unpadded(
++                            num_actual_tokens=num_tokens,
++                            num_actual_reqs=num_reqs,
++                        ),
++                        query_start_loc=query_start_loc,
++                        query_start_loc_cpu=query_start_loc_cpu,
++                        max_query_len=1,
++                    )
++
++        return self.build(
++            common_prefix_len=0,
++            common_attn_metadata=common_attn_metadata,
++            fast_build=True,
++        )
++
+ 
+ def build_prefill_chunk_metadata(
+     start_idx: int,

+ 58 - 0
docker-images/glm52-nf3-hybrid/patches/sm120-return-lse.patch

@@ -0,0 +1,58 @@
+--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
++++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
+@@ -32,6 +32,13 @@
+ class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):
+     """SM120 FlashInfer sparse-MLA implementation."""
+ 
++    # glm52-dcp patch: FlashInfer's trtllm_batch_decode_with_kv_cache_mla
++    # accepts return_lse (verified in flashinfer 0.6.14), so this impl CAN
++    # return the decode LSE that DCP's cross-rank merge requires. Upstream
++    # never plumbed it through; see forward_mqa below. Mirrors the parent
++    # FlashInferMLASparseImpl's LSE handling.
++    can_return_lse_for_decode = True
++
+     def __init__(
+         self,
+         num_heads: int,
+@@ -126,8 +133,12 @@
+             ),
+         )
+ 
++        # glm52-dcp patch: size the out buffer from the query's actual head
++        # count, not self.num_heads. Under DCP the MLA decode all-gathers
++        # queries across ranks, so each rank sees the FULL head count (64)
++        # rather than the TP-local slice (16) — self.num_heads is stale here.
+         output = q.new_empty(
+-            (num_actual_toks, self.num_heads, self.kv_lora_rank),
++            (num_actual_toks, q.shape[1], self.kv_lora_rank),
+             dtype=q.dtype,
+         )
+ 
+@@ -153,5 +164,27 @@
+             bmm2_scale=1.0,
+             sparse_mla_top_k=attn_metadata.topk_tokens,
+             kv_scale_format=self.kv_scale_format,
++            # glm52-dcp patch: ask the kernel for the decode LSE when DCP
++            # needs it for the cross-rank merge (base class sets this flag
++            # only when dcp_world_size > 1).
++            return_lse=self.need_to_return_lse_for_decode,
+         )
++        if self.need_to_return_lse_for_decode:
++            # With return_lse=True flashinfer returns (out, lse).
++            assert isinstance(out, tuple), (
++                "expected (out, lse) tuple from flashinfer when return_lse=True"
++            )
++            o, lse = out
++            o = o.view(-1, o.shape[-2], o.shape[-1])
++            # Reuse the parent impl's LSE normalization + empty-row masking
++            # conventions (the shared DCP reducer expects (num_tokens, num_heads)
++            # and -inf LSE on rows with no valid top-k indices).
++            from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
++                FlashInferMLASparseImpl as _ParentImpl,
++            )
++            lse = _ParentImpl._normalize_lse(lse, o.shape[0], o.shape[1])
++            empty_rows = (topk_indices_physical == -1).all(dim=-1)
++            o.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
++            lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
++            return o, lse
+         return out.squeeze(1), None