[Bug] Qwen3-235B开启AMXInt8时，kt服务在特定用户输入下必定崩溃

### 检查清单

- [x] 1. 我已经搜索过相关问题，但未能获得预期的帮助
- [x] 2. 该问题在最新版本中尚未修复
- [x] 3. 请注意，如果您提交的BUG相关 issue 缺少对应环境信息和最小可复现示例，我们将难以复现和定位问题，降低获得反馈的可能性
- [x] 4. 如果您提出的不是bug而是问题，请在讨论区发起讨论 https://github.com/kvcache-ai/ktransformers/discussions。否则该 issue 将被关闭
- [x] 5. 为方便社区交流，我将使用中文/英文或附上中文/英文翻译（如使用其他语言）。未附带翻译的非中文/英语内容可能会被关闭

### 问题描述

当用户通过api与开启AMXInt8的Qwen3-235B对话时，某些输入会导致kt服务报错并无法继续进行对话。
以下为一个案例：用户输入为 “redis怎样作为缓存，在基于fastapi的python程序中使用？为我提供在多工作进程间保持对同一缓存的读取和更新的示例”
kt服务报错：“RuntimeError: CUDA error: an illegal memory access was encountered
6月 09 14:00:55 root bash[2572385]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.”
我做了以下验证：

- [ ] 重复输入原问题，每次均能复现相同报错；
- [x] 输入“redis怎样作为缓存，在基于fastapi的python程序中使用？”，可以正常回答；
- [x] 输入“为我提供在多工作进程间保持对同一缓存的读取和更新的示例”，可以正常回答；
- [x] 输入“在基于fastapi的python程序中使用？为我提供在多工作进程间”，可以正常回答；

即可以确认只在该问题完整输入时复现该现象。

