feat(qwen3-a8w8): wire serving backend

Author: vegetabledoww

examples/model/qwen3_14b/npu_generate.py

@@ -10,6 +10,7 @@
 from __future__ import annotations
 
 import argparse
+import dataclasses
 import statistics
 import sys
 import time
@@ -33,11 +34,17 @@ def _bootstrap_package_root() -> None:
 
 from python.core import GenerateConfig, LLMEngine, RuntimeConfig
 from python.core.kv_cache import KvCacheManager
+from python.core.model_loader import ModelLoader
 from python.core.parallel import ParallelConfig, parse_device_ids
 from python.profile import get_profiler, merge_profile, profile_span
-from examples.model.qwen3_14b.runner.npu_executor import Qwen314BPyptoExecutor as PyptoExecutor
+from examples.model.qwen3_14b.runner.a8w8_loader import Qwen3A8W8DirectoryLoader
+from examples.model.qwen3_14b.runner.npu_executor import Qwen314BPyptoExecutor
+from examples.model.qwen3_14b.runner.npu_executor_a8w8 import Qwen314BA8W8PyptoExecutor
 from python.core.types import LoadedModel
-import dataclasses
+
+
+_QWEN3_BF16_FORMAT = "qwen3-14b"
+_QWEN3_A8W8_FORMAT = "qwen3-a8w8"
 
 
 # -----------------------------------------------------------------------------
@@ -70,24 +77,6 @@ def TimePhase(self, name: str):
         finally:
             self.phases[name] = self.phases.get(name, 0.0) + (time.perf_counter() - t0)
 
-    def WrapKernel(self, fn, name: str, *, group_by_decode_step: bool = False):
-        """Return a wrapper that records every call's duration under `name`."""
-
-        def wrapper(*args, **kwargs):
-            t0 = time.perf_counter()
-            try:
-                return fn(*args, **kwargs)
-            finally:
-                dt = time.perf_counter() - t0
-                self.kernel_times[name].append(dt)
-                if group_by_decode_step and self._decode_step_idx >= 0:
-                    bucket = self.kernel_per_decode_step[name]
-                    while len(bucket) <= self._decode_step_idx:
-                        bucket.append([])
-                    bucket[self._decode_step_idx].append(dt)
-
-        return wrapper
-
     def BeginDecodeStep(self) -> None:
         self._decode_step_idx += 1
 
