diff --git a/diffulex/config.py b/diffulex/config.py index 98a41f46..9113837c 100755 --- a/diffulex/config.py +++ b/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,6 +189,10 @@ 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}") @@ -191,6 +200,12 @@ def __post_init__(self): 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) diff --git a/diffulex/sampler/fast_dllm_v2.py b/diffulex/sampler/fast_dllm_v2.py index b6da9cd4..2ee3c5bc 100644 --- a/diffulex/sampler/fast_dllm_v2.py +++ b/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 diff --git a/diffulex/strategy/fast_dllm_v2/__init__.py b/diffulex/strategy/fast_dllm_v2/__init__.py new file mode 100644 index 00000000..d44b6a81 --- /dev/null +++ b/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", +] diff --git a/diffulex/strategy/fast_dllm_v2/attention/__init__.py b/diffulex/strategy/fast_dllm_v2/attention/__init__.py new file mode 100644 index 00000000..9234351d --- /dev/null +++ b/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", +] diff --git a/diffulex/strategy/fast_dllm_v2/attention/metadata.py b/diffulex/strategy/fast_dllm_v2/attention/metadata.py new file mode 100644 index 00000000..b75a0f2e --- /dev/null +++ b/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() diff --git a/diffulex/strategy/fast_dllm_v2/config.py b/diffulex/strategy/fast_dllm_v2/config.py new file mode 100644 index 00000000..8d9d7aed --- /dev/null +++ b/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, + ) diff --git a/diffulex/strategy/fast_dllm_v2/engine/__init__.py b/diffulex/strategy/fast_dllm_v2/engine/__init__.py new file mode 100644 index 00000000..8e2bcd2b --- /dev/null +++ b/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", +] diff --git a/diffulex/strategy/fast_dllm_v2/engine/kv_cache_manager.py b/diffulex/strategy/fast_dllm_v2/engine/kv_cache_manager.py new file mode 100644 index 00000000..96a06c58 --- /dev/null +++ b/diffulex/strategy/fast_dllm_v2/engine/kv_cache_manager.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from diffulex.config import Config +from diffulex.engine.kv_cache_manager import AutoKVCacheManager, KVCacheManagerBase + + +@AutoKVCacheManager.register("fast_dllm_v2") +class FastDLLMV2KVCacheManager(KVCacheManagerBase): + def __init__(self, config: Config): + super().__init__(config) + + def _missing_cache_pages(self, req) -> int: + if getattr(req, "is_decoding", False) and hasattr(req, "fdv2_read_cache_pages"): + return max(0, int(req.fdv2_read_cache_pages) - len(req.page_table)) + return super()._missing_cache_pages(req) + + def may_append(self, req) -> None: + if not (getattr(req, "is_decoding", False) and hasattr(req, "fdv2_read_cache_pages")): + return super().may_append(req) + + missing_pages = self._missing_cache_pages(req) + if missing_pages > len(self.free_page_ids): + raise RuntimeError( + "Insufficient free KV cache pages for Fast-dLLM v2 current block: " + f"missing_pages={missing_pages}, free_pages={len(self.free_page_ids)}, " + f"fdv2_read_cache_len={req.fdv2_read_cache_len}, req_id={req.req_id}" + ) + + for _ in range(missing_pages): + page_id = self.free_page_ids[0] + self._allocate_page(page_id) + req.page_table.append(page_id) diff --git a/diffulex/strategy/fast_dllm_v2/engine/model_runner.py b/diffulex/strategy/fast_dllm_v2/engine/model_runner.py new file mode 100644 index 00000000..b5b5ec38 --- /dev/null +++ b/diffulex/strategy/fast_dllm_v2/engine/model_runner.py @@ -0,0 +1,432 @@ +from __future__ import annotations + +import os +from multiprocessing.synchronize import Event + +import torch +from tqdm import tqdm + +from diffulex.attention.metadata import reset_warming_up, set_fetch_fn_for_attn_metadata, set_warming_up +from diffulex.config import Config +from diffulex.engine.model_runner import AutoModelRunner, ModelRunnerBase +from diffulex.profiling import record_function +from diffulex.strategy.fast_dllm_v2.attention.metadata import ( + fetch_fast_dllm_v2_attn_metadata, + reset_fast_dllm_v2_attn_metadata, + set_fast_dllm_v2_attn_metadata, +) +from diffulex.strategy.fast_dllm_v2.engine.request import FastDLLMV2Mode + +_FDV2_GRAPH_FULL_BUFFER = "full_buffer_init" +_FDV2_GRAPH_SUB_BLOCK = "sub_block_cache_only" +_FDV2_GRAPH_FINAL_COMMIT = "final_commit" + + +@AutoModelRunner.register("fast_dllm_v2") +class FastDLLMV2ModelRunner(ModelRunnerBase): + def __init__(self, config: Config, rank: int, event: Event | list[Event]): + set_fetch_fn_for_attn_metadata(fetch_fast_dllm_v2_attn_metadata) + self.init_attn_metadata_fn( + set_fast_dllm_v2_attn_metadata, + reset_fast_dllm_v2_attn_metadata, + fetch_fast_dllm_v2_attn_metadata, + ) + self.mask_token_id = config.mask_token_id + self.fdv2_attention_block_size = int(config.block_size) * int(config.buffer_size) + self.is_prefix_full = False + self.mask_prefix_hole = False + self.prefix_causal = False + super().__init__(config, rank, event) + + def _prepare_decode_req(self, req): + prepared = super()._prepare_decode_req(req) + + if getattr(req, "fdv2_mode", None) == FastDLLMV2Mode.SUB_BLOCK_REFINE: + if not bool(getattr(req, "fdv2_use_block_cache", True)): + prepared["slot_mapping"] = self._fdv2_buffer_slot_mapping(req) + prepared["fdv2_cache_only"] = False + return prepared + + block = req.fdv2_current_sub_block + valid_len = req.fdv2_block_valid_len(block) + input_ids = list(block.token_ids[:valid_len]) + positions = list(range(block.start, block.start + valid_len)) + slot_mapping = [] + for abs_pos in range(block.start, block.start + valid_len): + rel_page_id = abs_pos // req.page_size + if rel_page_id >= len(req.page_table): + slot_mapping.append(-1) + continue + page_id = req.page_table[rel_page_id] + slot_mapping.append(page_id * self.page_size + abs_pos % self.page_size) + + return dict( + input_ids=input_ids, + positions=positions, + context_len=req.fdv2_read_cache_len, + seqlen_q=valid_len, + seqlen_k=req.fdv2_read_cache_len, + valid_slice=valid_len, + slot_mapping=slot_mapping, + status=2, + prefix_len=0, + padded_prefix_len=0, + fdv2_cache_only=True, + ) + + if getattr(req, "fdv2_mode", None) == FastDLLMV2Mode.FULL_BUFFER_INIT: + prepared["slot_mapping"] = self._fdv2_buffer_slot_mapping(req) + prepared["fdv2_cache_only"] = False + return prepared + + if getattr(req, "fdv2_mode", None) == FastDLLMV2Mode.FINAL_COMMIT: + prepared["slot_mapping"] = self._fdv2_buffer_slot_mapping(req) + prepared["fdv2_cache_only"] = False + return prepared + + prepared["fdv2_cache_only"] = False + return prepared + + def _fdv2_buffer_slot_mapping(self, req) -> list[int]: + slots: list[int] = [] + for abs_pos in range(req.fdv2_buffer_start, req.fdv2_effective_buffer_end): + rel_page_id = abs_pos // req.page_size + if rel_page_id >= len(req.page_table): + slots.append(-1) + continue + page_id = req.page_table[rel_page_id] + slots.append(page_id * self.page_size + abs_pos % self.page_size) + return slots + + def prepare_chunked_prefill_multi_block(self, reqs): + input_ids, positions = super().prepare_chunked_prefill_multi_block(reqs) + if not reqs: + return input_ids, positions + attn_metadata = self.fetch_attn_metadata() + # Diffulex uses `block_size` as Fast-dLLM v2's sub-block scheduling + # unit. The adapted model's block-causal attention mask is still the + # full Fast-dLLM block (`small_block_size * num_small_blocks`). + attn_metadata.block_size = self.fdv2_attention_block_size + attn_metadata.fdv2_mode = int(getattr(reqs[0], "fdv2_mode", FastDLLMV2Mode.FULL_BUFFER_INIT)) + fdv2_cache_only = any( + getattr(req, "is_decoding", False) + and getattr(req, "fdv2_mode", None) == FastDLLMV2Mode.SUB_BLOCK_REFINE + and bool(getattr(req, "fdv2_use_block_cache", True)) + for req in reqs + ) + attn_metadata.fdv2_cache_only = bool(fdv2_cache_only) + return input_ids, positions + + @staticmethod + def _graph_seq_batch_sizes(max_num_seqs: int) -> list[int]: + # Fast-dLLM v2's greedy refinement is very sensitive to tiny bf16 + # differences. Replaying a larger padded CUDA graph changes GEMM shapes + # versus eager and can flip accepted tokens, so capture exact request + # counts for lossless graph replay. + return list(range(1, max_num_seqs + 1)) + + def _fdv2_graph_mode(self, attn_metadata) -> str: + if bool(getattr(attn_metadata, "fdv2_cache_only", False)): + return _FDV2_GRAPH_SUB_BLOCK + if int(getattr(attn_metadata, "fdv2_mode", FastDLLMV2Mode.FULL_BUFFER_INIT)) == int( + FastDLLMV2Mode.FINAL_COMMIT + ): + return _FDV2_GRAPH_FINAL_COMMIT + return _FDV2_GRAPH_FULL_BUFFER + + def _fdv2_graph_q_len(self, mode: str) -> int: + if mode == _FDV2_GRAPH_SUB_BLOCK: + return int(self.config.block_size) + if mode in (_FDV2_GRAPH_FULL_BUFFER, _FDV2_GRAPH_FINAL_COMMIT): + return int(self.fdv2_attention_block_size) + raise ValueError(f"Unknown Fast-dLLM v2 CUDA graph mode: {mode}") + + @staticmethod + def _fdv2_graph_mode_enabled(mode: str) -> bool: + raw = os.environ.get("DIFFULEX_FDV2_GRAPH_MODES") + if not raw: + return True + aliases = { + "full": _FDV2_GRAPH_FULL_BUFFER, + "full_buffer": _FDV2_GRAPH_FULL_BUFFER, + "full_buffer_init": _FDV2_GRAPH_FULL_BUFFER, + "sub": _FDV2_GRAPH_SUB_BLOCK, + "sub_block": _FDV2_GRAPH_SUB_BLOCK, + "sub_block_cache_only": _FDV2_GRAPH_SUB_BLOCK, + "commit": _FDV2_GRAPH_FINAL_COMMIT, + "final": _FDV2_GRAPH_FINAL_COMMIT, + "final_commit": _FDV2_GRAPH_FINAL_COMMIT, + } + enabled = {aliases.get(part.strip(), part.strip()) for part in raw.split(",") if part.strip()} + return mode in enabled + + def _fdv2_can_run_decode_graph(self, input_ids: torch.Tensor, attn_metadata) -> bool: + if self.enforce_eager or not bool(getattr(self.config, "enable_full_static_runner", True)): + return False + graphs = getattr(self, "graphs", None) + graph_bs = getattr(self, "graph_bs", None) + if not graphs or not graph_bs: + return False + mode = self._fdv2_graph_mode(attn_metadata) + if not self._fdv2_graph_mode_enabled(mode): + return False + q_len = self._fdv2_graph_q_len(mode) + num_tokens = int(input_ids.size(0)) + if num_tokens <= 0 or num_tokens % q_len != 0: + return False + cu_seqlens_q = getattr(attn_metadata, "cu_seqlens_q", None) + if cu_seqlens_q is not None and int(cu_seqlens_q.numel()) > 1: + q_diffs = cu_seqlens_q[1:] - cu_seqlens_q[:-1] + if bool(torch.any(q_diffs != q_len).item()): + return False + mode_bs = graph_bs.get(mode) + return bool(mode_bs) and num_tokens <= max(mode_bs) + + @torch.inference_mode() + def _fdv2_run_decode_graph(self, input_ids: torch.Tensor, positions: torch.Tensor, attn_metadata): + mode = self._fdv2_graph_mode(attn_metadata) + q_len = self._fdv2_graph_q_len(mode) + num_tokens = int(input_ids.size(0)) + captured_num_tokens = next(x for x in self.graph_bs[mode] if x >= num_tokens) + captured_num_seqs = captured_num_tokens // q_len + graph = self.graphs[(mode, captured_num_tokens)] + graph_vars = self.graph_vars + + num_reqs = attn_metadata.num_reqs + graph_capacity = int(graph_vars["context_lens"].size(0)) + if captured_num_seqs > graph_capacity: + raise RuntimeError( + "Captured Fast-dLLM v2 CUDA graph batch size exceeds graph buffer capacity: " + f"mode={mode}, captured_num_seqs={captured_num_seqs}, " + f"graph_capacity={graph_capacity}, captured_num_tokens={captured_num_tokens}, " + f"num_tokens={num_tokens}" + ) + if num_reqs > captured_num_seqs: + raise RuntimeError( + "Fast-dLLM v2 CUDA graph bucket cannot cover current request count: " + f"mode={mode}, num_reqs={num_reqs}, captured_num_seqs={captured_num_seqs}, " + f"num_tokens={num_tokens}, q_len={q_len}" + ) + + self._copy_common_graph_inputs( + graph_vars, + attn_metadata, + input_ids, + positions, + num_tokens, + num_reqs, + ) + + for i in range(num_reqs, captured_num_seqs): + graph_vars["cu_seqlens_q"][i + 1] = graph_vars["cu_seqlens_q"][i] + graph_vars["cu_seqlens_k"][i + 1] = graph_vars["cu_seqlens_k"][i] + + restore_fields = ( + "slot_mapping", + "context_lens", + "cu_seqlens_q", + "cu_seqlens_k", + "valid_slices", + "status_table", + "prefix_lens", + "padded_prefix_lens", + "page_tables", + "block_size", + "fdv2_cache_only", + "fdv2_mode", + ) + original_metadata = {field: getattr(attn_metadata, field, None) for field in restore_fields} + try: + attn_metadata.slot_mapping = graph_vars["slot_mapping"] + attn_metadata.context_lens = graph_vars["context_lens"] + attn_metadata.cu_seqlens_q = graph_vars["cu_seqlens_q"] + attn_metadata.cu_seqlens_k = graph_vars["cu_seqlens_k"] + attn_metadata.valid_slices = graph_vars["valid_slices"] + attn_metadata.status_table = graph_vars["status_table"] + attn_metadata.prefix_lens = graph_vars["prefix_lens"] + attn_metadata.padded_prefix_lens = graph_vars["padded_prefix_lens"] + attn_metadata.page_tables = graph_vars["page_tables"] + attn_metadata.block_size = self.fdv2_attention_block_size + attn_metadata.fdv2_cache_only = mode == _FDV2_GRAPH_SUB_BLOCK + attn_metadata.fdv2_mode = int( + FastDLLMV2Mode.SUB_BLOCK_REFINE + if mode == _FDV2_GRAPH_SUB_BLOCK + else FastDLLMV2Mode.FINAL_COMMIT + if mode == _FDV2_GRAPH_FINAL_COMMIT + else FastDLLMV2Mode.FULL_BUFFER_INIT + ) + self._bind_decode_graph_extra_metadata(attn_metadata, graph_vars, num_tokens) + graph.replay() + finally: + for field, value in original_metadata.items(): + setattr(attn_metadata, field, value) + + if bool(getattr(self, "graph_outputs_are_logits", False)): + return graph_vars["outputs"][:num_tokens] + return self.model.compute_logits(graph_vars["outputs"][:num_tokens]) + + @torch.inference_mode() + def run_model_multi_block(self, input_ids: torch.Tensor, positions: torch.Tensor): + with record_function("diffulex.fast_dllm_v2.model_forward"): + attn_metadata = self.fetch_attn_metadata() + if attn_metadata.has_prefill: + with record_function("diffulex.fast_dllm_v2.eager_prefill"): + return self.model.compute_logits(self.model(input_ids, positions)) + + if not self._fdv2_can_run_decode_graph(input_ids, attn_metadata): + with record_function("diffulex.fast_dllm_v2.eager_decode"): + return self.model.compute_logits(self.model(input_ids, positions)) + + with record_function("diffulex.fast_dllm_v2.cuda_graph_decode"): + return self._fdv2_run_decode_graph(input_ids, positions, attn_metadata) + + @torch.inference_mode() + def capture_cudagraph(self): + set_warming_up(True) + config = self.config + hf_config = config.hf_config + max_num_seqs = min(self.config.max_num_reqs, 512) + max_num_pages = (config.max_model_len + self.page_size - 1) // self.page_size + max_q_len = int(self.fdv2_attention_block_size) + max_num_tokens = max_num_seqs * max_q_len + device = self._cuda_graph_device() + capture_logits = self._can_capture_logits_in_graph() + self.graph_outputs_are_logits = capture_logits + output_size = self._model_logits_size() if capture_logits else hf_config.hidden_size + + input_ids = torch.zeros(max_num_tokens, dtype=torch.int64, device=device) + positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device) + slot_mapping = torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device) + context_lens = torch.zeros(max_num_seqs, dtype=torch.int32, device=device) + page_tables = torch.zeros(max_num_seqs, max_num_pages, dtype=torch.int32, device=device) + valid_slices = torch.zeros(max_num_seqs, dtype=torch.int32, device=device) + status_table = torch.zeros(max_num_seqs, dtype=torch.int32, device=device) + prefix_lens = torch.zeros(max_num_seqs, dtype=torch.int32, device=device) + padded_prefix_lens = torch.zeros(max_num_seqs, dtype=torch.int32, device=device) + outputs = torch.zeros(max_num_tokens, output_size, dtype=self._model_logits_dtype(), device=device) + + cu_seqlens_q = torch.zeros(max_num_seqs + 1, dtype=torch.int32, device=device) + cu_seqlens_k = torch.zeros(max_num_seqs + 1, dtype=torch.int32, device=device) + + graph_vars = dict( + input_ids=input_ids, + positions=positions, + slot_mapping=slot_mapping, + context_lens=context_lens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + page_tables=page_tables, + valid_slices=valid_slices, + status_table=status_table, + prefix_lens=prefix_lens, + padded_prefix_lens=padded_prefix_lens, + outputs=outputs, + ) + + seq_bs_list = self._graph_seq_batch_sizes(max_num_seqs) + self.graph_bs = { + mode: [num_seqs * self._fdv2_graph_q_len(mode) for num_seqs in seq_bs_list] + for mode in (_FDV2_GRAPH_FULL_BUFFER, _FDV2_GRAPH_SUB_BLOCK, _FDV2_GRAPH_FINAL_COMMIT) + } + self.graphs = {} + self.graph_pool = None + + try: + for mode in (_FDV2_GRAPH_FULL_BUFFER, _FDV2_GRAPH_SUB_BLOCK, _FDV2_GRAPH_FINAL_COMMIT): + if not self._fdv2_graph_mode_enabled(mode): + continue + q_len = self._fdv2_graph_q_len(mode) + is_sub_block = mode == _FDV2_GRAPH_SUB_BLOCK + is_final_commit = mode == _FDV2_GRAPH_FINAL_COMMIT + for num_tokens in tqdm( + reversed(self.graph_bs[mode]), + desc=f"Capturing Fast-dLLM v2 CUDA graphs ({mode})", + ): + num_seqs = num_tokens // q_len + input_ids.zero_() + positions.zero_() + slot_mapping.fill_(-1) + page_tables.fill_(0) + context_lens.zero_() + valid_slices.zero_() + status_table.zero_() + prefix_lens.zero_() + padded_prefix_lens.zero_() + + for i in range(max_num_seqs + 1): + cu_seqlens_q[i] = i * q_len + cu_seqlens_k[i] = i * config.max_model_len + for i in range(num_seqs): + valid_slices[i] = (i + 1) * q_len + # Capture decode graphs with a non-empty cache prefix. The + # Triton attention kernel must record the cache-attention + # path because real Fast-dLLM v2 decode replays against + # prefix/current-buffer KV cache even for full-buffer init + # and final commit. + context_lens[:num_seqs] = config.max_model_len + if is_sub_block: + status_table[:num_seqs] = 2 + else: + status_table[:num_seqs] = 1 + + self.set_attn_metadata( + False, + slot_mapping=slot_mapping[:num_tokens], + need_kv_cache_store=True, + context_lens=context_lens[:num_seqs], + cu_seqlens_q=cu_seqlens_q[: num_seqs + 1], + cu_seqlens_k=cu_seqlens_k[: num_seqs + 1], + max_seqlen_q=q_len, + max_seqlen_k=config.max_model_len, + page_size=config.kv_cache_page_size, + page_tables=page_tables[:num_seqs], + block_size=self.fdv2_attention_block_size, + kv_cache_layout=config.kv_cache_layout, + fdv2_cache_only=is_sub_block, + fdv2_mode=int( + FastDLLMV2Mode.SUB_BLOCK_REFINE + if is_sub_block + else FastDLLMV2Mode.FINAL_COMMIT + if is_final_commit + else FastDLLMV2Mode.FULL_BUFFER_INIT + ), + ) + attn_metadata = self.fetch_attn_metadata() + attn_metadata.init_multi_block( + valid_slices=valid_slices[:num_seqs], + buffer_size=config.buffer_size, + is_prefix_full=self.is_prefix_full, + status_table=status_table[:num_seqs], + prefix_lens=prefix_lens[:num_seqs], + padded_prefix_lens=padded_prefix_lens[:num_seqs], + mask_prefix_hole=bool(getattr(self, "mask_prefix_hole", False)), + prefix_causal=bool(getattr(self, "prefix_causal", False)), + ) + attn_metadata.block_size = self.fdv2_attention_block_size + attn_metadata.fdv2_cache_only = is_sub_block + attn_metadata.fdv2_mode = int( + FastDLLMV2Mode.SUB_BLOCK_REFINE + if is_sub_block + else FastDLLMV2Mode.FINAL_COMMIT + if is_final_commit + else FastDLLMV2Mode.FULL_BUFFER_INIT + ) + self._init_graph_capture_extra_metadata(attn_metadata, graph_vars, num_tokens) + + graph = self._capture_model_forward_graph( + input_ids, + positions, + outputs, + num_tokens, + allow_compile=True, + capture_logits=capture_logits, + ) + if self.graph_pool is None: + self.graph_pool = graph.pool() + self.graphs[(mode, num_tokens)] = graph + torch.cuda.synchronize() + self.reset_attn_metadata() + + self.graph_vars = graph_vars + finally: + reset_warming_up() diff --git a/diffulex/strategy/fast_dllm_v2/engine/request.py b/diffulex/strategy/fast_dllm_v2/engine/request.py new file mode 100644 index 00000000..050044f1 --- /dev/null +++ b/diffulex/strategy/fast_dllm_v2/engine/request.py @@ -0,0 +1,407 @@ +from __future__ import annotations + +from enum import IntEnum + +from diffulex.attention.metadata import is_warming_up +from diffulex.config import Config +from diffulex.engine.request import AutoReq, DllmReq +from diffulex.engine.dllm_block import DllmBlock, DllmBlockBuffer +from diffulex.engine.status import DllmBlockStatus +from diffulex.sampling_params import SamplingParams + + +class FastDLLMV2Mode(IntEnum): + FULL_BUFFER_INIT = 0 + SUB_BLOCK_REFINE = 1 + FINAL_COMMIT = 2 + + +@AutoReq.register("fast_dllm_v2") +class FastDLLMV2Req(DllmReq): + def __init__( + self, + token_ids: list[int], + sampling_params: SamplingParams = SamplingParams(), + config: Config | None = None, + ): + super().__init__(token_ids, sampling_params) + self.fdv2_current_buffer_initialized = False + self.fdv2_current_sub_block_idx = 0 + self.fdv2_mode = FastDLLMV2Mode.FULL_BUFFER_INIT + self.fdv2_prefix_cache_len = 0 + self.fdv2_pending_next_token_id: int | None = None + self.fdv2_use_block_cache = True + + @property + def fdv2_sequence_limit(self) -> int: + if not hasattr(self, "prefix_len"): + return int(getattr(self, "max_model_len", len(self.token_ids))) + return min(int(self.max_model_len), int(self.prefix_len) + int(self.max_new_tokens)) + + @property + def fdv2_committed_generated_len(self) -> int: + total = 0 + for block in getattr(self, "dllm_blocks", []): + if not block.is_in_cache: + continue + valid_len = int(getattr(block, "valid_commit_len", block.block_size)) + block_gen_start = max(int(self.prefix_len), int(block.start)) + block_gen_end = min(int(block.start) + valid_len, int(block.end), self.fdv2_sequence_limit) + total += max(0, block_gen_end - block_gen_start) + return total + + def _fdv2_block_valid_commit_len(self, block: DllmBlock, committed_generated_len: int | None = None) -> int: + if committed_generated_len is None: + committed_generated_len = self.fdv2_committed_generated_len + prompt_len = max(0, min(int(block.end), int(self.prefix_len)) - int(block.start)) + remaining_new = max(0, int(self.max_new_tokens) - int(committed_generated_len)) + remaining_model = max( + 0, + int(self.max_model_len) - int(self.prefix_len) - int(committed_generated_len), + ) + generated_len = min(int(block.block_size) - prompt_len, remaining_new, remaining_model) + return max(0, min(int(block.block_size), prompt_len + generated_len)) + + def fdv2_block_valid_len(self, block: DllmBlock) -> int: + if block.is_in_cache or block.is_to_cache: + return max(0, min(int(getattr(block, "valid_commit_len", block.block_size)), int(block.block_size))) + return max(0, min(int(block.end), self.fdv2_sequence_limit) - int(block.start)) + + def fdv2_block_mask_token_relative_ids(self, block: DllmBlock) -> list[int]: + valid_len = self.fdv2_block_valid_len(block) + return [idx for idx in block.mask_token_relative_ids if idx < valid_len] + + def _fdv2_block_complete(self, block: DllmBlock) -> bool: + valid_len = self.fdv2_block_valid_len(block) + if valid_len <= int(getattr(block, "editable_start", 0)): + return True + return all(token_id != self.mask_token_id for token_id in block.token_ids[:valid_len]) + + def init_multi_block(self, config: Config): + self.is_multi_block = True + self.status_history = [self.status] + self.completion_reason = None + self._resume_prefill_until = 0 + self._terminal_context_block_id: int | None = None + + self.block_size = config.block_size + self.buffer_size = config.buffer_size + strategy_config = getattr(config, "strategy", None) + self.fdv2_use_block_cache = bool( + getattr(strategy_config, "use_block_cache", getattr(config, "fdv2_use_block_cache", True)) + ) + self.mask_token_id = config.mask_token_id + self.thresholds = config.decoding_thresholds + self.eos_token_id = config.eos + self.max_model_len = config.max_model_len + self.max_new_tokens = self.max_tokens + self.auto_max_nfe_enabled = self.max_nfe is None + self.auto_max_nfe_warmup_steps = int(getattr(config, "auto_max_nfe_warmup_steps", 8)) + self.auto_max_nfe_tpf_floor = float(getattr(config, "auto_max_nfe_tpf_floor", 1.0)) + self.auto_max_nfe_token_count = 0 + self.auto_max_nfe_value = None + self.auto_max_nfe_avg_tpf = None + + self.dllm_blocks: list[DllmBlock] = [] + self.dllm_block_buffer: DllmBlockBuffer | None = None + + if self.max_model_len_reached and not is_warming_up(): + self.force_deactivate() + return + + self.prefix_len = len(self.token_ids) + self.padded_prefix_len = self.prefix_len + + fdv2_block_size = self.block_size * self.buffer_size + self.fdv2_prefix_cache_len = (self.prefix_len // fdv2_block_size) * fdv2_block_size + buffer_start = self.fdv2_prefix_cache_len + buffer_end = buffer_start + fdv2_block_size + + if len(self.token_ids) < buffer_end: + self.pad_tokens(buffer_end - len(self.token_ids)) + + for i in range(self.fdv2_prefix_cache_len // self.block_size): + start = i * self.block_size + end = start + self.block_size + dllm_block = DllmBlock( + block_id=i, + start=start, + end=end, + block_size=self.block_size, + mask_token_id=self.mask_token_id, + thresholds=self.thresholds, + status=None, + prev_block=None if i == 0 else self.dllm_blocks[-1], + ) + dllm_block.post_init_dllm_block(self, None) + self.dllm_blocks.append(dllm_block) + + for _ in range(self.buffer_size): + block_id = len(self.dllm_blocks) + start = buffer_start + (block_id - self.fdv2_prefix_cache_len // self.block_size) * self.block_size + end = start + self.block_size + editable_start = max(0, min(self.block_size, self.prefix_len - start)) + dllm_block = DllmBlock( + block_id=block_id, + start=start, + end=end, + block_size=self.block_size, + mask_token_id=self.mask_token_id, + thresholds=self.thresholds, + status=DllmBlockStatus.ACTIVE, + prev_block=None if block_id == 0 else self.dllm_blocks[-1], + editable_start=editable_start, + ) + dllm_block.post_init_dllm_block(self, None) + dllm_block.commit_ready = False + dllm_block.valid_commit_len = self.block_size + self.dllm_blocks.append(dllm_block) + + self.dllm_block_buffer = DllmBlockBuffer( + buffer_size=self.buffer_size, + dllm_blocks=self.dllm_blocks[-self.buffer_size :], + ) + self.dllm_block_buffer.post_init_dllm_block_buffer(self) + for block in self.fdv2_buffer_blocks: + block.commit_ready = False + + @property + def num_prefix_blocks(self) -> int: + return int(self.fdv2_prefix_cache_len) // int(self.block_size) + + @property + def num_prefix_pages(self) -> int: + return self.num_pages_with_seq_len(self.fdv2_prefix_cache_len) + + @property + def fdv2_buffer_blocks(self): + if self.dllm_block_buffer is None: + return [] + return self.dllm_block_buffer.dllm_blocks + + @property + def fdv2_current_sub_block(self): + return self.fdv2_buffer_blocks[self.fdv2_current_sub_block_idx] + + @property + def fdv2_buffer_start(self) -> int: + return self.dllm_block_buffer.first_running_block.start + + @property + def fdv2_buffer_end(self) -> int: + return self.dllm_block_buffer.last_running_block.end + + @property + def fdv2_effective_buffer_end(self) -> int: + return min(self.fdv2_buffer_end, self.fdv2_sequence_limit) + + @property + def fdv2_buffer_token_ids(self) -> list[int]: + return self.token_ids[self.fdv2_buffer_start : self.fdv2_effective_buffer_end] + + @property + def fdv2_buffer_position_ids(self) -> list[int]: + return list(range(self.fdv2_buffer_start, self.fdv2_effective_buffer_end)) + + @property + def fdv2_replace_position(self) -> int: + return int(self.fdv2_current_sub_block_idx) * int(self.block_size) + + @property + def fdv2_read_cache_len(self) -> int: + return min(self.fdv2_buffer_end, self.fdv2_sequence_limit) + + @property + def fdv2_read_cache_pages(self) -> int: + return self.num_pages_with_seq_len(self.fdv2_read_cache_len) + + @property + def fdv2_committed_prefix_len(self) -> int: + return self.contiguous_in_cache_prefix_len + + @property + def fdv2_sub_block_complete(self) -> bool: + return self._fdv2_block_complete(self.fdv2_current_sub_block) + + @property + def fdv2_current_sub_block_first_token_is_mask(self) -> bool: + return ( + self.fdv2_block_valid_len(self.fdv2_current_sub_block) > 0 + and self.fdv2_current_sub_block.token_ids[0] == self.mask_token_id + ) + + @property + def fdv2_buffer_complete(self) -> bool: + return all(self._fdv2_block_complete(block) for block in self.fdv2_buffer_blocks) + + @property + def max_new_tokens_reached(self) -> bool: + if hasattr(self, "prefix_len"): + return self.fdv2_committed_generated_len >= self.max_new_tokens + return super().max_new_tokens_reached + + @property + def max_model_len_reached(self) -> bool: + if hasattr(self, "prefix_len"): + return int(self.prefix_len) + self.fdv2_committed_generated_len >= self.max_model_len + return super().max_model_len_reached + + @property + def eos_token_generated(self) -> bool: + if self.ignore_eos: + return False + return self.eos_token_id in self.full_response + + @property + def full_response(self) -> list[int]: + tokens: list[int] = [] + for block in getattr(self, "dllm_blocks", []): + if not block.is_in_cache: + continue + valid_len = max(0, min(int(getattr(block, "valid_commit_len", block.block_size)), int(block.block_size))) + start = max(int(block.start), int(self.prefix_len)) + end = min(int(block.start) + valid_len, int(block.end), self.fdv2_sequence_limit) + if end > start: + tokens.extend(self.token_ids[start:end]) + return tokens + + @property + def truncated_response(self) -> list[int]: + generated_seq = self.full_response + if self.eos_token_id in generated_seq and not self.ignore_eos: + generated_seq = generated_seq[: generated_seq.index(self.eos_token_id)] + if self.max_new_tokens_reached: + generated_seq = generated_seq[: self.max_new_tokens] + return generated_seq + + def refresh_fdv2_mode(self) -> None: + if not self.is_decoding: + return + while ( + self.fdv2_current_sub_block_idx < self.buffer_size - 1 + and self._fdv2_block_complete(self.fdv2_current_sub_block) + ): + self.fdv2_current_sub_block_idx += 1 + if not self.fdv2_current_buffer_initialized: + self.fdv2_mode = FastDLLMV2Mode.FULL_BUFFER_INIT + return + if self.fdv2_buffer_complete: + self.fdv2_mode = FastDLLMV2Mode.FINAL_COMMIT + return + if self.fdv2_current_sub_block_first_token_is_mask: + self.fdv2_mode = FastDLLMV2Mode.FULL_BUFFER_INIT + else: + self.fdv2_mode = FastDLLMV2Mode.SUB_BLOCK_REFINE + + @property + def running_sequence(self) -> list[int]: + if self.is_prefilling: + return super().running_sequence + if self.is_decoding or self.is_completed: + if self.fdv2_mode == FastDLLMV2Mode.SUB_BLOCK_REFINE and self.fdv2_use_block_cache: + valid_len = self.fdv2_block_valid_len(self.fdv2_current_sub_block) + return self.fdv2_current_sub_block.token_ids[:valid_len] + return self.fdv2_buffer_token_ids + return [] + + @property + def running_position_ids(self) -> range | list[int]: + if self.is_prefilling: + return super().running_position_ids + if self.is_decoding or self.is_completed: + if self.fdv2_mode == FastDLLMV2Mode.SUB_BLOCK_REFINE and self.fdv2_use_block_cache: + valid_len = self.fdv2_block_valid_len(self.fdv2_current_sub_block) + return list(range(self.fdv2_current_sub_block.start, self.fdv2_current_sub_block.start + valid_len)) + return self.fdv2_buffer_position_ids + return [] + + @property + def running_len(self) -> int: + if self.is_prefilling: + return int(self.fdv2_prefix_cache_len) + if self.is_decoding: + if self.fdv2_mode == FastDLLMV2Mode.SUB_BLOCK_REFINE and self.fdv2_use_block_cache: + return self.fdv2_block_valid_len(self.fdv2_current_sub_block) + return len(self.fdv2_buffer_token_ids) + return super().running_len + + @property + def valid_len(self) -> int: + if self.is_prefilling: + return int(self.fdv2_prefix_cache_len) + if self.is_decoding: + if self.fdv2_mode == FastDLLMV2Mode.SUB_BLOCK_REFINE and self.fdv2_use_block_cache: + return self.fdv2_block_valid_len(self.fdv2_current_sub_block) + return len(self.fdv2_buffer_token_ids) + return super().valid_len + + @property + def cache_len(self) -> int: + if self.is_decoding: + return max(super().cache_len, self.fdv2_read_cache_len) + return super().cache_len + + def step(self): + if self.dllm_block_buffer is None: + self.lazy_activate() + return + super().step() + if not self.is_decoding: + return + for block in self.fdv2_buffer_blocks: + if block.is_dummy: + block.status = DllmBlockStatus.ACTIVE + self.refresh_fdv2_mode() + + def _seed_next_block_from_pending_token(self) -> None: + token_id = self.fdv2_pending_next_token_id + if token_id is None: + return + if self.max_new_tokens_reached or self.max_model_len_reached: + self.fdv2_pending_next_token_id = None + return + + for block in self.fdv2_buffer_blocks: + if int(getattr(block, "editable_start", 0)) > 0: + continue + if not block.token_ids or block.token_ids[0] != self.mask_token_id: + continue + block.write_token(int(token_id), 0) + self.new_tokens += 1 + self.fdv2_pending_next_token_id = None + return + + def postprocess(self): + if self.dllm_block_buffer is None: + return + + if self.is_decoding and self.fdv2_mode == FastDLLMV2Mode.FULL_BUFFER_INIT: + self.fdv2_current_buffer_initialized = True + self.refresh_fdv2_mode() + return + + if self.is_decoding and self.fdv2_mode == FastDLLMV2Mode.SUB_BLOCK_REFINE: + if self._fdv2_block_complete(self.fdv2_current_sub_block) and self.fdv2_current_sub_block_idx < self.buffer_size - 1: + self.fdv2_current_sub_block_idx += 1 + self.refresh_fdv2_mode() + return + + if self.is_decoding and self.fdv2_mode == FastDLLMV2Mode.FINAL_COMMIT: + committed_generated_len = self.fdv2_committed_generated_len + for block in self.fdv2_buffer_blocks: + block.valid_commit_len = self._fdv2_block_valid_commit_len(block, committed_generated_len) + block_gen_start = max(int(self.prefix_len), int(block.start)) + block_gen_end = min(int(block.start) + int(block.valid_commit_len), int(block.end)) + committed_generated_len += max(0, block_gen_end - block_gen_start) + if block.is_active: + block.commit_ready = True + block.status = DllmBlockStatus.TO_CACHE + super().postprocess() + self.fdv2_current_buffer_initialized = False + self.fdv2_current_sub_block_idx = 0 + self.fdv2_mode = FastDLLMV2Mode.FULL_BUFFER_INIT + self._seed_next_block_from_pending_token() + return + + super().postprocess() + self._seed_next_block_from_pending_token() diff --git a/diffulex/strategy/fast_dllm_v2/engine/scheduler.py b/diffulex/strategy/fast_dllm_v2/engine/scheduler.py new file mode 100644 index 00000000..74b0cd5f --- /dev/null +++ b/diffulex/strategy/fast_dllm_v2/engine/scheduler.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from diffulex.config import Config +from diffulex.engine.scheduler import AutoScheduler, SchedulerBase +from diffulex.engine.status import DllmReqStatus +from diffulex.strategy.fast_dllm_v2.engine.request import FastDLLMV2Mode + + +@AutoScheduler.register("fast_dllm_v2") +class FastDLLMV2Scheduler(SchedulerBase): + def __init__(self, config: Config): + super().__init__(config) + + @staticmethod + def _req_mode(req) -> FastDLLMV2Mode: + return getattr(req, "fdv2_mode", FastDLLMV2Mode.FULL_BUFFER_INIT) + + def schedule(self): + scheduled, is_prefill = super().schedule() + if is_prefill or len(scheduled) <= 1: + return scheduled, is_prefill + + target_mode = self._req_mode(scheduled[0]) + kept = [] + deferred = [] + for req in scheduled: + if self._req_mode(req) == target_mode: + kept.append(req) + else: + deferred.append(req) + + if not deferred: + return scheduled, is_prefill + + for req in deferred: + if req in self.running_reqs: + self.running_reqs.remove(req) + req.status = DllmReqStatus.DECODING + self.running_reqs.append(req) + + return kept, is_prefill diff --git a/diffulex_bench/arg_parser.py b/diffulex_bench/arg_parser.py index d1209a98..7a71b86f 100644 --- a/diffulex_bench/arg_parser.py +++ b/diffulex_bench/arg_parser.py @@ -19,7 +19,7 @@ "diffusion_gemma", ] -DECODING_STRATEGY_CHOICES = ["d2f", "multi_bd", "dmax"] +DECODING_STRATEGY_CHOICES = ["d2f", "multi_bd", "dmax", "fast_dllm_v2"] TOKEN_MERGE_MODE_CHOICES = ["dmax_topk", "iter_smooth_topk"] ATTN_IMPL_CHOICES = ["triton", "triton_grouped", "naive"] MOE_GEMM_IMPL_CHOICES = ["triton", "vllm", "vllm_modular", "naive"] @@ -106,7 +106,7 @@ def create_argument_parser() -> argparse.ArgumentParser: type=str, default="d2f", choices=DECODING_STRATEGY_CHOICES, - help="Decoding strategy (d2f, multi_bd, dmax)", + help="Decoding strategy (d2f, multi_bd, dmax, fast_dllm_v2)", ) parser.add_argument( "--sampling-mode", diff --git a/diffulex_bench/config.py b/diffulex_bench/config.py index 53f6095c..979e1245 100644 --- a/diffulex_bench/config.py +++ b/diffulex_bench/config.py @@ -109,7 +109,7 @@ class EngineConfig: model_path: str tokenizer_path: Optional[str] = None model_name: str = "dream" # Options: dream, sdar, fast_dllm_v2, llada - decoding_strategy: str = "d2f" # Options: d2f, multi_bd, dmax, diffusion_gemma + decoding_strategy: str = "d2f" # Options: d2f, multi_bd, dmax, fast_dllm_v2, diffusion_gemma sampling_mode: str = "naive" # Options: naive, edit max_post_edit_steps: int = 16 # max refinement steps after all masks filled mask_token_id: int = 151666 diff --git a/diffulex_bench/configs/fast_dllm_v2_gsm8k.yml b/diffulex_bench/configs/fast_dllm_v2_gsm8k.yml index 7b59a9b1..7a823b8d 100644 --- a/diffulex_bench/configs/fast_dllm_v2_gsm8k.yml +++ b/diffulex_bench/configs/fast_dllm_v2_gsm8k.yml @@ -1,12 +1,14 @@ -# Fast_dLLM V2 model - MultiBD strategy on GSM8K -# Aligned with test_diffulex_dry_run.yaml strategy_config.fast_dllm_v2 -# Fast_dLLM V2 uses multi_bd, buffer_size=1 +# Fast_dLLM V2 model - native dual-cache strategy on GSM8K using +# Diffulex's bundled gsm8k_diffulex task. +# Fast-dLLM v2 paper defaults map to Diffulex block_size=8, buffer_size=4: +# one Fast-dLLM sub-block is one Diffulex block, and one Fast-dLLM block is +# one Diffulex DllmBlockBuffer. engine: model_path: "Efficient-Large-Model/Fast_dLLM_v2_7B" tokenizer_path: null model_name: "fast_dllm_v2" - decoding_strategy: "multi_bd" + decoding_strategy: "fast_dllm_v2" mask_token_id: 151665 use_lora: false @@ -15,26 +17,29 @@ engine: tensor_parallel_size: 1 data_parallel_size: 1 gpu_memory_utilization: 0.4 - max_model_len: 1024 - max_num_batched_tokens: 1024 - max_num_reqs: 24 + max_model_len: 2048 + max_num_batched_tokens: 4096 + max_num_reqs: 32 enforce_eager: false kv_cache_layout: "unified" + # Fast_dLLM_v2_7B has 28 attention heads; the grouped Triton path currently + # requires power-of-two tile ranges during warmup for this shape. + attn_impl: "triton" decoding_thresholds: add_block_threshold: 0.1 semi_complete_threshold: 0.9 accept_threshold: 0.95 - block_size: 32 - buffer_size: 1 # Fast_dLLM V2 uses buffer_size=1 + block_size: 8 + buffer_size: 4 eval: dataset_name: "gsm8k_diffulex" dataset_split: "test" dataset_limit: null temperature: 0.0 - max_tokens: 256 - add_bos_token: true # Instruct model: chat template format + max_tokens: 2048 + add_bos_token: false output_dir: "benchmark_results" save_results: true diff --git a/diffulex_bench/configs/fast_dllm_v2_multibd_gsm8k.yml b/diffulex_bench/configs/fast_dllm_v2_multibd_gsm8k.yml new file mode 100644 index 00000000..8c02fcf7 --- /dev/null +++ b/diffulex_bench/configs/fast_dllm_v2_multibd_gsm8k.yml @@ -0,0 +1,42 @@ +# Fast_dLLM V2 model - training-free MultiBD on GSM8K using Diffulex's +# bundled gsm8k_diffulex task. This uses Diffulex's generic MultiBD runtime: +# a bounded block buffer of consecutive diffusion blocks decoded concurrently. + +engine: + model_path: "Efficient-Large-Model/Fast_dLLM_v2_7B" + tokenizer_path: null + model_name: "fast_dllm_v2" + decoding_strategy: "multi_bd" + mask_token_id: 151665 + + use_lora: false + lora_path: "" + + tensor_parallel_size: 1 + data_parallel_size: 1 + gpu_memory_utilization: 0.4 + max_model_len: 2048 + max_num_batched_tokens: 4096 + max_num_reqs: 32 + + enforce_eager: false + kv_cache_layout: "unified" + attn_impl: "triton" + + decoding_thresholds: + add_block_threshold: 0.1 + semi_complete_threshold: 0.9 + accept_threshold: 0.975 + token_stability_threshold: 0.5 + block_size: 4 + buffer_size: 6 + +eval: + dataset_name: "gsm8k_diffulex" + dataset_split: "test" + dataset_limit: null + temperature: 0.0 + max_tokens: 2048 + add_bos_token: false + output_dir: "benchmark_results" + save_results: true diff --git a/diffulex_bench/main.py b/diffulex_bench/main.py index 7280aad0..49a80158 100644 --- a/diffulex_bench/main.py +++ b/diffulex_bench/main.py @@ -539,7 +539,8 @@ def apply_engine_arg_overrides(engine: EngineConfig) -> None: ("remask_threshold", "remask_threshold"), ("token_stability_threshold", "token_stability_threshold"), ): - if was_provided(cli_key) and getattr(args, cli_key, None) is not None: + flag = f"--{cli_key.replace('_', '-')}" + if option_was_provided(flag) and getattr(args, cli_key, None) is not None: config.engine.decoding_thresholds[yaml_key] = getattr(args, cli_key) if was_provided("max_post_edit_steps") and getattr(args, "max_post_edit_steps", None) is not None: config.engine.max_post_edit_steps = args.max_post_edit_steps diff --git a/diffulex_kernel/python/chunked_prefill_grouped_triton.py b/diffulex_kernel/python/chunked_prefill_grouped_triton.py index ce3ef70e..8688f61d 100644 --- a/diffulex_kernel/python/chunked_prefill_grouped_triton.py +++ b/diffulex_kernel/python/chunked_prefill_grouped_triton.py @@ -75,6 +75,7 @@ def _chunked_prefill_grouped_attn_unified_kernel( IS_PREFIX_FULL: tl.constexpr, MASK_PREFIX_HOLE: tl.constexpr, PREFIX_CAUSAL: tl.constexpr, + FDV2_CACHE_ONLY: tl.constexpr, SLIDING_WINDOW: tl.constexpr, ): req_id = tl.program_id(0) @@ -186,6 +187,8 @@ def _chunked_prefill_grouped_attn_unified_kernel( loop_range = block_causal_range else: loop_range = full_range + if FDV2_CACHE_ONLY: + loop_range = 0 for kv_block_id in range(0, loop_range): kv_block_start = kv_block_id * BLOCK_N @@ -313,6 +316,7 @@ def chunked_prefill_attn_grouped_unified( IS_PREFIX_FULL=attn_metadata.is_prefix_full, MASK_PREFIX_HOLE=bool(getattr(attn_metadata, "mask_prefix_hole", False)), PREFIX_CAUSAL=bool(getattr(attn_metadata, "prefix_causal", False)), + FDV2_CACHE_ONLY=bool(getattr(attn_metadata, "fdv2_cache_only", False)), SLIDING_WINDOW=int(sliding_window or 0), ) return out diff --git a/diffulex_kernel/python/chunked_prefill_triton.py b/diffulex_kernel/python/chunked_prefill_triton.py index b01ca62e..cb7844b4 100644 --- a/diffulex_kernel/python/chunked_prefill_triton.py +++ b/diffulex_kernel/python/chunked_prefill_triton.py @@ -122,6 +122,7 @@ def _chunked_prefill_attn_unified_kernel( IS_PREFIX_FULL: tl.constexpr, MASK_PREFIX_HOLE: tl.constexpr, PREFIX_CAUSAL: tl.constexpr, + FDV2_CACHE_ONLY: tl.constexpr, SLIDING_WINDOW: tl.constexpr, ): req_id = tl.program_id(0) @@ -239,6 +240,8 @@ def _chunked_prefill_attn_unified_kernel( loop_range = block_causal_range else: loop_range = full_range + if FDV2_CACHE_ONLY: + loop_range = 0 for kv_block_id in range(0, loop_range): kv_block_start = kv_block_id * BLOCK_N @@ -372,6 +375,7 @@ def _run_chunked_prefill_attn_unified_kernel( IS_PREFIX_FULL=attn_metadata.is_prefix_full, MASK_PREFIX_HOLE=bool(getattr(attn_metadata, "mask_prefix_hole", False)), PREFIX_CAUSAL=bool(getattr(attn_metadata, "prefix_causal", False)), + FDV2_CACHE_ONLY=bool(getattr(attn_metadata, "fdv2_cache_only", False)), SLIDING_WINDOW=int(sliding_window or 0), ) return o diff --git a/test/python/engine/test_fast_dllm_v2_strategy.py b/test/python/engine/test_fast_dllm_v2_strategy.py new file mode 100644 index 00000000..dcf4632c --- /dev/null +++ b/test/python/engine/test_fast_dllm_v2_strategy.py @@ -0,0 +1,431 @@ +from types import SimpleNamespace + +import torch + +from diffulex.config import Config +from diffulex.engine.kv_cache_manager import AutoKVCacheManager +from diffulex.engine.model_runner import AutoModelRunner +from diffulex.engine.request import AutoReq +from diffulex.engine.scheduler import AutoScheduler +from diffulex.engine.status import DllmBlockStatus, DllmReqStatus +from diffulex.sampler.fast_dllm_v2 import FastdLLMV2Sampler +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 FastDLLMV2Mode, FastDLLMV2Req +from diffulex.strategy.fast_dllm_v2.engine.scheduler import FastDLLMV2Scheduler +from diffulex.strategy.multi_bd.config import MultiBDStrategyConfig +from diffulex.strategy.multi_bd.engine.kv_cache_manager import MultiBDKVCacheManager +from diffulex.strategy.multi_bd.engine.model_runner import MultiBDModelRunner +from diffulex.strategy.multi_bd.engine.request import MultiBDReq +from diffulex.strategy.multi_bd.engine.scheduler import MultiBDScheduler + + +def _runtime_config(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / "config.json").write_text( + """ + { + "model_type": "llama", + "hidden_size": 16, + "intermediate_size": 32, + "num_attention_heads": 2, + "num_hidden_layers": 1, + "num_key_value_heads": 2, + "vocab_size": 128, + "max_position_embeddings": 4096 + } + """, + encoding="utf-8", + ) + return Config( + str(model_dir), + model_name="fast_dllm_v2", + decoding_strategy="fast_dllm_v2", + block_size=8, + buffer_size=4, + tensor_parallel_size=1, + data_parallel_size=1, + device_ids=[0], + num_pages=8, + ) + + +def _req_config(): + return SimpleNamespace( + block_size=8, + buffer_size=4, + mask_token_id=99, + decoding_thresholds=SimpleNamespace( + add_block_threshold=0.1, + semi_complete_threshold=0.9, + accept_threshold=0.95, + token_stability_threshold=0.0, + ), + eos=-1, + max_model_len=256, + ) + + +def _init_decode_req() -> FastDLLMV2Req: + req = FastDLLMV2Req([1, 2, 3, 4, 5, 6, 7, 8]) + req.page_size = 8 + req.init_multi_block(_req_config()) + req.dllm_blocks[0].in_cache() + req.status = DllmReqStatus.DECODING + req.status_history.append(DllmReqStatus.DECODING) + return req + + +def test_fast_dllm_v2_strategy_registers_runtime_components(tmp_path): + cfg = _runtime_config(tmp_path) + + assert cfg.multi_block_prefix_full is False + assert cfg.enable_prefix_caching is True + assert cfg.strategy.name == "fast_dllm_v2" + assert cfg.strategy.sub_block_size == 8 + assert cfg.strategy.block_size == 32 + + assert isinstance(AutoReq.create(cfg, [1, 2, 3]), FastDLLMV2Req) + assert isinstance(AutoKVCacheManager.from_config(cfg), FastDLLMV2KVCacheManager) + assert isinstance(AutoScheduler.from_config(cfg), FastDLLMV2Scheduler) + assert AutoModelRunner._MODULE_MAPPING["fast_dllm_v2"] is FastDLLMV2ModelRunner + + auto_cfg = Config( + cfg.model, + model_name="fast_dllm_v2", + block_size=8, + buffer_size=4, + tensor_parallel_size=1, + data_parallel_size=1, + device_ids=[0], + num_pages=8, + ) + assert auto_cfg.decoding_strategy == "fast_dllm_v2" + + +def test_fast_dllm_v2_can_use_plain_multi_bd_strategy(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / "config.json").write_text( + """ + { + "model_type": "llama", + "hidden_size": 16, + "intermediate_size": 32, + "num_attention_heads": 2, + "num_hidden_layers": 1, + "num_key_value_heads": 2, + "vocab_size": 128, + "max_position_embeddings": 4096 + } + """, + encoding="utf-8", + ) + + cfg = Config( + str(model_dir), + model_name="fast_dllm_v2", + decoding_strategy="multi_bd", + block_size=8, + buffer_size=1, + tensor_parallel_size=1, + data_parallel_size=1, + device_ids=[0], + num_pages=8, + ) + + assert cfg.decoding_strategy == "multi_bd" + assert isinstance(cfg.strategy, MultiBDStrategyConfig) + assert isinstance(AutoReq.create(cfg, [1, 2, 3]), MultiBDReq) + assert isinstance(AutoKVCacheManager.from_config(cfg), MultiBDKVCacheManager) + assert isinstance(AutoScheduler.from_config(cfg), MultiBDScheduler) + assert AutoModelRunner._MODULE_MAPPING["multi_bd"] is MultiBDModelRunner + + +def test_fast_dllm_v2_request_activates_whole_buffer_and_advances_modes(): + req = _init_decode_req() + + req.step() + + assert req.fdv2_mode == FastDLLMV2Mode.FULL_BUFFER_INIT + assert all(block.is_active for block in req.dllm_block_buffer.dllm_blocks) + + req.postprocess() + assert req.fdv2_current_buffer_initialized is True + + first_block = req.fdv2_current_sub_block + for rel_idx in first_block.mask_token_relative_ids: + first_block.write_token(10 + rel_idx, rel_idx) + + req.step() + assert req.fdv2_mode == FastDLLMV2Mode.FULL_BUFFER_INIT + assert req.fdv2_current_sub_block_idx == 2 + + req.fdv2_current_sub_block.write_token(777, 0) + req.step() + assert req.fdv2_mode == FastDLLMV2Mode.SUB_BLOCK_REFINE + + for block in req.dllm_block_buffer.dllm_blocks: + for rel_idx in block.mask_token_relative_ids: + block.write_token(20 + rel_idx, rel_idx) + + req.step() + assert req.fdv2_mode == FastDLLMV2Mode.FINAL_COMMIT + + +def test_fast_dllm_v2_final_commit_seeds_next_block_first_token(): + req = _init_decode_req() + + for block in req.dllm_block_buffer.dllm_blocks: + for rel_idx in block.mask_token_relative_ids: + block.write_token(20 + rel_idx, rel_idx) + block.status = DllmBlockStatus.TO_CACHE + block.commit_ready = True + + req.fdv2_mode = FastDLLMV2Mode.FINAL_COMMIT + req.fdv2_pending_next_token_id = 777 + + req.postprocess() + + first_block = req.dllm_block_buffer.dllm_blocks[0] + assert first_block.block_id == 4 + assert req.token_ids[first_block.start] == 777 + assert req.fdv2_pending_next_token_id is None + assert req.new_tokens == 1 + + +def test_fast_dllm_v2_kv_manager_allocates_current_pages_without_hashing(): + manager = FastDLLMV2KVCacheManager(SimpleNamespace(num_pages=8, page_size=8, enable_prefix_caching=True)) + req = _init_decode_req() + req.status = DllmReqStatus.PREFILLING + req.page_table = [0] + manager._allocate_page(0) + prefix_page = req.token_ids[: req.page_size] + manager.pages[0].update(manager.compute_hash(prefix_page), prefix_page) + req.status = DllmReqStatus.DECODING + + manager.may_append(req) + + assert len(req.page_table) == req.fdv2_read_cache_pages + for page_id in req.page_table[1:]: + assert manager.pages[page_id].hash == -1 + + +def test_fast_dllm_v2_sampler_uses_current_sub_block_and_local_shift(monkeypatch): + req = _init_decode_req() + req.step() + req.postprocess() + req.fdv2_mode = FastDLLMV2Mode.SUB_BLOCK_REFINE + req.fdv2_current_sub_block_idx = 2 + + block = req.fdv2_current_sub_block + block.status = DllmBlockStatus.ACTIVE + token_ids = list(block.token_ids) + for rel_idx in range(4): + block.write_token(10 + rel_idx, rel_idx) + for rel_idx in range(4, 8): + req.token_ids[block.start + rel_idx] = block.mask_token_id + + sampler = FastdLLMV2Sampler() + monkeypatch.setattr( + sampler, + "fetch_attn_metadata", + lambda: SimpleNamespace( + cu_seqlens_q=torch.tensor([0, 8], dtype=torch.int32), + is_prefill=[False], + ), + ) + + logits = torch.zeros(8, 16) + for row in range(8): + logits[row, row] = 10.0 + + out = sampler([req], logits, torch.tensor([0.0])) + req_out = out.true_local_ids_map[str(req.req_id)] + + assert set(req_out.keys()) == {str(block.block_id)} + assert req_out[str(block.block_id)] + assert min(req_out[str(block.block_id)]) >= 4 + assert all(str(other.block_id) not in req_out for other in req.dllm_block_buffer.dllm_blocks if other is not block) + + req.token_ids[block.start : block.end] = token_ids + + +def test_fast_dllm_v2_native_sampler_forces_top1_without_prev_block_gate(): + sampler = FastdLLMV2Sampler() + block = SimpleNamespace( + thresholds=SimpleNamespace(accept_threshold=1.0), + prev_block=SimpleNamespace(is_semi_complete=False), + ) + block.req = SimpleNamespace(fdv2_current_sub_block=block) + confidence = torch.tensor([0.1, 0.7, 0.2]) + initial_confidence = torch.tensor([0.1, 0.7, 0.2]) + sampled_tokens = torch.tensor([10, 11, 12]) + + accepted = sampler._compute_accepted_ids(block, confidence, initial_confidence, sampled_tokens) + + assert accepted.tolist() == [1] + + +def test_fast_dllm_v2_plain_multibd_sampler_uses_prev_block_gate(): + sampler = FastdLLMV2Sampler() + block = SimpleNamespace( + thresholds=SimpleNamespace(accept_threshold=1.0), + prev_block=SimpleNamespace(is_semi_complete=False), + ) + confidence = torch.tensor([0.1, 0.7, 0.2]) + initial_confidence = torch.tensor([0.1, 0.7, 0.2]) + sampled_tokens = torch.tensor([10, 11, 12]) + + accepted = sampler._compute_accepted_ids(block, confidence, initial_confidence, sampled_tokens) + + assert accepted.tolist() == [] + + +class _FakeGraph: + def __init__(self, outputs, value): + self.outputs = outputs + self.value = value + self.replayed = False + + def replay(self): + self.outputs.fill_(self.value) + self.replayed = True + + +class _FakeModel: + @staticmethod + def compute_logits(hidden_states): + return hidden_states + 2 + + +def _fake_graph_runner(): + runner = FastDLLMV2ModelRunner.__new__(FastDLLMV2ModelRunner) + runner.config = SimpleNamespace( + block_size=8, + buffer_size=4, + enable_full_static_runner=True, + ) + runner.enforce_eager = False + runner.fdv2_attention_block_size = 32 + runner.model = _FakeModel() + outputs = torch.zeros(96, 3) + runner.graph_vars = { + "input_ids": torch.zeros(96, dtype=torch.int64), + "positions": torch.zeros(96, dtype=torch.int64), + "slot_mapping": torch.full((96,), -1, dtype=torch.int32), + "context_lens": torch.zeros(8, dtype=torch.int32), + "cu_seqlens_q": torch.zeros(9, dtype=torch.int32), + "cu_seqlens_k": torch.zeros(9, dtype=torch.int32), + "valid_slices": torch.zeros(8, dtype=torch.int32), + "status_table": torch.zeros(8, dtype=torch.int32), + "prefix_lens": torch.zeros(8, dtype=torch.int32), + "padded_prefix_lens": torch.zeros(8, dtype=torch.int32), + "page_tables": torch.full((8, 4), -1, dtype=torch.int32), + "outputs": outputs, + } + runner.graph_bs = { + "full_buffer_init": [32, 64, 96], + "sub_block_cache_only": [8, 16, 32], + "final_commit": [32, 64, 96], + } + runner.graphs = { + ("full_buffer_init", 64): _FakeGraph(outputs, 5), + ("sub_block_cache_only", 32): _FakeGraph(outputs, 7), + ("final_commit", 64): _FakeGraph(outputs, 11), + } + runner.graph_outputs_are_logits = False + return runner + + +def _fake_attn_metadata(*, cache_only: bool, num_reqs: int, q_len: int, mode: FastDLLMV2Mode | None = None): + if mode is None: + mode = FastDLLMV2Mode.SUB_BLOCK_REFINE if cache_only else FastDLLMV2Mode.FULL_BUFFER_INIT + return SimpleNamespace( + has_prefill=False, + fdv2_cache_only=cache_only, + fdv2_mode=int(mode), + block_size=32, + num_reqs=num_reqs, + slot_mapping=torch.arange(num_reqs * q_len, dtype=torch.int32), + context_lens=torch.arange(num_reqs, dtype=torch.int32) + 32, + cu_seqlens_q=torch.arange(num_reqs + 1, dtype=torch.int32) * q_len, + cu_seqlens_k=torch.arange(num_reqs + 1, dtype=torch.int32) * 64, + valid_slices=(torch.arange(num_reqs, dtype=torch.int32) + 1) * q_len, + status_table=torch.full((num_reqs,), 2 if cache_only else 1, dtype=torch.int32), + prefix_lens=torch.zeros(num_reqs, dtype=torch.int32), + padded_prefix_lens=torch.zeros(num_reqs, dtype=torch.int32), + page_tables=torch.zeros(num_reqs, 2, dtype=torch.int32), + ) + + +def test_fast_dllm_v2_cuda_graph_replay_uses_sub_block_bucket_and_static_metadata(): + runner = _fake_graph_runner() + attn_metadata = _fake_attn_metadata(cache_only=True, num_reqs=3, q_len=8) + input_ids = torch.arange(24, dtype=torch.int64) + positions = torch.arange(128, 152, dtype=torch.int64) + + assert runner._fdv2_can_run_decode_graph(input_ids, attn_metadata) + + logits = runner._fdv2_run_decode_graph(input_ids, positions, attn_metadata) + + assert runner.graphs[("sub_block_cache_only", 32)].replayed is True + assert runner.graphs[("full_buffer_init", 64)].replayed is False + assert runner.graphs[("final_commit", 64)].replayed is False + assert logits.shape == (24, 3) + assert torch.all(logits == 9) + assert torch.equal(runner.graph_vars["input_ids"][:24], input_ids) + assert torch.equal(runner.graph_vars["positions"][:24], positions) + assert runner.graph_vars["cu_seqlens_q"][4] == runner.graph_vars["cu_seqlens_q"][3] + assert runner.graph_vars["cu_seqlens_k"][4] == runner.graph_vars["cu_seqlens_k"][3] + assert attn_metadata.fdv2_cache_only is True + assert attn_metadata.block_size == 32 + + +def test_fast_dllm_v2_cuda_graph_replay_uses_full_buffer_bucket(): + runner = _fake_graph_runner() + attn_metadata = _fake_attn_metadata(cache_only=False, num_reqs=2, q_len=32) + input_ids = torch.arange(64, dtype=torch.int64) + positions = torch.arange(64, dtype=torch.int64) + + assert runner._fdv2_can_run_decode_graph(input_ids, attn_metadata) + + logits = runner._fdv2_run_decode_graph(input_ids, positions, attn_metadata) + + assert runner.graphs[("full_buffer_init", 64)].replayed is True + assert runner.graphs[("sub_block_cache_only", 32)].replayed is False + assert runner.graphs[("final_commit", 64)].replayed is False + assert logits.shape == (64, 3) + assert torch.all(logits == 7) + assert attn_metadata.fdv2_cache_only is False + assert attn_metadata.block_size == 32 + + +def test_fast_dllm_v2_cuda_graph_replay_uses_final_commit_bucket(): + runner = _fake_graph_runner() + attn_metadata = _fake_attn_metadata( + cache_only=False, + mode=FastDLLMV2Mode.FINAL_COMMIT, + num_reqs=2, + q_len=32, + ) + input_ids = torch.arange(64, dtype=torch.int64) + positions = torch.arange(64, dtype=torch.int64) + + assert runner._fdv2_can_run_decode_graph(input_ids, attn_metadata) + + logits = runner._fdv2_run_decode_graph(input_ids, positions, attn_metadata) + + assert runner.graphs[("final_commit", 64)].replayed is True + assert runner.graphs[("full_buffer_init", 64)].replayed is False + assert runner.graphs[("sub_block_cache_only", 32)].replayed is False + assert logits.shape == (64, 3) + assert torch.all(logits == 13) + assert attn_metadata.fdv2_cache_only is False + assert attn_metadata.fdv2_mode == int(FastDLLMV2Mode.FINAL_COMMIT) + + +def test_fast_dllm_v2_cuda_graph_uses_exact_batch_buckets(): + assert FastDLLMV2ModelRunner._graph_seq_batch_sizes(0) == [] + assert FastDLLMV2ModelRunner._graph_seq_batch_sizes(5) == [1, 2, 3, 4, 5] diff --git a/test/python/kernel/test_dllm_flash_attn_chunked_prefill_unified_kernel.py b/test/python/kernel/test_dllm_flash_attn_chunked_prefill_unified_kernel.py index e139161a..fcb000e0 100644 --- a/test/python/kernel/test_dllm_flash_attn_chunked_prefill_unified_kernel.py +++ b/test/python/kernel/test_dllm_flash_attn_chunked_prefill_unified_kernel.py @@ -154,6 +154,8 @@ def call_chunked_prefill_kernel( IS_PREFIX_FULL=is_prefix_full, MASK_PREFIX_HOLE=False, PREFIX_CAUSAL=False, + FDV2_CACHE_ONLY=False, + SLIDING_WINDOW=0, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, ) diff --git a/test/python/kernel/test_vllm_attention_perf.py b/test/python/kernel/test_vllm_attention_perf.py index 2448fc74..7385bfa2 100644 --- a/test/python/kernel/test_vllm_attention_perf.py +++ b/test/python/kernel/test_vllm_attention_perf.py @@ -338,6 +338,8 @@ def _diffulex_attention_fixed(data: AttentionBenchData, case: AttentionBenchCase IS_PREFIX_FULL=metadata.is_prefix_full, MASK_PREFIX_HOLE=bool(getattr(metadata, "mask_prefix_hole", False)), PREFIX_CAUSAL=bool(getattr(metadata, "prefix_causal", False)), + FDV2_CACHE_ONLY=False, + SLIDING_WINDOW=int(data.case.sliding_window or 0), ) return out