I tested the FP8 batch decode unit tests on MI300a and see the following failures:
test_batch_decode_with_paged_kv_cache[True-kv_dtype0-q_dtype0-True-0.0-NONE-NHD-256-32-4-16-32768-87] __________________________________________________________
batch_size = 87, kv_len = 32768, page_size = 16, num_kv_heads = 4, num_qo_heads = 32, head_dim = 256, kv_layout = 'NHD', pos_encoding_mode = 'NONE', logits_soft_cap = 0.0, return_lse = True, q_dtype = torch.float16
kv_dtype = torch.float8_e4m3fnuz, contiguous_kv = True
@pytest.mark.parametrize("batch_size", [12, 17, 64, 87])
@pytest.mark.parametrize("kv_len", [54, 97, 512, 1024, 2048, 4096, 8192, 16384, 32768])
@pytest.mark.parametrize("page_size", [1, 8, 16])
@pytest.mark.parametrize("num_kv_heads", [4])
@pytest.mark.parametrize("num_qo_heads", [4, 32])
@pytest.mark.parametrize("head_dim", [128, 256])
@pytest.mark.parametrize("kv_layout", ["NHD"])
@pytest.mark.parametrize("pos_encoding_mode", ["NONE"])
@pytest.mark.parametrize("logits_soft_cap", [0.0])
@pytest.mark.parametrize("return_lse", [True])
@pytest.mark.parametrize("q_dtype", [torch.float16])
@pytest.mark.parametrize("kv_dtype", [torch.float8_e4m3fnuz])
@pytest.mark.parametrize("contiguous_kv", [True])
def test_batch_decode_with_paged_kv_cache(
batch_size,
kv_len,
page_size,
num_kv_heads,
num_qo_heads,
head_dim,
kv_layout,
pos_encoding_mode,
logits_soft_cap,
return_lse,
q_dtype,
kv_dtype,
contiguous_kv,
):
q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=q_dtype)
num_pages_per_seq = (kv_len + page_size - 1) // page_size
total_num_pages = num_pages_per_seq * batch_size
if kv_layout == "HND":
kv_shape = [total_num_pages, 2, num_kv_heads, page_size, head_dim]
else:
kv_shape = [total_num_pages, 2, page_size, num_kv_heads, head_dim]
if not contiguous_kv:
tmp = [kv_shape[0]]
for v in kv_shape[1:]:
tmp.append(2)
tmp.append(v)
kv_shape = tmp
kv_data_fp32 = torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
kv_data = kv_data_fp32.to(kv_dtype)
kv_data = kv_data[:, 1, :, 1, :, 1, :, 1, :]
kv_data_fp32 = kv_data_fp32[:, 1, :, 1, :, 1, :, 1, :]
# actual data is stored in non-contiguous memory
assert (
kv_data.stride(-4)
!= kv_data.shape[-3] * kv_data.shape[-2] * kv_data.shape[-1]
)
else:
kv_data_fp32 = torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
kv_data = kv_data_fp32.to(kv_dtype)
kv_indptr = (
torch.arange(0, batch_size + 1, device="cuda:0", dtype=torch.int32)
* num_pages_per_seq
)
kv_indices = torch.arange(0, total_num_pages, device="cuda:0", dtype=torch.int32)
kv_last_page_len = torch.full(
(batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda:0"
)
workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.int8, device="cuda:0")
wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
workspace_buffer, kv_layout
)
wrapper.plan(
kv_indptr,
kv_indices,
kv_last_page_len,
num_qo_heads,
num_kv_heads,
head_dim,
page_size,
logits_soft_cap=logits_soft_cap,
pos_encoding_mode=pos_encoding_mode,
data_type=kv_dtype,
q_data_type=q_dtype,
)
if return_lse:
o, _ = wrapper.run(q, kv_data, return_lse=True)
else:
o = wrapper.run(q, kv_data)
for i in range(batch_size):
perm_dims = [0, 2, 1, 3] if kv_layout == "HND" else [0, 1, 2, 3]
perm_dims_last = [1, 0, 2] if kv_layout == "HND" else [0, 1, 2]
qi = q[i]
ki = torch.cat(
[
kv_data_fp32[kv_indptr[i] : kv_indptr[i + 1] - 1, 0]
.permute(*perm_dims)
.reshape(-1, num_kv_heads, head_dim),
(
kv_data_fp32[kv_indptr[i + 1] - 1, 0, :, : kv_last_page_len[i]]
if kv_layout == "HND"
else kv_data_fp32[kv_indptr[i + 1] - 1, 0, : kv_last_page_len[i], :]
)
.permute(*perm_dims_last)
.reshape(-1, num_kv_heads, head_dim),
],
dim=0,
).to(kv_dtype)
vi = torch.cat(
[
kv_data_fp32[kv_indptr[i] : kv_indptr[i + 1] - 1, 1]
.permute(*perm_dims)
.reshape(-1, num_kv_heads, head_dim),
(
kv_data_fp32[kv_indptr[i + 1] - 1, 1, :, : kv_last_page_len[i]]
if kv_layout == "HND"
else kv_data_fp32[kv_indptr[i + 1] - 1, 1, : kv_last_page_len[i], :]
)
.permute(*perm_dims_last)
.reshape(-1, num_kv_heads, head_dim),
],
dim=0,
).to(kv_dtype)
o_ref_i = flashinfer.decode.single_decode_with_kv_cache(
qi,
ki,
vi,
pos_encoding_mode=pos_encoding_mode,
logits_soft_cap=logits_soft_cap,
)
> torch.testing.assert_close(o[i], o_ref_i, rtol=1e-3, atol=1e-3)
E AssertionError: Tensor-likes are not close!
E
E Mismatched elements: 229 / 8192 (2.8%)
E Greatest absolute difference: 0.029266357421875 at index (26, 220) (up to 0.001 allowed)
E Greatest relative difference: 109.25 at index (26, 201) (up to 0.001 allowed)
test_batch_decode_kernels_hip_fp8.py:163: AssertionError
================================================================================================= short test summary info =================================================================================================
FAILED test_batch_decode_kernels_hip_fp8.py::test_batch_decode_with_paged_kv_cache[True-kv_dtype0-q_dtype0-True-0.0-NONE-NHD-256-32-4-1-8192-87] - AssertionError: Tensor-likes are not close!
FAILED test_batch_decode_kernels_hip_fp8.py::test_batch_decode_with_paged_kv_cache[True-kv_dtype0-q_dtype0-True-0.0-NONE-NHD-256-32-4-8-2048-87] - AssertionError: Tensor-likes are not close!
FAILED test_batch_decode_kernels_hip_fp8.py::test_batch_decode_with_paged_kv_cache[True-kv_dtype0-q_dtype0-True-0.0-NONE-NHD-256-32-4-8-16384-87] - AssertionError: Tensor-likes are not close!
FAILED test_batch_decode_kernels_hip_fp8.py::test_batch_decode_with_paged_kv_cache[True-kv_dtype0-q_dtype0-True-0.0-NONE-NHD-256-32-4-16-32768-87] - AssertionError: Tensor-likes are not close!
======================================================================================== 4 failed, 860 passed in 87.41s (0:01:27) =================================================
I tested the FP8 batch decode unit tests on MI300a and see the following failures: