@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0a309b5ee9

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 3, 22:45 UTC

[ROCm] Cap Triton paged attention block size to fix ROCm shared memory OOM (#38502)

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

Andreas Karatzas committed May 10, 2026 at 05:03 UTC 0a309b5ee9480d0d78e7f00946a418ffd7e3baf1

3 files changed +44 -18

.buildkite/test-amd.yaml

+9 -6

@@ -1803,9 +1803,10 @@ steps:

1803 - tests/models/multimodal/generation

1804 - tests/models/multimodal/test_mapping.py

1805 commands:

1806 - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git

1807 - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py

1808 - - pytest -v -s models/multimodal/test_mapping.py

1806 + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'

1807 + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'

1808 + - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB

1809 +

1810

1811 - label: Multi-Modal Models (Extended Generation 2) # TBD

1812 timeout_in_minutes: 180

@@ -1817,8 +1818,10 @@ steps:

1818 - vllm/

1819 - tests/models/multimodal/generation

1820 commands:

1820 - - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git

1821 - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'

1821 + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'

1822 + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'

1823 + - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

1824 +

1825

1826 - label: Multi-Modal Models (Extended Generation 3) # TBD

1827 timeout_in_minutes: 180

@@ -3043,7 +3046,7 @@ steps:

3046 - vllm/

3047 - tests/models/language/generation

3048 commands:

3046 - - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'

3049 + - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'

3050 - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'

3051 - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

3052

vllm/v1/attention/backends/rocm_attn.py

+8 -4

@@ -29,6 +29,7 @@ from vllm.v1.attention.backend import (

29 )

30 from vllm.v1.attention.ops.chunked_prefill_paged_decode import (

31 chunked_prefill_paged_decode,

32 + has_native_kv_cache_layout,

33 )

34 from vllm.v1.attention.ops.paged_attn import PagedAttention

35 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (

@@ -468,9 +469,10 @@ class RocmAttentionImpl(AttentionImpl):

469 # Get the actual block_size from value_cache

470 # value_cache shape: [num_blocks, num_heads, head_size, block_size]

471 block_size = value_cache.shape[3]

472 + has_native_layout = has_native_kv_cache_layout(key_cache, value_cache)

473

472 - if block_size in (16, 32):

473 - # Normal 16, 32, use vLLM native HIP C++ logic

474 + if block_size in (16, 32) and has_native_layout:

475 + # Normal 16, 32 with contiguous blocks: use vLLM native HIP C++ logic.

476 PagedAttention.write_to_paged_cache(

477 key,

478 value,

@@ -482,8 +484,10 @@ class RocmAttentionImpl(AttentionImpl):

484 layer._v_scale,

485 )

486 else:

485 - # Case B: Non-standard blocks (e.g., 64, 128, 544 in Qwen3Next or Qwen3.5 ),

486 - # force using our modified Triton logic

487 + # Non-standard blocks and hybrid attention/Mamba layouts need the

488 + # stride-aware Triton writer. The native reshape_and_cache kernel

489 + # assumes contiguous block storage and writes to the wrong hybrid

490 + # cache blocks.

491 triton_reshape_and_cache_flash(

492 key,

493 value,

vllm/v1/attention/ops/chunked_prefill_paged_decode.py

+27 -8

@@ -21,6 +21,22 @@ logger = init_logger(__name__)

21 float8_info = torch.finfo(current_platform.fp8_dtype())

24 +def has_native_kv_cache_layout(

25 + key_cache: torch.Tensor,

26 + value_cache: torch.Tensor,

27 +) -> bool:

28 + """Return whether KV cache blocks can use the native ROCm pairing.

29 +

30 + The native reshape_and_cache writer assumes packed blocks. If cache update

31 + needs reshape_and_cache_flash for a stride-padded hybrid layout, decode

32 + should use the matching Triton path too.

33 + """

34 + return (

35 + key_cache.stride(0) == key_cache.shape[1:].numel()

36 + and value_cache.stride(0) == value_cache.shape[1:].numel()

37 + )

38 +

39 +

40 @triton.jit

41 def cdiv_fn(x, y):

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

@@ -346,14 +362,12 @@ def chunked_prefill_paged_decode(

362 alibi_slopes,

363 sinks,

364 )

349 - # Triton is only forced when encountering a non-standard block

350 - # like Qwen3 with a size of 544.

351 - # 1. Check if block_size is a power of 2 (16, 32, 64...)

352 - # 2. If it's a power of 2, we trust the vLLM's native use_custom decision.

353 - # 3. If it's not a power of 2 (such as Qwen3's 544),

354 - # then our Triton path is forced.

365 + has_native_layout = has_native_kv_cache_layout(key_cache, value_cache)

366 + # Force Triton for non-standard blocks like Qwen3's 544 and for

367 + # stride-padded hybrid layouts. The latter use reshape_and_cache_flash

368 + # during cache update, so keep decode on the matching stride-aware path.

369 is_pow2 = block_size > 0 and (block_size & (block_size - 1) == 0)

356 - if not is_pow2:

370 + if not is_pow2 or not has_native_layout:

371 use_custom = False

372

373 if use_custom:

@@ -404,7 +418,12 @@ def chunked_prefill_paged_decode(

418 real_block_size = value_cache.shape[3]

419 # The standard model directly uses the original block_size.

420 # Non-standard 544 uses 32 to accommodate integer division logic.

407 - TRITON_BLOCK_SIZE = block_size if is_pow2 else 32

421 + # Cap at 128 to avoid exceeding GPU shared memory limits

422 + # (e.g. hybrid Mamba models inflate block_size to 2048).

423 + # The kernel handles TRITON_BLOCK_SIZE != PHYSICAL_BLOCK_SIZE

424 + # via the l_block_idx/internal_offsets addressing logic.

425 + MAX_TRITON_BLOCK_SIZE = 128

426 + TRITON_BLOCK_SIZE = min(block_size, MAX_TRITON_BLOCK_SIZE) if is_pow2 else 32

427 if is_block_table_ptr:

428 # Using the physical base address of tensors

429 kv_element_size = key_cache.element_size()