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.

Activity

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

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions