{"record":{"id":"d8b12f4331555373","repo":"xai-org/x-algorithm","slug":"expected-k-shape-to-be-kv-shape-inferred-fro","errorCode":null,"errorMessage":"Expected {k.shape=} to be {kv_shape} (inferred from q)","messagePattern":"Expected (.+?) to be (.+?) \\(inferred from q\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention_fa3.py","lineNumber":95,"sourceCode":"\ndef _attention_forward(\n    q,\n    k,\n    v,\n    config: TuningConfig,\n    save_residuals: bool = False,\n    bound=None,\n    sm_scale: float = 1.0,\n    cap: float = -1.0,\n    cap_method: str = \"tanh\",\n):\n    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:\n        raise ValueError(f\"q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.ndim=}\")\n    batch_size, q_seq_len, num_q_heads, head_dim = q.shape\n    _, kv_seq_len, num_kv_heads, _ = k.shape\n    kv_shape = (batch_size, kv_seq_len, num_kv_heads, head_dim)\n    if k.shape != kv_shape:\n        raise ValueError(f\"Expected {k.shape=} to be {kv_shape} (inferred from q)\")\n    if v.shape != kv_shape:\n        raise ValueError(f\"Expected {v.shape=} to be {kv_shape} (inferred from q)\")\n    if (dtype := q.dtype) != k.dtype or dtype != v.dtype:\n        raise ValueError(\n            f\"q, k, and v should all have the same dtype, got: {q.dtype}, {k.dtype}, {v.dtype}\"\n        )\n    if num_q_heads % num_kv_heads:\n        raise ValueError(f\"{num_q_heads=} must be divisible by and {num_kv_heads=}\")\n    q_heads_per_kv_head = num_q_heads // num_kv_heads\n    if head_dim % 64:\n        raise ValueError(f\"{head_dim=} must be divisible by 64\")\n    if jnp.dtype(dtype) not in map(jnp.dtype, [jnp.float16, jnp.bfloat16]):\n        raise NotImplementedError(f\"Only f16 and bf16 are supported, got dtype: {dtype}\")\n\n    max_concurrent_steps = min(config.max_concurrent_steps, kv_seq_len // config.block_kv)\n    block_q, block_kv = config.block_q, config.block_kv\n    if kv_seq_len % block_kv:\n        raise ValueError(f\"{kv_seq_len=} must be a multiple of {block_kv=}\")","sourceCodeStart":77,"sourceCodeEnd":113,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention_fa3.py#L77-L113","documentation":"After unpacking q's shape, _attention_forward infers the expected kv shape (batch_size, kv_seq_len, num_kv_heads, head_dim) from q and k, and verifies k matches exactly: batch, head_dim and kv-head count must agree with q; only seq_len may differ.","triggerScenarios":"k with a different batch size, head_dim, or num_kv_heads than q (e.g. k reshaped to (B, S_kv, H_q, D) instead of grouping kv heads for GQA).","commonSituations":"GQA setups where k/v still carry the full query-head count; mixed dtypes/layouts after rearrange; off-by-one head grouping in MQA.","solutions":["Group kv heads correctly: num_kv_heads = num_q_heads / group_size","Ensure batch size and head_dim match q exactly","Verify with an assert/print of q.shape and k.shape before the call"],"exampleFix":"# before\nk = k  # (B, S_kv, num_q_heads, D) — ungrouped\n# after\nk = k[:, :, :num_kv_heads, :]  # or repeat correctly for GQA: (B, S_kv, num_kv_heads, D)","handlingStrategy":"validation","validationCode":"B, _, Hq, D = q.shape\nassert k.shape == (B, k.shape[1], k.shape[2], D) and k.shape[2] in divisors(Hq)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Log q.shape/k.shape once before training","Write a GQA shape test comparing kv head grouping against config"],"tags":["jax","attention","gqa","shape-validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}