[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

@@ -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:
@@ -228,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])
+        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
@@ -467,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(
@@ -492,6 +508,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 +552,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")
@@ -555,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):

op_tests/test_flydsl_fp8_paged_mqa_logits.py

@@ -429,13 +429,24 @@ 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)
+    # 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()
-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 +466,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 +476,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,6 +689,107 @@ 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
+    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)
+    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.
+    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(),
+    ]
+    for bad in bad_inputs:
+        try:
+            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 bad context_lens rejected: ok")
+
+
 def main():
     if get_gfx() not in SUPPORTED_GFX:
         aiter.logger.warning(
@@ -726,6 +839,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,
@@ -801,11 +924,18 @@ def main():
     args = parser.parse_args()
 
     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_mixed_context_lens()
+    test_fp8_paged_mqa_logits_bad_context_lens()
     if compact:
         _summarize("fp8_paged_mqa_logits", compact)