biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix indexer documentation and safety issues

Browse Source 
1. score_topk.py: Fix docstring — K^IComp[s] is shared (MQA), not per-head K^IComp[s,h]
   Matches the .cu kernel and production Indexer.forward() einsum.
2. score_topk.py: Add WARNING about valid_lens broadcast being wrong for batched prefill
3. csa_indexer.py: Replace random weights with RuntimeError — CSAIndexer has no
   checkpoint loading. Production uses the Indexer class in single_shot_inference.py.
4. csa_indexer.py: Document RoPE assumption — indexer queries/keys have no RoPE.
   NEEDS VERIFICATION against HF reference.
This commit is contained in:
biondizzle
2026-06-02 19:08:40 +00:00
parent d770111cb1
commit b111525af4
2 changed files with 36 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
dsv4/kernels/indexer/csa_indexer.py
View File

@@ -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

12
dsv4/kernels/indexer/score_topk.py
View File

@@ -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()