JAX for inference only on 4090

[torchstation]

This keeps coming up in conversation so I figured I'd start a thread.

One thing I keep noticing is that people underestimate tail latency until they actually try to run this stuff for more than one afternoon.

Curious whether that lines up with what other people are seeing.

This is mostly for local work rather than public serving.

I expected the biggest difference to come from the model, but the stack around it mattered at least as much.

[southspark]

I think the bottleneck is somewhere else. The reported numbers can look compelling, but in ordinary usage I didn't see the same benefit.

[southden]

Interesting, I hadn't tried that. I don't think your conclusion is far off. The biggest change for me was simplifying the stack instead of adding one more optimization layer.

[plainscope50]

A couple things I'd check first. Before rebuilding anything, I'd test one clean baseline with Phi and compare from there. Once you know the plain path works, then add the faster path back in piece by piece.

I'm happy to test one more variant if you want a second data point.

[modgarden]

A couple things I'd check first. Before rebuilding anything, I'd test one clean baseline with Phi and compare from there. Once you know the plain path

What worked for me was pretty simple. On B200, the easiest win is often to back off the aggressive settings and see whether the problem survives. Once you know the plain path works, then add the faster path back in piece by piece.

[emberrunner]

Not sure I agree there. I've seen people blame INT4 when the real issue was much more mundane. Could be a workload difference, but I wouldn't generalize too fast from one setup.

[llmarch51]

Appreciate the detail. I don't think your conclusion is far off.

[oxideworks59]

Need a little more detail here. Is this local-only or are you routing anything externally?

[mangopeak72]

What worked for me was pretty simple. If you're on llama.cpp, I'd confirm the version first and then reduce variables until the behavior is obvious.

[turboscope]

I tested this tonight. On A6000 with MLX, I saw basically the same behavior. I expected the model choice to matter more than it did.