[ROCm][DSV4][Perf] Optimize Triton sparse-MLA decode on gfx950 (#52212)
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com> Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Cursor Grok 4.6 <cursoragent@cursor.com>
This commit is contained in:
Fangzhou Ai
co-authored by
Nick Hill
Cursor Grok 4.6
Nick Hill
Cursor Grok 4.6
parent
1f0e0bf612
commit
ef43e3101b
@@ -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)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user