@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 007dd90859

Fork 0

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

[gpt-oss] Enable gpt-oss on ampere (#22714)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

Yongye Zhu committed Aug 12, 2025 at 06:21 UTC 007dd90859cc0337510536677418a43d8f66e286

10 files changed +26 -17

tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py

+3 -2

@@ -25,5 +25,6 @@ class DummyPlatform(Platform):

25 compilation_config.custom_ops = ["all"]

27 def get_attn_backend_cls(self, backend_name, head_size, dtype,

28 - kv_cache_dtype, block_size, use_v1, use_mla):

29 - return "vllm_add_dummy_platform.dummy_attention_backend.DummyAttentionBackend" # noqa E501

\ No newline at end of file

28 + kv_cache_dtype, block_size, use_v1, use_mla,

29 + has_sink):

30 + return "vllm_add_dummy_platform.dummy_attention_backend.DummyAttentionBackend" # noqa E501

vllm/attention/layer.py

+3 -1

@@ -138,6 +138,7 @@ class Attention(nn.Module):

138 self.head_size = head_size

139 self.num_kv_heads = num_kv_heads

140 self.sliding_window = sliding_window

141 + self.has_sink = extra_impl_args.get("sinks") is not None

142

143 quant_method = quant_config.get_quant_method(

144 self, prefix=prefix) if quant_config else None

@@ -165,7 +166,8 @@ class Attention(nn.Module):

166 kv_cache_dtype,

167 block_size,

168 is_attention_free,

168 - use_mla=use_mla)

169 + use_mla=use_mla,

170 + has_sink=self.has_sink)

171 else:

172 self.attn_backend = attn_backend

173

vllm/attention/selector.py

+4 -1

@@ -144,6 +144,7 @@ def get_attn_backend(

144 block_size: int,

145 is_attention_free: bool = False,

146 use_mla: bool = False,

147 + has_sink: bool = False,

148 ) -> type[AttentionBackend]:

149 """Selects which attention backend to use and lazily imports it."""

150 # Accessing envs.* behind an @lru_cache decorator can cause the wrong

@@ -158,6 +159,7 @@ def get_attn_backend(

159 is_attention_free=is_attention_free,

160 use_v1=envs.VLLM_USE_V1,

161 use_mla=use_mla,

162 + has_sink=has_sink,

163 )

164

165

@@ -170,6 +172,7 @@ def _cached_get_attn_backend(

172 is_attention_free: bool,

173 use_v1: bool = False,

174 use_mla: bool = False,

175 + has_sink: bool = False,

176 ) -> type[AttentionBackend]:

177 # If there are no attention layers (e.g. we are running Mamba),

178 # use the placeholder NO_ATTENTION

@@ -201,7 +204,7 @@ def _cached_get_attn_backend(

204 # get device-specific attn_backend

205 attention_cls = current_platform.get_attn_backend_cls(

206 selected_backend, head_size, dtype, kv_cache_dtype, block_size, use_v1,

204 - use_mla)

207 + use_mla, has_sink)

208 if not attention_cls:

209 raise ValueError(

210 f"Invalid attention backend for {current_platform.device_name}")

vllm/model_executor/layers/quantization/mxfp4.py

+1 -1

@@ -42,7 +42,7 @@ class Mxfp4Config(QuantizationConfig):

43 @classmethod

44 def get_min_capability(cls) -> int:

45 - return 90

45 + return 80

47 @classmethod

48 def get_name(cls) -> QuantizationMethods:

vllm/platforms/cpu.py

+2 -2

@@ -91,8 +91,8 @@ class CpuPlatform(Platform):

91 @classmethod

92 def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,

93 dtype: torch.dtype, kv_cache_dtype: Optional[str],

94 - block_size: int, use_v1: bool,

95 - use_mla: bool) -> str:

94 + block_size: int, use_v1: bool, use_mla: bool,

95 + has_sink: bool) -> str:

96 if selected_backend and selected_backend != _Backend.TORCH_SDPA:

97 logger.info("Cannot use %s backend on CPU.", selected_backend)

98 if use_mla:

vllm/platforms/cuda.py

+5 -2

@@ -222,8 +222,8 @@ class CudaPlatformBase(Platform):

222

223 @classmethod

224 def get_attn_backend_cls(cls, selected_backend, head_size, dtype,

225 - kv_cache_dtype, block_size, use_v1,

226 - use_mla) -> str:

225 + kv_cache_dtype, block_size, use_v1, use_mla,

226 + has_sink) -> str:

227 if use_mla:

228 # TODO(lucas): refactor to be more concise

229 # we should probably consider factoring out V1 here

@@ -321,6 +321,9 @@ class CudaPlatformBase(Platform):

321

322 # FlashAttention is the default for SM 8.0+ GPUs

323 if cls.has_device_capability(80):

324 + if has_sink:

325 + logger.info_once("Using Triton backend on V1 engine.")

326 + return TRITON_ATTN_VLLM_V1

327 if is_default_backend_supported := is_attn_backend_supported(

328 FLASH_ATTN_V1, head_size, dtype,

329 allow_import_error=False):

vllm/platforms/interface.py

+2 -2

@@ -196,8 +196,8 @@ class Platform:

196 @classmethod

197 def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,

198 dtype: torch.dtype, kv_cache_dtype: Optional[str],

199 - block_size: int, use_v1: bool,

200 - use_mla: bool) -> str:

199 + block_size: int, use_v1: bool, use_mla: bool,

200 + has_sink: bool) -> str:

201 """Get the attention backend class of a device."""

202 return ""

203

vllm/platforms/rocm.py

+2 -2

@@ -188,8 +188,8 @@ class RocmPlatform(Platform):

188

189 @classmethod

190 def get_attn_backend_cls(cls, selected_backend, head_size, dtype,

191 - kv_cache_dtype, block_size, use_v1,

192 - use_mla) -> str:

191 + kv_cache_dtype, block_size, use_v1, use_mla,

192 + has_sink) -> str:

193 if use_mla:

194 from vllm.attention.backends.rocm_aiter_mla import (

195 is_aiter_mla_enabled)

vllm/platforms/tpu.py

+2 -2

@@ -46,8 +46,8 @@ class TpuPlatform(Platform):

46 @classmethod

47 def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,

48 dtype: torch.dtype, kv_cache_dtype: Optional[str],

49 - block_size: int, use_v1: bool,

50 - use_mla: bool) -> str:

49 + block_size: int, use_v1: bool, use_mla: bool,

50 + has_sink) -> str:

51 if (selected_backend != _Backend.PALLAS

52 and selected_backend != _Backend.PALLAS_VLLM_V1):

53 logger.info("Cannot use %s backend on TPU.", selected_backend)

vllm/platforms/xpu.py

+2 -2

@@ -35,8 +35,8 @@ class XPUPlatform(Platform):

35 @classmethod

36 def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,

37 dtype: torch.dtype, kv_cache_dtype: Optional[str],

38 - block_size: int, use_v1: bool,

39 - use_mla: bool) -> str:

38 + block_size: int, use_v1: bool, use_mla: bool,

39 + has_sink: bool) -> str:

40 if selected_backend is not None and selected_backend != _Backend.IPEX:

41 logger.info("Cannot use %s backend on XPU.", selected_backend)

42 use_v1 = envs.VLLM_USE_V1