fix for small cache-max-entry-count (InternLM#3221)

* fix for small cache-max-entry-count

* fix kernel

Author: grimoire

lmdeploy/pytorch/engine/engine.py

@@ -554,7 +554,7 @@ def _make_infer_outputs(self, next_token_ids: torch.LongTensor, running: SeqList
                 outputs[session_id].logits = logits.split(seq_length)[idx]
         return outputs
 
-    def _make_forward_inputs(self, prefill: bool = None):
+    def _make_forward_inputs(self, prefill: bool = None, enable_empty: bool = False):
         """make forward inputs."""
         prefill_interval = self.scheduler_config.prefill_interval
 
@@ -609,6 +609,10 @@ def __need_logits(seqs: SeqList):
         if prefill is None:
             prefill = self._do_prefill()
         scheduler_output = self.scheduler.schedule(is_prefill=prefill, prealloc_size=prefill_interval)
+
+        if enable_empty and len(scheduler_output.running) == 0:
+            return None
+
         # schedule decoding if no valid prefill reqs.
         if prefill and len(scheduler_output.running) == 0:
             prefill = False
@@ -709,9 +713,13 @@ async def _async_loop_main(self, resp_que: asyncio.Queue, has_runable_event: asy
         forward_inputs = None
         next_running = None
 
-        async def _send_next_inputs(prefill: bool = None):
+        async def _send_next_inputs(prefill: bool = None, enable_empty: bool = False):
             nonlocal forward_inputs, next_running
-            forward_inputs = self._make_forward_inputs(prefill)
+            forward_inputs = self._make_forward_inputs(prefill, enable_empty)
+            if forward_inputs is None:
+                forward_inputs = None
+                next_running = None
+                return
             next_running = forward_inputs.pop('running')
             await self.executor.forward_async(forward_inputs)
 
@@ -730,7 +738,7 @@ async def _prefetch_next_inputs():
 
             if enable:
                 # send next forward
-                await _send_next_inputs(prefill)
+                await _send_next_inputs(prefill, True)
 
         while True:
             if next_running is None:

lmdeploy/pytorch/kernels/cuda/blocked_fp8_fused_moe.py

@@ -116,15 +116,15 @@ def fused_moe_blocked_f8_kernel(
     as_ptrs = A_scale + offs_am * stride_asm
     bs_ptrs = B_scale + stride_bse * exp_id + offs_bsn * stride_bsn
 
-    acc_scale = tl.load(as_ptrs) * tl.load(bs_ptrs)
+    acc_scale = tl.load(as_ptrs, mask=mask_sid, other=1.0) * tl.load(bs_ptrs)
     acc_ratio = 1 / acc_scale
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         # load scales
         k_start = (k + 1) * BLOCK_SIZE_K
         offs_ksa = k_start // group_ak
         offs_ksb = k_start // group_bk
-        a_scale = tl.load(as_ptrs + offs_ksa * stride_ask, mask=k_start < K, other=1.0)
+        a_scale = tl.load(as_ptrs + offs_ksa * stride_ask, mask=mask_sid and k_start < K, other=1.0)
         b_scale = tl.load(bs_ptrs + offs_ksb * stride_bsk, mask=k_start < K, other=1.0)
 
         # load ab