From ef43e3101b8fda8cd2b52de150c76b4fc177fad2 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sun, 16 Aug 2026 15:16:46 -0400 Subject: [PATCH] [ROCm][DSV4][Perf] Optimize Triton sparse-MLA decode on gfx950 (#52212) Signed-off-by: Fangzhou Ai Signed-off-by: fai Co-authored-by: Nick Hill Co-authored-by: Cursor Grok 4.6 --- .../attention/test_rocm_triton_attn_dsv4.py | 481 +++++++++- tests/kernels/test_compressor_kv_cache.py | 243 ++++- vllm/models/deepseek_v4/amd/rocm.py | 51 +- .../common/ops/fused_compress_quant_cache.py | 23 +- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 846 ++++++++++++++++-- 5 files changed, 1550 insertions(+), 94 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 6fe2a3e775..7a98f3ec70 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,48 @@ 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 + 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" + ) + + +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_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM] + 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 : scale_base + 7].to(torch.float32) - 127.0 + cache_flat[scale_base[:, None] + scale_offsets].to(torch.float32) - 127.0 ) - 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]) + 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( @@ -158,19 +193,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: @@ -198,6 +244,46 @@ def _ragged_from_rows( ) +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 @@ -312,6 +398,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 ( @@ -366,6 +465,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 +478,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,9 +492,111 @@ 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) +@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)) + + +@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: @@ -408,6 +611,9 @@ 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 shared gfx942 selector retains its original 16-split ceiling. + assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16 + # 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): @@ -418,6 +624,16 @@ 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) + 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(512, 1, 128, 7812) == 1 + + @requires_split_decode_arch @pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8]) @pytest.mark.parametrize("with_extra", [True, False]) @@ -473,7 +689,14 @@ 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" + 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, @@ -503,6 +726,218 @@ def test_sparse_attn_decode_split_k_kernel( torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None: + device = torch.device("cuda") + 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 + + actual = _launch_sparse_decode_reduce(part_m, part_l, part_acc, True) + + assert torch.isfinite(actual).all() + assert torch.equal(actual, torch.full_like(actual, 2)) + + +@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 = 16 + num_heads = 16 + num_splits = 8 + extra_per_query = 65 * num_splits + max_extra_per_query = 8192 + 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) + 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) + + 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, + adaptive_splits=True, + ) + + run_decode() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_out = run_decode() + 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) + + extra_indices[: short_indices.numel()].copy_(short_indices) + extra_indptr.copy_(short_indptr) + graph.replay() + torch.accelerator.synchronize() + short_out = out.clone() + expected_short = _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=short_extra_rows, + ) + + 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) + + # --------------------------------------------------------------------------- # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) # --------------------------------------------------------------------------- diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index c1d9a2106b..cfdec96e61 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -28,16 +28,81 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( _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 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 +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 + + +@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)], @@ -91,6 +156,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"), [ @@ -402,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" @@ -410,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 ) @@ -431,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 6a457298dc..5e7609bd11 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -19,6 +19,7 @@ from vllm.models.deepseek_v4.sparse_mla import ( 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.attention.ops.rocm_aiter_mla_sparse import ( 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) @@ -294,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 @@ -306,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 @@ -370,6 +389,16 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBui 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() @@ -451,7 +480,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 @@ -605,6 +636,13 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): 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( @@ -615,6 +653,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): 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 @@ -666,6 +705,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): 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, + 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 a2085cd220..4c9f464ef0 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 @@ from typing import Any 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 94eeaba6af..5c9ab386d4 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 @@ -1211,6 +1303,90 @@ 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 _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, @@ -1685,6 +1861,407 @@ def _sparse_attn_decode_partial_kernel( ) +@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, + ADAPTIVE_SPLITS: tl.constexpr, + ONE_WAVE_SPLITS: 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 + 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) + 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) + + 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) + + 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) + + 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) + ( + 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, + ) + + if HAS_EXTRA: + 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) + + 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, + ) + 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) + 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, @@ -1701,6 +2278,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, @@ -1769,17 +2347,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) @@ -1986,6 +2574,49 @@ 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 >= 16 + and num_splits > 4 + and _decode_partial_iters(avg_main_len, avg_extra_len, 4, block_k) <= 3 + ): + return 4 + 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 + ) + 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, @@ -1998,6 +2629,9 @@ 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, + 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, ( @@ -2038,6 +2672,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 @@ -2059,7 +2697,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 @@ -2103,14 +2750,35 @@ 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 + ) + + 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( @@ -2123,48 +2791,88 @@ 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, + ADAPTIVE_SPLITS=adaptive_splits, + ONE_WAVE_SPLITS=one_wave_splits, + 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, @@ -2181,6 +2889,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, @@ -2206,6 +2915,9 @@ 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, + 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( @@ -2241,6 +2953,9 @@ def _rocm_sparse_attn_decode_triton( extra_cache=extra_cache, extra_indices=extra_ragged_indices, extra_indptr=extra_ragged_indptr, + out=out, + extra_cache_nan_free=extra_cache_nan_free, + adaptive_splits=adaptive_splits, ) @@ -2312,6 +3027,8 @@ def rocm_sparse_attn_decode( nope_head_dim: int, 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, " @@ -2341,6 +3058,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 attn_out = _rocm_sparse_attn_decode_triton( q=q, main_cache=swa_k_cache, @@ -2357,5 +3075,9 @@ 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, + extra_cache_nan_free=extra_cache_nan_free, + adaptive_splits=adaptive_splits, ) - output.copy_(attn_out.to(output.dtype)) + if direct_out is None: + output.copy_(attn_out.to(output.dtype))