日志如下：
```
6月 09 13:36:52 root bash[2572286]: INFO: Started server process [2572286]
6月 09 13:36:52 root bash[2572286]: INFO: Waiting for application startup.
6月 09 13:36:52 root bash[2572286]: INFO: Application startup complete.
6月 09 13:36:52 root bash[2572286]: INFO: Uvicorn running on http://0.0.0.0:6399 (Press CTRL+C to quit)
6月 09 13:58:48 root bash[2572286]: sched_rpc started with PID: 2574607
6月 09 13:58:48 root bash[2572286]: Queue Proxy Started
6月 09 13:58:48 root bash[2572286]: INFO: 127.0.0.1:54258 - "POST /api/embed HTTP/1.1" 404 Not Found
6月 09 13:59:33 root bash[2572286]: INFO: 192.168.5.5:36306 - "GET /v1/models HTTP/1.1" 200 OK
6月 09 13:59:53 root bash[2572286]: INFO: 127.0.0.1:49092 - "POST /v1/api/embed HTTP/1.1" 404 Not Found
6月 09 14:00:54 root bash[2572286]: INFO: 192.168.5.5:50648 - "GET /v1/models HTTP/1.1" 200 OK
6月 09 14:00:54 root bash[2572286]: /root/miniconda3/envs/kt/lib/python3.11/site-packages/pydantic/main.py:519: UserWarning: Pydantic serializer warnings:
6月 09 14:00:54 root bash[2572286]: PydanticSerializationUnexpectedValue(Expected `list[dict[str, any]]` - serialized value may not be as expected [input_value={}, input_type=dict])
6月 09 14:00:54 root bash[2572286]: return self.__pydantic_serializer__.to_json(
6月 09 14:00:54 root bash[2572286]: INFO: 192.168.5.5:50660 - "POST /v1/chat/completions HTTP/1.1" 200 OK
6月 09 14:00:54 root bash[2572286]: 2025-06-09 14:00:54,232 DEBUG /root/miniconda3/envs/kt/lib/python3.11/site-packages/ktransformers/server/backend/interfaces/balance_serve.py[418]: get input ids of shape torch.Size([1, 46])
6月 09 14:00:55 root bash[2572385]: 2025-06-09 14:00:55,439 - INFO - flashinfer.jit: Loading JIT ops: sampling
6月 09 14:00:55 root bash[2572385]: /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/utils/cpp_extension.py:2356: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
6月 09 14:00:55 root bash[2572385]: If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].
6月 09 14:00:55 root bash[2572385]: warnings.warn(
6月 09 14:00:55 root bash[2572385]: /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/utils/cpp_extension.py:2356: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
6月 09 14:00:55 root bash[2572385]: If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].
6月 09 14:00:55 root bash[2572385]: warnings.warn(
6月 09 14:00:55 root bash[2572385]: cuda_graph: 1/6, warmup finished.
6月 09 14:00:55 root bash[2572385]: capturing cuda graph 2 2
6月 09 14:00:55 root bash[2572385]: cuda_graph: 2/6, warmup finished.
6月 09 14:00:55 root bash[2572385]: capturing cuda graph 3 3
6月 09 14:00:55 root bash[2572385]: cuda_graph: 3/6, warmup finished.
6月 09 14:00:55 root bash[2572385]: capturing cuda graph 4 64
6月 09 14:00:55 root bash[2572385]: cuda_graph: 4/6, warmup finished.
6月 09 14:00:55 root bash[2572385]: capturing cuda graph 4 256
6月 09 14:00:55 root bash[2572385]: cuda_graph: 5/6, warmup finished.
6月 09 14:00:55 root bash[2572385]: capturing cuda graph 4 512
6月 09 14:00:55 root bash[2572385]: cuda_graph: 6/6, warmup finished.
6月 09 14:00:55 root bash[2572385]: add query id: 1, batch.query_lengths: 46, batch_query_tokens: torch.Size([32768]), batch.block_indexes: tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
6月 09 14:00:55 root bash[2572385]: 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27,
6月 09 14:00:55 root bash[2572385]: 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41,
6月 09 14:00:55 root bash[2572385]: 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55,
6月 09 14:00:55 root bash[2572385]: 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69,
6月 09 14:00:55 root bash[2572385]: 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83,
6月 09 14:00:55 root bash[2572385]: 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97,
6月 09 14:00:55 root bash[2572385]: 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
6月 09 14:00:55 root bash[2572385]: 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
6月 09 14:00:55 root bash[2572385]: 126, 127], dtype=torch.int32)
6月 09 14:00:55 root bash[2572385]: prefill_batch_i: 46,
6月 09 14:00:55 root bash[2572385]: Model execution time (GPU): 1175.037 ms, 0.851 tokens/s
6月 09 14:00:55 root bash[2572385]: 2025-06-09 14:00:55,457 - INFO - flashinfer.jit: Finished loading JIT ops: sampling
6月 09 14:00:55 root bash[2572385]: Process SpawnProcess-1:
6月 09 14:00:55 root bash[2572385]: Traceback (most recent call last):
6月 09 14:00:55 root bash[2572385]: File "/root/miniconda3/envs/kt/lib/python3.11/multiprocessing/process.py", line 314, in _bootstrap
6月 09 14:00:55 root bash[2572385]: self.run()
6月 09 14:00:55 root bash[2572385]: File "/root/miniconda3/envs/kt/lib/python3.11/multiprocessing/process.py", line 108, in run
6月 09 14:00:55 root bash[2572385]: self._target(*self._args, **self._kwargs)
6月 09 14:00:55 root bash[2572385]: File "/root/miniconda3/envs/kt/lib/python3.11/site-packages/ktransformers/server/backend/interfaces/balance_serve.py", line 282, in run_engine
6月 09 14:00:55 root bash[2572385]: engine.loop()
6月 09 14:00:55 root bash[2572385]: File "/root/miniconda3/envs/kt/lib/python3.11/site-packages/ktransformers/server/backend/interfaces/balance_serve.py", line 260, in loop
6月 09 14:00:55 root bash[2572385]: generated_tokens, probs = self.sampling( self.model_runner.output)
6月 09 14:00:55 root bash[2572385]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
6月 09 14:00:55 root bash[2572385]: File "/root/miniconda3/envs/kt/lib/python3.11/site-packages/ktransformers/server/backend/interfaces/balance_serve.py", line 224, in sampling
6月 09 14:00:55 root bash[2572385]: generated_tokens, probs=self.sampler(logit, sample_options)
6月 09 14:00:55 root bash[2572385]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
6月 09 14:00:55 root bash[2572385]: File "/root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
6月 09 14:00:55 root bash[2572385]: return self._call_impl(*args, **kwargs)
6月 09 14:00:55 root bash[2572385]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
6月 09 14:00:55 root bash[2572385]: File "/root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
6月 09 14:00:55 root bash[2572385]: return forward_call(*args, **kwargs)
6月 09 14:00:55 root bash[2572385]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
6月 09 14:00:55 root bash[2572385]: File "/root/miniconda3/envs/kt/lib/python3.11/site-packages/ktransformers/server/balance_serve/inference/sampling/sampler.py", line 97, in forward
6月 09 14:00:55 root bash[2572385]: temperature_0_idx = torch.where(sampling_config.temperatures == 0)[0]
6月 09 14:00:55 root bash[2572385]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
6月 09 14:00:55 root bash[2572385]: RuntimeError: CUDA error: an illegal memory access was encountered
6月 09 14:00:55 root bash[2572385]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
6月 09 14:00:55 root bash[2572385]: For debugging consider passing CUDA_LAUNCH_BLOCKING=1
6月 09 14:00:55 root bash[2572385]: Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
6月 09 14:00:55 root bash[2572385]: terminate called after throwing an instance of 'c10::Error'
6月 09 14:00:55 root bash[2572385]: what(): CUDA error: an illegal memory access was encountered
6月 09 14:00:55 root bash[2572385]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
6月 09 14:00:55 root bash[2572385]: For debugging consider passing CUDA_LAUNCH_BLOCKING=1
6月 09 14:00:55 root bash[2572385]: Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
6月 09 14:00:55 root bash[2572385]: Exception raised from c10_cuda_check_implementation at /pytorch/c10/cuda/CUDAException.cpp:43 (most recent call first):
6月 09 14:00:55 root bash[2572385]: frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string, std::allocator >) + 0x98 (0x754cb9d785e8 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libc10.so)
6月 09 14:00:55 root bash[2572385]: frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__cxx11::basic_string, std::allocator > const&) + 0xe0 (0x754cb9d0d4a2 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libc10.so)
6月 09 14:00:55 root bash[2572385]: frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x3c2 (0x754cba1b8422 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libc10_cuda.so)
6月 09 14:00:55 root bash[2572385]: frame #3: + 0x32953 (0x754cba194953 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libc10_cuda.so)
6月 09 14:00:55 root bash[2572385]: frame #4: + 0x32b01 (0x754cba194b01 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libc10_cuda.so)
6月 09 14:00:55 root bash[2572385]: frame #5: + 0x66c7f2 (0x754cb186c7f2 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libtorch_python.so)
6月 09 14:00:55 root bash[2572385]: frame #6: + 0x457608 (0x754cb1657608 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libtorch_python.so)
6月 09 14:00:55 root bash[2572385]: frame #7: + 0x7b385 (0x754cb9d59385 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libc10.so)
6月 09 14:00:55 root bash[2572385]: frame #8: c10::TensorImpl::~TensorImpl() + 0x9 (0x754cb9d52f39 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libc10.so)
6月 09 14:00:55 root bash[2572385]: frame #9: + 0x7182b8 (0x754cb19182b8 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libtorch_python.so)
6月 09 14:00:55 root bash[2572385]: frame #10: + 0x7186d1 (0x754cb19186d1 in /root/miniconda3/envs/kt/lib/python3.11/site-packages/torch/lib/libtorch_python.so)
6月 09 14:00:55 root bash[2572385]: frame #11: /root/miniconda3/envs/kt/bin/python() [0x4f5eb7]
6月 09 14:00:55 root bash[2572385]: frame #12: /root/miniconda3/envs/kt/bin/python() [0x5031bf]
6月 09 14:00:55 root bash[2572385]: frame #13: /root/miniconda3/envs/kt/bin/python() [0x55f246]
6月 09 14:00:55 root bash[2572385]: frame #14: /root/miniconda3/envs/kt/bin/python() [0x50302a]
6月 09 14:00:55 root bash[2572385]: frame #15: /root/miniconda3/envs/kt/bin/python() [0x5031bf]
6月 09 14:00:55 root bash[2572385]: frame #16: /root/miniconda3/envs/kt/bin/python() [0x55f246]
6月 09 14:00:55 root bash[2572385]: frame #17: /root/miniconda3/envs/kt/bin/python() [0x55ef87]
6月 09 14:00:55 root bash[2572385]: frame #18: /root/miniconda3/envs/kt/bin/python() [0x55f040]
6月 09 14:00:55 root bash[2572385]: frame #19: /root/miniconda3/envs/kt/bin/python() [0x541355]
6月 09 14:00:55 root bash[2572385]: frame #20: /root/miniconda3/envs/kt/bin/python() [0x53fd88]
6月 09 14:00:55 root bash[2572385]: frame #21: /root/miniconda3/envs/kt/bin/python() [0x53fdc4]
6月 09 14:00:55 root bash[2572385]: frame #22: /root/miniconda3/envs/kt/bin/python() [0x53fdc4]
6月 09 14:00:55 root bash[2572385]: frame #23: /root/miniconda3/envs/kt/bin/python() [0x54e833]
6月 09 14:00:55 root bash[2572385]: frame #24: _PyEval_EvalFrameDefault + 0x46f9 (0x515ac9 in /root/miniconda3/envs/kt/bin/python)
6月 09 14:00:55 root bash[2572385]: frame #25: /root/miniconda3/envs/kt/bin/python() [0x5cc3aa]
6月 09 14:00:55 root bash[2572385]: frame #26: PyEval_EvalCode + 0x9f (0x5cba7f in /root/miniconda3/envs/kt/bin/python)
6月 09 14:00:55 root bash[2572385]: frame #27: /root/miniconda3/envs/kt/bin/python() [0x5ecba7]
6月 09 14:00:55 root bash[2572385]: frame #28: /root/miniconda3/envs/kt/bin/python() [0x5e8740]
6月 09 14:00:55 root bash[2572385]: frame #29: PyRun_StringFlags + 0x5f (0x5db24f in /root/miniconda3/envs/kt/bin/python)
6月 09 14:00:55 root bash[2572385]: frame #30: PyRun_SimpleStringFlags + 0x3b (0x5daffb in /root/miniconda3/envs/kt/bin/python)
6月 09 14:00:55 root bash[2572385]: frame #31: Py_RunMain + 0x388 (0x5f7498 in /root/miniconda3/envs/kt/bin/python)
6月 09 14:00:55 root bash[2572385]: frame #32: Py_BytesMain + 0x39 (0x5bc149 in /root/miniconda3/envs/kt/bin/python)
6月 09 14:00:55 root bash[2572385]: frame #33: + 0x29d90 (0x754cbca29d90 in /lib/x86_64-linux-gnu/libc.so.6)
6月 09 14:00:55 root bash[2572385]: frame #34: __libc_start_main + 0x80 (0x754cbca29e40 in /lib/x86_64-linux-gnu/libc.so.6)
6月 09 14:00:55 root bash[2572385]: frame #35: /root/miniconda3/envs/kt/bin/python() [0x5bbf93]
```

### 复现步骤

启动命令：`python ktransformers/server/main.py --port 6399 --architectures Qwen3MoeForCausalLM --model_path /work/ktransformers/models/Q3Config/Qwen3-235B-A22B/ --gguf_path /work/ktransformers/models/Qwen3/Qwen3-235B-A22B-BF16 --model_name qwen3-235b-a22b --optimize_config_path ktransformers/optimize/optimize_rules/Qwen3Moe-serve-amx.yaml --max_new_tokens 32768 --cache_lens 32768 --chunk_size 512 --max_batch_size 1 --cpu_infer 46 --backend_type balance_serve`
修改配置文件内容：
```
diff --git a/ktransformers/optimize/optimize_rules/Qwen3Moe-serve-amx.yaml b/ktransformers/optimize/optimize_rules/Qwen3Moe-serve-amx.yaml
index 8607ca0..be39d59 100644
--- a/ktransformers/optimize/optimize_rules/Qwen3Moe-serve-amx.yaml
+++ b/ktransformers/optimize/optimize_rules/Qwen3Moe-serve-amx.yaml
@@ -56,7 +56,7 @@
generate_device: "cpu"
generate_op: "KExpertsCPU"
out_device: "cuda"
- backend: "AMXBF16" # or "AMXBF16" or "llamafile" (default)
+ backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
recursive: False # don't recursively inject submodules of this module
- match:
name: "^model\\.layers\\..*\\.self_attn$"
@@ -93,4 +93,4 @@
class: ktransformers.operators.mlp.KQwen2MoeMLP
kwargs:
generate_device: "cuda"
- prefill_device: "cuda"
\ No newline at end of file
+ prefill_device: "cuda"
```

### 环境信息

系统版本：Ubuntu 22.04.5 LTS
CPU：单路 Intel(R) Xeon(R) Platinum 8461V
GPU：3090 24G
ktransformer版本：0.3.1 commit 70719703396d4efabeb047f0f480436131f05918
```
commit 70719703396d4efabeb047f0f480436131f05918 (HEAD -> main, origin/main, origin/HEAD)
Merge: cc367f5 00949d5
Author: Atream <80757050+Atream@users.noreply.github.com>
Date: Mon Jun 2 16:07:04 2025 +0800

Merge pull request #1343 from zhouye/main

Mirror #1247 in server mode
```