feat(pymllm): support Qwen3 Jetson BF16, W4A16, and W8A8 serving

docs/index.rst

@@ -357,6 +357,11 @@ Documents
 
    cache/index
 
+.. toctree::
+   :maxdepth: 2
+
+   pymllm_runtime/index
+
 .. toctree::
    :maxdepth: 2
 

docs/pymllm_runtime/developer_guide.rst

@@ -0,0 +1,220 @@
+pymllm Developer Guide
+======================
+
+总览
+----------------------------------------
+
+本文档面向希望为 ``pymllm`` 增加模型、量化格式、kernel 或性能优化的开发者。当前代码处在
+快速演进阶段，推荐遵循“小步验证、边界清晰、先单测后服务级验证”的工作方式。
+
+开发环境建议
+----------------------------------------
+
+推荐使用 editable install，便于修改 Python 代码后直接验证：
+
+.. code-block:: bash
+
+   cd <repo-root>
+   SKBUILD_WHEEL_CMAKE=false python3 -m pip install -e .
+   python3 -m pip install -e <repo-root>/mllm-kernel --no-deps --no-build-isolation
+
+最小检查：
+
+.. code-block:: bash
+
+   python3 - <<'PY'
+   import pymllm
+   import mllm_kernel
+   print("ok")
+   PY
+
+``mllm-kernel`` 的 JIT 编译产物会写入 ``~/.cache/mllm_kernel``。正常修改后重新运行
+会触发相应 kernel 的加载或编译；只有在验证首次编译行为、排查失败缓存、或更换 CUTLASS
+等外部头文件来源时，才需要清理对应缓存：
+
+.. code-block:: bash
+
+   rm -rf ~/.cache/mllm_kernel/<kernel-cache-name>
+
+新增模型
+----------------------------------------
+
+新增模型时，优先复用现有 ``pymllm.layers`` 和 ``pymllm.executor`` 约定，而不是把
+HuggingFace 模型直接包进服务。
+
+推荐步骤：
+
+1. 新增 ``pymllm/models/<model_name>.py``。
+2. 在 ``pymllm/models/__init__.py`` 注册 architecture 字符串。
+3. 实现模型类，保持 ``forward(input_ids, positions, forward_batch)`` 风格。
+4. 所有 linear layer 都接受 ``quant_method``。
+5. 实现 ``load_weights``，处理 checkpoint key、stacked projection 和 tied embedding。
+6. 增加最小单测。
+7. 最后做服务级 smoke test。
+
+最小测试建议：
+
+.. code-block:: bash
+
+   pytest pymllm/tests/test_<model>_model_registry.py -q
+   pytest pymllm/tests/test_<model>_weight_loading.py -q
+   pytest pymllm/tests/test_<model>_forward_timing.py -q
+
+新增量化 scheme
+----------------------------------------
+
+新增量化路径时，不建议在模型文件里写格式判断。推荐保持以下分层：
+
+.. code-block:: text
+
+   QuantizationConfig
+       parses checkpoint config
+       decides whether a layer is quantized
+
+   LinearMethod
+       owns linear layer lifecycle
+
+   Scheme
+       owns checkpoint-facing params
+       owns post-load layout conversion
+       owns kernel apply path
+
+``create_weights`` 应注册 checkpoint-facing 参数名。``process_weights_after_loading`` 应作为
+checkpoint layout 到 runtime kernel layout 的唯一转换边界。``apply`` 中只做 forward 必需的
+runtime 计算，不应重复做权重 repack。
+
+新增量化路径至少需要覆盖：
+
+- config 解析测试。
+- ``ignore`` / prefix 匹配测试。
+- 参数注册 shape/dtype 测试。
+- post-load layout 转换测试。
+- forward correctness 或 smoke test。
+
+新增 CUDA JIT kernel
+----------------------------------------
+
+若 kernel 适合走 ``mllm-kernel`` 的 TVM-FFI JIT 路径，推荐结构如下：
+
+.. code-block:: text
+
+   mllm-kernel/mllm_kernel/cuda/csrc/<area>/<kernel>.cuh
+   mllm-kernel/mllm_kernel/cuda/jit/<kernel>.py
+   mllm-kernel/tests/test_<kernel>.py
+   mllm-kernel/benchmarks/bench_<kernel>.py
+
+Python wrapper 应负责：
+
+- 校验输入 shape、dtype、device。
+- 分配输出 tensor。
+- 调用 ``@jit`` 包装后的 compiled module。
+- 暴露稳定、简洁的 Python API。
+
+CUDA/C++ source 应尽量只表达 kernel 语义，不混入 checkpoint 配置解析或模型层逻辑。
+
+如果 kernel 依赖 CUTLASS 等重模板库，可以先做编译 spike。确认 Jetson 目标设备上的编译时间、
+缓存路径、include 来源和内存占用后，再决定使用 TVM-FFI JIT、torch extension JIT 或 AOT 构建。
+
+服务级验证
+----------------------------------------
+
+服务级 smoke test 应覆盖：
+
+- ``/v1/models`` 可返回。
+- 文本 ``/v1/chat/completions`` 可完成。
+- 图文模型能处理容器内图片绝对路径。
+- streaming 与 non-streaming 至少各测一次。
+- 中止请求或客户端断连不会泄漏 running request。
+
+示例：
+
+.. code-block:: bash
+
+   curl -s --noproxy '*' http://127.0.0.1:30000/v1/models ; echo
+
+   curl -s --noproxy '*' http://127.0.0.1:30000/v1/chat/completions \
+     -H "Content-Type: application/json" \
+     -d '{
+       "model": "default",
+       "messages": [{"role": "user", "content": "只回复 ok"}],
+       "max_tokens": 8,
+       "temperature": 0.0,
+       "stream": false
+     }' ; echo
+
+性能验证
+----------------------------------------
+
+性能数据需要固定口径，否则不同记录之间很难比较。建议记录：
+
+- commit hash。
+- JetPack / L4T 版本。
+- GPU 型号和 compute capability。
+- PyTorch、Triton、FlashInfer、CUDA 版本。
+- 模型路径和量化格式。
+- 启动命令。
+- prompt token 数、max tokens、temperature。
+- 是否启用 radix cache、CUDA Graph、shared queue。
+- 是否包含首次 JIT 编译。
+
+对服务级请求，建议丢弃第一次 warmup 结果，记录第 2/3 次请求的 prefill/decode 统计。
+对 kernel microbench，建议单独记录 warmup、重复次数、输入 shape 和 dtype。
+
+常见问题定位
+----------------------------------------
+
+启动失败
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+优先确认：
+
+- ``pymllm`` 和 ``mllm_kernel`` 是否来自预期源码目录或安装版本。
+- ``model_path`` 和 ``tokenizer_path`` 是否在容器内可见。
+- ``transformers`` 是否能读取目标 ``config.json``。
+- CUDA 是否可用，``torch.cuda.get_device_capability()`` 是否符合量化 kernel 要求。
+
+W8A8 编译失败
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+优先确认：
+
+- ``CUTLASS_HOME`` 是否设置正确。
+- ``flashinfer`` 是否包含 bundled CUTLASS。
+- ``~/.cache/mllm_kernel/cutlass_int8_scaled_mm/`` 是否存在旧的失败缓存。
+- 当前 GPU 是否为 SM80-SM89。
+
+请求卡住或 CPU 占用高
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+优先确认：
+
+- scheduler 是否启用了 idle sleep。
+- tokenizer / scheduler / detokenizer 子进程是否全部存活。
+- 是否有请求已经断连但未 abort。
+- ``max_total_tokens`` 是否过小导致 KV allocation 反复失败和 eviction。
+
+输出异常
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+优先确认：
+
+- tokenizer chat template 是否符合目标模型。
+- EOS token 是否从 config、generation_config 或 tokenizer 中正确解析。
+- 量化模型的 ``ignore`` 是否覆盖视觉分支、embedding、norm 和 lm_head 等不应量化模块。
+- ``process_weights_after_loading`` 是否已执行。
+
+贡献建议
+----------------------------------------
+
+开发时尽量保持以下边界：
+
+- 服务协议变化放在 ``pymllm/server``。
+- 请求/响应结构放在 ``pymllm/engine/io_struct.py``。
+- 调度策略放在 ``pymllm/orchestrator/scheduler_process.py``。
+- GPU 资源和 forward 逻辑放在 ``pymllm/executor``。
+- 模型结构放在 ``pymllm/models``。
+- 基础层放在 ``pymllm/layers``。
+- 量化格式放在 ``pymllm/quantization``。
+- 自定义 kernel 放在 ``mllm-kernel``。
+
+这样可以避免把一次模型适配写成跨层补丁，也方便后续把同一能力复用到更多模型和设备。

docs/pymllm_runtime/index.rst

@@ -0,0 +1,12 @@
+pymllm Runtime
+==============
+
+.. toctree::
+   :maxdepth: 2
+
+   setup_and_usage
+   runtime_design
+   models_and_quantization
+   kernels_and_acceleration
+   developer_guide
+