@@ -145,24 +134,23 @@ def InstallProfiling(engine: LLMEngine, model_id: str, collector: _TimingCollect
     """
     executor = engine._executor  # type: ignore[attr-defined]
     compiled = executor._compiled[model_id]  # type: ignore[attr-defined]
+    runner = executor._runners[model_id]  # type: ignore[attr-defined]
+    kernel_names = {
+        id(compiled.prefill): ("kernel.prefill_fwd", False),
+        id(compiled.decode): ("kernel.decode_layer", True),
+    }
 
-    # Kernels are dispatched by Qwen314BModelRunner.
-    if hasattr(compiled.prefill, "chip_callable") or hasattr(compiled.prefill, "compiled"):
-        runner = executor._runners[model_id]  # type: ignore[attr-defined]
-        orig_run_program = runner._run_distributed_program  # type: ignore[attr-defined]
-        kernel_names = {
-            id(compiled.prefill): ("kernel.prefill_fwd", False),
-            id(compiled.decode): ("kernel.decode_layer", True),
-        }
+    def install_runner_kernel_timing(method_name: str) -> None:
+        orig_run = getattr(runner, method_name)
 
-        def timed_run_program(callable_spec, *args, **kwargs):
+        def timed_run(callable_spec, *args, **kwargs):
             kernel_info = kernel_names.get(id(callable_spec))
             if kernel_info is None:
-                return orig_run_program(callable_spec, *args, **kwargs)
+                return orig_run(callable_spec, *args, **kwargs)
             name, group_by_decode_step = kernel_info
             t0 = time.perf_counter()
             try:
-                timing = orig_run_program(callable_spec, *args, **kwargs)
+                timing = orig_run(callable_spec, *args, **kwargs)
             finally:
                 dt = time.perf_counter() - t0
                 collector.kernel_times[name].append(dt)
@@ -174,14 +162,16 @@ def timed_run_program(callable_spec, *args, **kwargs):
             collector.RecordRunTiming(name, timing)
             return timing
 
-        runner._run_distributed_program = timed_run_program  # type: ignore[attr-defined]
+        setattr(runner, method_name, timed_run)
+
+    # A8W8 kernels are dispatched as L2 callables by Qwen314BA8W8ModelRunner.
+    if hasattr(compiled.prefill, "chip_callable"):
+        install_runner_kernel_timing("_run_l2_program")
+    # Original Qwen3-14B kernels are dispatched by the L3 model runner.
+    elif hasattr(compiled.prefill, "compiled"):
+        install_runner_kernel_timing("_run_distributed_program")
     else:
-        # Per-layer kernel wrappers. compiled.prefill / compiled.decode are invoked
-        # once per transformer layer inside run_prefill / run_decode respectively.
-        compiled.prefill = collector.WrapKernel(compiled.prefill, "kernel.prefill_fwd")
-        compiled.decode = collector.WrapKernel(
-            compiled.decode, "kernel.decode_layer", group_by_decode_step=True
-        )
+        raise TypeError("unsupported compiled kernel wrapper for profiling")
 
     # Top-level executor API wrappers.
     orig_prefill = executor.run_prefill
@@ -345,6 +335,15 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--model-dir", required=True, help="Local model directory, e.g. a Hugging Face snapshot.")
     parser.add_argument("--prompt", required=True, help="Prompt text.")
     parser.add_argument("--model-id", default="qwen3-14b-local")
+    parser.add_argument(
+        "--model-format",
+        default=_QWEN3_BF16_FORMAT,
+        choices=[_QWEN3_BF16_FORMAT, _QWEN3_A8W8_FORMAT],
+        help=(
+            "Qwen3-14B weight format. Use qwen3-14b for the original BF16/L3 path "
+            "or qwen3-a8w8 for compressed-tensors W8A8 checkpoints."
+        ),
+    )
     parser.add_argument("--platform", default="a2a3", choices=["a2a3sim", "a2a3", "a5sim", "a5"])
     parser.add_argument("--device-id", type=int, default=0, help="Default NPU device id when --devices is unset.")
     parser.add_argument(
@@ -367,12 +366,24 @@ def build_parser() -> argparse.ArgumentParser:
         help="Offline generation does not launch DP replicas; values > 1 fail fast.",
     )
     parser.add_argument("--max-seq-len", type=int, default=4096)
+    parser.add_argument("--max-batch-size", type=int, default=16)
     parser.add_argument("--max-new-tokens", type=int, default=32)
+    parser.add_argument(
+        "--decode-backend",
+        default="a8w8",
+        choices=["a8w8"],
+        help="For qwen3-a8w8 only: run the A8W8 prefill/decode backend.",
+    )
     parser.add_argument("--temperature", type=float, default=0.0)
     parser.add_argument("--top-p", type=float, default=1.0)
     parser.add_argument("--top-k", type=int, default=None)
     parser.add_argument("--stream", action="store_true", default=False)
     parser.add_argument("--save-kernels-dir", default=None)
+    parser.add_argument(
+        "--pto-isa-commit",
+        default=None,
+        help="For qwen3-a8w8 only: pin PyPTO compile/assemble to the installed runtime's pto-isa revision.",
+    )
     parser.add_argument(
         "--num-layers-override",
         type=int,
@@ -395,6 +406,22 @@ def build_parser() -> argparse.ArgumentParser:
     return parser
 
 
+def _model_loader_for_format(model_format: str) -> ModelLoader | None:
+    if model_format != _QWEN3_A8W8_FORMAT:
+        return None
+    model_loader = ModelLoader()
+    model_loader.register(Qwen3A8W8DirectoryLoader())
+    return model_loader
+
+
+def _executor_class_for_format(model_format: str):
+    if model_format == _QWEN3_A8W8_FORMAT:
+        return Qwen314BA8W8PyptoExecutor
+    if model_format == _QWEN3_BF16_FORMAT:
+        return Qwen314BPyptoExecutor
+    raise ValueError(f"unsupported model_format: {model_format!r}")
+
+
 def main() -> None:
     args = build_parser().parse_args()
     get_profiler(process_name="npu_generate")
@@ -417,14 +444,22 @@ def main() -> None:
     device_ids = parallel_config.replica_device_groups[0]
 
     kv_cache_manager = KvCacheManager()
-    executor = PyptoExecutor(
+    executor_cls = _executor_class_for_format(args.model_format)
+    executor_kwargs = {
+        "platform": args.platform,
+        "device_ids": device_ids,
+        "save_kernels_dir": args.save_kernels_dir,
+        "l3_trace": args.profile_verbose,
+    }
+    if args.model_format == _QWEN3_A8W8_FORMAT:
+        executor_kwargs["pto_isa_commit"] = args.pto_isa_commit
+    executor = executor_cls(
         kv_cache_manager,
-        platform=args.platform,
-        device_ids=device_ids,
-        save_kernels_dir=args.save_kernels_dir,
-        l3_trace=args.profile_verbose,
+        **executor_kwargs,
     )
+    model_loader = _model_loader_for_format(args.model_format)
     engine = LLMEngine(
+        model_loader=model_loader,
         kv_cache_manager=kv_cache_manager,
         executor=executor,
     )
@@ -437,15 +472,16 @@ def main() -> None:
         engine.init_model(
             model_id=args.model_id,
             model_dir=str(model_dir),
-            model_format="huggingface",
+            model_format=_QWEN3_A8W8_FORMAT if args.model_format == _QWEN3_A8W8_FORMAT else "huggingface",
+            decode_backend=args.decode_backend,
             runtime_config=RuntimeConfig(
                 page_size=128,
-                max_batch_size=16,
+                max_batch_size=args.max_batch_size,
                 max_seq_len=args.max_seq_len,
                 max_new_tokens=args.max_new_tokens,
                 device="cpu",
-                kv_dtype="bfloat16",
-                weight_dtype="float32",
+                kv_dtype="int8" if args.model_format == _QWEN3_A8W8_FORMAT else "bfloat16",
+                weight_dtype="bfloat16" if args.model_format == _QWEN3_A8W8_FORMAT else "float32",
             ),
         )
         if collector is not None: