From 0ab1722ad3376aba499aea2511e2ec70874956ab Mon Sep 17 00:00:00 2001 From: fengzhengxiong Date: Tue, 20 Oct 2026 11:22:19 +0800 Subject: [PATCH] fix(musa): support Gemma4 heterogeneous attention MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 修复 MUSA FA4 能力误判，隔离 TD/非 TD Triton kernel，并将 B200 宽 tile 调优限制到 CUDA，补充 Gemma4 head_dim=512 回归覆盖。 Assisted-by: Codex Signed-off-by: fengzhengxiong --- .../test_triton_unified_attention.py | 113 +++ ...test_musa_fa4_capability_false_positive.py | 135 ++++ ...est_triton_unified_attention_dual_entry.py | 660 ++++++++++++++++++ ...ified_attention_td_dependency_isolation.py | 144 ++++ .../attention/ops/triton_unified_attention.py | 403 ++++++++++- vllm/vllm_flash_attn/flash_attn_interface.py | 7 + 6 files changed, 1457 insertions(+), 5 deletions(-) create mode 100644 tests/v1/attention/test_musa_fa4_capability_false_positive.py create mode 100644 tests/v1/attention/test_triton_unified_attention_dual_entry.py create mode 100644 tests/v1/attention/test_triton_unified_attention_td_dependency_isolation.py diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index d3435ea665..8df0ae6821 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -643,3 +643,116 @@ def test_triton_unified_attn_use_td_tile_clamp( soft_cap=None, seq_threshold_3D=0, ) + + +# --------------------------------------------------------------------------- +# Gemma4 全局注意力层使用 head_dim=512，与滑动窗口层的 head_dim=256 异构。 +# 这里不把 512 加入共享的 HEAD_SIZES / HEAD_SIZES_USE_TD，因为 CUDA 的 +# 512 路径仅支持 FA4，加入共享矩阵会扩大无关回归范围。离线环境只验证参数 +# 和选路逻辑；数值用例仅在真实 MUSA 设备上执行。 +# --------------------------------------------------------------------------- +GEMMA4_GLOBAL_HEAD_SIZE = 512 + + +def test_gemma4_global_head_dim_512_selection_logic() -> None: + """验证 head_size=512 的 tile 选择和 2 的幂补齐逻辑。""" + from vllm.v1.attention.ops.triton_unified_attention import _get_tile_size + + # head_size=512 不属于 Gemma3 的 128/256 特例，使用通用规则：2 字节元素 + # 下 prefill 为 32，decode 为 16。 + assert _get_tile_size(GEMMA4_GLOBAL_HEAD_SIZE, 1024, 2, True) == 32 + assert _get_tile_size(GEMMA4_GLOBAL_HEAD_SIZE, 1024, 2, False) == 16 + # 512 已是 2 的幂，因此 HEAD_SIZE == HEAD_SIZE_PADDED。 + assert next_power_of_2(GEMMA4_GLOBAL_HEAD_SIZE) == GEMMA4_GLOBAL_HEAD_SIZE + + +def _run_head_size_512_non_td_case( + seq_lens: list[tuple[int, int]], + num_heads: tuple[int, int], + block_size: int, + num_blocks: int = 2048, +) -> None: + """以 Gemma4 全局层的 head_size=512 运行非 TD 指针 kernel。""" + torch.set_default_device(DEVICE_TYPE) + set_random_seed(0) + + head_size = GEMMA4_GLOBAL_HEAD_SIZE + num_seqs = len(seq_lens) + query_lens = [x[0] for x in seq_lens] + kv_lens = [x[1] for x in seq_lens] + num_query_heads, num_kv_heads = num_heads + assert num_query_heads % num_kv_heads == 0 + max_query_len = max(query_lens) + max_kv_len = max(kv_lens) + sliding_window = None + scale = head_size**-0.5 + + query = torch.randn( + sum(query_lens), num_query_heads, head_size, dtype=torch.bfloat16 + ) + key_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, dtype=torch.bfloat16 + ) + value_cache = torch.randn_like(key_cache) + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( + dim=0, dtype=torch.int32 + ) + kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32) + + max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size + block_tables = torch.randint( + 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 + ) + + output = torch.empty_like(query) + + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=cu_query_lens, + seqused_k=kv_lens_tensor, + max_seqlen_q=max_query_len, + max_seqlen_k=max_kv_len, + softmax_scale=scale, + causal=True, + window_size=(-1, -1), + block_table=block_tables, + softcap=0, + q_descale=None, + k_descale=None, + v_descale=None, + use_td=False, + ) + + ref_output = ref_paged_attn( + query=query, + key_cache=key_cache, + value_cache=value_cache, + query_lens=query_lens, + kv_lens=kv_lens, + block_tables=block_tables, + scale=scale, + sliding_window=sliding_window, + soft_cap=None, + ) + torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not current_platform.is_musa(), + reason="head_dim=512 non-TD coverage targets the MUSA pointer kernel; " + "other platforms keep their existing head-size matrix.", +) +@pytest.mark.parametrize("num_heads", [(4, 4), (8, 2)]) +@torch.inference_mode() +def test_triton_unified_attn_head_size_512_musa_non_td( + num_heads: tuple[int, int], +) -> None: + """验证 Gemma4 head_dim=512 非 TD 指针 kernel 的数值正确性。""" + _run_head_size_512_non_td_case( + seq_lens=[(1, 1328), (5, 18), (129, 463)], + num_heads=num_heads, + block_size=16, + ) diff --git a/tests/v1/attention/test_musa_fa4_capability_false_positive.py b/tests/v1/attention/test_musa_fa4_capability_false_positive.py new file mode 100644 index 0000000000..dc7b76db82 --- /dev/null +++ b/tests/v1/attention/test_musa_fa4_capability_false_positive.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""回归测试：MUSA 的伪设备能力不能被当作 CUDA FA4 支持证据。 + +MUSA 的 ``DeviceCapability(10, 0)`` 会命中 SM100 能力族，但目标环境实际 +只提供 FA3。本文件使用平台桩验证真实 ``_is_fa4_supported`` 的分支逻辑； +不执行 MUSA kernel，也不代表 head_dim=512 数值测试。 +""" + +import importlib.util +import pathlib +import sys +import types + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +_MODULE_PATH = _REPO_ROOT / "vllm" / "vllm_flash_attn" / "flash_attn_interface.py" +_PROBE_NAME = "_flash_attn_interface_test" + + +class _FakeCapability: + """复现 ``DeviceCapability`` 的整数映射。""" + + def __init__(self, major: int, minor: int) -> None: + self.major = major + self.minor = minor + + def to_int(self) -> int: + assert 0 <= self.minor < 10 + return self.major * 10 + self.minor + + +class _FakePlatform: + """复现真实设备能力判断的最小平台桩。""" + + def __init__(self, name: str, capability: _FakeCapability) -> None: + self._name = name + self._capability = capability + + def is_musa(self) -> bool: + return self._name == "musa" + + def get_device_capability(self, device_id: int = 0): + return self._capability + + def is_device_capability_family(self, capability: int, device_id: int = 0) -> bool: + # 与 Platform.is_device_capability_family 保持相同语义。 + current = self.get_device_capability(device_id=device_id) + if current is None: + return False + return (current.to_int() // 10) == (capability // 10) + + +def _query_fa4_support( + platform: _FakePlatform, fa4_available: bool = True +) -> tuple[bool, str | None]: + """用平台桩调用真实 ``_is_fa4_supported()``。""" + module_names = ("torch", "vllm", "vllm.platforms", _PROBE_NAME) + saved = {key: sys.modules.get(key) for key in module_names} + + # 被测函数不使用 torch；离线环境用空模块满足生产文件的顶层导入。 + fake_torch = types.ModuleType("torch") + fake_torch.bfloat16 = object() + fake_torch.int32 = object() + fake_torch.Tensor = object + fake_torch.ops = types.SimpleNamespace() + sys.modules["torch"] = fake_torch + fake_vllm = types.ModuleType("vllm") + fake_platforms = types.ModuleType("vllm.platforms") + fake_platforms.current_platform = platform + fake_vllm.platforms = fake_platforms + + sys.modules["vllm"] = fake_vllm + sys.modules["vllm.platforms"] = fake_platforms + + try: + spec = importlib.util.spec_from_file_location(_PROBE_NAME, _MODULE_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[_PROBE_NAME] = module + spec.loader.exec_module(module) + + # 目标环境包含 ``cute`` FA4 接口，源码工作树没有；这里避免提前返回。 + module.FA4_AVAILABLE = fa4_available + module.FA4_UNAVAILABLE_REASON = None + + return module._is_fa4_supported() + finally: + for key, value in saved.items(): + if value is None: + sys.modules.pop(key, None) + else: + sys.modules[key] = value + + +def test_musa_pseudo_capability_is_not_fa4_supported() -> None: + """MUSA capability 10.0 不能被判定为支持 FA4。""" + supported, reason = _query_fa4_support( + _FakePlatform("musa", _FakeCapability(10, 0)) + ) + assert supported is False, ( + "MUSA reported FA4 as supported solely because DeviceCapability(10, 0)" + ".to_int() == 100 collides with the CUDA SM100 bucket; the MUSA " + "FlashAttention build only provides FA3." + ) + assert reason, "an explicit unsupported reason must be returned for MUSA" + + +def test_cuda_sm100_fa4_still_supported() -> None: + """CUDA SM100 应保持原有 FA4 支持。""" + supported, reason = _query_fa4_support( + _FakePlatform("cuda", _FakeCapability(10, 0)) + ) + assert supported is True + assert reason is None + + +def test_cuda_sm90_fa4_still_supported() -> None: + """CUDA SM90 应保持原有 FA4 支持。""" + supported, _ = _query_fa4_support(_FakePlatform("cuda", _FakeCapability(9, 0))) + assert supported is True + + +def test_non_musa_family_100_is_not_blanket_disabled() -> None: + """只排除 MUSA，避免误伤其他 family 100 平台。""" + supported, _ = _query_fa4_support(_FakePlatform("xpu", _FakeCapability(10, 0))) + assert supported is True + + +def test_rocm_and_low_capability_unchanged() -> None: + """ROCm 和低能力平台应保持修改前结果。""" + for name in ("rocm", "xpu"): + supported, _ = _query_fa4_support( + _FakePlatform(name, _FakeCapability(8, 0)) + ) + assert supported is False diff --git a/tests/v1/attention/test_triton_unified_attention_dual_entry.py b/tests/v1/attention/test_triton_unified_attention_dual_entry.py new file mode 100644 index 0000000000..5fcce7d31d --- /dev/null +++ b/tests/v1/attention/test_triton_unified_attention_dual_entry.py @@ -0,0 +1,660 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""验证 Triton unified attention 的 TD/非 TD 双入口。 + +这里只验证主机侧选路及两份 kernel 公共逻辑的一致性。真实 Triton 编译和 +MUSA 数值正确性由 kernel 测试覆盖。 +""" + +import ast +import contextlib +import copy +import difflib +import importlib.util +import pathlib +import sys +import types + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +_OPS_PATH = ( + _REPO_ROOT / "vllm" / "v1" / "attention" / "ops" / "triton_unified_attention.py" +) + +_NON_TD_ENTRY = "kernel_unified_attention" +_TD_ENTRY = "kernel_unified_attention_td" +_TD_HELPERS = ("_load_q_td", "_load_kv_tile_td", "_store_output_td") +_TD_API_ATTR = "make_tensor_descriptor" +_TD_ARG_NAMES = ("USE_TD", "USE_TD_QO") +_TD_ONLY_NAME = "q_block_local_len" + +_MOD_NAME = "_triton_unified_attention_dual_entry_test" + + +# --------------------------------------------------------------------------- +# AST 辅助函数 +# --------------------------------------------------------------------------- + + +def _module_ast() -> ast.Module: + assert _OPS_PATH.is_file(), f"op file not found: {_OPS_PATH}" + return ast.parse(_OPS_PATH.read_text(encoding="utf-8")) + + +def _top_level_functions(tree: ast.Module) -> dict[str, ast.FunctionDef]: + return { + node.name: node + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + } + + +def _direct_callees(fn: ast.FunctionDef, known: set[str]) -> set[str]: + out: set[str] = set() + for node in ast.walk(fn): + if isinstance(node, ast.Call): + func = node.func + if isinstance(func, ast.Name) and func.id in known: + out.add(func.id) + return out + + +def _transitive_closure(fn_name: str, fns: dict[str, ast.FunctionDef]) -> set[str]: + known = set(fns) + seen: set[str] = set() + stack = [fn_name] + while stack: + cur = stack.pop() + for callee in _direct_callees(fns[cur], known): + if callee not in seen: + seen.add(callee) + if callee in fns: + stack.append(callee) + return seen + + +def _has_td_api(fn: ast.FunctionDef) -> bool: + return any( + isinstance(node, ast.Attribute) and node.attr == _TD_API_ATTR + for node in ast.walk(fn) + ) + + +# --------------------------------------------------------------------------- +# 双入口归一化比较 +# --------------------------------------------------------------------------- + +# 比较前会从两侧移除允许存在的差异：函数名、TD 专属 constexpr 参数、 +# USE_TD/USE_TD_QO 分支，以及仅用于 TD descriptor shape 的 +# q_block_local_len。剩余差异均视为非预期漂移。 +_ALLOWED_BRANCH_TESTS = frozenset(_TD_ARG_NAMES) + + +class _TDStripper(ast.NodeTransformer): + """移除 TD 专属语句和分支，保留指针回退。""" + + def visit_If(self, node: ast.If) -> object: + self.generic_visit(node) + test = node.test + if isinstance(test, ast.Name) and test.id in _ALLOWED_BRANCH_TESTS: + if node.orelse: + # 保留 ``if USE_TD_QO`` 的指针回退分支。 + return node.orelse + # 完全移除 ``if USE_TD: tl.static_assert(...)``。 + return None + return node + + def visit_Assign(self, node: ast.Assign) -> object: + self.generic_visit(node) + if any( + isinstance(t, ast.Name) and t.id == _TD_ONLY_NAME for t in node.targets + ): + return None + return node + + +def _normalize_kernel(fn: ast.FunctionDef) -> ast.FunctionDef: + """复制 ``fn`` 并移除 TD 专属结构。""" + out = copy.deepcopy(fn) + out.name = "" + + args = out.args + n_defaults = len(args.defaults) + first_default_idx = len(args.args) - n_defaults + new_args: list[ast.arg] = [] + new_defaults: list[ast.expr] = [] + for i, arg in enumerate(args.args): + if arg.arg in _TD_ARG_NAMES: + continue + new_args.append(arg) + if i >= first_default_idx: + new_defaults.append(args.defaults[i - first_default_idx]) + args.args = new_args + args.defaults = new_defaults + + wrapper = ast.Module(body=out.body, type_ignores=[]) + wrapper = _TDStripper().visit(wrapper) + ast.fix_missing_locations(wrapper) + out.body = wrapper.body + return out + + +def dual_entry_normalized_diff() -> tuple[bool, str, str]: + """返回归一化后两个 kernel 的相等状态、非 TD 源码和差异。""" + fns = _top_level_functions(_module_ast()) + assert _NON_TD_ENTRY in fns and _TD_ENTRY in fns, "both kernel entries must exist" + non_td_src = ast.unparse(_normalize_kernel(fns[_NON_TD_ENTRY])) + td_src = ast.unparse(_normalize_kernel(fns[_TD_ENTRY])) + diff = "\n".join( + difflib.unified_diff( + non_td_src.splitlines(), + td_src.splitlines(), + "non_td_normalized", + "td_normalized", + lineterm="", + ) + ) + return non_td_src == td_src, non_td_src, diff + + +def test_dual_entry_common_logic_matches_after_td_normalization() -> None: + """两个入口只能存在已允许的 TD 结构差异。""" + equal, _, diff = dual_entry_normalized_diff() + assert equal, ( + "non-TD and TD entries diverge beyond the allowed TD-only constructs; " + "the duplicated body has drifted:\n" + diff + ) + + +def test_td_entry_retains_td_helpers_and_api() -> None: + """TD 入口必须继续依赖 TD helper 和 API。""" + fns = _top_level_functions(_module_ast()) + closure = _transitive_closure(_TD_ENTRY, fns) | {_TD_ENTRY} + for helper in _TD_HELPERS: + assert helper in closure, f"{_TD_ENTRY} must still reach {helper}" + assert _has_td_api(fns[_TD_HELPERS[0]]), ( + f"{_TD_HELPERS[0]} must still call tl.{_TD_API_ATTR}" + ) + + +# --------------------------------------------------------------------------- +# 使用轻量桩模块执行真实 wrapper 逻辑 +# --------------------------------------------------------------------------- + + +class _StubPlatform: + """``current_platform`` 的最小桩实现。""" + + device_type = "cpu" + + def __init__(self, name: str = "cpu", capability_int: int = 80) -> None: + self._name = name + self._capability_int = capability_int + + def fp8_dtype(self): + import torch + + return getattr(torch, "float8_e4m3fn", torch.float16) + + def is_musa(self) -> bool: + return self._name == "musa" + + def is_xpu(self) -> bool: + return self._name == "xpu" + + def is_cuda(self) -> bool: + return self._name == "cuda" + + def is_rocm(self) -> bool: + return self._name == "rocm" + + def is_device_capability_family(self, capability: int, device_id: int = 0) -> bool: + return (self._capability_int // 10) == (capability // 10) + + +def _install_stubs(td_api_present: bool, platform: _StubPlatform) -> types.ModuleType: + """注册 ``vllm.*`` 桩模块并返回 ``tl`` 桩对象。""" + import logging + + def _mod(name: str) -> types.ModuleType: + m = types.ModuleType(name) + sys.modules[name] = m + return m + + for pkg in ("vllm", "vllm.v1", "vllm.v1.attention", "vllm.v1.attention.ops"): + _mod(pkg) + + envs = _mod("vllm.envs") + envs.VLLM_BATCH_INVARIANT = False + envs.VLLM_TRITON_ATTN_USE_TD = None + + logger_m = _mod("vllm.logger") + logger_m.init_logger = lambda name: logging.getLogger(name) + + platforms = _mod("vllm.platforms") + platforms.current_platform = platform + + triton_utils = _mod("vllm.triton_utils") + + class _TL: + # 只提供模块导入时需要的属性，不实现兜底 __getattr__，确保 + # hasattr(tl, "make_tensor_descriptor") 能真实反映 API 是否存在。 + constexpr = "constexpr" + int64 = "int64" + int32 = "int32" + int1 = "int1" + + class _Triton: + def jit(self, fn=None, **kw): + if fn is None: + return lambda f: f + return fn + + @staticmethod + def next_power_of_2(n) -> int: + return 1 << (max(int(n), 1) - 1).bit_length() + + tl = _TL() + if td_api_present: + tl.make_tensor_descriptor = lambda **kw: None + triton_utils.tl = tl + triton_utils.triton = _Triton() + + helpers = _mod("vllm.v1.attention.ops.triton_attention_helpers") + for name in ( + "apply_alibi_to_score", + "apply_softcap", + "cdiv_fn", + "compute_kv_seq_mask", + "compute_tile_loop_bounds", + "find_seq_idx", + "init_softmax_M", + "load_qq_bias_tile", + "resolve_seq_and_query_len", + "softmax_step", + "store_segm_reduce_scalars", + ): + setattr(helpers, name, lambda *a, **k: None) + + kv = _mod("vllm.v1.kv_cache_interface") + + class KVQuantMode: + NONE = 0 + FP8_PER_TENSOR = 1 + INT8_PER_TOKEN_HEAD = 2 + FP8_PER_TOKEN_HEAD = 3 + + kv.KVQuantMode = KVQuantMode + return tl + + +_STUB_NAMES = ( + "vllm", + "vllm.v1", + "vllm.v1.attention", + "vllm.v1.attention.ops", + "vllm.envs", + "vllm.logger", + "vllm.platforms", + "vllm.triton_utils", + "vllm.v1.attention.ops.triton_attention_helpers", + "vllm.v1.kv_cache_interface", + _MOD_NAME, +) + + +@contextlib.contextmanager +def _stubbed_vllm(td_api_present: bool, platform: _StubPlatform): + """在上下文内安装桩模块，并在退出时恢复。""" + saved = {name: sys.modules.get(name) for name in _STUB_NAMES} + try: + _install_stubs(td_api_present=td_api_present, platform=platform) + yield + finally: + for name, previous in saved.items(): + if previous is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + +def _load_op_module(td_api_present: bool = True, platform: _StubPlatform | None = None): + """使用全新桩模块加载真实算子模块。""" + with _stubbed_vllm( + td_api_present=td_api_present, platform=platform or _StubPlatform() + ): + spec = importlib.util.spec_from_file_location(_MOD_NAME, _OPS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[_MOD_NAME] = module + spec.loader.exec_module(module) + return module + + +class _Launch: + def __init__(self, recorder: "_KernelRecorder") -> None: + self._recorder = recorder + + def __call__(self, *args, **kwargs) -> None: + self._recorder.launches.append(kwargs) + + +class _KernelRecorder: + """记录 ``kernel[grid](...)`` 调用的 Triton kernel 桩。""" + + def __init__(self, name: str) -> None: + self.name = name + self.launches: list[dict] = [] + self.grids: list[tuple] = [] + + def __getitem__(self, grid): + self.grids.append(grid) + return _Launch(self) + + +class _TrapKernel: + """被索引即报错，用于模拟 Triton 构建 kernel cache key。""" + + def __init__(self, name: str) -> None: + self.name = name + + def __getitem__(self, grid): + raise AssertionError( + f"{self.name} was indexed (its Triton cache key would be evaluated) " + f"on a path that must not touch it" + ) + + +def _small_attention_inputs(module, head_size: int = 256): + """构造最小可用的二维启动参数。""" + import torch + + num_query_heads, num_kv_heads = 4, 4 + num_seqs = 1 + query = torch.randn(2, num_query_heads, head_size, dtype=torch.bfloat16) + key_cache = torch.randn(4, 16, num_kv_heads, head_size, dtype=torch.bfloat16) + value_cache = torch.randn_like(key_cache) + out = torch.empty_like(query) + cu_seqlens_q = torch.tensor([0, 2], dtype=torch.int32) + seqused_k = torch.tensor([2], dtype=torch.int32) + block_table = torch.zeros((num_seqs, 1), dtype=torch.int32) + return dict( + q=query, + k=key_cache, + v=value_cache, + out=out, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=2, + seqused_k=seqused_k, + max_seqlen_k=2, + softmax_scale=head_size**-0.5, + causal=True, + window_size=(-1, -1), + block_table=block_table, + softcap=0, + q_descale=None, + k_descale=None, + v_descale=None, + ) + + +def _require_torch() -> None: + try: + import torch # noqa: F401 + except ImportError: # pragma: no cover + import unittest + + raise unittest.SkipTest("torch is unavailable for the runtime selection tests") + + +def _run_selection(use_td: bool, td_api_present: bool = True): + """执行一次 wrapper，并返回入口选择及两个 recorder。""" + _require_torch() + platform = _StubPlatform("musa") + module = _load_op_module(td_api_present=td_api_present, platform=platform) + nontd = _KernelRecorder(_NON_TD_ENTRY) + td = _KernelRecorder(_TD_ENTRY) + module.kernel_unified_attention = nontd + module.kernel_unified_attention_td = td + module.unified_attention(**_small_attention_inputs(module), use_td=use_td) + chosen = _NON_TD_ENTRY if nontd.launches else _TD_ENTRY + return chosen, nontd, td + + +# ---- wrapper 入口选择 ---------------------------------------------------- + + +def test_wrapper_uses_non_td_entry_when_td_disabled() -> None: + chosen, nontd, td = _run_selection(use_td=False) + assert chosen == _NON_TD_ENTRY + assert not td.launches, "TD entry must not be launched when use_td=False" + + +def test_wrapper_uses_td_entry_when_use_td_true() -> None: + chosen, nontd, td = _run_selection(use_td=True) + assert chosen == _TD_ENTRY + assert not nontd.launches + assert td.launches[0]["USE_TD"] is True + assert "USE_TD_QO" in td.launches[0] + + +def test_wrapper_uses_td_entry_when_use_td_qo_true() -> None: + """验证 2 的幂 head size 下 use_td_qo 会选择 TD 入口。""" + chosen, _, td = _run_selection(use_td=True) + assert chosen == _TD_ENTRY + assert td.launches[0]["USE_TD_QO"] is True + + +def test_musa_default_selects_non_td_entry() -> None: + """MUSA 默认不启用 TD，因此应选择非 TD 入口。""" + chosen, nontd, _ = _run_selection(use_td=False) + assert chosen == _NON_TD_ENTRY + assert nontd.launches, "the non-TD entry must be the one launched" + + +def test_non_td_launch_never_indexes_td_entry() -> None: + """非 TD 路径不能索引 TD kernel。""" + _require_torch() + module = _load_op_module(td_api_present=True, platform=_StubPlatform("musa")) + nontd = _KernelRecorder(_NON_TD_ENTRY) + module.kernel_unified_attention = nontd + module.kernel_unified_attention_td = _TrapKernel(_TD_ENTRY) + module.unified_attention(**_small_attention_inputs(module), use_td=False) + assert len(nontd.launches) == 1 + + +def test_non_td_launch_does_not_pass_td_constexprs() -> None: + _, nontd, _ = _run_selection(use_td=False) + kwargs = nontd.launches[0] + assert "USE_TD" not in kwargs + assert "USE_TD_QO" not in kwargs + + +# ---- 旧 Triton API 兼容性 ----------------------------------------------- +def test_explicit_td_without_td_api_raises_clear_host_error() -> None: + """缺少 TD API 时，显式 TD 请求应在主机侧清晰报错。""" + _require_torch() + module = _load_op_module(td_api_present=False, platform=_StubPlatform("xpu")) + module.kernel_unified_attention = _KernelRecorder(_NON_TD_ENTRY) + module.kernel_unified_attention_td = _TrapKernel(_TD_ENTRY) + try: + module.unified_attention(**_small_attention_inputs(module), use_td=True) + except RuntimeError as exc: + message = str(exc) + else: # pragma: no cover - failure path + raise AssertionError( + "requesting TD without tl.make_tensor_descriptor must raise RuntimeError" + ) + lowered = message.lower() + assert "unavailable" in lowered, ( + "error must state tensor descriptors are unavailable" + ) + assert "does not support" in lowered, "error must state the current Triton lacks it" + assert "VLLM_TRITON_ATTN_USE_TD=0" in message, ( + "error must tell the user they can disable TD" + ) + assert "triton version" in lowered, ( + "error must offer a compatible Triton version as an alternative" + ) + + +def test_non_td_call_works_without_td_api() -> None: + """缺少 TD API 的 Triton 仍应能使用非 TD 路径。""" + _require_torch() + module = _load_op_module(td_api_present=False, platform=_StubPlatform("musa")) + nontd = _KernelRecorder(_NON_TD_ENTRY) + module.kernel_unified_attention = nontd + module.kernel_unified_attention_td = _TrapKernel(_TD_ENTRY) + module.unified_attention(**_small_attention_inputs(module), use_td=False) + assert len(nontd.launches) == 1 + + +# --------------------------------------------------------------------------- +# tuned_large_head 的 CUDA 平台门禁 +# --------------------------------------------------------------------------- +# +# ``tuned_large_head`` 会把 prefill KV tile 扩到 128。MUSA 的伪 SM100 +# 同样命中能力判断，但 128 KiB tile 会超过其 72 KiB 共享内存上限。以下 +# 测试执行真实主机选路逻辑并读取启动参数。 + + +def _attention_inputs( + head_size: int = 256, + num_query_heads: int = 4, + num_kv_heads: int = 4, + query_len: int = 17, + kv_len: int = 257, + block_size: int = 16, +): + """构造有效的二维 ``unified_attention`` 启动参数。""" + import torch + + num_blocks = (kv_len + block_size - 1) // block_size + query = torch.randn(query_len, num_query_heads, head_size, dtype=torch.bfloat16) + key_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, dtype=torch.bfloat16 + ) + value_cache = torch.randn_like(key_cache) + out = torch.empty_like(query) + cu_seqlens_q = torch.tensor([0, query_len], dtype=torch.int32) + seqused_k = torch.tensor([kv_len], dtype=torch.int32) + block_table = torch.zeros((1, num_blocks), dtype=torch.int32) + return dict( + q=query, + k=key_cache, + v=value_cache, + out=out, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=query_len, + seqused_k=seqused_k, + max_seqlen_k=kv_len, + softmax_scale=head_size**-0.5, + causal=True, + window_size=(-1, -1), + block_table=block_table, + softcap=0, + q_descale=None, + k_descale=None, + v_descale=None, + ) + + +def _capture_launch(platform: _StubPlatform, **inputs): + """执行一次真实 wrapper 并返回启动 constexpr 和 grid。""" + _require_torch() + module = _load_op_module(td_api_present=True, platform=platform) + nontd = _KernelRecorder(_NON_TD_ENTRY) + module.kernel_unified_attention = nontd + module.kernel_unified_attention_td = _TrapKernel(_TD_ENTRY) + module.unified_attention(**_attention_inputs(**inputs), use_td=False) + assert len(nontd.launches) == 1, "the non-TD entry must launch exactly once" + return nontd.launches[0], nontd.grids[0] + + +def test_musa_pseudo_sm100_prefill_skips_b200_tuned_large_head() -> None: + """MUSA 的伪 SM100 不能使用 B200 宽 tile。""" + kwargs, _grid = _capture_launch( + _StubPlatform("musa", 100), + head_size=256, + num_query_heads=4, + num_kv_heads=4, + query_len=17, + kv_len=257, + block_size=16, + ) + assert kwargs["TILE_SIZE"] == 16, ( + f"MUSA prefill TILE_SIZE must be 16, got {kwargs['TILE_SIZE']}" + ) + assert kwargs["BLOCK_M"] == 4, f"MUSA BLOCK_M must be 4, got {kwargs['BLOCK_M']}" + assert kwargs.get("num_warps") == 4, ( + f"MUSA num_warps must be 4, got {kwargs.get('num_warps')}" + ) + assert kwargs.get("num_stages") == 1, ( + f"MUSA num_stages must be 1, got {kwargs.get('num_stages')}" + ) + + +def test_cuda_sm100_prefill_keeps_b200_tuned_large_head() -> None: + """CUDA SM100 必须保留原有 B200 优化。""" + kwargs, _grid = _capture_launch( + _StubPlatform("cuda", 100), + head_size=256, + num_query_heads=4, + num_kv_heads=4, + query_len=17, + kv_len=257, + block_size=16, + ) + assert kwargs["TILE_SIZE"] == 128, ( + f"CUDA SM100 TILE_SIZE must stay 128, got {kwargs['TILE_SIZE']}" + ) + assert kwargs["BLOCK_M"] == 32, ( + f"CUDA SM100 BLOCK_M must be 32, got {kwargs['BLOCK_M']}" + ) + assert kwargs.get("num_warps") == 8, ( + f"CUDA SM100 num_warps must be 8, got {kwargs.get('num_warps')}" + ) + assert kwargs.get("num_stages") == 2, ( + f"CUDA SM100 num_stages must be 2, got {kwargs.get('num_stages')}" + ) + + +def test_cuda_non_sm100_prefill_skips_tuned_large_head() -> None: + """低于 SM100 的 CUDA 设备也不能使用宽 tile。""" + kwargs, _grid = _capture_launch( + _StubPlatform("cuda", 80), + head_size=256, + num_query_heads=4, + num_kv_heads=4, + query_len=17, + kv_len=257, + block_size=16, + ) + assert kwargs["TILE_SIZE"] == 32, ( + f"CUDA non-SM100 TILE_SIZE must be the default 32, got {kwargs['TILE_SIZE']}" + ) + assert kwargs["BLOCK_M"] == 16, ( + f"CUDA non-SM100 BLOCK_M must be 16, got {kwargs['BLOCK_M']}" + ) + assert "num_warps" not in kwargs, "CUDA non-SM100 must not pin num_warps" + assert "num_stages" not in kwargs, "CUDA non-SM100 must not pin num_stages" + + +def test_musa_decode_path_unchanged() -> None: + """MUSA decode 的启动参数必须保持不变。""" + kwargs, _grid = _capture_launch( + _StubPlatform("musa", 100), + head_size=256, + num_query_heads=4, + num_kv_heads=4, + query_len=1, + kv_len=257, + block_size=16, + ) + assert kwargs["TILE_SIZE"] == 16, ( + f"MUSA decode TILE_SIZE must be 16, got {kwargs['TILE_SIZE']}" + ) + assert kwargs["BLOCK_M"] == 4, ( + f"MUSA decode BLOCK_M must be 4, got {kwargs['BLOCK_M']}" + ) + assert kwargs.get("num_warps") == 4 + assert kwargs.get("num_stages") == 1 diff --git a/tests/v1/attention/test_triton_unified_attention_td_dependency_isolation.py b/tests/v1/attention/test_triton_unified_attention_td_dependency_isolation.py new file mode 100644 index 0000000000..c1f7a1850b --- /dev/null +++ b/tests/v1/attention/test_triton_unified_attention_td_dependency_isolation.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""回归保护：隔离 Triton tensor descriptor（TD）依赖。 + +Triton 在 constexpr 死代码消除前遍历调用图，因此非 TD kernel 的词法依赖 +不能触达 ``make_tensor_descriptor``。这里使用 AST 验证源码结构；真实 MUSA +kernel 和 head_dim=512 数值正确性由 kernel 测试覆盖。 +""" + +import ast +import pathlib + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +_KERNEL_PATH = ( + _REPO_ROOT / "vllm" / "v1" / "attention" / "ops" / "triton_unified_attention.py" +) + +# MUSA 在 Triton 3.2.0 上必须能够编译的入口。 +_NON_TD_ENTRY = "kernel_unified_attention" +# 会调用 ``tl.make_tensor_descriptor`` 的 helper。 +_TD_HELPERS = ("_load_q_td", "_load_kv_tile_td", "_store_output_td") +# 目标 Triton 版本缺少的 tensor descriptor API。 +_TD_API_ATTR = "make_tensor_descriptor" +# 确实共享且不依赖 TD 的 helper，用于正向对照。 +_TD_FREE_SHARED = ("_cast_kv_tile", "reduce_segments") + + +def _module_ast() -> ast.Module: + assert _KERNEL_PATH.is_file(), f"kernel file not found: {_KERNEL_PATH}" + return ast.parse(_KERNEL_PATH.read_text(encoding="utf-8")) + + +def _top_level_functions(tree: ast.Module) -> dict[str, ast.FunctionDef]: + return { + node.name: node + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + } + + +def _is_triton_jit(fn: ast.FunctionDef) -> bool: + for dec in fn.decorator_list: + if isinstance(dec, ast.Attribute) and dec.attr == "jit": + return True + if isinstance(dec, ast.Name) and dec.id == "jit": + return True + return False + + +def _direct_callees(fn: ast.FunctionDef, known: set[str]) -> set[str]: + """返回 ``fn`` 直接调用的模块级函数名。""" + out: set[str] = set() + for node in ast.walk(fn): + if isinstance(node, ast.Call): + func = node.func + if isinstance(func, ast.Name) and func.id in known: + out.add(func.id) + return out + + +def _transitive_closure(fn_name: str, fns: dict[str, ast.FunctionDef]) -> set[str]: + """返回从 ``fn_name`` 可达的词法调用闭包，不含自身。""" + known = set(fns) + seen: set[str] = set() + stack = [fn_name] + while stack: + cur = stack.pop() + for callee in _direct_callees(fns[cur], known): + if callee not in seen: + seen.add(callee) + if callee in fns: + stack.append(callee) + return seen + + +def _has_td_api(fn: ast.FunctionDef) -> bool: + return any( + isinstance(node, ast.Attribute) and node.attr == _TD_API_ATTR + for node in ast.walk(fn) + ) + + +def _referenced_names(fn: ast.FunctionDef) -> set[str]: + names: set[str] = set() + for node in ast.walk(fn): + if isinstance(node, ast.Name): + names.add(node.id) + elif isinstance(node, ast.Attribute): + names.add(node.attr) + return names + + +def test_td_helpers_exist_and_bind_the_td_api() -> None: + """确认 TD helper 存在、使用 @triton.jit 且调用 TD API。""" + fns = _top_level_functions(_module_ast()) + for name in _TD_HELPERS: + assert name in fns, f"{name} is missing from {_KERNEL_PATH.name}" + assert _is_triton_jit(fns[name]), f"{name} must be decorated with @triton.jit" + assert _has_td_api(fns[name]), f"{name} must call tl.{_TD_API_ATTR}" + + +def test_non_td_entry_closure_excludes_td_helpers() -> None: + """非 TD kernel 的词法闭包不能触达 TD helper。""" + fns = _top_level_functions(_module_ast()) + closure = _transitive_closure(_NON_TD_ENTRY, fns) + leaked = sorted(set(_TD_HELPERS) & closure) + assert not leaked, ( + f"{_NON_TD_ENTRY} lexically reaches TD helpers {leaked}; Triton's " + f"DependenciesFinder would recurse into them and hit " + f"tl.{_TD_API_ATTR} before constexpr dead-code elimination." + ) + + +def test_non_td_entry_closure_excludes_td_api() -> None: + """非 TD kernel 可达的函数不能引用 TD API。""" + fns = _top_level_functions(_module_ast()) + closure = _transitive_closure(_NON_TD_ENTRY, fns) | {_NON_TD_ENTRY} + offenders = sorted(name for name in closure if _has_td_api(fns[name])) + assert not offenders, ( + f"{_NON_TD_ENTRY} dependency closure contains TD API users {offenders}; " + f"the non-TD JIT entry must be physically free of " + f"tl.{_TD_API_ATTR}." + ) + + +def test_non_td_entry_body_has_no_td_names() -> None: + """非 TD kernel 主体中不能出现 TD 符号。""" + fns = _top_level_functions(_module_ast()) + forbidden = set(_TD_HELPERS) | {_TD_API_ATTR} + leaked = sorted(forbidden & _referenced_names(fns[_NON_TD_ENTRY])) + assert not leaked, ( + f"{_NON_TD_ENTRY} body still references TD names {leaked}; even " + f"constexpr-guarded references are resolved by DependenciesFinder." + ) + + +def test_shared_td_free_helpers_stay_clean() -> None: + """作为正向对照，共享 helper 必须保持无 TD 依赖。""" + fns = _top_level_functions(_module_ast()) + for name in _TD_FREE_SHARED: + assert name in fns, f"{name} is missing from {_KERNEL_PATH.name}" + assert not _has_td_api(fns[name]), f"{name} must not use tl.{_TD_API_ATTR}" + leaked = sorted(set(_TD_HELPERS) & _transitive_closure(name, fns)) + assert not leaked, f"{name} unexpectedly reaches TD helpers {leaked}" diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index a802e63167..c04e436a21 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -176,8 +176,378 @@ def _store_output_td( ) +# --------------------------------------------------------------------------- +# 非 TD 入口（指针路径），供 MUSA 及其他非 TD 调用方使用。 +# +# 此入口在词法依赖上不能引用 TD helper 或 ``tl.make_tensor_descriptor``。 +# Triton 会在 constexpr 死代码消除前构建 JIT cache key 并遍历调用闭包；若 +# 只用 ``USE_TD`` 保护单一 kernel，缺少 TD API 的 Triton 仍会解析 TD helper +# 并抛出 ``AttributeError``。因此 TD 与非 TD 入口必须在代码结构上分离。 +# --------------------------------------------------------------------------- @triton.jit def kernel_unified_attention( + # 二维路径的输出地址。三维模式将分段中间结果写入 ``segm_*``， + # ``output_ptr`` 不参与计算。 + output_ptr, + # 输入 + query_ptr, + key_cache_ptr, + value_cache_ptr, + sink_ptr, + block_tables_ptr, + seq_lens_ptr, + alibi_slopes_ptr, + qq_bias_ptr, + # 标量参数 + scale, + q_scale, + k_scale, + v_scale, + out_scale, + softcap, + num_query_heads: tl.constexpr, # int + num_queries_per_kv: 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 + qq_bias_stride_0: tl.int64, # int + BLOCK_SIZE: tl.constexpr, # int + TILE_SIZE: tl.constexpr, # int must be power of 2 + HEAD_SIZE: tl.constexpr, # int + HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2 + USE_ALIBI_SLOPES: tl.constexpr, # bool + USE_ALIBI_SQRT: tl.constexpr, # bool + USE_QQ_BIAS: tl.constexpr, # bool + USE_SOFTCAP: tl.constexpr, # bool + USE_SINKS: tl.constexpr, # bool + SLIDING_WINDOW: tl.constexpr, # int + USE_CAUSAL: tl.constexpr, # bool + USE_PER_SEQ_CAUSAL: tl.constexpr, # bool + per_seq_causal_ptr, # [num_seqs] bool, or None + USE_MM_PREFIX: tl.constexpr, # bool + MAX_MM_RANGES: tl.constexpr, # int + mm_prefix_range_ptr, + 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.constexpr, # 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.constexpr, # int + query_start_len_ptr, + BLOCK_Q: tl.constexpr, + num_seqs: tl.int32, + BLOCK_M: tl.constexpr, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, + USE_FP8: tl.constexpr, + # 切换二维/三维布局。二维路径在一个 tile 循环中处理完整序列并写入 + # ``output_ptr``；三维路径按 segment 限定循环并写入分段中间结果，最后 + # 由 ``reduce_segments`` 合并。 + IS_3D: tl.constexpr, + # 以下参数默认为 None，使 Triton 能在对应 constexpr 分支不可达时跳过实例化。 + # @quinnlp 在 intel/intel-xpu-backend-for-triton#6758 的评审意见 + # r3204641104 中指出了这里的性能回退。 + # 分段输出仅用于三维模式，二维模式不使用。 + segm_output_ptr=None, + segm_max_ptr=None, + segm_expsum_ptr=None, + # 每个 token/head 的 scale cache，仅在 KV_QUANT_MODE 为 2 或 3 时使用。 + k_scale_cache_ptr=None, + v_scale_cache_ptr=None, + stride_ks_blk: tl.int64 = None, + stride_ks_slot: tl.int64 = None, + stride_ks_head: tl.int64 = None, + stride_vs_blk: tl.int64 = None, + stride_vs_slot: tl.int64 = None, + stride_vs_head: tl.int64 = None, + # KV cache 量化模式由 constexpr 分支处理：NONE (0)、 + # FP8_PER_TENSOR (1)、INT8_PER_TOKEN_HEAD (2)、FP8_PER_TOKEN_HEAD (3)。 + KV_QUANT_MODE: tl.constexpr = 0, + FP8_MIN: tl.constexpr = float8_info.min, + FP8_MAX: tl.constexpr = float8_info.max, + # 分块/块内注意力。``CHUNK_LOOKBACK >= 0`` 启用 Gemma3 块内层使用的 + # 分块 mask，并在 helper 中优先于 ``SLIDING_WINDOW``；``-1`` 表示关闭。 + CHUNK_LOOKBACK: tl.constexpr = -1, + CHUNK_SIZE: tl.constexpr = -1, + Q_IS_FP8: tl.constexpr = False, +): + USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2 + USE_FP8_Q_DESCALE: tl.constexpr = KV_QUANT_MODE == 1 and Q_IS_FP8 + + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + segm_idx = tl.program_id(2) if IS_3D else 0 + + ( + seq_idx, + q_block_local_idx, + cur_batch_in_all_start_index, + cur_batch_query_len, + seq_len, + ) = resolve_seq_and_query_len( + query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q + ) + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + if IS_3D: + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len: + return + else: + tiles_per_segment = 0 + + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d[None, :] + ) + + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + # Q : (BLOCK_M, HEAD_SIZE_PADDED) + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=0.0, + ) + + block_table_offset = seq_idx * block_table_stride + + M = init_softmax_M( + sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D + ) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + # acc : (BLOCK_M, HEAD_SIZE_PADDED) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32) + score_scale = scale + value_scale = 1.0 + if USE_FP8_Q_DESCALE: + score_scale = scale * tl.load(q_scale) * tl.load(k_scale) + value_scale = tl.load(v_scale) + + context_len = seq_len - cur_batch_query_len + + if USE_ALIBI_SLOPES: + alibi_slope = tl.load( + alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0 + ) + + if USE_QQ_BIAS: + qq_bias_row_ptrs = qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0 + + loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds( + context_len, + seq_len, + cur_batch_query_len, + q_block_local_idx, + segm_idx, + tiles_per_segment, + TILE_SIZE, + BLOCK_M, + BLOCK_Q, + num_queries_per_kv, + SLIDING_WINDOW, + USE_MM_PREFIX, + IS_3D, + USE_CAUSAL, + USE_PER_SEQ_CAUSAL, + CHUNK_LOOKBACK, + CHUNK_SIZE, + ) + + # 遍历滑动窗口范围内的 tile。 + for j in range(loop_lo, loop_hi): + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + + v_offset = ( + physical_block_idx[:, None] * stride_v_cache_0 + + kv_head_idx * stride_v_cache_2 + + offs_d[None, :] * stride_v_cache_3 + + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1 + ) + k_offset = ( + physical_block_idx[None, :] * stride_k_cache_0 + + kv_head_idx * stride_k_cache_2 + + offs_d[:, None] * stride_k_cache_3 + + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1 + ) + # K : (HEAD_SIZE, TILE_SIZE) + K_load = tl.load( + key_cache_ptr + k_offset, + mask=dim_mask[:, None] & tile_mask[None, :], + other=0.0, + ) + # V : (TILE_SIZE, HEAD_SIZE) + V_load = tl.load( + value_cache_ptr + v_offset, + mask=dim_mask[None, :] & tile_mask[:, None], + other=0.0, + ) + K = _cast_kv_tile(K_load, Q, k_scale, KV_QUANT_MODE) + V = _cast_kv_tile(V_load, Q, v_scale, KV_QUANT_MODE) + + # INT8/FP8 per-token-head 模式使用的 scale。 + if USE_PER_TOKEN_HEAD_SCALES: + scale_idx = ( + physical_block_idx * stride_ks_blk + + (seq_offset % BLOCK_SIZE) * stride_ks_slot + + kv_head_idx * stride_ks_head + ) + k_token_head_scales = tl.load( + k_scale_cache_ptr + scale_idx, mask=tile_mask, other=1.0 + ) + v_scale_idx = ( + physical_block_idx * stride_vs_blk + + (seq_offset % BLOCK_SIZE) * stride_vs_slot + + kv_head_idx * stride_vs_head + ) + v_token_head_scales = tl.load( + v_scale_cache_ptr + v_scale_idx, mask=tile_mask, other=1.0 + ) + + query_abs_pos = context_len + query_pos[:, None] + seq_mask = compute_kv_seq_mask( + query_abs_pos, + seq_offset, + seq_idx, + seq_len, + mm_prefix_range_ptr, + SLIDING_WINDOW, + USE_MM_PREFIX, + MAX_MM_RANGES, + USE_CAUSAL, + USE_PER_SEQ_CAUSAL, + per_seq_causal_ptr, + CHUNK_LOOKBACK, + CHUNK_SIZE, + ) + + # S : (BLOCK_M, TILE_SIZE) + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32) + if USE_PER_TOKEN_HEAD_SCALES: + # per-token-head 量化将 softmax_scale 与 k_scale 融合，避免对 S + # 额外执行一次 BLOCK_M × TILE_SIZE 乘法。 + S += tl.dot(Q, K) * (score_scale * k_token_head_scales[None, :]) + else: + S += score_scale * tl.dot(Q, K) + + if USE_SOFTCAP: + S = apply_softcap(S, softcap) + + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf") + ) + + if USE_ALIBI_SLOPES: + S = apply_alibi_to_score( + S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT + ) + + if USE_QQ_BIAS: + S += load_qq_bias_tile( + qq_bias_row_ptrs, seq_offset, context_len, qq_bias_stride_0 + ) + + M, L, P, alpha = softmax_step(S, M, L) + acc = acc * alpha[:, None] + + if SLIDING_WINDOW: + qpos_lo = q_block_local_idx * BLOCK_Q + dist = context_len + qpos_lo - seq_offset[:, None] + if USE_PER_SEQ_CAUSAL: + is_causal_seq = tl.load(per_seq_causal_ptr + seq_idx) + sw_mask_v = tl.where( + is_causal_seq, + dist < SLIDING_WINDOW, + (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW), + ) + elif USE_CAUSAL: + sw_mask_v = dist < SLIDING_WINDOW + else: + sw_mask_v = (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW) + V = tl.where(sw_mask_v, V, 0.0) + if USE_PER_TOKEN_HEAD_SCALES: + # per-token-head 量化在 P 上应用 v_scale，而不是在 V 上应用。 + P_v = (P * v_token_head_scales[None, :]).to(V.dtype) + acc += tl.dot(P_v, V) + else: + acc += tl.dot(P.to(V.dtype), V) + + # ---- 收尾 ------------------------------------------------------------- + if IS_3D: + if USE_FP8_Q_DESCALE: + acc *= value_scale + # 保存分段中间结果，随后由 ``reduce_segments`` 完成合并。 + segm_output_offset = ( + query_offset_0[:, None].to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + segm_idx * HEAD_SIZE_PADDED + + tl.arange(0, HEAD_SIZE_PADDED)[None, :] + ) + tl.store( + segm_output_ptr + segm_output_offset, + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + store_segm_reduce_scalars( + segm_max_ptr, + segm_expsum_ptr, + query_offset_0, + query_offset_1, + segm_idx, + M, + L, + query_mask_0, + query_mask_1, + num_query_heads, + NUM_SEGMENTS_PER_SEQ, + ) + else: + acc = acc / L[:, None] + if USE_FP8_Q_DESCALE: + acc *= value_scale + if USE_FP8: + acc = acc * tl.load(out_scale) + acc = tl.clamp(acc, FP8_MIN, FP8_MAX) + output_offset = ( + query_offset_0[:, None] * output_stride_0 + + query_offset_1[:, None] * output_stride_1 + + offs_d[None, :] + ) + tl.store( + output_ptr + output_offset, + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + + +# --------------------------------------------------------------------------- +# TD 入口。这里保留独立的 ``@triton.jit`` 函数并重复上方指针路径主体， +# 避免非 TD 入口在词法依赖上触达 TD helper/API。仅当 ``use_td`` 或 +# ``use_td_qo`` 启用时由 wrapper 选择，并保留拆分前的 TD 分支和指针回退。 +# --------------------------------------------------------------------------- +@triton.jit +def kernel_unified_attention_td( # Output destination for the 2D path. In 3D mode per-segment partials # go to the ``segm_*`` tensors (see bottom of signature) and # ``output_ptr`` is unused (callers may pass any non-null pointer). @@ -874,9 +1244,12 @@ def unified_attention( # head_size 256 with many query rows per sequence (e.g. diffusion-gemma # bidirectional canvas passes) is prefill-shaped, but the decode-oriented # defaults (BLOCK_Q=8, TILE=32, 4 warps) under-tile it. A wider KV tile + - # more query rows per block + 8 warps is ~2x faster on B200. + # B200 上增加每个 block 的 query 行数并使用 8 个 warp 约快 2 倍。 + # MUSA 的伪 SM100 也会命中 capability-family(100)，但宽 tile 会超过其 + # 共享内存上限，因此此优化必须限定为 CUDA。 tuned_large_head = ( - head_size == 256 + current_platform.is_cuda() + and head_size == 256 and max_seqlen_q > 1 and num_queries_per_kv <= 16 and current_platform.is_device_capability_family(100) @@ -1027,7 +1400,28 @@ def unified_attention( if launch_num_stages is not None: launch_kwargs["num_stages"] = launch_num_stages - kernel_unified_attention[grid]( + # 选择 kernel 入口。非 TD 入口不能接收 TD 专属 constexpr；运行时缺少 + # ``tl.make_tensor_descriptor`` 时也不能启动 TD 入口。这里在主机侧提前 + # 报错，避免在构建 cache key 时才出现难以定位的 JIT AttributeError。 + wants_td = use_td or use_td_qo + if wants_td: + if not hasattr(tl, "make_tensor_descriptor"): + raise RuntimeError( + "TRITON_ATTN tensor-descriptor (TD) mode was requested, but " + "this Triton build does not provide " + "``tl.make_tensor_descriptor``: tensor descriptors are " + "unavailable and the current Triton does not support them. " + "Set VLLM_TRITON_ATTN_USE_TD=0 to use the pointer-path " + "kernel, or use a Triton version that provides tensor " + "descriptors." + ) + kernel = kernel_unified_attention_td + td_kwargs: dict[str, Any] = {"USE_TD": use_td, "USE_TD_QO": use_td_qo} + else: + kernel = kernel_unified_attention + td_kwargs = {} + + kernel[grid]( output_ptr=out, segm_output_ptr=segm_output_ptr, segm_max_ptr=segm_max_ptr, @@ -1097,8 +1491,7 @@ def unified_attention( Q_IS_FP8=(q.dtype == current_platform.fp8_dtype()), CHUNK_LOOKBACK=chunk_lookback, CHUNK_SIZE=chunk_size, - USE_TD=use_td, - USE_TD_QO=use_td_qo, + **td_kwargs, **launch_kwargs, ) diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index f015003367..74d89a5449 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -74,6 +74,13 @@ def _is_fa4_supported() -> tuple[bool, str | None]: return False, f"FA4 is unavailable due to: {FA4_UNAVAILABLE_REASON}" from vllm.platforms import current_platform + # MUSA 上报的伪 DeviceCapability(10, 0) 会被通用能力判断识别为 CUDA + # SM100，但这不能证明 MUSA 存在 CUDA FA4 kernel。MUSA FlashAttention + # 当前只提供 FA3，无法处理 Gemma4 的 global head_dim=512，因此需要显式 + # 排除 MUSA 并回退到现有 TRITON_ATTN 路径。CUDA、ROCm、XPU 行为不变。 + if current_platform.is_musa(): + return False, "FA4 is not supported on MUSA (the MUSA build provides FA3)" + if not ( current_platform.is_device_capability_family(90) or current_platform.is_device_capability_family(100) -- Gitee