diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index b91158f5ada..e5467f3ed30 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -37,35 +37,43 @@ f"got `{installed_flydsl_version}`." ) - from .gemm_kernels import flydsl_hgemm, flydsl_preshuffle_gemm_a8 - from .moe_kernels import flydsl_moe_stage1, flydsl_moe_stage2 from .fmha_kernels import flydsl_flash_attn_func - from .kernels.qk_norm_rope_quant import flydsl_qk_norm_rope_quant + from .gemm_kernels import flydsl_hgemm, flydsl_preshuffle_gemm_a8 + from .kernels.fp8_mqa_logits import ( + DEFAULT_VARIANT as FP8_MQA_LOGITS_DEFAULT_VARIANT, + ) + from .kernels.fp8_mqa_logits import ( + KERNEL_VARIANTS as FP8_MQA_LOGITS_VARIANTS, + ) + from .kernels.fp8_mqa_logits import ( + flydsl_fp8_mqa_logits, + ) from .kernels.pa_mqa_logits_fp4 import ( flydsl_pa_mqa_logits_fp4, ) from .kernels.pa_mqa_logits_fp4_prefill import ( + compute_varqlen_windows, flydsl_pa_mqa_logits_fp4_prefill, + flydsl_pa_mqa_logits_fp4_varqlen, ) - from .kernels.fp8_mqa_logits import ( - flydsl_fp8_mqa_logits, - KERNEL_VARIANTS as FP8_MQA_LOGITS_VARIANTS, - DEFAULT_VARIANT as FP8_MQA_LOGITS_DEFAULT_VARIANT, - ) + from .kernels.qk_norm_rope_quant import flydsl_qk_norm_rope_quant + from .moe_kernels import flydsl_moe_stage1, flydsl_moe_stage2 # from .linear_attention_kernels import flydsl_gdr_decode __all__ += [ - "flydsl_preshuffle_gemm_a8", + "FP8_MQA_LOGITS_DEFAULT_VARIANT", + "FP8_MQA_LOGITS_VARIANTS", + "compute_varqlen_windows", + "flydsl_flash_attn_func", + "flydsl_fp8_mqa_logits", + "flydsl_hgemm", "flydsl_moe_stage1", "flydsl_moe_stage2", - "flydsl_hgemm", - "flydsl_flash_attn_func", - "flydsl_qk_norm_rope_quant", "flydsl_pa_mqa_logits_fp4", "flydsl_pa_mqa_logits_fp4_prefill", - "flydsl_fp8_mqa_logits", - "FP8_MQA_LOGITS_VARIANTS", - "FP8_MQA_LOGITS_DEFAULT_VARIANT", + "flydsl_pa_mqa_logits_fp4_varqlen", + "flydsl_preshuffle_gemm_a8", + "flydsl_qk_norm_rope_quant", # "flydsl_gdr_decode", ] diff --git a/aiter/ops/flydsl/kernels/pa_mqa_logits_fp4.py b/aiter/ops/flydsl/kernels/pa_mqa_logits_fp4.py index bfababf4b0c..1e95fa6dae7 100644 --- a/aiter/ops/flydsl/kernels/pa_mqa_logits_fp4.py +++ b/aiter/ops/flydsl/kernels/pa_mqa_logits_fp4.py @@ -126,11 +126,14 @@ def compute_varctx_schedule( next_n=1, cta_info_out=None, ): + B = context_lens.shape[0] + if parallel_unit_num is None: + chunks_per_seq = max(1, (max_seq_len + block_k - 1) // block_k) + parallel_unit_num = B * next_n * chunks_per_seq P = parallel_unit_num if P % next_n != 0: raise ValueError(f"parallel_unit_num={P} must be a multiple of next_n={next_n}") S = P // next_n - B = context_lens.shape[0] if S < B: raise ValueError( f"compute_varctx_schedule: parallel_unit_num//next_n={S} < batches={B} " @@ -691,13 +694,20 @@ def flydsl_pa_mqa_logits_fp4( block_k: int = 256, kv_block_size: int = 64, num_warps: int = DEFAULT_NUM_WARPS, - parallel_unit_num: int = 512, + parallel_unit_num: Optional[int] = None, out: Optional[torch.Tensor] = None, cta_info: Optional[torch.Tensor] = None, total_ctas: Optional[int] = None, stream: Optional[torch.cuda.Stream] = None, ) -> torch.Tensor: - """Decode/varctx FP4 paged MQA logits (gfx950).""" + """Decode/varctx FP4 paged MQA logits (gfx950). + + ``parallel_unit_num`` is the persistent-grid CTA count; when ``None`` it is + auto-derived (cudagraph-safe, no device→host sync) as + ``batch * next_n * ceil(max_seq_len / block_k)``, which is a multiple of + ``next_n`` and ``>= batch*next_n`` by construction. Pass a smaller explicit + value to trade parallelism for fewer no-op CTAs. + """ batch_size, q_next_n, heads, head_dim_packed = q_fp4.shape head_dim = head_dim_packed * 2 max_blocks_per_seq = block_tables.shape[1] diff --git a/aiter/ops/flydsl/kernels/pa_mqa_logits_fp4_prefill.py b/aiter/ops/flydsl/kernels/pa_mqa_logits_fp4_prefill.py index 5771fd6b773..c817626258e 100644 --- a/aiter/ops/flydsl/kernels/pa_mqa_logits_fp4_prefill.py +++ b/aiter/ops/flydsl/kernels/pa_mqa_logits_fp4_prefill.py @@ -55,8 +55,14 @@ def compute_prefill_schedule( block_k, parallel_unit_num, max_seq_len, + cta_info_out=None, ): - """Compute the persistent-grid schedule for ragged-prefill MQA logits.""" + """Compute the persistent-grid schedule for ragged-prefill MQA logits. + + Pass `cta_info_out` (a fixed [parallel_unit_num, CTA_INFO_WIDTH] int32 buffer) + to write the schedule into a stable address (CUDAGraph decode: the captured + kernel replays from this pointer while `build()` refreshes its contents). + """ device = local_ends.device P = parallel_unit_num T = local_ends.shape[0] # fixed total_tokens (rows) @@ -94,7 +100,10 @@ def compute_prefill_schedule( # ── map each fixed slot → (row, split) + emit cta_info in ONE kernel ── # (the ~25 per-slot torch ops below were the bulk of the ~50-launch cost). - cta_info = torch.empty(P, CTA_INFO_WIDTH, dtype=torch.int32, device=device) + if cta_info_out is None: + cta_info = torch.empty(P, CTA_INFO_WIDTH, dtype=torch.int32, device=device) + else: + cta_info = cta_info_out safe_i32 = safe.reshape(1).to(torch.int32) total_splits_i32 = total_splits.reshape(1).to(torch.int32) BLOCK_P = 256 @@ -756,3 +765,143 @@ def flydsl_pa_mqa_logits_fp4_prefill( stream, ) return out + + +@triton.jit +def _varqlen_windows_kernel( + cu_ptr, # [B+1] int32, prefix-sum of per-batch qlen + ctx_ptr, # [B] int32, per-batch KV length + row_to_batch_ptr, # [total_q] int32 (out) + local_starts_ptr, # [total_q] int32 (out) + local_ends_ptr, # [total_q] int32 (out) + total_q, + B, + BLOCK: tl.constexpr, +): + """Fused build of ragged-row metadata for per-batch variable qlen (MTP).""" + pid = tl.program_id(0) + r = pid * BLOCK + tl.arange(0, BLOCK) + rmask = r < total_q + + lo = tl.zeros([BLOCK], tl.int32) + hi = tl.full([BLOCK], B, tl.int32) + for _ in tl.static_range(32): + mid = (lo + hi) // 2 + cu_mid = tl.load( + cu_ptr + 1 + tl.minimum(mid, B - 1), mask=(mid < B), other=2147483647 + ) + go_right = cu_mid <= r + lo = tl.where(go_right, mid + 1, lo) + hi = tl.where(go_right, hi, mid) + b = tl.minimum(lo, B - 1) + + cu_b = tl.load(cu_ptr + b, mask=rmask, other=0) + cu_b1 = tl.load(cu_ptr + b + 1, mask=rmask, other=0) + ctx_b = tl.load(ctx_ptr + b, mask=rmask, other=0) + n = r - cu_b + qlen = cu_b1 - cu_b + le = tl.maximum(ctx_b - qlen + n + 1, 0) + # Rows beyond the real total Σ (cu[B]) are FLAT tail-padding — force an empty + # window so the mqa kernel / top_k skip them (used when `total_q` is the padded + # count, e.g. the CUDAGraph decode path scores all padded rows in one shot). + real_total = tl.load(cu_ptr + B) + le = tl.where(r < real_total, le, tl.zeros([BLOCK], tl.int32)) + + tl.store(row_to_batch_ptr + r, b, mask=rmask) + tl.store(local_starts_ptr + r, tl.zeros([BLOCK], tl.int32), mask=rmask) + tl.store(local_ends_ptr + r, le, mask=rmask) + + +def compute_varqlen_windows(cu_seq_q, context_lens, total_q, *, out=None): + """Build ragged-row metadata for per-batch variable query length (MTP). + + Pass `out=(row_to_batch, local_starts, local_ends)` (fixed int32 buffers each + >= total_q long) to write into stable addresses — the CUDAGraph decode path + scores all padded rows, so top_k replays from these window pointers while + `build()` refreshes their contents. Rows past the real total (cu[B]) get an + empty window (local_ends == 0) so they are skipped. + """ + dev = cu_seq_q.device + cu = cu_seq_q.to(torch.int32).contiguous() + ctx = context_lens.to(torch.int32).contiguous() + B = ctx.shape[0] + if out is None: + row_to_batch = torch.empty(total_q, dtype=torch.int32, device=dev) + local_starts = torch.empty(total_q, dtype=torch.int32, device=dev) + local_ends = torch.empty(total_q, dtype=torch.int32, device=dev) + else: + row_to_batch, local_starts, local_ends = out + if total_q > 0: + BLOCK = 256 + grid = (triton.cdiv(total_q, BLOCK),) + _varqlen_windows_kernel[grid]( + cu, + ctx, + row_to_batch, + local_starts, + local_ends, + total_q, + B, + BLOCK=BLOCK, + ) + return row_to_batch, local_starts, local_ends + + +def flydsl_pa_mqa_logits_fp4_varqlen( + q_fp4: torch.Tensor, + q_scale: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: torch.Tensor, + block_tables: torch.Tensor, + weights: torch.Tensor, + max_seq_len: int, + *, + cu_seq_q: torch.Tensor | None = None, + context_lens: torch.Tensor | None = None, + windows: tuple | None = None, + weight_scale: float = 1.0, + block_k: int = 256, + kv_block_size: int = 64, + num_warps: int = DEFAULT_NUM_WARPS, + parallel_unit_num: int | None = None, + out: torch.Tensor | None = None, + cta_info: torch.Tensor | None = None, + n_ctas: int | None = None, + stream: torch.cuda.Stream | None = None, +) -> torch.Tensor: + """Variable-qlen (per-batch MTP) FP4 paged MQA logits (gfx950).""" + total_q = q_fp4.shape[0] + if windows is None: + if cu_seq_q is None or context_lens is None: + raise ValueError( + "flydsl_pa_mqa_logits_fp4_varqlen: pass windows=(row_to_batch, " + "local_starts, local_ends) built once via " + "compute_varqlen_windows, or both cu_seq_q and context_lens " + "to build them here." + ) + windows = compute_varqlen_windows(cu_seq_q, context_lens, total_q) + row_to_batch, local_starts, local_ends = windows + if parallel_unit_num is None: + chunks_per_seq = max(1, (max_seq_len + block_k - 1) // block_k) + parallel_unit_num = total_q * chunks_per_seq + return flydsl_pa_mqa_logits_fp4_prefill( + q_fp4, + q_scale, + kv_cache, + kv_scale, + block_tables, + weights, + row_to_batch, + local_starts, + local_ends, + max_seq_len, + weight_scale=weight_scale, + block_k=block_k, + kv_block_size=kv_block_size, + num_warps=num_warps, + parallel_unit_num=parallel_unit_num, + out=out, + cta_info=cta_info, + n_ctas=n_ctas, + stream=stream, + ) diff --git a/op_tests/test_flydsl_pa_mqa_logits_fp4.py b/op_tests/test_flydsl_pa_mqa_logits_fp4.py index 06532ff974d..6b0f9f41e7e 100644 --- a/op_tests/test_flydsl_pa_mqa_logits_fp4.py +++ b/op_tests/test_flydsl_pa_mqa_logits_fp4.py @@ -366,7 +366,7 @@ def test_pa_mqa_logits_fp4_qfp4_kvfp4( num_iters=20, num_warmup=3, num_warps=4, - parallel_unit_num=512, + parallel_unit_num=None, head_dim=DEFAULT_HEAD_DIM, bench=True, ): @@ -465,12 +465,16 @@ def test_pa_mqa_logits_fp4_qfp4_kvfp4( # The persistent-grid schedule has S = parallel_unit_num // next_n batch # slots; if batch_size exceeds S the surplus batches are silently dropped - # (their out stays -inf -> NaN cosine). Grow the grid so every (batch, - # next_n) gets at least one slot. - parallel_unit_num = max(parallel_unit_num, batch_size * next_n) + # (their out stays -inf -> NaN cosine). For an explicit value grow the grid + # so every (batch, next_n) gets at least one slot; when None, let the + # scheduler auto-derive a cudagraph-safe grid (= batch*next_n*ceil(t_max/ + # block_k)), which already satisfies both constraints. + if parallel_unit_num is not None: + parallel_unit_num = max(parallel_unit_num, batch_size * next_n) safe, cta_info, total_ctas = compute_varctx_schedule( context_lens, block_k, parallel_unit_num, t_max, next_n=next_n ) + parallel_unit_num = total_ctas print( f" schedule: parallel_unit={parallel_unit_num} num_warps={num_warps} " f"safe_chunks_per_cta={safe} total_ctas={total_ctas} " @@ -666,8 +670,9 @@ def main(): parser.add_argument( "--parallel_unit_num", type=int, - default=512, - help="target CTA count for host schedule (default 512)", + default=None, + help="target CTA count for host schedule " + "(default: auto = batch*next_n*ceil(max_seq_len/block_k))", ) parser.add_argument( "--next_n", diff --git a/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py b/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py index c39f2b224a1..a0d6ac9731e 100644 --- a/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py +++ b/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py @@ -2,13 +2,18 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""aiter op-test for ``flydsl_pa_mqa_logits_fp4_prefill``. +"""aiter op-test for ``flydsl_pa_mqa_logits_fp4_prefill`` and +``flydsl_pa_mqa_logits_fp4_varqlen``. Validates the ragged-prefill FP4 paged MQA logits kernel (gfx950) against a pure-torch reference. Each query row owns a seq-local window ``[local_start, local_end)`` into its sequence's paged FP4 KV cache, read straight from ``block_tables`` (no cp_gather staging). +The variable-qlen path (``flydsl_pa_mqa_logits_fp4_varqlen``) is the same kernel +driven by a ``cu_seq_q`` (per-batch qlen prefix-sum) adapter with MTP tail-causal +windows ``[0, ctx_b - qlen_b + n]``; those cases also report TFLOPS / GB/s. + Accuracy (torch ref): - vs exact FP4-dequant ref -> kernel correctness (cos ~ 1.0) - vs full-precision bf16 ref -> FP4 quant accuracy @@ -456,6 +461,241 @@ def atom_logits(): print() +# ── Variable-qlen (per-batch MTP) via cu_seq_q ─────────────────────── + +_VARQLEN_PERF_SUMMARY = [] + + +def _mtp_windows_ref(qlens, ctxs): + """Independent (python) construction of the MTP tail-causal windows, used + both for the torch reference and to unit-check the device-side builder.""" + rb, ls, le = [], [], [] + for b, (ql, ctx) in enumerate(zip(qlens, ctxs)): + for n in range(ql): + rb.append(b) + ls.append(0) + le.append(max(0, ctx - ql + n + 1)) + return rb, ls, le + + +def run_varqlen_case( + qlens, + ctxs, + heads=64, + head_dim=128, + kv_block_size=64, + block_k=256, + seed=0, + bench=True, + iters=50, + warmup=10, +): + from aiter.ops.flydsl import ( + flydsl_pa_mqa_logits_fp4_prefill, + flydsl_pa_mqa_logits_fp4_varqlen, + ) + from aiter.ops.flydsl.kernels.pa_mqa_logits_fp4_prefill import ( + compute_prefill_schedule, + compute_varqlen_windows, + ) + + torch.manual_seed(seed) + bs = len(qlens) + assert len(ctxs) == bs + total_q = int(sum(qlens)) + + max_end = max(ctxs) + max_blocks_per_seq = max( + (max_end + block_k - 1) // block_k * (block_k // kv_block_size), + block_k // kv_block_size, + ) + t_max = max_blocks_per_seq * kv_block_size + max_seq_len = t_max + num_blocks = max_blocks_per_seq * bs + + # ---- paged FP4 KV (dense per-batch cache, contiguous block_tables) ---- + kv_bf16 = torch.randn(bs, t_max, head_dim, dtype=torch.bfloat16, device=dev) + block_tables = torch.arange(num_blocks, dtype=torch.int32, device=dev).reshape( + bs, max_blocks_per_seq + ) + kv_fp4_d, kv_e8_d = fp4_quant_e2m1_with_e8m0(kv_bf16.reshape(-1, head_dim)) + kv_dq = fp4_dequant_e2m1_with_e8m0( + kv_fp4_d.reshape(bs, t_max, head_dim // 2), + kv_e8_d.reshape(bs, t_max, head_dim // 32), + ) + k_flat = kv_bf16.reshape(bs * t_max, head_dim) + tb = torch.arange(bs, device=dev).repeat_interleave(t_max) + tt = torch.arange(t_max, device=dev).repeat(bs) + phys = block_tables[tb, tt // kv_block_size].long() + slot_mapping = (phys * kv_block_size + (tt % kv_block_size)).to(torch.int32) + k_tiles = head_dim // 128 + kv_cache = torch.zeros( + num_blocks, k_tiles, 4, kv_block_size, 16, dtype=torch.uint8, device=dev + ) + kv_scale = torch.zeros( + num_blocks, k_tiles, 4, kv_block_size, dtype=torch.uint8, device=dev + ) + indexer_k_fp4_paged_preshuffle( + k_flat, slot_mapping, kv_cache, kv_scale, kv_block_size + ) + + # ---- packed (ragged) Q / weights: [total_q, ...] ---- + q_bf16 = torch.randn(total_q, heads, head_dim, dtype=torch.bfloat16, device=dev) + weights = (torch.randn(total_q, heads, dtype=torch.float32, device=dev) * 0.1).to( + torch.bfloat16 + ) + weight_scale = 1.5 + q_fp4, q_scale = quant_q_fp4_preshuffle(q_bf16) + q_e8 = fp4_quant_e2m1_with_e8m0(q_bf16.reshape(total_q * heads, head_dim))[ + 1 + ].reshape(total_q, heads, head_dim // 32) + q_dq = fp4_dequant_e2m1_with_e8m0( + q_fp4.reshape(total_q, heads, head_dim // 2), q_e8 + ) + + # ---- cu_seq_q (prefix-sum of per-batch qlen) + per-batch KV length ---- + cu = [0] + for ql in qlens: + cu.append(cu[-1] + ql) + cu_seq_q = torch.tensor(cu, dtype=torch.int32, device=dev) + context_lens = torch.tensor(ctxs, dtype=torch.int32, device=dev) + + # ---- unit-check the device builder against the independent python windows ---- + rb_ref, ls_ref, le_ref = _mtp_windows_ref(qlens, ctxs) + rb_k, ls_k, le_k = compute_varqlen_windows(cu_seq_q, context_lens, total_q) + assert rb_k.tolist() == rb_ref, f"row_to_batch: {rb_k.tolist()} != {rb_ref}" + assert ls_k.tolist() == ls_ref, "local_starts mismatch" + assert le_k.tolist() == le_ref, f"local_ends: {le_k.tolist()} != {le_ref}" + + ref_fp4 = ref_prefill_logits( + q_dq, kv_dq, weights, rb_ref, ls_ref, le_ref, max_seq_len, weight_scale + ) + ref_bf16 = ref_prefill_logits( + q_bf16, kv_bf16, weights, rb_ref, ls_ref, le_ref, max_seq_len, weight_scale + ) + + out = flydsl_pa_mqa_logits_fp4_varqlen( + q_fp4, + q_scale, + kv_cache, + kv_scale, + block_tables, + weights, + max_seq_len, + cu_seq_q=cu_seq_q, + context_lens=context_lens, + weight_scale=weight_scale, + block_k=block_k, + kv_block_size=kv_block_size, + ) + torch.cuda.synchronize() + + # "build once, reuse" path: passing prebuilt windows must match the + # cu_seq_q path bit-for-bit (windows are built once, kernel called many). + out_reuse = flydsl_pa_mqa_logits_fp4_varqlen( + q_fp4, + q_scale, + kv_cache, + kv_scale, + block_tables, + weights, + max_seq_len, + windows=(rb_k, ls_k, le_k), + weight_scale=weight_scale, + block_k=block_k, + kv_block_size=kv_block_size, + ) + torch.cuda.synchronize() + assert torch.equal(out, out_reuse), "windows= reuse path differs from cu_seq_q path" + + m = ~torch.isneginf(ref_fp4) + cos_exact = _cos(out[m], ref_fp4[m]).item() + cos_bf16 = _cos(out[m], ref_bf16[m]).item() + oob_ok = bool(torch.isneginf(out[~m]).all().item()) if (~m).any() else True + print( + f" qlens={qlens} ctxs={ctxs} heads={heads} total_q={total_q} " + f"cos_exact={cos_exact:.6f} cos_bf16={cos_bf16:.6f} oob_neginf={oob_ok}" + ) + assert cos_exact > 0.99, f"kernel vs FP4-dequant ref cos {cos_exact:.4f} < 0.99" + assert cos_bf16 > 0.95, f"kernel vs bf16 ref cos {cos_bf16:.4f} < 0.95" + assert oob_ok, "OOB cells were not left at -inf" + + if not bench: + return + + # ---- Perf: isolate the kernel launch (precompute the schedule once) ---- + pun = total_q * max(1, (max_seq_len + block_k - 1) // block_k) + _, cta_info, n_ctas = compute_prefill_schedule( + rb_k, ls_k, le_k, block_k, pun, max_seq_len + ) + out_buf = torch.full( + (total_q, max_seq_len), float("-inf"), dtype=torch.float32, device=dev + ) + + def run_fp4(): + flydsl_pa_mqa_logits_fp4_prefill( + q_fp4, + q_scale, + kv_cache, + kv_scale, + block_tables, + weights, + rb_k, + ls_k, + le_k, + max_seq_len, + weight_scale=weight_scale, + block_k=block_k, + kv_block_size=kv_block_size, + parallel_unit_num=pun, + out=out_buf, + cta_info=cta_info, + n_ctas=n_ctas, + ) + + _, us = run_perftest(run_fp4, num_iters=iters, num_warmup=warmup) + + # USEFUL work (exact ragged): (row, token) pairs = sum of window lengths. + # Bytes: Q/weights/out counted once; KV as the per-batch unique footprint + # (max window over the batch's rows), i.e. assuming cross-row KV reuse (L2). + head_dim_packed, head_dim_scales = head_dim // 2, head_dim // 32 + win = [int(e - s) for s, e in zip(ls_ref, le_ref)] + pairs = sum(win) + kv_uni = {} + for b, e in zip(rb_ref, le_ref): + kv_uni[b] = max(kv_uni.get(b, 0), int(e)) + kv_tokens = sum(kv_uni.values()) + flops = pairs * heads * (2 * head_dim + 3) + bytes_total = ( + total_q * heads * (head_dim_packed + head_dim_scales) # Q fp4 + scale + + kv_tokens * (head_dim_packed + head_dim_scales) # KV fp4 + scale + + total_q * heads * 2 # weights bf16 + + bs * max_blocks_per_seq * 4 # block_tables i32 + + pairs * 4 # out fp32 (written window) + ) + sec = us * 1e-6 + tflops = flops / sec / 1e12 + gbps = bytes_total / sec / 1e9 + _VARQLEN_PERF_SUMMARY.append((total_q, heads, pairs, kv_tokens, us, tflops, gbps)) + + +def _print_varqlen_perf_summary(): + print("\n" + "=" * 80) + print("Perf summary (flydsl varqlen FP4; kernel-only, useful FLOPs/bytes)") + print("=" * 80) + print( + f" {'total_q':>7} | {'heads':>5} | {'pairs':>8} | {'kv_tok':>7} | " + f"{'us':>9} | {'TFLOPS':>7} | {'GB/s':>7}" + ) + print(" " + "-" * 68) + for total_q, heads, pairs, kv_tokens, us, tflops, gbps in _VARQLEN_PERF_SUMMARY: + print( + f" {total_q:>7} | {heads:>5} | {pairs:>8} | {kv_tokens:>7} | " + f"{us:>9.2f} | {tflops:>7.2f} | {gbps:>7.1f}" + ) + print() + + def main(): ap = argparse.ArgumentParser() ap.add_argument("--bench", action="store_true") @@ -499,6 +739,31 @@ def main(): iters=args.iters, warmup=args.warmup, ) + + # ── Variable-qlen (per-batch MTP via cu_seq_q) sweep + TFLOPS/bandwidth ── + print("=" * 80) + print("[test] FP4 paged variable-qlen (per-batch MTP) MQA logits") + print("=" * 80) + varqlen_cfgs = [ + ([3, 1, 2], [512, 320, 768], 64, 0), + ([10, 1, 5, 2], [1024, 256, 2048, 512], 64, 1), # incl a qlen=10 batch + ([2, 0, 3], [384, 256, 640], 64, 2), # empty batch (qlen=0) + ([4, 3], [512, 768], 128, 3), + ([16, 8, 24, 4], [4096, 2048, 4096, 1024], 64, 4), # larger (repr. perf) + ] + for qlens, ctxs, vheads, seed in varqlen_cfgs: + run_varqlen_case( + qlens, + ctxs, + heads=vheads, + seed=seed, + bench=args.bench, + iters=args.iters, + warmup=args.warmup, + ) + if _VARQLEN_PERF_SUMMARY: + _print_varqlen_perf_summary() + print(" PASS")