WIP

huggingface · Jan 17, 2025 · 6570149 · 6570149
1 parent 653118d
commit 6570149
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/...ration-inference/server/text_generation_server/jetstream_pt_support/models/qwen2_model.py b/...ration-inference/server/text_generation_server/jetstream_pt_support/models/qwen2_model.py
@@ -426,6 +426,13 @@ def forward(
     ragged_batch_index: precomputed batch index for ragged attention
     ragged_block_index: precomputed block index for ragged attention
     """
+    # jax.debug.print("tokens", tokens)
+    # jax.debug.print("input_pos", input_pos)
+    # jax.debug.print("mask", mask)
+    # jax.debug.print("start", start)
+    # jax.debug.print("ragged_batch_index", ragged_batch_index)
+    # jax.debug.print("ragged_block_index", ragged_block_index)
+
     with jax.named_scope("transformer_tok"):
       seqlen = tokens.shape[-1]
       h = self.tok_embeddings(tokens)