feat: Initial block diffusion strategy support

.gitignore

@@ -24,10 +24,16 @@ TestResults.xml
 *.pyc
 *.egg-info
 lora_weight*
+*.log
 log/
 dist/
 build/
 cache/
 uv.lock
 ckpt/
-data/
+data/
+tilelang
+autotuner.log
+Fast-dLLM
+Discrete-Diffusion-Forcing
+position_explanation.md

.vscode/launch.json

@@ -6,11 +6,11 @@
     "configurations": [
 
 
-        
-
+
+    
         {
             "name": "Python Debugger: Current File",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${file}",
             "console": "integratedTerminal",
@@ -53,7 +53,7 @@
         },
         {
             "name": "PyDbg: `diffulex` Qwen3",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${workspaceFolder}/examples/test_qwen_dvllm.py",
             "console": "integratedTerminal",
@@ -64,7 +64,7 @@
         },
         {
             "name": "PyDbg: `diffulex` Dream `HumanEval`",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${workspaceFolder}/examples/test_dream_dvllm_human_eval.py",
             "console": "integratedTerminal",
@@ -76,7 +76,7 @@
         },
         {
             "name": "PyDbg: `diffulex` Dream `GSM8K`",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${workspaceFolder}/examples/test_dream_dvllm_gsm8k.py",
             "console": "integratedTerminal",
@@ -86,9 +86,21 @@
                 // "CUDA_VISIBLE_DEVICES": "0,1"
             }
         },
