AirLLMChatGLM is an adapter that maps AirLLM's streaming interface to ChatGLM's module structure and tensor layout, handling layer name routing, sequence positioning, rotary embeddings, and KV cache packing.
AirLLMChatGLM maps its layer names to ChatGLM's module structure via set_layer_names_dict, with keys: embed → transformer.embedding.word_embeddings, layer_prefix → transformer.encoder.layers, norm → transformer.encoder.final_layernorm, lm_head → transformer.output_layer, and an additional rotary_pos_emb key → transformer.rotary_pos_emb.[1]
ChatGLM uses a [seq_len, batch, heads, dim] tensor layout: AirLLMChatGLM.get_sequence_len reads sequence length from axis 0 (seq.shape[0]), and get_past_key_values_cache_seq_len likewise reads from axis 0 of the key tensor (past_key_values[0][0].shape[0]).[1]
AirLLMChatGLM.get_pos_emb_args computes rotary embeddings over the full configured seq_length, slices the result to the current sequence length, and transposes to [seq_len, 1, ...] layout before returning them under the rotary_pos_emb key.[1] AirLLMChatGLM.get_attention_mask_args always passes attention_mask=None, and get_position_ids_args returns an empty dict — ChatGLM relies on rotary embeddings and does not use explicit masks or position IDs during AirLLM layer streaming.[1]
AirLLMChatGLM.get_past_key_value_args packs the KV cache as {'kv_cache': (k_cache, v_cache)}, matching the kv_cache parameter name required by ChatGLM's layer forward signature.[1]
Sources