@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9f0247cfa4

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 13, 10:45 UTC

`VLLM_USE_TRITON_FLASH_ATTN` V0 variable deprecation (#27611)

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

Andreas Karatzas committed Nov 11, 2025 at 20:34 UTC 9f0247cfa40a52356aa7860c163c062eb086d266

15 files changed +11 -1587

.buildkite/scripts/hardware_ci/run-amd-test.sh

+2 -6

@@ -78,17 +78,13 @@ HF_MOUNT="/root/.cache/huggingface"

78 commands=$@

79 echo "Commands:$commands"

81 -if [[ $commands == *"pytest -v -s basic_correctness/test_basic_correctness.py"* ]]; then

82 - commands=${commands//"pytest -v -s basic_correctness/test_basic_correctness.py"/"VLLM_USE_TRITON_FLASH_ATTN=0 pytest -v -s basic_correctness/test_basic_correctness.py"}

83 -fi

81 +commands=${commands//"pytest -v -s basic_correctness/test_basic_correctness.py"/"pytest -v -s basic_correctness/test_basic_correctness.py"}

83 if [[ $commands == *"pytest -v -s models/test_registry.py"* ]]; then

84 commands=${commands//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}

85 fi

89 -if [[ $commands == *"pytest -v -s compile/test_basic_correctness.py"* ]]; then

90 - commands=${commands//"pytest -v -s compile/test_basic_correctness.py"/"VLLM_USE_TRITON_FLASH_ATTN=0 pytest -v -s compile/test_basic_correctness.py"}

91 -fi

87 +commands=${commands//"pytest -v -s compile/test_basic_correctness.py"/"pytest -v -s compile/test_basic_correctness.py"}

89 if [[ $commands == *"pytest -v -s lora"* ]]; then

90 commands=${commands//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}

tests/kernels/test_triton_flash_attention.py deleted

-516

@@ -1,516 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Tests for the triton_flash_attention kernel

4 -

5 -Run `pytest tests/kernels/test_triton_flash_attention.py`.

6 -"""

7 -

8 -import pytest

9 -import torch

10 -

11 -from vllm.attention.ops.triton_flash_attention import (

12 - SUPPORTED_LAYOUTS,

13 - MetaData,

14 - compute_alibi_tensor,

15 - scale_fp8,

16 - triton_attention_rocm,

17 -)

18 -from vllm.platforms import current_platform

19 -

20 -

21 -class ReferenceAttention:

22 - def __init__(

23 - self, Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, use_alibi, dtype, input_metadata

24 - ):

25 - self.Z = Z

26 - self.HQ = HQ

27 - self.HK = HK

28 - self.N_CTX_Q = N_CTX_Q

29 - self.N_CTX_K = N_CTX_K

30 - self.D_HEAD = D_HEAD

31 - self.use_alibi = use_alibi

32 - self.dtype = dtype

33 - self.input_metadata = input_metadata

34 -

35 - def fwd(self, q, k, v):

36 - scores = (

37 - torch.einsum("bhqd,bhkd->bhqk", q, k).float() * self.input_metadata.sm_scale

38 - )

39 - if self.input_metadata.causal:

40 - mask = torch.tril(

41 - torch.ones(self.N_CTX_Q, self.N_CTX_K, device="cuda"),

42 - diagonal=self.N_CTX_K - self.N_CTX_Q,

43 - )

44 - scores[:, :, mask == 0] = float("-inf")

45 -

46 - if self.input_metadata.bias is not None:

47 - scores += self.input_metadata.bias

48 -

49 - if self.use_alibi:

50 - scores += compute_alibi_tensor(

51 - self.input_metadata.alibi_slopes, self.N_CTX_Q, self.N_CTX_K

52 - )

53 -

54 - p = torch.softmax(scores, dim=-1)

55 - if self.input_metadata.causal:

56 - # If N_CTX_Q > N_CTX_K, there's at least one row of all -infs going

57 - # into softmax. This creates a row of NaNs as -inf - -inf == NaN.

58 - # So we fix this by converting the NaNs to 0s, which is what they

59 - # should be out of the softmax.

60 - nan_mask = torch.isnan(p)

61 - p[nan_mask == 1] = 0

62 - ref_out = torch.einsum("bhqk,bhkd->bhqd", p.to(self.dtype), v)

63 - # compare

64 - if self.input_metadata.layout == "bshd":

65 - ref_out = ref_out.transpose(1, 2).clone()

66 - return ref_out

67 -

68 - def fwd_fp8(self, q_quantized, k_quantized, v_quantized):

69 - q = (q_quantized.to(torch.float16) * self.input_metadata.q_descale).to(

70 - self.dtype

71 - )

72 - k = (k_quantized.to(torch.float16) * self.input_metadata.k_descale).to(

73 - self.dtype

74 - )

75 - v = (v_quantized.to(torch.float16) * self.input_metadata.v_descale).to(

76 - self.dtype

77 - )

78 - result = self.fwd(q, k, v)

79 - if self.input_metadata.o_scale is not None:

80 - result, _ = scale_fp8(result, self.input_metadata.o_scale)

81 - return result

82 -

83 - def fwd_fp8_kv(self, q, k_quantized, v_quantized):

84 - k_descale, v_descale = (

85 - self.input_metadata.k_descale,

86 - self.input_metadata.v_descale,

87 - )

88 - k_dequantized = (

89 - k_quantized.to(torch.float32) * k_descale.to(torch.float32)

90 - ).to(self.dtype)

91 - v_dequantized = (

92 - v_quantized.to(torch.float32) * v_descale.to(torch.float32)

93 - ).to(self.dtype)

94 - return self.fwd(q, k_dequantized, v_dequantized)

95 -

96 - def varlen_fwd(self, q, k, v, is_mqa=False):

97 - ref_out = torch.empty_like(q)

98 - if is_mqa:

99 - # Make KV look like HQ/HK "groups" of HK. Later, we will reshape so

100 - # the size aligns with Q.

101 - k_ref = k.view(k.shape[0], k.shape[1], 1, k.shape[2]).expand(

102 - -1, -1, self.HQ // self.HK, -1

103 - )

104 - v_ref = v.view(v.shape[0], v.shape[1], 1, v.shape[2]).expand(

105 - -1, -1, self.HQ // self.HK, -1

106 - )

107 - else:

108 - k_ref = k

109 - v_ref = v

110 -

111 - for i in range(0, self.input_metadata.num_contexts):

112 - start_q, start_k = (

113 - self.input_metadata.cu_seqlens_q[i],

114 - self.input_metadata.cu_seqlens_k[i],

115 - )

116 - end_q, end_k = (

117 - self.input_metadata.cu_seqlens_q[i + 1],

118 - self.input_metadata.cu_seqlens_k[i + 1],

119 - )

120 - k_curr = k_ref[start_k:end_k]

121 - v_curr = v_ref[start_k:end_k]

122 - if is_mqa:

123 - k_curr = k_curr.reshape(k_curr.shape[0], -1, k_curr.shape[3])

124 - v_curr = v_curr.reshape(v_curr.shape[0], -1, v_curr.shape[3])

125 - scores = torch.einsum("qhd,khd->qhk", q[start_q:end_q], k_curr).float()

126 - p = torch.softmax(scores * self.input_metadata.sm_scale, dim=-1).half()

127 - ref_out[start_q:end_q] = torch.einsum("qhk,khd->qhd", p, v_curr)

128 - return ref_out

129 -

130 -

131 -def quantize_input(q, k, v, fp8_kv=False, use_o_scale=False):

132 - q_descale = None

133 - if not fp8_kv:

134 - q, q_descale = scale_fp8(q)

135 - k, k_descale = scale_fp8(k)

136 - v, v_descale = scale_fp8(v)

137 -

138 - # In real world use case, the p scale would be a parameter trained by the

139 - # model.

140 - p_scale = None

141 -

142 - o_scale = torch.rand(1, device="cuda", requires_grad=False) if use_o_scale else None

143 -

144 - return q, k, v, q_descale, k_descale, v_descale, p_scale, o_scale

145 -

146 -

147 -def input_helper(

148 - Z,

149 - HQ,

150 - HK,

151 - N_CTX_Q,

152 - N_CTX_K,

153 - D_HEAD,

154 - dtype,

155 - layout=None,

156 - use_alibi=None,

157 - causal=None,

158 - is_fp8=False,

159 - fp8_kv=False,

160 - use_o_scale=False,

161 - use_bias=False,

162 -):

163 - assert layout in SUPPORTED_LAYOUTS, "Got unsupported layout."

164 -

165 - current_platform.seed_everything(0)

166 -

167 - # Initialize q, k, v

168 - if layout == "bhsd":

169 - q_tensor_shape = (Z, HQ, N_CTX_Q, D_HEAD)

170 - k_tensor_shape = (Z, HK, N_CTX_K, D_HEAD)

171 - elif layout == "bshd":

172 - q_tensor_shape = (Z, N_CTX_Q, HQ, D_HEAD)

173 - k_tensor_shape = (Z, N_CTX_K, HK, D_HEAD)

174 -

175 - if use_alibi:

176 - # for n heads the set of slopes is the geometric sequence that starts

177 - # 2^(-8/n)

178 - alibi_slopes = torch.tensor(

179 - [2 ** (-8 / HQ * i) for i in range(1, HQ + 1)],

180 - dtype=torch.float32,

181 - device="cuda",

182 - ).repeat(Z, 1)

183 - else:

184 - alibi_slopes = None

185 -

186 - if use_bias:

187 - bias = torch.randn(

188 - (1, HQ, N_CTX_Q, N_CTX_K), dtype=dtype, device="cuda", requires_grad=False

189 - )

190 - else:

191 - bias = None

192 -

193 - q = torch.randn(q_tensor_shape, dtype=dtype, device="cuda", requires_grad=False)

194 - k = torch.randn(k_tensor_shape, dtype=dtype, device="cuda", requires_grad=False)

195 - v = torch.randn(k_tensor_shape, dtype=dtype, device="cuda", requires_grad=False)

196 -

197 - if is_fp8:

198 - (q, k, v, q_descale, k_descale, v_descale, p_scale, o_scale) = quantize_input(

199 - q, k, v, use_o_scale=use_o_scale, fp8_kv=fp8_kv

200 - )

201 - else:

202 - q_descale = k_descale = v_descale = p_scale = o_scale = None

203 -

204 - input_metadata = MetaData(

205 - sm_scale=D_HEAD**-0.5,

206 - max_seqlens_q=N_CTX_Q,

207 - max_seqlens_k=N_CTX_K,

208 - layout=layout,

209 - alibi_slopes=alibi_slopes,

210 - alibi_batch=Z,

211 - alibi_nheads=HQ,

212 - q_descale=q_descale,

213 - k_descale=k_descale,

214 - v_descale=v_descale,

215 - p_scale=p_scale,

216 - o_scale=o_scale,

217 - bias=bias,

218 - seqlen_q=N_CTX_Q,

219 - seqlen_k=N_CTX_K,

220 - )

221 - return q, k, v, input_metadata

222 -

223 -

224 -def varlen_input_helper(

225 - Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, equal_seqlens=False

226 -):

227 - current_platform.seed_everything(0)

228 -

229 - # Random sequence lengths. Using N_CTX as kind of max of sum of individual

230 - # seqs

231 - if not equal_seqlens:

232 - max_seqlens_q = N_CTX_Q // Z

233 - max_seqlens_k = N_CTX_K // Z

234 - seqlens_q = torch.randint(1, max_seqlens_q + 1, (Z,), dtype=torch.int32)

235 - seqlens_k = torch.randint(1, max_seqlens_k + 1, (Z,), dtype=torch.int32)

236 - else:

237 - seqlens_q = torch.full((Z,), N_CTX_Q // Z)

238 - seqlens_k = torch.full((Z,), N_CTX_K // Z)

239 -

240 - # Calculate cumulative sequence lengths

241 - cu_seqlens_q = torch.cat(

242 - [

243 - torch.tensor([0], dtype=torch.int32),

244 - seqlens_q.cumsum(dim=0, dtype=torch.int32),

245 - ]

246 - )

247 - cu_seqlens_k = torch.cat(

248 - [

249 - torch.tensor([0], dtype=torch.int32),

250 - seqlens_k.cumsum(dim=0, dtype=torch.int32),

251 - ]

252 - )

253 - cu_seqlens_q = cu_seqlens_q.to(device="cuda")

254 - cu_seqlens_k = cu_seqlens_k.to(device="cuda")

255 -

256 - # Initialize q, k, v with variable lengths

257 - total_q = cu_seqlens_q[-1].item()

258 - total_k = cu_seqlens_k[-1].item()

259 - q = (

260 - torch.randn((total_q, HQ, D_HEAD), dtype=dtype, device="cuda")

261 - .normal_(mean=0.0, std=0.5)

262 - .requires_grad_()

263 - )

264 - k = (

265 - torch.randn((total_k, HK, D_HEAD), dtype=dtype, device="cuda")

266 - .normal_(mean=0.0, std=0.5)

267 - .requires_grad_()

268 - )

269 - v = (

270 - torch.randn((total_k, HK, D_HEAD), dtype=dtype, device="cuda")

271 - .normal_(mean=0.0, std=0.5)

272 - .requires_grad_()

273 - )

274 - sm_scale = D_HEAD**-0.5

275 - input_metadata = MetaData(sm_scale=sm_scale)

276 - input_metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k)

277 - return q, k, v, input_metadata

278 -

279 -

280 -@pytest.mark.parametrize(

281 - "Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD",

282 - [

283 - (1, 48, 12, 1, 1, 64),

284 - (4, 4, 4, 128, 128, 65),

285 - (16, 48, 48, 1, 1, 128),

286 - (64, 48, 24, 3, 3, 128),

287 - (4, 4, 4, 113, 123, 1),

288 - ],

289 -)

290 -@pytest.mark.parametrize("causal", [True, False])

291 -@pytest.mark.parametrize("use_alibi", [True, False])

292 -@pytest.mark.parametrize("layout", ["bshd"])

293 -def test_op_fwd(

294 - Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, causal, use_alibi, layout, dtype=torch.float16

295 -):

296 - current_platform.seed_everything(0)

297 - q, k, v, input_metadata = input_helper(

298 - Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, layout, use_alibi, causal

299 - )

300 -

301 - o = torch.empty_like(q)

302 -

303 - # triton implementation

304 - tri_out, _ = triton_attention_rocm(q, k, v, o, input_metadata)

305 -

306 - # Transpose here if layout is bshd so we have same reference code for all

307 - # layouts

308 - if layout == "bshd":

309 - q = q.transpose(1, 2).clone()

310 - k = k.transpose(1, 2).clone()

311 - v = v.transpose(1, 2).clone()

312 - # Replicate K and V if using MQA/GQA

313 - if HQ != HK:

314 - k = (

315 - k.view(k.shape[0], k.shape[1], -1, k.shape[2], k.shape[3])

316 - .expand(-1, -1, HQ // HK, -1, -1)

317 - .reshape(k.shape[0], -1, k.shape[2], k.shape[3])

318 - )

319 - v = (

320 - v.view(v.shape[0], v.shape[1], -1, v.shape[2], v.shape[3])

321 - .expand(-1, -1, HQ // HK, -1, -1)

322 - .reshape(v.shape[0], -1, v.shape[2], v.shape[3])

323 - )

324 -

325 - ref_impl = ReferenceAttention(

326 - Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, use_alibi, dtype, input_metadata

327 - )

328 - ref_out = ref_impl.fwd(q, k, v)

329 -

330 - torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)

331 -

332 -

333 -@pytest.mark.parametrize(

334 - "Z, H, N_CTX_Q, N_CTX_K, D_HEAD",

335 - [

336 - (4, 48, 1, 1, 64),

337 - (4, 48, 1, 1, 128),

338 - (4, 48, 3, 3, 128),

339 - (4, 4, 128, 128, 65),

340 - ],

341 -)

342 -@pytest.mark.parametrize("causal", [True, False])

343 -@pytest.mark.parametrize("layout", ["bhsd"])

344 -@pytest.mark.parametrize("use_o_scale", [True, False])

345 -@pytest.mark.skipif(

346 - torch.cuda.get_device_capability() < (9, 0),

347 - reason="Triton FP8 requires CUDA 9.0 or higher",

348 -)

349 -def test_op_fwd_fp8(

350 - Z, H, N_CTX_Q, N_CTX_K, D_HEAD, causal, layout, use_o_scale, dtype=torch.float32

351 -):

352 - current_platform.seed_everything(0)

353 -

354 - # Disable grad to save memory it won't run into OOM on CI machine.

355 - # q, k, v, input_metadata = input_helper(Z, H, H, N_CTX_Q, N_CTX_K, D_HEAD,

356 - # dtype, layout)

357 -

358 - q_quantized, k_quantized, v_quantized, input_metadata = input_helper(

359 - Z,

360 - H,

361 - H,

362 - N_CTX_Q,

363 - N_CTX_K,

364 - D_HEAD,

365 - dtype,

366 - causal=causal,

367 - layout=layout,

368 - is_fp8=True,

369 - use_o_scale=use_o_scale,

370 - )

371 -

372 - o = torch.empty_like(q_quantized) if use_o_scale else None

373 -

374 - tri_out, _ = triton_attention_rocm(

375 - q_quantized, k_quantized, v_quantized, o, input_metadata

376 - )

377 -

378 - ref_impl = ReferenceAttention(

379 - Z, H, H, N_CTX_Q, N_CTX_K, D_HEAD, False, dtype, input_metadata

380 - )

381 - ref_out = ref_impl.fwd_fp8(q_quantized, k_quantized, v_quantized)

382 -

383 - # compare

384 - torch.testing.assert_close(

385 - ref_out.to(torch.float32), tri_out.to(torch.float32), atol=7e-2, rtol=2e-1

386 - )

387 -

388 -

389 -@pytest.mark.parametrize(

390 - "Z, H, N_CTX_Q, N_CTX_K, D_HEAD",

391 - [

392 - (4, 48, 1, 1, 64),

393 - (4, 48, 1, 1, 128),

394 - (4, 48, 3, 3, 128),

395 - (4, 4, 128, 128, 65),

396 - (4, 4, 113, 123, 1),

397 - ],

398 -)

399 -@pytest.mark.parametrize("causal", [True, False])

400 -@pytest.mark.parametrize("layout", ["bhsd"])

401 -def test_op_fwd_fp8_kv(

402 - Z, H, N_CTX_Q, N_CTX_K, D_HEAD, causal, layout, dtype=torch.float32

403 -):

404 - current_platform.seed_everything(0)

405 -

406 - q, k_quantized, v_quantized, input_metadata = input_helper(

407 - Z,

408 - H,

409 - H,

410 - N_CTX_Q,

411 - N_CTX_K,

412 - D_HEAD,

413 - dtype,

414 - causal=causal,

415 - layout=layout,

416 - is_fp8=True,

417 - fp8_kv=True,

418 - )

419 -

420 - o = torch.empty_like(q)

421 -

422 - tri_out, _ = triton_attention_rocm(q, k_quantized, v_quantized, o, input_metadata)

423 -

424 - ref_impl = ReferenceAttention(

425 - Z, H, H, N_CTX_Q, N_CTX_K, D_HEAD, False, dtype, input_metadata

426 - )

427 - ref_out = ref_impl.fwd_fp8_kv(q, k_quantized, v_quantized)

428 -

429 - torch.testing.assert_close(ref_out, tri_out, atol=3e-2, rtol=8e-1)

430 -

431 -

432 -@pytest.mark.parametrize(

433 - "Z, H, N_CTX_Q, N_CTX_K, D_HEAD",

434 - [

435 - (4, 48, 1, 1, 64),

436 - (4, 48, 1, 1, 128),

437 - (4, 48, 3, 3, 128),

438 - (4, 4, 128, 128, 65),

439 - ],

440 -)

441 -@pytest.mark.parametrize("causal", [True, False])

442 -@pytest.mark.parametrize("use_bias", [True])

443 -@pytest.mark.parametrize("dtype", [torch.bfloat16])

444 -def test_op_fwd_bias(Z, H, N_CTX_Q, N_CTX_K, D_HEAD, causal, use_bias, dtype):

445 - current_platform.seed_everything(0)

446 - q, k, v, input_metadata = input_helper(

447 - Z,

448 - H,

449 - H,

450 - N_CTX_Q,

451 - N_CTX_K,

452 - D_HEAD,

453 - dtype,

454 - layout="bhsd",

455 - causal=causal,

456 - use_bias=use_bias,

457 - )

458 - o = torch.empty_like(q)

459 -

460 - # triton implementation

461 - tri_out, _ = triton_attention_rocm(q, k, v, o, input_metadata)

462 -

463 - ref_impl = ReferenceAttention(

464 - Z, H, H, N_CTX_Q, N_CTX_K, D_HEAD, False, dtype, input_metadata

465 - )

466 - ref_out = ref_impl.fwd(q, k, v)

467 -

468 - # compare

469 - torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)

470 -

471 -

472 -# NOTE: Uses thd layout, so also tests thd.

473 -@pytest.mark.parametrize(

474 - "Z, H, N_CTX, D_HEAD",

475 - [(1, 48, 256, 64), (4, 48, 512, 64), (16, 48, 512, 64), (64, 48, 128, 128)],

476 -)

477 -@pytest.mark.parametrize("causal", [True, False])

478 -def test_op_varlen_fwd(Z, H, N_CTX, D_HEAD, causal, dtype=torch.float16):

479 - q, k, v, input_metadata = varlen_input_helper(Z, H, H, N_CTX, N_CTX, D_HEAD, dtype)

480 -

481 - tri_out = torch.empty_like(q)

482 - triton_attention_rocm(q, k, v, tri_out, input_metadata)

483 -

484 - ref_impl = ReferenceAttention(

485 - Z, H, H, N_CTX, N_CTX, D_HEAD, False, dtype, input_metadata

486 - )

487 - ref_out = ref_impl.varlen_fwd(q, k, v, is_mqa=False)

488 -

489 - torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)

490 -

491 -

492 -# NOTE: Uses thd layout, so also tests thd.

493 -@pytest.mark.parametrize(

494 - "Z, HQ, HK, N_CTX, D_HEAD",

495 - [

496 - (2, 48, 24, 128, 64),

497 - (4, 48, 12, 256, 64),

498 - (4, 48, 4, 512, 64),

499 - (4, 64, 16, 128, 128),

500 - ],

501 -)

502 -@pytest.mark.parametrize("causal", [False])

503 -def test_op_varlen_mqa_fwd(Z, HQ, HK, N_CTX, D_HEAD, causal, dtype=torch.float16):

504 - q, k, v, input_metadata = varlen_input_helper(

505 - Z, HQ, HK, N_CTX, N_CTX, D_HEAD, dtype

506 - )

507 -

508 - tri_out = torch.empty_like(q)

509 - triton_attention_rocm(q, k, v, tri_out, input_metadata)

510 -

511 - ref_impl = ReferenceAttention(

512 - Z, HQ, HK, N_CTX, N_CTX, D_HEAD, False, dtype, input_metadata

513 - )

514 - ref_out = ref_impl.varlen_fwd(q, k, v, is_mqa=True)

515 -

516 - torch.testing.assert_close(ref_out, tri_out, atol=2e-2, rtol=2e-2)

tests/models/language/pooling/test_classification.py

-6

@@ -27,13 +27,7 @@ def test_models(

27 example_prompts,

28 model: str,

29 dtype: str,

30 - monkeypatch,

30 ) -> None:

32 - if current_platform.is_rocm():

33 - # ROCm Triton FA does not currently support sliding window attention

34 - # switch to use ROCm CK FA backend

35 - monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")

36 -

31 with vllm_runner(model, max_model_len=512, dtype=dtype) as vllm_model:

32 vllm_outputs = vllm_model.classify(example_prompts)

tests/models/language/pooling/test_embedding.py

-7

@@ -4,7 +4,6 @@

4 import pytest

6 from vllm.config import PoolerConfig

7 -from vllm.platforms import current_platform

8 from ...utils import check_embeddings_close

@@ -51,13 +50,7 @@ def test_models(

50 vllm_runner,

51 example_prompts,

52 model,

54 - monkeypatch,

53 ) -> None:

56 - if model == "BAAI/bge-multilingual-gemma2" and current_platform.is_rocm():

57 - # ROCm Triton FA does not currently support sliding window attention

58 - # switch to use ROCm CK FA backend

59 - monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")

60 -

54 vllm_extra_kwargs = {}

55 if model == "ssmits/Qwen2-7B-Instruct-embed-base":

56 vllm_extra_kwargs["pooler_config"] = PoolerConfig(

tests/models/language/pooling/test_mm_classifier_conversion.py

-13

@@ -2,18 +2,11 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from vllm.config.pooler import PoolerConfig

5 -from vllm.platforms import current_platform

7 def test_idefics_multimodal(

8 vllm_runner,

10 - monkeypatch,

9 ) -> None:

12 - if current_platform.is_rocm():

13 - # ROCm Triton FA does not currently support sliding window attention

14 - # switch to use ROCm CK FA backend

15 - monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")

16 -

10 prompts = [

11 "Hello, my name is",

12 "The president of the United States is",

@@ -59,13 +52,7 @@ def update_config(config):

53 def test_gemma_multimodal(

54 vllm_runner,

62 - monkeypatch,

55 ) -> None:

64 - if current_platform.is_rocm():

65 - # ROCm Triton FA does not currently support sliding window attention

66 - # switch to use ROCm CK FA backend

67 - monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")

68 -

56 messages = [

57 {

58 "role": "system",

tests/models/language/pooling/test_reward.py

-6

@@ -76,7 +76,6 @@ def test_prm_models(

76 math_step_prompts,

77 model: str,

78 dtype: str,

79 - monkeypatch,

79 ) -> None:

80 check_transformers_version(

81 "Qwen/Qwen2.5-Math-PRM-7B", max_transformers_version="4.53.2"

@@ -85,11 +84,6 @@ def test_prm_models(

84 if current_platform.is_cpu():

85 pytest.skip("CPU only supports V1")

88 - if current_platform.is_rocm():

89 - # ROCm Triton FA does not currently support sliding window attention

90 - # switch to use ROCm CK FA backend

91 - monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", "False")

92 -

87 with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:

88 vllm_outputs = vllm_model.reward(math_step_prompts)

tests/models/multimodal/generation/test_common.py

-8

@@ -5,7 +5,6 @@ image, embedding, and video support for different VLMs in vLLM.

5 """

7 import math

8 -import os

8 from collections import defaultdict

9 from pathlib import PosixPath

@@ -38,13 +37,6 @@ from .vlm_utils.types import (

37 VLMTestType,

38 )

41 -# This hack is needed for phi3v & paligemma models

42 -# ROCm Triton FA can run into shared memory issues with these models,

43 -# use other backends in the meantime

44 -# FIXME (mattwong, gshtrasb, hongxiayan)

45 -if current_platform.is_rocm():

46 - os.environ["VLLM_USE_TRITON_FLASH_ATTN"] = "0"

47 -

40 COMMON_BROADCAST_SETTINGS = {

41 "test_type": VLMTestType.IMAGE,

42 "dtype": "half",

tests/models/multimodal/generation/test_phi4_multimodal.py

-7

@@ -11,7 +11,6 @@ from huggingface_hub import snapshot_download

11 from vllm.assets.image import ImageAsset

12 from vllm.lora.request import LoRARequest

13 from vllm.multimodal.image import rescale_image_size

14 -from vllm.platforms import current_platform

15 from ....conftest import (

16 IMAGE_ASSETS,

@@ -46,12 +45,6 @@ models = [model_path]

46 target_dtype = "half"

49 -# ROCm Triton FA can run into shared memory issues with these models,

50 -# use other backends in the meantime

51 -# FIXME (mattwong, gshtrasb, hongxiayan)

52 -if current_platform.is_rocm():

53 - os.environ["VLLM_USE_TRITON_FLASH_ATTN"] = "0"

54 -

49 def run_test(

50 hf_runner: type[HfRunner],

tests/models/multimodal/generation/test_phi4mm.py

-7

@@ -14,7 +14,6 @@ from vllm.assets.image import ImageAsset

14 from vllm.logprobs import SampleLogprobs

15 from vllm.lora.request import LoRARequest

16 from vllm.multimodal.image import convert_image_mode, rescale_image_size

17 -from vllm.platforms import current_platform

18 from ....conftest import (

19 IMAGE_ASSETS,

@@ -68,12 +67,6 @@ def vllm_to_hf_output(

68 target_dtype = "half"

71 -# ROCm Triton FA can run into shared memory issues with these models,

72 -# use other backends in the meantime

73 -# FIXME (mattwong, gshtrasb, hongxiayan)

74 -if current_platform.is_rocm():

75 - os.environ["VLLM_USE_TRITON_FLASH_ATTN"] = "0"

76 -

71 def run_test(

72 hf_runner: type[HfRunner],

tests/quantization/test_quark.py

-5

@@ -8,7 +8,6 @@ See also `tests/kernels/moe/test_ocp_mx_moe.py`.

8 """

10 import importlib.metadata

11 -import os

11 from dataclasses import dataclass

12 from importlib.util import find_spec

@@ -246,8 +245,6 @@ def test_mxfp4_gsm8k_correctness(config: AccuracyTestConfig):

245 task = "gsm8k"

246 rtol = 0.03

247

249 - os.environ["VLLM_USE_TRITON_FLASH_ATTN"] = "0"

250 -

248 results = lm_eval.simple_evaluate(

249 model="vllm",

250 model_args=config.get_model_args(tp_size=8, model_max_len=38768),

@@ -263,8 +260,6 @@ def test_mxfp4_gsm8k_correctness(config: AccuracyTestConfig):

260 and measured_value + rtol > EXPECTED_VALUE

261 ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"

262

266 - del os.environ["VLLM_USE_TRITON_FLASH_ATTN"]

267 -

263

264 @pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")

265 @pytest.mark.parametrize("float_dtype", [torch.bfloat16, torch.float16])

vllm/attention/ops/triton_flash_attention.py deleted

-932

@@ -1,932 +0,0 @@

1 -#!/usr/bin/env python

2 -# SPDX-License-Identifier: Apache-2.0

3 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -"""

5 -Fused Attention

6 -===============

7 -

8 -This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao

9 -(https://tridao.me/publications/flash2/flash2.pdf)

10 -Credits: OpenAI kernel team, AMD ML Frameworks Triton team

11 -

12 -Features supported:

13 -

14 -1) Fwd with causal masking

15 -2) Any sequence lengths without padding (currently fwd kernel only)

16 -3) Support for different sequence lengths for q and k

17 -4) Nested tensor API currently does not support dropout or bias.

18 -

19 -Not currently supported:

20 -

21 -1) Non power of two head dims

22 -

23 -"""

24 -

25 -import torch

26 -

27 -from vllm.platforms import current_platform

28 -from vllm.triton_utils import tl, triton

29 -

30 -# Avoid misleading ROCm warning.

31 -if current_platform.is_rocm():

32 - from vllm.platforms.rocm import on_gfx1x

33 -else:

34 - on_gfx1x = lambda *args, **kwargs: False

35 -

36 -torch_dtype: tl.constexpr = torch.float16

37 -

38 -

39 -@triton.jit

40 -def cdiv_fn(x, y):

41 - return (x + y - 1) // y

42 -

43 -

44 -@triton.jit

45 -def max_fn(x, y):

46 - return tl.math.max(x, y)

47 -

48 -

49 -@triton.jit

50 -def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride):

51 - ms = tl.arange(0, m)

52 - ns = tl.arange(0, n)

53 - return philox_offset + ms[:, None] * stride + ns[None, :]

54 -

55 -

56 -@triton.jit

57 -def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride):

58 - rng_offsets = dropout_offsets(

59 - philox_seed, philox_offset, dropout_p, m, n, stride

60 - ).to(tl.uint32)

61 - # TODO: use tl.randint for better performance

62 - return tl.rand(philox_seed, rng_offsets)

63 -

64 -

65 -@triton.jit

66 -def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride):

67 - rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride)

68 - rng_keep = rng_output > dropout_p

69 - return rng_keep

70 -

71 -

72 -@triton.jit

73 -def load_fn(block_ptr, first, second, pad):

74 - if first and second:

75 - tensor = tl.load(block_ptr, boundary_check=(0, 1), padding_option=pad)

76 - elif first:

77 - tensor = tl.load(block_ptr, boundary_check=(0,), padding_option=pad)

78 - elif second:

79 - tensor = tl.load(block_ptr, boundary_check=(1,), padding_option=pad)

80 - else:

81 - tensor = tl.load(block_ptr)

82 - return tensor

83 -

84 -

85 -@triton.jit

86 -def _attn_fwd_inner(

87 - acc,

88 - l_i,

89 - m_i,

90 - q,

91 - K_block_ptr,

92 - V_block_ptr,

93 - start_m,

94 - actual_seqlen_k,

95 - dropout_p,

96 - philox_seed,

97 - batch_philox_offset,

98 - encoded_softmax_block_ptr,

99 - block_min,

100 - block_max,

101 - offs_n_causal,

102 - masked_blocks,

103 - n_extra_tokens,

104 - bias_ptr,

105 - IS_CAUSAL: tl.constexpr,

106 - BLOCK_M: tl.constexpr,

107 - BLOCK_DMODEL: tl.constexpr,

108 - BLOCK_N: tl.constexpr,

109 - OFFS_M: tl.constexpr,

110 - OFFS_N: tl.constexpr,

111 - PRE_LOAD_V: tl.constexpr,

112 - MASK_STEPS: tl.constexpr,

113 - ENABLE_DROPOUT: tl.constexpr,

114 - RETURN_ENCODED_SOFTMAX: tl.constexpr,

115 - PADDED_HEAD: tl.constexpr,

116 - USE_FP8: tl.constexpr,

117 - qk_scale,

118 - p_descale,

119 -):

120 - # loop over k, v, and update accumulator

121 - for start_n in range(block_min, block_max, BLOCK_N):

122 - # For padded blocks, we will overrun the tensor size if

123 - # we load all BLOCK_N. For others, the blocks are all within range.

124 - k = load_fn(

125 - K_block_ptr,

126 - PADDED_HEAD,

127 - MASK_STEPS and (n_extra_tokens != 0),

128 - "zero",

129 - )

130 - if PRE_LOAD_V:

131 - v = load_fn(

132 - V_block_ptr,

133 - MASK_STEPS and (n_extra_tokens != 0),

134 - PADDED_HEAD,

135 - "zero",

136 - )

137 - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)

138 - # We start from end of seqlen_k so only the first iteration would need

139 - # to be checked for padding if it is not a multiple of block_n

140 - # TODO: This can be optimized to only be true for the padded block.

141 - if MASK_STEPS: # noqa: SIM102

142 - # If this is the last block / iteration, we want to

143 - # mask if the sequence length is not a multiple of block size

144 - # a solution is to always do BLOCK_M // BLOCK_N + 1 steps

145 - # if not is_modulo_mn. last step might get wasted but that is okay.

146 - # check if this masking works for that case.

147 - if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0):

148 - boundary_m = tl.full([BLOCK_M], actual_seqlen_k, dtype=tl.int32)

149 - size_n = start_n + OFFS_N[None, :]

150 - mask = size_n < boundary_m[:, None]

151 - qk = tl.where(mask, qk, float("-inf"))

152 - if IS_CAUSAL:

153 - causal_boundary = start_n + offs_n_causal

154 - causal_mask = OFFS_M[:, None] >= causal_boundary[None, :]

155 - qk = tl.where(causal_mask, qk, float("-inf"))

156 - # -- compute qk ----

157 - qk += tl.dot(q, k)

158 - if USE_FP8:

159 - qk *= qk_scale

160 - if bias_ptr is not None:

161 - bias = load_fn(

162 - bias_ptr, False, MASK_STEPS and (n_extra_tokens != 0), "zero"

163 - )

164 - # While bias is added after multiplying qk with sm_scale, our

165 - # optimization to use 2^x instead of e^x results in an additional

166 - # scale factor of log2(e) which we must also multiply the bias with.

167 - qk += bias * 1.44269504089

168 - m_ij = tl.maximum(m_i, tl.max(qk, 1))

169 - qk = qk - m_ij[:, None]

170 - p = tl.math.exp2(qk)

171 -

172 - # CAVEAT: Must update l_ij before applying dropout

173 - l_ij = tl.sum(p, 1)

174 - if ENABLE_DROPOUT:

175 - philox_offset = (

176 - batch_philox_offset

177 - + start_m * BLOCK_M * actual_seqlen_k

178 - + start_n

179 - - BLOCK_N

180 - )

181 - keep = dropout_mask(

182 - philox_seed,

183 - philox_offset,

184 - dropout_p,

185 - BLOCK_M,

186 - BLOCK_N,

187 - actual_seqlen_k,

188 - )

189 - if RETURN_ENCODED_SOFTMAX:

190 - tl.store(

191 - encoded_softmax_block_ptr,

192 - tl.where(keep, p, -p).to(encoded_softmax_block_ptr.type.element_ty),

193 - )

194 - p = tl.where(keep, p, 0.0)

195 - elif RETURN_ENCODED_SOFTMAX:

196 - tl.store(

197 - encoded_softmax_block_ptr,

198 - p.to(encoded_softmax_block_ptr.type.element_ty),

199 - )

200 - # -- update output accumulator --

201 - alpha = tl.math.exp2(m_i - m_ij)

202 - acc = acc * alpha[:, None]

203 - if not PRE_LOAD_V:

204 - v = load_fn(

205 - V_block_ptr,

206 - MASK_STEPS and (n_extra_tokens != 0),

207 - PADDED_HEAD,

208 - "zero",

209 - )

210 - # -- update m_i and l_i

211 - l_i = l_i * alpha + l_ij

212 - # update m_i and l_i

213 - m_i = m_ij

214 -

215 - if USE_FP8:

216 - p *= p_descale

217 -

218 - acc += tl.dot(p.to(V_block_ptr.type.element_ty), v)

219 -

220 - V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))

221 - K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))

222 - if bias_ptr is not None:

223 - bias_ptr = tl.advance(bias_ptr, (0, BLOCK_N))

224 - if RETURN_ENCODED_SOFTMAX:

225 - encoded_softmax_block_ptr = tl.advance(

226 - encoded_softmax_block_ptr, (0, BLOCK_N)

227 - )

228 - return acc, l_i, m_i

229 -

230 -

231 -def get_cdna_autotune_configs():

232 - return [

233 - triton.Config(

234 - {"BLOCK_M": 256, "BLOCK_N": 64, "waves_per_eu": 2, "PRE_LOAD_V": False},

235 - num_stages=1,

236 - num_warps=8,

237 - ),

238 - triton.Config(

239 - {"BLOCK_M": 128, "BLOCK_N": 128, "waves_per_eu": 2, "PRE_LOAD_V": False},

240 - num_stages=1,

241 - num_warps=4,

242 - ),

243 - triton.Config(

244 - {"BLOCK_M": 256, "BLOCK_N": 128, "waves_per_eu": 2, "PRE_LOAD_V": False},

245 - num_stages=1,

246 - num_warps=8,

247 - ),

248 - triton.Config(

249 - {"BLOCK_M": 128, "BLOCK_N": 64, "waves_per_eu": 1, "PRE_LOAD_V": False},

250 - num_stages=1,

251 - num_warps=4,

252 - ),

253 - triton.Config(

254 - {"BLOCK_M": 128, "BLOCK_N": 64, "waves_per_eu": 3, "PRE_LOAD_V": True},

255 - num_stages=1,

256 - num_warps=4,

257 - ),

258 - triton.Config(

259 - {"BLOCK_M": 128, "BLOCK_N": 64, "waves_per_eu": 3, "PRE_LOAD_V": False},

260 - num_stages=1,

261 - num_warps=4,

262 - ),

263 - triton.Config(

264 - {"BLOCK_M": 64, "BLOCK_N": 64, "waves_per_eu": 4, "PRE_LOAD_V": False},

265 - num_stages=1,

266 - num_warps=8,

267 - ),

268 - triton.Config(

269 - {"BLOCK_M": 32, "BLOCK_N": 32, "waves_per_eu": 4, "PRE_LOAD_V": False},

270 - num_stages=1,

271 - num_warps=8,

272 - ),

273 - # TODO: This config fails with head_size not pow2 with data mismatches.

274 - # triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 1,

275 - # 'PRE_LOAD_V': False}, num_stages=1, num_warps=4),

276 - # Fails in AccelerateAMDMatmul (Triton) assert when using FP8:

277 - # triton.Config(

278 - # {

279 - # "BLOCK_M": 16,

280 - # "BLOCK_N": 16,

281 - # "waves_per_eu": 1,

282 - # "PRE_LOAD_V": False,

283 - # },

284 - # num_stages=1,

285 - # num_warps=4,

286 - # ),

287 - ], ["IS_CAUSAL", "dropout_p", "BLOCK_DMODEL", "USE_FP8"]

288 -

289 -

290 -def get_rdna_autotune_configs():

291 - return [

292 - triton.Config(

293 - {"BLOCK_M": 32, "BLOCK_N": 32, "waves_per_eu": 4, "PRE_LOAD_V": False},

294 - num_stages=1,

295 - num_warps=2,

296 - ),

297 - triton.Config(

298 - {"BLOCK_M": 32, "BLOCK_N": 32, "waves_per_eu": 2, "PRE_LOAD_V": False},

299 - num_stages=1,

300 - num_warps=2,

301 - ),

302 - triton.Config(

303 - {"BLOCK_M": 32, "BLOCK_N": 16, "waves_per_eu": 4, "PRE_LOAD_V": False},

304 - num_stages=1,

305 - num_warps=2,

306 - ),

307 - triton.Config(

308 - {"BLOCK_M": 32, "BLOCK_N": 16, "waves_per_eu": 2, "PRE_LOAD_V": False},

309 - num_stages=1,

310 - num_warps=2,

311 - ),

312 - # Fails in AccelerateAMDMatmul (Triton) assert when using FP8:

313 - # triton.Config(

314 - # {

315 - # 'BLOCK_M': 16,

316 - # 'BLOCK_N': 16,

317 - # 'waves_per_eu': 4,

318 - # 'PRE_LOAD_V': False

319 - # },

320 - # num_stages=1,

321 - # num_warps=2),

322 - # triton.Config(

323 - # {

324 - # 'BLOCK_M': 16,

325 - # 'BLOCK_N': 16,

326 - # 'waves_per_eu': 2,

327 - # 'PRE_LOAD_V': False

328 - # },

329 - # num_stages=1,

330 - # num_warps=2),

331 - # # Fall-back config.

332 - # triton.Config(

333 - # {

334 - # 'BLOCK_M': 16,

335 - # 'BLOCK_N': 16,

336 - # 'waves_per_eu': 1,

337 - # 'PRE_LOAD_V': False

338 - # },

339 - # num_stages=1,

340 - # num_warps=2),

341 - ], ["IS_CAUSAL", "dropout_p", "BLOCK_DMODEL", "USE_FP8"]

342 -

343 -

344 -def get_autotune_configs():

345 - if on_gfx1x():

346 - return get_rdna_autotune_configs()

347 - else:

348 - return get_cdna_autotune_configs()

349 -

350 -

351 -autotune_configs, autotune_keys = get_autotune_configs()

352 -

353 -float8_info = torch.finfo(current_platform.fp8_dtype())

354 -

355 -

356 -@triton.autotune(

357 - configs=autotune_configs,

358 - key=autotune_keys,

359 -)

360 -@triton.jit

361 -def attn_fwd(

362 - Q,

363 - K,

364 - V,

365 - bias,

366 - sm_scale,

367 - q_scale,

368 - k_scale,

369 - v_scale,

370 - p_scale,

371 - p_descale,

372 - o_descale,

373 - L,

374 - Out,

375 - stride_qz: tl.int64,

376 - stride_qh: tl.int64,

377 - stride_qm: tl.int64,

378 - stride_qk: tl.int64,

379 - stride_kz: tl.int64,

380 - stride_kh: tl.int64,

381 - stride_kn: tl.int64,

382 - stride_kk: tl.int64,

383 - stride_vz: tl.int64,

384 - stride_vh: tl.int64,

385 - stride_vk: tl.int64,

386 - stride_vn: tl.int64,

387 - stride_oz: tl.int64,

388 - stride_oh: tl.int64,

389 - stride_om: tl.int64,

390 - stride_on: tl.int64,

391 - stride_bz: tl.int64,

392 - stride_bh: tl.int64,

393 - stride_bm: tl.int64,

394 - stride_bn: tl.int64,

395 - cu_seqlens_q,

396 - cu_seqlens_k,

397 - dropout_p,

398 - philox_seed,

399 - philox_offset_base,

400 - encoded_softmax,

401 - HQ: tl.constexpr,

402 - HK: tl.constexpr,

403 - ACTUAL_BLOCK_DMODEL: tl.constexpr,

404 - MAX_SEQLENS_Q: tl.constexpr,

405 - MAX_SEQLENS_K: tl.constexpr,

406 - VARLEN: tl.constexpr,

407 - IS_CAUSAL: tl.constexpr,

408 - BLOCK_M: tl.constexpr,

409 - BLOCK_DMODEL: tl.constexpr,

410 - USE_FP8: tl.constexpr,

411 - USE_FP8_OUT: tl.constexpr,

412 - BLOCK_N: tl.constexpr,

413 - PRE_LOAD_V: tl.constexpr,

414 - BIAS_TYPE: tl.constexpr,

415 - ENABLE_DROPOUT: tl.constexpr,

416 - RETURN_ENCODED_SOFTMAX: tl.constexpr,

417 - FP8_MIN: tl.constexpr = float8_info.min,

418 - FP8_MAX: tl.constexpr = float8_info.max,

419 -):

420 - start_m = tl.program_id(0)

421 - off_h_q = tl.program_id(1)

422 - off_z = tl.program_id(2)

423 - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)

424 - offs_n = tl.arange(0, BLOCK_N)

425 - if VARLEN:

426 - cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)

427 - cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)

428 - seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start

429 - # We have a one-size-fits-all grid in id(0). Some seqlens might be too

430 - # small for all start_m so for those we return early.

431 - if start_m * BLOCK_M > seqlen_q:

432 - return

433 - cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)

434 - cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)

435 - seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start

436 - else:

437 - cu_seqlens_q_start = 0

438 - cu_seqlens_k_start = 0

439 - seqlen_q = MAX_SEQLENS_Q

440 - seqlen_k = MAX_SEQLENS_K

441 -

442 - # Now we compute whether we need to exit early due to causal masking.

443 - # This is because for seqlen_q > seqlen_k, M rows of the attn scores

444 - # are completely masked, resulting in 0s written to the output, and

445 - # inf written to LSE. We don't need to do any GEMMs in this case.

446 - # This block of code determines what N is, and if this WG is operating

447 - # on those M rows.

448 - n_blocks = cdiv_fn(seqlen_k, BLOCK_N)

449 - if IS_CAUSAL:

450 - # If seqlen_q == seqlen_k, the attn scores are a square matrix.

451 - # If seqlen_q != seqlen_k, attn scores are rectangular which means

452 - # the causal mask boundary is bottom right aligned, and ends at either

453 - # the top edge (seqlen_q < seqlen_k) or left edge.

454 - # This captures the decrease in n_blocks if we have a rectangular attn

455 - # matrix

456 - n_blocks_seqlen = cdiv_fn(

457 - (start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N

458 - )

459 - # This is what adjusts the block_max for the current WG, only

460 - # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks

461 - n_blocks = min(n_blocks, n_blocks_seqlen)

462 - # If we have no blocks after adjusting for seqlen deltas, this WG is

463 - # part of the blocks that are all 0. We exit early.

464 - if n_blocks <= 0:

465 - o_offset = (

466 - off_z * stride_oz + cu_seqlens_q_start * stride_om + off_h_q * stride_oh

467 - )

468 - O_block_ptr = tl.make_block_ptr(

469 - base=Out + o_offset,

470 - shape=(seqlen_q, BLOCK_DMODEL),

471 - strides=(stride_om, stride_on),

472 - offsets=(start_m * BLOCK_M, 0),

473 - block_shape=(BLOCK_M, BLOCK_DMODEL),

474 - order=(1, 0),

475 - )

476 - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty)

477 - # We still need to write 0s to the result

478 - # tl.store(O_block_ptr,

479 - # acc.to(Out.type.element_ty), boundary_check=(0,1))

480 - # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q

481 - # + offs_m

482 - # We store inf to LSE, not -inf because in the bwd pass,

483 - # we subtract this

484 - # from qk which makes it -inf, such that exp(qk - inf) = 0

485 - # for these masked blocks.

486 - # l = tl.full([BLOCK_M], value=float("inf"), dtype=tl.float32)

487 - # tl.store(l_ptrs, l)

488 - # TODO: Should dropout and return encoded softmax be handled here?

489 - return

490 -

491 - # If MQA / GQA, set the K and V head offsets appropriately.

492 - GROUP_SIZE: tl.constexpr = HQ // HK

493 - off_h_k = off_h_q // GROUP_SIZE if GROUP_SIZE != 1 else off_h_q

494 -

495 - n_extra_tokens = 0

496 - if seqlen_k < BLOCK_N:

497 - n_extra_tokens = BLOCK_N - seqlen_k

498 - elif seqlen_k % BLOCK_N:

499 - n_extra_tokens = seqlen_k % BLOCK_N

500 - padded_head = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL

501 -

502 - # Compute pointers for all the tensors used in this kernel.

503 - q_offset = off_z * stride_qz + off_h_q * stride_qh + cu_seqlens_q_start * stride_qm

504 - Q_block_ptr = tl.make_block_ptr(

505 - base=Q + q_offset,

506 - shape=(seqlen_q, ACTUAL_BLOCK_DMODEL),

507 - strides=(stride_qm, stride_qk),

508 - offsets=(start_m * BLOCK_M, 0),

509 - block_shape=(BLOCK_M, BLOCK_DMODEL),

510 - order=(1, 0),

511 - )

512 - k_offset = off_z * stride_kz + off_h_k * stride_kh + cu_seqlens_k_start * stride_kn

513 - K_block_ptr = tl.make_block_ptr(

514 - base=K + k_offset,

515 - shape=(ACTUAL_BLOCK_DMODEL, seqlen_k),

516 - strides=(stride_kk, stride_kn),

517 - offsets=(0, 0),

518 - block_shape=(BLOCK_DMODEL, BLOCK_N),

519 - order=(0, 1),

520 - )

521 - v_offset = off_z * stride_vz + off_h_k * stride_vh + cu_seqlens_k_start * stride_vk

522 - V_block_ptr = tl.make_block_ptr(

523 - base=V + v_offset,

524 - shape=(seqlen_k, ACTUAL_BLOCK_DMODEL),

525 - strides=(stride_vk, stride_vn),

526 - offsets=(0, 0),

527 - block_shape=(BLOCK_N, BLOCK_DMODEL),

528 - order=(1, 0),

529 - )

530 - if BIAS_TYPE != 0:

531 - bias_ptr = tl.make_block_ptr(

532 - base=bias + off_h_q * stride_bh,

533 - shape=(seqlen_q, seqlen_k),

534 - strides=(stride_bm, stride_bn),

535 - offsets=(start_m * BLOCK_M, 0),

536 - block_shape=(BLOCK_M, BLOCK_N),

537 - order=(1, 0),

538 - )

539 - else:

540 - bias_ptr = None

541 - if ENABLE_DROPOUT:

542 - batch_philox_offset = (

543 - philox_offset_base + (off_z * HQ + off_h_q) * seqlen_q * seqlen_k

544 - )

545 - else:

546 - batch_philox_offset = 0

547 - # We can ask to return the dropout mask without actually doing any dropout.

548 - # In this case, we return an invalid pointer so indicate the mask is not i

549 - # valid.

550 - # TODO: Fix encoded softmax. It currently uses just h_q in the base offset.

551 - if RETURN_ENCODED_SOFTMAX:

552 - encoded_softmax_block_ptr = tl.make_block_ptr(

553 - base=encoded_softmax + off_h_q * seqlen_q * seqlen_k,

554 - shape=(seqlen_q, seqlen_k),

555 - strides=(seqlen_k, 1),

556 - offsets=(start_m * BLOCK_M, 0),

557 - block_shape=(BLOCK_M, BLOCK_N),

558 - order=(1, 0),

559 - )

560 - else:

561 - encoded_softmax_block_ptr = 0

562 - # initialize pointer to m and l

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

564 - l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

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

566 - # scale sm_scale by log_2(e) and use 2^x in the loop as we do not

567 - # have native e^x support in HW.

568 - qk_scale = sm_scale * 1.44269504089

569 - # Q is loaded once at the beginning and shared by all N blocks.

570 - q = load_fn(Q_block_ptr, True, padded_head, "zero")

571 - if not USE_FP8:

572 - q = (q * qk_scale).to(Q_block_ptr.type.element_ty)

573 - acc_scale = 1.0

574 - else:

575 - qk_scale *= q_scale * k_scale

576 - acc_scale = p_scale * v_scale

577 -

578 - # Here we compute how many full and masked blocks we have.

579 - padded_block_k = n_extra_tokens != 0

580 - is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0)

581 - if IS_CAUSAL:

582 - # There are always at least BLOCK_M // BLOCK_N masked blocks.

583 - # Additionally there might be one more due to dissimilar seqlens.

584 - masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn)

585 - else:

586 - # Padding on Q does not need to be masked in the FA loop.

587 - masked_blocks = padded_block_k

588 - # if IS_CAUSAL, not is_modulo_mn does not always result in an additional

589 - # block. In this case we might exceed n_blocks so pick the min.

590 - masked_blocks = min(masked_blocks, n_blocks)

591 - n_full_blocks = n_blocks - masked_blocks

592 - block_min = 0

593 - block_max = n_blocks * BLOCK_N

594 - # Compute for full blocks. Here we set causal to false regardless of its

595 - # value because there is no masking. Similarly we do not need padding.

596 - if n_full_blocks > 0:

597 - block_max = (n_blocks - masked_blocks) * BLOCK_N

598 - acc, l_i, m_i = _attn_fwd_inner(

599 - acc,

600 - l_i,

601 - m_i,

602 - q,

603 - K_block_ptr,

604 - V_block_ptr,

605 - start_m,

606 - seqlen_k,

607 - dropout_p,

608 - philox_seed,

609 - batch_philox_offset,

610 - encoded_softmax_block_ptr,

611 - # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _

612 - block_min,

613 - block_max,

614 - 0,

615 - 0,

616 - 0,

617 - bias_ptr,

618 - # IS_CAUSAL, ....

619 - False,

620 - BLOCK_M,

621 - BLOCK_DMODEL,

622 - BLOCK_N,

623 - offs_m,

624 - offs_n,

625 - # _, MASK_STEPS, ...

626 - PRE_LOAD_V,

627 - False,

628 - ENABLE_DROPOUT,

629 - RETURN_ENCODED_SOFTMAX,

630 - padded_head,

631 - USE_FP8,

632 - qk_scale,

633 - p_descale,

634 - )

635 - block_min = block_max

636 - block_max = n_blocks * BLOCK_N

637 -

638 - tl.debug_barrier()

639 - # Remaining blocks, if any, are full / not masked.

640 - if masked_blocks > 0:

641 - offs_n_causal = offs_n + (seqlen_q - seqlen_k) if IS_CAUSAL else 0

642 - K_block_ptr = tl.advance(K_block_ptr, (0, n_full_blocks * BLOCK_N))

643 - V_block_ptr = tl.advance(V_block_ptr, (n_full_blocks * BLOCK_N, 0))

644 - if bias_ptr is not None:

645 - bias_ptr = tl.advance(bias_ptr, (0, n_full_blocks * BLOCK_N))

646 - if RETURN_ENCODED_SOFTMAX:

647 - encoded_softmax_block_ptr = tl.advance(

648 - encoded_softmax_block_ptr, (0, n_full_blocks)

649 - )

650 - acc, l_i, m_i = _attn_fwd_inner(

651 - acc,

652 - l_i,

653 - m_i,

654 - q,

655 - K_block_ptr,

656 - V_block_ptr,

657 - start_m,

658 - seqlen_k,

659 - dropout_p,

660 - philox_seed,

661 - batch_philox_offset,

662 - encoded_softmax_block_ptr,

663 - block_min,

664 - block_max,

665 - offs_n_causal,

666 - masked_blocks,

667 - n_extra_tokens,

668 - bias_ptr,

669 - IS_CAUSAL,

670 - BLOCK_M,

671 - BLOCK_DMODEL,

672 - BLOCK_N,

673 - offs_m,

674 - offs_n,

675 - # _, MASK_STEPS, ...

676 - PRE_LOAD_V,

677 - True,

678 - ENABLE_DROPOUT,

679 - RETURN_ENCODED_SOFTMAX,

680 - padded_head,

681 - USE_FP8,

682 - qk_scale,

683 - p_descale,

684 - )

685 - # epilogue

686 -

687 - if USE_FP8:

688 - acc *= acc_scale

689 - acc = acc / l_i[:, None]

690 - if ENABLE_DROPOUT:

691 - acc = acc / (1 - dropout_p)

692 - # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M,

693 - # then we have one block with a row of all NaNs which come from computing

694 - # softmax over a row of all -infs (-inf - inf = NaN). We check for that here

695 - # and store 0s where there are NaNs as these rows should've been zeroed out.

696 - end_m_idx = (start_m + 1) * BLOCK_M

697 - start_m_idx = start_m * BLOCK_M

698 - causal_start_idx = seqlen_q - seqlen_k

699 - if USE_FP8_OUT:

700 - acc *= o_descale

701 - acc = tl.clamp(acc, FP8_MIN, FP8_MAX)

702 - acc = acc.to(Out.type.element_ty)

703 - if IS_CAUSAL: # noqa: SIM102

704 - if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx:

705 - out_mask_boundary = tl.full(

706 - (BLOCK_DMODEL,), causal_start_idx, dtype=tl.int32

707 - )

708 - mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)

709 - out_ptrs_mask = mask_m_offsets[:, None] >= out_mask_boundary[None, :]

710 - z = tl.zeros((1,), tl.float32)

711 - acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty))

712 - # write back LSE

713 - # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q + offs_m

714 - # If seqlen_q not multiple of BLOCK_M, we need to mask out the last

715 - # few rows. This is only true for the last M block. For others,

716 - # overflow_size will be -ve

717 - # overflow_size = end_m_idx - seqlen_q

718 - # if overflow_size > 0:

719 - # boundary = tl.full((BLOCK_M,), BLOCK_M - overflow_size, dtype=tl.int32)

720 - # # This is a > check because mask being 0 blocks the store.

721 - # l_ptrs_mask = boundary > tl.arange(0, BLOCK_M)

722 - # tl.store(l_ptrs, m_i + tl.math.log2(l_i), mask=l_ptrs_mask)

723 - # else:

724 - # tl.store(l_ptrs, m_i + tl.math.log2(l_i))

725 -

726 - # write back O

727 - o_offset = off_z * stride_oz + cu_seqlens_q_start * stride_om + off_h_q * stride_oh

728 - O_block_ptr = tl.make_block_ptr(

729 - base=Out + o_offset,

730 - shape=(seqlen_q, ACTUAL_BLOCK_DMODEL),

731 - strides=(stride_om, stride_on),

732 - offsets=(start_m * BLOCK_M, 0),

733 - block_shape=(BLOCK_M, BLOCK_DMODEL),

734 - order=(1, 0),

735 - )

736 - # Need boundary check on this to make sure the padding from the

737 - # Q and KV tensors in both dims are not part of what we store back.

738 - # TODO: Do the boundary check optionally.

739 - tl.store(O_block_ptr, acc, boundary_check=(0, 1))

740 -

741 -

742 -def check_args(

743 - q,

744 - k,

745 - v,

746 - o,

747 - varlen=True,

748 - max_seqlens=None,

749 - cu_seqlens_q=None,

750 - cu_seqlens_k=None,

751 -):

752 - assert q.dim() == k.dim() and q.dim() == v.dim()

753 - if varlen:

754 - assert q.dim() == 3

755 - total_q, nheads_q, head_size = q.shape

756 - total_k, nheads_k, _ = k.shape

757 - assert cu_seqlens_q is not None

758 - assert cu_seqlens_k is not None

759 - assert len(cu_seqlens_q) == len(cu_seqlens_k)

760 - else:

761 - assert q.dim() == 4

762 - batch, nheads_q, seqlen_q, head_size = q.shape

763 - _, nheads_k, seqlen_k, _ = k.shape

764 - assert max_seqlens > 0

765 - assert k.shape == v.shape

766 - assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1]

767 - # TODO: Change assert if we support qkl f8 and v f16

768 - assert q.dtype == k.dtype and q.dtype == v.dtype

769 - assert head_size <= 256

770 - assert o.shape == q.shape

771 - assert (nheads_q % nheads_k) == 0

772 -

773 -

774 -class _attention(torch.autograd.Function):

775 - @staticmethod

776 - def forward(

777 - ctx,

778 - q,

779 - k,

780 - v,

781 - o,

782 - cu_seqlens_q,

783 - cu_seqlens_k,

784 - max_seqlens_q,

785 - max_seqlens_k,

786 - causal=False,

787 - sm_scale=1.0,

788 - bias=None,

789 - fp8_scales=None,

790 - fp8_out_scale=None,

791 - ):

792 - if fp8_scales is not None:

793 - use_fp8 = True

794 - (q_scale, k_scale, v_scale, p_scale) = fp8_scales

795 - float8 = current_platform.fp8_dtype()

796 -

797 - def check_and_convert(t, scale):

798 - if t.dtype != float8:

799 - descale = 1.0 / scale

800 - ts = (t * descale).clamp(min=float8_info.min, max=float8_info.max)

801 - return ts.to(float8)

802 - else:

803 - return t

804 -

805 - q = check_and_convert(q, q_scale)

806 - k = check_and_convert(k, k_scale)

807 - v = check_and_convert(v, v_scale)

808 - else:

809 - use_fp8 = False

810 - q_scale = k_scale = v_scale = p_scale = 1.0

811 -

812 - if o is None:

813 - o = torch.empty_like(q, dtype=v.dtype)

814 -

815 - check_args(

816 - q,

817 - k,

818 - v,

819 - o,

820 - varlen=True,

821 - cu_seqlens_q=cu_seqlens_q,

822 - cu_seqlens_k=cu_seqlens_k,

823 - )

824 - if True: # varlen

825 - total_q, nheads_q, head_size = q.shape

826 - total_k, nheads_k, _ = k.shape

827 - batch = len(cu_seqlens_q) - 1

828 - q_strides = (0, q.stride(1), q.stride(0), q.stride(2))

829 - k_strides = (0, k.stride(1), k.stride(0), k.stride(2))

830 - v_strides = (0, v.stride(1), v.stride(0), v.stride(2))

831 - o_strides = (0, o.stride(1), o.stride(0), o.stride(2))

832 - else:

833 - batch, seqlen_q, nheads_q, head_size = q.shape

834 - _, seqlen_k, nheads_k, _ = k.shape

835 - q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))

836 - k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))

837 - v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))

838 - o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))

839 -

840 - # Get closest power of 2 over or equal to 32.

841 - unpadded_head_dims = {32, 64, 128, 256}

842 - if head_size not in unpadded_head_dims:

843 - padded_d_model = None

844 - for i in unpadded_head_dims:

845 - if i > head_size:

846 - padded_d_model = i

847 - break

848 - assert padded_d_model is not None

849 - else:

850 - padded_d_model = head_size

851 -

852 - grid = lambda META: (

853 - triton.cdiv(max_seqlens_q, META["BLOCK_M"]),

854 - nheads_q,

855 - batch,

856 - )

857 -

858 - encoded_softmax = None

859 -

860 - # Seed the RNG so we get reproducible results for testing.

861 - philox_seed = 0x1BF52

862 - philox_offset = 0x1D4B42

863 -

864 - if bias is not None:

865 - bias_strides = (

866 - bias.stride(0),

867 - bias.stride(1),

868 - bias.stride(2),

869 - bias.stride(3),

870 - )

871 - else:

872 - bias_strides = (0, 0, 0, 0)

873 -

874 - p_descale = 1.0 / p_scale

875 - o_descale = 1.0 / fp8_out_scale.item() if fp8_out_scale is not None else 1.0

876 -

877 - arg_max_seqlens_q = 0 if on_gfx1x() else max_seqlens_q

878 - arg_max_seqlens_k = 0 if on_gfx1x() else max_seqlens_k

879 -

880 - attn_fwd[grid](

881 - q,

882 - k,

883 - v,

884 - bias,

885 - sm_scale,

886 - q_scale,

887 - k_scale,

888 - v_scale,

889 - p_scale,

890 - p_descale,

891 - o_descale,

892 - None,

893 - o,

894 - *q_strides,

895 - *k_strides,

896 - *v_strides,

897 - *o_strides,

898 - *bias_strides,

899 - cu_seqlens_q,

900 - cu_seqlens_k,

901 - dropout_p=0.0,

902 - philox_seed=philox_seed,

903 - philox_offset_base=philox_offset,

904 - encoded_softmax=encoded_softmax,

905 - HQ=nheads_q,

906 - HK=nheads_k,

907 - ACTUAL_BLOCK_DMODEL=head_size,

908 - MAX_SEQLENS_Q=arg_max_seqlens_q,

909 - MAX_SEQLENS_K=arg_max_seqlens_k,

910 - IS_CAUSAL=causal,

911 - VARLEN=True,

912 - BLOCK_DMODEL=padded_d_model,

913 - BIAS_TYPE=0 if bias is None else 1,

914 - ENABLE_DROPOUT=False,

915 - RETURN_ENCODED_SOFTMAX=False,

916 - USE_FP8=use_fp8,

917 - USE_FP8_OUT=fp8_out_scale is not None,

918 - )

919 -

920 - ctx.grid = grid

921 - ctx.sm_scale = sm_scale

922 - ctx.BLOCK_DMODEL = head_size

923 - ctx.causal = causal

924 - ctx.dropout_p = 0.0

925 - ctx.philox_seed = philox_seed

926 - ctx.philox_offset = philox_offset

927 - ctx.encoded_softmax = encoded_softmax

928 - ctx.return_encoded_softmax = False

929 - return o, encoded_softmax

930 -

931 -

932 -triton_attention = _attention.apply

vllm/envs.py

-6

@@ -18,7 +18,6 @@ if TYPE_CHECKING:

18 VLLM_RINGBUFFER_WARNING_INTERVAL: int = 60

19 VLLM_NCCL_SO_PATH: str | None = None

20 LD_LIBRARY_PATH: str | None = None

21 - VLLM_USE_TRITON_FLASH_ATTN: bool = True

21 VLLM_V1_USE_PREFILL_DECODE_ATTENTION: bool = False

22 VLLM_FLASH_ATTN_VERSION: int | None = None

23 LOCAL_RANK: int = 0

@@ -521,10 +520,6 @@ environment_variables: dict[str, Callable[[], Any]] = {

520 # when `VLLM_NCCL_SO_PATH` is not set, vllm will try to find the nccl

521 # library file in the locations specified by `LD_LIBRARY_PATH`

522 "LD_LIBRARY_PATH": lambda: os.environ.get("LD_LIBRARY_PATH", None),

524 - # flag to control if vllm should use triton flash attention

525 - "VLLM_USE_TRITON_FLASH_ATTN": lambda: (

526 - os.environ.get("VLLM_USE_TRITON_FLASH_ATTN", "True").lower() in ("true", "1")

527 - ),

523 # Use separate prefill and decode kernels for V1 attention instead of

524 # the unified triton kernel.

525 "VLLM_V1_USE_PREFILL_DECODE_ATTENTION": lambda: (

@@ -1554,7 +1549,6 @@ def compute_hash() -> str:

1549 "VLLM_PP_LAYER_PARTITION",

1550 "VLLM_MLA_DISABLE",

1551 "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH",

1557 - "VLLM_USE_TRITON_FLASH_ATTN",

1552 "VLLM_USE_TRITON_AWQ",

1553 "VLLM_DP_RANK",

1554 "VLLM_DP_SIZE",

vllm/platforms/rocm.py

+2 -19

@@ -49,25 +49,8 @@ _ROCM_UNSUPPORTED_MODELS: list[str] = []

50 # Models partially supported by ROCm.

51 # Architecture -> Reason.

52 -_ROCM_SWA_REASON = (

53 - "Sliding window attention (SWA) is not yet supported in "

54 - "Triton flash attention. For half-precision SWA support, "

55 - "please use CK flash attention by setting "

56 - "`VLLM_USE_TRITON_FLASH_ATTN=0`"

57 -)

58 -_ROCM_PARTIALLY_SUPPORTED_MODELS: dict[str, str] = {

59 - "Qwen2ForCausalLM": _ROCM_SWA_REASON,

60 - "MistralForCausalLM": _ROCM_SWA_REASON,

61 - "MixtralForCausalLM": _ROCM_SWA_REASON,

62 - "PaliGemmaForConditionalGeneration": (

63 - "ROCm flash attention does not yet fully support 32-bit precision on PaliGemma"

64 - ),

65 - "Phi3VForCausalLM": (

66 - "ROCm Triton flash attention may run into compilation errors due to "

67 - "excessive use of shared memory. If this happens, disable Triton FA "

68 - "by setting `VLLM_USE_TRITON_FLASH_ATTN=0`"

69 - ),

70 -}

52 +_ROCM_SWA_REASON = ()

53 +_ROCM_PARTIALLY_SUPPORTED_MODELS: dict[str, str] = {}

54 _ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = {

55 "0x74a0": "AMD_Instinct_MI300A",

56 "0x74a1": "AMD_Instinct_MI300X",

vllm/usage/usage_lib.py

-1

@@ -37,7 +37,6 @@ _GLOBAL_RUNTIME_DATA = dict[str, str | int | bool]()

38 _USAGE_ENV_VARS_TO_COLLECT = [

39 "VLLM_USE_MODELSCOPE",

40 - "VLLM_USE_TRITON_FLASH_ATTN",

40 "VLLM_ATTENTION_BACKEND",

41 "VLLM_USE_FLASHINFER_SAMPLER",

42 "VLLM_PP_LAYER_PARTITION",

vllm/v1/attention/backends/mla/triton_mla.py

+7 -48

@@ -5,22 +5,18 @@ from typing import ClassVar

6 import torch

8 -from vllm import envs

8 from vllm.attention.backends.abstract import (

9 AttentionLayer,

10 AttentionType,

11 is_quantized_kv_cache,

12 )

13 from vllm.attention.ops.triton_decode_attention import decode_attention_fwd

15 -from vllm.attention.ops.triton_flash_attention import triton_attention

14 from vllm.config.cache import CacheDType

15 from vllm.logger import init_logger

16 from vllm.model_executor.layers.batch_invariant import (

17 vllm_is_batch_invariant,

18 )

21 -from vllm.platforms import current_platform

19 from vllm.platforms.interface import DeviceCapability

23 -from vllm.triton_utils import HAS_TRITON

20 from vllm.v1.attention.backends.mla.common import (

21 MLACommonBackend,

22 MLACommonImpl,

@@ -99,55 +95,18 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

95 "TritonMLA V1 with FP8 KV cache not yet supported"

96 )

102 - self.use_triton_flash_attn = envs.VLLM_USE_TRITON_FLASH_ATTN

103 - self.triton_fa_func = triton_attention if HAS_TRITON else None

104 -

105 - def _flash_attn_varlen_diff_headdims_rocm(

106 - self, q, k, v, softmax_scale=None, **kwargs

98 + def _flash_attn_varlen_diff_headdims(

99 + self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

100 ):

108 - assert self.triton_fa_func is not None

109 -

110 - # Triton Attention requires a padded V

111 - padded_v = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]], value=0)

112 - # The output of triton_attention is a tuple of

113 - # [output_tensor, encoded_softmax] where encoded_softmax is always None

114 - output_tensor, _ = self.triton_fa_func(

101 + return super()._flash_attn_varlen_diff_headdims(

102 q,

103 k,

117 - padded_v,

118 - None, # output

119 - kwargs["cu_seqlens_q"],

120 - kwargs["cu_seqlens_k"],

121 - kwargs["max_seqlen_q"],

122 - kwargs["max_seqlen_k"],

123 - kwargs["causal"],

124 - softmax_scale,

125 - None, # bias

104 + v,

105 + return_softmax_lse=return_softmax_lse,

106 + softmax_scale=softmax_scale,

107 + **kwargs,

108 )

109

128 - return output_tensor

129 -

130 - def _flash_attn_varlen_diff_headdims(

131 - self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

132 - ):

133 - if (

134 - current_platform.is_rocm()

135 - and self.use_triton_flash_attn

136 - and not return_softmax_lse

137 - ):

138 - return self._flash_attn_varlen_diff_headdims_rocm(

139 - q, k, v, softmax_scale=softmax_scale, **kwargs

140 - )

141 - else:

142 - return super()._flash_attn_varlen_diff_headdims(

143 - q,

144 - k,

145 - v,

146 - return_softmax_lse=return_softmax_lse,

147 - softmax_scale=softmax_scale,

148 - **kwargs,

149 - )

150 -

110 def _forward_decode(

111 self,

112 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],