381 lines
13 KiB
Python
381 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Patch ROCm DFlash support into installed vLLM and AITER packages.
|
|
|
|
The reachable host uses newer `vllm/v1` attention paths than the older
|
|
one-off patch script from the learnings. This script applies the same logical
|
|
fixes against the actual installed package files inside the runtime container.
|
|
It is intentionally idempotent.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def locate_module_file(module_name: str) -> Path:
|
|
spec = importlib.util.find_spec(module_name)
|
|
if spec is None or spec.origin is None:
|
|
raise RuntimeError(f"Could not locate module: {module_name}")
|
|
return Path(spec.origin).resolve()
|
|
|
|
|
|
def first_existing(paths: list[Path]) -> Path:
|
|
for path in paths:
|
|
if path.exists():
|
|
return path
|
|
raise RuntimeError("Could not locate any expected path:\n" + "\n".join(map(str, paths)))
|
|
|
|
|
|
def replace_once(text: str, old: str, new: str, path: Path) -> str:
|
|
if new in text:
|
|
return text
|
|
if old not in text:
|
|
raise RuntimeError(f"Pattern not found in {path}: {old[:120]!r}")
|
|
return text.replace(old, new, 1)
|
|
|
|
|
|
def replace_all_regex(
|
|
text: str,
|
|
pattern: str,
|
|
repl: str,
|
|
path: Path,
|
|
*,
|
|
min_count: int = 1,
|
|
) -> str:
|
|
compiled = re.compile(pattern, re.MULTILINE)
|
|
matches = list(compiled.finditer(text))
|
|
marker = re.sub(r"\\(?:g<[^>]+>|[1-9][0-9]*)", "", repl)
|
|
if not matches:
|
|
if repl in text or (marker and marker in text):
|
|
return text
|
|
raise RuntimeError(f"Regex pattern not found in {path}: {pattern}")
|
|
if len(matches) < min_count:
|
|
updated = compiled.sub(repl, text)
|
|
if updated != text and marker and marker in updated:
|
|
return updated
|
|
raise RuntimeError(
|
|
f"Expected at least {min_count} matches in {path}, found {len(matches)}"
|
|
)
|
|
return compiled.sub(repl, text)
|
|
|
|
|
|
def patch_file(path: Path, transform) -> None:
|
|
original = path.read_text()
|
|
updated = transform(original, path)
|
|
if updated == original:
|
|
print(f"[skip] {path}")
|
|
return
|
|
path.write_text(updated)
|
|
print(f"[patch] {path}")
|
|
|
|
|
|
def patch_rocm_aiter_fa(text: str, path: Path) -> str:
|
|
text = replace_once(
|
|
text,
|
|
' @staticmethod\n def get_name() -> str:\n return "FLASH_ATTN"\n\n @staticmethod\n def get_impl_cls() -> type["AiterFlashAttentionImpl"]:\n',
|
|
' @staticmethod\n def get_name() -> str:\n return "FLASH_ATTN"\n\n @classmethod\n def supports_non_causal(cls) -> bool:\n return True\n\n @staticmethod\n def get_impl_cls() -> type["AiterFlashAttentionImpl"]:\n',
|
|
path,
|
|
)
|
|
text = replace_once(
|
|
text,
|
|
"class AiterFlashAttentionMetadata:\n",
|
|
"class AiterFlashAttentionMetadata:\n causal: bool\n",
|
|
path,
|
|
)
|
|
text = replace_once(
|
|
text,
|
|
" attn_metadata = AiterFlashAttentionMetadata(\n num_actual_tokens=common_attn_metadata.num_actual_tokens,\n",
|
|
" attn_metadata = AiterFlashAttentionMetadata(\n causal=common_attn_metadata.causal,\n num_actual_tokens=common_attn_metadata.num_actual_tokens,\n",
|
|
path,
|
|
)
|
|
text = replace_once(
|
|
text,
|
|
" return AiterFlashAttentionMetadata(\n num_actual_tokens=num_tokens,\n",
|
|
" return AiterFlashAttentionMetadata(\n causal=common_attn_metadata.causal,\n num_actual_tokens=num_tokens,\n",
|
|
path,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"(softmax_scale=self\.scale,\n)(\s*)causal=True,",
|
|
r"\1\2causal=attn_metadata.causal,",
|
|
path,
|
|
min_count=5,
|
|
)
|
|
return text
|
|
|
|
|
|
def patch_supports_non_causal(text: str, path: Path, backend_name: str) -> str:
|
|
insertion = (
|
|
f' @staticmethod\n def get_name() -> str:\n return "{backend_name}"\n\n'
|
|
" @classmethod\n def supports_non_causal(cls) -> bool:\n return True\n\n"
|
|
)
|
|
current = f' @staticmethod\n def get_name() -> str:\n return "{backend_name}"\n\n'
|
|
if "def supports_non_causal" in text:
|
|
return text
|
|
if current not in text:
|
|
raise RuntimeError(f"Could not find get_name block in {path}")
|
|
return text.replace(current, insertion, 1)
|
|
|
|
|
|
def patch_aiter_wrapper(text: str, path: Path) -> str:
|
|
if "IS_CAUSAL=causal" in text:
|
|
return text.replace(
|
|
' assert causal, "Only causal attention is supported"\n',
|
|
"",
|
|
)
|
|
text = replace_once(
|
|
text,
|
|
' assert causal, "Only causal attention is supported"\n',
|
|
"",
|
|
path,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"(ALL_DECODE=ALL_DECODE,\n)(\s*)(\*\*attn_config,)",
|
|
r"\1\2IS_CAUSAL=causal,\n\2\3",
|
|
path,
|
|
min_count=1,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"(ALL_DECODE=ALL_DECODE,\n)(\s*)(\*\*config,)",
|
|
r"\1\2IS_CAUSAL=causal,\n\2\3",
|
|
path,
|
|
min_count=1,
|
|
)
|
|
return text
|
|
|
|
|
|
def patch_aiter_kernel(text: str, path: Path) -> str:
|
|
if "IS_CAUSAL: tl.constexpr = True" in text:
|
|
text = text.replace(
|
|
"num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)",
|
|
"num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) if IS_CAUSAL else cdiv_fn(seq_len, TILE_SIZE)",
|
|
)
|
|
text = text.replace(
|
|
"seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1",
|
|
"seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1 if IS_CAUSAL else seq_offset[None, :] < seq_len",
|
|
)
|
|
return text
|
|
text = replace_all_regex(
|
|
text,
|
|
r"(ALL_DECODE: tl\.constexpr = False, # bool\n)(\):)",
|
|
r"\1 IS_CAUSAL: tl.constexpr = True, # bool\n\2",
|
|
path,
|
|
min_count=2,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"num_tiles = cdiv_fn\(max_seq_prefix_len, TILE_SIZE\)",
|
|
"num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) if IS_CAUSAL else cdiv_fn(seq_len, TILE_SIZE)",
|
|
path,
|
|
min_count=2,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"seq_mask = seq_offset\[None, :\] < context_len \+ query_pos\[:, None\] \+ 1",
|
|
"seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1 if IS_CAUSAL else seq_offset[None, :] < seq_len",
|
|
path,
|
|
min_count=2,
|
|
)
|
|
return text
|
|
|
|
|
|
def patch_vllm_triton_unified_attention(text: str, path: Path) -> str:
|
|
if "IS_CAUSAL=causal" in text:
|
|
return text.replace(
|
|
' assert causal, "Only causal attention is supported"\n',
|
|
"",
|
|
)
|
|
text = replace_once(
|
|
text,
|
|
' assert causal, "Only causal attention is supported"\n',
|
|
"",
|
|
path,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"num_tiles = cdiv_fn\(max_seq_prefix_len, TILE_SIZE\)",
|
|
"num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) if IS_CAUSAL else cdiv_fn(seq_len, TILE_SIZE)",
|
|
path,
|
|
min_count=2,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"seq_mask = seq_offset\[None, :\] <= query_abs_pos",
|
|
"seq_mask = seq_offset[None, :] <= query_abs_pos if IS_CAUSAL else seq_offset[None, :] < seq_len",
|
|
path,
|
|
min_count=2,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"USE_FP8: tl\.constexpr, # bool",
|
|
"USE_FP8: tl.constexpr, # bool\n IS_CAUSAL: tl.constexpr = True, # bool",
|
|
path,
|
|
min_count=2,
|
|
)
|
|
text = replace_all_regex(
|
|
text,
|
|
r"(BLOCK_M=BLOCK_M,\n)(\s*)",
|
|
r"\1\2IS_CAUSAL=causal,\n\2",
|
|
path,
|
|
min_count=2,
|
|
)
|
|
return text
|
|
|
|
|
|
def patch_vllm_dflash(text: str, path: Path) -> str:
|
|
return replace_once(
|
|
text,
|
|
' assert getattr(attn_metadata, "causal", None) is False, (\n',
|
|
' assert getattr(attn_metadata, "causal", None) in (True, False, None), (\n',
|
|
path,
|
|
)
|
|
|
|
|
|
def patch_vllm_selector(text: str, path: Path) -> str:
|
|
if "AttentionBackendEnum.TRITON_MLA" in text:
|
|
return text
|
|
text = replace_once(
|
|
text,
|
|
"from vllm.v1.attention.backends.registry import (\n"
|
|
" MAMBA_TYPE_TO_BACKEND_MAP,\n"
|
|
" MambaAttentionBackendEnum,\n"
|
|
")\n",
|
|
"from vllm.v1.attention.backends.registry import (\n"
|
|
" AttentionBackendEnum,\n"
|
|
" MAMBA_TYPE_TO_BACKEND_MAP,\n"
|
|
" MambaAttentionBackendEnum,\n"
|
|
")\n",
|
|
path,
|
|
)
|
|
new = (
|
|
" speculative_config = vllm_config.speculative_config\n"
|
|
" hf_config = vllm_config.model_config.hf_config\n"
|
|
" architectures = list(getattr(hf_config, \"architectures\", []) or [])\n"
|
|
" is_dflash_draft = any(\n"
|
|
" str(arch).startswith(\"DFlash\") for arch in architectures\n"
|
|
" )\n"
|
|
" use_non_causal = (\n"
|
|
" speculative_config is not None\n"
|
|
" and speculative_config.method == \"dflash\"\n"
|
|
" and is_dflash_draft\n"
|
|
" )\n"
|
|
"\n"
|
|
" backend = vllm_config.attention_config.backend\n"
|
|
" if (\n"
|
|
" speculative_config is not None\n"
|
|
" and speculative_config.method == \"dflash\"\n"
|
|
" and use_mla\n"
|
|
" and not is_dflash_draft\n"
|
|
" and backend is None\n"
|
|
" ):\n"
|
|
" backend = AttentionBackendEnum.TRITON_MLA\n"
|
|
)
|
|
old_variants = [
|
|
(
|
|
" speculative_config = vllm_config.speculative_config\n"
|
|
" use_non_causal = (\n"
|
|
" speculative_config is not None and speculative_config.method == \"dflash\"\n"
|
|
" )\n"
|
|
),
|
|
(
|
|
" speculative_config = vllm_config.speculative_config\n"
|
|
" hf_config = vllm_config.model_config.hf_config\n"
|
|
" architectures = list(getattr(hf_config, \"architectures\", []) or [])\n"
|
|
" use_non_causal = (\n"
|
|
" speculative_config is not None\n"
|
|
" and speculative_config.method == \"dflash\"\n"
|
|
" and any(str(arch).startswith(\"DFlash\") for arch in architectures)\n"
|
|
" )\n"
|
|
),
|
|
]
|
|
for old in old_variants:
|
|
if old in text:
|
|
text = text.replace(old, new, 1)
|
|
break
|
|
else:
|
|
if new not in text:
|
|
raise RuntimeError(f"Could not find selector speculative block in {path}")
|
|
return replace_once(
|
|
text,
|
|
" backend=vllm_config.attention_config.backend,\n",
|
|
" backend=backend,\n",
|
|
path,
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
vllm_root = locate_module_file("vllm").parent
|
|
site_packages = vllm_root.parent
|
|
|
|
rocm_aiter_fa = vllm_root / "v1" / "attention" / "backends" / "rocm_aiter_fa.py"
|
|
rocm_attn = vllm_root / "v1" / "attention" / "backends" / "rocm_attn.py"
|
|
rocm_aiter_unified = (
|
|
vllm_root / "v1" / "attention" / "backends" / "rocm_aiter_unified_attn.py"
|
|
)
|
|
triton_attn = vllm_root / "v1" / "attention" / "backends" / "triton_attn.py"
|
|
selector_path = vllm_root / "v1" / "attention" / "selector.py"
|
|
vllm_triton_ops = (
|
|
vllm_root / "v1" / "attention" / "ops" / "triton_unified_attention.py"
|
|
)
|
|
vllm_dflash = vllm_root / "v1" / "spec_decode" / "dflash.py"
|
|
aiter_wrapper = first_existing(
|
|
[
|
|
site_packages / "aiter" / "ops" / "triton" / "unified_attention.py",
|
|
site_packages
|
|
/ "aiter"
|
|
/ "ops"
|
|
/ "triton"
|
|
/ "attention"
|
|
/ "unified_attention.py",
|
|
]
|
|
)
|
|
aiter_kernel = first_existing(
|
|
[
|
|
site_packages
|
|
/ "aiter"
|
|
/ "ops"
|
|
/ "triton"
|
|
/ "_triton_kernels"
|
|
/ "unified_attention.py",
|
|
site_packages
|
|
/ "aiter"
|
|
/ "ops"
|
|
/ "triton"
|
|
/ "_triton_kernels"
|
|
/ "attention"
|
|
/ "unified_attention.py",
|
|
]
|
|
)
|
|
|
|
patch_file(rocm_aiter_fa, patch_rocm_aiter_fa)
|
|
patch_file(
|
|
rocm_attn,
|
|
lambda text, path: patch_supports_non_causal(text, path, "ROCM_ATTN"),
|
|
)
|
|
patch_file(
|
|
rocm_aiter_unified,
|
|
lambda text, path: patch_supports_non_causal(
|
|
text, path, "ROCM_AITER_UNIFIED_ATTN"
|
|
),
|
|
)
|
|
patch_file(
|
|
triton_attn,
|
|
lambda text, path: patch_supports_non_causal(text, path, "TRITON_ATTN"),
|
|
)
|
|
patch_file(selector_path, patch_vllm_selector)
|
|
patch_file(vllm_triton_ops, patch_vllm_triton_unified_attention)
|
|
patch_file(vllm_dflash, patch_vllm_dflash)
|
|
patch_file(aiter_wrapper, patch_aiter_wrapper)
|
|
patch_file(aiter_kernel, patch_aiter_kernel)
|
|
print("[done] ROCm DFlash patch applied")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|