From 2400c0a504125ecb1b0a3a3a380589ec9309debd Mon Sep 17 00:00:00 2001 From: Chengyue Wu Date: Fri, 17 Oct 2025 01:39:15 -0700 Subject: [PATCH 01/23] add fast_dllm_model --- .gitignore | 4 +- .vscode/launch.json | 40 ++- .../configuration_fast_dllm_v2.py | 90 +++++++ d2f_engine/models/fast_dllm_v2.py | 237 ++++++++++++++++++ examples/test_dream_dvllm_gsm8k.py | 12 +- 5 files changed, 364 insertions(+), 19 deletions(-) create mode 100755 d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py create mode 100755 d2f_engine/models/fast_dllm_v2.py diff --git a/.gitignore b/.gitignore index b35bc8ef..bc8329f2 100755 --- a/.gitignore +++ b/.gitignore @@ -28,4 +28,6 @@ log/ dist/ build/ cache/ -uv.lock \ No newline at end of file +uv.lock +ckpt/ +data/ \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json index 425df5f5..eda65b1a 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -3,25 +3,41 @@ // Hover to view descriptions of existing attributes. // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 "version": "0.2.0", - "configurations": [ - - - + "configurations": [ + { + "name": "Python Debugger: Current File", + "type": "debugpy", + "request": "launch", + "program": "${file}", + "console": "integratedTerminal", + "env": { + "HF_ALLOW_CODE_EVAL": "1", + "CUDA_VISIBLE_DEVICES": "0" + } + }, { "name": "PyDbg: `Dream` Accelerate Launch Debug", "type": "debugpy", "request": "launch", "module": "accelerate.commands.launch", "args": [ - "--main_process_port", "29520", - "--num_processes", "1", + "--main_process_port", + "29520", + "--num_processes", + "1", "eval_dream.py", - "--model", "dream_lora", - "--model_args", "pretrained=/data1/ckpts/Dream-org/Dream-v0-Base-7B,lora_path=/data1/xck/ckpt/wx_dream_base/Decoder-ddt_test-20k,max_new_tokens=256,diffusion_steps=256,temperature=0,add_bos_token=true,escape_until=true,block_size=32,block_add_threshold=0.9,skip_threshold=0.95,decoded_token_threshold=0.9,dtype=bfloat16,sampling_strategy=default,save_dir=evals_dream_single/Decoder-ddt_test-20k/humaneval-ns0-len256-temp0-limit10000-diffsteps256-block32-thresh0.9-decodethresh0.9-skip0.95-toppnone-dtypebfloat16-samplingdefault", - "--tasks", "humaneval", - "--num_fewshot", "0", - "--batch_size", "1", - "--output_path", "evals_dream_single/Decoder-ddt_test-20k/humaneval-ns0-len256-temp0-limit10000-diffsteps256-block32-thresh0.9-decodethresh0.9-skip0.95-toppnone-dtypebfloat16-samplingdefault", + "--model", + "dream_lora", + "--model_args", + "pretrained=/data1/ckpts/Dream-org/Dream-v0-Base-7B,lora_path=/data1/xck/ckpt/wx_dream_base/Decoder-ddt_test-20k,max_new_tokens=256,diffusion_steps=256,temperature=0,add_bos_token=true,escape_until=true,block_size=32,block_add_threshold=0.9,skip_threshold=0.95,decoded_token_threshold=0.9,dtype=bfloat16,sampling_strategy=default,save_dir=evals_dream_single/Decoder-ddt_test-20k/humaneval-ns0-len256-temp0-limit10000-diffsteps256-block32-thresh0.9-decodethresh0.9-skip0.95-toppnone-dtypebfloat16-samplingdefault", + "--tasks", + "humaneval", + "--num_fewshot", + "0", + "--batch_size", + "1", + "--output_path", + "evals_dream_single/Decoder-ddt_test-20k/humaneval-ns0-len256-temp0-limit10000-diffsteps256-block32-thresh0.9-decodethresh0.9-skip0.95-toppnone-dtypebfloat16-samplingdefault", "--log_samples", "--confirm_run_unsafe_code" ], diff --git a/d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py b/d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py new file mode 100755 index 00000000..ab484c64 --- /dev/null +++ b/d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py @@ -0,0 +1,90 @@ + +# coding=utf-8 +# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved. +# +# 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. +"""FastdLLM V2 model configuration""" + +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 FastdLLMV2Config(PretrainedConfig): + model_type = "FastdLLMV2" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=151936, + hidden_size=4096, + intermediate_size=22016, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=False, # cache not used in diffusion + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=28, + attention_dropout=0.0, + mask_token_id=151665, + pad_token_id=151643, + bd_size=32, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + 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 + + # for 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.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_dropout = attention_dropout + # Validate the correctness of rotary position embeddings parameters + # BC: if there is a 'type' field, move it to 'rope_type'. + 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, + **kwargs, + ) + self.mask_token_id = mask_token_id + self.pad_token_id = pad_token_id + self.bd_size = bd_size + diff --git a/d2f_engine/models/fast_dllm_v2.py b/d2f_engine/models/fast_dllm_v2.py new file mode 100755 index 00000000..0ca99d51 --- /dev/null +++ b/d2f_engine/models/fast_dllm_v2.py @@ -0,0 +1,237 @@ +import os +import torch +import torch.nn as nn +import torch.distributed as dist + +from d2f_engine.layers.activation import SiluAndMul +from d2f_engine.layers.attention.attention_v5 import Attention +from d2f_engine.layers.layernorm import RMSNorm +from d2f_engine.layers.linear import RowParallelLinear, ColumnParallelLinear +from d2f_engine.layers.rotary_embedding import get_rope +from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from d2f_engine.models.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config + + +if os.environ.get("TRITON_INTERPRET", None) == "1": + torch._dynamo.reset() + torch._dynamo.config.suppress_errors = True + torch.backends.optimized_mode = False + + +class FastdLLMV2RMSNorm(RMSNorm): + def __init__(self, hidden_size, eps=1e-6): + super().__init__(hidden_size, eps) + + +class FastdLLMV2Attention(nn.Module): + """FastdLLM V2 attention mechanism.""" + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int = 32768, + head_dim: int | None = None, + rms_norm_eps: float = 1e-6, + qkv_bias: bool = True, + rope_theta: float = 10000, + rope_scaling: tuple | None = None, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + self.head_dim = head_dim or 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 + + self.q_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim, + bias=qkv_bias, + ) + self.k_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.v_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + "diffusion_lm", # Dream uses full attention + ) + + 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) + + q, k = self.rotary_emb(positions, q, k) + o = self.attn(q, k, v, mask) + output = self.o_proj(o) + return output + + +class FastdLLMV2MLP(nn.Module): + """FastdLLM V2 MLP with SiLU activation.""" + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + ) + assert hidden_act == "silu" + self.act_fn = SiluAndMul() + + def forward(self, x): + gate = self.gate_proj(x) + up = self.up_proj(x) + x = self.act_fn(torch.cat([gate, up], dim=-1)) + x = self.down_proj(x) + return x + + +class FastdLLMV2DecoderLayer(nn.Module): + """FastdLLM V2 transformer decoder layer.""" + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.self_attn = FastdLLMV2Attention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + max_position=config.max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=True, # Dream uses bias in attention + head_dim=getattr(config, 'head_dim', None), + rope_theta=getattr(config, "rope_theta", 10000), + rope_scaling=getattr(config, "rope_scaling", None), + ) + self.mlp = FastdLLMV2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + self.input_layernorm = FastdLLMV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = FastdLLMV2RMSNorm(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 FastdLLMV2Model(nn.Module): + """FastdLLM V2 model for diffusion language modeling.""" + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([FastdLLMV2DecoderLayer(config) + for _ in range(config.num_hidden_layers)]) + self.norm = FastdLLMV2RMSNorm(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 enumerate(self.layers): + hidden_states, residual = layer(positions, hidden_states, residual, mask) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class FastdLLMV2ForDiffusionLM(nn.Module): + """FastdLLM V2 model for diffusion language modeling with LM head.""" + packed_modules_mapping = {} + + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.model = FastdLLMV2Model(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, model_type='diffusion_lm') + 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: + hidden_states = self.model(input_ids, positions, mask) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + logits = self.lm_head(hidden_states) + return logits diff --git a/examples/test_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index 1b6c478e..a03057b5 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -42,10 +42,10 @@ def summarize_profiling(csv_path: str) -> dict: if __name__ == "__main__": - model = "/root/autodl-fs/models/Dream-org/Dream-v0-Base-7B" + model = "ckpt/Dream-v0-Base-7B" LLM = LLM( model, - lora_path="/root/autodl-fs/models/SJTU-Deng-Lab/D2F_Dream_Base_7B_Lora", + lora_path="ckpt/D2F_Dream_Base_7B_Lora", use_lora=True, model_name="dream", model_type="diffusion_lm", @@ -64,7 +64,7 @@ def summarize_profiling(csv_path: str) -> dict: tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) sampling_params = SamplingParams(temperature=0.0, max_tokens=256) - dataset = load_dataset("/root/autodl-fs/datasets/openai/gsm8k", "main")['test']['question'][:] + dataset = load_dataset("data/gsm8k", "main")['test']['question'][:] prompts = [tokenizer.bos_token + FEW_SHOTS + p for p in tqdm(dataset)] output_file = "log/profiles/perf_dvllm_dream_7B.json" @@ -85,6 +85,6 @@ def summarize_profiling(csv_path: str) -> dict: 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 + # 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 From 241a071e8aae1eb12ae30b7818904347775f8acc Mon Sep 17 00:00:00 2001 From: Chengyue Wu Date: Fri, 17 Oct 2025 01:39:15 -0700 Subject: [PATCH 02/23] add fast_dllm_model --- .gitignore | 4 +- .vscode/launch.json | 40 ++- .../configuration_fast_dllm_v2.py | 90 +++++++ d2f_engine/models/fast_dllm_v2.py | 237 ++++++++++++++++++ examples/test_dream_dvllm_gsm8k.py | 12 +- 5 files changed, 364 insertions(+), 19 deletions(-) create mode 100755 d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py create mode 100755 d2f_engine/models/fast_dllm_v2.py diff --git a/.gitignore b/.gitignore index b35bc8ef..bc8329f2 100755 --- a/.gitignore +++ b/.gitignore @@ -28,4 +28,6 @@ log/ dist/ build/ cache/ -uv.lock \ No newline at end of file +uv.lock +ckpt/ +data/ \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json index 425df5f5..eda65b1a 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -3,25 +3,41 @@ // Hover to view descriptions of existing attributes. // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 "version": "0.2.0", - "configurations": [ - - - + "configurations": [ + { + "name": "Python Debugger: Current File", + "type": "debugpy", + "request": "launch", + "program": "${file}", + "console": "integratedTerminal", + "env": { + "HF_ALLOW_CODE_EVAL": "1", + "CUDA_VISIBLE_DEVICES": "0" + } + }, { "name": "PyDbg: `Dream` Accelerate Launch Debug", "type": "debugpy", "request": "launch", "module": "accelerate.commands.launch", "args": [ - "--main_process_port", "29520", - "--num_processes", "1", + "--main_process_port", + "29520", + "--num_processes", + "1", "eval_dream.py", - "--model", "dream_lora", - "--model_args", "pretrained=/data1/ckpts/Dream-org/Dream-v0-Base-7B,lora_path=/data1/xck/ckpt/wx_dream_base/Decoder-ddt_test-20k,max_new_tokens=256,diffusion_steps=256,temperature=0,add_bos_token=true,escape_until=true,block_size=32,block_add_threshold=0.9,skip_threshold=0.95,decoded_token_threshold=0.9,dtype=bfloat16,sampling_strategy=default,save_dir=evals_dream_single/Decoder-ddt_test-20k/humaneval-ns0-len256-temp0-limit10000-diffsteps256-block32-thresh0.9-decodethresh0.9-skip0.95-toppnone-dtypebfloat16-samplingdefault", - "--tasks", "humaneval", - "--num_fewshot", "0", - "--batch_size", "1", - "--output_path", "evals_dream_single/Decoder-ddt_test-20k/humaneval-ns0-len256-temp0-limit10000-diffsteps256-block32-thresh0.9-decodethresh0.9-skip0.95-toppnone-dtypebfloat16-samplingdefault", + "--model", + "dream_lora", + "--model_args", + "pretrained=/data1/ckpts/Dream-org/Dream-v0-Base-7B,lora_path=/data1/xck/ckpt/wx_dream_base/Decoder-ddt_test-20k,max_new_tokens=256,diffusion_steps=256,temperature=0,add_bos_token=true,escape_until=true,block_size=32,block_add_threshold=0.9,skip_threshold=0.95,decoded_token_threshold=0.9,dtype=bfloat16,sampling_strategy=default,save_dir=evals_dream_single/Decoder-ddt_test-20k/humaneval-ns0-len256-temp0-limit10000-diffsteps256-block32-thresh0.9-decodethresh0.9-skip0.95-toppnone-dtypebfloat16-samplingdefault", + "--tasks", + "humaneval", + "--num_fewshot", + "0", + "--batch_size", + "1", + "--output_path", + "evals_dream_single/Decoder-ddt_test-20k/humaneval-ns0-len256-temp0-limit10000-diffsteps256-block32-thresh0.9-decodethresh0.9-skip0.95-toppnone-dtypebfloat16-samplingdefault", "--log_samples", "--confirm_run_unsafe_code" ], diff --git a/d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py b/d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py new file mode 100755 index 00000000..ab484c64 --- /dev/null +++ b/d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py @@ -0,0 +1,90 @@ + +# coding=utf-8 +# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved. +# +# 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. +"""FastdLLM V2 model configuration""" + +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 FastdLLMV2Config(PretrainedConfig): + model_type = "FastdLLMV2" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=151936, + hidden_size=4096, + intermediate_size=22016, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=False, # cache not used in diffusion + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=28, + attention_dropout=0.0, + mask_token_id=151665, + pad_token_id=151643, + bd_size=32, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + 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 + + # for 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.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_dropout = attention_dropout + # Validate the correctness of rotary position embeddings parameters + # BC: if there is a 'type' field, move it to 'rope_type'. + 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, + **kwargs, + ) + self.mask_token_id = mask_token_id + self.pad_token_id = pad_token_id + self.bd_size = bd_size + diff --git a/d2f_engine/models/fast_dllm_v2.py b/d2f_engine/models/fast_dllm_v2.py new file mode 100755 index 00000000..0ca99d51 --- /dev/null +++ b/d2f_engine/models/fast_dllm_v2.py @@ -0,0 +1,237 @@ +import os +import torch +import torch.nn as nn +import torch.distributed as dist + +from d2f_engine.layers.activation import SiluAndMul +from d2f_engine.layers.attention.attention_v5 import Attention +from d2f_engine.layers.layernorm import RMSNorm +from d2f_engine.layers.linear import RowParallelLinear, ColumnParallelLinear +from d2f_engine.layers.rotary_embedding import get_rope +from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from d2f_engine.models.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config + + +if os.environ.get("TRITON_INTERPRET", None) == "1": + torch._dynamo.reset() + torch._dynamo.config.suppress_errors = True + torch.backends.optimized_mode = False + + +class FastdLLMV2RMSNorm(RMSNorm): + def __init__(self, hidden_size, eps=1e-6): + super().__init__(hidden_size, eps) + + +class FastdLLMV2Attention(nn.Module): + """FastdLLM V2 attention mechanism.""" + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int = 32768, + head_dim: int | None = None, + rms_norm_eps: float = 1e-6, + qkv_bias: bool = True, + rope_theta: float = 10000, + rope_scaling: tuple | None = None, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + self.head_dim = head_dim or 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 + + self.q_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim, + bias=qkv_bias, + ) + self.k_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.v_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + "diffusion_lm", # Dream uses full attention + ) + + 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) + + q, k = self.rotary_emb(positions, q, k) + o = self.attn(q, k, v, mask) + output = self.o_proj(o) + return output + + +class FastdLLMV2MLP(nn.Module): + """FastdLLM V2 MLP with SiLU activation.""" + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + ) + assert hidden_act == "silu" + self.act_fn = SiluAndMul() + + def forward(self, x): + gate = self.gate_proj(x) + up = self.up_proj(x) + x = self.act_fn(torch.cat([gate, up], dim=-1)) + x = self.down_proj(x) + return x + + +class FastdLLMV2DecoderLayer(nn.Module): + """FastdLLM V2 transformer decoder layer.""" + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.self_attn = FastdLLMV2Attention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + max_position=config.max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=True, # Dream uses bias in attention + head_dim=getattr(config, 'head_dim', None), + rope_theta=getattr(config, "rope_theta", 10000), + rope_scaling=getattr(config, "rope_scaling", None), + ) + self.mlp = FastdLLMV2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + self.input_layernorm = FastdLLMV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = FastdLLMV2RMSNorm(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 FastdLLMV2Model(nn.Module): + """FastdLLM V2 model for diffusion language modeling.""" + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([FastdLLMV2DecoderLayer(config) + for _ in range(config.num_hidden_layers)]) + self.norm = FastdLLMV2RMSNorm(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 enumerate(self.layers): + hidden_states, residual = layer(positions, hidden_states, residual, mask) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class FastdLLMV2ForDiffusionLM(nn.Module): + """FastdLLM V2 model for diffusion language modeling with LM head.""" + packed_modules_mapping = {} + + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.model = FastdLLMV2Model(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, model_type='diffusion_lm') + 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: + hidden_states = self.model(input_ids, positions, mask) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + logits = self.lm_head(hidden_states) + return logits diff --git a/examples/test_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index 1b6c478e..a03057b5 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -42,10 +42,10 @@ def summarize_profiling(csv_path: str) -> dict: if __name__ == "__main__": - model = "/root/autodl-fs/models/Dream-org/Dream-v0-Base-7B" + model = "ckpt/Dream-v0-Base-7B" LLM = LLM( model, - lora_path="/root/autodl-fs/models/SJTU-Deng-Lab/D2F_Dream_Base_7B_Lora", + lora_path="ckpt/D2F_Dream_Base_7B_Lora", use_lora=True, model_name="dream", model_type="diffusion_lm", @@ -64,7 +64,7 @@ def summarize_profiling(csv_path: str) -> dict: tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) sampling_params = SamplingParams(temperature=0.0, max_tokens=256) - dataset = load_dataset("/root/autodl-fs/datasets/openai/gsm8k", "main")['test']['question'][:] + dataset = load_dataset("data/gsm8k", "main")['test']['question'][:] prompts = [tokenizer.bos_token + FEW_SHOTS + p for p in tqdm(dataset)] output_file = "log/profiles/perf_dvllm_dream_7B.json" @@ -85,6 +85,6 @@ def summarize_profiling(csv_path: str) -> dict: 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 + # 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 From 08a1f14820d4cfadf7b03a645d439c6f9a92b552 Mon Sep 17 00:00:00 2001 From: drewjin Date: Tue, 4 Nov 2025 08:05:06 +0000 Subject: [PATCH 03/23] refactor: optimized the code style --- README.md | 4 - d2f_engine/config.py | 4 +- d2f_engine/engine/async_engine.py | 74 --------------- d2f_engine/engine/block_manager.py | 1 + d2f_engine/engine/dp_engine.py | 9 +- d2f_engine/engine/llm_engine.py | 6 +- d2f_engine/engine/model_runner.py | 6 +- d2f_engine/engine/scheduler.py | 2 +- d2f_engine/engine/sequence.py | 6 +- d2f_engine/layers/activation.py | 4 +- d2f_engine/layers/embed_head.py | 3 +- d2f_engine/layers/layernorm.py | 2 +- d2f_engine/layers/linear.py | 2 +- d2f_engine/layers/rotary_embedding.py | 5 +- d2f_engine/models/dream.py | 7 +- d2f_engine/models/fast_dllm_v2.py | 4 +- d2f_engine/models/llada.py | 6 +- d2f_engine/models/qwen3.py | 9 +- d2f_engine/serve/__init__.py | 4 - d2f_engine/serve/__main__.py | 130 -------------------------- 20 files changed, 45 insertions(+), 243 deletions(-) delete mode 100644 d2f_engine/engine/async_engine.py delete mode 100644 d2f_engine/serve/__init__.py delete mode 100644 d2f_engine/serve/__main__.py diff --git a/README.md b/README.md index 7584c9dd..621d2470 100755 --- a/README.md +++ b/README.md @@ -10,10 +10,6 @@ vLLM implementation for Diffusion LLMs, D2F is integrated as the core inference Based on [Nano-vLLM](https://github.com/GeeeekExplorer/nano-vllm). -## How We Implement - - - ## Easy Install D2F-vLLM ```shell diff --git a/d2f_engine/config.py b/d2f_engine/config.py index 91fcc271..ecc1408d 100755 --- a/d2f_engine/config.py +++ b/d2f_engine/config.py @@ -9,6 +9,8 @@ class Config: lora_path: str = "" model_name: str = "dream" model_type: str = "diffusion_lm" # "causal_lm" or "diffusion_lm" + decoding_strategy: str = "d2f" # "d2f", "fast-dllm-v2", "block-diffusion" + mask_token_id: int = 151666 diffusion_block_size: int = 32 @@ -59,4 +61,4 @@ def __post_init__(self): self.hf_config = AutoConfig.from_pretrained(self.model, trust_remote_code=True) cfg_max_model_len = self.hf_config.max_position_embeddings if hasattr(self.hf_config, "max_position_embeddings") else self.hf_config.max_sequence_length self.max_model_len = min(self.max_model_len, cfg_max_model_len) - assert self.max_num_batched_tokens >= self.max_model_len + assert self.max_num_batched_tokens >= self.max_model_len \ No newline at end of file diff --git a/d2f_engine/engine/async_engine.py b/d2f_engine/engine/async_engine.py deleted file mode 100644 index ba6613b8..00000000 --- a/d2f_engine/engine/async_engine.py +++ /dev/null @@ -1,74 +0,0 @@ -import asyncio -from typing import AsyncGenerator, Dict, List, Tuple, Any, DefaultDict -from collections import defaultdict - -from d2f_engine.llm import LLM -from d2f_engine.sampling_params import SamplingParams - - -class AsyncEngine: - """Async driver with streaming on top of the sync Engine. - - Maintains a single background stepping task and per-sequence subscriber queues. - """ - def __init__(self, model: str, **kwargs): - self._engine = LLM(model, **kwargs) - self._subs: DefaultDict[int, list[asyncio.Queue]] = defaultdict(list) - self._driver_task: asyncio.Task | None = None - self._lock = asyncio.Lock() - - def add_request(self, prompt: str | List[int], sampling_params: SamplingParams) -> int: - return self._engine.add_request(prompt, sampling_params) - - async def _ensure_driver(self): - async with self._lock: - if self._driver_task is None or self._driver_task.done(): - self._driver_task = asyncio.create_task(self._driver()) - - async def _driver(self): - # Simple cooperative loop; runs while there are any subscribers - while True: - # If no subscribers, pause a bit and check again - if not any(self._subs.values()): - await asyncio.sleep(0.005) - # Also stop if engine has no work - if self._engine.is_finished(): - # extra sleep to avoid a spin - await asyncio.sleep(0) - continue - continue - result = self._engine.step() - outputs, _num_tok, _is_prefill, _diff_steps, deltas = ( - result if len(result) == 5 else (*result, []) - ) - # Send token deltas - for sid, toks, fin in deltas: - queues = self._subs.get(sid, []) - for q in queues: - await q.put((sid, toks, fin)) - # Ensure completions delivered even if no last delta - for sid, full_toks in outputs: - queues = self._subs.get(sid, []) - for q in queues: - await q.put((sid, full_toks, True)) - await asyncio.sleep(0) - - async def stream(self, seq_id: int) -> AsyncGenerator[Tuple[List[int], bool], None]: - q: asyncio.Queue = asyncio.Queue() - self._subs[seq_id].append(q) - await self._ensure_driver() - try: - finished = False - while not finished: - sid, toks, fin = await q.get() - if sid != seq_id: - continue - finished = fin - yield toks, finished - finally: - # Remove subscription - subs = self._subs.get(seq_id, []) - if q in subs: - subs.remove(q) - if not subs: - self._subs.pop(seq_id, None) diff --git a/d2f_engine/engine/block_manager.py b/d2f_engine/engine/block_manager.py index 8cc6645e..376646b2 100755 --- a/d2f_engine/engine/block_manager.py +++ b/d2f_engine/engine/block_manager.py @@ -1,4 +1,5 @@ import xxhash + import numpy as np from collections import deque diff --git a/d2f_engine/engine/dp_engine.py b/d2f_engine/engine/dp_engine.py index e0b6e8ee..c61e0020 100755 --- a/d2f_engine/engine/dp_engine.py +++ b/d2f_engine/engine/dp_engine.py @@ -1,13 +1,14 @@ -import atexit -import multiprocessing as mp import os import sys +import torch +import atexit import traceback import faulthandler -import torch -from multiprocessing.connection import wait as mp_wait + +import multiprocessing as mp from typing import List, Any +from multiprocessing.connection import wait as mp_wait from d2f_engine.config import Config from d2f_engine.engine.llm_engine import LLMEngine diff --git a/d2f_engine/engine/llm_engine.py b/d2f_engine/engine/llm_engine.py index 38a43532..8a70a9d1 100755 --- a/d2f_engine/engine/llm_engine.py +++ b/d2f_engine/engine/llm_engine.py @@ -2,11 +2,11 @@ import torch.multiprocessing as mp -from dataclasses import fields -from time import perf_counter +from typing import List from tqdm.auto import tqdm +from time import perf_counter +from dataclasses import fields from transformers import AutoTokenizer -from typing import List from d2f_engine.config import Config from d2f_engine.sampling_params import SamplingParams diff --git a/d2f_engine/engine/model_runner.py b/d2f_engine/engine/model_runner.py index 2e6f9238..3f8a62eb 100755 --- a/d2f_engine/engine/model_runner.py +++ b/d2f_engine/engine/model_runner.py @@ -1,7 +1,7 @@ import time - -import pickle import torch +import pickle + import torch.distributed as dist from typing import List @@ -23,6 +23,7 @@ reset_context_diffusion_lm ) + class ModelRunnerBase(ABC): """Base class for model runners supporting different model types.""" def __init__(self, config: Config, rank: int, event: Event | List[Event]): @@ -361,6 +362,7 @@ def __init__(self, config: Config, rank: int, event: Event | List[Event]): super().__init__(config, rank, event) self.diffusion_block_size = config.diffusion_block_size self.mask_token_id = config.mask_token_id + self.decoding_strategy = config.decoding_strategy def warmup_model(self): # return diff --git a/d2f_engine/engine/scheduler.py b/d2f_engine/engine/scheduler.py index 00f158b2..6c756bdc 100755 --- a/d2f_engine/engine/scheduler.py +++ b/d2f_engine/engine/scheduler.py @@ -1,8 +1,8 @@ import torch from collections import deque -from typing import Tuple, List, Deque from abc import ABC, abstractmethod +from typing import Tuple, List, Deque from d2f_engine.config import Config from d2f_engine.engine.sequence import ( diff --git a/d2f_engine/engine/sequence.py b/d2f_engine/engine/sequence.py index f5e8bf5c..be779885 100755 --- a/d2f_engine/engine/sequence.py +++ b/d2f_engine/engine/sequence.py @@ -3,12 +3,13 @@ from copy import copy from enum import Enum, auto from itertools import count -from typing import List, Tuple, Any from dataclasses import dataclass +from typing import List, Tuple, Any from d2f_engine.config import Config from d2f_engine.sampling_params import SamplingParams + class SequenceStatus(Enum): WAITING = auto() RUNNING = auto() @@ -215,6 +216,7 @@ def __init__(self, token_ids: List[int], config: Config = None): super().__init__(token_ids, sampling_params) 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 @@ -256,6 +258,7 @@ 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, @@ -288,6 +291,7 @@ 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"] diff --git a/d2f_engine/layers/activation.py b/d2f_engine/layers/activation.py index 041ee200..49de8dea 100755 --- a/d2f_engine/layers/activation.py +++ b/d2f_engine/layers/activation.py @@ -1,10 +1,10 @@ import torch -from torch import nn + +import torch.nn as nn import torch.nn.functional as F class SiluAndMul(nn.Module): - def __init__(self): super().__init__() diff --git a/d2f_engine/layers/embed_head.py b/d2f_engine/layers/embed_head.py index b9c78903..3a11a98b 100755 --- a/d2f_engine/layers/embed_head.py +++ b/d2f_engine/layers/embed_head.py @@ -1,5 +1,6 @@ import torch -from torch import nn + +import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist diff --git a/d2f_engine/layers/layernorm.py b/d2f_engine/layers/layernorm.py index 32dcfa2c..88c2cf02 100755 --- a/d2f_engine/layers/layernorm.py +++ b/d2f_engine/layers/layernorm.py @@ -1,5 +1,5 @@ import torch -from torch import nn +import torch.nn as nn class RMSNorm(nn.Module): diff --git a/d2f_engine/layers/linear.py b/d2f_engine/layers/linear.py index dc0c06a6..cf14eb9f 100755 --- a/d2f_engine/layers/linear.py +++ b/d2f_engine/layers/linear.py @@ -1,5 +1,5 @@ import torch -from torch import nn +import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist diff --git a/d2f_engine/layers/rotary_embedding.py b/d2f_engine/layers/rotary_embedding.py index aef6c7d3..6b206332 100755 --- a/d2f_engine/layers/rotary_embedding.py +++ b/d2f_engine/layers/rotary_embedding.py @@ -1,6 +1,7 @@ -from functools import lru_cache import torch -from torch import nn +import torch.nn as nn + +from functools import lru_cache def apply_rotary_emb( diff --git a/d2f_engine/models/dream.py b/d2f_engine/models/dream.py index 522ae5a7..dc81ebd0 100755 --- a/d2f_engine/models/dream.py +++ b/d2f_engine/models/dream.py @@ -3,13 +3,14 @@ import torch.nn as nn import torch.distributed as dist +from d2f_engine.layers.layernorm import RMSNorm from d2f_engine.layers.activation import SiluAndMul +from d2f_engine.layers.rotary_embedding import get_rope from d2f_engine.layers.attention.attention_v5 import Attention -from d2f_engine.layers.layernorm import RMSNorm +from d2f_engine.models.config.dream.configuration_dream import DreamConfig from d2f_engine.layers.linear import RowParallelLinear, ColumnParallelLinear -from d2f_engine.layers.rotary_embedding import get_rope from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead -from d2f_engine.models.config.dream.configuration_dream import DreamConfig + if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/d2f_engine/models/fast_dllm_v2.py b/d2f_engine/models/fast_dllm_v2.py index 0ca99d51..b3732a92 100755 --- a/d2f_engine/models/fast_dllm_v2.py +++ b/d2f_engine/models/fast_dllm_v2.py @@ -3,11 +3,11 @@ import torch.nn as nn import torch.distributed as dist +from d2f_engine.layers.layernorm import RMSNorm from d2f_engine.layers.activation import SiluAndMul +from d2f_engine.layers.rotary_embedding import get_rope from d2f_engine.layers.attention.attention_v5 import Attention -from d2f_engine.layers.layernorm import RMSNorm from d2f_engine.layers.linear import RowParallelLinear, ColumnParallelLinear -from d2f_engine.layers.rotary_embedding import get_rope from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead from d2f_engine.models.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config diff --git a/d2f_engine/models/llada.py b/d2f_engine/models/llada.py index 5158ea69..942429c0 100755 --- a/d2f_engine/models/llada.py +++ b/d2f_engine/models/llada.py @@ -3,13 +3,13 @@ import torch.nn as nn import torch.distributed as dist +from d2f_engine.layers.layernorm import RMSNorm from d2f_engine.layers.activation import SiluAndMul +from d2f_engine.layers.rotary_embedding import get_rope from d2f_engine.layers.attention.attention_v5 import Attention -from d2f_engine.layers.layernorm import RMSNorm +from d2f_engine.models.config.llada.configuration_llada import LLaDAConfig from d2f_engine.layers.linear import RowParallelLinear, ColumnParallelLinear -from d2f_engine.layers.rotary_embedding import get_rope from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead -from d2f_engine.models.config.llada.configuration_llada import LLaDAConfig if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/d2f_engine/models/qwen3.py b/d2f_engine/models/qwen3.py index 59c53f69..9d22676d 100755 --- a/d2f_engine/models/qwen3.py +++ b/d2f_engine/models/qwen3.py @@ -1,14 +1,15 @@ import torch -from torch import nn +import torch.nn as nn import torch.distributed as dist + from transformers import Qwen3Config -from d2f_engine.layers.activation import SiluAndMul -from d2f_engine.layers.attention.attention_v4 import Attention from d2f_engine.layers.layernorm import RMSNorm -from d2f_engine.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear +from d2f_engine.layers.activation import SiluAndMul from d2f_engine.layers.rotary_embedding import get_rope +from d2f_engine.layers.attention.attention_v4 import Attention from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from d2f_engine.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear class Qwen3Attention(nn.Module): diff --git a/d2f_engine/serve/__init__.py b/d2f_engine/serve/__init__.py deleted file mode 100644 index 9ea748e0..00000000 --- a/d2f_engine/serve/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -"""FastAPI serving for D2fEngine. - -Run: python -m d2f_engine.serve --model /path/to/weights [--model-name mymodel] -""" diff --git a/d2f_engine/serve/__main__.py b/d2f_engine/serve/__main__.py deleted file mode 100644 index 9463ee5f..00000000 --- a/d2f_engine/serve/__main__.py +++ /dev/null @@ -1,130 +0,0 @@ -import argparse -import asyncio -import json -from typing import Dict, Any - -from fastapi import FastAPI, HTTPException -from fastapi.responses import StreamingResponse -import uvicorn - -from d2f_engine.engine.async_engine import AsyncEngine -from d2f_engine.sampling_params import SamplingParams - - -def build_app(engine: AsyncEngine, model_name: str) -> FastAPI: - app = FastAPI(title="D2fEngine Serve", version="0.1") - - @app.get("/health") - async def health(): - return {"status": "ok", "model_name": model_name} - - @app.post("/v1/stream") - async def stream(payload: Dict[str, Any]): - # Route by external model name - req_model = payload.get("model") - if req_model and req_model != model_name: - raise HTTPException(status_code=404, detail=f"model {req_model} not served here") - - prompt = payload.get("prompt") - if prompt is None: - raise HTTPException(status_code=400, detail="prompt is required") - - params = payload.get("sampling_params", {}) - sp = SamplingParams( - temperature=float(params.get("temperature", 1.0)), - max_tokens=int(params.get("max_tokens", 64)), - ignore_eos=bool(params.get("ignore_eos", False)), - ) - - seq_id = engine.add_request(prompt, sp) - - async def token_stream(): - async for toks, finished in engine.stream(seq_id): - data = {"tokens": toks, "finished": finished} - yield json.dumps(data, separators=(",", ":")) + "\n" - - return StreamingResponse(token_stream(), media_type="application/x-ndjson") - - @app.post("/v1/generate") - async def generate(payload: Dict[str, Any]): - req_model = payload.get("model") - if req_model and req_model != model_name: - raise HTTPException(status_code=404, detail=f"model {req_model} not served here") - prompt = payload.get("prompt") - if prompt is None: - raise HTTPException(status_code=400, detail="prompt is required") - params = payload.get("sampling_params", {}) - sp = SamplingParams( - temperature=float(params.get("temperature", 1.0)), - max_tokens=int(params.get("max_tokens", 64)), - ignore_eos=bool(params.get("ignore_eos", False)), - ) - # Use underlying synchronous engine for one-shot - out = engine._engine.generate([prompt], sp, use_tqdm=False)[0] - return out - - return app - - -def main(): - parser = argparse.ArgumentParser(description="Serve D2fEngine with FastAPI") - parser.add_argument("--model", required=True, help="Path to HF model repo or local folder") - parser.add_argument("--model-name", default="default", help="External model name exposed by server") - # Mirror Config fields (subset commonly used); all forwarded to AsyncEngine/LLM - parser.add_argument("--model-type", default="diffusion_lm", choices=["causal_lm", "diffusion_lm"]) - parser.add_argument("--tensor-parallel-size", type=int, default=2) - parser.add_argument("--data-parallel-size", type=int, default=1) - parser.add_argument("--max-num-batched-tokens", type=int, default=4096) - parser.add_argument("--max-num-seqs", type=int, default=128) - parser.add_argument("--max-model-len", type=int, default=2048) - parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) - parser.add_argument("--mask-token-id", type=int, default=151666) - parser.add_argument("--diffusion-block-size", type=int, default=32) - parser.add_argument("--accept-threshold", type=float, default=0.9) - parser.add_argument("--complete-threshold", type=float, default=0.95) - parser.add_argument("--add-new-block-threshold", type=float, default=0.1) - parser.add_argument("--kv-cache-layout", default="unified", choices=["unified", "distinct"]) - parser.add_argument("--kvcache-block-size", type=int, default=256) - parser.add_argument("--k-cache-hdim-split-factor-x", type=int, default=8) - parser.add_argument("--use-lora", action="store_false") - parser.add_argument("--lora-path", default="") - parser.add_argument("--master-addr", default="localhost") - parser.add_argument("--master-port", type=int, default=2333) - parser.add_argument("--device-start", type=int, default=0) - parser.add_argument("--enforce-eager", action="store_false") - parser.add_argument("--host", default="0.0.0.0") - parser.add_argument("--port", type=int, default=8000) - - args = parser.parse_args() - - # Build engine with aligned kwargs to Config - engine = AsyncEngine( - args.model, - model_type=args.model_type, - tensor_parallel_size=args.tensor_parallel_size, - data_parallel_size=args.data_parallel_size, - max_num_batched_tokens=args.max_num_batched_tokens, - max_num_seqs=args.max_num_seqs, - max_model_len=args.max_model_len, - gpu_memory_utilization=args.gpu_memory_utilization, - mask_token_id=args.mask_token_id, - diffusion_block_size=args.diffusion_block_size, - accept_threshold=args.accept_threshold, - complete_threshold=args.complete_threshold, - add_new_block_threshold=args.add_new_block_threshold, - use_lora=args.use_lora, - lora_path=args.lora_path, - kv_cache_layout=args.kv_cache_layout, - kvcache_block_size=args.kvcache_block_size, - k_cache_hdim_split_factor_x=args.k_cache_hdim_split_factor_x, - master_addr=args.master_addr, - master_port=args.master_port, - device_start=args.device_start, - enforce_eager=args.enforce_eager, - ) - - app = build_app(engine, args.model_name) - uvicorn.run(app, host=args.host, port=args.port) - -if __name__ == "__main__": - main() From 8224f99df5355013acbd8d21f90f7f9a48461d4e Mon Sep 17 00:00:00 2001 From: drewjin Date: Tue, 4 Nov 2025 08:42:21 +0000 Subject: [PATCH 04/23] refactor: rename the project to `diffuserve` and move the original `d2f_engine` into `diffuserve/legacy` --- d2f_engine/__init__.py | 3 --- d2f_engine/layers/attention/ops/__init__.py | 7 ------- diffuserve/legacy/__init__.py | 3 +++ {d2f_engine => diffuserve/legacy}/config.py | 0 .../legacy}/engine/block_manager.py | 4 ++-- .../legacy}/engine/dp_engine.py | 6 +++--- .../legacy}/engine/llm_engine.py | 10 +++++----- .../legacy}/engine/model_runner.py | 12 ++++++------ .../legacy}/engine/scheduler.py | 8 ++++---- .../legacy}/engine/sequence.py | 4 ++-- .../legacy}/layers/activation.py | 0 .../legacy}/layers/attention/attention_v1.py | 2 +- .../layers/attention/attention_v1_profile.py | 2 +- .../legacy}/layers/attention/attention_v2.py | 2 +- .../legacy}/layers/attention/attention_v2_dup.py | 4 ++-- .../layers/attention/attention_v2_profile.py | 2 +- .../legacy}/layers/attention/attention_v3.py | 4 ++-- .../legacy}/layers/attention/attention_v4.py | 4 ++-- .../legacy}/layers/attention/attention_v5.py | 4 ++-- diffuserve/legacy/layers/attention/ops/__init__.py | 7 +++++++ .../ops/chunked_prefill_decoding_unified_kernel.py | 2 +- .../layers/attention/ops/kv_cache_kernels.py | 4 ++-- .../legacy}/layers/attention/ops/prefix_prefill.py | 0 .../layers/attention/ops/tilus_decode_attn_dlm.py | 0 .../layers/attention/ops/triton_decode_attn_clm.py | 0 .../layers/attention/ops/triton_decode_attn_dlm.py | 4 ++-- .../layers/attention/ops/triton_flash_attention.py | 0 .../legacy}/layers/embed_head.py | 2 +- .../legacy}/layers/layernorm.py | 0 {d2f_engine => diffuserve/legacy}/layers/linear.py | 0 .../legacy}/layers/rotary_embedding.py | 0 .../legacy}/layers/sampler.py | 4 ++-- {d2f_engine => diffuserve/legacy}/llm.py | 6 +++--- .../legacy}/models/auto_model.py | 10 +++++----- .../models/config/dream/configuration_dream.py | 0 .../fast_dllm_v2/configuration_fast_dllm_v2.py | 0 .../models/config/llada/configuration_llada.py | 0 {d2f_engine => diffuserve/legacy}/models/dream.py | 14 +++++++------- .../legacy}/models/fast_dllm_v2.py | 14 +++++++------- {d2f_engine => diffuserve/legacy}/models/llada.py | 14 +++++++------- {d2f_engine => diffuserve/legacy}/models/qwen3.py | 12 ++++++------ .../legacy}/models/utils/check_config.py | 0 .../legacy}/sampling_params.py | 0 {d2f_engine => diffuserve/legacy}/utils/checker.py | 0 {d2f_engine => diffuserve/legacy}/utils/context.py | 2 +- {d2f_engine => diffuserve/legacy}/utils/loader.py | 2 +- document/site/css/brands.min.css | 2 +- examples/test_causal_lm_decoding_kernel.py | 2 +- examples/test_dllm_decoding_kernel.py | 2 +- examples/test_dllm_kv_cache_load.py | 2 +- examples/test_dllm_kv_cache_store.py | 2 +- examples/test_dream_dvllm_gsm8k.py | 2 +- examples/test_dream_dvllm_human_eval.py | 2 +- examples/test_dream_model_weight.py | 4 ++-- examples/test_dream_model_weight_fixed.py | 4 ++-- examples/test_llada_dvllm_human_eval.py | 2 +- examples/test_qwen_dvllm.py | 2 +- pyproject.toml | 2 +- 58 files changed, 103 insertions(+), 103 deletions(-) delete mode 100755 d2f_engine/__init__.py delete mode 100755 d2f_engine/layers/attention/ops/__init__.py create mode 100755 diffuserve/legacy/__init__.py rename {d2f_engine => diffuserve/legacy}/config.py (100%) rename {d2f_engine => diffuserve/legacy}/engine/block_manager.py (97%) rename {d2f_engine => diffuserve/legacy}/engine/dp_engine.py (98%) rename {d2f_engine => diffuserve/legacy}/engine/llm_engine.py (94%) rename {d2f_engine => diffuserve/legacy}/engine/model_runner.py (98%) rename {d2f_engine => diffuserve/legacy}/engine/scheduler.py (97%) rename {d2f_engine => diffuserve/legacy}/engine/sequence.py (99%) rename {d2f_engine => diffuserve/legacy}/layers/activation.py (100%) rename {d2f_engine => diffuserve/legacy}/layers/attention/attention_v1.py (99%) rename {d2f_engine => diffuserve/legacy}/layers/attention/attention_v1_profile.py (99%) rename {d2f_engine => diffuserve/legacy}/layers/attention/attention_v2.py (99%) rename {d2f_engine => diffuserve/legacy}/layers/attention/attention_v2_dup.py (98%) rename {d2f_engine => diffuserve/legacy}/layers/attention/attention_v2_profile.py (99%) rename {d2f_engine => diffuserve/legacy}/layers/attention/attention_v3.py (98%) rename {d2f_engine => diffuserve/legacy}/layers/attention/attention_v4.py (97%) rename {d2f_engine => diffuserve/legacy}/layers/attention/attention_v5.py (97%) create mode 100755 diffuserve/legacy/layers/attention/ops/__init__.py rename {d2f_engine => diffuserve/legacy}/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py (99%) rename {d2f_engine => diffuserve/legacy}/layers/attention/ops/kv_cache_kernels.py (99%) rename {d2f_engine => diffuserve/legacy}/layers/attention/ops/prefix_prefill.py (100%) rename {d2f_engine => diffuserve/legacy}/layers/attention/ops/tilus_decode_attn_dlm.py (100%) rename {d2f_engine => diffuserve/legacy}/layers/attention/ops/triton_decode_attn_clm.py (100%) rename {d2f_engine => diffuserve/legacy}/layers/attention/ops/triton_decode_attn_dlm.py (97%) rename {d2f_engine => diffuserve/legacy}/layers/attention/ops/triton_flash_attention.py (100%) rename {d2f_engine => diffuserve/legacy}/layers/embed_head.py (96%) rename {d2f_engine => diffuserve/legacy}/layers/layernorm.py (100%) rename {d2f_engine => diffuserve/legacy}/layers/linear.py (100%) rename {d2f_engine => diffuserve/legacy}/layers/rotary_embedding.py (100%) rename {d2f_engine => diffuserve/legacy}/layers/sampler.py (98%) rename {d2f_engine => diffuserve/legacy}/llm.py (64%) rename {d2f_engine => diffuserve/legacy}/models/auto_model.py (58%) rename {d2f_engine => diffuserve/legacy}/models/config/dream/configuration_dream.py (100%) rename {d2f_engine => diffuserve/legacy}/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py (100%) rename {d2f_engine => diffuserve/legacy}/models/config/llada/configuration_llada.py (100%) rename {d2f_engine => diffuserve/legacy}/models/dream.py (93%) rename {d2f_engine => diffuserve/legacy}/models/fast_dllm_v2.py (93%) rename {d2f_engine => diffuserve/legacy}/models/llada.py (94%) rename {d2f_engine => diffuserve/legacy}/models/qwen3.py (93%) rename {d2f_engine => diffuserve/legacy}/models/utils/check_config.py (100%) rename {d2f_engine => diffuserve/legacy}/sampling_params.py (100%) rename {d2f_engine => diffuserve/legacy}/utils/checker.py (100%) rename {d2f_engine => diffuserve/legacy}/utils/context.py (98%) rename {d2f_engine => diffuserve/legacy}/utils/loader.py (99%) diff --git a/d2f_engine/__init__.py b/d2f_engine/__init__.py deleted file mode 100755 index e50d09d8..00000000 --- a/d2f_engine/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from d2f_engine.llm import LLM -from d2f_engine.sampling_params import SamplingParams -from d2f_engine.engine.async_engine import AsyncEngine diff --git a/d2f_engine/layers/attention/ops/__init__.py b/d2f_engine/layers/attention/ops/__init__.py deleted file mode 100755 index c7c36f1c..00000000 --- a/d2f_engine/layers/attention/ops/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from d2f_engine.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd as causal_lm_flash_decoding -from d2f_engine.layers.attention.ops.triton_decode_attn_dlm import diffusion_lm_flash_decoding, CHECK_ATTENTION -from d2f_engine.layers.attention.ops.chunked_prefill_decoding_unified_kernel import chunked_prefill_paged_decode as diffusion_lm_parallel_flash_decoding -from d2f_engine.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/diffuserve/legacy/__init__.py b/diffuserve/legacy/__init__.py new file mode 100755 index 00000000..e2d13424 --- /dev/null +++ b/diffuserve/legacy/__init__.py @@ -0,0 +1,3 @@ +from diffuserve.legacy.llm import LLM +from diffuserve.legacy.sampling_params import SamplingParams +from diffuserve.legacy.engine.async_engine import AsyncEngine diff --git a/d2f_engine/config.py b/diffuserve/legacy/config.py similarity index 100% rename from d2f_engine/config.py rename to diffuserve/legacy/config.py diff --git a/d2f_engine/engine/block_manager.py b/diffuserve/legacy/engine/block_manager.py similarity index 97% rename from d2f_engine/engine/block_manager.py rename to diffuserve/legacy/engine/block_manager.py index 376646b2..cfaf13cd 100755 --- a/d2f_engine/engine/block_manager.py +++ b/diffuserve/legacy/engine/block_manager.py @@ -7,8 +7,8 @@ from dataclasses import dataclass, field from typing import List, Dict, Deque, Set -from d2f_engine.config import Config -from d2f_engine.engine.sequence import SequenceBase, SequenceForCausalLM, SequenceForDiffusionLM +from diffuserve.legacy.config import Config +from diffuserve.legacy.engine.sequence import SequenceBase, SequenceForCausalLM, SequenceForDiffusionLM @dataclass diff --git a/d2f_engine/engine/dp_engine.py b/diffuserve/legacy/engine/dp_engine.py similarity index 98% rename from d2f_engine/engine/dp_engine.py rename to diffuserve/legacy/engine/dp_engine.py index c61e0020..9de2cb13 100755 --- a/d2f_engine/engine/dp_engine.py +++ b/diffuserve/legacy/engine/dp_engine.py @@ -10,9 +10,9 @@ from typing import List, Any from multiprocessing.connection import wait as mp_wait -from d2f_engine.config import Config -from d2f_engine.engine.llm_engine import LLMEngine -from d2f_engine.sampling_params import SamplingParams +from diffuserve.legacy.config import Config +from diffuserve.legacy.engine.llm_engine import LLMEngine +from diffuserve.legacy.sampling_params import SamplingParams def _dp_child_entry(config: Config, dp_idx: int, local_devices: list[int], conn): diff --git a/d2f_engine/engine/llm_engine.py b/diffuserve/legacy/engine/llm_engine.py similarity index 94% rename from d2f_engine/engine/llm_engine.py rename to diffuserve/legacy/engine/llm_engine.py index 8a70a9d1..e05be015 100755 --- a/d2f_engine/engine/llm_engine.py +++ b/diffuserve/legacy/engine/llm_engine.py @@ -8,11 +8,11 @@ from dataclasses import fields from transformers import AutoTokenizer -from d2f_engine.config import Config -from d2f_engine.sampling_params import SamplingParams -from d2f_engine.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM -from d2f_engine.engine.scheduler import AutoScheduler, SchedulerBase -from d2f_engine.engine.model_runner import AutoModelRunner +from diffuserve.legacy.config import Config +from diffuserve.legacy.sampling_params import SamplingParams +from diffuserve.legacy.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM +from diffuserve.legacy.engine.scheduler import AutoScheduler, SchedulerBase +from diffuserve.legacy.engine.model_runner import AutoModelRunner class LLMEngine: diff --git a/d2f_engine/engine/model_runner.py b/diffuserve/legacy/engine/model_runner.py similarity index 98% rename from d2f_engine/engine/model_runner.py rename to diffuserve/legacy/engine/model_runner.py index 3f8a62eb..1667c132 100755 --- a/d2f_engine/engine/model_runner.py +++ b/diffuserve/legacy/engine/model_runner.py @@ -9,12 +9,12 @@ from multiprocessing.synchronize import Event from multiprocessing.shared_memory import SharedMemory -from d2f_engine.config import Config -from d2f_engine.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM, SequenceBase -from d2f_engine.models.auto_model import AutoModelLM -from d2f_engine.layers.sampler import AutoSampler -from d2f_engine.utils.checker import CHECK_SLOT_MAPPING -from d2f_engine.utils.context import ( +from diffuserve.legacy.config import Config +from diffuserve.legacy.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM, SequenceBase +from diffuserve.legacy.models.auto_model import AutoModelLM +from diffuserve.legacy.layers.sampler import AutoSampler +from diffuserve.legacy.utils.checker import CHECK_SLOT_MAPPING +from diffuserve.legacy.utils.context import ( set_context_causal_lm, get_context_causal_lm, reset_context_causal_lm, diff --git a/d2f_engine/engine/scheduler.py b/diffuserve/legacy/engine/scheduler.py similarity index 97% rename from d2f_engine/engine/scheduler.py rename to diffuserve/legacy/engine/scheduler.py index 6c756bdc..66be8fc0 100755 --- a/d2f_engine/engine/scheduler.py +++ b/diffuserve/legacy/engine/scheduler.py @@ -4,13 +4,13 @@ from abc import ABC, abstractmethod from typing import Tuple, List, Deque -from d2f_engine.config import Config -from d2f_engine.engine.sequence import ( +from diffuserve.legacy.config import Config +from diffuserve.legacy.engine.sequence import ( SequenceBase, SequenceStatus, SequenceForDiffusionLM, SequenceForCausalLM ) -from d2f_engine.layers.sampler import SampleOutputForDiffusionLM -from d2f_engine.engine.block_manager import AutoBlockManager +from diffuserve.legacy.layers.sampler import SampleOutputForDiffusionLM +from diffuserve.legacy.engine.block_manager import AutoBlockManager class SchedulerBase(ABC): diff --git a/d2f_engine/engine/sequence.py b/diffuserve/legacy/engine/sequence.py similarity index 99% rename from d2f_engine/engine/sequence.py rename to diffuserve/legacy/engine/sequence.py index be779885..e273cb44 100755 --- a/d2f_engine/engine/sequence.py +++ b/diffuserve/legacy/engine/sequence.py @@ -6,8 +6,8 @@ from dataclasses import dataclass from typing import List, Tuple, Any -from d2f_engine.config import Config -from d2f_engine.sampling_params import SamplingParams +from diffuserve.legacy.config import Config +from diffuserve.legacy.sampling_params import SamplingParams class SequenceStatus(Enum): diff --git a/d2f_engine/layers/activation.py b/diffuserve/legacy/layers/activation.py similarity index 100% rename from d2f_engine/layers/activation.py rename to diffuserve/legacy/layers/activation.py diff --git a/d2f_engine/layers/attention/attention_v1.py b/diffuserve/legacy/layers/attention/attention_v1.py similarity index 99% rename from d2f_engine/layers/attention/attention_v1.py rename to diffuserve/legacy/layers/attention/attention_v1.py index 6bdc446d..a6ed634e 100755 --- a/d2f_engine/layers/attention/attention_v1.py +++ b/diffuserve/legacy/layers/attention/attention_v1.py @@ -18,7 +18,7 @@ else: from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache -from d2f_engine.utils.context import ( +from diffuserve.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/d2f_engine/layers/attention/attention_v1_profile.py b/diffuserve/legacy/layers/attention/attention_v1_profile.py similarity index 99% rename from d2f_engine/layers/attention/attention_v1_profile.py rename to diffuserve/legacy/layers/attention/attention_v1_profile.py index 8e1f0353..877a9bdc 100755 --- a/d2f_engine/layers/attention/attention_v1_profile.py +++ b/diffuserve/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 d2f_engine.utils.context import ( +from diffuserve.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/d2f_engine/layers/attention/attention_v2.py b/diffuserve/legacy/layers/attention/attention_v2.py similarity index 99% rename from d2f_engine/layers/attention/attention_v2.py rename to diffuserve/legacy/layers/attention/attention_v2.py index 6131a235..4f271882 100755 --- a/d2f_engine/layers/attention/attention_v2.py +++ b/diffuserve/legacy/layers/attention/attention_v2.py @@ -17,7 +17,7 @@ else: from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache -from d2f_engine.utils.context import ( +from diffuserve.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/d2f_engine/layers/attention/attention_v2_dup.py b/diffuserve/legacy/layers/attention/attention_v2_dup.py similarity index 98% rename from d2f_engine/layers/attention/attention_v2_dup.py rename to diffuserve/legacy/layers/attention/attention_v2_dup.py index 6943e481..a8be77f2 100755 --- a/d2f_engine/layers/attention/attention_v2_dup.py +++ b/diffuserve/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 d2f_engine.engine.sequence import SequenceForDiffusionLM -from d2f_engine.utils.context import ( +from diffuserve.legacy.engine.sequence import SequenceForDiffusionLM +from diffuserve.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/d2f_engine/layers/attention/attention_v2_profile.py b/diffuserve/legacy/layers/attention/attention_v2_profile.py similarity index 99% rename from d2f_engine/layers/attention/attention_v2_profile.py rename to diffuserve/legacy/layers/attention/attention_v2_profile.py index d3e0ec43..e3b1b7cc 100755 --- a/d2f_engine/layers/attention/attention_v2_profile.py +++ b/diffuserve/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 d2f_engine.utils.context import ( +from diffuserve.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/d2f_engine/layers/attention/attention_v3.py b/diffuserve/legacy/layers/attention/attention_v3.py similarity index 98% rename from d2f_engine/layers/attention/attention_v3.py rename to diffuserve/legacy/layers/attention/attention_v3.py index 032a6127..7dee9682 100755 --- a/d2f_engine/layers/attention/attention_v3.py +++ b/diffuserve/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 d2f_engine.layers.attention.ops import causal_lm_flash_decoding -from d2f_engine.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm +from diffuserve.legacy.layers.attention.ops import causal_lm_flash_decoding +from diffuserve.legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm @triton.jit diff --git a/d2f_engine/layers/attention/attention_v4.py b/diffuserve/legacy/layers/attention/attention_v4.py similarity index 97% rename from d2f_engine/layers/attention/attention_v4.py rename to diffuserve/legacy/layers/attention/attention_v4.py index ebcdef79..3fdccc44 100755 --- a/d2f_engine/layers/attention/attention_v4.py +++ b/diffuserve/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 d2f_engine.layers.attention.ops import ( +from diffuserve.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 d2f_engine.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm +from diffuserve.legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm class Attention(nn.Module): diff --git a/d2f_engine/layers/attention/attention_v5.py b/diffuserve/legacy/layers/attention/attention_v5.py similarity index 97% rename from d2f_engine/layers/attention/attention_v5.py rename to diffuserve/legacy/layers/attention/attention_v5.py index 267e2986..4e3540cd 100644 --- a/d2f_engine/layers/attention/attention_v5.py +++ b/diffuserve/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 d2f_engine.layers.attention.ops import ( +from diffuserve.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 d2f_engine.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm +from diffuserve.legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm class Attention(nn.Module): diff --git a/diffuserve/legacy/layers/attention/ops/__init__.py b/diffuserve/legacy/layers/attention/ops/__init__.py new file mode 100755 index 00000000..8e202106 --- /dev/null +++ b/diffuserve/legacy/layers/attention/ops/__init__.py @@ -0,0 +1,7 @@ +from diffuserve.legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd as causal_lm_flash_decoding +from diffuserve.legacy.layers.attention.ops.triton_decode_attn_dlm import diffusion_lm_flash_decoding, CHECK_ATTENTION +from diffuserve.legacy.layers.attention.ops.chunked_prefill_decoding_unified_kernel import chunked_prefill_paged_decode as diffusion_lm_parallel_flash_decoding +from diffuserve.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/d2f_engine/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py b/diffuserve/legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py similarity index 99% rename from d2f_engine/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py rename to diffuserve/legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py index ecc9fe12..8cc41a72 100755 --- a/d2f_engine/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py +++ b/diffuserve/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 d2f_engine.layers.attention.ops.prefix_prefill import context_attention_fwd +from diffuserve.legacy.layers.attention.ops.prefix_prefill import context_attention_fwd @triton.jit diff --git a/d2f_engine/layers/attention/ops/kv_cache_kernels.py b/diffuserve/legacy/layers/attention/ops/kv_cache_kernels.py similarity index 99% rename from d2f_engine/layers/attention/ops/kv_cache_kernels.py rename to diffuserve/legacy/layers/attention/ops/kv_cache_kernels.py index 0f6a0a7f..6f7d70c4 100755 --- a/d2f_engine/layers/attention/ops/kv_cache_kernels.py +++ b/diffuserve/legacy/layers/attention/ops/kv_cache_kernels.py @@ -6,8 +6,8 @@ from typing import Tuple from einops import rearrange -from d2f_engine.utils.context import ContextForDiffusionLM -from d2f_engine.engine.sequence import SequenceForDiffusionLM +from diffuserve.legacy.utils.context import ContextForDiffusionLM +from diffuserve.legacy.engine.sequence import SequenceForDiffusionLM @triton.jit def store_kvcache_kernel_causal_lm( diff --git a/d2f_engine/layers/attention/ops/prefix_prefill.py b/diffuserve/legacy/layers/attention/ops/prefix_prefill.py similarity index 100% rename from d2f_engine/layers/attention/ops/prefix_prefill.py rename to diffuserve/legacy/layers/attention/ops/prefix_prefill.py diff --git a/d2f_engine/layers/attention/ops/tilus_decode_attn_dlm.py b/diffuserve/legacy/layers/attention/ops/tilus_decode_attn_dlm.py similarity index 100% rename from d2f_engine/layers/attention/ops/tilus_decode_attn_dlm.py rename to diffuserve/legacy/layers/attention/ops/tilus_decode_attn_dlm.py diff --git a/d2f_engine/layers/attention/ops/triton_decode_attn_clm.py b/diffuserve/legacy/layers/attention/ops/triton_decode_attn_clm.py similarity index 100% rename from d2f_engine/layers/attention/ops/triton_decode_attn_clm.py rename to diffuserve/legacy/layers/attention/ops/triton_decode_attn_clm.py diff --git a/d2f_engine/layers/attention/ops/triton_decode_attn_dlm.py b/diffuserve/legacy/layers/attention/ops/triton_decode_attn_dlm.py similarity index 97% rename from d2f_engine/layers/attention/ops/triton_decode_attn_dlm.py rename to diffuserve/legacy/layers/attention/ops/triton_decode_attn_dlm.py index a86514d4..8db75c0e 100755 --- a/d2f_engine/layers/attention/ops/triton_decode_attn_dlm.py +++ b/diffuserve/legacy/layers/attention/ops/triton_decode_attn_dlm.py @@ -12,7 +12,7 @@ import triton.language as tl -from d2f_engine.utils.context import ContextForDiffusionLM +from diffuserve.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 d2f_engine.layers.attention.ops import load_kvcache + from diffuserve.legacy.layers.attention.ops import load_kvcache torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False diff --git a/d2f_engine/layers/attention/ops/triton_flash_attention.py b/diffuserve/legacy/layers/attention/ops/triton_flash_attention.py similarity index 100% rename from d2f_engine/layers/attention/ops/triton_flash_attention.py rename to diffuserve/legacy/layers/attention/ops/triton_flash_attention.py diff --git a/d2f_engine/layers/embed_head.py b/diffuserve/legacy/layers/embed_head.py similarity index 96% rename from d2f_engine/layers/embed_head.py rename to diffuserve/legacy/layers/embed_head.py index 3a11a98b..e4fd553f 100755 --- a/d2f_engine/layers/embed_head.py +++ b/diffuserve/legacy/layers/embed_head.py @@ -4,7 +4,7 @@ import torch.nn.functional as F import torch.distributed as dist -from d2f_engine.utils.context import get_context_causal_lm, get_context_diffusion_lm +from diffuserve.legacy.utils.context import get_context_causal_lm, get_context_diffusion_lm class VocabParallelEmbedding(nn.Module): diff --git a/d2f_engine/layers/layernorm.py b/diffuserve/legacy/layers/layernorm.py similarity index 100% rename from d2f_engine/layers/layernorm.py rename to diffuserve/legacy/layers/layernorm.py diff --git a/d2f_engine/layers/linear.py b/diffuserve/legacy/layers/linear.py similarity index 100% rename from d2f_engine/layers/linear.py rename to diffuserve/legacy/layers/linear.py diff --git a/d2f_engine/layers/rotary_embedding.py b/diffuserve/legacy/layers/rotary_embedding.py similarity index 100% rename from d2f_engine/layers/rotary_embedding.py rename to diffuserve/legacy/layers/rotary_embedding.py diff --git a/d2f_engine/layers/sampler.py b/diffuserve/legacy/layers/sampler.py similarity index 98% rename from d2f_engine/layers/sampler.py rename to diffuserve/legacy/layers/sampler.py index 83e5f1f6..f8babe03 100644 --- a/d2f_engine/layers/sampler.py +++ b/diffuserve/legacy/layers/sampler.py @@ -8,8 +8,8 @@ from dataclasses import dataclass from easydict import EasyDict as edict -from d2f_engine.config import Config -from d2f_engine.utils.context import get_context_diffusion_lm +from diffuserve.legacy.config import Config +from diffuserve.legacy.utils.context import get_context_diffusion_lm class SamplerForCausalLM(nn.Module): diff --git a/d2f_engine/llm.py b/diffuserve/legacy/llm.py similarity index 64% rename from d2f_engine/llm.py rename to diffuserve/legacy/llm.py index 6beaf454..ce3ecea6 100755 --- a/d2f_engine/llm.py +++ b/diffuserve/legacy/llm.py @@ -1,6 +1,6 @@ -from d2f_engine.engine.llm_engine import LLMEngine -from d2f_engine.engine.dp_engine import DPEngine -from d2f_engine.config import Config +from diffuserve.legacy.engine.llm_engine import LLMEngine +from diffuserve.legacy.engine.dp_engine import DPEngine +from diffuserve.legacy.config import Config class LLM: def __new__(cls, model, **kwargs): diff --git a/d2f_engine/models/auto_model.py b/diffuserve/legacy/models/auto_model.py similarity index 58% rename from d2f_engine/models/auto_model.py rename to diffuserve/legacy/models/auto_model.py index fc6f9cae..d9edee28 100755 --- a/d2f_engine/models/auto_model.py +++ b/diffuserve/legacy/models/auto_model.py @@ -1,8 +1,8 @@ -from d2f_engine.config import Config -from d2f_engine.utils.loader import load_model -from d2f_engine.models.dream import DreamForDiffusionLM -from d2f_engine.models.qwen3 import Qwen3ForCausalLM -from d2f_engine.models.llada import LLaDAForDiffusionLM +from diffuserve.legacy.config import Config +from diffuserve.legacy.utils.loader import load_model +from diffuserve.legacy.models.dream import DreamForDiffusionLM +from diffuserve.legacy.models.qwen3 import Qwen3ForCausalLM +from diffuserve.legacy.models.llada import LLaDAForDiffusionLM class AutoModelLM: diff --git a/d2f_engine/models/config/dream/configuration_dream.py b/diffuserve/legacy/models/config/dream/configuration_dream.py similarity index 100% rename from d2f_engine/models/config/dream/configuration_dream.py rename to diffuserve/legacy/models/config/dream/configuration_dream.py diff --git a/d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py b/diffuserve/legacy/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py similarity index 100% rename from d2f_engine/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py rename to diffuserve/legacy/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py diff --git a/d2f_engine/models/config/llada/configuration_llada.py b/diffuserve/legacy/models/config/llada/configuration_llada.py similarity index 100% rename from d2f_engine/models/config/llada/configuration_llada.py rename to diffuserve/legacy/models/config/llada/configuration_llada.py diff --git a/d2f_engine/models/dream.py b/diffuserve/legacy/models/dream.py similarity index 93% rename from d2f_engine/models/dream.py rename to diffuserve/legacy/models/dream.py index dc81ebd0..9ac26007 100755 --- a/d2f_engine/models/dream.py +++ b/diffuserve/legacy/models/dream.py @@ -3,13 +3,13 @@ import torch.nn as nn import torch.distributed as dist -from d2f_engine.layers.layernorm import RMSNorm -from d2f_engine.layers.activation import SiluAndMul -from d2f_engine.layers.rotary_embedding import get_rope -from d2f_engine.layers.attention.attention_v5 import Attention -from d2f_engine.models.config.dream.configuration_dream import DreamConfig -from d2f_engine.layers.linear import RowParallelLinear, ColumnParallelLinear -from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffuserve.legacy.layers.layernorm import RMSNorm +from diffuserve.legacy.layers.activation import SiluAndMul +from diffuserve.legacy.layers.rotary_embedding import get_rope +from diffuserve.legacy.layers.attention.attention_v5 import Attention +from diffuserve.legacy.models.config.dream.configuration_dream import DreamConfig +from diffuserve.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear +from diffuserve.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead diff --git a/d2f_engine/models/fast_dllm_v2.py b/diffuserve/legacy/models/fast_dllm_v2.py similarity index 93% rename from d2f_engine/models/fast_dllm_v2.py rename to diffuserve/legacy/models/fast_dllm_v2.py index b3732a92..40a2e024 100755 --- a/d2f_engine/models/fast_dllm_v2.py +++ b/diffuserve/legacy/models/fast_dllm_v2.py @@ -3,13 +3,13 @@ import torch.nn as nn import torch.distributed as dist -from d2f_engine.layers.layernorm import RMSNorm -from d2f_engine.layers.activation import SiluAndMul -from d2f_engine.layers.rotary_embedding import get_rope -from d2f_engine.layers.attention.attention_v5 import Attention -from d2f_engine.layers.linear import RowParallelLinear, ColumnParallelLinear -from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead -from d2f_engine.models.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config +from diffuserve.legacy.layers.layernorm import RMSNorm +from diffuserve.legacy.layers.activation import SiluAndMul +from diffuserve.legacy.layers.rotary_embedding import get_rope +from diffuserve.legacy.layers.attention.attention_v5 import Attention +from diffuserve.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear +from diffuserve.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffuserve.legacy.models.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/d2f_engine/models/llada.py b/diffuserve/legacy/models/llada.py similarity index 94% rename from d2f_engine/models/llada.py rename to diffuserve/legacy/models/llada.py index 942429c0..e35ebfbb 100755 --- a/d2f_engine/models/llada.py +++ b/diffuserve/legacy/models/llada.py @@ -3,13 +3,13 @@ import torch.nn as nn import torch.distributed as dist -from d2f_engine.layers.layernorm import RMSNorm -from d2f_engine.layers.activation import SiluAndMul -from d2f_engine.layers.rotary_embedding import get_rope -from d2f_engine.layers.attention.attention_v5 import Attention -from d2f_engine.models.config.llada.configuration_llada import LLaDAConfig -from d2f_engine.layers.linear import RowParallelLinear, ColumnParallelLinear -from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffuserve.legacy.layers.layernorm import RMSNorm +from diffuserve.legacy.layers.activation import SiluAndMul +from diffuserve.legacy.layers.rotary_embedding import get_rope +from diffuserve.legacy.layers.attention.attention_v5 import Attention +from diffuserve.legacy.models.config.llada.configuration_llada import LLaDAConfig +from diffuserve.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear +from diffuserve.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/d2f_engine/models/qwen3.py b/diffuserve/legacy/models/qwen3.py similarity index 93% rename from d2f_engine/models/qwen3.py rename to diffuserve/legacy/models/qwen3.py index 9d22676d..76906123 100755 --- a/d2f_engine/models/qwen3.py +++ b/diffuserve/legacy/models/qwen3.py @@ -4,12 +4,12 @@ from transformers import Qwen3Config -from d2f_engine.layers.layernorm import RMSNorm -from d2f_engine.layers.activation import SiluAndMul -from d2f_engine.layers.rotary_embedding import get_rope -from d2f_engine.layers.attention.attention_v4 import Attention -from d2f_engine.layers.embed_head import VocabParallelEmbedding, ParallelLMHead -from d2f_engine.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear +from diffuserve.legacy.layers.layernorm import RMSNorm +from diffuserve.legacy.layers.activation import SiluAndMul +from diffuserve.legacy.layers.rotary_embedding import get_rope +from diffuserve.legacy.layers.attention.attention_v4 import Attention +from diffuserve.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffuserve.legacy.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear class Qwen3Attention(nn.Module): diff --git a/d2f_engine/models/utils/check_config.py b/diffuserve/legacy/models/utils/check_config.py similarity index 100% rename from d2f_engine/models/utils/check_config.py rename to diffuserve/legacy/models/utils/check_config.py diff --git a/d2f_engine/sampling_params.py b/diffuserve/legacy/sampling_params.py similarity index 100% rename from d2f_engine/sampling_params.py rename to diffuserve/legacy/sampling_params.py diff --git a/d2f_engine/utils/checker.py b/diffuserve/legacy/utils/checker.py similarity index 100% rename from d2f_engine/utils/checker.py rename to diffuserve/legacy/utils/checker.py diff --git a/d2f_engine/utils/context.py b/diffuserve/legacy/utils/context.py similarity index 98% rename from d2f_engine/utils/context.py rename to diffuserve/legacy/utils/context.py index 434136ed..c6e24f1f 100755 --- a/d2f_engine/utils/context.py +++ b/diffuserve/legacy/utils/context.py @@ -3,7 +3,7 @@ from typing import List from dataclasses import dataclass -from d2f_engine.engine.sequence import SequenceForDiffusionLM +from diffuserve.legacy.engine.sequence import SequenceForDiffusionLM @dataclass class ContextBase: diff --git a/d2f_engine/utils/loader.py b/diffuserve/legacy/utils/loader.py similarity index 99% rename from d2f_engine/utils/loader.py rename to diffuserve/legacy/utils/loader.py index 4e649d0d..c16a86bf 100755 --- a/d2f_engine/utils/loader.py +++ b/diffuserve/legacy/utils/loader.py @@ -7,7 +7,7 @@ from glob import glob from functools import partial from safetensors import safe_open -from d2f_engine.config import Config +from diffuserve.legacy.config import Config def load_lora_config(lora_path: str) -> dict: diff --git a/document/site/css/brands.min.css b/document/site/css/brands.min.css index 93a9a732..d68af04e 100644 --- a/document/site/css/brands.min.css +++ b/document/site/css/brands.min.css @@ -3,4 +3,4 @@ * License - https://fontawesome.com/license/free (Icons: CC BY 4.0, Fonts: SIL OFL 1.1, Code: MIT License) * Copyright 2023 Fonticons, Inc. */ -:host,:root{--fa-style-family-brands:"Font Awesome 6 Brands";--fa-font-brands:normal 400 1em/1 "Font Awesome 6 Brands"}@font-face{font-family:"Font Awesome 6 Brands";font-style:normal;font-weight:400;font-display:block;src:url(../webfonts/fa-brands-400.woff2) format("woff2"),url(../webfonts/fa-brands-400.ttf) format("truetype")}.fa-brands,.fab{font-weight:400}.fa-monero:before{content:"\f3d0"}.fa-hooli:before{content:"\f427"}.fa-yelp:before{content:"\f1e9"}.fa-cc-visa:before{content:"\f1f0"}.fa-lastfm:before{content:"\f202"}.fa-shopware:before{content:"\f5b5"}.fa-creative-commons-nc:before{content:"\f4e8"}.fa-aws:before{content:"\f375"}.fa-redhat:before{content:"\f7bc"}.fa-yoast:before{content:"\f2b1"}.fa-cloudflare:before{content:"\e07d"}.fa-ups:before{content:"\f7e0"}.fa-pixiv:before{content:"\e640"}.fa-wpexplorer:before{content:"\f2de"}.fa-dyalog:before{content:"\f399"}.fa-bity:before{content:"\f37a"}.fa-stackpath:before{content:"\f842"}.fa-buysellads:before{content:"\f20d"}.fa-first-order:before{content:"\f2b0"}.fa-modx:before{content:"\f285"}.fa-guilded:before{content:"\e07e"}.fa-vnv:before{content:"\f40b"}.fa-js-square:before,.fa-square-js:before{content:"\f3b9"}.fa-microsoft:before{content:"\f3ca"}.fa-qq:before{content:"\f1d6"}.fa-orcid:before{content:"\f8d2"}.fa-java:before{content:"\f4e4"}.fa-invision:before{content:"\f7b0"}.fa-creative-commons-pd-alt:before{content:"\f4ed"}.fa-centercode:before{content:"\f380"}.fa-glide-g:before{content:"\f2a6"}.fa-drupal:before{content:"\f1a9"}.fa-hire-a-helper:before{content:"\f3b0"}.fa-creative-commons-by:before{content:"\f4e7"}.fa-unity:before{content:"\e049"}.fa-whmcs:before{content:"\f40d"}.fa-rocketchat:before{content:"\f3e8"}.fa-vk:before{content:"\f189"}.fa-untappd:before{content:"\f405"}.fa-mailchimp:before{content:"\f59e"}.fa-css3-alt:before{content:"\f38b"}.fa-reddit-square:before,.fa-square-reddit:before{content:"\f1a2"}.fa-vimeo-v:before{content:"\f27d"}.fa-contao:before{content:"\f26d"}.fa-square-font-awesome:before{content:"\e5ad"}.fa-deskpro:before{content:"\f38f"}.fa-brave:before{content:"\e63c"}.fa-sistrix:before{content:"\f3ee"}.fa-instagram-square:before,.fa-square-instagram:before{content:"\e055"}.fa-battle-net:before{content:"\f835"}.fa-the-red-yeti:before{content:"\f69d"}.fa-hacker-news-square:before,.fa-square-hacker-news:before{content:"\f3af"}.fa-edge:before{content:"\f282"}.fa-threads:before{content:"\e618"}.fa-napster:before{content:"\f3d2"}.fa-snapchat-square:before,.fa-square-snapchat:before{content:"\f2ad"}.fa-google-plus-g:before{content:"\f0d5"}.fa-artstation:before{content:"\f77a"}.fa-markdown:before{content:"\f60f"}.fa-sourcetree:before{content:"\f7d3"}.fa-google-plus:before{content:"\f2b3"}.fa-diaspora:before{content:"\f791"}.fa-foursquare:before{content:"\f180"}.fa-stack-overflow:before{content:"\f16c"}.fa-github-alt:before{content:"\f113"}.fa-phoenix-squadron:before{content:"\f511"}.fa-pagelines:before{content:"\f18c"}.fa-algolia:before{content:"\f36c"}.fa-red-river:before{content:"\f3e3"}.fa-creative-commons-sa:before{content:"\f4ef"}.fa-safari:before{content:"\f267"}.fa-google:before{content:"\f1a0"}.fa-font-awesome-alt:before,.fa-square-font-awesome-stroke:before{content:"\f35c"}.fa-atlassian:before{content:"\f77b"}.fa-linkedin-in:before{content:"\f0e1"}.fa-digital-ocean:before{content:"\f391"}.fa-nimblr:before{content:"\f5a8"}.fa-chromecast:before{content:"\f838"}.fa-evernote:before{content:"\f839"}.fa-hacker-news:before{content:"\f1d4"}.fa-creative-commons-sampling:before{content:"\f4f0"}.fa-adversal:before{content:"\f36a"}.fa-creative-commons:before{content:"\f25e"}.fa-watchman-monitoring:before{content:"\e087"}.fa-fonticons:before{content:"\f280"}.fa-weixin:before{content:"\f1d7"}.fa-shirtsinbulk:before{content:"\f214"}.fa-codepen:before{content:"\f1cb"}.fa-git-alt:before{content:"\f841"}.fa-lyft:before{content:"\f3c3"}.fa-rev:before{content:"\f5b2"}.fa-windows:before{content:"\f17a"}.fa-wizards-of-the-coast:before{content:"\f730"}.fa-square-viadeo:before,.fa-viadeo-square:before{content:"\f2aa"}.fa-meetup:before{content:"\f2e0"}.fa-centos:before{content:"\f789"}.fa-adn:before{content:"\f170"}.fa-cloudsmith:before{content:"\f384"}.fa-opensuse:before{content:"\e62b"}.fa-pied-piper-alt:before{content:"\f1a8"}.fa-dribbble-square:before,.fa-square-dribbble:before{content:"\f397"}.fa-codiepie:before{content:"\f284"}.fa-node:before{content:"\f419"}.fa-mix:before{content:"\f3cb"}.fa-steam:before{content:"\f1b6"}.fa-cc-apple-pay:before{content:"\f416"}.fa-scribd:before{content:"\f28a"}.fa-debian:before{content:"\e60b"}.fa-openid:before{content:"\f19b"}.fa-instalod:before{content:"\e081"}.fa-expeditedssl:before{content:"\f23e"}.fa-sellcast:before{content:"\f2da"}.fa-square-twitter:before,.fa-twitter-square:before{content:"\f081"}.fa-r-project:before{content:"\f4f7"}.fa-delicious:before{content:"\f1a5"}.fa-freebsd:before{content:"\f3a4"}.fa-vuejs:before{content:"\f41f"}.fa-accusoft:before{content:"\f369"}.fa-ioxhost:before{content:"\f208"}.fa-fonticons-fi:before{content:"\f3a2"}.fa-app-store:before{content:"\f36f"}.fa-cc-mastercard:before{content:"\f1f1"}.fa-itunes-note:before{content:"\f3b5"}.fa-golang:before{content:"\e40f"}.fa-kickstarter:before{content:"\f3bb"}.fa-grav:before{content:"\f2d6"}.fa-weibo:before{content:"\f18a"}.fa-uncharted:before{content:"\e084"}.fa-firstdraft:before{content:"\f3a1"}.fa-square-youtube:before,.fa-youtube-square:before{content:"\f431"}.fa-wikipedia-w:before{content:"\f266"}.fa-rendact:before,.fa-wpressr:before{content:"\f3e4"}.fa-angellist:before{content:"\f209"}.fa-galactic-republic:before{content:"\f50c"}.fa-nfc-directional:before{content:"\e530"}.fa-skype:before{content:"\f17e"}.fa-joget:before{content:"\f3b7"}.fa-fedora:before{content:"\f798"}.fa-stripe-s:before{content:"\f42a"}.fa-meta:before{content:"\e49b"}.fa-laravel:before{content:"\f3bd"}.fa-hotjar:before{content:"\f3b1"}.fa-bluetooth-b:before{content:"\f294"}.fa-square-letterboxd:before{content:"\e62e"}.fa-sticker-mule:before{content:"\f3f7"}.fa-creative-commons-zero:before{content:"\f4f3"}.fa-hips:before{content:"\f452"}.fa-behance:before{content:"\f1b4"}.fa-reddit:before{content:"\f1a1"}.fa-discord:before{content:"\f392"}.fa-chrome:before{content:"\f268"}.fa-app-store-ios:before{content:"\f370"}.fa-cc-discover:before{content:"\f1f2"}.fa-wpbeginner:before{content:"\f297"}.fa-confluence:before{content:"\f78d"}.fa-shoelace:before{content:"\e60c"}.fa-mdb:before{content:"\f8ca"}.fa-dochub:before{content:"\f394"}.fa-accessible-icon:before{content:"\f368"}.fa-ebay:before{content:"\f4f4"}.fa-amazon:before{content:"\f270"}.fa-unsplash:before{content:"\e07c"}.fa-yarn:before{content:"\f7e3"}.fa-square-steam:before,.fa-steam-square:before{content:"\f1b7"}.fa-500px:before{content:"\f26e"}.fa-square-vimeo:before,.fa-vimeo-square:before{content:"\f194"}.fa-asymmetrik:before{content:"\f372"}.fa-font-awesome-flag:before,.fa-font-awesome-logo-full:before,.fa-font-awesome:before{content:"\f2b4"}.fa-gratipay:before{content:"\f184"}.fa-apple:before{content:"\f179"}.fa-hive:before{content:"\e07f"}.fa-gitkraken:before{content:"\f3a6"}.fa-keybase:before{content:"\f4f5"}.fa-apple-pay:before{content:"\f415"}.fa-padlet:before{content:"\e4a0"}.fa-amazon-pay:before{content:"\f42c"}.fa-github-square:before,.fa-square-github:before{content:"\f092"}.fa-stumbleupon:before{content:"\f1a4"}.fa-fedex:before{content:"\f797"}.fa-phoenix-framework:before{content:"\f3dc"}.fa-shopify:before{content:"\e057"}.fa-neos:before{content:"\f612"}.fa-square-threads:before{content:"\e619"}.fa-hackerrank:before{content:"\f5f7"}.fa-researchgate:before{content:"\f4f8"}.fa-swift:before{content:"\f8e1"}.fa-angular:before{content:"\f420"}.fa-speakap:before{content:"\f3f3"}.fa-angrycreative:before{content:"\f36e"}.fa-y-combinator:before{content:"\f23b"}.fa-empire:before{content:"\f1d1"}.fa-envira:before{content:"\f299"}.fa-google-scholar:before{content:"\e63b"}.fa-gitlab-square:before,.fa-square-gitlab:before{content:"\e5ae"}.fa-studiovinari:before{content:"\f3f8"}.fa-pied-piper:before{content:"\f2ae"}.fa-wordpress:before{content:"\f19a"}.fa-product-hunt:before{content:"\f288"}.fa-firefox:before{content:"\f269"}.fa-linode:before{content:"\f2b8"}.fa-goodreads:before{content:"\f3a8"}.fa-odnoklassniki-square:before,.fa-square-odnoklassniki:before{content:"\f264"}.fa-jsfiddle:before{content:"\f1cc"}.fa-sith:before{content:"\f512"}.fa-themeisle:before{content:"\f2b2"}.fa-page4:before{content:"\f3d7"}.fa-hashnode:before{content:"\e499"}.fa-react:before{content:"\f41b"}.fa-cc-paypal:before{content:"\f1f4"}.fa-squarespace:before{content:"\f5be"}.fa-cc-stripe:before{content:"\f1f5"}.fa-creative-commons-share:before{content:"\f4f2"}.fa-bitcoin:before{content:"\f379"}.fa-keycdn:before{content:"\f3ba"}.fa-opera:before{content:"\f26a"}.fa-itch-io:before{content:"\f83a"}.fa-umbraco:before{content:"\f8e8"}.fa-galactic-senate:before{content:"\f50d"}.fa-ubuntu:before{content:"\f7df"}.fa-draft2digital:before{content:"\f396"}.fa-stripe:before{content:"\f429"}.fa-houzz:before{content:"\f27c"}.fa-gg:before{content:"\f260"}.fa-dhl:before{content:"\f790"}.fa-pinterest-square:before,.fa-square-pinterest:before{content:"\f0d3"}.fa-xing:before{content:"\f168"}.fa-blackberry:before{content:"\f37b"}.fa-creative-commons-pd:before{content:"\f4ec"}.fa-playstation:before{content:"\f3df"}.fa-quinscape:before{content:"\f459"}.fa-less:before{content:"\f41d"}.fa-blogger-b:before{content:"\f37d"}.fa-opencart:before{content:"\f23d"}.fa-vine:before{content:"\f1ca"}.fa-signal-messenger:before{content:"\e663"}.fa-paypal:before{content:"\f1ed"}.fa-gitlab:before{content:"\f296"}.fa-typo3:before{content:"\f42b"}.fa-reddit-alien:before{content:"\f281"}.fa-yahoo:before{content:"\f19e"}.fa-dailymotion:before{content:"\e052"}.fa-affiliatetheme:before{content:"\f36b"}.fa-pied-piper-pp:before{content:"\f1a7"}.fa-bootstrap:before{content:"\f836"}.fa-odnoklassniki:before{content:"\f263"}.fa-nfc-symbol:before{content:"\e531"}.fa-mintbit:before{content:"\e62f"}.fa-ethereum:before{content:"\f42e"}.fa-speaker-deck:before{content:"\f83c"}.fa-creative-commons-nc-eu:before{content:"\f4e9"}.fa-patreon:before{content:"\f3d9"}.fa-avianex:before{content:"\f374"}.fa-ello:before{content:"\f5f1"}.fa-gofore:before{content:"\f3a7"}.fa-bimobject:before{content:"\f378"}.fa-brave-reverse:before{content:"\e63d"}.fa-facebook-f:before{content:"\f39e"}.fa-google-plus-square:before,.fa-square-google-plus:before{content:"\f0d4"}.fa-mandalorian:before{content:"\f50f"}.fa-first-order-alt:before{content:"\f50a"}.fa-osi:before{content:"\f41a"}.fa-google-wallet:before{content:"\f1ee"}.fa-d-and-d-beyond:before{content:"\f6ca"}.fa-periscope:before{content:"\f3da"}.fa-fulcrum:before{content:"\f50b"}.fa-cloudscale:before{content:"\f383"}.fa-forumbee:before{content:"\f211"}.fa-mizuni:before{content:"\f3cc"}.fa-schlix:before{content:"\f3ea"}.fa-square-xing:before,.fa-xing-square:before{content:"\f169"}.fa-bandcamp:before{content:"\f2d5"}.fa-wpforms:before{content:"\f298"}.fa-cloudversify:before{content:"\f385"}.fa-usps:before{content:"\f7e1"}.fa-megaport:before{content:"\f5a3"}.fa-magento:before{content:"\f3c4"}.fa-spotify:before{content:"\f1bc"}.fa-optin-monster:before{content:"\f23c"}.fa-fly:before{content:"\f417"}.fa-aviato:before{content:"\f421"}.fa-itunes:before{content:"\f3b4"}.fa-cuttlefish:before{content:"\f38c"}.fa-blogger:before{content:"\f37c"}.fa-flickr:before{content:"\f16e"}.fa-viber:before{content:"\f409"}.fa-soundcloud:before{content:"\f1be"}.fa-digg:before{content:"\f1a6"}.fa-tencent-weibo:before{content:"\f1d5"}.fa-letterboxd:before{content:"\e62d"}.fa-symfony:before{content:"\f83d"}.fa-maxcdn:before{content:"\f136"}.fa-etsy:before{content:"\f2d7"}.fa-facebook-messenger:before{content:"\f39f"}.fa-audible:before{content:"\f373"}.fa-think-peaks:before{content:"\f731"}.fa-bilibili:before{content:"\e3d9"}.fa-erlang:before{content:"\f39d"}.fa-x-twitter:before{content:"\e61b"}.fa-cotton-bureau:before{content:"\f89e"}.fa-dashcube:before{content:"\f210"}.fa-42-group:before,.fa-innosoft:before{content:"\e080"}.fa-stack-exchange:before{content:"\f18d"}.fa-elementor:before{content:"\f430"}.fa-pied-piper-square:before,.fa-square-pied-piper:before{content:"\e01e"}.fa-creative-commons-nd:before{content:"\f4eb"}.fa-palfed:before{content:"\f3d8"}.fa-superpowers:before{content:"\f2dd"}.fa-resolving:before{content:"\f3e7"}.fa-xbox:before{content:"\f412"}.fa-searchengin:before{content:"\f3eb"}.fa-tiktok:before{content:"\e07b"}.fa-facebook-square:before,.fa-square-facebook:before{content:"\f082"}.fa-renren:before{content:"\f18b"}.fa-linux:before{content:"\f17c"}.fa-glide:before{content:"\f2a5"}.fa-linkedin:before{content:"\f08c"}.fa-hubspot:before{content:"\f3b2"}.fa-deploydog:before{content:"\f38e"}.fa-twitch:before{content:"\f1e8"}.fa-ravelry:before{content:"\f2d9"}.fa-mixer:before{content:"\e056"}.fa-lastfm-square:before,.fa-square-lastfm:before{content:"\f203"}.fa-vimeo:before{content:"\f40a"}.fa-mendeley:before{content:"\f7b3"}.fa-uniregistry:before{content:"\f404"}.fa-figma:before{content:"\f799"}.fa-creative-commons-remix:before{content:"\f4ee"}.fa-cc-amazon-pay:before{content:"\f42d"}.fa-dropbox:before{content:"\f16b"}.fa-instagram:before{content:"\f16d"}.fa-cmplid:before{content:"\e360"}.fa-upwork:before{content:"\e641"}.fa-facebook:before{content:"\f09a"}.fa-gripfire:before{content:"\f3ac"}.fa-jedi-order:before{content:"\f50e"}.fa-uikit:before{content:"\f403"}.fa-fort-awesome-alt:before{content:"\f3a3"}.fa-phabricator:before{content:"\f3db"}.fa-ussunnah:before{content:"\f407"}.fa-earlybirds:before{content:"\f39a"}.fa-trade-federation:before{content:"\f513"}.fa-autoprefixer:before{content:"\f41c"}.fa-whatsapp:before{content:"\f232"}.fa-slideshare:before{content:"\f1e7"}.fa-google-play:before{content:"\f3ab"}.fa-viadeo:before{content:"\f2a9"}.fa-line:before{content:"\f3c0"}.fa-google-drive:before{content:"\f3aa"}.fa-servicestack:before{content:"\f3ec"}.fa-simplybuilt:before{content:"\f215"}.fa-bitbucket:before{content:"\f171"}.fa-imdb:before{content:"\f2d8"}.fa-deezer:before{content:"\e077"}.fa-raspberry-pi:before{content:"\f7bb"}.fa-jira:before{content:"\f7b1"}.fa-docker:before{content:"\f395"}.fa-screenpal:before{content:"\e570"}.fa-bluetooth:before{content:"\f293"}.fa-gitter:before{content:"\f426"}.fa-d-and-d:before{content:"\f38d"}.fa-microblog:before{content:"\e01a"}.fa-cc-diners-club:before{content:"\f24c"}.fa-gg-circle:before{content:"\f261"}.fa-pied-piper-hat:before{content:"\f4e5"}.fa-kickstarter-k:before{content:"\f3bc"}.fa-yandex:before{content:"\f413"}.fa-readme:before{content:"\f4d5"}.fa-html5:before{content:"\f13b"}.fa-sellsy:before{content:"\f213"}.fa-sass:before{content:"\f41e"}.fa-wirsindhandwerk:before,.fa-wsh:before{content:"\e2d0"}.fa-buromobelexperte:before{content:"\f37f"}.fa-salesforce:before{content:"\f83b"}.fa-octopus-deploy:before{content:"\e082"}.fa-medapps:before{content:"\f3c6"}.fa-ns8:before{content:"\f3d5"}.fa-pinterest-p:before{content:"\f231"}.fa-apper:before{content:"\f371"}.fa-fort-awesome:before{content:"\f286"}.fa-waze:before{content:"\f83f"}.fa-cc-jcb:before{content:"\f24b"}.fa-snapchat-ghost:before,.fa-snapchat:before{content:"\f2ab"}.fa-fantasy-flight-games:before{content:"\f6dc"}.fa-rust:before{content:"\e07a"}.fa-wix:before{content:"\f5cf"}.fa-behance-square:before,.fa-square-behance:before{content:"\f1b5"}.fa-supple:before{content:"\f3f9"}.fa-webflow:before{content:"\e65c"}.fa-rebel:before{content:"\f1d0"}.fa-css3:before{content:"\f13c"}.fa-staylinked:before{content:"\f3f5"}.fa-kaggle:before{content:"\f5fa"}.fa-space-awesome:before{content:"\e5ac"}.fa-deviantart:before{content:"\f1bd"}.fa-cpanel:before{content:"\f388"}.fa-goodreads-g:before{content:"\f3a9"}.fa-git-square:before,.fa-square-git:before{content:"\f1d2"}.fa-square-tumblr:before,.fa-tumblr-square:before{content:"\f174"}.fa-trello:before{content:"\f181"}.fa-creative-commons-nc-jp:before{content:"\f4ea"}.fa-get-pocket:before{content:"\f265"}.fa-perbyte:before{content:"\e083"}.fa-grunt:before{content:"\f3ad"}.fa-weebly:before{content:"\f5cc"}.fa-connectdevelop:before{content:"\f20e"}.fa-leanpub:before{content:"\f212"}.fa-black-tie:before{content:"\f27e"}.fa-themeco:before{content:"\f5c6"}.fa-python:before{content:"\f3e2"}.fa-android:before{content:"\f17b"}.fa-bots:before{content:"\e340"}.fa-free-code-camp:before{content:"\f2c5"}.fa-hornbill:before{content:"\f592"}.fa-js:before{content:"\f3b8"}.fa-ideal:before{content:"\e013"}.fa-git:before{content:"\f1d3"}.fa-dev:before{content:"\f6cc"}.fa-sketch:before{content:"\f7c6"}.fa-yandex-international:before{content:"\f414"}.fa-cc-amex:before{content:"\f1f3"}.fa-uber:before{content:"\f402"}.fa-github:before{content:"\f09b"}.fa-php:before{content:"\f457"}.fa-alipay:before{content:"\f642"}.fa-youtube:before{content:"\f167"}.fa-skyatlas:before{content:"\f216"}.fa-firefox-browser:before{content:"\e007"}.fa-replyd:before{content:"\f3e6"}.fa-suse:before{content:"\f7d6"}.fa-jenkins:before{content:"\f3b6"}.fa-twitter:before{content:"\f099"}.fa-rockrms:before{content:"\f3e9"}.fa-pinterest:before{content:"\f0d2"}.fa-buffer:before{content:"\f837"}.fa-npm:before{content:"\f3d4"}.fa-yammer:before{content:"\f840"}.fa-btc:before{content:"\f15a"}.fa-dribbble:before{content:"\f17d"}.fa-stumbleupon-circle:before{content:"\f1a3"}.fa-internet-explorer:before{content:"\f26b"}.fa-stubber:before{content:"\e5c7"}.fa-telegram-plane:before,.fa-telegram:before{content:"\f2c6"}.fa-old-republic:before{content:"\f510"}.fa-odysee:before{content:"\e5c6"}.fa-square-whatsapp:before,.fa-whatsapp-square:before{content:"\f40c"}.fa-node-js:before{content:"\f3d3"}.fa-edge-legacy:before{content:"\e078"}.fa-slack-hash:before,.fa-slack:before{content:"\f198"}.fa-medrt:before{content:"\f3c8"}.fa-usb:before{content:"\f287"}.fa-tumblr:before{content:"\f173"}.fa-vaadin:before{content:"\f408"}.fa-quora:before{content:"\f2c4"}.fa-square-x-twitter:before{content:"\e61a"}.fa-reacteurope:before{content:"\f75d"}.fa-medium-m:before,.fa-medium:before{content:"\f23a"}.fa-amilia:before{content:"\f36d"}.fa-mixcloud:before{content:"\f289"}.fa-flipboard:before{content:"\f44d"}.fa-viacoin:before{content:"\f237"}.fa-critical-role:before{content:"\f6c9"}.fa-sitrox:before{content:"\e44a"}.fa-discourse:before{content:"\f393"}.fa-joomla:before{content:"\f1aa"}.fa-mastodon:before{content:"\f4f6"}.fa-airbnb:before{content:"\f834"}.fa-wolf-pack-battalion:before{content:"\f514"}.fa-buy-n-large:before{content:"\f8a6"}.fa-gulp:before{content:"\f3ae"}.fa-creative-commons-sampling-plus:before{content:"\f4f1"}.fa-strava:before{content:"\f428"}.fa-ember:before{content:"\f423"}.fa-canadian-maple-leaf:before{content:"\f785"}.fa-teamspeak:before{content:"\f4f9"}.fa-pushed:before{content:"\f3e1"}.fa-wordpress-simple:before{content:"\f411"}.fa-nutritionix:before{content:"\f3d6"}.fa-wodu:before{content:"\e088"}.fa-google-pay:before{content:"\e079"}.fa-intercom:before{content:"\f7af"}.fa-zhihu:before{content:"\f63f"}.fa-korvue:before{content:"\f42f"}.fa-pix:before{content:"\e43a"}.fa-steam-symbol:before{content:"\f3f6"} \ No newline at end of file +:host,:root{--fa-style-family-brands:"Font Awesome 6 Brands";--fa-font-brands:normal 400 1em/1 "Font Awesome 6 Brands"}@font-face{font-family:"Font Awesome 6 Brands";font-style:normal;font-weight:400;font-display:block;src:url(../webfonts/fa-brands-400.woff2) format("woff2"),url(../webfonts/fa-brands-400.ttf) format("truetype")}.fa-brands,.fab{font-weight:400}.fa-monero:before{content:"\f3d0"}.fa-hooli:before{content:"\f427"}.fa-yelp:before{content:"\f1e9"}.fa-cc-visa:before{content:"\f1f0"}.fa-lastfm:before{content:"\f202"}.fa-shopware:before{content:"\f5b5"}.fa-creative-commons-nc:before{content:"\f4e8"}.fa-aws:before{content:"\f375"}.fa-redhat:before{content:"\f7bc"}.fa-yoast:before{content:"\f2b1"}.fa-cloudflare:before{content:"\e07d"}.fa-ups:before{content:"\f7e0"}.fa-pixiv:before{content:"\e640"}.fa-wpexplorer:before{content:"\f2de"}.fa-dyalog:before{content:"\f399"}.fa-bity:before{content:"\f37a"}.fa-stackpath:before{content:"\f842"}.fa-buysellads:before{content:"\f20d"}.fa-first-order:before{content:"\f2b0"}.fa-modx:before{content:"\f285"}.fa-guilded:before{content:"\e07e"}.fa-vnv:before{content:"\f40b"}.fa-js-square:before,.fa-square-js:before{content:"\f3b9"}.fa-microsoft:before{content:"\f3ca"}.fa-qq:before{content:"\f1d6"}.fa-orcid:before{content:"\f8d2"}.fa-java:before{content:"\f4e4"}.fa-invision:before{content:"\f7b0"}.fa-creative-commons-pd-alt:before{content:"\f4ed"}.fa-centercode:before{content:"\f380"}.fa-glide-g:before{content:"\f2a6"}.fa-drupal:before{content:"\f1a9"}.fa-hire-a-helper:before{content:"\f3b0"}.fa-creative-commons-by:before{content:"\f4e7"}.fa-unity:before{content:"\e049"}.fa-whmcs:before{content:"\f40d"}.fa-rocketchat:before{content:"\f3e8"}.fa-vk:before{content:"\f189"}.fa-untappd:before{content:"\f405"}.fa-mailchimp:before{content:"\f59e"}.fa-css3-alt:before{content:"\f38b"}.fa-reddit-square:before,.fa-square-reddit:before{content:"\f1a2"}.fa-vimeo-v:before{content:"\f27d"}.fa-contao:before{content:"\f26d"}.fa-square-font-awesome:before{content:"\e5ad"}.fa-deskpro:before{content:"\f38f"}.fa-brave:before{content:"\e63c"}.fa-sistrix:before{content:"\f3ee"}.fa-instagram-square:before,.fa-square-instagram:before{content:"\e055"}.fa-battle-net:before{content:"\f835"}.fa-the-red-yeti:before{content:"\f69d"}.fa-hacker-news-square:before,.fa-square-hacker-news:before{content:"\f3af"}.fa-edge:before{content:"\f282"}.fa-threads:before{content:"\e618"}.fa-napster:before{content:"\f3d2"}.fa-snapchat-square:before,.fa-square-snapchat:before{content:"\f2ad"}.fa-google-plus-g:before{content:"\f0d5"}.fa-artstation:before{content:"\f77a"}.fa-markdown:before{content:"\f60f"}.fa-sourcetree:before{content:"\f7d3"}.fa-google-plus:before{content:"\f2b3"}.fa-diaspora:before{content:"\f791"}.fa-foursquare:before{content:"\f180"}.fa-stack-overflow:before{content:"\f16c"}.fa-github-alt:before{content:"\f113"}.fa-phoenix-squadron:before{content:"\f511"}.fa-pagelines:before{content:"\f18c"}.fa-algolia:before{content:"\f36c"}.fa-red-river:before{content:"\f3e3"}.fa-creative-commons-sa:before{content:"\f4ef"}.fa-safari:before{content:"\f267"}.fa-google:before{content:"\f1a0"}.fa-font-awesome-alt:before,.fa-square-font-awesome-stroke:before{content:"\f35c"}.fa-atlassian:before{content:"\f77b"}.fa-linkedin-in:before{content:"\f0e1"}.fa-digital-ocean:before{content:"\f391"}.fa-nimblr:before{content:"\f5a8"}.fa-chromecast:before{content:"\f838"}.fa-evernote:before{content:"\f839"}.fa-hacker-news:before{content:"\f1d4"}.fa-creative-commons-sampling:before{content:"\f4f0"}.fa-adversal:before{content:"\f36a"}.fa-creative-commons:before{content:"\f25e"}.fa-watchman-monitoring:before{content:"\e087"}.fa-fonticons:before{content:"\f280"}.fa-weixin:before{content:"\f1d7"}.fa-shirtsinbulk:before{content:"\f214"}.fa-codepen:before{content:"\f1cb"}.fa-git-alt:before{content:"\f841"}.fa-lyft:before{content:"\f3c3"}.fa-rev:before{content:"\f5b2"}.fa-windows:before{content:"\f17a"}.fa-wizards-of-the-coast:before{content:"\f730"}.fa-square-viadeo:before,.fa-viadeo-square:before{content:"\f2aa"}.fa-meetup:before{content:"\f2e0"}.fa-centos:before{content:"\f789"}.fa-adn:before{content:"\f170"}.fa-cloudsmith:before{content:"\f384"}.fa-opensuse:before{content:"\e62b"}.fa-pied-piper-alt:before{content:"\f1a8"}.fa-dribbble-square:before,.fa-square-dribbble:before{content:"\f397"}.fa-codiepie:before{content:"\f284"}.fa-node:before{content:"\f419"}.fa-mix:before{content:"\f3cb"}.fa-steam:before{content:"\f1b6"}.fa-cc-apple-pay:before{content:"\f416"}.fa-scribd:before{content:"\f28a"}.fa-debian:before{content:"\e60b"}.fa-openid:before{content:"\f19b"}.fa-instalod:before{content:"\e081"}.fa-expeditedssl:before{content:"\f23e"}.fa-sellcast:before{content:"\f2da"}.fa-square-twitter:before,.fa-twitter-square:before{content:"\f081"}.fa-r-project:before{content:"\f4f7"}.fa-delicious:before{content:"\f1a5"}.fa-freebsd:before{content:"\f3a4"}.fa-vuejs:before{content:"\f41f"}.fa-accusoft:before{content:"\f369"}.fa-ioxhost:before{content:"\f208"}.fa-fonticons-fi:before{content:"\f3a2"}.fa-app-store:before{content:"\f36f"}.fa-cc-mastercard:before{content:"\f1f1"}.fa-itunes-note:before{content:"\f3b5"}.fa-golang:before{content:"\e40f"}.fa-kickstarter:before{content:"\f3bb"}.fa-grav:before{content:"\f2d6"}.fa-weibo:before{content:"\f18a"}.fa-uncharted:before{content:"\e084"}.fa-firstdraft:before{content:"\f3a1"}.fa-square-youtube:before,.fa-youtube-square:before{content:"\f431"}.fa-wikipedia-w:before{content:"\f266"}.fa-rendact:before,.fa-wpressr:before{content:"\f3e4"}.fa-angellist:before{content:"\f209"}.fa-galactic-republic:before{content:"\f50c"}.fa-nfc-directional:before{content:"\e530"}.fa-skype:before{content:"\f17e"}.fa-joget:before{content:"\f3b7"}.fa-fedora:before{content:"\f798"}.fa-stripe-s:before{content:"\f42a"}.fa-meta:before{content:"\e49b"}.fa-laravel:before{content:"\f3bd"}.fa-hotjar:before{content:"\f3b1"}.fa-bluetooth-b:before{content:"\f294"}.fa-square-letterboxd:before{content:"\e62e"}.fa-sticker-mule:before{content:"\f3f7"}.fa-creative-commons-zero:before{content:"\f4f3"}.fa-hips:before{content:"\f452"}.fa-behance:before{content:"\f1b4"}.fa-reddit:before{content:"\f1a1"}.fa-discord:before{content:"\f392"}.fa-chrome:before{content:"\f268"}.fa-app-store-ios:before{content:"\f370"}.fa-cc-discover:before{content:"\f1f2"}.fa-wpbeginner:before{content:"\f297"}.fa-confluence:before{content:"\f78d"}.fa-shoelace:before{content:"\e60c"}.fa-mdb:before{content:"\f8ca"}.fa-dochub:before{content:"\f394"}.fa-accessible-icon:before{content:"\f368"}.fa-ebay:before{content:"\f4f4"}.fa-amazon:before{content:"\f270"}.fa-unsplash:before{content:"\e07c"}.fa-yarn:before{content:"\f7e3"}.fa-square-steam:before,.fa-steam-square:before{content:"\f1b7"}.fa-500px:before{content:"\f26e"}.fa-square-vimeo:before,.fa-vimeo-square:before{content:"\f194"}.fa-asymmetrik:before{content:"\f372"}.fa-font-awesome-flag:before,.fa-font-awesome-logo-full:before,.fa-font-awesome:before{content:"\f2b4"}.fa-gratipay:before{content:"\f184"}.fa-apple:before{content:"\f179"}.fa-hive:before{content:"\e07f"}.fa-gitkraken:before{content:"\f3a6"}.fa-keybase:before{content:"\f4f5"}.fa-apple-pay:before{content:"\f415"}.fa-padlet:before{content:"\e4a0"}.fa-amazon-pay:before{content:"\f42c"}.fa-github-square:before,.fa-square-github:before{content:"\f092"}.fa-stumbleupon:before{content:"\f1a4"}.fa-fedex:before{content:"\f797"}.fa-phoenix-framework:before{content:"\f3dc"}.fa-shopify:before{content:"\e057"}.fa-neos:before{content:"\f612"}.fa-square-threads:before{content:"\e619"}.fa-hackerrank:before{content:"\f5f7"}.fa-researchgate:before{content:"\f4f8"}.fa-swift:before{content:"\f8e1"}.fa-angular:before{content:"\f420"}.fa-speakap:before{content:"\f3f3"}.fa-angrycreative:before{content:"\f36e"}.fa-y-combinator:before{content:"\f23b"}.fa-empire:before{content:"\f1d1"}.fa-envira:before{content:"\f299"}.fa-google-scholar:before{content:"\e63b"}.fa-gitlab-square:before,.fa-square-gitlab:before{content:"\e5ae"}.fa-studiovinari:before{content:"\f3f8"}.fa-pied-piper:before{content:"\f2ae"}.fa-wordpress:before{content:"\f19a"}.fa-product-hunt:before{content:"\f288"}.fa-firefox:before{content:"\f269"}.fa-linode:before{content:"\f2b8"}.fa-goodreads:before{content:"\f3a8"}.fa-odnoklassniki-square:before,.fa-square-odnoklassniki:before{content:"\f264"}.fa-jsfiddle:before{content:"\f1cc"}.fa-sith:before{content:"\f512"}.fa-themeisle:before{content:"\f2b2"}.fa-page4:before{content:"\f3d7"}.fa-hashnode:before{content:"\e499"}.fa-react:before{content:"\f41b"}.fa-cc-paypal:before{content:"\f1f4"}.fa-squarespace:before{content:"\f5be"}.fa-cc-stripe:before{content:"\f1f5"}.fa-creative-commons-share:before{content:"\f4f2"}.fa-bitcoin:before{content:"\f379"}.fa-keycdn:before{content:"\f3ba"}.fa-opera:before{content:"\f26a"}.fa-itch-io:before{content:"\f83a"}.fa-umbraco:before{content:"\f8e8"}.fa-galactic-senate:before{content:"\f50d"}.fa-ubuntu:before{content:"\f7df"}.fa-draft2digital:before{content:"\f396"}.fa-stripe:before{content:"\f429"}.fa-houzz:before{content:"\f27c"}.fa-gg:before{content:"\f260"}.fa-dhl:before{content:"\f790"}.fa-pinterest-square:before,.fa-square-pinterest:before{content:"\f0d3"}.fa-xing:before{content:"\f168"}.fa-blackberry:before{content:"\f37b"}.fa-creative-commons-pd:before{content:"\f4ec"}.fa-playstation:before{content:"\f3df"}.fa-quinscape:before{content:"\f459"}.fa-less:before{content:"\f41d"}.fa-blogger-b:before{content:"\f37d"}.fa-opencart:before{content:"\f23d"}.fa-vine:before{content:"\f1ca"}.fa-signal-messenger:before{content:"\e663"}.fa-paypal:before{content:"\f1ed"}.fa-gitlab:before{content:"\f296"}.fa-typo3:before{content:"\f42b"}.fa-reddit-alien:before{content:"\f281"}.fa-yahoo:before{content:"\f19e"}.fa-dailymotion:before{content:"\e052"}.fa-affiliatetheme:before{content:"\f36b"}.fa-pied-piper-pp:before{content:"\f1a7"}.fa-bootstrap:before{content:"\f836"}.fa-odnoklassniki:before{content:"\f263"}.fa-nfc-symbol:before{content:"\e531"}.fa-mintbit:before{content:"\e62f"}.fa-ethereum:before{content:"\f42e"}.fa-speaker-deck:before{content:"\f83c"}.fa-creative-commons-nc-eu:before{content:"\f4e9"}.fa-patreon:before{content:"\f3d9"}.fa-avianex:before{content:"\f374"}.fa-ello:before{content:"\f5f1"}.fa-gofore:before{content:"\f3a7"}.fa-bimobject:before{content:"\f378"}.fa-brave-reverse:before{content:"\e63d"}.fa-facebook-f:before{content:"\f39e"}.fa-google-plus-square:before,.fa-square-google-plus:before{content:"\f0d4"}.fa-mandalorian:before{content:"\f50f"}.fa-first-order-alt:before{content:"\f50a"}.fa-osi:before{content:"\f41a"}.fa-google-wallet:before{content:"\f1ee"}.fa-d-and-d-beyond:before{content:"\f6ca"}.fa-periscope:before{content:"\f3da"}.fa-fulcrum:before{content:"\f50b"}.fa-cloudscale:before{content:"\f383"}.fa-forumbee:before{content:"\f211"}.fa-mizuni:before{content:"\f3cc"}.fa-schlix:before{content:"\f3ea"}.fa-square-xing:before,.fa-xing-square:before{content:"\f169"}.fa-bandcamp:before{content:"\f2d5"}.fa-wpforms:before{content:"\f298"}.fa-cloudversify:before{content:"\f385"}.fa-usps:before{content:"\f7e1"}.fa-megaport:before{content:"\f5a3"}.fa-magento:before{content:"\f3c4"}.fa-spotify:before{content:"\f1bc"}.fa-optin-monster:before{content:"\f23c"}.fa-fly:before{content:"\f417"}.fa-aviato:before{content:"\f421"}.fa-itunes:before{content:"\f3b4"}.fa-cuttlefish:before{content:"\f38c"}.fa-blogger:before{content:"\f37c"}.fa-flickr:before{content:"\f16e"}.fa-viber:before{content:"\f409"}.fa-soundcloud:before{content:"\f1be"}.fa-digg:before{content:"\f1a6"}.fa-tencent-weibo:before{content:"\f1d5"}.fa-letterboxd:before{content:"\e62d"}.fa-symfony:before{content:"\f83d"}.fa-maxcdn:before{content:"\f136"}.fa-etsy:before{content:"\f2d7"}.fa-facebook-messenger:before{content:"\f39f"}.fa-audible:before{content:"\f373"}.fa-think-peaks:before{content:"\f731"}.fa-bilibili:before{content:"\e3d9"}.fa-erlang:before{content:"\f39d"}.fa-x-twitter:before{content:"\e61b"}.fa-cotton-bureau:before{content:"\f89e"}.fa-dashcube:before{content:"\f210"}.fa-42-group:before,.fa-innosoft:before{content:"\e080"}.fa-stack-exchange:before{content:"\f18d"}.fa-elementor:before{content:"\f430"}.fa-pied-piper-square:before,.fa-square-pied-piper:before{content:"\e01e"}.fa-creative-commons-nd:before{content:"\f4eb"}.fa-palfed:before{content:"\f3d8"}.fa-superpowers:before{content:"\f2dd"}.fa-resolving:before{content:"\f3e7"}.fa-xbox:before{content:"\f412"}.fa-searchengin:before{content:"\f3eb"}.fa-tiktok:before{content:"\e07b"}.fa-facebook-square:before,.fa-square-facebook:before{content:"\f082"}.fa-renren:before{content:"\f18b"}.fa-linux:before{content:"\f17c"}.fa-glide:before{content:"\f2a5"}.fa-linkedin:before{content:"\f08c"}.fa-hubspot:before{content:"\f3b2"}.fa-deploydog:before{content:"\f38e"}.fa-twitch:before{content:"\f1e8"}.fa-ravelry:before{content:"\f2d9"}.fa-mixer:before{content:"\e056"}.fa-lastfm-square:before,.fa-square-lastfm:before{content:"\f203"}.fa-vimeo:before{content:"\f40a"}.fa-mendeley:before{content:"\f7b3"}.fa-uniregistry:before{content:"\f404"}.fa-figma:before{content:"\f799"}.fa-creative-commons-remix:before{content:"\f4ee"}.fa-cc-amazon-pay:before{content:"\f42d"}.fa-dropbox:before{content:"\f16b"}.fa-instagram:before{content:"\f16d"}.fa-cmplid:before{content:"\e360"}.fa-upwork:before{content:"\e641"}.fa-facebook:before{content:"\f09a"}.fa-gripfire:before{content:"\f3ac"}.fa-jedi-order:before{content:"\f50e"}.fa-uikit:before{content:"\f403"}.fa-fort-awesome-alt:before{content:"\f3a3"}.fa-phabricator:before{content:"\f3db"}.fa-ussunnah:before{content:"\f407"}.fa-earlybirds:before{content:"\f39a"}.fa-trade-federation:before{content:"\f513"}.fa-autoprefixer:before{content:"\f41c"}.fa-whatsapp:before{content:"\f232"}.fa-slideshare:before{content:"\f1e7"}.fa-google-play:before{content:"\f3ab"}.fa-viadeo:before{content:"\f2a9"}.fa-line:before{content:"\f3c0"}.fa-google-drive:before{content:"\f3aa"}.fa-servicestack:before{content:"\f3ec"}.fa-simplybuilt:before{content:"\f215"}.fa-bitbucket:before{content:"\f171"}.fa-imdb:before{content:"\f2d8"}.fa-deezer:before{content:"\e077"}.fa-raspberry-pi:before{content:"\f7bb"}.fa-jira:before{content:"\f7b1"}.fa-docker:before{content:"\f395"}.fa-screenpal:before{content:"\e570"}.fa-bluetooth:before{content:"\f293"}.fa-gitter:before{content:"\f426"}.fa-d-and-d:before{content:"\f38d"}.fa-microblog:before{content:"\e01a"}.fa-cc-diners-club:before{content:"\f24c"}.fa-gg-circle:before{content:"\f261"}.fa-pied-piper-hat:before{content:"\f4e5"}.fa-kickstarter-k:before{content:"\f3bc"}.fa-yandex:before{content:"\f413"}.fa-readme:before{content:"\f4d5"}.fa-html5:before{content:"\f13b"}.fa-sellsy:before{content:"\f213"}.fa-sass:before{content:"\f41e"}.fa-wirsindhandwerk:before,.fa-wsh:before{content:"\e2d0"}.fa-buromobelexperte:before{content:"\f37f"}.fa-salesforce:before{content:"\f83b"}.fa-octopus-deploy:before{content:"\e082"}.fa-medapps:before{content:"\f3c6"}.fa-ns8:before{content:"\f3d5"}.fa-pinterest-p:before{content:"\f231"}.fa-apper:before{content:"\f371"}.fa-fort-awesome:before{content:"\f286"}.fa-waze:before{content:"\f83f"}.fa-cc-jcb:before{content:"\f24b"}.fa-snapchat-ghost:before,.fa-snapchat:before{content:"\f2ab"}.fa-fantasy-flight-games:before{content:"\f6dc"}.fa-rust:before{content:"\e07a"}.fa-wix:before{content:"\f5cf"}.fa-behance-square:before,.fa-square-behance:before{content:"\f1b5"}.fa-supple:before{content:"\f3f9"}.fa-webflow:before{content:"\e65c"}.fa-rebel:before{content:"\f1d0"}.fa-css3:before{content:"\f13c"}.fa-staylinked:before{content:"\f3f5"}.fa-kaggle:before{content:"\f5fa"}.fa-space-awesome:before{content:"\e5ac"}.fa-deviantart:before{content:"\f1bd"}.fa-cpanel:before{content:"\f388"}.fa-goodreads-g:before{content:"\f3a9"}.fa-git-square:before,.fa-square-git:before{content:"\f1d2"}.fa-square-tumblr:before,.fa-tumblr-square:before{content:"\f174"}.fa-trello:before{content:"\f181"}.fa-creative-commons-nc-jp:before{content:"\f4ea"}.fa-get-pocket:before{content:"\f265"}.fa-perbyte:before{content:"\e083"}.fa-grunt:before{content:"\f3ad"}.fa-weebly:before{content:"\f5cc"}.fa-connectdevelop:before{content:"\f20e"}.fa-leanpub:before{content:"\f212"}.fa-black-tie:before{content:"\f27e"}.fa-themeco:before{content:"\f5c6"}.fa-python:before{content:"\f3e2"}.fa-android:before{content:"\f17b"}.fa-bots:before{content:"\e340"}.fa-free-code-camp:before{content:"\f2c5"}.fa-hornbill:before{content:"\f592"}.fa-js:before{content:"\f3b8"}.fa-ideal:before{content:"\e013"}.fa-git:before{content:"\f1d3"}.fa-dev:before{content:"\f6cc"}.fa-sketch:before{content:"\f7c6"}.fa-yandex-international:before{content:"\f414"}.fa-cc-amex:before{content:"\f1f3"}.fa-uber:before{content:"\f402"}.fa-github:before{content:"\f09b"}.fa-php:before{content:"\f457"}.fa-alipay:before{content:"\f642"}.fa-youtube:before{content:"\f167"}.fa-skyatlas:before{content:"\f216"}.fa-firefox-browser:before{content:"\e007"}.fa-replyd:before{content:"\f3e6"}.fa-suse:before{content:"\f7d6"}.fa-jenkins:before{content:"\f3b6"}.fa-twitter:before{content:"\f099"}.fa-rockrms:before{content:"\f3e9"}.fa-pinterest:before{content:"\f0d2"}.fa-buffer:before{content:"\f837"}.fa-npm:before{content:"\f3d4"}.fa-yammer:before{content:"\f840"}.fa-btc:before{content:"\f15a"}.fa-dribbble:before{content:"\f17d"}.fa-stumbleupon-circle:before{content:"\f1a3"}.fa-internet-explorer:before{content:"\f26b"}.fa-stubber:before{content:"\e5c7"}.fa-telegram-plane:before,.fa-telegram:before{content:"\f2c6"}.fa-old-republic:before{content:"\f510"}.fa-odysee:before{content:"\e5c6"}.fa-square-whatsapp:before,.fa-whatsapp-square:before{content:"\f40c"}.fa-node-js:before{content:"\f3d3"}.fa-edge-diffuserve.legacy:before{content:"\e078"}.fa-slack-hash:before,.fa-slack:before{content:"\f198"}.fa-medrt:before{content:"\f3c8"}.fa-usb:before{content:"\f287"}.fa-tumblr:before{content:"\f173"}.fa-vaadin:before{content:"\f408"}.fa-quora:before{content:"\f2c4"}.fa-square-x-twitter:before{content:"\e61a"}.fa-reacteurope:before{content:"\f75d"}.fa-medium-m:before,.fa-medium:before{content:"\f23a"}.fa-amilia:before{content:"\f36d"}.fa-mixcloud:before{content:"\f289"}.fa-flipboard:before{content:"\f44d"}.fa-viacoin:before{content:"\f237"}.fa-critical-role:before{content:"\f6c9"}.fa-sitrox:before{content:"\e44a"}.fa-discourse:before{content:"\f393"}.fa-joomla:before{content:"\f1aa"}.fa-mastodon:before{content:"\f4f6"}.fa-airbnb:before{content:"\f834"}.fa-wolf-pack-battalion:before{content:"\f514"}.fa-buy-n-large:before{content:"\f8a6"}.fa-gulp:before{content:"\f3ae"}.fa-creative-commons-sampling-plus:before{content:"\f4f1"}.fa-strava:before{content:"\f428"}.fa-ember:before{content:"\f423"}.fa-canadian-maple-leaf:before{content:"\f785"}.fa-teamspeak:before{content:"\f4f9"}.fa-pushed:before{content:"\f3e1"}.fa-wordpress-simple:before{content:"\f411"}.fa-nutritionix:before{content:"\f3d6"}.fa-wodu:before{content:"\e088"}.fa-google-pay:before{content:"\e079"}.fa-intercom:before{content:"\f7af"}.fa-zhihu:before{content:"\f63f"}.fa-korvue:before{content:"\f42f"}.fa-pix:before{content:"\e43a"}.fa-steam-symbol:before{content:"\f3f6"} \ 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 086dda69..aab4ca4b 100755 --- a/examples/test_causal_lm_decoding_kernel.py +++ b/examples/test_causal_lm_decoding_kernel.py @@ -1,6 +1,6 @@ import torch -from d2f_engine.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd +from diffuserve.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 2e961a98..a0da171b 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 d2f_engine.layers.attention.ops import diffusion_lm_parallel_flash_decoding, diffusion_lm_flash_decoding +from diffuserve.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 8b9b4793..6d378040 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 d2f_engine.layers.attention.ops import store_kvcache_unified_layout, load_kvcache, CHECK_LOADING +from diffuserve.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 35cdabbf..2d88e5bf 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 d2f_engine.layers.attention.attention_v4 import store_kvcache_distinct_layout, store_kvcache_unified +from diffuserve.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 a03057b5..357a707c 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -9,7 +9,7 @@ from viztracer import VizTracer from transformers import AutoTokenizer -from d2f_engine import LLM, SamplingParams +from diffuserve.legacy import LLM, SamplingParams def summarize_profiling(csv_path: str) -> dict: diff --git a/examples/test_dream_dvllm_human_eval.py b/examples/test_dream_dvllm_human_eval.py index b0e7f2bd..fbc89c0c 100755 --- a/examples/test_dream_dvllm_human_eval.py +++ b/examples/test_dream_dvllm_human_eval.py @@ -8,7 +8,7 @@ from viztracer import VizTracer from transformers import AutoTokenizer -from d2f_engine import LLM, SamplingParams +from diffuserve.legacy import LLM, SamplingParams def summarize_profiling(csv_path: str) -> dict: diff --git a/examples/test_dream_model_weight.py b/examples/test_dream_model_weight.py index d0b413b8..aefece5f 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 d2f_engine.config import Config -from d2f_engine.models.auto_model import AutoModelLM +from diffuserve.legacy.config import Config +from diffuserve.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 a8bc6c95..da75a3ae 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 d2f_engine.config import Config -from d2f_engine.engine.model_runner import AutoModelRunner +from diffuserve.legacy.config import Config +from diffuserve.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_llada_dvllm_human_eval.py b/examples/test_llada_dvllm_human_eval.py index 45587b0e..502b9cad 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 d2f_engine import LLM, SamplingParams +from diffuserve.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 61411e51..fb5740f6 100755 --- a/examples/test_qwen_dvllm.py +++ b/examples/test_qwen_dvllm.py @@ -1,6 +1,6 @@ import os -from d2f_engine import LLM, SamplingParams +from diffuserve.legacy import LLM, SamplingParams from viztracer import VizTracer diff --git a/pyproject.toml b/pyproject.toml index 1a7141ce..b9c3460c 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ Repository = "https://zhijie-group.github.io/D2fEngine" "Organization" = "https://github.com/zhijie-group" [tool.setuptools.packages.find] -include = ["d2f_engine"] +include = ["diffuserve"] [[tool.uv.index]] url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" From f82abad0c9c05290b99103b5e0ac87eebe897536 Mon Sep 17 00:00:00 2001 From: drewjin Date: Wed, 5 Nov 2025 07:30:19 +0000 Subject: [PATCH 05/23] feat(model): finish model registration logics --- diffuserve/__init__.py | 2 + diffuserve/config.py | 63 + diffuserve/engine/block_manager.py | 174 +++ diffuserve/engine/dp_engine.py | 295 +++++ diffuserve/engine/llm_engine.py | 147 +++ diffuserve/engine/model_runner.py | 507 ++++++++ diffuserve/engine/scheduler.py | 234 ++++ diffuserve/engine/sequence.py | 552 +++++++++ diffuserve/layer/activation.py | 14 + diffuserve/layer/attention/attention_v5.py | 148 +++ diffuserve/layer/attention/ops/__init__.py | 7 + ...chunked_prefill_decoding_unified_kernel.py | 375 ++++++ .../layer/attention/ops/kv_cache_kernels.py | 395 ++++++ .../layer/attention/ops/prefix_prefill.py | 1090 +++++++++++++++++ .../attention/ops/tilus_decode_attn_dlm.py | 161 +++ .../attention/ops/triton_decode_attn_clm.py | 681 ++++++++++ .../attention/ops/triton_decode_attn_dlm.py | 120 ++ .../attention/ops/triton_flash_attention.py | 1022 ++++++++++++++++ diffuserve/layer/embed_head.py | 73 ++ diffuserve/layer/layernorm.py | 51 + diffuserve/layer/linear.py | 244 ++++ diffuserve/layer/rotary_embedding.py | 79 ++ diffuserve/layer/sampler.py | 217 ++++ diffuserve/legacy/__init__.py | 1 - diffuserve/llm.py | 10 + diffuserve/model/auto_model.py | 97 ++ .../model/config/dream/configuration_dream.py | 88 ++ .../configuration_fast_dllm_v2.py | 90 ++ .../model/config/llada/configuration_llada.py | 459 +++++++ diffuserve/model/diffucoder.py | 0 diffuserve/model/dream.py | 238 ++++ diffuserve/model/fast_dllm_v2.py | 239 ++++ diffuserve/model/llada.py | 266 ++++ diffuserve/model/llada2.py | 0 diffuserve/model/llada_moe.py | 0 diffuserve/model/sdar.py | 0 diffuserve/model/utils/check_config.py | 8 + diffuserve/sampling_params.py | 8 + diffuserve/utils/checker.py | 28 + diffuserve/utils/context.py | 112 ++ diffuserve/utils/loader.py | 204 +++ pyproject.toml | 8 +- 42 files changed, 8502 insertions(+), 5 deletions(-) create mode 100755 diffuserve/__init__.py create mode 100755 diffuserve/config.py create mode 100755 diffuserve/engine/block_manager.py create mode 100755 diffuserve/engine/dp_engine.py create mode 100755 diffuserve/engine/llm_engine.py create mode 100755 diffuserve/engine/model_runner.py create mode 100755 diffuserve/engine/scheduler.py create mode 100755 diffuserve/engine/sequence.py create mode 100755 diffuserve/layer/activation.py create mode 100644 diffuserve/layer/attention/attention_v5.py create mode 100755 diffuserve/layer/attention/ops/__init__.py create mode 100755 diffuserve/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py create mode 100755 diffuserve/layer/attention/ops/kv_cache_kernels.py create mode 100755 diffuserve/layer/attention/ops/prefix_prefill.py create mode 100755 diffuserve/layer/attention/ops/tilus_decode_attn_dlm.py create mode 100755 diffuserve/layer/attention/ops/triton_decode_attn_clm.py create mode 100755 diffuserve/layer/attention/ops/triton_decode_attn_dlm.py create mode 100755 diffuserve/layer/attention/ops/triton_flash_attention.py create mode 100755 diffuserve/layer/embed_head.py create mode 100755 diffuserve/layer/layernorm.py create mode 100755 diffuserve/layer/linear.py create mode 100755 diffuserve/layer/rotary_embedding.py create mode 100644 diffuserve/layer/sampler.py create mode 100755 diffuserve/llm.py create mode 100755 diffuserve/model/auto_model.py create mode 100755 diffuserve/model/config/dream/configuration_dream.py create mode 100755 diffuserve/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py create mode 100644 diffuserve/model/config/llada/configuration_llada.py create mode 100644 diffuserve/model/diffucoder.py create mode 100755 diffuserve/model/dream.py create mode 100755 diffuserve/model/fast_dllm_v2.py create mode 100755 diffuserve/model/llada.py create mode 100644 diffuserve/model/llada2.py create mode 100644 diffuserve/model/llada_moe.py create mode 100644 diffuserve/model/sdar.py create mode 100755 diffuserve/model/utils/check_config.py create mode 100755 diffuserve/sampling_params.py create mode 100755 diffuserve/utils/checker.py create mode 100755 diffuserve/utils/context.py create mode 100755 diffuserve/utils/loader.py diff --git a/diffuserve/__init__.py b/diffuserve/__init__.py new file mode 100755 index 00000000..82b9e51b --- /dev/null +++ b/diffuserve/__init__.py @@ -0,0 +1,2 @@ +from diffuserve.legacy.llm import LLM +from diffuserve.legacy.sampling_params import SamplingParams diff --git a/diffuserve/config.py b/diffuserve/config.py new file mode 100755 index 00000000..e5533921 --- /dev/null +++ b/diffuserve/config.py @@ -0,0 +1,63 @@ +import os +from dataclasses import dataclass +from transformers import AutoConfig + + +@dataclass +class Config: + model: str + lora_path: str = "" + model_name: str = "dream" + decoding_strategy: str = "d2f" # "d2f", "fast-dllm-v2", "block-diffusion" + + mask_token_id: int = 151666 + diffusion_block_size: int = 32 + + accept_threshold: float = 0.9 + complete_threshold: float = 0.95 + add_new_block_threshold: float = 0.1 + + use_lora: bool = False + max_num_batched_tokens: int = 4096 + max_num_seqs: int = 128 + max_model_len: int = 2048 + gpu_memory_utilization: float = 0.9 + + data_parallel_size: int = 1 + tensor_parallel_size: int = 2 + # Distributed comm (per tensor-parallel group). When using multiple DP + # replicas on one host, assign unique master_port per replica. + master_addr: str = "localhost" + master_port: int = 2333 + # Shared memory segment name for intra-TP RPC; must be unique per DP group. + shm_name: str = "diffuserve_shm" + # Start device index for this TP group (set by DP launcher). + device_start: int = 0 + + enforce_eager: bool = False + hf_config: AutoConfig | None = None + eos: int = -1 + kvcache_block_size: int = 256 + num_kvcache_blocks: int = -1 + k_cache_hdim_split_factor_x: int = 8 + kv_cache_layout: str = "unified" # "unified" or "distinct" + + def __post_init__(self): + assert os.path.isdir(self.model) + assert self.kvcache_block_size % 16 == 0 + assert 1 <= self.tensor_parallel_size <= 8 + assert 1 <= self.data_parallel_size <= 1024 + assert isinstance(self.master_port, int) and 0 < self.master_port < 65536 + assert isinstance(self.device_start, int) and self.device_start >= 0 + + # LoRA validation + if self.use_lora: + if not self.lora_path: + raise ValueError("lora_path must be provided when use_lora is True") + if not os.path.exists(self.lora_path): + print(f"Warning: LoRA path {self.lora_path} does not exist") + + self.hf_config = AutoConfig.from_pretrained(self.model, trust_remote_code=True) + cfg_max_model_len = self.hf_config.max_position_embeddings if hasattr(self.hf_config, "max_position_embeddings") else self.hf_config.max_sequence_length + self.max_model_len = min(self.max_model_len, cfg_max_model_len) + assert self.max_num_batched_tokens >= self.max_model_len \ No newline at end of file diff --git a/diffuserve/engine/block_manager.py b/diffuserve/engine/block_manager.py new file mode 100755 index 00000000..d822947b --- /dev/null +++ b/diffuserve/engine/block_manager.py @@ -0,0 +1,174 @@ +import xxhash + +import numpy as np + +from collections import deque +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import List, Dict, Deque, Set + +from diffuserve.config import Config +from diffuserve.engine.sequence import SequenceBase, SequenceForCausalLM, SequenceForDiffusionLM + + +@dataclass +class Block: + block_id: int + ref_count: int = 0 + hash: int = -1 + token_ids: List[int] = field(default_factory=list) + + def update(self, hash: int, token_ids: list[int]): + self.hash = hash + self.token_ids = token_ids + + def reset(self): + self.ref_count = 1 + self.hash = -1 + self.token_ids = [] + + +class BlockManagerBase(ABC): + def __init__(self, num_blocks: int, block_size: int): + assert num_blocks > 0 + self.block_size = block_size + self.blocks: List[Block] = [Block(block_id=i) for i in range(num_blocks)] + self.hash_to_block_id: Dict[int, int] = dict() + self.free_block_ids: Deque[int] = deque(range(num_blocks)) + self.used_block_ids: Set[int] = set() + + @classmethod + def compute_hash(cls, token_ids: List[int], prefix: int = -1): + h = xxhash.xxh64() + if prefix != -1: + h.update(prefix.to_bytes(8, "little")) + h.update(np.array(token_ids).tobytes()) + return h.intdigest() + + def _allocate_block(self, block_id: int) -> Block: + block = self.blocks[block_id] + assert block.ref_count == 0 + block.reset() + self.free_block_ids.remove(block_id) + self.used_block_ids.add(block_id) + return self.blocks[block_id] + + def _free_block(self, block_id: int) -> Block: + assert self.blocks[block_id].ref_count == 0 + self.used_block_ids.remove(block_id) + self.free_block_ids.append(block_id) + + def can_allocate(self, seq: SequenceBase) -> bool: + return len(self.free_block_ids) >= seq.num_blocks + + def allocate(self, seq: SequenceBase): + assert not seq.block_table + h = -1 + cache_miss = False + for i in range(seq.num_blocks): + token_ids = seq.block(i) + h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1 + block_id = self.hash_to_block_id.get(h, -1) + if block_id == -1 or self.blocks[block_id].token_ids != token_ids: + cache_miss = True + seq.block_cache_missed.append(cache_miss) + if cache_miss: + block_id = self.free_block_ids[0] + block = self._allocate_block(block_id) + else: + seq.num_cached_tokens += self.block_size + if block_id in self.used_block_ids: + block = self.blocks[block_id] + block.ref_count += 1 + else: + block = self._allocate_block(block_id) + if h != -1: + block.update(h, token_ids) + self.hash_to_block_id[h] = block_id + seq.block_table.append(block_id) + + def free(self, seq: SequenceBase): + for block_id in reversed(seq.block_table): + block = self.blocks[block_id] + block.ref_count -= 1 + if block.ref_count == 0: + self._free_block(block_id) + seq.num_cached_tokens = 0 + seq.block_table.clear() + + @abstractmethod + def can_append(self, seq: SequenceBase) -> bool: + pass + + @abstractmethod + def may_append(self, seq: SequenceBase): + pass + + +class BlockManagerForCausalLM(BlockManagerBase): + def can_append(self, seq: SequenceForCausalLM) -> bool: + return len(self.free_block_ids) >= (len(seq) % self.block_size == 1) + + def may_append(self, seq: SequenceBase): + block_table = seq.block_table + last_block = self.blocks[block_table[-1]] + if len(seq) % self.block_size == 1: + assert last_block.hash != -1 + block_id = self.free_block_ids[0] + self._allocate_block(block_id) + block_table.append(block_id) + elif len(seq) % self.block_size == 0: + assert last_block.hash == -1 + token_ids = seq.block(seq.num_blocks-1) + prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 + h = self.compute_hash(token_ids, prefix) + last_block.update(h, token_ids) + self.hash_to_block_id[h] = last_block.block_id + else: + assert last_block.hash == -1 + +class BlockManagerForDiffusionLM(BlockManagerBase): + def can_append(self, seq: SequenceForDiffusionLM) -> bool: + return len(self.free_block_ids) >= (seq.cached_or_caching_num_tokens % self.block_size == 1) + + def may_append(self, seq: SequenceForDiffusionLM): + # Handle edge case when no tokens are cached yet + if seq.cached_or_caching_num_tokens == 0: + return + + block_table = seq.block_table + if not block_table: + return + + last_block = self.blocks[block_table[-1]] + + if seq.cached_or_caching_num_tokens // self.block_size == len(seq.block_table): + if last_block.hash == -1: + prev_block_end_token = seq.cached_or_caching_num_tokens - seq.caching_num_tokens - 1 # 256th token (0-indexed: 255) + prev_block_idx = prev_block_end_token // self.block_size # block containing 255th token + + if prev_block_idx < seq.num_blocks: + # This block should be full, so set its hash + token_ids = seq.block(prev_block_idx) + prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 + h = self.compute_hash(token_ids, prefix) + last_block.update(h, token_ids) + self.hash_to_block_id[h] = last_block.block_id + + # Now allocate a new block + block_id = self.free_block_ids[0] + self._allocate_block(block_id) + block_table.append(block_id) + + +class AutoBlockManager(BlockManagerBase): + BLOCK_MANAGER_MAPPING = { + "causal_lm": BlockManagerForCausalLM, + "diffusion_lm": BlockManagerForDiffusionLM, + } + @classmethod + def from_config(cls, config: Config) -> BlockManagerBase: + block_manager_cls = cls.BLOCK_MANAGER_MAPPING.get(config.model_type) + if not block_manager_cls: + raise ValueError(f"Unsupported model type: {config.model_type}") + return block_manager_cls(config.num_kvcache_blocks, config.kvcache_block_size) \ No newline at end of file diff --git a/diffuserve/engine/dp_engine.py b/diffuserve/engine/dp_engine.py new file mode 100755 index 00000000..968ed43b --- /dev/null +++ b/diffuserve/engine/dp_engine.py @@ -0,0 +1,295 @@ +import os +import sys +import torch +import atexit +import traceback +import faulthandler + +import multiprocessing as mp + +from typing import List, Any +from multiprocessing.connection import wait as mp_wait + +from diffuserve.config import Config +from diffuserve.engine.llm_engine import LLMEngine +from diffuserve.sampling_params import SamplingParams + + +def _dp_child_entry(config: Config, dp_idx: int, local_devices: list[int], conn): + """Child process entry point: create an LLMEngine for this DP rank and serve RPC via Pipe.""" + try: + # Enable Python-level crash diagnostics for hard crashes (segfault, OOM kill signals, etc.). + try: + faulthandler.enable(all_threads=True) + except Exception: + pass + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(x) for x in local_devices) + cfg = Config( + model=config.model, + lora_path=config.lora_path, + model_name=config.model_name, + model_type=config.model_type, + mask_token_id=config.mask_token_id, + diffusion_block_size=config.diffusion_block_size, + accept_threshold=config.accept_threshold, + complete_threshold=config.complete_threshold, + add_new_block_threshold=config.add_new_block_threshold, + use_lora=config.use_lora, + max_num_batched_tokens=config.max_num_batched_tokens, + max_num_seqs=config.max_num_seqs, + max_model_len=config.max_model_len, + gpu_memory_utilization=config.gpu_memory_utilization, + data_parallel_size=1, + tensor_parallel_size=config.tensor_parallel_size, + master_addr=config.master_addr, + master_port=int(config.master_port) + dp_idx, + shm_name=f"{config.shm_name}_{dp_idx}", + enforce_eager=config.enforce_eager, + kvcache_block_size=config.kvcache_block_size, + num_kvcache_blocks=config.num_kvcache_blocks, + k_cache_hdim_split_factor_x=config.k_cache_hdim_split_factor_x, + kv_cache_layout=config.kv_cache_layout, + ) + setattr(cfg, "device_start", 0) + + engine = LLMEngine(cfg.model, **{k: getattr(cfg, k) for k in cfg.__dataclass_fields__.keys() if k != "model"}) + + while True: + msg = conn.recv() + if not msg: + continue + cmd, *args = msg + if cmd == "generate": + plist, sp_arg, use_tqdm = args + out = engine.generate(plist, sp_arg, use_tqdm) + conn.send(("ok", out)) + continue + if cmd == "exit": + engine.exit() + conn.send(("ok", None)) + break + elif cmd == "add_request": + prompt, sp = args + sid = engine.add_request(prompt, sp) + conn.send(("ok", sid)) + elif cmd == "step": + out = engine.step() + conn.send(("ok", out)) + elif cmd == "is_finished": + conn.send(("ok", engine.is_finished())) + else: + conn.send(("err", f"unknown_cmd:{cmd}")) + except Exception as e: + # Include full traceback for easier debugging and also print to stderr as a fallback. + tb = traceback.format_exc() + msg = f"{type(e).__name__}: {e}\n{tb}" + try: + conn.send(("err", msg)) + except Exception: + pass + try: + print(f"[DP Child {dp_idx}] Unhandled exception:\n{msg}", file=sys.stderr, flush=True) + except Exception: + pass + + +class DPEngine: + """Data-parallel wrapper that runs N independent TP groups as child processes and aggregates results.""" + def __init__(self, model, **kwargs): + config_fields = {f for f in Config.__dataclass_fields__.keys()} + config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} + self.config = cfg = Config(model, **config_kwargs) + self.dp_size = cfg.data_parallel_size + assert self.dp_size > 1, "Use LLMEngine directly when data_parallel_size == 1" + + ctx = mp.get_context("spawn") + self.conns: List[Any] = [] + self.ps: List[mp.Process] = [] + # Topology check and mapping + base_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if base_visible: + vis = [int(x) for x in base_visible.split(',') if x.strip() != ''] + else: + vis = list(range(torch.cuda.device_count())) + + need_gpus = self.dp_size * cfg.tensor_parallel_size + assert len(vis) >= need_gpus, f"Require {need_gpus} GPUs (dp={self.dp_size}, tp={cfg.tensor_parallel_size}), visible {len(vis)}" + + # Optional overrides: kwargs['device_ids'] or env D2F_DEVICE_MAP + override = None + if 'device_ids' in kwargs and kwargs['device_ids']: + override = list(kwargs['device_ids']) + elif os.environ.get('D2F_DEVICE_MAP'): + override = [int(x) for x in os.environ['D2F_DEVICE_MAP'].split(',') if x.strip() != ''] + if override is not None: + assert len(override) >= need_gpus, f"device_ids length {len(override)} < required {need_gpus}" + # All override devices must be in visible list + assert all(d in vis for d in override[:need_gpus]), "device_ids must be subset of CUDA_VISIBLE_DEVICES" + plan = override[:need_gpus] + else: + plan = vis[:need_gpus] + + tp = cfg.tensor_parallel_size + for dp_idx in range(self.dp_size): + local_devices = plan[dp_idx*tp:(dp_idx+1)*tp] + parent_conn, child_conn = ctx.Pipe() + p = ctx.Process(target=_dp_child_entry, args=(cfg, dp_idx, local_devices, child_conn)) + p.start() + self.ps.append(p) + self.conns.append(parent_conn) + self._rr = 0 # round-robin pointer + self._gid_counter = 0 + self._gid_map = {} # (replica, local_id) -> global_id + self._rev_gid_map = {} # global_id -> (replica, local_id) + atexit.register(self.exit) + + def _ask(self, replica: int, cmd: str, *args): + conn = self.conns[replica] + conn.send((cmd, *args)) + try: + tag, payload = conn.recv() + except EOFError: + p = self.ps[replica] + exitcode = p.exitcode + raise RuntimeError( + f"DP child #{replica} terminated unexpectedly (exitcode={exitcode}). " + f"This may indicate OOM or a native crash. Try setting env: " + f"PYTHONFAULTHANDLER=1 CUDA_LAUNCH_BLOCKING=1 TORCH_SHOW_CPP_STACKTRACES=1 to get more diagnostics." + ) + if tag == "ok": + return payload + raise RuntimeError(f"DP child #{replica} error: {payload}") + + def exit(self): + for i, p in enumerate(self.ps): + if p.is_alive(): + try: + self._ask(i, "exit") + except Exception: + pass + p.join(timeout=5) + + def add_request(self, prompt: str | List[int], sampling_params: SamplingParams): + target = self._rr + self._rr = (self._rr + 1) % self.dp_size + local_id = self._ask(target, "add_request", prompt, sampling_params) + gid = self._gid_counter + self._gid_counter += 1 + self._gid_map[(target, local_id)] = gid + self._rev_gid_map[gid] = (target, local_id) + return gid + + def step(self): + all_outputs = [] + total_tokens = 0 + any_prefill = False + merged_diff_steps = {} + merged_deltas = [] + for i in range(self.dp_size): + done = self._ask(i, "is_finished") + if done: + continue + outputs, num_tokens, is_prefill, n_diff_steps, deltas = self._ask(i, "step") + if outputs: + # remap local seq_ids to global ids + for sid, toks in outputs: + gid = self._gid_map.get((i, sid), None) + if gid is not None: + all_outputs.append((gid, toks)) + total_tokens += num_tokens + any_prefill = any_prefill or is_prefill + if n_diff_steps: + merged_diff_steps.update(n_diff_steps) + if deltas: + for sid, toks, fin in deltas: + gid = self._gid_map.get((i, sid), None) + if gid is not None: + merged_deltas.append((gid, toks, fin)) + return all_outputs, total_tokens, any_prefill, merged_diff_steps, merged_deltas + + def is_finished(self): + return all(self._ask(i, "is_finished") for i in range(self.dp_size)) + + def generate(self, prompts: List[str] | List[List[int]], sampling_params: SamplingParams | List[SamplingParams], use_tqdm: bool = True): + """Load-balanced generate with random shuffling and stable order restoration. + - Randomly shuffle inputs to balance load across DP replicas. + - Partition shuffled list evenly among replicas. + - Send to children, collect outputs, then unshuffle to original order. + """ + import random + n = len(prompts) + idxs = list(range(n)) + random.shuffle(idxs) + shuffled_prompts = [prompts[i] for i in idxs] + # Align sampling params with shuffled prompts + if isinstance(sampling_params, list): + if len(sampling_params) == n: + shuffled_sps = [sampling_params[i] for i in idxs] + elif len(sampling_params) == self.dp_size: + # per-shard SP; keep as-is and broadcast per-shard below + shuffled_sps = sampling_params + else: + shuffled_sps = [sampling_params[0]] * n + else: + shuffled_sps = sampling_params + + # Even partition of shuffled inputs + base = n // self.dp_size + rem = n % self.dp_size + slices = {} + start = 0 + for i in range(self.dp_size): + add = base + (1 if i < rem else 0) + end = start + add + if start < end: + slices[i] = (start, end) + start = end + + pending = {} + conn_to_idx = {} + collected = {} + for i, (s, e) in slices.items(): + if isinstance(shuffled_sps, list): + if len(shuffled_sps) == n: + sp_arg = shuffled_sps[s:e] + elif len(shuffled_sps) == self.dp_size: + sp_arg = shuffled_sps[i] + else: + sp_arg = shuffled_sps[0] + else: + sp_arg = shuffled_sps + conn = self.conns[i] + conn.send(("generate", shuffled_prompts[s:e], sp_arg, use_tqdm)) + pending[i] = True + conn_to_idx[conn] = i + # Collect + while pending: + ready = mp_wait([self.conns[i] for i in pending.keys()]) + for conn in ready: + try: + tag, payload = conn.recv() + except EOFError: + idx = conn_to_idx[conn] + p = self.ps[idx] + exitcode = p.exitcode + raise RuntimeError( + f"DP child #{idx} terminated unexpectedly during generate (exitcode={exitcode}). " + f"Enable envs: PYTHONFAULTHANDLER=1 CUDA_LAUNCH_BLOCKING=1 TORCH_SHOW_CPP_STACKTRACES=1 for more info." + ) + idx = conn_to_idx[conn] + if tag == "ok": + collected[idx] = payload + else: + raise RuntimeError(f"DP child #{idx} error: {payload}") + del pending[idx] + # Restore to original order + restored = [None] * n + for i, (s, e) in slices.items(): + outs = collected.get(i, []) + # outs are aligned with shuffled order s:e + for local_k, out in enumerate(outs): + global_pos = s + local_k + orig_idx = idxs[global_pos] + restored[orig_idx] = out + assert all(x is not None for x in restored), "Mismatch in outputs after DP collection" + return restored diff --git a/diffuserve/engine/llm_engine.py b/diffuserve/engine/llm_engine.py new file mode 100755 index 00000000..37daba57 --- /dev/null +++ b/diffuserve/engine/llm_engine.py @@ -0,0 +1,147 @@ +import atexit + +import torch.multiprocessing as mp + +from typing import List +from tqdm.auto import tqdm +from time import perf_counter +from dataclasses import fields +from transformers import AutoTokenizer + +from diffuserve.config import Config +from diffuserve.sampling_params import SamplingParams +from diffuserve.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM +from diffuserve.engine.scheduler import AutoScheduler, SchedulerBase +from diffuserve.engine.model_runner import AutoModelRunner + + +class LLMEngine: + def __init__(self, model, **kwargs): + config_fields = {field.name for field in fields(Config)} + config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} + self.config = config = Config(model, **config_kwargs) + self.engine_type = config.model_type + self.ps = [] + self.events = [] + ctx = mp.get_context("spawn") + for i in range(1, config.tensor_parallel_size): + event = ctx.Event() + process = ctx.Process(target=AutoModelRunner.from_config, args=(config, i, event)) + process.start() + self.ps.append(process) + self.events.append(event) + self.model_runner = AutoModelRunner.from_config(config, 0, self.events) + self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True, trust_remote_code=True) + config.eos = self.tokenizer.eos_token_id + self.scheduler: SchedulerBase = AutoScheduler.from_config(config) + self._exited = False + atexit.register(self.exit) + + def exit(self): + if getattr(self, "_exited", False): + return + self._exited = True + if hasattr(self, "model_runner") and self.model_runner is not None: + try: + self.model_runner.call("exit") + except Exception: + pass + try: + del self.model_runner + except Exception: + pass + for p in getattr(self, "ps", []): + try: + p.join() + except Exception: + pass + + def add_request(self, prompt: str | List[int], sampling_params: SamplingParams): + if isinstance(prompt, str): + prompt = self.tokenizer.encode(prompt) + + if self.engine_type == "causal_lm": + seq = SequenceForCausalLM(prompt, sampling_params) + elif self.engine_type == "diffusion_lm": + seq = SequenceForDiffusionLM(prompt, sampling_params, config=self.config) + else: + raise ValueError(f"Unsupported engine type: {self.engine_type}") + + seq.block_size = self.config.kvcache_block_size + self.scheduler.add(seq) + # Return seq_id so caller can build a stable mapping + return seq.seq_id + + def step(self): + seqs, is_prefill = self.scheduler.schedule() + 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] + if self.engine_type == "causal_lm": + num_tokens = sum(len(seq) for seq in seqs) if is_prefill else len(seqs) + # For streaming: provide per-seq deltas (newly appended token) on decode steps + if not is_prefill: + deltas = [(seq.seq_id, [seq.last_token], seq.is_finished) for seq in seqs] + else: + deltas = [] + else: + 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) + # Diffusion decoding modifies tokens in-place; we currently don't stream intermediate edits + deltas = [] + return outputs, num_tokens, is_prefill, n_diff_steps, deltas + + def is_finished(self): + return self.scheduler.is_finished() + + def generate( + self, + prompts: List[str] | List[List[int]], + sampling_params: SamplingParams | List[SamplingParams], + use_tqdm: bool = True, + ) -> List[str]: + if use_tqdm: + pbar = tqdm(total=len(prompts), desc="Generating", dynamic_ncols=True) + if not isinstance(sampling_params, list): + sampling_params = [sampling_params] * len(prompts) + # Map internal seq_id -> input index to keep output order stable + seqid_to_idx = {} + for idx, (prompt, sp) in enumerate(zip(prompts, sampling_params)): + sid = self.add_request(prompt, sp) + seqid_to_idx[sid] = idx + outputs = [None] * len(prompts) + prefill_throughput = decode_throughput = 0. + n_steps = 0 + n_diff_steps = [-1] * len(prompts) + while not self.is_finished(): + t = perf_counter() + n_steps += 1 + output, num_tokens, is_prefill, cur_n_diff_steps, _ = self.step() + if use_tqdm: + if is_prefill: + prefill_throughput = num_tokens / (perf_counter() - t) + else: + decode_throughput = num_tokens / (perf_counter() - t) + pbar.set_postfix({ + "Prefill": f"{int(prefill_throughput)}tok/s", + "Decode": f"{int(decode_throughput)}tok/s", + }) + if cur_n_diff_steps: + for seq_id, n_step in cur_n_diff_steps.items(): + if seq_id in seqid_to_idx and n_step >= 0: + n_diff_steps[seqid_to_idx[seq_id]] = n_step + for seq_id, token_ids in output: + if seq_id in seqid_to_idx: + outputs[seqid_to_idx[seq_id]] = token_ids + if use_tqdm: + pbar.update(1) + print(f"Finished in {n_steps} steps, prefill throughput: {prefill_throughput:.2f} tok/s, decode throughput: {decode_throughput:.2f} tok/s") + # Ensure all outputs are present + assert all(toks is not None for toks in outputs), "Some sequences did not produce outputs" + outputs = [{ + "text": self.tokenizer.decode(token_ids).split(self.tokenizer.eos_token)[0], + "token_ids": token_ids[:token_ids.index(self.config.eos)] if self.config.eos in token_ids else token_ids, + "n_diff_steps": n_diff_step, + } for token_ids, n_diff_step in zip(outputs, n_diff_steps)] + if use_tqdm: + pbar.close() + return outputs diff --git a/diffuserve/engine/model_runner.py b/diffuserve/engine/model_runner.py new file mode 100755 index 00000000..6cbe6c84 --- /dev/null +++ b/diffuserve/engine/model_runner.py @@ -0,0 +1,507 @@ +import time +import torch +import pickle + +import torch.distributed as dist + +from typing import List +from abc import ABC, abstractmethod +from multiprocessing.synchronize import Event +from multiprocessing.shared_memory import SharedMemory + +from diffuserve.config import Config +from diffuserve.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM, SequenceBase +from diffuserve.model.auto_model import AutoModelForDiffusionLM +from diffuserve.layer.sampler import AutoSampler +from diffuserve.utils.checker import CHECK_SLOT_MAPPING +from diffuserve.utils.context import ( + set_context_causal_lm, + get_context_causal_lm, + reset_context_causal_lm, + set_context_diffusion_lm, + get_context_diffusion_lm, + reset_context_diffusion_lm +) + + +class ModelRunnerBase(ABC): + """Base class for model runners supporting different model types.""" + def __init__(self, config: Config, rank: int, event: Event | List[Event]): + self.config = config + self.model_type = config.model_type + hf_config = config.hf_config + self.block_size = config.kvcache_block_size + self.enforce_eager = config.enforce_eager + self.world_size = config.tensor_parallel_size + self.rank = rank + self.event = event + + # Initialize model, sampler, and kv cache + init_method = f"tcp://{config.master_addr}:{config.master_port}" + dist.init_process_group("nccl", init_method, world_size=self.world_size, rank=rank) + device_id = (getattr(config, "device_start", 0) or 0) + rank + assert 0 <= device_id < torch.cuda.device_count(), f"Invalid device_id {device_id}." + torch.cuda.set_device(device_id) + default_dtype = torch.get_default_dtype() + default_dtype = (hf_config.torch_dtype if hasattr(hf_config, "torch_dtype") + and hf_config.torch_dtype else torch.bfloat16) + torch.set_default_dtype(default_dtype) + torch.set_default_device(f"cuda:{device_id}") + self.model = AutoModelForDiffusionLM.from_config(config) + self.sampler = AutoSampler.from_config(config) + self.warmup_model() + self.allocate_kv_cache() # NOCHANGE + if not self.enforce_eager: + self.capture_cudagraph() + + # Allocate shared memory for inter-process communication + # NOCHANGE + torch.set_default_device("cpu") + torch.set_default_dtype(default_dtype) + if self.world_size > 1: + if rank == 0: + try: + shm = SharedMemory(name=config.shm_name) + shm.close() + shm.unlink() + except FileNotFoundError: + pass + shm_size = 2**25 if self.model_type == "diffusion_lm" else 2**20 + self.shm = SharedMemory(name=config.shm_name, create=True, size=shm_size) + dist.barrier() + else: + dist.barrier() + self.shm = SharedMemory(name=config.shm_name) + self.loop() + + def exit(self): + if self.world_size > 1: + self.shm.close() + dist.barrier() + if self.rank == 0: + self.shm.unlink() + if not self.enforce_eager: + del self.graphs, self.graph_pool + torch.cuda.synchronize() + dist.destroy_process_group() + + def loop(self): + while True: + method_name, args = self.read_shm() + self.call(method_name, *args) + if method_name == "exit": + break + + def read_shm(self): + assert self.world_size > 1 and self.rank + self.event.wait() + n = int.from_bytes(self.shm.buf[0:4], "little") + method_name, *args = pickle.loads(self.shm.buf[4:n+4]) + self.event.clear() + return method_name, args + + def write_shm(self, method_name, *args): + assert self.world_size > 1 and not self.rank + data = pickle.dumps([method_name, *args]) + n = len(data) + + if n + 4 > len(self.shm.buf): + raise ValueError(f"Serialized data size ({n} bytes) exceeds shared memory buffer size ({len(self.shm.buf)} bytes). " + f"Consider increasing shared memory size or reducing batch size.") + + self.shm.buf[0:4] = n.to_bytes(4, "little") + self.shm.buf[4:n+4] = data + for event in self.event: + event.set() + + def call(self, method_name, *args): + if self.world_size > 1 and self.rank == 0: + self.write_shm(method_name, *args) + method = getattr(self, method_name, None) + return method(*args) + + @abstractmethod + def warmup_model(self): + """Model-specific warmup logic.""" + pass + + @abstractmethod + def allocate_kv_cache(self): + pass + + def prepare_block_tables(self, seqs: List[SequenceBase]): + max_len = max(len(seq.block_table) for seq in seqs) + block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs] + block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + return block_tables + + @abstractmethod + def prepare_prefill(self, seqs: List[SequenceBase]): + """Model-specific prefill preparation.""" + pass + + @abstractmethod + def prepare_decode(self, seqs: List[SequenceBase]): + """Model-specific decode preparation.""" + pass + + def prepare_sample(self, seqs: List[SequenceBase]): + temperatures = [] + for seq in seqs: + temperatures.append(seq.temperature) + temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True) + return temperatures + + @abstractmethod + @torch.inference_mode() + def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): + """Model-specific forward pass.""" + pass + + @abstractmethod + def run(self, seqs: List[SequenceBase], is_prefill: bool) -> List[int]: + """Main inference pipeline.""" + pass + + @abstractmethod + @torch.inference_mode() + def capture_cudagraph(self): + """Model-specific CUDA graph capture.""" + pass + + +class ModelRunnerForDiffusionLM(ModelRunnerBase): + """Model runner for Diffusion Language Models. TODO: Implement DLM-specific logic.""" + def __init__(self, config: Config, rank: int, event: Event | List[Event]): + super().__init__(config, rank, event) + self.diffusion_block_size = config.diffusion_block_size + self.mask_token_id = config.mask_token_id + self.decoding_strategy = config.decoding_strategy + + def warmup_model(self): + # return + print("Warming up model...") + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + max_num_batched_tokens, max_model_len = self.config.max_num_batched_tokens, self.config.max_model_len + num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs) + test_input_ids = [0] * max_model_len + seqs = [SequenceForDiffusionLM(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() + + 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 - # noqa: E731 + used - peak + current) // block_bytes + try: + num_kvcache_blocks = get_num_kvcache_blocks(config.gpu_memory_utilization) + assert num_kvcache_blocks > 0 + except: # noqa: E722 + gpu_memory_utilization = config.gpu_memory_utilization + while num_kvcache_blocks <= 200: + print(f"Warning: GPU memory utilization {gpu_memory_utilization} is too low to allocate kv cache. " + f"Automatically adding 0.05, which is {gpu_memory_utilization + 0.05:.2f} now.") + 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(f"Allocated {config.num_kvcache_blocks} blocks of size {self.block_size} for kv cache on rank {self.rank}.") + + if config.kv_cache_layout == "distinct": + # k_cache: [layer_id, block_id, head, head_dim // x, block_size(segmented seq_len), x] + # v_cache: [layer_id, block_id, head, head_dim, block_size(segmented seq_len)] + 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": + # [kv_separated, layer_id, block_id, block_size(segmented seq_len), head, head_dim] + 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(f"Unsupported kv_cache_layout: {config.kv_cache_layout}. " + f"Supported values are 'distinct' and 'unified'.") + + def prepare_prefill(self, seqs: List[SequenceForDiffusionLM]): + input_ids = [] + positions = [] + cu_seqlens_q = [0] + cu_seqlens_k = [0] + max_seqlen_q = 0 + max_seqlen_k = 0 + slot_mapping = [] + block_tables = None + context_lens = [] + seq_lens = [] + + for seq in seqs: + seq.next_diffusion_step(is_prefill=True) + + total_seqlen = len(seq) + # tokens and positions to run in this prefill step + input_ids.extend(seq[seq.cached_num_tokens:]) + positions.extend(list(range(seq.cached_num_tokens, total_seqlen))) + seq_lens.append(total_seqlen) + context_lens.append(0) + assert len(input_ids) == len(positions), ( + f"prepare_prefill(diffusion): len(input_ids) {len(input_ids)} != len(positions) {len(positions)}" + ) + + seqlen_q = total_seqlen - seq.cached_num_tokens + seqlen_k = total_seqlen + cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) + cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k) + + max_seqlen_q = max(seqlen_q, max_seqlen_q) + max_seqlen_k = max(seqlen_k, max_seqlen_k) + + if not seq.block_table: + continue + # build slot mapping for prefix cache prompt blocks + for i in range(0, seq.num_prompt_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 + else: + end = start + seq.last_block_prompt_num_tokens + slot_mapping.extend(list(range(start, end))) + else: + slot_mapping.extend([-1] * self.block_size) + # pad to a full diffusion block + slot_mapping.extend([-1] * seq.diffusion_block_size) + + # For diffusion prefill we always need block tables for prefix cache bookkeeping + block_tables = self.prepare_block_tables(seqs) + + input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) + positions = 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 = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + + # More checks to avoid downstream rotary errors + assert cu_seqlens_q[-1].item() == input_ids.numel(), ( + f"prepare_prefill(diffusion): cu_seqlens_q[-1]={cu_seqlens_q[-1].item()} != num_tokens={input_ids.numel()}" + ) + assert cu_seqlens_k[-1].item() == sum(seq_lens), ( + f"prepare_prefill(diffusion): cu_seqlens_k[-1]={cu_seqlens_k[-1].item()} != sum(seq_lens)={sum(seq_lens)}" + ) + + set_context_diffusion_lm( + True, + 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, + seqs=seqs, + kv_cache_layout=self.config.kv_cache_layout, + seq_lens=seq_lens, + seq_lens_ts=seq_lens_ts, + ) + return input_ids, positions + + def prepare_decode(self, seqs: List[SequenceForDiffusionLM]): + input_ids = [] + positions = [] + cu_seqlens_q = [0] + cu_seqlens_k = [0] + slot_mapping = [] + context_lens = [] + seq_lens = [] + seq_id_to_queue_id = {} + need_kv_cache_store = False + max_seqlen_q = 0 + max_seqlen_k = 0 + # if sum((sum(seq.active_blocks) + sum(seq.to_cache_blocks)) * seq.diffusion_block_size for seq in seqs) == 1536: + # pass + 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 + 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 + 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 + get_step = lambda diff_blk, start_idx: ( + diff_blk.remaining_length(start_idx) + if diff_blk.remaining_length(start_idx) + local_start_idx() <= seq.block_size + else 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(list(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: + # Covering all the after-active blocks + 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 + padding_slots = [-1] * (num_blocks_to_pad * seq.diffusion_block_size) + slot_mapping.extend(padding_slots) + break + assert len(input_ids) == len(positions), f"Input IDs length {len(input_ids)} does not match positions length {len(positions)}" + assert len(input_ids) == len(slot_mapping), f"Input IDs length {len(input_ids)} does not match slot mapping length {len(slot_mapping)}" + + # CHECK_SLOT_MAPPING(seqs, slot_mapping) + input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) + positions = 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 = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + block_tables = self.prepare_block_tables(seqs) + set_context_diffusion_lm(False, slot_mapping=slot_mapping, context_lens=context_lens, + cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, + 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, + kv_cache_layout=self.config.kv_cache_layout, need_kv_cache_store=need_kv_cache_store, + d2f_pp=True) + return input_ids, positions + + @torch.inference_mode() + 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)) + else: + bs = input_ids.size(0) + context = get_context_diffusion_lm() + graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] + graph_vars = self.graph_vars + for k, v in graph_vars.items(): + if k != "outputs": + v.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 + graph.replay() + return self.model.compute_logits(graph_vars["outputs"][:bs]) + + 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}") + s = 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() - s:.2f} seconds") + s = time.time() + logits = self.run_model(input_ids, positions, is_prefill) + print(f"Ran model in {time.time() - s:.2f} seconds") + s = time.time() + sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None + print(f"Sampled tokens in {time.time() - s:.2f} seconds") + reset_context_diffusion_lm() + 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 + logits = self.run_model(input_ids, positions, is_prefill) + sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None + reset_context_diffusion_lm() + 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.") + + +class AutoModelRunner: + @classmethod + def from_config(cls, config: Config, rank: int, event: Event | List[Event]): + """Factory method to create a model runner based on the model type.""" + return ModelRunnerForDiffusionLM(config, rank, event) \ No newline at end of file diff --git a/diffuserve/engine/scheduler.py b/diffuserve/engine/scheduler.py new file mode 100755 index 00000000..0018ec07 --- /dev/null +++ b/diffuserve/engine/scheduler.py @@ -0,0 +1,234 @@ +import torch + +from collections import deque +from abc import ABC, abstractmethod +from typing import Tuple, List, Deque + +from diffuserve.config import Config +from diffuserve.engine.sequence import ( + SequenceBase, SequenceStatus, + SequenceForDiffusionLM, SequenceForCausalLM +) +from diffuserve.layer.sampler import SampleOutputForDiffusionLM +from diffuserve.engine.block_manager import AutoBlockManager + + +class SchedulerBase(ABC): + def __init__(self, config: Config): + self.max_num_seqs = config.max_num_seqs + self.max_num_batched_tokens = config.max_num_batched_tokens + self.eos = config.eos + self.block_manager = AutoBlockManager.from_config(config) + self.waiting: Deque[SequenceBase] = deque() + self.running: Deque[SequenceBase] = deque() + + @abstractmethod + def is_finished(self) -> bool: + pass + + @abstractmethod + def add(self, seq: SequenceBase) -> None: + pass + + @abstractmethod + def schedule(self) -> Tuple[List[SequenceBase], bool]: + pass + + @abstractmethod + def preempt(self, seq: SequenceBase) -> None: + pass + + @abstractmethod + def postprocess(self, seqs: List[SequenceBase], token_ids: List[int]): + pass + + +class SchedulerForCausalLM(SchedulerBase): + def __init__(self, config: Config): + super().__init__(config) + + def is_finished(self) -> bool: + return not self.waiting and not self.running + + def add(self, seq: SequenceForCausalLM) -> None: + self.waiting.append(seq) + + def schedule(self) -> Tuple[List[SequenceForCausalLM], bool]: + # prefill + scheduled_seqs = [] + num_seqs = 0 + num_batched_tokens = 0 + while self.waiting and num_seqs < self.max_num_seqs: + seq = self.waiting[0] + if num_batched_tokens + len(seq) > self.max_num_batched_tokens \ + or not self.block_manager.can_allocate(seq): + break + num_seqs += 1 + self.block_manager.allocate(seq) + num_batched_tokens += len(seq) - seq.num_cached_tokens + seq.status = SequenceStatus.RUNNING + self.waiting.popleft() + self.running.append(seq) + scheduled_seqs.append(seq) + if scheduled_seqs: + return scheduled_seqs, True + + # decode + while self.running and num_seqs < self.max_num_seqs: + seq = self.running.popleft() + while not self.block_manager.can_append(seq): + if self.running: + self.preempt(self.running.pop()) + else: + self.preempt(seq) + break + else: + num_seqs += 1 + self.block_manager.may_append(seq) + scheduled_seqs.append(seq) + if not scheduled_seqs: + # Provide diagnostics to understand starvation/resource issues + diag = { + "phase": "decode", + "waiting": len(self.waiting), + "running": len(self.running), + "max_num_seqs": self.max_num_seqs, + "max_num_batched_tokens": self.max_num_batched_tokens, + } + # Probe a few candidates for can_append and lengths + candidates = list(self.running)[:3] + list(self.waiting)[:2] + infos = [] + for j, s in enumerate(candidates): + try: + cap = self.block_manager.can_append(s) + except Exception: + cap = "error" + infos.append( + f"[{j}] status={s.status.name}, len={len(s)}, new_tokens={getattr(s, 'num_completion_tokens', getattr(s, 'new_tokens', '?'))}, cached={getattr(s, 'num_cached_tokens', '?')}, can_append={cap}" + ) + raise RuntimeError(f"SchedulerForCausalLM: unable to schedule any sequence in decode; state={diag}; details={' | '.join(infos)}") + self.running.extendleft(reversed(scheduled_seqs)) + return scheduled_seqs, False + + def preempt(self, seq: SequenceForCausalLM) -> None: + seq.status = SequenceStatus.WAITING + self.block_manager.free(seq) + self.waiting.appendleft(seq) + + def postprocess(self, seqs: List[SequenceForCausalLM], token_ids: List[int]) -> None: + for seq, token_id in zip(seqs, token_ids): + seq.append_token(token_id) + if (not seq.ignore_eos and token_id == self.eos) \ + or seq.num_completion_tokens == seq.max_tokens: + seq.status = SequenceStatus.FINISHED + self.block_manager.free(seq) + self.running.remove(seq) + + +# TODO +class SchedulerForDiffusionLM(SchedulerBase): + def __init__(self, config: Config): + super().__init__(config) + self.diffusion_block_size = config.diffusion_block_size + + def is_finished(self) -> bool: + return not self.waiting and not self.running + + def add(self, seq: SequenceForDiffusionLM) -> None: + self.waiting.append(seq) + + def schedule(self): + # prefill + scheduled_seqs = [] + num_seqs = 0 + num_batched_tokens = 0 + while self.waiting and num_seqs < self.max_num_seqs: + seq = self.waiting[0] + if num_batched_tokens + len(seq) + seq.diffusion_block_size > self.max_num_batched_tokens or not self.block_manager.can_allocate(seq): + break + num_seqs += 1 + self.block_manager.allocate(seq) + num_batched_tokens += len(seq) + seq.diffusion_block_size - seq.num_cached_tokens + seq.status = SequenceStatus.RUNNING + self.waiting.popleft() + self.running.append(seq) + scheduled_seqs.append(seq) + if scheduled_seqs: + return scheduled_seqs, True + + # decode + while self.running and num_seqs < self.max_num_seqs: + seq = self.running.popleft() + while not self.block_manager.can_append(seq): + if self.running: + self.preempt(self.running.pop()) + else: + self.preempt(seq) + break + else: + num_seqs += 1 + self.block_manager.may_append(seq) + scheduled_seqs.append(seq) + if not scheduled_seqs: + diag = { + "phase": "decode", + "waiting": len(self.waiting), + "running": len(self.running), + "max_num_seqs": self.max_num_seqs, + "max_num_batched_tokens": self.max_num_batched_tokens, + "diffusion_block_size": getattr(self, 'diffusion_block_size', None), + } + candidates = list(self.running)[:3] + list(self.waiting)[:2] + infos = [] + for j, s in enumerate(candidates): + try: + cap = self.block_manager.can_append(s) + except Exception: + cap = "error" + infos.append( + f"[{j}] status={s.status.name}, len={len(s)}, diff_block={getattr(s, 'diffusion_block_size', '?')}, new_tokens={getattr(s, 'new_tokens', '?')}, cached={getattr(s, 'num_cached_tokens', '?')}, can_append={cap}" + ) + raise RuntimeError(f"SchedulerForDiffusionLM: unable to schedule any sequence in decode; state={diag}; details={' | '.join(infos)}") + self.running.extendleft(reversed(scheduled_seqs)) + return scheduled_seqs, False + + def preempt(self, seq: SequenceForDiffusionLM) -> None: + seq.status = SequenceStatus.WAITING + self.block_manager.free(seq) + self.waiting.appendleft(seq) + + def postprocess(self, seqs: List[SequenceForDiffusionLM], sample_output: SampleOutputForDiffusionLM) -> None: + n_diff_steps = {} + for seq in seqs: + seq.reset_new_tokens() + seq_id = str(seq.seq_id) + cur_true_local_ids_sub_map = sample_output.true_local_ids_map.get(seq_id, {}) + cur_accepted_ids_sub_map = sample_output.accepted_ids_map.get(seq_id, {}) + cur_sampled_tokens_sub_map = sample_output.sampled_tokens_map.get(seq_id, {}) + for block_id, accepted_ids in cur_accepted_ids_sub_map.items(): + if len(accepted_ids) > 0: + diffusion_block = seq.diffusion_blocks[int(block_id)] + sampled_tokens = cur_sampled_tokens_sub_map.get(block_id, []) + true_local_ids = cur_true_local_ids_sub_map.get(block_id, []) + + for true_local_id, accepted_id in zip(true_local_ids, accepted_ids): + diffusion_block.modify_token(true_local_id, sampled_tokens[accepted_id]) + if ((not seq.ignore_eos and sampled_tokens[accepted_id].item() == self.eos) + or seq.num_completion_tokens >= seq.max_tokens): + seq.meet_eos = True + if seq.meet_eos and seq.diffusion_blocks[-1].available_to_cache: + seq.status = SequenceStatus.FINISHED + self.block_manager.free(seq) + self.running.remove(seq) + n_diff_steps[seq.seq_id] = seq.n_steps + seq.post_process() + return n_diff_steps + +class AutoScheduler: + SCHEDULER_MAPPING = { + "causal_lm": SchedulerForCausalLM, + "diffusion_lm": SchedulerForDiffusionLM, + } + @classmethod + def from_config(cls, config: Config): + return cls.SCHEDULER_MAPPING[config.model_type](config) \ No newline at end of file diff --git a/diffuserve/engine/sequence.py b/diffuserve/engine/sequence.py new file mode 100755 index 00000000..cccf58e7 --- /dev/null +++ b/diffuserve/engine/sequence.py @@ -0,0 +1,552 @@ +import torch + +from copy import copy +from enum import Enum, auto +from itertools import count +from dataclasses import dataclass +from typing import List, Tuple, Any + +from diffuserve.config import Config +from diffuserve.sampling_params import SamplingParams + + +class SequenceStatus(Enum): + WAITING = auto() + RUNNING = auto() + FINISHED = auto() + + +class SequenceBase: + block_size = 256 + counter = count() + + def __init__(self, token_ids: List[int], sampling_params: SamplingParams = SamplingParams()): + self.seq_id = next(SequenceBase.counter) + 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 = [] + self.block_cache_missed = [] + self.temperature = sampling_params.temperature + self.max_tokens = sampling_params.max_tokens + self.ignore_eos = sampling_params.ignore_eos + + def __len__(self) -> int: + return self.num_tokens + + def __getitem__(self, key) -> int: + return self.token_ids[key] + + @property + def is_finished(self) -> bool: + return self.status == SequenceStatus.FINISHED + + @property + def prompt_token_ids(self) -> List[int]: + return self.token_ids[:self.num_prompt_tokens] + + @property + def num_blocks(self) -> int: + return (self.num_tokens + self.block_size - 1) // self.block_size + + @property + def last_block_num_tokens(self) -> int: + return self.num_tokens - (self.num_blocks - 1) * self.block_size + + def block(self, i) -> List[int]: + assert 0 <= i < self.num_blocks + return self.token_ids[i * self.block_size: (i + 1) * self.block_size] + + def append_token(self, token_id: int) -> None: + self.token_ids.append(token_id) + self.last_token = token_id + self.num_tokens += 1 + + +class SequenceForCausalLM(SequenceBase): + """Standard sequence implementation for Causal Language Models.""" + + def __init__(self, token_ids: List[int], sampling_params = SamplingParams()): + super().__init__(token_ids, sampling_params) + + def __repr__(self) -> str: + return (f"SequenceForCausalLM(block_size={self.block_size}, counter={self.counter}, " + f"seq_id={self.seq_id}, status={self.status.name}, num_tokens={self.num_tokens}, " + f"num_prompt_tokens={self.num_prompt_tokens}, num_cached_tokens={self.num_cached_tokens}, " + f"temperature={self.temperature}, max_tokens={self.max_tokens}, ignore_eos={self.ignore_eos})") + + def __getstate__(self) -> Tuple[int, int, int, List[int], int]: + return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table, + self.token_ids if self.num_completion_tokens == 0 else self.last_token) + + def __setstate__(self, state: Tuple[int, int, int, List[int], int]) -> None: + self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table = state[:-1] + if self.num_completion_tokens == 0: + self.token_ids = state[-1] + else: + self.last_token = state[-1] + + @property + def num_completion_tokens(self) -> int: + return self.num_tokens - self.num_prompt_tokens + + @property + def completion_token_ids(self) -> List[int]: + return self.token_ids[self.num_prompt_tokens:] + + @property + def num_cached_blocks(self) -> int: + return (self.num_cached_tokens + self.block_size - 1) // self.block_size + + +class DiffusionBlockStatus(Enum): + ACTIVE = auto() + TO_CACHE = auto() + IN_CACHE = auto() + + +@dataclass +class DiffusionBlock: + block_id: int = 0 + status: DiffusionBlockStatus = DiffusionBlockStatus.ACTIVE + + global_start_id: int = 0 + global_end_id: int | None = None + cursor: int = 0 + + mask_token_id: int = 151666 + size: int = 32 + is_prompt: bool = False + + accept_threshold: float = 0.95 # Threshold to accept a token in the diffusion block + add_new_block_threshold: float = 0.1 # Threshold to add a new block + complete_threshold: float = 0.9 # Can only be cached when the current diffusion block is completed + + seq: "SequenceForDiffusionLM" = None # Reference to the sequence this block belongs to + pre_block: "DiffusionBlock" = None # Create prefix linked list of diffusion blocks + suf_block: "DiffusionBlock" = None # Create suffix linked list of diffusion blocks + + def __post_init__(self): + self.global_end_id = self.global_start_id + self.size + + def __getitem__(self, key: int) -> int: + return self.seq[self.global_start_id + key] + + def __len__(self) -> int: + return self.size + + @property + def current_complete_ratio(self) -> float: + return ( + sum([token_id != self.mask_token_id for token_id in self.token_ids]) / self.size + ) if self.size > 0 else 0.0 + + @property + def available_to_cache(self) -> bool: + return self.current_complete_ratio == 1.0 + + @property + def is_active(self) -> bool: + return self.status == DiffusionBlockStatus.ACTIVE + + @property + def is_in_cache(self) -> bool: + return self.status == DiffusionBlockStatus.IN_CACHE + + @property + def is_to_cache(self) -> bool: + return self.status == DiffusionBlockStatus.TO_CACHE + + @property + def pre_block_complete(self) -> bool: + return self.pre_block.current_complete_ratio >= self.complete_threshold if self.pre_block is not None else True + + @property + def add_new_block(self) -> bool: + return self.current_complete_ratio >= self.add_new_block_threshold + + @property + def token_ids(self) -> torch.Tensor: + if self.seq is not None: + return self.seq.token_ids[self.global_start_id:self.global_end_id] + else: + raise RuntimeError("Sequence is not set for the diffusion block.") + + @property + def local_mask_tokens(self) -> List[bool]: + return [token_id == self.seq.mask_token_id for token_id in self.token_ids] + + @property + def local_mask_token_ids(self) -> List[int]: + return [idx for idx, mask_token in enumerate(self.local_mask_tokens) if mask_token] + + @property + def global_mask_token_ids(self) -> List[int]: + 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) + return [mask_token_id + offset for mask_token_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 = DiffusionBlockStatus.TO_CACHE + + def in_cache(self) -> None: + if self.is_to_cache: + self.status = DiffusionBlockStatus.IN_CACHE + + def modify_token(self, local_token_id: int, modified_to: int) -> None: + 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() + self.seq.new_tokens += 1 + + +class SequenceForDiffusionLM(SequenceBase): + """Sequence implementation for Diffusion Language Models.""" + + def __init__(self, token_ids: List[int], + sampling_params = SamplingParams(), + config: Config = None): + super().__init__(token_ids, sampling_params) + 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 = None + self.meet_eos = False + self.diffusion_blocks: List[DiffusionBlock] = [] + self.n_steps = 0 + + 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": getattr(self, "input_token_ids", []), + "input_num_tokens": getattr(self, "input_num_tokens", 0), + "input_num_prompt_tokens": getattr(self, "input_num_prompt_tokens", 0), + "new_tokens": getattr(self, "new_tokens", 0), + "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", None) + self.n_steps = state.get("n_steps", 0) + # Align tensor devices when sequence is reconstructed on a different rank + 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 = DiffusionBlock( + 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 + + def __repr__(self) -> str: + return (f"SequenceForDiffusionLM(block_size={self.block_size}, counter={self.counter}, " + f"seq_id={self.seq_id}, status={self.status.name}, num_tokens={self.num_tokens}, " + f"num_prompt_tokens={self.num_prompt_tokens}, num_cached_tokens={self.num_cached_tokens}, " + f"temperature={self.temperature}, max_tokens={self.max_tokens}, ignore_eos={self.ignore_eos}, " + f"diffusion_block_size={self.diffusion_block_size}, " + f"block_mask={(self.block_mask.shape if self.block_mask is not None else None)}, " + f"input_token_ids={getattr(self, 'input_token_ids', None)}, input_num_tokens={getattr(self, 'input_num_tokens', None)})") + + @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:] + + @property + def active_blocks(self) -> List[bool]: + return [block.is_active for block in self.diffusion_blocks] + + @property + def to_cache_blocks(self) -> List[bool]: + return [block.is_to_cache 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] + + @property + def num_prompt_blocks(self) -> int: + return (self.input_num_prompt_tokens + self.block_size - 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 + + @property + def updated_or_updating_kv_cache_block_ids(self) -> List[int]: + return [idx for idx, caching in enumerate(self.caching_blocks) if caching] + + @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] + + @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 cached_num_tokens - 1 + + @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) + + @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)]) # build up token-wise mapping + return mapping + + def token_to_diffusion_block_id(self, token_id: int) -> int: + if token_id < self.input_num_tokens: + return 0 + else: + return (token_id - self.input_num_tokens) // self.diffusion_block_size + 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 = [] + positions = [] + 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([token_id + temp_context_len for token_id in range(block.size)]) + temp_context_len += block.size + + return input_tokens, positions, context_len + + def reset_new_tokens(self) -> None: + self.new_tokens = 0 + + def post_process(self) -> None: + for diff_blk in self.diffusion_blocks: + diff_blk.cursor = 0 + if diff_blk.is_in_cache: + continue + + if diff_blk.is_to_cache: + diff_blk.in_cache() + elif diff_blk.is_active: + if diff_blk.available_to_cache: + diff_blk.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.kv_cache_layout == "distinct": + return self.block_mask[..., self.cached_num_tokens:, self.cached_num_tokens:] + else: + 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_wise_causal_mask = torch.zeros(mask_shape, dtype=torch.bool, device=torch.cuda.current_device()) + block_wise_causal_mask[..., :self.input_num_tokens, :self.input_num_tokens] = True + num_diffusion_blocks = (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 + start_w = 0 + end_w = end_h + block_wise_causal_mask[..., start_h:end_h, start_w:end_w] = True + self.block_mask = block_wise_causal_mask.clone() + else: + return + assert self.block_mask is not None, "block_mask must exist before incremental update" + 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.zeros(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: + # Take a snapshot of the original input state + 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( + DiffusionBlock( + block_id=len(self.diffusion_blocks), + status=DiffusionBlockStatus.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 self.diffusion_blocks[-1].add_new_block and not self.meet_eos: + added_num_tokens = ( + self.diffusion_block_size + if self.num_tokens + self.diffusion_block_size <= self.max_model_len + else self.max_model_len - self.num_tokens + ) + + diffusion_seq = [self.mask_token_id] * added_num_tokens + current_diffusion_block = DiffusionBlock( + block_id=len(self.diffusion_blocks), + status=DiffusionBlockStatus.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] if self.diffusion_blocks else None + ) + + self.diffusion_blocks[-1].suf_block = current_diffusion_block + self.token_ids += diffusion_seq + self.num_tokens += added_num_tokens + self.diffusion_blocks.append(current_diffusion_block) + + self.update_block_mask(is_prefill=is_prefill) \ No newline at end of file diff --git a/diffuserve/layer/activation.py b/diffuserve/layer/activation.py new file mode 100755 index 00000000..49de8dea --- /dev/null +++ b/diffuserve/layer/activation.py @@ -0,0 +1,14 @@ +import torch + +import torch.nn as nn +import torch.nn.functional as F + + +class SiluAndMul(nn.Module): + def __init__(self): + super().__init__() + + @torch.compile + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, y = x.chunk(2, -1) + return F.silu(x) * y diff --git a/diffuserve/layer/attention/attention_v5.py b/diffuserve/layer/attention/attention_v5.py new file mode 100644 index 00000000..60ea5017 --- /dev/null +++ b/diffuserve/layer/attention/attention_v5.py @@ -0,0 +1,148 @@ +import os +import torch + +import torch.nn as nn + +from typing import List +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 diffuserve.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 diffuserve.legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm + + +class Attention(nn.Module): + def __init__( + self, + num_heads, + head_dim, + scale, + num_kv_heads, + model_type='causal_lm' + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + self.scale = scale + self.num_kv_heads = num_kv_heads + self.k_cache = self.v_cache = torch.tensor([]) + self.causal = model_type == 'causal_lm' + self.model_type = model_type + 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] + + @lru_cache(maxsize=32) + def causal_lm_block_mask(self, cum_seq_lens: torch.Tensor, B: int, H: int, Q_LEN: int, KV_LEN: int, device: str): + cache_key = (B, H, Q_LEN, KV_LEN, device) + document_ids = torch.zeros((cum_seq_lens[-1],), dtype=torch.int32, device=device) + start_idx = 0 + for doc_idx, seq_len in enumerate(cum_seq_lens[1:]): + end_idx = seq_len + document_ids[start_idx:end_idx] = doc_idx + start_idx = end_idx + + def _mask_mod(batch, head, token_q, token_kv): + causal_mask = token_q >= token_kv + document_mask = document_ids[token_q] == document_ids[token_kv] + return causal_mask & document_mask + + 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] + + 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) + + context: ContextForDiffusionLM = get_context_causal_lm() if self.model_type == 'causal_lm' else get_context_diffusion_lm() + k_cache, v_cache = self.k_cache, self.v_cache + is_unified_layout = context.kv_cache_layout == "unified" + + # Fast Store KV cache + if k_cache.numel() and v_cache.numel(): + if not (self.model_type == 'diffusion_lm' and not context.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, context.slot_mapping, self.model_type, context) + # 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 + if context.is_prefill: + # Block PK + if context.block_tables is not None and self.model_type == 'causal_lm': + k, v = k_cache, v_cache + elif context.block_tables is not None and self.model_type == 'diffusion_lm': + # 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_fn = self.causal_lm_block_mask if self.model_type == 'causal_lm' else self.dllm_block_mask + input_obj = context.cu_seqlens_q if self.model_type == 'causal_lm' else context.block_mask + block_mask = block_mask_fn(input_obj, B, H, S, S, str(q.device)) + o = self.attention(q_t, k_t, v_t, block_mask=block_mask) + else: + config = context.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, context, k, v) + o = flash_attn_varlen_func(q, k_comb, v_comb, + context.cu_seqlens_q, context.cu_seqlens_k, + context.max_seqlen_q, context.max_seqlen_k, + softmax_scale=self.scale, block_table=None) + 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, + context.block_tables, context.cu_seqlens_q, context.total_lens, + max(context.total_lens), max(context.seq_lens), 1.0, 1.0, + diffusion_block_size, context.block_mask + ) + CHECK_ATTENTION(o, q, k, v, k_cache, v_cache, context) + + # Final reshape + if not context.is_prefill: + o = o.view(-1, self.num_heads * self.head_dim).contiguous() + elif context.is_prefill: + o = rearrange(o, '1 h s d -> s (h d)').contiguous() + + return o \ No newline at end of file diff --git a/diffuserve/layer/attention/ops/__init__.py b/diffuserve/layer/attention/ops/__init__.py new file mode 100755 index 00000000..8e202106 --- /dev/null +++ b/diffuserve/layer/attention/ops/__init__.py @@ -0,0 +1,7 @@ +from diffuserve.legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd as causal_lm_flash_decoding +from diffuserve.legacy.layers.attention.ops.triton_decode_attn_dlm import diffusion_lm_flash_decoding, CHECK_ATTENTION +from diffuserve.legacy.layers.attention.ops.chunked_prefill_decoding_unified_kernel import chunked_prefill_paged_decode as diffusion_lm_parallel_flash_decoding +from diffuserve.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/diffuserve/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py b/diffuserve/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py new file mode 100755 index 00000000..8cc41a72 --- /dev/null +++ b/diffuserve/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py @@ -0,0 +1,375 @@ +# 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 diffuserve.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/diffuserve/layer/attention/ops/kv_cache_kernels.py b/diffuserve/layer/attention/ops/kv_cache_kernels.py new file mode 100755 index 00000000..6f7d70c4 --- /dev/null +++ b/diffuserve/layer/attention/ops/kv_cache_kernels.py @@ -0,0 +1,395 @@ +import torch +import triton + +import triton.language as tl + +from typing import Tuple +from einops import rearrange + +from diffuserve.legacy.utils.context import ContextForDiffusionLM +from diffuserve.legacy.engine.sequence import SequenceForDiffusionLM + +@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) + + +@triton.jit +def store_kvcache_kernel_diffusion_lm( + key_ptr, + key_stride, + value_ptr, + value_stride, + k_cache_ptr, + v_cache_ptr, + slot_mapping_ptr, + D: tl.constexpr +): + token_idx = tl.program_id(0) + slot = tl.load(slot_mapping_ptr + token_idx) + if slot < 0: + return + key_offsets = token_idx * key_stride + tl.arange(0, D) + value_offsets = token_idx * value_stride + tl.arange(0, D) + key = tl.load(key_ptr + key_offsets) + value = tl.load(value_ptr + value_offsets) + cache_offsets = slot * D + tl.arange(0, D) + tl.store(k_cache_ptr + cache_offsets, key) + tl.store(v_cache_ptr + cache_offsets, value) + + +@triton.jit +def store_kvcache_kernel_diffusion_lm_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, + v_cache_stride_nblks, v_cache_stride_h, v_cache_stride_d, v_cache_stride_blk_sz, + nheads, hdim, blk_sz, + x: tl.constexpr, D: tl.constexpr +): + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: D2F + + # Translated from vLLM's CUDA kernel + # Referencing https://github.com/vllm-project/vllm/blob/main/csrc/cache_kernels.cu#L212 + # and https://github.com/vllm-project/vllm/blob/main/csrc/cache_kernels.cu#L415 + + # Organization: SJTU DENG Lab + # Author: Drew Jin (JIN. Yijie, @drewjin) + # Date: 2025-08-03 + # Email: drewjin0827@gmail.com + # All rights reserved. + + token_idx = tl.program_id(0) + slot_idx = tl.load(slot_mapping_ptr + token_idx) + if slot_idx < 0: + return + + blk_idx = slot_idx // blk_sz + off_blk = slot_idx % blk_sz + + offs_d = tl.arange(0, D) + offs_k = token_idx * k_stride + offs_d + offs_v = token_idx * v_stride + offs_d + k = tl.load(k_ptr + offs_k) + v = tl.load(v_ptr + offs_v) + + h_ids = offs_d // hdim + h_offs = offs_d % hdim + x_ids = h_offs // x + x_offs = h_offs % x + + k_cache_offs = (blk_idx * k_cache_stride_nblks + h_ids * k_cache_stride_h + + x_ids * k_cache_stride_dx + off_blk * k_cache_stride_blk_sz + + x_offs * k_cache_stride_x) + v_cache_offs = (blk_idx * v_cache_stride_nblks + h_ids * v_cache_stride_h + + h_offs * v_cache_stride_d + off_blk * v_cache_stride_blk_sz) + + tl.store(k_cache_ptr + k_cache_offs, k) + tl.store(v_cache_ptr + v_cache_offs, v) + + +def store_kvcache_distinct_layout(key: torch.Tensor, value: torch.Tensor, + k_cache: torch.Tensor, v_cache: torch.Tensor, + slot_mapping: torch.Tensor, model_type: str = 'causal_lm', + context: ContextForDiffusionLM = None) -> None: + + if model_type == 'causal_lm': + # k_cache: [num_blks, blk_sz, h, hdim] + # v_cache: [num_blks, blk_sz, h, hdim] + 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() + store_kvcache_kernel_causal_lm[(N,)]( + key, key.stride(0), + value, value.stride(0), + k_cache, v_cache, slot_mapping, D + ) + else: + # 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 + HDim = HDim_x * x + N = key.shape[0] + assert HDim == key.shape[-1] and NHeads == key.shape[1] + assert N == slot_mapping.numel() + + GRID = (N, ) + store_kvcache_kernel_diffusion_lm_distinct[GRID]( + key, value, + k_cache, v_cache, + slot_mapping, + key.stride(0), value.stride(0), + *k_cache.stride(), *v_cache.stride(), + NHeads, HDim, Blk_sz, + x, HDim * NHeads + ) + + +def store_kvcache_unified_layout(key: torch.Tensor, value: torch.Tensor, + k_cache: torch.Tensor, v_cache: torch.Tensor, + slot_mapping: torch.Tensor, model_type: str = 'causal_lm', + context: ContextForDiffusionLM = None) -> 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()}" + + if model_type == 'causal_lm': + store_kvcache_kernel_causal_lm[(N,)]( + key, key.stride(0), + value, value.stride(0), + k_cache, v_cache, slot_mapping, D + ) + elif model_type == 'diffusion_lm': + store_kvcache_kernel_diffusion_lm[(N,)]( + key, key.stride(0), + value, value.stride(0), + k_cache, v_cache, slot_mapping, D + ) + + +@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): + # BUG FIX + # SPDX-License-Identifier: Apache-2.0 + # SPDX-FileCopyrightText: D2F + + # Organization: SJTU DENG Lab + # Author: Drew Jin (JIN. Yijie, @drewjin) + # Date: 2025-08-01 + # Email: drewjin0827@gmail.com + # All rights reserved. + + seq_idx = tl.program_id(0) + local_blk_idx = tl.program_id(1) + kv_head_idx = tl.program_id(2) + + off_local_blk = seq_idx * block_table_stride_nseqs + local_blk_idx * block_table_stride_maxblks + global_blk_idx = tl.load(block_table_ptr + off_local_blk) + + if global_blk_idx != -1: + off_ctxlen = seq_idx * ctxlens_stride + global_ctxlen = tl.load(ctxlens_ptr + off_ctxlen) + cur_window_sz = (local_blk_idx + 1) * PAGE_SIZE + prev_window_sz = local_blk_idx * PAGE_SIZE + local_ctxlen = tl.where(global_ctxlen > cur_window_sz, PAGE_SIZE, global_ctxlen % PAGE_SIZE) + if global_ctxlen > prev_window_sz: + # Load KV cache + offs_kv_cache_seq = tl.arange(0, PAGE_SIZE) + offs_kv_cache_hdim = tl.arange(0, HEAD_DIM) + offs_kv_cache = ( # [NBlks, BlkSz, Hkv, Hdim] + global_blk_idx[None, :] * kv_cache_stride_nblks + # NBlks: BlkId + offs_kv_cache_seq[None, :] * kv_cache_stride_blk + # BlkSz: TokenIds + kv_head_idx * kv_cache_stride_h + # Hkv: HeadId + offs_kv_cache_hdim[:, None] * kv_cache_stride_d # Hdim: HeadDim Elems + ) + kv_cache_mask = offs_kv_cache_seq[None, :] < local_ctxlen + k_cache = tl.load(k_cache_ptr + offs_kv_cache, mask=kv_cache_mask, other=0.0) + v_cache = tl.load(v_cache_ptr + offs_kv_cache, mask=kv_cache_mask, other=0.0) + + # Store KV cache into output KV tensors + off_cu_seqlens_k = seq_idx * cu_seqlens_k_stride + kv_out_start_idx = tl.load(cu_seqlens_k_ptr + off_cu_seqlens_k) + cur_kv_cache_to_out_start_idx = kv_out_start_idx + prev_window_sz + offs_kv_cache_to_out = ( # [Seq, Hkv, Hdim] + (cur_kv_cache_to_out_start_idx + offs_kv_cache_seq[None, :]) * kv_out_stride_s + # Seq: TokenIds over Offset + kv_head_idx * kv_out_stride_h + # Hkv: HeadId + offs_kv_cache_hdim[:, None] * kv_out_stride_d # Hdim: HeadDim Elems + ) + tl.store(k_out_ptr + offs_kv_cache_to_out, k_cache, mask=kv_cache_mask) + tl.store(v_out_ptr + offs_kv_cache_to_out, v_cache, mask=kv_cache_mask) + + # Load and store active KV only once when first meet + if local_blk_idx == LAST_BLK_ID: + # Load KV new + off_cu_seqlens_q = seq_idx * cu_seqlens_q_stride + off_seqlens = seq_idx * seqlens_stride + kv_new_start_idx = tl.load(cu_seqlens_q_ptr + off_cu_seqlens_q) + active_seqlen = tl.load(seqlens_ptr + off_seqlens) + offs_kv_new_seq = tl.arange(0, DIFFUSION_BLOCK_SIZE) + offs_kv_new_hdim = tl.arange(0, HEAD_DIM) + + for diff_blk_idx in tl.range(active_seqlen // DIFFUSION_BLOCK_SIZE, loop_unroll_factor=KV_LOAD_UNROLL_FACTOR): + off_diff_blk = diff_blk_idx * DIFFUSION_BLOCK_SIZE + cur_kv_new_start_idx = kv_new_start_idx + off_diff_blk + offs_cur_kv_new_seq = ( # [Seq, Hkv, Hdim] + (cur_kv_new_start_idx + offs_kv_new_seq[None, :]) * kv_new_stride_s + # Seq: TokenIds over Offset + kv_head_idx * kv_new_stride_h + # Hkv: HeadId + offs_kv_new_hdim[:, None] * kv_new_stride_d # Hdim: HeadDim Elems + ) + k_new = tl.load(k_new_ptr + offs_cur_kv_new_seq) + v_new = tl.load(v_new_ptr + offs_cur_kv_new_seq) + + # Store KV new into output KV tensors + off_ctxlen = seq_idx * ctxlens_stride + off_cu_seqlens_k = seq_idx * cu_seqlens_k_stride + global_ctxlen = tl.load(ctxlens_ptr + off_ctxlen) + kv_out_start_idx = tl.load(cu_seqlens_k_ptr + off_cu_seqlens_k) + cur_kv_new_to_out_start_idx = global_ctxlen + kv_out_start_idx + off_diff_blk + offs_cur_kv_new_to_out = ( # [Seq, Hkv, Hdim] + (cur_kv_new_to_out_start_idx + offs_kv_new_seq[None, :]) * kv_out_stride_s + # Seq: TokenIds over Offset + kv_head_idx * kv_out_stride_h + # Hkv: HeadId + offs_kv_new_hdim[:, None] * kv_out_stride_d # Hdim: HeadDim Elems + ) + tl.store(k_out_ptr + offs_cur_kv_new_to_out, k_new) + tl.store(v_out_ptr + offs_cur_kv_new_to_out, v_new) + + +def load_kvcache(k_cache: torch.Tensor, v_cache: torch.Tensor, + context: ContextForDiffusionLM, + 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 + + ctxlens = context.context_lens + seqlens = context.seq_lens_ts + assert sum(seqlens) == k_new.shape[0] + DIFFUSION_BLOCK_SIZE = context.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 + 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 + + kv_output_shape = (sum(total_lens).item(), H_KV, HEAD_DIM) + k_output = torch.empty(kv_output_shape, device=k_cache.device, dtype=k_cache.dtype) + v_output = torch.empty_like(k_output) + + GRID = (NUM_SEQS, MAX_SEQ_BLOCKS, H_KV) + load_kvcache_kernel_kv[GRID]( + k_cache, v_cache, + k_new, v_new, + context.block_tables, + k_output, v_output, + seqlens, ctxlens, + cu_seqlens_q, cu_seqlens_k, + *k_cache.stride(), + *k_new.stride(), + *context.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, + 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: ContextForDiffusionLM) -> 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: ContextForDiffusionLM) -> 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 diff --git a/diffuserve/layer/attention/ops/prefix_prefill.py b/diffuserve/layer/attention/ops/prefix_prefill.py new file mode 100755 index 00000000..03cf31a8 --- /dev/null +++ b/diffuserve/layer/attention/ops/prefix_prefill.py @@ -0,0 +1,1090 @@ +# 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/diffuserve/layer/attention/ops/tilus_decode_attn_dlm.py b/diffuserve/layer/attention/ops/tilus_decode_attn_dlm.py new file mode 100755 index 00000000..fc7bc03b --- /dev/null +++ b/diffuserve/layer/attention/ops/tilus_decode_attn_dlm.py @@ -0,0 +1,161 @@ +# 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/diffuserve/layer/attention/ops/triton_decode_attn_clm.py b/diffuserve/layer/attention/ops/triton_decode_attn_clm.py new file mode 100755 index 00000000..71be2616 --- /dev/null +++ b/diffuserve/layer/attention/ops/triton_decode_attn_clm.py @@ -0,0 +1,681 @@ +# 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/diffuserve/layer/attention/ops/triton_decode_attn_dlm.py b/diffuserve/layer/attention/ops/triton_decode_attn_dlm.py new file mode 100755 index 00000000..8db75c0e --- /dev/null +++ b/diffuserve/layer/attention/ops/triton_decode_attn_dlm.py @@ -0,0 +1,120 @@ +# 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 diffuserve.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 diffuserve.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/diffuserve/layer/attention/ops/triton_flash_attention.py b/diffuserve/layer/attention/ops/triton_flash_attention.py new file mode 100755 index 00000000..37dd5356 --- /dev/null +++ b/diffuserve/layer/attention/ops/triton_flash_attention.py @@ -0,0 +1,1022 @@ +# 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/diffuserve/layer/embed_head.py b/diffuserve/layer/embed_head.py new file mode 100755 index 00000000..3b85e447 --- /dev/null +++ b/diffuserve/layer/embed_head.py @@ -0,0 +1,73 @@ +import torch + +import torch.nn as nn +import torch.nn.functional as F +import torch.distributed as dist + +from diffuserve.utils.context import get_context_causal_lm, get_context_diffusion_lm + + +class VocabParallelEmbedding(nn.Module): + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + ): + super().__init__() + self.tp_rank = dist.get_rank() + self.tp_size = dist.get_world_size() + assert num_embeddings % self.tp_size == 0 + self.num_embeddings = num_embeddings + self.num_embeddings_per_partition = self.num_embeddings // self.tp_size + self.vocab_start_idx = self.num_embeddings_per_partition * self.tp_rank + self.vocab_end_idx = self.vocab_start_idx + self.num_embeddings_per_partition + self.weight = nn.Parameter(torch.empty(self.num_embeddings_per_partition, embedding_dim)) + self.weight.weight_loader = self.weight_loader + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + param_data = param.data + shard_size = param_data.size(0) + start_idx = self.tp_rank * shard_size + loaded_weight = loaded_weight.narrow(0, start_idx, shard_size) + assert param_data.size() == loaded_weight.size() + param_data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor): + if self.tp_size > 1: + mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx) + x = mask * (x - self.vocab_start_idx) + y = F.embedding(x, self.weight) + if self.tp_size > 1: + y = mask.unsqueeze(1) * y + dist.all_reduce(y) + return y + + +class ParallelLMHead(VocabParallelEmbedding): + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + bias: bool = False, + model_type: str = 'causal_lm', + ): + super().__init__(num_embeddings, embedding_dim) + if bias: + self.bias = nn.Parameter(torch.empty(self.num_embeddings_per_partition)) + self.bias.weight_loader = self.weight_loader + else: + self.register_parameter("bias", None) + self.model_type = model_type + + def forward(self, x: torch.Tensor): + context = get_context_causal_lm() if self.model_type == 'causal_lm' else get_context_diffusion_lm() + if context.is_prefill and self.model_type == 'causal_lm': + last_indices = context.cu_seqlens_q[1:] - 1 + x = x[last_indices].contiguous() + logits = F.linear(x, self.weight, self.bias) + if self.tp_size > 1: + all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] if self.tp_rank == 0 else None + dist.gather(logits, all_logits, 0) + logits = torch.cat(all_logits, -1) if self.tp_rank == 0 else None + return logits \ No newline at end of file diff --git a/diffuserve/layer/layernorm.py b/diffuserve/layer/layernorm.py new file mode 100755 index 00000000..88c2cf02 --- /dev/null +++ b/diffuserve/layer/layernorm.py @@ -0,0 +1,51 @@ +import torch +import torch.nn as nn + + +class RMSNorm(nn.Module): + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + @torch.compile + def rms_forward( + self, + x: torch.Tensor, + ) -> torch.Tensor: + orig_dtype = x.dtype + x = x.to(torch.float32) + var = x.pow(2).mean(dim=-1, keepdim=True) + x.mul_(torch.rsqrt(var + self.eps)) + x = x.to(orig_dtype).mul_(self.weight) + return x + + @torch.compile + def add_rms_forward( + self, + x: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + orig_dtype = x.dtype + x = x.to(torch.float32).add_(residual.to(torch.float32)) + residual = x.to(orig_dtype) + var = x.pow(2).mean(dim=-1, keepdim=True) + x.mul_(torch.rsqrt(var + self.eps)) + x = x.to(orig_dtype).mul_(self.weight) + return x, residual + + def forward( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is None: + return self.rms_forward(x) + else: + return self.add_rms_forward(x, residual) diff --git a/diffuserve/layer/linear.py b/diffuserve/layer/linear.py new file mode 100755 index 00000000..cf14eb9f --- /dev/null +++ b/diffuserve/layer/linear.py @@ -0,0 +1,244 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributed as dist + + +def divide(numerator, denominator): + assert numerator % denominator == 0 + return numerator // denominator + + +class LoRAMixin: + """Mixin class to add LoRA support to existing linear layers.""" + def __init_lora__(self, r: int = 0, lora_alpha: float = 1.0, lora_dropout: float = 0.0): + if r > 0: + self.r = r + self.lora_alpha = lora_alpha + self.scaling = lora_alpha / r + + # Initialize LoRA parameters + if hasattr(self, 'output_size_per_partition'): + out_features = self.output_size_per_partition + else: + out_features = self.output_size + + if hasattr(self, 'input_size_per_partition'): + in_features = self.input_size_per_partition + else: + in_features = self.input_size + + self.lora_A = nn.Parameter(torch.zeros(r, in_features)) + self.lora_B = nn.Parameter(torch.zeros(out_features, r)) + self.lora_dropout = nn.Dropout(lora_dropout) if lora_dropout > 0 else nn.Identity() + self.merged = False + + # Initialize weights + nn.init.kaiming_uniform_(self.lora_A, a=5**0.5) + nn.init.zeros_(self.lora_B) + else: + self.r = 0 + self.merged = True + + def merge_lora(self): + """Merge LoRA weights into base weight.""" + if hasattr(self, 'r') and self.r > 0 and not self.merged: + self.weight.data += self.scaling * torch.mm(self.lora_B, self.lora_A) + self.merged = True + + def lora_forward(self, x: torch.Tensor, base_output: torch.Tensor) -> torch.Tensor: + """Apply LoRA forward pass.""" + if not hasattr(self, 'r') or self.r == 0 or self.merged: + return base_output + + lora_out = F.linear(self.lora_dropout(x), self.lora_A.T) + lora_out = F.linear(lora_out, self.lora_B.T) + return base_output + lora_out * self.scaling + + +class LinearBase(nn.Module): + + def __init__( + self, + input_size: int, + output_size: int, + tp_dim: int | None = None, + ): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.tp_dim = tp_dim + self.tp_rank = dist.get_rank() + self.tp_size = dist.get_world_size() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + +class ReplicatedLinear(LinearBase, LoRAMixin): + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = False, + r: int = 0, + lora_alpha: float = 1.0, + lora_dropout: float = 0.0, + ): + LinearBase.__init__(self, input_size, output_size) + self.weight = nn.Parameter(torch.empty(self.output_size, self.input_size)) + self.weight.weight_loader = self.weight_loader + if bias: + self.bias = nn.Parameter(torch.empty(self.output_size)) + self.bias.weight_loader = self.weight_loader + else: + self.register_parameter("bias", None) + + self.__init_lora__(r, lora_alpha, lora_dropout) + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + param.data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + base_out = F.linear(x, self.weight, self.bias) + return self.lora_forward(x, base_out) + + +class ColumnParallelLinear(LinearBase, LoRAMixin): + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = False, + r: int = 0, + lora_alpha: float = 1.0, + lora_dropout: float = 0.0, + ): + LinearBase.__init__(self, input_size, output_size, 0) + self.input_size_per_partition = input_size + self.output_size_per_partition = divide(output_size, self.tp_size) + + self.weight = nn.Parameter(torch.empty(self.output_size_per_partition, self.input_size)) + self.weight.weight_loader = self.weight_loader + if bias: + self.bias = nn.Parameter(torch.empty(self.output_size_per_partition)) + self.bias.weight_loader = self.weight_loader + else: + self.register_parameter("bias", None) + + self.__init_lora__(r, lora_alpha, lora_dropout) + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + param_data = param.data + shard_size = param_data.size(self.tp_dim) + start_idx = self.tp_rank * shard_size + loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size) + param_data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + base_out = F.linear(x, self.weight, self.bias) + return self.lora_forward(x, base_out) + + +class MergedColumnParallelLinear(ColumnParallelLinear): + + def __init__( + self, + input_size: int, + output_sizes: list[int], + bias: bool = False, + r: int = 0, + lora_alpha: float = 1.0, + lora_dropout: float = 0.0, + ): + self.output_sizes = output_sizes + super().__init__(input_size, sum(output_sizes), bias=bias, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout) + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int): + param_data = param.data + shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size + shard_size = self.output_sizes[loaded_shard_id] // self.tp_size + param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size) + loaded_weight = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank] + param_data.copy_(loaded_weight) + + +class QKVParallelLinear(ColumnParallelLinear): + + def __init__( + self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: int | None = None, + bias: bool = False, + r: int = 0, + lora_alpha: float = 1.0, + lora_dropout: float = 0.0, + ): + self.head_size = head_size + self.total_num_heads = total_num_heads + self.total_num_kv_heads = total_num_kv_heads or total_num_heads + tp_size = dist.get_world_size() + self.num_heads = divide(self.total_num_heads, tp_size) + self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) + input_size = hidden_size + output_size = (self.total_num_heads + 2 * self.total_num_kv_heads) * self.head_size + super().__init__(input_size, output_size, bias, r, lora_alpha, lora_dropout) + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str): + param_data = param.data + assert loaded_shard_id in ["q", "k", "v"] + if loaded_shard_id == "q": + shard_size = self.num_heads * self.head_size + shard_offset = 0 + elif loaded_shard_id == "k": + shard_size = self.num_kv_heads * self.head_size + shard_offset = self.num_heads * self.head_size + else: + shard_size = self.num_kv_heads * self.head_size + shard_offset = self.num_heads * self.head_size + self.num_kv_heads * self.head_size + param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size) + loaded_weight = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank] + param_data.copy_(loaded_weight) + + +class RowParallelLinear(LinearBase, LoRAMixin): + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = False, + r: int = 0, + lora_alpha: float = 1.0, + lora_dropout: float = 0.0, + ): + LinearBase.__init__(self, input_size, output_size, 1) + self.input_size_per_partition = divide(input_size, self.tp_size) + self.output_size_per_partition = output_size + + self.weight = nn.Parameter(torch.empty(self.output_size, self.input_size_per_partition)) + self.weight.weight_loader = self.weight_loader + if bias: + self.bias = nn.Parameter(torch.empty(self.output_size)) + self.bias.weight_loader = self.weight_loader + else: + self.register_parameter("bias", None) + + self.__init_lora__(r, lora_alpha, lora_dropout) + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + param_data = param.data + shard_size = param_data.size(self.tp_dim) + start_idx = self.tp_rank * shard_size + loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size) + param_data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = F.linear(x, self.weight, self.bias if self.tp_rank == 0 else None) + if self.tp_size > 1: + dist.all_reduce(y) + return self.lora_forward(x, y) diff --git a/diffuserve/layer/rotary_embedding.py b/diffuserve/layer/rotary_embedding.py new file mode 100755 index 00000000..6b206332 --- /dev/null +++ b/diffuserve/layer/rotary_embedding.py @@ -0,0 +1,79 @@ +import torch +import torch.nn as nn + +from functools import lru_cache + + +def apply_rotary_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> torch.Tensor: + cos = cos.unsqueeze(-2) + sin = sin.unsqueeze(-2) + x1, x2 = torch.chunk(x.to(torch.float32), 2, dim=-1) + y1 = x1 * cos - x2 * sin + y2 = x2 * cos + x1 * sin + return torch.cat((y1, y2), dim=-1).to(x.dtype) + + +class RotaryEmbedding(nn.Module): + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: float, + ) -> None: + super().__init__() + self.head_size = head_size + assert rotary_dim == head_size + inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)) + t = torch.arange(max_position_embeddings, dtype=torch.float) + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() + sin = freqs.sin() + cache = torch.cat((cos, sin), dim=-1) + self.register_buffer("cos_sin_cache", cache, persistent=False) + + @torch.compile + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Derive shapes from the tensors being viewed to keep SymInts consistent for torch.compile. + # This avoids FakeTensor failing to prove equality between independent symbolic dims + # coming from positions.size(0) vs query.size(0). + q_tokens = query.size(0) + k_tokens = key.size(0) + + cos_sin = self.cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + + # Reshape using only sizes from the target tensor for Dynamo friendliness + query_shape = query.shape + nheads_q = query_shape[-1] // self.head_size + query = query.view(q_tokens, nheads_q, self.head_size) + query = apply_rotary_emb(query, cos, sin).view(query_shape) + + key_shape = key.shape + nheads_k = key_shape[-1] // self.head_size + key = key.view(k_tokens, nheads_k, self.head_size) + key = apply_rotary_emb(key, cos, sin).view(key_shape) + return query, key + + +@lru_cache(1) +def get_rope( + head_size: int, + rotary_dim: int, + max_position: int, + base: float, + rope_scaling: dict | None = None, +): + assert rope_scaling is None + rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base) + return rotary_emb diff --git a/diffuserve/layer/sampler.py b/diffuserve/layer/sampler.py new file mode 100644 index 00000000..28dd284a --- /dev/null +++ b/diffuserve/layer/sampler.py @@ -0,0 +1,217 @@ +import torch + +import torch.nn as nn +import torch.nn.functional as F +import torch.distributions as dists + +from typing import List, Dict +from dataclasses import dataclass +from easydict import EasyDict as edict + +from diffuserve.config import Config +from diffuserve.utils.context import get_context_diffusion_lm + + +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 = get_context_diffusion_lm() + 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 = get_context_diffusion_lm() + 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/diffuserve/legacy/__init__.py b/diffuserve/legacy/__init__.py index e2d13424..82b9e51b 100755 --- a/diffuserve/legacy/__init__.py +++ b/diffuserve/legacy/__init__.py @@ -1,3 +1,2 @@ from diffuserve.legacy.llm import LLM from diffuserve.legacy.sampling_params import SamplingParams -from diffuserve.legacy.engine.async_engine import AsyncEngine diff --git a/diffuserve/llm.py b/diffuserve/llm.py new file mode 100755 index 00000000..9a86e4ce --- /dev/null +++ b/diffuserve/llm.py @@ -0,0 +1,10 @@ +from diffuserve.engine.llm_engine import LLMEngine +from diffuserve.engine.dp_engine import DPEngine +from diffuserve.config import Config + +class LLM: + def __new__(cls, model, **kwargs): + cfg = Config(model, **{k: v for k, v in kwargs.items() if k in Config.__dataclass_fields__.keys()}) + if cfg.data_parallel_size > 1: + return DPEngine(model, **kwargs) + return LLMEngine(model, **kwargs) diff --git a/diffuserve/model/auto_model.py b/diffuserve/model/auto_model.py new file mode 100755 index 00000000..f39c5cf1 --- /dev/null +++ b/diffuserve/model/auto_model.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from typing import Any, Callable + +from diffuserve.legacy.config import Config +from diffuserve.legacy.utils.loader import load_model +from diffuserve.legacy.models.dream import DreamForDiffusionLM +from diffuserve.legacy.models.llada import LLaDAForDiffusionLM + +_NOT_PROVIDED = object() +RegistryEntry = tuple[Callable[[Any], Any] | type | None, bool] + +class AutoModelForDiffusionLM: + """Factory and registry for diffusion language models.""" + + MODEL_MAPPING: dict[str, RegistryEntry] = {} + + @classmethod + def register( + cls, + model_name: str, + model_class: Callable[[Any], Any] | type | None = _NOT_PROVIDED, + *, + use_full_config: bool = False, + exist_ok: bool = False, + ): + """Register a model factory or class under ``model_name``. + + When ``model_class`` is omitted this method returns a decorator. + + Args: + model_name: Key used to retrieve the model. + model_class: Callable or class that builds the model 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(model_name, str) or not model_name: + raise ValueError("model_name must be a non-empty string.") + + if model_class is _NOT_PROVIDED: + def decorator(model_cls): + cls._register(model_name, model_cls, use_full_config=use_full_config, exist_ok=exist_ok) + return model_cls + + return decorator + + cls._register(model_name, model_class, use_full_config=use_full_config, exist_ok=exist_ok) + return model_class + + @classmethod + def _register( + cls, + model_name: str, + model_class: Callable[[Any], Any] | type | None, + *, + use_full_config: bool, + exist_ok: bool, + ) -> None: + if not exist_ok and model_name in cls.MODEL_MAPPING: + raise ValueError(f"Model '{model_name}' is already registered.") + cls.MODEL_MAPPING[model_name] = (model_class, use_full_config) + + @classmethod + def unregister(cls, model_name: str) -> None: + cls.MODEL_MAPPING.pop(model_name, None) + + @classmethod + def available_models(cls) -> tuple[str, ...]: + return tuple(sorted(cls.MODEL_MAPPING)) + + @classmethod + def from_config(cls, config: Config): + if not hasattr(config, "model_name"): + raise AttributeError("Config must define 'model_name' to build a model.") + + try: + factory, use_full_config = cls.MODEL_MAPPING[config.model_name] + except KeyError as err: + available = ", ".join(cls.available_models()) or "" + raise ValueError( + f"Model '{config.model_name}' is not registered. Available models: {available}." + ) from err + + if factory is None: + raise ValueError(f"Model '{config.model_name}' is reserved but not implemented yet.") + + init_arg = config if use_full_config else config.hf_config + if init_arg is None: + raise ValueError("Config.hf_config must be initialized before building the model.") + + model = factory(init_arg) + return load_model(model, config) + +# Backwards compatibility with the old name while callers migrate. +AutoModelLM = AutoModelForDiffusionLM \ No newline at end of file diff --git a/diffuserve/model/config/dream/configuration_dream.py b/diffuserve/model/config/dream/configuration_dream.py new file mode 100755 index 00000000..6a8c49df --- /dev/null +++ b/diffuserve/model/config/dream/configuration_dream.py @@ -0,0 +1,88 @@ + +# coding=utf-8 +# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved. +# +# 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. +"""Dream model configuration""" + +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 DreamConfig(PretrainedConfig): + model_type = "Dream" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=151936, + hidden_size=4096, + intermediate_size=22016, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=False, # cache not used in diffusion + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=28, + attention_dropout=0.0, + mask_token_id=151666, + pad_token_id=151643, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + 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 + + # for 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.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_dropout = attention_dropout + # Validate the correctness of rotary position embeddings parameters + # BC: if there is a 'type' field, move it to 'rope_type'. + 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, + **kwargs, + ) + self.mask_token_id = mask_token_id + self.pad_token_id = pad_token_id + diff --git a/diffuserve/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py b/diffuserve/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py new file mode 100755 index 00000000..ab484c64 --- /dev/null +++ b/diffuserve/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py @@ -0,0 +1,90 @@ + +# coding=utf-8 +# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved. +# +# 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. +"""FastdLLM V2 model configuration""" + +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 FastdLLMV2Config(PretrainedConfig): + model_type = "FastdLLMV2" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=151936, + hidden_size=4096, + intermediate_size=22016, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=False, # cache not used in diffusion + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=28, + attention_dropout=0.0, + mask_token_id=151665, + pad_token_id=151643, + bd_size=32, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + 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 + + # for 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.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_dropout = attention_dropout + # Validate the correctness of rotary position embeddings parameters + # BC: if there is a 'type' field, move it to 'rope_type'. + 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, + **kwargs, + ) + self.mask_token_id = mask_token_id + self.pad_token_id = pad_token_id + self.bd_size = bd_size + diff --git a/diffuserve/model/config/llada/configuration_llada.py b/diffuserve/model/config/llada/configuration_llada.py new file mode 100644 index 00000000..01cd50ef --- /dev/null +++ b/diffuserve/model/config/llada/configuration_llada.py @@ -0,0 +1,459 @@ +""" +LLaDA configuration +""" +from transformers import AutoConfig, PretrainedConfig + +from enum import Enum +from os import PathLike +from typing import Union +from dataclasses import asdict, dataclass, field +from glob import glob +from pathlib import Path +from typing import ( + Any, + Dict, + Iterable, + List, + Optional, + Tuple, + Type, + TypeVar, + Union, + cast, +) + + +__all__ = [ + "ActivationType", + "ActivationCheckpointingStrategy", + "BlockType", + "LayerNormType", + "InitFnType", + "ModelConfig", +] + +PathOrStr = Union[str, PathLike] + + +class StrEnum(str, Enum): + """ + This is equivalent to Python's :class:`enum.StrEnum` since version 3.11. + We include this here for compatibility with older version of Python. + """ + + def __str__(self) -> str: + return self.value + + def __repr__(self) -> str: + return f"'{str(self)}'" + + +class LayerNormType(StrEnum): + default = "default" + """ + The default LayerNorm implementation, equivalent to PyTorch's built-in version. + """ + + low_precision = "low_precision" + """ + A low-precision version of the default LayerNorm. + """ + + rms = "rms" + """ + An RMSNorm implementation. When using ``torch.compile`` this is + probably the fastest implementation. + """ + + gemma_rms = "gemma_rms" + """ + An RMSNorm implementation by gemmma. When using ``torch.compile`` this is + probably the fastest implementation. + """ + + amd_compatible = "amd_compatible" + """ + LayerNorm implemented manually to work around an issue with ROCm. + """ + + +class ActivationType(StrEnum): + gelu = "gelu" + relu = "relu" + silu = "silu" + swiglu = "swiglu" + + +class BlockType(StrEnum): + sequential = "sequential" + parallel = "parallel" + + llama = "llama" + """ + A block similar to the sequential block with slightly different + implementations of operations like attention to imitate the behavior of Llama. + """ + + +class InitFnType(StrEnum): + mitchell = "mitchell" + """ + The strategy suggested to us by Mitchell Wortsman from UW. + This uses a truncated normal distribution with an adaptive standard deviation that depends + on the size of the weights as well as the depth of the layer. + """ + + normal = "normal" + """ + All weights are initialized from the same normal distribution. + """ + + kaiming_normal = "kaiming_normal" + """ + All weights are initialized with the Kaiming method from a normal distribution. + Note this currently won't work with FSDP. + """ + + fan_in = "fan_in" + """ + "Fan-in variance scaling", i.e. normal with a standard deviation of ``1/sqrt(d_in)`` where ``d_in`` + is the input dimensionality of the kernel. + """ + + full_megatron = "full_megatron" + """ + This is what metaseq calls "full megatron init". It is the init used for Llama 2. + """ + + +@dataclass +class ModelConfig(): + """ + LLaDA (model) configuration. + """ + + # Note that the defaults for these attributes are equivalent to the base GPT2 model. + + d_model: int = 768 + """ + The hidden size of the model. + """ + + n_heads: int = 12 + """ + The number of self-attention heads. + """ + + n_kv_heads: Optional[int] = None + """ + The number of heads to use for keys and values. Defaults to `n_heads`. + Set this to ``None`` or ``n_heads`` for normal multi-head attention. + Set this to 1 for multi-query attention. + Set it to some in-between value for Llama2-style grouped query attention. + """ + + n_layers: int = 12 + """ + The number of layers/blocks. + """ + + mlp_ratio: int = 4 + """ + The ratio of the inner MLP dimensionality to ``d_model``. + This is only used when ``mlp_hidden_size`` is not set. + """ + + mlp_hidden_size: Optional[int] = None + """ + Set the exact hidden size for the MLP. Otherwise the inner MLP hidden size will be set to `mlp_ratio * d_model`. + """ + + activation_type: ActivationType = ActivationType.swiglu + """ + The activation function to use within the MLP layers. + """ + + block_type: BlockType = BlockType.sequential + """ + The transformer block implementation. + """ + + block_group_size: int = 1 + """ + The number of blocks to group together into a single parent block. + This has no affect on the number of parameters in the model and is only used to wrap groups + of blocks together with a single FSDP wrapper during training. + """ + + alibi: bool = False + """ + If ``True``, use ALiBi embeddings. Mutually exclusive with ``rope``. + """ + + alibi_bias_max: float = 8.0 + """ + Maximum absolute value of ALiBi bias. + """ + + rope: bool = False + """ + Use rotary positional embeddings (RoPE). Mutually exclusive with ``alibi``. + """ + + rope_full_precision: bool = True + """ + If ``True``, apply RoPE embeddings at full precision regardless of the input type. Otherwise, + apply RoPE at the precision of the input. + """ + + flash_attention: bool = False + """ + If ``True``, use ``FlashAttention``. + """ + + attention_dropout: float = 0.1 + """ + The dropout probability within the attention modules. + """ + + multi_query_attention: Optional[bool] = None + """ + Use the Multi-Query formulation of attention used in PaLM. This reduces the number of parameters + and is more efficient during inference. + """ + + attention_layer_norm: bool = False + """ + Apply layer norm to the keys and queries within the attention mechanism. + This can help stabilize training. + """ + + residual_dropout: float = 0.1 + """ + The dropout probability for the MLP and attention output within each block. + """ + + embedding_dropout: float = 0.1 + """ + The dropout probability for embeddings. + """ + + input_emb_norm: bool = False + """ + An input hidden_states norm implementation by gemmma. + """ + + layer_norm_type: LayerNormType = LayerNormType.default + """ + The layernorm implementation to use. + """ + + layer_norm_with_affine: bool = True + """ + Whether to include bias and weight parameters for the layer norms. + This only affects layer norms that are immediately followed by a linear layer in the forward pass, + so everything except QK-norms. To turn off affines for QK norms as well, set :attr:`attention_layer_norm_with_affine` + to ``False``. + """ + + rms_norm_eps: float = 1e-05 + """ + The rms layernorm eps param. + """ + + attention_layer_norm_with_affine: bool = True + """ + Toggle affine transform for the QK norms. + """ + + max_sequence_length: int = 1024 + """ + The maximum input sequence length supported by the model. + """ + + rope_theta: float = 10000.0 + """ + The rope base param. + """ + + include_qkv_bias: Optional[bool] = False + """ + Whether or not to include bias parameters in qkv linear layers. + """ + + include_bias: bool = False + """ + Whether or not to include bias parameters in linear layers. + In PaLM, they got rid of all bias terms because they found that large + models tend to have near 0 bias terms anyway. + """ + + bias_for_layer_norm: Optional[bool] = None + """ + Whether or not to include bias parameters in layer norm. + This is separate from the include_bias parameter, because of a ROCm crash when biases are disabled in + layer norm. + When this is None (the default), it inherits the setting from include_bias. + """ + + scale_logits: bool = False + """ + If ``True``, scale the output logits by ``1 / sqrt(d_model)``. + """ + + vocab_size: int = 50257 + """ + Vocabulary size of the model. + """ + + embedding_size: Optional[int] = 50304 + """ + The number of embeddings, i.e. the number of tokens. If set to ``None`` it will default + to ``vocab_size``. If ``vocab_size`` is not a multiple of 128, setting this to the + next multiple of 128 that's greater than ``vocab_size`` can improve throughput + substantially. + """ + + weight_tying: bool = True + """ + Whether to tie output linear weights to the input embedding. + """ + + eos_token_id: int = 50256 + """ + The ID of the end-of-sentence special token. + """ + + pad_token_id: int = 50256 + """ + The ID of the token to use for padding. Defaults to the ID of the EOS token. + """ + + mask_token_id: Optional[int] = 50256 + """ + The ID of the token to use for mask token. Defaults to the ID of the EOS token. + """ + + init_device: Optional[str] = None + """ + The torch device to use when initializing the model parameters, e.g. "cpu", "cuda:0", "meta". + """ + + init_fn: InitFnType = InitFnType.normal + """ + The weight initialization strategy. + """ + + init_std: float = 0.02 + """ + The standard deviation to use when initializing weights with a "fixed distribution" ``init_fn``, such + as "normal". + """ + + init_cutoff_factor: Optional[float] = None + """ + A positive factor used to scale the cutoff values when initializing weights with a "fixed distribution" ``init_fn``, such + as "normal". Setting this to None means values are not cutoff. + """ + + precision: Optional[str] = None + """ + Precision used to train/evaluate with. You shouldn't set this directly. + See :data:`TrainConfig.precision` instead. + """ + + @property + def effective_n_kv_heads(self) -> int: + if self.n_kv_heads is None: + if self.multi_query_attention is True: + return 1 + else: + return self.n_heads + else: + if self.multi_query_attention is None: + return self.n_kv_heads + if self.multi_query_attention: + n_kv_heads_should_be = 1 + else: + n_kv_heads_should_be = self.n_heads + if self.n_kv_heads == n_kv_heads_should_be: + return n_kv_heads_should_be + else: + raise Exception( + "You can't set `multi_query_attention` and `n_kv_heads` at the same time." + ) + +class ActivationCheckpointingStrategy(StrEnum): + whole_layer = "whole_layer" + """ + Checkpoint every transformer layer. + """ + + one_in_two = "one_in_two" + """ + Checkpoint one in two transformer layers. + """ + + one_in_three = "one_in_three" + """ + Checkpoint one in three transformer layers. + """ + + one_in_four = "one_in_four" + """ + Checkpoint one in four transformer layers. + """ + + two_in_three = "two_in_three" + """ + Checkpoint two out of every three transformer layers. + """ + + three_in_four = "three_in_four" + """ + Checkpoint three out of four of every transformer layers. + """ + + four_in_five = "four_in_five" + """ + Checkpoint four out of five of every transformer layers. + """ + + nine_in_ten = "nine_in_ten" + """ + Checkpoint nine out of ten of every transformer layers. + """ + + fine_grained = "fine_grained" + """ + Focus checkpointing on where it is cheap to recompute and saves most memory. + """ + + +class LLaDAConfig(PretrainedConfig): + model_type = "llada" + keys_to_ignore_at_inference = ["past_key_values"] # TODO: confirm + + def __init__(self, use_cache: bool = False, **kwargs): + model_config = ModelConfig() + all_kwargs = model_config.__dict__ + all_kwargs.update(kwargs) + all_kwargs.update({"use_cache": use_cache}) + all_kwargs.update( + { + "architectures": all_kwargs.get("architectures", ["LLaDAModelLM"]) + } + ) + super().__init__(**all_kwargs) + + @property + def num_attention_heads(self): + return self.n_heads + + @property + def num_hidden_layers(self): + return self.n_layers + + @property + def hidden_size(self): + return self.d_model \ No newline at end of file diff --git a/diffuserve/model/diffucoder.py b/diffuserve/model/diffucoder.py new file mode 100644 index 00000000..e69de29b diff --git a/diffuserve/model/dream.py b/diffuserve/model/dream.py new file mode 100755 index 00000000..bb2d7a12 --- /dev/null +++ b/diffuserve/model/dream.py @@ -0,0 +1,238 @@ +import os +import torch +import torch.nn as nn +import torch.distributed as dist + +from diffuserve.layer.layernorm import RMSNorm +from diffuserve.layer.activation import SiluAndMul +from diffuserve.layer.rotary_embedding import get_rope +from diffuserve.layer.attention.attention_v5 import Attention +from diffuserve.model.auto_model import AutoModelForDiffusionLM +from diffuserve.model.config.dream.configuration_dream import DreamConfig +from diffuserve.layer.linear import RowParallelLinear, ColumnParallelLinear +from diffuserve.layer.embed_head import VocabParallelEmbedding, ParallelLMHead + + +if os.environ.get("TRITON_INTERPRET", None) == "1": + torch._dynamo.reset() + torch._dynamo.config.suppress_errors = True + torch.backends.optimized_mode = False + + +class DreamRMSNorm(RMSNorm): + def __init__(self, hidden_size, eps=1e-6): + super().__init__(hidden_size, eps) + + +class DreamAttention(nn.Module): + """Dream attention mechanism.""" + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int = 32768, + head_dim: int | None = None, + rms_norm_eps: float = 1e-6, + qkv_bias: bool = True, + rope_theta: float = 10000, + rope_scaling: tuple | None = None, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + self.head_dim = head_dim or 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 + + self.q_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim, + bias=qkv_bias, + ) + self.k_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.v_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + "diffusion_lm", # Dream uses full attention + ) + + 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) + + q, k = self.rotary_emb(positions, q, k) + o = self.attn(q, k, v, mask) + output = self.o_proj(o) + return output + + +class DreamMLP(nn.Module): + """Dream MLP with SiLU activation.""" + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + ) + assert hidden_act == "silu" + self.act_fn = SiluAndMul() + + def forward(self, x): + gate = self.gate_proj(x) + up = self.up_proj(x) + x = self.act_fn(torch.cat([gate, up], dim=-1)) + x = self.down_proj(x) + return x + + +class DreamDecoderLayer(nn.Module): + """Dream transformer decoder layer.""" + def __init__( + self, + config: DreamConfig, + ) -> None: + super().__init__() + self.self_attn = DreamAttention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + max_position=config.max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=True, # Dream uses bias in attention + head_dim=getattr(config, 'head_dim', None), + rope_theta=getattr(config, "rope_theta", 10000), + rope_scaling=getattr(config, "rope_scaling", None), + ) + self.mlp = DreamMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + self.input_layernorm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = DreamRMSNorm(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 DreamModel(nn.Module): + """Dream model for diffusion language modeling.""" + def __init__( + self, + config: DreamConfig, + ) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([DreamDecoderLayer(config) + for _ in range(config.num_hidden_layers)]) + self.norm = DreamRMSNorm(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 enumerate(self.layers): + hidden_states, residual = layer(positions, hidden_states, residual, mask) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + +@AutoModelForDiffusionLM.register("dream") +class DreamForDiffusionLM(nn.Module): + """Dream model for diffusion language modeling with LM head.""" + packed_modules_mapping = {} + + def __init__( + self, + config: DreamConfig, + ) -> None: + super().__init__() + self.model = DreamModel(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, model_type='diffusion_lm') + 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: + hidden_states = self.model(input_ids, positions, mask) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + logits = self.lm_head(hidden_states) + return logits diff --git a/diffuserve/model/fast_dllm_v2.py b/diffuserve/model/fast_dllm_v2.py new file mode 100755 index 00000000..4fcc4740 --- /dev/null +++ b/diffuserve/model/fast_dllm_v2.py @@ -0,0 +1,239 @@ +import os +import torch +import torch.nn as nn +import torch.distributed as dist + +from diffuserve.layer.layernorm import RMSNorm +from diffuserve.layer.activation import SiluAndMul +from diffuserve.layer.rotary_embedding import get_rope +from diffuserve.layer.attention.attention_v5 import Attention +from diffuserve.model.auto_model import AutoModelForDiffusionLM +from diffuserve.layer.linear import RowParallelLinear, ColumnParallelLinear +from diffuserve.layer.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffuserve.model.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config + + +if os.environ.get("TRITON_INTERPRET", None) == "1": + torch._dynamo.reset() + torch._dynamo.config.suppress_errors = True + torch.backends.optimized_mode = False + + +class FastdLLMV2RMSNorm(RMSNorm): + def __init__(self, hidden_size, eps=1e-6): + super().__init__(hidden_size, eps) + + +class FastdLLMV2Attention(nn.Module): + """FastdLLM V2 attention mechanism.""" + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int = 32768, + head_dim: int | None = None, + rms_norm_eps: float = 1e-6, + qkv_bias: bool = True, + rope_theta: float = 10000, + rope_scaling: tuple | None = None, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + self.head_dim = head_dim or 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 + + self.q_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim, + bias=qkv_bias, + ) + self.k_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.v_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + "diffusion_lm", # Dream uses full attention + ) + + 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) + + q, k = self.rotary_emb(positions, q, k) + o = self.attn(q, k, v, mask) + output = self.o_proj(o) + return output + + +class FastdLLMV2MLP(nn.Module): + """FastdLLM V2 MLP with SiLU activation.""" + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + ) + assert hidden_act == "silu" + self.act_fn = SiluAndMul() + + def forward(self, x): + gate = self.gate_proj(x) + up = self.up_proj(x) + x = self.act_fn(torch.cat([gate, up], dim=-1)) + x = self.down_proj(x) + return x + + +class FastdLLMV2DecoderLayer(nn.Module): + """FastdLLM V2 transformer decoder layer.""" + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.self_attn = FastdLLMV2Attention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + max_position=config.max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=True, # Dream uses bias in attention + head_dim=getattr(config, 'head_dim', None), + rope_theta=getattr(config, "rope_theta", 10000), + rope_scaling=getattr(config, "rope_scaling", None), + ) + self.mlp = FastdLLMV2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + self.input_layernorm = FastdLLMV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = FastdLLMV2RMSNorm(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 FastdLLMV2Model(nn.Module): + """FastdLLM V2 model for diffusion language modeling.""" + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([FastdLLMV2DecoderLayer(config) + for _ in range(config.num_hidden_layers)]) + self.norm = FastdLLMV2RMSNorm(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 enumerate(self.layers): + hidden_states, residual = layer(positions, hidden_states, residual, mask) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +@AutoModelForDiffusionLM.register("fast_dllm_v2") +class FastdLLMV2ForDiffusionLM(nn.Module): + """FastdLLM V2 model for diffusion language modeling with LM head.""" + packed_modules_mapping = {} + + def __init__( + self, + config: FastdLLMV2Config, + ) -> None: + super().__init__() + self.model = FastdLLMV2Model(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, model_type='diffusion_lm') + 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: + hidden_states = self.model(input_ids, positions, mask) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + logits = self.lm_head(hidden_states) + return logits diff --git a/diffuserve/model/llada.py b/diffuserve/model/llada.py new file mode 100755 index 00000000..84a16703 --- /dev/null +++ b/diffuserve/model/llada.py @@ -0,0 +1,266 @@ +import os +import torch +import torch.nn as nn +import torch.distributed as dist + +from diffuserve.layer.layernorm import RMSNorm +from diffuserve.layer.activation import SiluAndMul +from diffuserve.layer.rotary_embedding import get_rope +from diffuserve.layer.attention.attention_v5 import Attention +from diffuserve.model.auto_model import AutoModelForDiffusionLM +from diffuserve.model.config.llada.configuration_llada import LLaDAConfig +from diffuserve.layer.linear import RowParallelLinear, ColumnParallelLinear +from diffuserve.layer.embed_head import VocabParallelEmbedding, ParallelLMHead + + +if os.environ.get("TRITON_INTERPRET", None) == "1": + torch._dynamo.reset() + torch._dynamo.config.suppress_errors = True + torch.backends.optimized_mode = False + + +class LLaDARMSNorm(RMSNorm): + def __init__(self, hidden_size, eps=1e-6): + super().__init__(hidden_size, eps) + + +class LLaDAAttention(nn.Module): + """LLaDA attention.""" + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int = 32768, + head_dim: int | None = None, + rms_norm_eps: float = 1e-6, + qkv_bias: bool = True, + rope_theta: float = 10000, + rope_scaling: tuple | None = None, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + self.head_dim = head_dim or 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 + + self.q_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim, + bias=qkv_bias, + ) + self.k_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.v_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + "diffusion_lm", + ) + + 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) + + q, k = self.rotary_emb(positions, q, k) + o = self.attn(q, k, v, mask) + output = self.o_proj(o) + return output + + +class LLaDAMLP(nn.Module): + """LLaDA MLP.""" + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + ) + assert hidden_act == "silu" + self.act_fn = SiluAndMul() + + def forward(self, x): + gate = self.gate_proj(x) + up = self.up_proj(x) + x = self.act_fn(torch.cat([gate, up], dim=-1)) + x = self.down_proj(x) + return x + + +class LLaDABlock(nn.Module): + """LLaDA transformer block.""" + def __init__( + self, + config, + ) -> None: + super().__init__() + self.self_attn = LLaDAAttention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.n_kv_heads, + max_position=config.max_sequence_length, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=True, + head_dim=getattr(config, 'head_dim', None), + rope_theta=getattr(config, "rope_theta", 10000), + rope_scaling=getattr(config, "rope_scaling", None), + ) + self.mlp = LLaDAMLP( + hidden_size=config.hidden_size, + intermediate_size=config.mlp_hidden_size, + hidden_act=config.activation_type, + ) + self.input_layernorm = LLaDARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = LLaDARMSNorm(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 LLaDAModel(nn.Module): + """LLaDA backbone.""" + def __init__( + self, + config: LLaDAConfig, + ) -> None: + super().__init__() + self.config = config + self.transformer = nn.ModuleDict( + dict( + wte=VocabParallelEmbedding( + config.embedding_size or config.vocab_size, config.d_model + ), + emb_drop=nn.Dropout(config.embedding_dropout), + ln_f=LLaDARMSNorm(config.hidden_size, config.rms_norm_eps) + ) + ) + + blocks = [LLaDABlock(config) for _ in range(config.n_layers)] + self.transformer.update({"blocks": nn.ModuleList(blocks)}) + + if not (self.config.alibi or self.config.rope): + self.transformer.update( + {"wpe": nn.Embedding(config.max_sequence_length, config.d_model, device=config.init_device)} + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + mask: torch.Tensor | None = None + ) -> torch.Tensor: + hidden_states = self.transformer.emb_drop(self.transformer.wte(input_ids)) + residual = None + for block_idx, block in enumerate(self.transformer.blocks): + hidden_states, residual = block(positions, hidden_states, residual, mask) + hidden_states, _ = self.transformer.ln_f(hidden_states, residual) + return hidden_states + + +@AutoModelForDiffusionLM.register("llada") +class LLaDAForDiffusionLM(nn.Module): + """LLaDA with LM head.""" + packed_modules_mapping = { + "q_proj": ("self_attn.q_proj", None), + "k_proj": ("self_attn.k_proj", None), + "v_proj": ("self_attn.v_proj", None), + "attn_out": ("self_attn.o_proj", None), + "attn_norm": ("input_layernorm", None), + "ff_norm": ("post_attention_layernorm", None), + + "ff_proj": ("mlp.gate_proj", None), + "up_proj": ("mlp.up_proj", None), + "ff_out": ("mlp.down_proj", None), + + "transformer.ff_out": ("lm_head", None) + } + + def __init__( + self, + config: LLaDAConfig, + ) -> None: + super().__init__() + self.model = LLaDAModel(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, model_type='diffusion_lm') + if getattr(config, 'weight_tying', False): + self.lm_head.weight.data = self.model.transformer.wte.weight.data + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + mask: torch.Tensor | None = None + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, mask) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + logits = self.lm_head(hidden_states) + return logits diff --git a/diffuserve/model/llada2.py b/diffuserve/model/llada2.py new file mode 100644 index 00000000..e69de29b diff --git a/diffuserve/model/llada_moe.py b/diffuserve/model/llada_moe.py new file mode 100644 index 00000000..e69de29b diff --git a/diffuserve/model/sdar.py b/diffuserve/model/sdar.py new file mode 100644 index 00000000..e69de29b diff --git a/diffuserve/model/utils/check_config.py b/diffuserve/model/utils/check_config.py new file mode 100755 index 00000000..9261e45b --- /dev/null +++ b/diffuserve/model/utils/check_config.py @@ -0,0 +1,8 @@ +from typing import Optional, Type + +def check_config_diff( + current_config: Optional[Type] = None, + default_config_cls: Optional[Type] = None, + compare_to: str = "default", # "default" or "existing" +): + pass \ No newline at end of file diff --git a/diffuserve/sampling_params.py b/diffuserve/sampling_params.py new file mode 100755 index 00000000..67d60bb4 --- /dev/null +++ b/diffuserve/sampling_params.py @@ -0,0 +1,8 @@ +from dataclasses import dataclass + + +@dataclass +class SamplingParams: + temperature: float = 1.0 + max_tokens: int = 64 + ignore_eos: bool = False diff --git a/diffuserve/utils/checker.py b/diffuserve/utils/checker.py new file mode 100755 index 00000000..e933806e --- /dev/null +++ b/diffuserve/utils/checker.py @@ -0,0 +1,28 @@ +def CHECK_SLOT_MAPPING(seqs, slot_mapping): + # check slot mapping layout + start_idx = 0 + for seq in seqs: + cur_ref_slot_mapping = [] + for idx in range(seq.num_diffusion_blocks): + if seq.active_blocks[idx]: + padding_num_tokens = (seq.num_diffusion_blocks - idx) * seq.diffusion_block_size + cur_ref_slot_mapping.extend([-1] * padding_num_tokens) + break + elif seq.to_cache_blocks[idx]: + cur_ref_slot_mapping.extend([0] * seq.diffusion_block_size) + cur_slot_mapping = slot_mapping[start_idx:start_idx + len(cur_ref_slot_mapping)] + for slot, ref_slot in zip(cur_slot_mapping, cur_ref_slot_mapping): + try: + if ref_slot == -1: + assert slot == -1 + elif ref_slot == 0: + assert slot != -1 + elif ref_slot is not None: + assert slot is not None + except AssertionError: + raise ValueError(f"Slot mapping mismatch: {slot} != {ref_slot}. " + f"Check the implementation of prepare_decode.\n" + f"slot_mapping: {cur_slot_mapping}\n" + f"ref_slot_mapping: {cur_ref_slot_mapping}\n" + f"diff: {[s - r for s, r in zip(cur_slot_mapping, cur_ref_slot_mapping)]}") + start_idx += len(cur_ref_slot_mapping) \ No newline at end of file diff --git a/diffuserve/utils/context.py b/diffuserve/utils/context.py new file mode 100755 index 00000000..a5ed2575 --- /dev/null +++ b/diffuserve/utils/context.py @@ -0,0 +1,112 @@ +import torch + +from dataclasses import dataclass + +from diffuserve.legacy.engine.sequence import SequenceForDiffusionLM + +@dataclass +class ContextBase: + 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 + + +# Global context for diffusion language model +@dataclass +class ContextForDiffusionLM(ContextBase): + seqs: list[SequenceForDiffusionLM] = None + seq_lens: list[int] = None + seq_lens_ts: torch.Tensor | None = None + kv_cache_layout: str = "unified" # "unified" or "distinct" + need_kv_cache_store: bool = True + block_mask: list[torch.Tensor] | None = None + + 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: + 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 block_mask_for_checking(self) -> torch.Tensor: + for seq in self.seqs: + seq.set_layout("unified") + + 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) + 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 + block_mask[start_row:end_row, start_col:end_col] = mask.clone() + start_row, start_col = end_row, start_col + + for seq in self.seqs: + seq.set_layout("distinct") + return block_mask.to(mask.device) + + @property + def total_num_seqs(self) -> int: + return len(self.seqs) if self.seqs is not None else 0 + +_CONTEXT_FOR_DIFFUSION_LM = ContextForDiffusionLM() + +def get_context_diffusion_lm() -> ContextForDiffusionLM: + return _CONTEXT_FOR_DIFFUSION_LM + +def set_context_diffusion_lm( + is_prefill, + cu_seqlens_q=None, cu_seqlens_k=None, + max_seqlen_q=0, max_seqlen_k=0, + slot_mapping=None, context_lens=None, block_tables=None, + seqs=None, seq_lens=None, seq_lens_ts=None, kv_cache_layout="unified", need_kv_cache_store=True, + d2f_pp=False +) -> None: + global _CONTEXT_FOR_DIFFUSION_LM + _CONTEXT_FOR_DIFFUSION_LM = ContextForDiffusionLM( + is_prefill, + cu_seqlens_q, cu_seqlens_k, + max_seqlen_q, max_seqlen_k, + slot_mapping, context_lens, block_tables, + seqs, seq_lens, seq_lens_ts, kv_cache_layout, need_kv_cache_store, d2f_pp + ) + +def reset_context_diffusion_lm() -> None: + global _CONTEXT_FOR_DIFFUSION_LM + _CONTEXT_FOR_DIFFUSION_LM = ContextForDiffusionLM() \ No newline at end of file diff --git a/diffuserve/utils/loader.py b/diffuserve/utils/loader.py new file mode 100755 index 00000000..c16a86bf --- /dev/null +++ b/diffuserve/utils/loader.py @@ -0,0 +1,204 @@ +import os +import json +import torch +import torch.nn as nn + +from tqdm import tqdm +from glob import glob +from functools import partial +from safetensors import safe_open +from diffuserve.legacy.config import Config + + +def load_lora_config(lora_path: str) -> dict: + """Load LoRA configuration from adapter_config.json.""" + config_path = os.path.join(lora_path, "adapter_config.json") + if os.path.exists(config_path): + with open(config_path, 'r') as f: + return json.load(f) + return {} + + +def enable_lora_for_model(model: nn.Module, lora_config: dict): + """Enable LoRA for existing linear layers in the model.""" + r = lora_config.get('r', 16) + lora_alpha = lora_config.get('lora_alpha', 32.0) + lora_dropout = lora_config.get('lora_dropout', 0.0) + target_modules = lora_config.get('target_modules', []) + + for name, module in model.named_modules(): + if hasattr(module, '__init_lora__'): + should_apply = True + if target_modules: + leaf = name.split('.')[-1] if name else name + should_apply = any(target == leaf for target in target_modules) + if should_apply: + module.__init_lora__(r, lora_alpha, lora_dropout) + return model + + +def default_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor): + param.data.copy_(loaded_weight) + + +def load_model(model: nn.Module, config: Config): + """Load model weights and optionally LoRA weights.""" + # Enable LoRA for linear layers if LoRA is enabled + if config.use_lora and config.lora_path: + lora_config = load_lora_config(config.lora_path) + if lora_config: + print(f"LoRA Config Loaded: {lora_config}") + model = enable_lora_for_model(model, lora_config) + else: + print("No adapter_config.json found, using default LoRA parameters") + default_config = {'r': 16, 'lora_alpha': 32.0, 'lora_dropout': 0.0} + model = enable_lora_for_model(model, default_config) + + # Load base model weights + packed_modules_mapping = getattr(model, "packed_modules_mapping", {}) + for file in tqdm(glob(os.path.join(config.model, "*.safetensors")), desc="Loading base model"): + with safe_open(file, "pt", "cpu") as f: + for weight_name in f.keys(): + for k in packed_modules_mapping: + if k in weight_name: + + if config.model_name == "llada" and k == "ff_out" and "transformer.ff_out" in weight_name: + continue + elif config.model_name == "llada" and k == "transformer.ff_out": + v, shard_id = packed_modules_mapping[k] + assert v == "lm_head" + param_name = "lm_head.weight" + else: + v, shard_id = packed_modules_mapping[k] + param_name = weight_name.replace(k, v) + + if "layernorm" in param_name: + param = model.get_parameter(param_name) + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, f.get_tensor(weight_name)) + else: + param = model.get_parameter(param_name) + weight_loader = partial(getattr(param, "weight_loader"), param, f.get_tensor(weight_name)) + if shard_id is None: + weight_loader() + else: + weight_loader(shard_id) + break + else: + param = model.get_parameter(weight_name) + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, f.get_tensor(weight_name)) + + # Load LoRA weights if enabled + if config.use_lora and config.lora_path: + if os.path.exists(config.lora_path): + print(f"Loading LoRA weights from {config.lora_path}") + load_lora_weights_fn = partial(load_lora_weights, model, config.lora_path) + packed_modules_mapping = packed_modules_mapping if config.model_name == "llada" else None + model = load_lora_weights_fn(packed_modules_mapping=packed_modules_mapping) + else: + print(f"Warning: LoRA path {config.lora_path} does not exist, skipping LoRA loading") + + return model + + +def load_lora_weights(model: nn.Module, lora_path: str, packed_modules_mapping: dict | None = None): + """Load LoRA weights into LoRA-enabled layers.""" + try: + lora_config = load_lora_config(lora_path) + target_modules = lora_config.get('target_modules', []) + + lora_weights = {} + + for file in tqdm(glob(os.path.join(lora_path, "*.safetensors")), desc="Loading LoRA"): + with safe_open(file, "pt", "cpu") as f: + for weight_name in f.keys(): + lora_weights[weight_name] = f.get_tensor(weight_name) + + applied_count = 0 + + modified_modules = None + if packed_modules_mapping is not None: + modified_modules = [v for k, (v, _) in packed_modules_mapping.items() if k in target_modules] + rev_mapping = {v: k for k, (v, _) in packed_modules_mapping.items()} + + for name, module in model.named_modules(): + if hasattr(module, 'lora_A') and hasattr(module, 'lora_B'): + should_apply = True + + if modified_modules is not None: + modified_module_type = '.'.join(name.split('.')[-2:]) + org_module_type = rev_mapping[modified_module_type] + org_name = name.replace(modified_module_type, org_module_type) + should_apply = any(target in modified_module_type for target in modified_modules) + elif target_modules: + module_type = name.split('.')[-1] if '.' in name else name + should_apply = any(target in module_type for target in target_modules) + + if not should_apply: + continue + + base_patterns = [ + name, + f"base_model.model.{name}", + f"model.{name}", + ] if modified_modules is None else [ + org_name, + f"base_model.model.{org_name}", + f"model.{org_name}", + ] + + found_a = found_b = None + for base_name in base_patterns: + lora_a_keys = [ + f"{base_name}.lora_A.weight", + f"{base_name}.lora_A.default.weight", + f"{base_name}.lora_A", + ] + lora_b_keys = [ + f"{base_name}.lora_B.weight", + f"{base_name}.lora_B.default.weight", + f"{base_name}.lora_B", + ] + + for key in lora_a_keys: + if key in lora_weights: + found_a = lora_weights[key] + break + for key in lora_b_keys: + if key in lora_weights: + found_b = lora_weights[key] + break + + if found_a is not None and found_b is not None: + break + + if found_a is not None and found_b is not None: + if hasattr(module, 'tp_size') and module.tp_size > 1: + if hasattr(module, 'tp_dim') and module.tp_dim == 0: + shard_size = found_b.size(0) // module.tp_size + start_idx = module.tp_rank * shard_size + found_b = found_b[start_idx:start_idx + shard_size] + elif hasattr(module, 'tp_dim') and module.tp_dim == 1: + shard_size = found_a.size(1) // module.tp_size + start_idx = module.tp_rank * shard_size + found_a = found_a[:, start_idx:start_idx + shard_size] + + try: + module.lora_A.data.copy_(found_a) + module.lora_B.data.copy_(found_b) + applied_count += 1 + except Exception as e: + print(f"Failed to load LoRA weights for {name}: {e}") + + for module in model.modules(): + if hasattr(module, 'merge_lora'): + module.merge_lora() + + print(f"LoRA weights applied to {applied_count} layers and merged") + + except Exception as e: + print(f"Error loading LoRA weights: {e}") + print("Continuing with base model only") + + return model diff --git a/pyproject.toml b/pyproject.toml index b9c3460c..3f9188a8 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,13 @@ [project] -name = "D2fEngine" +name = "DiffuServe" version = "0.0.1" authors = [ { name = "Drew Jin (Yijie Jin)", email = "drewjin0827@gmail.com"}, ] maintainers = [ - { name = "SJTU DENG Lab" } + { name = "DENG Lab @ SJTU" } ] -description = "vLLM implementation of Diffusion LLMs (based on D2F decoding paradigm)" +description = "Diffusion LLM serving engine supporting multiple variants of dLLM decoding strategies with high efficiency and low latency." readme = "README.md" requires-python = ">=3.12" license = { file = "LICENSE" } @@ -44,4 +44,4 @@ include = ["diffuserve"] [[tool.uv.index]] url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" -default = true +default = true \ No newline at end of file From f972d0122ecdbf3c8dc087ccad17c895a19704a1 Mon Sep 17 00:00:00 2001 From: drewjin Date: Wed, 12 Nov 2025 17:36:34 +0800 Subject: [PATCH 06/23] feat(strategy): add support for strategy registration; refactor: renaming the project to diffulex (diffusion + flex) --- README.md | 2 +- diffulex/__init__.py | 2 + {diffuserve => diffulex}/config.py | 1 + .../engine/block_manager.py | 153 +++--- {diffuserve => diffulex}/engine/dp_engine.py | 6 +- .../legacy => diffulex}/engine/llm_engine.py | 35 +- diffulex/engine/model_runner.py | 250 +++++++++ diffulex/engine/scheduler.py | 113 ++++ .../legacy => diffulex}/engine/sequence.py | 40 +- {diffuserve => diffulex}/layer/activation.py | 0 .../layer/attention/attention_v5.py | 4 +- diffulex/layer/attention/ops/__init__.py | 7 + ...chunked_prefill_decoding_unified_kernel.py | 2 +- .../layer}/attention/ops/kv_cache_kernels.py | 4 +- .../layer/attention/ops/prefix_prefill.py | 0 .../attention/ops/tilus_decode_attn_dlm.py | 0 .../attention/ops/triton_decode_attn_clm.py | 0 .../attention/ops/triton_decode_attn_dlm.py | 4 +- .../attention/ops/triton_flash_attention.py | 0 {diffuserve => diffulex}/layer/embed_head.py | 2 +- {diffuserve => diffulex}/layer/layernorm.py | 0 {diffuserve => diffulex}/layer/linear.py | 0 .../layer/rotary_embedding.py | 0 {diffuserve => diffulex}/layer/sampler.py | 4 +- diffulex/legacy/__init__.py | 2 + {diffuserve => diffulex}/legacy/config.py | 0 .../legacy}/engine/block_manager.py | 4 +- .../legacy/engine/dp_engine.py | 6 +- .../legacy}/engine/llm_engine.py | 10 +- .../legacy/engine/model_runner.py | 12 +- .../legacy}/engine/scheduler.py | 8 +- .../legacy}/engine/sequence.py | 4 +- .../legacy/layers/activation.py | 0 .../legacy/layers/attention/attention_v1.py | 2 +- .../layers/attention/attention_v1_profile.py | 2 +- .../legacy/layers/attention/attention_v2.py | 2 +- .../layers/attention/attention_v2_dup.py | 4 +- .../layers/attention/attention_v2_profile.py | 2 +- .../legacy/layers/attention/attention_v3.py | 4 +- .../legacy/layers/attention/attention_v4.py | 4 +- .../legacy/layers/attention/attention_v5.py | 4 +- .../legacy/layers/attention/ops/__init__.py | 7 + ...chunked_prefill_decoding_unified_kernel.py | 2 +- .../layers}/attention/ops/kv_cache_kernels.py | 4 +- .../layers/attention/ops/prefix_prefill.py | 0 .../attention/ops/tilus_decode_attn_dlm.py | 0 .../attention/ops/triton_decode_attn_clm.py | 0 .../attention/ops/triton_decode_attn_dlm.py | 4 +- .../attention/ops/triton_flash_attention.py | 0 .../legacy/layers/embed_head.py | 2 +- .../legacy/layers/layernorm.py | 0 .../legacy/layers/linear.py | 0 .../legacy/layers/rotary_embedding.py | 0 .../legacy/layers/sampler.py | 4 +- {diffuserve => diffulex}/legacy/llm.py | 6 +- .../legacy/models/auto_model.py | 10 +- .../config/dream/configuration_dream.py | 0 .../configuration_fast_dllm_v2.py | 0 .../config/llada/configuration_llada.py | 0 .../legacy/models/dream.py | 14 +- .../legacy/models/fast_dllm_v2.py | 14 +- .../legacy/models/llada.py | 14 +- .../legacy/models/qwen3.py | 12 +- .../legacy/models/utils/check_config.py | 0 .../legacy/sampling_params.py | 0 .../legacy/utils/checker.py | 0 .../legacy/utils/context.py | 2 +- .../legacy}/utils/loader.py | 2 +- {diffuserve => diffulex}/llm.py | 6 +- {diffuserve => diffulex}/model/auto_model.py | 8 +- .../model/config/dream/configuration_dream.py | 0 .../configuration_fast_dllm_v2.py | 0 .../model/config/llada/configuration_llada.py | 0 {diffuserve => diffulex}/model/diffucoder.py | 0 {diffuserve => diffulex}/model/dream.py | 16 +- .../model/fast_dllm_v2.py | 16 +- {diffuserve => diffulex}/model/llada.py | 16 +- {diffuserve => diffulex}/model/llada2.py | 0 {diffuserve => diffulex}/model/llada_moe.py | 0 {diffuserve => diffulex}/model/sdar.py | 0 .../model/utils/check_config.py | 0 {diffuserve => diffulex}/sampling_params.py | 0 diffulex/strategy/__init__.py | 7 + diffulex/strategy/d2f/__init__.py | 12 + diffulex/strategy/d2f/block_manager.py | 42 ++ diffulex/strategy/d2f/model_runner.py | 421 +++++++++++++++ diffulex/strategy/d2f/scheduler.py | 133 +++++ {diffuserve => diffulex}/utils/checker.py | 0 {diffuserve => diffulex}/utils/context.py | 2 +- .../legacy => diffulex}/utils/loader.py | 2 +- diffuserve/__init__.py | 2 - diffuserve/engine/model_runner.py | 507 ------------------ diffuserve/layer/attention/ops/__init__.py | 7 - diffuserve/legacy/__init__.py | 2 - diffuserve/legacy/engine/scheduler.py | 234 -------- .../legacy/layers/attention/ops/__init__.py | 7 - examples/test_causal_lm_decoding_kernel.py | 2 +- examples/test_dllm_decoding_kernel.py | 2 +- examples/test_dllm_kv_cache_load.py | 2 +- examples/test_dllm_kv_cache_store.py | 2 +- examples/test_dream_dvllm_gsm8k.py | 2 +- examples/test_dream_dvllm_human_eval.py | 2 +- examples/test_dream_model_weight.py | 4 +- examples/test_dream_model_weight_fixed.py | 4 +- examples/test_llada_dvllm_human_eval.py | 2 +- examples/test_qwen_dvllm.py | 2 +- pyproject.toml | 4 +- 107 files changed, 1232 insertions(+), 1032 deletions(-) create mode 100755 diffulex/__init__.py rename {diffuserve => diffulex}/config.py (99%) rename {diffuserve/legacy => diffulex}/engine/block_manager.py (50%) rename {diffuserve => diffulex}/engine/dp_engine.py (98%) rename {diffuserve/legacy => diffulex}/engine/llm_engine.py (79%) create mode 100755 diffulex/engine/model_runner.py create mode 100755 diffulex/engine/scheduler.py rename {diffuserve/legacy => diffulex}/engine/sequence.py (92%) rename {diffuserve => diffulex}/layer/activation.py (100%) rename {diffuserve => diffulex}/layer/attention/attention_v5.py (97%) create mode 100755 diffulex/layer/attention/ops/__init__.py rename {diffuserve => diffulex}/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py (99%) rename {diffuserve/legacy/layers => diffulex/layer}/attention/ops/kv_cache_kernels.py (99%) rename {diffuserve => diffulex}/layer/attention/ops/prefix_prefill.py (100%) rename {diffuserve => diffulex}/layer/attention/ops/tilus_decode_attn_dlm.py (100%) rename {diffuserve => diffulex}/layer/attention/ops/triton_decode_attn_clm.py (100%) rename {diffuserve/legacy/layers => diffulex/layer}/attention/ops/triton_decode_attn_dlm.py (97%) rename {diffuserve => diffulex}/layer/attention/ops/triton_flash_attention.py (100%) rename {diffuserve => diffulex}/layer/embed_head.py (96%) rename {diffuserve => diffulex}/layer/layernorm.py (100%) rename {diffuserve => diffulex}/layer/linear.py (100%) rename {diffuserve => diffulex}/layer/rotary_embedding.py (100%) rename {diffuserve => diffulex}/layer/sampler.py (98%) create mode 100755 diffulex/legacy/__init__.py rename {diffuserve => diffulex}/legacy/config.py (100%) rename {diffuserve => diffulex/legacy}/engine/block_manager.py (97%) rename {diffuserve => diffulex}/legacy/engine/dp_engine.py (98%) rename {diffuserve => diffulex/legacy}/engine/llm_engine.py (94%) rename {diffuserve => diffulex}/legacy/engine/model_runner.py (98%) rename {diffuserve => diffulex/legacy}/engine/scheduler.py (97%) rename {diffuserve => diffulex/legacy}/engine/sequence.py (99%) rename {diffuserve => diffulex}/legacy/layers/activation.py (100%) rename {diffuserve => diffulex}/legacy/layers/attention/attention_v1.py (99%) rename {diffuserve => diffulex}/legacy/layers/attention/attention_v1_profile.py (99%) rename {diffuserve => diffulex}/legacy/layers/attention/attention_v2.py (99%) rename {diffuserve => diffulex}/legacy/layers/attention/attention_v2_dup.py (98%) rename {diffuserve => diffulex}/legacy/layers/attention/attention_v2_profile.py (99%) rename {diffuserve => diffulex}/legacy/layers/attention/attention_v3.py (98%) rename {diffuserve => diffulex}/legacy/layers/attention/attention_v4.py (97%) rename {diffuserve => diffulex}/legacy/layers/attention/attention_v5.py (97%) create mode 100755 diffulex/legacy/layers/attention/ops/__init__.py rename {diffuserve => diffulex}/legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py (99%) rename {diffuserve/layer => diffulex/legacy/layers}/attention/ops/kv_cache_kernels.py (99%) rename {diffuserve => diffulex}/legacy/layers/attention/ops/prefix_prefill.py (100%) rename {diffuserve => diffulex}/legacy/layers/attention/ops/tilus_decode_attn_dlm.py (100%) rename {diffuserve => diffulex}/legacy/layers/attention/ops/triton_decode_attn_clm.py (100%) rename {diffuserve/layer => diffulex/legacy/layers}/attention/ops/triton_decode_attn_dlm.py (97%) rename {diffuserve => diffulex}/legacy/layers/attention/ops/triton_flash_attention.py (100%) rename {diffuserve => diffulex}/legacy/layers/embed_head.py (96%) rename {diffuserve => diffulex}/legacy/layers/layernorm.py (100%) rename {diffuserve => diffulex}/legacy/layers/linear.py (100%) rename {diffuserve => diffulex}/legacy/layers/rotary_embedding.py (100%) rename {diffuserve => diffulex}/legacy/layers/sampler.py (98%) rename {diffuserve => diffulex}/legacy/llm.py (64%) rename {diffuserve => diffulex}/legacy/models/auto_model.py (58%) rename {diffuserve => diffulex}/legacy/models/config/dream/configuration_dream.py (100%) rename {diffuserve => diffulex}/legacy/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py (100%) rename {diffuserve => diffulex}/legacy/models/config/llada/configuration_llada.py (100%) rename {diffuserve => diffulex}/legacy/models/dream.py (93%) rename {diffuserve => diffulex}/legacy/models/fast_dllm_v2.py (93%) rename {diffuserve => diffulex}/legacy/models/llada.py (94%) rename {diffuserve => diffulex}/legacy/models/qwen3.py (93%) rename {diffuserve => diffulex}/legacy/models/utils/check_config.py (100%) rename {diffuserve => diffulex}/legacy/sampling_params.py (100%) rename {diffuserve => diffulex}/legacy/utils/checker.py (100%) rename {diffuserve => diffulex}/legacy/utils/context.py (98%) rename {diffuserve => diffulex/legacy}/utils/loader.py (99%) rename {diffuserve => diffulex}/llm.py (67%) rename {diffuserve => diffulex}/model/auto_model.py (93%) rename {diffuserve => diffulex}/model/config/dream/configuration_dream.py (100%) rename {diffuserve => diffulex}/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py (100%) rename {diffuserve => diffulex}/model/config/llada/configuration_llada.py (100%) rename {diffuserve => diffulex}/model/diffucoder.py (100%) rename {diffuserve => diffulex}/model/dream.py (93%) rename {diffuserve => diffulex}/model/fast_dllm_v2.py (93%) rename {diffuserve => diffulex}/model/llada.py (94%) rename {diffuserve => diffulex}/model/llada2.py (100%) rename {diffuserve => diffulex}/model/llada_moe.py (100%) rename {diffuserve => diffulex}/model/sdar.py (100%) rename {diffuserve => diffulex}/model/utils/check_config.py (100%) rename {diffuserve => diffulex}/sampling_params.py (100%) create mode 100644 diffulex/strategy/__init__.py create mode 100644 diffulex/strategy/d2f/__init__.py create mode 100644 diffulex/strategy/d2f/block_manager.py create mode 100644 diffulex/strategy/d2f/model_runner.py create mode 100644 diffulex/strategy/d2f/scheduler.py rename {diffuserve => diffulex}/utils/checker.py (100%) rename {diffuserve => diffulex}/utils/context.py (98%) rename {diffuserve/legacy => diffulex}/utils/loader.py (99%) delete mode 100755 diffuserve/__init__.py delete mode 100755 diffuserve/engine/model_runner.py delete mode 100755 diffuserve/layer/attention/ops/__init__.py delete mode 100755 diffuserve/legacy/__init__.py delete mode 100755 diffuserve/legacy/engine/scheduler.py delete mode 100755 diffuserve/legacy/layers/attention/ops/__init__.py diff --git a/README.md b/README.md index 621d2470..cf5f6bb1 100755 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@

-# D2fEngine +# Diffulex: Flexible Diffusion LLM Inference Engine vLLM implementation for Diffusion LLMs, D2F is integrated as the core inference strategy, while also support training-free strategies like Fast-dLLM. diff --git a/diffulex/__init__.py b/diffulex/__init__.py new file mode 100755 index 00000000..c71384e5 --- /dev/null +++ b/diffulex/__init__.py @@ -0,0 +1,2 @@ +from diffulex.legacy.llm import LLM +from diffulex.legacy.sampling_params import SamplingParams diff --git a/diffuserve/config.py b/diffulex/config.py similarity index 99% rename from diffuserve/config.py rename to diffulex/config.py index e5533921..664b4cdc 100755 --- a/diffuserve/config.py +++ b/diffulex/config.py @@ -1,4 +1,5 @@ import os + from dataclasses import dataclass from transformers import AutoConfig diff --git a/diffuserve/legacy/engine/block_manager.py b/diffulex/engine/block_manager.py similarity index 50% rename from diffuserve/legacy/engine/block_manager.py rename to diffulex/engine/block_manager.py index cfaf13cd..6007d9de 100755 --- a/diffuserve/legacy/engine/block_manager.py +++ b/diffulex/engine/block_manager.py @@ -5,10 +5,10 @@ from collections import deque from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import List, Dict, Deque, Set +from typing import Callable, Dict, Deque, Iterable, List, Set -from diffuserve.legacy.config import Config -from diffuserve.legacy.engine.sequence import SequenceBase, SequenceForCausalLM, SequenceForDiffusionLM +from diffulex.config import Config +from diffulex.engine.sequence import SequenceBase @dataclass @@ -28,9 +28,16 @@ def reset(self): self.token_ids = [] +BlockManagerFactory = Callable[[Config], "BlockManagerBase"] +_NOT_PROVIDED = object() + + class BlockManagerBase(ABC): - def __init__(self, num_blocks: int, block_size: int): + def __init__(self, config: Config): + num_blocks = config.num_kvcache_blocks + block_size = config.kvcache_block_size assert num_blocks > 0 + self.config = config self.block_size = block_size self.blocks: List[Block] = [Block(block_id=i) for i in range(num_blocks)] self.hash_to_block_id: Dict[int, int] = dict() @@ -101,74 +108,76 @@ def can_append(self, seq: SequenceBase) -> bool: pass @abstractmethod - def may_append(self, seq: SequenceBase): + def may_append(self, seq: SequenceBase) -> None: pass - - -class BlockManagerForCausalLM(BlockManagerBase): - def can_append(self, seq: SequenceForCausalLM) -> bool: - return len(self.free_block_ids) >= (len(seq) % self.block_size == 1) - - def may_append(self, seq: SequenceBase): - block_table = seq.block_table - last_block = self.blocks[block_table[-1]] - if len(seq) % self.block_size == 1: - assert last_block.hash != -1 - block_id = self.free_block_ids[0] - self._allocate_block(block_id) - block_table.append(block_id) - elif len(seq) % self.block_size == 0: - assert last_block.hash == -1 - token_ids = seq.block(seq.num_blocks-1) - prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 - h = self.compute_hash(token_ids, prefix) - last_block.update(h, token_ids) - self.hash_to_block_id[h] = last_block.block_id - else: - assert last_block.hash == -1 - -class BlockManagerForDiffusionLM(BlockManagerBase): - def can_append(self, seq: SequenceForDiffusionLM) -> bool: - return len(self.free_block_ids) >= (seq.cached_or_caching_num_tokens % self.block_size == 1) - - def may_append(self, seq: SequenceForDiffusionLM): - # Handle edge case when no tokens are cached yet - if seq.cached_or_caching_num_tokens == 0: - return - - block_table = seq.block_table - if not block_table: - return - - last_block = self.blocks[block_table[-1]] - - if seq.cached_or_caching_num_tokens // self.block_size == len(seq.block_table): - if last_block.hash == -1: - prev_block_end_token = seq.cached_or_caching_num_tokens - seq.caching_num_tokens - 1 # 256th token (0-indexed: 255) - prev_block_idx = prev_block_end_token // self.block_size # block containing 255th token - - if prev_block_idx < seq.num_blocks: - # This block should be full, so set its hash - token_ids = seq.block(prev_block_idx) - prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 - h = self.compute_hash(token_ids, prefix) - last_block.update(h, token_ids) - self.hash_to_block_id[h] = last_block.block_id - - # Now allocate a new block - block_id = self.free_block_ids[0] - self._allocate_block(block_id) - block_table.append(block_id) - - -class AutoBlockManager(BlockManagerBase): - BLOCK_MANAGER_MAPPING = { - "causal_lm": BlockManagerForCausalLM, - "diffusion_lm": BlockManagerForDiffusionLM, - } + + +class AutoBlockManager: + """Registry-driven factory for block manager implementations.""" + + _BLOCK_MANAGER_MAPPING: Dict[str, BlockManagerFactory] = {} + _DEFAULT_KEY = "__default__" + + @classmethod + def register( + cls, + strategy_name: str, + factory: BlockManagerFactory | object = _NOT_PROVIDED, + *, + aliases: Iterable[str] = (), + is_default: bool = False, + exist_ok: bool = False, + ): + if not isinstance(strategy_name, str) or not strategy_name: + raise ValueError("strategy_name must be a non-empty string.") + if isinstance(aliases, str): + raise TypeError("aliases must be an iterable of strings, not a single string.") + + def decorator(factory_fn: BlockManagerFactory): + cls._register(strategy_name, factory_fn, exist_ok=exist_ok) + for alias in dict.fromkeys(aliases): + if not isinstance(alias, str) or not alias: + raise ValueError("aliases must contain non-empty strings.") + cls._register(alias, factory_fn, exist_ok=exist_ok) + if is_default: + cls._register(cls._DEFAULT_KEY, factory_fn, exist_ok=True) + return factory_fn + + if factory is _NOT_PROVIDED: + return decorator + return decorator(factory) + + @classmethod + def _register(cls, key: str, factory: BlockManagerFactory, *, exist_ok: bool) -> None: + if not exist_ok and key in cls._BLOCK_MANAGER_MAPPING and cls._BLOCK_MANAGER_MAPPING[key] is not factory: + raise ValueError(f"Block manager '{key}' is already registered.") + cls._BLOCK_MANAGER_MAPPING[key] = factory + + @classmethod + def unregister(cls, strategy_name: str) -> None: + cls._BLOCK_MANAGER_MAPPING.pop(strategy_name, None) + + @classmethod + def available_block_managers(cls) -> tuple[str, ...]: + return tuple(sorted(k for k in cls._BLOCK_MANAGER_MAPPING if k != cls._DEFAULT_KEY)) + @classmethod def from_config(cls, config: Config) -> BlockManagerBase: - block_manager_cls = cls.BLOCK_MANAGER_MAPPING.get(config.model_type) - if not block_manager_cls: - raise ValueError(f"Unsupported model type: {config.model_type}") - return block_manager_cls(config.num_kvcache_blocks, config.kvcache_block_size) \ No newline at end of file + candidates: List[str] = [] + for attr in ("decoding_strategy", "model_type"): + value = getattr(config, attr, None) + if isinstance(value, str) and value: + candidates.append(value) + candidates.append(cls._DEFAULT_KEY) + + for key in candidates: + factory = cls._BLOCK_MANAGER_MAPPING.get(key) + if factory is not None: + return factory(config) + + available = ", ".join(cls.available_block_managers()) or "" + raise ValueError( + "No block manager registered for decoding_strategy=" + f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" + f"'{getattr(config, 'model_type', None)}'. Available block managers: {available}." + ) \ No newline at end of file diff --git a/diffuserve/engine/dp_engine.py b/diffulex/engine/dp_engine.py similarity index 98% rename from diffuserve/engine/dp_engine.py rename to diffulex/engine/dp_engine.py index 968ed43b..04a6c905 100755 --- a/diffuserve/engine/dp_engine.py +++ b/diffulex/engine/dp_engine.py @@ -10,9 +10,9 @@ from typing import List, Any from multiprocessing.connection import wait as mp_wait -from diffuserve.config import Config -from diffuserve.engine.llm_engine import LLMEngine -from diffuserve.sampling_params import SamplingParams +from diffulex.config import Config +from diffulex.engine.llm_engine import LLMEngine +from diffulex.sampling_params import SamplingParams def _dp_child_entry(config: Config, dp_idx: int, local_devices: list[int], conn): diff --git a/diffuserve/legacy/engine/llm_engine.py b/diffulex/engine/llm_engine.py similarity index 79% rename from diffuserve/legacy/engine/llm_engine.py rename to diffulex/engine/llm_engine.py index e05be015..23576aa6 100755 --- a/diffuserve/legacy/engine/llm_engine.py +++ b/diffulex/engine/llm_engine.py @@ -8,11 +8,13 @@ from dataclasses import fields from transformers import AutoTokenizer -from diffuserve.legacy.config import Config -from diffuserve.legacy.sampling_params import SamplingParams -from diffuserve.legacy.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM -from diffuserve.legacy.engine.scheduler import AutoScheduler, SchedulerBase -from diffuserve.legacy.engine.model_runner import AutoModelRunner +import diffulex.strategy # noqa: F401 + +from diffulex.config import Config +from diffulex.sampling_params import SamplingParams +from diffulex.engine.sequence import SequenceForDiffusionLM +from diffulex.engine.scheduler import AutoScheduler, SchedulerBase +from diffulex.engine.model_runner import AutoModelRunner class LLMEngine: @@ -60,13 +62,7 @@ def add_request(self, prompt: str | List[int], sampling_params: SamplingParams): if isinstance(prompt, str): prompt = self.tokenizer.encode(prompt) - if self.engine_type == "causal_lm": - seq = SequenceForCausalLM(prompt, sampling_params) - elif self.engine_type == "diffusion_lm": - seq = SequenceForDiffusionLM(prompt, sampling_params, config=self.config) - else: - raise ValueError(f"Unsupported engine type: {self.engine_type}") - + seq = SequenceForDiffusionLM(prompt, sampling_params, config=self.config) seq.block_size = self.config.kvcache_block_size self.scheduler.add(seq) # Return seq_id so caller can build a stable mapping @@ -77,17 +73,9 @@ 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] - if self.engine_type == "causal_lm": - num_tokens = sum(len(seq) for seq in seqs) if is_prefill else len(seqs) - # For streaming: provide per-seq deltas (newly appended token) on decode steps - if not is_prefill: - deltas = [(seq.seq_id, [seq.last_token], seq.is_finished) for seq in seqs] - else: - deltas = [] - else: - 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) - # Diffusion decoding modifies tokens in-place; we currently don't stream intermediate edits - deltas = [] + 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) + # Diffusion decoding modifies tokens in-place; we currently don't stream intermediate edits + deltas = [] return outputs, num_tokens, is_prefill, n_diff_steps, deltas def is_finished(self): @@ -134,6 +122,7 @@ def generate( outputs[seqid_to_idx[seq_id]] = token_ids if use_tqdm: pbar.update(1) + print(f"Finished in {n_steps} steps, prefill throughput: {prefill_throughput:.2f} tok/s, decode throughput: {decode_throughput:.2f} tok/s") # Ensure all outputs are present assert all(toks is not None for toks in outputs), "Some sequences did not produce outputs" diff --git a/diffulex/engine/model_runner.py b/diffulex/engine/model_runner.py new file mode 100755 index 00000000..fbdb48c7 --- /dev/null +++ b/diffulex/engine/model_runner.py @@ -0,0 +1,250 @@ +import torch +import pickle + +import torch.distributed as dist + +from typing import Callable, Dict, Iterable, List +from abc import ABC, abstractmethod +from multiprocessing.synchronize import Event +from multiprocessing.shared_memory import SharedMemory + +from diffulex.config import Config +from diffulex.engine.sequence import SequenceBase +from diffulex.model.auto_model import AutoModelForDiffusionLM +from diffulex.layer.sampler import AutoSampler + +RunnerFactory = Callable[[Config, int, Event | List[Event]], "ModelRunnerBase"] +_NOT_PROVIDED = object() + + +class ModelRunnerBase(ABC): + """Base class for model runners supporting different model types.""" + def __init__(self, config: Config, rank: int, event: Event | List[Event]): + self.config = config + self.model_type = config.model_type + hf_config = config.hf_config + self.block_size = config.kvcache_block_size + self.enforce_eager = config.enforce_eager + self.world_size = config.tensor_parallel_size + self.rank = rank + self.event = event + + # Initialize model, sampler, and kv cache + init_method = f"tcp://{config.master_addr}:{config.master_port}" + dist.init_process_group("nccl", init_method, world_size=self.world_size, rank=rank) + device_id = (getattr(config, "device_start", 0) or 0) + rank + assert 0 <= device_id < torch.cuda.device_count(), f"Invalid device_id {device_id}." + torch.cuda.set_device(device_id) + default_dtype = torch.get_default_dtype() + default_dtype = (hf_config.torch_dtype if hasattr(hf_config, "torch_dtype") + and hf_config.torch_dtype else torch.bfloat16) + torch.set_default_dtype(default_dtype) + torch.set_default_device(f"cuda:{device_id}") + self.model = self.load_model(config) + self.sampler = self.load_sampler(config) + self.warmup_model() + self.allocate_kv_cache() # NOCHANGE + if not self.enforce_eager: + self.capture_cudagraph() + + # Allocate shared memory for inter-process communication + # NOCHANGE + torch.set_default_device("cpu") + torch.set_default_dtype(default_dtype) + if self.world_size > 1: + if rank == 0: + try: + shm = SharedMemory(name=config.shm_name) + shm.close() + shm.unlink() + except FileNotFoundError: + pass + shm_size = 2**25 if self.model_type == "diffusion_lm" else 2**20 + self.shm = SharedMemory(name=config.shm_name, create=True, size=shm_size) + dist.barrier() + else: + dist.barrier() + self.shm = SharedMemory(name=config.shm_name) + self.loop() + + def exit(self): + if self.world_size > 1: + self.shm.close() + dist.barrier() + if self.rank == 0: + self.shm.unlink() + if not self.enforce_eager: + del self.graphs, self.graph_pool + torch.cuda.synchronize() + dist.destroy_process_group() + + def loop(self): + while True: + method_name, args = self.read_shm() + self.call(method_name, *args) + if method_name == "exit": + break + + def read_shm(self): + assert self.world_size > 1 and self.rank + self.event.wait() + n = int.from_bytes(self.shm.buf[0:4], "little") + method_name, *args = pickle.loads(self.shm.buf[4:n+4]) + self.event.clear() + return method_name, args + + def write_shm(self, method_name, *args): + assert self.world_size > 1 and not self.rank + data = pickle.dumps([method_name, *args]) + n = len(data) + + if n + 4 > len(self.shm.buf): + raise ValueError(f"Serialized data size ({n} bytes) exceeds shared memory buffer size ({len(self.shm.buf)} bytes). " + f"Consider increasing shared memory size or reducing batch size.") + + self.shm.buf[0:4] = n.to_bytes(4, "little") + self.shm.buf[4:n+4] = data + for event in self.event: + event.set() + + def call(self, method_name, *args): + if self.world_size > 1 and self.rank == 0: + self.write_shm(method_name, *args) + method = getattr(self, method_name, None) + return method(*args) + + def load_model(self, config: Config): + """Instantiate the underlying model; override to customize.""" + return AutoModelForDiffusionLM.from_config(config) + + def load_sampler(self, config: Config): + """Instantiate the sampler implementation; override to customize.""" + return AutoSampler.from_config(config) + + @abstractmethod + def warmup_model(self): + """Model-specific warmup logic.""" + pass + + @abstractmethod + def allocate_kv_cache(self): + pass + + def prepare_block_tables(self, seqs: List[SequenceBase]): + max_len = max(len(seq.block_table) for seq in seqs) + block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs] + block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + return block_tables + + @abstractmethod + def prepare_prefill(self, seqs: List[SequenceBase]): + """Model-specific prefill preparation.""" + pass + + @abstractmethod + def prepare_decode(self, seqs: List[SequenceBase]): + """Model-specific decode preparation.""" + pass + + def prepare_sample(self, seqs: List[SequenceBase]): + temperatures = [] + for seq in seqs: + temperatures.append(seq.temperature) + temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True) + return temperatures + + @abstractmethod + @torch.inference_mode() + def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): + """Model-specific forward pass.""" + pass + + @abstractmethod + def run(self, seqs: List[SequenceBase], is_prefill: bool) -> List[int]: + """Main inference pipeline.""" + pass + + @abstractmethod + @torch.inference_mode() + def capture_cudagraph(self): + """Model-specific CUDA graph capture.""" + pass + + +class AutoModelRunner: + """Registry and factory that selects a ModelRunner implementation based on the configured decoding strategy. + + Example: + >>> @AutoModelRunner.register("my_strategy") + ... class MyRunner(ModelRunnerBase): + ... ... + + This allows `LLMEngine` to instantiate the appropriate runner using `Config.decoding_strategy`. + """ + + _RUNNER_MAPPING: Dict[str, RunnerFactory] = {} + _DEFAULT_KEY = "__default__" + + @classmethod + def register( + cls, + strategy_name: str, + runner: RunnerFactory | object = _NOT_PROVIDED, + *, + aliases: Iterable[str] = (), + is_default: bool = False, + exist_ok: bool = False, + ): + if not isinstance(strategy_name, str) or not strategy_name: + raise ValueError("strategy_name must be a non-empty string.") + if isinstance(aliases, str): + raise TypeError("aliases must be an iterable of strings, not a single string.") + + def decorator(factory: RunnerFactory): + cls._register(strategy_name, factory, exist_ok=exist_ok) + for alias in dict.fromkeys(aliases): + if not isinstance(alias, str) or not alias: + raise ValueError("aliases must contain non-empty strings.") + cls._register(alias, factory, exist_ok=exist_ok) + if is_default: + cls._register(cls._DEFAULT_KEY, factory, exist_ok=True) + return factory + + if runner is _NOT_PROVIDED: + return decorator + return decorator(runner) + + @classmethod + def _register(cls, key: str, factory: RunnerFactory, *, exist_ok: bool) -> None: + if not exist_ok and key in cls._RUNNER_MAPPING and cls._RUNNER_MAPPING[key] is not factory: + raise ValueError(f"Model runner '{key}' is already registered.") + cls._RUNNER_MAPPING[key] = factory + + @classmethod + def unregister(cls, strategy_name: str) -> None: + cls._RUNNER_MAPPING.pop(strategy_name, None) + + @classmethod + def available_runners(cls) -> tuple[str, ...]: + return tuple(sorted(k for k in cls._RUNNER_MAPPING if k != cls._DEFAULT_KEY)) + + @classmethod + def from_config(cls, config: Config, rank: int, event: Event | List[Event]): + candidates: List[str] = [] + for attr in ("decoding_strategy", "model_type"): + value = getattr(config, attr, None) + if isinstance(value, str) and value: + candidates.append(value) + candidates.append(cls._DEFAULT_KEY) + + for key in candidates: + factory = cls._RUNNER_MAPPING.get(key) + if factory is not None: + return factory(config, rank, event) + + available = ", ".join(cls.available_runners()) or "" + raise ValueError( + "No model runner registered for decoding_strategy=" + f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" + f"'{getattr(config, 'model_type', None)}'. Available runners: {available}." + ) diff --git a/diffulex/engine/scheduler.py b/diffulex/engine/scheduler.py new file mode 100755 index 00000000..393949f8 --- /dev/null +++ b/diffulex/engine/scheduler.py @@ -0,0 +1,113 @@ +from collections import deque +from abc import ABC, abstractmethod +from typing import Callable, Deque, Dict, Iterable, List, Tuple + +from diffulex.config import Config +from diffulex.engine.sequence import SequenceBase +from diffulex.engine.block_manager import AutoBlockManager + + +class SchedulerBase(ABC): + def __init__(self, config: Config): + self.config = config + self.max_num_seqs = config.max_num_seqs + self.max_num_batched_tokens = config.max_num_batched_tokens + self.eos = config.eos + self.block_manager = AutoBlockManager.from_config(config) + self.waiting: Deque[SequenceBase] = deque() + self.running: Deque[SequenceBase] = deque() + + @abstractmethod + def is_finished(self) -> bool: + pass + + @abstractmethod + def add(self, seq: SequenceBase) -> None: + pass + + @abstractmethod + def schedule(self) -> Tuple[List[SequenceBase], bool]: + pass + + @abstractmethod + def preempt(self, seq: SequenceBase) -> None: + pass + + @abstractmethod + def postprocess(self, seqs: List[SequenceBase], sampler_output): + pass + + +SchedulerFactory = Callable[[Config], "SchedulerBase"] +_NOT_PROVIDED = object() + + +class AutoScheduler: + """Registry-driven factory for scheduler implementations.""" + + _SCHEDULER_MAPPING: Dict[str, SchedulerFactory] = {} + _DEFAULT_KEY = "__default__" + + @classmethod + def register( + cls, + strategy_name: str, + factory: SchedulerFactory | object = _NOT_PROVIDED, + *, + aliases: Iterable[str] = (), + is_default: bool = False, + exist_ok: bool = False, + ): + if not isinstance(strategy_name, str) or not strategy_name: + raise ValueError("strategy_name must be a non-empty string.") + if isinstance(aliases, str): + raise TypeError("aliases must be an iterable of strings, not a single string.") + + def decorator(factory_fn: SchedulerFactory): + cls._register(strategy_name, factory_fn, exist_ok=exist_ok) + for alias in dict.fromkeys(aliases): + if not isinstance(alias, str) or not alias: + raise ValueError("aliases must contain non-empty strings.") + cls._register(alias, factory_fn, exist_ok=exist_ok) + if is_default: + cls._register(cls._DEFAULT_KEY, factory_fn, exist_ok=True) + return factory_fn + + if factory is _NOT_PROVIDED: + return decorator + return decorator(factory) + + @classmethod + def _register(cls, key: str, factory: SchedulerFactory, *, exist_ok: bool) -> None: + if not exist_ok and key in cls._SCHEDULER_MAPPING and cls._SCHEDULER_MAPPING[key] is not factory: + raise ValueError(f"Scheduler '{key}' is already registered.") + cls._SCHEDULER_MAPPING[key] = factory + + @classmethod + def unregister(cls, strategy_name: str) -> None: + cls._SCHEDULER_MAPPING.pop(strategy_name, None) + + @classmethod + def available_schedulers(cls) -> tuple[str, ...]: + return tuple(sorted(k for k in cls._SCHEDULER_MAPPING if k != cls._DEFAULT_KEY)) + + @classmethod + def from_config(cls, config: Config) -> SchedulerBase: + candidates: List[str] = [] + for attr in ("decoding_strategy", "model_type"): + value = getattr(config, attr, None) + if isinstance(value, str) and value: + candidates.append(value) + candidates.append(cls._DEFAULT_KEY) + + for key in candidates: + factory = cls._SCHEDULER_MAPPING.get(key) + if factory is not None: + return factory(config) + + available = ", ".join(cls.available_schedulers()) or "" + raise ValueError( + "No scheduler registered for decoding_strategy=" + f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" + f"'{getattr(config, 'model_type', None)}'. Available schedulers: {available}." + ) \ No newline at end of file diff --git a/diffuserve/legacy/engine/sequence.py b/diffulex/engine/sequence.py similarity index 92% rename from diffuserve/legacy/engine/sequence.py rename to diffulex/engine/sequence.py index e273cb44..653fea44 100755 --- a/diffuserve/legacy/engine/sequence.py +++ b/diffulex/engine/sequence.py @@ -6,8 +6,8 @@ from dataclasses import dataclass from typing import List, Tuple, Any -from diffuserve.legacy.config import Config -from diffuserve.legacy.sampling_params import SamplingParams +from diffulex.config import Config +from diffulex.sampling_params import SamplingParams class SequenceStatus(Enum): @@ -66,42 +66,6 @@ def append_token(self, token_id: int) -> None: self.num_tokens += 1 -class SequenceForCausalLM(SequenceBase): - """Standard sequence implementation for Causal Language Models.""" - - def __init__(self, token_ids: List[int], sampling_params = SamplingParams()): - super().__init__(token_ids, sampling_params) - - def __repr__(self) -> str: - return (f"SequenceForCausalLM(block_size={self.block_size}, counter={self.counter}, " - f"seq_id={self.seq_id}, status={self.status.name}, num_tokens={self.num_tokens}, " - f"num_prompt_tokens={self.num_prompt_tokens}, num_cached_tokens={self.num_cached_tokens}, " - f"temperature={self.temperature}, max_tokens={self.max_tokens}, ignore_eos={self.ignore_eos})") - - def __getstate__(self) -> Tuple[int, int, int, List[int], int]: - return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table, - self.token_ids if self.num_completion_tokens == 0 else self.last_token) - - def __setstate__(self, state: Tuple[int, int, int, List[int], int]) -> None: - self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table = state[:-1] - if self.num_completion_tokens == 0: - self.token_ids = state[-1] - else: - self.last_token = state[-1] - - @property - def num_completion_tokens(self) -> int: - return self.num_tokens - self.num_prompt_tokens - - @property - def completion_token_ids(self) -> List[int]: - return self.token_ids[self.num_prompt_tokens:] - - @property - def num_cached_blocks(self) -> int: - return (self.num_cached_tokens + self.block_size - 1) // self.block_size - - class DiffusionBlockStatus(Enum): ACTIVE = auto() TO_CACHE = auto() diff --git a/diffuserve/layer/activation.py b/diffulex/layer/activation.py similarity index 100% rename from diffuserve/layer/activation.py rename to diffulex/layer/activation.py diff --git a/diffuserve/layer/attention/attention_v5.py b/diffulex/layer/attention/attention_v5.py similarity index 97% rename from diffuserve/layer/attention/attention_v5.py rename to diffulex/layer/attention/attention_v5.py index 60ea5017..3e34f12e 100644 --- a/diffuserve/layer/attention/attention_v5.py +++ b/diffulex/layer/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 diffuserve.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 diffuserve.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/layer/attention/ops/__init__.py b/diffulex/layer/attention/ops/__init__.py new file mode 100755 index 00000000..579ccbfe --- /dev/null +++ b/diffulex/layer/attention/ops/__init__.py @@ -0,0 +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 ( + store_kvcache_distinct_layout, store_kvcache_unified_layout, load_kvcache, + CHECK_STORING, CHECK_LOADING +) \ No newline at end of file diff --git a/diffuserve/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py b/diffulex/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py similarity index 99% rename from diffuserve/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py rename to diffulex/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py index 8cc41a72..aed7e060 100755 --- a/diffuserve/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py +++ b/diffulex/layer/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 diffuserve.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/diffuserve/legacy/layers/attention/ops/kv_cache_kernels.py b/diffulex/layer/attention/ops/kv_cache_kernels.py similarity index 99% rename from diffuserve/legacy/layers/attention/ops/kv_cache_kernels.py rename to diffulex/layer/attention/ops/kv_cache_kernels.py index 6f7d70c4..a62e2757 100755 --- a/diffuserve/legacy/layers/attention/ops/kv_cache_kernels.py +++ b/diffulex/layer/attention/ops/kv_cache_kernels.py @@ -6,8 +6,8 @@ from typing import Tuple from einops import rearrange -from diffuserve.legacy.utils.context import ContextForDiffusionLM -from diffuserve.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/diffuserve/layer/attention/ops/prefix_prefill.py b/diffulex/layer/attention/ops/prefix_prefill.py similarity index 100% rename from diffuserve/layer/attention/ops/prefix_prefill.py rename to diffulex/layer/attention/ops/prefix_prefill.py diff --git a/diffuserve/layer/attention/ops/tilus_decode_attn_dlm.py b/diffulex/layer/attention/ops/tilus_decode_attn_dlm.py similarity index 100% rename from diffuserve/layer/attention/ops/tilus_decode_attn_dlm.py rename to diffulex/layer/attention/ops/tilus_decode_attn_dlm.py diff --git a/diffuserve/layer/attention/ops/triton_decode_attn_clm.py b/diffulex/layer/attention/ops/triton_decode_attn_clm.py similarity index 100% rename from diffuserve/layer/attention/ops/triton_decode_attn_clm.py rename to diffulex/layer/attention/ops/triton_decode_attn_clm.py diff --git a/diffuserve/legacy/layers/attention/ops/triton_decode_attn_dlm.py b/diffulex/layer/attention/ops/triton_decode_attn_dlm.py similarity index 97% rename from diffuserve/legacy/layers/attention/ops/triton_decode_attn_dlm.py rename to diffulex/layer/attention/ops/triton_decode_attn_dlm.py index 8db75c0e..e39ed1e0 100755 --- a/diffuserve/legacy/layers/attention/ops/triton_decode_attn_dlm.py +++ b/diffulex/layer/attention/ops/triton_decode_attn_dlm.py @@ -12,7 +12,7 @@ import triton.language as tl -from diffuserve.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 diffuserve.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/diffuserve/layer/attention/ops/triton_flash_attention.py b/diffulex/layer/attention/ops/triton_flash_attention.py similarity index 100% rename from diffuserve/layer/attention/ops/triton_flash_attention.py rename to diffulex/layer/attention/ops/triton_flash_attention.py diff --git a/diffuserve/layer/embed_head.py b/diffulex/layer/embed_head.py similarity index 96% rename from diffuserve/layer/embed_head.py rename to diffulex/layer/embed_head.py index 3b85e447..f46cb694 100755 --- a/diffuserve/layer/embed_head.py +++ b/diffulex/layer/embed_head.py @@ -4,7 +4,7 @@ import torch.nn.functional as F import torch.distributed as dist -from diffuserve.utils.context import get_context_causal_lm, get_context_diffusion_lm +from diffulex.utils.context import get_context_causal_lm, get_context_diffusion_lm class VocabParallelEmbedding(nn.Module): diff --git a/diffuserve/layer/layernorm.py b/diffulex/layer/layernorm.py similarity index 100% rename from diffuserve/layer/layernorm.py rename to diffulex/layer/layernorm.py diff --git a/diffuserve/layer/linear.py b/diffulex/layer/linear.py similarity index 100% rename from diffuserve/layer/linear.py rename to diffulex/layer/linear.py diff --git a/diffuserve/layer/rotary_embedding.py b/diffulex/layer/rotary_embedding.py similarity index 100% rename from diffuserve/layer/rotary_embedding.py rename to diffulex/layer/rotary_embedding.py diff --git a/diffuserve/layer/sampler.py b/diffulex/layer/sampler.py similarity index 98% rename from diffuserve/layer/sampler.py rename to diffulex/layer/sampler.py index 28dd284a..5d9f0c7a 100644 --- a/diffuserve/layer/sampler.py +++ b/diffulex/layer/sampler.py @@ -8,8 +8,8 @@ from dataclasses import dataclass from easydict import EasyDict as edict -from diffuserve.config import Config -from diffuserve.utils.context import get_context_diffusion_lm +from diffulex.config import Config +from diffulex.utils.context import get_context_diffusion_lm class SamplerForDiffusionLM(nn.Module): diff --git a/diffulex/legacy/__init__.py b/diffulex/legacy/__init__.py new file mode 100755 index 00000000..c71384e5 --- /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/diffuserve/legacy/config.py b/diffulex/legacy/config.py similarity index 100% rename from diffuserve/legacy/config.py rename to diffulex/legacy/config.py diff --git a/diffuserve/engine/block_manager.py b/diffulex/legacy/engine/block_manager.py similarity index 97% rename from diffuserve/engine/block_manager.py rename to diffulex/legacy/engine/block_manager.py index d822947b..7f12ce9b 100755 --- a/diffuserve/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 diffuserve.config import Config -from diffuserve.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/diffuserve/legacy/engine/dp_engine.py b/diffulex/legacy/engine/dp_engine.py similarity index 98% rename from diffuserve/legacy/engine/dp_engine.py rename to diffulex/legacy/engine/dp_engine.py index 9de2cb13..8fe47821 100755 --- a/diffuserve/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 diffuserve.legacy.config import Config -from diffuserve.legacy.engine.llm_engine import LLMEngine -from diffuserve.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/diffuserve/engine/llm_engine.py b/diffulex/legacy/engine/llm_engine.py similarity index 94% rename from diffuserve/engine/llm_engine.py rename to diffulex/legacy/engine/llm_engine.py index 37daba57..580bff08 100755 --- a/diffuserve/engine/llm_engine.py +++ b/diffulex/legacy/engine/llm_engine.py @@ -8,11 +8,11 @@ from dataclasses import fields from transformers import AutoTokenizer -from diffuserve.config import Config -from diffuserve.sampling_params import SamplingParams -from diffuserve.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM -from diffuserve.engine.scheduler import AutoScheduler, SchedulerBase -from diffuserve.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/diffuserve/legacy/engine/model_runner.py b/diffulex/legacy/engine/model_runner.py similarity index 98% rename from diffuserve/legacy/engine/model_runner.py rename to diffulex/legacy/engine/model_runner.py index 1667c132..4a881aee 100755 --- a/diffuserve/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 diffuserve.legacy.config import Config -from diffuserve.legacy.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM, SequenceBase -from diffuserve.legacy.models.auto_model import AutoModelLM -from diffuserve.legacy.layers.sampler import AutoSampler -from diffuserve.legacy.utils.checker import CHECK_SLOT_MAPPING -from diffuserve.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/diffuserve/engine/scheduler.py b/diffulex/legacy/engine/scheduler.py similarity index 97% rename from diffuserve/engine/scheduler.py rename to diffulex/legacy/engine/scheduler.py index 0018ec07..e1718bdf 100755 --- a/diffuserve/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 diffuserve.config import Config -from diffuserve.engine.sequence import ( +from diffulex.legacy.config import Config +from diffulex.legacy.engine.sequence import ( SequenceBase, SequenceStatus, SequenceForDiffusionLM, SequenceForCausalLM ) -from diffuserve.layer.sampler import SampleOutputForDiffusionLM -from diffuserve.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/diffuserve/engine/sequence.py b/diffulex/legacy/engine/sequence.py similarity index 99% rename from diffuserve/engine/sequence.py rename to diffulex/legacy/engine/sequence.py index cccf58e7..4f32c55c 100755 --- a/diffuserve/engine/sequence.py +++ b/diffulex/legacy/engine/sequence.py @@ -6,8 +6,8 @@ from dataclasses import dataclass from typing import List, Tuple, Any -from diffuserve.config import Config -from diffuserve.sampling_params import SamplingParams +from diffulex.legacy.config import Config +from diffulex.legacy.sampling_params import SamplingParams class SequenceStatus(Enum): diff --git a/diffuserve/legacy/layers/activation.py b/diffulex/legacy/layers/activation.py similarity index 100% rename from diffuserve/legacy/layers/activation.py rename to diffulex/legacy/layers/activation.py diff --git a/diffuserve/legacy/layers/attention/attention_v1.py b/diffulex/legacy/layers/attention/attention_v1.py similarity index 99% rename from diffuserve/legacy/layers/attention/attention_v1.py rename to diffulex/legacy/layers/attention/attention_v1.py index a6ed634e..4ca79c21 100755 --- a/diffuserve/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 diffuserve.legacy.utils.context import ( +from diffulex.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffuserve/legacy/layers/attention/attention_v1_profile.py b/diffulex/legacy/layers/attention/attention_v1_profile.py similarity index 99% rename from diffuserve/legacy/layers/attention/attention_v1_profile.py rename to diffulex/legacy/layers/attention/attention_v1_profile.py index 877a9bdc..6bb44fee 100755 --- a/diffuserve/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 diffuserve.legacy.utils.context import ( +from diffulex.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffuserve/legacy/layers/attention/attention_v2.py b/diffulex/legacy/layers/attention/attention_v2.py similarity index 99% rename from diffuserve/legacy/layers/attention/attention_v2.py rename to diffulex/legacy/layers/attention/attention_v2.py index 4f271882..5238fd13 100755 --- a/diffuserve/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 diffuserve.legacy.utils.context import ( +from diffulex.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffuserve/legacy/layers/attention/attention_v2_dup.py b/diffulex/legacy/layers/attention/attention_v2_dup.py similarity index 98% rename from diffuserve/legacy/layers/attention/attention_v2_dup.py rename to diffulex/legacy/layers/attention/attention_v2_dup.py index a8be77f2..43ce5e90 100755 --- a/diffuserve/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 diffuserve.legacy.engine.sequence import SequenceForDiffusionLM -from diffuserve.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/diffuserve/legacy/layers/attention/attention_v2_profile.py b/diffulex/legacy/layers/attention/attention_v2_profile.py similarity index 99% rename from diffuserve/legacy/layers/attention/attention_v2_profile.py rename to diffulex/legacy/layers/attention/attention_v2_profile.py index e3b1b7cc..2a98b209 100755 --- a/diffuserve/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 diffuserve.legacy.utils.context import ( +from diffulex.legacy.utils.context import ( ContextForCausalLM, ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm ) diff --git a/diffuserve/legacy/layers/attention/attention_v3.py b/diffulex/legacy/layers/attention/attention_v3.py similarity index 98% rename from diffuserve/legacy/layers/attention/attention_v3.py rename to diffulex/legacy/layers/attention/attention_v3.py index 7dee9682..cfd02bcc 100755 --- a/diffuserve/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 diffuserve.legacy.layers.attention.ops import causal_lm_flash_decoding -from diffuserve.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/diffuserve/legacy/layers/attention/attention_v4.py b/diffulex/legacy/layers/attention/attention_v4.py similarity index 97% rename from diffuserve/legacy/layers/attention/attention_v4.py rename to diffulex/legacy/layers/attention/attention_v4.py index 3fdccc44..e846fd82 100755 --- a/diffuserve/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 diffuserve.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 diffuserve.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/diffuserve/legacy/layers/attention/attention_v5.py b/diffulex/legacy/layers/attention/attention_v5.py similarity index 97% rename from diffuserve/legacy/layers/attention/attention_v5.py rename to diffulex/legacy/layers/attention/attention_v5.py index 4e3540cd..e019bca4 100644 --- a/diffuserve/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 diffuserve.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 diffuserve.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/ops/__init__.py b/diffulex/legacy/layers/attention/ops/__init__.py new file mode 100755 index 00000000..579ccbfe --- /dev/null +++ b/diffulex/legacy/layers/attention/ops/__init__.py @@ -0,0 +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 ( + store_kvcache_distinct_layout, store_kvcache_unified_layout, load_kvcache, + CHECK_STORING, CHECK_LOADING +) \ No newline at end of file diff --git a/diffuserve/legacy/layers/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 diffuserve/legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py rename to diffulex/legacy/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py index 8cc41a72..aed7e060 100755 --- a/diffuserve/legacy/layers/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 diffuserve.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/diffuserve/layer/attention/ops/kv_cache_kernels.py b/diffulex/legacy/layers/attention/ops/kv_cache_kernels.py similarity index 99% rename from diffuserve/layer/attention/ops/kv_cache_kernels.py rename to diffulex/legacy/layers/attention/ops/kv_cache_kernels.py index 6f7d70c4..a62e2757 100755 --- a/diffuserve/layer/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 diffuserve.legacy.utils.context import ContextForDiffusionLM -from diffuserve.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/diffuserve/legacy/layers/attention/ops/prefix_prefill.py b/diffulex/legacy/layers/attention/ops/prefix_prefill.py similarity index 100% rename from diffuserve/legacy/layers/attention/ops/prefix_prefill.py rename to diffulex/legacy/layers/attention/ops/prefix_prefill.py diff --git a/diffuserve/legacy/layers/attention/ops/tilus_decode_attn_dlm.py b/diffulex/legacy/layers/attention/ops/tilus_decode_attn_dlm.py similarity index 100% rename from diffuserve/legacy/layers/attention/ops/tilus_decode_attn_dlm.py rename to diffulex/legacy/layers/attention/ops/tilus_decode_attn_dlm.py diff --git a/diffuserve/legacy/layers/attention/ops/triton_decode_attn_clm.py b/diffulex/legacy/layers/attention/ops/triton_decode_attn_clm.py similarity index 100% rename from diffuserve/legacy/layers/attention/ops/triton_decode_attn_clm.py rename to diffulex/legacy/layers/attention/ops/triton_decode_attn_clm.py diff --git a/diffuserve/layer/attention/ops/triton_decode_attn_dlm.py b/diffulex/legacy/layers/attention/ops/triton_decode_attn_dlm.py similarity index 97% rename from diffuserve/layer/attention/ops/triton_decode_attn_dlm.py rename to diffulex/legacy/layers/attention/ops/triton_decode_attn_dlm.py index 8db75c0e..e39ed1e0 100755 --- a/diffuserve/layer/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 diffuserve.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 diffuserve.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/diffuserve/legacy/layers/attention/ops/triton_flash_attention.py b/diffulex/legacy/layers/attention/ops/triton_flash_attention.py similarity index 100% rename from diffuserve/legacy/layers/attention/ops/triton_flash_attention.py rename to diffulex/legacy/layers/attention/ops/triton_flash_attention.py diff --git a/diffuserve/legacy/layers/embed_head.py b/diffulex/legacy/layers/embed_head.py similarity index 96% rename from diffuserve/legacy/layers/embed_head.py rename to diffulex/legacy/layers/embed_head.py index e4fd553f..b781b2d1 100755 --- a/diffuserve/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 diffuserve.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/diffuserve/legacy/layers/layernorm.py b/diffulex/legacy/layers/layernorm.py similarity index 100% rename from diffuserve/legacy/layers/layernorm.py rename to diffulex/legacy/layers/layernorm.py diff --git a/diffuserve/legacy/layers/linear.py b/diffulex/legacy/layers/linear.py similarity index 100% rename from diffuserve/legacy/layers/linear.py rename to diffulex/legacy/layers/linear.py diff --git a/diffuserve/legacy/layers/rotary_embedding.py b/diffulex/legacy/layers/rotary_embedding.py similarity index 100% rename from diffuserve/legacy/layers/rotary_embedding.py rename to diffulex/legacy/layers/rotary_embedding.py diff --git a/diffuserve/legacy/layers/sampler.py b/diffulex/legacy/layers/sampler.py similarity index 98% rename from diffuserve/legacy/layers/sampler.py rename to diffulex/legacy/layers/sampler.py index f8babe03..fe7bb758 100644 --- a/diffuserve/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 diffuserve.legacy.config import Config -from diffuserve.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): diff --git a/diffuserve/legacy/llm.py b/diffulex/legacy/llm.py similarity index 64% rename from diffuserve/legacy/llm.py rename to diffulex/legacy/llm.py index ce3ecea6..c519d7a1 100755 --- a/diffuserve/legacy/llm.py +++ b/diffulex/legacy/llm.py @@ -1,6 +1,6 @@ -from diffuserve.legacy.engine.llm_engine import LLMEngine -from diffuserve.legacy.engine.dp_engine import DPEngine -from diffuserve.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/diffuserve/legacy/models/auto_model.py b/diffulex/legacy/models/auto_model.py similarity index 58% rename from diffuserve/legacy/models/auto_model.py rename to diffulex/legacy/models/auto_model.py index d9edee28..185aa0b0 100755 --- a/diffuserve/legacy/models/auto_model.py +++ b/diffulex/legacy/models/auto_model.py @@ -1,8 +1,8 @@ -from diffuserve.legacy.config import Config -from diffuserve.legacy.utils.loader import load_model -from diffuserve.legacy.models.dream import DreamForDiffusionLM -from diffuserve.legacy.models.qwen3 import Qwen3ForCausalLM -from diffuserve.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/diffuserve/legacy/models/config/dream/configuration_dream.py b/diffulex/legacy/models/config/dream/configuration_dream.py similarity index 100% rename from diffuserve/legacy/models/config/dream/configuration_dream.py rename to diffulex/legacy/models/config/dream/configuration_dream.py diff --git a/diffuserve/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 diffuserve/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/diffuserve/legacy/models/config/llada/configuration_llada.py b/diffulex/legacy/models/config/llada/configuration_llada.py similarity index 100% rename from diffuserve/legacy/models/config/llada/configuration_llada.py rename to diffulex/legacy/models/config/llada/configuration_llada.py diff --git a/diffuserve/legacy/models/dream.py b/diffulex/legacy/models/dream.py similarity index 93% rename from diffuserve/legacy/models/dream.py rename to diffulex/legacy/models/dream.py index 9ac26007..4f5bd36a 100755 --- a/diffuserve/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 diffuserve.legacy.layers.layernorm import RMSNorm -from diffuserve.legacy.layers.activation import SiluAndMul -from diffuserve.legacy.layers.rotary_embedding import get_rope -from diffuserve.legacy.layers.attention.attention_v5 import Attention -from diffuserve.legacy.models.config.dream.configuration_dream import DreamConfig -from diffuserve.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear -from diffuserve.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/diffuserve/legacy/models/fast_dllm_v2.py b/diffulex/legacy/models/fast_dllm_v2.py similarity index 93% rename from diffuserve/legacy/models/fast_dllm_v2.py rename to diffulex/legacy/models/fast_dllm_v2.py index 40a2e024..4739b412 100755 --- a/diffuserve/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 diffuserve.legacy.layers.layernorm import RMSNorm -from diffuserve.legacy.layers.activation import SiluAndMul -from diffuserve.legacy.layers.rotary_embedding import get_rope -from diffuserve.legacy.layers.attention.attention_v5 import Attention -from diffuserve.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear -from diffuserve.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead -from diffuserve.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/diffuserve/legacy/models/llada.py b/diffulex/legacy/models/llada.py similarity index 94% rename from diffuserve/legacy/models/llada.py rename to diffulex/legacy/models/llada.py index e35ebfbb..342a1c57 100755 --- a/diffuserve/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 diffuserve.legacy.layers.layernorm import RMSNorm -from diffuserve.legacy.layers.activation import SiluAndMul -from diffuserve.legacy.layers.rotary_embedding import get_rope -from diffuserve.legacy.layers.attention.attention_v5 import Attention -from diffuserve.legacy.models.config.llada.configuration_llada import LLaDAConfig -from diffuserve.legacy.layers.linear import RowParallelLinear, ColumnParallelLinear -from diffuserve.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/diffuserve/legacy/models/qwen3.py b/diffulex/legacy/models/qwen3.py similarity index 93% rename from diffuserve/legacy/models/qwen3.py rename to diffulex/legacy/models/qwen3.py index 76906123..f6803d9c 100755 --- a/diffuserve/legacy/models/qwen3.py +++ b/diffulex/legacy/models/qwen3.py @@ -4,12 +4,12 @@ from transformers import Qwen3Config -from diffuserve.legacy.layers.layernorm import RMSNorm -from diffuserve.legacy.layers.activation import SiluAndMul -from diffuserve.legacy.layers.rotary_embedding import get_rope -from diffuserve.legacy.layers.attention.attention_v4 import Attention -from diffuserve.legacy.layers.embed_head import VocabParallelEmbedding, ParallelLMHead -from diffuserve.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/diffuserve/legacy/models/utils/check_config.py b/diffulex/legacy/models/utils/check_config.py similarity index 100% rename from diffuserve/legacy/models/utils/check_config.py rename to diffulex/legacy/models/utils/check_config.py diff --git a/diffuserve/legacy/sampling_params.py b/diffulex/legacy/sampling_params.py similarity index 100% rename from diffuserve/legacy/sampling_params.py rename to diffulex/legacy/sampling_params.py diff --git a/diffuserve/legacy/utils/checker.py b/diffulex/legacy/utils/checker.py similarity index 100% rename from diffuserve/legacy/utils/checker.py rename to diffulex/legacy/utils/checker.py diff --git a/diffuserve/legacy/utils/context.py b/diffulex/legacy/utils/context.py similarity index 98% rename from diffuserve/legacy/utils/context.py rename to diffulex/legacy/utils/context.py index c6e24f1f..7d49ea38 100755 --- a/diffuserve/legacy/utils/context.py +++ b/diffulex/legacy/utils/context.py @@ -3,7 +3,7 @@ from typing import List from dataclasses import dataclass -from diffuserve.legacy.engine.sequence import SequenceForDiffusionLM +from diffulex.legacy.engine.sequence import SequenceForDiffusionLM @dataclass class ContextBase: diff --git a/diffuserve/utils/loader.py b/diffulex/legacy/utils/loader.py similarity index 99% rename from diffuserve/utils/loader.py rename to diffulex/legacy/utils/loader.py index c16a86bf..5dd07bd3 100755 --- a/diffuserve/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 diffuserve.legacy.config import Config +from diffulex.legacy.config import Config def load_lora_config(lora_path: str) -> dict: diff --git a/diffuserve/llm.py b/diffulex/llm.py similarity index 67% rename from diffuserve/llm.py rename to diffulex/llm.py index 9a86e4ce..7cd69bea 100755 --- a/diffuserve/llm.py +++ b/diffulex/llm.py @@ -1,6 +1,6 @@ -from diffuserve.engine.llm_engine import LLMEngine -from diffuserve.engine.dp_engine import DPEngine -from diffuserve.config import Config +from diffulex.config import Config +from diffulex.engine.dp_engine import DPEngine +from diffulex.engine.llm_engine import LLMEngine class LLM: def __new__(cls, model, **kwargs): diff --git a/diffuserve/model/auto_model.py b/diffulex/model/auto_model.py similarity index 93% rename from diffuserve/model/auto_model.py rename to diffulex/model/auto_model.py index f39c5cf1..65e9feef 100755 --- a/diffuserve/model/auto_model.py +++ b/diffulex/model/auto_model.py @@ -2,10 +2,10 @@ from typing import Any, Callable -from diffuserve.legacy.config import Config -from diffuserve.legacy.utils.loader import load_model -from diffuserve.legacy.models.dream import DreamForDiffusionLM -from diffuserve.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.llada import LLaDAForDiffusionLM _NOT_PROVIDED = object() RegistryEntry = tuple[Callable[[Any], Any] | type | None, bool] diff --git a/diffuserve/model/config/dream/configuration_dream.py b/diffulex/model/config/dream/configuration_dream.py similarity index 100% rename from diffuserve/model/config/dream/configuration_dream.py rename to diffulex/model/config/dream/configuration_dream.py diff --git a/diffuserve/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py b/diffulex/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py similarity index 100% rename from diffuserve/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py rename to diffulex/model/config/fast_dllm_v2/configuration_fast_dllm_v2.py diff --git a/diffuserve/model/config/llada/configuration_llada.py b/diffulex/model/config/llada/configuration_llada.py similarity index 100% rename from diffuserve/model/config/llada/configuration_llada.py rename to diffulex/model/config/llada/configuration_llada.py diff --git a/diffuserve/model/diffucoder.py b/diffulex/model/diffucoder.py similarity index 100% rename from diffuserve/model/diffucoder.py rename to diffulex/model/diffucoder.py diff --git a/diffuserve/model/dream.py b/diffulex/model/dream.py similarity index 93% rename from diffuserve/model/dream.py rename to diffulex/model/dream.py index bb2d7a12..8b75a883 100755 --- a/diffuserve/model/dream.py +++ b/diffulex/model/dream.py @@ -3,14 +3,14 @@ import torch.nn as nn import torch.distributed as dist -from diffuserve.layer.layernorm import RMSNorm -from diffuserve.layer.activation import SiluAndMul -from diffuserve.layer.rotary_embedding import get_rope -from diffuserve.layer.attention.attention_v5 import Attention -from diffuserve.model.auto_model import AutoModelForDiffusionLM -from diffuserve.model.config.dream.configuration_dream import DreamConfig -from diffuserve.layer.linear import RowParallelLinear, ColumnParallelLinear -from diffuserve.layer.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffulex.layer.layernorm import RMSNorm +from diffulex.layer.activation import SiluAndMul +from diffulex.layer.rotary_embedding import get_rope +from diffulex.layer.attention.attention_v5 import Attention +from diffulex.model.auto_model import AutoModelForDiffusionLM +from diffulex.model.config.dream.configuration_dream import DreamConfig +from diffulex.layer.linear import RowParallelLinear, ColumnParallelLinear +from diffulex.layer.embed_head import VocabParallelEmbedding, ParallelLMHead if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/diffuserve/model/fast_dllm_v2.py b/diffulex/model/fast_dllm_v2.py similarity index 93% rename from diffuserve/model/fast_dllm_v2.py rename to diffulex/model/fast_dllm_v2.py index 4fcc4740..ba3724b2 100755 --- a/diffuserve/model/fast_dllm_v2.py +++ b/diffulex/model/fast_dllm_v2.py @@ -3,14 +3,14 @@ import torch.nn as nn import torch.distributed as dist -from diffuserve.layer.layernorm import RMSNorm -from diffuserve.layer.activation import SiluAndMul -from diffuserve.layer.rotary_embedding import get_rope -from diffuserve.layer.attention.attention_v5 import Attention -from diffuserve.model.auto_model import AutoModelForDiffusionLM -from diffuserve.layer.linear import RowParallelLinear, ColumnParallelLinear -from diffuserve.layer.embed_head import VocabParallelEmbedding, ParallelLMHead -from diffuserve.model.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config +from diffulex.layer.layernorm import RMSNorm +from diffulex.layer.activation import SiluAndMul +from diffulex.layer.rotary_embedding import get_rope +from diffulex.layer.attention.attention_v5 import Attention +from diffulex.model.auto_model import AutoModelForDiffusionLM +from diffulex.layer.linear import RowParallelLinear, ColumnParallelLinear +from diffulex.layer.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffulex.model.config.fast_dllm_v2.configuration_fast_dllm_v2 import FastdLLMV2Config if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/diffuserve/model/llada.py b/diffulex/model/llada.py similarity index 94% rename from diffuserve/model/llada.py rename to diffulex/model/llada.py index 84a16703..b1c3e485 100755 --- a/diffuserve/model/llada.py +++ b/diffulex/model/llada.py @@ -3,14 +3,14 @@ import torch.nn as nn import torch.distributed as dist -from diffuserve.layer.layernorm import RMSNorm -from diffuserve.layer.activation import SiluAndMul -from diffuserve.layer.rotary_embedding import get_rope -from diffuserve.layer.attention.attention_v5 import Attention -from diffuserve.model.auto_model import AutoModelForDiffusionLM -from diffuserve.model.config.llada.configuration_llada import LLaDAConfig -from diffuserve.layer.linear import RowParallelLinear, ColumnParallelLinear -from diffuserve.layer.embed_head import VocabParallelEmbedding, ParallelLMHead +from diffulex.layer.layernorm import RMSNorm +from diffulex.layer.activation import SiluAndMul +from diffulex.layer.rotary_embedding import get_rope +from diffulex.layer.attention.attention_v5 import Attention +from diffulex.model.auto_model import AutoModelForDiffusionLM +from diffulex.model.config.llada.configuration_llada import LLaDAConfig +from diffulex.layer.linear import RowParallelLinear, ColumnParallelLinear +from diffulex.layer.embed_head import VocabParallelEmbedding, ParallelLMHead if os.environ.get("TRITON_INTERPRET", None) == "1": diff --git a/diffuserve/model/llada2.py b/diffulex/model/llada2.py similarity index 100% rename from diffuserve/model/llada2.py rename to diffulex/model/llada2.py diff --git a/diffuserve/model/llada_moe.py b/diffulex/model/llada_moe.py similarity index 100% rename from diffuserve/model/llada_moe.py rename to diffulex/model/llada_moe.py diff --git a/diffuserve/model/sdar.py b/diffulex/model/sdar.py similarity index 100% rename from diffuserve/model/sdar.py rename to diffulex/model/sdar.py diff --git a/diffuserve/model/utils/check_config.py b/diffulex/model/utils/check_config.py similarity index 100% rename from diffuserve/model/utils/check_config.py rename to diffulex/model/utils/check_config.py diff --git a/diffuserve/sampling_params.py b/diffulex/sampling_params.py similarity index 100% rename from diffuserve/sampling_params.py rename to diffulex/sampling_params.py diff --git a/diffulex/strategy/__init__.py b/diffulex/strategy/__init__.py new file mode 100644 index 00000000..d0b96e09 --- /dev/null +++ b/diffulex/strategy/__init__.py @@ -0,0 +1,7 @@ +"""Diffulex strategy package that imports built-in strategies to trigger registration.""" +from __future__ import annotations + +# Import built-in strategies so their registrations run at import time. +from . import d2f # noqa: F401 + +__all__ = ["d2f"] diff --git a/diffulex/strategy/d2f/__init__.py b/diffulex/strategy/d2f/__init__.py new file mode 100644 index 00000000..d25e1991 --- /dev/null +++ b/diffulex/strategy/d2f/__init__.py @@ -0,0 +1,12 @@ +"""D2F strategy component exports.""" +from __future__ import annotations + +from .block_manager import D2FBlockManager +from .model_runner import D2FModelRunner +from .scheduler import D2FScheduler + +__all__ = [ + "D2FBlockManager", + "D2FModelRunner", + "D2FScheduler", +] diff --git a/diffulex/strategy/d2f/block_manager.py b/diffulex/strategy/d2f/block_manager.py new file mode 100644 index 00000000..d348b453 --- /dev/null +++ b/diffulex/strategy/d2f/block_manager.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from typing import List + +from diffulex.config import Config +from diffulex.engine.block_manager import AutoBlockManager, BlockManagerBase +from diffulex.engine.sequence import SequenceForDiffusionLM + + +@AutoBlockManager.register( + "d2f", + aliases=("diffusion_lm",), + is_default=True, +) +class D2FBlockManager(BlockManagerBase): + def __init__(self, config: Config): + super().__init__(config) + + def can_append(self, seq: SequenceForDiffusionLM) -> bool: + required = 1 if seq.cached_or_caching_num_tokens % self.block_size == 1 else 0 + return len(self.free_block_ids) >= required + + def may_append(self, seq: SequenceForDiffusionLM) -> None: + if seq.cached_or_caching_num_tokens == 0: + return + block_table = seq.block_table + if not block_table: + return + last_block = self.blocks[block_table[-1]] + if seq.cached_or_caching_num_tokens // self.block_size == len(seq.block_table): + if last_block.hash == -1: + prev_end_token = seq.cached_or_caching_num_tokens - seq.caching_num_tokens - 1 + prev_block_idx = prev_end_token // self.block_size + if prev_block_idx < seq.num_blocks: + token_ids: List[int] = seq.block(prev_block_idx) + prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 + h = self.compute_hash(token_ids, prefix) + last_block.update(h, token_ids) + 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) diff --git a/diffulex/strategy/d2f/model_runner.py b/diffulex/strategy/d2f/model_runner.py new file mode 100644 index 00000000..d0b932c6 --- /dev/null +++ b/diffulex/strategy/d2f/model_runner.py @@ -0,0 +1,421 @@ +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.model_runner import AutoModelRunner, ModelRunnerBase +from diffulex.engine.sequence import SequenceForDiffusionLM, SequenceBase +from diffulex.utils.context import ( + get_context_diffusion_lm, + reset_context_diffusion_lm, + set_context_diffusion_lm, +) + + +@AutoModelRunner.register( + "d2f", + aliases=("diffusion_lm",), + 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) + self.diffusion_block_size = config.diffusion_block_size + self.mask_token_id = config.mask_token_id + self.decoding_strategy = config.decoding_strategy + + def warmup_model(self): + print("Warming up model...") + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + max_num_batched_tokens, max_model_len = ( + self.config.max_num_batched_tokens, + self.config.max_model_len, + ) + num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs) + test_input_ids = [0] * max_model_len + seqs = [SequenceForDiffusionLM(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() + + 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[SequenceForDiffusionLM]): + input_ids: List[int] = [] + positions: List[int] = [] + cu_seqlens_q = [0] + cu_seqlens_k = [0] + max_seqlen_q = 0 + max_seqlen_k = 0 + 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) + + 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 + cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) + cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k) + + max_seqlen_q = max(seqlen_q, max_seqlen_q) + max_seqlen_k = max(seqlen_k, max_seqlen_k) + + if not seq.block_table: + continue + for i in range(0, seq.num_prompt_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 + else: + end = start + seq.last_block_prompt_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_context_diffusion_lm( + True, + cu_seqlens_q=cu_seqlens_q_tensor, + cu_seqlens_k=cu_seqlens_k_tensor, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + slot_mapping=slot_mapping_tensor, + context_lens=context_lens_tensor, + block_tables=block_tables, + seqs=seqs, + kv_cache_layout=self.config.kv_cache_layout, + seq_lens=seq_lens, + seq_lens_ts=seq_lens_ts, + ) + return input_ids_tensor, positions_tensor + + def prepare_decode(self, seqs: List[SequenceForDiffusionLM]): + 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 + 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 + 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), + ) + ) + + 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_context_diffusion_lm( + False, + slot_mapping=slot_mapping_tensor, + context_lens=context_lens_tensor, + cu_seqlens_q=cu_seqlens_q_tensor, + cu_seqlens_k=cu_seqlens_k_tensor, + 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, + kv_cache_layout=self.config.kv_cache_layout, + need_kv_cache_store=need_kv_cache_store, + d2f_pp=True, + ) + return input_ids_tensor, positions_tensor + + @torch.inference_mode() + 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 = get_context_diffusion_lm() + graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] + 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 + graph.replay() + return self.model.compute_logits(graph_vars["outputs"][:bs]) + + 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_context_diffusion_lm() + 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 + logits = self.run_model(input_ids, positions, is_prefill) + sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None + reset_context_diffusion_lm() + 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.") diff --git a/diffulex/strategy/d2f/scheduler.py b/diffulex/strategy/d2f/scheduler.py new file mode 100644 index 00000000..7be134da --- /dev/null +++ b/diffulex/strategy/d2f/scheduler.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from typing import Dict, List, Tuple + +from diffulex.config import Config +from diffulex.engine.scheduler import AutoScheduler, SchedulerBase +from diffulex.engine.sequence import ( + SequenceBase, + SequenceForDiffusionLM, + SequenceStatus, +) +from diffulex.layer.sampler import SampleOutputForDiffusionLM + + +@AutoScheduler.register( + "d2f", + aliases=("diffusion_lm",), + is_default=True, +) +class D2FScheduler(SchedulerBase): + def __init__(self, config: Config): + super().__init__(config) + self.diffusion_block_size = config.diffusion_block_size + + def is_finished(self) -> bool: + return not self.waiting and not self.running + + def add(self, seq: SequenceForDiffusionLM) -> None: + self.waiting.append(seq) + + def schedule(self) -> Tuple[List[SequenceBase], bool]: + scheduled: List[SequenceBase] = [] + num_seqs = 0 + num_batched_tokens = 0 + while self.waiting and num_seqs < self.max_num_seqs: + seq = self.waiting[0] + projected = len(seq) + seq.diffusion_block_size + if ( + num_batched_tokens + projected > self.max_num_batched_tokens + or not self.block_manager.can_allocate(seq) + ): + break + num_seqs += 1 + self.block_manager.allocate(seq) + num_batched_tokens += projected - seq.num_cached_tokens + seq.status = SequenceStatus.RUNNING + self.waiting.popleft() + self.running.append(seq) + scheduled.append(seq) + if scheduled: + return scheduled, True + + while self.running and num_seqs < self.max_num_seqs: + seq = self.running.popleft() + while not self.block_manager.can_append(seq): + if self.running: + self.preempt(self.running.pop()) + else: + self.preempt(seq) + break + else: + num_seqs += 1 + self.block_manager.may_append(seq) + scheduled.append(seq) + if not scheduled: + diag = { + "phase": "decode", + "waiting": len(self.waiting), + "running": len(self.running), + "max_num_seqs": self.max_num_seqs, + "max_num_batched_tokens": self.max_num_batched_tokens, + "diffusion_block_size": self.diffusion_block_size, + } + candidates = list(self.running)[:3] + list(self.waiting)[:2] + details = [] + for idx, candidate in enumerate(candidates): + try: + can_append = self.block_manager.can_append(candidate) + except Exception: + can_append = "error" + details.append( + f"[{idx}] status={candidate.status.name}, len={len(candidate)}, " + f"diff_block={getattr(candidate, 'diffusion_block_size', '?')}, " + f"new_tokens={getattr(candidate, 'new_tokens', '?')}, " + f"cached={getattr(candidate, 'num_cached_tokens', '?')}, " + f"can_append={can_append}" + ) + raise RuntimeError( + "D2FScheduler: 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: SequenceForDiffusionLM) -> None: + seq.status = SequenceStatus.WAITING + self.block_manager.free(seq) + self.waiting.appendleft(seq) + + def postprocess( + self, + seqs: List[SequenceForDiffusionLM], + sample_output: SampleOutputForDiffusionLM, + ) -> Dict[int, int]: + n_diff_steps: Dict[int, int] = {} + for seq in seqs: + seq.reset_new_tokens() + seq_id = str(seq.seq_id) + true_ids_map = sample_output.true_local_ids_map.get(seq_id, {}) + accepted_ids_map = sample_output.accepted_ids_map.get(seq_id, {}) + sampled_tokens_map = sample_output.sampled_tokens_map.get(seq_id, {}) + for block_id, accepted_ids in accepted_ids_map.items(): + if not accepted_ids: + continue + diffusion_block = seq.diffusion_blocks[int(block_id)] + sampled_tokens = sampled_tokens_map.get(block_id, []) + true_local_ids = true_ids_map.get(block_id, []) + for true_local_id, accepted_id in zip(true_local_ids, accepted_ids): + token = sampled_tokens[accepted_id] + diffusion_block.modify_token(true_local_id, token) + if ( + (not seq.ignore_eos and token.item() == self.eos) + or seq.num_completion_tokens >= seq.max_tokens + ): + seq.meet_eos = True + if seq.meet_eos and seq.diffusion_blocks[-1].available_to_cache: + seq.status = SequenceStatus.FINISHED + self.block_manager.free(seq) + if seq in self.running: + self.running.remove(seq) + n_diff_steps[seq.seq_id] = seq.n_steps + seq.post_process() + return n_diff_steps diff --git a/diffuserve/utils/checker.py b/diffulex/utils/checker.py similarity index 100% rename from diffuserve/utils/checker.py rename to diffulex/utils/checker.py diff --git a/diffuserve/utils/context.py b/diffulex/utils/context.py similarity index 98% rename from diffuserve/utils/context.py rename to diffulex/utils/context.py index a5ed2575..f2261de2 100755 --- a/diffuserve/utils/context.py +++ b/diffulex/utils/context.py @@ -2,7 +2,7 @@ from dataclasses import dataclass -from diffuserve.legacy.engine.sequence import SequenceForDiffusionLM +from diffulex.legacy.engine.sequence import SequenceForDiffusionLM @dataclass class ContextBase: diff --git a/diffuserve/legacy/utils/loader.py b/diffulex/utils/loader.py similarity index 99% rename from diffuserve/legacy/utils/loader.py rename to diffulex/utils/loader.py index c16a86bf..5dd07bd3 100755 --- a/diffuserve/legacy/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 diffuserve.legacy.config import Config +from diffulex.legacy.config import Config def load_lora_config(lora_path: str) -> dict: diff --git a/diffuserve/__init__.py b/diffuserve/__init__.py deleted file mode 100755 index 82b9e51b..00000000 --- a/diffuserve/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from diffuserve.legacy.llm import LLM -from diffuserve.legacy.sampling_params import SamplingParams diff --git a/diffuserve/engine/model_runner.py b/diffuserve/engine/model_runner.py deleted file mode 100755 index 6cbe6c84..00000000 --- a/diffuserve/engine/model_runner.py +++ /dev/null @@ -1,507 +0,0 @@ -import time -import torch -import pickle - -import torch.distributed as dist - -from typing import List -from abc import ABC, abstractmethod -from multiprocessing.synchronize import Event -from multiprocessing.shared_memory import SharedMemory - -from diffuserve.config import Config -from diffuserve.engine.sequence import SequenceForCausalLM, SequenceForDiffusionLM, SequenceBase -from diffuserve.model.auto_model import AutoModelForDiffusionLM -from diffuserve.layer.sampler import AutoSampler -from diffuserve.utils.checker import CHECK_SLOT_MAPPING -from diffuserve.utils.context import ( - set_context_causal_lm, - get_context_causal_lm, - reset_context_causal_lm, - set_context_diffusion_lm, - get_context_diffusion_lm, - reset_context_diffusion_lm -) - - -class ModelRunnerBase(ABC): - """Base class for model runners supporting different model types.""" - def __init__(self, config: Config, rank: int, event: Event | List[Event]): - self.config = config - self.model_type = config.model_type - hf_config = config.hf_config - self.block_size = config.kvcache_block_size - self.enforce_eager = config.enforce_eager - self.world_size = config.tensor_parallel_size - self.rank = rank - self.event = event - - # Initialize model, sampler, and kv cache - init_method = f"tcp://{config.master_addr}:{config.master_port}" - dist.init_process_group("nccl", init_method, world_size=self.world_size, rank=rank) - device_id = (getattr(config, "device_start", 0) or 0) + rank - assert 0 <= device_id < torch.cuda.device_count(), f"Invalid device_id {device_id}." - torch.cuda.set_device(device_id) - default_dtype = torch.get_default_dtype() - default_dtype = (hf_config.torch_dtype if hasattr(hf_config, "torch_dtype") - and hf_config.torch_dtype else torch.bfloat16) - torch.set_default_dtype(default_dtype) - torch.set_default_device(f"cuda:{device_id}") - self.model = AutoModelForDiffusionLM.from_config(config) - self.sampler = AutoSampler.from_config(config) - self.warmup_model() - self.allocate_kv_cache() # NOCHANGE - if not self.enforce_eager: - self.capture_cudagraph() - - # Allocate shared memory for inter-process communication - # NOCHANGE - torch.set_default_device("cpu") - torch.set_default_dtype(default_dtype) - if self.world_size > 1: - if rank == 0: - try: - shm = SharedMemory(name=config.shm_name) - shm.close() - shm.unlink() - except FileNotFoundError: - pass - shm_size = 2**25 if self.model_type == "diffusion_lm" else 2**20 - self.shm = SharedMemory(name=config.shm_name, create=True, size=shm_size) - dist.barrier() - else: - dist.barrier() - self.shm = SharedMemory(name=config.shm_name) - self.loop() - - def exit(self): - if self.world_size > 1: - self.shm.close() - dist.barrier() - if self.rank == 0: - self.shm.unlink() - if not self.enforce_eager: - del self.graphs, self.graph_pool - torch.cuda.synchronize() - dist.destroy_process_group() - - def loop(self): - while True: - method_name, args = self.read_shm() - self.call(method_name, *args) - if method_name == "exit": - break - - def read_shm(self): - assert self.world_size > 1 and self.rank - self.event.wait() - n = int.from_bytes(self.shm.buf[0:4], "little") - method_name, *args = pickle.loads(self.shm.buf[4:n+4]) - self.event.clear() - return method_name, args - - def write_shm(self, method_name, *args): - assert self.world_size > 1 and not self.rank - data = pickle.dumps([method_name, *args]) - n = len(data) - - if n + 4 > len(self.shm.buf): - raise ValueError(f"Serialized data size ({n} bytes) exceeds shared memory buffer size ({len(self.shm.buf)} bytes). " - f"Consider increasing shared memory size or reducing batch size.") - - self.shm.buf[0:4] = n.to_bytes(4, "little") - self.shm.buf[4:n+4] = data - for event in self.event: - event.set() - - def call(self, method_name, *args): - if self.world_size > 1 and self.rank == 0: - self.write_shm(method_name, *args) - method = getattr(self, method_name, None) - return method(*args) - - @abstractmethod - def warmup_model(self): - """Model-specific warmup logic.""" - pass - - @abstractmethod - def allocate_kv_cache(self): - pass - - def prepare_block_tables(self, seqs: List[SequenceBase]): - max_len = max(len(seq.block_table) for seq in seqs) - block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs] - block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - return block_tables - - @abstractmethod - def prepare_prefill(self, seqs: List[SequenceBase]): - """Model-specific prefill preparation.""" - pass - - @abstractmethod - def prepare_decode(self, seqs: List[SequenceBase]): - """Model-specific decode preparation.""" - pass - - def prepare_sample(self, seqs: List[SequenceBase]): - temperatures = [] - for seq in seqs: - temperatures.append(seq.temperature) - temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True) - return temperatures - - @abstractmethod - @torch.inference_mode() - def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): - """Model-specific forward pass.""" - pass - - @abstractmethod - def run(self, seqs: List[SequenceBase], is_prefill: bool) -> List[int]: - """Main inference pipeline.""" - pass - - @abstractmethod - @torch.inference_mode() - def capture_cudagraph(self): - """Model-specific CUDA graph capture.""" - pass - - -class ModelRunnerForDiffusionLM(ModelRunnerBase): - """Model runner for Diffusion Language Models. TODO: Implement DLM-specific logic.""" - def __init__(self, config: Config, rank: int, event: Event | List[Event]): - super().__init__(config, rank, event) - self.diffusion_block_size = config.diffusion_block_size - self.mask_token_id = config.mask_token_id - self.decoding_strategy = config.decoding_strategy - - def warmup_model(self): - # return - print("Warming up model...") - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats() - max_num_batched_tokens, max_model_len = self.config.max_num_batched_tokens, self.config.max_model_len - num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs) - test_input_ids = [0] * max_model_len - seqs = [SequenceForDiffusionLM(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() - - 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 - # noqa: E731 - used - peak + current) // block_bytes - try: - num_kvcache_blocks = get_num_kvcache_blocks(config.gpu_memory_utilization) - assert num_kvcache_blocks > 0 - except: # noqa: E722 - gpu_memory_utilization = config.gpu_memory_utilization - while num_kvcache_blocks <= 200: - print(f"Warning: GPU memory utilization {gpu_memory_utilization} is too low to allocate kv cache. " - f"Automatically adding 0.05, which is {gpu_memory_utilization + 0.05:.2f} now.") - 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(f"Allocated {config.num_kvcache_blocks} blocks of size {self.block_size} for kv cache on rank {self.rank}.") - - if config.kv_cache_layout == "distinct": - # k_cache: [layer_id, block_id, head, head_dim // x, block_size(segmented seq_len), x] - # v_cache: [layer_id, block_id, head, head_dim, block_size(segmented seq_len)] - 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": - # [kv_separated, layer_id, block_id, block_size(segmented seq_len), head, head_dim] - 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(f"Unsupported kv_cache_layout: {config.kv_cache_layout}. " - f"Supported values are 'distinct' and 'unified'.") - - def prepare_prefill(self, seqs: List[SequenceForDiffusionLM]): - input_ids = [] - positions = [] - cu_seqlens_q = [0] - cu_seqlens_k = [0] - max_seqlen_q = 0 - max_seqlen_k = 0 - slot_mapping = [] - block_tables = None - context_lens = [] - seq_lens = [] - - for seq in seqs: - seq.next_diffusion_step(is_prefill=True) - - total_seqlen = len(seq) - # tokens and positions to run in this prefill step - input_ids.extend(seq[seq.cached_num_tokens:]) - positions.extend(list(range(seq.cached_num_tokens, total_seqlen))) - seq_lens.append(total_seqlen) - context_lens.append(0) - assert len(input_ids) == len(positions), ( - f"prepare_prefill(diffusion): len(input_ids) {len(input_ids)} != len(positions) {len(positions)}" - ) - - seqlen_q = total_seqlen - seq.cached_num_tokens - seqlen_k = total_seqlen - cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) - cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k) - - max_seqlen_q = max(seqlen_q, max_seqlen_q) - max_seqlen_k = max(seqlen_k, max_seqlen_k) - - if not seq.block_table: - continue - # build slot mapping for prefix cache prompt blocks - for i in range(0, seq.num_prompt_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 - else: - end = start + seq.last_block_prompt_num_tokens - slot_mapping.extend(list(range(start, end))) - else: - slot_mapping.extend([-1] * self.block_size) - # pad to a full diffusion block - slot_mapping.extend([-1] * seq.diffusion_block_size) - - # For diffusion prefill we always need block tables for prefix cache bookkeeping - block_tables = self.prepare_block_tables(seqs) - - input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) - positions = 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 = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - - # More checks to avoid downstream rotary errors - assert cu_seqlens_q[-1].item() == input_ids.numel(), ( - f"prepare_prefill(diffusion): cu_seqlens_q[-1]={cu_seqlens_q[-1].item()} != num_tokens={input_ids.numel()}" - ) - assert cu_seqlens_k[-1].item() == sum(seq_lens), ( - f"prepare_prefill(diffusion): cu_seqlens_k[-1]={cu_seqlens_k[-1].item()} != sum(seq_lens)={sum(seq_lens)}" - ) - - set_context_diffusion_lm( - True, - 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, - seqs=seqs, - kv_cache_layout=self.config.kv_cache_layout, - seq_lens=seq_lens, - seq_lens_ts=seq_lens_ts, - ) - return input_ids, positions - - def prepare_decode(self, seqs: List[SequenceForDiffusionLM]): - input_ids = [] - positions = [] - cu_seqlens_q = [0] - cu_seqlens_k = [0] - slot_mapping = [] - context_lens = [] - seq_lens = [] - seq_id_to_queue_id = {} - need_kv_cache_store = False - max_seqlen_q = 0 - max_seqlen_k = 0 - # if sum((sum(seq.active_blocks) + sum(seq.to_cache_blocks)) * seq.diffusion_block_size for seq in seqs) == 1536: - # pass - 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 - 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 - 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 - get_step = lambda diff_blk, start_idx: ( - diff_blk.remaining_length(start_idx) - if diff_blk.remaining_length(start_idx) + local_start_idx() <= seq.block_size - else 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(list(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: - # Covering all the after-active blocks - 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 - padding_slots = [-1] * (num_blocks_to_pad * seq.diffusion_block_size) - slot_mapping.extend(padding_slots) - break - assert len(input_ids) == len(positions), f"Input IDs length {len(input_ids)} does not match positions length {len(positions)}" - assert len(input_ids) == len(slot_mapping), f"Input IDs length {len(input_ids)} does not match slot mapping length {len(slot_mapping)}" - - # CHECK_SLOT_MAPPING(seqs, slot_mapping) - input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) - positions = 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 = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - block_tables = self.prepare_block_tables(seqs) - set_context_diffusion_lm(False, slot_mapping=slot_mapping, context_lens=context_lens, - cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, - 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, - kv_cache_layout=self.config.kv_cache_layout, need_kv_cache_store=need_kv_cache_store, - d2f_pp=True) - return input_ids, positions - - @torch.inference_mode() - 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)) - else: - bs = input_ids.size(0) - context = get_context_diffusion_lm() - graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] - graph_vars = self.graph_vars - for k, v in graph_vars.items(): - if k != "outputs": - v.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 - graph.replay() - return self.model.compute_logits(graph_vars["outputs"][:bs]) - - 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}") - s = 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() - s:.2f} seconds") - s = time.time() - logits = self.run_model(input_ids, positions, is_prefill) - print(f"Ran model in {time.time() - s:.2f} seconds") - s = time.time() - sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None - print(f"Sampled tokens in {time.time() - s:.2f} seconds") - reset_context_diffusion_lm() - 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 - logits = self.run_model(input_ids, positions, is_prefill) - sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None - reset_context_diffusion_lm() - 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.") - - -class AutoModelRunner: - @classmethod - def from_config(cls, config: Config, rank: int, event: Event | List[Event]): - """Factory method to create a model runner based on the model type.""" - return ModelRunnerForDiffusionLM(config, rank, event) \ No newline at end of file diff --git a/diffuserve/layer/attention/ops/__init__.py b/diffuserve/layer/attention/ops/__init__.py deleted file mode 100755 index 8e202106..00000000 --- a/diffuserve/layer/attention/ops/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from diffuserve.legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd as causal_lm_flash_decoding -from diffuserve.legacy.layers.attention.ops.triton_decode_attn_dlm import diffusion_lm_flash_decoding, CHECK_ATTENTION -from diffuserve.legacy.layers.attention.ops.chunked_prefill_decoding_unified_kernel import chunked_prefill_paged_decode as diffusion_lm_parallel_flash_decoding -from diffuserve.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/diffuserve/legacy/__init__.py b/diffuserve/legacy/__init__.py deleted file mode 100755 index 82b9e51b..00000000 --- a/diffuserve/legacy/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from diffuserve.legacy.llm import LLM -from diffuserve.legacy.sampling_params import SamplingParams diff --git a/diffuserve/legacy/engine/scheduler.py b/diffuserve/legacy/engine/scheduler.py deleted file mode 100755 index 66be8fc0..00000000 --- a/diffuserve/legacy/engine/scheduler.py +++ /dev/null @@ -1,234 +0,0 @@ -import torch - -from collections import deque -from abc import ABC, abstractmethod -from typing import Tuple, List, Deque - -from diffuserve.legacy.config import Config -from diffuserve.legacy.engine.sequence import ( - SequenceBase, SequenceStatus, - SequenceForDiffusionLM, SequenceForCausalLM -) -from diffuserve.legacy.layers.sampler import SampleOutputForDiffusionLM -from diffuserve.legacy.engine.block_manager import AutoBlockManager - - -class SchedulerBase(ABC): - def __init__(self, config: Config): - self.max_num_seqs = config.max_num_seqs - self.max_num_batched_tokens = config.max_num_batched_tokens - self.eos = config.eos - self.block_manager = AutoBlockManager.from_config(config) - self.waiting: Deque[SequenceBase] = deque() - self.running: Deque[SequenceBase] = deque() - - @abstractmethod - def is_finished(self) -> bool: - pass - - @abstractmethod - def add(self, seq: SequenceBase) -> None: - pass - - @abstractmethod - def schedule(self) -> Tuple[List[SequenceBase], bool]: - pass - - @abstractmethod - def preempt(self, seq: SequenceBase) -> None: - pass - - @abstractmethod - def postprocess(self, seqs: List[SequenceBase], token_ids: List[int]): - pass - - -class SchedulerForCausalLM(SchedulerBase): - def __init__(self, config: Config): - super().__init__(config) - - def is_finished(self) -> bool: - return not self.waiting and not self.running - - def add(self, seq: SequenceForCausalLM) -> None: - self.waiting.append(seq) - - def schedule(self) -> Tuple[List[SequenceForCausalLM], bool]: - # prefill - scheduled_seqs = [] - num_seqs = 0 - num_batched_tokens = 0 - while self.waiting and num_seqs < self.max_num_seqs: - seq = self.waiting[0] - if num_batched_tokens + len(seq) > self.max_num_batched_tokens \ - or not self.block_manager.can_allocate(seq): - break - num_seqs += 1 - self.block_manager.allocate(seq) - num_batched_tokens += len(seq) - seq.num_cached_tokens - seq.status = SequenceStatus.RUNNING - self.waiting.popleft() - self.running.append(seq) - scheduled_seqs.append(seq) - if scheduled_seqs: - return scheduled_seqs, True - - # decode - while self.running and num_seqs < self.max_num_seqs: - seq = self.running.popleft() - while not self.block_manager.can_append(seq): - if self.running: - self.preempt(self.running.pop()) - else: - self.preempt(seq) - break - else: - num_seqs += 1 - self.block_manager.may_append(seq) - scheduled_seqs.append(seq) - if not scheduled_seqs: - # Provide diagnostics to understand starvation/resource issues - diag = { - "phase": "decode", - "waiting": len(self.waiting), - "running": len(self.running), - "max_num_seqs": self.max_num_seqs, - "max_num_batched_tokens": self.max_num_batched_tokens, - } - # Probe a few candidates for can_append and lengths - candidates = list(self.running)[:3] + list(self.waiting)[:2] - infos = [] - for j, s in enumerate(candidates): - try: - cap = self.block_manager.can_append(s) - except Exception: - cap = "error" - infos.append( - f"[{j}] status={s.status.name}, len={len(s)}, new_tokens={getattr(s, 'num_completion_tokens', getattr(s, 'new_tokens', '?'))}, cached={getattr(s, 'num_cached_tokens', '?')}, can_append={cap}" - ) - raise RuntimeError(f"SchedulerForCausalLM: unable to schedule any sequence in decode; state={diag}; details={' | '.join(infos)}") - self.running.extendleft(reversed(scheduled_seqs)) - return scheduled_seqs, False - - def preempt(self, seq: SequenceForCausalLM) -> None: - seq.status = SequenceStatus.WAITING - self.block_manager.free(seq) - self.waiting.appendleft(seq) - - def postprocess(self, seqs: List[SequenceForCausalLM], token_ids: List[int]) -> None: - for seq, token_id in zip(seqs, token_ids): - seq.append_token(token_id) - if (not seq.ignore_eos and token_id == self.eos) \ - or seq.num_completion_tokens == seq.max_tokens: - seq.status = SequenceStatus.FINISHED - self.block_manager.free(seq) - self.running.remove(seq) - - -# TODO -class SchedulerForDiffusionLM(SchedulerBase): - def __init__(self, config: Config): - super().__init__(config) - self.diffusion_block_size = config.diffusion_block_size - - def is_finished(self) -> bool: - return not self.waiting and not self.running - - def add(self, seq: SequenceForDiffusionLM) -> None: - self.waiting.append(seq) - - def schedule(self): - # prefill - scheduled_seqs = [] - num_seqs = 0 - num_batched_tokens = 0 - while self.waiting and num_seqs < self.max_num_seqs: - seq = self.waiting[0] - if num_batched_tokens + len(seq) + seq.diffusion_block_size > self.max_num_batched_tokens or not self.block_manager.can_allocate(seq): - break - num_seqs += 1 - self.block_manager.allocate(seq) - num_batched_tokens += len(seq) + seq.diffusion_block_size - seq.num_cached_tokens - seq.status = SequenceStatus.RUNNING - self.waiting.popleft() - self.running.append(seq) - scheduled_seqs.append(seq) - if scheduled_seqs: - return scheduled_seqs, True - - # decode - while self.running and num_seqs < self.max_num_seqs: - seq = self.running.popleft() - while not self.block_manager.can_append(seq): - if self.running: - self.preempt(self.running.pop()) - else: - self.preempt(seq) - break - else: - num_seqs += 1 - self.block_manager.may_append(seq) - scheduled_seqs.append(seq) - if not scheduled_seqs: - diag = { - "phase": "decode", - "waiting": len(self.waiting), - "running": len(self.running), - "max_num_seqs": self.max_num_seqs, - "max_num_batched_tokens": self.max_num_batched_tokens, - "diffusion_block_size": getattr(self, 'diffusion_block_size', None), - } - candidates = list(self.running)[:3] + list(self.waiting)[:2] - infos = [] - for j, s in enumerate(candidates): - try: - cap = self.block_manager.can_append(s) - except Exception: - cap = "error" - infos.append( - f"[{j}] status={s.status.name}, len={len(s)}, diff_block={getattr(s, 'diffusion_block_size', '?')}, new_tokens={getattr(s, 'new_tokens', '?')}, cached={getattr(s, 'num_cached_tokens', '?')}, can_append={cap}" - ) - raise RuntimeError(f"SchedulerForDiffusionLM: unable to schedule any sequence in decode; state={diag}; details={' | '.join(infos)}") - self.running.extendleft(reversed(scheduled_seqs)) - return scheduled_seqs, False - - def preempt(self, seq: SequenceForDiffusionLM) -> None: - seq.status = SequenceStatus.WAITING - self.block_manager.free(seq) - self.waiting.appendleft(seq) - - def postprocess(self, seqs: List[SequenceForDiffusionLM], sample_output: SampleOutputForDiffusionLM) -> None: - n_diff_steps = {} - for seq in seqs: - seq.reset_new_tokens() - seq_id = str(seq.seq_id) - cur_true_local_ids_sub_map = sample_output.true_local_ids_map.get(seq_id, {}) - cur_accepted_ids_sub_map = sample_output.accepted_ids_map.get(seq_id, {}) - cur_sampled_tokens_sub_map = sample_output.sampled_tokens_map.get(seq_id, {}) - for block_id, accepted_ids in cur_accepted_ids_sub_map.items(): - if len(accepted_ids) > 0: - diffusion_block = seq.diffusion_blocks[int(block_id)] - sampled_tokens = cur_sampled_tokens_sub_map.get(block_id, []) - true_local_ids = cur_true_local_ids_sub_map.get(block_id, []) - - for true_local_id, accepted_id in zip(true_local_ids, accepted_ids): - diffusion_block.modify_token(true_local_id, sampled_tokens[accepted_id]) - if ((not seq.ignore_eos and sampled_tokens[accepted_id].item() == self.eos) - or seq.num_completion_tokens >= seq.max_tokens): - seq.meet_eos = True - if seq.meet_eos and seq.diffusion_blocks[-1].available_to_cache: - seq.status = SequenceStatus.FINISHED - self.block_manager.free(seq) - self.running.remove(seq) - n_diff_steps[seq.seq_id] = seq.n_steps - seq.post_process() - return n_diff_steps - -class AutoScheduler: - SCHEDULER_MAPPING = { - "causal_lm": SchedulerForCausalLM, - "diffusion_lm": SchedulerForDiffusionLM, - } - @classmethod - def from_config(cls, config: Config): - return cls.SCHEDULER_MAPPING[config.model_type](config) \ No newline at end of file diff --git a/diffuserve/legacy/layers/attention/ops/__init__.py b/diffuserve/legacy/layers/attention/ops/__init__.py deleted file mode 100755 index 8e202106..00000000 --- a/diffuserve/legacy/layers/attention/ops/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from diffuserve.legacy.layers.attention.ops.triton_decode_attn_clm import causal_lm_decode_attention_fwd as causal_lm_flash_decoding -from diffuserve.legacy.layers.attention.ops.triton_decode_attn_dlm import diffusion_lm_flash_decoding, CHECK_ATTENTION -from diffuserve.legacy.layers.attention.ops.chunked_prefill_decoding_unified_kernel import chunked_prefill_paged_decode as diffusion_lm_parallel_flash_decoding -from diffuserve.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/examples/test_causal_lm_decoding_kernel.py b/examples/test_causal_lm_decoding_kernel.py index aab4ca4b..f5f0e836 100755 --- a/examples/test_causal_lm_decoding_kernel.py +++ b/examples/test_causal_lm_decoding_kernel.py @@ -1,6 +1,6 @@ import torch -from diffuserve.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 a0da171b..c91925dd 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 diffuserve.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 6d378040..6096ba1c 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 diffuserve.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 2d88e5bf..7ee9c1b7 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 diffuserve.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 357a707c..dad28bde 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -9,7 +9,7 @@ from viztracer import VizTracer from transformers import AutoTokenizer -from diffuserve.legacy import LLM, SamplingParams +from diffulex.legacy import LLM, SamplingParams def summarize_profiling(csv_path: str) -> dict: diff --git a/examples/test_dream_dvllm_human_eval.py b/examples/test_dream_dvllm_human_eval.py index fbc89c0c..787b0275 100755 --- a/examples/test_dream_dvllm_human_eval.py +++ b/examples/test_dream_dvllm_human_eval.py @@ -8,7 +8,7 @@ from viztracer import VizTracer from transformers import AutoTokenizer -from diffuserve.legacy import LLM, SamplingParams +from diffulex.legacy import LLM, SamplingParams def summarize_profiling(csv_path: str) -> dict: diff --git a/examples/test_dream_model_weight.py b/examples/test_dream_model_weight.py index aefece5f..5ad412ee 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 diffuserve.legacy.config import Config -from diffuserve.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 da75a3ae..ae2b61b1 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 diffuserve.legacy.config import Config -from diffuserve.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_llada_dvllm_human_eval.py b/examples/test_llada_dvllm_human_eval.py index 502b9cad..5370026e 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 diffuserve.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 fb5740f6..14da0c48 100755 --- a/examples/test_qwen_dvllm.py +++ b/examples/test_qwen_dvllm.py @@ -1,6 +1,6 @@ import os -from diffuserve.legacy import LLM, SamplingParams +from diffulex.legacy import LLM, SamplingParams from viztracer import VizTracer diff --git a/pyproject.toml b/pyproject.toml index 3f9188a8..59288b0c 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [project] -name = "DiffuServe" +name = "Diffulex" version = "0.0.1" authors = [ { name = "Drew Jin (Yijie Jin)", email = "drewjin0827@gmail.com"}, @@ -7,7 +7,7 @@ authors = [ maintainers = [ { name = "DENG Lab @ SJTU" } ] -description = "Diffusion LLM serving engine supporting multiple variants of dLLM decoding strategies with high efficiency and low latency." +description = "Flexible Diffusion LLM serving engine supporting multiple variants of dLLM decoding strategies with high efficiency and low latency." readme = "README.md" requires-python = ">=3.12" license = { file = "LICENSE" } From 3792939476f475601d1be58c224c70d342739ce9 Mon Sep 17 00:00:00 2001 From: drewjin Date: Thu, 20 Nov 2025 15:24:28 +0800 Subject: [PATCH 07/23] refactor: add registry mechanism of diffulex (temporary upload) --- diffulex/attention/__init__.py | 2 + .../attn_impl.py} | 44 +- diffulex/attention/metadata.py | 25 + .../{layer => }/attention/ops/__init__.py | 0 ...chunked_prefill_decoding_unified_kernel.py | 0 .../attention/ops/kv_cache_kernels.py | 4 +- .../attention/ops/prefix_prefill.py | 0 .../attention/ops/tilus_decode_attn_dlm.py | 0 .../attention/ops/triton_decode_attn_clm.py | 0 .../attention/ops/triton_decode_attn_dlm.py | 0 .../attention/ops/triton_flash_attention.py | 0 diffulex/diffulex.py | 10 + .../engine/{dp_engine.py => dp_worker.py} | 16 +- .../{block_manager.py => kvcache_manager.py} | 79 +-- diffulex/engine/model_runner.py | 80 +-- diffulex/engine/scheduler.py | 71 +-- diffulex/engine/sequence.py | 507 ++---------------- diffulex/engine/strategy_registry.py | 56 ++ .../engine/{llm_engine.py => tp_worker.py} | 18 +- diffulex/layer/sampler.py | 15 +- diffulex/llm.py | 10 - .../model/config/llada/configuration_llada.py | 4 +- diffulex/model/dream.py | 4 +- diffulex/model/fast_dllm_v2.py | 4 +- diffulex/model/llada.py | 6 +- diffulex/model/utils/check_config.py | 8 - diffulex/strategy/__init__.py | 13 + diffulex/strategy/d2f/__init__.py | 14 +- diffulex/strategy/d2f/attention/metadata.py | 28 + .../kvcache_manager.py} | 18 +- .../strategy/d2f/{ => engine}/model_runner.py | 56 +- .../strategy/d2f/{ => engine}/scheduler.py | 23 +- diffulex/strategy/d2f/engine/sequence.py | 484 +++++++++++++++++ diffulex/utils/context.py | 112 ---- diffulex/utils/registry.py | 34 ++ examples/eval_llada.py | 4 +- .../model_cache/dream/generation_utils.py | 6 +- .../model_cache/llada/configuration_llada.py | 2 +- examples/model_cache/llada/modeling_llada.py | 4 +- examples/summary.py | 2 +- examples/test_dream_dvllm_gsm8k.py | 2 +- examples/test_dream_dvllm_human_eval.py | 2 +- examples/test_llada_dvllm_human_eval.py | 2 +- 43 files changed, 870 insertions(+), 899 deletions(-) create mode 100644 diffulex/attention/__init__.py rename diffulex/{layer/attention/attention_v5.py => attention/attn_impl.py} (77%) create mode 100644 diffulex/attention/metadata.py rename diffulex/{layer => }/attention/ops/__init__.py (100%) rename diffulex/{layer => }/attention/ops/chunked_prefill_decoding_unified_kernel.py (100%) rename diffulex/{layer => }/attention/ops/kv_cache_kernels.py (99%) rename diffulex/{layer => }/attention/ops/prefix_prefill.py (100%) rename diffulex/{layer => }/attention/ops/tilus_decode_attn_dlm.py (100%) rename diffulex/{layer => }/attention/ops/triton_decode_attn_clm.py (100%) rename diffulex/{layer => }/attention/ops/triton_decode_attn_dlm.py (100%) rename diffulex/{layer => }/attention/ops/triton_flash_attention.py (100%) create mode 100755 diffulex/diffulex.py rename diffulex/engine/{dp_engine.py => dp_worker.py} (95%) rename diffulex/engine/{block_manager.py => kvcache_manager.py} (59%) create mode 100644 diffulex/engine/strategy_registry.py rename diffulex/engine/{llm_engine.py => tp_worker.py} (92%) delete mode 100755 diffulex/llm.py delete mode 100755 diffulex/model/utils/check_config.py create mode 100644 diffulex/strategy/d2f/attention/metadata.py rename diffulex/strategy/d2f/{block_manager.py => engine/kvcache_manager.py} (75%) rename diffulex/strategy/d2f/{ => engine}/model_runner.py (93%) rename diffulex/strategy/d2f/{ => engine}/scheduler.py (91%) create mode 100644 diffulex/strategy/d2f/engine/sequence.py delete mode 100755 diffulex/utils/context.py create mode 100644 diffulex/utils/registry.py diff --git a/diffulex/attention/__init__.py b/diffulex/attention/__init__.py new file mode 100644 index 00000000..e38b5ff8 --- /dev/null +++ b/diffulex/attention/__init__.py @@ -0,0 +1,2 @@ +from .attn_impl import Attention +from .metadata import fetch_attn_metadata, set_fetch_fn_for_attn_metadata, AttnMetaDataBase \ No newline at end of file diff --git a/diffulex/layer/attention/attention_v5.py b/diffulex/attention/attn_impl.py similarity index 77% rename from diffulex/layer/attention/attention_v5.py rename to diffulex/attention/attn_impl.py index 3e34f12e..51ac0713 100644 --- a/diffulex/layer/attention/attention_v5.py +++ b/diffulex/attention/attn_impl.py @@ -3,19 +3,19 @@ import torch.nn as nn -from typing import List +from typing import list 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.legacy.layers.attention.ops import ( +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.legacy.utils.context import ContextForDiffusionLM, get_context_causal_lm, get_context_diffusion_lm +from diffulex.attention import AttnMetaDataBase, fetch_attn_metadata class Attention(nn.Module): @@ -83,30 +83,30 @@ def _mask_mod(batch, head, token_q, token_kv): return self._block_mask_cache[cache_key] def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - mask: List[torch.Tensor] | None = None) -> 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) - context: ContextForDiffusionLM = get_context_causal_lm() if self.model_type == 'causal_lm' else get_context_diffusion_lm() + attn_metadata: AttnMetaDataBase = fetch_attn_metadata() k_cache, v_cache = self.k_cache, self.v_cache - is_unified_layout = context.kv_cache_layout == "unified" + is_unified_layout = attn_metadata.kv_cache_layout == "unified" # Fast Store KV cache if k_cache.numel() and v_cache.numel(): - if not (self.model_type == 'diffusion_lm' and not context.need_kv_cache_store): + if not (self.model_type == 'diffusion_lm' and not 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, context.slot_mapping, self.model_type, context) + store_kvcache(k, v, k_cache, v_cache, attn_metadata.slot_mapping, self.model_type, 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 - if context.is_prefill: + if attn_metadata.is_prefill: # Block PK - if context.block_tables is not None and self.model_type == 'causal_lm': + if attn_metadata.block_tables is not None and self.model_type == 'causal_lm': k, v = k_cache, v_cache - elif context.block_tables is not None and self.model_type == 'diffusion_lm': + elif attn_metadata.block_tables is not None and self.model_type == 'diffusion_lm': # TODO: Implement Prefix Caching pass @@ -115,17 +115,17 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, B, H, S, _ = q_t.shape block_mask_fn = self.causal_lm_block_mask if self.model_type == 'causal_lm' else self.dllm_block_mask - input_obj = context.cu_seqlens_q if self.model_type == 'causal_lm' else context.block_mask + input_obj = attn_metadata.cu_seqlens_q if self.model_type == 'causal_lm' else attn_metadata.block_mask block_mask = block_mask_fn(input_obj, B, H, S, S, str(q.device)) o = self.attention(q_t, k_t, v_t, block_mask=block_mask) else: - config = context.seqs[0].config + 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, context, k, v) + 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, - context.cu_seqlens_q, context.cu_seqlens_k, - context.max_seqlen_q, context.max_seqlen_k, + 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) else: # FIXME: Kernel not ok... @@ -133,16 +133,16 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, 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, - context.block_tables, context.cu_seqlens_q, context.total_lens, - max(context.total_lens), max(context.seq_lens), 1.0, 1.0, - diffusion_block_size, context.block_mask + 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, context) + CHECK_ATTENTION(o, q, k, v, k_cache, v_cache, attn_metadata) # Final reshape - if not context.is_prefill: + if not attn_metadata.is_prefill: o = o.view(-1, self.num_heads * self.head_dim).contiguous() - elif context.is_prefill: + elif attn_metadata.is_prefill: o = rearrange(o, '1 h s d -> s (h d)').contiguous() return o \ No newline at end of file diff --git a/diffulex/attention/metadata.py b/diffulex/attention/metadata.py new file mode 100644 index 00000000..6b157e00 --- /dev/null +++ b/diffulex/attention/metadata.py @@ -0,0 +1,25 @@ +import torch + +from typing import Callable +from dataclasses import dataclass + + +@dataclass +class AttnMetaDataBase: + 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 + + +FN_TYPE_AttnMetaDataFetch = Callable[[], AttnMetaDataBase] + +fetch_attn_metadata: FN_TYPE_AttnMetaDataFetch = ... + +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 diff --git a/diffulex/layer/attention/ops/__init__.py b/diffulex/attention/ops/__init__.py similarity index 100% rename from diffulex/layer/attention/ops/__init__.py rename to diffulex/attention/ops/__init__.py diff --git a/diffulex/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py b/diffulex/attention/ops/chunked_prefill_decoding_unified_kernel.py similarity index 100% rename from diffulex/layer/attention/ops/chunked_prefill_decoding_unified_kernel.py rename to diffulex/attention/ops/chunked_prefill_decoding_unified_kernel.py diff --git a/diffulex/layer/attention/ops/kv_cache_kernels.py b/diffulex/attention/ops/kv_cache_kernels.py similarity index 99% rename from diffulex/layer/attention/ops/kv_cache_kernels.py rename to diffulex/attention/ops/kv_cache_kernels.py index a62e2757..1a6fc5ee 100755 --- a/diffulex/layer/attention/ops/kv_cache_kernels.py +++ b/diffulex/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.utils.context import ContextForDiffusionLM +from diffulex.strategy.d2f.sequence import D2FSequence @triton.jit def store_kvcache_kernel_causal_lm( diff --git a/diffulex/layer/attention/ops/prefix_prefill.py b/diffulex/attention/ops/prefix_prefill.py similarity index 100% rename from diffulex/layer/attention/ops/prefix_prefill.py rename to diffulex/attention/ops/prefix_prefill.py diff --git a/diffulex/layer/attention/ops/tilus_decode_attn_dlm.py b/diffulex/attention/ops/tilus_decode_attn_dlm.py similarity index 100% rename from diffulex/layer/attention/ops/tilus_decode_attn_dlm.py rename to diffulex/attention/ops/tilus_decode_attn_dlm.py diff --git a/diffulex/layer/attention/ops/triton_decode_attn_clm.py b/diffulex/attention/ops/triton_decode_attn_clm.py similarity index 100% rename from diffulex/layer/attention/ops/triton_decode_attn_clm.py rename to diffulex/attention/ops/triton_decode_attn_clm.py diff --git a/diffulex/layer/attention/ops/triton_decode_attn_dlm.py b/diffulex/attention/ops/triton_decode_attn_dlm.py similarity index 100% rename from diffulex/layer/attention/ops/triton_decode_attn_dlm.py rename to diffulex/attention/ops/triton_decode_attn_dlm.py diff --git a/diffulex/layer/attention/ops/triton_flash_attention.py b/diffulex/attention/ops/triton_flash_attention.py similarity index 100% rename from diffulex/layer/attention/ops/triton_flash_attention.py rename to diffulex/attention/ops/triton_flash_attention.py diff --git a/diffulex/diffulex.py b/diffulex/diffulex.py new file mode 100755 index 00000000..08612ba1 --- /dev/null +++ b/diffulex/diffulex.py @@ -0,0 +1,10 @@ +from diffulex.config import Config +from diffulex.engine.dp_worker import DiffulexDPWorker +from diffulex.engine.tp_worker import DiffulexTPWorker + +class Diffulex: + def __new__(cls, model, **kwargs): + cfg = Config(model, **{k: v for k, v in kwargs.items() if k in Config.__dataclass_fields__.keys()}) + if cfg.data_parallel_size > 1: + return DiffulexDPWorker(model, **kwargs) + return DiffulexTPWorker(model, **kwargs) \ No newline at end of file diff --git a/diffulex/engine/dp_engine.py b/diffulex/engine/dp_worker.py similarity index 95% rename from diffulex/engine/dp_engine.py rename to diffulex/engine/dp_worker.py index 04a6c905..b2366c16 100755 --- a/diffulex/engine/dp_engine.py +++ b/diffulex/engine/dp_worker.py @@ -7,11 +7,11 @@ import multiprocessing as mp -from typing import List, Any +from typing import list, Any from multiprocessing.connection import wait as mp_wait from diffulex.config import Config -from diffulex.engine.llm_engine import LLMEngine +from diffulex.engine.tp_worker import DiffulexTPWorker from diffulex.sampling_params import SamplingParams @@ -52,7 +52,7 @@ def _dp_child_entry(config: Config, dp_idx: int, local_devices: list[int], conn) ) setattr(cfg, "device_start", 0) - engine = LLMEngine(cfg.model, **{k: getattr(cfg, k) for k in cfg.__dataclass_fields__.keys() if k != "model"}) + engine = DiffulexTPWorker(cfg.model, **{k: getattr(cfg, k) for k in cfg.__dataclass_fields__.keys() if k != "model"}) while True: msg = conn.recv() @@ -93,7 +93,7 @@ def _dp_child_entry(config: Config, dp_idx: int, local_devices: list[int], conn) pass -class DPEngine: +class DiffulexDPWorker: """Data-parallel wrapper that runs N independent TP groups as child processes and aggregates results.""" def __init__(self, model, **kwargs): config_fields = {f for f in Config.__dataclass_fields__.keys()} @@ -103,8 +103,8 @@ def __init__(self, model, **kwargs): assert self.dp_size > 1, "Use LLMEngine directly when data_parallel_size == 1" ctx = mp.get_context("spawn") - self.conns: List[Any] = [] - self.ps: List[mp.Process] = [] + self.conns: list[Any] = [] + self.ps: list[mp.Process] = [] # Topology check and mapping base_visible = os.environ.get("CUDA_VISIBLE_DEVICES") if base_visible: @@ -169,7 +169,7 @@ def exit(self): pass p.join(timeout=5) - def add_request(self, prompt: str | List[int], sampling_params: SamplingParams): + def add_request(self, prompt: str | list[int], sampling_params: SamplingParams): target = self._rr self._rr = (self._rr + 1) % self.dp_size local_id = self._ask(target, "add_request", prompt, sampling_params) @@ -210,7 +210,7 @@ def step(self): def is_finished(self): return all(self._ask(i, "is_finished") for i in range(self.dp_size)) - def generate(self, prompts: List[str] | List[List[int]], sampling_params: SamplingParams | List[SamplingParams], use_tqdm: bool = True): + def generate(self, prompts: list[str] | list[list[int]], sampling_params: SamplingParams | list[SamplingParams], use_tqdm: bool = True): """Load-balanced generate with random shuffling and stable order restoration. - Randomly shuffle inputs to balance load across DP replicas. - Partition shuffled list evenly among replicas. diff --git a/diffulex/engine/block_manager.py b/diffulex/engine/kvcache_manager.py similarity index 59% rename from diffulex/engine/block_manager.py rename to diffulex/engine/kvcache_manager.py index 6007d9de..3f7e7661 100755 --- a/diffulex/engine/block_manager.py +++ b/diffulex/engine/kvcache_manager.py @@ -2,13 +2,14 @@ import numpy as np +from typing import Callable from collections import deque from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Callable, Dict, Deque, Iterable, List, Set from diffulex.config import Config from diffulex.engine.sequence import SequenceBase +from diffulex.engine.strategy_registry import DiffulexStrategyRegistry @dataclass @@ -16,7 +17,7 @@ class Block: block_id: int ref_count: int = 0 hash: int = -1 - token_ids: List[int] = field(default_factory=list) + token_ids: list[int] = field(default_factory=list) def update(self, hash: int, token_ids: list[int]): self.hash = hash @@ -28,24 +29,20 @@ def reset(self): self.token_ids = [] -BlockManagerFactory = Callable[[Config], "BlockManagerBase"] -_NOT_PROVIDED = object() - - -class BlockManagerBase(ABC): +class KVCacheManagerBase(ABC): def __init__(self, config: Config): num_blocks = config.num_kvcache_blocks block_size = config.kvcache_block_size assert num_blocks > 0 self.config = config self.block_size = block_size - self.blocks: List[Block] = [Block(block_id=i) for i in range(num_blocks)] - self.hash_to_block_id: Dict[int, int] = dict() - self.free_block_ids: Deque[int] = deque(range(num_blocks)) - self.used_block_ids: Set[int] = set() + self.blocks: list[Block] = [Block(block_id=i) for i in range(num_blocks)] + self.hash_to_block_id: dict[int, int] = dict() + self.free_block_ids: deque[int] = deque(range(num_blocks)) + self.used_block_ids: set[int] = set() @classmethod - def compute_hash(cls, token_ids: List[int], prefix: int = -1): + def compute_hash(cls, token_ids: list[int], prefix: int = -1): h = xxhash.xxh64() if prefix != -1: h.update(prefix.to_bytes(8, "little")) @@ -112,58 +109,16 @@ def may_append(self, seq: SequenceBase) -> None: pass -class AutoBlockManager: - """Registry-driven factory for block manager implementations.""" - - _BLOCK_MANAGER_MAPPING: Dict[str, BlockManagerFactory] = {} - _DEFAULT_KEY = "__default__" +KVCacheManagerFactory = Callable[[Config], "KVCacheManagerBase"] - @classmethod - def register( - cls, - strategy_name: str, - factory: BlockManagerFactory | object = _NOT_PROVIDED, - *, - aliases: Iterable[str] = (), - is_default: bool = False, - exist_ok: bool = False, - ): - if not isinstance(strategy_name, str) or not strategy_name: - raise ValueError("strategy_name must be a non-empty string.") - if isinstance(aliases, str): - raise TypeError("aliases must be an iterable of strings, not a single string.") - - def decorator(factory_fn: BlockManagerFactory): - cls._register(strategy_name, factory_fn, exist_ok=exist_ok) - for alias in dict.fromkeys(aliases): - if not isinstance(alias, str) or not alias: - raise ValueError("aliases must contain non-empty strings.") - cls._register(alias, factory_fn, exist_ok=exist_ok) - if is_default: - cls._register(cls._DEFAULT_KEY, factory_fn, exist_ok=True) - return factory_fn - - if factory is _NOT_PROVIDED: - return decorator - return decorator(factory) - @classmethod - def _register(cls, key: str, factory: BlockManagerFactory, *, exist_ok: bool) -> None: - if not exist_ok and key in cls._BLOCK_MANAGER_MAPPING and cls._BLOCK_MANAGER_MAPPING[key] is not factory: - raise ValueError(f"Block manager '{key}' is already registered.") - cls._BLOCK_MANAGER_MAPPING[key] = factory - - @classmethod - def unregister(cls, strategy_name: str) -> None: - cls._BLOCK_MANAGER_MAPPING.pop(strategy_name, None) - - @classmethod - def available_block_managers(cls) -> tuple[str, ...]: - return tuple(sorted(k for k in cls._BLOCK_MANAGER_MAPPING if k != cls._DEFAULT_KEY)) +class AutoKVCacheManager(DiffulexStrategyRegistry): + """Registry-driven factory for block manager implementations.""" @classmethod - def from_config(cls, config: Config) -> BlockManagerBase: - candidates: List[str] = [] + def from_config(cls, config: Config) -> KVCacheManagerBase: + cls._MODULE_MAPPING: dict[str, KVCacheManagerFactory] + candidates: list[str] = [] for attr in ("decoding_strategy", "model_type"): value = getattr(config, attr, None) if isinstance(value, str) and value: @@ -171,11 +126,11 @@ def from_config(cls, config: Config) -> BlockManagerBase: candidates.append(cls._DEFAULT_KEY) for key in candidates: - factory = cls._BLOCK_MANAGER_MAPPING.get(key) + factory = cls._MODULE_MAPPING.get(key) if factory is not None: return factory(config) - available = ", ".join(cls.available_block_managers()) or "" + available = ", ".join(cls.available_modules()) or "" raise ValueError( "No block manager registered for decoding_strategy=" f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" diff --git a/diffulex/engine/model_runner.py b/diffulex/engine/model_runner.py index fbdb48c7..e5dfd1d9 100755 --- a/diffulex/engine/model_runner.py +++ b/diffulex/engine/model_runner.py @@ -3,23 +3,21 @@ import torch.distributed as dist -from typing import Callable, Dict, Iterable, List +from typing import Callable from abc import ABC, abstractmethod from multiprocessing.synchronize import Event from multiprocessing.shared_memory import SharedMemory from diffulex.config import Config +from diffulex.layer.sampler import AutoSampler from diffulex.engine.sequence import SequenceBase from diffulex.model.auto_model import AutoModelForDiffusionLM -from diffulex.layer.sampler import AutoSampler - -RunnerFactory = Callable[[Config, int, Event | List[Event]], "ModelRunnerBase"] -_NOT_PROVIDED = object() +from diffulex.engine.strategy_registry import DiffulexStrategyRegistry class ModelRunnerBase(ABC): """Base class for model runners supporting different model types.""" - def __init__(self, config: Config, rank: int, event: Event | List[Event]): + def __init__(self, config: Config, rank: int, event: Event | list[Event]): self.config = config self.model_type = config.model_type hf_config = config.hf_config @@ -130,23 +128,23 @@ def warmup_model(self): def allocate_kv_cache(self): pass - def prepare_block_tables(self, seqs: List[SequenceBase]): + def prepare_block_tables(self, seqs: list[SequenceBase]): max_len = max(len(seq.block_table) for seq in seqs) block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs] block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) return block_tables @abstractmethod - def prepare_prefill(self, seqs: List[SequenceBase]): + def prepare_prefill(self, seqs: list[SequenceBase]): """Model-specific prefill preparation.""" pass @abstractmethod - def prepare_decode(self, seqs: List[SequenceBase]): + def prepare_decode(self, seqs: list[SequenceBase]): """Model-specific decode preparation.""" pass - def prepare_sample(self, seqs: List[SequenceBase]): + def prepare_sample(self, seqs: list[SequenceBase]): temperatures = [] for seq in seqs: temperatures.append(seq.temperature) @@ -160,7 +158,7 @@ def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill pass @abstractmethod - def run(self, seqs: List[SequenceBase], is_prefill: bool) -> List[int]: + def run(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: """Main inference pipeline.""" pass @@ -169,9 +167,12 @@ def run(self, seqs: List[SequenceBase], is_prefill: bool) -> List[int]: def capture_cudagraph(self): """Model-specific CUDA graph capture.""" pass + + +RunnerFactory = Callable[[Config, int, Event | list[Event]], "ModelRunnerBase"] -class AutoModelRunner: +class AutoModelRunner(DiffulexStrategyRegistry): """Registry and factory that selects a ModelRunner implementation based on the configured decoding strategy. Example: @@ -182,55 +183,10 @@ class AutoModelRunner: This allows `LLMEngine` to instantiate the appropriate runner using `Config.decoding_strategy`. """ - _RUNNER_MAPPING: Dict[str, RunnerFactory] = {} - _DEFAULT_KEY = "__default__" - - @classmethod - def register( - cls, - strategy_name: str, - runner: RunnerFactory | object = _NOT_PROVIDED, - *, - aliases: Iterable[str] = (), - is_default: bool = False, - exist_ok: bool = False, - ): - if not isinstance(strategy_name, str) or not strategy_name: - raise ValueError("strategy_name must be a non-empty string.") - if isinstance(aliases, str): - raise TypeError("aliases must be an iterable of strings, not a single string.") - - def decorator(factory: RunnerFactory): - cls._register(strategy_name, factory, exist_ok=exist_ok) - for alias in dict.fromkeys(aliases): - if not isinstance(alias, str) or not alias: - raise ValueError("aliases must contain non-empty strings.") - cls._register(alias, factory, exist_ok=exist_ok) - if is_default: - cls._register(cls._DEFAULT_KEY, factory, exist_ok=True) - return factory - - if runner is _NOT_PROVIDED: - return decorator - return decorator(runner) - - @classmethod - def _register(cls, key: str, factory: RunnerFactory, *, exist_ok: bool) -> None: - if not exist_ok and key in cls._RUNNER_MAPPING and cls._RUNNER_MAPPING[key] is not factory: - raise ValueError(f"Model runner '{key}' is already registered.") - cls._RUNNER_MAPPING[key] = factory - - @classmethod - def unregister(cls, strategy_name: str) -> None: - cls._RUNNER_MAPPING.pop(strategy_name, None) - - @classmethod - def available_runners(cls) -> tuple[str, ...]: - return tuple(sorted(k for k in cls._RUNNER_MAPPING if k != cls._DEFAULT_KEY)) - @classmethod - def from_config(cls, config: Config, rank: int, event: Event | List[Event]): - candidates: List[str] = [] + def from_config(cls, config: Config, rank: int, event: Event | list[Event]): + cls._MODULE_MAPPING: dict[str, RunnerFactory] + candidates: list[str] = [] for attr in ("decoding_strategy", "model_type"): value = getattr(config, attr, None) if isinstance(value, str) and value: @@ -238,11 +194,11 @@ def from_config(cls, config: Config, rank: int, event: Event | List[Event]): candidates.append(cls._DEFAULT_KEY) for key in candidates: - factory = cls._RUNNER_MAPPING.get(key) + factory = cls._MODULE_MAPPING.get(key) if factory is not None: return factory(config, rank, event) - available = ", ".join(cls.available_runners()) or "" + available = ", ".join(cls.available_modules()) or "" raise ValueError( "No model runner registered for decoding_strategy=" f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" diff --git a/diffulex/engine/scheduler.py b/diffulex/engine/scheduler.py index 393949f8..286a7cbd 100755 --- a/diffulex/engine/scheduler.py +++ b/diffulex/engine/scheduler.py @@ -1,10 +1,11 @@ +from typing import Callable from collections import deque from abc import ABC, abstractmethod -from typing import Callable, Deque, Dict, Iterable, List, Tuple from diffulex.config import Config from diffulex.engine.sequence import SequenceBase -from diffulex.engine.block_manager import AutoBlockManager +from diffulex.engine.kvcache_manager import AutoKVCacheManager +from diffulex.engine.strategy_registry import DiffulexStrategyRegistry class SchedulerBase(ABC): @@ -13,9 +14,9 @@ def __init__(self, config: Config): self.max_num_seqs = config.max_num_seqs self.max_num_batched_tokens = config.max_num_batched_tokens self.eos = config.eos - self.block_manager = AutoBlockManager.from_config(config) - self.waiting: Deque[SequenceBase] = deque() - self.running: Deque[SequenceBase] = deque() + self.block_manager = AutoKVCacheManager.from_config(config) + self.waiting: deque[SequenceBase] = deque() + self.running: deque[SequenceBase] = deque() @abstractmethod def is_finished(self) -> bool: @@ -26,7 +27,7 @@ def add(self, seq: SequenceBase) -> None: pass @abstractmethod - def schedule(self) -> Tuple[List[SequenceBase], bool]: + def schedule(self) -> tuple[list[SequenceBase], bool]: pass @abstractmethod @@ -34,66 +35,20 @@ def preempt(self, seq: SequenceBase) -> None: pass @abstractmethod - def postprocess(self, seqs: List[SequenceBase], sampler_output): + def postprocess(self, seqs: list[SequenceBase], sampler_output): pass SchedulerFactory = Callable[[Config], "SchedulerBase"] -_NOT_PROVIDED = object() -class AutoScheduler: +class AutoScheduler(DiffulexStrategyRegistry): """Registry-driven factory for scheduler implementations.""" - _SCHEDULER_MAPPING: Dict[str, SchedulerFactory] = {} - _DEFAULT_KEY = "__default__" - - @classmethod - def register( - cls, - strategy_name: str, - factory: SchedulerFactory | object = _NOT_PROVIDED, - *, - aliases: Iterable[str] = (), - is_default: bool = False, - exist_ok: bool = False, - ): - if not isinstance(strategy_name, str) or not strategy_name: - raise ValueError("strategy_name must be a non-empty string.") - if isinstance(aliases, str): - raise TypeError("aliases must be an iterable of strings, not a single string.") - - def decorator(factory_fn: SchedulerFactory): - cls._register(strategy_name, factory_fn, exist_ok=exist_ok) - for alias in dict.fromkeys(aliases): - if not isinstance(alias, str) or not alias: - raise ValueError("aliases must contain non-empty strings.") - cls._register(alias, factory_fn, exist_ok=exist_ok) - if is_default: - cls._register(cls._DEFAULT_KEY, factory_fn, exist_ok=True) - return factory_fn - - if factory is _NOT_PROVIDED: - return decorator - return decorator(factory) - - @classmethod - def _register(cls, key: str, factory: SchedulerFactory, *, exist_ok: bool) -> None: - if not exist_ok and key in cls._SCHEDULER_MAPPING and cls._SCHEDULER_MAPPING[key] is not factory: - raise ValueError(f"Scheduler '{key}' is already registered.") - cls._SCHEDULER_MAPPING[key] = factory - - @classmethod - def unregister(cls, strategy_name: str) -> None: - cls._SCHEDULER_MAPPING.pop(strategy_name, None) - - @classmethod - def available_schedulers(cls) -> tuple[str, ...]: - return tuple(sorted(k for k in cls._SCHEDULER_MAPPING if k != cls._DEFAULT_KEY)) - @classmethod def from_config(cls, config: Config) -> SchedulerBase: - candidates: List[str] = [] + cls._MODULE_MAPPING: dict[str, SchedulerFactory] + candidates: list[str] = [] for attr in ("decoding_strategy", "model_type"): value = getattr(config, attr, None) if isinstance(value, str) and value: @@ -101,11 +56,11 @@ def from_config(cls, config: Config) -> SchedulerBase: candidates.append(cls._DEFAULT_KEY) for key in candidates: - factory = cls._SCHEDULER_MAPPING.get(key) + factory = cls._MODULE_MAPPING.get(key) if factory is not None: return factory(config) - available = ", ".join(cls.available_schedulers()) or "" + available = ", ".join(cls.available_modules()) or "" raise ValueError( "No scheduler registered for decoding_strategy=" f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" diff --git a/diffulex/engine/sequence.py b/diffulex/engine/sequence.py index 653fea44..57c9aad6 100755 --- a/diffulex/engine/sequence.py +++ b/diffulex/engine/sequence.py @@ -1,13 +1,15 @@ -import torch +"""Sequence base class and registry.""" + +from __future__ import annotations from copy import copy from enum import Enum, auto from itertools import count -from dataclasses import dataclass -from typing import List, Tuple, Any +from typing import Callable from diffulex.config import Config from diffulex.sampling_params import SamplingParams +from diffulex.engine.strategy_registry import DiffulexStrategyRegistry class SequenceStatus(Enum): @@ -17,10 +19,12 @@ class SequenceStatus(Enum): class SequenceBase: + """Minimal base class that tracks prompt tokens and cache bookkeeping.""" + block_size = 256 counter = count() - - def __init__(self, token_ids: List[int], sampling_params: SamplingParams = SamplingParams()): + + def __init__(self, token_ids: list[int], sampling_params: SamplingParams = SamplingParams()): self.seq_id = next(SequenceBase.counter) self.status = SequenceStatus.WAITING self.token_ids = copy(token_ids) @@ -28,11 +32,12 @@ def __init__(self, token_ids: List[int], sampling_params: SamplingParams = Sampl self.num_tokens = len(token_ids) self.num_prompt_tokens = len(token_ids) self.num_cached_tokens = 0 - self.block_table = [] - self.block_cache_missed = [] + self.block_table: list[int] = [] + self.block_cache_missed: list[bool] = [] self.temperature = sampling_params.temperature self.max_tokens = sampling_params.max_tokens self.ignore_eos = sampling_params.ignore_eos + self.new_tokens = 0 def __len__(self) -> int: return self.num_tokens @@ -45,7 +50,7 @@ def is_finished(self) -> bool: return self.status == SequenceStatus.FINISHED @property - def prompt_token_ids(self) -> List[int]: + def prompt_token_ids(self) -> list[int]: return self.token_ids[:self.num_prompt_tokens] @property @@ -56,9 +61,9 @@ def num_blocks(self) -> int: def last_block_num_tokens(self) -> int: return self.num_tokens - (self.num_blocks - 1) * self.block_size - def block(self, i) -> List[int]: - assert 0 <= i < self.num_blocks - return self.token_ids[i * self.block_size: (i + 1) * self.block_size] + def block(self, index: int) -> list[int]: + assert 0 <= index < self.num_blocks + return self.token_ids[index * self.block_size : (index + 1) * self.block_size] def append_token(self, token_id: int) -> None: self.token_ids.append(token_id) @@ -66,451 +71,35 @@ def append_token(self, token_id: int) -> None: self.num_tokens += 1 -class DiffusionBlockStatus(Enum): - ACTIVE = auto() - TO_CACHE = auto() - IN_CACHE = auto() - - -@dataclass -class DiffusionBlock: - block_id: int = 0 - status: DiffusionBlockStatus = DiffusionBlockStatus.ACTIVE - - global_start_id: int = 0 - global_end_id: int | None = None - cursor: int = 0 - - mask_token_id: int = 151666 - size: int = 32 - is_prompt: bool = False - - accept_threshold: float = 0.95 # Threshold to accept a token in the diffusion block - add_new_block_threshold: float = 0.1 # Threshold to add a new block - complete_threshold: float = 0.9 # Can only be cached when the current diffusion block is completed - - seq: "SequenceForDiffusionLM" = None # Reference to the sequence this block belongs to - pre_block: "DiffusionBlock" = None # Create prefix linked list of diffusion blocks - suf_block: "DiffusionBlock" = None # Create suffix linked list of diffusion blocks - - def __post_init__(self): - self.global_end_id = self.global_start_id + self.size - - def __getitem__(self, key: int) -> int: - return self.seq[self.global_start_id + key] - - def __len__(self) -> int: - return self.size - - @property - def current_complete_ratio(self) -> float: - return ( - sum([token_id != self.mask_token_id for token_id in self.token_ids]) / self.size - ) if self.size > 0 else 0.0 - - @property - def available_to_cache(self) -> bool: - return self.current_complete_ratio == 1.0 - - @property - def is_active(self) -> bool: - return self.status == DiffusionBlockStatus.ACTIVE - - @property - def is_in_cache(self) -> bool: - return self.status == DiffusionBlockStatus.IN_CACHE - - @property - def is_to_cache(self) -> bool: - return self.status == DiffusionBlockStatus.TO_CACHE - - @property - def pre_block_complete(self) -> bool: - return self.pre_block.current_complete_ratio >= self.complete_threshold if self.pre_block is not None else True - - @property - def add_new_block(self) -> bool: - return self.current_complete_ratio >= self.add_new_block_threshold - - @property - def token_ids(self) -> torch.Tensor: - if self.seq is not None: - return self.seq.token_ids[self.global_start_id:self.global_end_id] - else: - raise RuntimeError("Sequence is not set for the diffusion block.") - - @property - def local_mask_tokens(self) -> List[bool]: - return [token_id == self.seq.mask_token_id for token_id in self.token_ids] - - @property - def local_mask_token_ids(self) -> List[int]: - return [idx for idx, mask_token in enumerate(self.local_mask_tokens) if mask_token] - - @property - def global_mask_token_ids(self) -> List[int]: - 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) - return [mask_token_id + offset for mask_token_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 = DiffusionBlockStatus.TO_CACHE - - def in_cache(self) -> None: - if self.is_to_cache: - self.status = DiffusionBlockStatus.IN_CACHE - - def modify_token(self, local_token_id: int, modified_to: int) -> None: - 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() - self.seq.new_tokens += 1 - - -class SequenceForDiffusionLM(SequenceBase): - """Sequence implementation for Diffusion Language Models.""" - - def __init__(self, token_ids: List[int], - sampling_params = SamplingParams(), - config: Config = None): - super().__init__(token_ids, sampling_params) - 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 = None - self.meet_eos = False - self.diffusion_blocks: List[DiffusionBlock] = [] - self.n_steps = 0 - - 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": getattr(self, "input_token_ids", []), - "input_num_tokens": getattr(self, "input_num_tokens", 0), - "input_num_prompt_tokens": getattr(self, "input_num_prompt_tokens", 0), - "new_tokens": getattr(self, "new_tokens", 0), - "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", None) - self.n_steps = state.get("n_steps", 0) - # Align tensor devices when sequence is reconstructed on a different rank - 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 = DiffusionBlock( - 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 - - def __repr__(self) -> str: - return (f"SequenceForDiffusionLM(block_size={self.block_size}, counter={self.counter}, " - f"seq_id={self.seq_id}, status={self.status.name}, num_tokens={self.num_tokens}, " - f"num_prompt_tokens={self.num_prompt_tokens}, num_cached_tokens={self.num_cached_tokens}, " - f"temperature={self.temperature}, max_tokens={self.max_tokens}, ignore_eos={self.ignore_eos}, " - f"diffusion_block_size={self.diffusion_block_size}, " - f"block_mask={(self.block_mask.shape if self.block_mask is not None else None)}, " - f"input_token_ids={getattr(self, 'input_token_ids', None)}, input_num_tokens={getattr(self, 'input_num_tokens', None)})") - - @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:] - - @property - def active_blocks(self) -> List[bool]: - return [block.is_active for block in self.diffusion_blocks] - - @property - def to_cache_blocks(self) -> List[bool]: - return [block.is_to_cache 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] - - @property - def num_prompt_blocks(self) -> int: - return (self.input_num_prompt_tokens + self.block_size - 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 - - @property - def updated_or_updating_kv_cache_block_ids(self) -> List[int]: - return [idx for idx, caching in enumerate(self.caching_blocks) if caching] - - @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] - - @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 cached_num_tokens - 1 - - @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) - - @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)]) # build up token-wise mapping - return mapping - - def token_to_diffusion_block_id(self, token_id: int) -> int: - if token_id < self.input_num_tokens: - return 0 - else: - return (token_id - self.input_num_tokens) // self.diffusion_block_size + 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 = [] - positions = [] - 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([token_id + temp_context_len for token_id in range(block.size)]) - temp_context_len += block.size - - return input_tokens, positions, context_len - - def reset_new_tokens(self) -> None: - self.new_tokens = 0 - - def post_process(self) -> None: - for diff_blk in self.diffusion_blocks: - diff_blk.cursor = 0 - if diff_blk.is_in_cache: - continue - - if diff_blk.is_to_cache: - diff_blk.in_cache() - elif diff_blk.is_active: - if diff_blk.available_to_cache: - diff_blk.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.kv_cache_layout == "distinct": - return self.block_mask[..., self.cached_num_tokens:, self.cached_num_tokens:] - else: - 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_wise_causal_mask = torch.zeros(mask_shape, dtype=torch.bool, device=torch.cuda.current_device()) - block_wise_causal_mask[..., :self.input_num_tokens, :self.input_num_tokens] = True - num_diffusion_blocks = (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 - start_w = 0 - end_w = end_h - block_wise_causal_mask[..., start_h:end_h, start_w:end_w] = True - self.block_mask = block_wise_causal_mask.clone() - else: - return - assert self.block_mask is not None, "block_mask must exist before incremental update" - 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.zeros(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: - # Take a snapshot of the original input state - 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( - DiffusionBlock( - block_id=len(self.diffusion_blocks), - status=DiffusionBlockStatus.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 self.diffusion_blocks[-1].add_new_block and not self.meet_eos: - added_num_tokens = ( - self.diffusion_block_size - if self.num_tokens + self.diffusion_block_size <= self.max_model_len - else self.max_model_len - self.num_tokens - ) - - diffusion_seq = [self.mask_token_id] * added_num_tokens - current_diffusion_block = DiffusionBlock( - block_id=len(self.diffusion_blocks), - status=DiffusionBlockStatus.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] if self.diffusion_blocks else None - ) - - self.diffusion_blocks[-1].suf_block = current_diffusion_block - self.token_ids += diffusion_seq - self.num_tokens += added_num_tokens - self.diffusion_blocks.append(current_diffusion_block) - - self.update_block_mask(is_prefill=is_prefill) \ No newline at end of file +SequenceFactory = Callable[[list[int], SamplingParams, Config], SequenceBase] + + +class AutoSequence(DiffulexStrategyRegistry): + """Registry-driven factory for sequence implementations.""" + + @classmethod + def create( + cls, + config: Config, + token_ids: list[int], + sampling_params: SamplingParams = SamplingParams(), + ) -> SequenceBase: + cls._MODULE_MAPPING: dict[str, SequenceFactory] + candidates: list[str] = [] + for attr in ("decoding_strategy", "model_type"): + value = getattr(config, attr, None) + if isinstance(value, str) and value: + candidates.append(value) + candidates.append(cls._DEFAULT_KEY) + + for key in candidates: + factory = cls._MODULE_MAPPING.get(key) + if factory is not None: + return factory(token_ids, sampling_params, config) + + available = ", ".join(cls.available_modules()) or "" + raise ValueError( + "No sequence registered for decoding_strategy=" + f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" + f"'{getattr(config, 'model_type', None)}'. Available sequences: {available}." + ) \ No newline at end of file diff --git a/diffulex/engine/strategy_registry.py b/diffulex/engine/strategy_registry.py new file mode 100644 index 00000000..b3cd2ed4 --- /dev/null +++ b/diffulex/engine/strategy_registry.py @@ -0,0 +1,56 @@ +from typing import Iterable + +from diffulex.utils.registry import fetch_factory_name + + +_NOT_PROVIDED = object() + + +class DiffulexStrategyRegistry: + """Registry-driven factory for module implementations.""" + + _MODULE_MAPPING: dict[str, object] = {} + _DEFAULT_KEY = "__default__" + + @classmethod + def register( + cls, + strategy_name: str, + factory: object = _NOT_PROVIDED, + *, + aliases: Iterable[str] = (), + is_default: bool = False, + exist_ok: bool = False, + ): + if not isinstance(strategy_name, str) or not strategy_name: + raise ValueError("strategy_name must be a non-empty string.") + if isinstance(aliases, str): + raise TypeError("aliases must be an iterable of strings, not a single string.") + + def decorator(factory_fn: object): + cls._register(strategy_name, factory_fn, exist_ok=exist_ok) + for alias in dict.fromkeys(aliases): + if not isinstance(alias, str) or not alias: + raise ValueError("aliases must contain non-empty strings.") + cls._register(alias, factory_fn, exist_ok=exist_ok) + if is_default: + cls._register(cls._DEFAULT_KEY, factory_fn, exist_ok=True) + return factory_fn + + if factory is _NOT_PROVIDED: + return decorator + return decorator(factory) + + @classmethod + def _register(cls, key: str, factory: object, *, exist_ok: bool) -> None: + if not exist_ok and key in cls._MODULE_MAPPING and cls._MODULE_MAPPING[key] is not factory: + raise ValueError(f"Module '{key}: {fetch_factory_name(factory)}' is already registered.") + cls._MODULE_MAPPING[key] = factory + + @classmethod + def unregister(cls, strategy_name: str) -> None: + cls._MODULE_MAPPING.pop(strategy_name, None) + + @classmethod + def available_modules(cls) -> tuple[str, ...]: + return tuple(sorted(k for k in cls._MODULE_MAPPING if k != cls._DEFAULT_KEY)) \ No newline at end of file diff --git a/diffulex/engine/llm_engine.py b/diffulex/engine/tp_worker.py similarity index 92% rename from diffulex/engine/llm_engine.py rename to diffulex/engine/tp_worker.py index 23576aa6..464befc8 100755 --- a/diffulex/engine/llm_engine.py +++ b/diffulex/engine/tp_worker.py @@ -2,22 +2,19 @@ import torch.multiprocessing as mp -from typing import List from tqdm.auto import tqdm from time import perf_counter from dataclasses import fields from transformers import AutoTokenizer -import diffulex.strategy # noqa: F401 - from diffulex.config import Config from diffulex.sampling_params import SamplingParams -from diffulex.engine.sequence import SequenceForDiffusionLM +from diffulex.engine.sequence import AutoSequence from diffulex.engine.scheduler import AutoScheduler, SchedulerBase from diffulex.engine.model_runner import AutoModelRunner -class LLMEngine: +class DiffulexTPWorker: def __init__(self, model, **kwargs): config_fields = {field.name for field in fields(Config)} config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} @@ -58,11 +55,10 @@ def exit(self): except Exception: pass - def add_request(self, prompt: str | List[int], sampling_params: SamplingParams): + def add_request(self, prompt: str | list[int], sampling_params: SamplingParams): if isinstance(prompt, str): prompt = self.tokenizer.encode(prompt) - - seq = SequenceForDiffusionLM(prompt, sampling_params, config=self.config) + seq = AutoSequence.create(self.config, prompt, sampling_params) seq.block_size = self.config.kvcache_block_size self.scheduler.add(seq) # Return seq_id so caller can build a stable mapping @@ -83,10 +79,10 @@ def is_finished(self): def generate( self, - prompts: List[str] | List[List[int]], - sampling_params: SamplingParams | List[SamplingParams], + prompts: list[str] | list[list[int]], + sampling_params: SamplingParams | list[SamplingParams], use_tqdm: bool = True, - ) -> List[str]: + ) -> list[str]: if use_tqdm: pbar = tqdm(total=len(prompts), desc="Generating", dynamic_ncols=True) if not isinstance(sampling_params, list): diff --git a/diffulex/layer/sampler.py b/diffulex/layer/sampler.py index 5d9f0c7a..b1b64a03 100644 --- a/diffulex/layer/sampler.py +++ b/diffulex/layer/sampler.py @@ -4,12 +4,11 @@ import torch.nn.functional as F import torch.distributions as dists -from typing import List, Dict from dataclasses import dataclass -from easydict import EasyDict as edict +from easydict import Easydict as edict from diffulex.config import Config -from diffulex.utils.context import get_context_diffusion_lm +from diffulex.attention import fetch_attn_metadata class SamplerForDiffusionLM(nn.Module): @@ -71,9 +70,9 @@ def sample_tokens(self, logits, temperature=0.0, top_p=None, top_k=None, @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]]] + 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) @@ -97,7 +96,7 @@ def _shift_logits(self, logits, last_logit=None): def forward(self, logits: torch.Tensor, temperatures: torch.Tensor, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False): - context = get_context_diffusion_lm() + 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 = {} @@ -154,7 +153,7 @@ def forward(self, logits: torch.Tensor, temperatures: torch.Tensor, 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 = get_context_diffusion_lm() + 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 = {} diff --git a/diffulex/llm.py b/diffulex/llm.py deleted file mode 100755 index 7cd69bea..00000000 --- a/diffulex/llm.py +++ /dev/null @@ -1,10 +0,0 @@ -from diffulex.config import Config -from diffulex.engine.dp_engine import DPEngine -from diffulex.engine.llm_engine import LLMEngine - -class LLM: - def __new__(cls, model, **kwargs): - cfg = Config(model, **{k: v for k, v in kwargs.items() if k in Config.__dataclass_fields__.keys()}) - if cfg.data_parallel_size > 1: - return DPEngine(model, **kwargs) - return LLMEngine(model, **kwargs) diff --git a/diffulex/model/config/llada/configuration_llada.py b/diffulex/model/config/llada/configuration_llada.py index 01cd50ef..716fd8b3 100644 --- a/diffulex/model/config/llada/configuration_llada.py +++ b/diffulex/model/config/llada/configuration_llada.py @@ -11,9 +11,9 @@ from pathlib import Path from typing import ( Any, - Dict, + dict, Iterable, - List, + list, Optional, Tuple, Type, diff --git a/diffulex/model/dream.py b/diffulex/model/dream.py index 8b75a883..96d43648 100755 --- a/diffulex/model/dream.py +++ b/diffulex/model/dream.py @@ -3,10 +3,10 @@ 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.attention.attention_v5 import Attention from diffulex.model.auto_model import AutoModelForDiffusionLM from diffulex.model.config.dream.configuration_dream import DreamConfig from diffulex.layer.linear import RowParallelLinear, ColumnParallelLinear @@ -189,7 +189,7 @@ def __init__( ) -> None: super().__init__() self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) - self.layers = nn.ModuleList([DreamDecoderLayer(config) + self.layers = nn.Modulelist([DreamDecoderLayer(config) for _ in range(config.num_hidden_layers)]) self.norm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/diffulex/model/fast_dllm_v2.py b/diffulex/model/fast_dllm_v2.py index ba3724b2..1647559c 100755 --- a/diffulex/model/fast_dllm_v2.py +++ b/diffulex/model/fast_dllm_v2.py @@ -3,10 +3,10 @@ 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.attention.attention_v5 import Attention from diffulex.model.auto_model import AutoModelForDiffusionLM from diffulex.layer.linear import RowParallelLinear, ColumnParallelLinear from diffulex.layer.embed_head import VocabParallelEmbedding, ParallelLMHead @@ -189,7 +189,7 @@ def __init__( ) -> None: super().__init__() self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) - self.layers = nn.ModuleList([FastdLLMV2DecoderLayer(config) + self.layers = nn.Modulelist([FastdLLMV2DecoderLayer(config) for _ in range(config.num_hidden_layers)]) self.norm = FastdLLMV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/diffulex/model/llada.py b/diffulex/model/llada.py index b1c3e485..5a1b79bb 100755 --- a/diffulex/model/llada.py +++ b/diffulex/model/llada.py @@ -3,10 +3,10 @@ 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.attention.attention_v5 import Attention from diffulex.model.auto_model import AutoModelForDiffusionLM from diffulex.model.config.llada.configuration_llada import LLaDAConfig from diffulex.layer.linear import RowParallelLinear, ColumnParallelLinear @@ -189,7 +189,7 @@ def __init__( ) -> None: super().__init__() self.config = config - self.transformer = nn.ModuleDict( + self.transformer = nn.Moduledict( dict( wte=VocabParallelEmbedding( config.embedding_size or config.vocab_size, config.d_model @@ -200,7 +200,7 @@ def __init__( ) blocks = [LLaDABlock(config) for _ in range(config.n_layers)] - self.transformer.update({"blocks": nn.ModuleList(blocks)}) + self.transformer.update({"blocks": nn.Modulelist(blocks)}) if not (self.config.alibi or self.config.rope): self.transformer.update( diff --git a/diffulex/model/utils/check_config.py b/diffulex/model/utils/check_config.py deleted file mode 100755 index 9261e45b..00000000 --- a/diffulex/model/utils/check_config.py +++ /dev/null @@ -1,8 +0,0 @@ -from typing import Optional, Type - -def check_config_diff( - current_config: Optional[Type] = None, - default_config_cls: Optional[Type] = None, - compare_to: str = "default", # "default" or "existing" -): - pass \ No newline at end of file diff --git a/diffulex/strategy/__init__.py b/diffulex/strategy/__init__.py index d0b96e09..e19f44c4 100644 --- a/diffulex/strategy/__init__.py +++ b/diffulex/strategy/__init__.py @@ -5,3 +5,16 @@ from . import d2f # noqa: F401 __all__ = ["d2f"] + +DECODING_STRATEGY = None + +def fetch_decoding_strategy() -> str | None: + return DECODING_STRATEGY + +def set_decoding_strategy(strategy: str) -> None: + global DECODING_STRATEGY + DECODING_STRATEGY = strategy + +def reset_decoding_strategy() -> None: + global DECODING_STRATEGY + DECODING_STRATEGY = None \ No newline at end of file diff --git a/diffulex/strategy/d2f/__init__.py b/diffulex/strategy/d2f/__init__.py index d25e1991..a83722c7 100644 --- a/diffulex/strategy/d2f/__init__.py +++ b/diffulex/strategy/d2f/__init__.py @@ -1,12 +1,14 @@ """D2F strategy component exports.""" from __future__ import annotations -from .block_manager import D2FBlockManager -from .model_runner import D2FModelRunner -from .scheduler import D2FScheduler +from .engine.kvcache_manager import D2FKVCacheManager +from .engine.model_runner import D2FModelRunner +from .engine.scheduler import D2FScheduler +from .engine.sequence import D2FSequence __all__ = [ - "D2FBlockManager", - "D2FModelRunner", - "D2FScheduler", + "D2FKVCacheManager", + "D2FModelRunner", + "D2FScheduler", + "D2FSequence", ] diff --git a/diffulex/strategy/d2f/attention/metadata.py b/diffulex/strategy/d2f/attention/metadata.py new file mode 100644 index 00000000..def3f344 --- /dev/null +++ b/diffulex/strategy/d2f/attention/metadata.py @@ -0,0 +1,28 @@ +import torch + +from dataclasses import dataclass + +from diffulex.attention.metadata import AttnMetaDataBase + + +@dataclass +class D2FAttnMetaData(AttnMetaDataBase): + seq_lens: list[int] = None + seq_lens_ts: torch.Tensor | None = None + d2f_pp: bool = False + block_mask: list[torch.Tensor] | None = None + + +D2F_ATTN_METADATA = D2FAttnMetaData() + +def fetch_d2f_attn_metadata() -> D2FAttnMetaData: + return D2F_ATTN_METADATA + +def set_d2f_attn_metadata() -> None: + # TODO + global D2F_ATTN_METADATA + D2F_ATTN_METADATA = D2FAttnMetaData() + +def reset_d2f_attn_metadata() -> None: + global D2F_ATTN_METADATA + D2F_ATTN_METADATA = D2FAttnMetaData() \ No newline at end of file diff --git a/diffulex/strategy/d2f/block_manager.py b/diffulex/strategy/d2f/engine/kvcache_manager.py similarity index 75% rename from diffulex/strategy/d2f/block_manager.py rename to diffulex/strategy/d2f/engine/kvcache_manager.py index d348b453..70c7f583 100644 --- a/diffulex/strategy/d2f/block_manager.py +++ b/diffulex/strategy/d2f/engine/kvcache_manager.py @@ -1,26 +1,28 @@ from __future__ import annotations -from typing import List +from typing import TYPE_CHECKING, list from diffulex.config import Config -from diffulex.engine.block_manager import AutoBlockManager, BlockManagerBase -from diffulex.engine.sequence import SequenceForDiffusionLM +from diffulex.engine.kvcache_manager import AutoKVCacheManager, KVCacheManagerBase +if TYPE_CHECKING: + from .sequence import D2FSequence -@AutoBlockManager.register( + +@AutoKVCacheManager.register( "d2f", aliases=("diffusion_lm",), is_default=True, ) -class D2FBlockManager(BlockManagerBase): +class D2FKVCacheManager(KVCacheManagerBase): def __init__(self, config: Config): super().__init__(config) - def can_append(self, seq: SequenceForDiffusionLM) -> bool: + 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 - def may_append(self, seq: SequenceForDiffusionLM) -> None: + def may_append(self, seq: "D2FSequence") -> None: if seq.cached_or_caching_num_tokens == 0: return block_table = seq.block_table @@ -32,7 +34,7 @@ def may_append(self, seq: SequenceForDiffusionLM) -> None: prev_end_token = seq.cached_or_caching_num_tokens - seq.caching_num_tokens - 1 prev_block_idx = prev_end_token // self.block_size if prev_block_idx < seq.num_blocks: - token_ids: List[int] = seq.block(prev_block_idx) + token_ids: list[int] = seq.block(prev_block_idx) prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 h = self.compute_hash(token_ids, prefix) last_block.update(h, token_ids) diff --git a/diffulex/strategy/d2f/model_runner.py b/diffulex/strategy/d2f/engine/model_runner.py similarity index 93% rename from diffulex/strategy/d2f/model_runner.py rename to diffulex/strategy/d2f/engine/model_runner.py index d0b932c6..481049b0 100644 --- a/diffulex/strategy/d2f/model_runner.py +++ b/diffulex/strategy/d2f/engine/model_runner.py @@ -1,19 +1,17 @@ from __future__ import annotations import time -from typing import List +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.d2f.engine.sequence import D2FSequence +from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata from diffulex.engine.model_runner import AutoModelRunner, ModelRunnerBase -from diffulex.engine.sequence import SequenceForDiffusionLM, SequenceBase -from diffulex.utils.context import ( - get_context_diffusion_lm, - reset_context_diffusion_lm, - set_context_diffusion_lm, -) +from diffulex.strategy.d2f.attention.metadata import fetch_d2f_attn_metadata, set_d2f_attn_metadata, reset_d2f_attn_metadata @AutoModelRunner.register( @@ -24,11 +22,12 @@ class D2FModelRunner(ModelRunnerBase): """Reference implementation of D2F decoding strategy.""" - def __init__(self, config: Config, rank: int, event: Event | List[Event]): + def __init__(self, config: Config, rank: int, event: Event | list[Event]): super().__init__(config, rank, event) 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) def warmup_model(self): print("Warming up model...") @@ -40,7 +39,7 @@ 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 = [SequenceForDiffusionLM(test_input_ids, config=self.config) for _ in range(num_seqs)] + seqs = [D2FSequence(test_input_ids, config=self.config) for _ in range(num_seqs)] self.run(seqs, True) for seq in seqs: seq.post_process() @@ -156,17 +155,17 @@ def allocate_kv_cache(self): ) ) - def prepare_prefill(self, seqs: List[SequenceForDiffusionLM]): - input_ids: List[int] = [] - positions: List[int] = [] + def prepare_prefill(self, seqs: list[D2FSequence]): + input_ids: list[int] = [] + positions: list[int] = [] cu_seqlens_q = [0] cu_seqlens_k = [0] max_seqlen_q = 0 max_seqlen_k = 0 - slot_mapping: List[int] = [] + slot_mapping: list[int] = [] block_tables = None - context_lens: List[int] = [] - seq_lens: List[int] = [] + context_lens: list[int] = [] + seq_lens: list[int] = [] for seq in seqs: seq.next_diffusion_step(is_prefill=True) @@ -228,7 +227,7 @@ def prepare_prefill(self, seqs: List[SequenceForDiffusionLM]): ) ) - set_context_diffusion_lm( + set_d2f_attn_metadata( True, cu_seqlens_q=cu_seqlens_q_tensor, cu_seqlens_k=cu_seqlens_k_tensor, @@ -244,14 +243,14 @@ def prepare_prefill(self, seqs: List[SequenceForDiffusionLM]): ) return input_ids_tensor, positions_tensor - def prepare_decode(self, seqs: List[SequenceForDiffusionLM]): - input_ids: List[int] = [] - positions: List[int] = [] + def prepare_decode(self, seqs: list[D2FSequence]): + 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] = [] + 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 @@ -351,7 +350,7 @@ def get_step(diff_blk, begin_idx): 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_context_diffusion_lm( + set_d2f_attn_metadata( False, slot_mapping=slot_mapping_tensor, context_lens=context_lens_tensor, @@ -374,7 +373,7 @@ def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill 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 = get_context_diffusion_lm() + context = fetch_d2f_attn_metadata() graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] graph_vars = self.graph_vars for key, value in graph_vars.items(): @@ -388,7 +387,8 @@ def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill graph.replay() return self.model.compute_logits(graph_vars["outputs"][:bs]) - def run_verbose(self, seqs: List[SequenceBase], is_prefill: bool) -> List[int]: + @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() @@ -401,15 +401,15 @@ def run_verbose(self, seqs: List[SequenceBase], is_prefill: bool) -> List[int]: 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_context_diffusion_lm() + reset_d2f_attn_metadata() return sample_output - def run(self, seqs: List[SequenceBase], is_prefill: bool) -> List[int]: + 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_context_diffusion_lm() + reset_d2f_attn_metadata() return sample_output @torch.inference_mode() diff --git a/diffulex/strategy/d2f/scheduler.py b/diffulex/strategy/d2f/engine/scheduler.py similarity index 91% rename from diffulex/strategy/d2f/scheduler.py rename to diffulex/strategy/d2f/engine/scheduler.py index 7be134da..dc650c69 100644 --- a/diffulex/strategy/d2f/scheduler.py +++ b/diffulex/strategy/d2f/engine/scheduler.py @@ -1,14 +1,9 @@ from __future__ import annotations -from typing import Dict, List, Tuple - from diffulex.config import Config from diffulex.engine.scheduler import AutoScheduler, SchedulerBase -from diffulex.engine.sequence import ( - SequenceBase, - SequenceForDiffusionLM, - SequenceStatus, -) +from diffulex.engine.sequence import SequenceBase, SequenceStatus +from .sequence import D2FSequence from diffulex.layer.sampler import SampleOutputForDiffusionLM @@ -25,11 +20,11 @@ def __init__(self, config: Config): def is_finished(self) -> bool: return not self.waiting and not self.running - def add(self, seq: SequenceForDiffusionLM) -> None: + 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[SequenceBase], bool]: + scheduled: list[SequenceBase] = [] num_seqs = 0 num_batched_tokens = 0 while self.waiting and num_seqs < self.max_num_seqs: @@ -92,17 +87,17 @@ def schedule(self) -> Tuple[List[SequenceBase], bool]: self.running.extendleft(reversed(scheduled)) return scheduled, False - def preempt(self, seq: SequenceForDiffusionLM) -> None: + def preempt(self, seq: D2FSequence) -> None: seq.status = SequenceStatus.WAITING self.block_manager.free(seq) self.waiting.appendleft(seq) def postprocess( self, - seqs: List[SequenceForDiffusionLM], + seqs: list[D2FSequence], sample_output: SampleOutputForDiffusionLM, - ) -> Dict[int, int]: - n_diff_steps: Dict[int, int] = {} + ) -> dict[int, int]: + n_diff_steps: dict[int, int] = {} for seq in seqs: seq.reset_new_tokens() seq_id = str(seq.seq_id) diff --git a/diffulex/strategy/d2f/engine/sequence.py b/diffulex/strategy/d2f/engine/sequence.py new file mode 100644 index 00000000..379439ee --- /dev/null +++ b/diffulex/strategy/d2f/engine/sequence.py @@ -0,0 +1,484 @@ +from __future__ import annotations + +import torch + +from dataclasses import dataclass +from enum import Enum, auto + +from diffulex.config import Config +from diffulex.engine.sequence import AutoSequence, SequenceBase +from diffulex.sampling_params import SamplingParams + + +class D2FDiffusionBlockStatus(Enum): + ACTIVE = auto() + TO_CACHE = auto() + IN_CACHE = auto() + + +@dataclass +class D2FDiffusionBlock: + block_id: int = 0 + status: D2FDiffusionBlockStatus = D2FDiffusionBlockStatus.ACTIVE + + global_start_id: int = 0 + global_end_id: int | None = None + cursor: int = 0 + + 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: "D2FSequence" | None = None + pre_block: "D2FDiffusionBlock" | None = None + suf_block: "D2FDiffusionBlock" | None = None + + def __post_init__(self) -> None: + self.global_end_id = self.global_start_id + self.size + + def __getitem__(self, key: int) -> int: + return self.seq[self.global_start_id + key] # type: ignore[index] + + def __len__(self) -> int: + return self.size + + @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 + + @property + def available_to_cache(self) -> bool: + return self.current_complete_ratio == 1.0 + + @property + def is_active(self) -> bool: + return self.status == D2FDiffusionBlockStatus.ACTIVE + + @property + def is_in_cache(self) -> bool: + return self.status == D2FDiffusionBlockStatus.IN_CACHE + + @property + def is_to_cache(self) -> bool: + return self.status == D2FDiffusionBlockStatus.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 + + @property + def add_new_block(self) -> bool: + return self.current_complete_ratio >= self.add_new_block_threshold + + @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] + + @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] + + @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) + 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 = D2FDiffusionBlockStatus.TO_CACHE + + def in_cache(self) -> None: + if self.is_to_cache: + self.status = D2FDiffusionBlockStatus.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( + "d2f", + aliases=("diffusion_lm",), + is_default=True, +) +class D2FSequence(SequenceBase): + """Sequence implementation tailored for diffusion-based decoding.""" + + def __init__( + self, + token_ids: list[int], + sampling_params: SamplingParams = SamplingParams(), + config: Config | None = None, + ): + super().__init__(token_ids, sampling_params) + 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] = [] + 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 = D2FDiffusionBlock( + 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 :] + + @property + def active_blocks(self) -> list[bool]: + return [block.is_active for block in self.diffusion_blocks] + + @property + def to_cache_blocks(self) -> list[bool]: + return [block.is_to_cache 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] + + @property + def num_prompt_blocks(self) -> int: + return (self.input_num_prompt_tokens + self.block_size - 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 + + @property + def updated_or_updating_kv_cache_block_ids(self) -> list[int]: + return [idx for idx, caching in enumerate(self.caching_blocks) if caching] + + @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] + + @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) + + @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) + + @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 + + @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 reset_new_tokens(self) -> None: + self.new_tokens = 0 + + def post_process(self) -> None: + for block in self.diffusion_blocks: + block.cursor = 0 + if block.is_in_cache: + continue + if block.is_to_cache: + block.in_cache() + elif block.is_active: + 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( + D2FDiffusionBlock( + block_id=len(self.diffusion_blocks), + status=D2FDiffusionBlockStatus.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 = D2FDiffusionBlock( + block_id=len(self.diffusion_blocks), + status=D2FDiffusionBlockStatus.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 diff --git a/diffulex/utils/context.py b/diffulex/utils/context.py deleted file mode 100755 index f2261de2..00000000 --- a/diffulex/utils/context.py +++ /dev/null @@ -1,112 +0,0 @@ -import torch - -from dataclasses import dataclass - -from diffulex.legacy.engine.sequence import SequenceForDiffusionLM - -@dataclass -class ContextBase: - 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 - - -# Global context for diffusion language model -@dataclass -class ContextForDiffusionLM(ContextBase): - seqs: list[SequenceForDiffusionLM] = None - seq_lens: list[int] = None - seq_lens_ts: torch.Tensor | None = None - kv_cache_layout: str = "unified" # "unified" or "distinct" - need_kv_cache_store: bool = True - block_mask: list[torch.Tensor] | None = None - - 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: - 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 block_mask_for_checking(self) -> torch.Tensor: - for seq in self.seqs: - seq.set_layout("unified") - - 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) - 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 - block_mask[start_row:end_row, start_col:end_col] = mask.clone() - start_row, start_col = end_row, start_col - - for seq in self.seqs: - seq.set_layout("distinct") - return block_mask.to(mask.device) - - @property - def total_num_seqs(self) -> int: - return len(self.seqs) if self.seqs is not None else 0 - -_CONTEXT_FOR_DIFFUSION_LM = ContextForDiffusionLM() - -def get_context_diffusion_lm() -> ContextForDiffusionLM: - return _CONTEXT_FOR_DIFFUSION_LM - -def set_context_diffusion_lm( - is_prefill, - cu_seqlens_q=None, cu_seqlens_k=None, - max_seqlen_q=0, max_seqlen_k=0, - slot_mapping=None, context_lens=None, block_tables=None, - seqs=None, seq_lens=None, seq_lens_ts=None, kv_cache_layout="unified", need_kv_cache_store=True, - d2f_pp=False -) -> None: - global _CONTEXT_FOR_DIFFUSION_LM - _CONTEXT_FOR_DIFFUSION_LM = ContextForDiffusionLM( - is_prefill, - cu_seqlens_q, cu_seqlens_k, - max_seqlen_q, max_seqlen_k, - slot_mapping, context_lens, block_tables, - seqs, seq_lens, seq_lens_ts, kv_cache_layout, need_kv_cache_store, d2f_pp - ) - -def reset_context_diffusion_lm() -> None: - global _CONTEXT_FOR_DIFFUSION_LM - _CONTEXT_FOR_DIFFUSION_LM = ContextForDiffusionLM() \ No newline at end of file diff --git a/diffulex/utils/registry.py b/diffulex/utils/registry.py new file mode 100644 index 00000000..4602ae90 --- /dev/null +++ b/diffulex/utils/registry.py @@ -0,0 +1,34 @@ +import inspect +import functools + +from typing import Any + + +def fetch_factory_name(factory: Any) -> str: + # unwrap decorated functions (inspect.unwrap works for functions with __wrapped__) + try: + orig = inspect.unwrap(factory) + except Exception: + orig = factory + + # handle functools.partial + if isinstance(orig, functools.partial): + return fetch_factory_name(orig.func) + + # class + if inspect.isclass(orig): + name = getattr(orig, "__qualname__", orig.__name__) + module = getattr(orig, "__module__", "") + # function + elif inspect.isfunction(orig): + name = getattr(orig, "__qualname__", orig.__name__) + module = getattr(orig, "__module__", "") + else: + # callable instance (object with __call__) + name = getattr(orig, "__name__", None) or orig.__class__.__name__ + module = getattr(orig, "__module__", getattr(orig.__class__, "__module__", "")) + + if module and module != "builtins": + return f"{module}.{name}" + return name + \ No newline at end of file diff --git a/examples/eval_llada.py b/examples/eval_llada.py index f3b66033..f9c2dd16 100755 --- a/examples/eval_llada.py +++ b/examples/eval_llada.py @@ -3,7 +3,7 @@ import json import time # add time module from datetime import timedelta -from typing import List, Optional, Tuple, Type, TypeVar, Union, Dict +from typing import List, Optional, Tuple, Type, TypeVar, Union, dict import torch import torch.nn.functional as F import torch.distributions as dists @@ -543,7 +543,7 @@ def create_from_arg_string( return cls(**args, **args2) def apply_chat_template( - self, chat_history: List[Dict[str, str]], add_generation_prompt: bool = True + self, chat_history: List[dict[str, str]], add_generation_prompt: bool = True ) -> str: """ Method to apply a chat template to a list of chat history between user and model. diff --git a/examples/model_cache/dream/generation_utils.py b/examples/model_cache/dream/generation_utils.py index ecc30e19..1d180884 100755 --- a/examples/model_cache/dream/generation_utils.py +++ b/examples/model_cache/dream/generation_utils.py @@ -16,7 +16,7 @@ import warnings import copy from dataclasses import dataclass -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any, dict, Optional, Tuple, Union import torch import torch.distributions as dists @@ -152,7 +152,7 @@ def _expand_inputs_for_generation( expand_size: int = 1, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.LongTensor] = None - ) -> Tuple[torch.LongTensor, Dict[str, Any]]: + ) -> Tuple[torch.LongTensor, dict[str, Any]]: """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]""" # Do not call torch.repeat_interleave if expand_size is 1 because it clones # the input tensor and thus requires more memory although no change is applied @@ -216,7 +216,7 @@ def _prepare_generated_length( return generation_config def _prepare_generation_config( - self, generation_config: Optional[DreamGenerationConfig], **kwargs: Dict + self, generation_config: Optional[DreamGenerationConfig], **kwargs: dict ) -> DreamGenerationConfig: """ Prepares the base generation config, then applies any generation configuration options from kwargs. This diff --git a/examples/model_cache/llada/configuration_llada.py b/examples/model_cache/llada/configuration_llada.py index 3556bdac..a503dc46 100755 --- a/examples/model_cache/llada/configuration_llada.py +++ b/examples/model_cache/llada/configuration_llada.py @@ -11,7 +11,7 @@ from pathlib import Path from typing import ( Any, - Dict, + dict, Iterable, List, Optional, diff --git a/examples/model_cache/llada/modeling_llada.py b/examples/model_cache/llada/modeling_llada.py index babc3dff..74979814 100755 --- a/examples/model_cache/llada/modeling_llada.py +++ b/examples/model_cache/llada/modeling_llada.py @@ -8,7 +8,7 @@ from functools import partial from typing import ( Callable, - Dict, + dict, Iterable, List, NamedTuple, @@ -1059,7 +1059,7 @@ def __init__(self, config: ModelConfig, init_params: bool = True): torch.backends.cuda.enable_flash_sdp(True) torch.backends.cuda.enable_mem_efficient_sdp(False) # this is super slow so make sure torch won't use it - self.transformer = nn.ModuleDict( + self.transformer = nn.Moduledict( dict( wte=nn.Embedding( config.embedding_size or config.vocab_size, config.d_model, device=config.init_device diff --git a/examples/summary.py b/examples/summary.py index fc58c2bd..871814a6 100755 --- a/examples/summary.py +++ b/examples/summary.py @@ -7,7 +7,7 @@ def summarize_profiling(csv_path: str) -> dict: total_nums = {} avgs = {} with open(csv_path, 'r', newline='') as f: - reader = csv.DictReader(f) + reader = csv.dictReader(f) for row in reader: for k, v in row.items(): try: diff --git a/examples/test_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index dad28bde..4c047e37 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -17,7 +17,7 @@ def summarize_profiling(csv_path: str) -> dict: total_nums = {} avgs = {} with open(csv_path, 'r', newline='') as f: - reader = csv.DictReader(f) + reader = csv.dictReader(f) for row in reader: for k, v in row.items(): try: diff --git a/examples/test_dream_dvllm_human_eval.py b/examples/test_dream_dvllm_human_eval.py index 787b0275..2d95f003 100755 --- a/examples/test_dream_dvllm_human_eval.py +++ b/examples/test_dream_dvllm_human_eval.py @@ -16,7 +16,7 @@ def summarize_profiling(csv_path: str) -> dict: total_nums = {} avgs = {} with open(csv_path, 'r', newline='') as f: - reader = csv.DictReader(f) + reader = csv.dictReader(f) for row in reader: for k, v in row.items(): try: diff --git a/examples/test_llada_dvllm_human_eval.py b/examples/test_llada_dvllm_human_eval.py index 5370026e..82127ac1 100755 --- a/examples/test_llada_dvllm_human_eval.py +++ b/examples/test_llada_dvllm_human_eval.py @@ -16,7 +16,7 @@ def summarize_profiling(csv_path: str) -> dict: total_nums = {} avgs = {} with open(csv_path, 'r', newline='') as f: - reader = csv.DictReader(f) + reader = csv.dictReader(f) for row in reader: for k, v in row.items(): try: From 82bef122dc4f6a7805b3e156a45480aeff823d31 Mon Sep 17 00:00:00 2001 From: drewjin Date: Mon, 1 Dec 2025 08:07:17 +0000 Subject: [PATCH 08/23] feat(strategy): implement D2F strategy components including KVCacheManager, ModelRunner, Scheduler, and Sequence classes --- diffulex/strategy/block_diffusion/__init__.py | 14 + .../block_diffusion/attention/metadata.py | 28 + .../block_diffusion/engine/kvcache_manager.py | 44 ++ .../block_diffusion/engine/model_runner.py | 421 +++++++++++++++ .../block_diffusion/engine/scheduler.py | 128 +++++ .../block_diffusion/engine/sequence.py | 484 ++++++++++++++++++ .../strategy/d2f/engine/kvcache_manager.py | 2 +- 7 files changed, 1120 insertions(+), 1 deletion(-) create mode 100644 diffulex/strategy/block_diffusion/__init__.py create mode 100644 diffulex/strategy/block_diffusion/attention/metadata.py create mode 100644 diffulex/strategy/block_diffusion/engine/kvcache_manager.py create mode 100644 diffulex/strategy/block_diffusion/engine/model_runner.py create mode 100644 diffulex/strategy/block_diffusion/engine/scheduler.py create mode 100644 diffulex/strategy/block_diffusion/engine/sequence.py diff --git a/diffulex/strategy/block_diffusion/__init__.py b/diffulex/strategy/block_diffusion/__init__.py new file mode 100644 index 00000000..a83722c7 --- /dev/null +++ b/diffulex/strategy/block_diffusion/__init__.py @@ -0,0 +1,14 @@ +"""D2F strategy component exports.""" +from __future__ import annotations + +from .engine.kvcache_manager import D2FKVCacheManager +from .engine.model_runner import D2FModelRunner +from .engine.scheduler import D2FScheduler +from .engine.sequence import D2FSequence + +__all__ = [ + "D2FKVCacheManager", + "D2FModelRunner", + "D2FScheduler", + "D2FSequence", +] diff --git a/diffulex/strategy/block_diffusion/attention/metadata.py b/diffulex/strategy/block_diffusion/attention/metadata.py new file mode 100644 index 00000000..def3f344 --- /dev/null +++ b/diffulex/strategy/block_diffusion/attention/metadata.py @@ -0,0 +1,28 @@ +import torch + +from dataclasses import dataclass + +from diffulex.attention.metadata import AttnMetaDataBase + + +@dataclass +class D2FAttnMetaData(AttnMetaDataBase): + seq_lens: list[int] = None + seq_lens_ts: torch.Tensor | None = None + d2f_pp: bool = False + block_mask: list[torch.Tensor] | None = None + + +D2F_ATTN_METADATA = D2FAttnMetaData() + +def fetch_d2f_attn_metadata() -> D2FAttnMetaData: + return D2F_ATTN_METADATA + +def set_d2f_attn_metadata() -> None: + # TODO + global D2F_ATTN_METADATA + D2F_ATTN_METADATA = D2FAttnMetaData() + +def reset_d2f_attn_metadata() -> None: + global D2F_ATTN_METADATA + D2F_ATTN_METADATA = D2FAttnMetaData() \ 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 new file mode 100644 index 00000000..70c7f583 --- /dev/null +++ b/diffulex/strategy/block_diffusion/engine/kvcache_manager.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, list + +from diffulex.config import Config +from diffulex.engine.kvcache_manager import AutoKVCacheManager, KVCacheManagerBase + +if TYPE_CHECKING: + from .sequence import D2FSequence + + +@AutoKVCacheManager.register( + "d2f", + aliases=("diffusion_lm",), + is_default=True, +) +class D2FKVCacheManager(KVCacheManagerBase): + 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 + + def may_append(self, seq: "D2FSequence") -> None: + if seq.cached_or_caching_num_tokens == 0: + return + block_table = seq.block_table + if not block_table: + return + last_block = self.blocks[block_table[-1]] + if seq.cached_or_caching_num_tokens // self.block_size == len(seq.block_table): + if last_block.hash == -1: + prev_end_token = seq.cached_or_caching_num_tokens - seq.caching_num_tokens - 1 + prev_block_idx = prev_end_token // self.block_size + if prev_block_idx < seq.num_blocks: + token_ids: list[int] = seq.block(prev_block_idx) + prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 + h = self.compute_hash(token_ids, prefix) + last_block.update(h, token_ids) + 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) diff --git a/diffulex/strategy/block_diffusion/engine/model_runner.py b/diffulex/strategy/block_diffusion/engine/model_runner.py new file mode 100644 index 00000000..481049b0 --- /dev/null +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -0,0 +1,421 @@ +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.d2f.engine.sequence import D2FSequence +from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata +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 + + +@AutoModelRunner.register( + "d2f", + aliases=("diffusion_lm",), + 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) + 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) + + def warmup_model(self): + print("Warming up model...") + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + max_num_batched_tokens, max_model_len = ( + self.config.max_num_batched_tokens, + self.config.max_model_len, + ) + num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs) + test_input_ids = [0] * max_model_len + seqs = [D2FSequence(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() + + 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[D2FSequence]): + input_ids: list[int] = [] + positions: list[int] = [] + cu_seqlens_q = [0] + cu_seqlens_k = [0] + max_seqlen_q = 0 + max_seqlen_k = 0 + 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) + + 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 + cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) + cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k) + + max_seqlen_q = max(seqlen_q, max_seqlen_q) + max_seqlen_k = max(seqlen_k, max_seqlen_k) + + if not seq.block_table: + continue + for i in range(0, seq.num_prompt_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 + else: + end = start + seq.last_block_prompt_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_d2f_attn_metadata( + True, + cu_seqlens_q=cu_seqlens_q_tensor, + cu_seqlens_k=cu_seqlens_k_tensor, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + slot_mapping=slot_mapping_tensor, + context_lens=context_lens_tensor, + block_tables=block_tables, + seqs=seqs, + kv_cache_layout=self.config.kv_cache_layout, + seq_lens=seq_lens, + seq_lens_ts=seq_lens_ts, + ) + return input_ids_tensor, positions_tensor + + def prepare_decode(self, seqs: list[D2FSequence]): + 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 + 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 + 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), + ) + ) + + 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_d2f_attn_metadata( + False, + slot_mapping=slot_mapping_tensor, + context_lens=context_lens_tensor, + cu_seqlens_q=cu_seqlens_q_tensor, + cu_seqlens_k=cu_seqlens_k_tensor, + 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, + kv_cache_layout=self.config.kv_cache_layout, + need_kv_cache_store=need_kv_cache_store, + d2f_pp=True, + ) + return input_ids_tensor, positions_tensor + + @torch.inference_mode() + 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_d2f_attn_metadata() + graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] + 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 + 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 + logits = self.run_model(input_ids, positions, is_prefill) + sample_output = self.sampler(logits, temperatures) if self.rank == 0 else None + reset_d2f_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.") diff --git a/diffulex/strategy/block_diffusion/engine/scheduler.py b/diffulex/strategy/block_diffusion/engine/scheduler.py new file mode 100644 index 00000000..dc650c69 --- /dev/null +++ b/diffulex/strategy/block_diffusion/engine/scheduler.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from diffulex.config import Config +from diffulex.engine.scheduler import AutoScheduler, SchedulerBase +from diffulex.engine.sequence import SequenceBase, SequenceStatus +from .sequence import D2FSequence +from diffulex.layer.sampler import SampleOutputForDiffusionLM + + +@AutoScheduler.register( + "d2f", + aliases=("diffusion_lm",), + is_default=True, +) +class D2FScheduler(SchedulerBase): + def __init__(self, config: Config): + super().__init__(config) + self.diffusion_block_size = config.diffusion_block_size + + def is_finished(self) -> bool: + return not self.waiting and not self.running + + def add(self, seq: D2FSequence) -> None: + self.waiting.append(seq) + + def schedule(self) -> tuple[list[SequenceBase], bool]: + scheduled: list[SequenceBase] = [] + num_seqs = 0 + num_batched_tokens = 0 + while self.waiting and num_seqs < self.max_num_seqs: + seq = self.waiting[0] + projected = len(seq) + seq.diffusion_block_size + if ( + num_batched_tokens + projected > self.max_num_batched_tokens + or not self.block_manager.can_allocate(seq) + ): + break + num_seqs += 1 + self.block_manager.allocate(seq) + num_batched_tokens += projected - seq.num_cached_tokens + seq.status = SequenceStatus.RUNNING + self.waiting.popleft() + self.running.append(seq) + scheduled.append(seq) + if scheduled: + return scheduled, True + + while self.running and num_seqs < self.max_num_seqs: + seq = self.running.popleft() + while not self.block_manager.can_append(seq): + if self.running: + self.preempt(self.running.pop()) + else: + self.preempt(seq) + break + else: + num_seqs += 1 + self.block_manager.may_append(seq) + scheduled.append(seq) + if not scheduled: + diag = { + "phase": "decode", + "waiting": len(self.waiting), + "running": len(self.running), + "max_num_seqs": self.max_num_seqs, + "max_num_batched_tokens": self.max_num_batched_tokens, + "diffusion_block_size": self.diffusion_block_size, + } + candidates = list(self.running)[:3] + list(self.waiting)[:2] + details = [] + for idx, candidate in enumerate(candidates): + try: + can_append = self.block_manager.can_append(candidate) + except Exception: + can_append = "error" + details.append( + f"[{idx}] status={candidate.status.name}, len={len(candidate)}, " + f"diff_block={getattr(candidate, 'diffusion_block_size', '?')}, " + f"new_tokens={getattr(candidate, 'new_tokens', '?')}, " + f"cached={getattr(candidate, 'num_cached_tokens', '?')}, " + f"can_append={can_append}" + ) + raise RuntimeError( + "D2FScheduler: 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: D2FSequence) -> None: + seq.status = SequenceStatus.WAITING + self.block_manager.free(seq) + self.waiting.appendleft(seq) + + def postprocess( + self, + seqs: list[D2FSequence], + sample_output: SampleOutputForDiffusionLM, + ) -> dict[int, int]: + n_diff_steps: dict[int, int] = {} + for seq in seqs: + seq.reset_new_tokens() + seq_id = str(seq.seq_id) + true_ids_map = sample_output.true_local_ids_map.get(seq_id, {}) + accepted_ids_map = sample_output.accepted_ids_map.get(seq_id, {}) + sampled_tokens_map = sample_output.sampled_tokens_map.get(seq_id, {}) + for block_id, accepted_ids in accepted_ids_map.items(): + if not accepted_ids: + continue + diffusion_block = seq.diffusion_blocks[int(block_id)] + sampled_tokens = sampled_tokens_map.get(block_id, []) + true_local_ids = true_ids_map.get(block_id, []) + for true_local_id, accepted_id in zip(true_local_ids, accepted_ids): + token = sampled_tokens[accepted_id] + diffusion_block.modify_token(true_local_id, token) + if ( + (not seq.ignore_eos and token.item() == self.eos) + or seq.num_completion_tokens >= seq.max_tokens + ): + seq.meet_eos = True + if seq.meet_eos and seq.diffusion_blocks[-1].available_to_cache: + seq.status = SequenceStatus.FINISHED + self.block_manager.free(seq) + if seq in self.running: + self.running.remove(seq) + n_diff_steps[seq.seq_id] = seq.n_steps + seq.post_process() + return n_diff_steps diff --git a/diffulex/strategy/block_diffusion/engine/sequence.py b/diffulex/strategy/block_diffusion/engine/sequence.py new file mode 100644 index 00000000..379439ee --- /dev/null +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -0,0 +1,484 @@ +from __future__ import annotations + +import torch + +from dataclasses import dataclass +from enum import Enum, auto + +from diffulex.config import Config +from diffulex.engine.sequence import AutoSequence, SequenceBase +from diffulex.sampling_params import SamplingParams + + +class D2FDiffusionBlockStatus(Enum): + ACTIVE = auto() + TO_CACHE = auto() + IN_CACHE = auto() + + +@dataclass +class D2FDiffusionBlock: + block_id: int = 0 + status: D2FDiffusionBlockStatus = D2FDiffusionBlockStatus.ACTIVE + + global_start_id: int = 0 + global_end_id: int | None = None + cursor: int = 0 + + 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: "D2FSequence" | None = None + pre_block: "D2FDiffusionBlock" | None = None + suf_block: "D2FDiffusionBlock" | None = None + + def __post_init__(self) -> None: + self.global_end_id = self.global_start_id + self.size + + def __getitem__(self, key: int) -> int: + return self.seq[self.global_start_id + key] # type: ignore[index] + + def __len__(self) -> int: + return self.size + + @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 + + @property + def available_to_cache(self) -> bool: + return self.current_complete_ratio == 1.0 + + @property + def is_active(self) -> bool: + return self.status == D2FDiffusionBlockStatus.ACTIVE + + @property + def is_in_cache(self) -> bool: + return self.status == D2FDiffusionBlockStatus.IN_CACHE + + @property + def is_to_cache(self) -> bool: + return self.status == D2FDiffusionBlockStatus.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 + + @property + def add_new_block(self) -> bool: + return self.current_complete_ratio >= self.add_new_block_threshold + + @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] + + @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] + + @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) + 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 = D2FDiffusionBlockStatus.TO_CACHE + + def in_cache(self) -> None: + if self.is_to_cache: + self.status = D2FDiffusionBlockStatus.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( + "d2f", + aliases=("diffusion_lm",), + is_default=True, +) +class D2FSequence(SequenceBase): + """Sequence implementation tailored for diffusion-based decoding.""" + + def __init__( + self, + token_ids: list[int], + sampling_params: SamplingParams = SamplingParams(), + config: Config | None = None, + ): + super().__init__(token_ids, sampling_params) + 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] = [] + 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 = D2FDiffusionBlock( + 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 :] + + @property + def active_blocks(self) -> list[bool]: + return [block.is_active for block in self.diffusion_blocks] + + @property + def to_cache_blocks(self) -> list[bool]: + return [block.is_to_cache 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] + + @property + def num_prompt_blocks(self) -> int: + return (self.input_num_prompt_tokens + self.block_size - 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 + + @property + def updated_or_updating_kv_cache_block_ids(self) -> list[int]: + return [idx for idx, caching in enumerate(self.caching_blocks) if caching] + + @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] + + @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) + + @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) + + @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 + + @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 reset_new_tokens(self) -> None: + self.new_tokens = 0 + + def post_process(self) -> None: + for block in self.diffusion_blocks: + block.cursor = 0 + if block.is_in_cache: + continue + if block.is_to_cache: + block.in_cache() + elif block.is_active: + 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( + D2FDiffusionBlock( + block_id=len(self.diffusion_blocks), + status=D2FDiffusionBlockStatus.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 = D2FDiffusionBlock( + block_id=len(self.diffusion_blocks), + status=D2FDiffusionBlockStatus.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 diff --git a/diffulex/strategy/d2f/engine/kvcache_manager.py b/diffulex/strategy/d2f/engine/kvcache_manager.py index 70c7f583..91a12ec3 100644 --- a/diffulex/strategy/d2f/engine/kvcache_manager.py +++ b/diffulex/strategy/d2f/engine/kvcache_manager.py @@ -41,4 +41,4 @@ def may_append(self, seq: "D2FSequence") -> 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 From cd3844c3e9190878c1fec2bd36ffd8b837155714 Mon Sep 17 00:00:00 2001 From: drewjin Date: Wed, 3 Dec 2025 12:35:55 +0000 Subject: [PATCH 09/23] refactor: update project structure and rename components to align with diffulex; adjust launch configurations and import paths --- .vscode/launch.json | 50 ++++++----- diffulex/__init__.py | 6 +- diffulex/attention/attn_impl.py | 39 ++------ diffulex/attention/ops/kv_cache_kernels.py | 89 +++++++------------ diffulex/engine/dp_worker.py | 3 +- diffulex/engine/kvcache_manager.py | 5 +- diffulex/engine/model_runner.py | 10 +-- diffulex/engine/scheduler.py | 5 +- diffulex/engine/sequence.py | 5 +- diffulex/engine/strategy_registry.py | 26 +++++- diffulex/engine/tp_worker.py | 1 - diffulex/layer/embed_head.py | 8 -- diffulex/layer/sampler.py | 2 +- diffulex/model/__init__.py | 11 +++ diffulex/model/auto_model.py | 7 +- .../model/config/llada/configuration_llada.py | 3 - diffulex/model/dream.py | 5 +- diffulex/model/fast_dllm_v2.py | 4 +- diffulex/model/llada.py | 4 +- diffulex/strategy/block_diffusion/__init__.py | 18 ++-- .../block_diffusion/attention/metadata.py | 22 ++--- .../block_diffusion/engine/kvcache_manager.py | 14 ++- .../block_diffusion/engine/model_runner.py | 34 ++++--- .../block_diffusion/engine/scheduler.py | 18 ++-- .../block_diffusion/engine/sequence.py | 42 ++++----- diffulex/strategy/d2f/attention/metadata.py | 81 ++++++++++++++++- .../strategy/d2f/engine/kvcache_manager.py | 8 +- diffulex/strategy/d2f/engine/model_runner.py | 8 +- diffulex/strategy/d2f/engine/scheduler.py | 6 +- diffulex/strategy/d2f/engine/sequence.py | 6 +- examples/test_dream_dvllm_gsm8k.py | 11 ++- examples/test_dream_model_weight.py | 1 - examples/test_dream_model_weight_fixed.py | 1 - pyproject.toml | 2 +- 34 files changed, 276 insertions(+), 279 deletions(-) create mode 100644 diffulex/model/__init__.py diff --git a/.vscode/launch.json b/.vscode/launch.json index eda65b1a..0224a397 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -4,9 +4,13 @@ // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 "version": "0.2.0", "configurations": [ + + + + { "name": "Python Debugger: Current File", - "type": "debugpy", + "type": "python", "request": "launch", "program": "${file}", "console": "integratedTerminal", @@ -17,7 +21,7 @@ }, { "name": "PyDbg: `Dream` Accelerate Launch Debug", - "type": "debugpy", + "type": "python", "request": "launch", "module": "accelerate.commands.launch", "args": [ @@ -48,10 +52,10 @@ "cwd": "${workspaceFolder}/demo" }, { - "name": "PyDbg: `d2f_vllm` Qwen3", - "type": "debugpy", + "name": "PyDbg: `diffulex` Qwen3", + "type": "python", "request": "launch", - "program": "${workspaceFolder}/demo/test_qwen_dvllm.py", + "program": "${workspaceFolder}/examples/test_qwen_dvllm.py", "console": "integratedTerminal", "env": { // "TRITON_INTERPRET": "1", @@ -59,8 +63,8 @@ } }, { - "name": "PyDbg: `d2f_vllm` Dream `HumanEval`", - "type": "debugpy", + "name": "PyDbg: `diffulex` Dream `HumanEval`", + "type": "python", "request": "launch", "program": "${workspaceFolder}/examples/test_dream_dvllm_human_eval.py", "console": "integratedTerminal", @@ -71,8 +75,8 @@ } }, { - "name": "PyDbg: `d2f_vllm` Dream `GSM8K`", - "type": "debugpy", + "name": "PyDbg: `diffulex` Dream `GSM8K`", + "type": "python", "request": "launch", "program": "${workspaceFolder}/examples/test_dream_dvllm_gsm8k.py", "console": "integratedTerminal", @@ -83,8 +87,8 @@ } }, { - "name": "PyDbg: `d2f_engine` LLaDA `HumanEval`", - "type": "debugpy", + "name": "PyDbg: `diffulex` LLaDA `HumanEval`", + "type": "python", "request": "launch", "program": "${workspaceFolder}/examples/test_llada_dvllm_human_eval.py", "console": "integratedTerminal", @@ -95,10 +99,10 @@ } }, { - "name": "PyDbg: `d2f_vllm` kernel func `load_kvcache_kernel`", - "type": "debugpy", + "name": "PyDbg: `diffulex` kernel func `load_kvcache_kernel`", + "type": "python", "request": "launch", - "program": "${workspaceFolder}/demo/test_dllm_kv_cache_load.py", + "program": "${workspaceFolder}/examples/test_dllm_kv_cache_load.py", "console": "integratedTerminal", "env": { "TRITON_INTERPRET": "1", @@ -106,10 +110,10 @@ } }, { - "name": "PyDbg: `d2f_vllm` kernel func `chunked_prefill_paged_decode`", - "type": "debugpy", + "name": "PyDbg: `diffulex` kernel func `chunked_prefill_paged_decode`", + "type": "python", "request": "launch", - "program": "${workspaceFolder}/demo/test_dllm_decoding_kernel.py", + "program": "${workspaceFolder}/examples/test_dllm_decoding_kernel.py", "console": "integratedTerminal", "env": { "TRITON_INTERPRET": "1", @@ -117,10 +121,10 @@ } }, { - "name": "PyDbg: `d2f_vllm` kernel func `causal_lm_decode_attention_fwd`", - "type": "debugpy", + "name": "PyDbg: `diffulex` kernel func `causal_lm_decode_attention_fwd`", + "type": "python", "request": "launch", - "program": "${workspaceFolder}/demo/test_causal_lm_decoding_kernel.py", + "program": "${workspaceFolder}/examples/test_causal_lm_decoding_kernel.py", "console": "integratedTerminal", "env": { "TRITON_INTERPRET": "1", @@ -128,10 +132,10 @@ } }, { - "name": "PyDbg: `d2f_vllm` kernel func `store_kvcache_kernel_diffusion_lm`", - "type": "debugpy", + "name": "PyDbg: `diffulex` kernel func `store_kvcache_kernel_diffusion_lm`", + "type": "python", "request": "launch", - "program": "${workspaceFolder}/demo/test_dllm_kv_cache_store.py", + "program": "${workspaceFolder}/examples/test_dllm_kv_cache_store.py", "console": "integratedTerminal", "env": { "TRITON_INTERPRET": "1", diff --git a/diffulex/__init__.py b/diffulex/__init__.py index c71384e5..23098a7e 100755 --- a/diffulex/__init__.py +++ b/diffulex/__init__.py @@ -1,2 +1,4 @@ -from diffulex.legacy.llm import LLM -from diffulex.legacy.sampling_params import SamplingParams +from diffulex.diffulex import Diffulex +from diffulex.sampling_params import SamplingParams +# Import strategies to trigger registration +from diffulex import strategy # noqa: F401 diff --git a/diffulex/attention/attn_impl.py b/diffulex/attention/attn_impl.py index 51ac0713..91907ea0 100644 --- a/diffulex/attention/attn_impl.py +++ b/diffulex/attention/attn_impl.py @@ -3,7 +3,6 @@ import torch.nn as nn -from typing import list from functools import lru_cache, partial from einops import rearrange from torch.nn.attention.flex_attention import create_block_mask @@ -15,7 +14,7 @@ store_kvcache_unified_layout, store_kvcache_distinct_layout, load_kvcache, CHECK_STORING, CHECK_LOADING, CHECK_ATTENTION ) -from diffulex.attention import AttnMetaDataBase, fetch_attn_metadata +from diffulex.attention.metadata import AttnMetaDataBase, fetch_attn_metadata class Attention(nn.Module): @@ -25,7 +24,6 @@ def __init__( head_dim, scale, num_kv_heads, - model_type='causal_lm' ): super().__init__() self.num_heads = num_heads @@ -33,8 +31,6 @@ def __init__( self.scale = scale self.num_kv_heads = num_kv_heads self.k_cache = self.v_cache = torch.tensor([]) - self.causal = model_type == 'causal_lm' - self.model_type = model_type is_rtx_xx90 = lambda x: "4090" in x or "3090" in x kernel_options = { "BLOCK_M": 64, @@ -61,27 +57,6 @@ def _mask_mod(batch, head, token_q, token_kv): ) return self._block_mask_cache[cache_key] - @lru_cache(maxsize=32) - def causal_lm_block_mask(self, cum_seq_lens: torch.Tensor, B: int, H: int, Q_LEN: int, KV_LEN: int, device: str): - cache_key = (B, H, Q_LEN, KV_LEN, device) - document_ids = torch.zeros((cum_seq_lens[-1],), dtype=torch.int32, device=device) - start_idx = 0 - for doc_idx, seq_len in enumerate(cum_seq_lens[1:]): - end_idx = seq_len - document_ids[start_idx:end_idx] = doc_idx - start_idx = end_idx - - def _mask_mod(batch, head, token_q, token_kv): - causal_mask = token_q >= token_kv - document_mask = document_ids[token_q] == document_ids[token_kv] - return causal_mask & document_mask - - 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] - def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: list[torch.Tensor] | None = None) -> torch.Tensor: # Reshape @@ -95,18 +70,16 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, # Fast Store KV cache if k_cache.numel() and v_cache.numel(): - if not (self.model_type == 'diffusion_lm' and not attn_metadata.need_kv_cache_store): + if not (not 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, self.model_type, attn_metadata) + 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 if attn_metadata.is_prefill: # Block PK - if attn_metadata.block_tables is not None and self.model_type == 'causal_lm': - k, v = k_cache, v_cache - elif attn_metadata.block_tables is not None and self.model_type == 'diffusion_lm': + if attn_metadata.block_tables is not None: # TODO: Implement Prefix Caching pass @@ -114,9 +87,7 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_t, k_t, v_t = [transpose_fn(t) for t in (q, k, v)] B, H, S, _ = q_t.shape - block_mask_fn = self.causal_lm_block_mask if self.model_type == 'causal_lm' else self.dllm_block_mask - input_obj = attn_metadata.cu_seqlens_q if self.model_type == 'causal_lm' else attn_metadata.block_mask - block_mask = block_mask_fn(input_obj, B, H, S, S, str(q.device)) + 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) else: config = attn_metadata.seqs[0].config diff --git a/diffulex/attention/ops/kv_cache_kernels.py b/diffulex/attention/ops/kv_cache_kernels.py index 1a6fc5ee..41a42ba6 100755 --- a/diffulex/attention/ops/kv_cache_kernels.py +++ b/diffulex/attention/ops/kv_cache_kernels.py @@ -3,11 +3,8 @@ import triton.language as tl -from typing import Tuple -from einops import rearrange +from typing import Any -from diffulex.utils.context import ContextForDiffusionLM -from diffulex.strategy.d2f.sequence import D2FSequence @triton.jit def store_kvcache_kernel_causal_lm( @@ -108,49 +105,32 @@ 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, model_type: str = 'causal_lm', - context: ContextForDiffusionLM = None) -> None: + slot_mapping: torch.Tensor, + context = None) -> None: + # 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 + HDim = HDim_x * x + N = key.shape[0] + assert HDim == key.shape[-1] and NHeads == key.shape[1] + assert N == slot_mapping.numel() - if model_type == 'causal_lm': - # k_cache: [num_blks, blk_sz, h, hdim] - # v_cache: [num_blks, blk_sz, h, hdim] - 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() - store_kvcache_kernel_causal_lm[(N,)]( - key, key.stride(0), - value, value.stride(0), - k_cache, v_cache, slot_mapping, D - ) - else: - # 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 - HDim = HDim_x * x - N = key.shape[0] - assert HDim == key.shape[-1] and NHeads == key.shape[1] - assert N == slot_mapping.numel() - - GRID = (N, ) - store_kvcache_kernel_diffusion_lm_distinct[GRID]( - key, value, - k_cache, v_cache, - slot_mapping, - key.stride(0), value.stride(0), - *k_cache.stride(), *v_cache.stride(), - NHeads, HDim, Blk_sz, - x, HDim * NHeads - ) + GRID = (N, ) + store_kvcache_kernel_diffusion_lm_distinct[GRID]( + key, value, + k_cache, v_cache, + slot_mapping, + key.stride(0), value.stride(0), + *k_cache.stride(), *v_cache.stride(), + NHeads, HDim, Blk_sz, + x, HDim * NHeads + ) def store_kvcache_unified_layout(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, - slot_mapping: torch.Tensor, model_type: str = 'causal_lm', - context: ContextForDiffusionLM = None) -> None: + slot_mapping: torch.Tensor, + context: Any = None) -> None: N, num_heads, head_dim = key.shape D = num_heads * head_dim assert key.stride(-1) == 1 and value.stride(-1) == 1 @@ -158,18 +138,11 @@ def store_kvcache_unified_layout(key: torch.Tensor, value: torch.Tensor, 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()}" - if model_type == 'causal_lm': - store_kvcache_kernel_causal_lm[(N,)]( - key, key.stride(0), - value, value.stride(0), - k_cache, v_cache, slot_mapping, D - ) - elif model_type == 'diffusion_lm': - store_kvcache_kernel_diffusion_lm[(N,)]( - key, key.stride(0), - value, value.stride(0), - k_cache, v_cache, slot_mapping, D - ) + store_kvcache_kernel_diffusion_lm[(N,)]( + key, key.stride(0), + value, value.stride(0), + k_cache, v_cache, slot_mapping, D + ) @triton.jit @@ -276,8 +249,8 @@ def load_kvcache_kernel_kv(k_cache_ptr, v_cache_ptr, def load_kvcache(k_cache: torch.Tensor, v_cache: torch.Tensor, - context: ContextForDiffusionLM, - k_new: torch.Tensor, v_new: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + context: Any, + 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 @@ -329,7 +302,7 @@ def load_kvcache(k_cache: torch.Tensor, v_cache: torch.Tensor, def CHECK_STORING(k_cache: torch.Tensor, v_cache: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - context: ContextForDiffusionLM) -> None: + 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 @@ -365,7 +338,7 @@ def CHECK_STORING(k_cache: torch.Tensor, v_cache: torch.Tensor, 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: ContextForDiffusionLM) -> Tuple[torch.Tensor, 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 = [] diff --git a/diffulex/engine/dp_worker.py b/diffulex/engine/dp_worker.py index b2366c16..0281930b 100755 --- a/diffulex/engine/dp_worker.py +++ b/diffulex/engine/dp_worker.py @@ -7,7 +7,7 @@ import multiprocessing as mp -from typing import list, Any +from typing import Any from multiprocessing.connection import wait as mp_wait from diffulex.config import Config @@ -28,7 +28,6 @@ def _dp_child_entry(config: Config, dp_idx: int, local_devices: list[int], conn) model=config.model, lora_path=config.lora_path, model_name=config.model_name, - model_type=config.model_type, mask_token_id=config.mask_token_id, diffusion_block_size=config.diffusion_block_size, accept_threshold=config.accept_threshold, diff --git a/diffulex/engine/kvcache_manager.py b/diffulex/engine/kvcache_manager.py index 3f7e7661..0b4b54a0 100755 --- a/diffulex/engine/kvcache_manager.py +++ b/diffulex/engine/kvcache_manager.py @@ -119,7 +119,7 @@ class AutoKVCacheManager(DiffulexStrategyRegistry): def from_config(cls, config: Config) -> KVCacheManagerBase: cls._MODULE_MAPPING: dict[str, KVCacheManagerFactory] candidates: list[str] = [] - for attr in ("decoding_strategy", "model_type"): + for attr in ("decoding_strategy",): value = getattr(config, attr, None) if isinstance(value, str) and value: candidates.append(value) @@ -133,6 +133,5 @@ def from_config(cls, config: Config) -> KVCacheManagerBase: available = ", ".join(cls.available_modules()) or "" raise ValueError( "No block manager registered for decoding_strategy=" - f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" - f"'{getattr(config, 'model_type', None)}'. Available block managers: {available}." + f"'{getattr(config, 'decoding_strategy', None)}'. Available block managers: {available}." ) \ No newline at end of file diff --git a/diffulex/engine/model_runner.py b/diffulex/engine/model_runner.py index e5dfd1d9..bdb54a34 100755 --- a/diffulex/engine/model_runner.py +++ b/diffulex/engine/model_runner.py @@ -11,7 +11,7 @@ from diffulex.config import Config from diffulex.layer.sampler import AutoSampler from diffulex.engine.sequence import SequenceBase -from diffulex.model.auto_model import AutoModelForDiffusionLM +from diffulex.model import AutoModelForDiffusionLM from diffulex.engine.strategy_registry import DiffulexStrategyRegistry @@ -19,7 +19,6 @@ class ModelRunnerBase(ABC): """Base class for model runners supporting different model types.""" def __init__(self, config: Config, rank: int, event: Event | list[Event]): self.config = config - self.model_type = config.model_type hf_config = config.hf_config self.block_size = config.kvcache_block_size self.enforce_eager = config.enforce_eager @@ -57,7 +56,7 @@ def __init__(self, config: Config, rank: int, event: Event | list[Event]): shm.unlink() except FileNotFoundError: pass - shm_size = 2**25 if self.model_type == "diffusion_lm" else 2**20 + shm_size = 2**25 self.shm = SharedMemory(name=config.shm_name, create=True, size=shm_size) dist.barrier() else: @@ -187,7 +186,7 @@ class AutoModelRunner(DiffulexStrategyRegistry): def from_config(cls, config: Config, rank: int, event: Event | list[Event]): cls._MODULE_MAPPING: dict[str, RunnerFactory] candidates: list[str] = [] - for attr in ("decoding_strategy", "model_type"): + for attr in ("decoding_strategy",): value = getattr(config, attr, None) if isinstance(value, str) and value: candidates.append(value) @@ -201,6 +200,5 @@ def from_config(cls, config: Config, rank: int, event: Event | list[Event]): available = ", ".join(cls.available_modules()) or "" raise ValueError( "No model runner registered for decoding_strategy=" - f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" - f"'{getattr(config, 'model_type', None)}'. Available runners: {available}." + f"'{getattr(config, 'decoding_strategy', None)}'. Available runners: {available}." ) diff --git a/diffulex/engine/scheduler.py b/diffulex/engine/scheduler.py index 286a7cbd..96d6c93f 100755 --- a/diffulex/engine/scheduler.py +++ b/diffulex/engine/scheduler.py @@ -49,7 +49,7 @@ class AutoScheduler(DiffulexStrategyRegistry): def from_config(cls, config: Config) -> SchedulerBase: cls._MODULE_MAPPING: dict[str, SchedulerFactory] candidates: list[str] = [] - for attr in ("decoding_strategy", "model_type"): + for attr in ("decoding_strategy",): value = getattr(config, attr, None) if isinstance(value, str) and value: candidates.append(value) @@ -63,6 +63,5 @@ def from_config(cls, config: Config) -> SchedulerBase: available = ", ".join(cls.available_modules()) or "" raise ValueError( "No scheduler registered for decoding_strategy=" - f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" - f"'{getattr(config, 'model_type', None)}'. Available schedulers: {available}." + f"'{getattr(config, 'decoding_strategy', None)}'. Available schedulers: {available}." ) \ No newline at end of file diff --git a/diffulex/engine/sequence.py b/diffulex/engine/sequence.py index 57c9aad6..6fd29d4e 100755 --- a/diffulex/engine/sequence.py +++ b/diffulex/engine/sequence.py @@ -86,7 +86,7 @@ def create( ) -> SequenceBase: cls._MODULE_MAPPING: dict[str, SequenceFactory] candidates: list[str] = [] - for attr in ("decoding_strategy", "model_type"): + for attr in ("decoding_strategy",): value = getattr(config, attr, None) if isinstance(value, str) and value: candidates.append(value) @@ -100,6 +100,5 @@ def create( available = ", ".join(cls.available_modules()) or "" raise ValueError( "No sequence registered for decoding_strategy=" - f"'{getattr(config, 'decoding_strategy', None)}' or model_type=" - f"'{getattr(config, 'model_type', None)}'. Available sequences: {available}." + f"'{getattr(config, 'decoding_strategy', None)}'. Available sequences: {available}." ) \ No newline at end of file diff --git a/diffulex/engine/strategy_registry.py b/diffulex/engine/strategy_registry.py index b3cd2ed4..66291e44 100644 --- a/diffulex/engine/strategy_registry.py +++ b/diffulex/engine/strategy_registry.py @@ -9,8 +9,12 @@ class DiffulexStrategyRegistry: """Registry-driven factory for module implementations.""" - _MODULE_MAPPING: dict[str, object] = {} - _DEFAULT_KEY = "__default__" + _DEFAULT_KEY = "__default__" + + def __init_subclass__(cls, **kwargs): + """Initialize a separate _MODULE_MAPPING for each subclass.""" + super().__init_subclass__(**kwargs) + cls._MODULE_MAPPING: dict[str, object] = {} @classmethod def register( @@ -43,8 +47,22 @@ def decorator(factory_fn: object): @classmethod def _register(cls, key: str, factory: object, *, exist_ok: bool) -> None: - if not exist_ok and key in cls._MODULE_MAPPING and cls._MODULE_MAPPING[key] is not factory: - raise ValueError(f"Module '{key}: {fetch_factory_name(factory)}' is already registered.") + # If the same factory is already registered, silently skip (idempotent registration) + if key in cls._MODULE_MAPPING: + existing = cls._MODULE_MAPPING[key] + # Check if it's the same factory object + if existing is factory: + return # Same factory already registered, no-op + # Check if it's the same class by name and module (handles module reload cases) + existing_name = fetch_factory_name(existing) + new_name = fetch_factory_name(factory) + if existing_name == new_name: + return # Same class already registered (possibly from module reload), no-op + if not exist_ok: + raise ValueError( + f"Module '{key}: {new_name}' is already registered as '{existing_name}'. " + f"Use exist_ok=True to override." + ) cls._MODULE_MAPPING[key] = factory @classmethod diff --git a/diffulex/engine/tp_worker.py b/diffulex/engine/tp_worker.py index 464befc8..9978dce9 100755 --- a/diffulex/engine/tp_worker.py +++ b/diffulex/engine/tp_worker.py @@ -19,7 +19,6 @@ def __init__(self, model, **kwargs): config_fields = {field.name for field in fields(Config)} config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} self.config = config = Config(model, **config_kwargs) - self.engine_type = config.model_type self.ps = [] self.events = [] ctx = mp.get_context("spawn") diff --git a/diffulex/layer/embed_head.py b/diffulex/layer/embed_head.py index f46cb694..b96acac7 100755 --- a/diffulex/layer/embed_head.py +++ b/diffulex/layer/embed_head.py @@ -4,8 +4,6 @@ import torch.nn.functional as F import torch.distributed as dist -from diffulex.utils.context import get_context_causal_lm, get_context_diffusion_lm - class VocabParallelEmbedding(nn.Module): @@ -50,7 +48,6 @@ def __init__( num_embeddings: int, embedding_dim: int, bias: bool = False, - model_type: str = 'causal_lm', ): super().__init__(num_embeddings, embedding_dim) if bias: @@ -58,13 +55,8 @@ def __init__( self.bias.weight_loader = self.weight_loader else: self.register_parameter("bias", None) - self.model_type = model_type def forward(self, x: torch.Tensor): - context = get_context_causal_lm() if self.model_type == 'causal_lm' else get_context_diffusion_lm() - if context.is_prefill and self.model_type == 'causal_lm': - last_indices = context.cu_seqlens_q[1:] - 1 - x = x[last_indices].contiguous() logits = F.linear(x, self.weight, self.bias) if self.tp_size > 1: all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] if self.tp_rank == 0 else None diff --git a/diffulex/layer/sampler.py b/diffulex/layer/sampler.py index b1b64a03..7afc8b8d 100644 --- a/diffulex/layer/sampler.py +++ b/diffulex/layer/sampler.py @@ -5,7 +5,7 @@ import torch.distributions as dists from dataclasses import dataclass -from easydict import Easydict as edict +from easydict import EasyDict as edict from diffulex.config import Config from diffulex.attention import fetch_attn_metadata diff --git a/diffulex/model/__init__.py b/diffulex/model/__init__.py new file mode 100644 index 00000000..61e71e9e --- /dev/null +++ b/diffulex/model/__init__.py @@ -0,0 +1,11 @@ +"""Diffulex model package that imports built-in models to trigger registration.""" +from __future__ import annotations + +# 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 + +__all__ = ["dream", "llada", "fast_dllm_v2"] + +from .auto_model import AutoModelForDiffusionLM \ No newline at end of file diff --git a/diffulex/model/auto_model.py b/diffulex/model/auto_model.py index 65e9feef..bb77064a 100755 --- a/diffulex/model/auto_model.py +++ b/diffulex/model/auto_model.py @@ -2,10 +2,9 @@ from typing import Any, Callable -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.llada import LLaDAForDiffusionLM +from diffulex.config import Config +from diffulex.utils.loader import load_model + _NOT_PROVIDED = object() RegistryEntry = tuple[Callable[[Any], Any] | type | None, bool] diff --git a/diffulex/model/config/llada/configuration_llada.py b/diffulex/model/config/llada/configuration_llada.py index 716fd8b3..b9f1f7bd 100644 --- a/diffulex/model/config/llada/configuration_llada.py +++ b/diffulex/model/config/llada/configuration_llada.py @@ -11,11 +11,8 @@ from pathlib import Path from typing import ( Any, - dict, Iterable, - list, Optional, - Tuple, Type, TypeVar, Union, diff --git a/diffulex/model/dream.py b/diffulex/model/dream.py index 96d43648..c7e3ac51 100755 --- a/diffulex/model/dream.py +++ b/diffulex/model/dream.py @@ -83,7 +83,6 @@ def __init__( self.head_dim, self.scaling, self.num_kv_heads, - "diffusion_lm", # Dream uses full attention ) def forward( @@ -189,7 +188,7 @@ def __init__( ) -> None: super().__init__() self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) - self.layers = nn.Modulelist([DreamDecoderLayer(config) + self.layers = nn.ModuleList([DreamDecoderLayer(config) for _ in range(config.num_hidden_layers)]) self.norm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -217,7 +216,7 @@ def __init__( ) -> None: super().__init__() self.model = DreamModel(config) - self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, model_type='diffusion_lm') + 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 diff --git a/diffulex/model/fast_dllm_v2.py b/diffulex/model/fast_dllm_v2.py index 1647559c..d707ebd8 100755 --- a/diffulex/model/fast_dllm_v2.py +++ b/diffulex/model/fast_dllm_v2.py @@ -189,7 +189,7 @@ def __init__( ) -> None: super().__init__() self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) - self.layers = nn.Modulelist([FastdLLMV2DecoderLayer(config) + self.layers = nn.ModuleList([FastdLLMV2DecoderLayer(config) for _ in range(config.num_hidden_layers)]) self.norm = FastdLLMV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -218,7 +218,7 @@ def __init__( ) -> None: super().__init__() self.model = FastdLLMV2Model(config) - self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, model_type='diffusion_lm') + 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 diff --git a/diffulex/model/llada.py b/diffulex/model/llada.py index 5a1b79bb..c3a5243c 100755 --- a/diffulex/model/llada.py +++ b/diffulex/model/llada.py @@ -200,7 +200,7 @@ def __init__( ) blocks = [LLaDABlock(config) for _ in range(config.n_layers)] - self.transformer.update({"blocks": nn.Modulelist(blocks)}) + self.transformer.update({"blocks": nn.ModuleList(blocks)}) if not (self.config.alibi or self.config.rope): self.transformer.update( @@ -245,7 +245,7 @@ def __init__( ) -> None: super().__init__() self.model = LLaDAModel(config) - self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, model_type='diffusion_lm') + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) if getattr(config, 'weight_tying', False): self.lm_head.weight.data = self.model.transformer.wte.weight.data diff --git a/diffulex/strategy/block_diffusion/__init__.py b/diffulex/strategy/block_diffusion/__init__.py index a83722c7..8dabc025 100644 --- a/diffulex/strategy/block_diffusion/__init__.py +++ b/diffulex/strategy/block_diffusion/__init__.py @@ -1,14 +1,14 @@ -"""D2F strategy component exports.""" +"""Block Diffusion strategy component exports.""" from __future__ import annotations -from .engine.kvcache_manager import D2FKVCacheManager -from .engine.model_runner import D2FModelRunner -from .engine.scheduler import D2FScheduler -from .engine.sequence import D2FSequence +from .engine.kvcache_manager import BlockDiffusionKVCacheManager +from .engine.model_runner import BlockDiffusionModelRunner +from .engine.scheduler import BlockDiffusionScheduler +from .engine.sequence import BlockDiffusionSequence __all__ = [ - "D2FKVCacheManager", - "D2FModelRunner", - "D2FScheduler", - "D2FSequence", + "BlockDiffusionKVCacheManager", + "BlockDiffusionModelRunner", + "BlockDiffusionScheduler", + "BlockDiffusionSequence", ] diff --git a/diffulex/strategy/block_diffusion/attention/metadata.py b/diffulex/strategy/block_diffusion/attention/metadata.py index def3f344..a8396b44 100644 --- a/diffulex/strategy/block_diffusion/attention/metadata.py +++ b/diffulex/strategy/block_diffusion/attention/metadata.py @@ -6,23 +6,23 @@ @dataclass -class D2FAttnMetaData(AttnMetaDataBase): +class BlockDiffusionAttnMetaData(AttnMetaDataBase): seq_lens: list[int] = None seq_lens_ts: torch.Tensor | None = None - d2f_pp: bool = False + block_diffusion_pp: bool = False block_mask: list[torch.Tensor] | None = None -D2F_ATTN_METADATA = D2FAttnMetaData() +BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() -def fetch_d2f_attn_metadata() -> D2FAttnMetaData: - return D2F_ATTN_METADATA +def fetch_block_diffusion_attn_metadata() -> BlockDiffusionAttnMetaData: + return BLOCK_DIFFUSION_ATTN_METADATA -def set_d2f_attn_metadata() -> None: +def set_block_diffusion_attn_metadata() -> None: # TODO - global D2F_ATTN_METADATA - D2F_ATTN_METADATA = D2FAttnMetaData() + global BLOCK_DIFFUSION_ATTN_METADATA + BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() -def reset_d2f_attn_metadata() -> None: - global D2F_ATTN_METADATA - D2F_ATTN_METADATA = D2FAttnMetaData() \ No newline at end of file +def reset_block_diffusion_attn_metadata() -> None: + global BLOCK_DIFFUSION_ATTN_METADATA + BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() \ 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 70c7f583..40413949 100644 --- a/diffulex/strategy/block_diffusion/engine/kvcache_manager.py +++ b/diffulex/strategy/block_diffusion/engine/kvcache_manager.py @@ -6,23 +6,19 @@ from diffulex.engine.kvcache_manager import AutoKVCacheManager, KVCacheManagerBase if TYPE_CHECKING: - from .sequence import D2FSequence + from .sequence import BlockDiffusionSequence -@AutoKVCacheManager.register( - "d2f", - aliases=("diffusion_lm",), - is_default=True, -) -class D2FKVCacheManager(KVCacheManagerBase): +@AutoKVCacheManager.register("block_diffusion", is_default=True) +class BlockDiffusionKVCacheManager(KVCacheManagerBase): def __init__(self, config: Config): super().__init__(config) - def can_append(self, seq: "D2FSequence") -> bool: + 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 may_append(self, seq: "D2FSequence") -> None: + def may_append(self, seq: "BlockDiffusionSequence") -> None: if seq.cached_or_caching_num_tokens == 0: return block_table = seq.block_table diff --git a/diffulex/strategy/block_diffusion/engine/model_runner.py b/diffulex/strategy/block_diffusion/engine/model_runner.py index 481049b0..2ff0d8c9 100644 --- a/diffulex/strategy/block_diffusion/engine/model_runner.py +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -8,26 +8,22 @@ from diffulex.config import Config from diffulex.engine.sequence import SequenceBase -from diffulex.strategy.d2f.engine.sequence import D2FSequence +from diffulex.strategy.block_diffusion.engine.sequence import BlockDiffusionSequence from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata 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 +from diffulex.strategy.block_diffusion.attention.metadata import fetch_block_diffusion_attn_metadata, set_block_diffusion_attn_metadata, reset_block_diffusion_attn_metadata -@AutoModelRunner.register( - "d2f", - aliases=("diffusion_lm",), - is_default=True, -) -class D2FModelRunner(ModelRunnerBase): - """Reference implementation of D2F decoding strategy.""" +@AutoModelRunner.register("block_diffusion", is_default=True) +class BlockDiffusionModelRunner(ModelRunnerBase): + """Reference implementation of Block Diffusion decoding strategy.""" def __init__(self, config: Config, rank: int, event: Event | list[Event]): super().__init__(config, rank, event) 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) + set_fetch_fn_for_attn_metadata(fetch_block_diffusion_attn_metadata) def warmup_model(self): print("Warming up model...") @@ -39,7 +35,7 @@ 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 = [D2FSequence(test_input_ids, config=self.config) for _ in range(num_seqs)] + seqs = [BlockDiffusionSequence(test_input_ids, config=self.config) for _ in range(num_seqs)] self.run(seqs, True) for seq in seqs: seq.post_process() @@ -155,7 +151,7 @@ def allocate_kv_cache(self): ) ) - def prepare_prefill(self, seqs: list[D2FSequence]): + def prepare_prefill(self, seqs: list[BlockDiffusionSequence]): input_ids: list[int] = [] positions: list[int] = [] cu_seqlens_q = [0] @@ -227,7 +223,7 @@ def prepare_prefill(self, seqs: list[D2FSequence]): ) ) - set_d2f_attn_metadata( + set_block_diffusion_attn_metadata( True, cu_seqlens_q=cu_seqlens_q_tensor, cu_seqlens_k=cu_seqlens_k_tensor, @@ -243,7 +239,7 @@ def prepare_prefill(self, seqs: list[D2FSequence]): ) return input_ids_tensor, positions_tensor - def prepare_decode(self, seqs: list[D2FSequence]): + def prepare_decode(self, seqs: list[BlockDiffusionSequence]): input_ids: list[int] = [] positions: list[int] = [] cu_seqlens_q = [0] @@ -350,7 +346,7 @@ def get_step(diff_blk, begin_idx): 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_d2f_attn_metadata( + set_block_diffusion_attn_metadata( False, slot_mapping=slot_mapping_tensor, context_lens=context_lens_tensor, @@ -364,7 +360,7 @@ 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, + block_diffusion_pp=True, ) return input_ids_tensor, positions_tensor @@ -373,7 +369,7 @@ def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill 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_d2f_attn_metadata() + context = fetch_block_diffusion_attn_metadata() graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] graph_vars = self.graph_vars for key, value in graph_vars.items(): @@ -401,7 +397,7 @@ def run_verbose(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: 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() + reset_block_diffusion_attn_metadata() return sample_output def run(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: @@ -409,7 +405,7 @@ def run(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: 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_d2f_attn_metadata() + reset_block_diffusion_attn_metadata() return sample_output @torch.inference_mode() diff --git a/diffulex/strategy/block_diffusion/engine/scheduler.py b/diffulex/strategy/block_diffusion/engine/scheduler.py index dc650c69..cc203af6 100644 --- a/diffulex/strategy/block_diffusion/engine/scheduler.py +++ b/diffulex/strategy/block_diffusion/engine/scheduler.py @@ -3,16 +3,12 @@ from diffulex.config import Config from diffulex.engine.scheduler import AutoScheduler, SchedulerBase from diffulex.engine.sequence import SequenceBase, SequenceStatus -from .sequence import D2FSequence +from .sequence import BlockDiffusionSequence from diffulex.layer.sampler import SampleOutputForDiffusionLM -@AutoScheduler.register( - "d2f", - aliases=("diffusion_lm",), - is_default=True, -) -class D2FScheduler(SchedulerBase): +@AutoScheduler.register("block_diffusion", is_default=True) +class BlockDiffusionScheduler(SchedulerBase): def __init__(self, config: Config): super().__init__(config) self.diffusion_block_size = config.diffusion_block_size @@ -20,7 +16,7 @@ def __init__(self, config: Config): def is_finished(self) -> bool: return not self.waiting and not self.running - def add(self, seq: D2FSequence) -> None: + def add(self, seq: BlockDiffusionSequence) -> None: self.waiting.append(seq) def schedule(self) -> tuple[list[SequenceBase], bool]: @@ -81,20 +77,20 @@ def schedule(self) -> tuple[list[SequenceBase], bool]: f"can_append={can_append}" ) raise RuntimeError( - "D2FScheduler: unable to schedule any sequence in decode; " + "BlockDiffusionScheduler: 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: D2FSequence) -> None: + def preempt(self, seq: BlockDiffusionSequence) -> None: seq.status = SequenceStatus.WAITING self.block_manager.free(seq) self.waiting.appendleft(seq) def postprocess( self, - seqs: list[D2FSequence], + seqs: list[BlockDiffusionSequence], sample_output: SampleOutputForDiffusionLM, ) -> dict[int, int]: n_diff_steps: dict[int, int] = {} diff --git a/diffulex/strategy/block_diffusion/engine/sequence.py b/diffulex/strategy/block_diffusion/engine/sequence.py index 379439ee..01d1ff9a 100644 --- a/diffulex/strategy/block_diffusion/engine/sequence.py +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -10,16 +10,16 @@ from diffulex.sampling_params import SamplingParams -class D2FDiffusionBlockStatus(Enum): +class BlockDiffusionBlockStatus(Enum): ACTIVE = auto() TO_CACHE = auto() IN_CACHE = auto() @dataclass -class D2FDiffusionBlock: +class BlockDiffusionBlock: block_id: int = 0 - status: D2FDiffusionBlockStatus = D2FDiffusionBlockStatus.ACTIVE + status: BlockDiffusionBlockStatus = BlockDiffusionBlockStatus.ACTIVE global_start_id: int = 0 global_end_id: int | None = None @@ -33,9 +33,9 @@ class D2FDiffusionBlock: add_new_block_threshold: float = 0.1 complete_threshold: float = 0.9 - seq: "D2FSequence" | None = None - pre_block: "D2FDiffusionBlock" | None = None - suf_block: "D2FDiffusionBlock" | None = None + seq: "BlockDiffusionSequence" | None = None + pre_block: "BlockDiffusionBlock" | None = None + suf_block: "BlockDiffusionBlock" | None = None def __post_init__(self) -> None: self.global_end_id = self.global_start_id + self.size @@ -58,15 +58,15 @@ def available_to_cache(self) -> bool: @property def is_active(self) -> bool: - return self.status == D2FDiffusionBlockStatus.ACTIVE + return self.status == BlockDiffusionBlockStatus.ACTIVE @property def is_in_cache(self) -> bool: - return self.status == D2FDiffusionBlockStatus.IN_CACHE + return self.status == BlockDiffusionBlockStatus.IN_CACHE @property def is_to_cache(self) -> bool: - return self.status == D2FDiffusionBlockStatus.TO_CACHE + return self.status == BlockDiffusionBlockStatus.TO_CACHE @property def pre_block_complete(self) -> bool: @@ -106,11 +106,11 @@ def remaining_length(self, start_idx: int) -> int: def to_cache(self) -> None: if self.available_to_cache and not self.is_in_cache: - self.status = D2FDiffusionBlockStatus.TO_CACHE + self.status = BlockDiffusionBlockStatus.TO_CACHE def in_cache(self) -> None: if self.is_to_cache: - self.status = D2FDiffusionBlockStatus.IN_CACHE + self.status = BlockDiffusionBlockStatus.IN_CACHE def modify_token(self, local_token_id: int, modified_to: int) -> None: if self.seq is None: @@ -121,12 +121,8 @@ def modify_token(self, local_token_id: int, modified_to: int) -> None: self.seq.new_tokens += 1 -@AutoSequence.register( - "d2f", - aliases=("diffusion_lm",), - is_default=True, -) -class D2FSequence(SequenceBase): +@AutoSequence.register("block_diffusion", is_default=True) +class BlockDiffusionSequence(SequenceBase): """Sequence implementation tailored for diffusion-based decoding.""" def __init__( @@ -147,7 +143,7 @@ def __init__( 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.diffusion_blocks: list[BlockDiffusionBlock] = [] self.n_steps = 0 self.input_token_ids: list[int] = [] self.input_num_tokens = 0 @@ -254,7 +250,7 @@ def __setstate__(self, state): self.diffusion_blocks = [] pre_block = None for block_state in state["diffusion_blocks_state"]: - block = D2FDiffusionBlock( + block = BlockDiffusionBlock( block_id=block_state["block_id"], status=block_state["status"], global_start_id=block_state["global_start_id"], @@ -442,9 +438,9 @@ def next_diffusion_step(self, is_prefill: bool = False) -> None: self.input_num_prompt_tokens = self.num_prompt_tokens self.num_prompt_tokens += self.diffusion_block_size self.diffusion_blocks.append( - D2FDiffusionBlock( + BlockDiffusionBlock( block_id=len(self.diffusion_blocks), - status=D2FDiffusionBlockStatus.TO_CACHE, + status=BlockDiffusionBlockStatus.TO_CACHE, global_start_id=0, mask_token_id=self.mask_token_id, size=len(self.input_token_ids), @@ -465,9 +461,9 @@ def next_diffusion_step(self, is_prefill: bool = False) -> None: return added_num_tokens = min(self.diffusion_block_size, remaining) diffusion_seq = [self.mask_token_id] * added_num_tokens - current_block = D2FDiffusionBlock( + current_block = BlockDiffusionBlock( block_id=len(self.diffusion_blocks), - status=D2FDiffusionBlockStatus.ACTIVE, + status=BlockDiffusionBlockStatus.ACTIVE, global_start_id=self.num_tokens, mask_token_id=self.mask_token_id, size=added_num_tokens, diff --git a/diffulex/strategy/d2f/attention/metadata.py b/diffulex/strategy/d2f/attention/metadata.py index def3f344..b9d4d3c7 100644 --- a/diffulex/strategy/d2f/attention/metadata.py +++ b/diffulex/strategy/d2f/attention/metadata.py @@ -1,8 +1,10 @@ import torch +from typing import List from dataclasses import dataclass from diffulex.attention.metadata import AttnMetaDataBase +from diffulex.strategy.d2f.engine.sequence import D2FSequence @dataclass @@ -10,18 +12,89 @@ class D2FAttnMetaData(AttnMetaDataBase): seq_lens: list[int] = None seq_lens_ts: torch.Tensor | None = None d2f_pp: bool = False - block_mask: list[torch.Tensor] | None = None + block_mask: torch.Tensor | None = None + seqs: List[D2FSequence] = None + kv_cache_layout: str = "unified" + need_kv_cache_store: bool = True + 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: + return len(self.seqs) if self.seqs is not None else 0 + D2F_ATTN_METADATA = D2FAttnMetaData() def fetch_d2f_attn_metadata() -> D2FAttnMetaData: return D2F_ATTN_METADATA -def set_d2f_attn_metadata() -> None: - # TODO +def set_d2f_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, + seqs: List[D2FSequence] | None = None, + seq_lens: list[int] | None = None, + 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, +) -> None: global D2F_ATTN_METADATA - D2F_ATTN_METADATA = D2FAttnMetaData() + D2F_ATTN_METADATA = D2FAttnMetaData( + 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, + 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, + ) def reset_d2f_attn_metadata() -> None: global D2F_ATTN_METADATA diff --git a/diffulex/strategy/d2f/engine/kvcache_manager.py b/diffulex/strategy/d2f/engine/kvcache_manager.py index 91a12ec3..119a3f0d 100644 --- a/diffulex/strategy/d2f/engine/kvcache_manager.py +++ b/diffulex/strategy/d2f/engine/kvcache_manager.py @@ -1,6 +1,6 @@ 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 @@ -9,11 +9,7 @@ from .sequence import D2FSequence -@AutoKVCacheManager.register( - "d2f", - aliases=("diffusion_lm",), - is_default=True, -) +@AutoKVCacheManager.register("d2f", is_default=True) class D2FKVCacheManager(KVCacheManagerBase): def __init__(self, config: Config): super().__init__(config) diff --git a/diffulex/strategy/d2f/engine/model_runner.py b/diffulex/strategy/d2f/engine/model_runner.py index 481049b0..6d45f7a8 100644 --- a/diffulex/strategy/d2f/engine/model_runner.py +++ b/diffulex/strategy/d2f/engine/model_runner.py @@ -1,7 +1,7 @@ from __future__ import annotations import time -from typing import list + from multiprocessing.synchronize import Event import torch @@ -14,11 +14,7 @@ from diffulex.strategy.d2f.attention.metadata import fetch_d2f_attn_metadata, set_d2f_attn_metadata, reset_d2f_attn_metadata -@AutoModelRunner.register( - "d2f", - aliases=("diffusion_lm",), - is_default=True, -) +@AutoModelRunner.register("d2f", is_default=True) class D2FModelRunner(ModelRunnerBase): """Reference implementation of D2F decoding strategy.""" diff --git a/diffulex/strategy/d2f/engine/scheduler.py b/diffulex/strategy/d2f/engine/scheduler.py index dc650c69..335b54d1 100644 --- a/diffulex/strategy/d2f/engine/scheduler.py +++ b/diffulex/strategy/d2f/engine/scheduler.py @@ -7,11 +7,7 @@ from diffulex.layer.sampler import SampleOutputForDiffusionLM -@AutoScheduler.register( - "d2f", - aliases=("diffusion_lm",), - is_default=True, -) +@AutoScheduler.register("d2f", is_default=True) class D2FScheduler(SchedulerBase): def __init__(self, config: Config): super().__init__(config) diff --git a/diffulex/strategy/d2f/engine/sequence.py b/diffulex/strategy/d2f/engine/sequence.py index 379439ee..4de92c74 100644 --- a/diffulex/strategy/d2f/engine/sequence.py +++ b/diffulex/strategy/d2f/engine/sequence.py @@ -121,11 +121,7 @@ def modify_token(self, local_token_id: int, modified_to: int) -> None: self.seq.new_tokens += 1 -@AutoSequence.register( - "d2f", - aliases=("diffusion_lm",), - is_default=True, -) +@AutoSequence.register("d2f", is_default=True) class D2FSequence(SequenceBase): """Sequence implementation tailored for diffusion-based decoding.""" diff --git a/examples/test_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index 4c047e37..92880e22 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -9,7 +9,7 @@ from viztracer import VizTracer from transformers import AutoTokenizer -from diffulex.legacy import LLM, SamplingParams +from diffulex import Diffulex, SamplingParams def summarize_profiling(csv_path: str) -> dict: @@ -42,15 +42,14 @@ def summarize_profiling(csv_path: str) -> dict: if __name__ == "__main__": - model = "ckpt/Dream-v0-Base-7B" - LLM = LLM( + model = "/data1/ckpts/Dream-org/Dream-v0-Base-7B" + LLM = Diffulex( model, - lora_path="ckpt/D2F_Dream_Base_7B_Lora", + lora_path="/data1/ckpts/SJTU-Deng-Lab/D2F_Dream_Base_7B_Lora", use_lora=True, model_name="dream", - model_type="diffusion_lm", enforce_eager=True, - data_parallel_size=8, + data_parallel_size=1, tensor_parallel_size=1, gpu_memory_utilization=0.25, max_num_batched_tokens=2048, diff --git a/examples/test_dream_model_weight.py b/examples/test_dream_model_weight.py index 5ad412ee..8455c2b8 100755 --- a/examples/test_dream_model_weight.py +++ b/examples/test_dream_model_weight.py @@ -43,7 +43,6 @@ lora_path=lora_path, use_lora=True, model_name="dream", - model_type="diffusion_lm", enforce_eager=True, tensor_parallel_size=1, accept_threshold=0.95, diff --git a/examples/test_dream_model_weight_fixed.py b/examples/test_dream_model_weight_fixed.py index ae2b61b1..a3d693e7 100755 --- a/examples/test_dream_model_weight_fixed.py +++ b/examples/test_dream_model_weight_fixed.py @@ -45,7 +45,6 @@ lora_path=lora_path, use_lora=True, model_name="dream", - model_type="diffusion_lm", enforce_eager=True, tensor_parallel_size=1, accept_threshold=0.95, diff --git a/pyproject.toml b/pyproject.toml index 59288b0c..188ae078 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ Repository = "https://zhijie-group.github.io/D2fEngine" "Organization" = "https://github.com/zhijie-group" [tool.setuptools.packages.find] -include = ["diffuserve"] +include = ["diffulex"] [[tool.uv.index]] url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" From 1beadc21ff79ee085b07416b6fba4add55ad74ae Mon Sep 17 00:00:00 2001 From: drewjin Date: Fri, 5 Dec 2025 12:33:03 +0000 Subject: [PATCH 10/23] feat(strategy): add block diffusion strategy implementation; fix: fix the launching errors of refactored d2f strategy --- .vscode/launch.json | 1 + diffulex/attention/__init__.py | 15 +- diffulex/attention/attn_impl.py | 51 ++--- diffulex/config.py | 2 +- diffulex/engine/model_runner.py | 2 +- diffulex/layer/sampler.py | 216 ------------------ diffulex/sampler/__init__.py | 10 + diffulex/sampler/auto_sampler.py | 90 ++++++++ diffulex/sampler/base.py | 78 +++++++ diffulex/sampler/dream.py | 83 +++++++ diffulex/sampler/llada.py | 69 ++++++ diffulex/strategy/block_diffusion/__init__.py | 16 +- .../block_diffusion/attention/metadata.py | 64 ++++-- .../block_diffusion/engine/kvcache_manager.py | 12 +- .../block_diffusion/engine/model_runner.py | 32 +-- .../block_diffusion/engine/scheduler.py | 15 +- .../block_diffusion/engine/sequence.py | 90 ++------ diffulex/strategy/d2f/attention/metadata.py | 35 --- diffulex/strategy/d2f/engine/model_runner.py | 6 +- diffulex/strategy/d2f/engine/scheduler.py | 3 +- diffulex/strategy/d2f/engine/sequence.py | 50 +--- examples/test_dream_dvllm_gsm8k.py | 4 +- 22 files changed, 483 insertions(+), 461 deletions(-) delete mode 100644 diffulex/layer/sampler.py create mode 100644 diffulex/sampler/__init__.py create mode 100644 diffulex/sampler/auto_sampler.py create mode 100644 diffulex/sampler/base.py create mode 100644 diffulex/sampler/dream.py create mode 100644 diffulex/sampler/llada.py diff --git a/.vscode/launch.json b/.vscode/launch.json index 0224a397..24c595e4 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -6,6 +6,7 @@ "configurations": [ + { diff --git a/diffulex/attention/__init__.py b/diffulex/attention/__init__.py index e38b5ff8..dfe9af02 100644 --- a/diffulex/attention/__init__.py +++ b/diffulex/attention/__init__.py @@ -1,2 +1,15 @@ 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() \ No newline at end of file diff --git a/diffulex/attention/attn_impl.py b/diffulex/attention/attn_impl.py index 91907ea0..8069abff 100644 --- a/diffulex/attention/attn_impl.py +++ b/diffulex/attention/attn_impl.py @@ -3,18 +3,14 @@ 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.attention.metadata import AttnMetaDataBase, fetch_attn_metadata +from diffulex.attention.metadata import AttnMetaDataBase class Attention(nn.Module): @@ -32,7 +28,7 @@ def __init__( 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 = { + self.kernel_options = { "BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_M1": 32, @@ -40,22 +36,11 @@ def __init__( "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] + + # 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: @@ -64,7 +49,7 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, k = k.view(-1, self.num_kv_heads, self.head_dim) v = v.view(-1, self.num_kv_heads, self.head_dim) - 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" @@ -75,20 +60,17 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, 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 = 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=self.scale, block_table=None) else: config = attn_metadata.seqs[0].config diffusion_block_size = config.diffusion_block_size @@ -111,9 +93,4 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, CHECK_ATTENTION(o, q, k, v, k_cache, v_cache, attn_metadata) # 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 o.view(-1, self.num_heads * self.head_dim).contiguous() \ 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..7fa852da 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 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/sampler/__init__.py b/diffulex/sampler/__init__.py new file mode 100644 index 00000000..8270318d --- /dev/null +++ b/diffulex/sampler/__init__.py @@ -0,0 +1,10 @@ +"""Diffulex sampler package that imports built-in samplers to trigger registration.""" +from __future__ import annotations + +# Import built-in samplers so their registrations run at import time. +from . import dream # noqa: F401 +from . import llada # noqa: F401 + +__all__ = ["dream", "llada"] + +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..d6f81ad7 --- /dev/null +++ b/diffulex/sampler/base.py @@ -0,0 +1,78 @@ +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 + + +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) \ No newline at end of file diff --git a/diffulex/sampler/dream.py b/diffulex/sampler/dream.py new file mode 100644 index 00000000..d8cd9517 --- /dev/null +++ b/diffulex/sampler/dream.py @@ -0,0 +1,83 @@ +import torch + +from dataclasses import dataclass + +from diffulex.sampler.auto_sampler import AutoSampler +from diffulex.sampler.base import SamplerBase, SampleOutputBase + + +@dataclass +class DreamSampleOutputForDiffusionLM(SampleOutputBase): + pass + + +@AutoSampler.register("dream") +class DreamSamplerForDiffusionLM(SamplerBase): + 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 = 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 = {} + 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 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/llada.py b/diffulex/sampler/llada.py new file mode 100644 index 00000000..45befcd0 --- /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 SamplerBase, SampleOutputBase + + +@dataclass +class LLaDASampleOutputForDiffusionLM(SampleOutputBase): + pass + + +@AutoSampler.register("llada") +class LLaDASamplerForDiffusionLM(SamplerBase): + 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/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..f61269fd 100644 --- a/diffulex/strategy/block_diffusion/attention/metadata.py +++ b/diffulex/strategy/block_diffusion/attention/metadata.py @@ -1,28 +1,66 @@ 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): +class BDAttnMetaData(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 + seqs: List[BDSequence] = None + kv_cache_layout: str = "unified" + need_kv_cache_store: bool = True + 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 + + @property + def total_num_seqs(self) -> int: + return len(self.seqs) if self.seqs is not None else 0 + -BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() +BD_ATTN_METADATA = BDAttnMetaData() -def fetch_block_diffusion_attn_metadata() -> BlockDiffusionAttnMetaData: - return BLOCK_DIFFUSION_ATTN_METADATA +def fetch_bd_attn_metadata() -> BDAttnMetaData: + return BD_ATTN_METADATA -def set_block_diffusion_attn_metadata() -> None: - # TODO - global BLOCK_DIFFUSION_ATTN_METADATA - BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() +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, + seqs: List[BDSequence] | None = None, + seq_lens: list[int] | None = None, + seq_lens_ts: torch.Tensor | None = None, + 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, + seq_lens=seq_lens, + seq_lens_ts=seq_lens_ts, + seqs=seqs, + kv_cache_layout=kv_cache_layout, + need_kv_cache_store=need_kv_cache_store, + ) -def reset_block_diffusion_attn_metadata() -> None: - global BLOCK_DIFFUSION_ATTN_METADATA - BLOCK_DIFFUSION_ATTN_METADATA = BlockDiffusionAttnMetaData() \ No newline at end of file +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..5c42789b 100644 --- a/diffulex/strategy/block_diffusion/engine/kvcache_manager.py +++ b/diffulex/strategy/block_diffusion/engine/kvcache_manager.py @@ -1,24 +1,24 @@ 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: + def can_append(self, seq: "BDSequence") -> bool: required = 1 if seq.cached_or_caching_num_tokens % self.block_size == 1 else 0 return len(self.free_block_ids) >= required - 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 +37,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..c69eb2d9 100644 --- a/diffulex/strategy/block_diffusion/engine/model_runner.py +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -1,29 +1,29 @@ 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.strategy.block_diffusion.engine.sequence import BDSequence from diffulex.attention.metadata import set_fetch_fn_for_attn_metadata 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]): + # Set fetch function BEFORE calling super().__init__ + set_fetch_fn_for_attn_metadata(fetch_bd_attn_metadata) + super().__init__(config, rank, event) 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) def warmup_model(self): print("Warming up model...") @@ -35,7 +35,7 @@ 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() @@ -151,7 +151,7 @@ def allocate_kv_cache(self): ) ) - 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] @@ -223,7 +223,7 @@ def prepare_prefill(self, seqs: list[BlockDiffusionSequence]): ) ) - set_block_diffusion_attn_metadata( + set_bd_attn_metadata( True, cu_seqlens_q=cu_seqlens_q_tensor, cu_seqlens_k=cu_seqlens_k_tensor, @@ -239,7 +239,7 @@ def prepare_prefill(self, seqs: list[BlockDiffusionSequence]): ) 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] @@ -346,7 +346,7 @@ def get_step(diff_blk, begin_idx): 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, @@ -360,7 +360,7 @@ 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, - block_diffusion_pp=True, + d2f_pp=True, ) return input_ids_tensor, positions_tensor @@ -369,7 +369,7 @@ def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill 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() + context = fetch_bd_attn_metadata() graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] graph_vars = self.graph_vars for key, value in graph_vars.items(): @@ -397,7 +397,7 @@ def run_verbose(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: 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() + reset_bd_attn_metadata() return sample_output def run(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: @@ -405,7 +405,7 @@ def run(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: 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() + reset_bd_attn_metadata() return sample_output @torch.inference_mode() 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..47e31873 100644 --- a/diffulex/strategy/block_diffusion/engine/sequence.py +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -2,24 +2,24 @@ 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 @@ -33,9 +33,9 @@ class BlockDiffusionBlock: 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 + pre_block: "BDDiffusionBlock" | None = None + suf_block: "BDDiffusionBlock" | None = None def __post_init__(self) -> None: self.global_end_id = self.global_start_id + self.size @@ -58,15 +58,15 @@ def available_to_cache(self) -> bool: @property def is_active(self) -> bool: - return self.status == BlockDiffusionBlockStatus.ACTIVE + return self.status == BDDiffusionBlockStatus.ACTIVE @property def is_in_cache(self) -> bool: - return self.status == BlockDiffusionBlockStatus.IN_CACHE + return self.status == BDDiffusionBlockStatus.IN_CACHE @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: @@ -106,11 +106,11 @@ def remaining_length(self, start_idx: int) -> int: def to_cache(self) -> None: if self.available_to_cache and not self.is_in_cache: - self.status = BlockDiffusionBlockStatus.TO_CACHE + self.status = BDDiffusionBlockStatus.TO_CACHE def in_cache(self) -> None: if self.is_to_cache: - self.status = BlockDiffusionBlockStatus.IN_CACHE + self.status = BDDiffusionBlockStatus.IN_CACHE def modify_token(self, local_token_id: int, modified_to: int) -> None: if self.seq is None: @@ -122,7 +122,7 @@ def modify_token(self, local_token_id: int, modified_to: int) -> None: @AutoSequence.register("block_diffusion", is_default=True) -class BlockDiffusionSequence(SequenceBase): +class BDSequence(SequenceBase): """Sequence implementation tailored for diffusion-based decoding.""" def __init__( @@ -133,17 +133,15 @@ 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_block_size = config.diffusion_block_size - self.block_mask: torch.Tensor | None = None self.meet_eos = False - self.diffusion_blocks: list[BlockDiffusionBlock] = [] + self.diffusion_blocks: list[BDDiffusionBlock] = [] self.n_steps = 0 self.input_token_ids: list[int] = [] self.input_num_tokens = 0 @@ -151,9 +149,9 @@ def __init__( def __repr__(self) -> str: return ( - "SequenceForDiffusionLM(seq_id={seq_id}, status={status}, num_tokens={num_tokens}, " + "BDSequence(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 +159,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 +194,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 +204,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 +225,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 +235,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(): @@ -250,7 +243,7 @@ def __setstate__(self, state): self.diffusion_blocks = [] pre_block = None for block_state in state["diffusion_blocks_state"]: - block = BlockDiffusionBlock( + block = BDDiffusionBlock( block_id=block_state["block_id"], status=block_state["status"], global_start_id=block_state["global_start_id"], @@ -396,40 +389,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: @@ -438,9 +397,9 @@ def next_diffusion_step(self, is_prefill: bool = False) -> None: self.input_num_prompt_tokens = self.num_prompt_tokens self.num_prompt_tokens += self.diffusion_block_size self.diffusion_blocks.append( - BlockDiffusionBlock( + BDDiffusionBlock( block_id=len(self.diffusion_blocks), - status=BlockDiffusionBlockStatus.TO_CACHE, + status=BDDiffusionBlockStatus.TO_CACHE, global_start_id=0, mask_token_id=self.mask_token_id, size=len(self.input_token_ids), @@ -461,9 +420,9 @@ def next_diffusion_step(self, is_prefill: bool = False) -> None: return added_num_tokens = min(self.diffusion_block_size, remaining) diffusion_seq = [self.mask_token_id] * added_num_tokens - current_block = BlockDiffusionBlock( + current_block = BDDiffusionBlock( block_id=len(self.diffusion_blocks), - status=BlockDiffusionBlockStatus.ACTIVE, + status=BDDiffusionBlockStatus.ACTIVE, global_start_id=self.num_tokens, mask_token_id=self.mask_token_id, size=added_num_tokens, @@ -476,5 +435,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/strategy/d2f/attention/metadata.py b/diffulex/strategy/d2f/attention/metadata.py index b9d4d3c7..12d4011d 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,6 @@ 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, ) -> None: global D2F_ATTN_METADATA D2F_ATTN_METADATA = D2FAttnMetaData( @@ -89,8 +56,6 @@ 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, diff --git a/diffulex/strategy/d2f/engine/model_runner.py b/diffulex/strategy/d2f/engine/model_runner.py index 6d45f7a8..a982d382 100644 --- a/diffulex/strategy/d2f/engine/model_runner.py +++ b/diffulex/strategy/d2f/engine/model_runner.py @@ -17,13 +17,13 @@ @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]): + # Set fetch function BEFORE calling super().__init__ + set_fetch_fn_for_attn_metadata(fetch_d2f_attn_metadata) + super().__init__(config, rank, event) 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) def warmup_model(self): print("Warming up model...") diff --git a/diffulex/strategy/d2f/engine/scheduler.py b/diffulex/strategy/d2f/engine/scheduler.py index 335b54d1..f5a4454c 100644 --- a/diffulex/strategy/d2f/engine/scheduler.py +++ b/diffulex/strategy/d2f/engine/scheduler.py @@ -4,7 +4,6 @@ from diffulex.engine.scheduler import AutoScheduler, SchedulerBase from diffulex.engine.sequence import SequenceBase, SequenceStatus from .sequence import D2FSequence -from diffulex.layer.sampler import SampleOutputForDiffusionLM @AutoScheduler.register("d2f", is_default=True) @@ -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..f01a824e 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,13 +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 @@ -153,7 +151,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 +159,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 +194,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 +204,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 +225,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 +235,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(): @@ -396,40 +389,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: @@ -476,5 +435,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/examples/test_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index 92880e22..4eafafcf 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -63,7 +63,7 @@ def summarize_profiling(csv_path: str) -> dict: 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")['test']['question'][:] prompts = [tokenizer.bos_token + FEW_SHOTS + p for p in tqdm(dataset)] output_file = "log/profiles/perf_dvllm_dream_7B.json" @@ -71,7 +71,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() From 6778a360dedd9d55aaf3d4266ac4ee67b882163d Mon Sep 17 00:00:00 2001 From: drewjin Date: Wed, 10 Dec 2025 12:41:21 +0000 Subject: [PATCH 11/23] refactor: update launch configurations to use debugpy and remove unused parameter in model_runner --- .vscode/launch.json | 23 +++++++++----------- diffulex/strategy/d2f/engine/model_runner.py | 1 - examples/test_dream_dvllm_gsm8k.py | 3 ++- 3 files changed, 12 insertions(+), 15 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index 24c595e4..730e3254 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -4,14 +4,11 @@ // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 "version": "0.2.0", "configurations": [ - - - { "name": "Python Debugger: Current File", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${file}", "console": "integratedTerminal", @@ -22,7 +19,7 @@ }, { "name": "PyDbg: `Dream` Accelerate Launch Debug", - "type": "python", + "type": "debugpy", "request": "launch", "module": "accelerate.commands.launch", "args": [ @@ -54,7 +51,7 @@ }, { "name": "PyDbg: `diffulex` Qwen3", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_qwen_dvllm.py", "console": "integratedTerminal", @@ -65,7 +62,7 @@ }, { "name": "PyDbg: `diffulex` Dream `HumanEval`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_dream_dvllm_human_eval.py", "console": "integratedTerminal", @@ -77,7 +74,7 @@ }, { "name": "PyDbg: `diffulex` Dream `GSM8K`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_dream_dvllm_gsm8k.py", "console": "integratedTerminal", @@ -89,7 +86,7 @@ }, { "name": "PyDbg: `diffulex` LLaDA `HumanEval`", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${workspaceFolder}/examples/test_llada_dvllm_human_eval.py", "console": "integratedTerminal", @@ -101,7 +98,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", @@ -112,7 +109,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", @@ -123,7 +120,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", @@ -134,7 +131,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/diffulex/strategy/d2f/engine/model_runner.py b/diffulex/strategy/d2f/engine/model_runner.py index a982d382..543d4d59 100644 --- a/diffulex/strategy/d2f/engine/model_runner.py +++ b/diffulex/strategy/d2f/engine/model_runner.py @@ -360,7 +360,6 @@ 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, ) return input_ids_tensor, positions_tensor diff --git a/examples/test_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index 4eafafcf..03f13b7a 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -63,7 +63,8 @@ def summarize_profiling(csv_path: str) -> dict: tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) sampling_params = SamplingParams(temperature=0.0, max_tokens=256) - dataset = load_dataset("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" From 06bc60542324227ba93ce8f5819702ce994d592d Mon Sep 17 00:00:00 2001 From: drewjin Date: Thu, 11 Dec 2025 18:36:23 +0000 Subject: [PATCH 12/23] feat(strategy): finish block diffusion implementation --- diffulex/engine/sequence.py | 7 +- .../block_diffusion/attention/metadata.py | 20 +- .../block_diffusion/engine/model_runner.py | 157 +----- .../block_diffusion/engine/sequence.py | 463 ++++-------------- 4 files changed, 142 insertions(+), 505 deletions(-) diff --git a/diffulex/engine/sequence.py b/diffulex/engine/sequence.py index 6fd29d4e..00c34ab6 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,16 @@ 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 - + 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/strategy/block_diffusion/attention/metadata.py b/diffulex/strategy/block_diffusion/attention/metadata.py index f61269fd..6f679f0e 100644 --- a/diffulex/strategy/block_diffusion/attention/metadata.py +++ b/diffulex/strategy/block_diffusion/attention/metadata.py @@ -9,19 +9,13 @@ @dataclass class BDAttnMetaData(AttnMetaDataBase): - seq_lens: list[int] = None - seq_lens_ts: torch.Tensor | None = None - seqs: List[BDSequence] = None + diffusion_block_size: int = 32, kv_cache_layout: str = "unified" need_kv_cache_store: bool = True 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 - - @property - def total_num_seqs(self) -> int: - return len(self.seqs) if self.seqs is not None else 0 + 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() @@ -38,9 +32,7 @@ def set_bd_attn_metadata( slot_mapping: torch.Tensor | None = None, context_lens: torch.Tensor | None = None, block_tables: torch.Tensor | None = None, - seqs: List[BDSequence] | None = None, - seq_lens: list[int] | None = None, - seq_lens_ts: torch.Tensor | None = None, + diffusion_block_size: int = 32, kv_cache_layout: str = "unified", need_kv_cache_store: bool = True, ) -> None: @@ -54,9 +46,7 @@ def set_bd_attn_metadata( slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables, - seq_lens=seq_lens, - seq_lens_ts=seq_lens_ts, - seqs=seqs, + diffusion_block_size=diffusion_block_size, kv_cache_layout=kv_cache_layout, need_kv_cache_store=need_kv_cache_store, ) diff --git a/diffulex/strategy/block_diffusion/engine/model_runner.py b/diffulex/strategy/block_diffusion/engine/model_runner.py index c69eb2d9..d3cb45b0 100644 --- a/diffulex/strategy/block_diffusion/engine/model_runner.py +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -161,22 +161,14 @@ def prepare_prefill(self, seqs: list[BDSequence]): 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,41 +180,30 @@ def prepare_prefill(self, seqs: list[BDSequence]): 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_bd_attn_metadata( True, cu_seqlens_q=cu_seqlens_q_tensor, @@ -232,10 +213,8 @@ def prepare_prefill(self, seqs: list[BDSequence]): 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, ) return input_ids_tensor, positions_tensor @@ -246,101 +225,37 @@ def prepare_decode(self, seqs: list[BDSequence]): 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 = len(seq) + seqlen_q = self.diffusion_block_size + seqlen_k = seqlen 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: + for i in range(0, seq.num_blocks_in_active_diffusion_block): + start = seq.block_table[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) @@ -355,12 +270,9 @@ 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, + diffusion_block_size=self.diffusion_block_size, kv_cache_layout=self.config.kv_cache_layout, need_kv_cache_store=need_kv_cache_store, - d2f_pp=True, ) return input_ids_tensor, positions_tensor @@ -383,23 +295,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_bd_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/block_diffusion/engine/sequence.py b/diffulex/strategy/block_diffusion/engine/sequence.py index 47e31873..dbf96a77 100644 --- a/diffulex/strategy/block_diffusion/engine/sequence.py +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -28,14 +28,8 @@ class BDDiffusionBlock: 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: "BDSequence" | None = None - pre_block: "BDDiffusionBlock" | None = None - suf_block: "BDDiffusionBlock" | None = None def __post_init__(self) -> None: self.global_end_id = self.global_start_id + self.size @@ -45,81 +39,39 @@ def __getitem__(self, key: int) -> int: def __len__(self) -> int: return self.size - + @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 == BDDiffusionBlockStatus.ACTIVE - - @property - def is_in_cache(self) -> bool: - return self.status == BDDiffusionBlockStatus.IN_CACHE - + @property def is_to_cache(self) -> bool: 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 - - @property - def add_new_block(self) -> bool: - return self.current_complete_ratio >= self.add_new_block_threshold - - @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 is_in_cache(self) -> bool: + return self.status == BDDiffusionBlockStatus.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] - + def available_to_cache(self) -> bool: + return not self.has_mask_token and self.is_active + @property - def local_mask_token_ids(self) -> list[int]: - return [idx for idx, is_mask in enumerate(self.local_mask_tokens) if is_mask] - + def available_in_cache(self) -> bool: + return self.is_to_cache + @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) - 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 = 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 - + def available_to_add_new_block(self) -> bool: + return self.is_in_cache + @AutoSequence.register("block_diffusion", is_default=True) class BDSequence(SequenceBase): @@ -134,305 +86,102 @@ def __init__( super().__init__(token_ids, sampling_params) if config is None: raise ValueError("BDSequence requires a Config instance.") + self.config = config - 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.meet_eos = False self.diffusion_blocks: list[BDDiffusionBlock] = [] - 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 ( - "BDSequence(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})" - ).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, - ) - - 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, - "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, - "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.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.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 = BDDiffusionBlock( - 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 :] - - @property - def active_blocks(self) -> list[bool]: - return [block.is_active for block in self.diffusion_blocks] - - @property - def to_cache_blocks(self) -> list[bool]: - return [block.is_to_cache 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] - - @property - def num_prompt_blocks(self) -> int: - return (self.input_num_prompt_tokens + self.block_size - 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 - - @property - def updated_or_updating_kv_cache_block_ids(self) -> list[int]: - return [idx for idx, caching in enumerate(self.caching_blocks) if caching] - - @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)] - + self.diffusion_block_size = config.diffusion_block_size + self.mask_token_id = config.mask_token_id + @property - def cached_block_ids(self) -> list[int]: - return [idx for idx, in_cache in enumerate(self.in_cache_blocks) if in_cache] - + def prefix_len_with_padding(self) -> int: + return self.prefix_len + self.pad_prefix_len + @property - def mask_tokens(self) -> list[bool]: - return [token_id == self.mask_token_id for token_id in self.token_ids] - + def diffusion_block_status(self) -> list[BDDiffusionBlockStatus]: + return [block.status for block in self.diffusion_blocks] + @property - def caching_num_tokens(self) -> int: - return sum(block.size for block in self.diffusion_blocks if block.is_to_cache) - + def num_prefix_blocks(self) -> int: + return (self.prefix_len + self.block_size - 1) // self.block_size + @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) - + def prefix_last_block_num_tokens(self) -> int: + return self.prefix_len - (self.num_prefix_blocks - 1) * self.block_size + @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) - + def active_block_token_ids(self) -> list[int]: + return self.diffusion_blocks[-1].token_ids + @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) - - @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 - - @property - def num_diffusion_blocks(self) -> int: - return len(self.diffusion_blocks) - + def num_blocks_in_active_diffusion_block(self) -> int: + return self.diffusion_block_size // self.block_size + 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 reset_new_tokens(self) -> None: - self.new_tokens = 0 - - def post_process(self) -> None: - for block in self.diffusion_blocks: - block.cursor = 0 - if block.is_in_cache: - continue - if block.is_to_cache: - block.in_cache() - elif block.is_active: - if block.available_to_cache: - block.to_cache() - else: - break - - def set_layout(self, layout: str) -> None: - self.kv_cache_layout = layout - - 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 + 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 = block_size - (self.prefix_len % block_size) + if self.prefix_len % block_size == 0: + self.pad_prefix_len = 0 + + # 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 + + 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.TO_CACHE, - global_start_id=0, + status=BDDiffusionBlockStatus.ACTIVE, + global_start_id=self.num_tokens, + size=self.diffusion_block_size, 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, + is_prompt=False, 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 = BDDiffusionBlock( - block_id=len(self.diffusion_blocks), - status=BDDiffusionBlockStatus.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) \ No newline at end of file + ) \ No newline at end of file From 2efbddc07585d35e57d22d8cd307a40df9bbca6f Mon Sep 17 00:00:00 2001 From: drewjin Date: Mon, 15 Dec 2025 11:52:39 +0000 Subject: [PATCH 13/23] fix(strategy): fix d2f strategy launching errors; pref(kernel): implement dllm prefill/decode tilelang kernels --- .gitignore | 3 +- .vscode/launch.json | 14 +- diffulex/attention/attn_impl.py | 40 +- diffulex/attention/metadata.py | 8 + diffulex/attention/ops/__init__.py | 7 - ...chunked_prefill_decoding_unified_kernel.py | 375 ------ diffulex/attention/ops/prefix_prefill.py | 1090 ----------------- .../attention/ops/tilus_decode_attn_dlm.py | 161 --- .../attention/ops/triton_decode_attn_clm.py | 681 ---------- .../attention/ops/triton_decode_attn_dlm.py | 120 -- .../attention/ops/triton_flash_attention.py | 1022 ---------------- diffulex/engine/model_runner.py | 110 +- diffulex/model/__init__.py | 22 +- diffulex/model/fast_dllm_v2.py | 1 - diffulex/sampler/__init__.py | 23 +- diffulex/sampler/fast_dllm_v2.py | 83 ++ diffulex/strategy/__init__.py | 22 +- .../block_diffusion/attention/metadata.py | 7 +- .../block_diffusion/engine/kvcache_manager.py | 3 +- .../block_diffusion/engine/model_runner.py | 212 ++-- .../block_diffusion/engine/sequence.py | 12 +- diffulex/strategy/d2f/attention/metadata.py | 6 + .../strategy/d2f/engine/kvcache_manager.py | 3 +- diffulex/strategy/d2f/engine/model_runner.py | 137 +-- diffulex/strategy/d2f/engine/scheduler.py | 6 +- diffulex/strategy/d2f/engine/sequence.py | 1 - diffulex_kernel/README.md | 0 diffulex_kernel/__init__.py | 2 + diffulex_kernel/python/dllm_flash_attn.py | 403 ++++++ .../python}/kv_cache_kernels.py | 169 +-- examples/ops/prefix_prefill.py | 814 ++++++++++++ examples/test_dream_dvllm_gsm8k.py | 9 +- ...py => test_dream_dvllm_human_eval copy.py} | 0 examples/test_fastdllmv2_diffulex_gsm8k.py | 83 ++ pyproject.toml | 5 +- 35 files changed, 1749 insertions(+), 3905 deletions(-) delete mode 100755 diffulex/attention/ops/__init__.py delete mode 100755 diffulex/attention/ops/chunked_prefill_decoding_unified_kernel.py delete mode 100755 diffulex/attention/ops/prefix_prefill.py delete mode 100755 diffulex/attention/ops/tilus_decode_attn_dlm.py delete mode 100755 diffulex/attention/ops/triton_decode_attn_clm.py delete mode 100755 diffulex/attention/ops/triton_decode_attn_dlm.py delete mode 100755 diffulex/attention/ops/triton_flash_attention.py create mode 100644 diffulex/sampler/fast_dllm_v2.py create mode 100644 diffulex_kernel/README.md create mode 100644 diffulex_kernel/__init__.py create mode 100644 diffulex_kernel/python/dllm_flash_attn.py rename {diffulex/attention/ops => diffulex_kernel/python}/kv_cache_kernels.py (60%) create mode 100644 examples/ops/prefix_prefill.py rename examples/{test_dream_dvllm_human_eval.py => test_dream_dvllm_human_eval copy.py} (100%) create mode 100755 examples/test_fastdllmv2_diffulex_gsm8k.py diff --git a/.gitignore b/.gitignore index bc8329f2..2e7d1668 100755 --- a/.gitignore +++ b/.gitignore @@ -30,4 +30,5 @@ build/ cache/ uv.lock ckpt/ -data/ \ No newline at end of file +data/ +tilelang \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json index 730e3254..1ea0c9bd 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -4,8 +4,6 @@ // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 "version": "0.2.0", "configurations": [ - - { "name": "Python Debugger: Current File", "type": "debugpy", @@ -84,6 +82,18 @@ // "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": "debugpy", diff --git a/diffulex/attention/attn_impl.py b/diffulex/attention/attn_impl.py index 8069abff..aaf03975 100644 --- a/diffulex/attention/attn_impl.py +++ b/diffulex/attention/attn_impl.py @@ -3,12 +3,11 @@ import torch.nn as nn -from flash_attn import flash_attn_varlen_func - -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 @@ -55,42 +54,21 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, # 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) # 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 - o = 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=self.scale, block_table=None) + 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 for decode mode") # Final reshape return o.view(-1, self.num_heads * self.head_dim).contiguous() \ No newline at end of file diff --git a/diffulex/attention/metadata.py b/diffulex/attention/metadata.py index 6b157e00..b71cb0c3 100644 --- a/diffulex/attention/metadata.py +++ b/diffulex/attention/metadata.py @@ -14,6 +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] diff --git a/diffulex/attention/ops/__init__.py b/diffulex/attention/ops/__init__.py deleted file mode 100755 index 579ccbfe..00000000 --- a/diffulex/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/attention/ops/chunked_prefill_decoding_unified_kernel.py b/diffulex/attention/ops/chunked_prefill_decoding_unified_kernel.py deleted file mode 100755 index aed7e060..00000000 --- a/diffulex/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/attention/ops/prefix_prefill.py b/diffulex/attention/ops/prefix_prefill.py deleted file mode 100755 index 03cf31a8..00000000 --- a/diffulex/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/attention/ops/tilus_decode_attn_dlm.py b/diffulex/attention/ops/tilus_decode_attn_dlm.py deleted file mode 100755 index fc7bc03b..00000000 --- a/diffulex/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/attention/ops/triton_decode_attn_clm.py b/diffulex/attention/ops/triton_decode_attn_clm.py deleted file mode 100755 index 71be2616..00000000 --- a/diffulex/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/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/attention/ops/triton_flash_attention.py b/diffulex/attention/ops/triton_flash_attention.py deleted file mode 100755 index 37dd5356..00000000 --- a/diffulex/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/engine/model_runner.py b/diffulex/engine/model_runner.py index 7fa852da..0316dd0c 100755 --- a/diffulex/engine/model_runner.py +++ b/diffulex/engine/model_runner.py @@ -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/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/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/sampler/__init__.py b/diffulex/sampler/__init__.py index 8270318d..8f561b5f 100644 --- a/diffulex/sampler/__init__.py +++ b/diffulex/sampler/__init__.py @@ -1,10 +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 samplers so their registrations run at import time. -from . import dream # noqa: F401 -from . import llada # noqa: F401 +# 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 = [] -__all__ = ["dream", "llada"] +_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/fast_dllm_v2.py b/diffulex/sampler/fast_dllm_v2.py new file mode 100644 index 00000000..2422da73 --- /dev/null +++ b/diffulex/sampler/fast_dllm_v2.py @@ -0,0 +1,83 @@ +import torch + +from dataclasses import dataclass + +from diffulex.sampler.auto_sampler import AutoSampler +from diffulex.sampler.base import SamplerBase, SampleOutputBase + + +@dataclass +class FastdLLMV2SampleOutputForDiffusionLM(SampleOutputBase): + pass + + +@AutoSampler.register("fast_dllm_v2") +class FastdLLMV2SamplerForDiffusionLM(SamplerBase): + 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 = 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 = {} + 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 FastdLLMV2SampleOutputForDiffusionLM( + 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/attention/metadata.py b/diffulex/strategy/block_diffusion/attention/metadata.py index 6f679f0e..436a69ee 100644 --- a/diffulex/strategy/block_diffusion/attention/metadata.py +++ b/diffulex/strategy/block_diffusion/attention/metadata.py @@ -9,7 +9,6 @@ @dataclass class BDAttnMetaData(AttnMetaDataBase): - diffusion_block_size: int = 32, kv_cache_layout: str = "unified" need_kv_cache_store: bool = True @@ -32,7 +31,10 @@ def set_bd_attn_metadata( 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 = "varlen", + attn_type: str = "full_attention", kv_cache_layout: str = "unified", need_kv_cache_store: bool = True, ) -> None: @@ -46,9 +48,12 @@ def set_bd_attn_metadata( 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: diff --git a/diffulex/strategy/block_diffusion/engine/kvcache_manager.py b/diffulex/strategy/block_diffusion/engine/kvcache_manager.py index 5c42789b..9659c10b 100644 --- a/diffulex/strategy/block_diffusion/engine/kvcache_manager.py +++ b/diffulex/strategy/block_diffusion/engine/kvcache_manager.py @@ -15,8 +15,7 @@ def __init__(self, config: Config): super().__init__(config) def can_append(self, seq: "BDSequence") -> 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: "BDSequence") -> None: if seq.cached_or_caching_num_tokens == 0: diff --git a/diffulex/strategy/block_diffusion/engine/model_runner.py b/diffulex/strategy/block_diffusion/engine/model_runner.py index d3cb45b0..37cc5362 100644 --- a/diffulex/strategy/block_diffusion/engine/model_runner.py +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -18,13 +18,12 @@ class BDModelRunner(ModelRunnerBase): """Reference implementation of Block Diffusion decoding strategy.""" def __init__(self, config: Config, rank: int, event: Event | list[Event]): - # Set fetch function BEFORE calling super().__init__ set_fetch_fn_for_attn_metadata(fetch_bd_attn_metadata) - - super().__init__(config, rank, event) self.diffusion_block_size = config.diffusion_block_size self.mask_token_id = config.mask_token_id - + + super().__init__(config, rank, event) + def warmup_model(self): print("Warming up model...") torch.cuda.empty_cache() @@ -41,116 +40,6 @@ def warmup_model(self): 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 - ) - ) - def prepare_prefill(self, seqs: list[BDSequence]): input_ids: list[int] = [] positions: list[int] = [] @@ -215,6 +104,8 @@ def prepare_prefill(self, seqs: list[BDSequence]): block_tables=block_tables, diffusion_block_size=self.diffusion_block_size, kv_cache_layout=self.config.kv_cache_layout, + attn_type="block_attention", + decode_mode="static", ) return input_ids_tensor, positions_tensor @@ -270,6 +161,7 @@ def prepare_decode(self, seqs: list[BDSequence]): max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, block_tables=block_tables, + page_block_size=self.config.kvcache_page_size, diffusion_block_size=self.diffusion_block_size, kv_cache_layout=self.config.kv_cache_layout, need_kv_cache_store=need_kv_cache_store, @@ -280,20 +172,24 @@ def prepare_decode(self, seqs: list[BDSequence]): 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) + num_tokens = input_ids.size(0) context = fetch_bd_attn_metadata() - graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] + 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]) + 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) @@ -305,8 +201,70 @@ def run(self, seqs: list[SequenceBase], is_prefill: bool) -> list[int]: @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/sequence.py b/diffulex/strategy/block_diffusion/engine/sequence.py index dbf96a77..7bc2c765 100644 --- a/diffulex/strategy/block_diffusion/engine/sequence.py +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -116,6 +116,14 @@ def active_block_token_ids(self) -> list[int]: def num_blocks_in_active_diffusion_block(self) -> int: return self.diffusion_block_size // self.block_size + @property + def cached_num_tokens(self) -> int: + return sum(block.size for block in self.diffusion_blocks if block.is_in_cache) + + @property + def cached_or_caching_num_tokens(self) -> int: + return sum(block.size for block in self.diffusion_blocks if block.is_to_cache or block.is_in_cache) + def diffusion_decoding_inputs(self) -> tuple[list[int], list[int], int]: return ( self.active_block_token_ids, @@ -133,9 +141,7 @@ def init_diffusion_blocks(self) -> None: # Calculate prefix blocks and padding num_prefix_blocks = self.prefix_len // block_size - self.pad_prefix_len = block_size - (self.prefix_len % block_size) - if self.prefix_len % block_size == 0: - self.pad_prefix_len = 0 + 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) diff --git a/diffulex/strategy/d2f/attention/metadata.py b/diffulex/strategy/d2f/attention/metadata.py index 12d4011d..523daf4d 100644 --- a/diffulex/strategy/d2f/attention/metadata.py +++ b/diffulex/strategy/d2f/attention/metadata.py @@ -43,6 +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, + diffusion_block_size: int = 32, + decode_mode: str = "varlen", + attn_type: str = "full_attention", ) -> None: global D2F_ATTN_METADATA D2F_ATTN_METADATA = D2FAttnMetaData( @@ -59,6 +62,9 @@ def set_d2f_attn_metadata( 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 543d4d59..81a2a84b 100644 --- a/diffulex/strategy/d2f/engine/model_runner.py +++ b/diffulex/strategy/d2f/engine/model_runner.py @@ -18,12 +18,12 @@ class D2FModelRunner(ModelRunnerBase): """Reference implementation of D2F decoding strategy.""" def __init__(self, config: Config, rank: int, event: Event | list[Event]): - # Set fetch function BEFORE calling super().__init__ set_fetch_fn_for_attn_metadata(fetch_d2f_attn_metadata) - super().__init__(config, rank, event) self.diffusion_block_size = config.diffusion_block_size self.mask_token_id = config.mask_token_id + + super().__init__(config, rank, event) def warmup_model(self): print("Warming up model...") @@ -41,116 +41,6 @@ def warmup_model(self): 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 - ) - ) - def prepare_prefill(self, seqs: list[D2FSequence]): input_ids: list[int] = [] positions: list[int] = [] @@ -236,6 +126,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,6 +253,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, + diffusion_block_size=self.diffusion_block_size, + decode_mode="varlen", + attn_type="full_attention", ) return input_ids_tensor, positions_tensor @@ -382,23 +278,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 f5a4454c..a4b8f29d 100644 --- a/diffulex/strategy/d2f/engine/scheduler.py +++ b/diffulex/strategy/d2f/engine/scheduler.py @@ -2,7 +2,7 @@ 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 @@ -18,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: diff --git a/diffulex/strategy/d2f/engine/sequence.py b/diffulex/strategy/d2f/engine/sequence.py index f01a824e..78ca0f0c 100644 --- a/diffulex/strategy/d2f/engine/sequence.py +++ b/diffulex/strategy/d2f/engine/sequence.py @@ -434,5 +434,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) \ No newline at end of file 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/dllm_flash_attn.py b/diffulex_kernel/python/dllm_flash_attn.py new file mode 100644 index 00000000..63bbee43 --- /dev/null +++ b/diffulex_kernel/python/dllm_flash_attn.py @@ -0,0 +1,403 @@ +import torch +import tilelang +import tilelang.language as T + +from flash_attn import flash_attn_varlen_func + +from diffulex_kernel.python.kv_cache_kernels import load_kvcache +from diffulex.attention.metadata import AttnMetaDataBase + + +@tilelang.jit(out_idx=[6], 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], 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 and kv_block_idx == loop_range - 1: + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + num_diffusion_blocks = T.min(i // DIFFUSION_BLOCK_SIZE + 1, BLOCK_M // DIFFUSION_BLOCK_SIZE) + acc_score[i, j] = T.if_then_else( + (kv_block_idx * BLOCK_N + j >= kv_block_idx * BLOCK_N + num_diffusion_blocks * DIFFUSION_BLOCK_SIZE) 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_block_idx * BLOCK_M, head_idx, d_idx] = O_shared[i, d_idx] + + return kernel + + +@tilelang.jit(out_idx=[10], 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], 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 + T.device_assert(cur_q_seqlen == DIFFUSION_BLOCK_SIZE, "cur_q_seqlen must be equal to DIFFUSION_BLOCK_SIZE") + T.device_assert(cur_kv_seqlen == DIFFUSION_BLOCK_SIZE, "cur_kv_seqlen must be equal to DIFFUSION_BLOCK_SIZE") + + 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)) + + # Fusion of Q/KVCache Cross-Attention and QKV Self-Attention (Full-Attention) + 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 == -1: + 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( + (q_start_idx + i >= cur_q_seqlen or + kv_start_idx + 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) + 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) + T.gemm(acc_score_kv_cast, V_shared, acc_output, policy=T.GemmWarpPolicy.FullRow) + + break + + 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( + (q_start_idx + 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) + for i, j in T.Parallel(BLOCK_M, HEAD_DIM): + acc_output[i, j] *= scores_scale[i] + + # Compute attention output + T.copy(V_Cache[page_block_idx_global, :, kv_head_idx, :], V_shared) + T.gemm(acc_score_kvcache_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_start_idx < 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": + attn_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], + True, + DIFFUSION_BLOCK_SIZE=attn_metadata.diffusion_block_size, + BLOCK_M=128, + BLOCK_N=128, + NUM_STAGES=2, + NUM_THREADS=256 + ) + return attn_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": + attn_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", + DIFFUSION_BLOCK_SIZE=attn_metadata.diffusion_block_size, + MAX_SEQ_NUM_BLOCKS=attn_metadata.block_tables.shape[1], + PAGE_BLOCK_SIZE=attn_metadata.page_block_size, + BLOCK_M=128, + BLOCK_N=128, + NUM_STAGES=2, + NUM_THREADS=256 + ) + return attn_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/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_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index 03f13b7a..66056272 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -58,7 +58,8 @@ 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) @@ -85,6 +86,6 @@ def summarize_profiling(csv_path: str) -> dict: 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 + 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_dream_dvllm_human_eval.py b/examples/test_dream_dvllm_human_eval copy.py similarity index 100% rename from examples/test_dream_dvllm_human_eval.py rename to examples/test_dream_dvllm_human_eval copy.py diff --git a/examples/test_fastdllmv2_diffulex_gsm8k.py b/examples/test_fastdllmv2_diffulex_gsm8k.py new file mode 100755 index 00000000..2dd4f3ff --- /dev/null +++ b/examples/test_fastdllmv2_diffulex_gsm8k.py @@ -0,0 +1,83 @@ +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" + +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 = [tokenizer.apply_chat_template(p, tokenize=False) for p 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/pyproject.toml b/pyproject.toml index 188ae078..826a717f 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,10 @@ Repository = "https://zhijie-group.github.io/D2fEngine" "Organization" = "https://github.com/zhijie-group" [tool.setuptools.packages.find] -include = ["diffulex"] +include = [ + "diffulex", + "diffulex_kernel", +] [[tool.uv.index]] url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" From c115d15aade5c74171a558929414d82abdbfd29b Mon Sep 17 00:00:00 2001 From: drewjin Date: Mon, 15 Dec 2025 15:13:22 +0000 Subject: [PATCH 14/23] fix(kernel): fix unable to compile bug and add autotuning into dllm kernels --- .gitignore | 4 +- diffulex/sampler/fast_dllm_v2.py | 13 +- .../block_diffusion/attention/metadata.py | 1 + .../block_diffusion/engine/model_runner.py | 2 +- .../block_diffusion/engine/sequence.py | 40 +++++- diffulex_kernel/python/auto_tuner.py | 54 +++++++ diffulex_kernel/python/dllm_flash_attn.py | 132 +++++++++++++----- 7 files changed, 204 insertions(+), 42 deletions(-) create mode 100644 diffulex_kernel/python/auto_tuner.py diff --git a/.gitignore b/.gitignore index 2e7d1668..2712c2f8 100755 --- a/.gitignore +++ b/.gitignore @@ -31,4 +31,6 @@ cache/ uv.lock ckpt/ data/ -tilelang \ No newline at end of file +tilelang +autotuner.log +Fast-dLLM \ No newline at end of file diff --git a/diffulex/sampler/fast_dllm_v2.py b/diffulex/sampler/fast_dllm_v2.py index 2422da73..38b5b9ad 100644 --- a/diffulex/sampler/fast_dllm_v2.py +++ b/diffulex/sampler/fast_dllm_v2.py @@ -4,6 +4,7 @@ from diffulex.sampler.auto_sampler import AutoSampler from diffulex.sampler.base import SamplerBase, SampleOutputBase +from diffulex.engine.sequence import SequenceBase @dataclass @@ -26,11 +27,13 @@ def _shift_logits(self, logits, last_logit=None): shifted_logits[0, ...] = 1.0 return shifted_logits - def forward(self, logits: torch.Tensor, temperatures: torch.Tensor, + def forward(self, seqs: list[SequenceBase], 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) + 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 = {} @@ -38,7 +41,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) + shifted_logits = self._shift_logits(seq_logits, seq.cached_or_caching_num_tokens - 1) for block_id, block in enumerate(seq.diffusion_blocks): if not block.is_active or sum(block.local_mask_tokens) == 0: continue diff --git a/diffulex/strategy/block_diffusion/attention/metadata.py b/diffulex/strategy/block_diffusion/attention/metadata.py index 436a69ee..1bda15d7 100644 --- a/diffulex/strategy/block_diffusion/attention/metadata.py +++ b/diffulex/strategy/block_diffusion/attention/metadata.py @@ -9,6 +9,7 @@ @dataclass class BDAttnMetaData(AttnMetaDataBase): + seqs: List[BDSequence] = None kv_cache_layout: str = "unified" need_kv_cache_store: bool = True diff --git a/diffulex/strategy/block_diffusion/engine/model_runner.py b/diffulex/strategy/block_diffusion/engine/model_runner.py index 37cc5362..67a95a91 100644 --- a/diffulex/strategy/block_diffusion/engine/model_runner.py +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -195,7 +195,7 @@ 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 + sample_output = self.sampler(seqs, logits, temperatures) if self.rank == 0 else None reset_bd_attn_metadata() return sample_output diff --git a/diffulex/strategy/block_diffusion/engine/sequence.py b/diffulex/strategy/block_diffusion/engine/sequence.py index 7bc2c765..bd482909 100644 --- a/diffulex/strategy/block_diffusion/engine/sequence.py +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -40,6 +40,14 @@ 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 + @property def token_ids(self) -> list[int]: return self.seq.token_ids[self.global_start_id: self.global_end_id] @@ -72,6 +80,23 @@ def available_in_cache(self) -> bool: 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.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) + return [mask_id + offset for mask_id in self.local_mask_token_ids] + @AutoSequence.register("block_diffusion", is_default=True) class BDSequence(SequenceBase): @@ -190,4 +215,17 @@ def next_diffusion_step(self) -> None: is_prompt=False, seq=self, ) - ) \ No newline at end of file + ) + + def post_process(self) -> None: + for block in self.diffusion_blocks: + block.cursor = 0 + if block.is_in_cache: + continue + if block.is_to_cache: + block.in_cache() + elif block.is_active: + if block.available_to_cache: + block.to_cache() + else: + break \ 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..23798dd9 --- /dev/null +++ b/diffulex_kernel/python/auto_tuner.py @@ -0,0 +1,54 @@ +import torch +import itertools + +def get_heuristic_config() -> dict: + # Get CUDA device properties + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + device = torch.cuda.current_device() + sm_major, sm_minor = torch.cuda.get_device_capability(device) + sm_version = sm_major * 10 + sm_minor + if sm_version >= 80 and sm_version < 90: + return { + "BLOCK_M": 128, + "BLOCK_N": 256, + "NUM_STAGES": 2, + "NUM_THREADS": 128, + } + elif sm_version >= 90 and sm_version < 100: + return { + "BLOCK_M": 128, + "BLOCK_N": 256, + "NUM_STAGES": 3, + "NUM_THREADS": 256, + } + else: + return { + "BLOCK_M": 128, + "BLOCK_N": 256, + "NUM_STAGES": 0, + "NUM_THREADS": 128, + } + +def build_configs(): + BLOCK_M_LIST = [64, 128, 256] + BLOCK_N_LIST = [64, 128, 256] + NUM_STAGES_LIST = [0, 1, 2, 3] + 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 index 63bbee43..a149a74f 100644 --- a/diffulex_kernel/python/dllm_flash_attn.py +++ b/diffulex_kernel/python/dllm_flash_attn.py @@ -2,13 +2,26 @@ import tilelang import tilelang.language as T +from tilelang.autotuner import set_autotune_inputs from flash_attn import flash_attn_varlen_func +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 -@tilelang.jit(out_idx=[6], pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}) +# Kernel缓存，避免重复autotune和编译 +_prefill_kernel_cache = {} +_decode_kernel_cache = {} + + +@tilelang.autotune( + configs=build_configs() +) +@tilelang.jit( + out_idx=[6], + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}, +) def dllm_flash_attn_prefill_kernel( NUM_SEQS: int, NUM_GROUPS: int, @@ -50,7 +63,7 @@ def kernel( 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], 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) @@ -144,8 +157,11 @@ def kernel( return kernel - -@tilelang.jit(out_idx=[10], pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}) +@tilelang.autotune(configs=build_configs()) +@tilelang.jit( + out_idx=[10], + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}, +) def dllm_flash_attn_decode_kernel( NUM_SEQS: int, NUM_GROUPS: int, @@ -201,7 +217,7 @@ def kernel( 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], 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) @@ -309,8 +325,8 @@ def kernel( acc_output[i, j] *= scores_scale[i] # Compute attention output - T.copy(V_Cache[page_block_idx_global, :, kv_head_idx, :], V_shared) - T.gemm(acc_score_kvcache_cast, V_shared, acc_output, policy=T.GemmWarpPolicy.FullRow) + 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) for i, j in T.Parallel(BLOCK_M, HEAD_DIM): acc_output[i, j] /= log_sum[i] @@ -338,20 +354,41 @@ def dllm_flash_attn_prefill( softmax_scale=scale, block_table=None ) elif attn_metadata.attn_type == "block_attention": - attn_kernel = dllm_flash_attn_prefill_kernel( + # 创建缓存键，基于kernel的参数 + cache_key = ( attn_metadata.num_seqs, - q.shape[1] // k.shape[1], - q.shape[0], - k.shape[0], - q.shape[1], - q.shape[2], - True, - DIFFUSION_BLOCK_SIZE=attn_metadata.diffusion_block_size, - BLOCK_M=128, - BLOCK_N=128, - NUM_STAGES=2, - NUM_THREADS=256 + q.shape[1] // k.shape[1], # NUM_GROUPS + q.shape[0], # Q_LEN + k.shape[0], # KV_LEN + q.shape[1], # NUM_HEADS + q.shape[2], # HEAD_DIM + attn_metadata.diffusion_block_size, # DIFFUSION_BLOCK_SIZE ) + + # 检查缓存 + if cache_key not in _prefill_kernel_cache: + # 使用set_autotune_inputs来触发autotune + # 这会在第一次调用时为所有配置测试性能并选择最佳配置 + with set_autotune_inputs([ + q, k, v, + attn_metadata.cu_seqlens_q, + attn_metadata.cu_seqlens_k, + attn_metadata.max_seqlen_q, + ]): + attn_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], + True, + attn_metadata.diffusion_block_size, + ) + _prefill_kernel_cache[cache_key] = attn_kernel + else: + attn_kernel = _prefill_kernel_cache[cache_key] + return attn_kernel( q, k, v, attn_metadata.cu_seqlens_q, @@ -370,23 +407,50 @@ def dllm_flash_attn_decode( attn_metadata: AttnMetaDataBase ) -> torch.Tensor: if attn_metadata.decode_mode == "static": - attn_kernel = dllm_flash_attn_decode_kernel( + # 创建缓存键，基于kernel的参数 + cache_key = ( 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", - DIFFUSION_BLOCK_SIZE=attn_metadata.diffusion_block_size, - MAX_SEQ_NUM_BLOCKS=attn_metadata.block_tables.shape[1], - PAGE_BLOCK_SIZE=attn_metadata.page_block_size, - BLOCK_M=128, - BLOCK_N=128, - NUM_STAGES=2, - NUM_THREADS=256 + q.shape[1] // k.shape[1], # NUM_GROUPS + k_cache.shape[0], # NUM_PAGE_BLOCKS + q.shape[0], # Q_LEN + k.shape[0], # KV_LEN + q.shape[1], # NUM_HEADS + q.shape[2], # HEAD_DIM + attn_metadata.attn_type == "block_attention", # IS_BLOCK_ATTN + attn_metadata.diffusion_block_size, # DIFFUSION_BLOCK_SIZE + attn_metadata.block_tables.shape[1], # MAX_SEQ_NUM_BLOCKS + attn_metadata.page_block_size, # PAGE_BLOCK_SIZE ) + + # 检查缓存 + if cache_key not in _decode_kernel_cache: + # 使用set_autotune_inputs来触发autotune + # 这会在第一次调用时为所有配置测试性能并选择最佳配置 + with set_autotune_inputs([ + 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, + ]): + attn_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, + ) + _decode_kernel_cache[cache_key] = attn_kernel + else: + attn_kernel = _decode_kernel_cache[cache_key] + return attn_kernel( q, k, v, k_cache, v_cache, attn_metadata.block_tables, From 5c0253cd54af759604f321c31b4f2c6b06e3d190 Mon Sep 17 00:00:00 2001 From: luozixin2 Date: Wed, 17 Dec 2025 07:57:07 +0000 Subject: [PATCH 15/23] feat: add SDAR model support for fast_dllm_v2 branch - Add SDARConfig in diffulex/model/config/sdar/configuration_sdar.py - Implement SDARForDiffusionLM with native Diffulex KV cache integration - Register SDAR model to AutoModelForDiffusionLM - Update model __init__.py to trigger SDAR registration --- diffulex/model/__init__.py | 3 +- .../model/config/sdar/configuration_sdar.py | 78 +++++++ diffulex/model/sdar.py | 210 ++++++++++++++++++ 3 files changed, 290 insertions(+), 1 deletion(-) create mode 100644 diffulex/model/config/sdar/configuration_sdar.py diff --git a/diffulex/model/__init__.py b/diffulex/model/__init__.py index 61e71e9e..dee516c4 100644 --- a/diffulex/model/__init__.py +++ b/diffulex/model/__init__.py @@ -5,7 +5,8 @@ from . import dream # noqa: F401 from . import llada # noqa: F401 from . import fast_dllm_v2 # noqa: F401 +from . import sdar # noqa: F401 -__all__ = ["dream", "llada", "fast_dllm_v2"] +__all__ = ["dream", "llada", "fast_dllm_v2", "sdar"] 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/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", +] From 9af55b9449cf54fffb87af56ed090842b464e72d Mon Sep 17 00:00:00 2001 From: luozixin2 Date: Wed, 17 Dec 2025 08:50:48 +0000 Subject: [PATCH 16/23] feat: add SDAR model test script - Add test_sdar_dvllm.py in examples/ - Support automatic weight key conversion from HF format to Diffulex-native - Test SDAR model forward pass with minimal attention metadata setup - Use .venv Python environment for execution --- examples/test_sdar_dvllm.py | 210 ++++++++++++++++++++++++++++++++++++ 1 file changed, 210 insertions(+) create mode 100644 examples/test_sdar_dvllm.py 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() + + From eb2a23383f5ca690d4f9cb96c36ef80ef5ae266a Mon Sep 17 00:00:00 2001 From: drewjin Date: Wed, 17 Dec 2025 15:15:49 +0000 Subject: [PATCH 17/23] refactor: update attention metadata handling and add warming up functionality; fix launch configurations and improve kernel integration; test: add dllm_flash_attn_decode_kernel test file --- .vscode/launch.json | 2 + diffulex/attention/__init__.py | 13 +- diffulex/attention/metadata.py | 16 +- diffulex/engine/sequence.py | 1 + diffulex/engine/tp_worker.py | 2 +- diffulex/sampler/fast_dllm_v2.py | 65 +-- .../block_diffusion/attention/metadata.py | 2 +- .../block_diffusion/engine/model_runner.py | 9 +- .../block_diffusion/engine/sequence.py | 23 +- diffulex/strategy/d2f/engine/model_runner.py | 4 +- diffulex/strategy/d2f/engine/sequence.py | 1 - diffulex_kernel/python/auto_tuner.py | 32 +- diffulex_kernel/python/dllm_flash_attn.py | 215 +++++----- examples/test_dream_dvllm_human_eval copy.py | 87 ---- examples/test_fastdllmv2_diffulex_gsm8k.py | 5 +- tests/.gitkeep | 0 .../test_dllm_flash_attn_decode_kernel.py | 399 ++++++++++++++++++ 17 files changed, 596 insertions(+), 280 deletions(-) delete mode 100755 examples/test_dream_dvllm_human_eval copy.py create mode 100644 tests/.gitkeep create mode 100644 tests/python/kernel/test_dllm_flash_attn_decode_kernel.py diff --git a/.vscode/launch.json b/.vscode/launch.json index 1ea0c9bd..783cad42 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -4,6 +4,8 @@ // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 "version": "0.2.0", "configurations": [ + + { "name": "Python Debugger: Current File", "type": "debugpy", diff --git a/diffulex/attention/__init__.py b/diffulex/attention/__init__.py index dfe9af02..a390a61d 100644 --- a/diffulex/attention/__init__.py +++ b/diffulex/attention/__init__.py @@ -1,4 +1,3 @@ -from .attn_impl import Attention from . import metadata from .metadata import set_fetch_fn_for_attn_metadata, AttnMetaDataBase @@ -12,4 +11,14 @@ def __call__(self, *args, **kwargs): def __repr__(self): return repr(metadata.fetch_attn_metadata) -fetch_attn_metadata = _FetchAttnMetadataProxy() \ No newline at end of file +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/metadata.py b/diffulex/attention/metadata.py index b71cb0c3..75c290ef 100644 --- a/diffulex/attention/metadata.py +++ b/diffulex/attention/metadata.py @@ -23,11 +23,23 @@ class AttnMetaDataBase: def num_seqs(self) -> int: return len(self.cu_seqlens_q) - 1 - FN_TYPE_AttnMetaDataFetch = Callable[[], AttnMetaDataBase] fetch_attn_metadata: FN_TYPE_AttnMetaDataFetch = ... 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/engine/sequence.py b/diffulex/engine/sequence.py index 00c34ab6..b467a906 100755 --- a/diffulex/engine/sequence.py +++ b/diffulex/engine/sequence.py @@ -37,6 +37,7 @@ 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 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/sampler/fast_dllm_v2.py b/diffulex/sampler/fast_dllm_v2.py index 38b5b9ad..ada75559 100644 --- a/diffulex/sampler/fast_dllm_v2.py +++ b/diffulex/sampler/fast_dllm_v2.py @@ -15,10 +15,15 @@ class FastdLLMV2SampleOutputForDiffusionLM(SampleOutputBase): @AutoSampler.register("fast_dllm_v2") class FastdLLMV2SamplerForDiffusionLM(SamplerBase): def _shift_logits(self, logits, last_logit=None): + """ + Shift logits to align with Fast-dLLM's prediction pattern. + 参考 generation_functions.py 中的 logits shift 逻辑（105, 112行） + """ if logits.shape[1] == 0: print("Warning: logits sequence length is 0, returning empty logits") raise Exception("logits sequence length is 0") + # 对应 generation_functions.py: logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1) shifted_logits = torch.zeros_like(logits) shifted_logits[1:, ...] = logits[:-1, ...] if last_logit is not None: @@ -28,48 +33,57 @@ def _shift_logits(self, logits, last_logit=None): return shifted_logits def forward(self, seqs: list[SequenceBase], logits: torch.Tensor, temperatures: torch.Tensor, - top_p=None, top_k=None, margin_confidence=False, neg_entropy=False): + 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 = {} + shifted_logits = self._shift_logits(seq_logits, seq.cached_or_caching_num_tokens - 1) + 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])) + if len(block.global_mask_token_ids) == 0: + continue + + 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") + ) + + 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: - 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()] + 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 @@ -82,5 +96,4 @@ def forward(self, seqs: list[SequenceBase], logits: torch.Tensor, temperatures: 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/strategy/block_diffusion/attention/metadata.py b/diffulex/strategy/block_diffusion/attention/metadata.py index 1bda15d7..d9832b9d 100644 --- a/diffulex/strategy/block_diffusion/attention/metadata.py +++ b/diffulex/strategy/block_diffusion/attention/metadata.py @@ -34,7 +34,7 @@ def set_bd_attn_metadata( block_tables: torch.Tensor | None = None, page_block_size: int = 32, diffusion_block_size: int = 32, - decode_mode: str = "varlen", + decode_mode: str = "static", attn_type: str = "full_attention", kv_cache_layout: str = "unified", need_kv_cache_store: bool = True, diff --git a/diffulex/strategy/block_diffusion/engine/model_runner.py b/diffulex/strategy/block_diffusion/engine/model_runner.py index 67a95a91..03b037ae 100644 --- a/diffulex/strategy/block_diffusion/engine/model_runner.py +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -9,7 +9,7 @@ from diffulex.config import Config from diffulex.engine.sequence import SequenceBase from diffulex.strategy.block_diffusion.engine.sequence import BDSequence -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.block_diffusion.attention.metadata import fetch_bd_attn_metadata, set_bd_attn_metadata, reset_bd_attn_metadata @@ -26,6 +26,7 @@ def __init__(self, config: Config, rank: int, event: Event | list[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 = ( @@ -39,6 +40,7 @@ def warmup_model(self): for seq in seqs: seq.post_process() torch.cuda.empty_cache() + reset_warming_up() def prepare_prefill(self, seqs: list[BDSequence]): input_ids: list[int] = [] @@ -129,9 +131,8 @@ def prepare_decode(self, seqs: list[BDSequence]): positions.extend(cur_positions) context_lens.append(cur_context_len) - seqlen = len(seq) seqlen_q = self.diffusion_block_size - seqlen_k = seqlen + 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) @@ -161,7 +162,7 @@ def prepare_decode(self, seqs: list[BDSequence]): max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, block_tables=block_tables, - page_block_size=self.config.kvcache_page_size, + 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, diff --git a/diffulex/strategy/block_diffusion/engine/sequence.py b/diffulex/strategy/block_diffusion/engine/sequence.py index bd482909..a1ac5e5e 100644 --- a/diffulex/strategy/block_diffusion/engine/sequence.py +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -47,6 +47,14 @@ def to_cache(self) -> None: 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 token_ids(self) -> list[int]: @@ -92,9 +100,7 @@ def local_mask_token_ids(self) -> list[int]: 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] @@ -145,10 +151,21 @@ def num_blocks_in_active_diffusion_block(self) -> int: 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_num_tokens(self) -> int: return sum(block.size for block in self.diffusion_blocks if block.is_to_cache or block.is_in_cache) + @property + 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, diff --git a/diffulex/strategy/d2f/engine/model_runner.py b/diffulex/strategy/d2f/engine/model_runner.py index 81a2a84b..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 @@ -27,6 +27,7 @@ def __init__(self, config: Config, rank: int, event: Event | list[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,6 +41,7 @@ def warmup_model(self): for seq in seqs: seq.post_process() torch.cuda.empty_cache() + reset_warming_up() def prepare_prefill(self, seqs: list[D2FSequence]): input_ids: list[int] = [] diff --git a/diffulex/strategy/d2f/engine/sequence.py b/diffulex/strategy/d2f/engine/sequence.py index 78ca0f0c..312b7126 100644 --- a/diffulex/strategy/d2f/engine/sequence.py +++ b/diffulex/strategy/d2f/engine/sequence.py @@ -140,7 +140,6 @@ def __init__( 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.meet_eos = False self.diffusion_blocks: list[D2FDiffusionBlock] = [] self.n_steps = 0 self.input_token_ids: list[int] = [] diff --git a/diffulex_kernel/python/auto_tuner.py b/diffulex_kernel/python/auto_tuner.py index 23798dd9..f9b5ea0d 100644 --- a/diffulex_kernel/python/auto_tuner.py +++ b/diffulex_kernel/python/auto_tuner.py @@ -1,39 +1,9 @@ -import torch import itertools -def get_heuristic_config() -> dict: - # Get CUDA device properties - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - sm_major, sm_minor = torch.cuda.get_device_capability(device) - sm_version = sm_major * 10 + sm_minor - if sm_version >= 80 and sm_version < 90: - return { - "BLOCK_M": 128, - "BLOCK_N": 256, - "NUM_STAGES": 2, - "NUM_THREADS": 128, - } - elif sm_version >= 90 and sm_version < 100: - return { - "BLOCK_M": 128, - "BLOCK_N": 256, - "NUM_STAGES": 3, - "NUM_THREADS": 256, - } - else: - return { - "BLOCK_M": 128, - "BLOCK_N": 256, - "NUM_STAGES": 0, - "NUM_THREADS": 128, - } - def build_configs(): BLOCK_M_LIST = [64, 128, 256] BLOCK_N_LIST = [64, 128, 256] - NUM_STAGES_LIST = [0, 1, 2, 3] + NUM_STAGES_LIST = [0, 1, 2] NUM_THREADS_LIST = [128, 256] CONFIGS = list( itertools.product( diff --git a/diffulex_kernel/python/dllm_flash_attn.py b/diffulex_kernel/python/dllm_flash_attn.py index a149a74f..37ca01be 100644 --- a/diffulex_kernel/python/dllm_flash_attn.py +++ b/diffulex_kernel/python/dllm_flash_attn.py @@ -2,24 +2,27 @@ import tilelang import tilelang.language as T -from tilelang.autotuner import set_autotune_inputs from flash_attn import flash_attn_varlen_func +from tilelang.autotuner import set_autotune_inputs + +from tilelang.engine.callback import register_cuda_postproc_callback 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 +from diffulex.attention.metadata import AttnMetaDataBase, is_warming_up +@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缓存，避免重复autotune和编译 -_prefill_kernel_cache = {} -_decode_kernel_cache = {} +kernel_config = None -@tilelang.autotune( - configs=build_configs() -) +@tilelang.autotune(configs=build_configs()) @tilelang.jit( - out_idx=[6], + out_idx=[-1], pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}, ) def dllm_flash_attn_prefill_kernel( @@ -157,9 +160,9 @@ def kernel( return kernel -@tilelang.autotune(configs=build_configs()) + @tilelang.jit( - out_idx=[10], + out_idx=[-1], pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,}, ) def dllm_flash_attn_decode_kernel( @@ -241,8 +244,6 @@ def kernel( cur_q_seqlen = q_end_idx - q_start_idx cur_kv_seqlen = kv_end_idx - kv_start_idx - T.device_assert(cur_q_seqlen == DIFFUSION_BLOCK_SIZE, "cur_q_seqlen must be equal to DIFFUSION_BLOCK_SIZE") - T.device_assert(cur_kv_seqlen == DIFFUSION_BLOCK_SIZE, "cur_kv_seqlen must be equal to DIFFUSION_BLOCK_SIZE") cur_context_len = context_lens[seq_idx] @@ -255,78 +256,79 @@ def kernel( T.fill(log_sum, 0) T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) - # Fusion of Q/KVCache Cross-Attention and QKV Self-Attention (Full-Attention) + # Q/KVCache Cross-Attention 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 == -1: - 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( + 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( (q_start_idx + i >= cur_q_seqlen or - kv_start_idx + j >= cur_kv_seqlen), -1e9, 0 + page_block_idx_local * PAGE_BLOCK_SIZE + j >= cur_context_len), -1e9, 0 ) - T.gemm(Q_shared, K_shared, acc_score_kv, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + # 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_kv, scores_max, dim=1, clear=False) + 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, 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, 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_kv, acc_score_kv_cast) + T.copy(acc_score_kvcache, acc_score_kvcache_cast) 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) - T.gemm(acc_score_kv_cast, V_shared, acc_output, policy=T.GemmWarpPolicy.FullRow) - - break + # Compute attention output + 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) - 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( - (q_start_idx + i >= cur_q_seqlen or - page_block_idx_local * PAGE_BLOCK_SIZE + j >= cur_context_len), -1e9, 0 + # QKV Self-Attention + loop_range = 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) + for i, j in T.Parallel(BLOCK_M, BLOCK_N): + acc_score_kv[i, j] = T.if_then_else( + (i >= cur_q_seqlen or + kv_block_idx * BLOCK_N + j >= cur_kv_seqlen), -1e9, 0 ) - # Compute attention scores - T.gemm(Q_shared, K_Cache_shared, acc_score_kvcache, transpose_b=True, policy=T.GemmWarpPolicy.FullRow) + T.gemm(Q_shared, K_shared, acc_score_kv, 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) + 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) - 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) + 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_kvcache, acc_score_kvcache_cast) + T.copy(acc_score_kv, acc_score_kv_cast) for i, j in T.Parallel(BLOCK_M, HEAD_DIM): acc_output[i, j] *= scores_scale[i] - # Compute attention output - 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) + T.copy(V[kv_start_idx : kv_start_idx + BLOCK_N, kv_head_idx, :], V_shared) + T.gemm(acc_score_kv_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] @@ -354,48 +356,50 @@ def dllm_flash_attn_prefill( softmax_scale=scale, block_table=None ) elif attn_metadata.attn_type == "block_attention": - # 创建缓存键，基于kernel的参数 - cache_key = ( - attn_metadata.num_seqs, - q.shape[1] // k.shape[1], # NUM_GROUPS - q.shape[0], # Q_LEN - k.shape[0], # KV_LEN - q.shape[1], # NUM_HEADS - q.shape[2], # HEAD_DIM - attn_metadata.diffusion_block_size, # DIFFUSION_BLOCK_SIZE - ) - - # 检查缓存 - if cache_key not in _prefill_kernel_cache: - # 使用set_autotune_inputs来触发autotune - # 这会在第一次调用时为所有配置测试性能并选择最佳配置 + 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, + q, k, v, + attn_metadata.cu_seqlens_q, + attn_metadata.cu_seqlens_k, + attn_metadata.max_seqlen_q, ]): - attn_kernel = dllm_flash_attn_prefill_kernel( + 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], - True, - attn_metadata.diffusion_block_size, + attn_metadata.attn_type == "block_attention", + attn_metadata.diffusion_block_size ) - _prefill_kernel_cache[cache_key] = attn_kernel + 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: - attn_kernel = _prefill_kernel_cache[cache_key] + 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, + ) - return attn_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, @@ -407,51 +411,22 @@ def dllm_flash_attn_decode( attn_metadata: AttnMetaDataBase ) -> torch.Tensor: if attn_metadata.decode_mode == "static": - # 创建缓存键，基于kernel的参数 - cache_key = ( + decode_kernel = dllm_flash_attn_decode_kernel( attn_metadata.num_seqs, - q.shape[1] // k.shape[1], # NUM_GROUPS - k_cache.shape[0], # NUM_PAGE_BLOCKS - q.shape[0], # Q_LEN - k.shape[0], # KV_LEN - q.shape[1], # NUM_HEADS - q.shape[2], # HEAD_DIM - attn_metadata.attn_type == "block_attention", # IS_BLOCK_ATTN - attn_metadata.diffusion_block_size, # DIFFUSION_BLOCK_SIZE - attn_metadata.block_tables.shape[1], # MAX_SEQ_NUM_BLOCKS - attn_metadata.page_block_size, # PAGE_BLOCK_SIZE + 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 ) - # 检查缓存 - if cache_key not in _decode_kernel_cache: - # 使用set_autotune_inputs来触发autotune - # 这会在第一次调用时为所有配置测试性能并选择最佳配置 - with set_autotune_inputs([ - 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, - ]): - attn_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, - ) - _decode_kernel_cache[cache_key] = attn_kernel - else: - attn_kernel = _decode_kernel_cache[cache_key] - - return attn_kernel( + return decode_kernel( q, k, v, k_cache, v_cache, attn_metadata.block_tables, attn_metadata.context_lens, diff --git a/examples/test_dream_dvllm_human_eval copy.py b/examples/test_dream_dvllm_human_eval copy.py deleted file mode 100755 index 2d95f003..00000000 --- a/examples/test_dream_dvllm_human_eval copy.py +++ /dev/null @@ -1,87 +0,0 @@ -import os -import csv -import time - -import pandas as pd - -from datasets import load_dataset -from viztracer import VizTracer -from transformers import AutoTokenizer - -from diffulex.legacy import LLM, 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) - - -if __name__ == "__main__": - WEIGHT_DIR = "/data1/ckpts" - DATA_DIR = "/data1/LargeData" - model = f"{WEIGHT_DIR}/Dream-org/Dream-v0-Base-7B" - LLM = LLM( - model, - lora_path=f"{WEIGHT_DIR}/SJTU-Deng-Lab/D2F_Dream_Base_7B_Lora", - use_lora=True, - model_name="dream", - model_type="diffusion_lm", - enforce_eager=True, - data_parallel_size=1, - tensor_parallel_size=1, - gpu_memory_utilization=0.30, - max_num_batched_tokens=1024, - max_num_seqs=1, - max_model_len=1024, - accept_threshold=0.95, - complete_threshold=0.9, - add_new_block_threshold=0.1, - kvcache_block_size=32, - kv_cache_layout="unified" - ) - tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) - sampling_params = SamplingParams(temperature=0.0, max_tokens=256) - - dataset = load_dataset(f"{DATA_DIR}/openai/openai_humaneval")["test"]['prompt'][:] - prompts = [tokenizer.bos_token + p for p in 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(20) - 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_fastdllmv2_diffulex_gsm8k.py b/examples/test_fastdllmv2_diffulex_gsm8k.py index 2dd4f3ff..efa4ff7d 100755 --- a/examples/test_fastdllmv2_diffulex_gsm8k.py +++ b/examples/test_fastdllmv2_diffulex_gsm8k.py @@ -58,7 +58,10 @@ def summarize_profiling(csv_path: str) -> dict: sampling_params = SamplingParams(temperature=0.0, max_tokens=256) dataset = load_dataset("gsm8k", "main", split="test")["question"][:10] - prompts = [tokenizer.apply_chat_template(p, tokenize=False) for p in tqdm(dataset)] + 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): 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..1be0d4a9 --- /dev/null +++ b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py @@ -0,0 +1,399 @@ +import torch +import tilelang +import tilelang.testing + +from diffulex_kernel.python.dllm_flash_attn import dllm_flash_attn_decode_kernel, dllm_flash_attn_decode + + +def naive_attention_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 + num_heads = q.shape[1] + num_kv_heads = k.shape[1] + head_dim = q.shape[2] + + 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 + + # Compute attention for each head + for head_idx in range(num_heads): + kv_head_idx = head_idx // num_groups + + q_head = q_seq[:, head_idx, :] # [seq_q_len, head_dim] + k_head = k_combined[:, kv_head_idx, :] # [total_kv_len, head_dim] + v_head = v_combined[:, kv_head_idx, :] # [total_kv_len, head_dim] + + # Compute attention scores + attn_scores = torch.matmul(q_head.float(), k_head.float().T) * scale # [seq_q_len, total_kv_len] + + # Apply softmax + attn_probs = torch.softmax(attn_scores, dim=-1) + + # Compute output + output_head = torch.matmul(attn_probs, v_head.float()) # [seq_q_len, head_dim] + output[q_start:q_end, head_idx, :] = output_head.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, + ) + + print(decode_kernel.get_kernel_source()) + + 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_attention_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_f16_single_seq(): + """Test with single sequence, float16.""" + 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="float16", + ) + + +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() From 9c87d28cccb92fbe0fa43c389e43a332a3565d9c Mon Sep 17 00:00:00 2001 From: drewjin Date: Thu, 18 Dec 2025 10:41:35 +0000 Subject: [PATCH 18/23] test: add test script of dllm_flash_attn_prefill_kernel; fix: full_attn prefilling kernel; refactor: renaming imgs to assets --- .gitignore | 1 + {imgs => assets}/logo_lr.png | Bin diffulex_kernel/python/dllm_flash_attn.py | 17 +- .../test_dllm_flash_attn_decode_kernel.py | 92 +++-- .../test_dllm_flash_attn_prefill_kernel.py | 314 ++++++++++++++++++ 5 files changed, 380 insertions(+), 44 deletions(-) rename {imgs => assets}/logo_lr.png (100%) create mode 100644 tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py diff --git a/.gitignore b/.gitignore index 2712c2f8..978c0fc1 100755 --- a/.gitignore +++ b/.gitignore @@ -24,6 +24,7 @@ TestResults.xml *.pyc *.egg-info lora_weight* +*.log log/ dist/ build/ 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_kernel/python/dllm_flash_attn.py b/diffulex_kernel/python/dllm_flash_attn.py index 37ca01be..473712e6 100644 --- a/diffulex_kernel/python/dllm_flash_attn.py +++ b/diffulex_kernel/python/dllm_flash_attn.py @@ -11,11 +11,11 @@ from diffulex_kernel.python.kv_cache_kernels import load_kvcache from diffulex.attention.metadata import AttnMetaDataBase, is_warming_up -@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 +# @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 @@ -57,8 +57,7 @@ def kernel( 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): + 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) @@ -156,7 +155,7 @@ def kernel( 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_block_idx * BLOCK_M, head_idx, d_idx] = O_shared[i, d_idx] + O[i + q_start_idx + q_block_idx * BLOCK_M, head_idx, d_idx] = O_shared[i, d_idx] return kernel @@ -244,6 +243,8 @@ def kernel( cur_q_seqlen = q_end_idx - q_start_idx cur_kv_seqlen = kv_end_idx - kv_start_idx + T.device_assert(cur_q_seqlen == DIFFUSION_BLOCK_SIZE, "cur_q_seqlen must be equal to DIFFUSION_BLOCK_SIZE") + T.device_assert(cur_kv_seqlen == DIFFUSION_BLOCK_SIZE, "cur_kv_seqlen must be equal to DIFFUSION_BLOCK_SIZE") cur_context_len = context_lens[seq_idx] diff --git a/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py index 1be0d4a9..b7d29863 100644 --- a/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py +++ b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py @@ -1,11 +1,15 @@ +import os +from pathlib import Path + import torch import tilelang import tilelang.testing +import torch.nn.functional as F -from diffulex_kernel.python.dllm_flash_attn import dllm_flash_attn_decode_kernel, dllm_flash_attn_decode +from diffulex_kernel.python.dllm_flash_attn import dllm_flash_attn_decode_kernel -def naive_attention_with_kvcache( +def naive_sdpa_with_kvcache( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -46,6 +50,8 @@ def naive_attention_with_kvcache( output = torch.zeros_like(q) + kv_indices = torch.arange(num_heads, device=q.device) // num_groups + for seq_idx in range(num_seqs): q_start = cu_seqlens_q[seq_idx].item() q_end = cu_seqlens_q[seq_idx + 1].item() @@ -84,23 +90,24 @@ def naive_attention_with_kvcache( k_combined = k_seq v_combined = v_seq - # Compute attention for each head - for head_idx in range(num_heads): - kv_head_idx = head_idx // num_groups - - q_head = q_seq[:, head_idx, :] # [seq_q_len, head_dim] - k_head = k_combined[:, kv_head_idx, :] # [total_kv_len, head_dim] - v_head = v_combined[:, kv_head_idx, :] # [total_kv_len, head_dim] - - # Compute attention scores - attn_scores = torch.matmul(q_head.float(), k_head.float().T) * scale # [seq_q_len, total_kv_len] - - # Apply softmax - attn_probs = torch.softmax(attn_scores, dim=-1) - - # Compute output - output_head = torch.matmul(attn_probs, v_head.float()) # [seq_q_len, head_dim] - output[q_start:q_end, head_idx, :] = output_head.to(output.dtype) + # Expand KV per head according to GQA groups and run SDPA once + k_per_head = k_combined[:, kv_indices, :] # [total_kv_len, num_heads, head_dim] + v_per_head = v_combined[:, kv_indices, :] + + q_sdpa = q_seq.transpose(0, 1).unsqueeze(0) # [1, num_heads, seq_q_len, head_dim] + k_sdpa = k_per_head.transpose(0, 1).unsqueeze(0) # [1, num_heads, total_kv_len, head_dim] + v_sdpa = v_per_head.transpose(0, 1).unsqueeze(0) + + attn_out = F.scaled_dot_product_attention( + q_sdpa, + k_sdpa, + v_sdpa, + dropout_p=0.0, + is_causal=False, + scale=scale, + ) # [1, num_heads, seq_q_len, head_dim] + + output[q_start:q_end] = attn_out.squeeze(0).transpose(0, 1).to(output.dtype) return output @@ -182,7 +189,20 @@ def run_dllm_flash_attn_decode( num_threads, ) - print(decode_kernel.get_kernel_source()) + 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, @@ -194,7 +214,7 @@ def run_dllm_flash_attn_decode( ) # Compute reference output - ref_output = naive_attention_with_kvcache( + ref_output = naive_sdpa_with_kvcache( q, k, v, k_cache, v_cache, block_tables, context_lens, cu_seqlens_q, cu_seqlens_k, @@ -360,21 +380,21 @@ def test_decode_bf16_diffusion_block_64(): ) -def test_decode_f16_single_seq(): - """Test with single sequence, float16.""" - 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="float16", - ) +# def test_decode_f16_single_seq(): +# """Test with single sequence, float16.""" +# 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="float16", +# ) def test_decode_bf16_varied_stages(): 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..453d8ec4 --- /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_start: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() From aa927ed188153b188e27298c42b02c573df9ccf4 Mon Sep 17 00:00:00 2001 From: drewjin Date: Thu, 18 Dec 2025 11:19:28 +0000 Subject: [PATCH 19/23] fix: dllm_flash_attn_prefill_kernel bug fixed, corresponding test script bug fixed, all test passed. --- diffulex_kernel/python/dllm_flash_attn.py | 6 +++--- tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/diffulex_kernel/python/dllm_flash_attn.py b/diffulex_kernel/python/dllm_flash_attn.py index 473712e6..f77cc14b 100644 --- a/diffulex_kernel/python/dllm_flash_attn.py +++ b/diffulex_kernel/python/dllm_flash_attn.py @@ -106,11 +106,11 @@ def kernel( 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 and kv_block_idx == loop_range - 1: + if IS_BLOCK_ATTN: for i, j in T.Parallel(BLOCK_M, BLOCK_N): - num_diffusion_blocks = T.min(i // DIFFUSION_BLOCK_SIZE + 1, BLOCK_M // DIFFUSION_BLOCK_SIZE) + num_diffusion_blocks = (q_block_idx * BLOCK_M + i) // DIFFUSION_BLOCK_SIZE + 1 acc_score[i, j] = T.if_then_else( - (kv_block_idx * BLOCK_N + j >= kv_block_idx * BLOCK_N + num_diffusion_blocks * DIFFUSION_BLOCK_SIZE) or + (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 ) diff --git a/tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py b/tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py index 453d8ec4..6bc9ba80 100644 --- a/tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py +++ b/tests/python/kernel/test_dllm_flash_attn_prefill_kernel.py @@ -59,7 +59,7 @@ def naive_sdpa_prefill( 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_start:block_end] = True + block_mask[..., block_start:block_end, :block_end] = True attn_out = F.scaled_dot_product_attention( q_sdpa, From bd2d26d9a11c96f0e4b9c2fec6b67a2801ec7369 Mon Sep 17 00:00:00 2001 From: drewjin Date: Thu, 18 Dec 2025 15:28:30 +0000 Subject: [PATCH 20/23] refactor: improve dllm_flash_attn_decode_kernel structure and logic; remove redundant assertions and enhance comments for clarity; update test to utilize einops for tensor rearrangement --- diffulex_kernel/python/dllm_flash_attn.py | 87 ++++++++++--------- .../test_dllm_flash_attn_decode_kernel.py | 37 ++------ 2 files changed, 53 insertions(+), 71 deletions(-) diff --git a/diffulex_kernel/python/dllm_flash_attn.py b/diffulex_kernel/python/dllm_flash_attn.py index f77cc14b..82ffed48 100644 --- a/diffulex_kernel/python/dllm_flash_attn.py +++ b/diffulex_kernel/python/dllm_flash_attn.py @@ -5,18 +5,18 @@ from flash_attn import flash_attn_varlen_func from tilelang.autotuner import set_autotune_inputs -from tilelang.engine.callback import register_cuda_postproc_callback - 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 @@ -243,8 +243,6 @@ def kernel( cur_q_seqlen = q_end_idx - q_start_idx cur_kv_seqlen = kv_end_idx - kv_start_idx - T.device_assert(cur_q_seqlen == DIFFUSION_BLOCK_SIZE, "cur_q_seqlen must be equal to DIFFUSION_BLOCK_SIZE") - T.device_assert(cur_kv_seqlen == DIFFUSION_BLOCK_SIZE, "cur_kv_seqlen must be equal to DIFFUSION_BLOCK_SIZE") cur_context_len = context_lens[seq_idx] @@ -257,14 +255,17 @@ def kernel( T.fill(log_sum, 0) T.fill(scores_max, -T.infinity(ACCUM_DTYPE)) - # Q/KVCache Cross-Attention + # ========================== + # 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( - (q_start_idx + i >= cur_q_seqlen or + (i >= cur_q_seqlen or page_block_idx_local * PAGE_BLOCK_SIZE + j >= cur_context_len), -1e9, 0 ) @@ -289,54 +290,60 @@ def kernel( 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] - # Compute attention output + # 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) - # QKV Self-Attention - loop_range = 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) - for i, j in T.Parallel(BLOCK_M, BLOCK_N): - acc_score_kv[i, j] = T.if_then_else( - (i >= cur_q_seqlen or - kv_block_idx * BLOCK_N + 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) + 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.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.gemm(Q_shared, K_shared, acc_score_kv, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - T.copy(acc_score_kv, acc_score_kv_cast) - 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) - T.gemm(acc_score_kv_cast, V_shared, acc_output, 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 + q_start_idx < cur_q_seqlen: + if i < cur_q_seqlen: O[i + q_start_idx, head_idx, d_idx] = O_shared[i, d_idx] return kernel diff --git a/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py index b7d29863..b08edf9f 100644 --- a/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py +++ b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py @@ -5,6 +5,7 @@ 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 @@ -44,14 +45,8 @@ def naive_sdpa_with_kvcache( output: [Q_LEN, NUM_HEADS, HEAD_DIM] """ num_seqs = len(cu_seqlens_q) - 1 - num_heads = q.shape[1] - num_kv_heads = k.shape[1] - head_dim = q.shape[2] output = torch.zeros_like(q) - - kv_indices = torch.arange(num_heads, device=q.device) // num_groups - for seq_idx in range(num_seqs): q_start = cu_seqlens_q[seq_idx].item() q_end = cu_seqlens_q[seq_idx + 1].item() @@ -89,14 +84,10 @@ def naive_sdpa_with_kvcache( else: k_combined = k_seq v_combined = v_seq - - # Expand KV per head according to GQA groups and run SDPA once - k_per_head = k_combined[:, kv_indices, :] # [total_kv_len, num_heads, head_dim] - v_per_head = v_combined[:, kv_indices, :] - q_sdpa = q_seq.transpose(0, 1).unsqueeze(0) # [1, num_heads, seq_q_len, head_dim] - k_sdpa = k_per_head.transpose(0, 1).unsqueeze(0) # [1, num_heads, total_kv_len, head_dim] - v_sdpa = v_per_head.transpose(0, 1).unsqueeze(0) + 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, @@ -105,9 +96,10 @@ def naive_sdpa_with_kvcache( 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] = attn_out.squeeze(0).transpose(0, 1).to(output.dtype) + output[q_start:q_end] = rearrange(attn_out, '1 h s d -> s h d').to(output.dtype) return output @@ -380,23 +372,6 @@ def test_decode_bf16_diffusion_block_64(): ) -# def test_decode_f16_single_seq(): -# """Test with single sequence, float16.""" -# 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="float16", -# ) - - def test_decode_bf16_varied_stages(): """Test with different pipeline stages.""" run_dllm_flash_attn_decode( From 35748d1e4823f091dd8e30eeddda86ceabfd2cb6 Mon Sep 17 00:00:00 2001 From: drewjin Date: Fri, 19 Dec 2025 08:27:40 +0000 Subject: [PATCH 21/23] fix: dllm_flash_attn_decode_kernel fixed --- diffulex_kernel/python/dllm_flash_attn.py | 187 ++++++++++++++++++ .../test_dllm_flash_attn_decode_kernel.py | 3 +- 2 files changed, 189 insertions(+), 1 deletion(-) diff --git a/diffulex_kernel/python/dllm_flash_attn.py b/diffulex_kernel/python/dllm_flash_attn.py index 82ffed48..099ed68b 100644 --- a/diffulex_kernel/python/dllm_flash_attn.py +++ b/diffulex_kernel/python/dllm_flash_attn.py @@ -192,6 +192,193 @@ def dllm_flash_attn_decode_kernel( 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), diff --git a/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py index b08edf9f..29200be6 100644 --- a/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py +++ b/tests/python/kernel/test_dllm_flash_attn_decode_kernel.py @@ -7,7 +7,8 @@ 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 +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( From e8547d031d43d1cec73bdde214f8676706e6821b Mon Sep 17 00:00:00 2001 From: drewjin Date: Sun, 21 Dec 2025 17:14:12 +0000 Subject: [PATCH 22/23] fix: block diffusion available to run, yet slow, and with buggy output; refactor: Move diffulex_legacy to a separate module to prevent import conflicts; --- .gitignore | 4 ++- .vscode/launch.json | 2 ++ diffulex/attention/attn_impl.py | 30 ++++++++--------- diffulex/legacy/__init__.py | 2 -- diffulex/sampler/base.py | 33 ++++++++++++++++++- diffulex/sampler/dream.py | 22 ++++--------- diffulex/sampler/fast_dllm_v2.py | 32 +++++------------- diffulex/sampler/llada.py | 4 +-- .../block_diffusion/engine/model_runner.py | 7 ++-- .../block_diffusion/engine/sequence.py | 31 ++++++++++++++--- diffulex/strategy/d2f/engine/sequence.py | 12 +++++++ diffulex/utils/loader.py | 2 +- diffulex_legacy/__init__.py | 2 ++ .../legacy => diffulex_legacy}/config.py | 0 .../engine/block_manager.py | 4 +-- .../engine/dp_engine.py | 6 ++-- .../engine/llm_engine.py | 10 +++--- .../engine/model_runner.py | 12 +++---- .../engine/scheduler.py | 8 ++--- .../engine/sequence.py | 4 +-- .../layers/activation.py | 0 .../layers/attention/attention_v1.py | 2 +- .../layers/attention/attention_v1_profile.py | 2 +- .../layers/attention/attention_v2.py | 2 +- .../layers/attention/attention_v2_dup.py | 4 +-- .../layers/attention/attention_v2_profile.py | 2 +- .../layers/attention/attention_v3.py | 4 +-- .../layers/attention/attention_v4.py | 4 +-- .../layers/attention/attention_v5.py | 4 +-- .../layers/attention/ops/__init__.py | 8 ++--- ...chunked_prefill_decoding_unified_kernel.py | 2 +- .../layers/attention/ops/kv_cache_kernels.py | 4 +-- .../layers/attention/ops/prefix_prefill.py | 0 .../attention/ops/tilus_decode_attn_dlm.py | 0 .../attention/ops/triton_decode_attn_clm.py | 0 .../attention/ops/triton_decode_attn_dlm.py | 4 +-- .../attention/ops/triton_flash_attention.py | 0 .../layers/embed_head.py | 2 +- .../layers/layernorm.py | 0 .../layers/linear.py | 0 .../layers/rotary_embedding.py | 0 .../layers/sampler.py | 5 +-- {diffulex/legacy => diffulex_legacy}/llm.py | 6 ++-- .../models/auto_model.py | 10 +++--- .../config/dream/configuration_dream.py | 0 .../configuration_fast_dllm_v2.py | 0 .../config/llada/configuration_llada.py | 0 .../models/dream.py | 14 ++++---- .../models/fast_dllm_v2.py | 14 ++++---- .../models/llada.py | 14 ++++---- .../models/qwen3.py | 12 +++---- .../models/utils/check_config.py | 0 .../sampling_params.py | 0 .../utils/checker.py | 0 .../utils/context.py | 2 +- .../utils/loader.py | 2 +- examples/test_causal_lm_decoding_kernel.py | 2 +- examples/test_dllm_decoding_kernel.py | 2 +- examples/test_dllm_kv_cache_load.py | 2 +- examples/test_dllm_kv_cache_store.py | 2 +- examples/test_dream_model_weight.py | 4 +-- examples/test_dream_model_weight_fixed.py | 4 +-- examples/test_fastdllmv2_diffulex_gsm8k.py | 9 ++--- examples/test_llada_dvllm_human_eval.py | 2 +- examples/test_qwen_dvllm.py | 2 +- pyproject.toml | 1 + 66 files changed, 214 insertions(+), 166 deletions(-) delete mode 100755 diffulex/legacy/__init__.py create mode 100755 diffulex_legacy/__init__.py rename {diffulex/legacy => diffulex_legacy}/config.py (100%) rename {diffulex/legacy => diffulex_legacy}/engine/block_manager.py (98%) rename {diffulex/legacy => diffulex_legacy}/engine/dp_engine.py (98%) rename {diffulex/legacy => diffulex_legacy}/engine/llm_engine.py (95%) rename {diffulex/legacy => diffulex_legacy}/engine/model_runner.py (99%) rename {diffulex/legacy => diffulex_legacy}/engine/scheduler.py (97%) rename {diffulex/legacy => diffulex_legacy}/engine/sequence.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/activation.py (100%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/attention_v1.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/attention_v1_profile.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/attention_v2.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/attention_v2_dup.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/attention_v2_profile.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/attention_v3.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/attention_v4.py (98%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/attention_v5.py (98%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/ops/__init__.py (56%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/ops/chunked_prefill_decoding_unified_kernel.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/ops/kv_cache_kernels.py (99%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/ops/prefix_prefill.py (100%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/ops/tilus_decode_attn_dlm.py (100%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/ops/triton_decode_attn_clm.py (100%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/ops/triton_decode_attn_dlm.py (97%) rename {diffulex/legacy => diffulex_legacy}/layers/attention/ops/triton_flash_attention.py (100%) rename {diffulex/legacy => diffulex_legacy}/layers/embed_head.py (97%) rename {diffulex/legacy => diffulex_legacy}/layers/layernorm.py (100%) rename {diffulex/legacy => diffulex_legacy}/layers/linear.py (100%) rename {diffulex/legacy => diffulex_legacy}/layers/rotary_embedding.py (100%) rename {diffulex/legacy => diffulex_legacy}/layers/sampler.py (98%) rename {diffulex/legacy => diffulex_legacy}/llm.py (65%) rename {diffulex/legacy => diffulex_legacy}/models/auto_model.py (59%) rename {diffulex/legacy => diffulex_legacy}/models/config/dream/configuration_dream.py (100%) rename {diffulex/legacy => diffulex_legacy}/models/config/fast_dllm_v2/configuration_fast_dllm_v2.py (100%) rename {diffulex/legacy => diffulex_legacy}/models/config/llada/configuration_llada.py (100%) rename {diffulex/legacy => diffulex_legacy}/models/dream.py (94%) rename {diffulex/legacy => diffulex_legacy}/models/fast_dllm_v2.py (94%) rename {diffulex/legacy => diffulex_legacy}/models/llada.py (95%) rename {diffulex/legacy => diffulex_legacy}/models/qwen3.py (94%) rename {diffulex/legacy => diffulex_legacy}/models/utils/check_config.py (100%) rename {diffulex/legacy => diffulex_legacy}/sampling_params.py (100%) rename {diffulex/legacy => diffulex_legacy}/utils/checker.py (100%) rename {diffulex/legacy => diffulex_legacy}/utils/context.py (98%) rename {diffulex/legacy => diffulex_legacy}/utils/loader.py (99%) diff --git a/.gitignore b/.gitignore index 978c0fc1..cc0144f9 100755 --- a/.gitignore +++ b/.gitignore @@ -34,4 +34,6 @@ ckpt/ data/ tilelang autotuner.log -Fast-dLLM \ No newline at end of file +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 783cad42..72ef3267 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -6,6 +6,8 @@ "configurations": [ + + { "name": "Python Debugger: Current File", "type": "debugpy", diff --git a/diffulex/attention/attn_impl.py b/diffulex/attention/attn_impl.py index aaf03975..9ec5f7fc 100644 --- a/diffulex/attention/attn_impl.py +++ b/diffulex/attention/attn_impl.py @@ -1,7 +1,6 @@ -import os import torch - import torch.nn as nn +from einops import rearrange from diffulex_kernel import ( store_kvcache_distinct_layout, @@ -26,16 +25,15 @@ 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 - self.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.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 @@ -44,9 +42,9 @@ def __init__( 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 = self.fetch_attn_metadata() k_cache, v_cache = self.k_cache, self.v_cache @@ -68,7 +66,7 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, if is_unified_layout: o = dllm_flash_attn_decode(q, k, v, k_cache, v_cache, self.scale, attn_metadata) else: - raise NotImplementedError("Distinct layout is not supported for decode mode") + raise NotImplementedError("Distinct layout is not supported yet...") # Final reshape - return o.view(-1, self.num_heads * self.head_dim).contiguous() \ 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/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/sampler/base.py b/diffulex/sampler/base.py index d6f81ad7..34f394fe 100644 --- a/diffulex/sampler/base.py +++ b/diffulex/sampler/base.py @@ -6,6 +6,8 @@ from dataclasses import dataclass from easydict import EasyDict as edict +from diffulex.engine.sequence import SequenceBase + class SamplerBase(nn.Module): def __init__(self): @@ -75,4 +77,33 @@ class SampleOutputBase: 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) \ No newline at end of file + 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 index d8cd9517..9f063408 100644 --- a/diffulex/sampler/dream.py +++ b/diffulex/sampler/dream.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from diffulex.sampler.auto_sampler import AutoSampler -from diffulex.sampler.base import SamplerBase, SampleOutputBase +from diffulex.sampler.base import SamplerShiftLogits, SampleOutputBase @dataclass @@ -12,20 +12,7 @@ class DreamSampleOutputForDiffusionLM(SampleOutputBase): @AutoSampler.register("dream") -class DreamSamplerForDiffusionLM(SamplerBase): - 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 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() @@ -38,7 +25,10 @@ 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) + + 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 diff --git a/diffulex/sampler/fast_dllm_v2.py b/diffulex/sampler/fast_dllm_v2.py index ada75559..ec323b50 100644 --- a/diffulex/sampler/fast_dllm_v2.py +++ b/diffulex/sampler/fast_dllm_v2.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from diffulex.sampler.auto_sampler import AutoSampler -from diffulex.sampler.base import SamplerBase, SampleOutputBase +from diffulex.sampler.base import SamplerShiftLogits, SampleOutputBase from diffulex.engine.sequence import SequenceBase @@ -13,25 +13,7 @@ class FastdLLMV2SampleOutputForDiffusionLM(SampleOutputBase): @AutoSampler.register("fast_dllm_v2") -class FastdLLMV2SamplerForDiffusionLM(SamplerBase): - def _shift_logits(self, logits, last_logit=None): - """ - Shift logits to align with Fast-dLLM's prediction pattern. - 参考 generation_functions.py 中的 logits shift 逻辑（105, 112行） - """ - if logits.shape[1] == 0: - print("Warning: logits sequence length is 0, returning empty logits") - raise Exception("logits sequence length is 0") - - # 对应 generation_functions.py: logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1) - 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 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() @@ -43,13 +25,14 @@ def forward(self, seqs: list[SequenceBase], logits: torch.Tensor, temperatures: 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_num_tokens - 1) + 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: @@ -58,7 +41,10 @@ def forward(self, seqs: list[SequenceBase], logits: torch.Tensor, temperatures: if len(block.global_mask_token_ids) == 0: continue - mask_token_logits = shifted_logits[block.global_mask_token_ids, ...] + 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, diff --git a/diffulex/sampler/llada.py b/diffulex/sampler/llada.py index 45befcd0..5202fa14 100644 --- a/diffulex/sampler/llada.py +++ b/diffulex/sampler/llada.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from diffulex.sampler.auto_sampler import AutoSampler -from diffulex.sampler.base import SamplerBase, SampleOutputBase +from diffulex.sampler.base import SamplerNoShiftLogits, SampleOutputBase @dataclass @@ -12,7 +12,7 @@ class LLaDASampleOutputForDiffusionLM(SampleOutputBase): @AutoSampler.register("llada") -class LLaDASamplerForDiffusionLM(SamplerBase): +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() diff --git a/diffulex/strategy/block_diffusion/engine/model_runner.py b/diffulex/strategy/block_diffusion/engine/model_runner.py index 03b037ae..d363ba4a 100644 --- a/diffulex/strategy/block_diffusion/engine/model_runner.py +++ b/diffulex/strategy/block_diffusion/engine/model_runner.py @@ -141,8 +141,11 @@ def prepare_decode(self, seqs: list[BDSequence]): if seq.diffusion_blocks[-1].is_active: slot_mapping.extend([-1] * self.diffusion_block_size) elif seq.diffusion_blocks[-1].is_to_cache: - for i in range(0, seq.num_blocks_in_active_diffusion_block): - start = seq.block_table[i] * self.block_size + 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)) diff --git a/diffulex/strategy/block_diffusion/engine/sequence.py b/diffulex/strategy/block_diffusion/engine/sequence.py index a1ac5e5e..936b2425 100644 --- a/diffulex/strategy/block_diffusion/engine/sequence.py +++ b/diffulex/strategy/block_diffusion/engine/sequence.py @@ -1,7 +1,5 @@ from __future__ import annotations -import torch - from enum import Enum, auto from dataclasses import dataclass @@ -122,6 +120,11 @@ def __init__( self.diffusion_blocks: list[BDDiffusionBlock] = [] self.diffusion_block_size = config.diffusion_block_size self.mask_token_id = config.mask_token_id + self.n_steps = 0 + + @property + def completion_token_ids(self) -> list[int]: + return self.token_ids[self.prefix_len : ] @property def prefix_len_with_padding(self) -> int: @@ -144,7 +147,7 @@ def active_block_token_ids(self) -> list[int]: return self.diffusion_blocks[-1].token_ids @property - def num_blocks_in_active_diffusion_block(self) -> int: + def num_page_blocks_in_active_diffusion_block(self) -> int: return self.diffusion_block_size // self.block_size @property @@ -155,9 +158,25 @@ def cached_num_tokens(self) -> int: 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: + 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 sum(block.size for block in self.diffusion_blocks if block.is_to_cache or block.is_in_cache) + return self.cached_or_caching_last_token_id + 1 + + @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_completion_tokens(self) -> int: @@ -217,6 +236,7 @@ def init_diffusion_blocks(self) -> None: ) 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.""" @@ -226,13 +246,14 @@ def next_diffusion_step(self) -> None: BDDiffusionBlock( block_id=len(self.diffusion_blocks), status=BDDiffusionBlockStatus.ACTIVE, - global_start_id=self.num_tokens, + 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: diff --git a/diffulex/strategy/d2f/engine/sequence.py b/diffulex/strategy/d2f/engine/sequence.py index 312b7126..db22bc89 100644 --- a/diffulex/strategy/d2f/engine/sequence.py +++ b/diffulex/strategy/d2f/engine/sequence.py @@ -325,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: 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_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/legacy/layers/attention/ops/__init__.py b/diffulex_legacy/layers/attention/ops/__init__.py similarity index 56% rename from diffulex/legacy/layers/attention/ops/__init__.py rename to diffulex_legacy/layers/attention/ops/__init__.py index 579ccbfe..d9a1b2ef 100755 --- a/diffulex/legacy/layers/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/legacy/layers/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/legacy/layers/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/legacy/layers/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/legacy/layers/attention/ops/prefix_prefill.py b/diffulex_legacy/layers/attention/ops/prefix_prefill.py similarity index 100% rename from diffulex/legacy/layers/attention/ops/prefix_prefill.py rename to diffulex_legacy/layers/attention/ops/prefix_prefill.py diff --git a/diffulex/legacy/layers/attention/ops/tilus_decode_attn_dlm.py b/diffulex_legacy/layers/attention/ops/tilus_decode_attn_dlm.py similarity index 100% rename from diffulex/legacy/layers/attention/ops/tilus_decode_attn_dlm.py rename to diffulex_legacy/layers/attention/ops/tilus_decode_attn_dlm.py diff --git a/diffulex/legacy/layers/attention/ops/triton_decode_attn_clm.py b/diffulex_legacy/layers/attention/ops/triton_decode_attn_clm.py similarity index 100% rename from diffulex/legacy/layers/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/legacy/layers/attention/ops/triton_flash_attention.py b/diffulex_legacy/layers/attention/ops/triton_flash_attention.py similarity index 100% rename from diffulex/legacy/layers/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/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_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 index efa4ff7d..1ca3152c 100755 --- a/examples/test_fastdllmv2_diffulex_gsm8k.py +++ b/examples/test_fastdllmv2_diffulex_gsm8k.py @@ -35,7 +35,8 @@ def summarize_profiling(csv_path: str) -> dict: 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<|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" @@ -81,6 +82,6 @@ def summarize_profiling(csv_path: str) -> dict: 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 + 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/pyproject.toml b/pyproject.toml index 826a717f..f2e26077 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ Repository = "https://zhijie-group.github.io/D2fEngine" include = [ "diffulex", "diffulex_kernel", + "diffulex_legacy", ] [[tool.uv.index]] From 19a61f7501b0d27033a842d104c0775858f11be3 Mon Sep 17 00:00:00 2001 From: drewjin Date: Mon, 22 Dec 2025 08:34:30 +0000 Subject: [PATCH 23/23] fix: correct usage of csv.DictReader in multiple example scripts and update model cache class definition --- examples/model_cache/llada/modeling_llada.py | 2 +- examples/summary.py | 2 +- examples/test_dream_dvllm_gsm8k.py | 2 +- examples/test_fastdllmv2_diffulex_gsm8k.py | 2 +- examples/test_llada_dvllm_human_eval.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/model_cache/llada/modeling_llada.py b/examples/model_cache/llada/modeling_llada.py index 74979814..87865e64 100755 --- a/examples/model_cache/llada/modeling_llada.py +++ b/examples/model_cache/llada/modeling_llada.py @@ -1059,7 +1059,7 @@ def __init__(self, config: ModelConfig, init_params: bool = True): torch.backends.cuda.enable_flash_sdp(True) torch.backends.cuda.enable_mem_efficient_sdp(False) # this is super slow so make sure torch won't use it - self.transformer = nn.Moduledict( + self.transformer = nn.ModuleDict( dict( wte=nn.Embedding( config.embedding_size or config.vocab_size, config.d_model, device=config.init_device diff --git a/examples/summary.py b/examples/summary.py index 871814a6..fc58c2bd 100755 --- a/examples/summary.py +++ b/examples/summary.py @@ -7,7 +7,7 @@ def summarize_profiling(csv_path: str) -> dict: total_nums = {} avgs = {} with open(csv_path, 'r', newline='') as f: - reader = csv.dictReader(f) + reader = csv.DictReader(f) for row in reader: for k, v in row.items(): try: diff --git a/examples/test_dream_dvllm_gsm8k.py b/examples/test_dream_dvllm_gsm8k.py index 66056272..4e25f8f7 100755 --- a/examples/test_dream_dvllm_gsm8k.py +++ b/examples/test_dream_dvllm_gsm8k.py @@ -17,7 +17,7 @@ def summarize_profiling(csv_path: str) -> dict: total_nums = {} avgs = {} with open(csv_path, 'r', newline='') as f: - reader = csv.dictReader(f) + reader = csv.DictReader(f) for row in reader: for k, v in row.items(): try: diff --git a/examples/test_fastdllmv2_diffulex_gsm8k.py b/examples/test_fastdllmv2_diffulex_gsm8k.py index 1ca3152c..3950537c 100755 --- a/examples/test_fastdllmv2_diffulex_gsm8k.py +++ b/examples/test_fastdllmv2_diffulex_gsm8k.py @@ -17,7 +17,7 @@ def summarize_profiling(csv_path: str) -> dict: total_nums = {} avgs = {} with open(csv_path, 'r', newline='') as f: - reader = csv.dictReader(f) + reader = csv.DictReader(f) for row in reader: for k, v in row.items(): try: diff --git a/examples/test_llada_dvllm_human_eval.py b/examples/test_llada_dvllm_human_eval.py index 5e3608f5..e52dc8ac 100755 --- a/examples/test_llada_dvllm_human_eval.py +++ b/examples/test_llada_dvllm_human_eval.py @@ -16,7 +16,7 @@ def summarize_profiling(csv_path: str) -> dict: total_nums = {} avgs = {} with open(csv_path, 'r', newline='') as f: - reader = csv.dictReader(f) + reader = csv.DictReader(f) for row in reader: for k, v in row.items(): try: