[fix]: address the second review pass

- Allow the frozen FP32 boundary modules through the FSDP mixed-dtype
  guard when a model declares replicate_fp32_modules.
- Make the tile-64 FP4 kernel and its reference handle a listed block
  with an empty query half instead of propagating NaN; add a CPU
  regression test.
- Compare generated row ids in the row-plan test.
- Consume method checkpoint/stop requests after the pre-loop eval, so
  best/ always refers to the evaluated weights.
- Remove the unused audio-stats module and its claim.
- Refresh the V2 model description and drop the dead compression script
  reference.

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py

@@ -260,9 +260,15 @@ def fp4_vsa_attention_reference(q: torch.Tensor,
                     if position == 0 and first_block_max_floor is not None:
                         group_max = group_max.clamp(min=first_block_max_floor)
                     new_max = torch.maximum(row_max, group_max.amax(dim=-1))
-                    m_scaled = new_max * scale_log2 + offset
-                    p_tilde = torch.exp2(scores * scale_log2 - m_scaled[:, None])
-                    alpha = torch.exp2((row_max - new_max) * scale_log2)
+                    # A listed block can have no selected keys for this query
+                    # half; both maxima are -inf and the usual exp2 terms are
+                    # NaN, so emit zeros and keep the running state unchanged.
+                    empty = torch.isinf(new_max) & (new_max < 0)
+                    m_scaled = torch.where(empty, torch.zeros_like(new_max), new_max * scale_log2 + offset)
+                    p_tilde = torch.where(empty[:, None], torch.zeros_like(scores),
+                                          torch.exp2(scores * scale_log2 - m_scaled[:, None]))
+                    alpha = torch.where(empty, torch.zeros_like(new_max),
+                                        torch.exp2((row_max - new_max) * scale_log2))
                     alpha = torch.where(torch.isinf(row_max), torch.zeros_like(alpha), alpha)
                     row_sum = row_sum * alpha + p_tilde.sum(dim=-1)
                     if quantize:
@@ -279,7 +285,9 @@ def fp4_vsa_attention_reference(q: torch.Tensor,
                 out[b, h, rows] = acc * inv[:, None]
                 if weights is not None:
                     for n, p_hat, block_max in visited:
-                        decay = torch.exp2((block_max - row_max) * scale_log2) * inv
+                        finite = torch.isfinite(block_max) & torch.isfinite(row_max)
+                        decay = torch.where(finite, torch.exp2((block_max - row_max) * scale_log2) * inv,
+                                            torch.zeros_like(inv))
                         weights[b, h, rows, n * BLOCK:(n + 1) * BLOCK] = p_hat * decay[:, None]
     if weights is not None:
         return out, weights

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py

@@ -175,13 +175,20 @@ def _fwd_kernel(Q, K, V, O, OHP, LSE, MRUN, Q2K_IDX, Q2K_NUM, KV_VALID, Q2K_QUAD
         if FLOOR:
             group_max = tl.where(it == 0, tl.maximum(group_max, FLOOR_VALUE), group_max)
         m_new = tl.maximum(m_i, tl.max(group_max, axis=1))
+        # A valid mask can list a block that has no selected keys for one query
+        # half (tile 64 pairs two halves per block). Both maxima are then -inf;
+        # exp2(-inf - -inf) is NaN, so keep the running state unchanged and
+        # emit zero probabilities for that half.
+        empty = m_new == float("-inf")
         m_scaled = m_new * scale_log2 + OFFSET
-        alpha = tl.math.exp2((m_i - m_new) * scale_log2)
+        alpha = tl.where(empty, 0.0, tl.math.exp2((m_i - m_new) * scale_log2))
         if QUANT:
             p_tilde, p_hat = _quantized_p(s, m_scaled, it == 0, scale_log2, FLOOR, FLOOR_VALUE, ROWS, BLOCK_N)
         else:
             p_tilde = tl.math.exp2(s * scale_log2 - m_scaled[:, None])
             p_hat = p_tilde
+        p_tilde = tl.where(empty[:, None], 0.0, p_tilde)
+        p_hat = tl.where(empty[:, None], 0.0, p_hat)
         l_i = l_i * alpha + tl.sum(p_tilde, axis=1)
         v = tl.load(V + base + keys[:, None] * D + offs_d[None, :])
         acc = acc * alpha[:, None] + tl.dot(p_hat.to(tl.bfloat16), v)
@@ -228,7 +235,10 @@ def _bwd_dq_kernel(Q, K, V, DO, LSE, DELTA, DQ, Q2K_IDX, Q2K_NUM, KV_VALID, Q2K_
             quad = tl.load(Q2K_QUAD + meta * list_width + (it // halves)).to(tl.int32)
             bit = (quad >> ((pid % halves) * 2 + (it % halves))) & 1
             ok = ok & (bit != 0)
-        p = tl.where(ok[None, :], tl.math.exp2(s * scale_log2 - lse[:, None]), 0.0)
+        # A row with no valid keys at all has lse = -inf; clamp it so an empty
+        # half contributes zero instead of exp2(-inf - -inf) = NaN.
+        lse_safe = tl.where(lse == float("-inf"), 0.0, lse)
+        p = tl.where(ok[None, :], tl.math.exp2(s * scale_log2 - lse_safe[:, None]), 0.0)
         dp = tl.dot(do, tl.trans(v))
         ds = p * (dp - delta[:, None])
         dq += tl.dot(ds.to(tl.bfloat16), k)
@@ -274,14 +284,16 @@ def _bwd_dkdv_kernel(Q, K, V, DO, LSE, DELTA, MRUN, DK, DV, K2Q_IDX, K2Q_SLOT, K
             s = tl.where(ok[None, :] & (bit != 0), s, float("-inf"))
         else:
             s = tl.where(ok[None, :], s, float("-inf"))
-        p = tl.math.exp2(s * scale_log2 - lse[:, None])
+        lse_safe = tl.where(lse == float("-inf"), 0.0, lse)
+        p = tl.math.exp2(s * scale_log2 - lse_safe[:, None])
         if QUANT:
             qmeta = bh * n_qblocks + m_block
             first_visit = slot == tl.load(Q2K_NUM + qmeta) - 1
             m_run = tl.load(MRUN + (qmeta * list_width + slot) * BLOCK_N + row0 + offs_r)
-            m_scaled = m_run * scale_log2 + OFFSET
+            m_ok = m_run != float("-inf")
+            m_scaled = tl.where(m_ok, m_run * scale_log2 + OFFSET, 0.0)
             _, p_hat = _quantized_p(s, m_scaled, first_visit, scale_log2, FLOOR, FLOOR_VALUE, ROWS, ROWS)
-            weight = p_hat * tl.math.exp2(m_scaled - lse)[:, None]
+            weight = tl.where(m_ok[:, None], p_hat * tl.math.exp2(m_scaled - lse_safe)[:, None], 0.0)
         else:
             weight = p
         dv += tl.dot(tl.trans(weight.to(tl.bfloat16)), do)