[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)
* Merge with private repo * Update README * Update README * Update README * Add PyTorch requirements * Fix sync scopes for MQA logits (#256) * Update README
This commit is contained in:
Chenggang Zhao
+225
-148
@@ -10,9 +10,9 @@ from deep_gemm.testing import (
|
||||
ignore_env, get_arch_major,
|
||||
test_filter
|
||||
)
|
||||
from deep_gemm.utils import ceil_div, per_custom_dims_cast_to_fp8
|
||||
from deep_gemm.utils import ceil_div, per_custom_dims_cast_to_fp8, per_token_cast_to_fp4, cast_back_from_fp4
|
||||
|
||||
from generators import generate_normal, get_ue8m0_usage, get_kernel_types, MajorTypeAB
|
||||
from generators import get_arch_major, generate_normal, get_ue8m0_usage, get_kernel_types, reset_seed, MajorTypeAB
|
||||
|
||||
|
||||
def apply_skip_head_mid(d: torch.Tensor, head_splits: Tuple[int, int, int]):
|
||||
@@ -53,40 +53,14 @@ def test_gemm_skip_head_mid() -> None:
|
||||
assert diff < 0.001, f'{m=}, {n=}, {k=}, {kernel_opt}, {diff:.5f}'
|
||||
|
||||
t = bench_kineto(lambda: deep_gemm.fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, disable_ue8m0_cast=disable_ue8m0_cast),
|
||||
'fp8_gemm', suppress_kineto_output=True)
|
||||
'gemm_', suppress_kineto_output=True)
|
||||
print(f' > Perf (m={m:5}, n={n:5}, k={k:5}, {kernel_opt}): '
|
||||
f'{t * 1e6:4.0f} us | '
|
||||
f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
|
||||
f'{(count_bytes(a, b, d)) / 1e9 / t:4.0f} GB/s')
|
||||
f'{t * 1e6:4.0f} us | '
|
||||
f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
|
||||
f'{(count_bytes(a, b, d)) / 1e9 / t:4.0f} GB/s')
|
||||
print()
|
||||
|
||||
|
||||
def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor:
|
||||
num_blocks, block_size, num_heads, head_dim = x.shape
|
||||
assert num_heads == 1
|
||||
x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4)
|
||||
sf = x_amax / 448.0
|
||||
x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn)
|
||||
x_fp8 = torch.empty((num_blocks, block_size * (head_dim + 4)), device=x.device, dtype=torch.uint8)
|
||||
x_fp8[ :, : block_size * head_dim] = x_scaled.view(num_blocks, block_size * head_dim).view(dtype=torch.uint8)
|
||||
x_fp8[ :, block_size * head_dim :] = sf.view(num_blocks, block_size).view(dtype=torch.uint8)
|
||||
return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4)
|
||||
|
||||
|
||||
def generate_cp_test_data(seq_len, seq_len_kv):
|
||||
assert seq_len_kv % seq_len == 0 and seq_len % 2 == 0
|
||||
chunk_size = seq_len // 2
|
||||
cp_size = seq_len_kv // seq_len
|
||||
# Select an arbitrary CP rank
|
||||
cp_id = cp_size // 3
|
||||
ks = torch.zeros(seq_len, dtype=torch.int, device='cuda')
|
||||
ke = torch.zeros(seq_len, dtype=torch.int, device='cuda')
|
||||
for i in range(chunk_size):
|
||||
ke[i] = cp_id * chunk_size + i
|
||||
ke[i + chunk_size] = (cp_size * 2 - 1 - cp_id) * chunk_size + i
|
||||
return ks, ke
|
||||
|
||||
|
||||
def ref_fp8_mqa_logits(q: torch.Tensor, kv: torch.Tensor, weights: torch.Tensor,
|
||||
cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, cost_only: bool = False):
|
||||
seq_len_kv = kv.shape[0]
|
||||
@@ -113,92 +87,137 @@ def ref_fp8_mqa_logits(q: torch.Tensor, kv: torch.Tensor, weights: torch.Tensor,
|
||||
return logits, cost
|
||||
|
||||
|
||||
@ignore_env('DG_JIT_PTXAS_CHECK', lambda: get_arch_major() == 10)
|
||||
def test_mqa_logits():
|
||||
|
||||
# Helper functions
|
||||
def generate_ks_ke_tests(seq_len: int, seq_len_kv: int, disable_cp: bool):
|
||||
if disable_cp:
|
||||
ks = torch.zeros(seq_len, dtype=torch.int, device='cuda')
|
||||
ke = torch.arange(seq_len, dtype=torch.int, device='cuda') + (seq_len_kv - seq_len)
|
||||
return ks, ke
|
||||
assert seq_len_kv % seq_len == 0 and seq_len % 2 == 0
|
||||
chunk_size = seq_len // 2
|
||||
cp_size = seq_len_kv // seq_len
|
||||
# Select an arbitrary CP rank
|
||||
cp_id = cp_size // 3
|
||||
ks = torch.zeros(seq_len, dtype=torch.int, device='cuda')
|
||||
ke = torch.zeros(seq_len, dtype=torch.int, device='cuda')
|
||||
for i in range(chunk_size):
|
||||
ke[i] = cp_id * chunk_size + i
|
||||
ke[i + chunk_size] = (cp_size * 2 - 1 - cp_id) * chunk_size + i
|
||||
return ks, ke
|
||||
|
||||
def enumerate_mqa_logits():
|
||||
for is_fp4 in ((True, False) if get_arch_major() == 10 else (False, )):
|
||||
for logits_dtype in (torch.float, torch.bfloat16):
|
||||
for compressed_logits, clean_logits in [(False, True), (True, False)]:
|
||||
for seq_len in (2048, 4096):
|
||||
for seq_len_kv in (4096, 8192):
|
||||
for num_heads, head_dim in [(64, 128)]:
|
||||
for disable_cp in (False, True):
|
||||
yield is_fp4, logits_dtype, compressed_logits, clean_logits, seq_len, seq_len_kv, num_heads, head_dim, disable_cp
|
||||
|
||||
print('Testing FP8 MQA Logits:')
|
||||
num_heads, head_dim = 64, 128
|
||||
for seq_len in (2048, 4096):
|
||||
for compressed_logits in (False, True):
|
||||
for seq_len_kv in (4096, 8192):
|
||||
for disable_cp in (False, True):
|
||||
q = torch.randn(seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)
|
||||
kv = torch.randn(seq_len_kv, head_dim, device='cuda', dtype=torch.bfloat16)
|
||||
weights = torch.randn(seq_len, num_heads, device='cuda', dtype=torch.float32)
|
||||
for is_fp4, logits_dtype, compressed_logits, clean_logits, seq_len, seq_len_kv, num_heads, head_dim, disable_cp in enumerate_mqa_logits():
|
||||
# Generate random inputs
|
||||
q = torch.randn(seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)
|
||||
kv = torch.randn(seq_len_kv, head_dim, device='cuda', dtype=torch.bfloat16)
|
||||
weights = torch.randn(seq_len, num_heads, device='cuda', dtype=torch.float32)
|
||||
ks, ke = generate_ks_ke_tests(seq_len, seq_len_kv, disable_cp)
|
||||
|
||||
if disable_cp:
|
||||
ks = torch.zeros(seq_len, dtype=torch.int, device='cuda')
|
||||
ke = torch.arange(seq_len, dtype=torch.int, device='cuda') + (seq_len_kv - seq_len)
|
||||
else:
|
||||
ks, ke = generate_cp_test_data(seq_len, seq_len_kv)
|
||||
# Calculate reference logits
|
||||
ref_logits, ref_cost = ref_fp8_mqa_logits(q, kv, weights, ks, ke)
|
||||
|
||||
q_fp8 = q.to(torch.float8_e4m3fn)
|
||||
kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0, ), False)
|
||||
# Quantize Q and KV to FP4 / FP8
|
||||
if is_fp4:
|
||||
q_fp4 = per_token_cast_to_fp4(q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
|
||||
q_in = (q_fp4[0].view(seq_len, num_heads, head_dim // 2), q_fp4[1].view(seq_len, num_heads))
|
||||
q_simulated = cast_back_from_fp4(q_fp4[0], q_fp4[1], gran_k=32, use_packed_ue8m0=True).view(seq_len, num_heads, head_dim).to(torch.bfloat16)
|
||||
|
||||
if compressed_logits:
|
||||
max_seqlen_k = (ke - ks).max().item()
|
||||
logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke, max_seqlen_k=max_seqlen_k, clean_logits=False)
|
||||
assert logits.size() == (seq_len, max_seqlen_k)
|
||||
tmp = torch.full((seq_len, seq_len_kv), float('-inf'), device='cuda')
|
||||
for i in range(seq_len):
|
||||
tmp[i, ks[i] : ke[i]] = logits[i, : ke[i] - ks[i]]
|
||||
logits = tmp
|
||||
else:
|
||||
logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke)
|
||||
kv_fp4 = per_token_cast_to_fp4(kv.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
|
||||
kv_in = (kv_fp4[0].view(seq_len_kv, head_dim // 2), kv_fp4[1].view(seq_len_kv))
|
||||
kv_simulated = cast_back_from_fp4(kv_fp4[0], kv_fp4[1], gran_k=32, use_packed_ue8m0=True).view(seq_len_kv, head_dim).to(torch.bfloat16)
|
||||
else:
|
||||
q_in = q.to(torch.float8_e4m3fn), None
|
||||
q_simulated = q_in[0].to(torch.bfloat16)
|
||||
kv_in = per_custom_dims_cast_to_fp8(kv, (0, ), False)
|
||||
kv_simulated = (kv_in[0].float() * kv_in[1].unsqueeze(1)).to(torch.bfloat16)
|
||||
|
||||
do_check = (seq_len_kv < 32768)
|
||||
if do_check:
|
||||
ref_logits, ref_cost = ref_fp8_mqa_logits(q=q, kv=kv, weights=weights, cu_seqlen_ks=ks, cu_seqlen_ke=ke)
|
||||
# Calculate reference logits
|
||||
simulated_logits, _ = ref_fp8_mqa_logits(q_simulated, kv_simulated, weights, ks, ke)
|
||||
|
||||
ref_neginf_mask = (ref_logits == float('-inf'))
|
||||
neginf_mask = (logits == float('-inf'))
|
||||
assert torch.equal(neginf_mask, ref_neginf_mask)
|
||||
# Prepare kwargs
|
||||
kernel_kwargs = dict(
|
||||
q=q_in, kv=kv_in, weights=weights,
|
||||
cu_seq_len_k_start=ks, cu_seq_len_k_end=ke,
|
||||
clean_logits=clean_logits, max_seqlen_k=0,
|
||||
logits_dtype=logits_dtype
|
||||
)
|
||||
if compressed_logits:
|
||||
max_seqlen_k = (ke - ks).max().item()
|
||||
kernel_kwargs['max_seqlen_k'] = max_seqlen_k
|
||||
|
||||
ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
|
||||
logits = logits.masked_fill(neginf_mask, 0)
|
||||
diff = calc_diff(logits, ref_logits)
|
||||
assert diff < 1e-3, f'{diff=}'
|
||||
else:
|
||||
ref_cost = ref_fp8_mqa_logits(q=q, kv=kv, weights=weights, cu_seqlen_ks=ks, cu_seqlen_ke=ke, cost_only=True)
|
||||
# Run kernel
|
||||
logits = deep_gemm.fp8_fp4_mqa_logits(**kernel_kwargs)
|
||||
|
||||
tflops = 2 * ref_cost * num_heads * head_dim / 1e12
|
||||
if compressed_logits:
|
||||
t = bench_kineto(lambda: deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke, max_seqlen_k=max_seqlen_k, clean_logits=False), 'fp8_mqa_logits')
|
||||
else:
|
||||
t, clean_t = bench_kineto(lambda: deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke), ('fp8_mqa_logits', 'clean_logits'))
|
||||
clean_bytes = (seq_len * seq_len_kv - ref_cost) * 4 + count_bytes(ks, ke)
|
||||
print(f' > S={seq_len:4}, SKV={seq_len_kv:6}, H={num_heads:3}, D={head_dim:3}, CP={0 if disable_cp else 1}: '
|
||||
f'{tflops / t:4.0f} TFLOPS, {t * 1e6:4.0f} us, '
|
||||
f'{(count_bytes(q_fp8, kv_fp8, weights, ks, ke) + ref_cost * 4) / t / 1e9:4.0f} GB/s', end='')
|
||||
# noinspection PyUnboundLocalVariable
|
||||
print(f' | clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s' if not compressed_logits else '')
|
||||
# Post process for compressed logits
|
||||
if compressed_logits:
|
||||
assert logits.size() == (seq_len, max_seqlen_k)
|
||||
tmp = torch.full((seq_len, seq_len_kv), float('-inf'), device='cuda')
|
||||
for i in range(seq_len):
|
||||
tmp[i, ks[i] : ke[i]] = logits[i, : ke[i] - ks[i]]
|
||||
logits = tmp
|
||||
|
||||
# Validation
|
||||
ref_neginf_mask = (ref_logits == float('-inf'))
|
||||
neginf_mask = (logits == float('-inf'))
|
||||
assert torch.equal(neginf_mask, ref_neginf_mask)
|
||||
|
||||
ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
|
||||
simulated_logits = simulated_logits.masked_fill(ref_neginf_mask, 0)
|
||||
logits = logits.masked_fill(ref_neginf_mask, 0)
|
||||
diff = calc_diff(logits, ref_logits)
|
||||
simulated_diff = calc_diff(logits, simulated_logits)
|
||||
assert diff < 0.02 if is_fp4 else 1e-3, f"Diff: {diff}"
|
||||
assert simulated_diff < 5e-6, f"Simulated Diff: {simulated_diff}"
|
||||
|
||||
# Profiling
|
||||
tflops = 2 * ref_cost * num_heads * head_dim / 1e12
|
||||
t, clean_t = bench_kineto(lambda: deep_gemm.fp8_fp4_mqa_logits(**kernel_kwargs), ('mqa_logits', 'clean_logits'))
|
||||
clean_bytes = (seq_len * seq_len_kv - ref_cost) * 4 + count_bytes(ks, ke)
|
||||
|
||||
print(f' > FP4={is_fp4}, BF16={logits_dtype == torch.bfloat16}, S={seq_len:4}, SKV={seq_len_kv:6}, H={num_heads:3}, D={head_dim:3}, CP={0 if disable_cp else 1}: '
|
||||
f'{tflops / t:4.0f} TFLOPS, {t * 1e6:4.0f} us, '
|
||||
f'{(count_bytes(q_in, kv_in, weights, ks, ke) + ref_cost * 4) / t / 1e9:4.0f} GB/s', end='')
|
||||
print(f' | clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s' if clean_logits else '')
|
||||
print()
|
||||
|
||||
|
||||
def ref_fp8_paged_mqa_logits(q: torch.Tensor, kv_cache: torch.Tensor,
|
||||
weights: torch.Tensor, context_lens: torch.Tensor, block_tables: torch.Tensor,
|
||||
max_model_len: int, is_context_lens_2d: bool):
|
||||
batch_size, next_n, heads, dim = q.size()
|
||||
def ref_paged_mqa_logits(q: torch.Tensor, kv_cache: torch.Tensor,
|
||||
weights: torch.Tensor, context_lens: torch.Tensor, block_tables: torch.Tensor,
|
||||
max_model_len: int, use_2d_context_lens: bool):
|
||||
batch_size, next_n, num_heads, dim = q.size()
|
||||
num_block, block_size, _, dim = kv_cache.size()
|
||||
logits = torch.full([batch_size * next_n, max_model_len], float('-inf'), device=q.device, dtype=torch.float32)
|
||||
context_lens = context_lens.tolist()
|
||||
for i in range(batch_size):
|
||||
context_len = context_lens[i]
|
||||
q_offsets = torch.full((next_n, ), context_len, device='cuda', dtype=torch.int32) if is_context_lens_2d \
|
||||
else torch.arange(context_len - next_n, context_len, device='cuda')
|
||||
q_offsets = torch.full((next_n, ), context_len, device='cuda', dtype=torch.int32) if use_2d_context_lens \
|
||||
else torch.arange(context_len - next_n, context_len, device='cuda')
|
||||
weight_slice = weights[i * next_n:(i + 1) * next_n, :].transpose(0, 1).contiguous()
|
||||
|
||||
num_blocks = (context_len + block_size - 1) // block_size
|
||||
block_idxs = block_tables[i][:num_blocks]
|
||||
kv_slice = kv_cache[block_idxs] # [num_blocks, block_size, kv_heads, dim]
|
||||
kx = kv_slice.permute(2, 3, 0, 1).reshape(kv_slice.size(2), dim, -1) # [kv_heads, dim, total_tokens]
|
||||
qx = q[i].transpose(0, 1) # q[i]: [next_n, heads, dim] -> [heads, next_n, dim]
|
||||
s = torch.matmul(qx, kx).to(logits.dtype) # [heads, next_n, dim] @ [1, dim, total_tokens] -> [heads, next_n, total_tokens]
|
||||
qx = q[i].transpose(0, 1) # q[i]: [next_n, num_heads, dim] -> [num_heads, next_n, dim]
|
||||
s = torch.matmul(qx, kx).to(logits.dtype) # [num_heads, next_n, dim] @ [1, dim, total_tokens] -> [num_heads, next_n, total_tokens]
|
||||
|
||||
total_len = num_blocks * block_size
|
||||
k_offsets = torch.arange(0, total_len, device=q.device)
|
||||
mask = (k_offsets[None, :] < context_len) & (k_offsets[None, :] <= q_offsets[:, None])
|
||||
s = torch.where(mask[None, :, :], s, float('-inf')) # mask shape: [1, next_n, total_tokens]
|
||||
s = torch.relu(s) * weight_slice[..., None] # weight_slice: [heads, next_n] -> [heads, next_n, 1]
|
||||
s = torch.relu(s) * weight_slice[..., None] # weight_slice: [num_heads, next_n] -> [num_heads, next_n, 1]
|
||||
s = s.sum(dim=0) # [next_n, total_tokens]
|
||||
logits[i * next_n:(i + 1) * next_n, :total_len] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float('-inf'))
|
||||
|
||||
@@ -206,70 +225,129 @@ def ref_fp8_paged_mqa_logits(q: torch.Tensor, kv_cache: torch.Tensor,
|
||||
|
||||
|
||||
def test_paged_mqa_logits():
|
||||
print('Testing FP8 Paged MQA Logits:')
|
||||
max_model_len = 111 * 1000
|
||||
for is_context_lens_2d in (False, True):
|
||||
for batch_size, next_n in [(64, 1), (64, 2), (128, 1)]:
|
||||
for heads, index_dim in [(64, 128)]:
|
||||
for avg_kv in (8192, 32768):
|
||||
num_blocks, blocksize = max_model_len * 3, 64
|
||||
|
||||
q = torch.randn((batch_size, next_n, heads, index_dim), device='cuda', dtype=torch.bfloat16)
|
||||
kv_cache = torch.randn((num_blocks, blocksize, 1, index_dim), device='cuda', dtype=torch.bfloat16)
|
||||
weights = torch.randn((batch_size * next_n, heads), device='cuda', dtype=torch.float32)
|
||||
q_fp8 = q.to(torch.float8_e4m3fn)
|
||||
kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
|
||||
# Helper functions
|
||||
def kv_cache_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
num_blocks, block_size, num_heads, head_dim = x.shape
|
||||
assert num_heads == 1
|
||||
x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4)
|
||||
sf = x_amax / 448.0
|
||||
x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn)
|
||||
x_cast_back = x_scaled.float() * sf
|
||||
|
||||
context_lens = torch.randint(int(0.7 * avg_kv), int(1.3 * avg_kv), (batch_size, )).cuda().to(torch.int32)
|
||||
context_lens_list = context_lens.tolist()
|
||||
max_block_len = (max(context_lens_list) + blocksize - 1) // blocksize * blocksize
|
||||
block_tables = torch.zeros((batch_size, max_block_len), device='cuda', dtype=torch.int32)
|
||||
x_fp8 = torch.empty((num_blocks, block_size * (head_dim + 4)), device=x.device, dtype=torch.uint8)
|
||||
x_fp8[ :, : block_size * head_dim] = x_scaled.view(num_blocks, block_size * head_dim).view(torch.uint8)
|
||||
x_fp8[ :, block_size * head_dim :] = sf.view(num_blocks, block_size).view(torch.uint8)
|
||||
return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4), x_cast_back.to(x.dtype)
|
||||
|
||||
counter, block_idx_pool = 0, torch.randperm(num_blocks, device='cuda', dtype=torch.int32)
|
||||
for i in range(batch_size):
|
||||
num_blocks = ceil_div(context_lens_list[i], blocksize)
|
||||
block_tables[i][:num_blocks] = block_idx_pool[counter: counter+num_blocks]
|
||||
counter += num_blocks
|
||||
def kv_cache_cast_to_fp4(x: torch.Tensor) -> torch.Tensor:
|
||||
num_blocks, block_size, num_heads, head_dim = x.shape
|
||||
assert num_heads == 1 and head_dim == 128
|
||||
x_scaled, sf = per_token_cast_to_fp4(x.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
|
||||
x_cast_back = cast_back_from_fp4(x_scaled, sf, gran_k=32, use_packed_ue8m0=True).view(num_blocks, block_size, 1, head_dim)
|
||||
|
||||
ref_logits = ref_fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_tables, max_model_len, is_context_lens_2d)
|
||||
positions = torch.arange(max_model_len, device='cuda').unsqueeze(0).expand(batch_size * next_n, -1)
|
||||
x_fp4 = torch.empty((num_blocks, block_size * (head_dim // 2 + 4)), device=x.device, dtype=torch.uint8)
|
||||
x_fp4[ :, : block_size * head_dim // 2] = x_scaled.view(num_blocks, block_size * head_dim // 2).view(torch.uint8)
|
||||
x_fp4[ :, block_size * head_dim // 2 :] = sf.view(num_blocks, block_size).view(torch.uint8)
|
||||
return x_fp4.view(num_blocks, block_size, num_heads, head_dim // 2 + 4), x_cast_back.to(x.dtype)
|
||||
|
||||
if is_context_lens_2d:
|
||||
context_lens_2d = ((context_lens.unsqueeze(1) + 1) * torch.rand(batch_size, next_n, device='cuda')).int()
|
||||
context_lens_2d[:, next_n-1] = context_lens
|
||||
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(context_lens_2d, blocksize, deep_gemm.get_num_sms())
|
||||
logits = deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens_2d, block_tables, schedule_metadata, max_model_len, clean_logits=False)
|
||||
ref_neginf_mask = ~(positions < context_lens_2d.view(-1).unsqueeze(1))
|
||||
else:
|
||||
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(context_lens, blocksize, deep_gemm.get_num_sms())
|
||||
logits = deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens, block_tables, schedule_metadata, max_model_len, clean_logits=True)
|
||||
row_indices = torch.arange(batch_size * next_n, device='cuda') // next_n
|
||||
next_n_offset = torch.arange(batch_size * next_n, device='cuda') % next_n
|
||||
ref_neginf_mask = ~(positions <= (context_lens[row_indices] - next_n + next_n_offset).unsqueeze(1))
|
||||
neginf_mask = (logits == float('-inf'))
|
||||
assert torch.equal(neginf_mask, ref_neginf_mask)
|
||||
def enumerate_paged_mqa_logits():
|
||||
arch_major = get_arch_major()
|
||||
for is_fp4 in ((True, False) if arch_major == 10 else (False, )):
|
||||
for logits_dtype in (torch.float, torch.bfloat16):
|
||||
for block_kv in ((32, 64) if arch_major == 10 else (64, )):
|
||||
for use_2d_context_lens, clean_logits in [(True, False)]:
|
||||
for batch_size in (256, ):
|
||||
for next_n in (1, 2, 4, 5, 6) if arch_major == 10 else (1, 2):
|
||||
for num_heads, head_dim in [(64, 128)]:
|
||||
for avg_kv in (8192, 32768):
|
||||
yield is_fp4, logits_dtype, block_kv, use_2d_context_lens, clean_logits, batch_size, next_n, num_heads, head_dim, avg_kv
|
||||
|
||||
logits = logits.masked_fill(ref_neginf_mask, 0)
|
||||
ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
|
||||
diff = calc_diff(logits, ref_logits)
|
||||
assert diff < 1e-3, f"{diff=}"
|
||||
|
||||
sum_lens = sum(context_lens.to(torch.int64))
|
||||
tflops = 2 * sum_lens * next_n * heads * index_dim / 1e12
|
||||
input_bytes = count_bytes(q_fp8, weights, context_lens) + sum_lens * (index_dim + 4) + (sum_lens / blocksize) * 4
|
||||
output_bytes = sum_lens * next_n * 4
|
||||
if is_context_lens_2d:
|
||||
t = bench_kineto(lambda: deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens_2d, block_tables, schedule_metadata, max_model_len, clean_logits=False),
|
||||
'fp8_paged_mqa_logits')
|
||||
else:
|
||||
t, clean_t = bench_kineto(lambda: deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens, block_tables, schedule_metadata, max_model_len, clean_logits=True),
|
||||
('fp8_paged_mqa_logits', 'clean_logits'))
|
||||
clean_bytes = (batch_size * next_n * max_model_len - neginf_mask.sum().item()) * 4 + count_bytes(context_lens)
|
||||
print(f' > BSZ={batch_size:3}, NextN={next_n:1}, H={heads:2}, D={index_dim:2}, L={avg_kv:6}: '
|
||||
f'{tflops / t:4.0f} TFLOPS, {t * 1e6:3.0f} us, '
|
||||
f'{(input_bytes + output_bytes) / t / 1e9:4.0f} GB/s', end='')
|
||||
# noinspection PyUnboundLocalVariable
|
||||
print(f' | clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s' if not is_context_lens_2d else '')
|
||||
print('Testing FP8/FP4 Paged MQA Logits:')
|
||||
max_model_len = 111 * 1024
|
||||
num_total_blocks = max_model_len * 5
|
||||
|
||||
for is_fp4, logits_dtype, block_kv, use_2d_context_lens, clean_logits, batch_size, next_n, num_heads, head_dim, avg_kv in enumerate_paged_mqa_logits():
|
||||
# Generate random inputs
|
||||
q = torch.randn((batch_size, next_n, num_heads, head_dim), device='cuda', dtype=torch.bfloat16)
|
||||
kv_cache = torch.randn((num_total_blocks, block_kv, 1, head_dim), device='cuda', dtype=torch.bfloat16)
|
||||
weights = torch.randn((batch_size * next_n, num_heads), device='cuda', dtype=torch.float)
|
||||
context_lens = torch.randint(int(0.7 * avg_kv), int(1.3 * avg_kv), (batch_size,), device='cuda', dtype=torch.int)
|
||||
|
||||
# Assign block tables
|
||||
num_blocks_per_query = ceil_div(context_lens, block_kv)
|
||||
block_table = torch.empty((batch_size, num_blocks_per_query.max().item()), device='cuda', dtype=torch.int)
|
||||
block_idx_pool = torch.randperm(num_total_blocks, device='cuda', dtype=torch.int)
|
||||
offset = 0
|
||||
for i, num_blocks in enumerate(num_blocks_per_query.tolist()):
|
||||
block_table[i, :num_blocks] = block_idx_pool[offset : offset + num_blocks]
|
||||
offset += num_blocks
|
||||
|
||||
# Calculate reference logits
|
||||
ref_logits = ref_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, max_model_len, use_2d_context_lens)
|
||||
|
||||
# Quantize Q and KV cache to FP4 / FP8
|
||||
if is_fp4:
|
||||
q_fp4 = per_token_cast_to_fp4(q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
|
||||
q_in = (q_fp4[0].view(batch_size, next_n, num_heads, head_dim // 2), q_fp4[1].view(batch_size, next_n, num_heads))
|
||||
q_simulated = cast_back_from_fp4(q_fp4[0], q_fp4[1], gran_k=32, use_packed_ue8m0=True).view(batch_size, next_n, num_heads, head_dim).to(torch.bfloat16)
|
||||
kv_in, kv_simulated = kv_cache_cast_to_fp4(kv_cache)
|
||||
else:
|
||||
q_in = q.to(torch.float8_e4m3fn), None
|
||||
q_simulated = q_in[0].to(torch.bfloat16)
|
||||
kv_in, kv_simulated = kv_cache_cast_to_fp8(kv_cache)
|
||||
|
||||
# Calculate simulated reference logits
|
||||
simulated_logits = ref_paged_mqa_logits(q_simulated, kv_simulated, weights, context_lens, block_table, max_model_len, use_2d_context_lens)
|
||||
|
||||
# Prepare masks and context lengths with NextN
|
||||
positions = torch.arange(max_model_len, device='cuda').unsqueeze(0).expand(batch_size * next_n, -1)
|
||||
if use_2d_context_lens:
|
||||
context_lens_nextn = ((context_lens.unsqueeze(1) + 1) * torch.rand(batch_size, next_n, device='cuda')).int()
|
||||
# Ensure last token matches actual length
|
||||
context_lens_nextn[:, -1] = context_lens
|
||||
ref_neginf_mask = ~(positions < context_lens_nextn.view(-1, 1))
|
||||
else:
|
||||
context_lens_nextn = context_lens
|
||||
offsets = torch.arange(batch_size * next_n, device='cuda')
|
||||
limits = (context_lens[offsets // next_n] - next_n + offsets % next_n).unsqueeze(1)
|
||||
ref_neginf_mask = ~(positions <= limits)
|
||||
|
||||
# Run Kernel
|
||||
kernel_kwargs = dict(
|
||||
q=q_in, kv_cache=kv_in, weights=weights,
|
||||
context_lens=context_lens_nextn, block_table=block_table,
|
||||
schedule_meta=deep_gemm.get_paged_mqa_logits_metadata(context_lens_nextn, block_kv, deep_gemm.get_num_sms()),
|
||||
max_context_len=max_model_len, clean_logits=clean_logits, logits_dtype=logits_dtype
|
||||
)
|
||||
logits = deep_gemm.fp8_fp4_paged_mqa_logits(**kernel_kwargs)
|
||||
|
||||
# Validation
|
||||
assert logits.dtype == logits_dtype
|
||||
logits = logits.to(torch.float)
|
||||
|
||||
if clean_logits:
|
||||
assert torch.equal(logits == float('-inf'), ref_neginf_mask), "Mask mismatch"
|
||||
|
||||
logits_masked = logits.masked_fill(ref_neginf_mask, 0)
|
||||
ref_masked = ref_logits.masked_fill(ref_neginf_mask, 0)
|
||||
simulated_masked = simulated_logits.masked_fill(ref_neginf_mask, 0)
|
||||
diff = calc_diff(logits_masked, ref_masked)
|
||||
simulated_diff = calc_diff(logits_masked, simulated_masked)
|
||||
assert diff < 0.02 if is_fp4 else 1e-3, f"Diff: {diff}"
|
||||
assert simulated_diff < 5e-6, f"Simulated Diff: {simulated_diff}"
|
||||
|
||||
# Profiling
|
||||
sum_lens = context_lens.sum().item()
|
||||
tflops_calc = 2 * sum_lens * next_n * num_heads * head_dim / 1e12
|
||||
kv_bytes_per_token = head_dim / (2 if is_fp4 else 1) + 4
|
||||
total_bytes = count_bytes(q, weights) + sum_lens * kv_bytes_per_token + (sum_lens * next_n * logits_dtype.itemsize)
|
||||
|
||||
t, clean_t = bench_kineto(lambda: deep_gemm.fp8_fp4_paged_mqa_logits(**kernel_kwargs), ('paged_mqa_logits', 'clean_logits'))
|
||||
print(f' > FP4={is_fp4}, BF16={logits_dtype == torch.bfloat16}, BLOCK_KV={block_kv}, BSZ={batch_size:3}, NextN={next_n:1}, H={num_heads:2}, D={head_dim:2}, L={avg_kv:6}: '
|
||||
f'{tflops_calc / t:4.0f} TFLOPS, {t * 1e6:3.0f} us, {total_bytes / t / 1e9:4.0f} GB/s', end='')
|
||||
print(f' | clean: {clean_t*1e6:3.0f} us' if clean_logits else '')
|
||||
print()
|
||||
|
||||
|
||||
@@ -280,6 +358,5 @@ if __name__ == '__main__':
|
||||
random.seed(0)
|
||||
|
||||
test_gemm_skip_head_mid()
|
||||
|
||||
test_mqa_logits()
|
||||
test_paged_mqa_logits()
|
||||
|
||||
Reference in New Issue
Block a user