From 299adc8124e93acd67ae8deb5be9f89a66257abc Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Mon, 12 Oct 2026 23:50:07 +0300 Subject: [PATCH 1/6] [FlyDSL] Read (B, next_n) context lengths in the paged MQA logits kernel --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 26 +++++++- op_tests/test_flydsl_fp8_paged_mqa_logits.py | 64 +++++++++++++++++++ 2 files changed, 87 insertions(+), 3 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 4bf0ec91081..c900570fc2f 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -205,6 +205,7 @@ def kernel( stride_q_batch: fx.Int32, max_block_len: fx.Int32, stride_out: fx.Int32, + ctx_stride: fx.Int32, ): tid = fx.Int32(fx.thread_idx.x) wave = udiv(tid, 64) @@ -228,7 +229,7 @@ def _imin(a, b): ragged_nn = _imin(fx.Int32(nn_t[pid_batch]), rows_per_batch) nn = (has_next_n_lens != 0).select(ragged_nn, rows_per_batch) - context_len = fx.Int32(context_t[pid_batch]) + context_len = fx.Int32(context_t[pid_batch * ctx_stride + ctx_stride - 1]) page_count = uceildiv(context_len, fx.Int32(KV_BLOCK_SIZE)) pages_per_split = uceildiv(page_count, split_kv) page_lo = pid_split * pages_per_split @@ -438,6 +439,7 @@ def launch( stride_q_batch, max_block_len, stride_out, + ctx_stride, stream: fx.Stream = fx.Stream(None), # noqa: B008 ): kernel._func.__name__ = kernel_name @@ -456,6 +458,7 @@ def launch( stride_q_batch, max_block_len, stride_out, + ctx_stride, ).launch( grid=(fx.Int64(grid_blocks), 1, 1), block=(THREADS, 1, 1), @@ -524,12 +527,28 @@ def flydsl_fp8_paged_mqa_logits( if out_logits.shape[0] != batch_size * next_n: raise ValueError(f"out_logits first dimension must be {batch_size * next_n}") - context_lens = context_lens.reshape(batch_size) + # The kernel reads sequence b's length at b * ctx_stride + ctx_stride - 1: + # element b of [B], or the last entry of row b of a (B, n) per-row table. + # A strided input is rejected rather than copied: a hidden copy allocates on + # every call, and the kernel would otherwise read it as dense. + if ( + context_lens.dim() not in (1, 2) + or not context_lens.is_contiguous() + or context_lens.shape[0] != batch_size + or (context_lens.dim() == 2 and context_lens.shape[1] < 1) + ): + raise ValueError( + f"context_lens must be a contiguous [B] or [B, n] with B={batch_size} " + f"and n >= 1, got shape {tuple(context_lens.shape)} with strides " + f"{context_lens.stride()}" + ) + ctx_stride = context_lens.shape[1] if context_lens.dim() == 2 else 1 + context_lens = context_lens.reshape(-1) max_block_len = kv_indices.shape[-1] kv_indices = kv_indices.reshape(batch_size, max_block_len) has_next_n_lens = next_n_lens is not None if next_n_lens is None: - next_n_lens = context_lens + next_n_lens = context_lens[:batch_size] next_n_lens = next_n_lens.reshape(batch_size) if next_n_lens.dtype != torch.int32 or next_n_lens.device != q_fp8.device: raise ValueError("next_n_lens must be int32 on the same device as q_fp8") @@ -576,6 +595,7 @@ def flydsl_fp8_paged_mqa_logits( int(q_fp8.stride(0)), int(max_block_len), int(out_logits.stride(0)), + int(ctx_stride), stream, ) return out_logits diff --git a/op_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/test_flydsl_fp8_paged_mqa_logits.py index c638d8d6f76..55b6b716d72 100644 --- a/op_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -677,6 +677,67 @@ def _summarize(name, rows): aiter.logger.info("%s summary (markdown):\n%s", name, table) +def test_fp8_paged_mqa_logits_2d_context_lens(batch, next_n, kv_len, split_kv): + """A (B, next_n) per-row table, the layout DeepGEMM takes and vLLM passes, + gives the same logits as the per-sequence lengths it ends with.""" + inp = _build_inputs( + batch, + next_n, + HEADS, + HEAD_DIM, + kv_len, + _E4M3_NATIVE, + block_size=KV_BLOCK_SIZE, + var_ratio=0.9, + ) + kv_cache_kernel, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM) + seq_lens = inp.context_lens + row_lens = ( + seq_lens[:, None] + - next_n + + 1 + + torch.arange(next_n, dtype=torch.int32, device="cuda") + ) + + def run(context_lens): + out.fill_(float("-inf")) + flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out, + context_lens, + inp.block_tables, + inp.max_model_len, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + SplitKV=_split_arg(split_kv), + ) + return out.clone() + + expected = run(seq_lens) + assert torch.equal(run(row_lens), expected) + assert torch.equal(run(seq_lens[:, None]), expected) + # Wrong row count, no columns, or a strided [B] or [B, n] are rejected. + bad_inputs = [row_lens[:-1], torch.cat([row_lens, row_lens[:1]]), row_lens[:, :0]] + if next_n > 1: + bad_inputs += [row_lens[:, -1], row_lens.t().contiguous().t()] + for bad in bad_inputs: + try: + run(bad) + except ValueError: + continue + raise AssertionError(f"context_lens of shape {tuple(bad.shape)} was accepted") + aiter.logger.info( + "fp8_paged_mqa_logits 2d context_lens ok: batch=%d next_n=%d kv_len=%d " + "split_kv=%d", + batch, + next_n, + kv_len, + split_kv, + ) + + def main(): if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning( @@ -800,6 +861,9 @@ def main(): ) args = parser.parse_args() + for next_n, split_kv in itertools.product((1, 3), (0, 3)): + test_fp8_paged_mqa_logits_2d_context_lens(4, next_n, 3000, split_kv) + compact = [ test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv, dtype) for dtype, batch, next_n, kv_len, split_kv in itertools.product( From ae01af65b76de3a34b14dbe726e832081a5dd413 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Tue, 13 Oct 2026 00:03:18 +0300 Subject: [PATCH 2/6] [FlyDSL] Reject non-int32 or off-device context_lens in paged MQA logits The kernel reads context_lens as int32 words, so an int64 table gave wrong lengths. Compare the 2D and 1D logits with checkAllclose at zero tolerance. --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 5 ++++ op_tests/test_flydsl_fp8_paged_mqa_logits.py | 26 +++++++++++++++---- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index c900570fc2f..f5318f21ccf 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -531,6 +531,11 @@ def flydsl_fp8_paged_mqa_logits( # element b of [B], or the last entry of row b of a (B, n) per-row table. # A strided input is rejected rather than copied: a hidden copy allocates on # every call, and the kernel would otherwise read it as dense. + if context_lens.dtype != torch.int32 or context_lens.device != q_fp8.device: + raise ValueError( + "context_lens must be int32 on the same device as q_fp8, got " + f"{context_lens.dtype} on {context_lens.device}" + ) if ( context_lens.dim() not in (1, 2) or not context_lens.is_contiguous() diff --git a/op_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/test_flydsl_fp8_paged_mqa_logits.py index 55b6b716d72..1134ed61ca5 100644 --- a/op_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -716,10 +716,24 @@ def run(context_lens): return out.clone() expected = run(seq_lens) - assert torch.equal(run(row_lens), expected) - assert torch.equal(run(seq_lens[:, None]), expected) - # Wrong row count, no columns, or a strided [B] or [B, n] are rejected. - bad_inputs = [row_lens[:-1], torch.cat([row_lens, row_lens[:1]]), row_lens[:, :0]] + for name, context_lens in ( + ("(B, next_n)", row_lens), + ("(B, 1)", seq_lens[:, None]), + ): + msg = f"2d context_lens {name} vs [B]: " + err = checkAllclose( + run(context_lens), expected, rtol=0, atol=0, tol_err_ratio=0, msg=msg + ) + assert err == 0, f"{msg}{err:.2%} of the logits differ" + # Wrong row count, no columns, a strided [B] or [B, n], a dtype other than + # int32, or a tensor on another device are rejected. + bad_inputs = [ + row_lens[:-1], + torch.cat([row_lens, row_lens[:1]]), + row_lens[:, :0], + row_lens.long(), + row_lens.cpu(), + ] if next_n > 1: bad_inputs += [row_lens[:, -1], row_lens.t().contiguous().t()] for bad in bad_inputs: @@ -727,7 +741,9 @@ def run(context_lens): run(bad) except ValueError: continue - raise AssertionError(f"context_lens of shape {tuple(bad.shape)} was accepted") + raise AssertionError( + f"context_lens {bad.dtype} {tuple(bad.shape)} on {bad.device} was accepted" + ) aiter.logger.info( "fp8_paged_mqa_logits 2d context_lens ok: batch=%d next_n=%d kv_len=%d " "split_kv=%d", From 178e842393246172ddbceac97f0e9ede71d1ef50 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Wed, 14 Oct 2026 20:18:11 +0300 Subject: [PATCH 3/6] [FlyDSL] Document context_lens in the wrapper; require 1 or next_n columns A (B, n) table with n other than next_n read the wrong column as the length. --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 23 ++++++++++++++----- op_tests/test_flydsl_fp8_paged_mqa_logits.py | 11 ++++++--- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index f5318f21ccf..46aa49c320a 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -495,6 +495,18 @@ def flydsl_fp8_paged_mqa_logits( ``next_n`` may be 1..8. The compiled kernel reserves LDS for eight rows; ``next_n_lens`` optionally selects fewer live rows per sequence. + + ``context_lens``: int32, contiguous, on ``q_fp8``'s device. Either + + - ``[B]``: the full length ``L_b`` of request ``b``, including its ``next_n`` + query tokens, or + - ``(B, next_n)``: row ``b`` is ``[L_b - next_n + 1, ..., L_b]``, the layout + vLLM builds. Only the last column is read: query row ``j`` + (``0 <= j < next_n``) attends to KV positions ``[0, L_b - next_n + 1 + j)``, + so other row values are not supported. + + ``(B, 1)`` is the same as ``[B]``. Anything else raises ``ValueError``; a + strided tensor is rejected rather than copied. """ if get_gfx() != _GFX950: raise RuntimeError(f"gfx950 kernel requested on {get_gfx()}") @@ -527,8 +539,7 @@ def flydsl_fp8_paged_mqa_logits( if out_logits.shape[0] != batch_size * next_n: raise ValueError(f"out_logits first dimension must be {batch_size * next_n}") - # The kernel reads sequence b's length at b * ctx_stride + ctx_stride - 1: - # element b of [B], or the last entry of row b of a (B, n) per-row table. + # The kernel reads request b's length at b * ctx_stride + ctx_stride - 1. # A strided input is rejected rather than copied: a hidden copy allocates on # every call, and the kernel would otherwise read it as dense. if context_lens.dtype != torch.int32 or context_lens.device != q_fp8.device: @@ -540,12 +551,12 @@ def flydsl_fp8_paged_mqa_logits( context_lens.dim() not in (1, 2) or not context_lens.is_contiguous() or context_lens.shape[0] != batch_size - or (context_lens.dim() == 2 and context_lens.shape[1] < 1) + or (context_lens.dim() == 2 and context_lens.shape[1] not in (1, next_n)) ): raise ValueError( - f"context_lens must be a contiguous [B] or [B, n] with B={batch_size} " - f"and n >= 1, got shape {tuple(context_lens.shape)} with strides " - f"{context_lens.stride()}" + f"context_lens must be a contiguous [B] or [B, n] with B={batch_size}, " + f"n in (1, next_n={next_n}), got shape {tuple(context_lens.shape)} " + f"with strides {context_lens.stride()}" ) ctx_stride = context_lens.shape[1] if context_lens.dim() == 2 else 1 context_lens = context_lens.reshape(-1) diff --git a/op_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/test_flydsl_fp8_paged_mqa_logits.py index 1134ed61ca5..aee7b6f80b4 100644 --- a/op_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -725,8 +725,9 @@ def run(context_lens): run(context_lens), expected, rtol=0, atol=0, tol_err_ratio=0, msg=msg ) assert err == 0, f"{msg}{err:.2%} of the logits differ" - # Wrong row count, no columns, a strided [B] or [B, n], a dtype other than - # int32, or a tensor on another device are rejected. + # Wrong row count, a column count other than 1 or next_n, a strided [B] or + # [B, n], a dtype other than int32, or a tensor on another device are + # rejected. bad_inputs = [ row_lens[:-1], torch.cat([row_lens, row_lens[:1]]), @@ -735,7 +736,11 @@ def run(context_lens): row_lens.cpu(), ] if next_n > 1: - bad_inputs += [row_lens[:, -1], row_lens.t().contiguous().t()] + bad_inputs += [ + row_lens[:, -1], + row_lens.t().contiguous().t(), + row_lens[:, 1:].contiguous(), + ] for bad in bad_inputs: try: run(bad) From 4f6a5c710f06311bb5625be6aa407e69fe581296 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Wed, 14 Oct 2026 20:45:05 +0300 Subject: [PATCH 4/6] [FlyDSL] Make the context_lens layout an axis of the paged MQA logits sweep -c/--ctx_lens (default 1d 2d) runs the 2D table through the existing sweep, graded against the fp32 reference and timed next to 1D. The rejection checks stay as a small separate test. --- op_tests/test_flydsl_fp8_paged_mqa_logits.py | 118 +++++++++---------- 1 file changed, 54 insertions(+), 64 deletions(-) diff --git a/op_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/test_flydsl_fp8_paged_mqa_logits.py index aee7b6f80b4..329bcc7493b 100644 --- a/op_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -429,13 +429,23 @@ def _time_candidates(candidates, ref, flops, nbytes, msg): return ret +def _context_lens_arg(context_lens, next_n, layout): + """``[B]`` lengths as the kernel takes them: ``1d`` as is, ``2d`` as the + ``(B, next_n)`` table vLLM builds (row ``b`` = ``[L_b - next_n + 1, ..., L_b]``).""" + if layout == "1d": + return context_lens + offsets = torch.arange(next_n, dtype=torch.int32, device=context_lens.device) + return context_lens[:, None] - next_n + 1 + offsets + + @benchmark() -def test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv, dtype): +def test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv, dtype, ctx_lens): q_dtype = _Q_DTYPE[dtype] inp = _build_inputs( batch, next_n, HEADS, HEAD_DIM, kv_len, q_dtype, block_size=KV_BLOCK_SIZE ) kv_cache_kernel, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM) + fly_context_lens = _context_lens_arg(inp.context_lens, next_n, ctx_lens) ref = run_torch( inp.q, inp.kv_cache_fp8, @@ -455,7 +465,7 @@ def flydsl(): kv_cache_kernel, inp.weights, out, - inp.context_lens, + fly_context_lens, inp.block_tables, inp.max_model_len, Preshuffle=True, @@ -465,8 +475,9 @@ def flydsl(): candidates = {"flydsl": (flydsl, out)} # Gluon is a real kernel under test, not the torch reference. It only - # accepts a launch-wide next_n = Q.shape[1], so it is compact-only. - if deepgemm_fp8_paged_mqa_logits is not None: + # accepts a launch-wide next_n = Q.shape[1], so it is compact-only, and it + # runs on 1D lengths only: its 2D support is a separate change. + if deepgemm_fp8_paged_mqa_logits is not None and ctx_lens == "1d": out_g = out.clone() def gluon(): @@ -677,9 +688,9 @@ def _summarize(name, rows): aiter.logger.info("%s summary (markdown):\n%s", name, table) -def test_fp8_paged_mqa_logits_2d_context_lens(batch, next_n, kv_len, split_kv): - """A (B, next_n) per-row table, the layout DeepGEMM takes and vLLM passes, - gives the same logits as the per-sequence lengths it ends with.""" +def test_fp8_paged_mqa_logits_bad_context_lens(next_n=3, kv_len=3000): + """Layouts the kernel would read wrongly raise ``ValueError``.""" + batch = 4 inp = _build_inputs( batch, next_n, @@ -691,72 +702,38 @@ def test_fp8_paged_mqa_logits_2d_context_lens(batch, next_n, kv_len, split_kv): var_ratio=0.9, ) kv_cache_kernel, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM) - seq_lens = inp.context_lens - row_lens = ( - seq_lens[:, None] - - next_n - + 1 - + torch.arange(next_n, dtype=torch.int32, device="cuda") - ) - - def run(context_lens): - out.fill_(float("-inf")) - flydsl_fp8_paged_mqa_logits( - inp.q_fp8, - kv_cache_kernel, - inp.weights, - out, - context_lens, - inp.block_tables, - inp.max_model_len, - Preshuffle=True, - KVBlockSize=KV_BLOCK_SIZE, - SplitKV=_split_arg(split_kv), - ) - return out.clone() - - expected = run(seq_lens) - for name, context_lens in ( - ("(B, next_n)", row_lens), - ("(B, 1)", seq_lens[:, None]), - ): - msg = f"2d context_lens {name} vs [B]: " - err = checkAllclose( - run(context_lens), expected, rtol=0, atol=0, tol_err_ratio=0, msg=msg - ) - assert err == 0, f"{msg}{err:.2%} of the logits differ" + row_lens = _context_lens_arg(inp.context_lens, next_n, "2d") # Wrong row count, a column count other than 1 or next_n, a strided [B] or - # [B, n], a dtype other than int32, or a tensor on another device are - # rejected. + # [B, n], a dtype other than int32, or a tensor on another device. bad_inputs = [ row_lens[:-1], torch.cat([row_lens, row_lens[:1]]), row_lens[:, :0], + row_lens[:, 1:].contiguous(), + row_lens[:, -1], + row_lens.t().contiguous().t(), row_lens.long(), row_lens.cpu(), ] - if next_n > 1: - bad_inputs += [ - row_lens[:, -1], - row_lens.t().contiguous().t(), - row_lens[:, 1:].contiguous(), - ] for bad in bad_inputs: try: - run(bad) + flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out, + bad, + inp.block_tables, + inp.max_model_len, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + ) except ValueError: continue raise AssertionError( f"context_lens {bad.dtype} {tuple(bad.shape)} on {bad.device} was accepted" ) - aiter.logger.info( - "fp8_paged_mqa_logits 2d context_lens ok: batch=%d next_n=%d kv_len=%d " - "split_kv=%d", - batch, - next_n, - kv_len, - split_kv, - ) + aiter.logger.info("fp8_paged_mqa_logits bad context_lens rejected: ok") def main(): @@ -808,6 +785,16 @@ def main(): default=[0], help="SplitKV; 0 = auto occupancy mapping.\n e.g.: --split-kv 0 1 4", ) + parser.add_argument( + "-c", + "--ctx_lens", + type=str, + nargs="*", + default=["1d", "2d"], + choices=["1d", "2d"], + help="context_lens layout: [B] (1d) or the (B, next_n) table vLLM " + "builds (2d).\n e.g.: -c 2d", + ) parser.add_argument( "--ragged-batch", type=int, @@ -882,15 +869,18 @@ def main(): ) args = parser.parse_args() - for next_n, split_kv in itertools.product((1, 3), (0, 3)): - test_fp8_paged_mqa_logits_2d_context_lens(4, next_n, 3000, split_kv) - compact = [ - test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv, dtype) - for dtype, batch, next_n, kv_len, split_kv in itertools.product( - args.dtype, args.batch, args.next_n, args.kv_len, args.split_kv + test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv, dtype, ctx_lens) + for dtype, batch, next_n, kv_len, split_kv, ctx_lens in itertools.product( + args.dtype, + args.batch, + args.next_n, + args.kv_len, + args.split_kv, + args.ctx_lens, ) ] + test_fp8_paged_mqa_logits_bad_context_lens() if compact: _summarize("fp8_paged_mqa_logits", compact) From b782862423c410b22a59eb8f96f0f4c9ad29fd80 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Wed, 14 Oct 2026 23:27:58 +0300 Subject: [PATCH 5/6] [FlyDSL] Make the context_lens stride a build-time constant Stride 1 builds the pre-2D kernel (same code and name); stride n builds a variant named ..._cs{n}. _compile() caches one build per stride. --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 32 +++++++++++++------ 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 46aa49c320a..0815add2cc4 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -179,11 +179,27 @@ def device_cu_count(device_index: int) -> int: return 256 -def _build_kernel(): +def _build_kernel(ctx_stride: int): + # ctx_stride is a build-time constant: 1 for a [B] context_lens (the + # pre-2D kernel and name), n for a (B, n) table (its own build and name). kernel_name = ( "fp8_paged_mqa_logits_gfx950_H32_D128_bkv64_kvb64_" "nn8_w8_lds2_mfma16_ntg2l_flydsl" ) + if ctx_stride != 1: + kernel_name += f"_cs{ctx_stride}" + + # Chosen at build time, outside the kernel body: request b's length is + # element b of a [B], or the last entry of row b of a (B, ctx_stride) table. + if ctx_stride == 1: + + def _context_len_index(pid_batch): + return pid_batch + + else: + + def _context_len_index(pid_batch): + return pid_batch * ctx_stride + (ctx_stride - 1) @fx.struct class SharedStorage: @@ -205,7 +221,6 @@ def kernel( stride_q_batch: fx.Int32, max_block_len: fx.Int32, stride_out: fx.Int32, - ctx_stride: fx.Int32, ): tid = fx.Int32(fx.thread_idx.x) wave = udiv(tid, 64) @@ -229,7 +244,7 @@ def _imin(a, b): ragged_nn = _imin(fx.Int32(nn_t[pid_batch]), rows_per_batch) nn = (has_next_n_lens != 0).select(ragged_nn, rows_per_batch) - context_len = fx.Int32(context_t[pid_batch * ctx_stride + ctx_stride - 1]) + context_len = fx.Int32(context_t[_context_len_index(pid_batch)]) page_count = uceildiv(context_len, fx.Int32(KV_BLOCK_SIZE)) pages_per_split = uceildiv(page_count, split_kv) page_lo = pid_split * pages_per_split @@ -439,7 +454,6 @@ def launch( stride_q_batch, max_block_len, stride_out, - ctx_stride, stream: fx.Stream = fx.Stream(None), # noqa: B008 ): kernel._func.__name__ = kernel_name @@ -458,7 +472,6 @@ def launch( stride_q_batch, max_block_len, stride_out, - ctx_stride, ).launch( grid=(fx.Int64(grid_blocks), 1, 1), block=(THREADS, 1, 1), @@ -470,9 +483,9 @@ def launch( return launch -@lru_cache(maxsize=1) -def _compile(): - return _build_kernel() +@lru_cache(maxsize=MAX_NN) +def _compile(ctx_stride: int = 1): + return _build_kernel(ctx_stride) def flydsl_fp8_paged_mqa_logits( @@ -590,7 +603,7 @@ def flydsl_fp8_paged_mqa_logits( split_kv = max(1, int(SplitKV)) grid_blocks = batch_size * split_kv - launcher = _compile() + launcher = _compile(int(ctx_stride)) launcher.compile_hints = dict(DEFAULT_COMPILE_HINTS) stream = stream or torch.cuda.current_stream(q_fp8.device) with torch.cuda.device(q_fp8.device.index): @@ -611,7 +624,6 @@ def flydsl_fp8_paged_mqa_logits( int(q_fp8.stride(0)), int(max_block_len), int(out_logits.stride(0)), - int(ctx_stride), stream, ) return out_logits From f476da79a1895fce933bcd9bd2a40ae85117185d Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Thu, 15 Oct 2026 23:13:06 +0300 Subject: [PATCH 6/6] [FlyDSL] Test paged MQA logits 2D context_lens with mixed lengths One case with B=5, next_n=3, distinct lengths and an empty (CUDA-graph padding) request, graded against the fp32 reference for 1D and 2D, with and without SplitKV. The 2D table clamps at 0 like vLLM. --- op_tests/test_flydsl_fp8_paged_mqa_logits.py | 57 +++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/op_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/test_flydsl_fp8_paged_mqa_logits.py index 329bcc7493b..16f4e8b6a4f 100644 --- a/op_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -435,7 +435,8 @@ def _context_lens_arg(context_lens, next_n, layout): if layout == "1d": return context_lens offsets = torch.arange(next_n, dtype=torch.int32, device=context_lens.device) - return context_lens[:, None] - next_n + 1 + offsets + # vLLM clamps at 0, so a padding request (L_b = 0) is an all-zero row. + return (context_lens[:, None] - next_n + 1 + offsets).clamp_(min=0) @benchmark() @@ -688,6 +689,59 @@ def _summarize(name, rows): aiter.logger.info("%s summary (markdown):\n%s", name, table) +def test_fp8_paged_mqa_logits_mixed_context_lens(batch=5, next_n=3, kv_len=3000): + """Distinct per-request lengths and one empty (CUDA-graph padding) request, + graded against the fp32 reference for both layouts. With uniform lengths a + kernel that read every request's length from row 0 would still pass.""" + inp = _build_inputs( + batch, + next_n, + HEADS, + HEAD_DIM, + kv_len, + _E4M3_NATIVE, + block_size=KV_BLOCK_SIZE, + var_ratio=0.9, + ) + kv_cache_kernel, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM) + seq_lens = inp.context_lens.clone() + seq_lens[1] = 0 + live = seq_lens[seq_lens > 0].tolist() + assert len(set(live)) == len(live), f"lengths are not distinct: {live}" + ref = run_torch( + inp.q, + inp.kv_cache_fp8, + inp.weights, + seq_lens, + inp.block_tables, + inp.max_model_len, + inp.fp8_dtype, + block_size=KV_BLOCK_SIZE, + ) + for layout, split_kv in itertools.product(("1d", "2d"), (0, 3)): + out.fill_(float("-inf")) + flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out, + _context_lens_arg(seq_lens, next_n, layout), + inp.block_tables, + inp.max_model_len, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + SplitKV=_split_arg(split_kv), + ) + _grade( + ref, + out, + f"mixed context_lens {layout} split_kv={split_kv} lengths={seq_lens.tolist()}", + ) + aiter.logger.info( + "fp8_paged_mqa_logits mixed context_lens ok: lengths=%s", seq_lens.tolist() + ) + + def test_fp8_paged_mqa_logits_bad_context_lens(next_n=3, kv_len=3000): """Layouts the kernel would read wrongly raise ``ValueError``.""" batch = 4 @@ -880,6 +934,7 @@ def main(): args.ctx_lens, ) ] + test_fp8_paged_mqa_logits_mixed_context_lens() test_fp8_paged_mqa_logits_bad_context_lens() if compact: _summarize("fp8_paged_mqa_logits", compact)