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 14, 05:30 UTC
Force TRTLLM attention for gpt-oss on SM100 (#22678)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 13, 2025 at 00:22 UTC
c6b928798e96f0a99a666945686c63b61bbbced4
4 files changed
+20
-9
vllm/model_executor/models/gpt_oss.py
+1
-4
@@ -8,7 +8,6 @@ import torch.distributed as dist
8
from torch import nn
9
from transformers import GptOssConfig
10
11
-from vllm import envs
11
from vllm.attention import Attention, AttentionType
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, VllmConfig
@@ -70,11 +69,9 @@ class OAIAttention(nn.Module):
69
70
tp_size = get_tensor_model_parallel_world_size()
71
73
- attention_sink_dtype = (torch.float32 if envs.VLLM_USE_TRTLLM_ATTENTION
74
- else torch.bfloat16)
72
self.sinks = torch.nn.Parameter(
73
torch.empty(config.num_attention_heads // tp_size,
77
- dtype=attention_sink_dtype,
74
+ dtype=torch.bfloat16,
75
requires_grad=False))
76
77
self.norm = RMSNorm(config.hidden_size, eps=1e-5)
vllm/utils/flashinfer.py
+8
@@ -154,6 +154,7 @@ def use_trtllm_attention(
154
num_qo_heads: Optional[int],
155
num_kv_heads: Optional[int],
156
attn_head_size: Optional[int],
157
+ has_sinks: bool = False,
158
) -> bool:
159
# Requires SM100 and NVIDIA artifactory to be accessible to download cubins
160
if not (current_platform.is_device_capability(100)
@@ -165,6 +166,13 @@ def use_trtllm_attention(
166
or num_qo_heads % num_kv_heads != 0):
167
return False
168
169
+ # If sinks are being used, we must use TRTLLM attention as it's
170
+ # the only backend that supports them
171
+ if has_sinks:
172
+ logger.info_once(
173
+ "Using TRTLLM attention (required for attention sinks).")
174
+ return True
175
+
176
env_value = envs.VLLM_USE_TRTLLM_ATTENTION
177
if env_value is not None:
178
logger.info_once("VLLM_USE_TRTLLM_ATTENTION is set to %s", env_value)
vllm/v1/attention/backends/flashinfer.py
+7
-4
@@ -523,14 +523,17 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
523
num_kv_heads = self.kv_cache_spec.num_kv_heads
524
head_dim = self.kv_cache_spec.head_size
525
526
+ # Check if any layer uses sinks (requires TRTLLM attention)
527
+ has_sinks = self.global_hyperparameters.has_sinks
528
+
529
# currently prefill trtllm attention does not support fp8 kv cache
530
prefill_use_trtllm = not cache_dtype.startswith("fp8") \
531
and use_trtllm_attention(
532
num_prefill_tokens, max_seq_len, cache_dtype,
530
- num_qo_heads, num_kv_heads, head_dim)
533
+ num_qo_heads, num_kv_heads, head_dim, has_sinks)
534
decode_use_trtllm = use_trtllm_attention(
535
num_decode_tokens, max_seq_len, cache_dtype,
533
- num_qo_heads, num_kv_heads, head_dim)
536
+ num_qo_heads, num_kv_heads, head_dim, has_sinks)
537
538
attn_metadata = FlashInferMetadata(
539
num_actual_tokens=num_actual_tokens,
@@ -642,9 +645,9 @@ class FlashInferImpl(AttentionImpl):
645
f"heads in the layer. Expected {num_heads}, but got "
646
f"{sinks.shape[0]}."
647
)
648
+ # Cast sinks to float32 if needed (FlashInfer requirement)
649
if sinks.dtype != torch.float32:
646
- raise ValueError("Sinks must be of type float32, but got "
647
- f"{sinks.dtype}.")
650
+ sinks = sinks.to(torch.float32)
651
self.sinks = sinks
652
653
def forward(
vllm/v1/attention/backends/utils.py
+4
-1
@@ -285,6 +285,7 @@ class PerLayerParameters:
285
window_left: int
286
logits_soft_cap: Optional[float]
287
sm_scale: float
288
+ has_sinks: bool = False
289
290
291
def get_per_layer_parameters(
@@ -307,9 +308,11 @@ def get_per_layer_parameters(
308
window_left = window_size[0] if window_size is not None else -1
309
logits_soft_cap = getattr(impl, "logits_soft_cap", None)
310
sm_scale = impl.scale
311
+ has_sinks = getattr(impl, "sinks", None) is not None
312
313
per_layer_params[key] = PerLayerParameters(window_left,
312
- logits_soft_cap, sm_scale)
314
+ logits_soft_cap, sm_scale,
315
+ has_sinks)
316
317
return per_layer_params
318