The 401(k) contribution is Federal tax free, when you make the contribution, and most likely State too. I believe that is true for California, specifically.
There was a court case some years ago about people making 401(k) or IRA contributions in New York, avoiding the New York state income tax. Then they moved to Florida (no income tax), and took the money out. New York sued, saying they had to pay the New York income tax that had been deferred, but the court said no. So you should be able to avoid California state income tax, and then later if you were to move to, for example, Texas (no income tax), have no state income tax liability.
At the Federal level, you will have different problems. You won't have the money; it will be held by the 401(k) trustee. When you try to access the money (cash the account out), you will have to pay the deferred taxes. Effectively, when you remove the money it becomes income in the year it is removed.
You can take the money out at any time, but if you are less than 59 1/2 at the time that you take it, there is a 10% penalty. The agreement is that the Feds let you defer paying the tax because it is going to finance your retirement, and they will tax it later. If you take it out before 59 1/2, they figure you are not retired yet, and are breaking your part of the agreement.
Of course you can generally leave the money in the 401(k) plan with your old employer and let it grow until you are 59.5, or roll it over into another 401(k) with a new employer (if they let you), or into an IRA. But if you have returned to your own country, having an account in the U.S. would introduce both investment risk and currency risk.
If you are in another country when you want the money, the question would be where your U.S. residence would be. If you live in California, then go to, say France, your U.S. residence would still be California, and you would still owe California income tax. If you move from California to Texas and then to France, your U.S. residence would be Texas.
This is pretty vague, as you might have heard in the Rahm Emanual case -- was he a resident of Chicago or Washington, D.C.? Same problem with Howard Hughes who was born in Texas, but then spent most his life in California, then to Nevada, then to Nicaragua, and the Bahamas. When he died Texas, California and Nevada all claimed him as a resident, for estate taxes. The important thing is to be able to make a reasonable case that you are a resident of where ever you want to be -- driver's license, mailing address, living quarters, and so on.
