Fix TRT-LLM MLA backend applying k_scale to BF16 KV cache in BMM1

[debo3]

When a model checkpoint contains KV cache scaling factors (k_scale/v_scale) but the KV cache dtype is BF16 (not FP8), the TRT-LLM MLA backend unconditionally applies k_scale in the BMM1 attention score computation. This is incorrect because k_scale is a quantization compensation factor that should only be applied when KV cache values are actually FP8-quantized.

For example, with k_scale=0.06, attention scores are scaled down by ~16x, producing garbage output (degenerate repetition/random tokens).

This fix gates k_scale behind a self.data_type == torch.float8_e4m3fn check in both the decode and extend/target_verify paths, so k_scale is only applied when the KV cache actually stores FP8 values.

Affected code paths:

• forward_decode: BMM1 scale computation (~L915)
• forward_extend: target_verify/draft_extend BMM1 scale computation (~L1037)

Motivation

When serving FP8-quantized DeepSeek-V3 checkpoints that contain per-layer k_scale/v_scale tensors (from KV cache calibration), SGLang's weight loader (BaseKVCacheMethod) loads these into layer.k_scale_float regardless of what KV cache dtype the server is using. The TRT-LLM MLA backend then unconditionally applies k_scale in BMM1:

k_scale = layer.k_scale_float if getattr(layer, "k_scale_float", None) is not None else 1.0
bmm1_scale = q_scale * k_scale * layer.scaling

With BF16 KV cache (the default), key values are stored unscaled, so applying k_scale (typical values 0.02-0.06) reduces attention scores by 16-50x. This produces completely degenerate output: repetitive tokens, garbage characters, or premature EOS.

Modifications

Single file change in python/sglang/srt/layers/attention/trtllm_mla_backend.py:

• forward_decode (~L915): Wrapped k_scale lookup in if self.data_type == torch.float8_e4m3fn, defaulting to k_scale = 1.0 for non-FP8 KV cache.
• forward_extend (~L1037): Same change for the target_verify/draft_extend path.

self.data_type is set from model_runner.kv_cache_dtype at init (line 288), so it correctly reflects the runtime KV cache dtype.

Accuracy Tests

Tested on 8x NVIDIA B200 GPUs, TP=8, DeepSeek-V3 671B (FP8 weights), temperature=0, 10 diverse prompts (factual, reasoning, code, creative, translation, QA).

Model with k_scale/v_scale in checkpoint + BF16 KV cache (the bug scenario):

Prompt	main (before)	This PR (after)
"The capital of France is"	Paris 1 is111,1,0000000000...	Paris. It is the largest city in France and the country's main center of culture and commerce.
"Explain quantum computing in simple terms:"	Quantum Quantum Quantum Quantum Quantum... (28x repetition)	Quantum computing is a type of computing that uses quantum bits, or qubits, which can represent and store information...
"If a train travels at 60 mph for 2.5 hours, it will cover"	a 0.0.0.0.0.111.0.0.0 0000000...	a distance of 150 miles. This is because 60 miles per hour multiplied by 2.5 hours equals 150 miles.
"def is_prime(n):"	\nn\n (5 tokens then EOS)	if n <= 1: return False for i in range(2, int(n**0.5) + 1): if n % i == 0: return False return True
"What is photosynthesis?"	Photos What is the process ofs,1. The // k k |...	Photosynthesis is the process by which green plants, algae, and some bacteria convert light energy...

Summary: main produces 0/10 coherent outputs; this PR produces 10/10 coherent outputs.

No regression (model without k_scale/v_scale in checkpoint):

Both main and this PR produce byte-identical outputs for all 10 prompts on a standard DeepSeek-V3 checkpoint without k_scale/v_scale tensors. When layer.k_scale_float is None, both branches default to k_scale = 1.0.

Benchmarking and Profiling

No performance impact. The change adds a single if self.data_type == torch.float8_e4m3fn check (a Python attribute comparison) before the existing getattr call. This is negligible compared to the GPU kernel execution time.

Measured throughput on B200 with DeepSeek-V3 671B TP=8: ~106-118 tok/s on both branches (identical within noise).

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Fridge003]

/tag-and-rerun-ci