A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[Bug] [ROCm] Fix Llama 4 Enablement Bug on ROCm: V0 ROCmFlashAttentionImpl and Triton Fused MoE bugs (#16198)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com> Co-authored-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: kliuae <kuanfu.liu@embeddedllm.com>
TJian committed
Apr 9, 2025 at 10:12 UTC
2976dc27e9dc2a799db8337cf9825b63a26eeac5
3 files changed
+15
-9
vllm/attention/backends/rocm_flash_attn.py
+4
-1
@@ -471,7 +471,10 @@ class ROCmFlashAttentionImpl(AttentionImpl):
471
if blocksparse_params is not None:
472
raise ValueError(
473
"ROCmFlashAttention does not support blocksparse attention.")
474
-
474
+ if use_irope:
475
+ logger.warning(
476
+ "Using irope in V0 is not supported yet, it will fall back "
477
+ "to global attention for long context.")
478
if logits_soft_cap is None:
479
# In flash-attn, setting logits_soft_cap as 0 means no soft cap.
480
self.logits_soft_cap = 0.0
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
@@ -1002,6 +1002,7 @@ direct_register_custom_op(
1002
op_func=inplace_fused_experts,
1003
mutates_args=["hidden_states"],
1004
fake_impl=inplace_fused_experts_fake,
1005
+ tags=(torch.Tag.needs_fixed_stride_order, ),
1006
)
1007
1008
@@ -1060,6 +1061,7 @@ direct_register_custom_op(
1061
op_func=outplace_fused_experts,
1062
mutates_args=[],
1063
fake_impl=outplace_fused_experts_fake,
1064
+ tags=(torch.Tag.needs_fixed_stride_order, ),
1065
)
1066
1067
vllm/utils.py
+9
-8
@@ -40,7 +40,7 @@ from dataclasses import dataclass, field
40
from functools import cache, lru_cache, partial, wraps
41
from types import MappingProxyType
42
from typing import (TYPE_CHECKING, Any, Callable, Generic, Literal, NamedTuple,
43
- Optional, Type, TypeVar, Union, cast, overload)
43
+ Optional, Tuple, Type, TypeVar, Union, cast, overload)
44
from uuid import uuid4
45
46
import cachetools
@@ -1935,12 +1935,13 @@ vllm_lib = Library("vllm", "FRAGMENT") # noqa
1935
1936
1937
def direct_register_custom_op(
1938
- op_name: str,
1939
- op_func: Callable,
1940
- mutates_args: list[str],
1941
- fake_impl: Optional[Callable] = None,
1942
- target_lib: Optional[Library] = None,
1943
- dispatch_key: str = "CUDA",
1938
+ op_name: str,
1939
+ op_func: Callable,
1940
+ mutates_args: list[str],
1941
+ fake_impl: Optional[Callable] = None,
1942
+ target_lib: Optional[Library] = None,
1943
+ dispatch_key: str = "CUDA",
1944
+ tags: Tuple[torch.Tag, ...] = (),
1945
):
1946
"""
1947
`torch.library.custom_op` can have significant overhead because it
@@ -1979,7 +1980,7 @@ def direct_register_custom_op(
1980
import torch._custom_op.impl
1981
schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
1982
my_lib = target_lib or vllm_lib
1982
- my_lib.define(op_name + schema_str)
1983
+ my_lib.define(op_name + schema_str, tags=tags)
1984
my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
1985
if fake_impl is not None:
1986
my_lib._register_fake(op_name, fake_impl)