{"record":{"id":"0f95236d3f35f4b5","repo":"jax-ml/jax","slug":"soft-cap-must-not-be-0-0","errorCode":null,"errorMessage":"{soft_cap=} must not be 0.0.","messagePattern":"(.+?) must not be 0\\.0\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py","lineNumber":274,"sourceCode":"        f\"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where\"\n        \" `max_num_seqs` is `page_indices.shape[0]`.\"\n    )\n  if (\n      kv_lens.dtype != jnp.int32\n      or page_indices.dtype != jnp.int32\n      or cu_q_lens.dtype != jnp.int32\n  ):\n    raise ValueError(\n        \"The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be\"\n        f\" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},\"\n        f\" {cu_q_lens.dtype=}.\"\n    )\n  if num_q_heads % num_kv_heads != 0:\n    raise ValueError(f\"{num_q_heads=} must be divisible by {num_kv_heads=}\")\n  if sliding_window is not None and sliding_window <= 0:\n    raise ValueError(f\"{sliding_window=} must be positive.\")\n  if soft_cap is not None and soft_cap == 0.0:\n    raise ValueError(f\"{soft_cap=} must not be 0.0.\")\n  if (\n      num_kv_pages_per_block is not None\n      and not 0 < num_kv_pages_per_block <= pages_per_seq\n  ):\n    raise ValueError(\n        f\"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}].\"\n    )\n  if num_queries_per_block is not None and num_queries_per_block <= 0:\n    raise ValueError(f\"{num_queries_per_block=} must be positive.\")\n  if vmem_limit_bytes is not None and vmem_limit_bytes <= 0:\n    raise ValueError(f\"{vmem_limit_bytes=} must be positive.\")\n  del sm_scale  # No constraints on sm_scale.\n  del mask_value  # No consstraints on mask_value.\n\n\ndef ragged_paged_attention_kernel(\n    # Prefetch\n    kv_lens_ref,  # [max_num_seqs]","sourceCodeStart":256,"sourceCodeEnd":292,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py#L256-L292","documentation":"In ragged_paged_attention, soft_cap (tanh attention logit capping, as used by Gemma-2) is optional, but exactly 0.0 is forbidden because dividing by a zero cap would blow up the logits; use None to disable capping. static_validate_inputs distinguishes None from 0.0 explicitly.","triggerScenarios":"Passing soft_cap=0.0 intending to turn off logit soft-capping.","commonSituations":"Config-driven code where a missing attn_logit_softcapping value defaults to 0.0 (common in Gemma-2 / HF configs) and is forwarded verbatim to the TPU kernel.","solutions":["Convert 0.0 to None before calling: soft_cap = soft_cap or None (careful: this also maps a genuine tiny cap; check == 0.0 explicitly)","Pass the model's real cap value (e.g. 50.0 for Gemma-2) when capping is desired","Normalize model configs at load time so 'disabled' is always None"],"exampleFix":"// before\nsoft_cap = cfg.attn_logit_softcapping  # 0.0 when disabled\nattn(..., soft_cap=soft_cap)\n// after\nsoft_cap = None if not cfg.attn_logit_softcapping else cfg.attn_logit_softcapping\nattn(..., soft_cap=soft_cap)","handlingStrategy":"validation","validationCode":"soft_cap = None if soft_cap is None or soft_cap == 0.0 else float(soft_cap)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Map 0.0 soft-caps from HF-style configs to None eagerly","Keep a config-normalization step for attention hyperparameters"],"tags":["jax","pallas","tpu","softcap","attention","config-validation"],"backgroundTag":"invalid-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}