Fix 4-bit split QKV projection

[abhinav-bellapu]

Fixes #737.

Summary

• Add a shared 4-bit Q/K/V projection helper for regular attention.
• Preserve the existing flattened 4-bit projection behavior for normal [batch, pos, d_model] inputs.
• Handle split Q/K/V inputs by selecting the matching input-head/output-head projection slice, matching complex_attn_linear.
• Add CPU-safe regression coverage with fake bitsandbytes matmul, so the split 4-bit path is tested without requiring a quantized 7B model locally.

Why

The old 4-bit path projected split inputs to [batch, pos, input_head, d_model] and then tried to reshape directly to [batch, pos, head, d_head], which reproduces the shape error in #737. Split Q/K/V inputs need head-specific projections: each input head should keep only the corresponding output-head slice.

Validation

• make format
• make check-format
• uv run --python 3.11 pytest tests/unit/components/test_attention.py -q
• uv run --python 3.11 mypy .
• make unit-test

[abhinav-bellapu]

@jlarson4 ready for review

[jlarson4]

Thanks for this fix! I verified it locally. The unit-tier, fake-bnb approach for CPU-safe coverage is a nice touch.

There are a couple organization changes outlined below that we should address before merging, but altogether the solution is sound.

Outside of the current edits in this PR, it would be appreciated if you could update tests/QUARANTINES.md.

The new test moves the two quarantined entries in tests/unit/components/test_attention.py (cited there as :48 and :83) down by ~72 lines. Please refresh those line references.

[jlarson4]

The 4D path use n_heads times the necessary compute and memory – please dequantize once and reuse complex_attn_linear instead.

_project_4bit_qkv projects every input head through all heads' weights and then keeps only the diagonal. On a 4-bit Llama-7B (32 heads) at pos=2048, that's a ~512 MiB fp16 (~1 GiB fp32) transient per projection, three times per layer, of which only the ~16 MiB diagonal survives. This is a real OOM risk on the memory-constrained setups people use 4bit models to avoid.

Since matmul_4bit dequantizes internally anyway, dequantizing the weight lets the split path share the already-tested einsum instead of maintaining a parallel implementation. Divergence between these two paths is how #737 happened in the first place.

[jlarson4]

Additionally, if you could possibly add a docstring comment to this function?

[jlarson4]

W_Q currently gets a static cast while W_K/W_V get the runtime isinstance + ValueError. A non-quantized W_Q under load_in_4bit=True dies with an opaque AttributeError.

Let's move the isinstance guard from the W_K/W_V cases into the helper and drop the Q/K/V asymmetry. Checking isinstance(weight, Params4bit) inside _project_4bit_qkv gives all three weights the clear error and collapses the three near-identical call-site blocks.