diff --git a/QEfficient/utils/test_utils.py b/QEfficient/utils/test_utils.py index c91248c3e8..598c8ddf96 100644 --- a/QEfficient/utils/test_utils.py +++ b/QEfficient/utils/test_utils.py @@ -75,8 +75,9 @@ def load_qeff_causal_lm_model( continuous_batching: bool = False, qaic_config: Dict = None, config: Optional[AutoConfig] = None, + torch_dtype: Optional[torch.dtype] = torch.float32, ): - kwargs = dict(continuous_batching=continuous_batching, qaic_config=qaic_config) + kwargs = dict(continuous_batching=continuous_batching, qaic_config=qaic_config, torch_dtype=torch_dtype) if config is None: if num_hidden_layers != -1: kwargs["num_hidden_layers"] = num_hidden_layers diff --git a/scripts/Jenkinsfile b/scripts/Jenkinsfile index f437a1521a..53554fb2f6 100644 --- a/scripts/Jenkinsfile +++ b/scripts/Jenkinsfile @@ -1,59 +1,62 @@ -def testFilter(String profile) { - switch (profile) { - case 'dummy_layers_model': - return '(not full_layers) and (not few_layers)' - case 'few_layers_model': - return '(not full_layers) and (not dummy_layers)' - case 'full_layers_model': - return '(not dummy_layers) and (not few_layers)' - default: - error "Unsupported TEST_PROFILE value: ${profile}" - } -} +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// Common ignore flags shared by every pytest invocation. +def PYTEST_IGNORE = '--ignore tests/unit_test --ignore tests/finetune --ignore tests/vllm --ignore tests/nightly_pipeline --ignore tests/transformers/sampler' + +// --------------------------------------------------------------------------- +// Pipeline +// --------------------------------------------------------------------------- pipeline { agent { node { label params.NODE_LABEL } } - options { disableConcurrentBuilds() } + options { + disableConcurrentBuilds() + } parameters { - string( - name: 'NODE_LABEL', - defaultValue: 'qeff_node', - description: 'Jenkins agent/node label to run this pipeline on' - ) - choice( - name: 'TEST_PROFILE', - choices: [ - 'dummy_layers_model', - 'few_layers_model', - 'full_layers_model' - ], - description: 'Select test profile' - ) - string( - name: 'SELECT_TEST_STAGES', - defaultValue: 'ALL', - description: 'Select which test stages you want to run (all run by default)' - ) - booleanParam(name: 'RUN_HL_APIS', defaultValue: true) - booleanParam(name: 'RUN_QAIC_FEATURE', defaultValue: true) - booleanParam(name: 'RUN_QAIC_MM', defaultValue: true) - booleanParam(name: 'RUN_QAIC_DIFFUSION', defaultValue: true) - booleanParam(name: 'RUN_CLI', defaultValue: true) - booleanParam(name: 'RUN_FINETUNE', defaultValue: false) + string(name: 'NODE_LABEL', defaultValue: 'qeff_node', + description: 'Jenkins agent/node label to run this pipeline on') + + choice(name: 'TEST_PROFILE', + choices: ['few_layers', 'tiny_model', 'full_model'], + description: 'CI test profile: tiny_model (fast, random-weight tinies) or full_model (real weights, nightly)') + + booleanParam(name: 'RUN_NON_QAIC_LLM', defaultValue: true, + description: 'Run Non-QAIC: LLM stage (Phase 1)') + + booleanParam(name: 'RUN_QAIC_LLM', defaultValue: true, + description: 'Run QAIC: LLM stage (Phase 2)') + + // Number of xdist workers / QAIC cards on the host. + string(name: 'QAIC_WORKERS', defaultValue: '4', + description: 'pytest-xdist worker count for QAIC stages (one card per worker)') } environment { - TEST_FILTER = testFilter(params.TEST_PROFILE) + QAIC_N = "${params.QAIC_WORKERS}" + QEFF_TEST_PROFILE = "${params.TEST_PROFILE}" + // Persistent QPC cache shared across CI runs to avoid cold-cache recompile. + QEFF_QPC_CACHE = "/home/ubuntu/.cache/qeff_qpcs" + DOCKER_QPC_CACHE = "/root/.cache/qeff_qpcs" } stages { + + // ── Install ────────────────────────────────────────────────────────── stage('Install QEfficient') { steps { sh ''' . ~/.bashrc - sudo docker run --privileged -dit --name ${BUILD_TAG} -e HF_TOKEN=${HF_TOKEN} -v ./:/efficient-transformers -v ${HF_PATH}:${DOCKER_HF_PATH} ${DOCKER_LATEST}:master_latest + mkdir -p ${QEFF_QPC_CACHE} + sudo docker run --privileged -dit --name ${BUILD_TAG} \ + -e HF_TOKEN=${HF_TOKEN} \ + -v ./:/efficient-transformers \ + -v ${HF_PATH}:${DOCKER_HF_PATH} \ + -v ${QEFF_QPC_CACHE}:${DOCKER_QPC_CACHE} \ + ${DOCKER_LATEST}:master_latest sudo docker exec ${BUILD_TAG} bash -c " cd /efficient-transformers && apt update && @@ -72,174 +75,58 @@ pipeline { } } - stage('HL API Tests') { - when { expression { params.RUN_HL_APIS } } - parallel { - stage('Export & Compile') { - steps { - timeout(time: params.TEST_PROFILE == 'full_layers_model' ? 0 : 60, unit: params.TEST_PROFILE == 'full_layers_model' ? 'HOURS' : 'MINUTES') { - sh ''' - sudo docker exec ${BUILD_TAG} bash -c " - cd /efficient-transformers && - . preflight_qeff/bin/activate && - mkdir -p $PWD/Non_cli_qaic && - export TOKENIZERS_PARALLELISM=false && - export QEFF_HOME=$PWD/Non_cli_qaic && - pytest tests -m '(not on_qaic) and (not finetune) and ${TEST_FILTER}' --ignore tests/vllm --ignore tests/unit_test --ignore tests/nightly_pipeline -n 4 --junitxml=tests/tests_log1.xml --durations=10 && - junitparser merge tests/tests_log1.xml tests/tests_log.xml && - deactivate" - ''' - } - } - } + // stage('Non-QAIC: LLM') { + // steps { + // catchError(buildResult: 'FAILURE', stageResult: 'FAILURE') { + // sh """ + // sudo docker exec ${BUILD_TAG} bash -c " + // cd /efficient-transformers && + // . preflight_qeff/bin/activate && + // export QEFF_TEST_PROFILE=${QEFF_TEST_PROFILE} && + // export TOKENIZERS_PARALLELISM=false && + // export QEFF_HOME=${DOCKER_QPC_CACHE}/llm && + // set -o pipefail && + // pytest tests/transformers/models/causal_lm_models -m 'llm and non_qaic' \ + // -n auto --dist worksteal \ + // --junitxml=tests/tests_log_non_qaic_llm.xml \ + // --log-level=ERROR --no-header && + // junitparser merge tests/tests_log_non_qaic_llm.xml tests/tests_log.xml && + // deactivate" + // """ + // } + // } + // } - stage('QAIC LLM') { - steps { - timeout(time: params.TEST_PROFILE == 'full_layers_model' ? 0 : 240, unit: params.TEST_PROFILE == 'full_layers_model' ? 'HOURS' : 'MINUTES') { - sh ''' - sudo docker exec ${BUILD_TAG} bash -c " - cd /efficient-transformers && - . preflight_qeff/bin/activate && - mkdir -p $PWD/Non_qaic_llm && - export TOKENIZERS_PARALLELISM=false && - export QEFF_HOME=$PWD/Non_qaic_llm && - pytest tests -m '(llm_model) and (not qnn) and ${TEST_FILTER}' --ignore tests/vllm --ignore tests/unit_test --ignore tests/nightly_pipeline --junitxml=tests/tests_log2.xml --durations=10 && - junitparser merge tests/tests_log2.xml tests/tests_log.xml && - deactivate" - ''' - } - } - } - - } - } - stage('QAIC FEATURE') { - when {expression { params.RUN_QAIC_FEATURE }} + stage('QAIC: LLM') { + when { expression { params.RUN_QAIC_LLM } } steps { - timeout(time: params.TEST_PROFILE == 'full_layers_model' ? 0 : 120, unit: params.TEST_PROFILE == 'full_layers_model' ? 'HOURS' : 'MINUTES') { - sh ''' - sudo docker exec ${BUILD_TAG} bash -c " - cd /efficient-transformers && - . preflight_qeff/bin/activate && - mkdir -p $PWD/Non_qaic_feature && - export TOKENIZERS_PARALLELISM=false && - export QEFF_HOME=$PWD/Non_qaic_feature && - pytest tests -m '(on_qaic) and (feature) and (not qnn) and ${TEST_FILTER}' --ignore tests/transformers/sampler --ignore tests/vllm --ignore tests/unit_test --ignore tests/nightly_pipeline --junitxml=tests/tests_log2_feature.xml --durations=10 && - junitparser merge tests/tests_log2_feature.xml tests/tests_log.xml && - deactivate" - ''' - } - } - } - stage('QAIC Multimodal') { - when {expression { params.RUN_QAIC_MM }} - steps { - timeout(time: params.TEST_PROFILE == 'full_layers_model' ? 0 : 180, unit: params.TEST_PROFILE == 'full_layers_model' ? 'HOURS' : 'MINUTES') { - sh ''' - sudo docker exec ${BUILD_TAG} bash -c " - cd /efficient-transformers && - . preflight_qeff/bin/activate && - mkdir -p $PWD/Non_cli_qaic_multimodal && - export TOKENIZERS_PARALLELISM=false && - export QEFF_HOME=$PWD/Non_cli_qaic_multimodal && - pytest tests -m '(multimodal) and (not qnn) and ${TEST_FILTER}' --ignore tests/vllm --ignore tests/unit_test --ignore tests/nightly_pipeline --ignore tests/transformers/models/reranker/test_reranker_mad.py --junitxml=tests/tests_log6.xml --durations=10 && - junitparser merge tests/tests_log6.xml tests/tests_log.xml && - deactivate" - ''' - } - } - } - stage('QAIC Reranker Tests') { - when { expression { params.RUN_QAIC_MM } } - steps { - timeout(time: 20, unit: 'MINUTES') { - sh ''' - sudo docker exec ${BUILD_TAG} bash -c " - cd /efficient-transformers && - . preflight_qeff/bin/activate && - mkdir -p $PWD/Non_cli_qaic_reranker && - export TOKENIZERS_PARALLELISM=false && - export QEFF_HOME=$PWD/Non_cli_qaic_reranker && - export QEFF_RERANKER_DOC_LIMIT=1 && - pytest -q tests/transformers/models/reranker/test_reranker_mad.py --maxfail=1 --junitxml=tests/tests_log_reranker.xml --durations=10 && - junitparser merge tests/tests_log_reranker.xml tests/tests_log.xml && - deactivate" - ''' - } - } - } - stage('QAIC Diffusion Models Tests') { - when { expression { params.RUN_QAIC_DIFFUSION } } - steps { - timeout(time: 120, unit: 'MINUTES') { - sh ''' - sudo docker exec ${BUILD_TAG} bash -c " - cd /efficient-transformers && - . preflight_qeff/bin/activate && - mkdir -p $PWD/Non_cli_qaic_diffusion && - export TOKENIZERS_PARALLELISM=false && - export QEFF_HOME=$PWD/Non_cli_qaic_diffusion && - export HF_HUB_CACHE=/huggingface_hub && - pytest tests -m 'diffusion_models' --ignore tests/vllm --ignore tests/unit_test --ignore tests/nightly_pipeline --junitxml=tests/tests_log_diffusion.xml --durations=10 && - junitparser merge tests/tests_log_diffusion.xml tests/tests_log.xml && - deactivate" - ''' - } - } - } - - stage('CLI Tests') { - when { expression { params.RUN_CLI } } - steps { - timeout(time: 120, unit: 'MINUTES') { - sh ''' - sudo docker exec ${BUILD_TAG} bash -c " - #source /qnn_sdk/bin/envsetup.sh && - #source /qnn_sdk/bin/envcheck -c && - cd /efficient-transformers && - . preflight_qeff/bin/activate && - mkdir -p $PWD/cli && - export TOKENIZERS_PARALLELISM=false && - export QEFF_HOME=$PWD/cli && - pytest tests -m '(cli and not qnn) and (not finetune)' --ignore tests/vllm --ignore tests/unit_test --ignore tests/nightly_pipeline --junitxml=tests/tests_log3.xml --durations=10 && - junitparser merge tests/tests_log3.xml tests/tests_log.xml && - deactivate" - ''' - } + catchError(buildResult: 'FAILURE', stageResult: 'FAILURE') { + sh """ + sudo docker exec ${BUILD_TAG} bash -c " + cd /efficient-transformers && + . preflight_qeff/bin/activate && + export QEFF_TEST_PROFILE=${QEFF_TEST_PROFILE} && + export TOKENIZERS_PARALLELISM=false && + export QEFF_HOME=${DOCKER_QPC_CACHE}/llm && + set -o pipefail && + pytest tests/transformers/models/causal_lm_models -m 'qaic and llm' \ + --junitxml=tests/tests_log_qaic_llm.xml \ + --durations=20 \ + --log-level=ERROR --no-header && + junitparser merge tests/tests_log_qaic_llm.xml tests/tests_log.xml && + deactivate" + """ + } } } - stage('Finetune Tests') { - when { expression { params.RUN_FINETUNE } } - steps { - timeout(time: 20, unit: 'MINUTES') { - sh ''' - sudo docker exec ${BUILD_TAG} bash -c " - cd /efficient-transformers && - . preflight_qeff/bin/activate && - # TODO: Update torch_qaic path to py312 when migrating to Python 3.12 - pip install /opt/qti-aic/integrations/torch_qaic/py312/torch_qaic-0.1.0-cp312-cp312-manylinux_2_34_x86_64.whl && - # pip install /opt/qti-aic/integrations/torch_qaic/py310/torch_qaic-0.1.0-cp310-cp310-linux_x86_64.whl && - pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cpu && - mkdir -p $PWD/cli_qaic_finetuning && - export TOKENIZERS_PARALLELISM=false && - export QEFF_HOME=$PWD/cli_qaic_finetuning && - pytest tests -m '(finetune)' --ignore tests/vllm --ignore tests/unit_test --ignore tests/nightly_pipeline --junitxml=tests/tests_log_finetune.xml --durations=10 && - junitparser merge tests/tests_log_finetune.xml tests/tests_log.xml && - deactivate" - ''' - } - } - } - } + } // end stages post { always { script { try { - sh ''' - sudo chown -R ubuntu . - ''' + sh 'sudo chown -R ubuntu .' } catch (error) { echo "Failed to change ownership: ${error}" } @@ -253,9 +140,7 @@ pipeline { } script { try { - sh ''' - sudo docker rm -f ${BUILD_TAG} - ''' + sh 'sudo docker rm -f ${BUILD_TAG}' } catch (error) { echo "Failed to delete container ${BUILD_TAG}: ${error}" } @@ -264,4 +149,4 @@ pipeline { deleteDir() } } -} +} \ No newline at end of file diff --git a/tests/configs/causal_model_configs.json b/tests/configs/causal_model_configs.json index 9854964702..ffee6c6271 100644 --- a/tests/configs/causal_model_configs.json +++ b/tests/configs/causal_model_configs.json @@ -1,723 +1,723 @@ { - "causal_lm_models": [ - { - "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "gpt2", - "model_type": "gpt2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 50257, - "num_key_value_heads": 1 - } - }, - { - "model_name": "allenai/OLMo-2-0425-1B", - "model_type": "olmo2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 100352, - "num_key_value_heads": 1 - } - }, - { - "model_name": "Salesforce/codegen-350M-mono", - "model_type": "codegen", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 51200, - "num_key_value_heads": 1, - "rotary_dim": 16 - } - }, - { - "model_name": "ibm-granite/granite-3.1-1b-a400m-base", - "model_type": "granitemoe", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49155, - "num_key_value_heads": 1 - } - }, - { - "model_name": "microsoft/Phi-3-mini-4k-instruct", - "model_type": "phi3", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32064, - "num_key_value_heads": 1 - } - }, - { - "model_name": "tiiuae/falcon-7b", - "model_type": "falcon", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 65024, - "num_key_value_heads": 1 - } - }, - { - "model_name": "Qwen/Qwen3-30B-A3B-Instruct-2507", - "model_type": "qwen3_moe", - "additional_params": { - "hidden_size": 256, - "intermediate_size": 256, - "max_position_embeddings": 128, - "max_window_layers": 48, - "moe_intermediate_size": 768, - "num_attention_heads": 2, - "num_experts": 4, - "num_experts_per_tok": 2, - "num_hidden_layers": 1, - "num_key_value_heads": 1, - "vocab_size": 151936 - } - }, - { - "model_name": "Qwen/Qwen2-0.5B", - "model_type": "qwen2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 151936, - "num_key_value_heads": 1 - } - }, - { - "model_name": "bigcode/starcoder2-3b", - "model_type": "starcoder2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49152, - "num_key_value_heads": 1 - } - }, - { - "model_name": "Felladrin/Minueza-32M-Base", - "model_type": "mistral", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32002, - "num_key_value_heads": 1 - } - }, - { - "model_name": "wtang06/mpt-125m-c4", - "model_type": "mpt", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 50368 - } - }, - { - "model_name": "hakurei/gpt-j-random-tinier", - "model_type": "gptj", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 50400, - "num_key_value_heads": 1, - "rotary_dim": 16 - } - }, - { - "model_name": "mistralai/Mixtral-8x7B-Instruct-v0.1", - "model_type": "mixtral", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "meta-llama/Llama-3.2-1B", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 128256, - "num_key_value_heads": 1, - "rope_scaling": { - "factor": 32.0, - "high_freq_factor": 4.0, - "low_freq_factor": 1.0, - "original_max_position_embeddings": 8192, - "rope_type": "llama3", - "rope_theta": 500000.0 + "causal_lm_models": [ + { + "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "gpt2", + "model_type": "gpt2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 50257, + "num_key_value_heads": 1 + } + }, + { + "model_name": "allenai/OLMo-2-0425-1B", + "model_type": "olmo2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 100352, + "num_key_value_heads": 1 + } + }, + { + "model_name": "Salesforce/codegen-350M-mono", + "model_type": "codegen", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 51200, + "num_key_value_heads": 1, + "rotary_dim": 16 + } + }, + { + "model_name": "ibm-granite/granite-3.1-1b-a400m-base", + "model_type": "granitemoe", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49155, + "num_key_value_heads": 1 + } + }, + { + "model_name": "microsoft/Phi-3-mini-4k-instruct", + "model_type": "phi3", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32064, + "num_key_value_heads": 1 + } + }, + { + "model_name": "tiiuae/falcon-7b", + "model_type": "falcon", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 65024, + "num_key_value_heads": 1 + } + }, + { + "model_name": "Qwen/Qwen3-30B-A3B-Instruct-2507", + "model_type": "qwen3_moe", + "additional_params": { + "hidden_size": 256, + "intermediate_size": 256, + "max_position_embeddings": 128, + "max_window_layers": 48, + "moe_intermediate_size": 768, + "num_attention_heads": 2, + "num_experts": 4, + "num_experts_per_tok": 2, + "num_hidden_layers": 1, + "num_key_value_heads": 1, + "vocab_size": 151936 + } + }, + { + "model_name": "Qwen/Qwen2-0.5B", + "model_type": "qwen2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 151936, + "num_key_value_heads": 1 + } + }, + { + "model_name": "bigcode/starcoder2-3b", + "model_type": "starcoder2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49152, + "num_key_value_heads": 1 + } + }, + { + "model_name": "Felladrin/Minueza-32M-Base", + "model_type": "mistral", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32002, + "num_key_value_heads": 1 + } + }, + { + "model_name": "wtang06/mpt-125m-c4", + "model_type": "mpt", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 50368 + } + }, + { + "model_name": "hakurei/gpt-j-random-tinier", + "model_type": "gptj", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 50400, + "num_key_value_heads": 1, + "rotary_dim": 16 + } + }, + { + "model_name": "mistralai/Mixtral-8x7B-Instruct-v0.1", + "model_type": "mixtral", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "meta-llama/Llama-3.2-1B", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 128256, + "num_key_value_heads": 1, + "rope_scaling": { + "factor": 32.0, + "high_freq_factor": 4.0, + "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + "rope_theta": 500000.0 + } + } + }, + { + "model_name": "unsloth/gemma-2b", + "model_type": "gemma", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 256000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "unsloth/gemma-2-2b", + "model_type": "gemma2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 256000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32003 + } + }, + { + "model_name": "TheBloke/Llama-2-7B-GPTQ", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32000 + } + }, + { + "model_name": "ibm-granite/granite-20b-code-base", + "model_type": "gpt_bigcode", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49152, + "num_key_value_heads": 1, + "activation_function": "gelu", + "architectures": [ + "GPTBigCodeForCausalLM" + ] + } + }, + { + "model_name": "neuralmagic/Llama-3.2-3B-Instruct-FP8", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 128256 + } + }, + { + "model_name": "neuralmagic/Qwen2-0.5B-Instruct-FP8", + "model_type": "qwen2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 151936 + } + }, + { + "model_name": "ibm-granite/granite-3.1-2b-instruct", + "model_type": "granite", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49155, + "num_key_value_heads": 1 + } + }, + { + "model_name": "hpcai-tech/grok-1", + "model_type": null, + "additional_params": { + "max_position_embeddings": 64, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 131072, + "num_key_value_heads": 1 + } + }, + { + "model_name": "Snowflake/Llama-3.1-SwiftKV-8B-Instruct", + "model_type": null, + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "hidden_size": 256, + "intermediate_size": 256, + "vocab_size": 128256, + "num_key_value_layers": 1, + "num_key_value_heads": 1, + "rope_scaling": { + "factor": 8.0, + "high_freq_factor": 4.0, + "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3" + } + } } - } - }, - { - "model_name": "unsloth/gemma-2b", - "model_type": "gemma", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 256000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "unsloth/gemma-2-2b", - "model_type": "gemma2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 256000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32003 - } - }, - { - "model_name": "TheBloke/Llama-2-7B-GPTQ", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32000 - } - }, - { - "model_name": "ibm-granite/granite-20b-code-base", - "model_type": "gpt_bigcode", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49152, - "num_key_value_heads": 1, - "activation_function": "gelu", - "architectures": [ - "GPTBigCodeForCausalLM" - ] - } - }, - { - "model_name": "neuralmagic/Llama-3.2-3B-Instruct-FP8", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 128256 - } - }, - { - "model_name": "neuralmagic/Qwen2-0.5B-Instruct-FP8", - "model_type": "qwen2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 2, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 151936 - } - }, - { - "model_name": "ibm-granite/granite-3.1-2b-instruct", - "model_type": "granite", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49155, - "num_key_value_heads": 1 - } - }, - { - "model_name": "hpcai-tech/grok-1", - "model_type": null, - "additional_params": { - "max_position_embeddings": 64, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 131072, - "num_key_value_heads": 1 - } - }, - { - "model_name": "Snowflake/Llama-3.1-SwiftKV-8B-Instruct", - "model_type": null, - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 2, - "num_attention_heads": 2, - "hidden_size": 256, - "intermediate_size": 256, - "vocab_size": 128256, - "num_key_value_layers": 1, - "num_key_value_heads": 1, - "rope_scaling": { - "factor": 8.0, - "high_freq_factor": 4.0, - "low_freq_factor": 1.0, - "original_max_position_embeddings": 8192, - "rope_type": "llama3" + ], + "causal_lm_fp16_test_models": [ + { + "model_name": "gpt2", + "model_type": "gpt2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 50257, + "num_key_value_heads": 1 + } + }, + { + "model_name": "hf-internal-testing/tiny-random-Gemma2ForCausalLM", + "model_type": "gemma2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 256000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "hf-internal-testing/tiny-random-GPTBigCodeForCausalLM", + "model_type": "gpt_bigcode", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49152, + "num_key_value_heads": 1, + "activation_function": "gelu", + "architectures": [ + "GPTBigCodeForCausalLM" + ] + } + }, + { + "model_name": "Qwen/Qwen2-0.5B", + "model_type": "qwen2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 151936, + "num_key_value_heads": 1 + } + }, + { + "model_name": "hf-internal-testing/tiny-random-MixtralForCausalLM", + "model_type": "mixtral", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 128256, + "num_key_value_heads": 1, + "rope_scaling": { + "factor": 32.0, + "high_freq_factor": 4.0, + "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + "rope_theta": 500000.0 + } + } + }, + { + "model_name": "allenai/OLMo-2-0425-1B", + "model_type": "olmo2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 100352, + "num_key_value_heads": 1 + } + }, + { + "model_name": "hf-internal-testing/tiny-random-GraniteForCausalLM", + "model_type": "granite", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49155, + "num_key_value_heads": 1 + } } - } - } - ], - "causal_lm_fp16_test_models": [ - { - "model_name": "gpt2", - "model_type": "gpt2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 50257, - "num_key_value_heads": 1 - } - }, - { - "model_name": "hf-internal-testing/tiny-random-Gemma2ForCausalLM", - "model_type": "gemma2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 256000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "hf-internal-testing/tiny-random-GPTBigCodeForCausalLM", - "model_type": "gpt_bigcode", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49152, - "num_key_value_heads": 1, - "activation_function": "gelu", - "architectures": [ - "GPTBigCodeForCausalLM" - ] - } - }, - { - "model_name": "Qwen/Qwen2-0.5B", - "model_type": "qwen2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 151936, - "num_key_value_heads": 1 - } - }, - { - "model_name": "hf-internal-testing/tiny-random-MixtralForCausalLM", - "model_type": "mixtral", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "hf-internal-testing/tiny-random-LlamaForCausalLM", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 128256, - "num_key_value_heads": 1, - "rope_scaling": { - "factor": 32.0, - "high_freq_factor": 4.0, - "low_freq_factor": 1.0, - "original_max_position_embeddings": 8192, - "rope_type": "llama3", - "rope_theta": 500000.0 + ], + "spd_causal_lm_models": [ + { + "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "Qwen/Qwen2-0.5B", + "model_type": "qwen2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 151936, + "num_key_value_heads": 1 + } } - } - }, - { - "model_name": "allenai/OLMo-2-0425-1B", - "model_type": "olmo2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 100352, - "num_key_value_heads": 1 - } - }, - { - "model_name": "hf-internal-testing/tiny-random-GraniteForCausalLM", - "model_type": "granite", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49155, - "num_key_value_heads": 1 - } - } - ], - "spd_causal_lm_models": [ - { - "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "Qwen/Qwen2-0.5B", - "model_type": "qwen2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 151936, - "num_key_value_heads": 1 - } - } - ], - "qnn_causal_lm_models": [ - { - "model_name": "mistralai/Mixtral-8x7B-Instruct-v0.1", - "model_type": "mixtral", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 32000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "meta-llama/Llama-3.2-1B", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 128256, - "num_key_value_heads": 1, - "rope_scaling": { - "factor": 32.0, - "high_freq_factor": 4.0, - "low_freq_factor": 1.0, - "original_max_position_embeddings": 8192, - "rope_type": "llama3", - "rope_theta": 500000.0 + ], + "qnn_causal_lm_models": [ + { + "model_name": "mistralai/Mixtral-8x7B-Instruct-v0.1", + "model_type": "mixtral", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 32000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "meta-llama/Llama-3.2-1B", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 128256, + "num_key_value_heads": 1, + "rope_scaling": { + "factor": 32.0, + "high_freq_factor": 4.0, + "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + "rope_theta": 500000.0 + } + } + }, + { + "model_name": "unsloth/gemma-2b", + "model_type": "gemma", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 256000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "ibm-granite/granite-guardian-3.1-2b", + "model_type": "granite", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49155, + "num_key_value_heads": 1 + } } - } - }, - { - "model_name": "unsloth/gemma-2b", - "model_type": "gemma", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 256000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "ibm-granite/granite-guardian-3.1-2b", - "model_type": "granite", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49155, - "num_key_value_heads": 1 - } - } - ], - "prefix_caching_models": [ - { - "model_name": "gpt2", - "model_type": "gpt2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 50257, - "num_key_value_heads": 1 - } - } - ], - "blockedKV_causal_lm_models": [ - { - "model_name": "unsloth/gemma-2b", - "model_type": "gemma", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 256000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "unsloth/gemma-2-2b", - "model_type": "gemma2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 256000, - "num_key_value_heads": 1 - } - }, - { - "model_name": "ibm-granite/granite-3.1-1b-a400m-base", - "model_type": "granitemoe", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49155, - "num_key_value_heads": 1 - } - }, - { - "model_name": "meta-llama/Llama-3.2-1B", - "model_type": "llama", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 128256, - "num_key_value_heads": 1, - "rope_scaling": { - "factor": 32.0, - "high_freq_factor": 4.0, - "low_freq_factor": 1.0, - "original_max_position_embeddings": 8192, - "rope_type": "llama3", - "rope_theta": 500000.0 + ], + "prefix_caching_models": [ + { + "model_name": "gpt2", + "model_type": "gpt2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 50257, + "num_key_value_heads": 1 + } } - } - }, - { - "model_name": "wtang06/mpt-125m-c4", - "model_type": "mpt", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 50368 - } - }, - { - "model_name": "bigcode/starcoder2-3b", - "model_type": "starcoder2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 49152, - "num_key_value_heads": 1 - } - } - ], - "disaggregated_dummy_models": [ - { - "model_name": "openai/gpt-oss-20b", - "model_type": "gpt_oss", - "tokenizer_id": "gpt2", - "additional_params": { - "num_hidden_layers": 2, - "hidden_size": 64, - "intermediate_size": 256, - "num_attention_heads": 2, - "num_key_value_heads": 1, - "num_local_experts": 4, - "head_dim": 32, - "max_position_embeddings": 512, - "vocab_size": 201088, - "sliding_window": 128 - } - }, - { - "model_name": "Qwen/Qwen3-30B-A3B-Instruct-2507", - "model_type": "qwen3_moe", - "additional_params": { - "hidden_size": 256, - "intermediate_size": 256, - "max_position_embeddings": 512, - "max_window_layers": 48, - "moe_intermediate_size": 768, - "num_attention_heads": 2, - "num_experts": 4, - "num_experts_per_tok": 2, - "num_hidden_layers": 2, - "num_key_value_heads": 1, - "vocab_size": 151936 - } - } - ], - "causal_lm_models_pl1": [ - { - "model_name": "gpt2", - "model_type": "gpt2", - "additional_params": { - "max_position_embeddings": 128, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "hidden_size": 64, - "intermediate_size": 256, - "vocab_size": 50257, - "num_key_value_heads": 1 - } - }, - { - "model_name": "openai/gpt-oss-20b", - "model_type": "gpt_oss", - "additional_params": { - "num_hidden_layers": 2, - "hidden_size": 64, - "intermediate_size": 256, - "num_attention_heads": 2, - "num_key_value_heads": 1, - "num_local_experts": 4 - } - } - ] + ], + "blockedKV_causal_lm_models": [ + { + "model_name": "unsloth/gemma-2b", + "model_type": "gemma", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 256000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "unsloth/gemma-2-2b", + "model_type": "gemma2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 256000, + "num_key_value_heads": 1 + } + }, + { + "model_name": "ibm-granite/granite-3.1-1b-a400m-base", + "model_type": "granitemoe", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49155, + "num_key_value_heads": 1 + } + }, + { + "model_name": "meta-llama/Llama-3.2-1B", + "model_type": "llama", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 128256, + "num_key_value_heads": 1, + "rope_scaling": { + "factor": 32.0, + "high_freq_factor": 4.0, + "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + "rope_theta": 500000.0 + } + } + }, + { + "model_name": "wtang06/mpt-125m-c4", + "model_type": "mpt", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 50368 + } + }, + { + "model_name": "bigcode/starcoder2-3b", + "model_type": "starcoder2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 49152, + "num_key_value_heads": 1 + } + } + ], + "disaggregated_dummy_models": [ + { + "model_name": "openai/gpt-oss-20b", + "model_type": "gpt_oss", + "tokenizer_id": "gpt2", + "additional_params": { + "num_hidden_layers": 2, + "hidden_size": 64, + "intermediate_size": 256, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "num_local_experts": 4, + "head_dim": 32, + "max_position_embeddings": 512, + "vocab_size": 201088, + "sliding_window": 128 + } + }, + { + "model_name": "Qwen/Qwen3-30B-A3B-Instruct-2507", + "model_type": "qwen3_moe", + "additional_params": { + "hidden_size": 256, + "intermediate_size": 256, + "max_position_embeddings": 512, + "max_window_layers": 48, + "moe_intermediate_size": 768, + "num_attention_heads": 2, + "num_experts": 4, + "num_experts_per_tok": 2, + "num_hidden_layers": 2, + "num_key_value_heads": 1, + "vocab_size": 151936 + } + } + ], + "causal_lm_models_pl1": [ + { + "model_name": "gpt2", + "model_type": "gpt2", + "additional_params": { + "max_position_embeddings": 128, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 64, + "intermediate_size": 256, + "vocab_size": 50257, + "num_key_value_heads": 1 + } + }, + { + "model_name": "openai/gpt-oss-20b", + "model_type": "gpt_oss", + "additional_params": { + "num_hidden_layers": 2, + "hidden_size": 64, + "intermediate_size": 256, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "num_local_experts": 4 + } + } + ] } \ No newline at end of file diff --git a/tests/transformers/caching/__init__.py b/tests/transformers/caching/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/tests/transformers/caching/test_prefix_caching.py b/tests/transformers/caching/test_prefix_caching.py deleted file mode 100644 index 40caa33f8f..0000000000 --- a/tests/transformers/caching/test_prefix_caching.py +++ /dev/null @@ -1,254 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import json -import os - -import numpy as np -import pytest -from transformers import AutoTokenizer - -from QEfficient.generation.text_generation_inference import TextGeneration -from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM -from QEfficient.utils._utils import create_json -from QEfficient.utils.constants import QnnConstants -from QEfficient.utils.test_utils import load_qeff_causal_lm_model - -CONFIG_PATH = os.path.join(os.path.dirname(__file__), "../../configs/causal_model_configs.json") -with open(CONFIG_PATH, "r") as f: - config_data = json.load(f) - prefix_caching_models = config_data["prefix_caching_models"] - -test_models = [model["model_name"] for model in prefix_caching_models] -model_config_dict = {model["model_name"]: model for model in prefix_caching_models} - - -def prefix_caching_inference(model_name, qpc_path): - prefixes = ["Once upon a time ", "Once upon a time "] - suffixes1 = ["in a land far away", "there was a small village"] - suffixes2 = ["a little girl", "in a bustling city"] - - tokenizer = AutoTokenizer.from_pretrained(model_name) - - generator = TextGeneration(tokenizer=tokenizer, qpc_path=qpc_path, full_batch_size=2, ctx_len=256) - - prompts = [pref + suff for pref, suff in zip(prefixes, suffixes1)] - - # generation for batch_indices = 0, 1 - prompts_exec_info = generator.generate(prompts) - ############################## - # generation for batch_indices - ############################## - # Run prefill for indices 2, 3 with same prompts - out2, pos2, gen_len2 = generator._qaic_model.run_prefill( - prompts[0], generation_len=None, decode_batch_id=np.array(2, dtype=np.int64).reshape(1, 1) - ) - out3, pos3, gen_len3 = generator._qaic_model.run_prefill( - prompts[1], generation_len=None, decode_batch_id=np.array(3, dtype=np.int64).reshape(1, 1) - ) - - # Run decode for batch indices 2, 3 - decode_inputs = { - "input_ids": np.array([[out2["logits"].argmax(2)[0][0]], [out3["logits"].argmax(2)[0][0]]]), - "position_ids": np.array([[pos2[0][0]], [pos3[0][0]]]), - "batch_index": np.array([[2], [3]], dtype=np.int64), - } - - # Set logits placeholder for decode - logits_out_placeholder = np.zeros( - ( - generator._qaic_model.full_batch_size, - generator._qaic_model._decode_seq_len, - generator._qaic_model._vocab_size, - ), - dtype=np.float32, - ) - generator._qaic_model._session.set_buffers({"logits": logits_out_placeholder}) - - generation_outputs = [] - for i in range(gen_len2): - generation_outputs.append(decode_inputs["input_ids"]) - outputs = generator._qaic_model._session.run(decode_inputs) - logits = outputs["logits"] - if len(logits.shape) == 2: - logits = np.expand_dims(logits, 1) - next_token_id = logits.argmax(2) - - decode_inputs["input_ids"] = next_token_id - decode_inputs["position_ids"] += 1 - - assert np.all(generator._qaic_model.generated_ids[0, :gen_len2] == [int(val[0, 0]) for val in generation_outputs]) - assert np.all(generator._qaic_model.generated_ids[1, :gen_len2] == [int(val[1, 0]) for val in generation_outputs]) - - ############################## - # Now rerun with cached prefix on 0th index with prompt3 and use -1 for 1st index - ############################## - - nprompts = [pref + suff for pref, suff in zip(prefixes, suffixes2)] - - ## Prefill run on index 0 - prompt = nprompts[0] - inputs = tokenizer(prompt, return_tensors="np", padding=True) - position_ids = inputs["attention_mask"].sum(1, keepdims=True) - padded_len = inputs["input_ids"].shape[1] - num_chunks = -(padded_len // -generator._qaic_model._prefill_seq_len) - padded_len = num_chunks * generator._qaic_model._prefill_seq_len # Convert to a multiple of prompt_len - - # Initialize variables specific to request - # Calculate the max generation length. - max_gen_len = generator._qaic_model._ctx_len - position_ids.max() - - # Set the prefill logic buffer - logits_out_placeholder = np.zeros((1, 1, generator._qaic_model._vocab_size), dtype=np.float32) - generator._qaic_model._session.set_buffers({"logits": logits_out_placeholder}) - inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) - inputs["position_ids"] = np.where(inputs.pop("attention_mask"), np.arange(padded_len), -1) - inputs.pop("token_type_ids", None) - inputs["batch_index"] = np.array([[0]], dtype=np.int64) - norm_outputs = generator._qaic_model._session.run(inputs) - inputs["input_ids"][:, :3] = inputs["input_ids"][:, 4:7] - inputs["input_ids"][:, 3:] = 50256 - inputs["position_ids"][:, :3] = inputs["position_ids"][:, 4:7] - inputs["position_ids"][:, 3:] = -1 - mod_outputs = generator._qaic_model._session.run(inputs) - assert (mod_outputs["logits"] == norm_outputs["logits"]).all() - decode_inputs = { - "input_ids": np.array([[mod_outputs["logits"].argmax(2)[0][0]], [0]]), - "position_ids": np.array([[position_ids[0][0]], [-1]]), - "batch_index": np.array([[0], [1]], dtype=np.int64), - } - - # Set logits placeholder for decode - logits_out_placeholder = np.zeros( - ( - generator._qaic_model.full_batch_size, - generator._qaic_model._decode_seq_len, - generator._qaic_model._vocab_size, - ), - dtype=np.float32, - ) - generator._qaic_model._session.set_buffers({"logits": logits_out_placeholder}) - - generation_outputs = [] - for i in range(max_gen_len): - generation_outputs.append(decode_inputs["input_ids"]) - outputs = generator._qaic_model._session.run(decode_inputs) - logits = outputs["logits"] - if len(logits.shape) == 2: - logits = np.expand_dims(logits, 1) - next_token_id = logits.argmax(2) - - decode_inputs["input_ids"] = next_token_id - decode_inputs["position_ids"][0][0] += 1 - - # TODO: add a check if this matches normal execution for same prompt - ############## - # Now run decode on 1st index again with mod_inputs and check if output is correct - ############## - decode_inputs = { - "input_ids": np.array([[0], [prompts_exec_info.generated_ids[1][0]]]), - "position_ids": np.array([[-1], [9]]), - "batch_index": np.array([[0], [1]], dtype=np.int64), - } - - # Set logits placeholder for decode - logits_out_placeholder = np.zeros( - ( - generator._qaic_model.full_batch_size, - generator._qaic_model._decode_seq_len, - generator._qaic_model._vocab_size, - ), - dtype=np.float32, - ) - generator._qaic_model._session.set_buffers({"logits": logits_out_placeholder}) - - generation_outputs_prefill_cached = [] - for i in range(max_gen_len): - generation_outputs_prefill_cached.append(decode_inputs["input_ids"]) - outputs = generator._qaic_model._session.run(decode_inputs) - logits = outputs["logits"] - if len(logits.shape) == 2: - logits = np.expand_dims(logits, 1) - next_token_id = logits.argmax(2) - - decode_inputs["input_ids"] = next_token_id - decode_inputs["position_ids"][1][0] += 1 - - assert np.all( - prompts_exec_info.generated_ids[1][:247] == [int(val[1, 0]) for val in generation_outputs_prefill_cached][:247] - ) - - -@pytest.mark.full_layers -@pytest.mark.on_qaic -@pytest.mark.feature -@pytest.mark.parametrize("model_name", test_models) -def test_full_simple_prefix_caching(model_name, manual_cleanup): - """ - The test should first generate output with some prefix+suffix1 or batch_id and then confirm that we are still able to execute of prefix+suffix2 on same batch id and getting correct output. - """ - qeff_model = load_qeff_causal_lm_model(model_name=model_name, continuous_batching=True) - qeff_model.compile( - prefill_seq_len=128, - ctx_len=256, - full_batch_size=2, - kv_cache_batch_size=4, - num_cores=16, - ) - prefix_caching_inference(model_name=model_name, qpc_path=qeff_model.qpc_path) - assert os.path.isfile(os.path.join(os.path.dirname(qeff_model.qpc_path), "qconfig.json")) - manual_cleanup(qeff_model.onnx_path) - - -@pytest.mark.on_qaic -@pytest.mark.feature -@pytest.mark.parametrize("model_name", test_models) -def test_simple_prefix_caching(model_name, manual_cleanup): - """ - The test should first generate output with some prefix+suffix1 or batch_id and then confirm that we are still able to execute of prefix+suffix2 on same batch id and getting correct output. - """ - qeff_model = load_qeff_causal_lm_model( - model_name=model_name, - continuous_batching=True, - num_hidden_layers=1, - ) - qeff_model.compile( - prefill_seq_len=128, - ctx_len=256, - full_batch_size=2, - kv_cache_batch_size=4, - num_cores=16, - ) - prefix_caching_inference(model_name=model_name, qpc_path=qeff_model.qpc_path) - assert os.path.isfile(os.path.join(os.path.dirname(qeff_model.qpc_path), "qconfig.json")) - manual_cleanup(qeff_model.onnx_path) - - -################################# QNN Tests ################################# - - -@pytest.mark.on_qaic -@pytest.mark.feature -@pytest.mark.qnn -@pytest.mark.parametrize("model_name", test_models) -def test_simple_prefix_caching_qnn(model_name): - qeff_model = QEFFAutoModelForCausalLM.from_pretrained(model_name, continuous_batching=True) - qnn_config_json_path = os.path.join(os.getcwd(), "qnn_config.json") - create_json(qnn_config_json_path, QnnConstants.QNN_SAMPLE_CONFIG) - - qeff_model.compile( - prefill_seq_len=128, - ctx_len=256, - full_batch_size=2, - kv_cache_batch_size=4, - num_cores=14, - enable_qnn=True, - qnn_config=qnn_config_json_path, - ) - prefix_caching_inference(model_name=model_name, qpc_path=qeff_model.qpc_path) - assert os.path.isfile(os.path.join(os.path.dirname(qeff_model.qpc_path), "qconfig.json")) - os.remove(qnn_config_json_path) diff --git a/tests/transformers/disaggregated/test_disagg_mode.py b/tests/transformers/disaggregated/test_disagg_mode.py deleted file mode 100644 index 1844261040..0000000000 --- a/tests/transformers/disaggregated/test_disagg_mode.py +++ /dev/null @@ -1,554 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import json -import os -import time - -import numpy as np -import pytest -import torch -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer -from transformers.cache_utils import DynamicCache - -from QEfficient import QEFFAutoModelForCausalLM -from QEfficient.generation.cloud_infer import QAICInferenceSession -from QEfficient.transformers.quantizers import replace_transformers_quantizers, undo_transformers_quantizers - -# Dummy model configs — loaded from the shared config file. -_CONFIG_FILE = os.path.join(os.path.dirname(__file__), "..", "..", "configs", "causal_model_configs.json") -with open(_CONFIG_FILE) as _f: - _raw = json.load(_f) - -_DISAGG_DUMMY_CONFIGS = { - entry["model_name"]: { - "model_type": entry["model_type"], - "tokenizer_id": entry.get("tokenizer_id", entry["model_name"]), - **entry["additional_params"], - } - for entry in _raw["disaggregated_dummy_models"] -} - -# Test parameters: model IDs to test (loaded from config) -# - model_id_blocking: models that use blocking/sliding window attention -# - model_id_chunking: models that use chunking -model_id_blocking = [name for name, cfg in _DISAGG_DUMMY_CONFIGS.items() if cfg["model_type"] == "gpt_oss"] -model_id_chunking = [name for name, cfg in _DISAGG_DUMMY_CONFIGS.items() if cfg["model_type"] == "qwen3_moe"] - -prompt2 = """ -Once upon a time, in a small town, there lived a young boy named Alex. Alex was a curious and adventurous child, always eager to explore the world around him. One day, while playing in the park, Alex stumbled upon a mysterious old book hidden beneath a pile of leaves. The book was filled with stories of distant lands, magical creatures, and extraordinary adventures. - -As Alex flipped through the pages, he discovered a map that led to a hidden treasure. Excited by the prospect of a real-life treasure hunt, Alex decided to embark on a thrilling journey. He packed his backpack with snacks, a flashlight, and a compass, and set off into the unknown. - -The path to the treasure was not an easy one. Alex had to navigate through dense forests, cross rickety bridges, and solve riddles that guarded the treasure's location. -""" -prompt1 = "Once upon a time" -prompts = [prompt1, prompt2] - - -def _make_dummy_model(model_id: str) -> AutoModelForCausalLM: - """Create a tiny model from a dummy config — no weight download required. - - A fixed seed ensures the weights are reproducible across test runs so that - the QAIC-compiled model (which may be cached on disk) always matches the - in-process PyTorch model used for reference comparisons. - - Weights are scaled to std≈0.02 (matching real transformer init) so that - intermediate activations stay small and float16 rounding errors on QAIC - remain within the 5e-2 tolerance used for logit accuracy checks. - """ - cfg = _DISAGG_DUMMY_CONFIGS[model_id] - model_type = cfg["model_type"] - params = {k: v for k, v in cfg.items() if k not in ("model_type", "tokenizer_id")} - config = AutoConfig.for_model(model_type, **params) - torch.manual_seed(42) - model = AutoModelForCausalLM.from_config(config, attn_implementation="eager") - with torch.no_grad(): - for param in model.parameters(): - param.mul_(0.02) - return model - - -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_id", model_id_blocking) -@pytest.mark.parametrize("prompt", prompts) -def test_disagg_mode_prefill(model_id, prompt): - # Run prefill - tokenizer_id = _DISAGG_DUMMY_CONFIGS[model_id].get("tokenizer_id", model_id) - tokenizer = AutoTokenizer.from_pretrained(tokenizer_id) - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - PREFILL_SEQ_LEN = 256 - CTX_LEN = 256 - - # Tokenize once; reuse for both reference and qeff model - raw_inputs = tokenizer(prompt, return_tensors="np", padding=True) - padded_len = raw_inputs["input_ids"].shape[1] - num_chunks = -(padded_len // -PREFILL_SEQ_LEN) # ceil divide without float - padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len - - replace_transformers_quantizers() - model = _make_dummy_model(model_id) - config = model.config - - raw_inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) - raw_inputs["position_ids"] = np.where(raw_inputs.pop("attention_mask"), np.arange(padded_len), -1) - raw_inputs.pop("token_type_ids", None) - - inputs = {k: torch.from_numpy(v).to(model.device) for k, v in raw_inputs.items()} - cache = DynamicCache(config=config) - ins = tokenizer(prompt, return_tensors="pt") - out = model(**ins, past_key_values=cache) - - undo_transformers_quantizers() - - qeff_model = QEFFAutoModelForCausalLM(model) - qeff_model.prefill(True) - config = qeff_model.model.config - - inputs = {k: torch.from_numpy(v) for k, v in raw_inputs.items()} - past_key_values = [] - for i in range(config.num_hidden_layers): - cache_len = config.sliding_window if i % 2 == 0 else PREFILL_SEQ_LEN - pad_shape = (1, config.num_key_value_heads, cache_len, config.head_dim) - past_key = torch.zeros(pad_shape, dtype=torch.float32) - past_value = torch.zeros(pad_shape, dtype=torch.float32) - past_key_values.append((past_key, past_value)) - inputs["past_key_values"] = past_key_values - - qeff_out = qeff_model.model(**inputs) - - # Check our pytorch implementation - assert (qeff_out.logits - out.logits[:, -1, :]).abs().max() < 1e-4 - - prefill_qpc_path = qeff_model.compile( - prefill_seq_len=PREFILL_SEQ_LEN, - ctx_len=CTX_LEN, - num_cores=16, - mxfp6_matmul=False, - mxint8_kv_cache=False, - num_devices=1, - mos=1, - aic_enable_depth_first=True, - num_speculative_tokens=None, - prefill_only=True, - ) - - prefill_session = QAICInferenceSession(prefill_qpc_path) - logits_out_placeholder = np.zeros((1, 1, config.vocab_size), dtype=np.float32) - prefill_session.set_buffers({"logits": logits_out_placeholder}) - inputs.pop("past_key_values") - inputs = {k: v.detach().numpy() for k, v in inputs.items()} - st = time.time() - qpc_out = prefill_session.run(inputs) - print(f"time for prefill_run={time.time() - st} sec\n") - del prefill_session - assert (torch.from_numpy(qpc_out["logits"]) - qeff_out.logits).abs().max() < 5e-2 - - -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_id", model_id_chunking) -@pytest.mark.parametrize("prompt", prompts) -def test_disagg_mode_prefill_chunked(model_id, prompt): - # Run prefill - tokenizer = AutoTokenizer.from_pretrained(model_id) - PREFILL_SEQ_LEN = 128 - CTX_LEN = 128 * 3 - - # Tokenize once; reuse for both reference and qeff model - raw_inputs = tokenizer(prompt, return_tensors="np", padding=True) - padded_len = raw_inputs["input_ids"].shape[1] - num_chunks = -(padded_len // -PREFILL_SEQ_LEN) # ceil divide without float - padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len - - replace_transformers_quantizers() - model = _make_dummy_model(model_id) - config = model.config - - raw_inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) - raw_inputs["position_ids"] = np.where(raw_inputs.pop("attention_mask"), np.arange(padded_len), -1) - raw_inputs.pop("token_type_ids", None) - - inputs = {k: torch.from_numpy(v).to(model.device) for k, v in raw_inputs.items()} - cache = DynamicCache(config=config) - ins = tokenizer(prompt, return_tensors="pt") - out = model(**ins, past_key_values=cache) - - undo_transformers_quantizers() - - # Reuse the already-loaded model — avoids a second full model load - qeff_model = QEFFAutoModelForCausalLM(model) - qeff_model.prefill(True, enable_chunking=True) - config = qeff_model.model.config - - # head_dim is explicit in gpt_oss but computed for qwen3_moe - head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - - inputs = {k: torch.from_numpy(v) for k, v in raw_inputs.items()} - past_key_values = [] - for i in range(config.num_hidden_layers): - pad_shape = (1, config.num_key_value_heads, CTX_LEN, head_dim) - past_key = torch.zeros(pad_shape, dtype=torch.float32) - past_value = torch.zeros(pad_shape, dtype=torch.float32) - past_key_values.append((past_key, past_value)) - inputs["past_key_values"] = past_key_values - - for i in range(num_chunks): - chunk_inputs = inputs.copy() - chunk_inputs["input_ids"] = inputs["input_ids"][:, i * PREFILL_SEQ_LEN : (i + 1) * PREFILL_SEQ_LEN] - chunk_inputs["position_ids"] = inputs["position_ids"][:, i * PREFILL_SEQ_LEN : (i + 1) * PREFILL_SEQ_LEN] - - qeff_out = qeff_model.model(**chunk_inputs) - inputs["past_key_values"] = qeff_out["past_key_values"] - - # Check our pytorch implementation - assert (qeff_out.logits - out.logits[:, -1, :]).abs().max() < 1e-4 - - prefill_qpc_path = qeff_model.compile( - prefill_seq_len=PREFILL_SEQ_LEN, - ctx_len=CTX_LEN, - num_cores=config.num_experts, - mxfp6_matmul=False, - mxint8_kv_cache=False, - num_devices=1, - mos=1, - aic_enable_depth_first=True, - num_speculative_tokens=None, - prefill_only=True, - enable_chunking=True, - ) - prefill_session = QAICInferenceSession(prefill_qpc_path) - prefill_session.skip_buffers( - [x for x in prefill_session.input_names + prefill_session.output_names if x.startswith("past_")] - ) - logits_out_placeholder = np.zeros((1, 1, config.vocab_size), dtype=np.float32) - prefill_session.set_buffers({"logits": logits_out_placeholder}) - inputs.pop("past_key_values") - inputs = {k: v.detach().numpy() for k, v in inputs.items()} - st = time.time() - for i in range(num_chunks): - chunk_inputs = inputs.copy() - chunk_inputs["input_ids"] = inputs["input_ids"][:, i * PREFILL_SEQ_LEN : (i + 1) * PREFILL_SEQ_LEN] - chunk_inputs["position_ids"] = inputs["position_ids"][:, i * PREFILL_SEQ_LEN : (i + 1) * PREFILL_SEQ_LEN] - qpc_out = prefill_session.run(chunk_inputs) - print(f"time for prefill_run={time.time() - st} sec\n") - del prefill_session - assert (torch.from_numpy(qpc_out["logits"]) - qeff_out.logits).abs().max() < 5e-2 - - -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_id", model_id_blocking) -@pytest.mark.parametrize("prompt", [prompt1]) -def test_disagg_mode_prefill_only_and_decode_only(model_id, prompt): - # Run prefill for original pytorch model - tokenizer_id = _DISAGG_DUMMY_CONFIGS[model_id].get("tokenizer_id", model_id) - tokenizer = AutoTokenizer.from_pretrained(tokenizer_id) - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - PREFILL_SEQ_LEN = 256 - CTX_LEN = 256 - - raw_inputs = tokenizer(prompt, return_tensors="np", padding=True) - padded_len = raw_inputs["input_ids"].shape[1] - num_chunks = -(padded_len // -PREFILL_SEQ_LEN) # ceil divide without float - padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len - - replace_transformers_quantizers() - model = _make_dummy_model(model_id) - config = model.config - - raw_inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) - raw_inputs["position_ids"] = np.where(raw_inputs.pop("attention_mask"), np.arange(padded_len), -1) - raw_inputs.pop("token_type_ids", None) - - inputs = {k: torch.from_numpy(v).to(model.device) for k, v in raw_inputs.items()} - cache = DynamicCache(config=config) - ins = tokenizer(prompt, return_tensors="pt") - orig_out = model(**ins, past_key_values=cache) - - position_ids = inputs["position_ids"] - generated_ids = [] - generation_len = 10 - out = orig_out - for _ in range(1, generation_len): - next_token_id = out["logits"][:, -1, :].argmax(-1).reshape(-1, 1) - generated_ids.append(next_token_id) - position_ids = position_ids.max(1, keepdim=True).values + 1 - decode_inputs = { - "input_ids": next_token_id, - "position_ids": position_ids, - "past_key_values": out["past_key_values"], - } - out = model(**decode_inputs) - - generated_ids.append(out["logits"][:, -1, :].argmax(-1).reshape(-1, 1)) - generated_ids = np.concatenate(generated_ids, axis=1) - predicted_string = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) - print("Original HF Model Outputs (Torch CPU): \n") - print("Prompt:", repr(prompt)) - print("Completion:", repr(predicted_string)) - - undo_transformers_quantizers() - - prefill_qeff_model = QEFFAutoModelForCausalLM(model) - prefill_qeff_model.prefill(enable=True) - config = prefill_qeff_model.model.config - - past_key_values = [] - for i in range(config.num_hidden_layers): - cache_len = config.sliding_window if i % 2 == 0 else PREFILL_SEQ_LEN - pad_shape = (1, config.num_key_value_heads, cache_len, config.head_dim) - past_key = torch.zeros(pad_shape, dtype=torch.float32) - past_value = torch.zeros(pad_shape, dtype=torch.float32) - past_key_values.append((past_key, past_value)) - inputs["past_key_values"] = past_key_values - - prefill_qeff_out = prefill_qeff_model.model(**inputs) - - # Check our pytorch implementation - assert (prefill_qeff_out.logits - orig_out.logits[:, -1, :]).abs().max() < 1e-4 - - decode_qeff_model = QEFFAutoModelForCausalLM(model) - decode_qeff_model.prefill(enable=False) - qeff_out = prefill_qeff_out - - position_ids = inputs["position_ids"] - qeff_generated_ids = [] - for _ in range(1, generation_len): - next_token_id = qeff_out["logits"][:, -1, :].argmax(-1).reshape(-1, 1) - qeff_generated_ids.append(next_token_id) - position_ids = position_ids.max(1, keepdim=True).values + 1 - decode_inputs = { - "input_ids": next_token_id, - "position_ids": position_ids, - "past_key_values": qeff_out["past_key_values"], - } - qeff_out = decode_qeff_model.model(**decode_inputs) - - qeff_generated_ids.append(out["logits"][:, -1, :].argmax(-1).reshape(-1, 1)) - qeff_generated_ids = np.concatenate(qeff_generated_ids, axis=1) - predicted_string = tokenizer.batch_decode(qeff_generated_ids, skip_special_tokens=True) - print("QEFF Transformed Model Outputs (Torch CPU): \n") - print("Prompt:", repr(prompt)) - print("Completion:", repr(predicted_string)) - - assert (qeff_generated_ids == generated_ids).all() - - prefill_qpc_path = prefill_qeff_model.compile( - prefill_seq_len=PREFILL_SEQ_LEN, - ctx_len=CTX_LEN, - num_cores=16, - mxfp6_matmul=False, - mxint8_kv_cache=False, - num_devices=1, - mos=1, - aic_enable_depth_first=True, - num_speculative_tokens=None, - prefill_only=True, - ) - - prefill_session = QAICInferenceSession(prefill_qpc_path) - logits_out_placeholder = np.zeros((1, 1, config.vocab_size), dtype=np.float32) - prefill_session.set_buffers({"logits": logits_out_placeholder}) - inputs.pop("past_key_values") - inputs = {k: v.detach().numpy() for k, v in inputs.items()} - qpc_out = prefill_session.run(inputs) - del prefill_session - assert (torch.from_numpy(qpc_out["logits"]) - prefill_qeff_out.logits).abs().max() < 5e-2 - - decode_qpc_path = decode_qeff_model.compile( - prefill_seq_len=1, - ctx_len=CTX_LEN, - num_cores=16, - mxfp6_matmul=False, - mxint8_kv_cache=False, - num_devices=1, - mos=1, - aic_enable_depth_first=True, - num_speculative_tokens=None, - offload_pt_weights=False, # Need the weights in memory for prefill-model export/compilation in the next step - ) - - qpc_outputs = [] - decode_session = QAICInferenceSession(decode_qpc_path) - decode_session.set_buffers({"logits": logits_out_placeholder}) - - decode_inputs = { - "input_ids": np.argmax(qpc_out["logits"]).reshape(1, 1), - "position_ids": np.max(inputs["position_ids"]).reshape(1, 1) + 1, - } - - qpc_outputs.append(decode_inputs["input_ids"][0][0]) - for i in range(config.num_hidden_layers): - if i % 2 == 0 and decode_inputs["position_ids"] >= config.sliding_window: - k = qpc_out[f"past_key.{i}_RetainedState"] - v = qpc_out[f"past_value.{i}_RetainedState"] - mod_pos_id = config.sliding_window - decode_inputs["position_ids"][0][0] % config.sliding_window - decode_inputs[f"past_key.{i}"] = np.concatenate((k[:, :, mod_pos_id:, :], k[:, :, :mod_pos_id, :]), axis=-2) - decode_inputs[f"past_value.{i}"] = np.concatenate( - (v[:, :, mod_pos_id:, :], v[:, :, :mod_pos_id, :]), axis=-2 - ) - else: - decode_inputs[f"past_key.{i}"] = qpc_out[f"past_key.{i}_RetainedState"] - decode_inputs[f"past_value.{i}"] = qpc_out[f"past_value.{i}_RetainedState"] - - decode_out = decode_session.run(decode_inputs) - decode_session.skip_buffers( - [x for x in decode_session.input_names + decode_session.output_names if x.startswith("past_")] - ) - pos_id = np.max(decode_inputs["position_ids"]).reshape(1, 1) + 1 - for i in range(generation_len - 1): - loop_decode_inputs = { - "input_ids": np.argmax(decode_out["logits"]).reshape(1, 1), - "position_ids": pos_id, - } - qpc_outputs.append(loop_decode_inputs["input_ids"][0][0]) - decode_out = decode_session.run(loop_decode_inputs) - pos_id += 1 - - print("QPC Outputs (AIC): \n") - print("Prompt:", repr(prompt)) - print("Completion:", repr(tokenizer.decode(qpc_outputs))) - - -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_id", model_id_blocking) -@pytest.mark.parametrize("prompt", [prompt1]) -def test_disagg_mode_prefix_caching(model_id, prompt): - PREFILL_SEQ_LEN = 128 - CTX_LEN = 128 * 3 - config = AutoConfig.from_pretrained(model_id, num_hidden_layers=2) - prefill_qeff_model = QEFFAutoModelForCausalLM.from_pretrained( - model_id, num_hidden_layers=2, continuous_batching=True - ) - prefill_qeff_model.prefill(enable=True, enable_chunking=True) - prefill_qpc_path = prefill_qeff_model.compile( - prefill_seq_len=PREFILL_SEQ_LEN, - ctx_len=CTX_LEN, - num_cores=16, - mxfp6_matmul=False, - mxint8_kv_cache=False, - num_devices=1, - mos=1, - aic_enable_depth_first=True, - num_speculative_tokens=None, - prefill_only=True, - enable_chunking=True, - full_batch_size=1, - kv_cache_batch_size=2, - ) - - decode_qeff_model = QEFFAutoModelForCausalLM.from_pretrained( - model_id, num_hidden_layers=2, continuous_batching=True - ) - decode_qeff_model.prefill(enable=False) - decode_qpc_path = decode_qeff_model.compile( - prefill_seq_len=1, - ctx_len=CTX_LEN, - num_cores=16, - mxfp6_matmul=False, - mxint8_kv_cache=False, - num_devices=1, - mos=1, - aic_enable_depth_first=True, - num_speculative_tokens=None, - offload_pt_weights=False, # Need the weights in memory for prefill-model export/compilation in the next step - full_batch_size=1, - kv_cache_batch_size=2, - retain_full_kv=True, - ) - - out1, ids1 = prefix_caching_inference(model_id, prefill_qpc_path, decode_qpc_path, prompt, decode_batch_id=0) - out2, ids2 = prefix_caching_inference(model_id, prefill_qpc_path, decode_qpc_path, prompt, decode_batch_id=1) - - for i in range(config.num_hidden_layers): - assert ( - np.abs( - out1[f"past_key.{i}_RetainedState"][0, :, :, :] - out2[f"past_key.{i}_RetainedState"][1, :, :, :] - ).max() - < 5e-2 - ) - assert ( - np.abs( - out1[f"past_value.{i}_RetainedState"][0, :, :, :] - out2[f"past_value.{i}_RetainedState"][1, :, :, :] - ).max() - < 5e-2 - ) - - -def prefix_caching_inference(model_id, prefill_qpc_path, decode_qpc_path, prompt, decode_batch_id): - PREFILL_SEQ_LEN = 128 - tokenizer = AutoTokenizer.from_pretrained(model_id) - config = AutoConfig.from_pretrained(model_id, num_hidden_layers=2) - inputs = tokenizer(prompt, return_tensors="np", padding=True) - padded_len = inputs["input_ids"].shape[1] - num_chunks = -(padded_len // -PREFILL_SEQ_LEN) # ceil divide without float - padded_len = num_chunks * PREFILL_SEQ_LEN # Convert to a multiple of prompt_len - - inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) - inputs["position_ids"] = np.where(inputs.pop("attention_mask"), np.arange(padded_len), -1) - inputs.pop("token_type_ids", None) - inputs["batch_index"] = np.array([[decode_batch_id]], dtype=np.int64) - - prefill_session = QAICInferenceSession(prefill_qpc_path) - logits_out_placeholder = np.zeros((1, 1, config.vocab_size), dtype=np.float32) - prefill_session.set_buffers({"logits": logits_out_placeholder}) - for i in range(num_chunks): - chunk_inputs = inputs.copy() - chunk_inputs["input_ids"] = inputs["input_ids"][:, i * PREFILL_SEQ_LEN : (i + 1) * PREFILL_SEQ_LEN] - chunk_inputs["position_ids"] = inputs["position_ids"][:, i * PREFILL_SEQ_LEN : (i + 1) * PREFILL_SEQ_LEN] - qpc_out = prefill_session.run(chunk_inputs) - del prefill_session - - qpc_outputs = [] - decode_inputs = { - "input_ids": np.argmax(qpc_out["logits"]).reshape(1, 1), - "position_ids": np.max(inputs["position_ids"]).reshape(1, 1) + 1, - "batch_index": inputs["batch_index"], - } - qpc_outputs.append(decode_inputs["input_ids"][0][0]) - - decode_session = QAICInferenceSession(decode_qpc_path) - decode_session.set_buffers({"logits": logits_out_placeholder}) - generation_len = 5 - - for i in range(config.num_hidden_layers): - if i % 2 == 0 and decode_inputs["position_ids"] >= config.sliding_window: - k = qpc_out[f"past_key.{i}_RetainedState"] - v = qpc_out[f"past_value.{i}_RetainedState"] - mod_pos_id = config.sliding_window - decode_inputs["position_ids"][0][0] % config.sliding_window - decode_inputs[f"past_key.{i}"] = np.concatenate((k[:, :, mod_pos_id:, :], k[:, :, :mod_pos_id, :]), axis=-2) - decode_inputs[f"past_value.{i}"] = np.concatenate( - (v[:, :, mod_pos_id:, :], v[:, :, :mod_pos_id, :]), axis=-2 - ) - else: - decode_inputs[f"past_key.{i}"] = qpc_out[f"past_key.{i}_RetainedState"] - decode_inputs[f"past_value.{i}"] = qpc_out[f"past_value.{i}_RetainedState"] - - decode_out = decode_session.run(decode_inputs) - pos_id = np.max(decode_inputs["position_ids"]).reshape(1, 1) + 1 - for i in range(generation_len - 1): - loop_decode_inputs = { - "input_ids": np.argmax(decode_out["logits"]).reshape(1, 1), - "position_ids": pos_id, - "batch_index": inputs["batch_index"], - } - for i in range(config.num_hidden_layers): - loop_decode_inputs[f"past_key.{i}"] = decode_out[f"past_key.{i}_RetainedState"] - loop_decode_inputs[f"past_value.{i}"] = decode_out[f"past_value.{i}_RetainedState"] - qpc_outputs.append(loop_decode_inputs["input_ids"][0][0]) - decode_out = decode_session.run(loop_decode_inputs) - pos_id += 1 - - print("QPC Outputs (AIC): \n") - print("Prompt:", repr(prompt)) - print("Completion:", repr(tokenizer.decode(qpc_outputs))) - return qpc_out, qpc_outputs diff --git a/tests/transformers/models/causal_lm_models/check_causal_models.py b/tests/transformers/models/causal_lm_models/check_causal_models.py index 9e64c2b571..e0e54ce5b2 100644 --- a/tests/transformers/models/causal_lm_models/check_causal_models.py +++ b/tests/transformers/models/causal_lm_models/check_causal_models.py @@ -4,63 +4,118 @@ # SPDX-License-Identifier: BSD-3-Clause # # ----------------------------------------------------------------------------- - import copy import os -from typing import Optional +from typing import Optional, Sequence import numpy as np import torch -from transformers import AutoConfig +from transformers import AutoConfig, AutoTokenizer +from QEfficient.generation.text_generation_inference import TextGeneration from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM from QEfficient.transformers.quantizers.auto import replace_transformers_quantizers from QEfficient.utils._utils import load_hf_tokenizer from QEfficient.utils.config_utils import get_first_config_value from QEfficient.utils.constants import ATTENTION_HEAD_CONFIG_KEYS, KV_HEAD_CONFIG_KEYS, Constants from QEfficient.utils.run_utils import ApiRunner -from QEfficient.utils.test_utils import ModelConfig, load_hf_causal_lm_model +from QEfficient.utils.test_utils import ModelConfig, load_hf_causal_lm_model, load_qeff_causal_lm_model + -from ..check_model_results import dump_and_compare_results +def _sequence_cosine_similarity(seq_a: Sequence[int], seq_b: Sequence[int], vocab_size: int) -> float: + """Compute cosine similarity between two token-id sequences. + Each sequence is represented as a bag-of-tokens count vector of length + ``vocab_size``. This gives a single scalar in [0, 1] that measures how + similar the generated outputs are without requiring exact token-by-token + alignment — useful when hardware quantisation introduces small divergences. -def get_custom_n_layers(model_name): + Args: + seq_a: Reference token-id sequence (e.g. PyTorch output). + seq_b: Candidate token-id sequence (e.g. AIC output). + vocab_size: Total vocabulary size used to size the count vectors. + + Returns: + Cosine similarity in [0, 1]. Returns 1.0 when both sequences are + identical and 0.0 when they share no tokens at all. """ - Function to set number layers of the variuos types of models such as swiftkv models and others - -------- + vec_a = np.zeros(vocab_size, dtype=np.float32) + vec_b = np.zeros(vocab_size, dtype=np.float32) + for token_id in seq_a: + vec_a[int(token_id)] += 1.0 + for token_id in seq_b: + vec_b[int(token_id)] += 1.0 + norm_a = np.linalg.norm(vec_a) + norm_b = np.linalg.norm(vec_b) + if norm_a == 0.0 or norm_b == 0.0: + return 1.0 if norm_a == norm_b else 0.0 + return float(np.dot(vec_a, vec_b) / (norm_a * norm_b)) + - :model_name: str +def get_custom_n_layers(model_name: str) -> int: + """Return the number of hidden layers to use when loading a model for testing. - :return n_layer + Most models are loaded with a single layer to keep tests fast. A small set + of architectures require at least two layers for correctness (e.g. models + with alternating sliding-window attention), and SwiftKV models must be + loaded at full depth so their key-value sharing logic is exercised. + + Args: + model_name: HuggingFace model identifier. + + Returns: + Number of hidden layers to pass to the model loader, or -1 to load + the full model without overriding the layer count. """ if model_name in {"microsoft/Phi-3-mini-4k-instruct", "neuralmagic/Qwen2-0.5B-Instruct-FP8", "openai/gpt-oss-20b"}: return 2 - elif model_name in ModelConfig.SWIFTKV_MODELS: + if model_name in ModelConfig.SWIFTKV_MODELS: return -1 return 1 def check_kv_repeat_causal_lm_pytorch_vs_ai100( model_name: str, - manual_cleanup: callable, - prompt_len: int = Constants.PROMPT_LEN, - ctx_len: int = Constants.CTX_LEN, n_layer: int = -1, config: Optional[AutoConfig] = None, + transform_params: Optional[dict] = None, + export_params: Optional[dict] = None, + compile_params: Optional[dict] = None, + generate_params: Optional[dict] = None, + export_compile_only: bool = False, + continuous_batching: bool = False, + cosine_similarity_threshold: float = 0.95, ): + """Validate causal LM export and compile with KV-head replication (GQA → MHA). + + Reads ``num_attention_heads`` and ``num_key_value_heads`` from the model + config, derives the replication factor, injects it into ``transform_params`` + as ``qaic_config={"num_replicate_kv_heads": factor}``, and delegates to + ``check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100``. + + Args: + model_name: HuggingFace model identifier. + n_layer: Number of hidden layers to load; -1 uses the model default. + config: Optional pre-built ``AutoConfig``; fetched from HF if ``None``. + transform_params: Extra transform kwargs merged with the KV-replicate config. + export_params: Keyword arguments forwarded to ``qeff_model.export()``. + compile_params: Keyword arguments forwarded to ``qeff_model.compile()``. + generate_params: Prompt and generation-length kwargs. + export_compile_only: When ``True``, stop after compilation. + continuous_batching: Whether to enable continuous-batching mode. + cosine_similarity_threshold: Minimum cosine similarity for output checks. """ - Validate causal LM flow with repeated KV heads configuration. - """ - if config is None: - model_config = AutoConfig.from_pretrained( + resolved_config = ( + config + if config is not None + else AutoConfig.from_pretrained( model_name, trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, ) - else: - model_config = config + ) - num_attention_heads = get_first_config_value(model_config, ATTENTION_HEAD_CONFIG_KEYS, default=1, cast_int=True) - num_key_value_heads = get_first_config_value(model_config, KV_HEAD_CONFIG_KEYS, default=None, cast_int=True) + num_attention_heads = get_first_config_value(resolved_config, ATTENTION_HEAD_CONFIG_KEYS, default=1, cast_int=True) + num_key_value_heads = get_first_config_value(resolved_config, KV_HEAD_CONFIG_KEYS, default=None, cast_int=True) if num_key_value_heads is None: num_key_value_heads = num_attention_heads if num_attention_heads < 1 or num_key_value_heads < 1: @@ -75,52 +130,82 @@ def check_kv_repeat_causal_lm_pytorch_vs_ai100( ) num_replicate_kv_heads = num_attention_heads // num_key_value_heads + kv_replicate_transform_params = { + **(transform_params or {}), + "qaic_config": {"num_replicate_kv_heads": num_replicate_kv_heads}, + } check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( model_name=model_name, - manual_cleanup=manual_cleanup, - prompt_len=prompt_len, - ctx_len=ctx_len, n_layer=n_layer, config=config, - qaic_config={"num_replicate_kv_heads": num_replicate_kv_heads}, + transform_params=kv_replicate_transform_params, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=export_compile_only, + continuous_batching=continuous_batching, + cosine_similarity_threshold=cosine_similarity_threshold, ) def check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( model_name: str, - manual_cleanup: callable, - num_devices: int = 1, continuous_batching: bool = False, - prompt_len: int = Constants.PROMPT_LEN, - ctx_len: int = Constants.CTX_LEN, n_layer: int = -1, - num_speculative_tokens: Optional[int] = None, - prefill_only: Optional[bool] = None, - enable_qnn: Optional[bool] = False, - qnn_config: Optional[str] = None, config: Optional[AutoConfig] = None, - pytorch_hf_tokens: Optional[list] = None, - qaic_config: Optional[dict] = None, - retain_full_kv: Optional[bool] = None, - compare_results: bool = False, - compile_only: bool = False, - mdp_num_partitions: Optional[int] = None, - mdp_strategy: Optional[str] = None, - use_onnx_subfunctions: bool = False, + transform_params: Optional[dict] = None, + export_params: Optional[dict] = None, + compile_params: Optional[dict] = None, + generate_params: Optional[dict] = None, + export_compile_only: bool = False, + cosine_similarity_threshold: float = 0.95, ): + """Run the full export → compile → generate pipeline and verify output quality. + + Loads the HuggingFace model and a QEff-transformed copy, runs both on + PyTorch CPU to collect reference token sequences, exports to ONNX, compiles + to a QPC, and (when ``export_compile_only=False``) runs inference on the + compiled QPC. Output quality is measured via cosine similarity between the + PyTorch reference sequences and the AIC-generated sequences. + + Args: + model_name: HuggingFace model identifier. + continuous_batching: Whether to enable continuous-batching mode. + n_layer: Number of hidden layers to load; -1 uses the model default. + config: Optional pre-built ``AutoConfig`` to pass to the model loader. + transform_params: Keyword arguments forwarded to ``QEFFAutoModelForCausalLM`` + and its ``transform()`` call (e.g. ``torch_dtype``, ``qaic_config``). + export_params: Keyword arguments forwarded to ``qeff_model.export()``. + compile_params: Keyword arguments forwarded to ``qeff_model.compile()`` + (e.g. ``prefill_seq_len``, ``ctx_len``, ``num_speculative_tokens``). + generate_params: Keyword arguments used to build the prompt and + generation length for the inference run. + export_compile_only: When ``True``, stop after compilation and skip the + generate + similarity check steps. + cosine_similarity_threshold: Minimum cosine similarity required between + the PyTorch reference and AIC output token sequences. + """ torch.manual_seed(42) replace_transformers_quantizers() - model_hf = load_hf_causal_lm_model(model_name, num_hidden_layers=n_layer, config=config) + torch_dtype = transform_params.get("torch_dtype", torch.float32) + model_hf = load_hf_causal_lm_model(model_name, num_hidden_layers=n_layer, config=config, torch_dtype=torch_dtype) + # print(model_hf) tokenizer = load_hf_tokenizer(pretrained_model_name_or_path=model_name) config = model_hf.config - batch_size = len(Constants.INPUT_STR) - prompts = Constants.INPUT_STR * 4 if continuous_batching else Constants.INPUT_STR + prompt = generate_params.pop("prompt", Constants.INPUT_STR) + prompt_len = compile_params.get("prefill_seq_len", Constants.PROMPT_LEN) + ctx_len = compile_params.get("ctx_len", Constants.CTX_LEN) + num_devices = compile_params.get("num_devices", 1) + batch_size = len(prompt) + prompts = prompt * 4 if continuous_batching else prompt full_batch_size = 4 - gen_len = 24 - is_tlm = False if num_speculative_tokens is None else True + num_speculative_tokens = compile_params.get("num_speculative_tokens", None) + is_tlm = num_speculative_tokens is not None + qaic_config = transform_params.get("qaic_config", None) + prefill_only = compile_params.get("prefill_only", None) + pytorch_hf_tokens = None pytorch_kv_tokens = None - ort_tokens = None qeff_model = QEFFAutoModelForCausalLM( copy.deepcopy(model_hf), @@ -136,37 +221,8 @@ def check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( num_devices=num_devices, qaic_config=qaic_config, ) - api_runner = ApiRunner( - batch_size, - tokenizer, - qeff_model.config, - prompts, - Constants.PROMPT_LEN, - Constants.CTX_LEN, - full_batch_size if continuous_batching else None, - ) - if continuous_batching is False: - pytorch_kv_tokens = api_runner.run_kv_model_on_pytorch(qeff_model.model) - - if model_name not in ModelConfig.SWIFTKV_MODELS and model_name not in ModelConfig.EXTERNAL_MODELS: - if continuous_batching: - pytorch_hf_tokens = api_runner.run_hf_model_on_pytorch_CB(model_hf) - pytorch_hf_tokens = np.vstack(pytorch_hf_tokens) - else: - pytorch_hf_tokens = api_runner.run_hf_model_on_pytorch(model_hf) - - onnx_model_path = qeff_model.export(use_onnx_subfunctions=use_onnx_subfunctions) - if continuous_batching is False: - ort_tokens = api_runner.run_kv_model_on_ort(onnx_model_path, is_tlm=is_tlm) - gen_len = ort_tokens.shape[-1] - - if pytorch_hf_tokens is not None and ort_tokens is not None: - assert (pytorch_hf_tokens == ort_tokens).all(), ( - "Tokens don't match for HF PyTorch model output and ONNXRT output." - ) - if pytorch_kv_tokens is not None and ort_tokens is not None: - assert (pytorch_kv_tokens == ort_tokens).all(), "Tokens don't match for ONNXRT output and PyTorch output." + _ = qeff_model.export(**export_params) compiler_options = {} if continuous_batching and prompt_len == 1: @@ -187,94 +243,307 @@ def check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( compiler_options["specializations"] = [prefill_spec, decode_spec] mdp_compile_kwargs = {} + mdp_num_partitions = compile_params.pop("mdp_num_partitions", None) + mdp_strategy = compile_params.pop("mdp_strategy", None) if mdp_num_partitions is not None: mdp_compile_kwargs["mdp_num_partitions"] = mdp_num_partitions if mdp_strategy is not None: mdp_compile_kwargs["mdp_strategy"] = mdp_strategy qpc_path = qeff_model.compile( - prefill_seq_len=prompt_len, - ctx_len=ctx_len, - num_devices=num_devices, - mxfp6=False, - aic_enable_depth_first=False, - num_speculative_tokens=num_speculative_tokens, - enable_qnn=enable_qnn, - qnn_config=qnn_config, - retain_full_kv=retain_full_kv, - prefill_only=prefill_only, + **compile_params, batch_size=batch_size if continuous_batching else 1, full_batch_size=full_batch_size if continuous_batching else None, - use_onnx_subfunctions=use_onnx_subfunctions, **compiler_options, **mdp_compile_kwargs, ) assert os.path.isfile(os.path.join(os.path.dirname(qpc_path), "qconfig.json")) - if compile_only: - manual_cleanup(onnx_model_path) + if export_compile_only: return - # Generate - exec_info = qeff_model.generate(tokenizer, prompts=prompts) + api_runner = ApiRunner( + batch_size, + tokenizer, + qeff_model.config, + prompts, + Constants.PROMPT_LEN, + Constants.CTX_LEN, + full_batch_size if continuous_batching else None, + ) + if not continuous_batching: + pytorch_kv_tokens = api_runner.run_kv_model_on_pytorch(qeff_model.model) + if model_name not in ModelConfig.SWIFTKV_MODELS and model_name not in ModelConfig.EXTERNAL_MODELS: + if continuous_batching: + pytorch_hf_tokens = api_runner.run_hf_model_on_pytorch_CB(model_hf) + pytorch_hf_tokens = np.vstack(pytorch_hf_tokens) + else: + pytorch_hf_tokens = api_runner.run_hf_model_on_pytorch(model_hf) + + exec_info = qeff_model.generate(tokenizer, prompts=prompts, **generate_params) + vocab_size = config.vocab_size if continuous_batching: - cloud_ai_100_tokens = exec_info.generated_ids - if cloud_ai_100_tokens is not None and ort_tokens is not None: - assert all( - [ - all(ort_token[:24] == cloud_token[:24]) - for ort_token, cloud_token in zip(ort_tokens, cloud_ai_100_tokens) - ] - ), "Tokens don't match for HF PyTorch model output and Cloud AI 100 output." - if pytorch_hf_tokens is not None and cloud_ai_100_tokens is not None: - assert all( - [ - all(pt_token[:24] == cloud_token[:24]) - for pt_token, cloud_token in zip(pytorch_hf_tokens, cloud_ai_100_tokens) - ] - ), "Tokens don't match for HF PyTorch model output and Cloud AI 100 output." + aic_tokens = exec_info.generated_ids + if aic_tokens is not None and pytorch_hf_tokens is not None: + gen_len = pytorch_hf_tokens.shape[-1] + for batch_idx, (pt_seq, aic_seq) in enumerate(zip(pytorch_hf_tokens, aic_tokens)): + similarity = _sequence_cosine_similarity(pt_seq[:gen_len], aic_seq[:gen_len], vocab_size) + assert similarity >= cosine_similarity_threshold, ( + f"Batch index {batch_idx}: cosine similarity between HF PyTorch and AIC output " + f"sequences is {similarity:.4f}, below threshold {cosine_similarity_threshold}." + ) else: - cloud_ai_100_tokens = exec_info.generated_ids[0][:, :gen_len] + gen_len = pytorch_kv_tokens.shape[-1] + aic_tokens = exec_info.generated_ids[0][:, :gen_len] if prefill_only: - assert (ort_tokens[0][0] == cloud_ai_100_tokens[0][0]).all(), ( - "prefill run output tokens don't match for ONNXRT output and Cloud AI 100 output." + # For prefill-only runs compare only the single next-token prediction. + pt_token = int(pytorch_hf_tokens[0][0]) + aic_token = int(aic_tokens[0][0]) + similarity = _sequence_cosine_similarity([pt_token], [aic_token], vocab_size) + assert similarity >= cosine_similarity_threshold, ( + f"Prefill next-token cosine similarity is {similarity:.4f}, " + f"below threshold {cosine_similarity_threshold}. " + f"HF PyTorch token: {pt_token}, AIC token: {aic_token}." ) else: - assert (ort_tokens == cloud_ai_100_tokens).all(), ( - "Tokens don't match for ONNXRT output and Cloud AI 100 output." + pytorch_hf_tokens = pytorch_hf_tokens.tolist() + aic_tokens = aic_tokens.flatten().tolist() + similarity = _sequence_cosine_similarity(pytorch_hf_tokens, aic_tokens, vocab_size) + assert similarity >= cosine_similarity_threshold, ( + f"Batch index {batch_idx}: cosine similarity between HF PyTorch and AIC output " + f"sequences is {similarity:.4f}, below threshold {cosine_similarity_threshold}." ) - manual_cleanup(onnx_model_path) # Clean up the model files after the tests are done. - if compare_results is False: + +def check_prefix_caching_inference( + model_name: str, + continuous_batching: bool = False, + n_layer: int = -1, + config: Optional[AutoConfig] = None, + transform_params: Optional[dict] = None, + export_params: Optional[dict] = None, + compile_params: Optional[dict] = None, + generate_params: Optional[dict] = None, + export_compile_only: bool = False, + cosine_similarity_threshold: float = 0.95, +): + """Compile a prefix-caching model and verify KV-cache reuse correctness. + + Compiles the model with a KV-cache batch size larger than the active batch + so that prefill results can be retained across requests. The test then: + + 1. Generates outputs for two prompts (prefix + suffix1) on batch slots 0 and 1. + 2. Manually runs prefill for the same prompts on slots 2 and 3 and verifies + that the step-by-step decode outputs match the session's ``generated_ids`` + via cosine similarity. + 3. Re-runs prefill on slot 0 with a modified input that shares the same + prefix but differs in the suffix, and verifies that the cached prefix + produces logits consistent with the unmodified run via cosine similarity. + 4. Re-runs decode on slot 1 starting from the cached prefix state and + verifies that the output sequence matches the original generation via + cosine similarity. + + Args: + model_name: HuggingFace model identifier. + continuous_batching: Whether to enable continuous-batching mode. + n_layer: Number of hidden layers to load; -1 uses the model default. + config: Optional pre-built ``AutoConfig`` to pass to the model loader. + transform_params: Keyword arguments forwarded to the model loader. + export_params: Keyword arguments forwarded to ``qeff_model.export()``. + compile_params: Keyword arguments forwarded to ``qeff_model.compile()`` + (must include ``full_batch_size`` and ``kv_cache_batch_size``). + generate_params: Prompt components: ``prefixes``, ``suffixes1``, + ``suffixes2``. + export_compile_only: When ``True``, stop after compilation and skip all + inference and similarity checks. + cosine_similarity_threshold: Minimum cosine similarity required for + each pairwise sequence comparison in this test. + """ + torch.manual_seed(42) + replace_transformers_quantizers() + torch_dtype = transform_params.get("torch_dtype", torch.float32) + qeff_model = load_qeff_causal_lm_model( + model_name=model_name, + num_hidden_layers=n_layer, + continuous_batching=continuous_batching, + config=config, + torch_dtype=torch_dtype, + ) + qeff_model.compile( + **compile_params, + ) + if export_compile_only: return - # Compare results for full model only. - compile_params = { - "prefill_seq_len": prompt_len, - "ctx_len": ctx_len, - "num_devices": num_devices, - "mxfp6": False, - "aic_enable_depth_first": False, - "num_speculative_tokens": num_speculative_tokens, - "enable_qnn": enable_qnn, - "qnn_config": qnn_config, - "retain_full_kv": retain_full_kv, - "prefill_only": prefill_only, - "batch_size": batch_size if continuous_batching else 1, - "full_batch_size": full_batch_size if continuous_batching else None, - "compiler_options": compiler_options, - "compile_only": compile_only, - "mdp_num_partitions": mdp_num_partitions, - "mdp_strategy": mdp_strategy, - "use_onnx_subfunctions": use_onnx_subfunctions, + + qpc_path = qeff_model.qpc_path + assert os.path.isfile(os.path.join(os.path.dirname(qpc_path), "qconfig.json")) + + full_batch_size = compile_params.get("full_batch_size", 2) + ctx_len = compile_params.get("ctx_len", Constants.CTX_LEN) + prefixes = generate_params.get("prefixes", ["Once upon a time ", "Once upon a time "]) + suffixes1 = generate_params.get("suffixes1", ["in a land far away", "there was a small village"]) + suffixes2 = generate_params.get("suffixes2", ["a little girl", "in a bustling city"]) + + tokenizer = AutoTokenizer.from_pretrained(model_name) + + generator = TextGeneration(tokenizer=tokenizer, qpc_path=qpc_path, full_batch_size=full_batch_size, ctx_len=ctx_len) + + prompts = [pref + suff for pref, suff in zip(prefixes, suffixes1)] + + # generation for batch_indices = 0, 1 + prompts_exec_info = generator.generate(prompts) + ############################## + # generation for batch_indices + ############################## + # Run prefill for indices 2, 3 with same prompts + out2, pos2, gen_len2 = generator._qaic_model.run_prefill( + prompts[0], generation_len=None, decode_batch_id=np.array(2, dtype=np.int64).reshape(1, 1) + ) + out3, pos3, gen_len3 = generator._qaic_model.run_prefill( + prompts[1], generation_len=None, decode_batch_id=np.array(3, dtype=np.int64).reshape(1, 1) + ) + + # Run decode for batch indices 2, 3 + decode_inputs = { + "input_ids": np.array([[out2["logits"].argmax(2)[0][0]], [out3["logits"].argmax(2)[0][0]]]), + "position_ids": np.array([[pos2[0][0]], [pos3[0][0]]]), + "batch_index": np.array([[2], [3]], dtype=np.int64), } - assert dump_and_compare_results( - model_name, - compile_params, - "causal_lm_model_results.json", - cloud_ai_100_tokens, - exec_info, - pytorch_hf_tokens, - pytorch_kv_tokens, - ort_tokens, + + # Set logits placeholder for decode + logits_out_placeholder = np.zeros( + ( + generator._qaic_model.full_batch_size, + generator._qaic_model._decode_seq_len, + generator._qaic_model._vocab_size, + ), + dtype=np.float32, + ) + generator._qaic_model._session.set_buffers({"logits": logits_out_placeholder}) + + generation_outputs = [] + for i in range(gen_len2): + generation_outputs.append(decode_inputs["input_ids"]) + outputs = generator._qaic_model._session.run(decode_inputs) + logits = outputs["logits"] + if len(logits.shape) == 2: + logits = np.expand_dims(logits, 1) + next_token_id = logits.argmax(2) + + decode_inputs["input_ids"] = next_token_id + decode_inputs["position_ids"] += 1 + + assert np.all(generator._qaic_model.generated_ids[0, :gen_len2] == [int(val[0, 0]) for val in generation_outputs]) + assert np.all(generator._qaic_model.generated_ids[1, :gen_len2] == [int(val[1, 0]) for val in generation_outputs]) + + ############################## + # Now rerun with cached prefix on 0th index with prompt3 and use -1 for 1st index + ############################## + + nprompts = [pref + suff for pref, suff in zip(prefixes, suffixes2)] + + ## Prefill run on index 0 + prompt = nprompts[0] + inputs = tokenizer(prompt, return_tensors="np", padding=True) + position_ids = inputs["attention_mask"].sum(1, keepdims=True) + padded_len = inputs["input_ids"].shape[1] + num_chunks = -(padded_len // -generator._qaic_model._prefill_seq_len) + padded_len = num_chunks * generator._qaic_model._prefill_seq_len # Convert to a multiple of prompt_len + + # Initialize variables specific to request + # Calculate the max generation length. + max_gen_len = generator._qaic_model._ctx_len - position_ids.max() + + # Set the prefill logic buffer + logits_out_placeholder = np.zeros((1, 1, generator._qaic_model._vocab_size), dtype=np.float32) + generator._qaic_model._session.set_buffers({"logits": logits_out_placeholder}) + inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) + inputs["position_ids"] = np.where(inputs.pop("attention_mask"), np.arange(padded_len), -1) + inputs.pop("token_type_ids", None) + inputs["batch_index"] = np.array([[0]], dtype=np.int64) + norm_outputs = generator._qaic_model._session.run(inputs) + inputs["input_ids"][:, :3] = inputs["input_ids"][:, 4:7] + inputs["input_ids"][:, 3:] = 50256 + inputs["position_ids"][:, :3] = inputs["position_ids"][:, 4:7] + inputs["position_ids"][:, 3:] = -1 + mod_outputs = generator._qaic_model._session.run(inputs) + + vocab_size = generator._qaic_model._vocab_size + unmodified_logit_tokens = norm_outputs["logits"].flatten().tolist() + modified_logit_tokens = mod_outputs["logits"].flatten().tolist() + prefix_kv_similarity = _sequence_cosine_similarity(unmodified_logit_tokens, modified_logit_tokens, vocab_size) + assert prefix_kv_similarity >= cosine_similarity_threshold, ( + f"Prefix-cache KV correctness: cosine similarity between unmodified and " + f"prefix-modified prefill logits is {prefix_kv_similarity:.4f}, " + f"below threshold {cosine_similarity_threshold}." + ) + decode_inputs = { + "input_ids": np.array([[mod_outputs["logits"].argmax(2)[0][0]], [0]]), + "position_ids": np.array([[position_ids[0][0]], [-1]]), + "batch_index": np.array([[0], [1]], dtype=np.int64), + } + + # Set logits placeholder for decode + logits_out_placeholder = np.zeros( + ( + generator._qaic_model.full_batch_size, + generator._qaic_model._decode_seq_len, + generator._qaic_model._vocab_size, + ), + dtype=np.float32, + ) + generator._qaic_model._session.set_buffers({"logits": logits_out_placeholder}) + + generation_outputs = [] + for i in range(max_gen_len): + generation_outputs.append(decode_inputs["input_ids"]) + outputs = generator._qaic_model._session.run(decode_inputs) + logits = outputs["logits"] + if len(logits.shape) == 2: + logits = np.expand_dims(logits, 1) + next_token_id = logits.argmax(2) + + decode_inputs["input_ids"] = next_token_id + decode_inputs["position_ids"][0][0] += 1 + + # TODO: add a check if this matches normal execution for same prompt + ############## + # Now run decode on 1st index again with mod_inputs and check if output is correct + ############## + decode_inputs = { + "input_ids": np.array([[0], [prompts_exec_info.generated_ids[1][0]]]), + "position_ids": np.array([[-1], [9]]), + "batch_index": np.array([[0], [1]], dtype=np.int64), + } + + # Set logits placeholder for decode + logits_out_placeholder = np.zeros( + ( + generator._qaic_model.full_batch_size, + generator._qaic_model._decode_seq_len, + generator._qaic_model._vocab_size, + ), + dtype=np.float32, + ) + generator._qaic_model._session.set_buffers({"logits": logits_out_placeholder}) + + generation_outputs_prefill_cached = [] + for i in range(max_gen_len): + generation_outputs_prefill_cached.append(decode_inputs["input_ids"]) + outputs = generator._qaic_model._session.run(decode_inputs) + logits = outputs["logits"] + if len(logits.shape) == 2: + logits = np.expand_dims(logits, 1) + next_token_id = logits.argmax(2) + + decode_inputs["input_ids"] = next_token_id + decode_inputs["position_ids"][1][0] += 1 + + baseline_seq = prompts_exec_info.generated_ids[1][:113].tolist() + cached_seq = [int(val[1, 0]) for val in generation_outputs_prefill_cached][:113] + cached_similarity = _sequence_cosine_similarity(baseline_seq, cached_seq, vocab_size) + assert cached_similarity >= cosine_similarity_threshold, ( + f"Prefix-cached decode: cosine similarity between baseline and prefix-cached " + f"output sequences is {cached_similarity:.4f}, below threshold {cosine_similarity_threshold}." ) diff --git a/tests/transformers/models/causal_lm_models/config.py b/tests/transformers/models/causal_lm_models/config.py new file mode 100644 index 0000000000..39bda805a2 --- /dev/null +++ b/tests/transformers/models/causal_lm_models/config.py @@ -0,0 +1,132 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import os + +import torch + +from QEfficient.utils.constants import Constants + +causal_lm_models_dict = { + # --- CodeGen --- + "Salesforce/codegen-350M-mono": "hf-internal-testing/tiny-random-CodeGenForCausalLM", + # --- Falcon --- + "tiiuae/falcon-7b": "hf-internal-testing/tiny-random-FalconForCausalLM", + "tiiuae/falcon-40b": "hf-internal-testing/tiny-random-FalconForCausalLM", + # --- Gemma --- + "unsloth/gemma-2b": "optimum-intel-internal-testing/tiny-random-GemmaForCausalLM", + "google/codegemma-2b": "optimum-intel-internal-testing/tiny-random-GemmaForCausalLM", + "google/codegemma-7b": "optimum-intel-internal-testing/tiny-random-GemmaForCausalLM", + "google/gemma-2b": "optimum-intel-internal-testing/tiny-random-GemmaForCausalLM", + "google/gemma-7b": "optimum-intel-internal-testing/tiny-random-GemmaForCausalLM", + # --- Gemma2 --- + "unsloth/gemma-2-2b": "trl-internal-testing/tiny-Gemma2ForCausalLM", + "google/gemma-2-2b": "trl-internal-testing/tiny-Gemma2ForCausalLM", + "google/gemma-2-9b": "trl-internal-testing/tiny-Gemma2ForCausalLM", + "google/gemma-2-27b": "trl-internal-testing/tiny-Gemma2ForCausalLM", + # --- GLM-4.5 MoE --- + "zai-org/GLM-4.5": "trl-internal-testing/tiny-Glm4MoeForCausalLM", + # --- GPT-2 --- + "openai-community/gpt2": "hf-internal-testing/tiny-random-GPT2LMHeadModel", + # --- GPT-J --- + "EleutherAI/gpt-j-6b": "hf-internal-testing/tiny-random-GPTJForCausalLM", + # --- GPT-OSS --- + "openai/gpt-oss-20b": "trl-internal-testing/tiny-GptOssForCausalLM", + # --- Granite MoE --- + "ibm-granite/granite-3.1-1b-a400m-base": "hf-internal-testing/tiny-random-GraniteMoeForCausalLM", + # --- GPTBigCode --- + "bigcode/starcoder": "hf-internal-testing/tiny-random-GPTBigCodeForCausalLM", + "ibm-granite/granite-20b-code-base": "hf-internal-testing/tiny-random-GPTBigCodeForCausalLM", + "ibm-granite/granite-20b-code-base-8k": "hf-internal-testing/tiny-random-GPTBigCodeForCausalLM", + "ibm-granite/granite-20b-code-instruct-8k": "hf-internal-testing/tiny-random-GPTBigCodeForCausalLM", + # --- Granite dense --- + "ibm-granite/granite-3.1-2b-instruct": "hf-internal-testing/tiny-random-GraniteForCausalLM", + "ibm-granite/granite-3.1-8b-instruct": "hf-internal-testing/tiny-random-GraniteForCausalLM", + "ibm-granite/granite-guardian-3.1-2b": "hf-internal-testing/tiny-random-GraniteForCausalLM", + "ibm-granite/granite-guardian-3.1-8b": "hf-internal-testing/tiny-random-GraniteForCausalLM", + # --- Grok-1 --- + "hpcai-tech/grok-1": "hpcai-tech/grok-1", # no tiny found + # --- Jais --- + "inceptionai/jais-adapted-7b": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "inceptionai/jais-adapted-13b-chat": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "inceptionai/jais-adapted-70b": "hf-internal-testing/tiny-random-LlamaForCausalLM", + # --- Llama --- + "TinyLlama/TinyLlama-1.1B-Chat-v1.0": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "codellama/CodeLlama-7b-hf": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "codellama/CodeLlama-13b-hf": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "codellama/CodeLlama-34b-hf": "hf-internal-testing/tiny-random-LlamaForCausalLM", + # "deepseek-ai/DeepSeek-R1-Distill-Llama-70B": "hf-internal-testing/tiny-random-LlamaForCausalLM", # can't open on HF due to QCOM compliance + "lmsys/vicuna-13b-delta-v0": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "lmsys/vicuna-13b-v1.3": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "lmsys/vicuna-13b-v1.5": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "meta-llama/Llama-2-7b-chat-hf": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "meta-llama/Llama-2-13b-chat-hf": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "meta-llama/Llama-2-70b-chat-hf": "hf-internal-testing/tiny-random-LlamaForCausalLM", + "meta-llama/Meta-Llama-3-8B": "trl-internal-testing/tiny-LlamaForCausalLM-3", + "meta-llama/Meta-Llama-3-70B": "trl-internal-testing/tiny-LlamaForCausalLM-3", + "meta-llama/Llama-3.1-8B": "trl-internal-testing/tiny-LlamaForCausalLM-3.1", + "meta-llama/Llama-3.1-70B": "trl-internal-testing/tiny-LlamaForCausalLM-3.1", + "meta-llama/Llama-3.2-1B": "trl-internal-testing/tiny-LlamaForCausalLM-3.2", + "meta-llama/Llama-3.2-3B": "trl-internal-testing/tiny-LlamaForCausalLM-3.2", + "meta-llama/Llama-3.3-70B-Instruct": "llamafactory/tiny-random-Llama-3", + # --- Quantized Llama models --- + "TheBloke/Llama-2-7B-GPTQ": "TheBloke/Llama-2-7B-GPTQ", # no tiny found + "TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ": "TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", # no tiny + "neuralmagic/Llama-3.2-3B-Instruct-FP8": "neuralmagic/Llama-3.2-3B-Instruct-FP8", # no tiny + # --- Llama SwiftKV --- + "Snowflake/Llama-3.1-SwiftKV-8B-Instruct": "Snowflake/Llama-3.1-SwiftKV-8B-Instruct", # no tiny + # --- Mistral --- + "Felladrin/Minueza-32M-Base": "hf-internal-testing/tiny-random-MistralForCausalLM", + "mistralai/Mistral-7B-Instruct-v0.1": "hf-internal-testing/tiny-random-MistralForCausalLM", + "mistralai/Codestral-22B-v0.1": "hf-internal-testing/tiny-random-MistralForCausalLM", + # --- Mixtral MoE --- + "mistralai/Mixtral-8x7B-v0.1": "hf-internal-testing/tiny-random-MixtralForCausalLM", + # --- MPT --- + "wtang06/mpt-125m-c4": "hf-internal-testing/tiny-random-MptForCausalLM", + # --- OLMo2 --- + "allenai/OLMo-2-0425-1B": "hf-internal-testing/tiny-random-Olmo2ForCausalLM", + # --- Phi3 --- + "microsoft/Phi-3-mini-4k-instruct": "optimum-intel-internal-testing/tiny-random-Phi3ForCausalLM", + # --- Qwen2 --- + "Qwen/Qwen2-0.5B": "peft-internal-testing/tiny-dummy-qwen2", + "Qwen/Qwen2-1.5B-Instruct": "peft-internal-testing/tiny-dummy-qwen2", + # "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", # can't open on HF due to QCOM compliance + "neuralmagic/Qwen2-0.5B-Instruct-FP8": "neuralmagic/Qwen2-0.5B-Instruct-FP8", # no tiny + # --- Qwen3 MoE --- + "Qwen/Qwen3-30B-A3B-Instruct-2507": "hf-internal-testing/tiny-random-Qwen3MoeForCausalLM", + # --- Starcoder --- + "bigcode/starcoder2-3b": "hf-internal-testing/tiny-random-Starcoder2ForCausalLM", + "bigcode/starcoder2-15b": "hf-internal-testing/tiny-random-Starcoder2ForCausalLM", +} + + +transform_params = {"torch_dtype": torch.float16} +export_params = {"use_onnx_subfunctions": True} +compile_params = { + "prefill_seq_len": 32, + "ctx_len": 128, + "use_onnx_subfunctions": True, +} +generate_params = {"prompt": Constants.INPUT_STR, "generation_len": 25} + +# Minimum cosine similarity between PyTorch and AIC output token sequences required +# for a generate test to pass. Sequences are compared as one-hot token-id vectors. +# Lower values tolerate more divergence; 0.95 is a practical floor for tiny models. +COSINE_SIMILARITY_THRESHOLD = float(os.environ.get("QEFF_COSINE_SIM_THRESHOLD", "0.90")) + + +QEFF_TEST_PROFILE = os.environ.get("QEFF_TEST_PROFILE", "").strip().lower() + +if QEFF_TEST_PROFILE == "tiny_model": + test_models_causal = set(causal_lm_models_dict.values()) +elif QEFF_TEST_PROFILE == "full_model": + test_models_causal = set(causal_lm_models_dict.keys()) + compile_params["num_devices"] = 4 + compile_params["mxfp6_matmul"] = True + compile_params["mxint8_kv_cache"] = True +else: + test_models_causal = set(causal_lm_models_dict.keys()) diff --git a/tests/transformers/models/causal_lm_models/test_bf16.py b/tests/transformers/models/causal_lm_models/test_bf16.py new file mode 100644 index 0000000000..ad6eb9825e --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_bf16.py @@ -0,0 +1,96 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + + +import pytest +import torch + +from .check_causal_models import check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100 +from .config import ( + COSINE_SIMILARITY_THRESHOLD, + QEFF_TEST_PROFILE, + causal_lm_models_dict, + compile_params, + export_params, + generate_params, + test_models_causal, +) + + +@pytest.mark.skip(reason="enable it later when all the models are fixed") +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_bf16_export_bf16_compile_ccl_cb(model_name): + """Verify that BF16 export and compilation with CCL and continuous batching succeed. + + Exports the model in BF16 with CCL enabled and continuous batching, compiles + with compute-context-length (CCL) specialisations for both prefill and decode, + and asserts that ``qconfig.json`` is present. Inference is not run; this test + validates the export-to-compile pipeline for the BF16 + CCL + CB configuration. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + qaic_config = {"ccl_enabled": True} + temp_transform_params = {"torch_dtype": torch.bfloat16, "qaic_config": qaic_config} + temp_compile_params = { + **compile_params, + "num_cores": 4, + "aic-hw-version": "ai200", + "comp_ctx_lengths_prefill": [256, 500], + "comp_ctx_lengths_decode": [512, 1024], + } + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=temp_transform_params, + export_params=export_params, + compile_params=temp_compile_params, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.skip(reason="BF16 + CCL end-to-end generate not yet supported on AIC") +@pytest.mark.llm +@pytest.mark.qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_bf16_export_bf16_compile_generate_ccl_cb(model_name): + """Verify end-to-end BF16 inference quality with CCL and continuous batching. + + Exports in BF16 with CCL enabled and continuous batching, compiles with CCL + specialisations, runs inference on the AIC device, and checks cosine similarity + between AIC output sequences and the HF PyTorch reference. + + Currently skipped: BF16 + CCL end-to-end generate is not yet supported on AIC. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + qaic_config = {"ccl_enabled": True} + temp_transform_params = {"torch_dtype": torch.bfloat16, "qaic_config": qaic_config} + temp_compile_params = { + **compile_params, + "num_cores": 4, + "aic-hw-version": "ai200", + "comp_ctx_lengths_prefill": [256, 500], + "comp_ctx_lengths_decode": [512, 1024], + } + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=temp_transform_params, + export_params=export_params, + compile_params=temp_compile_params, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) diff --git a/tests/transformers/models/causal_lm_models/test_blocking.py b/tests/transformers/models/causal_lm_models/test_blocking.py new file mode 100644 index 0000000000..35179a7887 --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_blocking.py @@ -0,0 +1,213 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + + +import pytest + +from .check_causal_models import check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100 +from .config import ( + COSINE_SIMILARITY_THRESHOLD, + QEFF_TEST_PROFILE, + causal_lm_models_dict, + compile_params, + export_params, + generate_params, + test_models_causal, + transform_params, +) + +# Blocking configuration constants used across both compile-only and generate tests. +HEAD_BLOCK_SIZE = 8 +NUM_KV_BLOCKS = 2 +NUM_Q_BLOCKS = 2 + + +@pytest.mark.skip(reason="enable it later when all the models are fixed") +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_blocking_cb(model_name): + """Verify that FP16 export and compilation succeed for all attention-blocking modes with continuous batching. + + Runs five sub-cases in sequence, each enabling a different combination of + head, KV, and Q blocking via ``qaic_config``: + + 1. Head blocking only (``head_block_size``). + 2. KV blocking only (``num_kv_blocks``). + 3. Q blocking only (``num_q_blocks``). + 4. Combined KV + Q blocking. + 5. Combined head + KV + Q blocking. + + Each sub-case asserts that ``qconfig.json`` is present after compilation. + Inference is not run. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + # head blocking only + qaic_config = dict(enable_blocking=True, head_block_size=HEAD_BLOCK_SIZE) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + # kv blocking only + qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + # q blocking only + qaic_config = dict(enable_blocking=True, num_q_blocks=NUM_Q_BLOCKS) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + # combined kv + q blocking + qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS, num_q_blocks=NUM_Q_BLOCKS) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + # combined head + kv + q blocking + qaic_config = dict( + enable_blocking=True, + head_block_size=HEAD_BLOCK_SIZE, + num_kv_blocks=NUM_KV_BLOCKS, + num_q_blocks=NUM_Q_BLOCKS, + ) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.skip(reason="enable it later when all the models are fixed") +@pytest.mark.llm +@pytest.mark.qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_generate_blocking_cb(model_name): + """Verify end-to-end FP16 inference quality for all attention-blocking modes with continuous batching. + + Runs five sub-cases in sequence, each enabling a different combination of + head, KV, and Q blocking via ``qaic_config``: + + 1. Head blocking only (``head_block_size``). + 2. KV blocking only (``num_kv_blocks``). + 3. Q blocking only (``num_q_blocks``). + 4. Combined KV + Q blocking. + 5. Combined head + KV + Q blocking. + + Each sub-case runs inference on the AIC device and checks that the cosine + similarity between the AIC output sequences and the HF PyTorch reference + sequences meets the configured threshold for every batch slot. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + # head blocking only + qaic_config = dict(enable_blocking=True, head_block_size=HEAD_BLOCK_SIZE) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + # kv blocking only + qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + # q blocking only + qaic_config = dict(enable_blocking=True, num_q_blocks=NUM_Q_BLOCKS) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + # combined kv + q blocking + qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS, num_q_blocks=NUM_Q_BLOCKS) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + # combined head + kv + q blocking + qaic_config = dict( + enable_blocking=True, + head_block_size=HEAD_BLOCK_SIZE, + num_kv_blocks=NUM_KV_BLOCKS, + num_q_blocks=NUM_Q_BLOCKS, + ) + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) diff --git a/tests/transformers/models/causal_lm_models/test_causal_lm_blocking_hqkv.py b/tests/transformers/models/causal_lm_models/test_causal_lm_blocking_hqkv.py deleted file mode 100644 index 0568939cd2..0000000000 --- a/tests/transformers/models/causal_lm_models/test_causal_lm_blocking_hqkv.py +++ /dev/null @@ -1,379 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import json -import os - -import pytest -from transformers import AutoConfig - -from QEfficient.utils.test_utils import ModelConfig - -from .check_causal_models import ( - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100, - get_custom_n_layers, -) - -CONFIG_PATH = os.path.join(os.path.dirname(__file__), "../../../configs/causal_model_configs.json") -with open(CONFIG_PATH, "r") as f: - config_data = json.load(f) - blockedKV_models = config_data["blockedKV_causal_lm_models"] -test_models_blockedKV = [model["model_name"] for model in blockedKV_models] -model_config_dict = {model["model_name"]: model for model in blockedKV_models} - - -@pytest.mark.full_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_blockedKV[:1]) -def test_full_causal_all_blocking_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - HEAD_BLOCK_SIZE = 8 - NUM_KV_BLOCKS = 2 - NUM_Q_BLOCKS = 2 - - # head blocking only - qaic_config = dict(enable_blocking=True, head_block_size=HEAD_BLOCK_SIZE) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, manual_cleanup=manual_cleanup, num_devices=4 - ) - - # kv blocking only - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, manual_cleanup=manual_cleanup, num_devices=4 - ) - - # q block only - qaic_config = dict(enable_blocking=True, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, manual_cleanup=manual_cleanup, num_devices=4 - ) - - # qkv blocking - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, manual_cleanup=manual_cleanup, num_devices=4 - ) - - # head qkv blocking - qaic_config = dict( - enable_blocking=True, - head_block_size=HEAD_BLOCK_SIZE, - num_kv_blocks=NUM_KV_BLOCKS, - num_q_blocks=NUM_Q_BLOCKS, - ) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, manual_cleanup=manual_cleanup, num_devices=4 - ) - - -@pytest.mark.few_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_blockedKV[:1]) -def test_few_causal_all_blocking_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - HEAD_BLOCK_SIZE = 8 - NUM_KV_BLOCKS = 2 - NUM_Q_BLOCKS = 2 - n_layer = get_custom_n_layers(model_name) - # head blocking only - qaic_config = dict(enable_blocking=True, head_block_size=HEAD_BLOCK_SIZE) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, manual_cleanup=manual_cleanup - ) - - # kv blocking only - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, manual_cleanup=manual_cleanup - ) - - # q block only - qaic_config = dict(enable_blocking=True, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, manual_cleanup=manual_cleanup - ) - - # qkv blocking - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, manual_cleanup=manual_cleanup - ) - - # head qkv blocking - qaic_config = dict( - enable_blocking=True, - head_block_size=HEAD_BLOCK_SIZE, - num_kv_blocks=NUM_KV_BLOCKS, - num_q_blocks=NUM_Q_BLOCKS, - ) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, manual_cleanup=manual_cleanup - ) - - -@pytest.mark.dummy_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_blockedKV[:1]) -def test_dummy_causal_all_blocking_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - HEAD_BLOCK_SIZE = 8 - NUM_KV_BLOCKS = 2 - NUM_Q_BLOCKS = 2 - - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **model_config_dict[model_name].get("additional_params", {}), - ) - n_layer = -1 - if model_name in ModelConfig.QUANTIZED_MODELS: - n_layer = get_custom_n_layers(model_name) - hf_config = None - # head blocking only - qaic_config = dict(enable_blocking=True, head_block_size=HEAD_BLOCK_SIZE) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, config=hf_config, manual_cleanup=manual_cleanup - ) - - # kv blocking only - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, config=hf_config, manual_cleanup=manual_cleanup - ) - - # q block only - qaic_config = dict(enable_blocking=True, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, config=hf_config, manual_cleanup=manual_cleanup - ) - - # qkv blocking - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, config=hf_config, manual_cleanup=manual_cleanup - ) - - # head qkv blocking - qaic_config = dict( - enable_blocking=True, - head_block_size=HEAD_BLOCK_SIZE, - num_kv_blocks=NUM_KV_BLOCKS, - num_q_blocks=NUM_Q_BLOCKS, - ) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, qaic_config=qaic_config, n_layer=n_layer, config=hf_config, manual_cleanup=manual_cleanup - ) - - -@pytest.mark.full_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_blockedKV[:1]) -def test_full_causal_all_blocking_pytorch_vs_kv_vs_ort_vs_ai100_CB(model_name, manual_cleanup): - HEAD_BLOCK_SIZE = 8 - NUM_KV_BLOCKS = 2 - NUM_Q_BLOCKS = 2 - - # head blocking only - qaic_config = dict(enable_blocking=True, head_block_size=HEAD_BLOCK_SIZE) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - num_devices=4, - ) - - # kv blocking only - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - num_devices=4, - ) - - # q block only - qaic_config = dict(enable_blocking=True, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - num_devices=4, - ) - - # qkv blocking - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - num_devices=4, - ) - - # head qkv blocking - qaic_config = dict( - enable_blocking=True, - head_block_size=HEAD_BLOCK_SIZE, - num_kv_blocks=NUM_KV_BLOCKS, - num_q_blocks=NUM_Q_BLOCKS, - ) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - num_devices=4, - ) - - -@pytest.mark.few_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_blockedKV[:1]) -def test_few_causal_all_blocking_pytorch_vs_kv_vs_ort_vs_ai100_CB(model_name, manual_cleanup): - HEAD_BLOCK_SIZE = 8 - NUM_KV_BLOCKS = 2 - NUM_Q_BLOCKS = 2 - n_layer = get_custom_n_layers(model_name) - # head blocking only - qaic_config = dict(enable_blocking=True, head_block_size=HEAD_BLOCK_SIZE) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - # kv blocking only - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - # q block only - qaic_config = dict(enable_blocking=True, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - # qkv blocking - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - # head qkv blocking - qaic_config = dict( - enable_blocking=True, - head_block_size=HEAD_BLOCK_SIZE, - num_kv_blocks=NUM_KV_BLOCKS, - num_q_blocks=NUM_Q_BLOCKS, - ) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - -@pytest.mark.dummy_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_blockedKV[:1]) -def test_dummy_causal_all_blocking_pytorch_vs_kv_vs_ort_vs_ai100_CB(model_name, manual_cleanup): - HEAD_BLOCK_SIZE = 8 - NUM_KV_BLOCKS = 2 - NUM_Q_BLOCKS = 2 - - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **model_config_dict[model_name].get("additional_params", {}), - ) - n_layer = -1 - if model_name in ModelConfig.QUANTIZED_MODELS: - n_layer = get_custom_n_layers(model_name) - hf_config = None - # head blocking only - qaic_config = dict(enable_blocking=True, head_block_size=HEAD_BLOCK_SIZE) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - config=hf_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - # kv blocking only - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - config=hf_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - # q block only - qaic_config = dict(enable_blocking=True, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - config=hf_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - # qkv blocking - qaic_config = dict(enable_blocking=True, num_kv_blocks=NUM_KV_BLOCKS, num_q_blocks=NUM_Q_BLOCKS) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - config=hf_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) - - # head qkv blocking - qaic_config = dict( - enable_blocking=True, - head_block_size=HEAD_BLOCK_SIZE, - num_kv_blocks=NUM_KV_BLOCKS, - num_q_blocks=NUM_Q_BLOCKS, - ) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - qaic_config=qaic_config, - n_layer=n_layer, - config=hf_config, - manual_cleanup=manual_cleanup, - continuous_batching=True, - ) diff --git a/tests/transformers/models/causal_lm_models/test_causal_lm_models.py b/tests/transformers/models/causal_lm_models/test_causal_lm_models.py deleted file mode 100644 index 88169fa3fb..0000000000 --- a/tests/transformers/models/causal_lm_models/test_causal_lm_models.py +++ /dev/null @@ -1,235 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import json -import os - -import pytest -from transformers import AutoConfig - -from QEfficient.utils._utils import create_json -from QEfficient.utils.constants import QnnConstants -from QEfficient.utils.test_utils import ModelConfig - -from .check_causal_models import ( - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100, - check_kv_repeat_causal_lm_pytorch_vs_ai100, - get_custom_n_layers, -) - -CONFIG_PATH = os.path.join(os.path.dirname(__file__), "../../../configs/causal_model_configs.json") -with open(CONFIG_PATH, "r") as f: - config_data = json.load(f) - causal_lm_models = config_data["causal_lm_models"] -test_models_causal = [model["model_name"] for model in causal_lm_models] -model_config_dict = {model["model_name"]: model for model in causal_lm_models} - - -@pytest.mark.full_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_causal) -def test_full_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - if model_name in ModelConfig.SKIPPED_MODELS: - pytest.skip("Test skipped for this model due to issues in HF.") - if model_name in ModelConfig.FULL_MODEL_TESTS_TO_SKIP: - pytest.skip(f"Skipping full model test for {model_name} due to resource constraints.") - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name, compare_results=True, manual_cleanup=manual_cleanup, num_devices=4 - ) - - -@pytest.mark.few_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_causal) -def test_few_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - if model_name in ModelConfig.SKIPPED_MODELS: - pytest.skip("Test skipped for this model due to issues in HF.") - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100(model_name=model_name, n_layer=n_layer, manual_cleanup=manual_cleanup) - - -@pytest.mark.few_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("use_onnx_subfunctions", [False, True]) -@pytest.mark.parametrize("model_name", test_models_causal) -def test_few_causal_lm_onnx_mdp_compile_only(model_name, use_onnx_subfunctions, manual_cleanup): - if model_name in ModelConfig.SKIPPED_MODELS: - pytest.skip("Test skipped for this model due to issues in HF.") - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - n_layer=n_layer, - manual_cleanup=manual_cleanup, - compile_only=True, - mdp_num_partitions=2, - mdp_strategy="onnx", - use_onnx_subfunctions=use_onnx_subfunctions, - ) - - -@pytest.mark.dummy_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_causal) -def test_dummy_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - if model_name in ModelConfig.SKIPPED_MODELS: - pytest.skip("Test skipped for this model due to issues in HF.") - custom_config = model_config_dict[model_name] - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **custom_config.get("additional_params", {}), - ) - if model_name in ModelConfig.QUANTIZED_MODELS: - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100(model_name, n_layer=n_layer, manual_cleanup=manual_cleanup) - else: - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100(model_name, config=hf_config, manual_cleanup=manual_cleanup) - - -@pytest.mark.dummy_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_causal) -def test_check_kv_repeat_custom_causal_lm_pytorch_vs_ai100(model_name, manual_cleanup): - """ - Test function to validate the PyTorch model and the Cloud AI 100 model with repeating original KV heads. - ``Mandatory`` Args: - :model_name (str): Hugging Face Model Card name, Example: ``gpt2`` - """ - if model_name in ModelConfig.SKIPPED_MODELS: - pytest.skip("Test skipped for this model due to issues in HF.") - custom_config = model_config_dict[model_name] - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **custom_config.get("additional_params", {}), - ) - if model_name in ModelConfig.REPEAT_KV_TEST_MODELS: - if model_name in ModelConfig.QUANTIZED_MODELS: - n_layer = get_custom_n_layers(model_name) - check_kv_repeat_causal_lm_pytorch_vs_ai100(model_name, manual_cleanup=manual_cleanup, n_layer=n_layer) - else: - check_kv_repeat_causal_lm_pytorch_vs_ai100(model_name, manual_cleanup=manual_cleanup, config=hf_config) - else: - pytest.skip(f"Skipping {model_name} as it is not in REPEAT_KV_TEST_MODELS") - - -@pytest.mark.full_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_causal) -def test_full_causal_lm_pytorch_vs_ort_vs_ai100_cb(model_name, manual_cleanup): - if model_name in ModelConfig.SKIPPED_MODELS: - pytest.skip("Test skipped for this model due to issues in HF.") - if model_name in ModelConfig.FULL_MODEL_TESTS_TO_SKIP: - pytest.skip(f"Skipping full model test for {model_name} due to resource constraints.") - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name, - continuous_batching=True, - manual_cleanup=manual_cleanup, - num_devices=4, - ) - - -@pytest.mark.few_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_causal) -def test_few_causal_lm_pytorch_vs_ort_vs_ai100_cb(model_name, manual_cleanup): - if model_name in ModelConfig.SKIPPED_MODELS: - pytest.skip("Test skipped for this model due to issues in HF.") - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - n_layer=n_layer, - continuous_batching=True, - manual_cleanup=manual_cleanup, - ) - - -@pytest.mark.dummy_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_causal) -def test_dummy_causal_lm_pytorch_vs_ort_vs_ai100_cb(model_name, manual_cleanup): - if model_name in ModelConfig.SKIPPED_MODELS: - pytest.skip("Test skipped for this model due to issues in HF.") - custom_config = model_config_dict[model_name] - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **custom_config.get("additional_params", {}), - ) - if model_name in ModelConfig.QUANTIZED_MODELS: - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name, - n_layer=n_layer, - continuous_batching=True, - manual_cleanup=manual_cleanup, - ) - else: - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name, - config=hf_config, - continuous_batching=True, - manual_cleanup=manual_cleanup, - ) - - -######################### QNN Tests ######################### - - -@pytest.mark.on_qaic -@pytest.mark.qnn -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_causal) -def test_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100_qnn(model_name, manual_cleanup): - """ - QNN Setup - Test function to validate the PyTorch model, the PyTorch model after KV changes, the ONNX model, and the Cloud AI 100 model, both with and without continuous batching. - ``Mandatory`` Args: - :model_name (str): Hugging Face Model Card name, Example: ``gpt2`` - """ - qnn_config_json_path = os.path.join(os.getcwd(), "qnn_config.json") - create_json(qnn_config_json_path, QnnConstants.QNN_SAMPLE_CONFIG) - n_layer = get_custom_n_layers(model_name) - - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - n_layer=n_layer, - enable_qnn=True, - qnn_config=qnn_config_json_path, - manual_cleanup=manual_cleanup, - ) - - -@pytest.mark.on_qaic -@pytest.mark.qnn -@pytest.mark.llm_model -def test_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100_pl1_qnn(manual_cleanup): - """ - Test function to validate the PyTorch model, the PyTorch model after KV changes, the ONNX model, and the Cloud AI 100 model for a prompt length of 1, both with and without continuous batching. - """ - model_name = "gpt2" - prompt_len = 1 - - qnn_config_json_path = os.path.join(os.getcwd(), "qnn_config.json") - create_json(qnn_config_json_path, QnnConstants.QNN_SAMPLE_CONFIG) - - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - prompt_len=prompt_len, - enable_qnn=True, - qnn_config=qnn_config_json_path, - manual_cleanup=manual_cleanup, - num_devices=4, - ) diff --git a/tests/transformers/models/causal_lm_models/test_causal_lm_pl1.py b/tests/transformers/models/causal_lm_models/test_causal_lm_pl1.py deleted file mode 100644 index f5f2384e67..0000000000 --- a/tests/transformers/models/causal_lm_models/test_causal_lm_pl1.py +++ /dev/null @@ -1,153 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import json -import os - -import pytest -import torch -from transformers import AutoConfig - -from QEfficient.utils.test_utils import ModelConfig - -from .check_causal_models import ( - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100, -) - -CONFIG_PATH = os.path.join(os.path.dirname(__file__), "../../../configs/causal_model_configs.json") -with open(CONFIG_PATH, "r") as f: - config_data = json.load(f) - causal_pl1_models = config_data["causal_lm_models_pl1"] -test_models_pl1 = [model["model_name"] for model in causal_pl1_models] -model_config_dict = {model["model_name"]: model for model in causal_pl1_models} - - -@pytest.mark.full_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_pl1) -@pytest.mark.parametrize("retain_full_kv", [True, False]) -def test_full_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100_pl1(model_name, retain_full_kv, manual_cleanup): - if model_name == "gpt2" and retain_full_kv: - pytest.skip("Skipping test for gpt2 with retain_full_kv=True as it is not supported.") - - torch.manual_seed(42) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - prompt_len=1, - retain_full_kv=retain_full_kv, - manual_cleanup=manual_cleanup, - num_devices=4, - ) - - -@pytest.mark.few_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_pl1) -@pytest.mark.parametrize("retain_full_kv", [True, False]) -def test_few_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100_pl1(model_name, retain_full_kv, manual_cleanup): - if model_name == "gpt2" and retain_full_kv: - pytest.skip("Skipping test for gpt2 with retain_full_kv=True as it is not supported.") - torch.manual_seed(42) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - n_layer=2, - prompt_len=1, - retain_full_kv=retain_full_kv, - manual_cleanup=manual_cleanup, - ) - - -@pytest.mark.dummy_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_pl1) -@pytest.mark.parametrize("retain_full_kv", [True, False]) -def test_dummy_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100_pl1(model_name, retain_full_kv, manual_cleanup): - if model_name == "gpt2" and retain_full_kv: - pytest.skip("Skipping test for gpt2 with retain_full_kv=True as it is not supported.") - - torch.manual_seed(42) - custom_config = model_config_dict[model_name] - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **custom_config.get("additional_params", {}), - ) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - prompt_len=1, - retain_full_kv=retain_full_kv, - config=hf_config, - manual_cleanup=manual_cleanup, - ) - - -@pytest.mark.full_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_pl1) -@pytest.mark.parametrize("retain_full_kv", [True, False]) -def test_full_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100_pl1_CB(model_name, retain_full_kv, manual_cleanup): - if model_name == "gpt2" and retain_full_kv: - pytest.skip("Skipping test for gpt2 with retain_full_kv=True as it is not supported.") - torch.manual_seed(42) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - continuous_batching=True, - prompt_len=1, - retain_full_kv=retain_full_kv, - manual_cleanup=manual_cleanup, - num_devices=4, - ) - - -@pytest.mark.few_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_pl1) -@pytest.mark.parametrize("retain_full_kv", [True, False]) -def test_few_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100_pl1_CB(model_name, retain_full_kv, manual_cleanup): - if model_name == "gpt2" and retain_full_kv: - pytest.skip("Skipping test for gpt2 with retain_full_kv=True as it is not supported.") - torch.manual_seed(42) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - n_layer=2, - continuous_batching=True, - prompt_len=1, - retain_full_kv=retain_full_kv, - manual_cleanup=manual_cleanup, - ) - - -@pytest.mark.dummy_layers -@pytest.mark.llm_model -@pytest.mark.on_qaic -@pytest.mark.parametrize("model_name", test_models_pl1) -@pytest.mark.parametrize("retain_full_kv", [True, False]) -def test_dummy_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100_pl1_CB(model_name, retain_full_kv, manual_cleanup): - if model_name == "gpt2" and retain_full_kv: - pytest.skip("Skipping test for gpt2 with retain_full_kv=True as it is not supported.") - - torch.manual_seed(42) - custom_config = model_config_dict[model_name] - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **custom_config.get("additional_params", {}), - ) - - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - continuous_batching=True, - prompt_len=1, - retain_full_kv=retain_full_kv, - config=hf_config, - manual_cleanup=manual_cleanup, - ) diff --git a/tests/transformers/models/causal_lm_models/test_causal_tlm_models.py b/tests/transformers/models/causal_lm_models/test_causal_tlm_models.py deleted file mode 100644 index 9d02acbd29..0000000000 --- a/tests/transformers/models/causal_lm_models/test_causal_tlm_models.py +++ /dev/null @@ -1,123 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import json -import os - -import pytest -from transformers import AutoConfig - -from QEfficient.utils.constants import Constants -from QEfficient.utils.test_utils import ModelConfig - -from .check_causal_models import ( - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100, - get_custom_n_layers, -) - -CONFIG_PATH = os.path.join(os.path.dirname(__file__), "../../../configs/causal_model_configs.json") -with open(CONFIG_PATH, "r") as f: - config_data = json.load(f) - spd_models = config_data["spd_causal_lm_models"] -test_models_spd = [model["model_name"] for model in spd_models] -model_config_dict = {model["model_name"]: model for model in spd_models} - - -@pytest.mark.full_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_spd) -def test_full_causal_tlm_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - num_speculative_tokens=Constants.NUM_SPECULATIVE_TOKENS, - manual_cleanup=manual_cleanup, - num_devices=4, - ) - - -@pytest.mark.few_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_spd) -def test_few_causal_tlm_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - num_speculative_tokens=Constants.NUM_SPECULATIVE_TOKENS, - n_layer=n_layer, - manual_cleanup=manual_cleanup, - ) - - -@pytest.mark.dummy_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_spd) -def test_dummy_causal_tlm_pytorch_vs_kv_vs_ort_vs_ai100(model_name, manual_cleanup): - custom_config = model_config_dict[model_name] - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **custom_config.get("additional_params", {}), - ) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - num_speculative_tokens=Constants.NUM_SPECULATIVE_TOKENS, - config=hf_config, - manual_cleanup=manual_cleanup, - ) - - -@pytest.mark.full_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_spd) -def test_full_causal_tlm_pytorch_vs_kv_vs_ort_vs_ai100_CB(model_name, manual_cleanup): - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - num_speculative_tokens=Constants.NUM_SPECULATIVE_TOKENS, - continuous_batching=True, - manual_cleanup=manual_cleanup, - num_devices=4, - ) - - -@pytest.mark.few_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_spd) -def test_few_causal_tlm_pytorch_vs_kv_vs_ort_vs_ai100_CB(model_name, manual_cleanup): - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - num_speculative_tokens=Constants.NUM_SPECULATIVE_TOKENS, - n_layer=n_layer, - continuous_batching=True, - manual_cleanup=manual_cleanup, - ) - - -@pytest.mark.dummy_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models_spd) -def test_dummy_causal_tlm_pytorch_vs_kv_vs_ort_vs_ai100_CB(model_name, manual_cleanup): - custom_config = model_config_dict[model_name] - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **custom_config.get("additional_params", {}), - ) - - check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( - model_name=model_name, - num_speculative_tokens=Constants.NUM_SPECULATIVE_TOKENS, - config=hf_config, - continuous_batching=True, - manual_cleanup=manual_cleanup, - ) diff --git a/tests/transformers/models/causal_lm_models/test_ccl.py b/tests/transformers/models/causal_lm_models/test_ccl.py new file mode 100644 index 0000000000..9496994f35 --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_ccl.py @@ -0,0 +1,146 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + + +import pytest +import torch + +from .check_causal_models import check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100 +from .config import ( + COSINE_SIMILARITY_THRESHOLD, + QEFF_TEST_PROFILE, + causal_lm_models_dict, + compile_params, + export_params, + generate_params, + test_models_causal, + transform_params, +) + + +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_ccl_cb(model_name): + """Verify that FP16 export and compilation succeed with CCL and continuous batching. + + Exports in FP16 with Compute-Context-Length (CCL) enabled and continuous + batching, compiles with separate prefill and decode CCL specialisations, and + asserts that ``qconfig.json`` is present. Inference is not run. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + qaic_config = {"ccl_enabled": True} + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params={ + **compile_params, + "comp_ctx_lengths_prefill": [256, 500], + "comp_ctx_lengths_decode": [512, 1024], + }, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_generate_ccl_cb(model_name): + """Verify end-to-end FP16 inference quality with CCL and continuous batching. + + Exports in FP16 with CCL enabled and continuous batching, compiles with CCL + specialisations, runs inference on the AIC device, and checks that the cosine + similarity between the AIC output sequences and the HF PyTorch reference + sequences meets the configured threshold for every batch slot. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + qaic_config = {"ccl_enabled": True} + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params={**transform_params, "qaic_config": qaic_config}, + export_params=export_params, + compile_params={ + **compile_params, + "comp_ctx_lengths_prefill": [256, 500], + "comp_ctx_lengths_decode": [512, 1024], + }, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp32_export_fp16_compile_ccl_cb(model_name): + """Verify that FP32 export followed by FP16 compilation succeeds with CCL and continuous batching. + + Exports the model in FP32 (no dtype cast at export time) with CCL enabled + and continuous batching, compiles in FP16 with CCL specialisations, and + asserts that ``qconfig.json`` is present. Inference is not run. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + qaic_config = {"ccl_enabled": True} + fp32_transform_params = {"torch_dtype": torch.float32, "qaic_config": qaic_config} + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=fp32_transform_params, + export_params=export_params, + compile_params={ + **compile_params, + "comp_ctx_lengths_prefill": [256, 500], + "comp_ctx_lengths_decode": [512, 1024], + }, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp32_export_fp16_compile_generate_ccl_cb(model_name): + """Verify end-to-end inference quality for FP32 export + FP16 compile with CCL and continuous batching. + + Exports in FP32 with CCL enabled and continuous batching, compiles in FP16 + with CCL specialisations, runs inference on the AIC device, and checks that + the cosine similarity between the AIC output sequences and the HF PyTorch + reference sequences meets the configured threshold for every batch slot. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + qaic_config = {"ccl_enabled": True} + fp32_transform_params = {"torch_dtype": torch.float32, "qaic_config": qaic_config} + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=fp32_transform_params, + export_params=export_params, + compile_params={ + **compile_params, + "comp_ctx_lengths_prefill": [256, 500], + "comp_ctx_lengths_decode": [512, 1024], + }, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) diff --git a/tests/transformers/models/causal_lm_models/test_default.py b/tests/transformers/models/causal_lm_models/test_default.py new file mode 100644 index 0000000000..e8ed7174e2 --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_default.py @@ -0,0 +1,121 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + + +import pytest + +from .check_causal_models import check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100 +from .config import ( + COSINE_SIMILARITY_THRESHOLD, + QEFF_TEST_PROFILE, + causal_lm_models_dict, + compile_params, + export_params, + generate_params, + test_models_causal, + transform_params, +) + + +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile(model_name): + """Verify that FP16 export and compilation succeed without errors. + + Exports the model to ONNX in FP16 with ONNX subfunctions enabled, compiles + it to a QPC, and asserts that the resulting ``qconfig.json`` is present. + Inference is not run; this test validates the export-to-compile pipeline only. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + transform_params=transform_params, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + continuous_batching=False, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_generate(model_name): + """Verify end-to-end FP16 inference quality against the HF PyTorch reference. + + Exports in FP16 with ONNX subfunctions, compiles, runs inference on the AIC + device, and checks that the cosine similarity between the AIC output token + sequences and the HF PyTorch reference sequences meets the configured threshold. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + transform_params=transform_params, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + continuous_batching=False, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_cb(model_name): + """Verify that FP16 export and compilation succeed in continuous-batching mode. + + Enables continuous batching with ONNX subfunctions, exports to ONNX, compiles + to a QPC, and asserts that ``qconfig.json`` is present. Inference is not run. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=transform_params, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_generate_cb(model_name): + """Verify end-to-end FP16 inference quality in continuous-batching mode. + + Enables continuous batching with ONNX subfunctions, compiles, runs inference + on the AIC device, and checks that the cosine similarity between the AIC + output token sequences and the HF PyTorch reference sequences meets the + configured threshold for every batch slot. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=transform_params, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) diff --git a/tests/transformers/models/causal_lm_models/test_disagg_mode.py b/tests/transformers/models/causal_lm_models/test_disagg_mode.py new file mode 100644 index 0000000000..0951d9bf57 --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_disagg_mode.py @@ -0,0 +1,315 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import os + +import numpy as np +import pytest +import torch +from transformers import AutoConfig +from transformers.cache_utils import DynamicCache + +from QEfficient import QEFFAutoModelForCausalLM +from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.transformers.quantizers import replace_transformers_quantizers, undo_transformers_quantizers + +from .utils import ( + _default_compile_kwargs, + _make_model, + _prefix_caching_inference, + _prepare_inputs, + _rotate_sliding_kv, + _run_hf_prefill, + _run_qeff_prefill, + _run_qpc_prefill, +) + +test_models_blocking_dict = {"openai/gpt-oss-20b": "tiny-random/gpt-oss-bf16"} +test_models_chunking_dict = {"Qwen/Qwen3-30B-A3B-Instruct-2507": "hf-internal-testing/tiny-random-Qwen3MoeForCausalLM"} + + +if os.environ.get("QEFF_TEST_PROFILE", "").strip().lower() == "tiny_model": + model_id_blocking = list(test_models_blocking_dict.values()) + model_id_chunking = list(test_models_chunking_dict.values()) +else: + model_id_blocking = list(test_models_blocking_dict.keys()) + model_id_chunking = list(test_models_chunking_dict.keys()) + + +prompt2 = """ +Once upon a time, in a small town, there lived a young boy named Alex. Alex was a curious and adventurous child, always eager to explore the world around him. One day, while playing in the park, Alex stumbled upon a mysterious old book hidden beneath a pile of leaves. The book was filled with stories of distant lands, magical creatures, and extraordinary adventures. + +As Alex flipped through the pages, he discovered a map that led to a hidden treasure. Excited by the prospect of a real-life treasure hunt, Alex decided to embark on a thrilling journey. He packed his backpack with snacks, a flashlight, and a compass, and set off into the unknown. + +The path to the treasure was not an easy one. Alex had to navigate through dense forests, cross rickety bridges, and solve riddles that guarded the treasure's location. +""" +prompt1 = "Once upon a time" +prompts = [prompt1, prompt2] + +TORCH_ATOL = 1e-4 +QAIC_ATOL = 5e-2 + + +@pytest.mark.qaic +@pytest.mark.llm # FIXME split into llm and vllm later +@pytest.mark.parametrize("model_id", model_id_blocking) +@pytest.mark.parametrize("prompt", prompts) +def test_disagg_mode_prefill(model_id, prompt): + """Verify prefill logit parity across HF PyTorch, QEff PyTorch, and the compiled QPC. + + Runs a single prefill pass through three execution paths and checks that + the last-token logits agree within tolerance at each step: + + 1. HF PyTorch reference model. + 2. QEff-transformed PyTorch model (sliding-window KV cache). + 3. Compiled prefill-only QPC on the AIC device. + """ + PREFILL_SEQ_LEN = 256 + CTX_LEN = 256 + _, hf_inputs, qeff_inputs, _, _ = _prepare_inputs(model_id, prompt, PREFILL_SEQ_LEN) + + # HF run + model, config, out = _run_hf_prefill(model_id, hf_inputs) + + # QEff PyTorch run + qeff_model, config, qeff_out = _run_qeff_prefill(model, qeff_inputs, PREFILL_SEQ_LEN, CTX_LEN, sliding_window=True) + assert (qeff_out.logits - out.logits[:, -1, :]).abs().max() < TORCH_ATOL + + # QPC run on QAIC + prefill_qpc_path = qeff_model.compile(**_default_compile_kwargs(PREFILL_SEQ_LEN, CTX_LEN, 16, prefill_only=True)) + qpc_out, _ = _run_qpc_prefill(qeff_inputs, prefill_qpc_path, config, PREFILL_SEQ_LEN) + assert (torch.from_numpy(qpc_out["logits"]) - qeff_out.logits).abs().max() < QAIC_ATOL + + +@pytest.mark.qaic +@pytest.mark.llm # FIXME split into llm and vllm later +@pytest.mark.parametrize("model_id", model_id_chunking) +@pytest.mark.parametrize("prompt", prompts) +def test_disagg_mode_prefill_chunked(model_id, prompt): + """Verify chunked prefill logit parity across HF PyTorch, QEff PyTorch, and the compiled QPC. + + Splits the prompt into ``prefill_seq_len``-sized chunks and feeds them + sequentially. Checks that the final-chunk last-token logits agree within + tolerance at each step: + + 1. HF PyTorch reference model (full sequence in one pass). + 2. QEff-transformed PyTorch model (chunked, with KV state carried across chunks). + 3. Compiled chunked prefill-only QPC on the AIC device. + """ + PREFILL_SEQ_LEN = 128 + CTX_LEN = 128 * 3 + _, hf_inputs, qeff_inputs, _, num_chunks = _prepare_inputs(model_id, prompt, PREFILL_SEQ_LEN) + + # HF run + model, config, out = _run_hf_prefill(model_id, hf_inputs) + + # QEff PyTorch run (chunked) + qeff_model, config, qeff_out = _run_qeff_prefill( + model, qeff_inputs, PREFILL_SEQ_LEN, CTX_LEN, sliding_window=False, enable_chunking=True, num_chunks=num_chunks + ) + assert (qeff_out.logits - out.logits[:, -1, :]).abs().max() < TORCH_ATOL + + # QPC run on QAIC (chunked) + prefill_qpc_path = qeff_model.compile( + **_default_compile_kwargs( + PREFILL_SEQ_LEN, + CTX_LEN, + config.num_experts, + prefill_only=True, + enable_chunking=True, + ) + ) + qpc_out, _ = _run_qpc_prefill( + qeff_inputs, + prefill_qpc_path, + config, + PREFILL_SEQ_LEN, + enable_chunking=True, + num_chunks=num_chunks, + skip_past_buffers=True, + ) + assert (torch.from_numpy(qpc_out["logits"]) - qeff_out.logits).abs().max() < QAIC_ATOL + + +@pytest.mark.qaic +@pytest.mark.llm # FIXME split into llm and vllm later +@pytest.mark.parametrize("model_id", model_id_blocking) +@pytest.mark.parametrize("prompt", [prompt1]) +def test_disagg_mode_prefill_only_and_decode_only(model_id, prompt): + """Verify disaggregated prefill-only and decode-only compilation and execution. + + Compiles separate prefill-only and decode-only QPCs and validates each stage: + + 1. HF PyTorch full generation loop (reference). + 2. QEff PyTorch prefill + decode loop — asserts token-level parity with the HF reference. + 3. AIC prefill QPC — asserts logit parity with the QEff PyTorch prefill output. + 4. AIC decode QPC — runs the decode loop and prints the completion (no token assertion; + end-to-end parity is covered by the PyTorch decode loop check above). + """ + PREFILL_SEQ_LEN = 256 + CTX_LEN = 256 + generation_len = 10 + tokenizer, hf_inputs, qeff_inputs, _, _ = _prepare_inputs(model_id, prompt, PREFILL_SEQ_LEN) + + # HF prefill run + replace_transformers_quantizers() + model = _make_model(model_id) + config = model.config + orig_out = model(**hf_inputs, past_key_values=DynamicCache(config=config)) + + # HF decode loop + position_ids = qeff_inputs["position_ids"] + generated_ids = [] + out = orig_out + for _ in range(1, generation_len): + next_token_id = out["logits"][:, -1, :].argmax(-1).reshape(-1, 1) + generated_ids.append(next_token_id) + position_ids = position_ids.max(1, keepdim=True).values + 1 + out = model(input_ids=next_token_id, position_ids=position_ids, past_key_values=out["past_key_values"]) + generated_ids.append(out["logits"][:, -1, :].argmax(-1).reshape(-1, 1)) + generated_ids = np.concatenate(generated_ids, axis=1) + print("Original HF Model Outputs (Torch CPU): \n") + print("Prompt:", repr(prompt)) + print("Completion:", repr(tokenizer.batch_decode(generated_ids, skip_special_tokens=True))) + undo_transformers_quantizers() + + # QEff PyTorch prefill + prefill_qeff_model, config, prefill_qeff_out = _run_qeff_prefill( + model, qeff_inputs, PREFILL_SEQ_LEN, CTX_LEN, sliding_window=True + ) + assert (prefill_qeff_out.logits - orig_out.logits[:, -1, :]).abs().max() < TORCH_ATOL + + # QEff PyTorch decode loop + decode_qeff_model = QEFFAutoModelForCausalLM(model) + decode_qeff_model.prefill(enable=False) + qeff_out = prefill_qeff_out + position_ids = qeff_inputs["position_ids"] + qeff_generated_ids = [] + for _ in range(1, generation_len): + next_token_id = qeff_out["logits"][:, -1, :].argmax(-1).reshape(-1, 1) + qeff_generated_ids.append(next_token_id) + position_ids = position_ids.max(1, keepdim=True).values + 1 + qeff_out = decode_qeff_model.model( + input_ids=next_token_id, position_ids=position_ids, past_key_values=qeff_out["past_key_values"] + ) + qeff_generated_ids.append(qeff_out["logits"][:, -1, :].argmax(-1).reshape(-1, 1)) + qeff_generated_ids = np.concatenate(qeff_generated_ids, axis=1) + print("QEFF Transformed Model Outputs (Torch CPU): \n") + print("Prompt:", repr(prompt)) + print("Completion:", repr(tokenizer.batch_decode(qeff_generated_ids, skip_special_tokens=True))) + assert (qeff_generated_ids == generated_ids).all() + + # QPC Compilation + decode_qpc_path = decode_qeff_model.compile( + **_default_compile_kwargs( + 1, + CTX_LEN, + 16, + offload_pt_weights=False, # weights must stay in memory for prefill compile + ) + ) + prefill_qpc_path = prefill_qeff_model.compile( + **_default_compile_kwargs(PREFILL_SEQ_LEN, CTX_LEN, 16, prefill_only=True) + ) + + # QPC prefill run on QAIC + qpc_out, qeff_inputs_np = _run_qpc_prefill(qeff_inputs, prefill_qpc_path, config, PREFILL_SEQ_LEN) + assert (torch.from_numpy(qpc_out["logits"]) - prefill_qeff_out.logits).abs().max() < QAIC_ATOL + + # QPC decode run on QAIC + decode_session = QAICInferenceSession(decode_qpc_path) + decode_session.set_buffers({"logits": np.zeros((1, 1, config.vocab_size), dtype=np.float32)}) + decode_inputs = { + "input_ids": np.argmax(qpc_out["logits"]).reshape(1, 1), + "position_ids": np.max(qeff_inputs_np["position_ids"]).reshape(1, 1) + 1, + } + qpc_outputs = [decode_inputs["input_ids"][0][0]] + _rotate_sliding_kv(qpc_out, decode_inputs, config) + decode_out = decode_session.run(decode_inputs) + decode_session.skip_buffers( + [x for x in decode_session.input_names + decode_session.output_names if x.startswith("past_")] + ) + pos_id = np.max(decode_inputs["position_ids"]).reshape(1, 1) + 1 + for _ in range(generation_len - 1): + loop_inputs = { + "input_ids": np.argmax(decode_out["logits"]).reshape(1, 1), + "position_ids": pos_id, + } + qpc_outputs.append(loop_inputs["input_ids"][0][0]) + decode_out = decode_session.run(loop_inputs) + pos_id += 1 + print("QPC Outputs (AIC): \n") + print("Prompt:", repr(prompt)) + print("Completion:", repr(tokenizer.decode(qpc_outputs))) + + +@pytest.mark.qaic +@pytest.mark.llm # FIXME split into llm and vllm later +@pytest.mark.parametrize("model_id", model_id_blocking) +@pytest.mark.parametrize("prompt", [prompt1]) +def test_disagg_mode_prefix_caching(model_id, prompt): + """Verify that prefix-caching KV state is identical across two batch slots. + + Compiles separate chunked prefill-only and decode-only QPCs with a KV-cache + batch size of 2. Runs the same prompt through both batch slots independently + and asserts that the retained KV state (``past_key`` and ``past_value``) for + every layer is numerically consistent between the two slots within ``QAIC_ATOL``. + """ + PREFILL_SEQ_LEN = 128 + CTX_LEN = 128 * 3 + config = AutoConfig.from_pretrained(model_id) + + # QPC prefill compile + prefill_qeff_model = QEFFAutoModelForCausalLM.from_pretrained(model_id, continuous_batching=True) + prefill_qeff_model.prefill(enable=True, enable_chunking=True) + prefill_qpc_path = prefill_qeff_model.compile( + **_default_compile_kwargs( + PREFILL_SEQ_LEN, + CTX_LEN, + 16, + prefill_only=True, + enable_chunking=True, + full_batch_size=1, + kv_cache_batch_size=2, + ) + ) + + # QPC decode compile + decode_qeff_model = QEFFAutoModelForCausalLM.from_pretrained(model_id, continuous_batching=True) + decode_qeff_model.prefill(enable=False) + decode_qpc_path = decode_qeff_model.compile( + **_default_compile_kwargs( + 1, + CTX_LEN, + 16, + offload_pt_weights=False, # weights must stay in memory for prefill compile + full_batch_size=1, + kv_cache_batch_size=2, + retain_full_kv=True, + ) + ) + + # QPC runs on QAIC for two different batch slots with the same prompt + kv_out_batch0, _ = _prefix_caching_inference( + model_id, config, prefill_qpc_path, decode_qpc_path, prompt, decode_batch_id=0, prefill_seq_len=PREFILL_SEQ_LEN + ) + kv_out_batch1, _ = _prefix_caching_inference( + model_id, config, prefill_qpc_path, decode_qpc_path, prompt, decode_batch_id=1, prefill_seq_len=PREFILL_SEQ_LEN + ) + for i in range(config.num_hidden_layers): + assert ( + np.abs( + kv_out_batch0[f"past_key.{i}_RetainedState"][0] - kv_out_batch1[f"past_key.{i}_RetainedState"][1] + ).max() + < QAIC_ATOL + ) + assert ( + np.abs( + kv_out_batch0[f"past_value.{i}_RetainedState"][0] - kv_out_batch1[f"past_value.{i}_RetainedState"][1] + ).max() + < QAIC_ATOL + ) diff --git a/tests/transformers/models/causal_lm_models/test_fp16_causal_lm.py b/tests/transformers/models/causal_lm_models/test_fp16_causal_lm.py deleted file mode 100644 index af8c3b70f0..0000000000 --- a/tests/transformers/models/causal_lm_models/test_fp16_causal_lm.py +++ /dev/null @@ -1,168 +0,0 @@ -# ----------------------------------------------------------------------------- -# -# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. -# SPDX-License-Identifier: BSD-3-Clause -# -# ----------------------------------------------------------------------------- - -import copy -import json -import os -from typing import Optional - -import pytest -import torch -from transformers import AutoConfig - -from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM -from QEfficient.transformers.quantizers.auto import replace_transformers_quantizers -from QEfficient.utils._utils import load_hf_tokenizer -from QEfficient.utils.constants import Constants -from QEfficient.utils.run_utils import ApiRunner -from QEfficient.utils.test_utils import ModelConfig, load_hf_causal_lm_model - -from .check_causal_models import ( - get_custom_n_layers, -) - -CONFIG_PATH = os.path.join(os.path.dirname(__file__), "../../../configs/causal_model_configs.json") -with open(CONFIG_PATH, "r") as f: - config_data = json.load(f) - causal_lm_fp16_test_models = config_data["causal_lm_fp16_test_models"] -test_models = [model["model_name"] for model in causal_lm_fp16_test_models] -model_config_dict = {model["model_name"]: model for model in causal_lm_fp16_test_models} - - -def check_causal_lm_pytorch_vs_kv_vs_ai100( - model_name: str, - manual_cleanup: callable, - prompt_len: int = Constants.PROMPT_LEN, - ctx_len: int = Constants.CTX_LEN, - n_layer: int = 1, - num_speculative_tokens: Optional[int] = None, - prefill_only: Optional[bool] = None, - enable_qnn: Optional[bool] = False, - qnn_config: Optional[str] = None, - config: Optional[AutoConfig] = None, - pytorch_hf_tokens: Optional[list] = None, - qaic_config: Optional[dict] = None, - retain_full_kv: Optional[bool] = None, - torch_dtype: Optional[torch.dtype] = torch.float32, -): - """ - Validate the PyTorch model, the PyTorch model after KV changes, the ONNX model, and the Cloud AI 100 model, both with and without continuous batching. - ``Mandatory`` Args: - :model_name (str): Hugging Face Model Card name, Example: ``gpt2`` - :prompt_len (int): Prompt length for the model to compile. - :ctx_len (int): Maximum context length to compile the model. - :n_layers (int): Number of layers for the Model. - """ - replace_transformers_quantizers() - - model_hf = load_hf_causal_lm_model(model_name, num_hidden_layers=n_layer, config=config, torch_dtype=torch_dtype) - tokenizer = load_hf_tokenizer(pretrained_model_name_or_path=model_name) - - config = model_hf.config - batch_size = len(Constants.INPUT_STR) - api_runner = ApiRunner( - batch_size, - tokenizer, - config, - Constants.INPUT_STR, - Constants.PROMPT_LEN, - Constants.CTX_LEN, - dtype=torch_dtype, - ) - - if model_name not in ModelConfig.SWIFTKV_MODELS and model_name not in ModelConfig.EXTERNAL_MODELS: - pytorch_hf_tokens = api_runner.run_hf_model_on_pytorch(model_hf) - - is_tlm = False if num_speculative_tokens is None else True - qeff_model = QEFFAutoModelForCausalLM( - copy.deepcopy(model_hf), is_tlm=is_tlm, pretrained_model_name_or_path=model_name, qaic_config=qaic_config - ) - - pytorch_kv_tokens = api_runner.run_kv_model_on_pytorch(qeff_model.model) - - if model_name not in ModelConfig.SWIFTKV_MODELS and model_name not in ModelConfig.EXTERNAL_MODELS: - assert (pytorch_hf_tokens == pytorch_kv_tokens).all(), ( - "Tokens don't match for HF PyTorch model output and KV PyTorch model output" - ) - qeff_model.export() - qpc_path = qeff_model.compile( - prefill_seq_len=prompt_len, - ctx_len=ctx_len, - num_cores=16, - mxfp6=False, - aic_hw_version="ai100", - aic_enable_depth_first=False, - num_speculative_tokens=num_speculative_tokens, - prefill_only=prefill_only, - enable_qnn=enable_qnn, - qnn_config=qnn_config, - ) - exec_info = qeff_model.generate(tokenizer, prompts=Constants.INPUT_STR) - gen_len = pytorch_kv_tokens.shape[-1] - cloud_ai_100_tokens = exec_info.generated_ids[0][ - :, :gen_len - ] # Because we always run for single input and single batch size - if prefill_only: - assert (pytorch_hf_tokens[0][0] == cloud_ai_100_tokens[0][0]).all(), ( - "prefill run output tokens don't match for ONNXRT output and Cloud AI 100 output." - ) - else: - assert (pytorch_hf_tokens == cloud_ai_100_tokens).all(), ( - "Tokens don't match for ONNXRT output and Cloud AI 100 output." - ) - assert os.path.isfile(os.path.join(os.path.dirname(qpc_path), "qconfig.json")) - if prefill_only is not None: - return - - assert os.path.isfile(os.path.join(os.path.dirname(qpc_path), "qconfig.json")) - manual_cleanup(qeff_model.onnx_path) - - -@pytest.mark.full_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models) -def test_full_fp16_causal_lm_pytorch_vs_kv_vs_ai100(model_name, manual_cleanup): - torch.manual_seed(42) - check_causal_lm_pytorch_vs_kv_vs_ai100( - model_name=model_name, torch_dtype=torch.float16, manual_cleanup=manual_cleanup - ) - - -@pytest.mark.few_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models) -def test_few_fp16_causal_lm_pytorch_vs_kv_vs_ai100(model_name, manual_cleanup): - torch.manual_seed(42) - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ai100( - model_name=model_name, n_layer=n_layer, torch_dtype=torch.float16, manual_cleanup=manual_cleanup - ) - - -@pytest.mark.dummy_layers -@pytest.mark.on_qaic -@pytest.mark.llm_model -@pytest.mark.parametrize("model_name", test_models) -def test_dummy_fp16_causal_lm_pytorch_vs_kv_vs_ai100(model_name, manual_cleanup): - torch.manual_seed(42) - custom_config = model_config_dict[model_name] - hf_config = AutoConfig.from_pretrained( - model_name, - trust_remote_code=model_name in ModelConfig.EXTERNAL_MODELS, - **custom_config.get("additional_params", {}), - ) - if model_name in ModelConfig.QUANTIZED_MODELS: - n_layer = get_custom_n_layers(model_name) - check_causal_lm_pytorch_vs_kv_vs_ai100( - model_name, n_layer=n_layer, torch_dtype=torch.float16, manual_cleanup=manual_cleanup - ) - else: - check_causal_lm_pytorch_vs_kv_vs_ai100( - model_name, config=hf_config, torch_dtype=torch.float16, manual_cleanup=manual_cleanup - ) diff --git a/tests/transformers/models/causal_lm_models/test_kv_replicate.py b/tests/transformers/models/causal_lm_models/test_kv_replicate.py new file mode 100644 index 0000000000..80909ce00d --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_kv_replicate.py @@ -0,0 +1,79 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + + +import pytest + +from .check_causal_models import check_kv_repeat_causal_lm_pytorch_vs_ai100 +from .config import ( + COSINE_SIMILARITY_THRESHOLD, + QEFF_TEST_PROFILE, + causal_lm_models_dict, + compile_params, + export_params, + generate_params, + test_models_causal, + transform_params, +) + + +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_kv_replicate_cb(model_name): + """Verify that FP16 export and compilation succeed with KV-head replication in continuous-batching mode. + + Applies ``ReplicateKVHeadTransform`` with continuous batching enabled, + exports in FP16, and compiles to a QPC. Asserts that ``qconfig.json`` is + present. Inference is not run. + + Args: + model_name: HuggingFace model identifier. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_kv_repeat_causal_lm_pytorch_vs_ai100( + model_name, + continuous_batching=True, + transform_params=transform_params, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_generate_kv_replicate_cb(model_name): + """Verify end-to-end FP16 inference quality with KV-head replication in continuous-batching mode. + + Applies ``ReplicateKVHeadTransform`` with continuous batching enabled, + exports in FP16, compiles to a QPC, runs inference on the AIC device, and + checks that the cosine similarity between the AIC output sequences and the + HF PyTorch reference sequences meets the configured threshold for every + batch slot. + + Args: + model_name: HuggingFace model identifier. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_kv_repeat_causal_lm_pytorch_vs_ai100( + model_name, + continuous_batching=True, + transform_params=transform_params, + export_params=export_params, + compile_params=compile_params, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) diff --git a/tests/transformers/models/causal_lm_models/test_mdp.py b/tests/transformers/models/causal_lm_models/test_mdp.py new file mode 100644 index 0000000000..7b5221508f --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_mdp.py @@ -0,0 +1,79 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + + +import pytest + +from .check_causal_models import check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100 +from .config import ( + COSINE_SIMILARITY_THRESHOLD, + QEFF_TEST_PROFILE, + causal_lm_models_dict, + compile_params, + export_params, + generate_params, + test_models_causal, + transform_params, +) + + +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_mdp(model_name): + """Verify that FP16 export and MDP compilation succeed for the given partition count. + + Exports the model to ONNX in FP16 with ONNX subfunctions enabled, then + compiles with ``mdp_num_partitions`` pipeline-parallel MDP partitions. + Asserts that ``qconfig.json`` is present. Inference is not run. + + Args: + model_name: HuggingFace model identifier. + mdp_num_partitions: Number of pipeline-parallel MDP partitions (2 or 4). + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + transform_params=transform_params, + export_params=export_params, + compile_params={**compile_params, "mdp_num_partitions": 2, "num_devices": 2}, + generate_params=generate_params, + continuous_batching=False, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_mdp_cb(model_name): + """Verify that FP16 export and MDP compilation succeed in continuous-batching mode. + + Exports in FP16 with ONNX subfunctions and continuous batching enabled, + compiles with ``mdp_num_partitions`` pipeline-parallel MDP partitions, and + asserts that ``qconfig.json`` is present. Inference is not run. + + Args: + model_name: HuggingFace model identifier. + mdp_num_partitions: Number of pipeline-parallel MDP partitions (2 or 4). + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=transform_params, + export_params=export_params, + compile_params={**compile_params, "mdp_num_partitions": 2, "num_devices": 2}, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) diff --git a/tests/transformers/models/causal_lm_models/test_prefix_caching.py b/tests/transformers/models/causal_lm_models/test_prefix_caching.py new file mode 100644 index 0000000000..1eaaf88053 --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_prefix_caching.py @@ -0,0 +1,97 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + + +import pytest + +from .check_causal_models import check_prefix_caching_inference +from .config import ( + COSINE_SIMILARITY_THRESHOLD, + QEFF_TEST_PROFILE, + causal_lm_models_dict, + compile_params, + export_params, + generate_params, + test_models_causal, + transform_params, +) + + +@pytest.mark.non_qaic +@pytest.mark.llm +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_prefix_caching_cb(model_name): + """Verify that FP16 export and compilation succeed for a prefix-caching continuous-batching model. + + Compiles the model with a KV-cache batch size larger than the active batch + (``full_batch_size=2``, ``kv_cache_batch_size=4``) so that prefill KV state + can be retained and reused across requests. Asserts that ``qconfig.json`` + is present. Inference is not run. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + temp_compile_params = { + **compile_params, + "full_batch_size": 2, + "kv_cache_batch_size": 4, + } + + check_prefix_caching_inference( + model_name, + transform_params=transform_params, + export_params=export_params, + compile_params=temp_compile_params, + generate_params=generate_params, + continuous_batching=True, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.qaic +@pytest.mark.llm +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_generate_prefix_caching_cb(model_name): + """Verify prefix-caching KV-reuse correctness end-to-end on the AIC device. + + Compiles with ``full_batch_size=2`` and ``kv_cache_batch_size=4``, then runs + a multi-stage inference scenario that exercises KV-cache reuse: + + 1. Generates outputs for two prompts (prefix + suffix1) on batch slots 0 and 1. + 2. Manually prefills the same prompts on slots 2 and 3 and verifies that the + step-by-step decode tokens match the session's accumulated ``generated_ids`` + via cosine similarity. + 3. Re-runs prefill on slot 0 with a modified input sharing the same prefix, + and verifies that the cached prefix produces consistent logits via cosine + similarity. + 4. Re-runs decode on slot 1 from the cached prefix state and verifies that + the output sequence matches the original baseline via cosine similarity. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + prefixes = ["Once upon a time ", "Once upon a time "] + suffixes1 = ["in a land far away", "there was a small village"] + suffixes2 = ["a little girl", "in a bustling city"] + temp_generate_params = {**generate_params, "prefixes": prefixes, "suffixes1": suffixes1, "suffixes2": suffixes2} + temp_compile_params = { + **compile_params, + "full_batch_size": 2, + "kv_cache_batch_size": 4, + } + + check_prefix_caching_inference( + model_name, + transform_params=transform_params, + export_params=export_params, + compile_params=temp_compile_params, + generate_params=temp_generate_params, + continuous_batching=True, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) diff --git a/tests/transformers/models/causal_lm_models/test_speculative.py b/tests/transformers/models/causal_lm_models/test_speculative.py new file mode 100644 index 0000000000..17cebf5401 --- /dev/null +++ b/tests/transformers/models/causal_lm_models/test_speculative.py @@ -0,0 +1,75 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + + +import pytest + +from QEfficient.utils.constants import Constants + +from .check_causal_models import check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100 +from .config import ( + COSINE_SIMILARITY_THRESHOLD, + QEFF_TEST_PROFILE, + causal_lm_models_dict, + compile_params, + export_params, + generate_params, + test_models_causal, + transform_params, +) + + +@pytest.mark.llm +@pytest.mark.non_qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_speculative_cb(model_name): + """Verify that FP16 export and compilation succeed with speculative decoding and continuous batching. + + Compiles the model as a Target Language Model (TLM) with the configured number + of speculative tokens and continuous batching enabled. Asserts that + ``qconfig.json`` is present. Inference is not run; this test validates the + export-to-compile pipeline for the speculative-decoding configuration. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=transform_params, + export_params=export_params, + compile_params={**compile_params, "num_speculative_tokens": Constants.NUM_SPECULATIVE_TOKENS}, + generate_params=generate_params, + export_compile_only=True, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) + + +@pytest.mark.llm +@pytest.mark.qaic +@pytest.mark.parametrize("model_name", test_models_causal) +def test_fp16_export_compile_generate_speculative_cb(model_name): + """Verify end-to-end FP16 inference quality with speculative decoding and continuous batching. + + Compiles the model as a TLM with the configured number of speculative tokens + and continuous batching enabled, runs inference on the AIC device, and checks + that the cosine similarity between the AIC output token sequences and the HF + PyTorch reference sequences meets the configured threshold for every batch slot. + """ + if causal_lm_models_dict.get(model_name, None) == model_name and QEFF_TEST_PROFILE == "tiny_model": + pytest.skip("Skipping it is not a tiny model and will run in nightly tests.") + + check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100( + model_name, + continuous_batching=True, + transform_params=transform_params, + export_params=export_params, + compile_params={**compile_params, "num_speculative_tokens": Constants.NUM_SPECULATIVE_TOKENS}, + generate_params=generate_params, + export_compile_only=False, + cosine_similarity_threshold=COSINE_SIMILARITY_THRESHOLD, + ) diff --git a/tests/transformers/models/causal_lm_models/utils.py b/tests/transformers/models/causal_lm_models/utils.py new file mode 100644 index 0000000000..47174a75cb --- /dev/null +++ b/tests/transformers/models/causal_lm_models/utils.py @@ -0,0 +1,238 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import time + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.cache_utils import DynamicCache + +from QEfficient import QEFFAutoModelForCausalLM +from QEfficient.generation.cloud_infer import QAICInferenceSession +from QEfficient.transformers.quantizers import replace_transformers_quantizers, undo_transformers_quantizers + + +def _make_model(model_id: str) -> AutoModelForCausalLM: + """Create a tiny model from a dummy config — no weight download required. + + A fixed seed ensures the weights are reproducible across test runs so that + the QAIC-compiled model (which may be cached on disk) always matches the + in-process PyTorch model used for reference comparisons. + + Weights are scaled to std≈0.02 (matching real transformer init) so that + intermediate activations stay small and float16 rounding errors on QAIC + remain within the 5e-2 tolerance used for logit accuracy checks. + """ + torch.manual_seed(42) + model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="eager", torch_dtype=torch.float32) + with torch.no_grad(): + for param in model.parameters(): + param.mul_(0.02) + return model + + +def _prepare_inputs(model_id, prompt, prefill_seq_len): + """Tokenize prompt for both HF reference and QAIC-formatted inputs.""" + tokenizer = AutoTokenizer.from_pretrained(model_id) + + _raw_padded = tokenizer(prompt, return_tensors="np", padding=True) + padded_len = _raw_padded["input_ids"].shape[1] + num_chunks = -(padded_len // -prefill_seq_len) # ceil divide + padded_len = num_chunks * prefill_seq_len + + raw_inputs = tokenizer(prompt, return_tensors="np", padding="max_length", max_length=padded_len) + raw_inputs["position_ids"] = np.where(raw_inputs.pop("attention_mask"), np.arange(padded_len), -1) + raw_inputs.pop("token_type_ids", None) + + hf_inputs = tokenizer(prompt, return_tensors="pt") + qeff_inputs = {k: torch.from_numpy(v) for k, v in raw_inputs.items()} + return tokenizer, hf_inputs, qeff_inputs, raw_inputs, num_chunks + + +def _make_kv_cache(config, prefill_seq_len, ctx_len, sliding_window=False): + """Build zeroed past_key_values list for QEff model forward pass.""" + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + past_key_values = [] + for i in range(config.num_hidden_layers): + cache_len = (config.sliding_window if i % 2 == 0 else prefill_seq_len) if sliding_window else ctx_len + shape = (1, config.num_key_value_heads, cache_len, head_dim) + past_key_values.append((torch.zeros(shape, dtype=torch.float32), torch.zeros(shape, dtype=torch.float32))) + return past_key_values + + +def _default_compile_kwargs(prefill_seq_len, ctx_len, num_cores, **overrides): + """Return compile kwargs with defaults, allowing per-call overrides.""" + kwargs = dict( + prefill_seq_len=prefill_seq_len, + ctx_len=ctx_len, + num_cores=num_cores, + mxfp6_matmul=False, + mxint8_kv_cache=False, + num_devices=1, + mos=1, + aic_enable_depth_first=True, + num_speculative_tokens=None, + ) + kwargs.update(overrides) + return kwargs + + +def _run_hf_prefill(model_id, hf_inputs): + """Run a single forward pass through the HF model and return (model, config, out).""" + replace_transformers_quantizers() + model = _make_model(model_id) + config = model.config + out = model(**hf_inputs, past_key_values=DynamicCache(config=config)) + undo_transformers_quantizers() + return model, config, out + + +def _run_qeff_prefill( + model, + qeff_inputs, + prefill_seq_len, + ctx_len, + *, + sliding_window, # use alternating sliding-window cache lengths (gpt_oss) + enable_chunking=False, # feed the sequence in prefill_seq_len-sized chunks sequentially + num_chunks=None, +): + """Apply disaggregated-prefill transforms and return (qeff_model, config, qeff_out).""" + + qeff_model = QEFFAutoModelForCausalLM(model) + qeff_model.prefill(enable=True, enable_chunking=enable_chunking) + config = qeff_model.model.config + qeff_inputs["past_key_values"] = _make_kv_cache(config, prefill_seq_len, ctx_len, sliding_window=sliding_window) + + if enable_chunking: + num_chunks = 1 if num_chunks is None else num_chunks + for i in range(num_chunks): + chunk_inputs = qeff_inputs.copy() + chunk_inputs["input_ids"] = qeff_inputs["input_ids"][:, i * prefill_seq_len : (i + 1) * prefill_seq_len] + chunk_inputs["position_ids"] = qeff_inputs["position_ids"][ + :, i * prefill_seq_len : (i + 1) * prefill_seq_len + ] + qeff_out = qeff_model.model(**chunk_inputs) + qeff_inputs["past_key_values"] = qeff_out["past_key_values"] + else: + qeff_out = qeff_model.model(**qeff_inputs) + return qeff_model, config, qeff_out + + +def _run_qpc_prefill( + qeff_inputs, + prefill_qpc_path, + config, + prefill_seq_len, + enable_chunking=False, # feed the sequence in prefill_seq_len-sized chunks sequentially + num_chunks=None, + skip_past_buffers=False, # skip host-side past_ I/O so the device retains KV state across chunks +): + """Run a compiled prefill QPC on QAIC and return (qpc_out, qeff_inputs_np).""" + + if isinstance(next(iter(qeff_inputs.values())), torch.Tensor): + qeff_inputs.pop("past_key_values", None) + qeff_inputs_np = {k: v.detach().numpy() for k, v in qeff_inputs.items()} + else: + qeff_inputs_np = {k: v for k, v in qeff_inputs.items()} + + session = QAICInferenceSession(prefill_qpc_path) + logits_out_placeholder = np.zeros((1, 1, config.vocab_size), dtype=np.float32) + session.set_buffers({"logits": logits_out_placeholder}) + + if enable_chunking: + if skip_past_buffers: + session.skip_buffers([x for x in session.input_names + session.output_names if x.startswith("past_")]) + session.set_buffers({"logits": np.zeros((1, 1, config.vocab_size), dtype=np.float32)}) + + t0 = time.time() + if enable_chunking: + num_chunks = 1 if num_chunks is None else num_chunks + for i in range(num_chunks): + chunk_inputs = qeff_inputs_np.copy() + chunk_inputs["input_ids"] = qeff_inputs_np["input_ids"][:, i * prefill_seq_len : (i + 1) * prefill_seq_len] + chunk_inputs["position_ids"] = qeff_inputs_np["position_ids"][ + :, i * prefill_seq_len : (i + 1) * prefill_seq_len + ] + qpc_out = session.run(chunk_inputs) + else: + qpc_out = session.run(qeff_inputs_np) + print(f"time for prefill_run={time.time() - t0} sec\n") + del session + return qpc_out, qeff_inputs_np + + +def _rotate_sliding_kv(qpc_out, decode_inputs, config): + """Copy prefill KV outputs into decode_inputs, rotating sliding-window layers. + + Even-indexed layers use sliding-window attention; their cache must be rotated + so that the oldest entries are at the front when the position exceeds the window. + """ + for i in range(config.num_hidden_layers): + if i % 2 == 0 and decode_inputs["position_ids"] >= config.sliding_window: + k = qpc_out[f"past_key.{i}_RetainedState"] + v = qpc_out[f"past_value.{i}_RetainedState"] + rot_offset = config.sliding_window - decode_inputs["position_ids"][0][0] % config.sliding_window + decode_inputs[f"past_key.{i}"] = np.concatenate((k[:, :, rot_offset:, :], k[:, :, :rot_offset, :]), axis=-2) + decode_inputs[f"past_value.{i}"] = np.concatenate( + (v[:, :, rot_offset:, :], v[:, :, :rot_offset, :]), axis=-2 + ) + else: + decode_inputs[f"past_key.{i}"] = qpc_out[f"past_key.{i}_RetainedState"] + decode_inputs[f"past_value.{i}"] = qpc_out[f"past_value.{i}_RetainedState"] + + +def _prefix_caching_inference( + model_id, + config, + prefill_qpc_path, + decode_qpc_path, + prompt, + decode_batch_id, # which slot in the decode KV cache batch to write into + prefill_seq_len, +): + generation_len = 5 + tokenizer, _, _, inputs, num_chunks = _prepare_inputs(model_id, prompt, prefill_seq_len) + inputs["batch_index"] = np.array([[decode_batch_id]], dtype=np.int64) + + # QPC prefill run on QAIC + qpc_out, inputs = _run_qpc_prefill( + inputs, prefill_qpc_path, config, prefill_seq_len, enable_chunking=True, num_chunks=num_chunks + ) + + decode_inputs = { + "input_ids": np.argmax(qpc_out["logits"]).reshape(1, 1), + "position_ids": np.max(inputs["position_ids"]).reshape(1, 1) + 1, + "batch_index": inputs["batch_index"], + } + qpc_outputs = [decode_inputs["input_ids"][0][0]] + _rotate_sliding_kv(qpc_out, decode_inputs, config) + + # QPC decode run on QAIC + decode_session = QAICInferenceSession(decode_qpc_path) + decode_session.set_buffers({"logits": np.zeros((1, 1, config.vocab_size), dtype=np.float32)}) + decode_out = decode_session.run(decode_inputs) + pos_id = np.max(decode_inputs["position_ids"]).reshape(1, 1) + 1 + + for _ in range(generation_len - 1): + loop_inputs = { + "input_ids": np.argmax(decode_out["logits"]).reshape(1, 1), + "position_ids": pos_id, + "batch_index": inputs["batch_index"], + } + for j in range(config.num_hidden_layers): + loop_inputs[f"past_key.{j}"] = decode_out[f"past_key.{j}_RetainedState"] + loop_inputs[f"past_value.{j}"] = decode_out[f"past_value.{j}_RetainedState"] + qpc_outputs.append(loop_inputs["input_ids"][0][0]) + decode_out = decode_session.run(loop_inputs) + pos_id += 1 + + print("QPC Outputs (AIC): \n") + print("Prompt:", repr(prompt)) + print("Completion:", repr(tokenizer.decode(qpc_outputs))) + return qpc_out, qpc_outputs