minor fix

Author: katjasrz

examples/jax_tvm_ffi/gemma3_flashinfer_jax.ipynb

@@ -1011,10 +1011,10 @@
     "# 1. Compile kernels (once)\n",
     "gelu_module          = gen_act_and_mul_module('gelu_tanh').build_and_load()\n",
     "rope_module          = gen_rope_module().build_and_load()\n",
+    "decode_local_module  = gen_single_decode_module(..., use_sliding_window=True, ...).build_and_load()\n",
+    "decode_global_module = gen_single_decode_module(..., use_sliding_window=False, ...).build_and_load()\n",
     "prefill_local        = gen_single_prefill_module('fa2', ..., use_sliding_window=True).build_and_load()\n",
     "prefill_global       = gen_single_prefill_module('fa2', ..., use_sliding_window=False).build_and_load()\n",
-    "decode_local_module  = gen_decode_jit_spec('bfloat16', 256, use_sliding_window=True).build_and_load()\n",
-    "decode_global_module = gen_decode_jit_spec('bfloat16', 256, use_sliding_window=False).build_and_load()\n",
     "\n",
     "# 2. Prefill: FlashInfer causal attention over all prompt tokens → KV-cache\n",
     "h_last, kv_caches = prefill(prompt_ids)\n",