+        {
+            "name": "PyDbg: `diffulex` Fast-DLLM-V2 `GSM8K`",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/examples/test_fastdllmv2_diffulex_gsm8k.py",
+            "console": "integratedTerminal",
+            "env": {
+                // "TORCHINDUCTOR_DISABLE": "1",
+                // "TRITON_INTERPRET": "1",
+                // "CUDA_VISIBLE_DEVICES": "0,1"
+            }
+        },
         {
             "name": "PyDbg: `diffulex` LLaDA `HumanEval`",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${workspaceFolder}/examples/test_llada_dvllm_human_eval.py",
             "console": "integratedTerminal",
@@ -100,7 +112,7 @@
         },
         {
             "name": "PyDbg: `diffulex` kernel func `load_kvcache_kernel`",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${workspaceFolder}/examples/test_dllm_kv_cache_load.py",
             "console": "integratedTerminal",
@@ -111,7 +123,7 @@
         },
         {
             "name": "PyDbg: `diffulex` kernel func `chunked_prefill_paged_decode`",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${workspaceFolder}/examples/test_dllm_decoding_kernel.py",
             "console": "integratedTerminal",
@@ -122,7 +134,7 @@
         },
         {
             "name": "PyDbg: `diffulex` kernel func `causal_lm_decode_attention_fwd`",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${workspaceFolder}/examples/test_causal_lm_decoding_kernel.py",
             "console": "integratedTerminal",
@@ -133,7 +145,7 @@
         },
         {
             "name": "PyDbg: `diffulex` kernel func `store_kvcache_kernel_diffusion_lm`",
-            "type": "python",
+            "type": "debugpy",
             "request": "launch",
             "program": "${workspaceFolder}/examples/test_dllm_kv_cache_store.py",
             "console": "integratedTerminal",

diffulex/attention/__init__.py

@@ -1,2 +1,24 @@
-from .attn_impl import Attention
-from .metadata import fetch_attn_metadata, set_fetch_fn_for_attn_metadata, AttnMetaDataBase
+from . import metadata
+from .metadata import set_fetch_fn_for_attn_metadata, AttnMetaDataBase
+
+# Create a proxy that dynamically accesses fetch_attn_metadata from the metadata module
+# This ensures we always get the current value, not a stale copy from __init__.py
+class _FetchAttnMetadataProxy:
+    """Proxy object that dynamically accesses fetch_attn_metadata from metadata module."""
+    def __call__(self, *args, **kwargs):
+        return metadata.fetch_attn_metadata(*args, **kwargs)
+
+    def __repr__(self):
+        return repr(metadata.fetch_attn_metadata)
+
+fetch_attn_metadata = _FetchAttnMetadataProxy()
+
+
+def __getattr__(name):
+    """Lazy import to avoid circular deps during module init."""
+    if name == "Attention":
+        from .attn_impl import Attention
+        return Attention
+    if name == "fetch_attn_metadata":
+        return metadata.fetch_attn_metadata
+    raise AttributeError(f"module {__name__} has no attribute {name}")

diffulex/attention/attn_impl.py

@@ -1,20 +1,14 @@
-import os
 import torch
-
 import torch.nn as nn
-
-from functools import lru_cache, partial
 from einops import rearrange
-from torch.nn.attention.flex_attention import create_block_mask 
-from flash_attn import flash_attn_varlen_func
-from transformers.integrations.flex_attention import compile_friendly_flex_attention as flex_attention
 
-from diffulex.attention.ops import (
-    causal_lm_flash_decoding, diffusion_lm_flash_decoding, diffusion_lm_parallel_flash_decoding,
-    store_kvcache_unified_layout, store_kvcache_distinct_layout, load_kvcache,
-    CHECK_STORING, CHECK_LOADING, CHECK_ATTENTION
+from diffulex_kernel import (
+    store_kvcache_distinct_layout, 
+    store_kvcache_unified_layout, 
+    dllm_flash_attn_decode, 
+    dllm_flash_attn_prefill
 )
-from diffulex.attention.metadata import AttnMetaDataBase, fetch_attn_metadata
+from diffulex.attention.metadata import AttnMetaDataBase
 
 
 class Attention(nn.Module):
@@ -31,89 +25,48 @@ def __init__(
         self.scale = scale
         self.num_kv_heads = num_kv_heads
         self.k_cache = self.v_cache = torch.tensor([])
-        is_rtx_xx90 = lambda x: "4090" in x or "3090" in x
-        kernel_options = {
-            "BLOCK_M": 64,
-            "BLOCK_N": 64,
-            "BLOCK_M1": 32,
-            "BLOCK_N1": 64,
-            "BLOCK_M2": 64,
-            "BLOCK_N2": 32,
-        } if is_rtx_xx90(torch.cuda.get_device_name(0)) else None
-        self.attention = torch.compile(
-            partial(flex_attention, kernel_options=kernel_options, enable_gqa=True, 
-                    return_lse=False, training=False), dynamic=True)
-        self._block_mask_cache = {}
-
-    @lru_cache(maxsize=32)
-    def dllm_block_mask(self, block_mask: torch.Tensor, 
-                        B: int, H: int, Q_LEN: int, KV_LEN: int, device: str):
-        cache_key = (B, H, Q_LEN, KV_LEN, device)
-        def _mask_mod(batch, head, token_q, token_kv):
-            return block_mask[token_q, token_kv]
-        if cache_key not in self._block_mask_cache:
-            self._block_mask_cache[cache_key] = create_block_mask(
-                _mask_mod, B, H, Q_LEN, KV_LEN, device=device
-            )
-        return self._block_mask_cache[cache_key]
+
+        self.q_shape = {
+            'nh': self.num_heads,
+            'hd': self.head_dim,
+        }
+        self.kv_shape = {
+            'nkvh': self.num_kv_heads,
+            'hd': self.head_dim,
+        }
+        # Import the specified fetch function
+        from diffulex.attention import fetch_attn_metadata
+        self.fetch_attn_metadata = fetch_attn_metadata
+
 
     def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                 mask: list[torch.Tensor] | None = None) -> torch.Tensor:
         # Reshape
-        q = q.view(-1, self.num_heads, self.head_dim)
-        k = k.view(-1, self.num_kv_heads, self.head_dim)
-        v = v.view(-1, self.num_kv_heads, self.head_dim)
+        q = rearrange(q, 's (nh hd) -> s nh hd', **self.q_shape)
+        k = rearrange(k, 's (nkvh hd) -> s nkvh hd', **self.kv_shape)
+        v = rearrange(v, 's (nkvh hd) -> s nkvh hd', **self.kv_shape)
 
-        attn_metadata: AttnMetaDataBase = fetch_attn_metadata()
+        attn_metadata: AttnMetaDataBase = self.fetch_attn_metadata()
         k_cache, v_cache = self.k_cache, self.v_cache
         is_unified_layout = attn_metadata.kv_cache_layout == "unified"
 
         # Fast Store KV cache
         if k_cache.numel() and v_cache.numel():
-            if not (not attn_metadata.need_kv_cache_store):
+            if attn_metadata.need_kv_cache_store:
                 store_kvcache = store_kvcache_unified_layout if is_unified_layout else store_kvcache_distinct_layout
                 store_kvcache(k, v, k_cache, v_cache, attn_metadata.slot_mapping, attn_metadata)
-                # CHECK_STORING(k_cache, v_cache, k, v, context)
 
-        transpose_fn = lambda x: rearrange(x, 's h d -> 1 h s d').contiguous()
-        # Prefill / Decode logic TODO: Replace the Flex Attention Prefilling
+        # Prefill / Decode logic
         if attn_metadata.is_prefill:
-            # Block PK
             if attn_metadata.block_tables is not None:
                 # TODO: Implement Prefix Caching
                 pass
-
-            # Attention computation
-            q_t, k_t, v_t = [transpose_fn(t) for t in (q, k, v)]
-
-            B, H, S, _ = q_t.shape
-            block_mask = self.dllm_block_mask(attn_metadata.block_mask, B, H, S, S, str(q.device))
-            o = self.attention(q_t, k_t, v_t, block_mask=block_mask)
+            o = dllm_flash_attn_prefill(q, k, v, self.scale, attn_metadata)
         else:
-            config = attn_metadata.seqs[0].config
-            diffusion_block_size = config.diffusion_block_size
             if is_unified_layout:
-                k_comb, v_comb = load_kvcache(self.k_cache, self.v_cache, attn_metadata, k, v)
-                o = flash_attn_varlen_func(q, k_comb, v_comb, 
-                                            attn_metadata.cu_seqlens_q, attn_metadata.cu_seqlens_k,
-                                            attn_metadata.max_seqlen_q, attn_metadata.max_seqlen_k,
-                                            softmax_scale=self.scale, block_table=None)
+                o = dllm_flash_attn_decode(q, k, v, k_cache, v_cache, self.scale, attn_metadata)
             else:
-                # FIXME: Kernel not ok...
-                o = torch.empty_like(q).to(q.device).to(q.dtype)
-                q, k, o, k_cache, v_cache = map(lambda x: x.to(torch.float32), (q, k, o, k_cache, v_cache))
-                diffusion_lm_parallel_flash_decoding(
-                    q, k, v, o, str(k_cache.dtype), k_cache, v_cache, 
-                    attn_metadata.block_tables, attn_metadata.cu_seqlens_q, attn_metadata.total_lens,
-                    max(attn_metadata.total_lens), max(attn_metadata.seq_lens), 1.0, 1.0,
-                    diffusion_block_size, attn_metadata.block_mask
-                )
-                CHECK_ATTENTION(o, q, k, v, k_cache, v_cache, attn_metadata)
+                raise NotImplementedError("Distinct layout is not supported yet...")
 
         # Final reshape
-        if not attn_metadata.is_prefill:
-            o = o.view(-1, self.num_heads * self.head_dim).contiguous()
-        elif attn_metadata.is_prefill:
-            o = rearrange(o, '1 h s d -> s (h d)').contiguous()
-
-        return o
+        return rearrange(o, 's nh hd -> s (nh hd)').contiguous()

diffulex/attention/metadata.py

@@ -14,12 +14,32 @@ 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]
 
 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
+    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