diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index f0ac4d8385..256b50d49a 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -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 + ) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 3d891c8ce0..6b1530d74e 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -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: