Feat/fast dllm v2 dual cache

diffulex/config.py

@@ -32,6 +32,7 @@
     "llada2_mini_dmax",
 }
 DIFFUSION_GEMMA_MODEL_NAMES = {"diffusion_gemma"}
+FAST_DLLM_V2_MODEL_NAMES = {"fast_dllm_v2"}
 
 def _token_content(token) -> str | None:
     if isinstance(token, dict):
@@ -78,7 +79,7 @@ class Config:
     model: str
     lora_path: str = ""
     model_name: str = "dream"
-    decoding_strategy: str = "d2f"  # "d2f", "multi_bd", "dmax", or "diffusion_gemma"
+    decoding_strategy: str = "d2f"  # "d2f", "multi_bd", "dmax", "fast_dllm_v2", or "diffusion_gemma"
 
     # Sampling Harness
     hf_config: Any | None = None
@@ -162,6 +163,10 @@ class Config:
     diffusion_gemma_confidence_threshold: float = 0.1
     diffusion_gemma_entropy_bound: float = 1.0
 
+    # Fast-dLLM v2 controls. When false, keep native sub-block refinement but
+    # recompute the full 32-token FDv2 block for each refinement step.
+    fdv2_use_block_cache: bool = True
+
     # Profiling
     profiler_config: ProfilerConfig | dict | None = None
 
@@ -184,13 +189,23 @@ def _validate_sampling_mode(self) -> None:
     def is_diffusion_gemma(self) -> bool:
         return self.model_name in DIFFUSION_GEMMA_MODEL_NAMES
 
+    @property
+    def is_fast_dllm_v2(self) -> bool:
+        return self.model_name in FAST_DLLM_V2_MODEL_NAMES
+
     def __post_init__(self):
         if not os.path.isdir(self.model):
             raise ValueError(f"model must be an existing directory, got: {self.model}")
         self.profiler_config = ProfilerConfig.from_value(self.profiler_config)
 
         if self.is_diffusion_gemma:
             self.decoding_strategy = "diffusion_gemma"
+        elif self.is_fast_dllm_v2:
+            # Fast_dLLM v2 supports two useful execution modes:
+            # the specialized dual-cache strategy, and the generic MultiBD
+            # runner for plain baseline comparisons.
+            if self.decoding_strategy not in ("fast_dllm_v2", "multi_bd"):
+                self.decoding_strategy = "fast_dllm_v2"
 
         self.strategy = StrategyConfigRegistry.normalize(self)
 

diffulex/sampler/fast_dllm_v2.py

@@ -1,11 +1,159 @@
 import torch
 
 from diffulex.sampler.auto_sampler import AutoSampler
+from diffulex.sampler.base import DllmSamplerNoShiftBase
 from diffulex.sampler.base import DllmSamplerShiftBase
 
+_FDV2_SUB_BLOCK_REFINE = 1
+_FDV2_FINAL_COMMIT = 2
+
 
 @AutoSampler.register("fast_dllm_v2")
 class FastdLLMV2Sampler(DllmSamplerShiftBase):
+    def _shift_logits(self, logits, last_logit=None):
+        del last_logit
+        if logits.shape[0] == 0:
+            return logits
+        shifted_logits = torch.empty_like(logits)
+        shifted_logits[0, ...] = logits[0, ...]
+        shifted_logits[1:, ...] = logits[:-1, ...]
+        return shifted_logits
+
+    def forward(
+        self,
+        reqs,
+        logits: torch.Tensor,
+        temperatures: torch.Tensor,
+        top_p=None,
+        top_k=None,
+        margin_confidence=False,
+        neg_entropy=False,
+        **kwargs,
+    ):
+        attn_metadata = self.fetch_attn_metadata()
+        split_logits = DllmSamplerNoShiftBase._split_logits_per_req(attn_metadata, reqs, logits)
+
+        accepted_ids_map = {}
+        sampled_tokens_map = {}
+        true_local_ids_map = {}
+        mask_token_rel_ids_map = {}
+        confidence_map = {}
+        initial_confidence_map = {}
+
+        for idx, (temperature, req, req_logits) in enumerate(zip(temperatures, reqs, split_logits)):
+            temperature_value = float(temperature.item()) if torch.is_tensor(temperature) else float(temperature)
+            req_id_str = str(req.req_id)
+            true_local_ids_sub_map = {}
+            accepted_ids_sub_map = {}
+            sampled_tokens_sub_map = {}
+            mask_token_rel_ids_sub_map = {}
+            confidence_sub_map = {}
+            initial_confidence_sub_map = {}
+
+            if req_logits.shape[0] == 0:
+                true_local_ids_map[req_id_str] = true_local_ids_sub_map
+                accepted_ids_map[req_id_str] = accepted_ids_sub_map
+                sampled_tokens_map[req_id_str] = sampled_tokens_sub_map
+                mask_token_rel_ids_map[req_id_str] = mask_token_rel_ids_sub_map
+                confidence_map[req_id_str] = confidence_sub_map
+                initial_confidence_map[req_id_str] = initial_confidence_sub_map
+                continue
+
+            if int(getattr(req, "fdv2_mode", -1)) == _FDV2_FINAL_COMMIT:
+                req.fdv2_pending_next_token_id = int(torch.argmax(req_logits[-1, ...]).item())
+                true_local_ids_map[req_id_str] = true_local_ids_sub_map
+                accepted_ids_map[req_id_str] = accepted_ids_sub_map
+                sampled_tokens_map[req_id_str] = sampled_tokens_sub_map
+                mask_token_rel_ids_map[req_id_str] = mask_token_rel_ids_sub_map
+                confidence_map[req_id_str] = confidence_sub_map
+                initial_confidence_map[req_id_str] = initial_confidence_sub_map
+                continue
+
+            is_fdv2_native_req = hasattr(req, "fdv2_current_sub_block") or int(
+                getattr(req, "fdv2_mode", -1)
+            ) in (_FDV2_SUB_BLOCK_REFINE, _FDV2_FINAL_COMMIT)
+            if is_fdv2_native_req:
+                shifted_logits = self._shift_logits(req_logits)
+            else:
+                last_logits = self._fetch_last_logits(req_logits, req)
+                shifted_logits = DllmSamplerShiftBase._shift_logits(self, req_logits, last_logits)
+
+            if attn_metadata.is_prefill[idx]:
+                candidate_blocks = []
+            elif not hasattr(req, "fdv2_current_sub_block"):
+                candidate_blocks = [block for block in req.dllm_blocks if block.is_active]
+            else:
+                candidate_blocks = [req.fdv2_current_sub_block]
+
+            for block in candidate_blocks:
+                block_mask_relative_ids = (
+                    req.fdv2_block_mask_token_relative_ids(block)
+                    if hasattr(req, "fdv2_block_mask_token_relative_ids")
+                    else list(block.mask_token_relative_ids)
+                )
+                if not block.is_active or not block_mask_relative_ids:
+                    continue
+
+                if attn_metadata.is_prefill[idx]:
+                    local_ids = DllmSamplerNoShiftBase._prefill_mask_token_local_ids(req, block, shifted_logits)
+                    mask_token_logits = shifted_logits[local_ids, ...]
+                elif int(getattr(req, "fdv2_mode", -1)) == _FDV2_SUB_BLOCK_REFINE:
+                    if bool(getattr(req, "fdv2_use_block_cache", True)):
+                        buf_ids = list(block_mask_relative_ids)
+                    else:
+                        buf_offset = int(block.start - req.dllm_block_buffer.first_running_block.start)
+                        buf_ids = [buf_offset + i for i in block_mask_relative_ids]
+                    mask_token_logits = shifted_logits[buf_ids, ...]
+                else:
+                    buf_offset = int(block.start - req.dllm_block_buffer.first_running_block.start)
+                    buf_ids = [buf_offset + i for i in block_mask_relative_ids]
+                    mask_token_logits = shifted_logits[buf_ids, ...]
+
+                confidence, sampled_tokens, initial_confidence = self.sample_tokens(
+                    mask_token_logits,
+                    temperature_value,
+                    top_p=top_p,
+                    top_k=top_k,
+                    neg_entropy=(neg_entropy == "neg_entropy"),
+                    margin_confidence=(margin_confidence == "margin_confidence"),
+                    forbidden_token_ids=[int(block.mask_token_id)],
+                )
+                block_id_str = str(block.block_id)
+                (
+                    accepted_ids_list,
+                    sampled_tokens_list,
+                    confidence_list,
+                    initial_confidence_list,
+                ) = self._materialize_sampled_block(
+                    block,
+                    confidence,
+                    sampled_tokens,
+                    initial_confidence,
+                    **kwargs,
+                )
+                true_local_ids_sub_map[block_id_str] = [block_mask_relative_ids[i] for i in accepted_ids_list]
+                accepted_ids_sub_map[block_id_str] = accepted_ids_list
+                sampled_tokens_sub_map[block_id_str] = sampled_tokens_list
+                mask_token_rel_ids_sub_map[block_id_str] = list(block_mask_relative_ids)
+                confidence_sub_map[block_id_str] = confidence_list
+                initial_confidence_sub_map[block_id_str] = initial_confidence_list
+
+            true_local_ids_map[req_id_str] = true_local_ids_sub_map
+            accepted_ids_map[req_id_str] = accepted_ids_sub_map
+            sampled_tokens_map[req_id_str] = sampled_tokens_sub_map
+            mask_token_rel_ids_map[req_id_str] = mask_token_rel_ids_sub_map
+            confidence_map[req_id_str] = confidence_sub_map
+            initial_confidence_map[req_id_str] = initial_confidence_sub_map
+
+        return self.output_cls(
+            true_local_ids_map=true_local_ids_map,
+            accepted_ids_map=accepted_ids_map,
+            sampled_tokens_map=sampled_tokens_map,
+            mask_token_rel_ids_map=mask_token_rel_ids_map,
+            confidence_map=confidence_map,
+            initial_confidence_map=initial_confidence_map,
+        )
+
     def _compute_accepted_ids(
         self,
         block,
@@ -17,15 +165,19 @@ def _compute_accepted_ids(
         **kwargs,
     ) -> torch.Tensor:
         accept_threshold = block.thresholds.accept_threshold
-        pre_block_complete = block.prev_block.is_semi_complete if block.prev_block else True
         high_conf_indices = torch.where(initial_confidence > accept_threshold)[0]
-        # Keep Dream's shifting behavior: only force a top-1 transfer token
-        # once the previous block is semi-complete (or for the initial block).
+        topk_idx = (
+            torch.topk(confidence, 1)[1]
+            if len(high_conf_indices) == 0
+            else torch.tensor([], device=confidence.device, dtype=torch.long)
+        )
+
+        req = getattr(block, "req", None)
+        is_fdv2_native_step = getattr(req, "fdv2_current_sub_block", None) is block
+        if is_fdv2_native_step:
+            return torch.unique(torch.cat([topk_idx, high_conf_indices]))
+
+        pre_block_complete = block.prev_block.is_semi_complete if block.prev_block else True
         if pre_block_complete:
-            topk_idx = (
-                torch.topk(confidence, 1)[1]
-                if len(high_conf_indices) == 0
-                else torch.tensor([], device=confidence.device, dtype=torch.long)
-            )
             return torch.unique(torch.cat([topk_idx, high_conf_indices]))
         return high_conf_indices

diffulex/strategy/fast_dllm_v2/__init__.py

@@ -0,0 +1,15 @@
+"""Fast-dLLM v2 hierarchical block decoding strategy."""
+
+from diffulex.strategy.fast_dllm_v2.config import FastDLLMV2StrategyConfig
+from diffulex.strategy.fast_dllm_v2.engine.kv_cache_manager import FastDLLMV2KVCacheManager
+from diffulex.strategy.fast_dllm_v2.engine.model_runner import FastDLLMV2ModelRunner
+from diffulex.strategy.fast_dllm_v2.engine.request import FastDLLMV2Req
+from diffulex.strategy.fast_dllm_v2.engine.scheduler import FastDLLMV2Scheduler
+
+__all__ = [
+    "FastDLLMV2StrategyConfig",
+    "FastDLLMV2KVCacheManager",
+    "FastDLLMV2ModelRunner",
+    "FastDLLMV2Req",
+    "FastDLLMV2Scheduler",
+]

diffulex/strategy/fast_dllm_v2/attention/__init__.py

@@ -0,0 +1,13 @@
+from diffulex.strategy.fast_dllm_v2.attention.metadata import (
+    FastDLLMV2AttnMetaData,
+    fetch_fast_dllm_v2_attn_metadata,
+    reset_fast_dllm_v2_attn_metadata,
+    set_fast_dllm_v2_attn_metadata,
+)
+
+__all__ = [
+    "FastDLLMV2AttnMetaData",
+    "fetch_fast_dllm_v2_attn_metadata",
+    "reset_fast_dllm_v2_attn_metadata",
+    "set_fast_dllm_v2_attn_metadata",
+]

diffulex/strategy/fast_dllm_v2/attention/metadata.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import torch
+
+from dataclasses import dataclass
+
+from diffulex.attention.metadata import infer_prefill_flags
+from diffulex.mixin.multi_block.attention_metadata import MultiBlockAttnMetaDataMixin
+
+
+@dataclass
+class FastDLLMV2AttnMetaData(MultiBlockAttnMetaDataMixin):
+    fdv2_cache_only: bool = False
+    fdv2_mode: int = 0
+
+
+_FAST_DLLM_V2_ATTN_METADATA = FastDLLMV2AttnMetaData()
+
+
+def set_fast_dllm_v2_attn_metadata(
+    is_prefill: list[bool] | bool,
+    cu_seqlens_q: torch.Tensor | None = None,
+    cu_seqlens_k: torch.Tensor | None = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    slot_mapping: torch.Tensor | None = None,
+    need_kv_cache_store: bool | None = None,
+    context_lens: torch.Tensor | None = None,
+    page_tables: torch.Tensor | None = None,
+    page_size: int = 32,
+    block_size: int = 32,
+    kv_cache_layout: str = "unified",
+    fdv2_cache_only: bool = False,
+    fdv2_mode: int = 0,
+) -> None:
+    global _FAST_DLLM_V2_ATTN_METADATA
+    has_prefill, all_prefill = infer_prefill_flags(is_prefill)
+    _FAST_DLLM_V2_ATTN_METADATA = FastDLLMV2AttnMetaData(
+        is_prefill=is_prefill if isinstance(is_prefill, list) else [bool(is_prefill)],
+        enforce_eager=False,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        slot_mapping=slot_mapping,
+        need_kv_cache_store_static=need_kv_cache_store,
+        has_prefill_static=has_prefill,
+        all_prefill_static=all_prefill,
+        context_lens=context_lens,
+        page_tables=page_tables,
+        page_size=page_size,
+        block_size=block_size,
+        kv_cache_layout=kv_cache_layout,
+        fdv2_cache_only=bool(fdv2_cache_only),
+        fdv2_mode=int(fdv2_mode),
+    )
+
+
+def fetch_fast_dllm_v2_attn_metadata() -> FastDLLMV2AttnMetaData:
+    return _FAST_DLLM_V2_ATTN_METADATA
+
+
+def reset_fast_dllm_v2_attn_metadata() -> None:
+    global _FAST_DLLM_V2_ATTN_METADATA
+    _FAST_DLLM_V2_ATTN_METADATA = FastDLLMV2AttnMetaData()

diffulex/strategy/fast_dllm_v2/config.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import os
+
+from dataclasses import dataclass
+
+from diffulex.engine.strategy_config_registry import StrategyConfigRegistry
+from diffulex.logger import get_logger
+
+
+logger = get_logger(__name__)
+
+
+@dataclass(frozen=True)
+class FastDLLMV2StrategyConfig:
+    name: str = "fast_dllm_v2"
+    sub_block_size: int = 8
+    block_size: int = 32
+    use_block_cache: bool = True
+
+
+@StrategyConfigRegistry.register("fast_dllm_v2")
+def normalize_fast_dllm_v2_config(config) -> FastDLLMV2StrategyConfig:
+    if config.block_size * config.buffer_size != 32:
+        logger.warning(
+            "Fast-dLLM v2 paper defaults map Diffulex block_size * buffer_size to 32; "
+            "got block_size=%s, buffer_size=%s.",
+            config.block_size,
+            config.buffer_size,
+        )
+    if config.buffer_size <= 1:
+        raise ValueError("decoding_strategy='fast_dllm_v2' requires buffer_size > 1.")
+
+    if config.multi_block_prefix_full:
+        logger.warning("Forcing multi_block_prefix_full=False for decoding_strategy=fast_dllm_v2.")
+    config.multi_block_prefix_full = False
+
+    if not config.enable_prefix_caching:
+        logger.info("Enabling prefix caching for decoding_strategy=fast_dllm_v2.")
+    config.enable_prefix_caching = True
+
+    use_block_cache = bool(getattr(config, "fdv2_use_block_cache", True))
+    env_value = os.environ.get("DIFFULEX_FDV2_USE_BLOCK_CACHE")
+    if env_value is not None:
+        use_block_cache = env_value.strip().lower() not in {"0", "false", "no", "off"}
+
+    return FastDLLMV2StrategyConfig(
+        sub_block_size=int(config.block_size),
+        block_size=int(config.block_size) * int(config.buffer_size),
+        use_block_cache=use_block_cache,
+    )

diffulex/strategy/fast_dllm_v2/engine/__init__.py

@@ -0,0 +1,11 @@
+from diffulex.strategy.fast_dllm_v2.engine.kv_cache_manager import FastDLLMV2KVCacheManager
+from diffulex.strategy.fast_dllm_v2.engine.model_runner import FastDLLMV2ModelRunner
+from diffulex.strategy.fast_dllm_v2.engine.request import FastDLLMV2Req
+from diffulex.strategy.fast_dllm_v2.engine.scheduler import FastDLLMV2Scheduler
+
+__all__ = [
+    "FastDLLMV2KVCacheManager",
+    "FastDLLMV2ModelRunner",
+    "FastDLLMV2Req",
+    "FastDLLMV2Scheduler",
+]