@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 26831949b4

[ROCm] Fix pooling startup workspace lock (#47912)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Jul 8, 2026 at 17:59 UTC 26831949b48a0d81fba379dcaf7e378206fd9087

4 files changed +40 -5

vllm/v1/attention/backends/rocm_attn.py

@@ -288,6 +288,8 @@ class RocmAttentionImpl(AttentionImpl):

288 self.alibi_slopes = alibi_slopes

289 if sliding_window is None:

290 self.sliding_window = (-1, -1)

291 + elif attn_type in (AttentionType.ENCODER, AttentionType.ENCODER_ONLY):

292 + self.sliding_window = (sliding_window - 1, sliding_window - 1)

293 else:

294 self.sliding_window = (sliding_window - 1, 0)

295 self.kv_cache_dtype = kv_cache_dtype

@@ -354,6 +356,7 @@ class RocmAttentionImpl(AttentionImpl):

356 softmax_scale=self.scale,

357 sliding_window_q=self.sliding_window[0],

358 sliding_window_k=self.sliding_window[1],

359 + sinks=self.sinks,

360 )

361 return output

362

vllm/v1/attention/backends/triton_attn.py

@@ -760,6 +760,7 @@ class TritonAttentionImpl(AttentionImpl):

760 softmax_scale=self.scale,

761 sliding_window_q=self.sliding_window[0],

762 sliding_window_k=self.sliding_window[1],

763 + sinks=self.sinks,

764 )

765 return output

766

vllm/v1/attention/ops/triton_prefill_attention.py

+14 -2

@@ -38,6 +38,7 @@ def _fwd_kernel(

38 Q,

39 K,

40 V,

41 + Sinks,

42 sm_scale,

43 B_Start_Loc,

44 B_Seqlen,

@@ -57,6 +58,7 @@ def _fwd_kernel(

58 IS_CAUSAL: tl.constexpr,

59 SLIDING_WINDOW_Q: tl.constexpr,

60 SLIDING_WINDOW_K: tl.constexpr,

61 + USE_SINKS: tl.constexpr,

62 Lk: tl.constexpr,

63 ):

64 cur_batch = tl.program_id(0)

@@ -94,8 +96,13 @@ def _fwd_kernel(

96 v_ptrs = V + off_v

98 # initialize pointer to m and l

97 - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")

98 - l_i = tl.zeros([BLOCK_M], dtype=tl.float32)

99 + if USE_SINKS:

100 + sink = tl.load(Sinks + cur_head) * 1.4426950408889634

101 + m_i = tl.full([BLOCK_M], sink, dtype=tl.float32)

102 + l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

103 + else:

104 + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")

105 + l_i = tl.zeros([BLOCK_M], dtype=tl.float32)

106 acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)

107

108 block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)

@@ -200,6 +207,7 @@ def context_attention_fwd(

207 softmax_scale: float | None = None,

208 sliding_window_q: int | None = None,

209 sliding_window_k: int | None = None,

210 + sinks: torch.Tensor | None = None,

211 ):

212 """

213 q, k, v: [b * s, head, head_dim]

@@ -217,6 +225,8 @@ def context_attention_fwd(

225

226 batch, head = b_seq_len.shape[0], q.shape[1]

227 kv_group_num = q.shape[1] // k.shape[1]

228 + if sinks is not None:

229 + assert sinks.shape[0] == head, "Sinks must be num_query_heads size"

230

231 grid = (batch, head, triton.cdiv(max_input_len, BLOCK))

232 num_warps = 4 if Lk <= 64 else 8

@@ -228,6 +238,7 @@ def context_attention_fwd(

238 q,

239 k,

240 v,

241 + sinks if sinks is not None else q,

242 sm_scale,

243 b_start_loc,

244 b_seq_len,

@@ -247,6 +258,7 @@ def context_attention_fwd(

258 IS_CAUSAL=is_causal,

259 SLIDING_WINDOW_Q=sliding_window_q,

260 SLIDING_WINDOW_K=sliding_window_k,

261 + USE_SINKS=sinks is not None,

262 num_warps=num_warps,

263 num_stages=1,

264 Lk=Lk,

vllm/v1/worker/gpu_worker.py

+22 -3

@@ -729,6 +729,11 @@ class Worker(WorkerBase):

729 @instrument(span_name="Warmup (GPU)")

730 def compile_or_warm_up_model(self) -> CompilationTimes:

731 warmup_sizes: list[int] = []

732 + cg_capture_sizes: list[int] = []

733 +

734 + if self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:

735 + cg_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes

736 + cg_capture_sizes = [] if cg_sizes is None else cg_sizes

737

738 if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE:

739 # warm up sizes that are not in cudagraph capture sizes,

@@ -736,11 +741,8 @@ class Worker(WorkerBase):

741 # e.g. for the max-num-batched token size in chunked prefill.

742 compile_sizes = self.vllm_config.compilation_config.compile_sizes

743 warmup_sizes = compile_sizes.copy() if compile_sizes is not None else [] # type: ignore[assignment]

739 - cg_capture_sizes: list[int] = []

744

745 if self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:

742 - cg_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes

743 - cg_capture_sizes = [] if cg_sizes is None else cg_sizes

746 warmup_sizes = [x for x in warmup_sizes if x not in cg_capture_sizes]

747

748 compile_ranges = self.vllm_config.compilation_config.get_compile_ranges()

@@ -753,6 +755,23 @@ class Worker(WorkerBase):

755 if not any(x in compile_range for x in all_sizes):

756 warmup_sizes.append(compile_range.end)

757

758 + # TODO(LucasWilkinson, akaratza): Remove when MRV1 is deprecated

759 + if (

760 + current_platform.is_rocm()

761 + and not self.use_v2_model_runner

762 + and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

763 + and get_pp_group().is_last_rank

764 + ):

765 + max_num_reqs = min(

766 + self.scheduler_config.max_num_seqs,

767 + self.scheduler_config.max_num_batched_tokens,

768 + )

769 + if (

770 + max_num_reqs not in cg_capture_sizes

771 + and max_num_reqs not in warmup_sizes

772 + ):

773 + warmup_sizes.append(max_num_reqs)

774 +

775 # We skip EPLB here since we don't want to record dummy metrics

776 for size in sorted(warmup_sizes, reverse=True):

777 logger.info("Compile and warming up model for size %d", size)