feat(dispatch): add tsingmicro txda backend for vllm 0.24.0 - #421
Open
tengqm wants to merge 7 commits into
Open
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| 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 |
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).
- 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)
User: txda uses tccl rather than gloo for communication layer.
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.
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.
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 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.
tengqm
force-pushed
the
port/txda-tsingmicro
branch
from
August 31, 2026 12:32
57b15c7 to
bd010ce
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Port the tsingmicro (TX8110) txda backend to vllm 0.24.0 and fix the KV-cache write path that vLLM 0.24's
Attention.forwardrequires.Changes
unified_kv_cache_updateonly when the backend'sforward_includes_kv_cache_updateisFalse.AttentionFLBackendinherited vLLM'sTruedefault, so the KV cache was never written and forward read an empty cache — silent garbage on every platform using the flag_gems attention backend. Set it toFalse. (The fix existed on the release/0.2 lineage but never reached current main.)flash_attn_varlen_func/reshape_and_cache_flashcompute silently wrong values on TX8110. This backend reuses the flag_gems KV layout / block table / slot mapping, computes attention with torch SDPA, and writes the cache via plain indexing with aslot_mapping >= 0guard against padded slots.platform.py,VENDOR_DEVICE_MAP,tsingmicro.yamlconfig, referencefused_moeimpl (flag_gems moe failed to import on this stack), rms_norm/rotary routed to reference impls.Validation
flagos-runtime-tsingmicro-tsm260610:2.1.2, Qwen3-4B, LLM API (dtype=bfloat16, enforce_eager=True, max_model_len=1024, gpu_memory_utilization=0.9, tensor_parallel_size=1, seed=0).Follow-up
attention_backend: [flagos, vendor]— flagos raises on TX8110 (no CUDA), falls back tovendor.txda.This PR was written in part with the assistance of generative AI.