[FlyDSL] Read (B, next_n) context lengths in the paged MQA logits kernel

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),
@@ -492,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()}")
@@ -524,12 +539,32 @@ 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 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:
+        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()
+        or context_lens.shape[0] != batch_size
+        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"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)
     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 +611,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

op_tests/test_flydsl_fp8_paged_mqa_logits.py

@@ -677,6 +677,88 @@ 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)
+    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, 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]]),
+        row_lens[:, :0],
+        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)
+        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,
+    )
+
+
 def main():
     if get_gfx() not in SUPPORTED_GFX:
         aiter.logger.warning(
@@ -800,6 +882,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(