From f3818657a2f0554b7486c9663f658ae342058374 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Thu, 13 Aug 2026 20:24:17 +0000 Subject: [PATCH 1/8] [ROCm][Perf] Optimize DSV4 sparse MLA decode for long contexts Improve low-batch long-context split-K parallelism, compile cache geometry into the Triton kernels, decode UE8M0 scales with exponent bits, and write directly to the caller output. Add production-shape coverage for 8192 selected KV rows and split count 32. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai --- .../attention/test_rocm_triton_attn_dsv4.py | 95 ++++++++++++++++++- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 71 +++++++++----- 2 files changed, 139 insertions(+), 27 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 6fe2a3e77587..775f1f5d1447 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -142,6 +142,34 @@ def _read_fp8_ds_mla_cache( return torch.cat([nope, rope]) +def _read_fp8_ds_mla_cache_rows( + cache: torch.Tensor, + slots: torch.Tensor, + block_size: int, + use_fnuz: bool, +) -> torch.Tensor: + cache_flat = cache.view(torch.uint8).flatten() + block_idx = slots // block_size + pos = slots % block_size + block_base = block_idx * cache.stride(0) + token_base = block_base + pos * 576 + scale_base = block_base + block_size * 576 + pos * 8 + + fp8_dtype = torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn + nope_offsets = torch.arange(NOPE_HEAD_DIM, device=cache.device) + nope_u8 = cache_flat[token_base[:, None] + nope_offsets] + nope = nope_u8.view(fp8_dtype).to(torch.float32) + scale_offsets = torch.arange(7, device=cache.device) + scales = torch.exp2( + cache_flat[scale_base[:, None] + scale_offsets].to(torch.float32) - 127.0 + ) + nope = nope * scales.repeat_interleave(64, dim=1) + rope_offsets = torch.arange(ROPE_HEAD_DIM * 2, device=cache.device) + rope_u8 = cache_flat[token_base[:, None] + NOPE_HEAD_DIM + rope_offsets] + rope = rope_u8.contiguous().view(torch.bfloat16).to(torch.float32) + return torch.cat([nope, rope], dim=1) + + def _ref_sparse_decode_ragged( q: torch.Tensor, main_cache: torch.Tensor, @@ -366,6 +394,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None: attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device) scale = HEAD_DIM**-0.5 + out = torch.empty_like(q) actual = _rocm_sparse_attn_decode_ragged_triton( q=q, main_cache=main_cache, @@ -378,6 +407,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None: extra_cache=extra_cache, extra_indices=extra_indices, extra_indptr=extra_indptr, + out=out, ) expected = _ref_sparse_decode_ragged( q=q, @@ -391,6 +421,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None: main_use_fnuz=main_use_fnuz, ) + assert actual.data_ptr() == out.data_ptr() torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) @@ -408,18 +439,23 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: # A tiny batch on a large device should split to add parallelism. assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1 - # The chosen count always stays within the searched [1, 16] range, and a + # Long C128A rows need 32 splits to fill a 256-CU device at low batch. + assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 32 + assert mod._decode_num_splits(8, 1, 128.0, 8192.0) == 32 + assert mod._decode_num_splits(64, 1, 128.0, 8192.0) == 4 + + # The chosen count always stays within the searched [1, 32] range, and a # zero-length workload never splits (no work to parallelize). for num_queries in (1, 4, 24, 224, 1024): splits = mod._decode_num_splits( num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0 ) - assert 1 <= splits <= 16 + assert 1 <= splits <= 32 assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 @requires_split_decode_arch -@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8]) +@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8, 32]) @pytest.mark.parametrize("with_extra", [True, False]) @pytest.mark.parametrize("with_sink", [True, False]) @torch.inference_mode() @@ -503,6 +539,59 @@ def test_sparse_attn_decode_split_k_kernel( torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_split_decode_arch +@torch.inference_mode() +def test_sparse_attn_decode_long_context(monkeypatch) -> None: + """Validate the production C128A maximum of 8192 selected KV rows.""" + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(11) + block_size = 64 + num_kv = 8192 + num_heads = 16 + q = torch.randn(1, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 + main_cache = _pack_fp8_ds_mla_cache( + torch.zeros(1, HEAD_DIM, dtype=torch.bfloat16, device=device), + block_size, + use_fnuz=current_platform.is_fp8_fnuz(), + ) + extra_cache = _pack_fp8_ds_mla_cache( + torch.randn(num_kv, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125, + block_size, + use_fnuz=False, + ) + empty_indices = torch.empty(0, dtype=torch.int32, device=device) + empty_indptr = torch.zeros(2, dtype=torch.int32, device=device) + extra_indices = torch.arange(num_kv, dtype=torch.int32, device=device) + extra_indptr = torch.tensor([0, num_kv], dtype=torch.int32, device=device) + scale = HEAD_DIM**-0.5 + + monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: 32) + actual = mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=main_cache, + main_indices=empty_indices, + main_indptr=empty_indptr, + scale=scale, + attn_sink=None, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + ) + + kv = _read_fp8_ds_mla_cache_rows( + extra_cache, extra_indices.to(torch.int64), block_size, use_fnuz=False + ) + scores = torch.matmul(q[0].float(), kv.T) * scale + expected = torch.matmul(torch.softmax(scores, dim=-1), kv) + torch.testing.assert_close( + actual[0], expected.to(torch.bfloat16), atol=2e-2, rtol=2e-2 + ) + + # --------------------------------------------------------------------------- # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) # --------------------------------------------------------------------------- diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index ed494ce2dae4..cf1d4d44ed2c 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1218,6 +1218,13 @@ def _sparse_attn_prefill_ragged_kernel( ) +@triton.jit +def _decode_e8m0_scales_triton(encoded_scales): + scale_bits = encoded_scales.to(tl.int32) << 23 + scale_bits = tl.where(encoded_scales == 0, 1 << 22, scale_bits) + return scale_bits.to(tl.float32, bitcast=True) + + @triton.jit def _sparse_attn_decode_ragged_kernel( q_ptr, @@ -1237,8 +1244,8 @@ def _sparse_attn_decode_ragged_kernel( extra_cache_stride0, main_num_rows, extra_num_rows, - main_block_size, - extra_block_size, + MAIN_BLOCK_SIZE: tl.constexpr, + EXTRA_BLOCK_SIZE: tl.constexpr, scale, num_heads, HAS_ATTN_SINK: tl.constexpr, @@ -1295,11 +1302,11 @@ def _sparse_attn_decode_ragged_kernel( valid = in_range & (slot >= 0) & (slot < main_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // main_block_size - pos_in_block = safe_slot % main_block_size + block_idx = safe_slot // MAIN_BLOCK_SIZE + pos_in_block = safe_slot % MAIN_BLOCK_SIZE cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 + token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 x_uint8 = tl.load( token_data_ptr[:, None] + nope_offsets[None, :], @@ -1315,7 +1322,7 @@ def _sparse_attn_decode_ragged_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + scales = _decode_e8m0_scales_triton(encoded_scales) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1359,14 +1366,14 @@ def _sparse_attn_decode_ragged_kernel( valid = in_range & (slot >= 0) & (slot < extra_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // extra_block_size - pos_in_block = safe_slot % extra_block_size + block_idx = safe_slot // EXTRA_BLOCK_SIZE + pos_in_block = safe_slot % EXTRA_BLOCK_SIZE cache_block_ptr = ( extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 ) token_data_ptr = cache_block_ptr + pos_in_block * 576 token_scale_ptr = ( - cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 + cache_block_ptr + EXTRA_BLOCK_SIZE * 576 + pos_in_block * 8 ) x_uint8 = tl.load( @@ -1383,7 +1390,7 @@ def _sparse_attn_decode_ragged_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + scales = _decode_e8m0_scales_triton(encoded_scales) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1477,8 +1484,8 @@ def _sparse_attn_decode_partial_kernel( pa_stride_h, main_num_rows, extra_num_rows, - main_block_size, - extra_block_size, + MAIN_BLOCK_SIZE: tl.constexpr, + EXTRA_BLOCK_SIZE: tl.constexpr, scale, num_heads, HAS_EXTRA: tl.constexpr, @@ -1545,11 +1552,11 @@ def _sparse_attn_decode_partial_kernel( valid = in_range & (slot >= 0) & (slot < main_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // main_block_size - pos_in_block = safe_slot % main_block_size + block_idx = safe_slot // MAIN_BLOCK_SIZE + pos_in_block = safe_slot % MAIN_BLOCK_SIZE cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 + token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 x_uint8 = tl.load( token_data_ptr[:, None] + nope_offsets[None, :], @@ -1565,7 +1572,7 @@ def _sparse_attn_decode_partial_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + scales = _decode_e8m0_scales_triton(encoded_scales) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1612,14 +1619,14 @@ def _sparse_attn_decode_partial_kernel( valid = in_range & (slot >= 0) & (slot < extra_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // extra_block_size - pos_in_block = safe_slot % extra_block_size + block_idx = safe_slot // EXTRA_BLOCK_SIZE + pos_in_block = safe_slot % EXTRA_BLOCK_SIZE cache_block_ptr = ( extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 ) token_data_ptr = cache_block_ptr + pos_in_block * 576 token_scale_ptr = ( - cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 + cache_block_ptr + EXTRA_BLOCK_SIZE * 576 + pos_in_block * 8 ) x_uint8 = tl.load( @@ -1636,7 +1643,7 @@ def _sparse_attn_decode_partial_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + scales = _decode_e8m0_scales_triton(encoded_scales) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1971,8 +1978,9 @@ def _decode_num_splits( mu = 0.04 best_splits = 1 best_cost = None - # Search up to 16 splits; beyond that the reduce/HBM overhead dominates. - for splits in range(1, 17): + # A full C128A row can contain 8K KV tokens. At low batch sizes, 32 splits + # keep those long rows to one device wave and halve the partial iterations. + for splits in range(1, 33): waves = (base * splits + cu - 1) // cu cost = waves * (1.0 / splits + mu) if best_cost is None or cost < best_cost - 1e-9: @@ -2005,6 +2013,7 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_cache: torch.Tensor | None = None, extra_indices: torch.Tensor | None = None, extra_indptr: torch.Tensor | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert main_cache.ndim == 3, ( @@ -2066,7 +2075,16 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32) block_h = 16 - out = torch.empty_like(q, dtype=torch.bfloat16) + if out is None: + out = torch.empty_like(q, dtype=torch.bfloat16) + else: + assert out.shape == q.shape, f"expected out shape {q.shape}, got {out.shape}" + assert out.device == q.device, ( + f"expected out on device {q.device}, got {out.device}" + ) + assert out.dtype == torch.bfloat16, ( + f"expected out dtype {torch.bfloat16}, got {out.dtype}" + ) heads_blocks = triton.cdiv(num_heads, block_h) nope_block = triton.next_power_of_2(nope_head_dim) comb_dim = nope_head_dim + rope_head_dim @@ -2213,6 +2231,7 @@ def _rocm_sparse_attn_decode_triton( main_ragged_indptr: torch.Tensor | None = None, extra_ragged_indices: torch.Tensor | None = None, extra_ragged_indptr: torch.Tensor | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: if main_ragged_indices is None or main_ragged_indptr is None: main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense( @@ -2248,6 +2267,7 @@ def _rocm_sparse_attn_decode_triton( extra_cache=extra_cache, extra_indices=extra_ragged_indices, extra_indptr=extra_ragged_indptr, + out=out, ) @@ -2348,6 +2368,7 @@ def rocm_sparse_attn_decode( if topk_indices is not None: extra_indices = topk_indices.reshape(topk_indices.shape[0], -1) + direct_out = output if output.dtype == torch.bfloat16 else None attn_out = _rocm_sparse_attn_decode_triton( q=q, main_cache=swa_k_cache, @@ -2364,5 +2385,7 @@ def rocm_sparse_attn_decode( main_ragged_indptr=swa_ragged_indptr, extra_ragged_indices=topk_ragged_indices, extra_ragged_indptr=topk_ragged_indptr, + out=direct_out, ) - output.copy_(attn_out.to(output.dtype)) + if direct_out is None: + output.copy_(attn_out.to(output.dtype)) From 38ef98046c5d976259a18f0e931a0ed24fe03884 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Fri, 14 Aug 2026 06:08:04 +0000 Subject: [PATCH 2/8] [ROCm][Perf] Add unified gfx950 sparse MLA decode kernel Use an in-tree Gluon-style gfx950 kernel with shared-LDS KV reuse, explicit MFMA layouts, pipelined gathers, direct S1 output, and reusable split-K scratch. Preserve the existing Triton kernels on other architectures. Validate eager and graph parity against the corrected AITER Gluon implementation across batch and long-context shapes. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai --- .../attention/test_rocm_triton_attn_dsv4.py | 176 +++- vllm/models/deepseek_v4/amd/rocm.py | 38 + .../v1/attention/ops/rocm_aiter_mla_sparse.py | 260 ++++-- .../ops/rocm_aiter_mla_sparse_gluon.py | 795 ++++++++++++++++++ 4 files changed, 1171 insertions(+), 98 deletions(-) create mode 100644 vllm/v1/attention/ops/rocm_aiter_mla_sparse_gluon.py diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 775f1f5d1447..40e4479f6cd0 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -24,18 +24,50 @@ def _on_split_decode_arch() -> bool: return False +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + try: + from vllm.platforms.rocm import _ON_GFX950 + + return bool(_ON_GFX950) + except Exception: + return False + + # The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other # architectures take the fallback decode kernel, so its tests are skipped there. requires_split_decode_arch = pytest.mark.skipif( not _on_split_decode_arch(), reason="split-K decode kernel is only tuned for AMD gfx942/gfx950", ) +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), reason="Gluon sparse decode kernel is only used on gfx950" +) NOPE_HEAD_DIM = 448 ROPE_HEAD_DIM = 64 HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM +def _make_split_k_buffers( + num_queries: int, + num_splits: int, + num_heads: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + partitions = num_queries * num_splits + return ( + torch.empty((partitions, num_heads), dtype=torch.float32, device=device), + torch.empty((partitions, num_heads), dtype=torch.float32, device=device), + torch.empty( + (partitions, num_heads, HEAD_DIM), + dtype=torch.float32, + device=device, + ), + ) + + def _ref_global_topk_ragged( topk_indices: torch.Tensor, token_to_req_indices: torch.Tensor, @@ -425,6 +457,66 @@ def test_sparse_attn_decode_ragged_kernel() -> None: torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_ue8m0_scale_groups() -> None: + """Each packed UE8M0 byte scales exactly one 64-column NoPE group.""" + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + block_size = 1 + cache = torch.zeros( + (1, block_size, 584), + dtype=torch.uint8, + device=device, + ) + cache_flat = cache.flatten() + one_bits = ( + torch.ones(NOPE_HEAD_DIM, dtype=torch.float32, device=device) + .to(torch.float8_e4m3fn) + .view(torch.uint8) + ) + cache_flat[:NOPE_HEAD_DIM].copy_(one_bits) + encoded_scales = torch.tensor( + [0, 125, 126, 127, 128, 129, 130], + dtype=torch.uint8, + device=device, + ) + cache_flat[576:583].copy_(encoded_scales) + cache_flat[583] = 0 + + q = torch.zeros((1, 1, HEAD_DIM), dtype=torch.bfloat16, device=device) + empty_indices = torch.empty(0, dtype=torch.int32, device=device) + empty_indptr = torch.zeros(2, dtype=torch.int32, device=device) + extra_indices = torch.zeros(1, dtype=torch.int32, device=device) + extra_indptr = torch.tensor([0, 1], dtype=torch.int32, device=device) + + actual = mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=cache, + main_indices=empty_indices, + main_indptr=empty_indptr, + scale=HEAD_DIM**-0.5, + attn_sink=None, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + ) + + expected_bits = encoded_scales.to(torch.int32) << 23 + expected_bits[0] = 1 << 22 + expected_scales = expected_bits.view(torch.float32) + expected = torch.cat( + [ + expected_scales.repeat_interleave(64), + torch.zeros(ROPE_HEAD_DIM, dtype=torch.float32, device=device), + ] + ).to(torch.bfloat16) + torch.testing.assert_close(actual[0, 0], expected, atol=1e-3, rtol=1e-3) + + @requires_split_decode_arch @torch.inference_mode() def test_decode_num_splits_heuristic(monkeypatch) -> None: @@ -439,28 +531,94 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: # A tiny batch on a large device should split to add parallelism. assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1 - # Long C128A rows need 32 splits to fill a 256-CU device at low batch. - assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 32 - assert mod._decode_num_splits(8, 1, 128.0, 8192.0) == 32 - assert mod._decode_num_splits(64, 1, 128.0, 8192.0) == 4 + # Preserve the standard-kernel limit while allowing gfx950's long-context + # path to opt into 32 splits. + assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16 + assert mod._decode_num_splits(1, 1, 128.0, 8192.0, max_splits=32) == 32 + assert mod._decode_num_splits(8, 1, 128.0, 8192.0, max_splits=32) == 32 + assert mod._decode_num_splits(64, 1, 128.0, 8192.0, max_splits=32) == 4 + max_queries = 512 + max_partitions = mod._max_decode_partitions(max_queries, 1, 32) + assert all( + queries + * mod._decode_num_splits( + queries, + 1, + avg_main_len=128.0, + avg_extra_len=8192.0, + max_splits=32, + ) + <= max_partitions + for queries in range(1, max_queries + 1) + ) - # The chosen count always stays within the searched [1, 32] range, and a + # The default count always stays within the standard [1, 16] range, and a # zero-length workload never splits (no work to parallelize). for num_queries in (1, 4, 24, 224, 1024): splits = mod._decode_num_splits( num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0 ) - assert 1 <= splits <= 32 + assert 1 <= splits <= 16 assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 +@requires_split_decode_arch +@torch.inference_mode() +def test_sparse_attn_decode_rejects_invalid_indices(monkeypatch) -> None: + """Full-tile sentinels and tail OOB slots are ignored.""" + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(17) + block_size = 4 + num_kv = 8 + num_heads = 3 + q = torch.randn(1, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) + kv = torch.randn(num_kv, HEAD_DIM, dtype=torch.bfloat16, device=device) + cache = _pack_fp8_ds_mla_cache( + kv, + block_size, + use_fnuz=current_platform.is_fp8_fnuz(), + ) + row = [i % num_kv for i in range(65)] + row[5] = -1 + row[64] = num_kv + indices, indptr = _ragged_from_rows([row], device) + + monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: 1) + actual = mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=cache, + main_indices=indices, + main_indptr=indptr, + scale=HEAD_DIM**-0.5, + attn_sink=None, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + ) + expected = _ref_sparse_decode_ragged( + q=q, + main_cache=cache, + main_rows=[[slot for slot in row if 0 <= slot < num_kv]], + scale=HEAD_DIM**-0.5, + attn_sink=None, + block_size=block_size, + main_use_fnuz=current_platform.is_fp8_fnuz(), + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + @requires_split_decode_arch @pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8, 32]) @pytest.mark.parametrize("with_extra", [True, False]) @pytest.mark.parametrize("with_sink", [True, False]) @torch.inference_mode() def test_sparse_attn_decode_split_k_kernel( - monkeypatch, num_splits: int, with_extra: bool, with_sink: bool + monkeypatch, + num_splits: int, + with_extra: bool, + with_sink: bool, ) -> None: """Flash-decode split-K decode path (partial + reduce kernels). @@ -510,6 +668,7 @@ def test_sparse_attn_decode_split_k_kernel( # Pin the split count so each parametrized value is exercised deterministically. monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits) + split_k_buffers = _make_split_k_buffers(num_queries, num_splits, num_heads, device) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, @@ -523,6 +682,7 @@ def test_sparse_attn_decode_split_k_kernel( extra_cache=extra_cache, extra_indices=extra_indices, extra_indptr=extra_indptr, + split_k_buffers=split_k_buffers, ) expected = _ref_sparse_decode_ragged( q=q, @@ -568,6 +728,7 @@ def test_sparse_attn_decode_long_context(monkeypatch) -> None: scale = HEAD_DIM**-0.5 monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: 32) + split_k_buffers = _make_split_k_buffers(1, 32, num_heads, device) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, main_cache=main_cache, @@ -580,6 +741,7 @@ def test_sparse_attn_decode_long_context(monkeypatch) -> None: extra_cache=extra_cache, extra_indices=extra_indices, extra_indptr=extra_indptr, + split_k_buffers=split_k_buffers, ) kv = _read_fp8_ds_mla_cache_rows( diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 23223acd2dba..4410d36b3f2f 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -19,6 +19,7 @@ DeepseekV4SparseMLAMetadataBuilder, ) from vllm.platforms import current_platform +from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( CommonAttentionMetadata, @@ -28,6 +29,8 @@ DeepseekSparseSWAMetadataBuilder, ) from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + SparseDecodeSplitKBuffers, + _max_decode_partitions, build_ragged_indices_from_dense, rocm_inv_rope_einsum, rocm_sparse_attn_decode, @@ -451,10 +454,44 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): backend_cls = DeepseekV4ROCMAiterMLASparseBackend def __init__(self, *args, **kwargs): + vllm_config = args[0] if args else kwargs["vllm_config"] super().__init__(*args, **kwargs) # Block scale for the preshuffled weight; None = not preshuffled. self._wqa_wkv_scale: torch.Tensor | None = None self._wo_b_scale: torch.Tensor | None = None + self._split_k_buffers: SparseDecodeSplitKBuffers | None = None + max_capture_size = ( + vllm_config.compilation_config.max_cudagraph_capture_size or 0 + ) + if _ON_GFX950 and not vllm_config.parallel_config.enable_dbo: + self._split_k_max_queries = min( + self.max_num_batched_tokens, max_capture_size + ) + heads_blocks = triton.cdiv(self.n_local_heads, 16) + self._split_k_max_partitions = ( + _max_decode_partitions(self._split_k_max_queries, heads_blocks, 32) + if self._split_k_max_queries > 0 + else 0 + ) + else: + self._split_k_max_queries = 0 + self._split_k_max_partitions = 0 + + def _get_split_k_buffers( + self, num_queries: int + ) -> SparseDecodeSplitKBuffers | None: + if num_queries > self._split_k_max_queries: + return None + if self._split_k_buffers is None: + partitions = self._split_k_max_partitions + heads = self.n_local_heads + part_m, part_l, part_acc = current_workspace_manager().get_simultaneous( + ((partitions, heads), torch.float32), + ((partitions, heads), torch.float32), + ((partitions, heads, self.head_dim), torch.float32), + ) + self._split_k_buffers = (part_m, part_l, part_acc) + return self._split_k_buffers @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: @@ -666,6 +703,7 @@ def _forward_decode( nope_head_dim=self.nope_head_dim, rope_head_dim=self.rope_head_dim, output=output, + split_k_buffers=self._get_split_k_buffers(q.shape[0]), ) def _forward_prefill( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index cf1d4d44ed2c..b1f3730216fb 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -24,6 +24,8 @@ _ON_GFX942 = False _ON_GFX950 = False +SparseDecodeSplitKBuffers = tuple[torch.Tensor, torch.Tensor, torch.Tensor] + @triton.jit def _indexer_k_quant_and_cache_kernel( @@ -1218,13 +1220,6 @@ def _sparse_attn_prefill_ragged_kernel( ) -@triton.jit -def _decode_e8m0_scales_triton(encoded_scales): - scale_bits = encoded_scales.to(tl.int32) << 23 - scale_bits = tl.where(encoded_scales == 0, 1 << 22, scale_bits) - return scale_bits.to(tl.float32, bitcast=True) - - @triton.jit def _sparse_attn_decode_ragged_kernel( q_ptr, @@ -1244,8 +1239,8 @@ def _sparse_attn_decode_ragged_kernel( extra_cache_stride0, main_num_rows, extra_num_rows, - MAIN_BLOCK_SIZE: tl.constexpr, - EXTRA_BLOCK_SIZE: tl.constexpr, + main_block_size, + extra_block_size, scale, num_heads, HAS_ATTN_SINK: tl.constexpr, @@ -1302,11 +1297,11 @@ def _sparse_attn_decode_ragged_kernel( valid = in_range & (slot >= 0) & (slot < main_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // MAIN_BLOCK_SIZE - pos_in_block = safe_slot % MAIN_BLOCK_SIZE + block_idx = safe_slot // main_block_size + pos_in_block = safe_slot % main_block_size cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 + token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 x_uint8 = tl.load( token_data_ptr[:, None] + nope_offsets[None, :], @@ -1322,7 +1317,7 @@ def _sparse_attn_decode_ragged_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1366,14 +1361,14 @@ def _sparse_attn_decode_ragged_kernel( valid = in_range & (slot >= 0) & (slot < extra_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // EXTRA_BLOCK_SIZE - pos_in_block = safe_slot % EXTRA_BLOCK_SIZE + block_idx = safe_slot // extra_block_size + pos_in_block = safe_slot % extra_block_size cache_block_ptr = ( extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 ) token_data_ptr = cache_block_ptr + pos_in_block * 576 token_scale_ptr = ( - cache_block_ptr + EXTRA_BLOCK_SIZE * 576 + pos_in_block * 8 + cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 ) x_uint8 = tl.load( @@ -1390,7 +1385,7 @@ def _sparse_attn_decode_ragged_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1484,8 +1479,8 @@ def _sparse_attn_decode_partial_kernel( pa_stride_h, main_num_rows, extra_num_rows, - MAIN_BLOCK_SIZE: tl.constexpr, - EXTRA_BLOCK_SIZE: tl.constexpr, + main_block_size, + extra_block_size, scale, num_heads, HAS_EXTRA: tl.constexpr, @@ -1552,11 +1547,11 @@ def _sparse_attn_decode_partial_kernel( valid = in_range & (slot >= 0) & (slot < main_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // MAIN_BLOCK_SIZE - pos_in_block = safe_slot % MAIN_BLOCK_SIZE + block_idx = safe_slot // main_block_size + pos_in_block = safe_slot % main_block_size cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 + token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 x_uint8 = tl.load( token_data_ptr[:, None] + nope_offsets[None, :], @@ -1572,7 +1567,7 @@ def _sparse_attn_decode_partial_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1619,14 +1614,14 @@ def _sparse_attn_decode_partial_kernel( valid = in_range & (slot >= 0) & (slot < extra_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // EXTRA_BLOCK_SIZE - pos_in_block = safe_slot % EXTRA_BLOCK_SIZE + block_idx = safe_slot // extra_block_size + pos_in_block = safe_slot % extra_block_size cache_block_ptr = ( extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 ) token_data_ptr = cache_block_ptr + pos_in_block * 576 token_scale_ptr = ( - cache_block_ptr + EXTRA_BLOCK_SIZE * 576 + pos_in_block * 8 + cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 ) x_uint8 = tl.load( @@ -1643,7 +1638,7 @@ def _sparse_attn_decode_partial_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1938,6 +1933,7 @@ def _decode_num_splits( avg_main_len: float = 0.0, avg_extra_len: float = 0.0, block_k: int = 32, + max_splits: int = 16, ) -> int: """Pick a flash-decode split count to keep the GPU busy across batch sizes. @@ -1978,9 +1974,7 @@ def _decode_num_splits( mu = 0.04 best_splits = 1 best_cost = None - # A full C128A row can contain 8K KV tokens. At low batch sizes, 32 splits - # keep those long rows to one device wave and halve the partial iterations. - for splits in range(1, 33): + for splits in range(1, max_splits + 1): waves = (base * splits + cu - 1) // cu cost = waves * (1.0 / splits + mu) if best_cost is None or cost < best_cost - 1e-9: @@ -2001,6 +1995,19 @@ def _decode_num_splits( return best_splits +@functools.cache +def _max_decode_partitions(max_queries: int, heads_blocks: int, max_splits: int) -> int: + return max( + queries + * _decode_num_splits( + queries, + heads_blocks, + max_splits=max_splits, + ) + for queries in range(1, max_queries + 1) + ) + + def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -2014,6 +2021,7 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indices: torch.Tensor | None = None, extra_indptr: torch.Tensor | None = None, out: torch.Tensor | None = None, + split_k_buffers: SparseDecodeSplitKBuffers | None = None, ) -> torch.Tensor: assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert main_cache.ndim == 3, ( @@ -2127,69 +2135,135 @@ def _rocm_sparse_attn_decode_ragged_triton( ) return out - block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950. + # Keep the split policy calibrated to the existing 32-token work estimate. + # The unified gfx950 Gluon body uses its fixed 64-token tile internally. + block_k = 32 # Average per-query segment lengths, read sync-free from the ragged index # sizes, let the split heuristic avoid over-splitting # main_indices/extra_indices are flat [nnz] int32. inv_q = 1.0 / max(1, num_queries) avg_main_len = main_indices.numel() * inv_q avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 - num_splits = _decode_num_splits( - num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k - ) + if _ON_GFX950: + num_splits = _decode_num_splits( + num_queries, + heads_blocks, + avg_main_len, + avg_extra_len, + block_k, + max_splits=32, + ) + else: + num_splits = _decode_num_splits( + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + ) - part_m = torch.empty( - (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device - ) - part_l = torch.empty_like(part_m) - part_acc = torch.empty( - (num_queries, num_splits, num_heads, comb_dim), - dtype=torch.float32, - device=q.device, - ) + if not _ON_GFX950: + split_k_buffers = None + required_partitions = num_queries * num_splits + if _ON_GFX950 and num_splits == 1: + part_m = part_l = part_acc = out + elif split_k_buffers is not None: + part_m, part_l, part_acc = split_k_buffers + if split_k_buffers is None and not (_ON_GFX950 and num_splits == 1): + part_shape = ( + (required_partitions, num_heads) + if _ON_GFX950 + else (num_queries, num_splits, num_heads) + ) + part_m = torch.empty(part_shape, dtype=torch.float32, device=q.device) + part_l = torch.empty_like(part_m) + acc_shape = ( + (required_partitions, num_heads, comb_dim) + if _ON_GFX950 + else (num_queries, num_splits, num_heads, comb_dim) + ) + part_acc = torch.empty( + acc_shape, + dtype=torch.float32, + device=q.device, + ) - _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - part_m, - part_l, - part_acc, - q.stride(0), - q.stride(1), - main_cache.stride(0), - extra_cache.stride(0), - part_m.stride(0), - part_m.stride(1), - part_acc.stride(0), - part_acc.stride(1), - part_acc.stride(2), - main_cache.shape[0] * main_cache.shape[1], - extra_cache.shape[0] * extra_cache.shape[1], - main_cache.shape[1], - extra_cache.shape[1], - scale, - num_heads, - HAS_EXTRA=has_extra, - NOPE_DIM=nope_head_dim, - NOPE_BLOCK=nope_block, - ROPE_DIM=rope_head_dim, - # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950). - # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere). - # Reading both with a single IS_FNUZ would decode one of them with the - # wrong FNUZ/OCP scale ratio (~1.87×). - IS_FNUZ_MAIN=is_fnuz, - IS_FNUZ_EXTRA=False, - BLOCK_H=block_h, - BLOCK_K=block_k, - NUM_SPLITS=num_splits, - NUM_STAGES=1, - num_warps=4, - ) + if _ON_GFX950: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse_gluon import ( + launch_sparse_attn_decode_partial_gfx950, + ) + + launch_sparse_attn_decode_partial_gfx950( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + attn_sink, + out, + part_m, + part_l, + part_acc, + scale, + num_heads, + has_extra, + has_attn_sink, + num_splits, + ) + else: + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + NOPE_BLOCK=nope_block, + ROPE_DIM=rope_head_dim, + IS_FNUZ_MAIN=is_fnuz, + IS_FNUZ_EXTRA=False, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + ) + + if _ON_GFX950 and num_splits == 1: + return out + + if _ON_GFX950: + pm_stride_s = part_m.stride(0) + pm_stride0 = num_splits * pm_stride_s + pa_stride_s = part_acc.stride(0) + pa_stride0 = num_splits * pa_stride_s + pa_stride_h = part_acc.stride(1) + else: + pm_stride0, pm_stride_s = part_m.stride(0), part_m.stride(1) + pa_stride0, pa_stride_s, pa_stride_h = ( + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + ) _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( part_m, @@ -2199,11 +2273,11 @@ def _rocm_sparse_attn_decode_ragged_triton( out, out.stride(0), out.stride(1), - part_m.stride(0), - part_m.stride(1), - part_acc.stride(0), - part_acc.stride(1), - part_acc.stride(2), + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, num_heads, HAS_ATTN_SINK=has_attn_sink, COMB_DIM=comb_dim, @@ -2232,6 +2306,7 @@ def _rocm_sparse_attn_decode_triton( extra_ragged_indices: torch.Tensor | None = None, extra_ragged_indptr: torch.Tensor | None = None, out: torch.Tensor | None = None, + split_k_buffers: SparseDecodeSplitKBuffers | None = None, ) -> torch.Tensor: if main_ragged_indices is None or main_ragged_indptr is None: main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense( @@ -2268,6 +2343,7 @@ def _rocm_sparse_attn_decode_triton( extra_indices=extra_ragged_indices, extra_indptr=extra_ragged_indptr, out=out, + split_k_buffers=split_k_buffers, ) @@ -2339,6 +2415,7 @@ def rocm_sparse_attn_decode( nope_head_dim: int, rope_head_dim: int, output: torch.Tensor, + split_k_buffers: SparseDecodeSplitKBuffers | None = None, ) -> None: assert swa_k_cache.dtype == torch.uint8, ( "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, " @@ -2368,7 +2445,7 @@ def rocm_sparse_attn_decode( if topk_indices is not None: extra_indices = topk_indices.reshape(topk_indices.shape[0], -1) - direct_out = output if output.dtype == torch.bfloat16 else None + direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None attn_out = _rocm_sparse_attn_decode_triton( q=q, main_cache=swa_k_cache, @@ -2386,6 +2463,7 @@ def rocm_sparse_attn_decode( extra_ragged_indices=topk_ragged_indices, extra_ragged_indptr=topk_ragged_indptr, out=direct_out, + split_k_buffers=split_k_buffers, ) if direct_out is None: output.copy_(attn_out.to(output.dtype)) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse_gluon.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse_gluon.py new file mode 100644 index 000000000000..816e53744e7d --- /dev/null +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse_gluon.py @@ -0,0 +1,795 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""gfx950 Gluon partial kernel for DeepSeek-V4 sparse-MLA decode. + +Adapted from AITER's MIT-licensed ``pa_decode_sparse`` gfx950 kernel. The vLLM +variant is intentionally limited to the packed two-segment DSv4 cache contract +and writes the existing Triton reducer's partial-buffer format. +""" + +import torch + +from vllm.triton_utils import gl, gluon + +_GFX950_BLOCK_H = 16 +_GFX950_BLOCK_K = 64 +_GFX950_HEAD_SIZE = 512 +_MAX_BUFFER_OFFSET = 2**31 - 1 + + +def _max_addressable_bytes(tensor: torch.Tensor) -> int: + """Return the byte span reachable from a tensor's data pointer.""" + span = 1 + for size, stride in zip(tensor.shape, tensor.stride()): + if size > 1: + span += (size - 1) * abs(stride) + return span * tensor.element_size() + + +@gluon.jit +def _cache_load(ptr, offsets, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): + if USE_BUFFER_LOAD: + return gl.amd.cdna4.buffer_load( + ptr=ptr, + offsets=offsets.to(gl.int32), + mask=mask, + other=other, + cache=".cg", + ) + return gl.load( + ptr + offsets.to(gl.int64), + mask=mask, + other=other, + cache_modifier=".cg", + ) + + +@gluon.jit +def _decode_e8m0_scales(encoded_scales): + scale_bits = encoded_scales.to(gl.int32) << 23 + scale_bits = gl.where(encoded_scales == 0, 1 << 22, scale_bits) + return scale_bits.to(gl.float32, bitcast=True) + + +@gluon.jit +def _slots( + indices_ptr, + segment_start, + k_pos, + segment_hi, + num_rows, + BLOCK_SIZE: gl.constexpr, + MASKED: gl.constexpr, +): + if MASKED: + in_range = k_pos < segment_hi + slot = gl.load(indices_ptr + segment_start + k_pos, mask=in_range, other=-1) + valid = in_range & (slot >= 0) & (slot < num_rows) + else: + slot = gl.load(indices_ptr + segment_start + k_pos) + valid = slot >= 0 + safe_slot = gl.where(valid, slot, 0) + return ( + (safe_slot // BLOCK_SIZE).to(gl.int32), + (safe_slot % BLOCK_SIZE).to(gl.int32), + valid, + ) + + +@gluon.jit +def _gather_fp8_tile( + cache_ptr, + cache_bf16_ptr, + indices_ptr, + segment_start, + k_start, + segment_hi, + cache_stride0, + num_rows, + full_offsets, + rope_offsets, + slot_offsets, + rope_slot_offsets, + gather_layout: gl.constexpr, + rope_gather_layout: gl.constexpr, + BLOCK_SIZE: gl.constexpr, + BLOCK_K: gl.constexpr, + MASKED: gl.constexpr, + USE_BUFFER_LOAD: gl.constexpr, +): + if not USE_BUFFER_LOAD: + cache_stride0 = cache_stride0.to(gl.int64) + + block, pos, valid = _slots( + indices_ptr, + segment_start, + k_start + slot_offsets, + segment_hi, + num_rows, + BLOCK_SIZE, + MASKED, + ) + block_g = gl.convert_layout(block, gl.SliceLayout(1, gather_layout)) + pos_g = gl.convert_layout(pos, gl.SliceLayout(1, gather_layout)) + valid_g = gl.convert_layout(valid, gl.SliceLayout(1, gather_layout)) + + data_offsets = (block_g * cache_stride0 + pos_g * 576)[:, None] + full_offsets[ + None, : + ] + scale_offsets = (block_g * cache_stride0 + BLOCK_SIZE * 576 + pos_g * 8)[ + :, None + ] + (full_offsets[None, :] // 64) + if MASKED: + data = _cache_load( + cache_ptr, + data_offsets, + USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0, + ) + encoded_scales = _cache_load( + cache_ptr, + scale_offsets, + USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=127, + ) + else: + data = _cache_load(cache_ptr, data_offsets, USE_BUFFER_LOAD) + encoded_scales = _cache_load(cache_ptr, scale_offsets, USE_BUFFER_LOAD) + nope = ( + data.to(gl.float8e4nv, bitcast=True).to(gl.float32) + * _decode_e8m0_scales(encoded_scales) + ).to(gl.bfloat16) + + rope_block, rope_pos, rope_valid = _slots( + indices_ptr, + segment_start, + k_start + rope_slot_offsets, + segment_hi, + num_rows, + BLOCK_SIZE, + MASKED, + ) + rope_offsets_global = (rope_block * (cache_stride0 // 2) + rope_pos * 288 + 224)[ + :, None + ] + rope_offsets[None, :] + if MASKED: + rope = _cache_load( + cache_bf16_ptr, + rope_offsets_global, + USE_BUFFER_LOAD, + mask=rope_valid[:, None], + other=0.0, + ) + else: + rope = _cache_load( + cache_bf16_ptr, + rope_offsets_global, + USE_BUFFER_LOAD, + ) + return nope, rope, valid + + +@gluon.jit +def _consume_fp8_tile( + nope, + rope, + valid, + q_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout: gl.constexpr, + pv_layout: gl.constexpr, + k_layout: gl.constexpr, + v_layout: gl.constexpr, + p_layout: gl.constexpr, + BLOCK_H: gl.constexpr, + BLOCK_K: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, +): + kv_smem.store(nope) + kv_smem.slice(448, 64, dim=1).store(rope) + + k = kv_smem.permute([1, 0]).load(k_layout) + scores = gl.amd.cdna4.mfma( + q_dot, + k, + gl.zeros([BLOCK_H, BLOCK_K], gl.float32, layout=qk_layout), + ) + column_mask = gl.convert_layout(valid, gl.SliceLayout(0, qk_layout))[None, :] + if not HEAD_ALIGNED: + column_mask = ( + gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[:, None] + & column_mask + ) + scores = gl.where(column_mask, scores, float("-inf")) + + block_max = gl.max(scores, axis=1) + m_new = gl.maximum(m_i, block_max) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + m_new_scaled = m_new * qk_scale + p = gl.exp2(scores * qk_scale - m_new_scaled[:, None]) + alpha = gl.exp2(m_i * qk_scale - m_new_scaled) + l_new = l_i * alpha + gl.sum(p, axis=1) + + v = kv_smem.load(v_layout) + p_dot = gl.convert_layout(p.to(gl.bfloat16), p_layout) + alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, pv_layout)) + acc = acc * alpha_pv[:, None] + acc = gl.amd.cdna4.mfma(p_dot, v, acc) + return m_new, l_new, acc + + +@gluon.jit +def _process_segment( + q_dot, + cache_ptr, + cache_bf16_ptr, + indices_ptr, + segment_start, + lo, + hi, + cache_stride0, + num_rows, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout: gl.constexpr, + pv_layout: gl.constexpr, + k_layout: gl.constexpr, + v_layout: gl.constexpr, + p_layout: gl.constexpr, + gather_layout: gl.constexpr, + rope_gather_layout: gl.constexpr, + slot_layout: gl.constexpr, + BLOCK_SIZE: gl.constexpr, + BLOCK_H: gl.constexpr, + BLOCK_K: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, + USE_BUFFER_LOAD: gl.constexpr, +): + full_offsets = gl.arange( + 0, + 512, + layout=gl.SliceLayout(0, gather_layout), + ) + rope_offsets = gl.arange( + 0, + 64, + layout=gl.SliceLayout(0, rope_gather_layout), + ) + slot_offsets = gl.arange(0, BLOCK_K, layout=slot_layout) + rope_slot_offsets = gl.arange( + 0, + BLOCK_K, + layout=gl.SliceLayout(1, rope_gather_layout), + ) + + full_hi = lo + ((hi - lo) // BLOCK_K) * BLOCK_K + num_full = (full_hi - lo) // BLOCK_K + if num_full > 0: + nope, rope, valid = _gather_fp8_tile( + cache_ptr, + cache_bf16_ptr, + indices_ptr, + segment_start, + lo, + hi, + cache_stride0, + num_rows, + full_offsets, + rope_offsets, + slot_offsets, + rope_slot_offsets, + gather_layout, + rope_gather_layout, + BLOCK_SIZE, + BLOCK_K, + False, + USE_BUFFER_LOAD, + ) + for tile_idx in range(1, num_full): + next_nope, next_rope, next_valid = _gather_fp8_tile( + cache_ptr, + cache_bf16_ptr, + indices_ptr, + segment_start, + lo + tile_idx * BLOCK_K, + hi, + cache_stride0, + num_rows, + full_offsets, + rope_offsets, + slot_offsets, + rope_slot_offsets, + gather_layout, + rope_gather_layout, + BLOCK_SIZE, + BLOCK_K, + False, + USE_BUFFER_LOAD, + ) + m_i, l_i, acc = _consume_fp8_tile( + nope, + rope, + valid, + q_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + BLOCK_H, + BLOCK_K, + HEAD_ALIGNED, + ) + nope, rope, valid = next_nope, next_rope, next_valid + m_i, l_i, acc = _consume_fp8_tile( + nope, + rope, + valid, + q_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + BLOCK_H, + BLOCK_K, + HEAD_ALIGNED, + ) + + if full_hi < hi: + nope, rope, valid = _gather_fp8_tile( + cache_ptr, + cache_bf16_ptr, + indices_ptr, + segment_start, + full_hi, + hi, + cache_stride0, + num_rows, + full_offsets, + rope_offsets, + slot_offsets, + rope_slot_offsets, + gather_layout, + rope_gather_layout, + BLOCK_SIZE, + BLOCK_K, + True, + USE_BUFFER_LOAD, + ) + m_i, l_i, acc = _consume_fp8_tile( + nope, + rope, + valid, + q_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + BLOCK_H, + BLOCK_K, + HEAD_ALIGNED, + ) + return m_i, l_i, acc + + +@gluon.jit +def _sparse_attn_decode_partial_gfx950_kernel( + q_ptr, + main_cache_ptr, + main_cache_bf16_ptr, + main_indices_ptr, + main_indptr_ptr, + extra_cache_ptr, + extra_cache_bf16_ptr, + extra_indices_ptr, + extra_indptr_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + scale: gl.constexpr, + q_stride0: gl.constexpr, + q_stride1: gl.constexpr, + main_cache_stride0, + extra_cache_stride0, + pm_stride0: gl.constexpr, + pm_stride_s: gl.constexpr, + pa_stride0: gl.constexpr, + pa_stride_s: gl.constexpr, + pa_stride_h: gl.constexpr, + main_num_rows, + extra_num_rows, + num_heads: gl.constexpr, + HAS_EXTRA: gl.constexpr, + HAS_ATTN_SINK: gl.constexpr, + MAIN_BLOCK_SIZE: gl.constexpr, + EXTRA_BLOCK_SIZE: gl.constexpr, + NUM_SPLITS: gl.constexpr, + HEAD_ALIGNED: gl.constexpr, + USE_BUFFER_LOAD: gl.constexpr, +): + num_warps: gl.constexpr = gl.num_warps() + query_idx = gl.program_id(0) + split_id = gl.program_id(1) + head_block = gl.program_id(2) + + qk_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[1, num_warps], + ) + pv_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[1, num_warps], + ) + q_layout: gl.constexpr = gl.DotOperandLayout(0, qk_layout, 8) + k_layout: gl.constexpr = gl.DotOperandLayout(1, qk_layout, 8) + p_layout: gl.constexpr = gl.DotOperandLayout(0, pv_layout, 8) + v_layout: gl.constexpr = gl.DotOperandLayout(1, pv_layout, 8) + + gather_layout: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 16], + threads_per_warp=[8, 8], + warps_per_cta=[1, num_warps], + order=[1, 0], + ) + rope_gather_layout: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[num_warps, 1], + order=[1, 0], + ) + slot_layout: gl.constexpr = gl.SliceLayout(1, gather_layout) + q_blocked_layout: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[1, num_warps], + order=[1, 0], + ) + kv_shared_layout: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 8]], + [64, 512], + [1, 0], + ) + + head_base = head_block * 16 + head_offsets_q = gl.arange( + 0, + 16, + layout=gl.SliceLayout(1, q_blocked_layout), + ) + dim_offsets_q = gl.arange( + 0, + 512, + layout=gl.SliceLayout(0, q_blocked_layout), + ) + heads_q = head_base + head_offsets_q + head_mask_q = heads_q < num_heads + q_offsets = ( + query_idx * q_stride0 + heads_q[:, None] * q_stride1 + dim_offsets_q[None, :] + ).to(gl.int32) + q = gl.amd.cdna4.buffer_load( + ptr=q_ptr, + offsets=q_offsets, + mask=head_mask_q[:, None], + other=0.0, + ) + q_dot = gl.convert_layout(q, q_layout) + + head_offsets_pv = gl.arange( + 0, + 16, + layout=gl.SliceLayout(1, pv_layout), + ) + heads_pv = head_base + head_offsets_pv + head_mask_pv = heads_pv < num_heads + + m_i = gl.full( + [16], + float("-inf"), + gl.float32, + layout=gl.SliceLayout(1, qk_layout), + ) + l_i = gl.zeros( + [16], + gl.float32, + layout=gl.SliceLayout(1, qk_layout), + ) + acc = gl.zeros( + [16, 512], + gl.float32, + layout=pv_layout, + ) + kv_smem = gl.allocate_shared_memory( + gl.bfloat16, + [64, 512], + kv_shared_layout, + ) + + main_start = gl.load(main_indptr_ptr + query_idx) + main_end = gl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS + main_lo = split_id * main_chunk + main_hi = gl.minimum(main_lo + main_chunk, main_len) + + rcp_ln2: gl.constexpr = 1.4426950408889634 + qk_scale: gl.constexpr = scale * rcp_ln2 + m_i, l_i, acc = _process_segment( + q_dot, + main_cache_ptr, + main_cache_bf16_ptr, + main_indices_ptr, + main_start, + main_lo, + main_hi, + main_cache_stride0, + main_num_rows, + m_i, + l_i, + acc, + head_mask_pv, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + gather_layout, + rope_gather_layout, + slot_layout, + MAIN_BLOCK_SIZE, + 16, + 64, + HEAD_ALIGNED, + USE_BUFFER_LOAD, + ) + if HAS_EXTRA: + extra_start = gl.load(extra_indptr_ptr + query_idx) + extra_end = gl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS + extra_lo = split_id * extra_chunk + extra_hi = gl.minimum(extra_lo + extra_chunk, extra_len) + m_i, l_i, acc = _process_segment( + q_dot, + extra_cache_ptr, + extra_cache_bf16_ptr, + extra_indices_ptr, + extra_start, + extra_lo, + extra_hi, + extra_cache_stride0, + extra_num_rows, + m_i, + l_i, + acc, + head_mask_pv, + qk_scale, + kv_smem, + qk_layout, + pv_layout, + k_layout, + v_layout, + p_layout, + gather_layout, + rope_gather_layout, + slot_layout, + EXTRA_BLOCK_SIZE, + 16, + 64, + HEAD_ALIGNED, + USE_BUFFER_LOAD, + ) + + m_pv = gl.convert_layout(m_i, gl.SliceLayout(1, pv_layout)) + l_pv = gl.convert_layout(l_i, gl.SliceLayout(1, pv_layout)) + if NUM_SPLITS == 1: + if HAS_ATTN_SINK: + m_scaled = m_pv * scale + sink = gl.amd.cdna4.buffer_load( + ptr=part_m_ptr, + offsets=heads_pv.to(gl.int32), + mask=head_mask_pv, + other=float("-inf"), + ).to(gl.float32) + m_final = gl.maximum(m_scaled, sink) + alpha = gl.exp2((m_scaled - m_final) * rcp_ln2) + l_final = l_pv * alpha + gl.exp2((sink - m_final) * rcp_ln2) + acc = acc * alpha[:, None] + else: + l_final = l_pv + denom = gl.maximum(l_final, 1.0e-30) + out = gl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0) + dim_offsets_out = gl.arange( + 0, + 512, + layout=gl.SliceLayout(0, pv_layout), + ) + out_offsets = ( + query_idx * pa_stride0 + + heads_pv[:, None] * pa_stride_h + + dim_offsets_out[None, :] + ).to(gl.int32) + gl.amd.cdna4.buffer_store( + out.to(part_acc_ptr.dtype.element_ty), + ptr=part_acc_ptr, + offsets=out_offsets, + mask=head_mask_pv[:, None], + ) + else: + m_natural = gl.where( + m_pv > float("-inf"), + m_pv * scale, + -3.4028234663852886e38, + ) + partial_row_base = query_idx * pm_stride0 + split_id * pm_stride_s + gl.amd.cdna4.buffer_store( + m_natural, + ptr=part_m_ptr + partial_row_base, + offsets=heads_pv.to(gl.int32), + mask=head_mask_pv, + ) + gl.amd.cdna4.buffer_store( + l_pv, + ptr=part_l_ptr + partial_row_base, + offsets=heads_pv.to(gl.int32), + mask=head_mask_pv, + ) + dim_offsets_acc = gl.arange( + 0, + 512, + layout=gl.SliceLayout(0, pv_layout), + ) + partial_acc_base = query_idx * pa_stride0 + split_id * pa_stride_s + partial_acc_offsets = ( + partial_acc_base + + heads_pv[:, None] * pa_stride_h + + dim_offsets_acc[None, :] + ).to(gl.int32) + gl.amd.cdna4.buffer_store( + acc, + ptr=part_acc_ptr, + offsets=partial_acc_offsets, + mask=head_mask_pv[:, None], + ) + + +def launch_sparse_attn_decode_partial_gfx950( + q: torch.Tensor, + main_cache: torch.Tensor, + main_indices: torch.Tensor, + main_indptr: torch.Tensor, + extra_cache: torch.Tensor, + extra_indices: torch.Tensor, + extra_indptr: torch.Tensor, + attn_sink: torch.Tensor, + out: torch.Tensor, + part_m: torch.Tensor, + part_l: torch.Tensor, + part_acc: torch.Tensor, + scale: float, + num_heads: int, + has_extra: bool, + has_attn_sink: bool, + num_splits: int, +) -> None: + """Launch the fixed gfx950 packed-cache partial kernel.""" + assert q.shape[-1] == _GFX950_HEAD_SIZE + assert main_cache.dtype == torch.uint8 + assert extra_cache.dtype == torch.uint8 + + use_buffer_load = ( + max( + _max_addressable_bytes(main_cache), + _max_addressable_bytes(extra_cache), + ) + < _MAX_BUFFER_OFFSET + ) + heads_blocks = (num_heads + _GFX950_BLOCK_H - 1) // _GFX950_BLOCK_H + grid = (q.shape[0], num_splits, heads_blocks) + if num_splits == 1: + part_m = attn_sink + part_l = out + part_acc = out + pm_stride0 = pm_stride_s = 0 + pa_stride0 = out.stride(0) + pa_stride_s = 0 + pa_stride_h = out.stride(1) + else: + pm_stride_s = part_m.stride(0) + pm_stride0 = num_splits * pm_stride_s + pa_stride_s = part_acc.stride(0) + pa_stride0 = num_splits * pa_stride_s + pa_stride_h = part_acc.stride(1) + _sparse_attn_decode_partial_gfx950_kernel[grid]( + q, + main_cache, + main_cache.view(torch.bfloat16), + main_indices, + main_indptr, + extra_cache, + extra_cache.view(torch.bfloat16), + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + scale, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + num_heads, + HAS_EXTRA=has_extra, + HAS_ATTN_SINK=has_attn_sink, + MAIN_BLOCK_SIZE=main_cache.shape[1], + EXTRA_BLOCK_SIZE=extra_cache.shape[1], + NUM_SPLITS=num_splits, + HEAD_ALIGNED=num_heads % _GFX950_BLOCK_H == 0, + USE_BUFFER_LOAD=use_buffer_load, + num_warps=4, + waves_per_eu=0, + ) From 656501b7a66f0bdb0a801b8141dfa31791fc39f0 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Fri, 14 Aug 2026 06:15:18 +0000 Subject: [PATCH 3/8] Revert "[ROCm][Perf] Add unified gfx950 sparse MLA decode kernel" This reverts commit 38ef98046c5d976259a18f0e931a0ed24fe03884. Signed-off-by: Fangzhou Ai --- .../attention/test_rocm_triton_attn_dsv4.py | 176 +--- vllm/models/deepseek_v4/amd/rocm.py | 38 - .../v1/attention/ops/rocm_aiter_mla_sparse.py | 260 ++---- .../ops/rocm_aiter_mla_sparse_gluon.py | 795 ------------------ 4 files changed, 98 insertions(+), 1171 deletions(-) delete mode 100644 vllm/v1/attention/ops/rocm_aiter_mla_sparse_gluon.py diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 40e4479f6cd0..775f1f5d1447 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -24,50 +24,18 @@ def _on_split_decode_arch() -> bool: return False -def _on_gfx950() -> bool: - if not current_platform.is_rocm(): - return False - try: - from vllm.platforms.rocm import _ON_GFX950 - - return bool(_ON_GFX950) - except Exception: - return False - - # The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other # architectures take the fallback decode kernel, so its tests are skipped there. requires_split_decode_arch = pytest.mark.skipif( not _on_split_decode_arch(), reason="split-K decode kernel is only tuned for AMD gfx942/gfx950", ) -requires_gfx950 = pytest.mark.skipif( - not _on_gfx950(), reason="Gluon sparse decode kernel is only used on gfx950" -) NOPE_HEAD_DIM = 448 ROPE_HEAD_DIM = 64 HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM -def _make_split_k_buffers( - num_queries: int, - num_splits: int, - num_heads: int, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - partitions = num_queries * num_splits - return ( - torch.empty((partitions, num_heads), dtype=torch.float32, device=device), - torch.empty((partitions, num_heads), dtype=torch.float32, device=device), - torch.empty( - (partitions, num_heads, HEAD_DIM), - dtype=torch.float32, - device=device, - ), - ) - - def _ref_global_topk_ragged( topk_indices: torch.Tensor, token_to_req_indices: torch.Tensor, @@ -457,66 +425,6 @@ def test_sparse_attn_decode_ragged_kernel() -> None: torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) -@requires_gfx950 -@torch.inference_mode() -def test_sparse_attn_decode_ue8m0_scale_groups() -> None: - """Each packed UE8M0 byte scales exactly one 64-column NoPE group.""" - from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod - - device = torch.device("cuda") - block_size = 1 - cache = torch.zeros( - (1, block_size, 584), - dtype=torch.uint8, - device=device, - ) - cache_flat = cache.flatten() - one_bits = ( - torch.ones(NOPE_HEAD_DIM, dtype=torch.float32, device=device) - .to(torch.float8_e4m3fn) - .view(torch.uint8) - ) - cache_flat[:NOPE_HEAD_DIM].copy_(one_bits) - encoded_scales = torch.tensor( - [0, 125, 126, 127, 128, 129, 130], - dtype=torch.uint8, - device=device, - ) - cache_flat[576:583].copy_(encoded_scales) - cache_flat[583] = 0 - - q = torch.zeros((1, 1, HEAD_DIM), dtype=torch.bfloat16, device=device) - empty_indices = torch.empty(0, dtype=torch.int32, device=device) - empty_indptr = torch.zeros(2, dtype=torch.int32, device=device) - extra_indices = torch.zeros(1, dtype=torch.int32, device=device) - extra_indptr = torch.tensor([0, 1], dtype=torch.int32, device=device) - - actual = mod._rocm_sparse_attn_decode_ragged_triton( - q=q, - main_cache=cache, - main_indices=empty_indices, - main_indptr=empty_indptr, - scale=HEAD_DIM**-0.5, - attn_sink=None, - nope_head_dim=NOPE_HEAD_DIM, - rope_head_dim=ROPE_HEAD_DIM, - extra_cache=cache, - extra_indices=extra_indices, - extra_indptr=extra_indptr, - ) - - expected_bits = encoded_scales.to(torch.int32) << 23 - expected_bits[0] = 1 << 22 - expected_scales = expected_bits.view(torch.float32) - expected = torch.cat( - [ - expected_scales.repeat_interleave(64), - torch.zeros(ROPE_HEAD_DIM, dtype=torch.float32, device=device), - ] - ).to(torch.bfloat16) - torch.testing.assert_close(actual[0, 0], expected, atol=1e-3, rtol=1e-3) - - @requires_split_decode_arch @torch.inference_mode() def test_decode_num_splits_heuristic(monkeypatch) -> None: @@ -531,94 +439,28 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: # A tiny batch on a large device should split to add parallelism. assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1 - # Preserve the standard-kernel limit while allowing gfx950's long-context - # path to opt into 32 splits. - assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16 - assert mod._decode_num_splits(1, 1, 128.0, 8192.0, max_splits=32) == 32 - assert mod._decode_num_splits(8, 1, 128.0, 8192.0, max_splits=32) == 32 - assert mod._decode_num_splits(64, 1, 128.0, 8192.0, max_splits=32) == 4 - max_queries = 512 - max_partitions = mod._max_decode_partitions(max_queries, 1, 32) - assert all( - queries - * mod._decode_num_splits( - queries, - 1, - avg_main_len=128.0, - avg_extra_len=8192.0, - max_splits=32, - ) - <= max_partitions - for queries in range(1, max_queries + 1) - ) + # Long C128A rows need 32 splits to fill a 256-CU device at low batch. + assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 32 + assert mod._decode_num_splits(8, 1, 128.0, 8192.0) == 32 + assert mod._decode_num_splits(64, 1, 128.0, 8192.0) == 4 - # The default count always stays within the standard [1, 16] range, and a + # The chosen count always stays within the searched [1, 32] range, and a # zero-length workload never splits (no work to parallelize). for num_queries in (1, 4, 24, 224, 1024): splits = mod._decode_num_splits( num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0 ) - assert 1 <= splits <= 16 + assert 1 <= splits <= 32 assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 -@requires_split_decode_arch -@torch.inference_mode() -def test_sparse_attn_decode_rejects_invalid_indices(monkeypatch) -> None: - """Full-tile sentinels and tail OOB slots are ignored.""" - from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod - - device = torch.device("cuda") - torch.manual_seed(17) - block_size = 4 - num_kv = 8 - num_heads = 3 - q = torch.randn(1, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) - kv = torch.randn(num_kv, HEAD_DIM, dtype=torch.bfloat16, device=device) - cache = _pack_fp8_ds_mla_cache( - kv, - block_size, - use_fnuz=current_platform.is_fp8_fnuz(), - ) - row = [i % num_kv for i in range(65)] - row[5] = -1 - row[64] = num_kv - indices, indptr = _ragged_from_rows([row], device) - - monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: 1) - actual = mod._rocm_sparse_attn_decode_ragged_triton( - q=q, - main_cache=cache, - main_indices=indices, - main_indptr=indptr, - scale=HEAD_DIM**-0.5, - attn_sink=None, - nope_head_dim=NOPE_HEAD_DIM, - rope_head_dim=ROPE_HEAD_DIM, - ) - expected = _ref_sparse_decode_ragged( - q=q, - main_cache=cache, - main_rows=[[slot for slot in row if 0 <= slot < num_kv]], - scale=HEAD_DIM**-0.5, - attn_sink=None, - block_size=block_size, - main_use_fnuz=current_platform.is_fp8_fnuz(), - ) - - torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) - - @requires_split_decode_arch @pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8, 32]) @pytest.mark.parametrize("with_extra", [True, False]) @pytest.mark.parametrize("with_sink", [True, False]) @torch.inference_mode() def test_sparse_attn_decode_split_k_kernel( - monkeypatch, - num_splits: int, - with_extra: bool, - with_sink: bool, + monkeypatch, num_splits: int, with_extra: bool, with_sink: bool ) -> None: """Flash-decode split-K decode path (partial + reduce kernels). @@ -668,7 +510,6 @@ def test_sparse_attn_decode_split_k_kernel( # Pin the split count so each parametrized value is exercised deterministically. monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits) - split_k_buffers = _make_split_k_buffers(num_queries, num_splits, num_heads, device) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, @@ -682,7 +523,6 @@ def test_sparse_attn_decode_split_k_kernel( extra_cache=extra_cache, extra_indices=extra_indices, extra_indptr=extra_indptr, - split_k_buffers=split_k_buffers, ) expected = _ref_sparse_decode_ragged( q=q, @@ -728,7 +568,6 @@ def test_sparse_attn_decode_long_context(monkeypatch) -> None: scale = HEAD_DIM**-0.5 monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: 32) - split_k_buffers = _make_split_k_buffers(1, 32, num_heads, device) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, main_cache=main_cache, @@ -741,7 +580,6 @@ def test_sparse_attn_decode_long_context(monkeypatch) -> None: extra_cache=extra_cache, extra_indices=extra_indices, extra_indptr=extra_indptr, - split_k_buffers=split_k_buffers, ) kv = _read_fp8_ds_mla_cache_rows( diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 4410d36b3f2f..23223acd2dba 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -19,7 +19,6 @@ DeepseekV4SparseMLAMetadataBuilder, ) from vllm.platforms import current_platform -from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( CommonAttentionMetadata, @@ -29,8 +28,6 @@ DeepseekSparseSWAMetadataBuilder, ) from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - SparseDecodeSplitKBuffers, - _max_decode_partitions, build_ragged_indices_from_dense, rocm_inv_rope_einsum, rocm_sparse_attn_decode, @@ -454,44 +451,10 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): backend_cls = DeepseekV4ROCMAiterMLASparseBackend def __init__(self, *args, **kwargs): - vllm_config = args[0] if args else kwargs["vllm_config"] super().__init__(*args, **kwargs) # Block scale for the preshuffled weight; None = not preshuffled. self._wqa_wkv_scale: torch.Tensor | None = None self._wo_b_scale: torch.Tensor | None = None - self._split_k_buffers: SparseDecodeSplitKBuffers | None = None - max_capture_size = ( - vllm_config.compilation_config.max_cudagraph_capture_size or 0 - ) - if _ON_GFX950 and not vllm_config.parallel_config.enable_dbo: - self._split_k_max_queries = min( - self.max_num_batched_tokens, max_capture_size - ) - heads_blocks = triton.cdiv(self.n_local_heads, 16) - self._split_k_max_partitions = ( - _max_decode_partitions(self._split_k_max_queries, heads_blocks, 32) - if self._split_k_max_queries > 0 - else 0 - ) - else: - self._split_k_max_queries = 0 - self._split_k_max_partitions = 0 - - def _get_split_k_buffers( - self, num_queries: int - ) -> SparseDecodeSplitKBuffers | None: - if num_queries > self._split_k_max_queries: - return None - if self._split_k_buffers is None: - partitions = self._split_k_max_partitions - heads = self.n_local_heads - part_m, part_l, part_acc = current_workspace_manager().get_simultaneous( - ((partitions, heads), torch.float32), - ((partitions, heads), torch.float32), - ((partitions, heads, self.head_dim), torch.float32), - ) - self._split_k_buffers = (part_m, part_l, part_acc) - return self._split_k_buffers @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: @@ -703,7 +666,6 @@ def _forward_decode( nope_head_dim=self.nope_head_dim, rope_head_dim=self.rope_head_dim, output=output, - split_k_buffers=self._get_split_k_buffers(q.shape[0]), ) def _forward_prefill( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index b1f3730216fb..cf1d4d44ed2c 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -24,8 +24,6 @@ _ON_GFX942 = False _ON_GFX950 = False -SparseDecodeSplitKBuffers = tuple[torch.Tensor, torch.Tensor, torch.Tensor] - @triton.jit def _indexer_k_quant_and_cache_kernel( @@ -1220,6 +1218,13 @@ def _sparse_attn_prefill_ragged_kernel( ) +@triton.jit +def _decode_e8m0_scales_triton(encoded_scales): + scale_bits = encoded_scales.to(tl.int32) << 23 + scale_bits = tl.where(encoded_scales == 0, 1 << 22, scale_bits) + return scale_bits.to(tl.float32, bitcast=True) + + @triton.jit def _sparse_attn_decode_ragged_kernel( q_ptr, @@ -1239,8 +1244,8 @@ def _sparse_attn_decode_ragged_kernel( extra_cache_stride0, main_num_rows, extra_num_rows, - main_block_size, - extra_block_size, + MAIN_BLOCK_SIZE: tl.constexpr, + EXTRA_BLOCK_SIZE: tl.constexpr, scale, num_heads, HAS_ATTN_SINK: tl.constexpr, @@ -1297,11 +1302,11 @@ def _sparse_attn_decode_ragged_kernel( valid = in_range & (slot >= 0) & (slot < main_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // main_block_size - pos_in_block = safe_slot % main_block_size + block_idx = safe_slot // MAIN_BLOCK_SIZE + pos_in_block = safe_slot % MAIN_BLOCK_SIZE cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 + token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 x_uint8 = tl.load( token_data_ptr[:, None] + nope_offsets[None, :], @@ -1317,7 +1322,7 @@ def _sparse_attn_decode_ragged_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + scales = _decode_e8m0_scales_triton(encoded_scales) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1361,14 +1366,14 @@ def _sparse_attn_decode_ragged_kernel( valid = in_range & (slot >= 0) & (slot < extra_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // extra_block_size - pos_in_block = safe_slot % extra_block_size + block_idx = safe_slot // EXTRA_BLOCK_SIZE + pos_in_block = safe_slot % EXTRA_BLOCK_SIZE cache_block_ptr = ( extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 ) token_data_ptr = cache_block_ptr + pos_in_block * 576 token_scale_ptr = ( - cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 + cache_block_ptr + EXTRA_BLOCK_SIZE * 576 + pos_in_block * 8 ) x_uint8 = tl.load( @@ -1385,7 +1390,7 @@ def _sparse_attn_decode_ragged_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + scales = _decode_e8m0_scales_triton(encoded_scales) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1479,8 +1484,8 @@ def _sparse_attn_decode_partial_kernel( pa_stride_h, main_num_rows, extra_num_rows, - main_block_size, - extra_block_size, + MAIN_BLOCK_SIZE: tl.constexpr, + EXTRA_BLOCK_SIZE: tl.constexpr, scale, num_heads, HAS_EXTRA: tl.constexpr, @@ -1547,11 +1552,11 @@ def _sparse_attn_decode_partial_kernel( valid = in_range & (slot >= 0) & (slot < main_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // main_block_size - pos_in_block = safe_slot % main_block_size + block_idx = safe_slot // MAIN_BLOCK_SIZE + pos_in_block = safe_slot % MAIN_BLOCK_SIZE cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 + token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 x_uint8 = tl.load( token_data_ptr[:, None] + nope_offsets[None, :], @@ -1567,7 +1572,7 @@ def _sparse_attn_decode_partial_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + scales = _decode_e8m0_scales_triton(encoded_scales) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1614,14 +1619,14 @@ def _sparse_attn_decode_partial_kernel( valid = in_range & (slot >= 0) & (slot < extra_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // extra_block_size - pos_in_block = safe_slot % extra_block_size + block_idx = safe_slot // EXTRA_BLOCK_SIZE + pos_in_block = safe_slot % EXTRA_BLOCK_SIZE cache_block_ptr = ( extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 ) token_data_ptr = cache_block_ptr + pos_in_block * 576 token_scale_ptr = ( - cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 + cache_block_ptr + EXTRA_BLOCK_SIZE * 576 + pos_in_block * 8 ) x_uint8 = tl.load( @@ -1638,7 +1643,7 @@ def _sparse_attn_decode_partial_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + scales = _decode_e8m0_scales_triton(encoded_scales) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1933,7 +1938,6 @@ def _decode_num_splits( avg_main_len: float = 0.0, avg_extra_len: float = 0.0, block_k: int = 32, - max_splits: int = 16, ) -> int: """Pick a flash-decode split count to keep the GPU busy across batch sizes. @@ -1974,7 +1978,9 @@ def _decode_num_splits( mu = 0.04 best_splits = 1 best_cost = None - for splits in range(1, max_splits + 1): + # A full C128A row can contain 8K KV tokens. At low batch sizes, 32 splits + # keep those long rows to one device wave and halve the partial iterations. + for splits in range(1, 33): waves = (base * splits + cu - 1) // cu cost = waves * (1.0 / splits + mu) if best_cost is None or cost < best_cost - 1e-9: @@ -1995,19 +2001,6 @@ def _decode_num_splits( return best_splits -@functools.cache -def _max_decode_partitions(max_queries: int, heads_blocks: int, max_splits: int) -> int: - return max( - queries - * _decode_num_splits( - queries, - heads_blocks, - max_splits=max_splits, - ) - for queries in range(1, max_queries + 1) - ) - - def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -2021,7 +2014,6 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indices: torch.Tensor | None = None, extra_indptr: torch.Tensor | None = None, out: torch.Tensor | None = None, - split_k_buffers: SparseDecodeSplitKBuffers | None = None, ) -> torch.Tensor: assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert main_cache.ndim == 3, ( @@ -2135,135 +2127,69 @@ def _rocm_sparse_attn_decode_ragged_triton( ) return out - # Keep the split policy calibrated to the existing 32-token work estimate. - # The unified gfx950 Gluon body uses its fixed 64-token tile internally. - block_k = 32 + block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950. # Average per-query segment lengths, read sync-free from the ragged index # sizes, let the split heuristic avoid over-splitting # main_indices/extra_indices are flat [nnz] int32. inv_q = 1.0 / max(1, num_queries) avg_main_len = main_indices.numel() * inv_q avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 - if _ON_GFX950: - num_splits = _decode_num_splits( - num_queries, - heads_blocks, - avg_main_len, - avg_extra_len, - block_k, - max_splits=32, - ) - else: - num_splits = _decode_num_splits( - num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k - ) - - if not _ON_GFX950: - split_k_buffers = None - required_partitions = num_queries * num_splits - if _ON_GFX950 and num_splits == 1: - part_m = part_l = part_acc = out - elif split_k_buffers is not None: - part_m, part_l, part_acc = split_k_buffers - if split_k_buffers is None and not (_ON_GFX950 and num_splits == 1): - part_shape = ( - (required_partitions, num_heads) - if _ON_GFX950 - else (num_queries, num_splits, num_heads) - ) - part_m = torch.empty(part_shape, dtype=torch.float32, device=q.device) - part_l = torch.empty_like(part_m) - acc_shape = ( - (required_partitions, num_heads, comb_dim) - if _ON_GFX950 - else (num_queries, num_splits, num_heads, comb_dim) - ) - part_acc = torch.empty( - acc_shape, - dtype=torch.float32, - device=q.device, - ) - - if _ON_GFX950: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse_gluon import ( - launch_sparse_attn_decode_partial_gfx950, - ) - - launch_sparse_attn_decode_partial_gfx950( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - attn_sink, - out, - part_m, - part_l, - part_acc, - scale, - num_heads, - has_extra, - has_attn_sink, - num_splits, - ) - else: - _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - part_m, - part_l, - part_acc, - q.stride(0), - q.stride(1), - main_cache.stride(0), - extra_cache.stride(0), - part_m.stride(0), - part_m.stride(1), - part_acc.stride(0), - part_acc.stride(1), - part_acc.stride(2), - main_cache.shape[0] * main_cache.shape[1], - extra_cache.shape[0] * extra_cache.shape[1], - main_cache.shape[1], - extra_cache.shape[1], - scale, - num_heads, - HAS_EXTRA=has_extra, - NOPE_DIM=nope_head_dim, - NOPE_BLOCK=nope_block, - ROPE_DIM=rope_head_dim, - IS_FNUZ_MAIN=is_fnuz, - IS_FNUZ_EXTRA=False, - BLOCK_H=block_h, - BLOCK_K=block_k, - NUM_SPLITS=num_splits, - NUM_STAGES=1, - num_warps=4, - ) + num_splits = _decode_num_splits( + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + ) - if _ON_GFX950 and num_splits == 1: - return out + part_m = torch.empty( + (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device + ) + part_l = torch.empty_like(part_m) + part_acc = torch.empty( + (num_queries, num_splits, num_heads, comb_dim), + dtype=torch.float32, + device=q.device, + ) - if _ON_GFX950: - pm_stride_s = part_m.stride(0) - pm_stride0 = num_splits * pm_stride_s - pa_stride_s = part_acc.stride(0) - pa_stride0 = num_splits * pa_stride_s - pa_stride_h = part_acc.stride(1) - else: - pm_stride0, pm_stride_s = part_m.stride(0), part_m.stride(1) - pa_stride0, pa_stride_s, pa_stride_h = ( - part_acc.stride(0), - part_acc.stride(1), - part_acc.stride(2), - ) + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + NOPE_BLOCK=nope_block, + ROPE_DIM=rope_head_dim, + # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950). + # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere). + # Reading both with a single IS_FNUZ would decode one of them with the + # wrong FNUZ/OCP scale ratio (~1.87×). + IS_FNUZ_MAIN=is_fnuz, + IS_FNUZ_EXTRA=False, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + ) _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( part_m, @@ -2273,11 +2199,11 @@ def _rocm_sparse_attn_decode_ragged_triton( out, out.stride(0), out.stride(1), - pm_stride0, - pm_stride_s, - pa_stride0, - pa_stride_s, - pa_stride_h, + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), num_heads, HAS_ATTN_SINK=has_attn_sink, COMB_DIM=comb_dim, @@ -2306,7 +2232,6 @@ def _rocm_sparse_attn_decode_triton( extra_ragged_indices: torch.Tensor | None = None, extra_ragged_indptr: torch.Tensor | None = None, out: torch.Tensor | None = None, - split_k_buffers: SparseDecodeSplitKBuffers | None = None, ) -> torch.Tensor: if main_ragged_indices is None or main_ragged_indptr is None: main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense( @@ -2343,7 +2268,6 @@ def _rocm_sparse_attn_decode_triton( extra_indices=extra_ragged_indices, extra_indptr=extra_ragged_indptr, out=out, - split_k_buffers=split_k_buffers, ) @@ -2415,7 +2339,6 @@ def rocm_sparse_attn_decode( nope_head_dim: int, rope_head_dim: int, output: torch.Tensor, - split_k_buffers: SparseDecodeSplitKBuffers | None = None, ) -> None: assert swa_k_cache.dtype == torch.uint8, ( "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, " @@ -2445,7 +2368,7 @@ def rocm_sparse_attn_decode( if topk_indices is not None: extra_indices = topk_indices.reshape(topk_indices.shape[0], -1) - direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None + direct_out = output if output.dtype == torch.bfloat16 else None attn_out = _rocm_sparse_attn_decode_triton( q=q, main_cache=swa_k_cache, @@ -2463,7 +2386,6 @@ def rocm_sparse_attn_decode( extra_ragged_indices=topk_ragged_indices, extra_ragged_indptr=topk_ragged_indptr, out=direct_out, - split_k_buffers=split_k_buffers, ) if direct_out is None: output.copy_(attn_out.to(output.dtype)) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse_gluon.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse_gluon.py deleted file mode 100644 index 816e53744e7d..000000000000 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse_gluon.py +++ /dev/null @@ -1,795 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""gfx950 Gluon partial kernel for DeepSeek-V4 sparse-MLA decode. - -Adapted from AITER's MIT-licensed ``pa_decode_sparse`` gfx950 kernel. The vLLM -variant is intentionally limited to the packed two-segment DSv4 cache contract -and writes the existing Triton reducer's partial-buffer format. -""" - -import torch - -from vllm.triton_utils import gl, gluon - -_GFX950_BLOCK_H = 16 -_GFX950_BLOCK_K = 64 -_GFX950_HEAD_SIZE = 512 -_MAX_BUFFER_OFFSET = 2**31 - 1 - - -def _max_addressable_bytes(tensor: torch.Tensor) -> int: - """Return the byte span reachable from a tensor's data pointer.""" - span = 1 - for size, stride in zip(tensor.shape, tensor.stride()): - if size > 1: - span += (size - 1) * abs(stride) - return span * tensor.element_size() - - -@gluon.jit -def _cache_load(ptr, offsets, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): - if USE_BUFFER_LOAD: - return gl.amd.cdna4.buffer_load( - ptr=ptr, - offsets=offsets.to(gl.int32), - mask=mask, - other=other, - cache=".cg", - ) - return gl.load( - ptr + offsets.to(gl.int64), - mask=mask, - other=other, - cache_modifier=".cg", - ) - - -@gluon.jit -def _decode_e8m0_scales(encoded_scales): - scale_bits = encoded_scales.to(gl.int32) << 23 - scale_bits = gl.where(encoded_scales == 0, 1 << 22, scale_bits) - return scale_bits.to(gl.float32, bitcast=True) - - -@gluon.jit -def _slots( - indices_ptr, - segment_start, - k_pos, - segment_hi, - num_rows, - BLOCK_SIZE: gl.constexpr, - MASKED: gl.constexpr, -): - if MASKED: - in_range = k_pos < segment_hi - slot = gl.load(indices_ptr + segment_start + k_pos, mask=in_range, other=-1) - valid = in_range & (slot >= 0) & (slot < num_rows) - else: - slot = gl.load(indices_ptr + segment_start + k_pos) - valid = slot >= 0 - safe_slot = gl.where(valid, slot, 0) - return ( - (safe_slot // BLOCK_SIZE).to(gl.int32), - (safe_slot % BLOCK_SIZE).to(gl.int32), - valid, - ) - - -@gluon.jit -def _gather_fp8_tile( - cache_ptr, - cache_bf16_ptr, - indices_ptr, - segment_start, - k_start, - segment_hi, - cache_stride0, - num_rows, - full_offsets, - rope_offsets, - slot_offsets, - rope_slot_offsets, - gather_layout: gl.constexpr, - rope_gather_layout: gl.constexpr, - BLOCK_SIZE: gl.constexpr, - BLOCK_K: gl.constexpr, - MASKED: gl.constexpr, - USE_BUFFER_LOAD: gl.constexpr, -): - if not USE_BUFFER_LOAD: - cache_stride0 = cache_stride0.to(gl.int64) - - block, pos, valid = _slots( - indices_ptr, - segment_start, - k_start + slot_offsets, - segment_hi, - num_rows, - BLOCK_SIZE, - MASKED, - ) - block_g = gl.convert_layout(block, gl.SliceLayout(1, gather_layout)) - pos_g = gl.convert_layout(pos, gl.SliceLayout(1, gather_layout)) - valid_g = gl.convert_layout(valid, gl.SliceLayout(1, gather_layout)) - - data_offsets = (block_g * cache_stride0 + pos_g * 576)[:, None] + full_offsets[ - None, : - ] - scale_offsets = (block_g * cache_stride0 + BLOCK_SIZE * 576 + pos_g * 8)[ - :, None - ] + (full_offsets[None, :] // 64) - if MASKED: - data = _cache_load( - cache_ptr, - data_offsets, - USE_BUFFER_LOAD, - mask=valid_g[:, None], - other=0, - ) - encoded_scales = _cache_load( - cache_ptr, - scale_offsets, - USE_BUFFER_LOAD, - mask=valid_g[:, None], - other=127, - ) - else: - data = _cache_load(cache_ptr, data_offsets, USE_BUFFER_LOAD) - encoded_scales = _cache_load(cache_ptr, scale_offsets, USE_BUFFER_LOAD) - nope = ( - data.to(gl.float8e4nv, bitcast=True).to(gl.float32) - * _decode_e8m0_scales(encoded_scales) - ).to(gl.bfloat16) - - rope_block, rope_pos, rope_valid = _slots( - indices_ptr, - segment_start, - k_start + rope_slot_offsets, - segment_hi, - num_rows, - BLOCK_SIZE, - MASKED, - ) - rope_offsets_global = (rope_block * (cache_stride0 // 2) + rope_pos * 288 + 224)[ - :, None - ] + rope_offsets[None, :] - if MASKED: - rope = _cache_load( - cache_bf16_ptr, - rope_offsets_global, - USE_BUFFER_LOAD, - mask=rope_valid[:, None], - other=0.0, - ) - else: - rope = _cache_load( - cache_bf16_ptr, - rope_offsets_global, - USE_BUFFER_LOAD, - ) - return nope, rope, valid - - -@gluon.jit -def _consume_fp8_tile( - nope, - rope, - valid, - q_dot, - m_i, - l_i, - acc, - head_mask, - qk_scale, - kv_smem, - qk_layout: gl.constexpr, - pv_layout: gl.constexpr, - k_layout: gl.constexpr, - v_layout: gl.constexpr, - p_layout: gl.constexpr, - BLOCK_H: gl.constexpr, - BLOCK_K: gl.constexpr, - HEAD_ALIGNED: gl.constexpr, -): - kv_smem.store(nope) - kv_smem.slice(448, 64, dim=1).store(rope) - - k = kv_smem.permute([1, 0]).load(k_layout) - scores = gl.amd.cdna4.mfma( - q_dot, - k, - gl.zeros([BLOCK_H, BLOCK_K], gl.float32, layout=qk_layout), - ) - column_mask = gl.convert_layout(valid, gl.SliceLayout(0, qk_layout))[None, :] - if not HEAD_ALIGNED: - column_mask = ( - gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[:, None] - & column_mask - ) - scores = gl.where(column_mask, scores, float("-inf")) - - block_max = gl.max(scores, axis=1) - m_new = gl.maximum(m_i, block_max) - m_new = gl.where(m_new > float("-inf"), m_new, 0.0) - m_new_scaled = m_new * qk_scale - p = gl.exp2(scores * qk_scale - m_new_scaled[:, None]) - alpha = gl.exp2(m_i * qk_scale - m_new_scaled) - l_new = l_i * alpha + gl.sum(p, axis=1) - - v = kv_smem.load(v_layout) - p_dot = gl.convert_layout(p.to(gl.bfloat16), p_layout) - alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, pv_layout)) - acc = acc * alpha_pv[:, None] - acc = gl.amd.cdna4.mfma(p_dot, v, acc) - return m_new, l_new, acc - - -@gluon.jit -def _process_segment( - q_dot, - cache_ptr, - cache_bf16_ptr, - indices_ptr, - segment_start, - lo, - hi, - cache_stride0, - num_rows, - m_i, - l_i, - acc, - head_mask, - qk_scale, - kv_smem, - qk_layout: gl.constexpr, - pv_layout: gl.constexpr, - k_layout: gl.constexpr, - v_layout: gl.constexpr, - p_layout: gl.constexpr, - gather_layout: gl.constexpr, - rope_gather_layout: gl.constexpr, - slot_layout: gl.constexpr, - BLOCK_SIZE: gl.constexpr, - BLOCK_H: gl.constexpr, - BLOCK_K: gl.constexpr, - HEAD_ALIGNED: gl.constexpr, - USE_BUFFER_LOAD: gl.constexpr, -): - full_offsets = gl.arange( - 0, - 512, - layout=gl.SliceLayout(0, gather_layout), - ) - rope_offsets = gl.arange( - 0, - 64, - layout=gl.SliceLayout(0, rope_gather_layout), - ) - slot_offsets = gl.arange(0, BLOCK_K, layout=slot_layout) - rope_slot_offsets = gl.arange( - 0, - BLOCK_K, - layout=gl.SliceLayout(1, rope_gather_layout), - ) - - full_hi = lo + ((hi - lo) // BLOCK_K) * BLOCK_K - num_full = (full_hi - lo) // BLOCK_K - if num_full > 0: - nope, rope, valid = _gather_fp8_tile( - cache_ptr, - cache_bf16_ptr, - indices_ptr, - segment_start, - lo, - hi, - cache_stride0, - num_rows, - full_offsets, - rope_offsets, - slot_offsets, - rope_slot_offsets, - gather_layout, - rope_gather_layout, - BLOCK_SIZE, - BLOCK_K, - False, - USE_BUFFER_LOAD, - ) - for tile_idx in range(1, num_full): - next_nope, next_rope, next_valid = _gather_fp8_tile( - cache_ptr, - cache_bf16_ptr, - indices_ptr, - segment_start, - lo + tile_idx * BLOCK_K, - hi, - cache_stride0, - num_rows, - full_offsets, - rope_offsets, - slot_offsets, - rope_slot_offsets, - gather_layout, - rope_gather_layout, - BLOCK_SIZE, - BLOCK_K, - False, - USE_BUFFER_LOAD, - ) - m_i, l_i, acc = _consume_fp8_tile( - nope, - rope, - valid, - q_dot, - m_i, - l_i, - acc, - head_mask, - qk_scale, - kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - BLOCK_H, - BLOCK_K, - HEAD_ALIGNED, - ) - nope, rope, valid = next_nope, next_rope, next_valid - m_i, l_i, acc = _consume_fp8_tile( - nope, - rope, - valid, - q_dot, - m_i, - l_i, - acc, - head_mask, - qk_scale, - kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - BLOCK_H, - BLOCK_K, - HEAD_ALIGNED, - ) - - if full_hi < hi: - nope, rope, valid = _gather_fp8_tile( - cache_ptr, - cache_bf16_ptr, - indices_ptr, - segment_start, - full_hi, - hi, - cache_stride0, - num_rows, - full_offsets, - rope_offsets, - slot_offsets, - rope_slot_offsets, - gather_layout, - rope_gather_layout, - BLOCK_SIZE, - BLOCK_K, - True, - USE_BUFFER_LOAD, - ) - m_i, l_i, acc = _consume_fp8_tile( - nope, - rope, - valid, - q_dot, - m_i, - l_i, - acc, - head_mask, - qk_scale, - kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - BLOCK_H, - BLOCK_K, - HEAD_ALIGNED, - ) - return m_i, l_i, acc - - -@gluon.jit -def _sparse_attn_decode_partial_gfx950_kernel( - q_ptr, - main_cache_ptr, - main_cache_bf16_ptr, - main_indices_ptr, - main_indptr_ptr, - extra_cache_ptr, - extra_cache_bf16_ptr, - extra_indices_ptr, - extra_indptr_ptr, - part_m_ptr, - part_l_ptr, - part_acc_ptr, - scale: gl.constexpr, - q_stride0: gl.constexpr, - q_stride1: gl.constexpr, - main_cache_stride0, - extra_cache_stride0, - pm_stride0: gl.constexpr, - pm_stride_s: gl.constexpr, - pa_stride0: gl.constexpr, - pa_stride_s: gl.constexpr, - pa_stride_h: gl.constexpr, - main_num_rows, - extra_num_rows, - num_heads: gl.constexpr, - HAS_EXTRA: gl.constexpr, - HAS_ATTN_SINK: gl.constexpr, - MAIN_BLOCK_SIZE: gl.constexpr, - EXTRA_BLOCK_SIZE: gl.constexpr, - NUM_SPLITS: gl.constexpr, - HEAD_ALIGNED: gl.constexpr, - USE_BUFFER_LOAD: gl.constexpr, -): - num_warps: gl.constexpr = gl.num_warps() - query_idx = gl.program_id(0) - split_id = gl.program_id(1) - head_block = gl.program_id(2) - - qk_layout: gl.constexpr = gl.amd.AMDMFMALayout( - version=4, - instr_shape=[16, 16, 16], - transposed=True, - warps_per_cta=[1, num_warps], - ) - pv_layout: gl.constexpr = gl.amd.AMDMFMALayout( - version=4, - instr_shape=[16, 16, 16], - transposed=True, - warps_per_cta=[1, num_warps], - ) - q_layout: gl.constexpr = gl.DotOperandLayout(0, qk_layout, 8) - k_layout: gl.constexpr = gl.DotOperandLayout(1, qk_layout, 8) - p_layout: gl.constexpr = gl.DotOperandLayout(0, pv_layout, 8) - v_layout: gl.constexpr = gl.DotOperandLayout(1, pv_layout, 8) - - gather_layout: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, 16], - threads_per_warp=[8, 8], - warps_per_cta=[1, num_warps], - order=[1, 0], - ) - rope_gather_layout: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, 8], - threads_per_warp=[8, 8], - warps_per_cta=[num_warps, 1], - order=[1, 0], - ) - slot_layout: gl.constexpr = gl.SliceLayout(1, gather_layout) - q_blocked_layout: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, 8], - threads_per_warp=[8, 8], - warps_per_cta=[1, num_warps], - order=[1, 0], - ) - kv_shared_layout: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( - [[512, 8]], - [64, 512], - [1, 0], - ) - - head_base = head_block * 16 - head_offsets_q = gl.arange( - 0, - 16, - layout=gl.SliceLayout(1, q_blocked_layout), - ) - dim_offsets_q = gl.arange( - 0, - 512, - layout=gl.SliceLayout(0, q_blocked_layout), - ) - heads_q = head_base + head_offsets_q - head_mask_q = heads_q < num_heads - q_offsets = ( - query_idx * q_stride0 + heads_q[:, None] * q_stride1 + dim_offsets_q[None, :] - ).to(gl.int32) - q = gl.amd.cdna4.buffer_load( - ptr=q_ptr, - offsets=q_offsets, - mask=head_mask_q[:, None], - other=0.0, - ) - q_dot = gl.convert_layout(q, q_layout) - - head_offsets_pv = gl.arange( - 0, - 16, - layout=gl.SliceLayout(1, pv_layout), - ) - heads_pv = head_base + head_offsets_pv - head_mask_pv = heads_pv < num_heads - - m_i = gl.full( - [16], - float("-inf"), - gl.float32, - layout=gl.SliceLayout(1, qk_layout), - ) - l_i = gl.zeros( - [16], - gl.float32, - layout=gl.SliceLayout(1, qk_layout), - ) - acc = gl.zeros( - [16, 512], - gl.float32, - layout=pv_layout, - ) - kv_smem = gl.allocate_shared_memory( - gl.bfloat16, - [64, 512], - kv_shared_layout, - ) - - main_start = gl.load(main_indptr_ptr + query_idx) - main_end = gl.load(main_indptr_ptr + query_idx + 1) - main_len = main_end - main_start - main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS - main_lo = split_id * main_chunk - main_hi = gl.minimum(main_lo + main_chunk, main_len) - - rcp_ln2: gl.constexpr = 1.4426950408889634 - qk_scale: gl.constexpr = scale * rcp_ln2 - m_i, l_i, acc = _process_segment( - q_dot, - main_cache_ptr, - main_cache_bf16_ptr, - main_indices_ptr, - main_start, - main_lo, - main_hi, - main_cache_stride0, - main_num_rows, - m_i, - l_i, - acc, - head_mask_pv, - qk_scale, - kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - gather_layout, - rope_gather_layout, - slot_layout, - MAIN_BLOCK_SIZE, - 16, - 64, - HEAD_ALIGNED, - USE_BUFFER_LOAD, - ) - if HAS_EXTRA: - extra_start = gl.load(extra_indptr_ptr + query_idx) - extra_end = gl.load(extra_indptr_ptr + query_idx + 1) - extra_len = extra_end - extra_start - extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS - extra_lo = split_id * extra_chunk - extra_hi = gl.minimum(extra_lo + extra_chunk, extra_len) - m_i, l_i, acc = _process_segment( - q_dot, - extra_cache_ptr, - extra_cache_bf16_ptr, - extra_indices_ptr, - extra_start, - extra_lo, - extra_hi, - extra_cache_stride0, - extra_num_rows, - m_i, - l_i, - acc, - head_mask_pv, - qk_scale, - kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - gather_layout, - rope_gather_layout, - slot_layout, - EXTRA_BLOCK_SIZE, - 16, - 64, - HEAD_ALIGNED, - USE_BUFFER_LOAD, - ) - - m_pv = gl.convert_layout(m_i, gl.SliceLayout(1, pv_layout)) - l_pv = gl.convert_layout(l_i, gl.SliceLayout(1, pv_layout)) - if NUM_SPLITS == 1: - if HAS_ATTN_SINK: - m_scaled = m_pv * scale - sink = gl.amd.cdna4.buffer_load( - ptr=part_m_ptr, - offsets=heads_pv.to(gl.int32), - mask=head_mask_pv, - other=float("-inf"), - ).to(gl.float32) - m_final = gl.maximum(m_scaled, sink) - alpha = gl.exp2((m_scaled - m_final) * rcp_ln2) - l_final = l_pv * alpha + gl.exp2((sink - m_final) * rcp_ln2) - acc = acc * alpha[:, None] - else: - l_final = l_pv - denom = gl.maximum(l_final, 1.0e-30) - out = gl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0) - dim_offsets_out = gl.arange( - 0, - 512, - layout=gl.SliceLayout(0, pv_layout), - ) - out_offsets = ( - query_idx * pa_stride0 - + heads_pv[:, None] * pa_stride_h - + dim_offsets_out[None, :] - ).to(gl.int32) - gl.amd.cdna4.buffer_store( - out.to(part_acc_ptr.dtype.element_ty), - ptr=part_acc_ptr, - offsets=out_offsets, - mask=head_mask_pv[:, None], - ) - else: - m_natural = gl.where( - m_pv > float("-inf"), - m_pv * scale, - -3.4028234663852886e38, - ) - partial_row_base = query_idx * pm_stride0 + split_id * pm_stride_s - gl.amd.cdna4.buffer_store( - m_natural, - ptr=part_m_ptr + partial_row_base, - offsets=heads_pv.to(gl.int32), - mask=head_mask_pv, - ) - gl.amd.cdna4.buffer_store( - l_pv, - ptr=part_l_ptr + partial_row_base, - offsets=heads_pv.to(gl.int32), - mask=head_mask_pv, - ) - dim_offsets_acc = gl.arange( - 0, - 512, - layout=gl.SliceLayout(0, pv_layout), - ) - partial_acc_base = query_idx * pa_stride0 + split_id * pa_stride_s - partial_acc_offsets = ( - partial_acc_base - + heads_pv[:, None] * pa_stride_h - + dim_offsets_acc[None, :] - ).to(gl.int32) - gl.amd.cdna4.buffer_store( - acc, - ptr=part_acc_ptr, - offsets=partial_acc_offsets, - mask=head_mask_pv[:, None], - ) - - -def launch_sparse_attn_decode_partial_gfx950( - q: torch.Tensor, - main_cache: torch.Tensor, - main_indices: torch.Tensor, - main_indptr: torch.Tensor, - extra_cache: torch.Tensor, - extra_indices: torch.Tensor, - extra_indptr: torch.Tensor, - attn_sink: torch.Tensor, - out: torch.Tensor, - part_m: torch.Tensor, - part_l: torch.Tensor, - part_acc: torch.Tensor, - scale: float, - num_heads: int, - has_extra: bool, - has_attn_sink: bool, - num_splits: int, -) -> None: - """Launch the fixed gfx950 packed-cache partial kernel.""" - assert q.shape[-1] == _GFX950_HEAD_SIZE - assert main_cache.dtype == torch.uint8 - assert extra_cache.dtype == torch.uint8 - - use_buffer_load = ( - max( - _max_addressable_bytes(main_cache), - _max_addressable_bytes(extra_cache), - ) - < _MAX_BUFFER_OFFSET - ) - heads_blocks = (num_heads + _GFX950_BLOCK_H - 1) // _GFX950_BLOCK_H - grid = (q.shape[0], num_splits, heads_blocks) - if num_splits == 1: - part_m = attn_sink - part_l = out - part_acc = out - pm_stride0 = pm_stride_s = 0 - pa_stride0 = out.stride(0) - pa_stride_s = 0 - pa_stride_h = out.stride(1) - else: - pm_stride_s = part_m.stride(0) - pm_stride0 = num_splits * pm_stride_s - pa_stride_s = part_acc.stride(0) - pa_stride0 = num_splits * pa_stride_s - pa_stride_h = part_acc.stride(1) - _sparse_attn_decode_partial_gfx950_kernel[grid]( - q, - main_cache, - main_cache.view(torch.bfloat16), - main_indices, - main_indptr, - extra_cache, - extra_cache.view(torch.bfloat16), - extra_indices, - extra_indptr, - part_m, - part_l, - part_acc, - scale, - q.stride(0), - q.stride(1), - main_cache.stride(0), - extra_cache.stride(0), - pm_stride0, - pm_stride_s, - pa_stride0, - pa_stride_s, - pa_stride_h, - main_cache.shape[0] * main_cache.shape[1], - extra_cache.shape[0] * extra_cache.shape[1], - num_heads, - HAS_EXTRA=has_extra, - HAS_ATTN_SINK=has_attn_sink, - MAIN_BLOCK_SIZE=main_cache.shape[1], - EXTRA_BLOCK_SIZE=extra_cache.shape[1], - NUM_SPLITS=num_splits, - HEAD_ALIGNED=num_heads % _GFX950_BLOCK_H == 0, - USE_BUFFER_LOAD=use_buffer_load, - num_warps=4, - waves_per_eu=0, - ) From f3d64de0abafcb0c2dfb0fb8321aa8a15d2bae09 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Fri, 14 Aug 2026 11:08:32 +0000 Subject: [PATCH 4/8] [ROCm][DSV4][Perf] Optimize Triton sparse MLA decode on gfx950 Add a gfx950-only standard Triton sparse-decode partial specialization with improved QK/PV dataflow, workload-aware split selection, and trusted NaN-free compressed-cache handling. Preserve the standard gfx942 and generic paths and the existing reducer. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai --- .../attention/test_rocm_triton_attn_dsv4.py | 526 ++++++++++- tests/kernels/test_compressor_kv_cache.py | 135 +++ vllm/models/deepseek_v4/amd/rocm.py | 21 + .../common/ops/fused_compress_quant_cache.py | 23 +- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 814 ++++++++++++++++-- 5 files changed, 1438 insertions(+), 81 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 775f1f5d1447..7ec9ed7c0371 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -24,12 +24,27 @@ def _on_split_decode_arch() -> bool: return False +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + try: + from vllm.platforms.rocm import _ON_GFX950 + + return _ON_GFX950 + except ImportError: + return False + + # The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other # architectures take the fallback decode kernel, so its tests are skipped there. requires_split_decode_arch = pytest.mark.skipif( not _on_split_decode_arch(), reason="split-K decode kernel is only tuned for AMD gfx942/gfx950", ) +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), + reason="optimized sparse decode partial is gfx950-only", +) NOPE_HEAD_DIM = 448 ROPE_HEAD_DIM = 64 @@ -118,28 +133,34 @@ def _pack_fp8_ds_mla_cache( return cache -def _read_fp8_ds_mla_cache( - cache: torch.Tensor, slot: int, block_size: int, use_fnuz: bool -) -> torch.Tensor: - cache_flat = cache.view(torch.uint8).flatten() +def _poison_fp8_ds_mla_cache_row( + cache: torch.Tensor, block_size: int, slot: int = 0 +) -> None: + flat = cache.flatten() block_idx = slot // block_size pos = slot % block_size block_base = block_idx * cache.stride(0) token_base = block_base + pos * 576 scale_base = block_base + block_size * 576 + pos * 8 - - fp8_dtype = torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn - nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM] - nope = nope_u8.view(fp8_dtype).to(torch.float32) - scales = torch.exp2( - cache_flat[scale_base : scale_base + 7].to(torch.float32) - 127.0 + flat[token_base] = 0x7F + flat[scale_base : scale_base + 7] = 255 + flat[token_base + NOPE_HEAD_DIM : token_base + 576].view(torch.bfloat16)[0] = float( + "nan" ) - nope = nope * scales.repeat_interleave(64) - rope_u8 = cache_flat[ - token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2 - ] - rope = rope_u8.view(torch.bfloat16).to(torch.float32) - return torch.cat([nope, rope]) + + +def _canonicalize_poisoned_fp8_ds_mla_cache_row( + cache: torch.Tensor, block_size: int, slot: int = 0 +) -> None: + flat = cache.flatten() + block_idx = slot // block_size + pos = slot % block_size + block_base = block_idx * cache.stride(0) + token_base = block_base + pos * 576 + scale_base = block_base + block_size * 576 + pos * 8 + flat[token_base] = 0 + flat[scale_base : scale_base + 7] = 254 + flat[token_base + NOPE_HEAD_DIM : token_base + 576].view(torch.bfloat16)[0] = 0 def _read_fp8_ds_mla_cache_rows( @@ -186,19 +207,30 @@ def _ref_sparse_decode_ragged( out = torch.empty_like(q_f32) for query_idx in range(q.shape[0]): - row_kv = [ - _read_fp8_ds_mla_cache(main_cache, int(slot), block_size, main_use_fnuz) - for slot in main_rows[query_idx] - ] - if extra_cache is not None and extra_rows is not None: - row_kv.extend( - _read_fp8_ds_mla_cache( - extra_cache, int(slot), block_size, extra_use_fnuz + row_kv = [] + if main_rows[query_idx]: + main_slots = torch.tensor( + main_rows[query_idx], dtype=torch.int64, device=q.device + ) + row_kv.append( + _read_fp8_ds_mla_cache_rows( + main_cache, main_slots, block_size, main_use_fnuz + ) + ) + if extra_cache is not None and extra_rows is not None and extra_rows[query_idx]: + extra_slots = torch.tensor( + extra_rows[query_idx], dtype=torch.int64, device=q.device + ) + row_kv.append( + _read_fp8_ds_mla_cache_rows( + extra_cache, extra_slots, block_size, extra_use_fnuz ) - for slot in extra_rows[query_idx] ) - kv = torch.stack(row_kv).to(q.device) + if not row_kv: + out[query_idx] = 0 + continue + kv = torch.cat(row_kv) for head_idx in range(q.shape[1]): scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale if attn_sink is not None: @@ -226,6 +258,78 @@ def _ragged_from_rows( ) +def _launch_gfx950_partial( + q: torch.Tensor, + main_cache: torch.Tensor, + main_indices: torch.Tensor, + main_indptr: torch.Tensor, + extra_cache: torch.Tensor, + extra_indices: torch.Tensor, + extra_indptr: torch.Tensor, + num_splits: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + num_queries, num_heads, _ = q.shape + part_m = torch.empty( + num_queries, + num_splits, + num_heads, + dtype=torch.float32, + device=q.device, + ) + part_l = torch.empty_like(part_m) + part_acc = torch.empty( + num_queries, + num_splits, + num_heads, + HEAD_DIM, + dtype=torch.float32, + device=q.device, + ) + assert part_m.is_contiguous() + assert part_l.is_contiguous() + assert part_acc.is_contiguous() + + mod._sparse_attn_decode_gfx950_partial_kernel[ + (num_queries, num_splits, (num_heads + 15) // 16) + ]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + HEAD_DIM**-0.5, + num_heads, + HAS_EXTRA=True, + NOPE_DIM=NOPE_HEAD_DIM, + ROPE_DIM=ROPE_HEAD_DIM, + IS_FNUZ_MAIN=current_platform.is_fp8_fnuz(), + IS_FNUZ_EXTRA=False, + TRUST_EXTRA_CACHE_NAN_FREE=True, + BLOCK_H=16, + BLOCK_K=32, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + waves_per_eu=0, + ) + return part_m, part_l, part_acc + + @torch.inference_mode() def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None: from vllm._aiter_ops import rocm_aiter_ops @@ -340,6 +444,19 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None: torch.testing.assert_close(actual_lens, expected_lens) +def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None: + from vllm.models.deepseek_v4.amd import rocm as mod + + monkeypatch.setattr(mod, "_ON_GFX950", True) + assert mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", True, True) + assert not mod._trust_dsv4_extra_cache_nan_free("bfloat16", False, True) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, False) + + monkeypatch.setattr(mod, "_ON_GFX950", False) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True) + + @torch.inference_mode() def test_sparse_attn_prefill_ragged_kernel() -> None: from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( @@ -425,6 +542,85 @@ def test_sparse_attn_decode_ragged_kernel() -> None: torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_scrubs_untrusted_cache_by_default() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_triton, + ) + + device = torch.device("cuda") + block_size = 4 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + extra_cache = torch.zeros_like(main_cache) + _poison_fp8_ds_mla_cache_row(main_cache, block_size) + _poison_fp8_ds_mla_cache_row(extra_cache, block_size) + indices = torch.zeros(1, dtype=torch.int32, device=device) + indptr = torch.tensor([0, 1], dtype=torch.int32, device=device) + + actual = _rocm_sparse_attn_decode_ragged_triton( + q=torch.ones(1, 1, HEAD_DIM, dtype=torch.bfloat16, device=device), + main_cache=main_cache, + main_indices=indices, + main_indptr=indptr, + scale=HEAD_DIM**-0.5, + attn_sink=None, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=indices, + extra_indptr=indptr, + ) + + assert not torch.isnan(actual).any() + assert torch.equal(actual, torch.zeros_like(actual)) + + +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_trusted_extra_matches_legacy_scrub() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_triton, + ) + + device = torch.device("cuda") + block_size = 4 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + legacy_extra = torch.zeros_like(main_cache) + _poison_fp8_ds_mla_cache_row(main_cache, block_size) + _poison_fp8_ds_mla_cache_row(legacy_extra, block_size) + canonical_extra = legacy_extra.clone() + _canonicalize_poisoned_fp8_ds_mla_cache_row(canonical_extra, block_size) + indices = torch.zeros(1, dtype=torch.int32, device=device) + indptr = torch.tensor([0, 1], dtype=torch.int32, device=device) + q = torch.ones(1, 1, HEAD_DIM, dtype=torch.bfloat16, device=device) + kwargs = { + "q": q, + "main_cache": main_cache, + "main_indices": indices, + "main_indptr": indptr, + "scale": HEAD_DIM**-0.5, + "attn_sink": None, + "nope_head_dim": NOPE_HEAD_DIM, + "rope_head_dim": ROPE_HEAD_DIM, + "extra_indices": indices, + "extra_indptr": indptr, + } + + legacy = _rocm_sparse_attn_decode_ragged_triton( + **kwargs, + extra_cache=legacy_extra, + ) + trusted = _rocm_sparse_attn_decode_ragged_triton( + **kwargs, + extra_cache=canonical_extra, + extra_cache_nan_free=True, + ) + + assert torch.equal(trusted, legacy) + assert torch.equal(trusted, torch.zeros_like(trusted)) + + @requires_split_decode_arch @torch.inference_mode() def test_decode_num_splits_heuristic(monkeypatch) -> None: @@ -454,6 +650,24 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 +@torch.inference_mode() +def test_decode_num_splits_gfx950(monkeypatch) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256) + for num_queries in (1, 8): + assert mod._decode_gfx950_num_splits(num_queries, 1, 128, 32) == 32 + assert mod._decode_gfx950_num_splits(num_queries, 1, 128, 7812) == 32 + + assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 31 + for extra_rows in (32, 256): + assert mod._decode_gfx950_num_splits(64, 1, 128, extra_rows) == 4 + assert mod._decode_gfx950_num_splits(64, 1, 128, 781) == 7 + for extra_rows in (3906, 7812): + assert mod._decode_gfx950_num_splits(64, 1, 128, extra_rows) == 8 + assert mod._decode_gfx950_num_splits(512, 1, 128, 7812) == 1 + + @requires_split_decode_arch @pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8, 32]) @pytest.mark.parametrize("with_extra", [True, False]) @@ -509,7 +723,8 @@ def test_sparse_attn_decode_split_k_kernel( scale = HEAD_DIM**-0.5 # Pin the split count so each parametrized value is exercised deterministically. - monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits) + split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits" + monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: num_splits) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, @@ -539,6 +754,260 @@ def test_sparse_attn_decode_split_k_kernel( torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@pytest.mark.parametrize("num_splits", [1, 2, 3, 8, 32]) +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_partial_buffer_layout(num_splits: int) -> None: + device = torch.device("cuda") + block_size = 64 + num_heads = 16 + entries_per_split = 65 + num_extra = num_splits * entries_per_split + q = torch.full( + (2, num_heads, HEAD_DIM), + 0.125, + dtype=torch.bfloat16, + device=device, + ) + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + main_indices = torch.empty(0, dtype=torch.int32, device=device) + main_indptr = torch.zeros(3, dtype=torch.int32, device=device) + extra_cache = _pack_fp8_ds_mla_cache( + torch.full( + (num_extra, HEAD_DIM), + 0.125, + dtype=torch.bfloat16, + device=device, + ), + block_size, + use_fnuz=False, + ) + extra_indices = torch.arange(num_extra, dtype=torch.int32, device=device) + extra_indptr = torch.tensor( + [0, num_extra, num_extra], dtype=torch.int32, device=device + ) + + batch = _launch_gfx950_partial( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + num_splits, + ) + single_nonempty = _launch_gfx950_partial( + q[:1], + main_cache, + main_indices, + torch.zeros(2, dtype=torch.int32, device=device), + extra_cache, + extra_indices, + torch.tensor([0, num_extra], dtype=torch.int32, device=device), + num_splits, + ) + single_empty = _launch_gfx950_partial( + q[1:], + main_cache, + main_indices, + torch.zeros(2, dtype=torch.int32, device=device), + extra_cache, + torch.empty(0, dtype=torch.int32, device=device), + torch.zeros(2, dtype=torch.int32, device=device), + num_splits, + ) + + for batch_part, nonempty_part, empty_part in zip( + batch, single_nonempty, single_empty + ): + assert torch.equal(batch_part[:1], nonempty_part) + assert torch.equal(batch_part[1:], empty_part) + + part_m, part_l, part_acc = batch + expected_m = HEAD_DIM * 0.125**2 * HEAD_DIM**-0.5 + torch.testing.assert_close( + part_m[0], + torch.full_like(part_m[0], expected_m), + rtol=1e-4, + atol=1e-4, + ) + assert torch.equal(part_l[0], torch.full_like(part_l[0], float(entries_per_split))) + torch.testing.assert_close( + part_acc[0], + torch.full_like(part_acc[0], entries_per_split * 0.125), + rtol=0, + atol=0, + ) + assert torch.equal( + part_m[1], torch.full_like(part_m[1], torch.finfo(torch.float32).min) + ) + assert torch.equal(part_l[1], torch.zeros_like(part_l[1])) + assert torch.equal(part_acc[1], torch.zeros_like(part_acc[1])) + + +@requires_gfx950 +@pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65]) +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_outer64_boundaries( + monkeypatch, extra_len: int +) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(13) + block_size = 4 + num_heads = 16 + num_extra_rows = 80 + q = torch.randn(2, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) + q *= 0.125 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + main_indices = torch.empty(0, dtype=torch.int32, device=device) + main_indptr = torch.zeros(3, dtype=torch.int32, device=device) + extra_cache = _pack_fp8_ds_mla_cache( + torch.randn(num_extra_rows, HEAD_DIM, dtype=torch.bfloat16, device=device) + * 0.125, + block_size, + use_fnuz=False, + ) + _poison_fp8_ds_mla_cache_row(extra_cache, block_size) + + raw_row = list(range(1, extra_len + 1)) + if extra_len > 3: + raw_row[3] = -1 + if extra_len > 40: + raw_row[40] = num_extra_rows + if extra_len > 64: + raw_row[64] = num_extra_rows + 1024 + extra_indices, extra_indptr = _ragged_from_rows([raw_row, []], device) + valid_row = [slot for slot in raw_row if 0 <= slot < num_extra_rows] + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) + + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1) + actual = mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=main_cache, + main_indices=main_indices, + main_indptr=main_indptr, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + ) + expected = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=[[], []], + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=[valid_row, []], + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + assert torch.equal(actual[1], torch.zeros_like(actual[1])) + + +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(17) + block_size = 64 + num_queries = 8 + num_heads = 16 + num_splits = 8 + extra_per_query = 65 * num_splits + q = ( + torch.randn( + num_queries, + num_heads, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + main_cache = _pack_fp8_ds_mla_cache( + torch.randn(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125, + block_size, + use_fnuz=False, + ) + extra_cache = _pack_fp8_ds_mla_cache( + torch.randn( + num_queries * extra_per_query, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125, + block_size, + use_fnuz=False, + ) + main_rows = [[query_idx] for query_idx in range(num_queries)] + extra_rows = [ + list(range(query_idx * extra_per_query, (query_idx + 1) * extra_per_query)) + for query_idx in range(num_queries) + ] + main_indices, main_indptr = _ragged_from_rows(main_rows, device) + extra_indices, extra_indptr = _ragged_from_rows(extra_rows, device) + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) + out = torch.empty_like(q) + + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: num_splits) + + def run_decode() -> torch.Tensor: + return mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=main_cache, + main_indices=main_indices, + main_indptr=main_indptr, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + out=out, + extra_cache_nan_free=True, + ) + + run_decode() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_out = run_decode() + captured = out.clone() + + q.add_(0.03125) + replacement = extra_rows[0][-1] + extra_indices[0] = replacement + extra_rows[0][0] = replacement + graph.replay() + torch.accelerator.synchronize() + + assert captured_out.data_ptr() == out.data_ptr() + assert not torch.equal(out, captured) + expected = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=main_rows, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=extra_rows, + ) + torch.testing.assert_close(out, expected, atol=2e-2, rtol=2e-2) + + @requires_split_decode_arch @torch.inference_mode() def test_sparse_attn_decode_long_context(monkeypatch) -> None: @@ -567,7 +1036,8 @@ def test_sparse_attn_decode_long_context(monkeypatch) -> None: extra_indptr = torch.tensor([0, num_kv], dtype=torch.int32, device=device) scale = HEAD_DIM**-0.5 - monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: 32) + split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits" + monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: 32) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, main_cache=main_cache, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index d8878a0aea90..445e598f2186 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -28,6 +28,7 @@ _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, _launch_two_stage_sparse_attn_compressor, + compress_norm_rope_store_triton, ) from vllm.models.deepseek_v4.compressor import _get_c128_boundary from vllm.platforms import current_platform @@ -35,6 +36,17 @@ from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + try: + from vllm.platforms.rocm import _ON_GFX950 + + return _ON_GFX950 + except Exception: + return False + + def test_compute_global_topk_reuses_output_buffers(): device = "cuda" topk_indices = torch.tensor( @@ -78,6 +90,129 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float): return x_fp8, scales +def _decode_dsv4_cache_row( + cache: torch.Tensor, block_size: int, scrub_nan: bool +) -> torch.Tensor: + flat = cache.flatten() + nope = flat[:448].view(torch.float8_e4m3fn).to(torch.bfloat16) + encoded = flat[block_size * 576 : block_size * 576 + 7] + scales = torch.exp2(encoded.to(torch.float32) - 127.0).to(torch.bfloat16) + nope = nope * scales.repeat_interleave(64) + rope = flat[448:576].view(torch.bfloat16) + decoded = torch.cat((nope, rope)) + if scrub_nan: + decoded = torch.where(decoded == decoded, decoded, 0.0) + return decoded + + +def _assert_nan_free_cache_matches_legacy_scrub( + cache: torch.Tensor, block_size: int +) -> None: + flat = cache.flatten() + scale_base = block_size * 576 + scale_codes = flat[scale_base : scale_base + 8] + assert scale_codes[0].item() == 254 + assert scale_codes[1].item() == 247 + assert scale_codes[:7].max().item() <= 254 + nope_bytes = flat[:448] + assert not ((nope_bytes == 0x7F) | (nope_bytes == 0xFF)).any() + + rope = flat[448:576].view(torch.bfloat16) + assert not torch.isnan(rope).any() + assert torch.isposinf(rope[0]) + assert torch.equal(rope[1:4], torch.zeros_like(rope[1:4])) + + legacy_cache = cache.clone() + legacy_flat = legacy_cache.flatten() + legacy_flat[scale_base] = 255 + legacy_rope = legacy_flat[448:576].view(torch.bfloat16) + legacy_rope[1:4] = float("nan") + canonical = _decode_dsv4_cache_row(cache, block_size, scrub_nan=False) + legacy = _decode_dsv4_cache_row(legacy_cache, block_size, scrub_nan=True) + torch.testing.assert_close(canonical, legacy, rtol=0, atol=0) + assert torch.isinf(canonical[0]) + assert torch.isposinf(canonical[64]) + + +@pytest.mark.skipif( + not _on_gfx950(), + reason="NaN-free fp8_ds_mla compressed-cache contract is gfx950-only", +) +@pytest.mark.parametrize("writer", ["single_pass", "two_stage_finalizer"]) +def test_gfx950_compressed_cache_canonicalizes_nonfinite(writer: str) -> None: + head_dim = 512 + rope_dim = 64 + block_size = 4 + device = "cuda" + + positions = torch.zeros(1, dtype=torch.int64, device=device) + slot_mapping = torch.zeros(1, dtype=torch.int64, device=device) + rms_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) + rms_weight[0] = float("inf") + rms_weight[64] = torch.finfo(torch.bfloat16).max + rms_weight[448] = float("inf") + rms_weight[450] = float("nan") + cos_sin_cache = torch.zeros(1, rope_dim, dtype=torch.float32, device=device) + cos_sin_cache[:, : rope_dim // 2] = 1.0 + cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + + state_cache = torch.zeros(1, 1, 2 * head_dim, dtype=torch.float32, device=device) + state_cache[..., :head_dim] = 1.0 + token_to_req = torch.zeros(1, dtype=torch.int32, device=device) + block_table = torch.zeros(1, 1, dtype=torch.int32, device=device) + + if writer == "single_pass": + compress_norm_rope_store_triton( + state_cache=state_cache, + num_actual=1, + token_to_req_indices=token_to_req, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=1, + state_width=head_dim, + cos_sin_cache=cos_sin_cache, + kv_cache=cache, + k_cache_metadata=SimpleNamespace(slot_mapping=slot_mapping), + pdl_kwargs={}, + head_dim=head_dim, + rope_head_dim=rope_dim, + compress_ratio=1, + overlap=False, + use_fp4_cache=False, + rms_norm_weight=rms_weight, + rms_norm_eps=1e-6, + quant_block=64, + token_stride=576, + scale_dim=8, + ) + else: + _launch_two_stage_sparse_attn_compressor( + state_cache, + token_to_req, + positions, + slot_mapping, + block_table, + 1, + head_dim, + 1, + cos_sin_cache, + cache, + slot_mapping, + rms_weight, + 1e-6, + 64, + 576, + 8, + head_dim, + rope_dim, + 1, + torch.empty(1, head_dim, dtype=torch.float32, device=device), + ) + + _assert_nan_free_cache_matches_legacy_scrub(cache, block_size) + + @pytest.mark.parametrize( ("starts", "query_start_loc", "expected"), [ diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 23223acd2dba..ae0b907ebeeb 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -19,6 +19,7 @@ DeepseekV4SparseMLAMetadataBuilder, ) from vllm.platforms import current_platform +from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( CommonAttentionMetadata, @@ -36,6 +37,19 @@ from vllm.v1.worker.workspace import current_workspace_manager +def _trust_dsv4_extra_cache_nan_free( + kv_cache_dtype: str, + has_kv_transfer: bool, + has_extra_cache: bool, +) -> bool: + return ( + _ON_GFX950 + and kv_cache_dtype == "fp8_ds_mla" + and not has_kv_transfer + and has_extra_cache + ) + + def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: lengths = lengths.to(dtype=torch.int32).contiguous() indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device) @@ -451,7 +465,9 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): backend_cls = DeepseekV4ROCMAiterMLASparseBackend def __init__(self, *args, **kwargs): + vllm_config = args[0] if args else kwargs["vllm_config"] super().__init__(*args, **kwargs) + self._has_kv_transfer = vllm_config.kv_transfer_config is not None # Block scale for the preshuffled weight; None = not preshuffled. self._wqa_wkv_scale: torch.Tensor | None = None self._wo_b_scale: torch.Tensor | None = None @@ -666,6 +682,11 @@ def _forward_decode( nope_head_dim=self.nope_head_dim, rope_head_dim=self.rope_head_dim, output=output, + extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free( + self.kv_cache_dtype, + self._has_kv_transfer, + not swa_only and kv_cache is not None, + ), ) def _forward_prefill( diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index a2085cd220f1..4c9f464ef067 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -24,8 +24,14 @@ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +if current_platform.is_rocm(): + from vllm.platforms.rocm import _ON_GFX950 +else: + _ON_GFX950 = False + from .fused_indexer_q import _fp32x2_to_fp4x2 @@ -61,12 +67,15 @@ def compress_norm_rope_store_triton( if head_dim == 512: kernel = _fused_kv_compress_norm_rope_insert_sparse_attn num_warps = 4 + kernel_kwargs = {"SANITIZE_CACHE_NANS": _ON_GFX950} elif use_fp4_cache: kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn num_warps = 1 + kernel_kwargs = {} else: kernel = _fused_kv_compress_norm_rope_insert_indexer_attn num_warps = 1 + kernel_kwargs = {} kernel[(num_actual,)]( # state cache @@ -103,6 +112,7 @@ def compress_norm_rope_store_triton( SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), num_warps=num_warps, + **kernel_kwargs, **pdl_kwargs, ) @@ -145,6 +155,7 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( TOKEN_STRIDE: tl.constexpr, # 576 for DeepseekV4 SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad) KV_BLOCK_STRIDE: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, ): """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope). @@ -261,7 +272,8 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( scale_idx = tl.arange(0, N_QUANT_BLOCKS) encoded = exponents + 127.0 - encoded = tl.maximum(tl.minimum(encoded, 255.0), 0.0) + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.maximum(tl.minimum(encoded, max_encoded), 0.0) tl.store( scale_ptr + scale_idx, encoded.to(tl.uint8), @@ -289,6 +301,8 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( new_even = even * cos_v - odd * sin_v new_odd = odd * cos_v + even * sin_v result = tl.interleave(new_even, new_odd) # [TRITON_BLOCK_SIZE] fp32 + if SANITIZE_CACHE_NANS: + result = tl.where(result == result, result, 0.0) # Store rotated rope portion as bf16 into the cache's bf16 area. bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) @@ -417,6 +431,7 @@ def _finalize_norm_rope_quant_store_sparse_attn( TOKEN_STRIDE: tl.constexpr, SCALE_DIM: tl.constexpr, KV_BLOCK_STRIDE: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, ): """Stage 2: read compressed_kv[512] from scratch buffer, then RMSNorm + FP8 quant (nope) + RoPE + bf16 store @@ -474,7 +489,8 @@ def _finalize_norm_rope_quant_store_sparse_attn( tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM) scale_idx = tl.arange(0, N_QUANT_BLOCKS) - encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0) + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.maximum(tl.minimum(exponents + 127.0, max_encoded), 0.0) tl.store( scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS ) @@ -494,6 +510,8 @@ def _finalize_norm_rope_quant_store_sparse_attn( new_even = even * cos_v - odd * sin_v new_odd = odd * cos_v + even * sin_v result = tl.interleave(new_even, new_odd) + if SANITIZE_CACHE_NANS: + result = tl.where(result == result, result, 0.0) bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) rope_local = block - NOPE_HEAD_DIM is_rope = (block >= NOPE_HEAD_DIM) & mask @@ -564,6 +582,7 @@ def _launch_two_stage_sparse_attn_compressor( TOKEN_STRIDE=token_stride, SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), + SANITIZE_CACHE_NANS=_ON_GFX950, ) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index cf1d4d44ed2c..74d03cda7048 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1225,6 +1225,83 @@ def _decode_e8m0_scales_triton(encoded_scales): return scale_bits.to(tl.float32, bitcast=True) +@triton.jit +def _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + CHUNK_START: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, +): + offsets = CHUNK_START + tl.arange(0, CHUNK_SIZE) + x_uint8 = tl.load( + token_data_ptr[:, None] + offsets[None, :], + mask=valid[:, None], + other=0, + ) + scale_offsets = CHUNK_START // 64 + tl.arange(0, CHUNK_SIZE // 64) + encoded_scales = tl.load( + token_scale_ptr[:, None] + scale_offsets[None, :], + mask=valid[:, None], + other=127, + ) + scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, CHUNK_SIZE // 64, 64)) + scales = tl.reshape(scales, (BLOCK_K, CHUNK_SIZE)) + if IS_FNUZ: + x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32) + else: + x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32) + value = (x_f32 * scales).to(tl.bfloat16) + zero = tl.zeros((BLOCK_K, CHUNK_SIZE), dtype=tl.bfloat16) + return tl.where(valid[:, None], value, zero) + + +@triton.jit +def _load_fp8_ds_mla_gfx950_tail128( + token_data_ptr, + token_scale_ptr, + valid, + NOPE_DIM: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, +): + tail_offsets = 384 + tl.arange(0, 128) + nope_mask = tail_offsets < NOPE_DIM + x_uint8 = tl.load( + token_data_ptr[:, None] + tail_offsets[None, :], + mask=valid[:, None] & nope_mask[None, :], + other=0, + ) + scale_offsets = 6 + tl.arange(0, 2) + scale_mask = scale_offsets < NOPE_DIM // 64 + encoded_scales = tl.load( + token_scale_ptr[:, None] + scale_offsets[None, :], + mask=valid[:, None] & scale_mask[None, :], + other=127, + ) + scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, 2, 64)) + scales = tl.reshape(scales, (BLOCK_K, 128)) + if IS_FNUZ: + x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32) + else: + x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32) + nope = (x_f32 * scales).to(tl.bfloat16) + + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + rope = tl.load( + rope_ptr[:, None] + (tail_offsets[None, :] - NOPE_DIM), + mask=valid[:, None] & ~nope_mask[None, :], + other=0.0, + ) + value = tl.where(nope_mask[None, :], nope, rope) + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) + return tl.where(valid[:, None], value, zero) + + @triton.jit def _sparse_attn_decode_ragged_kernel( q_ptr, @@ -1699,6 +1776,547 @@ def _sparse_attn_decode_partial_kernel( ) +@triton.jit +def _sparse_attn_decode_gfx950_partial_tail_tile( + q_combined, + cache_ptr, + indices_ptr, + index_start, + k_start, + k_hi, + cache_stride0, + num_rows, + scale: tl.constexpr, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + BLOCK_SIZE: tl.constexpr, + NOPE_DIM: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, +): + k_offsets = tl.arange(0, BLOCK_K) + k_pos = k_start + k_offsets + in_range = k_pos < k_hi + slot = tl.load(indices_ptr + index_start + k_pos, mask=in_range, other=-1) + valid = in_range & (slot >= 0) & (slot < num_rows) + safe_slot = tl.where(valid, slot, 0) + + block_idx = safe_slot // BLOCK_SIZE + pos_in_block = safe_slot % BLOCK_SIZE + cache_block_ptr = cache_ptr + block_idx.to(tl.int64) * cache_stride0 + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = cache_block_ptr + BLOCK_SIZE * 576 + pos_in_block * 8 + k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 0, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 128, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 256, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_tail = _load_fp8_ds_mla_gfx950_tail128( + token_data_ptr, + token_scale_ptr, + valid, + NOPE_DIM, + BLOCK_K, + IS_FNUZ, + ) + if not TRUST_EXTRA_CACHE_NAN_FREE: + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) + k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero) + k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero) + k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero) + k_tail = tl.where(k_tail == k_tail, k_tail, zero) + k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1) + k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1) + k_combined = tl.cat(k_nope_0, k_tail_256, dim=1) + + scores = tl.dot(q_combined, tl.trans(k_combined)) + scores *= scale * 1.4426950408889634 + scores = tl.where( + head_mask[:, None] & valid[None, :], + scores, + -3.4028234663852886e38, + ) + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + p_bf16 = p.to(k_nope_0a.dtype) + acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a) + acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b) + acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1) + acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail) + return ( + m_new, + l_new, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) + + +@triton.jit +def _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + cache_ptr, + slot, + valid, + cache_stride0, + scale: tl.constexpr, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + BLOCK_SIZE: tl.constexpr, + NOPE_DIM: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, +): + safe_slot = tl.where(valid, slot, 0) + block_idx = safe_slot // BLOCK_SIZE + pos_in_block = safe_slot % BLOCK_SIZE + cache_block_ptr = cache_ptr + block_idx.to(tl.int64) * cache_stride0 + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = cache_block_ptr + BLOCK_SIZE * 576 + pos_in_block * 8 + k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 0, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 128, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 256, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_tail = _load_fp8_ds_mla_gfx950_tail128( + token_data_ptr, + token_scale_ptr, + valid, + NOPE_DIM, + BLOCK_K, + IS_FNUZ, + ) + if not TRUST_EXTRA_CACHE_NAN_FREE: + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) + k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero) + k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero) + k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero) + k_tail = tl.where(k_tail == k_tail, k_tail, zero) + k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1) + k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1) + k_combined = tl.cat(k_nope_0, k_tail_256, dim=1) + + scores = tl.dot(q_combined, tl.trans(k_combined)) + scores *= scale * 1.4426950408889634 + scores = tl.where( + head_mask[:, None] & valid[None, :], + scores, + -3.4028234663852886e38, + ) + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + p_bf16 = p.to(k_nope_0a.dtype) + acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a) + acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b) + acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1) + acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail) + return ( + m_new, + l_new, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) + + +@triton.jit +def _sparse_attn_decode_gfx950_partial_kernel( + q_ptr, + main_cache_ptr, + main_indices_ptr, + main_indptr_ptr, + extra_cache_ptr, + extra_indices_ptr, + extra_indptr_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + q_stride0: tl.constexpr, + q_stride1: tl.constexpr, + main_cache_stride0: tl.constexpr, + extra_cache_stride0: tl.constexpr, + main_num_rows, + extra_num_rows, + MAIN_BLOCK_SIZE: tl.constexpr, + EXTRA_BLOCK_SIZE: tl.constexpr, + scale: tl.constexpr, + num_heads: tl.constexpr, + HAS_EXTRA: tl.constexpr, + NOPE_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + IS_FNUZ_MAIN: tl.constexpr, + IS_FNUZ_EXTRA: tl.constexpr, + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + query_idx = tl.program_id(0) + split_id = tl.program_id(1) + pid_h = tl.program_id(2) + + tl.static_assert(NOPE_DIM == 448) + tl.static_assert(ROPE_DIM == 64) + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + if num_heads % BLOCK_H == 0: + head_mask = tl.full((BLOCK_H,), True, tl.int1) + else: + head_mask = head_offsets < num_heads + nope_offsets_0a = tl.arange(0, 128) + nope_offsets_0b = 128 + tl.arange(0, 128) + nope_offsets_0 = tl.arange(0, 256) + tail_offsets = 256 + tl.arange(0, 256) + nope_offsets_1 = 256 + tl.arange(0, 128) + tail_offsets_128 = 384 + tl.arange(0, 128) + + q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1 + q_nope_0 = tl.load( + q_row_ptr + nope_offsets_0[None, :], + mask=head_mask[:, None], + other=0.0, + ) + q_tail = tl.load( + q_row_ptr + tail_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) + q_combined = tl.cat(q_nope_0, q_tail, dim=1) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc_nope_0a = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_nope_0b = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_nope_1 = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_tail = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + k_offsets = tl.arange(0, BLOCK_K) + + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS + main_lo = split_id * main_chunk + main_hi = tl.minimum(main_lo + main_chunk, main_len) + + for k_start in tl.range( + main_lo, + main_hi, + BLOCK_K, + num_stages=NUM_STAGES, + ): + k_pos = k_start + k_offsets + in_range = k_pos < main_hi + slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1) + valid = in_range & (slot >= 0) & (slot < main_num_rows) + safe_slot = tl.where(valid, slot, 0) + + block_idx = safe_slot // MAIN_BLOCK_SIZE + pos_in_block = safe_slot % MAIN_BLOCK_SIZE + cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 + + k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 0, + 128, + BLOCK_K, + IS_FNUZ_MAIN, + ) + k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 128, + 128, + BLOCK_K, + IS_FNUZ_MAIN, + ) + k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 256, + 128, + BLOCK_K, + IS_FNUZ_MAIN, + ) + k_tail = _load_fp8_ds_mla_gfx950_tail128( + token_data_ptr, + token_scale_ptr, + valid, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_MAIN, + ) + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) + k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero) + k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero) + k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero) + k_tail = tl.where(k_tail == k_tail, k_tail, zero) + k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1) + k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1) + k_combined = tl.cat(k_nope_0, k_tail_256, dim=1) + + scores = tl.dot(q_combined, tl.trans(k_combined)) + scores *= scale * 1.4426950408889634 + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + + p_bf16 = p.to(k_nope_0a.dtype) + acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a) + acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b) + acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1) + acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail) + m_i = m_new + l_i = l_new + + if HAS_EXTRA: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS + extra_lo = split_id * extra_chunk + extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len) + + outer_block_k: tl.constexpr = 2 * BLOCK_K + outer_k_offsets = tl.arange(0, outer_block_k) + extra_hi_full = ( + extra_lo + ((extra_hi - extra_lo) // outer_block_k) * outer_block_k + ) + for k_start in tl.range( + extra_lo, + extra_hi_full, + outer_block_k, + num_stages=NUM_STAGES, + ): + slot = tl.load(extra_indices_ptr + extra_start + k_start + outer_k_offsets) + valid = (slot >= 0) & (slot < extra_num_rows) + slot_pairs = tl.trans(tl.reshape(slot, (2, BLOCK_K))) + valid_pairs = tl.trans(tl.reshape(valid, (2, BLOCK_K))) + slot_lo, slot_hi = tl.split(slot_pairs) + valid_lo, valid_hi = tl.split(valid_pairs) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot_lo, + valid_lo, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot_hi, + valid_hi, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + if extra_hi_full < extra_hi: + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_tail_tile( + q_combined, + extra_cache_ptr, + extra_indices_ptr, + extra_start, + extra_hi_full, + extra_hi, + extra_cache_stride0, + extra_num_rows, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + if extra_hi_full + BLOCK_K < extra_hi: + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_tail_tile( + q_combined, + extra_cache_ptr, + extra_indices_ptr, + extra_start, + extra_hi_full + BLOCK_K, + extra_hi, + extra_cache_stride0, + extra_num_rows, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + + pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets + m_store = tl.where(l_i > 0.0, m_i * 0.6931471805599453, neg_large) + tl.store(part_m_ptr + pm_base, m_store, mask=head_mask) + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask) + acc_base = part_acc_ptr + ( + (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets[:, None] + ) * (NOPE_DIM + ROPE_DIM) + tl.store( + acc_base + nope_offsets_0a[None, :], + acc_nope_0a, + mask=head_mask[:, None], + ) + tl.store( + acc_base + nope_offsets_0b[None, :], + acc_nope_0b, + mask=head_mask[:, None], + ) + tl.store( + acc_base + nope_offsets_1[None, :], + acc_nope_1, + mask=head_mask[:, None], + ) + tl.store( + acc_base + tail_offsets_128[None, :], + acc_tail, + mask=head_mask[:, None], + ) + + @triton.jit def _sparse_attn_decode_reduce_kernel( part_m_ptr, @@ -2001,6 +2619,42 @@ def _decode_num_splits( return best_splits +def _decode_gfx950_num_splits( + num_queries: int, + heads_blocks: int, + avg_main_len: float = 0.0, + avg_extra_len: float = 0.0, + block_k: int = 32, +) -> int: + base = max(1, num_queries * heads_blocks) + cu = max(1, _decode_cu_count()) + target_workgroups = 2 * cu + num_splits = min( + 32, + max( + 1, + math.ceil(target_workgroups / base), + ), + ) + if ( + base >= 64 + and num_splits > 4 + and _decode_partial_iters(avg_main_len, avg_extra_len, 4, block_k) <= 3 + ): + return 4 + if base >= 64 and num_splits > 1: + target_waves = (base * num_splits + cu - 1) // cu + target_iters = _decode_partial_iters( + avg_main_len, avg_extra_len, num_splits, block_k + ) + for splits in range(1, num_splits): + waves = (base * splits + cu - 1) // cu + iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k) + if waves == target_waves and iters == target_iters: + return splits + return num_splits + + def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -2014,6 +2668,7 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indices: torch.Tensor | None = None, extra_indptr: torch.Tensor | None = None, out: torch.Tensor | None = None, + extra_cache_nan_free: bool = False, ) -> torch.Tensor: assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert main_cache.ndim == 3, ( @@ -2054,6 +2709,10 @@ def _rocm_sparse_attn_decode_ragged_triton( and extra_indices is not None and extra_indptr is not None ) + assert not extra_cache_nan_free or (_ON_GFX950 and has_extra), ( + "extra_cache_nan_free requires a gfx950 compressed cache with trusted " + "canonical-writer provenance" + ) if has_extra: assert extra_cache is not None assert extra_indices is not None @@ -2128,15 +2787,26 @@ def _rocm_sparse_attn_decode_ragged_triton( return out block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950. - # Average per-query segment lengths, read sync-free from the ragged index - # sizes, let the split heuristic avoid over-splitting - # main_indices/extra_indices are flat [nnz] int32. - inv_q = 1.0 / max(1, num_queries) - avg_main_len = main_indices.numel() * inv_q - avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 - num_splits = _decode_num_splits( - num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k - ) + if _ON_GFX950: + inv_q = 1.0 / max(1, num_queries) + avg_main_len = main_indices.numel() * inv_q + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 + num_splits = _decode_gfx950_num_splits( + num_queries, + heads_blocks, + avg_main_len, + avg_extra_len, + block_k, + ) + else: + # Average per-query segment lengths, read sync-free from the ragged + # index sizes, let the split heuristic avoid over-splitting. + inv_q = 1.0 / max(1, num_queries) + avg_main_len = main_indices.numel() * inv_q + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 + num_splits = _decode_num_splits( + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + ) part_m = torch.empty( (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device @@ -2148,48 +2818,86 @@ def _rocm_sparse_attn_decode_ragged_triton( device=q.device, ) - _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - part_m, - part_l, - part_acc, - q.stride(0), - q.stride(1), - main_cache.stride(0), - extra_cache.stride(0), - part_m.stride(0), - part_m.stride(1), - part_acc.stride(0), - part_acc.stride(1), - part_acc.stride(2), - main_cache.shape[0] * main_cache.shape[1], - extra_cache.shape[0] * extra_cache.shape[1], - main_cache.shape[1], - extra_cache.shape[1], - scale, - num_heads, - HAS_EXTRA=has_extra, - NOPE_DIM=nope_head_dim, - NOPE_BLOCK=nope_block, - ROPE_DIM=rope_head_dim, - # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950). - # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere). - # Reading both with a single IS_FNUZ would decode one of them with the - # wrong FNUZ/OCP scale ratio (~1.87×). - IS_FNUZ_MAIN=is_fnuz, - IS_FNUZ_EXTRA=False, - BLOCK_H=block_h, - BLOCK_K=block_k, - NUM_SPLITS=num_splits, - NUM_STAGES=1, - num_warps=4, - ) + if _ON_GFX950: + _sparse_attn_decode_gfx950_partial_kernel[ + (num_queries, num_splits, heads_blocks) + ]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + ROPE_DIM=rope_head_dim, + IS_FNUZ_MAIN=is_fnuz, + IS_FNUZ_EXTRA=False, + TRUST_EXTRA_CACHE_NAN_FREE=extra_cache_nan_free, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + waves_per_eu=0, + ) + else: + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + NOPE_BLOCK=nope_block, + ROPE_DIM=rope_head_dim, + # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950). + # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere). + # Reading both with a single IS_FNUZ would decode one of them with the + # wrong FNUZ/OCP scale ratio (~1.87×). + IS_FNUZ_MAIN=is_fnuz, + IS_FNUZ_EXTRA=False, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + ) _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( part_m, @@ -2232,6 +2940,7 @@ def _rocm_sparse_attn_decode_triton( extra_ragged_indices: torch.Tensor | None = None, extra_ragged_indptr: torch.Tensor | None = None, out: torch.Tensor | None = None, + extra_cache_nan_free: bool = False, ) -> torch.Tensor: if main_ragged_indices is None or main_ragged_indptr is None: main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense( @@ -2268,6 +2977,7 @@ def _rocm_sparse_attn_decode_triton( extra_indices=extra_ragged_indices, extra_indptr=extra_ragged_indptr, out=out, + extra_cache_nan_free=extra_cache_nan_free, ) @@ -2339,6 +3049,7 @@ def rocm_sparse_attn_decode( nope_head_dim: int, rope_head_dim: int, output: torch.Tensor, + extra_cache_nan_free: bool = False, ) -> None: assert swa_k_cache.dtype == torch.uint8, ( "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, " @@ -2386,6 +3097,7 @@ def rocm_sparse_attn_decode( extra_ragged_indices=topk_ragged_indices, extra_ragged_indptr=topk_ragged_indptr, out=direct_out, + extra_cache_nan_free=extra_cache_nan_free, ) if direct_out is None: output.copy_(attn_out.to(output.dtype)) From bc6925fc9e691e84463ff9b68be6a97f146d3419 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Fri, 14 Aug 2026 19:48:49 +0000 Subject: [PATCH 5/8] [ROCm][DSV4][Perf] Balance gfx950 sparse decode splits Broaden the gfx950 split policy across T=16-48 so short, mid, and long contexts retain the same modeled occupancy without excess reducer work. Extend the focused policy test with representative one- and two-wave cases. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai --- tests/kernels/attention/test_rocm_triton_attn_dsv4.py | 4 +++- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 11 +++++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 7ec9ed7c0371..79244947df15 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -659,7 +659,9 @@ def test_decode_num_splits_gfx950(monkeypatch) -> None: assert mod._decode_gfx950_num_splits(num_queries, 1, 128, 32) == 32 assert mod._decode_gfx950_num_splits(num_queries, 1, 128, 7812) == 32 - assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 31 + assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 4 + assert mod._decode_gfx950_num_splits(16, 1, 128, 781) == 13 + assert mod._decode_gfx950_num_splits(48, 1, 128, 3906) == 10 for extra_rows in (32, 256): assert mod._decode_gfx950_num_splits(64, 1, 128, extra_rows) == 4 assert mod._decode_gfx950_num_splits(64, 1, 128, 781) == 7 diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 74d03cda7048..eb07d5fd21b3 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -2637,12 +2637,19 @@ def _decode_gfx950_num_splits( ), ) if ( - base >= 64 + base >= 16 and num_splits > 4 and _decode_partial_iters(avg_main_len, avg_extra_len, 4, block_k) <= 3 ): return 4 - if base >= 64 and num_splits > 1: + if 16 <= base < 64: + one_wave_splits = max(1, cu // base) + one_wave_iters = _decode_partial_iters( + avg_main_len, avg_extra_len, one_wave_splits, block_k + ) + target_waves = 1 if one_wave_iters <= 9 else 2 + num_splits = min(num_splits, max(1, target_waves * cu // base)) + if base >= 16 and num_splits > 1: target_waves = (base * num_splits + cu - 1) // cu target_iters = _decode_partial_iters( avg_main_len, avg_extra_len, num_splits, block_k From 5dfdac01ed82d02e2bfc3e26fa01448bde70dced Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Sun, 16 Aug 2026 00:32:08 +0000 Subject: [PATCH 6/8] Optimize DeepSeek V4 sparse decode on gfx950 Adapt captured sparse-decode work to live ragged lengths and avoid per-shape cp-gather compilation on gfx950. Preserve the legacy selector, kernels, and launch specialization on other ROCm architectures. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai --- benchmarks/kernels/benchmark_cp_gather.py | 124 ++++++++- .../attention/test_rocm_triton_attn_dsv4.py | 263 ++++++++++++++++-- tests/kernels/test_compressor_kv_cache.py | 108 ++++++- vllm/models/deepseek_v4/amd/rocm.py | 30 +- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 242 +++++++++++++--- 5 files changed, 683 insertions(+), 84 deletions(-) diff --git a/benchmarks/kernels/benchmark_cp_gather.py b/benchmarks/kernels/benchmark_cp_gather.py index 5e73d6d6e95c..903e7756267a 100644 --- a/benchmarks/kernels/benchmark_cp_gather.py +++ b/benchmarks/kernels/benchmark_cp_gather.py @@ -7,8 +7,12 @@ import torch from vllm import _custom_ops as ops +from vllm.platforms import current_platform from vllm.triton_utils import triton from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + cp_gather_indexer_k_quant_cache_triton, +) SCENARIOS = { "single-60k": [60_000], @@ -31,6 +35,15 @@ "bfloat16": torch.bfloat16, "float32": torch.float32, } +INDEXER_TOKEN_COUNTS = (1638, 1737, 1738, 1739, 2029, 2031, 3217, 4078) + + +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import _ON_GFX950 + + return _ON_GFX950 def make_page_table( @@ -172,6 +185,65 @@ def run() -> None: return run, bytes_moved +def make_indexer_quant_gather( + seq_lens: list[int], + block_size: int, +) -> tuple[Callable[[], None], int]: + head_dim = 128 + cache_entry_bytes = head_dim + 4 + max_blocks_per_seq = 4096 + num_cache_blocks = 8192 + block_table = torch.full( + (len(seq_lens), max_blocks_per_seq), + -1, + dtype=torch.int32, + device="cuda", + ) + next_block = 0 + for req_id, seq_len in enumerate(seq_lens): + req_blocks = math.ceil(seq_len / block_size) + block_table[req_id, :req_blocks] = torch.arange( + next_block, + next_block + req_blocks, + dtype=torch.int32, + device="cuda", + ) + next_block += req_blocks + assert next_block <= num_cache_blocks + + cu_seq_lens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device="cuda") + cu_seq_lens[1:] = torch.tensor(seq_lens, dtype=torch.int32, device="cuda").cumsum( + dim=0 + ) + token_to_seq = torch.repeat_interleave( + torch.arange(len(seq_lens), dtype=torch.int32, device="cuda"), + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + ) + kv_cache = torch.zeros( + (num_cache_blocks, block_size, cache_entry_bytes), + dtype=torch.uint8, + device="cuda", + ) + num_tokens = sum(seq_lens) + dst_k = torch.empty( + (num_tokens, head_dim), dtype=current_platform.fp8_dtype(), device="cuda" + ) + dst_scale = torch.empty((num_tokens, 4), dtype=torch.uint8, device="cuda") + + def run() -> None: + cp_gather_indexer_k_quant_cache_triton( + kv_cache, + dst_k, + dst_scale, + block_table, + cu_seq_lens, + token_to_seq, + ) + + bytes_moved = 2 * num_tokens * cache_entry_bytes + return run, bytes_moved + + @torch.inference_mode() def run_scenario( variant: str, @@ -187,17 +259,21 @@ def run_scenario( run, bytes_moved = make_cache_gather(seq_lens, block_size, entry_size, dtype) elif variant == "fp8-upconvert": run, bytes_moved = make_fp8_upconvert(seq_lens, block_size) + elif variant == "indexer-quant": + run, bytes_moved = make_indexer_quant_gather(seq_lens, block_size) else: run, bytes_moved = make_maybe_dequant_gather(seq_lens, block_size, entry_size) - latency_ms = triton.testing.do_bench( - run, warmup=warmup_ms, rep=rep_ms, return_mode="median" + latency_ms, p99_ms, p999_ms = triton.testing.do_bench( + run, warmup=warmup_ms, rep=rep_ms, quantiles=[0.5, 0.99, 0.999] ) bandwidth_gbps = bytes_moved / latency_ms / 1e6 lengths = ",".join(str(seq_len) for seq_len in seq_lens) print( f"{variant:15s} {name:10s} batch={len(seq_lens):2d} " f"total={sum(seq_lens):7d} latency={latency_ms * 1e3:9.2f} us " + f"p99={p99_ms * 1e3:9.2f} us " + f"p99.9={p999_ms * 1e3:9.2f} us " f"bandwidth={bandwidth_gbps:8.1f} GB/s lengths=[{lengths}]" ) @@ -206,10 +282,12 @@ def main() -> None: parser = FlexibleArgumentParser(description="Benchmark cp_gather variants") parser.add_argument( "--variant", - choices=["all", "cache", "fp8-upconvert", "maybe-dequant"], + choices=["all", "cache", "fp8-upconvert", "maybe-dequant", "indexer-quant"], default="all", ) - parser.add_argument("--scenario", choices=["all", *SCENARIOS], default="all") + parser.add_argument( + "--scenario", choices=["all", *SCENARIOS, "indexer-8k1k"], default="all" + ) parser.add_argument("--dtype", choices=DTYPES, default="fp8") parser.add_argument("--block-size", type=int, default=64) parser.add_argument("--entry-size", type=int, default=576) @@ -217,19 +295,37 @@ def main() -> None: parser.add_argument("--rep-ms", type=int, default=100) parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() + if ( + args.variant == "indexer-quant" or args.scenario == "indexer-8k1k" + ) and not _on_gfx950(): + parser.error("indexer-quant is only available on gfx950") torch.manual_seed(args.seed) - variants = ( - ("cache", "fp8-upconvert", "maybe-dequant") - if args.variant == "all" - else (args.variant,) - ) - scenarios = ( - SCENARIOS - if args.scenario == "all" - else {args.scenario: SCENARIOS[args.scenario]} - ) + if args.variant == "all": + variants = ["cache", "fp8-upconvert", "maybe-dequant"] + if _on_gfx950(): + variants.append("indexer-quant") + else: + variants = [args.variant] for variant in variants: + if variant == "indexer-quant": + if args.scenario not in ("all", "indexer-8k1k"): + continue + scenarios = {} + for num_batches in (1, 2, 3): + for num_tokens in INDEXER_TOKEN_COUNTS: + base, remainder = divmod(num_tokens, num_batches) + scenarios[f"b{num_batches}-{num_tokens}"] = [ + base + (req_id < remainder) for req_id in range(num_batches) + ] + else: + if args.scenario == "indexer-8k1k": + continue + scenarios = ( + SCENARIOS + if args.scenario == "all" + else {args.scenario: SCENARIOS[args.scenario]} + ) for name, seq_lens in scenarios.items(): run_scenario( variant, diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 79244947df15..7683f699c760 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -267,6 +267,9 @@ def _launch_gfx950_partial( extra_indices: torch.Tensor, extra_indptr: torch.Tensor, num_splits: int, + adaptive_splits: bool = False, + poison_scratch: bool = False, + one_wave_splits: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod @@ -287,6 +290,11 @@ def _launch_gfx950_partial( dtype=torch.float32, device=q.device, ) + if poison_scratch: + part_m.fill_(float("nan")) + part_l.fill_(float("nan")) + part_acc.fill_(float("nan")) + part_acc[:, 1::2].fill_(float("inf")) assert part_m.is_contiguous() assert part_l.is_contiguous() assert part_acc.is_contiguous() @@ -320,6 +328,8 @@ def _launch_gfx950_partial( IS_FNUZ_MAIN=current_platform.is_fp8_fnuz(), IS_FNUZ_EXTRA=False, TRUST_EXTRA_CACHE_NAN_FREE=True, + ADAPTIVE_SPLITS=adaptive_splits, + ONE_WAVE_SPLITS=one_wave_splits or num_splits, BLOCK_H=16, BLOCK_K=32, NUM_SPLITS=num_splits, @@ -330,6 +340,46 @@ def _launch_gfx950_partial( return part_m, part_l, part_acc +def _launch_sparse_decode_reduce( + part_m: torch.Tensor, + part_l: torch.Tensor, + part_acc: torch.Tensor, + adaptive_splits: bool, +) -> torch.Tensor: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + num_queries, num_splits, num_heads = part_m.shape + out = torch.empty( + (num_queries, num_heads, HEAD_DIM), + dtype=torch.bfloat16, + device=part_m.device, + ) + attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device) + mod._sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( + part_m, + part_l, + part_acc, + attn_sink, + out, + out.stride(0), + out.stride(1), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + num_heads, + HAS_ATTN_SINK=False, + ADAPTIVE_SPLITS=adaptive_splits, + COMB_DIM=HEAD_DIM, + BLOCK_H=1, + NUM_SPLITS=num_splits, + SPLITS_PAD=1 << (num_splits - 1).bit_length(), + num_warps=4, + ) + return out + + @torch.inference_mode() def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None: from vllm._aiter_ops import rocm_aiter_ops @@ -621,6 +671,73 @@ def test_sparse_attn_decode_trusted_extra_matches_legacy_scrub() -> None: assert torch.equal(trusted, torch.zeros_like(trusted)) +@pytest.mark.parametrize("on_gfx950", [False, True]) +@torch.inference_mode() +def test_rocm_ragged_graph_buffer_view_tracks_source_width( + monkeypatch, on_gfx950: bool +) -> None: + from vllm.models.deepseek_v4.amd import rocm as rocm_mod + + monkeypatch.setattr(rocm_mod, "_ON_GFX950", on_gfx950) + + indices_buffer = torch.full((16,), -1, dtype=torch.int32) + indptr_buffer = torch.full((3,), -1, dtype=torch.int32) + first_indices = torch.tensor([3, 5, 7], dtype=torch.int32) + first_indptr = torch.tensor([0, 1, 3], dtype=torch.int32) + first_view, first_indptr_view = rocm_mod._copy_ragged_to_graph_buffers( + first_indices, + first_indptr, + indices_buffer, + indptr_buffer, + num_rows=2, + max_entries_per_row=8, + ) + + second_indices = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32) + second_indptr = torch.tensor([0, 2, 6], dtype=torch.int32) + second_view, second_indptr_view = rocm_mod._copy_ragged_to_graph_buffers( + second_indices, + second_indptr, + indices_buffer, + indptr_buffer, + num_rows=2, + max_entries_per_row=8, + ) + + expected_first_entries = ( + first_indices.numel() if on_gfx950 else indices_buffer.numel() + ) + expected_second_entries = ( + second_indices.numel() if on_gfx950 else indices_buffer.numel() + ) + assert first_view.numel() == expected_first_entries + assert second_view.numel() == expected_second_entries + assert first_view.data_ptr() == second_view.data_ptr() == indices_buffer.data_ptr() + assert first_indptr_view.data_ptr() == second_indptr_view.data_ptr() + assert torch.equal(second_view[: second_indices.numel()], second_indices) + assert torch.equal(second_indptr_view, second_indptr) + + +def test_rocm_capture_metadata_sets_adaptive_marker(monkeypatch) -> None: + from vllm.models.deepseek_v4.amd import rocm as rocm_mod + from vllm.models.deepseek_v4.sparse_mla import ( + DeepseekV4SparseMLAMetadataBuilder, + ) + + metadata = SimpleNamespace(for_cudagraph_capture=False) + monkeypatch.setattr( + DeepseekV4SparseMLAMetadataBuilder, + "build_for_cudagraph_capture", + lambda *_: metadata, + ) + builder = object.__new__(rocm_mod.DeepseekV4ROCMAiterMLASparseMetadataBuilder) + + actual = builder.build_for_cudagraph_capture(SimpleNamespace()) + + assert actual is metadata + assert actual.for_cudagraph_capture is _on_gfx950() + + @requires_split_decode_arch @torch.inference_mode() def test_decode_num_splits_heuristic(monkeypatch) -> None: @@ -635,18 +752,18 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: # A tiny batch on a large device should split to add parallelism. assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1 - # Long C128A rows need 32 splits to fill a 256-CU device at low batch. - assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 32 - assert mod._decode_num_splits(8, 1, 128.0, 8192.0) == 32 + # The shared gfx942 selector retains its original 16-split ceiling. + assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16 + assert mod._decode_num_splits(8, 1, 128.0, 8192.0) == 16 assert mod._decode_num_splits(64, 1, 128.0, 8192.0) == 4 - # The chosen count always stays within the searched [1, 32] range, and a + # The chosen count always stays within the searched [1, 16] range, and a # zero-length workload never splits (no work to parallelize). for num_queries in (1, 4, 24, 224, 1024): splits = mod._decode_num_splits( num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0 ) - assert 1 <= splits <= 32 + assert 1 <= splits <= 16 assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 @@ -671,7 +788,7 @@ def test_decode_num_splits_gfx950(monkeypatch) -> None: @requires_split_decode_arch -@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8, 32]) +@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8]) @pytest.mark.parametrize("with_extra", [True, False]) @pytest.mark.parametrize("with_sink", [True, False]) @torch.inference_mode() @@ -848,6 +965,79 @@ def test_sparse_attn_decode_gfx950_partial_buffer_layout(num_splits: int) -> Non assert torch.equal(part_acc[1], torch.zeros_like(part_acc[1])) +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None: + device = torch.device("cuda") + torch.manual_seed(19) + block_size = 64 + num_heads = 16 + num_splits = 32 + main_use_fnuz = current_platform.is_fp8_fnuz() + q = torch.randn(3, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) + q *= 0.125 + main_kv = torch.randn(3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 + main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size, main_use_fnuz) + main_rows = [[0], [1], [2]] + main_indices, main_indptr = _ragged_from_rows(main_rows, device) + extra_kv = torch.randn(2112, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 + extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, use_fnuz=False) + extra_rows = [ + list(range(64)), + list(range(64, 576)), + list(range(576, 2112)), + ] + extra_indices, extra_indptr = _ragged_from_rows(extra_rows, device) + + part_m, part_l, part_acc = _launch_gfx950_partial( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + num_splits, + adaptive_splits=True, + poison_scratch=True, + one_wave_splits=16, + ) + + neg_large = torch.finfo(torch.float32).min + assert torch.equal(part_m[0, 4:], torch.full_like(part_m[0, 4:], neg_large)) + assert torch.equal(part_l[0, 4:], torch.zeros_like(part_l[0, 4:])) + assert torch.isnan(part_acc[0, 4::2]).all() + assert torch.isinf(part_acc[0, 5::2]).all() + assert (part_l[1, :16] > 0).all() + assert torch.equal(part_m[1, 16:], torch.full_like(part_m[1, 16:], neg_large)) + assert torch.equal(part_l[1, 16:], torch.zeros_like(part_l[1, 16:])) + assert torch.isnan(part_acc[1, 16::2]).all() + assert torch.isinf(part_acc[1, 17::2]).all() + assert (part_l[2] > 0).all() + assert torch.isfinite(part_acc[2]).all() + + actual = _launch_sparse_decode_reduce( + part_m, + part_l, + part_acc, + adaptive_splits=True, + ) + expected = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=main_rows, + scale=HEAD_DIM**-0.5, + attn_sink=None, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=extra_rows, + main_use_fnuz=main_use_fnuz, + ) + + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + @requires_gfx950 @pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65]) @torch.inference_mode() @@ -922,10 +1112,11 @@ def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None: device = torch.device("cuda") torch.manual_seed(17) block_size = 64 - num_queries = 8 + num_queries = 16 num_heads = 16 num_splits = 8 extra_per_query = 65 * num_splits + max_extra_per_query = 8192 q = ( torch.randn( num_queries, @@ -958,7 +1149,18 @@ def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None: for query_idx in range(num_queries) ] main_indices, main_indptr = _ragged_from_rows(main_rows, device) - extra_indices, extra_indptr = _ragged_from_rows(extra_rows, device) + short_extra_rows = [row[:64] for row in extra_rows] + long_indices, long_indptr = _ragged_from_rows(extra_rows, device) + short_indices, short_indptr = _ragged_from_rows(short_extra_rows, device) + extra_indices = torch.full( + (num_queries * max_extra_per_query,), + -1, + dtype=torch.int32, + device=device, + ) + extra_indices[: long_indices.numel()].copy_(long_indices) + extra_indptr = long_indptr.clone() + extra_indices_ptr = extra_indices.data_ptr() attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) out = torch.empty_like(q) @@ -979,6 +1181,7 @@ def run_decode() -> torch.Tensor: extra_indptr=extra_indptr, out=out, extra_cache_nan_free=True, + adaptive_splits=True, ) run_decode() @@ -986,18 +1189,26 @@ def run_decode() -> torch.Tensor: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): captured_out = run_decode() - captured = out.clone() + torch.accelerator.synchronize() + captured_long = out.clone() + expected_long = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=main_rows, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=extra_rows, + ) + torch.testing.assert_close(captured_long, expected_long, atol=2e-2, rtol=2e-2) - q.add_(0.03125) - replacement = extra_rows[0][-1] - extra_indices[0] = replacement - extra_rows[0][0] = replacement + extra_indices[: short_indices.numel()].copy_(short_indices) + extra_indptr.copy_(short_indptr) graph.replay() torch.accelerator.synchronize() - - assert captured_out.data_ptr() == out.data_ptr() - assert not torch.equal(out, captured) - expected = _ref_sparse_decode_ragged( + short_out = out.clone() + expected_short = _ref_sparse_decode_ragged( q=q, main_cache=main_cache, main_rows=main_rows, @@ -1005,9 +1216,20 @@ def run_decode() -> torch.Tensor: attn_sink=attn_sink, block_size=block_size, extra_cache=extra_cache, - extra_rows=extra_rows, + extra_rows=short_extra_rows, ) - torch.testing.assert_close(out, expected, atol=2e-2, rtol=2e-2) + + assert captured_out.data_ptr() == out.data_ptr() + assert extra_indices.data_ptr() == extra_indices_ptr + assert extra_indices.numel() == num_queries * max_extra_per_query + assert not torch.equal(short_out, captured_long) + torch.testing.assert_close(short_out, expected_short, atol=2e-2, rtol=2e-2) + + extra_indices[: long_indices.numel()].copy_(long_indices) + extra_indptr.copy_(long_indptr) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2) @requires_split_decode_arch @@ -1039,7 +1261,8 @@ def test_sparse_attn_decode_long_context(monkeypatch) -> None: scale = HEAD_DIM**-0.5 split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits" - monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: 32) + num_splits = 32 if _on_gfx950() else 16 + monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: num_splits) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, main_cache=main_cache, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 663fd5a16bdf..cfdec96e61d0 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -35,6 +35,10 @@ from vllm.v1.attention.backends.mla.compressor_utils import ( get_dspark_swa_index_width, ) +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + cp_gather_indexer_k_quant_cache_triton, + indexer_k_quant_and_cache_triton, +) from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 @@ -50,6 +54,55 @@ def _on_gfx950() -> bool: return False +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") +def test_cp_gather_despecialized_kernel_is_gfx950_only(monkeypatch): + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + class FakeKernel: + def __init__(self): + self.calls = [] + + def __getitem__(self, grid): + def launch(*args): + self.calls.append((grid, args)) + + return launch + + legacy_kernel = FakeKernel() + gfx950_kernel = FakeKernel() + monkeypatch.setattr(mod, "_cp_gather_indexer_quant_cache_kernel", legacy_kernel) + monkeypatch.setattr( + mod, + "_cp_gather_indexer_quant_cache_gfx950_kernel", + gfx950_kernel, + ) + + k_cache = torch.zeros((4, 1, 132), dtype=torch.uint8) + k_fp8 = torch.empty((5, 128), dtype=current_platform.fp8_dtype()) + k_scale = torch.empty((5, 4), dtype=torch.uint8) + block_table = torch.zeros((2, 7), dtype=torch.int32) + cu_seqlen = torch.tensor([0, 2, 5], dtype=torch.int32) + token_to_seq = torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32) + args = (k_cache, k_fp8, k_scale, block_table, cu_seqlen, token_to_seq) + + monkeypatch.setattr(mod, "_ON_GFX950", True) + mod.cp_gather_indexer_k_quant_cache_triton(*args) + assert len(gfx950_kernel.calls) == 1 + assert not legacy_kernel.calls + gfx950_grid, gfx950_args = gfx950_kernel.calls[0] + assert gfx950_grid == (5,) + assert len(gfx950_args) == 18 + assert gfx950_args[-3:] == (2, 7, 4) + + monkeypatch.setattr(mod, "_ON_GFX950", False) + mod.cp_gather_indexer_k_quant_cache_triton(*args) + assert len(legacy_kernel.calls) == 1 + legacy_grid, legacy_args = legacy_kernel.calls[0] + assert legacy_grid == (5,) + assert len(legacy_args) == 19 + assert legacy_args[-4:] == (5, 2, 7, 4) + + @pytest.mark.parametrize( ("window_size", "num_speculative_tokens", "expected"), [(128, 5, 192), (512, 5, 576), (1024, 0, 1024)], @@ -537,7 +590,8 @@ def test_indexer_gather_accepts_upper_bound_output(): valid_tokens = 9 upper_bound_tokens = 13 block_size = 16 - num_blocks = 2 + num_seqs = 3 + num_blocks = num_seqs sentinel = 123 device = "cuda" @@ -545,13 +599,15 @@ def test_indexer_gather_accepts_upper_bound_output(): kv_cache = torch.zeros( num_blocks, block_size, cache_stride, dtype=torch.uint8, device=device ) - slot_mapping = torch.arange(valid_tokens, dtype=torch.int64, device=device) + slot_mapping = torch.tensor( + [0, 1, 2, 16, 17, 18, 32, 33, 34], dtype=torch.int64, device=device + ) ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, quant_block_size, "ue8m0") block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze( - 0 + 1 ) - cu_seq_lens = torch.tensor([0, valid_tokens], dtype=torch.int32, device=device) + cu_seq_lens = torch.tensor([0, 3, 6, 9], dtype=torch.int32, device=device) dst_k = torch.full( (upper_bound_tokens, head_dim), sentinel, dtype=torch.uint8, device=device ) @@ -566,8 +622,52 @@ def test_indexer_gather_accepts_upper_bound_output(): ops.cp_gather_indexer_k_quant_cache( kv_cache, dst_k, dst_scale, block_table, cu_seq_lens ) + + if current_platform.is_rocm(): + triton_kv_cache = torch.zeros_like(kv_cache) + indexer_k_quant_and_cache_triton( + k, + triton_kv_cache, + slot_mapping, + quant_block_size, + "ue8m0", + ) + triton_dst_k = torch.full_like(dst_k, sentinel) + triton_dst_scale = torch.full_like(dst_scale, sentinel) + token_to_seq = torch.cat( + ( + torch.repeat_interleave( + torch.arange(num_seqs, dtype=torch.int32, device=device), 3 + ), + torch.full( + (upper_bound_tokens - valid_tokens,), + -1, + dtype=torch.int32, + device=device, + ), + ) + ) + cp_gather_indexer_k_quant_cache_triton( + triton_kv_cache, + triton_dst_k.view(current_platform.fp8_dtype()), + triton_dst_scale, + block_table, + cu_seq_lens, + token_to_seq, + ) torch.accelerator.synchronize() + if current_platform.is_rocm(): + triton_recovered = triton_dst_k[:valid_tokens].view( + current_platform.fp8_dtype() + ).float() * triton_dst_scale[:valid_tokens].view(torch.float32) + triton_error = (triton_recovered - k.float()).abs().amax(dim=1) + max_triton_error = ( + 16.0 * triton_dst_scale[:valid_tokens].view(torch.float32).flatten() + ) + assert torch.all(triton_error <= max_triton_error) + assert torch.all(triton_dst_k[valid_tokens:] == sentinel) + assert torch.all(triton_dst_scale[valid_tokens:] == sentinel) k_recovered = dst_k[:valid_tokens].view(torch.float8_e4m3fn).float() * dst_scale[ :valid_tokens ].view(torch.float32) diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 09eb0ece6fe0..5e7609bd1191 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -308,9 +308,13 @@ def _copy_ragged_to_graph_buffers( max_entries = max(num_rows * max_entries_per_row, 1) ragged_out = ragged_indices_buffer[:max_entries] - nnz = ragged_indices.numel() - if nnz > 0: - ragged_out[:nnz].copy_(ragged_indices, non_blocking=True) + source_entries = ragged_indices.numel() + if source_entries > 0: + ragged_out[:source_entries].copy_(ragged_indices, non_blocking=True) + if _ON_GFX950: + # Preserve the graph-stable base pointer while exposing source capacity + # to the sync-free split selector; indptr still carries the true NNZ. + ragged_out = ragged_out[: max(source_entries, 1)] return ragged_out, indptr_out @@ -320,6 +324,7 @@ class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata): c128a_decode_topk_ragged_indices: torch.Tensor | None = None c128a_decode_topk_ragged_indptr: torch.Tensor | None = None + for_cudagraph_capture: bool = False @dataclass @@ -384,6 +389,16 @@ def build( c128a_decode_topk_ragged_indptr=ragged_indptr, ) + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ) -> DeepseekV4ROCMAiterMLASparseMetadata: + metadata = cast( + DeepseekV4ROCMAiterMLASparseMetadata, + super().build_for_cudagraph_capture(common_attn_metadata), + ) + metadata.for_cudagraph_capture = _ON_GFX950 + return metadata + class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): # Keep fused multi-step decode disabled until update_draft_decode_metadata() @@ -621,6 +636,13 @@ def forward_mqa( attn_metadata=rocm_metadata, swa_only=swa_only, output=output[:num_decode_tokens], + adaptive_splits=( + _ON_GFX950 + and not swa_only + and self.compress_ratio == 128 + and rocm_metadata is not None + and rocm_metadata.for_cudagraph_capture + ), ) def _forward_decode( @@ -631,6 +653,7 @@ def _forward_decode( attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, swa_only: bool, output: torch.Tensor, + adaptive_splits: bool, ) -> None: num_decodes = swa_metadata.num_decodes num_decode_tokens = swa_metadata.num_decode_tokens @@ -682,6 +705,7 @@ def _forward_decode( nope_head_dim=self.nope_head_dim, rope_head_dim=self.rope_head_dim, output=output, + adaptive_splits=adaptive_splits, extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free( self.kv_cache_dtype, self._has_kv_transfer, diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index eb07d5fd21b3..8a5de5d8b4b6 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -214,6 +214,87 @@ def _cp_gather_indexer_quant_cache_kernel( tl.store(dst_k_ptr + offset, val, mask=valid_block) +@triton.jit(do_not_specialize=["num_batches"]) +def _cp_gather_indexer_quant_cache_gfx950_kernel( + kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B] + # [n_blks, blk_size, head_dim] + kv_cache_scale_ptr, # [n_blks, blk_size] + k_fp8_ptr, # [num_tokens, head_dim] + k_scale_ptr, # [num_tokens] + block_table_ptr, # [batch_size, block_table_stride] + cu_seqlen_ptr, # [batch_size + 1] + token_to_seq_ptr, # [num_tokens] + block_size, + block_table_stride, + kv_cache_stride, + kv_cache_scale_stride, + LAYOUT: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_TILE_SIZE: tl.constexpr, + HEAD_TILE_SIZE: tl.constexpr, + num_batches, + BLOCK_TABLE_WIDTH: tl.constexpr, + NUM_BLOCKS: tl.constexpr, +): + tid = tl.program_id(0) + offset = tl.arange(0, HEAD_DIM) + batch_id = tl.load(token_to_seq_ptr + tid) + valid_batch = (batch_id >= 0) & (batch_id < num_batches) + safe_batch_id = tl.where(valid_batch, batch_id, 0) + batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0) + batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0) + batch_offset = tid - batch_start + valid_token = valid_batch & (tid >= batch_start) & (tid < batch_end) + if not valid_token: + return + block_table_id = batch_offset // block_size + block_offset = batch_offset % block_size + valid_block_table = ( + valid_token + & (block_table_id >= 0) + & (block_table_id < BLOCK_TABLE_WIDTH) + & (block_offset >= 0) + & (block_offset < block_size) + ) + safe_block_table_id = tl.where(valid_block_table, block_table_id, 0) + block_table_offset = safe_batch_id * block_table_stride + safe_block_table_id + block_id = tl.load( + block_table_ptr + block_table_offset, mask=valid_block_table, other=-1 + ) + valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS) + # The packed KV layout makes per-block strides large + # enough that block_id * stride can exceed 32-bit range. + safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64) + safe_block_offset = tl.where(valid_block, block_offset, 0) + tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE + if LAYOUT == "SHUFFLE": + src_cache_offset = ( + safe_block_id * kv_cache_stride + + (safe_block_offset // BLOCK_TILE_SIZE) * HEAD_DIM * BLOCK_TILE_SIZE + + tiled_block_offset * HEAD_TILE_SIZE + ) + else: + src_cache_offset = ( + safe_block_id * kv_cache_stride + safe_block_offset * HEAD_DIM + ) + src_scale_offset = safe_block_id * kv_cache_scale_stride + safe_block_offset + dst_offset = tid * HEAD_DIM + src_scale_ptr = kv_cache_scale_ptr + src_scale_offset + src_cache_ptr = kv_cache_ptr + src_cache_offset + dst_k_ptr = k_fp8_ptr + dst_offset + scale_val = tl.load(src_scale_ptr, mask=valid_block, other=0.0) + tl.store(k_scale_ptr + tid, scale_val) + if LAYOUT == "SHUFFLE": + tiled_src_offset = ( + offset // HEAD_TILE_SIZE * HEAD_TILE_SIZE * BLOCK_TILE_SIZE + + offset % HEAD_TILE_SIZE + ) + else: + tiled_src_offset = offset + val = tl.load(src_cache_ptr + tiled_src_offset) + tl.store(dst_k_ptr + offset, val, mask=valid_block) + + def cp_gather_indexer_k_quant_cache_triton( k_cache: torch.Tensor, # [num_blocks, block_size, head_dim + 4] k_fp8: torch.Tensor, @@ -237,7 +318,7 @@ def cp_gather_indexer_k_quant_cache_triton( grid = (num_tokens,) k_fp8_scale = k_fp8_scale.view(torch.float32) layout = "NORMAL" if block_size == 1 else "SHUFFLE" - _cp_gather_indexer_quant_cache_kernel[grid]( + kernel_args = ( k_cache_value, k_cache_scale, k_fp8, @@ -253,11 +334,22 @@ def cp_gather_indexer_k_quant_cache_triton( head_dim, block_tile_size, head_tile_size, - num_tokens, - cu_seqlen.shape[0] - 1, - block_table.shape[1], - num_blocks, ) + if _ON_GFX950: + _cp_gather_indexer_quant_cache_gfx950_kernel[grid]( + *kernel_args, + cu_seqlen.shape[0] - 1, + block_table.shape[1], + num_blocks, + ) + else: + _cp_gather_indexer_quant_cache_kernel[grid]( + *kernel_args, + num_tokens, + cu_seqlen.shape[0] - 1, + block_table.shape[1], + num_blocks, + ) # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L156 @@ -1321,8 +1413,8 @@ def _sparse_attn_decode_ragged_kernel( extra_cache_stride0, main_num_rows, extra_num_rows, - MAIN_BLOCK_SIZE: tl.constexpr, - EXTRA_BLOCK_SIZE: tl.constexpr, + main_block_size, + extra_block_size, scale, num_heads, HAS_ATTN_SINK: tl.constexpr, @@ -1379,11 +1471,11 @@ def _sparse_attn_decode_ragged_kernel( valid = in_range & (slot >= 0) & (slot < main_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // MAIN_BLOCK_SIZE - pos_in_block = safe_slot % MAIN_BLOCK_SIZE + block_idx = safe_slot // main_block_size + pos_in_block = safe_slot % main_block_size cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 + token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 x_uint8 = tl.load( token_data_ptr[:, None] + nope_offsets[None, :], @@ -1399,7 +1491,7 @@ def _sparse_attn_decode_ragged_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1443,14 +1535,14 @@ def _sparse_attn_decode_ragged_kernel( valid = in_range & (slot >= 0) & (slot < extra_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // EXTRA_BLOCK_SIZE - pos_in_block = safe_slot % EXTRA_BLOCK_SIZE + block_idx = safe_slot // extra_block_size + pos_in_block = safe_slot % extra_block_size cache_block_ptr = ( extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 ) token_data_ptr = cache_block_ptr + pos_in_block * 576 token_scale_ptr = ( - cache_block_ptr + EXTRA_BLOCK_SIZE * 576 + pos_in_block * 8 + cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 ) x_uint8 = tl.load( @@ -1467,7 +1559,7 @@ def _sparse_attn_decode_ragged_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1561,8 +1653,8 @@ def _sparse_attn_decode_partial_kernel( pa_stride_h, main_num_rows, extra_num_rows, - MAIN_BLOCK_SIZE: tl.constexpr, - EXTRA_BLOCK_SIZE: tl.constexpr, + main_block_size, + extra_block_size, scale, num_heads, HAS_EXTRA: tl.constexpr, @@ -1629,11 +1721,11 @@ def _sparse_attn_decode_partial_kernel( valid = in_range & (slot >= 0) & (slot < main_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // MAIN_BLOCK_SIZE - pos_in_block = safe_slot % MAIN_BLOCK_SIZE + block_idx = safe_slot // main_block_size + pos_in_block = safe_slot % main_block_size cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 + token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 x_uint8 = tl.load( token_data_ptr[:, None] + nope_offsets[None, :], @@ -1649,7 +1741,7 @@ def _sparse_attn_decode_partial_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -1696,14 +1788,14 @@ def _sparse_attn_decode_partial_kernel( valid = in_range & (slot >= 0) & (slot < extra_num_rows) safe_slot = tl.where(valid, slot, 0) - block_idx = safe_slot // EXTRA_BLOCK_SIZE - pos_in_block = safe_slot % EXTRA_BLOCK_SIZE + block_idx = safe_slot // extra_block_size + pos_in_block = safe_slot % extra_block_size cache_block_ptr = ( extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 ) token_data_ptr = cache_block_ptr + pos_in_block * 576 token_scale_ptr = ( - cache_block_ptr + EXTRA_BLOCK_SIZE * 576 + pos_in_block * 8 + cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 ) x_uint8 = tl.load( @@ -1720,7 +1812,7 @@ def _sparse_attn_decode_partial_kernel( mask=valid[:, None] & nope_mask[None, :], other=127, ) - scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) @@ -2013,6 +2105,8 @@ def _sparse_attn_decode_gfx950_partial_kernel( IS_FNUZ_MAIN: tl.constexpr, IS_FNUZ_EXTRA: tl.constexpr, TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, + ADAPTIVE_SPLITS: tl.constexpr, + ONE_WAVE_SPLITS: tl.constexpr, BLOCK_H: tl.constexpr, BLOCK_K: tl.constexpr, NUM_SPLITS: tl.constexpr, @@ -2024,11 +2118,44 @@ def _sparse_attn_decode_gfx950_partial_kernel( tl.static_assert(NOPE_DIM == 448) tl.static_assert(ROPE_DIM == 64) + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) if num_heads % BLOCK_H == 0: head_mask = tl.full((BLOCK_H,), True, tl.int1) else: head_mask = head_offsets < num_heads + neg_large = -3.4028234663852886e38 + + if ADAPTIVE_SPLITS: + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + if HAS_EXTRA: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + else: + extra_start = 0 + extra_len = 0 + split4_span: tl.constexpr = 4 * BLOCK_K + split4_iters = (main_len + split4_span - 1) // split4_span + split4_iters += (extra_len + split4_span - 1) // split4_span + use_four_splits = split4_iters <= 3 + work_splits = NUM_SPLITS + if ONE_WAVE_SPLITS > 4 and ONE_WAVE_SPLITS < NUM_SPLITS: + one_wave_span: tl.constexpr = ONE_WAVE_SPLITS * BLOCK_K + one_wave_iters = (main_len + one_wave_span - 1) // one_wave_span + one_wave_iters += (extra_len + one_wave_span - 1) // one_wave_span + work_splits = tl.where(one_wave_iters <= 3, ONE_WAVE_SPLITS, work_splits) + work_splits = tl.where(use_four_splits, 4, work_splits) + if split_id >= work_splits: + pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets + tl.store(part_m_ptr + pm_base, neg_large, mask=head_mask) + tl.store(part_l_ptr + pm_base, 0.0, mask=head_mask) + return + else: + work_splits = NUM_SPLITS + nope_offsets_0a = tl.arange(0, 128) nope_offsets_0b = 128 + tl.arange(0, 128) nope_offsets_0 = tl.arange(0, 256) @@ -2049,7 +2176,6 @@ def _sparse_attn_decode_gfx950_partial_kernel( ) q_combined = tl.cat(q_nope_0, q_tail, dim=1) - neg_large = -3.4028234663852886e38 m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) acc_nope_0a = tl.zeros((BLOCK_H, 128), dtype=tl.float32) @@ -2058,10 +2184,11 @@ def _sparse_attn_decode_gfx950_partial_kernel( acc_tail = tl.zeros((BLOCK_H, 128), dtype=tl.float32) k_offsets = tl.arange(0, BLOCK_K) - main_start = tl.load(main_indptr_ptr + query_idx) - main_end = tl.load(main_indptr_ptr + query_idx + 1) - main_len = main_end - main_start - main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS + if not ADAPTIVE_SPLITS: + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + main_chunk = (main_len + work_splits - 1) // work_splits main_lo = split_id * main_chunk main_hi = tl.minimum(main_lo + main_chunk, main_len) @@ -2147,10 +2274,11 @@ def _sparse_attn_decode_gfx950_partial_kernel( l_i = l_new if HAS_EXTRA: - extra_start = tl.load(extra_indptr_ptr + query_idx) - extra_end = tl.load(extra_indptr_ptr + query_idx + 1) - extra_len = extra_end - extra_start - extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS + if not ADAPTIVE_SPLITS: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + extra_chunk = (extra_len + work_splits - 1) // work_splits extra_lo = split_id * extra_chunk extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len) @@ -2333,6 +2461,7 @@ def _sparse_attn_decode_reduce_kernel( pa_stride_h, num_heads, HAS_ATTN_SINK: tl.constexpr, + ADAPTIVE_SPLITS: tl.constexpr, COMB_DIM: tl.constexpr, BLOCK_H: tl.constexpr, NUM_SPLITS: tl.constexpr, @@ -2401,17 +2530,27 @@ def _sparse_attn_decode_reduce_kernel( other=neg_large, ) w_s = tl.exp(m_s - m_final) + if ADAPTIVE_SPLITS: + active_split = m_s > neg_large + w_s = tl.where(head_mask & active_split, w_s, 0.0) acc_base = ( part_acc_ptr + query_idx * pa_stride0 + s * pa_stride_s + head_offsets[:, None] * pa_stride_h ) - acc_s = tl.load( - acc_base + comb_offsets[None, :], - mask=head_mask[:, None], - other=0.0, - ) + if ADAPTIVE_SPLITS: + acc_s = tl.load( + acc_base + comb_offsets[None, :], + mask=head_mask[:, None] & active_split[:, None], + other=0.0, + ) + else: + acc_s = tl.load( + acc_base + comb_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) acc += w_s[:, None] * acc_s out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0) @@ -2596,9 +2735,8 @@ def _decode_num_splits( mu = 0.04 best_splits = 1 best_cost = None - # A full C128A row can contain 8K KV tokens. At low batch sizes, 32 splits - # keep those long rows to one device wave and halve the partial iterations. - for splits in range(1, 33): + # Search up to 16 splits; beyond that the reduce/HBM overhead dominates. + for splits in range(1, 17): waves = (base * splits + cu - 1) // cu cost = waves * (1.0 / splits + mu) if best_cost is None or cost < best_cost - 1e-9: @@ -2676,6 +2814,7 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indptr: torch.Tensor | None = None, out: torch.Tensor | None = None, extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> torch.Tensor: assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert main_cache.ndim == 3, ( @@ -2815,6 +2954,16 @@ def _rocm_sparse_attn_decode_ragged_triton( num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k ) + base_workgroups = num_queries * heads_blocks + adaptive_splits = ( + _ON_GFX950 and adaptive_splits and base_workgroups >= 16 and num_splits > 4 + ) + one_wave_splits = ( + max(1, _decode_cu_count() // base_workgroups) + if adaptive_splits and 16 <= base_workgroups < 64 + else num_splits + ) + part_m = torch.empty( (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device ) @@ -2855,6 +3004,8 @@ def _rocm_sparse_attn_decode_ragged_triton( IS_FNUZ_MAIN=is_fnuz, IS_FNUZ_EXTRA=False, TRUST_EXTRA_CACHE_NAN_FREE=extra_cache_nan_free, + ADAPTIVE_SPLITS=adaptive_splits, + ONE_WAVE_SPLITS=one_wave_splits, BLOCK_H=block_h, BLOCK_K=block_k, NUM_SPLITS=num_splits, @@ -2921,6 +3072,7 @@ def _rocm_sparse_attn_decode_ragged_triton( part_acc.stride(2), num_heads, HAS_ATTN_SINK=has_attn_sink, + ADAPTIVE_SPLITS=adaptive_splits, COMB_DIM=comb_dim, BLOCK_H=1, NUM_SPLITS=num_splits, @@ -2948,6 +3100,7 @@ def _rocm_sparse_attn_decode_triton( extra_ragged_indptr: torch.Tensor | None = None, out: torch.Tensor | None = None, extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> torch.Tensor: if main_ragged_indices is None or main_ragged_indptr is None: main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense( @@ -2985,6 +3138,7 @@ def _rocm_sparse_attn_decode_triton( extra_indptr=extra_ragged_indptr, out=out, extra_cache_nan_free=extra_cache_nan_free, + adaptive_splits=adaptive_splits, ) @@ -3057,6 +3211,7 @@ def rocm_sparse_attn_decode( rope_head_dim: int, output: torch.Tensor, extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> None: assert swa_k_cache.dtype == torch.uint8, ( "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, " @@ -3086,7 +3241,7 @@ def rocm_sparse_attn_decode( if topk_indices is not None: extra_indices = topk_indices.reshape(topk_indices.shape[0], -1) - direct_out = output if output.dtype == torch.bfloat16 else None + direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None attn_out = _rocm_sparse_attn_decode_triton( q=q, main_cache=swa_k_cache, @@ -3105,6 +3260,7 @@ def rocm_sparse_attn_decode( extra_ragged_indptr=topk_ragged_indptr, out=direct_out, extra_cache_nan_free=extra_cache_nan_free, + adaptive_splits=adaptive_splits, ) if direct_out is None: output.copy_(attn_out.to(output.dtype)) From d7d33e576ce9a955b1d2bf2e54e74cfe5b801662 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai Date: Sun, 16 Aug 2026 01:44:08 +0000 Subject: [PATCH 7/8] Simplify gfx950 sparse decode change Remove the one-off benchmark extension and redundant private-layout tests. Reuse the loaded-tile helper in the gfx950 partial kernel while preserving graph replay and stale-scratch coverage. Assisted-by: OpenAI Codex Signed-off-by: Fangzhou Ai --- benchmarks/kernels/benchmark_cp_gather.py | 124 +----- .../attention/test_rocm_triton_attn_dsv4.py | 381 +----------------- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 296 +++----------- 3 files changed, 90 insertions(+), 711 deletions(-) diff --git a/benchmarks/kernels/benchmark_cp_gather.py b/benchmarks/kernels/benchmark_cp_gather.py index 903e7756267a..5e73d6d6e95c 100644 --- a/benchmarks/kernels/benchmark_cp_gather.py +++ b/benchmarks/kernels/benchmark_cp_gather.py @@ -7,12 +7,8 @@ import torch from vllm import _custom_ops as ops -from vllm.platforms import current_platform from vllm.triton_utils import triton from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - cp_gather_indexer_k_quant_cache_triton, -) SCENARIOS = { "single-60k": [60_000], @@ -35,15 +31,6 @@ "bfloat16": torch.bfloat16, "float32": torch.float32, } -INDEXER_TOKEN_COUNTS = (1638, 1737, 1738, 1739, 2029, 2031, 3217, 4078) - - -def _on_gfx950() -> bool: - if not current_platform.is_rocm(): - return False - from vllm.platforms.rocm import _ON_GFX950 - - return _ON_GFX950 def make_page_table( @@ -185,65 +172,6 @@ def run() -> None: return run, bytes_moved -def make_indexer_quant_gather( - seq_lens: list[int], - block_size: int, -) -> tuple[Callable[[], None], int]: - head_dim = 128 - cache_entry_bytes = head_dim + 4 - max_blocks_per_seq = 4096 - num_cache_blocks = 8192 - block_table = torch.full( - (len(seq_lens), max_blocks_per_seq), - -1, - dtype=torch.int32, - device="cuda", - ) - next_block = 0 - for req_id, seq_len in enumerate(seq_lens): - req_blocks = math.ceil(seq_len / block_size) - block_table[req_id, :req_blocks] = torch.arange( - next_block, - next_block + req_blocks, - dtype=torch.int32, - device="cuda", - ) - next_block += req_blocks - assert next_block <= num_cache_blocks - - cu_seq_lens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device="cuda") - cu_seq_lens[1:] = torch.tensor(seq_lens, dtype=torch.int32, device="cuda").cumsum( - dim=0 - ) - token_to_seq = torch.repeat_interleave( - torch.arange(len(seq_lens), dtype=torch.int32, device="cuda"), - torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), - ) - kv_cache = torch.zeros( - (num_cache_blocks, block_size, cache_entry_bytes), - dtype=torch.uint8, - device="cuda", - ) - num_tokens = sum(seq_lens) - dst_k = torch.empty( - (num_tokens, head_dim), dtype=current_platform.fp8_dtype(), device="cuda" - ) - dst_scale = torch.empty((num_tokens, 4), dtype=torch.uint8, device="cuda") - - def run() -> None: - cp_gather_indexer_k_quant_cache_triton( - kv_cache, - dst_k, - dst_scale, - block_table, - cu_seq_lens, - token_to_seq, - ) - - bytes_moved = 2 * num_tokens * cache_entry_bytes - return run, bytes_moved - - @torch.inference_mode() def run_scenario( variant: str, @@ -259,21 +187,17 @@ def run_scenario( run, bytes_moved = make_cache_gather(seq_lens, block_size, entry_size, dtype) elif variant == "fp8-upconvert": run, bytes_moved = make_fp8_upconvert(seq_lens, block_size) - elif variant == "indexer-quant": - run, bytes_moved = make_indexer_quant_gather(seq_lens, block_size) else: run, bytes_moved = make_maybe_dequant_gather(seq_lens, block_size, entry_size) - latency_ms, p99_ms, p999_ms = triton.testing.do_bench( - run, warmup=warmup_ms, rep=rep_ms, quantiles=[0.5, 0.99, 0.999] + latency_ms = triton.testing.do_bench( + run, warmup=warmup_ms, rep=rep_ms, return_mode="median" ) bandwidth_gbps = bytes_moved / latency_ms / 1e6 lengths = ",".join(str(seq_len) for seq_len in seq_lens) print( f"{variant:15s} {name:10s} batch={len(seq_lens):2d} " f"total={sum(seq_lens):7d} latency={latency_ms * 1e3:9.2f} us " - f"p99={p99_ms * 1e3:9.2f} us " - f"p99.9={p999_ms * 1e3:9.2f} us " f"bandwidth={bandwidth_gbps:8.1f} GB/s lengths=[{lengths}]" ) @@ -282,12 +206,10 @@ def main() -> None: parser = FlexibleArgumentParser(description="Benchmark cp_gather variants") parser.add_argument( "--variant", - choices=["all", "cache", "fp8-upconvert", "maybe-dequant", "indexer-quant"], + choices=["all", "cache", "fp8-upconvert", "maybe-dequant"], default="all", ) - parser.add_argument( - "--scenario", choices=["all", *SCENARIOS, "indexer-8k1k"], default="all" - ) + parser.add_argument("--scenario", choices=["all", *SCENARIOS], default="all") parser.add_argument("--dtype", choices=DTYPES, default="fp8") parser.add_argument("--block-size", type=int, default=64) parser.add_argument("--entry-size", type=int, default=576) @@ -295,37 +217,19 @@ def main() -> None: parser.add_argument("--rep-ms", type=int, default=100) parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() - if ( - args.variant == "indexer-quant" or args.scenario == "indexer-8k1k" - ) and not _on_gfx950(): - parser.error("indexer-quant is only available on gfx950") torch.manual_seed(args.seed) - if args.variant == "all": - variants = ["cache", "fp8-upconvert", "maybe-dequant"] - if _on_gfx950(): - variants.append("indexer-quant") - else: - variants = [args.variant] + variants = ( + ("cache", "fp8-upconvert", "maybe-dequant") + if args.variant == "all" + else (args.variant,) + ) + scenarios = ( + SCENARIOS + if args.scenario == "all" + else {args.scenario: SCENARIOS[args.scenario]} + ) for variant in variants: - if variant == "indexer-quant": - if args.scenario not in ("all", "indexer-8k1k"): - continue - scenarios = {} - for num_batches in (1, 2, 3): - for num_tokens in INDEXER_TOKEN_COUNTS: - base, remainder = divmod(num_tokens, num_batches) - scenarios[f"b{num_batches}-{num_tokens}"] = [ - base + (req_id < remainder) for req_id in range(num_batches) - ] - else: - if args.scenario == "indexer-8k1k": - continue - scenarios = ( - SCENARIOS - if args.scenario == "all" - else {args.scenario: SCENARIOS[args.scenario]} - ) for name, seq_lens in scenarios.items(): run_scenario( variant, diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 7683f699c760..7a98f3ec7005 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -149,20 +149,6 @@ def _poison_fp8_ds_mla_cache_row( ) -def _canonicalize_poisoned_fp8_ds_mla_cache_row( - cache: torch.Tensor, block_size: int, slot: int = 0 -) -> None: - flat = cache.flatten() - block_idx = slot // block_size - pos = slot % block_size - block_base = block_idx * cache.stride(0) - token_base = block_base + pos * 576 - scale_base = block_base + block_size * 576 + pos * 8 - flat[token_base] = 0 - flat[scale_base : scale_base + 7] = 254 - flat[token_base + NOPE_HEAD_DIM : token_base + 576].view(torch.bfloat16)[0] = 0 - - def _read_fp8_ds_mla_cache_rows( cache: torch.Tensor, slots: torch.Tensor, @@ -258,88 +244,6 @@ def _ragged_from_rows( ) -def _launch_gfx950_partial( - q: torch.Tensor, - main_cache: torch.Tensor, - main_indices: torch.Tensor, - main_indptr: torch.Tensor, - extra_cache: torch.Tensor, - extra_indices: torch.Tensor, - extra_indptr: torch.Tensor, - num_splits: int, - adaptive_splits: bool = False, - poison_scratch: bool = False, - one_wave_splits: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod - - num_queries, num_heads, _ = q.shape - part_m = torch.empty( - num_queries, - num_splits, - num_heads, - dtype=torch.float32, - device=q.device, - ) - part_l = torch.empty_like(part_m) - part_acc = torch.empty( - num_queries, - num_splits, - num_heads, - HEAD_DIM, - dtype=torch.float32, - device=q.device, - ) - if poison_scratch: - part_m.fill_(float("nan")) - part_l.fill_(float("nan")) - part_acc.fill_(float("nan")) - part_acc[:, 1::2].fill_(float("inf")) - assert part_m.is_contiguous() - assert part_l.is_contiguous() - assert part_acc.is_contiguous() - - mod._sparse_attn_decode_gfx950_partial_kernel[ - (num_queries, num_splits, (num_heads + 15) // 16) - ]( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - part_m, - part_l, - part_acc, - q.stride(0), - q.stride(1), - main_cache.stride(0), - extra_cache.stride(0), - main_cache.shape[0] * main_cache.shape[1], - extra_cache.shape[0] * extra_cache.shape[1], - main_cache.shape[1], - extra_cache.shape[1], - HEAD_DIM**-0.5, - num_heads, - HAS_EXTRA=True, - NOPE_DIM=NOPE_HEAD_DIM, - ROPE_DIM=ROPE_HEAD_DIM, - IS_FNUZ_MAIN=current_platform.is_fp8_fnuz(), - IS_FNUZ_EXTRA=False, - TRUST_EXTRA_CACHE_NAN_FREE=True, - ADAPTIVE_SPLITS=adaptive_splits, - ONE_WAVE_SPLITS=one_wave_splits or num_splits, - BLOCK_H=16, - BLOCK_K=32, - NUM_SPLITS=num_splits, - NUM_STAGES=1, - num_warps=4, - waves_per_eu=0, - ) - return part_m, part_l, part_acc - - def _launch_sparse_decode_reduce( part_m: torch.Tensor, part_l: torch.Tensor, @@ -626,51 +530,6 @@ def test_sparse_attn_decode_scrubs_untrusted_cache_by_default() -> None: assert torch.equal(actual, torch.zeros_like(actual)) -@requires_gfx950 -@torch.inference_mode() -def test_sparse_attn_decode_trusted_extra_matches_legacy_scrub() -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _rocm_sparse_attn_decode_ragged_triton, - ) - - device = torch.device("cuda") - block_size = 4 - main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) - legacy_extra = torch.zeros_like(main_cache) - _poison_fp8_ds_mla_cache_row(main_cache, block_size) - _poison_fp8_ds_mla_cache_row(legacy_extra, block_size) - canonical_extra = legacy_extra.clone() - _canonicalize_poisoned_fp8_ds_mla_cache_row(canonical_extra, block_size) - indices = torch.zeros(1, dtype=torch.int32, device=device) - indptr = torch.tensor([0, 1], dtype=torch.int32, device=device) - q = torch.ones(1, 1, HEAD_DIM, dtype=torch.bfloat16, device=device) - kwargs = { - "q": q, - "main_cache": main_cache, - "main_indices": indices, - "main_indptr": indptr, - "scale": HEAD_DIM**-0.5, - "attn_sink": None, - "nope_head_dim": NOPE_HEAD_DIM, - "rope_head_dim": ROPE_HEAD_DIM, - "extra_indices": indices, - "extra_indptr": indptr, - } - - legacy = _rocm_sparse_attn_decode_ragged_triton( - **kwargs, - extra_cache=legacy_extra, - ) - trusted = _rocm_sparse_attn_decode_ragged_triton( - **kwargs, - extra_cache=canonical_extra, - extra_cache_nan_free=True, - ) - - assert torch.equal(trusted, legacy) - assert torch.equal(trusted, torch.zeros_like(trusted)) - - @pytest.mark.parametrize("on_gfx950", [False, True]) @torch.inference_mode() def test_rocm_ragged_graph_buffer_view_tracks_source_width( @@ -754,8 +613,6 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: # The shared gfx942 selector retains its original 16-split ceiling. assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16 - assert mod._decode_num_splits(8, 1, 128.0, 8192.0) == 16 - assert mod._decode_num_splits(64, 1, 128.0, 8192.0) == 4 # The chosen count always stays within the searched [1, 16] range, and a # zero-length workload never splits (no work to parallelize). @@ -772,18 +629,8 @@ def test_decode_num_splits_gfx950(monkeypatch) -> None: from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256) - for num_queries in (1, 8): - assert mod._decode_gfx950_num_splits(num_queries, 1, 128, 32) == 32 - assert mod._decode_gfx950_num_splits(num_queries, 1, 128, 7812) == 32 - + assert mod._decode_gfx950_num_splits(1, 1, 128, 8192) == 32 assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 4 - assert mod._decode_gfx950_num_splits(16, 1, 128, 781) == 13 - assert mod._decode_gfx950_num_splits(48, 1, 128, 3906) == 10 - for extra_rows in (32, 256): - assert mod._decode_gfx950_num_splits(64, 1, 128, extra_rows) == 4 - assert mod._decode_gfx950_num_splits(64, 1, 128, 781) == 7 - for extra_rows in (3906, 7812): - assert mod._decode_gfx950_num_splits(64, 1, 128, extra_rows) == 8 assert mod._decode_gfx950_num_splits(512, 1, 128, 7812) == 1 @@ -843,7 +690,13 @@ def test_sparse_attn_decode_split_k_kernel( # Pin the split count so each parametrized value is exercised deterministically. split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits" + other_split_fn = ( + "_decode_num_splits" if _on_gfx950() else "_decode_gfx950_num_splits" + ) monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: num_splits) + monkeypatch.setattr( + mod, other_split_fn, lambda *args, **kwargs: pytest.fail("wrong selector") + ) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, @@ -873,169 +726,22 @@ def test_sparse_attn_decode_split_k_kernel( torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) -@requires_gfx950 -@pytest.mark.parametrize("num_splits", [1, 2, 3, 8, 32]) -@torch.inference_mode() -def test_sparse_attn_decode_gfx950_partial_buffer_layout(num_splits: int) -> None: - device = torch.device("cuda") - block_size = 64 - num_heads = 16 - entries_per_split = 65 - num_extra = num_splits * entries_per_split - q = torch.full( - (2, num_heads, HEAD_DIM), - 0.125, - dtype=torch.bfloat16, - device=device, - ) - main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) - main_indices = torch.empty(0, dtype=torch.int32, device=device) - main_indptr = torch.zeros(3, dtype=torch.int32, device=device) - extra_cache = _pack_fp8_ds_mla_cache( - torch.full( - (num_extra, HEAD_DIM), - 0.125, - dtype=torch.bfloat16, - device=device, - ), - block_size, - use_fnuz=False, - ) - extra_indices = torch.arange(num_extra, dtype=torch.int32, device=device) - extra_indptr = torch.tensor( - [0, num_extra, num_extra], dtype=torch.int32, device=device - ) - - batch = _launch_gfx950_partial( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - num_splits, - ) - single_nonempty = _launch_gfx950_partial( - q[:1], - main_cache, - main_indices, - torch.zeros(2, dtype=torch.int32, device=device), - extra_cache, - extra_indices, - torch.tensor([0, num_extra], dtype=torch.int32, device=device), - num_splits, - ) - single_empty = _launch_gfx950_partial( - q[1:], - main_cache, - main_indices, - torch.zeros(2, dtype=torch.int32, device=device), - extra_cache, - torch.empty(0, dtype=torch.int32, device=device), - torch.zeros(2, dtype=torch.int32, device=device), - num_splits, - ) - - for batch_part, nonempty_part, empty_part in zip( - batch, single_nonempty, single_empty - ): - assert torch.equal(batch_part[:1], nonempty_part) - assert torch.equal(batch_part[1:], empty_part) - - part_m, part_l, part_acc = batch - expected_m = HEAD_DIM * 0.125**2 * HEAD_DIM**-0.5 - torch.testing.assert_close( - part_m[0], - torch.full_like(part_m[0], expected_m), - rtol=1e-4, - atol=1e-4, - ) - assert torch.equal(part_l[0], torch.full_like(part_l[0], float(entries_per_split))) - torch.testing.assert_close( - part_acc[0], - torch.full_like(part_acc[0], entries_per_split * 0.125), - rtol=0, - atol=0, - ) - assert torch.equal( - part_m[1], torch.full_like(part_m[1], torch.finfo(torch.float32).min) - ) - assert torch.equal(part_l[1], torch.zeros_like(part_l[1])) - assert torch.equal(part_acc[1], torch.zeros_like(part_acc[1])) - - @requires_gfx950 @torch.inference_mode() def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None: device = torch.device("cuda") - torch.manual_seed(19) - block_size = 64 - num_heads = 16 - num_splits = 32 - main_use_fnuz = current_platform.is_fp8_fnuz() - q = torch.randn(3, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) - q *= 0.125 - main_kv = torch.randn(3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 - main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size, main_use_fnuz) - main_rows = [[0], [1], [2]] - main_indices, main_indptr = _ragged_from_rows(main_rows, device) - extra_kv = torch.randn(2112, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 - extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, use_fnuz=False) - extra_rows = [ - list(range(64)), - list(range(64, 576)), - list(range(576, 2112)), - ] - extra_indices, extra_indptr = _ragged_from_rows(extra_rows, device) - - part_m, part_l, part_acc = _launch_gfx950_partial( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - num_splits, - adaptive_splits=True, - poison_scratch=True, - one_wave_splits=16, - ) + part_m = torch.full((1, 8, 1), torch.finfo(torch.float32).min, device=device) + part_l = torch.zeros_like(part_m) + part_acc = torch.full((1, 8, 1, HEAD_DIM), float("nan"), device=device) + part_m[:, :2] = 0 + part_l[:, :2] = 1 + part_acc[:, 0] = 1 + part_acc[:, 1] = 3 - neg_large = torch.finfo(torch.float32).min - assert torch.equal(part_m[0, 4:], torch.full_like(part_m[0, 4:], neg_large)) - assert torch.equal(part_l[0, 4:], torch.zeros_like(part_l[0, 4:])) - assert torch.isnan(part_acc[0, 4::2]).all() - assert torch.isinf(part_acc[0, 5::2]).all() - assert (part_l[1, :16] > 0).all() - assert torch.equal(part_m[1, 16:], torch.full_like(part_m[1, 16:], neg_large)) - assert torch.equal(part_l[1, 16:], torch.zeros_like(part_l[1, 16:])) - assert torch.isnan(part_acc[1, 16::2]).all() - assert torch.isinf(part_acc[1, 17::2]).all() - assert (part_l[2] > 0).all() - assert torch.isfinite(part_acc[2]).all() - - actual = _launch_sparse_decode_reduce( - part_m, - part_l, - part_acc, - adaptive_splits=True, - ) - expected = _ref_sparse_decode_ragged( - q=q, - main_cache=main_cache, - main_rows=main_rows, - scale=HEAD_DIM**-0.5, - attn_sink=None, - block_size=block_size, - extra_cache=extra_cache, - extra_rows=extra_rows, - main_use_fnuz=main_use_fnuz, - ) + actual = _launch_sparse_decode_reduce(part_m, part_l, part_acc, True) assert torch.isfinite(actual).all() - torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + assert torch.equal(actual, torch.full_like(actual, 2)) @requires_gfx950 @@ -1232,61 +938,6 @@ def run_decode() -> torch.Tensor: torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2) -@requires_split_decode_arch -@torch.inference_mode() -def test_sparse_attn_decode_long_context(monkeypatch) -> None: - """Validate the production C128A maximum of 8192 selected KV rows.""" - from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod - - device = torch.device("cuda") - torch.manual_seed(11) - block_size = 64 - num_kv = 8192 - num_heads = 16 - q = torch.randn(1, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 - main_cache = _pack_fp8_ds_mla_cache( - torch.zeros(1, HEAD_DIM, dtype=torch.bfloat16, device=device), - block_size, - use_fnuz=current_platform.is_fp8_fnuz(), - ) - extra_cache = _pack_fp8_ds_mla_cache( - torch.randn(num_kv, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125, - block_size, - use_fnuz=False, - ) - empty_indices = torch.empty(0, dtype=torch.int32, device=device) - empty_indptr = torch.zeros(2, dtype=torch.int32, device=device) - extra_indices = torch.arange(num_kv, dtype=torch.int32, device=device) - extra_indptr = torch.tensor([0, num_kv], dtype=torch.int32, device=device) - scale = HEAD_DIM**-0.5 - - split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits" - num_splits = 32 if _on_gfx950() else 16 - monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: num_splits) - actual = mod._rocm_sparse_attn_decode_ragged_triton( - q=q, - main_cache=main_cache, - main_indices=empty_indices, - main_indptr=empty_indptr, - scale=scale, - attn_sink=None, - nope_head_dim=NOPE_HEAD_DIM, - rope_head_dim=ROPE_HEAD_DIM, - extra_cache=extra_cache, - extra_indices=extra_indices, - extra_indptr=extra_indptr, - ) - - kv = _read_fp8_ds_mla_cache_rows( - extra_cache, extra_indices.to(torch.int64), block_size, use_fnuz=False - ) - scores = torch.matmul(q[0].float(), kv.T) * scale - expected = torch.matmul(torch.softmax(scores, dim=-1), kv) - torch.testing.assert_close( - actual[0], expected.to(torch.bfloat16), atol=2e-2, rtol=2e-2 - ) - - # --------------------------------------------------------------------------- # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) # --------------------------------------------------------------------------- diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 8a5de5d8b4b6..1582071e08b7 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1868,115 +1868,6 @@ def _sparse_attn_decode_partial_kernel( ) -@triton.jit -def _sparse_attn_decode_gfx950_partial_tail_tile( - q_combined, - cache_ptr, - indices_ptr, - index_start, - k_start, - k_hi, - cache_stride0, - num_rows, - scale: tl.constexpr, - head_mask, - m_i, - l_i, - acc_nope_0a, - acc_nope_0b, - acc_nope_1, - acc_tail, - BLOCK_SIZE: tl.constexpr, - NOPE_DIM: tl.constexpr, - BLOCK_K: tl.constexpr, - IS_FNUZ: tl.constexpr, - TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, -): - k_offsets = tl.arange(0, BLOCK_K) - k_pos = k_start + k_offsets - in_range = k_pos < k_hi - slot = tl.load(indices_ptr + index_start + k_pos, mask=in_range, other=-1) - valid = in_range & (slot >= 0) & (slot < num_rows) - safe_slot = tl.where(valid, slot, 0) - - block_idx = safe_slot // BLOCK_SIZE - pos_in_block = safe_slot % BLOCK_SIZE - cache_block_ptr = cache_ptr + block_idx.to(tl.int64) * cache_stride0 - token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + BLOCK_SIZE * 576 + pos_in_block * 8 - k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk( - token_data_ptr, - token_scale_ptr, - valid, - 0, - 128, - BLOCK_K, - IS_FNUZ, - ) - k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk( - token_data_ptr, - token_scale_ptr, - valid, - 128, - 128, - BLOCK_K, - IS_FNUZ, - ) - k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk( - token_data_ptr, - token_scale_ptr, - valid, - 256, - 128, - BLOCK_K, - IS_FNUZ, - ) - k_tail = _load_fp8_ds_mla_gfx950_tail128( - token_data_ptr, - token_scale_ptr, - valid, - NOPE_DIM, - BLOCK_K, - IS_FNUZ, - ) - if not TRUST_EXTRA_CACHE_NAN_FREE: - zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) - k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero) - k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero) - k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero) - k_tail = tl.where(k_tail == k_tail, k_tail, zero) - k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1) - k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1) - k_combined = tl.cat(k_nope_0, k_tail_256, dim=1) - - scores = tl.dot(q_combined, tl.trans(k_combined)) - scores *= scale * 1.4426950408889634 - scores = tl.where( - head_mask[:, None] & valid[None, :], - scores, - -3.4028234663852886e38, - ) - m_block = tl.max(scores, axis=1) - m_new = tl.maximum(m_i, m_block) - alpha = tl.exp2(m_i - m_new) - p = tl.exp2(scores - m_new[:, None]) - p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) - l_new = l_i * alpha + tl.sum(p, axis=1) - p_bf16 = p.to(k_nope_0a.dtype) - acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a) - acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b) - acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1) - acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail) - return ( - m_new, - l_new, - acc_nope_0a, - acc_nope_0b, - acc_nope_1, - acc_tail, - ) - - @triton.jit def _sparse_attn_decode_gfx950_partial_loaded_tile( q_combined, @@ -2202,76 +2093,33 @@ def _sparse_attn_decode_gfx950_partial_kernel( in_range = k_pos < main_hi slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1) valid = in_range & (slot >= 0) & (slot < main_num_rows) - safe_slot = tl.where(valid, slot, 0) - - block_idx = safe_slot // MAIN_BLOCK_SIZE - pos_in_block = safe_slot % MAIN_BLOCK_SIZE - cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 - token_data_ptr = cache_block_ptr + pos_in_block * 576 - token_scale_ptr = cache_block_ptr + MAIN_BLOCK_SIZE * 576 + pos_in_block * 8 - - k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk( - token_data_ptr, - token_scale_ptr, - valid, - 0, - 128, - BLOCK_K, - IS_FNUZ_MAIN, - ) - k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk( - token_data_ptr, - token_scale_ptr, - valid, - 128, - 128, - BLOCK_K, - IS_FNUZ_MAIN, - ) - k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk( - token_data_ptr, - token_scale_ptr, - valid, - 256, - 128, - BLOCK_K, - IS_FNUZ_MAIN, - ) - k_tail = _load_fp8_ds_mla_gfx950_tail128( - token_data_ptr, - token_scale_ptr, + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + main_cache_ptr, + slot, valid, + main_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + MAIN_BLOCK_SIZE, NOPE_DIM, BLOCK_K, IS_FNUZ_MAIN, + False, ) - zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) - k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero) - k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero) - k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero) - k_tail = tl.where(k_tail == k_tail, k_tail, zero) - k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1) - k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1) - k_combined = tl.cat(k_nope_0, k_tail_256, dim=1) - - scores = tl.dot(q_combined, tl.trans(k_combined)) - scores *= scale * 1.4426950408889634 - scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) - - m_block = tl.max(scores, axis=1) - m_new = tl.maximum(m_i, m_block) - alpha = tl.exp2(m_i - m_new) - p = tl.exp2(scores - m_new[:, None]) - p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) - l_new = l_i * alpha + tl.sum(p, axis=1) - - p_bf16 = p.to(k_nope_0a.dtype) - acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a) - acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b) - acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1) - acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail) - m_i = m_new - l_i = l_new if HAS_EXTRA: if not ADAPTIVE_SPLITS: @@ -2353,68 +2201,44 @@ def _sparse_attn_decode_gfx950_partial_kernel( IS_FNUZ_EXTRA, TRUST_EXTRA_CACHE_NAN_FREE, ) - if extra_hi_full < extra_hi: - ( - m_i, - l_i, - acc_nope_0a, - acc_nope_0b, - acc_nope_1, - acc_tail, - ) = _sparse_attn_decode_gfx950_partial_tail_tile( - q_combined, - extra_cache_ptr, - extra_indices_ptr, - extra_start, - extra_hi_full, - extra_hi, - extra_cache_stride0, - extra_num_rows, - scale, - head_mask, - m_i, - l_i, - acc_nope_0a, - acc_nope_0b, - acc_nope_1, - acc_tail, - EXTRA_BLOCK_SIZE, - NOPE_DIM, - BLOCK_K, - IS_FNUZ_EXTRA, - TRUST_EXTRA_CACHE_NAN_FREE, - ) - if extra_hi_full + BLOCK_K < extra_hi: - ( - m_i, - l_i, - acc_nope_0a, - acc_nope_0b, - acc_nope_1, - acc_tail, - ) = _sparse_attn_decode_gfx950_partial_tail_tile( - q_combined, - extra_cache_ptr, - extra_indices_ptr, - extra_start, - extra_hi_full + BLOCK_K, - extra_hi, - extra_cache_stride0, - extra_num_rows, - scale, - head_mask, - m_i, - l_i, - acc_nope_0a, - acc_nope_0b, - acc_nope_1, - acc_tail, - EXTRA_BLOCK_SIZE, - NOPE_DIM, - BLOCK_K, - IS_FNUZ_EXTRA, - TRUST_EXTRA_CACHE_NAN_FREE, - ) + for tail_idx in tl.static_range(2): + tail_start = extra_hi_full + tail_idx * BLOCK_K + if tail_start < extra_hi: + k_pos = tail_start + k_offsets + in_range = k_pos < extra_hi + slot = tl.load( + extra_indices_ptr + extra_start + k_pos, + mask=in_range, + other=-1, + ) + valid = in_range & (slot >= 0) & (slot < extra_num_rows) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot, + valid, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets m_store = tl.where(l_i > 0.0, m_i * 0.6931471805599453, neg_large) From 9c4d637d92d6341c8dd4ad85be4443748ca50995 Mon Sep 17 00:00:00 2001 From: fai Date: Sun, 16 Aug 2026 03:26:43 +0000 Subject: [PATCH 8/8] [ROCm] Keep DeepSeek V4 on MRV1 with the wide eager attention region #51430/#51768 moved DSV4 to MRV2 and a narrow eager region. That is a large decode TPOT regression on ROCm. Default ROCm back to MRV1, wrap the full attention body in the eager break for MRV1 only, and drop the MRV1+PIECEWISE rejection. CUDA keeps MRV2 and the narrow region. Co-authored-by: Nick Hill Co-authored-by: Cursor Grok 4.6 Signed-off-by: fai --- tests/test_config.py | 61 ++++++++++-------------- vllm/config/vllm.py | 32 +++---------- vllm/models/deepseek_v4/attention.py | 69 ++++++++++++++++++++++++++-- 3 files changed, 96 insertions(+), 66 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 93bcdbc258d6..70c8728d25d4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -66,41 +66,18 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): assert envs.VLLM_USE_V2_MODEL_RUNNER is expected -@pytest.mark.parametrize( - "cudagraph_mode", - [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE], -) -def test_deepseek_v4_rejects_mrv1_piecewise_cudagraph(cudagraph_mode): - config = SimpleNamespace( - use_v2_model_runner=False, - model_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]), - compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), - ) - - with pytest.raises(ValueError, match="DeepSeek V4 does not support PIECEWISE"): - VllmConfig._validate_mrv1_piecewise_cudagraph(config) - - -@pytest.mark.parametrize( - ("use_v2_model_runner", "architecture", "cudagraph_mode"), - [ - (True, "DeepseekV4ForCausalLM", CUDAGraphMode.PIECEWISE), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.NONE), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL_DECODE_ONLY), - (False, "LlamaForCausalLM", CUDAGraphMode.PIECEWISE), - ], -) -def test_mrv1_piecewise_cudagraph_allowed( - use_v2_model_runner, architecture, cudagraph_mode -): - config = SimpleNamespace( - use_v2_model_runner=use_v2_model_runner, - model_config=SimpleNamespace(architectures=[architecture]), - compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), - ) - - VllmConfig._validate_mrv1_piecewise_cudagraph(config) +def test_rocm_defaults_deepseek_v4_to_mrv1(monkeypatch): + """ROCm keeps DeepSeek V4 on MRV1, which is still faster there.""" + from vllm.config.vllm import default_v2_model_runner_architectures + from vllm.platforms import current_platform + + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + # The lookup is lru_cached against a fixed platform. + default_v2_model_runner_architectures.cache_clear() + try: + assert "DeepseekV4ForCausalLM" not in default_v2_model_runner_architectures() + finally: + default_v2_model_runner_architectures.cache_clear() @pytest.mark.parametrize( @@ -330,10 +307,20 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( ), ], ) -def test_is_default_v2_model_runner_model(model_config, expected): +def test_is_default_v2_model_runner_model(model_config, expected, monkeypatch): + from vllm.config.vllm import default_v2_model_runner_architectures + from vllm.platforms import current_platform + + # The expectations below are the platform-independent defaults; ROCm's + # DeepSeek V4 carve-out is covered by test_rocm_defaults_deepseek_v4_to_mrv1. + monkeypatch.setattr(current_platform, "is_rocm", lambda: False) + default_v2_model_runner_architectures.cache_clear() config = SimpleNamespace(model_config=model_config) - assert VllmConfig._is_default_v2_model_runner_model(config) is expected + try: + assert VllmConfig._is_default_v2_model_runner_model(config) is expected + finally: + default_v2_model_runner_architectures.cache_clear() @pytest.mark.skip_global_cleanup diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 2b5e512ccacf..0a4c5cad8838 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -66,10 +66,6 @@ logger = init_logger(__name__) -MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES = frozenset( - {"DeepseekV4ForCausalLM"} -) - DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( { "DeepseekV2ForCausalLM", @@ -87,6 +83,13 @@ @lru_cache def default_v2_model_runner_architectures() -> frozenset[str]: """Architectures defaulting to the V2 model runner on this platform.""" + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + # TODO(rocm): DeepSeek V4 is still faster on MRV1 on ROCm. The + # attention layer picks the eager cudagraph region MRV1 needs, so + # this is a perf default only; drop it once MRV2 catches up. + return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {"DeepseekV4ForCausalLM"} return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES @@ -691,25 +694,6 @@ def _is_default_v2_model_runner_model(self) -> bool: return False return is_default_v2_architecture or not model_config.is_moe - def _validate_mrv1_piecewise_cudagraph(self) -> None: - if self.use_v2_model_runner: - return - model_config = self.model_config - if model_config is None: - return - if not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs(): - return - architectures = getattr(model_config, "architectures", []) - if any( - arch in MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES - for arch in architectures - ): - raise ValueError( - "DeepSeek V4 does not support PIECEWISE CUDA graphs with " - "Model Runner V1. Use Model Runner V2 or disable PIECEWISE " - "CUDA graphs." - ) - @property def needs_dp_coordinator(self) -> bool: """ @@ -1644,8 +1628,6 @@ def has_blocked_weights(): "pipeline parallelism", ) - self._validate_mrv1_piecewise_cudagraph() - # final check of cudagraph mode after all possible updates if current_platform.is_cuda_alike(): if ( diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 74a481e4b5f6..e2dc1bb45a7a 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -298,6 +298,13 @@ def __init__( eager_scratch_pool=eager_scratch_pool, ) + self._prepare_and_attn_fn = self._prepare_and_attn + if not vllm_config.use_v2_model_runner: + # MRV1's piecewise capture only tolerates the wide eager region: with + # the narrow one the attention input preparation stays in the captured + # graph and MRV1 produces garbage (#51430). + self._prepare_and_attn_fn = self._prepare_and_attn_eager + # Will be None on ROCm for now. self.aux_stream_list = aux_stream_list # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; @@ -379,6 +386,64 @@ def forward( self.eps, ) + self._prepare_and_attn_fn( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + @eager_break_during_capture + def _prepare_and_attn_eager( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly. + + The nested ``_sparse_indexer_and_attn`` break runs inline, since + ``add_eager`` clears ``_capturing`` before invoking this. + """ + self._prepare_and_attn( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + + def _prepare_and_attn( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Attention input preparation followed by the sparse indexer and MLA. + + Only the latter runs in the eager break. + """ attn_metadata = get_forward_context().attn_metadata indexer = self.indexer compressor = self.compressor @@ -438,10 +503,6 @@ def project_query_and_cache_kv() -> torch.Tensor: positions, o_padded, ) - o = o_padded[:, : self.n_local_heads, :] - - # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). - return self._o_proj(o, positions) def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: # Override point: the ROCm layer preshuffles this weight in place, so