From f69e8070f494556033799e399e34967e503d2115 Mon Sep 17 00:00:00 2001 From: Qiming Teng Date: Sun, 30 Aug 2026 16:39:07 +0800 Subject: [PATCH 1/7] feat(dispatch): add tsingmicro txda backend for vllm 0.24.0 Port the tsingmicro (txda) backend from 8236c0a onto main: - utils.py: VENDOR_DEVICE_MAP tsingmicro -> txda, DeviceInfo.supported_device - platform.py: is_pin_memory_available + get_device_capability txda branches - dispatch/backends/vendor/txda: attention_backend-only registration (8236c0a registered 4 ops that are not on the Backend ABC and were silently dropped) Basing on main drops the 8236c0a router import (vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe) in favor of the 0.24.0 paths (vllm._aiter_ops / fused_moe.experts.rocm_aiter_moe). --- .../dispatch/backends/vendor/txda/__init__.py | 11 +++ .../backends/vendor/txda/register_ops.py | 53 +++++++++++++ vllm_fl/dispatch/backends/vendor/txda/txda.py | 76 +++++++++++++++++++ vllm_fl/platform.py | 5 +- vllm_fl/utils.py | 4 +- 5 files changed, 147 insertions(+), 2 deletions(-) create mode 100644 vllm_fl/dispatch/backends/vendor/txda/__init__.py create mode 100644 vllm_fl/dispatch/backends/vendor/txda/register_ops.py create mode 100644 vllm_fl/dispatch/backends/vendor/txda/txda.py diff --git a/vllm_fl/dispatch/backends/vendor/txda/__init__.py b/vllm_fl/dispatch/backends/vendor/txda/__init__.py new file mode 100644 index 000000000..2cbff9d7e --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/txda/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +""" +Txda (tsingmicro) backend for vllm-plugin-FL dispatch. + +This backend provides operator implementations for Tsingmicro TX devices. +""" + +from .txda import TxdaBackend + +__all__ = ["TxdaBackend"] diff --git a/vllm_fl/dispatch/backends/vendor/txda/register_ops.py b/vllm_fl/dispatch/backends/vendor/txda/register_ops.py new file mode 100644 index 000000000..de33ec956 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/txda/register_ops.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +""" +Txda (tsingmicro) backend operator registrations. + +This module registers the VENDOR (txda) implementations for the dispatch +system. Only the attention backend is registered: tsingmicro TX devices use +the flag_gems implementations for fused ops (silu_and_mul, rms_norm, +rotary_embedding), so no vendor overrides are needed for those. +""" + +from __future__ import annotations + +import functools + +from vllm_fl.dispatch.types import OpImpl, BackendImplKind, BackendPriority + + +def _bind_is_available(fn, is_available_fn): + """Wrap a function and bind _is_available attribute for OpImpl.""" + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + wrapper._is_available = is_available_fn + return wrapper + + +def register_builtins(registry) -> None: + """ + Register the txda (tsingmicro) VENDOR operator implementations. + + Args: + registry: Registry to register into + """ + from .txda import TxdaBackend + + backend = TxdaBackend() + is_avail = backend.is_available + + impls = [ + OpImpl( + op_name="attention_backend", + impl_id="vendor.txda", + kind=BackendImplKind.VENDOR, + fn=_bind_is_available(backend.attention_backend, is_avail), + vendor="txda", + priority=BackendPriority.VENDOR, + ), + ] + + registry.register_many(impls) diff --git a/vllm_fl/dispatch/backends/vendor/txda/txda.py b/vllm_fl/dispatch/backends/vendor/txda/txda.py new file mode 100644 index 000000000..1df9dcfff --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/txda/txda.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +""" +Txda (tsingmicro) backend implementation. + +This backend provides operator implementations for Tsingmicro TX devices. +For attention it uses the FlagGems attention backend. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from vllm_fl.dispatch.backends.base import Backend + + +class TxdaBackend(Backend): + """ + Txda (tsingmicro) backend for operator implementations. + + Tsingmicro TX devices use torch_txda (a PrivateUse1-based runtime), so the + CUDA-style fused ops are not applicable; dispatch falls back to the + flag_gems implementations. Only the attention backend is registered here. + """ + + _available: Optional[bool] = None + + @property + def name(self) -> str: + return "txda" + + @property + def vendor(self) -> Optional[str]: + return "txda" + + def is_available(self) -> bool: + """ + Check if tsingmicro TX hardware is available. + + Detection is based on the torch_txda runtime. + """ + if TxdaBackend._available is None: + try: + import torch_txda # noqa: F401 + TxdaBackend._available = ( + torch.txda.is_available() and torch.txda.device_count() > 0 + ) + except Exception: + TxdaBackend._available = False + return TxdaBackend._available + + # ==================== Operator Implementations ==================== + + def attention_backend( + self, use_mla: bool = False, use_sparse: bool = False + ) -> str: + """ + Get the attention backend class path for tsingmicro TX. + + Returns the FlagGems attention backend (MLA-aware). + + Args: + use_mla: Whether to use Multi-head Latent Attention (MLA) + use_sparse: Whether to use Deepseek Sparse Attention (DSA) + + Returns: + Fully qualified class path string + """ + if use_mla: + return "vllm_fl.dispatch.backends.flaggems.impl.mla.MLAFLBackend" + return ( + "vllm_fl.dispatch.backends.flaggems.impl.attention." + "AttentionFLBackend" + ) diff --git a/vllm_fl/platform.py b/vllm_fl/platform.py index c85b84b84..a59fcf275 100644 --- a/vllm_fl/platform.py +++ b/vllm_fl/platform.py @@ -133,7 +133,7 @@ def get_device_name(cls, device_id: int = 0) -> str: ### TODO(lms): change pin_memory depend device @classmethod def is_pin_memory_available(cls): - if cls.device_type in ["cuda", "xpu", "npu", "musa"]: + if cls.device_type in ["cuda", "xpu", "npu", "musa", "txda"]: return True return False @@ -457,6 +457,9 @@ def get_device_capability(cls, device_id: int = 0) -> DeviceCapability: # TODO: For PTPU/Sunrise devices, return None if cls.device_type == "ptpu": return None + if cls.device_type == "txda": + major, minor = torch.txda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) major, minor = torch.cuda.get_device_capability(device_id) return DeviceCapability(major=major, minor=minor) diff --git a/vllm_fl/utils.py b/vllm_fl/utils.py index 632dccc73..5d36614ef 100644 --- a/vllm_fl/utils.py +++ b/vllm_fl/utils.py @@ -49,6 +49,8 @@ "hygon": {"device_type": "cuda", "device_name": "cuda"}, # Registered backend: vendor/thead (PPU) "thead": {"device_type": "cuda", "device_name": "thead"}, + # Registered backend: vendor/txda (tsingmicro) + "tsingmicro": {"device_type": "txda", "device_name": "txda"}, } # Keep the vLLM base-class no-op for platforms not validated by this change. @@ -237,7 +239,7 @@ def get_op_config() -> Optional[dict[str, str]]: class DeviceInfo: def __init__(self): self.device = DeviceDetector() - self.supported_device = ["nvidia", "ascend", "metax", "mthreads", "sunrise", "thead"] + self.supported_device = ["nvidia", "ascend", "metax", "mthreads", "sunrise", "thead", "tsingmicro"] backend.set_torch_backend_device_fn(self.device.vendor_name) @property From 5ae34b68e35a4341d47afe69c075460a5da7645a Mon Sep 17 00:00:00 2001 From: Qiming Teng Date: Sun, 30 Aug 2026 22:01:23 +0800 Subject: [PATCH 2/7] fix(dispatch): tsingmicro txda runtime fixes for vllm 0.24.0 - model_runner: use a CPU generator for txda seeded sampling, since torch.Generator does not support the txda device - platform: map the txda device to the gloo distributed backend - graph: select torch.txda.TXDAGraph for graph capture on txda - config: add tsingmicro.yaml flagos blacklist (masked_fill/to_copy/copy triton kernels hang on TX8110 in the top-k sampler path) --- vllm_fl/compilation/graph.py | 2 ++ vllm_fl/dispatch/config/tsingmicro.yaml | 42 +++++++++++++++++++++++++ vllm_fl/platform.py | 1 + vllm_fl/worker/model_runner.py | 8 +++-- 4 files changed, 51 insertions(+), 2 deletions(-) create mode 100644 vllm_fl/dispatch/config/tsingmicro.yaml diff --git a/vllm_fl/compilation/graph.py b/vllm_fl/compilation/graph.py index 2ec96382c..b4352f097 100644 --- a/vllm_fl/compilation/graph.py +++ b/vllm_fl/compilation/graph.py @@ -49,6 +49,8 @@ class Graph: graph = torch.musa.MUSAGraph elif current_platform.device_type == "ptpu": graph = torch.ptpu.PTPUGraph + elif current_platform.device_type == "txda": + graph = torch.txda.TXDAGraph else: raise NotImplementedError("not support graph") diff --git a/vllm_fl/dispatch/config/tsingmicro.yaml b/vllm_fl/dispatch/config/tsingmicro.yaml new file mode 100644 index 000000000..c37166c23 --- /dev/null +++ b/vllm_fl/dispatch/config/tsingmicro.yaml @@ -0,0 +1,42 @@ +# vLLM-FL Dispatch Configuration for TSINGMICRO (TX8110 / TSM Runtime) +# Auto-loaded when running on TSINGMICRO txda hardware + +# Preferred default backend type: flagos, vendor, reference +prefer: flagos + +# Strict Mode: +# true = Raise an error immediately on failure; do not attempt other backends. +# false = Attempt the next available backend in sequence upon failure (Default). +strict: false + +# Per-operator backend execution order (Optional) +# Only the backends listed here will be attempted, in the order specified. +op_backends: + # attention_backend: prioritize flagos (Triton attention) + attention_backend: + - flagos + - reference + rms_norm: + - flagos + - reference + silu_and_mul: + - flagos + - reference + rotary_embedding: + - flagos + - reference + +# FlagOS operator blacklist +# to_copy/copy_/copy hang the TX8110 triton copy kernel on dtype-changing +# casts (e.g. int32 -> int64) in topk_topp sampling; mask these off. +flagos_blacklist: + - masked_fill + - masked_fill_ + - to_copy + - copy_ + - copy + +# OOT (Out-of-Tree) operator blacklist +oot_blacklist: + - masked_fill + - masked_fill_ diff --git a/vllm_fl/platform.py b/vllm_fl/platform.py index a59fcf275..a758970d5 100644 --- a/vllm_fl/platform.py +++ b/vllm_fl/platform.py @@ -49,6 +49,7 @@ "npu": "hccl", "cuda": "nccl", "musa": "mccl", + "txda": "gloo", } diff --git a/vllm_fl/worker/model_runner.py b/vllm_fl/worker/model_runner.py index 343dfdd45..9be0945f5 100644 --- a/vllm_fl/worker/model_runner.py +++ b/vllm_fl/worker/model_runner.py @@ -1267,7 +1267,9 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None sampling_params and sampling_params.sampling_type == SamplingType.RANDOM_SEED ): - generator = torch.Generator(device=self.device) + generator = torch.Generator( + device="cpu" if self.device.type == "txda" else self.device + ) generator.manual_seed(sampling_params.seed) else: generator = None @@ -6186,7 +6188,9 @@ def _dummy_sampler_run( sampling_metadata=replace( dummy_metadata, generators={ - 0: torch.Generator(device=self.device).manual_seed(0) + 0: torch.Generator( + device="cpu" if self.device.type == "txda" else self.device + ).manual_seed(0) }, ), ) From a62a8fc6a7d2e89c09f4b68f7cf64badfc5d732f Mon Sep 17 00:00:00 2001 From: Qiming Teng Date: Sun, 30 Aug 2026 22:16:56 +0800 Subject: [PATCH 3/7] fix(platform): txda uses tccl not gloo for communication User: txda uses tccl rather than gloo for communication layer. --- vllm_fl/platform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_fl/platform.py b/vllm_fl/platform.py index a758970d5..e51ef5970 100644 --- a/vllm_fl/platform.py +++ b/vllm_fl/platform.py @@ -49,7 +49,7 @@ "npu": "hccl", "cuda": "nccl", "musa": "mccl", - "txda": "gloo", + "txda": "tccl", } From d055dfb7465d436a750d228b38d4f5898b15ad4f Mon Sep 17 00:00:00 2001 From: Qiming Teng Date: Mon, 31 Aug 2026 07:53:16 +0800 Subject: [PATCH 4/7] fix(dispatch): route tsingmicro attention_backend to vendor fallback flagos attention_backend is CUDA-gated and raises RuntimeError on txda; reference backend registration is broken upstream (ReferenceBackend lacks moe_align_block_size/moe_sum/topk_softmax/grouped_topk) so [flagos, reference] yields exactly one candidate that always fails. vendor.txda returns AttentionFLBackend (flag_gems), the intended txda impl. --- vllm_fl/dispatch/config/tsingmicro.yaml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm_fl/dispatch/config/tsingmicro.yaml b/vllm_fl/dispatch/config/tsingmicro.yaml index c37166c23..d71dccb25 100644 --- a/vllm_fl/dispatch/config/tsingmicro.yaml +++ b/vllm_fl/dispatch/config/tsingmicro.yaml @@ -12,10 +12,13 @@ strict: false # Per-operator backend execution order (Optional) # Only the backends listed here will be attempted, in the order specified. op_backends: - # attention_backend: prioritize flagos (Triton attention) + # attention_backend: flagos (Triton attention) is CUDA-gated and raises on + # txda, so fall back to vendor.txda -> AttentionFLBackend (flag_gems). + # reference is not listed: its registration is broken upstream (missing + # ReferenceBackend methods) and its impl returns CUDA-only FLASH_ATTN. attention_backend: - flagos - - reference + - vendor rms_norm: - flagos - reference From b04a08b975c0bb550a09ddf950b02b45e32e97e9 Mon Sep 17 00:00:00 2001 From: Qiming Teng Date: Mon, 31 Aug 2026 12:35:15 +0800 Subject: [PATCH 5/7] fix(dispatch): repair reference backend registration register_builtins eagerly resolved backend.moe_align_block_size / moe_sum / topk_softmax / grouped_topk, none of which exist on ReferenceBackend. The AttributeError aborted the whole registration, so the 'reference' token matched no impls on any platform. Add the four missing MoE reference implementations (impl/fused_moe.py, exact flag_gems semantics) and make registration per-op fault tolerant: a missing method logs a warning and skips just that op instead of aborting everything. --- .../backends/reference/impl/fused_moe.py | 215 ++++++++++++++++++ .../dispatch/backends/reference/reference.py | 67 ++++++ .../backends/reference/register_ops.py | 146 ++++-------- 3 files changed, 325 insertions(+), 103 deletions(-) create mode 100644 vllm_fl/dispatch/backends/reference/impl/fused_moe.py diff --git a/vllm_fl/dispatch/backends/reference/impl/fused_moe.py b/vllm_fl/dispatch/backends/reference/impl/fused_moe.py new file mode 100644 index 000000000..5bb3b2d1c --- /dev/null +++ b/vllm_fl/dispatch/backends/reference/impl/fused_moe.py @@ -0,0 +1,215 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +""" +Reference fused-MoE operator implementations using PyTorch. + +The flag_gems (flagtree-compiled) MoE kernels compute wrong values on +TX8110 without raising, so these ops are routed to pure-torch fallbacks +that are numerically fine on txda (see config/tsingmicro.yaml). +""" + +from __future__ import annotations + +import torch + +from vllm.utils.math_utils import round_up + + +def moe_align_block_size_torch( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure-torch port of flag_gems' moe_align_block_size_triton semantics. + + Allocation math matches flaggems (numel + num_experts*(block_size-1), + round_up, numel= 0) & (expert_of_sorted < num_experts) + order = order[keep] + expert_of_sorted = expert_of_sorted[keep] + num_valid = order.numel() + + # cdiv(counts, bs)*bs per expert, then exclusive-scan into offsets. + counts = torch.zeros(num_experts, dtype=torch.int32, device=topk_ids.device) + ones = torch.ones(num_valid, dtype=torch.int32, device=topk_ids.device) + counts.scatter_add_(0, expert_of_sorted, ones) # bincount absent on txda + aligned = ((counts + block_size - 1) // block_size) * block_size + offsets = aligned.cumsum(0, dtype=torch.int64) + starts = offsets - aligned + + total = int(offsets[-1].item()) + num_tokens_post_pad[0] = total + max_num_m_blocks = (total + block_size - 1) // block_size + expert_ids = torch.full( + (max_num_m_blocks,), -1, dtype=torch.int32, device=topk_ids.device + ) + + # Scatter tokens into each expert's aligned span: the triton kernel + # reads tokens at block-aligned offsets, so per-expert padding must be + # materialized in the output (a plain prefix fill shifts later experts). + # starts[e] is the padded offset of expert e; within-expert tokens stay + # in original (stable) order, so position = starts[e] + within-index. + starts = (offsets - aligned).to(torch.int64) + run_start = torch.zeros(num_experts, dtype=torch.int64, device=topk_ids.device) + present = counts > 0 + run_start[present] = counts.cumsum(0)[present] - counts[present] + within = torch.arange(num_valid, device=topk_ids.device) - run_start[expert_of_sorted] + positions = starts[expert_of_sorted] + within + sorted_ids[positions] = order.to(torch.int32) + # Each block belongs to the expert that owns its padded span. Aligned + # counts are multiples of block_size, so a block never straddles experts + # and expert_of_padded[block_start] is unambiguous; there is no partial + # final block. (Indexing the valid-only list by padded positions, as an + # earlier version did, labels the wrong blocks whenever an expert's + # aligned span exceeds one block.) + expert_of_padded = torch.repeat_interleave( + torch.arange(num_experts, dtype=torch.int32, device=topk_ids.device), aligned + ) + block_starts = torch.arange(max_num_m_blocks, device=topk_ids.device) * block_size + expert_ids[:] = expert_of_padded[block_starts] + + if expert_map is not None: + # Padding blocks (-1) are never read by the moe kernel; map only + # real experts so -1 cannot wrap to a valid map entry. + expert_ids = torch.where( + expert_ids == -1, -1, expert_map[expert_ids.clamp(min=0)] + ) + + return sorted_ids, expert_ids, num_tokens_post_pad + + +def moe_sum_torch(inp: torch.Tensor, out: torch.Tensor) -> None: + """Sum over the intermediate-cache hidden dim, written into out in-place.""" + reduced = inp.float().sum(dim=1) + out.copy_(reduced.to(dtype=out.dtype)) + + +def topk_softmax_torch( + topk_weights: torch.Tensor, + topk_indices: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Full-row softmax -> topk, matching vllm csrc topk_softmax. + + Fills the three pre-allocated output tensors in-place; weights are the + raw (post-softmax) scores at the selected experts. + """ + scores = torch.softmax(gating_output.float(), dim=-1) + vals, idx = torch.topk(scores, k=topk_weights.size(-1), dim=-1, sorted=False) + if renormalize: + vals = vals / vals.sum(-1, keepdim=True) + topk_weights.copy_(vals) + topk_indices.copy_(idx.to(topk_indices.dtype)) + token_expert_indices.copy_(idx.to(torch.int32)) + return topk_weights, topk_indices + + +def grouped_topk_torch( + scores: torch.Tensor, + n_group: int, + topk_group: int, + topk: int, + renormalize: bool, + routed_scaling_factor: float, + bias: torch.Tensor | None, + scoring_func: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pure-torch replica of flag_gems' grouped_topk semantics. + + Group score = top-2 sum of (score+bias) within the group (flag_gems + max1+max2 convention), groups selected by that score, then top-k by + selection score (score+bias, -inf outside selected groups). Output + weights are the RAW processed scores (no bias), float32, matching + flag_gems so the F (flagtree) and T (triton) paths route identically. + """ + scores = scores.float() + M, num_experts = scores.shape + assert num_experts % n_group == 0 + assert scoring_func in (0, 1) + + if bias is None: + bias = torch.zeros(num_experts, dtype=scores.dtype, device=scores.device) + else: + bias = bias.to(dtype=scores.dtype, device=scores.device).reshape(-1) + assert bias.numel() == num_experts + + scores_processed = torch.sigmoid(scores) if scoring_func == 1 else scores + scored = scores_processed + bias + + # Top-2 sum of the scored values within each group (flag_gems convention). + group_scores = scored.view(M, n_group, -1).topk(2, dim=-1).values.sum(-1) + # NaN group scores count as non-finite (kernel converts them to -inf). + finite = torch.isfinite(group_scores) + group_scores = torch.where( + finite, group_scores, torch.full_like(group_scores, -float("inf")) + ) + + # Select topk_group groups, top-k by scored value inside them, output the + # RAW processed score (no bias). Rows whose max group score is -inf fall + # back to the flag_gems default (1/topk, 0..topk-1), applied after + # renormalize/scaling so defaults stay unscaled. + group_idx = group_scores.topk(topk_group, dim=-1, sorted=False).indices + mask = torch.zeros(M, n_group, dtype=torch.bool, device=scores.device) + mask.scatter_(1, group_idx, True) + mask = mask.repeat_interleave(num_experts // n_group, dim=1) + selection = torch.where(mask, scored, torch.full_like(scored, -float("inf"))) + topk_values, topk_indices = selection.topk(topk, dim=-1, sorted=False) + topk_weights = scores_processed.gather(-1, topk_indices) + + if renormalize: + topk_weights = ( + topk_weights / (topk_weights.sum(-1, keepdim=True) + 1e-20) + ) * routed_scaling_factor + else: + topk_weights = topk_weights * routed_scaling_factor + + if_proceed = group_scores.max(-1).values != -float("inf") + default_vals = torch.full( + (M, topk), 1.0 / topk, dtype=torch.float32, device=scores.device + ) + default_idx = torch.arange(topk, dtype=torch.int32, device=scores.device) + default_idx = default_idx.expand(M, -1) + topk_weights = torch.where( + if_proceed.unsqueeze(-1), topk_weights, default_vals + ) + topk_indices = torch.where( + if_proceed.unsqueeze(-1), topk_indices.to(torch.int32), default_idx + ) + + return topk_weights.to(torch.float32), topk_indices.to(torch.int32) diff --git a/vllm_fl/dispatch/backends/reference/reference.py b/vllm_fl/dispatch/backends/reference/reference.py index 015dd2afc..61c8911c8 100644 --- a/vllm_fl/dispatch/backends/reference/reference.py +++ b/vllm_fl/dispatch/backends/reference/reference.py @@ -167,6 +167,73 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> return AttentionBackendEnum.FLASHMLA.get_path() return AttentionBackendEnum.FLASH_ATTN.get_path() + def moe_align_block_size( + self, + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: Optional[torch.Tensor] = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from .impl.fused_moe import moe_align_block_size_torch + + return moe_align_block_size_torch( + topk_ids, + block_size, + num_experts, + expert_map=expert_map, + pad_sorted_ids=pad_sorted_ids, + ignore_invalid_experts=ignore_invalid_experts, + ) + + def moe_sum(self, inp: torch.Tensor, out: torch.Tensor) -> None: + from .impl.fused_moe import moe_sum_torch + + moe_sum_torch(inp, out) + + def topk_softmax( + self, + topk_weights: torch.Tensor, + topk_indices: torch.Tensor, + token_expert_indices: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + from .impl.fused_moe import topk_softmax_torch + + return topk_softmax_torch( + topk_weights, + topk_indices, + token_expert_indices, + gating_output, + renormalize, + ) + + def grouped_topk( + self, + scores: torch.Tensor, + n_group: int, + topk_group: int, + topk: int, + renormalize: bool, + routed_scaling_factor: float, + bias: Optional[torch.Tensor] = None, + scoring_func: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + from .impl.fused_moe import grouped_topk_torch + + return grouped_topk_torch( + scores, + n_group, + topk_group, + topk, + renormalize, + routed_scaling_factor, + bias, + scoring_func, + ) + def invoke_fused_moe_triton_kernel( self, A, diff --git a/vllm_fl/dispatch/backends/reference/register_ops.py b/vllm_fl/dispatch/backends/reference/register_ops.py index 8eaf19f83..ba0fc2e98 100644 --- a/vllm_fl/dispatch/backends/reference/register_ops.py +++ b/vllm_fl/dispatch/backends/reference/register_ops.py @@ -9,9 +9,12 @@ from __future__ import annotations import functools +import logging from vllm_fl.dispatch.types import BackendImplKind, BackendPriority, OpImpl +logger = logging.getLogger(__name__) + def _bind_is_available(fn, is_available_fn): """Wrap a function and bind _is_available attribute for OpImpl.is_available() check.""" @@ -24,6 +27,25 @@ def wrapper(*args, **kwargs): return wrapper +# (op_name, backend method name) pairs. getattr is resolved lazily below so a +# single missing method logs a warning and skips that op instead of aborting +# the whole reference registration (which would leave every `reference` token +# in a platform yaml matching nothing). +_REFERENCE_OPS = [ + ("dynamic_per_token_quant_int8", "dynamic_per_token_quant_int8"), + ("silu_and_mul", "silu_and_mul"), + ("gelu_and_mul", "gelu_and_mul"), + ("rms_norm", "rms_norm"), + ("rotary_embedding", "rotary_embedding"), + ("attention_backend", "attention_backend"), + ("moe_align_block_size", "moe_align_block_size"), + ("moe_sum", "moe_sum"), + ("topk_softmax", "topk_softmax"), + ("invoke_fused_moe_triton_kernel", "invoke_fused_moe_triton_kernel"), + ("grouped_topk", "grouped_topk"), +] + + def register_builtins(registry) -> None: """ Register all PyTorch (REFERENCE) operator implementations. @@ -36,108 +58,26 @@ def register_builtins(registry) -> None: backend = ReferenceBackend() is_avail = backend.is_available - impls = [ - # Quantization - OpImpl( - op_name="dynamic_per_token_quant_int8", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available( - backend.dynamic_per_token_quant_int8, - is_avail, - ), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # Activation - OpImpl( - op_name="silu_and_mul", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.silu_and_mul, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - OpImpl( - op_name="gelu_and_mul", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.gelu_and_mul, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # Normalization - OpImpl( - op_name="rms_norm", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.rms_norm, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # Rotary Embedding - OpImpl( - op_name="rotary_embedding", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.rotary_embedding, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # Attention Backend - OpImpl( - op_name="attention_backend", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.attention_backend, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # MoE align - OpImpl( - op_name="moe_align_block_size", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.moe_align_block_size, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # MoE sum - OpImpl( - op_name="moe_sum", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.moe_sum, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # topk softmax - OpImpl( - op_name="topk_softmax", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.topk_softmax, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # invoke fused moe triton kernel - OpImpl( - op_name="invoke_fused_moe_triton_kernel", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.invoke_fused_moe_triton_kernel, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # grouped topk - OpImpl( - op_name="grouped_topk", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.grouped_topk, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - ] + impls = [] + for op_name, method_name in _REFERENCE_OPS: + try: + method = getattr(backend, method_name) + except AttributeError: + logger.warning( + "Reference backend missing %s; skipping op %s", + method_name, + op_name, + ) + continue + impls.append( + OpImpl( + op_name=op_name, + impl_id="reference.torch", + kind=BackendImplKind.REFERENCE, + fn=_bind_is_available(method, is_avail), + vendor=None, + priority=BackendPriority.REFERENCE, + ) + ) registry.register_many(impls) From 25b50d533a9a2cc9275e31a62c79fb58fe0a6d59 Mon Sep 17 00:00:00 2001 From: Qiming Teng Date: Mon, 31 Aug 2026 12:35:15 +0800 Subject: [PATCH 6/7] fix(dispatch): route tsingmicro rms_norm/rotary to reference impls flag_gems kernels compiled with flagtree compute silently wrong on TX8110 (rms_norm maxrel ~7228.7, rotary ~768) while reference.torch is accurate (maxrel < 1%). Route rms_norm and rotary_embedding to reference first on tsingmicro; keep silu_and_mul and attention_backend on flagos. --- vllm_fl/dispatch/config/tsingmicro.yaml | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/vllm_fl/dispatch/config/tsingmicro.yaml b/vllm_fl/dispatch/config/tsingmicro.yaml index d71dccb25..42fba91f5 100644 --- a/vllm_fl/dispatch/config/tsingmicro.yaml +++ b/vllm_fl/dispatch/config/tsingmicro.yaml @@ -19,15 +19,20 @@ op_backends: attention_backend: - flagos - vendor + # rms_norm/rotary_embedding: flag_gems (flagtree-compiled) kernels compute + # WRONG values on TX8110 without raising (probe: rms_norm maxrel ~7200, + # rope maxabs ~1.7-2.9 vs float32 ref), so reference-first: strict:false only + # falls back on exception, and wrong output never triggers it. Reference.torch + # impls are numerically fine on txda (maxrel < 1%, bf16 precision). rms_norm: - - flagos - reference + - flagos silu_and_mul: - flagos - reference rotary_embedding: - - flagos - reference + - flagos # FlagOS operator blacklist # to_copy/copy_/copy hang the TX8110 triton copy kernel on dtype-changing From bd010cec12d094bb19439bd0a856e09bb89d5821 Mon Sep 17 00:00:00 2001 From: Qiming Teng Date: Mon, 31 Aug 2026 20:32:03 +0800 Subject: [PATCH 7/7] fix(attention): set forward_includes_kv_cache_update=False vLLM 0.24 calls unified_kv_cache_update only when the backend flag is False; AttentionFLBackend inherited the True default, so the KV cache was never written and forward read an empty cache (garbage on every flag_gems attention platform). Add TxdaSDPAAttentionBackend: flag_gems flash_attn kernels compute wrong values on TX8110, so compute attention with torch SDPA instead, reusing the flag_gems KV layout/metadata. Writes the cache via plain indexing with a slot_mapping >= 0 guard. --- .../backends/flaggems/impl/attention.py | 7 + .../backends/vendor/txda/impl/__init__.py | 0 .../backends/vendor/txda/impl/attention.py | 262 ++++++++++++++++++ vllm_fl/dispatch/backends/vendor/txda/txda.py | 9 +- 4 files changed, 275 insertions(+), 3 deletions(-) create mode 100644 vllm_fl/dispatch/backends/vendor/txda/impl/__init__.py create mode 100644 vllm_fl/dispatch/backends/vendor/txda/impl/attention.py diff --git a/vllm_fl/dispatch/backends/flaggems/impl/attention.py b/vllm_fl/dispatch/backends/flaggems/impl/attention.py index e95a5ced3..d05414faa 100644 --- a/vllm_fl/dispatch/backends/flaggems/impl/attention.py +++ b/vllm_fl/dispatch/backends/flaggems/impl/attention.py @@ -52,6 +52,13 @@ class AttentionFLBackend(AttentionBackend): accept_output_buffer: bool = True supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] + # do_kv_cache_update is invoked by vLLM's unified_kv_cache_update custom + # op before forward() -- which Attention.forward only calls when this flag + # is False. The inherited True default silently skips the KV write and + # leaves forward() reading an empty cache (garbage outputs on every + # platform using this backend). + forward_includes_kv_cache_update: bool = False + @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: vllm_config = get_current_vllm_config() diff --git a/vllm_fl/dispatch/backends/vendor/txda/impl/__init__.py b/vllm_fl/dispatch/backends/vendor/txda/impl/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_fl/dispatch/backends/vendor/txda/impl/attention.py b/vllm_fl/dispatch/backends/vendor/txda/impl/attention.py new file mode 100644 index 000000000..2481d5098 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/txda/impl/attention.py @@ -0,0 +1,262 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +""" +Txda (tsingmicro) SDPA reference attention backend. + +The flag_gems flash_attn_varlen_func kernel computes silently wrong values on +TX8110 (probe: maxrel=inf/nan/37.7 for causal/noncausal varlen cases), so the +flag_gems AttentionFLBackend cannot be used. This backend reuses the flag_gems +metadata machinery (KV layout, block table, slot mapping) but computes attention +with torch SDPA, which is numerically correct on txda. Compiler-independent: +works under both the flagtree and triton compilers. +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch + +from vllm.v1.attention.backend import AttentionType +from vllm_fl.dispatch.backends.flaggems.impl.attention import ( + AttentionFLBackend, + AttentionFLImpl, +) + +_DEBUG = os.environ.get("FL_DEBUG_TXDA_ATTN") == "1" +_PRINTED = [0] + + +class TxdaSDPAAttentionBackend(AttentionFLBackend): + """Attention backend for tsingmicro TX devices using torch SDPA.""" + + # AttentionFLBackend inherits forward_includes_kv_cache_update=True from + # vllm's AttentionBackend default, but its do_kv_cache_update is designed + # to be called separately by vLLM's unified_kv_cache_update custom op -- + # which Attention.forward only invokes when this flag is False. With the + # True default the KV cache is never written and forward reads zeros + # (silent garbage). Every vLLM backend overrides it to False; so do we. + forward_includes_kv_cache_update: bool = False + + # get_name is inherited: vLLM requires the name to be a member of + # AttentionBackendEnum ("CUSTOM"), which AttentionFLBackend already returns. + + @staticmethod + def get_impl_cls() -> type["TxdaSDPAAttentionImpl"]: + return TxdaSDPAAttentionImpl + + +class TxdaSDPAAttentionImpl(AttentionFLImpl): + """ + SDPA-based attention impl for TX8110. + + do_kv_cache_update and forward are overridden to avoid the flag_gems + kernels (reshape_and_cache_flash / flash_attn_varlen_func) that compute + silently wrong values on TX8110. The KV cache layout + (2, num_blocks, block_size, num_kv_heads, head_size) and metadata are + unchanged, so AttentionFLMetadataBuilder is reused as-is. + """ + + def do_kv_cache_update( + self, + layer: torch.nn.Module, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ): + """Write key/value into the paged KV cache via plain indexing. + + Avoids flag_gems reshape_and_cache_flash (wrong on TX8110). Indexing + by (block_id, offset) is layout-independent. + """ + if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): + return + + key_cache, value_cache = kv_cache.unbind(0) + block_size = key_cache.shape[1] + block_ids = slot_mapping // block_size + offsets = slot_mapping % block_size + # Padded slots carry slot_mapping == -1; without this guard they index + # the last block (floor division: -1 // block_size == -1) and corrupt + # the cache with padded garbage. + valid = slot_mapping >= 0 + key_cache[block_ids[valid], offsets[valid]] = key[valid] + value_cache[block_ids[valid], offsets[valid]] = value[valid] + + if _DEBUG and _PRINTED[0] < 400: + _PRINTED[0] += 1 + print( + f"[txda-debug] kv_update#{_PRINTED[0]} n={key.shape[0]} " + f"k0={key[0].reshape(-1)[:4].tolist()} " + f"v0={value[0].reshape(-1)[:4].tolist()} " + f"slot0={slot_mapping[0].item()} slotN={slot_mapping[-1].item()} " + f"bids0={block_ids[:4].tolist()} offs0={offsets[:4].tolist()} " + f"block_size={block_size}", + flush=True, + ) + + def forward( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata, + output: Optional[torch.Tensor] = None, + output_scale: Optional[torch.Tensor] = None, + output_block_scale: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward with per-request torch SDPA on the paged KV cache.""" + assert output is not None, "Output tensor must be provided." + + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError( + "fused output quantization is not supported for TxdaSDPAAttentionImpl" + ) + + if attn_metadata is None: + # Profiling run. + return output.fill_(0) + + if self.alibi_slopes is not None: + raise NotImplementedError("alibi not supported on TXDA_SDPA") + if self.logits_soft_cap: + raise NotImplementedError("logits soft cap not supported on TXDA_SDPA") + if attn_metadata.use_cascade: + raise NotImplementedError("cascade attention not supported on TXDA_SDPA") + + num_actual_tokens = attn_metadata.num_actual_tokens + query = query[:num_actual_tokens] + output = output[:num_actual_tokens] + + if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): + return self._forward_encoder(query, key, value, output, attn_metadata) + + key_cache, value_cache = kv_cache.unbind(0) + cu_seqlens_q = attn_metadata.query_start_loc + seq_lens = attn_metadata.seq_lens + block_table = attn_metadata.block_table + + num_reqs = cu_seqlens_q.shape[0] - 1 + window_left = self.sliding_window[0] # -1 means no sliding window + for i in range(num_reqs): + qs, qe = cu_seqlens_q[i].item(), cu_seqlens_q[i + 1].item() + q_len = qe - qs + seq_len = seq_lens[i].item() + if q_len == 0 or seq_len == 0: + output[qs:qe] = 0 + continue + + # Gather the request's KV from the paged cache. Advanced indexing + # (key_cache[blocks]) has no PrivateUse1 kernel on txda, so it falls + # back to a CPU copy of the whole cache and hangs at engine scale. + # Gather via per-block slices + cat instead (probe-verified). + blocks = block_table[i].tolist() + k = torch.cat([key_cache[b] for b in blocks], dim=0).reshape( + -1, self.num_kv_heads, self.head_size + )[:seq_len] + v = torch.cat([value_cache[b] for b in blocks], dim=0).reshape( + -1, self.num_kv_heads, self.head_size + )[:seq_len] + q = query[qs:qe] + + out_i = self._sdpa(q, k, v, seq_len, window_left) + # output is [num_tokens, num_heads, head_size]; out_i matches directly. + output[qs:qe] = out_i + + if _DEBUG and i == 0 and _PRINTED[0] < 400: + _PRINTED[0] += 1 + print( + f"[txda-debug] fwd#{_PRINTED[0]} layer={getattr(layer, 'name', '?')} " + f"n={num_actual_tokens} cu_q={cu_seqlens_q.tolist()} " + f"seq_lens={seq_lens.tolist()} reqs={num_reqs} " + f"bt0={blocks[:4]} seq_len={seq_len} q_len={q_len} " + f"k_rb0={k[0].reshape(-1)[:4].tolist()} " + f"q0={q[0].reshape(-1)[:4].tolist()} " + f"out0={out_i[0].reshape(-1)[:4].tolist()}", + flush=True, + ) + + return output + + def _sdpa( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + seq_len: int, + window_left: int, + ) -> torch.Tensor: + """SDPA for one request. + + q: [q_len, num_heads, head_size]; k/v: [seq_len, num_kv_heads, head_size]. + Returns [q_len, num_heads, head_size]. + """ + q_len = q.shape[0] + # Query tokens are the last q_len of the request's sequence. + q_start = seq_len - q_len + + q = q.permute(1, 0, 2).unsqueeze(0) # [1, H, q_len, D] + kk = k.permute(1, 0, 2).unsqueeze(0) # [1, kv_h, seq, D] + vv = v.permute(1, 0, 2).unsqueeze(0) + if self.num_queries_per_kv > 1: + kk = kk.repeat_interleave(self.num_queries_per_kv, dim=1) + vv = vv.repeat_interleave(self.num_queries_per_kv, dim=1) + + # Mask: key position j is visible to query row i iff + # j <= q_start + i, plus the sliding-window left bound. + causal = False + attn_mask = None + if q_len == 1: + # Decode: the single new token attends all keys; no mask needed. + pass + elif q_start == 0 and window_left < 0: + causal = True # Full prefill: plain causal. + else: + rows = torch.arange(q_len, device=q.device).unsqueeze(1) + cols = torch.arange(seq_len, device=q.device).unsqueeze(0) + visible = cols <= (q_start + rows) + if window_left >= 0: + visible &= cols >= (q_start + rows - window_left) + attn_mask = visible + + out = torch.nn.functional.scaled_dot_product_attention( + q, + kk, + vv, + attn_mask=attn_mask, + is_causal=causal, + scale=self.scale, + ) + return out.permute(0, 2, 1, 3) # [1, q_len, H, D] + + def _forward_encoder( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + attn_metadata, + ) -> torch.Tensor: + """Encoder attention over contiguous q/k/v (no paged cache).""" + cu_seqlens_q = attn_metadata.query_start_loc + num_reqs = cu_seqlens_q.shape[0] - 1 + for i in range(num_reqs): + qs, qe = cu_seqlens_q[i].item(), cu_seqlens_q[i + 1].item() + q_len = qe - qs + if q_len == 0: + continue + q = query[qs:qe].permute(1, 0, 2).unsqueeze(0) + k = key[qs:qe].permute(1, 0, 2).unsqueeze(0) + v = value[qs:qe].permute(1, 0, 2).unsqueeze(0) + if self.num_queries_per_kv > 1: + k = k.repeat_interleave(self.num_queries_per_kv, dim=1) + v = v.repeat_interleave(self.num_queries_per_kv, dim=1) + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, is_causal=False, scale=self.scale + ) + output[qs:qe] = out.permute(0, 2, 1, 3).squeeze(0) + return output diff --git a/vllm_fl/dispatch/backends/vendor/txda/txda.py b/vllm_fl/dispatch/backends/vendor/txda/txda.py index 1df9dcfff..c52fa58b1 100644 --- a/vllm_fl/dispatch/backends/vendor/txda/txda.py +++ b/vllm_fl/dispatch/backends/vendor/txda/txda.py @@ -59,7 +59,10 @@ def attention_backend( """ Get the attention backend class path for tsingmicro TX. - Returns the FlagGems attention backend (MLA-aware). + Returns the txda SDPA backend (reuses the flag_gems metadata machinery + but computes attention with torch SDPA, which is numerically correct on + TX8110 where flag_gems flash_attn_varlen_func is not). The MLA branch + still points at the flag_gems MLA backend; MLA is unverified on TX8110. Args: use_mla: Whether to use Multi-head Latent Attention (MLA) @@ -71,6 +74,6 @@ def attention_backend( if use_mla: return "vllm_fl.dispatch.backends.flaggems.impl.mla.MLAFLBackend" return ( - "vllm_fl.dispatch.backends.flaggems.impl.attention." - "AttentionFLBackend" + "vllm_fl.dispatch.backends.vendor.txda.impl.attention." + "TxdaSDPAAttentionBackend" )