[Bug]: DSA indexer wk FP8 block scale is dropped at load on the PyTorch backend (DeepSeek-V3.2 / GLM-5 FP8 checkpoints)

System Info

• GPU: 8x NVIDIA H200 NVL (SM90; two 4-GPU NVLink islands, cross-island over PCIe/UPI), driver 595.71.05
• Host: Ubuntu 24.04.4, 2 TiB RAM, Xeon 6747P
• Container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc25 (CUDA 13.2, Python 3.12.3, torch 2.12.0a0+nv26.5, transformers 5.5.4)
• TensorRT-LLM: 1.3.0rc25 (tag v1.3.0rc25, 785c948); the code paths cited below were also checked on main @ a6616d6f (2026-09-03)
• Backend: PyTorch (trtllm-serve --backend pytorch)

Who can help?

@xwang233 (DSA indexer projection, #18264 / #12055) @NVShreyas (GLM-5 support, #11990)

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Every FP8 block-scaled DSA checkpoint stores the lightning-indexer key projection as FP8 with a 128x128 block scale: model.layers.L.self_attn.indexer.wk.weight (F8_E4M3 [128, hidden]) plus model.layers.L.self_attn.indexer.wk.weight_scale_inv (F32 [1, hidden/128]). This is the case for zai-org/GLM-5.3, zai-org/GLM-5.2-FP8, zai-org/GLM-5.1-FP8, deepseek-ai/DeepSeek-V3.2-Exp and W4AFP8 derivatives such as PhalaCloud/GLM-5.3-W4AFP8.

TRT-LLM deliberately keeps wk (and weights_proj) unquantized: Indexer.__init__ builds it as an fp32 Linear with quant_config=None (tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py, self.wk = Linear(..., dtype=torch.float32, quant_config=None, ...)). When the DeepSeek-V3 weight loader reaches that module it takes the generic branch of DeepseekV3WeightLoader.load_weights (tensorrt_llm/_torch/models/modeling_deepseekv3.py, module_weights = filter_weights(name, weights); module.load_weights(weights=[module_weights])), and UnquantizedLinearMethod.load_weights_vanilla -> copy_weight (tensorrt_llm/_torch/modules/linear.py) does src.to(dst.dtype): the FP8 codes are cast to fp32 and weight_scale_inv is never applied. wq_b is loaded correctly (FP8 block-scale Linear) and kv_b_proj has an explicit dequant path, but wk has neither.

Numeric reproduction (run inside the rc25 container with one GPU; check_indexer_wk_load.py builds the same fp32 Linear(quant_config=None), feeds it {"weight": fp8, "weight_scale_inv": scale} through Linear.load_weights, and compares):

checkpoint: model.layers.6.self_attn.indexer.wk.weight torch.float8_e4m3fn (128, 6144); scale torch.float32 (1, 48) min=0.000556 max=0.0034
loaded vs raw-cast  : rel-diff 0.000e+00  max-abs 0.000e+00
loaded vs dequant   : rel-diff 1.087e+03  max-abs 4.478e+02
|loaded| mean 39.39  |dequant| mean 0.03813  ratio 1033.0x
row cosine(loaded, dequant): min 0.8847 mean 0.9073
VERDICT: RAW CAST (scale dropped)

The per-block scales span 0.0004..0.0035 (a 4-6x spread across the 48 column blocks), so after k_norm the indexer keys are not only mis-scaled but directionally wrong (row cosine 0.91 vs the real weight). This only affects contexts longer than index_topk (2048 tokens), where the indexer actually selects, so short-prompt accuracy tests do not see it.

Behavioural reproduction (8x H200 NVL, PhalaCloud/GLM-5.3-W4AFP8 loaded through a MIXED_PRECISION quant config + #18393, TP8/EP8, greedy, --reasoning_parser deepseek-r1 --tool_parser glm47): with a ~7k-token system prompt and no tools the model emitted <tool_call>getWeather(...) repeated until max_tokens; with a ~52k-token prompt it exhausted max_tokens inside reasoning. After dequantizing wk at load (fix below) both prompts terminate normally (46-char and 2k-char completions), matching SGLang's first line on the same checkpoint. This looks like the same family as #15295 (GLM-5.1-FP8 corruption under tool use on H200), which also only appears with long prompts on the DSA path.

import argparse, json
from pathlib import Path
import torch
from safetensors.torch import load_file
from tensorrt_llm._torch.modules.linear import Linear
from tensorrt_llm._torch.models.modeling_deepseekv3 import weight_dequant

ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True); ap.add_argument("--layer", type=int, default=6)
args = ap.parse_args(); root = Path(args.model)
idx = json.loads((root / "model.safetensors.index.json").read_text())["weight_map"]
wname = f"model.layers.{args.layer}.self_attn.indexer.wk.weight"; sname = wname + "_scale_inv"
w = load_file(str(root / idx[wname]))[wname]; s = load_file(str(root / idx[sname]))[sname]
lin = Linear(w.shape[1], w.shape[0], bias=False, dtype=torch.float32, quant_config=None,
             skip_create_weights_in_init=True, use_custom_cublas_mm=True)
lin.create_weights(); lin = lin.cuda()
lin.load_weights(weights=[{"weight": w, "weight_scale_inv": s}])
loaded = lin.weight.data.float().cpu(); raw = w.float(); deq = weight_dequant(w.cuda(), s.cuda()).float().cpu()
rel = lambda a, b: ((a - b).norm() / b.norm()).item()
print("loaded vs raw-cast:", rel(loaded, raw), "loaded vs dequant:", rel(loaded, deq))
print("row cosine(loaded, dequant):", torch.nn.functional.cosine_similarity(loaded, deq, dim=1).mean().item())

Expected behavior

The indexer wk parameter equals weight_dequant(weight, weight_scale_inv) (cast to the module dtype), as it does for every other FP8 block-scaled projection in the checkpoint.

actual behavior

The parameter equals the raw FP8 codes cast to fp32 (bit-for-bit), i.e. the block scale is dropped, and long-context outputs degrade.

additional notes

• Fix proposed in a PR (linked below): dequantize FP8 block-scaled checkpoint tensors before they are handed to an unquantized Linear in the DeepSeek-V3 weight loader, reusing the existing weight_dequant triton kernel. It is a no-op for every other module.
• With TRTLLM_DSA_INDEXER_BF16=1 (#18264) the same raw cast happens into a bf16 parameter.
• Checked on main @ a6616d6f: the indexer is still built with quant_config=None and the generic loader branch is unchanged, so the defect is not rc25-specific.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Start with tensorrt_llm/_torch/models/modeling_deepseekv3.py and tensorrt_llm/_torch/modules/linear.py, then inspect the indexer definition in tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py. Run check_indexer_wk_load.py against a checkpoint to reproduce the raw-cast versus dequantized values. Done means the loaded wk parameter matches weight_dequant(weight, weight_scale_inv) and the long-context behavior no longer degrades.