diff --git a/.gitignore b/.gitignore index bc8329f2..cc0144f9 100755 --- a/.gitignore +++ b/.gitignore @@ -24,10 +24,16 @@ TestResults.xml *.pyc *.egg-info lora_weight* +*.log log/ dist/ build/ cache/ uv.lock ckpt/ -data/ \ No newline at end of file +data/ +tilelang +autotuner.log +Fast-dLLM +Discrete-Diffusion-Forcing +position_explanation.md \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json index 0224a397..61675775 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -6,11 +6,11 @@ "configurations": [ - - + + { "name": "Python Debugger: Current File", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${file}", "console": "integratedTerminal", @@ -53,7 +53,7 @@ }, { "name": "PyDbg: `diffulex` Qwen3", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_qwen_dvllm.py", "console": "integratedTerminal", @@ -64,7 +64,7 @@ }, { "name": "PyDbg: `diffulex` Dream `HumanEval`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_dream_dvllm_human_eval.py", "console": "integratedTerminal", @@ -76,7 +76,7 @@ }, { "name": "PyDbg: `diffulex` Dream `GSM8K`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_dream_dvllm_gsm8k.py", "console": "integratedTerminal", @@ -86,9 +86,21 @@ // "CUDA_VISIBLE_DEVICES": "0,1" } }, + { + "name": "PyDbg: `diffulex` Fast-DLLM-V2 `GSM8K`", + "type": "debugpy", + "request": "launch", + "program": "${workspaceFolder}/examples/test_fastdllmv2_diffulex_gsm8k.py", + "console": "integratedTerminal", + "env": { + // "TORCHINDUCTOR_DISABLE": "1", + // "TRITON_INTERPRET": "1", + // "CUDA_VISIBLE_DEVICES": "0,1" + } + }, { "name": "PyDbg: `diffulex` LLaDA `HumanEval`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_llada_dvllm_human_eval.py", "console": "integratedTerminal", @@ -100,7 +112,7 @@ }, { "name": "PyDbg: `diffulex` kernel func `load_kvcache_kernel`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_dllm_kv_cache_load.py", "console": "integratedTerminal", @@ -111,7 +123,7 @@ }, { "name": "PyDbg: `diffulex` kernel func `chunked_prefill_paged_decode`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_dllm_decoding_kernel.py", "console": "integratedTerminal", @@ -122,7 +134,7 @@ }, { "name": "PyDbg: `diffulex` kernel func `causal_lm_decode_attention_fwd`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_causal_lm_decoding_kernel.py", "console": "integratedTerminal", @@ -133,7 +145,7 @@ }, { "name": "PyDbg: `diffulex` kernel func `store_kvcache_kernel_diffusion_lm`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_dllm_kv_cache_store.py", "console": "integratedTerminal", diff --git a/imgs/logo_lr.png b/assets/logo_lr.png similarity index 100% rename from imgs/logo_lr.png rename to assets/logo_lr.png diff --git a/diffulex/attention/__init__.py b/diffulex/attention/__init__.py index e38b5ff8..a390a61d 100644 --- a/diffulex/attention/__init__.py +++ b/diffulex/attention/__init__.py @@ -1,2 +1,24 @@ -from .attn_impl import Attention -from .metadata import fetch_attn_metadata, set_fetch_fn_for_attn_metadata, AttnMetaDataBase \ No newline at end of file +from . import metadata +from .metadata import set_fetch_fn_for_attn_metadata, AttnMetaDataBase + +# Create a proxy that dynamically accesses fetch_attn_metadata from the metadata module +# This ensures we always get the current value, not a stale copy from __init__.py +class _FetchAttnMetadataProxy: + """Proxy object that dynamically accesses fetch_attn_metadata from metadata module.""" + def __call__(self, *args, **kwargs): + return metadata.fetch_attn_metadata(*args, **kwargs) + + def __repr__(self): + return repr(metadata.fetch_attn_metadata) + +fetch_attn_metadata = _FetchAttnMetadataProxy() + + +def __getattr__(name): + """Lazy import to avoid circular deps during module init.""" + if name == "Attention": + from .attn_impl import Attention + return Attention + if name == "fetch_attn_metadata": + return metadata.fetch_attn_metadata + raise AttributeError(f"module {__name__} has no attribute {name}") \ No newline at end of file diff --git a/diffulex/attention/attn_impl.py b/diffulex/attention/attn_impl.py index 91907ea0..9ec5f7fc 100644 --- a/diffulex/attention/attn_impl.py +++ b/diffulex/attention/attn_impl.py @@ -1,20 +1,14 @@ -import os import torch - import torch.nn as nn - -from functools import lru_cache, partial from einops import rearrange -from torch.nn.attention.flex_attention import create_block_mask -from flash_attn import flash_attn_varlen_func -from transformers.integrations.flex_attention import compile_friendly_flex_attention as flex_attention -from diffulex.attention.ops import ( - causal_lm_flash_decoding, diffusion_lm_flash_decoding, diffusion_lm_parallel_flash_decoding, - store_kvcache_unified_layout, store_kvcache_distinct_layout, load_kvcache, - CHECK_STORING, CHECK_LOADING, CHECK_ATTENTION +from diffulex_kernel import ( + store_kvcache_distinct_layout, + store_kvcache_unified_layout, + dllm_flash_attn_decode, + dllm_flash_attn_prefill ) -from diffulex.attention.metadata import AttnMetaDataBase, fetch_attn_metadata +from diffulex.attention.metadata import AttnMetaDataBase class Attention(nn.Module): @@ -31,89 +25,48 @@ def __init__( self.scale = scale self.num_kv_heads = num_kv_heads self.k_cache = self.v_cache = torch.tensor([]) - is_rtx_xx90 = lambda x: "4090" in x or "3090" in x - kernel_options = { - "BLOCK_M": 64, - "BLOCK_N": 64, - "BLOCK_M1": 32, - "BLOCK_N1": 64, - "BLOCK_M2": 64, - "BLOCK_N2": 32, - } if is_rtx_xx90(torch.cuda.get_device_name(0)) else None - self.attention = torch.compile( - partial(flex_attention, kernel_options=kernel_options, enable_gqa=True, - return_lse=False, training=False), dynamic=True) - self._block_mask_cache = {} - - @lru_cache(maxsize=32) - def dllm_block_mask(self, block_mask: torch.Tensor, - B: int, H: int, Q_LEN: int, KV_LEN: int, device: str): - cache_key = (B, H, Q_LEN, KV_LEN, device) - def _mask_mod(batch, head, token_q, token_kv): - return block_mask[token_q, token_kv] - if cache_key not in self._block_mask_cache: - self._block_mask_cache[cache_key] = create_block_mask( - _mask_mod, B, H, Q_LEN, KV_LEN, device=device - ) - return self._block_mask_cache[cache_key] + + self.q_shape = { + 'nh': self.num_heads, + 'hd': self.head_dim, + } + self.kv_shape = { + 'nkvh': self.num_kv_heads, + 'hd': self.head_dim, + } + # Import the specified fetch function + from diffulex.attention import fetch_attn_metadata + self.fetch_attn_metadata = fetch_attn_metadata + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: list[torch.Tensor] | None = None) -> torch.Tensor: # Reshape - q = q.view(-1, self.num_heads, self.head_dim) - k = k.view(-1, self.num_kv_heads, self.head_dim) - v = v.view(-1, self.num_kv_heads, self.head_dim) + q = rearrange(q, 's (nh hd) -> s nh hd', **self.q_shape) + k = rearrange(k, 's (nkvh hd) -> s nkvh hd', **self.kv_shape) + v = rearrange(v, 's (nkvh hd) -> s nkvh hd', **self.kv_shape) - attn_metadata: AttnMetaDataBase = fetch_attn_metadata() + attn_metadata: AttnMetaDataBase = self.fetch_attn_metadata() k_cache, v_cache = self.k_cache, self.v_cache is_unified_layout = attn_metadata.kv_cache_layout == "unified" # Fast Store KV cache if k_cache.numel() and v_cache.numel(): - if not (not attn_metadata.need_kv_cache_store): + if attn_metadata.need_kv_cache_store: store_kvcache = store_kvcache_unified_layout if is_unified_layout else store_kvcache_distinct_layout store_kvcache(k, v, k_cache, v_cache, attn_metadata.slot_mapping, attn_metadata) - # CHECK_STORING(k_cache, v_cache, k, v, context) - transpose_fn = lambda x: rearrange(x, 's h d -> 1 h s d').contiguous() - # Prefill / Decode logic TODO: Replace the Flex Attention Prefilling + # Prefill / Decode logic if attn_metadata.is_prefill: - # Block PK if attn_metadata.block_tables is not None: # TODO: Implement Prefix Caching pass - - # Attention computation - q_t, k_t, v_t = [transpose_fn(t) for t in (q, k, v)] - - B, H, S, _ = q_t.shape - block_mask = self.dllm_block_mask(attn_metadata.block_mask, B, H, S, S, str(q.device)) - o = self.attention(q_t, k_t, v_t, block_mask=block_mask) + o = dllm_flash_attn_prefill(q, k, v, self.scale, attn_metadata) else: - config = attn_metadata.seqs[0].config - diffusion_block_size = config.diffusion_block_size if is_unified_layout: - k_comb, v_comb = load_kvcache(self.k_cache, self.v_cache, attn_metadata, k, v) - o = flash_attn_varlen_func(q, k_comb, v_comb, - attn_metadata.cu_seqlens_q, attn_metadata.cu_seqlens_k, - attn_metadata.max_seqlen_q, attn_metadata.max_seqlen_k, - softmax_scale=self.scale, block_table=None) + o = dllm_flash_attn_decode(q, k, v, k_cache, v_cache, self.scale, attn_metadata) else: - # FIXME: Kernel not ok... - o = torch.empty_like(q).to(q.device).to(q.dtype) - q, k, o, k_cache, v_cache = map(lambda x: x.to(torch.float32), (q, k, o, k_cache, v_cache)) - diffusion_lm_parallel_flash_decoding( - q, k, v, o, str(k_cache.dtype), k_cache, v_cache, - attn_metadata.block_tables, attn_metadata.cu_seqlens_q, attn_metadata.total_lens, - max(attn_metadata.total_lens), max(attn_metadata.seq_lens), 1.0, 1.0, - diffusion_block_size, attn_metadata.block_mask - ) - CHECK_ATTENTION(o, q, k, v, k_cache, v_cache, attn_metadata) + raise NotImplementedError("Distinct layout is not supported yet...") # Final reshape - if not attn_metadata.is_prefill: - o = o.view(-1, self.num_heads * self.head_dim).contiguous() - elif attn_metadata.is_prefill: - o = rearrange(o, '1 h s d -> s (h d)').contiguous() - - return o \ No newline at end of file + return rearrange(o, 's nh hd -> s (nh hd)').contiguous() \ No newline at end of file diff --git a/diffulex/attention/metadata.py b/diffulex/attention/metadata.py index 6b157e00..75c290ef 100644 --- a/diffulex/attention/metadata.py +++ b/diffulex/attention/metadata.py @@ -14,7 +14,14 @@ class AttnMetaDataBase: slot_mapping: torch.Tensor | None = None context_lens: torch.Tensor | None = None block_tables: torch.Tensor | None = None - + page_block_size: int = 32 + attn_type: str = "block_attention" + diffusion_block_size: int = 32 + decode_mode: str = "static" + + @property + def num_seqs(self) -> int: + return len(self.cu_seqlens_q) - 1 FN_TYPE_AttnMetaDataFetch = Callable[[], AttnMetaDataBase] @@ -22,4 +29,17 @@ class AttnMetaDataBase: def set_fetch_fn_for_attn_metadata(fn: FN_TYPE_AttnMetaDataFetch) -> None: global fetch_attn_metadata - fetch_attn_metadata = fn \ No newline at end of file + fetch_attn_metadata = fn + +WARMING_UP = False + +def set_warming_up(is_warming_up: bool) -> None: + global WARMING_UP + WARMING_UP = is_warming_up + +def is_warming_up() -> bool: + return WARMING_UP + +def reset_warming_up() -> None: + global WARMING_UP + WARMING_UP = False \ No newline at end of file diff --git a/diffulex/attention/ops/triton_decode_attn_dlm.py b/diffulex/attention/ops/triton_decode_attn_dlm.py deleted file mode 100755 index e39ed1e0..00000000 --- a/diffulex/attention/ops/triton_decode_attn_dlm.py +++ /dev/null @@ -1,120 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: D2F - -# Organization: SJTU DENG Lab -# Author: Drew Jin (JIN. Yijie, @drewjin) -# Date: 2025-08-07 -# Email: drewjin0827@gmail.com -# All rights reserved. - -import torch -import triton - -import triton.language as tl - -from diffulex.legacy.utils.context import ContextForDiffusionLM - - -def CHECK_ATTENTION(o: torch.Tensor, q: torch.Tensor, k_new: torch.Tensor, v_new: torch.Tensor, - k_cache: torch.Tensor, v_cache: torch.Tensor, context: ContextForDiffusionLM): - """ - Check the attention output against the input tensors. - """ - from einops import rearrange - from torch.nn.functional import scaled_dot_product_attention as sdpa - from torch.nn.attention import SDPBackend, sdpa_kernel - - from diffulex.legacy.layers.attention.ops import load_kvcache - - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - - h_dim = v_cache.shape[-2] - x = k_cache.shape[-1] - k_cache_unified = rearrange(k_cache, "b h n s x -> b s h (n x)", n=h_dim // x, x=x).contiguous() - v_cache_unified = rearrange(v_cache, "b h d s -> b s h d").contiguous() - - transpose_fn = lambda x: rearrange(x, 's h d -> 1 h s d').contiguous() - k, v = load_kvcache(k_cache_unified, v_cache_unified, context, k_new, v_new) - q, k, v = map(transpose_fn, (q, k, v)) - mask = context.block_mask_for_checking - with sdpa_kernel(SDPBackend.MATH): - ref_o = sdpa(q, k, v, attn_mask=mask, enable_gqa=True) - - ref_o = rearrange(ref_o, '1 h s d -> s h d') - assert torch.allclose(o, ref_o, atol=1e-3, rtol=1e-3), "Attention output does not match reference!" - - -@triton.jit -def dlm_flash_decoding_kernel(q_ptr, k_ptr, v_ptr, o_ptr, mask_ptr, softmax_scale, - k_cache_ptr, v_cache_ptr, block_tables_ptr, - cu_seqlens_q_ptr, total_lens_ptr, ctx_lens_ptr, - q_stride_m, q_stride_nh, q_stride_d, - k_stride_n, k_stride_nh, k_stride_d, - v_stride_n, v_stride_nh, v_stride_d, - o_stride_m, o_stride_nh, o_stride_d, - mask_stride_m, mask_stride_n, - k_cache_stride_nblks, k_cache_stride_h, k_cache_stride_dx, k_cache_stride_blk_sz, k_cache_stride_x, - v_cache_stride_nblks, v_cache_stride_h, v_cache_stride_d, v_cache_stride_blk_sz, - block_tables_stride_nseqs, block_tables_stride_nblks, - cu_seqlens_q_ptr_stride, total_lens_ptr_stride, ctx_lens_ptr_stride, - NUM_HEADS: tl.constexpr, HEAD_DIM: tl.constexpr, KV_HEAD_GROUP_SIZE: tl.constexpr, - BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, x: tl.constexpr, BLOCK_SIZE: tl.constexpr, - NUM_UNROLL_CACHE: tl.constexpr = 4, NUM_UNROLL_Q: tl.constexpr = 1): - pass - - -def diffusion_lm_flash_decoding(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor, - k_cache: torch.Tensor, v_cache: torch.Tensor, block_tables: torch.Tensor, - cu_seqlens_q: torch.Tensor, seq_lens: torch.Tensor, total_lens: torch.Tensor, ctx_lens: torch.Tensor, - max_total_len: int | None = None, max_seq_len: int | None = None, - diffusion_block_size: int = 32): - ''' - FIXME - q: [TotalInputLength, NumHeads, HeadDim] - k: [TotalInputLength, NumHeads, HeadDim] - v: [TotalInputLength, NumHeads, HeadDim] - mask: [TotalInputLength, TotalInputLength] - k_cache: [NumBlocks, NumHeads, HeadDim // x, BlockSize, x] - v_cache: [NumBlocks, NumHeads, HeadDim, BlockSize] - block_tables: [NumSeqs, MaxSeqNumBlocks] # NumSeqs == BatchSize - ... - ''' - is_pow_of_2 = lambda x: (x & (x - 1)) == 0 and x > 0 - assert k_cache.shape[-2] == v_cache.shape[-1], "BLOCK_SIZE between k_cache and v_cache must match" - assert k.shape == v.shape, "k, v must have the same shape" - assert k.shape[1] == k_cache.shape[1] == v_cache.shape[1], "Number of heads must match" - assert q.shape[1] % k.shape[1] == 0, "Number of heads in q must be a multiple of the number of heads in k and v" - assert k_cache.shape[-3] * k_cache.shape[-1] == v_cache.shape[-2] == q.shape[-1], "Head dimension must match" - assert is_pow_of_2(q.shape[-1]) and is_pow_of_2(k_cache.shape[-3] * k_cache.shape[-1]), \ - "Head dimension must be a multiple of 2 for triton kernel compatibility" - assert len(seq_lens) == len(ctx_lens) == len(total_lens) == len(cu_seqlens_q) - 1 == len(block_tables), \ - "Number of sequences must match across all inputs" - - BLOCK_SIZE = k_cache.shape[-2] # BLOCK_SIZE or PAGE_SIZE of paged kv cache - NUM_SEQS = len(ctx_lens) - NUM_HEADS = q.shape[1] - o = torch.empty_like(q).to(q.device).to(q.dtype) - x = k_cache.shape[-1] - max_seq_len = max_seq_len if max_seq_len is not None else max(seq_lens) - max_total_len = max_total_len if max_total_len is not None else max(total_lens) - softmax_scale = 1.0 / (k.shape[-1] ** 0.5) - - KV_HEAD_GROUP_SIZE = q.shape[1] // k.shape[1] - HEAD_DIM = q.shape[-1] - BLOCK_M = BLOCK_N = diffusion_block_size * 2 - GRID = (NUM_SEQS, NUM_HEADS, triton.cdiv(max_seq_len, BLOCK_M)) - - dlm_flash_decoding_kernel[GRID]( - q, k, v, o, mask, softmax_scale, k_cache, v_cache, block_tables, - cu_seqlens_q, total_lens, ctx_lens, - *q.stride(), *k.stride(), *v.stride(), *o.stride(), *mask.stride(), - *k_cache.stride(), *v_cache.stride(), *block_tables.stride(), - cu_seqlens_q.stride(0), total_lens.stride(0), ctx_lens.stride(0), - NUM_HEADS=NUM_HEADS, HEAD_DIM=HEAD_DIM, - KV_HEAD_GROUP_SIZE=KV_HEAD_GROUP_SIZE, - BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, x=x, - BLOCK_SIZE=BLOCK_SIZE, - NUM_UNROLL_CACHE=4, NUM_UNROLL_Q=1 - ) - return o \ No newline at end of file diff --git a/diffulex/config.py b/diffulex/config.py index 664b4cdc..96af47ce 100755 --- a/diffulex/config.py +++ b/diffulex/config.py @@ -38,7 +38,7 @@ class Config: enforce_eager: bool = False hf_config: AutoConfig | None = None eos: int = -1 - kvcache_block_size: int = 256 + kvcache_block_size: int = 32 num_kvcache_blocks: int = -1 k_cache_hdim_split_factor_x: int = 8 kv_cache_layout: str = "unified" # "unified" or "distinct" diff --git a/diffulex/engine/model_runner.py b/diffulex/engine/model_runner.py index bdb54a34..0316dd0c 100755 --- a/diffulex/engine/model_runner.py +++ b/diffulex/engine/model_runner.py @@ -9,7 +9,7 @@ from multiprocessing.shared_memory import SharedMemory from diffulex.config import Config -from diffulex.layer.sampler import AutoSampler +from diffulex.sampler import AutoSampler from diffulex.engine.sequence import SequenceBase from diffulex.model import AutoModelForDiffusionLM from diffulex.engine.strategy_registry import DiffulexStrategyRegistry @@ -123,9 +123,115 @@ def warmup_model(self): """Model-specific warmup logic.""" pass - @abstractmethod def allocate_kv_cache(self): - pass + config = self.config + hf_config = config.hf_config + free, total = torch.cuda.mem_get_info() + used = total - free + peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"] + current = torch.cuda.memory_stats()["allocated_bytes.all.current"] + num_kv_heads = getattr( + hf_config, + "num_key_value_heads", + getattr(hf_config, "n_kv_heads", None), + ) // self.world_size + + if hasattr(hf_config, "head_dim"): + head_dim = hf_config.head_dim + elif hasattr(hf_config, "hidden_size") and hasattr(hf_config, "num_attention_heads"): + head_dim = hf_config.hidden_size // hf_config.num_attention_heads + else: + raise AttributeError(f"Cannot determine head_dim from config: {type(hf_config)}") + + dtype = ( + hf_config.torch_dtype + if hasattr(hf_config, "torch_dtype") and hf_config.torch_dtype + else torch.bfloat16 + ) + block_bytes = ( + 2 + * hf_config.num_hidden_layers + * self.block_size + * num_kv_heads + * head_dim + * dtype.itemsize + ) + get_num_kvcache_blocks = ( + lambda gpu_memory_utilization: int(total * gpu_memory_utilization - used - peak + current) + // block_bytes + ) + try: + num_kvcache_blocks = get_num_kvcache_blocks(config.gpu_memory_utilization) + assert num_kvcache_blocks > 0 + except Exception: + gpu_memory_utilization = config.gpu_memory_utilization + while num_kvcache_blocks <= 200: + print( + "Warning: GPU memory utilization " + f"{gpu_memory_utilization} is too low to allocate kv cache. " + "Automatically adding 0.05." + ) + gpu_memory_utilization += 0.05 + num_kvcache_blocks = get_num_kvcache_blocks(gpu_memory_utilization) + print( + f"Set gpu_memory_utilization to {gpu_memory_utilization:.2f} " + "to allocate kv cache." + ) + config.gpu_memory_utilization = gpu_memory_utilization + + config.num_kvcache_blocks = num_kvcache_blocks + print( + "Allocated {num_blocks} blocks of size {block_size} for kv cache on rank {rank}.".format( + num_blocks=config.num_kvcache_blocks, + block_size=self.block_size, + rank=self.rank, + ) + ) + + if config.kv_cache_layout == "distinct": + x = config.k_cache_hdim_split_factor_x + self.k_cache = torch.zeros( + hf_config.num_hidden_layers, + config.num_kvcache_blocks, + num_kv_heads, + head_dim // x, + self.block_size, + x, + ) + self.v_cache = torch.zeros( + hf_config.num_hidden_layers, + config.num_kvcache_blocks, + num_kv_heads, + head_dim, + self.block_size, + ) + layer_id = 0 + for module in self.model.modules(): + if hasattr(module, "k_cache") and hasattr(module, "v_cache"): + module.k_cache = self.k_cache[layer_id] + module.v_cache = self.v_cache[layer_id] + layer_id += 1 + elif config.kv_cache_layout == "unified": + self.kv_cache = torch.zeros( + 2, + hf_config.num_hidden_layers, + config.num_kvcache_blocks, + self.block_size, + num_kv_heads, + head_dim, + ) + layer_id = 0 + for module in self.model.modules(): + if hasattr(module, "k_cache") and hasattr(module, "v_cache"): + module.k_cache = self.kv_cache[0, layer_id] + module.v_cache = self.kv_cache[1, layer_id] + layer_id += 1 + else: + raise ValueError( + "Unsupported kv_cache_layout: {layout}. Supported values are 'distinct' and 'unified'.".format( + layout=config.kv_cache_layout + ) + ) def prepare_block_tables(self, seqs: list[SequenceBase]): max_len = max(len(seq.block_table) for seq in seqs) diff --git a/diffulex/engine/sequence.py b/diffulex/engine/sequence.py index 6fd29d4e..b467a906 100755 --- a/diffulex/engine/sequence.py +++ b/diffulex/engine/sequence.py @@ -29,7 +29,6 @@ def __init__(self, token_ids: list[int], sampling_params: SamplingParams = Sampl self.status = SequenceStatus.WAITING self.token_ids = copy(token_ids) self.last_token = token_ids[-1] - self.num_tokens = len(token_ids) self.num_prompt_tokens = len(token_ids) self.num_cached_tokens = 0 self.block_table: list[int] = [] @@ -38,12 +37,17 @@ def __init__(self, token_ids: list[int], sampling_params: SamplingParams = Sampl self.max_tokens = sampling_params.max_tokens self.ignore_eos = sampling_params.ignore_eos self.new_tokens = 0 - + self.meet_eos = False + def __len__(self) -> int: return self.num_tokens def __getitem__(self, key) -> int: return self.token_ids[key] + + @property + def num_tokens(self) -> int: + return len(self.token_ids) @property def is_finished(self) -> bool: diff --git a/diffulex/engine/tp_worker.py b/diffulex/engine/tp_worker.py index 9978dce9..3ea53c56 100755 --- a/diffulex/engine/tp_worker.py +++ b/diffulex/engine/tp_worker.py @@ -68,7 +68,7 @@ def step(self): sample_output = self.model_runner.call("run", seqs, is_prefill) n_diff_steps = self.scheduler.postprocess(seqs, sample_output) outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished] - num_tokens = sum(seq.input_num_tokens + seq.new_tokens for seq in seqs) if is_prefill else sum(seq.new_tokens for seq in seqs) + num_tokens = sum(seq.num_tokens for seq in seqs) if is_prefill else sum(seq.new_tokens for seq in seqs) # Diffusion decoding modifies tokens in-place; we currently don't stream intermediate edits deltas = [] return outputs, num_tokens, is_prefill, n_diff_steps, deltas diff --git a/diffulex/layer/sampler.py b/diffulex/layer/sampler.py deleted file mode 100644 index 7afc8b8d..00000000 --- a/diffulex/layer/sampler.py +++ /dev/null @@ -1,216 +0,0 @@ -import torch - -import torch.nn as nn -import torch.nn.functional as F -import torch.distributions as dists - -from dataclasses import dataclass -from easydict import EasyDict as edict - -from diffulex.config import Config -from diffulex.attention import fetch_attn_metadata - - -class SamplerForDiffusionLM(nn.Module): - def __init__(self): - super().__init__() - - def top_p_logits(self, logits, top_p): - sorted_logits, sorted_indices = torch.sort(logits, descending=True) - cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) - sorted_indices_to_remove = cumulative_probs > top_p - sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() - sorted_indices_to_remove[..., 0] = 0 - - mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device) - mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove) - logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min) - return logits - - def top_k_logits(self, logits, top_k): - top_k = min(top_k, logits.size(-1)) - indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] - logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min) - return logits - - def sample_tokens(self, logits, temperature=0.0, top_p=None, top_k=None, - margin_confidence=False, neg_entropy=False): - if temperature > 0: - logits = logits / temperature - if top_p is not None and top_p < 1: - logits = self.top_p_logits(logits, top_p) - if top_k is not None: - logits = self.top_k_logits(logits, top_k) - probs = torch.softmax(logits, dim=-1) - - if temperature > 0: - try: - x0 = dists.Categorical(probs=probs).sample() - initial_confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1) - except: - initial_confidence, x0 = probs.max(dim=-1) - else: - initial_confidence, x0 = probs.max(dim=-1) - - confidence = initial_confidence.clone() - - if margin_confidence: - sorted_probs, _ = torch.sort(probs, dim=-1, descending=True) - top1_probs = sorted_probs[:, 0] - top2_probs = sorted_probs[:, 1] - confidence = top1_probs - top2_probs - - if neg_entropy: - epsilon = 1e-10 - log_probs = torch.log(probs + epsilon) - confidence = torch.sum(probs * log_probs, dim=-1) - - return confidence, x0, initial_confidence - - -@dataclass -class SampleOutputForDiffusionLM: - true_local_ids_map: dict[str, dict[str, list[int]]] - accepted_ids_map: dict[str, list[int]] - sampled_tokens_map: dict[str, dict[str, list[int]]] - - def __post_init__(self): - self.accepted_ids_map = edict(self.accepted_ids_map) - self.sampled_tokens_map = edict(self.sampled_tokens_map) - self.true_local_ids_map = edict(self.true_local_ids_map) - - -class SamplerForDream(SamplerForDiffusionLM): - def _shift_logits(self, logits, last_logit=None): - if logits.shape[1] == 0: - print("Warning: logits sequence length is 0, returning empty logits") - raise Exception("logits sequence length is 0") - - shifted_logits = torch.zeros_like(logits) - shifted_logits[1:, ...] = logits[:-1, ...] - if last_logit is not None: - shifted_logits[0, ...] = last_logit - return shifted_logits - shifted_logits[0, ...] = 1.0 - return shifted_logits - - def forward(self, logits: torch.Tensor, temperatures: torch.Tensor, - top_p=None, top_k=None, margin_confidence=False, neg_entropy=False): - context = fetch_attn_metadata() - seqs = context.seqs - split_logits = torch.split(logits, [len(seq) for seq in seqs] if context.is_prefill else context.seq_lens, dim=0) - accepted_ids_map = {} - sampled_tokens_map = {} - true_local_ids_map = {} - for temperature, seq, seq_logits in zip(temperatures, seqs, split_logits): - true_local_ids_sub_map = {} - accepted_ids_sub_map = {} - sampled_tokens_sub_map = {} - shifted_logits = self._shift_logits(seq_logits, seq.cached_or_caching_last_token_id) - for block_id, block in enumerate(seq.diffusion_blocks): - if not block.is_active or sum(block.local_mask_tokens) == 0: - continue - - if len(block.global_mask_token_ids) > 0: - mask_token_logits = shifted_logits[block.global_mask_token_ids, ...] - confidence, sampled_tokens, initial_confidence = self.sample_tokens( - mask_token_logits, - temperature, - top_p=top_p, - top_k=top_k, - neg_entropy=(neg_entropy == "neg_entropy"), - margin_confidence=(margin_confidence == "margin_confidence") - ) - - if block.pre_block_complete: - high_conf_indices = torch.where(initial_confidence > block.accept_threshold)[0] - if len(high_conf_indices) == 0: - number_transfer_tokens = 1 - _, transfer_index = torch.topk(confidence, number_transfer_tokens) - else: - transfer_index = torch.tensor([], device=sampled_tokens.device, dtype=torch.long) - accepted_ids = torch.unique(torch.cat([transfer_index, high_conf_indices])) - else: - high_conf_indices = torch.where(initial_confidence > block.accept_threshold)[0] - accepted_ids = high_conf_indices - - true_local_ids_sub_map[str(block_id)] = [block.local_mask_token_ids[accepted_id] for accepted_id in accepted_ids.tolist()] - accepted_ids_sub_map[str(block_id)] = accepted_ids.tolist() - sampled_tokens_sub_map[str(block_id)] = sampled_tokens - - seq_idx = str(seq.seq_id) - true_local_ids_map[seq_idx] = true_local_ids_sub_map - accepted_ids_map[seq_idx] = accepted_ids_sub_map - sampled_tokens_map[seq_idx] = sampled_tokens_sub_map - - return SampleOutputForDiffusionLM( - true_local_ids_map=true_local_ids_map, - accepted_ids_map=accepted_ids_map, - sampled_tokens_map=sampled_tokens_map - ) - - -class SamplerForLLaDA(SamplerForDiffusionLM): - def forward(self, logits: torch.Tensor, temperatures: torch.Tensor, - top_p=None, top_k=None, margin_confidence=False, neg_entropy=False): - context = fetch_attn_metadata() - seqs = context.seqs - split_logits = torch.split(logits, [len(seq) for seq in seqs] if context.is_prefill else context.seq_lens, dim=0) - accepted_ids_map = {} - sampled_tokens_map = {} - true_local_ids_map = {} - for temperature, seq, seq_logits in zip(temperatures, seqs, split_logits): - true_local_ids_sub_map = {} - accepted_ids_sub_map = {} - sampled_tokens_sub_map = {} - for block_id, block in enumerate(seq.diffusion_blocks): - if not block.is_active or sum(block.local_mask_tokens) == 0: - continue - - if len(block.global_mask_token_ids) > 0: - mask_token_logits = seq_logits[block.global_mask_token_ids, ...] - confidence, sampled_tokens, initial_confidence = self.sample_tokens( - mask_token_logits, - temperature, - top_p=top_p, - top_k=top_k, - neg_entropy=(neg_entropy == "neg_entropy"), - margin_confidence=(margin_confidence == "margin_confidence") - ) - - if block.pre_block_complete: - high_conf_indices = torch.where(initial_confidence > block.accept_threshold)[0] - if len(high_conf_indices) == 0: - number_transfer_tokens = 1 - _, transfer_index = torch.topk(confidence, number_transfer_tokens) - else: - transfer_index = torch.tensor([], device=sampled_tokens.device, dtype=torch.long) - accepted_ids = torch.unique(torch.cat([transfer_index, high_conf_indices])) - else: - high_conf_indices = torch.where(initial_confidence > block.accept_threshold)[0] - accepted_ids = high_conf_indices - - true_local_ids_sub_map[str(block_id)] = [block.local_mask_token_ids[accepted_id] for accepted_id in accepted_ids.tolist()] - accepted_ids_sub_map[str(block_id)] = accepted_ids.tolist() - sampled_tokens_sub_map[str(block_id)] = sampled_tokens - - seq_idx = str(seq.seq_id) - true_local_ids_map[seq_idx] = true_local_ids_sub_map - accepted_ids_map[seq_idx] = accepted_ids_sub_map - sampled_tokens_map[seq_idx] = sampled_tokens_sub_map - - return SampleOutputForDiffusionLM( - true_local_ids_map=true_local_ids_map, - accepted_ids_map=accepted_ids_map, - sampled_tokens_map=sampled_tokens_map - ) - - -class AutoSampler: - MODEL_MAPPING = { - "dream": SamplerForDream, - "llada": SamplerForLLaDA - } - @classmethod - def from_config(cls, config: Config): - return cls.MODEL_MAPPING[config.model_name]() \ No newline at end of file diff --git a/diffulex/legacy/__init__.py b/diffulex/legacy/__init__.py deleted file mode 100755 index c71384e5..00000000 --- a/diffulex/legacy/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from diffulex.legacy.llm import LLM -from diffulex.legacy.sampling_params import SamplingParams diff --git a/diffulex/legacy/layers/attention/ops/__init__.py b/diffulex/legacy/layers/attention/ops/__init__.py deleted file mode 100755 index 579ccbfe..00000000 --- a/diffulex/legacy/layers/attention/ops/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from diffulex.legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd as causal_lm_flash_decoding -from diffulex.legacy.layers.attention.ops.triton_decode_attn_dlm import diffusion_lm_flash_decoding, CHECK_ATTENTION -from diffulex.legacy.layers.attention.ops.chunked_prefill_decoding_unified_kernel import chunked_prefill_paged_decode as diffusion_lm_parallel_flash_decoding -from diffulex.legacy.layers.attention.ops.kv_cache_kernels import ( - store_kvcache_distinct_layout, store_kvcache_unified_layout, load_kvcache, - CHECK_STORING, CHECK_LOADING -) \ No newline at end of file diff --git a/diffulex/legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py b/diffulex/legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py deleted file mode 100755 index aed7e060..00000000 --- a/diffulex/legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py +++ /dev/null @@ -1,375 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# type: ignore -# This file is adapted from the vLLM project: -# https://github.com/vllm-project/vllm/blob/main/vllm/attention/ops/chunked_prefill_paged_decode.py - -# Authors: -# - Burkhard Ringlein -# - Jan van Lunteren -# - Chih-Chieh Yang -# - Thomas Parnell - -import torch - -from vllm import _custom_ops as ops -from vllm.platforms import current_platform -from vllm.platforms.rocm import use_rocm_custom_paged_attention -from vllm.triton_utils import tl, triton - -from diffulex.legacy.layers.attention.ops.prefix_prefill import context_attention_fwd - - -@triton.jit -def cdiv_fn(x, y): - return (x + y - 1) // y - - -@triton.jit -def kernel_paged_attention_2d( - output_ptr, # [num_tokens, num_query_heads, head_size] - query_ptr, # [num_tokens, num_query_heads, head_size] - key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x] - value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size] - block_tables_ptr, # [num_seqs, max_num_blocks_per_seq] - seq_lens_ptr, # [num_seqs] - alibi_slopes_ptr, # [num_query_heads] - scale, # float32 - k_scale, # float32 - v_scale, # float32 - num_query_heads: tl.constexpr, # int - num_queries_per_kv: tl.constexpr, # int - num_queries_per_kv_padded: tl.constexpr, # int - block_table_stride: tl.int64, # int - query_stride_0: tl.int64, # int - query_stride_1: tl.int64, # int, should be equal to head_size - output_stride_0: tl.int64, # int - output_stride_1: tl.int64, # int, should be equal to head_size - BLOCK_SIZE: tl.constexpr, # int - HEAD_SIZE: tl.constexpr, # int - HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2 - USE_ALIBI_SLOPES: tl.constexpr, # bool - SLIDING_WINDOW: tl.constexpr, # int - x: tl.constexpr, # int - stride_k_cache_0: tl.int64, # int - stride_k_cache_1: tl.int64, # int - stride_k_cache_2: tl.int64, # int - stride_k_cache_3: tl.int64, # int - stride_k_cache_4: tl.int64, # int - stride_v_cache_0: tl.int64, # int - stride_v_cache_1: tl.int64, # int - stride_v_cache_2: tl.int64, # int - stride_v_cache_3: tl.int64, # int - filter_by_query_len: tl.constexpr, # bool - query_start_len_ptr, # [num_seqs+1] -): - seq_idx = tl.program_id(0) - kv_head_idx = tl.program_id(1) - - if filter_by_query_len: - cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) - cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + - 1) - cur_batch_query_len = cur_batch_in_all_stop_index \ - - cur_batch_in_all_start_index - if cur_batch_query_len > 1: - return - else: - cur_batch_in_all_start_index = seq_idx - - query_head_idx = kv_head_idx * num_queries_per_kv + tl.arange( - 0, num_queries_per_kv_padded) - - query_offset = (cur_batch_in_all_start_index * query_stride_0 + - query_head_idx[:, None] * query_stride_1) - - head_mask = query_head_idx < (kv_head_idx + 1) * num_queries_per_kv - head_mask = head_mask & (query_head_idx < num_query_heads) - - dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1, - 0).to(tl.int1) - - # Q : (num_queries_per_kv, HEAD_SIZE,) - Q = tl.load( - query_ptr + query_offset + tl.arange(0, HEAD_SIZE_PADDED)[None, :], - mask=dim_mask[None, :] & head_mask[:, None], - other=0.0, - ) - - block_table_offset = seq_idx * block_table_stride - - M = tl.full([num_queries_per_kv_padded], float("-inf"), dtype=tl.float32) - L = tl.full([num_queries_per_kv_padded], 1.0, dtype=tl.float32) - acc = tl.zeros([num_queries_per_kv_padded, HEAD_SIZE_PADDED], - dtype=tl.float32) - - # sequence len for this particular sequence - seq_len = tl.load(seq_lens_ptr + seq_idx) - - # alibi slope for this head - if USE_ALIBI_SLOPES: - alibi_slope = tl.load(alibi_slopes_ptr + query_head_idx, - mask=head_mask, - other=0.0) - - num_blocks = cdiv_fn(seq_len, BLOCK_SIZE) - - # iterate through tiles - for j in range(0, num_blocks): - - physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j) - - offs_n = tl.arange(0, BLOCK_SIZE) - offs_d = tl.arange(0, HEAD_SIZE_PADDED) - - v_offset = (physical_block_idx * stride_v_cache_0 + - kv_head_idx * stride_v_cache_1 + - offs_d[None, :] * stride_v_cache_2 + - offs_n[:, None] * stride_v_cache_3) - - k_offset = (physical_block_idx * stride_k_cache_0 + - kv_head_idx * stride_k_cache_1 + - (offs_d[:, None] // x) * stride_k_cache_2 + - offs_n[None, :] * stride_k_cache_3 + - (offs_d[:, None] % x) * stride_k_cache_4) - - # K : (HEAD_SIZE, BLOCK_SIZE) - K_load = tl.load(key_cache_ptr + k_offset, - mask=dim_mask[:, None], - other=0.0) - - if K_load.dtype.is_fp8(): - K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype) - else: - K = K_load - - # V : (BLOCK_SIZE, HEAD_SIZE) - V_load = tl.load(value_cache_ptr + v_offset, - mask=dim_mask[None, :], - other=0.0) - - if V_load.dtype.is_fp8(): - V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype) - else: - V = V_load - - seq_offset = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - boundary = tl.full([BLOCK_SIZE], seq_len, dtype=tl.int32) - seq_mask = seq_offset[None, :] < boundary - - # S : (num_queries_per_kv, BLOCK_SIZE,) - S = tl.where(head_mask[:, None] & seq_mask, 0.0, - float("-inf")).to(tl.float32) - S += scale * tl.dot(Q, K) - - context_len = seq_len - 1 - - if SLIDING_WINDOW > 0: - S = tl.where((context_len - seq_offset) < SLIDING_WINDOW, S, - -10000) - - if USE_ALIBI_SLOPES: - S += alibi_slope[:, None] * (seq_offset - context_len) - - # compute running maximum - # m_j : (num_queries_per_kv,) - m_j = tl.maximum(M, tl.max(S, axis=1)) - - # P : (num_queries_per_kv, BLOCK_SIZE,) - P = tl.exp(S - m_j[:, None]) - - # l_j : (num_queries_per_kv,) - l_j = tl.sum(P, axis=1) - - # alpha : (num_queries_per_kv, ) - alpha = tl.exp(M - m_j) - - # acc : (num_queries_per_kv, BLOCK_SIZE,) - acc = acc * alpha[:, None] - - # update constants - L = L * alpha + l_j - M = m_j - - # acc : (num_queries_per_kv, BLOCK_SIZE,) - acc += tl.dot(P.to(V.dtype), V) - - # epilogue - acc = acc / L[:, None] - - output_offset = (cur_batch_in_all_start_index * output_stride_0 + - query_head_idx * output_stride_1) - - tl.store( - output_ptr + output_offset[:, None] + - tl.arange(0, HEAD_SIZE_PADDED)[None, :], - acc, - mask=dim_mask[None, :] & head_mask[:, None], - ) - - -def chunked_prefill_paged_decode( - query, - key, - value, - output, - kv_cache_dtype, - key_cache, - value_cache, - block_table, - query_start_loc, - seq_lens, - max_seq_len, - max_query_len, - k_scale, - v_scale, - diffusion_blk_sz=32, - alibi_slopes=None, - sliding_window=None, - sm_scale=None, - mask=None, -): - if sm_scale is None: - sm_scale = 1.0 / (query.shape[1]**0.5) - - use_alibi_slopes = alibi_slopes is not None - - if sliding_window is None or sliding_window <= 0: - sliding_window = 0 - - if max_query_len > 1: - context_attention_fwd( - q=query, - k=key, - v=value, - o=output, - kv_cache_dtype=kv_cache_dtype, - k_cache=key_cache, - v_cache=value_cache, - b_loc=block_table, - b_start_loc=query_start_loc, - b_seq_len=seq_lens, - max_seq_len=max_seq_len, - max_input_len=max_query_len, - k_scale=k_scale, - v_scale=v_scale, - diffusion_blk_sz=diffusion_blk_sz, - alibi_slopes=alibi_slopes, - sliding_window=sliding_window, - sm_scale=sm_scale, - skip_decode=True, - mask=mask - ) - return - - block_size = value_cache.shape[3] - num_seqs = len(seq_lens) - num_query_heads = query.shape[1] - num_kv_heads = key.shape[1] - num_queries_per_kv = query.shape[1] // key.shape[1] - head_size = query.shape[2] - - # Conversion of FP8 Tensor from uint8 storage to - # appropriate torch.dtype for interpretation by Triton - if "fp8" in kv_cache_dtype: - assert key_cache.dtype in [torch.uint8, current_platform.fp8_dtype()] - assert value_cache.dtype in [torch.uint8, current_platform.fp8_dtype()] - - if kv_cache_dtype in ("fp8", "fp8_e4m3"): - target_dtype = current_platform.fp8_dtype() - elif kv_cache_dtype == "fp8_e5m2": - target_dtype = torch.float8_e5m2 - else: - raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype) - - key_cache = key_cache.view(target_dtype) - value_cache = value_cache.view(target_dtype) - - num_queries_per_kv_padded = max(triton.next_power_of_2(num_queries_per_kv), 16) - - use_custom = use_rocm_custom_paged_attention(query.dtype, head_size, - block_size, - num_queries_per_kv, - max_seq_len, sliding_window, - kv_cache_dtype, alibi_slopes) - if use_custom: - _PARTITION_SIZE_ROCM = 256 - max_num_partitions = ((max_seq_len + _PARTITION_SIZE_ROCM - 1) // - _PARTITION_SIZE_ROCM) - assert _PARTITION_SIZE_ROCM % block_size == 0 - total_num_seq = block_table.shape[0] - tmp_output = torch.empty( - size=(total_num_seq, num_query_heads, max_num_partitions, - head_size), - dtype=output.dtype, - device=output.device, - ) - exp_sums = torch.empty( - size=(total_num_seq, num_query_heads, max_num_partitions), - dtype=torch.float32, - device=output.device, - ) - max_logits = torch.empty_like(exp_sums) - - ops.paged_attention_rocm( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale=sm_scale, - block_tables=block_table, - seq_lens=seq_lens, - query_start_loc=query_start_loc, - block_size=block_size, - max_seq_len=max_seq_len, - alibi_slopes=alibi_slopes, - kv_cache_dtype=kv_cache_dtype, - k_scale=k_scale, - v_scale=v_scale, - ) - else: - kernel_paged_attention_2d[( - num_seqs, - num_kv_heads, - )]( - output_ptr=output, - query_ptr=query, - key_cache_ptr=key_cache, - value_cache_ptr=value_cache, - block_tables_ptr=block_table, - seq_lens_ptr=seq_lens, - alibi_slopes_ptr=alibi_slopes, - scale=sm_scale, - k_scale=k_scale, - v_scale=v_scale, - num_query_heads=num_query_heads, - num_queries_per_kv=num_queries_per_kv, - num_queries_per_kv_padded=num_queries_per_kv_padded, - block_table_stride=block_table.stride(0), - query_stride_0=query.stride(0), - query_stride_1=query.stride(1), - output_stride_0=output.stride(0), - output_stride_1=output.stride(1), - BLOCK_SIZE=block_size, - HEAD_SIZE=head_size, - HEAD_SIZE_PADDED=triton.next_power_of_2(head_size), - USE_ALIBI_SLOPES=use_alibi_slopes, - SLIDING_WINDOW=sliding_window, - x=key_cache.shape[4], - stride_k_cache_0=key_cache.stride(0), - stride_k_cache_1=key_cache.stride(1), - stride_k_cache_2=key_cache.stride(2), - stride_k_cache_3=key_cache.stride(3), - stride_k_cache_4=key_cache.stride(4), - stride_v_cache_0=value_cache.stride(0), - stride_v_cache_1=value_cache.stride(1), - stride_v_cache_2=value_cache.stride(2), - stride_v_cache_3=value_cache.stride(3), - filter_by_query_len=True, - query_start_len_ptr=query_start_loc, - ) \ No newline at end of file diff --git a/diffulex/legacy/layers/attention/ops/prefix_prefill.py b/diffulex/legacy/layers/attention/ops/prefix_prefill.py deleted file mode 100755 index 03cf31a8..00000000 --- a/diffulex/legacy/layers/attention/ops/prefix_prefill.py +++ /dev/null @@ -1,1090 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# type: ignore -# This file is adapted from the vLLM project: -# https://github.com/vllm-project/vllm/blob/main/vllm/attention/ops/prefix_prefill.py -# The kernels in this file are originally adapted from LightLLM's context_attention_fwd: -# https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py - -import torch -import triton - -import triton.language as tl - -from vllm.platforms import current_platform - -# Static kernels parameters -BASE_BLOCK = 128 if current_platform.has_device_capability(80) else 64 -NUM_WARPS = 4 if current_platform.is_rocm() else 8 - -# To check compatibility -IS_TURING = current_platform.get_device_capability() == (7, 5) - - -@triton.jit -def _fwd_kernel_d2f(Q, K, V, Mask, - K_cache, V_cache, - B_Loc, - sm_scale, k_scale, v_scale, - B_Start_Loc, - B_Seqlen, - x: tl.constexpr, - Out, - stride_b_loc_b, stride_b_loc_s, - stride_qbs, stride_qh, stride_qd, - stride_kbs, stride_kh, stride_kd, - stride_vbs, stride_vh, stride_vd, - stride_obs, stride_oh, stride_od, - stride_k_cache_bs, stride_k_cache_h, stride_k_cache_d, stride_k_cache_bl: tl.constexpr, stride_k_cache_x, - stride_v_cache_bs, stride_v_cache_h, stride_v_cache_d, stride_v_cache_bl, - stride_mask_m, stride_mask_n, - num_queries_per_kv: tl.constexpr, - IN_PRECISION: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_DMODEL_PADDED: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BLOCK_N: tl.constexpr, - SLIDING_WINDOW: tl.constexpr, - num_unroll_cache: tl.constexpr, - num_unroll_request: tl.constexpr, - SKIP_DECODE: tl.constexpr, - DIFFUSION_BLK_SZ: tl.constexpr, - MAX_Q_LEN: tl.constexpr = 0, - MAX_CTX_LEN: tl.constexpr = 0): - cur_batch = tl.program_id(0) - cur_head = tl.program_id(1) - start_m = tl.program_id(2) - - tl.device_print("=" * 60, cur_batch) - tl.device_print("Program Start", cur_batch) - tl.device_print("=" * 60, cur_batch) - tl.device_print("cur_batch", cur_batch) - tl.device_print("cur_head", cur_head) - tl.device_print("start_m", start_m) - - cur_kv_head = cur_head // num_queries_per_kv - - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) - cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1) - cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index - cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len - - if SKIP_DECODE and cur_batch_query_len == 1: - return - - # start position inside of the query - # generally, N goes over kv, while M goes over query_len - block_start_loc = BLOCK_M * start_m - - # initialize offsets - # [BLOCK_SIZE]; starts at 0 - offs_bs_n = tl.arange(0, BLOCK_SIZE) - # [N]; starts at 0 - offs_n = tl.arange(0, BLOCK_N) - # [D]; starts at 0 - offs_d = tl.arange(0, BLOCK_DMODEL_PADDED) - # [M]; starts at current position in query - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - # [M,D] - offs_q = (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + cur_head * stride_qh + offs_d[None, :] * stride_qd - dim_mask = tl.where(tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(tl.int1) # [D] - q = tl.load(Q + offs_q, mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len), other=0.0) # [M,D] - - # initialize pointer to m and l - m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) - l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) # [M,D] - - # compute query against context (no causal mask here) - for start_n in tl.range(0, cur_batch_ctx_len, BLOCK_SIZE, loop_unroll_factor=num_unroll_cache): - start_n = tl.multiple_of(start_n, BLOCK_SIZE) - # ---- compute qk ---- - bn = tl.load(B_Loc + cur_batch * stride_b_loc_b + (start_n // BLOCK_SIZE) * stride_b_loc_s) - tl.device_print("[CTX] start_n=", start_n) - tl.device_print("[CTX] bn=", bn) - tl.device_print("[CTX] ctx_len=", cur_batch_ctx_len) - # [D,BLOCK_SIZE] - offs_k = (bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h + - (offs_d[:, None] // x) * stride_k_cache_d + - ((start_n + offs_bs_n[None, :]) % BLOCK_SIZE) * stride_k_cache_bl + - (offs_d[:, None] % x) * stride_k_cache_x) - - # [BLOCK_SIZE,D] - offs_v = (bn[:, None] * stride_v_cache_bs + cur_kv_head * stride_v_cache_h + - offs_d[None, :] * stride_v_cache_d + offs_bs_n[:, None] * stride_v_cache_bl) - - if start_n + BLOCK_SIZE > cur_batch_ctx_len or BLOCK_DMODEL != BLOCK_DMODEL_PADDED: - k_load = tl.load(K_cache + offs_k, - mask=dim_mask[:, None] & ((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len), - other=0.0) # [D,N] - else: - k_load = tl.load(K_cache + offs_k) - - if k_load.dtype.is_fp8(): - k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype) - else: - k = k_load - - qk = tl.zeros([BLOCK_M, BLOCK_SIZE], dtype=tl.float32) # [M,N] - qk += tl.dot(q, k, input_precision=IN_PRECISION) - qk_mask = ((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len) & (offs_m[:, None] < cur_batch_query_len) - qk = tl.where(qk_mask, qk, float("-inf")) - - qk *= sm_scale - if SLIDING_WINDOW > 0: - # (cur_batch_ctx_len + offs_m[:, None]) are the positions of - # Q entries in sequence - # (start_n + offs_bs_n[None, :]) are the positions of - # KV entries in sequence - # So the condition makes sure each entry in Q only attends - # to KV entries not more than SLIDING_WINDOW away. - # - # We can't use -inf here, because the - # sliding window may lead to the entire row being masked. - # This then makes m_ij contain -inf, which causes NaNs in - # exp(). - qk = tl.where((cur_batch_ctx_len + offs_m[:, None]) - (start_n + offs_bs_n[None, :]) < SLIDING_WINDOW, qk, -10000) - - # compute running maximum - m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp(qk - m_ij[:, None]) - l_ij = tl.sum(p, axis=1) - alpha = tl.exp(m_i - m_ij) - acc = acc * alpha[:, None] - - # update acc - if start_n + BLOCK_SIZE > cur_batch_ctx_len or BLOCK_DMODEL != BLOCK_DMODEL_PADDED: - v_load = tl.load(V_cache + offs_v, - mask=dim_mask[None, :] & ((start_n + offs_bs_n[:, None]) < cur_batch_ctx_len), - other=0.0) # [N,D] - else: - v_load = tl.load(V_cache + offs_v) - - if v_load.dtype.is_fp8(): - v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype) - else: - v = v_load - p = p.to(v.dtype) - - acc += tl.dot(p, v, input_precision=IN_PRECISION) - # # update m_i and l_i - l_i = l_i * alpha + l_ij - m_i = m_ij - - offs_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd - offs_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd - k_ptrs = K + offs_k - v_ptrs = V + offs_v - - # block_mask is 0 when we're already past the current query length - block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0) - - # compute query against itself (with custom dense mask) - for start_n in tl.range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N, loop_unroll_factor=num_unroll_request): - start_n = tl.multiple_of(start_n, BLOCK_N) - tl.device_print("[SELF] start_n=", start_n) - tl.device_print("[SELF] q_len=", cur_batch_query_len) - tl.device_print("[SELF] block_mask=", block_mask) - # ---- compute qk ---- - k = tl.load(k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, - mask=dim_mask[:, None] & ((start_n + offs_n[None, :]) < cur_batch_query_len), - other=0.0) - - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk += tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) - qk *= sm_scale - - # apply causal mask - # qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, - # float("-inf")) - - # TODO apply block-wise causal mask - offs_mask = offs_m[:, None] * stride_mask_m + (start_n + offs_n[None, :]) * stride_mask_n - mask_ptrs = Mask + offs_mask - m_mask = (offs_m[:, None] < cur_batch_query_len) & ((start_n + offs_n[None, :]) < cur_batch_query_len) - mask = tl.load(mask_ptrs, mask=m_mask, other=False) - qk = tl.where(mask, qk, float("-inf")) - valid_cnt = tl.sum(mask, axis=1) - tl.device_print("[SELF] valid per-row row0=", valid_cnt) - if SLIDING_WINDOW > 0: - qk = tl.where(offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW, qk, -10000) - - # compute running maximum - m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp(qk - m_ij[:, None]) - l_ij = tl.sum(p, axis=1) - alpha = tl.exp(m_i - m_ij) - acc = acc * alpha[:, None] - - # update acc - v = tl.load(v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, - mask=dim_mask[None, :] & ((start_n + offs_n[:, None]) < cur_batch_query_len), - other=0.0) - p = p.to(v.dtype) - - acc += tl.dot(p, v, input_precision=IN_PRECISION) - # update m_i and l_i - l_i = l_i * alpha + l_ij - m_i = m_ij - - acc = acc / l_i[:, None] - - # initialize pointers to output - off_o = (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + cur_head * stride_oh + offs_d[None, :] * stride_od - out_ptrs = Out + off_o - tl.store(out_ptrs, acc, mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len)) - tl.device_print("\n\n", cur_batch) - return - - -@triton.jit -def _fwd_kernel(Q, K, V, - K_cache, V_cache, - B_Loc, - sm_scale, k_scale, v_scale, - B_Start_Loc, - B_Seqlen, - x: tl.constexpr, - Out, - stride_b_loc_b, stride_b_loc_s, - stride_qbs, stride_qh, stride_qd, - stride_kbs, stride_kh, stride_kd, - stride_vbs, stride_vh, stride_vd, - stride_obs, stride_oh, stride_od, - stride_k_cache_bs, stride_k_cache_h, stride_k_cache_d, stride_k_cache_bl: tl.constexpr, stride_k_cache_x, - stride_v_cache_bs, stride_v_cache_h, stride_v_cache_d, stride_v_cache_bl, - num_queries_per_kv: tl.constexpr, - IN_PRECISION: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_DMODEL_PADDED: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BLOCK_N: tl.constexpr, - SLIDING_WINDOW: tl.constexpr, - num_unroll_cache: tl.constexpr, - num_unroll_request: tl.constexpr, - SKIP_DECODE: tl.constexpr, - MAX_Q_LEN: tl.constexpr = 0, - MAX_CTX_LEN: tl.constexpr = 0): - cur_batch = tl.program_id(0) - cur_head = tl.program_id(1) - start_m = tl.program_id(2) - - cur_kv_head = cur_head // num_queries_per_kv - - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) - cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1) - cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index - cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len - - if SKIP_DECODE and cur_batch_query_len == 1: - return - - # start position inside of the query - # generally, N goes over kv, while M goes over query_len - block_start_loc = BLOCK_M * start_m - - # initialize offsets - # [BLOCK_SIZE]; starts at 0 - offs_bs_n = tl.arange(0, BLOCK_SIZE) - # [N]; starts at 0 - offs_n = tl.arange(0, BLOCK_N) - # [D]; starts at 0 - offs_d = tl.arange(0, BLOCK_DMODEL_PADDED) - # [M]; starts at current position in query - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - # [M,D] - off_q = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + - cur_head * stride_qh + offs_d[None, :] * stride_qd) - - dim_mask = tl.where(tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(tl.int1) # [D] - - q = tl.load(Q + off_q, - mask=dim_mask[None, :] & - (offs_m[:, None] < cur_batch_query_len), - other=0.0) # [M,D] - - # initialize pointer to m and l - m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) - l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) # [M,D] - - # compute query against context (no causal mask here) - for start_n in tl.range(0, cur_batch_ctx_len, BLOCK_SIZE, \ - loop_unroll_factor=num_unroll_cache): - start_n = tl.multiple_of(start_n, BLOCK_SIZE) - # -- compute qk ---- - bn = tl.load(B_Loc + cur_batch * stride_b_loc_b + - (start_n // BLOCK_SIZE) * stride_b_loc_s) - # [D,BLOCK_SIZE] - off_k = (bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h + - (offs_d[:, None] // x) * stride_k_cache_d + - ((start_n + offs_bs_n[None, :]) % BLOCK_SIZE) * stride_k_cache_bl + - (offs_d[:, None] % x) * stride_k_cache_x) - - # [BLOCK_SIZE,D] - off_v = (bn[:, None] * stride_v_cache_bs + - cur_kv_head * stride_v_cache_h + - offs_d[None, :] * stride_v_cache_d + - offs_bs_n[:, None] * stride_v_cache_bl) - - if start_n + BLOCK_SIZE > cur_batch_ctx_len or \ - BLOCK_DMODEL != BLOCK_DMODEL_PADDED: - k_load = tl.load( - K_cache + off_k, - mask=dim_mask[:, None] & - ((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len), - other=0.0) # [D,N] - else: - k_load = tl.load(K_cache + off_k) - - if k_load.dtype.is_fp8(): - k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype) - else: - k = k_load - - qk = tl.zeros([BLOCK_M, BLOCK_SIZE], dtype=tl.float32) # [M,N] - qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) - qk = tl.where((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len, qk, - float("-inf")) - qk *= sm_scale - if SLIDING_WINDOW > 0: - # (cur_batch_ctx_len + offs_m[:, None]) are the positions of - # Q entries in sequence - # (start_n + offs_bs_n[None, :]) are the positions of - # KV entries in sequence - # So the condition makes sure each entry in Q only attends - # to KV entries not more than SLIDING_WINDOW away. - # - # We can't use -inf here, because the - # sliding window may lead to the entire row being masked. - # This then makes m_ij contain -inf, which causes NaNs in - # exp(). - qk = tl.where((cur_batch_ctx_len + offs_m[:, None]) - - (start_n + offs_bs_n[None, :]) < SLIDING_WINDOW, qk, - -10000) - - # compute running maximum - m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp(qk - m_ij[:, None]) - l_ij = tl.sum(p, axis=1) - alpha = tl.exp(m_i - m_ij) - acc = acc * alpha[:, None] - - # update acc - if start_n + BLOCK_SIZE > cur_batch_ctx_len or \ - BLOCK_DMODEL != BLOCK_DMODEL_PADDED: - v_load = tl.load( - V_cache + off_v, - mask=dim_mask[None, :] & - ((start_n + offs_bs_n[:, None]) < cur_batch_ctx_len), - other=0.0) # [N,D] - else: - v_load = tl.load(V_cache + off_v) - - if v_load.dtype.is_fp8(): - v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype) - else: - v = v_load - p = p.to(v.dtype) - - acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION) - # # update m_i and l_i - l_i = l_i * alpha + l_ij - m_i = m_ij - - off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd - off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd - k_ptrs = K + off_k - v_ptrs = V + off_v - - # block_mask is 0 when we're already past the current query length - block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0) - - # compute query against itself (with causal mask) - for start_n in tl.range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N, loop_unroll_factor=num_unroll_request): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - k = tl.load(k_ptrs + - (cur_batch_in_all_start_index + start_n) * stride_kbs, - mask=dim_mask[:, None] & - ((start_n + offs_n[None, :]) < cur_batch_query_len), - other=0.0) - - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) - qk *= sm_scale - # apply causal mask - qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) - if SLIDING_WINDOW > 0: - qk = tl.where( - offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW, - qk, -10000) - - # compute running maximum - m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp(qk - m_ij[:, None]) - l_ij = tl.sum(p, axis=1) - alpha = tl.exp(m_i - m_ij) - acc = acc * alpha[:, None] - - # update acc - v = tl.load(v_ptrs + - (cur_batch_in_all_start_index + start_n) * stride_vbs, - mask=dim_mask[None, :] & - ((start_n + offs_n[:, None]) < cur_batch_query_len), - other=0.0) - p = p.to(v.dtype) - - acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION) - # update m_i and l_i - l_i = l_i * alpha + l_ij - m_i = m_ij - - acc = acc / l_i[:, None] - - # initialize pointers to output - off_o = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + - cur_head * stride_oh + offs_d[None, :] * stride_od) - out_ptrs = Out + off_o - tl.store(out_ptrs, - acc, - mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len)) - return - - -@triton.jit -def _fwd_kernel_flash_attn_v2( - Q, - K, - V, - K_cache, - V_cache, - B_Loc, - sm_scale, - B_Start_Loc, - B_Seqlen, - B_Ctxlen, - block_size, - x, - Out, - stride_b_loc_b, - stride_b_loc_s, - stride_qbs, - stride_qh, - stride_qd, - stride_kbs, - stride_kh, - stride_kd, - stride_vbs, - stride_vh, - stride_vd, - stride_obs, - stride_oh, - stride_od, - stride_k_cache_bs, - stride_k_cache_h, - stride_k_cache_d, - stride_k_cache_bl, - stride_k_cache_x, - stride_v_cache_bs, - stride_v_cache_h, - stride_v_cache_d, - stride_v_cache_bl, - num_queries_per_kv: int, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, -): - cur_batch = tl.program_id(0) - cur_head = tl.program_id(1) - start_m = tl.program_id(2) - - cur_kv_head = cur_head // num_queries_per_kv - - cur_batch_ctx_len = tl.load(B_Ctxlen + cur_batch) - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) - - block_start_loc = BLOCK_M * start_m - - # initialize offsets - offs_n = tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_DMODEL) - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - off_q = (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + cur_head * stride_qh + offs_d[None, :] * stride_qd - - q = tl.load(Q + off_q, mask=offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len, other=0.0) - - # # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) - - for start_n in range(0, cur_batch_ctx_len, BLOCK_N): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - bn = tl.load(B_Loc + cur_batch * stride_b_loc_b + - ((start_n + offs_n) // block_size) * stride_b_loc_s, - mask=(start_n + offs_n) < cur_batch_ctx_len, - other=0) - off_k = ( - bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h + - (offs_d[:, None] // x) * stride_k_cache_d + - ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl + - (offs_d[:, None] % x) * stride_k_cache_x) - off_v = (bn[:, None] * stride_v_cache_bs + - cur_kv_head * stride_v_cache_h + - offs_d[None, :] * stride_v_cache_d + - (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl) - k = tl.load(K_cache + off_k, - mask=(start_n + offs_n[None, :]) < cur_batch_ctx_len, - other=0.0) - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk += tl.dot(q, k) - qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk, - float("-inf")) - qk *= sm_scale - - # -- compute m_ij, p, l_ij - m_ij = tl.max(qk, 1) - m_i_new = tl.maximum(m_i, m_ij) - p = tl.math.exp(qk - m_i_new[:, None]) - l_ij = tl.sum(p, 1) - # -- update m_i and l_i - - alpha = tl.math.exp(m_i - m_i_new) - l_i_new = alpha * l_i + l_ij - # -- update output accumulator -- - # scale p - # scale acc - acc_scale = alpha - # acc_scale = l_i / l_i_new * alpha - acc = acc * acc_scale[:, None] - # update acc - v = tl.load(V_cache + off_v, - mask=(start_n + offs_n[:, None]) < cur_batch_ctx_len, - other=0.0) - - p = p.to(v.dtype) - acc += tl.dot(p, v) - # update m_i and l_i - l_i = l_i_new - m_i = m_i_new - - off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + - offs_d[:, None] * stride_kd) - off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + - offs_d[None, :] * stride_vd) - k_ptrs = K + off_k - v_ptrs = V + off_v - - block_mask = tl.where( - block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0) - - for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - k = tl.load(k_ptrs + - (cur_batch_in_all_start_index + start_n) * stride_kbs, - mask=(start_n + offs_n[None, :]) - < cur_batch_seq_len - cur_batch_ctx_len, - other=0.0) - - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk += tl.dot(q, k) - qk *= sm_scale - qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, - float("-inf")) - - # -- compute m_ij, p, l_ij - m_ij = tl.max(qk, 1) - m_i_new = tl.maximum(m_i, m_ij) - p = tl.math.exp(qk - m_i_new[:, None]) - l_ij = tl.sum(p, 1) - # -- update m_i and l_i - - alpha = tl.math.exp(m_i - m_i_new) - l_i_new = alpha * l_i + l_ij - # -- update output accumulator -- - # scale p - # scale acc - acc_scale = alpha - # acc_scale = l_i / l_i_new * alpha - acc = acc * acc_scale[:, None] - # update acc - v = tl.load(v_ptrs + - (cur_batch_in_all_start_index + start_n) * stride_vbs, - mask=(start_n + offs_n[:, None]) - < cur_batch_seq_len - cur_batch_ctx_len, - other=0.0) - - p = p.to(v.dtype) - acc += tl.dot(p, v) - # update m_i and l_i - l_i = l_i_new - m_i = m_i_new - - # acc /= l_i[:, None] - # initialize pointers to output - off_o = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + - cur_head * stride_oh + offs_d[None, :] * stride_od) - out_ptrs = Out + off_o - tl.store(out_ptrs, - acc, - mask=offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len) - return - - -@triton.jit -def _fwd_kernel_alibi( - Q, - K, - V, - K_cache, - V_cache, - B_Loc, - sm_scale, - k_scale, - v_scale, - B_Start_Loc, - B_Seqlen, - Alibi_slopes, - block_size, - x, - Out, - stride_b_loc_b, - stride_b_loc_s, - stride_qbs, - stride_qh, - stride_qd, - stride_kbs, - stride_kh, - stride_kd, - stride_vbs, - stride_vh, - stride_vd, - stride_obs, - stride_oh, - stride_od, - stride_k_cache_bs, - stride_k_cache_h, - stride_k_cache_d, - stride_k_cache_bl, - stride_k_cache_x, - stride_v_cache_bs, - stride_v_cache_h, - stride_v_cache_d, - stride_v_cache_bl, - num_queries_per_kv: int, - IN_PRECISION: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, # head size - BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2 - BLOCK_N: tl.constexpr, - SKIP_DECODE: tl.constexpr, -): - # attn_bias[] - cur_batch = tl.program_id(0) - cur_head = tl.program_id(1) - start_m = tl.program_id(2) - - cur_kv_head = cur_head // num_queries_per_kv - - # cur_batch_seq_len: the length of prompts - # cur_batch_ctx_len: the length of prefix - # cur_batch_in_all_start_index: the start id of the dim=0 - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) - cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1) - cur_batch_query_len = (cur_batch_in_all_stop_index - - cur_batch_in_all_start_index) - cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len - - if SKIP_DECODE and cur_batch_query_len == 1: - return - - block_start_loc = BLOCK_M * start_m - - # initialize offsets - offs_n = tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_DMODEL_PADDED) - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - off_q = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + - cur_head * stride_qh + offs_d[None, :] * stride_qd) - - dim_mask = tl.where( - tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to(tl.int1) - - q = tl.load(Q + off_q, - mask=dim_mask[None, :] & - (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len), - other=0.0) - - # # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) - - alibi_slope = tl.load(Alibi_slopes + cur_head) - alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len - alibi_start_k = 0 - for start_n in range(0, cur_batch_ctx_len, BLOCK_N): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - bn = tl.load(B_Loc + cur_batch * stride_b_loc_b + - ((start_n + offs_n) // block_size) * stride_b_loc_s, - mask=(start_n + offs_n) < cur_batch_ctx_len, - other=0) - off_k = ( - bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h + - (offs_d[:, None] // x) * stride_k_cache_d + - ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl + - (offs_d[:, None] % x) * stride_k_cache_x) - off_v = (bn[:, None] * stride_v_cache_bs + - cur_kv_head * stride_v_cache_h + - offs_d[None, :] * stride_v_cache_d + - (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl) - k_load = tl.load(K_cache + off_k, - mask=dim_mask[:, None] & - ((start_n + offs_n[None, :]) < cur_batch_ctx_len), - other=0.0) # [D,N] - - if k_load.dtype.is_fp8(): - k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype) - else: - k = k_load - - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) - qk = tl.where((start_n + offs_n[None, :]) < cur_batch_ctx_len, qk, - float("-inf")) - qk *= sm_scale - - # load alibi - alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k - - alibi_start_q[:, None]) * alibi_slope - alibi = tl.where( - (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len), alibi, - float("-inf")) - qk += alibi - alibi_start_k += BLOCK_N - - # -- compute m_ij, p, l_ij - m_ij = tl.max(qk, 1) - m_i_new = tl.maximum(m_i, m_ij) - p = tl.math.exp(qk - m_i_new[:, None]) - l_ij = tl.sum(p, 1) - # -- update m_i and l_i - - alpha = tl.math.exp(m_i - m_i_new) - l_i_new = alpha * l_i + l_ij - # -- update output accumulator -- - # scale p - # scale acc - acc_scale = alpha - # acc_scale = l_i / l_i_new * alpha - acc = acc * acc_scale[:, None] - # update acc - v_load = tl.load(V_cache + off_v, - mask=dim_mask[None, :] & - ((start_n + offs_n[:, None]) < cur_batch_ctx_len), - other=0.0) - if v_load.dtype.is_fp8(): - v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype) - else: - v = v_load - p = p.to(v.dtype) - - acc = tl.dot(p, v, acc=acc, input_precision='ieee') - # update m_i and l_i - l_i = l_i_new - m_i = m_i_new - - off_k = (offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + - offs_d[:, None] * stride_kd) - off_v = (offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + - offs_d[None, :] * stride_vd) - k_ptrs = K + off_k - v_ptrs = V + off_v - - block_mask = tl.where( - block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0) - - # init alibi - alibi_slope = tl.load(Alibi_slopes + cur_head) - alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len - alibi_start_k = cur_batch_ctx_len - # # init debugger - # offset_db_q = tl.arange(0, BLOCK_M) + block_start_loc - # offset_db_k = tl.arange(0, BLOCK_N) - # calc q[BLOCK_M, BLOCK_MODEL] mul k[prefix_len: , BLOCK_DMODEL] - for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - k = tl.load( - k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, - mask=dim_mask[:, None] & ((start_n + offs_n[None, :]) - < cur_batch_seq_len - cur_batch_ctx_len), - other=0.0) - - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk = tl.dot(q, k, acc=qk, input_precision='ieee') - qk *= sm_scale - qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, - float("-inf")) - - # load alibi - alibi = (tl.arange(0, BLOCK_N)[None, :] + alibi_start_k - - alibi_start_q[:, None]) * alibi_slope - alibi = tl.where( - (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len), alibi, - float("-inf")) - qk += alibi - alibi_start_k += BLOCK_N - - # -- compute m_ij, p, l_ij - m_ij = tl.max(qk, 1) - m_i_new = tl.maximum(m_i, m_ij) - p = tl.math.exp(qk - m_i_new[:, None]) - l_ij = tl.sum(p, 1) - # -- update m_i and l_i - - alpha = tl.math.exp(m_i - m_i_new) - l_i_new = alpha * l_i + l_ij - # -- update output accumulator -- - # scale p - # scale acc - acc_scale = alpha - # acc_scale = l_i / l_i_new * alpha - acc = acc * acc_scale[:, None] - # update acc - v = tl.load( - v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, - mask=dim_mask[None, :] & ((start_n + offs_n[:, None]) - < cur_batch_seq_len - cur_batch_ctx_len), - other=0.0) - p = p.to(v.dtype) - - acc = tl.dot(p, v, acc=acc, input_precision='ieee') - # update m_i and l_i - l_i = l_i_new - m_i = m_i_new - - acc = acc / l_i[:, None] - - # initialize pointers to output - off_o = ((cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + - cur_head * stride_oh + offs_d[None, :] * stride_od) - out_ptrs = Out + off_o - tl.store(out_ptrs, - acc, - mask=dim_mask[None, :] & - (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len)) - return - - -@torch.inference_mode() -def context_attention_fwd(q, - k, - v, - o, - kv_cache_dtype: str, - k_cache, - v_cache, - b_loc, - b_start_loc, - b_seq_len, - max_seq_len, - max_input_len, - k_scale: torch.Tensor, - v_scale: torch.Tensor, - diffusion_blk_sz=None, - alibi_slopes=None, - sliding_window=None, - sm_scale=None, - skip_decode=False, - mask: torch.Tensor=None): - - q_dtype_is_f32 = q.dtype is torch.float32 - - # Turing does have tensor core for float32 multiplication - # use ieee as fallback for triton kernels work. There is also - # warning on vllm/config.py to inform users this fallback - # implementation - IN_PRECISION = 'ieee' if IS_TURING and q_dtype_is_f32 else None - - # Conversion of FP8 Tensor from uint8 storage to - # appropriate torch.dtype for interpretation by Triton - if "fp8" in kv_cache_dtype: - assert k_cache.dtype in [torch.uint8, current_platform.fp8_dtype()] - assert v_cache.dtype in [torch.uint8, current_platform.fp8_dtype()] - - if kv_cache_dtype in ("fp8", "fp8_e4m3"): - target_dtype = current_platform.fp8_dtype() - elif kv_cache_dtype == "fp8_e5m2": - target_dtype = torch.float8_e5m2 - else: - raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype) - - k_cache = k_cache.view(target_dtype) - v_cache = v_cache.view(target_dtype) - - if (k_cache.dtype == torch.uint8 - or v_cache.dtype == torch.uint8 and kv_cache_dtype == "auto"): - raise ValueError("kv_cache_dtype='auto' unsupported for\ - FP8 KV Cache prefill kernel") - - # shape constraints - Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] - assert Lq == Lk and Lk == Lv - # round up Lk to a power of 2 - this is required for Triton block size - Lk_padded = triton.next_power_of_2(Lk) - - if sm_scale is None: - sm_scale = 1.0 / (Lq**0.5) - batch, head = b_seq_len.shape[0], q.shape[1] - num_queries_per_kv = q.shape[1] // k.shape[1] - - assert batch + 1 == len(b_start_loc) - - # 0 means "disable" - if sliding_window is None or sliding_window <= 0: - sliding_window = 0 - - if alibi_slopes is not None: - # need to reduce num. blocks when using fp32 - # due to increased use of GPU shared memory - # if q.dtype is torch.float32: - BLOCK = BASE_BLOCK // 2 if q_dtype_is_f32 else BASE_BLOCK - # batch, head, - grid = (batch, head, triton.cdiv(max_input_len, BLOCK)) - _fwd_kernel_alibi[grid]( - q, - k, - v, - k_cache, - v_cache, - b_loc, - sm_scale, - k_scale, - v_scale, - b_start_loc, - b_seq_len, - alibi_slopes, - v_cache.shape[3], - k_cache.shape[4], - o, - b_loc.stride(0), - b_loc.stride(1), - q.stride(0), - q.stride(1), - q.stride(2), - k.stride(0), - k.stride(1), - k.stride(2), - v.stride(0), - v.stride(1), - v.stride(2), - o.stride(0), - o.stride(1), - o.stride(2), - k_cache.stride(0), - k_cache.stride(1), - k_cache.stride(2), - k_cache.stride(3), - k_cache.stride(4), #[num_blocks, num_kv_heads, head_size/x, block_size, x] - v_cache.stride(0), - v_cache.stride(1), - v_cache.stride(2), - v_cache.stride(3), #[num_blocks, num_kv_heads, head_size, block_size] - num_queries_per_kv=num_queries_per_kv, - IN_PRECISION=IN_PRECISION, - BLOCK_M=BLOCK, - BLOCK_DMODEL=Lk, - BLOCK_DMODEL_PADDED=Lk_padded, - BLOCK_N=BLOCK, - SKIP_DECODE=skip_decode, - num_warps=NUM_WARPS, - num_stages=1, - ) - return - - max_seq_len = 0 if max_seq_len is None else max_seq_len - extra_kargs = {} - if current_platform.is_rocm(): - extra_kargs = {"kpack": 2, "waves_per_eu": 2} - - if diffusion_blk_sz is None: - grid = lambda META: (batch, head, triton.cdiv(max_input_len, META["BLOCK_M"])) - _fwd_kernel[grid]( - q, k, v, - k_cache, v_cache, - b_loc, - sm_scale, k_scale, v_scale, - b_start_loc, b_seq_len, - k_cache.shape[4], - o, - b_loc.stride(0), b_loc.stride(1), - q.stride(0), q.stride(1), q.stride(2), - k.stride(0), k.stride(1), k.stride(2), - v.stride(0), v.stride(1), v.stride(2), - o.stride(0), o.stride(1), o.stride(2), - #[num_blocks, num_kv_heads, head_size/x, block_size, x] - k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.stride(3), k_cache.stride(4), - #[num_blocks, num_kv_heads, head_size, block_size] - v_cache.stride(0), v_cache.stride(1), v_cache.stride(2), v_cache.stride(3), - BLOCK_SIZE=v_cache.shape[3], - num_queries_per_kv=num_queries_per_kv, - IN_PRECISION=IN_PRECISION, - BLOCK_DMODEL=Lk, - BLOCK_DMODEL_PADDED=Lk_padded, - SLIDING_WINDOW=sliding_window, - SKIP_DECODE=skip_decode, - BLOCK_M=128, - BLOCK_N=64, - num_unroll_cache=4, - num_unroll_request=1, - num_warps=4, - num_stages=1, - **extra_kargs) - else: - # FIXME: computation not correct - BLOCK_M = BLOCK_N = diffusion_blk_sz * 2 - GRID = (batch, head, triton.cdiv(max_input_len, BLOCK_M)) - _fwd_kernel_d2f[GRID]( - q, k, v, mask, - k_cache, v_cache, - b_loc, - sm_scale, k_scale, v_scale, - b_start_loc, b_seq_len, - k_cache.shape[-1], - o, - *b_loc.stride(), - *q.stride(), - *k.stride(), - *v.stride(), - *o.stride(), - *k_cache.stride(), #[num_blocks, num_kv_heads, head_size/x, block_size, x] - *v_cache.stride(), #[num_blocks, num_kv_heads, head_size, block_size] - *mask.stride(), - BLOCK_SIZE=v_cache.shape[-1], - num_queries_per_kv=num_queries_per_kv, - IN_PRECISION=IN_PRECISION, - BLOCK_DMODEL=Lk, - BLOCK_DMODEL_PADDED=Lk_padded, - SLIDING_WINDOW=sliding_window, - SKIP_DECODE=skip_decode, - BLOCK_M=BLOCK_M, - BLOCK_N=BLOCK_N, - DIFFUSION_BLK_SZ=diffusion_blk_sz, - num_unroll_cache=4, - num_unroll_request=1, - num_warps=4, - num_stages=1, - **extra_kargs) - return \ No newline at end of file diff --git a/diffulex/legacy/layers/attention/ops/tilus_decode_attn_dlm.py b/diffulex/legacy/layers/attention/ops/tilus_decode_attn_dlm.py deleted file mode 100755 index fc7bc03b..00000000 --- a/diffulex/legacy/layers/attention/ops/tilus_decode_attn_dlm.py +++ /dev/null @@ -1,161 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: D2F -# type: ignore - -# Organization: SJTU DENG Lab -# Author: Drew Jin (JIN. Yijie, @drewjin) -# Date: 2025-08-15 -# Email: drewjin0827@gmail.com -# All rights reserved. - -import tilus -import torch - -import numpy as np - -from hidet.ir import DataType -from tilus.utils import cdiv -from tilus import boolean, f32, int32, int64, void_p - - -tilus.option.cache_dir("./cache") - - -class TilusDecodeAttnForDifusionLM(tilus.Script): - """ - Fusing kvcache loading, attention against kvcache, self-attention, - and self-attention custom mask applying all together - """ - def __init__(self, dtype: DataType, num_heads: int, num_kv_heads: int, - head_dim: int, num_warps: int, diffusion_block_size: int, - page_size: int = 256, x: int = 8): - super().__init__() - self.dtype: DataType = dtype - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads - self.head_dim = head_dim - self.x = x - self.head_dim_x = head_dim // x - self.num_warps = num_warps - self.block_q = diffusion_block_size * 2 - self.block_kv = diffusion_block_size * 2 - self.block_kvc = self.page_size = page_size - self.score_scale = float(1.0 / np.sqrt(head_dim)) - self.group_size = num_heads // num_kv_heads - - # For attn against kvcache - self.qkc_config = self.cuda.resolve_dot_config( - dtype, - f32, - m=self.block_q, - n=self.block_kv, - k=self.head_dim, - warp_m=self.num_warps, - warp_n=1, - ) - self.pvc_config = self.cuda.resolve_dot_config( - dtype, - f32, - m=self.block_q, - n=self.head_dim, - k=self.block_kvc, - warp_m=self.num_warps, - warp_n=1, - ) - - # For self-attn - self.qk_config = self.cuda.resolve_dot_config( - dtype, - f32, - m=self.block_q, - n=self.block_kv, - k=self.head_dim, - warp_m=self.num_warps, - warp_n=1, - ) - self.pv_config = self.cuda.resolve_dot_config( - dtype, - f32, - m=self.block_q, - n=self.head_dim, - k=self.block_kv, - warp_m=self.num_warps, - warp_n=1, - ) - assert self.qk_config.lc == self.pv_config.la - - - def __call__(self, q_ptr: void_p, k_ptr: void_p, v_ptr: void_p, o_ptr: void_p, - k_cache_ptr: void_p, v_cache_ptr: void_p, page_table_ptr: void_p, - cu_seqlens_q_ptr: void_p, total_lens_ptr: void_p, ctxlens_ptr: void_p, - num_seqs: int, max_seqlen: int, q_len: int, kv_len: int, num_pages: int, max_seq_pages: int): - # TODO - # Setup Grid - self.attrs.warps = self.num_warps - self.attrs.blocks = (cdiv(max_seqlen, self.block_q), self.num_heads, num_seqs) - - # Get programs ids - start_m = self.blockIdx.x - head = self.blockIdx.y - seq = self.blockIdx.z - - # build-up global_views - global_q = self.global_view(q_ptr, dtype=self.dtype, shape=[q_len, self.num_heads, self.head_dim]) - global_k = self.global_view(k_ptr, dtype=self.dtype, shape=[kv_len, self.num_kv_heads, self.head_dim]) - global_v = self.global_view(v_ptr, dtype=self.dtype, shape=[kv_len, self.num_kv_heads, self.head_dim]) - global_o = self.global_view(o_ptr, dtype=self.dtype, shape=[q_len, self.num_heads, self.head_dim]) - global_k_cache = self.global_view(k_cache_ptr, dtype=self.dtype, shape=[num_pages, self.num_kv_heads, - self.head_dim_x, self.page_size, self.x]) - global_v_cache = self.global_view(v_cache_ptr, dtype=self.dtype, shape=[num_pages, self.num_kv_heads, - self.head_dim, self.page_size]) - global_page_table = self.global_view(page_table_ptr, dtype=int64, shape=[num_seqs, max_seq_pages]) - global_cu_seqlens_q = self.global_view(cu_seqlens_q_ptr, dtype=int32, shape=[num_seqs + 1]) - global_total_lens = self.global_view(total_lens_ptr, dtype=int32, shape=[num_seqs]) - global_ctxlens = self.global_view(ctxlens_ptr, dtype=int32, shape=[num_seqs]) - - # Allocate registers for q_start_idx, total_len, ctxlen - shared_q_start_idx = self.shared_tensor(dtype=int32, shape=[1]) - shared_total_len = self.shared_tensor(dtype=int32, shape=[1]) - shared_ctxlen = self.shared_tensor(dtype=int32, shape=[1]) - load_q_start_idx = self.load_global(global_cu_seqlens_q, offsets=[seq], shape=[1], dims=[0]) - load_total_len = self.load_global(global_total_lens, offsets=[seq], shape=[1], dims=[0]) - load_ctxlen = self.load_global(global_ctxlens, offsets=[seq], shape=[1], dims=[0]) - self.store_shared(shared_q_start_idx, load_q_start_idx) - self.store_shared(shared_total_len, load_total_len) - self.store_shared(shared_ctxlen, load_ctxlen) - self.sync() - q_start_idx = self.load_shared(shared_q_start_idx) - total_len = self.load_shared(shared_total_len) - ctxlen = self.load_shared(shared_ctxlen) - self.sync() - self.free_shared(shared_q_start_idx) - self.free_shared(shared_total_len) - self.free_shared(shared_ctxlen) - - # Load q tile into register - off_q = start_m * self.block_q + q_start_idx - shared_q = self.shared_tensor(dtype=self.dtype, shape=[self.block_q, self.head_dim]) - load_q = self.load_global(global_q, offsets=[off_q, head, 0], shape=[self.block_q, self.head_dim], dims=[0, 2]) - self.store_shared(shared_q, load_q) - self.sync() - q = self.load_shared(shared_q) - self.sync() - self.free_shared(shared_q) - - # Allocate shared memory for k, v, k_cache, and v_cache - shared_k = self.shared_tensor(dtype=self.dtype, shape=[self.block_kv, self.head_dim]) - shared_v = self.shared_tensor(dtype=self.dtype, shape=[self.block_kv, self.head_dim]) - shared_k_cache = self.shared_tensor(dtype=self.dtype, shape=[self.page_size, self.head_dim]) - shared_v_cache = self.shared_tensor(dtype=self.dtype, shape=[self.page_size, self.head_dim]) - shared_page_table = self.shared_tensor(dtype=int64, shape=[1]) - - # Init accumulators - acc = self.register_tensor(dtype=f32, shape=[self.block_q, self.head_dim], init=0.0) - m_i = self.register_tensor(dtype=f32, shape=[self.block_q, 1], init=-1e6) # rowmax(attn_score) - l_i = self.register_tensor(dtype=f32, shape=[self.block_q, 1], init=0.0) # rowsum(exp(attn_score - m_i)) - - # Pre-launch async copy for K Cache - self.copy_async(global_k_cache, shared_k_cache, - offsets=[seq_first_page, head // self.group_size, 0, 0, 0], dims=[2, 3, 4]) - self.copy_async_commit_group() - \ No newline at end of file diff --git a/diffulex/legacy/layers/attention/ops/triton_decode_attn_clm.py b/diffulex/legacy/layers/attention/ops/triton_decode_attn_clm.py deleted file mode 100755 index 71be2616..00000000 --- a/diffulex/legacy/layers/attention/ops/triton_decode_attn_clm.py +++ /dev/null @@ -1,681 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# type: ignore - -# Adapted from vllm -# https://github.com/vllm-project/vllm/blob/main/vllm/attention/ops/triton_decode_attention.py -# formerly adapted from -# https://github.com/sgl-project/sglang/blob/9f635ea50de920aa507f486daafba26a5b837574/python/sglang/srt/layers/attention/triton_ops/decode_attention.py -# which was originally adapted from -# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py -# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py - -# Changes: -# - Add support for page size >= 1. - -# Copyright 2025 vLLM Team -# Copyright 2023-2024 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -""" -Memory-efficient attention for decoding. -It supports page size >= 1. -""" - -import torch -import logging - -from vllm.platforms import current_platform -from vllm.triton_utils import tl, triton - -is_hip_ = current_platform.is_rocm() - -logger = logging.getLogger(__name__) - -# Only print the following warnings when triton version < 3.2.0. -# The issue won't affect performance or accuracy. -if triton.__version__ < '3.2.0': - logger.warning( - "The following error message 'operation scheduled before its operands' " - "can be ignored.") - - -@triton.jit -def tanh(x): - # Tanh is just a scaled sigmoid - return 2 * tl.sigmoid(2 * x) - 1 - - -@triton.jit -def _fwd_kernel_stage1( - Q, - K_Buffer, - V_Buffer, - sm_scale, - Req_to_tokens, - B_Seqlen, - Att_Out, - stride_req_to_tokens_b, - stride_qbs, - stride_qh, - stride_buf_kbs, - stride_buf_kh, - stride_buf_vbs, - stride_buf_vh, - stride_mid_ob, - stride_mid_oh, - stride_mid_os, - kv_group_num: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_DV: tl.constexpr, - BLOCK_N: tl.constexpr, - NUM_KV_SPLITS: tl.constexpr, - PAGE_SIZE: tl.constexpr, - logit_cap: tl.constexpr, - Lk: tl.constexpr, - Lv: tl.constexpr, -): - cur_batch = tl.program_id(0) - cur_head = tl.program_id(1) - split_kv_id = tl.program_id(2) - - cur_kv_head = cur_head // kv_group_num - - offs_d = tl.arange(0, BLOCK_DMODEL) - offs_dv = tl.arange(0, BLOCK_DV) - mask_d = offs_d < Lk - mask_dv = offs_dv < Lv - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - cur_batch_req_idx = cur_batch - - off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d - q = tl.load(Q + off_q, mask=mask_d, other=0.0) - - kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) - split_kv_start = kv_len_per_split * split_kv_id - split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, - cur_batch_seq_len) - - e_max = -float("inf") - e_sum = 0.0 - acc = tl.zeros([BLOCK_DV], dtype=tl.float32) - - if split_kv_end > split_kv_start: - for start_n in range(split_kv_start, split_kv_end, BLOCK_N): - offs_n = start_n + tl.arange(0, BLOCK_N) - kv_page_number = tl.load( - Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + - offs_n // PAGE_SIZE, - mask=offs_n < split_kv_end, - other=0, - ) - kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE - offs_buf_k = (kv_loc[:, None] * stride_buf_kbs + - cur_kv_head * stride_buf_kh + offs_d[None, :]) - k = tl.load( - K_Buffer + offs_buf_k, - mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]), - other=0.0, - ) - qk = tl.sum(q[None, :] * k, 1) - qk *= sm_scale - - if logit_cap > 0: - qk = logit_cap * tanh(qk / logit_cap) - - qk = tl.where(offs_n < split_kv_end, qk, float("-inf")) - - offs_buf_v = (kv_loc[:, None] * stride_buf_vbs + - cur_kv_head * stride_buf_vh + offs_dv[None, :]) - v = tl.load( - V_Buffer + offs_buf_v, - mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), - other=0.0, - ) - - n_e_max = tl.maximum(tl.max(qk, 0), e_max) - re_scale = tl.exp(e_max - n_e_max) - p = tl.exp(qk - n_e_max) - acc *= re_scale - acc += tl.sum(p[:, None] * v, 0) - - e_sum = e_sum * re_scale + tl.sum(p, 0) - e_max = n_e_max - - offs_mid_o = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh + - split_kv_id * stride_mid_os + offs_dv) - - tl.store( - Att_Out + offs_mid_o, - acc / e_sum, - mask=(mask_dv), - ) - - offs_mid_o_1 = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh + - split_kv_id * stride_mid_os + Lv) - - tl.store( - Att_Out + offs_mid_o_1, - e_max + tl.log(e_sum), - ) - - -def _decode_attn_m_fwd( - q, - k_buffer, - v_buffer, - att_out, - Req_to_tokens, - B_Seqlen, - num_kv_splits, - sm_scale, - page_size, - logit_cap, -): - BLOCK = 64 if not is_hip_ else 8 - - NUM_KV_SPLITS = num_kv_splits - Lk = k_buffer.shape[-1] - Lv = v_buffer.shape[-1] - - batch, head_num = q.shape[0], q.shape[1] - - grid = (batch, head_num, NUM_KV_SPLITS) - kv_group_num = q.shape[1] // k_buffer.shape[-2] - - num_warps = 4 - if kv_group_num != 1: - num_warps = 1 if is_hip_ else 2 - - BLOCK_DMODEL = triton.next_power_of_2(Lk) - BLOCK_DV = triton.next_power_of_2(Lv) - - _fwd_kernel_stage1[grid]( - q, - k_buffer, - v_buffer, - sm_scale, - Req_to_tokens, - B_Seqlen, - att_out, - Req_to_tokens.stride(0), - q.stride(0), - q.stride(1), - k_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) - k_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) - v_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) - v_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) - att_out.stride(0), - att_out.stride(1), - att_out.stride(2), - kv_group_num=kv_group_num, - BLOCK_DMODEL=BLOCK_DMODEL, - BLOCK_DV=BLOCK_DV, - BLOCK_N=BLOCK, - NUM_KV_SPLITS=NUM_KV_SPLITS, - PAGE_SIZE=page_size, - logit_cap=logit_cap, - num_warps=num_warps, - num_stages=2, - Lk=Lk, - Lv=Lv, - ) - - -@triton.jit -def _fwd_grouped_kernel_stage1( - Q, - K_Buffer, - V_Buffer, - sm_scale, - Req_to_tokens, - B_Seqlen, - Att_Out, - stride_req_to_tokens_b, - stride_qbs, - stride_qh, - stride_buf_kbs, - stride_buf_kh, - stride_buf_vbs, - stride_buf_vh, - stride_mid_ob, - stride_mid_oh, - stride_mid_os, - kv_group_num: tl.constexpr, - q_head_num: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_DPE: tl.constexpr, - BLOCK_DV: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_H: tl.constexpr, - NUM_KV_SPLITS: tl.constexpr, - PAGE_SIZE: tl.constexpr, - logit_cap: tl.constexpr, - Lk: tl.constexpr, - Lv: tl.constexpr, -): - cur_batch = tl.program_id(0) - cur_head_id = tl.program_id(1) - cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H) - split_kv_id = tl.program_id(2) - - if kv_group_num > BLOCK_H: - VALID_BLOCK_H: tl.constexpr = BLOCK_H - else: - VALID_BLOCK_H: tl.constexpr = kv_group_num - cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) - mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H - mask_h = mask_h & (cur_head < q_head_num) - - offs_d = tl.arange(0, BLOCK_DMODEL) - offs_dv = tl.arange(0, BLOCK_DV) - mask_d = offs_d < Lk - mask_dv = offs_dv < Lv - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - cur_batch_req_idx = cur_batch - - offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :] - q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0) - - if BLOCK_DPE > 0: - offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) - mask_dpe = offs_dpe < Lk - off_qpe = (cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]) - qpe = tl.load(Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0) - - kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) - split_kv_start = kv_len_per_split * split_kv_id - split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) - - e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") - e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) - acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) - - if split_kv_end > split_kv_start: - for start_n in range(split_kv_start, split_kv_end, BLOCK_N): - offs_n = start_n + tl.arange(0, BLOCK_N) - kv_page_number = tl.load( - Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n // PAGE_SIZE, - mask=offs_n < split_kv_end, other=0, - ) - kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE - offs_buf_k = (kv_loc[None, :] * stride_buf_kbs + cur_kv_head * stride_buf_kh + offs_d[:, None]) - k = tl.load(K_Buffer + offs_buf_k, mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]), other=0.0) - qk = tl.dot(q, k.to(q.dtype)) - if BLOCK_DPE > 0: - offs_buf_kpe = kv_loc[None, :] * stride_buf_kbs + cur_kv_head * stride_buf_kh + offs_dpe[:, None] - kpe = tl.load(K_Buffer + offs_buf_kpe, mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]), other=0.0) - qk += tl.dot(qpe, kpe.to(qpe.dtype)) - qk *= sm_scale - - if logit_cap > 0: - qk = logit_cap * tanh(qk / logit_cap) - - qk = tl.where(mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf")) - - offs_buf_v = kv_loc[:, None] * stride_buf_vbs + cur_kv_head * stride_buf_vh + offs_dv[None, :] - v = tl.load(V_Buffer + offs_buf_v, mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), other=0.0) - - n_e_max = tl.maximum(tl.max(qk, 1), e_max) - re_scale = tl.exp(e_max - n_e_max) - p = tl.exp(qk - n_e_max[:, None]) - acc *= re_scale[:, None] - acc += tl.dot(p.to(v.dtype), v) - - e_sum = e_sum * re_scale + tl.sum(p, 1) - e_max = n_e_max - - offs_mid_o = cur_batch * stride_mid_ob + cur_head[:, None] * stride_mid_oh + split_kv_id * stride_mid_os + offs_dv[None, :] - tl.store(Att_Out + offs_mid_o, acc / e_sum[:, None], mask=(mask_h[:, None]) & (mask_dv[None, :])) - offs_mid_o_1 = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + split_kv_id * stride_mid_os + Lv - - tl.store(Att_Out + offs_mid_o_1, e_max + tl.log(e_sum), mask=mask_h) - - -def _decode_grouped_attn_m_fwd( - q, - k_cache, - v_cache, - attn_out, - Req_to_tokens, - B_Seqlen, - num_kv_splits, - sm_scale, - page_size, - logit_cap, -): - BLOCK = 32 - Lk = k_cache.shape[-1] - Lv = v_cache.shape[-1] - - # [TODO] work around shmem limit on MI3xx - if is_hip_ and Lk >= 576: - BLOCK = 16 - - if Lk == 576: - BLOCK_DMODEL = 512 - BLOCK_DPE = 64 - elif Lk == 288: - BLOCK_DMODEL = 256 - BLOCK_DPE = 32 - else: - BLOCK_DMODEL = triton.next_power_of_2(Lk) - BLOCK_DPE = 0 - BLOCK_DV = triton.next_power_of_2(Lv) - - batch, head_num = q.shape[0], q.shape[1] - kv_group_num = q.shape[1] // k_cache.shape[-2] - - BLOCK_H = 16 - NUM_KV_SPLITS = num_kv_splits - grid = ( - batch, - triton.cdiv(head_num, min(BLOCK_H, kv_group_num)), - NUM_KV_SPLITS, - ) - - extra_kargs = {} - num_stages = 2 - if is_hip_: - # https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization/workload.html#mi300x-triton-kernel-performance-optimization - # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py - extra_kargs = { - "waves_per_eu": 1, - "matrix_instr_nonkdim": 16, - "kpack": 2 - } - num_stages = 1 - - _fwd_grouped_kernel_stage1[grid]( - q, - k_cache, - v_cache, - sm_scale, - Req_to_tokens, - B_Seqlen, - attn_out, - Req_to_tokens.stride(0), - q.stride(0), - q.stride(1), - k_cache.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) - k_cache.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) - v_cache.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) - v_cache.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM) - attn_out.stride(0), - attn_out.stride(1), - attn_out.stride(2), - kv_group_num=kv_group_num, - q_head_num=head_num, - BLOCK_DMODEL=BLOCK_DMODEL, - BLOCK_DPE=BLOCK_DPE, - BLOCK_DV=BLOCK_DV, - BLOCK_N=BLOCK, - BLOCK_H=BLOCK_H, - NUM_KV_SPLITS=NUM_KV_SPLITS, - PAGE_SIZE=page_size, - logit_cap=logit_cap, - num_warps=4, - num_stages=num_stages, - Lk=Lk, - Lv=Lv, - **extra_kargs, - ) - - -@triton.jit -def _fwd_kernel_stage2( - Mid_O, - o, - B_Seqlen, - stride_mid_ob, - stride_mid_oh, - stride_mid_os, - stride_obs, - stride_oh, - NUM_KV_SPLITS: tl.constexpr, - BLOCK_DV: tl.constexpr, - Lv: tl.constexpr, -): - cur_batch = tl.program_id(0) - cur_head = tl.program_id(1) - - cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) - - offs_d = tl.arange(0, BLOCK_DV) - mask_d = offs_d < Lv - - e_sum = 0.0 - e_max = -float("inf") - acc = tl.zeros([BLOCK_DV], dtype=tl.float32) - - offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d - offs_logic = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + Lv - - for split_kv_id in range(0, NUM_KV_SPLITS): - kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) - split_kv_start = kv_len_per_split * split_kv_id - split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, - cur_batch_seq_len) - - if split_kv_end > split_kv_start: - tv = tl.load(Mid_O + offs_v + split_kv_id * stride_mid_os, - mask=mask_d, - other=0.0) - tlogic = tl.load(Mid_O + offs_logic + split_kv_id * stride_mid_os) - n_e_max = tl.maximum(tlogic, e_max) - - old_scale = tl.exp(e_max - n_e_max) - acc *= old_scale - exp_logic = tl.exp(tlogic - n_e_max) - acc += exp_logic * tv - - e_sum = e_sum * old_scale + exp_logic - e_max = n_e_max - - tl.store( - o + cur_batch * stride_obs + cur_head * stride_oh + offs_d, - acc / e_sum, - mask=mask_d, - ) - - -def _decode_softmax_reducev_fwd( - logits, - q, - o, - v_buffer, - b_seq_len, - num_kv_splits, -): - batch, head_num = q.shape[0], q.shape[1] - Lv = v_buffer.shape[-1] - BLOCK_DV = triton.next_power_of_2(Lv) - - NUM_KV_SPLITS = num_kv_splits - - extra_kargs = {} - if is_hip_: - # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html - # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py - extra_kargs = { - "waves_per_eu": 4, - "matrix_instr_nonkdim": 16, - "kpack": 2 - } - - grid = (batch, head_num) - _fwd_kernel_stage2[grid]( - logits, - o, - b_seq_len, - logits.stride(0), - logits.stride(1), - logits.stride(2), - o.stride(0), - o.stride(1), - NUM_KV_SPLITS=NUM_KV_SPLITS, - BLOCK_DV=BLOCK_DV, - Lv=Lv, - num_warps=4, - num_stages=2, - **extra_kargs, - ) - - -def decode_attention_fwd_normal( - q, - k_buffer, - v_buffer, - o, - req_to_token, - b_seq_len, - attn_logits, - num_kv_splits, - sm_scale, - page_size, - logit_cap=0.0, -): - _decode_attn_m_fwd( - q, - k_buffer, - v_buffer, - attn_logits, - req_to_token, - b_seq_len, - num_kv_splits, - sm_scale, - page_size, - logit_cap, - ) - _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, b_seq_len, - num_kv_splits) - - -def decode_attention_fwd_grouped( - q, - k_cache, - v_cache, - o, - req_to_token, - b_seq_len, - attn_logits, - num_kv_splits, - softmax_scale, - page_size, - logit_cap=0.0, -): - _decode_grouped_attn_m_fwd( - q, - k_cache, - v_cache, - attn_logits, - req_to_token, - b_seq_len, - num_kv_splits, - softmax_scale, - page_size, - logit_cap, - ) - _decode_softmax_reducev_fwd( - attn_logits, - q, - o, - v_cache, - b_seq_len, - num_kv_splits - ) - - -def causal_lm_decode_attention_fwd( - q, - k_cache, - v_cache, - block_tables, - cache_seqlens, - o=None, - attn_logits=None, - softmax_scale=None, - num_kv_splits=1, - page_size=1, - logit_cap=0.0, -): - """ - Forward pass for decode attention using Triton kernels. - - Args: - q: Query tensor of shape [batch_size, num_heads, head_dim]. - Contains the query vectors for the current decoding step. - k_cache: Key cache tensor storing all previous key vectors. - Shape depends on page_size but generally [..., page_size, num_kv_heads, head_dim]. - v_cache: Value cache tensor storing all previous value vectors. - Shape depends on page_size but generally [..., page_size, num_kv_heads, head_dim]. - o: Output tensor of shape [batch_size, num_heads, head_dim]. - Will store the computed attention output. - block_tables: Token mapping tensor that maps request indices to token positions - in the paged memory layout. Shape [batch_size, max_seq_len // page_size]. - cache_seqlens: Batch sequence lengths tensor of shape [batch_size]. - Contains the actual sequence length for each batch item. - attn_logits: Intermediate attention logits tensor used for computation splits. - Shape [batch_size, num_heads, num_kv_splits, head_dim + 1]. - The extra "+1" dimension stores log-sum-exp values (e_max + log(e_sum)) - at index head_dim, while indices 0:head_dim store the attention outputs - for each split. This is needed for numerically stable softmax reduction - across splits in the second stage. - num_kv_splits: Number of splits for KV cache processing to manage memory usage. - Higher values reduce memory but may increase computation overhead. - softmax_scale: Scaling factor applied to attention scores before softmax. - Typically 1/sqrt(head_dim) for scaled dot-product attention. - page_size: Size of each page in the paged attention memory layout. Default is 1. - Larger page sizes can improve memory efficiency. - logit_cap: Optional logit capping value. If > 0, applies tanh-based capping to - attention logits to prevent overflow. Default is 0.0 (no capping). - """ - kv_group_num = q.shape[1] // v_cache.shape[-2] - - o = o if o is not None else torch.empty_like(q).to(q.device, q.dtype) - batch_size, num_heads, head_dim = q.shape # In CausalLM: batch_size = num_seqs - attn_logits_shape = (batch_size, num_heads, num_kv_splits, head_dim + 1) - attn_logits = attn_logits if attn_logits is not None else torch.empty(attn_logits_shape).to(q.device, q.dtype) - softmax_scale = q.shape[-1] ** (-0.5) if softmax_scale is None else softmax_scale - assert num_kv_splits == attn_logits.shape[2] - if kv_group_num == 1: - # MHA - decode_attention_fwd_normal( - q, - k_cache, - v_cache, - o, - block_tables, - cache_seqlens, - attn_logits, - num_kv_splits, - softmax_scale, - page_size, - logit_cap, - ) - else: - # GQA/MQA/MLA - decode_attention_fwd_grouped( - q, - k_cache, - v_cache, - o, - block_tables, - cache_seqlens, - attn_logits, - num_kv_splits, - softmax_scale, - page_size, - logit_cap, - ) - return o \ No newline at end of file diff --git a/diffulex/legacy/layers/attention/ops/triton_flash_attention.py b/diffulex/legacy/layers/attention/ops/triton_flash_attention.py deleted file mode 100755 index 37dd5356..00000000 --- a/diffulex/legacy/layers/attention/ops/triton_flash_attention.py +++ /dev/null @@ -1,1022 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from vLLM -# https://github.com/vllm-project/vllm/blob/main/vllm/attention/ops/triton_flash_attention.py -# type: ignore -""" -Fused Attention -=============== - -This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao -(https://tridao.me/publications/flash2/flash2.pdf) -Credits: OpenAI kernel team, AMD ML Frameworks Triton team - -Features supported: - -1) Fwd with causal masking -2) Any sequence lengths without padding (currently fwd kernel only) -3) Support for different sequence lengths for q and k -4) Nested tensor API currently does not support dropout or bias. - -Not currently supported: - -1) Non power of two head dims - -""" - -import torch - -from vllm.platforms import current_platform -from vllm.triton_utils import tl, triton - -# Avoid misleading ROCm warning. -if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx1x -else: - on_gfx1x = lambda *args, **kwargs: False - -torch_dtype: tl.constexpr = torch.float16 - - -@triton.jit -def cdiv_fn(x, y): - return (x + y - 1) // y - - -@triton.jit -def max_fn(x, y): - return tl.math.max(x, y) - - -@triton.jit -def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride): - ms = tl.arange(0, m) - ns = tl.arange(0, n) - return philox_offset + ms[:, None] * stride + ns[None, :] - - -@triton.jit -def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride): - rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, - stride).to(tl.uint32) - # TODO: use tl.randint for better performance - return tl.rand(philox_seed, rng_offsets) - - -@triton.jit -def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride): - rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n, - stride) - rng_keep = rng_output > dropout_p - return rng_keep - - -@triton.jit -def load_fn(block_ptr, first, second, pad): - if first and second: - tensor = tl.load(block_ptr, boundary_check=(0, 1), padding_option=pad) - elif first: - tensor = tl.load(block_ptr, boundary_check=(0, ), padding_option=pad) - elif second: - tensor = tl.load(block_ptr, boundary_check=(1, ), padding_option=pad) - else: - tensor = tl.load(block_ptr) - return tensor - - -@triton.jit -def _attn_fwd_inner( - acc, - l_i, - m_i, - q, - K_block_ptr, - V_block_ptr, - start_m, - actual_seqlen_k, - dropout_p, - philox_seed, - batch_philox_offset, - encoded_softmax_block_ptr, - block_min, - block_max, - offs_n_causal, - masked_blocks, - n_extra_tokens, - bias_ptr, - IS_CAUSAL: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, - OFFS_M: tl.constexpr, - OFFS_N: tl.constexpr, - PRE_LOAD_V: tl.constexpr, - MASK_STEPS: tl.constexpr, - ENABLE_DROPOUT: tl.constexpr, - RETURN_ENCODED_SOFTMAX: tl.constexpr, - PADDED_HEAD: tl.constexpr, - USE_FP8: tl.constexpr, - qk_scale, - p_descale, -): - # loop over k, v, and update accumulator - for start_n in range(block_min, block_max, BLOCK_N): - # For padded blocks, we will overrun the tensor size if - # we load all BLOCK_N. For others, the blocks are all within range. - k = load_fn( - K_block_ptr, - PADDED_HEAD, - MASK_STEPS and (n_extra_tokens != 0), - "zero", - ) - if PRE_LOAD_V: - v = load_fn( - V_block_ptr, - MASK_STEPS and (n_extra_tokens != 0), - PADDED_HEAD, - "zero", - ) - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - # We start from end of seqlen_k so only the first iteration would need - # to be checked for padding if it is not a multiple of block_n - # TODO: This can be optimized to only be true for the padded block. - if MASK_STEPS: # noqa: SIM102 - # If this is the last block / iteration, we want to - # mask if the sequence length is not a multiple of block size - # a solution is to always do BLOCK_M // BLOCK_N + 1 steps - # if not is_modulo_mn. last step might get wasted but that is okay. - # check if this masking works for that case. - if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0): - boundary_m = tl.full([BLOCK_M], - actual_seqlen_k, - dtype=tl.int32) - size_n = start_n + OFFS_N[None, :] - mask = size_n < boundary_m[:, None] - qk = tl.where(mask, qk, float("-inf")) - if IS_CAUSAL: - causal_boundary = start_n + offs_n_causal - causal_mask = OFFS_M[:, None] >= causal_boundary[None, :] - qk = tl.where(causal_mask, qk, float("-inf")) - # -- compute qk ---- - qk += tl.dot(q, k) - if USE_FP8: - qk *= qk_scale - if bias_ptr is not None: - bias = load_fn(bias_ptr, False, MASK_STEPS - and (n_extra_tokens != 0), "zero") - # While bias is added after multiplying qk with sm_scale, our - # optimization to use 2^x instead of e^x results in an additional - # scale factor of log2(e) which we must also multiply the bias with. - qk += bias * 1.44269504089 - m_ij = tl.maximum(m_i, tl.max(qk, 1)) - qk = qk - m_ij[:, None] - p = tl.math.exp2(qk) - - # CAVEAT: Must update l_ij before applying dropout - l_ij = tl.sum(p, 1) - if ENABLE_DROPOUT: - philox_offset = (batch_philox_offset + - start_m * BLOCK_M * actual_seqlen_k + start_n - - BLOCK_N) - keep = dropout_mask( - philox_seed, - philox_offset, - dropout_p, - BLOCK_M, - BLOCK_N, - actual_seqlen_k, - ) - if RETURN_ENCODED_SOFTMAX: - tl.store( - encoded_softmax_block_ptr, - tl.where(keep, p, - -p).to(encoded_softmax_block_ptr.type.element_ty), - ) - p = tl.where(keep, p, 0.0) - elif RETURN_ENCODED_SOFTMAX: - tl.store( - encoded_softmax_block_ptr, - p.to(encoded_softmax_block_ptr.type.element_ty), - ) - # -- update output accumulator -- - alpha = tl.math.exp2(m_i - m_ij) - acc = acc * alpha[:, None] - if not PRE_LOAD_V: - v = load_fn( - V_block_ptr, - MASK_STEPS and (n_extra_tokens != 0), - PADDED_HEAD, - "zero", - ) - # -- update m_i and l_i - l_i = l_i * alpha + l_ij - # update m_i and l_i - m_i = m_ij - - if USE_FP8: - p *= p_descale - - acc += tl.dot(p.to(V_block_ptr.type.element_ty), v) - - V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) - K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N)) - if bias_ptr is not None: - bias_ptr = tl.advance(bias_ptr, (0, BLOCK_N)) - if RETURN_ENCODED_SOFTMAX: - encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr, - (0, BLOCK_N)) - return acc, l_i, m_i - - -def get_cdna_autotune_configs(): - return [ - triton.Config( - { - 'BLOCK_M': 256, - 'BLOCK_N': 64, - 'waves_per_eu': 2, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=8), - triton.Config( - { - 'BLOCK_M': 128, - 'BLOCK_N': 128, - 'waves_per_eu': 2, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=4), - triton.Config( - { - 'BLOCK_M': 256, - 'BLOCK_N': 128, - 'waves_per_eu': 2, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=8), - triton.Config( - { - 'BLOCK_M': 128, - 'BLOCK_N': 64, - 'waves_per_eu': 1, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=4), - triton.Config( - { - 'BLOCK_M': 128, - 'BLOCK_N': 64, - 'waves_per_eu': 3, - 'PRE_LOAD_V': True - }, - num_stages=1, - num_warps=4), - triton.Config( - { - 'BLOCK_M': 128, - 'BLOCK_N': 64, - 'waves_per_eu': 3, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=4), - triton.Config( - { - 'BLOCK_M': 64, - 'BLOCK_N': 64, - 'waves_per_eu': 4, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=8), - triton.Config( - { - 'BLOCK_M': 32, - 'BLOCK_N': 32, - 'waves_per_eu': 4, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=8), - # TODO: This config fails with head_size not pow2 with data mismatches. - # triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 1, - # 'PRE_LOAD_V': False}, num_stages=1, num_warps=4), - - # Fails in AccelerateAMDMatmul (Triton) assert when using FP8: - # triton.Config( - # { - # "BLOCK_M": 16, - # "BLOCK_N": 16, - # "waves_per_eu": 1, - # "PRE_LOAD_V": False, - # }, - # num_stages=1, - # num_warps=4, - # ), - ], ['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL', 'USE_FP8'] - - -def get_rdna_autotune_configs(): - return [ - triton.Config( - { - 'BLOCK_M': 32, - 'BLOCK_N': 32, - 'waves_per_eu': 4, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=2), - triton.Config( - { - 'BLOCK_M': 32, - 'BLOCK_N': 32, - 'waves_per_eu': 2, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=2), - triton.Config( - { - 'BLOCK_M': 32, - 'BLOCK_N': 16, - 'waves_per_eu': 4, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=2), - triton.Config( - { - 'BLOCK_M': 32, - 'BLOCK_N': 16, - 'waves_per_eu': 2, - 'PRE_LOAD_V': False - }, - num_stages=1, - num_warps=2), - # Fails in AccelerateAMDMatmul (Triton) assert when using FP8: - # triton.Config( - # { - # 'BLOCK_M': 16, - # 'BLOCK_N': 16, - # 'waves_per_eu': 4, - # 'PRE_LOAD_V': False - # }, - # num_stages=1, - # num_warps=2), - # triton.Config( - # { - # 'BLOCK_M': 16, - # 'BLOCK_N': 16, - # 'waves_per_eu': 2, - # 'PRE_LOAD_V': False - # }, - # num_stages=1, - # num_warps=2), - # # Fall-back config. - # triton.Config( - # { - # 'BLOCK_M': 16, - # 'BLOCK_N': 16, - # 'waves_per_eu': 1, - # 'PRE_LOAD_V': False - # }, - # num_stages=1, - # num_warps=2), - ], ['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL', 'USE_FP8'] - - -def get_autotune_configs(): - if on_gfx1x(): - return get_rdna_autotune_configs() - else: - return get_cdna_autotune_configs() - - -autotune_configs, autotune_keys = get_autotune_configs() - -float8_info = torch.finfo(current_platform.fp8_dtype()) - - -@triton.autotune( - configs=autotune_configs, - key=autotune_keys, -) -@triton.jit -def attn_fwd( - Q, - K, - V, - bias, - sm_scale, - q_scale, - k_scale, - v_scale, - p_scale, - p_descale, - o_descale, - L, - Out, - stride_qz: tl.int64, - stride_qh: tl.int64, - stride_qm: tl.int64, - stride_qk: tl.int64, - stride_kz: tl.int64, - stride_kh: tl.int64, - stride_kn: tl.int64, - stride_kk: tl.int64, - stride_vz: tl.int64, - stride_vh: tl.int64, - stride_vk: tl.int64, - stride_vn: tl.int64, - stride_oz: tl.int64, - stride_oh: tl.int64, - stride_om: tl.int64, - stride_on: tl.int64, - stride_bz: tl.int64, - stride_bh: tl.int64, - stride_bm: tl.int64, - stride_bn: tl.int64, - cu_seqlens_q, - cu_seqlens_k, - dropout_p, - philox_seed, - philox_offset_base, - encoded_softmax, - HQ: tl.constexpr, - HK: tl.constexpr, - ACTUAL_BLOCK_DMODEL: tl.constexpr, - MAX_SEQLENS_Q: tl.constexpr, - MAX_SEQLENS_K: tl.constexpr, - VARLEN: tl.constexpr, - IS_CAUSAL: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - USE_FP8: tl.constexpr, - USE_FP8_OUT: tl.constexpr, - BLOCK_N: tl.constexpr, - PRE_LOAD_V: tl.constexpr, - BIAS_TYPE: tl.constexpr, - ENABLE_DROPOUT: tl.constexpr, - RETURN_ENCODED_SOFTMAX: tl.constexpr, - FP8_MIN: tl.constexpr = float8_info.min, - FP8_MAX: tl.constexpr = float8_info.max, -): - start_m = tl.program_id(0) - off_h_q = tl.program_id(1) - off_z = tl.program_id(2) - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, BLOCK_N) - if VARLEN: - cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z) - cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1) - seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start - # We have a one-size-fits-all grid in id(0). Some seqlens might be too - # small for all start_m so for those we return early. - if start_m * BLOCK_M > seqlen_q: - return - cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z) - cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1) - seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start - else: - cu_seqlens_q_start = 0 - cu_seqlens_k_start = 0 - seqlen_q = MAX_SEQLENS_Q - seqlen_k = MAX_SEQLENS_K - - # Now we compute whether we need to exit early due to causal masking. - # This is because for seqlen_q > seqlen_k, M rows of the attn scores - # are completely masked, resulting in 0s written to the output, and - # inf written to LSE. We don't need to do any GEMMs in this case. - # This block of code determines what N is, and if this WG is operating - # on those M rows. - n_blocks = cdiv_fn(seqlen_k, BLOCK_N) - if IS_CAUSAL: - # If seqlen_q == seqlen_k, the attn scores are a square matrix. - # If seqlen_q != seqlen_k, attn scores are rectangular which means - # the causal mask boundary is bottom right aligned, and ends at either - # the top edge (seqlen_q < seqlen_k) or left edge. - # This captures the decrease in n_blocks if we have a rectangular attn - # matrix - n_blocks_seqlen = cdiv_fn( - (start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N) - # This is what adjusts the block_max for the current WG, only - # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks - n_blocks = min(n_blocks, n_blocks_seqlen) - # If we have no blocks after adjusting for seqlen deltas, this WG is - # part of the blocks that are all 0. We exit early. - if n_blocks <= 0: - o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om + - off_h_q * stride_oh) - O_block_ptr = tl.make_block_ptr( - base=Out + o_offset, - shape=(seqlen_q, BLOCK_DMODEL), - strides=(stride_om, stride_on), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_DMODEL), - order=(1, 0), - ) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty) - # We still need to write 0s to the result - # tl.store(O_block_ptr, - # acc.to(Out.type.element_ty), boundary_check=(0,1)) - # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q - # + offs_m - # We store inf to LSE, not -inf because in the bwd pass, - # we subtract this - # from qk which makes it -inf, such that exp(qk - inf) = 0 - # for these masked blocks. - # l = tl.full([BLOCK_M], value=float("inf"), dtype=tl.float32) - # tl.store(l_ptrs, l) - # TODO: Should dropout and return encoded softmax be handled here? - return - - # If MQA / GQA, set the K and V head offsets appropriately. - GROUP_SIZE: tl.constexpr = HQ // HK - off_h_k = off_h_q // GROUP_SIZE if GROUP_SIZE != 1 else off_h_q - - n_extra_tokens = 0 - if seqlen_k < BLOCK_N: - n_extra_tokens = BLOCK_N - seqlen_k - elif seqlen_k % BLOCK_N: - n_extra_tokens = seqlen_k % BLOCK_N - padded_head = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL - - # Compute pointers for all the tensors used in this kernel. - q_offset = (off_z * stride_qz + off_h_q * stride_qh + - cu_seqlens_q_start * stride_qm) - Q_block_ptr = tl.make_block_ptr( - base=Q + q_offset, - shape=(seqlen_q, ACTUAL_BLOCK_DMODEL), - strides=(stride_qm, stride_qk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_DMODEL), - order=(1, 0), - ) - k_offset = (off_z * stride_kz + off_h_k * stride_kh + - cu_seqlens_k_start * stride_kn) - K_block_ptr = tl.make_block_ptr( - base=K + k_offset, - shape=(ACTUAL_BLOCK_DMODEL, seqlen_k), - strides=(stride_kk, stride_kn), - offsets=(0, 0), - block_shape=(BLOCK_DMODEL, BLOCK_N), - order=(0, 1), - ) - v_offset = (off_z * stride_vz + off_h_k * stride_vh + - cu_seqlens_k_start * stride_vk) - V_block_ptr = tl.make_block_ptr( - base=V + v_offset, - shape=(seqlen_k, ACTUAL_BLOCK_DMODEL), - strides=(stride_vk, stride_vn), - offsets=(0, 0), - block_shape=(BLOCK_N, BLOCK_DMODEL), - order=(1, 0), - ) - if BIAS_TYPE != 0: - bias_ptr = tl.make_block_ptr( - base=bias + off_h_q * stride_bh, - shape=(seqlen_q, seqlen_k), - strides=(stride_bm, stride_bn), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_N), - order=(1, 0), - ) - else: - bias_ptr = None - if ENABLE_DROPOUT: - batch_philox_offset = philox_offset_base \ - + (off_z * HQ + off_h_q) \ - * seqlen_q * seqlen_k - else: - batch_philox_offset = 0 - # We can ask to return the dropout mask without actually doing any dropout. - # In this case, we return an invalid pointer so indicate the mask is not i - # valid. - # TODO: Fix encoded softmax. It currently uses just h_q in the base offset. - if RETURN_ENCODED_SOFTMAX: - encoded_softmax_block_ptr = tl.make_block_ptr( - base=encoded_softmax + off_h_q * seqlen_q * seqlen_k, - shape=(seqlen_q, seqlen_k), - strides=(seqlen_k, 1), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_N), - order=(1, 0), - ) - else: - encoded_softmax_block_ptr = 0 - # initialize pointer to m and l - m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) - l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) - # scale sm_scale by log_2(e) and use 2^x in the loop as we do not - # have native e^x support in HW. - qk_scale = sm_scale * 1.44269504089 - # Q is loaded once at the beginning and shared by all N blocks. - q = load_fn(Q_block_ptr, True, padded_head, "zero") - if not USE_FP8: - q = (q * qk_scale).to(Q_block_ptr.type.element_ty) - acc_scale = 1.0 - else: - qk_scale *= q_scale * k_scale - acc_scale = p_scale * v_scale - - # Here we compute how many full and masked blocks we have. - padded_block_k = n_extra_tokens != 0 - is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0) - if IS_CAUSAL: - # There are always at least BLOCK_M // BLOCK_N masked blocks. - # Additionally there might be one more due to dissimilar seqlens. - masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn) - else: - # Padding on Q does not need to be masked in the FA loop. - masked_blocks = padded_block_k - # if IS_CAUSAL, not is_modulo_mn does not always result in an additional - # block. In this case we might exceed n_blocks so pick the min. - masked_blocks = min(masked_blocks, n_blocks) - n_full_blocks = n_blocks - masked_blocks - block_min = 0 - block_max = n_blocks * BLOCK_N - # Compute for full blocks. Here we set causal to false regardless of its - # value because there is no masking. Similarly we do not need padding. - if n_full_blocks > 0: - block_max = (n_blocks - masked_blocks) * BLOCK_N - acc, l_i, m_i = _attn_fwd_inner( - acc, - l_i, - m_i, - q, - K_block_ptr, - V_block_ptr, - start_m, - seqlen_k, - dropout_p, - philox_seed, - batch_philox_offset, - encoded_softmax_block_ptr, - # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _ - block_min, - block_max, - 0, - 0, - 0, - bias_ptr, - # IS_CAUSAL, .... - False, - BLOCK_M, - BLOCK_DMODEL, - BLOCK_N, - offs_m, - offs_n, - # _, MASK_STEPS, ... - PRE_LOAD_V, - False, - ENABLE_DROPOUT, - RETURN_ENCODED_SOFTMAX, - padded_head, - USE_FP8, - qk_scale, - p_descale, - ) - block_min = block_max - block_max = n_blocks * BLOCK_N - - tl.debug_barrier() - # Remaining blocks, if any, are full / not masked. - if masked_blocks > 0: - offs_n_causal = offs_n + (seqlen_q - seqlen_k) if IS_CAUSAL else 0 - K_block_ptr = tl.advance(K_block_ptr, (0, n_full_blocks * BLOCK_N)) - V_block_ptr = tl.advance(V_block_ptr, (n_full_blocks * BLOCK_N, 0)) - if bias_ptr is not None: - bias_ptr = tl.advance(bias_ptr, (0, n_full_blocks * BLOCK_N)) - if RETURN_ENCODED_SOFTMAX: - encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr, - (0, n_full_blocks)) - acc, l_i, m_i = _attn_fwd_inner( - acc, - l_i, - m_i, - q, - K_block_ptr, - V_block_ptr, - start_m, - seqlen_k, - dropout_p, - philox_seed, - batch_philox_offset, - encoded_softmax_block_ptr, - block_min, - block_max, - offs_n_causal, - masked_blocks, - n_extra_tokens, - bias_ptr, - IS_CAUSAL, - BLOCK_M, - BLOCK_DMODEL, - BLOCK_N, - offs_m, - offs_n, - # _, MASK_STEPS, ... - PRE_LOAD_V, - True, - ENABLE_DROPOUT, - RETURN_ENCODED_SOFTMAX, - padded_head, - USE_FP8, - qk_scale, - p_descale, - ) - # epilogue - - if USE_FP8: - acc *= acc_scale - acc = acc / l_i[:, None] - if ENABLE_DROPOUT: - acc = acc / (1 - dropout_p) - # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M, - # then we have one block with a row of all NaNs which come from computing - # softmax over a row of all -infs (-inf - inf = NaN). We check for that here - # and store 0s where there are NaNs as these rows should've been zeroed out. - end_m_idx = (start_m + 1) * BLOCK_M - start_m_idx = start_m * BLOCK_M - causal_start_idx = seqlen_q - seqlen_k - if USE_FP8_OUT: - acc *= o_descale - acc = tl.clamp(acc, FP8_MIN, FP8_MAX) - acc = acc.to(Out.type.element_ty) - if IS_CAUSAL: # noqa: SIM102 - if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx: - out_mask_boundary = tl.full((BLOCK_DMODEL, ), - causal_start_idx, - dtype=tl.int32) - mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M) - out_ptrs_mask = (mask_m_offsets[:, None] - >= out_mask_boundary[None, :]) - z = tl.zeros((1, ), tl.float32) - acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty)) - # write back LSE - # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q + offs_m - # If seqlen_q not multiple of BLOCK_M, we need to mask out the last - # few rows. This is only true for the last M block. For others, - # overflow_size will be -ve - # overflow_size = end_m_idx - seqlen_q - # if overflow_size > 0: - # boundary = tl.full((BLOCK_M,), BLOCK_M - overflow_size, dtype=tl.int32) - # # This is a > check because mask being 0 blocks the store. - # l_ptrs_mask = boundary > tl.arange(0, BLOCK_M) - # tl.store(l_ptrs, m_i + tl.math.log2(l_i), mask=l_ptrs_mask) - # else: - # tl.store(l_ptrs, m_i + tl.math.log2(l_i)) - - # write back O - o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om + - off_h_q * stride_oh) - O_block_ptr = tl.make_block_ptr( - base=Out + o_offset, - shape=(seqlen_q, ACTUAL_BLOCK_DMODEL), - strides=(stride_om, stride_on), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_DMODEL), - order=(1, 0), - ) - # Need boundary check on this to make sure the padding from the - # Q and KV tensors in both dims are not part of what we store back. - # TODO: Do the boundary check optionally. - tl.store(O_block_ptr, acc, boundary_check=(0, 1)) - - -def check_args( - q, - k, - v, - o, - varlen=True, - max_seqlens=None, - cu_seqlens_q=None, - cu_seqlens_k=None, -): - assert q.dim() == k.dim() and q.dim() == v.dim() - if varlen: - assert q.dim() == 3 - total_q, nheads_q, head_size = q.shape - total_k, nheads_k, _ = k.shape - assert cu_seqlens_q is not None - assert cu_seqlens_k is not None - assert len(cu_seqlens_q) == len(cu_seqlens_k) - else: - assert q.dim() == 4 - batch, nheads_q, seqlen_q, head_size = q.shape - _, nheads_k, seqlen_k, _ = k.shape - assert max_seqlens > 0 - assert k.shape == v.shape - assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1] - # TODO: Change assert if we support qkl f8 and v f16 - assert q.dtype == k.dtype and q.dtype == v.dtype - assert head_size <= 256 - assert o.shape == q.shape - assert (nheads_q % nheads_k) == 0 - - -class _attention(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - q, - k, - v, - o, - cu_seqlens_q, - cu_seqlens_k, - max_seqlens_q, - max_seqlens_k, - causal=False, - sm_scale=1.0, - bias=None, - fp8_scales=None, - fp8_out_scale=None, - block_table=None, - ): - if block_table is not None: - raise NotImplementedError( - "Prefix Caching is not supported in this version, " - "block_table can only be None." - ) - if fp8_scales is not None: - use_fp8 = True - (q_scale, k_scale, v_scale, p_scale) = fp8_scales - float8 = current_platform.fp8_dtype() - - def check_and_convert(t, scale): - if t.dtype != float8: - descale = 1.0 / scale - ts = (t * descale).clamp(min=float8_info.min, - max=float8_info.max) - return ts.to(float8) - else: - return t - - q = check_and_convert(q, q_scale) - k = check_and_convert(k, k_scale) - v = check_and_convert(v, v_scale) - else: - use_fp8 = False - q_scale = k_scale = v_scale = p_scale = 1.0 - - if o is None: - o = torch.empty_like(q, dtype=v.dtype) - - check_args( - q, - k, - v, - o, - varlen=True, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - ) - if True: # varlen - total_q, nheads_q, head_size = q.shape - total_k, nheads_k, _ = k.shape - batch = len(cu_seqlens_q) - 1 - q_strides = (0, q.stride(1), q.stride(0), q.stride(2)) - k_strides = (0, k.stride(1), k.stride(0), k.stride(2)) - v_strides = (0, v.stride(1), v.stride(0), v.stride(2)) - o_strides = (0, o.stride(1), o.stride(0), o.stride(2)) - else: - batch, seqlen_q, nheads_q, head_size = q.shape - _, seqlen_k, nheads_k, _ = k.shape - q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3)) - k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3)) - v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3)) - o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3)) - - # Get closest power of 2 over or equal to 32. - unpadded_head_dims = {32, 64, 128, 256} - if head_size not in unpadded_head_dims: - padded_d_model = None - for i in unpadded_head_dims: - if i > head_size: - padded_d_model = i - break - assert padded_d_model is not None - else: - padded_d_model = head_size - - grid = lambda META: ( - triton.cdiv(max_seqlens_q, META["BLOCK_M"]), - nheads_q, - batch, - ) - - encoded_softmax = None - - # Seed the RNG so we get reproducible results for testing. - philox_seed = 0x1BF52 - philox_offset = 0x1D4B42 - - if bias is not None: - bias_strides = ( - bias.stride(0), - bias.stride(1), - bias.stride(2), - bias.stride(3), - ) - else: - bias_strides = (0, 0, 0, 0) - - p_descale = 1.0 / p_scale - o_descale = 1.0 / fp8_out_scale.item( - ) if fp8_out_scale is not None else 1.0 - - arg_max_seqlens_q = 0 if on_gfx1x() else max_seqlens_q - arg_max_seqlens_k = 0 if on_gfx1x() else max_seqlens_k - - attn_fwd[grid]( - q, - k, - v, - bias, - sm_scale, - q_scale, - k_scale, - v_scale, - p_scale, - p_descale, - o_descale, - None, - o, - *q_strides, - *k_strides, - *v_strides, - *o_strides, - *bias_strides, - cu_seqlens_q, - cu_seqlens_k, - dropout_p=0.0, - philox_seed=philox_seed, - philox_offset_base=philox_offset, - encoded_softmax=encoded_softmax, - HQ=nheads_q, - HK=nheads_k, - ACTUAL_BLOCK_DMODEL=head_size, - MAX_SEQLENS_Q=arg_max_seqlens_q, - MAX_SEQLENS_K=arg_max_seqlens_k, - IS_CAUSAL=causal, - VARLEN=True, - BLOCK_DMODEL=padded_d_model, - BIAS_TYPE=0 if bias is None else 1, - ENABLE_DROPOUT=False, - RETURN_ENCODED_SOFTMAX=False, - USE_FP8=use_fp8, - USE_FP8_OUT=fp8_out_scale is not None, - ) - - ctx.grid = grid - ctx.sm_scale = sm_scale - ctx.BLOCK_DMODEL = head_size - ctx.causal = causal - ctx.dropout_p = 0.0 - ctx.philox_seed = philox_seed - ctx.philox_offset = philox_offset - ctx.encoded_softmax = encoded_softmax - ctx.return_encoded_softmax = False - return o, encoded_softmax - -def triton_flash_attention( - q, - k, - v, - o, - cu_seqlens_q, - cu_seqlens_k, - max_seqlens_q, - max_seqlens_k, - causal=False, - softmax_scale=1.0, - bias=None, - fp8_scales=None, - fp8_out_scale=None, - block_table=None, -): - _attention.apply( - q, - k, - v, - o, - cu_seqlens_q, - cu_seqlens_k, - max_seqlens_q, - max_seqlens_k, - causal, - softmax_scale, - bias, - fp8_scales, - fp8_out_scale, - block_table, - ) \ No newline at end of file diff --git a/diffulex/model/__init__.py b/diffulex/model/__init__.py index 61e71e9e..12581e27 100644 --- a/diffulex/model/__init__.py +++ b/diffulex/model/__init__.py @@ -1,11 +1,25 @@ """Diffulex model package that imports built-in models to trigger registration.""" from __future__ import annotations +import importlib +from pathlib import Path # Import built-in models so their registrations run at import time. -from . import dream # noqa: F401 -from . import llada # noqa: F401 -from . import fast_dllm_v2 # noqa: F401 +# Automatically import all Python files except auto_model and __init__ +_excluded_modules = {"auto_model", "__init__"} +_model_modules = [] -__all__ = ["dream", "llada", "fast_dllm_v2"] +_current_dir = Path(__file__).parent +for py_file in _current_dir.glob("*.py"): + module_name = py_file.stem + if module_name not in _excluded_modules: + try: + importlib.import_module(f".{module_name}", __name__) + _model_modules.append(module_name) + except Exception as e: + # Skip modules that fail to import + import warnings + warnings.warn(f"Failed to import {module_name}: {e}", ImportWarning) + +__all__ = _model_modules.copy() from .auto_model import AutoModelForDiffusionLM \ No newline at end of file diff --git a/diffulex/model/config/sdar/configuration_sdar.py b/diffulex/model/config/sdar/configuration_sdar.py new file mode 100644 index 00000000..f2014181 --- /dev/null +++ b/diffulex/model/config/sdar/configuration_sdar.py @@ -0,0 +1,78 @@ +# coding=utf-8 +"""SDAR model configuration (Diffulex native).""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_rope_utils import rope_config_validation +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + + +class SDARConfig(PretrainedConfig): + model_type = "sdar" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 151936, + hidden_size: int = 4096, + intermediate_size: int = 22016, + num_hidden_layers: int = 32, + num_attention_heads: int = 32, + num_key_value_heads: int | None = 32, + head_dim: int | None = 128, + hidden_act: str = "silu", + max_position_embeddings: int = 32768, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-6, + use_cache: bool = False, # Diffulex uses its own KV cache path. + tie_word_embeddings: bool = False, + rope_theta: float = 10000.0, + rope_scaling=None, + attention_bias: bool = False, + use_sliding_window: bool = False, + sliding_window: int = 4096, + max_window_layers: int = 28, + attention_dropout: float = 0.0, + pad_token_id: int = 151643, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + # Backward compatibility. + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.tie_word_embeddings = tie_word_embeddings + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + + self.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window if use_sliding_window else None + self.max_window_layers = max_window_layers + self.attention_dropout = attention_dropout + + # Validate rotary position embedding parameters (Transformers helper). + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + rope_config_validation(self) + + super().__init__(tie_word_embeddings=tie_word_embeddings, pad_token_id=pad_token_id, **kwargs) + + +__all__ = ["SDARConfig"] + + diff --git a/diffulex/model/fast_dllm_v2.py b/diffulex/model/fast_dllm_v2.py index d707ebd8..126705b9 100755 --- a/diffulex/model/fast_dllm_v2.py +++ b/diffulex/model/fast_dllm_v2.py @@ -83,7 +83,6 @@ def __init__( self.head_dim, self.scaling, self.num_kv_heads, - "diffusion_lm", # Dream uses full attention ) def forward( diff --git a/diffulex/model/sdar.py b/diffulex/model/sdar.py index e69de29b..a733c453 100644 --- a/diffulex/model/sdar.py +++ b/diffulex/model/sdar.py @@ -0,0 +1,210 @@ +import os + +import torch +import torch.nn as nn +import torch.distributed as dist + +from diffulex.attention import Attention +from diffulex.layer.layernorm import RMSNorm +from diffulex.layer.activation import SiluAndMul +from diffulex.layer.rotary_embedding import get_rope +from diffulex.layer.linear import RowParallelLinear, ColumnParallelLinear +from diffulex.layer.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffulex.model.auto_model import AutoModelForDiffusionLM +from diffulex.model.config.sdar.configuration_sdar import SDARConfig + + +if os.environ.get("TRITON_INTERPRET", None) == "1": + torch._dynamo.reset() + torch._dynamo.config.suppress_errors = True + torch.backends.optimized_mode = False + + +class SDARAttention(nn.Module): + """SDAR attention (Diffulex native KV cache path). + + Compatible with Diffulex runner KV cache injection: + runner sets `self.attn.k_cache` / `self.attn.v_cache` by assigning to modules + that expose these attributes (see `diffulex/attention/attn_impl.py`). + """ + + def __init__(self, config: SDARConfig) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + + self.total_num_kv_heads = config.num_key_value_heads + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + + head_dim = getattr(config, "head_dim", None) + self.head_dim = head_dim or (config.hidden_size // self.total_num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + bias = getattr(config, "attention_bias", False) + self.q_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_heads * self.head_dim, + bias=bias, + ) + self.k_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=bias, + ) + self.v_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=bias, + ) + + # SDAR uses q/k per-head RMSNorm. + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=config.max_position_embeddings, + base=getattr(config, "rope_theta", 10000), + rope_scaling=getattr(config, "rope_scaling", None), + ) + + # Diffulex Attention implements KV cache store/load via injected k_cache/v_cache. + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + q = self.q_proj(hidden_states) + k = self.k_proj(hidden_states) + v = self.v_proj(hidden_states) + + # Per-head norm. + q_by_head = q.view(-1, self.num_heads, self.head_dim) + q_by_head = self.q_norm(q_by_head) + q = q_by_head.view(q.shape) + + k_by_head = k.view(-1, self.num_kv_heads, self.head_dim) + k_by_head = self.k_norm(k_by_head) + k = k_by_head.view(k.shape) + + q, k = self.rotary_emb(positions, q, k) + o = self.attn(q, k, v, mask) + return self.o_proj(o) + + +class SDARMLP(nn.Module): + """SDAR MLP: SiLU(gate) * up -> down.""" + + def __init__(self, config: SDARConfig) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear(config.hidden_size, config.intermediate_size, bias=False) + self.up_proj = ColumnParallelLinear(config.hidden_size, config.intermediate_size, bias=False) + self.down_proj = RowParallelLinear(config.intermediate_size, config.hidden_size, bias=False) + assert getattr(config, "hidden_act", "silu") == "silu" + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = self.gate_proj(x) + up = self.up_proj(x) + x = self.act_fn(torch.cat([gate, up], dim=-1)) + return self.down_proj(x) + + +class SDARDecoderLayer(nn.Module): + def __init__(self, config: SDARConfig) -> None: + super().__init__() + self.self_attn = SDARAttention(config) + self.mlp = SDARMLP(config) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self.self_attn(positions, hidden_states, mask) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class SDARModel(nn.Module): + def __init__(self, config: SDARConfig) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([SDARDecoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) + residual = None + for layer in self.layers: + hidden_states, residual = layer(positions, hidden_states, residual, mask) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +@AutoModelForDiffusionLM.register("sdar") +class SDARForDiffusionLM(nn.Module): + packed_modules_mapping = {} + + def __init__(self, config: SDARConfig) -> None: + super().__init__() + self.model = SDARModel(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + if getattr(config, "tie_word_embeddings", False): + self.lm_head.weight.data = self.model.embed_tokens.weight.data + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, mask) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lm_head(hidden_states) + + +__all__ = [ + "SDARConfig", + "SDARAttention", + "SDARMLP", + "SDARDecoderLayer", + "SDARModel", + "SDARForDiffusionLM", +] diff --git a/diffulex/sampler/__init__.py b/diffulex/sampler/__init__.py new file mode 100644 index 00000000..8f561b5f --- /dev/null +++ b/diffulex/sampler/__init__.py @@ -0,0 +1,25 @@ +"""Diffulex sampler package that imports built-in samplers to trigger registration.""" +from __future__ import annotations +import importlib +from pathlib import Path + +# Import built-in models so their registrations run at import time. +# Automatically import all Python files except auto_model and __init__ +_excluded_modules = {"auto_sampler", "__init__"} +_model_modules = [] + +_current_dir = Path(__file__).parent +for py_file in _current_dir.glob("*.py"): + module_name = py_file.stem + if module_name not in _excluded_modules: + try: + importlib.import_module(f".{module_name}", __name__) + _model_modules.append(module_name) + except Exception as e: + # Skip modules that fail to import + import warnings + warnings.warn(f"Failed to import {module_name}: {e}", ImportWarning) + +__all__ = _model_modules.copy() + +from .auto_sampler import AutoSampler \ No newline at end of file diff --git a/diffulex/sampler/auto_sampler.py b/diffulex/sampler/auto_sampler.py new file mode 100644 index 00000000..63641383 --- /dev/null +++ b/diffulex/sampler/auto_sampler.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from typing import Any, Callable + +from diffulex.config import Config + + +_NOT_PROVIDED = object() +RegistryEntry = tuple[Callable[[Any], Any] | type | None, bool] + + +class AutoSampler: + """Factory and registry for diffusion language model samplers.""" + + SAMPLER_MAPPING: dict[str, RegistryEntry] = {} + + @classmethod + def register( + cls, + sampler_name: str, + sampler_class: Callable[[Any], Any] | type | None = _NOT_PROVIDED, + *, + use_full_config: bool = False, + exist_ok: bool = False, + ): + """Register a sampler factory or class under ``sampler_name``. + + When ``sampler_class`` is omitted this method returns a decorator. + + Args: + sampler_name: Key used to retrieve the sampler. + sampler_class: Callable or class that builds the sampler instance. + use_full_config: Pass the entire :class:`Config` to the factory + instead of ``config.hf_config``. + exist_ok: Allow overriding an existing registration. + """ + + if not isinstance(sampler_name, str) or not sampler_name: + raise ValueError("sampler_name must be a non-empty string.") + + if sampler_class is _NOT_PROVIDED: + def decorator(sampler_cls): + cls._register(sampler_name, sampler_cls, use_full_config=use_full_config, exist_ok=exist_ok) + return sampler_cls + + return decorator + + cls._register(sampler_name, sampler_class, use_full_config=use_full_config, exist_ok=exist_ok) + return sampler_class + + @classmethod + def _register( + cls, + sampler_name: str, + sampler_class: Callable[[Any], Any] | type | None, + *, + use_full_config: bool, + exist_ok: bool, + ) -> None: + if not exist_ok and sampler_name in cls.SAMPLER_MAPPING: + raise ValueError(f"Sampler '{sampler_name}' is already registered.") + cls.SAMPLER_MAPPING[sampler_name] = (sampler_class, use_full_config) + + @classmethod + def unregister(cls, sampler_name: str) -> None: + cls.SAMPLER_MAPPING.pop(sampler_name, None) + + @classmethod + def available_samplers(cls) -> tuple[str, ...]: + return tuple(sorted(cls.SAMPLER_MAPPING)) + + @classmethod + def from_config(cls, config: Config): + if not hasattr(config, "model_name"): + raise AttributeError("Config must define 'model_name' to build a sampler.") + + try: + factory, use_full_config = cls.SAMPLER_MAPPING[config.model_name] + except KeyError as err: + available = ", ".join(cls.available_samplers()) or "" + raise ValueError( + f"Sampler '{config.model_name}' is not registered. Available samplers: {available}." + ) from err + + if factory is None: + raise ValueError(f"Sampler '{config.model_name}' is reserved but not implemented yet.") + + # Samplers don't require initialization arguments, they are nn.Module subclasses + sampler = factory() + return sampler \ No newline at end of file diff --git a/diffulex/sampler/base.py b/diffulex/sampler/base.py new file mode 100644 index 00000000..34f394fe --- /dev/null +++ b/diffulex/sampler/base.py @@ -0,0 +1,109 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributions as dists + +from dataclasses import dataclass +from easydict import EasyDict as edict + +from diffulex.engine.sequence import SequenceBase + + +class SamplerBase(nn.Module): + def __init__(self): + super().__init__() + from diffulex.attention import fetch_attn_metadata + self.fetch_attn_metadata = fetch_attn_metadata + + def top_p_logits(self, logits, top_p): + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() + sorted_indices_to_remove[..., 0] = 0 + + mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device) + mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove) + logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min) + return logits + + def top_k_logits(self, logits, top_k): + top_k = min(top_k, logits.size(-1)) + indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min) + return logits + + def sample_tokens(self, logits, temperature=0.0, top_p=None, top_k=None, + margin_confidence=False, neg_entropy=False): + if temperature > 0: + logits = logits / temperature + if top_p is not None and top_p < 1: + logits = self.top_p_logits(logits, top_p) + if top_k is not None: + logits = self.top_k_logits(logits, top_k) + probs = torch.softmax(logits, dim=-1) + + if temperature > 0: + try: + x0 = dists.Categorical(probs=probs).sample() + initial_confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1) + except: + initial_confidence, x0 = probs.max(dim=-1) + else: + initial_confidence, x0 = probs.max(dim=-1) + + confidence = initial_confidence.clone() + + if margin_confidence: + sorted_probs, _ = torch.sort(probs, dim=-1, descending=True) + top1_probs = sorted_probs[:, 0] + top2_probs = sorted_probs[:, 1] + confidence = top1_probs - top2_probs + + if neg_entropy: + epsilon = 1e-10 + log_probs = torch.log(probs + epsilon) + confidence = torch.sum(probs * log_probs, dim=-1) + + return confidence, x0, initial_confidence + + +@dataclass +class SampleOutputBase: + true_local_ids_map: dict[str, dict[str, list[int]]] + accepted_ids_map: dict[str, dict[str, list[int]]] + sampled_tokens_map: dict[str, dict[str, list[int]]] + + def __post_init__(self): + self.accepted_ids_map = edict(self.accepted_ids_map) + self.sampled_tokens_map = edict(self.sampled_tokens_map) + self.true_local_ids_map = edict(self.true_local_ids_map) + + +class SamplerShiftLogits(SamplerBase): + def __init__(self): + super().__init__() + self.seq_last_logits_map: dict[str, torch.Tensor] = {} + + def _fetch_last_logits(self, logits: torch.Tensor, seq: SequenceBase) -> torch.Tensor: + if seq.has_to_cache_block: + last_logits = logits[seq.to_cache_last_token_id] + self.seq_last_logits_map[seq.seq_id] = last_logits + return self.seq_last_logits_map[seq.seq_id] + + def _shift_logits(self, logits, last_logit=None): + if logits.shape[1] == 0: + print("Warning: logits sequence length is 0, returning empty logits") + raise Exception("logits sequence length is 0") + + shifted_logits = torch.zeros_like(logits) + shifted_logits[1:, ...] = logits[:-1, ...] + if last_logit is not None: + shifted_logits[0, ...] = last_logit + return shifted_logits + shifted_logits[0, ...] = 1.0 + return shifted_logits + + +class SamplerNoShiftLogits(SamplerBase): + pass \ No newline at end of file diff --git a/diffulex/sampler/dream.py b/diffulex/sampler/dream.py new file mode 100644 index 00000000..9f063408 --- /dev/null +++ b/diffulex/sampler/dream.py @@ -0,0 +1,73 @@ +import torch + +from dataclasses import dataclass + +from diffulex.sampler.auto_sampler import AutoSampler +from diffulex.sampler.base import SamplerShiftLogits, SampleOutputBase + + +@dataclass +class DreamSampleOutputForDiffusionLM(SampleOutputBase): + pass + + +@AutoSampler.register("dream") +class DreamSamplerForDiffusionLM(SamplerShiftLogits): + def forward(self, logits: torch.Tensor, temperatures: torch.Tensor, + top_p=None, top_k=None, margin_confidence=False, neg_entropy=False): + context = self.fetch_attn_metadata() + seqs = context.seqs + split_logits = torch.split(logits, [len(seq) for seq in seqs] if context.is_prefill else context.seq_lens, dim=0) + accepted_ids_map = {} + sampled_tokens_map = {} + true_local_ids_map = {} + for temperature, seq, seq_logits in zip(temperatures, seqs, split_logits): + true_local_ids_sub_map = {} + accepted_ids_sub_map = {} + sampled_tokens_sub_map = {} + + last_logits = self._fetch_last_logits(seq_logits, seq) + + shifted_logits = self._shift_logits(seq_logits, last_logits) + for block_id, block in enumerate(seq.diffusion_blocks): + if not block.is_active or sum(block.local_mask_tokens) == 0: + continue + + if len(block.global_mask_token_ids) > 0: + mask_token_logits = shifted_logits[block.global_mask_token_ids, ...] + confidence, sampled_tokens, initial_confidence = self.sample_tokens( + mask_token_logits, + temperature, + top_p=top_p, + top_k=top_k, + neg_entropy=(neg_entropy == "neg_entropy"), + margin_confidence=(margin_confidence == "margin_confidence") + ) + + if block.pre_block_complete: + high_conf_indices = torch.where(initial_confidence > block.accept_threshold)[0] + if len(high_conf_indices) == 0: + number_transfer_tokens = 1 + _, transfer_index = torch.topk(confidence, number_transfer_tokens) + else: + transfer_index = torch.tensor([], device=sampled_tokens.device, dtype=torch.long) + accepted_ids = torch.unique(torch.cat([transfer_index, high_conf_indices])) + else: + high_conf_indices = torch.where(initial_confidence > block.accept_threshold)[0] + accepted_ids = high_conf_indices + + true_local_ids_sub_map[str(block_id)] = [block.local_mask_token_ids[accepted_id] for accepted_id in accepted_ids.tolist()] + accepted_ids_sub_map[str(block_id)] = accepted_ids.tolist() + sampled_tokens_sub_map[str(block_id)] = sampled_tokens + + seq_idx = str(seq.seq_id) + true_local_ids_map[seq_idx] = true_local_ids_sub_map + accepted_ids_map[seq_idx] = accepted_ids_sub_map + sampled_tokens_map[seq_idx] = sampled_tokens_sub_map + + return DreamSampleOutputForDiffusionLM( + true_local_ids_map=true_local_ids_map, + accepted_ids_map=accepted_ids_map, + sampled_tokens_map=sampled_tokens_map + ) + diff --git a/diffulex/sampler/fast_dllm_v2.py b/diffulex/sampler/fast_dllm_v2.py new file mode 100644 index 00000000..ec323b50 --- /dev/null +++ b/diffulex/sampler/fast_dllm_v2.py @@ -0,0 +1,85 @@ +import torch + +from dataclasses import dataclass + +from diffulex.sampler.auto_sampler import AutoSampler +from diffulex.sampler.base import SamplerShiftLogits, SampleOutputBase +from diffulex.engine.sequence import SequenceBase + + +@dataclass +class FastdLLMV2SampleOutputForDiffusionLM(SampleOutputBase): + pass + + +@AutoSampler.register("fast_dllm_v2") +class FastdLLMV2SamplerForDiffusionLM(SamplerShiftLogits): + def forward(self, seqs: list[SequenceBase], logits: torch.Tensor, temperatures: torch.Tensor, + top_p=None, top_k=None, margin_confidence=False, neg_entropy=False, threshold=0.95): + attn_metadata = self.fetch_attn_metadata() + split_logits = torch.split( + logits, [len(seq) for seq in seqs] if attn_metadata.is_prefill + else [attn_metadata.diffusion_block_size] * len(seqs), dim=0 + ) + + accepted_ids_map = {} + sampled_tokens_map = {} + true_local_ids_map = {} + for temperature, seq, seq_logits in zip(temperatures, seqs, split_logits): + true_local_ids_sub_map = {} + accepted_ids_sub_map = {} + sampled_tokens_sub_map = {} + + last_logits = self._fetch_last_logits(seq_logits, seq) + + shifted_logits = self._shift_logits(seq_logits, last_logits) + + for block_id, block in enumerate(seq.diffusion_blocks): + if not block.is_active or sum(block.local_mask_tokens) == 0: + continue + + if len(block.global_mask_token_ids) == 0: + continue + + if attn_metadata.is_prefill: + mask_token_logits = shifted_logits[block.global_mask_token_ids, ...] + else: + mask_token_logits = shifted_logits[block.local_mask_token_ids, ...] + + confidence, sampled_tokens, initial_confidence = self.sample_tokens( + mask_token_logits, + temperature, + top_p=top_p, + top_k=top_k, + neg_entropy=(neg_entropy == "neg_entropy"), + margin_confidence=(margin_confidence == "margin_confidence") + ) + + high_conf_indices = torch.where(initial_confidence > threshold)[0] + + if len(high_conf_indices) == 0: + max_prob_idx = initial_confidence.argmax() + accepted_ids = torch.tensor([max_prob_idx], device=sampled_tokens.device, dtype=torch.long) + else: + max_prob_idx = initial_confidence.argmax() + accepted_ids = torch.unique(torch.cat([ + high_conf_indices, + torch.tensor([max_prob_idx], device=sampled_tokens.device, dtype=torch.long) + ])) + + true_local_ids_sub_map[str(block_id)] = [ + block.local_mask_token_ids[accepted_id] for accepted_id in accepted_ids.tolist() + ] + accepted_ids_sub_map[str(block_id)] = accepted_ids.tolist() + sampled_tokens_sub_map[str(block_id)] = sampled_tokens + + seq_idx = str(seq.seq_id) + true_local_ids_map[seq_idx] = true_local_ids_sub_map + accepted_ids_map[seq_idx] = accepted_ids_sub_map + sampled_tokens_map[seq_idx] = sampled_tokens_sub_map + + return FastdLLMV2SampleOutputForDiffusionLM( + true_local_ids_map=true_local_ids_map, + accepted_ids_map=accepted_ids_map, + sampled_tokens_map=sampled_tokens_map + ) \ No newline at end of file diff --git a/diffulex/sampler/llada.py b/diffulex/sampler/llada.py new file mode 100644 index 00000000..5202fa14 --- /dev/null +++ b/diffulex/sampler/llada.py @@ -0,0 +1,69 @@ +import torch + +from dataclasses import dataclass + +from diffulex.sampler.auto_sampler import AutoSampler +from diffulex.sampler.base import SamplerNoShiftLogits, SampleOutputBase + + +@dataclass +class LLaDASampleOutputForDiffusionLM(SampleOutputBase): + pass + + +@AutoSampler.register("llada") +class LLaDASamplerForDiffusionLM(SamplerNoShiftLogits): + def forward(self, logits: torch.Tensor, temperatures: torch.Tensor, + top_p=None, top_k=None, margin_confidence=False, neg_entropy=False): + context = self.fetch_attn_metadata() + seqs = context.seqs + split_logits = torch.split(logits, [len(seq) for seq in seqs] if context.is_prefill else context.seq_lens, dim=0) + accepted_ids_map = {} + sampled_tokens_map = {} + true_local_ids_map = {} + for temperature, seq, seq_logits in zip(temperatures, seqs, split_logits): + true_local_ids_sub_map = {} + accepted_ids_sub_map = {} + sampled_tokens_sub_map = {} + for block_id, block in enumerate(seq.diffusion_blocks): + if not block.is_active or sum(block.local_mask_tokens) == 0: + continue + + if len(block.global_mask_token_ids) > 0: + mask_token_logits = seq_logits[block.global_mask_token_ids, ...] + confidence, sampled_tokens, initial_confidence = self.sample_tokens( + mask_token_logits, + temperature, + top_p=top_p, + top_k=top_k, + neg_entropy=(neg_entropy == "neg_entropy"), + margin_confidence=(margin_confidence == "margin_confidence") + ) + + if block.pre_block_complete: + high_conf_indices = torch.where(initial_confidence > block.accept_threshold)[0] + if len(high_conf_indices) == 0: + number_transfer_tokens = 1 + _, transfer_index = torch.topk(confidence, number_transfer_tokens) + else: + transfer_index = torch.tensor([], device=sampled_tokens.device, dtype=torch.long) + accepted_ids = torch.unique(torch.cat([transfer_index, high_conf_indices])) + else: + high_conf_indices = torch.where(initial_confidence > block.accept_threshold)[0] + accepted_ids = high_conf_indices + + true_local_ids_sub_map[str(block_id)] = [block.local_mask_token_ids[accepted_id] for accepted_id in accepted_ids.tolist()] + accepted_ids_sub_map[str(block_id)] = accepted_ids.tolist() + sampled_tokens_sub_map[str(block_id)] = sampled_tokens + + seq_idx = str(seq.seq_id) + true_local_ids_map[seq_idx] = true_local_ids_sub_map + accepted_ids_map[seq_idx] = accepted_ids_sub_map + sampled_tokens_map[seq_idx] = sampled_tokens_sub_map + + return LLaDASampleOutputForDiffusionLM( + true_local_ids_map=true_local_ids_map, + accepted_ids_map=accepted_ids_map, + sampled_tokens_map=sampled_tokens_map + ) + diff --git a/diffulex/strategy/__init__.py b/diffulex/strategy/__init__.py index e19f44c4..34e7614f 100644 --- a/diffulex/strategy/__init__.py +++ b/diffulex/strategy/__init__.py @@ -1,10 +1,28 @@ """Diffulex strategy package that imports built-in strategies to trigger registration.""" from __future__ import annotations +import importlib +from pathlib import Path # Import built-in strategies so their registrations run at import time. -from . import d2f # noqa: F401 +# Automatically import all subdirectory packages in the current directory +_excluded_dirs = {"__pycache__", "__init__"} +_strategy_modules = [] -__all__ = ["d2f"] +_current_dir = Path(__file__).parent +for item in _current_dir.iterdir(): + if item.is_dir() and not item.name.startswith("_") and item.name not in _excluded_dirs: + # Check if it's a Python package (has __init__.py) + init_file = item / "__init__.py" + if init_file.exists(): + try: + importlib.import_module(f".{item.name}", __name__) + _strategy_modules.append(item.name) + except Exception as e: + # Skip packages that fail to import + import warnings + warnings.warn(f"Failed to import strategy {item.name}: {e}", ImportWarning) + +__all__ = _strategy_modules.copy() DECODING_STRATEGY = None diff --git a/diffulex/strategy/block_diffusion/__init__.py b/diffulex/strategy/block_diffusion/__init__.py index 8dabc025..845afa2a 100644 --- a/diffulex/strategy/block_diffusion/__init__.py +++ b/diffulex/strategy/block_diffusion/__init__.py @@ -1,14 +1,14 @@ """Block Diffusion strategy component exports.""" from __future__ import annotations -from .engine.kvcache_manager import BlockDiffusionKVCacheManager -from .engine.model_runner import BlockDiffusionModelRunner -from .engine.scheduler import BlockDiffusionScheduler -from .engine.sequence import BlockDiffusionSequence +from .engine.kvcache_manager import BDKVCacheManager +from .engine.model_runner import BDModelRunner +from .engine.scheduler import BDScheduler +from .engine.sequence import BDSequence __all__ = [ - "BlockDiffusionKVCacheManager", - "BlockDiffusionModelRunner", - "BlockDiffusionScheduler", - "BlockDiffusionSequence", + "BDKVCacheManager", + "BDModelRunner", + "BDScheduler", + "BDSequence", ] diff --git a/diffulex/strategy/block_diffusion/attention/metadata.py b/diffulex/strategy/block_diffusion/attention/metadata.py index a8396b44..d9832b9d 100644 --- a/diffulex/strategy/block_diffusion/attention/metadata.py +++ b/diffulex/strategy/block_diffusion/attention/metadata.py @@ -1,28 +1,62 @@ import torch +from typing import List from dataclasses import dataclass from diffulex.attention.metadata import AttnMetaDataBase +from diffulex.strategy.block_diffusion.engine.sequence import BDSequence @dataclass -class BlockDiffusionAttnMetaData(AttnMetaDataBase): - seq_lens: list[int] = None - seq_lens_ts: torch.Tensor | None = None - block_diffusion_pp: bool = False - block_mask: list[torch.Tensor] | None = None +class BDAttnMetaData(AttnMetaDataBase): + seqs: List[BDSequence] = None + kv_cache_layout: str = "unified" + need_kv_cache_store: bool = True - -BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() - -def fetch_block_diffusion_attn_metadata() -> BlockDiffusionAttnMetaData: - return BLOCK_DIFFUSION_ATTN_METADATA - -def set_block_diffusion_attn_metadata() -> None: - # TODO - global BLOCK_DIFFUSION_ATTN_METADATA - BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() - -def reset_block_diffusion_attn_metadata() -> None: - global BLOCK_DIFFUSION_ATTN_METADATA - BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() \ No newline at end of file + def __post_init__(self): + if self.context_lens is not None and sum(self.context_lens) > 0: + self.total_lens = self.diffusion_block_size + self.context_lens + + +BD_ATTN_METADATA = BDAttnMetaData() + +def fetch_bd_attn_metadata() -> BDAttnMetaData: + return BD_ATTN_METADATA + +def set_bd_attn_metadata( + is_prefill: bool = False, + 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, + context_lens: torch.Tensor | None = None, + block_tables: torch.Tensor | None = None, + page_block_size: int = 32, + diffusion_block_size: int = 32, + decode_mode: str = "static", + attn_type: str = "full_attention", + kv_cache_layout: str = "unified", + need_kv_cache_store: bool = True, +) -> None: + global BD_ATTN_METADATA + BD_ATTN_METADATA = BDAttnMetaData( + is_prefill=is_prefill, + 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, + context_lens=context_lens, + block_tables=block_tables, + page_block_size=page_block_size, + diffusion_block_size=diffusion_block_size, + kv_cache_layout=kv_cache_layout, + need_kv_cache_store=need_kv_cache_store, + decode_mode=decode_mode, + attn_type=attn_type, + ) + +def reset_bd_attn_metadata() -> None: + global BD_ATTN_METADATA + BD_ATTN_METADATA = BDAttnMetaData() \ No newline at end of file diff --git a/diffulex/strategy/block_diffusion/engine/kvcache_manager.py b/diffulex/strategy/block_diffusion/engine/kvcache_manager.py index 40413949..9659c10b 100644 --- a/diffulex/strategy/block_diffusion/engine/kvcache_manager.py +++ b/diffulex/strategy/block_diffusion/engine/kvcache_manager.py @@ -1,24 +1,23 @@ from __future__ import annotations -from typing import TYPE_CHECKING, list +from typing import TYPE_CHECKING from diffulex.config import Config from diffulex.engine.kvcache_manager import AutoKVCacheManager, KVCacheManagerBase if TYPE_CHECKING: - from .sequence import BlockDiffusionSequence + from .sequence import BDSequence @AutoKVCacheManager.register("block_diffusion", is_default=True) -class BlockDiffusionKVCacheManager(KVCacheManagerBase): +class BDKVCacheManager(KVCacheManagerBase): def __init__(self, config: Config): super().__init__(config) - def can_append(self, seq: "BlockDiffusionSequence") -> bool: - required = 1 if seq.cached_or_caching_num_tokens % self.block_size == 1 else 0 - return len(self.free_block_ids) >= required + def can_append(self, seq: "BDSequence") -> bool: + return len(self.free_block_ids) >= (seq.cached_or_caching_num_tokens % self.block_size == 1) - def may_append(self, seq: "BlockDiffusionSequence") -> None: + def may_append(self, seq: "BDSequence") -> None: if seq.cached_or_caching_num_tokens == 0: return block_table = seq.block_table @@ -37,4 +36,4 @@ def may_append(self, seq: "BlockDiffusionSequence") -> None: self.hash_to_block_id[h] = last_block.block_id block_id = self.free_block_ids[0] self._allocate_block(block_id) - block_table.append(block_id) + block_table.append(block_id) \ No newline at end of file diff --git a/diffulex/strategy/block_diffusion/engine/model_runner.py b/diffulex/strategy/block_diffusion/engine/model_runner.py index 2ff0d8c9..d363ba4a 100644 --- a/diffulex/strategy/block_diffusion/engine/model_runner.py +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -1,32 +1,32 @@ from __future__ import annotations import time -from typing import list + from multiprocessing.synchronize import Event import torch from diffulex.config import Config from diffulex.engine.sequence import SequenceBase -from diffulex.strategy.block_diffusion.engine.sequence import BlockDiffusionSequence -from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata +from diffulex.strategy.block_diffusion.engine.sequence import BDSequence +from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata, set_warming_up, reset_warming_up from diffulex.engine.model_runner import AutoModelRunner, ModelRunnerBase -from diffulex.strategy.block_diffusion.attention.metadata import fetch_block_diffusion_attn_metadata, set_block_diffusion_attn_metadata, reset_block_diffusion_attn_metadata +from diffulex.strategy.block_diffusion.attention.metadata import fetch_bd_attn_metadata, set_bd_attn_metadata, reset_bd_attn_metadata @AutoModelRunner.register("block_diffusion", is_default=True) -class BlockDiffusionModelRunner(ModelRunnerBase): +class BDModelRunner(ModelRunnerBase): """Reference implementation of Block Diffusion decoding strategy.""" - def __init__(self, config: Config, rank: int, event: Event | list[Event]): - super().__init__(config, rank, event) + set_fetch_fn_for_attn_metadata(fetch_bd_attn_metadata) self.diffusion_block_size = config.diffusion_block_size self.mask_token_id = config.mask_token_id - self.decoding_strategy = config.decoding_strategy - set_fetch_fn_for_attn_metadata(fetch_block_diffusion_attn_metadata) - + + super().__init__(config, rank, event) + def warmup_model(self): print("Warming up model...") + set_warming_up(True) torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() max_num_batched_tokens, max_model_len = ( @@ -35,123 +35,14 @@ def warmup_model(self): ) num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs) test_input_ids = [0] * max_model_len - seqs = [BlockDiffusionSequence(test_input_ids, config=self.config) for _ in range(num_seqs)] + seqs = [BDSequence(test_input_ids, config=self.config) for _ in range(num_seqs)] self.run(seqs, True) for seq in seqs: seq.post_process() torch.cuda.empty_cache() + reset_warming_up() - def allocate_kv_cache(self): - config = self.config - hf_config = config.hf_config - free, total = torch.cuda.mem_get_info() - used = total - free - peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"] - current = torch.cuda.memory_stats()["allocated_bytes.all.current"] - num_kv_heads = getattr( - hf_config, - "num_key_value_heads", - getattr(hf_config, "n_kv_heads", None), - ) // self.world_size - - if hasattr(hf_config, "head_dim"): - head_dim = hf_config.head_dim - elif hasattr(hf_config, "hidden_size") and hasattr(hf_config, "num_attention_heads"): - head_dim = hf_config.hidden_size // hf_config.num_attention_heads - else: - raise AttributeError(f"Cannot determine head_dim from config: {type(hf_config)}") - - dtype = ( - hf_config.torch_dtype - if hasattr(hf_config, "torch_dtype") and hf_config.torch_dtype - else torch.bfloat16 - ) - block_bytes = ( - 2 - * hf_config.num_hidden_layers - * self.block_size - * num_kv_heads - * head_dim - * dtype.itemsize - ) - get_num_kvcache_blocks = ( - lambda gpu_memory_utilization: int(total * gpu_memory_utilization - used - peak + current) - // block_bytes - ) - try: - num_kvcache_blocks = get_num_kvcache_blocks(config.gpu_memory_utilization) - assert num_kvcache_blocks > 0 - except Exception: - gpu_memory_utilization = config.gpu_memory_utilization - while num_kvcache_blocks <= 200: - print( - "Warning: GPU memory utilization " - f"{gpu_memory_utilization} is too low to allocate kv cache. " - "Automatically adding 0.05." - ) - gpu_memory_utilization += 0.05 - num_kvcache_blocks = get_num_kvcache_blocks(gpu_memory_utilization) - print( - f"Set gpu_memory_utilization to {gpu_memory_utilization:.2f} " - "to allocate kv cache." - ) - config.gpu_memory_utilization = gpu_memory_utilization - - config.num_kvcache_blocks = num_kvcache_blocks - print( - "Allocated {num_blocks} blocks of size {block_size} for kv cache on rank {rank}.".format( - num_blocks=config.num_kvcache_blocks, - block_size=self.block_size, - rank=self.rank, - ) - ) - - if config.kv_cache_layout == "distinct": - x = config.k_cache_hdim_split_factor_x - self.k_cache = torch.zeros( - hf_config.num_hidden_layers, - config.num_kvcache_blocks, - num_kv_heads, - head_dim // x, - self.block_size, - x, - ) - self.v_cache = torch.zeros( - hf_config.num_hidden_layers, - config.num_kvcache_blocks, - num_kv_heads, - head_dim, - self.block_size, - ) - layer_id = 0 - for module in self.model.modules(): - if hasattr(module, "k_cache") and hasattr(module, "v_cache"): - module.k_cache = self.k_cache[layer_id] - module.v_cache = self.v_cache[layer_id] - layer_id += 1 - elif config.kv_cache_layout == "unified": - self.kv_cache = torch.zeros( - 2, - hf_config.num_hidden_layers, - config.num_kvcache_blocks, - self.block_size, - num_kv_heads, - head_dim, - ) - layer_id = 0 - for module in self.model.modules(): - if hasattr(module, "k_cache") and hasattr(module, "v_cache"): - module.k_cache = self.kv_cache[0, layer_id] - module.v_cache = self.kv_cache[1, layer_id] - layer_id += 1 - else: - raise ValueError( - "Unsupported kv_cache_layout: {layout}. Supported values are 'distinct' and 'unified'.".format( - layout=config.kv_cache_layout - ) - ) - - def prepare_prefill(self, seqs: list[BlockDiffusionSequence]): + def prepare_prefill(self, seqs: list[BDSequence]): input_ids: list[int] = [] positions: list[int] = [] cu_seqlens_q = [0] @@ -161,22 +52,14 @@ def prepare_prefill(self, seqs: list[BlockDiffusionSequence]): slot_mapping: list[int] = [] block_tables = None context_lens: list[int] = [] - seq_lens: list[int] = [] for seq in seqs: - seq.next_diffusion_step(is_prefill=True) + seq.init_diffusion_blocks() total_seqlen = len(seq) input_ids.extend(seq[seq.cached_num_tokens:]) positions.extend(range(seq.cached_num_tokens, total_seqlen)) - seq_lens.append(total_seqlen) context_lens.append(0) - assert len(input_ids) == len(positions), ( - "prepare_prefill(diffusion): len(input_ids) {len_ids} != len(positions) {len_pos}".format( - len_ids=len(input_ids), - len_pos=len(positions), - ) - ) seqlen_q = total_seqlen - seq.cached_num_tokens seqlen_k = total_seqlen @@ -188,42 +71,31 @@ def prepare_prefill(self, seqs: list[BlockDiffusionSequence]): if not seq.block_table: continue - for i in range(0, seq.num_prompt_blocks): + has_padding_mask = seq.pad_prefix_len > 0 + for i in range(0, seq.num_prefix_blocks): if seq.block_cache_missed[i]: - start = seq.block_table[i] * self.block_size - if i != seq.num_prompt_blocks - 1: - end = start + self.block_size + if has_padding_mask and i == seq.num_prefix_blocks - 1: + slot_mapping.extend([-1] * self.block_size) else: - end = start + seq.last_block_prompt_num_tokens - slot_mapping.extend(range(start, end)) + start = seq.block_table[i] * self.block_size + if i != seq.num_prefix_blocks - 1: + end = start + self.block_size + else: + end = start + seq.prefix_last_block_num_tokens + slot_mapping.extend(range(start, end)) else: slot_mapping.extend([-1] * self.block_size) - slot_mapping.extend([-1] * seq.diffusion_block_size) block_tables = self.prepare_block_tables(seqs) input_ids_tensor = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) positions_tensor = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) - seq_lens_ts = torch.tensor(seq_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) context_lens_tensor = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) cu_seqlens_q_tensor = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) cu_seqlens_k_tensor = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) slot_mapping_tensor = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - assert cu_seqlens_q_tensor[-1].item() == input_ids_tensor.numel(), ( - "prepare_prefill(diffusion): cu_seqlens_q[-1]={cq} != num_tokens={nt}".format( - cq=cu_seqlens_q_tensor[-1].item(), - nt=input_ids_tensor.numel(), - ) - ) - assert cu_seqlens_k_tensor[-1].item() == sum(seq_lens), ( - "prepare_prefill(diffusion): cu_seqlens_k[-1]={ck} != sum(seq_lens)={sl}".format( - ck=cu_seqlens_k_tensor[-1].item(), - sl=sum(seq_lens), - ) - ) - - set_block_diffusion_attn_metadata( + set_bd_attn_metadata( True, cu_seqlens_q=cu_seqlens_q_tensor, cu_seqlens_k=cu_seqlens_k_tensor, @@ -232,121 +104,59 @@ def prepare_prefill(self, seqs: list[BlockDiffusionSequence]): slot_mapping=slot_mapping_tensor, context_lens=context_lens_tensor, block_tables=block_tables, - seqs=seqs, + diffusion_block_size=self.diffusion_block_size, kv_cache_layout=self.config.kv_cache_layout, - seq_lens=seq_lens, - seq_lens_ts=seq_lens_ts, + attn_type="block_attention", + decode_mode="static", ) return input_ids_tensor, positions_tensor - def prepare_decode(self, seqs: list[BlockDiffusionSequence]): + def prepare_decode(self, seqs: list[BDSequence]): input_ids: list[int] = [] positions: list[int] = [] cu_seqlens_q = [0] cu_seqlens_k = [0] slot_mapping: list[int] = [] context_lens: list[int] = [] - seq_lens: list[int] = [] - seq_id_to_queue_id: dict[int, int] = {} need_kv_cache_store = False max_seqlen_q = 0 max_seqlen_k = 0 - for seq_idx_in_queue, seq in enumerate(seqs): - seq_id = seq.seq_id - seq_id_to_queue_id[seq_id] = seq_idx_in_queue + for seq in seqs: seq.next_diffusion_step() + cur_input_ids, cur_positions, cur_context_len = seq.diffusion_decoding_inputs() - seq_lens.append(len(cur_input_ids)) input_ids.extend(cur_input_ids) positions.extend(cur_positions) context_lens.append(cur_context_len) - total_seqlen = len(seq) - seqlen_q = total_seqlen - seq.cached_num_tokens - seqlen_k = total_seqlen + seqlen_q = self.diffusion_block_size + seqlen_k = self.diffusion_block_size max_seqlen_q = max(seqlen_q, max_seqlen_q) max_seqlen_k = max(seqlen_k, max_seqlen_k) cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k) - mem_block_to_diffusion_blocks_map = seq.mem_block_to_diffusion_blocks_map - context_len = context_lens[seq_id_to_queue_id[seq_id]] - for mem_block_idx in range(0, seq.num_blocks): - start_idx = mem_block_idx * seq.block_size - end_idx = start_idx + seq.block_size - cur_map = mem_block_to_diffusion_blocks_map[mem_block_idx] - is_last_block = False - meet_active_block = False - while start_idx < end_idx and not is_last_block and not meet_active_block: - local_start_idx = lambda: start_idx % seq.block_size - diffusion_block = seq.diffusion_blocks[cur_map[local_start_idx()]] - if diffusion_block.block_id == 0 and diffusion_block.cursor != start_idx: - diffusion_block.cursor = start_idx - if cur_map[local_start_idx()] == seq.num_diffusion_blocks - 1: - is_last_block = True - - def get_step(diff_blk, begin_idx): - remaining = diff_blk.remaining_length(begin_idx) - if remaining + local_start_idx() <= seq.block_size: - return remaining - return seq.block_size - local_start_idx() - - if diffusion_block.is_in_cache: - step = get_step(diffusion_block, start_idx) - diffusion_block.cursor += step - start_idx += step - elif diffusion_block.is_to_cache: - step = get_step(diffusion_block, start_idx) - diffusion_block.cursor += step - cur_diffusion_block_start = 0 - cur_diffusion_block_end = step - start_idx += step - mem_block_start = ( - seq.block_table[mem_block_idx] * self.block_size - + context_len % seq.block_size - ) - context_len += step - slot_mapping.extend( - range( - mem_block_start + cur_diffusion_block_start, - mem_block_start + cur_diffusion_block_end, - ) - ) - need_kv_cache_store = True - elif diffusion_block.is_active: - meet_active_block = True - - if meet_active_block: - active = seq.active_blocks - first_active_idx = next((i for i, v in enumerate(active) if v), None) - if first_active_idx is not None: - num_blocks_to_pad = len(active) - first_active_idx - slot_mapping.extend([-1] * (num_blocks_to_pad * seq.diffusion_block_size)) - break - assert len(input_ids) == len(positions), ( - "Input IDs length {len_ids} does not match positions length {len_pos}".format( - len_ids=len(input_ids), - len_pos=len(positions), - ) - ) - assert len(input_ids) == len(slot_mapping), ( - "Input IDs length {len_ids} does not match slot mapping length {len_slot}".format( - len_ids=len(input_ids), - len_slot=len(slot_mapping), - ) - ) - + if seq.diffusion_blocks[-1].is_active: + slot_mapping.extend([-1] * self.diffusion_block_size) + elif seq.diffusion_blocks[-1].is_to_cache: + need_kv_cache_store = True + num_pages_storing = seq.num_page_blocks_in_active_diffusion_block + total_num_pages = len(seq.block_table) + for i in range(0, num_pages_storing): + start = seq.block_table[total_num_pages - num_pages_storing + i] * self.block_size + end = start + self.block_size + slot_mapping.extend(range(start, end)) + input_ids_tensor = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) positions_tensor = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) - seq_lens_ts = torch.tensor(seq_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) cu_seqlens_q_tensor = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) cu_seqlens_k_tensor = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) slot_mapping_tensor = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) context_lens_tensor = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) block_tables = self.prepare_block_tables(seqs) - set_block_diffusion_attn_metadata( + set_bd_attn_metadata( False, slot_mapping=slot_mapping_tensor, context_lens=context_lens_tensor, @@ -355,12 +165,10 @@ def get_step(diff_blk, begin_idx): max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, block_tables=block_tables, - seqs=seqs, - seq_lens=seq_lens, - seq_lens_ts=seq_lens_ts, + page_block_size=self.config.kvcache_block_size, + diffusion_block_size=self.diffusion_block_size, kv_cache_layout=self.config.kv_cache_layout, need_kv_cache_store=need_kv_cache_store, - block_diffusion_pp=True, ) return input_ids_tensor, positions_tensor @@ -368,50 +176,99 @@ def get_step(diff_blk, begin_idx): def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): if is_prefill or self.enforce_eager or input_ids.size(0) > 512: return self.model.compute_logits(self.model(input_ids, positions)) - bs = input_ids.size(0) - context = fetch_block_diffusion_attn_metadata() - graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] + num_tokens = input_ids.size(0) + context = fetch_bd_attn_metadata() + graph = self.graphs[next(x for x in self.graph_bs if x >= num_tokens)] graph_vars = self.graph_vars for key, value in graph_vars.items(): if key != "outputs": value.zero_() - graph_vars["input_ids"][:bs] = input_ids - graph_vars["positions"][:bs] = positions - graph_vars["slot_mapping"][:bs] = context.slot_mapping - graph_vars["context_lens"][:bs] = context.context_lens - graph_vars["block_tables"][:bs, : context.block_tables.size(1)] = context.block_tables + + num_seqs = len(context.context_lens) + graph_vars["input_ids"][:num_tokens] = input_ids + graph_vars["positions"][:num_tokens] = positions + graph_vars["slot_mapping"][:num_tokens] = context.slot_mapping + graph_vars["context_lens"][:num_seqs] = context.context_lens + graph_vars["cu_seqlens_q"][:num_seqs + 1] = context.cu_seqlens_q + graph_vars["cu_seqlens_k"][:num_seqs + 1] = context.cu_seqlens_k + graph_vars["block_tables"][:num_seqs, : context.block_tables.size(1)] = context.block_tables graph.replay() - return self.model.compute_logits(graph_vars["outputs"][:bs]) - - @torch.inference_mode() - def run_verbose(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: - print("= =" * 20) - print(f"Running {'prefill' if is_prefill else 'decode'} for {len(seqs)} sequences on rank {self.rank}") - start = time.time() - input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs) - temperatures = self.prepare_sample(seqs) if self.rank == 0 else None - print(f"Prepared input in {time.time() - start:.2f} seconds") - start = time.time() - logits = self.run_model(input_ids, positions, is_prefill) - print(f"Ran model in {time.time() - start:.2f} seconds") - start = time.time() - sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None - print(f"Sampled tokens in {time.time() - start:.2f} seconds") - reset_block_diffusion_attn_metadata() - return sample_output + return self.model.compute_logits(graph_vars["outputs"][:num_tokens]) def run(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs) temperatures = self.prepare_sample(seqs) if self.rank == 0 else None logits = self.run_model(input_ids, positions, is_prefill) - sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None - reset_block_diffusion_attn_metadata() + sample_output = self.sampler(seqs, logits, temperatures) if self.rank == 0 else None + reset_bd_attn_metadata() return sample_output @torch.inference_mode() def capture_cudagraph(self): - """ - TODO: Varlen decoding does not support CUDA graph capture yet. - Can be implemented, but requires drastically high overhead. - """ - raise NotImplementedError("CUDA graph capture for DiffusionLM is not implemented yet.") + config = self.config + hf_config = config.hf_config + max_num_seqs = min(self.config.max_num_seqs, 512) + max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size + diffusion_block_size = self.diffusion_block_size + + max_num_tokens = max_num_seqs * diffusion_block_size + + input_ids = torch.zeros(max_num_tokens, dtype=torch.int64) + positions = torch.zeros(max_num_tokens, dtype=torch.int64) + slot_mapping = torch.zeros(max_num_tokens, dtype=torch.int32) + context_lens = torch.zeros(max_num_seqs, dtype=torch.int32) + block_tables = torch.zeros(max_num_seqs, max_num_blocks, dtype=torch.int32) + outputs = torch.zeros(max_num_tokens, hf_config.hidden_size) + + cu_seqlens_q = torch.zeros(max_num_seqs + 1, dtype=torch.int32) + for i in range(max_num_seqs + 1): + cu_seqlens_q[i] = i * diffusion_block_size + + cu_seqlens_k = torch.zeros(max_num_seqs + 1, dtype=torch.int32) + for i in range(max_num_seqs + 1): + cu_seqlens_k[i] = i * config.max_model_len + + self.graph_bs = [] + seq_bs_list = [1, 2, 4, 8] + list(range(16, max_num_seqs + 1, 16)) + for num_seqs in seq_bs_list: + self.graph_bs.append(num_seqs * diffusion_block_size) + self.graphs = {} + self.graph_pool = None + + for num_tokens in reversed(self.graph_bs): + num_seqs = num_tokens // diffusion_block_size + graph = torch.cuda.CUDAGraph() + + set_bd_attn_metadata( + False, + slot_mapping=slot_mapping[:num_tokens], + 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=diffusion_block_size, + max_seqlen_k=config.max_model_len, + block_tables=block_tables[:num_seqs], + diffusion_block_size=diffusion_block_size, + kv_cache_layout=self.config.kv_cache_layout, + need_kv_cache_store=True, + ) + + outputs[:num_tokens] = self.model(input_ids[:num_tokens], positions[:num_tokens]) # warmup + with torch.cuda.graph(graph, self.graph_pool): + outputs[:num_tokens] = self.model(input_ids[:num_tokens], positions[:num_tokens]) # capture + if self.graph_pool is None: + self.graph_pool = graph.pool() + self.graphs[num_tokens] = graph + torch.cuda.synchronize() + reset_bd_attn_metadata() + + self.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, + block_tables=block_tables, + outputs=outputs, + ) diff --git a/diffulex/strategy/block_diffusion/engine/scheduler.py b/diffulex/strategy/block_diffusion/engine/scheduler.py index cc203af6..947b0133 100644 --- a/diffulex/strategy/block_diffusion/engine/scheduler.py +++ b/diffulex/strategy/block_diffusion/engine/scheduler.py @@ -3,12 +3,11 @@ from diffulex.config import Config from diffulex.engine.scheduler import AutoScheduler, SchedulerBase from diffulex.engine.sequence import SequenceBase, SequenceStatus -from .sequence import BlockDiffusionSequence -from diffulex.layer.sampler import SampleOutputForDiffusionLM +from .sequence import BDSequence @AutoScheduler.register("block_diffusion", is_default=True) -class BlockDiffusionScheduler(SchedulerBase): +class BDScheduler(SchedulerBase): def __init__(self, config: Config): super().__init__(config) self.diffusion_block_size = config.diffusion_block_size @@ -16,7 +15,7 @@ def __init__(self, config: Config): def is_finished(self) -> bool: return not self.waiting and not self.running - def add(self, seq: BlockDiffusionSequence) -> None: + def add(self, seq: BDSequence) -> None: self.waiting.append(seq) def schedule(self) -> tuple[list[SequenceBase], bool]: @@ -77,21 +76,21 @@ def schedule(self) -> tuple[list[SequenceBase], bool]: f"can_append={can_append}" ) raise RuntimeError( - "BlockDiffusionScheduler: unable to schedule any sequence in decode; " + "BDScheduler: unable to schedule any sequence in decode; " f"state={diag}; details={' | '.join(details)}" ) self.running.extendleft(reversed(scheduled)) return scheduled, False - def preempt(self, seq: BlockDiffusionSequence) -> None: + def preempt(self, seq: BDSequence) -> None: seq.status = SequenceStatus.WAITING self.block_manager.free(seq) self.waiting.appendleft(seq) def postprocess( self, - seqs: list[BlockDiffusionSequence], - sample_output: SampleOutputForDiffusionLM, + seqs: list[BDSequence], + sample_output, ) -> dict[int, int]: n_diff_steps: dict[int, int] = {} for seq in seqs: diff --git a/diffulex/strategy/block_diffusion/engine/sequence.py b/diffulex/strategy/block_diffusion/engine/sequence.py index 01d1ff9a..936b2425 100644 --- a/diffulex/strategy/block_diffusion/engine/sequence.py +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -1,25 +1,23 @@ from __future__ import annotations -import torch - -from dataclasses import dataclass from enum import Enum, auto +from dataclasses import dataclass from diffulex.config import Config -from diffulex.engine.sequence import AutoSequence, SequenceBase from diffulex.sampling_params import SamplingParams +from diffulex.engine.sequence import AutoSequence, SequenceBase -class BlockDiffusionBlockStatus(Enum): +class BDDiffusionBlockStatus(Enum): ACTIVE = auto() TO_CACHE = auto() IN_CACHE = auto() @dataclass -class BlockDiffusionBlock: +class BDDiffusionBlock: block_id: int = 0 - status: BlockDiffusionBlockStatus = BlockDiffusionBlockStatus.ACTIVE + status: BDDiffusionBlockStatus = BDDiffusionBlockStatus.ACTIVE global_start_id: int = 0 global_end_id: int | None = None @@ -28,14 +26,8 @@ class BlockDiffusionBlock: mask_token_id: int = 151666 size: int = 32 is_prompt: bool = False - - accept_threshold: float = 0.95 - add_new_block_threshold: float = 0.1 - complete_threshold: float = 0.9 - - seq: "BlockDiffusionSequence" | None = None - pre_block: "BlockDiffusionBlock" | None = None - suf_block: "BlockDiffusionBlock" | None = None + + seq: "BDSequence" | None = None def __post_init__(self) -> None: self.global_end_id = self.global_start_id + self.size @@ -45,84 +37,73 @@ def __getitem__(self, key: int) -> int: def __len__(self) -> int: return self.size - + + def to_cache(self) -> None: + if self.available_to_cache and not self.is_in_cache: + self.status = BDDiffusionBlockStatus.TO_CACHE + + def in_cache(self) -> None: + if self.is_to_cache: + self.status = BDDiffusionBlockStatus.IN_CACHE + + def modify_token(self, local_token_id: int, modified_to: int) -> None: + if self.seq is None: + raise RuntimeError("Diffusion block is not attached to a sequence.") + target_id = local_token_id + self.global_start_id + assert self.seq.token_ids[target_id] == self.mask_token_id + self.seq.token_ids[target_id] = modified_to.item() # type: ignore[assignment] + self.seq.new_tokens += 1 + @property - def current_complete_ratio(self) -> float: - if self.size == 0: - return 0.0 - return sum(token_id != self.mask_token_id for token_id in self.token_ids) / self.size - + def token_ids(self) -> list[int]: + return self.seq.token_ids[self.global_start_id: self.global_end_id] + @property - def available_to_cache(self) -> bool: - return self.current_complete_ratio == 1.0 - + def has_mask_token(self) -> bool: + return any(token == self.mask_token_id for token in self.token_ids) + @property def is_active(self) -> bool: - return self.status == BlockDiffusionBlockStatus.ACTIVE - - @property - def is_in_cache(self) -> bool: - return self.status == BlockDiffusionBlockStatus.IN_CACHE - + return self.status == BDDiffusionBlockStatus.ACTIVE + @property def is_to_cache(self) -> bool: - return self.status == BlockDiffusionBlockStatus.TO_CACHE - + return self.status == BDDiffusionBlockStatus.TO_CACHE + @property - def pre_block_complete(self) -> bool: - if self.pre_block is None: - return True - return self.pre_block.current_complete_ratio >= self.complete_threshold - + def is_in_cache(self) -> bool: + return self.status == BDDiffusionBlockStatus.IN_CACHE + @property - def add_new_block(self) -> bool: - return self.current_complete_ratio >= self.add_new_block_threshold - + def available_to_cache(self) -> bool: + return not self.has_mask_token and self.is_active + @property - def token_ids(self) -> list[int]: - if self.seq is None: - raise RuntimeError("Diffusion block is not attached to a sequence.") - return self.seq.token_ids[self.global_start_id : self.global_end_id] - + def available_in_cache(self) -> bool: + return self.is_to_cache + + @property + def available_to_add_new_block(self) -> bool: + return self.is_in_cache + @property def local_mask_tokens(self) -> list[bool]: - return [token_id == self.seq.mask_token_id for token_id in self.token_ids] # type: ignore[arg-type] - + return [token_id == self.mask_token_id for token_id in self.token_ids] + @property def local_mask_token_ids(self) -> list[int]: return [idx for idx, is_mask in enumerate(self.local_mask_tokens) if is_mask] - + @property def global_mask_token_ids(self) -> list[int]: if self.seq is None: return [] - offset = self.global_start_id - in_cache_blocks = list(range(sum(self.seq.in_cache_blocks))) - offset -= sum(self.seq.diffusion_blocks[block_id].size for block_id in in_cache_blocks) + offset = self.global_start_id - self.size * sum(block.is_to_cache for block in self.seq.diffusion_blocks) return [mask_id + offset for mask_id in self.local_mask_token_ids] - - def remaining_length(self, start_idx: int) -> int: - return self.size - self.cursor - - def to_cache(self) -> None: - if self.available_to_cache and not self.is_in_cache: - self.status = BlockDiffusionBlockStatus.TO_CACHE - - def in_cache(self) -> None: - if self.is_to_cache: - self.status = BlockDiffusionBlockStatus.IN_CACHE - - def modify_token(self, local_token_id: int, modified_to: int) -> None: - if self.seq is None: - raise RuntimeError("Diffusion block is not attached to a sequence.") - target_id = local_token_id + self.global_start_id - assert self.seq.token_ids[target_id] == self.mask_token_id - self.seq.token_ids[target_id] = modified_to.item() # type: ignore[assignment] - self.seq.new_tokens += 1 - + @AutoSequence.register("block_diffusion", is_default=True) -class BlockDiffusionSequence(SequenceBase): +class BDSequence(SequenceBase): """Sequence implementation tailored for diffusion-based decoding.""" def __init__( @@ -133,253 +114,147 @@ def __init__( ): super().__init__(token_ids, sampling_params) if config is None: - raise ValueError("SequenceForDiffusionLM requires a Config instance.") + raise ValueError("BDSequence requires a Config instance.") + self.config = config - self.decoding_strategy = config.decoding_strategy - self.kv_cache_layout = config.kv_cache_layout - self.eos_token_id = config.eos - self.max_model_len = config.max_model_len - self.mask_token_id = config.mask_token_id + self.diffusion_blocks: list[BDDiffusionBlock] = [] self.diffusion_block_size = config.diffusion_block_size - self.block_mask: torch.Tensor | None = None - self.meet_eos = False - self.diffusion_blocks: list[BlockDiffusionBlock] = [] + self.mask_token_id = config.mask_token_id self.n_steps = 0 - self.input_token_ids: list[int] = [] - self.input_num_tokens = 0 - self.input_num_prompt_tokens = 0 - - def __repr__(self) -> str: - return ( - "SequenceForDiffusionLM(seq_id={seq_id}, status={status}, num_tokens={num_tokens}, " - "num_prompt_tokens={num_prompt_tokens}, num_cached_tokens={num_cached_tokens}, " - "diffusion_block_size={diffusion_block_size}, mask_shape={mask_shape})" - ).format( - seq_id=self.seq_id, - status=self.status.name, - num_tokens=self.num_tokens, - num_prompt_tokens=self.num_prompt_tokens, - num_cached_tokens=self.num_cached_tokens, - diffusion_block_size=self.diffusion_block_size, - mask_shape=self.block_mask.shape if self.block_mask is not None else None, - ) - - def __getstate__(self): - diffusion_blocks_state = [] - for block in self.diffusion_blocks: - diffusion_blocks_state.append( - { - "block_id": block.block_id, - "status": block.status, - "global_start_id": block.global_start_id, - "global_end_id": block.global_end_id, - "cursor": block.cursor, - "mask_token_id": block.mask_token_id, - "size": block.size, - "is_prompt": block.is_prompt, - "accept_threshold": block.accept_threshold, - "add_new_block_threshold": block.add_new_block_threshold, - "complete_threshold": block.complete_threshold, - } - ) - - state = { - "seq_id": self.seq_id, - "status": self.status, - "token_ids": self.token_ids, - "last_token": self.last_token, - "num_tokens": self.num_tokens, - "num_prompt_tokens": self.num_prompt_tokens, - "num_cached_tokens": self.num_cached_tokens, - "block_table": self.block_table, - "block_cache_missed": self.block_cache_missed, - "temperature": self.temperature, - "max_tokens": self.max_tokens, - "ignore_eos": self.ignore_eos, - "config": self.config, - "decoding_strategy": self.decoding_strategy, - "kv_cache_layout": self.kv_cache_layout, - "eos_token_id": self.eos_token_id, - "max_model_len": self.max_model_len, - "mask_token_id": self.mask_token_id, - "diffusion_block_size": self.diffusion_block_size, - "diffusion_blocks_state": diffusion_blocks_state, - "input_token_ids": self.input_token_ids, - "input_num_tokens": self.input_num_tokens, - "input_num_prompt_tokens": self.input_num_prompt_tokens, - "new_tokens": self.new_tokens, - "block_mask": self.block_mask, - "meet_eos": self.meet_eos, - "n_steps": self.n_steps, - } - return state - - def __setstate__(self, state): - self.seq_id = state["seq_id"] - self.status = state["status"] - self.token_ids = state["token_ids"] - self.last_token = state["last_token"] - self.num_tokens = state["num_tokens"] - self.num_prompt_tokens = state["num_prompt_tokens"] - self.num_cached_tokens = state["num_cached_tokens"] - self.block_table = state["block_table"] - self.block_cache_missed = state["block_cache_missed"] - self.temperature = state["temperature"] - self.max_tokens = state["max_tokens"] - self.ignore_eos = state["ignore_eos"] - self.meet_eos = state["meet_eos"] - - self.config = state["config"] - self.decoding_strategy = state.get("decoding_strategy", getattr(self.config, "decoding_strategy", None)) - self.kv_cache_layout = state.get("kv_cache_layout", getattr(self.config, "kv_cache_layout", None)) - self.eos_token_id = state["eos_token_id"] - self.max_model_len = state["max_model_len"] - self.mask_token_id = state["mask_token_id"] - self.diffusion_block_size = state["diffusion_block_size"] - - self.input_token_ids = state.get("input_token_ids", []) - self.input_num_tokens = state.get("input_num_tokens", 0) - self.input_num_prompt_tokens = state.get("input_num_prompt_tokens", 0) - self.new_tokens = state.get("new_tokens", 0) - self.block_mask = state.get("block_mask") - self.n_steps = state.get("n_steps", 0) - - if self.block_mask is not None and self.block_mask.device.index != torch.cuda.current_device(): - self.block_mask = self.block_mask.to(torch.cuda.current_device()) - - self.diffusion_blocks = [] - pre_block = None - for block_state in state["diffusion_blocks_state"]: - block = BlockDiffusionBlock( - block_id=block_state["block_id"], - status=block_state["status"], - global_start_id=block_state["global_start_id"], - global_end_id=block_state["global_end_id"], - cursor=block_state.get("cursor", 0), - mask_token_id=block_state["mask_token_id"], - size=block_state["size"], - is_prompt=block_state["is_prompt"], - accept_threshold=block_state.get("accept_threshold", 0.95), - add_new_block_threshold=block_state.get("add_new_block_threshold", 0.1), - complete_threshold=block_state.get("complete_threshold", 0.9), - seq=self, - pre_block=pre_block, - ) - if pre_block is not None: - pre_block.suf_block = block - self.diffusion_blocks.append(block) - pre_block = block - - @property - def num_completion_tokens(self) -> int: - return self.num_tokens - self.input_num_tokens - + @property def completion_token_ids(self) -> list[int]: - return self.token_ids[self.input_num_prompt_tokens :] - + return self.token_ids[self.prefix_len : ] + @property - def active_blocks(self) -> list[bool]: - return [block.is_active for block in self.diffusion_blocks] - + def prefix_len_with_padding(self) -> int: + return self.prefix_len + self.pad_prefix_len + @property - def to_cache_blocks(self) -> list[bool]: - return [block.is_to_cache for block in self.diffusion_blocks] - + def diffusion_block_status(self) -> list[BDDiffusionBlockStatus]: + return [block.status for block in self.diffusion_blocks] + @property - def in_cache_blocks(self) -> list[bool]: - return [block.is_in_cache for block in self.diffusion_blocks] - + def num_prefix_blocks(self) -> int: + return (self.prefix_len + self.block_size - 1) // self.block_size + @property - def num_prompt_blocks(self) -> int: - return (self.input_num_prompt_tokens + self.block_size - 1) // self.block_size - + def prefix_last_block_num_tokens(self) -> int: + return self.prefix_len - (self.num_prefix_blocks - 1) * self.block_size + @property - def last_block_prompt_num_tokens(self) -> int: - return self.input_num_prompt_tokens - (self.num_prompt_blocks - 1) * self.block_size - + def active_block_token_ids(self) -> list[int]: + return self.diffusion_blocks[-1].token_ids + @property - def updated_or_updating_kv_cache_block_ids(self) -> list[int]: - return [idx for idx, caching in enumerate(self.caching_blocks) if caching] - + def num_page_blocks_in_active_diffusion_block(self) -> int: + return self.diffusion_block_size // self.block_size + @property - def caching_blocks(self) -> list[bool]: - return [to_cache or in_cache for to_cache, in_cache in zip(self.to_cache_blocks, self.in_cache_blocks)] - - @property - def cached_block_ids(self) -> list[int]: - return [idx for idx, in_cache in enumerate(self.in_cache_blocks) if in_cache] - - @property - def mask_tokens(self) -> list[bool]: - return [token_id == self.mask_token_id for token_id in self.token_ids] - + def cached_num_tokens(self) -> int: + return sum(block.size for block in self.diffusion_blocks if block.is_in_cache) + @property def caching_num_tokens(self) -> int: return sum(block.size for block in self.diffusion_blocks if block.is_to_cache) - + @property def cached_or_caching_last_token_id(self) -> int: - cached_num_tokens = 0 - for block_id in self.updated_or_updating_kv_cache_block_ids: - block = self.diffusion_blocks[block_id] - cached_num_tokens += block.size - return max(cached_num_tokens - 1, 0) - + return max(sum(block.size for block in self.diffusion_blocks if block.is_to_cache or block.is_in_cache) - 1, 0) + @property def cached_or_caching_num_tokens(self) -> int: return self.cached_or_caching_last_token_id + 1 - + @property - def cached_num_tokens(self) -> int: - return sum(block.size for block in self.diffusion_blocks if block.is_in_cache) - - @property - def num_cached_blocks(self) -> int: - return (self.num_cached_tokens + self.block_size - 1) // self.block_size - - @property - def diffusion_num_tokens(self) -> int: - return sum(self.mask_tokens) - + def has_to_cache_block(self) -> bool: + return any(block.is_to_cache for block in self.diffusion_blocks) + @property - def mem_block_to_diffusion_blocks_map(self) -> list[list[int]]: - mapping = [] - for block_id in range(self.num_blocks): - window_start = block_id * self.block_size - window_length = self.block_size if block_id < self.num_blocks - 1 else self.last_block_num_tokens - mapping.append( - [self.token_to_diffusion_block_id(token_id) for token_id in range(window_start, window_start + window_length)] - ) - return mapping - - def token_to_diffusion_block_id(self, token_id: int) -> int: - if token_id < self.input_num_tokens: - return 0 - return (token_id - self.input_num_tokens) // self.diffusion_block_size + 1 - + def to_cache_last_token_id(self) -> int: + to_cache_num_tokens = 0 + for block in self.diffusion_blocks: + if block.is_to_cache: + to_cache_num_tokens += block.size + return to_cache_num_tokens - 1 + @property - def num_diffusion_blocks(self) -> int: - return len(self.diffusion_blocks) - - def diffusion_decoding_inputs(self) -> tuple[list[int], list[int], int]: - to_cache_and_active_blocks = self.diffusion_blocks[self.cached_block_ids[-1] + 1 :] - assert len(to_cache_and_active_blocks) == sum(self.active_blocks) + sum(self.to_cache_blocks) - - input_tokens: list[int] = [] - positions: list[int] = [] - context_len = sum(self.diffusion_blocks[block_id].size for block_id in self.cached_block_ids) - temp_context_len = context_len - for block in to_cache_and_active_blocks: - input_tokens.extend(block.token_ids) - positions.extend(range(temp_context_len, temp_context_len + block.size)) - temp_context_len += block.size - - return input_tokens, positions, context_len - + def num_completion_tokens(self) -> int: + return self.num_tokens - self.num_prompt_tokens + def reset_new_tokens(self) -> None: self.new_tokens = 0 - + + def diffusion_decoding_inputs(self) -> tuple[list[int], list[int], int]: + return ( + self.active_block_token_ids, + list(range(self.num_tokens - self.diffusion_block_size, self.num_tokens)), + self.num_tokens - self.diffusion_block_size, + ) + + def extend_mask_tokens(self, extend_len: int) -> None: + self.token_ids.extend([self.mask_token_id] * extend_len) + + def init_diffusion_blocks(self) -> None: + """Initialize diffusion blocks: prefix blocks are TO_CACHE, last block with mask tokens is ACTIVE.""" + self.prefix_len = len(self.token_ids) + block_size = self.diffusion_block_size + + # Calculate prefix blocks and padding + num_prefix_blocks = self.prefix_len // block_size + self.pad_prefix_len = 0 if self.prefix_len % block_size == 0 else block_size - (self.prefix_len % block_size) + + # Add mask tokens for the last prefix block + self.extend_mask_tokens(self.pad_prefix_len) + + # Calculate total blocks needed + total_num_blocks = num_prefix_blocks if self.pad_prefix_len == 0 else num_prefix_blocks + 1 + + # Create all blocks + current_pos = 0 + for block_id in range(total_num_blocks): + # Determine block status + block_tokens = self.token_ids[current_pos:current_pos + block_size] + has_mask_token = any(token == self.mask_token_id for token in block_tokens) + is_last_prefix_block = (block_id == num_prefix_blocks) + + if block_id < num_prefix_blocks: + status = BDDiffusionBlockStatus.TO_CACHE + elif is_last_prefix_block: + status = BDDiffusionBlockStatus.ACTIVE if has_mask_token else BDDiffusionBlockStatus.TO_CACHE + else: + status = BDDiffusionBlockStatus.TO_CACHE + + block = BDDiffusionBlock( + block_id=block_id, + status=status, + global_start_id=current_pos, + size=block_size, + mask_token_id=self.mask_token_id, + is_prompt=(block_id <= num_prefix_blocks), + seq=self, + ) + self.diffusion_blocks.append(block) + current_pos += block_size + self.n_steps += 1 + + def next_diffusion_step(self) -> None: + """Append new diffusion block if needed.""" + if self.diffusion_blocks[-1].available_to_add_new_block: + self.extend_mask_tokens(self.diffusion_block_size) + self.diffusion_blocks.append( + BDDiffusionBlock( + block_id=len(self.diffusion_blocks), + status=BDDiffusionBlockStatus.ACTIVE, + global_start_id=self.num_tokens - self.diffusion_block_size, + size=self.diffusion_block_size, + mask_token_id=self.mask_token_id, + is_prompt=False, + seq=self, + ) + ) + self.n_steps += 1 + def post_process(self) -> None: for block in self.diffusion_blocks: block.cursor = 0 @@ -391,90 +266,4 @@ def post_process(self) -> None: if block.available_to_cache: block.to_cache() else: - break - - def set_layout(self, layout: str) -> None: - self.kv_cache_layout = layout - - @property - def current_block_mask(self) -> torch.Tensor: - if self.block_mask is None: - raise RuntimeError("Block mask not initialized.") - if self.kv_cache_layout == "distinct": - return self.block_mask[..., self.cached_num_tokens :, self.cached_num_tokens :] - return self.block_mask[..., self.cached_num_tokens :, :] - - def update_block_mask(self, is_prefill: bool = False) -> None: - if is_prefill: - num_tokens = self.num_tokens - mask_shape = (1, 1, num_tokens, num_tokens) - block_mask = torch.zeros(mask_shape, dtype=torch.bool, device=torch.cuda.current_device()) - block_mask[..., : self.input_num_tokens, : self.input_num_tokens] = True - num_diffusion_blocks = ( - self.num_tokens - self.input_num_tokens + self.diffusion_block_size - 1 - ) // self.diffusion_block_size - for block_id in range(num_diffusion_blocks): - start_h = self.input_num_tokens + block_id * self.diffusion_block_size - end_h = start_h + self.diffusion_block_size - block_mask[..., start_h:end_h, :end_h] = True - self.block_mask = block_mask.clone() - return - - if self.block_mask is None: - raise RuntimeError("Prefill block mask must be created before decode updates.") - dev = self.block_mask.device - left_shape = (1, 1, self.num_tokens - self.diffusion_block_size, self.diffusion_block_size) - down_shape = (1, 1, self.diffusion_block_size, self.num_tokens) - left_cat_tensor = torch.zeros(left_shape, dtype=torch.bool, device=dev) - down_cat_tensor = torch.ones(down_shape, dtype=torch.bool, device=dev) - self.block_mask = torch.cat([self.block_mask, left_cat_tensor], dim=-1) - self.block_mask = torch.cat([self.block_mask, down_cat_tensor], dim=-2) - - def next_diffusion_step(self, is_prefill: bool = False) -> None: - self.n_steps += 1 - if is_prefill: - self.input_token_ids = self.token_ids.copy() - self.input_num_tokens = self.num_tokens - self.input_num_prompt_tokens = self.num_prompt_tokens - self.num_prompt_tokens += self.diffusion_block_size - self.diffusion_blocks.append( - BlockDiffusionBlock( - block_id=len(self.diffusion_blocks), - status=BlockDiffusionBlockStatus.TO_CACHE, - global_start_id=0, - mask_token_id=self.mask_token_id, - size=len(self.input_token_ids), - accept_threshold=self.config.accept_threshold, - add_new_block_threshold=self.config.add_new_block_threshold, - complete_threshold=self.config.complete_threshold, - is_prompt=True, - seq=self, - ) - ) - - if not self.diffusion_blocks: - return - - if self.diffusion_blocks[-1].add_new_block and not self.meet_eos: - remaining = self.max_model_len - self.num_tokens - if remaining <= 0: - return - added_num_tokens = min(self.diffusion_block_size, remaining) - diffusion_seq = [self.mask_token_id] * added_num_tokens - current_block = BlockDiffusionBlock( - block_id=len(self.diffusion_blocks), - status=BlockDiffusionBlockStatus.ACTIVE, - global_start_id=self.num_tokens, - mask_token_id=self.mask_token_id, - size=added_num_tokens, - accept_threshold=self.config.accept_threshold, - add_new_block_threshold=self.config.add_new_block_threshold, - complete_threshold=self.config.complete_threshold, - seq=self, - pre_block=self.diffusion_blocks[-1], - ) - self.diffusion_blocks[-1].suf_block = current_block - self.token_ids += diffusion_seq - self.num_tokens += added_num_tokens - self.diffusion_blocks.append(current_block) - self.update_block_mask(is_prefill=is_prefill) \ No newline at end of file + break \ No newline at end of file diff --git a/diffulex/strategy/d2f/attention/metadata.py b/diffulex/strategy/d2f/attention/metadata.py index b9d4d3c7..523daf4d 100644 --- a/diffulex/strategy/d2f/attention/metadata.py +++ b/diffulex/strategy/d2f/attention/metadata.py @@ -11,8 +11,6 @@ class D2FAttnMetaData(AttnMetaDataBase): seq_lens: list[int] = None seq_lens_ts: torch.Tensor | None = None - d2f_pp: bool = False - block_mask: torch.Tensor | None = None seqs: List[D2FSequence] = None kv_cache_layout: str = "unified" need_kv_cache_store: bool = True @@ -20,35 +18,6 @@ class D2FAttnMetaData(AttnMetaDataBase): def __post_init__(self): if self.seq_lens_ts is not None and self.context_lens is not None: self.total_lens = self.seq_lens_ts + self.context_lens - if not self.is_prefill and self.d2f_pp: - return - if self.seqs is not None and len(self.seqs) > 0: - if self.is_prefill: - masks = [seq.current_block_mask for seq in self.seqs] - total_len = sum(mask.size(-1) for mask in masks) - self.block_mask = torch.zeros(total_len, total_len, dtype=torch.bool) - - start_idx = 0 - for mask in masks: - seq_len = mask.size(-1) - end_idx = start_idx + seq_len - self.block_mask[start_idx:end_idx, start_idx:end_idx] = mask.clone() - start_idx = end_idx - self.block_mask = self.block_mask.to(mask.device) - else: - masks = [seq.current_block_mask for seq in self.seqs] - total_height = sum(mask.size(-2) for mask in masks) - total_width = sum(mask.size(-1) for mask in masks) - self.block_mask = torch.zeros(total_height, total_width, dtype=torch.bool) - start_row = 0 - start_col = 0 - for mask in masks: - height, width = mask.size(-2), mask.size(-1) - end_row = start_row + height - end_col = start_col + width - self.block_mask[start_row:end_row, start_col:end_col] = mask.clone() - start_row, start_col = end_row, end_col - self.block_mask = self.block_mask.to(mask.device) @property def total_num_seqs(self) -> int: @@ -74,8 +43,9 @@ def set_d2f_attn_metadata( seq_lens_ts: torch.Tensor | None = None, kv_cache_layout: str = "unified", need_kv_cache_store: bool = True, - d2f_pp: bool = False, - block_mask: torch.Tensor | None = None, + diffusion_block_size: int = 32, + decode_mode: str = "varlen", + attn_type: str = "full_attention", ) -> None: global D2F_ATTN_METADATA D2F_ATTN_METADATA = D2FAttnMetaData( @@ -89,11 +59,12 @@ def set_d2f_attn_metadata( block_tables=block_tables, seq_lens=seq_lens, seq_lens_ts=seq_lens_ts, - d2f_pp=d2f_pp, - block_mask=block_mask, seqs=seqs, kv_cache_layout=kv_cache_layout, need_kv_cache_store=need_kv_cache_store, + diffusion_block_size=diffusion_block_size, + decode_mode=decode_mode, + attn_type=attn_type, ) def reset_d2f_attn_metadata() -> None: diff --git a/diffulex/strategy/d2f/engine/kvcache_manager.py b/diffulex/strategy/d2f/engine/kvcache_manager.py index 119a3f0d..f3eeb730 100644 --- a/diffulex/strategy/d2f/engine/kvcache_manager.py +++ b/diffulex/strategy/d2f/engine/kvcache_manager.py @@ -15,8 +15,7 @@ def __init__(self, config: Config): super().__init__(config) def can_append(self, seq: "D2FSequence") -> bool: - required = 1 if seq.cached_or_caching_num_tokens % self.block_size == 1 else 0 - return len(self.free_block_ids) >= required + return len(self.free_block_ids) >= (seq.cached_or_caching_num_tokens % self.block_size == 1) def may_append(self, seq: "D2FSequence") -> None: if seq.cached_or_caching_num_tokens == 0: diff --git a/diffulex/strategy/d2f/engine/model_runner.py b/diffulex/strategy/d2f/engine/model_runner.py index 6d45f7a8..7d736ab6 100644 --- a/diffulex/strategy/d2f/engine/model_runner.py +++ b/diffulex/strategy/d2f/engine/model_runner.py @@ -9,7 +9,7 @@ from diffulex.config import Config from diffulex.engine.sequence import SequenceBase from diffulex.strategy.d2f.engine.sequence import D2FSequence -from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata +from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata, set_warming_up, reset_warming_up from diffulex.engine.model_runner import AutoModelRunner, ModelRunnerBase from diffulex.strategy.d2f.attention.metadata import fetch_d2f_attn_metadata, set_d2f_attn_metadata, reset_d2f_attn_metadata @@ -17,16 +17,17 @@ @AutoModelRunner.register("d2f", is_default=True) class D2FModelRunner(ModelRunnerBase): """Reference implementation of D2F decoding strategy.""" - def __init__(self, config: Config, rank: int, event: Event | list[Event]): - super().__init__(config, rank, event) + set_fetch_fn_for_attn_metadata(fetch_d2f_attn_metadata) + self.diffusion_block_size = config.diffusion_block_size self.mask_token_id = config.mask_token_id - self.decoding_strategy = config.decoding_strategy - set_fetch_fn_for_attn_metadata(fetch_d2f_attn_metadata) + + super().__init__(config, rank, event) def warmup_model(self): print("Warming up model...") + set_warming_up(True) torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() max_num_batched_tokens, max_model_len = ( @@ -40,116 +41,7 @@ def warmup_model(self): for seq in seqs: seq.post_process() torch.cuda.empty_cache() - - def allocate_kv_cache(self): - config = self.config - hf_config = config.hf_config - free, total = torch.cuda.mem_get_info() - used = total - free - peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"] - current = torch.cuda.memory_stats()["allocated_bytes.all.current"] - num_kv_heads = getattr( - hf_config, - "num_key_value_heads", - getattr(hf_config, "n_kv_heads", None), - ) // self.world_size - - if hasattr(hf_config, "head_dim"): - head_dim = hf_config.head_dim - elif hasattr(hf_config, "hidden_size") and hasattr(hf_config, "num_attention_heads"): - head_dim = hf_config.hidden_size // hf_config.num_attention_heads - else: - raise AttributeError(f"Cannot determine head_dim from config: {type(hf_config)}") - - dtype = ( - hf_config.torch_dtype - if hasattr(hf_config, "torch_dtype") and hf_config.torch_dtype - else torch.bfloat16 - ) - block_bytes = ( - 2 - * hf_config.num_hidden_layers - * self.block_size - * num_kv_heads - * head_dim - * dtype.itemsize - ) - get_num_kvcache_blocks = ( - lambda gpu_memory_utilization: int(total * gpu_memory_utilization - used - peak + current) - // block_bytes - ) - try: - num_kvcache_blocks = get_num_kvcache_blocks(config.gpu_memory_utilization) - assert num_kvcache_blocks > 0 - except Exception: - gpu_memory_utilization = config.gpu_memory_utilization - while num_kvcache_blocks <= 200: - print( - "Warning: GPU memory utilization " - f"{gpu_memory_utilization} is too low to allocate kv cache. " - "Automatically adding 0.05." - ) - gpu_memory_utilization += 0.05 - num_kvcache_blocks = get_num_kvcache_blocks(gpu_memory_utilization) - print( - f"Set gpu_memory_utilization to {gpu_memory_utilization:.2f} " - "to allocate kv cache." - ) - config.gpu_memory_utilization = gpu_memory_utilization - - config.num_kvcache_blocks = num_kvcache_blocks - print( - "Allocated {num_blocks} blocks of size {block_size} for kv cache on rank {rank}.".format( - num_blocks=config.num_kvcache_blocks, - block_size=self.block_size, - rank=self.rank, - ) - ) - - if config.kv_cache_layout == "distinct": - x = config.k_cache_hdim_split_factor_x - self.k_cache = torch.zeros( - hf_config.num_hidden_layers, - config.num_kvcache_blocks, - num_kv_heads, - head_dim // x, - self.block_size, - x, - ) - self.v_cache = torch.zeros( - hf_config.num_hidden_layers, - config.num_kvcache_blocks, - num_kv_heads, - head_dim, - self.block_size, - ) - layer_id = 0 - for module in self.model.modules(): - if hasattr(module, "k_cache") and hasattr(module, "v_cache"): - module.k_cache = self.k_cache[layer_id] - module.v_cache = self.v_cache[layer_id] - layer_id += 1 - elif config.kv_cache_layout == "unified": - self.kv_cache = torch.zeros( - 2, - hf_config.num_hidden_layers, - config.num_kvcache_blocks, - self.block_size, - num_kv_heads, - head_dim, - ) - layer_id = 0 - for module in self.model.modules(): - if hasattr(module, "k_cache") and hasattr(module, "v_cache"): - module.k_cache = self.kv_cache[0, layer_id] - module.v_cache = self.kv_cache[1, layer_id] - layer_id += 1 - else: - raise ValueError( - "Unsupported kv_cache_layout: {layout}. Supported values are 'distinct' and 'unified'.".format( - layout=config.kv_cache_layout - ) - ) + reset_warming_up() def prepare_prefill(self, seqs: list[D2FSequence]): input_ids: list[int] = [] @@ -236,6 +128,9 @@ def prepare_prefill(self, seqs: list[D2FSequence]): kv_cache_layout=self.config.kv_cache_layout, seq_lens=seq_lens, seq_lens_ts=seq_lens_ts, + diffusion_block_size=self.diffusion_block_size, + decode_mode="varlen", + attn_type="full_attention", ) return input_ids_tensor, positions_tensor @@ -360,7 +255,9 @@ def get_step(diff_blk, begin_idx): seq_lens_ts=seq_lens_ts, kv_cache_layout=self.config.kv_cache_layout, need_kv_cache_store=need_kv_cache_store, - d2f_pp=True, + diffusion_block_size=self.diffusion_block_size, + decode_mode="varlen", + attn_type="full_attention", ) return input_ids_tensor, positions_tensor @@ -383,23 +280,6 @@ def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill graph.replay() return self.model.compute_logits(graph_vars["outputs"][:bs]) - @torch.inference_mode() - def run_verbose(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: - print("= =" * 20) - print(f"Running {'prefill' if is_prefill else 'decode'} for {len(seqs)} sequences on rank {self.rank}") - start = time.time() - input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs) - temperatures = self.prepare_sample(seqs) if self.rank == 0 else None - print(f"Prepared input in {time.time() - start:.2f} seconds") - start = time.time() - logits = self.run_model(input_ids, positions, is_prefill) - print(f"Ran model in {time.time() - start:.2f} seconds") - start = time.time() - sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None - print(f"Sampled tokens in {time.time() - start:.2f} seconds") - reset_d2f_attn_metadata() - return sample_output - def run(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs) temperatures = self.prepare_sample(seqs) if self.rank == 0 else None diff --git a/diffulex/strategy/d2f/engine/scheduler.py b/diffulex/strategy/d2f/engine/scheduler.py index 335b54d1..a4b8f29d 100644 --- a/diffulex/strategy/d2f/engine/scheduler.py +++ b/diffulex/strategy/d2f/engine/scheduler.py @@ -2,9 +2,8 @@ from diffulex.config import Config from diffulex.engine.scheduler import AutoScheduler, SchedulerBase -from diffulex.engine.sequence import SequenceBase, SequenceStatus +from diffulex.engine.sequence import SequenceStatus from .sequence import D2FSequence -from diffulex.layer.sampler import SampleOutputForDiffusionLM @AutoScheduler.register("d2f", is_default=True) @@ -19,8 +18,8 @@ def is_finished(self) -> bool: def add(self, seq: D2FSequence) -> None: self.waiting.append(seq) - def schedule(self) -> tuple[list[SequenceBase], bool]: - scheduled: list[SequenceBase] = [] + def schedule(self) -> tuple[list[D2FSequence], bool]: + scheduled: list[D2FSequence] = [] num_seqs = 0 num_batched_tokens = 0 while self.waiting and num_seqs < self.max_num_seqs: @@ -91,7 +90,7 @@ def preempt(self, seq: D2FSequence) -> None: def postprocess( self, seqs: list[D2FSequence], - sample_output: SampleOutputForDiffusionLM, + sample_output, ) -> dict[int, int]: n_diff_steps: dict[int, int] = {} for seq in seqs: diff --git a/diffulex/strategy/d2f/engine/sequence.py b/diffulex/strategy/d2f/engine/sequence.py index 4de92c74..db22bc89 100644 --- a/diffulex/strategy/d2f/engine/sequence.py +++ b/diffulex/strategy/d2f/engine/sequence.py @@ -2,12 +2,12 @@ import torch -from dataclasses import dataclass from enum import Enum, auto +from dataclasses import dataclass from diffulex.config import Config -from diffulex.engine.sequence import AutoSequence, SequenceBase from diffulex.sampling_params import SamplingParams +from diffulex.engine.sequence import AutoSequence, SequenceBase class D2FDiffusionBlockStatus(Enum): @@ -135,14 +135,11 @@ def __init__( if config is None: raise ValueError("SequenceForDiffusionLM requires a Config instance.") self.config = config - self.decoding_strategy = config.decoding_strategy self.kv_cache_layout = config.kv_cache_layout self.eos_token_id = config.eos self.max_model_len = config.max_model_len self.mask_token_id = config.mask_token_id self.diffusion_block_size = config.diffusion_block_size - self.block_mask: torch.Tensor | None = None - self.meet_eos = False self.diffusion_blocks: list[D2FDiffusionBlock] = [] self.n_steps = 0 self.input_token_ids: list[int] = [] @@ -153,7 +150,7 @@ def __repr__(self) -> str: return ( "SequenceForDiffusionLM(seq_id={seq_id}, status={status}, num_tokens={num_tokens}, " "num_prompt_tokens={num_prompt_tokens}, num_cached_tokens={num_cached_tokens}, " - "diffusion_block_size={diffusion_block_size}, mask_shape={mask_shape})" + "diffusion_block_size={diffusion_block_size})" ).format( seq_id=self.seq_id, status=self.status.name, @@ -161,7 +158,6 @@ def __repr__(self) -> str: num_prompt_tokens=self.num_prompt_tokens, num_cached_tokens=self.num_cached_tokens, diffusion_block_size=self.diffusion_block_size, - mask_shape=self.block_mask.shape if self.block_mask is not None else None, ) def __getstate__(self): @@ -197,7 +193,6 @@ def __getstate__(self): "max_tokens": self.max_tokens, "ignore_eos": self.ignore_eos, "config": self.config, - "decoding_strategy": self.decoding_strategy, "kv_cache_layout": self.kv_cache_layout, "eos_token_id": self.eos_token_id, "max_model_len": self.max_model_len, @@ -208,7 +203,6 @@ def __getstate__(self): "input_num_tokens": self.input_num_tokens, "input_num_prompt_tokens": self.input_num_prompt_tokens, "new_tokens": self.new_tokens, - "block_mask": self.block_mask, "meet_eos": self.meet_eos, "n_steps": self.n_steps, } @@ -230,7 +224,6 @@ def __setstate__(self, state): self.meet_eos = state["meet_eos"] self.config = state["config"] - self.decoding_strategy = state.get("decoding_strategy", getattr(self.config, "decoding_strategy", None)) self.kv_cache_layout = state.get("kv_cache_layout", getattr(self.config, "kv_cache_layout", None)) self.eos_token_id = state["eos_token_id"] self.max_model_len = state["max_model_len"] @@ -241,7 +234,6 @@ def __setstate__(self, state): self.input_num_tokens = state.get("input_num_tokens", 0) self.input_num_prompt_tokens = state.get("input_num_prompt_tokens", 0) self.new_tokens = state.get("new_tokens", 0) - self.block_mask = state.get("block_mask") self.n_steps = state.get("n_steps", 0) if self.block_mask is not None and self.block_mask.device.index != torch.cuda.current_device(): @@ -333,6 +325,18 @@ def cached_or_caching_num_tokens(self) -> int: @property def cached_num_tokens(self) -> int: return sum(block.size for block in self.diffusion_blocks if block.is_in_cache) + + @property + def has_to_cache_block(self) -> bool: + return any(block.is_to_cache for block in self.diffusion_blocks) + + @property + def to_cache_last_token_id(self) -> int: + to_cache_num_tokens = 0 + for block in self.diffusion_blocks: + if block.is_to_cache: + to_cache_num_tokens += block.size + return to_cache_num_tokens - 1 @property def num_cached_blocks(self) -> int: @@ -396,40 +400,6 @@ def post_process(self) -> None: def set_layout(self, layout: str) -> None: self.kv_cache_layout = layout - @property - def current_block_mask(self) -> torch.Tensor: - if self.block_mask is None: - raise RuntimeError("Block mask not initialized.") - if self.kv_cache_layout == "distinct": - return self.block_mask[..., self.cached_num_tokens :, self.cached_num_tokens :] - return self.block_mask[..., self.cached_num_tokens :, :] - - def update_block_mask(self, is_prefill: bool = False) -> None: - if is_prefill: - num_tokens = self.num_tokens - mask_shape = (1, 1, num_tokens, num_tokens) - block_mask = torch.zeros(mask_shape, dtype=torch.bool, device=torch.cuda.current_device()) - block_mask[..., : self.input_num_tokens, : self.input_num_tokens] = True - num_diffusion_blocks = ( - self.num_tokens - self.input_num_tokens + self.diffusion_block_size - 1 - ) // self.diffusion_block_size - for block_id in range(num_diffusion_blocks): - start_h = self.input_num_tokens + block_id * self.diffusion_block_size - end_h = start_h + self.diffusion_block_size - block_mask[..., start_h:end_h, :end_h] = True - self.block_mask = block_mask.clone() - return - - if self.block_mask is None: - raise RuntimeError("Prefill block mask must be created before decode updates.") - dev = self.block_mask.device - left_shape = (1, 1, self.num_tokens - self.diffusion_block_size, self.diffusion_block_size) - down_shape = (1, 1, self.diffusion_block_size, self.num_tokens) - left_cat_tensor = torch.zeros(left_shape, dtype=torch.bool, device=dev) - down_cat_tensor = torch.ones(down_shape, dtype=torch.bool, device=dev) - self.block_mask = torch.cat([self.block_mask, left_cat_tensor], dim=-1) - self.block_mask = torch.cat([self.block_mask, down_cat_tensor], dim=-2) - def next_diffusion_step(self, is_prefill: bool = False) -> None: self.n_steps += 1 if is_prefill: @@ -475,6 +445,4 @@ def next_diffusion_step(self, is_prefill: bool = False) -> None: ) self.diffusion_blocks[-1].suf_block = current_block self.token_ids += diffusion_seq - self.num_tokens += added_num_tokens - self.diffusion_blocks.append(current_block) - self.update_block_mask(is_prefill=is_prefill) \ No newline at end of file + self.diffusion_blocks.append(current_block) \ No newline at end of file diff --git a/diffulex/utils/loader.py b/diffulex/utils/loader.py index 5dd07bd3..b2e7cbe9 100755 --- a/diffulex/utils/loader.py +++ b/diffulex/utils/loader.py @@ -7,7 +7,7 @@ from glob import glob from functools import partial from safetensors import safe_open -from diffulex.legacy.config import Config +from diffulex.config import Config def load_lora_config(lora_path: str) -> dict: diff --git a/diffulex_kernel/README.md b/diffulex_kernel/README.md new file mode 100644 index 00000000..e69de29b diff --git a/diffulex_kernel/__init__.py b/diffulex_kernel/__init__.py new file mode 100644 index 00000000..2369bb62 --- /dev/null +++ b/diffulex_kernel/__init__.py @@ -0,0 +1,2 @@ +from diffulex_kernel.python.dllm_flash_attn import dllm_flash_attn_decode, dllm_flash_attn_prefill +from diffulex_kernel.python.kv_cache_kernels import store_kvcache_distinct_layout, store_kvcache_unified_layout \ No newline at end of file diff --git a/diffulex_kernel/python/auto_tuner.py b/diffulex_kernel/python/auto_tuner.py new file mode 100644 index 00000000..f9b5ea0d --- /dev/null +++ b/diffulex_kernel/python/auto_tuner.py @@ -0,0 +1,24 @@ +import itertools + +def build_configs(): + BLOCK_M_LIST = [64, 128, 256] + BLOCK_N_LIST = [64, 128, 256] + NUM_STAGES_LIST = [0, 1, 2] + NUM_THREADS_LIST = [128, 256] + CONFIGS = list( + itertools.product( + BLOCK_M_LIST, + BLOCK_N_LIST, + NUM_STAGES_LIST, + NUM_THREADS_LIST, + ) + ) + + return [ + { + "BLOCK_M": c[0], + "BLOCK_N": c[1], + "NUM_STAGES": c[2], + "NUM_THREADS": c[3], + } for c in CONFIGS + ] \ No newline at end of file diff --git a/diffulex_kernel/python/dllm_flash_attn.py b/diffulex_kernel/python/dllm_flash_attn.py new file mode 100644 index 00000000..099ed68b --- /dev/null +++ b/diffulex_kernel/python/dllm_flash_attn.py @@ -0,0 +1,637 @@ +import torch +import tilelang +import tilelang.language as T + +from flash_attn import flash_attn_varlen_func +from tilelang.autotuner import set_autotune_inputs + +from diffulex_kernel.python.auto_tuner import build_configs +from diffulex_kernel.python.kv_cache_kernels import load_kvcache +from diffulex.attention.metadata import AttnMetaDataBase, is_warming_up + +# from tilelang.engine.callback import register_cuda_postproc_callback +# @register_cuda_postproc_callback +# def tilelang_callback_cuda_postproc(code, _): +# code = "// tilelang_callback_cuda_postproc: generated CUDA code by TileLang\n" + code +# print(code) +# return code + + +kernel_config = None + + +@tilelang.autotune(configs=build_configs()) +@tilelang.jit( + out_idx=[-1], + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}, +) +def dllm_flash_attn_prefill_kernel( + NUM_SEQS: int, + NUM_GROUPS: int, + Q_LEN: int, + KV_LEN: int, + NUM_HEADS: int, + HEAD_DIM: int, + IS_BLOCK_ATTN: bool, + DIFFUSION_BLOCK_SIZE: int, + BLOCK_M: int = 64, + BLOCK_N: int = 64, + NUM_STAGES: int = 1, + NUM_THREADS: int = 128, +): + SCALE = (1.0 / HEAD_DIM)**0.5 * 1.44269504 # log2(e) + NUM_KV_HEADS = NUM_HEADS // NUM_GROUPS + Q_SHAPE = [Q_LEN, NUM_HEADS, HEAD_DIM] + KV_SHAPE = [KV_LEN, NUM_KV_HEADS, HEAD_DIM] + O_SHAPE = [Q_LEN, NUM_HEADS, HEAD_DIM] + DTYPE = "bfloat16" + ACCUM_DTYPE = "float" + + @T.prim_func + def kernel( + Q: T.Tensor(Q_SHAPE, DTYPE), + K: T.Tensor(KV_SHAPE, DTYPE), + V: T.Tensor(KV_SHAPE, DTYPE), + cu_seqlens_q: T.Tensor(NUM_SEQS + 1, "int32"), + cu_seqlens_k: T.Tensor(NUM_SEQS + 1, "int32"), + max_seqlen_q: T.int32, + O: T.Tensor(O_SHAPE, DTYPE), + ): + with T.Kernel(T.ceildiv(max_seqlen_q, BLOCK_M), NUM_HEADS, NUM_SEQS, threads=NUM_THREADS) as (bx, by, bz): + Q_shared = T.alloc_shared([BLOCK_M, HEAD_DIM], DTYPE) + K_shared = T.alloc_shared([BLOCK_N, HEAD_DIM], DTYPE) + V_shared = T.alloc_shared([BLOCK_N, HEAD_DIM], DTYPE) + O_shared = T.alloc_shared([BLOCK_M, HEAD_DIM], DTYPE) + + acc_score = T.alloc_fragment([BLOCK_M, BLOCK_N], ACCUM_DTYPE) + acc_score_cast = T.alloc_fragment([BLOCK_M, BLOCK_N], DTYPE) + acc_output = T.alloc_fragment([BLOCK_M, HEAD_DIM], ACCUM_DTYPE) + scores_max = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_max_prev = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_scale = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_sum = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + log_sum = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + + T.annotate_layout({ + Q_shared: tilelang.layout.make_swizzled_layout(Q_shared), + O_shared: tilelang.layout.make_swizzled_layout(O_shared), + }) + + q_block_idx = bx + seq_idx = bz + head_idx = by + kv_head_idx = head_idx // NUM_GROUPS + + q_start_idx = cu_seqlens_q[seq_idx] + kv_start_idx = cu_seqlens_k[seq_idx] + q_end_idx = cu_seqlens_q[seq_idx + 1] + kv_end_idx = cu_seqlens_k[seq_idx + 1] + + cur_q_seqlen = q_end_idx - q_start_idx + cur_kv_seqlen = kv_end_idx - kv_start_idx + + T.copy(Q[q_start_idx + q_block_idx * BLOCK_M : q_start_idx + (q_block_idx + 1) * BLOCK_M, head_idx, :], Q_shared) + + T.fill(acc_output, 0) + T.fill(acc_score, 0) + T.fill(log_sum, 0) + T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) + + # The same boundary condition as naive causal mask + loop_range = ( + T.min(T.ceildiv(cur_q_seqlen + (q_block_idx + 1) * BLOCK_M, BLOCK_N), T.ceildiv(cur_kv_seqlen, BLOCK_N)) + if IS_BLOCK_ATTN else T.ceildiv(cur_kv_seqlen, BLOCK_N) + ) + for kv_block_idx in T.Pipelined(loop_range, num_stages=NUM_STAGES): + T.copy(K[kv_start_idx + kv_block_idx * BLOCK_N : kv_start_idx + (kv_block_idx + 1) * BLOCK_N, kv_head_idx, :], K_shared) + + # Initialize acc_score with mask + if IS_BLOCK_ATTN: + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + num_diffusion_blocks = (q_block_idx * BLOCK_M + i) // DIFFUSION_BLOCK_SIZE + 1 + acc_score[i, j] = T.if_then_else( + (num_diffusion_blocks * DIFFUSION_BLOCK_SIZE <= kv_block_idx * BLOCK_N + j) or + (q_block_idx * BLOCK_M + i >= cur_q_seqlen or + kv_block_idx * BLOCK_N + j >= cur_kv_seqlen), -1e9, 0 + ) + else: + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + acc_score[i, j] = T.if_then_else( + (q_block_idx * BLOCK_M + i >= cur_q_seqlen or + kv_block_idx * BLOCK_N + j >= cur_kv_seqlen), -1e9, 0 + ) + + # Compute attention scores + T.gemm(Q_shared, K_shared, acc_score, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + # Compute online softmax + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) + T.reduce_max(acc_score, scores_max, dim=1, clear=False) # T.reduce_max(acc_score, scores_max, dim=1, clear=True) # TODO: check if this is correct + for i in T.Parallel(BLOCK_M): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + + for i in T.parallel(BLOCK_M): + scores_scale[i] = T.exp2(scores_max_prev[i] * SCALE - scores_max[i] * SCALE) + + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + acc_score[i, j] = T.exp2(acc_score[i, j] * SCALE - scores_max[i] * SCALE) + + T.reduce_sum(acc_score, scores_sum, dim=1) + for i in T.Parallel(BLOCK_M): + log_sum[i] = log_sum[i] * scores_scale[i] + scores_sum[i] + + T.copy(acc_score, acc_score_cast) + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] *= scores_scale[i] + + # Compute attention output + T.copy(V[kv_start_idx + kv_block_idx * BLOCK_N : kv_start_idx + (kv_block_idx + 1) * BLOCK_N, kv_head_idx, :], V_shared) + T.gemm(acc_score_cast, V_shared, acc_output, policy=T.GemmWarpPolicy.FullRow) + + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] /= log_sum[i] + + T.copy(acc_output, O_shared) + for i, d_idx in T.Parallel(BLOCK_M, HEAD_DIM): + if i + q_block_idx * BLOCK_M < cur_q_seqlen: + O[i + q_start_idx + q_block_idx * BLOCK_M, head_idx, d_idx] = O_shared[i, d_idx] + + return kernel + + +@tilelang.jit( + out_idx=[-1], + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}, +) +def dllm_flash_attn_decode_kernel( + NUM_SEQS: int, + NUM_GROUPS: int, + NUM_PAGE_BLOCKS: int, + Q_LEN: int, + KV_LEN: int, + NUM_HEADS: int, + HEAD_DIM: int, + IS_BLOCK_ATTN: bool, + DIFFUSION_BLOCK_SIZE: int, + MAX_SEQ_NUM_BLOCKS: int, + PAGE_BLOCK_SIZE: int = 32, + BLOCK_M: int = 64, + BLOCK_N: int = 64, + NUM_STAGES: int = 1, + NUM_THREADS: int = 128, +): + SCALE = (1.0 / HEAD_DIM)**0.5 * 1.44269504 # log2(e) + NUM_KV_HEADS = NUM_HEADS // NUM_GROUPS + Q_SHAPE = [Q_LEN, NUM_HEADS, HEAD_DIM] + KV_SHAPE = [KV_LEN, NUM_KV_HEADS, HEAD_DIM] + O_SHAPE = [Q_LEN, NUM_HEADS, HEAD_DIM] + K_CACHE_SHAPE = [NUM_PAGE_BLOCKS, PAGE_BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM] + V_CACHE_SHAPE = [NUM_PAGE_BLOCKS, PAGE_BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM] + BLOCK_TABLE_SHAPE = [NUM_SEQS, MAX_SEQ_NUM_BLOCKS] + DTYPE = "bfloat16" + ACCUM_DTYPE = "float" + + @T.prim_func + def kernel( + Q: T.Tensor(Q_SHAPE, DTYPE), + K: T.Tensor(KV_SHAPE, DTYPE), + V: T.Tensor(KV_SHAPE, DTYPE), + K_Cache: T.Tensor(K_CACHE_SHAPE, DTYPE), + V_Cache: T.Tensor(V_CACHE_SHAPE, DTYPE), + block_tables: T.Tensor(BLOCK_TABLE_SHAPE, "int32"), + context_lens: T.Tensor(NUM_SEQS, "int32"), + cu_seqlens_q: T.Tensor(NUM_SEQS + 1, "int32"), + cu_seqlens_k: T.Tensor(NUM_SEQS + 1, "int32"), + max_seqlen_q: T.int32, + O: T.Tensor(O_SHAPE, DTYPE), + ): + with T.Kernel(NUM_SEQS, NUM_HEADS, threads=NUM_THREADS) as (bx, by): + Q_shared = T.alloc_shared([BLOCK_M, HEAD_DIM], DTYPE) + K_shared = T.alloc_shared([BLOCK_N, HEAD_DIM], DTYPE) + V_shared = T.alloc_shared([BLOCK_N, HEAD_DIM], DTYPE) + O_shared = T.alloc_shared([BLOCK_M, HEAD_DIM], DTYPE) + K_Cache_shared = T.alloc_shared([PAGE_BLOCK_SIZE, HEAD_DIM], DTYPE) + V_Cache_shared = T.alloc_shared([PAGE_BLOCK_SIZE, HEAD_DIM], DTYPE) + + acc_score_kv = T.alloc_fragment([BLOCK_M, BLOCK_N], ACCUM_DTYPE) + acc_score_kv_cast = T.alloc_fragment([BLOCK_M, BLOCK_N], DTYPE) + acc_score_kvcache = T.alloc_fragment([BLOCK_M, PAGE_BLOCK_SIZE], ACCUM_DTYPE) + acc_score_kvcache_cast = T.alloc_fragment([BLOCK_M, PAGE_BLOCK_SIZE], DTYPE) + + acc_output = T.alloc_fragment([BLOCK_M, HEAD_DIM], ACCUM_DTYPE) + scores_max = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_max_prev = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_scale = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_sum = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + log_sum = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + + T.annotate_layout({ + Q_shared: tilelang.layout.make_swizzled_layout(Q_shared), + O_shared: tilelang.layout.make_swizzled_layout(O_shared), + }) + + seq_idx = bx + head_idx = by + kv_head_idx = head_idx // NUM_GROUPS + + q_start_idx = cu_seqlens_q[seq_idx] + kv_start_idx = cu_seqlens_k[seq_idx] + q_end_idx = cu_seqlens_q[seq_idx + 1] + kv_end_idx = cu_seqlens_k[seq_idx + 1] + + cur_q_seqlen = q_end_idx - q_start_idx + cur_kv_seqlen = kv_end_idx - kv_start_idx + + cur_context_len = context_lens[seq_idx] + + T.copy(Q[q_start_idx : q_start_idx + BLOCK_M, head_idx, :], Q_shared) + + T.fill(acc_output, 0) + T.fill(acc_score_kv, 0) + T.fill(acc_score_kvcache, 0) + T.fill(log_sum, 0) + T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) + + # ========================== + # Stage 1: KV Cache Attention (Context) + # ========================== + for page_block_idx_local in T.Pipelined(MAX_SEQ_NUM_BLOCKS, num_stages=NUM_STAGES): + page_block_idx_global = block_tables[seq_idx, page_block_idx_local] + if page_block_idx_global >= 0: + T.copy(K_Cache[page_block_idx_global, :, kv_head_idx, :], K_Cache_shared) + + for i, j in T.Parallel(BLOCK_M, PAGE_BLOCK_SIZE): + acc_score_kvcache[i, j] = T.if_then_else( + (i >= cur_q_seqlen or + page_block_idx_local * PAGE_BLOCK_SIZE + j >= cur_context_len), -1e9, 0 + ) + + # Compute attention scores + T.gemm(Q_shared, K_Cache_shared, acc_score_kvcache, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + # Compute online softmax + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) + T.reduce_max(acc_score_kvcache, scores_max, dim=1, clear=False) + for i in T.Parallel(BLOCK_M): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + + for i in T.Parallel(BLOCK_M): + scores_scale[i] = T.exp2(scores_max_prev[i] * SCALE - scores_max[i] * SCALE) + + for i, j in T.Parallel(BLOCK_M, PAGE_BLOCK_SIZE): + acc_score_kvcache[i, j] = T.exp2(acc_score_kvcache[i, j] * SCALE - scores_max[i] * SCALE) + + T.reduce_sum(acc_score_kvcache, scores_sum, dim=1) + for i in T.Parallel(BLOCK_M): + log_sum[i] = log_sum[i] * scores_scale[i] + scores_sum[i] + + T.copy(acc_score_kvcache, acc_score_kvcache_cast) + + # Scale previous output accumulator + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] *= scores_scale[i] + + # Accumulate current V_cache contribution + T.copy(V_Cache[page_block_idx_global, :, kv_head_idx, :], V_Cache_shared) + T.gemm(acc_score_kvcache_cast, V_Cache_shared, acc_output, policy=T.GemmWarpPolicy.FullRow) + + if page_block_idx_local == MAX_SEQ_NUM_BLOCKS - 1: + # ========================== + # Stage 2: Fresh KV Attention (Self-Attn) + # ========================== + T.copy(K[kv_start_idx : kv_start_idx + BLOCK_N, kv_head_idx, :], K_shared) + + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + acc_score_kv[i, j] = T.if_then_else(i >= cur_q_seqlen or j >= cur_kv_seqlen, -1e9, 0) + + T.gemm(Q_shared, K_shared, acc_score_kv, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) + T.reduce_max(acc_score_kv, scores_max, dim=1, clear=False) + for i in T.Parallel(BLOCK_M): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + + for i in T.Parallel(BLOCK_M): + scores_scale[i] = T.exp2(scores_max_prev[i] * SCALE - scores_max[i] * SCALE) + + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + acc_score_kv[i, j] = T.exp2(acc_score_kv[i, j] * SCALE - scores_max[i] * SCALE) + + T.reduce_sum(acc_score_kv, scores_sum, dim=1) + for i in T.Parallel(BLOCK_M): + log_sum[i] = log_sum[i] * scores_scale[i] + scores_sum[i] + + T.copy(acc_score_kv, acc_score_kv_cast) + + # Scale previous output + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] *= scores_scale[i] + + T.copy(V[kv_start_idx : kv_start_idx + BLOCK_N, kv_head_idx, :], V_shared) + + # Accumulate current V contribution + T.gemm(acc_score_kv_cast, V_shared, acc_output, policy=T.GemmWarpPolicy.FullRow) + + # Finalize + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] /= log_sum[i] + + T.copy(acc_output, O_shared) + for i, d_idx in T.Parallel(BLOCK_M, HEAD_DIM): + if i < cur_q_seqlen: + O[i + q_start_idx, head_idx, d_idx] = O_shared[i, d_idx] + + return kernel + + +@tilelang.jit( + out_idx=[-1], + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}, +) +def dllm_flash_attn_decode_kernel_legacy( + NUM_SEQS: int, + NUM_GROUPS: int, + NUM_PAGE_BLOCKS: int, + Q_LEN: int, + KV_LEN: int, + NUM_HEADS: int, + HEAD_DIM: int, + IS_BLOCK_ATTN: bool, + DIFFUSION_BLOCK_SIZE: int, + MAX_SEQ_NUM_BLOCKS: int, + PAGE_BLOCK_SIZE: int = 32, + BLOCK_M: int = 64, + BLOCK_N: int = 64, + NUM_STAGES: int = 1, + NUM_THREADS: int = 128, +): + SCALE = (1.0 / HEAD_DIM)**0.5 * 1.44269504 # log2(e) + NUM_KV_HEADS = NUM_HEADS // NUM_GROUPS + Q_SHAPE = [Q_LEN, NUM_HEADS, HEAD_DIM] + KV_SHAPE = [KV_LEN, NUM_KV_HEADS, HEAD_DIM] + O_SHAPE = [Q_LEN, NUM_HEADS, HEAD_DIM] + K_CACHE_SHAPE = [NUM_PAGE_BLOCKS, PAGE_BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM] + V_CACHE_SHAPE = [NUM_PAGE_BLOCKS, PAGE_BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM] + BLOCK_TABLE_SHAPE = [NUM_SEQS, MAX_SEQ_NUM_BLOCKS] + DTYPE = "bfloat16" + ACCUM_DTYPE = "float" + + @T.prim_func + def kernel( + Q: T.Tensor(Q_SHAPE, DTYPE), + K: T.Tensor(KV_SHAPE, DTYPE), + V: T.Tensor(KV_SHAPE, DTYPE), + K_Cache: T.Tensor(K_CACHE_SHAPE, DTYPE), + V_Cache: T.Tensor(V_CACHE_SHAPE, DTYPE), + block_tables: T.Tensor(BLOCK_TABLE_SHAPE, "int32"), + context_lens: T.Tensor(NUM_SEQS, "int32"), + cu_seqlens_q: T.Tensor(NUM_SEQS + 1, "int32"), + cu_seqlens_k: T.Tensor(NUM_SEQS + 1, "int32"), + max_seqlen_q: T.int32, + O: T.Tensor(O_SHAPE, DTYPE), + ): + with T.Kernel(NUM_SEQS, NUM_HEADS, threads=NUM_THREADS) as (bx, by): + Q_shared = T.alloc_shared([BLOCK_M, HEAD_DIM], DTYPE) + K_shared = T.alloc_shared([BLOCK_N, HEAD_DIM], DTYPE) + V_shared = T.alloc_shared([BLOCK_N, HEAD_DIM], DTYPE) + O_shared = T.alloc_shared([BLOCK_M, HEAD_DIM], DTYPE) + K_Cache_shared = T.alloc_shared([PAGE_BLOCK_SIZE, HEAD_DIM], DTYPE) + V_Cache_shared = T.alloc_shared([PAGE_BLOCK_SIZE, HEAD_DIM], DTYPE) + + acc_score_kv = T.alloc_fragment([BLOCK_M, BLOCK_N], ACCUM_DTYPE) + acc_score_kv_cast = T.alloc_fragment([BLOCK_M, BLOCK_N], DTYPE) + acc_score_kvcache = T.alloc_fragment([BLOCK_M, PAGE_BLOCK_SIZE], ACCUM_DTYPE) + acc_score_kvcache_cast = T.alloc_fragment([BLOCK_M, PAGE_BLOCK_SIZE], DTYPE) + + acc_output = T.alloc_fragment([BLOCK_M, HEAD_DIM], ACCUM_DTYPE) + scores_max = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_max_prev = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_scale = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + scores_sum = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + log_sum = T.alloc_fragment([BLOCK_M], ACCUM_DTYPE) + block_table = T.alloc_fragment([MAX_SEQ_NUM_BLOCKS], "int32") + + T.annotate_layout({ + Q_shared: tilelang.layout.make_swizzled_layout(Q_shared), + O_shared: tilelang.layout.make_swizzled_layout(O_shared), + }) + + seq_idx = bx + head_idx = by + kv_head_idx = head_idx // NUM_GROUPS + + q_start_idx = cu_seqlens_q[seq_idx] + kv_start_idx = cu_seqlens_k[seq_idx] + q_end_idx = cu_seqlens_q[seq_idx + 1] + kv_end_idx = cu_seqlens_k[seq_idx + 1] + + cur_q_seqlen = q_end_idx - q_start_idx + cur_kv_seqlen = kv_end_idx - kv_start_idx + + cur_context_len = context_lens[seq_idx] + + T.copy(block_tables[seq_idx, :], block_table) + T.copy(Q[q_start_idx : q_start_idx + BLOCK_M, head_idx, :], Q_shared) + + T.fill(acc_output, 0) + T.fill(acc_score_kv, 0) + T.fill(acc_score_kvcache, 0) + T.fill(log_sum, 0) + T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) + + # ========================== + # Stage 1: KV Cache Attention (Context) + # ========================== + for page_block_idx_local in T.Pipelined(MAX_SEQ_NUM_BLOCKS, num_stages=NUM_STAGES): + page_block_idx_global = block_table[page_block_idx_local] + if page_block_idx_global >= 0: + T.copy(K_Cache[page_block_idx_global, :, kv_head_idx, :], K_Cache_shared) + + for i, j in T.Parallel(BLOCK_M, PAGE_BLOCK_SIZE): + acc_score_kvcache[i, j] = T.if_then_else( + (i >= cur_q_seqlen or + page_block_idx_local * PAGE_BLOCK_SIZE + j >= cur_context_len), -1e9, 0 + ) + + # Compute attention scores + T.gemm(Q_shared, K_Cache_shared, acc_score_kvcache, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + # Compute online softmax + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) + T.reduce_max(acc_score_kvcache, scores_max, dim=1, clear=False) + for i in T.Parallel(BLOCK_M): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + + for i in T.Parallel(BLOCK_M): + scores_scale[i] = T.exp2(scores_max_prev[i] * SCALE - scores_max[i] * SCALE) + + for i, j in T.Parallel(BLOCK_M, PAGE_BLOCK_SIZE): + acc_score_kvcache[i, j] = T.exp2(acc_score_kvcache[i, j] * SCALE - scores_max[i] * SCALE) + + T.reduce_sum(acc_score_kvcache, scores_sum, dim=1) + for i in T.Parallel(BLOCK_M): + log_sum[i] = log_sum[i] * scores_scale[i] + scores_sum[i] + + T.copy(acc_score_kvcache, acc_score_kvcache_cast) + + # Scale previous output accumulator + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] *= scores_scale[i] + + # Accumulate current V_cache contribution + T.copy(V_Cache[page_block_idx_global, :, kv_head_idx, :], V_Cache_shared) + T.gemm(acc_score_kvcache_cast, V_Cache_shared, acc_output, policy=T.GemmWarpPolicy.FullRow) + + if page_block_idx_local == MAX_SEQ_NUM_BLOCKS - 1: + # ========================== + # Stage 2: Fresh KV Attention (Self-Attn) + # ========================== + T.copy(K[kv_start_idx : kv_start_idx + BLOCK_N, kv_head_idx, :], K_shared) + + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + acc_score_kv[i, j] = T.if_then_else(i >= cur_q_seqlen or j >= cur_kv_seqlen, -1e9, 0) + + T.gemm(Q_shared, K_shared, acc_score_kv, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) + T.reduce_max(acc_score_kv, scores_max, dim=1, clear=False) + for i in T.Parallel(BLOCK_M): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + + for i in T.Parallel(BLOCK_M): + scores_scale[i] = T.exp2(scores_max_prev[i] * SCALE - scores_max[i] * SCALE) + + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + acc_score_kv[i, j] = T.exp2(acc_score_kv[i, j] * SCALE - scores_max[i] * SCALE) + + T.reduce_sum(acc_score_kv, scores_sum, dim=1) + for i in T.Parallel(BLOCK_M): + log_sum[i] = log_sum[i] * scores_scale[i] + scores_sum[i] + + T.copy(acc_score_kv, acc_score_kv_cast) + + # Scale previous output + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] *= scores_scale[i] + + T.copy(V[kv_start_idx : kv_start_idx + BLOCK_N, kv_head_idx, :], V_shared) + + # Accumulate current V contribution + T.gemm(acc_score_kv_cast, V_shared, acc_output, policy=T.GemmWarpPolicy.FullRow) + + # Finalize + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] /= log_sum[i] + + T.copy(acc_output, O_shared) + for i, d_idx in T.Parallel(BLOCK_M, HEAD_DIM): + if i < cur_q_seqlen: + O[i + q_start_idx, head_idx, d_idx] = O_shared[i, d_idx] + + return kernel + + +def dllm_flash_attn_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float, + attn_metadata: AttnMetaDataBase +) -> torch.Tensor: + if attn_metadata.attn_type == "full_attention": + return flash_attn_varlen_func( + q, k, v, + attn_metadata.cu_seqlens_q, attn_metadata.cu_seqlens_k, + attn_metadata.max_seqlen_q, attn_metadata.max_seqlen_k, + softmax_scale=scale, block_table=None + ) + elif attn_metadata.attn_type == "block_attention": + if is_warming_up(): + global kernel_config + with set_autotune_inputs([ + q, k, v, + attn_metadata.cu_seqlens_q, + attn_metadata.cu_seqlens_k, + attn_metadata.max_seqlen_q, + ]): + prefill_kernel = dllm_flash_attn_prefill_kernel( + attn_metadata.num_seqs, + q.shape[1] // k.shape[1], + q.shape[0], + k.shape[0], + q.shape[1], + q.shape[2], + attn_metadata.attn_type == "block_attention", + attn_metadata.diffusion_block_size + ) + kernel_config = prefill_kernel.config + return prefill_kernel( + q, k, v, + attn_metadata.cu_seqlens_q, + attn_metadata.cu_seqlens_k, + attn_metadata.max_seqlen_q, + ) + else: + prefill_kernel = dllm_flash_attn_prefill_kernel( + attn_metadata.num_seqs, + q.shape[1] // k.shape[1], + q.shape[0], + k.shape[0], + q.shape[1], + q.shape[2], + attn_metadata.attn_type == "block_attention", + attn_metadata.diffusion_block_size, + **kernel_config + ) + return prefill_kernel( + q, k, v, + attn_metadata.cu_seqlens_q, + attn_metadata.cu_seqlens_k, + attn_metadata.max_seqlen_q, + ) + + +def dllm_flash_attn_decode( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + scale: float, + attn_metadata: AttnMetaDataBase +) -> torch.Tensor: + if attn_metadata.decode_mode == "static": + decode_kernel = dllm_flash_attn_decode_kernel( + attn_metadata.num_seqs, + q.shape[1] // k.shape[1], + k_cache.shape[0], + q.shape[0], + k.shape[0], + q.shape[1], + q.shape[2], + attn_metadata.attn_type == "block_attention", + attn_metadata.diffusion_block_size, + attn_metadata.block_tables.shape[1], + attn_metadata.page_block_size, + **kernel_config + ) + + return decode_kernel( + q, k, v, k_cache, v_cache, + attn_metadata.block_tables, + attn_metadata.context_lens, + attn_metadata.cu_seqlens_q, + attn_metadata.cu_seqlens_k, + attn_metadata.max_seqlen_q, + ) + elif attn_metadata.decode_mode == "varlen": + k_comb, v_comb = load_kvcache(k_cache, v_cache, attn_metadata, k, v) + return flash_attn_varlen_func(q, k_comb, v_comb, + attn_metadata.cu_seqlens_q, attn_metadata.cu_seqlens_k, + attn_metadata.max_seqlen_q, attn_metadata.max_seqlen_k, + softmax_scale=scale, block_table=None) \ No newline at end of file diff --git a/diffulex/attention/ops/kv_cache_kernels.py b/diffulex_kernel/python/kv_cache_kernels.py similarity index 60% rename from diffulex/attention/ops/kv_cache_kernels.py rename to diffulex_kernel/python/kv_cache_kernels.py index 41a42ba6..b235f83a 100755 --- a/diffulex/attention/ops/kv_cache_kernels.py +++ b/diffulex_kernel/python/kv_cache_kernels.py @@ -3,33 +3,13 @@ import triton.language as tl -from typing import Any +from typing import Tuple - -@triton.jit -def store_kvcache_kernel_causal_lm( - key_ptr, - key_stride, - value_ptr, - value_stride, - k_cache_ptr, - v_cache_ptr, - slot_mapping_ptr, - D: tl.constexpr -): - idx = tl.program_id(0) - key_offsets = idx * key_stride + tl.arange(0, D) - value_offsets = idx * value_stride + tl.arange(0, D) - key = tl.load(key_ptr + key_offsets) - value = tl.load(value_ptr + value_offsets) - slot = tl.load(slot_mapping_ptr + idx) - cache_offsets = slot * D + tl.arange(0, D) - tl.store(k_cache_ptr + cache_offsets, key) - tl.store(v_cache_ptr + cache_offsets, value) +from diffulex.attention.metadata import AttnMetaDataBase @triton.jit -def store_kvcache_kernel_diffusion_lm( +def dllm_store_kvcache_kernel_unified( key_ptr, key_stride, value_ptr, @@ -53,7 +33,7 @@ def store_kvcache_kernel_diffusion_lm( @triton.jit -def store_kvcache_kernel_diffusion_lm_distinct( +def dllm_store_kvcache_kernel_distinct( k_ptr, v_ptr, k_cache_ptr, v_cache_ptr, slot_mapping_ptr, k_stride, v_stride, k_cache_stride_nblks, k_cache_stride_h, k_cache_stride_dx, k_cache_stride_blk_sz, k_cache_stride_x, @@ -105,8 +85,8 @@ def store_kvcache_kernel_diffusion_lm_distinct( def store_kvcache_distinct_layout(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, - slot_mapping: torch.Tensor, - context = None) -> None: + slot_mapping: torch.Tensor, attn_metadata: AttnMetaDataBase) -> None: + # TODO: implement diffusion lm kv cache store # k_cache: [num_blks, h, hdim // x, blk_sz, x] # v_cache: [num_blks, h, hdim, blk_sz] NBlks, NHeads, HDim_x, Blk_sz, x = k_cache.shape @@ -116,7 +96,7 @@ def store_kvcache_distinct_layout(key: torch.Tensor, value: torch.Tensor, assert N == slot_mapping.numel() GRID = (N, ) - store_kvcache_kernel_diffusion_lm_distinct[GRID]( + dllm_store_kvcache_kernel_distinct[GRID]( key, value, k_cache, v_cache, slot_mapping, @@ -129,16 +109,15 @@ def store_kvcache_distinct_layout(key: torch.Tensor, value: torch.Tensor, def store_kvcache_unified_layout(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, - slot_mapping: torch.Tensor, - context: Any = None) -> None: + slot_mapping: torch.Tensor, attn_metadata: AttnMetaDataBase) -> None: N, num_heads, head_dim = key.shape D = num_heads * head_dim assert key.stride(-1) == 1 and value.stride(-1) == 1 assert key.stride(1) == head_dim and value.stride(1) == head_dim assert k_cache.stride(1) == D and v_cache.stride(1) == D assert N == slot_mapping.numel(), f"`N`: {N}, `slot_mapping.numel()`: {slot_mapping.numel()}" - - store_kvcache_kernel_diffusion_lm[(N,)]( + + dllm_store_kvcache_kernel_unified[(N,)]( key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D @@ -146,23 +125,23 @@ def store_kvcache_unified_layout(key: torch.Tensor, value: torch.Tensor, @triton.jit -def load_kvcache_kernel_kv(k_cache_ptr, v_cache_ptr, - k_new_ptr, v_new_ptr, - block_table_ptr, - k_out_ptr, v_out_ptr, - seqlens_ptr, ctxlens_ptr, - cu_seqlens_q_ptr, cu_seqlens_k_ptr, - kv_cache_stride_nblks, kv_cache_stride_blk, kv_cache_stride_h, kv_cache_stride_d, - kv_new_stride_s, kv_new_stride_h, kv_new_stride_d, - block_table_stride_nseqs, block_table_stride_maxblks, - kv_out_stride_s, kv_out_stride_h, kv_out_stride_d, - ctxlens_stride, seqlens_stride, - cu_seqlens_q_stride, cu_seqlens_k_stride, - LAST_BLK_ID: tl.constexpr, - HEAD_DIM: tl.constexpr, - PAGE_SIZE: tl.constexpr, - DIFFUSION_BLOCK_SIZE: tl.constexpr, - KV_LOAD_UNROLL_FACTOR: tl.constexpr): +def load_kvcache_kernel(k_cache_ptr, v_cache_ptr, + k_new_ptr, v_new_ptr, + block_table_ptr, + k_out_ptr, v_out_ptr, + seqlens_ptr, ctxlens_ptr, + cu_seqlens_q_ptr, cu_seqlens_k_ptr, + kv_cache_stride_nblks, kv_cache_stride_blk, kv_cache_stride_h, kv_cache_stride_d, + kv_new_stride_s, kv_new_stride_h, kv_new_stride_d, + block_table_stride_nseqs, block_table_stride_maxblks, + kv_out_stride_s, kv_out_stride_h, kv_out_stride_d, + ctxlens_stride, seqlens_stride, + cu_seqlens_q_stride, cu_seqlens_k_stride, + LAST_BLK_ID: tl.constexpr, + HEAD_DIM: tl.constexpr, + PAGE_SIZE: tl.constexpr, + DIFFUSION_BLOCK_SIZE: tl.constexpr, + KV_LOAD_UNROLL_FACTOR: tl.constexpr): # BUG FIX # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: D2F @@ -249,23 +228,23 @@ def load_kvcache_kernel_kv(k_cache_ptr, v_cache_ptr, def load_kvcache(k_cache: torch.Tensor, v_cache: torch.Tensor, - context: Any, - k_new: torch.Tensor, v_new: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + attn_metadata: AttnMetaDataBase, + k_new: torch.Tensor, v_new: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert k_cache.shape == v_cache.shape assert k_new.shape == v_new.shape N_BLOCKS, PAGE_SIZE, H_KV, HEAD_DIM = k_cache.shape - NUM_SEQS, MAX_SEQ_BLOCKS = context.block_tables.shape + NUM_SEQS, MAX_SEQ_BLOCKS = attn_metadata.block_tables.shape - ctxlens = context.context_lens - seqlens = context.seq_lens_ts + ctxlens = attn_metadata.context_lens + seqlens = attn_metadata.seq_lens_ts assert sum(seqlens) == k_new.shape[0] - DIFFUSION_BLOCK_SIZE = context.seqs[0].diffusion_block_size + DIFFUSION_BLOCK_SIZE = attn_metadata.seqs[0].diffusion_block_size MAX_DIFFUSION_BLOCK_SIZE = max(seqlens) assert MAX_DIFFUSION_BLOCK_SIZE % DIFFUSION_BLOCK_SIZE == 0 total_lens = ctxlens + seqlens - cu_seqlens_q = context.cu_seqlens_q - cu_seqlens_k = context.cu_seqlens_k + cu_seqlens_q = attn_metadata.cu_seqlens_q + cu_seqlens_k = attn_metadata.cu_seqlens_k assert sum(total_lens) == cu_seqlens_k[-1] assert cu_seqlens_q.shape == cu_seqlens_k.shape assert cu_seqlens_q.shape[0] == NUM_SEQS + 1 @@ -275,94 +254,26 @@ def load_kvcache(k_cache: torch.Tensor, v_cache: torch.Tensor, v_output = torch.empty_like(k_output) GRID = (NUM_SEQS, MAX_SEQ_BLOCKS, H_KV) - load_kvcache_kernel_kv[GRID]( + load_kvcache_kernel[GRID]( k_cache, v_cache, k_new, v_new, - context.block_tables, + attn_metadata.block_tables, k_output, v_output, seqlens, ctxlens, cu_seqlens_q, cu_seqlens_k, *k_cache.stride(), *k_new.stride(), - *context.block_tables.stride(), + *attn_metadata.block_tables.stride(), *k_output.stride(), ctxlens.stride(0), seqlens.stride(0), cu_seqlens_q.stride(0), cu_seqlens_k.stride(0), - LAST_BLK_ID=context.block_tables.shape[-1] - 1, + LAST_BLK_ID=attn_metadata.block_tables.shape[-1] - 1, HEAD_DIM=HEAD_DIM, PAGE_SIZE=PAGE_SIZE, DIFFUSION_BLOCK_SIZE=DIFFUSION_BLOCK_SIZE, KV_LOAD_UNROLL_FACTOR=2 ) - return k_output, v_output - - -def CHECK_STORING(k_cache: torch.Tensor, v_cache: torch.Tensor, - k: torch.Tensor, v: torch.Tensor, - context) -> None: - k_list, v_list = [torch.split(tensor, context.seq_lens, dim=0) for tensor in (k, v)] - for seq_idx, seq in enumerate(context.seqs): - cached_num_tokens = seq.cached_num_tokens - caching_num_tokens = seq.caching_num_tokens - block_size = seq.block_size - if caching_num_tokens == 0: - continue - - k_cache_list, v_cache_list = [], [] - for local_mem_blk_idx, global_mem_blk_idx in enumerate(context.block_tables[seq_idx]): - if caching_num_tokens == 0: - break - - if global_mem_blk_idx.item() == -1: - continue - - if cached_num_tokens > block_size: - cached_num_tokens -= block_size - continue - - cur_start_idx = cached_num_tokens % block_size - remain_num_tokens = min(block_size - cur_start_idx, caching_num_tokens) - k_cache_list.append(k_cache[global_mem_blk_idx, cur_start_idx:cur_start_idx + remain_num_tokens]) - v_cache_list.append(v_cache[global_mem_blk_idx, cur_start_idx:cur_start_idx + remain_num_tokens]) - cached_num_tokens += remain_num_tokens - caching_num_tokens -= remain_num_tokens - k_cache_temp = torch.cat(k_cache_list, dim=0) - v_cache_temp = torch.cat(v_cache_list, dim=0) - assert torch.allclose(k_cache_temp, k_list[seq_idx][:seq.caching_num_tokens], atol=1e-5), f"K cache mismatch for seq {seq_idx}!" - assert torch.allclose(v_cache_temp, v_list[seq_idx][:seq.caching_num_tokens], atol=1e-5), f"V cache mismatch for seq {seq_idx}!" - - -def CHECK_LOADING(k_comb: torch.Tensor, v_comb: torch.Tensor, - k_new: torch.Tensor, v_new: torch.Tensor, - k_cache: torch.Tensor, v_cache: torch.Tensor, - context: Any) -> tuple[torch.Tensor, torch.Tensor]: - try: - k_list, v_list = [torch.split(tensor, context.seq_lens, dim=0) for tensor in (k_new, v_new)] - cat_k_list = [] - cat_v_list = [] - for seq_idx, (k, v) in enumerate(zip(k_list, v_list)): - cur_ctxlen = context.context_lens[seq_idx] - k_cache_temp, v_cache_temp = None, None - for mem_block_idx in context.block_tables[seq_idx]: - if mem_block_idx.item() == -1: - continue - k_mem_block, v_mem_block = k_cache[mem_block_idx], v_cache[mem_block_idx] - mem_block_size = k_cache.shape[1] - cur_window = mem_block_size if mem_block_size <= cur_ctxlen else cur_ctxlen % mem_block_size - cur_ctxlen = cur_ctxlen - cur_window - k_cache_temp = k_mem_block[:cur_window] if k_cache_temp is None \ - else torch.cat((k_cache_temp, k_mem_block[:cur_window]), dim=0) - v_cache_temp = v_mem_block[:cur_window] if v_cache_temp is None \ - else torch.cat((v_cache_temp, v_mem_block[:cur_window]), dim=0) - cat_k_list.extend([k_cache_temp, k]) - cat_v_list.extend([v_cache_temp, v]) - k_cache_check, v_cache_check = torch.cat(cat_k_list, dim=0), torch.cat(cat_v_list, dim=0) - assert torch.allclose(k_comb, k_cache_check, atol=1e-5), "K cache mismatch!" - assert torch.allclose(v_comb, v_cache_check, atol=1e-5), "V cache mismatch!" - return k_comb, v_comb - except AssertionError as e: - raise ValueError(f"KV cache loading check failed: {e}") - # return k_cache_check, v_cache_check \ No newline at end of file + return k_output, v_output \ No newline at end of file diff --git a/diffulex_legacy/__init__.py b/diffulex_legacy/__init__.py new file mode 100755 index 00000000..9923ac00 --- /dev/null +++ b/diffulex_legacy/__init__.py @@ -0,0 +1,2 @@ +from diffulex_legacy.llm import LLM +from diffulex_legacy.sampling_params import SamplingParams diff --git a/diffulex/legacy/config.py b/diffulex_legacy/config.py similarity index 100% rename from diffulex/legacy/config.py rename to diffulex_legacy/config.py diff --git a/diffulex/legacy/engine/block_manager.py b/diffulex_legacy/engine/block_manager.py similarity index 98% rename from diffulex/legacy/engine/block_manager.py rename to diffulex_legacy/engine/block_manager.py index 7f12ce9b..b8413bc8 100755 --- a/diffulex/legacy/engine/block_manager.py +++ b/diffulex_legacy/engine/block_manager.py @@ -7,8 +7,8 @@ from dataclasses import dataclass, field from typing import List, Dict, Deque, Set -from diffulex.legacy.config import Config -from diffulex.legacy.engine.sequence import SequenceBase, SequenceForCausalLM, SequenceForDiffusionLM +from diffulex_legacy.config import Config +from diffulex_legacy.engine.sequence import SequenceBase, SequenceForCausalLM, SequenceForDiffusionLM @dataclass diff --git a/diffulex/legacy/engine/dp_engine.py b/diffulex_legacy/engine/dp_engine.py similarity index 98% rename from diffulex/legacy/engine/dp_engine.py rename to diffulex_legacy/engine/dp_engine.py index 8fe47821..70f8e829 100755 --- a/diffulex/legacy/engine/dp_engine.py +++ b/diffulex_legacy/engine/dp_engine.py @@ -10,9 +10,9 @@ from typing import List, Any from multiprocessing.connection import wait as mp_wait -from diffulex.legacy.config import Config -from diffulex.legacy.engine.llm_engine import LLMEngine -from diffulex.legacy.sampling_params import SamplingParams +from diffulex_legacy.config import Config +from diffulex_legacy.engine.llm_engine import LLMEngine +from diffulex_legacy.sampling_params import SamplingParams def _dp_child_entry(config: Config, dp_idx: int, local_devices: list[int], conn): diff --git a/diffulex/legacy/engine/llm_engine.py b/diffulex_legacy/engine/llm_engine.py similarity index 95% rename from diffulex/legacy/engine/llm_engine.py rename to diffulex_legacy/engine/llm_engine.py index 580bff08..3db08306 100755 --- a/diffulex/legacy/engine/llm_engine.py +++ b/diffulex_legacy/engine/llm_engine.py @@ -8,11 +8,11 @@ from dataclasses import fields from transformers import AutoTokenizer -from diffulex.legacy.config import Config -from diffulex.legacy.sampling_params import SamplingParams -from diffulex.legacy.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM -from diffulex.legacy.engine.scheduler import AutoScheduler, SchedulerBase -from diffulex.legacy.engine.model_runner import AutoModelRunner +from diffulex_legacy.config import Config +from diffulex_legacy.sampling_params import SamplingParams +from diffulex_legacy.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM +from diffulex_legacy.engine.scheduler import AutoScheduler, SchedulerBase +from diffulex_legacy.engine.model_runner import AutoModelRunner class LLMEngine: diff --git a/diffulex/legacy/engine/model_runner.py b/diffulex_legacy/engine/model_runner.py similarity index 99% rename from diffulex/legacy/engine/model_runner.py rename to diffulex_legacy/engine/model_runner.py index 4a881aee..e7fcd0e9 100755 --- a/diffulex/legacy/engine/model_runner.py +++ b/diffulex_legacy/engine/model_runner.py @@ -9,12 +9,12 @@ from multiprocessing.synchronize import Event from multiprocessing.shared_memory import SharedMemory -from diffulex.legacy.config import Config -from diffulex.legacy.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM, SequenceBase -from diffulex.legacy.models.auto_model import AutoModelLM -from diffulex.legacy.layers.sampler import AutoSampler -from diffulex.legacy.utils.checker import CHECK_SLOT_MAPPING -from diffulex.legacy.utils.context import ( +from diffulex_legacy.config import Config +from diffulex_legacy.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM, SequenceBase +from diffulex_legacy.models.auto_model import AutoModelLM +from diffulex_legacy.layers.sampler import AutoSampler +from diffulex_legacy.utils.checker import CHECK_SLOT_MAPPING +from diffulex_legacy.utils.context import ( set_context_causal_lm, get_context_causal_lm, reset_context_causal_lm, diff --git a/diffulex/legacy/engine/scheduler.py b/diffulex_legacy/engine/scheduler.py similarity index 97% rename from diffulex/legacy/engine/scheduler.py rename to diffulex_legacy/engine/scheduler.py index e1718bdf..469f8470 100755 --- a/diffulex/legacy/engine/scheduler.py +++ b/diffulex_legacy/engine/scheduler.py @@ -4,13 +4,13 @@ from abc import ABC, abstractmethod from typing import Tuple, List, Deque -from diffulex.legacy.config import Config -from diffulex.legacy.engine.sequence import ( +from diffulex_legacy.config import Config +from diffulex_legacy.engine.sequence import ( SequenceBase, SequenceStatus, SequenceForDiffusionLM, SequenceForCausalLM ) -from diffulex.legacy.layers.sampler import SampleOutputForDiffusionLM -from diffulex.legacy.engine.block_manager import AutoBlockManager +from diffulex_legacy.layers.sampler import SampleOutputForDiffusionLM +from diffulex_legacy.engine.block_manager import AutoBlockManager class SchedulerBase(ABC): diff --git a/diffulex/legacy/engine/sequence.py b/diffulex_legacy/engine/sequence.py similarity index 99% rename from diffulex/legacy/engine/sequence.py rename to diffulex_legacy/engine/sequence.py index 4f32c55c..8b41e9dc 100755 --- a/diffulex/legacy/engine/sequence.py +++ b/diffulex_legacy/engine/sequence.py @@ -6,8 +6,8 @@ from dataclasses import dataclass from typing import List, Tuple, Any -from diffulex.legacy.config import Config -from diffulex.legacy.sampling_params import SamplingParams +from diffulex_legacy.config import Config +from diffulex_legacy.sampling_params import SamplingParams class SequenceStatus(Enum): diff --git a/diffulex/legacy/layers/activation.py b/diffulex_legacy/layers/activation.py similarity index 100% rename from diffulex/legacy/layers/activation.py rename to diffulex_legacy/layers/activation.py diff --git a/diffulex/legacy/layers/attention/attention_v1.py b/diffulex_legacy/layers/attention/attention_v1.py similarity index 99% rename from diffulex/legacy/layers/attention/attention_v1.py rename to diffulex_legacy/layers/attention/attention_v1.py index 4ca79c21..dbcbd655 100755 --- a/diffulex/legacy/layers/attention/attention_v1.py +++ b/diffulex_legacy/layers/attention/attention_v1.py @@ -18,7 +18,7 @@ else: from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache -from diffulex.legacy.utils.context import ( +from diffulex_legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffulex/legacy/layers/attention/attention_v1_profile.py b/diffulex_legacy/layers/attention/attention_v1_profile.py similarity index 99% rename from diffulex/legacy/layers/attention/attention_v1_profile.py rename to diffulex_legacy/layers/attention/attention_v1_profile.py index 6bb44fee..f3e0f5de 100755 --- a/diffulex/legacy/layers/attention/attention_v1_profile.py +++ b/diffulex_legacy/layers/attention/attention_v1_profile.py @@ -20,7 +20,7 @@ else: from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache -from diffulex.legacy.utils.context import ( +from diffulex_legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffulex/legacy/layers/attention/attention_v2.py b/diffulex_legacy/layers/attention/attention_v2.py similarity index 99% rename from diffulex/legacy/layers/attention/attention_v2.py rename to diffulex_legacy/layers/attention/attention_v2.py index 5238fd13..970ac03b 100755 --- a/diffulex/legacy/layers/attention/attention_v2.py +++ b/diffulex_legacy/layers/attention/attention_v2.py @@ -17,7 +17,7 @@ else: from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache -from diffulex.legacy.utils.context import ( +from diffulex_legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffulex/legacy/layers/attention/attention_v2_dup.py b/diffulex_legacy/layers/attention/attention_v2_dup.py similarity index 99% rename from diffulex/legacy/layers/attention/attention_v2_dup.py rename to diffulex_legacy/layers/attention/attention_v2_dup.py index 43ce5e90..f5afbbab 100755 --- a/diffulex/legacy/layers/attention/attention_v2_dup.py +++ b/diffulex_legacy/layers/attention/attention_v2_dup.py @@ -17,8 +17,8 @@ else: from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache -from diffulex.legacy.engine.sequence import SequenceForDiffusionLM -from diffulex.legacy.utils.context import ( +from diffulex_legacy.engine.sequence import SequenceForDiffusionLM +from diffulex_legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffulex/legacy/layers/attention/attention_v2_profile.py b/diffulex_legacy/layers/attention/attention_v2_profile.py similarity index 99% rename from diffulex/legacy/layers/attention/attention_v2_profile.py rename to diffulex_legacy/layers/attention/attention_v2_profile.py index 2a98b209..8817c1d6 100755 --- a/diffulex/legacy/layers/attention/attention_v2_profile.py +++ b/diffulex_legacy/layers/attention/attention_v2_profile.py @@ -20,7 +20,7 @@ else: from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache -from diffulex.legacy.utils.context import ( +from diffulex_legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffulex/legacy/layers/attention/attention_v3.py b/diffulex_legacy/layers/attention/attention_v3.py similarity index 99% rename from diffulex/legacy/layers/attention/attention_v3.py rename to diffulex_legacy/layers/attention/attention_v3.py index cfd02bcc..ec438f3b 100755 --- a/diffulex/legacy/layers/attention/attention_v3.py +++ b/diffulex_legacy/layers/attention/attention_v3.py @@ -10,8 +10,8 @@ from torch.nn.attention.flex_attention import flex_attention, create_block_mask from flash_attn import flash_attn_with_kvcache -from diffulex.legacy.layers.attention.ops import causal_lm_flash_decoding -from diffulex.legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm +from diffulex_legacy.layers.attention.ops import causal_lm_flash_decoding +from diffulex_legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm @triton.jit diff --git a/diffulex/legacy/layers/attention/attention_v4.py b/diffulex_legacy/layers/attention/attention_v4.py similarity index 98% rename from diffulex/legacy/layers/attention/attention_v4.py rename to diffulex_legacy/layers/attention/attention_v4.py index e846fd82..88e624f5 100755 --- a/diffulex/legacy/layers/attention/attention_v4.py +++ b/diffulex_legacy/layers/attention/attention_v4.py @@ -9,12 +9,12 @@ from torch.nn.attention.flex_attention import create_block_mask from transformers.integrations.flex_attention import compile_friendly_flex_attention as flex_attention -from diffulex.legacy.layers.attention.ops import ( +from diffulex_legacy.layers.attention.ops import ( causal_lm_flash_decoding, diffusion_lm_flash_decoding, diffusion_lm_parallel_flash_decoding, store_kvcache_unified_layout, store_kvcache_distinct_layout, load_kvcache, CHECK_STORING, CHECK_LOADING, CHECK_ATTENTION ) -from diffulex.legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm +from diffulex_legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm class Attention(nn.Module): diff --git a/diffulex/legacy/layers/attention/attention_v5.py b/diffulex_legacy/layers/attention/attention_v5.py similarity index 98% rename from diffulex/legacy/layers/attention/attention_v5.py rename to diffulex_legacy/layers/attention/attention_v5.py index e019bca4..4ac0727f 100644 --- a/diffulex/legacy/layers/attention/attention_v5.py +++ b/diffulex_legacy/layers/attention/attention_v5.py @@ -10,12 +10,12 @@ from flash_attn import flash_attn_varlen_func from transformers.integrations.flex_attention import compile_friendly_flex_attention as flex_attention -from diffulex.legacy.layers.attention.ops import ( +from diffulex_legacy.layers.attention.ops import ( causal_lm_flash_decoding, diffusion_lm_flash_decoding, diffusion_lm_parallel_flash_decoding, store_kvcache_unified_layout, store_kvcache_distinct_layout, load_kvcache, CHECK_STORING, CHECK_LOADING, CHECK_ATTENTION ) -from diffulex.legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm +from diffulex_legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm class Attention(nn.Module): diff --git a/diffulex/attention/ops/__init__.py b/diffulex_legacy/layers/attention/ops/__init__.py similarity index 56% rename from diffulex/attention/ops/__init__.py rename to diffulex_legacy/layers/attention/ops/__init__.py index 579ccbfe..d9a1b2ef 100755 --- a/diffulex/attention/ops/__init__.py +++ b/diffulex_legacy/layers/attention/ops/__init__.py @@ -1,7 +1,7 @@ -from diffulex.legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd as causal_lm_flash_decoding -from diffulex.legacy.layers.attention.ops.triton_decode_attn_dlm import diffusion_lm_flash_decoding, CHECK_ATTENTION -from diffulex.legacy.layers.attention.ops.chunked_prefill_decoding_unified_kernel import chunked_prefill_paged_decode as diffusion_lm_parallel_flash_decoding -from diffulex.legacy.layers.attention.ops.kv_cache_kernels import ( +from diffulex_legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd as causal_lm_flash_decoding +from diffulex_legacy.layers.attention.ops.triton_decode_attn_dlm import diffusion_lm_flash_decoding, CHECK_ATTENTION +from diffulex_legacy.layers.attention.ops.chunked_prefill_decoding_unified_kernel import chunked_prefill_paged_decode as diffusion_lm_parallel_flash_decoding +from diffulex_legacy.layers.attention.ops.kv_cache_kernels import ( store_kvcache_distinct_layout, store_kvcache_unified_layout, load_kvcache, CHECK_STORING, CHECK_LOADING ) \ No newline at end of file diff --git a/diffulex/attention/ops/chunked_prefill_decoding_unified_kernel.py b/diffulex_legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py similarity index 99% rename from diffulex/attention/ops/chunked_prefill_decoding_unified_kernel.py rename to diffulex_legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py index aed7e060..7c863c49 100755 --- a/diffulex/attention/ops/chunked_prefill_decoding_unified_kernel.py +++ b/diffulex_legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py @@ -18,7 +18,7 @@ from vllm.platforms.rocm import use_rocm_custom_paged_attention from vllm.triton_utils import tl, triton -from diffulex.legacy.layers.attention.ops.prefix_prefill import context_attention_fwd +from diffulex_legacy.layers.attention.ops.prefix_prefill import context_attention_fwd @triton.jit diff --git a/diffulex/legacy/layers/attention/ops/kv_cache_kernels.py b/diffulex_legacy/layers/attention/ops/kv_cache_kernels.py similarity index 99% rename from diffulex/legacy/layers/attention/ops/kv_cache_kernels.py rename to diffulex_legacy/layers/attention/ops/kv_cache_kernels.py index a62e2757..fcd6c223 100755 --- a/diffulex/legacy/layers/attention/ops/kv_cache_kernels.py +++ b/diffulex_legacy/layers/attention/ops/kv_cache_kernels.py @@ -6,8 +6,8 @@ from typing import Tuple from einops import rearrange -from diffulex.legacy.utils.context import ContextForDiffusionLM -from diffulex.legacy.engine.sequence import SequenceForDiffusionLM +from diffulex_legacy.utils.context import ContextForDiffusionLM +from diffulex_legacy.engine.sequence import SequenceForDiffusionLM @triton.jit def store_kvcache_kernel_causal_lm( diff --git a/diffulex/attention/ops/prefix_prefill.py b/diffulex_legacy/layers/attention/ops/prefix_prefill.py similarity index 100% rename from diffulex/attention/ops/prefix_prefill.py rename to diffulex_legacy/layers/attention/ops/prefix_prefill.py diff --git a/diffulex/attention/ops/tilus_decode_attn_dlm.py b/diffulex_legacy/layers/attention/ops/tilus_decode_attn_dlm.py similarity index 100% rename from diffulex/attention/ops/tilus_decode_attn_dlm.py rename to diffulex_legacy/layers/attention/ops/tilus_decode_attn_dlm.py diff --git a/diffulex/attention/ops/triton_decode_attn_clm.py b/diffulex_legacy/layers/attention/ops/triton_decode_attn_clm.py similarity index 100% rename from diffulex/attention/ops/triton_decode_attn_clm.py rename to diffulex_legacy/layers/attention/ops/triton_decode_attn_clm.py diff --git a/diffulex/legacy/layers/attention/ops/triton_decode_attn_dlm.py b/diffulex_legacy/layers/attention/ops/triton_decode_attn_dlm.py similarity index 97% rename from diffulex/legacy/layers/attention/ops/triton_decode_attn_dlm.py rename to diffulex_legacy/layers/attention/ops/triton_decode_attn_dlm.py index e39ed1e0..00706874 100755 --- a/diffulex/legacy/layers/attention/ops/triton_decode_attn_dlm.py +++ b/diffulex_legacy/layers/attention/ops/triton_decode_attn_dlm.py @@ -12,7 +12,7 @@ import triton.language as tl -from diffulex.legacy.utils.context import ContextForDiffusionLM +from diffulex_legacy.utils.context import ContextForDiffusionLM def CHECK_ATTENTION(o: torch.Tensor, q: torch.Tensor, k_new: torch.Tensor, v_new: torch.Tensor, @@ -24,7 +24,7 @@ def CHECK_ATTENTION(o: torch.Tensor, q: torch.Tensor, k_new: torch.Tensor, v_new from torch.nn.functional import scaled_dot_product_attention as sdpa from torch.nn.attention import SDPBackend, sdpa_kernel - from diffulex.legacy.layers.attention.ops import load_kvcache + from diffulex_legacy.layers.attention.ops import load_kvcache torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False diff --git a/diffulex/attention/ops/triton_flash_attention.py b/diffulex_legacy/layers/attention/ops/triton_flash_attention.py similarity index 100% rename from diffulex/attention/ops/triton_flash_attention.py rename to diffulex_legacy/layers/attention/ops/triton_flash_attention.py diff --git a/diffulex/legacy/layers/embed_head.py b/diffulex_legacy/layers/embed_head.py similarity index 97% rename from diffulex/legacy/layers/embed_head.py rename to diffulex_legacy/layers/embed_head.py index b781b2d1..1c49bbb1 100755 --- a/diffulex/legacy/layers/embed_head.py +++ b/diffulex_legacy/layers/embed_head.py @@ -4,7 +4,7 @@ import torch.nn.functional as F import torch.distributed as dist -from diffulex.legacy.utils.context import get_context_causal_lm, get_context_diffusion_lm +from diffulex_legacy.utils.context import get_context_causal_lm, get_context_diffusion_lm class VocabParallelEmbedding(nn.Module): diff --git a/diffulex/legacy/layers/layernorm.py b/diffulex_legacy/layers/layernorm.py similarity index 100% rename from diffulex/legacy/layers/layernorm.py rename to diffulex_legacy/layers/layernorm.py diff --git a/diffulex/legacy/layers/linear.py b/diffulex_legacy/layers/linear.py similarity index 100% rename from diffulex/legacy/layers/linear.py rename to diffulex_legacy/layers/linear.py diff --git a/diffulex/legacy/layers/rotary_embedding.py b/diffulex_legacy/layers/rotary_embedding.py similarity index 100% rename from diffulex/legacy/layers/rotary_embedding.py rename to diffulex_legacy/layers/rotary_embedding.py diff --git a/diffulex/legacy/layers/sampler.py b/diffulex_legacy/layers/sampler.py similarity index 98% rename from diffulex/legacy/layers/sampler.py rename to diffulex_legacy/layers/sampler.py index fe7bb758..0c4ef21a 100644 --- a/diffulex/legacy/layers/sampler.py +++ b/diffulex_legacy/layers/sampler.py @@ -8,8 +8,8 @@ from dataclasses import dataclass from easydict import EasyDict as edict -from diffulex.legacy.config import Config -from diffulex.legacy.utils.context import get_context_diffusion_lm +from diffulex_legacy.config import Config +from diffulex_legacy.utils.context import get_context_diffusion_lm class SamplerForCausalLM(nn.Module): @@ -122,6 +122,7 @@ def forward(self, logits: torch.Tensor, temperatures: torch.Tensor, true_local_ids_sub_map = {} accepted_ids_sub_map = {} sampled_tokens_sub_map = {} + shifted_logits = self._shift_logits(seq_logits, seq.cached_or_caching_last_token_id) for block_id, block in enumerate(seq.diffusion_blocks): if not block.is_active or sum(block.local_mask_tokens) == 0: diff --git a/diffulex/legacy/llm.py b/diffulex_legacy/llm.py similarity index 65% rename from diffulex/legacy/llm.py rename to diffulex_legacy/llm.py index c519d7a1..69c1f098 100755 --- a/diffulex/legacy/llm.py +++ b/diffulex_legacy/llm.py @@ -1,6 +1,6 @@ -from diffulex.legacy.engine.llm_engine import LLMEngine -from diffulex.legacy.engine.dp_engine import DPEngine -from diffulex.legacy.config import Config +from diffulex_legacy.engine.llm_engine import LLMEngine +from diffulex_legacy.engine.dp_engine import DPEngine +from diffulex_legacy.config import Config class LLM: def __new__(cls, model, **kwargs): diff --git a/diffulex/legacy/models/auto_model.py b/diffulex_legacy/models/auto_model.py similarity index 59% rename from diffulex/legacy/models/auto_model.py rename to diffulex_legacy/models/auto_model.py index 185aa0b0..0b6c65db 100755 --- a/diffulex/legacy/models/auto_model.py +++ b/diffulex_legacy/models/auto_model.py @@ -1,8 +1,8 @@ -from diffulex.legacy.config import Config -from diffulex.legacy.utils.loader import load_model -from diffulex.legacy.models.dream import DreamForDiffusionLM -from diffulex.legacy.models.qwen3 import Qwen3ForCausalLM -from diffulex.legacy.models.llada import LLaDAForDiffusionLM +from diffulex_legacy.config import Config +from diffulex_legacy.utils.loader import load_model +from diffulex_legacy.models.dream import DreamForDiffusionLM +from diffulex_legacy.models.qwen3 import Qwen3ForCausalLM +from diffulex_legacy.models.llada import LLaDAForDiffusionLM class AutoModelLM: diff --git a/diffulex/legacy/models/config/dream/configuration_dream.py b/diffulex_legacy/models/config/dream/configuration_dream.py similarity index 100% rename from diffulex/legacy/models/config/dream/configuration_dream.py rename to diffulex_legacy/models/config/dream/configuration_dream.py diff --git a/diffulex/legacy/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py b/diffulex_legacy/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py similarity index 100% rename from diffulex/legacy/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py rename to diffulex_legacy/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py diff --git a/diffulex/legacy/models/config/llada/configuration_llada.py b/diffulex_legacy/models/config/llada/configuration_llada.py similarity index 100% rename from diffulex/legacy/models/config/llada/configuration_llada.py rename to diffulex_legacy/models/config/llada/configuration_llada.py diff --git a/diffulex/legacy/models/dream.py b/diffulex_legacy/models/dream.py similarity index 94% rename from diffulex/legacy/models/dream.py rename to diffulex_legacy/models/dream.py index 4f5bd36a..7f197a6e 100755 --- a/diffulex/legacy/models/dream.py +++ b/diffulex_legacy/models/dream.py @@ -3,13 +3,13 @@ import torch.nn as nn import torch.distributed as dist -from diffulex.legacy.layers.layernorm import RMSNorm -from diffulex.legacy.layers.activation import SiluAndMul -from diffulex.legacy.layers.rotary_embedding import get_rope -from diffulex.legacy.layers.attention.attention_v5 import Attention -from diffulex.legacy.models.config.dream.configuration_dream import DreamConfig -from diffulex.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear -from diffulex.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffulex_legacy.layers.layernorm import RMSNorm +from diffulex_legacy.layers.activation import SiluAndMul +from diffulex_legacy.layers.rotary_embedding import get_rope +from diffulex_legacy.layers.attention.attention_v5 import Attention +from diffulex_legacy.models.config.dream.configuration_dream import DreamConfig +from diffulex_legacy.layers.linear import RowParallelLinear, ColumnParallelLinear +from diffulex_legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead diff --git a/diffulex/legacy/models/fast_dllm_v2.py b/diffulex_legacy/models/fast_dllm_v2.py similarity index 94% rename from diffulex/legacy/models/fast_dllm_v2.py rename to diffulex_legacy/models/fast_dllm_v2.py index 4739b412..ac905b82 100755 --- a/diffulex/legacy/models/fast_dllm_v2.py +++ b/diffulex_legacy/models/fast_dllm_v2.py @@ -3,13 +3,13 @@ import torch.nn as nn import torch.distributed as dist -from diffulex.legacy.layers.layernorm import RMSNorm -from diffulex.legacy.layers.activation import SiluAndMul -from diffulex.legacy.layers.rotary_embedding import get_rope -from diffulex.legacy.layers.attention.attention_v5 import Attention -from diffulex.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear -from diffulex.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead -from diffulex.legacy.models.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config +from diffulex_legacy.layers.layernorm import RMSNorm +from diffulex_legacy.layers.activation import SiluAndMul +from diffulex_legacy.layers.rotary_embedding import get_rope +from diffulex_legacy.layers.attention.attention_v5 import Attention +from diffulex_legacy.layers.linear import RowParallelLinear, ColumnParallelLinear +from diffulex_legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffulex_legacy.models.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/diffulex/legacy/models/llada.py b/diffulex_legacy/models/llada.py similarity index 95% rename from diffulex/legacy/models/llada.py rename to diffulex_legacy/models/llada.py index 342a1c57..ff8a491a 100755 --- a/diffulex/legacy/models/llada.py +++ b/diffulex_legacy/models/llada.py @@ -3,13 +3,13 @@ import torch.nn as nn import torch.distributed as dist -from diffulex.legacy.layers.layernorm import RMSNorm -from diffulex.legacy.layers.activation import SiluAndMul -from diffulex.legacy.layers.rotary_embedding import get_rope -from diffulex.legacy.layers.attention.attention_v5 import Attention -from diffulex.legacy.models.config.llada.configuration_llada import LLaDAConfig -from diffulex.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear -from diffulex.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffulex_legacy.layers.layernorm import RMSNorm +from diffulex_legacy.layers.activation import SiluAndMul +from diffulex_legacy.layers.rotary_embedding import get_rope +from diffulex_legacy.layers.attention.attention_v5 import Attention +from diffulex_legacy.models.config.llada.configuration_llada import LLaDAConfig +from diffulex_legacy.layers.linear import RowParallelLinear, ColumnParallelLinear +from diffulex_legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/diffulex/legacy/models/qwen3.py b/diffulex_legacy/models/qwen3.py similarity index 94% rename from diffulex/legacy/models/qwen3.py rename to diffulex_legacy/models/qwen3.py index f6803d9c..ddded093 100755 --- a/diffulex/legacy/models/qwen3.py +++ b/diffulex_legacy/models/qwen3.py @@ -4,12 +4,12 @@ from transformers import Qwen3Config -from diffulex.legacy.layers.layernorm import RMSNorm -from diffulex.legacy.layers.activation import SiluAndMul -from diffulex.legacy.layers.rotary_embedding import get_rope -from diffulex.legacy.layers.attention.attention_v4 import Attention -from diffulex.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead -from diffulex.legacy.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear +from diffulex_legacy.layers.layernorm import RMSNorm +from diffulex_legacy.layers.activation import SiluAndMul +from diffulex_legacy.layers.rotary_embedding import get_rope +from diffulex_legacy.layers.attention.attention_v4 import Attention +from diffulex_legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffulex_legacy.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear class Qwen3Attention(nn.Module): diff --git a/diffulex/legacy/models/utils/check_config.py b/diffulex_legacy/models/utils/check_config.py similarity index 100% rename from diffulex/legacy/models/utils/check_config.py rename to diffulex_legacy/models/utils/check_config.py diff --git a/diffulex/legacy/sampling_params.py b/diffulex_legacy/sampling_params.py similarity index 100% rename from diffulex/legacy/sampling_params.py rename to diffulex_legacy/sampling_params.py diff --git a/diffulex/legacy/utils/checker.py b/diffulex_legacy/utils/checker.py similarity index 100% rename from diffulex/legacy/utils/checker.py rename to diffulex_legacy/utils/checker.py diff --git a/diffulex/legacy/utils/context.py b/diffulex_legacy/utils/context.py similarity index 98% rename from diffulex/legacy/utils/context.py rename to diffulex_legacy/utils/context.py index 7d49ea38..89862763 100755 --- a/diffulex/legacy/utils/context.py +++ b/diffulex_legacy/utils/context.py @@ -3,7 +3,7 @@ from typing import List from dataclasses import dataclass -from diffulex.legacy.engine.sequence import SequenceForDiffusionLM +from diffulex_legacy.engine.sequence import SequenceForDiffusionLM @dataclass class ContextBase: diff --git a/diffulex/legacy/utils/loader.py b/diffulex_legacy/utils/loader.py similarity index 99% rename from diffulex/legacy/utils/loader.py rename to diffulex_legacy/utils/loader.py index 5dd07bd3..733898b1 100755 --- a/diffulex/legacy/utils/loader.py +++ b/diffulex_legacy/utils/loader.py @@ -7,7 +7,7 @@ from glob import glob from functools import partial from safetensors import safe_open -from diffulex.legacy.config import Config +from diffulex_legacy.config import Config def load_lora_config(lora_path: str) -> dict: diff --git a/examples/ops/prefix_prefill.py b/examples/ops/prefix_prefill.py new file mode 100644 index 00000000..079fc744 --- /dev/null +++ b/examples/ops/prefix_prefill.py @@ -0,0 +1,814 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# The kernels in this file are adapted from LightLLM's context_attention_fwd: +# https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +# Static kernels parameters +BASE_BLOCK = 128 if current_platform.has_device_capability(80) else 64 +NUM_WARPS = 4 if current_platform.is_rocm() else 8 + +# To check compatibility +IS_TURING = current_platform.get_device_capability() == (7, 5) +float8_info = torch.finfo(current_platform.fp8_dtype()) + + +# Here's an example autotuner config for this kernel. This config does provide +# a performance improvement, but dramatically increases first call latency in +# triton 3.2. Because of this tradeoff, it's currently commented out. +# @triton.autotune( +# configs=[ +# triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, \ +# "num_unroll_cache": 4, \ +# "num_unroll_request": 1 } | \ +# ({"kpack": 2, "waves_per_eu": 2} \ +# if current_platform.is_rocm() else {}), \ +# num_warps=4, \ +# num_stages=1) +# ], +# key=["BLOCK_SIZE", "MAX_Q_LEN", "MAX_CTX_LEN"] +# ) +@triton.jit +def _fwd_kernel( + Q, + K, + V, + K_cache, + V_cache, + sink_ptr, + B_Loc, + sm_scale, + k_scale, + v_scale, + out_scale_inv, + B_Start_Loc, + B_Seqlen, + x: tl.constexpr, + Out, + stride_b_loc_b, + stride_b_loc_s, + stride_qbs, + stride_qh, + stride_qd, + stride_kbs, + stride_kh, + stride_kd, + stride_vbs, + stride_vh, + stride_vd, + stride_obs, + stride_oh, + stride_od, + stride_k_cache_bs, + stride_k_cache_h, + stride_k_cache_d, + stride_k_cache_bl: tl.constexpr, + stride_k_cache_x, + stride_v_cache_bs, + stride_v_cache_h, + stride_v_cache_d, + stride_v_cache_bl, + num_queries_per_kv: tl.constexpr, + IN_PRECISION: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DMODEL_PADDED: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, + SLIDING_WINDOW: tl.constexpr, + num_unroll_cache: tl.constexpr, + num_unroll_request: tl.constexpr, + SKIP_DECODE: tl.constexpr, + USE_SINKS: tl.constexpr, + USE_FP8: tl.constexpr, + MAX_Q_LEN: tl.constexpr = 0, + MAX_CTX_LEN: tl.constexpr = 0, + FP8_MIN: tl.constexpr = float8_info.min, + FP8_MAX: tl.constexpr = float8_info.max, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + start_m = tl.program_id(2) + + cur_kv_head = cur_head // num_queries_per_kv + + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) + cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1) + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len + + if SKIP_DECODE and cur_batch_query_len == 1: + return + + # start position inside of the query + # generally, N goes over kv, while M goes over query_len + block_start_loc = BLOCK_M * start_m + + # initialize offsets + # [BLOCK_SIZE]; starts at 0 + offs_bs_n = tl.arange(0, BLOCK_SIZE) + # [N]; starts at 0 + offs_n = tl.arange(0, BLOCK_N) + # [D]; starts at 0 + offs_d = tl.arange(0, BLOCK_DMODEL_PADDED) + # [M]; starts at current position in query + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + # [M,D] + off_q = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] * stride_qd + ) + + dim_mask = tl.where(tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to( + tl.int1 + ) # [D] + + q = tl.load( + Q + off_q, + mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len), + other=0.0, + ) # [M,D] + + # initialize pointer to m and l + if not USE_SINKS: + m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + else: + m_i = tl.load( + sink_ptr + tl.full([BLOCK_M], cur_head, dtype=tl.int64), + mask=(offs_m < cur_batch_query_len), + other=float("-inf"), + ).to(dtype=tl.float32) + + l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) # [M,D] + + # compute query against context (no causal mask here) + for start_n in tl.range( + 0, cur_batch_ctx_len, BLOCK_SIZE, loop_unroll_factor=num_unroll_cache + ): + start_n = tl.multiple_of(start_n, BLOCK_SIZE) + # -- compute qk ---- + bn = tl.load( + B_Loc + + cur_batch * stride_b_loc_b + + (start_n // BLOCK_SIZE) * stride_b_loc_s + ).to(tl.int64) + # [D,BLOCK_SIZE] + off_k = ( + bn[None, :] * stride_k_cache_bs + + cur_kv_head * stride_k_cache_h + + (offs_d[:, None] // x) * stride_k_cache_d + + ((start_n + offs_bs_n[None, :]) % BLOCK_SIZE) * stride_k_cache_bl + + (offs_d[:, None] % x) * stride_k_cache_x + ) + + # [BLOCK_SIZE,D] + off_v = ( + bn[:, None] * stride_v_cache_bs + + cur_kv_head * stride_v_cache_h + + offs_d[None, :] * stride_v_cache_d + + offs_bs_n[:, None] * stride_v_cache_bl + ) + + if ( + start_n + BLOCK_SIZE > cur_batch_ctx_len + or BLOCK_DMODEL != BLOCK_DMODEL_PADDED + ): + k_load = tl.load( + K_cache + off_k, + mask=dim_mask[:, None] + & ((start_n + offs_bs_n[None, :]) < cur_batch_ctx_len), + other=0.0, + ) # [D,N] + else: + k_load = tl.load(K_cache + off_k) + + if k_load.dtype.is_fp8(): + k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype) + else: + k = k_load + + qk = tl.zeros([BLOCK_M, BLOCK_SIZE], dtype=tl.float32) # [M,N] + qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) + qk = tl.where( + (start_n + offs_bs_n[None, :]) < cur_batch_ctx_len, qk, float("-inf") + ) + qk *= sm_scale + if SLIDING_WINDOW > 0: + # (cur_batch_ctx_len + offs_m[:, None]) are the positions of + # Q entries in sequence + # (start_n + offs_bs_n[None, :]) are the positions of + # KV entries in sequence + # So the condition makes sure each entry in Q only attends + # to KV entries not more than SLIDING_WINDOW away. + # + # We can't use -inf here, because the + # sliding window may lead to the entire row being masked. + # This then makes m_ij contain -inf, which causes NaNs in + # exp(). + qk = tl.where( + (cur_batch_ctx_len + offs_m[:, None]) - (start_n + offs_bs_n[None, :]) + < SLIDING_WINDOW, + qk, + -10000, + ) + + # compute running maximum + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + alpha = tl.exp(m_i - m_ij) + acc = acc * alpha[:, None] + + # update acc + if ( + start_n + BLOCK_SIZE > cur_batch_ctx_len + or BLOCK_DMODEL != BLOCK_DMODEL_PADDED + ): + v_load = tl.load( + V_cache + off_v, + mask=dim_mask[None, :] + & ((start_n + offs_bs_n[:, None]) < cur_batch_ctx_len), + other=0.0, + ) # [N,D] + else: + v_load = tl.load(V_cache + off_v) + + if v_load.dtype.is_fp8(): + v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype) + else: + v = v_load + p = p.to(v.dtype) + + acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION) + # # update m_i and l_i + l_i = l_i * alpha + l_ij + m_i = m_ij + + off_k = ( + offs_n[None, :] * stride_kbs + + cur_kv_head * stride_kh + + offs_d[:, None] * stride_kd + ) + off_v = ( + offs_n[:, None] * stride_vbs + + cur_kv_head * stride_vh + + offs_d[None, :] * stride_vd + ) + k_ptrs = K + off_k + v_ptrs = V + off_v + + # block_mask is 0 when we're already past the current query length + block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0) + + # compute query against itself (with causal mask) + for start_n in tl.range( + 0, + block_mask * (start_m + 1) * BLOCK_M, + BLOCK_N, + loop_unroll_factor=num_unroll_request, + ): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = tl.load( + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, + mask=dim_mask[:, None] + & ((start_n + offs_n[None, :]) < cur_batch_query_len), + other=0.0, + ) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) + qk *= sm_scale + # apply causal mask + qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) + if SLIDING_WINDOW > 0: + qk = tl.where( + offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW, + qk, + -10000, + ) + + # compute running maximum + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + alpha = tl.exp(m_i - m_ij) + acc = acc * alpha[:, None] + + # update acc + v = tl.load( + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, + mask=dim_mask[None, :] + & ((start_n + offs_n[:, None]) < cur_batch_query_len), + other=0.0, + ) + p = p.to(v.dtype) + + acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION) + # update m_i and l_i + l_i = l_i * alpha + l_ij + m_i = m_ij + + acc = acc / l_i[:, None] + + # initialize pointers to output + off_o = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + + cur_head * stride_oh + + offs_d[None, :] * stride_od + ) + out_ptrs = Out + off_o + if USE_FP8: + acc = acc * tl.load(out_scale_inv) + acc = tl.clamp(acc, FP8_MIN, FP8_MAX) + tl.store( + out_ptrs, acc, mask=dim_mask[None, :] & (offs_m[:, None] < cur_batch_query_len) + ) + return + + +@triton.jit +def _fwd_kernel_alibi( + Q, + K, + V, + K_cache, + V_cache, + B_Loc, + sm_scale, + k_scale, + v_scale, + B_Start_Loc, + B_Seqlen, + Alibi_slopes, + block_size, + x, + Out, + stride_b_loc_b, + stride_b_loc_s, + stride_qbs, + stride_qh, + stride_qd, + stride_kbs, + stride_kh, + stride_kd, + stride_vbs, + stride_vh, + stride_vd, + stride_obs, + stride_oh, + stride_od, + stride_k_cache_bs, + stride_k_cache_h, + stride_k_cache_d, + stride_k_cache_bl, + stride_k_cache_x, + stride_v_cache_bs, + stride_v_cache_h, + stride_v_cache_d, + stride_v_cache_bl, + num_queries_per_kv: int, + IN_PRECISION: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, # head size + BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2 + BLOCK_N: tl.constexpr, + SKIP_DECODE: tl.constexpr, +): + # attn_bias[] + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + start_m = tl.program_id(2) + + cur_kv_head = cur_head // num_queries_per_kv + + # cur_batch_seq_len: the length of prompts + # cur_batch_ctx_len: the length of prefix + # cur_batch_in_all_start_index: the start id of the dim=0 + cur_batch_seq_len = tl.load(B_Seqlen + cur_batch) + cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch) + cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1) + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len + + if SKIP_DECODE and cur_batch_query_len == 1: + return + + block_start_loc = BLOCK_M * start_m + + # initialize offsets + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL_PADDED) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + off_q = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] * stride_qd + ) + + dim_mask = tl.where(tl.arange(0, BLOCK_DMODEL_PADDED) < BLOCK_DMODEL, 1, 0).to( + tl.int1 + ) + + q = tl.load( + Q + off_q, + mask=dim_mask[None, :] + & (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len), + other=0.0, + ) + + # # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) + + alibi_slope = tl.load(Alibi_slopes + cur_head) + alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len + alibi_start_k = 0 + for start_n in range(0, cur_batch_ctx_len, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + bn = tl.load( + B_Loc + + cur_batch * stride_b_loc_b + + ((start_n + offs_n) // block_size) * stride_b_loc_s, + mask=(start_n + offs_n) < cur_batch_ctx_len, + other=0, + ).to(tl.int64) + off_k = ( + bn[None, :] * stride_k_cache_bs + + cur_kv_head * stride_k_cache_h + + (offs_d[:, None] // x) * stride_k_cache_d + + ((start_n + offs_n[None, :]) % block_size) * stride_k_cache_bl + + (offs_d[:, None] % x) * stride_k_cache_x + ) + off_v = ( + bn[:, None] * stride_v_cache_bs + + cur_kv_head * stride_v_cache_h + + offs_d[None, :] * stride_v_cache_d + + (start_n + offs_n[:, None]) % block_size * stride_v_cache_bl + ) + k_load = tl.load( + K_cache + off_k, + mask=dim_mask[:, None] & ((start_n + offs_n[None, :]) < cur_batch_ctx_len), + other=0.0, + ) # [D,N] + + if k_load.dtype.is_fp8(): + k = (k_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype) + else: + k = k_load + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION) + qk = tl.where( + (start_n + offs_n[None, :]) < cur_batch_ctx_len, qk, float("-inf") + ) + qk *= sm_scale + + # load alibi + alibi = ( + tl.arange(0, BLOCK_N)[None, :] + alibi_start_k - alibi_start_q[:, None] + ) * alibi_slope + alibi = tl.where( + (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len), + alibi, + float("-inf"), + ) + qk += alibi + alibi_start_k += BLOCK_N + + # -- compute m_ij, p, l_ij + m_ij = tl.max(qk, 1) + m_i_new = tl.maximum(m_i, m_ij) + p = tl.math.exp(qk - m_i_new[:, None]) + l_ij = tl.sum(p, 1) + # -- update m_i and l_i + + alpha = tl.math.exp(m_i - m_i_new) + l_i_new = alpha * l_i + l_ij + # -- update output accumulator -- + # scale p + # scale acc + acc_scale = alpha + # acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + # update acc + v_load = tl.load( + V_cache + off_v, + mask=dim_mask[None, :] & ((start_n + offs_n[:, None]) < cur_batch_ctx_len), + other=0.0, + ) + if v_load.dtype.is_fp8(): + v = (v_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype) + else: + v = v_load + p = p.to(v.dtype) + + acc = tl.dot(p, v, acc=acc, input_precision="ieee") + # update m_i and l_i + l_i = l_i_new + m_i = m_i_new + + off_k = ( + offs_n[None, :] * stride_kbs + + cur_kv_head * stride_kh + + offs_d[:, None] * stride_kd + ) + off_v = ( + offs_n[:, None] * stride_vbs + + cur_kv_head * stride_vh + + offs_d[None, :] * stride_vd + ) + k_ptrs = K + off_k + v_ptrs = V + off_v + + block_mask = tl.where(block_start_loc < cur_batch_seq_len - cur_batch_ctx_len, 1, 0) + + # init alibi + alibi_slope = tl.load(Alibi_slopes + cur_head) + alibi_start_q = tl.arange(0, BLOCK_M) + block_start_loc + cur_batch_ctx_len + alibi_start_k = cur_batch_ctx_len + # # init debugger + # offset_db_q = tl.arange(0, BLOCK_M) + block_start_loc + # offset_db_k = tl.arange(0, BLOCK_N) + # calc q[BLOCK_M, BLOCK_MODEL] mul k[prefix_len: , BLOCK_DMODEL] + for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = tl.load( + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs, + mask=dim_mask[:, None] + & ((start_n + offs_n[None, :]) < cur_batch_seq_len - cur_batch_ctx_len), + other=0.0, + ) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk = tl.dot(q, k, acc=qk, input_precision="ieee") + qk *= sm_scale + qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) + + # load alibi + alibi = ( + tl.arange(0, BLOCK_N)[None, :] + alibi_start_k - alibi_start_q[:, None] + ) * alibi_slope + alibi = tl.where( + (alibi <= 0) & (alibi_start_q[:, None] < cur_batch_seq_len), + alibi, + float("-inf"), + ) + qk += alibi + alibi_start_k += BLOCK_N + + # -- compute m_ij, p, l_ij + m_ij = tl.max(qk, 1) + m_i_new = tl.maximum(m_i, m_ij) + p = tl.math.exp(qk - m_i_new[:, None]) + l_ij = tl.sum(p, 1) + # -- update m_i and l_i + + alpha = tl.math.exp(m_i - m_i_new) + l_i_new = alpha * l_i + l_ij + # -- update output accumulator -- + # scale p + # scale acc + acc_scale = alpha + # acc_scale = l_i / l_i_new * alpha + acc = acc * acc_scale[:, None] + # update acc + v = tl.load( + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs, + mask=dim_mask[None, :] + & ((start_n + offs_n[:, None]) < cur_batch_seq_len - cur_batch_ctx_len), + other=0.0, + ) + p = p.to(v.dtype) + + acc = tl.dot(p, v, acc=acc, input_precision="ieee") + # update m_i and l_i + l_i = l_i_new + m_i = m_i_new + + acc = acc / l_i[:, None] + + # initialize pointers to output + off_o = ( + (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs + + cur_head * stride_oh + + offs_d[None, :] * stride_od + ) + out_ptrs = Out + off_o + tl.store( + out_ptrs, + acc, + mask=dim_mask[None, :] + & (offs_m[:, None] < cur_batch_seq_len - cur_batch_ctx_len), + ) + return + + +@torch.inference_mode() +def context_attention_fwd( + q, + k, + v, + o, + kv_cache_dtype: str, + k_cache, + v_cache, + b_loc, + b_start_loc, + b_seq_len, + max_seq_len, + max_input_len, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + alibi_slopes=None, + sliding_window=None, + sm_scale=None, + skip_decode=False, + fp8_out_scale=None, + sinks=None, +): + q_dtype_is_f32 = q.dtype is torch.float32 + + # Turing does have tensor core for float32 multiplication + # use ieee as fallback for triton kernels work. There is also + # warning on vllm/config.py to inform users this fallback + # implementation + IN_PRECISION = "ieee" if IS_TURING and q_dtype_is_f32 else None + + # Conversion of FP8 Tensor from uint8 storage to + # appropriate torch.dtype for interpretation by Triton + if "fp8" in kv_cache_dtype: + assert k_cache.dtype in [torch.uint8, current_platform.fp8_dtype()] + assert v_cache.dtype in [torch.uint8, current_platform.fp8_dtype()] + + if kv_cache_dtype in ("fp8", "fp8_e4m3"): + target_dtype = current_platform.fp8_dtype() + elif kv_cache_dtype == "fp8_e5m2": + target_dtype = torch.float8_e5m2 + else: + raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype) + + k_cache = k_cache.view(target_dtype) + v_cache = v_cache.view(target_dtype) + + if ( + k_cache.dtype == torch.uint8 + or v_cache.dtype == torch.uint8 + and kv_cache_dtype == "auto" + ): + raise ValueError( + "kv_cache_dtype='auto' unsupported for\ + FP8 KV Cache prefill kernel" + ) + + # shape constraints + Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] + assert Lq == Lk and Lk == Lv + # round up Lk to a power of 2 - this is required for Triton block size + Lk_padded = triton.next_power_of_2(Lk) + + if sm_scale is None: + sm_scale = 1.0 / (Lq**0.5) + batch, head = b_seq_len.shape[0], q.shape[1] + num_queries_per_kv = q.shape[1] // k.shape[1] + + assert batch + 1 == len(b_start_loc) + + # 0 means "disable" + if sliding_window is None or sliding_window <= 0: + sliding_window = 0 + + if alibi_slopes is not None: + assert sinks is None, "Sinks arg is not supported with alibi" + assert fp8_out_scale is None, "FP8 output not supported with alibi" + # need to reduce num. blocks when using fp32 + # due to increased use of GPU shared memory + # if q.dtype is torch.float32: + BLOCK = BASE_BLOCK // 2 if q_dtype_is_f32 else BASE_BLOCK + # batch, head, + grid = (batch, head, triton.cdiv(max_input_len, BLOCK)) + _fwd_kernel_alibi[grid]( + q, + k, + v, + k_cache, + v_cache, + b_loc, + sm_scale, + k_scale, + v_scale, + b_start_loc, + b_seq_len, + alibi_slopes, + v_cache.shape[3], + k_cache.shape[4], + o, + b_loc.stride(0), + b_loc.stride(1), + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + o.stride(0), + o.stride(1), + o.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + k_cache.stride(3), + k_cache.stride(4), # [num_blocks, num_kv_heads, head_size/x, block_size, x] + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + v_cache.stride(3), # [num_blocks, num_kv_heads, head_size, block_size] + num_queries_per_kv=num_queries_per_kv, + IN_PRECISION=IN_PRECISION, + BLOCK_M=BLOCK, + BLOCK_DMODEL=Lk, + BLOCK_DMODEL_PADDED=Lk_padded, + BLOCK_N=BLOCK, + SKIP_DECODE=skip_decode, + num_warps=NUM_WARPS, + num_stages=1, + ) + return + + max_seq_len = 0 if max_seq_len is None else max_seq_len + extra_kargs = {} + if current_platform.is_rocm(): + extra_kargs = {"kpack": 1, "waves_per_eu": 2} + + grid = lambda META: (batch, head, triton.cdiv(max_input_len, META["BLOCK_M"])) + _fwd_kernel[grid]( + q, + k, + v, + k_cache, + v_cache, + sinks, + b_loc, + sm_scale, + k_scale, + v_scale, + 1.0 / fp8_out_scale if fp8_out_scale is not None else 1.0, + b_start_loc, + b_seq_len, + k_cache.shape[4], + o, + b_loc.stride(0), + b_loc.stride(1), + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + o.stride(0), + o.stride(1), + o.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + k_cache.stride(3), + k_cache.stride(4), # [num_blocks, num_kv_heads, head_size/x, block_size, x] + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + v_cache.stride(3), # [num_blocks, num_kv_heads, head_size, block_size] + BLOCK_SIZE=v_cache.shape[3], + num_queries_per_kv=num_queries_per_kv, + IN_PRECISION=IN_PRECISION, + BLOCK_DMODEL=Lk, + BLOCK_DMODEL_PADDED=Lk_padded, + SLIDING_WINDOW=sliding_window, + SKIP_DECODE=skip_decode, + USE_FP8=fp8_out_scale is not None, + BLOCK_M=128, + BLOCK_N=64, + num_unroll_cache=4, + num_unroll_request=1, + num_warps=4, + num_stages=1, + USE_SINKS=sinks is not None, + **extra_kargs, + ) + return \ No newline at end of file diff --git a/examples/test_causal_lm_decoding_kernel.py b/examples/test_causal_lm_decoding_kernel.py index f5f0e836..4c7e440f 100755 --- a/examples/test_causal_lm_decoding_kernel.py +++ b/examples/test_causal_lm_decoding_kernel.py @@ -1,6 +1,6 @@ import torch -from diffulex.legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd +from diffulex_legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd if __name__ == "__main__": torch.random.manual_seed(114514) diff --git a/examples/test_dllm_decoding_kernel.py b/examples/test_dllm_decoding_kernel.py index c91925dd..4c6178a4 100755 --- a/examples/test_dllm_decoding_kernel.py +++ b/examples/test_dllm_decoding_kernel.py @@ -5,7 +5,7 @@ from einops import rearrange from torch.nn.functional import scaled_dot_product_attention -from diffulex.legacy.layers.attention.ops import diffusion_lm_parallel_flash_decoding, diffusion_lm_flash_decoding +from diffulex_legacy.layers.attention.ops import diffusion_lm_parallel_flash_decoding, diffusion_lm_flash_decoding if __name__ == "__main__": diff --git a/examples/test_dllm_kv_cache_load.py b/examples/test_dllm_kv_cache_load.py index 6096ba1c..80d1616c 100755 --- a/examples/test_dllm_kv_cache_load.py +++ b/examples/test_dllm_kv_cache_load.py @@ -4,7 +4,7 @@ from dataclasses import dataclass from mimic_data.mimic_slot_mapping import slot_mapping -from diffulex.legacy.layers.attention.ops import store_kvcache_unified_layout, load_kvcache, CHECK_LOADING +from diffulex_legacy.layers.attention.ops import store_kvcache_unified_layout, load_kvcache, CHECK_LOADING @dataclass class MimicSequenceForDiffusionLM: diff --git a/examples/test_dllm_kv_cache_store.py b/examples/test_dllm_kv_cache_store.py index 7ee9c1b7..b2b41130 100755 --- a/examples/test_dllm_kv_cache_store.py +++ b/examples/test_dllm_kv_cache_store.py @@ -3,7 +3,7 @@ from einops import rearrange -from diffulex.legacy.layers.attention.attention_v4 import store_kvcache_distinct_layout, store_kvcache_unified +from diffulex_legacy.layers.attention.attention_v4 import store_kvcache_distinct_layout, store_kvcache_unified if __name__ == "__main__": diff --git a/examples/test_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index 92880e22..1affedbc 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -58,12 +58,14 @@ def summarize_profiling(csv_path: str) -> dict: accept_threshold=0.95, complete_threshold=0.9, add_new_block_threshold=0.1, - kv_cache_layout="unified" + kv_cache_layout="unified", + decoding_strategy="d2f" ) tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) sampling_params = SamplingParams(temperature=0.0, max_tokens=256) - dataset = load_dataset("data/gsm8k", "main")['test']['question'][:] + dataset = load_dataset( + "gsm8k", "main", split="test")["question"][:10] prompts = [tokenizer.bos_token + FEW_SHOTS + p for p in tqdm(dataset)] output_file = "log/profiles/perf_dvllm_dream_7B.json" @@ -71,7 +73,7 @@ def summarize_profiling(csv_path: str) -> dict: os.remove(output_file) # with VizTracer(output_file=output_file, file_info=True) as tracer: # outputs = llm.generate(prompts[:5], sampling_params) - time.sleep(60) + # time.sleep(60) s = time.time() outputs = LLM.generate(prompts, sampling_params) e = time.time() diff --git a/examples/test_dream_model_weight.py b/examples/test_dream_model_weight.py index 8455c2b8..fb8ce12b 100755 --- a/examples/test_dream_model_weight.py +++ b/examples/test_dream_model_weight.py @@ -4,8 +4,8 @@ from peft import PeftModel, PeftConfig from lm_eval.models.utils import get_dtype -from diffulex.legacy.config import Config -from diffulex.legacy.models.auto_model import AutoModelLM +from diffulex_legacy.config import Config +from diffulex_legacy.models.auto_model import AutoModelLM from model_cache.dream.model_dream import DreamModel from model_cache.dream.configuration_dream import DreamConfig diff --git a/examples/test_dream_model_weight_fixed.py b/examples/test_dream_model_weight_fixed.py index a3d693e7..d09b8fe7 100755 --- a/examples/test_dream_model_weight_fixed.py +++ b/examples/test_dream_model_weight_fixed.py @@ -4,8 +4,8 @@ from peft import PeftModel, PeftConfig from lm_eval.models.utils import get_dtype -from diffulex.legacy.config import Config -from diffulex.legacy.engine.model_runner import AutoModelRunner +from diffulex_legacy.config import Config +from diffulex_legacy.engine.model_runner import AutoModelRunner from model_cache.dream.model_dream import DreamModel from model_cache.dream.configuration_dream import DreamConfig diff --git a/examples/test_fastdllmv2_diffulex_gsm8k.py b/examples/test_fastdllmv2_diffulex_gsm8k.py new file mode 100755 index 00000000..3950537c --- /dev/null +++ b/examples/test_fastdllmv2_diffulex_gsm8k.py @@ -0,0 +1,87 @@ +import os +import csv +import time + +import pandas as pd + +from tqdm import tqdm +from datasets import load_dataset +from viztracer import VizTracer +from transformers import AutoTokenizer + +from diffulex import Diffulex, SamplingParams + + +def summarize_profiling(csv_path: str) -> dict: + totals = {} + total_nums = {} + avgs = {} + with open(csv_path, 'r', newline='') as f: + reader = csv.DictReader(f) + for row in reader: + for k, v in row.items(): + try: + val = float(v) + except ValueError: + continue + if val != 0.0: + total_nums[k] = total_nums.get(k, 0) + 1 + totals[k] = totals.get(k, 0.0) + val + print(pd.DataFrame([totals]).T) + for k, v in totals.items(): + if k in total_nums and total_nums[k] > 0: + avgs[k] = v / total_nums[k] + else: + avgs[k] = 0.0 + print(pd.DataFrame([avgs]).T) + +# FEW_SHOTS = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nQuestion: Jen and Tyler are gymnasts practicing flips. Jen is practicing the triple-flip while Tyler is practicing the double-flip. Jen did sixteen triple-flips during practice. Tyler flipped in the air half the number of times Jen did. How many double-flips did Tyler do?\nAnswer:<|im_end|>\n<|im_start|>assistant\nJen did 16 triple-flips, so she did 16 * 3 = <<16*3=48>>48 flips.\nTyler did half the number of flips, so he did 48 / 2 = <<48/2=24>>24 flips.\nA double flip has two flips, so Tyler did 24 / 2 = <<24/2=12>>12 double-flips.\n#### 12<|im_end|>\n<|im_start|>user\nQuestion: Four people in a law firm are planning a party. Mary will buy a platter of pasta for $20 and a loaf of bread for $2. Elle and Andrea will split the cost for buying 4 cans of soda which cost $1.50 each, and chicken wings for $10. Joe will buy a cake that costs $5. How much more will Mary spend than the rest of the firm put together?\nAnswer:<|im_end|>\n<|im_start|>assistant\nMary will spend $20 + $2 = $<<20+2=22>>22.\nElle and Andrea will spend $1.5 x 4 = $<<1.5*4=6>>6 for the soda.\nElle and Andrea will spend $6 + $10 = $<<6+10=16>>16 for the soda and chicken wings.\nElle, Andrea, and Joe together will spend $16 + $5 = $<<16+5=21>>21.\nSo, Mary will spend $22 - $21 = $<<22-21=1>>1 more than all of them combined.\n#### 1<|im_end|>\n<|im_start|>user\nQuestion: A charcoal grill burns fifteen coals to ash every twenty minutes of grilling. The grill ran for long enough to burn three bags of coals. Each bag of coal contains 60 coals. How long did the grill run?\nAnswer:<|im_end|>\n<|im_start|>assistant\nThe grill burned 3 * 60 = <<3*60=180>>180 coals.\nIt takes 20 minutes to burn 15 coals, so the grill ran for 180 / 15 * 20 = <<180/15*20=240>>240 minutes.\n#### 240<|im_end|>\n<|im_start|>user\nQuestion: A bear is preparing to hibernate for the winter and needs to gain 1000 pounds. At the end of summer, the bear feasts on berries and small woodland animals. During autumn, it devours acorns and salmon. It gained a fifth of the weight it needed from berries during summer, and during autumn, it gained twice that amount from acorns. Salmon made up half of the remaining weight it had needed to gain. How many pounds did it gain eating small animals?\nAnswer:<|im_end|>\n<|im_start|>assistant\nThe bear gained 1 / 5 * 1000 = <<1/5*1000=200>>200 pounds from berries.\nIt gained 2 * 200 = <<2*200=400>>400 pounds from acorns.\nIt still needed 1000 - 200 - 400 = <<1000-200-400=400>>400 pounds.\nThus, it gained 400 / 2 = <<400/2=200>>200 pounds from salmon.\nTherefore, the bear gained 400 - 200 = <<400-200=200>>200 pounds from small animals.\n#### 200<|im_end|>\n<|im_start|>user\nQuestion: Janet’s ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?\nAnswer:<|im_end|>\n<|im_start|>assistant\n" +FEW_SHOTS = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + +if __name__ == "__main__": + model = "/data1/ckpts/Efficient-Large-Model/Fast_dLLM_v2_7B" + LLM = Diffulex( + model, + use_lora=False, + model_name="fast_dllm_v2", + enforce_eager=True, + data_parallel_size=1, + tensor_parallel_size=1, + gpu_memory_utilization=0.25, + max_num_batched_tokens=2048, + max_num_seqs=20, + max_model_len=2048, + kv_cache_layout="unified", + decoding_strategy="block_diffusion", + mask_token_id=151665, + ) + tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) + sampling_params = SamplingParams(temperature=0.0, max_tokens=256) + + dataset = load_dataset("gsm8k", "main", split="test")["question"][:10] + prompts = [ + FEW_SHOTS + f"<|im_start|>user\nQuestion: {question}\nAnswer:<|im_end|>\n<|im_start|>assistant\n" + for question in tqdm(dataset) + ] + + output_file = "log/profiles/perf_dvllm_dream_7B.json" + if os.path.exists(output_file): + os.remove(output_file) + # with VizTracer(output_file=output_file, file_info=True) as tracer: + # outputs = llm.generate(prompts[:5], sampling_params) + # time.sleep(60) + s = time.time() + outputs = LLM.generate(prompts, sampling_params) + e = time.time() + print("=*=" * 30, + "\nProfiling Results\n", + "=*=" * 30, "\n" + f"Generated {len(outputs)} outputs.\n" + f"Total tokens: {sum(len(o['token_ids']) for o in outputs)}\n" + f"Total time: {e - s:.2f} seconds.\n" + f"Avg TPS: {sum(len(o['token_ids']) for o in outputs) / (e - s):.2f} tok/s.\n" + f"AVG Number of Diffusion Steps: {sum(o['n_diff_steps'] for o in outputs) / len(outputs):.2f}\n", + "=*=" * 30) + for idx, o in enumerate(outputs): + print("\n", "=*=" * 30) + print(f"[Prompt {idx} Result] \n{prompts[idx] + "\n----------\n" + o['text']}\n") \ No newline at end of file diff --git a/examples/test_llada_dvllm_human_eval.py b/examples/test_llada_dvllm_human_eval.py index 82127ac1..5e3608f5 100755 --- a/examples/test_llada_dvllm_human_eval.py +++ b/examples/test_llada_dvllm_human_eval.py @@ -8,7 +8,7 @@ from viztracer import VizTracer from transformers import AutoTokenizer -from diffulex.legacy import LLM, SamplingParams +from diffulex_legacy import LLM, SamplingParams def summarize_profiling(csv_path: str) -> dict: diff --git a/examples/test_qwen_dvllm.py b/examples/test_qwen_dvllm.py index 14da0c48..bfd06aa9 100755 --- a/examples/test_qwen_dvllm.py +++ b/examples/test_qwen_dvllm.py @@ -1,6 +1,6 @@ import os -from diffulex.legacy import LLM, SamplingParams +from diffulex_legacy import LLM, SamplingParams from viztracer import VizTracer diff --git a/examples/test_sdar_dvllm.py b/examples/test_sdar_dvllm.py new file mode 100644 index 00000000..78fbbd7b --- /dev/null +++ b/examples/test_sdar_dvllm.py @@ -0,0 +1,210 @@ +import argparse +import os +import shutil +from pathlib import Path +import sys + + +# Ensure we import Diffulex from THIS repo (fast_dllm_v2 workspace), not an installed copy. +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _convert_safetensors_keys(src_safetensors: Path, dst_safetensors: Path) -> None: + """Convert HF-style SDAR weight keys to Diffulex-native names. + + HF checkpoint keys are like: + - model.embed_tokens.weight + - model.layers.0.self_attn.q_proj.weight + - lm_head.weight + + Diffulex native SDAR (this repo) expects: + - model.embed_tokens.weight + - model.layers.0.self_attn.q_proj.weight + - lm_head.weight + + So conversion is only needed when the source checkpoint is missing the leading + "model." prefix (some export pipelines do that). + """ + from safetensors.torch import safe_open, save_file + + tensors = {} + with safe_open(str(src_safetensors), framework="pt", device="cpu") as f: + for k in f.keys(): + new_k = k + # Add missing prefix for backbone weights. + if not new_k.startswith("model.") and new_k != "lm_head.weight": + new_k = "model." + new_k + tensors[new_k] = f.get_tensor(k) + + dst_safetensors.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(dst_safetensors)) + + +def ensure_converted_model_dir(src_model_dir: Path, out_dir: Path) -> Path: + """Create a converted model dir that Diffulex-native SDAR can load, if needed.""" + marker = out_dir / ".diffulex_sdar_converted" + dst_safetensors = out_dir / "model.safetensors" + src_safetensors = src_model_dir / "model.safetensors" + + if not src_safetensors.exists(): + raise FileNotFoundError(f"Missing {src_safetensors}") + + # If the checkpoint already matches Diffulex module names, use it directly. + from safetensors.torch import safe_open + + with safe_open(str(src_safetensors), framework="pt", device="cpu") as f: + keys = set(f.keys()) + if "model.embed_tokens.weight" in keys: + return src_model_dir + + if marker.exists() and dst_safetensors.exists(): + return out_dir + + out_dir.mkdir(parents=True, exist_ok=True) + + # Copy non-weight artifacts required by AutoConfig/AutoTokenizer. + for name in [ + "config.json", + "configuration_sdar.py", + "modeling_sdar.py", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "vocab.json", + "merges.txt", + "added_tokens.json", + "generation_config.json", + "chat_template.jinja", + "README.md", + "tokenization_qwen2.py", + "tokenization_qwen2_fast.py", + ]: + src = src_model_dir / name + if src.exists(): + shutil.copy2(src, out_dir / name) + + # Convert weights. + _convert_safetensors_keys(src_safetensors, dst_safetensors) + + marker.write_text(f"converted_from={src_model_dir}\n", encoding="utf-8") + return out_dir + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + type=str, + default="/home/lzx/SDAR/training/model/SDAR-1.7B-Chat", + help="SDAR HF model directory (contains config.json + model.safetensors).", + ) + parser.add_argument("--device", type=int, default=0) + parser.add_argument( + "--converted-dir", + type=str, + default="/home/lzx/tmp/diffulex_sdar_converted", + help="Output directory for converted checkpoint keys (Diffulex-native).", + ) + parser.add_argument("--prompt", type=str, default="你好，请用一句话介绍 SDAR。") + parser.add_argument("--max-len", type=int, default=128) + args = parser.parse_args() + + src_model_dir = Path(args.model) + converted_dir = Path(args.converted_dir) + model_dir = ensure_converted_model_dir(src_model_dir, converted_dir) + + # IMPORTANT: do not import diffulex before conversion; it may eagerly load config. + import socket + import torch + import torch.distributed as dist + from transformers import AutoTokenizer + + # Minimal single-process distributed init (required by Diffulex TP layers). + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", str(port)) + + torch.cuda.set_device(args.device) + dist.init_process_group("nccl", rank=0, world_size=1) + + # Build Config + load model weights using Diffulex loader. + from diffulex.config import Config + from diffulex.model.auto_model import AutoModelForDiffusionLM + + cfg = Config( + model=str(model_dir), + model_name="sdar", + tensor_parallel_size=1, + data_parallel_size=1, + enforce_eager=True, + ) + + dtype = getattr(cfg.hf_config, "torch_dtype", None) or torch.bfloat16 + torch.set_default_dtype(dtype) + torch.set_default_device(f"cuda:{args.device}") + + model = AutoModelForDiffusionLM.from_config(cfg).eval() + + tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True, use_fast=True) + ids = tokenizer.encode(args.prompt, add_special_tokens=True)[: args.max_len] + input_ids = torch.tensor(ids, dtype=torch.int64, device=f"cuda:{args.device}") + positions = torch.arange(input_ids.numel(), dtype=torch.int64, device=f"cuda:{args.device}") + + # Provide minimal attention metadata so Diffulex Attention can run in "prefill" mode. + # In the real engine this is provided by strategy-specific runners. + from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata + from types import SimpleNamespace + + n = int(input_ids.numel()) + cu = torch.tensor([0, n], dtype=torch.int32, device=f"cuda:{args.device}") + + def _fetch_attn_metadata(): + return SimpleNamespace( + # Core fields used by attn_impl.Attention + is_prefill=True, + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=n, + max_seqlen_k=n, + block_tables=None, + slot_mapping=None, + # KV cache controls + kv_cache_layout="unified", + need_kv_cache_store=False, + # Fields referenced in decode path (kept for completeness) + seqs=[], + total_lens=[], + seq_lens=[], + seq_lens_ts=None, + block_mask=None, + ) + + set_fetch_fn_for_attn_metadata(_fetch_attn_metadata) + + with torch.inference_mode(): + hs = model(input_ids, positions) + logits = model.compute_logits(hs) + next_id = int(logits[-1].argmax().item()) + + print("=" * 80) + print(f"[model_dir] {model_dir}") + print(f"[prompt] {args.prompt}") + print(f"[input_len] {len(ids)}") + print(f"[next_token_id] {next_id}") + print(f"[next_token] {tokenizer.decode([next_id])!r}") + + dist.destroy_process_group() + + +if __name__ == "__main__": + # Avoid tokenizer parallel warnings in multi-proc. + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + main() + + diff --git a/pyproject.toml b/pyproject.toml index 188ae078..f2e26077 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,11 @@ Repository = "https://zhijie-group.github.io/D2fEngine" "Organization" = "https://github.com/zhijie-group" [tool.setuptools.packages.find] -include = ["diffulex"] +include = [ + "diffulex", + "diffulex_kernel", + "diffulex_legacy", +] [[tool.uv.index]] url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" diff --git a/tests/.gitkeep b/tests/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py new file mode 100644 index 00000000..29200be6 --- /dev/null +++ b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py @@ -0,0 +1,395 @@ +import os +from pathlib import Path + +import torch +import tilelang +import tilelang.testing +import torch.nn.functional as F +from einops import rearrange + +# from diffulex_kernel.python.dllm_flash_attn import dllm_flash_attn_decode_kernel +from diffulex_kernel.python.dllm_flash_attn import dllm_flash_attn_decode_kernel_legacy as dllm_flash_attn_decode_kernel + + +def naive_sdpa_with_kvcache( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + scale: float, + num_groups: int, + page_block_size: int, +) -> torch.Tensor: + """ + Naive attention reference implementation with KV cache support. + + Args: + q: [Q_LEN, NUM_HEADS, HEAD_DIM] + k: [KV_LEN, NUM_KV_HEADS, HEAD_DIM] + v: [KV_LEN, NUM_KV_HEADS, HEAD_DIM] + k_cache: [NUM_PAGE_BLOCKS, PAGE_BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM] + v_cache: [NUM_PAGE_BLOCKS, PAGE_BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM] + block_tables: [NUM_SEQS, MAX_SEQ_NUM_BLOCKS] + context_lens: [NUM_SEQS] + cu_seqlens_q: [NUM_SEQS + 1] + cu_seqlens_k: [NUM_SEQS + 1] + scale: attention scale + num_groups: number of GQA groups + page_block_size: page block size + + Returns: + output: [Q_LEN, NUM_HEADS, HEAD_DIM] + """ + num_seqs = len(cu_seqlens_q) - 1 + + output = torch.zeros_like(q) + for seq_idx in range(num_seqs): + q_start = cu_seqlens_q[seq_idx].item() + q_end = cu_seqlens_q[seq_idx + 1].item() + kv_start = cu_seqlens_k[seq_idx].item() + kv_end = cu_seqlens_k[seq_idx + 1].item() + + q_seq = q[q_start:q_end] # [seq_q_len, num_heads, head_dim] + k_seq = k[kv_start:kv_end] # [seq_kv_len, num_kv_heads, head_dim] + v_seq = v[kv_start:kv_end] # [seq_kv_len, num_kv_heads, head_dim] + + context_len = context_lens[seq_idx].item() + + # Load KV cache for this sequence + k_cache_seq_list = [] + v_cache_seq_list = [] + + for block_idx in range(block_tables.shape[1]): + page_block_idx = block_tables[seq_idx, block_idx].item() + if page_block_idx >= 0: + # Calculate how many tokens to take from this block + block_start = block_idx * page_block_size + if block_start < context_len: + block_end = min(block_start + page_block_size, context_len) + num_tokens = block_end - block_start + k_cache_seq_list.append(k_cache[page_block_idx, :num_tokens]) + v_cache_seq_list.append(v_cache[page_block_idx, :num_tokens]) + + if k_cache_seq_list: + k_cache_seq = torch.cat(k_cache_seq_list, dim=0) # [context_len, num_kv_heads, head_dim] + v_cache_seq = torch.cat(v_cache_seq_list, dim=0) # [context_len, num_kv_heads, head_dim] + + # Combine KV cache and current KV + k_combined = torch.cat([k_cache_seq, k_seq], dim=0) + v_combined = torch.cat([v_cache_seq, v_seq], dim=0) + else: + k_combined = k_seq + v_combined = v_seq + + q_sdpa = rearrange(q_seq, 's h d -> 1 h s d') # [1, num_heads, seq_q_len, head_dim] + k_sdpa = rearrange(k_combined, 's h d -> 1 h s d') # [1, num_heads, total_kv_len, head_dim] + v_sdpa = rearrange(v_combined, 's h d -> 1 h s d') # [1, num_heads, total_kv_len, head_dim] + + attn_out = F.scaled_dot_product_attention( + q_sdpa, + k_sdpa, + v_sdpa, + dropout_p=0.0, + is_causal=False, + scale=scale, + enable_gqa=True, + ) # [1, num_heads, seq_q_len, head_dim] + + output[q_start:q_end] = rearrange(attn_out, '1 h s d -> s h d').to(output.dtype) + + return output + + +def run_dllm_flash_attn_decode( + num_seqs: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + max_q_len: int, + max_kv_len: int, + context_len: int, + page_block_size: int, + diffusion_block_size: int, + is_block_attn: bool, + dtype: str = "bfloat16", + block_m: int = 64, + block_n: int = 64, + num_stages: int = 1, + num_threads: int = 128, +): + """ + Run DLLM flash attention decode kernel test. + """ + torch_dtype = getattr(torch, dtype) + device = "cuda" + + num_groups = num_heads // num_kv_heads + + # Decode phase: each sequence decodes exactly one block; length equals block size + total_q_len = num_seqs * diffusion_block_size + total_kv_len = num_seqs * diffusion_block_size + + # Calculate number of page blocks needed + num_blocks_per_seq = (context_len + page_block_size - 1) // page_block_size + max_seq_num_blocks = num_blocks_per_seq + num_page_blocks = num_seqs * num_blocks_per_seq + + # Generate input tensors + q = torch.randn(total_q_len, num_heads, head_dim, dtype=torch_dtype, device=device) + k = torch.randn(total_kv_len, num_kv_heads, head_dim, dtype=torch_dtype, device=device) + v = torch.randn(total_kv_len, num_kv_heads, head_dim, dtype=torch_dtype, device=device) + + # KV cache + k_cache = torch.randn(num_page_blocks, page_block_size, num_kv_heads, head_dim, dtype=torch_dtype, device=device) + v_cache = torch.randn(num_page_blocks, page_block_size, num_kv_heads, head_dim, dtype=torch_dtype, device=device) + + # Block tables - assign page blocks sequentially for each sequence + block_tables = torch.zeros(num_seqs, max_seq_num_blocks, dtype=torch.int32, device=device) + for seq_idx in range(num_seqs): + for block_idx in range(num_blocks_per_seq): + block_tables[seq_idx, block_idx] = seq_idx * num_blocks_per_seq + block_idx + + # Context lengths + context_lens = torch.full((num_seqs,), context_len, dtype=torch.int32, device=device) + + # Cumulative sequence lengths + cu_seqlens_q = torch.arange(0, (num_seqs + 1) * diffusion_block_size, diffusion_block_size, dtype=torch.int32, device=device) + cu_seqlens_k = torch.arange(0, (num_seqs + 1) * diffusion_block_size, diffusion_block_size, dtype=torch.int32, device=device) + + scale = 1.0 / (head_dim ** 0.5) + + # Run kernel + decode_kernel = dllm_flash_attn_decode_kernel( + num_seqs, + num_groups, + num_page_blocks, + total_q_len, + total_kv_len, + num_heads, + head_dim, + is_block_attn, + diffusion_block_size, + max_seq_num_blocks, + page_block_size, + block_m, + block_n, + num_stages, + num_threads, + ) + + kernel_source = decode_kernel.get_kernel_source() + + cuda_cache_dir = os.getenv("CUDA_CACHE_DIR", "/data1/jyj/Diffulex/cuda_cache") + cache_root = Path(cuda_cache_dir) / "test_dllm_flash_attn_decode_kernel" + case_dir = cache_root / ( + f"seq{num_seqs}_heads{num_heads}_kv{num_kv_heads}_hd{head_dim}_" + f"ctx{context_len}_pbs{page_block_size}_dbs{diffusion_block_size}_" + f"block{int(is_block_attn)}_dtype{dtype}_bm{block_m}_bn{block_n}_" + f"stg{num_stages}_thr{num_threads}_mq{max_q_len}_mk{max_kv_len}" + ) + case_dir.mkdir(parents=True, exist_ok=True) + kernel_path = case_dir / "kernel.cu" + kernel_path.write_text(kernel_source) + print(f"Kernel source saved to {kernel_path}") + + output = decode_kernel( + q, k, v, k_cache, v_cache, + block_tables, + context_lens, + cu_seqlens_q, + cu_seqlens_k, + max_q_len, + ) + + # Compute reference output + ref_output = naive_sdpa_with_kvcache( + q, k, v, k_cache, v_cache, + block_tables, context_lens, + cu_seqlens_q, cu_seqlens_k, + scale, num_groups, page_block_size, + ) + + # Compare outputs + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) + print(f"Test passed! Shape: {output.shape}") + + +# ==================== Kernel Tests ==================== +def test_decode_bf16_single_seq(): + """Test with single sequence, bfloat16.""" + run_dllm_flash_attn_decode( + num_seqs=1, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=128, + page_block_size=32, + diffusion_block_size=32, + is_block_attn=False, + dtype="bfloat16", + ) + + +def test_decode_bf16_multi_seq(): + """Test with multiple sequences, bfloat16.""" + run_dllm_flash_attn_decode( + num_seqs=4, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=256, + page_block_size=32, + diffusion_block_size=32, + is_block_attn=False, + dtype="bfloat16", + ) + + +def test_decode_bf16_block_attn(): + """Test with block attention enabled.""" + run_dllm_flash_attn_decode( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=128, + page_block_size=32, + diffusion_block_size=32, + is_block_attn=True, + dtype="bfloat16", + ) + + +def test_decode_bf16_gqa_4(): + """Test with GQA ratio 4.""" + run_dllm_flash_attn_decode( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=128, + page_block_size=32, + diffusion_block_size=32, + is_block_attn=False, + dtype="bfloat16", + ) + + +def test_decode_bf16_gqa_8(): + """Test with GQA ratio 8.""" + run_dllm_flash_attn_decode( + num_seqs=2, + num_heads=32, + num_kv_heads=4, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=128, + page_block_size=32, + diffusion_block_size=32, + is_block_attn=False, + dtype="bfloat16", + ) + + +def test_decode_bf16_head_dim_64(): + """Test with head dimension 64.""" + run_dllm_flash_attn_decode( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=64, + max_q_len=64, + max_kv_len=64, + context_len=128, + page_block_size=32, + diffusion_block_size=32, + is_block_attn=False, + dtype="bfloat16", + ) + + +def test_decode_bf16_large_context(): + """Test with larger context length.""" + run_dllm_flash_attn_decode( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=512, + page_block_size=32, + diffusion_block_size=32, + is_block_attn=False, + dtype="bfloat16", + ) + + +def test_decode_bf16_page_block_64(): + """Test with page block size 64.""" + run_dllm_flash_attn_decode( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=256, + page_block_size=64, + diffusion_block_size=32, + is_block_attn=False, + dtype="bfloat16", + ) + + +def test_decode_bf16_diffusion_block_64(): + """Test with diffusion block size 64.""" + run_dllm_flash_attn_decode( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=128, + page_block_size=32, + diffusion_block_size=64, + is_block_attn=True, + dtype="bfloat16", + ) + + +def test_decode_bf16_varied_stages(): + """Test with different pipeline stages.""" + run_dllm_flash_attn_decode( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=64, + max_kv_len=64, + context_len=128, + page_block_size=32, + diffusion_block_size=32, + is_block_attn=False, + dtype="bfloat16", + num_stages=2, + ) + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py b/tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py new file mode 100644 index 00000000..6bc9ba80 --- /dev/null +++ b/tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py @@ -0,0 +1,314 @@ +import os +from pathlib import Path + +import torch +import tilelang +import tilelang.testing +import torch.nn.functional as F +from einops import rearrange + +from diffulex_kernel.python.dllm_flash_attn import dllm_flash_attn_prefill_kernel + + +def naive_sdpa_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + scale: float, + diffusion_block_size: int, + is_block_attn: bool, +) -> torch.Tensor: + """ + Naive prefill attention reference to verify TileLang kernel. + """ + num_seqs = len(cu_seqlens_q) - 1 + + output = torch.zeros_like(q) + for seq_idx in range(num_seqs): + q_start = cu_seqlens_q[seq_idx].item() + q_end = cu_seqlens_q[seq_idx + 1].item() + kv_start = cu_seqlens_k[seq_idx].item() + kv_end = cu_seqlens_k[seq_idx + 1].item() + + q_seq = q[q_start:q_end] + k_seq = k[kv_start:kv_end] + v_seq = v[kv_start:kv_end] + + q_len = q_seq.shape[0] + kv_len = k_seq.shape[0] + + q_sdpa = rearrange(q_seq, 's h d -> 1 h s d') # [1, num_heads, q_len, head_dim] + k_sdpa = rearrange(k_seq, 's h d -> 1 h s d') # [1, num_heads, kv_len, head_dim] + v_sdpa = rearrange(v_seq, 's h d -> 1 h s d') # [1, num_heads, kv_len, head_dim] + + if not is_block_attn: + attn_out = F.scaled_dot_product_attention( + q_sdpa, + k_sdpa, + v_sdpa, + dropout_p=0.0, + is_causal=False, + scale=scale, + enable_gqa=True, + ) + else: + block_mask = torch.zeros((1, 1, q_len, kv_len), dtype=q.dtype, device=q.device).bool() + num_diffusion_blocks = (kv_len + diffusion_block_size - 1) // diffusion_block_size + for block_idx in range(num_diffusion_blocks): + block_start = block_idx * diffusion_block_size + block_end = min(block_start + diffusion_block_size, kv_len) + block_mask[..., block_start:block_end, :block_end] = True + + attn_out = F.scaled_dot_product_attention( + q_sdpa, + k_sdpa, + v_sdpa, + attn_mask=block_mask, + dropout_p=0.0, + is_causal=False, + scale=scale, + enable_gqa=True, + ) + + output[q_start:q_end] = rearrange(attn_out, '1 h s d -> s h d').to(output.dtype) + + return output + + +def run_dllm_flash_attn_prefill( + num_seqs: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + max_q_len: int, + max_kv_len: int, + is_block_attn: bool, + diffusion_block_size: int, + dtype: str = "bfloat16", + block_m: int = 64, + block_n: int = 64, + num_stages: int = 1, + num_threads: int = 128, +): + """Run prefill kernel and compare with naive reference.""" + torch_dtype = getattr(torch, dtype) + device = "cuda" + num_groups = num_heads // num_kv_heads + + # Use uniform seq length per sequence to cover block mask branches + cu_seqlens_q = torch.arange(0, (num_seqs + 1) * max_q_len, max_q_len, dtype=torch.int32, device=device) + cu_seqlens_k = torch.arange(0, (num_seqs + 1) * max_kv_len, max_kv_len, dtype=torch.int32, device=device) + + total_q_len = cu_seqlens_q[-1].item() + total_kv_len = cu_seqlens_k[-1].item() + + q = torch.randn(total_q_len, num_heads, head_dim, dtype=torch_dtype, device=device) + k = torch.randn(total_kv_len, num_kv_heads, head_dim, dtype=torch_dtype, device=device) + v = torch.randn_like(k) + + prefill_kernel = dllm_flash_attn_prefill_kernel( + num_seqs, + num_groups, + total_q_len, + total_kv_len, + num_heads, + head_dim, + is_block_attn, + diffusion_block_size, + block_m, + block_n, + num_stages, + num_threads, + ) + + kernel_source = prefill_kernel.get_kernel_source() + cuda_cache_dir = os.getenv("CUDA_CACHE_DIR", "/data1/jyj/Diffulex/cuda_cache") + cache_root = Path(cuda_cache_dir) / "test_dllm_flash_attn_prefill_kernel" + case_dir = cache_root / ( + f"seq{num_seqs}_heads{num_heads}_kv{num_kv_heads}_hd{head_dim}_" + f"mq{max_q_len}_mk{max_kv_len}_block{int(is_block_attn)}_" + f"dbs{diffusion_block_size}_dtype{dtype}_bm{block_m}_bn{block_n}_" + f"stg{num_stages}_thr{num_threads}" + ) + case_dir.mkdir(parents=True, exist_ok=True) + kernel_path = case_dir / "kernel.cu" + kernel_path.write_text(kernel_source) + print(f"Kernel source saved to {kernel_path}") + + output = prefill_kernel( + q, k, v, + cu_seqlens_q, + cu_seqlens_k, + max_q_len, + ) + + scale = 1.0 / (head_dim ** 0.5) + ref_output = naive_sdpa_prefill( + q, k, v, + cu_seqlens_q, + cu_seqlens_k, + scale, + diffusion_block_size, + is_block_attn, + ) + + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) + print(f"Test passed! Shape: {output.shape}") + + +# ==================== Kernel Tests ==================== +def test_prefill_bf16_single_seq(): + """Single sequence, bfloat16.""" + run_dllm_flash_attn_prefill( + num_seqs=1, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=192, + max_kv_len=192, + is_block_attn=False, + diffusion_block_size=32, + dtype="bfloat16", + ) + + +def test_prefill_bf16_multi_seq(): + """Multiple sequences, bfloat16.""" + run_dllm_flash_attn_prefill( + num_seqs=4, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=256, + max_kv_len=256, + is_block_attn=False, + diffusion_block_size=32, + dtype="bfloat16", + ) + + +def test_prefill_bf16_block_attn(): + """Block attention, bfloat16.""" + run_dllm_flash_attn_prefill( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=256, + max_kv_len=256, + is_block_attn=True, + diffusion_block_size=32, + dtype="bfloat16", + ) + + +def test_prefill_bf16_block_attn_multi_seq_long_ctx(): + """Block attention, more sequences and longer context.""" + run_dllm_flash_attn_prefill( + num_seqs=3, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=320, + max_kv_len=320, + is_block_attn=True, + diffusion_block_size=32, + dtype="bfloat16", + ) + + +def test_prefill_bf16_block_attn_diffusion_64(): + """Block attention, diffusion block 64 to hit mask branch.""" + run_dllm_flash_attn_prefill( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=256, + max_kv_len=256, + is_block_attn=True, + diffusion_block_size=64, + dtype="bfloat16", + ) + + +def test_prefill_bf16_gqa_4(): + """GQA ratio = 4.""" + run_dllm_flash_attn_prefill( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=192, + max_kv_len=192, + is_block_attn=False, + diffusion_block_size=32, + dtype="bfloat16", + ) + + +def test_prefill_bf16_gqa_8(): + """GQA ratio = 8.""" + run_dllm_flash_attn_prefill( + num_seqs=2, + num_heads=32, + num_kv_heads=4, + head_dim=128, + max_q_len=192, + max_kv_len=192, + is_block_attn=False, + diffusion_block_size=32, + dtype="bfloat16", + ) + + +def test_prefill_bf16_block_attn_gqa_8(): + """Block attention with GQA ratio = 8.""" + run_dllm_flash_attn_prefill( + num_seqs=2, + num_heads=32, + num_kv_heads=4, + head_dim=128, + max_q_len=256, + max_kv_len=256, + is_block_attn=True, + diffusion_block_size=32, + dtype="bfloat16", + ) + + +def test_prefill_bf16_head_dim_64(): + """Head dim = 64.""" + run_dllm_flash_attn_prefill( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=64, + max_q_len=192, + max_kv_len=192, + is_block_attn=False, + diffusion_block_size=32, + dtype="bfloat16", + ) + + +def test_prefill_bf16_varied_stages(): + """Multiple pipeline stages.""" + run_dllm_flash_attn_prefill( + num_seqs=2, + num_heads=32, + num_kv_heads=8, + head_dim=128, + max_q_len=192, + max_kv_len=192, + is_block_attn=False, + diffusion_block_size=32, + dtype="bfloat16", + num_stages=2, + ) + + +if __name__ == "__main__": + tilelang.testing.main()