iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[ROCm][DSv4] Fix sparse-indexer logits collapse on gfx950/gfx942 (#50455)

Browse Source 
Signed-off-by: Frida Andersson <fanderss@amd.com>
Signed-off-by: frida-andersson <fanderss@amd.com>
This commit is contained in:
frida-andersson
2026-09-18 16:57:44 +00:00
committed by GitHub
parent 8a5cf54387
commit 71fc70d3ae
2 changed files with 373 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+186
View File
@@ -1430,3 +1430,189 @@ def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:
# Second call returns the same cached object.
assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out
@requires_split_decode_arch
@torch.inference_mode()
def test_fp8_paged_mqa_logits_triton_matches_torch_ref() -> None:
"""Block-flat Triton decode logits vs the eager torch reference.
Compare only valid key positions: the kernel skips the -inf pad that
the torch ref writes past seq_len.
"""
from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
from vllm.v1.worker.workspace import init_workspace_manager
device = torch.device("cuda")
init_workspace_manager(device)
torch.manual_seed(0)
fp8_dtype = current_platform.fp8_dtype()
batch_size = 2
next_n = 1
num_heads = 4
head_size = 128
block_size = 64
max_model_len = 256
seq_lens = [180, 64]
num_pages = [(seq_len + block_size - 1) // block_size for seq_len in seq_lens]
num_blocks = sum(num_pages)
q = torch.randn(
batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16
).to(fp8_dtype)
weights = torch.rand(
batch_size * next_n, num_heads, device=device, dtype=torch.float32
)
context_lens = torch.tensor(seq_lens, device=device, dtype=torch.int32)
block_tables = torch.full(
(batch_size, max(num_pages)), 0, device=device, dtype=torch.int32
)
page = 0
for i, n_pages in enumerate(num_pages):
block_tables[i, :n_pages] = torch.arange(
page, page + n_pages, device=device, dtype=torch.int32
)
page += n_pages
kv_cache = torch.empty(
num_blocks, block_size, 1, head_size + 4, device=device, dtype=torch.uint8
)
values = torch.randn(
num_blocks, block_size, head_size, device=device, dtype=torch.bfloat16
).to(fp8_dtype)
scales = torch.rand(
num_blocks, block_size, device=device, dtype=torch.float32
).clamp(min=1e-3)
kv_flat = kv_cache.view(num_blocks, -1)
scale_off = block_size * head_size
kv_flat[:, :scale_off] = values.reshape(num_blocks, -1).view(torch.uint8)
kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_blocks, -1)
got = mod.rocm_fp8_paged_mqa_logits_triton(
q, kv_cache, weights, context_lens, block_tables, max_model_len
)
ref = mod.fp8_paged_mqa_logits_torch(
q, kv_cache, weights, context_lens, block_tables, max_model_len
)
for i, seq_len in enumerate(seq_lens):
torch.testing.assert_close(
got[i, :seq_len], ref[i, :seq_len], atol=2e-3, rtol=2e-3
)
def test_indexer_k_is_c4a_block_flat_gate() -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_indexer_k_is_c4a_block_flat,
)
assert not _indexer_k_is_c4a_block_flat(1)
assert not _indexer_k_is_c4a_block_flat(2)
assert _indexer_k_is_c4a_block_flat(4)
assert not _indexer_k_is_c4a_block_flat(128)
@requires_split_decode_arch
@torch.inference_mode()
def test_paged_mqa_logits_gate_v41_shuffle_vs_c4a_flat() -> None:
"""Ratio 1/2 stay on AITER; ratio 4 takes Triton."""
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
fp8_paged_mqa_logits_torch,
indexer_k_quant_and_cache_triton,
rocm_fp8_paged_mqa_logits,
rocm_fp8_paged_mqa_logits_triton,
)
from vllm.v1.worker.workspace import init_workspace_manager
device = torch.device("cuda")
init_workspace_manager(device)
torch.manual_seed(1)
fp8_dtype = current_platform.fp8_dtype()
batch_size = 1
next_n = 1
num_heads = 32
head_size = 128
block_size = 64
seq_len = 180
max_model_len = 256
num_pages = (seq_len + block_size - 1) // block_size
num_tokens = num_pages * block_size
q = torch.randn(
batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16
).to(fp8_dtype)
weights = torch.rand(
batch_size * next_n, num_heads, device=device, dtype=torch.float32
)
context_lens = torch.tensor([seq_len], device=device, dtype=torch.int32)
block_tables = torch.arange(num_pages, device=device, dtype=torch.int32).view(1, -1)
dummy_sched = torch.zeros(8, 2, device=device, dtype=torch.int32)
k_bf16 = torch.randn(num_tokens, head_size, device=device, dtype=torch.bfloat16)
slot = torch.arange(num_tokens, device=device, dtype=torch.int32)
cache_shuffle = torch.zeros(
num_pages, block_size, head_size + 4, device=device, dtype=torch.uint8
)
indexer_k_quant_and_cache_triton(k_bf16, cache_shuffle, slot, 128, "ue8m0")
cache_shuffle = cache_shuffle.unsqueeze(2)
aiter = rocm_fp8_paged_mqa_logits(
q,
cache_shuffle,
weights,
context_lens,
block_tables,
dummy_sched,
max_model_len,
compress_ratio=1,
)
for ratio in (1, 2):
got = rocm_fp8_paged_mqa_logits(
q,
cache_shuffle,
weights,
context_lens,
block_tables,
dummy_sched,
max_model_len,
compress_ratio=ratio,
)
torch.testing.assert_close(got[:, :seq_len], aiter[:, :seq_len], atol=0, rtol=0)
cache_flat = torch.empty(
num_pages, block_size, 1, head_size + 4, device=device, dtype=torch.uint8
)
values = torch.randn(
num_pages, block_size, head_size, device=device, dtype=torch.bfloat16
).to(fp8_dtype)
scales = torch.rand(
num_pages, block_size, device=device, dtype=torch.float32
).clamp(min=1e-3)
kv_flat = cache_flat.view(num_pages, -1)
scale_off = block_size * head_size
kv_flat[:, :scale_off] = values.reshape(num_pages, -1).view(torch.uint8)
kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_pages, -1)
gated_c4a = rocm_fp8_paged_mqa_logits(
q,
cache_flat,
weights,
context_lens,
block_tables,
dummy_sched,
max_model_len,
compress_ratio=4,
)
triton = rocm_fp8_paged_mqa_logits_triton(
q, cache_flat, weights, context_lens, block_tables, max_model_len
)
ref = fp8_paged_mqa_logits_torch(
q, cache_flat, weights, context_lens, block_tables, max_model_len
)
torch.testing.assert_close(
gated_c4a[:, :seq_len], triton[:, :seq_len], atol=0, rtol=0
)
torch.testing.assert_close(
gated_c4a[:, :seq_len], ref[:, :seq_len], atol=2e-3, rtol=2e-3
)
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+187 -5
View File
@@ -66,6 +66,11 @@ _GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024
_GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024
def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool:
"""V4.0 C4A is block-flat (NORMAL). Ratio 1 and 2 are 16×16 SHUFFLE."""
return compress_ratio == 4
def _get_aiter_top_k_kernel(
*,
is_prefill: bool,
@@ -458,6 +463,7 @@ def cp_gather_indexer_k_quant_cache_triton(
token_to_seq: torch.Tensor,
block_tile_size: int = 16,
head_tile_size: int = 16,
cache_layout: str | None = None,
):
num_tokens = k_fp8.size(0)
block_size = k_cache.size(1)
@@ -470,7 +476,10 @@ def cp_gather_indexer_k_quant_cache_triton(
k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32)
grid = (num_tokens,)
k_fp8_scale = k_fp8_scale.view(torch.float32)
layout = "NORMAL" if block_size == 1 else "SHUFFLE"
if cache_layout is None:
layout = "NORMAL" if block_size == 1 else "SHUFFLE"
else:
layout = cache_layout
kernel_args = (
k_cache_value,
k_cache_scale,
@@ -607,6 +616,158 @@ def fp8_paged_mqa_logits_torch(
return logits
@triton.jit
def _fp8_paged_mqa_logits_decode_kernel(
q_ptr, # fp8 [B, NEXT_N, H, D]
kv_val_ptr, # fp8, block-flat: [num_blocks, block_size*(D+4)]
kv_scale_ptr, # fp32, block-flat: [num_blocks, block_size*(D+4)//4]
weights_ptr, # fp32 [B*NEXT_N, H]
ctx_lens_ptr, # int32 [B*NEXT_N] if CTX_PER_ROW else [B]
block_tables_ptr, # int32 [B, max_blocks]
logits_ptr, # fp32 [B*NEXT_N, max_model_len]
stride_q_b,
stride_q_n,
stride_q_h,
stride_w_row,
stride_kvblk_fp8,
stride_kvblk_f32,
scale_region_off, # block_size*D // 4 (fp32 offset of scale region within a block)
stride_bt_b,
stride_logits_row,
max_blocks, # block_tables width; guards the block-table gather
max_model_len,
NUM_HEADS: tl.constexpr,
HEAD_SIZE: tl.constexpr,
BLOCK_SIZE: tl.constexpr, # paged-cache page size
BLOCK_KV: tl.constexpr, # positions per tile; must divide BLOCK_SIZE
N_SPLITS: tl.constexpr, # KV-tile parallelism factor (grid dim 1)
NEXT_N: tl.constexpr, # query positions per batch (1 = decode; >1 = MTP verify)
CTX_PER_ROW: tl.constexpr, # ctx_lens is already per-(b, n)
):
# scale_region_off = block_size*D//4 needs D % 4 == 0 (D=128).
tl.static_assert(HEAD_SIZE % 4 == 0)
# Grid (B*NEXT_N, N_SPLITS): disjoint KV tiles per program, no host sync.
row = tl.program_id(0)
split = tl.program_id(1)
b = row // NEXT_N
n = row % NEXT_N
if CTX_PER_ROW:
seq_len = tl.load(ctx_lens_ptr + row)
else:
seq_len = tl.load(ctx_lens_ptr + b) - NEXT_N + n + 1
seq_len = tl.minimum(tl.maximum(seq_len, 0), max_model_len)
h = tl.arange(0, NUM_HEADS)
d = tl.arange(0, HEAD_SIZE)
# keep q/kv fp8 -> fp8 MFMA (f32 accum), not the slow f32 path
q = tl.load(
q_ptr + b * stride_q_b + n * stride_q_n + h[:, None] * stride_q_h + d[None, :]
)
w = tl.load(weights_ptr + row * stride_w_row + h).to(tl.float32) # [H]
kv_col = tl.arange(0, BLOCK_KV)
for kv_start in tl.range(split * BLOCK_KV, seq_len, N_SPLITS * BLOCK_KV):
pos = kv_start + kv_col
# Tiles are page-aligned (BLOCK_KV | BLOCK_SIZE), so a tile is one page.
logical_blk = kv_start // BLOCK_SIZE
blk_ok = logical_blk < max_blocks # guard against a wild page -> fault
mask_pos = (pos < seq_len) & blk_ok
page = tl.load(
block_tables_ptr + b * stride_bt_b + logical_blk, mask=blk_ok, other=0
).to(tl.int64)
pos_in_blk = pos - logical_blk * BLOCK_SIZE # [BLOCK_KV] == pos % BLOCK_SIZE
# block-flat layout: values region, then scales region.
val_off = page * stride_kvblk_fp8 + pos_in_blk[:, None] * HEAD_SIZE + d[None, :]
kv = tl.load(
kv_val_ptr + val_off, mask=mask_pos[:, None], other=0.0
) # [BLOCK_KV, D] fp8
sc = tl.load(
kv_scale_ptr + page * stride_kvblk_f32 + scale_region_off + pos_in_blk,
mask=mask_pos,
other=0.0,
) # [BLOCK_KV]
# fp8 upcasts exactly; bit-identical on gfx942, ~1e-5 rel on gfx950
# (accum order), top-k unchanged.
s = tl.dot(kv, tl.trans(q)) # [BLOCK_KV, H]
s = tl.maximum(s, 0.0)
s = s * w[None, :]
s = tl.sum(s, axis=1) # [BLOCK_KV]
s = s * sc
tl.store(logits_ptr + row * stride_logits_row + pos, s, mask=mask_pos)
def rocm_fp8_paged_mqa_logits_triton(
q_fp8: torch.Tensor,
kv_cache_fp8: torch.Tensor,
weights: torch.Tensor,
context_lens: torch.Tensor,
block_tables: torch.Tensor,
max_model_len: int,
) -> torch.Tensor:
"""Triton paged MQA-logits for decode and MTP; matches the torch ref but
has no host sync, so it is safe to capture under a full CUDA graph."""
batch_size, next_n, num_heads, head_size = q_fp8.shape
block_size = kv_cache_fp8.shape[1]
BLOCK_KV = 64
assert block_size % BLOCK_KV == 0
fp8_dtype = current_platform.fp8_dtype()
num_blocks = kv_cache_fp8.shape[0]
kv_flat = kv_cache_fp8.reshape(
num_blocks, -1
) # uint8 [num_blocks, block_size*(D+4)]
kv_val = kv_flat.view(fp8_dtype) # [num_blocks, block_size*(D+4)] fp8
kv_scale = kv_flat.view(torch.float32) # [num_blocks, block_size*(D+4)//4] fp32
cl = context_lens.reshape(-1)
ctx_per_row = not (next_n > 1 and cl.numel() == batch_size)
max_blocks = block_tables.shape[1]
(out_logits,) = current_workspace_manager().get_simultaneous(
((batch_size * next_n, max_model_len), torch.float32),
)
# Memory-bound over the KV range: split each row's keys across programs so
# few-row / long-context launches still fill the GPU. Cap splits at the
# device CU count (304 on gfx942, 256 on gfx950) rather than a gfx950-sized
# constant. All terms are static at launch, so the grid stays CUDA-graph-safe.
rows = batch_size * next_n
tiles_cap = (max_model_len + BLOCK_KV - 1) // BLOCK_KV
N_SPLITS = max(1, min(max(1, _decode_cu_count()), tiles_cap, 1024 // rows))
_fp8_paged_mqa_logits_decode_kernel[(rows, N_SPLITS)](
q_fp8,
kv_val,
kv_scale,
weights,
cl,
block_tables,
out_logits,
q_fp8.stride(0),
q_fp8.stride(1),
q_fp8.stride(2),
weights.stride(0),
kv_val.stride(0),
kv_scale.stride(0),
(block_size * head_size) // 4,
block_tables.stride(0),
out_logits.stride(0),
max_blocks,
max_model_len,
NUM_HEADS=num_heads,
HEAD_SIZE=head_size,
BLOCK_SIZE=block_size,
BLOCK_KV=BLOCK_KV,
N_SPLITS=N_SPLITS,
NEXT_N=next_n,
CTX_PER_ROW=ctx_per_row,
num_warps=4,
num_stages=2,
)
return out_logits
@functools.lru_cache
def paged_mqa_logits_module():
paged_mqa_logits_module_path = None
@@ -632,6 +793,8 @@ def rocm_fp8_paged_mqa_logits(
block_tables: torch.Tensor,
schedule_metadata: torch.Tensor,
max_model_len: int,
*,
compress_ratio: int = 1,
) -> torch.Tensor:
"""Compute FP8 MQA logits using paged KV-cache.
@@ -639,8 +802,7 @@ def rocm_fp8_paged_mqa_logits(
q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to
`torch.float8_e4m3fn` by caller.
kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape
[num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last
4 bytes per (block,pos) store the `float` dequant scale.
[num_blocks, block_size, 1, D+4], dtype `torch.uint8`.
weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
context_lens: Tensor of shape [B], dtype int32; effective context length
for each batch element.
@@ -649,6 +811,7 @@ def rocm_fp8_paged_mqa_logits(
schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
used to distribute work across SMs.
max_model_len: Maximum sequence length used to size the logits output.
compress_ratio: C4A (4) takes block-flat Triton; 1 and 2 stay on AITER.
Returns:
Logits tensor of shape [B * next_n, max_model_len], dtype
@@ -657,11 +820,26 @@ def rocm_fp8_paged_mqa_logits(
"""
from vllm._aiter_ops import rocm_aiter_ops
aiter_paged_mqa_logits_module = None
# if rocm_aiter_ops.is_enabled():
batch_size, next_n = q_fp8.shape[:2]
block_size = kv_cache_fp8.shape[1]
# C4A only: Flash/DSv3.2 also skip insert but still write SHUFFLE.
if (
(_ON_GFX950 or _ON_GFX942)
and _indexer_k_is_c4a_block_flat(compress_ratio)
and block_size > 1
):
if block_size % 64 == 0:
return rocm_fp8_paged_mqa_logits_triton(
q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
)
# Non 64-aligned page size (not used in prod): eager torch ref.
return fp8_paged_mqa_logits_torch(
q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
)
aiter_paged_mqa_logits_module = None
if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():
aiter_paged_mqa_logits_module = paged_mqa_logits_module()
@@ -1128,6 +1306,9 @@ def rocm_aiter_sparse_attn_indexer(
chunk.block_table,
chunk.cu_seq_lens,
token_to_seq=chunk.token_to_seq,
cache_layout=(
"NORMAL" if _indexer_k_is_c4a_block_flat(compress_ratio) else None
),
)
logits = rocm_fp8_mqa_logits(
q_fp8[chunk.token_start : chunk.token_end],
@@ -1225,6 +1406,7 @@ def rocm_aiter_sparse_attn_indexer(
decode_metadata.block_table,
decode_metadata.schedule_metadata,
max_model_len=max_model_len,
compress_ratio=compress_ratio,
)
if candidate_blocks is not None: