gistfile1.txt

1. Start the benchmark server in vscode as [this](https://gist.github.com/vanbasten23/dd4f3cbb314a7b9cf6c003103c23c019). Select the correct python intepreter.
2. Then start the vllm server in debugger.
3. After the server is up and running.
4. Add the breakpoint (remember to turn of dynamo and jax jit)
5. Use the [script](https://gist.github.com/vanbasten23/726b28f072993fb7587482672b9c96a9) to send benchmarking request. Make sure to use the correct conda/python.
6. Then dump the input and output.


=========================
pip install flatbuffers
# for old jax
import torchax
import pickle
weights, jax_m = torchax.extract_jax(self.self_attn)
with open("/home/xiowei_google_com/old_weights.pkl", "wb") as f:
    pickle.dump(weights, f)

import jax
from jax import export
import jax.numpy as jnp
exported = export.export(jax.jit(jax_m))(jax.eval_shape(lambda x: x, weights), (jax.ShapeDtypeStruct((16,), jnp.int32), jax.ShapeDtypeStruct((16, 1536), jnp.bfloat16)))
serialized = exported.serialize()
with open("/home/xiowei_google_com/old_exports.pkl", "wb") as f:
    pickle.dump(serialized, f)



# For new jax
import torchax
import pickle
weights, jax_m = torchax.extract_jax(self.self_attn)
with open("/home/xiowei_google_com/new_weights.pkl", "wb") as f:
    pickle.dump(weights, f)

import jax
from jax import export
import jax.numpy as jnp
exported = export.export(jax.jit(jax_m))(jax.eval_shape(lambda x: x, weights), (jax.ShapeDtypeStruct((16,), jnp.int32), jax.ShapeDtypeStruct((16, 1536), jnp.bfloat16)))
serialized = exported.serialize()
with open("/home/xiowei_google_com/new_exports.pkl", "wb") as f:
    pickle.dump(serialized, f)