fix a bug in native_backend

Kipsora · Kipsora · commit 0725846d8fe1 · 2025-10-13T22:33:46.000-04:00
diff --git a/python/sgl_jax/srt/layers/attention/native_backend.py b/python/sgl_jax/srt/layers/attention/native_backend.py
@@ -133,7 +133,7 @@ def _get_and_update_kv_cache(
             v = v.reshape(num_tokens, kv_heads, head_dim)
 
         if forward_batch.forward_mode == ForwardMode.EXTEND:
-            forward_batch.token_to_kv_pool.set_kv_buffer(
+            token_to_kv_pool.set_kv_buffer(
                 layer_id, forward_batch.out_cache_loc, k, v, is_decode=False
             )
         else:

Original file line number	Diff line number	Diff line change
`@@ -133,7 +133,7 @@ def _get_and_update_kv_cache(`
`133`	`133`	`v = v.reshape(num_tokens, kv_heads, head_dim)`
`134`	`134`
`135`	`135`	`if forward_batch.forward_mode == ForwardMode.EXTEND:`
`136`		`- forward_batch.token_to_kv_pool.set_kv_buffer(`
	`136`	`+ token_to_kv_pool.set_kv_buffer(`
`137`	`137`	`layer_id, forward_batch.out_cache_loc, k, v, is_decode=False`
`138`	`138`	`)`
`139`	`139`	`else:`