NVIDIA / NVIDIA/TensorRT-LLM

Align locality-domain MoE quantization detection with Linear planning

Open
#18,377 2 comments 0 reactions 1 assignee View on GitHub

@peaceh-nv is already working on this.

Since Sep 17, 2026.

Model optimization
Dominant language
Python
Stars
14.7k
Forks
2.8k
Avg merge
2d 23h
Merged PRs (30d)
489

Description

Summary

Align quantization detection in LocalityDomainExecutionPlanner.plan_moe with LocalityDomainExecutionPlanner.plan_linear.

Rationale

plan_moe currently reads quant_config.quant_mode and calls has_any_quant() without excluding KV-cache-only quantization. plan_linear reads the intended layer quantization source and excludes KV-cache-only quantization. The two planners can therefore classify equivalent configurations differently. In particular, a KV-cache-quantized BF16 MoE can be rejected as unsupported instead of using the BF16 locality-domain MoE path.

Affected area

  • tensorrt_llm/_torch/locality_domain/policy.py
  • LocalityDomainExecutionPlanner.plan_moe

Required changes

  • Confirm and use the intended MoE layer quantization attribute.
  • Use the same quantization source for NVFP4 detection and general quantization detection.
  • Exclude KV-cache-only quantization when determining whether weight quantization is present.
  • Classify KV-cache-quantized BF16 MoE configurations as BF16 when all other BF16 requirements are met.

Acceptance criteria

  • plan_moe and plan_linear apply consistent effective-weight quantization semantics.
  • A KV-cache-only-quantized BF16 MoE does not reach the unsupported-quantization branch solely because of KV-cache quantization.
  • Unit coverage verifies the expected classification and disabled reasons.

Backlinks

Requested by: @farazkh80

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.