Skip to content

GPU Lloyd-Max TqAttention writes compressed-region V aggregate in the rotated basis (no inverse WHT) #435

Description

@pekkah

GPU Lloyd-Max TqAttention writes compressed-region V aggregate in the rotated basis (no inverse WHT)

Summary

Found while investigating #432. Both GPU Lloyd-Max TurboQuant attention kernels aggregate the compressed-region V contribution in the rotated/sign-flipped basis and never apply the inverse transform:

  • CUDA llm_tq_attention (src/SharpInference.Cuda/CudaTextKernels.cs, Phase 3): acc += weight * cbook[idx] * norm per output dim, then out[out_off + d] = acc — the TQ V codes are stored rotated (llm_tq_kv_append applies WHT + sign flip to both K and V), so this sum is a rotated-domain vector. It is added together with the unrotated FP32-window V contribution and written straight to the attention output, which feeds the Wo matmul.
  • Vulkan TqAttention shader (src/SharpInference.Vulkan/Shaders.cs): identical structure, identical omission.

CudaTurboQuantTests.TqAttention_NeedleInHaystack documents this as a convention — "The output is in the rotated/sign-flipped basis (the V-side of TqAttention does not un-rotate — same convention as the Vulkan TqAttention shader)" — and compares the kernel output against the rotated needle value, so the test suite enshrines the wrong basis rather than catching it.

Why this is a bug

The rotation R = D·H is orthonormal, so K-scores computed in the rotated domain are fine (llm_tq_rotate_query rotates q the same way). But the V aggregate must be brought back: the correct output is R⁻¹·Σ w·norm·centroid + Σ w·v_fp32. Omitting R⁻¹ means the compressed-region part of the attention output is scrambled by a fixed rotation before hitting the output projection. This corrupts generation for any model once context exceeds the FP32 window (256), on CUDA and Vulkan alike, independently of the quantizer-distortion issue in #432.

The correct pattern already exists in-tree: the CPU path (TurboQuantKvCache.ComputeVAggregation Phase 3) accumulates in the rotated domain and applies the deferred per-head sign-flip + inverse WHT once; the KVarN CUDA kernel (llm_kvarn_attention Phase 3) does ONE inverse WHT per head in shared memory and then adds the FP32-window contribution in the original domain.

Fix sketch

  1. In both kernels, keep the TQ V aggregate in a per-head shared-memory accumulator, apply sign-flip + inverse WHT (butterfly over head_dim, mirroring llm_kvarn_attention), then add the FP32-window aggregate in the original domain.
  2. Rewrite TqAttention_NeedleInHaystack (and its Vulkan twin) to compare against the needle value in the original basis.
  3. Regenerate the precompiled SPIR-V table (scripts/gen-spirv.ps1) for the Vulkan shader change.
  4. Re-validate with the perplexity harness on -g -1 (and a Vulkan run): Lloyd-Max GPU numbers should move from garbage toward the CPU Lloyd-Max numbers reported in TurboQuant Lloyd-Max 3-bit collapses on Qwen3 (PPL 61x, degenerate decode) — pre-existing on master #432.

Severity / exposure

After #432's fix the default --tq resolves to KVarN wherever supported, so this path is now opt-in (--tq-mode lloydmax) or a fallback (Vulkan backend, partial CUDA offload, MoE-on-GPU, SnapKV) — but those fallback paths currently produce corrupted attention output at depth, which is worse than the quantizer-distortion collapse described in #432.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions