[FLYDSL] Support paged mqa logits fp4 varqlen kernel

aiter/ops/flydsl/__init__.py

@@ -46,6 +46,7 @@
     )
     from .kernels.pa_mqa_logits_fp4_prefill import (
         flydsl_pa_mqa_logits_fp4_prefill,
+        flydsl_pa_mqa_logits_fp4_varqlen,
     )
     from .kernels.fp8_mqa_logits import (
         flydsl_fp8_mqa_logits,
@@ -64,6 +65,7 @@
         "flydsl_qk_norm_rope_quant",
         "flydsl_pa_mqa_logits_fp4",
         "flydsl_pa_mqa_logits_fp4_prefill",
+        "flydsl_pa_mqa_logits_fp4_varqlen",
         "flydsl_fp8_mqa_logits",
         "FP8_MQA_LOGITS_VARIANTS",
         "FP8_MQA_LOGITS_DEFAULT_VARIANT",

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]

aiter/ops/flydsl/kernels/pa_mqa_logits_fp4_prefill.py

@@ -756,3 +756,70 @@ def flydsl_pa_mqa_logits_fp4_prefill(
         stream,
     )
     return out
+
+
+def _build_varqlen_mtp_windows(cu_seq_q, context_lens, total_q):
+    """Build ragged-row metadata for per-batch variable query length (MTP)."""
+    dev = cu_seq_q.device
+    cu = cu_seq_q.to(torch.int32)
+    B = context_lens.shape[0]
+    r = torch.arange(total_q, device=dev, dtype=torch.int32)
+    # row -> batch = number of batch-ends (cu[1:]) that are <= r.
+    row_to_batch = torch.searchsorted(cu[1:].contiguous(), r, right=True).to(
+        torch.int32
+    )
+    row_to_batch = row_to_batch.clamp_(max=max(B - 1, 0))
+    n = r - cu[:-1][row_to_batch]  # in-batch position
+    qlen = (cu[1:] - cu[:-1])[row_to_batch]
+    ctx = context_lens.to(torch.int32)[row_to_batch]
+    local_starts = torch.zeros(total_q, device=dev, dtype=torch.int32)
+    local_ends = (ctx - qlen + n + 1).clamp_(min=0).to(torch.int32)
+    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,
+    cu_seq_q: torch.Tensor,
+    context_lens: torch.Tensor,
+    max_seq_len: int,
+    *,
+    weight_scale: float = 1.0,
+    block_k: int = 256,
+    kv_block_size: int = 64,
+    num_warps: int = DEFAULT_NUM_WARPS,
+    parallel_unit_num: Optional[int] = None,
+    out: Optional[torch.Tensor] = None,
+    stream: Optional[torch.cuda.Stream] = None,
+) -> torch.Tensor:
+    """Variable-qlen (per-batch MTP) FP4 paged MQA logits (gfx950)."""
+    total_q = q_fp4.shape[0]
+    row_to_batch, local_starts, local_ends = _build_varqlen_mtp_windows(
+        cu_seq_q, context_lens, total_q
+    )
+    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,
+        stream=stream,
+    )

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",