[dev] current vllm .layers_output is very likely wrong

[Butanium]

ndif-team/nnsight#642
@khaiwang here noted that vllm output is (mlp_out, residual), so we should make layer output return the sum probably? The tricky question is how to handle value editing? should we just set the mlp out in the tuple to 0 and store the edited version in residual? That feels like the safe bet, I don't think it could lead to weird behaviors to the end user.

[khaiwang]

Yes, we decide to keep it as-is to keep nnsight general enough. On nnterp, you might want to return the sum if users want the hidden_states output as in transformers. For value editing, I think setting either value to 0 and apply the other one is fine (I have tested this and they are the same, you can confirm it by testing more cases), as they will be summed up right after input to the next layer.

[Butanium]

@khaiwang does .source work with vllm? Have you checked if there is an easy .source hook point we could use for the sum? Nw if you don't know yet I can check that later this week

[khaiwang]

.source works case by case, which depends on whether vllm wraps the module into cuda kernel or raw python. Source can only access python level signatures and cannot access any wrapped kernels.

[khaiwang]

Just realized you are asking this specific case. Short answer is no. mlp output and residual in vllm are passed as two parameters to the cuda kernal fused_add_norm, and this wrapped function returns normed hidden states and residual (which is the sum of old mlp output and residual), this is the only point we can get the sum, but unfortunately at this point we cannot patch or change it. So for now the only way of having the sum and supporting patching is add them together manually, and split them back to the next layer if there's an update.