{"record":{"id":"4e69e25ede9a0cd3","repo":"sgl-project/sglang","slug":"write-pos-must-have-shape-batch","errorCode":null,"errorMessage":"`write_pos` must have shape {(batch,)}.","messagePattern":"`write_pos` must have shape (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_replayssm.py","lineNumber":726,"sourceCode":"    if dt_bias.ndim not in (1, 2) or not dt_bias.is_contiguous():\n        raise ValueError(\"KDA `dt_bias` must be a contiguous 1D or 2D tensor.\")\n    flat_a = a.view(batch, -1)\n    flat_dt_bias = dt_bias.view(-1)\n    _, num_q_heads, num_v_heads, key_dim, value_dim = validate_packed_decode_inputs(\n        mixed_qkv,\n        flat_a,\n        b,\n        A_log,\n        flat_dt_bias,\n        initial_state,\n        out,\n        ssm_state_indices,\n    )\n\n    if write_pos.ndim != 1 or write_pos.dtype is not torch.int32:\n        raise ValueError(\"`write_pos` must be a 1D int32 tensor.\")\n    if write_pos.shape != (batch,):\n        raise ValueError(f\"`write_pos` must have shape {(batch,)}.\")\n    if force_flush is not None and (\n        force_flush.ndim != 1\n        or force_flush.dtype is not torch.int32\n        or force_flush.shape != (batch,)\n    ):\n        raise ValueError(\"`force_flush` must be a length-B int32 tensor or None.\")\n\n    cache_length = d_cache.size(2)\n    if cache_length < 1:\n        raise ValueError(\"ReplaySSM cache length must be at least 1.\")\n    if d_cache.shape[1:] != (num_v_heads, cache_length, value_dim):\n        raise ValueError(\"`d_cache` must have shape [slots, HV, L, V].\")\n    if k_cache.shape[1:] != (num_q_heads, cache_length, key_dim):\n        raise ValueError(\"`k_cache` must have shape [slots, H, L, K].\")\n    if g_cache.shape[1:] != (num_v_heads, cache_length, key_dim):\n        raise ValueError(\"`g_cache` must have shape [slots, HV, L, K].\")\n    if g_cache.dtype is not torch.float32:\n        raise ValueError(\"`g_cache` must have dtype torch.float32.\")","sourceCodeStart":708,"sourceCodeEnd":744,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_replayssm.py#L708-L744","documentation":"write_pos must have exactly one entry per batch row — shape (batch,) where batch = mixed_qkv.size(0) — because the kernel indexes it per row to place the flushed output in the replay cache. A mismatched length means rows would read out-of-bounds positions.","triggerScenarios":"Passing write_pos sized for a different batch than mixed_qkv, e.g. after slicing qkv rows without slicing write_pos, or sharing a cached write_pos tensor across calls with varying batch.","commonSituations":"Continuous batching where the running batch size changes between steps but a stale write_pos is reused; tests reusing fixtures across batch sizes; prefix slicing of the batch dim on some tensors but not others.","solutions":["Regenerate write_pos per call: torch.full((batch,), pos, dtype=torch.int32, device=...)","Slice consistently: if you index mixed_qkv rows, index write_pos with the same mask","Assert write_pos.shape == (mixed_qkv.size(0),) before the call in debug builds"],"exampleFix":"// before\nwrite_pos = cached_write_pos  # built for batch=8, now batch=5\n// after\nwrite_pos = torch.full((mixed_qkv.size(0),), pos, dtype=torch.int32, device=mixed_qkv.device)","handlingStrategy":"validation","validationCode":"batch = mixed_qkv.size(0)\nassert write_pos.shape == (batch,), (write_pos.shape, batch)","typeGuard":"def write_pos_matches(qkv: torch.Tensor, wp: torch.Tensor) -> bool:\n    return wp.shape == (qkv.size(0),)","tryCatchPattern":null,"preventionTips":["Rebuild write_pos each decode step from the current batch size","Apply the same row mask/slice to all per-row tensors"],"tags":["kda","replayssm","shape-validation","batch-mismatch"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}