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