524 lines
19 KiB
Python
524 lines
19 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,
|
|
)
|
|
# Ensure AiterFlashAttentionMetadata dataclass has causal: bool field
|
|
if " causal: bool" not in text.split("class AiterFlashAttentionMetadata")[1].split("def ")[0]:
|
|
text = replace_once(
|
|
text,
|
|
"class AiterFlashAttentionMetadata:\n",
|
|
"class AiterFlashAttentionMetadata:\n causal: bool\n",
|
|
path,
|
|
)
|
|
|
|
# Ensure AiterFlashAttentionMetadata() constructors have causal= kwarg.
|
|
# Use a robust approach: find each constructor call, ensure causal= appears
|
|
# exactly once (first arg). This handles all upstream variations.
|
|
def _ensure_causal_in_constructor(m: re.Match) -> str:
|
|
call_text = m.group(0)
|
|
if re.search(r'\bcausal=', call_text):
|
|
return call_text # already has causal=, leave it
|
|
# Insert causal= as first kwarg after the opening paren
|
|
return call_text.replace(
|
|
"AiterFlashAttentionMetadata(\n",
|
|
"AiterFlashAttentionMetadata(\n causal=common_attn_metadata.causal,\n",
|
|
)
|
|
|
|
text = re.sub(
|
|
r'(?:attn_metadata = |return )AiterFlashAttentionMetadata\([^)]*?\)',
|
|
_ensure_causal_in_constructor,
|
|
text,
|
|
flags=re.DOTALL,
|
|
)
|
|
|
|
# Replace hardcoded causal=True with dynamic causal=attn_metadata.causal
|
|
# in flash attention calls. Skip if upstream already uses a dynamic causal.
|
|
text = re.sub(
|
|
r"(softmax_scale=self\.scale,\n)(\s*)causal=True,",
|
|
r"\1\2causal=attn_metadata.causal,",
|
|
text,
|
|
flags=re.MULTILINE,
|
|
)
|
|
|
|
# Safety: remove any duplicate causal= kwargs that may have been created
|
|
# by patching an upstream that already had causal= in constructors.
|
|
# This finds lines like `causal=...,` that appear more than once in the
|
|
# same parenthesized call and removes the extras (keeps first occurrence).
|
|
def _dedup_causal(m: re.Match) -> str:
|
|
block = m.group(0)
|
|
causal_lines = [i for i, line in enumerate(block.split('\n')) if re.match(r'\s*causal=', line)]
|
|
if len(causal_lines) <= 1:
|
|
return block
|
|
lines = block.split('\n')
|
|
# Keep only the first causal= line
|
|
seen = False
|
|
result = []
|
|
for line in lines:
|
|
if re.match(r'\s*causal=', line):
|
|
if seen:
|
|
continue # drop duplicate
|
|
seen = True
|
|
result.append(line)
|
|
return '\n'.join(result)
|
|
|
|
text = re.sub(
|
|
r'AiterFlashAttentionMetadata\([^)]*\)',
|
|
_dedup_causal,
|
|
text,
|
|
flags=re.DOTALL,
|
|
)
|
|
|
|
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 patch_mla_dynamic_fp8_scales(text: str, path: Path) -> str:
|
|
"""Patch MLA do_kv_cache_update to compute FP8 scales from actual data.
|
|
|
|
Without calibrated scales, FP8 KV cache clips values and produces garbage.
|
|
This patch computes the scale dynamically from each batch of KV data
|
|
before writing to cache, ensuring the full FP8 E4M3 range is used.
|
|
"""
|
|
marker = "# [DFlash patch] dynamic FP8 scale"
|
|
if marker in text:
|
|
return text
|
|
|
|
old = (
|
|
" def do_kv_cache_update(\n"
|
|
" self,\n"
|
|
" kv_c_normed: torch.Tensor,\n"
|
|
" k_pe: torch.Tensor,\n"
|
|
" kv_cache: torch.Tensor,\n"
|
|
" slot_mapping: torch.Tensor,\n"
|
|
" kv_cache_dtype: str,\n"
|
|
" k_scale: torch.Tensor,\n"
|
|
" ) -> None:\n"
|
|
" if kv_cache.numel() == 0:\n"
|
|
" return\n"
|
|
" from vllm import _custom_ops as ops\n"
|
|
"\n"
|
|
" ops.concat_and_cache_mla(\n"
|
|
" kv_c_normed,\n"
|
|
" k_pe.squeeze(1),\n"
|
|
" kv_cache,\n"
|
|
" slot_mapping.flatten(),\n"
|
|
" kv_cache_dtype=kv_cache_dtype,\n"
|
|
" scale=k_scale,\n"
|
|
" )"
|
|
)
|
|
|
|
new = (
|
|
" _fp8_call_count = 0\n"
|
|
" _fp8_scale_log = {}\n"
|
|
"\n"
|
|
" def do_kv_cache_update(\n"
|
|
" self,\n"
|
|
" kv_c_normed: torch.Tensor,\n"
|
|
" k_pe: torch.Tensor,\n"
|
|
" kv_cache: torch.Tensor,\n"
|
|
" slot_mapping: torch.Tensor,\n"
|
|
" kv_cache_dtype: str,\n"
|
|
" k_scale: torch.Tensor,\n"
|
|
" ) -> None:\n"
|
|
" if kv_cache.numel() == 0:\n"
|
|
" return\n"
|
|
" from vllm import _custom_ops as ops\n"
|
|
"\n"
|
|
" " + marker + "\n"
|
|
" if kv_cache_dtype in ('fp8', 'fp8_e4m3', 'fp8_e5m2'):\n"
|
|
" import os as _os, json as _json\n"
|
|
" import torch as _torch\n"
|
|
" _fp8_max = _torch.finfo(_torch.float8_e4m3fnuz).max\n"
|
|
" _abs_max = max(\n"
|
|
" kv_c_normed.abs().max().item(),\n"
|
|
" k_pe.abs().max().item(),\n"
|
|
" 1e-12,\n"
|
|
" )\n"
|
|
" _scale = _abs_max / _fp8_max\n"
|
|
" _static_file = _os.environ.get('FP8_KV_SCALES_FILE')\n"
|
|
" if _static_file and _os.path.exists(_static_file):\n"
|
|
" if not hasattr(self, '_fp8_static_scales'):\n"
|
|
" self._fp8_static_scales = _json.load(open(_static_file))\n"
|
|
" _s = self._fp8_static_scales.get('global_scale', _scale)\n"
|
|
" k_scale.fill_(_s)\n"
|
|
" else:\n"
|
|
" k_scale.copy_(_torch.max(k_scale, _torch.tensor(_scale, device=k_scale.device)))\n"
|
|
" type(self)._fp8_call_count += 1\n"
|
|
" _cur = k_scale.item()\n"
|
|
" type(self)._fp8_scale_log[id(k_scale)] = _cur\n"
|
|
" if type(self)._fp8_call_count % 5000 == 0:\n"
|
|
" _vals = list(type(self)._fp8_scale_log.values())\n"
|
|
" _dump = {'global_scale': max(_vals), 'num_layers': len(_vals),\n"
|
|
" 'min_scale': min(_vals), 'max_scale': max(_vals),\n"
|
|
" 'call_count': type(self)._fp8_call_count}\n"
|
|
" _out = '/tmp/fp8_kv_scales.json'\n"
|
|
" with open(_out, 'w') as _f:\n"
|
|
" _json.dump(_dump, _f, indent=2)\n"
|
|
"\n"
|
|
" ops.concat_and_cache_mla(\n"
|
|
" kv_c_normed,\n"
|
|
" k_pe.squeeze(1),\n"
|
|
" kv_cache,\n"
|
|
" slot_mapping.flatten(),\n"
|
|
" kv_cache_dtype=kv_cache_dtype,\n"
|
|
" scale=k_scale,\n"
|
|
" )"
|
|
)
|
|
|
|
if old not in text:
|
|
if marker in text:
|
|
return text
|
|
raise RuntimeError(f"Could not find do_kv_cache_update in {path}")
|
|
return text.replace(old, new, 1)
|
|
|
|
|
|
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",
|
|
]
|
|
)
|
|
|
|
mla_backend = vllm_root / "v1" / "attention" / "backend.py"
|
|
patch_file(mla_backend, patch_mla_dynamic_fp8_scales)
|
|
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())
|