From d75cd5d76fab898474d75fcede6ea8e04ce5578b Mon Sep 17 00:00:00 2001
From: Patrick Kappeler <58693036+citizenu03bb@users.noreply.github.com>
Date: Mon, 29 Jun 2026 16:34:26 +0200
Subject: [PATCH] fix(inference): clamp dense walk top-k to feature count

---
 crates/larql-inference/src/vindex/walk_ffn/mod.rs | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/crates/larql-inference/src/vindex/walk_ffn/mod.rs b/crates/larql-inference/src/vindex/walk_ffn/mod.rs
index cb5f99eee..193864a43 100644
--- a/crates/larql-inference/src/vindex/walk_ffn/mod.rs
+++ b/crates/larql-inference/src/vindex/walk_ffn/mod.rs
@@ -212,7 +212,9 @@ fn walk_trace_env_enabled() -> bool {
 
 impl<'a> WalkFfn<'a> {
     fn top_k_for(&self, layer: usize) -> usize {
-        self.config.k_for(layer).unwrap_or(usize::MAX)
+        self.config
+            .k_for(layer)
+            .unwrap_or_else(|| self.index.num_features(layer))
     }
 
     // ── Legacy constructors (stable public API) ──
@@ -786,6 +788,15 @@ mod dispatch_tests {
         assert_eq!(out.shape(), &[1, weights.hidden_size]);
     }
 
+    #[test]
+    fn dense_top_k_for_clamps_to_layer_feature_count() {
+        let weights = shared_weights();
+        let idx = mock_index(weights);
+        let ffn = WalkFfn::new_unlimited(weights, &idx);
+
+        assert_eq!(ffn.top_k_for(0), weights.intermediate_size);
+    }
+
     /// Variant of `MockGateIndex` that yields a `FeatureMeta` for every
     /// `(layer, feature)` query. This is what `take_trace` needs to
     /// promote a residual into a populated `WalkHit` — without it, the