Fix paged-attention KV cache dtype + size accounting (issue #119)

[LxYuan0420]

This PR is:

• To align the Metal paged-attention KV cache dtype with the model's dtype (fixes batched decode parity for Metal paged-attention parity mismatch vs standard path #119).
• To compute KV cache byte sizes via torch.dtype.itemsize instead of allocating temporary tensors.

Notes:

• tests/test_metal_kernel_paged.py::test_batched_decode_matches now passes.
• tests/test_metal_kernel_paged.py::test_greedy_output_matches remains xfailed (tracked in Metal paged-attention parity mismatch vs standard path #119). This is a remaining single-request greedy parity mismatch between the paged-kernel path and the standard path; fixing it likely requires deeper kernel/offset semantics work, so I'm keeping it out of this PR to keep scope tight.

Quick manual smoke test:

Terminal 1:

vllm serve Qwen/Qwen3-0.6B --host 127.0.0.1 --port 8000 --max-model-len 2048

Terminal 2 (single request):

curl -fsS http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"model":"Qwen/Qwen3-0.6B","messages":[{"role":"user","content":"Write a 2-sentence apple story."}],"max_tokens":512,"temperature":0.8}' \
| jq -r '.choices[0].message.content'

Terminal 2 (concurrent 4 requests):

for i in 1 2 3 4; do
  (
    echo "===== req $i ====="
    curl -fsS http://127.0.0.1:8000/v1/chat/completions \
      -H 'Content-Type: application/json' \
      -d "{\"model\":\"Qwen/Qwen3-0.6B\",\"messages\":[{\"role\":\"user\",\"content\":\"Write a 2-sentence apple story (${i}).\"}],\"max_tokens\":256,\"temperature\":0.8}" \
    | jq -r '.choices[0].message.content'
    echo
  ) &
done
wait

Related: #119

[WindChimeRan]

@LxYuan0420
Thanks for tracking this down and fixing it!

The dtype fix looks good to me. The root cause is clear (hardcoded float16 when the model's actual dtype is bfloat16), The kernel already supports all three float types natively, so this was purely a Python-side plumbing issue.

One suggestion: we now have a diagnostic tool in #127 (tools/avg_gen_length.py) that runs offline inference on ShareGPT prompts and reports response length statistics (mean/std). It would be great to run it before and after this fix to quantify the improvement, specifically comparing the paged path (VLLM_METAL_USE_PAGED_ATTENTION=1) against the non-paged baseline. If the distributions align more closely after the fix, that's a strong quantitative signal beyond the existing test assertions.

[LxYuan0420]

Good point. I ran a small end-to-end sanity check with tools/avg_gen_length.py (ShareGPT, fixed seed) and it completes without crashing on my side. For meaningful numbers, could you paste the exact command + summary table you’re using so we can standardize on that as the reference?
@WindChimeRan

[WindChimeRan]

Findings

• --max-num-seqs appears to not take effect on the mlx_lm path. In my runs, the script seems to ignore this flag and proceeds with its own batching behavior.
• The response-length distributions for main-branch mlx_lm and paged KV are currently very similar under this setup.
• With the current experimental setting, the results are inconclusive as a signal for quality/regression detection.

This may indicate we need a more discriminative setup (e.g., a different dataset, longer decoding lengths, or settings that amplify divergence if it exists).

As a simple short-term check, we could also manually inspect a small sample of responses side by side. If the response quality improves after the fix, or if the paged and non-paged paths look qualitatively similar, that would still be useful supporting evidence.

e.g.,

what's the capital of France?
write a quicksort in python
...

Experiment

Without this patch, on the main branch:

# mlx_lm path
python tools/avg_gen_length.py

============================================================
  max_num_seqs      N   Mean tokens        Std
------------------------------------------------------------
             1    100         244.8       33.4
             8    100         243.6       36.5
============================================================

# paged kv  path
# no batch size 1, taking too long
VLLM_METAL_USE_PAGED_ATTENTION=1 VLLM_METAL_MEMORY_FRACTION=0.4 python tools/avg_gen_length.py --max-num-seqs 8

============================================================
  max_num_seqs      N   Mean tokens        Std
------------------------------------------------------------
             8    100         243.5       36.6
============================================================

[LxYuan0420]

@WindChimeRan Was the "before" run using the same settings like VLLM_METAL_MEMORY_FRACTION=0.4 ?

[LxYuan0420]

I think the current result is good enough to show stability under pressure. This PR is a bug fix, not a performance improvement: it fixes the block-exhaustion failure (RuntimeError: Not enough free blocks) by aligning allocation/preemption behavior with scheduler-driven recompute

@ericcurtin PTAL

[ericcurtin]

LGTM, the dtype fix is clean and the root cause (hardcoded float16 vs actual model dtype) is clear. The KvCacheDtypeInference abstraction is well-scoped and the fallback behavior is sensible. Merging.