[Fix] ChatModule incorrect temperature buffer shape (mlc-ai#1070)

PR mlc-ai#1048 updated the signature of softmax in the built model library and changed the temperature buffer shape in ChatModule. This causes some existing demo unable to run since we did not do a round of model library update. This PR reverts the ChatModule change, and adds back the softmax function in non-batching case. With this PR, the regression should be fixed.
octoml · masahi · Oct 30, 2023 · Oct 8, 2023 · Oct 8, 2023 · Oct 8, 2023
commit 204860b786460247f77717053cd3a4e1e0069b06
diff --git a/cpp/llm_chat.cc b/cpp/llm_chat.cc
@@ -481,7 +481,7 @@ class LLMChat {
     // Step 6. KV cache creation.
     this->kv_cache_ = ft_.create_kv_cache_func_();
     // Step 7. Pre-allocate fixed size ndarray
-    this->temperature_arr_ = NDArray::Empty({1}, DataType::Float(32), device_);
+    this->temperature_arr_ = NDArray::Empty({}, DataType::Float(32), device_);
     float temperature = static_cast<float>(this->temperature_);
     this->temperature_arr_.CopyFromBytes(&temperature, sizeof(float));
     if (ft_.use_disco) {

diff --git a/mlc_llm/relax_model/llama.py b/mlc_llm/relax_model/llama.py
@@ -1068,7 +1068,20 @@ def create_paged_kv_cache_func(bb: relax.BlockBuilder, config: LlamaConfig) -> N
         bb.emit_func_output(cache)
 
 
-def create_softmax_func(bb: relax.BlockBuilder, config: LlamaConfig) -> None:
+def create_softmax_func_for_single_seq(bb: relax.BlockBuilder, config: LlamaConfig) -> None:
+    with bb.function("softmax_with_temperature"):
+        logits = nn.Placeholder(
+            (1, 1, tvm.tir.Var("vocab_size", "int64")), dtype="float32", name="logits"
+        )
+        temperature = nn.Placeholder((), dtype="float32", name="temperature")
+        with bb.dataflow():
+            div = bb.emit(relax.op.divide(logits, temperature))
+            softmax = bb.emit(relax.op.nn.softmax(div, axis=-1))
+            gv = bb.emit_output(softmax)
+        bb.emit_func_output(gv, [logits, temperature])
+
+
+def create_softmax_func_for_batching(bb: relax.BlockBuilder, config: LlamaConfig) -> None:
     with bb.function("softmax_with_temperature"):
         bsz = tvm.tir.Var("nseq", "int64")
         logits = nn.Placeholder(
@@ -1206,12 +1219,13 @@ def get_model(args, hf_config):
         create_prefill_func_for_batching(bb, param_manager, config, args.quantization)
         create_decoding_func_for_batching(bb, param_manager, config, args.quantization)
         create_paged_kv_cache_func(bb, config)
+        create_softmax_func_for_batching(bb, config)
     else:
         create_prefill_func_for_single_seq(bb, param_manager, config, args.quantization, sep_embed)
         create_decoding_func_for_single_seq(bb, param_manager, config, args.quantization)
         create_kv_cache_func(bb, config)
+        create_softmax_func_for_single_seq(bb, config)
 
-    create_softmax_func(bb, config)
     create_metadata_func(
         bb,
         model_name=model_name,