--- baked-ours/models/deepseek_v4_1/nvidia/model_state.py	2026-09-10 20:13:20.000000000 +0200
+++ baked/models/deepseek_v4_1/nvidia/model_state.py	2026-09-11 10:13:47.000000000 +0200
@@ -113,6 +113,8 @@
                 input_batch.query_start_loc[: input_batch.num_reqs + 1],
                 window,
                 input_batch.num_tokens,
+                input_batch=input_batch,
+                req_states=req_states,
             )
         return model_inputs