Skip to content

Decode attention overflows shared memory on 64 KiB GPUs (Turing) for head_dim 512 (Gemma-4) #839

Closed

Description

@festeraeb

On an RTX 2060 (sm_75, 64 KiB opt-in shared memory per block), serving Gemma-4 26B-A4B fails in the first decode step:

triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 81920, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.
  File ".../freetoken/kernel/triton/attention.py", in decode_paged_attention

decode_paged_attention sizes the stage-1 tile from the K and V tiles only (2 * block_n * block_d) and never goes below _MIN_BLOCK_KV (32). Gemma-4's full-attention layers have head_dim 512, so with a 16-head Q tile the kernel needs 16 KiB (Q) + 2 x 32 KiB (K, V) = 80 KiB, more than a 64 KiB part has.

Related: #446 (Turing support) and #24 (sm_75 changes; it fixes the extend-attention tile but not this decode tile).

  • GPU: RTX 2060 12 GB (sm_75). Host: Xeon, 117 GB RAM, Linux. NVIDIA driver 580.178, CUDA 13.3
  • FreeToken: main @ a2ce636, built from source (uv pip install -e ".[accel]"), torch 2.14.1+cu130, sglang-kernel uninstalled (its wheels have no sm_75 build; fix(cuda): support Turing GPUs #24 gates it)
  • Checkpoint: google_gemma-4-26B-A4B-it-IQ4_XS.gguf, an imatrix IQ4_XS GGUF of google/gemma-4-26B-A4B-it (the quantizer's repo is not recorded in the file's metadata)
  • Command: CUDA_VISIBLE_DEVICES=<2060> ft serve --model google_gemma-4-26B-A4B-it-IQ4_XS.gguf --dtype float16 --moe-strategy offload --moe-cache-auto --memory-ratio 0.85 --gpu 0
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    legacy-gpuPre-Ampere NVIDIA GPUs (<sm80)

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions