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, 06:00 UTC
[AMD] [Quantization] Add override flag for attention dtype instead of using kv_cache_dtype trigger (#17331)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Jun 11, 2025 at 14:53 UTC
c7ea0b56cd9a5e607f30cc637b4c800ce47bafca
3 files changed
+21
-1
vllm/attention/backends/rocm_flash_attn.py
+9
-1
@@ -17,6 +17,7 @@ from vllm.attention.backends.utils import (CommonAttentionState,
17
CommonMetadataBuilder)
18
from vllm.attention.ops.paged_attn import (PagedAttention,
19
PagedAttentionMetadata)
20
+from vllm.config import get_current_vllm_config
21
from vllm.logger import init_logger
22
from vllm.platforms import current_platform
23
from vllm.platforms.rocm import use_rocm_custom_paged_attention
@@ -584,6 +585,10 @@ class ROCmFlashAttentionImpl(AttentionImpl):
585
logger.debug("Using naive (SDPA) attention in ROCmBackend")
586
587
self.aiter_kv_scales_initialized = False
588
+ self.force_fp8_attention = (
589
+ get_current_vllm_config() is not None
590
+ and get_current_vllm_config().model_config.override_attention_dtype
591
+ == "fp8")
592
593
def repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor:
594
"""torch.repeat_interleave(x, dim=1, repeats=n_rep)"""
@@ -770,9 +775,12 @@ class ROCmFlashAttentionImpl(AttentionImpl):
775
query.dtype,
776
seq_lens,
777
make_attn_mask=causal_mask) # type: ignore
778
+
779
use_fp8_scales = (layer._q_scale and layer._k_scale
780
and layer._v_scale and layer._prob_scale
775
- and self.kv_cache_dtype == "fp8")
781
+ and (self.kv_cache_dtype == "fp8"
782
+ or self.force_fp8_attention))
783
+
784
full_scales = (
785
layer._q_scale.item(), layer._k_scale.item(),
786
layer._v_scale.item(),
vllm/config.py
+8
@@ -417,6 +417,8 @@ class ModelConfig:
417
available.\n
418
- "vllm" will use the vLLM model implementation.\n
419
- "transformers" will use the Transformers model implementation."""
420
+ override_attention_dtype: Optional[str] = None
421
+ """Override dtype for attention"""
422
423
def compute_hash(self) -> str:
424
"""
@@ -517,6 +519,12 @@ class ModelConfig:
519
520
from vllm.platforms import current_platform
521
522
+ if (self.override_attention_dtype is not None
523
+ and not current_platform.is_rocm()):
524
+ warnings.warn(
525
+ "override-attention-dtype is set but not using ROCm platform",
526
+ stacklevel=2)
527
+
528
if (self.enable_sleep_mode
529
and not current_platform.is_sleep_mode_available()):
530
raise ValueError(
vllm/engine/arg_utils.py
+4
@@ -429,6 +429,7 @@ class EngineArgs:
429
override_generation_config: dict[str, Any] = \
430
get_field(ModelConfig, "override_generation_config")
431
model_impl: str = ModelConfig.model_impl
432
+ override_attention_dtype: str = ModelConfig.override_attention_dtype
433
434
calculate_kv_scales: bool = CacheConfig.calculate_kv_scales
435
@@ -549,6 +550,8 @@ class EngineArgs:
550
model_group.add_argument("--model-impl",
551
choices=[f.value for f in ModelImpl],
552
**model_kwargs["model_impl"])
553
+ model_group.add_argument("--override-attention-dtype",
554
+ **model_kwargs["override_attention_dtype"])
555
556
# Model loading arguments
557
load_kwargs = get_kwargs(LoadConfig)
@@ -946,6 +949,7 @@ class EngineArgs:
949
override_generation_config=self.override_generation_config,
950
enable_sleep_mode=self.enable_sleep_mode,
951
model_impl=self.model_impl,
952
+ override_attention_dtype=self.override_attention_dtype,
953
)
954
955
def create_load_config(self) -> LoadConfig: