{"record":{"id":"7af10e11063ec2a9","repo":"sgl-project/sglang","slug":"varlen-kda-requires-batch-size-1","errorCode":null,"errorMessage":"varlen KDA requires batch size 1","messagePattern":"varlen KDA requires batch size 1","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_prefill.py","lineNumber":327,"sourceCode":"    flat_a_log = (\n        a_log.reshape(-1)\n        if a_log is not None\n        else torch.empty(1, device=g.device, dtype=torch.float32)\n    )\n    flat_bias = (\n        dt_bias.reshape(-1)\n        if dt_bias is not None\n        else torch.empty(1, device=g.device, dtype=torch.float32)\n    )\n    activate = a_log is not None\n    has_bias = dt_bias is not None\n    use_lower_bound = lower_bound is not None\n    lower_bound_value = 0.0 if lower_bound is None else lower_bound\n\n    is_varlen = cu_seqlens is not None\n    if is_varlen:\n        if g.size(0) != 1:\n            raise ValueError(\"varlen KDA requires batch size 1\")\n        if chunk_indices is None:\n            chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE)\n        metadata = cu_seqlens\n        gate_kernel = _gate_cumsum_operands_varlen\n    else:\n        metadata = torch.empty(0, device=g.device, dtype=torch.int32)\n        chunk_indices = torch.empty(0, 2, device=g.device, dtype=torch.long)\n        gate_kernel = _gate_cumsum_operands\n\n    return gate_kernel(\n        g,\n        q,\n        k,\n        beta,\n        flat_a_log,\n        flat_bias,\n        metadata,\n        chunk_indices,","sourceCodeStart":309,"sourceCodeEnd":345,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_prefill.py#L309-L345","documentation":"In varlen mode (cu_seqlens provided), the Helion KDA prefill chunked gate/cumsum path only supports a single logical batch: g.size(0) must be 1. gate_chunk_cumsum_operands raises this when it sees a varlen request whose gate tensor still carries a batch dimension greater than 1.","triggerScenarios":"Calling chunk_kda with cu_seqlens set (varlen batch) but g shaped [B>1, T, ...] — e.g. feeding a padded batch tensor together with cumulative-sequence-length metadata.","commonSituations":"Mixing padded-batch tensors with varlen metadata after a scheduler change; tests that pass a batch of 2 with cu_seqlens; leftover batching from a prefill aggregation path that should have flattened to [1, total_tokens].","solutions":["Flatten the batch into [1, total_tokens] and pass matching cu_seqlens covering all packed sequences","If you truly need batch > 1, don't use varlen — drop cu_seqlens and pad instead","Check the caller (e.g. attention backend prefill) flattens requests before invoking chunk_kda"],"exampleFix":"// before\ng = g  # [2, T, HV, K] with cu_seqlens\nchunk_kda(..., g=g, cu_seqlens=cu)\n// after\ng = g.reshape(1, -1, *g.shape[2:])  # [1, 2T, HV, K]\nchunk_kda(..., g=g, cu_seqlens=cu)  # cu covers both seqs","handlingStrategy":"validation","validationCode":"if cu_seqlens is not None:\n    assert g.size(0) == 1, \"flatten varlen inputs to [1, total_tokens]\"\n    g = g.reshape(1, -1, *g.shape[2:])","typeGuard":"def valid_varlen_gates(g: torch.Tensor, cu: torch.Tensor | None) -> bool:\n    return cu is None or g.size(0) == 1","tryCatchPattern":null,"preventionTips":["Flatten all packed prefill tensors to batch dim 1 together with cu_seqlens","Keep one code path that prepares varlen inputs, used by all kernels"],"tags":["kda","helion","varlen","batch-shape"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}