diff --git a/dsv4/kernels/indexer/csa_indexer.py b/dsv4/kernels/indexer/csa_indexer.py index 6b5cf021..b736b663 100644 --- a/dsv4/kernels/indexer/csa_indexer.py +++ b/dsv4/kernels/indexer/csa_indexer.py @@ -4,9 +4,19 @@ Paper §2.3.1, eq. 13–17: c_Q = h_t · W_DQ (shared with main queries) q^I_t = c_Q · W_IUQ (low-rank indexer queries) w^I_t = h_t · W_w (per-head weights) - I[t,s] = Σ_h w^I_t,h · ReLU(q^I_t,h · K^IComp[s]) + I[t,s] = Σ_h w^I_t,h · ReLU(q^I_t,h · K^IComp[s]) (MQA: shared key K) Selected = TopK(I[t,:]) +Key layout: K^IComp[s] is shared across indexer heads (MQA, NOT per-head). +The dot product is: q^I_t,h (per-head) · K^IComp[s] (shared). +This matches the production Indexer.forward() einsum 'tnd,cd->tnc'. + +RoPE: Neither indexer queries nor keys have RoPE applied. +The indexer is a lightweight scoring mechanism for block selection, +not a full attention layer. If the HF reference applies RoPE to +indexer keys, the stored FP4 keys would need it baked in at +compression time. VERIFY THIS AGAINST THE REFERENCE BEFORE PRODUCTION. + The indexer only exists in CSA layers. HCA and SWA layers don't have an indexer (they do dense attention). """ @@ -47,14 +57,22 @@ class CSAIndexer: # For now, use a simple torch linear; will swap to Nvfp4Linear # with FP4 output in Phase 2. if not hasattr(self, '_q_up_weight'): - # Lazy init — weights would be loaded from checkpoint - d_c = self.config.query_compression_dim - n_ih = self.config.indexer_num_heads - c_i = self.config.indexer_head_dim - self._q_up_weight = torch.randn( - d_c, n_ih * c_i, dtype=torch.bfloat16, device='cuda') * 0.02 - self._w_head_weight = torch.randn( - self.config.hidden_size, n_ih, dtype=torch.bfloat16, device='cuda') * 0.02 + # WARNING: USING RANDOM WEIGHTS — csa_indexer.py has NO weight loading. + # The production path uses the Indexer class in single_shot_inference.py + # which loads real weights from the checkpoint via Nvfp4Linear. + # This CSAIndexer class should NOT be used for production inference. + # If you see this message, you need to wire up checkpoint weight loading + # or use the production Indexer instead. + raise RuntimeError( + "CSAIndexer has no checkpoint weight loading. " + "Use the production Indexer class (single_shot_inference.py) instead, " + "or implement weight loading for CSAIndexer.") + # Old code (random weights — removed to prevent silent incorrect behavior): + # d_c = self.config.query_compression_dim + # n_ih = self.config.indexer_num_heads + # c_i = self.config.indexer_head_dim + # self._q_up_weight = torch.randn(d_c, n_ih * c_i, ...) * 0.02 + # self._w_head_weight = torch.randn(hidden_size, n_ih, ...) * 0.02 q_I = torch.nn.functional.linear(c_Q, self._q_up_weight.T) # [T, n_ih * c_i] BF16 w_h = torch.nn.functional.linear(h_t, self._w_head_weight.T).float() # [T, n_ih] FP32 diff --git a/dsv4/kernels/indexer/score_topk.py b/dsv4/kernels/indexer/score_topk.py index 9b0a8227..517d926b 100644 --- a/dsv4/kernels/indexer/score_topk.py +++ b/dsv4/kernels/indexer/score_topk.py @@ -39,10 +39,14 @@ def run_indexer_score_topk( ) -> torch.Tensor: """Returns [T, top_k] int32 of selected compressed entry indices. - The kernel computes: - I[t,s] = Σ_h w_h[t,h] · ReLU(q_I[t,h] · K^IComp[s,h]) + The kernel computes (MQA — shared key across indexer heads): + I[t,s] = Σ_h w_h[t,h] · ReLU(q_I[t,h] · K^IComp[s]) topk_indices = argtopk(I[t,:], k=top_k) + Note: K^IComp[s] is shared across heads (MQA), NOT per-head K^IComp[s,h]. + This matches the .cu kernel and the production Indexer.forward() einsum. + The paper (eq. 16) uses the shared-key form. + q_I is passed as BF16 and dequantized to FP32 before the kernel. The indexer keys are stored FP4 in the cache and dequantized inside the kernel. @@ -61,7 +65,9 @@ def run_indexer_score_topk( # Simplification: assume T == B for now (one token per request in decode). if valid_lens.shape[0] != T: # Prefill: T > B. We need to map tokens to requests. - # For now, broadcast the first request's valid_lens. + # WARNING: broadcasting request 0's valid_lens is WRONG for batched + # or multi-request prefill — it selects from wrong key ranges per token. + # This is only correct for single-request bring-up. # TODO: proper per-token valid_lens from request_ids mapping. valid_lens = valid_lens[:1].expand(T).contiguous()